diff --git a/passes/opt/wreduce.cc b/passes/opt/wreduce.cc
index 532aad96c..9450cb329 100644
--- a/passes/opt/wreduce.cc
+++ b/passes/opt/wreduce.cc
@@ -469,28 +469,74 @@ struct WreduceWorker
 					keep_bits.insert(bit);
 		}
 
-		for (auto c : module->selected_cells())
-			work_queue_cells.insert(c);
-
-		while (!work_queue_cells.empty())
-		{
-			work_queue_bits.clear();
-			for (auto c : work_queue_cells)
-				run_cell(c);
-
-			work_queue_cells.clear();
-			for (auto bit : work_queue_bits)
-			for (auto port : mi.query_ports(bit))
-				if (module->selected(port.cell))
-					work_queue_cells.insert(port.cell);
-		}
-
 		pool<SigSpec> complete_wires;
 		for (auto w : module->wires())
 			complete_wires.insert(mi.sigmap(w));
 
+		// Build bit_drivers_db for cells
+		dict<SigBit, tuple<IdString,IdString>> bit_drivers_db;
+		for (auto cell : module->cells()) {
+			for (auto conn : cell->connections()) {
+				if (!cell->output(conn.first)) continue;
+				for (int i = 0; i < GetSize(conn.second); i++) {
+					SigBit bit(mi.sigmap(conn.second[i]));
+					bit_drivers_db[bit] = tuple<IdString,IdString>(cell->name, conn.first);
+				}
+			}
+		}
+
+		// Build wire mapping for dependency tracking
+		dict<SigBit, Wire*> bit_to_wire_map;
+		for (auto w : module->wires())
+			for (auto bit : mi.sigmap(w))
+				bit_to_wire_map[bit] = w;
+
+		// Create unified topological sort for both cells and wires
+		TopoSort<IdString, RTLIL::sort_by_id_str> unified_toposort;
+		
+		// Add all cells and processable wires as nodes
+		for (auto cell : module->selected_cells())
+			unified_toposort.node(cell->name);
 		for (auto w : module->selected_wires())
-		{
+			unified_toposort.node(w->name);
+		
+		// Build edges between cells and wires based on signal flow
+		// In topological sort: edge(A, B) means A should be processed before B
+		for (auto cell : module->selected_cells()) {
+			for (auto &conn : cell->connections()) {
+				bool is_output = cell->output(conn.first);
+				for (auto bit : mi.sigmap(conn.second)) {
+					Wire *wire = bit_to_wire_map.count(bit) ? bit_to_wire_map[bit] : nullptr;
+					if (!wire || !unified_toposort.has_node(wire->name)) continue;
+					
+					if (is_output) {
+						// Cell drives wire: process cell before wire
+						// Cell reduction may affect wire, so cell -> wire
+						unified_toposort.edge(cell->name, wire->name);
+					} else {
+						// Wire drives cell: process wire before cell
+						// Wire reduction may affect cell, so wire -> cell
+						unified_toposort.edge(wire->name, cell->name);
+					}
+				}
+			}
+		}
+		unified_toposort.analyze_loops = false;
+		unified_toposort.sort();
+
+		// Process cells and wires together in unified topological order
+		for (auto name : unified_toposort.sorted) {
+			Cell *c = module->cell(name);
+			Wire *w = module->wire(name);
+			
+			if (c && module->selected(c)) {
+				run_cell(c);
+				continue;
+			}
+
+			if (!(w && module->selected(w)))
+				continue;
+
 			int unused_top_bits = 0;
 
 			if (w->port_id > 0 || count_nontrivial_wire_attrs(w) > 0)
@@ -514,6 +560,7 @@ struct WreduceWorker
 			Wire *nw = module->addWire(module->uniquify(IdString(w->name.str() + "_wreduce")), GetSize(w) - unused_top_bits);
 			module->connect(nw, SigSpec(w).extract(0, GetSize(nw)));
 			module->swap_names(w, nw);
+			mi.reload_module(); // TODO: SILIMATE: CAN WE SPEED THIS UP?
 		}
 	}
 };
