/*
 *  yosys -- Yosys Open SYnthesis Suite
 *
 *  Copyright (C) 2026  Akash Levy <akash@silimate.com>
 *
 *  Permission to use, copy, modify, and/or distribute this software for any
 *  purpose with or without fee is hereby granted, provided that the above
 *  copyright notice and this permission notice appear in all copies.
 *
 *  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 *  WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 *  MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 *  ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 *  WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 *  ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 *  OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 */

#include "kernel/celltypes.h"
#include "kernel/register.h"
#include "kernel/rtlil.h"
#include "kernel/sigtools.h"

USING_YOSYS_NAMESPACE
PRIVATE_NAMESPACE_BEGIN

struct BoundaryConeWorker {
	Module *child, *parent;
	SigMap child_sigmap;
	int max_cells;
	int max_bits;

	dict<SigBit, Cell*> bit_driver;
	dict<SigBit, SigBit> input_map;
	dict<SigBit, SigBit> copied_bits;
	dict<Cell*, Cell*> copied_cells;
	std::vector<Wire*> created_wires;
	std::vector<Cell*> created_cells;
	pool<Cell*> active_cells;
	int copied_cell_count = 0;
	int materialized_bit_count = 0;
	bool failed = false;

	BoundaryConeWorker(Module *child, Module *parent, Cell *instance, int max_cells, int max_bits)
		: child(child), parent(parent), child_sigmap(child), max_cells(max_cells), max_bits(max_bits)
	{
		for (auto wire : child->wires()) {
			if (!wire->port_input || wire->port_output)
				continue;
			if (!instance->connections_.count(wire->name))
				continue;
			SigSpec conn = instance->connections_.at(wire->name);
			if (GetSize(conn) != wire->width) {
				failed = true;
				continue;
			}
			for (int i = 0; i < wire->width; i++)
				input_map[child_sigmap(SigBit(wire, i))] = conn[i];
		}

		for (auto cell : child->cells()) {
			if (!yosys_celltypes.cell_evaluable(cell->type) || cell->has_keep_attr())
				continue;
			for (auto &conn : cell->connections()) {
				if (!yosys_celltypes.cell_output(cell->type, conn.first))
					continue;
				for (auto bit : conn.second)
					if (bit.is_wire())
						bit_driver[child_sigmap(bit)] = cell;
			}
		}
	}

	SigSpec materialize(SigSpec sig)
	{
		SigSpec result;
		for (auto bit : sig) {
			if (++materialized_bit_count > max_bits) {
				failed = true;
				break;
			}
			result.append(materialize(bit));
			if (failed)
				break;
		}
		return result;
	}

	SigBit materialize(SigBit bit)
	{
		bit = child_sigmap(bit);

		if (!bit.is_wire())
			return bit;

		if (input_map.count(bit))
			return input_map.at(bit);

		if (copied_bits.count(bit))
			return copied_bits.at(bit);

		if (!bit_driver.count(bit)) {
			failed = true;
			return RTLIL::Sx;
		}

		Cell *driver = bit_driver.at(bit);
		if (active_cells.count(driver) || copied_cell_count >= max_cells) {
			failed = true;
			return RTLIL::Sx;
		}

		copy_driver(driver);
		if (failed || !copied_bits.count(bit)) {
			failed = true;
			return RTLIL::Sx;
		}
		return copied_bits.at(bit);
	}

	void copy_driver(Cell *driver)
	{
		if (copied_cells.count(driver))
			return;

		active_cells.insert(driver);
		dict<IdString, SigSpec> new_connections;

		for (auto &conn : driver->connections()) {
			if (yosys_celltypes.cell_input(driver->type, conn.first)) {
				SigSpec mapped = materialize(conn.second);
				if (failed)
					break;
				new_connections[conn.first] = mapped;
				continue;
			}

			if (yosys_celltypes.cell_output(driver->type, conn.first)) {
				Wire *wire = parent->addWire(NEW_ID_SUFFIX("opt_boundary"), GetSize(conn.second));
				created_wires.push_back(wire);
				SigSpec mapped = wire;
				new_connections[conn.first] = mapped;
				for (int i = 0; i < GetSize(conn.second); i++) {
					if (conn.second[i].is_wire())
						copied_bits[child_sigmap(conn.second[i])] = mapped[i];
				}
				continue;
			}

			failed = true;
			break;
		}

		if (!failed && copied_cell_count >= max_cells)
			failed = true;

		if (!failed) {
			Cell *copy = parent->addCell(NEW_ID_SUFFIX("opt_boundary"), driver);
			for (auto &conn : new_connections)
				copy->setPort(conn.first, conn.second);
			copied_cells[driver] = copy;
			created_cells.push_back(copy);
			copied_cell_count++;
		}

		active_cells.erase(driver);
	}

