(** The checker: AST → typed IR.

    Two passes, because top-level names in a package are order-independent
    (plan.org, Modules): the first collects every type, signature and global,
    the second checks bodies against them. Mutually recursive functions need no
    forward declaration, and a struct may be used above where it is declared.

    Checking is *bidirectional*. An expression is checked against an expected
    type when there is one and inferred when there is not, which is what makes
    [None], a bare [0] and a struct literal work without any inference engine:
    the expected type flows in from the function's return type, the parameter
    it is being passed to, or the field it is being stored in.

    The rule from the two misparse bugs applies here too: *anything not yet
    implemented is rejected by name*, never approximated. Milestone 2 is
    calc-me.flan and nothing more (plan.org, Build sequence), so [Vec], [Map],
    [Result]/[try], user unions, closures, [dotimes], [defer], generics and
    cross-package imports are all errors with a message that says which
    milestone they belong to. *)

let fail = Loc.fail

(* [List.map]'s evaluation order is unspecified, and checking allocates frame
   slots as a side effect. Left-to-right is required, not a preference: a later
   let binding sees an earlier one, and slot numbering must be reproducible. *)
let rec map_lr f = function
  | [] -> []
  | x :: rest -> let y = f x in y :: map_lr f rest

let rec map2_lr f xs ys =
  match xs, ys with
  | [], [] -> []
  | x :: xs, y :: ys -> let z = f x y in z :: map2_lr f xs ys
  | _ -> invalid_arg "map2_lr"

(* ── Environments ──────────────────────────────────────────────────── *)

type binding = {
  slot       : int;
  bty        : Types.t;
  assignable : bool;   (* locals are places; parameters are not — spec-memory *)
}

type env = {
  structs : (string, Tast.structure) Hashtbl.t;
  unions  : (string, Tast.union) Hashtbl.t;
  aliases : (string, Ast.texpr) Hashtbl.t;
  consts  : (string, int64) Hashtbl.t;         (* compile-time array lengths *)
  locs    : (string, Loc.t) Hashtbl.t;         (* where each type was declared *)
  fns     : (string, Types.t list * Types.t) Hashtbl.t;
  globals : (string, Types.t * bool) Hashtbl.t;   (* type, is a constant *)
}

let new_env () = {
  structs = Hashtbl.create 16;
  unions  = Hashtbl.create 16;
  aliases = Hashtbl.create 16;
  consts  = Hashtbl.create 16;
  locs    = Hashtbl.create 16;
  fns     = Hashtbl.create 32;
  globals = Hashtbl.create 16;
}

(* Per-function state. Slots are never reused, so [slots] is also the frame
   size — the interpreter allocates one array of this length per call. *)
type ctx = {
  env           : env;
  ret           : Types.t;
  mutable slots : int;
  (* The type of each slot, newest first. A backend needs it to size the
     frame — nothing else records it, since the IR refers to slots by index. *)
  mutable slot_tys : Types.t list;
  mutable scope : (string * binding) list;   (* innermost first *)
}

let fresh_slot ctx ty =
  let s = ctx.slots in
  ctx.slots <- s + 1;
  ctx.slot_tys <- ty :: ctx.slot_tys;
  s

let bind ctx name bty ~assignable =
  let slot = fresh_slot ctx bty in
  ctx.scope <- (name, { slot; bty; assignable }) :: ctx.scope;
  slot

let lookup ctx name = List.assoc_opt name ctx.scope

let scoped ctx f =
  let saved = ctx.scope in
  let r = f () in
  ctx.scope <- saved;
  r

(* ── Type resolution ───────────────────────────────────────────────── *)

let unimplemented loc what milestone =
  fail loc "%s is not implemented yet — milestone %d (see plan.org)"
    what milestone

let rec resolve env ?(seen = []) (t : Ast.texpr) : Types.t =
  let loc = t.Ast.tloc in
  match t.Ast.t with
  | Ast.Tname n -> resolve_name env ~seen loc n
  | Ast.Tslice e -> Types.Slice (resolve env ~seen e)
  | Ast.Tarray (l, e) -> Types.Array (array_len env loc l, resolve env ~seen e)
  | Ast.Tmap _ -> unimplemented loc "the Map type" 6
  | Ast.Tfn (ps, r) ->
    Types.Fn (List.map (resolve env ~seen) ps, resolve env ~seen r)
  | Ast.Tapp (name, args) ->
    (match name, args with
     | "Ptr", [ a ] -> Types.Ptr (resolve env ~seen a)
     | "Option", [ a ] -> Types.Option (resolve env ~seen a)
     | ("Ptr" | "Option"), _ -> fail loc "(%s T) takes exactly one type" name
     | "Vec", _ -> unimplemented loc "(Vec T)" 6
     | "Map", _ -> unimplemented loc "(Map K V)" 6
     | "Result", _ -> unimplemented loc "(Result T E)" 6
     | "Handle", _ -> unimplemented loc "(Handle T)" 6
     | _ ->
       fail loc
         "%s takes no type arguments — generics are milestone 5" name)

(* One edit away from a type that exists — a substitution, an insertion, a
   deletion or a transposition of neighbours. Bounded at one, because two edits
   is no longer a typo, it is a guess. *)
