#include "kernel/sigtools.h"
#include "kernel/yosys.h"
#include <set>

USING_YOSYS_NAMESPACE
PRIVATE_NAMESPACE_BEGIN

// Signal cell driver(s), precompute a cell output signal to a cell map
void sigCellDrivers(RTLIL::Module *module, SigMap &sigmap, dict<RTLIL::SigSpec, std::set<Cell *>> &sig2CellsInFanout,
		    dict<RTLIL::SigSpec, std::set<Cell *>> &sig2CellsInFanin)
{
	for (auto cell : module->selected_cells()) {
		for (auto &conn : cell->connections()) {
			IdString portName = conn.first;
			RTLIL::SigSpec actual = conn.second;
			if (cell->output(portName)) {
				sig2CellsInFanin[actual].insert(cell);
				for (int i = 0; i < actual.size(); i++) {
					SigSpec bit_sig = actual.extract(i, 1);
					sig2CellsInFanin[sigmap(bit_sig)].insert(cell);
				}
			} else {
				sig2CellsInFanout[sigmap(actual)].insert(cell);
				for (int i = 0; i < actual.size(); i++) {
					SigSpec bit_sig = actual.extract(i, 1);
					if (!bit_sig.is_fully_const()) {
						sig2CellsInFanout[sigmap(bit_sig)].insert(cell);
					}
				}
			}
		}
	}
}

// Assign statements fanin, fanout, traces the lhs2rhs and rhs2lhs sigspecs and precompute maps
void lhs2rhs_rhs2lhs(RTLIL::Module *module, SigMap &sigmap, dict<RTLIL::SigSpec, std::set<RTLIL::SigSpec>> &rhsSig2LhsSig,
		     dict<RTLIL::SigSpec, RTLIL::SigSpec> &lhsSig2rhsSig)
{
	for (auto it = module->connections().begin(); it != module->connections().end(); ++it) {
		RTLIL::SigSpec lhs = it->first;
		RTLIL::SigSpec rhs = it->second;
		std::vector<SigSpec> lhsBits;
		for (int i = 0; i < lhs.size(); i++) {
			SigSpec bit_sig = lhs.extract(i, 1);
			lhsBits.push_back(bit_sig);
		}
		std::vector<SigSpec> rhsBits;
		for (int i = 0; i < rhs.size(); i++) {
			SigSpec bit_sig = rhs.extract(i, 1);
			if (bit_sig.is_fully_const()) {
				continue;
			}
			rhsBits.push_back(bit_sig);
		}

		for (uint32_t i = 0; i < lhsBits.size(); i++) {
			if (i < rhsBits.size()) {
				rhsSig2LhsSig[sigmap(rhsBits[i])].insert(sigmap(lhsBits[i]));
				lhsSig2rhsSig[sigmap(lhsBits[i])] = sigmap(rhsBits[i]);
			}
		}
	}
}

void collectTransitiveFanin(RTLIL::SigSpec &sig, dict<RTLIL::SigSpec, std::set<Cell *>> &sig2CellsInFanin,
			    dict<RTLIL::SigSpec, RTLIL::SigSpec> &lhsSig2RhsSig, std::set<Cell *> &visitedCells,
			    std::set<RTLIL::SigSpec> &visitedSigSpec)
{
	if (sig.is_fully_const()) {
		return;
	}
	if (visitedSigSpec.count(sig)) {
		return;
	}
	visitedSigSpec.insert(sig);

	if (sig2CellsInFanin.count(sig)) {
		for (Cell *cell : sig2CellsInFanin[sig]) {
			if (visitedCells.count(cell)) {
				continue;
			}
			visitedCells.insert(cell);
			for (auto &conn : cell->connections()) {
				IdString portName = conn.first;
				RTLIL::SigSpec actual = conn.second;
				if (cell->input(portName)) {
					if (!actual.is_chunk()) {
						for (auto it = actual.chunks().rbegin(); it != actual.chunks().rend(); ++it) {
							RTLIL::SigSpec sub_actual = *it;
							collectTransitiveFanin(sub_actual, sig2CellsInFanin, lhsSig2RhsSig, visitedCells,
									       visitedSigSpec);
						}
					} else {
						collectTransitiveFanin(actual, sig2CellsInFanin, lhsSig2RhsSig, visitedCells, visitedSigSpec);
					}
				}
			}
		}
	}
	if (lhsSig2RhsSig.count(sig)) {
		RTLIL::SigSpec rhs = lhsSig2RhsSig[sig];
		collectTransitiveFanin(rhs, sig2CellsInFanin, lhsSig2RhsSig, visitedCells, visitedSigSpec);
	}
}

