Implement TrieDB::is_complete

Author: remagpie

util/merkle/src/triedb.rs

@@ -14,7 +14,7 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-use ccrypto::blake256;
+use ccrypto::{blake256, BLAKE_NULL_RLP};
 use hashdb::HashDB;
 use primitives::H256;
 
@@ -108,7 +108,22 @@ impl<'db> TrieDB<'db> {
 
     /// Check if every leaf of the trie exists
     pub fn is_complete(&self) -> bool {
-        unimplemented!()
+        *self.root == BLAKE_NULL_RLP || 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().flatten().all(|child| self.is_complete_aux(child))
+                }
+                Some(RlpNode::Leaf(..)) => true,
+                None => false,
+            }
+        } else {
+            false
+        }
     }
 }
 
@@ -146,4 +161,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());
+    }
 }