Refactored to upgradeable contracts (2/2) (Dispute, Royalties)

contracts/AccessController.sol

@@ -52,7 +52,7 @@ contract AccessController is IAccessController, GovernableUpgradeable, UUPSUpgra
         _disableInitializers();
     }
 
-    /// @notice Initializes implementation contract
+    /// @notice initializer for this implementation contract
     /// @param governance The address of the governance contract
     function initialize(address governance) external initializer {
         __GovernableUpgradeable_init(governance);

contracts/access/AccessControlled.sol

@@ -43,7 +43,7 @@ abstract contract AccessControlled {
         _;
     }
 
-    /// @dev Constructor contract by setting the ACCESS_CONTROLLER and IP_ACCOUNT_REGISTRY addresses.
+    /// @dev Constructor
     /// @param accessController The address of the AccessController contract.
     /// @param ipAccountRegistry The address of the IPAccountRegistry contract.
     /// @custom:oz-upgrades-unsafe-allow constructor

contracts/governance/GovernableUpgradeable.sol

@@ -45,6 +45,8 @@ abstract contract GovernableUpgradeable is IGovernable, Initializable {
         _disableInitializers();
     }
 
+    /// @notice initializer for this implementation contract
+    /// @param governance_ The address of the governance.
     function __GovernableUpgradeable_init(address governance_) internal onlyInitializing {
         if (governance_ == address(0)) revert Errors.Governance__ZeroAddress();
         _getGovernableUpgradeableStorage().governance = governance_;

contracts/interfaces/modules/royalty/IRoyaltyModule.sol

@@ -31,7 +31,7 @@ interface IRoyaltyModule is IModule {
     event LicenseMintingFeePaid(address receiverIpId, address payerAddress, address token, uint256 amount);
 
     /// @notice Returns the licensing module address
-    function LICENSING_MODULE() external view returns (address);
+    function licensingModule() external view returns (address);
 
     /// @notice Indicates if a royalty policy is whitelisted
     /// @param royaltyPolicy The address of the royalty policy

contracts/interfaces/modules/royalty/policies/IRoyaltyPolicyLAP.sol

@@ -46,7 +46,7 @@ interface IRoyaltyPolicyLAP is IRoyaltyPolicy {
     function LIQUID_SPLIT_MAIN() external view returns (address);
 
     /// @notice Returns the Ancestors Vault Implementation address
-    function ANCESTORS_VAULT_IMPL() external view returns (address);
+    function ancestorsVaultImpl() external view returns (address);
 
     /// @notice Distributes funds internally so that accounts holding the royalty nfts at distribution moment can
     /// claim afterwards

contracts/modules/dispute/policies/ArbitrationPolicySP.sol

@@ -3,23 +3,27 @@ pragma solidity 0.8.23;
 
 import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
 
-import { Governable } from "../../../governance/Governable.sol";
+import { GovernableUpgradeable } from "../../../governance/GovernableUpgradeable.sol";
 import { IDisputeModule } from "../../../interfaces/modules/dispute/IDisputeModule.sol";
 import { IArbitrationPolicy } from "../../../interfaces/modules/dispute/policies/IArbitrationPolicy.sol";
 import { Errors } from "../../../lib/Errors.sol";
 
 /// @title Story Protocol Arbitration Policy
 /// @notice The Story Protocol arbitration policy is a simple policy that requires the dispute initiator to pay a fixed
 ///         amount of tokens to raise a dispute and refunds that amount if the dispute initiator wins the dispute.
-contract ArbitrationPolicySP is IArbitrationPolicy, Governable {
+contract ArbitrationPolicySP is IArbitrationPolicy, GovernableUpgradeable, UUPSUpgradeable {
     using SafeERC20 for IERC20;
 
     /// @notice Returns the dispute module address
+    /// @custom:oz-upgrades-unsafe-allow state-variable-immutable
     address public immutable DISPUTE_MODULE;
     /// @notice Returns the payment token address
+    /// @custom:oz-upgrades-unsafe-allow state-variable-immutable
     address public immutable PAYMENT_TOKEN;
     /// @notice Returns the arbitration price
+    /// @custom:oz-upgrades-unsafe-allow state-variable-immutable
     uint256 public immutable ARBITRATION_PRICE;
 
     /// @notice Restricts the calls to the DisputeModule
@@ -28,12 +32,12 @@ contract ArbitrationPolicySP is IArbitrationPolicy, Governable {
         _;
     }
 
-    constructor(
-        address _disputeModule,
-        address _paymentToken,
-        uint256 _arbitrationPrice,
-        address _governable
-    ) Governable(_governable) {
+    /// Constructor
+    /// @param _disputeModule The dispute module address
+    /// @param _paymentToken The ERC20 payment token address
+    /// @param _arbitrationPrice The arbitration price
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor(address _disputeModule, address _paymentToken, uint256 _arbitrationPrice) {
         if (_disputeModule == address(0)) revert Errors.ArbitrationPolicySP__ZeroDisputeModule();
         if (_paymentToken == address(0)) revert Errors.ArbitrationPolicySP__ZeroPaymentToken();
 
@@ -42,6 +46,13 @@ contract ArbitrationPolicySP is IArbitrationPolicy, Governable {
         ARBITRATION_PRICE = _arbitrationPrice;
     }
 
+    /// @notice initializer for this implementation contract
+    /// @param governance The address of the governance contract
+    function initialize(address governance) public initializer {
+        __GovernableUpgradeable_init(governance);
+        __UUPSUpgradeable_init();
+    }
+
     /// @notice Executes custom logic on raising dispute.
     /// @dev Enforced to be only callable by the DisputeModule.
     /// @param caller Address of the caller
@@ -74,8 +85,12 @@ contract ArbitrationPolicySP is IArbitrationPolicy, Governable {
     /// @dev Enforced to be only callable by the governance protocol admin.
     function governanceWithdraw() external onlyProtocolAdmin {
         uint256 balance = IERC20(PAYMENT_TOKEN).balanceOf(address(this));
-        IERC20(PAYMENT_TOKEN).safeTransfer(governance, balance);
+        IERC20(PAYMENT_TOKEN).safeTransfer(msg.sender, balance);
 
         emit GovernanceWithdrew(balance);
     }
+
+    /// @notice Hook that is called before any upgrade
+    /// @param newImplementation Address of the new implementation
+    function _authorizeUpgrade(address newImplementation) internal override onlyProtocolAdmin {}
 }

contracts/modules/licensing/BasePolicyFrameworkManager.sol

@@ -47,7 +47,7 @@ abstract contract BasePolicyFrameworkManager is IPolicyFrameworkManager, ERC165,
         LICENSING_MODULE = ILicensingModule(licensing);
     }
 
-    /// @notice Initializes the BasePolicyFrameworkManager contract as per the Initializable contract.
+    /// @notice initializer for this implementation contract
     /// @param _name The name of the policy framework manager
     /// @param _licenseTextUrl The URL to the off chain legal agreement template text
     function __BasePolicyFrameworkManager_init(

contracts/modules/licensing/LicensingModule.sol

@@ -117,6 +117,8 @@ contract LicensingModule is
         _disableInitializers();
     }
 
+    /// @notice initializer for this implementation contract
+    /// @param governance The address of the governance contract
     function initialize(address governance) public initializer {
         __ReentrancyGuard_init();
         __UUPSUpgradeable_init();

contracts/modules/royalty/RoyaltyModule.sol

@@ -1,12 +1,13 @@
 // SPDX-License-Identifier: BUSL-1.1
 pragma solidity 0.8.23;
 
-import { ReentrancyGuard } from "@openzeppelin/contracts/utils/ReentrancyGuard.sol";
+import { ReentrancyGuardUpgradeable } from "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardUpgradeable.sol";
+import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
 import { ERC165Checker } from "@openzeppelin/contracts/utils/introspection/ERC165Checker.sol";
 import { IERC165 } from "@openzeppelin/contracts/utils/introspection/IERC165.sol";
 
 import { BaseModule } from "../BaseModule.sol";
-import { Governable } from "../../governance/Governable.sol";
+import { GovernableUpgradeable } from "../../governance/GovernableUpgradeable.sol";
 import { IRoyaltyModule } from "../../interfaces/modules/royalty/IRoyaltyModule.sol";
 import { IRoyaltyPolicy } from "../../interfaces/modules/royalty/policies/IRoyaltyPolicy.sol";
 import { Errors } from "../../lib/Errors.sol";
@@ -16,37 +17,61 @@ import { BaseModule } from "../BaseModule.sol";
 /// @title Story Protocol Royalty Module
 /// @notice The Story Protocol royalty module allows to set royalty policies an IP asset and pay royalties as a
 ///         derivative IP.
-contract RoyaltyModule is IRoyaltyModule, Governable, ReentrancyGuard, BaseModule {
+contract RoyaltyModule is
+    IRoyaltyModule,
+    GovernableUpgradeable,
+    ReentrancyGuardUpgradeable,
+    BaseModule,
+    UUPSUpgradeable
+{
     using ERC165Checker for address;
 
-    string public constant override name = ROYALTY_MODULE_KEY;
-
-    /// @notice Returns the licensing module address
-    address public LICENSING_MODULE;
+    /// @dev Storage structure for the RoyaltyModule
+    /// @param licensingModule The address of the licensing module
+    /// @param isWhitelistedRoyaltyPolicy Indicates if a royalty policy is whitelisted
+    /// @param isWhitelistedRoyaltyToken Indicates if a royalty token is whitelisted
+    /// @param royaltyPolicies Indicates the royalty policy for a given IP asset
+    /// @custom:storage-location erc7201:story-protocol.RoyaltyModule
+    struct RoyaltyModuleStorage {
+        address licensingModule;
+        mapping(address royaltyPolicy => bool isWhitelisted) isWhitelistedRoyaltyPolicy;
+        mapping(address token => bool) isWhitelistedRoyaltyToken;
+        mapping(address ipId => address royaltyPolicy) royaltyPolicies;
+    }
 
-    /// @notice Indicates if a royalty policy is whitelisted
-    mapping(address royaltyPolicy => bool isWhitelisted) public isWhitelistedRoyaltyPolicy;
+    // keccak256(abi.encode(uint256(keccak256("story-protocol.RoyaltyModule")) - 1)) & ~bytes32(uint256(0xff));
+    bytes32 private constant RoyaltyModuleStorageLocation =
+        0x98dd2c34f21d19fd1d178ed731f3db3d03e0b4e39f02dbeb040e80c9427a0300;
 
-    /// @notice Indicates if a royalty token is whitelisted
-    mapping(address token => bool) public isWhitelistedRoyaltyToken;
+    string public constant override name = ROYALTY_MODULE_KEY;
 
-    /// @notice Indicates the royalty policy for a given IP asset
-    mapping(address ipId => address royaltyPolicy) public royaltyPolicies;
+    /// @notice Constructor
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor() {
+        _disableInitializers();
+    }
 
-    constructor(address governance) Governable(governance) {}
+    /// @notice initializer for this implementation contract
+    /// @param _governance The address of the governance contract
+    function initialize(address _governance) external initializer {
+        __GovernableUpgradeable_init(_governance);
+        __ReentrancyGuard_init();
+        __UUPSUpgradeable_init();
+    }
 
     /// @notice Modifier to enforce that the caller is the licensing module
     modifier onlyLicensingModule() {
-        if (msg.sender != LICENSING_MODULE) revert Errors.RoyaltyModule__NotAllowedCaller();
+        RoyaltyModuleStorage storage $ = _getRoyaltyModuleStorage();
+        if (msg.sender != $.licensingModule) revert Errors.RoyaltyModule__NotAllowedCaller();
         _;
     }
 
     /// @notice Sets the license registry
     /// @dev Enforced to be only callable by the protocol admin
-    /// @param licensingModule The address of the license registry
-    function setLicensingModule(address licensingModule) external onlyProtocolAdmin {
-        if (licensingModule == address(0)) revert Errors.RoyaltyModule__ZeroLicensingModule();
-        LICENSING_MODULE = licensingModule;
+    /// @param licensing The address of the license registry
+    function setLicensingModule(address licensing) external onlyProtocolAdmin {
+        if (licensing == address(0)) revert Errors.RoyaltyModule__ZeroLicensingModule();
+        _getRoyaltyModuleStorage().licensingModule = licensing;
     }
 
     /// @notice Whitelist a royalty policy
@@ -56,7 +81,8 @@ contract RoyaltyModule is IRoyaltyModule, Governable, ReentrancyGuard, BaseModul
     function whitelistRoyaltyPolicy(address royaltyPolicy, bool allowed) external onlyProtocolAdmin {
         if (royaltyPolicy == address(0)) revert Errors.RoyaltyModule__ZeroRoyaltyPolicy();
 
-        isWhitelistedRoyaltyPolicy[royaltyPolicy] = allowed;
+        RoyaltyModuleStorage storage $ = _getRoyaltyModuleStorage();
+        $.isWhitelistedRoyaltyPolicy[royaltyPolicy] = allowed;
 
         emit RoyaltyPolicyWhitelistUpdated(royaltyPolicy, allowed);
     }
@@ -68,7 +94,8 @@ contract RoyaltyModule is IRoyaltyModule, Governable, ReentrancyGuard, BaseModul
     function whitelistRoyaltyToken(address token, bool allowed) external onlyProtocolAdmin {
         if (token == address(0)) revert Errors.RoyaltyModule__ZeroRoyaltyToken();
 
-        isWhitelistedRoyaltyToken[token] = allowed;
+        RoyaltyModuleStorage storage $ = _getRoyaltyModuleStorage();
+        $.isWhitelistedRoyaltyToken[token] = allowed;
 
         emit RoyaltyTokenWhitelistUpdated(token, allowed);
     }
@@ -85,9 +112,11 @@ contract RoyaltyModule is IRoyaltyModule, Governable, ReentrancyGuard, BaseModul
         bytes calldata licenseData,
         bytes calldata externalData
     ) external nonReentrant onlyLicensingModule {
-        if (!isWhitelistedRoyaltyPolicy[royaltyPolicy]) revert Errors.RoyaltyModule__NotWhitelistedRoyaltyPolicy();
+        RoyaltyModuleStorage storage $ = _getRoyaltyModuleStorage();
 
-        address royaltyPolicyIpId = royaltyPolicies[ipId];
+        if (!$.isWhitelistedRoyaltyPolicy[royaltyPolicy]) revert Errors.RoyaltyModule__NotWhitelistedRoyaltyPolicy();
+
+        address royaltyPolicyIpId = $.royaltyPolicies[ipId];
 
         // if the node is a root node, then royaltyPolicyIpId will be address(0) and any type of royalty type can be
         // selected to mint a license if the node is a derivative node, then the any minted licenses by the derivative
@@ -113,11 +142,12 @@ contract RoyaltyModule is IRoyaltyModule, Governable, ReentrancyGuard, BaseModul
         bytes[] memory licenseData,
         bytes calldata externalData
     ) external nonReentrant onlyLicensingModule {
-        if (!isWhitelistedRoyaltyPolicy[royaltyPolicy]) revert Errors.RoyaltyModule__NotWhitelistedRoyaltyPolicy();
+        RoyaltyModuleStorage storage $ = _getRoyaltyModuleStorage();
+        if (!$.isWhitelistedRoyaltyPolicy[royaltyPolicy]) revert Errors.RoyaltyModule__NotWhitelistedRoyaltyPolicy();
         if (parentIpIds.length == 0) revert Errors.RoyaltyModule__NoParentsOnLinking();
 
         for (uint32 i = 0; i < parentIpIds.length; i++) {
-            address parentRoyaltyPolicy = royaltyPolicies[parentIpIds[i]];
+            address parentRoyaltyPolicy = $.royaltyPolicies[parentIpIds[i]];
             // if the parent node has a royalty policy set, then the derivative node should have the same royalty
             // policy if the parent node does not have a royalty policy set, then the derivative node can set any type
             // of royalty policy as long as the children ip obtained and is burning all licenses with that royalty type
@@ -126,7 +156,7 @@ contract RoyaltyModule is IRoyaltyModule, Governable, ReentrancyGuard, BaseModul
                 revert Errors.RoyaltyModule__IncompatibleRoyaltyPolicy();
         }
 
-        royaltyPolicies[ipId] = royaltyPolicy;
+        $.royaltyPolicies[ipId] = royaltyPolicy;
 
         IRoyaltyPolicy(royaltyPolicy).onLinkToParents(ipId, parentIpIds, licenseData, externalData);
     }
@@ -142,13 +172,15 @@ contract RoyaltyModule is IRoyaltyModule, Governable, ReentrancyGuard, BaseModul
         address token,
         uint256 amount
     ) external nonReentrant {
-        if (!isWhitelistedRoyaltyToken[token]) revert Errors.RoyaltyModule__NotWhitelistedRoyaltyToken();
+        RoyaltyModuleStorage storage $ = _getRoyaltyModuleStorage();
+        if (!$.isWhitelistedRoyaltyToken[token]) revert Errors.RoyaltyModule__NotWhitelistedRoyaltyToken();
 
-        address payerRoyaltyPolicy = royaltyPolicies[payerIpId];
+        address payerRoyaltyPolicy = $.royaltyPolicies[payerIpId];
         // if the payer does not have a royalty policy set, then the payer is not a derivative ip and does not pay
         // royalties the receiver ip can have a zero royalty policy since that could mean it is an ip a root
         if (payerRoyaltyPolicy == address(0)) revert Errors.RoyaltyModule__NoRoyaltyPolicySet();
-        if (!isWhitelistedRoyaltyPolicy[payerRoyaltyPolicy]) revert Errors.RoyaltyModule__NotWhitelistedRoyaltyPolicy();
+        if (!$.isWhitelistedRoyaltyPolicy[payerRoyaltyPolicy])
+            revert Errors.RoyaltyModule__NotWhitelistedRoyaltyPolicy();
 
         IRoyaltyPolicy(payerRoyaltyPolicy).onRoyaltyPayment(msg.sender, receiverIpId, token, amount);
 
@@ -168,19 +200,57 @@ contract RoyaltyModule is IRoyaltyModule, Governable, ReentrancyGuard, BaseModul
         address token,
         uint256 amount
     ) external onlyLicensingModule {
-        if (!isWhitelistedRoyaltyToken[token]) revert Errors.RoyaltyModule__NotWhitelistedRoyaltyToken();
+        RoyaltyModuleStorage storage $ = _getRoyaltyModuleStorage();
+        if (!$.isWhitelistedRoyaltyToken[token]) revert Errors.RoyaltyModule__NotWhitelistedRoyaltyToken();
 
         if (licenseRoyaltyPolicy == address(0)) revert Errors.RoyaltyModule__NoRoyaltyPolicySet();
-        if (!isWhitelistedRoyaltyPolicy[licenseRoyaltyPolicy])
+        if (!$.isWhitelistedRoyaltyPolicy[licenseRoyaltyPolicy])
             revert Errors.RoyaltyModule__NotWhitelistedRoyaltyPolicy();
 
         IRoyaltyPolicy(licenseRoyaltyPolicy).onRoyaltyPayment(payerAddress, receiverIpId, token, amount);
 
         emit LicenseMintingFeePaid(receiverIpId, payerAddress, token, amount);
     }
 
+    /// @notice Returns the licensing module address
+    function licensingModule() external view returns (address) {
+        return _getRoyaltyModuleStorage().licensingModule;
+    }
+
+    /// @notice Indicates if a royalty policy is whitelisted
+    /// @param royaltyPolicy The address of the royalty policy
+    /// @return isWhitelisted True if the royalty policy is whitelisted
+    function isWhitelistedRoyaltyPolicy(address royaltyPolicy) external view returns (bool) {
+        return _getRoyaltyModuleStorage().isWhitelistedRoyaltyPolicy[royaltyPolicy];
+    }
+
+    /// @notice Indicates if a royalty token is whitelisted
+    /// @param token The address of the royalty token
+    /// @return isWhitelisted True if the royalty token is whitelisted
+    function isWhitelistedRoyaltyToken(address token) external view returns (bool) {
+        return _getRoyaltyModuleStorage().isWhitelistedRoyaltyToken[token];
+    }
+
+    /// @notice Indicates the royalty policy for a given IP asset
+    /// @param ipId The ID of IP asset
+    /// @return royaltyPolicy The address of the royalty policy
+    function royaltyPolicies(address ipId) external view returns (address) {
+        return _getRoyaltyModuleStorage().royaltyPolicies[ipId];
+    }
+
     /// @notice IERC165 interface support.
     function supportsInterface(bytes4 interfaceId) public view virtual override(BaseModule, IERC165) returns (bool) {
         return interfaceId == type(IRoyaltyModule).interfaceId || super.supportsInterface(interfaceId);
     }
+
+    /// @notice Hook that is called before any upgrade for authorization
+    /// @param newImplementation Address of the new implementation
+    function _authorizeUpgrade(address newImplementation) internal override onlyProtocolAdmin {}
+
+    /// @dev Returns the storage struct of RoyaltyModule.
+    function _getRoyaltyModuleStorage() private pure returns (RoyaltyModuleStorage storage $) {
+        assembly {
+            $.slot := RoyaltyModuleStorageLocation
+        }
+    }
 }