(* SPDX-License-Identifier: AGPL-3.0-or-later *)
(* Copyright Â© 2021-2024 OCamlPro *)
(* Written by the Owi programmers *)

open Types
open Binary
open Syntax
open Fmt

type typ =
  | Num_type of num_type
  | Ref_type of binary heap_type
  | Any
  | Something

let typ_equal t1 t2 =
  match (t1, t2) with
  | Num_type t1, Num_type t2 -> Types.num_type_eq t1 t2
  | Ref_type t1, Ref_type t2 -> Types.heap_type_eq t1 t2
  | Any, _ | _, Any -> true
  | Something, _ | _, Something -> true
  | _, _ -> false

let pp_typ fmt = function
  | Num_type t -> pp_num_type fmt t
  | Ref_type t -> pp_heap_type fmt t
  | Any -> string fmt "any"
  | Something -> string fmt "something"

let pp_typ_list fmt l = list ~sep:sp pp_typ fmt l

let typ_of_val_type = function
  | Types.Ref_type (_null, t) -> Ref_type t
  | Num_type t -> Num_type t

let typ_of_pt pt = typ_of_val_type @@ snd pt

module Index = struct
  module M = Int
  module Map = Map.Make (Int)
  include M
end

let check_mem modul n =
  if n >= Array.length modul.mem then Error (`Unknown_memory (Raw n)) else Ok ()

let check_data modul n =
  if n >= Array.length modul.data then Error (`Unknown_data (Raw n)) else Ok ()

