From 8b6925c5b022613a884afa3ef1c563408570a080 Mon Sep 17 00:00:00 2001
From: Natalia Kokoromyti <nataliakokoromyti@gmail.com>
Date: Tue, 13 Jan 2026 14:20:11 -0800
Subject: [PATCH 1/5] Add opt_balance_tree pass for timing optimization

This pass converts cascaded chains of arithmetic and logic cells ($add,
$mul, $and, $or, $xor) into balanced binary trees to improve timing
performance in hardware synthesis.

The optimization uses a breadth-first search approach to identify chains
of compatible cells, then recursively constructs balanced trees that
reduce the critical path depth.

Features:
- Supports arithmetic cells: $add, $mul
- Supports logic cells: $and, $or, $xor
- Command-line options: -arith (arithmetic only), -logic (logic only)
- Preserves signed/unsigned semantics
- Comprehensive test suite with 30 test cases

Original implementation by Akash Levy <akash@silimate.com> for Silimate.
Upstreamed from https://github.com/Silimate/yosys
---
 passes/opt/Makefile.inc        |    1 +
 passes/opt/opt_balance_tree.cc |  355 +++++++++
 tests/opt/opt_balance_tree.ys  | 1356 ++++++++++++++++++++++++++++++++
 3 files changed, 1712 insertions(+)
 create mode 100644 passes/opt/opt_balance_tree.cc
 create mode 100644 tests/opt/opt_balance_tree.ys

diff --git a/passes/opt/Makefile.inc b/passes/opt/Makefile.inc
index 426d9a79a..5dee824ff 100644
--- a/passes/opt/Makefile.inc
+++ b/passes/opt/Makefile.inc
@@ -23,6 +23,7 @@ OBJS += passes/opt/opt_lut_ins.o
 OBJS += passes/opt/opt_ffinv.o
 OBJS += passes/opt/pmux2shiftx.o
 OBJS += passes/opt/muxpack.o
+OBJS += passes/opt/opt_balance_tree.o
 
 OBJS += passes/opt/peepopt.o
 GENFILES += passes/opt/peepopt_pm.h
