utils: add BitGrouper for shared bit-partition logic

kernel/unstable/patch.cc

@@ -97,6 +97,8 @@ void Patch::gc(Cell* old_cell, bool track) {
 		auto dir = old_cell->port_dir(port_name);
 		// Unknown port direction (e.g. user module instance whose interface
 		// isn't registered): can't decide input vs output, so don't gc it.
+		// TODO: should be log_assert once PD_UNKNOWN is fixed at the source
+		// (see claude-notes.md).
 		if (dir == PD_UNKNOWN)
 			return;
 		// TODO only running GC through whole connections?

kernel/utils.h

@@ -22,6 +22,7 @@
 
 #include "kernel/yosys.h"
 #include <iterator>
+#include <optional>
 
 #ifndef UTILS_H
 #define UTILS_H
@@ -125,6 +126,91 @@ public:
 };
 
 
+// ---------------------------------------------------
+// BitGrouper — partition output bits by a per-bit key
+// ---------------------------------------------------
+//
+// Many passes that split a multi-bit cell or word-level FF into smaller
+// pieces share the same shape:
+//
+//   1. iterate bits of the output
+//   2. for each bit, compute a key describing where it should go (or
+//      decide it stays where it is)
+//   3. emit one piece per distinct key, covering exactly that group's bits
+//   4. either consume the un-grouped "remaining" bits as a smaller original,
+//      or report that everything was grouped.
+//
+// BitGrouper handles (1)+(2)+(3)'s bookkeeping; callers do (3)'s
+// per-group emission and (4)'s remainder handling.
+//
+// Construction is one-shot: pass a width and a `key_fn(int) -> optional<Key>`.
+// `key_fn` returning std::nullopt leaves the bit in `remaining()`.
+//
+// Group iteration order matches first appearance, so emission is
+// deterministic across runs.
+
+template<typename Key>
+class BitGrouper
+{
+public:
+	struct Group {
+		Key key;
+		std::vector<int> indices;
+
+		// True when `indices` is exactly [front, front+size). Callers can
+		// use this to pick the cheap contiguous slice op when available.
+		bool is_contiguous() const {
+			if (indices.empty())
+				return true;
+			int start = indices.front();
+			for (size_t i = 0; i < indices.size(); i++)
+				if (indices[i] != start + (int)i)
+					return false;
+			return true;
+		}
+	};
+
+	template<typename KeyFn>
+	BitGrouper(int width, KeyFn key_fn) {
+		std::map<Key, size_t> idx;
+		for (int i = 0; i < width; i++) {
+			std::optional<Key> k = key_fn(i);
+			if (!k) {
+				remaining_.push_back(i);
+				continue;
+			}
+			auto it = idx.find(*k);
+			if (it == idx.end()) {
+				idx.emplace(*k, groups_.size());
+				groups_.push_back({*k, {i}});
+			} else {
+				groups_[it->second].indices.push_back(i);
+			}
+		}
+	}
+
+	const std::vector<Group> &groups() const { return groups_; }
+	const std::vector<int> &remaining() const { return remaining_; }
+	bool fully_grouped() const { return remaining_.empty(); }
+	bool nothing_grouped() const { return groups_.empty(); }
+
+	// Pick the bits of `sig` at the group's indices. Uses contiguous-range
+	// slicing when possible.
+	static RTLIL::SigSpec extract(const RTLIL::SigSpec &sig, const Group &g) {
+		if (g.is_contiguous() && !g.indices.empty())
+			return sig.extract(g.indices.front(), (int)g.indices.size());
+		RTLIL::SigSpec out;
+		for (int i : g.indices)
+			out.append(sig[i]);
+		return out;
+	}
+
+private:
+	std::vector<Group> groups_;
+	std::vector<int> remaining_;
+};
+
+
 // ---------------------------------------------------
 // Best-effort topological sorting with loop detection
 // ---------------------------------------------------

passes/opt/opt_dff.cc

@@ -23,6 +23,7 @@
 #include "kernel/rtlil.h"
 #include "kernel/qcsat.h"
 #include "kernel/modtools.h"
+#include "kernel/utils.h"
 #include "kernel/sigtools.h"
 #include "kernel/ffinit.h"
 #include "kernel/ff.h"