void observabilityClean(RTLIL::Module *module, dict<RTLIL::SigSpec, std::set<Cell *>> &sig2CellsInFanin,
			dict<RTLIL::SigSpec, RTLIL::SigSpec> &lhsSig2RhsSig)
{
	if (module->get_bool_attribute(ID::keep))
		return;
	std::set<Cell *> visitedCells;
	std::set<RTLIL::SigSpec> visitedSigSpec;

	for (auto elt : sig2CellsInFanin) {
		RTLIL::SigSpec po = elt.first;
		RTLIL::Wire *w = po[0].wire;
		if (w && !w->port_output) {
			continue;
		}
		collectTransitiveFanin(po, sig2CellsInFanin, lhsSig2RhsSig, visitedCells, visitedSigSpec);
	}

	for (auto elt : lhsSig2RhsSig) {
		RTLIL::SigSpec po = elt.first;
		RTLIL::Wire *w = po[0].wire;
		if (w && !w->port_output) {
			continue;
		}
		collectTransitiveFanin(po, sig2CellsInFanin, lhsSig2RhsSig, visitedCells, visitedSigSpec);
	}

	pool<RTLIL::SigSig> newConnections;
	for (auto it = module->connections().begin(); it != module->connections().end(); ++it) {
		RTLIL::SigSpec lhs = it->first;
		if (visitedSigSpec.count(lhs)) {
			newConnections.insert(*it);
			continue;
		}
	}

	module->connections_.clear();
	for (auto conn : newConnections) {
		module->connect(conn);
	}

	pool<RTLIL::Wire *> wiresToRemove;
	for (auto wire : module->wires()) {
		RTLIL::SigSpec sig = wire;
		if (visitedSigSpec.count(sig)) {
			continue;
		}
		RTLIL::Wire *w = sig[0].wire;
		if (w->port_id) {
			continue;
		}
		if (w->get_bool_attribute(ID::keep))
			continue;
		wiresToRemove.insert(w);
	}

	module->remove(wiresToRemove);

	std::set<Cell *> cellsToRemove;
	for (auto cell : module->cells()) {
		if (visitedCells.count(cell)) {
			continue;
		}
		if (cell->has_keep_attr())
			continue;
		cellsToRemove.insert(cell);
	}

	for (auto cell : cellsToRemove) {
		module->remove(cell);
	}
}

struct ObsClean : public ScriptPass {
	ObsClean() : ScriptPass("obs_clean", "Observability-based cleanup") {}
	void script() override {}

	void execute(std::vector<std::string>, RTLIL::Design *design) override
	{
		if (design == nullptr) {
			log_error("No design object");
			return;
		}
		log("Running obs_clean pass\n");
		log_flush();
		for (auto module : design->selected_modules()) {
			SigMap sigmap(module);
			// Precompute cell output sigspec to cell map
			dict<RTLIL::SigSpec, std::set<Cell *>> sig2CellsInFanin;
			dict<RTLIL::SigSpec, std::set<Cell *>> sig2CellsInFanout;
			sigCellDrivers(module, sigmap, sig2CellsInFanout, sig2CellsInFanin);
			// Precompute lhs2rhs and rhs2lhs sigspec map
			dict<RTLIL::SigSpec, RTLIL::SigSpec> lhsSig2RhsSig;
			dict<RTLIL::SigSpec, std::set<RTLIL::SigSpec>> rhsSig2LhsSig;
			lhs2rhs_rhs2lhs(module, sigmap, rhsSig2LhsSig, lhsSig2RhsSig);
			// Actual cleanup
			observabilityClean(module, sig2CellsInFanin, lhsSig2RhsSig);
		}
		log("End obs_clean pass\n");
		log_flush();
	}
} SplitNetlist;

PRIVATE_NAMESPACE_END