and near_miss env n =
  let one_edit a b =
    let la = String.length a and lb = String.length b in
    if abs (la - lb) > 1 then false
    else begin
      (* Walk both until they diverge, then require the tails to match with the
         single edit applied. *)
      let i = ref 0 in
      while !i < la && !i < lb && a.[!i] = b.[!i] do incr i done;
      let ta s k = String.sub s k (String.length s - k) in
      if la = lb then
        !i < la
        && (ta a (!i + 1) = ta b (!i + 1)
            (* stirng/string: two neighbours swapped. *)
            || (!i + 1 < la && a.[!i] = b.[!i + 1] && a.[!i + 1] = b.[!i]
                && ta a (!i + 2) = ta b (!i + 2)))
      else if la < lb then ta a !i = ta b (!i + 1)
      else ta a (!i + 1) = ta b !i
    end
  in
  let candidates =
    Types.primitive_names
    @ Hashtbl.fold (fun k _ acc -> k :: acc) env.aliases []
    @ Hashtbl.fold (fun k _ acc -> k :: acc) env.structs []
    @ Hashtbl.fold (fun k _ acc -> k :: acc) env.unions []
  in
  List.find_opt (fun c -> c <> n && one_edit n c) candidates

and resolve_name env ~seen loc n =
  match Types.ikind_of_name n with
  | Some k -> Types.Int k
  | None ->
    match Types.fkind_of_name n with
    | Some k -> Types.Float k
    | None ->
      match n with
      | "bool" -> Types.Bool
      | "string" -> Types.String
      | "Unit" -> Types.Unit
      | "Never" -> Types.Never
      | _ when Hashtbl.mem env.aliases n ->
        if List.mem n seen then
          fail loc "the type alias %s is defined in terms of itself" n
        else resolve env ~seen:(n :: seen) (Hashtbl.find env.aliases n)
      | _ when Hashtbl.mem env.structs n || Hashtbl.mem env.unions n ->
        Types.Named n
      (* A typo in a primitive is lowercase too, and the type-variable rule
         below would otherwise report [f65] as unimplemented generics and send
         you to plan.org instead of to the character you mistyped. *)
      | _ when near_miss env n <> None ->
        fail loc "unknown type %s — did you mean %s?" n
          (Option.get (near_miss env n))
      (* Lowercase is a type variable, Capitalized is concrete — no sigil
         (plan.org, Types). A variable parses, but nothing at milestone 2 can
         give a value one, so it is rejected here rather than later. *)
      | _ when n <> "" && n.[0] = Char.lowercase_ascii n.[0] ->
        unimplemented loc
          (Printf.sprintf "generic code over the type variable %s" n) 5
      | _ -> fail loc "unknown type %s" n

and array_len env loc = function
  | Ast.Lint n -> n
  | Ast.Lname n ->
    (match Hashtbl.find_opt env.consts n with
     | Some v -> v
     | None ->
       fail loc "%s is not a compile-time integer constant, so it cannot be \
                 an array length" n)

(* ── Small helpers over the AST ────────────────────────────────────── *)

(* Untyped literals: their machine type comes from context, so when one is an
   operand of a binary operator we look at the *other* operand first. *)
let is_literal (e : Ast.expr) =
  match e.Ast.e with Ast.Int _ | Ast.Float _ | Ast.Byte _ -> true | _ -> false

(* [addr] takes the address of a place, but the parser only builds places for
   [set]. Recover one from the expression it parsed instead. *)
let place_of_expr (e : Ast.expr) : Ast.place option =
  match e.Ast.e with
  | Ast.Var s -> Some (Ast.Pvar s)
  | Ast.Field (t, f) -> Some (Ast.Pfield (t, f))
  | Ast.Call ({ Ast.e = Ast.Var "at"; _ }, t :: idx) when idx <> [] ->
    Some (Ast.Pindex (t, idx))
  | Ast.Call ({ Ast.e = Ast.Var "get"; _ }, [ m; k ]) -> Some (Ast.Pkey (m, k))
  | Ast.Call ({ Ast.e = Ast.Var "deref"; _ }, [ p ]) -> Some (Ast.Pderef p)
  | _ -> None

let mk loc ty e : Tast.expr = { Tast.e; ty; loc }

let unit_at loc = mk loc Types.Unit Tast.Unit

(* Every integer index into an array or slice is i32 at milestone 2. *)
let index_ty = Types.Int Types.I32

let expect loc ~want (got : Tast.expr) =
  match want with
  | None -> got
  | Some w ->
    if Types.fits ~expected:w ~actual:got.Tast.ty then got
    else
      fail loc "expected %s, found %s" (Types.to_string w)
        (Types.to_string got.Tast.ty)

(* ── Expressions ───────────────────────────────────────────────────── *)

