From 7d79d5041f14e86d963322edf64c2c765f07237f Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <y.imai@proof-ninja.co.jp>
Date: Wed, 20 Jul 2022 10:07:35 +0900
Subject: [PATCH 01/15] Protocol Storage : Improve the speed of contract
 deletion

---
 .../lib_protocol/contract_storage.ml          |  15 ++-
 src/proto_alpha/lib_protocol/raw_context.ml   | 112 ++++++++++++++++++
 src/proto_alpha/lib_protocol/raw_context.mli  |   3 +
 .../lib_protocol/raw_context_intf.ml          |  26 ++++
 src/proto_alpha/lib_protocol/storage.ml       |   3 +
 src/proto_alpha/lib_protocol/storage.mli      |  18 ++-
 .../lib_protocol/storage_functors.ml          |  32 ++++-
 src/proto_alpha/lib_protocol/storage_sigs.ml  |  25 +++-
 8 files changed, 222 insertions(+), 12 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/contract_storage.ml b/src/proto_alpha/lib_protocol/contract_storage.ml
index 4e17ed75ccb6..64e4c0e21f1b 100644
--- a/src/proto_alpha/lib_protocol/contract_storage.ml
+++ b/src/proto_alpha/lib_protocol/contract_storage.ml
@@ -441,10 +441,17 @@ let delete c contract =
          do not exist). An implicit contract deletion should not cost
          extra gas. *)
       Contract_delegate_storage.unlink c contract >>=? fun c ->
-      Storage.Contract.Spendable_balance.remove_existing c contract
-      >>=? fun c ->
-      Contract_manager_storage.remove_existing c contract >>=? fun c ->
-      Storage.Contract.Counter.remove_existing c contract
+      Storage.Contract.with_local_context
+        ~for_write:true
+        c
+        contract
+        (fun local ->
+          Storage.Contract.Spendable_balance.Local.remove_existing local
+          >>=? fun local ->
+          Storage.Contract.Manager.Local.remove_existing local >>=? fun local ->
+          Storage.Contract.Counter.Local.remove_existing local >|=? fun local ->
+          (local, ()))
+      >|=? fun (c, ()) -> c
 
 let allocated c contract = Storage.Contract.Spendable_balance.mem c contract
 
diff --git a/src/proto_alpha/lib_protocol/raw_context.ml b/src/proto_alpha/lib_protocol/raw_context.ml
index a34b0145576c..b00460d5b37a 100644
--- a/src/proto_alpha/lib_protocol/raw_context.ml
+++ b/src/proto_alpha/lib_protocol/raw_context.ml
@@ -1021,12 +1021,23 @@ type value = bytes
 
 type tree = Context.tree
 
+(* The type for relative context accesses instead from the root. In order for
+   the carbonated storage functions to consume the gas, this has gas infomation
+*)
+type local_context = {
+  tree : tree;
+  path : key;
+  remaining_operation_gas : Gas_limit_repr.Arith.fp;
+  unlimited_operation_gas : bool;
+}
+
 module type T =
   Raw_context_intf.T
     with type root := root
      and type key := key
      and type value := value
      and type tree := tree
+     and type local_context := local_context
 
 let mem ctxt k = Context.mem (context ctxt) k
 
@@ -1553,3 +1564,104 @@ module Dal = struct
 
   let shards ctxt ~endorser = compute_shards ~index:0 ctxt ~endorser
 end
+
+let with_local_context ~for_write ctxt key f =
+  (find_tree ctxt key >|= function None -> Tree.empty ctxt | Some tree -> tree)
+  >>= fun tree ->
+  let local_ctxt =
+    {
+      tree;
+      path = key;
+      remaining_operation_gas = remaining_operation_gas ctxt;
+      unlimited_operation_gas = unlimited_operation_gas ctxt;
+    }
+  in
+  f local_ctxt >>=? fun (local_ctxt, res) ->
+  (if for_write then add_tree ctxt key local_ctxt.tree else Lwt.return ctxt)
+  >|= fun ctxt ->
+  update_remaining_operation_gas ctxt local_ctxt.remaining_operation_gas
+  |> fun ctxt ->
+  update_unlimited_operation_gas ctxt local_ctxt.unlimited_operation_gas
+  |> fun ctxt -> ok (ctxt, res)
+
+module Local_context : sig
+  include
+    Raw_context_intf.VIEW
+      with type t = local_context
+       and type key := key
+       and type value := value
+       and type tree := tree
+
+  val consume_gas :
+    local_context -> Gas_limit_repr.cost -> local_context tzresult
+
+  val absolute_key : local_context -> key -> key
+end = struct
+  type t = local_context
+
+  let consume_gas local cost =
+    match Gas_limit_repr.raw_consume local.remaining_operation_gas cost with
+    | Some gas_counter -> Ok {local with remaining_operation_gas = gas_counter}
+    | None ->
+        if local.unlimited_operation_gas then ok local
+        else error Operation_quota_exceeded
+
+  let tree local = local.tree
+
+  let update_root_tree local tree = {local with tree}
+
+  let absolute_key local key = local.path @ key
+
+  let find local = Tree.find (tree local)
+
+  let find_tree local = Tree.find_tree (tree local)
+
+  let mem local = Tree.mem (tree local)
+
+  let mem_tree local = Tree.mem_tree (tree local)
+
+  let get local = Tree.get (tree local)
+
+  let get_tree local = Tree.get_tree (tree local)
+
+  let update local key b =
+    Tree.update (tree local) key b >|=? update_root_tree local
+
+  let update_tree local key b =
+    Tree.update_tree (tree local) key b >|=? update_root_tree local
+
+  let init local key b =
+    Tree.init (tree local) key b >|=? update_root_tree local
+
+  let init_tree local key t =
+    Tree.init_tree (tree local) key t >|=? update_root_tree local
+
+  let add local i b = Tree.add (tree local) i b >|= update_root_tree local
+
+  let add_tree local i t =
+    Tree.add_tree (tree local) i t >|= update_root_tree local
+
+  let remove local i = Tree.remove (tree local) i >|= update_root_tree local
+
+  let remove_existing local key =
+    Tree.remove_existing (tree local) key >|=? update_root_tree local
+
+  let remove_existing_tree local key =
+    Tree.remove_existing_tree (tree local) key >|=? update_root_tree local
+
+  let add_or_remove local key vopt =
+    Tree.add_or_remove (tree local) key vopt >|= update_root_tree local
+
+  let add_or_remove_tree local key topt =
+    Tree.add_or_remove_tree (tree local) key topt >|= update_root_tree local
+
+  let fold ?depth local key ~order ~init ~f =
+    Tree.fold ?depth (tree local) key ~order ~init ~f
+
+  let list local ?offset ?length key =
+    Tree.list (tree local) ?offset ?length key
+
+  let config local = Tree.config (tree local)
+
+  let length local i = Tree.length (tree local) i
+end
diff --git a/src/proto_alpha/lib_protocol/raw_context.mli b/src/proto_alpha/lib_protocol/raw_context.mli
index eef60d8c3117..8c44aaf2752e 100644
--- a/src/proto_alpha/lib_protocol/raw_context.mli
+++ b/src/proto_alpha/lib_protocol/raw_context.mli
@@ -181,12 +181,15 @@ type value = bytes
 
 type tree
 
+type local_context
+
 module type T =
   Raw_context_intf.T
     with type root := root
      and type key := key
      and type value := value
      and type tree := tree
+     and type local_context := local_context
 
 include T with type t := t
 
diff --git a/src/proto_alpha/lib_protocol/raw_context_intf.ml b/src/proto_alpha/lib_protocol/raw_context_intf.ml
index 39c8b058d78a..a721b371e556 100644
--- a/src/proto_alpha/lib_protocol/raw_context_intf.ml
+++ b/src/proto_alpha/lib_protocol/raw_context_intf.ml
@@ -457,6 +457,11 @@ module type T = sig
 
   include VIEW
 
+  (** The type for relative context accesses instead from the root. In order for
+      the carbonated storage functions to consume the gas, this has gas infomation
+  *)
+  type local_context
+
   module Tree :
     TREE
       with type t := t
@@ -559,4 +564,25 @@ module type T = sig
   val check_enough_gas : t -> Gas_limit_repr.cost -> unit tzresult
 
   val description : t Storage_description.t