let check_align memarg_align align =
  if Int32.ge memarg_align align then Error `Alignment_too_large else Ok ()

module Env = struct
  type t =
    { locals : typ Index.Map.t
    ; result_type : binary result_type
    ; blocks : typ list list
    ; modul : Binary.modul
    ; refs : (int, unit) Hashtbl.t
    }

  let local_get i env =
    match Index.Map.find_opt i env.locals with
    | None -> Error (`Unknown_local (Raw i))
    | Some v -> Ok v

  let global_get i modul =
    if i >= Array.length modul.global then Error (`Unknown_global (Raw i))
    else
      match modul.global.(i) with
      | Runtime.Local { typ = desc; _ } | Imported { desc; _ } -> Ok desc

  let func_get i modul =
    if i >= Array.length modul.func then Error (`Unknown_func (Raw i))
    else
      match modul.func.(i) with
      | Runtime.Local { type_f = Bt_raw (_, t); _ }
      | Imported { desc = Bt_raw (_, t); _ } ->
        Ok t

  let block_type_get i env =
    match List.nth_opt env.blocks i with
    | None -> Error (`Unknown_label (Raw i))
    | Some bt -> Ok bt

  let table_type_get i (modul : Binary.modul) =
    if i >= Array.length modul.table then Error (`Unknown_table (Raw i))
    else
      match modul.table.(i) with
      | Runtime.Local (_, (_, t)) | Imported { desc = _, t; _ } -> Ok t

  let elem_type_get i modul =
    if i >= Array.length modul.elem then Error (`Unknown_elem (Raw i))
    else match modul.elem.(i) with value -> Ok value.typ

  let make ~params ~locals ~modul ~result_type ~refs =
    let l = List.mapi (fun i v -> (i, v)) (params @ locals) in
    let locals =
      List.fold_left
        (fun locals (i, (_, typ)) ->
          let typ = typ_of_val_type typ in
          Index.Map.add i typ locals )
        Index.Map.empty l
    in
    { locals; modul; result_type; blocks = []; refs }
end

type stack = typ list

let i32 = Num_type I32

let i64 = Num_type I64

let f32 = Num_type F32

let f64 = Num_type F64

let i31 = Ref_type I31_ht

let any = Any

let itype = function S32 -> i32 | S64 -> i64

let ftype = function S32 -> f32 | S64 -> f64

let arraytype _modul _i = (* TODO *) assert false

module Stack : sig
  type t = typ list

  val drop : t -> t Result.t

  val pop : t -> t -> t Result.t

  val push : t -> t -> t Result.t

  val pop_push : binary block_type -> t -> t Result.t

  val pop_ref : t -> t Result.t

  val equal : t -> t -> bool

  val match_ref_type : binary heap_type -> binary heap_type -> bool

  val match_types : typ -> typ -> bool

  val pp : formatter -> t -> unit

  val match_prefix : prefix:t -> stack:t -> t option
end = struct
  type t = typ list

  let pp fmt (s : stack) = pf fmt "[%a]" pp_typ_list s

  let match_num_type (required : num_type) (got : num_type) =
    match (required, got) with
    | I32, I32 -> true
    | I64, I64 -> true
    | F32, F32 -> true
    | F64, F64 -> true
    | _, _ -> false

  let match_ref_type required got =
    match (required, got) with
    | Any_ht, _ -> true
    | None_ht, None_ht -> true
    | Eq_ht, Eq_ht -> true
    | I31_ht, I31_ht -> true
    | Struct_ht, Struct_ht -> true
    | Array_ht, Array_ht -> true
    | No_func_ht, No_func_ht -> true
    | Func_ht, Func_ht -> true
    | Extern_ht, Extern_ht -> true
    | No_extern_ht, No_extern_ht -> true
    | _ ->
      (* TODO: complete this *)
      false

  let match_types required got =
    match (required, got) with
    | Something, _ | _, Something -> true
    | Any, _ | _, Any -> true
    | Num_type required, Num_type got -> match_num_type required got
    | Ref_type required, Ref_type got -> match_ref_type required got
    | Num_type _, Ref_type _ | Ref_type _, Num_type _ -> false

  let rec equal s s' =
    match (s, s') with
    | [], s | s, [] -> List.for_all (function Any -> true | _ -> false) s
    | Any :: tl, Any :: tl' -> equal tl s' || equal s tl'
    | Any :: tl, hd :: tl' | hd :: tl', Any :: tl ->
      equal tl (hd :: tl') || equal (Any :: tl) tl'
    | hd :: tl, hd' :: tl' -> match_types hd hd' && equal tl tl'

  let ( ||| ) l r = match (l, r) with None, v | v, None -> v | _l, r -> r

  let rec match_prefix ~prefix ~stack =
    match (prefix, stack) with
    | [], stack -> Some stack
    | _hd :: _tl, [] -> None
    | _hd :: tl, Any :: tl' ->
      match_prefix ~prefix ~stack:tl' ||| match_prefix ~prefix:tl ~stack
    | hd :: tl, hd' :: tl' ->
      if match_types hd hd' then match_prefix ~prefix:tl ~stack:tl' else None

  let pop required stack =
    match match_prefix ~prefix:required ~stack with
    | None -> Error (`Type_mismatch "pop")
    | Some stack -> Ok stack

  let pop_ref = function
    | (Something | Ref_type _) :: tl -> Ok tl
    | Any :: _ as stack -> Ok stack
    | _ -> Error (`Type_mismatch "pop_ref")

  let drop stack =
    match stack with
    | [] -> Error (`Type_mismatch "drop")
    | Any :: _ -> Ok [ Any ]
    | _ :: tl -> Ok tl

  let push t stack = ok @@ t @ stack

  let pop_push (Bt_raw (_, (pt, rt))) stack =
    let pt, rt = (List.rev_map typ_of_pt pt, List.rev_map typ_of_val_type rt) in
    let* stack = pop pt stack in
    push rt stack
end

