lisp.cr

module Lisp

  class List
    property list : Array(Lisp::Object)
    def initialize(@list = [] of Lisp::Object)
    end
    def to_s
      "(" + @list.map(&.to_s).join(" ") + ")"
    end
    forward_missing_to list
  end
  alias Object = (String | Int64 | List | Nil | Bool | Builtin | Lisp::Symbol)

  class ParseError < Exception; end

  def 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 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("INT", buffer)
      else
        Token.new("IDENT", buffer)
      end
    end

    def lex_number
      if peek == '-'
        @i += 1
      end
      buffer = ""
      while @i < @s.size && peek.number?
        buffer += peek
        @i += 1
      end
      Token.new("INT", buffer)
    end
  end

  class Parser
    class ParseError < Exception; end

    @buffer = [] of Lisp::Object
    @tokens = [] of Token
    def parse(@tokens) : Lisp::Object
      @buffer = [] of Lisp::Object
      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"
        Lisp::Symbol.new tok.value
      when "STRING"
        tok.value
      when "INT"
        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 Lisp::Object
      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 => Lisp::Object
    end

    def [](key)
      if @env.has_key? key
        @env[key]
      elsif @parent
        @parent.as(Env)[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(Lisp::Object), Lisp::Object?)

  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

  struct Lisp::Symbol
    property value
    def initialize(@value : String)
    end
  end

  DefaultEnv = Env.new

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

  macro typecheck(val, type)
    {{val}}.is_a? {{type}}
  end

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

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

  builtin("print", -1) do |i, args|
    lprint = ->(val : Lisp::Object) do
      if typecheck(val, String)
        puts val
      else
        puts i.write(val)
      end
    end
    args.each { |a| lprint.call a }
  end

  builtin("size", 1) do |i, args|
    a = args.first
    if a.is_a?(List)
      a.size.to_i64
    elsif a.is_a?(String)
      a.size.to_i64
    else
      type = a.class.to_s
      error("type error: size expects list or string, got #{a} : #{type}")
    end
  end

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

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

  builtin("at", 2) do |i, args|
    expect List, args[0]
    expect Int, args[1]
    idx = args[1].as(Int64)
    list = args[0].as(List)
    error("index out of bounds") unless list.size > idx
    list[idx]
  end

  builtin("map", 2) do |i, args|
    expect List, args[0]
    expect Builtin, args[1]
    list = args[0].as(List)
    fn = args[1]
    list = list.map { |v| i.eval_call_builtin(fn, [v]) }
    List.new(list)
  end

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

  builtin("write", 1) do |i, args|
    i.write args.first
  end

  class Interpreter
    property env : Env
    def initialize()
      @env = DefaultEnv
    end

    def eval(val : Lisp::Object)
      if self_evaluating? val
        return val
      elsif typecheck(val, List)
        eval_list val
      else
        eval_ident val.as(Lisp::Symbol)
      end
    end

    def eval_list(list)
      if list.size == 0
        return list
      end
      fn = list.list.first
      if keyword? fn
        eval_kw(list[0].as(Lisp::Symbol).value, list[1..-1])
      elsif typecheck(fn, Lisp::Symbol)
        eval_call fn.value, list[1..-1]
      else
        error("cannot call non-callable #{fn}")
      end
    end

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

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

    def keyword?(k)
      typecheck(k, Symbol) && ["if", "set!", "def"].includes?(k.value)
    end

    def eval_kw(kw : String, args : Array(Lisp::Object))
      if kw == "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
          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 == "def" || kw == "set!"
        error("def and set! require exactly 2 args") unless args.size == 2
        if args[0].is_a? Lisp::Symbol
          key = args[0].as(Lisp::Symbol).value
        else
          error("def and set require identifier as first arg")
        end
        val = self.eval(args[1])
        if kw == "def"
          @env[key] = val
        else
          @env.set! key, val
        end
      else
        error("unimplemented keyword #{kw}")
      end
    end

    class RuntimeError < Exception; end

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

    def argcheck(fn, args)
      fn = fn.as(Builtin)
      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

    def self_evaluating?(val)
      [Int64, Bool, String, Nil].includes? val.class
    end

    def write(v : Lisp::Object)
      case v
      when Bool
        v ? "#t" : "#f"
      when String
        v.inspect
      when Nil
        "nil"
      else
        v.to_s
      end
    end
  end
end