From 28a5ea1fa8c7a0af7362c73959283f1261f407f0 Mon Sep 17 00:00:00 2001
From: Mllw Chrry <mariia.zhvanko@gmail.com>
Date: Wed, 12 Jun 2024 02:33:21 -0700
Subject: [PATCH] fixed AvlTree test coverage (#105)

* fixed AvlTree test coverage

* fixed array out-of-bounds issue
---
 .../mock/libs/data-structures/AvlTreeMock.sol | 46 ++++++++++++++-----
 1 file changed, 34 insertions(+), 12 deletions(-)

diff --git a/contracts/mock/libs/data-structures/AvlTreeMock.sol b/contracts/mock/libs/data-structures/AvlTreeMock.sol
index 5c39384d..66378e13 100644
--- a/contracts/mock/libs/data-structures/AvlTreeMock.sol
+++ b/contracts/mock/libs/data-structures/AvlTreeMock.sol
@@ -178,13 +178,24 @@ contract AvlTreeMock {
     }
 
     function traverseAddress() external view returns (uint256[] memory, address[] memory) {
-        (uint256[] memory keys_, bytes32[] memory values_) = _traverseAll(
-            _addressTree.first(),
-            _addressTree.size(),
-            true
-        );
+        Traversal.Iterator memory iterator_ = _addressTree.first();
+
+        bytes32[] memory keys_ = new bytes32[](_addressTree.size());
+        bytes32[] memory values_ = new bytes32[](keys_.length);
+
+        if (keys_.length != 0) {
+            (keys_[0], values_[0]) = iterator_.value();
+        }
+
+        uint256 index_ = 1;
+
+        while (iterator_.hasNext()) {
+            (keys_[index_], values_[index_]) = iterator_.next();
+
+            index_++;
+        }
 
-        return (keys_, values_.asAddressArray());
+        return (keys_.asUint256Array(), values_.asAddressArray());
     }
 
     function backwardsTraversalAddress()
@@ -192,13 +203,24 @@ contract AvlTreeMock {
         view
         returns (uint256[] memory, address[] memory)
     {
-        (uint256[] memory keys_, bytes32[] memory values_) = _traverseAll(
-            _addressTree.last(),
-            _addressTree.size(),
-            false
-        );
+        Traversal.Iterator memory iterator_ = _addressTree.last();
+
+        bytes32[] memory keys_ = new bytes32[](_addressTree.size());
+        bytes32[] memory values_ = new bytes32[](keys_.length);
+
+        if (keys_.length != 0) {
+            (keys_[0], values_[0]) = iterator_.value();
+        }
+
+        uint256 index_ = 1;
+
+        while (iterator_.hasPrev()) {
+            (keys_[index_], values_[index_]) = iterator_.prev();
+
+            index_++;
+        }
 
-        return (keys_, values_.asAddressArray());
+        return (keys_.asUint256Array(), values_.asAddressArray());
     }
 
     function nextOnLast() external view returns (uint256, uint256) {