	void rollback()
	{
		for (auto it = created_cells.rbegin(); it != created_cells.rend(); ++it)
			parent->remove(*it);
		for (auto it = created_wires.rbegin(); it != created_wires.rend(); ++it)
			parent->remove(pool<Wire*>{*it});
		created_cells.clear();
		created_wires.clear();
		copied_cells.clear();
		copied_bits.clear();
		copied_cell_count = 0;
		materialized_bit_count = 0;
	}
};

static bool protected_module(Module *module)
{
	return RTLIL::unescape_id(module->name).compare(0, 2, "DW") == 0 ||
			module->get_blackbox_attribute() ||
			module->get_bool_attribute(ID::keep) ||
			module->get_bool_attribute(ID::keep_hierarchy);
}

struct ParentUsage {
	Design *design;
	SigMap sigmap;
	SigPool used;

	ParentUsage(Module *module, Design *design) : design(design), sigmap(module)
	{
		auto count_usage = [&](const SigSpec &signal) {
			for (auto bit : signal)
				used.add(sigmap(bit));
		};

		for (auto wire : module->wires()) {
			if (wire->port_output)
				count_usage(wire);
		}

		for (auto [_, process] : module->processes)
			process->rewrite_sigspecs(count_usage);

		for (auto cell : module->cells()) {
			Module *cell_module = design->module(cell->type);
			for (auto &conn : cell->connections()) {
				if (yosys_celltypes.cell_known(cell->type)) {
					if (yosys_celltypes.cell_input(cell->type, conn.first))
						count_usage(conn.second);
					continue;
				}

				if (cell_module != nullptr) {
					Wire *port = cell_module->wire(conn.first);
					if (port != nullptr && port->port_input)
						count_usage(conn.second);
					continue;
				}

				// Unknown cells may observe any connection.
				count_usage(conn.second);
			}
		}
	}

	bool check(SigBit bit)
	{
		return bit.is_wire() && used.check(sigmap(bit));
	}
};

struct OptBoundaryPass : Pass {
	OptBoundaryPass() : Pass("opt_boundary", "perform conservative parent-side cross-boundary cone optimization") {}

	void help() override
	{
		log("\n");
		log("    opt_boundary [options] [selection]\n");
		log("\n");
		log("This pass performs a conservative form of hierarchical boundary optimization.\n");
		log("For each selected parent module, it looks through instances of non-blackbox,\n");
		log("non-keep child modules and copies small evaluable combinational cones that\n");
		log("drive child output ports into the parent.  The original child module body is\n");
		log("left unchanged; optimized instance outputs are disconnected only after an\n");
		log("equivalent parent-side cone has been created.\n");
		log("\n");
		log("    -max_cells <N>\n");
		log("        maximum number of child cells to copy for one output bit. Default: 8.\n");
		log("\n");
		log("    -max_bits <N>\n");
		log("        maximum number of child cone bits to inspect for one output bit. Default: 4096.\n");
		log("\n");
		log("    -no_disconnect\n");
		log("        copy eligible cones into the parent but leave instance output ports\n");
		log("        connected to their original nets.\n");
		log("\n");
	}

