diff --git a/passes/opt/opt_balance_tree.cc b/passes/opt/opt_balance_tree.cc
index 915a6fb39..aa1a44de2 100644
--- a/passes/opt/opt_balance_tree.cc
+++ b/passes/opt/opt_balance_tree.cc
@@ -21,7 +21,7 @@
 
 #include "kernel/yosys.h"
 #include "kernel/sigtools.h"
-#include <deque>
+#include "kernel/utils.h"
 
 USING_YOSYS_NAMESPACE
 PRIVATE_NAMESPACE_BEGIN
@@ -35,6 +35,11 @@ struct OptBalanceTreeWorker {
 	// Counts of each cell type that are getting balanced
 	dict<IdString, int> cell_count;
 
+	// Per cell type netlist indexes, rebuilt for each balanced cell type
+	dict<SigSpec, Cell*> sig_to_driver;
+	pool<SigSpec> input_port_sigs;
+	pool<Cell*> consumed_cells;
+
 	// 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)
@@ -64,6 +69,18 @@ struct OptBalanceTreeWorker {
 		return y_width >= natural_width;
 	}
 
+	// Get the driver of a cell input port if it continues the chain, else nullptr
+	Cell *chain_driver(Cell *cell, IdString port, IdString cell_type) {
+		auto sig = sigmap(cell->getPort(port));
+		Cell *drv = sig_to_driver[sig];
+		if (!drv || !is_right_type(drv, cell_type))
+			return nullptr;
+		for (auto bit : sig)
+			if (input_port_sigs.count(bit) && !consumed_cells.count(drv))
+				return nullptr;
+		return drv;
+	}
+
 	// 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
@@ -142,7 +159,7 @@ struct OptBalanceTreeWorker {
 		// 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;
+			sig_to_driver.clear();
 			dict<SigSpec, pool<Cell*>> sig_to_sink;
 			for (auto cell : module->selected_cells())
 			{
@@ -162,7 +179,7 @@ struct OptBalanceTreeWorker {
 			}
 
 			// Need to check if any wires connect to module ports
-			pool<SigSpec> input_port_sigs;
+			input_port_sigs.clear();
 			pool<SigSpec> output_port_sigs;
 			for (auto wire : module->selected_wires())
 				if (wire->port_input || wire->port_output) {
@@ -176,7 +193,7 @@ struct OptBalanceTreeWorker {
 				}
 
 			// Actual logic starts here
-			pool<Cell*> consumed_cells;
+			consumed_cells.clear();
 			for (auto cell : module->selected_cells())
 			{
 				// If consumed or not the correct type, skip
@@ -189,11 +206,15 @@ struct OptBalanceTreeWorker {
 				pool<Cell*> sinks;
 				pool<Cell*> current_loads = sig_to_sink[y];
 				pool<Cell*> next_loads;
+				pool<Cell*> visited_loads;
 				while (!current_loads.empty())
 				{
 					// Find each sink and see what they are
 					for (auto x : current_loads)
 					{
+						if (!visited_loads.insert(x).second)
+							continue;
+
 						// 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))
@@ -246,31 +267,45 @@ struct OptBalanceTreeWorker {
 					if (consumed_cells.count(head_cell))
 						continue;
 
-					// Get sources of the chain
+					// Collect the chain cone into a topological sort
+					TopoSort<Cell*, IdString::compare_ptr_by_name<Cell>> toposort;
+					toposort.analyze_loops = false;
+					toposort.node(head_cell);
+					vector<Cell*> queue = {head_cell};
+					while (!queue.empty())
+					{
+						Cell *x = queue.back();
+						queue.pop_back();
+						for (IdString port: {ID::A, ID::B})
+							if (Cell *drv = chain_driver(x, port, cell_type)) {
+								if (!toposort.has_node(drv))
+									queue.push_back(drv);
+								toposort.edge(drv, x);
+							}
+					}
+
+					// Abandon chains containing combinational loops, since
+					// rebalancing them is not sound (and would not terminate)
+					if (!toposort.sort())
+						continue;
+
+					// Get sources of the chain: process cells from head to
+					// drivers, counting the paths leading back to the head so
+					// reconvergent sources are counted with multiplicity
 					dict<SigSpec, int> sources;
 					dict<SigSpec, bool> signeds;
-					int inner_cells = 0;
-					std::deque<Cell*> bfs_queue = {head_cell};
-					while (bfs_queue.size())
+					int inner_cells = GetSize(toposort.sorted) - 1;
+					dict<Cell*, int> reach;
+					reach[head_cell] = 1;
+					for (int i = GetSize(toposort.sorted); i-- > 0; )
 					{
-						Cell* x = bfs_queue.front();
-						bfs_queue.pop_front();
-
+						Cell* x = toposort.sorted[i];
 						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);
+							if (Cell *drv = chain_driver(x, port, cell_type)) {
+								reach[drv] += reach[x];
 							} else {
-								sources[sig]++;
+								auto sig = sigmap(x->getPort(port));
+								sources[sig] += reach[x];
 								signeds[sig] = x->getParam(port == ID::A ? ID::A_SIGNED : ID::B_SIGNED).as_bool();
 							}
 						}
diff --git a/tests/opt/opt_balance_tree.ys b/tests/opt/opt_balance_tree.ys
index 030449731..dd95576c3 100644
--- a/tests/opt/opt_balance_tree.ys
+++ b/tests/opt/opt_balance_tree.ys
@@ -1244,3 +1244,96 @@ design -load postopt
 
 design -reset
 log -pop
+
+
+# Test 31
+log -header "Combinational loop is left untouched"
+log -push
+design -reset
+read_verilog <<EOF
+module top(input a, output b);
+  assign b = a | b;
+endmodule
+EOF
+hierarchy -auto-top
+proc
+select t:$or -assert-count 1
+opt_balance_tree
+select t:$or -assert-count 1
+
+design -reset
+log -pop
+
+
+# Test 32
+log -header "Multi-cell combinational loop is left untouched"
+log -push
+design -reset
+read_verilog <<EOF
+module top(input a, input c, output b);
+  wire w;
+  assign w = a | b;
+  assign b = w | c;
+endmodule
+EOF
+hierarchy -auto-top
+proc
+select t:$or -assert-count 2
+opt_balance_tree
+select t:$or -assert-count 2
+
+design -reset
+log -pop
+
+
+# Test 33
+log -header "Clean chain beside a combinational loop is still balanced"
+log -push
+design -reset
+read_verilog <<EOF
+module top(input a, input [7:0] v, output b, output x);
+  assign b = a | b;
+  assign x = v[0] | v[1] | v[2] | v[3] | v[4] | v[5] | v[6] | v[7];
+endmodule
+EOF
+hierarchy -auto-top
+proc
+select t:$or -assert-count 8
+opt_balance_tree
+select t:$or -assert-count 8
+
+# Loop cell is untouched
+select o:b %ci2 t:$or %i -assert-count 1
+
+# Chain was balanced to depth 3
+select i:v %co6 o:x %i -assert-count 1
+
+design -reset
+log -pop
+
+
+# Test 34
+log -header "Loop through a different cell type does not block balancing"
+log -push
+design -reset
+read_verilog <<EOF
+module top(input a, c, input [3:0] v, output b, output x);
+  wire w;
+  assign w = b & a;
+  assign b = w | c;
+  assign x = v[0] | v[1] | v[2] | v[3];
+endmodule
+EOF
+hierarchy -auto-top
+proc
+select t:$and -assert-count 1
+select t:$or -assert-count 4
+opt_balance_tree
+select t:$and -assert-count 1
+select t:$or -assert-count 4
+
+# Chain was balanced to depth 2
+select i:v %co4 o:x %i -assert-count 1
+
+design -reset
+log -pop