let rec typecheck_instr (env : Env.t) (stack : stack) (instr : binary instr) :
  stack Result.t =
  match instr with
  | Nop -> Ok stack
  | Drop -> Stack.drop stack
  | Return ->
    let+ _stack =
      Stack.pop (List.rev_map typ_of_val_type env.result_type) stack
    in
    [ any ]
  | Unreachable -> Ok [ any ]
  | I32_const _ -> Stack.push [ i32 ] stack
  | I64_const _ -> Stack.push [ i64 ] stack
  | F32_const _ -> Stack.push [ f32 ] stack
  | F64_const _ -> Stack.push [ f64 ] stack
  | I_unop (s, _op) ->
    let t = itype s in
    let* stack = Stack.pop [ t ] stack in
    Stack.push [ t ] stack
  | I_binop (s, _op) ->
    let t = itype s in
    let* stack = Stack.pop [ t; t ] stack in
    Stack.push [ t ] stack
  | F_unop (s, _op) ->
    let t = ftype s in
    let* stack = Stack.pop [ t ] stack in
    Stack.push [ t ] stack
  | F_binop (s, _op) ->
    let t = ftype s in
    let* stack = Stack.pop [ t; t ] stack in
    Stack.push [ t ] stack
  | I_testop (nn, _) ->
    let* stack = Stack.pop [ itype nn ] stack in
    Stack.push [ i32 ] stack
  | I_relop (nn, _) ->
    let t = itype nn in
    let* stack = Stack.pop [ t; t ] stack in
    Stack.push [ i32 ] stack
  | F_relop (nn, _) ->
    let t = ftype nn in
    let* stack = Stack.pop [ t; t ] stack in
    Stack.push [ i32 ] stack
  | Local_get (Raw i) ->
    let* t = Env.local_get i env in
    Stack.push [ t ] stack
  | Local_set (Raw i) ->
    let* t = Env.local_get i env in
    Stack.pop [ t ] stack
  | Local_tee (Raw i) ->
    let* t = Env.local_get i env in
    let* stack = Stack.pop [ t ] stack in
    Stack.push [ t ] stack
  | Global_get (Raw i) ->
    let* _mut, t = Env.global_get i env.modul in
    let t = typ_of_val_type t in
    Stack.push [ t ] stack
  | Global_set (Raw i) ->
    let* mut, t = Env.global_get i env.modul in
    let* () =
      match mut with Var -> Ok () | Const -> Error `Global_is_immutable
    in
    let t = typ_of_val_type t in
    Stack.pop [ t ] stack
  | If_else (_id, block_type, e1, e2) ->
    let* stack = Stack.pop [ i32 ] stack in
    let* stack_e1 = typecheck_expr env e1 ~is_loop:false block_type ~stack in
    let+ _stack_e2 = typecheck_expr env e2 ~is_loop:false block_type ~stack in
    stack_e1
  | I_load8 (nn, _, memarg) ->
    let* () = check_mem env.modul 0 in
    let* () = check_align memarg.align 1l in
    let* stack = Stack.pop [ i32 ] stack in
    Stack.push [ itype nn ] stack
  | I_load16 (nn, _, memarg) ->
    let* () = check_mem env.modul 0 in
    let* () = check_align memarg.align 2l in
    let* stack = Stack.pop [ i32 ] stack in
    Stack.push [ itype nn ] stack
  | I_load (nn, memarg) ->
    let* () = check_mem env.modul 0 in
    let max_allowed = match nn with S32 -> 4l | S64 -> 8l in
    let* () = check_align memarg.align max_allowed in
    let* stack = Stack.pop [ i32 ] stack in
    Stack.push [ itype nn ] stack
  | I64_load32 (_, memarg) ->
    let* () = check_mem env.modul 0 in
    let* () = check_align memarg.align 4l in
    let* stack = Stack.pop [ i32 ] stack in
    Stack.push [ i64 ] stack
  | I_store8 (nn, memarg) ->
    let* () = check_mem env.modul 0 in
    let* () = check_align memarg.align 1l in
    Stack.pop [ itype nn; i32 ] stack
  | I_store16 (nn, memarg) ->
    let* () = check_mem env.modul 0 in
    let* () = check_align memarg.align 2l in
    Stack.pop [ itype nn; i32 ] stack
  | I_store (nn, memarg) ->
    let* () = check_mem env.modul 0 in
    let max_allowed = match nn with S32 -> 4l | S64 -> 8l in
    let* () = check_align memarg.align max_allowed in
    Stack.pop [ itype nn; i32 ] stack
  | I64_store32 memarg ->
    let* () = check_mem env.modul 0 in
    let* () = check_align memarg.align 4l in
    Stack.pop [ i64; i32 ] stack
  | F_load (nn, memarg) ->
    let* () = check_mem env.modul 0 in
    let max_allowed = match nn with S32 -> 4l | S64 -> 8l in
    let* () = check_align memarg.align max_allowed in
    let* stack = Stack.pop [ i32 ] stack in
    Stack.push [ ftype nn ] stack
  | F_store (nn, memarg) ->
    let* () = check_mem env.modul 0 in
    let max_allowed = match nn with S32 -> 4l | S64 -> 8l in
    let* () = check_align memarg.align max_allowed in
    Stack.pop [ ftype nn; i32 ] stack
  | I_reinterpret_f (inn, fnn) ->
    let* stack = Stack.pop [ ftype fnn ] stack in
    Stack.push [ itype inn ] stack
  | F_reinterpret_i (fnn, inn) ->
    let* stack = Stack.pop [ itype inn ] stack in
    Stack.push [ ftype fnn ] stack
  | F32_demote_f64 ->
    let* stack = Stack.pop [ f64 ] stack in
    Stack.push [ f32 ] stack
  | F64_promote_f32 ->
    let* stack = Stack.pop [ f32 ] stack in
    Stack.push [ f64 ] stack
  | F_convert_i (fnn, inn, _) ->
    let* stack = Stack.pop [ itype inn ] stack in
    Stack.push [ ftype fnn ] stack
  | I_trunc_f (inn, fnn, _) | I_trunc_sat_f (inn, fnn, _) ->
    let* stack = Stack.pop [ ftype fnn ] stack in
    Stack.push [ itype inn ] stack
  | I32_wrap_i64 ->
    let* stack = Stack.pop [ i64 ] stack in
    Stack.push [ i32 ] stack
  | I_extend8_s nn | I_extend16_s nn ->
    let t = itype nn in
    let* stack = Stack.pop [ t ] stack in
    Stack.push [ t ] stack
  | I64_extend32_s ->
    let* stack = Stack.pop [ i64 ] stack in
    Stack.push [ i64 ] stack
  | I64_extend_i32 _ ->
    let* stack = Stack.pop [ i32 ] stack in
    Stack.push [ i64 ] stack
  | Memory_grow ->
    let* () = check_mem env.modul 0 in
    let* stack = Stack.pop [ i32 ] stack in
    Stack.push [ i32 ] stack
  | Memory_size ->
    let* () = check_mem env.modul 0 in
    Stack.push [ i32 ] stack
  | Memory_copy | Memory_fill ->
    let* () = check_mem env.modul 0 in
    Stack.pop [ i32; i32; i32 ] stack
  | Memory_init (Raw id) ->
    let* () = check_mem env.modul 0 in
    let* () = check_data env.modul id in
    Stack.pop [ i32; i32; i32 ] stack
  | Block (_, bt, expr) -> typecheck_expr env expr ~is_loop:false bt ~stack
  | Loop (_, bt, expr) -> typecheck_expr env expr ~is_loop:true bt ~stack
  | Call_indirect (Raw tbl_id, bt) ->
    let* _tbl_type = Env.table_type_get tbl_id env.modul in
    let* stack = Stack.pop [ i32 ] stack in
    Stack.pop_push bt stack
  | Call (Raw i) ->
    let* pt, rt = Env.func_get i env.modul in
    let* stack = Stack.pop (List.rev_map typ_of_pt pt) stack in
    Stack.push (List.rev_map typ_of_val_type rt) stack
  | Call_ref _t ->
    let+ stack = Stack.pop_ref stack in
    (* TODO:
       let bt = Env.type_get t env in
         Stack.pop_push (Some bt) stack
    *)
    stack
  | Return_call (Raw i) ->
    let* pt, rt = Env.func_get i env.modul in
    if
      not
        (Stack.equal
           (List.rev_map typ_of_val_type rt)
           (List.rev_map typ_of_val_type env.result_type) )
    then Error (`Type_mismatch "return_call")
    else
      let+ _stack = Stack.pop (List.rev_map typ_of_pt pt) stack in
      [ any ]
  | Return_call_indirect (Raw tbl_id, Bt_raw (_, (pt, rt))) ->
    let* _tbl_type = Env.table_type_get tbl_id env.modul in
    if
      not
        (Stack.equal
           (List.rev_map typ_of_val_type rt)
           (List.rev_map typ_of_val_type env.result_type) )
    then Error (`Type_mismatch "return_call_indirect")
    else
      let* stack = Stack.pop [ i32 ] stack in
      let+ _stack = Stack.pop (List.rev_map typ_of_pt pt) stack in
      [ any ]
  | Return_call_ref (Bt_raw (_, (pt, rt))) ->
    if
      not
        (Stack.equal
           (List.rev_map typ_of_val_type rt)
           (List.rev_map typ_of_val_type env.result_type) )
    then Error (`Type_mismatch "return_call_ref")
    else
      let* stack = Stack.pop_ref stack in
      let+ _stack = Stack.pop (List.rev_map typ_of_pt pt) stack in
      [ any ]
  | Data_drop (Raw id) ->
    let* () = check_data env.modul id in
    Ok stack
  | Table_init (Raw ti, Raw ei) ->
    let* table_typ = Env.table_type_get ti env.modul in
    let* elem_typ = Env.elem_type_get ei env.modul in
    if not @@ Stack.match_ref_type (snd table_typ) (snd elem_typ) then
      Error (`Type_mismatch "table_init")
    else Stack.pop [ i32; i32; i32 ] stack
  | Table_copy (Raw i, Raw i') ->
    let* typ = Env.table_type_get i env.modul in
    let* typ' = Env.table_type_get i' env.modul in
    if not @@ Types.ref_type_eq typ typ' then Error (`Type_mismatch "table_copy")
    else Stack.pop [ i32; i32; i32 ] stack
  | Table_fill (Raw i) ->
    let* _null, t = Env.table_type_get i env.modul in
    Stack.pop [ i32; Ref_type t; i32 ] stack
  | Table_grow (Raw i) ->
    let* _null, t = Env.table_type_get i env.modul in
    let* stack = Stack.pop [ i32; Ref_type t ] stack in
    Stack.push [ i32 ] stack
  | Table_size (Raw i) ->
    let* _null, _t = Env.table_type_get i env.modul in
    Stack.push [ i32 ] stack
  | Ref_is_null ->
    let* stack = Stack.pop_ref stack in
    Stack.push [ i32 ] stack
  | Ref_null rt -> Stack.push [ Ref_type rt ] stack
  | Elem_drop (Raw id) ->
    let* _elem_typ = Env.elem_type_get id env.modul in
    Ok stack
  | Select t ->
    let* stack = Stack.pop [ i32 ] stack in
    begin
      match t with
      | None -> begin
        match stack with
        | Ref_type _ :: _tl -> Error (`Type_mismatch "select implicit")
        | Any :: _ -> Ok [ Something; Any ]
        | hd :: Any :: _ -> ok @@ (hd :: [ Any ])
        | hd :: hd' :: tl when Stack.match_types hd hd' -> ok @@ (hd :: tl)
        | _ -> Error (`Type_mismatch "select")
      end
      | Some t ->
        let t = List.map typ_of_val_type t in
        let* stack = Stack.pop t stack in
        let* stack = Stack.pop t stack in
        Stack.push t stack
    end
  | Ref_func (Raw i) ->
    if not @@ Hashtbl.mem env.refs i then Error `Undeclared_function_reference
    else Stack.push [ Ref_type Func_ht ] stack
  | Br (Raw i) ->
    let* jt = Env.block_type_get i env in
    let* _stack = Stack.pop jt stack in
    Ok [ any ]
  | Br_if (Raw i) ->
    let* stack = Stack.pop [ i32 ] stack in
    let* jt = Env.block_type_get i env in
    let* stack = Stack.pop jt stack in
    Stack.push jt stack
  | Br_table (branches, Raw i) ->
    let* stack = Stack.pop [ i32 ] stack in
    let* default_jt = Env.block_type_get i env in
    let* _stack = Stack.pop default_jt stack in
    let* () =
      array_iter
        (fun (Raw i : binary indice) ->
          let* jt = Env.block_type_get i env in
          if not (List.length jt = List.length default_jt) then
            Error (`Type_mismatch "br_table")
          else
            let* _stack = Stack.pop jt stack in
            Ok () )
        branches
    in
    Ok [ any ]
  | Table_get (Raw i) ->
    let* _null, t = Env.table_type_get i env.modul in
    let* stack = Stack.pop [ i32 ] stack in
    Stack.push [ Ref_type t ] stack
  | Table_set (Raw i) ->
    let* _null, t = Env.table_type_get i env.modul in
    Stack.pop [ Ref_type t; i32 ] stack
  | Array_len ->
    (* TODO: fixme, Something is not right *)
    let* stack = Stack.pop [ Something ] stack in
    Stack.push [ i32 ] stack
  | Ref_i31 ->
    let* stack = Stack.pop [ i32 ] stack in
    Stack.push [ i31 ] stack
  | I31_get_s | I31_get_u ->
    let* stack = Stack.pop [ i31 ] stack in
    Stack.push [ i32 ] stack
  | ( Array_new_data _ | Array_new _ | Array_new_default _ | Array_new_elem _
    | Array_new_fixed _ | Array_get _ | Array_get_u _ | Array_set _
    | Struct_get _ | Struct_get_s _ | Struct_set _ | Struct_new _
    | Struct_new_default _ | Extern_externalize | Extern_internalize
    | Ref_as_non_null | Ref_cast _ | Ref_test _ | Br_on_non_null _
    | Br_on_null _ | Br_on_cast _ | Br_on_cast_fail _ | Ref_eq ) as i ->
    Log.debug2 "TODO (typecheck instr) %a" pp_instr i;
    assert false

