From bdb5e0f29c3e913fb4e701317105363064b9a7d3 Mon Sep 17 00:00:00 2001
From: Eddie Hung <eddie@fpgeh.com>
Date: Wed, 11 Sep 2019 13:36:37 -0700
Subject: [PATCH] Cope with presence of reset muxes too

---
 passes/pmgen/peepopt_dffmux.pmg | 29 ++++++++++++++++++++----
 tests/various/peepopt.ys        | 39 +++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/passes/pmgen/peepopt_dffmux.pmg b/passes/pmgen/peepopt_dffmux.pmg
index 9b4fef76f..60a708616 100644
--- a/passes/pmgen/peepopt_dffmux.pmg
+++ b/passes/pmgen/peepopt_dffmux.pmg
@@ -1,16 +1,34 @@
 pattern dffmux
 
-state <IdString> cemuxAB
+state <IdString> cemuxAB rstmuxBA
+state <SigSpec> sigD
 
 match dff
 	select dff->type == $dff
 	select GetSize(port(dff, \D)) > 1
 endmatch
 
+match rstmux
+	select rstmux->type == $mux
+	select GetSize(port(rstmux, \Y)) > 1
+	index <SigSpec> port(rstmux, \Y) === port(dff, \D)
+	choice <IdString> BA {\B, \A}
+	select port(rstmux, BA).is_fully_const()
+	set rstmuxBA BA
+	optional
+endmatch
+
+code sigD
+	if (rstmux)
+		sigD = port(rstmux, rstmuxBA == \B ? \A : \B);
+	else
+		sigD = port(dff, \D);
+endcode
+
 match cemux
 	select cemux->type == $mux
 	select GetSize(port(cemux, \Y)) > 1
-	index <SigSpec> port(cemux, \Y) === port(dff, \D)
+	index <SigSpec> port(cemux, \Y) === sigD
 	choice <IdString> AB {\A, \B}
 	index <SigSpec> port(cemux, AB) === port(dff, \Q)
 	set cemuxAB AB
@@ -19,6 +37,9 @@ endmatch
 code
 	SigSpec &D = cemux->connections_.at(cemuxAB == \A ? \B : \A);
 	SigSpec &Q = dff->connections_.at(\Q);
+	Const rst;
+	if (rstmux)
+		rst = port(rstmux, rstmuxBA).as_const();
 	int width = GetSize(D);
 
 	if (D[width-1] == D[width-2]) {
@@ -30,12 +51,12 @@ code
 		for (i = width-1; i >= 2; i--) {
 			if (!is_signed) {
 				module->connect(Q[i], sign);
-				if (D[i-1] != 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)
+				if (D[i-2] != sign || (rst.size() && rst[i-1] != rst[width-1]))
 					break;
 			}
 		}
diff --git a/tests/various/peepopt.ys b/tests/various/peepopt.ys
index 8dce679ff..886c8cd9d 100644
--- a/tests/various/peepopt.ys
+++ b/tests/various/peepopt.ys
@@ -110,3 +110,42 @@ design -load postopt
 select -assert-count 1 t:$dff r:WIDTH=5 %i
 select -assert-count 1 t:$mux r:WIDTH=5 %i
 select -assert-count 0 t:$dff t:$mux %% t:* %D
+
+####################
+
+design -reset
+read_verilog <<EOT
+module peepopt_dffmuxext_unsigned_rst(input clk, ce, rst, input [1:0] i, output reg [3:0] o);
+    always @(posedge clk) if (rst) o <= 0; else if (ce) o <= i;
+endmodule
+EOT
+
+proc
+equiv_opt -assert peepopt
+design -load postopt
+wreduce
+select -assert-count 1 t:$dff r:WIDTH=2 %i
+select -assert-count 2 t:$mux
+select -assert-count 2 t:$mux r:WIDTH=2 %i
+select -assert-count 0 t:$dff t:$mux %% t:* %D
+
+####################
+
+design -reset
+read_verilog <<EOT
+module peepopt_dffmuxext_signed_rst(input clk, ce, rstn, input signed [1:0] i, output reg signed [3:0] o);
+    always @(posedge clk) begin
+        if (ce) o <= i;
+        if (!rstn) o <= 4'b1111;
+    end
+endmodule
+EOT
+
+proc
+equiv_opt -assert peepopt
+design -load postopt
+wreduce
+select -assert-count 1 t:$dff r:WIDTH=2 %i
+select -assert-count 2 t:$mux
+select -assert-count 2 t:$mux r:WIDTH=2 %i
+select -assert-count 0 t:$logic_not t:$dff t:$mux %% t:* %D
-- 
2.30.2