From d50d5e426d11972a3ca5f004cb2f77ec32812214 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 28 Mar 2024 14:05:06 -0400
Subject: [PATCH 01/55] initial ideas

---
 Cargo.lock                                    |   1 +
 substrate/primitives/runtime/Cargo.toml       |   2 +
 substrate/primitives/runtime/src/lib.rs       |   1 +
 .../primitives/runtime/src/proving_trie.rs    | 111 ++++++++++++++++++
 4 files changed, 115 insertions(+)
 create mode 100644 substrate/primitives/runtime/src/proving_trie.rs

diff --git a/Cargo.lock b/Cargo.lock
index 413bd28abe03..5e61a544ee04 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -18979,6 +18979,7 @@ dependencies = [
  "sp-state-machine",
  "sp-std 14.0.0",
  "sp-tracing 16.0.0",
+ "sp-trie",
  "sp-weights",
  "substrate-test-runtime-client",
  "zstd 0.12.4",
diff --git a/substrate/primitives/runtime/Cargo.toml b/substrate/primitives/runtime/Cargo.toml
index cacfc0597229..778bb7d23e32 100644
--- a/substrate/primitives/runtime/Cargo.toml
+++ b/substrate/primitives/runtime/Cargo.toml
@@ -31,6 +31,7 @@ sp-arithmetic = { path = "../arithmetic", default-features = false }
 sp-core = { path = "../core", default-features = false }
 sp-io = { path = "../io", default-features = false }
 sp-std = { path = "../std", default-features = false }
+sp-trie = { path = "../trie", default-features = false }
 sp-weights = { path = "../weights", default-features = false }
 docify = { version = "0.2.7" }
 
@@ -66,6 +67,7 @@ std = [
 	"sp-state-machine/std",
 	"sp-std/std",
 	"sp-tracing/std",
+	"sp-trie/std",
 	"sp-weights/std",
 ]
 
diff --git a/substrate/primitives/runtime/src/lib.rs b/substrate/primitives/runtime/src/lib.rs
index 44bf3c969e54..455d6e7b9a50 100644
--- a/substrate/primitives/runtime/src/lib.rs
+++ b/substrate/primitives/runtime/src/lib.rs
@@ -85,6 +85,7 @@ pub mod generic;
 pub mod legacy;
 mod multiaddress;
 pub mod offchain;
+pub mod proving_trie;
 pub mod runtime_logger;
 mod runtime_string;
 #[cfg(feature = "std")]
diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
new file mode 100644
index 000000000000..cb9bd97d3eb5
--- /dev/null
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -0,0 +1,111 @@
+use sp_trie::{
+	trie_types::{TrieDBBuilder, TrieDBMutBuilderV0},
+	LayoutV0, MemoryDB, Recorder, Trie, TrieMut, EMPTY_PREFIX,
+};
+
+use crate::{traits::HashOutput, Decode, Encode, KeyTypeId};
+
+/// A trie instance for checking and generating proofs.
+pub struct ProvingTrie<Hashing, Hash, Item>
+where
+	Hashing: sp_core::Hasher<Out = Hash>,
+	Hash: HashOutput,
+	Item: Encode + Decode,
+{
+	db: MemoryDB<Hashing>,
+	root: Hash,
+	_phantom: core::marker::PhantomData<Item>,
+}
+
+impl<Hashing, Hash, Item> ProvingTrie<Hashing, Hash, Item>
+where
+	Hashing: sp_core::Hasher<Out = Hash>,
+	Hash: HashOutput,
+	Item: Encode + Decode,
+{
+	/// Access the underlying trie root.
+	pub fn root(&self) -> &Hash {
+		&self.root
+	}
+
+	// Check a proof contained within the current memory-db. Returns `None` if the
+	// nodes within the current `MemoryDB` are insufficient to query the item.
+	fn query(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Item> {
+		let trie = TrieDBBuilder::new(&self.db, &self.root).build();
+		let val_idx = (key_id, key_data)
+			.using_encoded(|s| trie.get(s))
+			.ok()?
+			.and_then(|raw| u32::decode(&mut &*raw).ok())?;
+
+		val_idx
+			.using_encoded(|s| trie.get(s))
+			.ok()?
+			.and_then(|raw| Item::decode(&mut &*raw).ok())
+	}
+
+	/// Prove the full verification data for a given key and key ID.
+	pub fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>> {
+		let mut recorder = Recorder::<LayoutV0<Hashing>>::new();
+		{
+			let trie =
+				TrieDBBuilder::new(&self.db, &self.root).with_recorder(&mut recorder).build();
+			let val_idx = (key_id, key_data).using_encoded(|s| {
+				trie.get(s).ok()?.and_then(|raw| u32::decode(&mut &*raw).ok())
+			})?;
+
+			val_idx.using_encoded(|s| {
+				trie.get(s).ok()?.and_then(|raw| Item::decode(&mut &*raw).ok())
+			})?;
+		}
+
+		Some(recorder.drain().into_iter().map(|r| r.data).collect())
+	}
+
+	fn from_nodes(root: Hash, nodes: &[Vec<u8>]) -> Self {
+		use sp_trie::HashDBT;
+
+		let mut memory_db = MemoryDB::default();
+		for node in nodes {
+			HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]);
+		}
+
+		ProvingTrie { db: memory_db, root, _phantom: Default::default() }
+	}
+
+	// fn generate_for<I>(items: I) -> Result<Self, &'static str>
+	// where
+	// 	I: IntoIterator<Item = Item>,
+	// {
+	// 	let mut db = MemoryDB::default();
+	// 	let mut root = Default::default();
+
+	// 	{
+	// 		let mut trie = TrieDBMutBuilderV0::new(&mut db, &mut root).build();
+	// 		for (i, (validator, full_id)) in validators.into_iter().enumerate() {
+	// 			let i = i as u32;
+	// 			let keys = match <Session<T>>::load_keys(&validator) {
+	// 				None => continue,
+	// 				Some(k) => k,
+	// 			};
+
+	// 			let full_id = (validator, full_id);
+
+	// 			// map each key to the owner index.
+	// 			for key_id in T::Keys::key_ids() {
+	// 				let key = keys.get_raw(*key_id);
+	// 				let res =
+	// 					(key_id, key).using_encoded(|k| i.using_encoded(|v| trie.insert(k, v)));
+
+	// 				let _ = res.map_err(|_| "failed to insert into trie")?;
+	// 			}
+
+	// 			// map each owner index to the full identification.
+	// 			let _ = i
+	// 				.using_encoded(|k| full_id.using_encoded(|v| trie.insert(k, v)))
+	// 				.map_err(|_| "failed to insert into trie")?;
+	// 		}
+	// 	}
+
+	// 	Ok(ProvingTrie { db, root })
+	// }
+}

From 7f4422e090ca583bbb4ab20c9555e420d081046f Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 28 Mar 2024 14:21:20 -0400
Subject: [PATCH 02/55] Update proving_trie.rs

---
 .../primitives/runtime/src/proving_trie.rs    | 85 ++++++++++---------
 1 file changed, 45 insertions(+), 40 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index cb9bd97d3eb5..b4307f9164e9 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -1,3 +1,22 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! A generic implementation of a simple merkle trie used for making and verifying proofs.
+
 use sp_trie::{
 	trie_types::{TrieDBBuilder, TrieDBMutBuilderV0},
 	LayoutV0, MemoryDB, Recorder, Trie, TrieMut, EMPTY_PREFIX,
@@ -28,9 +47,9 @@ where
 		&self.root
 	}
 
-	// Check a proof contained within the current memory-db. Returns `None` if the
-	// nodes within the current `MemoryDB` are insufficient to query the item.
-	fn query(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Item> {
+	/// Check a proof contained within the current memory-db. Returns `None` if the
+	/// nodes within the current `MemoryDB` are insufficient to query the item.
+	pub fn query(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Item> {
 		let trie = TrieDBBuilder::new(&self.db, &self.root).build();
 		let val_idx = (key_id, key_data)
 			.using_encoded(|s| trie.get(s))
@@ -61,7 +80,8 @@ where
 		Some(recorder.drain().into_iter().map(|r| r.data).collect())
 	}
 
-	fn from_nodes(root: Hash, nodes: &[Vec<u8>]) -> Self {
+	/// Create a new instance of a `ProvingTrie` using a set of raw nodes.
+	pub fn from_nodes(root: Hash, nodes: &[Vec<u8>]) -> Self {
 		use sp_trie::HashDBT;
 
 		let mut memory_db = MemoryDB::default();
@@ -72,40 +92,25 @@ where
 		ProvingTrie { db: memory_db, root, _phantom: Default::default() }
 	}
 
-	// fn generate_for<I>(items: I) -> Result<Self, &'static str>
-	// where
-	// 	I: IntoIterator<Item = Item>,
-	// {
-	// 	let mut db = MemoryDB::default();
-	// 	let mut root = Default::default();
-
-	// 	{
-	// 		let mut trie = TrieDBMutBuilderV0::new(&mut db, &mut root).build();
-	// 		for (i, (validator, full_id)) in validators.into_iter().enumerate() {
-	// 			let i = i as u32;
-	// 			let keys = match <Session<T>>::load_keys(&validator) {
-	// 				None => continue,
-	// 				Some(k) => k,
-	// 			};
-
-	// 			let full_id = (validator, full_id);
-
-	// 			// map each key to the owner index.
-	// 			for key_id in T::Keys::key_ids() {
-	// 				let key = keys.get_raw(*key_id);
-	// 				let res =
-	// 					(key_id, key).using_encoded(|k| i.using_encoded(|v| trie.insert(k, v)));
-
-	// 				let _ = res.map_err(|_| "failed to insert into trie")?;
-	// 			}
-
-	// 			// map each owner index to the full identification.
-	// 			let _ = i
-	// 				.using_encoded(|k| full_id.using_encoded(|v| trie.insert(k, v)))
-	// 				.map_err(|_| "failed to insert into trie")?;
-	// 		}
-	// 	}
-
-	// 	Ok(ProvingTrie { db, root })
-	// }
+	/// Create a new instance of a `ProvingTrie` using an iterator of items in the trie.
+	pub fn generate_for<I>(items: I) -> Result<Self, &'static str>
+	where
+		I: IntoIterator<Item = Item>,
+	{
+		let mut db = MemoryDB::default();
+		let mut root = Default::default();
+
+		{
+			let mut trie = TrieDBMutBuilderV0::new(&mut db, &mut root).build();
+			for (i, item) in items.into_iter().enumerate() {
+				let i = i as u32;
+
+				// insert each item into the trie
+				i.using_encoded(|k| item.using_encoded(|v| trie.insert(k, v)))
+					.map_err(|_| "failed to insert into trie")?;
+			}
+		}
+
+		Ok(ProvingTrie { db, root, _phantom: Default::default() })
+	}
 }

From bdc0c8422ae02d42df9ab0cf4cbdf820925348e8 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 28 Mar 2024 14:31:47 -0400
Subject: [PATCH 03/55] create trait

---
 .../primitives/runtime/src/proving_trie.rs    | 35 +++++++++++++++----
 1 file changed, 28 insertions(+), 7 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index b4307f9164e9..bd24a25eb65c 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -15,17 +15,38 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//! A generic implementation of a simple merkle trie used for making and verifying proofs.
+//! Types for a simple merkle trie used for checking and generating proofs.
 
 use sp_trie::{
 	trie_types::{TrieDBBuilder, TrieDBMutBuilderV0},
 	LayoutV0, MemoryDB, Recorder, Trie, TrieMut, EMPTY_PREFIX,
 };
 
-use crate::{traits::HashOutput, Decode, Encode, KeyTypeId};
+use crate::{traits::HashOutput, Decode, DispatchError, Encode, KeyTypeId};
+
+/// A trait for creating a merkle trie for checking and generating merkle proofs.
+pub trait ProvingTrie<Hashing, Hash, Item>
+where
+	Hashing: sp_core::Hasher<Out = Hash>,
+	Hash: HashOutput,
+	Item: Encode + Decode,
+	Self: Sized,
+{
+	/// Access the underlying trie root.
+	fn root(&self) -> &Hash;
+	/// Check a proof contained within the current memory-db. Returns `None` if the
+	/// nodes within the current `MemoryDB` are insufficient to query the item.
+	fn query(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Item>;
+	/// Prove the full verification data for a given key and key ID.
+	fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>>;
+	/// Create a new instance of a `ProvingTrie` using an iterator of items in the trie.
+	fn generate<I>(items: I) -> Result<Self, DispatchError>
+	where
+		I: IntoIterator<Item = Item>;
+}
 
 /// A trie instance for checking and generating proofs.
-pub struct ProvingTrie<Hashing, Hash, Item>
+pub struct BasicProvingTrie<Hashing, Hash, Item>
 where
 	Hashing: sp_core::Hasher<Out = Hash>,
 	Hash: HashOutput,
@@ -36,7 +57,7 @@ where
 	_phantom: core::marker::PhantomData<Item>,
 }
 
-impl<Hashing, Hash, Item> ProvingTrie<Hashing, Hash, Item>
+impl<Hashing, Hash, Item> BasicProvingTrie<Hashing, Hash, Item>
 where
 	Hashing: sp_core::Hasher<Out = Hash>,
 	Hash: HashOutput,
@@ -89,11 +110,11 @@ where
 			HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]);
 		}
 
-		ProvingTrie { db: memory_db, root, _phantom: Default::default() }
+		Self { db: memory_db, root, _phantom: Default::default() }
 	}
 
 	/// Create a new instance of a `ProvingTrie` using an iterator of items in the trie.
-	pub fn generate_for<I>(items: I) -> Result<Self, &'static str>
+	pub fn generate_for<I>(items: I) -> Result<Self, DispatchError>
 	where
 		I: IntoIterator<Item = Item>,
 	{
@@ -111,6 +132,6 @@ where
 			}
 		}
 
-		Ok(ProvingTrie { db, root, _phantom: Default::default() })
+		Ok(Self { db, root, _phantom: Default::default() })
 	}
 }

From da385ab4390ce36dee639014ddf0c6e0ead25f8a Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 28 Mar 2024 14:38:05 -0400
Subject: [PATCH 04/55] use trait

---
 .../primitives/runtime/src/proving_trie.rs    | 41 +++++++++++--------
 1 file changed, 24 insertions(+), 17 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index bd24a25eb65c..346a425d279b 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -57,20 +57,20 @@ where
 	_phantom: core::marker::PhantomData<Item>,
 }
 
-impl<Hashing, Hash, Item> BasicProvingTrie<Hashing, Hash, Item>
+impl<Hashing, Hash, Item> ProvingTrie<Hashing, Hash, Item> for BasicProvingTrie<Hashing, Hash, Item>
 where
 	Hashing: sp_core::Hasher<Out = Hash>,
 	Hash: HashOutput,
 	Item: Encode + Decode,
 {
 	/// Access the underlying trie root.
-	pub fn root(&self) -> &Hash {
+	fn root(&self) -> &Hash {
 		&self.root
 	}
 
 	/// Check a proof contained within the current memory-db. Returns `None` if the
 	/// nodes within the current `MemoryDB` are insufficient to query the item.
-	pub fn query(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Item> {
+	fn query(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Item> {
 		let trie = TrieDBBuilder::new(&self.db, &self.root).build();
 		let val_idx = (key_id, key_data)
 			.using_encoded(|s| trie.get(s))
@@ -84,7 +84,7 @@ where
 	}
 
 	/// Prove the full verification data for a given key and key ID.
-	pub fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>> {
+	fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>> {
 		let mut recorder = Recorder::<LayoutV0<Hashing>>::new();
 		{
 			let trie =
@@ -101,20 +101,8 @@ where
 		Some(recorder.drain().into_iter().map(|r| r.data).collect())
 	}
 
-	/// Create a new instance of a `ProvingTrie` using a set of raw nodes.
-	pub fn from_nodes(root: Hash, nodes: &[Vec<u8>]) -> Self {
-		use sp_trie::HashDBT;
-
-		let mut memory_db = MemoryDB::default();
-		for node in nodes {
-			HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]);
-		}
-
-		Self { db: memory_db, root, _phantom: Default::default() }
-	}
-
 	/// Create a new instance of a `ProvingTrie` using an iterator of items in the trie.
-	pub fn generate_for<I>(items: I) -> Result<Self, DispatchError>
+	fn generate<I>(items: I) -> Result<Self, DispatchError>
 	where
 		I: IntoIterator<Item = Item>,
 	{
@@ -135,3 +123,22 @@ where
 		Ok(Self { db, root, _phantom: Default::default() })
 	}
 }
+
+impl<Hashing, Hash, Item> BasicProvingTrie<Hashing, Hash, Item>
+where
+	Hashing: sp_core::Hasher<Out = Hash>,
+	Hash: HashOutput,
+	Item: Encode + Decode,
+{
+	/// Create a new instance of a `ProvingTrie` using a set of raw nodes.
+	pub fn from_nodes(root: Hash, nodes: &[Vec<u8>]) -> Self {
+		use sp_trie::HashDBT;
+
+		let mut memory_db = MemoryDB::default();
+		for node in nodes {
+			HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]);
+		}
+
+		Self { db: memory_db, root, _phantom: Default::default() }
+	}
+}

From 6ab4bfbc6aa46a10cfefc5bdc5ebb856313d2728 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 28 Mar 2024 15:09:45 -0400
Subject: [PATCH 05/55] clean up trait and basic trie further

---
 .../primitives/runtime/src/proving_trie.rs    | 34 +++----------------
 1 file changed, 5 insertions(+), 29 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index 346a425d279b..9974ca41100a 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -19,17 +19,14 @@
 
 use sp_trie::{
 	trie_types::{TrieDBBuilder, TrieDBMutBuilderV0},
-	LayoutV0, MemoryDB, Recorder, Trie, TrieMut, EMPTY_PREFIX,
+	LayoutV0, MemoryDB, Recorder, Trie, TrieMut,
 };
 
-use crate::{traits::HashOutput, Decode, DispatchError, Encode, KeyTypeId};
+use crate::{Decode, DispatchError, Encode, KeyTypeId};
 
 /// A trait for creating a merkle trie for checking and generating merkle proofs.
 pub trait ProvingTrie<Hashing, Hash, Item>
 where
-	Hashing: sp_core::Hasher<Out = Hash>,
-	Hash: HashOutput,
-	Item: Encode + Decode,
 	Self: Sized,
 {
 	/// Access the underlying trie root.
@@ -40,7 +37,7 @@ where
 	/// Prove the full verification data for a given key and key ID.
 	fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>>;
 	/// Create a new instance of a `ProvingTrie` using an iterator of items in the trie.
-	fn generate<I>(items: I) -> Result<Self, DispatchError>
+	fn generate_for<I>(items: I) -> Result<Self, DispatchError>
 	where
 		I: IntoIterator<Item = Item>;
 }
@@ -49,8 +46,6 @@ where
 pub struct BasicProvingTrie<Hashing, Hash, Item>
 where
 	Hashing: sp_core::Hasher<Out = Hash>,
-	Hash: HashOutput,
-	Item: Encode + Decode,
 {
 	db: MemoryDB<Hashing>,
 	root: Hash,
@@ -60,7 +55,7 @@ where
 impl<Hashing, Hash, Item> ProvingTrie<Hashing, Hash, Item> for BasicProvingTrie<Hashing, Hash, Item>
 where
 	Hashing: sp_core::Hasher<Out = Hash>,
-	Hash: HashOutput,
+	Hash: Default,
 	Item: Encode + Decode,
 {
 	/// Access the underlying trie root.
@@ -102,7 +97,7 @@ where
 	}
 
 	/// Create a new instance of a `ProvingTrie` using an iterator of items in the trie.
-	fn generate<I>(items: I) -> Result<Self, DispatchError>
+	fn generate_for<I>(items: I) -> Result<Self, DispatchError>
 	where
 		I: IntoIterator<Item = Item>,
 	{
@@ -123,22 +118,3 @@ where
 		Ok(Self { db, root, _phantom: Default::default() })
 	}
 }
-
-impl<Hashing, Hash, Item> BasicProvingTrie<Hashing, Hash, Item>
-where
-	Hashing: sp_core::Hasher<Out = Hash>,
-	Hash: HashOutput,
-	Item: Encode + Decode,
-{
-	/// Create a new instance of a `ProvingTrie` using a set of raw nodes.
-	pub fn from_nodes(root: Hash, nodes: &[Vec<u8>]) -> Self {
-		use sp_trie::HashDBT;
-
-		let mut memory_db = MemoryDB::default();
-		for node in nodes {
-			HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]);
-		}
-
-		Self { db: memory_db, root, _phantom: Default::default() }
-	}
-}

From 1cfb29f69dd02076c1ba051786a36096e354e6c1 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 28 Mar 2024 15:09:58 -0400
Subject: [PATCH 06/55] use trait in session historical

---
 substrate/frame/session/src/historical/mod.rs | 47 ++++++++++---------
 .../frame/session/src/historical/offchain.rs  | 11 +++--
 2 files changed, 32 insertions(+), 26 deletions(-)

diff --git a/substrate/frame/session/src/historical/mod.rs b/substrate/frame/session/src/historical/mod.rs
index b9cecea1a7f7..3c30a3126660 100644
--- a/substrate/frame/session/src/historical/mod.rs
+++ b/substrate/frame/session/src/historical/mod.rs
@@ -32,8 +32,9 @@ mod shared;
 
 use codec::{Decode, Encode};
 use sp_runtime::{
+	proving_trie::ProvingTrie,
 	traits::{Convert, OpaqueKeys},
-	KeyTypeId,
+	DispatchError, KeyTypeId,
 };
 use sp_session::{MembershipProof, ValidatorCount};
 use sp_staking::SessionIndex;
