Parametrizing SDK Storage by a Host

stylus-proc/src/macros/entrypoint.rs

@@ -126,7 +126,7 @@ fn top_level_storage_impl(item: &syn::ItemStruct) -> syn::ItemImpl {
 
 fn struct_entrypoint_fn(name: &Ident) -> syn::ItemFn {
     parse_quote! {
-        fn #STRUCT_ENTRYPOINT_FN<H: Host>(input: alloc::vec::Vec<u8>, host: Rc<H>) -> stylus_sdk::ArbResult {
+        fn #STRUCT_ENTRYPOINT_FN<H: Host>(input: alloc::vec::Vec<u8>, host: &H) -> stylus_sdk::ArbResult {
             stylus_sdk::abi::router_entrypoint::<#name<H>, #name<H>, H>(input, host)
         }
     }
@@ -155,12 +155,12 @@ fn user_entrypoint_fn(user_fn: Ident) -> syn::ItemFn {
     parse_quote! {
         #[no_mangle]
         pub extern "C" fn user_entrypoint(len: usize) -> usize {
-            let host = Rc::new(WasmHost{});
+            let host = WasmHost{};
             #deny_reentrant
             host.pay_for_memory_grow(0);
 
             let input = host.args(len);
-            let (data, status) = match #user_fn(input, Rc::clone(&host)) {
+            let (data, status) = match #user_fn(input, &host) {
                 Ok(data) => (data, 0),
                 Err(data) => (data, 1),
             };

stylus-proc/src/macros/storage.rs

@@ -73,7 +73,7 @@ impl Storage {
                 const SLOT_BYTES: usize = 32;
                 const REQUIRED_SLOTS: usize = Self::required_slots();
 
-                unsafe fn new(mut root: stylus_sdk::alloy_primitives::U256, offset: u8, host: Rc<H>) -> Self {
+                unsafe fn new(mut root: stylus_sdk::alloy_primitives::U256, offset: u8, host: *const H) -> Self {
                     use stylus_sdk::{storage, alloy_primitives};
                     debug_assert!(offset == 0);
 
@@ -96,6 +96,20 @@ impl Storage {
             }
         }
     }
+
+    fn impl_host_access(&self) -> syn::ItemImpl {
+        let name = &self.name;
+        let (impl_generics, ty_generics, where_clause) = self.generics.split_for_impl();
+        parse_quote! {
+            impl #impl_generics stylus_sdk::host::HostAccess<H> for #name #ty_generics #where_clause {
+                fn get_host(&self) -> &H {
+                    // SAFETY: Host is guaranteed to be valid and non-null for the lifetime of the storage
+                    // as injected by the Stylus entrypoint function.
+                    unsafe { &*self.host }
+                }
+            }
+        }
+    }
 }
 
 impl From<&mut syn::ItemStruct> for Storage {
@@ -126,6 +140,7 @@ impl ToTokens for Storage {
     fn to_tokens(&self, tokens: &mut TokenStream) {
         self.item_impl().to_tokens(tokens);
         self.impl_storage_type().to_tokens(tokens);
+        self.impl_host_access().to_tokens(tokens);
         for field in &self.fields {
             field.impl_borrow(&self.name).to_tokens(tokens);
             field.impl_borrow_mut(&self.name).to_tokens(tokens);
@@ -166,7 +181,7 @@ impl StorageField {
             return None;
         };
         let ty = &self.ty;
-        if ty.to_token_stream().to_string() == "Rc < H >".to_string() {
+        if ty.to_token_stream().to_string() == "*const H".to_string() {
             return None;
         }
         Some(parse_quote! {
@@ -180,7 +195,7 @@ impl StorageField {
                 space -= bytes;
 
                 let root = root + alloy_primitives::U256::from(slot);
-                let field = <#ty as storage::StorageType<H>>::new(root, space as u8, Rc::clone(&host));
+                let field = <#ty as storage::StorageType<H>>::new(root, space as u8, host);
                 if words > 0 {
                     slot += words;
                     space = 32;
@@ -192,7 +207,8 @@ impl StorageField {
 
     fn size(&self) -> TokenStream {
         let ty = &self.ty;
-        if ty.to_token_stream().to_string() == "Rc < H >".to_string() {
+        println!("{}", ty.to_token_stream().to_string());
+        if ty.to_token_stream().to_string() == "*const H".to_string() {
             return quote! {};
         }
         quote! {

stylus-sdk/src/abi/mod.rs

@@ -18,7 +18,6 @@
 use alloc::vec::Vec;
 use alloy_primitives::U256;
 use core::borrow::BorrowMut;
-use rclite::Rc;
 
 use alloy_sol_types::{abi::TokenSeq, private::SolTypeValue, SolType};
 
@@ -98,7 +97,7 @@ where
 //    if no value is received in the transaction. It is implicitly payable.
 //  - Fallback is called when no other function matches a selector. If a receive function is not
 //    defined, then calls with no input calldata will be routed to the fallback function.
-pub fn router_entrypoint<R, S, H>(input: alloc::vec::Vec<u8>, host: Rc<H>) -> ArbResult
+pub fn router_entrypoint<R, S, H>(input: alloc::vec::Vec<u8>, host: &H) -> ArbResult
 where
     R: Router<S>,
     S: StorageType<H> + TopLevelStorage + BorrowMut<R::Storage>,

stylus-sdk/src/host/mod.rs

@@ -27,6 +27,10 @@ pub trait Host:
 {
 }
 
+pub trait HostAccess<H: Host> {
+    fn get_host(&self) -> &H;
+}
+
 /// TODO
 pub trait CryptographyAccess {
     /// TODO

stylus-sdk/src/storage/array.rs

@@ -1,18 +1,17 @@
 // Copyright 2023-2024, Offchain Labs, Inc.
 // For licensing, see https://github.com/OffchainLabs/stylus-sdk-rs/blob/main/licenses/COPYRIGHT.md
 
-use crate::host::Host;
+use crate::host::{Host, HostAccess};
 
 use super::{Erase, StorageGuard, StorageGuardMut, StorageType};
 use alloy_primitives::U256;
 use core::marker::PhantomData;
-use rclite::Rc;
 
 /// Accessor for a storage-backed array.
 pub struct StorageArray<H: Host, S: StorageType<H>, const N: usize> {
     slot: U256,
     marker: PhantomData<S>,
-    host: Rc<H>,
+    host: *const H,
 }
 
 impl<H: Host, S: StorageType<H>, const N: usize> StorageType<H> for StorageArray<H, S, N> {
@@ -27,7 +26,7 @@ impl<H: Host, S: StorageType<H>, const N: usize> StorageType<H> for StorageArray
 
     const REQUIRED_SLOTS: usize = Self::required_slots();
 
-    unsafe fn new(slot: U256, offset: u8, host: Rc<H>) -> Self {
+    unsafe fn new(slot: U256, offset: u8, host: *const H) -> Self {
         debug_assert!(offset == 0);
         Self {
             slot,
@@ -45,6 +44,14 @@ impl<H: Host, S: StorageType<H>, const N: usize> StorageType<H> for StorageArray
     }
 }
 
+impl<H: Host, S: StorageType<H>, const N: usize> HostAccess<H> for StorageArray<H, S, N> {
+    fn get_host(&self) -> &H {
+        // SAFETY: Host is guaranteed to be valid and non-null for the lifetime of the storage
+        // as injected by the Stylus entrypoint function.
+        unsafe { &*self.host }
+    }
+}
+
 impl<H: Host, S: StorageType<H>, const N: usize> StorageArray<H, S, N> {
     /// Gets the number of elements stored.
     ///
@@ -82,7 +89,7 @@ impl<H: Host, S: StorageType<H>, const N: usize> StorageArray<H, S, N> {
             return None;
         }
         let (slot, offset) = self.index_slot(index);
-        Some(S::new(slot, offset, Rc::clone(&self.host)))
+        Some(S::new(slot, offset, self.get_host()))
     }
 
     /// Gets the underlying accessor to the element at a given index, even if out of bounds.
@@ -92,7 +99,7 @@ impl<H: Host, S: StorageType<H>, const N: usize> StorageArray<H, S, N> {
     /// Enables aliasing. UB if out of bounds.
     unsafe fn accessor_unchecked(&self, index: usize) -> S {
         let (slot, offset) = self.index_slot(index);
-        S::new(slot, offset, Rc::clone(&self.host))
+        S::new(slot, offset, self.get_host())
     }
 
     /// Gets the element at the given index, if it exists.

stylus-sdk/src/storage/bytes.rs

@@ -2,20 +2,22 @@
 // For licensing, see https://github.com/OffchainLabs/stylus-sdk-rs/blob/main/licenses/COPYRIGHT.md
 
 use super::{Erase, GlobalStorage, Storage, StorageB8, StorageGuard, StorageGuardMut, StorageType};
-use crate::{crypto, host::Host};
+use crate::{
+    crypto,
+    host::{Host, HostAccess},
+};
 use alloc::{
     string::{String, ToString},
     vec::Vec,
 };
 use alloy_primitives::{U256, U8};
 use core::cell::OnceCell;
-use rclite::Rc;
 
 /// Accessor for storage-backed bytes.
 pub struct StorageBytes<H: Host> {
     root: U256,
     base: OnceCell<U256>,
-    host: Rc<H>,
+    host: *const H,
 }
 
 impl<H: Host> StorageType<H> for StorageBytes<H> {
@@ -28,7 +30,7 @@ impl<H: Host> StorageType<H> for StorageBytes<H> {
     where
         Self: 'a;
 
-    unsafe fn new(root: U256, offset: u8, host: Rc<H>) -> Self {
+    unsafe fn new(root: U256, offset: u8, host: *const H) -> Self {
         debug_assert!(offset == 0);
         Self {
             root,
@@ -46,6 +48,14 @@ impl<H: Host> StorageType<H> for StorageBytes<H> {
     }
 }
 
+impl<H: Host> HostAccess<H> for StorageBytes<H> {
+    fn get_host(&self) -> &H {
+        // SAFETY: Host is guaranteed to be valid and non-null for the lifetime of the storage
+        // as injected by the Stylus entrypoint function.
+        unsafe { &*self.host }
+    }
+}
+
 impl<H: Host> StorageBytes<H> {
     /// Returns `true` if the collection contains no elements.
     pub fn is_empty(&self) -> bool {
@@ -54,7 +64,7 @@ impl<H: Host> StorageBytes<H> {
 
     /// Gets the number of bytes stored.
     pub fn len(&self) -> usize {
-        let word = Storage::get_word(Rc::clone(&self.host), self.root);
+        let word = Storage::get_word(self.get_host(), self.root);
 
         // check if the data is short
         let slot: &[u8] = word.as_ref();
@@ -83,30 +93,26 @@ impl<H: Host> StorageBytes<H> {
 
         // if shrinking, pull data in
         if (len < 32) && (old > 32) {
-            let word = Storage::get_word(Rc::clone(&self.host), *self.base());
-            Storage::set_word(Rc::clone(&self.host), self.root, word);
+            let word = Storage::get_word(self.get_host(), *self.base());
+            Storage::set_word(self.get_host(), self.root, word);
             return self.write_len(len);
         }
 
         // if growing, push data out
-        let mut word = Storage::get_word(Rc::clone(&self.host), self.root);
+        let mut word = Storage::get_word(self.get_host(), self.root);
         word[31] = 0; // clear len byte
-        Storage::set_word(Rc::clone(&self.host), *self.base(), word);
+        Storage::set_word(self.get_host(), *self.base(), word);
         self.write_len(len)
     }
 
     /// Updates the length while being conscious of representation.
     unsafe fn write_len(&mut self, len: usize) {
         if len < 32 {
             // place the len in the last byte of the root with the long bit low
-            Storage::set_uint(Rc::clone(&self.host), self.root, 31, U8::from(len * 2));
+            Storage::set_uint(self.get_host(), self.root, 31, U8::from(len * 2));
         } else {
             // place the len in the root with the long bit high
-            Storage::set_word(
-                Rc::clone(&self.host),
-                self.root,
-                U256::from(len * 2 + 1).into(),
-            )
+            Storage::set_word(self.get_host(), self.root, U256::from(len * 2 + 1).into())
         }
     }
 
@@ -118,7 +124,7 @@ impl<H: Host> StorageBytes<H> {
         macro_rules! assign {
             ($slot:expr) => {
                 unsafe {
-                    Storage::set_uint(Rc::clone(&self.host), $slot, index % 32, value); // pack value
+                    Storage::set_uint(self.get_host(), $slot, index % 32, value); // pack value
                     self.write_len(index + 1);
                 }
             };
@@ -131,8 +137,8 @@ impl<H: Host> StorageBytes<H> {
         // convert to multi-word representation
         if index == 31 {
             // copy content over (len byte will be overwritten)
-            let word = Storage::get_word(Rc::clone(&self.host), self.root);
-            unsafe { Storage::set_word(Rc::clone(&self.host), *self.base(), word) };
+            let word = Storage::get_word(self.get_host(), self.root);
+            unsafe { Storage::set_word(self.get_host(), *self.base(), word) };
         }
 
         let slot = self.base() + U256::from(index / 32);
@@ -152,13 +158,13 @@ impl<H: Host> StorageBytes<H> {
         let clean = index % 32 == 0;
         let byte = self.get(index)?;
 
-        let clear = |slot| unsafe { Storage::clear_word(Rc::clone(&self.host), slot) };
+        let clear = |slot| unsafe { Storage::clear_word(self.get_host(), slot) };
 
         // convert to single-word representation
         if len == 32 {
             // copy content over
-            let word = Storage::get_word(Rc::clone(&self.host), *self.base());
-            unsafe { Storage::set_word(Rc::clone(&self.host), self.root, word) };
+            let word = Storage::get_word(self.get_host(), *self.base());
+            unsafe { Storage::set_word(self.get_host(), self.root, word) };
             clear(*self.base());
         }
 
@@ -169,7 +175,7 @@ impl<H: Host> StorageBytes<H> {
 
         // clear the value
         if len < 32 {
-            unsafe { Storage::set_byte(Rc::clone(&self.host), self.root, index, 0) };
+            unsafe { Storage::set_byte(self.get_host(), self.root, index, 0) };
         }
 
         // set the new length
@@ -193,7 +199,7 @@ impl<H: Host> StorageBytes<H> {
             return None;
         }
         let (slot, offset) = self.index_slot(index);
-        let value = unsafe { StorageB8::new(slot, offset, Rc::clone(&self.host)) };
+        let value = unsafe { StorageB8::new(slot, offset, self.get_host()) };
         Some(StorageGuardMut::new(value))
     }
 
@@ -204,7 +210,7 @@ impl<H: Host> StorageBytes<H> {
     /// UB if index is out of bounds.
     pub unsafe fn get_unchecked(&self, index: usize) -> u8 {
         let (slot, offset) = self.index_slot(index);
-        unsafe { Storage::get_byte(Rc::clone(&self.host), slot, offset.into()) }
+        unsafe { Storage::get_byte(self.get_host(), slot, offset.into()) }
     }
 
     /// Gets the full contents of the collection.
@@ -247,11 +253,11 @@ impl<H: Host> Erase<H> for StorageBytes<H> {
         if len > 31 {
             while len > 0 {
                 let slot = self.index_slot(len as usize - 1).0;
-                unsafe { Storage::clear_word(Rc::clone(&self.host), slot) };
+                unsafe { Storage::clear_word(self.get_host(), slot) };
                 len -= 32;
             }
         }
-        unsafe { Storage::clear_word(Rc::clone(&self.host), self.root) };
+        unsafe { Storage::clear_word(self.get_host(), self.root) };
     }
 }
 
@@ -284,7 +290,7 @@ impl<H: Host> StorageType<H> for StorageString<H> {
     where
         Self: 'a;
 
-    unsafe fn new(slot: U256, offset: u8, host: Rc<H>) -> Self {
+    unsafe fn new(slot: U256, offset: u8, host: *const H) -> Self {
         Self(StorageBytes::new(slot, offset, host))
     }