From 513da7af297b1aed301d4b9ae98a9c93a431252b Mon Sep 17 00:00:00 2001
From: Thomas Gazagnaire <thomas@gazagnaire.org>
Date: Thu, 24 Mar 2022 13:33:21 +0100
Subject: [PATCH 1/2] lib_context: expose Tree.kinded_key

---
 src/lib_context/helpers/context.ml  |  5 +++++
 src/lib_context/helpers/context.mli |  2 ++
 src/lib_context/memory/context.ml   | 10 ++++++++++
 src/lib_context/sigs/context.ml     |  2 ++
 4 files changed, 19 insertions(+)

diff --git a/src/lib_context/helpers/context.ml b/src/lib_context/helpers/context.ml
index 73e1b3cbc455..634be779dc68 100644
--- a/src/lib_context/helpers/context.ml
+++ b/src/lib_context/helpers/context.ml
@@ -199,6 +199,11 @@ module Make_tree (Conf : Conf) (Store : DB) = struct
     in
     Store.Tree.shallow repo kinded_hash
 
+  let kinded_key t =
+    match Store.Tree.key t with
+    | (None | Some (`Node _)) as r -> r
+    | Some (`Contents (v, ())) -> Some (`Value v)
+
   let is_shallow tree =
     match Store.Tree.inspect tree with
     | `Node `Key -> true
diff --git a/src/lib_context/helpers/context.mli b/src/lib_context/helpers/context.mli
index eecf97408441..3b8c99151fe9 100644
--- a/src/lib_context/helpers/context.mli
+++ b/src/lib_context/helpers/context.mli
@@ -58,6 +58,8 @@ module Make_tree (Conf : Conf) (DB : DB) : sig
 
   val shallow : DB.repo -> kinded_key -> DB.tree
 
+  val kinded_key : DB.tree -> kinded_key option
+
   val is_shallow : DB.tree -> bool
 
   (** Exception raised by [find_tree] and [add_tree] when applied to shallow
diff --git a/src/lib_context/memory/context.ml b/src/lib_context/memory/context.ml
index f60b8176110f..7eea8599552d 100644
--- a/src/lib_context/memory/context.ml
+++ b/src/lib_context/memory/context.ml
@@ -40,12 +40,22 @@ module Make (Encoding : module type of Tezos_context_encoding.Context) = struct
       match t with
       | `Node hash -> `Node (Hash.of_context_hash hash)
       | `Value hash -> `Contents (Hash.of_context_hash hash, ())
+
+    let of_irmin_key t : kinded_key =
+      match t with
+      | `Node hash -> `Node (Hash.to_context_hash hash)
+      | `Contents (hash, ()) -> `Value (Hash.to_context_hash hash)
   end
 
   module Tree = struct
     include Tezos_context_helpers.Context.Make_tree (Conf) (Store)
 
     let shallow repo key = Store.Tree.shallow repo (Kinded_key.to_irmin_key key)
+
+    let kinded_key tree =
+      match Store.Tree.key tree with
+      | None -> None
+      | Some h -> Some (Kinded_key.of_irmin_key h)
   end
 
   include Tree
diff --git a/src/lib_context/sigs/context.ml b/src/lib_context/sigs/context.ml
index 341c915a8e9d..7a636de2698d 100644
--- a/src/lib_context/sigs/context.ml
+++ b/src/lib_context/sigs/context.ml
@@ -453,6 +453,8 @@ module type S = sig
     val shallow : repo -> kinded_key -> tree
 
     val is_shallow : tree -> bool
+
+    val kinded_key : tree -> kinded_key option
   end
 
   (** [produce r h f] runs [f] on top of a real store [r], producing a proof and
-- 
GitLab


From f3b1b9ba33ae411bbe419d3f476b62c3cc41ec25 Mon Sep 17 00:00:00 2001
From: Thomas Gazagnaire <thomas@gazagnaire.org>
Date: Thu, 24 Mar 2022 15:00:27 +0100
Subject: [PATCH 2/2] lib_context: expose Context.flush

---
 src/lib_context/context.mli | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/lib_context/context.mli b/src/lib_context/context.mli
index 169cc6812be0..9109d21d056a 100644
--- a/src/lib_context/context.mli
+++ b/src/lib_context/context.mli
@@ -72,6 +72,8 @@ module Make (Encoding : module type of Tezos_context_encoding.Context) : sig
     Does not fail when the context is not in read-only mode. *)
   val sync : index -> unit Lwt.t
 
+  val flush : t -> t Lwt.t
+
   val compute_testchain_chain_id : Block_hash.t -> Chain_id.t
 
   val compute_testchain_genesis : Block_hash.t -> Block_hash.t
-- 
GitLab