From f2da9435e0f52afb36b52dcc3b02dc5edc6cd7bb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ole=20Kr=C3=BCger?= <ole.kruger@trili.tech>
Date: Fri, 12 Aug 2022 15:38:21 +0100
Subject: [PATCH 1/4] SCORU: WASM: Eliminate module_ref in favour of explicit
 module_reg passing

---
 src/lib_scoru_wasm/test/ast_generators.ml     |  20 +--
 src/lib_scoru_wasm/test/ast_printer.ml        |   2 +-
 src/lib_scoru_wasm/test/test_input.ml         |  10 +-
 src/lib_scoru_wasm/test/test_wasm_encoding.ml |  34 +---
 src/lib_scoru_wasm/wasm_encoding.ml           | 129 +++++++-------
 src/lib_scoru_wasm/wasm_encoding.mli          |  48 ++----
 src/lib_scoru_wasm/wasm_pvm.ml                |  78 ++++-----
 src/lib_webassembly/bin/script/run.ml         |  22 ++-
 src/lib_webassembly/exec/eval.ml              | 160 +++++++++---------
 src/lib_webassembly/exec/eval.mli             |  12 +-
 src/lib_webassembly/runtime/instance.ml       |  20 +--
 11 files changed, 231 insertions(+), 304 deletions(-)

diff --git a/src/lib_scoru_wasm/test/ast_generators.ml b/src/lib_scoru_wasm/test/ast_generators.ml
index 7b6580e17394..87c7171fc274 100644
--- a/src/lib_scoru_wasm/test/ast_generators.ml
+++ b/src/lib_scoru_wasm/test/ast_generators.ml
@@ -412,22 +412,20 @@ let allocations_gen =
   let+ datas = datas_table_gen in
   Ast.{blocks; datas}
 
-let module_ref_and_instance_gen ?module_reg () =
+let module_key_and_instance_gen ?module_reg () =
   let module_reg =
     match module_reg with
     | None -> Instance.ModuleMap.create ()
     | Some module_reg -> module_reg
   in
   let* module_name = string_printable in
-  let module_ref =
-    Instance.(alloc_module_ref (Module_key module_name) module_reg)
-  in
+  let module_key = Instance.Module_key module_name in
   let* types = vector_gen func_type_gen in
-  let* funcs = vector_gen @@ func_gen module_ref in
+  let* funcs = vector_gen @@ func_gen module_key in
   let* tables = vector_gen table_gen in
   let* memories = vector_gen memory_gen in
   let* globals = vector_gen global_gen in
-  let* exports = map_gen (extern_gen module_ref) in
+  let* exports = map_gen (extern_gen module_key) in
   let* elems = vector_gen elems_gen in
   let* datas = vector_gen datas_gen in
   let* allocations = allocations_gen in
@@ -444,14 +442,14 @@ let module_ref_and_instance_gen ?module_reg () =
       allocations;
     }
   in
-  Instance.update_module_ref module_ref module_ ;
-  return (module_ref, module_)
+  Instance.update_module_ref module_reg module_key module_ ;
+  return (module_key, module_)
 
 let module_gen ?module_reg () =
-  map snd (module_ref_and_instance_gen ?module_reg ())
+  map snd (module_key_and_instance_gen ?module_reg ())
 
 let frame_gen ~module_reg =
-  let* inst, _ = module_ref_and_instance_gen ~module_reg () in
+  let* inst, _ = module_key_and_instance_gen ~module_reg () in
   let+ locals = small_list (map ref value_gen) in
   Eval.{inst; locals}
 
@@ -471,7 +469,7 @@ let rec admin_instr'_gen ~module_reg depth =
     Refer ref_
   in
   let invoke_gen =
-    let* inst, _ = module_ref_and_instance_gen ~module_reg () in
+    let* inst, _ = module_key_and_instance_gen ~module_reg () in
     let+ func = func_gen inst in
     Invoke func
   in
diff --git a/src/lib_scoru_wasm/test/ast_printer.ml b/src/lib_scoru_wasm/test/ast_printer.ml
index 92b74d2fb1d2..f2e1b8a4b08d 100644
--- a/src/lib_scoru_wasm/test/ast_printer.ml
+++ b/src/lib_scoru_wasm/test/ast_printer.ml
@@ -541,7 +541,7 @@ let pp_module out
 
 let pp_frame out frame =
   let open Eval in
-  let (Module_key key) = frame.inst.key in
+  let (Module_key key) = frame.inst in
   Format.fprintf
     out
     "@[<v 2>{module = %s;@;locals = %a;@;}@]"
diff --git a/src/lib_scoru_wasm/test/test_input.ml b/src/lib_scoru_wasm/test/test_input.ml
index 52529774d74d..d25e9cbe07c9 100644
--- a/src/lib_scoru_wasm/test/test_input.ml
+++ b/src/lib_scoru_wasm/test/test_input.ml
@@ -156,19 +156,19 @@ let test_host_fun () =
   Host_funcs.register_host_funcs host_funcs_registry ;
 
   let module_reg = Instance.ModuleMap.create () in
-  let module_ref =
-    Instance.(alloc_module_ref (Module_key "test") ~module_inst module_reg)
-  in
+  let module_key = Instance.Module_key "test" in
+  Instance.update_module_ref module_reg module_key module_inst ;
 
   let* result =
     Eval.invoke
-      ~caller:module_ref
+      ~module_reg
+      ~caller:module_key
       host_funcs_registry
       ~input
       Host_funcs.Internal_for_tests.read_input
       values
   in
-  let* module_inst = Instance.resolve_module_ref module_ref in
+  let* module_inst = Instance.resolve_module_ref module_reg module_key in
   let* memory = Lazy_vector.LwtInt32Vector.get 0l module_inst.memories in
   assert (Input_buffer.num_elements input = Z.zero) ;
   let* m = Memory.load_bytes memory 0l 1 in
diff --git a/src/lib_scoru_wasm/test/test_wasm_encoding.ml b/src/lib_scoru_wasm/test/test_wasm_encoding.ml
index 6dad0f1250db..003459980d01 100644
--- a/src/lib_scoru_wasm/test/test_wasm_encoding.ml
+++ b/src/lib_scoru_wasm/test/test_wasm_encoding.ml
@@ -103,7 +103,6 @@ let test_module_roundtrip () =
        is not important when encoding or decoding. *)
     Instance.ModuleMap.create ()
   in
-  let lazy_dummy_module_reg = Lazy.from_val dummy_module_reg in
 
   qcheck
     ~print
@@ -112,17 +111,11 @@ let test_module_roundtrip () =
       (* We need to print here in order to force lazy bindings to be evaluated. *)
       let module1_str = print module1 in
       let*! module2 =
-        encode_decode
-          (Wasm_encoding.module_instance_encoding
-             ~module_reg:lazy_dummy_module_reg)
-          module1
+        encode_decode Wasm_encoding.module_instance_encoding module1
       in
       let module2_str = print module2 in
       let*! module3 =
-        encode_decode
-          (Wasm_encoding.module_instance_encoding
-             ~module_reg:lazy_dummy_module_reg)
-          module2
+        encode_decode Wasm_encoding.module_instance_encoding module2
       in
       let module3_str = print module3 in
       (* Check that modules match. *)
@@ -142,30 +135,19 @@ let test_generic_tree ~pp ~gen ~encoding =
        is not important when encoding or decoding. *)
     Instance.ModuleMap.create ()
   in
-  let lazy_dummy_module_reg = Lazy.from_val dummy_module_reg in
   let host_funcs = Host_funcs.empty () in
   qcheck ~print (gen ~host_funcs ~module_reg:dummy_module_reg) (fun value1 ->
       let*! empty_tree = empty_tree () in
       (* We need to print here in order to force lazy bindings to be evaluated. *)
       let _ = print value1 in
       let*! tree1 =
-        Tree_encoding.encode
-          (encoding ~host_funcs ~module_reg:lazy_dummy_module_reg)
-          value1
-          empty_tree
-      in
-      let*! value2 =
-        Tree_encoding.decode
-          (encoding ~host_funcs ~module_reg:lazy_dummy_module_reg)
-          tree1
+        Tree_encoding.encode (encoding ~host_funcs) value1 empty_tree
       in
+      let*! value2 = Tree_encoding.decode (encoding ~host_funcs) tree1 in
       (* We need to print here in order to force lazy bindings to be evaluated. *)
       let _ = print value2 in
       let*! tree2 =
-        Tree_encoding.encode
-          (encoding ~host_funcs ~module_reg:lazy_dummy_module_reg)
-          value2
-          empty_tree
+        Tree_encoding.encode (encoding ~host_funcs) value2 empty_tree
       in
       assert (Tree.equal tree1 tree2) ;
       return_unit)
