rename to store_by_key and make it a store of a one length array

Co-authored-by: Richard Pringle <richard.pringle@avalabs.org> Signed-off-by: Franfran <51274081+iFrostizz@users.noreply.github.com>
ava-labs · Jul 5, 2024 · 6ab3bc0 · 6ab3bc0
1 parent 7dfad96
commit 6ab3bc0
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 35 deletions.
diff --git a/x/programs/rust/examples/automated-market-maker/src/lib.rs b/x/programs/rust/examples/automated-market-maker/src/lib.rs
@@ -31,15 +31,11 @@ pub fn add_liquidity(context: Context<StateKeys>, amount_x: u64, amount_y: u64)
 
     program
         .state()
-        .store(StateKeys::ReserveX, &(reserve_x + amount_x))
-        .unwrap();
-    program
-        .state()
-        .store(StateKeys::ReserveY, &(reserve_y + amount_y))
-        .unwrap();
-    program
-        .state()
-        .store(StateKeys::TotalySupply, &(total_supply + minted))
+        .store([
+            (StateKeys::ReserveX, &(reserve_x + amount_x)),
+            (StateKeys::ReserveY, &(reserve_y + amount_y)),
+            (StateKeys::TotalySupply, &(total_supply + minted)),
+        ])
         .unwrap();
 
     minted
@@ -57,15 +53,11 @@ pub fn remove_liquidity(context: Context<StateKeys>, shares: u64) -> (u64, u64)
 
     program
         .state()
-        .store(StateKeys::ReserveX, &(reserve_x - amount_x))
-        .unwrap();
-    program
-        .state()
-        .store(StateKeys::ReserveY, &(reserve_y - amount_y))
-        .unwrap();
-    program
-        .state()
-        .store(StateKeys::TotalySupply, &(total_supply - shares))
+        .store([
+            (StateKeys::ReserveX, &(reserve_x - amount_x)),
+            (StateKeys::ReserveY, &(reserve_y - amount_y)),
+            (StateKeys::TotalySupply, &(total_supply - shares)),
+        ])
         .unwrap();
 
     (amount_x, amount_y)
@@ -95,11 +87,10 @@ pub fn swap(context: Context<StateKeys>, amount_in: u64, x_to_y: bool) -> u64 {
 
     program
         .state()
-        .store(StateKeys::ReserveX, &reserve_x)
-        .unwrap();
-    program
-        .state()
-        .store(StateKeys::ReserveY, &reserve_y)
+        .store([
+            (StateKeys::ReserveX, &reserve_x),
+            (StateKeys::ReserveY, &reserve_y),
+        ])
         .unwrap();
 
     out

diff --git a/x/programs/rust/examples/counter/src/lib.rs b/x/programs/rust/examples/counter/src/lib.rs
@@ -22,7 +22,7 @@ pub fn inc(context: Context<StateKeys>, to: Address, amount: Count) -> bool {
     context
         .program()
         .state()
-        .store_single(StateKeys::Counter(to), &counter)
+        .store_by_key(StateKeys::Counter(to), &counter)
         .expect("failed to store counter");
 
     true

diff --git a/x/programs/rust/examples/token/src/lib.rs b/x/programs/rust/examples/token/src/lib.rs
@@ -26,19 +26,19 @@ pub fn init(context: Context<StateKeys>) {
     // set total supply
     program
         .state()
-        .store_single(StateKeys::TotalSupply, &INITIAL_SUPPLY)
+        .store_by_key(StateKeys::TotalSupply, &INITIAL_SUPPLY)
         .expect("failed to store total supply");
 
     // set token name
     program
         .state()
-        .store_single(StateKeys::Name, b"WasmCoin")
+        .store_by_key(StateKeys::Name, b"WasmCoin")
         .expect("failed to store coin name");
 
     // set token symbol
     program
         .state()
-        .store_single(StateKeys::Symbol, b"WACK")
+        .store_by_key(StateKeys::Symbol, b"WACK")
         .expect("failed to store symbol");
 }
 
@@ -76,7 +76,7 @@ fn mint_to_internal(
 
     program
         .state()
-        .store_single(StateKeys::Balance(recipient), &(balance + amount))
+        .store_by_key(StateKeys::Balance(recipient), &(balance + amount))
         .expect("failed to store balance");
 
     context

diff --git a/x/programs/rust/wasmlanche-sdk/src/state.rs b/x/programs/rust/wasmlanche-sdk/src/state.rs
@@ -51,9 +51,16 @@ impl<'a, K: Key> State<'a, K> {
         pairs: Pairs,
     ) -> Result<(), Error> {
         let cache = &mut self.cache.borrow_mut();
-        for (key, value) in pairs {
-            Self::store_internal(cache, key, value)?;
-        }
+
+        pairs
+            .into_iter()
+            .map(|(k, v)| borsh::to_vec(&v).map(|bytes| (k, Some(bytes))))
+            .try_for_each(|result| {
+                result.map(|(k, v)| {
+                    cache.insert(k, v);
+                })
+            })
+            .map_err(|_| StateError::Serialization)?;
 
         Ok(())
     }
@@ -63,9 +70,8 @@ impl<'a, K: Key> State<'a, K> {
     /// # Errors
     /// Returns an [`Error`] if the key or value cannot be
     /// serialized or if the host fails to handle the operation.
-    pub fn store_single<V: BorshSerialize>(self, key: K, value: &V) -> Result<(), Error> {
-        let cache = &mut self.cache.borrow_mut();
-        Self::store_internal(cache, key, value)
+    pub fn store_by_key<V: BorshSerialize>(self, key: K, value: &V) -> Result<(), Error> {
+        self.store([(key, value)])
     }
 
     fn store_internal<V: BorshSerialize>(

diff --git a/x/programs/test/programs/state_access/src/lib.rs b/x/programs/test/programs/state_access/src/lib.rs
@@ -12,7 +12,7 @@ pub fn put(context: Context<StateKeys>, value: i64) {
     context
         .program()
         .state()
-        .store_single(StateKeys::State, &value)
+        .store_by_key(StateKeys::State, &value)
         .expect("failed to store state");
 }