On all tails helper (#235)

Alex6323 · Sep 15, 2020 · 455b914 · 455b914
1 parent 8173e56
commit 455b914
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 18 deletions.
diff --git a/bee-protocol/src/tangle/helper.rs b/bee-protocol/src/tangle/helper.rs
@@ -11,7 +11,7 @@
 
 use crate::tangle::TransactionMetadata;
 
-use bee_tangle::{traversal, Tangle};
+use bee_tangle::{traversal, Tangle, TransactionRef};
 
 use bee_crypto::ternary::Hash;
 
@@ -39,3 +39,18 @@ pub(crate) fn find_tail_of_bundle(tangle: &Tangle<TransactionMetadata>, root: Ha
 
     tail
 }
+
+pub fn on_all_tails<Apply: FnMut(&Hash, &TransactionRef, &TransactionMetadata)>(
+    tangle: &Tangle<TransactionMetadata>,
+    root: Hash,
+    apply: Apply,
+) {
+    traversal::visit_parents_depth_first(
+        tangle,
+        root,
+        |_, _, metadata| !metadata.flags.is_tail(),
+        |_, _, _| {},
+        apply,
+        |_| {},
+    );
+}
diff --git a/bee-protocol/src/tangle/mod.rs b/bee-protocol/src/tangle/mod.rs
@@ -71,10 +71,6 @@ impl MsTangle {
         opt
     }
 
-    pub fn get_metadata(&self, hash: &Hash) -> Option<TransactionMetadata> {
-        self.inner.get_metadata(hash)
-    }
-
     pub fn add_milestone(&self, index: MilestoneIndex, hash: Hash) {
         // TODO: only insert if vacant
         self.milestones.insert(index, hash);

diff --git a/bee-snapshot/src/pruning/mod.rs b/bee-snapshot/src/pruning/mod.rs
@@ -16,7 +16,10 @@ pub use config::{PruningConfig, PruningConfigBuilder};
 use crate::local::{LocalSnapshotConfig, LocalSnapshotMetadata};
 
 use bee_crypto::ternary::Hash;
-use bee_protocol::{tangle::tangle, MilestoneIndex};
+use bee_protocol::{
+    tangle::{helper, tangle},
+    MilestoneIndex,
+};
 use bee_tangle::traversal;
 
 use dashmap::DashMap;
@@ -104,15 +107,9 @@ pub fn get_new_solid_entry_points(target_index: MilestoneIndex) -> Result<DashMa
                 // Find all tails
                 let mut tail_hashes: Vec<Hash> = Vec::new();
 
-                // Get all the tails
-                traversal::visit_parents_depth_first(
-                    tangle(),
-                    approvee,
-                    |_hash, _tx, metadata| !metadata.flags.is_tail(),
-                    |_hash, _tx, _metadata| {},
-                    |hash, _tx, _metadata| tail_hashes.push(hash.clone()),
-                    |_hash| {},
-                );
+                helper::on_all_tails(tangle(), approvee, |hash, _tx, _metadata| {
+                    tail_hashes.push(hash.clone())
+                });
 
                 for tail_hash in tail_hashes {
                     match tangle().get_metadata(&tail_hash) {

diff --git a/bee-tangle/src/traversal.rs b/bee-tangle/src/traversal.rs
@@ -162,10 +162,7 @@ pub fn visit_children_depth_first<Metadata, Match, Apply, ElseApply>(
                 }
             }
             None => {
-                // NOTE: this has to be dealt at the protocol level now ;)
-                // if !tangle.solid_entry_points.contains(hash) {
                 else_apply(hash);
-                //}
                 visited.insert(*hash);
                 children.pop();
             }