From c34b538cd2bc45da4544e887180184e30957904a Mon Sep 17 00:00:00 2001
From: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Date: Thu, 22 Dec 2022 01:14:45 +0100
Subject: [PATCH] WIP

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
---
 frame/support/src/storage/mod.rs         |  6 ++++++
 frame/support/src/storage/types/value.rs | 27 +++++++++++++++++-------
 frame/support/src/traits/storage.rs      |  8 +------
 3 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs
index 8c0d6207c3f4d..28d41dfd31f76 100644
--- a/frame/support/src/storage/mod.rs
+++ b/frame/support/src/storage/mod.rs
@@ -1290,6 +1290,12 @@ pub trait StorageDecodeLength: private::Sealed + codec::DecodeLength {
 	}
 }
 
+#[derive(codec::Encode, codec::Decode, crate::RuntimeDebug, Eq, PartialEq, Clone)]
+pub enum ProofSizeMode {
+	MaxEncodedLen,
+	Measured,
+}
+
 /// Provides `Sealed` trait to prevent implementing trait `StorageAppend` & `StorageDecodeLength`
 /// & `EncodeLikeTuple` outside of this crate.
 mod private {
diff --git a/frame/support/src/storage/types/value.rs b/frame/support/src/storage/types/value.rs
index dc3cc51b87ff1..e11a1b2281072 100644
--- a/frame/support/src/storage/types/value.rs
+++ b/frame/support/src/storage/types/value.rs
@@ -30,18 +30,28 @@ use codec::{Decode, Encode, EncodeLike, FullCodec, MaxEncodedLen};
 use sp_arithmetic::traits::SaturatedConversion;
 use sp_std::prelude::*;
 
+use crate::storage::ProofSizeMode;
+
 /// A type that allow to store a value.
 ///
 /// Each value is stored at:
 /// ```nocompile
 /// Twox128(Prefix::pallet_prefix()) ++ Twox128(Prefix::STORAGE_PREFIX)
 /// ```
-pub struct StorageValue<Prefix, Value, QueryKind = OptionQuery, OnEmpty = GetDefault>(
-	core::marker::PhantomData<(Prefix, Value, QueryKind, OnEmpty)>,
+pub struct StorageValue<Prefix, Value, QueryKind = OptionQuery, OnEmpty = GetDefault, ProofSize = MeasureProofSize>(
+	core::marker::PhantomData<(Prefix, Value, QueryKind, OnEmpty, ProofSize)>,
 );
 
-impl<Prefix, Value, QueryKind, OnEmpty> crate::storage::generator::StorageValue<Value>
-	for StorageValue<Prefix, Value, QueryKind, OnEmpty>
+pub struct MeasureProofSize;
+
+impl crate::pallet_prelude::Get<Option<ProofSizeMode>> for MeasureProofSize {
+	fn get() -> Option<ProofSizeMode> {
+		None
+	}
+}
+
+impl<Prefix, Value, QueryKind, OnEmpty, ProofSize> crate::storage::generator::StorageValue<Value>
+	for StorageValue<Prefix, Value, QueryKind, OnEmpty, ProofSize>
 where
 	Prefix: StorageInstance,
 	Value: FullCodec,
@@ -63,7 +73,7 @@ where
 	}
 }
 
-impl<Prefix, Value, QueryKind, OnEmpty> StorageValue<Prefix, Value, QueryKind, OnEmpty>
+impl<Prefix, Value, QueryKind, OnEmpty, ProofSize> StorageValue<Prefix, Value, QueryKind, OnEmpty, ProofSize>
 where
 	Prefix: StorageInstance,
 	Value: FullCodec,
@@ -224,13 +234,14 @@ where
 	}
 }
 
-impl<Prefix, Value, QueryKind, OnEmpty> crate::traits::StorageInfoTrait
-	for StorageValue<Prefix, Value, QueryKind, OnEmpty>
+impl<Prefix, Value, QueryKind, OnEmpty, ProofSize> crate::traits::StorageInfoTrait
+	for StorageValue<Prefix, Value, QueryKind, OnEmpty, ProofSize>
 where
 	Prefix: StorageInstance,
 	Value: FullCodec + MaxEncodedLen,
 	QueryKind: QueryKindTrait<Value, OnEmpty>,
 	OnEmpty: crate::traits::Get<QueryKind::Query> + 'static,
+	ProofSize: crate::traits::Get<Option<crate::storage::ProofSizeMode>>,
 {
 	fn storage_info() -> Vec<StorageInfo> {
 		vec![StorageInfo {
@@ -239,7 +250,7 @@ where
 			prefix: Self::hashed_key().to_vec(),
 			max_values: Some(1),
 			max_size: Some(Value::max_encoded_len().saturated_into()),
-			proof_size: None,
+			proof_size: ProofSize::get(),
 		}]
 	}
 }
diff --git a/frame/support/src/traits/storage.rs b/frame/support/src/traits/storage.rs
index cb1f81d687a57..f31cd7956c590 100644
--- a/frame/support/src/traits/storage.rs
+++ b/frame/support/src/traits/storage.rs
@@ -65,13 +65,7 @@ pub struct StorageInfo {
 	pub max_values: Option<u32>,
 	/// The maximum size of key/values in the storage, or none if no maximum specified.
 	pub max_size: Option<u32>,
-	pub proof_size: Option<ProofSizeMode>,
-}
-
-#[derive(codec::Encode, codec::Decode, crate::RuntimeDebug, Eq, PartialEq, Clone)]
-pub enum ProofSizeMode {
-	MaxEncodedLen,
-	Measured,
+	pub proof_size: Option<crate::storage::ProofSizeMode>,
 }
 
 /// A trait to give information about storage.