parser.lang

import {
  Let, Id, Eq, Num, OpenParen, CloseParen, Comma,
  Command, JsOp, Def, Return, End, DataClass, New, Dot,
  Class, Get, Str, Bang, OpenSquare, CloseSquare,
  If, Else, PlusEq, While, Do, Regex, Continue,
  Break, Export, Default, Spread, Arrow, Is, Bind,
  Not, For, Of, OpenBrace, CloseBrace, Colon, Import, From,
  Type
} from "./lexer"

export dataclass NamedLet(name: str, expr, type)
export dataclass NumExpr(value: num)
export dataclass RegexNode(value: num)
export dataclass StrExpr(value: str)
export dataclass FunctionCall(lhs_expr, args)
export dataclass IdLookup(name: str)
export dataclass CommandExpr(name: str, expr)
export dataclass JsOpExpr(lhs, type: str, rhs)
export dataclass FunctionDef(name: str, args, body, type)
export dataclass ReturnExpr(expr)
export dataclass ContinueStatement
export dataclass BreakStatement
export dataclass NewExpr(expr)
export dataclass NotExpr(expr)
export dataclass DataClassDef(name: str, properties)
export dataclass DotAccess(lhs, property)
export dataclass ClassDef(name: str, properties, entries)
export dataclass ClassInstanceEntry(name: str, expr, type)
export dataclass ClassGetterExpr(name: str, expr)
export dataclass PrefixDotLookup(name: str)
export dataclass PrefixBindLookup(name: str)
export dataclass ArrayLiteral(elements)
export dataclass ObjLit(entries)
export dataclass SimpleObjEntry(name: str, expr)
export dataclass IfStatement(branches)
export dataclass IfBranch(test_expr, body)
export dataclass ElseIfBranch(test_expr, body)
export dataclass ElseBranch(body)
export dataclass NodeAssignment(lhs_expr, rhs_expr)
export dataclass NodePlusAssignment(lhs_expr, rhs_expr)
export dataclass WhileStatement(test_expr, body)
export dataclass PropertyLookup(lhs, property)
export dataclass ExportDefault(expr)
export dataclass ExportStatement(statement)
export dataclass SpreadExpr(expr)
export dataclass SpreadArg(name: str)
export dataclass SimpleArg(name: str, type)
export dataclass SimpleDefaultArg(name: str, expr)
export dataclass ArrowFn(arg_name: str, return_expr)
export dataclass IsOperator(lhs, rhs)
export dataclass IsNotOperator(lhs, rhs)
export dataclass BoundFunctionDef(name: str, args, body)
export dataclass ForLoop(iter_name: str, iterable_expr, body)
export dataclass ParenExpr(expr)
export dataclass RenamedProperty(old_name: str, new_name: str)
export dataclass RegularObjectProperty(name: str)
export dataclass LetObjectDeconstruction(entries, rhs)
export dataclass LetArrDeconstruction(entries, rhs)
export dataclass ArrComma
export dataclass ArrNameEntry(name: str)
export dataclass ImportStatement(imports, path)
export dataclass DefaultImport(name: str, path: str)
export dataclass NamedClassArg(name: str, type)
export dataclass DefaultNamedClassArg(name: str, expr)
export dataclass ObjClassArg(entries)
export dataclass ObjClassArgEntry(name: str)
export dataclass DefaultObjClassArg(name: str, expr)
export dataclass NumT
export dataclass StrT
export dataclass BoolT
export dataclass ArrayT(type)
export dataclass TypeDef(name: str, type)
export dataclass TypeIdLookup(name: str)

