(* MODIFIED by Conrad Watt *)

open Values
open Types
open Instance
open Ast
open Source


(* Errors *)

module Link = Error.Make ()
module Trap = Error.Make ()
module Crash = Error.Make ()
module Exhaustion = Error.Make ()

exception Link = Link.Error
exception Trap = Trap.Error
exception Crash = Crash.Error (* failure that cannot happen in valid code *)
exception Exhaustion = Exhaustion.Error

let lookup category list x =
  try Lib.List32.nth list x.it with Failure _ ->
    Crash.error x.at ("undefined " ^ category ^ " " ^ Int32.to_string x.it)

let type_ (inst : instance) x = lookup "type" inst.module_.it.types x
let func (inst : instance) x = lookup "function" inst.funcs x
let table (inst : instance) x = lookup "table" inst.tables x
let memory (inst : instance) x = lookup "memory" inst.memories x
let global (inst : instance) x = lookup "global" inst.globals x

let func_type_of = function
  | AstFunc (inst, f) -> (lookup "type" (!inst).module_.it.types f.it.ftype).it
  | HostFunc (t, _) -> t

(* Functions & Constants *)

let update_globs gs ls = List.map2 (fun g l -> g := Ast_convert.unconvert_value (InterpreterAux.Wasm.g_val l)) gs ls

let resolve_res at = function
	| InterpreterAux.Wasm_Interpreter.RTrap -> Trap.error at "trap!"
	| InterpreterAux.Wasm_Interpreter.RValue res -> List.map Ast_convert.unconvert_value res
	| InterpreterAux.Wasm_Interpreter.RCrash InterpreterAux.Wasm_Interpreter.CExhaustion -> Exhaustion.error at "call stack exhausted"
	| _ -> Crash.error at "crash!"

let invoke (clos : closure) (vs : value list) : value list =
  let (inst, at) = match clos with AstFunc (finst, f) -> (finst, f.at) | HostFunc _ -> (ref (instance (empty_module @@ no_region)), no_region) in
  let FuncType (ins, out) = func_type_of clos in
  if List.length vs <> List.length ins then
    Crash.error at "wrong number of arguments";
