From 98003430d672af05cff7d950e08a42526e766f74 Mon Sep 17 00:00:00 2001
From: =?utf8?q?Marcelina=20Ko=C5=9Bcielnicka?= <mwk@0x04.net>
Date: Sun, 8 Aug 2021 00:33:31 +0200
Subject: [PATCH] opt_merge: Use FfInitVals.

Partial #2920 fix.
---
 passes/opt/opt_merge.cc     | 35 +++++++------------------------
 tests/opt/bug2920.ys        | 42 +++++++++++++++++++++++++++++++++++++
 tests/opt/opt_merge_init.ys |  2 +-
 3 files changed, 51 insertions(+), 28 deletions(-)
 create mode 100644 tests/opt/bug2920.ys

diff --git a/passes/opt/opt_merge.cc b/passes/opt/opt_merge.cc
index f27277574..d9861f49b 100644
--- a/passes/opt/opt_merge.cc
+++ b/passes/opt/opt_merge.cc
@@ -18,6 +18,7 @@
  */
 
 #include "kernel/register.h"
+#include "kernel/ffinit.h"
 #include "kernel/sigtools.h"
 #include "kernel/log.h"
 #include "kernel/celltypes.h"
@@ -35,7 +36,7 @@ struct OptMergeWorker
 	RTLIL::Design *design;
 	RTLIL::Module *module;
 	SigMap assign_map;
-	SigMap dff_init_map;
+	FfInitVals initvals;
 	bool mode_share_all;
 
 	CellTypes ct;
@@ -121,8 +122,7 @@ struct OptMergeWorker
 				if (it.first == ID::Q && RTLIL::builtin_ff_cell_types().count(cell->type)) {
 					// For the 'Q' output of state elements,
 					//   use its (* init *) attribute value
-					for (const auto &b : dff_init_map(it.second))
-						sig.append(b.wire ? State::Sx : b);
+					sig = initvals(it.second);
 				}
 				else
 					continue;
@@ -176,12 +176,8 @@ struct OptMergeWorker
 				if (it.first == ID::Q && RTLIL::builtin_ff_cell_types().count(cell1->type)) {
 					// For the 'Q' output of state elements,
 					//   use the (* init *) attribute value
-					auto &sig1 = conn1[it.first];
-					for (const auto &b : dff_init_map(it.second))
-						sig1.append(b.wire ? State::Sx : b);
-					auto &sig2 = conn2[it.first];
-					for (const auto &b : dff_init_map(cell2->getPort(it.first)))
-						sig2.append(b.wire ? State::Sx : b);
+					conn1[it.first] = initvals(it.second);
+					conn2[it.first] = initvals(cell2->getPort(it.first));
 				}
 				else {
 					conn1[it.first] = RTLIL::SigSpec();
@@ -247,14 +243,7 @@ struct OptMergeWorker
 		log("Finding identical cells in module `%s'.\n", module->name.c_str());
 		assign_map.set(module);
 
-		dff_init_map.set(module);
-		for (auto &it : module->wires_)
-			if (it.second->attributes.count(ID::init) != 0) {
-				Const initval = it.second->attributes.at(ID::init);
-				for (int i = 0; i < GetSize(initval) && i < GetSize(it.second); i++)
-					if (initval[i] == State::S0 || initval[i] == State::S1)
-						dff_init_map.add(SigBit(it.second, i), initval[i]);
-			}
+		initvals.set(&assign_map, module);
 
 		bool did_something = true;
 		while (did_something)
@@ -296,16 +285,8 @@ struct OptMergeWorker
 								module->connect(RTLIL::SigSig(it.second, other_sig));
 								assign_map.add(it.second, other_sig);
 
-								if (it.first == ID::Q && RTLIL::builtin_ff_cell_types().count(cell->type)) {
-									for (auto c : it.second.chunks()) {
-										auto jt = c.wire->attributes.find(ID::init);
-										if (jt == c.wire->attributes.end())
-											continue;
-										for (int i = c.offset; i < c.offset + c.width; i++)
-											jt->second[i] = State::Sx;
-									}
-									dff_init_map.add(it.second, Const(State::Sx, GetSize(it.second)));
-								}
+								if (it.first == ID::Q && RTLIL::builtin_ff_cell_types().count(cell->type))
+									initvals.remove_init(it.second);
 							}
 						}
 						log_debug("    Removing %s cell `%s' from module `%s'.\n", cell->type.c_str(), cell->name.c_str(), module->name.c_str());
diff --git a/tests/opt/bug2920.ys b/tests/opt/bug2920.ys
new file mode 100644
index 000000000..a8281a73a
--- /dev/null
+++ b/tests/opt/bug2920.ys
@@ -0,0 +1,42 @@
+read_ilang <<EOT
+
+module \mod
+  wire input 1 \clk
+  attribute \init 2'00
+  wire width 2 $q1
+  attribute \init 2'00
+  wire width 2 $q2
+  wire output 2 width 4 \q
+  cell $dff $ff1
+    parameter \CLK_POLARITY 1'1
+    parameter \WIDTH 1
+    connect \CLK \clk
+    connect \D 1'0
+    connect \Q $q1 [0]
+  end
+  cell $dff $ff2
+    parameter \CLK_POLARITY 1'1
+    parameter \WIDTH 1
+    connect \CLK \clk
+    connect \D 1'0
+    connect \Q $q2 [0]
+  end
+  cell $dff $ff3
+    parameter \CLK_POLARITY 1'1
+    parameter \WIDTH 2
+    connect \CLK \clk
+    connect \D 2'00
+    connect \Q { $q1 [1] $q2 [1] }
+  end
+  connect \q [0] $q1 [0]
+  connect \q [1] $q2 [0]
+  connect \q [2] $q1 [1]
+  connect \q [3] $q2 [1]
+end
+
+EOT
+
+opt_clean
+opt_merge
+opt_dff
+opt_clean
diff --git a/tests/opt/opt_merge_init.ys b/tests/opt/opt_merge_init.ys
index 0176f09c7..20b6cabee 100644
--- a/tests/opt/opt_merge_init.ys
+++ b/tests/opt/opt_merge_init.ys
@@ -48,7 +48,7 @@ EOT
 
 opt_merge
 select -assert-count 1 t:$dff
-select -assert-count 1 a:init=2'bx1
+select -assert-count 1 a:init=2'bx1 a:init=2'b1x
 
 
 design -reset
-- 
2.30.2