From d10e42c4bf3c1e739f699db8599cef7143a7ba4e Mon Sep 17 00:00:00 2001
From: williamzhu17 <williammzhu@gmail.com>
Date: Thu, 8 May 2025 17:36:35 -0700
Subject: [PATCH 1/6] added some tests

---
 tests/silimate/opt_balance_tree.ys | 153 +++++++++++++++++++++++++++++
 1 file changed, 153 insertions(+)
 create mode 100644 tests/silimate/opt_balance_tree.ys

diff --git a/tests/silimate/opt_balance_tree.ys b/tests/silimate/opt_balance_tree.ys
new file mode 100644
index 000000000..5d15f92ca
--- /dev/null
+++ b/tests/silimate/opt_balance_tree.ys
@@ -0,0 +1,153 @@
+log -header "Should not be turned into a tree"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a,
+  input wire b,
+  input wire c,
+  input wire d,
+  output wire x,
+  output wire y,
+  output wire z
+);
+  assign x = a & b;
+  assign y = x & c;
+  assign z = y & d;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_expand
+# equiv_opt -assert opt_balance_tree
+opt_balance_tree
+
+# TODO check design to make sure that it is not turned into a tree
+# design -load postopt
+
+design -reset
+log -pop
+
+log -header "With a cell"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a,
+  input wire b,
+  input wire c,
+  input wire d,
+  output wire x,
+  output wire y,
+  output wire z
+);
+  wire temp;
+  assign y = ~temp;
+  assign x = a & b;
+  assign temp = x & c;
+  assign z = temp & d;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_expand
+# equiv_opt -assert opt_balance_tree
+opt_balance_tree
+
+# TODO check design to make sure that it is not turned into a tree
+# Note: this one already worked, just including here for completeness
+# design -load postopt
+
+design -reset
+log -pop
+
+log -header "Word out port"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a,
+  input wire b,
+  input wire c,
+  input wire d,
+  output wire [2:0] x
+);
+  assign x[0] = a & b;
+  assign x[1] = x[0] & c;
+  assign x[2] = x[1] & d;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_expand
+# equiv_opt -assert opt_balance_tree
+opt_balance_tree
+
+# TODO check design to make sure that it is not turned into a tree
+# Note: this one already worked, just including here for completeness
+# design -load postopt
+
+design -reset
+log -pop
+
+log -header "Fanout going to multiple outputs"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a,
+  input wire b,
+  input wire c,
+  input wire d,
+  output wire [2:0] x,
+  output wire y
+);
+  assign x[0] = a & b;
+  assign x[1] = x[0] & c;
+  assign x[2] = x[1] & d;
+
+  assign y = x[1];
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_expand
+# equiv_opt -assert opt_balance_tree
+opt_balance_tree
+
+# TODO check design to make sure that it is not turned into a tree
+# Note: this one already worked, just including here for completeness
+# design -load postopt
+
+design -reset
+log -pop
+
+log -header "Fanout going to multiple of the same word"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a,
+  input wire b,
+  input wire c,
+  input wire d,
+  output wire [3:0] x,
+);
+  assign x[0] = a & b;
+  assign x[1] = x[0] & c;
+  assign x[2] = x[1] & d;
+  assign x[3] = x[1];
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_expand
+# equiv_opt -assert opt_balance_tree
+opt_balance_tree
+
+# TODO check design to make sure that it is not turned into a tree
+# Note: this one already worked, just including here for completeness
+# design -load postopt
+
+design -reset
+log -pop
\ No newline at end of file

From 8d20ed06372e21764f2f2c519ca43541cdda6937 Mon Sep 17 00:00:00 2001
From: williamzhu17 <williammzhu@gmail.com>
Date: Thu, 8 May 2025 17:38:42 -0700
Subject: [PATCH 2/6] updated tests a bit

---
 tests/silimate/opt_balance_tree.ys | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/tests/silimate/opt_balance_tree.ys b/tests/silimate/opt_balance_tree.ys
