type_checker.lang

import {
  IdLookup, NamedLet, NumExpr, FunctionCall, CommandExpr,
  JsOpExpr, FunctionDef, ReturnExpr, DataClassDef,
  NewExpr, DotAccess, ClassDef, ClassInstanceEntry,
  ClassGetterExpr, PrefixDotLookup, StrExpr,
  NotExpr, ArrayLiteral, IfStatement, NodeAssignment,
  NodePlusAssignment, WhileStatement, RegexNode,
  ContinueStatement, BreakStatement, IfBranch,
  ElseIfBranch, ElseBranch, PropertyLookup, ExportDefault,
  ExportStatement, SpreadExpr, SimpleArg, SpreadArg,
  ArrowFn, IsOperator, BoundFunctionDef, ForLoop,
  IsNotOperator, ParenExpr, LetObjectDeconstruction,
  RegularObjectProperty, RenamedProperty, ImportStatement,
  DefaultImport, LetArrDeconstruction, ArrNameEntry,
  ArrComma, DefaultObjClassArg, NamedClassArg, ObjClassArg,
  SimpleDefaultArg, ObjLit, SimpleObjEntry, PrefixBindLookup,
  NumT, StrT, ArrayT, TypeDef, TypeIdLookup, BoolT
} from "./parser"
import Lexer from "./lexer"
import Parser from "./parser"
import fs from "fs"

dataclass FnT(args, return_type)
dataclass AnyT
dataclass NullT
dataclass ObjT(properties)
dataclass UnionT(types)
dataclass DataClassT(properties)
dataclass ClassT(methods, getters, properties)
dataclass CombinedT(root_t, properties)
dataclass RegexT

let BUILTIN_TYPES = {
  console: new ObjT({ log: new FnT(new AnyT(), new NullT()) })
  process: new ObjT({ argv: new ArrayT(new StrT()) }),
  RegExp: new RegexT(),
  true: new BoolT(),
  false: new BoolT()
}

let Buffer = new ObjT({ toString: new FnT([], new StrT()) })

let BUILTIN_PACKAGES = {
  fs: new ObjT({
    readFileSync: new FnT([{ type: new StrT() }], Buffer)
  })
}

