Remove partial key size limit from trie codec (paritytech#12566)

* remove size limit from trie codec * test previous upper limit is not enforced anymore * fmt * restore test
gear-tech · Feb 27, 2023 · ddd081e · ddd081e
1 parent f59726d
commit ddd081e
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 9 deletions.
diff --git a/primitives/trie/src/lib.rs b/primitives/trie/src/lib.rs
@@ -527,7 +527,6 @@ where
 /// Constants used into trie simplification codec.
 mod trie_constants {
 	const FIRST_PREFIX: u8 = 0b_00 << 6;
-	pub const NIBBLE_SIZE_BOUND: usize = u16::max_value() as usize;
 	pub const LEAF_PREFIX_MASK: u8 = 0b_01 << 6;
 	pub const BRANCH_WITHOUT_MASK: u8 = 0b_10 << 6;
 	pub const BRANCH_WITH_MASK: u8 = 0b_11 << 6;
@@ -987,6 +986,21 @@ mod tests {
 		assert_eq!(first_storage_root, second_storage_root);
 	}
 
+	#[test]
+	fn big_key() {
+		let check = |keysize: usize| {
+			let mut memdb = PrefixedMemoryDB::<Blake2Hasher>::default();
+			let mut root = Default::default();
+			let mut t = TrieDBMutBuilder::<LayoutV1>::new(&mut memdb, &mut root).build();
+			t.insert(&vec![0x01u8; keysize][..], &[0x01u8, 0x23]).unwrap();
+			std::mem::drop(t);
+			let t = TrieDBBuilder::<LayoutV1>::new(&memdb, &root).build();
+			assert_eq!(t.get(&vec![0x01u8; keysize][..]).unwrap(), Some(vec![0x01u8, 0x23]));
+		};
+		check(u16::MAX as usize / 2); // old limit
+		check(u16::MAX as usize / 2 + 1); // value over old limit still works
+	}
+
 	#[test]
 	fn node_with_no_children_fail_decoding() {
 		let branch = NodeCodec::<Blake2Hasher>::branch_node_nibbled(

diff --git a/primitives/trie/src/node_codec.rs b/primitives/trie/src/node_codec.rs
@@ -279,8 +279,6 @@ fn partial_from_iterator_encode<I: Iterator<Item = u8>>(
 	nibble_count: usize,
 	node_kind: NodeKind,
 ) -> Vec<u8> {
-	let nibble_count = sp_std::cmp::min(trie_constants::NIBBLE_SIZE_BOUND, nibble_count);
-
 	let mut output = Vec::with_capacity(4 + (nibble_count / nibble_ops::NIBBLE_PER_BYTE));
 	match node_kind {
 		NodeKind::Leaf => NodeHeader::Leaf(nibble_count).encode_to(&mut output),

diff --git a/primitives/trie/src/node_header.rs b/primitives/trie/src/node_header.rs
@@ -117,8 +117,6 @@ pub(crate) fn size_and_prefix_iterator(
 	prefix: u8,
 	prefix_mask: usize,
 ) -> impl Iterator<Item = u8> {
-	let size = sp_std::cmp::min(trie_constants::NIBBLE_SIZE_BOUND, size);
-
 	let max_value = 255u8 >> prefix_mask;
 	let l1 = sp_std::cmp::min((max_value as usize).saturating_sub(1), size);
 	let (first_byte, mut rem) = if size == l1 {
@@ -165,12 +163,11 @@ fn decode_size(
 		return Ok(result)
 	}
 	result -= 1;
-	while result <= trie_constants::NIBBLE_SIZE_BOUND {
+	loop {
 		let n = input.read_byte()? as usize;
 		if n < 255 {
 			return Ok(result + n + 1)
 		}
 		result += 255;
 	}
-	Ok(trie_constants::NIBBLE_SIZE_BOUND)
 }
diff --git a/primitives/trie/src/trie_stream.rs b/primitives/trie/src/trie_stream.rs
@@ -54,8 +54,7 @@ fn branch_node_bit_mask(has_children: impl Iterator<Item = bool>) -> (u8, u8) {
 
 /// Create a leaf/branch node, encoding a number of nibbles.
 fn fuse_nibbles_node(nibbles: &[u8], kind: NodeKind) -> impl Iterator<Item = u8> + '_ {
-	let size = sp_std::cmp::min(trie_constants::NIBBLE_SIZE_BOUND, nibbles.len());
-
+	let size = nibbles.len();
 	let iter_start = match kind {
 		NodeKind::Leaf => size_and_prefix_iterator(size, trie_constants::LEAF_PREFIX_MASK, 2),
 		NodeKind::BranchNoValue =>