From 6e99dbd351fbb7dc99118adb7c650b4eb546cfe4 Mon Sep 17 00:00:00 2001
From: "Emil J. Tywoniak" <emil@tywoniak.eu>
Date: Mon, 3 Feb 2025 18:31:33 +0100
Subject: [PATCH] fixup! abstract: -state MVP

---
 passes/cmds/abstract.cc   | 112 ++++++++++++++++++++++++++++----------
 tests/various/abstract.ys |  62 ++++++++++++++++++++-
 2 files changed, 142 insertions(+), 32 deletions(-)

diff --git a/passes/cmds/abstract.cc b/passes/cmds/abstract.cc
index 3b20fb36c..2bb9a6839 100644
--- a/passes/cmds/abstract.cc
+++ b/passes/cmds/abstract.cc
@@ -6,35 +6,79 @@ USING_YOSYS_NAMESPACE
 PRIVATE_NAMESPACE_BEGIN
 
 bool abstract_state(Module* mod, Cell* cell, Wire* enable, bool enable_pol) {
+	FfData ff(nullptr, cell);
+	if (ff.has_sr)
+		log_cmd_error("SR not supported\n");
+
+	// Normalize to simpler FF
+	ff.unmap_ce();
+	ff.unmap_srst();
+	if (ff.has_arst)
+		ff.arst_to_aload();
+
+	// Construct abstract value
+	auto anyseq = mod->Anyseq(NEW_ID, ff.width);
+
+	if (ff.has_aload) {
+		// aload := enable ? anyseq : aload
+		// ff.pol_aload = enable_pol;
+		Wire* abstracted_ad = mod->addWire(NEW_ID, ff.sig_ad.size());
+		SigSpec mux_a, mux_b;
+		if (enable_pol) {
+			mux_a = ff.sig_ad;
+			mux_b = anyseq;
+		} else {
+			mux_a = anyseq;
+			mux_b = ff.sig_ad;
+		}
+		(void)mod->addMux(NEW_ID,
+			mux_a,
+			mux_b,
+			enable,
+			abstracted_ad);
+		ff.sig_ad = abstracted_ad;
+	}
+	Wire* abstracted_d = mod->addWire(NEW_ID, ff.sig_d.size());
+	SigSpec mux_a, mux_b;
+	if (enable_pol) {
+		mux_a = ff.sig_d;
+		mux_b = anyseq;
+	} else {
+		mux_a = anyseq;
+		mux_b = ff.sig_d;
+	}
+	(void)mod->addMux(NEW_ID,
+		mux_a,
+		mux_b,
+		enable,
+		abstracted_d);
+	ff.sig_d = abstracted_d;
+	(void)ff.emit();
+	return true;
+}
+
+bool abstract_value(Module* mod, Wire* wire, Wire* enable, bool enable_pol) {
+	// (void)mod->addMux(NEW_ID,
+	// 	mux_a,
+	// 	mux_b,
+	// 	enable,
+	// 	abstracted);
+	// 	cell->setPort(ID::D, SigSpec(abstracted));
+	return false;
+}
+
+bool abstract_init(Module* mod, Cell* cell) {
 	CellTypes ct;
 	ct.setup_internals_ff();
 	if (!ct.cell_types.count(cell->type))
 		return false;
-	FfData ff(nullptr, cell);
-	// Doesn't matter if there was an enable signal already
-	// we discard it and mux with symbolic value
-	ff.has_ce = false;
-	Wire* inp = cell->getPort(ID::D).as_wire();
-	cell = ff.emit();
+	// TODO figure out memory cells?
 
-	Wire* abstracted = mod->addWire(NEW_ID, inp->width);
-	SigSpec mux_a, mux_b;
-	if (enable_pol) {
-		mux_a = inp;
-		mux_b = mod->Anyseq(NEW_ID, inp->width);
-	} else {
-		mux_a = mod->Anyseq(NEW_ID, inp->width);
-		mux_b = inp;
-	}
-	(void)mod->addMux(NEW_ID,
-	mux_a,
-	mux_b,
-	enable,
-	abstracted);
-	cell->setPort(ID::D, SigSpec(abstracted));
+	cell->unsetParam(ID::init);
 	return true;
 }
 
