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)
//

state <SigSpec> add_a add_b add_y add_a_ext
state <Const> add_a_signed
state <IdString> add_a_id add_b_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
	set add_b_id B
endmatch

code add_y add_a add_b add_a_ext
	// Get adder signals
	add_a = port(add, \A);
	add_b = port(add, \B);
	add_y = port(add, \Y);
	add_a_ext = SigSpec(add_a);
	add_a_ext.extend_u0(GetSize(add_y), param(add, \A_SIGNED).as_bool());
	// 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) === add_a_ext
	index <SigSpec> port(mux, \B) === add_y
endmatch

code
	// Get mux signal
	SigSpec mux_y = port(mux, \Y);
	SigSpec mux_a = port(mux, \A);
	SigSpec mux_b = port(mux, \B);
	// Prevent multiple branches to edit the same pair of mux/add
	if (mux_b == add_y || mux_a == add_y) {
	  // Create new mid wire
	  SigSpec mid = module->addWire(NEW_ID, GetSize(add_b));
	  // Rewire
	  add->setPort(\B, mid);
	  add->setPort(\A, add_a);
	  add->setPort(\Y, add_y);
	  mux->setPort(add_a_id, Const(State::S0, GetSize(add_b)));
	  mux->setPort(add_b_id, add_b);
	  mux->setPort(\Y, mid);
	  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