1 files changed, 15 insertions, 7 deletions

diff --git a/passes/memory/memory_dff.cc b/passes/memory/memory_dff.cc
index 32df1917..32b97f27 100644
--- a/passes/memory/memory_dff.cc
+++ b/passes/memory/memory_dff.cc
@@ -17,6 +17,7 @@
  *
  */
 
+#include <algorithm>
 #include "kernel/yosys.h"
 #include "kernel/sigtools.h"
 
@@ -41,7 +42,7 @@ struct MemoryDffWorker
 			if (wire->attributes.count("\\init") == 0)
 				continue;
 			SigSpec sig = sigmap(wire);
-			Const initval = wire->attributes.count("\\init");
+			Const initval = wire->attributes.at("\\init");
 			for (int i = 0; i < GetSize(sig) && i < GetSize(initval); i++)
 				if (initval[i] == State::S0 || initval[i] == State::S1)
 					init_bits.insert(sig[i]);
@@ -182,20 +183,27 @@ struct MemoryDffWorker
 
 		if (mux_cells_a.count(sig_data) || mux_cells_b.count(sig_data))
 		{
-			bool enable_invert = mux_cells_a.count(sig_data) != 0;
-			Cell *mux = enable_invert ? mux_cells_a.at(sig_data) : mux_cells_b.at(sig_data);
-			SigSpec check_q = sigmap(mux->getPort(enable_invert ? "\\B" : "\\A"));
+			RTLIL::SigSpec en;
+			std::vector<RTLIL::SigSpec> check_q;
+
+			do {
+				bool enable_invert = mux_cells_a.count(sig_data) != 0;
+				Cell *mux = enable_invert ? mux_cells_a.at(sig_data) : mux_cells_b.at(sig_data);
+				check_q.push_back(sigmap(mux->getPort(enable_invert ? "\\B" : "\\A")));
+				sig_data = sigmap(mux->getPort("\\Y"));
+				en.append(enable_invert ? module->LogicNot(NEW_ID, mux->getPort("\\S")) : mux->getPort("\\S"));
+			} while (mux_cells_a.count(sig_data) || mux_cells_b.count(sig_data));
 
-			sig_data = sigmap(mux->getPort("\\Y"));
 			for (auto bit : sig_data)
 				if (sigbit_users_count[bit] > 1)
 					goto skip_ff_after_read_merging;
 
-			if (find_sig_before_dff(sig_data, clk_data, clk_polarity, true) && clk_data != RTLIL::SigSpec(RTLIL::State::Sx) && sig_data == check_q)
+			if (find_sig_before_dff(sig_data, clk_data, clk_polarity, true) && clk_data != RTLIL::SigSpec(RTLIL::State::Sx) &&
+					std::all_of(check_q.begin(), check_q.end(), [&](const SigSpec &cq) {return cq == sig_data; }))
 			{
 				disconnect_dff(sig_data);
 				cell->setPort("\\CLK", clk_data);
-				cell->setPort("\\EN", enable_invert ? module->LogicNot(NEW_ID, mux->getPort("\\S")) : mux->getPort("\\S"));
+				cell->setPort("\\EN", en.size() > 1 ? module->ReduceAnd(NEW_ID, en) : en);
 				cell->setPort("\\DATA", sig_data);
 				cell->parameters["\\CLK_ENABLE"] = RTLIL::Const(1);
 				cell->parameters["\\CLK_POLARITY"] = RTLIL::Const(clk_polarity);