/*------------------------------------------------------------------------------
| Part of tweedledum.  This file is distributed under the MIT License.
| See accompanying file /LICENSE for details.
*-----------------------------------------------------------------------------*/
#pragma once

#include "BaseStrategy.h"

#include <algorithm>
#include <mockturtle/networks/klut.hpp>
#include <mockturtle/utils/node_map.hpp>
#include <vector>

namespace tweedledum {

class EagerStrategy : public BaseStrategy {
	using LogicNetwork = mockturtle::klut_network;
	using Node = LogicNetwork::node;

public:
	bool compute_steps(LogicNetwork const& klut) override;

private:
	void cleanup(LogicNetwork const& klut,
	    mockturtle::node_map<uint32_t, LogicNetwork>& ref_counts, Node node);
};

inline bool EagerStrategy::compute_steps(LogicNetwork const& klut)
{
	using namespace mockturtle;
	// Initialize reference counters
	node_map<uint32_t, LogicNetwork> ref_counts(klut, 0);
	klut.foreach_node([&](auto node) {
		ref_counts[node] = klut.fanout_size(node);
	});

	klut.clear_visited();
	std::vector<Node> outputs;
	klut.foreach_po([&](auto const& signal) {
		auto node = klut.get_node(signal);
		outputs.push_back(node);
		klut.set_visited(node, 1u);
	});

	this->steps_.reserve(klut.size() * 2);
	klut.foreach_node([&](auto node) {
		if (klut.is_constant(node) || klut.is_pi(node)) {
			return true;
		}

		steps_.emplace_back(Action::compute, node);
		if (!klut.visited(node)) {
			cleanup(klut, ref_counts, node);
		}
		return true;
	});
	return true;
}

inline void EagerStrategy::cleanup(LogicNetwork const& klut,
    mockturtle::node_map<uint32_t, LogicNetwork>& ref_counts, Node node)
{
	klut.foreach_fanin(node, [&](auto const& input) {
		const auto child = klut.get_node(input);
		if (klut.is_constant(child) || klut.is_pi(child)) {
			return;
		}

		if (--ref_counts[child] == 0u) {
			steps_.emplace_back(Action::cleanup, node);
			cleanup(klut, ref_counts, child);
		}
	});
}

} // namespace tweedledum