diff --git a/passes/opt/opt_balance_tree.cc b/passes/opt/opt_balance_tree.cc
new file mode 100644
index 000000000..273a74af3
--- /dev/null
+++ b/passes/opt/opt_balance_tree.cc
@@ -0,0 +1,355 @@
+/*
+ *  yosys -- Yosys Open SYnthesis Suite
+ *
+ *  Copyright (C) 2012  Claire Xenia Wolf <claire@yosyshq.com>
+ *                2019  Eddie Hung        <eddie@fpgeh.com>
+ *                2024  Akash Levy        <akash@silimate.com>
+ *
+ *  Permission to use, copy, modify, and/or distribute this software for any
+ *  purpose with or without fee is hereby granted, provided that the above
+ *  copyright notice and this permission notice appear in all copies.
+ *
+ *  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ *  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ *  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+ *  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ *  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+ *  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+ *  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+ *
+ */
+
+#include "kernel/yosys.h"
+#include "kernel/sigtools.h"
+
+USING_YOSYS_NAMESPACE
+PRIVATE_NAMESPACE_BEGIN
+
+
+struct OptBalanceTreeWorker {
+	// Module and signal map
+	Module *module;
+	SigMap sigmap;
+
+	// Counts of each cell type that are getting balanced
+	dict<IdString, int> cell_count;
+
+	// Check if cell is of the right type and has matching input/output widths
+	bool is_right_type(Cell* cell, IdString cell_type) {
+		return cell->type == cell_type &&
+		       cell->getParam(ID::Y_WIDTH).as_int() >= cell->getParam(ID::A_WIDTH).as_int() &&
+		       cell->getParam(ID::Y_WIDTH).as_int() >= cell->getParam(ID::B_WIDTH).as_int();
+	}
+
+	// Create a balanced binary tree from a vector of source signals
+	SigSpec create_balanced_tree(vector<SigSpec> &sources, IdString cell_type, Cell* cell) {
+		// Base case: if we have no sources, return an empty signal
+		if (sources.size() == 0)
+			return SigSpec();
+
+		// Base case: if we have only one source, return it
+		if (sources.size() == 1)
+			return sources[0];
+		
+		// Base case: if we have two sources, create a single cell
+		if (sources.size() == 2) {
+			// Create a new cell of the same type
+			Cell* new_cell = module->addCell(NEW_ID, cell_type);
+			
+			// Copy attributes from reference cell
+			new_cell->attributes = cell->attributes;
+			
+			// Create output wire
+			int out_width = cell->getParam(ID::Y_WIDTH).as_int();
+			if (cell_type == ID($add))
+				out_width = max(sources[0].size(), sources[1].size()) + 1;
+			else if (cell_type == ID($mul))
+				out_width = sources[0].size() + sources[1].size();
+			Wire* out_wire = module->addWire(NEW_ID, out_width);
+			
+			// Connect ports and fix up parameters
+			new_cell->setPort(ID::A, sources[0]);
+			new_cell->setPort(ID::B, sources[1]);
+			new_cell->setPort(ID::Y, out_wire);
+			new_cell->fixup_parameters();
+			new_cell->setParam(ID::A_SIGNED, cell->getParam(ID::A_SIGNED));
+			new_cell->setParam(ID::B_SIGNED, cell->getParam(ID::B_SIGNED));
+			
+			// Update count and return output wire
+			cell_count[cell_type]++;
+			return out_wire;
+		}
+		
+		// Recursive case: split sources into two groups and create subtrees
+		int mid = (sources.size() + 1) / 2;
+		vector<SigSpec> left_sources(sources.begin(), sources.begin() + mid);
+		vector<SigSpec> right_sources(sources.begin() + mid, sources.end());
+		
+		SigSpec left_tree = create_balanced_tree(left_sources, cell_type, cell);
+		SigSpec right_tree = create_balanced_tree(right_sources, cell_type, cell);
+		
+		// Create a cell to combine the two subtrees
+		Cell* new_cell = module->addCell(NEW_ID, cell_type);
+		
+		// Copy attributes from reference cell
+		new_cell->attributes = cell->attributes;
+		
+		// Create output wire
+		int out_width = cell->getParam(ID::Y_WIDTH).as_int();
+		if (cell_type == ID($add))
+			out_width = max(left_tree.size(), right_tree.size()) + 1;
+		else if (cell_type == ID($mul))
+			out_width = left_tree.size() + right_tree.size();
+		Wire* out_wire = module->addWire(NEW_ID, out_width);
+		
+		// Connect ports and fix up parameters
+		new_cell->setPort(ID::A, left_tree);
+		new_cell->setPort(ID::B, right_tree);
+		new_cell->setPort(ID::Y, out_wire);
+		new_cell->fixup_parameters();
+		new_cell->setParam(ID::A_SIGNED, cell->getParam(ID::A_SIGNED));
+		new_cell->setParam(ID::B_SIGNED, cell->getParam(ID::B_SIGNED));
+		
+		// Update count and return output wire
+		cell_count[cell_type]++;
+		return out_wire;
+	}
+
+	OptBalanceTreeWorker(Module *module, const vector<IdString> cell_types) : module(module), sigmap(module) {
+		// Do for each cell type
+		for (auto cell_type : cell_types) {
+			// Index all of the nets in the module
+			dict<SigSpec, Cell*> sig_to_driver;
+			dict<SigSpec, pool<Cell*>> sig_to_sink;
+			for (auto cell : module->selected_cells())
+			{
+				for (auto &conn : cell->connections())
+				{
+					if (cell->output(conn.first))
+						sig_to_driver[sigmap(conn.second)] = cell;
+
+					if (cell->input(conn.first))
+					{
+						SigSpec sig = sigmap(conn.second);
+						if (sig_to_sink.count(sig) == 0)
+							sig_to_sink[sig] = pool<Cell*>();
+						sig_to_sink[sig].insert(cell);
+					}
+				}
+			}
+
+			// Need to check if any wires connect to module ports
+			pool<SigSpec> input_port_sigs;
+			pool<SigSpec> output_port_sigs;
+			for (auto wire : module->selected_wires())
+				if (wire->port_input || wire->port_output) {
+					SigSpec sig = sigmap(wire);
+					for (auto bit : sig) {
+						if (wire->port_input)
+							input_port_sigs.insert(bit);
+						if (wire->port_output)
+							output_port_sigs.insert(bit);
+					}
+				}
+
+			// Actual logic starts here
+			pool<Cell*> consumed_cells;
+			for (auto cell : module->selected_cells())
+			{
+				// If consumed or not the correct type, skip
+				if (consumed_cells.count(cell) || !is_right_type(cell, cell_type))
+					continue;
+
+				// BFS, following all chains until they hit a cell of a different type
+				// Pick the longest one
+				auto y = sigmap(cell->getPort(ID::Y));
+				pool<Cell*> sinks;
+				pool<Cell*> current_loads = sig_to_sink[y];
+				pool<Cell*> next_loads;
+				while (!current_loads.empty())
+				{
+					// Find each sink and see what they are
+					for (auto x : current_loads)
+					{
+						// If not the correct type, don't follow any further
+						// (but add the originating cell to the list of sinks)
+						if (!is_right_type(x, cell_type))
+						{
+							sinks.insert(cell);
+							continue;
+						}
+
+						auto xy = sigmap(x->getPort(ID::Y));
+
+						// If this signal drives a port, add it to the sinks
+						// (even though it may not be the end of a chain)
+						for (auto bit : xy) {
+							if (output_port_sigs.count(bit) && !consumed_cells.count(x)) {
+								sinks.insert(x);
+								break;
+							}
+						}
+
+						// Search signal's fanout
+						auto& next = sig_to_sink[xy];
+						for (auto z : next)
+							next_loads.insert(z);
+					}
+
+					// If we couldn't find any downstream loads, stop.
+					// Create a reduction for each of the max-length chains we found
+					if (next_loads.empty())
+					{
+						for (auto s : current_loads)
+						{
+							// Not one of our gates? Don't follow any further
+							if (!is_right_type(s, cell_type))
+								continue;
+
+							sinks.insert(s);
+						}
+						break;
+					}
+
+					// Otherwise, continue down the chain
+					current_loads = next_loads;
+					next_loads.clear();
+				}
+
+				// We have our list of sinks, now go tree balance the chains
+				for (auto head_cell : sinks)
+				{
+					// Avoid duplication if we already were covered
+					if (consumed_cells.count(head_cell))
+						continue;
+
+					// Get sources of the chain
+					dict<SigSpec, int> sources;
+					dict<SigSpec, bool> signeds;
+					int inner_cells = 0;
+					std::deque<Cell*> bfs_queue = {head_cell};
+					while (bfs_queue.size())
+					{
+						Cell* x = bfs_queue.front();
+						bfs_queue.pop_front();
+
+						for (IdString port: {ID::A, ID::B}) {
+							auto sig = sigmap(x->getPort(port));
+							Cell* drv = sig_to_driver[sig];
+							bool drv_ok = drv && is_right_type(drv, cell_type);
+							for (auto bit : sig) {
+								if (input_port_sigs.count(bit) && !consumed_cells.count(drv)) {
+									drv_ok = false;
+									break;
+								}
+							}
+							if (drv_ok) {
+								inner_cells++;
+								bfs_queue.push_back(drv);
+							} else {
+								sources[sig]++;
+								signeds[sig] = x->getParam(port == ID::A ? ID::A_SIGNED : ID::B_SIGNED).as_bool();
+							}
+						}
+					}
+
+					if (inner_cells)
+					{
+						// Create a tree
+						log_debug("  Creating tree for %s with %d sources and %d inner cells...\n", log_id(head_cell), GetSize(sources), inner_cells);
+						
+						// Build a vector of all source signals
+						vector<SigSpec> source_signals;
+						vector<bool> signed_flags;
+						for (auto &source : sources) {
+							for (int i = 0; i < source.second; i++) {
+								source_signals.push_back(source.first);
+								signed_flags.push_back(signeds[source.first]);
+							}
+						}
+
+						// If not all signed flags are the same, do not balance
+						if (!std::all_of(signed_flags.begin(), signed_flags.end(), [&](bool flag) { return flag == signed_flags[0]; })) {
+							continue;
+						}
+						
+						// Create the balanced tree
+						SigSpec tree_output = create_balanced_tree(source_signals, cell_type, head_cell);
+						
+						// Connect the tree output to the head cell's output
+						SigSpec head_output = sigmap(head_cell->getPort(ID::Y));
+						int connect_width = std::min(head_output.size(), tree_output.size());
+						module->connect(head_output.extract(0, connect_width), tree_output.extract(0, connect_width));
+						if (head_output.size() > tree_output.size()) {
+							SigBit sext_bit = head_cell->getParam(ID::A_SIGNED).as_bool() ? head_output[connect_width - 1] : State::S0;
+							module->connect(head_output.extract(connect_width, head_output.size() - connect_width), SigSpec(sext_bit, head_output.size() - connect_width));
+						}
+
+						// Mark consumed cell for removal
+						consumed_cells.insert(head_cell);
+					}
+				}
+			}
+			
+			// Remove all consumed cells, which now have been replaced by trees
+			for (auto cell : consumed_cells)
+				module->remove(cell);
+		}
+	}
+};
+
+struct OptBalanceTreePass : public Pass {
+	OptBalanceTreePass() : Pass("opt_balance_tree", "$and/$or/$xor/$add/$mul cascades to trees") { }
+	void help() override {
+		//   |---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|
+		log("\n");
+		log("    opt_balance_tree [options] [selection]\n");
+		log("\n");
+		log("This pass converts cascaded chains of $and/$or/$xor/$add/$mul cells into\n");
+		log("trees of cells to improve timing.\n");
+		log("\n");
+		log("    -arith\n");
+		log("        only convert arithmetic cells.\n");
+		log("\n");
+		log("    -logic\n");
+		log("        only convert logic cells.\n");
+		log("\n");
+	}
+	void execute(std::vector<std::string> args, RTLIL::Design *design) override {
+		log_header(design, "Executing OPT_BALANCE_TREE pass (cell cascades to trees).\n");
+
+		// Handle arguments
+		size_t argidx;
+		vector<IdString> cell_types = {ID($and), ID($or), ID($xor), ID($add), ID($mul)};
+		for (argidx = 1; argidx < args.size(); argidx++) {
+			if (args[argidx] == "-arith") {
+				cell_types = {ID($add), ID($mul)};
+				continue;
+			}
+			if (args[argidx] == "-logic") {
+				cell_types = {ID($and), ID($or), ID($xor)};
+				continue;
+			}
+			break;
+		}
+		extra_args(args, argidx, design);
+
+		// Count of all cells that were packed
+		dict<IdString, int> cell_count;
+		for (auto module : design->selected_modules()) {
+			OptBalanceTreeWorker worker(module, cell_types);
+			for (auto cell : worker.cell_count) {
+				cell_count[cell.first] += cell.second;
+			}
+		}
+
+		// Log stats
+		for (auto cell_type : cell_types)
+			log("Converted %d %s cells into trees.\n", cell_count[cell_type], log_id(cell_type));
+
+		// Clean up
+		Yosys::run_pass("clean -purge");
+	}
+} OptBalanceTreePass;
+
+PRIVATE_NAMESPACE_END
diff --git a/tests/opt/opt_balance_tree.ys b/tests/opt/opt_balance_tree.ys
new file mode 100644
index 000000000..508f5fc24
--- /dev/null
+++ b/tests/opt/opt_balance_tree.ys
@@ -0,0 +1,1356 @@
+# Test 1
+log -header "Simple AND chain"
+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
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Get selections
+select i:a %co2 t:$and %i -set a_cells
+select i:b %co2 t:$and %i -set b_cells
+select i:c %co2 t:$and %i -set c_cells
+select i:d %co2 t:$and %i -set d_cells
+select o:x %ci2 t:$and %i -set x_cell
+select @a_cells @b_cells %i -set a_and_b_cells
+select @c_cells @d_cells %i -set c_and_d_cells
+select @a_cells -assert-count 1
+select @b_cells -assert-count 1
+select @c_cells -assert-count 1
+select @d_cells -assert-count 1
+select @x_cell -assert-count 1
+
+# Check that x is (a & b) & (c & d)
+select @a_and_b_cells %co3 @x_cell %i -assert-count 1
+select @c_and_d_cells %co3 @x_cell %i -assert-count 1
+
+design -reset
+log -pop
+
+
+
+# Test 2
+log -header "AND chain with intermediate outputs"
+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_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Merge should remove two redundant nodes for (a & b)
+select -assert-count 6 t:$and
+equiv_opt -assert opt_merge
+design -load postopt
+select -assert-count 4 t:$and
+
+# Get selections
+select i:a %co2 t:$and %i -set a_cells
+select i:b %co2 t:$and %i -set b_cells
+select i:c %co2 t:$and %i -set c_cells
+select i:d %co2 t:$and %i -set d_cells
+select o:x %ci2 t:$and %i -set x_cell
+select o:y %ci2 t:$and %i -set y_cell
+select o:z %ci2 t:$and %i -set z_cell
+select @c_cells @d_cells %i -set c_and_d_cells
+select @a_cells -assert-count 1
+select @b_cells -assert-count 1
+select @c_cells -assert-count 2
+select @d_cells -assert-count 1
+select @x_cell -assert-count 1
+select @y_cell -assert-count 1
+select @z_cell -assert-count 1
+select @c_and_d_cells -assert-count 1
+
+# Check that x is a & b
+select @a_cells @b_cells %i @x_cell %i -assert-count 1
+
+# Check that y is (a & b) & c
+select @c_cells @y_cell %i -assert-count 1
+select @x_cell %co3 @y_cell %i -assert-count 1
+
+# Check that z is (a & b) & (c & d)
+select @x_cell %co3 @z_cell %i -assert-count 1
+select @c_and_d_cells %co3 @z_cell %i -assert-count 1
+
+design -reset
+log -pop
+
+
+
+# Test 3
+log -header "AND chain with intermediate outputs and one inverter"
+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_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Merge should remove two redundant nodes for (a & b)
+select -assert-count 6 t:$and
+equiv_opt -assert opt_merge
+design -load postopt
+select -assert-count 4 t:$and
+
+# Get selections
+select i:a %co2 t:$and %i -set a_cells
+select i:b %co2 t:$and %i -set b_cells
+select i:c %co2 t:$and %i -set c_cells
+select i:d %co2 t:$and %i -set d_cells
+select o:x %ci2 t:$and %i -set x_cell
+select o:y %ci2 t:$not %i -set y_cell
+select o:z %ci2 t:$and %i -set z_cell
+select @y_cell %ci3 t:$and %i -set y_pre_cell
+select @c_cells @d_cells %i -set c_and_d_cells
+select @a_cells -assert-count 1
+select @b_cells -assert-count 1
+select @c_cells -assert-count 2
+select @d_cells -assert-count 1
+select @x_cell -assert-count 1
+select @y_cell -assert-count 1
+select @z_cell -assert-count 1
+select @y_pre_cell -assert-count 1
+select @c_and_d_cells -assert-count 1
+
+# Check that x is a & b
+select @a_cells @b_cells %i @x_cell %i -assert-count 1
+
+# Check that y_pre (y before inversion) is (a & b) & c
+select @c_cells @y_pre_cell %i -assert-count 1
+select @x_cell %co3 @y_pre_cell %i -assert-count 1
+
+# Check that z is (a & b) & (c & d)
+select @x_cell %co3 @z_cell %i -assert-count 1
+select @c_and_d_cells %co3 @z_cell %i -assert-count 1
+
+design -reset
+log -pop
+
+
+
+# Test 4
+log -header "AND chain with intermediate outputs to a 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_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Merge should remove two redundant nodes for (a & b)
+select -assert-count 6 t:$and
+equiv_opt -assert opt_merge
+design -load postopt
+select -assert-count 4 t:$and
+
+# Get selections
+select i:a %co2 t:$and %i -set a_cells
+select i:b %co2 t:$and %i -set b_cells
+select i:c %co2 t:$and %i -set c_cells
+select i:d %co2 t:$and %i -set d_cells
+select o:x %ci2 t:$and %i -set x_cells
+select @a_cells @b_cells %i -set a_and_b_cells
+select @c_cells @d_cells %i -set c_and_d_cells
+select @a_cells -assert-count 1
+select @b_cells -assert-count 1
+select @c_cells -assert-count 2
+select @d_cells -assert-count 1
+select @x_cells -assert-count 3
+select @a_and_b_cells -assert-count 1
+select @c_and_d_cells -assert-count 1
+
+# Check that x[0] is a & b
+select @a_and_b_cells @x_cells %i -assert-count 1
+
+# Check that x[1] is (a & b) & c
+select @c_cells @x_cells %i -assert-count 1
+
+# Check that x[2] is (a & b) & (c & d)
+select @a_and_b_cells %co3 @x_cells %i -assert-count 3
+select @c_and_d_cells %co3 @x_cells %i -assert-count 1
+
+design -reset
+log -pop
+
+
+
+# Test 5
+log -header "AND chain with intermediate outputs to a word out port and one other fanout"
+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
+
+# Merge should remove two redundant nodes for (a & b)
+select -assert-count 6 t:$and
+equiv_opt -assert opt_merge
+design -load postopt
+select -assert-count 4 t:$and
+
+# Get selections
+select i:a %co2 t:$and %i -set a_cells
+select i:b %co2 t:$and %i -set b_cells
+select i:c %co2 t:$and %i -set c_cells
+select i:d %co2 t:$and %i -set d_cells
+select o:x %ci2 t:$and %i -set x_cells
+select o:y %ci2 t:$and %i -set y_cell
+select @a_cells @b_cells %i -set a_and_b_cells
+select @c_cells @d_cells %i -set c_and_d_cells
+select @a_cells -assert-count 1
+select @b_cells -assert-count 1
+select @c_cells -assert-count 2
+select @d_cells -assert-count 1
+select @x_cells -assert-count 3
+select @y_cell -assert-count 1
+select @a_and_b_cells -assert-count 1
+select @c_and_d_cells -assert-count 1
+
+# Check that x[0] is a & b
+select @a_and_b_cells @x_cells %i -assert-count 1
+
+# Check that x[1] is (a & b) & c
+select @c_cells @x_cells %i -assert-count 1
+
+# Check that x[2] is (a & b) & (c & d)
+select @a_and_b_cells %co3 @x_cells %i -assert-count 3
+select @c_and_d_cells %co3 @x_cells %i -assert-count 1
+
+# Check that y is part of x
+select @x_cells @y_cell %i -assert-count 1
+
+design -reset
+log -pop
+
+
+
+# Test 6
+log -header "Pre-existing partial tree"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a, b, c, d,
+  input wire e, f, g, h,
+  input wire k, l, m, n,
+  output wire x,
+);
+  wire i, j;
+
+  assign i = a & b & c & d;
+  assign j = e & f & g & h;
+
+  assign x = i & j & k & l & m & n;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Assert that there are 11 AND gates
+select t:$and -assert-count 11
+
+# Assert that logic depth is 4
+# Check that 4 fanin-traversals results in all AND gates selected
+select o:x %ci2 %ci2 %ci2 %ci2 t:$and %i -assert-count 11
+
+design -reset
+log -pop
+
+
+
+# Test 7
+log -header "Pre-existing partial tree with intermediate outputs"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a, b, c, d,
+  input wire e, f, g, h,
+  input wire k, l, m, n,
+  output wire i, j,
+  output wire x,
+);
+
+  assign i = a & b & c & d;
+  assign j = e & f & g & h;
+
+  assign x = i & j & k & l & m & n;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Merge removes some redundancy
+equiv_opt -assert opt_merge
+design -load postopt
+
+# Assert that logic depth is 2 for i and j
+# Check that 2 fanin-traversals results in all AND gates selected
+select o:i %ci2 %ci2 t:$and %i -assert-count 3
+select o:j %ci2 %ci2 t:$and %i -assert-count 3
+
+# Assert that logic depth is 4 for x
+# Check that 4 fanin-traversals results in 11 AND gates selected
+select o:x %ci2 %ci2 %ci2 %ci2 t:$and %i -assert-count 11
+
+design -reset
+log -pop
+
+
+
+# Test 8
+log -header "Simple 1-bit ADD chain"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a,
+  input wire b,
+  input wire c,
+  input wire d,
+  input wire e,
+  input wire f,
+  input wire g,
+  input wire h,
+  output wire [3:0] x,
+);
+  assign x = a + b + c + d + e + f + g + h;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+# Get selections
+select i:a %co2 t:$add %i -set a_cell
+select i:b %co2 t:$add %i -set b_cell
+select i:c %co2 t:$add %i -set c_cell
+select i:d %co2 t:$add %i -set d_cell
+select i:e %co2 t:$add %i -set e_cell
+select i:f %co2 t:$add %i -set f_cell
+select i:g %co2 t:$add %i -set g_cell
+select i:h %co2 t:$add %i -set h_cell
+select o:x %ci2 t:$add %i -set x_cell
+select @a_cell @b_cell %i -set a_plus_b_cell
+select @c_cell @d_cell %i -set c_plus_d_cell
+select @e_cell @f_cell %i -set e_plus_f_cell
+select @g_cell @h_cell %i -set g_plus_h_cell
+select @a_cell -assert-count 1
+select @b_cell -assert-count 1
+select @c_cell -assert-count 1
+select @d_cell -assert-count 1
+select @e_cell -assert-count 1
+select @f_cell -assert-count 1
+select @g_cell -assert-count 1
+select @h_cell -assert-count 1
+select @x_cell -assert-count 1
+select @a_plus_b_cell -assert-count 1
+select @c_plus_d_cell -assert-count 1
+select @e_plus_f_cell -assert-count 1
+select @g_plus_h_cell -assert-count 1
+
+# Check intermediate cells
+select @a_plus_b_cell %co3 -set a_plus_b_fanout
+select @c_plus_d_cell %co3 -set c_plus_d_fanout
+select @e_plus_f_cell %co3 -set e_plus_f_fanout
+select @g_plus_h_cell %co3 -set g_plus_h_fanout
+select @a_plus_b_fanout @c_plus_d_fanout %i -set a_plus_b_plus_c_plus_d_cell
+select @e_plus_f_fanout @g_plus_h_fanout %i -set e_plus_f_plus_g_plus_h_cell
+select @a_plus_b_plus_c_plus_d_cell t:$add %i -assert-count 1
+select @e_plus_f_plus_g_plus_h_cell t:$add %i -assert-count 1
+
+# Check that x is ((a + b) + (c + d)) + ((e + f) + (g + h))
+select @a_plus_b_plus_c_plus_d_cell %co3 @x_cell %i -assert-count 1
+select @e_plus_f_plus_g_plus_h_cell %co3 @x_cell %i -assert-count 1
+
+design -reset
+log -pop
+
+
+
+# Test 9
+log -header "Add chain with intermediate outputs"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a,
+  input wire b,
+  input wire c,
+  input wire d,
+  output wire [1:0] x,
+  output wire [1:0] y,
+  output wire [1:0] z
+);
+  assign x = a + b;
+  assign y = x + c;
+  assign z = y + d;
+endmodule
+EOF
+check -assert
+
+# Proc for flops
+proc
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+# Merge should remove two redundant nodes for (a + b)
+select -assert-count 6 t:$add
+equiv_opt -assert opt_merge
+design -load postopt
+select -assert-count 4 t:$add
+
+# Get selections
+select i:a %co2 t:$add %i -set a_cells
+select i:b %co2 t:$add %i -set b_cells
+select i:c %co2 t:$add %i -set c_cells
+select i:d %co2 t:$add %i -set d_cells
+select o:x %ci2 t:$add %i -set x_cell
+select o:y %ci2 t:$add %i -set y_cell
+select o:z %ci2 t:$add %i -set z_cell
+select @c_cells @d_cells %i -set c_plus_d_cells
+select @a_cells -assert-count 1
+select @b_cells -assert-count 1
+select @c_cells -assert-count 2
+select @d_cells -assert-count 1
+select @x_cell -assert-count 1
+select @y_cell -assert-count 1
+select @z_cell -assert-count 1
+select @c_plus_d_cells -assert-count 1
+
+# Check that x is a + b
+select @a_cells @b_cells %i @x_cell %i -assert-count 1
+
+# Check that y is (a + b) + c
+select @c_cells @y_cell %i -assert-count 1
+select @x_cell %co3 @y_cell %i -assert-count 1
+
+# Check that z is (a + b) + (c + d)
+select @x_cell %co3 @z_cell %i -assert-count 1
+select @c_plus_d_cells %co3 @z_cell %i -assert-count 1
+
+design -reset
+log -pop
+
+
+
+# Test 10
+log -header "Add chain with intermediate outputs and one inverter"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a,
+  input wire b,
+  input wire c,
+  input wire d,
+  output wire [1:0] x,
+  output wire [1:0] y,
+  output wire [1:0] z
+);
+  wire [1:0] temp;
+  assign y = ~temp;
+  assign x = a + b;
+  assign temp = x + c;
+  assign z = temp + d;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+# Merge should remove two redundant nodes for (a + b)
+select -assert-count 6 t:$add
+equiv_opt -assert opt_merge
+design -load postopt
+select -assert-count 4 t:$add
+
+# Get selections
+select i:a %co2 t:$add %i -set a_cells
+select i:b %co2 t:$add %i -set b_cells
+select i:c %co2 t:$add %i -set c_cells
+select i:d %co2 t:$add %i -set d_cells
+select o:x %ci2 t:$add %i -set x_cell
+select o:y %ci2 t:$not %i -set y_cell
+select o:z %ci2 t:$add %i -set z_cell
+select @y_cell %ci3 t:$add %i -set y_pre_cell
+select @c_cells @d_cells %i -set c_plus_d_cells
+select @a_cells -assert-count 1
+select @b_cells -assert-count 1
+select @c_cells -assert-count 2
+select @d_cells -assert-count 1
+select @x_cell -assert-count 1
+select @y_cell -assert-count 1
+select @z_cell -assert-count 1
+select @y_pre_cell -assert-count 1
+select @c_plus_d_cells -assert-count 1
+
+# Check that x is a + b
+select @a_cells @b_cells %i @x_cell %i -assert-count 1
+
+# Check that y_pre (y before inversion) is (a + b) + c
+select @c_cells @y_pre_cell %i -assert-count 1
+select @x_cell %co3 @y_pre_cell %i -assert-count 1
+
+# Check that z is (a + b) + (c + d)
+select @x_cell %co3 @z_cell %i -assert-count 1
+select @c_plus_d_cells %co3 @z_cell %i -assert-count 1
+
+design -reset
+log -pop
+
+
+
+# Test 11
+log -header "Add chain with intermediate outputs to a 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 [5:0] x
+);
+  assign x[1:0] = a + b;
+  assign x[3:2] = x[1:0] + c;
+  assign x[5:4] = x[3:2] + d;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+# Merge should remove two redundant nodes for (a + b)
+select -assert-count 6 t:$add
+equiv_opt -assert opt_merge
+design -load postopt
+select -assert-count 4 t:$add
+
+# Get selections
+select i:a %co2 t:$add %i -set a_cells
+select i:b %co2 t:$add %i -set b_cells
+select i:c %co2 t:$add %i -set c_cells
+select i:d %co2 t:$add %i -set d_cells
+select o:x %ci2 t:$add %i -set x_cells
+select @a_cells @b_cells %i -set a_plus_b_cells
+select @c_cells @d_cells %i -set c_plus_d_cells
+select @a_cells -assert-count 1
+select @b_cells -assert-count 1
+select @c_cells -assert-count 2
+select @d_cells -assert-count 1
+select @x_cells -assert-count 3
+select @a_plus_b_cells -assert-count 1
+select @c_plus_d_cells -assert-count 1
+
+# Check that x[1:0] is a + b
+select @a_plus_b_cells @x_cells %i -assert-count 1
+
+# Check that x[3:2] is (a + b) + c
+select @c_cells @x_cells %i -assert-count 1
+
+# Check that x[5:4] is (a + b) + (c + d)
+select @a_plus_b_cells %co3 @x_cells %i -assert-count 3
+select @c_plus_d_cells %co3 @x_cells %i -assert-count 1
+
+design -reset
+log -pop
+
+
+
+# Test 12
+log -header "Add chain with intermediate outputs to a word out port and one other fanout"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a,
+  input wire b,
+  input wire c,
+  input wire d,
+  output wire [5:0] x,
+  output wire y
+);
+  assign x[1:0] = a + b;
+  assign x[3:2] = x[1:0] + c;
+  assign x[5:4] = x[3:2] + d;
+
+  assign y = x[3];
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+# Merge should remove two redundant nodes for (a + b)
+select -assert-count 6 t:$add
+equiv_opt -assert opt_merge
+design -load postopt
+select -assert-count 4 t:$add
+
+# Get selections
+select i:a %co2 t:$add %i -set a_cells
+select i:b %co2 t:$add %i -set b_cells
+select i:c %co2 t:$add %i -set c_cells
+select i:d %co2 t:$add %i -set d_cells
+select o:x %ci2 t:$add %i -set x_cells
+select @a_cells @b_cells %i -set a_plus_b_cells
+select @c_cells @d_cells %i -set c_plus_d_cells
+select @a_cells -assert-count 1
+select @b_cells -assert-count 1
+select @c_cells -assert-count 2
+select @d_cells -assert-count 1
+select @x_cells -assert-count 3
+select @a_plus_b_cells -assert-count 1
+select @c_plus_d_cells -assert-count 1
+
+# Check that x[1:0] is a + b
+select @a_plus_b_cells @x_cells %i -assert-count 1
+
+# Check that x[3:2] is (a + b) + c
+select @c_cells @x_cells %i -assert-count 1
+
+# Check that x[5:4] is (a + b) + (c + d)
+select @a_plus_b_cells %co3 @x_cells %i -assert-count 3
+select @c_plus_d_cells %co3 @x_cells %i -assert-count 1
+
+design -reset
+log -pop
+
+
+
+# Test 13
+log -header "Pre-existing partial tree"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a, b, c, d,
+  input wire e, f, g, h,
+  input wire k, l, m, n,
+  output wire [3:0] x,
+);
+  wire [2:0] i, j;
+
+  assign i = a + b + c + d;
+  assign j = e + f + g + h;
+
+  assign x = i + j + k + l + m + n;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+# Assert that there are 11 adders
+select t:$add -assert-count 11
+
+# Assert that logic depth is 4
+# Check that 4 fanin-traversals results in all adders selected
+select o:x %ci2 %ci2 %ci2 %ci2 t:$add %i -assert-count 11
+
+design -reset
+log -pop
+
+
+
+# Test 14
+log -header "Pre-existing partial tree with intermediate outputs"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a, b, c, d,
+  input wire e, f, g, h,
+  input wire k, l, m, n,
+  output wire [2:0] i, j,
+  output wire [3:0] x,
+);
+
+  assign i = a + b + c + d;
+  assign j = e + f + g + h;
+
+  assign x = i + j + k + l + m + n;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+# Merge removes some redundancy
+equiv_opt -assert opt_merge
+design -load postopt
+
+# Assert that logic depth is 2 for i and j
+# Check that 2 fanin-traversals results in all adders selected
+select o:i %ci2 %ci2 t:$add %i -assert-count 3
+select o:j %ci2 %ci2 t:$add %i -assert-count 3
+
+# Assert that logic depth is 4 for x
+# Check that 4 fanin-traversals results in 11 adders selected
+select o:x %ci2 %ci2 %ci2 %ci2 t:$add %i -assert-count 11
+
+design -reset
+log -pop
+
+
+
+# Test 15
+log -header "Mixed signed/unsigned adders with different bit widths"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire [3:0] a,
+  input wire signed [4:0] b,
+  input wire [2:0] c,
+  input wire signed [3:0] d,
+  input wire [1:0] e,
+  input wire signed [2:0] f,
+  output wire signed [7:0] x,
+);
+  assign x = a + b + c + d + e + f;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 16
+log -header "Adder chain with fanout to other logic (multiplexer)"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire [2:0] a, b, c, d,
+  input wire sel,
+  output wire [4:0] sum,
+  output wire [2:0] mux_out
+);
+  wire [3:0] partial_sum;
+  assign partial_sum = a + b;
+  assign sum = partial_sum + c + d;
+  assign mux_out = sel ? partial_sum[2:0] : c;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 17
+log -header "Adder chain with fanin from other logic (shift operation)"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire [3:0] a, b, c,
+  input wire [1:0] shift_amt,
+  output wire [5:0] result
+);
+  wire [4:0] shifted_a;
+  assign shifted_a = a << shift_amt;
+  assign result = shifted_a + b + c;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 18
+log -header "Bit-split fanout - different bits go to different outputs"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire [2:0] a, b, c, d,
+  output wire [3:0] sum,
+  output wire low_bit,
+  output wire [1:0] mid_bits,
+  output wire high_bit
+);
+  assign sum = a + b + c + d;
+  assign low_bit = sum[0];
+  assign mid_bits = sum[2:1];
+  assign high_bit = sum[3];
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 19
+log -header "Bit-split fanin - different bits come from different sources"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire [7:0] data_in,
+  input wire [1:0] a, b, c,
+  output wire [4:0] result
+);
+  wire [1:0] low_bits;
+  wire [2:0] mid_bits;
+  wire [1:0] high_bits;
+  
+  assign low_bits = data_in[1:0];
+  assign mid_bits = data_in[4:2];
+  assign high_bits = data_in[7:6];
+  
+  assign result = {1'b0, low_bits} + {1'b0, mid_bits} + {3'b0, high_bits} + {2'b0, a} + {2'b0, b} + {2'b0, c};
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 20
+log -header "Converge/reconverge network - diamond pattern"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire [2:0] a, b, c, d,
+  output wire [4:0] out1,
+  output wire [4:0] out2
+);
+  wire [3:0] sum1, sum2;
+  
+  assign sum1 = a + b;
+  assign sum2 = c + d;
+  assign out1 = sum1 + sum2;
+  assign out2 = sum1 + c;  // Reconvergence: sum1 used again
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 21
+log -header "Same input added multiple times in different paths"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire [2:0] a, b, c,
+  output wire [5:0] result
+);
+  wire [4:0] sum1, sum2;
+  
+  assign sum1 = a + b + a;  // 'a' appears twice
+  assign sum2 = a + c;      // 'a' appears again
+  assign result = sum1 + sum2;  // Final sum includes 'a' three times total
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 22
+log -header "Complex scenario combining multiple test patterns"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire [3:0] a,
+  input wire signed [4:0] b,
+  input wire [2:0] c, d, e,
+  input wire sel,
+  output wire signed [7:0] main_result,
+  output wire [3:0] side_result,
+  output wire overflow
+);
+  wire [4:0] partial1, partial2;
+  wire [5:0] intermediate;
+  
+  assign partial1 = a + c;
+  assign partial2 = d + e + c;  // 'c' used twice
+  assign intermediate = partial1 + partial2;
+  assign main_result = intermediate + b;
+  assign side_result = sel ? partial1[3:0] : partial2[3:0];
+  assign overflow = intermediate[5];
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 23
+log -header "Wide adders with varying input widths"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire [15:0] wide_a,
+  input wire [7:0] med_b,
+  input wire [3:0] small_c,
+  input wire [1:0] tiny_d,
+  input wire single_e,
+  output wire [16:0] wide_result
+);
+  assign wide_result = wide_a + med_b + small_c + tiny_d + single_e;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 24
+log -header "Mixed arithmetic operations (add/sub) in tree structure"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire [4:0] a, b, c, d, e, f,
+  output wire [6:0] result
+);
+  wire [6:0] sum_part, diff_part;
+  
+  assign sum_part = a + b + c;
+  assign diff_part = d - e + f;  // Mix of add and subtract
+  assign result = sum_part + diff_part;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 25
+log -header "Unsigned inputs with signed intermediate and output"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire [3:0] a, b, c, d,  // All unsigned
+  output wire signed [6:0] result  // Signed output
+);
+  wire signed [4:0] intermediate1, intermediate2;
+  
+  assign intermediate1 = $signed(a) + $signed(b);  // Convert to signed
+  assign intermediate2 = $signed(c) + $signed(d);  // Convert to signed
+  assign result = intermediate1 + intermediate2;   // Signed addition
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 26
+log -header "Mixed signed/unsigned with negative values possible"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire signed [4:0] a, b,    // Signed inputs (can be negative)
+  input wire [3:0] c, d,           // Unsigned inputs
+  output wire signed [7:0] result  // Signed output
+);
+  wire signed [5:0] sum_signed, sum_unsigned;
+  
+  assign sum_signed = a + b;                    // Signed + Signed
+  assign sum_unsigned = $signed({1'b0, c}) + $signed({1'b0, d});  // Force unsigned to signed safely
+  assign result = sum_signed + sum_unsigned;   // Mixed result
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 27
+log -header "Sign extension in adder tree with different input widths"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire signed [2:0] narrow_a,     // 3-bit signed
+  input wire signed [4:0] wide_b,       // 5-bit signed  
+  input wire [3:0] unsigned_c,          // 4-bit unsigned
+  input wire signed [3:0] signed_d,     // 4-bit signed
+  output wire signed [8:0] result       // Wide signed output
+);
+  // Let Verilog handle sign extension automatically in the adder tree
+  assign result = narrow_a + wide_b + $signed({1'b0, unsigned_c}) + signed_d;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 28
+log -header "Unsigned overflow vs signed extension in tree balancing"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire [2:0] u_a, u_b, u_c,           // Small unsigned inputs
+  input wire signed [2:0] s_a, s_b, s_c,    // Small signed inputs (same width)
+  output wire [5:0] unsigned_result,        // Unsigned result
+  output wire signed [5:0] signed_result    // Signed result
+);
+  // Test how signedness affects tree balancing for same-width inputs
+  assign unsigned_result = u_a + u_b + u_c;  // All unsigned arithmetic
+  assign signed_result = s_a + s_b + s_c;    // All signed arithmetic
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 29
+log -header "Signedness with bit selection and concatenation"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire signed [7:0] data_a, data_b,
+  input wire [3:0] unsigned_offset,
+  output wire signed [9:0] result
+);
+  wire signed [8:0] sum_full;
+  wire signed [6:0] sum_high, sum_low;
+  
+  assign sum_full = data_a + data_b;
+  assign sum_high = data_a[7:2] + data_b[7:2];           // High bits (signed)
+  assign sum_low = $signed({1'b0, data_a[5:0]}) + $signed({1'b0, data_b[5:0]});  // Low bits as unsigned->signed
+  
+  // Combine with unsigned offset
+  assign result = sum_full + $signed({6'b0, unsigned_offset});
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop
+
+
+
+# Test 30
+log -header "Complex signedness with conditional sign extension"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire signed [4:0] base_val,
+  input wire [3:0] pos_offset1, pos_offset2,
+  input wire signed [3:0] signed_offset,
+  input wire extend_sign,
+  output wire signed [8:0] result
+);
+  wire signed [5:0] extended_base;
+  wire signed [6:0] offset_sum;
+  
+  // Conditional sign extension based on input
+  assign extended_base = extend_sign ? 
+    {{1{base_val[4]}}, base_val} :     // Sign extend
+    {1'b0, base_val};                  // Zero extend
+  
+  // Mix of signed and unsigned offsets
+  assign offset_sum = $signed({2'b0, pos_offset1}) + 
+                     $signed({2'b0, pos_offset2}) + 
+                     {{2{signed_offset[3]}}, signed_offset};
+  
+  assign result = extended_base + offset_sum;
+endmodule
+EOF
+check -assert
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+design -load postopt
+
+# Width reduction
+equiv_opt -assert wreduce
+design -load postopt
+
+design -reset
+log -pop

From 6aef8ea8ab28702b4429147ddbea825ca4715582 Mon Sep 17 00:00:00 2001
From: Natalia Kokoromyti <nataliakokoromyti@gmail.com>
Date: Tue, 13 Jan 2026 15:31:46 -0800
Subject: [PATCH 2/5] Add missing <deque> include for MSVC compatibility

---
 passes/opt/opt_balance_tree.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/passes/opt/opt_balance_tree.cc b/passes/opt/opt_balance_tree.cc
index 273a74af3..8811b1331 100644
--- a/passes/opt/opt_balance_tree.cc
+++ b/passes/opt/opt_balance_tree.cc
@@ -21,6 +21,7 @@
 
 #include "kernel/yosys.h"
 #include "kernel/sigtools.h"
+#include <deque>
 
 USING_YOSYS_NAMESPACE
 PRIVATE_NAMESPACE_BEGIN

From 60ac3670cbcf320275fc7cf31ece78638fa6e87f Mon Sep 17 00:00:00 2001
From: Natalia <nataliakokoromyti@gmail.com>
Date: Wed, 14 Jan 2026 13:12:55 -0800
Subject: [PATCH 3/5] Fix truncation issue in opt_balance_tree pass

Only allow rebalancing of cells with "natural" output widths (no truncation).
This prevents equivalence failures when moving operands between adders
with different intermediate truncation points.

For each operation type, the natural width is:
- Addition: max(A_WIDTH, B_WIDTH) + 1 (for carry bit)
- Multiplication: A_WIDTH + B_WIDTH
- Logic ops: max(A_WIDTH, B_WIDTH)

Fixes widlarizer's counterexample in YosysHQ/yosys#5605 where an 8-bit
intermediate wire was intentionally truncating adder results, and
rebalancing would change where that truncation occurred.
---
 passes/opt/opt_balance_tree.cc | 27 ++++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/passes/opt/opt_balance_tree.cc b/passes/opt/opt_balance_tree.cc
index 8811b1331..9e4fa0620 100644
--- a/passes/opt/opt_balance_tree.cc
+++ b/passes/opt/opt_balance_tree.cc
@@ -36,10 +36,31 @@ struct OptBalanceTreeWorker {
 	dict<IdString, int> cell_count;
 
 	// Check if cell is of the right type and has matching input/output widths
+	// Only allow cells with "natural" output widths (no truncation) to prevent
+	// equivalence issues when rebalancing (see YosysHQ/yosys#5605)
 	bool is_right_type(Cell* cell, IdString cell_type) {
-		return cell->type == cell_type &&
-		       cell->getParam(ID::Y_WIDTH).as_int() >= cell->getParam(ID::A_WIDTH).as_int() &&
-		       cell->getParam(ID::Y_WIDTH).as_int() >= cell->getParam(ID::B_WIDTH).as_int();
+		if (cell->type != cell_type)
+			return false;
+
+		int y_width = cell->getParam(ID::Y_WIDTH).as_int();
+		int a_width = cell->getParam(ID::A_WIDTH).as_int();
+		int b_width = cell->getParam(ID::B_WIDTH).as_int();
+
+		// Calculate the "natural" output width for this operation
+		int natural_width;
+		if (cell_type == ID($add)) {
+			// Addition produces max(A_WIDTH, B_WIDTH) + 1 (for carry bit)
+			natural_width = std::max(a_width, b_width) + 1;
+		} else if (cell_type == ID($mul)) {
+			// Multiplication produces A_WIDTH + B_WIDTH
+			natural_width = a_width + b_width;
+		} else {
+			// Logic operations ($and/$or/$xor) produce max(A_WIDTH, B_WIDTH)
+			natural_width = std::max(a_width, b_width);
+		}
+
+		// Only allow cells where Y_WIDTH equals the natural width (no truncation)
+		return y_width == natural_width;
 	}
 
 	// Create a balanced binary tree from a vector of source signals

From 305b6c81d7a39a4ee46f94eaded2162e2dd74eb3 Mon Sep 17 00:00:00 2001
From: Natalia <nataliakokoromyti@gmail.com>
Date: Wed, 14 Jan 2026 14:58:53 -0800
Subject: [PATCH 4/5] Refine width check to allow Y_WIDTH >= natural width

Change from equality check to >= to allow cells where output
is wider than natural width (zero-extended). Only reject cells
with Y_WIDTH < natural width (truncated).

This fixes test failures while still preventing the truncation
issue identified in widlarizer's feedback.
---
 passes/opt/opt_balance_tree.cc | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/passes/opt/opt_balance_tree.cc b/passes/opt/opt_balance_tree.cc
index 9e4fa0620..5c19e3975 100644
--- a/passes/opt/opt_balance_tree.cc
+++ b/passes/opt/opt_balance_tree.cc
@@ -59,8 +59,9 @@ struct OptBalanceTreeWorker {
 			natural_width = std::max(a_width, b_width);
 		}
 
-		// Only allow cells where Y_WIDTH equals the natural width (no truncation)
-		return y_width == natural_width;
+		// Only allow cells where Y_WIDTH >= natural width (no truncation)
+		// This prevents rebalancing chains where truncation semantics matter
+		return y_width >= natural_width;
 	}
 
 	// Create a balanced binary tree from a vector of source signals

From d5e1647d1167a0943a2b2c488ce63291c35d3972 Mon Sep 17 00:00:00 2001
From: Natalia <nataliakokoromyti@gmail.com>
Date: Wed, 14 Jan 2026 17:06:01 -0800
Subject: [PATCH 5/5] fix tests with truncation issues

---
 tests/opt/opt_balance_tree.ys | 154 +++++-----------------------------
 1 file changed, 22 insertions(+), 132 deletions(-)

diff --git a/tests/opt/opt_balance_tree.ys b/tests/opt/opt_balance_tree.ys
index 508f5fc24..6f8b8b711 100644
--- a/tests/opt/opt_balance_tree.ys
+++ b/tests/opt/opt_balance_tree.ys
@@ -385,7 +385,7 @@ log -pop
 
 
 # Test 8
-log -header "Simple 1-bit ADD chain"
+log -header "Simple 1-bit ADD chain (4 inputs)"
 log -push
 design -reset
 read_verilog <<EOF
@@ -394,13 +394,13 @@ module top (
   input wire b,
   input wire c,
   input wire d,
-  input wire e,
-  input wire f,
-  input wire g,
-  input wire h,
-  output wire [3:0] x,
+  output wire [3:0] x
 );
-  assign x = a + b + c + d + e + f + g + h;
+  wire [1:0] ab;
+  wire [2:0] abc;
+  assign ab = a + b;
+  assign abc = ab + c;
+  assign x = abc + d;
 endmodule
 EOF
 check -assert
@@ -418,42 +418,20 @@ select i:a %co2 t:$add %i -set a_cell
 select i:b %co2 t:$add %i -set b_cell
 select i:c %co2 t:$add %i -set c_cell
 select i:d %co2 t:$add %i -set d_cell
-select i:e %co2 t:$add %i -set e_cell
-select i:f %co2 t:$add %i -set f_cell
-select i:g %co2 t:$add %i -set g_cell
-select i:h %co2 t:$add %i -set h_cell
 select o:x %ci2 t:$add %i -set x_cell
 select @a_cell @b_cell %i -set a_plus_b_cell
 select @c_cell @d_cell %i -set c_plus_d_cell
-select @e_cell @f_cell %i -set e_plus_f_cell
-select @g_cell @h_cell %i -set g_plus_h_cell
 select @a_cell -assert-count 1
 select @b_cell -assert-count 1
 select @c_cell -assert-count 1
 select @d_cell -assert-count 1
-select @e_cell -assert-count 1
-select @f_cell -assert-count 1
-select @g_cell -assert-count 1
-select @h_cell -assert-count 1
 select @x_cell -assert-count 1
 select @a_plus_b_cell -assert-count 1
 select @c_plus_d_cell -assert-count 1
-select @e_plus_f_cell -assert-count 1
-select @g_plus_h_cell -assert-count 1
 
-# Check intermediate cells
-select @a_plus_b_cell %co3 -set a_plus_b_fanout
-select @c_plus_d_cell %co3 -set c_plus_d_fanout
-select @e_plus_f_cell %co3 -set e_plus_f_fanout
-select @g_plus_h_cell %co3 -set g_plus_h_fanout
-select @a_plus_b_fanout @c_plus_d_fanout %i -set a_plus_b_plus_c_plus_d_cell
-select @e_plus_f_fanout @g_plus_h_fanout %i -set e_plus_f_plus_g_plus_h_cell
-select @a_plus_b_plus_c_plus_d_cell t:$add %i -assert-count 1
-select @e_plus_f_plus_g_plus_h_cell t:$add %i -assert-count 1
-
-# Check that x is ((a + b) + (c + d)) + ((e + f) + (g + h))
-select @a_plus_b_plus_c_plus_d_cell %co3 @x_cell %i -assert-count 1
-select @e_plus_f_plus_g_plus_h_cell %co3 @x_cell %i -assert-count 1
+# Check that x is (a + b) + (c + d)
+select @a_plus_b_cell %co3 @x_cell %i -assert-count 1
+select @c_plus_d_cell %co3 @x_cell %i -assert-count 1
 
 design -reset
 log -pop
@@ -471,8 +449,8 @@ module top (
   input wire c,
   input wire d,
   output wire [1:0] x,
-  output wire [1:0] y,
-  output wire [1:0] z
+  output wire [2:0] y,
+  output wire [3:0] z
 );
   assign x = a + b;
   assign y = x + c;
@@ -543,10 +521,10 @@ module top (
   input wire c,
   input wire d,
   output wire [1:0] x,
-  output wire [1:0] y,
-  output wire [1:0] z
+  output wire [2:0] y,
+  output wire [3:0] z
 );
-  wire [1:0] temp;
+  wire [2:0] temp;
   assign y = ~temp;
   assign x = a + b;
   assign temp = x + c;
@@ -615,11 +593,11 @@ module top (
   input wire b,
   input wire c,
   input wire d,
-  output wire [5:0] x
+  output wire [8:0] x
 );
   assign x[1:0] = a + b;
-  assign x[3:2] = x[1:0] + c;
-  assign x[5:4] = x[3:2] + d;
+  assign x[4:2] = x[1:0] + c;
+  assign x[8:5] = x[4:2] + d;
 endmodule
 EOF
 check -assert
@@ -679,14 +657,14 @@ module top (
   input wire b,
   input wire c,
   input wire d,
-  output wire [5:0] x,
+  output wire [8:0] x,
   output wire y
 );
   assign x[1:0] = a + b;
-  assign x[3:2] = x[1:0] + c;
-  assign x[5:4] = x[3:2] + d;
+  assign x[4:2] = x[1:0] + c;
+  assign x[8:5] = x[4:2] + d;
 
-  assign y = x[3];
+  assign y = x[4];
 endmodule
 EOF
 check -assert
@@ -736,94 +714,6 @@ log -pop
 
 
 
-# Test 13
-log -header "Pre-existing partial tree"
-log -push
-design -reset
-read_verilog <<EOF
-module top (
-  input wire a, b, c, d,
-  input wire e, f, g, h,
-  input wire k, l, m, n,
-  output wire [3:0] x,
-);
-  wire [2:0] i, j;
-
-  assign i = a + b + c + d;
-  assign j = e + f + g + h;
-
-  assign x = i + j + k + l + m + n;
-endmodule
-EOF
-check -assert
-
-# Check equivalence after opt_balance_tree
-equiv_opt -assert opt_balance_tree
-design -load postopt
-
-# Width reduction
-equiv_opt -assert wreduce
-design -load postopt
-
-# Assert that there are 11 adders
-select t:$add -assert-count 11
-
-# Assert that logic depth is 4
-# Check that 4 fanin-traversals results in all adders selected
-select o:x %ci2 %ci2 %ci2 %ci2 t:$add %i -assert-count 11
-
-design -reset
-log -pop
-
-
-
-# Test 14
-log -header "Pre-existing partial tree with intermediate outputs"
-log -push
-design -reset
-read_verilog <<EOF
-module top (
-  input wire a, b, c, d,
-  input wire e, f, g, h,
-  input wire k, l, m, n,
-  output wire [2:0] i, j,
-  output wire [3:0] x,
-);
-
-  assign i = a + b + c + d;
-  assign j = e + f + g + h;
-
-  assign x = i + j + k + l + m + n;
-endmodule
-EOF
-check -assert
-
-# Check equivalence after opt_balance_tree
-equiv_opt -assert opt_balance_tree
-design -load postopt
-
-# Width reduction
-equiv_opt -assert wreduce
-design -load postopt
-
-# Merge removes some redundancy
-equiv_opt -assert opt_merge
-design -load postopt
-
-# Assert that logic depth is 2 for i and j
-# Check that 2 fanin-traversals results in all adders selected
-select o:i %ci2 %ci2 t:$add %i -assert-count 3
-select o:j %ci2 %ci2 t:$add %i -assert-count 3
-
-# Assert that logic depth is 4 for x
-# Check that 4 fanin-traversals results in 11 adders selected
-select o:x %ci2 %ci2 %ci2 %ci2 t:$add %i -assert-count 11
-
-design -reset
-log -pop
-
-
-
 # Test 15
 log -header "Mixed signed/unsigned adders with different bit widths"
 log -push