@@ -190,8 +172,7 @@ let test_input_buffer_tree () =
   test_generic_tree
     ~pp:Ast_printer.pp_input_buffer
     ~gen:(fun ~host_funcs:_ ~module_reg:_ -> Ast_generators.input_buffer_gen)
-    ~encoding:(fun ~host_funcs:_ ~module_reg:_ ->
-      Wasm_encoding.input_buffer_encoding)
+    ~encoding:(fun ~host_funcs:_ -> Wasm_encoding.input_buffer_encoding)
 
 (** Test serialize/deserialize values and compare trees. *)
 let test_values_tree () =
@@ -199,8 +180,7 @@ let test_values_tree () =
     ~pp:(Format.pp_print_list Ast_printer.pp_value)
     ~gen:(fun ~host_funcs:_ ~module_reg:_ ->
       QCheck2.Gen.list Ast_generators.value_gen)
-    ~encoding:(fun ~host_funcs:_ ~module_reg ->
-      Wasm_encoding.(values_encoding ~module_reg))
+    ~encoding:(fun ~host_funcs:_ -> Wasm_encoding.values_encoding)
 
 (** Test serialize/deserialize administrative instructions and compare trees. *)
 let test_admin_instr_tree () =
diff --git a/src/lib_scoru_wasm/wasm_encoding.ml b/src/lib_scoru_wasm/wasm_encoding.ml
index e99ad353b7b0..6e501300a647 100644
--- a/src/lib_scoru_wasm/wasm_encoding.ml
+++ b/src/lib_scoru_wasm/wasm_encoding.ml
@@ -453,14 +453,13 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
             "type_result"
             (value [] Interpreter_encodings.Types.value_type_encoding)))
 
-  let module_ref_encoding ~module_reg =
+  let module_ref_encoding =
     conv
-      (fun key ->
-        Instance.{key = Module_key key; registry = Lazy.force module_reg})
-      (fun Instance.{key = Module_key key; _} -> key)
+      (fun key -> Instance.Module_key key)
+      (fun (Instance.Module_key key) -> key)
       (value [] Data_encoding.string)
 
-  let function_encoding ~module_reg =
+  let function_encoding =
     tagged_union
       string_tag
       [
@@ -479,7 +478,7 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
           (tup5
              ~flatten:false
              function_type_encoding
-             (scope ["module"] (module_ref_encoding ~module_reg))
+             (scope ["module"] module_ref_encoding)
              (value ["ftype"] Interpreter_encodings.Ast.var_encoding)
              (lazy_vector_encoding
                 "locals"
@@ -497,13 +496,13 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
             Func.AstFunc (type_, module_, func));
       ]
 
-  let value_ref_encoding ~module_reg =
+  let value_ref_encoding =
     tagged_union
       string_tag
       [
         case
           "FuncRef"
-          (function_encoding ~module_reg)
+          function_encoding
           (fun val_ref ->
             match val_ref with
             | Instance.FuncRef func_inst -> Some func_inst
@@ -521,7 +520,7 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
           (fun v -> Values.NullRef v);
       ]
 
-  let value_encoding ~module_reg =
+  let value_encoding =
     tagged_union
       string_tag
       [
@@ -537,12 +536,12 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
           (fun v -> Values.Vec v);
         case
           "RefType"
-          (value_ref_encoding ~module_reg)
+          value_ref_encoding
           (function Values.Ref r -> Some r | _ -> None)
           (fun r -> Values.Ref r);
       ]
 
-  let values_encoding ~module_reg = list_encoding (value_encoding ~module_reg)
+  let values_encoding = list_encoding value_encoding
 
   let memory_encoding =
     conv
@@ -558,7 +557,7 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
          (value_option ["max"] Data_encoding.int32)
          (scope ["chunks"] chunked_byte_vector))
 
-  let table_encoding ~module_reg =
+  let table_encoding =
     conv
       (fun (min, max, vector, ref_type) ->
         let table_type = Types.TableType ({min; max}, ref_type) in
@@ -570,10 +569,10 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
          ~flatten:false
          (value ["min"] Data_encoding.int32)
          (value_option ["max"] Data_encoding.int32)
-         (lazy_vector_encoding "refs" (value_ref_encoding ~module_reg))
+         (lazy_vector_encoding "refs" value_ref_encoding)
          (value ["ref-type"] Interpreter_encodings.Types.ref_type_encoding))
 
-  let global_encoding ~module_reg =
+  let global_encoding =
     conv
       (fun (type_, value) ->
         let ty = Types.GlobalType (Values.type_of_value value, type_) in
@@ -585,41 +584,38 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
       (tup2
          ~flatten:false
          (value ["type"] Interpreter_encodings.Types.mutability_encoding)
-         (scope ["value"] (value_encoding ~module_reg)))
+         (scope ["value"] value_encoding))
 
   let memory_instance_encoding = lazy_vector_encoding "memories" memory_encoding
 
-  let table_vector_encoding ~module_reg =
-    lazy_vector_encoding "tables" (table_encoding ~module_reg)
+  let table_vector_encoding = lazy_vector_encoding "tables" table_encoding
 
-  let global_vector_encoding ~module_reg =
-    lazy_vector_encoding "globals" (global_encoding ~module_reg)
+  let global_vector_encoding = lazy_vector_encoding "globals" global_encoding
 
   let data_label_ref_encoding =
     conv (fun x -> ref x) (fun r -> !r) data_label_encoding
 
-  let function_vector_encoding ~module_reg =
-    lazy_vector_encoding "functions" (function_encoding ~module_reg)
+  let function_vector_encoding =
+    lazy_vector_encoding "functions" function_encoding
 
   let function_type_vector_encoding =
     lazy_vector_encoding "types" function_type_encoding
 
-  let value_ref_vector_encoding ~module_reg =
-    lazy_vector_encoding "refs" (value_ref_encoding ~module_reg)
+  let value_ref_vector_encoding = lazy_vector_encoding "refs" value_ref_encoding
 
-  let extern_map_encoding ~module_reg =
+  let extern_map_encoding =
     NameMap.lazy_map
       (tagged_union
          string_tag
          [
            case
              "ExternFunc"
-             (function_encoding ~module_reg)
+             function_encoding
              (function Instance.ExternFunc x -> Some x | _ -> None)
              (fun x -> Instance.ExternFunc x);
            case
              "ExternTable"
-             (table_encoding ~module_reg)
+             table_encoding
              (function Instance.ExternTable x -> Some x | _ -> None)
              (fun x -> Instance.ExternTable x);
            case
@@ -629,18 +625,15 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
              (fun x -> Instance.ExternMemory x);
            case
              "ExternGlobal"
-             (global_encoding ~module_reg)
+             global_encoding
              (function Instance.ExternGlobal x -> Some x | _ -> None)
              (fun x -> Instance.ExternGlobal x);
          ])
 
-  let value_ref_vector_vector_encoding ~module_reg =
+  let value_ref_vector_vector_encoding =
     lazy_vector_encoding
       "elements"
-      (conv
-         (fun x -> ref x)
-         (fun r -> !r)
-         (value_ref_vector_encoding ~module_reg))
+      (conv (fun x -> ref x) (fun r -> !r) value_ref_vector_encoding)
 
   let data_instance_encoding =
     lazy_vector_encoding "datas" data_label_ref_encoding
@@ -659,7 +652,7 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
       (fun {blocks; datas} -> (blocks, datas))
       (tup2 ~flatten:false block_table_encoding datas_table_encoding)
 
-  let module_instance_encoding ~module_reg =
+  let module_instance_encoding =
     conv
       (fun ( types,
              funcs,
@@ -704,37 +697,32 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
       (tup9
          ~flatten:false
          function_type_vector_encoding
-         (function_vector_encoding ~module_reg)
-         (table_vector_encoding ~module_reg)
+         function_vector_encoding
+         table_vector_encoding
          memory_instance_encoding
-         (global_vector_encoding ~module_reg)
-         (extern_map_encoding ~module_reg)
-         (value_ref_vector_vector_encoding ~module_reg)
+         global_vector_encoding
+         extern_map_encoding
+         value_ref_vector_vector_encoding
          data_instance_encoding
          allocations_encoding)
 
   let module_instances_encoding =
-    with_self_reference (fun module_reg ->
-        conv
-          Instance.ModuleMap.of_immutable
-          Instance.ModuleMap.snapshot
-          (scope
-             ["modules"]
-             (ModuleMap.lazy_map (module_instance_encoding ~module_reg))))
-
-  let frame_encoding ~module_reg =
-    let locals_encoding =
-      list_encoding @@ conv ref ( ! ) @@ value_encoding ~module_reg
-    in
+    conv
+      Instance.ModuleMap.of_immutable
+      Instance.ModuleMap.snapshot
+      (scope ["modules"] (ModuleMap.lazy_map module_instance_encoding))
+
+  let frame_encoding =
+    let locals_encoding = list_encoding @@ conv ref ( ! ) @@ value_encoding in
     conv
       (fun (inst, locals) -> Eval.{inst; locals})
       (fun Eval.{inst; locals} -> (inst, locals))
       (tup2
          ~flatten:true
-         (scope ["module"] (module_ref_encoding ~module_reg))
+         (scope ["module"] module_ref_encoding)
          (scope ["locals"] locals_encoding))
 
-  let rec admin_instr'_encoding ~module_reg =
+  let rec admin_instr'_encoding () =
     let open Eval in
     tagged_union
       string_tag
@@ -755,12 +743,12 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
           (fun x -> Plain x);
         case
           "Refer"
-          (value_ref_encoding ~module_reg)
+          value_ref_encoding
           (function Refer x -> Some x | _ -> None)
           (fun x -> Refer x);
         case
           "Invoke"
-          (function_encoding ~module_reg)
+          function_encoding
           (function Invoke x -> Some x | _ -> None)
           (fun x -> Invoke x);
         case
@@ -770,15 +758,12 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
           (fun x -> Trapping x);
         case
           "Returning"
-          (values_encoding ~module_reg)
+          values_encoding
           (function Returning x -> Some x | _ -> None)
           (fun x -> Returning x);
         case
           "Breaking"
-          (tup2
-             ~flatten:false
-             (value [] Data_encoding.int32)
-             (values_encoding ~module_reg))
+          (tup2 ~flatten:false (value [] Data_encoding.int32) values_encoding)
           (function
             | Breaking (index, values) -> Some (index, values) | _ -> None)
           (fun (index, values) -> Breaking (index, values));
@@ -788,8 +773,8 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
              ~flatten:false
              (value [] Data_encoding.int32)
              (list_encoding instruction_encoding)
-             (values_encoding ~module_reg)
-             (list_encoding (admin_instr_encoding ~module_reg)))
+             values_encoding
+             (list_encoding (admin_instr_encoding ())))
           (function
             | Label (index, final_instrs, (values, instrs)) ->
                 Some (index, final_instrs, values, instrs)
@@ -801,9 +786,9 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
           (tup4
              ~flatten:false
              (value [] Data_encoding.int32)
-             (frame_encoding ~module_reg)
-             (values_encoding ~module_reg)
-             (list_encoding (admin_instr_encoding ~module_reg)))
+             frame_encoding
+             values_encoding
+             (list_encoding (admin_instr_encoding ())))
           (function
             | Frame (index, frame, (values, instrs)) ->
                 Some (index, frame, values, instrs)
@@ -812,11 +797,13 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
             Frame (index, frame, (values, instrs)));
       ]
 
