Define constructor guard trait

OffchainLabs · Feb 5, 2025 · 62c39f2 · 62c39f2
1 parent 6b89c84
commit 62c39f2
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 20 deletions.
diff --git a/stylus-core/src/host.rs b/stylus-core/src/host.rs
@@ -55,6 +55,11 @@ pub trait ValueDenier {
     fn deny_value(&self, method_name: &str) -> Result<(), Vec<u8>>;
 }
 
+/// Defines a trait that guards whether the constructor was already called.
+pub trait ConstructorGuard {
+    fn check_constructor_slot(&self) -> Result<(), Vec<u8>>;
+}
+
 /// Provides access to native cryptography extensions provided by
 /// a Stylus contract host, such as keccak256.
 pub trait CryptographyAccess {

diff --git a/stylus-proc/src/macros/public/types.rs b/stylus-proc/src/macros/public/types.rs
@@ -110,7 +110,7 @@ impl PublicImpl {
         parse_quote! {
             impl<S, #generic_params> #Router<S> for #self_ty
             where
-                S: stylus_sdk::stylus_core::storage::TopLevelStorage + core::borrow::BorrowMut<Self> + stylus_sdk::stylus_core::ValueDenier,
+                S: stylus_sdk::stylus_core::storage::TopLevelStorage + core::borrow::BorrowMut<Self> + stylus_sdk::stylus_core::ValueDenier + stylus_sdk::stylus_core::ConstructorGuard,
                 #(
                     S: core::borrow::BorrowMut<#inheritance>,
                 )*
@@ -352,7 +352,7 @@ impl<E: FnExtension> PublicFn<E> {
         parse_quote!({
             use stylus_sdk::abi::{internal, internal::EncodableReturnType};
             #deny_value
-            if let Err(e) = internal::constructor_guard() {
+            if let Err(e) = storage.check_constructor_slot() {
                 return Some(Err(e));
             }
             let args = match <#decode_inputs as #SolType>::abi_decode_params(input, true) {

diff --git a/stylus-proc/src/macros/storage.rs b/stylus-proc/src/macros/storage.rs
@@ -171,6 +171,29 @@ impl Storage {
             }
         }
     }
+    fn impl_constructor_guard(&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::stylus_core::host::ConstructorGuard for #name #ty_generics #where_clause {
+                fn check_constructor_slot(&self) -> Result<(), Vec<u8>> {
+                    let mut slot = unsafe {
+                        stylus_sdk::storage::StorageBool::new(
+                            stylus_sdk::abi::internal::CONSTRUCTOR_EXECUTED_SLOT,
+                            0,
+                            self.__stylus_host.clone()
+                        )
+                    };
+                    if slot.get() {
+                        stylus_sdk::console!("constructor already called");
+                        return Err(alloc::vec![]);
+                    }
+                    slot.set(true);
+                    Ok(())
+                }
+            }
+        }
+    }
     fn impl_from_vm(&self) -> syn::ItemImpl {
         let name = &self.name;
         let (_, ty_generics, where_clause) = self.generics.split_for_impl();
@@ -229,6 +252,7 @@ impl ToTokens for Storage {
         self.impl_storage_type().to_tokens(tokens);
         self.impl_host_access().to_tokens(tokens);
         self.impl_value_denier().to_tokens(tokens);
+        self.impl_constructor_guard().to_tokens(tokens);
         self.impl_from_vm().to_tokens(tokens);
         for field in &self.fields {
             field.impl_borrow(&self.name).to_tokens(tokens);

diff --git a/stylus-sdk/src/abi/internal.rs b/stylus-sdk/src/abi/internal.rs
@@ -4,13 +4,8 @@
 //! This module provides functions for code generated by `stylus-sdk-proc`.
 //! Most users shouldn't call these.
 
-use crate::{
-    abi::AbiType,
-    console,
-    storage::{StorageBool, StorageType},
-    ArbResult,
-};
-use alloc::{vec, vec::Vec};
+use crate::{abi::AbiType, console, ArbResult};
+use alloc::vec::Vec;
 use alloy_primitives::U256;
 use alloy_sol_types::SolType;
 use core::fmt;
@@ -68,17 +63,6 @@ pub fn failed_to_decode_arguments(err: alloy_sol_types::Error) {
     console!("failed to decode arguments: {err}");
 }
 
-#[allow(unused)]
-pub fn constructor_guard() -> Result<(), Vec<u8>> {
-    let mut slot = unsafe { StorageBool::new(CONSTRUCTOR_EXECUTED_SLOT, 0) };
-    if slot.get() {
-        console!("constructor already called");
-        return Err(vec![]);
-    }
-    slot.set(true);
-    Ok(())
-}
-
 pub trait AbiResult {
     type OkType;
 }