
    Part of the Austral project , under the Apache License v2.0 with LLVM Exceptions . 
    See LICENSE file for details . 

    SPDX - License - Identifier : Apache-2.0 WITH LLVM - exception 

   Part of the Austral project, under the Apache License v2.0 with LLVM Exceptions.
   See LICENSE file for details.

   SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
*)

open Identifier
open Common
open Type
open TypeSystem
open TypeBindings
open TypeMatch
open TypeVarSet
open TypeParser
open TypeParameter
open TypeParameters
open Region
open Id
open LexEnv
open Env
open EnvTypes
open EnvUtils
open EnvExtras
open Ast
open Tast
open TastUtil
open Linked
open Util
open Reporter
open BuiltIn
open Error

module Errors = TypeErrors

(* Since the extraction pass has already happened, we can simplify the call to
   `parse_type` by passing an empty list of local type signatures. *)
let parse_typespec (env: env) (rm: region_map) (typarams: typarams) (ty: qtypespec): ty =
  parse_type env [] rm typarams ty

let rec augment_expr (module_name: module_name) (env: env) (rm: region_map) (typarams: typarams) (lexenv: lexenv) (asserted_ty: ty option) (expr: aexpr): texpr =
  with_frame ("Augment expression: " ^ (expr_kind expr))
    (fun _ ->
      let aug = augment_expr module_name env rm typarams lexenv None in
      match expr with
      | NilConstant ->
         TNilConstant
      | BoolConstant b ->
         TBoolConstant b
      | IntConstant i ->
         TIntConstant i
      | FloatConstant f ->
         TFloatConstant f
      | StringConstant s ->
         TStringConstant s
      | Variable name ->
         (match get_decl_by_name env (qident_to_sident name) with
          | Some (Function { id; typarams; value_params; rt; _ }) ->
             let arg_tys: ty list = List.map (fun (ValueParameter (_, ty)) -> ty) value_params in
             let fn_ty: ty = FnPtr (arg_tys, rt) in
             if (typarams_size typarams > 0) then
               (* Function is generic, need an asserted type. *)
               (match asserted_ty with
                | Some asserted_ty ->
                   let bindings: type_bindings = match_type (env, module_name) fn_ty asserted_ty in
                   let rt: ty = replace_variables bindings rt
                   and arg_tys: ty list = List.map (replace_variables bindings) arg_tys in
                   let effective_fn_ty: ty = FnPtr (arg_tys, rt) in
                   TFunVar (id, effective_fn_ty, bindings)
                | None ->
                   Errors.unconstrained_generic_function (original_name name))
             else
               (* Function is concrete. *)
               TFunVar (id, fn_ty, empty_bindings)
          | _ ->
             (match get_variable env lexenv name with
              | Some (ty, src) ->
                 (match src with
                  | VarConstant ->
                     TConstVar (name, ty)
                  | VarParam ->
                     TParamVar ((original_name name), ty)
                  | VarLocal ->
                     TLocalVar ((original_name name), ty))
              | None ->
                 Errors.unknown_name
                   ~kind:"variable"
                   ~name:(original_name name)))
      | FunctionCall (name, args) ->
         augment_call module_name env lexenv asserted_ty name (augment_arglist module_name env rm typarams lexenv None args)
      | ArithmeticExpression (op, lhs, rhs) ->
         (* Augment operands. *)
         let lhs' = aug lhs
         and rhs' = aug rhs in
         (* Get the types. *)
         let lhs_ty = get_type lhs'
         and rhs_ty = get_type rhs'
         in
         let augment_arithmetic_call _ =
             (* Because method calls don't implement the lenient type checking
                rules for arithmetic, we have to cast the arguments to their
                expected types, so that constants are allowed to be used as
                arguments. This means we have to transform a `ty` into a type
                specifier. *)
           let op_name =
             match op with
             | Add ->
                "trappingAdd"
             | Subtract ->
                "trappingSubtract"
             | Multiply ->
                "trappingMultiply"
             | Divide ->
                "trappingDivide"
           in
           let op_qname = make_qident (pervasive_module_name, make_ident op_name, make_ident op_name) in
           let args: typed_arglist = TPositionalArglist [
                          TCast (lhs', lhs_ty);
                          TCast (rhs', lhs_ty);
                        ]
           in
           augment_call module_name env lexenv asserted_ty op_qname args
         in
         (* Are the types the same? *)
         if lhs_ty = rhs_ty then
           (* If the types are the same type, check it is a numeric type. *)
           if (is_numeric lhs_ty) then
             augment_arithmetic_call ()
           else
             Errors.arithmetic_not_numeric lhs_ty
         else
  If the types are different , check if at least one operator is a constant .
           let are_int_constants = (is_int_constant lhs) || (is_int_constant rhs)
           and are_float_constants = (is_float_constant lhs) || (is_float_constant rhs) in
           if (are_int_constants || are_float_constants) then
             (* If either operand is a constant, let it pass *)
             augment_arithmetic_call ()
           else
             Errors.arithmetic_incompatible_types ~lhs:lhs_ty ~rhs:rhs_ty
      | Comparison (op, lhs, rhs) ->
         let lhs' = aug lhs
         and rhs' = aug rhs in
         let _ = match_type_with_value (env, module_name) (get_type lhs') rhs' in
         TComparison (op, lhs', rhs')
      | Conjunction (lhs, rhs) ->
         let lhs' = aug lhs
         and rhs' = aug rhs in
         if ((is_bool lhs') && (is_bool rhs')) then
           TConjunction (lhs', rhs')
         else
           Errors.logical_operands_not_boolean
             ~operator:"conjunction"
             ~types:[get_type lhs'; get_type rhs']
      | Disjunction (lhs, rhs) ->
         let lhs' = aug lhs
         and rhs' = aug rhs in
         if ((is_bool lhs') && (is_bool rhs')) then
           TDisjunction (lhs', rhs')
         else
           Errors.logical_operands_not_boolean
             ~operator:"disjunction"
             ~types:[get_type lhs'; get_type rhs']
      | Negation e ->
         let e' = aug e in
         if is_bool e' then
           TNegation e'
         else
           Errors.logical_operands_not_boolean
             ~operator:"negation"
             ~types:[get_type e']
      | IfExpression (c, t, f) ->
         let c' = aug c
         and t' = aug t
         and f' = aug f in
         if is_bool c' then
           if (get_type t') = (get_type f') then
             TIfExpression (c', t', f')
           else
             Errors.if_inequal
               ~lhs:(get_type t')
               ~rhs:(get_type f')
         else
           Errors.condition_not_boolean
             ~kind:"if"
             ~form:"expression"
             ~ty:(get_type c')
      | Path (e, elems) ->
  Path rules : 

             1 . Path that begins in a reference ends in a reference . 
             2 . All paths end in a type in the free universe . 
          

            1. Path that begins in a reference ends in a reference.
            2. All paths end in a type in the free universe.
          *)
         let e' = aug e in
         (* Is the initial expression of a path a read reference or write reference or no reference? *)
         let (is_read, region): (bool * ty option) =
           (match (get_type e') with
            | ReadRef (_, r) ->
               (true, Some r)
            | WriteRef (_, r) ->
               (false, Some r)
            | _ ->
               (false, None))
         in
         let elems' = augment_path env module_name rm typarams lexenv (get_type e') elems in
         let path_ty: ty = get_path_ty_from_elems elems' in
         let path_ty: ty =
           (* If the path starts in a reference it should end in one. *)
           (match region with
            | Some reg ->
               if is_read then
                 ReadRef (path_ty, reg)
               else
                 WriteRef (path_ty, reg)
            | None ->
               path_ty)
         in
         let universe = type_universe path_ty in
         if universe = FreeUniverse then
           TPath {
               head = e';
               elems = elems';
               ty = path_ty
             }
         else
           Errors.path_not_free path_ty
      | Embed (ty, expr, args) ->
         TEmbed (
             parse_typespec env rm typarams ty,
             expr,
             List.map aug args
           )
      | Deref expr ->
         (* The type of the expression being dereferenced must be either a read-only
            reference or a write reference. *)
         let expr' = aug expr in
         let ty = get_type expr' in
         (match ty with
          | ReadRef _ ->
             TDeref expr'
          | WriteRef _ ->
             TDeref expr'
          | _ ->
             Errors.dereference_non_reference ty)
      | Typecast (expr, ty) ->
  The typecast operator has four uses : 

             1 . Clarifying the type of integer and floating point literals . 

             2 . Converting write references to read references . 

             3 . Clarifying the type of return type polymorphic functions . 

          

            1. Clarifying the type of integer and floating point literals.

            2. Converting write references to read references.

            3. Clarifying the type of return type polymorphic functions.

          *)
         let is_int_expr = function
           | TIntConstant _ -> true
           | _ -> false
         and is_float_expr = function
           | TFloatConstant _ -> true
           | _ -> false

         and is_int_type = function
           | Integer _ -> true
           | _ -> false

         and is_float_type = function
           | SingleFloat -> true
           | DoubleFloat -> true
           | _ -> false
         in
         let target_type = parse_typespec env rm typarams ty in
  By passing target_type as the asserted type we 're doing point 3 .
         let expr' = augment_expr module_name env rm typarams lexenv (Some target_type) expr in
  For 1 : if the expression is an int literal and the target type is an 
             int type , do the conversion . Similarly , if the expression is a float 
             constant and the target type is a float , do the conversion .
            int type, do the conversion. Similarly, if the expression is a float
            constant and the target type is a float, do the conversion. *)
         if (is_int_expr expr') && (is_int_type target_type) then
           TCast (expr', target_type)
         else
           if (is_float_expr expr') && (is_float_type target_type) then
             TCast (expr', target_type)
           else
  Otherwise , if the expression has a mutable reference type and the 
                 target type is a read reference , and they point to the same underlying 
                 type and underlying region , just convert it to a read ref . This does 
                 point 2 .
                target type is a read reference, and they point to the same underlying
                type and underlying region, just convert it to a read ref. This does
                point 2. *)
             (match (get_type expr') with
              | WriteRef (underlying_ty, region) ->
                 (match target_type with
                  | ReadRef (underlying_ty', region') ->
                     if ((equal_ty underlying_ty underlying_ty') && (equal_ty region region')) then
                       TCast (expr', target_type)
                     else
                       Errors.cast_different_references
                         ~different_types:(not (equal_ty underlying_ty underlying_ty'))
                         ~different_regions:(not (equal_ty region region'))
                  | _ ->
                     Errors.cast_write_ref_to_non_ref ())
              | _ ->
                 (try
                    let _ = match_type (env, module_name) target_type (get_type expr') in
                    expr'
                  with Austral_error _ ->
                    Errors.cast_invalid
                      ~target:target_type
                      ~source:(get_type expr')))
      | SizeOf ty ->
         TSizeOf (parse_typespec env rm typarams ty)
      | BorrowExpr (mode, name) ->
         (match get_variable env lexenv name with
          | Some (ty, src) ->
             (* TODO: check if `ty` is linear? *)
             (match src with
              | VarConstant ->
                 Errors.borrow_constant ()
              | VarParam ->
                 let name: identifier = original_name name
                 and reg: region = fresh_region ()
                 in
                 TBorrowExpr (mode, name, reg, ty)
              | VarLocal ->
                 let name: identifier = original_name name
                 and reg: region = fresh_region ()
                 in
                 TBorrowExpr (mode, name, reg, ty))
          | None ->
             Errors.unknown_name
               ~kind:"variable"
               ~name:(original_name name)))

and get_path_ty_from_elems (elems: typed_path_elem list): ty =
  assert ((List.length elems) > 0);
  let last = List.nth elems ((List.length elems) - 1) in
  path_elem_type last

and is_bool e =
  match get_type e with
  | Boolean -> true
  | _ -> false

and is_int_constant e =
  match e with
  | IntConstant _ -> true
  | _ -> false

and is_float_constant e =
  match e with
  | FloatConstant _ -> true
  | _ -> false

and augment_arglist (module_name: module_name) (env: env) (rm: region_map) (typarams: typarams) (lexenv: lexenv) (asserted_ty: ty option) (args: abstract_arglist): typed_arglist =
  let aug = augment_expr module_name env rm typarams lexenv asserted_ty in
  match args with
  | Positional args' ->
     TPositionalArglist (List.map aug args')
  | Named pairs ->
     TNamedArglist (List.map (fun (n, v) -> (n, aug v)) pairs)

and augment_path (env: env) (module_name: module_name) (rm: region_map) (typarams: typarams) (lexenv: lexenv) (head_ty: ty) (elems: path_elem list): typed_path_elem list =
  match elems with
  | [elem] ->
     [augment_path_elem env module_name rm typarams lexenv head_ty elem]
  | elem::rest ->
     let elem' = augment_path_elem env module_name rm typarams lexenv head_ty elem in
     let rest' = augment_path env module_name rm typarams lexenv (path_elem_type elem') rest in
     elem' :: rest'
  | [] ->
     err "Path is empty"

and augment_path_elem (env: env) (module_name: module_name) (rm: region_map) (typarams: typarams) (lexenv: lexenv) (head_ty: ty) (elem: path_elem): typed_path_elem =
  match elem with
  | SlotAccessor slot_name ->
     (match head_ty with
      | NamedType (name, args, _) ->
         augment_slot_accessor_elem env module_name slot_name name args
      | _ ->
         Errors.path_not_record (type_string head_ty))
  | PointerSlotAccessor slot_name ->
     (match head_ty with
      | Pointer pointed_to ->
         (* TODO: Addresses should not be indexable. *)
         augment_pointer_slot_accessor_elem env module_name slot_name pointed_to
      | ReadRef (ty, _) ->
         (match ty with
          | NamedType (name, args, _) ->
             augment_reference_slot_accessor_elem env module_name slot_name name args
          | _ ->
             Errors.path_not_record (type_string ty))
      | WriteRef (ty, _) ->
         (match ty with
          | NamedType (name, args, _) ->
             augment_reference_slot_accessor_elem env module_name slot_name name args
          | _ ->
             Errors.path_not_record (type_string ty))
      | _ ->
         Errors.path_not_record (type_string head_ty))
  | ArrayIndex ie ->
     let ie' = augment_expr module_name env rm typarams lexenv None ie in
     (match head_ty with
      | StaticArray elem_ty ->
         TArrayIndex (ie', elem_ty)
      | _ ->
         Errors.array_indexing_disallowed head_ty)

and augment_slot_accessor_elem (env: env) (module_name: module_name) (slot_name: identifier) (type_name: qident) (type_args: ty list) =
  (* Check: e' is a public record type *)
  let (source_module, vis, typarams, slots) = get_record_definition env type_name in
  if (vis = TypeVisPublic) || (module_name = source_module) then
    (* Check: the given slot name must exist in this record type. *)
    let (TypedSlot (_, slot_ty)) = get_slot_with_name type_name slots slot_name in
    let bindings = match_typarams (env, module_name) typarams type_args in
    let slot_ty' = replace_variables bindings slot_ty in
    TSlotAccessor (slot_name, slot_ty')
  else
    Errors.path_not_public
      ~type_name:(original_name type_name)
      ~slot_name

and augment_pointer_slot_accessor_elem (env: env) (module_name: module_name) (slot_name: identifier) (pointed_to: ty): typed_path_elem =
  match pointed_to with
  | NamedType (type_name, type_args, _) ->
     (* Check arg is a public record *)
     let (source_module, vis, typarams, slots) = get_record_definition env type_name in
     if (vis = TypeVisPublic) || (module_name = source_module) then
       (* Check: the given slot name must exist in this record type. *)
       let (TypedSlot (_, slot_ty)) = get_slot_with_name type_name slots slot_name in
       let bindings = match_typarams (env, module_name) typarams type_args in
       let slot_ty' = replace_variables bindings slot_ty in
       TPointerSlotAccessor (slot_name, slot_ty')
     else
       Errors.path_not_public
         ~type_name:(original_name type_name)
         ~slot_name
  | _ ->
     Errors.path_not_record (type_string pointed_to)

and augment_reference_slot_accessor_elem (env: env) (module_name: module_name) (slot_name: identifier) (type_name: qident) (type_args: ty list) =
  (* Check: e' is a public record type *)
  let (source_module, vis, typarams, slots) = get_record_definition env type_name in
  if (vis = TypeVisPublic) || (module_name = source_module) then
    (* Check: the given slot name must exist in this record type. *)
    let (TypedSlot (_, slot_ty)) = get_slot_with_name type_name slots slot_name in
    let bindings = match_typarams (env, module_name) typarams type_args in
    let slot_ty' = replace_variables bindings slot_ty in
    TPointerSlotAccessor (slot_name, slot_ty')
  else
    Errors.path_not_public
      ~type_name:(original_name type_name)
      ~slot_name

and get_record_definition (env: env) (name: qident): (module_name * type_vis * typarams * typed_slot list) =
  match get_decl_by_name env (qident_to_sident name) with
  | (Some (Record { mod_id; vis; typarams; slots; _ })) ->
     let mod_name: module_name = module_name_from_id env mod_id in
     (mod_name, vis, typarams, slots)
  | Some _ ->
     Errors.path_not_record (original_name name |> ident_string)
  | None ->
     err ("No record with this name: " ^ (ident_string (original_name name)))

and get_slot_with_name type_name slots slot_name =
  match List.find_opt (fun (TypedSlot (n, _)) -> n = slot_name) slots with
  | Some s -> s
  | None -> 
     Errors.no_such_slot
       ~type_name:(original_name type_name)
       ~slot_name

and augment_call (module_name: module_name) (env: env) (lexenv: lexenv) (asserted_ty: ty option) (name: qident) (args: typed_arglist): texpr =
  First , check if the callable name is a variable .
  match get_var lexenv (local_name name) with
  | Some (ty, _) ->
     augment_fptr_call module_name env (local_name name) ty asserted_ty args
  | None ->
     (* Otherwise, find a callable from the environment. *)
     (match get_callable env module_name (qident_to_sident name) with
      | Some callable ->
         augment_callable module_name env name callable asserted_ty args
      | None ->
         Errors.unknown_name
           ~kind:"callable"
           ~name:(original_name name))

and augment_fptr_call (module_name: module_name) (env: env) (name: identifier) (fn_ptr_ty: ty) (asserted_ty: ty option) (args: typed_arglist): texpr =
  (* Because function pointers don't preserve value parameter names, we can't
     accept named argument lists. *)
  let args: texpr list =
    match args with
    | TPositionalArglist args -> args
    | TNamedArglist _ ->
       Errors.fun_pointer_named_args ()
  in
  (* We extract the parameter types and return type from the function pointer
     type. *)
  let (paramtys, rt): (ty list * ty) =
    match fn_ptr_ty with
    | FnPtr (paramtys, rt) -> (paramtys, rt)
    | _ -> Errors.call_non_callable fn_ptr_ty
  in
  (* Then, we synthesize a parameter list from the function pointer type. The
     parameter names are fake, but this is so we can reuse the arity-checking,
     value parameter list-checking infrastructure we use elsewhere. *)
  let params: value_parameter list = List.map (fun paramty -> ValueParameter (make_ident "fake", paramty)) paramtys in
  (* Check synthetic parameter list against argument list. *)
  let bindings = check_argument_list env module_name name params args in
  (* Use the bindings to get the effective return type *)
  let rt' = replace_variables bindings rt in
  let (bindings', rt'') = handle_return_type_polymorphism env module_name name params rt' asserted_ty bindings in
  let arguments = cast_arguments bindings' params args in
  TFptrCall (name, arguments, rt'')

and augment_callable (module_name: module_name) (env: env) (name: qident) (callable: callable) (asserted_ty: ty option) (args: typed_arglist) =
  with_frame "Augment callable"
    (fun _ ->
      match callable with
      | FunctionCallable (id, typarams, params, rt) ->
         augment_function_call env module_name id name typarams params rt asserted_ty args
      | RecordConstructor (_, typarams, universe, slots) ->
         augment_record_constructor env module_name name typarams universe slots asserted_ty args
      | UnionConstructor { union_id; type_params; universe; case } ->
         let type_name: qident = (match get_decl_by_id env union_id with
                                  | Some (Union { name; mod_id; _ }) ->
                                     (* TODO: constructing a fake qident *)
                                     make_qident (module_name_from_id env mod_id, name, name)
                                  | _ ->
                                     internal_err ("no type name found for ID `" ^ (show_decl_id union_id) ^ "`"))
         in
         augment_union_constructor env module_name type_name type_params universe case asserted_ty args
      | MethodCallable { typeclass_id; value_parameters; return_type; _ } ->
         let param = (match get_decl_by_id env typeclass_id with
                      | Some (TypeClass { param; _ }) ->
                         param
                      | _ ->
                         internal_err ("no type parameter found for ID `" ^ (show_decl_id typeclass_id) ^ "`"))
         in
         augment_method_call env module_name typeclass_id param name value_parameters return_type asserted_ty args)

and augment_function_call (env: env) (module_name: module_name) (id: decl_id) name typarams params rt asserted_ty args =
  with_frame "Augment function call"
    (fun _ ->
      pqi ("Name", name);
      (* For simplicity, and to reduce duplication of code, we convert the argument
         list to a positional list. *)
      let param_names = List.map (fun (ValueParameter (n, _)) -> n) params in
      let arguments: texpr list = arglist_to_positional (args, param_names) in
      (* Check the list of params against the list of arguments *)
      let bindings = check_argument_list env module_name (original_name name) params arguments in
      (* Use the bindings to get the effective return type *)
      let rt' = replace_variables bindings rt in
      let (bindings', rt'') = handle_return_type_polymorphism env module_name (local_name name) params rt' asserted_ty bindings in
      (* Check: the set of bindings equals the set of type parameters *)
      let bindings'' = merge_bindings bindings bindings' in
      check_bindings typarams bindings'';
      let arguments' = cast_arguments bindings'' params arguments in
      let substs = make_substs bindings'' typarams in
      ps ("Bindings", show_bindings bindings'');
      pt ("Return type", rt'');
      TFuncall (id, name, arguments', rt'', substs))

and augment_record_constructor (env: env) (module_name: module_name) (name: qident) (typarams: typarams) (universe: universe) (slots: typed_slot list) (asserted_ty: ty option) (args: typed_arglist) =
  (* Check: the argument list must be named *)
  let args' = (match args with
               | TPositionalArglist l ->
                  if l = [] then
                    []
                  else
                    Errors.constructor_not_named "record"
               | TNamedArglist a ->
                  a) in
  (* Check: the set of slots matches the set of param names *)
  let slot_names: identifier list = List.map (fun (TypedSlot (name, _)) -> name) slots in
  let argument_names: identifier list = List.map (fun (n, _) -> n) args' in
  if not (ident_set_eq slot_names argument_names) then
    Errors.constructor_wrong_args "record"
  else
    (* Convert args to positional *)
    let arguments = arglist_to_positional (args, slot_names) in
    (* Check the list of params against the list of arguments *)
    let params = List.map (fun (TypedSlot (n, t)) -> ValueParameter (n, t)) slots in
    let bindings = check_argument_list env module_name (original_name name) params arguments in
    (* Use the bindings to get the effective return type *)
    let rt = NamedType (name,
                        List.map (fun tp -> TyVar (typaram_to_tyvar tp)) (typarams_as_list typarams),
                        universe)
    in
    let rt' = replace_variables bindings rt in
    let (bindings', rt'') = handle_return_type_polymorphism env module_name (local_name name) params rt' asserted_ty bindings in
    (* Check: the set of bindings equals the set of type parameters *)
    check_bindings typarams (merge_bindings bindings bindings');
    (* Check the resulting type is in the correct universe *)
    if universe_compatible universe (type_universe rt'') then
      TRecordConstructor (rt'', List.map2 (fun a b -> (a, b)) slot_names arguments)
    else
      err "Universe mismatch"

and augment_union_constructor (env: env) (module_name: module_name) (type_name: qident) (typarams: typarams) (universe: universe) (case: typed_case) (asserted_ty: ty option) (args: typed_arglist) =
  with_frame "Augment union constructor"
    (fun _ ->
      pqi ("Type name", type_name);
      let args' = (match args with
                   | TPositionalArglist l ->
                      if l = [] then
                        []
                      else
                        Errors.constructor_not_named "union"
                   | TNamedArglist a ->
                      a) in
      (* Check: the set of slots matches the set of param names *)
      let (TypedCase (case_name, slots)) = case in
      pi ("Case name", case_name);
      let slot_names: identifier list = List.map (fun (TypedSlot (name, _)) -> name) slots in
      let argument_names: identifier list = List.map (fun (n, _) -> n) args' in
      if not (ident_set_eq slot_names argument_names) then
        Errors.constructor_wrong_args "union"
      else
        (* Convert args to positional *)
        let arguments = arglist_to_positional (args, slot_names) in
        (* Check the list of params against the list of arguments *)
        let params = List.map (fun (TypedSlot (n, t)) -> ValueParameter (n, t)) slots in
        ps ("Params", String.concat "\n" (List.map (fun (ValueParameter (n, t)) -> (ident_string n) ^ ": " ^ (type_string t)) params));
        ps ("Arguments", String.concat "\n" (List.map show_texpr arguments));
        let bindings = check_argument_list env module_name case_name params arguments in
        (* Use the bindings to get the effective return type *)
        let rt = NamedType (type_name,
                            List.map (fun tp -> TyVar (typaram_to_tyvar tp)) (typarams_as_list typarams),
                            universe)
        in
        pt ("Type", rt);
        ps ("Bindings", show_bindings bindings);
        let rt' = replace_variables bindings rt in
        pt ("Type'", rt');
        let (bindings', rt'') = handle_return_type_polymorphism env module_name case_name params rt' asserted_ty bindings in
        (* Check: the set of bindings equals the set of type parameters *)
        check_bindings typarams (merge_bindings bindings bindings');
        (* Check the resulting type is in the correct universe *)
        if universe_compatible universe (type_universe rt'') then
          let _ = pt ("Type''", rt'')
          in
          TUnionConstructor (rt'', case_name, List.map2 (fun a b -> (a, b)) slot_names arguments)
        else
          err "Universe mismatch")

and augment_method_call (env: env) (source_module_name: module_name) (typeclass_id: decl_id) (typaram: type_parameter) (callable_name: qident) (params: value_parameter list) (rt: ty) (asserted_ty: ty option) (args: typed_arglist): texpr =
  with_frame "Augmenting method call"
    (fun _ ->
      pqi ("Callable name", callable_name);
      (* At this point, we know the method's name and the typeclass it belongs
         to. We pull the parameter list from the typeclass, and compare that against
         the argument list. *)
      (* For simplicity, and to reduce duplication of code, we convert the argument
         list to a positional list. *)
      let param_names = List.map (fun (ValueParameter (n, _)) -> n) params in
      let arguments: texpr list = arglist_to_positional (args, param_names) in
      (* Check the list of params against the list of arguments *)
      let bindings = check_argument_list env source_module_name (original_name callable_name) params arguments in
      (* Use the bindings to get the effective return type *)
      let rt' = replace_variables bindings rt in
      let (bindings', rt'') = handle_return_type_polymorphism env source_module_name (local_name callable_name) params rt' asserted_ty bindings in
      let bindings'' = merge_bindings bindings bindings' in
      ps ("Bindings", show_bindings bindings'');
      (* Check: the set of bindings equals the set of type parameters *)
      check_bindings (typarams_from_list [typaram]) bindings'';
      match get_binding bindings'' typaram with
      | (Some dispatch_ty) ->
         pt ("Dispatch Type", dispatch_ty);
         (* Is the dispatch type a type variable? *)
         (match dispatch_ty with
          | TyVar (TypeVariable (typaram, _, _, constraints)) ->
             (* If so, check if the constraints say that it implements the typeclass. *)
             let tc_name: sident = get_decl_sident_or_die env typeclass_id in
             if List.exists (fun c -> equal_sident tc_name c) constraints then
               (* If it's a tyvar that implements the typeclass, we have to
                  proceed differently. Essentially we can't go on with instance
                  resolution, because we don't have a type to resolve an
                  instance for. So we have to freeze the process. *)
               TVarMethodCall {
                   source_module_name = source_module_name;
                   typeclass_id = typeclass_id;
                   params = params;
                   method_name = callable_name;
                   args = arguments;
                   dispatch_ty = dispatch_ty;
                   rt = rt'';
                   bindings = bindings'';
                 }
             else
               (* If it doesn't, that's an error *)
               Errors.unconstrained_type_parameter
                 ~typeclass:(sident_name tc_name)
                 ~typaram
          | _ ->
             (* If it's not a type variable, continue normal instance resolution. *)
             let (instance, instance_bindings): decl * type_bindings =
               (match get_instance env source_module_name dispatch_ty typeclass_id with
                | Some (i, b) -> (i, b)
                | None ->
                   let tc_name: sident = get_decl_sident_or_die env typeclass_id in
                   Errors.no_suitable_instance
                     ~typeclass:(sident_name tc_name)
                     ~ty:dispatch_ty)
             in
             ps ("Instance bindings", show_bindings instance_bindings);
             let params' = List.map (fun (ValueParameter (n, t)) -> ValueParameter (n, replace_variables instance_bindings t)) params in
             let arguments' = cast_arguments instance_bindings params' arguments in
             let typarams = (match instance with
                             | Instance { typarams; _ } -> typarams
                             | _ -> internal_err "Couldn't get instance while getting type parameters")
             in
             let instance_id: decl_id = decl_id instance in
             let meth_id: ins_meth_id =
               (match get_instance_method_from_instance_id_and_method_name env instance_id (original_name callable_name) with
                | Some (InsMethRec { id; _ }) -> id
                | None -> internal_err ("Couldn't get instance method for `"
                                        ^ (ident_string (original_name callable_name))
                                        ^ "`"))
             in
             ps ("Bindings", show_bindings bindings'');
             let substs = make_substs instance_bindings typarams in
             TMethodCall (meth_id, callable_name, typarams, arguments', rt'', substs))
      | None ->
         internal_err "couldn't extract dispatch type.")

(* Given a list of type parameters, and a list of arguments, check that the
   lists have the same length and each argument satisfies each corresponding
   parameter. Return the resulting set of type bindings. *)
and check_argument_list (env: env) (module_name: module_name) (name: identifier) (params: value_parameter list) (args: texpr list): type_bindings =
  (* Check that the number of parameters is the same as the number of
     arguments. *)
  check_arity name params args;
  (* Check arguments against parameters *)
  check_argument_list' env module_name empty_bindings params args

and check_arity (name: identifier) (params: value_parameter list) (args: texpr list): unit =
  let nparams = List.length params
  and nargs = List.length args in
  if nparams = nargs then
    ()
  else
    Errors.call_wrong_arity ~name ~expected:nparams ~actual:nargs

(* Precondition: both lists have the same length. *)
and check_argument_list' (env: env) (module_name: module_name) (bindings: type_bindings) (params: value_parameter list) (args: texpr list): type_bindings =
  match (params, args) with
  | ((first_param::rest_params), (first_arg::rest_args)) ->
     let bindings' = merge_bindings bindings (match_parameter env module_name first_param first_arg) in
     check_argument_list' env module_name bindings' rest_params rest_args
  | ([], []) ->
     bindings
  | _ ->
     internal_err ("couldn't check argument list in module `"
                   ^ (mod_name_string module_name)
                   ^ "` with bindings: "
                   ^ (show_type_bindings bindings)
                   ^ ",\n params: `{ "
                   ^ (String.concat ", " (List.map show_value_parameter params))
                   ^ " }`,\n and args: `{ "
                   ^ (String.concat ", " (List.map show_texpr args))
                   ^ "}`")

and match_parameter (env: env) (module_name: module_name) (param: value_parameter) (arg: texpr): type_bindings =
  let (ValueParameter (_, ty)) = param in
  match_type_with_value (env, module_name) ty arg

and cast_arguments (bindings: type_bindings) (params: value_parameter list) (arguments: texpr list): texpr list =
  let f (ValueParameter (_, expected)) value =
    let expected' = replace_variables bindings expected in
    match expected' with
    | Integer _ ->
       (* Cast integer types *)
       TCast (value, expected')
    | _ ->
       value
  in
  List.map2 f params arguments

and make_substs (bindings: type_bindings) (typarams: typarams): type_bindings =
  let f (tp: type_parameter): (type_parameter * ty) option =
    if (typaram_universe tp) = RegionUniverse then
      Some (tp, RegionTy static_region)
    else
      match get_binding bindings tp with
      | Some ty ->
         Some (tp, ty)
      | None ->
         None
  in
  bindings_from_list (List.filter_map f (typarams_as_list typarams))

and handle_return_type_polymorphism (env: env) (module_name: module_name) (name: identifier) (params: value_parameter list) (rt: ty) (asserted_ty: ty option) (bindings: type_bindings): (type_bindings * ty) =
  if is_return_type_polymorphic params rt then
    match asserted_ty with
    | (Some asserted_ty') ->
       let bindings' = match_type (env, module_name) rt asserted_ty' in
       let bindings'' = merge_bindings bindings bindings' in
       (bindings'', replace_variables bindings'' rt)
    | None ->
       Errors.unconstrained_generic_function name
  else
    (empty_bindings, replace_variables bindings rt)

(* Given a function's parameter list and return type, check if the return type
   of a function is polymorphic. *)
and is_return_type_polymorphic (params: value_parameter list) (rt: ty): bool =
  (* Find the set of type variables in the return type. *)
  let rt_type_vars: TypeVarSet.t = type_variables rt in
  (* Find the set of type variables in each of the parameters. *)
  let param_type_vars: TypeVarSet.t list = List.map (fun (ValueParameter (_, ty)) -> type_variables ty) params in
  (* Union the type variables of the parameters into one set of all type
     variables in the parameter list. *)
  let param_type_vars: TypeVarSet.t = List.fold_left TypeVarSet.union TypeVarSet.empty param_type_vars in
  (* Find the set of type variables that are in the return type set but not in
     the parameter set. *)
  let vars_not_params: TypeVarSet.t =
    let param_ty_vars: type_var list = TypeVarSet.elements param_type_vars
    in
    let var_name_in_typarams (n: identifier): bool =
      List.exists (fun (TypeVariable (n', _, _, _)) -> equal_identifier n n') param_ty_vars
    in
    TypeVarSet.of_list (List.filter (fun (TypeVariable (n, _, _, _)) -> not (var_name_in_typarams n)) (TypeVarSet.elements rt_type_vars))
  in
  (TypeVarSet.cardinal vars_not_params) > 0

(* Check that there are as many bindings as there are type parameters, and that
   every type parameter is satisfied. *)
and check_bindings (typarams: typarams) (bindings: type_bindings): unit =
  if (typarams_size typarams) = (binding_count bindings) then
    let check (tp: type_parameter): unit =
      let n = typaram_name tp
      and u = typaram_universe tp
      in
      (match get_binding bindings tp with
       | Some ty ->
          if universe_compatible u (type_universe ty) then
            ()
          else
            err ("Mismatched universes: expected "
                 ^ (show_universe u)
                 ^ " and got "
                 ^ (show_universe (type_universe ty)))
       | None ->
          err ("No binding for this parameter: " ^ (ident_string n)))
    in
    let _ = List.map check (typarams_as_list typarams) in
    ()
  else
    (* I think this should not be an error *)
err ( " Not the same number of bindings and parameters . Bindings : " 
       ^ ( show_bindings bindings ) 
       ^ " . Parameters : " 
       ^ ( String.concat " , " ( List.map ( fun ( TypeParameter ( n , u ) ) - > ( ident_string n ) ^ " : " ^ ( universe_string u ) ) typarams ) ) )
      ^ (show_bindings bindings)
      ^ ". Parameters: "
      ^ (String.concat ", " (List.map (fun (TypeParameter (n, u)) -> (ident_string n) ^ " : " ^ (universe_string u)) typarams)))*)
    ()

let is_boolean = function
  | Boolean -> true
  | _ -> false

let is_compatible_with_index_type = function
  | TIntConstant _ ->
     true
  | e ->
     (get_type e) = index_type

type stmt_ctx = StmtCtx of module_name * env * region_map * typarams * lexenv * ty

let update_lexenv (StmtCtx (mn, menv, rm, typarams, _, rt)) (lexenv: lexenv): stmt_ctx =
  StmtCtx (mn, menv, rm, typarams, lexenv, rt)

let update_rm (StmtCtx (mn, menv, _, typarams, lexenv, rt)) (rm: region_map): stmt_ctx =
  StmtCtx (mn, menv, rm, typarams, lexenv, rt)

let stmt_kind (stmt: astmt): string =
  match stmt with
  | ASkip _ -> "skip"
  | ALet _ -> "let"
  | ADestructure _ -> "let (destructure)"
  | AAssign _ -> "assign"
  | AIf _ -> "if"
  | AWhen _ -> "if (no else)"
  | ACase _ -> "case"
  | AWhile _ -> "while"
  | AFor _ -> "for"
  | ABorrow _ -> "borrow"
  | ABlock _ -> "block"
  | ADiscarding _ -> "discard"
  | AReturn _ -> "return"

let rec augment_stmt (ctx: stmt_ctx) (stmt: astmt): tstmt =
  with_frame ("Augment statement: " ^ (stmt_kind stmt))
    (fun _ ->
      let (StmtCtx (module_name, env, rm, typarams, lexenv, rt)) = ctx in
      match stmt with
      | ASkip span ->
         TSkip span
      | ALet (span, name, ty, value, body) ->
         pi ("Name", name);
         adorn_error_with_span span
           (fun _ ->
             let expected_ty = parse_typespec env rm typarams ty in
             pt ("Expected type", expected_ty);
             let value' = augment_expr module_name env rm typarams lexenv (Some expected_ty) value in
             let bindings = match_type_with_value (env, module_name) expected_ty value' in
             ps ("Bindings", show_bindings bindings);
             let ty = replace_variables bindings expected_ty in
             pt ("Type", ty);
             pt ("Value type", get_type value');
             let lexenv' = push_var lexenv name ty VarLocal in
             let body' = augment_stmt (update_lexenv ctx lexenv') body in
             TLet (span, name, ty, value', body'))
      | ADestructure (span, bindings, value, body) ->
         adorn_error_with_span span
           (fun _ ->
             let value' = augment_expr module_name env rm typarams lexenv None value in
             (* Check: the value must be a public record type *)
             let rec_ty = get_type value' in
             (match rec_ty with
              | (NamedType (name, _, u)) ->
                 let (source_module, vis, record_typarams, slots) = get_record_definition env name in
                 let orig_type = NamedType (
                                     make_qident (source_module, original_name name, original_name name),
                                     List.map (fun tp -> TyVar (typaram_to_tyvar tp)) (typarams_as_list record_typarams),
                                     u
                                   )
                 in
                 let typebindings = match_type (env, module_name) orig_type rec_ty in
                 if (vis = TypeVisPublic) || (module_name = source_module) then
                   (* Find the set of slot names and the set of binding names, and compare them *)
                   let binding_names = List.map (fun (QBinding { name; _}) -> name) bindings
                   and slot_names = List.map (fun (TypedSlot (n, _)) -> n) slots in
                   if ident_set_eq binding_names slot_names then
                     let bindings': (identifier * qtypespec * ty * identifier) list =
                       group_bindings_slots bindings slots
                     in
                     let bindings'': (identifier * ty * ty * identifier) list
                       = List.map (fun (n, ty, actual, rename) -> (n, parse_typespec env rm typarams ty, replace_variables typebindings actual, rename)) bindings'
                     in
                     let newvars: (identifier * ty * var_source) list =
                       List.map (fun (_, ty, actual, rename) ->
                           let _ = match_type (env, module_name) ty actual in
                           (rename, ty, VarLocal))
                         bindings''
                     in
                     let lexenv' = push_vars lexenv newvars in
                     let body' = augment_stmt (update_lexenv ctx lexenv') body in
                     TDestructure (
                         span,
                         List.map (fun (name, _, ty, rename) -> TypedBinding { name; ty; rename; }) bindings'',
                         value',
                         body'
                       )
                   else
                     Errors.destructure_wrong_slots rec_ty
                 else
                   Errors.destructure_not_public rec_ty
              | _ ->
                 Errors.destructure_non_record rec_ty))
      | AAssign (span, LValue (var, elems), value) ->
         adorn_error_with_span span
           (fun _ ->
             (match get_var lexenv var with
              | Some (var_ty, _) ->
                 (match elems with
                  | [] ->
                     (* Assigning to a variable. *)
                     let value = augment_expr module_name env rm typarams lexenv None value in
                     let universe = type_universe (get_type value) in
                     if universe = FreeUniverse then
                       let _ = match_type_with_value (env, module_name) var_ty value in
                       TAssign (span, TypedLValue (var, []), value)
                     else
                       Errors.lvalue_not_free ()
                  | elems ->
                     (* Assigning to a path. *)
                     let elems = augment_lvalue_path env module_name rm typarams lexenv var_ty elems in
                     let value = augment_expr module_name env rm typarams lexenv None value in
                     let ty = get_path_ty_from_elems elems in
                     let _ = match_type_with_value (env, module_name) ty value in
                     let path = TPath {
                                    head = value;
                                    elems = elems;
                                    ty = ty
                                  }
                     in
                     let universe = type_universe (get_type path) in
                     if universe = FreeUniverse then
                       TAssign (span, TypedLValue (var, elems), value)
                     else
                       Errors.lvalue_not_free ())
              | None ->
                 Errors.unknown_name ~kind:"variable" ~name:var))
      | AIf (span, c, t, f) ->
         adorn_error_with_span span
           (fun _ ->
             let c' = augment_expr module_name env rm typarams lexenv None c in
             if is_boolean (get_type c') then
               TIf (span, c', augment_stmt ctx t, augment_stmt ctx f)
             else
               Errors.condition_not_boolean
                 ~kind:"if"
                 ~form:"statement"
                 ~ty:(get_type c'))
      | AWhen (span, c, t) ->
         adorn_error_with_span span
           (fun _ ->
             let c = augment_expr module_name env rm typarams lexenv None c in
             if is_boolean (get_type c) then
               TIf (span, c, augment_stmt ctx t, TSkip span)
             else
               Errors.condition_not_boolean
                 ~kind:"if"
                 ~form:"statement"
                 ~ty:(get_type c))
      | ACase (span, expr, whens) ->
  Type checking a case statement : 

             1 . Ensure the value is of a union type . 
             2 . Ensure the union type is public or it is defined in this module . 
             3 . Ensure the set of case names in the case statement equals the set of cases in the union definition . 
             4 . Iterate over the cases , and ensure the bindings are correct . 
          

            1. Ensure the value is of a union type.
            2. Ensure the union type is public or it is defined in this module.
            3. Ensure the set of case names in the case statement equals the set of cases in the union definition.
            4. Iterate over the cases, and ensure the bindings are correct.
          *)
         adorn_error_with_span span
           (fun _ ->
             let expr' = augment_expr module_name env rm typarams lexenv None expr in
             let ty = get_type expr' in
             pt ("Case type", ty);
             let (union_ty, cases) = get_union_type_definition module_name env (get_type expr') in
             let typebindings = match_type (env, module_name) union_ty ty in
             let case_names = List.map (fun (TypedCase (n, _)) -> n) (List.map union_case_to_typed_case cases) in
             let when_names = List.map (fun (AbstractWhen (n, _, _)) -> n) whens in
             if ident_set_eq case_names when_names then
               (* Group the cases and whens *)
               let whens' = group_cases_whens cases whens in
               let whens'' = List.map (fun (c, w) -> augment_when ctx typebindings w c) whens' in
               TCase (span, expr', whens'')
             else
               Errors.case_non_exhaustive ())
      | AWhile (span, c, body) ->
         adorn_error_with_span span
           (fun _ ->
             let c' = augment_expr module_name env rm typarams lexenv None c in
             if is_boolean (get_type c') then
               TWhile (span, c', augment_stmt ctx body)
             else
               Errors.condition_not_boolean
                 ~kind:"while"
                 ~form:"statement"
                 ~ty:(get_type c'))
      | AFor { span; name; initial; final; body; } ->
         adorn_error_with_span span
           (fun _ ->
             let i' = augment_expr module_name env rm typarams lexenv None initial
             and f' = augment_expr module_name env rm typarams lexenv None final in
             if is_compatible_with_index_type i' then
               if is_compatible_with_index_type f' then
                 let lexenv' = push_var lexenv name index_type VarLocal in
                 let b' = augment_stmt (update_lexenv ctx lexenv') body in
                 TFor (span, name, i', f', b')
               else
                 Errors.for_bounds_non_integral ~bound:"final" ~ty:(get_type f')
             else
               Errors.for_bounds_non_integral ~bound:"initial" ~ty:(get_type i'))
      | ABorrow { span; original; rename; region; body; mode; } ->
         adorn_error_with_span span
           (fun _ ->
             (match get_var lexenv original with
              | (Some (orig_ty, _)) ->
                 let u = type_universe orig_ty in
                 if ((u = LinearUniverse) || (u = TypeUniverse)) then
                   let region_obj = fresh_region () in
                   let refty =
                     (match mode with
                      | ReadBorrow ->
                         ReadRef (orig_ty, RegionTy region_obj)
                      | WriteBorrow ->
                         WriteRef (orig_ty, RegionTy region_obj))
                   in
                   let lexenv' = push_var lexenv rename refty VarLocal in
                   let rm' = add_region rm region region_obj in
                   let ctx' = update_lexenv ctx lexenv' in
                   let ctx''= update_rm ctx' rm' in
                   TBorrow {
                       span=span;
                       original=original;
                       rename=rename;
                       region=region;
                       orig_type=orig_ty;
                       ref_type=refty;
                       body=augment_stmt ctx'' body;
                       mode=mode
                     }
                 else
                   Errors.borrow_non_linear orig_ty
              | None ->
                 Errors.unknown_name ~kind:"variable" ~name:original))
      | ABlock (span, f, r) ->
         let a = augment_stmt ctx f in
         let b = augment_stmt ctx r in
         TBlock (span, a, b)
      | ADiscarding (span, e) ->
         adorn_error_with_span span
           (fun _ ->
             let e' = augment_expr module_name env rm typarams lexenv None e in
             ps ("Expression", show_texpr e');
             let u = type_universe (get_type e') in
             if ((u = LinearUniverse) || (u = TypeUniverse)) then
               Errors.discard_linear u
             else
               TDiscarding (span, e'))
      | AReturn (span, e) ->
         adorn_error_with_span span
           (fun _ ->
             let e' = augment_expr module_name env rm typarams lexenv (Some rt) e in
             pt ("Type", get_type e');
             let _ = match_type_with_value (env, module_name) rt e' in
             TReturn (span, e')))

and augment_lvalue_path (env: env) (module_name: module_name) (rm: region_map) (typarams: typarams) (lexenv: lexenv) (head_ty: ty) (elems: path_elem list): typed_path_elem list =
  match elems with
  | [elem] ->
     [augment_lvalue_path_elem env module_name rm typarams lexenv head_ty elem]
  | elem::rest ->
     let elem' = augment_lvalue_path_elem env module_name rm typarams lexenv head_ty elem in
     let rest' = augment_lvalue_path env module_name rm typarams lexenv (path_elem_type elem') rest in
     elem' :: rest'
  | [] ->
     err "Path is empty"

and augment_lvalue_path_elem (env: env) (module_name: module_name) (rm: region_map) (typarams: typarams) (lexenv: lexenv) (head_ty: ty) (elem: path_elem): typed_path_elem =
  match elem with
  | SlotAccessor slot_name ->
     (match head_ty with
      | NamedType (name, args, _) ->
         augment_slot_accessor_elem env module_name slot_name name args
      | _ ->
         Errors.path_not_record (type_string head_ty))
  | PointerSlotAccessor slot_name ->
     (match head_ty with
      | Pointer pointed_to ->
         (* TODO: addresses should not be indexable *)
         augment_pointer_slot_accessor_elem env module_name slot_name pointed_to
      | WriteRef (ty, _) ->
         (match ty with
          | NamedType (name, args, _) ->
             augment_reference_slot_accessor_elem env module_name slot_name name args
          | _ ->
             Errors.path_not_record (type_string ty))
      | _ ->
         Errors.path_not_record (type_string head_ty))
  | ArrayIndex ie ->
     let ie' = augment_expr module_name env rm typarams lexenv None ie in
     let _ = ie' in
     (match head_ty with
      | _ ->
         Errors.lvalue_index ())


and get_union_type_definition (importing_module: module_name) (env: env) (ty: ty): (ty * decl list) =
  let name: qident = (match ty with
                      | NamedType (n, _, _) ->
                         n
                      | _ ->
                         Errors.case_non_union ty) in
  match get_decl_by_name env (qident_to_sident name) with
  | Some (Union { id; vis; name; mod_id; typarams; universe; _ }) ->
     let module_name = module_name_from_id env mod_id in
     if (vis = TypeVisPublic) || (importing_module = module_name) then
       let n = make_qident (module_name, name, name)
       and args = List.map (fun tp -> TyVar (typaram_to_tyvar tp)) (typarams_as_list typarams)
       in
       let ty = NamedType (n, args, universe)
       and cases = get_union_cases env id in
       (ty, cases)
     else
       Errors.case_non_public ()
  | _ ->
     Errors.case_non_union ty

and group_cases_whens (cases: decl list) (whens: abstract_when list): (typed_case * abstract_when) list =
  List.map (fun (TypedCase (n, s)) -> (TypedCase (n, s), List.find (fun (AbstractWhen (n', _, _)) -> n = n') whens))
    (List.map union_case_to_typed_case cases)

and group_bindings_slots (bindings: qbinding list) (slots: typed_slot list): (identifier * qtypespec * ty * identifier) list =
  let f (binding: qbinding): (identifier * qtypespec * ty * identifier) =
    let QBinding { name; ty; rename; } = binding in
    let (TypedSlot (_, ty')) = List.find (fun (TypedSlot (n', _)) -> equal_identifier name n') slots in
    (name, ty, ty', rename)
  in
  List.map f bindings

and augment_when (ctx: stmt_ctx) (typebindings: type_bindings) (w: abstract_when) (c: typed_case): typed_when =
  with_frame "Augment when"
    (fun _ ->
      let (StmtCtx (_, menv, rm, typarams, lexenv, _)) = ctx in
      let (AbstractWhen (name, bindings, body)) = w
      and (TypedCase (_, slots)) = c in
      pi ("Case name", name);
      (* Check the set of binding names is the same as the set of slots *)
      let binding_names = List.map (fun (QBinding { name; _ }) -> name) bindings
      and slot_names = List.map (fun (TypedSlot (n, _)) -> n) slots in
      if ident_set_eq binding_names slot_names then
        (* Check the type of each binding matches the type of the slot *)
        let bindings' = group_bindings_slots bindings slots in
        let bindings'' = List.map (fun (n, ty, actual, rename) -> (n, parse_typespec menv rm typarams ty, replace_variables typebindings actual, rename)) bindings' in
        let newvars = List.map (fun (_, ty, actual, rename) ->
                          if equal_ty ty actual then
                            let _ = pi ("Binding name", rename)
                            in
                            pt ("Binding type",  ty);
                            (rename, ty, VarLocal)
                          else
                            Errors.slot_wrong_type
                              ~name:rename
                              ~expected:actual
                              ~actual:ty)
                        bindings'' in
        let lexenv' = push_vars lexenv newvars in
        let body' = augment_stmt (update_lexenv ctx lexenv') body in
        TypedWhen (name, List.map (fun (name, _, ty, rename) -> TypedBinding { name; ty; rename}) bindings'', body')
      else
        Errors.case_wrong_slots ())

let rec validate_constant_expression (expr: texpr): unit =
  if is_constant expr then
    ()
  else
    Errors.expression_not_constant ()

and is_constant = function
  | TNilConstant ->
     true
  | TBoolConstant _ ->
     true
  | TIntConstant _ ->
     true
  | TFloatConstant _ ->
     true
  | TStringConstant _ ->
     true
  | TConstVar _ ->
     true
  | TParamVar _ ->
     false
  | TLocalVar _ ->
     false
  | TFunVar _ ->
     true
  | TFuncall _ ->
     false
  | TMethodCall _ ->
     false
  | TVarMethodCall _ ->
     false
  | TFptrCall _ ->
     false
  | TCast _ ->
     true
  | TComparison (_, lhs, rhs) ->
     (is_constant lhs) && (is_constant rhs)
  | TConjunction (lhs, rhs) ->
     (is_constant lhs) && (is_constant rhs)
  | TDisjunction (lhs, rhs) ->
     (is_constant lhs) && (is_constant rhs)
  | TNegation e ->
     is_constant e
  | TIfExpression (c, t, f) ->
     (is_constant c) && (is_constant t) && (is_constant f)
  | TRecordConstructor (_, values) ->
     List.for_all (fun (_, v) -> is_constant v) values
  | TUnionConstructor (_, _, values) ->
     List.for_all (fun (_, v) -> is_constant v) values
  | TPath { head; elems; _ } ->
     (is_constant head) && (List.for_all is_path_elem_constant elems)
  | TEmbed _ ->
     true
  | TDeref _ ->
     false
  | TSizeOf _ ->
     true
  | TBorrowExpr _ ->
     false

and is_path_elem_constant = function
  | TSlotAccessor _ ->
     true
  | TPointerSlotAccessor _ ->
     true
  | TArrayIndex (e, _) ->
     is_constant e

let rec augment_decl (module_name: module_name) (kind: module_kind) (env: env) (decl: linked_definition): typed_decl =
  with_frame "Augment declaration"
    (fun _ ->
      let rm = empty_region_map in
      match decl with
      | LConstant (decl_id, vis, name, ty, expr, doc) ->
         ps ("Kind", "Constant");
         let expr' = augment_expr module_name env rm empty_typarams empty_lexenv (Some ty) expr in
         let _ = match_type_with_value (env, module_name) ty expr' in
         let _ = validate_constant_expression expr' in
         TConstant (decl_id, vis, name, ty, expr', doc)
      | LRecord (decl_id, vis, name, typarams, universe, slots, doc) ->
         ps ("Kind", "Record");
         TRecord (decl_id, vis, name, typarams, universe, slots, doc)
      | LUnion (decl_id, vis, name, typarams, universe, cases, doc) ->
         ps ("Kind", "Union");
         TUnion (decl_id, vis, name, typarams, universe, cases, doc)
      | LFunction (decl_id, vis, name, typarams, params, rt, body, doc, pragmas) ->
         ps ("Kind", "Function");
         pi ("Name", name);
         (match pragmas with
          | [ForeignImportPragma s] ->
             if typarams_size typarams = 0 then
               if kind = UnsafeModule then
                 TForeignFunction (decl_id, vis, name, params, rt, s, doc)
               else
                 Errors.foreign_in_safe_module ()
             else
               Errors.foreign_type_parameters ()
          | [ForeignExportPragma _] ->
             let ctx = StmtCtx (module_name, env, rm, typarams, (lexenv_from_params params), rt) in
             let body' = augment_stmt ctx body in
             TFunction (decl_id, vis, name, typarams, params, rt, body', doc)
          | [] ->
             let ctx = StmtCtx (module_name, env, rm, typarams, (lexenv_from_params params), rt) in
             let body' = augment_stmt ctx body in
             TFunction (decl_id, vis, name, typarams, params, rt, body', doc)
          | _ ->
             Errors.fun_invalid_pragmas ())
      | LTypeclass (decl_id, vis, name, typaram, methods, doc) ->
         ps ("Kind", "Typeclass");
         TTypeClass (decl_id, vis, name, typaram, List.map (augment_method_decl env rm typaram) methods, doc)
      | LInstance (decl_id, vis, name, typarams, arg, methods, doc) ->
         ps ("Kind", "Instance");
         (* TODO: the universe of the type parameter matches the universe of the type argument *)
         (* TODO: Check the methods in the instance match the methods in the class *)
         TInstance (decl_id, vis, name, typarams, arg, List.map (augment_method_def module_name env rm typarams) methods, doc))

and lexenv_from_params (params: value_parameter list): lexenv =
  match params with
  | (ValueParameter (n, t))::rest ->
     push_var (lexenv_from_params rest) n t VarParam
  | [] ->
     empty_lexenv

and augment_method_decl _ _ _ (LMethodDecl (decl_id, name, params, rt, _)) =
  TypedMethodDecl (decl_id, name, params, rt)

and augment_method_def module_name menv rm typarams (LMethodDef (ins_meth_id, name, params, rt, _, body)) =
  let ctx = StmtCtx (module_name, menv, rm, typarams, (lexenv_from_params params), rt) in
  let body' = augment_stmt ctx body in
  TypedMethodDef (ins_meth_id, name, params, rt, body')

let augment_module menv (LinkedModule { name; decls; kind; _ }) =
  with_frame "Typing pass"
    (fun _ ->
      let decls' = List.map (augment_decl name kind menv) decls in
      TypedModule (name, decls'))

 Since the extraction pass has already happened, we can simplify the call to
   `parse_type` by passing an empty list of local type signatures. 
 Function is generic, need an asserted type. 
 Function is concrete. 
 Augment operands. 
 Get the types. 
 Because method calls don't implement the lenient type checking
                rules for arithmetic, we have to cast the arguments to their
                expected types, so that constants are allowed to be used as
                arguments. This means we have to transform a `ty` into a type
                specifier. 
 Are the types the same? 
 If the types are the same type, check it is a numeric type. 
 If either operand is a constant, let it pass 
 Is the initial expression of a path a read reference or write reference or no reference? 
 If the path starts in a reference it should end in one. 
 The type of the expression being dereferenced must be either a read-only
            reference or a write reference. 
 TODO: check if `ty` is linear? 
 TODO: Addresses should not be indexable. 
 Check: e' is a public record type 
 Check: the given slot name must exist in this record type. 
 Check arg is a public record 
 Check: the given slot name must exist in this record type. 
 Check: e' is a public record type 
 Check: the given slot name must exist in this record type. 
 Otherwise, find a callable from the environment. 
 Because function pointers don't preserve value parameter names, we can't
     accept named argument lists. 
 We extract the parameter types and return type from the function pointer
     type. 
 Then, we synthesize a parameter list from the function pointer type. The
     parameter names are fake, but this is so we can reuse the arity-checking,
     value parameter list-checking infrastructure we use elsewhere. 
 Check synthetic parameter list against argument list. 
 Use the bindings to get the effective return type 
 TODO: constructing a fake qident 
 For simplicity, and to reduce duplication of code, we convert the argument
         list to a positional list. 
 Check the list of params against the list of arguments 
 Use the bindings to get the effective return type 
 Check: the set of bindings equals the set of type parameters 
 Check: the argument list must be named 
 Check: the set of slots matches the set of param names 
 Convert args to positional 
 Check the list of params against the list of arguments 
 Use the bindings to get the effective return type 
 Check: the set of bindings equals the set of type parameters 
 Check the resulting type is in the correct universe 
 Check: the set of slots matches the set of param names 
 Convert args to positional 
 Check the list of params against the list of arguments 
 Use the bindings to get the effective return type 
 Check: the set of bindings equals the set of type parameters 
 Check the resulting type is in the correct universe 
 At this point, we know the method's name and the typeclass it belongs
         to. We pull the parameter list from the typeclass, and compare that against
         the argument list. 
 For simplicity, and to reduce duplication of code, we convert the argument
         list to a positional list. 
 Check the list of params against the list of arguments 
 Use the bindings to get the effective return type 
 Check: the set of bindings equals the set of type parameters 
 Is the dispatch type a type variable? 
 If so, check if the constraints say that it implements the typeclass. 
 If it's a tyvar that implements the typeclass, we have to
                  proceed differently. Essentially we can't go on with instance
                  resolution, because we don't have a type to resolve an
                  instance for. So we have to freeze the process. 
 If it doesn't, that's an error 
 If it's not a type variable, continue normal instance resolution. 
 Given a list of type parameters, and a list of arguments, check that the
   lists have the same length and each argument satisfies each corresponding
   parameter. Return the resulting set of type bindings. 
 Check that the number of parameters is the same as the number of
     arguments. 
 Check arguments against parameters 
 Precondition: both lists have the same length. 
 Cast integer types 
 Given a function's parameter list and return type, check if the return type
   of a function is polymorphic. 
 Find the set of type variables in the return type. 
 Find the set of type variables in each of the parameters. 
 Union the type variables of the parameters into one set of all type
     variables in the parameter list. 
 Find the set of type variables that are in the return type set but not in
     the parameter set. 
 Check that there are as many bindings as there are type parameters, and that
   every type parameter is satisfied. 
 I think this should not be an error 
 Check: the value must be a public record type 
 Find the set of slot names and the set of binding names, and compare them 
 Assigning to a variable. 
 Assigning to a path. 
 Group the cases and whens 
 TODO: addresses should not be indexable 
 Check the set of binding names is the same as the set of slots 
 Check the type of each binding matches the type of the slot 
 TODO: the universe of the type parameter matches the universe of the type argument 
 TODO: Check the methods in the instance match the methods in the class 

    Part of the Austral project , under the Apache License v2.0 with LLVM Exceptions . 
    See LICENSE file for details . 

    SPDX - License - Identifier : Apache-2.0 WITH LLVM - exception 

   Part of the Austral project, under the Apache License v2.0 with LLVM Exceptions.
   See LICENSE file for details.

   SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
*)

open Identifier
open Common
open Type
open TypeSystem
open TypeBindings
open TypeMatch
open TypeVarSet
open TypeParser
open TypeParameter
open TypeParameters
open Region
open Id
open LexEnv
open Env
open EnvTypes
open EnvUtils
open EnvExtras
open Ast
open Tast
open TastUtil
open Linked
open Util
open Reporter
open BuiltIn
open Error

module Errors = TypeErrors

let parse_typespec (env: env) (rm: region_map) (typarams: typarams) (ty: qtypespec): ty =
  parse_type env [] rm typarams ty

let rec augment_expr (module_name: module_name) (env: env) (rm: region_map) (typarams: typarams) (lexenv: lexenv) (asserted_ty: ty option) (expr: aexpr): texpr =
  with_frame ("Augment expression: " ^ (expr_kind expr))
    (fun _ ->
      let aug = augment_expr module_name env rm typarams lexenv None in
      match expr with
      | NilConstant ->
         TNilConstant
      | BoolConstant b ->
         TBoolConstant b
      | IntConstant i ->
         TIntConstant i
      | FloatConstant f ->
         TFloatConstant f
      | StringConstant s ->
         TStringConstant s
      | Variable name ->
         (match get_decl_by_name env (qident_to_sident name) with
          | Some (Function { id; typarams; value_params; rt; _ }) ->
             let arg_tys: ty list = List.map (fun (ValueParameter (_, ty)) -> ty) value_params in
             let fn_ty: ty = FnPtr (arg_tys, rt) in
             if (typarams_size typarams > 0) then
               (match asserted_ty with
                | Some asserted_ty ->
                   let bindings: type_bindings = match_type (env, module_name) fn_ty asserted_ty in
                   let rt: ty = replace_variables bindings rt
                   and arg_tys: ty list = List.map (replace_variables bindings) arg_tys in
                   let effective_fn_ty: ty = FnPtr (arg_tys, rt) in
                   TFunVar (id, effective_fn_ty, bindings)
                | None ->
                   Errors.unconstrained_generic_function (original_name name))
             else
               TFunVar (id, fn_ty, empty_bindings)
          | _ ->
             (match get_variable env lexenv name with
              | Some (ty, src) ->
                 (match src with
                  | VarConstant ->
                     TConstVar (name, ty)
                  | VarParam ->
                     TParamVar ((original_name name), ty)
                  | VarLocal ->
                     TLocalVar ((original_name name), ty))
              | None ->
                 Errors.unknown_name
                   ~kind:"variable"
                   ~name:(original_name name)))
      | FunctionCall (name, args) ->
         augment_call module_name env lexenv asserted_ty name (augment_arglist module_name env rm typarams lexenv None args)
      | ArithmeticExpression (op, lhs, rhs) ->
         let lhs' = aug lhs
         and rhs' = aug rhs in
         let lhs_ty = get_type lhs'
         and rhs_ty = get_type rhs'
         in
         let augment_arithmetic_call _ =
           let op_name =
             match op with
             | Add ->
                "trappingAdd"
             | Subtract ->
                "trappingSubtract"
             | Multiply ->
                "trappingMultiply"
             | Divide ->
                "trappingDivide"
           in
           let op_qname = make_qident (pervasive_module_name, make_ident op_name, make_ident op_name) in
           let args: typed_arglist = TPositionalArglist [
                          TCast (lhs', lhs_ty);
                          TCast (rhs', lhs_ty);
                        ]
           in
           augment_call module_name env lexenv asserted_ty op_qname args
         in
         if lhs_ty = rhs_ty then
           if (is_numeric lhs_ty) then
             augment_arithmetic_call ()
           else
             Errors.arithmetic_not_numeric lhs_ty
         else
  If the types are different , check if at least one operator is a constant .
           let are_int_constants = (is_int_constant lhs) || (is_int_constant rhs)
           and are_float_constants = (is_float_constant lhs) || (is_float_constant rhs) in
           if (are_int_constants || are_float_constants) then
             augment_arithmetic_call ()
           else
             Errors.arithmetic_incompatible_types ~lhs:lhs_ty ~rhs:rhs_ty
      | Comparison (op, lhs, rhs) ->
         let lhs' = aug lhs
         and rhs' = aug rhs in
         let _ = match_type_with_value (env, module_name) (get_type lhs') rhs' in
         TComparison (op, lhs', rhs')
      | Conjunction (lhs, rhs) ->
         let lhs' = aug lhs
         and rhs' = aug rhs in
         if ((is_bool lhs') && (is_bool rhs')) then
           TConjunction (lhs', rhs')
         else
           Errors.logical_operands_not_boolean
             ~operator:"conjunction"
             ~types:[get_type lhs'; get_type rhs']
      | Disjunction (lhs, rhs) ->
         let lhs' = aug lhs
         and rhs' = aug rhs in
         if ((is_bool lhs') && (is_bool rhs')) then
           TDisjunction (lhs', rhs')
         else
           Errors.logical_operands_not_boolean
             ~operator:"disjunction"
             ~types:[get_type lhs'; get_type rhs']
      | Negation e ->
         let e' = aug e in
         if is_bool e' then
           TNegation e'
         else
           Errors.logical_operands_not_boolean
             ~operator:"negation"
             ~types:[get_type e']
      | IfExpression (c, t, f) ->
         let c' = aug c
         and t' = aug t
         and f' = aug f in
         if is_bool c' then
           if (get_type t') = (get_type f') then
             TIfExpression (c', t', f')
           else
             Errors.if_inequal
               ~lhs:(get_type t')
               ~rhs:(get_type f')
         else
           Errors.condition_not_boolean
             ~kind:"if"
             ~form:"expression"
             ~ty:(get_type c')
      | Path (e, elems) ->
  Path rules : 

             1 . Path that begins in a reference ends in a reference . 
             2 . All paths end in a type in the free universe . 
          

            1. Path that begins in a reference ends in a reference.
            2. All paths end in a type in the free universe.
          *)
         let e' = aug e in
         let (is_read, region): (bool * ty option) =
           (match (get_type e') with
            | ReadRef (_, r) ->
               (true, Some r)
            | WriteRef (_, r) ->
               (false, Some r)
            | _ ->
               (false, None))
         in
         let elems' = augment_path env module_name rm typarams lexenv (get_type e') elems in
         let path_ty: ty = get_path_ty_from_elems elems' in
         let path_ty: ty =
           (match region with
            | Some reg ->
               if is_read then
                 ReadRef (path_ty, reg)
               else
                 WriteRef (path_ty, reg)
            | None ->
               path_ty)
         in
         let universe = type_universe path_ty in
         if universe = FreeUniverse then
           TPath {
               head = e';
               elems = elems';
               ty = path_ty
             }
         else
           Errors.path_not_free path_ty
      | Embed (ty, expr, args) ->
         TEmbed (
             parse_typespec env rm typarams ty,
             expr,
             List.map aug args
           )
      | Deref expr ->
         let expr' = aug expr in
         let ty = get_type expr' in
         (match ty with
          | ReadRef _ ->
             TDeref expr'
          | WriteRef _ ->
             TDeref expr'
          | _ ->
             Errors.dereference_non_reference ty)
      | Typecast (expr, ty) ->
  The typecast operator has four uses : 

             1 . Clarifying the type of integer and floating point literals . 

             2 . Converting write references to read references . 

             3 . Clarifying the type of return type polymorphic functions . 

          

            1. Clarifying the type of integer and floating point literals.

            2. Converting write references to read references.

            3. Clarifying the type of return type polymorphic functions.

          *)
         let is_int_expr = function
           | TIntConstant _ -> true
           | _ -> false
         and is_float_expr = function
           | TFloatConstant _ -> true
           | _ -> false

         and is_int_type = function
           | Integer _ -> true
           | _ -> false

         and is_float_type = function
           | SingleFloat -> true
           | DoubleFloat -> true
           | _ -> false
         in
         let target_type = parse_typespec env rm typarams ty in
  By passing target_type as the asserted type we 're doing point 3 .
         let expr' = augment_expr module_name env rm typarams lexenv (Some target_type) expr in
  For 1 : if the expression is an int literal and the target type is an 
             int type , do the conversion . Similarly , if the expression is a float 
             constant and the target type is a float , do the conversion .
            int type, do the conversion. Similarly, if the expression is a float
            constant and the target type is a float, do the conversion. *)
         if (is_int_expr expr') && (is_int_type target_type) then
           TCast (expr', target_type)
         else
           if (is_float_expr expr') && (is_float_type target_type) then
             TCast (expr', target_type)
           else
  Otherwise , if the expression has a mutable reference type and the 
                 target type is a read reference , and they point to the same underlying 
                 type and underlying region , just convert it to a read ref . This does 
                 point 2 .
                target type is a read reference, and they point to the same underlying
                type and underlying region, just convert it to a read ref. This does
                point 2. *)
             (match (get_type expr') with
              | WriteRef (underlying_ty, region) ->
                 (match target_type with
                  | ReadRef (underlying_ty', region') ->
                     if ((equal_ty underlying_ty underlying_ty') && (equal_ty region region')) then
                       TCast (expr', target_type)
                     else
                       Errors.cast_different_references
                         ~different_types:(not (equal_ty underlying_ty underlying_ty'))
                         ~different_regions:(not (equal_ty region region'))
                  | _ ->
                     Errors.cast_write_ref_to_non_ref ())
              | _ ->
                 (try
                    let _ = match_type (env, module_name) target_type (get_type expr') in
                    expr'
                  with Austral_error _ ->
                    Errors.cast_invalid
                      ~target:target_type
                      ~source:(get_type expr')))
      | SizeOf ty ->
         TSizeOf (parse_typespec env rm typarams ty)
      | BorrowExpr (mode, name) ->
         (match get_variable env lexenv name with
          | Some (ty, src) ->
             (match src with
              | VarConstant ->
                 Errors.borrow_constant ()
              | VarParam ->
                 let name: identifier = original_name name
                 and reg: region = fresh_region ()
                 in
                 TBorrowExpr (mode, name, reg, ty)
              | VarLocal ->
                 let name: identifier = original_name name
                 and reg: region = fresh_region ()
                 in
                 TBorrowExpr (mode, name, reg, ty))
          | None ->
             Errors.unknown_name
               ~kind:"variable"
               ~name:(original_name name)))

and get_path_ty_from_elems (elems: typed_path_elem list): ty =
  assert ((List.length elems) > 0);
  let last = List.nth elems ((List.length elems) - 1) in
  path_elem_type last

and is_bool e =
  match get_type e with
  | Boolean -> true
  | _ -> false

and is_int_constant e =
  match e with
  | IntConstant _ -> true
  | _ -> false

and is_float_constant e =
  match e with
  | FloatConstant _ -> true
  | _ -> false

and augment_arglist (module_name: module_name) (env: env) (rm: region_map) (typarams: typarams) (lexenv: lexenv) (asserted_ty: ty option) (args: abstract_arglist): typed_arglist =
  let aug = augment_expr module_name env rm typarams lexenv asserted_ty in
  match args with
  | Positional args' ->
     TPositionalArglist (List.map aug args')
  | Named pairs ->
     TNamedArglist (List.map (fun (n, v) -> (n, aug v)) pairs)

and augment_path (env: env) (module_name: module_name) (rm: region_map) (typarams: typarams) (lexenv: lexenv) (head_ty: ty) (elems: path_elem list): typed_path_elem list =
  match elems with
  | [elem] ->
     [augment_path_elem env module_name rm typarams lexenv head_ty elem]
  | elem::rest ->
     let elem' = augment_path_elem env module_name rm typarams lexenv head_ty elem in
     let rest' = augment_path env module_name rm typarams lexenv (path_elem_type elem') rest in
     elem' :: rest'
  | [] ->
     err "Path is empty"

and augment_path_elem (env: env) (module_name: module_name) (rm: region_map) (typarams: typarams) (lexenv: lexenv) (head_ty: ty) (elem: path_elem): typed_path_elem =
  match elem with
  | SlotAccessor slot_name ->
     (match head_ty with
      | NamedType (name, args, _) ->
         augment_slot_accessor_elem env module_name slot_name name args
      | _ ->
         Errors.path_not_record (type_string head_ty))
  | PointerSlotAccessor slot_name ->
     (match head_ty with
      | Pointer pointed_to ->
         augment_pointer_slot_accessor_elem env module_name slot_name pointed_to
      | ReadRef (ty, _) ->
         (match ty with
          | NamedType (name, args, _) ->
             augment_reference_slot_accessor_elem env module_name slot_name name args
          | _ ->
             Errors.path_not_record (type_string ty))
      | WriteRef (ty, _) ->
         (match ty with
          | NamedType (name, args, _) ->
             augment_reference_slot_accessor_elem env module_name slot_name name args
          | _ ->
             Errors.path_not_record (type_string ty))
      | _ ->
         Errors.path_not_record (type_string head_ty))
  | ArrayIndex ie ->
     let ie' = augment_expr module_name env rm typarams lexenv None ie in
     (match head_ty with
      | StaticArray elem_ty ->
         TArrayIndex (ie', elem_ty)
      | _ ->
         Errors.array_indexing_disallowed head_ty)

and augment_slot_accessor_elem (env: env) (module_name: module_name) (slot_name: identifier) (type_name: qident) (type_args: ty list) =
  let (source_module, vis, typarams, slots) = get_record_definition env type_name in
  if (vis = TypeVisPublic) || (module_name = source_module) then
    let (TypedSlot (_, slot_ty)) = get_slot_with_name type_name slots slot_name in
    let bindings = match_typarams (env, module_name) typarams type_args in
    let slot_ty' = replace_variables bindings slot_ty in
    TSlotAccessor (slot_name, slot_ty')
  else
    Errors.path_not_public
      ~type_name:(original_name type_name)
      ~slot_name

and augment_pointer_slot_accessor_elem (env: env) (module_name: module_name) (slot_name: identifier) (pointed_to: ty): typed_path_elem =
  match pointed_to with
  | NamedType (type_name, type_args, _) ->
     let (source_module, vis, typarams, slots) = get_record_definition env type_name in
     if (vis = TypeVisPublic) || (module_name = source_module) then
       let (TypedSlot (_, slot_ty)) = get_slot_with_name type_name slots slot_name in
       let bindings = match_typarams (env, module_name) typarams type_args in
       let slot_ty' = replace_variables bindings slot_ty in
       TPointerSlotAccessor (slot_name, slot_ty')
     else
       Errors.path_not_public
         ~type_name:(original_name type_name)
         ~slot_name
  | _ ->
     Errors.path_not_record (type_string pointed_to)

and augment_reference_slot_accessor_elem (env: env) (module_name: module_name) (slot_name: identifier) (type_name: qident) (type_args: ty list) =
  let (source_module, vis, typarams, slots) = get_record_definition env type_name in
  if (vis = TypeVisPublic) || (module_name = source_module) then
    let (TypedSlot (_, slot_ty)) = get_slot_with_name type_name slots slot_name in
    let bindings = match_typarams (env, module_name) typarams type_args in
    let slot_ty' = replace_variables bindings slot_ty in
    TPointerSlotAccessor (slot_name, slot_ty')
  else
    Errors.path_not_public
      ~type_name:(original_name type_name)
      ~slot_name

and get_record_definition (env: env) (name: qident): (module_name * type_vis * typarams * typed_slot list) =
  match get_decl_by_name env (qident_to_sident name) with
  | (Some (Record { mod_id; vis; typarams; slots; _ })) ->
     let mod_name: module_name = module_name_from_id env mod_id in
     (mod_name, vis, typarams, slots)
  | Some _ ->
     Errors.path_not_record (original_name name |> ident_string)
  | None ->
     err ("No record with this name: " ^ (ident_string (original_name name)))

and get_slot_with_name type_name slots slot_name =
  match List.find_opt (fun (TypedSlot (n, _)) -> n = slot_name) slots with
  | Some s -> s
  | None -> 
     Errors.no_such_slot
       ~type_name:(original_name type_name)
       ~slot_name

and augment_call (module_name: module_name) (env: env) (lexenv: lexenv) (asserted_ty: ty option) (name: qident) (args: typed_arglist): texpr =
  First , check if the callable name is a variable .
  match get_var lexenv (local_name name) with
  | Some (ty, _) ->
     augment_fptr_call module_name env (local_name name) ty asserted_ty args
  | None ->
     (match get_callable env module_name (qident_to_sident name) with
      | Some callable ->
         augment_callable module_name env name callable asserted_ty args
      | None ->
         Errors.unknown_name
           ~kind:"callable"
           ~name:(original_name name))

and augment_fptr_call (module_name: module_name) (env: env) (name: identifier) (fn_ptr_ty: ty) (asserted_ty: ty option) (args: typed_arglist): texpr =
  let args: texpr list =
    match args with
    | TPositionalArglist args -> args
    | TNamedArglist _ ->
       Errors.fun_pointer_named_args ()
  in
  let (paramtys, rt): (ty list * ty) =
    match fn_ptr_ty with
    | FnPtr (paramtys, rt) -> (paramtys, rt)
    | _ -> Errors.call_non_callable fn_ptr_ty
  in
  let params: value_parameter list = List.map (fun paramty -> ValueParameter (make_ident "fake", paramty)) paramtys in
  let bindings = check_argument_list env module_name name params args in
  let rt' = replace_variables bindings rt in
  let (bindings', rt'') = handle_return_type_polymorphism env module_name name params rt' asserted_ty bindings in
  let arguments = cast_arguments bindings' params args in
  TFptrCall (name, arguments, rt'')

and augment_callable (module_name: module_name) (env: env) (name: qident) (callable: callable) (asserted_ty: ty option) (args: typed_arglist) =
  with_frame "Augment callable"
    (fun _ ->
      match callable with
      | FunctionCallable (id, typarams, params, rt) ->
         augment_function_call env module_name id name typarams params rt asserted_ty args
      | RecordConstructor (_, typarams, universe, slots) ->
         augment_record_constructor env module_name name typarams universe slots asserted_ty args
      | UnionConstructor { union_id; type_params; universe; case } ->
         let type_name: qident = (match get_decl_by_id env union_id with
                                  | Some (Union { name; mod_id; _ }) ->
                                     make_qident (module_name_from_id env mod_id, name, name)
                                  | _ ->
                                     internal_err ("no type name found for ID `" ^ (show_decl_id union_id) ^ "`"))
         in
         augment_union_constructor env module_name type_name type_params universe case asserted_ty args
      | MethodCallable { typeclass_id; value_parameters; return_type; _ } ->
         let param = (match get_decl_by_id env typeclass_id with
                      | Some (TypeClass { param; _ }) ->
                         param
                      | _ ->
                         internal_err ("no type parameter found for ID `" ^ (show_decl_id typeclass_id) ^ "`"))
         in
         augment_method_call env module_name typeclass_id param name value_parameters return_type asserted_ty args)

and augment_function_call (env: env) (module_name: module_name) (id: decl_id) name typarams params rt asserted_ty args =
  with_frame "Augment function call"
    (fun _ ->
      pqi ("Name", name);
      let param_names = List.map (fun (ValueParameter (n, _)) -> n) params in
      let arguments: texpr list = arglist_to_positional (args, param_names) in
      let bindings = check_argument_list env module_name (original_name name) params arguments in
      let rt' = replace_variables bindings rt in
      let (bindings', rt'') = handle_return_type_polymorphism env module_name (local_name name) params rt' asserted_ty bindings in
      let bindings'' = merge_bindings bindings bindings' in
      check_bindings typarams bindings'';
      let arguments' = cast_arguments bindings'' params arguments in
      let substs = make_substs bindings'' typarams in
      ps ("Bindings", show_bindings bindings'');
      pt ("Return type", rt'');
      TFuncall (id, name, arguments', rt'', substs))

and augment_record_constructor (env: env) (module_name: module_name) (name: qident) (typarams: typarams) (universe: universe) (slots: typed_slot list) (asserted_ty: ty option) (args: typed_arglist) =
  let args' = (match args with
               | TPositionalArglist l ->
                  if l = [] then
                    []
                  else
                    Errors.constructor_not_named "record"
               | TNamedArglist a ->
                  a) in
  let slot_names: identifier list = List.map (fun (TypedSlot (name, _)) -> name) slots in
  let argument_names: identifier list = List.map (fun (n, _) -> n) args' in
  if not (ident_set_eq slot_names argument_names) then
    Errors.constructor_wrong_args "record"
  else
    let arguments = arglist_to_positional (args, slot_names) in
    let params = List.map (fun (TypedSlot (n, t)) -> ValueParameter (n, t)) slots in
    let bindings = check_argument_list env module_name (original_name name) params arguments in
    let rt = NamedType (name,
                        List.map (fun tp -> TyVar (typaram_to_tyvar tp)) (typarams_as_list typarams),
                        universe)
    in
    let rt' = replace_variables bindings rt in
    let (bindings', rt'') = handle_return_type_polymorphism env module_name (local_name name) params rt' asserted_ty bindings in
    check_bindings typarams (merge_bindings bindings bindings');
    if universe_compatible universe (type_universe rt'') then
      TRecordConstructor (rt'', List.map2 (fun a b -> (a, b)) slot_names arguments)
    else
      err "Universe mismatch"

and augment_union_constructor (env: env) (module_name: module_name) (type_name: qident) (typarams: typarams) (universe: universe) (case: typed_case) (asserted_ty: ty option) (args: typed_arglist) =
  with_frame "Augment union constructor"
    (fun _ ->
      pqi ("Type name", type_name);
      let args' = (match args with
                   | TPositionalArglist l ->
                      if l = [] then
                        []
                      else
                        Errors.constructor_not_named "union"
                   | TNamedArglist a ->
                      a) in
      let (TypedCase (case_name, slots)) = case in
      pi ("Case name", case_name);
      let slot_names: identifier list = List.map (fun (TypedSlot (name, _)) -> name) slots in
      let argument_names: identifier list = List.map (fun (n, _) -> n) args' in
      if not (ident_set_eq slot_names argument_names) then
        Errors.constructor_wrong_args "union"
      else
        let arguments = arglist_to_positional (args, slot_names) in
        let params = List.map (fun (TypedSlot (n, t)) -> ValueParameter (n, t)) slots in
        ps ("Params", String.concat "\n" (List.map (fun (ValueParameter (n, t)) -> (ident_string n) ^ ": " ^ (type_string t)) params));
        ps ("Arguments", String.concat "\n" (List.map show_texpr arguments));
        let bindings = check_argument_list env module_name case_name params arguments in
        let rt = NamedType (type_name,
                            List.map (fun tp -> TyVar (typaram_to_tyvar tp)) (typarams_as_list typarams),
                            universe)
        in
        pt ("Type", rt);
        ps ("Bindings", show_bindings bindings);
        let rt' = replace_variables bindings rt in
        pt ("Type'", rt');
        let (bindings', rt'') = handle_return_type_polymorphism env module_name case_name params rt' asserted_ty bindings in
        check_bindings typarams (merge_bindings bindings bindings');
        if universe_compatible universe (type_universe rt'') then
          let _ = pt ("Type''", rt'')
          in
          TUnionConstructor (rt'', case_name, List.map2 (fun a b -> (a, b)) slot_names arguments)
        else
          err "Universe mismatch")

and augment_method_call (env: env) (source_module_name: module_name) (typeclass_id: decl_id) (typaram: type_parameter) (callable_name: qident) (params: value_parameter list) (rt: ty) (asserted_ty: ty option) (args: typed_arglist): texpr =
  with_frame "Augmenting method call"
    (fun _ ->
      pqi ("Callable name", callable_name);
      let param_names = List.map (fun (ValueParameter (n, _)) -> n) params in
      let arguments: texpr list = arglist_to_positional (args, param_names) in
      let bindings = check_argument_list env source_module_name (original_name callable_name) params arguments in
      let rt' = replace_variables bindings rt in
      let (bindings', rt'') = handle_return_type_polymorphism env source_module_name (local_name callable_name) params rt' asserted_ty bindings in
      let bindings'' = merge_bindings bindings bindings' in
      ps ("Bindings", show_bindings bindings'');
      check_bindings (typarams_from_list [typaram]) bindings'';
      match get_binding bindings'' typaram with
      | (Some dispatch_ty) ->
         pt ("Dispatch Type", dispatch_ty);
         (match dispatch_ty with
          | TyVar (TypeVariable (typaram, _, _, constraints)) ->
             let tc_name: sident = get_decl_sident_or_die env typeclass_id in
             if List.exists (fun c -> equal_sident tc_name c) constraints then
               TVarMethodCall {
                   source_module_name = source_module_name;
                   typeclass_id = typeclass_id;
                   params = params;
                   method_name = callable_name;
                   args = arguments;
                   dispatch_ty = dispatch_ty;
                   rt = rt'';
                   bindings = bindings'';
                 }
             else
               Errors.unconstrained_type_parameter
                 ~typeclass:(sident_name tc_name)
                 ~typaram
          | _ ->
             let (instance, instance_bindings): decl * type_bindings =
               (match get_instance env source_module_name dispatch_ty typeclass_id with
                | Some (i, b) -> (i, b)
                | None ->
                   let tc_name: sident = get_decl_sident_or_die env typeclass_id in
                   Errors.no_suitable_instance
                     ~typeclass:(sident_name tc_name)
                     ~ty:dispatch_ty)
             in
             ps ("Instance bindings", show_bindings instance_bindings);
             let params' = List.map (fun (ValueParameter (n, t)) -> ValueParameter (n, replace_variables instance_bindings t)) params in
             let arguments' = cast_arguments instance_bindings params' arguments in
             let typarams = (match instance with
                             | Instance { typarams; _ } -> typarams
                             | _ -> internal_err "Couldn't get instance while getting type parameters")
             in
             let instance_id: decl_id = decl_id instance in
             let meth_id: ins_meth_id =
               (match get_instance_method_from_instance_id_and_method_name env instance_id (original_name callable_name) with
                | Some (InsMethRec { id; _ }) -> id
                | None -> internal_err ("Couldn't get instance method for `"
                                        ^ (ident_string (original_name callable_name))
                                        ^ "`"))
             in
             ps ("Bindings", show_bindings bindings'');
             let substs = make_substs instance_bindings typarams in
             TMethodCall (meth_id, callable_name, typarams, arguments', rt'', substs))
      | None ->
         internal_err "couldn't extract dispatch type.")

and check_argument_list (env: env) (module_name: module_name) (name: identifier) (params: value_parameter list) (args: texpr list): type_bindings =
  check_arity name params args;
  check_argument_list' env module_name empty_bindings params args

and check_arity (name: identifier) (params: value_parameter list) (args: texpr list): unit =
  let nparams = List.length params
  and nargs = List.length args in
  if nparams = nargs then
    ()
  else
    Errors.call_wrong_arity ~name ~expected:nparams ~actual:nargs

and check_argument_list' (env: env) (module_name: module_name) (bindings: type_bindings) (params: value_parameter list) (args: texpr list): type_bindings =
  match (params, args) with
  | ((first_param::rest_params), (first_arg::rest_args)) ->
     let bindings' = merge_bindings bindings (match_parameter env module_name first_param first_arg) in
     check_argument_list' env module_name bindings' rest_params rest_args
  | ([], []) ->
     bindings
  | _ ->
     internal_err ("couldn't check argument list in module `"
                   ^ (mod_name_string module_name)
                   ^ "` with bindings: "
                   ^ (show_type_bindings bindings)
                   ^ ",\n params: `{ "
                   ^ (String.concat ", " (List.map show_value_parameter params))
                   ^ " }`,\n and args: `{ "
                   ^ (String.concat ", " (List.map show_texpr args))
                   ^ "}`")

and match_parameter (env: env) (module_name: module_name) (param: value_parameter) (arg: texpr): type_bindings =
  let (ValueParameter (_, ty)) = param in
  match_type_with_value (env, module_name) ty arg

and cast_arguments (bindings: type_bindings) (params: value_parameter list) (arguments: texpr list): texpr list =
  let f (ValueParameter (_, expected)) value =
    let expected' = replace_variables bindings expected in
    match expected' with
    | Integer _ ->
       TCast (value, expected')
    | _ ->
       value
  in
  List.map2 f params arguments

and make_substs (bindings: type_bindings) (typarams: typarams): type_bindings =
  let f (tp: type_parameter): (type_parameter * ty) option =
    if (typaram_universe tp) = RegionUniverse then
      Some (tp, RegionTy static_region)
    else
      match get_binding bindings tp with
      | Some ty ->
         Some (tp, ty)
      | None ->
         None
  in
  bindings_from_list (List.filter_map f (typarams_as_list typarams))

and handle_return_type_polymorphism (env: env) (module_name: module_name) (name: identifier) (params: value_parameter list) (rt: ty) (asserted_ty: ty option) (bindings: type_bindings): (type_bindings * ty) =
  if is_return_type_polymorphic params rt then
    match asserted_ty with
    | (Some asserted_ty') ->
       let bindings' = match_type (env, module_name) rt asserted_ty' in
       let bindings'' = merge_bindings bindings bindings' in
       (bindings'', replace_variables bindings'' rt)
    | None ->
       Errors.unconstrained_generic_function name
  else
    (empty_bindings, replace_variables bindings rt)

and is_return_type_polymorphic (params: value_parameter list) (rt: ty): bool =
  let rt_type_vars: TypeVarSet.t = type_variables rt in
  let param_type_vars: TypeVarSet.t list = List.map (fun (ValueParameter (_, ty)) -> type_variables ty) params in
  let param_type_vars: TypeVarSet.t = List.fold_left TypeVarSet.union TypeVarSet.empty param_type_vars in
  let vars_not_params: TypeVarSet.t =
    let param_ty_vars: type_var list = TypeVarSet.elements param_type_vars
    in
    let var_name_in_typarams (n: identifier): bool =
      List.exists (fun (TypeVariable (n', _, _, _)) -> equal_identifier n n') param_ty_vars
    in
    TypeVarSet.of_list (List.filter (fun (TypeVariable (n, _, _, _)) -> not (var_name_in_typarams n)) (TypeVarSet.elements rt_type_vars))
  in
  (TypeVarSet.cardinal vars_not_params) > 0

and check_bindings (typarams: typarams) (bindings: type_bindings): unit =
  if (typarams_size typarams) = (binding_count bindings) then
    let check (tp: type_parameter): unit =
      let n = typaram_name tp
      and u = typaram_universe tp
      in
      (match get_binding bindings tp with
       | Some ty ->
          if universe_compatible u (type_universe ty) then
            ()
          else
            err ("Mismatched universes: expected "
                 ^ (show_universe u)
                 ^ " and got "
                 ^ (show_universe (type_universe ty)))
       | None ->
          err ("No binding for this parameter: " ^ (ident_string n)))
    in
    let _ = List.map check (typarams_as_list typarams) in
    ()
  else
err ( " Not the same number of bindings and parameters . Bindings : " 
       ^ ( show_bindings bindings ) 
       ^ " . Parameters : " 
       ^ ( String.concat " , " ( List.map ( fun ( TypeParameter ( n , u ) ) - > ( ident_string n ) ^ " : " ^ ( universe_string u ) ) typarams ) ) )
      ^ (show_bindings bindings)
      ^ ". Parameters: "
      ^ (String.concat ", " (List.map (fun (TypeParameter (n, u)) -> (ident_string n) ^ " : " ^ (universe_string u)) typarams)))*)
    ()

let is_boolean = function
  | Boolean -> true
  | _ -> false

let is_compatible_with_index_type = function
  | TIntConstant _ ->
     true
  | e ->
     (get_type e) = index_type

type stmt_ctx = StmtCtx of module_name * env * region_map * typarams * lexenv * ty

let update_lexenv (StmtCtx (mn, menv, rm, typarams, _, rt)) (lexenv: lexenv): stmt_ctx =
  StmtCtx (mn, menv, rm, typarams, lexenv, rt)

let update_rm (StmtCtx (mn, menv, _, typarams, lexenv, rt)) (rm: region_map): stmt_ctx =
  StmtCtx (mn, menv, rm, typarams, lexenv, rt)

let stmt_kind (stmt: astmt): string =
  match stmt with
  | ASkip _ -> "skip"
  | ALet _ -> "let"
  | ADestructure _ -> "let (destructure)"
  | AAssign _ -> "assign"
  | AIf _ -> "if"
  | AWhen _ -> "if (no else)"
  | ACase _ -> "case"
  | AWhile _ -> "while"
  | AFor _ -> "for"
  | ABorrow _ -> "borrow"
  | ABlock _ -> "block"
  | ADiscarding _ -> "discard"
  | AReturn _ -> "return"

let rec augment_stmt (ctx: stmt_ctx) (stmt: astmt): tstmt =
  with_frame ("Augment statement: " ^ (stmt_kind stmt))
    (fun _ ->
      let (StmtCtx (module_name, env, rm, typarams, lexenv, rt)) = ctx in
      match stmt with
      | ASkip span ->
         TSkip span
      | ALet (span, name, ty, value, body) ->
         pi ("Name", name);
         adorn_error_with_span span
           (fun _ ->
             let expected_ty = parse_typespec env rm typarams ty in
             pt ("Expected type", expected_ty);
             let value' = augment_expr module_name env rm typarams lexenv (Some expected_ty) value in
             let bindings = match_type_with_value (env, module_name) expected_ty value' in
             ps ("Bindings", show_bindings bindings);
             let ty = replace_variables bindings expected_ty in
             pt ("Type", ty);
             pt ("Value type", get_type value');
             let lexenv' = push_var lexenv name ty VarLocal in
             let body' = augment_stmt (update_lexenv ctx lexenv') body in
             TLet (span, name, ty, value', body'))
      | ADestructure (span, bindings, value, body) ->
         adorn_error_with_span span
           (fun _ ->
             let value' = augment_expr module_name env rm typarams lexenv None value in
             let rec_ty = get_type value' in
             (match rec_ty with
              | (NamedType (name, _, u)) ->
                 let (source_module, vis, record_typarams, slots) = get_record_definition env name in
                 let orig_type = NamedType (
                                     make_qident (source_module, original_name name, original_name name),
                                     List.map (fun tp -> TyVar (typaram_to_tyvar tp)) (typarams_as_list record_typarams),
                                     u
                                   )
                 in
                 let typebindings = match_type (env, module_name) orig_type rec_ty in
                 if (vis = TypeVisPublic) || (module_name = source_module) then
                   let binding_names = List.map (fun (QBinding { name; _}) -> name) bindings
                   and slot_names = List.map (fun (TypedSlot (n, _)) -> n) slots in
                   if ident_set_eq binding_names slot_names then
                     let bindings': (identifier * qtypespec * ty * identifier) list =
                       group_bindings_slots bindings slots
                     in
                     let bindings'': (identifier * ty * ty * identifier) list
                       = List.map (fun (n, ty, actual, rename) -> (n, parse_typespec env rm typarams ty, replace_variables typebindings actual, rename)) bindings'
                     in
                     let newvars: (identifier * ty * var_source) list =
                       List.map (fun (_, ty, actual, rename) ->
                           let _ = match_type (env, module_name) ty actual in
                           (rename, ty, VarLocal))
                         bindings''
                     in
                     let lexenv' = push_vars lexenv newvars in
                     let body' = augment_stmt (update_lexenv ctx lexenv') body in
                     TDestructure (
                         span,
                         List.map (fun (name, _, ty, rename) -> TypedBinding { name; ty; rename; }) bindings'',
                         value',
                         body'
                       )
                   else
                     Errors.destructure_wrong_slots rec_ty
                 else
                   Errors.destructure_not_public rec_ty
              | _ ->
                 Errors.destructure_non_record rec_ty))
      | AAssign (span, LValue (var, elems), value) ->
         adorn_error_with_span span
           (fun _ ->
             (match get_var lexenv var with
              | Some (var_ty, _) ->
                 (match elems with
                  | [] ->
                     let value = augment_expr module_name env rm typarams lexenv None value in
                     let universe = type_universe (get_type value) in
                     if universe = FreeUniverse then
                       let _ = match_type_with_value (env, module_name) var_ty value in
                       TAssign (span, TypedLValue (var, []), value)
                     else
                       Errors.lvalue_not_free ()
                  | elems ->
                     let elems = augment_lvalue_path env module_name rm typarams lexenv var_ty elems in
                     let value = augment_expr module_name env rm typarams lexenv None value in
                     let ty = get_path_ty_from_elems elems in
                     let _ = match_type_with_value (env, module_name) ty value in
                     let path = TPath {
                                    head = value;
                                    elems = elems;
                                    ty = ty
                                  }
                     in
                     let universe = type_universe (get_type path) in
                     if universe = FreeUniverse then
                       TAssign (span, TypedLValue (var, elems), value)
                     else
                       Errors.lvalue_not_free ())
              | None ->
                 Errors.unknown_name ~kind:"variable" ~name:var))
      | AIf (span, c, t, f) ->
         adorn_error_with_span span
           (fun _ ->
             let c' = augment_expr module_name env rm typarams lexenv None c in
             if is_boolean (get_type c') then
               TIf (span, c', augment_stmt ctx t, augment_stmt ctx f)
             else
               Errors.condition_not_boolean
                 ~kind:"if"
                 ~form:"statement"
                 ~ty:(get_type c'))
      | AWhen (span, c, t) ->
         adorn_error_with_span span
           (fun _ ->
             let c = augment_expr module_name env rm typarams lexenv None c in
             if is_boolean (get_type c) then
               TIf (span, c, augment_stmt ctx t, TSkip span)
             else
               Errors.condition_not_boolean
                 ~kind:"if"
                 ~form:"statement"
                 ~ty:(get_type c))
      | ACase (span, expr, whens) ->
  Type checking a case statement : 

             1 . Ensure the value is of a union type . 
             2 . Ensure the union type is public or it is defined in this module . 
             3 . Ensure the set of case names in the case statement equals the set of cases in the union definition . 
             4 . Iterate over the cases , and ensure the bindings are correct . 
          

            1. Ensure the value is of a union type.
            2. Ensure the union type is public or it is defined in this module.
            3. Ensure the set of case names in the case statement equals the set of cases in the union definition.
            4. Iterate over the cases, and ensure the bindings are correct.
          *)
         adorn_error_with_span span
           (fun _ ->
             let expr' = augment_expr module_name env rm typarams lexenv None expr in
             let ty = get_type expr' in
             pt ("Case type", ty);
             let (union_ty, cases) = get_union_type_definition module_name env (get_type expr') in
             let typebindings = match_type (env, module_name) union_ty ty in
             let case_names = List.map (fun (TypedCase (n, _)) -> n) (List.map union_case_to_typed_case cases) in
             let when_names = List.map (fun (AbstractWhen (n, _, _)) -> n) whens in
             if ident_set_eq case_names when_names then
               let whens' = group_cases_whens cases whens in
               let whens'' = List.map (fun (c, w) -> augment_when ctx typebindings w c) whens' in
               TCase (span, expr', whens'')
             else
               Errors.case_non_exhaustive ())
      | AWhile (span, c, body) ->
         adorn_error_with_span span
           (fun _ ->
             let c' = augment_expr module_name env rm typarams lexenv None c in
             if is_boolean (get_type c') then
               TWhile (span, c', augment_stmt ctx body)
             else
               Errors.condition_not_boolean
                 ~kind:"while"
                 ~form:"statement"
                 ~ty:(get_type c'))
      | AFor { span; name; initial; final; body; } ->
         adorn_error_with_span span
           (fun _ ->
             let i' = augment_expr module_name env rm typarams lexenv None initial
             and f' = augment_expr module_name env rm typarams lexenv None final in
             if is_compatible_with_index_type i' then
               if is_compatible_with_index_type f' then
                 let lexenv' = push_var lexenv name index_type VarLocal in
                 let b' = augment_stmt (update_lexenv ctx lexenv') body in
                 TFor (span, name, i', f', b')
               else
                 Errors.for_bounds_non_integral ~bound:"final" ~ty:(get_type f')
             else
               Errors.for_bounds_non_integral ~bound:"initial" ~ty:(get_type i'))
      | ABorrow { span; original; rename; region; body; mode; } ->
         adorn_error_with_span span
           (fun _ ->
             (match get_var lexenv original with
              | (Some (orig_ty, _)) ->
                 let u = type_universe orig_ty in
                 if ((u = LinearUniverse) || (u = TypeUniverse)) then
                   let region_obj = fresh_region () in
                   let refty =
                     (match mode with
                      | ReadBorrow ->
                         ReadRef (orig_ty, RegionTy region_obj)
                      | WriteBorrow ->
                         WriteRef (orig_ty, RegionTy region_obj))
                   in
                   let lexenv' = push_var lexenv rename refty VarLocal in
                   let rm' = add_region rm region region_obj in
                   let ctx' = update_lexenv ctx lexenv' in
                   let ctx''= update_rm ctx' rm' in
                   TBorrow {
                       span=span;
                       original=original;
                       rename=rename;
                       region=region;
                       orig_type=orig_ty;
                       ref_type=refty;
                       body=augment_stmt ctx'' body;
                       mode=mode
                     }
                 else
                   Errors.borrow_non_linear orig_ty
              | None ->
                 Errors.unknown_name ~kind:"variable" ~name:original))
      | ABlock (span, f, r) ->
         let a = augment_stmt ctx f in
         let b = augment_stmt ctx r in
         TBlock (span, a, b)
      | ADiscarding (span, e) ->
         adorn_error_with_span span
           (fun _ ->
             let e' = augment_expr module_name env rm typarams lexenv None e in
             ps ("Expression", show_texpr e');
             let u = type_universe (get_type e') in
             if ((u = LinearUniverse) || (u = TypeUniverse)) then
               Errors.discard_linear u
             else
               TDiscarding (span, e'))
      | AReturn (span, e) ->
         adorn_error_with_span span
           (fun _ ->
             let e' = augment_expr module_name env rm typarams lexenv (Some rt) e in
             pt ("Type", get_type e');
             let _ = match_type_with_value (env, module_name) rt e' in
             TReturn (span, e')))

and augment_lvalue_path (env: env) (module_name: module_name) (rm: region_map) (typarams: typarams) (lexenv: lexenv) (head_ty: ty) (elems: path_elem list): typed_path_elem list =
  match elems with
  | [elem] ->
     [augment_lvalue_path_elem env module_name rm typarams lexenv head_ty elem]
  | elem::rest ->
     let elem' = augment_lvalue_path_elem env module_name rm typarams lexenv head_ty elem in
     let rest' = augment_lvalue_path env module_name rm typarams lexenv (path_elem_type elem') rest in
     elem' :: rest'
  | [] ->
     err "Path is empty"

and augment_lvalue_path_elem (env: env) (module_name: module_name) (rm: region_map) (typarams: typarams) (lexenv: lexenv) (head_ty: ty) (elem: path_elem): typed_path_elem =
  match elem with
  | SlotAccessor slot_name ->
     (match head_ty with
      | NamedType (name, args, _) ->
         augment_slot_accessor_elem env module_name slot_name name args
      | _ ->
         Errors.path_not_record (type_string head_ty))
  | PointerSlotAccessor slot_name ->
     (match head_ty with
      | Pointer pointed_to ->
         augment_pointer_slot_accessor_elem env module_name slot_name pointed_to
      | WriteRef (ty, _) ->
         (match ty with
          | NamedType (name, args, _) ->
             augment_reference_slot_accessor_elem env module_name slot_name name args
          | _ ->
             Errors.path_not_record (type_string ty))
      | _ ->
         Errors.path_not_record (type_string head_ty))
  | ArrayIndex ie ->
     let ie' = augment_expr module_name env rm typarams lexenv None ie in
     let _ = ie' in
     (match head_ty with
      | _ ->
         Errors.lvalue_index ())


and get_union_type_definition (importing_module: module_name) (env: env) (ty: ty): (ty * decl list) =
  let name: qident = (match ty with
                      | NamedType (n, _, _) ->
                         n
                      | _ ->
                         Errors.case_non_union ty) in
  match get_decl_by_name env (qident_to_sident name) with
  | Some (Union { id; vis; name; mod_id; typarams; universe; _ }) ->
     let module_name = module_name_from_id env mod_id in
     if (vis = TypeVisPublic) || (importing_module = module_name) then
       let n = make_qident (module_name, name, name)
       and args = List.map (fun tp -> TyVar (typaram_to_tyvar tp)) (typarams_as_list typarams)
       in
       let ty = NamedType (n, args, universe)
       and cases = get_union_cases env id in
       (ty, cases)
     else
       Errors.case_non_public ()
  | _ ->
     Errors.case_non_union ty

and group_cases_whens (cases: decl list) (whens: abstract_when list): (typed_case * abstract_when) list =
  List.map (fun (TypedCase (n, s)) -> (TypedCase (n, s), List.find (fun (AbstractWhen (n', _, _)) -> n = n') whens))
    (List.map union_case_to_typed_case cases)

and group_bindings_slots (bindings: qbinding list) (slots: typed_slot list): (identifier * qtypespec * ty * identifier) list =
  let f (binding: qbinding): (identifier * qtypespec * ty * identifier) =
    let QBinding { name; ty; rename; } = binding in
    let (TypedSlot (_, ty')) = List.find (fun (TypedSlot (n', _)) -> equal_identifier name n') slots in
    (name, ty, ty', rename)
  in
  List.map f bindings

and augment_when (ctx: stmt_ctx) (typebindings: type_bindings) (w: abstract_when) (c: typed_case): typed_when =
  with_frame "Augment when"
    (fun _ ->
      let (StmtCtx (_, menv, rm, typarams, lexenv, _)) = ctx in
      let (AbstractWhen (name, bindings, body)) = w
      and (TypedCase (_, slots)) = c in
      pi ("Case name", name);
      let binding_names = List.map (fun (QBinding { name; _ }) -> name) bindings
      and slot_names = List.map (fun (TypedSlot (n, _)) -> n) slots in
      if ident_set_eq binding_names slot_names then
        let bindings' = group_bindings_slots bindings slots in
        let bindings'' = List.map (fun (n, ty, actual, rename) -> (n, parse_typespec menv rm typarams ty, replace_variables typebindings actual, rename)) bindings' in
        let newvars = List.map (fun (_, ty, actual, rename) ->
                          if equal_ty ty actual then
                            let _ = pi ("Binding name", rename)
                            in
                            pt ("Binding type",  ty);
                            (rename, ty, VarLocal)
                          else
                            Errors.slot_wrong_type
                              ~name:rename
                              ~expected:actual
                              ~actual:ty)
                        bindings'' in
        let lexenv' = push_vars lexenv newvars in
        let body' = augment_stmt (update_lexenv ctx lexenv') body in
        TypedWhen (name, List.map (fun (name, _, ty, rename) -> TypedBinding { name; ty; rename}) bindings'', body')
      else
        Errors.case_wrong_slots ())

let rec validate_constant_expression (expr: texpr): unit =
  if is_constant expr then
    ()
  else
    Errors.expression_not_constant ()

and is_constant = function
  | TNilConstant ->
     true
  | TBoolConstant _ ->
     true
  | TIntConstant _ ->
     true
  | TFloatConstant _ ->
     true
  | TStringConstant _ ->
     true
  | TConstVar _ ->
     true
  | TParamVar _ ->
     false
  | TLocalVar _ ->
     false
  | TFunVar _ ->
     true
  | TFuncall _ ->
     false
  | TMethodCall _ ->
     false
  | TVarMethodCall _ ->
     false
  | TFptrCall _ ->
     false
  | TCast _ ->
     true
  | TComparison (_, lhs, rhs) ->
     (is_constant lhs) && (is_constant rhs)
  | TConjunction (lhs, rhs) ->
     (is_constant lhs) && (is_constant rhs)
  | TDisjunction (lhs, rhs) ->
     (is_constant lhs) && (is_constant rhs)
  | TNegation e ->
     is_constant e
  | TIfExpression (c, t, f) ->
     (is_constant c) && (is_constant t) && (is_constant f)
  | TRecordConstructor (_, values) ->
     List.for_all (fun (_, v) -> is_constant v) values
  | TUnionConstructor (_, _, values) ->
     List.for_all (fun (_, v) -> is_constant v) values
  | TPath { head; elems; _ } ->
     (is_constant head) && (List.for_all is_path_elem_constant elems)
  | TEmbed _ ->
     true
  | TDeref _ ->
     false
  | TSizeOf _ ->
     true
  | TBorrowExpr _ ->
     false

and is_path_elem_constant = function
  | TSlotAccessor _ ->
     true
  | TPointerSlotAccessor _ ->
     true
  | TArrayIndex (e, _) ->
     is_constant e

let rec augment_decl (module_name: module_name) (kind: module_kind) (env: env) (decl: linked_definition): typed_decl =
  with_frame "Augment declaration"
    (fun _ ->
      let rm = empty_region_map in
      match decl with
      | LConstant (decl_id, vis, name, ty, expr, doc) ->
         ps ("Kind", "Constant");
         let expr' = augment_expr module_name env rm empty_typarams empty_lexenv (Some ty) expr in
         let _ = match_type_with_value (env, module_name) ty expr' in
         let _ = validate_constant_expression expr' in
         TConstant (decl_id, vis, name, ty, expr', doc)
      | LRecord (decl_id, vis, name, typarams, universe, slots, doc) ->
         ps ("Kind", "Record");
         TRecord (decl_id, vis, name, typarams, universe, slots, doc)
      | LUnion (decl_id, vis, name, typarams, universe, cases, doc) ->
         ps ("Kind", "Union");
         TUnion (decl_id, vis, name, typarams, universe, cases, doc)
      | LFunction (decl_id, vis, name, typarams, params, rt, body, doc, pragmas) ->
         ps ("Kind", "Function");
         pi ("Name", name);
         (match pragmas with
          | [ForeignImportPragma s] ->
             if typarams_size typarams = 0 then
               if kind = UnsafeModule then
                 TForeignFunction (decl_id, vis, name, params, rt, s, doc)
               else
                 Errors.foreign_in_safe_module ()
             else
               Errors.foreign_type_parameters ()
          | [ForeignExportPragma _] ->
             let ctx = StmtCtx (module_name, env, rm, typarams, (lexenv_from_params params), rt) in
             let body' = augment_stmt ctx body in
             TFunction (decl_id, vis, name, typarams, params, rt, body', doc)
          | [] ->
             let ctx = StmtCtx (module_name, env, rm, typarams, (lexenv_from_params params), rt) in
             let body' = augment_stmt ctx body in
             TFunction (decl_id, vis, name, typarams, params, rt, body', doc)
          | _ ->
             Errors.fun_invalid_pragmas ())
      | LTypeclass (decl_id, vis, name, typaram, methods, doc) ->
         ps ("Kind", "Typeclass");
         TTypeClass (decl_id, vis, name, typaram, List.map (augment_method_decl env rm typaram) methods, doc)
      | LInstance (decl_id, vis, name, typarams, arg, methods, doc) ->
         ps ("Kind", "Instance");
         TInstance (decl_id, vis, name, typarams, arg, List.map (augment_method_def module_name env rm typarams) methods, doc))

and lexenv_from_params (params: value_parameter list): lexenv =
  match params with
  | (ValueParameter (n, t))::rest ->
     push_var (lexenv_from_params rest) n t VarParam
  | [] ->
     empty_lexenv

and augment_method_decl _ _ _ (LMethodDecl (decl_id, name, params, rt, _)) =
  TypedMethodDecl (decl_id, name, params, rt)

and augment_method_def module_name menv rm typarams (LMethodDef (ins_meth_id, name, params, rt, _, body)) =
  let ctx = StmtCtx (module_name, menv, rm, typarams, (lexenv_from_params params), rt) in
  let body' = augment_stmt ctx body in
  TypedMethodDef (ins_meth_id, name, params, rt, body')

let augment_module menv (LinkedModule { name; decls; kind; _ }) =
  with_frame "Typing pass"
    (fun _ ->
      let decls' = List.map (augment_decl name kind menv) decls in
      TypedModule (name, decls'))