-  and admin_instr_encoding ~module_reg =
+  and admin_instr_encoding () =
     conv
       Source.(at no_region)
       Source.(fun x -> x.it)
-      (delayed @@ fun () -> admin_instr'_encoding ~module_reg)
+      (delayed admin_instr'_encoding)
+
+  let admin_instr_encoding = admin_instr_encoding ()
 
   let input_buffer_message_encoding =
     conv_lwt
@@ -857,7 +844,7 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
                input_buffer_message_encoding))
          (value ["num-messages"] Data_encoding.z))
 
-  let config_encoding ~host_funcs ~module_reg =
+  let config_encoding ~host_funcs =
     conv
       (fun (frame, input, instrs, values, budget) ->
         Eval.{frame; input; code = (values, instrs); host_funcs; budget})
@@ -865,11 +852,9 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
         (frame, input, instrs, values, budget))
       (tup5
          ~flatten:true
-         (scope ["frame"] (frame_encoding ~module_reg))
+         (scope ["frame"] frame_encoding)
          (scope ["input"] input_buffer_encoding)
-         (scope
-            ["instructions"]
-            (list_encoding (admin_instr_encoding ~module_reg)))
-         (scope ["values"] (values_encoding ~module_reg))
+         (scope ["instructions"] (list_encoding admin_instr_encoding))
+         (scope ["values"] values_encoding)
          (value ["budget"] Data_encoding.int31))
 end
diff --git a/src/lib_scoru_wasm/wasm_encoding.mli b/src/lib_scoru_wasm/wasm_encoding.mli
index 54ae92015932..9155161f6f40 100644
--- a/src/lib_scoru_wasm/wasm_encoding.mli
+++ b/src/lib_scoru_wasm/wasm_encoding.mli
@@ -36,50 +36,37 @@ module Make (M : Tree_encoding.S) : sig
 
   val instruction_encoding : Ast.instr t
 
-  val function_encoding :
-    module_reg:Instance.module_reg Lazy.t -> Instance.func_inst t
+  val function_encoding : Instance.func_inst t
 
-  val value_ref_encoding :
-    module_reg:Instance.module_reg Lazy.t -> Values.ref_ t
+  val value_ref_encoding : Values.ref_ t
 
-  val value_encoding : module_reg:Instance.module_reg Lazy.t -> Values.value t
+  val value_encoding : Values.value t
 
-  val values_encoding :
-    module_reg:Instance.module_reg Lazy.t -> Values.value list t
+  val values_encoding : Values.value list t
 
   val memory_encoding : Partial_memory.memory t
 
-  val table_encoding :
-    module_reg:Instance.module_reg Lazy.t -> Partial_table.table t
+  val table_encoding : Partial_table.table t
 
-  val global_encoding : module_reg:Instance.module_reg Lazy.t -> Global.global t
+  val global_encoding : Global.global t
 
   val memory_instance_encoding : Partial_memory.memory Instance.Vector.t t
 
-  val table_vector_encoding :
-    module_reg:Instance.module_reg Lazy.t ->
-    Partial_table.table Instance.Vector.t t
+  val table_vector_encoding : Partial_table.table Instance.Vector.t t
 
-  val global_vector_encoding :
-    module_reg:Instance.module_reg Lazy.t -> Global.global Instance.Vector.t t
+  val global_vector_encoding : Global.global Instance.Vector.t t
 
   val data_label_ref_encoding : Ast.data_label ref t
 
-  val function_vector_encoding :
-    module_reg:Instance.module_reg Lazy.t ->
-    Instance.func_inst Instance.Vector.t t
+  val function_vector_encoding : Instance.func_inst Instance.Vector.t t
 
   val function_type_vector_encoding : Types.func_type Instance.Vector.t t
 
-  val value_ref_vector_encoding :
-    module_reg:Instance.module_reg Lazy.t -> Values.ref_ Instance.Vector.t t
+  val value_ref_vector_encoding : Values.ref_ Instance.Vector.t t
 
-  val extern_map_encoding :
-    module_reg:Instance.module_reg Lazy.t ->
-    Instance.extern Instance.NameMap.t t
+  val extern_map_encoding : Instance.extern Instance.NameMap.t t
 
   val value_ref_vector_vector_encoding :
-    module_reg:Instance.module_reg Lazy.t ->
     Values.ref_ Instance.Vector.t ref Instance.Vector.t t
 
   val block_table_encoding : Ast.block_table t
@@ -88,20 +75,15 @@ module Make (M : Tree_encoding.S) : sig
 
   val allocations_encoding : Ast.allocations t
 
-  val module_instance_encoding :
-    module_reg:Instance.module_reg Lazy.t -> Instance.module_inst t
+  val module_instance_encoding : Instance.module_inst t
 
   val module_instances_encoding : Instance.module_reg t
 
   val input_buffer_encoding : Input_buffer.t t
 
-  val admin_instr_encoding :
-    module_reg:Instance.module_reg Lazy.t -> Eval.admin_instr t
+  val admin_instr_encoding : Eval.admin_instr t
 
-  val frame_encoding : module_reg:Instance.module_reg Lazy.t -> Eval.frame t
+  val frame_encoding : Eval.frame t
 
-  val config_encoding :
-    host_funcs:Host_funcs.registry ->
-    module_reg:Instance.module_reg Lazy.t ->
-    Eval.config t
+  val config_encoding : host_funcs:Host_funcs.registry -> Eval.config t
 end
diff --git a/src/lib_scoru_wasm/wasm_pvm.ml b/src/lib_scoru_wasm/wasm_pvm.ml
index a946fd0a5cf7..9dc694ba488f 100644
--- a/src/lib_scoru_wasm/wasm_pvm.ml
+++ b/src/lib_scoru_wasm/wasm_pvm.ml
@@ -31,6 +31,7 @@ type pvm_state = {
   kernel : Lazy_containers.Chunked_byte_vector.Lwt.t;
   current_tick : Z.t;
   last_input_info : Wasm_pvm_sig.input_info option;
+  module_reg : Wasm.Instance.module_reg;
   tick : tick_state;
 }
 