+
 struct AbstractPass : public Pass {
 	AbstractPass() : Pass("abstract", "extract clock gating out of flip flops") { }
 	void help() override {
@@ -84,31 +128,41 @@ struct AbstractPass : public Pass {
 		}
 		extra_args(args, argidx, design);
 
-		unsigned int changed_cells = 0;
+		unsigned int changed = 0;
 		if ((mode == State) || (mode == Value)) {
 			if (!enable_name.length())
 				log_cmd_error("Unspecified enable wire\n");
+			CellTypes ct;
+			if (mode == State)
+				ct.setup_internals_ff();
 			for (auto mod : design->selected_modules()) {
-				log("module %s\n", mod->name.c_str());
+				log_debug("module %s\n", mod->name.c_str());
 				Wire *enable_wire = mod->wire("\\" + enable_name);
 				if (!enable_wire)
 					log_cmd_error("Enable wire %s not found in module %s\n", enable_name.c_str(), mod->name.c_str());
 				if (mode == State) {
-					for (auto cell : mod->selected_cells()) {
-						log("cell %s\n", cell->name.c_str());
-						changed_cells += abstract_state(mod, cell, enable_wire, enable_pol);
-					}
+					for (auto cell : mod->selected_cells())
+						if (ct.cell_types.count(cell->type))
+							changed += abstract_state(mod, cell, enable_wire, enable_pol);
 				} else {
 					// Value
+					for (auto wire : mod->selected_wires()) {
+						changed += abstract_value(mod, wire, enable_wire, enable_pol);
+					}
 					log_cmd_error("Unsupported (TODO)\n");
 				}
 			}
+			log("Abstracted %d cells.\n", changed);
 		} else if (mode == Initial) {
-			log_cmd_error("Unsupported\n");
+			for (auto mod : design->selected_modules()) {
+				for (auto cell : mod->selected_cells()) {
+					changed += abstract_init(mod, cell);
+				}
+			}
+			log("Abstracted %d wires.\n", changed);
 		} else {
 			log_cmd_error("No mode selected, see help message\n");
 		}
-		log("Abstracted %d cell(s).\n", changed_cells);
 	}
 } AbstractPass;
 
diff --git a/tests/various/abstract.ys b/tests/various/abstract.ys
index f9ff9abc4..8536960ca 100644
--- a/tests/various/abstract.ys
+++ b/tests/various/abstract.ys
@@ -11,14 +11,14 @@ endmodule
 EOT
 proc
 # dump
+# show -prefix before_base
 abstract -state -enablen magic
 check
+# show -prefix after_base
 # dump
 
 design -reset
-
 read_verilog <<EOT
-
 module half_clock_en (CLK, E, Q);
 	input CLK;
 	input E;
@@ -28,14 +28,70 @@ module half_clock_en (CLK, E, Q);
         if (E)
             Q <= ~Q;
 endmodule
-
 EOT
+
 proc
 opt_expr
 opt_dff
 # show
 # dump
+# show -prefix before_en
 abstract -state -enablen magic
 check
+# show -prefix after_en
+
+design -reset
+read_verilog <<EOT
+module half_clock (CLK, Q);
+	input CLK;
+	output reg Q = 1'b1;
+	always @(posedge CLK)
+		Q <= ~Q;
+endmodule
+EOT
+
+proc
+opt_expr
+opt_dff
+dump
+abstract -init
+check
+dump
+
+
+design -reset
+read_verilog <<EOT
+module this_adff (CLK, ARST, D, Q, magic);
+
+parameter WIDTH = 2;
+parameter CLK_POLARITY = 1'b1;
+parameter ARST_POLARITY = 1'b1;
+parameter ARST_VALUE = 0;
+
+input CLK, ARST, magic;
+input [WIDTH-1:0] D;
+output reg [WIDTH-1:0] Q;
+wire pos_clk = CLK == CLK_POLARITY;
+wire pos_arst = ARST == ARST_POLARITY;
+
+always @(posedge pos_clk, posedge pos_arst) begin
+	if (pos_arst)
+		Q <= ARST_VALUE;
+	else
+		Q <= D;
+end
+
+endmodule
+EOT
+
+proc
+opt_expr
+opt_dff
+# show
+# dump
+show -prefix before_a
+abstract -state -enablen magic
+check
+show -prefix after_a
 # opt_clean
 # show
\ No newline at end of file