1 files changed, 18 insertions, 10 deletions

diff --git a/passes/memory/memory_dff.cc b/passes/memory/memory_dff.cc
index 6e036397d..be4b3c100 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);
@@ -254,7 +262,7 @@ struct MemoryDffWorker
 				mux_cells_a[sigmap(cell->getPort("\\A"))] = cell;
 				mux_cells_b[sigmap(cell->getPort("\\B"))] = cell;
 			}
-			if (cell->type == "$not" || cell->type == "$_NOT_" || (cell->type == "$logic_not" && GetSize(cell->getPort("\\A")) == 1)) {
+			if (cell->type.in("$not", "$_NOT_") || (cell->type == "$logic_not" && GetSize(cell->getPort("\\A")) == 1)) {
 				SigSpec sig_a = cell->getPort("\\A");
 				SigSpec sig_y = cell->getPort("\\Y");
 				if (cell->type == "$not")
@@ -283,7 +291,7 @@ struct MemoryDffWorker
 
 struct MemoryDffPass : public Pass {
 	MemoryDffPass() : Pass("memory_dff", "merge input/output DFFs into memories") { }
-	virtual void help()
+	void help() YS_OVERRIDE
 	{
 		//   |---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|---v---|
 		log("\n");
@@ -297,7 +305,7 @@ struct MemoryDffPass : public Pass {
 		log("        do not merge registers on read ports\n");
 		log("\n");
 	}
-	virtual void execute(std::vector<std::string> args, RTLIL::Design *design)
+	void execute(std::vector<std::string> args, RTLIL::Design *design) YS_OVERRIDE
 	{
 		bool flag_wr_only = false;