diff --git a/passes/silimate/opt_balance_tree.cc b/passes/silimate/opt_balance_tree.cc
index b575b9172..f40d75b46 100644
--- a/passes/silimate/opt_balance_tree.cc
+++ b/passes/silimate/opt_balance_tree.cc
@@ -34,317 +34,303 @@ struct OptBalanceTreeWorker {
 	// Counts of each cell type that are getting balanced
 	dict<IdString, int> cell_count;
 
-	// Cells to remove
-	pool<Cell*> remove_cells;
-
-	// Signal chain data structures
-	dict<SigSpec, Cell*> sig_chain_next;
-	dict<SigSpec, Cell*> sig_chain_prev;
-	pool<SigBit> sigbit_with_non_chain_users;
-	pool<Cell*> chain_start_cells;
-	pool<Cell*> candidate_cells;
-
-	void make_sig_chain_next_prev(IdString cell_type) {
-		// Mark all wires with keep attribute or output port as having non-chain users
-		for (auto wire : module->wires()) {
-			if (wire->get_bool_attribute(ID::keep) || wire->port_output) {
-				for (auto bit : sigmap(wire))
-					sigbit_with_non_chain_users.insert(bit);
-			}
-		}
-
-		// Iterate over all cells in module
-		for (auto cell : module->cells()) {
-			// If cell matches and not marked as keep
-			if (cell->type == cell_type && !cell->get_bool_attribute(ID::keep)) {
-				// Get signals for cell ports
-				SigSpec a_sig = sigmap(cell->getPort(ID::A));
-				SigSpec b_sig = sigmap(cell->getPort(ID::B));
-				SigSpec y_sig = sigmap(cell->getPort(ID::Y));
-   
-	 			// If a_sig already has a chain user, mark its bits as having non-chain users
-				if (sig_chain_next.count(a_sig))
-					for (auto a_bit : a_sig.bits())
-						sigbit_with_non_chain_users.insert(a_bit);
-				// Otherwise, mark cell as the next in the chain relative to a_sig
-				else {
-					sig_chain_next[a_sig] = cell;
-				}
-
-				if (!b_sig.empty()) {
-					// If b_sig already has a chain user, mark its bits as having non-chain users
-					if (sig_chain_next.count(b_sig))
-						for (auto b_bit : b_sig.bits())
-							sigbit_with_non_chain_users.insert(b_bit);
-					// Otherwise, mark cell as the next in the chain relative to b_sig
-					else {
-						sig_chain_next[b_sig] = cell;
-					}
-				}
-				
-				// Add cell as candidate
-				candidate_cells.insert(cell);
-
-				// Mark cell as the previous in the chain relative to y_sig
-				sig_chain_prev[y_sig] = cell;
-				for (auto bit : y_sig.bits())
-					sig_chain_prev[bit] = cell;
-			}
-			// If cell is not matching type, mark all cell input signals as being non-chain users
-			else {
-				for (auto conn : cell->connections())
-					if (cell->input(conn.first))
-						for (auto bit : sigmap(conn.second))
-							sigbit_with_non_chain_users.insert(bit);
-			}
-		}
+	// 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();
 	}
 
-	void find_chain_start_cells() {
-		for (auto cell : candidate_cells) {
-			// Log candidate cell
-			log_debug("Considering %s (%s)\n", log_id(cell), log_id(cell->type));
+	// 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();
 
-			// Get signals for cell ports
-			SigSpec a_sig = sigmap(cell->getPort(ID::A));
-			SigSpec b_sig = sigmap(cell->getPort(ID::B));
-			SigSpec prev_sig = sig_chain_prev.count(a_sig) ? a_sig : b_sig;
+		// 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_ID2_SUFFIX("tree"), cell_type);
 			
-			// This is a start cell if there was no previous cell in the chain for a_sig or b_sig
-			if (sig_chain_prev.count(a_sig) + sig_chain_prev.count(b_sig) != 1) {
-				chain_start_cells.insert(cell);
-				continue;
-			}
-
-			// If any bits in previous cell signal have non-chain users, this is a start cell
-			for (auto bit : prev_sig.bits())
-				if (sigbit_with_non_chain_users.count(bit)) {
-					chain_start_cells.insert(cell);
-					continue;
-				}
+			// 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_ID2_SUFFIX("tree_out"), 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;
 		}
-	}
-
-	vector<Cell*> create_chain(Cell *start_cell) {
-		// Chain of cells
-		vector<Cell*> chain;
-
-		// Current cell
-		Cell *c = start_cell;
-
-		// Iterate over cells and add to chain
-		while (c != nullptr) {
-			chain.push_back(c);
-
-			SigSpec y_sig = sigmap(c->getPort(ID::Y));
-
-			if (sig_chain_next.count(y_sig) == 0)
-				break;
-
-			c = sig_chain_next.at(y_sig);
-			if (chain_start_cells.count(c) != 0)
-				break;
-		}
-
-		// Return chain of cells
-		return chain;
-	}
-
-	void wreduce(Cell *cell) {
-		// If cell is arithmetic, remove leading zeros from inputs, then clean up outputs
-		if (cell->type.in(ID($add), ID($mul))) {
-			// Remove leading zeros from inputs
-			for (auto inport : {ID::A, ID::B}) {
-				// Record number of bits removed
-				int bits_removed = 0;
-				IdString inport_signed = (inport == ID::A) ? ID::A_SIGNED : ID::B_SIGNED;
-				IdString inport_width = (inport == ID::A) ? ID::A_WIDTH : ID::B_WIDTH;
-				SigSpec inport_sig = sigmap(cell->getPort(inport));
-				cell->unsetPort(inport);
-				if (cell->getParam((inport == ID::A) ? ID::A_SIGNED : ID::B_SIGNED).as_bool()) {
-					while (GetSize(inport_sig) > 1 && inport_sig[GetSize(inport_sig)-1] == State::S0 && inport_sig[GetSize(inport_sig)-2] == State::S0) {
-						inport_sig.remove(GetSize(inport_sig)-1, 1);
-						bits_removed++;
-					}
-				} else {
-					while (GetSize(inport_sig) > 0 && inport_sig[GetSize(inport_sig)-1] == State::S0) {
-						inport_sig.remove(GetSize(inport_sig)-1, 1);
-						bits_removed++;
-					}
-				}
-				cell->setPort(inport, inport_sig);
-				cell->setParam(inport_width, GetSize(inport_sig));
-				log_debug("Width reduced %s/%s by %d bits\n", log_id(cell), log_id(inport), bits_removed);
-			}
-
-			// Record number of bits removed from output
-			int bits_removed = 0;
-
-			// Remove unnecessary bits from output
-			SigSpec y_sig = sigmap(cell->getPort(ID::Y));
-			cell->unsetPort(ID::Y);
-			int width;
-			if (cell->type == ID($add))
-				width = std::max(cell->getParam(ID::A_WIDTH).as_int(), cell->getParam(ID::B_WIDTH).as_int()) + 1;
-			else if (cell->type == ID($mul))
-				width = cell->getParam(ID::A_WIDTH).as_int() + cell->getParam(ID::B_WIDTH).as_int();
-			else log_abort();
-			for (int i = GetSize(y_sig) - 1; i >= width; i--) {
-				module->connect(y_sig[i], State::S0);
-				y_sig.remove(i, 1);
-				bits_removed++;
-			}
-			cell->setPort(ID::Y, y_sig);
-			cell->setParam(ID::Y_WIDTH, GetSize(y_sig));
-			log_debug("Width reduced %s/Y by %d bits\n", log_id(cell), bits_removed);
-		}
-
-		cell->fixup_parameters();
-	}
-
-	bool process_chain(vector<Cell*> &chain) {
-		// If chain size is less than 3, no balancing needed
-		if (GetSize(chain) < 3)
-			return false;
-
-		// Get mid, midnext (at index mid+1) and end of chain
-		Cell *mid_cell = chain[GetSize(chain) / 2];
-		Cell *cell = mid_cell; // SILIMATE: Set cell to mid_cell for better naming
-		Cell *midnext_cell = chain[GetSize(chain) / 2 + 1];
-		Cell *end_cell = chain.back();
-		log_debug("Balancing chain of %d cells: mid=%s, midnext=%s, endcell=%s\n",
-		    GetSize(chain), log_id(mid_cell), log_id(midnext_cell), log_id(end_cell));
-
-		// Get mid signals
-		SigSpec mid_a_sig = sigmap(mid_cell->getPort(ID::A));
-		SigSpec mid_b_sig = sigmap(mid_cell->getPort(ID::B));
-		SigSpec mid_non_chain_sig = sig_chain_prev.count(mid_a_sig) ? mid_b_sig : mid_a_sig;
-		IdString mid_non_chain_port = sig_chain_prev.count(mid_a_sig) ? ID::B : ID::A;
-
-		// Get midnext signals
-		SigSpec midnext_a_sig = sigmap(midnext_cell->getPort(ID::A));
-		SigSpec midnext_b_sig = sigmap(midnext_cell->getPort(ID::B));
-		IdString midnext_chain_port = sig_chain_prev.count(midnext_a_sig) ? ID::A : ID::B;
-
-		// Get output signal
-		SigSpec end_y_sig = sigmap(end_cell->getPort(ID::Y));
-
-		// Create new mid wire
-		Wire *mid_wire = module->addWire(NEW_ID2_SUFFIX("mid"), GetSize(end_y_sig)); // SILIMATE: Improve the naming
-
-		// Perform rotation
-		mid_cell->setPort(mid_non_chain_port, mid_wire);
-		mid_cell->setPort(ID::Y, end_y_sig);
-		midnext_cell->setPort(midnext_chain_port, mid_non_chain_sig);
-		end_cell->setPort(ID::Y, mid_wire);
-
-		// Recreate sigmap
-		sigmap.set(module);
-
-		// Get subtrees
-		vector<Cell*> left_chain(chain.begin(), chain.begin() + GetSize(chain) / 2);
-		vector<Cell*> right_chain(chain.begin() + GetSize(chain) / 2 + 1, chain.end());
-
-		// Recurse on subtrees
-		process_chain(left_chain);
-		process_chain(right_chain);
 		
-		// Width reduce left subtree
-		for (auto c : left_chain)
-			wreduce(c);
+		// 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());
 		
-		// Width reduce right subtree
-		for (auto c : right_chain)
-			wreduce(c);
-
-		// Recreate sigmap
-		sigmap.set(module);
-
-		// Width reduce mid cell
-		wreduce(mid_cell);
-
-		return true;
-	}
-
-	void cleanup() {
-		// Remove cells
-		for (auto cell : remove_cells)
-			module->remove(cell);
-
-		// Fix ports
-		module->fixup_ports();
-
-		// Clear data structures
-		remove_cells.clear();
-		sig_chain_next.clear();
-		sig_chain_prev.clear();
-		sigbit_with_non_chain_users.clear();
-		chain_start_cells.clear();
-		candidate_cells.clear();
+		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_ID2_SUFFIX("tree"), 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_ID2_SUFFIX("tree_out"), 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) {
-			// Find chains of ops
-			make_sig_chain_next_prev(cell_type);
-			find_chain_start_cells();
+			// 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;
 
-			// 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);
-				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"));
+					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);
 					}
