Refactor peepopt_dffmux and be sensitive to \init when trimming

1 files changed, 63 insertions, 32 deletions

diff --git a/passes/pmgen/peepopt_dffmux.pmg b/passes/pmgen/peepopt_dffmux.pmg
index c88a52226..2ec504cb4 100644
--- a/passes/pmgen/peepopt_dffmux.pmg
+++ b/passes/pmgen/peepopt_dffmux.pmg
@@ -8,21 +8,23 @@ match dff
 	select GetSize(port(dff, \D)) > 1
 endmatch
 
+code sigD
+	sigD = port(dff, \D);
+endcode
+
 match rstmux
 	select rstmux->type == $mux
 	select GetSize(port(rstmux, \Y)) > 1
-	index <SigSpec> port(rstmux, \Y) === port(dff, \D)
+	index <SigSpec> port(rstmux, \Y) === sigD
 	choice <IdString> BA {\B, \A}
 	select port(rstmux, BA).is_fully_const()
 	set rstmuxBA BA
-	optional
+	semioptional
 endmatch
 
 code sigD
 	if (rstmux)
 		sigD = port(rstmux, rstmuxBA == \B ? \A : \B);
-	else
-		sigD = port(dff, \D);
 endcode
 
 match cemux
@@ -32,45 +34,70 @@ match cemux
 	choice <IdString> AB {\A, \B}
 	index <SigSpec> port(cemux, AB) === port(dff, \Q)
 	set cemuxAB AB
+	semioptional
 endmatch
 
 code
-	SigSpec D = port(cemux, cemuxAB == \A ? \B : \A);
-	SigSpec Q = port(dff, \Q);
+	if (!cemux && !rstmux)
+		reject;
+endcode
+
+code
 	Const rst;
-	if (rstmux)
+	SigSpec D;
+	if (cemux) {
+		D = port(cemux, cemuxAB == \A ? \B : \A);
+		if (rstmux)
+			rst = port(rstmux, rstmuxBA).as_const();
+		else
+			rst = Const(State::Sx, GetSize(D));
+	}
+	else {
+		log_assert(rstmux);
+		D = port(rstmux, rstmuxBA  == \B ? \A : \B);
 		rst = port(rstmux, rstmuxBA).as_const();
+	}
+	SigSpec Q = port(dff, \Q);
 	int width = GetSize(D);
 
-	SigSpec &ceA = cemux->connections_.at(\A);
-	SigSpec &ceB = cemux->connections_.at(\B);
-	SigSpec &ceY = cemux->connections_.at(\Y);
 	SigSpec &dffD = dff->connections_.at(\D);
 	SigSpec &dffQ = dff->connections_.at(\Q);
+	Const init;
+	for (const auto &b : Q) {
+		auto it = b.wire->attributes.find(\init);
+		init.bits.push_back(it == b.wire->attributes.end() ? State::Sx : it->second[b.offset]);
+	}
 
-	if (D[width-1] == D[width-2]) {
-		did_something = true;
+	auto cmpx = [=](State lhs, State rhs) {
+		if (lhs == State::Sx || rhs == State::Sx)
+			return true;
+		return lhs == rhs;
+	};
 
-		SigBit sign = D[width-1];
-		bool is_signed = sign.wire;
-		int i;
-		for (i = width-1; i >= 2; i--) {
-			if (!is_signed) {
-				module->connect(Q[i], sign);
-				if (D[i-1] != sign || (rst.size() && rst[i-1] != rst[width-1]))
-					break;
-			}
-			else {
-				module->connect(Q[i], Q[i-1]);
-				if (D[i-2] != sign || (rst.size() && rst[i-1] != rst[width-1]))
-					break;
-			}
+	int i = width;
+	while (i > 2) {
+		i--;
+		if (D[i] != D[i-1])
+			break;
+		if (!cmpx(rst[i], rst[i-1]))
+			break;
+		if (!cmpx(init[i], init[i-1]))
+			break;
+		if (!cmpx(rst[i], init[i]))
+			break;
+		module->connect(Q[i], Q[i-1]);
+		did_something = true;
+	}
+	if (i < width-1) {
+		if (cemux) {
+			SigSpec &ceA = cemux->connections_.at(\A);
+			SigSpec &ceB = cemux->connections_.at(\B);
+			SigSpec &ceY = cemux->connections_.at(\Y);
+			ceA.remove(i, width-i);
+			ceB.remove(i, width-i);
+			ceY.remove(i, width-i);
+			cemux->fixup_parameters();
 		}
-
-		ceA.remove(i, width-i);
-		ceB.remove(i, width-i);
-		ceY.remove(i, width-i);
-		cemux->fixup_parameters();
 		dffD.remove(i, width-i);
 		dffQ.remove(i, width-i);
 		dff->fixup_parameters();
@@ -78,7 +105,11 @@ code
 		log("dffcemux pattern in %s: dff=%s, cemux=%s; removed top %d bits.\n", log_id(module), log_id(dff), log_id(cemux), width-i);
 		accept;
 	}
-	else {
+	else if (cemux) {
+		SigSpec &ceA = cemux->connections_.at(\A);
+		SigSpec &ceB = cemux->connections_.at(\B);
+		SigSpec &ceY = cemux->connections_.at(\Y);
+
 		int count = 0;
 		for (int i = width-1; i >= 0; i--) {
 			if (D[i].wire)