pattern muxadd
//
// Authored by Akash Levy and Alain Dargelas of Silimate, Inc. under ISC license.
// Transforms add->mux into mux->add:
// y = s ? (a + b) : a   ===>   y = a + (s ? b : 0)
// or
// y = s ? a : (a + b)   ===>   y = a + (s ? 0 : b)

state <SigSpec> add_a add_b add_y add_a_ext mux_a mux_b mux_y
state <Const> add_a_signed
state <IdString> add_a_id add_b_id mux_a_id mux_b_id

match add
	// Select adder
	select add->type == $add

	// Set ports, allowing A and B to be swapped
	choice <IdString> A {\A, \B}
	define <IdString> B (A == \A ? \B : \A)
	set add_a port(add, A)
	set add_b port(add, B)
	set add_y port(add, \Y)

	// Get signedness
	set add_a_signed param(add, (A == \A) ? \A_SIGNED : \B_SIGNED)

	// Choice ids
	set add_a_id A
	set add_b_id B
endmatch

code add_y add_a add_b add_a_ext
	// Get adder signals
	add_a_ext = SigSpec(port(add, add_a_id));
	add_a_ext.extend_u0(GetSize(add_y), add_a_signed.as_bool());

	// Fanout of each adder Y bit should be 1 (no bit-split)
	if (nusers(add_y) != 2)
		reject;
endcode

match mux 
	// Select mux of form: s ? (a + b) : a
	// Allow leading 0s when A_WIDTH != Y_WIDTH or s ? a : (a + b)
	select mux->type == $mux
	choice <IdString> AB {\A, \B}
	define <IdString> BA (AB == \A ? \B : \A)
	set mux_y port(mux, \Y)
	set mux_a port(mux, AB)
	set mux_b port(mux, BA)
	set mux_a_id AB
	set mux_b_id BA
	index <SigSpec> port(mux, AB) === add_a_ext
	index <SigSpec> port(mux, BA) === add_y
endmatch

code add_y add_a add_b add_a_ext add_a_id add_b_id  mux_y mux_a mux_b mux_a_id mux_b_id
	// Get mux signal
	SigSpec mid; 
	std::string adder_y_name;
	if (add_y.is_wire()) 
		adder_y_name = add_y.as_wire()->name.c_str();
	else
		adder_y_name = add_y.as_string();

	// Start by renaming the LHS of an eventual assign statement
	// where the RHS is the adder output (that is getting rewired).
	// Renaming the signal allows equiv_opt to function as it would
	// otherwise try to match the functionality which would fail
	// as the LHS signal has indeed changed function.

	// Adder output could be assigned
	for (auto it = module->connections().begin(); it != module->connections().end(); ++it) {
		RTLIL::SigSpec rhs = it->second;
		if (rhs.is_wire()) {
			const std::string& rhs_name = rhs.as_wire()->name.c_str();
			if (rhs_name == adder_y_name) {
				RTLIL::SigSpec lhs = it->first;
				if (lhs.is_wire()) {
					const std::string& lhs_name =  lhs.as_wire()->name.c_str();
					module->rename(lhs_name, module->uniquify("$" + lhs_name));
					break;
				}
			}
		}
	}
	// Alternatively, the port name could be a wire name
	if (add_y.is_wire()) {
		if (GetSize(adder_y_name)) {
			if (adder_y_name[0] != '$') {
				module->rename(adder_y_name, module->uniquify("$" + adder_y_name));
			}
		}
	} else {
		for (auto chunk : add_y.chunks()) {
			if (chunk.is_wire()) {
				const std::string& name = chunk.wire->name.c_str();
				if (name[0] != '$') {
					module->rename(name, module->uniquify("$" + name));
				}
			}
		}
	}

	// Create new mid wire
	Cell *cell = mux;
	mid = module->addWire(NEW_ID2_SUFFIX("mid"), GetSize(add_b));

	// Connect ports
	add->setPort(add_b_id, mid);
	add->setPort(add_a_id, add_a);
	add->setPort(\Y, add_y);
	cell = add;
	module->rename(add, NEW_ID2_SUFFIX("rot"));
	mux->setPort(mux_a_id, Const(State::S0, GetSize(add_b)));
	mux->setPort(mux_b_id, add_b);
	mux->setPort(\Y, mid);
	cell = mux;
	module->rename(mux, NEW_ID2_SUFFIX("rot"));
	module->connect(mux_y, add_y);

	// Log, fixup, accept
	log("muxadd pattern in %s: mux=%s, add=%s\n", log_id(module), log_id(mux), log_id(add));
	add->fixup_parameters();
	mux->fixup_parameters();
	did_something = true;
	accept;
endcode