From a897706a91028b8c7bad5bb706bc58c8b7d6b264 Mon Sep 17 00:00:00 2001
From: timofey <tim@ongoteam.yaconnect.com>
Date: Fri, 16 Jun 2023 13:31:58 +0000
Subject: [PATCH] `BalancedTrees`

---
 src/flow/binary/balancing.rs | 105 +++++++++++++++++++++++++++++++++++
 1 file changed, 105 insertions(+)

diff --git a/src/flow/binary/balancing.rs b/src/flow/binary/balancing.rs
index e69de29..fca81a6 100644
--- a/src/flow/binary/balancing.rs
+++ b/src/flow/binary/balancing.rs
@@ -0,0 +1,105 @@
+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 balancing_error<T: 'a>(&self, error: BalancingError) -> BTWrap<'a, Self, T>;
+
+    fn balancing_bind<A: 'a, B: 'a>(
+        &self,
+        ra: Result<A, BalancingError>,
+        f: impl FnOnce(A) -> BTWrap<'a, Self, B>,
+    ) -> BTWrap<'a, Self, B> {
+        match ra {
+            Ok(a) => f(a),
+            Err(e) => self.balancing_error(e),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct BalancedTrees<BT>(BT);
+
+impl<'a, BT: FunctorContext<'a>> FunctorContext<'a> for BalancedTrees<BT> {
+    type T = BT::T;
+}
+
+pub enum BalancingError {
+    LeafHeight(u64),
+    NodeHeight,
+    Balance(u64, u64),
+    HeightOverflow,
+    HeightMismatch { children: (u64, u64), parent: u64 },
+}
+
+fn balanced(hl: u64, hr: u64) -> Result<(), BalancingError> {
+    if hl.abs_diff(hr) > 1 {
+        Err(BalancingError::Balance(hl, hr))
+    } else {
+        Ok(())
+    }
+}
+
+fn parent_height(hl: u64, hr: u64) -> Result<u64, BalancingError> {
+    balanced(hl, hr)?;
+    std::cmp::max(hl, hr)
+        .checked_add(1)
+        .ok_or(BalancingError::HeightOverflow)
+}
+
+fn matches_height(hl: u64, hr: u64, hp: u64) -> Result<(), BalancingError> {
+    if parent_height(hl, hr)? == hp {
+        Ok(())
+    } else {
+        Err(BalancingError::HeightMismatch {
+            children: (hl, hr),
+            parent: hp,
+        })
+    }
+}
+
+impl<'a, BT: BinaryTreesUnbalanced<'a>> BinaryTrees<'a> for BalancedTrees<BT> {
+    type Node = BT::Node;
+
+    type Reference = BT::Reference;
+
+    type Tree = BT::Tree;
+
+    type Key = BT::Key;
+
+    type Comparator = BT::Comparator;
+
+    type _Tm = Self::T;
+
+    fn comparator(&self) -> &Self::Comparator {
+        self.0.comparator()
+    }
+
+    fn split(&self, node: &Self::Node) -> Split<'a, Self> {
+        self.0.split(node)
+    }
+
+    fn resolve(&self, reference: &Self::Reference) -> BTWrap<'a, Self, Self::Node> {
+        let hp = self.0.height_r(reference);
+        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));
+            ctx.balancing_bind(matches_height(hl, hr, hp), |_| Self::pure(node))
+        })
+    }
+
+    fn equal(&self, rl: &Self::Reference, rr: &Self::Reference) -> bool {
+        self.0.height_r(rl) == self.0.height_r(rr) && self.0.equal(rl, rr)
+    }
+
+    fn refer(&self, tree: &Self::Tree) -> Option<Self::Reference> {
+        self.0.refer(tree)
+    }
+}