From 47c2257f82f25f75175e0a41e9ca1b37d653bf6a Mon Sep 17 00:00:00 2001
From: abhinavputhran <aputhran@andrew.cmu.edu>
Date: Sun, 8 Mar 2026 19:41:31 -0400
Subject: [PATCH] setundef: more tests! and wire selection in -init mode

---
 passes/cmds/setundef.cc                      | 10 +++++++++-
 tests/various/setundef_selection.ys          | 17 +++++++++++------
 tests/various/setundef_selection_undriven.il |  4 ++++
 3 files changed, 24 insertions(+), 7 deletions(-)
 create mode 100644 tests/various/setundef_selection_undriven.il

diff --git a/passes/cmds/setundef.cc b/passes/cmds/setundef.cc
index cdd7f3446..99a223bdc 100644
--- a/passes/cmds/setundef.cc
+++ b/passes/cmds/setundef.cc
@@ -364,11 +364,19 @@ struct SetundefPass : public Pass {
 				pool<SigBit> ffbits;
 				pool<Wire*> initwires;
 
-				for (auto cell : module->selected_cells())
+				for (auto cell : module->cells())
 				{
 					if (!cell->is_builtin_ff())
 						continue;
 
+					bool cell_selected = design->selected(module, cell);
+   	 			bool wire_selected = false;
+					for (auto bit : sigmap(cell->getPort(ID::Q)))
+						if (bit.wire && design->selected(module, bit.wire))
+							wire_selected = true;
+					if (!cell_selected && !wire_selected)
+						continue;
+
 					for (auto bit : sigmap(cell->getPort(ID::Q)))
 						ffbits.insert(bit);
 				}
diff --git a/tests/various/setundef_selection.ys b/tests/various/setundef_selection.ys
index 67df8bed9..64713a742 100644
--- a/tests/various/setundef_selection.ys
+++ b/tests/various/setundef_selection.ys
@@ -7,21 +7,26 @@ endmodule
 EOT
 setundef -zero w:a
 sat -prove a 0
+sat -enable_undef -prove b 0 -falsify
 design -reset
 
 # Test that setundef -undriven -zero respects wire selection
-read_verilog <<EOT
-module test(output wire a, output wire b);
-    assign a = 1'b0;
-endmodule
-EOT
+read_rtlil setundef_selection_undriven.il
 setundef -undriven -zero w:b
 sat -prove b 0
+sat -enable_undef -prove a 0 -falsify
 design -reset
 
 # Test that setundef -init respects cell selection: only selected FF gets init set
 read_rtlil setundef_selection_ff.il
 setundef -init -zero c:myff_a
-# only wire a should have init attribute, not b
 select -assert-count 1 w:* a:init %i
+select -assert-count 0 w:b a:init %i
+design -reset
+
+# Test that setundef -init works with wire selection
+read_rtlil setundef_selection_ff.il
+setundef -init -zero w:a
+select -assert-count 1 w:* a:init %i
+select -assert-count 0 w:b a:init %i
 design -reset
diff --git a/tests/various/setundef_selection_undriven.il b/tests/various/setundef_selection_undriven.il
new file mode 100644
index 000000000..f77809eb2
--- /dev/null
+++ b/tests/various/setundef_selection_undriven.il
@@ -0,0 +1,4 @@
+module \test
+  wire output 1 \a
+  wire output 2 \b
+end