/*
	This file is part of solidity.

	solidity is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	solidity is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with solidity.  If not, see <http://www.gnu.org/licenses/>.
*/
// SPDX-License-Identifier: GPL-3.0
/**
 * @file CommonSubexpressionEliminator.cpp
 * @author Christian <c@ethdev.com>
 * @date 2015
 * Optimizer step for common subexpression elimination and stack reorganisation.
 */

#include <functional>
#include <libsolutil/Keccak256.h>
#include <libevmasm/CommonSubexpressionEliminator.h>
#include <libevmasm/AssemblyItem.h>
#include <libsolutil/StackTooDeepString.h>

#include <range/v3/view/reverse.hpp>

using namespace solidity;
using namespace solidity::evmasm;
using namespace solidity::langutil;

std::vector<AssemblyItem> CommonSubexpressionEliminator::getOptimizedItems()
{
	optimizeBreakingItem();

	KnownState nextInitialState = m_state;
	if (m_breakingItem)
		nextInitialState.feedItem(*m_breakingItem);
	KnownState nextState = nextInitialState;

	ScopeGuard reset([&]()
	{
		m_breakingItem = nullptr;
		m_storeOperations.clear();
		m_initialState = std::move(nextInitialState);
		m_state = std::move(nextState);
	});

	std::map<int, Id> initialStackContents;
	std::map<int, Id> targetStackContents;
	int minHeight = m_state.stackHeight() + 1;
	if (!m_state.stackElements().empty())
		minHeight = std::min(minHeight, m_state.stackElements().begin()->first);
	for (int height = minHeight; height <= m_initialState.stackHeight(); ++height)
		initialStackContents[height] = m_initialState.stackElement(height, SourceLocation());
	for (int height = minHeight; height <= m_state.stackHeight(); ++height)
		targetStackContents[height] = m_state.stackElement(height, SourceLocation());

	AssemblyItems items = CSECodeGenerator(m_state.expressionClasses(), m_storeOperations).generateCode(
		m_initialState.sequenceNumber(),
		m_initialState.stackHeight(),
		initialStackContents,
		targetStackContents
	);
	if (m_breakingItem)
		items.push_back(*m_breakingItem);

	return items;
}

void CommonSubexpressionEliminator::feedItem(AssemblyItem const& _item, bool _copyItem)
{
	StoreOperation op = m_state.feedItem(_item, _copyItem);
	if (op.isValid())
		m_storeOperations.push_back(op);
}

void CommonSubexpressionEliminator::optimizeBreakingItem()
{
	if (!m_breakingItem)
		return;

	ExpressionClasses& classes = m_state.expressionClasses();
	SourceLocation const& itemLocation = m_breakingItem->location();
	if (*m_breakingItem == AssemblyItem(Instruction::JUMPI))
	{
		AssemblyItem::JumpType jumpType = m_breakingItem->getJumpType();

		Id condition = m_state.stackElement(m_state.stackHeight() - 1, itemLocation);
		if (classes.knownNonZero(condition))
		{
			feedItem(AssemblyItem(Instruction::SWAP1, itemLocation), true);
			feedItem(AssemblyItem(Instruction::POP, itemLocation), true);

			AssemblyItem item(Instruction::JUMP, itemLocation);
			item.setJumpType(jumpType);
			m_breakingItem = classes.storeItem(item);
		}
		else if (classes.knownZero(condition))
		{
			AssemblyItem it(Instruction::POP, itemLocation);
			feedItem(it, true);
			feedItem(it, true);
			m_breakingItem = nullptr;
		}
	}
	else if (*m_breakingItem == AssemblyItem(Instruction::RETURN))
	{
		Id size = m_state.stackElement(m_state.stackHeight() - 1, itemLocation);
		if (classes.knownZero(size))
		{
			feedItem(AssemblyItem(Instruction::POP, itemLocation), true);
			feedItem(AssemblyItem(Instruction::POP, itemLocation), true);
			AssemblyItem item(Instruction::STOP, itemLocation);
			m_breakingItem = classes.storeItem(item);
		}
	}
}