@@ -557,11 +558,9 @@ struct OptDffWorker
 
 	bool try_merge_srst(FfDataSigMapped &ff, Cell *cell, bool &changed)
 	{
-		std::map<ctrls_t, std::vector<int>> groups;
-		std::vector<int> remaining_indices;
 		Const::Builder val_srst_builder(ff.width);
 
-		for (int i = 0; i < ff.width; i++) {
+		BitGrouper<ctrls_t> grouper(ff.width, [&](int i) -> std::optional<ctrls_t> {
 			ctrls_t resets;
 			State reset_val = ff.has_srst ? ff.val_srst[i] : State::Sx;
 
@@ -593,29 +592,26 @@ struct OptDffWorker
 				}
 			}
 
-			if (!resets.empty()) {
-				if (ff.has_srst)
-					resets.insert(ctrl_t(ff.sig_srst, ff.pol_srst));
-
-				groups[resets].push_back(i);
-			} else {
-				remaining_indices.push_back(i);
-			}
-
 			val_srst_builder.push_back(reset_val);
-		}
+
+			if (resets.empty())
+				return std::nullopt;
+			if (ff.has_srst)
+				resets.insert(ctrl_t(ff.sig_srst, ff.pol_srst));
+			return resets;
+		});
 
 		Const val_srst = val_srst_builder.build();
 
-		for (auto &it : groups) {
+		for (auto &g : grouper.groups()) {
-			FfDataSigMapped new_ff = ff.slice(it.second);
+			FfDataSigMapped new_ff = ff.slice(g.indices);
 			Const::Builder new_val_srst_builder(new_ff.width);
 			for (int i = 0; i < new_ff.width; i++)
-				new_val_srst_builder.push_back(val_srst[it.second[i]]);
+				new_val_srst_builder.push_back(val_srst[g.indices[i]]);
 
 			new_ff.val_srst = new_val_srst_builder.build();
 
-			ctrl_t srst = combine_resets(it.first, ff.is_fine);
+			ctrl_t srst = combine_resets(g.key, ff.is_fine);
 			new_ff.has_srst = true;
 			new_ff.sig_srst = srst.first;
 			new_ff.pol_srst = srst.second;
@@ -631,13 +627,13 @@ struct OptDffWorker
 					log_signal(new_ff.sig_d), log_signal(new_ff.sig_q), log_signal(new_ff.val_srst));
 		}
 