@@ -47,7 +48,7 @@ module Make (T : Tree_encoding.TREE) :
       Host_funcs.register_host_funcs registry ;
       registry
 
-    let tick_state_encoding ~module_reg =
+    let tick_state_encoding =
       let open Tree_encoding in
       tagged_union
         ~default:Decode
@@ -60,26 +61,25 @@ module Make (T : Tree_encoding.TREE) :
             (fun () -> Decode);
           case
             "eval"
-            (Wasm_encoding.config_encoding
-               ~host_funcs
-               ~module_reg:(Lazy.from_val module_reg))
+            (Wasm_encoding.config_encoding ~host_funcs)
             (function Eval eval_config -> Some eval_config | _ -> None)
             (fun eval_config -> Eval eval_config);
         ]
 
-    let pvm_state_encoding ~module_reg =
+    let pvm_state_encoding =
       let open Tree_encoding in
       conv
-        (fun (current_tick, kernel, last_input_info, tick) ->
-          {current_tick; kernel; last_input_info; tick})
-        (fun {current_tick; kernel; last_input_info; tick} ->
-          (current_tick, kernel, last_input_info, tick))
-        (tup4
+        (fun (current_tick, kernel, last_input_info, tick, module_reg) ->
+          {current_tick; kernel; last_input_info; tick; module_reg})
+        (fun {current_tick; kernel; last_input_info; tick; module_reg} ->
+          (current_tick, kernel, last_input_info, tick, module_reg))
+        (tup5
            ~flatten:true
            (value ~default:Z.zero ["wasm"; "current_tick"] Data_encoding.n)
            (scope ["durable"; "kernel"; "boot.wasm"] chunked_byte_vector)
            (value_option ["wasm"; "input"] Wasm_pvm_sig.input_info_encoding)
-           (scope ["wasm"] (tick_state_encoding ~module_reg)))
+           (scope ["wasm"] tick_state_encoding)
+           (scope ["modules"] Wasm_encoding.module_instances_encoding))
 
     let status_encoding =
       Tree_encoding.value ["input"; "consuming"] Data_encoding.bool
@@ -92,21 +92,24 @@ module Make (T : Tree_encoding.TREE) :
        kernel to expose a function named [kernel_next]. *)
     let wasm_entrypoint = "kernel_next"
 
-    let next_state ~module_reg state =
+    let next_state state =
       let open Lwt_syntax in
       match state.tick with
       | Decode ->
           let* ast_module =
             Wasm.Decode.decode ~name:wasm_main_module_name ~bytes:state.kernel
           in
-          let self =
-            Wasm.Instance.alloc_module_ref
-              (Wasm.Instance.Module_key wasm_main_module_name)
-              module_reg
-          in
+          let self = Wasm.Instance.Module_key wasm_main_module_name in
           (* The module instance is registered in [self] that contains the
              module registry, why we can ignore the result here. *)
-          let* _module_inst = Wasm.Eval.init ~self host_funcs ast_module [] in
+          let* _module_inst =
+            Wasm.Eval.init
+              ~module_reg:state.module_reg
+              ~self
+              host_funcs
+              ast_module
+              []
+          in
           let eval_config = Wasm.Eval.config host_funcs self [] [] in
           Lwt.return {state with tick = Eval eval_config}
       | Eval ({Wasm.Eval.frame; code; _} as eval_config) -> (
@@ -115,7 +118,9 @@ module Make (T : Tree_encoding.TREE) :
               (* We have an empty set of admin instructions so we create one
                  that invokes the main function. *)
               let* module_inst =
-                Wasm.Instance.ModuleMap.get wasm_main_module_name module_reg
+                Wasm.Instance.ModuleMap.get
+                  wasm_main_module_name
+                  state.module_reg
               in
               let* main_name =
                 Wasm.Instance.Vector.to_list @@ Wasm.Utf8.decode wasm_entrypoint
@@ -154,7 +159,7 @@ module Make (T : Tree_encoding.TREE) :
               Lwt.return {state with tick = Eval eval_config}
           | _ ->
               (* Continue execution. *)
-              let* eval_config = Wasm.Eval.step eval_config in
+              let* eval_config = Wasm.Eval.step state.module_reg eval_config in
               Lwt.return {state with tick = Eval eval_config})
 
     let module_reg_encoding =
@@ -162,31 +167,16 @@ module Make (T : Tree_encoding.TREE) :
         ["module-registry"]
         Wasm_encoding.module_instances_encoding
 
-    let module_reg_from_tree tree =
-      let open Lwt_syntax in
-      try
-        let* module_reg = Tree_encoding.decode module_reg_encoding tree in
-        return (Some module_reg)
-      with _ -> return None
-
-    let decode_state tree =
-      let open Lwt_syntax in
-      (* Try to decode the module registry. *)
-      let* module_reg_opt = module_reg_from_tree tree in
-      let module_reg =
-        Option.value_f ~default:Wasm.Instance.ModuleMap.create module_reg_opt
-      in
-      let+ state = Tree_encoding.decode (pvm_state_encoding ~module_reg) tree in
-      (state, module_reg)
-
     let compute_step tree =
       let open Lwt_syntax in
-      let* state, module_reg = decode_state tree in
-      let* state = next_state state ~module_reg in
+      let* state = Tree_encoding.decode pvm_state_encoding tree in
+      let* state = next_state state in
       let state = {state with current_tick = Z.succ state.current_tick} in
       (* Write the module registry to the tree in case it did not exist
          before. *)
-      let* tree = Tree_encoding.encode module_reg_encoding module_reg tree in
+      let* tree =
+        Tree_encoding.encode module_reg_encoding state.module_reg tree
+      in
       let want_more_input =
         match state.tick with
         | Eval {code = _, []; _} ->
@@ -196,7 +186,7 @@ module Make (T : Tree_encoding.TREE) :
         | _ -> false
       in
       let* tree = Tree_encoding.encode status_encoding want_more_input tree in
-      Tree_encoding.encode (pvm_state_encoding ~module_reg) state tree
+      Tree_encoding.encode pvm_state_encoding state tree
 
     let get_output _ _ = Lwt.return ""
 
@@ -259,7 +249,7 @@ module Make (T : Tree_encoding.TREE) :
       let level = Int32.to_string raw_level in
       let id = Z.to_string message_counter in
       let* current_tick = Tree_encoding.decode current_tick_encoding tree in
-      let* state, module_reg = decode_state tree in
+      let* state = Tree_encoding.decode pvm_state_encoding tree in
       let* () =
         match state.tick with
         | Eval config ->
@@ -281,9 +271,7 @@ module Make (T : Tree_encoding.TREE) :
             *)
             assert false
       in
-      let* tree =
-        Tree_encoding.encode (pvm_state_encoding ~module_reg) state tree
-      in
+      let* tree = Tree_encoding.encode pvm_state_encoding state tree in
       let* tree =
         Tree_encoding.encode current_tick_encoding (Z.succ current_tick) tree
       in
diff --git a/src/lib_webassembly/bin/script/run.ml b/src/lib_webassembly/bin/script/run.ml
index 84142a9529c5..07f35009918b 100644
--- a/src/lib_webassembly/bin/script/run.ml
+++ b/src/lib_webassembly/bin/script/run.ml
@@ -333,10 +333,10 @@ let modules : Ast.module_ Map.t ref = ref Map.empty
 (* NOTE: See [instantiate_module] below on why this exists. *)
 let unnamed_instance_counter = ref 0
 
-(* NOTE: See [instantiate_module] below on why this exists. *)
-let unnamed_instances : Instance.module_reg = Instance.ModuleMap.create ()
-
-let instances : Instance.module_reg = Instance.ModuleMap.create ()
+let instances : Instance.module_reg =
+  let module_reg = Instance.ModuleMap.create () in
+  Instance.ModuleMap.set "__empty" Instance.empty_module_inst module_reg ;
+  module_reg
 
 let registry : Instance.module_inst Map.t ref = ref Map.empty
 
@@ -358,11 +358,10 @@ let instantiate_module x_opt m imports =
         *)
         let index = !unnamed_instance_counter in
         unnamed_instance_counter := !unnamed_instance_counter + 1 ;
-        Instance.
-          {registry = unnamed_instances; key = Module_key (Int.to_string index)}
-    | Some name -> Instance.{registry = instances; key = Module_key name.it}
+        Instance.Module_key (Printf.sprintf "__unnamed_%i" index)
+    | Some name -> Instance.Module_key name.it
   in
-  Eval.init ~self host_funcs_registry m imports
+  Eval.init ~module_reg:instances ~self host_funcs_registry m imports
 
 let bind_lazy module_reg name instance =
   Option.iter
