Separated context for balancing

src/algo/add.rs

@@ -68,8 +68,11 @@ impl<T> Stack<T> {
         return self._get::<1>();
     }
 
-    fn getGrandGrandParent(&self) -> &T {
+    fn getGrandGrandParent(&self) -> Option<&T> {
-        return self._get::<2>();
+        match self.hasGrandGrandParent() {
+            false => None,
+            true => Some(self._get::<2>()),
+        }
     }
 
     fn pop(&mut self) -> T {
@@ -173,7 +176,23 @@ fn _add_recursive<Ctx: AddRecContext>(
     unsafe { balance(ctx, stack) }
 }
 
-unsafe fn balance<Ctx: AddRecContext>(ctx: &mut Ctx, stack: &mut Stack<Frame<Ctx::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>;
+}
+
+unsafe fn balance<Ctx: BalancingContext>(ctx: &mut Ctx, stack: &mut Stack<Frame<Ctx::NodeRef>>) {
     loop {
         if ctx.getColor(stack.getParent().node.clone()) == Color::BLACK {
             break;
@@ -183,10 +202,24 @@ unsafe fn balance<Ctx: AddRecContext>(ctx: &mut Ctx, stack: &mut Stack<Frame<Ctx
             break;
         }
 
-        let mut symCtx =
+        match stack.getGrandParent().dir {
-            DynamicDirectedBinaryTreeContextFromContext::wrap_mut(ctx, stack.getGrandParent().dir);
+            BinaryTreeDirection::LEFT => {
+                if _directed(&mut ctx.getLeftDirectedContext(), stack) {
+                    break;
+                }
+            }
+            BinaryTreeDirection::RIGHT => {
+                if _directed(&mut ctx.getRightDirectedContext(), stack) {
+                    break;
+                }
+            }
+        }
 
-        let uncle = symCtx.getOppositeChild(stack.getGrandParent().node.clone());
+        fn _directed<Ctx: DirectedBalancingContext>(
+            ctx: &mut Ctx,
+            stack: &mut Stack<Frame<Ctx::NodeRef>>,
+        ) -> bool {
+            let uncle = ctx.getOppositeChild(stack.getGrandParent().node.clone());
             if let Some(uncle) = uncle {
                 if ctx.getColor(uncle.clone()) == Color::RED {
                     ctx.setColor(stack.getParent().node.clone(), Color::BLACK);
@@ -194,25 +227,32 @@ unsafe fn balance<Ctx: AddRecContext>(ctx: &mut Ctx, stack: &mut Stack<Frame<Ctx
                     ctx.setColor(stack.getGrandParent().node.clone(), Color::RED);
 
                     if !stack.hasGrandGrandParent() {
-                    break;
+                        return true;
                     }
                     stack.popDiscard();
                     stack.popDiscard();
-                continue;
+                    return false;
                 }
             }
 
             if stack.getGrandParent().dir != stack.getParent().dir {
-            todo!("__rotateForward_Forward(parent, grandParent)");
+                ctx.rotateForward_Forward(
+                    stack.getGrandParent().node.clone(),
+                    stack.getParent().node.clone(),
+                );
                 stack.replaceTop(Frame {
-                node: stack.getParent().node,
+                    node: stack.getParent().node.clone(),
                     dir: stack.getGrandParent().dir,
                 })
             }
             ctx.setColor(stack.getParent().node.clone(), Color::BLACK);
             ctx.setColor(stack.getGrandParent().node.clone(), Color::RED);
-        todo!("__rotateOpposite_Switch(grandParent)");
+            ctx.rotateOpposite(
+                stack.getGrandGrandParent().map(|f| (f.node.clone(), f.dir)),
+                stack.getGrandParent().node.clone(),
+            );
 
-        break;
+            return true;
+        }
     }
 }