Skip to content

Commit

Permalink
Merge pull request #93 from privacy-scaling-explorations/refactor/sol…
Browse files Browse the repository at this point in the history
…idity-errors

New Solidity error syntax (revert)
  • Loading branch information
0xjei authored Dec 4, 2023
2 parents 6df4769 + bced68b commit 4502ae5
Show file tree
Hide file tree
Showing 5 changed files with 120 additions and 71 deletions.
65 changes: 47 additions & 18 deletions packages/imt.sol/contracts/BinaryIMT.sol
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,15 @@ struct BinaryIMTData {
bool useDefaultZeroes;
}

error ValueGreaterThanSnarkScalarField();
error DepthNotSupported();
error WrongDefaultZeroIndex();
error TreeIsFull();
error NewLeafCannotEqualOldLeaf();
error LeafDoesNotExist();
error LeafIndexOutOfRange();
error WrongMerkleProofPath();

/// @title Incremental binary Merkle tree.
/// @dev The incremental tree allows to calculate the root hash each time a leaf is added, ensuring
/// the integrity of the tree.
Expand Down Expand Up @@ -91,16 +100,20 @@ library BinaryIMT {
if (index == 30) return Z_30;
if (index == 31) return Z_31;
if (index == 32) return Z_32;
revert("IncrementalBinaryTree: defaultZero bad index");

revert WrongDefaultZeroIndex();
}