and typecheck_expr env expr ~is_loop (block_type : binary block_type option)
  ~stack:previous_stack : stack Result.t =
  let pt, rt =
    Option.fold ~none:([], [])
      ~some:(fun (Bt_raw (_, (pt, rt)) : binary block_type) ->
        (List.rev_map typ_of_pt pt, List.rev_map typ_of_val_type rt) )
      block_type
  in
  let jump_type = if is_loop then pt else rt in
  let env = { env with blocks = jump_type :: env.blocks } in
  let* stack = list_fold_left (typecheck_instr env) pt expr in
  if not (Stack.equal rt stack) then Error (`Type_mismatch "typecheck_expr 1")
  else
    match Stack.match_prefix ~prefix:pt ~stack:previous_stack with
    | None ->
      Error
        (`Type_mismatch
          (Fmt.str "expected a prefix of %a but stack has type %a" Stack.pp pt
             Stack.pp previous_stack ) )
    | Some stack_to_push -> Stack.push rt stack_to_push

let typecheck_function (modul : modul) func refs =
  match func with
  | Runtime.Imported _ -> Ok ()
  | Local func ->
    let (Bt_raw (_, (params, result))) = func.type_f in
    let env =
      Env.make ~params ~modul ~locals:func.locals ~result_type:result ~refs
    in
    let* stack =
      typecheck_expr env func.body ~is_loop:false
        (Some (Bt_raw (None, ([], result))))
        ~stack:[]
    in
    let required = List.rev_map typ_of_val_type result in
    if not @@ Stack.equal required stack then
      Error (`Type_mismatch "typecheck_function")
    else Ok ()

let typecheck_const_instr (modul : modul) refs stack = function
  | I32_const _ -> Stack.push [ i32 ] stack
  | I64_const _ -> Stack.push [ i64 ] stack
  | F32_const _ -> Stack.push [ f32 ] stack
  | F64_const _ -> Stack.push [ f64 ] stack
  | Ref_null t -> Stack.push [ Ref_type t ] stack
  | Ref_func (Raw i) ->
    let* _t = Env.func_get i modul in
    Hashtbl.add refs i ();
    Stack.push [ Ref_type Func_ht ] stack
  | Global_get (Raw i) ->
    if i >= Array.length modul.global then Error (`Unknown_global (Raw i))
    else
      let* mut, typ =
        match modul.global.(i) with
        | Runtime.Local _ -> Error (`Unknown_global (Raw i))
        | Imported { desc; _ } -> Ok desc
      in
      let* () =
        match mut with
        | Const -> Ok ()
        | Var -> Error `Constant_expression_required
      in
      Stack.push [ typ_of_val_type typ ] stack
  | I_binop (t, _op) ->
    let t = itype t in
    let* stack = Stack.pop [ t; t ] stack in
    Stack.push [ t ] stack
  | Array_new t ->
    let t = arraytype modul t in
    let* stack = Stack.pop [ i32; t ] stack in
    Stack.push [ Ref_type Array_ht ] stack
  | Array_new_default _i -> assert false
  | Ref_i31 ->
    let* stack = Stack.pop [ i32 ] stack in
    Stack.push [ i31 ] stack
  | _ -> Error `Constant_expression_required

let typecheck_const_expr (modul : modul) refs =
  list_fold_left (typecheck_const_instr modul refs) []

let typecheck_global (modul : modul) refs
  (global : (global, binary global_type) Runtime.t) =
  match global with
  | Imported _ -> Ok ()
  | Local { typ; init; _ } -> (
    let* real_type = typecheck_const_expr modul refs init in
    match real_type with
    | [ real_type ] ->
      let expected = typ_of_val_type @@ snd typ in
      if not @@ typ_equal expected real_type then
        Error (`Type_mismatch "typecheck global 1")
      else Ok ()
    | _whatever -> Error (`Type_mismatch "typecheck_global 2") )

let typecheck_elem modul refs (elem : elem) =
  let _null, expected_type = elem.typ in
  let* () =
    list_iter
      (fun init ->
        let* real_type = typecheck_const_expr modul refs init in
        match real_type with
        | [ real_type ] ->
          if not @@ typ_equal (Ref_type expected_type) real_type then
            Error (`Type_mismatch "typecheck_elem 1")
          else Ok ()
        | _whatever -> Error (`Type_mismatch "typecheck elem 2") )
      elem.init
  in
  match elem.mode with
  | Elem_passive | Elem_declarative -> Ok ()
  | Elem_active (None, _e) -> assert false
  | Elem_active (Some tbl_i, e) -> (
    let* _null, tbl_type = Env.table_type_get tbl_i modul in
    if not @@ Types.heap_type_eq tbl_type expected_type then
      Error (`Type_mismatch "typecheck elem 3")
    else
      let* t = typecheck_const_expr modul refs e in
      match t with
      | [ Ref_type t ] ->
        if not @@ Types.heap_type_eq t tbl_type then
          Error (`Type_mismatch "typecheck_elem 4")
        else Ok ()
      | [ _t ] -> Ok ()
      | _whatever -> Error (`Type_mismatch "typecheck_elem 5") )

