From eccaea6e35e82ccd6697234337698c2609ccb059 Mon Sep 17 00:00:00 2001
From: Thomas Letan <lthms@nomadic-labs.com>
Date: Wed, 10 Aug 2022 10:46:11 +0200
Subject: [PATCH 1/5] =?UTF-8?q?WASM:=20Fix=20helpers=20that=20checks=20WAS?=
 =?UTF-8?q?M=20PVM=20step=E2=80=99s=20proof=20size?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../lib_protocol/test/integration/test_sc_rollup_wasm.ml  | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/test/integration/test_sc_rollup_wasm.ml b/src/proto_alpha/lib_protocol/test/integration/test_sc_rollup_wasm.ml
index 0f51d7bfa3eb..0ed8a0a0049f 100644
--- a/src/proto_alpha/lib_protocol/test/integration/test_sc_rollup_wasm.ml
+++ b/src/proto_alpha/lib_protocol/test/integration/test_sc_rollup_wasm.ml
@@ -130,16 +130,16 @@ let check_proof_size ~loc context input_opt s =
 (* Like [eval] but also checks the proof size. *)
 let checked_eval ~loc context s =
   let open Lwt_result_syntax in
+  let* () = check_proof_size ~loc context None s in
   let*! s = Prover.eval s in
-  let+ () = check_proof_size ~loc context None s in
-  s
+  return s
 
 (* Like [set_input] but also checks the proof size. *)
 let checked_set_input ~loc context input s =
   let open Lwt_result_syntax in
+  let* () = check_proof_size ~loc context (Some input) s in
   let*! s = Prover.set_input input s in
-  let+ () = check_proof_size ~loc context (Some input) s in
-  s
+  return s
 
 let complete_boot_sector sector :
     Tezos_scoru_wasm.Gather_floppies.origination_message =
-- 
GitLab


From 1f888e276b23bfaafc6ec25e2da38314dd45c1ae Mon Sep 17 00:00:00 2001
From: Thomas Letan <lthms@nomadic-labs.com>
Date: Wed, 10 Aug 2022 20:13:01 +0200
Subject: [PATCH 2/5] WASM: Fix existing lazy-containers to support move of
 subtrees

---
 .../chunked_byte_vector.ml                    | 23 ++++-
 .../chunked_byte_vector.mli                   | 26 ++++-
 src/lib_lazy_containers/lazy_map.ml           | 57 ++++++-----
 src/lib_lazy_containers/lazy_map.mli          | 60 ++++++------
 src/lib_lazy_containers/lazy_vector.ml        | 95 +++++++++----------
 src/lib_lazy_containers/lazy_vector.mli       | 65 +++++++++----
 .../test/chunked_byte_vector_tests.ml         |  3 +-
 .../test/lazy_vector_tests.ml                 |  4 +-
 .../environment_V6.ml                         | 12 ++-
 .../environment_V7.ml                         | 12 ++-
 src/lib_scoru_wasm/test/test_input.ml         |  8 ++
 src/lib_scoru_wasm/test/test_wasm_encoding.ml |  8 ++
 src/lib_tree_encoding/decoding.ml             |  7 ++
 src/lib_tree_encoding/decoding.mli            |  4 +
 src/lib_tree_encoding/encoding.ml             | 16 +++-
 src/lib_tree_encoding/encoding.mli            |  6 ++
 src/lib_tree_encoding/test/test_encoding.ml   |  8 ++
 src/lib_tree_encoding/tree.ml                 | 10 ++
 src/lib_tree_encoding/tree.mli                | 11 +++
 src/lib_tree_encoding/tree_encoding.ml        | 51 ++++++----
 src/lib_tree_encoding/tree_encoding.mli       | 11 +++
 src/lib_webassembly/bin/text/parser.mly       | 34 +++----
 src/lib_webassembly/exec/eval.ml              | 50 +++++-----
 src/lib_webassembly/runtime/partial_memory.ml |  2 +-
 src/lib_webassembly/runtime/partial_table.ml  |  2 +-
 src/lib_webassembly/syntax/ast.ml             |  2 +-
 src/lib_webassembly/valid/valid.ml            |  2 +-
 27 files changed, 381 insertions(+), 208 deletions(-)

diff --git a/src/lib_lazy_containers/chunked_byte_vector.ml b/src/lib_lazy_containers/chunked_byte_vector.ml
index 759712c57340..510d894b7a7a 100644
--- a/src/lib_lazy_containers/chunked_byte_vector.ml
+++ b/src/lib_lazy_containers/chunked_byte_vector.ml
@@ -108,7 +108,12 @@ module type S = sig
 
   type t
 
-  val create : ?get_chunk:(int64 -> Chunk.t effect) -> int64 -> t
+  val create :
+    ?origin:Lazy_map.tree -> ?get_chunk:(int64 -> Chunk.t effect) -> int64 -> t
+
+  val origin : t -> Lazy_map.tree option
+
+  val allocate : int64 -> t
 
   val of_string : string -> t
 
@@ -140,12 +145,14 @@ module Make (Effect : Effect.S) : S with type 'a effect = 'a Effect.t = struct
 
   let def_get_chunk _ = Effect.return (Chunk.alloc ())
 
-  let create ?(get_chunk = def_get_chunk) length =
+  let create ?origin ?(get_chunk = def_get_chunk) length =
     let chunks =
-      Vector.create ~produce_value:get_chunk (Chunk.num_needed length)
+      Vector.create ?origin ~produce_value:get_chunk (Chunk.num_needed length)
     in
     {length; chunks}
 
+  let origin vector = Vector.origin vector.chunks
+
   let grow vector size_delta =
     if 0L < size_delta then (
       let new_size = Int64.add vector.length size_delta in
@@ -158,9 +165,17 @@ module Make (Effect : Effect.S) : S with type 'a effect = 'a Effect.t = struct
            error in case of absent value (which is the case when
            growing the chunked byte vector requires to allocate new
            chunks). *)
-        Vector.grow ~produce_value:def_get_chunk chunk_count_delta vector.chunks ;
+        Vector.grow
+          ~default:(fun () -> Chunk.alloc ())
+          chunk_count_delta
+          vector.chunks ;
       vector.length <- new_size)
 
+  let allocate length =
+    let res = create 0L in
+    grow res length ;
+    res
+
   let length vector = vector.length
 
   let load_byte vector address =
diff --git a/src/lib_lazy_containers/chunked_byte_vector.mli b/src/lib_lazy_containers/chunked_byte_vector.mli
index 077f756e923e..c50f8341dde5 100644
--- a/src/lib_lazy_containers/chunked_byte_vector.mli
+++ b/src/lib_lazy_containers/chunked_byte_vector.mli
@@ -60,8 +60,25 @@ module type S = sig
   type t
 
   (** [create length] creates a chunked byte vector that has capacity for [length]
-      bytes. *)
-  val create : ?get_chunk:(int64 -> Chunk.t effect) -> int64 -> t
+      bytes.
+
+      {b Note:} This function is expected to be use only by the
+      tree-encoding library. To create a brand new chunked byte
+      vector, use {!allocate}.  *)
+  val create :
+    ?origin:Lazy_map.tree -> ?get_chunk:(int64 -> Chunk.t effect) -> int64 -> t
+
+  (** [origin vec] returns the tree of origin of the vector, if it exists.
+
+      {b Note:} The sole consumer of this function is expected to be
+      the tree-encoding library. *)
+  val origin : t -> Lazy_map.tree option
+
+  (** [allocate len] creates a new zeroed chunked byte vector.
+
+      {b Note:} This function may be dangerous to use in a tick if
+      [len] is too large. *)
+  val allocate : int64 -> t
 
   (** [of_string str] creates a chunked byte vector from the given [str]. *)
   val of_string : string -> t
@@ -79,7 +96,10 @@ module type S = sig
   val to_bytes : t -> bytes effect
 
   (** [grow vector length_delta] increases the byte vector length by
-      [length_delta]. *)
+      [length_delta] and initializes the memory with empty chunks.
+
+      {b Note:} This function may be dangerous to use in a tick if
+      [length_delta] is too large. *)
   val grow : t -> int64 -> unit
 
   (** [length vector] returns the length of [vector] in bytes. *)
diff --git a/src/lib_lazy_containers/lazy_map.ml b/src/lib_lazy_containers/lazy_map.ml
index 30c09ac83419..90e49fee12be 100644
--- a/src/lib_lazy_containers/lazy_map.ml
+++ b/src/lib_lazy_containers/lazy_map.ml
@@ -23,6 +23,8 @@
 (*                                                                           *)
 (*****************************************************************************)
 
+type tree = ..
+
 module Effect = struct
   module type S = sig
     type 'a t
