From 65e22517f6c4e8dba32e450db980f2325874d273 Mon Sep 17 00:00:00 2001
From: Franklin Delehelle <franklin.delehelle@odena.eu>
Date: Wed, 2 Oct 2024 22:58:05 +0200
Subject: [PATCH] refactor: more ergonomic `WorkplanItem`

---
 mp2-v1/tests/common/cases/query.rs |  2 +-
 mp2-v1/tests/common/celltree.rs    |  2 +-
 mp2-v1/tests/common/index_tree.rs  |  2 +-
 mp2-v1/tests/common/rowtree.rs     |  2 +-
 ryhope/src/lib.rs                  |  2 +-
 ryhope/src/storage/updatetree.rs   | 36 ++++++++++++++++++++----------
 6 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/mp2-v1/tests/common/cases/query.rs b/mp2-v1/tests/common/cases/query.rs
index f0e31e763..eed1739a0 100644
--- a/mp2-v1/tests/common/cases/query.rs
+++ b/mp2-v1/tests/common/cases/query.rs
@@ -565,7 +565,7 @@ where
         // closure performing all the operations necessary beofre jumping to the next iteration
         let mut end_iteration = |proven_nodes: &mut HashSet<K>| -> Result<()> {
             proven_nodes.insert(k.clone());
-            workplan.done(&wk)?;
+            workplan.done(wk.k())?;
             Ok(())
         };
         // since epoch starts at genesis now, we can directly give the value of the block
diff --git a/mp2-v1/tests/common/celltree.rs b/mp2-v1/tests/common/celltree.rs
index b40dabf6b..64ca88c7a 100644
--- a/mp2-v1/tests/common/celltree.rs
+++ b/mp2-v1/tests/common/celltree.rs
@@ -178,7 +178,7 @@ impl TestContext {
                         .unwrap()
                 )
             );
-            workplan.done(&wk).unwrap();
+            workplan.done(wk.k()).unwrap();
         }
         let root = tree.root().await.unwrap();
         let root_data = tree.root_data().await.unwrap();
diff --git a/mp2-v1/tests/common/index_tree.rs b/mp2-v1/tests/common/index_tree.rs
index 6a469ca67..0d3c7ce8d 100644
--- a/mp2-v1/tests/common/index_tree.rs
+++ b/mp2-v1/tests/common/index_tree.rs
@@ -200,7 +200,7 @@ impl TestContext {
                 .store_proof(ProofKey::Index(proof_key), proof)
                 .expect("unable to store index tree proof");
 
-            workplan.done(&wk).unwrap();
+            workplan.done(wk.k()).unwrap();
         }
         let root = t.root().await.unwrap();
         let root_proof_key = IndexProofIdentifier {
diff --git a/mp2-v1/tests/common/rowtree.rs b/mp2-v1/tests/common/rowtree.rs
index 1b6a70946..7d9a096a8 100644
--- a/mp2-v1/tests/common/rowtree.rs
+++ b/mp2-v1/tests/common/rowtree.rs
@@ -238,7 +238,7 @@ impl TestContext {
                 new_proof_key,
                 hex::encode(extract_hash_from_proof(&proof).unwrap().to_bytes())
             );
-            workplan.done(&wk).unwrap();
+            workplan.done(wk.k()).unwrap();
         }
         let root = t.root().await.unwrap();
         let row = table.row.fetch(&root).await;
diff --git a/ryhope/src/lib.rs b/ryhope/src/lib.rs
index a091bdd4c..8ed71d628 100644
--- a/ryhope/src/lib.rs
+++ b/ryhope/src/lib.rs
@@ -166,7 +166,7 @@ where
 
             let mut payload = self.storage.data().fetch(item.k()).await;
             payload.aggregate(child_data.into_iter());
-            plan.done(&item)?;
+            plan.done(&item.k())?;
             self.storage
                 .data_mut()
                 .store(item.k().to_owned(), payload)
diff --git a/ryhope/src/storage/updatetree.rs b/ryhope/src/storage/updatetree.rs
index 99502f3f7..509ee12b0 100644
--- a/ryhope/src/storage/updatetree.rs
+++ b/ryhope/src/storage/updatetree.rs
@@ -383,6 +383,22 @@ impl<K: Clone + Hash + Eq + Debug> WorkplanItem<K> {
             WorkplanItem::Subtree { k, .. } | WorkplanItem::Node { k, .. } => k,
         }
     }
+
+    pub fn as_subtree(&self) -> &UpdateTree<K> {
+        if let WorkplanItem::Subtree { subtree, .. } = self {
+            subtree
+        } else {
+            unreachable!()
+        }
+    }
+
+    pub fn as_node(&self) -> bool {
+        if let WorkplanItem::Node { is_path_end, .. } = self {
+            *is_path_end
+        } else {
+            unreachable!()
+        }
+    }
 }
 
 /// An update plan to recompute all the hashes of the touched nodes stored in a
@@ -419,15 +435,11 @@ impl<T: Debug + Clone + Hash + Eq> UpdatePlan<T> {
 
     /// Mark the given item as having been completed. Its dependent will not be
     /// generated by the iterator until the item has been marked as completed.
-    pub fn done(&mut self, item: &WorkplanItem<T>) -> Result<()> {
-        let i = *self
-            .t
-            .idx
-            .get(&item.k())
-            .ok_or_else(|| anyhow!("unknwown key"))?;
+    pub fn done(&mut self, k: &T) -> Result<()> {
+        let i = *self.t.idx.get(k).ok_or_else(|| anyhow!("unknwown key"))?;
 
         // May happen when restarting a plan
-        self.anchors.retain(|k| k != item.k());
+        self.anchors.retain(|_k| _k != k);
 
         // Root node is hardcoded to 0
         if i == 0 {
@@ -541,14 +553,14 @@ mod tests {
 
         loop {
             let mut done = vec![];
-            while let Some(Next::Ready(k)) = workplan.next() {
-                println!("Doing {}", k.k());
-                done.push(k);
+            while let Some(Next::Ready(item)) = workplan.next() {
+                println!("Doing {}", item.k());
+                done.push(item);
                 workplan.t.print();
             }
 
             for d in done.iter() {
-                workplan.done(d).unwrap();
+                workplan.done(d.k()).unwrap();
             }
             if done.is_empty() {
                 break;
@@ -599,7 +611,7 @@ mod tests {
                         count_done += 1;
                     }
                 }
-                workplan.done(&item).unwrap();
+                workplan.done(item.k()).unwrap();
             }
 
             assert_eq!(count_done, mt.nodes.len());