refactor: fix restrict logic in mock forest

Signed-off-by: MASDXI <sirawitt42@gmail.com>

Author: MASDXI

contracts/abstracts/ForestToken.sol

@@ -32,6 +32,30 @@ abstract contract ForestToken is ERC20, IForestERC20 {
         return _dag.getTx(tokenId);
     }
 
+    function transactionLevel(bytes32 tokenId) internal view returns (uint256) {
+        return _dag.getTxLevel(tokenId);
+    }
+
+    function transactionParent(bytes32 tokenId) internal view returns (bytes32) {
+        return _dag.getTxParent(tokenId);
+    }
+
+    function transactionRoot(bytes32 tokenId) internal view returns (bytes32) {
+        return _dag.getTxRoot(tokenId);
+    }
+
+    function transactionValue(bytes32 tokenId) internal view returns (uint256) {
+        return _dag.getTxValue(tokenId);
+    }
+
+    function transactionCount(address account) internal view returns (uint256) {
+        return _dag.getTxCount(account);
+    }
+
+    function transactionHierarchy(bytes32 tokenId) internal view returns (uint256) {
+        return _dag.getTxHierarchy(tokenId);
+    }
+
     /**
      * @dev Internal function to transfer tokens and manage transactions within the forest.
      * @param from The sender address.

contracts/abstracts/extensions/FreezeAddress.sol

@@ -9,21 +9,11 @@ pragma solidity >=0.8.0 <0.9.0;
 abstract contract FreezeAddress {
     mapping(address => bool) private _frozen;
 
-    /**
-     * @notice Error thrown when an operation is attempted on a frozen address.
-     */
+    /** errors */
     error AddressFrozen();
-
-    /**
-     * @notice Error thrown when an operation is attempted to unfreeze an address that is not frozen.
-     */
     error AddressNotFrozen();
 
