From 41dd7491b20f51fd5bc741854ded1383eb4eb380 Mon Sep 17 00:00:00 2001
From: mllwchrry <mariia.zhvanko@gmail.com>
Date: Wed, 12 Jun 2024 11:36:58 +0300
Subject: [PATCH] fixed AvlTree test coverage

---
 .../mock/libs/data-structures/AvlTreeMock.sol | 42 +++++++++++++------
 1 file changed, 30 insertions(+), 12 deletions(-)

diff --git a/contracts/mock/libs/data-structures/AvlTreeMock.sol b/contracts/mock/libs/data-structures/AvlTreeMock.sol
index 5c39384d..469f45a9 100644
--- a/contracts/mock/libs/data-structures/AvlTreeMock.sol
+++ b/contracts/mock/libs/data-structures/AvlTreeMock.sol
@@ -178,13 +178,22 @@ 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);
+
+        (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 +201,22 @@ 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);
+
+        (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) {