Add try_mutate_exists

frame/balances/src/mock.rs

@@ -20,7 +20,7 @@ use sp_runtime::{Perbill, traits::{ConvertInto, IdentityLookup}, testing::Header
 use sp_core::H256;
 use sp_io;
 use frame_support::{impl_outer_origin, parameter_types};
-use frame_support::traits::Get;
+use frame_support::traits::{Get, StorageMapShim};
 use frame_support::weights::{Weight, DispatchInfo};
 use std::cell::RefCell;
 use crate::{GenesisConfig, Module, Trait};
@@ -71,6 +71,9 @@ impl frame_system::Trait for Test {
 	type AvailableBlockRatio = AvailableBlockRatio;
 	type Version = ();
 	type ModuleToIndex = ();
+	type AccountData = ();
+	type OnNewAccount = ();
+	type OnReapAccount = Balances;
 }
 parameter_types! {
 	pub const TransactionBaseFee: u64 = 0;
@@ -86,13 +89,17 @@ impl pallet_transaction_payment::Trait for Test {
 }
 impl Trait for Test {
 	type Balance = u64;
-	type OnReapAccount = System;
-	type OnNewAccount = ();
-	type Event = ();
-	type DustRemoval = ();
 	type TransferPayment = ();
+	type DustRemoval = ();
+	type Event = ();
 	type ExistentialDeposit = ExistentialDeposit;
 	type CreationFee = CreationFee;
+	type AccountStore = StorageMapShim<
+		super::Account<Test>,
+		system::CallOnCreatedAccount<Test>,
+		system::CallKillAccount<Test>,
+		u64, super::AccountData<u64>
+	>;
 }
 
 pub struct ExtBuilder {

frame/support/src/storage/generator/map.rs

@@ -123,6 +123,35 @@ impl<K: FullEncode, V: FullCodec, G: StorageMap<K, V>> storage::StorageMap<K, V>
 		ret
 	}
 
+	fn mutate_exists<KeyArg: EncodeLike<K>, R, F: FnOnce(&mut Option<V>) -> R>(key: KeyArg, f: F) -> R {
+		let final_key = Self::storage_map_final_key(key);
+		let mut val = unhashed::get(final_key.as_ref());
+
+		let ret = f(&mut val);
+		match val {
+			Some(ref val) => unhashed::put(final_key.as_ref(), &val.borrow()),
+			None => unhashed::kill(final_key.as_ref()),
+		}
+		ret
+	}
+
+	fn try_mutate_exists<KeyArg: EncodeLike<K>, R, E, F: FnOnce(&mut Option<V>) -> Result<R, E>>(
+		key: KeyArg,
+		f: F
+	) -> Result<R, E> {
+		let final_key = Self::storage_map_final_key(key);
+		let mut val = unhashed::get(final_key.as_ref());
+
+		let ret = f(&mut val);
+		if ret.is_ok() {
+			match val {
+				Some(ref val) => unhashed::put(final_key.as_ref(), &val.borrow()),
+				None => unhashed::kill(final_key.as_ref()),
+			}
+		}
+		ret
+	}
+
 	fn take<KeyArg: EncodeLike<K>>(key: KeyArg) -> Self::Query {
 		let key = Self::storage_map_final_key(key);
 		let value = unhashed::take(key.as_ref());

frame/support/src/storage/migration.rs

@@ -79,6 +79,15 @@ pub fn get_storage_value<T: Decode + Sized>(module: &[u8], item: &[u8], hash: &[
 	frame_support::storage::unhashed::get::<T>(&key)
 }
 
+/// Get a particular value in storage by the `module`, the map's `item` name and the key `hash`.
+pub fn take_storage_value<T: Decode + Sized>(module: &[u8], item: &[u8], hash: &[u8]) -> Option<T> {
+	let mut key = vec![0u8; 32 + hash.len()];
+	key[0..16].copy_from_slice(&Twox128::hash(module));
+	key[16..32].copy_from_slice(&Twox128::hash(item));
+	key[32..].copy_from_slice(hash);
+	frame_support::storage::unhashed::take::<T>(&key)
+}
+
 /// Put a particular value into storage by the `module`, the map's `item` name and the key `hash`.
 pub fn put_storage_value<T: Encode>(module: &[u8], item: &[u8], hash: &[u8], value: T) {
 	let mut key = vec![0u8; 32 + hash.len()];

frame/support/src/storage/mod.rs

@@ -148,14 +148,22 @@ pub trait StorageMap<K: FullEncode, V: FullCodec> {
 	/// Mutate the value under a key.
 	fn mutate<KeyArg: EncodeLike<K>, R, F: FnOnce(&mut Self::Query) -> R>(key: KeyArg, f: F) -> R;
 
+	/// Mutate the value under a key.
+	fn mutate_exists<KeyArg: EncodeLike<K>, R, F: FnOnce(&mut Option<V>) -> R>(key: KeyArg, f: F) -> R;
+
+	/// Mutate the item, only if an `Ok` value is returned.
+	fn try_mutate_exists<KeyArg: EncodeLike<K>, R, E, F: FnOnce(&mut Option<V>) -> Result<R, E>>(
+		key: KeyArg,
+		f: F,
+	) -> Result<R, E>;
+
 	/// Take the value under a key.
 	fn take<KeyArg: EncodeLike<K>>(key: KeyArg) -> Self::Query;
 
 	/// Append the given items to the value in the storage.
 	///
 	/// `V` is required to implement `codec::EncodeAppend`.
-	fn append<Items, Item, EncodeLikeItem, KeyArg>(key: KeyArg, items: Items) -> Result<(), &'static str>
-	where
+	fn append<Items, Item, EncodeLikeItem, KeyArg>(key: KeyArg, items: Items) -> Result<(), &'static str> where
 		KeyArg: EncodeLike<K>,
 		Item: Encode,
 		EncodeLikeItem: EncodeLike<Item>,
@@ -167,8 +175,7 @@ pub trait StorageMap<K: FullEncode, V: FullCodec> {
 	/// old (presumably corrupt) value is replaced with the given `items`.
 	///
 	/// `V` is required to implement `codec::EncodeAppend`.
-	fn append_or_insert<Items, Item, EncodeLikeItem, KeyArg>(key: KeyArg, items: Items)
-	where
+	fn append_or_insert<Items, Item, EncodeLikeItem, KeyArg>(key: KeyArg, items: Items) where
 		KeyArg: EncodeLike<K>,
 		Item: Encode,
 		EncodeLikeItem: EncodeLike<Item>,

frame/support/src/traits.rs

@@ -27,28 +27,104 @@ use sp_runtime::{
 };
 
 use crate::dispatch::Parameter;
-use crate::storage::StorageValue;
+use crate::storage::StorageMap;
 
 /// An abstraction of a value stored within storage, but possibly as part of a larger composite
 /// item.
-pub trait StoredValue<T> {
-	/// Get the item.
-	fn get() -> T;
+pub trait StoredMap<K, T> {
+	/// Get the item, or its default if it doesn't yet exist; we make no distinction between the
+	/// two.
+	fn get(k: &K) -> T;
 	/// Mutate the item.
-	fn mutate<R>(f: impl FnOnce(&mut T) -> R) -> R;
+	fn mutate<R>(k: &K, f: impl FnOnce(&mut T) -> R) -> R;
+	/// Mutate the item, removing or resetting to default value if it has been mutated to `None`.
+	fn mutate_exists<R>(k: &K, f: impl FnOnce(&mut Option<T>) -> R) -> R;
+	/// Maybe mutate the item only if an `Ok` value is returned from `f`. Do nothing if an `Err` is
+	/// returned. It is removed or reset to default value if it has been mutated to `None`
+	fn try_mutate_exists<R, E>(k: &K, f: impl FnOnce(&mut Option<T>) -> Result<R, E>) -> Result<R, E>;
 	/// Set the item to something new.
-	fn put(t: T) { Self::mutate(|i| *i = t); }
+	fn insert(k: &K, t: T) { Self::mutate(k, |i| *i = t); }
+	/// Remove the item or otherwise replace it with its default value; we don't care which.
+	fn remove(k: &K);
+}
+
+/// A simple, generic one-parameter event notifier/handler.
+pub trait Happened<T> {
+	/// The thing happened.
+	fn happened(t: &T);
 }
 
 /// A shim for placing around a storage item in order to use it as a `StoredValue`. Ideally this
 /// wouldn't be needed as `StorageValue`s should blanket implement `StoredValue`s, however this
 /// would break the ability to have custom impls of `StoredValue`. The other workaround is to
-/// implement it directly in the macro, but that's left for later.
-pub struct StorageValueShim<S>(S);
-impl<S: StorageValue<T, Query=T>, T: FullCodec> StoredValue<T> for StorageValueShim<S> {
-	fn get() -> T { S::get() }
-	fn mutate<R>(f: impl FnOnce(&mut T) -> R) -> R { S::mutate(f) }
-	fn put(t: T) { S::put(t) }
+/// implement it directly in the macro.
+///
+/// This form has the advantage that two additional types are provides, `Created` and `Removed`,
+/// which are both generic events that can be tied to handlers to do something in the case of being
+/// about to create an account where one didn't previously exist (at all; not just where it used to
+/// be the default value), or where the account is being removed or reset back to the default value
+/// where previously it did exist (though may have been in a default state). This works well with
+/// system module's `CallOnCreatedAccount` and `CallKillAccount`.
+pub struct StorageMapShim<
+	S,
+	Created,
+	Removed,
+	K,
+	T
+>(sp_std::marker::PhantomData<(S, Created, Removed, K, T)>);
+impl<
+	S: StorageMap<K, T, Query=T>,
+	Created: Happened<K>,
+	Removed: Happened<K>,
+	K: FullCodec,
+	T: FullCodec
+> StoredMap<K, T> for StorageMapShim<S, Created, Removed, K, T> {
+	fn get(k: &K) -> T { S::get(k) }
+	fn insert(k: &K, t: T) {
+		S::insert(k, t);
+		if !S::exists(&k) {
+			Created::happened(k);
+		}
+	}
+	fn remove(k: &K) {
+		if S::exists(&k) {
+			Removed::happened(&k);
+		}
+		S::remove(k);
+	}
+	fn mutate<R>(k: &K, f: impl FnOnce(&mut T) -> R) -> R {
+		let r = S::mutate(k, f);
+		if !S::exists(&k) {
+			Created::happened(k);
+		}
+		r
+	}
+	fn mutate_exists<R>(k: &K, f: impl FnOnce(&mut Option<T>) -> R) -> R {
+		let (existed, exists, r) = S::mutate_exists(k, |maybe_value| {
+			let existed = maybe_value.is_some();
+			let r = f(maybe_value);
+			(existed, maybe_value.is_some(), r)
+		});
+		if !existed && exists {
+			Created::happened(k);
+		} else if existed && !exists {
+			Removed::happened(k);
+		}
+		r
+	}
+	fn try_mutate_exists<R, E>(k: &K, f: impl FnOnce(&mut Option<T>) -> Result<R, E>) -> Result<R, E> {
+		S::try_mutate_exists(k, |maybe_value| {
+			let existed = maybe_value.is_some();
+			f(maybe_value).map(|v| (existed, maybe_value.is_some(), v))
+		}).map(|(existed, exists, v)| {
+			if !existed && exists {
+				Created::happened(k);
+			} else if existed && !exists {
+				Removed::happened(k);
+			}
+			v
+		})
+	}
 }
 
 /// Anything that can have a `::len()` method.