Skip to content

Commit

Permalink
Cleanup, more explicit naming
Browse files Browse the repository at this point in the history
  • Loading branch information
@kanshi
kanshi committed Feb 5, 2024
1 parent dfe5cd1 commit 21c0f76
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 83 deletions.
84 changes: 36 additions & 48 deletions contracts/Registrator.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,99 +20,87 @@ contract Registrator is Initializable, PausableUpgradeable, AccessControlUpgrade

IERC20 public tokenContract;

struct LockData {
uint256[] amount;
uint256[] unlockAt;
address[] unlockTo;
struct Data {
uint256 amount;
uint256 unlockAt;
address unlockTo;
}
struct Registration {
uint256 penalty;
Data[] data;
}

mapping(address => LockData) locks;
mapping(address => uint256) public penalties;
mapping(address => Registration) registrations;

event LockRegistered(address indexed _account);
event Registered(address indexed _account);

function getLock(address _address) public view returns (LockData memory) {
return locks[_address];
function getRegistration(address _address) public view returns (Registration memory) {
return registrations[_address];
}

function lock() external whenNotPaused {
function register(address _address) external whenNotPaused {
require(tokenContract.transferFrom(msg.sender, address(this), currentLockSize));
locks[msg.sender].amount.push(currentLockSize);
locks[msg.sender].unlockAt.push(block.number + lockBlocks);
locks[msg.sender].unlockTo.push(msg.sender);
emit LockRegistered(msg.sender);
}

function lockFor(address _address) external whenNotPaused {
require(tokenContract.transferFrom(msg.sender, address(this), currentLockSize));
locks[_address].amount.push(currentLockSize);
locks[_address].unlockAt.push(block.number + lockBlocks);
locks[_address].unlockTo.push(msg.sender);
emit LockRegistered(_address);

registrations[_address].data.push(Data(currentLockSize, block.number + lockBlocks, msg.sender));

emit Registered(_address);
}

function unlock(address _address, uint256 _upto) external whenNotPaused {
function unregister(address _address, uint256 _upto) external whenNotPaused {
require(_upto > 0, "UpTo parameter must be greater than 0");

uint256 requested = _upto;
uint256 unlocked = 0;
uint consumedSize = 0;
uint[] memory consumed = new uint[](locks[_address].amount.length);
for (uint i = 0; i < locks[_address].amount.length; i++) {
if ((locks[_address].unlockAt[i] < block.number) && (locks[_address].unlockTo[i] == msg.sender)) {
uint[] memory consumed = new uint[](registrations[_address].data.length);
for (uint i = 0; i < registrations[_address].data.length; i++) {

if ((registrations[_address].data[i].unlockAt < block.number) && (registrations[_address].data[i].unlockTo == msg.sender)) {
if (requested > 0) {
if (locks[_address].amount[i] > requested) {
locks[_address].amount[i] -= requested;
if (registrations[_address].data[i].amount > requested) {
registrations[_address].data[i].amount -= requested;
unlocked += requested;
requested = 0;
} else {
requested -= locks[_address].amount[i];
unlocked += locks[_address].amount[i];
requested -= registrations[_address].data[i].amount;
unlocked += registrations[_address].data[i].amount;

consumed[consumedSize] = i;
consumedSize++;
}
}
}

}

require(unlocked > 0, "No unlockables found");

require(tokenContract.balanceOf(address(this)) >= unlocked, "Insufficient contract token balance");
require(tokenContract.transfer(msg.sender, unlocked), "Token transfer failed");

require(locks[_address].amount.length == locks[_address].unlockAt.length, "Data consistency failure 1");
require(locks[_address].amount.length == locks[_address].unlockTo.length, "Data consistency failure 2");

uint removed = 0;
for (uint i = 0; i < consumedSize; i++) {
uint adjustedIndex = consumed[i] - removed;
require(adjustedIndex < locks[_address].amount.length, "Index out of bounds of amount");
require(adjustedIndex < locks[_address].unlockAt.length, "Index out of bounds of unlockAt");
require(adjustedIndex < locks[_address].unlockTo.length, "Index out of bounds of unlockTo");

for (uint j = adjustedIndex; j < locks[_address].amount.length - 1; j++) {
locks[_address].amount[j] = locks[_address].amount[j + 1];
locks[_address].unlockAt[j] = locks[_address].unlockAt[j + 1];
locks[_address].unlockTo[j] = locks[_address].unlockTo[j + 1];
require(adjustedIndex < registrations[_address].data.length, "Index out of bounds");

for (uint j = adjustedIndex; j < registrations[_address].data.length - 1; j++) {
registrations[_address].data[j] = registrations[_address].data[j + 1];
}
locks[_address].amount.pop();
locks[_address].unlockAt.pop();
locks[_address].unlockTo.pop();
registrations[_address].data.pop();
removed++;
}

if (locks[_address].amount.length == 0) {
delete locks[_address];
if (registrations[_address].data.length == 0) {
delete registrations[_address];
}
}

function setPenalty(address _account, uint256 _value)
function setPenalty(address _address, uint256 _value)
external
whenNotPaused
onlyRole(OPERATOR_ROLE)
{
penalties[_account] = _value;
registrations[_address].penalty = _value;
}

function setLockSize(uint256 _value)
Expand Down
70 changes: 35 additions & 35 deletions test/test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,13 +51,13 @@ describe("Registrator contract", function () {
// @ts-ignore
await token.connect(tester).approve(registratorAddress, lockAmount)
// @ts-ignore
const result = await registrator.connect(tester).lock()
const result = await registrator.connect(tester).register(tester.address)

expect(await token.balanceOf(registratorAddress)).to.equal(lockAmount)
expect(await token.balanceOf(tester.address)).to.equal(lockAmount)

const data = await registrator.getLock(tester.address)
expect(data[0][0]).to.equal(lockAmount)
const data = await registrator.getRegistration(tester.address)
expect(data.data[0].amount).to.equal(lockAmount)
})

it('Allows locking tokens for a specific address', async () => {
Expand All @@ -71,13 +71,13 @@ describe("Registrator contract", function () {
// @ts-ignore
await token.connect(tester).approve(registratorAddress, lockAmount)
// @ts-ignore
const result = await registrator.connect(tester).lockFor(receiver.address)
const result = await registrator.connect(tester).register(receiver.address)

expect(await token.balanceOf(registratorAddress)).to.equal(lockAmount)
expect(await token.balanceOf(tester.address)).to.equal(lockAmount)

const data = await registrator.getLock(receiver.address)
expect(data[0][0]).to.equal(lockAmount)
const data = await registrator.getRegistration(receiver.address)
expect(data.data[0].amount).to.equal(lockAmount)
})

it('Allows setting only non-zero lock lengths', async () => {
Expand Down Expand Up @@ -113,13 +113,13 @@ describe("Registrator contract", function () {
// @ts-ignore
await token.connect(tester).approve(registratorAddress, lockAmount)
// @ts-ignore
const result = await registrator.connect(tester).lock()
const result = await registrator.connect(tester).register(tester.address)

const data = await registrator.getLock(tester.address)
expect(data[0][0]).to.equal(lockAmount)
const data = await registrator.getRegistration(tester.address)
expect(data.data[0].amount).to.equal(lockAmount)

// @ts-ignore
await expect(registrator.connect(tester).unlock(tester.address, lockAmount)).to.be.revertedWith("No unlockables found")
await expect(registrator.connect(tester).unregister(tester.address, lockAmount)).to.be.revertedWith("No unlockables found")

})

Expand All @@ -134,15 +134,15 @@ describe("Registrator contract", function () {
// @ts-ignore
await token.connect(tester).approve(registratorAddress, lockAmount)
// @ts-ignore
const result = await registrator.connect(tester).lock()
const result = await registrator.connect(tester).register(tester.address)
expect(await token.balanceOf(tester.address)).to.equal(0)

for (let i = 0; i < defaultLockBlocks; i++) {
await network.provider.send("evm_mine")
}

// @ts-ignore
await registrator.connect(tester).unlock(tester.address, lockAmount)
await registrator.connect(tester).unregister(tester.address, lockAmount)
expect(await token.balanceOf(tester.address)).to.equal(lockAmount)
})

Expand All @@ -158,45 +158,45 @@ describe("Registrator contract", function () {
// @ts-ignore
await token.connect(tester).approve(registratorAddress, lockAmount)
// @ts-ignore
await registrator.connect(tester).lock()
await registrator.connect(tester).register(tester.address)
// @ts-ignore
await registrator.connect(operator).setLockBlocks(defaultLockBlocks * 2n)
// @ts-ignore
await token.connect(tester).approve(registratorAddress, lockAmount)
// @ts-ignore
await registrator.connect(tester).lock()
await registrator.connect(tester).register(tester.address)
// @ts-ignore
await registrator.connect(operator).setLockBlocks(defaultLockBlocks * 3n)
// @ts-ignore
await token.connect(tester).approve(registratorAddress, lockAmount)
// @ts-ignore
await registrator.connect(tester).lock()
await registrator.connect(tester).register(tester.address)
// @ts-ignore
await registrator.connect(operator).setLockBlocks(defaultLockBlocks * 4n)
// @ts-ignore
await token.connect(tester).approve(registratorAddress, lockAmount)
// @ts-ignore
await registrator.connect(tester).lock()
await registrator.connect(tester).register(tester.address)
// @ts-ignore
await registrator.connect(operator).setLockBlocks(defaultLockBlocks * 5n)
// @ts-ignore
await token.connect(tester).approve(registratorAddress, lockAmount)
// @ts-ignore
await registrator.connect(tester).lock()
await registrator.connect(tester).register(tester.address)


for (let i = 0; i < defaultLockBlocks * 3n; i++) {
await network.provider.send("evm_mine")
}

// @ts-ignore
await registrator.connect(tester).unlock(tester.address, lockAmount * 5n)
await registrator.connect(tester).unregister(tester.address, lockAmount * 5n)

expect(await token.balanceOf(tester.address)).to.equal(lockAmount * 3n)

const data = await registrator.getLock(tester.address)
const data = await registrator.getRegistration(tester.address)

expect(data[0][0]).to.equal(lockAmount)
expect(data.data[0].amount).to.equal(lockAmount)
})

it('Unlocks partial amount of tokens among multiple locks', async () => {
Expand All @@ -211,54 +211,54 @@ describe("Registrator contract", function () {
// @ts-ignore
await token.connect(tester).approve(registratorAddress, lockAmount)
// @ts-ignore
await registrator.connect(tester).lock()
await registrator.connect(tester).register(tester.address)
// @ts-ignore
await registrator.connect(operator).setLockBlocks(defaultLockBlocks * 2n)
// @ts-ignore
await token.connect(tester).approve(registratorAddress, lockAmount)
// @ts-ignore
await registrator.connect(tester).lock()
await registrator.connect(tester).register(tester.address)
// @ts-ignore
await registrator.connect(operator).setLockBlocks(defaultLockBlocks * 3n)
// @ts-ignore
await token.connect(tester).approve(registratorAddress, lockAmount)
// @ts-ignore
await registrator.connect(tester).lock()
await registrator.connect(tester).register(tester.address)

for (let i = 0; i < defaultLockBlocks * 2n; i++) {
await network.provider.send("evm_mine")
}

// @ts-ignore
await registrator.connect(tester).unlock(tester.address, lockAmount + (lockAmount / 2n))
await registrator.connect(tester).unregister(tester.address, lockAmount + (lockAmount / 2n))

expect(await token.balanceOf(tester.address)).to.equal(lockAmount + (lockAmount / 2n))

const data = await registrator.getLock(tester.address)
const data = await registrator.getRegistration(tester.address)

expect(data[0][0]).to.equal(lockAmount / 2n)
expect(data.data[0].amount).to.equal(lockAmount / 2n)
})

it('Allows setting penalties per address', async () => {
const { registrator, operator, tester } = await loadFixture(deploy)

expect(await registrator.penalties(tester.address)).to.equal(0)
expect((await registrator.getRegistration(tester.address)).penalty).to.equal(0)
// @ts-ignore
await registrator.connect(operator).setPenalty(tester.address, 100)

expect(await registrator.penalties(tester.address)).to.equal(100)
expect((await registrator.getRegistration(tester.address)).penalty).to.equal(100)
// @ts-ignore
await registrator.connect(operator).setPenalty(tester.address, 0)

expect(await registrator.penalties(tester.address)).to.equal(0)
expect((await registrator.getRegistration(tester.address)).penalty).to.equal(0)
})

it('Allows locking tokens after clearing data', async () => {
const { admin, registrator, tester, token, registratorAddress, operator } =
await loadFixture(deploy)

const preData = await registrator.getLock(tester.address)
expect(preData[0].length).to.equal(0)
const preData = await registrator.getRegistration(tester.address)
expect(preData.data.length).to.equal(0)

const lockAmount = defaultLockSize

Expand All @@ -268,18 +268,18 @@ describe("Registrator contract", function () {
// @ts-ignore
await token.connect(tester).approve(registratorAddress, lockAmount)
// @ts-ignore
await registrator.connect(tester).lock()
await registrator.connect(tester).register(tester.address)

for (let i = 0; i < defaultLockBlocks; i++) {
await network.provider.send("evm_mine")
}

// @ts-ignore
await registrator.connect(tester).unlock(tester.address, lockAmount)
await registrator.connect(tester).unregister(tester.address, lockAmount)

expect(await token.balanceOf(tester.address)).to.equal(lockAmount)

const postData = await registrator.getLock(tester.address)
expect(postData[0].length).to.equal(0)
const postData = await registrator.getRegistration(tester.address)
expect(postData.data.length).to.equal(0)
})
});

0 comments on commit 21c0f76

Please sign in to comment.