Beacon Proxy

use openzeppelin_stylus::proxy::{
    beacon::{proxy::BeaconProxy, IBeacon},
    erc1967,
    IProxy,
};
use stylus_sdk::{
    abi::Bytes,
    alloy_primitives::Address,
    prelude::*,
    ArbResult,
};

#[entrypoint]
#[storage]
struct MyBeaconProxy {
    beacon_proxy: BeaconProxy,
}

#[public]
impl MyBeaconProxy {
    #[constructor]
    fn constructor(
        &mut self,
        beacon: Address,
        data: Bytes,
    ) -> Result<(), erc1967::utils::Error> {
        self.beacon_proxy.constructor(beacon, &data)
    }

    /// Get the beacon address
    fn get_beacon(&self) -> Address {
        self.beacon_proxy.get_beacon()
    }

    /// Get the current implementation address from the beacon
    fn implementation(&self) -> Result<Address, Vec<u8>> {
        self.beacon_proxy.implementation()
    }

    /// Fallback function that delegates all calls to the implementation
    #[fallback]
    fn fallback(&mut self, calldata: &[u8]) -> ArbResult {
        unsafe { self.do_fallback(calldata) }
    }
}

unsafe impl IProxy for MyBeaconProxy {
    fn implementation(&self) -> Result<Address, Vec<u8>> {
        self.beacon_proxy.implementation()
    }
}

use openzeppelin_stylus::{
    access::ownable::{IOwnable, Ownable},
    proxy::beacon::{IBeacon, IUpgradeableBeacon, UpgradeableBeacon},
};
use stylus_sdk::{
    alloy_primitives::Address,
    prelude::*,
};

#[entrypoint]
#[storage]
struct MyUpgradeableBeacon {
    beacon: UpgradeableBeacon,
}

#[public]
impl MyUpgradeableBeacon {
    #[constructor]
    fn constructor(
        &mut self,
        implementation: Address,
        initial_owner: Address,
    ) -> Result<(), beacon::Error> {
        self.beacon.constructor(implementation, initial_owner)
    }

    /// Upgrade to a new implementation (only owner)
    fn upgrade_to(
        &mut self,
        new_implementation: Address,
    ) -> Result<(), beacon::Error> {
        self.beacon.upgrade_to(new_implementation)
    }
}

#[public]
impl IBeacon for MyUpgradeableBeacon {
    fn implementation(&self) -> Result<Address, Vec<u8>> {
        self.beacon.implementation()
    }
}

#[public]
impl IOwnable for MyUpgradeableBeacon {
    fn owner(&self) -> Address {
        self.beacon.owner()
    }

    fn transfer_ownership(&mut self, new_owner: Address) -> Result<(), Vec<u8>> {
        self.beacon.transfer_ownership(new_owner)
    }

    fn renounce_ownership(&mut self) -> Result<(), Vec<u8>> {
        self.beacon.renounce_ownership()
    }
}

#[public]
impl IUpgradeableBeacon for MyUpgradeableBeacon {
    fn upgrade_to(&mut self, new_implementation: Address) -> Result<(), Vec<u8>> {
        Ok(self.beacon.upgrade_to(new_implementation)?)
    }
}

use openzeppelin_stylus::proxy::beacon::IBeacon;
use stylus_sdk::{
    alloy_primitives::Address,
    prelude::*,
    storage::StorageAddress,
};

#[entrypoint]
#[storage]
struct MyCustomBeacon {
    implementation: StorageAddress,
    admin: StorageAddress,
}

#[public]
impl MyCustomBeacon {
    #[constructor]
    fn constructor(&mut self, implementation: Address, admin: Address) {
        self.implementation.set(implementation);
        self.admin.set(admin);
    }

    /// Upgrade implementation (only admin)
    fn upgrade_implementation(&mut self, new_implementation: Address) -> Result<(), Vec<u8>> {
        if self.admin.get() != msg::sender() {
            return Err("Only admin can upgrade".abi_encode());
        }

        if !new_implementation.has_code() {
            return Err("Invalid implementation".abi_encode());
        }

        self.implementation.set(new_implementation);
        Ok(())
    }
}

#[public]
impl IBeacon for MyCustomBeacon {
    fn implementation(&self) -> Result<Address, Vec<u8>> {
        Ok(self.implementation.get())
    }
}

impl MyBeaconProxy {
    #[constructor]
    fn constructor(
        &mut self,
        beacon: Address,
        data: Bytes,
    ) -> Result<(), erc1967::utils::Error> {
        // If data is provided, it will be passed to the implementation
        // returned by the beacon during construction via delegatecall
        self.beacon_proxy.constructor(beacon, &data)
    }
}