Skip to content

Commit

Permalink
Get leaves from memory in processMultiProofCalldata (#5140)
Browse files Browse the repository at this point in the history 
Signed-off-by: Hadrien Croubois <hadrien.croubois@gmail.com>
  • Loading branch information
@Amxx
Amxx committed Oct 2, 2024
1 parent a818284 commit de66e2c
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 37 deletions.
60 changes: 32 additions & 28 deletions contracts/utils/cryptography/MerkleProof.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,7 @@ library MerkleProof {
* This version handles proofs in calldata with the default hashing function.
*/
function verifyCalldata(bytes32[] calldata proof, bytes32 root, bytes32 leaf) internal pure returns (bool) {
return processProof(proof, leaf) == root;
return processProofCalldata(proof, leaf) == root;
}

/**
Expand Down Expand Up @@ -138,7 +138,7 @@ library MerkleProof {
bytes32 leaf,
function(bytes32, bytes32) view returns (bytes32) hasher
) internal view returns (bool) {
return processProof(proof, leaf, hasher) == root;
return processProofCalldata(proof, leaf, hasher) == root;
}

/**
Expand Down Expand Up @@ -200,15 +200,16 @@ library MerkleProof {
// `hashes` array. At the end of the process, the last hash in the `hashes` array should contain the root of
// the Merkle tree.
uint256 leavesLen = leaves.length;
uint256 proofFlagsLen = proofFlags.length;

// Check proof validity.
if (leavesLen + proof.length != proofFlags.length + 1) {
if (leavesLen + proof.length != proofFlagsLen + 1) {
revert MerkleProofInvalidMultiproof();
}

// The xxxPos values are "pointers" to the next value to consume in each array. All accesses are done using
// `xxx[xxxPos++]`, which return the current value and increment the pointer, thus mimicking a queue's "pop".
bytes32[] memory hashes = new bytes32[](proofFlags.length);
bytes32[] memory hashes = new bytes32[](proofFlagsLen);
uint256 leafPos = 0;
uint256 hashPos = 0;
uint256 proofPos = 0;
Expand All @@ -217,20 +218,20 @@ library MerkleProof {
// get the next hash.
// - depending on the flag, either another value from the "main queue" (merging branches) or an element from the
// `proof` array.
for (uint256 i = 0; i < proofFlags.length; i++) {
for (uint256 i = 0; i < proofFlagsLen; i++) {
bytes32 a = leafPos < leavesLen ? leaves[leafPos++] : hashes[hashPos++];
bytes32 b = proofFlags[i]
? (leafPos < leavesLen ? leaves[leafPos++] : hashes[hashPos++])
: proof[proofPos++];
hashes[i] = Hashes.commutativeKeccak256(a, b);
}

if (proofFlags.length > 0) {
if (proofFlagsLen > 0) {
if (proofPos != proof.length) {
revert MerkleProofInvalidMultiproof();
}
unchecked {
return hashes[proofFlags.length - 1];
return hashes[proofFlagsLen - 1];
}
} else if (leavesLen > 0) {
return leaves[0];
Expand Down Expand Up @@ -280,15 +281,16 @@ library MerkleProof {
// `hashes` array. At the end of the process, the last hash in the `hashes` array should contain the root of
// the Merkle tree.
uint256 leavesLen = leaves.length;
uint256 proofFlagsLen = proofFlags.length;

// Check proof validity.
if (leavesLen + proof.length != proofFlags.length + 1) {
if (leavesLen + proof.length != proofFlagsLen + 1) {
revert MerkleProofInvalidMultiproof();
}

// The xxxPos values are "pointers" to the next value to consume in each array. All accesses are done using
// `xxx[xxxPos++]`, which return the current value and increment the pointer, thus mimicking a queue's "pop".
bytes32[] memory hashes = new bytes32[](proofFlags.length);
bytes32[] memory hashes = new bytes32[](proofFlagsLen);
uint256 leafPos = 0;
uint256 hashPos = 0;
uint256 proofPos = 0;
Expand All @@ -297,20 +299,20 @@ library MerkleProof {
// get the next hash.
// - depending on the flag, either another value from the "main queue" (merging branches) or an element from the
// `proof` array.
for (uint256 i = 0; i < proofFlags.length; i++) {
for (uint256 i = 0; i < proofFlagsLen; i++) {
bytes32 a = leafPos < leavesLen ? leaves[leafPos++] : hashes[hashPos++];
bytes32 b = proofFlags[i]
? (leafPos < leavesLen ? leaves[leafPos++] : hashes[hashPos++])
: proof[proofPos++];
hashes[i] = hasher(a, b);
}

if (proofFlags.length > 0) {
if (proofFlagsLen > 0) {
if (proofPos != proof.length) {
revert MerkleProofInvalidMultiproof();
}
unchecked {
return hashes[proofFlags.length - 1];
return hashes[proofFlagsLen - 1];
}
} else if (leavesLen > 0) {
return leaves[0];
Expand All @@ -331,9 +333,9 @@ library MerkleProof {
bytes32[] calldata proof,
bool[] calldata proofFlags,
bytes32 root,
bytes32[] calldata leaves
bytes32[] memory leaves
) internal pure returns (bool) {
return processMultiProof(proof, proofFlags, leaves) == root;
return processMultiProofCalldata(proof, proofFlags, leaves) == root;
}

/**
Expand All @@ -351,22 +353,23 @@ library MerkleProof {
function processMultiProofCalldata(
bytes32[] calldata proof,
bool[] calldata proofFlags,
bytes32[] calldata leaves
bytes32[] memory leaves
) internal pure returns (bytes32 merkleRoot) {
// This function rebuilds the root hash by traversing the tree up from the leaves. The root is rebuilt by
// consuming and producing values on a queue. The queue starts with the `leaves` array, then goes onto the
// `hashes` array. At the end of the process, the last hash in the `hashes` array should contain the root of
// the Merkle tree.
uint256 leavesLen = leaves.length;
uint256 proofFlagsLen = proofFlags.length;

// Check proof validity.
if (leavesLen + proof.length != proofFlags.length + 1) {
if (leavesLen + proof.length != proofFlagsLen + 1) {
revert MerkleProofInvalidMultiproof();
}

// The xxxPos values are "pointers" to the next value to consume in each array. All accesses are done using
// `xxx[xxxPos++]`, which return the current value and increment the pointer, thus mimicking a queue's "pop".
bytes32[] memory hashes = new bytes32[](proofFlags.length);
bytes32[] memory hashes = new bytes32[](proofFlagsLen);
uint256 leafPos = 0;
uint256 hashPos = 0;
uint256 proofPos = 0;
Expand All @@ -375,20 +378,20 @@ library MerkleProof {
// get the next hash.
// - depending on the flag, either another value from the "main queue" (merging branches) or an element from the
// `proof` array.
for (uint256 i = 0; i < proofFlags.length; i++) {
for (uint256 i = 0; i < proofFlagsLen; i++) {
bytes32 a = leafPos < leavesLen ? leaves[leafPos++] : hashes[hashPos++];
bytes32 b = proofFlags[i]
? (leafPos < leavesLen ? leaves[leafPos++] : hashes[hashPos++])
: proof[proofPos++];
hashes[i] = Hashes.commutativeKeccak256(a, b);
}

if (proofFlags.length > 0) {
if (proofFlagsLen > 0) {
if (proofPos != proof.length) {
revert MerkleProofInvalidMultiproof();
}
unchecked {
return hashes[proofFlags.length - 1];
return hashes[proofFlagsLen - 1];
}
} else if (leavesLen > 0) {
return leaves[0];
Expand All @@ -409,10 +412,10 @@ library MerkleProof {
bytes32[] calldata proof,
bool[] calldata proofFlags,
bytes32 root,
bytes32[] calldata leaves,
bytes32[] memory leaves,
function(bytes32, bytes32) view returns (bytes32) hasher
) internal view returns (bool) {
return processMultiProof(proof, proofFlags, leaves, hasher) == root;
return processMultiProofCalldata(proof, proofFlags, leaves, hasher) == root;
}

/**
Expand All @@ -430,23 +433,24 @@ library MerkleProof {
function processMultiProofCalldata(
bytes32[] calldata proof,
bool[] calldata proofFlags,
bytes32[] calldata leaves,
bytes32[] memory leaves,
function(bytes32, bytes32) view returns (bytes32) hasher
) internal view returns (bytes32 merkleRoot) {
// This function rebuilds the root hash by traversing the tree up from the leaves. The root is rebuilt by
// consuming and producing values on a queue. The queue starts with the `leaves` array, then goes onto the
// `hashes` array. At the end of the process, the last hash in the `hashes` array should contain the root of
// the Merkle tree.
uint256 leavesLen = leaves.length;
uint256 proofFlagsLen = proofFlags.length;

// Check proof validity.
if (leavesLen + proof.length != proofFlags.length + 1) {
if (leavesLen + proof.length != proofFlagsLen + 1) {
revert MerkleProofInvalidMultiproof();
}

// The xxxPos values are "pointers" to the next value to consume in each array. All accesses are done using
// `xxx[xxxPos++]`, which return the current value and increment the pointer, thus mimicking a queue's "pop".
bytes32[] memory hashes = new bytes32[](proofFlags.length);
bytes32[] memory hashes = new bytes32[](proofFlagsLen);
uint256 leafPos = 0;
uint256 hashPos = 0;
uint256 proofPos = 0;
Expand All @@ -455,20 +459,20 @@ library MerkleProof {
// get the next hash.
// - depending on the flag, either another value from the "main queue" (merging branches) or an element from the
// `proof` array.
for (uint256 i = 0; i < proofFlags.length; i++) {
for (uint256 i = 0; i < proofFlagsLen; i++) {
bytes32 a = leafPos < leavesLen ? leaves[leafPos++] : hashes[hashPos++];
bytes32 b = proofFlags[i]
? (leafPos < leavesLen ? leaves[leafPos++] : hashes[hashPos++])
: proof[proofPos++];
hashes[i] = hasher(a, b);
}

if (proofFlags.length > 0) {
if (proofFlagsLen > 0) {
if (proofPos != proof.length) {
revert MerkleProofInvalidMultiproof();
}
unchecked {
return hashes[proofFlags.length - 1];
return hashes[proofFlagsLen - 1];
}
} else if (leavesLen > 0) {
return leaves[0];
Expand Down
19 changes: 10 additions & 9 deletions scripts/generate/templates/MerkleProof.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ function verify${suffix}(${(hash ? formatArgsMultiline : formatArgsSingleLine)(
'bytes32 leaf',
hash && `function(bytes32, bytes32) view returns (bytes32) ${hash}`,
)}) internal ${visibility} returns (bool) {
return processProof(proof, leaf${hash ? `, ${hash}` : ''}) == root;
return processProof${suffix}(proof, leaf${hash ? `, ${hash}` : ''}) == root;
}
/**
Expand Down Expand Up @@ -93,10 +93,10 @@ function multiProofVerify${suffix}(${formatArgsMultiline(
`bytes32[] ${location} proof`,
`bool[] ${location} proofFlags`,
'bytes32 root',
`bytes32[] ${location} leaves`,
`bytes32[] memory leaves`,
hash && `function(bytes32, bytes32) view returns (bytes32) ${hash}`,
)}) internal ${visibility} returns (bool) {
return processMultiProof(proof, proofFlags, leaves${hash ? `, ${hash}` : ''}) == root;
return processMultiProof${suffix}(proof, proofFlags, leaves${hash ? `, ${hash}` : ''}) == root;
}
/**
Expand All @@ -114,23 +114,24 @@ function multiProofVerify${suffix}(${formatArgsMultiline(
function processMultiProof${suffix}(${formatArgsMultiline(
`bytes32[] ${location} proof`,
`bool[] ${location} proofFlags`,
`bytes32[] ${location} leaves`,
`bytes32[] memory leaves`,
hash && `function(bytes32, bytes32) view returns (bytes32) ${hash}`,
)}) internal ${visibility} returns (bytes32 merkleRoot) {
// This function rebuilds the root hash by traversing the tree up from the leaves. The root is rebuilt by
// consuming and producing values on a queue. The queue starts with the \`leaves\` array, then goes onto the
// \`hashes\` array. At the end of the process, the last hash in the \`hashes\` array should contain the root of
// the Merkle tree.
uint256 leavesLen = leaves.length;
uint256 proofFlagsLen = proofFlags.length;
// Check proof validity.
if (leavesLen + proof.length != proofFlags.length + 1) {
if (leavesLen + proof.length != proofFlagsLen + 1) {
revert MerkleProofInvalidMultiproof();
}
// The xxxPos values are "pointers" to the next value to consume in each array. All accesses are done using
// \`xxx[xxxPos++]\`, which return the current value and increment the pointer, thus mimicking a queue's "pop".
bytes32[] memory hashes = new bytes32[](proofFlags.length);
bytes32[] memory hashes = new bytes32[](proofFlagsLen);
uint256 leafPos = 0;
uint256 hashPos = 0;
uint256 proofPos = 0;
Expand All @@ -139,20 +140,20 @@ function processMultiProof${suffix}(${formatArgsMultiline(
// get the next hash.
// - depending on the flag, either another value from the "main queue" (merging branches) or an element from the
// \`proof\` array.
for (uint256 i = 0; i < proofFlags.length; i++) {
for (uint256 i = 0; i < proofFlagsLen; i++) {
bytes32 a = leafPos < leavesLen ? leaves[leafPos++] : hashes[hashPos++];
bytes32 b = proofFlags[i]
? (leafPos < leavesLen ? leaves[leafPos++] : hashes[hashPos++])
: proof[proofPos++];
hashes[i] = ${hash ?? DEFAULT_HASH}(a, b);
}
if (proofFlags.length > 0) {
if (proofFlagsLen > 0) {
if (proofPos != proof.length) {
revert MerkleProofInvalidMultiproof();
}
unchecked {
return hashes[proofFlags.length - 1];
return hashes[proofFlagsLen - 1];
}
} else if (leavesLen > 0) {
return leaves[0];
Expand Down

0 comments on commit de66e2c

Please sign in to comment.