From 1dd64086db81e58c7810fbbe191dce2c4f15bddb Mon Sep 17 00:00:00 2001
From: Andrew Golovashevich <landgrafhomyak@gmail.com>
Date: Mon, 29 Dec 2025 20:57:25 +0300
Subject: [PATCH] Redesigned balancing without explicit using stack

---
 src/algo/add.rs | 90 ++++++++++++++++++++++++++-----------------------
 1 file changed, 47 insertions(+), 43 deletions(-)

diff --git a/src/algo/add.rs b/src/algo/add.rs
index 0ddc882..073819b 100644
--- a/src/algo/add.rs
+++ b/src/algo/add.rs
@@ -173,43 +173,53 @@ fn _add_recursive<Ctx: AddRecContext>(
         }
     }
 
-    unsafe { balance(ctx, stack) }
+    unsafe { balance(ctx) }
+}
+
+trait BalancingContext: RedBlackTreeColorContext {
+    fn getLeftDirectedContext(&mut self) -> impl DirectedBalancingContext<NodeRef = Self::NodeRef>;
+    fn getRightDirectedContext(&mut self)
+    -> impl DirectedBalancingContext<NodeRef = Self::NodeRef>;
+
+    fn stackPopParentInfo(&self) -> Option<Frame<Self::NodeRef>>;
 }
 
 trait DirectedBalancingContext:
     RedBlackTreeColorContext + DirectedBinaryTreeChildrenGetterContext
 {
     fn rotateForward_Forward(&mut self, parent: Self::NodeRef, pivot: Self::NodeRef);
-    fn rotateOpposite(
-        &mut self,
-        parent: Option<(Self::NodeRef, BinaryTreeDirection)>,
-        pivot: Self::NodeRef,
-    );
-}
-trait BalancingContext: RedBlackTreeColorContext {
-    fn getLeftDirectedContext(&mut self) -> impl DirectedBalancingContext<NodeRef = Self::NodeRef>;
-    fn getRightDirectedContext(&mut self)
-    -> impl DirectedBalancingContext<NodeRef = Self::NodeRef>;
+    fn rotateOpposite_GrandParent(&mut self, pivot: Self::NodeRef);
 }
 
-unsafe fn balance<Ctx: BalancingContext>(ctx: &mut Ctx, stack: &mut Stack<Frame<Ctx::NodeRef>>) {
+fn balance<Ctx: BalancingContext>(ctx: &mut Ctx) {
     loop {
-        if ctx.getColor(stack.getParent().node.clone()) == Color::BLACK {
-            break;
+        let parent;
+        match ctx.stackPopParentInfo() {
+            None => return,
+            Some(_parent) => parent = _parent,
         }
-        if !stack.hasGrandParent() {
-            ctx.setColor(stack.getParent().node.clone(), Color::BLACK);
+
+        if ctx.getColor(parent.node.clone()) == Color::BLACK {
             break;
         }
 
-        match stack.getGrandParent().dir {
+        let grandparent;
+        match ctx.stackPopParentInfo() {
+            None => {
+                ctx.setColor(parent.node.clone(), Color::BLACK);
+                break;
+            }
+            Some(_grandparent) => grandparent = _grandparent,
+        }
+
+        match grandparent.dir {
             BinaryTreeDirection::LEFT => {
-                if _directed(&mut ctx.getLeftDirectedContext(), stack) {
+                if _directed(&mut ctx.getLeftDirectedContext(), grandparent, parent) {
                     break;
                 }
             }
             BinaryTreeDirection::RIGHT => {
-                if _directed(&mut ctx.getRightDirectedContext(), stack) {
+                if _directed(&mut ctx.getRightDirectedContext(), grandparent, parent) {
                     break;
                 }
             }
@@ -217,40 +227,34 @@ unsafe fn balance<Ctx: BalancingContext>(ctx: &mut Ctx, stack: &mut Stack<Frame<
 
         fn _directed<Ctx: DirectedBalancingContext>(
             ctx: &mut Ctx,
-            stack: &mut Stack<Frame<Ctx::NodeRef>>,
+            grandparent: Frame<Ctx::NodeRef>,
+            parent: Frame<Ctx::NodeRef>,
         ) -> bool {
-            let uncle = ctx.getOppositeChild(stack.getGrandParent().node.clone());
+            let uncle = ctx.getOppositeChild(grandparent.node.clone());
             if let Some(uncle) = uncle {
                 if ctx.getColor(uncle.clone()) == Color::RED {
-                    ctx.setColor(stack.getParent().node.clone(), Color::BLACK);
+                    ctx.setColor(parent.node.clone(), Color::BLACK);
                     ctx.setColor(uncle, Color::BLACK);
-                    ctx.setColor(stack.getGrandParent().node.clone(), Color::RED);
+                    ctx.setColor(grandparent.node.clone(), Color::RED);
 
-                    if !stack.hasGrandGrandParent() {
-                        return true;
-                    }
-                    stack.popDiscard();
-                    stack.popDiscard();
                     return false;
                 }
             }
 
-            if stack.getGrandParent().dir != stack.getParent().dir {
-                ctx.rotateForward_Forward(
-                    stack.getGrandParent().node.clone(),
-                    stack.getParent().node.clone(),
-                );
-                stack.replaceTop(Frame {
-                    node: stack.getParent().node.clone(),
-                    dir: stack.getGrandParent().dir,
-                })
+            let normalizedParent;
+            if grandparent.dir != parent.dir {
+                normalizedParent = Frame {
+                    node: ctx.getOppositeChild(parent.node.clone()).unwrap(),
+                    dir: grandparent.dir,
+                };
+                ctx.rotateForward_Forward(grandparent.node.clone(), parent.node.clone());
+                let _ = parent;
+            } else {
+                normalizedParent = parent;
             }
-            ctx.setColor(stack.getParent().node.clone(), Color::BLACK);
-            ctx.setColor(stack.getGrandParent().node.clone(), Color::RED);
-            ctx.rotateOpposite(
-                stack.getGrandGrandParent().map(|f| (f.node.clone(), f.dir)),
-                stack.getGrandParent().node.clone(),
-            );
+            ctx.setColor(normalizedParent.node.clone(), Color::BLACK);
+            ctx.setColor(grandparent.node.clone(), Color::RED);
+            ctx.rotateOpposite_GrandParent(grandparent.node.clone());
 
             return true;
         }