/// @dev Initializes a tree.
/// @param self: Tree data.
/// @param depth: Depth of the tree.
/// @param zero: Zero value to be used.
function init(BinaryIMTData storage self, uint256 depth, uint256 zero) public {
require(zero < SNARK_SCALAR_FIELD, "BinaryIMT: leaf must be < SNARK_SCALAR_FIELD");
require(depth > 0 && depth <= MAX_DEPTH, "BinaryIMT: tree depth must be between 1 and 32");
if (zero >= SNARK_SCALAR_FIELD) {
revert ValueGreaterThanSnarkScalarField();
} else if (depth <= 0 || depth > MAX_DEPTH) {
revert DepthNotSupported();
}

self.depth = depth;

Expand All @@ -117,7 +130,9 @@ library BinaryIMT {
}

function initWithDefaultZeroes(BinaryIMTData storage self, uint256 depth) public {
require(depth > 0 && depth <= MAX_DEPTH, "BinaryIMT: tree depth must be between 1 and 32");
if (depth <= 0 || depth > MAX_DEPTH) {
revert DepthNotSupported();
}

self.depth = depth;
self.useDefaultZeroes = true;
Expand All @@ -131,8 +146,11 @@ library BinaryIMT {
function insert(BinaryIMTData storage self, uint256 leaf) public returns (uint256) {
uint256 depth = self.depth;

require(leaf < SNARK_SCALAR_FIELD, "BinaryIMT: leaf must be < SNARK_SCALAR_FIELD");
require(self.numberOfLeaves < 2 ** depth, "BinaryIMT: tree is full");
if (leaf >= SNARK_SCALAR_FIELD) {
revert ValueGreaterThanSnarkScalarField();
} else if (self.numberOfLeaves >= 2 ** depth) {
revert TreeIsFull();
}

uint256 index = self.numberOfLeaves;
uint256 hash = leaf;
Expand All @@ -155,6 +173,7 @@ library BinaryIMT {

self.root = hash;
self.numberOfLeaves += 1;

return hash;
}

Expand All @@ -171,9 +190,13 @@ library BinaryIMT {
uint256[] calldata proofSiblings,
uint8[] calldata proofPathIndices
) public {
require(newLeaf != leaf, "BinaryIMT: new leaf cannot be the same as the old one");
require(newLeaf < SNARK_SCALAR_FIELD, "BinaryIMT: new leaf must be < SNARK_SCALAR_FIELD");
require(verify(self, leaf, proofSiblings, proofPathIndices), "BinaryIMT: leaf is not part of the tree");
if (newLeaf == leaf) {
revert NewLeafCannotEqualOldLeaf();
} else if (newLeaf >= SNARK_SCALAR_FIELD) {
revert ValueGreaterThanSnarkScalarField();
} else if (!verify(self, leaf, proofSiblings, proofPathIndices)) {
revert LeafDoesNotExist();
}

uint256 depth = self.depth;
uint256 hash = newLeaf;
Expand All @@ -200,7 +223,10 @@ library BinaryIMT {
++i;
}
}
require(updateIndex < self.numberOfLeaves, "BinaryIMT: leaf index out of range");

if (updateIndex >= self.numberOfLeaves) {
revert LeafIndexOutOfRange();
}

self.root = hash;
}
Expand Down Expand Up @@ -231,19 +257,22 @@ library BinaryIMT {
uint256[] calldata proofSiblings,
uint8[] calldata proofPathIndices
) private view returns (bool) {
require(leaf < SNARK_SCALAR_FIELD, "BinaryIMT: leaf must be < SNARK_SCALAR_FIELD");
uint256 depth = self.depth;
require(
proofPathIndices.length == depth && proofSiblings.length == depth,
"BinaryIMT: length of path is not correct"
);

if (leaf >= SNARK_SCALAR_FIELD) {
revert ValueGreaterThanSnarkScalarField();
} else if (proofPathIndices.length != depth || proofSiblings.length != depth) {
revert WrongMerkleProofPath();
}

uint256 hash = leaf;

for (uint8 i = 0; i < depth; ) {
require(proofSiblings[i] < SNARK_SCALAR_FIELD, "BinaryIMT: sibling node must be < SNARK_SCALAR_FIELD");

require(proofPathIndices[i] == 1 || proofPathIndices[i] == 0, "BinaryIMT: path index is neither 0 nor 1");
if (proofSiblings[i] >= SNARK_SCALAR_FIELD) {
revert ValueGreaterThanSnarkScalarField();
} else if (proofPathIndices[i] != 1 && proofPathIndices[i] != 0) {
revert WrongMerkleProofPath();
}

if (proofPathIndices[i] == 0) {
hash = PoseidonT3.hash([hash, proofSiblings[i]]);
Expand Down
13 changes: 7 additions & 6 deletions packages/imt.sol/contracts/LeanIMT.sol
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,13 @@ error LeafCannotBeZero();
error LeafAlreadyExists();
error LeafDoesNotExist();

// The LeanIMT is an optimized version of the BinaryIMT.
// This implementation eliminates the use of zeroes, and make the tree depth dynamic.
// When a node doesn't have the right child, instead of using a zero hash as in the BinaryIMT,
// the node's value becomes that of its left child. Furthermore, rather than utilizing a static tree depth,
// it is updated based on the number of leaves in the tree. This approach
// results in the calculation of significantly fewer hashes, making the tree more efficient.
/// @title Lean Incremental binary Merkle tree.
/// @dev The LeanIMT is an optimized version of the BinaryIMT.
/// This implementation eliminates the use of zeroes, and make the tree depth dynamic.
/// When a node doesn't have the right child, instead of using a zero hash as in the BinaryIMT,
/// the node's value becomes that of its left child. Furthermore, rather than utilizing a static tree depth,
/// it is updated based on the number of leaves in the tree. This approach
/// results in the calculation of significantly fewer hashes, making the tree more efficient.
library LeanIMT {
uint256 public constant SNARK_SCALAR_FIELD =
21888242871839275222246405745257275088548364400416034343698204186575808495617;
Expand Down
69 changes: 44 additions & 25 deletions packages/imt.sol/contracts/QuinaryIMT.sol
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,14 @@ struct QuinaryIMTData {
mapping(uint256 => uint256[5]) lastSubtrees; // Caching these values is essential to efficient appends.
}

error ValueGreaterThanSnarkScalarField();
error DepthNotSupported();
error TreeIsFull();
error NewLeafCannotEqualOldLeaf();
error LeafDoesNotExist();
error LeafIndexOutOfRange();
error WrongMerkleProofPath();

/// @title Incremental quinary Merkle tree.
/// @dev The incremental tree allows to calculate the root hash each time a leaf is added, ensuring
/// the integrity of the tree.
Expand All @@ -27,8 +35,11 @@ library QuinaryIMT {
/// @param depth: Depth of the tree.
/// @param zero: Zero value to be used.
function init(QuinaryIMTData storage self, uint256 depth, uint256 zero) public {
require(zero < SNARK_SCALAR_FIELD, "QuinaryIMT: leaf must be < SNARK_SCALAR_FIELD");
require(depth > 0 && depth <= MAX_DEPTH, "QuinaryIMT: tree depth must be between 1 and 32");
if (zero >= SNARK_SCALAR_FIELD) {
revert ValueGreaterThanSnarkScalarField();
} else if (depth <= 0 || depth > MAX_DEPTH) {
revert DepthNotSupported();
}

self.depth = depth;

Expand Down Expand Up @@ -59,8 +70,11 @@ library QuinaryIMT {
function insert(QuinaryIMTData storage self, uint256 leaf) public {
uint256 depth = self.depth;

require(leaf < SNARK_SCALAR_FIELD, "QuinaryIMT: leaf must be < SNARK_SCALAR_FIELD");
require(self.numberOfLeaves < 5 ** depth, "QuinaryIMT: tree is full");
if (leaf >= SNARK_SCALAR_FIELD) {
revert ValueGreaterThanSnarkScalarField();
} else if (self.numberOfLeaves >= 5 ** depth) {
revert TreeIsFull();
}

uint256 index = self.numberOfLeaves;
uint256 hash = leaf;
Expand Down Expand Up @@ -104,9 +118,13 @@ library QuinaryIMT {
uint256[4][] calldata proofSiblings,
uint8[] calldata proofPathIndices
) public {
require(newLeaf != leaf, "QuinaryIMT: new leaf cannot be the same as the old one");
require(newLeaf < SNARK_SCALAR_FIELD, "QuinaryIMT: new leaf must be < SNARK_SCALAR_FIELD");
require(verify(self, leaf, proofSiblings, proofPathIndices), "QuinaryIMT: leaf is not part of the tree");
if (newLeaf == leaf) {
revert NewLeafCannotEqualOldLeaf();
} else if (newLeaf >= SNARK_SCALAR_FIELD) {
revert ValueGreaterThanSnarkScalarField();
} else if (!verify(self, leaf, proofSiblings, proofPathIndices)) {
revert LeafDoesNotExist();
}

uint256 depth = self.depth;
uint256 hash = newLeaf;
Expand Down Expand Up @@ -139,7 +157,10 @@ library QuinaryIMT {
++i;
}
}
require(updateIndex < self.numberOfLeaves, "QuinaryIMT: leaf index out of range");

if (updateIndex >= self.numberOfLeaves) {
revert LeafIndexOutOfRange();
}

self.root = hash;
}
Expand Down Expand Up @@ -170,38 +191,36 @@ library QuinaryIMT {
uint256[4][] calldata proofSiblings,
uint8[] calldata proofPathIndices
) private view returns (bool) {
require(leaf < SNARK_SCALAR_FIELD, "QuinaryIMT: leaf must be < SNARK_SCALAR_FIELD");
uint256 depth = self.depth;
require(
proofPathIndices.length == depth && proofSiblings.length == depth,
"QuinaryIMT: length of path is not correct"
);

if (leaf >= SNARK_SCALAR_FIELD) {
revert ValueGreaterThanSnarkScalarField();
} else if (proofPathIndices.length != depth || proofSiblings.length != depth) {
revert WrongMerkleProofPath();
}

uint256 hash = leaf;

for (uint8 i = 0; i < depth; ) {
uint256[5] memory nodes;

require(
proofPathIndices[i] >= 0 && proofPathIndices[i] < 5,
"QuinaryIMT: path index is not between 0 and 4"
);
if (proofPathIndices[i] < 0 || proofPathIndices[i] >= 5) {
revert WrongMerkleProofPath();
}

for (uint8 j = 0; j < 5; ) {
if (j < proofPathIndices[i]) {
require(
proofSiblings[i][j] < SNARK_SCALAR_FIELD,
"QuinaryIMT: sibling node must be < SNARK_SCALAR_FIELD"
);
if (proofSiblings[i][j] >= SNARK_SCALAR_FIELD) {
revert ValueGreaterThanSnarkScalarField();
}

nodes[j] = proofSiblings[i][j];
} else if (j == proofPathIndices[i]) {
nodes[j] = hash;
} else {
require(
proofSiblings[i][j - 1] < SNARK_SCALAR_FIELD,
"QuinaryIMT: sibling node must be < SNARK_SCALAR_FIELD"
);
if (proofSiblings[i][j - 1] >= SNARK_SCALAR_FIELD) {
revert ValueGreaterThanSnarkScalarField();
}

nodes[j] = proofSiblings[i][j - 1];
}
Expand Down
22 changes: 11 additions & 11 deletions packages/imt.sol/test/BinaryIMT.ts
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ describe("BinaryIMT", () => {
it("Should not create a tree with a depth > 32", async () => {
const transaction = binaryIMTTest.init(33)

await expect(transaction).to.be.revertedWith("BinaryIMT: tree depth must be between 1 and 32")
await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "DepthNotSupported")
})

it("Should create a tree", async () => {
Expand All @@ -49,7 +49,7 @@ describe("BinaryIMT", () => {

const transaction = binaryIMTTest.insert(leaf)

await expect(transaction).to.be.revertedWith("BinaryIMT: leaf must be < SNARK_SCALAR_FIELD")
await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "ValueGreaterThanSnarkScalarField")
})

it("Should insert a leaf in a tree", async () => {
Expand Down Expand Up @@ -115,7 +115,7 @@ describe("BinaryIMT", () => {

const transaction = binaryIMTTest.insert(3)

await expect(transaction).to.be.revertedWith("BinaryIMT: tree is full")
await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "TreeIsFull")
})
})

Expand All @@ -126,7 +126,7 @@ describe("BinaryIMT", () => {

const transaction = binaryIMTTest.update(1, 1, [0, 1], [0, 1])

await expect(transaction).to.be.revertedWith("BinaryIMT: new leaf cannot be the same as the old one")
await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "NewLeafCannotEqualOldLeaf")
})

it("Should not update a leaf if its new value is > SNARK_SCALAR_FIELD", async () => {
Expand All @@ -137,7 +137,7 @@ describe("BinaryIMT", () => {

const transaction = binaryIMTTest.update(1, newLeaf, [0, 1], [0, 1])

await expect(transaction).to.be.revertedWith("BinaryIMT: new leaf must be < SNARK_SCALAR_FIELD")
await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "ValueGreaterThanSnarkScalarField")
})

it("Should not update a leaf if its original value is > SNARK_SCALAR_FIELD", async () => {
Expand All @@ -148,7 +148,7 @@ describe("BinaryIMT", () => {

const transaction = binaryIMTTest.update(oldLeaf, 2, [0, 1], [0, 1])

await expect(transaction).to.be.revertedWith("BinaryIMT: leaf must be < SNARK_SCALAR_FIELD")
await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "ValueGreaterThanSnarkScalarField")
})

it("Should not update a leaf if the path indices are wrong", async () => {
Expand All @@ -169,7 +169,7 @@ describe("BinaryIMT", () => {
pathIndices
)

await expect(transaction).to.be.revertedWith("BinaryIMT: path index is neither 0 nor 1")
await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "WrongMerkleProofPath")
})

it("Should not update a leaf if the old leaf is wrong", async () => {
Expand All @@ -188,7 +188,7 @@ describe("BinaryIMT", () => {
pathIndices
)

await expect(transaction).to.be.revertedWith("BinaryIMT: leaf is not part of the tree")
await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "LeafDoesNotExist")
})

it("Should update a leaf", async () => {
Expand Down Expand Up @@ -251,7 +251,7 @@ describe("BinaryIMT", () => {
pathIndices
)

await expect(transaction).to.be.revertedWith("BinaryIMT: leaf index out of range")
await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "LeafIndexOutOfRange")
})
})

Expand All @@ -261,7 +261,7 @@ describe("BinaryIMT", () => {

const transaction = binaryIMTTest.remove(leaf, [0, 1], [0, 1])

await expect(transaction).to.be.revertedWith("BinaryIMT: leaf must be < SNARK_SCALAR_FIELD")
await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "ValueGreaterThanSnarkScalarField")
})

it("Should not remove a leaf that does not exist", async () => {
Expand All @@ -279,7 +279,7 @@ describe("BinaryIMT", () => {
pathIndices
)

await expect(transaction).to.be.revertedWith("BinaryIMT: leaf is not part of the tree")
await expect(transaction).to.be.revertedWithCustomError(binaryIMT, "LeafDoesNotExist")
})

it("Should remove a leaf", async () => {
Expand Down
Loading

0 comments on commit 4502ae5

Please sign in to comment.