From 719aa79140b889d5ca7cc7e72f5d42a065f8032d Mon Sep 17 00:00:00 2001
From: cheme <emericchevalier.pro@gmail.com>
Date: Sun, 13 Aug 2023 13:28:12 +0200
Subject: [PATCH] fixing triedbmut lookup, added some testing in test. (#198)

---
 trie-db/src/triedbmut.rs      |  9 +++---
 trie-db/test/src/triedbmut.rs | 53 ++++++++++++++++++++++++++++++++++-
 2 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/trie-db/src/triedbmut.rs b/trie-db/src/triedbmut.rs
index d4eb934a..bacc7f51 100644
--- a/trie-db/src/triedbmut.rs
+++ b/trie-db/src/triedbmut.rs
@@ -975,6 +975,7 @@ where
 		handle: &NodeHandle<TrieHash<L>, L::Location>,
 	) -> Result<Option<DBValue>, TrieHash<L>, CError<L>> {
 		let mut handle = handle;
+		// prefix only use for value node access, so this is always correct.
 		let prefix = (full_key, None);
 		loop {
 			let (mid, child) = match handle {
@@ -994,8 +995,8 @@ where
 				},
 				NodeHandle::InMemory(handle) => match &self.storage[handle] {
 					Node::Empty => return Ok(None),
-					Node::Leaf(key, value) =>
-						if NibbleSlice::from_stored(key) == partial {
+					Node::Leaf(slice, value) =>
+						if NibbleSlice::from_stored(slice) == partial {
 							return Ok(value.in_memory_fetched_value(
 								prefix,
 								self.db,
@@ -1046,7 +1047,7 @@ where
 								None
 							})
 						} else if partial.starts_with(&slice) {
-							let idx = partial.at(0);
+							let idx = partial.at(slice.len());
 							match children[idx as usize].as_ref() {
 								Some(child) => (1 + slice.len(), child),
 								None => return Ok(None),
@@ -2016,7 +2017,7 @@ where
 
 	/// Cache the given `encoded` node.
 	fn cache_node(&mut self, hash: TrieHash<L>) {
-		// If we have a cache, cache our node directly.
+		// Mark the node as new so that it is removed from the shared cache.
 		if let Some(cache) = self.cache.as_mut() {
 			cache.insert_new_node(&hash);
 		}
diff --git a/trie-db/test/src/triedbmut.rs b/trie-db/test/src/triedbmut.rs
index 2a9bdd59..371236dc 100644
--- a/trie-db/test/src/triedbmut.rs
+++ b/trie-db/test/src/triedbmut.rs
@@ -95,7 +95,15 @@ fn playpen_internal<T: TrieLayout>() {
 		let real = reference_trie_root::<T, _, _, _>(x.clone());
 		let mut memdb = PrefixedMemoryDB::<T>::default();
 		let memtrie = populate_trie::<T>(&memdb, &x);
-		let root = memtrie.commit().apply_to(&mut memdb);
+		// avoid duplicate
+		let value_set: std::collections::BTreeMap<&[u8], &[u8]> =
+			x.iter().map(|(k, v)| (k.as_slice(), v.as_slice())).collect();
+		for (k, v) in value_set {
+			assert_eq!(memtrie.get(k).unwrap().unwrap(), v);
+		}
+		let commit = memtrie.commit();
+		let root = commit.apply_to(&mut memdb);
+
 		if root != real {
 			println!("TRIE MISMATCH");
 			println!();
@@ -804,3 +812,46 @@ fn test_insert_remove_data_with_cache_internal<T: TrieLayout>() {
 		assert!(matches!(cache.lookup_value_for_key(key).unwrap(), CachedValue::NonExisting));
 	}
 }
+
+#[test]
+fn test_two_assets_memory_db() {
+	test_two_assets_memory_db_inner_1::<HashedValueNoExtThreshold<1>>();
+	test_two_assets_memory_db_inner_2::<HashedValueNoExtThreshold<1>>();
+}
+fn test_two_assets_memory_db_inner_1<T: TrieLayout>() {
+	let memdb = PrefixedMemoryDB::<T>::new(&[0u8]);
+	let mut state = TrieDBMutBuilder::<T>::new(&memdb).build();
+
+	let key1 = [1u8; 3];
+	let data1 = [1u8; 2];
+	state.insert(key1.as_ref(), &data1).unwrap();
+	assert_eq!(state.get(key1.as_ref()).unwrap().unwrap(), data1); //PASSING
+	let key2 = [2u8; 3];
+	let data2 = [2u8; 2];
+	state.insert(key2.as_ref(), &data2).unwrap();
+	assert_eq!(state.get(key1.as_ref()).unwrap().unwrap(), data1);
+
+	state.commit();
+}
+
+fn test_two_assets_memory_db_inner_2<T: TrieLayout>() {
+	let memdb = PrefixedMemoryDB::<T>::new(&[0u8]);
+	let mut state = TrieDBMutBuilder::<T>::new(&memdb).build();
+
+	let key1 = [1u8];
+	let data1 = [1u8; 2];
+	state.insert(key1.as_ref(), &data1).unwrap();
+	assert_eq!(state.get(key1.as_ref()).unwrap().unwrap(), data1);
+	let key2 = [1u8, 2];
+	let data2 = [2u8; 2];
+	state.insert(key2.as_ref(), &data2).unwrap();
+	assert_eq!(state.get(key1.as_ref()).unwrap().unwrap(), data1);
+	assert_eq!(state.get(key2.as_ref()).unwrap().unwrap(), data2);
+
+	let key3 = [1u8, 3];
+	let data3 = [3u8; 2];
+	state.insert(key3.as_ref(), &data3).unwrap();
+	assert_eq!(state.get(key1.as_ref()).unwrap().unwrap(), data1);
+	assert_eq!(state.get(key2.as_ref()).unwrap().unwrap(), data2);
+	assert_eq!(state.get(key3.as_ref()).unwrap().unwrap(), data3);
+}