@@ -66,27 +68,25 @@ module type S = sig
 
   type 'a t
 
+  val origin : 'a t -> tree option
+
   val string_of_key : key -> string
 
   val pp : (Format.formatter -> 'a -> unit) -> Format.formatter -> 'a t -> unit
 
   val to_string : ('a -> string) -> 'a t -> string
 
-  val create : ?values:'a Map.t -> ?produce_value:'a producer -> unit -> 'a t
+  val create :
+    ?values:'a Map.t ->
+    ?produce_value:'a producer ->
+    ?origin:tree ->
+    unit ->
+    'a t
 
   val get : key -> 'a t -> 'a effect
 
   val set : key -> 'a -> 'a t -> 'a t
 
-  val merge_into :
-    ?map_key:(key -> key) ->
-    ?choose_producer:('a producer -> 'a producer -> 'a producer) ->
-    'a t ->
-    'a t ->
-    'a t
-
-  val with_producer : ('a producer -> 'a producer) -> 'a t -> 'a t
-
   val loaded_bindings : 'a t -> (key * 'a) list
 end
 
@@ -102,7 +102,13 @@ module Make (Effect : Effect.S) (Key : KeyS) :
 
   type 'a producer = key -> 'a effect
 
-  type 'a t = {produce_value : 'a producer; mutable values : 'a Map.t}
+  type 'a t = {
+    origin : tree option;
+    produce_value : 'a producer;
+    mutable values : 'a Map.t;
+  }
+
+  let origin {origin; _} = origin
 
   let string_of_key = Key.to_string
 
@@ -130,8 +136,9 @@ module Make (Effect : Effect.S) (Key : KeyS) :
 
   let def_produce_value _ = raise UnexpectedAccess
 
-  let create ?(values = Map.empty) ?(produce_value = def_produce_value) () =
-    {produce_value; values}
+  let create ?(values = Map.empty) ?(produce_value = def_produce_value) ?origin
+      () =
+    {produce_value; values; origin}
 
   let get key map =
     let open Effect in
@@ -145,20 +152,6 @@ module Make (Effect : Effect.S) (Key : KeyS) :
 
   let set key value map = {map with values = Map.add key value map.values}
 
-  let merge_into ?(map_key = Fun.id) ?(choose_producer = fun _ dest -> dest) src
-      dest =
-    let produce_value = choose_producer src.produce_value dest.produce_value in
-    let values =
-      Map.fold
-        (fun src_key -> Map.add (map_key src_key))
-        src.values (* fold subject *)
-        dest.values (* accumulator *)
-    in
-    {produce_value; values}
-
-  let with_producer morph map =
-    {map with produce_value = morph map.produce_value}
-
   let loaded_bindings m = Map.bindings m.values
 end
 
@@ -182,7 +175,11 @@ module Mutable = struct
     val of_immutable : 'a Map.t -> 'a t
 
     val create :
-      ?values:'a Map.Map.t -> ?produce_value:'a Map.producer -> unit -> 'a t
+      ?values:'a Map.Map.t ->
+      ?produce_value:'a Map.producer ->
+      ?origin:tree ->
+      unit ->
+      'a t
 
     val get : key -> 'a t -> 'a Map.effect
 
@@ -203,8 +200,8 @@ module Mutable = struct
 
     let of_immutable = ref
 
-    let create ?values ?produce_value unit =
-      of_immutable (Map.create ?values ?produce_value unit)
+    let create ?values ?produce_value ?origin unit =
+      of_immutable (Map.create ?values ?produce_value ?origin unit)
 
     let get key map_ref = Map.get key !map_ref
 
diff --git a/src/lib_lazy_containers/lazy_map.mli b/src/lib_lazy_containers/lazy_map.mli
index 321cf3e42362..b50577b4690c 100644
--- a/src/lib_lazy_containers/lazy_map.mli
+++ b/src/lib_lazy_containers/lazy_map.mli
@@ -23,6 +23,13 @@
 (*                                                                           *)
 (*****************************************************************************)
 
+(** An extensible type to record the type of trees used as backend for
+    the lazy map.
+
+    {b Note:} If you use the 'tree-encoding' library, then its functor
+    takes care of adding a new constructor for the expected tree. *)
+type tree = ..
+
 (**
   A lazy map is a key-value association where each value is created dynamically.
 *)
@@ -64,6 +71,12 @@ module type S = sig
 
   type 'a t
 
+  (** [origin map] returns the tree of origin of the map, if it exists.
+
+      {b Note:} The sole consumer of this function is expected to be
+      the tree-encoding library. *)
+  val origin : 'a t -> tree option
+
   (** [string_of_key key] turns the given [key] into a string. *)
   val string_of_key : key -> string
 
@@ -76,10 +89,21 @@ module type S = sig
       mutation. *)
   val to_string : ('a -> string) -> 'a t -> string
 
-  (** [create ?values num_elements produce_value] produces a lazy map with
-      [num_elements] entries where each is created using [produce_value].
-      [values] may be provided to supply an initial set of entries. *)
-  val create : ?values:'a Map.t -> ?produce_value:'a producer -> unit -> 'a t
+  (** [create ?values ?produce_value ?origin num_elements] produces a
+      lazy map with [num_elements] entries where each is created using
+      [produce_value]. [values] may be provided to supply an initial
+      set of entries.
+
+      {b Note:} This function is intended to be used [produce_value]
+      should only be used by the tree-encoding library. If you want to
+      fill a newly created map with some value, use [values] or
+      [set] on the empty map. *)
+  val create :
+    ?values:'a Map.t ->
+    ?produce_value:'a producer ->
+    ?origin:tree ->
+    unit ->
+    'a t
 
   (** [get key map] retrieves the element at [key].
 
@@ -91,28 +115,6 @@ module type S = sig
       @raises Exn.Bounds when trying to set an invalid key *)
   val set : key -> 'a -> 'a t -> 'a t
 
-  (** [merge_into ?choose_producer ?map_key source dest] produces a new lazy map
-      by merging [source] into [dest].
-
-      The keys of [source] can be transformed using [map_key] prior to merging.
-      By default no keys will be transformed.
-
-      If a [key] was instantiated in [dest] it will also be present in the
-      resulting lazy map.
-
-      [choose_producer source_producer dest_producer] will be
-      used to determine the new producer function - by default the [dest]
-      producer is chosen. *)
-  val merge_into :
-    ?map_key:(key -> key) ->
-    ?choose_producer:('a producer -> 'a producer -> 'a producer) ->
-    'a t ->
-    'a t ->
-    'a t
-
-  (** [with_producer morph] lifts a morphism for a [producer] to one on [t]. *)
-  val with_producer : ('a producer -> 'a producer) -> 'a t -> 'a t
-
   (** [loaded_bindings map] returns the [(key * 'a) list] representation of the
       map [map] containing only the loaded values, in order of increasing keys.
       This function is a witness of internal mutations. *)
@@ -153,7 +155,11 @@ module Mutable : sig
     val of_immutable : 'a Map.t -> 'a t
 
     val create :
-      ?values:'a Map.Map.t -> ?produce_value:'a Map.producer -> unit -> 'a t
+      ?values:'a Map.Map.t ->
+      ?produce_value:'a Map.producer ->
+      ?origin:tree ->
+      unit ->
+      'a t
 
     val get : key -> 'a t -> 'a Map.effect
 
diff --git a/src/lib_lazy_containers/lazy_vector.ml b/src/lib_lazy_containers/lazy_vector.ml
index 7dd732d175fd..68a1441ae163 100644
--- a/src/lib_lazy_containers/lazy_vector.ml
+++ b/src/lib_lazy_containers/lazy_vector.ml
@@ -84,9 +84,12 @@ module type S = sig
     ?first_key:key ->
     ?values:'a Map.Map.t ->
     ?produce_value:'a producer ->
+    ?origin:Lazy_map.tree ->
     key ->
     'a t
 
+  val origin : 'a t -> Lazy_map.tree option
+
   val empty : unit -> 'a t
 
   val singleton : 'a -> 'a t
@@ -99,11 +102,13 @@ module type S = sig
 
   val cons : 'a -> 'a t -> 'a t
 
-  val grow : ?produce_value:'a producer -> key -> 'a t -> 'a t
+  val grow : ?default:(unit -> 'a) -> key -> 'a t -> 'a t
 
   val append : 'a -> 'a t -> 'a t * key
 
-  val concat : 'a t -> 'a t -> 'a t
+  val concat : 'a t -> 'a t -> 'a t effect
+
+  val unsafe_concat : 'a t -> 'a t -> 'a t
 
   val to_list : 'a t -> 'a list effect
 
@@ -154,10 +159,13 @@ module Make (Effect : Effect.S) (Key : KeyS) :
 
   let num_elements map = map.num_elements
 
-  let create ?(first_key = Key.zero) ?values ?produce_value num_elements =
-    let values = Map.create ?values ?produce_value () in
+  let create ?(first_key = Key.zero) ?values ?produce_value ?origin num_elements
+      =
+    let values = Map.create ?values ?produce_value ?origin () in
     {first = first_key; num_elements; values}
 
+  let origin {values; _} = Map.origin values
+
   let empty () = create Key.zero
 
   let of_list values =
@@ -187,50 +195,21 @@ module Make (Effect : Effect.S) (Key : KeyS) :
     let num_elements = Key.succ map.num_elements in
     {first; values; num_elements}
 
-  let grow ?produce_value delta map =
-    if
-      Key.unsigned_compare (Key.add delta map.num_elements) map.num_elements < 0
-    then raise Exn.SizeOverflow ;
-
-    let map_produce_value old_produce_value =
-      match produce_value with
-      | Some produce_new_value ->
-          let boundary = Key.add map.num_elements map.first in
-          fun key ->
-            if Key.compare key boundary >= 0 then
-              (* Normalize the key so that it is relative to the boundary.
-                  The first new value will be produced with
-                  [produce_value Key.zero]. *)
-              let key = Key.sub key boundary in
-              produce_new_value key
-            else old_produce_value key
-      | None -> old_produce_value
+  let append_opt elt map =
+    let num_elements = map.num_elements in
+    let map = {map with num_elements = Key.succ num_elements} in
+    let map =
+      match elt with Some elt -> set num_elements elt map | None -> map
     in
-    let values = Map.with_producer map_produce_value map.values in
-    let num_elements = Key.add map.num_elements delta in
-    {map with values; num_elements}
+    (map, num_elements)
 
-  let append elt map =
-    let i = num_elements map in
-    (map |> grow Key.(succ zero) |> set i elt, i)
+  let append elt map = append_opt (Some elt) map
 
-  let concat lhs rhs =
-    let boundary = Key.add lhs.first lhs.num_elements in
-    let choose_producer rhs_produce_value lhs_produce_value key =
-      if Key.compare key boundary >= 0 then
-        rhs_produce_value (Key.sub key boundary |> Key.add rhs.first)
-      else lhs_produce_value key
-    in
-    let num_elements = Key.add lhs.num_elements rhs.num_elements in
-    let rhs_offset = Key.sub boundary rhs.first in
-    let values =
-      Map.merge_into
-        ~choose_producer
-        ~map_key:(Key.add rhs_offset)
-        rhs.values
-        lhs.values
-    in
-    {lhs with num_elements; values}
+  let rec grow ?default delta map =
+    if Key.(delta <= zero) then map
+    else
+      let map, _ = append_opt (Option.map (fun f -> f ()) default) map in
+      grow ?default Key.(pred delta) map
 
   let to_list map =
     let open Effect in
@@ -248,8 +227,19 @@ module Make (Effect : Effect.S) (Key : KeyS) :
     if map.num_elements = Key.zero then return []
     else (unroll [@ocaml.tailcall]) [] (Key.pred map.num_elements)
 
+  let concat lhs rhs =
+    let open Effect in
+    let* lhs = to_list lhs in
+    let+ rhs = to_list rhs in
+    of_list (lhs @ rhs)
+
   let loaded_bindings m = Map.loaded_bindings m.values
 
+  let unsafe_concat lhs rhs =
+    let lhs = loaded_bindings lhs |> List.map snd in
+    let rhs = loaded_bindings rhs |> List.map snd in
+    of_list (lhs @ rhs)
+
   let first_key vector = vector.first
 end
 
@@ -286,14 +276,17 @@ module Mutable = struct
     val create :
       ?values:'a Vector.Map.Map.t ->
       ?produce_value:'a Vector.producer ->
+      ?origin:Lazy_map.tree ->
       key ->
       'a t
 
+    val origin : 'a t -> Lazy_map.tree option
+
     val get : key -> 'a t -> 'a Vector.effect
 
     val set : key -> 'a -> 'a t -> unit
 
-    val grow : ?produce_value:'a Vector.producer -> key -> 'a t -> unit
+    val grow : ?default:(unit -> 'a) -> key -> 'a t -> unit
 
     val append : 'a -> 'a t -> key
 
@@ -319,15 +312,17 @@ module Mutable = struct
 
     let of_immutable = ref
 
-    let create ?values ?produce_value num_elements =
-      of_immutable (Vector.create ?values ?produce_value num_elements)
+    let create ?values ?produce_value ?origin num_elements =
+      of_immutable (Vector.create ?values ?produce_value ?origin num_elements)
+
+    let origin vector = Vector.origin !vector
 
     let get key map_ref = Vector.get key !map_ref
 
     let set key value map_ref = map_ref := Vector.set key value !map_ref
 
-    let grow ?produce_value delta map_ref =
-      map_ref := Vector.grow ?produce_value delta !map_ref
+    let grow ?default delta map_ref =
+      map_ref := Vector.grow ?default delta !map_ref
 
     let append elt map_ref =
       let new_map, i = Vector.append elt !map_ref in
diff --git a/src/lib_lazy_containers/lazy_vector.mli b/src/lib_lazy_containers/lazy_vector.mli
index 2d15b5f33972..5aa1fbbb1e81 100644
--- a/src/lib_lazy_containers/lazy_vector.mli
+++ b/src/lib_lazy_containers/lazy_vector.mli
@@ -94,18 +94,30 @@ module type S = sig
       vector. *)
   val num_elements : 'a t -> key
 
-  (** [create ?first_key ?values ?produce_value num_elements] produces a lazy
-      vector with [num_elements] entries where each is created using
-      [produce_value]. [values] may be provided to supply an initial set of
-      entries. [first_key] specifies the first index of the vector if given and
-      defaults to zero. *)
+  (** [create ?first_key ?values ?produce_value ?origin num_elements]
+      produces a lazy vector with [num_elements] entries where each is
+      created using [produce_value]. [values] may be provided to
+      supply an initial set of entries. [first_key] specifies the
+      first index of the vector if given and defaults to zero.
+
+      {b Note:} The [produce_value] and [origin] arguments are
+      expected to be used by the 'tree-encoding' library. If you want
+      to pre-fill your vector, creates an empty vector and use [grow]
+      or [set]. *)
   val create :
     ?first_key:key ->
     ?values:'a Map.Map.t ->
     ?produce_value:'a producer ->
+    ?origin:Lazy_map.tree ->
     key ->
     'a t
 
+  (** [origin vec] returns the tree of origin of the vector, if it exists.
+
+      {b Note:} The sole consumer of this function is expected to be
+      the tree-encoding library. *)
+  val origin : 'a t -> Lazy_map.tree option
+
   (** [empty ()] creates a vector of size zero. This is used in conjunction with
       {!cons} to model list-like structure. *)
   val empty : unit -> 'a t
@@ -115,7 +127,10 @@ module type S = sig
 
   (** [of_list values] creates a vector where each association is the index in
       the list to its value. The first item's key is [zero], the second is
-      [succ zero] and so on. *)
+      [succ zero] and so on.
+
+      {b Note:} This function may be dangerous to use in a tick, if
+      the size of [of_list] is unbounded. *)
   val of_list : 'a list -> 'a t
 
   (** [get key vector] retrieves the element at [key].
@@ -129,16 +144,18 @@ module type S = sig
   val set : key -> 'a -> 'a t -> 'a t
 
   (** [cons value vector] prepends a value to the front and grows the vector by
-      one. That value can then be accessed using the [zero] key.
-
-      Time complexity: O(log(instantiated_elements_in_vector)) *)
+      one. That value can then be accessed using the [zero] key. *)
   val cons : 'a -> 'a t -> 'a t
 
-  (** [grow delta ?produce_value vector] creates a new lazy vector that has
+  (** [grow delta ?default vector] creates a new lazy vector that has
       [delta] more items than [vector]. This also retains all values that have
-      previously been created. New values will be created with [produce_values]
-      if provided, starting with [Key.zero] for the new values. *)
-  val grow : ?produce_value:'a producer -> key -> 'a t -> 'a t
+      previously been created. New values will be created with [default]
+      if provided.
+
+      {b Note:} This function may be dangerous to use in a tick, if
+      [delta] is unbounded, or if the result of [default] is
+      large. *)
+  val grow : ?default:(unit -> 'a) -> key -> 'a t -> 'a t
 
   (** [append elt vector] creates a new lazy vector that has one
       more item than [vector] whose value is [elt]. This is a shortcut
@@ -146,11 +163,22 @@ module type S = sig
       Also returns the key of the added element. *)
   val append : 'a -> 'a t -> 'a t * key
 
-  (** [concat lhs rhs] Concatenates two lazy vectors. *)
-  val concat : 'a t -> 'a t -> 'a t
+  (** [concat lhs rhs] concatenates two lazy vectors.
+
+      {b Note: This function maybe dangerous to use in a tick because
+      {i every} entries of both [lhs] and [rhs] will be loaded in
+      memory. *)
+  val concat : 'a t -> 'a t -> 'a t effect
+
+  (** [unsafe_concat] concatenates two lazy vectors, {b assuming every
+      entries of both vectors are already loaded in memory}. *)
+  val unsafe_concat : 'a t -> 'a t -> 'a t
 
   (** [to_list vector] extracts all values of the given [vector] and
-      collects them in a list.  *)
+      collects them in a list.
+
+      {b Note:} This function may be dangerous to use in a tick
+      because all entries of the vector are loaded in memory. *)
   val to_list : 'a t -> 'a list effect
 
   (** [loaded_bindings vector] returns the [(key * 'a) list] representation of
@@ -200,14 +228,17 @@ module Mutable : sig
     val create :
       ?values:'a Vector.Map.Map.t ->
       ?produce_value:'a Vector.producer ->
+      ?origin:Lazy_map.tree ->
       key ->
       'a t
 
+    val origin : 'a t -> Lazy_map.tree option
+
     val get : key -> 'a t -> 'a Vector.effect
 
     val set : key -> 'a -> 'a t -> unit
 
-    val grow : ?produce_value:'a Vector.producer -> key -> 'a t -> unit
+    val grow : ?default:(unit -> 'a) -> key -> 'a t -> unit
 
     val append : 'a -> 'a t -> key
 
diff --git a/src/lib_lazy_containers/test/chunked_byte_vector_tests.ml b/src/lib_lazy_containers/test/chunked_byte_vector_tests.ml
index 97a7b0363b7d..06a35ea087a7 100644
--- a/src/lib_lazy_containers/test/chunked_byte_vector_tests.ml
+++ b/src/lib_lazy_containers/test/chunked_byte_vector_tests.ml
@@ -49,8 +49,9 @@ let store_load_byte_works =
 let grow_works =
   Test.make
     ~name:"grow works"
-    Gen.(pair string ui64)
+    Gen.(pair string small_int)
     (fun (init_str, grow_len) ->
+      let grow_len = Int64.of_int grow_len in
       let vector = of_string init_str in
       let check_contents () =
         List.init (String.length init_str) (fun i ->
diff --git a/src/lib_lazy_containers/test/lazy_vector_tests.ml b/src/lib_lazy_containers/test/lazy_vector_tests.ml
index 00c908e74ae3..96c178f0f316 100644
--- a/src/lib_lazy_containers/test/lazy_vector_tests.ml
+++ b/src/lib_lazy_containers/test/lazy_vector_tests.ml
@@ -73,7 +73,7 @@ let grow_works =
     ~name:"grow works"
     Gen.(pair (gen int) nat)
     (fun (map, len) ->
-      let map2 = IntVector.grow ~produce_value:(fun x -> x * 2) len map in
+      let map2 = IntVector.grow ~default:(fun () -> 2) len map in
       let check1 =
         List.init (IntVector.num_elements map) (fun i ->
             IntVector.get i map2 = IntVector.get i map)
@@ -82,7 +82,7 @@ let grow_works =
       let check2 =
         List.init len (fun i ->
             let key = i + IntVector.num_elements map in
-            IntVector.get key map2 = i * 2)
+            IntVector.get key map2 = 2)
         |> List.for_all Fun.id
       in
       let check3 =
diff --git a/src/lib_protocol_environment/environment_V6.ml b/src/lib_protocol_environment/environment_V6.ml
index 76353e57d065..97e363dcdf85 100644
--- a/src/lib_protocol_environment/environment_V6.ml
+++ b/src/lib_protocol_environment/environment_V6.ml
@@ -1122,7 +1122,17 @@ struct
     module Make
         (Tree : Context.TREE with type key = string list and type value = bytes) =
     struct
-      module Wasm = Tezos_scoru_wasm.Wasm_pvm.Make (Tree)
+      type Lazy_containers.Lazy_map.tree += PVM_tree of Tree.tree
+
+      module Wasm = Tezos_scoru_wasm.Wasm_pvm.Make (struct
+        include Tree
+
+        let select = function
+          | PVM_tree t -> t
+          | _ -> raise Tree_encoding.Incorrect_tree_type
+
+        let wrap t = PVM_tree t
+      end)
 
       (* TODO: https://gitlab.com/tezos/tezos/-/issues/3214
          The rest of the module is pure boilerplate converting between
diff --git a/src/lib_protocol_environment/environment_V7.ml b/src/lib_protocol_environment/environment_V7.ml
index d4c8a061d97a..e6070c4d8e97 100644
--- a/src/lib_protocol_environment/environment_V7.ml
+++ b/src/lib_protocol_environment/environment_V7.ml
@@ -1122,7 +1122,17 @@ struct
     module Make
         (Tree : Context.TREE with type key = string list and type value = bytes) =
     struct
-      module Wasm = Tezos_scoru_wasm.Wasm_pvm.Make (Tree)
+      type Lazy_containers.Lazy_map.tree += PVM_tree of Tree.tree
+
+      module Wasm = Tezos_scoru_wasm.Wasm_pvm.Make (struct
+        include Tree
+
+        let select = function
+          | PVM_tree t -> t
+          | _ -> raise Tree_encoding.Incorrect_tree_type
+
+        let wrap t = PVM_tree t
+      end)
 
       (* TODO: https://gitlab.com/tezos/tezos/-/issues/3214
          The rest of the module is pure boilerplate converting between
diff --git a/src/lib_scoru_wasm/test/test_input.ml b/src/lib_scoru_wasm/test/test_input.ml
index 6e9f1882a95d..52529774d74d 100644
--- a/src/lib_scoru_wasm/test/test_input.ml
+++ b/src/lib_scoru_wasm/test/test_input.ml
@@ -191,10 +191,18 @@ let empty_tree () =
   let empty_store = Context.empty index in
   return @@ Context.Tree.empty empty_store
 
+type Lazy_containers.Lazy_map.tree += Tree of Context.tree
+
 module Tree : Tree_encoding.TREE with type tree = Context.tree = struct
   type tree = Context.tree
 
   include Context.Tree
+
+  let select = function
+    | Tree t -> t
+    | _ -> raise Tree_encoding.Incorrect_tree_type
+
+  let wrap t = Tree t
 end
 
 module Wasm = Wasm_pvm.Make (Tree)
diff --git a/src/lib_scoru_wasm/test/test_wasm_encoding.ml b/src/lib_scoru_wasm/test/test_wasm_encoding.ml
index 3ab3d4292ee4..6dad0f1250db 100644
--- a/src/lib_scoru_wasm/test/test_wasm_encoding.ml
+++ b/src/lib_scoru_wasm/test/test_wasm_encoding.ml
@@ -47,6 +47,8 @@ let qcheck ?count ?print gen f =
 (* Use context-binary for testing. *)
 module Context = Tezos_context_memory.Context_binary
 
+type Lazy_containers.Lazy_map.tree += Tree of Context.tree
+
 module Tree = struct
   type t = Context.t
 
@@ -57,6 +59,12 @@ module Tree = struct
   type value = Context.value
 
   include Context.Tree
+
+  let select = function
+    | Tree t -> t
+    | _ -> raise Tree_encoding.Incorrect_tree_type
+
+  let wrap t = Tree t
 end
 
 module Tree_encoding = Tree_encoding.Make (Tree)
diff --git a/src/lib_tree_encoding/decoding.ml b/src/lib_tree_encoding/decoding.ml
index 834c7f2274dc..8a216615f5bd 100644
--- a/src/lib_tree_encoding/decoding.ml
+++ b/src/lib_tree_encoding/decoding.ml
@@ -44,6 +44,8 @@ module type S = sig
 
   val value_option : key -> 'a Data_encoding.t -> 'a option t
 
+  val subtree : Lazy_containers.Lazy_map.tree option t
+
   val value : ?default:'a -> key -> 'a Data_encoding.t -> 'a t
 
   val scope : key -> 'a t -> 'a t
@@ -160,6 +162,11 @@ module Make (T : Tree.S) : S with type tree = T.tree = struct
     | None, Some default -> return default
     | None, None -> raise (Key_not_found (prefix key))
 
+  let subtree tree prefix =
+    let open Lwt_syntax in
+    let+ tree = T.find_tree tree (prefix []) in
+    Option.map T.wrap tree
+
   let scope key dec tree prefix = dec tree (append_key prefix key)
 
   let lazy_mapping to_key field_enc input_tree input_prefix =
diff --git a/src/lib_tree_encoding/decoding.mli b/src/lib_tree_encoding/decoding.mli
index cbe3a15549cd..9e3a325f7602 100644
--- a/src/lib_tree_encoding/decoding.mli
+++ b/src/lib_tree_encoding/decoding.mli
@@ -62,6 +62,10 @@ module type S = sig
   *)
   val value_option : key -> 'a Data_encoding.t -> 'a option t
 
+  (** [subtree] returns the subtree located at the prefix tree under
+      which it is called. *)
+  val subtree : Lazy_containers.Lazy_map.tree option t
+
   (** [value ?default key data_encoding] retrieves the value at a
       given [key] by decoding its raw value using the provided
       [data_encoding].
diff --git a/src/lib_tree_encoding/encoding.ml b/src/lib_tree_encoding/encoding.ml
index 89751f81ea3d..4a1375b56ae9 100644
--- a/src/lib_tree_encoding/encoding.ml
+++ b/src/lib_tree_encoding/encoding.ml
@@ -42,6 +42,9 @@ module type S = sig
 
   val run : 'a t -> 'a -> tree -> tree Lwt.t
 
+  val with_subtree :
+    ('a -> Lazy_containers.Lazy_map.tree option) -> 'a t -> 'a t
+
   val raw : key -> bytes t
 
   val value_option : key -> 'a Data_encoding.t -> 'a option t
@@ -83,6 +86,15 @@ module Make (T : Tree.S) = struct
 
   let run enc value tree = enc value Fun.id tree
 
+  let with_subtree get_subtree enc value prefix input_tree =
+    let open Lwt.Syntax in
+    match get_subtree value with
+    | Some tree ->
+        let* input_tree = T.remove input_tree (prefix []) in
+        let* input_tree = T.add_tree input_tree (prefix []) (T.select tree) in
+        enc value prefix input_tree
+    | None -> enc value prefix input_tree
+
   let lwt enc value prefix tree =
     let open Lwt_syntax in
     let* v = value in
@@ -110,8 +122,8 @@ module Make (T : Tree.S) = struct
 
   let raw suffix bytes prefix tree = T.add tree (prefix suffix) bytes
 
-  let value suffix enc =
-    contramap (Data_encoding.Binary.to_bytes_exn enc) (raw suffix)
+  let value suffix enc v prefix tree =
+    contramap (Data_encoding.Binary.to_bytes_exn enc) (raw suffix) v prefix tree
 
   let value_option key encoding v prefix tree =
     match v with
diff --git a/src/lib_tree_encoding/encoding.mli b/src/lib_tree_encoding/encoding.mli
index 6bf13e349a75..c4e280aa4dba 100644
--- a/src/lib_tree_encoding/encoding.mli
+++ b/src/lib_tree_encoding/encoding.mli
@@ -59,6 +59,12 @@ module type S = sig
       [No_tag_matched] exception. *)
   val run : 'a t -> 'a -> tree -> tree Lwt.t
 
+  (** [with_subtree get_subtree enc] will use [get_subtree] to fetch
+      the tree of origin of a value to be encoded with [enc], to place
+      it in the targeted tree before using [enc]. *)
+  val with_subtree :
+    ('a -> Lazy_containers.Lazy_map.tree option) -> 'a t -> 'a t
+
   (** [raw key] returns an encoder that encodes raw bytes at the given key. *)
   val raw : key -> bytes t
 
diff --git a/src/lib_tree_encoding/test/test_encoding.ml b/src/lib_tree_encoding/test/test_encoding.ml
index b7e3220f2d35..fd5041270dec 100644
--- a/src/lib_tree_encoding/test/test_encoding.ml
+++ b/src/lib_tree_encoding/test/test_encoding.ml
@@ -37,10 +37,18 @@ open Lazy_containers
 (* Use context-binary for testing. *)
 module Context = Tezos_context_memory.Context_binary
 
+type Lazy_containers.Lazy_map.tree += Tree of Context.tree
+
 module Tree : Tree_encoding.TREE with type tree = Context.tree = struct
   type tree = Context.tree
 
   include Context.Tree
+
+  let select = function
+    | Tree t -> t
+    | _ -> raise Tree_encoding.Incorrect_tree_type
+
+  let wrap t = Tree t
 end
 
 module Map =
diff --git a/src/lib_tree_encoding/tree.ml b/src/lib_tree_encoding/tree.ml
index 7ee7a6d1e8cd..e8a8d4555ee1 100644
--- a/src/lib_tree_encoding/tree.ml
+++ b/src/lib_tree_encoding/tree.ml
@@ -23,6 +23,8 @@
 (*                                                                           *)
 (*****************************************************************************)
 
+exception Incorrect_tree_type
+
 module type S = sig
   type tree
 
@@ -30,9 +32,17 @@ module type S = sig
 
   type value := bytes
 
+  val select : Lazy_containers.Lazy_map.tree -> tree
+
+  val wrap : tree -> Lazy_containers.Lazy_map.tree
+
   val remove : tree -> key -> tree Lwt.t
 
   val add : tree -> key -> value -> tree Lwt.t
 
+  val add_tree : tree -> key -> tree -> tree Lwt.t
+
   val find : tree -> key -> value option Lwt.t
+
+  val find_tree : tree -> key -> tree option Lwt.t
 end
diff --git a/src/lib_tree_encoding/tree.mli b/src/lib_tree_encoding/tree.mli
index d449352c4fb0..231f3cdf9314 100644
--- a/src/lib_tree_encoding/tree.mli
+++ b/src/lib_tree_encoding/tree.mli
@@ -25,6 +25,8 @@
 
 (** Exposes a module type {!S} representing trees. *)
 
+exception Incorrect_tree_type
+
 (** An immutable tree API. *)
 module type S = sig
   type tree
@@ -33,9 +35,18 @@ module type S = sig
 
   type value := bytes
 
+  (** @raise Incorrect_tree_type *)
+  val select : Lazy_containers.Lazy_map.tree -> tree
+
+  val wrap : tree -> Lazy_containers.Lazy_map.tree
+
   val remove : tree -> key -> tree Lwt.t
 
   val add : tree -> key -> value -> tree Lwt.t
 
+  val add_tree : tree -> key -> tree -> tree Lwt.t
+
   val find : tree -> key -> value option Lwt.t
+
+  val find_tree : tree -> key -> tree option Lwt.t
 end
diff --git a/src/lib_tree_encoding/tree_encoding.ml b/src/lib_tree_encoding/tree_encoding.ml
index 03709df43b80..e922bdd5e08c 100644
--- a/src/lib_tree_encoding/tree_encoding.ml
+++ b/src/lib_tree_encoding/tree_encoding.ml
@@ -24,8 +24,9 @@
 (*****************************************************************************)
 
 open Lazy_containers
+include Tree
 
-module type TREE = Tree.S
+module type TREE = S
 
 module type Lwt_vector = Lazy_vector.S with type 'a effect = 'a Lwt.t
 
@@ -342,7 +343,11 @@ module Make (T : Tree.S) : S with type tree = T.tree = struct
       let lazy_map value =
         let to_key k = [Map.string_of_key k] in
         let encode =
-          E.contramap Map.loaded_bindings (E.lazy_mapping to_key value.encode)
+          E.with_subtree
+            Map.origin
+            (E.contramap
+               Map.loaded_bindings
+               (E.lazy_mapping to_key value.encode))
         in
         let decode =
           D.map
@@ -367,23 +372,25 @@ module Make (T : Tree.S) : S with type tree = T.tree = struct
         let open Vector in
         let to_key k = [string_of_key k] in
         let encode =
-          E.contramap
-            (fun vector ->
-              (loaded_bindings vector, num_elements vector, first_key vector))
-            (E.tup3
-               (E.lazy_mapping to_key value.encode)
-               (E.scope ["length"] with_key.encode)
-               (E.scope ["head"] with_key.encode))
+          E.with_subtree Vector.origin
+          @@ E.contramap
+               (fun vector ->
+                 (loaded_bindings vector, num_elements vector, first_key vector))
+               (E.tup3
+                  (E.lazy_mapping to_key value.encode)
+                  (E.scope ["length"] with_key.encode)
+                  (E.scope ["head"] with_key.encode))
         in
         let decode =
           D.map
-            (fun (produce_value, len, head) ->
-              create ~produce_value ~first_key:head len)
+            (fun (origin, produce_value, len, head) ->
+              create ~produce_value ~first_key:head ?origin len)
             (let open D.Syntax in
-            let+ x = D.lazy_mapping to_key value.decode
+            let+ origin = D.subtree
+            and+ x = D.lazy_mapping to_key value.decode
             and+ y = D.scope ["length"] with_key.decode
             and+ z = D.scope ["head"] with_key.decode in
-            (x, y, z))
+            (origin, x, y, z))
         in
         {encode; decode}
     end
@@ -401,19 +408,21 @@ module Make (T : Tree.S) : S with type tree = T.tree = struct
     let open Chunked_byte_vector.Lwt in
     let to_key k = [Int64.to_string k] in
     let encode =
-      E.contramap
-        (fun vector -> (loaded_chunks vector, length vector))
-        (E.tup2
-           (E.lazy_mapping to_key chunk.encode)
-           (E.value ["length"] Data_encoding.int64))
+      E.with_subtree Chunked_byte_vector.Lwt.origin
+      @@ E.contramap
+           (fun vector -> (loaded_chunks vector, length vector))
+           (E.tup2
+              (E.lazy_mapping to_key chunk.encode)
+              (E.value ["length"] Data_encoding.int64))
     in
     let decode =
       D.map
-        (fun (get_chunk, len) -> create ~get_chunk len)
+        (fun (origin, get_chunk, len) -> create ?origin ~get_chunk len)
         (let open D.Syntax in
-        let+ x = D.lazy_mapping to_key chunk.decode
+        let+ origin = D.subtree
+        and+ x = D.lazy_mapping to_key chunk.decode
         and+ y = D.value ["length"] Data_encoding.int64 in
-        (x, y))
+        (origin, x, y))
     in
     {encode; decode}
 
diff --git a/src/lib_tree_encoding/tree_encoding.mli b/src/lib_tree_encoding/tree_encoding.mli
index 0c9fd02a8512..5a29797c60ce 100644
--- a/src/lib_tree_encoding/tree_encoding.mli
+++ b/src/lib_tree_encoding/tree_encoding.mli
@@ -23,6 +23,8 @@
 (*                                                                           *)
 (*****************************************************************************)
 
+exception Incorrect_tree_type
+
 module type TREE = sig
   type tree
 
@@ -30,11 +32,20 @@ module type TREE = sig
 
   type value := bytes
 
+  (** @raise Incorrect_tree_type *)
+  val select : Lazy_containers.Lazy_map.tree -> tree
+
+  val wrap : tree -> Lazy_containers.Lazy_map.tree
+
   val remove : tree -> key -> tree Lwt.t
 
   val add : tree -> key -> value -> tree Lwt.t
 
+  val add_tree : tree -> key -> tree -> tree Lwt.t
+
   val find : tree -> key -> value option Lwt.t
+
+  val find_tree : tree -> key -> tree option Lwt.t
 end
 
 module type Lwt_vector =
diff --git a/src/lib_webassembly/bin/text/parser.mly b/src/lib_webassembly/bin/text/parser.mly
index e199ca9a2107..06813e31a56d 100644
--- a/src/lib_webassembly/bin/text/parser.mly
+++ b/src/lib_webassembly/bin/text/parser.mly
@@ -429,7 +429,7 @@ func_type :
   | func_type_result
     { FuncType (Vector.empty (), $1) }
   | LPAR PARAM value_type_vector RPAR func_type
-    { let FuncType (ins, out) = $5 in FuncType (Vector.concat $3 ins, out) }
+    { let FuncType (ins, out) = $5 in FuncType (Vector.unsafe_concat $3 ins, out) }
   | LPAR PARAM bind_var value_type RPAR func_type  /* Sugar */
     { let FuncType (ins, out) = $6 in FuncType (Vector.cons
  $4 ins, out) }
@@ -438,7 +438,7 @@ func_type_result :
   | /* empty */
     { Vector.empty () }
   | LPAR RESULT value_type_vector RPAR func_type_result
-    { Vector.concat $3 $5 }
+    { Vector.unsafe_concat $3 $5 }
 
 table_type :
   | limits ref_type { TableType ($1, $2) }
@@ -629,13 +629,13 @@ call_instr_type :
 call_instr_params :
   | LPAR PARAM value_type_vector RPAR call_instr_params
     { fun c -> let FuncType (ts1, ts2) = $5 c in
-               FuncType (Vector.concat $3 ts1, ts2) }
+               FuncType (Vector.unsafe_concat $3 ts1, ts2) }
   | call_instr_results
     { fun c -> FuncType (Vector.empty (), $1 c) }
 
 call_instr_results :
   | LPAR RESULT value_type_vector RPAR call_instr_results
-    { fun c -> Vector.concat $3 ($5 c) }
+    { fun c -> Vector.unsafe_concat $3 ($5 c) }
   | /* empty */
     { fun c -> Vector.empty () }
 
@@ -662,13 +662,13 @@ call_instr_type_instr :
 call_instr_params_instr :
   | LPAR PARAM value_type_vector RPAR call_instr_params_instr
     { fun c ->
-      let FuncType (ts1, ts2), es = $5 c in FuncType (Vector.concat $3 ts1, ts2), es }
+      let FuncType (ts1, ts2), es = $5 c in FuncType (Vector.unsafe_concat $3 ts1, ts2), es }
   | call_instr_results_instr
     { fun c -> let ts, es = $1 c in FuncType (Vector.empty (), ts), es }
 
 call_instr_results_instr :
   | LPAR RESULT value_type_vector RPAR call_instr_results_instr
-    { fun c -> let ts, es = $5 c in Vector.concat $3 ts, es }
+    { fun c -> let ts, es = $5 c in Vector.unsafe_concat $3 ts, es }
   | instr
     { fun c -> Vector.empty (), $1 c }
 
@@ -708,13 +708,13 @@ block_param_body :
   | block_result_body { $1 }
   | LPAR PARAM value_type_vector RPAR block_param_body
     { let FuncType (ins, out) = fst $5 in
-      FuncType (Vector.concat $3 ins, out), snd $5 }
+      FuncType (Vector.unsafe_concat $3 ins, out), snd $5 }
 
 block_result_body :
   | instr_list { FuncType (Vector.empty (), Vector.empty ()), $1 }
   | LPAR RESULT value_type_vector RPAR block_result_body
     { let FuncType (ins, out) = fst $5 in
-      FuncType (ins, Vector.concat $3 out), snd $5 }
+      FuncType (ins, Vector.unsafe_concat $3 out), snd $5 }
 
 
 expr :  /* Sugar */
@@ -759,13 +759,13 @@ call_expr_params :
   | LPAR PARAM value_type_vector RPAR call_expr_params
     { fun c ->
       let FuncType (ts1, ts2), es = $5 c in
-      FuncType (Vector.concat $3 ts1, ts2), es }
+      FuncType (Vector.unsafe_concat $3 ts1, ts2), es }
   | call_expr_results
     { fun c -> let ts, es = $1 c in FuncType (Vector.empty (), ts), es }
 
 call_expr_results :
   | LPAR RESULT value_type_vector RPAR call_expr_results
-    { fun c -> let ts, es = $5 c in Vector.concat $3 ts, es }
+    { fun c -> let ts, es = $5 c in Vector.unsafe_concat $3 ts, es }
   | expr_list
     { fun c -> Vector.empty (), $1 c }
 
@@ -794,13 +794,13 @@ if_block_param_body :
   | if_block_result_body { $1 }
   | LPAR PARAM value_type_vector RPAR if_block_param_body
     { let FuncType (ins, out) = fst $5 in
-      FuncType (Vector.concat $3 ins, out), snd $5 }
+      FuncType (Vector.unsafe_concat $3 ins, out), snd $5 }
 
 if_block_result_body :
   | if_ { FuncType (Vector.empty (), Vector.empty ()), $1 }
   | LPAR RESULT value_type_vector RPAR if_block_result_body
     { let FuncType (ins, out) = fst $5 in
-      FuncType (ins, Vector.concat $3 out), snd $5 }
+      FuncType (ins, Vector.unsafe_concat $3 out), snd $5 }
 
 if_ :
   | expr if_
@@ -862,7 +862,7 @@ func_fields :
 func_fields_import :  /* Sugar */
   | func_fields_import_result { $1 }
   | LPAR PARAM value_type_vector RPAR func_fields_import
-    { let FuncType (ins, out) = $5 in FuncType (Vector.concat $3 ins, out) }
+    { let FuncType (ins, out) = $5 in FuncType (Vector.unsafe_concat $3 ins, out) }
   | LPAR PARAM bind_var value_type RPAR func_fields_import  /* Sugar */
     { let FuncType (ins, out) = $6 in FuncType (Vector.cons
  $4 ins, out) }
@@ -870,13 +870,13 @@ func_fields_import :  /* Sugar */
 func_fields_import_result :  /* Sugar */
   | /* empty */ { FuncType (Vector.empty (), Vector.empty ()) }
   | LPAR RESULT value_type_vector RPAR func_fields_import_result
-    { let FuncType (ins, out) = $5 in FuncType (ins, Vector.concat $3 out) }
+    { let FuncType (ins, out) = $5 in FuncType (ins, Vector.unsafe_concat $3 out) }
 
 func_fields_body :
   | func_result_body { $1 }
   | LPAR PARAM value_type_vector RPAR func_fields_body
     { let FuncType (ins, out) = fst $5 in
-      FuncType (Vector.concat $3 ins, out),
+      FuncType (Vector.unsafe_concat $3 ins, out),
       fun c -> anon_locals c (lazy (vec_to_list $3)); snd $5 c }
   | LPAR PARAM bind_var value_type RPAR func_fields_body  /* Sugar */
     { let FuncType (ins, out) = fst $6 in
@@ -888,7 +888,7 @@ func_result_body :
   | func_body { FuncType (Vector.empty (), Vector.empty ()), $1 }
   | LPAR RESULT value_type_vector RPAR func_result_body
     { let FuncType (ins, out) = fst $5 in
-      FuncType (ins, Vector.concat $3 out), snd $5 }
+      FuncType (ins, Vector.unsafe_concat $3 out), snd $5 }
 
 func_body :
   | instr_list
@@ -896,7 +896,7 @@ func_body :
       {ftype = -1l @@ at(); locals = Vector.empty (); body = alloc_block c ($1 c')} }
   | LPAR LOCAL value_type_vector RPAR func_body
     { fun c -> anon_locals_vector c (lazy $3); let f = $5 c in
-      {f with locals = Vector.concat $3 f.locals} }
+      {f with locals = Vector.unsafe_concat $3 f.locals} }
   | LPAR LOCAL bind_var value_type RPAR func_body  /* Sugar */
     { fun c -> ignore (bind_local c $3); let f = $6 c in
       {f with locals = Vector.cons $4 f.locals} }
diff --git a/src/lib_webassembly/exec/eval.ml b/src/lib_webassembly/exec/eval.ml
index 20043cb64e24..24f4067afd6b 100644
--- a/src/lib_webassembly/exec/eval.ml
+++ b/src/lib_webassembly/exec/eval.ml
@@ -1090,38 +1090,32 @@ let init ~self host_funcs (m : module_) (exts : extern list) : module_inst Lwt.t
   update_module_ref self inst0 ;
 
   let* fs = TzStdLib.List.map_s (create_func self) funcs in
-  let inst1 =
-    {
-      inst0 with
-      (* TODO: #3076
-         [fs]/[funcs] should be a lazy structure so we can avoid traversing it
-         completely. *)
-      funcs = Vector.concat inst0.funcs (Vector.of_list fs);
-    }
-  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 ;
 
   let* new_globals = TzStdLib.List.map_s (create_global self) globals in
-  let inst2 =
-    {
-      inst1 with
-      tables =
-        (* TODO: #3076
-           [tables] should be a lazy structure. *)
-        List.map (create_table inst1) tables
-        |> Vector.of_list |> Vector.concat inst1.tables;
-      memories =
-        (* TODO: #3076
-           [memories] should be a lazy structure. *)
-        List.map (create_memory inst1) memories
-        |> Vector.of_list
-        |> Vector.concat inst1.memories;
-      globals =
-        (* TODO: #3076
-           [new_globals]/[globals] should be lazy structures. *)
-        Vector.concat inst1.globals (Vector.of_list new_globals);
-    }
+  (* TODO: #3076
+     [tables] should be a lazy structure. *)
+  let* tables =
+    Vector.concat
+      inst1.tables
+      (Vector.of_list (List.map (create_table inst1) tables))
   in
+  (* TODO: #3076
+     [memories] should be a lazy structure. *)
+  let* memories =
+    Vector.concat
+      inst1.memories
+      (Vector.of_list (List.map (create_memory inst1) memories))
+  in
+  (* TODO: #3076
+     [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 ;
 
   let* new_exports = TzStdLib.List.map_s (create_export inst2) exports in
diff --git a/src/lib_webassembly/runtime/partial_memory.ml b/src/lib_webassembly/runtime/partial_memory.ml
index 3214c9eb36e4..935ff7f4724d 100644
--- a/src/lib_webassembly/runtime/partial_memory.ml
+++ b/src/lib_webassembly/runtime/partial_memory.ml
@@ -24,7 +24,7 @@ module Chunked = struct
 
   let length_from_pages pages = Int64.(mul (of_int32 pages) page_size)
 
-  let create pages = Backend.create (length_from_pages pages)
+  let create pages = Backend.allocate (length_from_pages pages)
 
   let grow delta_pages pages =
     Backend.grow pages (length_from_pages delta_pages)
diff --git a/src/lib_webassembly/runtime/partial_table.ml b/src/lib_webassembly/runtime/partial_table.ml
index 98cfaccfc1c2..09370e27df70 100644
--- a/src/lib_webassembly/runtime/partial_table.ml
+++ b/src/lib_webassembly/runtime/partial_table.ml
@@ -53,7 +53,7 @@ let grow tab delta r =
   else
     let lim' = {lim with min = new_size} in
     if not (valid_limits lim') then raise SizeLimit
-    else Vector.grow delta ~produce_value:(fun _ -> Lwt.return r) tab.content ;
+    else Vector.grow delta ~default:(fun () -> r) tab.content ;
     tab.ty <- TableType (lim', t) ;
     ()
 
diff --git a/src/lib_webassembly/syntax/ast.ml b/src/lib_webassembly/syntax/ast.ml
index 7d324ab6f7f3..e873d44eb5fb 100644
--- a/src/lib_webassembly/syntax/ast.ml
+++ b/src/lib_webassembly/syntax/ast.ml
@@ -450,7 +450,7 @@ let add_to_block allocs (Block_label b) instr =
 
 let alloc_data (allocs : allocations) len =
   let datas, d =
-    Vector.append (Chunked_byte_vector.Lwt.create len) allocs.datas
+    Vector.append (Chunked_byte_vector.Lwt.allocate len) allocs.datas
   in
   allocs.datas <- datas ;
   Data_label d
diff --git a/src/lib_webassembly/valid/valid.ml b/src/lib_webassembly/valid/valid.ml
index 306e482daf52..88b14b31dab7 100644
--- a/src/lib_webassembly/valid/valid.ml
+++ b/src/lib_webassembly/valid/valid.ml
@@ -624,7 +624,7 @@ let check_func (c : context) (f : func) =
   let c' =
     {
       c with
-      locals = Lazy_vector.LwtInt32Vector.concat ts1 locals |> vec_to_list;
+      locals = vec_to_list ts1 @ vec_to_list locals;
       results = vec_to_list ts2;
       labels = [ts2];
     }
-- 
GitLab


From d4a2df55667b10cddf4b5a880f70c9fa401de711 Mon Sep 17 00:00:00 2001
From: Thomas Letan <lthms@nomadic-labs.com>
Date: Wed, 10 Aug 2022 16:56:54 +0200
Subject: [PATCH 3/5] WASM: Change the PVM to always request inputs when the
 kernel yields

---
 src/lib_scoru_wasm/wasm_pvm.ml | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/lib_scoru_wasm/wasm_pvm.ml b/src/lib_scoru_wasm/wasm_pvm.ml
index 0164d6fa3ea7..a946fd0a5cf7 100644
--- a/src/lib_scoru_wasm/wasm_pvm.ml
+++ b/src/lib_scoru_wasm/wasm_pvm.ml
@@ -189,11 +189,10 @@ module Make (T : Tree_encoding.TREE) :
       let* tree = Tree_encoding.encode module_reg_encoding module_reg tree in
       let want_more_input =
         match state.tick with
-        | Eval {input; code = _, []; _} ->
+        | Eval {code = _, []; _} ->
             (* Ask for more input if the kernel has yielded (empty admin
-               instructions) and there are no element in the input buffer any
-               more. *)
-            Z.(lt (Wasm.Input_buffer.num_elements input) one)
+               instructions). *)
+            true
         | _ -> false
       in
       let* tree = Tree_encoding.encode status_encoding want_more_input tree in
-- 
GitLab


From 2bb178bd73c30085c375012e93319941ff973abe Mon Sep 17 00:00:00 2001
From: Thomas Letan <lthms@nomadic-labs.com>
Date: Wed, 10 Aug 2022 20:11:12 +0200
Subject: [PATCH 4/5] WASM: Provide a new test-case to check the WASM PVM can
 evaluate code

---
 .../test/integration/test_sc_rollup_wasm.ml   | 82 +++++++++++++------
 1 file changed, 59 insertions(+), 23 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/test/integration/test_sc_rollup_wasm.ml b/src/proto_alpha/lib_protocol/test/integration/test_sc_rollup_wasm.ml
index 0ed8a0a0049f..0b2174edee14 100644
--- a/src/proto_alpha/lib_protocol/test/integration/test_sc_rollup_wasm.ml
+++ b/src/proto_alpha/lib_protocol/test/integration/test_sc_rollup_wasm.ml
@@ -223,29 +223,28 @@ let should_boot_complete_boot_sector boot_sector () =
   let* () = check_chunks_count s boot_sector_len in
   return_unit
 
-let floppy_input i operator chunk =
-  let open Lwt_result_syntax in
-  let signature = Signature.sign operator.Account.sk chunk in
-  let floppy = Tezos_scoru_wasm.Gather_floppies.{chunk; signature} in
-  match
-    Sc_rollup.Inbox.Message.serialize
-      (External
-         (Data_encoding.Binary.to_string_exn
-            Tezos_scoru_wasm.Gather_floppies.floppy_encoding
-            floppy))
-  with
+let arbitrary_input i payload =
+  match Sc_rollup.Inbox.Message.serialize (External payload) with
   | Ok payload ->
-      return
-        Sc_rollup.
-          {
-            inbox_level = Raw_level.of_int32_exn 0l;
-            message_counter = Z.of_int i;
-            payload;
-          }
+      Sc_rollup.
+        {
+          inbox_level = Raw_level.of_int32_exn 0l;
+          message_counter = Z.of_int i;
+          payload;
+        }
   | Error err ->
       Format.printf "%a@," Environment.Error_monad.pp_trace err ;
       assert false
 
+let floppy_input i operator chunk =
+  let signature = Signature.sign operator.Account.sk chunk in
+  let floppy = Tezos_scoru_wasm.Gather_floppies.{chunk; signature} in
+  arbitrary_input
+    i
+    (Data_encoding.Binary.to_string_exn
+       Tezos_scoru_wasm.Gather_floppies.floppy_encoding
+       floppy)
+
 let should_interpret_empty_chunk () =
   let open Lwt_result_syntax in
   let op = operator () in
@@ -256,7 +255,7 @@ let should_interpret_empty_chunk () =
     @@ incomplete_boot_sector (String.make chunk_size 'a') op
   in
   let chunk = Bytes.empty in
-  let* correct_input = floppy_input 0 op chunk in
+  let correct_input = floppy_input 0 op chunk in
 
   (* Init the PVM *)
   let*! index = Context_binary.init "/tmp" in
@@ -285,8 +284,8 @@ let should_refuse_chunks_with_incorrect_signature () =
     @@ incomplete_boot_sector (String.make chunk_size 'a') good_op
   in
   let chunk = Bytes.make chunk_size 'b' in
-  let* incorrect_input = floppy_input 0 bad_op chunk in
-  let* correct_input = floppy_input 0 good_op chunk in
+  let incorrect_input = floppy_input 0 bad_op chunk in
+  let correct_input = floppy_input 0 good_op chunk in
 
   (* Init the PVM *)
   let*! index = Context_binary.init "/tmp" in
@@ -365,7 +364,7 @@ let should_boot_incomplete_boot_sector kernel () =
       (fun i s chunk ->
         (* We are installing the [i+2]th chunk ([i] starts at 0, and
            the first chunk is not part of the list). *)
-        let* input = floppy_input i operator chunk in
+        let input = floppy_input i operator chunk in
         let* s = checked_set_input ~loc:__LOC__ context input s in
         (* We have [i+2] chunks. *)
         let* () = check_chunks_count s ((i + 2) * chunk_size) in
@@ -375,7 +374,7 @@ let should_boot_incomplete_boot_sector kernel () =
   in
   (* Up until the very last one, where the status of the PVM change. *)
   let len = List.length chunks in
-  let* input = floppy_input len operator final_chunk in
+  let input = floppy_input len operator final_chunk in
   let* s = checked_set_input ~loc:__LOC__ context input s in
   let*! () = check_status s (Some Not_gathering_floppies) in
   let* () =
@@ -396,6 +395,39 @@ let read_kernel name =
    9863 bytes long - will be split into 3 chunks. *)
 let computation_kernel () = read_kernel "computation"
 
+let rec eval_until_set_input context s =
+  let open Lwt_result_syntax in
+  let*! info = Prover.get_status s in
+  match info with
+  | Computing ->
+      let* s = checked_eval ~loc:__LOC__ context s in
+      eval_until_set_input context s
+  | Waiting_for_input_message -> return s
+
+let should_boot_computation_kernel () =
+  let open Lwt_result_syntax in
+  let boot_sector =
+    Data_encoding.Binary.to_string_exn
+      Tezos_scoru_wasm.Gather_floppies.origination_message_encoding
+      (complete_boot_sector (String.to_bytes (computation_kernel ())))
+  in
+  let*! index = Context_binary.init "/tmp" in
+  let context = Context_binary.empty index in
+  let*! s = Prover.initial_state context in
+  let*! s = Prover.install_boot_sector s boot_sector in
+  (* installing the boot kernel *)
+  let* s = checked_eval ~loc:__LOC__ context s in
+  (* make the first tick of the WASM PVM, to switch it to “waiting for
+     input” mode *)
+  let* s = checked_eval ~loc:__LOC__ context s in
+  (* Feeding it with one input *)
+  let* s =
+    checked_set_input ~loc:__LOC__ context (arbitrary_input 0 "test") s
+  in
+  (* running until waiting for input *)
+  let* _s = eval_until_set_input context s in
+  return_unit
+
 let tests =
   [
     Tztest.tztest "should boot a complete boot sector" `Quick
@@ -419,4 +451,8 @@ let tests =
       "should refuse chunks with an incorrect signature"
       `Quick
       should_refuse_chunks_with_incorrect_signature;
+    Tztest.tztest
+      "should boot a valid kernel until reading inputs"
+      `Quick
+      should_boot_computation_kernel;
   ]
-- 
GitLab


From 6600c0b3e6d9a3aad3b62cb95a0a425f6b2aeaee Mon Sep 17 00:00:00 2001
From: Joel Bjornson <joel.bjornson@trili.tech>
Date: Wed, 10 Aug 2022 20:13:47 +0200
Subject: [PATCH 5/5] WASM/Test: Add test for swapping vectors

---
 src/lib_tree_encoding/test/test_encoding.ml | 40 +++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/src/lib_tree_encoding/test/test_encoding.ml b/src/lib_tree_encoding/test/test_encoding.ml
index fd5041270dec..fd4ba86ee651 100644
--- a/src/lib_tree_encoding/test/test_encoding.ml
+++ b/src/lib_tree_encoding/test/test_encoding.ml
@@ -431,6 +431,45 @@ let test_return () =
   assert (v = "K") ;
   return_unit
 
+let test_swap_vectors () =
+  let open Tree_encoding in
+  let open Lwt_result_syntax in
+  let int_vec_enc =
+    lazy_vector (value [] Data_encoding.int31) (value [] Data_encoding.int31)
+  in
+  let enc = tup2 ~flatten:false int_vec_enc int_vec_enc in
+  let*! tree = empty_tree () in
+  let assert_value_at_index ~ix vec expected =
+    let*! value = Lazy_containers.Lazy_vector.LwtIntVector.get ix vec in
+    assert (value = expected) ;
+    return_unit
+  in
+  (* Create a pair of vectors. *)
+  let vec_pair =
+    ( Lazy_containers.Lazy_vector.LwtIntVector.create
+        ~produce_value:(fun ix -> Lwt.return ix)
+        10,
+      Lazy_containers.Lazy_vector.LwtIntVector.create
+        ~produce_value:(fun ix -> Lwt.return (100 + ix))
+        10 )
+  in
+  (* Check elements/force evaluation of one element from each vector. *)
+  let* () = assert_value_at_index ~ix:1 (fst vec_pair) 1 in
+  let* () = assert_value_at_index ~ix:2 (snd vec_pair) 102 in
+  (* Encode the lazy vector to the tree. *)
+  let*! tree = encode enc vec_pair tree in
+  (* Decode the vector. *)
+  let*! vec_pair = decode enc tree in
+  (* Encode a new pair where the elements have been swapped. *)
+  let swapped_vec_pair = (snd vec_pair, fst vec_pair) in
+  let*! tree = encode enc swapped_vec_pair tree in
+  (* Decode the swapped version. *)
+  let*! swapped_vec_pair = decode enc tree in
+  (* Check that it's possible to access the elements of both vectors. *)
+  let* () = assert_value_at_index ~ix:1 (snd swapped_vec_pair) 1 in
+  let* () = assert_value_at_index ~ix:2 (fst swapped_vec_pair) 102 in
+  return_unit
+
 let tests =
   [
     tztest "String" `Quick test_string;
@@ -454,4 +493,5 @@ let tests =
     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;
   ]
-- 
GitLab