let rec check ctx ?want (e : Ast.expr) : Tast.expr =
  let loc = e.Ast.loc in
  match e.Ast.e with
  | Ast.Int n -> int_literal loc ~want n
  | Ast.Byte b -> int_literal loc ~want ~default:Types.U8 (Int64.of_int b)
  | Ast.Float x ->
    let k =
      match want with
      | Some (Types.Float k) -> k
      | Some other when other <> Types.Never ->
        fail loc "expected %s, found the float literal %g"
          (Types.to_string other) x
      | _ -> Types.F64
    in
    mk loc (Types.Float k) (Tast.Float (x, k))
  | Ast.Str s -> expect loc ~want (mk loc Types.String (Tast.Str s))
  | Ast.Kw _ ->
    unimplemented loc "a keyword at a call site (keyword->enum coercion)" 4
  | Ast.Quote _ ->
    unimplemented loc "a quoted symbol (restart names)" 6
  | Ast.Var name -> var ctx loc ~want name
  | Ast.Do body -> block ctx ?want loc body
  | Ast.Let (bs, body) -> check_let ctx ?want loc bs body
  | Ast.If (c, t, e') -> check_if ctx ?want loc c t e'
  | Ast.While (c, body) ->
    let c = check ctx ~want:Types.Bool c in
    let body = scoped ctx (fun () -> map_lr (fun b -> check ctx b) body) in
    expect loc ~want (mk loc Types.Unit (Tast.While (c, body)))
  | Ast.Return v ->
    let v =
      match v with
      | None ->
        if not (Types.equal ctx.ret Types.Unit) then
          fail loc "this function returns %s, so return needs a value"
            (Types.to_string ctx.ret);
        None
      | Some v -> Some (check ctx ~want:ctx.ret v)
    in
    mk loc Types.Never (Tast.Return v)
  | Ast.Set (p, v) ->
    let p, pty = check_place ctx loc p in
    let v = check ctx ~want:pty v in
    expect loc ~want (mk loc Types.Unit (Tast.Set (p, v)))
  | Ast.Field (target, name) ->
    let target, sname = struct_target ctx target in
    let s = Hashtbl.find ctx.env.structs sname in
    (match Tast.field_index s name with
     | None -> fail loc "%s has no field %s" sname name
     | Some i ->
       let fty = (List.nth s.Tast.fields i).Tast.fty in
       expect loc ~want (mk loc fty (Tast.Field (target, i))))
  | Ast.Struct (name, kvs) -> check_struct ctx ~want loc name kvs
  | Ast.Arr items -> check_arr ctx ~want loc items
  | Ast.Match (scrutinee, arms) -> check_match ctx ?want loc scrutinee arms
  | Ast.Call (head, args) -> check_call ctx ~want loc head args
  | Ast.Unwrap (Ast.Usome, v) ->
    (* Unwrap Some, else early-return None from the enclosing function, so the
       enclosing function must itself return an Option (plan.org). *)
    (match ctx.ret with
     | Types.Option _ ->
       let v = check ctx v in
       (match v.Tast.ty with
        | Types.Option t ->
          expect loc ~want (mk loc t (Tast.UnwrapSome v))
        | other ->
          fail loc "some takes an (Option T), found %s" (Types.to_string other))
     | other ->
       fail loc
         "some early-returns None, so the enclosing function must return an \
          Option; this one returns %s" (Types.to_string other))
  | Ast.Unwrap (Ast.Utry, _) -> unimplemented loc "try (Result)" 6
  | Ast.Fn _ -> unimplemented loc "fn values" 5
  | Ast.Dotimes _ -> unimplemented loc "dotimes" 4
  | Ast.Defer _ -> unimplemented loc "defer" 4

and int_literal loc ~want ?(default = Types.I32) n =
  match want with
  | Some (Types.Int k) -> mk loc (Types.Int k) (Tast.Int (in_range loc k n, k))
  (* An untyped integer constant is usable where a float is wanted, as in
     Odin. A float literal is never usable where an integer is wanted. *)
  | Some (Types.Float k) ->
    mk loc (Types.Float k) (Tast.Float (Int64.to_float n, k))
  | Some other when other <> Types.Never ->
    fail loc "expected %s, found the integer literal %Ld"
      (Types.to_string other) n
  | _ -> mk loc (Types.Int default) (Tast.Int (in_range loc default n, default))

(* Arithmetic wraps, but a literal that does not fit its type is a typo, not a
   wrap — 300 is never what someone meant by a u8. *)
and in_range loc k n =
  let bits = Types.bits k in
  let ok =
    if Types.signed k then
      bits = 64
      || (Int64.compare n (Int64.neg (Int64.shift_left 1L (bits - 1))) >= 0
          && Int64.compare n (Int64.shift_left 1L (bits - 1)) < 0)
    else
      Int64.compare n 0L >= 0
      && (bits = 64 || Int64.compare n (Int64.shift_left 1L bits) < 0)
  in
  if ok then n
  else fail loc "%Ld does not fit in %s" n (Types.ikind_name k)

and var ctx loc ~want name =
  match name with
  | "true" | "false" ->
    expect loc ~want (mk loc Types.Bool (Tast.Bool (name = "true")))
  | "None" ->
    (match want with
     | Some (Types.Option t) -> mk loc (Types.Option t) Tast.None_
     | Some other when other <> Types.Never ->
       fail loc "expected %s, found None" (Types.to_string other)
     | _ ->
       fail loc
         "nothing here says what None is an Option of — annotate the \
          function's return type or the binding")
  | _ ->
    match lookup ctx name with
    | Some b -> expect loc ~want (mk loc b.bty (Tast.Local b.slot))
    | None ->
      match Hashtbl.find_opt ctx.env.globals name with
      | Some (ty, _) -> expect loc ~want (mk loc ty (Tast.Global name))
      | None ->
        if Hashtbl.mem ctx.env.fns name then
          unimplemented loc
            (Printf.sprintf "the function value %s (a name used as a value)" name) 5
        else fail loc "unknown name %s" name

and block ctx ?want loc body =
  match body with
  | [] -> expect loc ~want (unit_at loc)
  | _ ->
    let rec go = function
      | [ last ] -> let l = check ctx ?want last in [ l ], l.Tast.ty
      | x :: rest -> let x = check ctx x in
        let rest, ty = go rest in x :: rest, ty
      | [] -> assert false
    in
    let body, ty = go body in
    mk loc ty (Tast.Do body)

and check_let ctx ?want loc bs body =
  scoped ctx (fun () ->
    let bs =
      map_lr
        (fun (b : Ast.binding) ->
           let want = Option.map (resolve ctx.env) b.Ast.bty in
           let v = check ctx ?want b.Ast.bval in
           (match v.Tast.ty with
            | Types.Unit | Types.Never ->
              fail b.Ast.bloc "%s would be bound to %s, which is not a value"
                b.Ast.bname (Types.to_string v.Tast.ty)
            | _ -> ());
           (* Locals are assignable places; parameters are not. *)
           let slot = bind ctx b.Ast.bname v.Tast.ty ~assignable:true in
           (slot, v))
        bs
    in
    let body = block ctx ?want loc body in
    mk loc body.Tast.ty (Tast.Let (bs, [ body ])))

and check_if ctx ?want loc c t e =
  let c = check ctx ~want:Types.Bool c in
  match e with
  | None ->
    (* A one-armed if produces Unit whatever the branch evaluates to: there is
       no value on the missing side. `when` desugars to this. *)
    let t = scoped ctx (fun () -> check ctx t) in
    expect loc ~want (mk loc Types.Unit (Tast.If (c, t, unit_at loc)))
  | Some e ->
    let t = scoped ctx (fun () -> check ctx ?want t) in
    (* With no expectation the then-branch supplies one for the else-branch,
       unless it diverges, in which case the else-branch decides. *)
    let ewant =
      match want with
      | Some _ -> want
      | None -> if t.Tast.ty = Types.Never then None else Some t.Tast.ty
    in
    let e = scoped ctx (fun () -> check ctx ?want:ewant e) in
    let ty =
      if t.Tast.ty = Types.Never then e.Tast.ty
      else if e.Tast.ty = Types.Never then t.Tast.ty
      else if Types.equal t.Tast.ty e.Tast.ty then t.Tast.ty
      else
        fail loc "the branches of this if have different types: %s and %s"
          (Types.to_string t.Tast.ty) (Types.to_string e.Tast.ty)
    in
    mk loc ty (Tast.If (c, t, e))

and check_struct ctx ~want loc name kvs =
  match Hashtbl.find_opt ctx.env.structs name with
  | None ->
    if Hashtbl.mem ctx.env.unions name then
      unimplemented loc "constructing a union value" 6
    else fail loc "unknown struct %s" name
  | Some s ->
    let seen = Hashtbl.create 8 in
    List.iter
      (fun (k, (v : Ast.expr)) ->
         if Hashtbl.mem seen k then
           fail v.Ast.loc "field %s is given twice" k;
         if Tast.field_index s k = None then
           fail v.Ast.loc "%s has no field %s" name k;
         Hashtbl.add seen k v)
      kvs;
    (* Omitted fields are zeroed — ZII, the same rule as a declaration with no
       initialiser (plan.org, Data model). Every field is present from here on,
       in declaration order, so no backend has to know about omission. *)
    let fields =
      map_lr
        (fun (f : Tast.field) ->
           match Hashtbl.find_opt seen f.Tast.fname with
           | Some v -> check ctx ~want:f.Tast.fty v
           | None -> mk loc f.Tast.fty (Tast.Zero f.Tast.fty))
        s.Tast.fields
    in
    expect loc ~want (mk loc (Types.Named name) (Tast.Make (name, fields)))

and check_arr ctx ~want loc items =
  let elem_want =
    match want with
    | Some (Types.Array (_, t)) -> Some t
    | Some (Types.Slice t) -> Some t
    | _ -> None
  in
  let items = map_lr (fun i -> check ctx ?want:elem_want i) items in
  let n = Int64.of_int (List.length items) in
  let elem =
    match elem_want, items with
    | Some t, _ -> t
    | None, first :: _ -> first.Tast.ty
    | None, [] ->
      fail loc "an empty array literal needs a type — annotate the binding"
  in
  List.iter
    (fun (i : Tast.expr) ->
       if not (Types.fits ~expected:elem ~actual:i.Tast.ty) then
         fail i.Tast.loc "this array's elements are %s, but this one is %s"
           (Types.to_string elem) (Types.to_string i.Tast.ty))
    items;
  (match want with
   | Some (Types.Array (m, _)) when not (Int64.equal m n) ->
     fail loc "expected %Ld elements, found %Ld" m n
   | _ -> ());
  (* [n T] and [T] are distinct in type and in ownership (spec-memory.md), so
     an array literal does not satisfy a slice expectation. *)
  expect loc ~want (mk loc (Types.Array (n, elem)) (Tast.Arr items))

and check_match ctx ?want loc scrutinee arms =
  let s = check ctx scrutinee in
  let elem =
    match s.Tast.ty with
    | Types.Option t -> t
    | other ->
      (* Union matching arrives with unions themselves, at milestone 6. *)
      fail loc "match works on an Option at milestone 2, not on %s"
        (Types.to_string other)
  in
  let want = ref want in
  let saw_some = ref false and saw_none = ref false and saw_wild = ref false in
  let arms =
    map_lr
      (fun (a : Ast.arm) ->
         let ctor, binds =
           match a.Ast.pat with
           | Ast.Pwild -> saw_wild := true; None, []
           | Ast.Pctor ("Some", [ x ]) -> saw_some := true; Some "Some", [ x ]
           | Ast.Pctor ("Some", _) ->
             fail a.Ast.aloc "the Some pattern binds exactly one name"
           | Ast.Pctor ("None", []) -> saw_none := true; Some "None", []
           | Ast.Pctor ("None", _) -> fail a.Ast.aloc "None binds no names"
           | Ast.Pctor (c, _) ->
             fail a.Ast.aloc
               "%s is not a case of Option — the cases are Some and None" c
         in
         scoped ctx (fun () ->
           let binds = List.map (fun n -> bind ctx n elem ~assignable:false) binds in
           let body = block ctx ?want:!want a.Ast.aloc a.Ast.body in
           if !want = None && body.Tast.ty <> Types.Never then
             want := Some body.Tast.ty;
           { Tast.acase = ctor; binds; abody = [ body ] }))
      arms
  in
  if not (!saw_wild || (!saw_some && !saw_none)) then
    fail loc
      "this match is not exhaustive — Option needs both Some and None, or a \
       _ arm";
  let ty = match !want with Some t -> t | None -> Types.Never in
  mk loc ty (Tast.Match (s, arms))

(* ── Places ────────────────────────────────────────────────────────── *)

(* The target of [.field] is a struct, or one level of pointer to one. The
   auto-deref is inserted here as a real node, so no backend re-derives it. *)
and struct_target ctx (target : Ast.expr) : Tast.expr * string =
  let t = check ctx target in
  match t.Tast.ty with
  | Types.Named n when Hashtbl.mem ctx.env.structs n -> t, n
  | Types.Ptr (Types.Named n) when Hashtbl.mem ctx.env.structs n ->
    mk t.Tast.loc (Types.Named n) (Tast.Deref t), n
  | other ->
    fail target.Ast.loc "%s is not a struct, so it has no fields"
      (Types.to_string other)

and check_place ctx loc (p : Ast.place) : Tast.place * Types.t =
  match p with
  | Ast.Pvar name ->
    (match lookup ctx name with
     | Some b ->
       if not b.assignable then
         fail loc
           "%s is a parameter, and parameters are not assignable places \
            (spec-memory.md) — bind a local with let" name;
       Tast.Plocal b.slot, b.bty
     | None ->
       match Hashtbl.find_opt ctx.env.globals name with
       | Some (_, true) -> fail loc "%s is a constant" name
       | Some (ty, false) -> Tast.Pglobal name, ty
       | None -> fail loc "unknown name %s" name)
  | Ast.Pfield (target, name) ->
    let target, sname = struct_target ctx target in
    let s = Hashtbl.find ctx.env.structs sname in
    (match Tast.field_index s name with
     | None -> fail loc "%s has no field %s" sname name
     | Some i -> Tast.Pfield (target, i), (List.nth s.Tast.fields i).Tast.fty)
  | Ast.Pindex (target, idx) ->
    let target = check ctx target in
    let idx, ty = indexed ctx target idx in
    Tast.Pindex (target, idx), ty
  | Ast.Pkey _ -> unimplemented loc "(get m k) as a place — Map" 6
  | Ast.Pderef target ->
    let target = check ctx target in
    (match target.Tast.ty with
     | Types.Ptr t -> Tast.Pderef target, t
     | other ->
       fail loc "deref takes a (Ptr T), found %s" (Types.to_string other))

(* An index or a slice bound that is a literal is known now, so it is an error
   now rather than a trap later. Only literals: a [defconst] is a global in the
   typed IR, not a folded constant, so [(at arr size)] still traps at runtime —
   which is what the emitted bounds check is for. A negative literal is wrong
   whatever the target, but a length is static only for [n T]. *)
and static_index loc (ty : Types.t) ~past_end what k =
  if k < 0L then
    fail loc "%s %Ld is negative — indices count from 0" what k;
  match ty with
  (* [past_end] is the difference between an index and a slice bound: the last
     valid index is len - 1, but a slice may end at len. *)
  | Types.Array (n, _) when if past_end then k > n else k >= n ->
    fail loc "%s %Ld is out of bounds for length %Ld" what k n
  | _ -> ()

(* The literal value of a checked expression, if it has one. *)
and literal (e : Tast.expr) =
  match e.Tast.e with Tast.Int (k, _) -> Some k | _ -> None

(* [(at a i)] and [(at grid row col)]: one index per dimension. *)
and indexed ctx (target : Tast.expr) (idx : Ast.expr list) =
  let rec go ty = function
    | [] -> [], ty
    | i :: rest ->
      let elem =
        match ty with
        | Types.Array (_, t) | Types.Slice t -> t
        | other ->
          fail i.Ast.loc "%s cannot be indexed" (Types.to_string other)
      in
      let loc = i.Ast.loc in
      let i = check ctx ~want:index_ty i in
      (match literal i with
       | Some k -> static_index loc ty ~past_end:false "index" k
       | None -> ());
      let rest, ty = go elem rest in
      i :: rest, ty
  in
  go target.Tast.ty idx

(* ── Calls ─────────────────────────────────────────────────────────── *)

and check_call ctx ~want loc (head : Ast.expr) (args : Ast.expr list) =
  match head.Ast.e with
  | Ast.Var name -> named_call ctx ~want loc name args
  | _ ->
    unimplemented loc "calling something other than a named function" 5

and arity loc name n args =
  if List.length args <> n then
    fail loc "%s takes %d argument%s, given %d" name n
      (if n = 1 then "" else "s") (List.length args)

and named_call ctx ~want loc name args =
  let prim p ty args = expect loc ~want (mk loc ty (Tast.Prim (p, args))) in
  match name with
  (* ── arithmetic and comparison ─────────────────────────────────── *)
  | "+" | "-" | "*" | "/" | "%" ->
    let p = match name with
      | "+" -> Tast.Add | "-" -> Tast.Sub | "*" -> Tast.Mul
      | "/" -> Tast.Div | _ -> Tast.Rem
    in
    arity loc name 2 args;
    let a, b = binary ctx name loc ~want:(numeric_want want) args in
    if not (Types.is_numeric a.Tast.ty) then
      fail loc "%s takes numbers, found %s" name (Types.to_string a.Tast.ty);
    prim p a.Tast.ty [ a; b ]
  | "=" | "!=" | "<" | "<=" | ">" | ">=" ->
    let p = match name with
      | "=" -> Tast.Eq | "!=" -> Tast.Ne | "<" -> Tast.Lt
      | "<=" -> Tast.Le | ">" -> Tast.Gt | _ -> Tast.Ge
    in
    arity loc name 2 args;
    let a, b = binary ctx name loc ~want:None args in
    if not (Types.is_comparable a.Tast.ty) then
      fail loc
        "%s compares machine numbers; %s has no built-in comparison \
         (plan.org, Types)" name (Types.to_string a.Tast.ty);
    prim p Types.Bool [ a; b ]
  | "not" ->
    arity loc name 1 args;
    prim Tast.Not Types.Bool [ check ctx ~want:Types.Bool (List.hd args) ]

  (* ── containers ────────────────────────────────────────────────── *)
  | "len" ->
    arity loc name 1 args;
    let a = check ctx (List.hd args) in
    (match a.Tast.ty with
     | Types.Array _ | Types.Slice _ | Types.String -> ()
     | other -> fail loc "len takes an array, a slice or a string, found %s"
                  (Types.to_string other));
    prim Tast.Len index_ty [ a ]
  | "at" | "nth" ->
    (match args with
     | target :: idx when idx <> [] ->
       let target = check ctx target in
       let idx, ty = indexed ctx target idx in
       prim Tast.At ty (target :: idx)
     | _ -> fail loc "%s is (%s collection index ...)" name name)
  | "slice" ->
    arity loc name 3 args;
    (match args with
     | [ target; lo; hi ] ->
       let target = check ctx target in
       let elem = match target.Tast.ty with
         | Types.Array (_, t) | Types.Slice t -> t
         | other -> fail loc "slice takes an array or a slice, found %s"
                      (Types.to_string other)
       in
       prim Tast.Slice (Types.Slice elem)
         (let lo_loc = lo.Ast.loc and hi_loc = hi.Ast.loc in
          let lo = check ctx ~want:index_ty lo in
          let hi = check ctx ~want:index_ty hi in
          let ty = target.Tast.ty in
          (* A bound may sit one past the end, so the length is checked against
             lo and hi both, not against the last valid index. *)
          (match literal lo with
           | Some k -> static_index lo_loc ty ~past_end:true "slice bound" k
           | None -> ());
          (match literal hi with
           | Some k -> static_index hi_loc ty ~past_end:true "slice bound" k
           | None -> ());
          (match literal lo, literal hi with
           | Some a, Some b when a > b ->
             fail loc "slice [%Ld %Ld) runs backwards — lo must not exceed hi" a b
           | _ -> ());
          [ target; lo; hi ])
     | _ -> assert false)

  (* ── pointers ──────────────────────────────────────────────────── *)
  | "addr" ->
    arity loc name 1 args;
    let a = List.hd args in
    (match place_of_expr a with
     | None ->
       fail a.Ast.loc
         "addr takes the address of a place — a name, (.field x), (at a i) \
          or (deref p)"
     | Some p ->
       let p, ty = check_place ctx a.Ast.loc p in
       expect loc ~want (mk loc (Types.Ptr ty) (Tast.Addr p)))
  | "deref" ->
    arity loc name 1 args;
    let a = check ctx (List.hd args) in
    (match a.Tast.ty with
     | Types.Ptr t -> expect loc ~want (mk loc t (Tast.Deref a))
     | other -> fail loc "deref takes a (Ptr T), found %s"
                  (Types.to_string other))

  (* ── Option ────────────────────────────────────────────────────── *)
  | "Some" ->
    arity loc name 1 args;
    let inner = match want with Some (Types.Option t) -> Some t | _ -> None in
    let a = check ctx ?want:inner (List.hd args) in
    expect loc ~want (mk loc (Types.Option a.Tast.ty) (Tast.Some_ a))

  (* ── the milestone-2 host primitives (plan.org) ────────────────── *)
  | "bytes" ->
    arity loc name 1 args;
    prim Tast.Bytes (Types.Slice (Types.Int Types.U8))
      [ check ctx ~want:Types.String (List.hd args) ]
  | "bytes->f64" ->
    arity loc name 1 args;
    prim Tast.BytesToF64 (Types.Float Types.F64) [ byte_slice ctx (List.hd args) ]
  | "bytes->i64" ->
    arity loc name 1 args;
    prim Tast.BytesToI64 (Types.Int Types.I64) [ byte_slice ctx (List.hd args) ]
  | "f64->bytes" ->
    arity loc name 1 args;
    prim Tast.F64ToBytes (Types.Slice (Types.Int Types.U8))
      [ check ctx ~want:(Types.Float Types.F64) (List.hd args) ]
  | "i64->bytes" ->
    arity loc name 1 args;
    prim Tast.I64ToBytes (Types.Slice (Types.Int Types.U8))
      [ check ctx ~want:(Types.Int Types.I64) (List.hd args) ]
  | "write-stdout" ->
    arity loc name 1 args;
    prim Tast.WriteStdout Types.Unit [ byte_slice ctx (List.hd args) ]
  | "exit" ->
    arity loc name 1 args;
    prim Tast.Exit Types.Never [ check ctx ~want:index_ty (List.hd args) ]
  | "argv" ->
    arity loc name 0 args;
    prim Tast.Argv (Types.Slice Types.String) []

  (* ── casts: (i32 x), (f64 x) ───────────────────────────────────── *)
  | _ when is_cast name && List.length args = 1 ->
    let target = resolve_name ctx.env ~seen:[] loc name in
    let a = check ctx (List.hd args) in
    if not (Types.is_numeric a.Tast.ty) then
      fail loc "%s converts a number, found %s" name
        (Types.to_string a.Tast.ty);
    prim (Tast.Cast target) target [ a ]

  (* ── ordinary calls ────────────────────────────────────────────── *)
  | _ ->
    match Hashtbl.find_opt ctx.env.fns name with
    | Some (params, ret) ->
      if List.length args <> List.length params then
        fail loc "%s takes %d argument%s, given %d" name
          (List.length params)
          (if List.length params = 1 then "" else "s")
          (List.length args);
      let args = map2_lr (fun p a -> check ctx ~want:p a) params args in
      expect loc ~want (mk loc ret (Tast.Call (name, args)))
    | None ->
      if Hashtbl.mem ctx.env.structs name || Hashtbl.mem ctx.env.unions name
      then
        fail loc
          "%s is a type — a struct value is written (%s {:field value ...})"
          name name
      else if String.contains name '/' then
        unimplemented loc
          (Printf.sprintf "the call %s into an imported package" name) 4
      else fail loc "unknown function %s" name

and is_cast name =
  Types.ikind_of_name name <> None || Types.fkind_of_name name <> None

and byte_slice ctx (a : Ast.expr) =
  check ctx ~want:(Types.Slice (Types.Int Types.U8)) a

and numeric_want want =
  match want with Some (Types.Int _ | Types.Float _) -> want | _ -> None

(* Both operands of a binary operator have one type, and there is no implicit
   widening, so one side has to decide it. Check the side that carries the most
   information first: a non-literal over a literal, and a float literal over an
   integer one, since an integer constant converts to a float and not back. *)
and binary ctx name loc ~want args =
  match args with
  | [ x; y ] ->
    let y_decides =
      (is_literal x && not (is_literal y))
      || (match x.Ast.e, y.Ast.e with
          | (Ast.Int _ | Ast.Byte _), Ast.Float _ -> true
          | _ -> false)
    in
    if y_decides then begin
      let b = check ctx ?want y in
      let a = check ctx ~want:b.Tast.ty x in
      a, b
    end else begin
      let a = check ctx ?want x in
      let b = check ctx ~want:a.Tast.ty y in
      a, b
    end
  | _ -> fail loc "%s takes two arguments" name

(* ── Declarations: pass 1, collect ─────────────────────────────────── *)

(* Constant folding, only over integers and only for defconst — enough for an
   array length like (/ screen-height cell-size). *)
let rec const_int env (e : Ast.expr) : int64 option =
  match e.Ast.e with
  | Ast.Int n -> Some n
  | Ast.Byte b -> Some (Int64.of_int b)
  | Ast.Var n -> Hashtbl.find_opt env.consts n
  | Ast.Call ({ Ast.e = Ast.Var op; _ }, [ x; y ]) ->
    (match const_int env x, const_int env y with
     | Some a, Some b ->
       (match op with
        | "+" -> Some (Int64.add a b)
        | "-" -> Some (Int64.sub a b)
        | "*" -> Some (Int64.mul a b)
        | "/" when b <> 0L -> Some (Int64.div a b)
        | "%" when b <> 0L -> Some (Int64.rem a b)
        | _ -> None)
     | _ -> None)
  | _ -> None

let collect env (decls : Ast.decl list) =
  (* Names first, so a struct may mention one declared below it. *)
  List.iter
    (fun (d : Ast.decl) ->
       match d.Ast.d with
       | Ast.Defstruct (n, _) ->
         Hashtbl.replace env.locs n d.Ast.dloc;
         Hashtbl.replace env.structs n { Tast.sname = n; fields = [] }
       | Ast.Defunion (n, _) ->
         Hashtbl.replace env.locs n d.Ast.dloc;
         Hashtbl.replace env.unions n { Tast.uname = n; cases = [] }
       | Ast.Defalias (n, t) -> Hashtbl.replace env.aliases n t
       | _ -> ())
    decls;
  (* Compile-time integer constants next, to a fixpoint, because an array
     length may name a constant declared below it — top-level names in a
     package are order-independent (plan.org, Modules). *)
  let fold_consts () =
    let progress = ref false in
    List.iter
      (fun (d : Ast.decl) ->
         match d.Ast.d with
         | Ast.Defconst (n, _, v) when not (Hashtbl.mem env.consts n) ->
           (match const_int env v with
            | Some i -> Hashtbl.replace env.consts n i; progress := true
            | None -> ())
         | _ -> ())
      decls;
    !progress
  in
  while fold_consts () do () done;
  let field (f : Ast.field) : Tast.field =
    { Tast.fname = f.Ast.fname; fty = resolve env f.Ast.fty }
  in
  (* Constants with no declared type are inferred from their value, which needs
     every other signature in hand — so they are deferred to a pass of their
     own below. *)
  let untyped = ref [] in
  List.iter
    (fun (d : Ast.decl) ->
       let loc = d.Ast.dloc in
       match d.Ast.d with
       | Ast.Package _ -> ()
       | Ast.Import (alias, _) ->
         unimplemented loc
           (Printf.sprintf "the cross-package import (import %s ...)" alias) 4
       | Ast.Defalias _ -> ()
       | Ast.Defstruct (n, fs) ->
         let names = List.map (fun (f : Ast.field) -> f.Ast.fname) fs in
         if List.length (List.sort_uniq compare names) <> List.length names then
           fail loc "%s declares the same field twice" n;
         Hashtbl.replace env.structs n
           { Tast.sname = n; fields = List.map field fs }
       | Ast.Defunion (n, vs) ->
         Hashtbl.replace env.unions n
           { Tast.uname = n;
             cases = List.map (fun (v : Ast.variant) ->
               { Tast.vname = v.Ast.vname;
                 vfields = List.map field v.Ast.vfields }) vs }
       | Ast.Defn fn ->
         if Hashtbl.mem env.fns fn.Ast.name then
           fail loc "%s is defined twice" fn.Ast.name;
         let params =
           List.map (fun (p : Ast.field) -> resolve env p.Ast.fty) fn.Ast.params
         in
         let ret =
           match fn.Ast.ret with None -> Types.Unit | Some t -> resolve env t
         in
         Hashtbl.replace env.fns fn.Ast.name (params, ret)
       | Ast.Defvar (n, t, _) ->
         let ty = match t with
           | Some t -> resolve env t
           | None -> fail loc "defvar %s needs a type" n
         in
         Hashtbl.replace env.globals n (ty, false)
       | Ast.Defconst (n, Some t, _) ->
         Hashtbl.replace env.globals n (resolve env t, true)
       | Ast.Defconst (n, None, v) -> untyped := (n, v) :: !untyped)
    decls;
  (* Also to a fixpoint, and for the same reason: one untyped constant may be
     defined in terms of another declared after it. A constant that still does
     not check once no progress is left has a real error, so the last round is
     run without swallowing it. *)
  let infer (_, v) =
    (check { env; ret = Types.Unit; slots = 0; slot_tys = []; scope = [] } v).Tast.ty
  in
  let pending = ref (List.rev !untyped) in
  let rec settle () =
    let left =
      List.filter
        (fun ((n, _) as c) ->
           match infer c with
           | ty -> Hashtbl.replace env.globals n (ty, true); false
           | exception Loc.Error _ -> true)
        !pending
    in
    let progressed = List.length left < List.length !pending in
    pending := left;
    if progressed && left <> [] then settle ()
  in
  settle ();
  List.iter (fun c -> ignore (infer c)) !pending

(* A type that contains itself by value has no finite size. [(Ptr T)] and a
   slice are indirections and break the cycle; a fixed array does not, because
   it is inline. Caught here rather than when a backend tries to lay the type
   out or a zero value is built for it — which would not fail, it would hang. *)
let check_finite env =
  let rec walk seen name =
    if List.mem name seen then
      fail (Option.value (Hashtbl.find_opt env.locs name) ~default:Loc.unknown)
        "%s contains itself by value, so it has no size — go through (Ptr %s)"
        name name;
    let seen = name :: seen in
    match Hashtbl.find_opt env.structs name with
    | Some s -> List.iter (fun (f : Tast.field) -> ty seen f.Tast.fty) s.Tast.fields
    | None ->
      match Hashtbl.find_opt env.unions name with
      | None -> ()
      | Some u ->
        List.iter
          (fun (c : Tast.variant) ->
             List.iter (fun (f : Tast.field) -> ty seen f.Tast.fty) c.Tast.vfields)
          u.Tast.cases
  and ty seen = function
    | Types.Named n -> walk seen n
    | Types.Array (_, e) | Types.Option e -> ty seen e
    | _ -> ()
  in
  Hashtbl.iter (fun n _ -> walk [] n) env.structs;
  Hashtbl.iter (fun n _ -> walk [] n) env.unions

(* ── Declarations: pass 2, check bodies ────────────────────────────── *)

let check_fn env (fn : Ast.fn) : Tast.fn =
  let params, ret = Hashtbl.find env.fns fn.Ast.name in
  let ctx = { env; ret; slots = 0; slot_tys = []; scope = [] } in
  List.iter2
    (fun (p : Ast.field) ty ->
       if List.mem_assoc p.Ast.fname ctx.scope then
         fail p.Ast.floc "%s has two parameters named %s" fn.Ast.name p.Ast.fname;
       ignore (bind ctx p.Ast.fname ty ~assignable:false))
    fn.Ast.params params;
  let body =
    match fn.Ast.fbody with
    | [] ->
      if Types.equal ret Types.Unit then []
      else fail fn.Ast.nloc "%s returns %s but has no body" fn.Ast.name
             (Types.to_string ret)
    | body ->
      (* The last form is the return value, unless the function returns Unit,
         in which case whatever it evaluates to is discarded. *)
      let want = if Types.equal ret Types.Unit then None else Some ret in
      let rec go = function
        | [ last ] -> [ check ctx ?want last ]
        | x :: rest -> check ctx x :: go rest
        | [] -> assert false
      in
      go body
  in
  { Tast.name = fn.Ast.name; params;
    slots = Array.of_list (List.rev ctx.slot_tys);
    ret; body; floc = fn.Ast.nloc }

let check_global env (d : Ast.decl) : Tast.global option =
  let ctx () = { env; ret = Types.Unit; slots = 0; slot_tys = []; scope = [] } in
  match d.Ast.d with
  | Ast.Defvar (n, _, init) ->
    let ty, _ = Hashtbl.find env.globals n in
    let ginit =
      match init with
      | Ast.Zeroed -> { Tast.e = Tast.Zero ty; ty; loc = d.Ast.dloc }
      | Ast.Uninit -> { Tast.e = Tast.Uninit ty; ty; loc = d.Ast.dloc }
      | Ast.Init v -> check (ctx ()) ~want:ty v
    in
    Some { Tast.gname = n; gty = ty; ginit; gconst = false }
  | Ast.Defconst (n, _, v) ->
    let ty, _ = Hashtbl.find env.globals n in
    Some { Tast.gname = n; gty = ty; ginit = check (ctx ()) ~want:ty v;
           gconst = true }
  | _ -> None

(* The entry point, plan.org: (defn main [args [string]] i32), with both the
   parameter and the return type optional. *)
let check_main env =
  match Hashtbl.find_opt env.fns "main" with
  | None -> ()   (* a library, or a file being checked on its own *)
  | Some (params, ret) ->
    let ok_params =
      match params with
      | [] -> true
      | [ Types.Slice Types.String ] -> true
      | _ -> false
    in
    if not ok_params then
      fail Loc.unknown
        "main takes no parameters or one [string], not (%s)"
        (String.concat " " (List.map Types.to_string params));
    if not (Types.equal ret Types.Unit || Types.equal ret (Types.Int Types.I32))
    then
      fail Loc.unknown "main returns i32 or nothing, not %s"
        (Types.to_string ret)

let program (decls : Ast.decl list) : Tast.program =
  let env = new_env () in
  let decls = Parse.program (Prelude.forms ()) @ decls in
  collect env decls;
  check_finite env;
  check_main env;
  let globals = List.filter_map (check_global env) decls in
  let fns =
    List.filter_map
      (fun (d : Ast.decl) ->
         match d.Ast.d with
         | Ast.Defn fn -> Some (check_fn env fn)
         | _ -> None)
      decls
  in
  (* Sorted, so the emitted IR is reproducible build to build: a Hashtbl's
     fold order is not. *)
  let values name tbl =
    Hashtbl.fold (fun _ v acc -> v :: acc) tbl []
    |> List.sort (fun a b -> String.compare (name a) (name b))
  in
  { Tast.structs = values (fun (s : Tast.structure) -> s.Tast.sname) env.structs;
    unions = values (fun (u : Tast.union) -> u.Tast.uname) env.unions;
    globals; fns }