index 5d15f92ca..50795226a 100644
--- a/tests/silimate/opt_balance_tree.ys
+++ b/tests/silimate/opt_balance_tree.ys
@@ -19,8 +19,7 @@ EOF
 check -assert
 
 # Check equivalence after opt_expand
-# equiv_opt -assert opt_balance_tree
-opt_balance_tree
+equiv_opt -assert opt_balance_tree
 
 # TODO check design to make sure that it is not turned into a tree
 # design -load postopt
@@ -51,8 +50,7 @@ EOF
 check -assert
 
 # Check equivalence after opt_expand
-# equiv_opt -assert opt_balance_tree
-opt_balance_tree
+equiv_opt -assert opt_balance_tree
 
 # TODO check design to make sure that it is not turned into a tree
 # Note: this one already worked, just including here for completeness
@@ -80,8 +78,7 @@ EOF
 check -assert
 
 # Check equivalence after opt_expand
-# equiv_opt -assert opt_balance_tree
-opt_balance_tree
+equiv_opt -assert opt_balance_tree
 
 # TODO check design to make sure that it is not turned into a tree
 # Note: this one already worked, just including here for completeness
@@ -112,8 +109,7 @@ EOF
 check -assert
 
 # Check equivalence after opt_expand
-# equiv_opt -assert opt_balance_tree
-opt_balance_tree
+equiv_opt -assert opt_balance_tree
 
 # TODO check design to make sure that it is not turned into a tree
 # Note: this one already worked, just including here for completeness
@@ -142,8 +138,7 @@ EOF
 check -assert
 
 # Check equivalence after opt_expand
-# equiv_opt -assert opt_balance_tree
-opt_balance_tree
+equiv_opt -assert opt_balance_tree
 
 # TODO check design to make sure that it is not turned into a tree
 # Note: this one already worked, just including here for completeness

From b265ea9dcfd549c762124bd718a6b3938030f2a7 Mon Sep 17 00:00:00 2001
From: williamzhu17 <williammzhu@gmail.com>
Date: Thu, 8 May 2025 17:39:16 -0700
Subject: [PATCH 3/6] removed comma

---
 tests/silimate/opt_balance_tree.ys | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/silimate/opt_balance_tree.ys b/tests/silimate/opt_balance_tree.ys
index 50795226a..ed683cb9b 100644
--- a/tests/silimate/opt_balance_tree.ys
+++ b/tests/silimate/opt_balance_tree.ys
@@ -127,7 +127,7 @@ module top (
   input wire b,
   input wire c,
   input wire d,
-  output wire [3:0] x,
+  output wire [3:0] x
 );
   assign x[0] = a & b;
   assign x[1] = x[0] & c;

From 51a951d6147cb038b627566d4c8824af49fd0022 Mon Sep 17 00:00:00 2001
From: williamzhu17 <williammzhu@gmail.com>
Date: Sun, 11 May 2025 10:17:08 -0700
Subject: [PATCH 4/6] wip tests

---
 tests/silimate/opt_balance_tree.ys | 82 +++++++++++++++++++++++-------
 1 file changed, 63 insertions(+), 19 deletions(-)

diff --git a/tests/silimate/opt_balance_tree.ys b/tests/silimate/opt_balance_tree.ys
index ed683cb9b..8427d25cb 100644
--- a/tests/silimate/opt_balance_tree.ys
+++ b/tests/silimate/opt_balance_tree.ys
@@ -1,3 +1,44 @@
+log -header "Should be turned into a tree"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a,
+  input wire b,
+  input wire c,
+  input wire d,
+  output wire x,
+);
+  assign x = a & b & c & d;
+endmodule
+EOF
+check -assert
+
+autoname
+write_json pre.json
+exec -- netlistsvg pre.json -o pre.svg
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+
+design -load postopt
+# opt_balance_tree
+
+autoname
+write_json post.json
+exec -- netlistsvg post.json -o post.svg
+
+# Checks if y is still wired up to the correct gate
+# select -set y_wires o:y %ci
+# select -set y_driver @y_wires %ci
+# select -set and_y_cell t:$and @y_driver %i
+# select @and_y_cell -assert-count 1
+# select -set inputs @and_y_cell %ci
+# select -assert-count 1 @inputs i:c %i
+
+design -reset
+log -pop
+
 log -header "Should not be turned into a tree"
 log -push
 design -reset