@@ -176,7 +177,7 @@ impl<T: Config, I: SessionManager<T::ValidatorId, T::FullIdentification>> NoteHi
 
 		if let Some(new_validators) = new_validators_and_id {
 			let count = new_validators.len() as ValidatorCount;
-			match ProvingTrie::<T>::generate_for(new_validators) {
+			match ValidatorProvingTrie::<T>::generate_for(new_validators) {
 				Ok(trie) => <HistoricalSessions<T>>::insert(new_index, &(trie.root, count)),
 				Err(reason) => {
 					print("Failed to generate historical ancestry-inclusion proof.");
@@ -221,13 +222,15 @@ pub type IdentificationTuple<T> =
 	(<T as pallet_session::Config>::ValidatorId, <T as Config>::FullIdentification);
 
 /// A trie instance for checking and generating proofs.
-pub struct ProvingTrie<T: Config> {
+pub struct ValidatorProvingTrie<T: Config> {
 	db: MemoryDB<T::Hashing>,
 	root: T::Hash,
 }
 
-impl<T: Config> ProvingTrie<T> {
-	fn generate_for<I>(validators: I) -> Result<Self, &'static str>
+impl<T: Config> ProvingTrie<T::Hashing, T::Hash, IdentificationTuple<T>>
+	for ValidatorProvingTrie<T>
+{
+	fn generate_for<I>(validators: I) -> Result<Self, DispatchError>
 	where
 		I: IntoIterator<Item = (T::ValidatorId, T::FullIdentification)>,
 	{
@@ -261,22 +264,11 @@ impl<T: Config> ProvingTrie<T> {
 			}
 		}
 
-		Ok(ProvingTrie { db, root })
-	}
-
-	fn from_nodes(root: T::Hash, nodes: &[Vec<u8>]) -> Self {
-		use sp_trie::HashDBT;
-
-		let mut memory_db = MemoryDB::default();
-		for node in nodes {
-			HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]);
-		}
-
-		ProvingTrie { db: memory_db, root }
+		Ok(Self { db, root })
 	}
 
 	/// Prove the full verification data for a given key and key ID.
-	pub fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>> {
+	fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>> {
 		let mut recorder = Recorder::<LayoutV0<T::Hashing>>::new();
 		{
 			let trie =
@@ -296,7 +288,7 @@ impl<T: Config> ProvingTrie<T> {
 	}
 
 	/// Access the underlying trie root.
-	pub fn root(&self) -> &T::Hash {
+	fn root(&self) -> &T::Hash {
 		&self.root
 	}
 
@@ -316,6 +308,19 @@ impl<T: Config> ProvingTrie<T> {
 	}
 }
 
+impl<T: Config> ValidatorProvingTrie<T> {
+	fn from_nodes(root: T::Hash, nodes: &[Vec<u8>]) -> Self {
+		use sp_trie::HashDBT;
+
+		let mut memory_db = MemoryDB::default();
+		for node in nodes {
+			HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]);
+		}
+
+		Self { db: memory_db, root }
+	}
+}
+
 impl<T: Config, D: AsRef<[u8]>> KeyOwnerProofSystem<(KeyTypeId, D)> for Pallet<T> {
 	type Proof = MembershipProof;
 	type IdentificationTuple = IdentificationTuple<T>;
@@ -332,7 +337,7 @@ impl<T: Config, D: AsRef<[u8]>> KeyOwnerProofSystem<(KeyTypeId, D)> for Pallet<T
 
 		let count = validators.len() as ValidatorCount;
 
-		let trie = ProvingTrie::<T>::generate_for(validators).ok()?;
+		let trie = ValidatorProvingTrie::<T>::generate_for(validators).ok()?;
 
 		let (id, data) = key;
 		trie.prove(id, data.as_ref()).map(|trie_nodes| MembershipProof {
@@ -364,7 +369,7 @@ impl<T: Config, D: AsRef<[u8]>> KeyOwnerProofSystem<(KeyTypeId, D)> for Pallet<T
 				return None
 			}
 
-			let trie = ProvingTrie::<T>::from_nodes(root, &proof.trie_nodes);
+			let trie = ValidatorProvingTrie::<T>::from_nodes(root, &proof.trie_nodes);
 			trie.query(id, data.as_ref())
 		}
 	}
diff --git a/substrate/frame/session/src/historical/offchain.rs b/substrate/frame/session/src/historical/offchain.rs
index 95f4d762949e..5a5a1114b9bf 100644
--- a/substrate/frame/session/src/historical/offchain.rs
+++ b/substrate/frame/session/src/historical/offchain.rs
@@ -20,17 +20,18 @@
 //! Validator-set extracting an iterator from an off-chain worker stored list containing historical
 //! validator-sets. Based on the logic of historical slashing, but the validation is done off-chain.
 //! Use [`fn store_current_session_validator_set_to_offchain()`](super::onchain) to store the
-//! required data to the offchain validator set. This is used in conjunction with [`ProvingTrie`]
-//! and the off-chain indexing API.
+//! required data to the offchain validator set. This is used in conjunction with
+//! [`ValidatorProvingTrie`] and the off-chain indexing API.
 
 use sp_runtime::{
 	offchain::storage::{MutateStorageError, StorageRetrievalError, StorageValueRef},
+	proving_trie::ProvingTrie,
 	KeyTypeId,
 };
 use sp_session::MembershipProof;
 use sp_std::prelude::*;
 
-use super::{shared, Config, IdentificationTuple, ProvingTrie};
+use super::{shared, Config, IdentificationTuple, ValidatorProvingTrie};
 use crate::{Pallet as SessionModule, SessionIndex};
 
 /// A set of validators, which was used for a fixed session index.
@@ -59,7 +60,7 @@ impl<T: Config> ValidatorSet<T> {
 }
 
 /// Implement conversion into iterator for usage
-/// with [ProvingTrie](super::ProvingTrie::generate_for).
+/// with [ValidatorProvingTrie](super::ValidatorProvingTrie::generate_for).
 impl<T: Config> sp_std::iter::IntoIterator for ValidatorSet<T> {
 	type Item = (T::ValidatorId, T::FullIdentification);
 	type IntoIter = sp_std::vec::IntoIter<Self::Item>;
@@ -79,7 +80,7 @@ pub fn prove_session_membership<T: Config, D: AsRef<[u8]>>(
 ) -> Option<MembershipProof> {
 	let validators = ValidatorSet::<T>::load_from_offchain_db(session_index)?;
 	let count = validators.len() as u32;
-	let trie = ProvingTrie::<T>::generate_for(validators.into_iter()).ok()?;
+	let trie = ValidatorProvingTrie::<T>::generate_for(validators.into_iter()).ok()?;
 
 	let (id, data) = session_key;
 	trie.prove(id, data.as_ref()).map(|trie_nodes| MembershipProof {

From 3080c7b3bd5ec3860da49eaa64bc9b45c261ed4d Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 28 Mar 2024 17:44:42 -0400
Subject: [PATCH 07/55] fix api, add basic end to end test

---
 .../primitives/runtime/src/proving_trie.rs    | 162 ++++++++++++------
 1 file changed, 110 insertions(+), 52 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index 9974ca41100a..520d6379a0b2 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -17,104 +17,162 @@
 
 //! Types for a simple merkle trie used for checking and generating proofs.
 
+use crate::{Decode, DispatchError, Encode};
+
+use sp_std::vec::Vec;
 use sp_trie::{
 	trie_types::{TrieDBBuilder, TrieDBMutBuilderV0},
-	LayoutV0, MemoryDB, Recorder, Trie, TrieMut,
+	LayoutV0, MemoryDB, Recorder, Trie, TrieMut, EMPTY_PREFIX,
 };
 
-use crate::{Decode, DispatchError, Encode, KeyTypeId};
-
 /// A trait for creating a merkle trie for checking and generating merkle proofs.
-pub trait ProvingTrie<Hashing, Hash, Item>
+pub trait ProvingTrie<Hashing, Hash, Key, Value>
 where
 	Self: Sized,
 {
+	/// Create a new instance of a `ProvingTrie` using an iterator of key/value pairs.
+	fn generate_for<I>(items: I) -> Result<Self, DispatchError>
+	where
+		I: IntoIterator<Item = (Key, Value)>;
 	/// Access the underlying trie root.
 	fn root(&self) -> &Hash;
 	/// Check a proof contained within the current memory-db. Returns `None` if the
 	/// nodes within the current `MemoryDB` are insufficient to query the item.
-	fn query(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Item>;
-	/// Prove the full verification data for a given key and key ID.
-	fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>>;
-	/// Create a new instance of a `ProvingTrie` using an iterator of items in the trie.
-	fn generate_for<I>(items: I) -> Result<Self, DispatchError>
-	where
-		I: IntoIterator<Item = Item>;
+	fn query(&self, key: Key) -> Option<Value>;
+	/// Create the full verification data needed to prove a key and its value in the trie. Returns
+	/// `None` if the nodes nodes within the current `MemoryDB` are insufficient to create a proof.
+	fn create_proof(&self, key: Key) -> Option<Vec<Vec<u8>>>;
+	/// Create a new instance of `ProvingTrie` from raw nodes. Nodes can be generated using the
+	/// `create_proof` function.
+	fn from_nodes(root: Hash, nodes: &[Vec<u8>]) -> Self;
 }
 
-/// A trie instance for checking and generating proofs.
-pub struct BasicProvingTrie<Hashing, Hash, Item>
+/// A basic trie implementation for checking and generating proofs for a key / value pair.
+pub struct BasicProvingTrie<Hashing, Hash, Key, Value>
 where
 	Hashing: sp_core::Hasher<Out = Hash>,
 {
 	db: MemoryDB<Hashing>,
 	root: Hash,
-	_phantom: core::marker::PhantomData<Item>,
+	_phantom: core::marker::PhantomData<(Key, Value)>,
 }
 
-impl<Hashing, Hash, Item> ProvingTrie<Hashing, Hash, Item> for BasicProvingTrie<Hashing, Hash, Item>
+impl<Hashing, Hash, Key, Value> ProvingTrie<Hashing, Hash, Key, Value>
+	for BasicProvingTrie<Hashing, Hash, Key, Value>
 where
 	Hashing: sp_core::Hasher<Out = Hash>,
-	Hash: Default,
-	Item: Encode + Decode,
+	Hash: Default + Send + Sync,
+	Key: Encode,
+	Value: Encode + Decode,
 {
-	/// Access the underlying trie root.
+	fn generate_for<I>(items: I) -> Result<Self, DispatchError>
+	where
+		I: IntoIterator<Item = (Key, Value)>,
+	{
+		let mut db = MemoryDB::default();
+		let mut root = Default::default();
+
+		{
+			let mut trie = TrieDBMutBuilderV0::new(&mut db, &mut root).build();
+			for (key, value) in items.into_iter() {
+				key.using_encoded(|k| value.using_encoded(|v| trie.insert(k, v)))
+					.map_err(|_| "failed to insert into trie")?;
+			}
+		}
+
+		Ok(Self { db, root, _phantom: Default::default() })
+	}
+
 	fn root(&self) -> &Hash {
 		&self.root
 	}
 
-	/// Check a proof contained within the current memory-db. Returns `None` if the
-	/// nodes within the current `MemoryDB` are insufficient to query the item.
-	fn query(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Item> {
+	fn query(&self, key: Key) -> Option<Value> {
 		let trie = TrieDBBuilder::new(&self.db, &self.root).build();
-		let val_idx = (key_id, key_data)
-			.using_encoded(|s| trie.get(s))
+		key.using_encoded(|s| trie.get(s))
 			.ok()?
-			.and_then(|raw| u32::decode(&mut &*raw).ok())?;
-
-		val_idx
-			.using_encoded(|s| trie.get(s))
-			.ok()?
-			.and_then(|raw| Item::decode(&mut &*raw).ok())
+			.and_then(|raw| Value::decode(&mut &*raw).ok())
 	}
 
-	/// Prove the full verification data for a given key and key ID.
-	fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>> {
+	fn create_proof(&self, key: Key) -> Option<Vec<Vec<u8>>> {
 		let mut recorder = Recorder::<LayoutV0<Hashing>>::new();
+
 		{
 			let trie =
 				TrieDBBuilder::new(&self.db, &self.root).with_recorder(&mut recorder).build();
-			let val_idx = (key_id, key_data).using_encoded(|s| {
-				trie.get(s).ok()?.and_then(|raw| u32::decode(&mut &*raw).ok())
-			})?;
 
-			val_idx.using_encoded(|s| {
-				trie.get(s).ok()?.and_then(|raw| Item::decode(&mut &*raw).ok())
+			key.using_encoded(|k| {
+				trie.get(k).ok()?.and_then(|raw| Value::decode(&mut &*raw).ok())
 			})?;
 		}
 
 		Some(recorder.drain().into_iter().map(|r| r.data).collect())
 	}
 
-	/// Create a new instance of a `ProvingTrie` using an iterator of items in the trie.
-	fn generate_for<I>(items: I) -> Result<Self, DispatchError>
-	where
-		I: IntoIterator<Item = Item>,
-	{
-		let mut db = MemoryDB::default();
-		let mut root = Default::default();
+	fn from_nodes(root: Hash, nodes: &[Vec<u8>]) -> Self {
+		use sp_trie::HashDBT;
 
-		{
-			let mut trie = TrieDBMutBuilderV0::new(&mut db, &mut root).build();
-			for (i, item) in items.into_iter().enumerate() {
-				let i = i as u32;
+		let mut memory_db = MemoryDB::default();
+		for node in nodes {
+			HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]);
+		}
 
-				// insert each item into the trie
-				i.using_encoded(|k| item.using_encoded(|v| trie.insert(k, v)))
-					.map_err(|_| "failed to insert into trie")?;
-			}
+		Self { db: memory_db, root, _phantom: Default::default() }
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use crate::traits::BlakeTwo256;
+	use sp_core::H256;
+	use sp_std::{collections::btree_map::BTreeMap, str::FromStr};
+
+	// A trie which simulates a trie of accounts (u32) and balances (u128).
+	type BalanceTrie = BasicProvingTrie<BlakeTwo256, H256, u32, u128>;
+
+	// The expected root hash for an empty trie.
+	fn empty_root() -> H256 {
+		H256::from_str("0x03170a2e7597b7b7e3d84c05391d139a62b157e78786d8c082f29dcf4c111314")
+			.unwrap()
+	}
+
+	#[test]
+	fn empty_trie_works() {
+		let empty_trie = BalanceTrie::generate_for(Vec::new()).unwrap();
+		assert_eq!(*empty_trie.root(), empty_root());
+	}
+
+	#[test]
+	fn basic_end_to_end() {
+		// Create a map of users and their balances.
+		let mut map = BTreeMap::<u32, u128>::new();
+		for i in 0..10u32 {
+			map.insert(i, i.into());
 		}
 
-		Ok(Self { db, root, _phantom: Default::default() })
+		// Put items into the trie.
+		let balance_trie = BalanceTrie::generate_for(map).unwrap();
+
+		// Root is changed.
+		let root = *balance_trie.root();
+		assert!(root != empty_root());
+
+		// Assert valid key is queryable.
+		assert_eq!(balance_trie.query(6u32), Some(6u128));
+		assert_eq!(balance_trie.query(9u32), Some(9u128));
+		// Invalid key returns none.
+		assert_eq!(balance_trie.query(69u32), None);
+
+		// Create a proof for a valid key.
+		let proof = balance_trie.prove(6u32).unwrap();
+		// Can't create proof for invalid key.
+		assert_eq!(balance_trie.prove(69u32), None);
+
+		// Create a new proving trie from the proof.
+		let new_balance_trie = BalanceTrie::from_nodes(root, &proof);
+
+		// Assert valid key is queryable.
+		assert_eq!(new_balance_trie.query(6u32), Some(6u128));
 	}
 }

From 39b0fca3553e4c49ce49548d39339a825cb1e2a8 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 28 Mar 2024 17:46:31 -0400
Subject: [PATCH 08/55] fix test

---
 substrate/primitives/runtime/src/proving_trie.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index 520d6379a0b2..3e72e89bc229 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -165,9 +165,9 @@ mod tests {
 		assert_eq!(balance_trie.query(69u32), None);
 
 		// Create a proof for a valid key.
-		let proof = balance_trie.prove(6u32).unwrap();
+		let proof = balance_trie.create_proof(6u32).unwrap();
 		// Can't create proof for invalid key.
-		assert_eq!(balance_trie.prove(69u32), None);
+		assert_eq!(balance_trie.create_proof(69u32), None);
 
 		// Create a new proving trie from the proof.
 		let new_balance_trie = BalanceTrie::from_nodes(root, &proof);

From 599f5761d552a1575e39a8ccf4abb279fa418959 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 28 Mar 2024 17:56:23 -0400
Subject: [PATCH 09/55] Revert "use trait in session historical"

This reverts commit 1cfb29f69dd02076c1ba051786a36096e354e6c1.
---
 substrate/frame/session/src/historical/mod.rs | 47 +++++++++----------
 .../frame/session/src/historical/offchain.rs  | 11 ++---
 2 files changed, 26 insertions(+), 32 deletions(-)

diff --git a/substrate/frame/session/src/historical/mod.rs b/substrate/frame/session/src/historical/mod.rs
index 3c30a3126660..b9cecea1a7f7 100644
--- a/substrate/frame/session/src/historical/mod.rs
+++ b/substrate/frame/session/src/historical/mod.rs
@@ -32,9 +32,8 @@ mod shared;
 
 use codec::{Decode, Encode};
 use sp_runtime::{
-	proving_trie::ProvingTrie,
 	traits::{Convert, OpaqueKeys},
-	DispatchError, KeyTypeId,
+	KeyTypeId,
 };
 use sp_session::{MembershipProof, ValidatorCount};
 use sp_staking::SessionIndex;
@@ -177,7 +176,7 @@ impl<T: Config, I: SessionManager<T::ValidatorId, T::FullIdentification>> NoteHi
 
 		if let Some(new_validators) = new_validators_and_id {
 			let count = new_validators.len() as ValidatorCount;
-			match ValidatorProvingTrie::<T>::generate_for(new_validators) {
+			match ProvingTrie::<T>::generate_for(new_validators) {
 				Ok(trie) => <HistoricalSessions<T>>::insert(new_index, &(trie.root, count)),
 				Err(reason) => {
 					print("Failed to generate historical ancestry-inclusion proof.");
@@ -222,15 +221,13 @@ pub type IdentificationTuple<T> =
 	(<T as pallet_session::Config>::ValidatorId, <T as Config>::FullIdentification);
 
 /// A trie instance for checking and generating proofs.
-pub struct ValidatorProvingTrie<T: Config> {
+pub struct ProvingTrie<T: Config> {
 	db: MemoryDB<T::Hashing>,
 	root: T::Hash,
 }
 
-impl<T: Config> ProvingTrie<T::Hashing, T::Hash, IdentificationTuple<T>>
-	for ValidatorProvingTrie<T>
-{
-	fn generate_for<I>(validators: I) -> Result<Self, DispatchError>
+impl<T: Config> ProvingTrie<T> {
+	fn generate_for<I>(validators: I) -> Result<Self, &'static str>
 	where
 		I: IntoIterator<Item = (T::ValidatorId, T::FullIdentification)>,
 	{
@@ -264,11 +261,22 @@ impl<T: Config> ProvingTrie<T::Hashing, T::Hash, IdentificationTuple<T>>
 			}
 		}
 
-		Ok(Self { db, root })
+		Ok(ProvingTrie { db, root })
+	}
+
+	fn from_nodes(root: T::Hash, nodes: &[Vec<u8>]) -> Self {
+		use sp_trie::HashDBT;
+
+		let mut memory_db = MemoryDB::default();
+		for node in nodes {
+			HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]);
+		}
+
+		ProvingTrie { db: memory_db, root }
 	}
 
 	/// Prove the full verification data for a given key and key ID.
-	fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>> {
+	pub fn prove(&self, key_id: KeyTypeId, key_data: &[u8]) -> Option<Vec<Vec<u8>>> {
 		let mut recorder = Recorder::<LayoutV0<T::Hashing>>::new();
 		{
 			let trie =
@@ -288,7 +296,7 @@ impl<T: Config> ProvingTrie<T::Hashing, T::Hash, IdentificationTuple<T>>
 	}
 
 	/// Access the underlying trie root.
-	fn root(&self) -> &T::Hash {
+	pub fn root(&self) -> &T::Hash {
 		&self.root
 	}
 
@@ -308,19 +316,6 @@ impl<T: Config> ProvingTrie<T::Hashing, T::Hash, IdentificationTuple<T>>
 	}
 }
 
-impl<T: Config> ValidatorProvingTrie<T> {
-	fn from_nodes(root: T::Hash, nodes: &[Vec<u8>]) -> Self {
-		use sp_trie::HashDBT;
-
-		let mut memory_db = MemoryDB::default();
-		for node in nodes {
-			HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]);
-		}
-
-		Self { db: memory_db, root }
-	}
-}
-
 impl<T: Config, D: AsRef<[u8]>> KeyOwnerProofSystem<(KeyTypeId, D)> for Pallet<T> {
 	type Proof = MembershipProof;
 	type IdentificationTuple = IdentificationTuple<T>;
@@ -337,7 +332,7 @@ impl<T: Config, D: AsRef<[u8]>> KeyOwnerProofSystem<(KeyTypeId, D)> for Pallet<T
 
 		let count = validators.len() as ValidatorCount;
 
-		let trie = ValidatorProvingTrie::<T>::generate_for(validators).ok()?;
+		let trie = ProvingTrie::<T>::generate_for(validators).ok()?;
 
 		let (id, data) = key;
 		trie.prove(id, data.as_ref()).map(|trie_nodes| MembershipProof {
@@ -369,7 +364,7 @@ impl<T: Config, D: AsRef<[u8]>> KeyOwnerProofSystem<(KeyTypeId, D)> for Pallet<T
 				return None
 			}
 
-			let trie = ValidatorProvingTrie::<T>::from_nodes(root, &proof.trie_nodes);
+			let trie = ProvingTrie::<T>::from_nodes(root, &proof.trie_nodes);
 			trie.query(id, data.as_ref())
 		}
 	}
diff --git a/substrate/frame/session/src/historical/offchain.rs b/substrate/frame/session/src/historical/offchain.rs
index 5a5a1114b9bf..95f4d762949e 100644
--- a/substrate/frame/session/src/historical/offchain.rs
+++ b/substrate/frame/session/src/historical/offchain.rs
@@ -20,18 +20,17 @@
 //! Validator-set extracting an iterator from an off-chain worker stored list containing historical
 //! validator-sets. Based on the logic of historical slashing, but the validation is done off-chain.
 //! Use [`fn store_current_session_validator_set_to_offchain()`](super::onchain) to store the
-//! required data to the offchain validator set. This is used in conjunction with
-//! [`ValidatorProvingTrie`] and the off-chain indexing API.
+//! required data to the offchain validator set. This is used in conjunction with [`ProvingTrie`]
+//! and the off-chain indexing API.
 
 use sp_runtime::{
 	offchain::storage::{MutateStorageError, StorageRetrievalError, StorageValueRef},
-	proving_trie::ProvingTrie,
 	KeyTypeId,
 };
 use sp_session::MembershipProof;
 use sp_std::prelude::*;
 
-use super::{shared, Config, IdentificationTuple, ValidatorProvingTrie};
+use super::{shared, Config, IdentificationTuple, ProvingTrie};
 use crate::{Pallet as SessionModule, SessionIndex};
 
 /// A set of validators, which was used for a fixed session index.
@@ -60,7 +59,7 @@ impl<T: Config> ValidatorSet<T> {
 }
 
 /// Implement conversion into iterator for usage
-/// with [ValidatorProvingTrie](super::ValidatorProvingTrie::generate_for).
+/// with [ProvingTrie](super::ProvingTrie::generate_for).
 impl<T: Config> sp_std::iter::IntoIterator for ValidatorSet<T> {
 	type Item = (T::ValidatorId, T::FullIdentification);
 	type IntoIter = sp_std::vec::IntoIter<Self::Item>;
@@ -80,7 +79,7 @@ pub fn prove_session_membership<T: Config, D: AsRef<[u8]>>(
 ) -> Option<MembershipProof> {
 	let validators = ValidatorSet::<T>::load_from_offchain_db(session_index)?;
 	let count = validators.len() as u32;
-	let trie = ValidatorProvingTrie::<T>::generate_for(validators.into_iter()).ok()?;
+	let trie = ProvingTrie::<T>::generate_for(validators.into_iter()).ok()?;
 
 	let (id, data) = session_key;
 	trie.prove(id, data.as_ref()).map(|trie_nodes| MembershipProof {

From ead7951b3b701cded49789839157bee6619e897d Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 4 Apr 2024 13:28:36 -0400
Subject: [PATCH 10/55] Update substrate/primitives/runtime/src/proving_trie.rs

Co-authored-by: Ankan <10196091+Ank4n@users.noreply.github.com>
---
 substrate/primitives/runtime/src/proving_trie.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index 3e72e89bc229..ba126db28809 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -40,7 +40,7 @@ where
 	/// nodes within the current `MemoryDB` are insufficient to query the item.
 	fn query(&self, key: Key) -> Option<Value>;
 	/// Create the full verification data needed to prove a key and its value in the trie. Returns
-	/// `None` if the nodes nodes within the current `MemoryDB` are insufficient to create a proof.
+	/// `None` if the nodes within the current `MemoryDB` are insufficient to create a proof.
 	fn create_proof(&self, key: Key) -> Option<Vec<Vec<u8>>>;
 	/// Create a new instance of `ProvingTrie` from raw nodes. Nodes can be generated using the
 	/// `create_proof` function.

From acfe734f8e149db8124e87cc09b181d6c3132093 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 15 Aug 2024 14:28:47 -0400
Subject: [PATCH 11/55] fix some feedback

---
 substrate/primitives/runtime/Cargo.toml       |  1 +
 .../primitives/runtime/src/proving_trie.rs    | 44 +++++--------------
 2 files changed, 12 insertions(+), 33 deletions(-)

diff --git a/substrate/primitives/runtime/Cargo.toml b/substrate/primitives/runtime/Cargo.toml
index 2476d027f50c..800bf4bd0737 100644
--- a/substrate/primitives/runtime/Cargo.toml
+++ b/substrate/primitives/runtime/Cargo.toml
@@ -32,6 +32,7 @@ sp-arithmetic = { workspace = true }
 sp-core = { workspace = true }
 sp-io = { workspace = true }
 sp-std = { workspace = true }
+sp-trie = { workspace = true }
 sp-weights = { workspace = true }
 docify = { workspace = true }
 tracing = { workspace = true, features = ["log"], default-features = false }
diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index ba126db28809..0110cd2051e9 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -22,46 +22,24 @@ use crate::{Decode, DispatchError, Encode};
 use sp_std::vec::Vec;
 use sp_trie::{
 	trie_types::{TrieDBBuilder, TrieDBMutBuilderV0},
-	LayoutV0, MemoryDB, Recorder, Trie, TrieMut, EMPTY_PREFIX,
+	LayoutV1, MemoryDB, Recorder, Trie, TrieMut, EMPTY_PREFIX,
 };
 
-/// A trait for creating a merkle trie for checking and generating merkle proofs.
-pub trait ProvingTrie<Hashing, Hash, Key, Value>
-where
-	Self: Sized,
-{
-	/// Create a new instance of a `ProvingTrie` using an iterator of key/value pairs.
-	fn generate_for<I>(items: I) -> Result<Self, DispatchError>
-	where
-		I: IntoIterator<Item = (Key, Value)>;
-	/// Access the underlying trie root.
-	fn root(&self) -> &Hash;
-	/// Check a proof contained within the current memory-db. Returns `None` if the
-	/// nodes within the current `MemoryDB` are insufficient to query the item.
-	fn query(&self, key: Key) -> Option<Value>;
-	/// Create the full verification data needed to prove a key and its value in the trie. Returns
-	/// `None` if the nodes within the current `MemoryDB` are insufficient to create a proof.
-	fn create_proof(&self, key: Key) -> Option<Vec<Vec<u8>>>;
-	/// Create a new instance of `ProvingTrie` from raw nodes. Nodes can be generated using the
-	/// `create_proof` function.
-	fn from_nodes(root: Hash, nodes: &[Vec<u8>]) -> Self;
-}
+type HashOf<Hashing> = <Hashing as sp_core::Hasher>::Out;
 
 /// A basic trie implementation for checking and generating proofs for a key / value pair.
-pub struct BasicProvingTrie<Hashing, Hash, Key, Value>
+pub struct BasicProvingTrie<Hashing, Key, Value>
 where
-	Hashing: sp_core::Hasher<Out = Hash>,
+	Hashing: sp_core::Hasher,
 {
 	db: MemoryDB<Hashing>,
-	root: Hash,
+	root: HashOf<Hashing>,
 	_phantom: core::marker::PhantomData<(Key, Value)>,
 }
 
-impl<Hashing, Hash, Key, Value> ProvingTrie<Hashing, Hash, Key, Value>
-	for BasicProvingTrie<Hashing, Hash, Key, Value>
+impl<Hashing, Key, Value> BasicProvingTrie<Hashing, Key, Value>
 where
-	Hashing: sp_core::Hasher<Out = Hash>,
-	Hash: Default + Send + Sync,
+	Hashing: sp_core::Hasher,
 	Key: Encode,
 	Value: Encode + Decode,
 {
@@ -83,7 +61,7 @@ where
 		Ok(Self { db, root, _phantom: Default::default() })
 	}
 
-	fn root(&self) -> &Hash {
+	fn root(&self) -> &HashOf<Hashing> {
 		&self.root
 	}
 
@@ -95,7 +73,7 @@ where
 	}
 
 	fn create_proof(&self, key: Key) -> Option<Vec<Vec<u8>>> {
-		let mut recorder = Recorder::<LayoutV0<Hashing>>::new();
+		let mut recorder = Recorder::<LayoutV1<Hashing>>::new();
 
 		{
 			let trie =
@@ -109,7 +87,7 @@ where
 		Some(recorder.drain().into_iter().map(|r| r.data).collect())
 	}
 
-	fn from_nodes(root: Hash, nodes: &[Vec<u8>]) -> Self {
+	fn from_nodes(root: HashOf<Hashing>, nodes: &[Vec<u8>]) -> Self {
 		use sp_trie::HashDBT;
 
 		let mut memory_db = MemoryDB::default();
@@ -129,7 +107,7 @@ mod tests {
 	use sp_std::{collections::btree_map::BTreeMap, str::FromStr};
 
 	// A trie which simulates a trie of accounts (u32) and balances (u128).
-	type BalanceTrie = BasicProvingTrie<BlakeTwo256, H256, u32, u128>;
+	type BalanceTrie = BasicProvingTrie<BlakeTwo256, u32, u128>;
 
 	// The expected root hash for an empty trie.
 	fn empty_root() -> H256 {

From 45f42877d91a56e672180725666de66607420ae3 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 15 Aug 2024 14:30:55 -0400
Subject: [PATCH 12/55] update name

---
 substrate/primitives/runtime/src/proving_trie.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index 0110cd2051e9..09447a49a091 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -72,7 +72,7 @@ where
 			.and_then(|raw| Value::decode(&mut &*raw).ok())
 	}
 
-	fn create_proof(&self, key: Key) -> Option<Vec<Vec<u8>>> {
+	fn create_single_value_proof(&self, key: Key) -> Option<Vec<Vec<u8>>> {
 		let mut recorder = Recorder::<LayoutV1<Hashing>>::new();
 
 		{
@@ -143,9 +143,9 @@ mod tests {
 		assert_eq!(balance_trie.query(69u32), None);
 
 		// Create a proof for a valid key.
-		let proof = balance_trie.create_proof(6u32).unwrap();
+		let proof = balance_trie.create_single_value_proof(6u32).unwrap();
 		// Can't create proof for invalid key.
-		assert_eq!(balance_trie.create_proof(69u32), None);
+		assert_eq!(balance_trie.create_single_value_proof(69u32), None);
 
 		// Create a new proving trie from the proof.
 		let new_balance_trie = BalanceTrie::from_nodes(root, &proof);

From efbe1c93e9031de5844fd240996f3b853486a0d4 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 15 Aug 2024 14:58:26 -0400
Subject: [PATCH 13/55] docs and multi value proof

---
 .../primitives/runtime/src/proving_trie.rs    | 41 ++++++++++++++++---
 1 file changed, 36 insertions(+), 5 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index 09447a49a091..4879b663b9cd 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -43,7 +43,8 @@ where
 	Key: Encode,
 	Value: Encode + Decode,
 {
-	fn generate_for<I>(items: I) -> Result<Self, DispatchError>
+	/// Create a new instance of a `ProvingTrie` using an iterator of key/value pairs.
+	pub fn generate_for<I>(items: I) -> Result<Self, DispatchError>
 	where
 		I: IntoIterator<Item = (Key, Value)>,
 	{
@@ -61,18 +62,46 @@ where
 		Ok(Self { db, root, _phantom: Default::default() })
 	}
 
-	fn root(&self) -> &HashOf<Hashing> {
+	/// Access the underlying trie root.
+	pub fn root(&self) -> &HashOf<Hashing> {
 		&self.root
 	}
 
-	fn query(&self, key: Key) -> Option<Value> {
+	/// Check a proof contained within the current memory-db. Returns `None` if the
+	/// nodes within the current `MemoryDB` are insufficient to query the item.
+	pub fn query(&self, key: Key) -> Option<Value> {
 		let trie = TrieDBBuilder::new(&self.db, &self.root).build();
 		key.using_encoded(|s| trie.get(s))
 			.ok()?
 			.and_then(|raw| Value::decode(&mut &*raw).ok())
 	}
 
-	fn create_single_value_proof(&self, key: Key) -> Option<Vec<Vec<u8>>> {
+	/// Create the full verification data needed to prove all `keys` and their values in the trie.
+	/// Returns `None` if the nodes within the current `MemoryDB` are insufficient to create a
+	/// proof.
+	pub fn create_proof(&self, keys: Vec<Key>) -> Option<Vec<Vec<u8>>> {
+		let mut recorder = Recorder::<LayoutV1<Hashing>>::new();
+
+		{
+			let trie =
+				TrieDBBuilder::new(&self.db, &self.root).with_recorder(&mut recorder).build();
+
+			keys.iter()
+				.map(|key| {
+					key.using_encoded(|k| {
+						trie.get(k).ok()?.and_then(|raw| Value::decode(&mut &*raw).ok())
+					})
+				})
+				.collect::<Option<Vec<_>>>()?;
+		}
+
+		Some(recorder.drain().into_iter().map(|r| r.data).collect())
+	}
+
+	/// Create the full verification data needed to prove a single key and its value in the trie.
+	/// Returns `None` if the nodes within the current `MemoryDB` are insufficient to create a
+	/// proof.
+	pub fn create_single_value_proof(&self, key: Key) -> Option<Vec<Vec<u8>>> {
 		let mut recorder = Recorder::<LayoutV1<Hashing>>::new();
 
 		{
@@ -87,7 +116,9 @@ where
 		Some(recorder.drain().into_iter().map(|r| r.data).collect())
 	}
 
-	fn from_nodes(root: HashOf<Hashing>, nodes: &[Vec<u8>]) -> Self {
+	/// Create a new instance of `ProvingTrie` from raw nodes. Nodes can be generated using the
+	/// `create_proof` function.
+	pub fn from_nodes(root: HashOf<Hashing>, nodes: &[Vec<u8>]) -> Self {
 		use sp_trie::HashDBT;
 
 		let mut memory_db = MemoryDB::default();

From cfa62ce6e8b34438d8544ebd17fd693dcc087432 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 15 Aug 2024 15:09:47 -0400
Subject: [PATCH 14/55] improve test

---
 .../primitives/runtime/src/proving_trie.rs    | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index 4879b663b9cd..5d11e3731955 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -67,7 +67,7 @@ where
 		&self.root
 	}
 
-	/// Check a proof contained within the current memory-db. Returns `None` if the
+	/// Check a proof contained within the current `MemoryDB`. Returns `None` if the
 	/// nodes within the current `MemoryDB` are insufficient to query the item.
 	pub fn query(&self, key: Key) -> Option<Value> {
 		let trie = TrieDBBuilder::new(&self.db, &self.root).build();
@@ -153,10 +153,10 @@ mod tests {
 	}
 
 	#[test]
-	fn basic_end_to_end() {
+	fn basic_end_to_end_single_value() {
 		// Create a map of users and their balances.
 		let mut map = BTreeMap::<u32, u128>::new();
-		for i in 0..10u32 {
+		for i in 0..100u32 {
 			map.insert(i, i.into());
 		}
 
@@ -167,21 +167,28 @@ mod tests {
 		let root = *balance_trie.root();
 		assert!(root != empty_root());
 
-		// Assert valid key is queryable.
+		// Assert valid keys are queryable.
 		assert_eq!(balance_trie.query(6u32), Some(6u128));
 		assert_eq!(balance_trie.query(9u32), Some(9u128));
+		assert_eq!(balance_trie.query(69u32), Some(69u128));
 		// Invalid key returns none.
-		assert_eq!(balance_trie.query(69u32), None);
+		assert_eq!(balance_trie.query(6969u32), None);
 
 		// Create a proof for a valid key.
 		let proof = balance_trie.create_single_value_proof(6u32).unwrap();
 		// Can't create proof for invalid key.
-		assert_eq!(balance_trie.create_single_value_proof(69u32), None);
+		assert_eq!(balance_trie.create_single_value_proof(6969u32), None);
 
 		// Create a new proving trie from the proof.
 		let new_balance_trie = BalanceTrie::from_nodes(root, &proof);
 
 		// Assert valid key is queryable.
 		assert_eq!(new_balance_trie.query(6u32), Some(6u128));
+		// A "neighbor" key is queryable, by happenstance.
+		assert_eq!(new_balance_trie.query(9u32), Some(9u128));
+		// A "non-neighbor" key is not queryable.
+		assert_eq!(new_balance_trie.query(69u32), None);
+		// An invalid key is not queryable.
+		assert_eq!(new_balance_trie.query(6969u32), None);
 	}
 }

From 885fd942a1ea156eaaeb85771506ea33eb775d26 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 15 Aug 2024 15:14:36 -0400
Subject: [PATCH 15/55] add multi-value query test

---
 .../primitives/runtime/src/proving_trie.rs    | 41 +++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index 5d11e3731955..3050394aabc8 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -191,4 +191,45 @@ mod tests {
 		// An invalid key is not queryable.
 		assert_eq!(new_balance_trie.query(6969u32), None);
 	}
+
+	#[test]
+	fn basic_end_to_end_multi_value() {
+		// Create a map of users and their balances.
+		let mut map = BTreeMap::<u32, u128>::new();
+		for i in 0..100u32 {
+			map.insert(i, i.into());
+		}
+
+		// Put items into the trie.
+		let balance_trie = BalanceTrie::generate_for(map).unwrap();
+
+		// Root is changed.
+		let root = *balance_trie.root();
+		assert!(root != empty_root());
+
+		// Assert valid keys are queryable.
+		assert_eq!(balance_trie.query(6u32), Some(6u128));
+		assert_eq!(balance_trie.query(9u32), Some(9u128));
+		assert_eq!(balance_trie.query(69u32), Some(69u128));
+		// Invalid key returns none.
+		assert_eq!(balance_trie.query(6969u32), None);
+
+		// Create a proof for a valid key.
+		let proof = balance_trie.create_proof(vec![6u32, 69u32]).unwrap();
+		// Can't create proof for invalid key.
+		assert_eq!(balance_trie.create_proof(vec![6u32, 69u32, 6969u32]), None);
+
+		// Create a new proving trie from the proof.
+		let new_balance_trie = BalanceTrie::from_nodes(root, &proof);
+
+		// Assert valid keys are queryable.
+		assert_eq!(new_balance_trie.query(6u32), Some(6u128));
+		assert_eq!(new_balance_trie.query(69u32), Some(69u128));
+		// A "neighbor" key is queryable, by happenstance.
+		assert_eq!(new_balance_trie.query(9u32), Some(9u128));
+		// A "non-neighbor" key is not queryable.
+		assert_eq!(new_balance_trie.query(20u32), None);
+		// An invalid key is not queryable.
+		assert_eq!(new_balance_trie.query(6969u32), None);
+	}
 }

From 5e3e51897a2b88e0b85197c70dcd8979079669ae Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 15 Aug 2024 15:19:32 -0400
Subject: [PATCH 16/55] Create pr_3881.prdoc

---
 prdoc/pr_3881.prdoc | 15 +++++++++++++++
 1 file changed, 15 insertions(+)
 create mode 100644 prdoc/pr_3881.prdoc

diff --git a/prdoc/pr_3881.prdoc b/prdoc/pr_3881.prdoc
new file mode 100644
index 000000000000..57e11d91c885
--- /dev/null
+++ b/prdoc/pr_3881.prdoc
@@ -0,0 +1,15 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: Introduce a Generic Proving Trie
+
+doc:
+  - audience: Runtime Dev
+    description: |
+      This PR introduces a Proving Trie object which can be used inside the runtime. This can allow
+      for things like airdrops where a single hash is stored on chain representing the whole airdrop
+      and individuals present a proof of their inclusion in the airdrop.
+
+crates:
+  - name: sp-runtime
+    bump: minor

From e6c375906b51965d57563fcc29410ff1f3572c41 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Sat, 17 Aug 2024 05:34:45 +0200
Subject: [PATCH 17/55] use v1

---
 substrate/primitives/runtime/src/proving_trie.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index 3050394aabc8..1fe29e9b3882 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -21,7 +21,7 @@ use crate::{Decode, DispatchError, Encode};
 
 use sp_std::vec::Vec;
 use sp_trie::{
-	trie_types::{TrieDBBuilder, TrieDBMutBuilderV0},
+	trie_types::{TrieDBBuilder, TrieDBMutBuilderV1},
 	LayoutV1, MemoryDB, Recorder, Trie, TrieMut, EMPTY_PREFIX,
 };
 
@@ -52,7 +52,7 @@ where
 		let mut root = Default::default();
 
 		{
-			let mut trie = TrieDBMutBuilderV0::new(&mut db, &mut root).build();
+			let mut trie = TrieDBMutBuilderV1::new(&mut db, &mut root).build();
 			for (key, value) in items.into_iter() {
 				key.using_encoded(|k| value.using_encoded(|v| trie.insert(k, v)))
 					.map_err(|_| "failed to insert into trie")?;

From b40c96a41d1f2c19c95c77e2cb8b6dc0ba6ebf71 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 19 Aug 2024 03:13:29 +0200
Subject: [PATCH 18/55] initial idea

---
 substrate/frame/assets/src/functions.rs | 25 +++++++++++++
 substrate/frame/assets/src/lib.rs       | 49 +++++++++++++++++++++++++
 substrate/frame/assets/src/weights.rs   | 10 +++++
 3 files changed, 84 insertions(+)

diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index c218c4ddc952..6176a8373a60 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -438,6 +438,31 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		Ok(())
 	}
 
+	/// Increases the asset `id` balance of `beneficiary` by `amount`.
+	///
+	/// This alters the registered supply of the asset and emits an event.
+	///
+	/// Will return an error or will increase the amount by exactly `amount`.
+	pub(super) fn do_mint_distribution(
+		id: T::AssetId,
+		merkle_root: T::Hash,
+		maybe_check_issuer: Option<T::AccountId>,
+	) -> DispatchResult {
+		let details = Asset::<T, I>::get(&id).ok_or(Error::<T, I>::Unknown)?;
+		ensure!(details.status == AssetStatus::Live, Error::<T, I>::AssetNotLive);
+
+		if let Some(check_issuer) = maybe_check_issuer {
+			ensure!(check_issuer == details.issuer, Error::<T, I>::NoPermission);
+		}
+
+		let distribution_count: u32 = MerklizedDistribution::<T, I>::count();
+		MerklizedDistribution::<T, I>::insert(&distribution_count, (id.clone(), merkle_root.clone()));
+
+		Self::deposit_event(Event::DistributionIssued { asset_id: id, merkle_root: merkle_root });
+
+		Ok(())
+	}
+
 	/// Increases the asset `id` balance of `beneficiary` by `amount`.
 	///
 	/// LOW-LEVEL: Does not alter the supply of asset or emit an event. Use `do_mint` if you need
diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index e909932bfc82..18b584bbce31 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -456,6 +456,30 @@ pub mod pallet {
 		ValueQuery,
 	>;
 
+	type DistributionCounter = u32;
+
+	#[pallet::storage]
+	/// Merklized distribution of an asset.
+	pub(super) type MerklizedDistribution<T: Config<I>, I: 'static = ()> = CountedStorageMap<
+		_,
+		Blake2_128Concat,
+		DistributionCounter,
+		(T::AssetId, T::Hash),
+		OptionQuery,
+	>;
+
+	#[pallet::storage]
+	/// Tracks the merklized distribution of an asset so that assets are only claimed once.
+	pub(super) type MerklizedDistributionTracker<T: Config<I>, I: 'static = ()> = StorageDoubleMap<
+		_,
+		Blake2_128Concat,
+		DistributionCounter,
+		Blake2_128Concat,
+		T::AccountId,
+		(),
+		OptionQuery,
+	>;
+
 	/// The asset ID enforced for the next asset creation, if any present. Otherwise, this storage
 	/// item has no effect.
 	///
@@ -639,6 +663,8 @@ pub mod pallet {
 		Deposited { asset_id: T::AssetId, who: T::AccountId, amount: T::Balance },
 		/// Some assets were withdrawn from the account (e.g. for transaction fees).
 		Withdrawn { asset_id: T::AssetId, who: T::AccountId, amount: T::Balance },
+		/// A distribution of assets were issued.
+		DistributionIssued { asset_id: T::AssetId, merkle_root: T::Hash },
 	}
 
 	#[pallet::error]
@@ -1798,6 +1824,29 @@ pub mod pallet {
 			)?;
 			Ok(())
 		}
+
+		/// Mint a distribution of assets of a particular class.
+		///
+		/// The origin must be Signed and the sender must be the Issuer of the asset `id`.
+		///
+		/// - `id`: The identifier of the asset to have some amount minted.
+		/// - `merkle_root`: The merkle root of a binary trie used to authorize minting.
+		///
+		/// Emits `DistributionIssued` event when successful.
+		///
+		/// Weight: `O(1)`
+		/// Modes: Pre-existing balance of `beneficiary`; Account pre-existence of `beneficiary`.
+		#[pallet::call_index(33)]
+		pub fn mint_distribution(
+			origin: OriginFor<T>,
+			id: T::AssetIdParameter,
+			merkle_root: T::Hash,
+		) -> DispatchResult {
+			let origin = ensure_signed(origin)?;
+			let id: T::AssetId = id.into();
+			Self::do_mint_distribution(id, merkle_root, Some(origin))?;
+			Ok(())
+		}
 	}
 
 	/// Implements [`AccountTouch`] trait.
diff --git a/substrate/frame/assets/src/weights.rs b/substrate/frame/assets/src/weights.rs
index 57f7e951b73c..509566c77bb6 100644
--- a/substrate/frame/assets/src/weights.rs
+++ b/substrate/frame/assets/src/weights.rs
@@ -84,6 +84,7 @@ pub trait WeightInfo {
 	fn refund_other() -> Weight;
 	fn block() -> Weight;
 	fn transfer_all() -> Weight;
+	fn mint_distribution() -> Weight;
 }
 
 /// Weights for `pallet_assets` using the Substrate node and recommended hardware.
@@ -541,6 +542,11 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+		fn mint_distribution() -> Weight {
+			Weight::default()
+		}
+
 }
 
 // For backwards compatibility and tests.
@@ -997,4 +1003,8 @@ impl WeightInfo for () {
 			.saturating_add(RocksDbWeight::get().reads(1_u64))
 			.saturating_add(RocksDbWeight::get().writes(1_u64))
 	}
+
+			fn mint_distribution() -> Weight {
+			Weight::default()
+		}
 }

From 37052ee698a8e36b516285213e523cd8448dd80d Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 19 Aug 2024 03:42:50 +0200
Subject: [PATCH 19/55] some bs code to make things compile

---
 Cargo.lock                                    |  3 ++
 substrate/frame/assets/Cargo.toml             |  1 +
 substrate/frame/assets/src/functions.rs       | 16 ++++++---
 substrate/frame/assets/src/lib.rs             | 28 +++++++++++++--
 substrate/frame/assets/src/weights.rs         | 21 +++++++----
 substrate/utils/binary-merkle-tree/Cargo.toml |  6 +++-
 substrate/utils/binary-merkle-tree/src/lib.rs | 35 ++++++++++++++++++-
 7 files changed, 95 insertions(+), 15 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4148b660cd3c..fc057f165ff8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1481,6 +1481,8 @@ dependencies = [
  "array-bytes",
  "hash-db",
  "log",
+ "parity-scale-codec",
+ "scale-info",
  "sp-core",
  "sp-runtime",
  "sp-tracing 16.0.0",
@@ -9851,6 +9853,7 @@ dependencies = [
 name = "pallet-assets"
 version = "29.1.0"
 dependencies = [
+ "binary-merkle-tree",
  "frame-benchmarking",
  "frame-support",
  "frame-system",
diff --git a/substrate/frame/assets/Cargo.toml b/substrate/frame/assets/Cargo.toml
index e20b576d0836..6605c65aa075 100644
--- a/substrate/frame/assets/Cargo.toml
+++ b/substrate/frame/assets/Cargo.toml
@@ -20,6 +20,7 @@ codec = { workspace = true }
 impl-trait-for-tuples = "0.2.2"
 log = { workspace = true }
 scale-info = { features = ["derive"], workspace = true }
+binary-merkle-tree = { workspace = true }
 # Needed for various traits. In our case, `OnFinalize`.
 sp-runtime = { workspace = true }
 # Needed for type-safe access to storage DB.
diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index 6176a8373a60..17254e2298b0 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -438,11 +438,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		Ok(())
 	}
 
-	/// Increases the asset `id` balance of `beneficiary` by `amount`.
-	///
-	/// This alters the registered supply of the asset and emits an event.
-	///
-	/// Will return an error or will increase the amount by exactly `amount`.
+	/// Creates a distribution in storage for asset `id`, which can be claimed via `do_claim_distribution`.
 	pub(super) fn do_mint_distribution(
 		id: T::AssetId,
 		merkle_root: T::Hash,
@@ -463,6 +459,16 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		Ok(())
 	}
 
+	/// A wrapper around `do_mint`, allowing a `merkle_proof` to control the amount minted and to whom.
+	pub(super) fn do_claim_distribution(
+		distribution_id: DistributionCounter,
+		merkle_proof: MerkleProof<T::Hash, (T::AccountId, T::Balance)>,
+	) -> DispatchResult {
+		let (asset_id, merkle_root) = MerklizedDistribution::<T, I>::get(distribution_id).ok_or(Error::<T, I>::Unknown)?;
+
+		Ok(())
+	}
+
 	/// Increases the asset `id` balance of `beneficiary` by `amount`.
 	///
 	/// LOW-LEVEL: Does not alter the supply of asset or emit an event. Use `do_mint` if you need
diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index 18b584bbce31..4af709cbe93f 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -194,6 +194,8 @@ use frame_support::{
 };
 use frame_system::Config as SystemConfig;
 
+use binary_merkle_tree::MerkleProof;
+
 pub use pallet::*;
 pub use weights::WeightInfo;
 
@@ -456,7 +458,7 @@ pub mod pallet {
 		ValueQuery,
 	>;
 
-	type DistributionCounter = u32;
+	pub type DistributionCounter = u32;
 
 	#[pallet::storage]
 	/// Merklized distribution of an asset.
@@ -1830,7 +1832,7 @@ pub mod pallet {
 		/// The origin must be Signed and the sender must be the Issuer of the asset `id`.
 		///
 		/// - `id`: The identifier of the asset to have some amount minted.
-		/// - `merkle_root`: The merkle root of a binary trie used to authorize minting.
+		/// - `merkle_root`: The merkle root of a binary tree used to authorize minting.
 		///
 		/// Emits `DistributionIssued` event when successful.
 		///
@@ -1847,6 +1849,28 @@ pub mod pallet {
 			Self::do_mint_distribution(id, merkle_root, Some(origin))?;
 			Ok(())
 		}
+
+		/// Claim a distribution of assets of a particular class.
+		///
+		/// Any signed origin may call this function.
+		///
+		/// - `distribution_id`: The identifier of the distribution.
+		/// - `merkle_proof`: The merkle proof of the account and balance in a binary tree used to authorize minting.
+		///
+		/// Emits `Issued` event when successful.
+		///
+		/// Weight: `O(1)`
+		/// Modes: Pre-existing balance of `beneficiary`; Account pre-existence of `beneficiary`.
+		#[pallet::call_index(34)]
+		pub fn claim_distribution(
+			origin: OriginFor<T>,
+			distribution_id: DistributionCounter,
+			merkle_proof: MerkleProof<T::Hash, (T::AccountId, T::Balance)>,
+		) -> DispatchResult {
+			ensure_signed(origin)?;
+			Self::do_claim_distribution(distribution_id, merkle_proof)?;
+			Ok(())
+		}
 	}
 
 	/// Implements [`AccountTouch`] trait.
diff --git a/substrate/frame/assets/src/weights.rs b/substrate/frame/assets/src/weights.rs
index 509566c77bb6..bb05c449af7d 100644
--- a/substrate/frame/assets/src/weights.rs
+++ b/substrate/frame/assets/src/weights.rs
@@ -85,6 +85,7 @@ pub trait WeightInfo {
 	fn block() -> Weight;
 	fn transfer_all() -> Weight;
 	fn mint_distribution() -> Weight;
+	fn claim_distribution() -> Weight;
 }
 
 /// Weights for `pallet_assets` using the Substrate node and recommended hardware.
@@ -543,9 +544,13 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
 
-		fn mint_distribution() -> Weight {
-			Weight::default()
-		}
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
 
 }
 
@@ -1004,7 +1009,11 @@ impl WeightInfo for () {
 			.saturating_add(RocksDbWeight::get().writes(1_u64))
 	}
 
-			fn mint_distribution() -> Weight {
-			Weight::default()
-		}
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
 }
diff --git a/substrate/utils/binary-merkle-tree/Cargo.toml b/substrate/utils/binary-merkle-tree/Cargo.toml
index 087ec5fd6c6d..5002c028da01 100644
--- a/substrate/utils/binary-merkle-tree/Cargo.toml
+++ b/substrate/utils/binary-merkle-tree/Cargo.toml
@@ -15,6 +15,10 @@ workspace = true
 array-bytes = { optional = true, workspace = true, default-features = true }
 log = { optional = true, workspace = true }
 hash-db = { workspace = true }
+codec = { features = [
+	"derive",
+], workspace = true }
+scale-info = { features = ["derive"], workspace = true }
 
 [dev-dependencies]
 array-bytes = { workspace = true, default-features = true }
@@ -25,4 +29,4 @@ sp-runtime = { workspace = true, default-features = true }
 [features]
 debug = ["array-bytes", "log"]
 default = ["debug", "std"]
-std = ["hash-db/std", "log/std", "sp-core/std", "sp-runtime/std"]
+std = ["codec/std", "scale-info/std", "hash-db/std", "log/std", "sp-core/std", "sp-runtime/std"]
diff --git a/substrate/utils/binary-merkle-tree/src/lib.rs b/substrate/utils/binary-merkle-tree/src/lib.rs
index f2d338cf028e..682df0a4aa07 100644
--- a/substrate/utils/binary-merkle-tree/src/lib.rs
+++ b/substrate/utils/binary-merkle-tree/src/lib.rs
@@ -38,6 +38,8 @@ use alloc::vec;
 use alloc::vec::Vec;
 
 use hash_db::Hasher;
+use codec::{Encode, Decode};
+use scale_info::TypeInfo;
 
 /// Construct a root hash of a Binary Merkle Tree created from given leaves.
 ///
@@ -87,7 +89,8 @@ where
 /// A generated merkle proof.
 ///
 /// The structure contains all necessary data to later on verify the proof and the leaf itself.
-#[derive(Debug, PartialEq, Eq)]
+// #[derive(Encode, Decode, Debug, PartialEq, Eq, TypeInfo)]
+#[derive(Debug, PartialEq, Eq, Clone)]
 pub struct MerkleProof<H, L> {
 	/// Root hash of generated merkle tree.
 	pub root: H,
@@ -107,6 +110,36 @@ pub struct MerkleProof<H, L> {
 	pub leaf: L,
 }
 
+// USIZE breaks these needed traits
+impl<H: 'static, L: 'static> TypeInfo for MerkleProof<H, L> {
+	type Identity = Self;
+
+	fn type_info() -> scale_info::Type {
+		scale_info::Type::builder()
+			.path(scale_info::Path::new("Dummy TODO", module_path!()))
+			.composite(
+				scale_info::build::Fields::unnamed()
+					.field(|f| f.ty::<u8>().docs(&["Dummy TODO"])),
+			)
+	}
+}
+
+impl<H, L> Encode for MerkleProof<H, L> {
+	fn size_hint(&self) -> usize {
+		true.size_hint()
+	}
+
+	fn using_encoded<R, F: FnOnce(&[u8]) -> R>(&self, f: F) -> R {
+		true.using_encoded(f)
+	}
+}
+
+impl<H, L> Decode for MerkleProof<H, L> {
+	fn decode<I: codec::Input>(input: &mut I) -> Result<Self, codec::Error> {
+		Err("DUMMY TODO".into())
+	}
+}
+
 /// A trait of object inspecting merkle root creation.
 ///
 /// It can be passed to [`merkelize_row`] or [`merkelize`] functions and will be notified

From 08f2b355a9249dc9311df6a33b4672a463597151 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 19 Aug 2024 04:12:48 +0200
Subject: [PATCH 20/55] more stuff

---
 substrate/frame/assets/src/functions.rs       | 13 +++++++++++--
 substrate/frame/assets/src/lib.rs             | 10 ++++++----
 substrate/utils/binary-merkle-tree/src/lib.rs | 16 ++++++++++++++++
 3 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index 17254e2298b0..808fc653b717 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -441,7 +441,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 	/// Creates a distribution in storage for asset `id`, which can be claimed via `do_claim_distribution`.
 	pub(super) fn do_mint_distribution(
 		id: T::AssetId,
-		merkle_root: T::Hash,
+		merkle_root: H256,
 		maybe_check_issuer: Option<T::AccountId>,
 	) -> DispatchResult {
 		let details = Asset::<T, I>::get(&id).ok_or(Error::<T, I>::Unknown)?;
@@ -462,10 +462,19 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 	/// A wrapper around `do_mint`, allowing a `merkle_proof` to control the amount minted and to whom.
 	pub(super) fn do_claim_distribution(
 		distribution_id: DistributionCounter,
-		merkle_proof: MerkleProof<T::Hash, (T::AccountId, T::Balance)>,
+		merkle_proof: DistributionProof<T, I>,
 	) -> DispatchResult {
 		let (asset_id, merkle_root) = MerklizedDistribution::<T, I>::get(distribution_id).ok_or(Error::<T, I>::Unknown)?;
 
+		ensure!(merkle_root == merkle_proof.root, "TODO ERROR");
+
+		// binary_merkle_tree::verify_proof::<_, _, _>(
+		// 	&merkle_proof.root,
+		// 	merkle_proof.proof,
+		// 	merkle_proof.number_of_leaves,
+		// 	merkle_proof.leaf_index,
+		// 	merkle_proof.leaf,
+		// );
 		Ok(())
 	}
 
diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index 4af709cbe93f..bfc46e9615d5 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -193,6 +193,7 @@ use frame_support::{
 	},
 };
 use frame_system::Config as SystemConfig;
+use sp_core::H256;
 
 use binary_merkle_tree::MerkleProof;
 
@@ -459,6 +460,7 @@ pub mod pallet {
 	>;
 
 	pub type DistributionCounter = u32;
+	pub type DistributionProof<T, I> = MerkleProof<H256, (<T as frame_system::Config>::AccountId, <T as Config<I>>::Balance)>;
 
 	#[pallet::storage]
 	/// Merklized distribution of an asset.
@@ -466,7 +468,7 @@ pub mod pallet {
 		_,
 		Blake2_128Concat,
 		DistributionCounter,
-		(T::AssetId, T::Hash),
+		(T::AssetId, H256),
 		OptionQuery,
 	>;
 
@@ -666,7 +668,7 @@ pub mod pallet {
 		/// Some assets were withdrawn from the account (e.g. for transaction fees).
 		Withdrawn { asset_id: T::AssetId, who: T::AccountId, amount: T::Balance },
 		/// A distribution of assets were issued.
-		DistributionIssued { asset_id: T::AssetId, merkle_root: T::Hash },
+		DistributionIssued { asset_id: T::AssetId, merkle_root: H256 },
 	}
 
 	#[pallet::error]
@@ -1842,7 +1844,7 @@ pub mod pallet {
 		pub fn mint_distribution(
 			origin: OriginFor<T>,
 			id: T::AssetIdParameter,
-			merkle_root: T::Hash,
+			merkle_root: H256,
 		) -> DispatchResult {
 			let origin = ensure_signed(origin)?;
 			let id: T::AssetId = id.into();
@@ -1865,7 +1867,7 @@ pub mod pallet {
 		pub fn claim_distribution(
 			origin: OriginFor<T>,
 			distribution_id: DistributionCounter,
-			merkle_proof: MerkleProof<T::Hash, (T::AccountId, T::Balance)>,
+			merkle_proof: DistributionProof<T, I>,
 		) -> DispatchResult {
 			ensure_signed(origin)?;
 			Self::do_claim_distribution(distribution_id, merkle_proof)?;
diff --git a/substrate/utils/binary-merkle-tree/src/lib.rs b/substrate/utils/binary-merkle-tree/src/lib.rs
index 682df0a4aa07..28f46546eec3 100644
--- a/substrate/utils/binary-merkle-tree/src/lib.rs
+++ b/substrate/utils/binary-merkle-tree/src/lib.rs
@@ -260,6 +260,22 @@ impl<'a, H, T: AsRef<[u8]>> From<&'a T> for Leaf<'a, H> {
 	}
 }
 
+// pub fn verify_merkle_proof<'a, H, L>(
+// 	merkle_proof: &'a MerkleProof<H::Out, L>,
+// ) -> bool
+// where
+// 	H: Hasher,
+// 	L: Into<Leaf<'a, H::Out>>,
+// {
+// 	verify_proof::<H, Vec<H::Out>, L>(
+// 		&merkle_proof.root,
+// 		merkle_proof.proof,
+// 		merkle_proof.number_of_leaves,
+// 		merkle_proof.leaf_index,
+// 		merkle_proof.leaf,
+// 	)
+// }
+
 /// Verify Merkle Proof correctness versus given root hash.
 ///
 /// The proof is NOT expected to contain leaf hash as the first

From 118a51c7359bf59d6b6ff6350316c5b582c597ec Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 19 Aug 2024 04:25:00 +0200
Subject: [PATCH 21/55] complete logic

---
 substrate/frame/assets/src/functions.rs | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index 808fc653b717..62e551252898 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -467,14 +467,23 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		let (asset_id, merkle_root) = MerklizedDistribution::<T, I>::get(distribution_id).ok_or(Error::<T, I>::Unknown)?;
 
 		ensure!(merkle_root == merkle_proof.root, "TODO ERROR");
+		let (beneficiary, amount) = merkle_proof.leaf;
+		let already_claimed = MerklizedDistributionTracker::<T, I>::contains_key(&distribution_id, &beneficiary);
+		ensure!(!already_claimed, "TODO ERROR");
 
-		// binary_merkle_tree::verify_proof::<_, _, _>(
+		use sp_runtime::traits::BlakeTwo256;
+		let verified = true;
+		// let verified = binary_merkle_tree::verify_proof::<BlakeTwo256, _, (T::AccountId, T::Balance)>(
 		// 	&merkle_proof.root,
 		// 	merkle_proof.proof,
 		// 	merkle_proof.number_of_leaves,
 		// 	merkle_proof.leaf_index,
 		// 	merkle_proof.leaf,
 		// );
+
+		Self::do_mint(asset_id, &beneficiary, amount, None)?;
+		MerklizedDistributionTracker::<T, I>::insert(&distribution_id, &beneficiary, ());
+
 		Ok(())
 	}
 

From 44bc5bb8e03a8109f4d3f39bcde9a5fbc6dab98b Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 19 Aug 2024 04:28:52 +0200
Subject: [PATCH 22/55] add verification check

---
 substrate/frame/assets/src/functions.rs | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index 62e551252898..c982595b3921 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -481,6 +481,8 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		// 	merkle_proof.leaf,
 		// );
 
+		ensure!(verified, "TODO ERROR");
+
 		Self::do_mint(asset_id, &beneficiary, amount, None)?;
 		MerklizedDistributionTracker::<T, I>::insert(&distribution_id, &beneficiary, ());
 

From 89abcd0cdab76de9ea23943818419cc9352f78c4 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 19 Aug 2024 03:13:29 +0200
Subject: [PATCH 23/55] initial idea

---
 substrate/frame/assets/src/functions.rs | 25 +++++++++++++
 substrate/frame/assets/src/lib.rs       | 49 +++++++++++++++++++++++++
 substrate/frame/assets/src/weights.rs   | 10 +++++
 3 files changed, 84 insertions(+)

diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index c218c4ddc952..6176a8373a60 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -438,6 +438,31 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		Ok(())
 	}
 
+	/// Increases the asset `id` balance of `beneficiary` by `amount`.
+	///
+	/// This alters the registered supply of the asset and emits an event.
+	///
+	/// Will return an error or will increase the amount by exactly `amount`.
+	pub(super) fn do_mint_distribution(
+		id: T::AssetId,
+		merkle_root: T::Hash,
+		maybe_check_issuer: Option<T::AccountId>,
+	) -> DispatchResult {
+		let details = Asset::<T, I>::get(&id).ok_or(Error::<T, I>::Unknown)?;
+		ensure!(details.status == AssetStatus::Live, Error::<T, I>::AssetNotLive);
+
+		if let Some(check_issuer) = maybe_check_issuer {
+			ensure!(check_issuer == details.issuer, Error::<T, I>::NoPermission);
+		}
+
+		let distribution_count: u32 = MerklizedDistribution::<T, I>::count();
+		MerklizedDistribution::<T, I>::insert(&distribution_count, (id.clone(), merkle_root.clone()));
+
+		Self::deposit_event(Event::DistributionIssued { asset_id: id, merkle_root: merkle_root });
+
+		Ok(())
+	}
+
 	/// Increases the asset `id` balance of `beneficiary` by `amount`.
 	///
 	/// LOW-LEVEL: Does not alter the supply of asset or emit an event. Use `do_mint` if you need
diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index e909932bfc82..18b584bbce31 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -456,6 +456,30 @@ pub mod pallet {
 		ValueQuery,
 	>;
 
+	type DistributionCounter = u32;
+
+	#[pallet::storage]
+	/// Merklized distribution of an asset.
+	pub(super) type MerklizedDistribution<T: Config<I>, I: 'static = ()> = CountedStorageMap<
+		_,
+		Blake2_128Concat,
+		DistributionCounter,
+		(T::AssetId, T::Hash),
+		OptionQuery,
+	>;
+
+	#[pallet::storage]
+	/// Tracks the merklized distribution of an asset so that assets are only claimed once.
+	pub(super) type MerklizedDistributionTracker<T: Config<I>, I: 'static = ()> = StorageDoubleMap<
+		_,
+		Blake2_128Concat,
+		DistributionCounter,
+		Blake2_128Concat,
+		T::AccountId,
+		(),
+		OptionQuery,
+	>;
+
 	/// The asset ID enforced for the next asset creation, if any present. Otherwise, this storage
 	/// item has no effect.
 	///
@@ -639,6 +663,8 @@ pub mod pallet {
 		Deposited { asset_id: T::AssetId, who: T::AccountId, amount: T::Balance },
 		/// Some assets were withdrawn from the account (e.g. for transaction fees).
 		Withdrawn { asset_id: T::AssetId, who: T::AccountId, amount: T::Balance },
+		/// A distribution of assets were issued.
+		DistributionIssued { asset_id: T::AssetId, merkle_root: T::Hash },
 	}
 
 	#[pallet::error]
@@ -1798,6 +1824,29 @@ pub mod pallet {
 			)?;
 			Ok(())
 		}
+
+		/// Mint a distribution of assets of a particular class.
+		///
+		/// The origin must be Signed and the sender must be the Issuer of the asset `id`.
+		///
+		/// - `id`: The identifier of the asset to have some amount minted.
+		/// - `merkle_root`: The merkle root of a binary trie used to authorize minting.
+		///
+		/// Emits `DistributionIssued` event when successful.
+		///
+		/// Weight: `O(1)`
+		/// Modes: Pre-existing balance of `beneficiary`; Account pre-existence of `beneficiary`.
+		#[pallet::call_index(33)]
+		pub fn mint_distribution(
+			origin: OriginFor<T>,
+			id: T::AssetIdParameter,
+			merkle_root: T::Hash,
+		) -> DispatchResult {
+			let origin = ensure_signed(origin)?;
+			let id: T::AssetId = id.into();
+			Self::do_mint_distribution(id, merkle_root, Some(origin))?;
+			Ok(())
+		}
 	}
 
 	/// Implements [`AccountTouch`] trait.
diff --git a/substrate/frame/assets/src/weights.rs b/substrate/frame/assets/src/weights.rs
index 57f7e951b73c..509566c77bb6 100644
--- a/substrate/frame/assets/src/weights.rs
+++ b/substrate/frame/assets/src/weights.rs
@@ -84,6 +84,7 @@ pub trait WeightInfo {
 	fn refund_other() -> Weight;
 	fn block() -> Weight;
 	fn transfer_all() -> Weight;
+	fn mint_distribution() -> Weight;
 }
 
 /// Weights for `pallet_assets` using the Substrate node and recommended hardware.
@@ -541,6 +542,11 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+		fn mint_distribution() -> Weight {
+			Weight::default()
+		}
+
 }
 
 // For backwards compatibility and tests.
@@ -997,4 +1003,8 @@ impl WeightInfo for () {
 			.saturating_add(RocksDbWeight::get().reads(1_u64))
 			.saturating_add(RocksDbWeight::get().writes(1_u64))
 	}
+
+			fn mint_distribution() -> Weight {
+			Weight::default()
+		}
 }

From 8ad3381dfa0c9b0519ed09ed056fc334531a78de Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 19 Aug 2024 03:42:50 +0200
Subject: [PATCH 24/55] some bs code to make things compile

---
 Cargo.lock                                    |  3 ++
 substrate/frame/assets/Cargo.toml             |  1 +
 substrate/frame/assets/src/functions.rs       | 16 ++++++---
 substrate/frame/assets/src/lib.rs             | 28 +++++++++++++--
 substrate/frame/assets/src/weights.rs         | 21 +++++++----
 substrate/utils/binary-merkle-tree/Cargo.toml |  6 +++-
 substrate/utils/binary-merkle-tree/src/lib.rs | 35 ++++++++++++++++++-
 7 files changed, 95 insertions(+), 15 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 599faf3d4df8..d55d400e98c9 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1473,6 +1473,8 @@ dependencies = [
  "array-bytes",
  "hash-db",
  "log",
+ "parity-scale-codec",
+ "scale-info",
  "sp-core",
  "sp-runtime",
  "sp-tracing 16.0.0",
@@ -9847,6 +9849,7 @@ dependencies = [
 name = "pallet-assets"
 version = "29.1.0"
 dependencies = [
+ "binary-merkle-tree",
  "frame-benchmarking",
  "frame-support",
  "frame-system",
diff --git a/substrate/frame/assets/Cargo.toml b/substrate/frame/assets/Cargo.toml
index e20b576d0836..6605c65aa075 100644
--- a/substrate/frame/assets/Cargo.toml
+++ b/substrate/frame/assets/Cargo.toml
@@ -20,6 +20,7 @@ codec = { workspace = true }
 impl-trait-for-tuples = "0.2.2"
 log = { workspace = true }
 scale-info = { features = ["derive"], workspace = true }
+binary-merkle-tree = { workspace = true }
 # Needed for various traits. In our case, `OnFinalize`.
 sp-runtime = { workspace = true }
 # Needed for type-safe access to storage DB.
diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index 6176a8373a60..17254e2298b0 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -438,11 +438,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		Ok(())
 	}
 
-	/// Increases the asset `id` balance of `beneficiary` by `amount`.
-	///
-	/// This alters the registered supply of the asset and emits an event.
-	///
-	/// Will return an error or will increase the amount by exactly `amount`.
+	/// Creates a distribution in storage for asset `id`, which can be claimed via `do_claim_distribution`.
 	pub(super) fn do_mint_distribution(
 		id: T::AssetId,
 		merkle_root: T::Hash,
@@ -463,6 +459,16 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		Ok(())
 	}
 
+	/// A wrapper around `do_mint`, allowing a `merkle_proof` to control the amount minted and to whom.
+	pub(super) fn do_claim_distribution(
+		distribution_id: DistributionCounter,
+		merkle_proof: MerkleProof<T::Hash, (T::AccountId, T::Balance)>,
+	) -> DispatchResult {
+		let (asset_id, merkle_root) = MerklizedDistribution::<T, I>::get(distribution_id).ok_or(Error::<T, I>::Unknown)?;
+
+		Ok(())
+	}
+
 	/// Increases the asset `id` balance of `beneficiary` by `amount`.
 	///
 	/// LOW-LEVEL: Does not alter the supply of asset or emit an event. Use `do_mint` if you need
diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index 18b584bbce31..4af709cbe93f 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -194,6 +194,8 @@ use frame_support::{
 };
 use frame_system::Config as SystemConfig;
 
+use binary_merkle_tree::MerkleProof;
+
 pub use pallet::*;
 pub use weights::WeightInfo;
 
@@ -456,7 +458,7 @@ pub mod pallet {
 		ValueQuery,
 	>;
 
-	type DistributionCounter = u32;
+	pub type DistributionCounter = u32;
 
 	#[pallet::storage]
 	/// Merklized distribution of an asset.
@@ -1830,7 +1832,7 @@ pub mod pallet {
 		/// The origin must be Signed and the sender must be the Issuer of the asset `id`.
 		///
 		/// - `id`: The identifier of the asset to have some amount minted.
-		/// - `merkle_root`: The merkle root of a binary trie used to authorize minting.
+		/// - `merkle_root`: The merkle root of a binary tree used to authorize minting.
 		///
 		/// Emits `DistributionIssued` event when successful.
 		///
@@ -1847,6 +1849,28 @@ pub mod pallet {
 			Self::do_mint_distribution(id, merkle_root, Some(origin))?;
 			Ok(())
 		}
+
+		/// Claim a distribution of assets of a particular class.
+		///
+		/// Any signed origin may call this function.
+		///
+		/// - `distribution_id`: The identifier of the distribution.
+		/// - `merkle_proof`: The merkle proof of the account and balance in a binary tree used to authorize minting.
+		///
+		/// Emits `Issued` event when successful.
+		///
+		/// Weight: `O(1)`
+		/// Modes: Pre-existing balance of `beneficiary`; Account pre-existence of `beneficiary`.
+		#[pallet::call_index(34)]
+		pub fn claim_distribution(
+			origin: OriginFor<T>,
+			distribution_id: DistributionCounter,
+			merkle_proof: MerkleProof<T::Hash, (T::AccountId, T::Balance)>,
+		) -> DispatchResult {
+			ensure_signed(origin)?;
+			Self::do_claim_distribution(distribution_id, merkle_proof)?;
+			Ok(())
+		}
 	}
 
 	/// Implements [`AccountTouch`] trait.
diff --git a/substrate/frame/assets/src/weights.rs b/substrate/frame/assets/src/weights.rs
index 509566c77bb6..bb05c449af7d 100644
--- a/substrate/frame/assets/src/weights.rs
+++ b/substrate/frame/assets/src/weights.rs
@@ -85,6 +85,7 @@ pub trait WeightInfo {
 	fn block() -> Weight;
 	fn transfer_all() -> Weight;
 	fn mint_distribution() -> Weight;
+	fn claim_distribution() -> Weight;
 }
 
 /// Weights for `pallet_assets` using the Substrate node and recommended hardware.
@@ -543,9 +544,13 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
 
-		fn mint_distribution() -> Weight {
-			Weight::default()
-		}
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
 
 }
 
@@ -1004,7 +1009,11 @@ impl WeightInfo for () {
 			.saturating_add(RocksDbWeight::get().writes(1_u64))
 	}
 
-			fn mint_distribution() -> Weight {
-			Weight::default()
-		}
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
 }
diff --git a/substrate/utils/binary-merkle-tree/Cargo.toml b/substrate/utils/binary-merkle-tree/Cargo.toml
index 087ec5fd6c6d..5002c028da01 100644
--- a/substrate/utils/binary-merkle-tree/Cargo.toml
+++ b/substrate/utils/binary-merkle-tree/Cargo.toml
@@ -15,6 +15,10 @@ workspace = true
 array-bytes = { optional = true, workspace = true, default-features = true }
 log = { optional = true, workspace = true }
 hash-db = { workspace = true }
+codec = { features = [
+	"derive",
+], workspace = true }
+scale-info = { features = ["derive"], workspace = true }
 
 [dev-dependencies]
 array-bytes = { workspace = true, default-features = true }
@@ -25,4 +29,4 @@ sp-runtime = { workspace = true, default-features = true }
 [features]
 debug = ["array-bytes", "log"]
 default = ["debug", "std"]
-std = ["hash-db/std", "log/std", "sp-core/std", "sp-runtime/std"]
+std = ["codec/std", "scale-info/std", "hash-db/std", "log/std", "sp-core/std", "sp-runtime/std"]
diff --git a/substrate/utils/binary-merkle-tree/src/lib.rs b/substrate/utils/binary-merkle-tree/src/lib.rs
index f2d338cf028e..682df0a4aa07 100644
--- a/substrate/utils/binary-merkle-tree/src/lib.rs
+++ b/substrate/utils/binary-merkle-tree/src/lib.rs
@@ -38,6 +38,8 @@ use alloc::vec;
 use alloc::vec::Vec;
 
 use hash_db::Hasher;
+use codec::{Encode, Decode};
+use scale_info::TypeInfo;
 
 /// Construct a root hash of a Binary Merkle Tree created from given leaves.
 ///
@@ -87,7 +89,8 @@ where
 /// A generated merkle proof.
 ///
 /// The structure contains all necessary data to later on verify the proof and the leaf itself.
-#[derive(Debug, PartialEq, Eq)]
+// #[derive(Encode, Decode, Debug, PartialEq, Eq, TypeInfo)]
+#[derive(Debug, PartialEq, Eq, Clone)]
 pub struct MerkleProof<H, L> {
 	/// Root hash of generated merkle tree.
 	pub root: H,
@@ -107,6 +110,36 @@ pub struct MerkleProof<H, L> {
 	pub leaf: L,
 }
 
+// USIZE breaks these needed traits
+impl<H: 'static, L: 'static> TypeInfo for MerkleProof<H, L> {
+	type Identity = Self;
+
+	fn type_info() -> scale_info::Type {
+		scale_info::Type::builder()
+			.path(scale_info::Path::new("Dummy TODO", module_path!()))
+			.composite(
+				scale_info::build::Fields::unnamed()
+					.field(|f| f.ty::<u8>().docs(&["Dummy TODO"])),
+			)
+	}
+}
+
+impl<H, L> Encode for MerkleProof<H, L> {
+	fn size_hint(&self) -> usize {
+		true.size_hint()
+	}
+
+	fn using_encoded<R, F: FnOnce(&[u8]) -> R>(&self, f: F) -> R {
+		true.using_encoded(f)
+	}
+}
+
+impl<H, L> Decode for MerkleProof<H, L> {
+	fn decode<I: codec::Input>(input: &mut I) -> Result<Self, codec::Error> {
+		Err("DUMMY TODO".into())
+	}
+}
+
 /// A trait of object inspecting merkle root creation.
 ///
 /// It can be passed to [`merkelize_row`] or [`merkelize`] functions and will be notified

From ae6fab7d81a5e242dae37abaa8e500b4802543e3 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 19 Aug 2024 04:12:48 +0200
Subject: [PATCH 25/55] more stuff

---
 substrate/frame/assets/src/functions.rs       | 13 +++++++++++--
 substrate/frame/assets/src/lib.rs             | 10 ++++++----
 substrate/utils/binary-merkle-tree/src/lib.rs | 16 ++++++++++++++++
 3 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index 17254e2298b0..808fc653b717 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -441,7 +441,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 	/// Creates a distribution in storage for asset `id`, which can be claimed via `do_claim_distribution`.
 	pub(super) fn do_mint_distribution(
 		id: T::AssetId,
-		merkle_root: T::Hash,
+		merkle_root: H256,
 		maybe_check_issuer: Option<T::AccountId>,
 	) -> DispatchResult {
 		let details = Asset::<T, I>::get(&id).ok_or(Error::<T, I>::Unknown)?;
@@ -462,10 +462,19 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 	/// A wrapper around `do_mint`, allowing a `merkle_proof` to control the amount minted and to whom.
 	pub(super) fn do_claim_distribution(
 		distribution_id: DistributionCounter,
-		merkle_proof: MerkleProof<T::Hash, (T::AccountId, T::Balance)>,
+		merkle_proof: DistributionProof<T, I>,
 	) -> DispatchResult {
 		let (asset_id, merkle_root) = MerklizedDistribution::<T, I>::get(distribution_id).ok_or(Error::<T, I>::Unknown)?;
 
+		ensure!(merkle_root == merkle_proof.root, "TODO ERROR");
+
+		// binary_merkle_tree::verify_proof::<_, _, _>(
+		// 	&merkle_proof.root,
+		// 	merkle_proof.proof,
+		// 	merkle_proof.number_of_leaves,
+		// 	merkle_proof.leaf_index,
+		// 	merkle_proof.leaf,
+		// );
 		Ok(())
 	}
 
diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index 4af709cbe93f..bfc46e9615d5 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -193,6 +193,7 @@ use frame_support::{
 	},
 };
 use frame_system::Config as SystemConfig;
+use sp_core::H256;
 
 use binary_merkle_tree::MerkleProof;
 
@@ -459,6 +460,7 @@ pub mod pallet {
 	>;
 
 	pub type DistributionCounter = u32;
+	pub type DistributionProof<T, I> = MerkleProof<H256, (<T as frame_system::Config>::AccountId, <T as Config<I>>::Balance)>;
 
 	#[pallet::storage]
 	/// Merklized distribution of an asset.
@@ -466,7 +468,7 @@ pub mod pallet {
 		_,
 		Blake2_128Concat,
 		DistributionCounter,
-		(T::AssetId, T::Hash),
+		(T::AssetId, H256),
 		OptionQuery,
 	>;
 
@@ -666,7 +668,7 @@ pub mod pallet {
 		/// Some assets were withdrawn from the account (e.g. for transaction fees).
 		Withdrawn { asset_id: T::AssetId, who: T::AccountId, amount: T::Balance },
 		/// A distribution of assets were issued.
-		DistributionIssued { asset_id: T::AssetId, merkle_root: T::Hash },
+		DistributionIssued { asset_id: T::AssetId, merkle_root: H256 },
 	}
 
 	#[pallet::error]
@@ -1842,7 +1844,7 @@ pub mod pallet {
 		pub fn mint_distribution(
 			origin: OriginFor<T>,
 			id: T::AssetIdParameter,
-			merkle_root: T::Hash,
+			merkle_root: H256,
 		) -> DispatchResult {
 			let origin = ensure_signed(origin)?;
 			let id: T::AssetId = id.into();
@@ -1865,7 +1867,7 @@ pub mod pallet {
 		pub fn claim_distribution(
 			origin: OriginFor<T>,
 			distribution_id: DistributionCounter,
-			merkle_proof: MerkleProof<T::Hash, (T::AccountId, T::Balance)>,
+			merkle_proof: DistributionProof<T, I>,
 		) -> DispatchResult {
 			ensure_signed(origin)?;
 			Self::do_claim_distribution(distribution_id, merkle_proof)?;
diff --git a/substrate/utils/binary-merkle-tree/src/lib.rs b/substrate/utils/binary-merkle-tree/src/lib.rs
index 682df0a4aa07..28f46546eec3 100644
--- a/substrate/utils/binary-merkle-tree/src/lib.rs
+++ b/substrate/utils/binary-merkle-tree/src/lib.rs
@@ -260,6 +260,22 @@ impl<'a, H, T: AsRef<[u8]>> From<&'a T> for Leaf<'a, H> {
 	}
 }
 
+// pub fn verify_merkle_proof<'a, H, L>(
+// 	merkle_proof: &'a MerkleProof<H::Out, L>,
+// ) -> bool
+// where
+// 	H: Hasher,
+// 	L: Into<Leaf<'a, H::Out>>,
+// {
+// 	verify_proof::<H, Vec<H::Out>, L>(
+// 		&merkle_proof.root,
+// 		merkle_proof.proof,
+// 		merkle_proof.number_of_leaves,
+// 		merkle_proof.leaf_index,
+// 		merkle_proof.leaf,
+// 	)
+// }
+
 /// Verify Merkle Proof correctness versus given root hash.
 ///
 /// The proof is NOT expected to contain leaf hash as the first

From 115f89b683d71c4dbc6c77cfd3e5f5a07f2c5cdb Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 19 Aug 2024 04:25:00 +0200
Subject: [PATCH 26/55] complete logic

---
 substrate/frame/assets/src/functions.rs | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index 808fc653b717..62e551252898 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -467,14 +467,23 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		let (asset_id, merkle_root) = MerklizedDistribution::<T, I>::get(distribution_id).ok_or(Error::<T, I>::Unknown)?;
 
 		ensure!(merkle_root == merkle_proof.root, "TODO ERROR");
+		let (beneficiary, amount) = merkle_proof.leaf;
+		let already_claimed = MerklizedDistributionTracker::<T, I>::contains_key(&distribution_id, &beneficiary);
+		ensure!(!already_claimed, "TODO ERROR");
 
-		// binary_merkle_tree::verify_proof::<_, _, _>(
+		use sp_runtime::traits::BlakeTwo256;
+		let verified = true;
+		// let verified = binary_merkle_tree::verify_proof::<BlakeTwo256, _, (T::AccountId, T::Balance)>(
 		// 	&merkle_proof.root,
 		// 	merkle_proof.proof,
 		// 	merkle_proof.number_of_leaves,
 		// 	merkle_proof.leaf_index,
 		// 	merkle_proof.leaf,
 		// );
+
+		Self::do_mint(asset_id, &beneficiary, amount, None)?;
+		MerklizedDistributionTracker::<T, I>::insert(&distribution_id, &beneficiary, ());
+
 		Ok(())
 	}
 

From e5f005622d3c41dac945ae2188f5537b101249cb Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 19 Aug 2024 04:28:52 +0200
Subject: [PATCH 27/55] add verification check

---
 substrate/frame/assets/src/functions.rs | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index 62e551252898..c982595b3921 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -481,6 +481,8 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		// 	merkle_proof.leaf,
 		// );
 
+		ensure!(verified, "TODO ERROR");
+
 		Self::do_mint(asset_id, &beneficiary, amount, None)?;
 		MerklizedDistributionTracker::<T, I>::insert(&distribution_id, &beneficiary, ());
 

From b80589cc164fa227564e675de96e63e6a3be1c8a Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 2 Sep 2024 22:01:28 -0400
Subject: [PATCH 28/55] large refactor to proving trie

---
 substrate/primitives/runtime/src/lib.rs       |  15 +
 .../primitives/runtime/src/proving_trie.rs    | 309 +++++++++++++-----
 2 files changed, 234 insertions(+), 90 deletions(-)

diff --git a/substrate/primitives/runtime/src/lib.rs b/substrate/primitives/runtime/src/lib.rs
index e36c3ae70bba..ba1ea3769724 100644
--- a/substrate/primitives/runtime/src/lib.rs
+++ b/substrate/primitives/runtime/src/lib.rs
@@ -104,6 +104,8 @@ pub use crate::runtime_string::*;
 // Re-export Multiaddress
 pub use multiaddress::MultiAddress;
 
+use proving_trie::TrieError;
+
 /// Re-export these since they're only "kind of" generic.
 pub use generic::{Digest, DigestItem};
 
@@ -593,6 +595,8 @@ pub enum DispatchError {
 	Unavailable,
 	/// Root origin is not allowed.
 	RootNotAllowed,
+	/// An error with tries.
+	Trie(TrieError),
 }
 
 /// Result of a `Dispatchable` which contains the `DispatchResult` and additional information about
@@ -698,6 +702,12 @@ impl From<ArithmeticError> for DispatchError {
 	}
 }
 
+impl From<TrieError> for DispatchError {
+	fn from(e: TrieError) -> DispatchError {
+		Self::Trie(e)
+	}
+}
+
 impl From<&'static str> for DispatchError {
 	fn from(err: &'static str) -> DispatchError {
 		Self::Other(err)
@@ -722,6 +732,7 @@ impl From<DispatchError> for &'static str {
 			Corruption => "State corrupt",
 			Unavailable => "Resource unavailable",
 			RootNotAllowed => "Root not allowed",
+			Trie(e) => e.into(),
 		}
 	}
 }
@@ -769,6 +780,10 @@ impl traits::Printable for DispatchError {
 			Corruption => "State corrupt".print(),
 			Unavailable => "Resource unavailable".print(),
 			RootNotAllowed => "Root not allowed".print(),
+			Trie(e) => {
+				"Trie error: ".print();
+				<&'static str>::from(*e).print();
+			},
 		}
 	}
 }
diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index 1fe29e9b3882..1bd5f7b2c299 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -17,16 +17,107 @@
 
 //! Types for a simple merkle trie used for checking and generating proofs.
 
-use crate::{Decode, DispatchError, Encode};
+use crate::{Decode, DispatchError, Encode, MaxEncodedLen, TypeInfo};
 
 use sp_std::vec::Vec;
 use sp_trie::{
-	trie_types::{TrieDBBuilder, TrieDBMutBuilderV1},
-	LayoutV1, MemoryDB, Recorder, Trie, TrieMut, EMPTY_PREFIX,
+	trie_types::{TrieDBBuilder, TrieDBMutBuilderV1, TrieError as SpTrieError},
+	LayoutV1, MemoryDB, Trie, TrieMut, VerifyError, EMPTY_PREFIX,
 };
 
+#[cfg(feature = "serde")]
+use crate::{Deserialize, Serialize};
+
 type HashOf<Hashing> = <Hashing as sp_core::Hasher>::Out;
 
+/// A runtime friendly error type for tries.
+#[derive(Eq, PartialEq, Clone, Copy, Encode, Decode, Debug, TypeInfo, MaxEncodedLen)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum TrieError {
+	/* From TrieError */
+	/// Attempted to create a trie with a state root not in the DB.
+	InvalidStateRoot,
+	/// Trie item not found in the database,
+	IncompleteDatabase,
+	/// A value was found in the trie with a nibble key that was not byte-aligned.
+	ValueAtIncompleteKey,
+	/// Corrupt Trie item.
+	DecoderError,
+	/// Hash is not value.
+	InvalidHash,
+	/* From VerifyError */
+	/// The statement being verified contains multiple key-value pairs with the same key. The
+	/// parameter is the duplicated key.
+	DuplicateKey,
+	/// The proof contains at least one extraneous node.
+	ExtraneousNode,
+	/// The proof contains at least one extraneous value which should have been omitted from the
+	/// proof.
+	ExtraneousValue,
+	/// The proof contains at least one extraneous hash reference the should have been omitted.
+	ExtraneousHashReference,
+	/// The proof contains an invalid child reference that exceeds the hash length.
+	InvalidChildReference,
+	/// The proof indicates that an expected value was not found in the trie.
+	ValueMismatch,
+	/// The proof is missing trie nodes required to verify.
+	IncompleteProof,
+	/// The root hash computed from the proof is incorrect.
+	RootMismatch,
+	/// One of the proof nodes could not be decoded.
+	DecodeError,
+}
+
+impl<T> From<SpTrieError<T>> for TrieError {
+	fn from(error: SpTrieError<T>) -> Self {
+		match error {
+			SpTrieError::InvalidStateRoot(..) => Self::InvalidStateRoot,
+			SpTrieError::IncompleteDatabase(..) => Self::IncompleteDatabase,
+			SpTrieError::ValueAtIncompleteKey(..) => Self::ValueAtIncompleteKey,
+			SpTrieError::DecoderError(..) => Self::DecoderError,
+			SpTrieError::InvalidHash(..) => Self::InvalidHash,
+		}
+	}
+}
+
+impl<T, U> From<VerifyError<T, U>> for TrieError {
+	fn from(error: VerifyError<T, U>) -> Self {
+		match error {
+			VerifyError::DuplicateKey(..) => Self::DuplicateKey,
+			VerifyError::ExtraneousNode => Self::ExtraneousNode,
+			VerifyError::ExtraneousValue(..) => Self::ExtraneousValue,
+			VerifyError::ExtraneousHashReference(..) => Self::ExtraneousHashReference,
+			VerifyError::InvalidChildReference(..) => Self::InvalidChildReference,
+			VerifyError::ValueMismatch(..) => Self::ValueMismatch,
+			VerifyError::IncompleteProof => Self::IncompleteProof,
+			VerifyError::RootMismatch(..) => Self::RootMismatch,
+			VerifyError::DecodeError(..) => Self::DecodeError,
+		}
+	}
+}
+
+impl From<TrieError> for &'static str {
+	fn from(e: TrieError) -> &'static str {
+		match e {
+			TrieError::InvalidStateRoot => "The state root is not in the database.",
+			TrieError::IncompleteDatabase => "A trie item was not found in the database.",
+			TrieError::ValueAtIncompleteKey =>
+				"A value was found with a key that is not byte-aligned.",
+			TrieError::DecoderError => "A corrupt trie item was encountered.",
+			TrieError::InvalidHash => "The hash does not match the expected value.",
+			TrieError::DuplicateKey => "The proof contains duplicate keys.",
+			TrieError::ExtraneousNode => "The proof contains extraneous nodes.",
+			TrieError::ExtraneousValue => "The proof contains extraneous values.",
+			TrieError::ExtraneousHashReference => "The proof contains extraneous hash references.",
+			TrieError::InvalidChildReference => "The proof contains an invalid child reference.",
+			TrieError::ValueMismatch => "The proof indicates a value mismatch.",
+			TrieError::IncompleteProof => "The proof is incomplete.",
+			TrieError::RootMismatch => "The root hash computed from the proof is incorrect.",
+			TrieError::DecodeError => "One of the proof nodes could not be decoded.",
+		}
+	}
+}
+
 /// A basic trie implementation for checking and generating proofs for a key / value pair.
 pub struct BasicProvingTrie<Hashing, Key, Value>
 where
@@ -76,49 +167,50 @@ where
 			.and_then(|raw| Value::decode(&mut &*raw).ok())
 	}
 
+	// /// Create the full verification data needed to prove all `keys` and their values in the
+	// trie. /// Returns `None` if the nodes within the current `MemoryDB` are insufficient to
+	// create a /// proof.
+	// pub fn create_proof(&self, keys: Vec<Key>) -> Option<RawStorageProof> {
+	// 	let mut recorder = Recorder::<LayoutV1<Hashing>>::new();
+
+	// 	{
+	// 		let trie =
+	// 			TrieDBBuilder::new(&self.db, &self.root).with_recorder(&mut recorder).build();
+
+	// 		keys.iter()
+	// 			.map(|key| {
+	// 				key.using_encoded(|k| {
+	// 					trie.get(k).ok()?.and_then(|raw| Value::decode(&mut &*raw).ok())
+	// 				})
+	// 			})
+	// 			.collect::<Option<Vec<_>>>()?;
+	// 	}
+
+	// 	Some(recorder.into_raw_storage_proof())
+	// }
+
 	/// Create the full verification data needed to prove all `keys` and their values in the trie.
 	/// Returns `None` if the nodes within the current `MemoryDB` are insufficient to create a
 	/// proof.
-	pub fn create_proof(&self, keys: Vec<Key>) -> Option<Vec<Vec<u8>>> {
-		let mut recorder = Recorder::<LayoutV1<Hashing>>::new();
-
-		{
-			let trie =
-				TrieDBBuilder::new(&self.db, &self.root).with_recorder(&mut recorder).build();
-
-			keys.iter()
-				.map(|key| {
-					key.using_encoded(|k| {
-						trie.get(k).ok()?.and_then(|raw| Value::decode(&mut &*raw).ok())
-					})
-				})
-				.collect::<Option<Vec<_>>>()?;
-		}
-
-		Some(recorder.drain().into_iter().map(|r| r.data).collect())
+	pub fn create_proof(&self, keys: &[Key]) -> Result<Vec<Vec<u8>>, DispatchError> {
+		sp_trie::generate_trie_proof::<LayoutV1<Hashing>, _, _, _>(
+			&self.db,
+			self.root,
+			&keys.into_iter().map(|k| k.encode()).collect::<Vec<Vec<u8>>>(),
+		)
+		.map_err(|err| TrieError::from(*err).into())
 	}
 
 	/// Create the full verification data needed to prove a single key and its value in the trie.
 	/// Returns `None` if the nodes within the current `MemoryDB` are insufficient to create a
 	/// proof.
-	pub fn create_single_value_proof(&self, key: Key) -> Option<Vec<Vec<u8>>> {
-		let mut recorder = Recorder::<LayoutV1<Hashing>>::new();
-
-		{
-			let trie =
-				TrieDBBuilder::new(&self.db, &self.root).with_recorder(&mut recorder).build();
-
-			key.using_encoded(|k| {
-				trie.get(k).ok()?.and_then(|raw| Value::decode(&mut &*raw).ok())
-			})?;
-		}
-
-		Some(recorder.drain().into_iter().map(|r| r.data).collect())
+	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<Vec<u8>>, DispatchError> {
+		self.create_proof(&[key])
 	}
 
 	/// Create a new instance of `ProvingTrie` from raw nodes. Nodes can be generated using the
-	/// `create_proof` function.
-	pub fn from_nodes(root: HashOf<Hashing>, nodes: &[Vec<u8>]) -> Self {
+	/// `create_proof` function. Only use internally.
+	fn from_nodes(root: HashOf<Hashing>, nodes: &[Vec<u8>]) -> Self {
 		use sp_trie::HashDBT;
 
 		let mut memory_db = MemoryDB::default();
@@ -130,6 +222,42 @@ where
 	}
 }
 
+/// Verify the existence or non-existence of `key` and `value` in a trie proof.
+pub fn verify_single_value_proof<Hashing, Key, Value>(
+	root: HashOf<Hashing>,
+	proof: &[Vec<u8>],
+	key: Key,
+	maybe_value: Option<Value>,
+) -> Result<(), DispatchError>
+where
+	Hashing: sp_core::Hasher,
+	Key: Encode,
+	Value: Encode,
+{
+	sp_trie::verify_trie_proof::<LayoutV1<Hashing>, _, _, _>(
+		&root,
+		proof,
+		&[(key.encode(), maybe_value.map(|value| value.encode()))],
+	)
+	.map_err(|err| TrieError::from(err).into())
+}
+
+/// Verify a proof which contains multiple keys and values.
+pub fn verify_proof<'a, Hashing, I, K, V>(
+	root: HashOf<Hashing>,
+	proof: &[Vec<u8>],
+	items: I,
+) -> Result<(), DispatchError>
+where
+	Hashing: sp_core::Hasher,
+	I: IntoIterator<Item = &'a (K, Option<V>)>,
+	K: 'a + AsRef<[u8]>,
+	V: 'a + AsRef<[u8]>,
+{
+	sp_trie::verify_trie_proof::<LayoutV1<Hashing>, _, _, _>(&root, proof, items)
+		.map_err(|err| TrieError::from(err).into())
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;
@@ -146,14 +274,7 @@ mod tests {
 			.unwrap()
 	}
 
-	#[test]
-	fn empty_trie_works() {
-		let empty_trie = BalanceTrie::generate_for(Vec::new()).unwrap();
-		assert_eq!(*empty_trie.root(), empty_root());
-	}
-
-	#[test]
-	fn basic_end_to_end_single_value() {
+	fn create_balance_trie() -> BalanceTrie {
 		// Create a map of users and their balances.
 		let mut map = BTreeMap::<u32, u128>::new();
 		for i in 0..100u32 {
@@ -174,62 +295,70 @@ mod tests {
 		// Invalid key returns none.
 		assert_eq!(balance_trie.query(6969u32), None);
 
+		balance_trie
+	}
+
+	#[test]
+	fn empty_trie_works() {
+		let empty_trie = BalanceTrie::generate_for(Vec::new()).unwrap();
+		assert_eq!(*empty_trie.root(), empty_root());
+	}
+
+	#[test]
+	fn basic_end_to_end_single_value() {
+		let balance_trie = create_balance_trie();
+		let root = *balance_trie.root();
+
 		// Create a proof for a valid key.
 		let proof = balance_trie.create_single_value_proof(6u32).unwrap();
-		// Can't create proof for invalid key.
-		assert_eq!(balance_trie.create_single_value_proof(6969u32), None);
-
-		// Create a new proving trie from the proof.
-		let new_balance_trie = BalanceTrie::from_nodes(root, &proof);
-
-		// Assert valid key is queryable.
-		assert_eq!(new_balance_trie.query(6u32), Some(6u128));
-		// A "neighbor" key is queryable, by happenstance.
-		assert_eq!(new_balance_trie.query(9u32), Some(9u128));
-		// A "non-neighbor" key is not queryable.
-		assert_eq!(new_balance_trie.query(69u32), None);
-		// An invalid key is not queryable.
-		assert_eq!(new_balance_trie.query(6969u32), None);
+
+		// Assert key is provable, all other keys are invalid.
+		for i in 0..200u32 {
+			if i == 6 {
+				assert_eq!(
+					verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, Some(6u128)),
+					Ok(())
+				);
+				// Wrong value is invalid.
+				assert_eq!(
+					verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, Some(7u128)),
+					Err(DispatchError::Trie)
+				);
+			} else {
+				assert_eq!(
+					verify_single_value_proof::<BlakeTwo256, _, _>(
+						root,
+						&proof,
+						i,
+						Some(u128::from(i))
+					),
+					Err(DispatchError::Trie)
+				);
+				assert_eq!(
+					verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, i, None::<u128>),
+					Err(DispatchError::Trie)
+				);
+			}
+		}
 	}
 
 	#[test]
 	fn basic_end_to_end_multi_value() {
-		// Create a map of users and their balances.
-		let mut map = BTreeMap::<u32, u128>::new();
-		for i in 0..100u32 {
-			map.insert(i, i.into());
-		}
-
-		// Put items into the trie.
-		let balance_trie = BalanceTrie::generate_for(map).unwrap();
-
-		// Root is changed.
+		let balance_trie = create_balance_trie();
 		let root = *balance_trie.root();
-		assert!(root != empty_root());
 
-		// Assert valid keys are queryable.
-		assert_eq!(balance_trie.query(6u32), Some(6u128));
-		assert_eq!(balance_trie.query(9u32), Some(9u128));
-		assert_eq!(balance_trie.query(69u32), Some(69u128));
-		// Invalid key returns none.
-		assert_eq!(balance_trie.query(6969u32), None);
+		// Create a proof for a valid and invalid key.
+		let proof = balance_trie.create_proof(&[6u32, 69u32, 6969u32]).unwrap();
 
-		// Create a proof for a valid key.
-		let proof = balance_trie.create_proof(vec![6u32, 69u32]).unwrap();
-		// Can't create proof for invalid key.
-		assert_eq!(balance_trie.create_proof(vec![6u32, 69u32, 6969u32]), None);
-
-		// Create a new proving trie from the proof.
-		let new_balance_trie = BalanceTrie::from_nodes(root, &proof);
+		let items = [(6u32, Some(6u128)), (69u32, Some(69u128)), (6969u32, None)];
+		let items_encoded = items
+			.into_iter()
+			.map(|(key, maybe_value)| (key.encode(), maybe_value.map(|value| value.encode())))
+			.collect::<Vec<(Vec<u8>, Option<Vec<u8>>)>>();
 
-		// Assert valid keys are queryable.
-		assert_eq!(new_balance_trie.query(6u32), Some(6u128));
-		assert_eq!(new_balance_trie.query(69u32), Some(69u128));
-		// A "neighbor" key is queryable, by happenstance.
-		assert_eq!(new_balance_trie.query(9u32), Some(9u128));
-		// A "non-neighbor" key is not queryable.
-		assert_eq!(new_balance_trie.query(20u32), None);
-		// An invalid key is not queryable.
-		assert_eq!(new_balance_trie.query(6969u32), None);
+		assert_eq!(verify_proof::<BlakeTwo256, _, _, _>(root, &proof, &items_encoded), Ok(()));
 	}
+
+	#[test]
+	fn proof_fails_with_bad_data() {}
 }

From 83c66e73a8fc68d7aa2ba97499f26465cd2865e9 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 2 Sep 2024 22:31:50 -0400
Subject: [PATCH 29/55] fix tests

---
 .../primitives/runtime/src/proving_trie.rs    | 38 ++++++++++++-------
 1 file changed, 24 insertions(+), 14 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index 1bd5f7b2c299..bc473c41aec6 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -46,8 +46,7 @@ pub enum TrieError {
 	/// Hash is not value.
 	InvalidHash,
 	/* From VerifyError */
-	/// The statement being verified contains multiple key-value pairs with the same key. The
-	/// parameter is the duplicated key.
+	/// The statement being verified contains multiple key-value pairs with the same key.
 	DuplicateKey,
 	/// The proof contains at least one extraneous node.
 	ExtraneousNode,
@@ -316,28 +315,39 @@ mod tests {
 		for i in 0..200u32 {
 			if i == 6 {
 				assert_eq!(
-					verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, Some(6u128)),
+					verify_single_value_proof::<BlakeTwo256, _, _>(
+						root,
+						&proof,
+						i,
+						Some(u128::from(i))
+					),
 					Ok(())
 				);
 				// Wrong value is invalid.
-				assert_eq!(
-					verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, Some(7u128)),
-					Err(DispatchError::Trie)
-				);
-			} else {
 				assert_eq!(
 					verify_single_value_proof::<BlakeTwo256, _, _>(
 						root,
 						&proof,
 						i,
-						Some(u128::from(i))
+						Some(u128::from(i + 1))
 					),
-					Err(DispatchError::Trie)
-				);
-				assert_eq!(
-					verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, i, None::<u128>),
-					Err(DispatchError::Trie)
+					Err(TrieError::RootMismatch.into())
 				);
+			} else {
+				assert!(verify_single_value_proof::<BlakeTwo256, _, _>(
+					root,
+					&proof,
+					i,
+					Some(u128::from(i))
+				)
+				.is_err());
+				assert!(verify_single_value_proof::<BlakeTwo256, _, _>(
+					root,
+					&proof,
+					i,
+					None::<u128>
+				)
+				.is_err());
 			}
 		}
 	}

From 5a4ff43e6b5f33e2bb96d13e890e9456c7458955 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 2 Sep 2024 22:43:31 -0400
Subject: [PATCH 30/55] use basic proving trie

---
 Cargo.lock                              |  1 -
 substrate/frame/assets/Cargo.toml       |  1 -
 substrate/frame/assets/src/functions.rs | 50 +++++++++++++------------
 substrate/frame/assets/src/lib.rs       | 23 ++++++------
 4 files changed, 38 insertions(+), 37 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index d55d400e98c9..c4cc4b110802 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -9849,7 +9849,6 @@ dependencies = [
 name = "pallet-assets"
 version = "29.1.0"
 dependencies = [
- "binary-merkle-tree",
  "frame-benchmarking",
  "frame-support",
  "frame-system",
diff --git a/substrate/frame/assets/Cargo.toml b/substrate/frame/assets/Cargo.toml
index 6605c65aa075..e20b576d0836 100644
--- a/substrate/frame/assets/Cargo.toml
+++ b/substrate/frame/assets/Cargo.toml
@@ -20,7 +20,6 @@ codec = { workspace = true }
 impl-trait-for-tuples = "0.2.2"
 log = { workspace = true }
 scale-info = { features = ["derive"], workspace = true }
-binary-merkle-tree = { workspace = true }
 # Needed for various traits. In our case, `OnFinalize`.
 sp-runtime = { workspace = true }
 # Needed for type-safe access to storage DB.
diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index c982595b3921..65b503208b2a 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -438,10 +438,11 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		Ok(())
 	}
 
-	/// Creates a distribution in storage for asset `id`, which can be claimed via `do_claim_distribution`.
+	/// Creates a distribution in storage for asset `id`, which can be claimed via
+	/// `do_claim_distribution`.
 	pub(super) fn do_mint_distribution(
 		id: T::AssetId,
-		merkle_root: H256,
+		merkle_root: T::Hash,
 		maybe_check_issuer: Option<T::AccountId>,
 	) -> DispatchResult {
 		let details = Asset::<T, I>::get(&id).ok_or(Error::<T, I>::Unknown)?;
@@ -452,36 +453,37 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		}
 
 		let distribution_count: u32 = MerklizedDistribution::<T, I>::count();
-		MerklizedDistribution::<T, I>::insert(&distribution_count, (id.clone(), merkle_root.clone()));
+		MerklizedDistribution::<T, I>::insert(
+			&distribution_count,
+			(id.clone(), merkle_root.clone()),
+		);
 
-		Self::deposit_event(Event::DistributionIssued { asset_id: id, merkle_root: merkle_root });
+		Self::deposit_event(Event::DistributionIssued { asset_id: id, merkle_root });
 
 		Ok(())
 	}
 
-	/// A wrapper around `do_mint`, allowing a `merkle_proof` to control the amount minted and to whom.
+	/// A wrapper around `do_mint`, allowing a `merkle_proof` to control the amount minted and to
+	/// whom.
 	pub(super) fn do_claim_distribution(
 		distribution_id: DistributionCounter,
-		merkle_proof: DistributionProof<T, I>,
+		beneficiary: T::AccountId,
+		amount: T::Balance,
+		merkle_proof: DistributionProof,
 	) -> DispatchResult {
-		let (asset_id, merkle_root) = MerklizedDistribution::<T, I>::get(distribution_id).ok_or(Error::<T, I>::Unknown)?;
-
-		ensure!(merkle_root == merkle_proof.root, "TODO ERROR");
-		let (beneficiary, amount) = merkle_proof.leaf;
-		let already_claimed = MerklizedDistributionTracker::<T, I>::contains_key(&distribution_id, &beneficiary);
-		ensure!(!already_claimed, "TODO ERROR");
-
-		use sp_runtime::traits::BlakeTwo256;
-		let verified = true;
-		// let verified = binary_merkle_tree::verify_proof::<BlakeTwo256, _, (T::AccountId, T::Balance)>(
-		// 	&merkle_proof.root,
-		// 	merkle_proof.proof,
-		// 	merkle_proof.number_of_leaves,
-		// 	merkle_proof.leaf_index,
-		// 	merkle_proof.leaf,
-		// );
-
-		ensure!(verified, "TODO ERROR");
+		if amount.is_zero() {
+			return Err(Error::<T, I>::BalanceLow.into())
+		}
+
+		let (asset_id, merkle_root) =
+			MerklizedDistribution::<T, I>::get(distribution_id).ok_or(Error::<T, I>::Unknown)?;
+
+		sp_runtime::proving_trie::verify_single_value_proof::<T::Hashing, _, _>(
+			merkle_root,
+			&merkle_proof,
+			&beneficiary,
+			Some(amount),
+		)?;
 
 		Self::do_mint(asset_id, &beneficiary, amount, None)?;
 		MerklizedDistributionTracker::<T, I>::insert(&distribution_id, &beneficiary, ());
diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index bfc46e9615d5..525491267540 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -193,9 +193,6 @@ use frame_support::{
 	},
 };
 use frame_system::Config as SystemConfig;
-use sp_core::H256;
-
-use binary_merkle_tree::MerkleProof;
 
 pub use pallet::*;
 pub use weights::WeightInfo;
@@ -460,7 +457,7 @@ pub mod pallet {
 	>;
 
 	pub type DistributionCounter = u32;
-	pub type DistributionProof<T, I> = MerkleProof<H256, (<T as frame_system::Config>::AccountId, <T as Config<I>>::Balance)>;
+	pub type DistributionProof = Vec<Vec<u8>>;
 
 	#[pallet::storage]
 	/// Merklized distribution of an asset.
@@ -468,7 +465,7 @@ pub mod pallet {
 		_,
 		Blake2_128Concat,
 		DistributionCounter,
-		(T::AssetId, H256),
+		(T::AssetId, T::Hash),
 		OptionQuery,
 	>;
 
@@ -668,7 +665,7 @@ pub mod pallet {
 		/// Some assets were withdrawn from the account (e.g. for transaction fees).
 		Withdrawn { asset_id: T::AssetId, who: T::AccountId, amount: T::Balance },
 		/// A distribution of assets were issued.
-		DistributionIssued { asset_id: T::AssetId, merkle_root: H256 },
+		DistributionIssued { asset_id: T::AssetId, merkle_root: T::Hash },
 	}
 
 	#[pallet::error]
@@ -1834,7 +1831,8 @@ pub mod pallet {
 		/// The origin must be Signed and the sender must be the Issuer of the asset `id`.
 		///
 		/// - `id`: The identifier of the asset to have some amount minted.
-		/// - `merkle_root`: The merkle root of a binary tree used to authorize minting.
+		/// - `merkle_root`: The merkle root of a compact base-16 merkle trie used to authorize
+		///   minting.
 		///
 		/// Emits `DistributionIssued` event when successful.
 		///
@@ -1844,7 +1842,7 @@ pub mod pallet {
 		pub fn mint_distribution(
 			origin: OriginFor<T>,
 			id: T::AssetIdParameter,
-			merkle_root: H256,
+			merkle_root: T::Hash,
 		) -> DispatchResult {
 			let origin = ensure_signed(origin)?;
 			let id: T::AssetId = id.into();
@@ -1857,7 +1855,8 @@ pub mod pallet {
 		/// Any signed origin may call this function.
 		///
 		/// - `distribution_id`: The identifier of the distribution.
-		/// - `merkle_proof`: The merkle proof of the account and balance in a binary tree used to authorize minting.
+		/// - `merkle_proof`: The merkle proof of the account and balance in a compact base-16
+		///   merkle trie used to authorize minting.
 		///
 		/// Emits `Issued` event when successful.
 		///
@@ -1867,10 +1866,12 @@ pub mod pallet {
 		pub fn claim_distribution(
 			origin: OriginFor<T>,
 			distribution_id: DistributionCounter,
-			merkle_proof: DistributionProof<T, I>,
+			beneficiary: T::AccountId,
+			amount: T::Balance,
+			merkle_proof: DistributionProof,
 		) -> DispatchResult {
 			ensure_signed(origin)?;
-			Self::do_claim_distribution(distribution_id, merkle_proof)?;
+			Self::do_claim_distribution(distribution_id, beneficiary, amount, merkle_proof)?;
 			Ok(())
 		}
 	}

From 84574e1df7c3ae2587f53282cd21bddd94491545 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 2 Sep 2024 22:47:06 -0400
Subject: [PATCH 31/55] undo changes to binary merkle tree

---
 substrate/utils/binary-merkle-tree/Cargo.toml |  6 +--
 substrate/utils/binary-merkle-tree/src/lib.rs | 51 +------------------
 2 files changed, 2 insertions(+), 55 deletions(-)

diff --git a/substrate/utils/binary-merkle-tree/Cargo.toml b/substrate/utils/binary-merkle-tree/Cargo.toml
index 5002c028da01..087ec5fd6c6d 100644
--- a/substrate/utils/binary-merkle-tree/Cargo.toml
+++ b/substrate/utils/binary-merkle-tree/Cargo.toml
@@ -15,10 +15,6 @@ workspace = true
 array-bytes = { optional = true, workspace = true, default-features = true }
 log = { optional = true, workspace = true }
 hash-db = { workspace = true }
-codec = { features = [
-	"derive",
-], workspace = true }
-scale-info = { features = ["derive"], workspace = true }
 
 [dev-dependencies]
 array-bytes = { workspace = true, default-features = true }
@@ -29,4 +25,4 @@ sp-runtime = { workspace = true, default-features = true }
 [features]
 debug = ["array-bytes", "log"]
 default = ["debug", "std"]
-std = ["codec/std", "scale-info/std", "hash-db/std", "log/std", "sp-core/std", "sp-runtime/std"]
+std = ["hash-db/std", "log/std", "sp-core/std", "sp-runtime/std"]
diff --git a/substrate/utils/binary-merkle-tree/src/lib.rs b/substrate/utils/binary-merkle-tree/src/lib.rs
index 28f46546eec3..f2d338cf028e 100644
--- a/substrate/utils/binary-merkle-tree/src/lib.rs
+++ b/substrate/utils/binary-merkle-tree/src/lib.rs
@@ -38,8 +38,6 @@ use alloc::vec;
 use alloc::vec::Vec;
 
 use hash_db::Hasher;
-use codec::{Encode, Decode};
-use scale_info::TypeInfo;
 
 /// Construct a root hash of a Binary Merkle Tree created from given leaves.
 ///
@@ -89,8 +87,7 @@ where
 /// A generated merkle proof.
 ///
 /// The structure contains all necessary data to later on verify the proof and the leaf itself.
-// #[derive(Encode, Decode, Debug, PartialEq, Eq, TypeInfo)]
-#[derive(Debug, PartialEq, Eq, Clone)]
+#[derive(Debug, PartialEq, Eq)]
 pub struct MerkleProof<H, L> {
 	/// Root hash of generated merkle tree.
 	pub root: H,
@@ -110,36 +107,6 @@ pub struct MerkleProof<H, L> {
 	pub leaf: L,
 }
 
-// USIZE breaks these needed traits
-impl<H: 'static, L: 'static> TypeInfo for MerkleProof<H, L> {
-	type Identity = Self;
-
-	fn type_info() -> scale_info::Type {
-		scale_info::Type::builder()
-			.path(scale_info::Path::new("Dummy TODO", module_path!()))
-			.composite(
-				scale_info::build::Fields::unnamed()
-					.field(|f| f.ty::<u8>().docs(&["Dummy TODO"])),
-			)
-	}
-}
-
-impl<H, L> Encode for MerkleProof<H, L> {
-	fn size_hint(&self) -> usize {
-		true.size_hint()
-	}
-
-	fn using_encoded<R, F: FnOnce(&[u8]) -> R>(&self, f: F) -> R {
-		true.using_encoded(f)
-	}
-}
-
-impl<H, L> Decode for MerkleProof<H, L> {
-	fn decode<I: codec::Input>(input: &mut I) -> Result<Self, codec::Error> {
-		Err("DUMMY TODO".into())
-	}
-}
-
 /// A trait of object inspecting merkle root creation.
 ///
 /// It can be passed to [`merkelize_row`] or [`merkelize`] functions and will be notified
@@ -260,22 +227,6 @@ impl<'a, H, T: AsRef<[u8]>> From<&'a T> for Leaf<'a, H> {
 	}
 }
 
-// pub fn verify_merkle_proof<'a, H, L>(
-// 	merkle_proof: &'a MerkleProof<H::Out, L>,
-// ) -> bool
-// where
-// 	H: Hasher,
-// 	L: Into<Leaf<'a, H::Out>>,
-// {
-// 	verify_proof::<H, Vec<H::Out>, L>(
-// 		&merkle_proof.root,
-// 		merkle_proof.proof,
-// 		merkle_proof.number_of_leaves,
-// 		merkle_proof.leaf_index,
-// 		merkle_proof.leaf,
-// 	)
-// }
-
 /// Verify Merkle Proof correctness versus given root hash.
 ///
 /// The proof is NOT expected to contain leaf hash as the first

From 444b1f7c619910bcb9aceb6857c7a0e1d7f538d5 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 2 Sep 2024 22:52:50 -0400
Subject: [PATCH 32/55] remove comment code

---
 .../primitives/runtime/src/proving_trie.rs    | 25 ++-----------------
 1 file changed, 2 insertions(+), 23 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index bc473c41aec6..de567386a1e7 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -15,7 +15,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//! Types for a simple merkle trie used for checking and generating proofs.
+//! Types for a compact base-16 merkle trie used for checking and generating proofs within the
+//! runtime.
 
 use crate::{Decode, DispatchError, Encode, MaxEncodedLen, TypeInfo};
 
@@ -166,28 +167,6 @@ where
 			.and_then(|raw| Value::decode(&mut &*raw).ok())
 	}
 
-	// /// Create the full verification data needed to prove all `keys` and their values in the
-	// trie. /// Returns `None` if the nodes within the current `MemoryDB` are insufficient to
-	// create a /// proof.
-	// pub fn create_proof(&self, keys: Vec<Key>) -> Option<RawStorageProof> {
-	// 	let mut recorder = Recorder::<LayoutV1<Hashing>>::new();
-
-	// 	{
-	// 		let trie =
-	// 			TrieDBBuilder::new(&self.db, &self.root).with_recorder(&mut recorder).build();
-
-	// 		keys.iter()
-	// 			.map(|key| {
-	// 				key.using_encoded(|k| {
-	// 					trie.get(k).ok()?.and_then(|raw| Value::decode(&mut &*raw).ok())
-	// 				})
-	// 			})
-	// 			.collect::<Option<Vec<_>>>()?;
-	// 	}
-
-	// 	Some(recorder.into_raw_storage_proof())
-	// }
-
 	/// Create the full verification data needed to prove all `keys` and their values in the trie.
 	/// Returns `None` if the nodes within the current `MemoryDB` are insufficient to create a
 	/// proof.

From f47b12eec2d89cb16ed7f1c79a5717349fe24131 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 2 Sep 2024 23:05:34 -0400
Subject: [PATCH 33/55] make api more runtime friendly

---
 Cargo.lock                                    |  2 -
 .../primitives/runtime/src/proving_trie.rs    | 38 ++++++-------------
 2 files changed, 12 insertions(+), 28 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6517a10547f8..6d37d9115435 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1476,8 +1476,6 @@ dependencies = [
  "array-bytes",
  "hash-db",
  "log",
- "parity-scale-codec",
- "scale-info",
  "sp-core",
  "sp-runtime",
  "sp-tracing 16.0.0",
diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index de567386a1e7..f3df7236b21f 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -23,7 +23,7 @@ use crate::{Decode, DispatchError, Encode, MaxEncodedLen, TypeInfo};
 use sp_std::vec::Vec;
 use sp_trie::{
 	trie_types::{TrieDBBuilder, TrieDBMutBuilderV1, TrieError as SpTrieError},
-	LayoutV1, MemoryDB, Trie, TrieMut, VerifyError, EMPTY_PREFIX,
+	LayoutV1, MemoryDB, Trie, TrieMut, VerifyError,
 };
 
 #[cfg(feature = "serde")]
@@ -185,19 +185,6 @@ where
 	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<Vec<u8>>, DispatchError> {
 		self.create_proof(&[key])
 	}
-
-	/// Create a new instance of `ProvingTrie` from raw nodes. Nodes can be generated using the
-	/// `create_proof` function. Only use internally.
-	fn from_nodes(root: HashOf<Hashing>, nodes: &[Vec<u8>]) -> Self {
-		use sp_trie::HashDBT;
-
-		let mut memory_db = MemoryDB::default();
-		for node in nodes {
-			HashDBT::insert(&mut memory_db, EMPTY_PREFIX, &node[..]);
-		}
-
-		Self { db: memory_db, root, _phantom: Default::default() }
-	}
 }
 
 /// Verify the existence or non-existence of `key` and `value` in a trie proof.
@@ -221,18 +208,22 @@ where
 }
 
 /// Verify a proof which contains multiple keys and values.
-pub fn verify_proof<'a, Hashing, I, K, V>(
+pub fn verify_proof<'a, Hashing, Key, Value>(
 	root: HashOf<Hashing>,
 	proof: &[Vec<u8>],
-	items: I,
+	items: &[(Key, Option<Value>)],
 ) -> Result<(), DispatchError>
 where
 	Hashing: sp_core::Hasher,
-	I: IntoIterator<Item = &'a (K, Option<V>)>,
-	K: 'a + AsRef<[u8]>,
-	V: 'a + AsRef<[u8]>,
+	Key: Encode,
+	Value: Encode,
 {
-	sp_trie::verify_trie_proof::<LayoutV1<Hashing>, _, _, _>(&root, proof, items)
+	let items_encoded = items
+		.into_iter()
+		.map(|(key, maybe_value)| (key.encode(), maybe_value.as_ref().map(|value| value.encode())))
+		.collect::<Vec<(Vec<u8>, Option<Vec<u8>>)>>();
+
+	sp_trie::verify_trie_proof::<LayoutV1<Hashing>, _, _, _>(&root, proof, &items_encoded)
 		.map_err(|err| TrieError::from(err).into())
 }
 
@@ -338,14 +329,9 @@ mod tests {
 
 		// Create a proof for a valid and invalid key.
 		let proof = balance_trie.create_proof(&[6u32, 69u32, 6969u32]).unwrap();
-
 		let items = [(6u32, Some(6u128)), (69u32, Some(69u128)), (6969u32, None)];
-		let items_encoded = items
-			.into_iter()
-			.map(|(key, maybe_value)| (key.encode(), maybe_value.map(|value| value.encode())))
-			.collect::<Vec<(Vec<u8>, Option<Vec<u8>>)>>();
 
-		assert_eq!(verify_proof::<BlakeTwo256, _, _, _>(root, &proof, &items_encoded), Ok(()));
+		assert_eq!(verify_proof::<BlakeTwo256, _, _>(root, &proof, &items), Ok(()));
 	}
 
 	#[test]

From ee90711281d5adc6949dda149313c0543de26363 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Fri, 6 Sep 2024 18:16:07 -0400
Subject: [PATCH 34/55] add basic test, and fix missing check for tracking
 distribution

---
 substrate/frame/assets/src/functions.rs |  5 ++
 substrate/frame/assets/src/lib.rs       |  2 +
 substrate/frame/assets/src/mock.rs      |  5 +-
 substrate/frame/assets/src/tests.rs     | 70 +++++++++++++++++++++++++
 4 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index 65b503208b2a..b8a784cca64f 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -478,6 +478,11 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		let (asset_id, merkle_root) =
 			MerklizedDistribution::<T, I>::get(distribution_id).ok_or(Error::<T, I>::Unknown)?;
 
+		ensure!(
+			!MerklizedDistributionTracker::<T, I>::contains_key(distribution_id, &beneficiary),
+			Error::<T, I>::AlreadyClaimed
+		);
+
 		sp_runtime::proving_trie::verify_single_value_proof::<T::Hashing, _, _>(
 			merkle_root,
 			&merkle_proof,
diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index 525491267540..2c113e958969 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -715,6 +715,8 @@ pub mod pallet {
 		CallbackFailed,
 		/// The asset ID must be equal to the [`NextAssetId`].
 		BadAssetId,
+		/// The asset distribution was already claimed!
+		AlreadyClaimed,
 	}
 
 	#[pallet::call(weight(<T as Config<I>>::WeightInfo))]
diff --git a/substrate/frame/assets/src/mock.rs b/substrate/frame/assets/src/mock.rs
index 2c160840e147..a425425f8c42 100644
--- a/substrate/frame/assets/src/mock.rs
+++ b/substrate/frame/assets/src/mock.rs
@@ -39,8 +39,9 @@ construct_runtime!(
 	}
 );
 
-type AccountId = u64;
-type AssetId = u32;
+pub(crate) type AccountId = u64;
+pub(crate) type AssetId = u32;
+pub(crate) type Balance = u64;
 
 #[derive_impl(frame_system::config_preludes::TestDefaultConfig)]
 impl frame_system::Config for Test {
diff --git a/substrate/frame/assets/src/tests.rs b/substrate/frame/assets/src/tests.rs
index af605c5a3c64..93dd59921691 100644
--- a/substrate/frame/assets/src/tests.rs
+++ b/substrate/frame/assets/src/tests.rs
@@ -1921,3 +1921,73 @@ fn asset_id_cannot_be_reused() {
 		assert!(Asset::<Test>::contains_key(7));
 	});
 }
+
+#[test]
+fn merklized_distribution_works() {
+	new_test_ext().execute_with(|| {
+		use alloc::collections::BTreeMap;
+		use sp_runtime::proving_trie::BasicProvingTrie;
+
+		// Create asset id 0 controlled by user 1, sufficient so it does not need ED.
+		assert_ok!(Assets::force_create(RuntimeOrigin::root(), 0, 1, true, 1));
+
+		// Offchain, user 1 creates a distribution of tokens.
+		type DistributionTrie =
+			BasicProvingTrie<<Test as frame_system::Config>::Hashing, AccountId, Balance>;
+		let mut distribution = BTreeMap::<AccountId, Balance>::new();
+		for i in 0..100u64 {
+			distribution.insert(i, i.into());
+		}
+
+		// Maybe the owner gives himself a little extra ;)
+		distribution.insert(1, 1337);
+
+		let distribution_trie = DistributionTrie::generate_for(distribution).unwrap();
+		let root = *distribution_trie.root();
+
+		// Use this trie root for the distribution
+		assert_ok!(Assets::mint_distribution(RuntimeOrigin::signed(1), 0, root));
+
+		// Now users claim their distributions permissionlessly with a proof.
+		let proof_for_1 = distribution_trie.create_single_value_proof(1).unwrap();
+		let amount_for_1 = distribution_trie.query(1).unwrap();
+		assert_ok!(Assets::claim_distribution(
+			RuntimeOrigin::signed(1),
+			0,
+			1,
+			amount_for_1,
+			proof_for_1
+		));
+		assert_eq!(Assets::balance(0, 1), 1337);
+
+		// Other users can claim their tokens.
+		let proof_for_69 = distribution_trie.create_single_value_proof(69).unwrap();
+		let amount_for_69 = distribution_trie.query(69).unwrap();
+		assert_ok!(Assets::claim_distribution(
+			RuntimeOrigin::signed(55),
+			0,
+			69,
+			amount_for_69,
+			proof_for_69
+		));
+		assert_eq!(Assets::balance(0, 69), 69);
+
+		// Owner (or anyone) can also distribute on behalf of the other users.
+		let proof_for_6 = distribution_trie.create_single_value_proof(6).unwrap();
+		let amount_for_6 = distribution_trie.query(6).unwrap();
+		assert_ok!(Assets::claim_distribution(
+			RuntimeOrigin::signed(1),
+			0,
+			6,
+			amount_for_6,
+			proof_for_6.clone()
+		));
+		assert_eq!(Assets::balance(0, 6), 6);
+
+		// You cannot double claim.
+		assert_noop!(
+			Assets::claim_distribution(RuntimeOrigin::signed(6), 0, 6, amount_for_6, proof_for_6),
+			Error::<Test>::AlreadyClaimed
+		);
+	});
+}

From 61162b634f99122a94013d735354dd505396eef0 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Fri, 6 Sep 2024 18:25:43 -0400
Subject: [PATCH 35/55] introduce distribution info and active

---
 substrate/frame/assets/src/functions.rs | 14 +++++++++-----
 substrate/frame/assets/src/lib.rs       |  7 +++----
 substrate/frame/assets/src/types.rs     | 13 +++++++++++++
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index b8a784cca64f..39018f72cabe 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -452,11 +452,14 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 			ensure!(check_issuer == details.issuer, Error::<T, I>::NoPermission);
 		}
 
+		let info = DistributionInfo {
+			asset_id: id.clone(),
+			merkle_root: merkle_root.clone(),
+			active: true,
+		};
+
 		let distribution_count: u32 = MerklizedDistribution::<T, I>::count();
-		MerklizedDistribution::<T, I>::insert(
-			&distribution_count,
-			(id.clone(), merkle_root.clone()),
-		);
+		MerklizedDistribution::<T, I>::insert(&distribution_count, info);
 
 		Self::deposit_event(Event::DistributionIssued { asset_id: id, merkle_root });
 
@@ -475,9 +478,10 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 			return Err(Error::<T, I>::BalanceLow.into())
 		}
 
-		let (asset_id, merkle_root) =
+		let DistributionInfo { asset_id, merkle_root, active } =
 			MerklizedDistribution::<T, I>::get(distribution_id).ok_or(Error::<T, I>::Unknown)?;
 
+		ensure!(active, Error::<T, I>::DistributionEnded);
 		ensure!(
 			!MerklizedDistributionTracker::<T, I>::contains_key(distribution_id, &beneficiary),
 			Error::<T, I>::AlreadyClaimed
diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index 2c113e958969..41ad22e0c917 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -456,16 +456,13 @@ pub mod pallet {
 		ValueQuery,
 	>;
 
-	pub type DistributionCounter = u32;
-	pub type DistributionProof = Vec<Vec<u8>>;
-
 	#[pallet::storage]
 	/// Merklized distribution of an asset.
 	pub(super) type MerklizedDistribution<T: Config<I>, I: 'static = ()> = CountedStorageMap<
 		_,
 		Blake2_128Concat,
 		DistributionCounter,
-		(T::AssetId, T::Hash),
+		DistributionInfo<T::AssetId, T::Hash>,
 		OptionQuery,
 	>;
 
@@ -717,6 +714,8 @@ pub mod pallet {
 		BadAssetId,
 		/// The asset distribution was already claimed!
 		AlreadyClaimed,
+		/// THe asset distribution is no longer active.
+		DistributionEnded,
 	}
 
 	#[pallet::call(weight(<T as Config<I>>::WeightInfo))]
diff --git a/substrate/frame/assets/src/types.rs b/substrate/frame/assets/src/types.rs
index 11edc7d3fcb5..648be0150eab 100644
--- a/substrate/frame/assets/src/types.rs
+++ b/substrate/frame/assets/src/types.rs
@@ -317,3 +317,16 @@ where
 			.saturating_mul_int(balance))
 	}
 }
+
+pub type DistributionCounter = u32;
+pub type DistributionProof = Vec<Vec<u8>>;
+
+#[derive(Eq, PartialEq, Copy, Clone, RuntimeDebug, Encode, Decode, TypeInfo, MaxEncodedLen)]
+pub struct DistributionInfo<AssetId, Hash> {
+	// The asset id we are distributing.
+	pub asset_id: AssetId,
+	// The merkle root which represents all the balances to distribute.
+	pub merkle_root: Hash,
+	// Whether the distribution is still active.
+	pub active: bool,
+}

From ae2ff1811e8ad9c49ffe80ab7de394bd6b44664c Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Fri, 6 Sep 2024 18:43:09 -0400
Subject: [PATCH 36/55] introduce end_distribution

---
 substrate/frame/assets/src/functions.rs | 32 ++++++++++++++++++++++---
 substrate/frame/assets/src/lib.rs       | 31 ++++++++++++++++++++----
 substrate/frame/assets/src/weights.rs   |  9 +++++++
 3 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index 39018f72cabe..8fb1df3d8753 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -458,10 +458,14 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 			active: true,
 		};
 
-		let distribution_count: u32 = MerklizedDistribution::<T, I>::count();
-		MerklizedDistribution::<T, I>::insert(&distribution_count, info);
+		let distribution_id: u32 = MerklizedDistribution::<T, I>::count();
+		MerklizedDistribution::<T, I>::insert(&distribution_id, info);
 
-		Self::deposit_event(Event::DistributionIssued { asset_id: id, merkle_root });
+		Self::deposit_event(Event::DistributionIssued {
+			distribution_id,
+			asset_id: id,
+			merkle_root,
+		});
 
 		Ok(())
 	}
@@ -500,6 +504,28 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		Ok(())
 	}
 
+	/// Ends the asset distribution of `distribution_id`.
+	pub(super) fn do_end_distribution(
+		distribution_id: DistributionCounter,
+		maybe_check_issuer: Option<T::AccountId>,
+	) -> DispatchResult {
+		let mut info =
+			MerklizedDistribution::<T, I>::get(&distribution_id).ok_or(Error::<T, I>::Unknown)?;
+		let details = Asset::<T, I>::get(&info.asset_id).ok_or(Error::<T, I>::Unknown)?;
+
+		if let Some(check_issuer) = maybe_check_issuer {
+			ensure!(check_issuer == details.issuer, Error::<T, I>::NoPermission);
+		}
+
+		info.active = false;
+
+		MerklizedDistribution::<T, I>::insert(&distribution_id, info);
+
+		Self::deposit_event(Event::DistributionEnded { distribution_id });
+
+		Ok(())
+	}
+
 	/// Increases the asset `id` balance of `beneficiary` by `amount`.
 	///
 	/// LOW-LEVEL: Does not alter the supply of asset or emit an event. Use `do_mint` if you need
diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index 41ad22e0c917..56a3ca6686a7 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -662,7 +662,13 @@ pub mod pallet {
 		/// Some assets were withdrawn from the account (e.g. for transaction fees).
 		Withdrawn { asset_id: T::AssetId, who: T::AccountId, amount: T::Balance },
 		/// A distribution of assets were issued.
-		DistributionIssued { asset_id: T::AssetId, merkle_root: T::Hash },
+		DistributionIssued {
+			distribution_id: DistributionCounter,
+			asset_id: T::AssetId,
+			merkle_root: T::Hash,
+		},
+		/// A distribution has ended.
+		DistributionEnded { distribution_id: DistributionCounter },
 	}
 
 	#[pallet::error]
@@ -1838,7 +1844,6 @@ pub mod pallet {
 		/// Emits `DistributionIssued` event when successful.
 		///
 		/// Weight: `O(1)`
-		/// Modes: Pre-existing balance of `beneficiary`; Account pre-existence of `beneficiary`.
 		#[pallet::call_index(33)]
 		pub fn mint_distribution(
 			origin: OriginFor<T>,
@@ -1861,8 +1866,7 @@ pub mod pallet {
 		///
 		/// Emits `Issued` event when successful.
 		///
-		/// Weight: `O(1)`
-		/// Modes: Pre-existing balance of `beneficiary`; Account pre-existence of `beneficiary`.
+		/// Weight: `O(P)` where `P` is the size of the merkle proof.
 		#[pallet::call_index(34)]
 		pub fn claim_distribution(
 			origin: OriginFor<T>,
@@ -1875,6 +1879,25 @@ pub mod pallet {
 			Self::do_claim_distribution(distribution_id, beneficiary, amount, merkle_proof)?;
 			Ok(())
 		}
+
+		/// End the distribution of assets by distribution id.
+		///
+		/// The origin must be Signed and the sender must be the Issuer of the asset `id`.
+		///
+		/// - `distribution_id`: The identifier of the distribution.
+		///
+		/// Emits `DistributionEnded` event when successful.
+		///
+		/// Weight: `O(1)`
+		#[pallet::call_index(35)]
+		pub fn end_distribution(
+			origin: OriginFor<T>,
+			distribution_id: DistributionCounter,
+		) -> DispatchResult {
+			let origin = ensure_signed(origin)?;
+			Self::do_end_distribution(distribution_id, Some(origin))?;
+			Ok(())
+		}
 	}
 
 	/// Implements [`AccountTouch`] trait.
diff --git a/substrate/frame/assets/src/weights.rs b/substrate/frame/assets/src/weights.rs
index bb05c449af7d..4845ae742dd4 100644
--- a/substrate/frame/assets/src/weights.rs
+++ b/substrate/frame/assets/src/weights.rs
@@ -86,6 +86,7 @@ pub trait WeightInfo {
 	fn transfer_all() -> Weight;
 	fn mint_distribution() -> Weight;
 	fn claim_distribution() -> Weight;
+	fn end_distribution() -> Weight;
 }
 
 /// Weights for `pallet_assets` using the Substrate node and recommended hardware.
@@ -552,6 +553,10 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 		Weight::default()
 	}
 
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
+
 }
 
 // For backwards compatibility and tests.
@@ -1016,4 +1021,8 @@ impl WeightInfo for () {
 	fn claim_distribution() -> Weight {
 		Weight::default()
 	}
+
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
 }

From fb63207c8eb7202fd3d89579e771042756cb8036 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Fri, 6 Sep 2024 20:49:18 -0400
Subject: [PATCH 37/55] add clean distribution

---
 substrate/frame/assets/src/functions.rs | 37 +++++++++++++++++++++++++
 substrate/frame/assets/src/lib.rs       | 30 +++++++++++++++++++-
 substrate/frame/assets/src/weights.rs   |  9 ++++++
 3 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index 8fb1df3d8753..870d37aa0a20 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -526,6 +526,43 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		Ok(())
 	}
 
+	/// Ends the asset distribution of `distribution_id`.
+	pub(super) fn do_clean_distribution(
+		distribution_id: DistributionCounter,
+	) -> DispatchResultWithPostInfo {
+		let info =
+			MerklizedDistribution::<T, I>::get(&distribution_id).ok_or(Error::<T, I>::Unknown)?;
+
+		ensure!(!info.active, Error::<T, I>::DistributionActive);
+
+		let mut refund_count = 0u32;
+		let distribution_iterator =
+			MerklizedDistributionTracker::<T, I>::iter_key_prefix(&distribution_id);
+
+		let mut all_refunded = true;
+		for who in distribution_iterator {
+			if refund_count >= 100 {
+				// TODO: T::RemoveKeysLimit::get() {
+				// Not everyone was able to be refunded this time around.
+				all_refunded = false;
+				break
+			}
+
+			MerklizedDistributionTracker::<T, I>::remove(&distribution_id, &who);
+			refund_count += 1;
+		}
+
+		if all_refunded {
+			Self::deposit_event(Event::<T, I>::DistributionPartiallyCleaned { distribution_id });
+			// Refund weight only the amount we actually used.
+			Ok(Some(T::WeightInfo::clean_distribution(refund_count)).into())
+		} else {
+			Self::deposit_event(Event::<T, I>::DistributionCleaned { distribution_id });
+			// No weight to refund since we did not finish the loop.
+			Ok(().into())
+		}
+	}
+
 	/// Increases the asset `id` balance of `beneficiary` by `amount`.
 	///
 	/// LOW-LEVEL: Does not alter the supply of asset or emit an event. Use `do_mint` if you need
diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index 56a3ca6686a7..bf4c00bc6eb4 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -669,6 +669,10 @@ pub mod pallet {
 		},
 		/// A distribution has ended.
 		DistributionEnded { distribution_id: DistributionCounter },
+		/// A distribution has been partially cleaned. There are still more items to clean up.
+		DistributionPartiallyCleaned { distribution_id: DistributionCounter },
+		/// A distribution has been fully cleaned.
+		DistributionCleaned { distribution_id: DistributionCounter },
 	}
 
 	#[pallet::error]
@@ -720,8 +724,10 @@ pub mod pallet {
 		BadAssetId,
 		/// The asset distribution was already claimed!
 		AlreadyClaimed,
-		/// THe asset distribution is no longer active.
+		/// The asset distribution is no longer active.
 		DistributionEnded,
+		/// The asset distribution is still active.
+		DistributionActive,
 	}
 
 	#[pallet::call(weight(<T as Config<I>>::WeightInfo))]
@@ -1898,6 +1904,28 @@ pub mod pallet {
 			Self::do_end_distribution(distribution_id, Some(origin))?;
 			Ok(())
 		}
+
+		/// Clean up the distribution tracker of an ended distribution. This function might need to
+		/// be called multiple times to remove all the items from the distribution tracker.
+		///
+		/// Any signed origin may call this function.
+		///
+		/// - `distribution_id`: The identifier of the distribution to clean. It cannot be active.
+		///
+		/// Emits `DistributionPartiallyCleaned` event when some elements have been removed, but
+		/// there are still some left. Emits `DistributionCleaned` when all of the distribution
+		/// history has been removed.
+		///
+		/// Weight: `O(N)` where `N` is the maximum number of elements that can be removed at once.
+		#[pallet::call_index(36)]
+		#[pallet::weight(T::WeightInfo::clean_distribution(100u32))] // TODO
+		pub fn clean_distribution(
+			origin: OriginFor<T>,
+			distribution_id: DistributionCounter,
+		) -> DispatchResultWithPostInfo {
+			ensure_signed(origin)?;
+			Self::do_clean_distribution(distribution_id)
+		}
 	}
 
 	/// Implements [`AccountTouch`] trait.
diff --git a/substrate/frame/assets/src/weights.rs b/substrate/frame/assets/src/weights.rs
index 4845ae742dd4..ac01d5d412a1 100644
--- a/substrate/frame/assets/src/weights.rs
+++ b/substrate/frame/assets/src/weights.rs
@@ -87,6 +87,7 @@ pub trait WeightInfo {
 	fn mint_distribution() -> Weight;
 	fn claim_distribution() -> Weight;
 	fn end_distribution() -> Weight;
+	fn clean_distribution(n: u32) -> Weight;
 }
 
 /// Weights for `pallet_assets` using the Substrate node and recommended hardware.
@@ -557,6 +558,10 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 		Weight::default()
 	}
 
+	fn clean_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
+
 }
 
 // For backwards compatibility and tests.
@@ -1025,4 +1030,8 @@ impl WeightInfo for () {
 	fn end_distribution() -> Weight {
 		Weight::default()
 	}
+
+	fn clean_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
 }

From 988b201e11d90f5e0719d9c26ba9da85bc302cc8 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Fri, 6 Sep 2024 22:01:37 -0400
Subject: [PATCH 38/55] weight stuff

---
 .../src/weights/pallet_assets_foreign.rs         | 16 ++++++++++++++++
 .../src/weights/pallet_assets_local.rs           | 16 ++++++++++++++++
 .../src/weights/pallet_assets_pool.rs            | 16 ++++++++++++++++
 .../src/weights/pallet_assets_foreign.rs         | 16 ++++++++++++++++
 .../src/weights/pallet_assets_local.rs           | 16 ++++++++++++++++
 .../src/weights/pallet_assets_pool.rs            | 16 ++++++++++++++++
 6 files changed, 96 insertions(+)

diff --git a/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_foreign.rs b/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_foreign.rs
index c76c1137335a..f099987cabc2 100644
--- a/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_foreign.rs
+++ b/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_foreign.rs
@@ -541,4 +541,20 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn clean_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
 }
diff --git a/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_local.rs b/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_local.rs
index cf4f60042bc6..5e2e9e421fc5 100644
--- a/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_local.rs
+++ b/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_local.rs
@@ -538,4 +538,20 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn clean_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
 }
diff --git a/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_pool.rs b/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_pool.rs
index 2cd85de00989..dabb5a103ac5 100644
--- a/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_pool.rs
+++ b/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_pool.rs
@@ -538,4 +538,20 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn clean_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
 }
diff --git a/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_foreign.rs b/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_foreign.rs
index 2692de9aeb50..60829e471df7 100644
--- a/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_foreign.rs
+++ b/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_foreign.rs
@@ -547,4 +547,20 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn clean_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
 }
diff --git a/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_local.rs b/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_local.rs
index d2e12549a45c..a25eac274fb0 100644
--- a/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_local.rs
+++ b/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_local.rs
@@ -545,4 +545,20 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn clean_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
 }
diff --git a/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_pool.rs b/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_pool.rs
index 8368f6e583cc..7676ffa5aaab 100644
--- a/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_pool.rs
+++ b/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_pool.rs
@@ -539,4 +539,20 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
+
+	fn mint_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn claim_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn end_distribution() -> Weight {
+		Weight::default()
+	}
+
+	fn clean_distribution(_n: u32) -> Weight {
+		Weight::default()
+	}
 }

From 233d33989b13d3207f8b3535b80ae8e752da3474 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Sat, 7 Sep 2024 13:02:29 -0400
Subject: [PATCH 39/55] temp fix for ui

---
 substrate/frame/assets/src/lib.rs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index bf4c00bc6eb4..8d37959bddbf 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -1879,10 +1879,11 @@ pub mod pallet {
 			distribution_id: DistributionCounter,
 			beneficiary: T::AccountId,
 			amount: T::Balance,
-			merkle_proof: DistributionProof,
+			merkle_proof: Vec<u8>, // TODO temp
 		) -> DispatchResult {
 			ensure_signed(origin)?;
-			Self::do_claim_distribution(distribution_id, beneficiary, amount, merkle_proof)?;
+			let proof: DistributionProof = Decode::decode(&mut &merkle_proof[..]).map_err(|_| Error::<T, I>::BadAssetId)?;
+			Self::do_claim_distribution(distribution_id, beneficiary, amount, proof)?;
 			Ok(())
 		}
 

From 122f9cc733b0fe69e224501ae1dbdfb4be174db4 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 16 Sep 2024 20:52:56 -0400
Subject: [PATCH 40/55] make work with binary tree

---
 Cargo.lock                                    |  1 +
 substrate/frame/assets/Cargo.toml             |  1 +
 substrate/frame/assets/src/functions.rs       | 27 +++++----
 substrate/frame/assets/src/lib.rs             | 18 ++++--
 substrate/frame/assets/src/tests.rs           | 60 +++++++++----------
 substrate/frame/assets/src/types.rs           |  3 +-
 substrate/frame/support/src/traits/proving.rs | 18 +++++-
 7 files changed, 73 insertions(+), 55 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1478a543d14c..db41718cb753 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -9857,6 +9857,7 @@ dependencies = [
 name = "pallet-assets"
 version = "29.1.0"
 dependencies = [
+ "binary-merkle-tree",
  "frame-benchmarking",
  "frame-support",
  "frame-system",
diff --git a/substrate/frame/assets/Cargo.toml b/substrate/frame/assets/Cargo.toml
index e20b576d0836..5f1a876a5dc3 100644
--- a/substrate/frame/assets/Cargo.toml
+++ b/substrate/frame/assets/Cargo.toml
@@ -32,6 +32,7 @@ sp-core = { workspace = true }
 [dev-dependencies]
 sp-io = { workspace = true, default-features = true }
 pallet-balances = { workspace = true, default-features = true }
+binary-merkle-tree = { workspace = true }
 
 [features]
 default = ["std"]
diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index 870d37aa0a20..02822ad46f64 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -474,29 +474,32 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 	/// whom.
 	pub(super) fn do_claim_distribution(
 		distribution_id: DistributionCounter,
-		beneficiary: T::AccountId,
-		amount: T::Balance,
-		merkle_proof: DistributionProof,
+		merkle_proof: Vec<u8>,
 	) -> DispatchResult {
-		if amount.is_zero() {
-			return Err(Error::<T, I>::BalanceLow.into())
-		}
+		let proof =
+			codec::Decode::decode(&mut &merkle_proof[..]).map_err(|_| Error::<T, I>::BadProof)?;
 
 		let DistributionInfo { asset_id, merkle_root, active } =
 			MerklizedDistribution::<T, I>::get(distribution_id).ok_or(Error::<T, I>::Unknown)?;
 
 		ensure!(active, Error::<T, I>::DistributionEnded);
+
+		let leaf = T::VerifyExistenceProof::verify_proof(proof, &merkle_root)
+			.map_err(|()| Error::<T, I>::BadProof)?;
+		let (beneficiary, amount) =
+			codec::Decode::decode(&mut &leaf[..]).map_err(|_| Error::<T, I>::CannotDecodeLeaf)?;
+
 		ensure!(
 			!MerklizedDistributionTracker::<T, I>::contains_key(distribution_id, &beneficiary),
 			Error::<T, I>::AlreadyClaimed
 		);
 
-		sp_runtime::proving_trie::verify_single_value_proof::<T::Hashing, _, _>(
-			merkle_root,
-			&merkle_proof,
-			&beneficiary,
-			Some(amount),
-		)?;
+		// sp_runtime::proving_trie::verify_single_value_proof::<T::Hashing, _, _>(
+		// 	merkle_root,
+		// 	&merkle_proof,
+		// 	&beneficiary,
+		// 	Some(amount),
+		// )?;
 
 		Self::do_mint(asset_id, &beneficiary, amount, None)?;
 		MerklizedDistributionTracker::<T, I>::insert(&distribution_id, &beneficiary, ());
diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index 8d37959bddbf..2f8e1a78a1ae 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -189,7 +189,8 @@ use frame_support::{
 			WithdrawConsequence,
 		},
 		BalanceStatus::Reserved,
-		Currency, EnsureOriginWithArg, Incrementable, ReservableCurrency, StoredMap,
+		BinaryMerkleTreeProver, Currency, EnsureOriginWithArg, Incrementable, ReservableCurrency,
+		StoredMap, VerifyExistenceProof,
 	},
 };
 use frame_system::Config as SystemConfig;
@@ -298,6 +299,7 @@ pub mod pallet {
 			type Extra = ();
 			type CallbackHandle = ();
 			type WeightInfo = ();
+			type VerifyExistenceProof = BinaryMerkleTreeProver<Self::Hashing>;
 			#[cfg(feature = "runtime-benchmarks")]
 			type BenchmarkHelper = ();
 		}
@@ -404,6 +406,9 @@ pub mod pallet {
 		/// used to set up auto-incrementing asset IDs for this collection.
 		type CallbackHandle: AssetsCallback<Self::AssetId, Self::AccountId>;
 
+		/// A type used to verify merkle proofs used for distributions.
+		type VerifyExistenceProof: VerifyExistenceProof<Hash = Self::Hash>;
+
 		/// Weight information for extrinsics in this pallet.
 		type WeightInfo: WeightInfo;
 
@@ -728,6 +733,10 @@ pub mod pallet {
 		DistributionEnded,
 		/// The asset distribution is still active.
 		DistributionActive,
+		/// The proof provided could not be verified.
+		BadProof,
+		/// The a leaf node was extracted from the proof, but it did not match the expected format.
+		CannotDecodeLeaf,
 	}
 
 	#[pallet::call(weight(<T as Config<I>>::WeightInfo))]
@@ -1877,13 +1886,10 @@ pub mod pallet {
 		pub fn claim_distribution(
 			origin: OriginFor<T>,
 			distribution_id: DistributionCounter,
-			beneficiary: T::AccountId,
-			amount: T::Balance,
-			merkle_proof: Vec<u8>, // TODO temp
+			merkle_proof: Vec<u8>,
 		) -> DispatchResult {
 			ensure_signed(origin)?;
-			let proof: DistributionProof = Decode::decode(&mut &merkle_proof[..]).map_err(|_| Error::<T, I>::BadAssetId)?;
-			Self::do_claim_distribution(distribution_id, beneficiary, amount, proof)?;
+			Self::do_claim_distribution(distribution_id, merkle_proof)?;
 			Ok(())
 		}
 
diff --git a/substrate/frame/assets/src/tests.rs b/substrate/frame/assets/src/tests.rs
index 93dd59921691..21a98a958910 100644
--- a/substrate/frame/assets/src/tests.rs
+++ b/substrate/frame/assets/src/tests.rs
@@ -19,6 +19,7 @@
 
 use super::*;
 use crate::{mock::*, Error};
+use codec::Encode;
 use frame_support::{
 	assert_noop, assert_ok,
 	dispatch::GetDispatchInfo,
@@ -1926,14 +1927,11 @@ fn asset_id_cannot_be_reused() {
 fn merklized_distribution_works() {
 	new_test_ext().execute_with(|| {
 		use alloc::collections::BTreeMap;
-		use sp_runtime::proving_trie::BasicProvingTrie;
 
 		// Create asset id 0 controlled by user 1, sufficient so it does not need ED.
 		assert_ok!(Assets::force_create(RuntimeOrigin::root(), 0, 1, true, 1));
 
 		// Offchain, user 1 creates a distribution of tokens.
-		type DistributionTrie =
-			BasicProvingTrie<<Test as frame_system::Config>::Hashing, AccountId, Balance>;
 		let mut distribution = BTreeMap::<AccountId, Balance>::new();
 		for i in 0..100u64 {
 			distribution.insert(i, i.into());
@@ -1942,51 +1940,47 @@ fn merklized_distribution_works() {
 		// Maybe the owner gives himself a little extra ;)
 		distribution.insert(1, 1337);
 
-		let distribution_trie = DistributionTrie::generate_for(distribution).unwrap();
-		let root = *distribution_trie.root();
+		let flat_distribution: Vec<Vec<u8>> =
+			distribution.into_iter().map(|item| item.encode()).collect();
+
+		let root = binary_merkle_tree::merkle_root::<<Test as frame_system::Config>::Hashing, _>(
+			flat_distribution.clone(),
+		);
+
+		let proof_for_69 = binary_merkle_tree::merkle_proof::<
+			<Test as frame_system::Config>::Hashing,
+			_,
+			_,
+		>(flat_distribution.clone(), 69);
+		let proof_for_1 = binary_merkle_tree::merkle_proof::<
+			<Test as frame_system::Config>::Hashing,
+			_,
+			_,
+		>(flat_distribution.clone(), 1);
+		let proof_for_6 = binary_merkle_tree::merkle_proof::<
+			<Test as frame_system::Config>::Hashing,
+			_,
+			_,
+		>(flat_distribution, 6);
 
 		// Use this trie root for the distribution
 		assert_ok!(Assets::mint_distribution(RuntimeOrigin::signed(1), 0, root));
 
 		// Now users claim their distributions permissionlessly with a proof.
-		let proof_for_1 = distribution_trie.create_single_value_proof(1).unwrap();
-		let amount_for_1 = distribution_trie.query(1).unwrap();
-		assert_ok!(Assets::claim_distribution(
-			RuntimeOrigin::signed(1),
-			0,
-			1,
-			amount_for_1,
-			proof_for_1
-		));
+		assert_ok!(Assets::claim_distribution(RuntimeOrigin::signed(1), 0, proof_for_1.encode()));
 		assert_eq!(Assets::balance(0, 1), 1337);
 
 		// Other users can claim their tokens.
-		let proof_for_69 = distribution_trie.create_single_value_proof(69).unwrap();
-		let amount_for_69 = distribution_trie.query(69).unwrap();
-		assert_ok!(Assets::claim_distribution(
-			RuntimeOrigin::signed(55),
-			0,
-			69,
-			amount_for_69,
-			proof_for_69
-		));
+		assert_ok!(Assets::claim_distribution(RuntimeOrigin::signed(55), 0, proof_for_69.encode()));
 		assert_eq!(Assets::balance(0, 69), 69);
 
 		// Owner (or anyone) can also distribute on behalf of the other users.
-		let proof_for_6 = distribution_trie.create_single_value_proof(6).unwrap();
-		let amount_for_6 = distribution_trie.query(6).unwrap();
-		assert_ok!(Assets::claim_distribution(
-			RuntimeOrigin::signed(1),
-			0,
-			6,
-			amount_for_6,
-			proof_for_6.clone()
-		));
+		assert_ok!(Assets::claim_distribution(RuntimeOrigin::signed(1), 0, proof_for_6.encode()));
 		assert_eq!(Assets::balance(0, 6), 6);
 
 		// You cannot double claim.
 		assert_noop!(
-			Assets::claim_distribution(RuntimeOrigin::signed(6), 0, 6, amount_for_6, proof_for_6),
+			Assets::claim_distribution(RuntimeOrigin::signed(6), 0, proof_for_6.encode()),
 			Error::<Test>::AlreadyClaimed
 		);
 	});
diff --git a/substrate/frame/assets/src/types.rs b/substrate/frame/assets/src/types.rs
index 648be0150eab..44b09652d46b 100644
--- a/substrate/frame/assets/src/types.rs
+++ b/substrate/frame/assets/src/types.rs
@@ -319,7 +319,8 @@ where
 }
 
 pub type DistributionCounter = u32;
-pub type DistributionProof = Vec<Vec<u8>>;
+pub type DistributionProofOf<T, I> =
+	<<T as Config<I>>::VerifyExistenceProof as VerifyExistenceProof>::Proof;
 
 #[derive(Eq, PartialEq, Copy, Clone, RuntimeDebug, Encode, Decode, TypeInfo, MaxEncodedLen)]
 pub struct DistributionInfo<AssetId, Hash> {
diff --git a/substrate/frame/support/src/traits/proving.rs b/substrate/frame/support/src/traits/proving.rs
index dc44f4cd68e7..549f8f719ba4 100644
--- a/substrate/frame/support/src/traits/proving.rs
+++ b/substrate/frame/support/src/traits/proving.rs
@@ -24,9 +24,9 @@ use sp_core::Hasher;
 /// Something that can verify the existence of some data in a given proof.
 pub trait VerifyExistenceProof {
 	/// The proof type.
-	type Proof;
+	type Proof: Encode + Decode;
 	/// The hash type.
-	type Hash;
+	type Hash: Encode + Decode;
 
 	/// Verify the given `proof`.
 	///
@@ -77,7 +77,10 @@ pub struct SixteenPatriciaMerkleTreeExistenceProof {
 /// Implements [`VerifyExistenceProof`] using a 16-patricia merkle tree.
 pub struct SixteenPatriciaMerkleTreeProver<H>(core::marker::PhantomData<H>);
 
-impl<H: Hasher> VerifyExistenceProof for SixteenPatriciaMerkleTreeProver<H> {
+impl<H: Hasher> VerifyExistenceProof for SixteenPatriciaMerkleTreeProver<H>
+where
+	H::Out: Decode + Encode,
+{
 	type Proof = SixteenPatriciaMerkleTreeExistenceProof;
 	type Hash = H::Out;
 
@@ -92,6 +95,15 @@ impl<H: Hasher> VerifyExistenceProof for SixteenPatriciaMerkleTreeProver<H> {
 	}
 }
 
+/// An implementation which always returns an error when this feature is unavailable.
+impl VerifyExistenceProof for () {
+	type Proof = ();
+	type Hash = ();
+	fn verify_proof(_proof: Self::Proof, _root: &Self::Hash) -> Result<Vec<u8>, ()> {
+		Err(())
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;

From 59fdc634ec5234431b741a8352147a14144c18d3 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 16 Sep 2024 21:20:34 -0400
Subject: [PATCH 41/55] initial stuff

---
 Cargo.lock                                    |   1 +
 substrate/primitives/runtime/Cargo.toml       |   2 +
 .../primitives/runtime/src/proving_trie.rs    |   2 +
 .../runtime/src/proving_trie/binary.rs        | 271 ++++++++++++++++++
 4 files changed, 276 insertions(+)
 create mode 100644 substrate/primitives/runtime/src/proving_trie/binary.rs

diff --git a/Cargo.lock b/Cargo.lock
index db41718cb753..f68b848f82f2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -20630,6 +20630,7 @@ dependencies = [
 name = "sp-runtime"
 version = "31.0.1"
 dependencies = [
+ "binary-merkle-tree",
  "docify",
  "either",
  "hash256-std-hasher",
diff --git a/substrate/primitives/runtime/Cargo.toml b/substrate/primitives/runtime/Cargo.toml
index 800bf4bd0737..eebd3a2b3322 100644
--- a/substrate/primitives/runtime/Cargo.toml
+++ b/substrate/primitives/runtime/Cargo.toml
@@ -36,6 +36,7 @@ sp-trie = { workspace = true }
 sp-weights = { workspace = true }
 docify = { workspace = true }
 tracing = { workspace = true, features = ["log"], default-features = false }
+binary-merkle-tree = { workspace = true }
 
 simple-mermaid = { version = "0.1.1", optional = true }
 
@@ -73,6 +74,7 @@ std = [
 	"sp-trie/std",
 	"sp-weights/std",
 	"tracing/std",
+	"binary-merkle-tree/std",
 ]
 
 # Serde support without relying on std features.
diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie.rs
index 9a423f18284f..0043f80d7ef7 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie.rs
@@ -28,6 +28,8 @@ use crate::{Decode, DispatchError, Encode, MaxEncodedLen, TypeInfo};
 #[cfg(feature = "serde")]
 use crate::{Deserialize, Serialize};
 
+mod binary;
+
 use sp_std::vec::Vec;
 use sp_trie::{
 	trie_types::{TrieDBBuilder, TrieDBMutBuilderV1, TrieError as SpTrieError},
diff --git a/substrate/primitives/runtime/src/proving_trie/binary.rs b/substrate/primitives/runtime/src/proving_trie/binary.rs
new file mode 100644
index 000000000000..2524ca97c255
--- /dev/null
+++ b/substrate/primitives/runtime/src/proving_trie/binary.rs
@@ -0,0 +1,271 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Types for a compact base-16 merkle trie used for checking and generating proofs within the
+//! runtime. The `sp-trie` crate exposes all of these same functionality (and more), but this
+//! library is designed to work more easily with runtime native types, which simply need to
+//! implement `Encode`/`Decode`. It also exposes a runtime friendly `TrieError` type which can be
+//! use inside of a FRAME Pallet.
+//!
+//! Proofs are created with latest substrate trie format (`LayoutV1`), and are not compatible with
+//! proofs using `LayoutV0`.
+
+use crate::{Decode, DispatchError, Encode, MaxEncodedLen, TypeInfo};
+#[cfg(feature = "serde")]
+use crate::{Deserialize, Serialize};
+
+use super::*;
+use sp_std::{collections::btree_map::BTreeMap, vec::Vec};
+use sp_trie::{
+	trie_types::{TrieDBBuilder, TrieDBMutBuilderV1, TrieError as SpTrieError},
+	LayoutV1, MemoryDB, Trie, TrieMut, VerifyError,
+};
+
+type HashOf<Hashing> = <Hashing as sp_core::Hasher>::Out;
+
+/// A helper structure for building a basic base-16 merkle trie and creating compact proofs for that
+/// trie. Proofs are created with latest substrate trie format (`LayoutV1`), and are not compatible
+/// with proofs using `LayoutV0`.
+pub struct BasicProvingTrie<Hashing, Key, Value>
+where
+	Hashing: sp_core::Hasher,
+{
+	db: Vec<(Key, Value)>,
+	root: HashOf<Hashing>,
+	_phantom: core::marker::PhantomData<(Key, Value)>,
+}
+
+impl<Hashing, Key, Value> BasicProvingTrie<Hashing, Key, Value>
+where
+	Hashing: sp_core::Hasher,
+	Key: Encode + Ord,
+	Value: Encode,
+{
+	/// Create a new instance of a `ProvingTrie` using an iterator of key/value pairs.
+	pub fn generate_for<I>(items: I) -> Result<Self, DispatchError>
+	where
+		I: IntoIterator<Item = (Key, Value)>,
+	{
+		let mut db_map = BTreeMap::default();
+		for (key, value) in items.into_iter() {
+			db_map.insert(key, value);
+		}
+
+		let db_map_len = db_map.len();
+
+		let db: Vec<(Key, Value)> = db_map.into_iter().collect();
+
+		if db.len() != db_map_len {
+			return Err("duplicate item key".into())
+		}
+
+		let root =
+			binary_merkle_tree::merkle_root::<Hashing, _>(db.iter().map(|item| item.encode()));
+
+		Ok(Self { db, root, _phantom: Default::default() })
+	}
+
+	/// Access the underlying trie root.
+	pub fn root(&self) -> &HashOf<Hashing> {
+		&self.root
+	}
+
+	/// Query a value contained within the current trie. Returns `None` if the
+	/// nodes within the current `MemoryDB` are insufficient to query the item.
+	pub fn query(&self, key: Key) -> Option<Value>
+	where
+		Value: Decode,
+	{
+		unimplemented!()
+
+		// let trie = TrieDBBuilder::new(&self.db, &self.root).build();
+		// key.using_encoded(|s| trie.get(s))
+		// 	.ok()?
+		// 	.and_then(|raw| Value::decode(&mut &*raw).ok())
+	}
+
+	/// Create a compact merkle proof needed to prove a single key and its value are in the trie.
+	/// Returns `None` if the nodes within the current `MemoryDB` are insufficient to create a
+	/// proof.
+	///
+	/// This function makes a proof with latest substrate trie format (`LayoutV1`), and is not
+	/// compatible with `LayoutV0`.
+	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<Vec<u8>>, DispatchError> {
+		unimplemented!()
+	}
+}
+
+/// Verify the existence or non-existence of `key` and `value` in a given trie root and proof.
+///
+/// Proofs must be created with latest substrate trie format (`LayoutV1`).
+pub fn verify_single_value_proof<Hashing, Key, Value>(
+	root: HashOf<Hashing>,
+	proof: &[Vec<u8>],
+	key: Key,
+	maybe_value: Option<Value>,
+) -> Result<(), DispatchError>
+where
+	Hashing: sp_core::Hasher,
+	Key: Encode,
+	Value: Encode,
+{
+	sp_trie::verify_trie_proof::<LayoutV1<Hashing>, _, _, _>(
+		&root,
+		proof,
+		&[(key.encode(), maybe_value.map(|value| value.encode()))],
+	)
+	.map_err(|err| TrieError::from(err).into())
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use crate::traits::BlakeTwo256;
+	use sp_core::H256;
+	use sp_std::collections::btree_map::BTreeMap;
+
+	// A trie which simulates a trie of accounts (u32) and balances (u128).
+	type BalanceTrie = BasicProvingTrie<BlakeTwo256, u32, u128>;
+
+	// The expected root hash for an empty trie.
+	fn empty_root() -> H256 {
+		sp_trie::empty_trie_root::<LayoutV1<BlakeTwo256>>()
+	}
+
+	fn create_balance_trie() -> BalanceTrie {
+		// Create a map of users and their balances.
+		let mut map = BTreeMap::<u32, u128>::new();
+		for i in 0..100u32 {
+			map.insert(i, i.into());
+		}
+
+		// Put items into the trie.
+		let balance_trie = BalanceTrie::generate_for(map).unwrap();
+
+		// Root is changed.
+		let root = *balance_trie.root();
+		assert!(root != empty_root());
+
+		// Assert valid keys are queryable.
+		assert_eq!(balance_trie.query(6u32), Some(6u128));
+		assert_eq!(balance_trie.query(9u32), Some(9u128));
+		assert_eq!(balance_trie.query(69u32), Some(69u128));
+		// Invalid key returns none.
+		assert_eq!(balance_trie.query(6969u32), None);
+
+		balance_trie
+	}
+
+	#[test]
+	fn empty_trie_works() {
+		let empty_trie = BalanceTrie::generate_for(Vec::new()).unwrap();
+		assert_eq!(*empty_trie.root(), empty_root());
+	}
+
+	#[test]
+	fn basic_end_to_end_single_value() {
+		let balance_trie = create_balance_trie();
+		let root = *balance_trie.root();
+
+		// Create a proof for a valid key.
+		let proof = balance_trie.create_single_value_proof(6u32).unwrap();
+
+		// Assert key is provable, all other keys are invalid.
+		for i in 0..200u32 {
+			if i == 6 {
+				assert_eq!(
+					verify_single_value_proof::<BlakeTwo256, _, _>(
+						root,
+						&proof,
+						i,
+						Some(u128::from(i))
+					),
+					Ok(())
+				);
+				// Wrong value is invalid.
+				assert_eq!(
+					verify_single_value_proof::<BlakeTwo256, _, _>(
+						root,
+						&proof,
+						i,
+						Some(u128::from(i + 1))
+					),
+					Err(TrieError::RootMismatch.into())
+				);
+			} else {
+				assert!(verify_single_value_proof::<BlakeTwo256, _, _>(
+					root,
+					&proof,
+					i,
+					Some(u128::from(i))
+				)
+				.is_err());
+				assert!(verify_single_value_proof::<BlakeTwo256, _, _>(
+					root,
+					&proof,
+					i,
+					None::<u128>
+				)
+				.is_err());
+			}
+		}
+	}
+
+	#[test]
+	fn basic_end_to_end_multi_value() {
+		let balance_trie = create_balance_trie();
+		let root = *balance_trie.root();
+
+		// Create a proof for a valid and invalid key.
+		let proof = balance_trie.create_proof(&[6u32, 69u32, 6969u32]).unwrap();
+		let items = [(6u32, Some(6u128)), (69u32, Some(69u128)), (6969u32, None)];
+
+		assert_eq!(verify_proof::<BlakeTwo256, _, _>(root, &proof, &items), Ok(()));
+	}
+
+	#[test]
+	fn proof_fails_with_bad_data() {
+		let balance_trie = create_balance_trie();
+		let root = *balance_trie.root();
+
+		// Create a proof for a valid key.
+		let proof = balance_trie.create_single_value_proof(6u32).unwrap();
+
+		// Correct data verifies successfully
+		assert_eq!(
+			verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, Some(6u128)),
+			Ok(())
+		);
+
+		// Fail to verify proof with wrong root
+		assert_eq!(
+			verify_single_value_proof::<BlakeTwo256, _, _>(
+				Default::default(),
+				&proof,
+				6u32,
+				Some(6u128)
+			),
+			Err(TrieError::RootMismatch.into())
+		);
+
+		// Fail to verify proof with wrong data
+		assert_eq!(
+			verify_single_value_proof::<BlakeTwo256, _, _>(root, &[], 6u32, Some(6u128)),
+			Err(TrieError::IncompleteProof.into())
+		);
+	}
+}

From 1f587ef7b45436d4073d1c6b2f0f4b7befa19eae Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 16 Sep 2024 22:12:01 -0400
Subject: [PATCH 42/55] more

---
 .../runtime/src/proving_trie/binary.rs        | 34 ++++++++++++++-----
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie/binary.rs b/substrate/primitives/runtime/src/proving_trie/binary.rs
index 2524ca97c255..4b5548dbbdc4 100644
--- a/substrate/primitives/runtime/src/proving_trie/binary.rs
+++ b/substrate/primitives/runtime/src/proving_trie/binary.rs
@@ -44,6 +44,7 @@ pub struct BasicProvingTrie<Hashing, Key, Value>
 where
 	Hashing: sp_core::Hasher,
 {
+	// Deduplicated and flattened list of key value pairs.
 	db: Vec<(Key, Value)>,
 	root: HashOf<Hashing>,
 	_phantom: core::marker::PhantomData<(Key, Value)>,
@@ -88,14 +89,9 @@ where
 	/// nodes within the current `MemoryDB` are insufficient to query the item.
 	pub fn query(&self, key: Key) -> Option<Value>
 	where
-		Value: Decode,
+		Value: Decode + Clone,
 	{
-		unimplemented!()
-
-		// let trie = TrieDBBuilder::new(&self.db, &self.root).build();
-		// key.using_encoded(|s| trie.get(s))
-		// 	.ok()?
-		// 	.and_then(|raw| Value::decode(&mut &*raw).ok())
+		self.db.iter().find(|(k, _)| *k == key).map(|(_, v)| v.clone())
 	}
 
 	/// Create a compact merkle proof needed to prove a single key and its value are in the trie.
@@ -104,8 +100,28 @@ where
 	///
 	/// This function makes a proof with latest substrate trie format (`LayoutV1`), and is not
 	/// compatible with `LayoutV0`.
-	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<Vec<u8>>, DispatchError> {
-		unimplemented!()
+	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError> {
+		let mut encoded = Vec::with_capacity(self.db.len()); // Pre-allocate the vector
+		let mut found_index = None;
+
+		// Find the index of our key, and encode the (key, value) pair.
+		for (i, (k, v)) in self.db.into_iter().enumerate() {
+			// If we found the key we are looking for, save it.
+			if k == key {
+				found_index = Some(i);
+			}
+
+			encoded.push((k, v).encode());
+		}
+
+		let index = found_index.ok_or("couldnt find")?;
+
+		let proof = binary_merkle_tree::merkle_proof::<
+			Hashing,
+			_,
+			_,
+		>(encoded, index as u32);
+		Ok(Encode::encode(&proof))
 	}
 }
 

From f3a1af025a4d430e775783381bab6d0d6193b9a6 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 16 Sep 2024 21:20:34 -0400
Subject: [PATCH 43/55] initial stuff

---
 .../runtime/src/proving_trie/binary.rs        | 34 +++++--------------
 1 file changed, 9 insertions(+), 25 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie/binary.rs b/substrate/primitives/runtime/src/proving_trie/binary.rs
index 4b5548dbbdc4..2524ca97c255 100644
--- a/substrate/primitives/runtime/src/proving_trie/binary.rs
+++ b/substrate/primitives/runtime/src/proving_trie/binary.rs
@@ -44,7 +44,6 @@ pub struct BasicProvingTrie<Hashing, Key, Value>
 where
 	Hashing: sp_core::Hasher,
 {
-	// Deduplicated and flattened list of key value pairs.
 	db: Vec<(Key, Value)>,
 	root: HashOf<Hashing>,
 	_phantom: core::marker::PhantomData<(Key, Value)>,
@@ -89,9 +88,14 @@ where
 	/// nodes within the current `MemoryDB` are insufficient to query the item.
 	pub fn query(&self, key: Key) -> Option<Value>
 	where
-		Value: Decode + Clone,
+		Value: Decode,
 	{
-		self.db.iter().find(|(k, _)| *k == key).map(|(_, v)| v.clone())
+		unimplemented!()
+
+		// let trie = TrieDBBuilder::new(&self.db, &self.root).build();
+		// key.using_encoded(|s| trie.get(s))
+		// 	.ok()?
+		// 	.and_then(|raw| Value::decode(&mut &*raw).ok())
 	}
 
 	/// Create a compact merkle proof needed to prove a single key and its value are in the trie.
@@ -100,28 +104,8 @@ where
 	///
 	/// This function makes a proof with latest substrate trie format (`LayoutV1`), and is not
 	/// compatible with `LayoutV0`.
-	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError> {
-		let mut encoded = Vec::with_capacity(self.db.len()); // Pre-allocate the vector
-		let mut found_index = None;
-
-		// Find the index of our key, and encode the (key, value) pair.
-		for (i, (k, v)) in self.db.into_iter().enumerate() {
-			// If we found the key we are looking for, save it.
-			if k == key {
-				found_index = Some(i);
-			}
-
-			encoded.push((k, v).encode());
-		}
-
-		let index = found_index.ok_or("couldnt find")?;
-
-		let proof = binary_merkle_tree::merkle_proof::<
-			Hashing,
-			_,
-			_,
-		>(encoded, index as u32);
-		Ok(Encode::encode(&proof))
+	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<Vec<u8>>, DispatchError> {
+		unimplemented!()
 	}
 }
 

From daf8f5072027a89edfa74313c700431bc0e122f2 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Mon, 16 Sep 2024 22:12:01 -0400
Subject: [PATCH 44/55] more

---
 .../runtime/src/proving_trie/binary.rs        | 34 ++++++++++++++-----
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie/binary.rs b/substrate/primitives/runtime/src/proving_trie/binary.rs
index 2524ca97c255..4b5548dbbdc4 100644
--- a/substrate/primitives/runtime/src/proving_trie/binary.rs
+++ b/substrate/primitives/runtime/src/proving_trie/binary.rs
@@ -44,6 +44,7 @@ pub struct BasicProvingTrie<Hashing, Key, Value>
 where
 	Hashing: sp_core::Hasher,
 {
+	// Deduplicated and flattened list of key value pairs.
 	db: Vec<(Key, Value)>,
 	root: HashOf<Hashing>,
 	_phantom: core::marker::PhantomData<(Key, Value)>,
@@ -88,14 +89,9 @@ where
 	/// nodes within the current `MemoryDB` are insufficient to query the item.
 	pub fn query(&self, key: Key) -> Option<Value>
 	where
-		Value: Decode,
+		Value: Decode + Clone,
 	{
-		unimplemented!()
-
-		// let trie = TrieDBBuilder::new(&self.db, &self.root).build();
-		// key.using_encoded(|s| trie.get(s))
-		// 	.ok()?
-		// 	.and_then(|raw| Value::decode(&mut &*raw).ok())
+		self.db.iter().find(|(k, _)| *k == key).map(|(_, v)| v.clone())
 	}
 
 	/// Create a compact merkle proof needed to prove a single key and its value are in the trie.
@@ -104,8 +100,28 @@ where
 	///
 	/// This function makes a proof with latest substrate trie format (`LayoutV1`), and is not
 	/// compatible with `LayoutV0`.
-	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<Vec<u8>>, DispatchError> {
-		unimplemented!()
+	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError> {
+		let mut encoded = Vec::with_capacity(self.db.len()); // Pre-allocate the vector
+		let mut found_index = None;
+
+		// Find the index of our key, and encode the (key, value) pair.
+		for (i, (k, v)) in self.db.into_iter().enumerate() {
+			// If we found the key we are looking for, save it.
+			if k == key {
+				found_index = Some(i);
+			}
+
+			encoded.push((k, v).encode());
+		}
+
+		let index = found_index.ok_or("couldnt find")?;
+
+		let proof = binary_merkle_tree::merkle_proof::<
+			Hashing,
+			_,
+			_,
+		>(encoded, index as u32);
+		Ok(Encode::encode(&proof))
 	}
 }
 

From 133a4403ce5f70e1791e9639b5a6f306725ea4f9 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Tue, 17 Sep 2024 19:50:14 -0400
Subject: [PATCH 45/55] fix up binary tree

---
 .../runtime/src/proving_trie/binary.rs        | 130 ++++++++----------
 1 file changed, 54 insertions(+), 76 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie/binary.rs b/substrate/primitives/runtime/src/proving_trie/binary.rs
index 4b5548dbbdc4..53ea653ce2d7 100644
--- a/substrate/primitives/runtime/src/proving_trie/binary.rs
+++ b/substrate/primitives/runtime/src/proving_trie/binary.rs
@@ -15,14 +15,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//! Types for a compact base-16 merkle trie used for checking and generating proofs within the
-//! runtime. The `sp-trie` crate exposes all of these same functionality (and more), but this
-//! library is designed to work more easily with runtime native types, which simply need to
-//! implement `Encode`/`Decode`. It also exposes a runtime friendly `TrieError` type which can be
-//! use inside of a FRAME Pallet.
-//!
-//! Proofs are created with latest substrate trie format (`LayoutV1`), and are not compatible with
-//! proofs using `LayoutV0`.
+//! Types for a base-2 merkle tree used for checking and generating proofs within the
+//! runtime. The `binary-merkle-tree` crate exposes all of these same functionality (and more), but
+//! this library is designed to work more easily with runtime native types, which simply need to
+//! implement `Encode`/`Decode`.
 
 use crate::{Decode, DispatchError, Encode, MaxEncodedLen, TypeInfo};
 #[cfg(feature = "serde")]
@@ -63,17 +59,13 @@ where
 	{
 		let mut db_map = BTreeMap::default();
 		for (key, value) in items.into_iter() {
-			db_map.insert(key, value);
+			if db_map.insert(key, value).is_some() {
+				return Err(TrieError::DuplicateKey.into());
+			}
 		}
 
-		let db_map_len = db_map.len();
-
 		let db: Vec<(Key, Value)> = db_map.into_iter().collect();
 
-		if db.len() != db_map_len {
-			return Err("duplicate item key".into())
-		}
-
 		let root =
 			binary_merkle_tree::merkle_root::<Hashing, _>(db.iter().map(|item| item.encode()));
 
@@ -100,14 +92,17 @@ where
 	///
 	/// This function makes a proof with latest substrate trie format (`LayoutV1`), and is not
 	/// compatible with `LayoutV0`.
-	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError> {
+	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError>
+	where
+		Hashing::Out: Encode,
+	{
 		let mut encoded = Vec::with_capacity(self.db.len()); // Pre-allocate the vector
 		let mut found_index = None;
 
 		// Find the index of our key, and encode the (key, value) pair.
-		for (i, (k, v)) in self.db.into_iter().enumerate() {
+		for (i, (k, v)) in self.db.iter().enumerate() {
 			// If we found the key we are looking for, save it.
-			if k == key {
+			if *k == key {
 				found_index = Some(i);
 			}
 
@@ -116,35 +111,48 @@ where
 
 		let index = found_index.ok_or("couldnt find")?;
 
-		let proof = binary_merkle_tree::merkle_proof::<
-			Hashing,
-			_,
-			_,
-		>(encoded, index as u32);
-		Ok(Encode::encode(&proof))
+		let proof = binary_merkle_tree::merkle_proof::<Hashing, Vec<Vec<u8>>, Vec<u8>>(
+			encoded,
+			index as u32,
+		);
+		Ok(proof.encode())
 	}
 }
 
-/// Verify the existence or non-existence of `key` and `value` in a given trie root and proof.
-///
-/// Proofs must be created with latest substrate trie format (`LayoutV1`).
+/// Verify the existence of `key` and `value` in a given trie root and proof.
 pub fn verify_single_value_proof<Hashing, Key, Value>(
 	root: HashOf<Hashing>,
-	proof: &[Vec<u8>],
+	proof: &[u8],
 	key: Key,
-	maybe_value: Option<Value>,
+	value: Value,
 ) -> Result<(), DispatchError>
 where
 	Hashing: sp_core::Hasher,
-	Key: Encode,
-	Value: Encode,
+	Hashing::Out: Decode,
+	Key: Encode + Decode,
+	Value: Encode + Decode,
 {
-	sp_trie::verify_trie_proof::<LayoutV1<Hashing>, _, _, _>(
-		&root,
-		proof,
-		&[(key.encode(), maybe_value.map(|value| value.encode()))],
-	)
-	.map_err(|err| TrieError::from(err).into())
+	let decoded_proof: binary_merkle_tree::MerkleProof<Hashing::Out, Vec<u8>> =
+		Decode::decode(&mut &proof[..]).map_err(|_| TrieError::IncompleteProof)?;
+	if root != decoded_proof.root {
+		return Err(TrieError::RootMismatch.into());
+	}
+
+	if (&key, &value).encode() != decoded_proof.leaf {
+		return Err(TrieError::ValueMismatch.into());
+	}
+
+	if binary_merkle_tree::verify_proof::<Hashing, _, _>(
+		&decoded_proof.root,
+		decoded_proof.proof,
+		decoded_proof.number_of_leaves,
+		decoded_proof.leaf_index,
+		&decoded_proof.leaf,
+	) {
+		Ok(())
+	} else {
+		Err(TrieError::IncompleteProof.into())
+	}
 }
 
 #[cfg(test)]
@@ -159,7 +167,8 @@ mod tests {
 
 	// The expected root hash for an empty trie.
 	fn empty_root() -> H256 {
-		sp_trie::empty_trie_root::<LayoutV1<BlakeTwo256>>()
+		let tree = BalanceTrie::generate_for(Vec::new()).unwrap();
+		*tree.root()
 	}
 
 	fn create_balance_trie() -> BalanceTrie {
@@ -180,8 +189,6 @@ mod tests {
 		assert_eq!(balance_trie.query(6u32), Some(6u128));
 		assert_eq!(balance_trie.query(9u32), Some(9u128));
 		assert_eq!(balance_trie.query(69u32), Some(69u128));
-		// Invalid key returns none.
-		assert_eq!(balance_trie.query(6969u32), None);
 
 		balance_trie
 	}
@@ -204,12 +211,7 @@ mod tests {
 		for i in 0..200u32 {
 			if i == 6 {
 				assert_eq!(
-					verify_single_value_proof::<BlakeTwo256, _, _>(
-						root,
-						&proof,
-						i,
-						Some(u128::from(i))
-					),
+					verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, i, u128::from(i)),
 					Ok(())
 				);
 				// Wrong value is invalid.
@@ -218,41 +220,22 @@ mod tests {
 						root,
 						&proof,
 						i,
-						Some(u128::from(i + 1))
+						u128::from(i + 1)
 					),
-					Err(TrieError::RootMismatch.into())
+					Err(TrieError::ValueMismatch.into())
 				);
 			} else {
 				assert!(verify_single_value_proof::<BlakeTwo256, _, _>(
 					root,
 					&proof,
 					i,
-					Some(u128::from(i))
-				)
-				.is_err());
-				assert!(verify_single_value_proof::<BlakeTwo256, _, _>(
-					root,
-					&proof,
-					i,
-					None::<u128>
+					u128::from(i)
 				)
 				.is_err());
 			}
 		}
 	}
 
-	#[test]
-	fn basic_end_to_end_multi_value() {
-		let balance_trie = create_balance_trie();
-		let root = *balance_trie.root();
-
-		// Create a proof for a valid and invalid key.
-		let proof = balance_trie.create_proof(&[6u32, 69u32, 6969u32]).unwrap();
-		let items = [(6u32, Some(6u128)), (69u32, Some(69u128)), (6969u32, None)];
-
-		assert_eq!(verify_proof::<BlakeTwo256, _, _>(root, &proof, &items), Ok(()));
-	}
-
 	#[test]
 	fn proof_fails_with_bad_data() {
 		let balance_trie = create_balance_trie();
@@ -263,24 +246,19 @@ mod tests {
 
 		// Correct data verifies successfully
 		assert_eq!(
-			verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, Some(6u128)),
+			verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, 6u128),
 			Ok(())
 		);
 
 		// Fail to verify proof with wrong root
 		assert_eq!(
-			verify_single_value_proof::<BlakeTwo256, _, _>(
-				Default::default(),
-				&proof,
-				6u32,
-				Some(6u128)
-			),
+			verify_single_value_proof::<BlakeTwo256, _, _>(Default::default(), &proof, 6u32, 6u128),
 			Err(TrieError::RootMismatch.into())
 		);
 
 		// Fail to verify proof with wrong data
 		assert_eq!(
-			verify_single_value_proof::<BlakeTwo256, _, _>(root, &[], 6u32, Some(6u128)),
+			verify_single_value_proof::<BlakeTwo256, _, _>(root, &[], 6u32, 6u128),
 			Err(TrieError::IncompleteProof.into())
 		);
 	}

From 813caae50a3239cf76f7787a2efdb511d9c18317 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Tue, 17 Sep 2024 20:29:45 -0400
Subject: [PATCH 46/55] refactor

---
 .../base16.rs}                                | 129 ++++--------------
 .../src/proving_trie/{binary.rs => base2.rs}  |  52 ++++---
 .../runtime/src/proving_trie/mod.rs           | 113 +++++++++++++++
 3 files changed, 171 insertions(+), 123 deletions(-)
 rename substrate/primitives/runtime/src/{proving_trie.rs => proving_trie/base16.rs} (68%)
 rename substrate/primitives/runtime/src/proving_trie/{binary.rs => base2.rs} (88%)
 create mode 100644 substrate/primitives/runtime/src/proving_trie/mod.rs

diff --git a/substrate/primitives/runtime/src/proving_trie.rs b/substrate/primitives/runtime/src/proving_trie/base16.rs
similarity index 68%
rename from substrate/primitives/runtime/src/proving_trie.rs
rename to substrate/primitives/runtime/src/proving_trie/base16.rs
index 0043f80d7ef7..c164de619ef9 100644
--- a/substrate/primitives/runtime/src/proving_trie.rs
+++ b/substrate/primitives/runtime/src/proving_trie/base16.rs
@@ -24,107 +24,14 @@
 //! Proofs are created with latest substrate trie format (`LayoutV1`), and are not compatible with
 //! proofs using `LayoutV0`.
 
-use crate::{Decode, DispatchError, Encode, MaxEncodedLen, TypeInfo};
-#[cfg(feature = "serde")]
-use crate::{Deserialize, Serialize};
-
-mod binary;
-
+use super::TrieError;
+use crate::{Decode, DispatchError, Encode};
 use sp_std::vec::Vec;
 use sp_trie::{
-	trie_types::{TrieDBBuilder, TrieDBMutBuilderV1, TrieError as SpTrieError},
-	LayoutV1, MemoryDB, Trie, TrieMut, VerifyError,
+	trie_types::{TrieDBBuilder, TrieDBMutBuilderV1},
+	LayoutV1, MemoryDB, Trie, TrieMut,
 };
 
-type HashOf<Hashing> = <Hashing as sp_core::Hasher>::Out;
-
-/// A runtime friendly error type for tries.
-#[derive(Eq, PartialEq, Clone, Copy, Encode, Decode, Debug, TypeInfo, MaxEncodedLen)]
-#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
-pub enum TrieError {
-	/* From TrieError */
-	/// Attempted to create a trie with a state root not in the DB.
-	InvalidStateRoot,
-	/// Trie item not found in the database,
-	IncompleteDatabase,
-	/// A value was found in the trie with a nibble key that was not byte-aligned.
-	ValueAtIncompleteKey,
-	/// Corrupt Trie item.
-	DecoderError,
-	/// Hash is not value.
-	InvalidHash,
-	/* From VerifyError */
-	/// The statement being verified contains multiple key-value pairs with the same key.
-	DuplicateKey,
-	/// The proof contains at least one extraneous node.
-	ExtraneousNode,
-	/// The proof contains at least one extraneous value which should have been omitted from the
-	/// proof.
-	ExtraneousValue,
-	/// The proof contains at least one extraneous hash reference the should have been omitted.
-	ExtraneousHashReference,
-	/// The proof contains an invalid child reference that exceeds the hash length.
-	InvalidChildReference,
-	/// The proof indicates that an expected value was not found in the trie.
-	ValueMismatch,
-	/// The proof is missing trie nodes required to verify.
-	IncompleteProof,
-	/// The root hash computed from the proof is incorrect.
-	RootMismatch,
-	/// One of the proof nodes could not be decoded.
-	DecodeError,
-}
-
-impl<T> From<SpTrieError<T>> for TrieError {
-	fn from(error: SpTrieError<T>) -> Self {
-		match error {
-			SpTrieError::InvalidStateRoot(..) => Self::InvalidStateRoot,
-			SpTrieError::IncompleteDatabase(..) => Self::IncompleteDatabase,
-			SpTrieError::ValueAtIncompleteKey(..) => Self::ValueAtIncompleteKey,
-			SpTrieError::DecoderError(..) => Self::DecoderError,
-			SpTrieError::InvalidHash(..) => Self::InvalidHash,
-		}
-	}
-}
-
-impl<T, U> From<VerifyError<T, U>> for TrieError {
-	fn from(error: VerifyError<T, U>) -> Self {
-		match error {
-			VerifyError::DuplicateKey(..) => Self::DuplicateKey,
-			VerifyError::ExtraneousNode => Self::ExtraneousNode,
-			VerifyError::ExtraneousValue(..) => Self::ExtraneousValue,
-			VerifyError::ExtraneousHashReference(..) => Self::ExtraneousHashReference,
-			VerifyError::InvalidChildReference(..) => Self::InvalidChildReference,
-			VerifyError::ValueMismatch(..) => Self::ValueMismatch,
-			VerifyError::IncompleteProof => Self::IncompleteProof,
-			VerifyError::RootMismatch(..) => Self::RootMismatch,
-			VerifyError::DecodeError(..) => Self::DecodeError,
-		}
-	}
-}
-
-impl From<TrieError> for &'static str {
-	fn from(e: TrieError) -> &'static str {
-		match e {
-			TrieError::InvalidStateRoot => "The state root is not in the database.",
-			TrieError::IncompleteDatabase => "A trie item was not found in the database.",
-			TrieError::ValueAtIncompleteKey =>
-				"A value was found with a key that is not byte-aligned.",
-			TrieError::DecoderError => "A corrupt trie item was encountered.",
-			TrieError::InvalidHash => "The hash does not match the expected value.",
-			TrieError::DuplicateKey => "The proof contains duplicate keys.",
-			TrieError::ExtraneousNode => "The proof contains extraneous nodes.",
-			TrieError::ExtraneousValue => "The proof contains extraneous values.",
-			TrieError::ExtraneousHashReference => "The proof contains extraneous hash references.",
-			TrieError::InvalidChildReference => "The proof contains an invalid child reference.",
-			TrieError::ValueMismatch => "The proof indicates a value mismatch.",
-			TrieError::IncompleteProof => "The proof is incomplete.",
-			TrieError::RootMismatch => "The root hash computed from the proof is incorrect.",
-			TrieError::DecodeError => "One of the proof nodes could not be decoded.",
-		}
-	}
-}
-
 /// A helper structure for building a basic base-16 merkle trie and creating compact proofs for that
 /// trie. Proofs are created with latest substrate trie format (`LayoutV1`), and are not compatible
 /// with proofs using `LayoutV0`.
@@ -133,7 +40,7 @@ where
 	Hashing: sp_core::Hasher,
 {
 	db: MemoryDB<Hashing>,
-	root: HashOf<Hashing>,
+	root: Hashing::Out,
 	_phantom: core::marker::PhantomData<(Key, Value)>,
 }
 
@@ -163,7 +70,7 @@ where
 	}
 
 	/// Access the underlying trie root.
-	pub fn root(&self) -> &HashOf<Hashing> {
+	pub fn root(&self) -> &Hashing::Out {
 		&self.root
 	}
 
@@ -213,7 +120,7 @@ where
 ///
 /// Proofs must be created with latest substrate trie format (`LayoutV1`).
 pub fn verify_single_value_proof<Hashing, Key, Value>(
-	root: HashOf<Hashing>,
+	root: Hashing::Out,
 	proof: &[Vec<u8>],
 	key: Key,
 	maybe_value: Option<Value>,
@@ -235,7 +142,7 @@ where
 ///
 /// Proofs must be created with latest substrate trie format (`LayoutV1`).
 pub fn verify_proof<Hashing, Key, Value>(
-	root: HashOf<Hashing>,
+	root: Hashing::Out,
 	proof: &[Vec<u8>],
 	items: &[(Key, Option<Value>)],
 ) -> Result<(), DispatchError>
@@ -390,4 +297,24 @@ mod tests {
 			Err(TrieError::IncompleteProof.into())
 		);
 	}
+
+	#[test]
+	fn inserting_the_same_key_twice() {
+		// Create a map of users and their balances.
+		let mut flat = Vec::<(u32, u128)>::new();
+		for i in 0..100u32 {
+			flat.push((i, i.into()));
+		}
+
+		// Duplicate key
+		flat.push((50, 1337.into()));
+
+		let Err(error) = BalanceTrie::generate_for(flat) else {
+			panic!("expected balance trie error");
+		};
+		assert_eq!(
+			error,
+			TrieError::DuplicateKey.into(),
+		);
+	}
 }
diff --git a/substrate/primitives/runtime/src/proving_trie/binary.rs b/substrate/primitives/runtime/src/proving_trie/base2.rs
similarity index 88%
rename from substrate/primitives/runtime/src/proving_trie/binary.rs
rename to substrate/primitives/runtime/src/proving_trie/base2.rs
index 53ea653ce2d7..b32248d3437c 100644
--- a/substrate/primitives/runtime/src/proving_trie/binary.rs
+++ b/substrate/primitives/runtime/src/proving_trie/base2.rs
@@ -20,18 +20,10 @@
 //! this library is designed to work more easily with runtime native types, which simply need to
 //! implement `Encode`/`Decode`.
 
-use crate::{Decode, DispatchError, Encode, MaxEncodedLen, TypeInfo};
-#[cfg(feature = "serde")]
-use crate::{Deserialize, Serialize};
-
-use super::*;
+use super::TrieError;
+use crate::{Decode, DispatchError, Encode};
+use binary_merkle_tree::{merkle_proof, merkle_root, verify_proof, MerkleProof};
 use sp_std::{collections::btree_map::BTreeMap, vec::Vec};
-use sp_trie::{
-	trie_types::{TrieDBBuilder, TrieDBMutBuilderV1, TrieError as SpTrieError},
-	LayoutV1, MemoryDB, Trie, TrieMut, VerifyError,
-};
-
-type HashOf<Hashing> = <Hashing as sp_core::Hasher>::Out;
 
 /// A helper structure for building a basic base-16 merkle trie and creating compact proofs for that
 /// trie. Proofs are created with latest substrate trie format (`LayoutV1`), and are not compatible
@@ -42,7 +34,7 @@ where
 {
 	// Deduplicated and flattened list of key value pairs.
 	db: Vec<(Key, Value)>,
-	root: HashOf<Hashing>,
+	root: Hashing::Out,
 	_phantom: core::marker::PhantomData<(Key, Value)>,
 }
 
@@ -66,14 +58,13 @@ where
 
 		let db: Vec<(Key, Value)> = db_map.into_iter().collect();
 
-		let root =
-			binary_merkle_tree::merkle_root::<Hashing, _>(db.iter().map(|item| item.encode()));
+		let root = merkle_root::<Hashing, _>(db.iter().map(|item| item.encode()));
 
 		Ok(Self { db, root, _phantom: Default::default() })
 	}
 
 	/// Access the underlying trie root.
-	pub fn root(&self) -> &HashOf<Hashing> {
+	pub fn root(&self) -> &Hashing::Out {
 		&self.root
 	}
 
@@ -111,17 +102,14 @@ where
 
 		let index = found_index.ok_or("couldnt find")?;
 
-		let proof = binary_merkle_tree::merkle_proof::<Hashing, Vec<Vec<u8>>, Vec<u8>>(
-			encoded,
-			index as u32,
-		);
+		let proof = merkle_proof::<Hashing, Vec<Vec<u8>>, Vec<u8>>(encoded, index as u32);
 		Ok(proof.encode())
 	}
 }
 
 /// Verify the existence of `key` and `value` in a given trie root and proof.
 pub fn verify_single_value_proof<Hashing, Key, Value>(
-	root: HashOf<Hashing>,
+	root: Hashing::Out,
 	proof: &[u8],
 	key: Key,
 	value: Value,
@@ -132,7 +120,7 @@ where
 	Key: Encode + Decode,
 	Value: Encode + Decode,
 {
-	let decoded_proof: binary_merkle_tree::MerkleProof<Hashing::Out, Vec<u8>> =
+	let decoded_proof: MerkleProof<Hashing::Out, Vec<u8>> =
 		Decode::decode(&mut &proof[..]).map_err(|_| TrieError::IncompleteProof)?;
 	if root != decoded_proof.root {
 		return Err(TrieError::RootMismatch.into());
@@ -142,7 +130,7 @@ where
 		return Err(TrieError::ValueMismatch.into());
 	}
 
-	if binary_merkle_tree::verify_proof::<Hashing, _, _>(
+	if verify_proof::<Hashing, _, _>(
 		&decoded_proof.root,
 		decoded_proof.proof,
 		decoded_proof.number_of_leaves,
@@ -262,4 +250,24 @@ mod tests {
 			Err(TrieError::IncompleteProof.into())
 		);
 	}
+
+	#[test]
+	fn duplicate_key_fails_to_build() {
+		// Create a map of users and their balances.
+		let mut flat = Vec::<(u32, u128)>::new();
+		for i in 0..100u32 {
+			flat.push((i, i.into()));
+		}
+
+		// Duplicate key
+		flat.push((50, 1337.into()));
+
+		let Err(error) = BalanceTrie::generate_for(flat) else {
+			panic!("expected balance trie error");
+		};
+		assert_eq!(
+			error,
+			TrieError::DuplicateKey.into(),
+		);
+	}
 }
diff --git a/substrate/primitives/runtime/src/proving_trie/mod.rs b/substrate/primitives/runtime/src/proving_trie/mod.rs
new file mode 100644
index 000000000000..60f11645a48f
--- /dev/null
+++ b/substrate/primitives/runtime/src/proving_trie/mod.rs
@@ -0,0 +1,113 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Types for merkle tries compatible with the runtime.
+
+pub mod base16;
+pub mod base2;
+
+use crate::{Decode, Encode, MaxEncodedLen, TypeInfo};
+#[cfg(feature = "serde")]
+use crate::{Deserialize, Serialize};
+use sp_trie::{trie_types::TrieError as SpTrieError, VerifyError};
+
+/// A runtime friendly error type for tries.
+#[derive(Eq, PartialEq, Clone, Copy, Encode, Decode, Debug, TypeInfo, MaxEncodedLen)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum TrieError {
+	/* From TrieError */
+	/// Attempted to create a trie with a state root not in the DB.
+	InvalidStateRoot,
+	/// Trie item not found in the database,
+	IncompleteDatabase,
+	/// A value was found in the trie with a nibble key that was not byte-aligned.
+	ValueAtIncompleteKey,
+	/// Corrupt Trie item.
+	DecoderError,
+	/// Hash is not value.
+	InvalidHash,
+	/* From VerifyError */
+	/// The statement being verified contains multiple key-value pairs with the same key.
+	DuplicateKey,
+	/// The proof contains at least one extraneous node.
+	ExtraneousNode,
+	/// The proof contains at least one extraneous value which should have been omitted from the
+	/// proof.
+	ExtraneousValue,
+	/// The proof contains at least one extraneous hash reference the should have been omitted.
+	ExtraneousHashReference,
+	/// The proof contains an invalid child reference that exceeds the hash length.
+	InvalidChildReference,
+	/// The proof indicates that an expected value was not found in the trie.
+	ValueMismatch,
+	/// The proof is missing trie nodes required to verify.
+	IncompleteProof,
+	/// The root hash computed from the proof is incorrect.
+	RootMismatch,
+	/// One of the proof nodes could not be decoded.
+	DecodeError,
+}
+
+impl<T> From<SpTrieError<T>> for TrieError {
+	fn from(error: SpTrieError<T>) -> Self {
+		match error {
+			SpTrieError::InvalidStateRoot(..) => Self::InvalidStateRoot,
+			SpTrieError::IncompleteDatabase(..) => Self::IncompleteDatabase,
+			SpTrieError::ValueAtIncompleteKey(..) => Self::ValueAtIncompleteKey,
+			SpTrieError::DecoderError(..) => Self::DecoderError,
+			SpTrieError::InvalidHash(..) => Self::InvalidHash,
+		}
+	}
+}
+
+impl<T, U> From<VerifyError<T, U>> for TrieError {
+	fn from(error: VerifyError<T, U>) -> Self {
+		match error {
+			VerifyError::DuplicateKey(..) => Self::DuplicateKey,
+			VerifyError::ExtraneousNode => Self::ExtraneousNode,
+			VerifyError::ExtraneousValue(..) => Self::ExtraneousValue,
+			VerifyError::ExtraneousHashReference(..) => Self::ExtraneousHashReference,
+			VerifyError::InvalidChildReference(..) => Self::InvalidChildReference,
+			VerifyError::ValueMismatch(..) => Self::ValueMismatch,
+			VerifyError::IncompleteProof => Self::IncompleteProof,
+			VerifyError::RootMismatch(..) => Self::RootMismatch,
+			VerifyError::DecodeError(..) => Self::DecodeError,
+		}
+	}
+}
+
+impl From<TrieError> for &'static str {
+	fn from(e: TrieError) -> &'static str {
+		match e {
+			TrieError::InvalidStateRoot => "The state root is not in the database.",
+			TrieError::IncompleteDatabase => "A trie item was not found in the database.",
+			TrieError::ValueAtIncompleteKey =>
+				"A value was found with a key that is not byte-aligned.",
+			TrieError::DecoderError => "A corrupt trie item was encountered.",
+			TrieError::InvalidHash => "The hash does not match the expected value.",
+			TrieError::DuplicateKey => "The proof contains duplicate keys.",
+			TrieError::ExtraneousNode => "The proof contains extraneous nodes.",
+			TrieError::ExtraneousValue => "The proof contains extraneous values.",
+			TrieError::ExtraneousHashReference => "The proof contains extraneous hash references.",
+			TrieError::InvalidChildReference => "The proof contains an invalid child reference.",
+			TrieError::ValueMismatch => "The proof indicates a value mismatch.",
+			TrieError::IncompleteProof => "The proof is incomplete.",
+			TrieError::RootMismatch => "The root hash computed from the proof is incorrect.",
+			TrieError::DecodeError => "One of the proof nodes could not be decoded.",
+		}
+	}
+}

From e2702758b813f29162838e7abe2e468a44a39188 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Tue, 17 Sep 2024 21:35:26 -0400
Subject: [PATCH 47/55] fixes

---
 .../runtime/src/proving_trie/base16.rs        | 24 ---------
 .../runtime/src/proving_trie/base2.rs         | 51 ++++---------------
 2 files changed, 10 insertions(+), 65 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie/base16.rs b/substrate/primitives/runtime/src/proving_trie/base16.rs
index c164de619ef9..5ced5f6ec2ec 100644
--- a/substrate/primitives/runtime/src/proving_trie/base16.rs
+++ b/substrate/primitives/runtime/src/proving_trie/base16.rs
@@ -87,8 +87,6 @@ where
 	}
 
 	/// Create a compact merkle proof needed to prove all `keys` and their values are in the trie.
-	/// Returns `None` if the nodes within the current `MemoryDB` are insufficient to create a
-	/// proof.
 	///
 	/// This function makes a proof with latest substrate trie format (`LayoutV1`), and is not
 	/// compatible with `LayoutV0`.
@@ -106,8 +104,6 @@ where
 	}
 
 	/// Create a compact merkle proof needed to prove a single key and its value are in the trie.
-	/// Returns `None` if the nodes within the current `MemoryDB` are insufficient to create a
-	/// proof.
 	///
 	/// This function makes a proof with latest substrate trie format (`LayoutV1`), and is not
 	/// compatible with `LayoutV0`.
@@ -297,24 +293,4 @@ mod tests {
 			Err(TrieError::IncompleteProof.into())
 		);
 	}
-
-	#[test]
-	fn inserting_the_same_key_twice() {
-		// Create a map of users and their balances.
-		let mut flat = Vec::<(u32, u128)>::new();
-		for i in 0..100u32 {
-			flat.push((i, i.into()));
-		}
-
-		// Duplicate key
-		flat.push((50, 1337.into()));
-
-		let Err(error) = BalanceTrie::generate_for(flat) else {
-			panic!("expected balance trie error");
-		};
-		assert_eq!(
-			error,
-			TrieError::DuplicateKey.into(),
-		);
-	}
 }
diff --git a/substrate/primitives/runtime/src/proving_trie/base2.rs b/substrate/primitives/runtime/src/proving_trie/base2.rs
index b32248d3437c..216910f2c7d0 100644
--- a/substrate/primitives/runtime/src/proving_trie/base2.rs
+++ b/substrate/primitives/runtime/src/proving_trie/base2.rs
@@ -25,15 +25,14 @@ use crate::{Decode, DispatchError, Encode};
 use binary_merkle_tree::{merkle_proof, merkle_root, verify_proof, MerkleProof};
 use sp_std::{collections::btree_map::BTreeMap, vec::Vec};
 
-/// A helper structure for building a basic base-16 merkle trie and creating compact proofs for that
-/// trie. Proofs are created with latest substrate trie format (`LayoutV1`), and are not compatible
-/// with proofs using `LayoutV0`.
+/// A helper structure for building a basic base-2 merkle trie and creating compact proofs for that
+/// trie.
 pub struct BasicProvingTrie<Hashing, Key, Value>
 where
 	Hashing: sp_core::Hasher,
 {
 	// Deduplicated and flattened list of key value pairs.
-	db: Vec<(Key, Value)>,
+	db: BTreeMap<Key, Value>,
 	root: Hashing::Out,
 	_phantom: core::marker::PhantomData<(Key, Value)>,
 }
@@ -49,17 +48,11 @@ where
 	where
 		I: IntoIterator<Item = (Key, Value)>,
 	{
-		let mut db_map = BTreeMap::default();
+		let mut db = BTreeMap::default();
 		for (key, value) in items.into_iter() {
-			if db_map.insert(key, value).is_some() {
-				return Err(TrieError::DuplicateKey.into());
-			}
+			db.insert(key, value);
 		}
-
-		let db: Vec<(Key, Value)> = db_map.into_iter().collect();
-
 		let root = merkle_root::<Hashing, _>(db.iter().map(|item| item.encode()));
-
 		Ok(Self { db, root, _phantom: Default::default() })
 	}
 
@@ -69,25 +62,22 @@ where
 	}
 
 	/// Query a value contained within the current trie. Returns `None` if the
-	/// nodes within the current `MemoryDB` are insufficient to query the item.
+	/// nodes within the current `db` are insufficient to query the item.
 	pub fn query(&self, key: Key) -> Option<Value>
 	where
 		Value: Decode + Clone,
 	{
-		self.db.iter().find(|(k, _)| *k == key).map(|(_, v)| v.clone())
+		self.db.get(&key).cloned()
 	}
 
 	/// Create a compact merkle proof needed to prove a single key and its value are in the trie.
-	/// Returns `None` if the nodes within the current `MemoryDB` are insufficient to create a
+	/// Returns `None` if the nodes within the current `db` are insufficient to create a
 	/// proof.
-	///
-	/// This function makes a proof with latest substrate trie format (`LayoutV1`), and is not
-	/// compatible with `LayoutV0`.
 	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError>
 	where
 		Hashing::Out: Encode,
 	{
-		let mut encoded = Vec::with_capacity(self.db.len()); // Pre-allocate the vector
+		let mut encoded = Vec::with_capacity(self.db.len());
 		let mut found_index = None;
 
 		// Find the index of our key, and encode the (key, value) pair.
@@ -100,8 +90,7 @@ where
 			encoded.push((k, v).encode());
 		}
 
-		let index = found_index.ok_or("couldnt find")?;
-
+		let index = found_index.ok_or(TrieError::IncompleteDatabase)?;
 		let proof = merkle_proof::<Hashing, Vec<Vec<u8>>, Vec<u8>>(encoded, index as u32);
 		Ok(proof.encode())
 	}
@@ -250,24 +239,4 @@ mod tests {
 			Err(TrieError::IncompleteProof.into())
 		);
 	}
-
-	#[test]
-	fn duplicate_key_fails_to_build() {
-		// Create a map of users and their balances.
-		let mut flat = Vec::<(u32, u128)>::new();
-		for i in 0..100u32 {
-			flat.push((i, i.into()));
-		}
-
-		// Duplicate key
-		flat.push((50, 1337.into()));
-
-		let Err(error) = BalanceTrie::generate_for(flat) else {
-			panic!("expected balance trie error");
-		};
-		assert_eq!(
-			error,
-			TrieError::DuplicateKey.into(),
-		);
-	}
 }

From a679bc961583905722d85b26776c334b630d4ef6 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Tue, 17 Sep 2024 22:19:29 -0400
Subject: [PATCH 48/55] update api for a single opaque blob

---
 .../runtime/src/proving_trie/base16.rs        | 32 +++++++++++++------
 1 file changed, 22 insertions(+), 10 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie/base16.rs b/substrate/primitives/runtime/src/proving_trie/base16.rs
index 5ced5f6ec2ec..69dbe77fa9d6 100644
--- a/substrate/primitives/runtime/src/proving_trie/base16.rs
+++ b/substrate/primitives/runtime/src/proving_trie/base16.rs
@@ -94,20 +94,21 @@ where
 	/// When verifying the proof created by this function, you must include all of the keys and
 	/// values of the proof, else the verifier will complain that extra nodes are provided in the
 	/// proof that are not needed.
-	pub fn create_proof(&self, keys: &[Key]) -> Result<Vec<Vec<u8>>, DispatchError> {
+	pub fn create_proof(&self, keys: &[Key]) -> Result<Vec<u8>, DispatchError> {
 		sp_trie::generate_trie_proof::<LayoutV1<Hashing>, _, _, _>(
 			&self.db,
 			self.root,
 			&keys.into_iter().map(|k| k.encode()).collect::<Vec<Vec<u8>>>(),
 		)
 		.map_err(|err| TrieError::from(*err).into())
+		.map(|structured_proof| structured_proof.encode())
 	}
 
 	/// Create a compact merkle proof needed to prove a single key and its value are in the trie.
 	///
 	/// This function makes a proof with latest substrate trie format (`LayoutV1`), and is not
 	/// compatible with `LayoutV0`.
-	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<Vec<u8>>, DispatchError> {
+	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError> {
 		self.create_proof(&[key])
 	}
 }
@@ -117,7 +118,7 @@ where
 /// Proofs must be created with latest substrate trie format (`LayoutV1`).
 pub fn verify_single_value_proof<Hashing, Key, Value>(
 	root: Hashing::Out,
-	proof: &[Vec<u8>],
+	proof: &[u8],
 	key: Key,
 	maybe_value: Option<Value>,
 ) -> Result<(), DispatchError>
@@ -126,9 +127,11 @@ where
 	Key: Encode,
 	Value: Encode,
 {
+	let structured_proof: Vec<Vec<u8>> =
+		Decode::decode(&mut &proof[..]).map_err(|_| TrieError::DecodeError)?;
 	sp_trie::verify_trie_proof::<LayoutV1<Hashing>, _, _, _>(
 		&root,
-		proof,
+		&structured_proof,
 		&[(key.encode(), maybe_value.map(|value| value.encode()))],
 	)
 	.map_err(|err| TrieError::from(err).into())
@@ -139,7 +142,7 @@ where
 /// Proofs must be created with latest substrate trie format (`LayoutV1`).
 pub fn verify_proof<Hashing, Key, Value>(
 	root: Hashing::Out,
-	proof: &[Vec<u8>],
+	proof: &[u8],
 	items: &[(Key, Option<Value>)],
 ) -> Result<(), DispatchError>
 where
@@ -147,13 +150,19 @@ where
 	Key: Encode,
 	Value: Encode,
 {
+	let structured_proof: Vec<Vec<u8>> =
+		Decode::decode(&mut &proof[..]).map_err(|_| TrieError::DecodeError)?;
 	let items_encoded = items
 		.into_iter()
 		.map(|(key, maybe_value)| (key.encode(), maybe_value.as_ref().map(|value| value.encode())))
 		.collect::<Vec<(Vec<u8>, Option<Vec<u8>>)>>();
 
-	sp_trie::verify_trie_proof::<LayoutV1<Hashing>, _, _, _>(&root, proof, &items_encoded)
-		.map_err(|err| TrieError::from(err).into())
+	sp_trie::verify_trie_proof::<LayoutV1<Hashing>, _, _, _>(
+		&root,
+		&structured_proof,
+		&items_encoded,
+	)
+	.map_err(|err| TrieError::from(err).into())
 }
 
 #[cfg(test)]
@@ -287,10 +296,13 @@ mod tests {
 			Err(TrieError::RootMismatch.into())
 		);
 
-		// Fail to verify proof with wrong data
+		// Crete a bad proof.
+		let bad_proof = balance_trie.create_single_value_proof(99u32).unwrap();
+
+		// Fail to verify data with the wrong proof
 		assert_eq!(
-			verify_single_value_proof::<BlakeTwo256, _, _>(root, &[], 6u32, Some(6u128)),
-			Err(TrieError::IncompleteProof.into())
+			verify_single_value_proof::<BlakeTwo256, _, _>(root, &bad_proof, 6u32, Some(6u128)),
+			Err(TrieError::ExtraneousHashReference.into())
 		);
 	}
 }

From 65d2b98a2d0b816de132ac57509538dc893e612e Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Wed, 18 Sep 2024 11:23:56 -0400
Subject: [PATCH 49/55] undo changes in assets

---
 substrate/frame/assets/Cargo.toml       |   1 -
 substrate/frame/assets/src/functions.rs | 128 ----------------------
 substrate/frame/assets/src/lib.rs       | 137 +-----------------------
 substrate/frame/assets/src/mock.rs      |   5 +-
 substrate/frame/assets/src/tests.rs     |  64 -----------
 substrate/frame/assets/src/types.rs     |  14 ---
 substrate/frame/assets/src/weights.rs   |  37 -------
 7 files changed, 3 insertions(+), 383 deletions(-)

diff --git a/substrate/frame/assets/Cargo.toml b/substrate/frame/assets/Cargo.toml
index 5f1a876a5dc3..e20b576d0836 100644
--- a/substrate/frame/assets/Cargo.toml
+++ b/substrate/frame/assets/Cargo.toml
@@ -32,7 +32,6 @@ sp-core = { workspace = true }
 [dev-dependencies]
 sp-io = { workspace = true, default-features = true }
 pallet-balances = { workspace = true, default-features = true }
-binary-merkle-tree = { workspace = true }
 
 [features]
 default = ["std"]
diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs
index 02822ad46f64..c218c4ddc952 100644
--- a/substrate/frame/assets/src/functions.rs
+++ b/substrate/frame/assets/src/functions.rs
@@ -438,134 +438,6 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		Ok(())
 	}
 
-	/// Creates a distribution in storage for asset `id`, which can be claimed via
-	/// `do_claim_distribution`.
-	pub(super) fn do_mint_distribution(
-		id: T::AssetId,
-		merkle_root: T::Hash,
-		maybe_check_issuer: Option<T::AccountId>,
-	) -> DispatchResult {
-		let details = Asset::<T, I>::get(&id).ok_or(Error::<T, I>::Unknown)?;
-		ensure!(details.status == AssetStatus::Live, Error::<T, I>::AssetNotLive);
-
-		if let Some(check_issuer) = maybe_check_issuer {
-			ensure!(check_issuer == details.issuer, Error::<T, I>::NoPermission);
-		}
-
-		let info = DistributionInfo {
-			asset_id: id.clone(),
-			merkle_root: merkle_root.clone(),
-			active: true,
-		};
-
-		let distribution_id: u32 = MerklizedDistribution::<T, I>::count();
-		MerklizedDistribution::<T, I>::insert(&distribution_id, info);
-
-		Self::deposit_event(Event::DistributionIssued {
-			distribution_id,
-			asset_id: id,
-			merkle_root,
-		});
-
-		Ok(())
-	}
-
-	/// A wrapper around `do_mint`, allowing a `merkle_proof` to control the amount minted and to
-	/// whom.
-	pub(super) fn do_claim_distribution(
-		distribution_id: DistributionCounter,
-		merkle_proof: Vec<u8>,
-	) -> DispatchResult {
-		let proof =
-			codec::Decode::decode(&mut &merkle_proof[..]).map_err(|_| Error::<T, I>::BadProof)?;
-
-		let DistributionInfo { asset_id, merkle_root, active } =
-			MerklizedDistribution::<T, I>::get(distribution_id).ok_or(Error::<T, I>::Unknown)?;
-
-		ensure!(active, Error::<T, I>::DistributionEnded);
-
-		let leaf = T::VerifyExistenceProof::verify_proof(proof, &merkle_root)
-			.map_err(|()| Error::<T, I>::BadProof)?;
-		let (beneficiary, amount) =
-			codec::Decode::decode(&mut &leaf[..]).map_err(|_| Error::<T, I>::CannotDecodeLeaf)?;
-
-		ensure!(
-			!MerklizedDistributionTracker::<T, I>::contains_key(distribution_id, &beneficiary),
-			Error::<T, I>::AlreadyClaimed
-		);
-
-		// sp_runtime::proving_trie::verify_single_value_proof::<T::Hashing, _, _>(
-		// 	merkle_root,
-		// 	&merkle_proof,
-		// 	&beneficiary,
-		// 	Some(amount),
-		// )?;
-
-		Self::do_mint(asset_id, &beneficiary, amount, None)?;
-		MerklizedDistributionTracker::<T, I>::insert(&distribution_id, &beneficiary, ());
-
-		Ok(())
-	}
-
-	/// Ends the asset distribution of `distribution_id`.
-	pub(super) fn do_end_distribution(
-		distribution_id: DistributionCounter,
-		maybe_check_issuer: Option<T::AccountId>,
-	) -> DispatchResult {
-		let mut info =
-			MerklizedDistribution::<T, I>::get(&distribution_id).ok_or(Error::<T, I>::Unknown)?;
-		let details = Asset::<T, I>::get(&info.asset_id).ok_or(Error::<T, I>::Unknown)?;
-
-		if let Some(check_issuer) = maybe_check_issuer {
-			ensure!(check_issuer == details.issuer, Error::<T, I>::NoPermission);
-		}
-
-		info.active = false;
-
-		MerklizedDistribution::<T, I>::insert(&distribution_id, info);
-
-		Self::deposit_event(Event::DistributionEnded { distribution_id });
-
-		Ok(())
-	}
-
-	/// Ends the asset distribution of `distribution_id`.
-	pub(super) fn do_clean_distribution(
-		distribution_id: DistributionCounter,
-	) -> DispatchResultWithPostInfo {
-		let info =
-			MerklizedDistribution::<T, I>::get(&distribution_id).ok_or(Error::<T, I>::Unknown)?;
-
-		ensure!(!info.active, Error::<T, I>::DistributionActive);
-
-		let mut refund_count = 0u32;
-		let distribution_iterator =
-			MerklizedDistributionTracker::<T, I>::iter_key_prefix(&distribution_id);
-
-		let mut all_refunded = true;
-		for who in distribution_iterator {
-			if refund_count >= 100 {
-				// TODO: T::RemoveKeysLimit::get() {
-				// Not everyone was able to be refunded this time around.
-				all_refunded = false;
-				break
-			}
-
-			MerklizedDistributionTracker::<T, I>::remove(&distribution_id, &who);
-			refund_count += 1;
-		}
-
-		if all_refunded {
-			Self::deposit_event(Event::<T, I>::DistributionPartiallyCleaned { distribution_id });
-			// Refund weight only the amount we actually used.
-			Ok(Some(T::WeightInfo::clean_distribution(refund_count)).into())
-		} else {
-			Self::deposit_event(Event::<T, I>::DistributionCleaned { distribution_id });
-			// No weight to refund since we did not finish the loop.
-			Ok(().into())
-		}
-	}
-
 	/// Increases the asset `id` balance of `beneficiary` by `amount`.
 	///
 	/// LOW-LEVEL: Does not alter the supply of asset or emit an event. Use `do_mint` if you need
diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs
index 2f8e1a78a1ae..e909932bfc82 100644
--- a/substrate/frame/assets/src/lib.rs
+++ b/substrate/frame/assets/src/lib.rs
@@ -189,8 +189,7 @@ use frame_support::{
 			WithdrawConsequence,
 		},
 		BalanceStatus::Reserved,
-		BinaryMerkleTreeProver, Currency, EnsureOriginWithArg, Incrementable, ReservableCurrency,
-		StoredMap, VerifyExistenceProof,
+		Currency, EnsureOriginWithArg, Incrementable, ReservableCurrency, StoredMap,
 	},
 };
 use frame_system::Config as SystemConfig;
@@ -299,7 +298,6 @@ pub mod pallet {
 			type Extra = ();
 			type CallbackHandle = ();
 			type WeightInfo = ();
-			type VerifyExistenceProof = BinaryMerkleTreeProver<Self::Hashing>;
 			#[cfg(feature = "runtime-benchmarks")]
 			type BenchmarkHelper = ();
 		}
@@ -406,9 +404,6 @@ pub mod pallet {
 		/// used to set up auto-incrementing asset IDs for this collection.
 		type CallbackHandle: AssetsCallback<Self::AssetId, Self::AccountId>;
 
-		/// A type used to verify merkle proofs used for distributions.
-		type VerifyExistenceProof: VerifyExistenceProof<Hash = Self::Hash>;
-
 		/// Weight information for extrinsics in this pallet.
 		type WeightInfo: WeightInfo;
 
@@ -461,28 +456,6 @@ pub mod pallet {
 		ValueQuery,
 	>;
 
-	#[pallet::storage]
-	/// Merklized distribution of an asset.
-	pub(super) type MerklizedDistribution<T: Config<I>, I: 'static = ()> = CountedStorageMap<
-		_,
-		Blake2_128Concat,
-		DistributionCounter,
-		DistributionInfo<T::AssetId, T::Hash>,
-		OptionQuery,
-	>;
-
-	#[pallet::storage]
-	/// Tracks the merklized distribution of an asset so that assets are only claimed once.
-	pub(super) type MerklizedDistributionTracker<T: Config<I>, I: 'static = ()> = StorageDoubleMap<
-		_,
-		Blake2_128Concat,
-		DistributionCounter,
-		Blake2_128Concat,
-		T::AccountId,
-		(),
-		OptionQuery,
-	>;
-
 	/// The asset ID enforced for the next asset creation, if any present. Otherwise, this storage
 	/// item has no effect.
 	///
@@ -666,18 +639,6 @@ pub mod pallet {
 		Deposited { asset_id: T::AssetId, who: T::AccountId, amount: T::Balance },
 		/// Some assets were withdrawn from the account (e.g. for transaction fees).
 		Withdrawn { asset_id: T::AssetId, who: T::AccountId, amount: T::Balance },
-		/// A distribution of assets were issued.
-		DistributionIssued {
-			distribution_id: DistributionCounter,
-			asset_id: T::AssetId,
-			merkle_root: T::Hash,
-		},
-		/// A distribution has ended.
-		DistributionEnded { distribution_id: DistributionCounter },
-		/// A distribution has been partially cleaned. There are still more items to clean up.
-		DistributionPartiallyCleaned { distribution_id: DistributionCounter },
-		/// A distribution has been fully cleaned.
-		DistributionCleaned { distribution_id: DistributionCounter },
 	}
 
 	#[pallet::error]
@@ -727,16 +688,6 @@ pub mod pallet {
 		CallbackFailed,
 		/// The asset ID must be equal to the [`NextAssetId`].
 		BadAssetId,
-		/// The asset distribution was already claimed!
-		AlreadyClaimed,
-		/// The asset distribution is no longer active.
-		DistributionEnded,
-		/// The asset distribution is still active.
-		DistributionActive,
-		/// The proof provided could not be verified.
-		BadProof,
-		/// The a leaf node was extracted from the proof, but it did not match the expected format.
-		CannotDecodeLeaf,
 	}
 
 	#[pallet::call(weight(<T as Config<I>>::WeightInfo))]
@@ -1847,92 +1798,6 @@ pub mod pallet {
 			)?;
 			Ok(())
 		}
-
-		/// Mint a distribution of assets of a particular class.
-		///
-		/// The origin must be Signed and the sender must be the Issuer of the asset `id`.
-		///
-		/// - `id`: The identifier of the asset to have some amount minted.
-		/// - `merkle_root`: The merkle root of a compact base-16 merkle trie used to authorize
-		///   minting.
-		///
-		/// Emits `DistributionIssued` event when successful.
-		///
-		/// Weight: `O(1)`
-		#[pallet::call_index(33)]
-		pub fn mint_distribution(
-			origin: OriginFor<T>,
-			id: T::AssetIdParameter,
-			merkle_root: T::Hash,
-		) -> DispatchResult {
-			let origin = ensure_signed(origin)?;
-			let id: T::AssetId = id.into();
-			Self::do_mint_distribution(id, merkle_root, Some(origin))?;
-			Ok(())
-		}
-
-		/// Claim a distribution of assets of a particular class.
-		///
-		/// Any signed origin may call this function.
-		///
-		/// - `distribution_id`: The identifier of the distribution.
-		/// - `merkle_proof`: The merkle proof of the account and balance in a compact base-16
-		///   merkle trie used to authorize minting.
-		///
-		/// Emits `Issued` event when successful.
-		///
-		/// Weight: `O(P)` where `P` is the size of the merkle proof.
-		#[pallet::call_index(34)]
-		pub fn claim_distribution(
-			origin: OriginFor<T>,
-			distribution_id: DistributionCounter,
-			merkle_proof: Vec<u8>,
-		) -> DispatchResult {
-			ensure_signed(origin)?;
-			Self::do_claim_distribution(distribution_id, merkle_proof)?;
-			Ok(())
-		}
-
-		/// End the distribution of assets by distribution id.
-		///
-		/// The origin must be Signed and the sender must be the Issuer of the asset `id`.
-		///
-		/// - `distribution_id`: The identifier of the distribution.
-		///
-		/// Emits `DistributionEnded` event when successful.
-		///
-		/// Weight: `O(1)`
-		#[pallet::call_index(35)]
-		pub fn end_distribution(
-			origin: OriginFor<T>,
-			distribution_id: DistributionCounter,
-		) -> DispatchResult {
-			let origin = ensure_signed(origin)?;
-			Self::do_end_distribution(distribution_id, Some(origin))?;
-			Ok(())
-		}
-
-		/// Clean up the distribution tracker of an ended distribution. This function might need to
-		/// be called multiple times to remove all the items from the distribution tracker.
-		///
-		/// Any signed origin may call this function.
-		///
-		/// - `distribution_id`: The identifier of the distribution to clean. It cannot be active.
-		///
-		/// Emits `DistributionPartiallyCleaned` event when some elements have been removed, but
-		/// there are still some left. Emits `DistributionCleaned` when all of the distribution
-		/// history has been removed.
-		///
-		/// Weight: `O(N)` where `N` is the maximum number of elements that can be removed at once.
-		#[pallet::call_index(36)]
-		#[pallet::weight(T::WeightInfo::clean_distribution(100u32))] // TODO
-		pub fn clean_distribution(
-			origin: OriginFor<T>,
-			distribution_id: DistributionCounter,
-		) -> DispatchResultWithPostInfo {
-			ensure_signed(origin)?;
-			Self::do_clean_distribution(distribution_id)
-		}
 	}
 
 	/// Implements [`AccountTouch`] trait.
diff --git a/substrate/frame/assets/src/mock.rs b/substrate/frame/assets/src/mock.rs
index a425425f8c42..2c160840e147 100644
--- a/substrate/frame/assets/src/mock.rs
+++ b/substrate/frame/assets/src/mock.rs
@@ -39,9 +39,8 @@ construct_runtime!(
 	}
 );
 
-pub(crate) type AccountId = u64;
-pub(crate) type AssetId = u32;
-pub(crate) type Balance = u64;
+type AccountId = u64;
+type AssetId = u32;
 
 #[derive_impl(frame_system::config_preludes::TestDefaultConfig)]
 impl frame_system::Config for Test {
diff --git a/substrate/frame/assets/src/tests.rs b/substrate/frame/assets/src/tests.rs
index 21a98a958910..af605c5a3c64 100644
--- a/substrate/frame/assets/src/tests.rs
+++ b/substrate/frame/assets/src/tests.rs
@@ -19,7 +19,6 @@
 
 use super::*;
 use crate::{mock::*, Error};
-use codec::Encode;
 use frame_support::{
 	assert_noop, assert_ok,
 	dispatch::GetDispatchInfo,
@@ -1922,66 +1921,3 @@ fn asset_id_cannot_be_reused() {
 		assert!(Asset::<Test>::contains_key(7));
 	});
 }
-
-#[test]
-fn merklized_distribution_works() {
-	new_test_ext().execute_with(|| {
-		use alloc::collections::BTreeMap;
-
-		// Create asset id 0 controlled by user 1, sufficient so it does not need ED.
-		assert_ok!(Assets::force_create(RuntimeOrigin::root(), 0, 1, true, 1));
-
-		// Offchain, user 1 creates a distribution of tokens.
-		let mut distribution = BTreeMap::<AccountId, Balance>::new();
-		for i in 0..100u64 {
-			distribution.insert(i, i.into());
-		}
-
-		// Maybe the owner gives himself a little extra ;)
-		distribution.insert(1, 1337);
-
-		let flat_distribution: Vec<Vec<u8>> =
-			distribution.into_iter().map(|item| item.encode()).collect();
-
-		let root = binary_merkle_tree::merkle_root::<<Test as frame_system::Config>::Hashing, _>(
-			flat_distribution.clone(),
-		);
-
-		let proof_for_69 = binary_merkle_tree::merkle_proof::<
-			<Test as frame_system::Config>::Hashing,
-			_,
-			_,
-		>(flat_distribution.clone(), 69);
-		let proof_for_1 = binary_merkle_tree::merkle_proof::<
-			<Test as frame_system::Config>::Hashing,
-			_,
-			_,
-		>(flat_distribution.clone(), 1);
-		let proof_for_6 = binary_merkle_tree::merkle_proof::<
-			<Test as frame_system::Config>::Hashing,
-			_,
-			_,
-		>(flat_distribution, 6);
-
-		// Use this trie root for the distribution
-		assert_ok!(Assets::mint_distribution(RuntimeOrigin::signed(1), 0, root));
-
-		// Now users claim their distributions permissionlessly with a proof.
-		assert_ok!(Assets::claim_distribution(RuntimeOrigin::signed(1), 0, proof_for_1.encode()));
-		assert_eq!(Assets::balance(0, 1), 1337);
-
-		// Other users can claim their tokens.
-		assert_ok!(Assets::claim_distribution(RuntimeOrigin::signed(55), 0, proof_for_69.encode()));
-		assert_eq!(Assets::balance(0, 69), 69);
-
-		// Owner (or anyone) can also distribute on behalf of the other users.
-		assert_ok!(Assets::claim_distribution(RuntimeOrigin::signed(1), 0, proof_for_6.encode()));
-		assert_eq!(Assets::balance(0, 6), 6);
-
-		// You cannot double claim.
-		assert_noop!(
-			Assets::claim_distribution(RuntimeOrigin::signed(6), 0, proof_for_6.encode()),
-			Error::<Test>::AlreadyClaimed
-		);
-	});
-}
diff --git a/substrate/frame/assets/src/types.rs b/substrate/frame/assets/src/types.rs
index 44b09652d46b..11edc7d3fcb5 100644
--- a/substrate/frame/assets/src/types.rs
+++ b/substrate/frame/assets/src/types.rs
@@ -317,17 +317,3 @@ where
 			.saturating_mul_int(balance))
 	}
 }
-
-pub type DistributionCounter = u32;
-pub type DistributionProofOf<T, I> =
-	<<T as Config<I>>::VerifyExistenceProof as VerifyExistenceProof>::Proof;
-
-#[derive(Eq, PartialEq, Copy, Clone, RuntimeDebug, Encode, Decode, TypeInfo, MaxEncodedLen)]
-pub struct DistributionInfo<AssetId, Hash> {
-	// The asset id we are distributing.
-	pub asset_id: AssetId,
-	// The merkle root which represents all the balances to distribute.
-	pub merkle_root: Hash,
-	// Whether the distribution is still active.
-	pub active: bool,
-}
diff --git a/substrate/frame/assets/src/weights.rs b/substrate/frame/assets/src/weights.rs
index ac01d5d412a1..57f7e951b73c 100644
--- a/substrate/frame/assets/src/weights.rs
+++ b/substrate/frame/assets/src/weights.rs
@@ -84,10 +84,6 @@ pub trait WeightInfo {
 	fn refund_other() -> Weight;
 	fn block() -> Weight;
 	fn transfer_all() -> Weight;
-	fn mint_distribution() -> Weight;
-	fn claim_distribution() -> Weight;
-	fn end_distribution() -> Weight;
-	fn clean_distribution(n: u32) -> Weight;
 }
 
 /// Weights for `pallet_assets` using the Substrate node and recommended hardware.
@@ -545,23 +541,6 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
-
-	fn mint_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn claim_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn end_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn clean_distribution(_n: u32) -> Weight {
-		Weight::default()
-	}
-
 }
 
 // For backwards compatibility and tests.
@@ -1018,20 +997,4 @@ impl WeightInfo for () {
 			.saturating_add(RocksDbWeight::get().reads(1_u64))
 			.saturating_add(RocksDbWeight::get().writes(1_u64))
 	}
-
-	fn mint_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn claim_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn end_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn clean_distribution(_n: u32) -> Weight {
-		Weight::default()
-	}
 }

From ed30142ad4d5d4627d5c77385eaf680c50720e80 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Wed, 18 Sep 2024 11:25:14 -0400
Subject: [PATCH 50/55] undo weight

---
 .../src/weights/pallet_assets_foreign.rs         | 16 ----------------
 .../src/weights/pallet_assets_local.rs           | 16 ----------------
 .../src/weights/pallet_assets_pool.rs            | 16 ----------------
 .../src/weights/pallet_assets_foreign.rs         | 16 ----------------
 .../src/weights/pallet_assets_local.rs           | 16 ----------------
 .../src/weights/pallet_assets_pool.rs            | 16 ----------------
 6 files changed, 96 deletions(-)

diff --git a/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_foreign.rs b/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_foreign.rs
index f099987cabc2..c76c1137335a 100644
--- a/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_foreign.rs
+++ b/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_foreign.rs
@@ -541,20 +541,4 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
-
-	fn mint_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn claim_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn end_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn clean_distribution(_n: u32) -> Weight {
-		Weight::default()
-	}
 }
diff --git a/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_local.rs b/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_local.rs
index 5e2e9e421fc5..cf4f60042bc6 100644
--- a/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_local.rs
+++ b/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_local.rs
@@ -538,20 +538,4 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
-
-	fn mint_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn claim_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn end_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn clean_distribution(_n: u32) -> Weight {
-		Weight::default()
-	}
 }
diff --git a/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_pool.rs b/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_pool.rs
index dabb5a103ac5..2cd85de00989 100644
--- a/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_pool.rs
+++ b/cumulus/parachains/runtimes/assets/asset-hub-rococo/src/weights/pallet_assets_pool.rs
@@ -538,20 +538,4 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
-
-	fn mint_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn claim_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn end_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn clean_distribution(_n: u32) -> Weight {
-		Weight::default()
-	}
 }
diff --git a/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_foreign.rs b/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_foreign.rs
index 60829e471df7..2692de9aeb50 100644
--- a/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_foreign.rs
+++ b/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_foreign.rs
@@ -547,20 +547,4 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
-
-	fn mint_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn claim_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn end_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn clean_distribution(_n: u32) -> Weight {
-		Weight::default()
-	}
 }
diff --git a/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_local.rs b/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_local.rs
index a25eac274fb0..d2e12549a45c 100644
--- a/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_local.rs
+++ b/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_local.rs
@@ -545,20 +545,4 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
-
-	fn mint_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn claim_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn end_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn clean_distribution(_n: u32) -> Weight {
-		Weight::default()
-	}
 }
diff --git a/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_pool.rs b/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_pool.rs
index 7676ffa5aaab..8368f6e583cc 100644
--- a/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_pool.rs
+++ b/cumulus/parachains/runtimes/assets/asset-hub-westend/src/weights/pallet_assets_pool.rs
@@ -539,20 +539,4 @@ impl<T: frame_system::Config> pallet_assets::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(1_u64))
 			.saturating_add(T::DbWeight::get().writes(1_u64))
 	}
-
-	fn mint_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn claim_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn end_distribution() -> Weight {
-		Weight::default()
-	}
-
-	fn clean_distribution(_n: u32) -> Weight {
-		Weight::default()
-	}
 }

From 39e78d73c1e44042254a0125bfccd84c8b2a2f64 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Wed, 18 Sep 2024 11:27:18 -0400
Subject: [PATCH 51/55] Update substrate/frame/support/src/traits/proving.rs

---
 substrate/frame/support/src/traits/proving.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/substrate/frame/support/src/traits/proving.rs b/substrate/frame/support/src/traits/proving.rs
index 549f8f719ba4..9e4033c46b14 100644
--- a/substrate/frame/support/src/traits/proving.rs
+++ b/substrate/frame/support/src/traits/proving.rs
@@ -79,7 +79,7 @@ pub struct SixteenPatriciaMerkleTreeProver<H>(core::marker::PhantomData<H>);
 
 impl<H: Hasher> VerifyExistenceProof for SixteenPatriciaMerkleTreeProver<H>
 where
-	H::Out: Decode + Encode,
+	H::Out: Encode + Decode,
 {
 	type Proof = SixteenPatriciaMerkleTreeExistenceProof;
 	type Hash = H::Out;

From 450b6c6fff040fdce49030d51a390b8e627aa0c1 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 19 Sep 2024 17:03:31 -0400
Subject: [PATCH 52/55] consolidate apis

---
 Cargo.lock                                    |  1 -
 .../runtime/src/proving_trie/base16.rs        | 35 ++++++++-----------
 2 files changed, 14 insertions(+), 22 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index dfc9e6a9e5fd..f1abb2e0da35 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -10728,7 +10728,6 @@ dependencies = [
 name = "pallet-assets"
 version = "29.1.0"
 dependencies = [
- "binary-merkle-tree",
  "frame-benchmarking",
  "frame-support",
  "frame-system",
diff --git a/substrate/primitives/runtime/src/proving_trie/base16.rs b/substrate/primitives/runtime/src/proving_trie/base16.rs
index 69dbe77fa9d6..3d52c0f6e78b 100644
--- a/substrate/primitives/runtime/src/proving_trie/base16.rs
+++ b/substrate/primitives/runtime/src/proving_trie/base16.rs
@@ -113,14 +113,14 @@ where
 	}
 }
 
-/// Verify the existence or non-existence of `key` and `value` in a given trie root and proof.
+/// Verify the existence of `key` and `value` in a given trie root and proof.
 ///
 /// Proofs must be created with latest substrate trie format (`LayoutV1`).
 pub fn verify_single_value_proof<Hashing, Key, Value>(
 	root: Hashing::Out,
 	proof: &[u8],
 	key: Key,
-	maybe_value: Option<Value>,
+	value: Value,
 ) -> Result<(), DispatchError>
 where
 	Hashing: sp_core::Hasher,
@@ -132,18 +132,18 @@ where
 	sp_trie::verify_trie_proof::<LayoutV1<Hashing>, _, _, _>(
 		&root,
 		&structured_proof,
-		&[(key.encode(), maybe_value.map(|value| value.encode()))],
+		&[(key.encode(), Some(value.encode()))],
 	)
 	.map_err(|err| TrieError::from(err).into())
 }
 
-/// Verify the existence or non-existence of multiple `items` in a given trie root and proof.
+/// Verify the existence of multiple `items` in a given trie root and proof.
 ///
 /// Proofs must be created with latest substrate trie format (`LayoutV1`).
 pub fn verify_proof<Hashing, Key, Value>(
 	root: Hashing::Out,
 	proof: &[u8],
-	items: &[(Key, Option<Value>)],
+	items: &[(Key, Value)],
 ) -> Result<(), DispatchError>
 where
 	Hashing: sp_core::Hasher,
@@ -154,7 +154,7 @@ where
 		Decode::decode(&mut &proof[..]).map_err(|_| TrieError::DecodeError)?;
 	let items_encoded = items
 		.into_iter()
-		.map(|(key, maybe_value)| (key.encode(), maybe_value.as_ref().map(|value| value.encode())))
+		.map(|(key, value)| (key.encode(), Some(value.encode())))
 		.collect::<Vec<(Vec<u8>, Option<Vec<u8>>)>>();
 
 	sp_trie::verify_trie_proof::<LayoutV1<Hashing>, _, _, _>(
@@ -226,7 +226,7 @@ mod tests {
 						root,
 						&proof,
 						i,
-						Some(u128::from(i))
+						u128::from(i)
 					),
 					Ok(())
 				);
@@ -236,7 +236,7 @@ mod tests {
 						root,
 						&proof,
 						i,
-						Some(u128::from(i + 1))
+						u128::from(i + 1)
 					),
 					Err(TrieError::RootMismatch.into())
 				);
@@ -245,14 +245,7 @@ mod tests {
 					root,
 					&proof,
 					i,
-					Some(u128::from(i))
-				)
-				.is_err());
-				assert!(verify_single_value_proof::<BlakeTwo256, _, _>(
-					root,
-					&proof,
-					i,
-					None::<u128>
+					u128::from(i)
 				)
 				.is_err());
 			}
@@ -265,8 +258,8 @@ mod tests {
 		let root = *balance_trie.root();
 
 		// Create a proof for a valid and invalid key.
-		let proof = balance_trie.create_proof(&[6u32, 69u32, 6969u32]).unwrap();
-		let items = [(6u32, Some(6u128)), (69u32, Some(69u128)), (6969u32, None)];
+		let proof = balance_trie.create_proof(&[6u32, 9u32, 69u32]).unwrap();
+		let items = [(6u32, 6u128), (9u32, 9u128), (69u32, 69u128)];
 
 		assert_eq!(verify_proof::<BlakeTwo256, _, _>(root, &proof, &items), Ok(()));
 	}
@@ -281,7 +274,7 @@ mod tests {
 
 		// Correct data verifies successfully
 		assert_eq!(
-			verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, Some(6u128)),
+			verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, 6u128),
 			Ok(())
 		);
 
@@ -291,7 +284,7 @@ mod tests {
 				Default::default(),
 				&proof,
 				6u32,
-				Some(6u128)
+				6u128
 			),
 			Err(TrieError::RootMismatch.into())
 		);
@@ -301,7 +294,7 @@ mod tests {
 
 		// Fail to verify data with the wrong proof
 		assert_eq!(
-			verify_single_value_proof::<BlakeTwo256, _, _>(root, &bad_proof, 6u32, Some(6u128)),
+			verify_single_value_proof::<BlakeTwo256, _, _>(root, &bad_proof, 6u32, 6u128),
 			Err(TrieError::ExtraneousHashReference.into())
 		);
 	}

From 240a77739c1c5dd7b9261efa789735b085f2e531 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 19 Sep 2024 17:36:39 -0400
Subject: [PATCH 53/55] create a proving trie trait

---
 .../runtime/src/proving_trie/base16.rs        | 113 +++++++++---------
 .../runtime/src/proving_trie/mod.rs           |  31 ++++-
 2 files changed, 84 insertions(+), 60 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie/base16.rs b/substrate/primitives/runtime/src/proving_trie/base16.rs
index 3d52c0f6e78b..a3f2d49dc636 100644
--- a/substrate/primitives/runtime/src/proving_trie/base16.rs
+++ b/substrate/primitives/runtime/src/proving_trie/base16.rs
@@ -24,7 +24,7 @@
 //! Proofs are created with latest substrate trie format (`LayoutV1`), and are not compatible with
 //! proofs using `LayoutV0`.
 
-use super::TrieError;
+use super::{ProvingTrie, TrieError};
 use crate::{Decode, DispatchError, Encode};
 use sp_std::vec::Vec;
 use sp_trie::{
@@ -45,13 +45,38 @@ where
 }
 
 impl<Hashing, Key, Value> BasicProvingTrie<Hashing, Key, Value>
+where
+	Hashing: sp_core::Hasher,
+	Key: Encode,
+	Value: Encode,
+{
+	/// Create a compact merkle proof needed to prove all `keys` and their values are in the trie.
+	///
+	/// This function makes a proof with latest substrate trie format (`LayoutV1`), and is not
+	/// compatible with `LayoutV0`.
+	///
+	/// When verifying the proof created by this function, you must include all of the keys and
+	/// values of the proof, else the verifier will complain that extra nodes are provided in the
+	/// proof that are not needed.
+	pub fn create_multi_value_proof(&self, keys: &[Key]) -> Result<Vec<u8>, DispatchError> {
+		sp_trie::generate_trie_proof::<LayoutV1<Hashing>, _, _, _>(
+			&self.db,
+			self.root,
+			&keys.into_iter().map(|k| k.encode()).collect::<Vec<Vec<u8>>>(),
+		)
+		.map_err(|err| TrieError::from(*err).into())
+		.map(|structured_proof| structured_proof.encode())
+	}
+}
+
+impl<Hashing, Key, Value> ProvingTrie<Hashing, Key, Value> for BasicProvingTrie<Hashing, Key, Value>
 where
 	Hashing: sp_core::Hasher,
 	Key: Encode,
 	Value: Encode,
 {
 	/// Create a new instance of a `ProvingTrie` using an iterator of key/value pairs.
-	pub fn generate_for<I>(items: I) -> Result<Self, DispatchError>
+	fn generate_for<I>(items: I) -> Result<Self, DispatchError>
 	where
 		I: IntoIterator<Item = (Key, Value)>,
 	{
@@ -70,13 +95,13 @@ where
 	}
 
 	/// Access the underlying trie root.
-	pub fn root(&self) -> &Hashing::Out {
+	fn root(&self) -> &Hashing::Out {
 		&self.root
 	}
 
 	/// Query a value contained within the current trie. Returns `None` if the
 	/// nodes within the current `MemoryDB` are insufficient to query the item.
-	pub fn query(&self, key: Key) -> Option<Value>
+	fn query(&self, key: Key) -> Option<Value>
 	where
 		Value: Decode,
 	{
@@ -86,37 +111,31 @@ where
 			.and_then(|raw| Value::decode(&mut &*raw).ok())
 	}
 
-	/// Create a compact merkle proof needed to prove all `keys` and their values are in the trie.
+	/// Create a compact merkle proof needed to prove a single key and its value are in the trie.
 	///
 	/// This function makes a proof with latest substrate trie format (`LayoutV1`), and is not
 	/// compatible with `LayoutV0`.
-	///
-	/// When verifying the proof created by this function, you must include all of the keys and
-	/// values of the proof, else the verifier will complain that extra nodes are provided in the
-	/// proof that are not needed.
-	pub fn create_proof(&self, keys: &[Key]) -> Result<Vec<u8>, DispatchError> {
-		sp_trie::generate_trie_proof::<LayoutV1<Hashing>, _, _, _>(
-			&self.db,
-			self.root,
-			&keys.into_iter().map(|k| k.encode()).collect::<Vec<Vec<u8>>>(),
-		)
-		.map_err(|err| TrieError::from(*err).into())
-		.map(|structured_proof| structured_proof.encode())
+	fn create_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError> {
+		self.create_multi_value_proof(&[key])
 	}
 
-	/// Create a compact merkle proof needed to prove a single key and its value are in the trie.
+	/// Verify the existence of `key` and `value` in a given trie root and proof.
 	///
-	/// This function makes a proof with latest substrate trie format (`LayoutV1`), and is not
-	/// compatible with `LayoutV0`.
-	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError> {
-		self.create_proof(&[key])
+	/// Proofs must be created with latest substrate trie format (`LayoutV1`).
+	fn verify_proof(
+		root: Hashing::Out,
+		proof: &[u8],
+		key: Key,
+		value: Value,
+	) -> Result<(), DispatchError> {
+		verify_proof::<Hashing, Key, Value>(root, proof, key, value)
 	}
 }
 
 /// Verify the existence of `key` and `value` in a given trie root and proof.
 ///
 /// Proofs must be created with latest substrate trie format (`LayoutV1`).
-pub fn verify_single_value_proof<Hashing, Key, Value>(
+pub fn verify_proof<Hashing, Key, Value>(
 	root: Hashing::Out,
 	proof: &[u8],
 	key: Key,
@@ -140,7 +159,7 @@ where
 /// Verify the existence of multiple `items` in a given trie root and proof.
 ///
 /// Proofs must be created with latest substrate trie format (`LayoutV1`).
-pub fn verify_proof<Hashing, Key, Value>(
+pub fn verify_multi_value_proof<Hashing, Key, Value>(
 	root: Hashing::Out,
 	proof: &[u8],
 	items: &[(Key, Value)],
@@ -216,38 +235,22 @@ mod tests {
 		let root = *balance_trie.root();
 
 		// Create a proof for a valid key.
-		let proof = balance_trie.create_single_value_proof(6u32).unwrap();
+		let proof = balance_trie.create_proof(6u32).unwrap();
 
 		// Assert key is provable, all other keys are invalid.
 		for i in 0..200u32 {
 			if i == 6 {
 				assert_eq!(
-					verify_single_value_proof::<BlakeTwo256, _, _>(
-						root,
-						&proof,
-						i,
-						u128::from(i)
-					),
+					verify_proof::<BlakeTwo256, _, _>(root, &proof, i, u128::from(i)),
 					Ok(())
 				);
 				// Wrong value is invalid.
 				assert_eq!(
-					verify_single_value_proof::<BlakeTwo256, _, _>(
-						root,
-						&proof,
-						i,
-						u128::from(i + 1)
-					),
+					verify_proof::<BlakeTwo256, _, _>(root, &proof, i, u128::from(i + 1)),
 					Err(TrieError::RootMismatch.into())
 				);
 			} else {
-				assert!(verify_single_value_proof::<BlakeTwo256, _, _>(
-					root,
-					&proof,
-					i,
-					u128::from(i)
-				)
-				.is_err());
+				assert!(verify_proof::<BlakeTwo256, _, _>(root, &proof, i, u128::from(i)).is_err());
 			}
 		}
 	}
@@ -258,10 +261,10 @@ mod tests {
 		let root = *balance_trie.root();
 
 		// Create a proof for a valid and invalid key.
-		let proof = balance_trie.create_proof(&[6u32, 9u32, 69u32]).unwrap();
+		let proof = balance_trie.create_multi_value_proof(&[6u32, 9u32, 69u32]).unwrap();
 		let items = [(6u32, 6u128), (9u32, 9u128), (69u32, 69u128)];
 
-		assert_eq!(verify_proof::<BlakeTwo256, _, _>(root, &proof, &items), Ok(()));
+		assert_eq!(verify_multi_value_proof::<BlakeTwo256, _, _>(root, &proof, &items), Ok(()));
 	}
 
 	#[test]
@@ -270,31 +273,23 @@ mod tests {
 		let root = *balance_trie.root();
 
 		// Create a proof for a valid key.
-		let proof = balance_trie.create_single_value_proof(6u32).unwrap();
+		let proof = balance_trie.create_proof(6u32).unwrap();
 
 		// Correct data verifies successfully
-		assert_eq!(
-			verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, 6u128),
-			Ok(())
-		);
+		assert_eq!(verify_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, 6u128), Ok(()));
 
 		// Fail to verify proof with wrong root
 		assert_eq!(
-			verify_single_value_proof::<BlakeTwo256, _, _>(
-				Default::default(),
-				&proof,
-				6u32,
-				6u128
-			),
+			verify_proof::<BlakeTwo256, _, _>(Default::default(), &proof, 6u32, 6u128),
 			Err(TrieError::RootMismatch.into())
 		);
 
 		// Crete a bad proof.
-		let bad_proof = balance_trie.create_single_value_proof(99u32).unwrap();
+		let bad_proof = balance_trie.create_proof(99u32).unwrap();
 
 		// Fail to verify data with the wrong proof
 		assert_eq!(
-			verify_single_value_proof::<BlakeTwo256, _, _>(root, &bad_proof, 6u32, 6u128),
+			verify_proof::<BlakeTwo256, _, _>(root, &bad_proof, 6u32, 6u128),
 			Err(TrieError::ExtraneousHashReference.into())
 		);
 	}
diff --git a/substrate/primitives/runtime/src/proving_trie/mod.rs b/substrate/primitives/runtime/src/proving_trie/mod.rs
index 60f11645a48f..e1764bb851b9 100644
--- a/substrate/primitives/runtime/src/proving_trie/mod.rs
+++ b/substrate/primitives/runtime/src/proving_trie/mod.rs
@@ -20,9 +20,10 @@
 pub mod base16;
 pub mod base2;
 
-use crate::{Decode, Encode, MaxEncodedLen, TypeInfo};
+use crate::{Decode, DispatchError, Encode, MaxEncodedLen, TypeInfo};
 #[cfg(feature = "serde")]
 use crate::{Deserialize, Serialize};
+use sp_std::vec::Vec;
 use sp_trie::{trie_types::TrieError as SpTrieError, VerifyError};
 
 /// A runtime friendly error type for tries.
@@ -111,3 +112,31 @@ impl From<TrieError> for &'static str {
 		}
 	}
 }
+
+/// An interface for creating, interacting with, and creating proofs in a merkle trie.
+pub trait ProvingTrie<Hashing, Key, Value>
+where
+	Self: Sized,
+	Hashing: sp_core::Hasher,
+{
+	/// Create a new instance of a `ProvingTrie` using an iterator of key/value pairs.
+	fn generate_for<I>(items: I) -> Result<Self, DispatchError>
+	where
+		I: IntoIterator<Item = (Key, Value)>;
+	/// Access the underlying trie root.
+	fn root(&self) -> &Hashing::Out;
+	/// Query a value contained within the current trie. Returns `None` if the
+	/// the value does not exist in the trie.
+	fn query(&self, key: Key) -> Option<Value>
+	where
+		Value: Decode;
+	/// Create a proof that can be used to verify a key and its value are in the trie.
+	fn create_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError>;
+	/// Verify the existence of `key` and `value` in a given trie root and proof.
+	fn verify_proof(
+		root: Hashing::Out,
+		proof: &[u8],
+		key: Key,
+		value: Value,
+	) -> Result<(), DispatchError>;
+}

From 0f1e28bc4f9cb8e5bb3b0c85b9ae1d2a5a33fdd7 Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 19 Sep 2024 17:46:23 -0400
Subject: [PATCH 54/55] implement trait for base2

---
 .../runtime/src/proving_trie/base16.rs        |  7 +---
 .../runtime/src/proving_trie/base2.rs         | 39 +++++++++++--------
 .../runtime/src/proving_trie/mod.rs           |  4 +-
 3 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie/base16.rs b/substrate/primitives/runtime/src/proving_trie/base16.rs
index a3f2d49dc636..c1379ff4949a 100644
--- a/substrate/primitives/runtime/src/proving_trie/base16.rs
+++ b/substrate/primitives/runtime/src/proving_trie/base16.rs
@@ -73,7 +73,7 @@ impl<Hashing, Key, Value> ProvingTrie<Hashing, Key, Value> for BasicProvingTrie<
 where
 	Hashing: sp_core::Hasher,
 	Key: Encode,
-	Value: Encode,
+	Value: Encode + Decode,
 {
 	/// Create a new instance of a `ProvingTrie` using an iterator of key/value pairs.
 	fn generate_for<I>(items: I) -> Result<Self, DispatchError>
@@ -101,10 +101,7 @@ where
 
 	/// Query a value contained within the current trie. Returns `None` if the
 	/// nodes within the current `MemoryDB` are insufficient to query the item.
-	fn query(&self, key: Key) -> Option<Value>
-	where
-		Value: Decode,
-	{
+	fn query(&self, key: Key) -> Option<Value> {
 		let trie = TrieDBBuilder::new(&self.db, &self.root).build();
 		key.using_encoded(|s| trie.get(s))
 			.ok()?
diff --git a/substrate/primitives/runtime/src/proving_trie/base2.rs b/substrate/primitives/runtime/src/proving_trie/base2.rs
index 216910f2c7d0..51086fa9b073 100644
--- a/substrate/primitives/runtime/src/proving_trie/base2.rs
+++ b/substrate/primitives/runtime/src/proving_trie/base2.rs
@@ -20,9 +20,9 @@
 //! this library is designed to work more easily with runtime native types, which simply need to
 //! implement `Encode`/`Decode`.
 
-use super::TrieError;
+use super::{ProvingTrie, TrieError};
 use crate::{Decode, DispatchError, Encode};
-use binary_merkle_tree::{merkle_proof, merkle_root, verify_proof, MerkleProof};
+use binary_merkle_tree::{merkle_proof, merkle_root, MerkleProof};
 use sp_std::{collections::btree_map::BTreeMap, vec::Vec};
 
 /// A helper structure for building a basic base-2 merkle trie and creating compact proofs for that
@@ -37,14 +37,15 @@ where
 	_phantom: core::marker::PhantomData<(Key, Value)>,
 }
 
-impl<Hashing, Key, Value> BasicProvingTrie<Hashing, Key, Value>
+impl<Hashing, Key, Value> ProvingTrie<Hashing, Key, Value> for BasicProvingTrie<Hashing, Key, Value>
 where
 	Hashing: sp_core::Hasher,
-	Key: Encode + Ord,
-	Value: Encode,
+	Hashing::Out: Encode + Decode,
+	Key: Encode + Decode + Ord,
+	Value: Encode + Decode + Clone,
 {
 	/// Create a new instance of a `ProvingTrie` using an iterator of key/value pairs.
-	pub fn generate_for<I>(items: I) -> Result<Self, DispatchError>
+	fn generate_for<I>(items: I) -> Result<Self, DispatchError>
 	where
 		I: IntoIterator<Item = (Key, Value)>,
 	{
@@ -57,26 +58,20 @@ where
 	}
 
 	/// Access the underlying trie root.
-	pub fn root(&self) -> &Hashing::Out {
+	fn root(&self) -> &Hashing::Out {
 		&self.root
 	}
 
 	/// Query a value contained within the current trie. Returns `None` if the
 	/// nodes within the current `db` are insufficient to query the item.
-	pub fn query(&self, key: Key) -> Option<Value>
-	where
-		Value: Decode + Clone,
-	{
+	fn query(&self, key: Key) -> Option<Value> {
 		self.db.get(&key).cloned()
 	}
 
 	/// Create a compact merkle proof needed to prove a single key and its value are in the trie.
 	/// Returns `None` if the nodes within the current `db` are insufficient to create a
 	/// proof.
-	pub fn create_single_value_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError>
-	where
-		Hashing::Out: Encode,
-	{
+	fn create_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError> {
 		let mut encoded = Vec::with_capacity(self.db.len());
 		let mut found_index = None;
 
@@ -94,10 +89,20 @@ where
 		let proof = merkle_proof::<Hashing, Vec<Vec<u8>>, Vec<u8>>(encoded, index as u32);
 		Ok(proof.encode())
 	}
+
+	/// Verify the existence of `key` and `value` in a given trie root and proof.
+	fn verify_proof(
+		root: Hashing::Out,
+		proof: &[u8],
+		key: Key,
+		value: Value,
+	) -> Result<(), DispatchError> {
+		verify_proof::<Hashing, Key, Value>(root, proof, key, value)
+	}
 }
 
 /// Verify the existence of `key` and `value` in a given trie root and proof.
-pub fn verify_single_value_proof<Hashing, Key, Value>(
+pub fn verify_proof<Hashing, Key, Value>(
 	root: Hashing::Out,
 	proof: &[u8],
 	key: Key,
@@ -119,7 +124,7 @@ where
 		return Err(TrieError::ValueMismatch.into());
 	}
 
-	if verify_proof::<Hashing, _, _>(
+	if binary_merkle_tree::verify_proof::<Hashing, _, _>(
 		&decoded_proof.root,
 		decoded_proof.proof,
 		decoded_proof.number_of_leaves,
diff --git a/substrate/primitives/runtime/src/proving_trie/mod.rs b/substrate/primitives/runtime/src/proving_trie/mod.rs
index e1764bb851b9..267796546e44 100644
--- a/substrate/primitives/runtime/src/proving_trie/mod.rs
+++ b/substrate/primitives/runtime/src/proving_trie/mod.rs
@@ -127,9 +127,7 @@ where
 	fn root(&self) -> &Hashing::Out;
 	/// Query a value contained within the current trie. Returns `None` if the
 	/// the value does not exist in the trie.
-	fn query(&self, key: Key) -> Option<Value>
-	where
-		Value: Decode;
+	fn query(&self, key: Key) -> Option<Value>;
 	/// Create a proof that can be used to verify a key and its value are in the trie.
 	fn create_proof(&self, key: Key) -> Result<Vec<u8>, DispatchError>;
 	/// Verify the existence of `key` and `value` in a given trie root and proof.

From 375dc12eed329abf70dd5fc96916f65d3492d6cf Mon Sep 17 00:00:00 2001
From: Shawn Tabrizi <shawntabrizi@gmail.com>
Date: Thu, 19 Sep 2024 18:09:26 -0400
Subject: [PATCH 55/55] fix tests

---
 .../runtime/src/proving_trie/base2.rs         | 16 +++++-----
 .../runtime/src/proving_trie/mod.rs           | 32 +++++++++++++++++++
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/substrate/primitives/runtime/src/proving_trie/base2.rs b/substrate/primitives/runtime/src/proving_trie/base2.rs
index 51086fa9b073..a8c1e952dfac 100644
--- a/substrate/primitives/runtime/src/proving_trie/base2.rs
+++ b/substrate/primitives/runtime/src/proving_trie/base2.rs
@@ -187,18 +187,18 @@ mod tests {
 		let root = *balance_trie.root();
 
 		// Create a proof for a valid key.
-		let proof = balance_trie.create_single_value_proof(6u32).unwrap();
+		let proof = balance_trie.create_proof(6u32).unwrap();
 
 		// Assert key is provable, all other keys are invalid.
 		for i in 0..200u32 {
 			if i == 6 {
 				assert_eq!(
-					verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, i, u128::from(i)),
+					verify_proof::<BlakeTwo256, _, _>(root, &proof, i, u128::from(i)),
 					Ok(())
 				);
 				// Wrong value is invalid.
 				assert_eq!(
-					verify_single_value_proof::<BlakeTwo256, _, _>(
+					verify_proof::<BlakeTwo256, _, _>(
 						root,
 						&proof,
 						i,
@@ -207,7 +207,7 @@ mod tests {
 					Err(TrieError::ValueMismatch.into())
 				);
 			} else {
-				assert!(verify_single_value_proof::<BlakeTwo256, _, _>(
+				assert!(verify_proof::<BlakeTwo256, _, _>(
 					root,
 					&proof,
 					i,
@@ -224,23 +224,23 @@ mod tests {
 		let root = *balance_trie.root();
 
 		// Create a proof for a valid key.
-		let proof = balance_trie.create_single_value_proof(6u32).unwrap();
+		let proof = balance_trie.create_proof(6u32).unwrap();
 
 		// Correct data verifies successfully
 		assert_eq!(
-			verify_single_value_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, 6u128),
+			verify_proof::<BlakeTwo256, _, _>(root, &proof, 6u32, 6u128),
 			Ok(())
 		);
 
 		// Fail to verify proof with wrong root
 		assert_eq!(
-			verify_single_value_proof::<BlakeTwo256, _, _>(Default::default(), &proof, 6u32, 6u128),
+			verify_proof::<BlakeTwo256, _, _>(Default::default(), &proof, 6u32, 6u128),
 			Err(TrieError::RootMismatch.into())
 		);
 
 		// Fail to verify proof with wrong data
 		assert_eq!(
-			verify_single_value_proof::<BlakeTwo256, _, _>(root, &[], 6u32, 6u128),
+			verify_proof::<BlakeTwo256, _, _>(root, &[], 6u32, 6u128),
 			Err(TrieError::IncompleteProof.into())
 		);
 	}
diff --git a/substrate/primitives/runtime/src/proving_trie/mod.rs b/substrate/primitives/runtime/src/proving_trie/mod.rs
index 267796546e44..8f8c4084ac98 100644
--- a/substrate/primitives/runtime/src/proving_trie/mod.rs
+++ b/substrate/primitives/runtime/src/proving_trie/mod.rs
@@ -138,3 +138,35 @@ where
 		value: Value,
 	) -> Result<(), DispatchError>;
 }
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use crate::traits::BlakeTwo256;
+	use sp_core::H256;
+	use sp_std::collections::btree_map::BTreeMap;
+
+	// A trie which simulates a trie of accounts (u32) and balances (u128).
+	type BalanceTrie2 = base2::BasicProvingTrie<BlakeTwo256, u32, u128>;
+	type BalanceTrie16 = base16::BasicProvingTrie<BlakeTwo256, u32, u128>;
+
+	#[test]
+	fn basic_api_usage_base_2() {
+		let balance_trie = BalanceTrie2::generate_for((0..100u32).map(|i| (i, i.into()))).unwrap();
+		let root = *balance_trie.root();
+		assert_eq!(balance_trie.query(69), Some(69));
+		assert_eq!(balance_trie.query(6969), None);
+		let proof = balance_trie.create_proof(69u32).unwrap();
+		assert_eq!(BalanceTrie2::verify_proof(root, &proof, 69u32, 69u128), Ok(()));
+	}
+
+	#[test]
+	fn basic_api_usage_base_16() {
+		let balance_trie = BalanceTrie16::generate_for((0..100u32).map(|i| (i, i.into()))).unwrap();
+		let root = *balance_trie.root();
+		assert_eq!(balance_trie.query(69), Some(69));
+		assert_eq!(balance_trie.query(6969), None);
+		let proof = balance_trie.create_proof(69u32).unwrap();
+		assert_eq!(BalanceTrie16::verify_proof(root, &proof, 69u32, 69u128), Ok(()));
+	}
+}