class TypeChecker(ast, types = BUILTIN_TYPES, self = {})
  return_types = []

  def check
    for let statement of .ast do
      .check_statement(statement)
    end
    null
  end

  def check_statement(node, is_exported = false)  
    if node is NamedLet
      .check_named_let(node, is_exported)
    else if node is FunctionDef
      .check_function_def(node)
    else if node is ReturnExpr
      .check_expr(node.expr)
    else if node is FunctionCall
      .check_function_call(node)
    else if node is LetArrDeconstruction
      .check_let_arr_deconstruction(node)
    else if node is DataClassDef
      .check_data_class_def(node, is_exported)
    else if node is TypeDef
      .check_type_def(node)
    else if node is DefaultImport
      .check_default_import(node)
    else if node is ExportStatement
      .check_export_statement(node)
    else if node is ClassDef
      .check_class_def(node)
    else if node is IfStatement
      .check_if_statement(node)
    else if node is NodePlusAssignment
      .check_node_plus_assignment(node)
    else if node is NodeAssignment
      .check_node_assignment(node)
    else
      assert_not_reached! "Unknown statement " + node.constructor.name
    end
  end

  def check_node_assignment({ lhs_expr, rhs_expr })
    let lhs_t = .infer(lhs_expr)
    let rhs_t = .infer(rhs_expr)
    console.log(lhs_expr, rhs_expr, lhs_t, rhs_t)
    assert! .is_match(lhs_t, rhs_t)
    null
  end

  def check_node_plus_assignment({ lhs_expr, rhs_expr })
    let lhs_t = .infer(lhs_expr)
    let rhs_t = .infer(rhs_expr)
    assert! lhs_t is NumT
    assert! rhs_t is NumT
    null
  end

  def check_if_branch({ test_expr, body })
    .check_expr(test_expr)
    assert! .infer(test_expr) is BoolT

    let tc = new TypeChecker(body, .types, .self)
    tc.check()
    .return_types = .return_types.concat(tc.return_types)
    null
  end

  def check_if_statement_branch(branch)
    if branch is IfBranch
      .check_if_branch(branch)
    else if branch is ElseBranch
      assert_not_reached! "unimplemented"
    else if branch is ElseIfBranch
      assert_not_reached! "unimplemented"
    else
      console.log(branch)
      assert_not_reached! "unknown if branch type"
    end
  end

  def check_if_statement({ branches })
    branches.map(::check_if_statement_branch)
  end

  def check_class_arg(node)
    assert! node is NamedClassArg
    node.type || new AnyT()
  end

  def check_class_instance_entry({ name, expr, type })
    type = type || .infer(expr)
    [name, type]
  end

  def check_class_getter_expr({ name, expr })  
    [name, .infer(expr)]
  end

  def check_class_entry(entry)
    if entry is ClassInstanceEntry
      return .check_class_instance_entry(entry)      
    else if entry is ClassGetterExpr
      return .check_class_getter_expr(entry)
    else if entry is FunctionDef
      let result = .check_function_def(entry)
      assert! result.length === 1
      return [entry.name, result[0]]
    else
      console.log(entry)
      assert_not_reached! "unknown class entry"
    end
  end

  def infer_prefix_dot_lookup({ name })
    assert! .self
    .self[name]
  end

  def check_class_def({ name, properties, entries })
    for let property of properties do
      .self[property.name] = .check_class_arg(property)
    end
    for let entry of entries do
      let [name, type] = .check_class_entry(entry)
      .self[name] = type
    end
    
    assert_not_reached! "class def not implemented"
  end

  def check_export_statement({ statement })
    assert! statement is DataClassDef
    .check_statement(statement, true)
  end

  def check_default_import({ name, path })
    if BUILTIN_PACKAGES[path]
      .types[name] = BUILTIN_PACKAGES[path]
    else if path.startsWith("./")
      let code = fs.readFileSync(path + ".lang").toString()
      let tokens = new Lexer(code).tokenize()
      let ast = new Parser(tokens).parse()
      let tc = new TypeChecker(ast).check()
      console.log(tc.types)
      assert! false
    else
      assert_not_reached! "unknown package " + path
    end
  end

  def check_type_def({ name, type })
    .types[name] = type
    null
  end

  def check_data_class_def({ name, properties }, is_exported)
    assert! properties.every(p => p is NamedClassArg)
    let property_types = properties.map(p => [p.name, p.type])
    .types[name] = new DataClassT(property_types)
    .types[name].exported = is_exported
    null
  end

  def check_let_arr_deconstruction({ entries, rhs })
    let { type } = .infer(rhs)
    for let entry of entries do
      continue if entry is ArrComma
      if entry is ArrNameEntry
        .types[entry.name] = type
      else
        console.log(entry)
        assert_not_reached! "unknown array entry " + entry.constructor.name
      end
    end
    null
  end

  def check_function_def({ name, args, body, type })
    let sub_types = Object.assign({}, .types)
    let tc = new TypeChecker(body, sub_types, .self)
    # infer the returned type
    # TODO: look at all return statements, not just the last one
    type = tc.infer(body.at(-1)) if !type
    # store the type
    .types[name] = new FnT(args, type)
    # store the arg types
    for let arg of args do
      sub_types[arg.name] = arg.type || new AnyT()
    end
    # type check the function body
    tc.check()
    return tc.return_types
  end

  def infer_id_lookup({ name })
    return new NullT() if name === "null"
    return new BoolT() if name === "true" || name === "false"
    return .types[name] if .types[name]
    assert_not_reached! "unknown type for " ++ name
  end

  def infer_str_property(method_name)
    if method_name === "slice"
      return new FnT([{ type: new NumT() }], new StrT())
    else if method_name === "match"
      let ret_t = new CombinedT(new ArrayT(new StrT()), {
        index: new NumT(),
        input: new StrT(),
        groups: new NullT()
      }) 
      return new FnT([{ type: new RegexT() }], ret_t)
    else if method_name === "length"
      return new NumT()
    else
      assert_not_reached! "unknown str method " + method_name
    end
  end

  def infer_dot_access({ lhs, property }, lhs_t = null)
    lhs_t = lhs_t || .infer(lhs)
    return lhs_t if lhs_t is AnyT
    if lhs_t is ArrayT
      return .infer_array_method(lhs_t, property)
    else if lhs_t is ObjT
      return lhs_t.properties[property]
    else if lhs_t is NumT
      return .infer_number_method(property)
    else if lhs_t is StrT
      return .infer_str_property(property)
    else if lhs_t is DataClassT
      let p = lhs_t.properties.find(p => p[0] === property)
      assert! p !== undefined
      return p[1]
    else if lhs_t is CombinedT
      let { root_t, properties } = lhs_t
      let t = properties[property]
      if t
        return t
      else
        return .infer_dot_access({ lhs: null, property: property }, root_t)
      end
    else
      console.log(lhs, property)
      assert_not_reached! "unknown lhs of dot access " ++ lhs_t.constructor.name
    end
  end

  def infer_new_expr({ expr })
    assert! expr is FunctionCall
    let { lhs_expr } = expr
    let class_t = .infer(lhs_expr)
    assert! class_t is DataClassT
    return class_t
  end

  def infer_property_lookup({ lhs, property }, lhs_t = null)
    lhs_t = lhs_t || .infer(lhs)
    let property_t = .infer(property)
    if lhs_t is StrT
      return new StrT() if property_t is NumT
      assert_not_reached! "unknown property `" ++ property ++ "` on str"
    else if lhs_t is ArrayT
      return lhs_t.type if property_t is NumT
      assert_not_reached! "unknown property `" ++ property ++ "` on array"
    else if lhs_t is CombinedT
      let { root_t, properties } = lhs_t
      let t = properties[property]
      if t
        return t
      else
        return .infer_property_lookup({ lhs: null, property: property }, root_t)
      end
    else
      console.log(lhs_t)
      assert_not_reached! "property lookup"
    end
  end

  def infer(expr)
    if expr is NumExpr
      return new NumT()
    else if expr is StrExpr
      return new StrT()
    else if expr is ReturnExpr
      let type = .infer(expr.expr)
      .return_types.push(type)
      return type
    else if expr is JsOpExpr
      return .infer_js_op(expr.type, expr.lhs, expr.rhs).return_type
    else if expr is IdLookup
      return .infer_id_lookup(expr)
    else if expr is FunctionCall
      return .infer(expr.lhs_expr).return_type
    else if expr is DotAccess
      return .infer_dot_access(expr)
    else if expr is ArrayLiteral
      return .infer_array_literal(expr)
    else if expr is NewExpr
      return .infer_new_expr(expr)
    else if expr is PrefixDotLookup
      return .infer_prefix_dot_lookup(expr)
    else if expr is PropertyLookup
      return .infer_property_lookup(expr)
    else if expr is NotExpr
      return new BoolT()
    else
      console.log(expr)
      assert_not_reached! "Cant infer " + expr.constructor.name
    end
  end

  def infer_number_method(property)
    if property === "toString"
      return new FnT([], new StrT())
    else
      assert_not_reached! "unknown property `" + property + "` on num"
    end
  end

  def infer_array_method(arr_t, property)
    if property === "push"
      return new FnT([{ type: arr_t.type }], new NumT())
    else
      assert_not_reached! "unknown array method " + property
    end
  end

  def infer_array_literal({ elements })
    let types = elements.map(::infer).uniq_by(::is_match)
    if types.length === 1
      return new ArrayT(types[0])
    else
      return new ArrayT(new UnionT(types))
    end
  end

  def panic_mismatch(name, expected, got)
    assert_not_reached! "type mismatch: expected `" + name + "` to be a " + expected + " but it was a " + got 
  end

  def resolve(type)
    return type if type is not TypeIdLookup
    .types[type.name]
  end

  def is_match(a, b)
    a = .resolve(a)
    b = .resolve(b)
    if a is UnionT
      assert_not_reached! "fixme: union types can't be unwrapped like this"
      return a.types.some(type => .is_match(type, b))
    else if b is UnionT
      return .is_match(b, a)
    end
    a.constructor === b.constructor
  end

  def infer_js_op(type, lhs, rhs)
    if type === "+"
      return new FnT([{ type: new NumT() }, { type: new NumT() }], new NumT())
    else if type === "++"
      return new FnT([{ type: new StrT() }, { type: new StrT() }], new StrT())
    else if type === "||"
      let lhs_t = .infer(lhs)
      let rhs_t = .infer(rhs)
      return new FnT([{ type: lhs_t }, { type: rhs_t }], new UnionT([lhs_t, rhs_t]))
    else if type === "!=="
      let lhs_t = .infer(lhs)
      let rhs_t = .infer(rhs)
      return new FnT([{ type: lhs_t }, { type: rhs_t }], new BoolT())
    else
      assert_not_reached! "unknown js op type " + type
    end
  end

  def check_js_op_expr({ lhs, rhs, type })
      let lhs_t = .infer(lhs)
      let rhs_t = .infer(rhs)
      if !.is_match(lhs_t, rhs_t)
        console.log(lhs_t, rhs_t)
        assert_not_reached! "js operands don't match"
      end
      let { args, return_type } = .infer_js_op(type, lhs, rhs)
      if !.is_match(lhs_t, args[0].type)
        console.log(type, return_type)
        console.log(lhs_t, args[0].type)
        assert_not_reached! "js op arg type mismatch"
      end
      assert_not_reached! "js op return type mismatch" if !.is_match(lhs_t, return_type)
  end

  def check_expr(expr)
    if expr is FunctionCall
      .check_function_call(expr)
    else if expr is NumExpr
      null
    else if expr is StrExpr
      null
    else if expr is JsOpExpr
      .check_js_op_expr(expr)
    else if expr is ArrayLiteral
      .check_array_literal(expr)
    else if expr is IdLookup
      assert_not_reached! "cant find `" + expr.name + "`" if !.types[expr.name]
    else if expr is NewExpr
      .check_new_expr(expr)
    else if expr is DotAccess
      .check_dot_access(expr)
    else if expr is NotExpr
      .check_not_expr(expr)
    else
      assert_not_reached! "check_expr: Unknown expr " + expr.constructor.name
    end
  end

  def check_not_expr({ expr })
    # console.log(expr, .infer(expr))
    # assert! .infer(expr) is BoolT
    # null

    # anything can be not'd
  end

  def check_dot_access({ lhs, property })
    let lhs_t = .infer(lhs)
    if lhs_t is DataClassT
      assert! lhs_t.properties.some(p => p[0] === property)
    else
      assert_not_reached! "Not able to dot access on " + lhs_t.constructor.name
    end
  end

  def check_new_expr({ expr })
    assert! expr is FunctionCall
    let { lhs_expr, args } = expr
    let class_t = .infer(lhs_expr)
    assert! class_t is DataClassT
    let { properties } = class_t
    for let iter of args.zip(properties) do
      let [arg, arg_t] = iter
      let [, type] = arg_t
      assert! .is_match(.infer(arg), type)
    end

    null
  end

  def check_array_literal({ elements })
    elements.map(::check_expr)
  end

  def pretty(obj)
    obj.constructor.name + "(" + JSON.stringify(obj) + ")"
  end

  def check_function_call({ lhs_expr, args })
    for let arg of args do
      .check_expr(arg)
    end
    assert! .infer(lhs_expr) is FnT
    let { args: fn_arg_types } = .infer(lhs_expr)
    return null if fn_arg_types is AnyT
    assert_not_reached! "args length mismatch" if fn_arg_types.length !== args.length

    for let iter of args.zip(fn_arg_types) do
      let [value, {type}] = iter
      if !.is_match(.infer(value), type)
        .panic_mismatch(.pretty(value), type.constructor.name, .infer(value).constructor.name)
      end
    end
    null
  end

  def check_named_let({ name, expr, type }, is_exported)
    if .types[name]
      assert_not_reached! "`" ++ name ++ "` already declared"
    end
    type = .resolve(type)
    .check_expr(expr)
    let inferred_type = .infer(expr)
    if !type
      .types[name] = inferred_type
    else if .is_match(type, .infer(expr))
      .types[name] = type
    else
      console.log(name, type, .infer(expr))
      console.log(expr)
      .panic_mismatch(name, type.constructor.name, .infer(expr).constructor.name)
    end
    .types[name].exported = is_exported
    null
  end
end

export default TypeChecker