From 5f79a6e868118bb9f91a3d72f23621a5b1f40f90 Mon Sep 17 00:00:00 2001
From: Jannis Harder <me@jix.one>
Date: Wed, 3 Sep 2025 15:36:42 +0200
Subject: [PATCH] Clean up $buf with 'z inputs, $input_port and $connect cells

This ensures that entering and leaving bufnorm followed by `opt_clean`
is equivalent to just running `opt_clean`.

Also make sure that 'z-$buf cells get techmapped in a compatible way.
---
 passes/opt/opt_clean.cc     | 29 +++++++++++++++++++++++++++--
 passes/techmap/simplemap.cc | 17 ++++++++++++++++-
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/passes/opt/opt_clean.cc b/passes/opt/opt_clean.cc
index d9bd4c027..a254dc915 100644
--- a/passes/opt/opt_clean.cc
+++ b/passes/opt/opt_clean.cc
@@ -600,15 +600,40 @@ void rmunused_module(RTLIL::Module *module, bool purge_mode, bool verbose, bool
 		log("Finding unused cells or wires in module %s..\n", module->name);
 
 	std::vector<RTLIL::Cell*> delcells;
-	for (auto cell : module->cells())
+	for (auto cell : module->cells()) {
 		if (cell->type.in(ID($pos), ID($_BUF_), ID($buf)) && !cell->has_keep_attr()) {
 			bool is_signed = cell->type == ID($pos) && cell->getParam(ID::A_SIGNED).as_bool();
 			RTLIL::SigSpec a = cell->getPort(ID::A);
 			RTLIL::SigSpec y = cell->getPort(ID::Y);
 			a.extend_u0(GetSize(y), is_signed);
-			module->connect(y, a);
+
+			if (a.has_const(State::Sz)) {
+				SigSpec new_a;
+				SigSpec new_y;
+				for (int i = 0; i < GetSize(a); ++i) {
+					SigBit b = a[i];
+					if (b == State::Sz)
+						continue;
+					new_a.append(b);
+					new_y.append(y[i]);
+				}
+				a = std::move(new_a);
+				y = std::move(new_y);
+			}
+			if (!y.empty())
+				module->connect(y, a);
+			delcells.push_back(cell);
+		} else if (cell->type.in(ID($connect)) && !cell->has_keep_attr()) {
+			RTLIL::SigSpec a = cell->getPort(ID::A);
+			RTLIL::SigSpec b = cell->getPort(ID::B);
+			if (a.has_const() && !b.has_const())
+				std::swap(a, b);
+			module->connect(a, b);
+			delcells.push_back(cell);
+		} else if (cell->type.in(ID($input_port)) && !cell->has_keep_attr()) {
 			delcells.push_back(cell);
 		}
+	}
 	for (auto cell : delcells) {
 		if (verbose)
 			log_debug("  removing buffer cell `%s': %s = %s\n", cell->name.c_str(),
diff --git a/passes/techmap/simplemap.cc b/passes/techmap/simplemap.cc
index b23985770..ddfd7e03e 100644
--- a/passes/techmap/simplemap.cc
+++ b/passes/techmap/simplemap.cc
@@ -47,7 +47,22 @@ void simplemap_buf(RTLIL::Module *module, RTLIL::Cell *cell)
 	RTLIL::SigSpec sig_a = cell->getPort(ID::A);
 	RTLIL::SigSpec sig_y = cell->getPort(ID::Y);
 
-	module->connect(RTLIL::SigSig(sig_y, sig_a));
+	if (sig_a.has_const(State::Sz)) {
+		SigSpec new_a;
+		SigSpec new_y;
+		for (int i = 0; i < GetSize(sig_a); ++i) {
+			SigBit b = sig_a[i];
+			if (b == State::Sz)
+				continue;
+			new_a.append(b);
+			new_y.append(sig_y[i]);
+		}
+		sig_a = std::move(new_a);
+		sig_y = std::move(new_y);
+	}
+
+	if (!sig_y.empty())
+		module->connect(RTLIL::SigSig(sig_y, sig_a));
 }
 
 void simplemap_pos(RTLIL::Module *module, RTLIL::Cell *cell)