@@ -18,11 +59,18 @@ endmodule
 EOF
 check -assert
 
-# Check equivalence after opt_expand
+# Check equivalence after opt_balance_tree
 equiv_opt -assert opt_balance_tree
 
-# TODO check design to make sure that it is not turned into a tree
-# design -load postopt
+design -load postopt
+
+# Checks if y is still wired up to the correct gate
+select -set y_wires o:y %ci
+select -set y_driver @y_wires %ci
+select -set and_y_cell t:$and @y_driver %i
+select @and_y_cell -assert-count 1
+select -set inputs @and_y_cell %ci
+select -assert-count 1 @inputs i:c %i
 
 design -reset
 log -pop
@@ -49,12 +97,11 @@ endmodule
 EOF
 check -assert
 
-# Check equivalence after opt_expand
+# Check equivalence after opt_balance_tree
 equiv_opt -assert opt_balance_tree
 
-# TODO check design to make sure that it is not turned into a tree
-# Note: this one already worked, just including here for completeness
-# design -load postopt
+design -load postopt
+select -assert-count 3 t:$and
 
 design -reset
 log -pop
@@ -77,12 +124,11 @@ endmodule
 EOF
 check -assert
 
-# Check equivalence after opt_expand
+# Check equivalence after opt_balance_tree
 equiv_opt -assert opt_balance_tree
 
-# TODO check design to make sure that it is not turned into a tree
-# Note: this one already worked, just including here for completeness
-# design -load postopt
+design -load postopt
+select -assert-count 3 t:$and
 
 design -reset
 log -pop
@@ -108,12 +154,11 @@ endmodule
 EOF
 check -assert
 
-# Check equivalence after opt_expand
+# Check equivalence after opt_balance_tree
 equiv_opt -assert opt_balance_tree
 
-# TODO check design to make sure that it is not turned into a tree
-# Note: this one already worked, just including here for completeness
-# design -load postopt
+design -load postopt
+select -assert-count 3 t:$and
 
 design -reset
 log -pop
@@ -137,12 +182,11 @@ endmodule
 EOF
 check -assert
 
-# Check equivalence after opt_expand
+# Check equivalence after opt_balance_tree
 equiv_opt -assert opt_balance_tree
 
-# TODO check design to make sure that it is not turned into a tree
-# Note: this one already worked, just including here for completeness
-# design -load postopt
+design -load postopt
+select -assert-count 3 t:$and
 
 design -reset
 log -pop
\ No newline at end of file

From 39be4f29bde53331c60d669c1f3d6ae73aa2b38a Mon Sep 17 00:00:00 2001
From: williamzhu17 <williammzhu@gmail.com>
Date: Sun, 11 May 2025 10:58:29 -0700
Subject: [PATCH 5/6] opt_balance_tree test cases

---
 passes/silimate/opt_balance_tree.cc |  25 ++++++-
 tests/silimate/opt_balance_tree.ys  | 107 ++++++++++++++++++++++++----
 2 files changed, 114 insertions(+), 18 deletions(-)

diff --git a/passes/silimate/opt_balance_tree.cc b/passes/silimate/opt_balance_tree.cc
index da308e582..b575b9172 100644
--- a/passes/silimate/opt_balance_tree.cc
+++ b/passes/silimate/opt_balance_tree.cc
@@ -202,10 +202,10 @@ struct OptBalanceTreeWorker {
 		cell->fixup_parameters();
 	}
 
