From 6c1fa45995f690ba1c8f8412ab09d938f83d9130 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Povi=C5=A1er?= <povik@cutebit.org>
Date: Mon, 16 Sep 2024 17:20:34 +0200
Subject: [PATCH] aiger2: Ingest `$pmux`

---
 backends/aiger2/aiger.cc | 41 +++++++++++++++++++++++++++++++++++++++-
 tests/various/aiger2.ys  | 28 +++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/backends/aiger2/aiger.cc b/backends/aiger2/aiger.cc
index d300827d3..c0b9a0d7a 100644
--- a/backends/aiger2/aiger.cc
+++ b/backends/aiger2/aiger.cc
@@ -42,7 +42,7 @@ PRIVATE_NAMESPACE_BEGIN
 // TODO
 //#define ARITH_OPS ID($add), ID($sub), ID($neg)
 
-#define KNOWN_OPS BITWISE_OPS, REDUCE_OPS, LOGIC_OPS, GATE_OPS, ID($pos), CMP_OPS /*, ARITH_OPS*/
+#define KNOWN_OPS BITWISE_OPS, REDUCE_OPS, LOGIC_OPS, GATE_OPS, ID($pos), CMP_OPS, ID($pmux) /*, ARITH_OPS*/
 
 template<typename Writer, typename Lit>
 struct Index {
@@ -199,6 +199,28 @@ struct Index {
 		return OR(AND(a, b), AND(c, OR(a, b)));
 	}
 
+	Lit REDUCE(std::vector<Lit> lits, bool op_xor=false)
+	{
+		std::vector<Lit> next;
+		while (lits.size() > 1) {
+			next.clear();
+			for (int i = 0; i < lits.size(); i += 2) {
+				if (i + 1 >= lits.size()) {
+					next.push_back(lits[i]);
+				} else {
+					Lit a = lits[i], b = lits[i + 1];
+					next.push_back(op_xor ? XOR(a, b) : AND(a, b));
+				}
+			}
+			next.swap(lits);
+		}
+
+		if (lits.empty())
+			return op_xor ? CFALSE : CTRUE;
+		else
+			return lits.front();
+	}
+
 	Lit impl_op(HierCursor &cursor, Cell *cell, IdString oport, int obit)
 	{
 		if (cell->type.in(REDUCE_OPS, LOGIC_OPS, CMP_OPS) && obit != 0) {
@@ -360,6 +382,23 @@ struct Index {
 					log_abort();
 				}
 			}
+		} else if (cell->type == ID($pmux)) {
+			SigSpec aport = cell->getPort(ID::A);
+			SigSpec bport = cell->getPort(ID::B);
+			SigSpec sport = cell->getPort(ID::S);
+			int width = aport.size();
+
+			Lit a = visit(cursor, aport[obit]);
+
+			std::vector<Lit> bar, sels;
+			for (int i = 0; i < sport.size(); i++) {
+				Lit s = visit(cursor, sport[i]);
+				Lit b = visit(cursor, bport[width * i + obit]);
+				bar.push_back(NOT(AND(s, b)));
+				sels.push_back(NOT(s));
+			}
+
+			return OR(AND(REDUCE(sels), a), NOT(REDUCE(bar)));
 		} else {
 			log_abort();
 		}
diff --git a/tests/various/aiger2.ys b/tests/various/aiger2.ys
index 5d8c6e848..0efb5dd76 100644
--- a/tests/various/aiger2.ys
+++ b/tests/various/aiger2.ys
@@ -165,3 +165,31 @@ read_aiger -module_name test aiger2_ops.aig
 select -assert-none test/t:$_AND_ test/t:$_NOT_ %% test/c:* %D
 miter -equiv -flatten gold test miter
 sat -verify -prove trigger 0 miter
+
+design -reset
+read_verilog -icells <<EOF
+module test();
+wire [1:0] pmux_a, pmux_s, pmux_y;
+wire [3:0] pmux_b;
+\$pmux #(
+	.S_WIDTH(2),
+	.WIDTH(2)
+) pmux(.A(pmux_a), .B(pmux_b), .S(pmux_s), .Y(pmux_y));
+endmodule
+EOF
+
+expose -input c:* %ci* w:* %i
+expose c:* %co* w:* %i
+splitnets -ports
+opt_clean
+copy test gold
+select test
+write_aiger2 aiger2_xmodel.aig
+select -clear
+delete test
+read_aiger -module_name test aiger2_xmodel.aig
+select -assert-none test/t:$_AND_ test/t:$_NOT_ %% test/c:* %D
+
+equiv_make gold test equiv
+equiv_induct -undef equiv
+equiv_status -assert equiv