-					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"));
-							}
-						}
+				}
+			}
+
+			// 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);
 					}
 				}
 
-				cell_count[cell_type] += GetSize(chain);
-			}
+			// 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;
 
-			// Clean up
-			cleanup();
+				// 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 (auto 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("  Creating tree for %d sources and %d inner cells...\n", 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());
+						log("  Connecting %s to %s\n", log_signal(head_output), log_signal(tree_output));
+						log_flush();
+						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/$xnor/$add/$mul cascades to trees") { }
+	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/$xnor/$add/$mul cells into\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");
 	}
 	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++) {
-			// No arguments yet
+			if (args[argidx] == "-arith") {
+				cell_types = {ID($add), ID($mul)};
+				continue;
+			}
 			break;
 		}
 		extra_args(args, argidx, design);
 
 		// Count of all cells that were packed
 		dict<IdString, int> cell_count;
-		const vector<IdString> cell_types = {ID($and), ID($or), ID($xor), ID($xnor), ID($add), ID($mul)};
 		for (auto module : design->selected_modules()) {
 			OptBalanceTreeWorker worker(module, cell_types);
 			for (auto cell : worker.cell_count) {
@@ -355,6 +341,9 @@ struct OptBalanceTreePass : public Pass {
 		// 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;
 
diff --git a/tests/alumacc/macc_infer_n_unmap.ys b/tests/alumacc/macc_infer_n_unmap.ys
index 569511b64..b05a5fe34 100644
--- a/tests/alumacc/macc_infer_n_unmap.ys
+++ b/tests/alumacc/macc_infer_n_unmap.ys
@@ -11,7 +11,7 @@ prep
 design -save gold
 alumacc
 opt_clean
-select -assert-count 1 t:$macc_v2
+select -assert-count 2 t:$macc_v2
 maccmap -unmap
 design -copy-from gold -as gold gate
 equiv_make gold gate equiv
diff --git a/tests/silimate/opt_balance_tree.ys b/tests/silimate/opt_balance_tree.ys
index aa42426eb..508f5fc24 100644
--- a/tests/silimate/opt_balance_tree.ys
+++ b/tests/silimate/opt_balance_tree.ys
@@ -1,4 +1,5 @@
-log -header "Should be turned into a tree"
+# Test 1
+log -header "Simple AND chain"
 log -push
 design -reset
 read_verilog <<EOF
@@ -16,30 +17,33 @@ 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
-select -set driven_by_a i:a %co %co
-select -set and_a_cell t:$and @driven_by_a %i
+# 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
 
-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 driven_by_c i:c %co %co
-select -set and_c_cell t:$and @driven_by_c %i
-
-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
+# 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
 
-log -header "Should not be turned into a tree"
+
+
+# Test 2
+log -header "AND chain with intermediate outputs"
 log -push
 design -reset
 read_verilog <<EOF
@@ -61,20 +65,50 @@ check -assert
 
 # Check equivalence after opt_balance_tree
 equiv_opt -assert opt_balance_tree
-
 design -load postopt
 
-# Checks if y is still wired up to the correct gate
-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
-select -assert-count 1 @inputs i:c %i
+# 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
 
-log -header "With a cell"
+
+
+# Test 3
+log -header "AND chain with intermediate outputs and one inverter"
 log -push
 design -reset
 read_verilog <<EOF
@@ -98,21 +132,52 @@ 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 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
+# 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
 
-log -header "Word out port"
+
+
+# Test 4
+log -header "AND chain with intermediate outputs to a word out port"
 log -push
 design -reset
 read_verilog <<EOF
@@ -132,20 +197,47 @@ 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 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
+# 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
 
-log -header "Fanout going to multiple outputs"
+
+
+# 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
@@ -168,111 +260,1097 @@ check -assert
 
 # Check equivalence after opt_balance_tree
 equiv_opt -assert opt_balance_tree
-
 design -load postopt
-select -assert-count 3 t:$and
+
+# 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
 
-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
-
-# 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
-
-log -header "Interesting tree situation"
+# 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 & i;
+  assign j = e & f & g & h;
 
-  assign x = i & j;
+  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
 
-# Checks if inputs to and gates has been rewired
+# Assert that there are 11 AND gates
+select t:$and -assert-count 11
 
-# a and b
-select -set driven_by_a i:a %co %co
-select -set and_a_cell t:$and @driven_by_a %i
-
-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 driven_by_c i:c %co %co
-select -set and_c_cell t:$and @driven_by_c %i
-
-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 driven_by_e i:e %co %co
-select -set and_e_cell t:$and @driven_by_e %i
-
-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 driven_by_g i:g %co %co
-select -set and_g_cell t:$and @driven_by_g %i
-
-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
+# 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
\ No newline at end of file
+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