@@ -440,7 +439,12 @@ let run_action act : Values.value list Lwt.t =
             vs
             ins_l ;
           let+ result =
-            Eval.invoke host_funcs_registry f (List.map (fun v -> v.it) vs)
+            Eval.invoke
+              ~module_reg:instances
+              ~caller:(Module_key "__empty")
+              host_funcs_registry
+              f
+              (List.map (fun v -> v.it) vs)
           in
           result
       | Some _ -> Assert.error act.at "export is not a function"
diff --git a/src/lib_webassembly/exec/eval.ml b/src/lib_webassembly/exec/eval.ml
index 24f4067afd6b..f384c4083b0e 100644
--- a/src/lib_webassembly/exec/eval.ml
+++ b/src/lib_webassembly/exec/eval.ml
@@ -55,7 +55,7 @@ let numeric_error at = function
 
 type 'a stack = 'a list
 
-type frame = {inst : module_ref; locals : value ref list}
+type frame = {inst : module_key; locals : value ref list}
 
 type code = value stack * admin_instr list
 
@@ -166,40 +166,40 @@ let drop n (vs : 'a stack) at =
  *   c : config
  *)
 
-let mem_oob frame x i n =
-  let* inst = resolve_module_ref frame.inst in
+let mem_oob module_reg frame x i n =
+  let* inst = resolve_module_ref module_reg frame.inst in
   let+ mem = memory inst x in
   I64.gt_u
     (I64.add (I64_convert.extend_i32_u i) (I64_convert.extend_i32_u n))
     (Memory.bound mem)
 
-let data_oob frame x i n =
-  let* inst = resolve_module_ref frame.inst in
+let data_oob module_reg frame x i n =
+  let* inst = resolve_module_ref module_reg frame.inst in
   let* data_label = data inst x in
   let+ data = Ast.get_data !data_label inst.allocations.datas in
   I64.gt_u
     (I64.add (I64_convert.extend_i32_u i) (I64_convert.extend_i32_u n))
     (Chunked_byte_vector.Lwt.length data)
 
-let table_oob frame x i n =
-  let* inst = resolve_module_ref frame.inst in
+let table_oob module_reg frame x i n =
+  let* inst = resolve_module_ref module_reg frame.inst in
   let+ tbl = table inst x in
   I64.gt_u
     (I64.add (I64_convert.extend_i32_u i) (I64_convert.extend_i32_u n))
     (I64_convert.extend_i32_u (Table.size tbl))
 
-let elem_oob frame x i n =
-  let* inst = resolve_module_ref frame.inst in
+let elem_oob module_reg frame x i n =
+  let* inst = resolve_module_ref module_reg frame.inst in
   let+ elem = elem inst x in
   I64.gt_u
     (I64.add (I64_convert.extend_i32_u i) (I64_convert.extend_i32_u n))
     (Int64.of_int32 (Instance.Vector.num_elements !elem))
 
-let rec step (c : config) : config Lwt.t =
+let rec step (module_reg : module_reg) (c : config) : config Lwt.t =
   let {frame; code = vs, es; _} = c in
   match es with
   | {it = From_block (Block_label b, i); at} :: es ->
-      let* inst = resolve_module_ref frame.inst in
+      let* inst = resolve_module_ref module_reg frame.inst in
       let* block = Vector.get b inst.allocations.blocks in
       let length = Vector.num_elements block in
       if i = length then Lwt.return {c with code = (vs, es)}
@@ -210,11 +210,11 @@ let rec step (c : config) : config Lwt.t =
             ( Plain instr.it @@ instr.at,
               {it = From_block (Block_label b, Int32.succ i); at} :: es )
         in
-        step_resolved c frame vs e es
-  | e :: es -> step_resolved c frame vs e es
+        step_resolved module_reg c frame vs e es
+  | e :: es -> step_resolved module_reg c frame vs e es
   | [] -> Lwt.return c
 
-and step_resolved (c : config) frame vs e es : config Lwt.t =
+and step_resolved module_reg (c : config) frame vs e es : config Lwt.t =
   let+ vs', es' =
     match (e.it, vs) with
     | From_block _, _ -> assert false (* resolved by [step] *)
