Add compute_unique_stats

broxus · Sep 1, 2023 · ffebd52 · ffebd52
1 parent 2452d93
commit ffebd52
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 9 deletions.
diff --git a/src/cell/mod.rs b/src/cell/mod.rs
@@ -144,6 +144,8 @@ pub trait CellImpl {
 
     /// Returns the sum of all bits and cells of all elements in the cell tree
     /// (including this cell).
+    ///
+    /// NOTE: identical cells are counted each time they occur in the tree.
     #[cfg(feature = "stats")]
     fn stats(&self) -> CellTreeStats;
 }
@@ -235,6 +237,44 @@ impl DynCell {
         CellSlice::new_unchecked(self)
     }
 
+    /// Recursively computes the count of distinct cells returning
+    /// the total storage used by this dag taking into account the
+    /// identification of equal cells.
+    pub fn compute_unique_stats(&self, limit: usize) -> Option<CellTreeStats> {
+        let mut visited = ahash::HashSet::<&HashBytes>::default();
+        let mut stack = vec![self.references()];
+
+        let mut stats = CellTreeStats {
+            bit_count: self.bit_len() as u64,
+            cell_count: 1,
+        };
+
+        'outer: while let Some(item) = stack.last_mut() {
+            for cell in item.by_ref() {
+                if !visited.insert(cell.repr_hash()) {
+                    continue;
+                }
+
+                if stats.cell_count >= limit as u64 {
+                    return None;
+                }
+
+                stats.bit_count += cell.bit_len() as u64;
+                stats.cell_count += 1;
+
+                let next = cell.references();
+                if next.max > 0 {
+                    stack.push(next);
+                    continue 'outer;
+                }
+            }
+
+            stack.pop();
+        }
+
+        Some(stats)
+    }
+
     /// Returns an object that implements [`Debug`] for printing only
     /// the root cell of the cell tree.
     ///
@@ -1062,18 +1102,14 @@ impl Iterator for LevelMaskIter {
 }
 
 /// Cell tree storage stats.
-///
-/// NOTE: identical cells are counted each time they occur in the tree.
 #[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
-#[cfg(feature = "stats")]
 pub struct CellTreeStats {
     /// Total number of bits in tree.
     pub bit_count: u64,
     /// Total number of cells in tree.
     pub cell_count: u64,
 }
 
-#[cfg(feature = "stats")]
 impl std::ops::Add for CellTreeStats {
     type Output = Self;
 
@@ -1086,7 +1122,6 @@ impl std::ops::Add for CellTreeStats {
     }
 }
 
