Refactor all storage access to be parametric in the new IO trait

engine-precompiles/src/native.rs

@@ -193,8 +193,11 @@ impl ExitToNear {
 
 #[cfg(feature = "contract")]
 fn get_nep141_from_erc20(erc20_token: &[u8]) -> AccountId {
+    use sdk::io::{StorageIntermediate, IO};
     AccountId::from_utf8(
-        sdk::read_storage(bytes_to_key(KeyPrefix::Erc20Nep141Map, erc20_token).as_slice())
+        sdk::near_runtime::Runtime
+            .read_storage(bytes_to_key(KeyPrefix::Erc20Nep141Map, erc20_token).as_slice())
+            .map(|s| s.to_vec())
             .expect(ERR_TARGET_TOKEN_NOT_FOUND),
     )
     .unwrap()

engine-sdk/src/error.rs

@@ -1,3 +1,4 @@
+#[derive(Debug)]
 pub struct BorshDeserializeError;
 
 impl AsRef<[u8]> for BorshDeserializeError {
@@ -6,6 +7,7 @@ impl AsRef<[u8]> for BorshDeserializeError {
     }
 }
 
+#[derive(Debug)]
 pub struct IncorrectInputLength;
 
 impl AsRef<[u8]> for IncorrectInputLength {
@@ -14,6 +16,7 @@ impl AsRef<[u8]> for IncorrectInputLength {
     }
 }
 
+#[derive(Debug)]
 pub enum ReadU64Error {
     InvalidU64,
     MissingValue,
@@ -27,3 +30,18 @@ impl AsRef<[u8]> for ReadU64Error {
         }
     }
 }
+
+#[derive(Debug)]
+pub enum ReadU256Error {
+    InvalidU256,
+    MissingValue,
+}
+
+impl AsRef<[u8]> for ReadU256Error {
+    fn as_ref(&self) -> &[u8] {
+        match self {
+            Self::InvalidU256 => b"ERR_NOT_U256",
+            Self::MissingValue => b"ERR_U256_NOT_FOUND",
+        }
+    }
+}

engine-sdk/src/io.rs

@@ -1,6 +1,7 @@
 use crate::error;
 use crate::prelude::{vec, Vec};
-use borsh::BorshDeserialize;
+use aurora_engine_types::U256;
+use borsh::{BorshDeserialize, BorshSerialize};
 
 /// The purpose of this trait is to represent a reference to a value that
 /// could be obtained by IO, but without eagerly loading it into memory.
@@ -109,4 +110,29 @@ pub trait IO {
         value.copy_to_slice(&mut result);
         Ok(u64::from_le_bytes(result))
     }
+
+    /// Convenience function to read a 256-bit unsigned integer from storage
+    /// (assumes big-endian encoding).
+    fn read_u256(&self, key: &[u8]) -> Result<U256, error::ReadU256Error> {
+        let value = self
+            .read_storage(key)
+            .ok_or(error::ReadU256Error::MissingValue)?;
+
+        if value.len() != 32 {
+            return Err(error::ReadU256Error::InvalidU256);
+        }
+
+        let mut result = [0u8; 32];
+        value.copy_to_slice(&mut result);
+        Ok(U256::from_big_endian(&result))
+    }
+
+    fn write_borsh<T: BorshSerialize>(
+        &mut self,
+        key: &[u8],
+        value: &T,
+    ) -> Option<Self::StorageValue> {
+        let bytes = value.try_to_vec().ok()?;
+        self.write_storage(key, &bytes)
+    }
 }

engine-sdk/src/lib.rs

@@ -3,10 +3,7 @@
 #![cfg_attr(not(feature = "std"), feature(alloc_error_handler))]
 #![cfg_attr(feature = "log", feature(panic_info_message))]
 
-use crate::prelude::{
-    vec, Address, BorshDeserialize, BorshSerialize, PromiseResult, Vec, H256,
-    STORAGE_PRICE_PER_BYTE,
-};
+use crate::prelude::{vec, Address, PromiseResult, Vec, H256, STORAGE_PRICE_PER_BYTE};
 pub use types::keccak;
 
 pub mod error;
@@ -15,7 +12,6 @@ pub mod near_runtime;
 mod prelude;
 pub mod types;
 
-use io::{StorageIntermediate, IO};
 use near_runtime::exports;
 
 const ECRECOVER_MESSAGE_SIZE: u64 = 32;
@@ -24,70 +20,6 @@ const ECRECOVER_MALLEABILITY_FLAG: u64 = 1;
 
 const GAS_FOR_STATE_MIGRATION: u64 = 100_000_000_000_000;
 