let typecheck_data modul refs (data : data) =
  match data.mode with
  | Data_passive -> Ok ()
  | Data_active (n, e) -> (
    let* () = check_mem modul n in
    let* t = typecheck_const_expr modul refs e in
    match t with
    | [ _t ] -> Ok ()
    | _whatever -> Error (`Type_mismatch "typecheck_data") )

let typecheck_start { start; func; _ } =
  match start with
  | None -> Ok ()
  | Some idx -> (
    let* f =
      if idx >= Array.length func then Error (`Unknown_func (Raw idx))
      else Ok func.(idx)
    in
    match f with
    | Local { type_f = Bt_raw (_, ([], [])); _ }
    | Imported { desc = Bt_raw (_, ([], [])); _ } ->
      Ok ()
    | _ -> Error `Start_function )

let validate_exports modul =
  let* () =
    list_iter
      (fun { id; name = _ } ->
        let* _t = Env.func_get id modul in
        Ok () )
      modul.exports.func
  in
  let* () =
    list_iter
      (fun { id; name = _ } ->
        let* _t = Env.table_type_get id modul in
        Ok () )
      modul.exports.table
  in
  let* () =
    list_iter
      (fun { id; name = _ } ->
        let* _t = Env.global_get id modul in
        Ok () )
      modul.exports.global
  in
  list_iter
    (fun { id; name = _ } ->
      let* () = check_mem modul id in
      Ok () )
    modul.exports.mem

let check_limit { min; max } =
  match max with
  | None -> Ok ()
  | Some max ->
    if min > max then Error `Size_minimum_greater_than_maximum else Ok ()

let validate_tables modul =
  array_iter
    (function
      | Runtime.Local (_, (limits, _)) | Imported { desc = limits, _; _ } ->
        check_limit limits )
    modul.table

let validate_mem modul =
  array_iter
    (function
      | Runtime.Local (_, desc) | Imported { desc; _ } ->
        let* () =
          if desc.min > 65536 then Error `Memory_size_too_large
          else
            match desc.max with
            | Some max when max > 65536 -> Error `Memory_size_too_large
            | Some _ | None -> Ok ()
        in
        check_limit desc )
    modul.mem

let modul (modul : modul) =
  Log.debug0 "typechecking ...@\n";
  let refs = Hashtbl.create 512 in
  let* () = array_iter (typecheck_global modul refs) modul.global in
  let* () = array_iter (typecheck_elem modul refs) modul.elem in
  let* () = array_iter (typecheck_data modul refs) modul.data in
  let* () = typecheck_start modul in
  let* () = validate_exports modul in
  let* () = validate_tables modul in
  let* () = validate_mem modul in
  List.iter
    (fun (export : export) -> Hashtbl.add refs export.id ())
    modul.exports.func;
  array_iter (fun func -> typecheck_function modul func refs) modul.func