-    /**
-     * @notice Event emitted when an address is frozen or unfrozen.
-     * @param account The address that was frozen or unfrozen.
-     * @param auth The boolean status indicating if the address is frozen (true) or unfrozen (false).
-     */
+    /** events */
     event FrozeAddress(address indexed account, bool indexed auth);
 
     modifier checkFrozenAddress(address from, address to) {

contracts/abstracts/extensions/FreezeBalance.sol

@@ -9,23 +9,11 @@ pragma solidity >=0.8.0 <0.9.0;
 abstract contract FreezeBalance {
     mapping(address => uint256) private _frozenBalance;
 
-    /**
-     * @notice Error thrown when an operation causes a balance overflow.
-     */
+    /** errors */
     error BalanceOverflow();
-
-    /**
-     * @notice Error thrown when a frozen balance restriction is violated.
-     * @param balance The total balance of the account.
-     * @param frozenBalance The frozen balance of the account.
-     */
     error BalanceFrozen(uint256 balance, uint256 frozenBalance);
 
-    /**
-     * @notice Event emitted when an account's balance is frozen or unfrozen.
-     * @param account The address of the account.
-     * @param value The amount of balance that is frozen.
-     */
+    /** events */
     event FrozenBalance(address indexed account, uint256 value);
 
     /**

contracts/abstracts/extensions/FreezeToken.sol

@@ -9,21 +9,11 @@ pragma solidity >=0.8.0 <0.9.0;
 abstract contract FreezeToken {
     mapping(bytes32 => bool) private _frozenToken;
 
-    /**
-     * @notice Error thrown when a token operation is attempted on a frozen token.
-     */
+    /** errors */
     error TokenFrozen();
-
-    /**
-     * @notice Error thrown when an operation is attempted to unfreeze a token that is not frozen.
-     */
     error TokenNotFrozen();
 
-    /**
-     * @notice Event emitted when a token is frozen or unfrozen.
-     * @param tokenId The identifier of the token.
-     * @param auth The status of the token, true if frozen, false if unfrozen.
-     */
+    /** events */
     event frozenToken(bytes32 indexed tokenId, bool indexed auth);
 
     /**

contracts/libraries/Forest.sol

@@ -68,7 +68,7 @@ library Forest {
     function createTx(DAG storage self, Tx memory newTx, address spender) internal {
         if (newTx.value == 0) revert TransactionZeroValue();
         bytes32 newId = calcTxHash(spender, self.nonces[spender]);
-        self.txs[newId] = newTx;
+        self.txs[newId] = Tx(newId, newTx.parent, newTx.value, newTx.level, newTx.owner);
         unchecked {
             self.nonces[spender]++;
         }

contracts/mocks/MockForest.sol

@@ -7,6 +7,17 @@ import "../abstracts/extensions/FreezeBalance.sol";
 import "../abstracts/extensions/FreezeToken.sol";
 
 contract MockForest is ForestToken, FreezeAddress, FreezeBalance, FreezeToken {
+    enum RESTRICT_TYPES { NULL, EQUAL, LESS, GREATER, BETWEEN }
+
+    struct Restrict {
+        RESTRICT_TYPES types;
+        bool enable;
+        uint256 start;
+        uint256 end;
+    }
+
+    mapping(bytes32 => Restrict) private _restricts;
+
     /// @custom:event for keep tracking token from root.
     event Transfer(address from, address to, bytes32 indexed root, bytes32 indexed parent, uint256 value);
 
@@ -20,30 +31,51 @@ contract MockForest is ForestToken, FreezeAddress, FreezeBalance, FreezeToken {
         _;
     }
 
-    // TODO
-    // modifier checkFrozenLevel(bytes32 tokenId) {
-        // check root equal check equal
-        // _;
-    // }
-
-    // TODO
-    // modifier checkFrozenBeforeLevel(bytes32 tokenId)  
-        // check root equal and check less than
-        // _;
-    // }
+    modifier checkFrozenLevel(bytes32 tokenId) {
+        Restrict memory restrict = getPartition(tokenId);
+        if (restrict.types == RESTRICT_TYPES.EQUAL && (restrict.enable)) {
+            if (transactionLevel(tokenId) == restrict.start) {
+                revert TokenFrozen();
+            }
+        }
+        _;
+    }
 
-    // TODO
-    // modifier checkFrozenAfterLevel(bytes32 tokenId)  
-        // check root equal check greater than
-        // _;
-    // }
+    modifier checkFrozenBeforeLevel(bytes32 tokenId) {
+        Restrict memory restrict = getPartition(tokenId);
+        if (restrict.types == RESTRICT_TYPES.LESS && (restrict.enable)) {
+            if (transactionLevel(tokenId) < restrict.start) {
+                revert TokenFrozen();
+            }
+        }
+        _;
+    }
+    
+    modifier checkFrozenAfterLevel(bytes32 tokenId) {
+        Restrict memory restrict = getPartition(tokenId);
+        uint256 txLevel = transactionLevel(tokenId);
+        if (restrict.types == RESTRICT_TYPES.GREATER && (restrict.enable)) {
+            if (transactionLevel(tokenId) > restrict.start && txLevel < restrict.end) {
+                revert TokenFrozen();
+            }
+        }
+        _;
+    }
 
-    // TODO
-    // modifier checkFrozenInBetweenLevel(bytes32 tokenId)
+    /** @dev restrict in partitioning style */
+    modifier checkFrozenInBetweenLevel(bytes32 tokenId) {
         // check root equal check greater than 'x' and less than 'y'
-        // _;
-    // }
-
+        Restrict memory restrict = getPartition(tokenId);
+        uint256 txLevel = transactionLevel(tokenId);
+        if (restrict.types == RESTRICT_TYPES.BETWEEN && (restrict.enable)) {
+            if (txLevel > restrict.start && txLevel < restrict.end) {
+                revert TokenFrozen();
+            }
+        }
+        _;
+    }
+    
+    /** @notice ERC20 Transfer also emit. */ 
     function _transfer(
         address from,
         address to,
@@ -58,7 +90,7 @@ contract MockForest is ForestToken, FreezeAddress, FreezeBalance, FreezeToken {
         checkFrozenRootOrParent(tokenId)
         checkFrozenToken(tokenId)
     {
-        /// @notice ERC20 Transfer also emit.
+        
         super._transfer(from, to, tokenId, value);
         Forest.Tx memory txn = _transaction(tokenId);
         emit Transfer(from, to, txn.root, txn.parent, value);
@@ -71,4 +103,20 @@ contract MockForest is ForestToken, FreezeAddress, FreezeBalance, FreezeToken {
     function burn(address account, bytes32 tokenId, uint256 value) public {
         _burnTransaction(account, tokenId, value);
     }
+
+    function setPartition(bytes32 tokenId, uint256 start, uint256 end, RESTRICT_TYPES restrict) public {
+        bytes32 rootTokenId = transactionRoot(tokenId);
+        _restricts[rootTokenId].types = restrict;
+        _restricts[rootTokenId].enable = true;
+        _restricts[rootTokenId].start = start;
+        _restricts[rootTokenId].end = end;
+    }
+
+    function clearPartition(bytes32 tokenId) public {
+        delete _restricts[transactionRoot(tokenId)];
+    }
+
+    function getPartition(bytes32 tokenId) public view returns (Restrict memory) {
+        return _restricts[transactionRoot(tokenId)];
+    }
 }

test/contracts/ERC20/ERC20.test.js

@@ -1,17 +1,14 @@
 const {loadFixture} = require("@nomicfoundation/hardhat-toolbox/network-helpers");
-const {network} = require("hardhat");
-const {anyValue} = require("@nomicfoundation/hardhat-chai-matchers/withArgs");
 const {expect} = require("chai");
-const {ZeroAddress} = require("ethers");
+const {amount, frozenAmount, tokenMetadata} = require("../../utils/constant");
 
+// skipping to test the ERC20 behavior because inherit from @openzeppelin/contracts
 describe("ERC20", function () {
-  const amount = 1000n;
-  const frozenAmount = 100n;
-
   async function deployTokenFixture() {
     const [owner, alice, bob, charlie, dave, otherAccount] = await ethers.getSigners();
     const contract = await ethers.getContractFactory("MockERC20");
-    const token = await contract.deploy("United States dollar", "USD");
+    const token = await contract.deploy(tokenMetadata.name, tokenMetadata.symbol);
+
     return {token, owner, alice, bob, charlie, dave, otherAccount};
   }
 
@@ -20,7 +17,7 @@ describe("ERC20", function () {
       const {token, alice, bob} = await loadFixture(deployTokenFixture);
       const aliceAddress = alice.address;
       const bobAddress = bob.address;
-      await token.mint(alice, amount);
+      await token.mint(aliceAddress, amount);
       await token.freezeAddress(aliceAddress);
       expect(await token.isFrozen(aliceAddress)).to.equal(true);
       await expect(token.connect(alice).transfer(bobAddress, amount)).to.be.reverted;
@@ -30,19 +27,18 @@ describe("ERC20", function () {
       const {token, owner, alice, bob} = await loadFixture(deployTokenFixture);
       const spenderAddress = owner.address;
       const aliceAddress = alice.address;
-      const bobAddress = bob.address;
-      await token.mint(alice, amount);
+      await token.mint(aliceAddress, amount);
       await token.connect(alice).approve(spenderAddress, amount);
       await token.freezeAddress(aliceAddress);
       expect(await token.isFrozen(aliceAddress)).to.equal(true);
-      await expect(token.connect(owner).transferFrom(aliceAddress, bobAddress, amount)).to.be.reverted;
+      await expect(token.connect(owner).transferFrom(aliceAddress, bob.address, amount)).to.be.reverted;
     });
 
     it("Freeze Alice Balance and transfer", async function () {
       const {token, alice, bob} = await loadFixture(deployTokenFixture);
       const aliceAddress = alice.address;
       const bobAddress = bob.address;
-      await token.mint(alice, amount);
+      await token.mint(aliceAddress, amount);
       await token.setFreezeBalance(aliceAddress, frozenAmount);
       expect(await token.getFrozenBalance(aliceAddress)).to.equal(frozenAmount);
       await expect(token.connect(alice).transfer(bobAddress, amount - frozenAmount)).not.to.be.reverted;

test/contracts/Forest/Forest.test.js

@@ -3,56 +3,56 @@ const {anyValue} = require("@nomicfoundation/hardhat-chai-matchers/withArgs");
 const {network} = require("hardhat");
 const {expect} = require("chai");
 const {encodeBytes32String, ZeroAddress, solidityPackedKeccak256, getBytes} = require("ethers");
+const {amount, freezeAmount, transferFrom, transfer, tokenMetadata} = require("../../utils/constant");
 
 describe("Forest", function () {
-  const transferMethod = "transfer(address,bytes32,uint256)";
-
   async function deployTokenFixture() {
     const [owner, otherAccount] = await ethers.getSigners();
     const contract = await ethers.getContractFactory("MockForest");
-    const token = await contract.deploy("United States dollar", "USD");
+    const token = await contract.deploy(tokenMetadata.name, tokenMetadata.symbol);
+
     return {token, owner, otherAccount};
   }
 
   describe("Scenarios", function () {
     it("Freeze Alice Account and transfer", async function () {
-      //  TODO
+      const {token, alice, bob} = await loadFixture(deployTokenFixture);
     });
 
     it("Freeze Alice Account and transferFrom", async function () {
-      //  TODO
+      const {token, alice, bob} = await loadFixture(deployTokenFixture);
     });
 
     it("Freeze Alice Balance and transfer", async function () {
-      //  TODO
+      const {token, alice, bob} = await loadFixture(deployTokenFixture);
     });
 
     it("Freeze Alice Balance and transferFrom", async function () {
-      //  TODO
+      const {token, alice, bob} = await loadFixture(deployTokenFixture);
     });
 
     it("Freeze Alice Token and transfer", async function () {
-      //  TODO
+      const {token, alice, bob} = await loadFixture(deployTokenFixture);
     });
 
     it("Freeze Alice Token and transferFrom", async function () {
-      //  TODO
+      const {token, alice, bob} = await loadFixture(deployTokenFixture);
     });
 
     it("Freeze at root and transfer", async function () {
-      //  TODO
+      const {token, alice, bob} = await loadFixture(deployTokenFixture);
     });
 
     it("Freeze at root and transferFrom", async function () {
-      //  TODO
+      const {token, alice, bob} = await loadFixture(deployTokenFixture);
     });
 
     it("Freeze at level and transfer", async function () {
-      //  TODO
+      const {token, alice, bob} = await loadFixture(deployTokenFixture);
     });
 
     it("Freeze at level and transferFrom", async function () {
-      //  TODO
+      const {token, alice, bob} = await loadFixture(deployTokenFixture);
     });
   });
 });