-#[cfg(feature = "stats")]
 impl std::ops::AddAssign for CellTreeStats {
     #[inline]
     fn add_assign(&mut self, rhs: Self) {

diff --git a/src/lib.rs b/src/lib.rs
@@ -149,7 +149,7 @@ pub mod error;
 
 #[cfg(test)]
 mod tests {
-    use crate::cell::MAX_BIT_LEN;
+    use crate::cell::{CellTreeStats, MAX_BIT_LEN};
     use crate::prelude::*;
     use crate::util::decode_base64;
 
@@ -174,7 +174,16 @@ mod tests {
     #[test]
     fn big_cell_deserialization() {
         let data = decode_base64("te6ccgIDAAwAAQAAAACIAAAEBAABAAEAAQABAAEEBAACAAIAAgACAAIEBAADAAMAAwADAAMEBAAEAAQABAAEAAQEBAAFAAUABQAFAAUEBAAGAAYABgAGAAYEBAAHAAcABwAHAAcEBAAIAAgACAAIAAgEBAAJAAkACQAJAAkEBAAKAAoACgAKAAoEBAALAAsACwALAAsABAAA").unwrap();
-        _ = Boc::decode(data).unwrap();
+        let cell = Boc::decode(data).unwrap();
+
+        let stats = cell.compute_unique_stats(1 << 22).unwrap();
+        assert_eq!(
+            stats,
+            CellTreeStats {
+                bit_count: 192,
+                cell_count: 12
+            }
+        );
     }
 
     #[test]
@@ -250,6 +259,15 @@ mod tests {
 
     #[test]
     fn test_tx() {
-        Boc::decode_base64("te6ccgICAQoAAQAADGkAAAO3ea37gczcXLp00bkP3eA1txaTwX6TyzGtowSuHiFwobmgAAF3fHG0RBrAoqQhyfVHKxY+b4xigHnXHqftp9X5vfYVKuY58i4/cAABd3p8EkwWJgK1gAA0gEVmAigABQAEAAECEQyBbEYb1mwEQAADAAIAb8mHoSBMFFhAAAAAAAACAAAAAAADMQg15pv/2PjjbqZFi59+K/39f1kPXUGLckkscjpa2sJAUBYMAJ1D7gMTiAAAAAAAAAAANAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgAIJyl+oF61WYJFz0URNA5vMfkcc7dxHYfH6w0cmoXG2Ro2za6+U+LRtB2aSLAAMVTmTPucTOeWBEjz1nOjURo9Gg/wIB4AAIAAYBAd8ABwCxSAE1v3A5m4uXTpo3Ifu8Brbi0ngv0nlmNbRglcPELhQ3NQAxLah1y23nqb6T3ERREC7LXfYeMu26LwYH1Ht6c3lDQZDuaygABhRYYAAALu+ONoiExMBWsEABRYgBNb9wOZuLl06aNyH7vAa24tJ4L9J5ZjW0YJXDxC4UNzQMAAkB4fZ7eRCTQYwyOQPFDYjRpK0QMs7JDtGuaerLBmn2TDLl25hSY50SC7Nnc6gIFU3xYshpJ4j3tGtYPCPCMXRuJgTPXNlw4YdSq3zWEWMJOr0f83TQcuo2IkFjiPQacwNzkMAAAGAR6lJjmJgK5JM7mRsgAAoBZYAYltQ65bbz1N9J7iIoiBdlrvsPGXbdF4MD6j29ObyhoMAAAAAAAAAAAAAAAAdzWUAAOAALBAAADAAMAAwADAQAAA0ADQANAA0EAAAOAA4ADgAOBAAADwAPAA8ADwQAABAAEAAQABAEAAARABEAEQARBAAAEgASABIAEgQAABMAEwATABMEAAAUABQAFAAUBAAAFQAVABUAFQQAABYAFgAWABYEAAAXABcAFwAXBAAAGAAYABgAGAQAABkAGQAZABkEAAAaABoAGgAaBAAAGwAbABsAGwQAABwAHAAcABwEAAAdAB0AHQAdBAAAHgAeAB4AHgQAAB8AHwAfAB8EAAAgACAAIAAgBAAAIQAhACEAIQQAACIAIgAiACIEAAAjACMAIwAjBAAAJAAkACQAJAQAACUAJQAlACUEAAAmACYAJgAmBAAAJwAnACcAJwQAACgAKAAoACgEAAApACkAKQApBAAAKgAqACoAKgQAACsAKwArACsEAAAsACwALAAsBAAALQAtAC0ALQQAAC4ALgAuAC4EAAAvAC8ALwAvBAAAMAAwADAAMAQAADEAMQAxADEEAAAyADIAMgAyBAAAMwAzADMAMwQAADQANAA0ADQEAAA1ADUANQA1BAAANgA2ADYANgQAADcANwA3ADcEAAA4ADgAOAA4BAAAOQA5ADkAOQQAADoAOgA6ADoEAAA7ADsAOwA7BAAAPAA8ADwAPAQAAD0APQA9AD0EAAA+AD4APgA+BAAAPwA/AD8APwQAAEAAQABAAEAEAABBAEEAQQBBBAAAQgBCAEIAQgQAAEMAQwBDAEMEAABEAEQARABEBAAARQBFAEUARQQAAEYARgBGAEYEAABHAEcARwBHBAAASABIAEgASAQAAEkASQBJAEkEAABKAEoASgBKBAAASwBLAEsASwQAAEwATABMAEwEAABNAE0ATQBNBAAATgBOAE4ATgQAAE8ATwBPAE8EAABQAFAAUABQBAAAUQBRAFEAUQQAAFIAUgBSAFIEAABTAFMAUwBTBAAAVABUAFQAVAQAAFUAVQBVAFUEAABWAFYAVgBWBAAAVwBXAFcAVwQAAFgAWABYAFgEAABZAFkAWQBZBAAAWgBaAFoAWgQAAFsAWwBbAFsEAABcAFwAXABcBAAAXQBdAF0AXQQAAF4AXgBeAF4EAABfAF8AXwBfBAAAYABgAGAAYAQAAGEAYQBhAGEEAABiAGIAYgBiBAAAYwBjAGMAYwQAAGQAZABkAGQEAABlAGUAZQBlBAAAZgBmAGYAZgQAAGcAZwBnAGcEAABoAGgAaABoBAAAaQBpAGkAaQQAAGoAagBqAGoEAABrAGsAawBrBAAAbABsAGwAbAQAAG0AbQBtAG0EAABuAG4AbgBuBAAAbwBvAG8AbwQAAHAAcABwAHAEAABxAHEAcQBxBAAAcgByAHIAcgQAAHMAcwBzAHMEAAB0AHQAdAB0BAAAdQB1AHUAdQQAAHYAdgB2AHYEAAB3AHcAdwB3BAAAeAB4AHgAeAQAAHkAeQB5AHkEAAB6AHoAegB6BAAAewB7AHsAewQAAHwAfAB8AHwEAAB9AH0AfQB9BAAAfgB+AH4AfgQAAH8AfwB/AH8EAACAAIAAgACABAAAgQCBAIEAgQQAAIIAggCCAIIEAACDAIMAgwCDBAAAhACEAIQAhAQAAIUAhQCFAIUEAACGAIYAhgCGBAAAhwCHAIcAhwQAAIgAiACIAIgEAACJAIkAiQCJBAAAigCKAIoAigQAAIsAiwCLAIsEAACMAIwAjACMBAAAjQCNAI0AjQQAAI4AjgCOAI4EAACPAI8AjwCPBAAAkACQAJAAkAQAAJEAkQCRAJEEAACSAJIAkgCSBAAAkwCTAJMAkwQAAJQAlACUAJQEAACVAJUAlQCVBAAAlgCWAJYAlgQAAJcAlwCXAJcEAACYAJgAmACYBAAAmQCZAJkAmQQAAJoAmgCaAJoEAACbAJsAmwCbBAAAnACcAJwAnAQAAJ0AnQCdAJ0EAACeAJ4AngCeBAAAnwCfAJ8AnwQAAKAAoACgAKAEAAChAKEAoQChBAAAogCiAKIAogQAAKMAowCjAKMEAACkAKQApACkBAAApQClAKUApQQAAKYApgCmAKYEAACnAKcApwCnBAAAqACoAKgAqAQAAKkAqQCpAKkEAACqAKoAqgCqBAAAqwCrAKsAqwQAAKwArACsAKwEAACtAK0ArQCtBAAArgCuAK4ArgQAAK8ArwCvAK8EAACwALAAsACwBAAAsQCxALEAsQQAALIAsgCyALIEAACzALMAswCzBAAAtAC0ALQAtAQAALUAtQC1ALUEAAC2ALYAtgC2BAAAtwC3ALcAtwQAALgAuAC4ALgEAAC5ALkAuQC5BAAAugC6ALoAugQAALsAuwC7ALsEAAC8ALwAvAC8BAAAvQC9AL0AvQQAAL4AvgC+AL4EAAC/AL8AvwC/BAAAwADAAMAAwAQAAMEAwQDBAMEEAADCAMIAwgDCBAAAwwDDAMMAwwQAAMQAxADEAMQEAADFAMUAxQDFBAAAxgDGAMYAxgQAAMcAxwDHAMcEAADIAMgAyADIBAAAyQDJAMkAyQQAAMoAygDKAMoEAADLAMsAywDLBAAAzADMAMwAzAQAAM0AzQDNAM0EAADOAM4AzgDOBAAAzwDPAM8AzwQAANAA0ADQANAEAADRANEA0QDRBAAA0gDSANIA0gQAANMA0wDTANMEAADUANQA1ADUBAAA1QDVANUA1QQAANYA1gDWANYEAADXANcA1wDXBAAA2ADYANgA2AQAANkA2QDZANkEAADaANoA2gDaBAAA2wDbANsA2wQAANwA3ADcANwEAADdAN0A3QDdBAAA3gDeAN4A3gQAAN8A3wDfAN8EAADgAOAA4ADgBAAA4QDhAOEA4QQAAOIA4gDiAOIEAADjAOMA4wDjBAAA5ADkAOQA5AQAAOUA5QDlAOUEAADmAOYA5gDmBAAA5wDnAOcA5wQAAOgA6ADoAOgEAADpAOkA6QDpBAAA6gDqAOoA6gQAAOsA6wDrAOsEAADsAOwA7ADsBAAA7QDtAO0A7QQAAO4A7gDuAO4EAADvAO8A7wDvBAAA8ADwAPAA8AQAAPEA8QDxAPEEAADyAPIA8gDyBAAA8wDzAPMA8wQAAPQA9AD0APQEAAD1APUA9QD1BAAA9gD2APYA9gQAAPcA9wD3APcEAAD4APgA+AD4BAAA+QD5APkA+QQAAPoA+gD6APoEAAD7APsA+wD7BAAA/AD8APwA/AQAAP0A/QD9AP0EAAD+AP4A/gD+BAAA/wD/AP8A/wQAAQABAAEAAQAEAAEBAQEBAQEBBAABAgECAQIBAgQAAQMBAwEDAQMEAAEEAQQBBAEEBAABBQEFAQUBBQQAAQYBBgEGAQYEAAEHAQcBBwEHBAABCAEIAQgBCAQAAQkBCQEJAQkAAA==").unwrap();
+        let cell = Boc::decode_base64("te6ccgICAQoAAQAADGkAAAO3ea37gczcXLp00bkP3eA1txaTwX6TyzGtowSuHiFwobmgAAF3fHG0RBrAoqQhyfVHKxY+b4xigHnXHqftp9X5vfYVKuY58i4/cAABd3p8EkwWJgK1gAA0gEVmAigABQAEAAECEQyBbEYb1mwEQAADAAIAb8mHoSBMFFhAAAAAAAACAAAAAAADMQg15pv/2PjjbqZFi59+K/39f1kPXUGLckkscjpa2sJAUBYMAJ1D7gMTiAAAAAAAAAAANAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgAIJyl+oF61WYJFz0URNA5vMfkcc7dxHYfH6w0cmoXG2Ro2za6+U+LRtB2aSLAAMVTmTPucTOeWBEjz1nOjURo9Gg/wIB4AAIAAYBAd8ABwCxSAE1v3A5m4uXTpo3Ifu8Brbi0ngv0nlmNbRglcPELhQ3NQAxLah1y23nqb6T3ERREC7LXfYeMu26LwYH1Ht6c3lDQZDuaygABhRYYAAALu+ONoiExMBWsEABRYgBNb9wOZuLl06aNyH7vAa24tJ4L9J5ZjW0YJXDxC4UNzQMAAkB4fZ7eRCTQYwyOQPFDYjRpK0QMs7JDtGuaerLBmn2TDLl25hSY50SC7Nnc6gIFU3xYshpJ4j3tGtYPCPCMXRuJgTPXNlw4YdSq3zWEWMJOr0f83TQcuo2IkFjiPQacwNzkMAAAGAR6lJjmJgK5JM7mRsgAAoBZYAYltQ65bbz1N9J7iIoiBdlrvsPGXbdF4MD6j29ObyhoMAAAAAAAAAAAAAAAAdzWUAAOAALBAAADAAMAAwADAQAAA0ADQANAA0EAAAOAA4ADgAOBAAADwAPAA8ADwQAABAAEAAQABAEAAARABEAEQARBAAAEgASABIAEgQAABMAEwATABMEAAAUABQAFAAUBAAAFQAVABUAFQQAABYAFgAWABYEAAAXABcAFwAXBAAAGAAYABgAGAQAABkAGQAZABkEAAAaABoAGgAaBAAAGwAbABsAGwQAABwAHAAcABwEAAAdAB0AHQAdBAAAHgAeAB4AHgQAAB8AHwAfAB8EAAAgACAAIAAgBAAAIQAhACEAIQQAACIAIgAiACIEAAAjACMAIwAjBAAAJAAkACQAJAQAACUAJQAlACUEAAAmACYAJgAmBAAAJwAnACcAJwQAACgAKAAoACgEAAApACkAKQApBAAAKgAqACoAKgQAACsAKwArACsEAAAsACwALAAsBAAALQAtAC0ALQQAAC4ALgAuAC4EAAAvAC8ALwAvBAAAMAAwADAAMAQAADEAMQAxADEEAAAyADIAMgAyBAAAMwAzADMAMwQAADQANAA0ADQEAAA1ADUANQA1BAAANgA2ADYANgQAADcANwA3ADcEAAA4ADgAOAA4BAAAOQA5ADkAOQQAADoAOgA6ADoEAAA7ADsAOwA7BAAAPAA8ADwAPAQAAD0APQA9AD0EAAA+AD4APgA+BAAAPwA/AD8APwQAAEAAQABAAEAEAABBAEEAQQBBBAAAQgBCAEIAQgQAAEMAQwBDAEMEAABEAEQARABEBAAARQBFAEUARQQAAEYARgBGAEYEAABHAEcARwBHBAAASABIAEgASAQAAEkASQBJAEkEAABKAEoASgBKBAAASwBLAEsASwQAAEwATABMAEwEAABNAE0ATQBNBAAATgBOAE4ATgQAAE8ATwBPAE8EAABQAFAAUABQBAAAUQBRAFEAUQQAAFIAUgBSAFIEAABTAFMAUwBTBAAAVABUAFQAVAQAAFUAVQBVAFUEAABWAFYAVgBWBAAAVwBXAFcAVwQAAFgAWABYAFgEAABZAFkAWQBZBAAAWgBaAFoAWgQAAFsAWwBbAFsEAABcAFwAXABcBAAAXQBdAF0AXQQAAF4AXgBeAF4EAABfAF8AXwBfBAAAYABgAGAAYAQAAGEAYQBhAGEEAABiAGIAYgBiBAAAYwBjAGMAYwQAAGQAZABkAGQEAABlAGUAZQBlBAAAZgBmAGYAZgQAAGcAZwBnAGcEAABoAGgAaABoBAAAaQBpAGkAaQQAAGoAagBqAGoEAABrAGsAawBrBAAAbABsAGwAbAQAAG0AbQBtAG0EAABuAG4AbgBuBAAAbwBvAG8AbwQAAHAAcABwAHAEAABxAHEAcQBxBAAAcgByAHIAcgQAAHMAcwBzAHMEAAB0AHQAdAB0BAAAdQB1AHUAdQQAAHYAdgB2AHYEAAB3AHcAdwB3BAAAeAB4AHgAeAQAAHkAeQB5AHkEAAB6AHoAegB6BAAAewB7AHsAewQAAHwAfAB8AHwEAAB9AH0AfQB9BAAAfgB+AH4AfgQAAH8AfwB/AH8EAACAAIAAgACABAAAgQCBAIEAgQQAAIIAggCCAIIEAACDAIMAgwCDBAAAhACEAIQAhAQAAIUAhQCFAIUEAACGAIYAhgCGBAAAhwCHAIcAhwQAAIgAiACIAIgEAACJAIkAiQCJBAAAigCKAIoAigQAAIsAiwCLAIsEAACMAIwAjACMBAAAjQCNAI0AjQQAAI4AjgCOAI4EAACPAI8AjwCPBAAAkACQAJAAkAQAAJEAkQCRAJEEAACSAJIAkgCSBAAAkwCTAJMAkwQAAJQAlACUAJQEAACVAJUAlQCVBAAAlgCWAJYAlgQAAJcAlwCXAJcEAACYAJgAmACYBAAAmQCZAJkAmQQAAJoAmgCaAJoEAACbAJsAmwCbBAAAnACcAJwAnAQAAJ0AnQCdAJ0EAACeAJ4AngCeBAAAnwCfAJ8AnwQAAKAAoACgAKAEAAChAKEAoQChBAAAogCiAKIAogQAAKMAowCjAKMEAACkAKQApACkBAAApQClAKUApQQAAKYApgCmAKYEAACnAKcApwCnBAAAqACoAKgAqAQAAKkAqQCpAKkEAACqAKoAqgCqBAAAqwCrAKsAqwQAAKwArACsAKwEAACtAK0ArQCtBAAArgCuAK4ArgQAAK8ArwCvAK8EAACwALAAsACwBAAAsQCxALEAsQQAALIAsgCyALIEAACzALMAswCzBAAAtAC0ALQAtAQAALUAtQC1ALUEAAC2ALYAtgC2BAAAtwC3ALcAtwQAALgAuAC4ALgEAAC5ALkAuQC5BAAAugC6ALoAugQAALsAuwC7ALsEAAC8ALwAvAC8BAAAvQC9AL0AvQQAAL4AvgC+AL4EAAC/AL8AvwC/BAAAwADAAMAAwAQAAMEAwQDBAMEEAADCAMIAwgDCBAAAwwDDAMMAwwQAAMQAxADEAMQEAADFAMUAxQDFBAAAxgDGAMYAxgQAAMcAxwDHAMcEAADIAMgAyADIBAAAyQDJAMkAyQQAAMoAygDKAMoEAADLAMsAywDLBAAAzADMAMwAzAQAAM0AzQDNAM0EAADOAM4AzgDOBAAAzwDPAM8AzwQAANAA0ADQANAEAADRANEA0QDRBAAA0gDSANIA0gQAANMA0wDTANMEAADUANQA1ADUBAAA1QDVANUA1QQAANYA1gDWANYEAADXANcA1wDXBAAA2ADYANgA2AQAANkA2QDZANkEAADaANoA2gDaBAAA2wDbANsA2wQAANwA3ADcANwEAADdAN0A3QDdBAAA3gDeAN4A3gQAAN8A3wDfAN8EAADgAOAA4ADgBAAA4QDhAOEA4QQAAOIA4gDiAOIEAADjAOMA4wDjBAAA5ADkAOQA5AQAAOUA5QDlAOUEAADmAOYA5gDmBAAA5wDnAOcA5wQAAOgA6ADoAOgEAADpAOkA6QDpBAAA6gDqAOoA6gQAAOsA6wDrAOsEAADsAOwA7ADsBAAA7QDtAO0A7QQAAO4A7gDuAO4EAADvAO8A7wDvBAAA8ADwAPAA8AQAAPEA8QDxAPEEAADyAPIA8gDyBAAA8wDzAPMA8wQAAPQA9AD0APQEAAD1APUA9QD1BAAA9gD2APYA9gQAAPcA9wD3APcEAAD4APgA+AD4BAAA+QD5APkA+QQAAPoA+gD6APoEAAD7APsA+wD7BAAA/AD8APwA/AQAAP0A/QD9AP0EAAD+AP4A/gD+BAAA/wD/AP8A/wQAAQABAAEAAQAEAAEBAQEBAQEBBAABAgECAQIBAgQAAQMBAwEDAQMEAAEEAQQBBAEEBAABBQEFAQUBBQQAAQYBBgEGAQYEAAEHAQcBBwEHBAABCAEIAQgBCAQAAQkBCQEJAQkAAA==").unwrap();
+
+        let stats = cell.compute_unique_stats(1 << 22).unwrap();
+        assert_eq!(
+            stats,
+            CellTreeStats {
+                bit_count: 4681,
+                cell_count: 266
+            }
+        );
     }
 }
diff --git a/src/merkle/tests/mod.rs b/src/merkle/tests/mod.rs
@@ -1,5 +1,5 @@
 use super::*;
-use crate::cell::EMPTY_CELL_HASH;
+use crate::cell::{CellTreeStats, EMPTY_CELL_HASH};
 use crate::error::Error;
 use crate::prelude::*;
 
@@ -39,6 +39,15 @@ fn create_proof_for_deep_cell() {
         cell = builder.build().unwrap();
     }
 
+    let stats = cell.compute_unique_stats(1 << 22).unwrap();
+    assert_eq!(
+        stats,
+        CellTreeStats {
+            bit_count: 65000 * 32,
+            cell_count: 65001
+        }
+    );
+
     {
         let encoded = Boc::encode_base64(cell.as_ref());
         let decoded = Boc::decode_base64(encoded).unwrap();