From 42c0f9f3a6bf42b6d886047560b5f2550878d79f Mon Sep 17 00:00:00 2001
From: "Emil J. Tywoniak" <emil@tywoniak.eu>
Date: Tue, 8 Apr 2025 23:49:15 +0200
Subject: [PATCH] opt_clean: handle undriven and x-bit driven bits consistently

---
 passes/opt/opt_clean.cc  | 30 ++++++++++++++++++++++++------
 tests/opt/opt_clean_x.ys | 14 ++++++++++++++
 2 files changed, 38 insertions(+), 6 deletions(-)
 create mode 100644 tests/opt/opt_clean_x.ys

diff --git a/passes/opt/opt_clean.cc b/passes/opt/opt_clean.cc
index c37c03607..eb57132cb 100644
--- a/passes/opt/opt_clean.cc
+++ b/passes/opt/opt_clean.cc
@@ -31,6 +31,24 @@ PRIVATE_NAMESPACE_BEGIN
 
 using RTLIL::id2cstr;
 
+struct CleanerPool : SigPool
+{
+	bool check_all_def(const RTLIL::SigSpec &sig) const
+	{
+		for (auto &bit : sig) {
+			if (bit.wire != NULL && bits.count(bit) == 0)
+				return false;
+			if (bit.wire == NULL && bit.data == RTLIL::State::Sx)
+				return false;
+		}
+		return true;
+	}
+	bool check_def(const RTLIL::SigBit &bit) const
+	{
+		return (bit.wire != NULL && bits.count(bit)) || (bit.wire == NULL && bit.data != RTLIL::State::Sx);
+	}
+};
+
 struct keep_cache_t
 {
 	Design *design;
@@ -356,7 +374,7 @@ bool rmunused_module_signals(RTLIL::Module *module, bool purge_mode, bool verbos
 	// used signals pre-sigmapped
 	SigPool raw_used_signals;
 	// used signals sigmapped, ignoring drivers (we keep track of this to set `unused_bits`)
-	SigPool used_signals_nodrivers;
+	CleanerPool used_signals_nodrivers;
 
 	// gather the usage information for cells
 	for (auto &it : module->cells_) {
@@ -472,14 +490,14 @@ bool rmunused_module_signals(RTLIL::Module *module, bool purge_mode, bool verbos
 				module->connect(new_conn);
 			}
 
-			if (!used_signals_nodrivers.check_all(s2)) {
+			if (!used_signals_nodrivers.check_all_def(s2)) {
 				std::string unused_bits;
 				for (int i = 0; i < GetSize(s2); i++) {
-					if (s2[i].wire == NULL)
-						continue;
-					if (!used_signals_nodrivers.check(s2[i])) {
+					if ((s2[i].wire == NULL) && (s2[i].data != RTLIL::State::Sx))
+					continue;
+					if (!used_signals_nodrivers.check_def(s2[i])) {
 						if (!unused_bits.empty())
-							unused_bits += " ";
+						unused_bits += " ";
 						unused_bits += stringf("%d", i);
 					}
 				}
diff --git a/tests/opt/opt_clean_x.ys b/tests/opt/opt_clean_x.ys
new file mode 100644
index 000000000..68a1ad4fb
--- /dev/null
+++ b/tests/opt/opt_clean_x.ys
@@ -0,0 +1,14 @@
+read_verilog <<EOT
+module alu(
+);
+	wire [1:0] p1, p2;
+	assign p1 = 2'bx1;
+	assign p2[0] = 1'b1;
+endmodule
+EOT
+
+proc
+opt_clean
+dump
+select -assert-count 1 w:p1 a:unused_bits=1 %i
+select -assert-count 1 w:p2 a:unused_bits=1 %i