-pub fn read_input() -> Vec<u8> {
-    near_runtime::Runtime.read_input().to_vec()
-}
-
-#[cfg_attr(not(feature = "contract"), allow(dead_code))]
-pub fn read_input_borsh<T: BorshDeserialize>() -> Result<T, error::BorshDeserializeError> {
-    near_runtime::Runtime.read_input_borsh()
-}
-
-#[cfg_attr(not(feature = "contract"), allow(dead_code))]
-pub fn read_input_arr20() -> Result<[u8; 20], error::IncorrectInputLength> {
-    near_runtime::Runtime.read_input_arr20()
-}
-
-/// Reads current input and stores in the given key keeping data in the runtime.
-pub fn read_input_and_store(key: &[u8]) {
-    near_runtime::Runtime.read_input_and_store(key);
-}
-
-pub fn return_output(value: &[u8]) {
-    near_runtime::Runtime.return_output(value);
-}
-
-#[allow(dead_code)]
-pub fn read_storage(key: &[u8]) -> Option<Vec<u8>> {
-    near_runtime::Runtime.read_storage(key).map(|s| s.to_vec())
-}
-
-pub fn read_storage_len(key: &[u8]) -> Option<usize> {
-    near_runtime::Runtime.read_storage_len(key)
-}
-
-/// Read u64 from storage at given key.
-pub fn read_u64(key: &[u8]) -> Result<u64, error::ReadU64Error> {
-    near_runtime::Runtime.read_u64(key)
-}
-
-pub fn write_storage(key: &[u8], value: &[u8]) {
-    near_runtime::Runtime.write_storage(key, value);
-}
-
-pub fn remove_storage(key: &[u8]) {
-    near_runtime::Runtime.remove_storage(key);
-}
-
-/// Returns the size of the blob stored in the given register.
-/// * If register is used, then returns the size, which can potentially be zero;
-/// * If register is not used, returns `u64::MAX`
-pub fn register_len(register_id: u64) -> Option<u64> {
-    let len = unsafe { exports::register_len(register_id) };
-
-    if len == u64::MAX {
-        None
-    } else {
-        Some(len)
-    }
-}
-
-pub fn remove_storage_with_result(key: &[u8]) -> Option<Vec<u8>> {
-    near_runtime::Runtime
-        .remove_storage(key)
-        .map(|s| s.to_vec())
-}
-
 #[allow(dead_code)]
 pub fn block_timestamp() -> u64 {
     // NEAR timestamp is in nanoseconds
@@ -228,10 +160,6 @@ pub fn self_deploy(code_key: &[u8]) {
     }
 }
 
-pub fn save_contract<T: BorshSerialize>(key: &[u8], data: &T) {
-    write_storage(key, &data.try_to_vec().unwrap()[..]);
-}
-
 #[allow(dead_code)]
 pub fn log(data: &str) {
     log_utf8(data.as_bytes())
@@ -373,11 +301,6 @@ pub fn promise_batch_action_function_call(
     }
 }
 
-#[allow(dead_code)]
-pub fn storage_has_key(key: &[u8]) -> bool {
-    near_runtime::Runtime.storage_has_key(key)
-}
-
 pub struct ECRecoverErr;
 
 impl ECRecoverErr {

engine-sdk/src/near_runtime.rs

@@ -3,6 +3,7 @@ pub struct RegisterIndex(u64);
 
 /// Singleton type used to implement the IO traits in the case of using NEAR's
 /// runtime (i.e. for wasm contracts).
+#[derive(Copy, Clone, Default)]
 pub struct Runtime;
 
 impl Runtime {

engine-sdk/src/types.rs

@@ -1,5 +1,6 @@
+use crate::io::IO;
+use crate::panic_utf8;
 use crate::prelude::{Address, H256};
-use crate::{panic_utf8, return_output};
 
 #[cfg(not(feature = "contract"))]
 use sha3::{Digest, Keccak256};
@@ -99,7 +100,7 @@ pub trait SdkProcess<T> {
 impl<T: AsRef<[u8]>, E: AsRef<[u8]>> SdkProcess<T> for Result<T, E> {
     fn sdk_process(self) {
         match self {
-            Ok(r) => return_output(r.as_ref()),
+            Ok(r) => crate::near_runtime::Runtime.return_output(r.as_ref()),
             Err(e) => panic_utf8(e.as_ref()),
         }
     }

engine-tests/src/test_utils/mod.rs

@@ -215,7 +215,7 @@ impl AuroraRunner {
             &[crate::prelude::storage::EthConnectorStorageId::FungibleToken as u8],
         );
         let ft_value = {
-            let mut current_ft: FungibleToken = trie
+            let mut current_ft: FungibleToken<aurora_engine_sdk::near_runtime::Runtime> = trie
                 .get(&ft_key)
                 .map(|bytes| FungibleToken::try_from_slice(&bytes).unwrap())
                 .unwrap_or_default();

engine-tests/src/tests/sanity.rs

@@ -482,8 +482,7 @@ fn test_block_hash() {
         crate::prelude::u256_to_arr(&number)
     };
     let account_id = runner.aurora_account_id;
-    let block_hash =
-        aurora_engine::engine::Engine::compute_block_hash(chain_id, 10, account_id.as_bytes());
+    let block_hash = aurora_engine::engine::compute_block_hash(chain_id, 10, account_id.as_bytes());
 
     assert_eq!(
         hex::encode(block_hash.0).as_str(),