Skip to content

Commit

Permalink
add type getter to entity
Browse files Browse the repository at this point in the history
  • Loading branch information
1kresh committed Jul 10, 2024
1 parent b106839 commit 45d2302
Show file tree
Hide file tree
Showing 23 changed files with 190 additions and 85 deletions.
8 changes: 7 additions & 1 deletion src/contracts/common/Entity.sol
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,16 @@ abstract contract Entity is Initializable, IEntity {
*/
address public immutable FACTORY;

constructor(address factory) {
/**
* @inheritdoc IEntity
*/
uint64 public immutable TYPE;

constructor(address factory, uint64 type_) {
_disableInitializers();

FACTORY = factory;
TYPE = type_;
}

/**
Expand Down
12 changes: 6 additions & 6 deletions src/contracts/common/Factory.sol
Original file line number Diff line number Diff line change
Expand Up @@ -21,22 +21,22 @@ contract Factory is Registry, Ownable, IFactory {
/**
* @inheritdoc IFactory
*/
function totalImplementations() public view returns (uint64) {
function totalTypes() public view returns (uint64) {
return uint64(_whitelistedImplementations.length());
}

/**
* @inheritdoc IFactory
*/
function implementation(uint64 index) public view returns (address) {
return _whitelistedImplementations.at(index);
function implementation(uint64 type_) public view returns (address) {
return _whitelistedImplementations.at(type_);
}

/**
* @inheritdoc IFactory
*/
function whitelist(address implementation_) external onlyOwner {
if (IEntity(implementation_).FACTORY() != address(this)) {
if (IEntity(implementation_).FACTORY() != address(this) || IEntity(implementation_).TYPE() != totalTypes()) {
revert InvalidImplementation();
}
if (!_whitelistedImplementations.add(implementation_)) {
Expand All @@ -47,8 +47,8 @@ contract Factory is Registry, Ownable, IFactory {
/**
* @inheritdoc IFactory
*/
function create(uint64 index, bool withInitialize, bytes memory data) external returns (address entity_) {
entity_ = implementation(index).cloneDeterministic(keccak256(abi.encode(totalEntities(), index, data)));
function create(uint64 type_, bool withInitialize, bytes memory data) external returns (address entity_) {
entity_ = implementation(type_).cloneDeterministic(keccak256(abi.encode(totalEntities(), type_, data)));

_addEntity(entity_);

Expand Down
5 changes: 3 additions & 2 deletions src/contracts/delegator/BaseDelegator.sol
Original file line number Diff line number Diff line change
Expand Up @@ -67,8 +67,9 @@ contract BaseDelegator is Entity, AccessControlUpgradeable, IBaseDelegator {
address vaultFactory,
address operatorVaultOptInService,
address operatorNetworkOptInService,
address delegatorFactory
) Entity(delegatorFactory) {
address delegatorFactory,
uint64 entityType
) Entity(delegatorFactory, entityType) {
NETWORK_REGISTRY = networkRegistry;
VAULT_FACTORY = vaultFactory;
OPERATOR_VAULT_OPT_IN_SERVICE = operatorVaultOptInService;
Expand Down
6 changes: 4 additions & 2 deletions src/contracts/delegator/FullRestakeDelegator.sol
Original file line number Diff line number Diff line change
Expand Up @@ -38,14 +38,16 @@ contract FullRestakeDelegator is BaseDelegator, IFullRestakeDelegator {
address vaultFactory,
address operatorVaultOptInService,
address operatorNetworkOptInService,
address delegatorFactory
address delegatorFactory,
uint64 entityType
)
BaseDelegator(
networkRegistry,
vaultFactory,
operatorVaultOptInService,
operatorNetworkOptInService,
delegatorFactory
delegatorFactory,
entityType
)
{}

Expand Down
6 changes: 4 additions & 2 deletions src/contracts/delegator/NetworkRestakeDelegator.sol
Original file line number Diff line number Diff line change
Expand Up @@ -38,14 +38,16 @@ contract NetworkRestakeDelegator is BaseDelegator, INetworkRestakeDelegator {
address vaultFactory,
address operatorVaultOptInService,
address operatorNetworkOptInService,
address delegatorFactory
address delegatorFactory,
uint64 entityType
)
BaseDelegator(
networkRegistry,
vaultFactory,
operatorVaultOptInService,
operatorNetworkOptInService,
delegatorFactory
delegatorFactory,
entityType
)
{}

Expand Down
5 changes: 3 additions & 2 deletions src/contracts/slasher/BaseSlasher.sol
Original file line number Diff line number Diff line change
Expand Up @@ -64,8 +64,9 @@ abstract contract BaseSlasher is Entity, IBaseSlasher {
address networkVaultOptInService,
address operatorVaultOptInService,
address operatorNetworkOptInService,
address slasherFactory
) Entity(slasherFactory) {
address slasherFactory,
uint64 entityType
) Entity(slasherFactory, entityType) {
VAULT_FACTORY = vaultFactory;
NETWORK_MIDDLEWARE_SERVICE = networkMiddlewareService;
NETWORK_VAULT_OPT_IN_SERVICE = networkVaultOptInService;
Expand Down
6 changes: 4 additions & 2 deletions src/contracts/slasher/Slasher.sol
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,17 @@ contract Slasher is BaseSlasher, ISlasher {
address networkVaultOptInService,
address operatorVaultOptInService,
address operatorNetworkOptInService,
address slasherFactory
address slasherFactory,
uint64 entityType
)
BaseSlasher(
vaultFactory,
networkMiddlewareService,
networkVaultOptInService,
operatorVaultOptInService,
operatorNetworkOptInService,
slasherFactory
slasherFactory,
entityType
)
{}

Expand Down
6 changes: 4 additions & 2 deletions src/contracts/slasher/VetoSlasher.sol
Original file line number Diff line number Diff line change
Expand Up @@ -64,15 +64,17 @@ contract VetoSlasher is BaseSlasher, AccessControlUpgradeable, IVetoSlasher {
address operatorVaultOptInService,
address operatorNetworkOptInService,
address networkRegistry,
address slasherFactory
address slasherFactory,
uint64 entityType
)
BaseSlasher(
vaultFactory,
networkMiddlewareService,
networkVaultOptInService,
operatorVaultOptInService,
operatorNetworkOptInService,
slasherFactory
slasherFactory,
entityType
)
{
NETWORK_REGISTRY = networkRegistry;
Expand Down
6 changes: 6 additions & 0 deletions src/interfaces/common/IEntity.sol
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,12 @@ interface IEntity {
*/
function FACTORY() external view returns (address);

/**
* @notice Get the entity's type.
* @return type of the entity
*/
function TYPE() external view returns (uint64);

/**
* @notice Initialize this entity contract using a given data.
* @param data some data to use
Expand Down
18 changes: 9 additions & 9 deletions src/interfaces/common/IFactory.sol
Original file line number Diff line number Diff line change
Expand Up @@ -8,30 +8,30 @@ interface IFactory is IRegistry {
error InvalidImplementation();

/**
* @notice Get the total number of whitelisted implementations.
* @return total number of implementations
* @notice Get the total number of whitelisted types.
* @return total number of types
*/
function totalImplementations() external view returns (uint64);
function totalTypes() external view returns (uint64);

/**
* @notice Get the implementation for a given index.
* @param index position to get the implementation at
* @notice Get the implementation for a given type.
* @param type_ position to get the implementation at
* @return address of the implementation
*/
function implementation(uint64 index) external view returns (address);
function implementation(uint64 type_) external view returns (address);

/**
* @notice Whitelist a new implementation for entities.
* @notice Whitelist a new type of entities.
* @param implementation address of the new implementation
*/
function whitelist(address implementation) external;

/**
* @notice Create a new entity at the factory.
* @param index `index`-th implementation to use
* @param type_ type's implementation to use
* @param withInitialize whether to call initialize on the entity
* @param data initial data for the entity creation
* @return address of the entity
*/
function create(uint64 index, bool withInitialize, bytes memory data) external returns (address);
function create(uint64 type_, bool withInitialize, bytes memory data) external returns (address);
}
12 changes: 8 additions & 4 deletions test/DelegatorFactory.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,8 @@ contract DelegatorFactoryTest is Test {
address(vaultFactory),
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(delegatorFactory)
address(delegatorFactory),
delegatorFactory.totalTypes()
)
);
delegatorFactory.whitelist(networkRestakeDelegatorImpl);
Expand All @@ -87,7 +88,8 @@ contract DelegatorFactoryTest is Test {
address(vaultFactory),
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(delegatorFactory)
address(delegatorFactory),
delegatorFactory.totalTypes()
)
);
delegatorFactory.whitelist(fullRestakeDelegatorImpl);
Expand All @@ -99,7 +101,8 @@ contract DelegatorFactoryTest is Test {
address(networkVaultOptInService),
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(slasherFactory)
address(slasherFactory),
slasherFactory.totalTypes()
)
);
slasherFactory.whitelist(slasherImpl);
Expand All @@ -112,7 +115,8 @@ contract DelegatorFactoryTest is Test {
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(networkRegistry),
address(slasherFactory)
address(slasherFactory),
slasherFactory.totalTypes()
)
);
slasherFactory.whitelist(vetoSlasherImpl);
Expand Down
12 changes: 8 additions & 4 deletions test/SlasherFactory.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,8 @@ contract SlasherFactoryTest is Test {
address(vaultFactory),
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(delegatorFactory)
address(delegatorFactory),
delegatorFactory.totalTypes()
)
);
delegatorFactory.whitelist(networkRestakeDelegatorImpl);
Expand All @@ -88,7 +89,8 @@ contract SlasherFactoryTest is Test {
address(vaultFactory),
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(delegatorFactory)
address(delegatorFactory),
delegatorFactory.totalTypes()
)
);
delegatorFactory.whitelist(fullRestakeDelegatorImpl);
Expand All @@ -100,7 +102,8 @@ contract SlasherFactoryTest is Test {
address(networkVaultOptInService),
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(slasherFactory)
address(slasherFactory),
slasherFactory.totalTypes()
)
);
slasherFactory.whitelist(slasherImpl);
Expand All @@ -113,7 +116,8 @@ contract SlasherFactoryTest is Test {
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(networkRegistry),
address(slasherFactory)
address(slasherFactory),
slasherFactory.totalTypes()
)
);
slasherFactory.whitelist(vetoSlasherImpl);
Expand Down
12 changes: 8 additions & 4 deletions test/VaultConfigurator.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,8 @@ contract VaultConfiguratorTest is Test {
address(vaultFactory),
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(delegatorFactory)
address(delegatorFactory),
delegatorFactory.totalTypes()
)
);
delegatorFactory.whitelist(networkRestakeDelegatorImpl);
Expand All @@ -91,7 +92,8 @@ contract VaultConfiguratorTest is Test {
address(vaultFactory),
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(delegatorFactory)
address(delegatorFactory),
delegatorFactory.totalTypes()
)
);
delegatorFactory.whitelist(fullRestakeDelegatorImpl);
Expand All @@ -103,7 +105,8 @@ contract VaultConfiguratorTest is Test {
address(networkVaultOptInService),
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(slasherFactory)
address(slasherFactory),
slasherFactory.totalTypes()
)
);
slasherFactory.whitelist(slasherImpl);
Expand All @@ -116,7 +119,8 @@ contract VaultConfiguratorTest is Test {
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(networkRegistry),
address(slasherFactory)
address(slasherFactory),
slasherFactory.totalTypes()
)
);
slasherFactory.whitelist(vetoSlasherImpl);
Expand Down
12 changes: 8 additions & 4 deletions test/VaultFactory.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,8 @@ contract VaultFactoryTest is Test {
address(vaultFactory),
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(delegatorFactory)
address(delegatorFactory),
delegatorFactory.totalTypes()
)
);
delegatorFactory.whitelist(networkRestakeDelegatorImpl);
Expand All @@ -87,7 +88,8 @@ contract VaultFactoryTest is Test {
address(vaultFactory),
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(delegatorFactory)
address(delegatorFactory),
delegatorFactory.totalTypes()
)
);
delegatorFactory.whitelist(fullRestakeDelegatorImpl);
Expand All @@ -99,7 +101,8 @@ contract VaultFactoryTest is Test {
address(networkVaultOptInService),
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(slasherFactory)
address(slasherFactory),
slasherFactory.totalTypes()
)
);
slasherFactory.whitelist(slasherImpl);
Expand All @@ -112,7 +115,8 @@ contract VaultFactoryTest is Test {
address(operatorVaultOptInService),
address(operatorNetworkOptInService),
address(networkRegistry),
address(slasherFactory)
address(slasherFactory),
slasherFactory.totalTypes()
)
);
slasherFactory.whitelist(vetoSlasherImpl);
Expand Down
10 changes: 9 additions & 1 deletion test/common/Entity.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -28,12 +28,20 @@ contract EntityTest is Test {
}

function test_Create() public {
address impl = address(new SimpleEntity(address(factory)));
address impl = address(new SimpleEntity(address(factory), factory.totalTypes()));
assertEq(IEntity(impl).FACTORY(), address(factory));
factory.whitelist(impl);

address entity = factory.create(0, true, "");
assertEq(IEntity(entity).FACTORY(), address(factory));
assertEq(IEntity(entity).TYPE(), 0);

impl = address(new SimpleEntity(address(factory), factory.totalTypes()));
factory.whitelist(impl);

entity = factory.create(1, true, "");
assertEq(IEntity(entity).FACTORY(), address(factory));
assertEq(IEntity(entity).TYPE(), 1);

vm.expectRevert();
IEntity(entity).initialize("");
Expand Down
Loading

0 comments on commit 45d2302

Please sign in to comment.