	void execute(std::vector<std::string> args, RTLIL::Design *design) override
	{
		log_header(design, "Executing OPT_BOUNDARY pass.\n");

		int max_cells = 8;
		int max_bits = 4096;
		bool no_disconnect = false;
		size_t argidx;
		for (argidx = 1; argidx < args.size(); argidx++) {
			if (args[argidx] == "-max_cells" && argidx + 1 < args.size()) {
				max_cells = atoi(args[++argidx].c_str());
				continue;
			}
			if (args[argidx] == "-max_bits" && argidx + 1 < args.size()) {
				max_bits = atoi(args[++argidx].c_str());
				continue;
			}
			if (args[argidx] == "-no_disconnect") {
				no_disconnect = true;
				continue;
			}
			break;
		}
		extra_args(args, argidx, design);

		if (max_cells < 1)
			log_cmd_error("The -max_cells value must be positive.\n");
		if (max_bits < 1)
			log_cmd_error("The -max_bits value must be positive.\n");

		bool did_something = false;
		for (auto parent : design->selected_modules(RTLIL::SELECT_WHOLE_ONLY, RTLIL::SB_UNBOXED_CMDERR)) {
			if (protected_module(parent)) {
				log_debug("opt_boundary: skipping protected parent module %s\n", log_id(parent));
				continue;
			}

			ParentUsage parent_usage(parent, design);

			for (auto instance : parent->cells().to_vector()) {
				if (instance->has_keep_attr()) {
					log_debug("opt_boundary: skipping kept instance %s in %s\n", log_id(instance), log_id(parent));
					continue;
				}

				Module *child = design->module(instance->type);
				if (child == nullptr) {
					log_debug("opt_boundary: skipping non-module cell %s (type %s) in %s\n",
							log_id(instance), log_id(instance->type), log_id(parent));
					continue;
				}
				if (protected_module(child)) {
					log_debug("opt_boundary: skipping protected child module %s for instance %s in %s\n",
							log_id(child), log_id(instance), log_id(parent));
					continue;
				}

				for (auto &conn : instance->connections_) {
					Wire *port = child->wire(conn.first);
					if (port == nullptr || !port->port_output || port->port_input)
						continue;
					if (port->width != GetSize(conn.second))
						log_error("Port %s connected on instance %s not found in module %s or width is not matching\n",
								log_id(conn.first), log_id(instance), log_id(child));

					SigSpec new_conn = conn.second;
					bool changed_port = false;
					log_debug("opt_boundary: checking output port %s (%d bits) on instance %s in %s\n",
							log_id(conn.first), port->width, log_id(instance), log_id(parent));
					for (int i = 0; i < port->width; i++) {
						if (!conn.second[i].is_wire()) {
							log_debug("opt_boundary: skipping %s[%d] on %s because parent connection is constant\n",
									log_id(port), i, log_id(instance));
							continue;
						}
						if (!parent_usage.check(conn.second[i])) {
							log_debug("opt_boundary: skipping %s[%d] on %s because parent net %s is unobserved\n",
									log_id(port), i, log_id(instance), log_signal(conn.second[i]));
							continue;
						}

						BoundaryConeWorker worker(child, parent, instance, max_cells, max_bits);
						SigBit replacement = worker.materialize(SigBit(port, i));
						if (worker.failed) {
							log_debug("opt_boundary: failed to materialize %s[%d] of instance %s after inspecting %d bits; rolling back\n",
									log_id(port), i, log_id(instance), worker.materialized_bit_count);
							worker.rollback();
							continue;
						}
						if (replacement == conn.second[i]) {
							log_debug("opt_boundary: skipping %s[%d] on %s because replacement is identical\n",
									log_id(port), i, log_id(instance));
							worker.rollback();
							continue;
						}
						if (parent_usage.sigmap(replacement) == parent_usage.sigmap(conn.second[i])) {
							log_debug("opt_boundary: skipping %s[%d] on %s because replacement is already equivalent in parent\n",
									log_id(port), i, log_id(instance));
							worker.rollback();
							continue;
						}
						if (no_disconnect && worker.copied_cell_count == 0) {
							log_debug("opt_boundary: skipping zero-cell bypass for %s[%d] on %s in -no_disconnect mode\n",
									log_id(port), i, log_id(instance));
							continue;
						}

						if (!no_disconnect) {
							parent->connect(conn.second[i], replacement);
							Wire *dummy = parent->addWire(NEW_ID_SUFFIX("opt_boundary_output"));
							new_conn[i] = SigBit(dummy, 0);
							changed_port = true;
						}
						did_something = true;

						if (worker.copied_cell_count > 0)
							log("Copied %d cells from cone driving %s[%d] of instance '%s' (type '%s') into '%s'\n",
									worker.copied_cell_count, log_id(port), i, log_id(instance), log_id(instance->type), log_id(parent));
						else
							log("Bypassed cone driving %s[%d] of instance '%s' (type '%s') in '%s'\n",
									log_id(port), i, log_id(instance), log_id(instance->type), log_id(parent));
					}

					if (changed_port)
						conn.second = new_conn;
				}
			}
		}

		if (did_something)
			design->scratchpad_set_bool("opt.did_something", true);
	}
} OptBoundaryPass;

PRIVATE_NAMESPACE_END