CSECodeGenerator::CSECodeGenerator(
	ExpressionClasses& _expressionClasses,
	std::vector<CSECodeGenerator::StoreOperation> const& _storeOperations
):
	m_expressionClasses(_expressionClasses)
{
	for (auto const& store: _storeOperations)
		m_storeOperations[std::make_pair(store.target, store.slot)].push_back(store);
}

AssemblyItems CSECodeGenerator::generateCode(
	unsigned _initialSequenceNumber,
	int _initialStackHeight,
	std::map<int, Id> const& _initialStack,
	std::map<int, Id> const& _targetStackContents
)
{
	m_stackHeight = _initialStackHeight;
	m_stack = _initialStack;
	m_targetStack = _targetStackContents;
	for (auto const& item: m_stack)
		m_classPositions[item.second].insert(item.first);

	// generate the dependency graph starting from final storage and memory writes and target stack contents
	for (auto const& p: m_storeOperations)
		addDependencies(p.second.back().expression);
	for (auto const& targetItem: m_targetStack)
	{
		m_finalClasses.insert(targetItem.second);
		addDependencies(targetItem.second);
	}

	// store all needed sequenced expressions
	std::set<std::pair<unsigned, Id>> sequencedExpressions;
	for (auto const& p: m_neededBy)
		for (auto id: {p.first, p.second})
			if (unsigned seqNr = m_expressionClasses.representative(id).sequenceNumber)
			{
				// Invalid sequenced operation.
				// @todo quick fix for now. Proper fix needs to choose representative with higher
				// sequence number during dependency analysis.
				assertThrow(seqNr >= _initialSequenceNumber, StackTooDeepException, util::stackTooDeepString);
				sequencedExpressions.insert(std::make_pair(seqNr, id));
			}

	// Perform all operations on storage and memory in order, if they are needed.
	for (auto const& seqAndId: sequencedExpressions)
		if (!m_classPositions.count(seqAndId.second))
			generateClassElement(seqAndId.second, true);

	// generate the target stack elements
	for (auto const& targetItem: m_targetStack)
	{
		if (m_stack.count(targetItem.first) && m_stack.at(targetItem.first) == targetItem.second)
			continue; // already there
		generateClassElement(targetItem.second);
		assertThrow(!m_classPositions[targetItem.second].empty(), OptimizerException, "");
		if (m_classPositions[targetItem.second].count(targetItem.first))
			continue;
		SourceLocation sourceLocation;
		if (m_expressionClasses.representative(targetItem.second).item)
			sourceLocation = m_expressionClasses.representative(targetItem.second).item->location();
		int position = classElementPosition(targetItem.second);
		if (position < targetItem.first)
			// it is already at its target, we need another copy
			appendDup(position, sourceLocation);
		else
			appendOrRemoveSwap(position, sourceLocation);
		appendOrRemoveSwap(targetItem.first, sourceLocation);
	}

	// remove surplus elements
	while (removeStackTopIfPossible())
	{
		// no-op
	}

	// check validity
	int finalHeight = 0;
	if (!m_targetStack.empty())
		// have target stack, so its height should be the final height
		finalHeight = (--m_targetStack.end())->first;
	else if (!_initialStack.empty())
		// no target stack, only erase the initial stack
		finalHeight = _initialStack.begin()->first - 1;
	else
		// neither initial no target stack, no change in height
		finalHeight = _initialStackHeight;
	assertThrow(finalHeight == m_stackHeight, OptimizerException, "Incorrect final stack height.");

	return m_generatedItems;
}