class Parser(tokens)
  index = 0
  get cur_token = .tokens[.index]
  get prev_token = .tokens[.index - 1]

  def consume(TokenClass)
    if .cur_token is TokenClass
      .index += 1
      return .prev_token
    else
      assert_not_reached! "[lang:" + .cur_token.line + "] Expected " + TokenClass + ", got " + .cur_token.constructor
    end
  end


  def clone_and_parse_until(...end_tokens)
    let cloned = new Parser(.tokens)
    cloned.index = .index
    let result = cloned.parse(...end_tokens)
    .index = cloned.index
    return result
  end


  def parse(...end_tokens)
    let ast = []

    def ::has_reached_end_token
      return false if end_tokens.length === 0
      return end_tokens.some(EndTokenClass => .cur_token is EndTokenClass)
    end


    while .index < .tokens.length && !has_reached_end_token() do
      let statement_or_expr = .parse_statement()
      
      if !statement_or_expr
        statement_or_expr = .parse_expr()
        if !statement_or_expr
          console.log(.prev_token, .cur_token, .tokens[.index + 1])
          assert_not_reached! "can't parse statement or expr"
        end
      end

      if .scan(If) && .prev_token.line === .cur_token.line
        statement_or_expr = .parse_postfix_if(statement_or_expr)
      end

      ast.push(statement_or_expr)
    end

    return ast
  end


  def scan(...TokenClasses)
    let i = 0
    for let TokenClass of TokenClasses do
      if .tokens[.index + i] is not TokenClass
        return null
      end
      i += 1
    end
    return .tokens[.index + i - 1]
  end

  def parse_statement
    if .scan(Let, OpenBrace)
      return .parse_let_obj()
    else if .scan(Let, OpenSquare)
      return .parse_let_arr()
    else if .scan(Let)
      return .parse_let()
    else if .scan(Def, Bind)
      return .parse_bound_def()
    else if .scan(Def)
      return .parse_def()
    else if .scan(DataClass)
      return .parse_data_class_def()
    else if .scan(Return)
      return .parse_return()
    else if .scan(Continue)
      return .parse_continue()
    else if .scan(Break)
      return .parse_break()
    else if .scan(Class)
      return .parse_class()
    else if .scan(If)
      return .parse_if()
    else if .scan(While)
      return .parse_while()
    else if .scan(Export, Default)
      return .parse_export_default()
    else if .scan(Export)
      return .parse_export()
    else if .scan(Import, Id)
      return .parse_default_import()
    else if .scan(Import)
      return .parse_obj_import()
    else if .scan(For)
      return .parse_for_loop()
    end
  end

  def first_parse_expr
    if .scan(Num)
      return .parse_num()
    else if .scan(Regex)
      return .parse_regex()
    else if .scan(Str)
      return .parse_str()
    else if .scan(OpenParen)
      return .parse_paren_expr()
    else if .scan(Id, Arrow)
      return .parse_arrow_fn()
    else if .scan(Id)
      return .parse_id_lookup()
    else if .scan(Command)
      return .parse_command()
    else if .scan(New)
      return .parse_new_expr()
    else if .scan(Bang)
      return .parse_not_expr()
    else if .scan(OpenSquare)
      return .parse_array()
    else if .scan(OpenBrace)
      return .parse_obj_lit()
    else if .scan(Spread)
      return .parse_spread()
    else if .scan(Bind, Id)
      return .parse_prefix_bind_lookup()
    else if .scan(Dot, Id)
      return .parse_prefix_dot_lookup()
    else if .scan(Type)
      return .parse_type_def()
    end
  end

  def parse_type_def
    .consume(Type)
    let { name } = .consume(Id)
    .consume(Eq)
    let type = .parse_type_expr()
    new TypeDef(name, type)
  end

  ASSIGNABLE_NODES = [DotAccess, PrefixDotLookup, IdLookup, PropertyLookup]
  def parse_node_assignment(lhs_expr)
    .consume(Eq)
    let rhs_expr = .parse_expr()
    new NodeAssignment(lhs_expr, rhs_expr)
  end

  def parse_node_plus_assignment(lhs_expr)
    .consume(PlusEq)
    let rhs_expr = .parse_expr()
    new NodePlusAssignment(lhs_expr, rhs_expr)
  end

  def can_assign(expr)
    return (.scan(Eq) || .scan(PlusEq)) && 
      .ASSIGNABLE_NODES.includes(expr.constructor)
  end

  def parse_expr_assignment(expr)
    if .scan(Eq)
      return .parse_node_assignment(expr)
    else if .scan(PlusEq)
      return .parse_node_plus_assignment(expr)
    else
      assert_not_reached! "cannot assign"
    end
  end

  def is_dot_access = .scan(Dot) && .prev_token.line === .cur_token.line
  def is_property_lookup = .scan(OpenSquare) && .prev_token.line === .cur_token.line

  def parse_expr
    let expr = .first_parse_expr()
    return undefined if !expr

    while true do
      if .scan(JsOp)
        expr = .parse_js_op(expr)
      else if .is_dot_access()
        expr = .parse_dot_access(expr)
      else if .is_property_lookup()
        expr = .parse_property_lookup(expr)
      else if .scan(OpenParen)
        expr = .parse_function_call(expr)
      else if .scan(Is, Not)
        expr = .parse_is_not_operator(expr)
      else if .scan(Is)
        expr = .parse_is_operator(expr)
      else if .can_assign(expr)
        return .parse_expr_assignment(expr)
      else
        break
      end
    end

    return expr
  end

  def parse_default_import
    .consume(Import)
    let { name } = .consume(Id)
    .consume(From)
    let { value: path} = .consume(Str)
    new DefaultImport(name, path)
  end

  def parse_obj_import
    .consume(Import)
    .consume(OpenBrace)
    let imports = []
    while !.scan(CloseBrace) do
      let { name } = .consume(Id)
      imports.push(name)
      .consume(Comma) if !.scan(CloseBrace)
    end
    .consume(CloseBrace)
    .consume(From)
    let { value: path } = .consume(Str)
    new ImportStatement(imports, path)
  end

  def parse_paren_expr
    .consume(OpenParen)
    let expr = .parse_expr()
    .consume(CloseParen)
    new ParenExpr(expr)
  end

  def parse_for_loop
    .consume(For)
    .consume(Let)
    let { name: iter_name } = .consume(Id)
    .consume(Of)
    let iterable_expr = .parse_expr()
    .consume(Do)
    let body = .clone_and_parse_until(End)
    .consume(End)
    new ForLoop(iter_name, iterable_expr, body)
  end

  def parse_is_not_operator(lhs)
    .consume(Is)
    .consume(Not)
    let rhs = .parse_expr()
    new IsNotOperator(lhs, rhs)
  end

  def parse_is_operator(lhs)
    .consume(Is)
    let rhs = .parse_expr()
    new IsOperator(lhs, rhs)
  end

  def parse_arrow_fn
    let { name: arg_name } = .consume(Id)
    .consume(Arrow)
    let return_expr = .parse_statement() || .parse_expr()
    new ArrowFn(arg_name, return_expr)
  end

  def parse_spread
    .consume(Spread)
    let expr = .parse_expr()
    new SpreadExpr(expr)
  end

  def parse_export_default
    .consume(Export)
    .consume(Default)
    let expr = .parse_expr()
    new ExportDefault(expr)
  end
  
  def parse_export
    .consume(Export)
    let statement = .parse_statement()
    new ExportStatement(statement)
  end

  def parse_property_lookup(lhs)
    .consume(OpenSquare)
    let property = .parse_expr()
    .consume(CloseSquare)
    new PropertyLookup(lhs, property)
  end

  def parse_while
    .consume(While)
    let test_expr = .parse_expr()
    .consume(Do)
    let body = .clone_and_parse_until(End)
    .consume(End)
    new WhileStatement(test_expr, body)
  end

  def parse_postfix_if(lhs)
    .consume(If)
    let test_expr = .parse_expr()
    new IfStatement([new IfBranch(test_expr, [lhs])])
  end

  def parse_else_if
    .consume(Else)
    .consume(If)
    let test_expr = .parse_expr()
    let body = .clone_and_parse_until(Else, End)
    new ElseIfBranch(test_expr, body)
  end

  def parse_else
    .consume(Else)
    let body = .clone_and_parse_until(End)
    new ElseBranch(body)
  end

  def parse_pass_branch
    .consume(If)
    let test_expr = .parse_expr()
    let pass_body = .clone_and_parse_until(Else, End)
    new IfBranch(test_expr, pass_body)
  end

  def parse_if
    let branches = []
    branches.push(.parse_pass_branch())

    while !.scan(End) do
      if .scan(Else, If)
        branches.push(.parse_else_if())
      else if .scan(Else)
        branches.push(.parse_else())
      end
    end

    .consume(End)
    new IfStatement(branches)
  end

  def parse_obj_lit
    let entries = []
    .consume(OpenBrace)
    while !.scan(CloseBrace) do
      let { name } = .consume(Id)
      .consume(Colon)
      let expr = .parse_expr()
      entries.push(new SimpleObjEntry(name, expr))
      .consume(Comma) if .scan(Comma)
    end
    .consume(CloseBrace)
    new ObjLit(entries)
  end

  def parse_array
    let elements = []
    .consume(OpenSquare)
    while !.scan(CloseSquare) do
      elements.push(.parse_expr())
      .consume(Comma) if !.scan(CloseSquare)
    end
    .consume(CloseSquare)
    new ArrayLiteral(elements)
  end

  def parse_not_expr
    .consume(Bang)
    let expr = .parse_expr()
    new NotExpr(expr)
  end

  def parse_prefix_bind_lookup
    .consume(Bind)
    let { name } = .consume(Id)
    new PrefixBindLookup(name)
  end

  def parse_prefix_dot_lookup
    .consume(Dot)
    let { name } = .consume(Id)
    new PrefixDotLookup(name)
  end

  def parse_class_instance_entry
    let { name } = .consume(Id)
    let type = .parse_type_annotation()
    .consume(Eq)
    let expr = .parse_expr()
    new ClassInstanceEntry(name, expr, type)
  end

  def parse_getter
    .consume(Get)
    let { name } = .consume(Id)
    .consume(Eq)
    let expr = .parse_expr()
    new ClassGetterExpr(name, expr)
  end

  def parse_class_entry
    if .scan(Id, Eq) || .scan(Id, Colon)
      return .parse_class_instance_entry()
    else if .scan(Get)
      return .parse_getter()
    else if .scan(Def)
      return .parse_def()
    else
      assert_not_reached! "no class entry found"
    end
  end

  def parse_obj_args
    .consume(OpenBrace)
    let entries = []
    while !.scan(CloseBrace) do
      if .scan(Id, Eq)
        let { name } = .consume(Id)
        .consume(Eq)
        let default_value = .parse_expr()
        entries.push(new DefaultObjClassArg(name, default_value))
      else if .scan(Id)
        let { name } = .consume(Id)
        entries.push(new ObjClassArgEntry(name))
      end
      .consume(Comma) if .scan(Comma)
    end
    .consume(CloseBrace)
    new ObjClassArg(entries)
  end

  def parse_class_arg
    if .scan(Id, Eq)
      let { name } = .consume(Id)
      .consume(Eq)
      let expr = .parse_expr()
      return new DefaultNamedClassArg(name, expr)
    else if .scan(Id)
      let { name } = .consume(Id)
      let type = .parse_type_annotation()
      return new NamedClassArg(name, type)
    else if .scan(OpenBrace)
      return .parse_obj_args()
    end
  end

  def parse_class_args
    .consume(OpenParen)
    let properties = []
    while !.scan(CloseParen) do
      properties.push(.parse_class_arg())
      .consume(Comma) if .scan(Comma)
    end
    .consume(CloseParen)
    return properties
  end

  def parse_class
    .consume(Class)
    let { name } = .consume(Id)
    let properties = null
    properties = .parse_class_args() if .scan(OpenParen)
    let entries = []
    while !.scan(End) && .cur_token do
      entries.push(.parse_class_entry())
    end
    .consume(End)
    new ClassDef(name, properties, entries)
  end

  def parse_dot_access(lhs)
    .consume(Dot)
    let { name: property } = .consume(Id)
    new DotAccess(lhs, property)
  end

  def parse_new_expr
    .consume(New)
    let expr = .parse_expr()
    new NewExpr(expr)
  end

  def parse_arg_defs
    .consume(OpenParen)
    let args = []
    while .cur_token is not CloseParen do
      if .scan(Id, Eq)
        let { name: arg_name } = .consume(Id)
        .consume(Eq)
        let expr = .parse_expr()
        args.push(new SimpleDefaultArg(arg_name, expr))
      else if .scan(Id)
        let { name: arg_name } = .consume(Id)
        let type = .parse_type_annotation()
        args.push(new SimpleArg(arg_name, type))
      else if .scan(Spread)
        .consume(Spread)
        let { name: arg_name } = .consume(Id)
        args.push(new SpreadArg(arg_name))
      else if .scan(OpenBrace)
        args.push(.parse_obj_args())
      else
        assert_not_reached! "arg type not supported"
      end
      break if !.scan(Comma)
      .consume(Comma)
    end
    .consume(CloseParen)
    return args
  end

  def parse_arg_names
    .consume(OpenParen)
    let args = []
    while !.scan(CloseParen) do
      let { name } = .consume(Id)
      let type = .parse_type_annotation()
      args.push(new NamedClassArg(name, type))
      .consume(Comma) if !.scan(CloseParen)
    end
    .consume(CloseParen)
    return args
  end

  def parse_data_class_def
    .consume(DataClass)
    let { name } = .consume(Id)
    let properties = null
    properties = .parse_arg_names() if .scan(OpenParen)
    new DataClassDef(name, properties)
  end

  def parse_return
    .consume(Return)
    let expr = .parse_expr()
    new ReturnExpr(expr)
  end

  def parse_continue
    .consume(Continue)
    new ContinueStatement()
  end

  def parse_break
    .consume(Break)
    new BreakStatement()
  end

  def parse_bound_def
    .consume(Def)
    .consume(Bind)
    let { name } = .consume(Id)
    let args = []
    args = .parse_arg_defs() if .scan(OpenParen)
    let body = .clone_and_parse_until(End)
    .consume(End)
    new BoundFunctionDef(name, args, body)
  end

  def parse_def
    .consume(Def)
    let { name } = .consume(Id)
    let args = []
    args = .parse_arg_defs() if .scan(OpenParen)
    let type = .parse_type_annotation()
    let body = []
    if .scan(Eq)
      .consume(Eq)
      let expr = .parse_expr()
      expr = new ReturnExpr(expr) if expr is not ReturnExpr
      body.push(expr)
    else
      body = .clone_and_parse_until(End)
      let non_returnable = [
        ReturnExpr,
        IfStatement,
        NamedLet
      ]
      if body.length > 0 && non_returnable.every(Klass => body.at(-1) is not Klass)
        body = [...body.slice(0, -1), new ReturnExpr(body.at(-1))]
      end
      .consume(End)
    end
    new FunctionDef(name, args, body, type)
  end

  def parse_js_op(lhs_expr)
    let { op } = .consume(JsOp)
    let rhs_expr = .parse_expr()
    new JsOpExpr(lhs_expr, op, rhs_expr)
  end

  def parse_command
    let { name } = .consume(Command)
    let expr = .parse_expr()
    assert_not_reached! "no expr" if !expr
    new CommandExpr(name, expr)
  end

  def parse_id_lookup
    let { name } = .consume(Id)
    new IdLookup(name)
  end

  def parse_function_call(expr)
    .consume(OpenParen)
    let args = []
    while .cur_token is not CloseParen do
      let arg = .parse_expr()
      args.push(arg)
      break if !.scan(Comma)
      .consume(Comma)
    end
    .consume(CloseParen)
    new FunctionCall(expr, args)
  end

  def parse_str
    let { value } = .consume(Str)
    new StrExpr(value)
  end

  def parse_regex
    let { value } = .consume(Regex)
    new RegexNode(value)
  end

  def parse_num
    let { value } = .consume(Num)
    new NumExpr(value)
  end

  def parse_renamed_property
    let { name: old_name } = .consume(Id)
    .consume(Colon)
    let { name: new_name } = .consume(Id)
    new RenamedProperty(old_name, new_name)
  end

  def parse_regular_object_property
    let { name } = .consume(Id)
    new RegularObjectProperty(name)
  end

  def parse_obj_deconstruction_entry
    if .scan(Id, Colon)
      return .parse_renamed_property()
    else if .scan(Id)
      return .parse_regular_object_property()
    else 
      assert_not_reached! "unknown object deconstruction entry"
    end
  end

  def parse_let_arr
    .consume(Let)
    .consume(OpenSquare)
    let entries = []
    while !.scan(CloseSquare) do
      if .scan(Comma)
        .consume(Comma)
        entries.push(new ArrComma())
      end
      continue if .scan(Comma)
      if .scan(Id)
        let { name } = .consume(Id)
        entries.push(new ArrNameEntry(name))
      else
        entries.push(.parse_obj_deconstruction())
      end
    end
    .consume(CloseSquare)
    .consume(Eq)
    let rhs = .parse_expr()
    new LetArrDeconstruction(entries, rhs)
  end

  def parse_obj_deconstruction
    .consume(OpenBrace)
    let entries = []
    while !.scan(CloseBrace) do
      entries.push(.parse_obj_deconstruction_entry())
      .consume(Comma) if !.scan(CloseBrace)
    end
    .consume(CloseBrace)
    return entries
  end

  def parse_let_obj
    .consume(Let)
    let entries = .parse_obj_deconstruction()
    .consume(Eq)
    let rhs = .parse_expr()
    new LetObjectDeconstruction(entries, rhs)
  end

  def parse_type_expr
    assert! .scan(Id)
    if .scan(Id).name === "num"
      .consume(Id)
      return new NumT()
    else if .scan(Id).name === "str"
      .consume(Id)
      return new StrT()
    else if .scan(Id).name === "bool"
      .consume(Id)
      return new BoolT()
    else if .scan(Id).name === "Array"
      .consume(Id)
      let [type] = .parse_type_params()
      return new ArrayT(type)
    else if .scan(Id)
      let { name } = .consume(Id)
      return new TypeIdLookup(name)
    else
      assert_not_reached! "type not implemented " + this.cur_token.name
    end
  end

  def parse_type_annotation
    return null if !.scan(Colon)
    .consume(Colon)

    .parse_type_expr()
  end

  def parse_type_params
    let types = []
    .consume(OpenSquare)
    while !.scan(CloseSquare) do
      types.push(.parse_type_expr())
      .consume(Comma) if .scan(Comma)
    end
    .consume(CloseSquare)

    types
  end

  def parse_let
    .consume(Let)
    let { name } = .consume(Id)
    let type = .parse_type_annotation()
    .consume(Eq)
    let expr = .parse_expr()
    new NamedLet(name, expr, type)
  end
end

export default Parser