+
+  val with_local_context :
+    for_write:bool ->
+    t ->
+    key ->
+    (local_context -> (local_context * 'a) tzresult Lwt.t) ->
+    (t * 'a) tzresult Lwt.t
+
+  module Local_context : sig
+    include
+      VIEW
+        with type t = local_context
+         and type tree := tree
+         and type key := key
+         and type value := value
+
+    val consume_gas :
+      local_context -> Gas_limit_repr.cost -> local_context tzresult
+
+    val absolute_key : local_context -> key -> key
+  end
 end
diff --git a/src/proto_alpha/lib_protocol/storage.ml b/src/proto_alpha/lib_protocol/storage.ml
index 637404e7263f..2c77dc861a75 100644
--- a/src/proto_alpha/lib_protocol/storage.ml
+++ b/src/proto_alpha/lib_protocol/storage.ml
@@ -188,6 +188,9 @@ module Contract = struct
 
   let list = Indexed_context.keys
 
+  let with_local_context ~for_write =
+    Indexed_context.with_local_context ~for_write
+
   module Spendable_balance =
     Indexed_context.Make_map
       (struct
diff --git a/src/proto_alpha/lib_protocol/storage.mli b/src/proto_alpha/lib_protocol/storage.mli
index c4c00f28d58f..178f271e4ef2 100644
--- a/src/proto_alpha/lib_protocol/storage.mli
+++ b/src/proto_alpha/lib_protocol/storage.mli
@@ -69,12 +69,20 @@ module Contract : sig
 
   val list : Raw_context.t -> Contract_repr.t list Lwt.t
 
+  val with_local_context :
+    for_write:bool ->
+    Raw_context.t ->
+    Contract_repr.t ->
+    (Raw_context.local_context ->
+    (Raw_context.local_context * 'a) tzresult Lwt.t) ->
+    (Raw_context.t * 'a) tzresult Lwt.t
+
   (** The tez possessed by a contract and that can be used. A contract
      may also possess tez in frozen deposits. Empty balances (of zero
      tez) are only allowed for originated contracts, not for implicit
      ones. *)
   module Spendable_balance :
-    Indexed_data_storage
+    Indexed_data_storage_with_local_context
       with type key = Contract_repr.t
        and type value = Tez_repr.t
        and type t := Raw_context.t
@@ -99,7 +107,7 @@ module Contract : sig
 
   (** The manager of a contract *)
   module Manager :
-    Indexed_data_storage
+    Indexed_data_storage_with_local_context
       with type key = Contract_repr.t
        and type value = Manager_repr.t
        and type t := Raw_context.t
@@ -148,7 +156,7 @@ module Contract : sig
        and type t := Raw_context.t
 
   module Counter :
-    Indexed_data_storage
+    Indexed_data_storage_with_local_context
       with type key = Contract_repr.t
        and type value = Z.t
        and type t := Raw_context.t
@@ -168,14 +176,14 @@ module Contract : sig
   (** Current storage space in bytes.
       Includes code, global storage and big map elements. *)
   module Used_storage_space :
-    Indexed_data_storage
+    Indexed_data_storage_with_local_context
       with type key = Contract_repr.t
        and type value = Z.t
        and type t := Raw_context.t
 
   (** Maximal space available without needing to burn new fees. *)
   module Paid_storage_space :
-    Indexed_data_storage
+    Indexed_data_storage_with_local_context
       with type key = Contract_repr.t
        and type value = Z.t
        and type t := Raw_context.t
diff --git a/src/proto_alpha/lib_protocol/storage_functors.ml b/src/proto_alpha/lib_protocol/storage_functors.ml
index b07dcbe31d9e..34122b7f2de2 100644
--- a/src/proto_alpha/lib_protocol/storage_functors.ml
+++ b/src/proto_alpha/lib_protocol/storage_functors.ml
@@ -142,6 +142,11 @@ module Make_subcontext (R : REGISTER) (C : Raw_context.T) (N : NAME) :
     Storage_description.register_named_subcontext description N.name
 
   let length = C.length
+
+  let with_local_context ~for_write ctxt k f =
+    C.with_local_context ~for_write ctxt (to_key k) f
+
+  module Local_context = C.Local_context
 end
 
 module Make_single_data_storage
@@ -824,8 +829,20 @@ module Make_indexed_subcontext (C : Raw_context.T) (I : INDEX) :
     let length c =
       let t, _i = unpack c in
       C.length t
+
+    let with_local_context ~for_write c k f =
+      let t, i = unpack c in
+      C.with_local_context ~for_write t (to_key i k) f >|=? fun (t, res) ->
+      (pack t i, res)
+
+    module Local_context = C.Local_context
   end
 
+  let with_local_context ~for_write s i f =
+    Raw_context.with_local_context ~for_write (pack s i) [] f >|=? fun (c, x) ->
+    let s, _ = unpack c in
+    (s, x)
+
   module Make_set (R : REGISTER) (N : NAME) :
     Data_set_storage with type t = t and type elt = key = struct
     type t = C.t
@@ -878,8 +895,10 @@ module Make_indexed_subcontext (C : Raw_context.T) (I : INDEX) :
   end
 
   module Make_map (N : NAME) (V : VALUE) :
-    Indexed_data_storage with type t = t and type key = key and type value = V.t =
-  struct
+    Indexed_data_storage_with_local_context
+      with type t = t
+       and type key = key
+       and type value = V.t = struct
     type t = C.t
 
     type context = t
@@ -967,6 +986,15 @@ module Make_indexed_subcontext (C : Raw_context.T) (I : INDEX) :
           find c k)
         (register_named_subcontext Raw_context.description N.name)
         V.encoding
+
+    module Local = struct
+      type context = Raw_context.Local_context.t
+
+      let remove_existing local =
+        Raw_context.Local_context.remove_existing local N.name
+
+      let remove local = Raw_context.Local_context.remove local N.name
+    end
   end
 
   module Make_carbonated_map (N : NAME) (V : VALUE) :
diff --git a/src/proto_alpha/lib_protocol/storage_sigs.ml b/src/proto_alpha/lib_protocol/storage_sigs.ml
index d6c0a2886083..e504346b1ed8 100644
--- a/src/proto_alpha/lib_protocol/storage_sigs.ml
+++ b/src/proto_alpha/lib_protocol/storage_sigs.ml
@@ -269,6 +269,18 @@ module type Indexed_data_storage = sig
     'a Lwt.t
 end
 
+module type Indexed_data_storage_with_local_context = sig
+  include Indexed_data_storage
+
+  module Local : sig
+    type context = Raw_context.local_context
+
+    val remove_existing : context -> context tzresult Lwt.t
+
+    val remove : context -> context Lwt.t
+  end
+end
+
 module type Indexed_data_snapshotable_storage = sig
   type snapshot
 
@@ -408,11 +420,22 @@ module type Indexed_raw_context = sig
 
   val copy : context -> from:key -> to_:key -> context tzresult Lwt.t
 
+  val with_local_context :
+    for_write:bool ->
+    context ->
+    key ->
+    (Raw_context.local_context ->
+    (Raw_context.local_context * 'a) tzresult Lwt.t) ->
+    (context * 'a) tzresult Lwt.t
+
   module Make_set (_ : REGISTER) (_ : NAME) :
     Data_set_storage with type t = t and type elt = key
 
   module Make_map (_ : NAME) (V : VALUE) :
-    Indexed_data_storage with type t = t and type key = key and type value = V.t
+    Indexed_data_storage_with_local_context
+      with type t = t
+       and type key = key
+       and type value = V.t
 
   module Make_carbonated_map (_ : NAME) (V : VALUE) :
     Non_iterable_indexed_carbonated_data_storage
-- 
GitLab


From f2b823cc5b0fe9b063e7babd6f22f22e4abed74d Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <y.imai@proof-ninja.co.jp>
Date: Wed, 27 Jul 2022 12:23:44 +0900
Subject: [PATCH 02/15] Protocol Storage: small changes for reviews *
 https://gitlab.com/tezos/tezos/-/merge_requests/5922#note_1039716605 *
 https://gitlab.com/tezos/tezos/-/merge_requests/5922#note_1039719299 *
 https://gitlab.com/tezos/tezos/-/merge_requests/5922#note_1039719695

---
 src/proto_alpha/lib_protocol/contract_storage.ml |  2 +-
 src/proto_alpha/lib_protocol/raw_context.ml      |  4 ++--
 src/proto_alpha/lib_protocol/raw_context_intf.ml | 12 ++++++------
 src/proto_alpha/lib_protocol/storage.ml          |  3 +--
 src/proto_alpha/lib_protocol/storage.mli         |  2 +-
 src/proto_alpha/lib_protocol/storage_functors.ml | 12 ++++++------
 src/proto_alpha/lib_protocol/storage_sigs.ml     |  2 +-
 7 files changed, 18 insertions(+), 19 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/contract_storage.ml b/src/proto_alpha/lib_protocol/contract_storage.ml
index 64e4c0e21f1b..05b6cdba17d1 100644
--- a/src/proto_alpha/lib_protocol/contract_storage.ml
+++ b/src/proto_alpha/lib_protocol/contract_storage.ml
@@ -442,7 +442,7 @@ let delete c contract =
          extra gas. *)
       Contract_delegate_storage.unlink c contract >>=? fun c ->
       Storage.Contract.with_local_context
-        ~for_write:true
+        ~add_back:true
         c
         contract
         (fun local ->
diff --git a/src/proto_alpha/lib_protocol/raw_context.ml b/src/proto_alpha/lib_protocol/raw_context.ml
index b00460d5b37a..79caf278d4d3 100644
--- a/src/proto_alpha/lib_protocol/raw_context.ml
+++ b/src/proto_alpha/lib_protocol/raw_context.ml
@@ -1565,7 +1565,7 @@ module Dal = struct
   let shards ctxt ~endorser = compute_shards ~index:0 ctxt ~endorser
 end
 
-let with_local_context ~for_write ctxt key f =
+let with_local_context ~add_back ctxt key f =
   (find_tree ctxt key >|= function None -> Tree.empty ctxt | Some tree -> tree)
   >>= fun tree ->
   let local_ctxt =
@@ -1577,7 +1577,7 @@ let with_local_context ~for_write ctxt key f =
     }
   in
   f local_ctxt >>=? fun (local_ctxt, res) ->
-  (if for_write then add_tree ctxt key local_ctxt.tree else Lwt.return ctxt)
+  (if add_back then add_tree ctxt key local_ctxt.tree else Lwt.return ctxt)
   >|= fun ctxt ->
   update_remaining_operation_gas ctxt local_ctxt.remaining_operation_gas
   |> fun ctxt ->
diff --git a/src/proto_alpha/lib_protocol/raw_context_intf.ml b/src/proto_alpha/lib_protocol/raw_context_intf.ml
index a721b371e556..f6534819ea33 100644
--- a/src/proto_alpha/lib_protocol/raw_context_intf.ml
+++ b/src/proto_alpha/lib_protocol/raw_context_intf.ml
@@ -457,11 +457,6 @@ module type T = sig
 
   include VIEW
 
-  (** The type for relative context accesses instead from the root. In order for
-      the carbonated storage functions to consume the gas, this has gas infomation
-  *)
-  type local_context
-
   module Tree :
     TREE
       with type t := t
@@ -565,8 +560,13 @@ module type T = sig
 
   val description : t Storage_description.t
 
+  (** The type for relative context accesses instead from the root. In order for
+      the carbonated storage functions to consume the gas, this has gas infomation
+  *)
+  type local_context
+
   val with_local_context :
-    for_write:bool ->
+    add_back:bool ->
     t ->
     key ->
     (local_context -> (local_context * 'a) tzresult Lwt.t) ->
diff --git a/src/proto_alpha/lib_protocol/storage.ml b/src/proto_alpha/lib_protocol/storage.ml
index 2c77dc861a75..ce38e6f01a0b 100644
--- a/src/proto_alpha/lib_protocol/storage.ml
+++ b/src/proto_alpha/lib_protocol/storage.ml
@@ -188,8 +188,7 @@ module Contract = struct
 
   let list = Indexed_context.keys
 
-  let with_local_context ~for_write =
-    Indexed_context.with_local_context ~for_write
+  let with_local_context = Indexed_context.with_local_context
 
   module Spendable_balance =
     Indexed_context.Make_map
diff --git a/src/proto_alpha/lib_protocol/storage.mli b/src/proto_alpha/lib_protocol/storage.mli
index 178f271e4ef2..929f77d7ed92 100644
--- a/src/proto_alpha/lib_protocol/storage.mli
+++ b/src/proto_alpha/lib_protocol/storage.mli
@@ -70,7 +70,7 @@ module Contract : sig
   val list : Raw_context.t -> Contract_repr.t list Lwt.t
 
   val with_local_context :
-    for_write:bool ->
+    add_back:bool ->
     Raw_context.t ->
     Contract_repr.t ->
     (Raw_context.local_context ->
diff --git a/src/proto_alpha/lib_protocol/storage_functors.ml b/src/proto_alpha/lib_protocol/storage_functors.ml
index 34122b7f2de2..4ef03e9ca993 100644
--- a/src/proto_alpha/lib_protocol/storage_functors.ml
+++ b/src/proto_alpha/lib_protocol/storage_functors.ml
@@ -143,8 +143,8 @@ module Make_subcontext (R : REGISTER) (C : Raw_context.T) (N : NAME) :
 
   let length = C.length
 
-  let with_local_context ~for_write ctxt k f =
-    C.with_local_context ~for_write ctxt (to_key k) f
+  let with_local_context ~add_back ctxt k f =
+    C.with_local_context ~add_back ctxt (to_key k) f
 
   module Local_context = C.Local_context
 end
@@ -830,16 +830,16 @@ module Make_indexed_subcontext (C : Raw_context.T) (I : INDEX) :
       let t, _i = unpack c in
       C.length t
 
-    let with_local_context ~for_write c k f =
+    let with_local_context ~add_back c k f =
       let t, i = unpack c in
-      C.with_local_context ~for_write t (to_key i k) f >|=? fun (t, res) ->
+      C.with_local_context ~add_back t (to_key i k) f >|=? fun (t, res) ->
       (pack t i, res)
 
     module Local_context = C.Local_context
   end
 
-  let with_local_context ~for_write s i f =
-    Raw_context.with_local_context ~for_write (pack s i) [] f >|=? fun (c, x) ->
+  let with_local_context ~add_back s i f =
+    Raw_context.with_local_context ~add_back (pack s i) [] f >|=? fun (c, x) ->
     let s, _ = unpack c in
     (s, x)
 
diff --git a/src/proto_alpha/lib_protocol/storage_sigs.ml b/src/proto_alpha/lib_protocol/storage_sigs.ml
index e504346b1ed8..d5065899e3bd 100644
--- a/src/proto_alpha/lib_protocol/storage_sigs.ml
+++ b/src/proto_alpha/lib_protocol/storage_sigs.ml
@@ -421,7 +421,7 @@ module type Indexed_raw_context = sig
   val copy : context -> from:key -> to_:key -> context tzresult Lwt.t
 
   val with_local_context :
-    for_write:bool ->
+    add_back:bool ->
     context ->
     key ->
     (Raw_context.local_context ->
-- 
GitLab


From aaa7432aab1a5cc10a2ac15abc180577d1a666c3 Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <y.imai@proof-ninja.co.jp>
Date: Wed, 27 Jul 2022 14:22:43 +0900
Subject: [PATCH 03/15] Protocol Storage : Make local_type a specia type *
 https://gitlab.com/tezos/tezos/-/merge_requests/5922#note_1039734752

---
 src/proto_alpha/lib_protocol/storage.ml          |  2 ++
 src/proto_alpha/lib_protocol/storage.mli         |  8 ++++++--
 src/proto_alpha/lib_protocol/storage_functors.ml |  7 ++++++-
 src/proto_alpha/lib_protocol/storage_sigs.ml     | 10 +++++++---
 4 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/storage.ml b/src/proto_alpha/lib_protocol/storage.ml
index ce38e6f01a0b..0902183524b8 100644
--- a/src/proto_alpha/lib_protocol/storage.ml
+++ b/src/proto_alpha/lib_protocol/storage.ml
@@ -188,6 +188,8 @@ module Contract = struct
 
   let list = Indexed_context.keys
 
+  type local_context = Indexed_context.local_context
+
   let with_local_context = Indexed_context.with_local_context
 
   module Spendable_balance =
diff --git a/src/proto_alpha/lib_protocol/storage.mli b/src/proto_alpha/lib_protocol/storage.mli
index 929f77d7ed92..bc1845e62d1c 100644
--- a/src/proto_alpha/lib_protocol/storage.mli
+++ b/src/proto_alpha/lib_protocol/storage.mli
@@ -69,12 +69,13 @@ module Contract : sig
 
   val list : Raw_context.t -> Contract_repr.t list Lwt.t
 
+  type local_context
+
   val with_local_context :
     add_back:bool ->
     Raw_context.t ->
     Contract_repr.t ->
-    (Raw_context.local_context ->
-    (Raw_context.local_context * 'a) tzresult Lwt.t) ->
+    (local_context -> (local_context * 'a) tzresult Lwt.t) ->
     (Raw_context.t * 'a) tzresult Lwt.t
 
   (** The tez possessed by a contract and that can be used. A contract
@@ -86,6 +87,7 @@ module Contract : sig
       with type key = Contract_repr.t
        and type value = Tez_repr.t
        and type t := Raw_context.t
+       and type local_context := local_context
 
   (** If the value is not set, the delegate didn't miss any endorsing
      opportunity.  If it is set, this value is a record of type
@@ -111,6 +113,7 @@ module Contract : sig
       with type key = Contract_repr.t
        and type value = Manager_repr.t
        and type t := Raw_context.t
+       and type local_context := local_context
 
   (** The delegate of a contract, if any. *)
   module Delegate :
@@ -160,6 +163,7 @@ module Contract : sig
       with type key = Contract_repr.t
        and type value = Z.t
        and type t := Raw_context.t
+       and type local_context := local_context
 
   module Code :
     Non_iterable_indexed_carbonated_data_storage
diff --git a/src/proto_alpha/lib_protocol/storage_functors.ml b/src/proto_alpha/lib_protocol/storage_functors.ml
index 4ef03e9ca993..56b095df73b6 100644
--- a/src/proto_alpha/lib_protocol/storage_functors.ml
+++ b/src/proto_alpha/lib_protocol/storage_functors.ml
@@ -668,6 +668,8 @@ module Make_indexed_subcontext (C : Raw_context.T) (I : INDEX) :
 
   type 'a ipath = 'a I.ipath
 
+  type local_context = Raw_context.local_context
+
   let clear t = C.remove t [] >|= fun t -> C.project t
 
   let fold_keys t ~order ~init ~f =
@@ -898,7 +900,8 @@ module Make_indexed_subcontext (C : Raw_context.T) (I : INDEX) :
     Indexed_data_storage_with_local_context
       with type t = t
        and type key = key
-       and type value = V.t = struct
+       and type value = V.t
+       and type local_context = local_context = struct
     type t = C.t
 
     type context = t
@@ -907,6 +910,8 @@ module Make_indexed_subcontext (C : Raw_context.T) (I : INDEX) :
 
     type value = V.t
 
+    type nonrec local_context = local_context
+
     include Make_encoder (V)
 
     let mem s i = Raw_context.mem (pack s i) N.name
diff --git a/src/proto_alpha/lib_protocol/storage_sigs.ml b/src/proto_alpha/lib_protocol/storage_sigs.ml
index d5065899e3bd..0fd43629c3a3 100644
--- a/src/proto_alpha/lib_protocol/storage_sigs.ml
+++ b/src/proto_alpha/lib_protocol/storage_sigs.ml
@@ -272,8 +272,10 @@ end
 module type Indexed_data_storage_with_local_context = sig
   include Indexed_data_storage
 
+  type local_context
+
   module Local : sig
-    type context = Raw_context.local_context
+    type context = local_context
 
     val remove_existing : context -> context tzresult Lwt.t
 
@@ -405,6 +407,8 @@ module type Indexed_raw_context = sig
 
   type 'a ipath
 
+  type local_context
+
   val clear : context -> Raw_context.t Lwt.t
 
   val fold_keys :
@@ -424,8 +428,7 @@ module type Indexed_raw_context = sig
     add_back:bool ->
     context ->
     key ->
-    (Raw_context.local_context ->
-    (Raw_context.local_context * 'a) tzresult Lwt.t) ->
+    (local_context -> (local_context * 'a) tzresult Lwt.t) ->
     (context * 'a) tzresult Lwt.t
 
   module Make_set (_ : REGISTER) (_ : NAME) :
@@ -436,6 +439,7 @@ module type Indexed_raw_context = sig
       with type t = t
        and type key = key
        and type value = V.t
+       and type local_context = local_context
 
   module Make_carbonated_map (_ : NAME) (V : VALUE) :
     Non_iterable_indexed_carbonated_data_storage
-- 
GitLab


From 438d3514be88bfc0254d45ab775fdf44c6dc7d29 Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <y.imai@proof-ninja.co.jp>
Date: Wed, 27 Jul 2022 19:11:56 +0900
Subject: [PATCH 04/15] Protocol Storage : Add other functions to Local module

---
 .../lib_protocol/storage_functors.ml          | 28 +++++++++++++++
 src/proto_alpha/lib_protocol/storage_sigs.ml  | 35 +++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/src/proto_alpha/lib_protocol/storage_functors.ml b/src/proto_alpha/lib_protocol/storage_functors.ml
index 56b095df73b6..717816e4d286 100644
--- a/src/proto_alpha/lib_protocol/storage_functors.ml
+++ b/src/proto_alpha/lib_protocol/storage_functors.ml
@@ -995,6 +995,34 @@ module Make_indexed_subcontext (C : Raw_context.T) (I : INDEX) :
     module Local = struct
       type context = Raw_context.Local_context.t
 
+      let mem local = Raw_context.Local_context.mem local N.name
+
+      let get local =
+        Raw_context.Local_context.get local N.name >|= fun r ->
+        let key () = Raw_context.Local_context.absolute_key local N.name in
+        r >>? of_bytes ~key
+
+      let find local =
+        Raw_context.Local_context.find local N.name >|= function
+        | None -> Result.return_none
+        | Some b ->
+            let key () = Raw_context.Local_context.absolute_key local N.name in
+            of_bytes ~key b >|? fun v -> Some v
+
+      let init local v =
+        Raw_context.Local_context.init local N.name (to_bytes v)
+
+      let update local v =
+        Raw_context.Local_context.update local N.name (to_bytes v)
+
+      let add local v = Raw_context.Local_context.add local N.name (to_bytes v)
+
+      let add_or_remove local vo =
+        Raw_context.Local_context.add_or_remove
+          local
+          N.name
+          (Option.map to_bytes vo)
+
       let remove_existing local =
         Raw_context.Local_context.remove_existing local N.name
 
diff --git a/src/proto_alpha/lib_protocol/storage_sigs.ml b/src/proto_alpha/lib_protocol/storage_sigs.ml
index 0fd43629c3a3..3921221ad5c3 100644
--- a/src/proto_alpha/lib_protocol/storage_sigs.ml
+++ b/src/proto_alpha/lib_protocol/storage_sigs.ml
@@ -277,8 +277,43 @@ module type Indexed_data_storage_with_local_context = sig
   module Local : sig
     type context = local_context
 
+    (** Tells if the data is already defined *)
+    val mem : context -> bool Lwt.t
+
+    (** Retrieve the value from the storage bucket ; returns a
+        {!Storage_error} if the key is not set or if the deserialisation
+        fails *)
+    val get : context -> value tzresult Lwt.t
+
+    (** Retrieves the value from the storage bucket ; returns [None] if
+        the data is not initialized, or {!Storage_helpers.Storage_error}
+        if the deserialisation fails *)
+    val find : context -> value option tzresult Lwt.t
+
+    (** Allocates the storage bucket and initializes it ; returns a
+        {!Storage_error Existing_key} if the bucket exists *)
+    val init : context -> value -> context tzresult Lwt.t
+
+    (** Updates the content of the bucket ; returns a {!Storage_Error
+        Missing_key} if the value does not exists *)
+    val update : context -> value -> context tzresult Lwt.t
+
+    (** Allocates the data and initializes it with a value ; just
+        updates it if the bucket exists *)
+    val add : context -> value -> context Lwt.t
+
+    (** When the value is [Some v], allocates the data and initializes
+        it with [v] ; just updates it if the bucket exists. When the
+        value is [None], delete the storage bucket when the value ; does
+        nothing if the bucket does not exists. *)
+    val add_or_remove : context -> value option -> context Lwt.t
+
+    (** Delete the storage bucket ; returns a {!Storage_error
+        Missing_key} if the bucket does not exists *)
     val remove_existing : context -> context tzresult Lwt.t
 
+    (** Removes the storage bucket and its contents ; does nothing if
+        the bucket does not exists *)
     val remove : context -> context Lwt.t
   end
 end
-- 
GitLab


From 3ed85996227a184084e6fdbb21d7d1c6bd44ef64 Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <yoshihiro503@proofcafe.org>
Date: Thu, 17 Jun 2021 10:32:55 +0900
Subject: [PATCH 05/15] Storage Protocol: optimize Big_map updating using local
 tree

Protocol Storage: introduce gas_info

Protocol Storage: show tree type

Protocol Storage: Add local access functions to Make_indexed_carbonated_data_storage

Protocol Storage: local remove and add for Big_map.Constants

Protocol Storage: update gas

Storage Protocol : gas_info set/get from constexts

Protocol Storage : using local api

Protocol Storage : apply update diff using local access

Protocol Storage : fmt

Protocol Storage: rebase on master

Storage protocol : choose improved one

Storage protocol: licenses

Protocol Storage: @fmt

Storage protocol: ..

Storage Protocol: introduce local context

Storage Protocol: fix the problem of add local

Protocol Storage: rename functions for reviews:
* https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_806513139
* https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_806514454
* https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_806515367
* https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_806536934

Protocol Storage: hide tree type

Protocol Storage: move local_context to the interface Raw_context.T

Protocol storage: add function apis for local access to Non_iterable_indexed_carbonated_storage_with_local_context

Protocol Storage: rename

Protocol Storage: Remove unneeded type interfaces
https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_827410999

Protocol Storage: Expose only abstract function to treat local context
(thanks the review: https://gitlab.com/tezos/tezos/-/merge_requests/3685)

Protocol Storage: introduce Local sub module
Thanks for the review https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_835825685

Protocol Storage: revert an unnecessary change
review: https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_835788002

Storage protocol : refactor for reviews:
* using Option.map: https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_842760659
* rename to Local_context : https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_842761340
* remove a unnecessary comment https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_842763436
* fix a comment: https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_842763559

Protocol Storage: refactor
* rename using_... to with_...: https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_842776474

Protocol Storage: remove unneccessary rename
for https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_862461337

Protocol Storage: renamed
for the review https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_862452452

Storage Protocol: rename variables
for the review: https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_862490858

Protocol Storage: delete unnecessary lines

Protocol Storage : Some functions are defined locally to prevent mosusing.
Thanks to the review: https://gitlab.com/tezos/tezos/-/merge_requests/3685/diffs#note_866531180

Protocol Storage : fix comments.
* https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_881720668
* https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_881727142
* https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_881772339

Protocol Storage : move lines of Local module
* https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_882018721

Protocol Storage : avoid hardcording of ["contents"] directory path.
* https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_881694117
* https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_881693055

Protocol Storage : fix for reviews
* https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_882418153
* https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_882367554
* https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_882329904

Protocol Storage : deserializing gas cost
* for the review: https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_882333499

Protocol Storage : move the function `with_local_context` to more generic signature
* for the review: https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_882337575

Protocol Storage : fix a comment
review: https://gitlab.com/tezos/tezos/-/merge_requests/3685/diffs#note_882329243

Protocol Storage : rebased
---
 .../lib_protocol/lazy_storage_diff.ml         |  29 ++--
 src/proto_alpha/lib_protocol/raw_context.ml   |  50 ++++--
 .../lib_protocol/raw_context_intf.ml          |  14 +-
 src/proto_alpha/lib_protocol/storage.ml       |  28 ++++
 src/proto_alpha/lib_protocol/storage.mli      |   2 +-
 .../lib_protocol/storage_functors.ml          | 146 +++++++++++++++++-
 .../lib_protocol/storage_functors.mli         |   2 +-
 src/proto_alpha/lib_protocol/storage_sigs.ml  |  35 ++++-
 8 files changed, 271 insertions(+), 35 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/lazy_storage_diff.ml b/src/proto_alpha/lib_protocol/lazy_storage_diff.ml
index 791b3f828448..53cdddde6397 100644
--- a/src/proto_alpha/lib_protocol/lazy_storage_diff.ml
+++ b/src/proto_alpha/lib_protocol/lazy_storage_diff.ml
@@ -118,7 +118,8 @@ module Big_map = struct
     Storage.Big_map.Key_type.init ctxt id key_type >>=? fun ctxt ->
     Storage.Big_map.Value_type.init ctxt id value_type
 
-  let apply_update ctxt ~id
+  (* [local_ctxt] points to the directory of the big_map *)
+  let apply_update local_ctxt
       {
         key = _key_is_shown_only_on_the_receipt_in_print_big_map_diff;
         key_hash;
@@ -126,27 +127,29 @@ module Big_map = struct
       } =
     match value with
     | None ->
-        Storage.Big_map.Contents.remove (ctxt, id) key_hash
-        >|=? fun (ctxt, freed, existed) ->
+        Storage.Big_map.Contents.Local.remove local_ctxt key_hash
+        >|=? fun (local_ctxt, freed, existed) ->
         let freed =
           if existed then freed + bytes_size_for_big_map_key else freed
         in
-        (ctxt, Z.of_int ~-freed)
+        (local_ctxt, Z.of_int ~-freed)
     | Some v ->
-        Storage.Big_map.Contents.add (ctxt, id) key_hash v
-        >|=? fun (ctxt, size_diff, existed) ->
+        Storage.Big_map.Contents.Local.add local_ctxt key_hash v
+        >|=? fun (local_ctxt, size_diff, existed) ->
         let size_diff =
           if existed then size_diff else size_diff + bytes_size_for_big_map_key
         in
-        (ctxt, Z.of_int size_diff)
+        (local_ctxt, Z.of_int size_diff)
 
   let apply_updates ctxt ~id updates =
-    List.fold_left_es
-      (fun (ctxt, size) update ->
-        apply_update ctxt ~id update >|=? fun (ctxt, added_size) ->
-        (ctxt, Z.add size added_size))
-      (ctxt, Z.zero)
-      updates
+    Storage.Big_map.Contents.with_local_context (ctxt, id) (fun local_ctxt ->
+        List.fold_left_es
+          (fun (local_ctxt, size) update ->
+            apply_update local_ctxt update >|=? fun (local_ctxt, added_size) ->
+            (local_ctxt, Z.add size added_size))
+          (local_ctxt, Z.zero)
+          updates)
+    >|=? fun ((ctxt, _id), size_diff) -> (ctxt, size_diff)
 
   include Storage.Big_map
 end
diff --git a/src/proto_alpha/lib_protocol/raw_context.ml b/src/proto_alpha/lib_protocol/raw_context.ml
index 79caf278d4d3..c4072bb1762c 100644
--- a/src/proto_alpha/lib_protocol/raw_context.ml
+++ b/src/proto_alpha/lib_protocol/raw_context.ml
@@ -1584,19 +1584,43 @@ let with_local_context ~add_back ctxt key f =
   update_unlimited_operation_gas ctxt local_ctxt.unlimited_operation_gas
   |> fun ctxt -> ok (ctxt, res)
 
-module Local_context : sig
-  include
-    Raw_context_intf.VIEW
-      with type t = local_context
-       and type key := key
-       and type value := value
-       and type tree := tree
-
-  val consume_gas :
-    local_context -> Gas_limit_repr.cost -> local_context tzresult
-
-  val absolute_key : local_context -> key -> key
-end = struct
+let make_local_context ctxt path =
+  {
+    tree = Tree.empty ctxt;
+    path;
+    remaining_operation_gas = remaining_operation_gas ctxt;
+    unlimited_operation_gas = unlimited_operation_gas ctxt;
+  }
+
+let with_local_context ctxt key f =
+  let find_local_context ctxt key =
+    find_tree ctxt key
+    >|= Option.map (fun tree ->
+            {
+              tree;
+              path = key;
+              remaining_operation_gas = remaining_operation_gas ctxt;
+              unlimited_operation_gas = unlimited_operation_gas ctxt;
+            })
+  in
+  (find_local_context ctxt key >>= function
+   | Some local_ctxt ->
+       f local_ctxt >>=? fun (local_ctxt, res) ->
+       update_tree ctxt key local_ctxt.tree >|=? fun ctxt ->
+       (ctxt, local_ctxt, res)
+   | None ->
+       let local_ctxt = make_local_context ctxt key in
+       f local_ctxt >>=? fun (local_ctxt, res) ->
+       add_tree ctxt key local_ctxt.tree >|= fun ctxt ->
+       ok (ctxt, local_ctxt, res))
+  >|=? fun (ctxt, local_ctxt, res) ->
+  update_remaining_operation_gas ctxt local_ctxt.remaining_operation_gas
+  |> fun ctxt ->
+  update_unlimited_operation_gas ctxt local_ctxt.unlimited_operation_gas
+  |> fun ctxt -> (ctxt, res)
+
+
+module Local_context = struct
   type t = local_context
 
   let consume_gas local cost =
diff --git a/src/proto_alpha/lib_protocol/raw_context_intf.ml b/src/proto_alpha/lib_protocol/raw_context_intf.ml
index f6534819ea33..4adc751406ec 100644
--- a/src/proto_alpha/lib_protocol/raw_context_intf.ml
+++ b/src/proto_alpha/lib_protocol/raw_context_intf.ml
@@ -457,6 +457,9 @@ module type T = sig
 
   include VIEW
 
+  (** The type for relative context accesses instead from the root *)
+  type local_context
+
   module Tree :
     TREE
       with type t := t
@@ -580,9 +583,14 @@ module type T = sig
          and type key := key
          and type value := value
 
-    val consume_gas :
-      local_context -> Gas_limit_repr.cost -> local_context tzresult
+    type t = local_context
+
+    val consume_gas : t -> Gas_limit_repr.cost -> t tzresult
+
+    val tree : t -> tree
+
+    val update_tree : t -> tree -> t
 
-    val absolute_key : local_context -> key -> key
+    val absolute_key : t -> key -> key
   end
 end
diff --git a/src/proto_alpha/lib_protocol/storage.ml b/src/proto_alpha/lib_protocol/storage.ml
index 0902183524b8..1226237990a2 100644
--- a/src/proto_alpha/lib_protocol/storage.ml
+++ b/src/proto_alpha/lib_protocol/storage.ml
@@ -542,7 +542,35 @@ module Big_map = struct
       | Some value ->
           consume_deserialize_gas ctxt value >|? fun ctxt -> (ctxt, value_opt)
 
+
     let keys_unaccounted = I.keys_unaccounted
+
+    let with_local_context = I.with_local_context
+
+    module Local = struct
+      include I.Local
+
+      let consume_deserialize_gas local_ctxt value =
+        Raw_context.Local_context.consume_gas
+          local_ctxt
+          (Script_repr.deserialized_cost value)
+
+      let get local_ctxt contract =
+        I.Local.get local_ctxt contract >>=? fun (local_ctxt, value) ->
+        Lwt.return
+          ( consume_deserialize_gas local_ctxt value >|? fun local_ctxt ->
+            (local_ctxt, value) )
+
+      let find local_ctxt contract =
+        I.Local.find local_ctxt contract >>=? fun (local_ctxt, value_opt) ->
+        Lwt.return
+        @@
+        match value_opt with
+        | None -> ok (local_ctxt, None)
+        | Some value ->
+            consume_deserialize_gas local_ctxt value >|? fun local_ctxt ->
+            (local_ctxt, value_opt)
+    end
   end
 end
 
diff --git a/src/proto_alpha/lib_protocol/storage.mli b/src/proto_alpha/lib_protocol/storage.mli
index bc1845e62d1c..64469afb3855 100644
--- a/src/proto_alpha/lib_protocol/storage.mli
+++ b/src/proto_alpha/lib_protocol/storage.mli
@@ -244,7 +244,7 @@ module Big_map : sig
 
   module Contents : sig
     include
-      Non_iterable_indexed_carbonated_data_storage
+      Non_iterable_indexed_carbonated_data_storage_with_local_context
         with type key = Script_expr_hash.t
          and type value = Script_repr.expr
          and type t := key
diff --git a/src/proto_alpha/lib_protocol/storage_functors.ml b/src/proto_alpha/lib_protocol/storage_functors.ml
index 717816e4d286..2ee5b028543f 100644
--- a/src/proto_alpha/lib_protocol/storage_functors.ml
+++ b/src/proto_alpha/lib_protocol/storage_functors.ml
@@ -146,7 +146,17 @@ module Make_subcontext (R : REGISTER) (C : Raw_context.T) (N : NAME) :
   let with_local_context ~add_back ctxt k f =
     C.with_local_context ~add_back ctxt (to_key k) f
 
-  module Local_context = C.Local_context
+  module Local_context = struct
+    type t = C.Local_context.t
+
+    let consume_gas = C.Local_context.consume_gas
+
+    let tree = C.Local_context.tree
+
+    let update_tree = C.Local_context.update_tree
+
+    let absolute_key local key = C.Local_context.absolute_key local (to_key key)
+  end
 end
 
 module Make_single_data_storage
@@ -468,6 +478,122 @@ module Make_indexed_carbonated_data_storage_INTERNAL
   let add_or_remove s i v =
     match v with None -> remove s i | Some v -> add s i v
 
+  let with_local_context s f = C.with_local_context s [] f
+
+  module Local = struct
+    type context = Raw_context.local_context
+
+    let consume_mem_gas local key =
+      C.Local_context.consume_gas
+        local
+        (Storage_costs.read_access ~path_length:(List.length key) ~read_bytes:0)
+
+    let existing_size tree i =
+      C.Tree.find tree (len_key i) >|= function
+      | None -> ok (0, false)
+      | Some len -> decode_len_value (len_key i) len >|? fun len -> (len, true)
+
+    let consume_read_gas get local i =
+      let len_key = len_key i in
+      get (C.Local_context.tree local) len_key >>=? fun len ->
+      Lwt.return
+        ( decode_len_value len_key len >>? fun read_bytes ->
+          let cost =
+            Storage_costs.read_access
+              ~path_length:(List.length len_key)
+              ~read_bytes
+          in
+          C.Local_context.consume_gas local cost )
+
+    let consume_serialize_write_gas set local i v =
+      let bytes = to_bytes v in
+      let len = Bytes.length bytes in
+      C.Local_context.consume_gas local (Gas_limit_repr.alloc_mbytes_cost len)
+      >>?= fun local ->
+      let cost = Storage_costs.write_access ~written_bytes:len in
+      C.Local_context.consume_gas local cost >>?= fun local ->
+      set (C.Local_context.tree local) (len_key i) (encode_len_value bytes)
+      >|=? fun tree -> (C.Local_context.update_tree local tree, bytes)
+
+    let consume_remove_gas del local i =
+      C.Local_context.consume_gas
+        local
+        (Storage_costs.write_access ~written_bytes:0)
+      >>?= fun local ->
+      del (C.Local_context.tree local) (len_key i) >|=? fun tree ->
+      C.Local_context.update_tree local tree
+
+    let mem local i =
+      let key = data_key i in
+      consume_mem_gas local key >>?= fun local ->
+      let tree = C.Local_context.tree local in
+      C.Tree.mem tree key >|= fun exists -> ok (local, exists)
+
+    let get local i =
+      let get tree i = C.Tree.get tree i in
+      consume_read_gas get local i >>=? fun local ->
+      get (C.Local_context.tree local) (data_key i) >>=? fun b ->
+      let key () = C.Local_context.absolute_key local (data_key i) in
+      Lwt.return (of_bytes ~key b >|? fun v -> (local, v))
+
+    let find local i =
+      let key = data_key i in
+      consume_mem_gas local key >>?= fun local ->
+      let tree = C.Local_context.tree local in
+      C.Tree.mem tree key >>= fun exists ->
+      if exists then get local i >|=? fun (local, v) -> (local, Some v)
+      else return (local, None)
+
+    let update local i v =
+      let tree = C.Local_context.tree local in
+      existing_size tree i >>=? fun (prev_size, _) ->
+      let update tree key b = C.Tree.update tree key b in
+      consume_serialize_write_gas update local i v >>=? fun (local, bytes) ->
+      update tree (data_key i) bytes >|=? fun tree ->
+      let size_diff = Bytes.length bytes - prev_size in
+      let local = C.Local_context.update_tree local tree in
+      (local, size_diff)
+
+    let init local i v =
+      let init tree key b = C.Tree.init tree key b in
+      consume_serialize_write_gas init local i v >>=? fun (local, bytes) ->
+      let tree = C.Local_context.tree local in
+      init tree (data_key i) bytes >|=? fun tree ->
+      let local = C.Local_context.update_tree local tree in
+      let size = Bytes.length bytes in
+      (local, size)
+
+    let add local i v =
+      let add tree i b = C.Tree.add tree i b >|= fun tree -> ok tree in
+      let tree = C.Local_context.tree local in
+      existing_size tree i >>=? fun (prev_size, existed) ->
+      consume_serialize_write_gas add local i v >>=? fun (local, bytes) ->
+      add (C.Local_context.tree local) (data_key i) bytes >|=? fun tree ->
+      let local = C.Local_context.update_tree local tree in
+      let size_diff = Bytes.length bytes - prev_size in
+      (local, size_diff, existed)
+
+    let remove local i =
+      let remove tree i = C.Tree.remove tree i >|= ok in
+      let tree = C.Local_context.tree local in
+      existing_size tree i >>=? fun (prev_size, existed) ->
+      consume_remove_gas remove local i >>=? fun local ->
+      remove tree (data_key i) >|=? fun tree ->
+      let local = C.Local_context.update_tree local tree in
+      (local, prev_size, existed)
+
+    let remove_existing local i =
+      existing_size (C.Local_context.tree local) i >>=? fun (prev_size, _) ->
+      let remove_existing tree key = C.Tree.remove_existing tree key in
+      consume_remove_gas remove_existing local i >>=? fun local ->
+      remove_existing (C.Local_context.tree local) (data_key i) >|=? fun tree ->
+      let local = C.Local_context.update_tree local tree in
+      (local, prev_size)
+
+    let add_or_remove local i v =
+      match v with None -> remove local i | Some v -> add local i v
+  end
+
   (** Because big map values are not stored under some common key,
       we have no choice but to fold over all nodes with a path of length
       [I.path_length] to retrieve actual keys and then paginate.
@@ -550,7 +676,7 @@ module Make_indexed_carbonated_data_storage : functor
   (I : INDEX)
   (V : VALUE)
   ->
-  Non_iterable_indexed_carbonated_data_storage_with_values
+  Non_iterable_indexed_carbonated_data_storage_with_local_context
     with type t = C.t
      and type key = I.t
      and type value = V.t =
@@ -837,7 +963,21 @@ module Make_indexed_subcontext (C : Raw_context.T) (I : INDEX) :
       C.with_local_context ~add_back t (to_key i k) f >|=? fun (t, res) ->
       (pack t i, res)
 
-    module Local_context = C.Local_context
+    let with_local_context c k f =
+      let (t, i) = unpack c in
+      C.with_local_context t (to_key i k) f >|=? fun (t, res) -> (pack t i, res)
+
+    module Local_context = struct
+      type t = C.Local_context.t
+
+      let consume_gas local gas = C.Local_context.consume_gas local gas
+
+      let tree = C.Local_context.tree
+
+      let update_tree = C.Local_context.update_tree
+
+      let absolute_key = C.Local_context.absolute_key
+    end
   end
 
   let with_local_context ~add_back s i f =
diff --git a/src/proto_alpha/lib_protocol/storage_functors.mli b/src/proto_alpha/lib_protocol/storage_functors.mli
index e83a1f9da791..0b1dcbb8edf5 100644
--- a/src/proto_alpha/lib_protocol/storage_functors.mli
+++ b/src/proto_alpha/lib_protocol/storage_functors.mli
@@ -84,7 +84,7 @@ module Make_indexed_carbonated_data_storage
     (C : Raw_context.T)
     (I : INDEX)
     (V : VALUE) :
-  Non_iterable_indexed_carbonated_data_storage_with_values
+  Non_iterable_indexed_carbonated_data_storage_with_local_context
     with type t = C.t
      and type key = I.t
      and type value = V.t
diff --git a/src/proto_alpha/lib_protocol/storage_sigs.ml b/src/proto_alpha/lib_protocol/storage_sigs.ml
index 3921221ad5c3..2d57bb85a16d 100644
--- a/src/proto_alpha/lib_protocol/storage_sigs.ml
+++ b/src/proto_alpha/lib_protocol/storage_sigs.ml
@@ -226,9 +226,42 @@ module type Non_iterable_indexed_carbonated_data_storage_with_values = sig
     (Raw_context.t * (key * value) list) tzresult Lwt.t
 end
 
-module type Non_iterable_indexed_carbonated_data_storage_INTERNAL = sig
+module type Non_iterable_indexed_carbonated_data_storage_with_local_context = sig
   include Non_iterable_indexed_carbonated_data_storage_with_values
 
+  val with_local_context :
+    context ->
+    (Raw_context.local_context ->
+    (Raw_context.local_context * 'a) tzresult Lwt.t) ->
+    (context * 'a) tzresult Lwt.t
+
+  module Local : sig
+    type context = Raw_context.local_context
+
+    val mem : context -> key -> (context * bool) tzresult Lwt.t
+
+    val get : context -> key -> (context * value) tzresult Lwt.t
+
+    val find : context -> key -> (context * value option) tzresult Lwt.t
+
+    val update : context -> key -> value -> (context * int) tzresult Lwt.t
+
+    val init : context -> key -> value -> (context * int) tzresult Lwt.t
+
+    val add : context -> key -> value -> (context * int * bool) tzresult Lwt.t
+
+    val add_or_remove :
+      context -> key -> value option -> (context * int * bool) tzresult Lwt.t
+
+    val remove_existing : context -> key -> (context * int) tzresult Lwt.t
+
+    val remove : context -> key -> (context * int * bool) tzresult Lwt.t
+  end
+end
+
+module type Non_iterable_indexed_carbonated_data_storage_INTERNAL = sig
+  include Non_iterable_indexed_carbonated_data_storage_with_local_context
+
   val fold_keys_unaccounted :
     context ->
     order:[`Sorted | `Undefined] ->
-- 
GitLab


From 5a9b4ec8793741a93100e9027791191dd74a1623 Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <yoshihiro503@proofcafe.org>
Date: Fri, 1 Apr 2022 14:51:10 +0900
Subject: [PATCH 06/15] Protocol Storage: Fixed a bug that `len` file was not
 removed

---
 src/proto_alpha/lib_protocol/storage_functors.ml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/proto_alpha/lib_protocol/storage_functors.ml b/src/proto_alpha/lib_protocol/storage_functors.ml
index 2ee5b028543f..e9341519b6ef 100644
--- a/src/proto_alpha/lib_protocol/storage_functors.ml
+++ b/src/proto_alpha/lib_protocol/storage_functors.ml
@@ -578,6 +578,7 @@ module Make_indexed_carbonated_data_storage_INTERNAL
       let tree = C.Local_context.tree local in
       existing_size tree i >>=? fun (prev_size, existed) ->
       consume_remove_gas remove local i >>=? fun local ->
+      let tree = C.Local_context.tree local in
       remove tree (data_key i) >|=? fun tree ->
       let local = C.Local_context.update_tree local tree in
       (local, prev_size, existed)
-- 
GitLab


From 2437799b828cbc1c299113a0b2c69695a8ed83b3 Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <yoshihiro503@proofcafe.org>
Date: Fri, 1 Apr 2022 20:42:35 +0900
Subject: [PATCH 07/15] Protocol Storage : refactored `with_local_context` *
 review: https://gitlab.com/tezos/tezos/-/merge_requests/3685#note_882367554

---
 src/proto_alpha/lib_protocol/raw_context.ml | 33 ++++++++-------------
 1 file changed, 12 insertions(+), 21 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/raw_context.ml b/src/proto_alpha/lib_protocol/raw_context.ml
index c4072bb1762c..d6cb191dc61b 100644
--- a/src/proto_alpha/lib_protocol/raw_context.ml
+++ b/src/proto_alpha/lib_protocol/raw_context.ml
@@ -1593,31 +1593,22 @@ let make_local_context ctxt path =
   }
 
 let with_local_context ctxt key f =
-  let find_local_context ctxt key =
-    find_tree ctxt key
-    >|= Option.map (fun tree ->
-            {
-              tree;
-              path = key;
-              remaining_operation_gas = remaining_operation_gas ctxt;
-              unlimited_operation_gas = unlimited_operation_gas ctxt;
-            })
+  (find_tree ctxt key >|= function None -> Tree.empty ctxt | Some tree -> tree)
+  >>= fun tree ->
+  let local_ctxt =
+    {
+      tree;
+      path = key;
+      remaining_operation_gas = remaining_operation_gas ctxt;
+      unlimited_operation_gas = unlimited_operation_gas ctxt;
+    }
   in
-  (find_local_context ctxt key >>= function
-   | Some local_ctxt ->
-       f local_ctxt >>=? fun (local_ctxt, res) ->
-       update_tree ctxt key local_ctxt.tree >|=? fun ctxt ->
-       (ctxt, local_ctxt, res)
-   | None ->
-       let local_ctxt = make_local_context ctxt key in
-       f local_ctxt >>=? fun (local_ctxt, res) ->
-       add_tree ctxt key local_ctxt.tree >|= fun ctxt ->
-       ok (ctxt, local_ctxt, res))
-  >|=? fun (ctxt, local_ctxt, res) ->
+  f local_ctxt >>=? fun (local_ctxt, res) ->
+  add_tree ctxt key local_ctxt.tree >|= fun ctxt ->
   update_remaining_operation_gas ctxt local_ctxt.remaining_operation_gas
   |> fun ctxt ->
   update_unlimited_operation_gas ctxt local_ctxt.unlimited_operation_gas
-  |> fun ctxt -> (ctxt, res)
+  |> fun ctxt -> ok (ctxt, res)
 
 
 module Local_context = struct
-- 
GitLab


From d1911cf11cb8633641e4f7c17a652a6c197f0198 Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <yoshihiro503@proofcafe.org>
Date: Fri, 1 Apr 2022 22:37:33 +0900
Subject: [PATCH 08/15] Protocol Storage: Take into account the gas cost of
 local tree access.

---
 src/proto_alpha/lib_protocol/storage_functors.ml | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/proto_alpha/lib_protocol/storage_functors.ml b/src/proto_alpha/lib_protocol/storage_functors.ml
index e9341519b6ef..534a2ae1a9b7 100644
--- a/src/proto_alpha/lib_protocol/storage_functors.ml
+++ b/src/proto_alpha/lib_protocol/storage_functors.ml
@@ -478,7 +478,16 @@ module Make_indexed_carbonated_data_storage_INTERNAL
   let add_or_remove s i v =
     match v with None -> remove s i | Some v -> add s i v
 
-  let with_local_context s f = C.with_local_context s [] f
+  let with_local_context s f =
+    (* The gas cost for using C.with_local_context are:
+       - find_tree : To access the directory
+       - add_tree  : `write_access` does not charge the path length. Then no need
+                     to reduce the gas *)
+    let consume_find_tree_gas c =
+      let path_length = List.length @@ C.absolute_key c [] in
+      C.consume_gas c (Storage_costs.read_access ~path_length ~read_bytes:0)
+    in
+    consume_find_tree_gas s >>?= fun s -> C.with_local_context s [] f
 
   module Local = struct
     type context = Raw_context.local_context
-- 
GitLab


From 6c5655ca559afc543ea56d13034c8d5d06a68562 Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <yoshihiro503@proofcafe.org>
Date: Sat, 2 Apr 2022 01:59:59 +0900
Subject: [PATCH 09/15] Storage Protocol: refactoring that internal tree field
 of local_context is now encapsulated and cannot be referenced directly.

---
 src/proto_alpha/lib_protocol/raw_context.ml   |   1 +
 .../lib_protocol/raw_context_intf.ml          |  18 ++-
 .../lib_protocol/storage_functors.ml          | 107 ++++++++++--------
 3 files changed, 74 insertions(+), 52 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/raw_context.ml b/src/proto_alpha/lib_protocol/raw_context.ml
index d6cb191dc61b..77dbe6c664f5 100644
--- a/src/proto_alpha/lib_protocol/raw_context.ml
+++ b/src/proto_alpha/lib_protocol/raw_context.ml
@@ -1679,4 +1679,5 @@ module Local_context = struct
   let config local = Tree.config (tree local)
 
   let length local i = Tree.length (tree local) i
+
 end
diff --git a/src/proto_alpha/lib_protocol/raw_context_intf.ml b/src/proto_alpha/lib_protocol/raw_context_intf.ml
index 4adc751406ec..e43f62077594 100644
--- a/src/proto_alpha/lib_protocol/raw_context_intf.ml
+++ b/src/proto_alpha/lib_protocol/raw_context_intf.ml
@@ -587,10 +587,22 @@ module type T = sig
 
     val consume_gas : t -> Gas_limit_repr.cost -> t tzresult
 
-    val tree : t -> tree
+    val absolute_key : t -> key -> key
 
-    val update_tree : t -> tree -> t
+    val mem : t -> key -> bool Lwt.t
 
-    val absolute_key : t -> key -> key
+    val get : t -> key -> value tzresult Lwt.t
+
+    val find : t -> key -> value option Lwt.t
+
+    val init : t -> key -> value -> t tzresult Lwt.t
+
+    val update : t -> key -> value -> t tzresult Lwt.t
+
+    val add : t -> key -> value -> t Lwt.t
+
+    val remove : t -> key -> t Lwt.t
+
+    val remove_existing : t -> key -> t tzresult Lwt.t
   end
 end
diff --git a/src/proto_alpha/lib_protocol/storage_functors.ml b/src/proto_alpha/lib_protocol/storage_functors.ml
index 534a2ae1a9b7..a0a6a046df80 100644
--- a/src/proto_alpha/lib_protocol/storage_functors.ml
+++ b/src/proto_alpha/lib_protocol/storage_functors.ml
@@ -151,11 +151,24 @@ module Make_subcontext (R : REGISTER) (C : Raw_context.T) (N : NAME) :
 
     let consume_gas = C.Local_context.consume_gas
 
-    let tree = C.Local_context.tree
+    let absolute_key local key = C.Local_context.absolute_key local (to_key key)
 
-    let update_tree = C.Local_context.update_tree
+    let mem local key = C.Local_context.mem local (to_key key)
 
-    let absolute_key local key = C.Local_context.absolute_key local (to_key key)
+    let get local key = C.Local_context.get local (to_key key)
+
+    let find local key = C.Local_context.find local (to_key key)
+
+    let init local key = C.Local_context.init local (to_key key)
+
+    let update local key = C.Local_context.update local (to_key key)
+
+    let add local key = C.Local_context.add local (to_key key)
+
+    let remove local key = C.Local_context.remove local (to_key key)
+
+    let remove_existing local key =
+      C.Local_context.remove_existing local (to_key key)
   end
 end
 
@@ -497,14 +510,14 @@ module Make_indexed_carbonated_data_storage_INTERNAL
         local
         (Storage_costs.read_access ~path_length:(List.length key) ~read_bytes:0)
 
-    let existing_size tree i =
-      C.Tree.find tree (len_key i) >|= function
+    let existing_size local i =
+      C.Local_context.find local (len_key i) >|= function
       | None -> ok (0, false)
       | Some len -> decode_len_value (len_key i) len >|? fun len -> (len, true)
 
     let consume_read_gas get local i =
       let len_key = len_key i in
-      get (C.Local_context.tree local) len_key >>=? fun len ->
+      get local len_key >>=? fun len ->
       Lwt.return
         ( decode_len_value len_key len >>? fun read_bytes ->
           let cost =
@@ -521,83 +534,67 @@ module Make_indexed_carbonated_data_storage_INTERNAL
       >>?= fun local ->
       let cost = Storage_costs.write_access ~written_bytes:len in
       C.Local_context.consume_gas local cost >>?= fun local ->
-      set (C.Local_context.tree local) (len_key i) (encode_len_value bytes)
-      >|=? fun tree -> (C.Local_context.update_tree local tree, bytes)
+      set local (len_key i) (encode_len_value bytes) >|=? fun local ->
+      (local, bytes)
 
     let consume_remove_gas del local i =
       C.Local_context.consume_gas
         local
         (Storage_costs.write_access ~written_bytes:0)
-      >>?= fun local ->
-      del (C.Local_context.tree local) (len_key i) >|=? fun tree ->
-      C.Local_context.update_tree local tree
+      >>?= fun local -> del local (len_key i)
 
     let mem local i =
       let key = data_key i in
       consume_mem_gas local key >>?= fun local ->
-      let tree = C.Local_context.tree local in
-      C.Tree.mem tree key >|= fun exists -> ok (local, exists)
+      C.Local_context.mem local key >|= fun exists -> ok (local, exists)
 
     let get local i =
-      let get tree i = C.Tree.get tree i in
-      consume_read_gas get local i >>=? fun local ->
-      get (C.Local_context.tree local) (data_key i) >>=? fun b ->
+      consume_read_gas C.Local_context.get local i >>=? fun local ->
+      C.Local_context.get local (data_key i) >>=? fun b ->
       let key () = C.Local_context.absolute_key local (data_key i) in
       Lwt.return (of_bytes ~key b >|? fun v -> (local, v))
 
     let find local i =
       let key = data_key i in
       consume_mem_gas local key >>?= fun local ->
-      let tree = C.Local_context.tree local in
-      C.Tree.mem tree key >>= fun exists ->
+      C.Local_context.mem local key >>= fun exists ->
       if exists then get local i >|=? fun (local, v) -> (local, Some v)
       else return (local, None)
 
     let update local i v =
-      let tree = C.Local_context.tree local in
-      existing_size tree i >>=? fun (prev_size, _) ->
-      let update tree key b = C.Tree.update tree key b in
-      consume_serialize_write_gas update local i v >>=? fun (local, bytes) ->
-      update tree (data_key i) bytes >|=? fun tree ->
+      existing_size local i >>=? fun (prev_size, _) ->
+      consume_serialize_write_gas C.Local_context.update local i v
+      >>=? fun (local, bytes) ->
+      C.Local_context.update local (data_key i) bytes >|=? fun local ->
       let size_diff = Bytes.length bytes - prev_size in
-      let local = C.Local_context.update_tree local tree in
       (local, size_diff)
 
     let init local i v =
-      let init tree key b = C.Tree.init tree key b in
-      consume_serialize_write_gas init local i v >>=? fun (local, bytes) ->
-      let tree = C.Local_context.tree local in
-      init tree (data_key i) bytes >|=? fun tree ->
-      let local = C.Local_context.update_tree local tree in
+      consume_serialize_write_gas C.Local_context.init local i v
+      >>=? fun (local, bytes) ->
+      C.Local_context.init local (data_key i) bytes >|=? fun local ->
       let size = Bytes.length bytes in
       (local, size)
 
     let add local i v =
-      let add tree i b = C.Tree.add tree i b >|= fun tree -> ok tree in
-      let tree = C.Local_context.tree local in
-      existing_size tree i >>=? fun (prev_size, existed) ->
+      let add local i v = C.Local_context.add local i v >|= ok in
+      existing_size local i >>=? fun (prev_size, existed) ->
       consume_serialize_write_gas add local i v >>=? fun (local, bytes) ->
-      add (C.Local_context.tree local) (data_key i) bytes >|=? fun tree ->
-      let local = C.Local_context.update_tree local tree in
+      add local (data_key i) bytes >|=? fun local ->
       let size_diff = Bytes.length bytes - prev_size in
       (local, size_diff, existed)
 
     let remove local i =
-      let remove tree i = C.Tree.remove tree i >|= ok in
-      let tree = C.Local_context.tree local in
-      existing_size tree i >>=? fun (prev_size, existed) ->
+      let remove local i = C.Local_context.remove local i >|= ok in
+      existing_size local i >>=? fun (prev_size, existed) ->
       consume_remove_gas remove local i >>=? fun local ->
-      let tree = C.Local_context.tree local in
-      remove tree (data_key i) >|=? fun tree ->
-      let local = C.Local_context.update_tree local tree in
-      (local, prev_size, existed)
+      remove local (data_key i) >|=? fun local -> (local, prev_size, existed)
 
     let remove_existing local i =
-      existing_size (C.Local_context.tree local) i >>=? fun (prev_size, _) ->
-      let remove_existing tree key = C.Tree.remove_existing tree key in
-      consume_remove_gas remove_existing local i >>=? fun local ->
-      remove_existing (C.Local_context.tree local) (data_key i) >|=? fun tree ->
-      let local = C.Local_context.update_tree local tree in
+      existing_size local i >>=? fun (prev_size, _) ->
+      consume_remove_gas C.Local_context.remove_existing local i
+      >>=? fun local ->
+      C.Local_context.remove_existing local (data_key i) >|=? fun local ->
       (local, prev_size)
 
     let add_or_remove local i v =
@@ -982,11 +979,23 @@ module Make_indexed_subcontext (C : Raw_context.T) (I : INDEX) :
 
       let consume_gas local gas = C.Local_context.consume_gas local gas
 
-      let tree = C.Local_context.tree
+      let absolute_key = C.Local_context.absolute_key
 
-      let update_tree = C.Local_context.update_tree
+      let mem local k = C.Local_context.mem local k
 
-      let absolute_key = C.Local_context.absolute_key
+      let get local k = C.Local_context.get local k
+
+      let find local k = C.Local_context.find local k
+
+      let init local k = C.Local_context.init local k
+
+      let update local k = C.Local_context.update local k
+
+      let add local k = C.Local_context.add local k
+
+      let remove local k = C.Local_context.remove local k
+
+      let remove_existing local k = C.Local_context.remove_existing local k
     end
   end
 
-- 
GitLab


From d3330d89764f14631e8061beedbd190c68949300 Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <yoshihiro503@proofcafe.org>
Date: Tue, 26 Apr 2022 12:19:31 +0900
Subject: [PATCH 10/15] Protocol Storage : use VIEW interface for local context
 module

---
 src/proto_alpha/lib_protocol/raw_context.ml   |  1 -
 .../lib_protocol/raw_context_intf.ml          | 21 +------
 .../lib_protocol/storage_functors.ml          | 56 +++----------------
 3 files changed, 11 insertions(+), 67 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/raw_context.ml b/src/proto_alpha/lib_protocol/raw_context.ml
index 77dbe6c664f5..6cfb27054f89 100644
--- a/src/proto_alpha/lib_protocol/raw_context.ml
+++ b/src/proto_alpha/lib_protocol/raw_context.ml
@@ -1610,7 +1610,6 @@ let with_local_context ctxt key f =
   update_unlimited_operation_gas ctxt local_ctxt.unlimited_operation_gas
   |> fun ctxt -> ok (ctxt, res)
 
-
 module Local_context = struct
   type t = local_context
 
diff --git a/src/proto_alpha/lib_protocol/raw_context_intf.ml b/src/proto_alpha/lib_protocol/raw_context_intf.ml
index e43f62077594..6222e671963a 100644
--- a/src/proto_alpha/lib_protocol/raw_context_intf.ml
+++ b/src/proto_alpha/lib_protocol/raw_context_intf.ml
@@ -585,24 +585,9 @@ module type T = sig
 
     type t = local_context
 
-    val consume_gas : t -> Gas_limit_repr.cost -> t tzresult
+    val consume_gas :
+      local_context -> Gas_limit_repr.cost -> local_context tzresult
 
-    val absolute_key : t -> key -> key
-
-    val mem : t -> key -> bool Lwt.t
-
-    val get : t -> key -> value tzresult Lwt.t
-
-    val find : t -> key -> value option Lwt.t
-
-    val init : t -> key -> value -> t tzresult Lwt.t
-
-    val update : t -> key -> value -> t tzresult Lwt.t
-
-    val add : t -> key -> value -> t Lwt.t
-
-    val remove : t -> key -> t Lwt.t
-
-    val remove_existing : t -> key -> t tzresult Lwt.t
+    val absolute_key : local_context -> key -> key
   end
 end
diff --git a/src/proto_alpha/lib_protocol/storage_functors.ml b/src/proto_alpha/lib_protocol/storage_functors.ml
index a0a6a046df80..c05ef73febc2 100644
--- a/src/proto_alpha/lib_protocol/storage_functors.ml
+++ b/src/proto_alpha/lib_protocol/storage_functors.ml
@@ -146,30 +146,7 @@ module Make_subcontext (R : REGISTER) (C : Raw_context.T) (N : NAME) :
   let with_local_context ~add_back ctxt k f =
     C.with_local_context ~add_back ctxt (to_key k) f
 
-  module Local_context = struct
-    type t = C.Local_context.t
-
-    let consume_gas = C.Local_context.consume_gas
-
-    let absolute_key local key = C.Local_context.absolute_key local (to_key key)
-
-    let mem local key = C.Local_context.mem local (to_key key)
-
-    let get local key = C.Local_context.get local (to_key key)
-
-    let find local key = C.Local_context.find local (to_key key)
-
-    let init local key = C.Local_context.init local (to_key key)
-
-    let update local key = C.Local_context.update local (to_key key)
-
-    let add local key = C.Local_context.add local (to_key key)
-
-    let remove local key = C.Local_context.remove local (to_key key)
-
-    let remove_existing local key =
-      C.Local_context.remove_existing local (to_key key)
-  end
+  module Local_context = C.Local_context
 end
 
 module Make_single_data_storage
@@ -500,7 +477,12 @@ module Make_indexed_carbonated_data_storage_INTERNAL
       let path_length = List.length @@ C.absolute_key c [] in
       C.consume_gas c (Storage_costs.read_access ~path_length ~read_bytes:0)
     in
-    consume_find_tree_gas s >>?= fun s -> C.with_local_context s [] f
+    let consume_add_tree_gas c =
+      C.consume_gas c (Storage_costs.write_access ~written_bytes:0)
+    in
+    consume_find_tree_gas s >>?= fun s ->
+    C.with_local_context s [] f >>=? fun (s, x) ->
+    consume_add_tree_gas s >>?= fun s -> return (s, x)
 
   module Local = struct
     type context = Raw_context.local_context
@@ -974,29 +956,7 @@ module Make_indexed_subcontext (C : Raw_context.T) (I : INDEX) :
       let (t, i) = unpack c in
       C.with_local_context t (to_key i k) f >|=? fun (t, res) -> (pack t i, res)
 
-    module Local_context = struct
-      type t = C.Local_context.t
-
-      let consume_gas local gas = C.Local_context.consume_gas local gas
-
-      let absolute_key = C.Local_context.absolute_key
-
-      let mem local k = C.Local_context.mem local k
-
-      let get local k = C.Local_context.get local k
-
-      let find local k = C.Local_context.find local k
-
-      let init local k = C.Local_context.init local k
-
-      let update local k = C.Local_context.update local k
-
-      let add local k = C.Local_context.add local k
-
-      let remove local k = C.Local_context.remove local k
-
-      let remove_existing local k = C.Local_context.remove_existing local k
-    end
+    module Local_context = C.Local_context
   end
 
   let with_local_context ~add_back s i f =
-- 
GitLab


From 3583a0ab3db0302c59c2816029177ab960f0bb1c Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <yoshihiro503@proofcafe.org>
Date: Tue, 26 Apr 2022 16:18:23 +0900
Subject: [PATCH 11/15] Protocol Storage : introduce length function for local
 context

---
 src/proto_alpha/lib_protocol/raw_context.ml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/proto_alpha/lib_protocol/raw_context.ml b/src/proto_alpha/lib_protocol/raw_context.ml
index 6cfb27054f89..457136091a7e 100644
--- a/src/proto_alpha/lib_protocol/raw_context.ml
+++ b/src/proto_alpha/lib_protocol/raw_context.ml
@@ -1678,5 +1678,4 @@ module Local_context = struct
   let config local = Tree.config (tree local)
 
   let length local i = Tree.length (tree local) i
-
 end
-- 
GitLab


From fa00809c7dc0b8ec0abab43455d185f2b7320264 Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <yoshihiro503@proofcafe.org>
Date: Tue, 17 May 2022 19:37:38 +0900
Subject: [PATCH 12/15] Protocol Storage : To the consumption costs remain the
 same, the additional gas cost when for_write is true is set to zero.

---
 .../lib_protocol/lazy_storage_diff.ml            |  5 ++++-
 src/proto_alpha/lib_protocol/storage_functors.ml | 16 ++++++++--------
 src/proto_alpha/lib_protocol/storage_sigs.ml     |  1 +
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/lazy_storage_diff.ml b/src/proto_alpha/lib_protocol/lazy_storage_diff.ml
index 53cdddde6397..5b3613dbfcf0 100644
--- a/src/proto_alpha/lib_protocol/lazy_storage_diff.ml
+++ b/src/proto_alpha/lib_protocol/lazy_storage_diff.ml
@@ -142,7 +142,10 @@ module Big_map = struct
         (local_ctxt, Z.of_int size_diff)
 
   let apply_updates ctxt ~id updates =
-    Storage.Big_map.Contents.with_local_context (ctxt, id) (fun local_ctxt ->
+    Storage.Big_map.Contents.with_local_context
+      ~for_write:true
+      (ctxt, id)
+      (fun local_ctxt ->
         List.fold_left_es
           (fun (local_ctxt, size) update ->
             apply_update local_ctxt update >|=? fun (local_ctxt, added_size) ->
diff --git a/src/proto_alpha/lib_protocol/storage_functors.ml b/src/proto_alpha/lib_protocol/storage_functors.ml
index c05ef73febc2..852502a630fa 100644
--- a/src/proto_alpha/lib_protocol/storage_functors.ml
+++ b/src/proto_alpha/lib_protocol/storage_functors.ml
@@ -468,11 +468,14 @@ module Make_indexed_carbonated_data_storage_INTERNAL
   let add_or_remove s i v =
     match v with None -> remove s i | Some v -> add s i v
 
-  let with_local_context s f =
+  let with_local_context ~for_write s f =
     (* The gas cost for using C.with_local_context are:
        - find_tree : To access the directory
        - add_tree  : `write_access` does not charge the path length. Then no need
                      to reduce the gas *)
+    (* Remark: To the consumption costs remain the same, the additional gas cost
+       when for_write is true is set to zero.
+    *)
     let consume_find_tree_gas c =
       let path_length = List.length @@ C.absolute_key c [] in
       C.consume_gas c (Storage_costs.read_access ~path_length ~read_bytes:0)
@@ -480,9 +483,10 @@ module Make_indexed_carbonated_data_storage_INTERNAL
     let consume_add_tree_gas c =
       C.consume_gas c (Storage_costs.write_access ~written_bytes:0)
     in
-    consume_find_tree_gas s >>?= fun s ->
-    C.with_local_context s [] f >>=? fun (s, x) ->
-    consume_add_tree_gas s >>?= fun s -> return (s, x)
+    (if not for_write then consume_find_tree_gas s else ok s) >>?= fun s ->
+    C.with_local_context ~for_write s [] f >>=? fun (s, x) ->
+    (if not for_write then consume_add_tree_gas s else ok s) >>?= fun s ->
+    return (s, x)
 
   module Local = struct
     type context = Raw_context.local_context
@@ -952,10 +956,6 @@ module Make_indexed_subcontext (C : Raw_context.T) (I : INDEX) :
       C.with_local_context ~add_back t (to_key i k) f >|=? fun (t, res) ->
       (pack t i, res)
 
-    let with_local_context c k f =
-      let (t, i) = unpack c in
-      C.with_local_context t (to_key i k) f >|=? fun (t, res) -> (pack t i, res)
-
     module Local_context = C.Local_context
   end
 
diff --git a/src/proto_alpha/lib_protocol/storage_sigs.ml b/src/proto_alpha/lib_protocol/storage_sigs.ml
index 2d57bb85a16d..ffebed95c5ec 100644
--- a/src/proto_alpha/lib_protocol/storage_sigs.ml
+++ b/src/proto_alpha/lib_protocol/storage_sigs.ml
@@ -230,6 +230,7 @@ module type Non_iterable_indexed_carbonated_data_storage_with_local_context = si
   include Non_iterable_indexed_carbonated_data_storage_with_values
 
   val with_local_context :
+    for_write:bool ->
     context ->
     (Raw_context.local_context ->
     (Raw_context.local_context * 'a) tzresult Lwt.t) ->
-- 
GitLab


From cddbb2d4742367ed1d8ffbeb4a329711eb88fbe5 Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <y.imai@proof-ninja.co.jp>
Date: Fri, 29 Jul 2022 10:26:26 +0900
Subject: [PATCH 13/15] Protocol Storage: rebased

---
 .../lib_protocol/lazy_storage_diff.ml         |  2 +-
 src/proto_alpha/lib_protocol/raw_context.ml   | 26 -------------------
 .../lib_protocol/raw_context_intf.ml          |  3 ---
 src/proto_alpha/lib_protocol/storage.ml       |  1 -
 .../lib_protocol/storage_functors.ml          | 10 +++----
 src/proto_alpha/lib_protocol/storage_sigs.ml  |  2 +-
 6 files changed, 7 insertions(+), 37 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/lazy_storage_diff.ml b/src/proto_alpha/lib_protocol/lazy_storage_diff.ml
index 5b3613dbfcf0..7ad9016e67eb 100644
--- a/src/proto_alpha/lib_protocol/lazy_storage_diff.ml
+++ b/src/proto_alpha/lib_protocol/lazy_storage_diff.ml
@@ -143,7 +143,7 @@ module Big_map = struct
 
   let apply_updates ctxt ~id updates =
     Storage.Big_map.Contents.with_local_context
-      ~for_write:true
+      ~add_back:true
       (ctxt, id)
       (fun local_ctxt ->
         List.fold_left_es
diff --git a/src/proto_alpha/lib_protocol/raw_context.ml b/src/proto_alpha/lib_protocol/raw_context.ml
index 457136091a7e..742e3e4faf04 100644
--- a/src/proto_alpha/lib_protocol/raw_context.ml
+++ b/src/proto_alpha/lib_protocol/raw_context.ml
@@ -1584,32 +1584,6 @@ let with_local_context ~add_back ctxt key f =
   update_unlimited_operation_gas ctxt local_ctxt.unlimited_operation_gas
   |> fun ctxt -> ok (ctxt, res)
 
-let make_local_context ctxt path =
-  {
-    tree = Tree.empty ctxt;
-    path;
-    remaining_operation_gas = remaining_operation_gas ctxt;
-    unlimited_operation_gas = unlimited_operation_gas ctxt;
-  }
-
-let with_local_context ctxt key f =
-  (find_tree ctxt key >|= function None -> Tree.empty ctxt | Some tree -> tree)
-  >>= fun tree ->
-  let local_ctxt =
-    {
-      tree;
-      path = key;
-      remaining_operation_gas = remaining_operation_gas ctxt;
-      unlimited_operation_gas = unlimited_operation_gas ctxt;
-    }
-  in
-  f local_ctxt >>=? fun (local_ctxt, res) ->
-  add_tree ctxt key local_ctxt.tree >|= fun ctxt ->
-  update_remaining_operation_gas ctxt local_ctxt.remaining_operation_gas
-  |> fun ctxt ->
-  update_unlimited_operation_gas ctxt local_ctxt.unlimited_operation_gas
-  |> fun ctxt -> ok (ctxt, res)
-
 module Local_context = struct
   type t = local_context
 
diff --git a/src/proto_alpha/lib_protocol/raw_context_intf.ml b/src/proto_alpha/lib_protocol/raw_context_intf.ml
index 6222e671963a..6c131a7717d3 100644
--- a/src/proto_alpha/lib_protocol/raw_context_intf.ml
+++ b/src/proto_alpha/lib_protocol/raw_context_intf.ml
@@ -457,9 +457,6 @@ module type T = sig
 
   include VIEW
 
-  (** The type for relative context accesses instead from the root *)
-  type local_context
-
   module Tree :
     TREE
       with type t := t
diff --git a/src/proto_alpha/lib_protocol/storage.ml b/src/proto_alpha/lib_protocol/storage.ml
index 1226237990a2..ed912f88612d 100644
--- a/src/proto_alpha/lib_protocol/storage.ml
+++ b/src/proto_alpha/lib_protocol/storage.ml
@@ -542,7 +542,6 @@ module Big_map = struct
       | Some value ->
           consume_deserialize_gas ctxt value >|? fun ctxt -> (ctxt, value_opt)
 
-
     let keys_unaccounted = I.keys_unaccounted
 
     let with_local_context = I.with_local_context
diff --git a/src/proto_alpha/lib_protocol/storage_functors.ml b/src/proto_alpha/lib_protocol/storage_functors.ml
index 852502a630fa..912e2e8a4156 100644
--- a/src/proto_alpha/lib_protocol/storage_functors.ml
+++ b/src/proto_alpha/lib_protocol/storage_functors.ml
@@ -468,13 +468,13 @@ module Make_indexed_carbonated_data_storage_INTERNAL
   let add_or_remove s i v =
     match v with None -> remove s i | Some v -> add s i v
 
-  let with_local_context ~for_write s f =
+  let with_local_context ~add_back s f =
     (* The gas cost for using C.with_local_context are:
        - find_tree : To access the directory
        - add_tree  : `write_access` does not charge the path length. Then no need
                      to reduce the gas *)
     (* Remark: To the consumption costs remain the same, the additional gas cost
-       when for_write is true is set to zero.
+       when add_back is true is set to zero.
     *)
     let consume_find_tree_gas c =
       let path_length = List.length @@ C.absolute_key c [] in
@@ -483,9 +483,9 @@ module Make_indexed_carbonated_data_storage_INTERNAL
     let consume_add_tree_gas c =
       C.consume_gas c (Storage_costs.write_access ~written_bytes:0)
     in
-    (if not for_write then consume_find_tree_gas s else ok s) >>?= fun s ->
-    C.with_local_context ~for_write s [] f >>=? fun (s, x) ->
-    (if not for_write then consume_add_tree_gas s else ok s) >>?= fun s ->
+    (if not add_back then consume_find_tree_gas s else ok s) >>?= fun s ->
+    C.with_local_context ~add_back s [] f >>=? fun (s, x) ->
+    (if not add_back then consume_add_tree_gas s else ok s) >>?= fun s ->
     return (s, x)
 
   module Local = struct
diff --git a/src/proto_alpha/lib_protocol/storage_sigs.ml b/src/proto_alpha/lib_protocol/storage_sigs.ml
index ffebed95c5ec..b3bdcba53ac6 100644
--- a/src/proto_alpha/lib_protocol/storage_sigs.ml
+++ b/src/proto_alpha/lib_protocol/storage_sigs.ml
@@ -230,7 +230,7 @@ module type Non_iterable_indexed_carbonated_data_storage_with_local_context = si
   include Non_iterable_indexed_carbonated_data_storage_with_values
 
   val with_local_context :
-    for_write:bool ->
+    add_back:bool ->
     context ->
     (Raw_context.local_context ->
     (Raw_context.local_context * 'a) tzresult Lwt.t) ->
-- 
GitLab


From 30d41737efd4697f4d2aa776f2d02f78b035d953 Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <yoshihiro503@proofcafe.org>
Date: Tue, 26 Apr 2022 12:19:31 +0900
Subject: [PATCH 14/15] Protocol Storage : use VIEW interface for local context
 module

---
 src/proto_alpha/lib_protocol/raw_context.ml | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/proto_alpha/lib_protocol/raw_context.ml b/src/proto_alpha/lib_protocol/raw_context.ml
index 742e3e4faf04..79caf278d4d3 100644
--- a/src/proto_alpha/lib_protocol/raw_context.ml
+++ b/src/proto_alpha/lib_protocol/raw_context.ml
@@ -1584,7 +1584,19 @@ let with_local_context ~add_back ctxt key f =
   update_unlimited_operation_gas ctxt local_ctxt.unlimited_operation_gas
   |> fun ctxt -> ok (ctxt, res)
 
-module Local_context = struct
+module Local_context : sig
+  include
+    Raw_context_intf.VIEW
+      with type t = local_context
+       and type key := key
+       and type value := value
+       and type tree := tree
+
+  val consume_gas :
+    local_context -> Gas_limit_repr.cost -> local_context tzresult
+
+  val absolute_key : local_context -> key -> key
+end = struct
   type t = local_context
 
   let consume_gas local cost =
-- 
GitLab


From adfefabb0c005915b1ac3b98eea50469d3ebd37e Mon Sep 17 00:00:00 2001
From: Yoshihiro Imai <yoshihiro503@proofcafe.org>
Date: Wed, 9 Feb 2022 11:56:52 +0900
Subject: [PATCH 15/15] Protocol Storage: improve expanding global constants

---
 .../lib_protocol/global_constants_storage.ml  | 58 ++++++++++++-------
 src/proto_alpha/lib_protocol/storage.ml       |  2 +-
 src/proto_alpha/lib_protocol/storage.mli      |  2 +-
 3 files changed, 38 insertions(+), 24 deletions(-)

diff --git a/src/proto_alpha/lib_protocol/global_constants_storage.ml b/src/proto_alpha/lib_protocol/global_constants_storage.ml
index be7b9dc62f46..885ad9f7358c 100644
--- a/src/proto_alpha/lib_protocol/global_constants_storage.ml
+++ b/src/proto_alpha/lib_protocol/global_constants_storage.ml
@@ -143,6 +143,13 @@ let get context hash =
   | None -> fail Nonexistent_global
   | Some value -> return (context, value)
 
+let get_local local hash =
+  Storage.Global_constants.Map.Local.find local hash
+  >>=? fun (context, value) ->
+  match value with
+  | None -> fail Nonexistent_global
+  | Some value -> return (context, value)
+
 let expr_to_address_in_context context expr =
   let lexpr = Script_repr.lazy_expr expr in
   Raw_context.consume_gas context @@ Script_repr.force_bytes_cost lexpr
@@ -160,27 +167,28 @@ let node_too_large node =
     nodes > Constants_repr.max_micheline_node_count
     || string_bytes + z_bytes > Constants_repr.max_micheline_bytes_limit)
 
-let expand_node context node =
+let expand_node_local local node =
   (* We charge for traversing the top-level node at the beginning.
      Inside the loop, we charge for traversing each new constant
      that gets expanded. *)
-  Raw_context.consume_gas
-    context
+  Raw_context.Local_context.consume_gas
+    local
     (Gas_costs.expand_no_constants_branch_cost node)
-  >>?= fun context ->
+  >>?= fun local ->
   bottom_up_fold_cps
     (* We carry a Boolean representing whether we
        had to do any expansions or not. *)
-    (context, Expr_hash_map.empty, false)
+    (local, Expr_hash_map.empty, false)
     node
-    (fun (context, _, did_expansion) node ->
-      return (context, node, did_expansion))
-    (fun (context, map, did_expansion) node k ->
+    (fun (local, _, did_expansion) node -> return (local, node, did_expansion))
+    (fun (local, map, did_expansion) node k ->
       match node with
       | Prim (_, H_constant, args, annot) -> (
           (* Charge for validating the b58check hash. *)
-          Raw_context.consume_gas context Gas_costs.expand_constants_branch_cost
-          >>?= fun context ->
+          Raw_context.Local_context.consume_gas
+            local
+            Gas_costs.expand_constants_branch_cost
+          >>?= fun local ->
           match (args, annot) with
           (* A constant Prim should always have a single String argument,
               being a properly formatted hash. *)
@@ -191,30 +199,36 @@ let expand_node context node =
                   match Expr_hash_map.find hash map with
                   | Some node ->
                       (* Charge traversing the newly retrieved node *)
-                      Raw_context.consume_gas
-                        context
+                      Raw_context.Local_context.consume_gas
+                        local
                         (Gas_costs.expand_no_constants_branch_cost node)
-                      >>?= fun context -> k (context, map, true) node
+                      >>?= fun local -> k (local, map, true) node
                   | None ->
-                      get context hash >>=? fun (context, expr) ->
+                      get_local local hash >>=? fun (local, expr) ->
                       (* Charge traversing the newly retrieved node *)
                       let node = root expr in
-                      Raw_context.consume_gas
-                        context
+                      Raw_context.Local_context.consume_gas
+                        local
                         (Gas_costs.expand_no_constants_branch_cost node)
-                      >>?= fun context ->
-                      k (context, Expr_hash_map.add hash node map, true) node))
+                      >>?= fun local ->
+                      k (local, Expr_hash_map.add hash node map, true) node))
           | _ -> fail Badly_formed_constant_expression)
       | Int _ | String _ | Bytes _ | Prim _ | Seq _ ->
-          k (context, map, did_expansion) node)
-  >>=? fun (context, node, did_expansion) ->
+          k (local, map, did_expansion) node)
+  >>=? fun (local, node, did_expansion) ->
   if did_expansion then
     (* Gas charged during expansion is at least proportional to the size of the
        resulting node so the execution time of [node_too_large] is already
        covered. *)
     if node_too_large node then fail Expression_too_large
-    else return (context, node)
-  else return (context, node)
+    else return (local, node)
+  else return (local, node)
+
+let expand_node context node =
+  Storage.Global_constants.Map.with_local_context
+    ~add_back:false
+    context
+    (fun local -> expand_node_local local node)
 
 let expand context expr =
   expand_node context (root expr) >|=? fun (context, node) ->
diff --git a/src/proto_alpha/lib_protocol/storage.ml b/src/proto_alpha/lib_protocol/storage.ml
index ed912f88612d..201296561bc3 100644
--- a/src/proto_alpha/lib_protocol/storage.ml
+++ b/src/proto_alpha/lib_protocol/storage.ml
@@ -391,7 +391,7 @@ end
 
 module Global_constants = struct
   module Map :
-    Non_iterable_indexed_carbonated_data_storage
+    Non_iterable_indexed_carbonated_data_storage_with_local_context
       with type t := Raw_context.t
        and type key = Script_expr_hash.t
        and type value = Script_repr.expr =
diff --git a/src/proto_alpha/lib_protocol/storage.mli b/src/proto_alpha/lib_protocol/storage.mli
index 64469afb3855..2020083f1adc 100644
--- a/src/proto_alpha/lib_protocol/storage.mli
+++ b/src/proto_alpha/lib_protocol/storage.mli
@@ -585,7 +585,7 @@ end
     [Michelson_v1_primitives.H_constant]. *)
 module Global_constants : sig
   module Map :
-    Non_iterable_indexed_carbonated_data_storage
+    Non_iterable_indexed_carbonated_data_storage_with_local_context
       with type t := Raw_context.t
        and type key = Script_expr_hash.t
        and type value = Script_repr.expr
-- 
GitLab