lilunar/src/typecheck/program.mbt

///|
pub(all) suberror TypeCheckError String derive(Show)

///|
pub(all) struct Env {
  local_ : Map[String, Type]
  parent : Env?
}

///|
pub fn Env::new(parent? : Env? = None) -> Env {
  { local_: Map::new(), parent }
}

///|
pub fn Env::get(self : Env, name : String) -> Type? {
  match self.local_.get(name) {
    Some(t) => Some(t)
    None =>
      match self.parent {
        Some(p) => p.get(name)
        None => None
      }
  }
}

///|
pub fn Env::set(self : Env, name : String, t : Type) -> Unit {
  self.local_.set(name, t)
}

///|
pub(all) struct Context {
  mut type_env : Env
  type_vars : Map[Int, TypeKind]
  struct_defs : Map[String, StructDef]
  func_types : Map[String, TypeKind]
  mut current_func_ret_ty : TypeKind?
}

///|
pub fn Context::new() -> Context {
  {
    type_env: Env::new(),
    type_vars: Map::new(),
    struct_defs: Map::new(),
    func_types: Map::new(),
    current_func_ret_ty: None,
  }
}

///|
pub fn Context::new_type_var(self : Context) -> TypeKind {
  let type_var_id = self.type_vars.length()
  self.type_vars.set(type_var_id, TypeVar(type_var_id))
  TypeVar(type_var_id)
}

///|
pub fn Context::lookup_type(self : Context, name : String) -> Type? {
  let looked = self.type_env.get(name)
  if looked is None {
    return None
  }
  let { kind, mutable } = looked.unwrap()
  let deref_kind = self.deref_type_var(kind)
  Some({ kind: deref_kind, mutable })
}

///|
pub fn Context::enter_scope(self : Self) -> Unit {
  let sub_env = Env::new(parent=Some(self.type_env))
  self.type_env = sub_env
}

///|
pub fn Context::exit_scope(self : Context) -> Unit {
  self.type_env = match self.type_env.parent {
    Some(p) => p
    None => self.type_env
  }
}

///|
pub fn Context::set_current_func_ret_ty(self : Context, ty : TypeKind) -> Unit {
  self.current_func_ret_ty = Some(ty)
}

///|
pub(all) struct Program {
  top_lets : Map[String, TopLet]
  top_functions : Map[String, TopFunction]
  struct_defs : Map[String, StructDef]
} derive(Show)

///|
pub fn Context::check_program(
  self : Context,
  program : @parser.Program,
) -> Program raise TypeCheckError {
  // 收集结构体名称
  for name, _ in program.struct_defs {
    if self.struct_defs.contains(name) {
      raise TypeCheckError("Duplicate struct definition: \{name}")
    }
    self.struct_defs.set(name, { name, fields: [] })
  }
  // 此时除泛型外，所有存在的类型都已确定
  // 收集结构体定义
  let struct_defs : Map[String, StructDef] = Map::new()
  for name, struct_def in program.struct_defs {
    let checked_struct_def = self.check_struct_def(struct_def)
    struct_defs.set(name, checked_struct_def)
    self.struct_defs.set(name, checked_struct_def)
  }
  // 收集顶层函数声明
  for name, func in program.top_functions {
    self.func_types.set(name, self.check_top_function_type_decl(func))
  }
  // 补充内建函数类型
  self.add_intrinsic_functions()
  // 此时顶层变量初始化表达式可引用的类型都已确定
  // 收集顶层变量类型和表达式
  let top_lets : Map[String, TopLet] = Map::new()
  for name, let_def in program.top_lets {
    let checked_let = self.check_top_let(let_def)
    top_lets.set(name, checked_let)
    self.type_env.set(name, checked_let.ty)
  }
  // 此时所有顶层标识符的类型都已确定
  // 检查顶层函数定义
  let top_functions : Map[String, TopFunction] = Map::new()
  for name, func in program.top_functions {
    top_functions.set(name, self.check_top_function_body(func))
  }
  { top_lets, top_functions, struct_defs }
}

///|
pub fn Context::add_intrinsic_functions(self : Context) -> Unit {
  let map = Map::of([
    ("read_int", Function([], Int)),
    ("print_int", Function([Int], Unit)),
    ("read_char", Function([], Int)),
    ("print_char", Function([Int], Unit)),
    ("print_endline", Function([], Unit)),
    ("int_of_float", Function([Double], Int)),
    ("float_of_int", Function([Int], Double)),
    ("truncate", Function([Double], Int)),
    ("floor", Function([Double], Double)),
    ("abs_float", Function([Double], Double)),
    ("sqrt", Function([Double], Double)),
    ("sin", Function([Double], Double)),
    ("cos", Function([Double], Double)),
    ("atan", Function([Double], Double)),
  ])
  for name, ty in map {
    self.type_env.set(name, { kind: ty, mutable: false })
  }
}