void CSECodeGenerator::addDependencies(Id _c)
{
	if (m_classPositions.count(_c))
		return; // it is already on the stack
	if (m_neededBy.find(_c) != m_neededBy.end())
		return; // we already computed the dependencies for _c
	ExpressionClasses::Expression expr = m_expressionClasses.representative(_c);
	assertThrow(expr.item, OptimizerException, "");
	// If this exception happens, we need to find a different way to generate the
	// compound expression.
	assertThrow(expr.item->type() != UndefinedItem, ItemNotAvailableException, "Undefined item requested but not available.");
	for (Id argument: expr.arguments)
	{
		addDependencies(argument);
		m_neededBy.insert(std::make_pair(argument, _c));
	}
	if (expr.item && expr.item->type() == Operation && (
		expr.item->instruction() == Instruction::SLOAD ||
		expr.item->instruction() == Instruction::MLOAD ||
		expr.item->instruction() == Instruction::KECCAK256
	))
	{
		// this loads an unknown value from storage or memory and thus, in addition to its
		// arguments, depends on all store operations to addresses where we do not know that
		// they are different that occur before this load
		StoreOperation::Target target = expr.item->instruction() == Instruction::SLOAD ?
			StoreOperation::Storage : StoreOperation::Memory;
		Id slotToLoadFrom = expr.arguments.at(0);
		for (auto const& p: m_storeOperations)
		{
			if (p.first.first != target)
				continue;
			Id slot = p.first.second;
			StoreOperations const& storeOps = p.second;
			if (storeOps.front().sequenceNumber > expr.sequenceNumber)
				continue;
			bool knownToBeIndependent = false;
			switch (expr.item->instruction())
			{
			case Instruction::SLOAD:
				knownToBeIndependent = m_expressionClasses.knownToBeDifferent(slot, slotToLoadFrom);
				break;
			case Instruction::MLOAD:
				knownToBeIndependent = m_expressionClasses.knownToBeDifferentBy32(slot, slotToLoadFrom);
				break;
			case Instruction::KECCAK256:
			{
				Id length = expr.arguments.at(1);
				AssemblyItem offsetInstr(Instruction::SUB, expr.item->location());
				Id offsetToStart = m_expressionClasses.find(offsetInstr, {slot, slotToLoadFrom});
				u256 const* o = m_expressionClasses.knownConstant(offsetToStart);
				u256 const* l = m_expressionClasses.knownConstant(length);
				if (l && *l == 0)
					knownToBeIndependent = true;
				else if (o)
				{
					// We could get problems here if both *o and *l are larger than 2**254
					// but it is probably ok for the optimizer to produce wrong code for such cases
					// which cannot be executed anyway because of the non-payable price.
					if (u2s(*o) <= -32)
						knownToBeIndependent = true;
					else if (l && u2s(*o) >= 0 && *o >= *l)
						knownToBeIndependent = true;
				}
				break;
			}
			default:
				break;
			}
			if (knownToBeIndependent)
				continue;

			// note that store and load never have the same sequence number
			Id latestStore = storeOps.front().expression;
			for (auto it = ++storeOps.begin(); it != storeOps.end(); ++it)
				if (it->sequenceNumber < expr.sequenceNumber)
					latestStore = it->expression;
			addDependencies(latestStore);
			m_neededBy.insert(std::make_pair(latestStore, _c));
		}
	}
}

