From 93c591d5761923b04609e630d849ae109eae6b1e Mon Sep 17 00:00:00 2001
From: timofey <tim@ongoteam.yaconnect.com>
Date: Fri, 16 Jun 2023 13:51:00 +0000
Subject: [PATCH] `BinaryTreesTreeOf` for `BalancedTrees`

---
 src/flow/binary/balancing.rs | 25 ++++++++++++++++---------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/src/flow/binary/balancing.rs b/src/flow/binary/balancing.rs
index 6ec91d6..9acff06 100644
--- a/src/flow/binary/balancing.rs
+++ b/src/flow/binary/balancing.rs
@@ -1,13 +1,7 @@
 use super::*;
 
 pub trait BinaryTreesUnbalanced<'a>: BinaryTreesHeight<'a> {
-    fn tree_with_height(
-        &self,
-        tl: Self::Tree,
-        key: Self::Key,
-        tr: Self::Tree,
-        height: u64,
-    ) -> BTWrap<'a, Self, Self::Tree>;
+    fn tree_of_with_height(&self, node: Self::Node, height: u64) -> BTWrap<'a, Self, Self::Tree>;
 
     fn balancing_error<T: 'a>(&self, error: BalancingError) -> BTWrap<'a, Self, T>;
 
@@ -21,6 +15,11 @@ pub trait BinaryTreesUnbalanced<'a>: BinaryTreesHeight<'a> {
             Err(e) => self.balancing_error(e),
         }
     }
+
+    fn node_heights(&self, node: &Self::Node) -> (u64, u64) {
+        let (tl, tr, _) = self.split(node);
+        (self.height(&tl), self.height(&tr))
+    }
 }
 
 #[derive(Clone)]
@@ -89,8 +88,7 @@ impl<'a, BT: BinaryTreesUnbalanced<'a>> BinaryTrees<'a> for BalancedTrees<BT> {
         let hp = *hp;
         let ctx = self.0.clone();
         Self::bind(self.0.resolve(reference), move |node| {
-            let (tl, tr, _) = ctx.split(&node);
-            let (hl, hr) = (ctx.height(&tl), ctx.height(&tr));
+            let (hl, hr) = ctx.node_heights(&node);
             ctx.balancing_bind(matches_height(hl, hr, hp), |_| Self::pure(node))
         })
     }
@@ -103,3 +101,12 @@ impl<'a, BT: BinaryTreesUnbalanced<'a>> BinaryTrees<'a> for BalancedTrees<BT> {
         Some((self.0.refer(tree)?, self.0.height(tree)))
     }
 }
+
+impl<'a, BT: BinaryTreesUnbalanced<'a>> BinaryTreesTreeOf<'a> for BalancedTrees<BT> {
+    fn tree_of(&self, node: Self::Node) -> BTWrap<'a, Self, Self::Tree> {
+        let (hl, hr) = self.0.node_heights(&node);
+        self.0.balancing_bind(parent_height(hl, hr), |height| {
+            self.0.tree_of_with_height(node, height)
+        })
+    }
+}