alisp.cr

module Lisp
  alias Value = Symbol | String | Int64 | Bool | Nil | Builtin | List

  record Symbol, value : String

  struct List
    getter list : Array(Value)

    def initialize(@list = Array(Value).new)
    end

    forward_missing_to list

    def to_s(io)
      io << '('
      @list.each { |value| io << value }
      io << '('
    end
  end

  class ParseError < Exception; end

  def self.parse_error(msg : String)
    raise ParseError.new(msg)
  end

  class Token
    property type, value

    def initialize(@type : String, @value : String)
    end

    def to_s
      "Token(#{@type}, #{@value})"
    end
  end

  class Lexer
    class LexError < Exception; end

    @i = 0_i64
    @s = ""

    def initialize
    end

    def skip_ws
      while @i < @s.size && @s[@i].ascii_whitespace?
        @i += 1
      end
    end

    def peek
      @s[@i]
    end

    def parse_error(msg)
      Lisp.parse_error(msg)
    end

    def lex(s : String) : Array(Token)
      @s, @i = s, 0_i64
      skip_ws
      raise LexError.new("empty input") if @i == @s.size
      buffer = [] of Token
      while @i < @s.size
        skip_ws
        break if @i == @s.size
        char = peek
        if char == '('
          buffer << Token.new("LPAREN", "(")
          @i += 1
        elsif char == ')'
          buffer << Token.new("RPAREN", ")")
          @i += 1
        elsif is_initial? char
          buffer << lex_ident
        elsif char.number?
          buffer << lex_number
        elsif char == '"'
          buffer << lex_string
        else
          raise LexError.new("unknown token")
        end
      end
      buffer
    end

    def lex_string
      # TODO: add support for escape sequences
      @i += 1
      closed = false
      buffer = ""
      while @i < @s.size
        char = peek
        if char == '"'
          @i += 1
          closed = true
          break
        else
          buffer += char
          @i += 1
        end
      end
      raise LexError.new("unterminated string") unless closed
      Token.new("STRING", buffer)
    end

    def is_initial?(c : Char)
      # Is this a valid initial character in a symbol?
      c.ascii_letter? || "+-*/%!?#".includes?(c)
    end

    def lex_ident
      buffer = ""
      buffer += peek
      @i += 1
      while @i < @s.size && (is_initial?(peek) || peek.number?)
        buffer += peek
        @i += 1
      end
      if buffer[0] == '#'
        parse_error("# is not a valid symbol") if buffer.size == 1
        case buffer[1]
        when 't'
          return Token.new("TRUE", "true")
        when 'f'
          return Token.new("FALSE", "false")
        else
          parse_error("invalid sequence #{buffer}")
        end
      elsif buffer == "nil"
        return Token.new("NIL", "nil")
      elsif buffer[0] == '-' && buffer[1] && buffer[1].number?
        return Token.new("NEGINT", buffer)
      else
        Token.new("IDENT", buffer)
      end
    end

    def lex_number
      if peek == '-'
        neg = true
        @i += 1
      else
        neg = false
      end
      buffer = ""
      while @i < @s.size && peek.number?
        buffer += peek
        @i += 1
      end
      if neg
        Token.new("NEGINT", buffer)
      else
        Token.new("POSINT", buffer)
      end
    end
  end

  class Parser
    class ParseError < Exception; end

    @buffer = [] of Value
    @tokens = [] of Token

    def parse(@tokens) : Value
      @buffer = [] of Value
      until @tokens.empty?
        @buffer << parse_elem
      end
      if @buffer.size == 1
        @buffer.first
      else
        List.new(@buffer)
      end
    end

    def parse_elem
      tok = @tokens.shift
      case tok.type
      when "IDENT"
        Symbol.new(tok.value)
      when "STRING"
        tok.value
      when "POSINT"
        tok.value.to_i64
      when "TRUE"
        true
      when "FALSE"
        false
      when "NIL"
        nil
      when "NEGINT"
        tok.value.to_i64
      when "LPAREN"
        parse_list
      else
        parse_error("unknown token #{tok.to_s}")
      end
    end

    def parse_error(msg)
      raise ParseError.new(msg)
    end

    def parse_list
      buffer = [] of Value
      closing = false
      until @tokens.empty?
        tok = @tokens.first
        if tok.type == "RPAREN"
          closing = true
          @tokens.shift
          break
        end
        buffer << parse_elem
      end
      parse_error("unclosed list") unless closing
      List.new(buffer)
    end
  end

  class Env
    property parent, env

    def initialize(@parent : Env? = nil)
      @env = {} of String => Value
    end

    def [](key)
      if @env.has_key? key
        @env[key]
      elsif parent = @parent
        parent[key]
      else
        raise KeyError.new("unbound variable #{key}")
      end
    end

    def []=(key, value)
      @env[key] = value
    end

    def set!(key, value)
      # The difference here is that set! will update any existing binding
      # in the current or enclosing env, and only introduce a new binding
      # if none exists, while def introduces a new binding
      # in the current scope and doesn't touch enclosing scopes.
      env = self
      while env
        if env.env.has_key? key
          env[key] = value
          return
        else
          env = env.parent
        end
      end
      self[key] = value
    end
  end

  alias BuiltinProc = Proc(Interpreter, Array(Value), Value?)

  class Builtin
    property fn : BuiltinProc
    property arity : Int32
    property name : String

    def initialize(@name, @arity, @fn)
    end

    def to_s
      arity = @arity == -1 ? "variadic" : @arity
      "#<builtin fn #{name}(arity: #{arity})>"
    end
  end

  DefaultEnv = Env.new

  macro expect(type, arg)
    unless {{arg}}.is_a?({{type}})
      error("Type error: expected " + {{type.stringify}} + ", got " + {{arg}}.class.to_s)
    end
  end

  def self.error(msg)
    raise Interpreter::RuntimeError.new(msg)
  end

  def self.builtin(name, arity, &fn : BuiltinProc)
    f = Builtin.new(name, arity, fn)
    DefaultEnv[name] = f
  end

  builtin("print", -1) do |i, args|
    args.each { |a| puts i.eval(a).to_s }
  end

  builtin("size", 1) do |i, args|
    first = args.first
    expect List, first
    first.size.to_i64
  end

  builtin("list", -1) do |i, args|
    List.new(args)
  end

  builtin("+", -1) do |i, args|
    sum = 0_i64
    args.map do |a|
      expect Int64, a
      sum += a
    end
    sum
  end

  builtin("at", 2) do |i, args|
    list = args[0]
    idx = args[1]
    expect List, list
    expect Int, idx

    error("index out of bounds") unless list.size > idx
    list[idx]
  end

  builtin("map", 2) do |i, args|
    list = args[0]
    fn = args[1]
    expect List, list
    expect Builtin, fn

    List.new(list.map { |v| i.eval_call_builtin(fn, [v]) })
  end

  builtin("read", 1) do |i, args|
    str = args[0]
    expect String, str
    value = Parser.new.parse(Lexer.new.lex(str))
  end

  class Interpreter
    property env : Env

    def initialize
      @env = DefaultEnv
    end

    def error(msg)
      raise Interpreter::RuntimeError.new(msg)
    end

    def eval(val : Value) : Value
      case val
      when Int64, Bool, Nil, String
        val
      when List
        eval_list val
      when Symbol
        eval_ident val
      end
    end

    def eval_list(val)
      if val.size == 0
        return val
      end
      case fn = val.first
      when Symbol
        return eval_kw(fn, val[1..-1]) if keyword?(fn)
      when String
        eval_call fn, val.list[1..-1]
      else
        raise "Invalid list head"
      end
    end

    def eval_call(fn : String, args) : Value
      fn = @env[fn]
      eval_call_builtin fn, args
    end

    def eval_call_builtin(fn : Value, args : Array(Value))
      args.map! { |x| eval(x) }
      apply fn, args
    end

    def keyword?(k)
      k.is_a?(Symbol) && {"if", "set!", "def"}.includes? k.value
    end

    def eval_kw(kw : Symbol, args : Array(Value))
      if kw.value == "if"
        error("if requires at least 2 args") if args.size < 2
        _else = args.size == 2 ? nil : args[2]
        cond = self.eval(args[0])
        if cond != nil && cond != false
          return self.eval args[1]
        elsif args.size == 3
          return self.eval _else
        else
          error("syntax error: if takes 2 or 3 args, got #{args.size}")
        end
      elsif kw.value == "def" || kw.value == "set!"
        error("def and set! require exactly 2 args") unless args.size == 2
        key = args[0]
        error("identifier must be symbol") unless key.is_a? Symbol

        val = eval(args[1])
        if kw == "def"
          @env[key.value] = val
        else
          @env.set! key.value, val
        end
      else
        error("unimplemented keyword #{kw}")
      end
    end

    class RuntimeError < Exception; end

    def apply(fn, args)
      if fn.is_a? Builtin
        argcheck(fn, args)
        return fn.fn.call self, args
      else
        error("cannot apply non-function #{fn}")
      end
    end

    def argcheck(fn, args)
      return true if fn.arity == -1
      return true if fn.arity == args.size
      error("function #{fn.name} expects #{fn.arity} argument(s), got #{args.size}")
    end

    def eval_ident(x)
      @env[x.value]
    end
  end
end