void CSECodeGenerator::generateClassElement(Id _c, bool _allowSequenced)
{
	for (auto const& it: m_classPositions)
		for (int p: it.second)
			if (p > m_stackHeight)
			{
				assertThrow(false, OptimizerException, "");
			}
	// do some cleanup
	removeStackTopIfPossible();

	if (m_classPositions.count(_c))
	{
		assertThrow(
			!m_classPositions[_c].empty(),
			OptimizerException,
			"Element already removed but still needed."
		);
		return;
	}
	ExpressionClasses::Expression const& expr = m_expressionClasses.representative(_c);
	assertThrow(
		_allowSequenced || expr.sequenceNumber == 0,
		OptimizerException,
		"Sequence constrained operation requested out of sequence."
	);
	assertThrow(expr.item, OptimizerException, "Non-generated expression without item.");
	assertThrow(
		expr.item->type() != UndefinedItem,
		OptimizerException,
		"Undefined item requested but not available."
	);
	std::vector<Id> const& arguments = expr.arguments;
	for (Id arg: arguments | ranges::views::reverse)
		generateClassElement(arg);

	SourceLocation const& itemLocation = expr.item->location();
	// The arguments are somewhere on the stack now, so it remains to move them at the correct place.
	// This is quite difficult as sometimes, the values also have to removed in this process
	// (if canBeRemoved() returns true) and the two arguments can be equal. For now, this is
	// implemented for every single case for combinations of up to two arguments manually.
	if (arguments.size() == 1)
	{
		if (canBeRemoved(arguments[0], _c))
			appendOrRemoveSwap(classElementPosition(arguments[0]), itemLocation);
		else
			appendDup(classElementPosition(arguments[0]), itemLocation);
	}
	else if (arguments.size() == 2)
	{
		if (canBeRemoved(arguments[1], _c))
		{
			appendOrRemoveSwap(classElementPosition(arguments[1]), itemLocation);
			if (arguments[0] == arguments[1])
				appendDup(m_stackHeight, itemLocation);
			else if (canBeRemoved(arguments[0], _c))
			{
				appendOrRemoveSwap(m_stackHeight - 1, itemLocation);
				appendOrRemoveSwap(classElementPosition(arguments[0]), itemLocation);
			}
			else
				appendDup(classElementPosition(arguments[0]), itemLocation);
		}
		else
		{
			if (arguments[0] == arguments[1])
			{
				appendDup(classElementPosition(arguments[0]), itemLocation);
				appendDup(m_stackHeight, itemLocation);
			}
			else if (canBeRemoved(arguments[0], _c))
			{
				appendOrRemoveSwap(classElementPosition(arguments[0]), itemLocation);
				appendDup(classElementPosition(arguments[1]), itemLocation);
				appendOrRemoveSwap(m_stackHeight - 1, itemLocation);
			}
			else
			{
				appendDup(classElementPosition(arguments[1]), itemLocation);
				appendDup(classElementPosition(arguments[0]), itemLocation);
			}
		}
	}
	else
		assertThrow(
			arguments.size() <= 2,
			OptimizerException,
			"Opcodes with more than two arguments not implemented yet."
		);
	for (size_t i = 0; i < arguments.size(); ++i)
		assertThrow(
			m_stack[m_stackHeight - static_cast<int>(i)] == arguments[i],
			OptimizerException,
			"Expected arguments not present."
		);

	while (SemanticInformation::isCommutativeOperation(*expr.item) &&
			!m_generatedItems.empty() &&
			m_generatedItems.back() == AssemblyItem(Instruction::SWAP1))
		// this will not append a swap but remove the one that is already there
		appendOrRemoveSwap(m_stackHeight - 1, itemLocation);
	for (size_t i = 0; i < arguments.size(); ++i)
	{
		m_classPositions[m_stack[m_stackHeight - static_cast<int>(i)]].erase(m_stackHeight - static_cast<int>(i));
		m_stack.erase(m_stackHeight - static_cast<int>(i));
	}
	appendItem(*expr.item);
	if (expr.item->type() != Operation || instructionInfo(expr.item->instruction(), EVMVersion()).ret == 1)
	{
		m_stack[m_stackHeight] = _c;
		m_classPositions[_c].insert(m_stackHeight);
	}
	else
	{
		assertThrow(
			instructionInfo(expr.item->instruction(), EVMVersion()).ret == 0,
			OptimizerException,
			"Invalid number of return values."
		);
		m_classPositions[_c]; // ensure it is created to mark the expression as generated
	}
}

int CSECodeGenerator::classElementPosition(Id _id) const
{
	assertThrow(
		m_classPositions.count(_id) && !m_classPositions.at(_id).empty(),
		OptimizerException,
		"Element requested but is not present."
	);
	return *max_element(m_classPositions.at(_id).begin(), m_classPositions.at(_id).end());
}

