Skip to content

Commit

Permalink
clean up
Browse files Browse the repository at this point in the history
  • Loading branch information
Arvolear committed Jun 9, 2024
1 parent 6a83bd6 commit c0f9adf
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 180 deletions.
240 changes: 66 additions & 174 deletions contracts/mock/libs/data-structures/AvlTreeMock.sol
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.4;

import {TypeCaster} from "../../../libs/utils/TypeCaster.sol";

import {Traversal, AvlTree} from "../../../libs/data-structures/AvlTree.sol";

contract AvlTreeMock {
using TypeCaster for *;
using AvlTree for *;
using Traversal for *;

Expand Down Expand Up @@ -96,95 +99,30 @@ contract AvlTreeMock {
}

function traverseUint() external view returns (uint256[] memory, uint256[] memory) {
Traversal.Iterator memory iterator_ = _uintTree.first();

uint256[] memory keys_ = new uint256[](_uintTree.size());
uint256[] memory values_ = new uint256[](_uintTree.size());

uint256 index_ = 0;

while (iterator_.isValid()) {
(bytes32 bytesKey_, bytes32 bytesValue_) = iterator_.value();

keys_[index_] = uint256(bytesKey_);
values_[index_] = uint256(bytesValue_);

iterator_.next();

index_++;
}
(uint256[] memory keys_, bytes32[] memory values_) = _traverseAll(
_uintTree.first(),
_uintTree.size(),
true
);

return (keys_, values_);
return (keys_, values_.asUint256Array());
}

function traverseFirstThreeUint() external view returns (uint256[] memory, uint256[] memory) {
Traversal.Iterator memory iterator_ = _uintTree.first();

uint256[] memory keys_ = new uint256[](3);
uint256[] memory values_ = new uint256[](3);

uint256 index_ = 0;

while (iterator_.isValid()) {
(bytes32 bytesKey_, bytes32 bytesValue_) = iterator_.value();

keys_[index_] = uint256(bytesKey_);
values_[index_] = uint256(bytesValue_);

if (index_ == 2) {
break;
}

iterator_.next();

index_++;
}
function backwardsTraversalUint() external view returns (uint256[] memory, uint256[] memory) {
(uint256[] memory keys_, bytes32[] memory values_) = _traverseAll(
_uintTree.last(),
_uintTree.size(),
false
);

return (keys_, values_);
return (keys_, values_.asUint256Array());
}

function brokenTraversalUint() external view {
Traversal.Iterator memory iterator_ = _uintTree.first();

uint256[] memory keys_ = new uint256[](_uintTree.size());
uint256[] memory values_ = new uint256[](_uintTree.size());

uint256 index_ = 0;

while (iterator_.isValid()) {
iterator_.currentNode = 0;

(bytes32 bytesKey_, bytes32 bytesValue_) = iterator_.value();

keys_[index_] = uint256(bytesKey_);
values_[index_] = uint256(bytesValue_);

iterator_.next();

index_++;
}
}

function backwardsTraversalUint() external view returns (uint256[] memory, uint256[] memory) {
Traversal.Iterator memory iterator_ = _uintTree.last();

uint256[] memory keys_ = new uint256[](_uintTree.size());
uint256[] memory values_ = new uint256[](_uintTree.size());

uint256 index_ = 0;

while (iterator_.isValid()) {
(bytes32 bytesKey_, bytes32 bytesValue_) = iterator_.value();

keys_[index_] = uint256(bytesKey_);
values_[index_] = uint256(bytesValue_);

iterator_.prev();

index_++;
}

return (keys_, values_);
iterator_.currentNode = 0;
iterator_.next();
}

function backAndForthTraverseUint()
Expand Down Expand Up @@ -227,126 +165,56 @@ contract AvlTreeMock {
return (keys_, values_);
}

function nextOnLast() external view returns (uint256, uint256) {
Traversal.Iterator memory iterator_ = _uintTree.last();

(bytes32 key_, bytes32 value_) = iterator_.next();

return (uint256(key_), uint256(value_));
}

function prevOnFirst() external view returns (uint256, uint256) {
Traversal.Iterator memory iterator_ = _uintTree.first();

(bytes32 key_, bytes32 value_) = iterator_.prev();

return (uint256(key_), uint256(value_));
}

function traverseBytes32() external view returns (uint256[] memory, bytes32[] memory) {
Traversal.Iterator memory iterator_ = _bytes32Tree.first();

uint256[] memory keys_ = new uint256[](_bytes32Tree.size());
bytes32[] memory values_ = new bytes32[](_bytes32Tree.size());

uint256 index_ = 0;

while (iterator_.isValid()) {
(bytes32 bytesKey_, bytes32 bytesValue_) = iterator_.value();

keys_[index_] = uint256(bytesKey_);
values_[index_] = bytesValue_;

iterator_.next();

index_++;
}

return (keys_, values_);
return _traverseAll(_bytes32Tree.first(), _bytes32Tree.size(), true);
}

function backwardsTraversalBytes32()
external
view
returns (uint256[] memory, bytes32[] memory)
{
Traversal.Iterator memory iterator_ = _bytes32Tree.last();

uint256[] memory keys_ = new uint256[](_bytes32Tree.size());
bytes32[] memory values_ = new bytes32[](_bytes32Tree.size());

uint256 index_ = 0;

while (iterator_.isValid()) {
(bytes32 bytesKey_, bytes32 bytesValue_) = iterator_.value();

keys_[index_] = uint256(bytesKey_);
values_[index_] = bytesValue_;

iterator_.prev();

index_++;
}

return (keys_, values_);
return _traverseAll(_bytes32Tree.last(), _bytes32Tree.size(), false);
}

function traverseAddress() external view returns (uint256[] memory, address[] memory) {
Traversal.Iterator memory iterator_ = _addressTree.first();

uint256[] memory keys_ = new uint256[](_addressTree.size());
address[] memory values_ = new address[](_addressTree.size());

if (keys_.length != 0) {
(bytes32 bytesKey_, bytes32 bytesValue_) = iterator_.value();

keys_[0] = uint256(bytesKey_);
values_[0] = address(uint160(uint256(bytesValue_)));
}

uint256 index_ = 1;

while (iterator_.hasNext()) {
(bytes32 bytesKey_, bytes32 bytesValue_) = iterator_.next();

keys_[index_] = uint256(bytesKey_);
values_[index_] = address(uint160(uint256(bytesValue_)));

index_++;
}
(uint256[] memory keys_, bytes32[] memory values_) = _traverseAll(
_addressTree.first(),
_addressTree.size(),
true
);

return (keys_, values_);
return (keys_, values_.asAddressArray());
}

function backwardsTraversalAddress()
external
view
returns (uint256[] memory, address[] memory)
{
Traversal.Iterator memory iterator_ = _addressTree.last();
(uint256[] memory keys_, bytes32[] memory values_) = _traverseAll(
_addressTree.last(),
_addressTree.size(),
false
);

uint256[] memory keys_ = new uint256[](_addressTree.size());
address[] memory values_ = new address[](_addressTree.size());

if (keys_.length != 0) {
(bytes32 bytesKey_, bytes32 bytesValue_) = iterator_.value();
return (keys_, values_.asAddressArray());
}

keys_[0] = uint256(bytesKey_);
values_[0] = address(uint160(uint256(bytesValue_)));
}
function nextOnLast() external view returns (uint256, uint256) {
Traversal.Iterator memory iterator_ = _uintTree.last();

uint256 index_ = 1;
(bytes32 key_, bytes32 value_) = iterator_.next();

while (iterator_.hasPrev()) {
(bytes32 bytesKey_, bytes32 bytesValue_) = iterator_.prev();
return (uint256(key_), uint256(value_));
}

keys_[index_] = uint256(bytesKey_);
values_[index_] = address(uint160(uint256(bytesValue_)));
function prevOnFirst() external view returns (uint256, uint256) {
Traversal.Iterator memory iterator_ = _uintTree.first();

index_++;
}
(bytes32 key_, bytes32 value_) = iterator_.prev();

return (keys_, values_);
return (uint256(key_), uint256(value_));
}

function isCustomComparatorSetUint() external view returns (bool) {
Expand All @@ -361,6 +229,30 @@ contract AvlTreeMock {
return _addressTree.isCustomComparatorSet();
}

function _traverseAll(
Traversal.Iterator memory iterator_,
uint256 size_,
bool direction_
) private view returns (uint256[] memory, bytes32[] memory) {
uint256[] memory keys_ = new uint256[](size_);
bytes32[] memory values_ = new bytes32[](size_);

uint256 index_;

while (iterator_.isValid()) {
(bytes32 bytesKey_, bytes32 bytesValue_) = iterator_.value();

keys_[index_] = uint256(bytesKey_);
values_[index_] = bytesValue_;

direction_ ? iterator_.next() : iterator_.prev();

index_++;
}

return (keys_, values_);
}

function _descComparator(bytes32 key1_, bytes32 key2_) private pure returns (int256) {
if (key1_ > key2_) {
return -1;
Expand Down
11 changes: 5 additions & 6 deletions test/libs/data-structures/AvlTree.test.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,13 @@
import { ethers } from "hardhat";
import { encodeBytes32String } from "ethers";
import { expect } from "chai";
import { Reverter } from "@/test/helpers/reverter";

import { SignerWithAddress } from "@nomicfoundation/hardhat-ethers/signers";

import { AvlTreeMock } from "@ethers-v6";

import { ZERO_ADDR } from "@/scripts/utils/constants";
import { SignerWithAddress } from "@nomicfoundation/hardhat-ethers/signers";
import { encodeBytes32String } from "ethers";

describe("AvlTree", () => {
const reverter = new Reverter();
Expand Down Expand Up @@ -327,10 +330,6 @@ describe("AvlTree", () => {
expect(fullTraversal[0]).to.deep.equal([1, 2, 4, 6, 7]);
expect(fullTraversal[1]).to.deep.equal([12, 10, 11, 13, 14]);

let firstThreeTraversal = await avlTree.traverseFirstThreeUint();
expect(firstThreeTraversal[0]).to.deep.equal([1, 2, 4]);
expect(firstThreeTraversal[1]).to.deep.equal([12, 10, 11]);

let backwardsTraversal = await avlTree.backwardsTraversalUint();
expect(backwardsTraversal[0]).to.deep.equal([7, 6, 4, 2, 1]);
expect(backwardsTraversal[1]).to.deep.equal([14, 13, 11, 10, 12]);
Expand Down

0 comments on commit c0f9adf

Please sign in to comment.