@@ -224,7 +224,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
             Lwt.return (vs, [Trapping "unreachable executed" @@ e.at])
         | Nop, vs -> Lwt.return (vs, [])
         | Block (bt, es'), vs ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let+ (FuncType (ts1, ts2)) = block_type inst bt in
             let n1 = Lazy_vector.LwtInt32Vector.num_elements ts1 in
             let n2 = Lazy_vector.LwtInt32Vector.num_elements ts2 in
@@ -233,7 +233,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
               [Label (n2, [], (args, [From_block (es', 0l) @@ e.at])) @@ e.at]
             )
         | Loop (bt, es'), vs ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let+ (FuncType (ts1, ts2)) = block_type inst bt in
             let n1 = Lazy_vector.LwtInt32Vector.num_elements ts1 in
             let args, vs' = (take n1 vs e.at, drop n1 vs e.at) in
@@ -257,11 +257,11 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
               else (vs', [Plain (Br (Lib.List32.nth xs i)) @@ e.at]))
         | Return, vs -> Lwt.return ([], [Returning vs @@ e.at])
         | Call x, vs ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let+ func = func inst x in
             (vs, [Invoke func @@ e.at])
         | CallIndirect (x, y), Num (I32 i) :: vs ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let* func = func_ref inst x i e.at and* type_ = type_ inst y in
             let+ check_eq = Types.func_types_equal type_ (Func.type_of func) in
             if not check_eq then
@@ -280,14 +280,14 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
               (local frame x := v ;
                (v :: vs', []))
         | GlobalGet x, vs ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let+ glob = global inst x in
             let value = Global.load glob in
             (value :: vs, [])
         | GlobalSet x, v :: vs' ->
             Lwt.catch
               (fun () ->
-                let* inst = resolve_module_ref frame.inst in
+                let* inst = resolve_module_ref module_reg frame.inst in
                 let+ glob = global inst x in
                 Global.store glob v ;
                 (vs', []))
@@ -300,7 +300,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
         | TableGet x, Num (I32 i) :: vs' ->
             Lwt.catch
               (fun () ->
-                let* inst = resolve_module_ref frame.inst in
+                let* inst = resolve_module_ref module_reg frame.inst in
                 let* tbl = table inst x in
                 let+ value = Table.load tbl i in
                 (Ref value :: vs', []))
@@ -309,18 +309,18 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
         | TableSet x, Ref r :: Num (I32 i) :: vs' ->
             Lwt.catch
               (fun () ->
-                let* inst = resolve_module_ref frame.inst in
+                let* inst = resolve_module_ref module_reg frame.inst in
                 let+ tbl = table inst x in
                 Table.store tbl i r ;
                 (vs', []))
               (fun exn ->
                 Lwt.return (vs', [Trapping (table_error e.at exn) @@ e.at]))
         | TableSize x, vs ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let+ tbl = table inst x in
             (Num (I32 (Table.size tbl)) :: vs, [])
         | TableGrow x, Num (I32 delta) :: Ref r :: vs' ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let+ tab = table inst x in
             let old_size = Table.size tab in
             let result =
@@ -332,7 +332,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
             in
             (Num (I32 result) :: vs', [])
         | TableFill x, Num (I32 n) :: Ref r :: Num (I32 i) :: vs' ->
-            let+ oob = table_oob frame x i n in
+            let+ oob = table_oob module_reg frame x i n in
             if oob then (vs', [Trapping (table_error e.at Table.Bounds) @@ e.at])
             else if n = 0l then (vs', [])
             else
@@ -350,8 +350,8 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
                     Plain (TableFill x);
                   ] )
         | TableCopy (x, y), Num (I32 n) :: Num (I32 s) :: Num (I32 d) :: vs' ->
-            let+ oob_d = table_oob frame x d n
-            and+ oob_s = table_oob frame y s n in
+            let+ oob_d = table_oob module_reg frame x d n
+            and+ oob_s = table_oob module_reg frame y s n in
             if oob_d || oob_s then
               (vs', [Trapping (table_error e.at Table.Bounds) @@ e.at])
             else if n = 0l then (vs', [])
@@ -385,14 +385,14 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
                     Plain (TableSet x);
                   ] )
         | TableInit (x, y), Num (I32 n) :: Num (I32 s) :: Num (I32 d) :: vs' ->
-            let* oob_d = table_oob frame x d n in
-            let* oob_s = elem_oob frame y s n in
+            let* oob_d = table_oob module_reg frame x d n in
+            let* oob_s = elem_oob module_reg frame y s n in
             if oob_d || oob_s then
               Lwt.return
                 (vs', [Trapping (table_error e.at Table.Bounds) @@ e.at])
             else if n = 0l then Lwt.return (vs', [])
             else
-              let* inst = resolve_module_ref frame.inst in
+              let* inst = resolve_module_ref module_reg frame.inst in
               let* seg = elem inst y in
               let+ value = Instance.Vector.get s !seg in
               ( vs',
@@ -412,12 +412,12 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
                     Plain (TableInit (x, y));
                   ] )
         | ElemDrop x, vs ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let+ seg = elem inst x in
             seg := Instance.Vector.create 0l ;
             (vs, [])
         | Load {offset; ty; pack; _}, Num (I32 i) :: vs' ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let* mem = memory inst (0l @@ e.at) in
             Lwt.catch
               (fun () ->
@@ -431,7 +431,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
               (fun exn ->
                 Lwt.return (vs', [Trapping (memory_error e.at exn) @@ e.at]))
         | Store {offset; pack; _}, Num n :: Num (I32 i) :: vs' ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let* mem = memory inst (0l @@ e.at) in
             Lwt.catch
               (fun () ->
@@ -444,7 +444,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
               (fun exn ->
                 Lwt.return (vs', [Trapping (memory_error e.at exn) @@ e.at]))
         | VecLoad {offset; ty; pack; _}, Num (I32 i) :: vs' ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let* mem = memory inst (0l @@ e.at) in
             Lwt.catch
               (fun () ->
@@ -458,7 +458,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
               (fun exn ->
                 Lwt.return (vs', [Trapping (memory_error e.at exn) @@ e.at]))
         | VecStore {offset; _}, Vec v :: Num (I32 i) :: vs' ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let* mem = memory inst (0l @@ e.at) in
             Lwt.catch
               (fun () ->
@@ -468,7 +468,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
                 Lwt.return (vs', [Trapping (memory_error e.at exn) @@ e.at]))
         | ( VecLoadLane ({offset; ty; pack; _}, j),
             Vec (V128 v) :: Num (I32 i) :: vs' ) ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let* mem = memory inst (0l @@ e.at) in
             Lwt.catch
               (fun () ->
@@ -496,7 +496,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
                 Lwt.return (vs', [Trapping (memory_error e.at exn) @@ e.at]))
         | ( VecStoreLane ({offset; ty; pack; _}, j),
             Vec (V128 v) :: Num (I32 i) :: vs' ) ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let* mem = memory inst (0l @@ e.at) in
             Lwt.catch
               (fun () ->
@@ -533,11 +533,11 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
               (fun exn ->
                 Lwt.return (vs', [Trapping (memory_error e.at exn) @@ e.at]))
         | MemorySize, vs ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let+ mem = memory inst (0l @@ e.at) in
             (Num (I32 (Memory.size mem)) :: vs, [])
         | MemoryGrow, Num (I32 delta) :: vs' ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let+ mem = memory inst (0l @@ e.at) in
             let old_size = Memory.size mem in
             let result =
@@ -550,7 +550,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
             in
             (Num (I32 result) :: vs', [])
         | MemoryFill, Num (I32 n) :: Num k :: Num (I32 i) :: vs' ->
-            let+ oob = mem_oob frame (0l @@ e.at) i n in
+            let+ oob = mem_oob module_reg frame (0l @@ e.at) i n in
             if oob then
               (vs', [Trapping (memory_error e.at Memory.Bounds) @@ e.at])
             else if n = 0l then (vs', [])
@@ -575,8 +575,8 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
                     Plain MemoryFill;
                   ] )
         | MemoryCopy, Num (I32 n) :: Num (I32 s) :: Num (I32 d) :: vs' ->
-            let+ oob_s = mem_oob frame (0l @@ e.at) s n
-            and+ oob_d = mem_oob frame (0l @@ e.at) d n in
+            let+ oob_s = mem_oob module_reg frame (0l @@ e.at) s n
+            and+ oob_d = mem_oob module_reg frame (0l @@ e.at) d n in
             if oob_s || oob_d then
               (vs', [Trapping (memory_error e.at Memory.Bounds) @@ e.at])
             else if n = 0l then (vs', [])
@@ -638,14 +638,14 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
                          });
                   ] )
         | MemoryInit x, Num (I32 n) :: Num (I32 s) :: Num (I32 d) :: vs' ->
-            let* mem_oob = mem_oob frame (0l @@ e.at) d n in
-            let* data_oob = data_oob frame x s n in
+            let* mem_oob = mem_oob module_reg frame (0l @@ e.at) d n in
+            let* data_oob = data_oob module_reg frame x s n in
             if mem_oob || data_oob then
               Lwt.return
                 (vs', [Trapping (memory_error e.at Memory.Bounds) @@ e.at])
             else if n = 0l then Lwt.return (vs', [])
             else
-              let* inst = resolve_module_ref frame.inst in
+              let* inst = resolve_module_ref module_reg frame.inst in
               let* seg = data inst x in
               let* seg = Ast.get_data !seg inst.allocations.datas in
               let+ b =
@@ -672,7 +672,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
                     Plain (MemoryInit x);
                   ] )
         | DataDrop x, vs ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let+ seg = data inst x in
             seg := Data_label 0l ;
             (vs, [])
@@ -683,7 +683,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
               | NullRef _ -> (Num (I32 1l) :: vs', [])
               | _ -> (Num (I32 0l) :: vs', []))
         | RefFunc x, vs' ->
-            let* inst = resolve_module_ref frame.inst in
+            let* inst = resolve_module_ref module_reg frame.inst in
             let+ f = func inst x in
             (Ref (FuncRef f) :: vs', [])
         | Const n, vs -> Lwt.return (Num n.it :: vs, [])
@@ -787,7 +787,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
     | Label (n, es0, (vs', {it = Breaking (k, vs0); at} :: es')), vs ->
         Lwt.return (vs, [Breaking (Int32.sub k 1l, vs0) @@ at])
     | Label (n, es0, code'), vs ->
-        let+ c' = step {c with code = code'} in
+        let+ c' = step module_reg {c with code = code'} in
         (vs, [Label (n, es0, c'.code) @@ e.at])
     | Frame (n, frame', (vs', [])), vs -> Lwt.return (vs' @ vs, [])
     | Frame (n, frame', (vs', {it = Trapping msg; at} :: es')), vs ->
@@ -797,6 +797,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
     | Frame (n, frame', code'), vs ->
         let+ c' =
           step
+            module_reg
             {
               frame = frame';
               code = code';
@@ -837,7 +838,7 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
                 let (Host_funcs.Host_func f) =
                   Host_funcs.lookup ~global_name c.host_funcs
                 in
-                let* inst = resolve_module_ref frame.inst in
+                let* inst = resolve_module_ref module_reg frame.inst in
                 let+ res = f c.input inst (List.rev args) in
                 (List.rev res @ vs', []))
               (function
@@ -845,17 +846,17 @@ and step_resolved (c : config) frame vs e es : config Lwt.t =
   in
   {c with code = (vs', es' @ es)}
 
-let rec eval (c : config) : value stack Lwt.t =
+let rec eval module_reg (c : config) : value stack Lwt.t =
   match c.code with
   | vs, [] -> Lwt.return vs
   | vs, {it = Trapping msg; at} :: _ -> Trap.error at msg
   | vs, es ->
-      let* c = step c in
-      eval c
+      let* c = step module_reg c in
+      eval module_reg c
 
 (* Functions & Constants *)
 
-let invoke ?caller ?(input = Input_buffer.alloc ()) host_funcs
+let invoke ~module_reg ~caller ?(input = Input_buffer.alloc ()) host_funcs
     (func : func_inst) (vs : value list) : value list Lwt.t =
   let at = match func with Func.AstFunc (_, _, f) -> f.at | _ -> no_region in
   let (FuncType (ins, out)) = Func.type_of func in
@@ -870,34 +871,31 @@ let invoke ?caller ?(input = Input_buffer.alloc ()) host_funcs
   let inst =
     match func with
     | Func.AstFunc (_, inst, _) -> inst
-    | Func.HostFunc (_, _) -> (
-        (* We create a fresh module reference to an empty module. [invoke] is
-           called at the top-level where we are not in a module context, yet.
-           Hence we need to create one for this call. *)
-        match caller with None -> empty_module_ref () | Some inst -> inst)
+    | Func.HostFunc (_, _) -> caller
   in
   let c = config ~input host_funcs inst (List.rev vs) [Invoke func @@ at] in
   Lwt.catch
     (fun () ->
-      let+ values = eval c in
+      let+ values = eval module_reg c in
       List.rev values)
     (function
       | Stack_overflow -> Exhaustion.error at "call stack exhausted"
       | exn -> Lwt.fail exn)
 
-let eval_const (inst : module_ref) (const : const) : value Lwt.t =
+let eval_const module_reg (inst : module_key) (const : const) : value Lwt.t =
   let c =
     config (Host_funcs.empty ()) inst [] [From_block (const.it, 0l) @@ const.at]
   in
-  let+ vs = eval c in
+  let+ vs = eval module_reg c in
   match vs with
   | [v] -> v
   | vs -> Crash.error const.at "wrong number of results on stack"
 
 (* Modules *)
 
-let create_func (inst_ref : module_ref) (f : func) : func_inst Lwt.t =
-  let* inst = resolve_module_ref inst_ref in
+let create_func module_reg (inst_ref : module_key) (f : func) : func_inst Lwt.t
+    =
+  let* inst = resolve_module_ref module_reg inst_ref in
   let+ type_ = type_ inst f.it.ftype in
   Func.alloc type_ inst_ref f
 
@@ -910,9 +908,10 @@ let create_memory (inst : module_inst) (mem : memory) : memory_inst =
   let {mtype} = mem.it in
   Memory.alloc mtype
 
-let create_global (inst : module_ref) (glob : global) : global_inst Lwt.t =
+let create_global module_reg (inst : module_key) (glob : global) :
+    global_inst Lwt.t =
   let {gtype; ginit} = glob.it in
-  let+ v = eval_const inst ginit in
+  let+ v = eval_const module_reg inst ginit in
   Global.alloc gtype v
 
 let create_export (inst : module_inst) (ex : export) : export_inst Lwt.t =
@@ -934,7 +933,8 @@ let create_export (inst : module_inst) (ex : export) : export_inst Lwt.t =
   in
   (name, ext)
 
-let create_elem (inst : module_ref) (seg : elem_segment) : elem_inst Lwt.t =
+let create_elem module_reg (inst : module_key) (seg : elem_segment) :
+    elem_inst Lwt.t =
   let {etype; einit; _} = seg.it in
   (* TODO: #3076
      [einit] should be changed to a lazy structure. We want to avoid traversing
@@ -943,7 +943,7 @@ let create_elem (inst : module_ref) (seg : elem_segment) : elem_inst Lwt.t =
   let+ init =
     TzStdLib.List.map_s
       (fun v ->
-        let+ r = eval_const inst v in
+        let+ r = eval_const module_reg inst v in
         as_ref r)
       einit
   in
@@ -1022,8 +1022,8 @@ let run_data (inst : module_inst) i data =
 
 let run_start start = List.map plain [Call start.it.sfunc @@ start.at]
 
-let init ~self host_funcs (m : module_) (exts : extern list) : module_inst Lwt.t
-    =
+let init ~module_reg ~self host_funcs (m : module_) (exts : extern list) :
+    module_inst Lwt.t =
   let open Lwt.Syntax in
   let {
     imports;
@@ -1042,7 +1042,7 @@ let init ~self host_funcs (m : module_) (exts : extern list) : module_inst Lwt.t
   in
 
   (* Initialize as empty module. *)
-  update_module_ref self empty_module_inst ;
+  update_module_ref module_reg self empty_module_inst ;
 
   (* TODO: #3076
 
@@ -1074,7 +1074,7 @@ let init ~self host_funcs (m : module_) (exts : extern list) : module_inst Lwt.t
     | Error () ->
         Link.error m.at "wrong number of imports provided for initialisation"
   in
-  update_module_ref self init_inst0 ;
+  update_module_ref module_reg self init_inst0 ;
 
   let inst0 =
     {
@@ -1087,17 +1087,19 @@ let init ~self host_funcs (m : module_) (exts : extern list) : module_inst Lwt.t
       allocations;
     }
   in
-  update_module_ref self inst0 ;
+  update_module_ref module_reg self inst0 ;
 
-  let* fs = TzStdLib.List.map_s (create_func self) funcs in
+  let* fs = TzStdLib.List.map_s (create_func module_reg self) funcs in
   (* TODO: #3076
      [fs]/[funcs] should be a lazy structure so we can avoid traversing it
      completely. *)
   let* funcs = Vector.concat inst0.funcs (Vector.of_list fs) in
   let inst1 = {inst0 with funcs} in
-  update_module_ref self inst1 ;
+  update_module_ref module_reg self inst1 ;
 
-  let* new_globals = TzStdLib.List.map_s (create_global self) globals in
+  let* new_globals =
+    TzStdLib.List.map_s (create_global module_reg self) globals
+  in
   (* TODO: #3076
      [tables] should be a lazy structure. *)
   let* tables =
@@ -1116,10 +1118,10 @@ let init ~self host_funcs (m : module_) (exts : extern list) : module_inst Lwt.t
      [new_globals]/[globals] should be lazy structures. *)
   let* globals = Vector.concat inst1.globals (Vector.of_list new_globals) in
   let inst2 = {inst1 with tables; memories; globals} in
-  update_module_ref self inst2 ;
+  update_module_ref module_reg self inst2 ;
 
   let* new_exports = TzStdLib.List.map_s (create_export inst2) exports in
-  let* new_elems = TzStdLib.List.map_s (create_elem self) elems in
+  let* new_elems = TzStdLib.List.map_s (create_elem module_reg self) elems in
   let new_datas = List.map (create_data inst2) datas in
   let* exports =
     (* TODO: #3076
@@ -1145,14 +1147,14 @@ let init ~self host_funcs (m : module_) (exts : extern list) : module_inst Lwt.t
         Vector.of_list new_datas;
     }
   in
-  update_module_ref self inst ;
+  update_module_ref module_reg self inst ;
 
   let es_elem = List.concat (Lib.List32.mapi (run_elem inst) elems) in
   let* datas = Lib.List32.mapi_s (run_data inst) datas in
   let es_data = TzStdLib.List.concat datas in
   let es_start = Lib.Option.get (Lib.Option.map run_start start) [] in
   let+ (_ : Values.value stack) =
-    eval (config host_funcs self [] (es_elem @ es_data @ es_start))
+    eval module_reg (config host_funcs self [] (es_elem @ es_data @ es_start))
   in
 
   inst
diff --git a/src/lib_webassembly/exec/eval.mli b/src/lib_webassembly/exec/eval.mli
index 640ecb85743d..3c6068e66855 100644
--- a/src/lib_webassembly/exec/eval.mli
+++ b/src/lib_webassembly/exec/eval.mli
@@ -10,21 +10,23 @@ exception Crash of Source.region * string
 exception Exhaustion of Source.region * string
 
 val init :
-  self:module_ref ->
+  module_reg:module_reg ->
+  self:module_key ->
   Host_funcs.registry ->
   Ast.module_ ->
   extern list ->
   module_inst Lwt.t (* raises Link, Trap *)
 
 val invoke :
-  ?caller:module_ref ->
+  module_reg:module_reg ->
+  caller:module_key ->
   ?input:Input_buffer.t ->
   Host_funcs.registry ->
   func_inst ->
   value list ->
   value list Lwt.t (* raises Trap *)
 
-type frame = {inst : module_ref; locals : value ref list}
+type frame = {inst : module_key; locals : value ref list}
 
 type code = value list * admin_instr list
 
@@ -49,12 +51,12 @@ type config = {
   budget : int; (* to model stack overflow *)
 }
 
-val step : config -> config Lwt.t
+val step : module_reg -> config -> config Lwt.t
 
 val config :
   ?input:input_inst ->
   Host_funcs.registry ->
-  module_ref ->
+  module_key ->
   value list ->
   admin_instr list ->
   config
diff --git a/src/lib_webassembly/runtime/instance.ml b/src/lib_webassembly/runtime/instance.ml
index b7ab3e6a9f22..34be3468a7d2 100644
--- a/src/lib_webassembly/runtime/instance.ml
+++ b/src/lib_webassembly/runtime/instance.ml
@@ -36,7 +36,7 @@ type module_inst = {
   allocations : Ast.allocations;
 }
 
-and func_inst = module_ref Func.t
+and func_inst = module_key Func.t
 
 and table_inst = Table.t
 
@@ -60,8 +60,6 @@ and extern =
 
 and module_reg = module_inst ModuleMap.t
 
-and module_ref = {registry : module_reg; key : module_key}
-
 (* Reference types *)
 
 type Values.ref_ += FuncRef of func_inst
@@ -99,22 +97,10 @@ let empty_module_inst =
     allocations = Ast.empty_allocations ();
   }
 
-let empty_module_ref () =
-  let registry = ModuleMap.create () in
-  let key = "empty" in
-  ModuleMap.set key empty_module_inst registry ;
-  {registry; key = Module_key key}
-
-let alloc_module_ref (Module_key mkey as key) ?(module_inst = empty_module_inst)
-    registry =
-  ModuleMap.set mkey module_inst registry ;
-  {registry; key}
-
-let update_module_ref {registry; key = Module_key key} module_inst =
+let update_module_ref registry (Module_key key) module_inst =
   ModuleMap.set key module_inst registry
 
-let resolve_module_ref {registry; key = Module_key key} =
-  ModuleMap.get key registry
+let resolve_module_ref registry (Module_key key) = ModuleMap.get key registry
 
 let extern_type_of = function
   | ExternFunc func -> ExternFuncType (Func.type_of func)
-- 
GitLab


From 55fdab7473738e258249f0c52d0533f1653b69a1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ole=20Kr=C3=BCger?= <ole.kruger@trili.tech>
Date: Fri, 12 Aug 2022 17:02:01 +0100
Subject: [PATCH 2/4] fixup: Rename module_ref_encoding -> module_key_encoding

---
 src/lib_scoru_wasm/wasm_encoding.ml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/lib_scoru_wasm/wasm_encoding.ml b/src/lib_scoru_wasm/wasm_encoding.ml
index 6e501300a647..dc4afc1d8121 100644
--- a/src/lib_scoru_wasm/wasm_encoding.ml
+++ b/src/lib_scoru_wasm/wasm_encoding.ml
@@ -453,7 +453,7 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
             "type_result"
             (value [] Interpreter_encodings.Types.value_type_encoding)))
 
-  let module_ref_encoding =
+  let module_key_encoding =
     conv
       (fun key -> Instance.Module_key key)
       (fun (Instance.Module_key key) -> key)
@@ -478,7 +478,7 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
           (tup5
              ~flatten:false
              function_type_encoding
-             (scope ["module"] module_ref_encoding)
+             (scope ["module"] module_key_encoding)
              (value ["ftype"] Interpreter_encodings.Ast.var_encoding)
              (lazy_vector_encoding
                 "locals"
@@ -719,7 +719,7 @@ module Make (Tree_encoding : Tree_encoding.S) = struct
       (fun Eval.{inst; locals} -> (inst, locals))
       (tup2
          ~flatten:true
-         (scope ["module"] module_ref_encoding)
+         (scope ["module"] module_key_encoding)
          (scope ["locals"] locals_encoding))
 
   let rec admin_instr'_encoding () =
-- 
GitLab


From b843424c1df76e8da7f037ab7f5ec06dd1966eaa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ole=20Kr=C3=BCger?= <ole.kruger@trili.tech>
Date: Fri, 12 Aug 2022 17:04:07 +0100
Subject: [PATCH 3/4] fixup: Remove with_self_reference from tree-encoding

---
 src/lib_tree_encoding/test/test_encoding.ml | 22 ---------------------
 src/lib_tree_encoding/tree_encoding.ml      | 22 ---------------------
 src/lib_tree_encoding/tree_encoding.mli     |  6 ------
 3 files changed, 50 deletions(-)

diff --git a/src/lib_tree_encoding/test/test_encoding.ml b/src/lib_tree_encoding/test/test_encoding.ml
index fd4ba86ee651..b74490907ccb 100644
--- a/src/lib_tree_encoding/test/test_encoding.ml
+++ b/src/lib_tree_encoding/test/test_encoding.ml
@@ -375,27 +375,6 @@ let test_value_option () =
 
 type cyclic = {name : string; self : unit -> cyclic}
 
-let test_with_self_ref () =
-  let open Tree_encoding in
-  let open Lwt_result_syntax in
-  let enc () =
-    with_self_reference (fun cycle ->
-        conv
-          (fun name -> {name; self = (fun () -> Lazy.force cycle)})
-          (fun {name; _} -> name)
-          (value [] Data_encoding.string))
-  in
-  (* A cycle is a value with a (lazy) self-reference. *)
-  let rec cycle = {name = "Cycle"; self = (fun () -> cycle)} in
-  (* Encode using an encoder and an empty tree. *)
-  let*! empty_tree = empty_tree () in
-  let*! tree = Tree_encoding.encode (enc ()) cycle empty_tree in
-  (* Decode using a new encoder value and the tree from above. *)
-  let*! ({name; self} as cycle) = Tree_encoding.decode (enc ()) tree in
-  assert (name = "Cycle") ;
-  assert (cycle == self ()) ;
-  return_unit
-
 let test_delayed () =
   let open Tree_encoding in
   let open Lwt_result_syntax in
@@ -490,7 +469,6 @@ let tests =
     tztest "Option" `Quick test_option;
     tztest "Value ~default" `Quick test_value_default;
     tztest "Value-option" `Quick test_value_option;
-    tztest "Self ref" `Quick test_with_self_ref;
     tztest "Delayed" `Quick test_delayed;
     tztest "Return" `Quick test_return;
     tztest "Swap vectors" `Quick test_swap_vectors;
diff --git a/src/lib_tree_encoding/tree_encoding.ml b/src/lib_tree_encoding/tree_encoding.ml
index e922bdd5e08c..52f8ad619489 100644
--- a/src/lib_tree_encoding/tree_encoding.ml
+++ b/src/lib_tree_encoding/tree_encoding.ml
@@ -174,8 +174,6 @@ module type S = sig
 
   val option : 'a t -> 'a option t
 
-  val with_self_reference : ('a Lazy.t -> 'a t) -> 'a t
-
   val delayed : (unit -> 'a t) -> 'a t
 end
 
@@ -475,26 +473,6 @@ module Make (T : Tree.S) : S with type tree = T.tree = struct
           (fun () -> None);
       ]
 
-  let with_self_reference f =
-    (* Mutable reference to the current value. *)
-    let current = ref None in
-    (* Sets the current value. *)
-    let set_current value =
-      current := Some value ;
-      value
-    in
-    (* Gets the current value from the ref. This should only be called once
-       the  encoding/decoding steps have already constructed a value and the ref
-       has been updated. *)
-    let get_current () =
-      match !current with
-      | Some value -> value
-      | None -> raise Uninitialized_self_ref
-    in
-    (* Intercepts the encoding and decoding steps to update the reference to the
-       current module. *)
-    conv set_current set_current (f (lazy (get_current ())))
-
   let delayed f =
     let enc = lazy (f ()) in
     let encode =
diff --git a/src/lib_tree_encoding/tree_encoding.mli b/src/lib_tree_encoding/tree_encoding.mli
index 5a29797c60ce..e27f3fbecde8 100644
--- a/src/lib_tree_encoding/tree_encoding.mli
+++ b/src/lib_tree_encoding/tree_encoding.mli
@@ -295,12 +295,6 @@ module type S = sig
       optional values. *)
   val option : 'a t -> 'a option t
 
-  (** [with_self_reference f] creates an encoder that allows accessing the
-      encoded/decoded value itself. It's useful for encoding cyclic
-      data-structures. Here, [f] is a function that takes the (lazy)
-      self-reference as an argument and constructs an encoder. *)
-  val with_self_reference : ('a Lazy.t -> 'a t) -> 'a t
-
   (** [delayed f] produces a tree encoder/decoder that delays evaluation of
       [f ()] until the encoder or decoder is actually needed. This is required
       to allow for directly recursive encoders/decoders. *)
-- 
GitLab


From e22aa3225251fd3260d525b7d9990a7a3078355c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ole=20Kr=C3=BCger?= <ole.kruger@trili.tech>
Date: Fri, 12 Aug 2022 17:20:49 +0100
Subject: [PATCH 4/4] fixup: Remove redundant module_reg_encoding

---
 src/lib_scoru_wasm/wasm_pvm.ml | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/src/lib_scoru_wasm/wasm_pvm.ml b/src/lib_scoru_wasm/wasm_pvm.ml
index 9dc694ba488f..ebf0c4195e98 100644
--- a/src/lib_scoru_wasm/wasm_pvm.ml
+++ b/src/lib_scoru_wasm/wasm_pvm.ml
@@ -162,21 +162,11 @@ module Make (T : Tree_encoding.TREE) :
               let* eval_config = Wasm.Eval.step state.module_reg eval_config in
               Lwt.return {state with tick = Eval eval_config})
 
-    let module_reg_encoding =
-      Tree_encoding.scope
-        ["module-registry"]
-        Wasm_encoding.module_instances_encoding
-
     let compute_step tree =
       let open Lwt_syntax in
       let* state = Tree_encoding.decode pvm_state_encoding tree in
       let* state = next_state state in
       let state = {state with current_tick = Z.succ state.current_tick} in
-      (* Write the module registry to the tree in case it did not exist
-         before. *)
-      let* tree =
-        Tree_encoding.encode module_reg_encoding state.module_reg tree
-      in
       let want_more_input =
         match state.tick with
         | Eval {code = _, []; _} ->
-- 
GitLab