bool CSECodeGenerator::canBeRemoved(Id _element, Id _result, int _fromPosition)
{
	// Default for _fromPosition is the canonical position of the element.
	if (_fromPosition == c_invalidPosition)
		_fromPosition = classElementPosition(_element);

	bool haveCopy = m_classPositions.at(_element).size() > 1;
	if (m_finalClasses.count(_element))
		// It is part of the target stack. It can be removed if it is a copy that is not in the target position.
		return haveCopy && (!m_targetStack.count(_fromPosition) || m_targetStack[_fromPosition] != _element);
	else if (!haveCopy)
	{
		// Can be removed unless it is needed by a class that has not been computed yet.
		// Note that m_classPositions also includes classes that were deleted in the meantime.
		auto range = m_neededBy.equal_range(_element);
		for (auto it = range.first; it != range.second; ++it)
			if (it->second != _result && !m_classPositions.count(it->second))
				return false;
	}
	return true;
}

bool CSECodeGenerator::removeStackTopIfPossible()
{
	if (m_stack.empty())
		return false;
	assertThrow(m_stack.count(m_stackHeight) > 0, OptimizerException, "");
	Id top = m_stack[m_stackHeight];
	if (!canBeRemoved(top, Id(-1), m_stackHeight))
		return false;
	m_classPositions[m_stack[m_stackHeight]].erase(m_stackHeight);
	m_stack.erase(m_stackHeight);
	appendItem(AssemblyItem(Instruction::POP));
	return true;
}

void CSECodeGenerator::appendDup(int _fromPosition, SourceLocation const& _location)
{
	assertThrow(_fromPosition != c_invalidPosition, OptimizerException, "");
	int instructionNum = 1 + m_stackHeight - _fromPosition;
	assertThrow(instructionNum <= 16, StackTooDeepException, util::stackTooDeepString);
	assertThrow(1 <= instructionNum, OptimizerException, "Invalid stack access.");
	appendItem(AssemblyItem(dupInstruction(static_cast<unsigned>(instructionNum)), _location));
	m_stack[m_stackHeight] = m_stack[_fromPosition];
	m_classPositions[m_stack[m_stackHeight]].insert(m_stackHeight);
}

void CSECodeGenerator::appendOrRemoveSwap(int _fromPosition, SourceLocation const& _location)
{
	assertThrow(_fromPosition != c_invalidPosition, OptimizerException, "");
	if (_fromPosition == m_stackHeight)
		return;
	int instructionNum = m_stackHeight - _fromPosition;
	assertThrow(instructionNum <= 16, StackTooDeepException, util::stackTooDeepString);
	assertThrow(1 <= instructionNum, OptimizerException, "Invalid stack access.");
	appendItem(AssemblyItem(swapInstruction(static_cast<unsigned>(instructionNum)), _location));

	if (m_stack[m_stackHeight] != m_stack[_fromPosition])
	{
		m_classPositions[m_stack[m_stackHeight]].erase(m_stackHeight);
		m_classPositions[m_stack[m_stackHeight]].insert(_fromPosition);
		m_classPositions[m_stack[_fromPosition]].erase(_fromPosition);
		m_classPositions[m_stack[_fromPosition]].insert(m_stackHeight);
		std::swap(m_stack[m_stackHeight], m_stack[_fromPosition]);
	}
	if (m_generatedItems.size() >= 2 &&
		SemanticInformation::isSwapInstruction(m_generatedItems.back()) &&
		*(m_generatedItems.end() - 2) == m_generatedItems.back())
	{
		m_generatedItems.pop_back();
		m_generatedItems.pop_back();
	}
}

void CSECodeGenerator::appendItem(AssemblyItem const& _item)
{
	m_generatedItems.push_back(_item);
	m_stackHeight += static_cast<int>(_item.deposit());
}