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

state <SigSpec> add_a add_b add_y
state <bool> add_a_signed
state <IdString> add_a_id

match add
	// Select adder
	select add->type == $add
	choice <IdString> A {\A, \B}
	define <IdString> B (A == \A ? \B : \A)
	set add_y port(add, \Y)
	set add_a port(add, A)
	set add_b port(add, B)
	set add_a_signed param(add, (A == \A) ? \A_SIGNED : \B_SIGNED)
	set add_a_id A
endmatch

code add_y add_a add_b
	// Get adder signals
	add_a = port(add, \A);
	add_b = port(add, \B);
	add_y = port(add, \Y);

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

	// A and B can be interchanged
	branch;
	std::swap(add_a, add_b);
endcode

match mux
	// Select mux of form s ? (a + b) : a, allow leading 0s when A_WIDTH != Y_WIDTH
	select mux->type == $mux
	index <SigSpec> port(mux, \A) === SigSpec({Const(!param(add, \A_SIGNED).bool() ? () : State::S0, GetSize(add_y)-GetSize(add_a)), add_a})
	index <SigSpec> port(mux, \B) === add_y
endmatch

code
	// Get mux signal
	SigSpec mux_y = port(mux, \Y);

	// Create new mid wire
	SigSpec mid = module->addWire(NEW_ID, GetSize(add_b));

	// Rewire
	mux->setPort(\A, Const(State::S0, GetSize(add_b)));
	mux->setPort(\B, add_b);
	mux->setPort(\Y, mid);
	add->setPort(\B, mid);
	add->setPort(\Y, mux_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();
	accept;
endcode