From 43ade1120b6eee18bc1b4e30127925b5b4c1337d Mon Sep 17 00:00:00 2001
From: Matej Focko <mfocko@redhat.com>
Date: Mon, 16 May 2022 18:36:31 +0200
Subject: [PATCH] fix: Ensure all affected nodes are being highlighted

Fixes #9

Signed-off-by: Matej Focko <mfocko@redhat.com>
---
 avl.js         | 20 +++++++++++++-------
 ranked_tree.js | 15 ++++++++-------
 ravl.js        |  1 +
 recorder.js    |  4 ++--
 wavl.js        | 14 ++++++++++----
 5 files changed, 34 insertions(+), 20 deletions(-)

diff --git a/avl.js b/avl.js
index 95344c9..eb78565 100644
--- a/avl.js
+++ b/avl.js
@@ -41,15 +41,16 @@ class AVLTree extends RankedTree {
             if (rotatingAroundRoot) {
                 this.root = newRoot;
             }
-            this.record("Fixing 0-child by single rotation: by z", z);
+            this.record("Fixing 0-child by single rotation: by z", z, newRoot);
 
             z.demote();
             this.record("Fixing 0-child by single rotation: demoting z", z);
         } else if (nodeDifference(y) == 1) {
-            rotateLeft(x);
+            let intermediateRoot = rotateLeft(x);
             this.record(
                 "Fixing 0-child by double-rotation: first rotation by x",
-                x
+                x,
+                intermediateRoot
             );
 
             newRoot = rotateRight(z);
@@ -58,7 +59,8 @@ class AVLTree extends RankedTree {
             }
             this.record(
                 "Fixing 0-child by double-rotation: second rotation by z",
-                z
+                z,
+                newRoot
             );
 
             y.promote();
@@ -129,8 +131,12 @@ class AVLTree extends RankedTree {
                 newRoot = rotateLeft(x);
                 break;
             default:
-                rotateRight(y);
-                this.record("AVL deletion, fixing by double-rotation: by y", y);
+                let intermediateRoot = rotateRight(y);
+                this.record(
+                    "AVL deletion, fixing by double-rotation: by y",
+                    y,
+                    intermediateRoot
+                );
 
                 newRoot = rotateLeft(x);
                 break;
@@ -138,7 +144,7 @@ class AVLTree extends RankedTree {
         if (rotatingAroundRoot) {
             this.root = newRoot;
         }
-        this.record("AVL deletion, fixing by rotation by x", x);
+        this.record("AVL deletion, fixing by rotation by x", x, newRoot);
 
         [newRoot.left, newRoot.right, newRoot]
             .filter((x) => x)
diff --git a/ranked_tree.js b/ranked_tree.js
index df2c5f7..6963715 100644
--- a/ranked_tree.js
+++ b/ranked_tree.js
@@ -4,7 +4,7 @@ class RankedTree {
         this.recorder = null;
     }
 
-    _getUnwrappedGraph(highlightedNode) {
+    _getUnwrappedGraph(...highlightedNodes) {
         let result = [];
 
         let q = new Queue();
@@ -18,8 +18,9 @@ class RankedTree {
             }
 
             let [value, rank] = [node.value, node.rank];
-            let highlight =
-                node === highlightedNode ? ', color="blue", penwidth=3' : "";
+            let highlight = highlightedNodes.includes(node)
+                ? ', color="blue", penwidth=3'
+                : "";
             result.push(
                 `\t"Node(${value})" [label="${value}, ${rank}"${highlight}];\n`
             );
@@ -46,10 +47,10 @@ class RankedTree {
         return result;
     }
 
-    toDot(highlightedNode) {
+    toDot(...highlightedNodes) {
         return [
             "digraph {\n",
-            ...this._getUnwrappedGraph(highlightedNode),
+            ...this._getUnwrappedGraph(...highlightedNodes),
             "}\n",
         ];
     }
@@ -58,11 +59,11 @@ class RankedTree {
         return this.toDot().join("\n");
     }
 
-    record(message, highlightedNode) {
+    record(message, ...highlightedNodes) {
         if (!this.recorder) {
             return;
         }
-        this.recorder.record(this, message, highlightedNode);
+        this.recorder.record(this, message, ...highlightedNodes);
     }
 
     rank() {
diff --git a/ravl.js b/ravl.js
index 8bd1e68..6be279e 100644
--- a/ravl.js
+++ b/ravl.js
@@ -16,6 +16,7 @@ class RAVLTree extends WAVLTree {
 
     deleteRebalance(node, parent) {
         // no-op
+        this.record("No rebalancing occurs");
         return;
     }
 }
diff --git a/recorder.js b/recorder.js
index 28e26a7..2c8585c 100644
--- a/recorder.js
+++ b/recorder.js
@@ -84,10 +84,10 @@ class Recorder {
         throw "not implemented!";
     }
 
-    record(tree, message, highlightedNode) {
+    record(tree, message, ...highlightedNodes) {
         // TODO: adjust join if necessary
         this.states.push(
-            new Record(tree.toDot(highlightedNode).join(""), message)
+            new Record(tree.toDot(...highlightedNodes).join(""), message)
         );
         this.render();
     }
diff --git a/wavl.js b/wavl.js
index 33a09f7..6e13911 100644
--- a/wavl.js
+++ b/wavl.js
@@ -32,13 +32,15 @@ class WAVLTree extends AVLTree {
         let wDiff = nodeDifference(w, y);
 
         if (wDiff == 1 && y.parent) {
+            let oldRoot = y.parent;
             newRoot = rotateLeft(y.parent);
             if (rotatingAroundRoot) {
                 this.root = newRoot;
             }
             this.record(
                 "Final step of deletion rebalance: single rotation by parent of y",
-                y.parent
+                oldRoot,
+                newRoot
             );
 
             y.promote();
@@ -55,19 +57,23 @@ class WAVLTree extends AVLTree {
                 );
             }
         } else if (wDiff == 2 && v && v.parent) {
-            rotateRight(v.parent);
+            let oldRoot = v.parent;
+            let intermediateRoot = rotateRight(v.parent);
             this.record(
                 "Final step of deletion rebalance: first of double rotation by parent of v",
-                v.parent
+                oldRoot,
+                intermediateRoot
             );
 
+            oldRoot = v.parent;
             newRoot = rotateLeft(v.parent);
             if (rotatingAroundRoot) {
                 this.root = newRoot;
             }
             this.record(
                 "Final step of deletion rebalance: second of double rotation by parent of v",
-                v.parent
+                oldRoot,
+                newRoot
             );
 
             v.promote().promote();