EnumerableStorageMap (#1763)

* Refactor decl_storage a bit to allow easier impl of linked map. * A bunch of refactorings for storage generation. - Rename StorageMap and ChildrenStorageMap to avoid confusion with generator::StorageMap. - Separate implementation from the procedural macro code to clean it up. - Make sure that genesis is initialised using the `StorageValue/StorageMap` generated implementations instead of going RAW. * WiP: Writing test. * Basic implementation. * Implement enumeration. * Fix non-std issues. * fix warning * Fix test-client. * Address review grumbles - part 1 * Avoid cloning the key, relax Storage requirements. * Rebuild runtime. * Remove dangling todo.
2026-06-12 21:41:12 +00:00 · 2019-02-13 08:52:52 +01:00
parent 6e26c52191
commit 9e2710246f
20 changed files with 803 additions and 243 deletions
@@ -0,0 +1,428 @@
+// Copyright 2019 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+use proc_macro2::TokenStream as TokenStream2;
+use syn;
+use quote::quote;
+use crate::storage::transformation::DeclStorageTypeInfos;
+
+pub fn option_unwrap(is_option: bool) -> TokenStream2 {
+	if !is_option {
+		// raw type case
+		quote!( unwrap_or_else )
+	} else {
+		// Option<> type case
+		quote!( or_else )
+	}
+}
+
+pub(crate) struct Impls<'a> {
+	pub scrate: &'a TokenStream2,
+	pub visibility: &'a syn::Visibility,
+	pub traitinstance: &'a syn::Ident,
+	pub traittype: &'a syn::TypeParamBound,
+	pub type_infos: DeclStorageTypeInfos<'a>,
+	pub fielddefault: TokenStream2,
+	pub prefix: String,
+	pub name: &'a syn::Ident,
+}
+
+impl<'a> Impls<'a> {
+	pub fn simple_value(self) -> TokenStream2 {
+		let Self {
+			scrate,
+			visibility,
+			traitinstance,
+			traittype,
+			type_infos,
+			fielddefault,
+			prefix,
+			name,
+		} = self;
+		let DeclStorageTypeInfos { typ, value_type, is_option, .. } = type_infos;
+		let option_simple_1 = option_unwrap(is_option);
+
+		let mutate_impl = if !is_option {
+			quote!{
+				<Self as #scrate::storage::generator::StorageValue<#typ>>::put(&val, storage)
+			}
+		} else {
+			quote!{
+				match val {
+					Some(ref val) => <Self as #scrate::storage::generator::StorageValue<#typ>>::put(&val, storage),
+					None => <Self as #scrate::storage::generator::StorageValue<#typ>>::kill(storage),
+				}
+			}
+		};
+
+		// generator for value
+		quote!{
+
+			#visibility struct #name<#traitinstance: #traittype>(#scrate::storage::generator::PhantomData<#traitinstance>);
+
+			impl<#traitinstance: #traittype> #scrate::storage::generator::StorageValue<#typ> for #name<#traitinstance> {
+				type Query = #value_type;
+
+				/// Get the storage key.
+				fn key() -> &'static [u8] {
+					#prefix.as_bytes()
+				}
+
+				/// Load the value from the provided storage instance.
+				fn get<S: #scrate::GenericStorage>(storage: &S) -> Self::Query {
+					storage.get(<Self as #scrate::storage::generator::StorageValue<#typ>>::key())
+						.#option_simple_1(|| #fielddefault)
+				}
+
+				/// Take a value from storage, removing it afterwards.
+				fn take<S: #scrate::GenericStorage>(storage: &S) -> Self::Query {
+					storage.take(<Self as #scrate::storage::generator::StorageValue<#typ>>::key())
+						.#option_simple_1(|| #fielddefault)
+				}
+
+				/// Mutate the value under a key.
+				fn mutate<R, F: FnOnce(&mut Self::Query) -> R, S: #scrate::GenericStorage>(f: F, storage: &S) -> R {
+					let mut val = <Self as #scrate::storage::generator::StorageValue<#typ>>::get(storage);
+
+					let ret = f(&mut val);
+					#mutate_impl ;
+					ret
+				}
+			}
+
+		}
+	}
+
+	pub fn map(self, kty: &syn::Type) -> TokenStream2 {
+		let Self {
+			scrate,
+			visibility,
+			traitinstance,
+			traittype,
+			type_infos,
+			fielddefault,
+			prefix,
+			name,
+		} = self;
+		let DeclStorageTypeInfos { typ, value_type, is_option, .. } = type_infos;
+		let option_simple_1 = option_unwrap(is_option);
+
+		let mutate_impl = if !is_option {
+			quote!{
+				<Self as #scrate::storage::generator::StorageMap<#kty, #typ>>::insert(key, &val, storage)
+			}
+		} else {
+			quote!{
+				match val {
+					Some(ref val) => <Self as #scrate::storage::generator::StorageMap<#kty, #typ>>::insert(key, &val, storage),
+					None => <Self as #scrate::storage::generator::StorageMap<#kty, #typ>>::remove(key, storage),
+				}
+			}
+		};
+		// generator for map
+		quote!{
+			#visibility struct #name<#traitinstance: #traittype>(#scrate::storage::generator::PhantomData<#traitinstance>);
+
+			impl<#traitinstance: #traittype> #scrate::storage::generator::StorageMap<#kty, #typ> for #name<#traitinstance> {
+				type Query = #value_type;
+
+				/// Get the prefix key in storage.
+				fn prefix() -> &'static [u8] {
+					#prefix.as_bytes()
+				}
+
+				/// Get the storage key used to fetch a value corresponding to a specific key.
+				fn key_for(x: &#kty) -> #scrate::rstd::vec::Vec<u8> {
+					let mut key = #prefix.as_bytes().to_vec();
+					#scrate::codec::Encode::encode_to(x, &mut key);
+					key
+				}
+
+				/// Load the value associated with the given key from the map.
+				fn get<S: #scrate::GenericStorage>(key: &#kty, storage: &S) -> Self::Query {
+					let key = <Self as #scrate::storage::generator::StorageMap<#kty, #typ>>::key_for(key);
+					storage.get(&key[..]).#option_simple_1(|| #fielddefault)
+				}
+
+				/// Take the value, reading and removing it.
+				fn take<S: #scrate::GenericStorage>(key: &#kty, storage: &S) -> Self::Query {
+					let key = <Self as #scrate::storage::generator::StorageMap<#kty, #typ>>::key_for(key);
+					storage.take(&key[..]).#option_simple_1(|| #fielddefault)
+				}
+
+				/// Mutate the value under a key
+				fn mutate<R, F: FnOnce(&mut Self::Query) -> R, S: #scrate::GenericStorage>(key: &#kty, f: F, storage: &S) -> R {
+					let mut val = <Self as #scrate::storage::generator::StorageMap<#kty, #typ>>::get(key, storage);
+
+					let ret = f(&mut val);
+					#mutate_impl ;
+					ret
+				}
+
+			}
+		}
+	}
+
+	pub fn linked_map(self, kty: &syn::Type) -> TokenStream2 {
+		let Self {
+			scrate,
+			visibility,
+			traitinstance,
+			traittype,
+			type_infos,
+			fielddefault,
+			prefix,
+			name,
+		} = self;
+		let DeclStorageTypeInfos { typ, value_type, is_option, .. } = type_infos;
+		let option_simple_1 = option_unwrap(is_option);
+		// make sure to use different prefix for head and elements.
+		let head_key = format!("head of {}", prefix);
+		let prefix = format!("{}", prefix);
+		let name_lowercase = name.to_string().to_lowercase();
+		let key_for = syn::Ident::new(&format!("key_for_{}", name_lowercase), name.span());
+		let internal_module = syn::Ident::new(&format!("__internal_do_not_use_{}", name_lowercase), name.span());
+		let linkage = syn::Ident::new(&format!("Linkage{}", name), name.span());
+		let borrowing_linkage = syn::Ident::new(&format!("Borrowing{}", linkage), name.span());
+		let enumerator = syn::Ident::new(&format!("Enumerator{}", name), name.span());
+		let put_or_insert = quote! {
+			match linkage {
+				Some(linkage) => storage.put(key_for, &(val, linkage)),
+				None => <Self as #scrate::storage::generator::StorageMap<#kty, #typ>>::insert(key, &val, storage),
+			}
+		};
+		let mutate_impl = if !type_infos.is_option {
+			put_or_insert
+		} else {
+			quote! {
+				match val {
+					Some(ref val) => #put_or_insert,
+					None => <Self as #scrate::storage::generator::StorageMap<#kty, #typ>>::remove(key, storage),
+				}
+			}
+		};
+
+		// generator for linked map
+		quote! {
+
+			mod #internal_module {
+				use super::*;
+
+				#[derive(Default, parity_codec_derive::Encode, parity_codec_derive::Decode)]
+				pub struct #linkage {
+					/// Previous element key in storage (None for the first element)
+					previous: Option<#kty>,
+					/// Next element key in storage (None for the last element)
+					next: Option<#kty>,
+				}
+
+				/// A helper struct to avoid unnecessary key cloning.
+				///
+				/// NOTE It has to have exact same parity-codec encoding as #linkage!
+				#[derive(parity_codec_derive::Encode)]
+				struct #borrowing_linkage<'a> {
+					previous: Option<&'a #kty>,
+					next: Option<&'a #kty>,
+				}
+
+				impl #linkage {
+					/// Update linkage when this element is removed.
+					///
+					/// Takes care of updating previous and next elements points
+					/// as well as updates head if the element is first or last.
+					pub fn remove<S: #scrate::GenericStorage>(
+						self,
+						storage: &S,
+					) {
+						let next_key = self.next.as_ref().map(|x| #key_for(x));
+						let prev_key = self.previous.as_ref().map(|x| #key_for(x));
+
+						if let Some(prev_key) = prev_key {
+							// Retrieve previous element and update `next`
+							let mut res = Self::read(storage, &*prev_key)
+								.expect("Linkage is updated in case entry is removed; it always points to existing keys; qed");
+							res.1.next = self.next;
+							storage.put(&*prev_key, &res);
+						} else {
+							// we were first so let's update the head
+							Self::write_head(storage, self.next.as_ref());
+						}
+
+						if let Some(next_key) = next_key {
+							// Update previous of next element
+							let mut res = Self::read(storage, &*next_key)
+								.expect("Linkage is updated in case entry is removed; it always points to existing keys; qed");
+							res.1.previous = self.previous;
+							storage.put(&*next_key, &res);
+						}
+					}
+
+					/// Read the contained data and it's linkage.
+					pub fn read<S: #scrate::GenericStorage>(storage: &S, key: &[u8]) -> Option<(#value_type, #linkage)> {
+						storage.get(key)
+					}
+
+					/// Generate linkage for newly inserted element.
+					///
+					/// Takes care of updating head and previous head's pointer.
+					pub fn insert_new_head<S: #scrate::GenericStorage>(
+						storage: &S,
+						key: &#kty,
+					) -> Self {
+						if let Some(head) = Self::read_head(storage) {
+							// update previous head predecessor
+							{
+								let head_key = #key_for(&head);
+								let (data, linkage) = Self::read(storage, &*head_key).expect(r#"
+									head is set when first element is inserted and unset when last element is removed;
+									if head is Some then it points to existing key; qed
+								"#);
+								storage.put(&*head_key, &(data, #borrowing_linkage {
+									next: linkage.next.as_ref(),
+									previous: Some(key),
+								}));
+							}
+							// update to current head
+							Self::write_head(storage, Some(key));
+							// return linkage with pointer to previous head
+							let mut linkage = Self::default();
+							linkage.next = Some(head);
+							linkage
+						} else {
+							// we are first - update the head and produce empty linkage
+							Self::write_head(storage, Some(key));
+							Self::default()
+						}
+					}
+
+					/// Read current head pointer.
+					pub fn read_head<S: #scrate::GenericStorage>(storage: &S) -> Option<#kty> {
+						storage.get(#head_key.as_bytes())
+					}
+
+					/// Overwrite current head pointer.
+					///
+					/// If `None` is given head is removed from storage.
+					fn write_head<S: #scrate::GenericStorage>(storage: &S, head: Option<&#kty>) {
+						match head {
+							Some(head) => storage.put(#head_key.as_bytes(), head),
+							None => storage.kill(#head_key.as_bytes()),
+						}
+					}
+				}
+
+				pub struct #enumerator<'a, S> {
+					pub storage: &'a S,
+					pub next: Option<#kty>,
+				}
+
+				impl<'a, S: #scrate::GenericStorage> Iterator for #enumerator<'a, S> {
+					type Item = (#kty, #typ);
+
+					fn next(&mut self) -> Option<Self::Item> {
+						let next = self.next.take()?;
+						let key_for = #key_for(&next);
+						let (val, linkage) = #linkage::read(self.storage, &*key_for)
+							.expect("previous/next only contain existing entires; we enumerate using next; entry exists; qed");
+						self.next = linkage.next;
+						Some((next, val))
+					}
+				}
+			}
+
+			fn #key_for(key: &#kty) -> #scrate::rstd::vec::Vec<u8> {
+				let mut key_for = #prefix.as_bytes().to_vec();
+				#scrate::codec::Encode::encode_to(&key, &mut key_for);
+				key_for
+			}
+
+			#visibility struct #name<#traitinstance: #traittype>(#scrate::storage::generator::PhantomData<#traitinstance>);
+
+			impl<#traitinstance: #traittype> #scrate::storage::generator::StorageMap<#kty, #typ> for #name<#traitinstance> {
+				type Query = #value_type;
+
+				/// Get the prefix key in storage.
+				fn prefix() -> &'static [u8] {
+					#prefix.as_bytes()
+				}
+
+				/// Get the storage key used to fetch a value corresponding to a specific key.
+				fn key_for(x: &#kty) -> #scrate::rstd::vec::Vec<u8> {
+					#key_for(x)
+				}
+
+				/// Load the value associated with the given key from the map.
+				fn get<S: #scrate::GenericStorage>(key: &#kty, storage: &S) -> Self::Query {
+					storage.get(&*#key_for(key)).#option_simple_1(|| #fielddefault)
+				}
+
+				/// Take the value, reading and removing it.
+				fn take<S: #scrate::GenericStorage>(key: &#kty, storage: &S) -> Self::Query {
+					let res: Option<(#value_type, self::#internal_module::#linkage)> = storage.take(&*#key_for(key));
+					match res {
+						Some((data, linkage)) => {
+							linkage.remove(storage);
+							data
+						},
+						None => #fielddefault
+					}
+				}
+
+				/// Remove the value under a key.
+				fn remove<S: #scrate::GenericStorage>(key: &#kty, storage: &S) {
+					<Self as #scrate::storage::generator::StorageMap<#kty, #typ>>::take(key, storage);
+				}
+
+				/// Store a value to be associated with the given key from the map.
+				fn insert<S: #scrate::GenericStorage>(key: &#kty, val: &#typ, storage: &S) {
+					let key_for = &*#key_for(key);
+					let linkage = match self::#internal_module::#linkage::read(storage, key_for) {
+						// overwrite but reuse existing linkage
+						Some((_data, linkage)) => linkage,
+						// create new linkage
+						None => self::#internal_module::#linkage::insert_new_head(storage, key),
+					};
+					storage.put(key_for, &(*val, linkage))
+				}
+
+				/// Mutate the value under a key
+				fn mutate<R, F: FnOnce(&mut Self::Query) -> R, S: #scrate::GenericStorage>(key: &#kty, f: F, storage: &S) -> R {
+					let key_for = &*#key_for(key);
+					let (mut val, linkage) = self::#internal_module::#linkage::read(storage, key_for)
+						.map(|(data, linkage)| (data, Some(linkage)))
+						.#option_simple_1(|| (#fielddefault, None));
+
+					let ret = f(&mut val);
+					#mutate_impl ;
+					ret
+				}
+			}
+
+			impl<#traitinstance: #traittype> #scrate::storage::generator::EnumerableStorageMap<#kty, #typ> for #name<#traitinstance> {
+				fn head<S: #scrate::GenericStorage>(storage: &S) -> Option<#kty> {
+					self::#internal_module::#linkage::read_head(storage)
+				}
+
+				fn enumerate<'a, S: #scrate::GenericStorage>(storage: &'a S) -> #scrate::storage::generator::Box<dyn Iterator<Item = (#kty, #typ)> + 'a> {
+					#scrate::storage::generator::Box::new(self::#internal_module::#enumerator {
+						next: self::#internal_module::#linkage::read_head(storage),
+						storage,
+					})
+				}
+			}
+		}
+	}
+}