(*  let inst = instance (empty_module @@ at) in
  let c = config (List.rev vs) [Invoke clos @@ at] in *)
  try
		let (ss, gs, store) = Ast_convert.create_store inst in
		let (s', res) = Interpreter.Wasm_Interpreter_Printing.run InterpreterAux.Arith.zero_nat (store, Ast_convert.empty_config ss vs clos) in
		let _ = update_globs gs (InterpreterAux.Wasm.globs s') in
		resolve_res at res
  with Stack_overflow -> Exhaustion.error at "call stack exhausted"

let eval_const (inst : instance) (const : const) : value =
	let body = Ast_convert.convert_b_es_to_es (Ast_convert.convert_instrs const.it) in
	let at = const.at in
	let (ss, gs, store) = Ast_convert.create_store (ref inst) in
	let (s', res) = Interpreter.Wasm_Interpreter_Printing.run InterpreterAux.Arith.zero_nat (store, ([], body)) in
	let _ = update_globs gs (InterpreterAux.Wasm.globs s') in
	let rs_values = resolve_res at res in
  match rs_values with
  | [v] -> v
  | vs -> Crash.error const.at "wrong number of results on stack"

let i32 (v : value) at =
  match v with
  | I32 i -> i
  | _ -> Crash.error at "type error: i32 value expected"


(* Modules *)

let create_closure (m : module_) (f : func) =
  AstFunc (ref (instance m), f)

let create_table (tab : table) =
  let {ttype = TableType (lim, t)} = tab.it in
  Table.create t lim

let create_memory (mem : memory) =
  let {mtype = MemoryType lim} = mem.it in
  Memory.create lim

let create_global (glob : global) =
  let {gtype = GlobalType (t, _); _} = glob.it in
  ref (default_value t)

let init_closure (inst : instance) (clos : closure) =
  match clos with
  | AstFunc (inst_ref, _) -> inst_ref := inst
  | _ -> assert false

let init_table (inst : instance) (seg : table_segment) =
  let {index; offset = const; init} = seg.it in
  let tab = table inst index in
  let offset = i32 (eval_const inst const) const.at in
  let end_ = Int32.(add offset (of_int (List.length init))) in
  let bound = Table.size tab in
  if I32.lt_u bound end_ || I32.lt_u end_ offset then
    Link.error seg.at "elements segment does not fit table";
  fun () -> Table.blit tab offset (List.map (fun x -> Func (func inst x)) init)

let init_memory (inst : instance) (seg : memory_segment) =
  let {index; offset = const; init} = seg.it in
  let mem = memory inst index in
  let offset' = i32 (eval_const inst const) const.at in
  let offset = I64_convert.extend_u_i32 offset' in
  let end_ = Int64.(add offset (of_int (String.length init))) in
  let bound = Memory.bound mem in
  if I64.lt_u bound end_ || I64.lt_u end_ offset then
    Link.error seg.at "data segment does not fit memory";
  fun () -> Memory.blit mem offset init

let init_global (inst : instance) (ref : value ref) (glob : global) =
  let {value; _} = glob.it in
  ref := eval_const inst value

let check_limits actual expected at =
  if I32.lt_u actual.min expected.min then
    Link.error at "actual size smaller than declared";
  if
    match actual.max, expected.max with
    | _, None -> false
    | None, Some _ -> true
    | Some i, Some j -> I32.gt_u i j
  then Link.error at "maximum size larger than declared"

let add_import (ext : extern) (im : import) (inst : instance) : instance =
  let {idesc; _} = im.it in
  match ext, idesc.it with
  | ExternalFunc clos, FuncImport x when func_type_of clos = (type_ inst x).it ->
    {inst with funcs = clos :: inst.funcs}
  | ExternalTable tab, TableImport (TableType (lim, t))
    when Table.elem_type tab = t ->
    check_limits (Table.limits tab) lim idesc.at;
    {inst with tables = tab :: inst.tables}
  | ExternalMemory mem, MemoryImport (MemoryType lim) ->
    check_limits (Memory.limits mem) lim idesc.at;
    {inst with memories = mem :: inst.memories}
  | ExternalGlobal v, GlobalImport (GlobalType (t, _)) when type_of v = t ->
    {inst with globals = ref v :: inst.globals}
  | _ ->
    Link.error idesc.at "type mismatch"

let add_export (inst : instance) (ex : export) (map : extern ExportMap.t)
  : extern ExportMap.t =
  let {name; edesc} = ex.it in
  let ext =
    match edesc.it with
    | FuncExport x -> ExternalFunc (func inst x)
    | TableExport x -> ExternalTable (table inst x)
    | MemoryExport x -> ExternalMemory (memory inst x)
    | GlobalExport x -> ExternalGlobal !(global inst x)
  in ExportMap.add name ext map

let init (m : module_) (exts : extern list) : instance =
  let
    { imports; tables; memories; globals; funcs;
      exports; elems; data; start; _
    } = m.it
  in
  if List.length exts <> List.length imports then
    Link.error m.at "wrong number of imports provided for initialisation";
  let fs = List.map (create_closure m) funcs in
  let gs = List.map create_global globals in
  let inst =
    List.fold_right2 add_import exts imports
      { (instance m) with
        funcs = fs;
        tables = List.map create_table tables;
        memories = List.map create_memory memories;
        globals = gs;
      }
  in
  List.iter2 (init_global inst) gs globals;
  List.iter (init_closure inst) fs;
  let init_elems = List.map (init_table inst) elems in
  let init_datas = List.map (init_memory inst) data in
  List.iter (fun f -> f ()) init_elems;
  List.iter (fun f -> f ()) init_datas;
  Lib.Option.app (fun x -> ignore (invoke (func inst x) [])) start;
  {inst with exports = List.fold_right (add_export inst) exports inst.exports}
