Implement TrieDB::is_complete

Author: remagpie

util/merkle/src/triedb.rs

@@ -108,7 +108,23 @@ impl<'db> TrieDB<'db> {
 
     /// Check if every leaf of the trie exists
     pub fn is_complete(&self) -> bool {
-        unimplemented!()
+        self.is_complete_aux(self.root)
+    }
+
+    /// Check if every leaf of the trie starting from `hash` exists
+    pub fn is_complete_aux(&self, hash: &H256) -> bool {
+        if let Some(node_rlp) = self.db.get(hash) {
+            match RlpNode::decoded(node_rlp.as_ref()) {
+                Some(RlpNode::Branch(.., children)) => children.iter().all(|child| match child {
+                    Some(h) => self.is_complete_aux(h),
+                    None => true,
+                }),
+                Some(RlpNode::Leaf(..)) => true,
+                None => false,
+            }
+        } else {
+            false
+        }
     }
 }
 
@@ -146,4 +162,33 @@ mod tests {
         assert_eq!(t.get(b"B"), Ok(Some(DBValue::from_slice(b"ABCBA"))));
         assert_eq!(t.get(b"C"), Ok(None));
     }
+
+    #[test]
+    fn is_complete_success() {
+        let mut memdb = MemoryDB::new();
+        let mut root = H256::new();
+        {
+            let mut t = TrieDBMut::new(&mut memdb, &mut root);
+            t.insert(b"A", b"ABC").unwrap();
+            t.insert(b"B", b"ABCBA").unwrap();
+        }
+
+        let t = TrieDB::try_new(&memdb, &root).unwrap();
+        assert!(t.is_complete());
+    }
+
+    #[test]
+    fn is_complete_fail() {
+        let mut memdb = MemoryDB::new();
+        let mut root = H256::new();
+        {
+            let mut t = TrieDBMut::new(&mut memdb, &mut root);
+            t.insert(b"A", b"ABC").unwrap();
+            t.insert(b"B", b"ABCBA").unwrap();
+        }
+        memdb.remove(memdb.keys().keys().next().unwrap());
+
+        let t = TrieDB::try_new(&memdb, &root).unwrap();
+        assert!(!t.is_complete());
+    }
 }