From 52300a114242c2d2949db20bc1e4ab4a9c12fb14 Mon Sep 17 00:00:00 2001
From: Pierre Krieger <pierre.krieger1708@gmail.com>
Date: Mon, 5 Sep 2022 22:56:18 +0200
Subject: [PATCH] Add some tests for remove_prefix (#2714)

Adds some tests regarding the code added in https://github.com/paritytech/smoldot/pull/2663

I wanted to maybe figure out why https://github.com/paritytech/smoldot/pull/2666 was failing, but didn't manage to find a failing case.

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
---
 src/trie/trie_structure/tests.rs | 281 +++++++++++++++++++++++++++++++
 1 file changed, 281 insertions(+)

diff --git a/src/trie/trie_structure/tests.rs b/src/trie/trie_structure/tests.rs
index 798f6f864d..5de123e83f 100644
--- a/src/trie/trie_structure/tests.rs
+++ b/src/trie/trie_structure/tests.rs
@@ -325,6 +325,287 @@ fn insert_branch() {
         .has_storage_value());
 }
 
+#[test]
+fn remove_prefix_basic() {
+    let mut trie = TrieStructure::new();
+
+    trie.node([Nibble::try_from(1).unwrap()].iter().cloned())
+        .into_vacant()
+        .unwrap()
+        .insert_storage_value()
+        .insert((), ());
+    trie.node(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(3).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    )
+    .into_vacant()
+    .unwrap()
+    .insert_storage_value()
+    .insert((), ());
+    trie.node(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(4).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    )
+    .into_vacant()
+    .unwrap()
+    .insert_storage_value()
+    .insert((), ());
+    trie.node(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(4).unwrap(),
+            Nibble::try_from(5).unwrap(),
+            Nibble::try_from(6).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    )
+    .into_vacant()
+    .unwrap()
+    .insert_storage_value()
+    .insert((), ());
+
+    trie.remove_prefix(
+        [Nibble::try_from(1).unwrap(), Nibble::try_from(2).unwrap()]
+            .iter()
+            .cloned(),
+    );
+
+    assert_eq!(trie.len(), 1);
+    assert!(trie
+        .node([Nibble::try_from(1).unwrap(),].iter().cloned(),)
+        .into_occupied()
+        .unwrap()
+        .has_storage_value());
+}
+
+#[test]
+fn remove_prefix_clear_all() {
+    let mut trie = TrieStructure::new();
+
+    trie.node(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(3).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    )
+    .into_vacant()
+    .unwrap()
+    .insert_storage_value()
+    .insert((), ());
+    trie.node(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(4).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    )
+    .into_vacant()
+    .unwrap()
+    .insert_storage_value()
+    .insert((), ());
+    trie.node(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(4).unwrap(),
+            Nibble::try_from(5).unwrap(),
+            Nibble::try_from(6).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    )
+    .into_vacant()
+    .unwrap()
+    .insert_storage_value()
+    .insert((), ());
+
+    trie.remove_prefix(
+        [Nibble::try_from(1).unwrap(), Nibble::try_from(2).unwrap()]
+            .iter()
+            .cloned(),
+    );
+
+    assert!(trie.is_empty());
+}
+
+#[test]
+fn remove_prefix_exact() {
+    let mut trie = TrieStructure::new();
+
+    trie.node([Nibble::try_from(1).unwrap()].iter().cloned())
+        .into_vacant()
+        .unwrap()
+        .insert_storage_value()
+        .insert((), ());
+
+    trie.node(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(3).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    )
+    .into_vacant()
+    .unwrap()
+    .insert_storage_value()
+    .insert((), ());
+    trie.node(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(3).unwrap(),
+            Nibble::try_from(4).unwrap(),
+            Nibble::try_from(5).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    )
+    .into_vacant()
+    .unwrap()
+    .insert_storage_value()
+    .insert((), ());
+    trie.node(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(3).unwrap(),
+            Nibble::try_from(4).unwrap(),
+            Nibble::try_from(6).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    )
+    .into_vacant()
+    .unwrap()
+    .insert_storage_value()
+    .insert((), ());
+
+    trie.remove_prefix(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(3).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    );
+
+    assert_eq!(trie.len(), 1);
+    assert!(trie
+        .node([Nibble::try_from(1).unwrap(),].iter().cloned(),)
+        .into_occupied()
+        .unwrap()
+        .has_storage_value());
+}
+
+#[test]
+fn remove_prefix_doesnt_match_anything() {
+    let mut trie = TrieStructure::new();
+
+    trie.node(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(3).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    )
+    .into_vacant()
+    .unwrap()
+    .insert_storage_value()
+    .insert((), ());
+
+    trie.remove_prefix(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(5).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    );
+
+    assert_eq!(trie.len(), 1);
+    assert!(trie
+        .node(
+            [
+                Nibble::try_from(1).unwrap(),
+                Nibble::try_from(2).unwrap(),
+                Nibble::try_from(3).unwrap(),
+            ]
+            .iter()
+            .cloned(),
+        )
+        .into_occupied()
+        .unwrap()
+        .has_storage_value());
+}
+
+#[test]
+fn remove_prefix_nothing_to_remove() {
+    let mut trie = TrieStructure::new();
+
+    trie.node(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(3).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    )
+    .into_vacant()
+    .unwrap()
+    .insert_storage_value()
+    .insert((), ());
+
+    trie.remove_prefix(
+        [
+            Nibble::try_from(1).unwrap(),
+            Nibble::try_from(2).unwrap(),
+            Nibble::try_from(3).unwrap(),
+            Nibble::try_from(4).unwrap(),
+        ]
+        .iter()
+        .cloned(),
+    );
+
+    assert_eq!(trie.len(), 1);
+    assert!(trie
+        .node(
+            [
+                Nibble::try_from(1).unwrap(),
+                Nibble::try_from(2).unwrap(),
+                Nibble::try_from(3).unwrap(),
+            ]
+            .iter()
+            .cloned(),
+        )
+        .into_occupied()
+        .unwrap()
+        .has_storage_value());
+}
+
 #[test]
 fn fuzzing() {
     fn uniform_sample(min: u8, max: u8) -> u8 {