Storage Map with Counters

[kianenigma]

For works such as staking, we want to be able to easily track the count of items of a map without the need to go through the expensive ::iter().count().

Historically, the pattern to do this was to manually do this. I think now is a good time to start elaborating on how we can achieve a nicer way to do this. The ideal syntax for me would be to do:

#[pallet::storage]
#[counter = validators_counter]
type Validators = StorageMap<_, _, _>

decl_storage! {
     pub Validators get(fn validators) counter(fn validators_counter): map ...
}

Essentially, this would use the name of the counter function given to generate a new unique storage key:

hash(pallet_name) ++ hash(storage_name ++ counter_name)

And store a new u32 under this key. Then we need to somehow override insert, mutate, remove and all other mutating calls, or have hooks for them, that would update the counter accordingly. I started prototyping this, and came up with something along these lines:

diff --git a/frame/support/src/storage/generator/map.rs b/frame/support/src/storage/generator/map.rs
index 9abc78839..9c9cbf755 100644
--- a/frame/support/src/storage/generator/map.rs
+++ b/frame/support/src/storage/generator/map.rs
@@ -42,6 +42,8 @@ pub trait StorageMap<K: FullEncode, V: FullCodec> {
 	/// Hasher. Used for generating final key.
 	type Hasher: StorageHasher;
 
+	type Hooks: storage::types::map::StorageMapHooks;
+
 	/// Module prefix. Used for generating final key.
 	fn module_prefix() -> &'static [u8];
 
diff --git a/frame/support/src/storage/types/map.rs b/frame/support/src/storage/types/map.rs
index 4af28a77c..83871eb96 100644
--- a/frame/support/src/storage/types/map.rs
+++ b/frame/support/src/storage/types/map.rs
@@ -23,12 +23,61 @@ use crate::{
 	storage::{
 		StorageAppend, StorageDecodeLength,
 		types::{OptionQuery, QueryKindTrait, OnEmptyGetter},
+		unhashed,
 	},
 	traits::{GetDefault, StorageInstance},
+	hash::{Twox128, StorageHasher},
 };
 use frame_metadata::{DefaultByteGetter, StorageEntryModifier};
 use sp_std::prelude::*;
 
+pub trait StorageMapHooks {
+	/// A key is being inserted; do something about it.
+	fn on_insert<I: StorageInstance>();
+	/// A key is being removed; do something about it.
+	fn on_remove<I: StorageInstance>();
+}
+
+/// A counter that is... no counter 🧠.
+pub struct Nothing;
+impl StorageMapHooks for Nothing {
+	fn on_insert<I: StorageInstance>() {}
+	fn on_remove<I: StorageInstance>() {}
+}
+
+/// Have a counter implemented as a hook.
+pub struct StorageMapCounter;
+impl StorageMapCounter {
+	const SUFFIX: &'static [u8] = b"__COUNTER__";
+	fn key_for<I: StorageInstance>() -> Vec<u8> {
+		let pallet = Twox128::hash(I::pallet_prefix().as_bytes());
+		let mut map_name = I::STORAGE_PREFIX.as_bytes().to_vec();
+		map_name.extend(Self::SUFFIX);
+		let suffix = Twox128::hash(&map_name);
+
+		let mut final_key = Vec::with_capacity(pallet.len() + suffix.len());
+		final_key.extend_from_slice(&pallet[..]);
+		final_key.extend_from_slice(&suffix[..]);
+
+		final_key
+	}
+}
+
+impl StorageMapHooks for StorageMapCounter {
+	fn on_insert<I: StorageInstance>() {
+		let key = Self::key_for::<I>();
+		let mut counter = unhashed::get::<u32>(&key).unwrap_or_default();
+		counter += 1u32;
+		unhashed::put(&key, &counter);
+	}
+	fn on_remove<I: StorageInstance>() {
+		let key = Self::key_for::<I>();
+		let mut counter = unhashed::get::<u32>(&key).unwrap_or_default();
+		counter -= 1;
+		unhashed::put(&key, &counter);
+	}
+}
+
 /// A type that allow to store value for given key. Allowing to insert/remove/iterate on values.
 ///
 /// Each value is stored at:
@@ -42,23 +91,25 @@ use sp_std::prelude::*;
 ///
 /// If the keys are not trusted (e.g. can be set by a user), a cryptographic `hasher` such as
 /// `blake2_128_concat` must be used.  Otherwise, other values in storage can be compromised.
-pub struct StorageMap<Prefix, Hasher, Key, Value, QueryKind=OptionQuery, OnEmpty=GetDefault>(
-	core::marker::PhantomData<(Prefix, Hasher, Key, Value, QueryKind, OnEmpty)>
+pub struct StorageMap<Prefix, Hasher, Key, Value, QueryKind=OptionQuery, OnEmpty=GetDefault, Hooks=Nothing>(
+	core::marker::PhantomData<(Prefix, Hasher, Key, Value, QueryKind, OnEmpty, Hooks)>
 );
 
-impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty>
+impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty, Hooks>
 	crate::storage::generator::StorageMap<Key, Value>
-	for StorageMap<Prefix, Hasher, Key, Value, QueryKind, OnEmpty>
+	for StorageMap<Prefix, Hasher, Key, Value, QueryKind, OnEmpty, Hooks>
 where
 	Prefix: StorageInstance,
 	Hasher: crate::hash::StorageHasher,
 	Key: FullCodec,
 	Value: FullCodec,
 	QueryKind: QueryKindTrait<Value, OnEmpty>,
+	Hooks: StorageMapHooks,
 	OnEmpty: crate::traits::Get<QueryKind::Query> + 'static,
 {
 	type Query = QueryKind::Query;
 	type Hasher = Hasher;
+	type Hooks = Hooks;
 	fn module_prefix() -> &'static [u8] {
 		Prefix::pallet_prefix().as_bytes()
 	}
diff --git a/frame/support/src/storage/types/mod.rs b/frame/support/src/storage/types/mod.rs
index 5bb6684b7..fdc815411 100644
--- a/frame/support/src/storage/types/mod.rs
+++ b/frame/support/src/storage/types/mod.rs
@@ -22,7 +22,7 @@ use codec::FullCodec;
 use frame_metadata::{DefaultByte, StorageEntryModifier};
 
 mod value;
-mod map;
+pub(crate) mod map;
 mod double_map;
 
 pub use value::{StorageValue, StorageValueMetadata};

And then we'd need to sprinkle the usage of Hooks into the impl of StorageMap and call it in appropriate places.

---

But I haven't got to work with either of decl or the pallet macros, so not sure if it is correct. And even so, this is aligned with the opinion that a counter-map should be just an extension (like what I named Hooks) to the normal StorageMap. Another idea is to have a make generator::StorageMap trait a super-trait, have a sub-trait called CountedStorageMap or something like this. In this case, the counter map would be a new storage type, so you probably will see something like this:

#[pallet::storage]
type Validators = CountedStorageMap<_, _, _>

decl_storage! {
     pub Validators get(fn validators): map_counted ...
}

[kianenigma]

What I like about this idea, implemented as a generic post-insert-hook and post-remove-hook is that we can probably easily extend it to paritytech/polkadot-sdk#328 as well.

[thiolliere]

I like the idea to introduce a Hooks to extend possibilities on insert and on remove, then we can introduce a type alias to be able to create some storage with counters with:

pub type StorageMapCounted<P, H, V, Q, E> = StorageMap<Prefix, Hasher, Key, Value, QueryKind=OptionQuery, OnEmpty=GetDefault, Hooks=Counter>;


#[pallet::storage]
type MyStorage<T: Config> = StorageMapCounted<_, Twox128Concat, u32, 32>;

Note that pallet macro should support this, or with minor change. Though in the metadata it will be seen as a storage map.
In order improve the metadata we could force Hooks to implement a trait HooksMetadata, and we would put this in the metadata.

[coriolinus]

hash(pallet_name) ++ hash(storage_name ++ counter_name)

Why not

hash(pallet_name) ++ hash(storage_name) ++ hash(counter_name)

? I'm not aware of any length limit for keys, and this would let us use common-prefix methods to manage the counter with the storage.

[edit] Oh, it looks like that's what you already implemented in the diff.

[thiolliere]

IMO any key with the prefix hash(pallet_name) ++ hash(storage_name) should be consider part of the storage map.
otherwise hash(counter) could write on a value (e.g. if storage hasher is identity). Also when we iterate we iterate on every key with the prefix and decode value. So any other key make the iteration difficult.

Better to store the counter in another place with the name storage_name ++ counter_name

[kianenigma]

Why not

hash(pallet_name) ++ hash(storage_name) ++ hash(counter_name)

I don't think there's any restriction on the length, I was just trying to stick to the size of the rest of the keys in the map. Both this, and

storage_name ++ counter_name

Work fine.

---

I am working on a draft here: https://github.com/paritytech/substrate/compare/kiz-bounded-staking?expand=1

It is not ready for a draft PR ready yet, but please provide feedback if you have any. So far I've tried to not change any of the base trait (StorageMap and generator::StorageMap) and keep the hooks as only an extension. But it might be helpful to actually make the hooks a first-class citizen of any storage map.

[apopiak]

Seems very sensible to me. I agree with Gui that the hooks seems more flexible and that we can always typedef a CountedMap

[kianenigma]

Current working branch is pretty generic over this and just provides post-hooks, one implementation of which is a Counter and the default is Nothing.

[KiChjang]

@kianenigma Is the Key type parameter for StorageMapHooks only going to be specified during implementation? If so, it would seem to me that a better way is to use an associated type instead of a type parameter to achieve the same goal.

I say this because you'll be able to implement StorageMapHooks directly on (). I know Nothing is more wordy and thus more self-documenting, but we already have a convention in the code base where () for Config traits means that the configuration should be left empty, or that a default implementation should suffice without any custom behaviours.

[kianenigma]

I'd be fine with switching the generic for an associated type if it fits better. Admittedly, this work in that branch is pretty much WIP. I really like the whole idea: don't just implement a counter, implement generic hooks, one function of which can be a counter. But the details can vary.

[thiolliere]

I am trying to see how this can be implemented, one difficulty is that this Hook would need to be triggered in lots of places and the storage traits in frame_support::storage is currently not very well designed for this change.
For example TryAppendMap trait is implemented automatically for any type implementing StorageMap, and in its implementation it directly calls: sp_io::storage::append(&key, item.encode());, if we add the hook we will need to modify this call and trait

Also we currently have 3 implementation of maps: map, double-map and n-map. Ideally map and double-map should directly use n-map implementation so that we don't need to duplicate this implementation 3 times.

(sidenote: to factorize the code I first thought about doing:

StorageMap<Prefix, K, V, ..> = StorageNMap<Prefix, Key<K, H>, ...>

but we can't do this if we want to keep the interface of storage map (like being able to call MyStorage::insert(k, v) instead of MyStorage::insert((k,), v) or MyDoubleMap::insert(k1, k2, v) instead of MyDoubleMap::insert((k1, k2), v)).)

So I think we can reorganize the code in frame_support::storage::types as such:

• do not use generators to implement map, double-map and n-map: because generators doesn't have knowledge of Hooks, if we want to continue to use generators, then we must put Hooks into them.
• do not use storage::StoragePrefixedMap for storage::types::StorageMap and storage::types::StorageDoubleMap: because storage::StoragePrefixedMap doesn't have knowledge of hooks, if we want to continue use it, we need to put Hooks into them.
• do not make implementation of storage::TryAppendMap for any implementation of storage::StorageMap: same storage::StorageTryAppend doesn't have knowledge of hooks.
• same for storage::TryAppendDoubleMap.
• copy method of storage::StoragePrefixedMap into storage::StorageMap, storage::StorageDoubleMap and storage::StorageNMap to avoid losing them.
• make storage::types::StorageMap implements: storage::StorageMap by calling storage::StorageNMap in its implementation
• make storage::types::StorageMap implements: storage::IterableStorageMap by calling storage::IterableStorageNMap in its implementation
• make storage::types::StorageMap implements: storage::TryAppendNMap by calling storage::TryAppendNMap in its implementation
• same for storage::types::StorageDoubleMap.
• make storage::types::StorageNMap implements storage::StorageNMap, storage::IterableStorageNMap and storage::TryAppendNMap directly taking care of calling the Hooks accordingly.

once we remove decl_storage, we can delete all the code in storage::generator, thus leaving no debt code.

[kianenigma]

Thanks for the detailed explanations @thiolliere

Before reading further, I want to add my current 2 cents that I think however we implement this, for now adding manual counters is worthwhile since it will be easy to migrate them later to a more macro-enhanced approach.

and I agree that the hook prototype was quite incomplete, and while expanding the idea I was also quite worried that if we miss triggering the hook somewhere, it would be quite a hard bug to track down.

Ideally, the hooks must be somewhere inside sp_io::xxx calls, in some other low level storage functions, and the macro code only calls this low level module or sp_io, and then it is less likely to miss a hook trigger.

[stale]

Hey, is anyone still working on this? Due to the inactivity this issue has been automatically marked as stale. It will be closed if no further activity occurs. Thank you for your contributions.

[thiolliere]

The more I think of it the less I'm convinced it is doable, for instance there is stuff implemented for the StorageMap trait like https://github.com/open-web3-stack/open-runtime-module-library/blob/443ee91bc2ca5f1fc155c0378eef6e89b67e2e97/utilities/src/iterator.rs#L81
Which directly calls into the storage to retrieve and remove keys/value. All those abstraction won't update the counters.

[kianenigma]

This is already done.