-		if (remaining_indices.empty()) {
+		if (grouper.fully_grouped()) {
 			module->remove(cell);
 			return true;
 		}
 
-		if (GetSize(remaining_indices) != ff.width) {
+		if ((int)grouper.remaining().size() != ff.width) {
-			ff = ff.slice(remaining_indices);
+			ff = ff.slice(grouper.remaining());
 			ff.cell = cell;
 			changed = true;
 		}
@@ -647,10 +643,8 @@ struct OptDffWorker
 
 	bool try_merge_ce(FfDataSigMapped &ff, Cell *cell, bool &changed)
 	{
-		std::map<std::pair<patterns_t, ctrls_t>, std::vector<int>> groups;
+		using CeKey = std::pair<patterns_t, ctrls_t>;
-		std::vector<int> remaining_indices;
+		BitGrouper<CeKey> grouper(ff.width, [&](int i) -> std::optional<CeKey> {
-
-		for (int i = 0; i < ff.width; i++) {
 			ctrls_t enables;
 
 			while (bit2mux.count(ff.sig_d[i]) && bitusers[ff.sig_d[i]] == 1) {
@@ -677,19 +671,17 @@ struct OptDffWorker
 			if (!opt.simple_dffe)
 				patterns = find_muxtree_feedback_patterns(ff.sig_d[i], ff.sig_q[i], pattern_t());
 
-			if (!patterns.empty() || !enables.empty()) {
+			if (patterns.empty() && enables.empty())
-				if (ff.has_ce)
+				return std::nullopt;
-					enables.insert(ctrl_t(ff.sig_ce, ff.pol_ce));
+			if (ff.has_ce)
-				simplify_patterns(patterns);
+				enables.insert(ctrl_t(ff.sig_ce, ff.pol_ce));
-				groups[std::make_pair(patterns, enables)].push_back(i);
+			simplify_patterns(patterns);
-			} else {
+			return std::make_pair(patterns, enables);
-				remaining_indices.push_back(i);
+		});
-			}
-		}
 
-		for (auto &it : groups) {
+		for (auto &g : grouper.groups()) {
-			FfDataSigMapped new_ff = ff.slice(it.second);
+			FfDataSigMapped new_ff = ff.slice(g.indices);
-			ctrl_t en = make_patterns_logic(it.first.first, it.first.second, ff.is_fine);
+			ctrl_t en = make_patterns_logic(g.key.first, g.key.second, ff.is_fine);
 
 			new_ff.has_ce = true;
 			new_ff.sig_ce = en.first;
@@ -705,13 +697,13 @@ struct OptDffWorker
 					log_signal(new_ff.sig_d), log_signal(new_ff.sig_q));
 		}
 
-		if (remaining_indices.empty()) {
+		if (grouper.fully_grouped()) {
 			module->remove(cell);
 			return true;
 		}
 
-		if (GetSize(remaining_indices) != ff.width) {
+		if ((int)grouper.remaining().size() != ff.width) {
-			ff = ff.slice(remaining_indices);
+			ff = ff.slice(grouper.remaining());
 			ff.cell = cell;
 			changed = true;
 		}

passes/opt/opt_expr.cc

@@ -170,10 +170,11 @@ bool group_cell_inputs(RTLIL::Module *module, RTLIL::Cell *cell, bool commutativ
 	std::vector<RTLIL::SigBit> bits_a = sig_a, bits_b = sig_b, bits_y = sig_y;
 
 	enum { GRP_DYN, GRP_CONST_A, GRP_CONST_B, GRP_CONST_AB, GRP_N };
-	std::map<std::pair<RTLIL::SigBit, RTLIL::SigBit>, std::set<RTLIL::SigBit>> grouped_bits[GRP_N];
+	// Two-level partition: outer key is the rewrite kind (GRP_*), inner key
-
+	// is the (a, b) pair that selects this bit's slot inside the new cell.
-	for (int i = 0; i < GetSize(bits_y); i++)
+	// Two original bits sharing (kind, a, b) collapse to the same output bit.
-	{
+	using Key = std::tuple<int, RTLIL::SigBit, RTLIL::SigBit>;
+	BitGrouper<Key> grouper(GetSize(bits_y), [&](int i) -> std::optional<Key> {
 		int group_idx = GRP_DYN;
 		RTLIL::SigBit bit_a = bits_a[i], bit_b = bits_b[i];
 
@@ -208,11 +209,16 @@ bool group_cell_inputs(RTLIL::Module *module, RTLIL::Cell *cell, bool commutativ
 				group_idx = GRP_CONST_B;
 		}
 
-		grouped_bits[group_idx][std::pair<RTLIL::SigBit, RTLIL::SigBit>(bit_a, bit_b)].insert(bits_y[i]);
+		return std::make_tuple(group_idx, bit_a, bit_b);
-	}
+	});
 
-	for (int i = 0; i < GRP_N; i++)
+	// If every original bit ended up with its own unique (a, b) slot within
-		if (GetSize(grouped_bits[i]) == GetSize(bits_y))
+	// some single kind, splitting would just rename without shrinking.
+	int slots_per_kind[GRP_N] = {};
+	for (auto &g : grouper.groups())
+		slots_per_kind[std::get<0>(g.key)]++;
+	for (int kind = 0; kind < GRP_N; kind++)
+		if (slots_per_kind[kind] == GetSize(bits_y))
 			return false;
 
 	log_debug("Replacing %s cell `%s' in module `%s' with cells using grouped bits:\n",
@@ -223,23 +229,33 @@ bool group_cell_inputs(RTLIL::Module *module, RTLIL::Cell *cell, bool commutativ
 	// Replacement bit for each original sig_y bit.
 	dict<SigBit, SigBit> bit_map;
 
-	for (int i = 0; i < GRP_N; i++)
+	// Pre-collect groups per kind, preserving BitGrouper's insertion order
+	// within each kind.
+	std::vector<const BitGrouper<Key>::Group*> per_kind[GRP_N];
+	for (auto &g : grouper.groups())
+		per_kind[std::get<0>(g.key)].push_back(&g);
+
+	for (int kind = 0; kind < GRP_N; kind++)
 	{
-		if (grouped_bits[i].empty())
+		if (per_kind[kind].empty())
 			continue;
 
-		int group_size = GetSize(grouped_bits[i]);
+		int group_size = GetSize(per_kind[kind]);
 		RTLIL::SigSpec new_y = patcher.addWire(NEW_ID, group_size);
 		RTLIL::SigSpec new_a, new_b;
 
-		for (auto &it : grouped_bits[i]) {
+		int slot = 0;
-			for (auto &bit : it.second)
+		for (auto *g : per_kind[kind]) {
-				bit_map[bit] = new_y[new_a.size()];
+			SigBit bit_a = std::get<1>(g->key);
-			new_a.append(it.first.first);
+			SigBit bit_b = std::get<2>(g->key);
-			new_b.append(it.first.second);
+			new_a.append(bit_a);
+			new_b.append(bit_b);
+			for (int i : g->indices)
+				bit_map[bits_y[i]] = new_y[slot];
+			slot++;
 		}
 
-		if (cell->type.in(ID($and), ID($or)) && i == GRP_CONST_A) {
+		if (cell->type.in(ID($and), ID($or)) && kind == GRP_CONST_A) {
 			if (!keepdc) {
 				if (cell->type == ID($and))
 					new_a.replace(dict<SigBit,SigBit>{{State::Sx, State::S0}, {State::Sz, State::S0}}, &new_b);
@@ -259,34 +275,34 @@ bool group_cell_inputs(RTLIL::Module *module, RTLIL::Cell *cell, bool commutativ
 			continue;
 		}
 
-		if (cell->type.in(ID($xor), ID($xnor)) && i == GRP_CONST_A) {
+		if (cell->type.in(ID($xor), ID($xnor)) && kind == GRP_CONST_A) {
 			SigSpec undef_a, undef_y, undef_b;
 			SigSpec def_y, def_a, def_b;
-			for (int i = 0; i < GetSize(new_y); i++) {
+			for (int j = 0; j < GetSize(new_y); j++) {
-				bool undef = new_a[i] == State::Sx || new_a[i] == State::Sz;
+				bool undef = new_a[j] == State::Sx || new_a[j] == State::Sz;
-				if (!keepdc && (undef || new_a[i] == new_b[i])) {
+				if (!keepdc && (undef || new_a[j] == new_b[j])) {
-					undef_a.append(new_a[i]);
+					undef_a.append(new_a[j]);
 					if (cell->type == ID($xor))
 						undef_b.append(State::S0);
 					// For consistency since simplemap does $xnor -> $_XOR_ + $_NOT_
 					else if (cell->type == ID($xnor))
 						undef_b.append(State::S1);
 					else log_abort();
-					undef_y.append(new_y[i]);
+					undef_y.append(new_y[j]);
 				}
-				else if (new_a[i] == State::S0 || new_a[i] == State::S1) {
+				else if (new_a[j] == State::S0 || new_a[j] == State::S1) {
-					undef_a.append(new_a[i]);
+					undef_a.append(new_a[j]);
 					if (cell->type == ID($xor))
-						undef_b.append(new_a[i] == State::S1 ? patcher.Not(NEW_ID, new_b[i]).as_bit() : new_b[i]);
+						undef_b.append(new_a[j] == State::S1 ? patcher.Not(NEW_ID, new_b[j]).as_bit() : new_b[j]);
 					else if (cell->type == ID($xnor))
-						undef_b.append(new_a[i] == State::S1 ? new_b[i] : patcher.Not(NEW_ID, new_b[i]).as_bit());
+						undef_b.append(new_a[j] == State::S1 ? new_b[j] : patcher.Not(NEW_ID, new_b[j]).as_bit());
 					else log_abort();
-					undef_y.append(new_y[i]);
+					undef_y.append(new_y[j]);
 				}
 				else {
-					def_a.append(new_a[i]);
+					def_a.append(new_a[j]);
-					def_b.append(new_b[i]);
+					def_b.append(new_b[j]);
-					def_y.append(new_y[i]);
+					def_y.append(new_y[j]);
 				}
 			}
 			if (!undef_y.empty()) {