-	void process_chain(vector<Cell*> &chain) {
+	bool process_chain(vector<Cell*> &chain) {
 		// If chain size is less than 3, no balancing needed
 		if (GetSize(chain) < 3)
-			return;
+			return false;
 
 		// Get mid, midnext (at index mid+1) and end of chain
 		Cell *mid_cell = chain[GetSize(chain) / 2];
@@ -262,6 +262,8 @@ struct OptBalanceTreeWorker {
 
 		// Width reduce mid cell
 		wreduce(mid_cell);
+
+		return true;
 	}
 
 	void cleanup() {
@@ -291,7 +293,24 @@ struct OptBalanceTreeWorker {
 			// For each chain, if len >= 3, convert to tree via "rotation" and recurse on subtrees
 			for (auto c : chain_start_cells) {
 				vector<Cell*> chain = create_chain(c);
-				process_chain(chain);
+				bool processed = process_chain(chain);
+
+				if (processed) {
+					// Rename cells and wires for formal check to pass as cells signals have changed functionalities post rotation
+					for (Cell *cell : chain) {
+						module->rename(cell, NEW_ID2_SUFFIX("rot_cell"));
+					}
+					for (Cell *cell : chain) {
+						SigSpec y_sig = sigmap(cell->getPort(ID::Y));
+						if (y_sig.is_wire()) {
+							Wire *wire = y_sig.as_wire();
+							if (wire && !wire->port_input && !wire->port_output) {
+								module->rename(y_sig.as_wire(), NEW_ID2_SUFFIX("rot_wire"));
+							}
+						}
+					}
+				}
+
 				cell_count[cell_type] += GetSize(chain);
 			}
 
diff --git a/tests/silimate/opt_balance_tree.ys b/tests/silimate/opt_balance_tree.ys
index 8427d25cb..29b16b627 100644
--- a/tests/silimate/opt_balance_tree.ys
+++ b/tests/silimate/opt_balance_tree.ys
@@ -14,27 +14,31 @@ endmodule
 EOF
 check -assert
 
-autoname
-write_json pre.json
-exec -- netlistsvg pre.json -o pre.svg
-
 # Check equivalence after opt_balance_tree
 equiv_opt -assert opt_balance_tree
 
 design -load postopt
-# opt_balance_tree
 
-autoname
-write_json post.json
-exec -- netlistsvg post.json -o post.svg
+# Checks if inputs to and gates has been rewired
+select -set a_wires i:a %co
+select -set driven_by_a @a_wires %co
+select -set and_a_cell t:$and @driven_by_a %i
 
-# Checks if y is still wired up to the correct gate
-# select -set y_wires o:y %ci
-# select -set y_driver @y_wires %ci
-# select -set and_y_cell t:$and @y_driver %i
-# select @and_y_cell -assert-count 1
-# select -set inputs @and_y_cell %ci
-# select -assert-count 1 @inputs i:c %i
+select -set b_wires i:b %co
+select -set driven_by_b @b_wires %co
+select -set and_b_cell t:$and @driven_by_b %i
+
+select -assert-none @and_a_cell @and_b_cell %d
+
+select -set c_wires i:c %co
+select -set driven_by_c @c_wires %co
+select -set and_c_cell t:$and @driven_by_c %i
+
+select -set d_wires i:d %co
+select -set driven_by_d @d_wires %co
+select -set and_d_cell t:$and @driven_by_d %i
+
+select -assert-none @and_c_cell @and_d_cell %d
 
 design -reset
 log -pop
@@ -188,5 +192,78 @@ equiv_opt -assert opt_balance_tree
 design -load postopt
 select -assert-count 3 t:$and
 
+design -reset
+log -pop
+
+log -header "Interesting tree situation"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a, b, c, d,
+  input wire e, f, g, h,
+  output wire x,
+);
+  wire i, j;
+
+  assign i = a & b & c & d;
+  assign j = e & f & g & h & i;
+
+  assign x = i & j;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+
+design -load postopt
+
+# Checks if inputs to and gates has been rewired
+
+# a and b
+select -set a_wires i:a %co
+select -set driven_by_a @a_wires %co
+select -set and_a_cell t:$and @driven_by_a %i
+
+select -set b_wires i:b %co
+select -set driven_by_b @b_wires %co
+select -set and_b_cell t:$and @driven_by_b %i
+
+select -assert-none @and_a_cell @and_b_cell %d
+
+# c and d
+select -set c_wires i:c %co
+select -set driven_by_c @c_wires %co
+select -set and_c_cell t:$and @driven_by_c %i
+
+select -set d_wires i:d %co
+select -set driven_by_d @d_wires %co
+select -set and_d_cell t:$and @driven_by_d %i
+
+select -assert-none @and_c_cell @and_d_cell %d
+
+# e and f
+select -set e_wires i:e %co
+select -set driven_by_e @e_wires %co
+select -set and_e_cell t:$and @driven_by_e %i
+
+select -set f_wires i:f %co
+select -set driven_by_f @f_wires %co
+select -set and_f_cell t:$and @driven_by_f %i
+
+select -assert-none @and_e_cell @and_f_cell %d
+
+# g and h
+select -set g_wires i:g %co
+select -set driven_by_g @g_wires %co
+select -set and_g_cell t:$and @driven_by_g %i
+
+select -set h_wires i:h %co
+select -set driven_by_h @h_wires %co
+select -set and_h_cell t:$and @driven_by_h %i
+
+select -assert-none @and_g_cell @and_h_cell %d
+
 design -reset
 log -pop
\ No newline at end of file

From fc86bd8e74744a0201c77f6b0c65470795d36d81 Mon Sep 17 00:00:00 2001
From: williamzhu17 <williammzhu@gmail.com>
Date: Sun, 11 May 2025 11:16:50 -0700
Subject: [PATCH 6/6] finalized tests

---
 tests/silimate/opt_balance_tree.ys | 115 ++++++++++++++++-------------
 1 file changed, 62 insertions(+), 53 deletions(-)

diff --git a/tests/silimate/opt_balance_tree.ys b/tests/silimate/opt_balance_tree.ys
index 29b16b627..aa42426eb 100644
--- a/tests/silimate/opt_balance_tree.ys
+++ b/tests/silimate/opt_balance_tree.ys
@@ -20,22 +20,18 @@ equiv_opt -assert opt_balance_tree
 design -load postopt
 
 # Checks if inputs to and gates has been rewired
-select -set a_wires i:a %co
-select -set driven_by_a @a_wires %co
+select -set driven_by_a i:a %co %co
 select -set and_a_cell t:$and @driven_by_a %i
 
-select -set b_wires i:b %co
-select -set driven_by_b @b_wires %co
+select -set driven_by_b i:b %co %co
 select -set and_b_cell t:$and @driven_by_b %i
 
 select -assert-none @and_a_cell @and_b_cell %d
 
-select -set c_wires i:c %co
-select -set driven_by_c @c_wires %co
+select -set driven_by_c i:c %co %co
 select -set and_c_cell t:$and @driven_by_c %i
 
-select -set d_wires i:d %co
-select -set driven_by_d @d_wires %co
+select -set driven_by_d i:d %co %co
 select -set and_d_cell t:$and @driven_by_d %i
 
 select -assert-none @and_c_cell @and_d_cell %d
@@ -69,8 +65,7 @@ equiv_opt -assert opt_balance_tree
 design -load postopt
 
 # Checks if y is still wired up to the correct gate
-select -set y_wires o:y %ci
-select -set y_driver @y_wires %ci
+select -set y_driver o:y %ci %ci
 select -set and_y_cell t:$and @y_driver %i
 select @and_y_cell -assert-count 1
 select -set inputs @and_y_cell %ci
@@ -107,6 +102,13 @@ equiv_opt -assert opt_balance_tree
 design -load postopt
 select -assert-count 3 t:$and
 
+# Checks if temp is still wired up to the correct gate
+select -set temp_driver w:temp %ci %ci
+select -set and_cell t:$and @temp_driver %i
+select @and_cell -assert-count 1
+select -set inputs @and_cell %ci
+select -assert-count 1 @inputs i:c %i
+
 design -reset
 log -pop
 
@@ -134,6 +136,42 @@ equiv_opt -assert opt_balance_tree
 design -load postopt
 select -assert-count 3 t:$and
 
+# Checks if x[1] is still wired up to the correct gate
+select -set target_drivers o:x %ci %ci
+select -set target_cells t:$and @target_drivers %i
+select -set inputs @target_cells %ci
+select -assert-count 1 @inputs i:c %i
+
+design -reset
+log -pop
+
+log -header "Fanout going to multiple outputs"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a,
+  input wire b,
+  input wire c,
+  input wire d,
+  output wire [2:0] x,
+  output wire y
+);
+  assign x[0] = a & b;
+  assign x[1] = x[0] & c;
+  assign x[2] = x[1] & d;
+
+  assign y = x[1];
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+
+design -load postopt
+select -assert-count 3 t:$and
+
 design -reset
 log -pop
 
@@ -164,33 +202,12 @@ equiv_opt -assert opt_balance_tree
 design -load postopt
 select -assert-count 3 t:$and
 
-design -reset
-log -pop
-
-log -header "Fanout going to multiple of the same word"
-log -push
-design -reset
-read_verilog <<EOF
-module top (
-  input wire a,
-  input wire b,
-  input wire c,
-  input wire d,
-  output wire [3:0] x
-);
-  assign x[0] = a & b;
-  assign x[1] = x[0] & c;
-  assign x[2] = x[1] & d;
-  assign x[3] = x[1];
-endmodule
-EOF
-check -assert
-
-# Check equivalence after opt_balance_tree
-equiv_opt -assert opt_balance_tree
-
-design -load postopt
-select -assert-count 3 t:$and
+# Checks if y is still wired up to the correct gate
+select -set y_wires o:y %ci
+select -set y_driver @y_wires %ci %ci
+select -set and_y_cell t:$and @y_driver %i
+select -set inputs @and_y_cell %ci
+select -assert-count 1 @inputs i:c %i
 
 design -reset
 log -pop
@@ -222,45 +239,37 @@ design -load postopt
 # Checks if inputs to and gates has been rewired
 
 # a and b
-select -set a_wires i:a %co
-select -set driven_by_a @a_wires %co
+select -set driven_by_a i:a %co %co
 select -set and_a_cell t:$and @driven_by_a %i
 
-select -set b_wires i:b %co
-select -set driven_by_b @b_wires %co
+select -set driven_by_b i:b %co %co
 select -set and_b_cell t:$and @driven_by_b %i
 
 select -assert-none @and_a_cell @and_b_cell %d
 
 # c and d
-select -set c_wires i:c %co
-select -set driven_by_c @c_wires %co
+select -set driven_by_c i:c %co %co
 select -set and_c_cell t:$and @driven_by_c %i
 
-select -set d_wires i:d %co
-select -set driven_by_d @d_wires %co
+select -set driven_by_d i:d %co %co
 select -set and_d_cell t:$and @driven_by_d %i
 
 select -assert-none @and_c_cell @and_d_cell %d
 
 # e and f
-select -set e_wires i:e %co
-select -set driven_by_e @e_wires %co
+select -set driven_by_e i:e %co %co
 select -set and_e_cell t:$and @driven_by_e %i
 
-select -set f_wires i:f %co
-select -set driven_by_f @f_wires %co
+select -set driven_by_f i:f %co %co
 select -set and_f_cell t:$and @driven_by_f %i
 
 select -assert-none @and_e_cell @and_f_cell %d
 
 # g and h
-select -set g_wires i:g %co
-select -set driven_by_g @g_wires %co
+select -set driven_by_g i:g %co %co
 select -set and_g_cell t:$and @driven_by_g %i
 
-select -set h_wires i:h %co
-select -set driven_by_h @h_wires %co
+select -set driven_by_h i:h %co %co
 select -set and_h_cell t:$and @driven_by_h %i
 
 select -assert-none @and_g_cell @and_h_cell %d