Reuse stack slots in Yul to EVM code generation.

4 files changed, 277 insertions, 32 deletions

diff --git a/libyul/backends/evm/EVMCodeTransform.cpp b/libyul/backends/evm/EVMCodeTransform.cpp
index 12abd754..9d8e9a06 100644
--- a/libyul/backends/evm/EVMCodeTransform.cpp
+++ b/libyul/backends/evm/EVMCodeTransform.cpp
@@ -20,6 +20,7 @@
 
 #include <libyul/backends/evm/EVMCodeTransform.h>
 
+#include <libyul/optimiser/NameCollector.h>
 #include <libyul/AsmAnalysisInfo.h>
 #include <libyul/AsmData.h>
 
@@ -32,6 +33,144 @@ using namespace dev;
 using namespace yul;
 using namespace dev::solidity;
 
+void VariableReferenceCounter::operator()(Identifier const& _identifier)
+{
+	increaseRefIfFound(_identifier.name);
+}
+
+void VariableReferenceCounter::operator()(FunctionDefinition const& _function)
+{
+	Scope* originalScope = m_scope;
+
+	solAssert(m_info.virtualBlocks.at(&_function), "");
+	m_scope = m_info.scopes.at(m_info.virtualBlocks.at(&_function).get()).get();
+	solAssert(m_scope, "Variable scope does not exist.");
+
+	for (auto const& v: _function.returnVariables)
+		increaseRefIfFound(v.name);
+
+	VariableReferenceCounter{m_context, m_info}(_function.body);
+
+	m_scope = originalScope;
+}
+
+void VariableReferenceCounter::operator()(ForLoop const& _forLoop)
+{
+	Scope* originalScope = m_scope;
+	// Special scoping rules.
+	m_scope = m_info.scopes.at(&_forLoop.pre).get();
+
+	walkVector(_forLoop.pre.statements);
+	visit(*_forLoop.condition);
+	(*this)(_forLoop.body);
+	(*this)(_forLoop.post);
+
+	m_scope = originalScope;
+}
+
+
+void VariableReferenceCounter::operator()(Block const& _block)
+{
+	Scope* originalScope = m_scope;
+	m_scope = m_info.scopes.at(&_block).get();
+
+	ASTWalker::operator()(_block);
+
+	m_scope = originalScope;
+}
+
+void VariableReferenceCounter::increaseRefIfFound(YulString _variableName)
+{
+	m_scope->lookup(_variableName, Scope::Visitor(
+		[=](Scope::Variable const& _var)
+		{
+			++m_context.variableReferences[&_var];
+		},
+		[=](Scope::Label const&) { },
+		[=](Scope::Function const&) { }
+	));
+}
+
+
+CodeTransform::CodeTransform(
+	AbstractAssembly& _assembly,
+	AsmAnalysisInfo& _analysisInfo,
+	Block const& _block,
+	bool _allowStackOpt,
+	bool _yul,
+	bool _evm15,
+	ExternalIdentifierAccess const& _identifierAccess,
+	bool _useNamedLabelsForFunctions,
+	int _stackAdjustment,
+	shared_ptr<Context> _context
+):
+	m_assembly(_assembly),
+	m_info(_analysisInfo),
+	m_allowStackOpt(_allowStackOpt),
+	m_yul(_yul),
+	m_evm15(_evm15),
+	m_useNamedLabelsForFunctions(_useNamedLabelsForFunctions),
+	m_identifierAccess(_identifierAccess),
+	m_stackAdjustment(_stackAdjustment),
+	m_context(_context)
+{
+	if (!m_context)
+	{
+		// initialize
+		m_context = make_shared<Context>();
+		if (m_allowStackOpt)
+			VariableReferenceCounter{*m_context, m_info}(_block);
+	}
+}
+
+void CodeTransform::decreaseReference(YulString, Scope::Variable const& _var)
+{
+	if (!m_allowStackOpt)
+		return;
+
+	unsigned& ref = m_context->variableReferences.at(&_var);
+	solAssert(ref >= 1, "");
+	--ref;
+	if (ref == 0)
+		m_variablesScheduledForDeletion.insert(&_var);
+}
+
+bool CodeTransform::unreferenced(Scope::Variable const& _var) const
+{
+	return !m_context->variableReferences.count(&_var) || m_context->variableReferences[&_var] == 0;
+}
+
+void CodeTransform::freeUnusedVariables()
+{
+	if (!m_allowStackOpt)
+		return;
+
+	for (auto const& identifier: m_scope->identifiers)
+		if (identifier.second.type() == typeid(Scope::Variable))
+		{
+			Scope::Variable const& var = boost::get<Scope::Variable>(identifier.second);
+			if (m_variablesScheduledForDeletion.count(&var))
+				deleteVariable(var);
+		}
+
+	while (m_unusedStackSlots.count(m_assembly.stackHeight() - 1))
+	{
+		solAssert(m_unusedStackSlots.erase(m_assembly.stackHeight() - 1), "");
+		m_assembly.appendInstruction(solidity::Instruction::POP);
+		--m_stackAdjustment;
+	}
+}
+
+void CodeTransform::deleteVariable(Scope::Variable const& _var)
+{
+	solAssert(m_allowStackOpt, "");
+	solAssert(m_context->variableStackHeights.count(&_var) > 0, "");
+	m_unusedStackSlots.insert(m_context->variableStackHeights[&_var]);
+	m_context->variableStackHeights.erase(&_var);
+	m_context->variableReferences.erase(&_var);
+	m_variablesScheduledForDeletion.erase(&_var);
+}
+
 void CodeTransform::operator()(VariableDeclaration const& _varDecl)
 {
 	solAssert(m_scope, "");
@@ -49,10 +188,40 @@ void CodeTransform::operator()(VariableDeclaration const& _varDecl)
 		while (variablesLeft--)
 			m_assembly.appendConstant(u256(0));
 	}
-	for (auto const& variable: _varDecl.variables)
+
+	bool atTopOfStack = true;
+	for (int varIndex = numVariables - 1; varIndex >= 0; --varIndex)
 	{
-		auto& var = boost::get<Scope::Variable>(m_scope->identifiers.at(variable.name));
-		m_context->variableStackHeights[&var] = height++;
+		auto& var = boost::get<Scope::Variable>(m_scope->identifiers.at(_varDecl.variables[varIndex].name));
+		m_context->variableStackHeights[&var] = height + varIndex;
+		if (!m_allowStackOpt)
+			continue;
+
+		if (unreferenced(var))
+		{
+			if (atTopOfStack)
+			{
+				m_context->variableStackHeights.erase(&var);
+				m_assembly.setSourceLocation(_varDecl.location);
+				m_assembly.appendInstruction(solidity::Instruction::POP);
+				--m_stackAdjustment;
+			}
+			else
+				m_variablesScheduledForDeletion.insert(&var);
+		}
+		else if (m_unusedStackSlots.empty())
+			atTopOfStack = false;
+		else
+		{
+			int slot = *m_unusedStackSlots.begin();
+			m_unusedStackSlots.erase(m_unusedStackSlots.begin());
+			m_context->variableStackHeights[&var] = slot;
+			m_assembly.setSourceLocation(_varDecl.location);
+			if (int heightDiff = variableHeightDiff(var, true))
+				m_assembly.appendInstruction(solidity::swapInstruction(heightDiff - 1));
+			m_assembly.appendInstruction(solidity::Instruction::POP);
+			--m_stackAdjustment;
+		}
 	}
 	checkStackHeight(&_varDecl);
 }
@@ -70,6 +239,7 @@ void CodeTransform::operator()(Assignment const& _assignment)
 
 void CodeTransform::operator()(StackAssignment const& _assignment)
 {
+	solAssert(!m_allowStackOpt, "");
 	m_assembly.setSourceLocation(_assignment.location);
 	generateAssignment(_assignment.variableName);
 	checkStackHeight(&_assignment);
@@ -84,6 +254,7 @@ void CodeTransform::operator()(ExpressionStatement const& _statement)
 
 void CodeTransform::operator()(Label const& _label)
 {
+	solAssert(!m_allowStackOpt, "");
 	m_assembly.setSourceLocation(_label.location);
 	solAssert(m_scope, "");
 	solAssert(m_scope->identifiers.count(_label.name), "");
@@ -169,11 +340,14 @@ void CodeTransform::operator()(Identifier const& _identifier)
 	if (m_scope->lookup(_identifier.name, Scope::NonconstVisitor(
 		[=](Scope::Variable& _var)
 		{
+			// TODO: opportunity for optimization: Do not DUP if this is the last reference
+			// to the top most element of the stack
 			if (int heightDiff = variableHeightDiff(_var, false))
 				m_assembly.appendInstruction(solidity::dupInstruction(heightDiff));
 			else
 				// Store something to balance the stack
 				m_assembly.appendConstant(u256(0));
+			decreaseReference(_identifier.name, _var);
 		},
 		[=](Scope::Label& _label)
 		{
@@ -217,6 +391,7 @@ void CodeTransform::operator()(Literal const& _literal)
 
 void CodeTransform::operator()(yul::Instruction const& _instruction)
 {
+	solAssert(!m_allowStackOpt, "");
 	solAssert(!m_evm15 || _instruction.instruction != solidity::Instruction::JUMP, "Bare JUMP instruction used for EVM1.5");
 	solAssert(!m_evm15 || _instruction.instruction != solidity::Instruction::JUMPI, "Bare JUMPI instruction used for EVM1.5");
 	m_assembly.setSourceLocation(_instruction.location);
@@ -329,6 +504,8 @@ void CodeTransform::operator()(FunctionDefinition const& _function)
 	CodeTransform(
 		m_assembly,
 		m_info,
+		_function.body,
+		m_allowStackOpt,
 		m_yul,
 		m_evm15,
 		m_identifierAccess,
@@ -350,6 +527,7 @@ void CodeTransform::operator()(FunctionDefinition const& _function)
 		if (!m_evm15)
 			stackLayout.push_back(_function.returnVariables.size()); // Move return label to the top
 		stackLayout += vector<int>(_function.parameters.size(), -1); // discard all arguments
+
 		for (size_t i = 0; i < _function.returnVariables.size(); ++i)
 			stackLayout.push_back(i); // Move return values down, but keep order.
 
@@ -475,20 +653,37 @@ void CodeTransform::visitExpression(Expression const& _expression)
 void CodeTransform::visitStatements(vector<Statement> const& _statements)
 {
 	for (auto const& statement: _statements)
+	{
+		freeUnusedVariables();
 		boost::apply_visitor(*this, statement);
+	}
+	freeUnusedVariables();
 }
 
 void CodeTransform::finalizeBlock(Block const& _block, int blockStartStackHeight)
 {
 	m_assembly.setSourceLocation(_block.location);
 
+	freeUnusedVariables();
+
 	// pop variables
 	solAssert(m_info.scopes.at(&_block).get() == m_scope, "");
-	for (size_t i = 0; i < m_scope->numberOfVariables(); ++i)
-		m_assembly.appendInstruction(solidity::Instruction::POP);
+	for (auto const& id: m_scope->identifiers)
+		if (id.second.type() == typeid(Scope::Variable))
+		{
+			Scope::Variable const& var = boost::get<Scope::Variable>(id.second);
+			if (m_allowStackOpt)
+			{
+				solAssert(!m_context->variableStackHeights.count(&var), "");
+				solAssert(!m_context->variableReferences.count(&var), "");
+				m_stackAdjustment++;
+			}
+			else
+				m_assembly.appendInstruction(solidity::Instruction::POP);
+		}
 
 	int deposit = m_assembly.stackHeight() - blockStartStackHeight;
-	solAssert(deposit == 0, "Invalid stack height at end of block.");
+	solAssert(deposit == 0, "Invalid stack height at end of block: " + to_string(deposit));
 	checkStackHeight(&_block);
 }
 
@@ -502,13 +697,13 @@ void CodeTransform::generateMultiAssignment(vector<Identifier> const& _variableN
 void CodeTransform::generateAssignment(Identifier const& _variableName)
 {
 	solAssert(m_scope, "");
-	auto var = m_scope->lookup(_variableName.name);
-	if (var)
+	if (auto var = m_scope->lookup(_variableName.name))
 	{
 		Scope::Variable const& _var = boost::get<Scope::Variable>(*var);
 		if (int heightDiff = variableHeightDiff(_var, true))
 			m_assembly.appendInstruction(solidity::swapInstruction(heightDiff - 1));
 		m_assembly.appendInstruction(solidity::Instruction::POP);
+		decreaseReference(_variableName.name, _var);
 	}
 	else
 	{
diff --git a/libyul/backends/evm/EVMCodeTransform.h b/libyul/backends/evm/EVMCodeTransform.h
index d559f85a..8927e999 100644
--- a/libyul/backends/evm/EVMCodeTransform.h
+++ b/libyul/backends/evm/EVMCodeTransform.h
@@ -20,6 +20,7 @@
 
 #include <libyul/backends/evm/EVMAssembly.h>
 
+#include <libyul/optimiser/ASTWalker.h>
 #include <libyul/AsmDataForward.h>
 
 #include <libyul/AsmScope.h>
@@ -37,6 +38,46 @@ namespace yul
 struct AsmAnalysisInfo;
 class EVMAssembly;
 
+struct CodeTransformContext
+{
+	std::map<Scope::Label const*, AbstractAssembly::LabelID> labelIDs;
+	std::map<Scope::Function const*, AbstractAssembly::LabelID> functionEntryIDs;
+	std::map<Scope::Variable const*, int> variableStackHeights;
+	std::map<Scope::Variable const*, unsigned> variableReferences;
+};
+
+/**
+ * Counts the number of references to a variable. This includes actual (read) references
+ * but also assignments to the variable. It does not include the declaration itself or
+ * function parameters, but it does include function return parameters.
+ *
+ * This component can handle multiple variables of the same name.
+ *
+ * Can only be applied to strict assembly.
+ */
+class VariableReferenceCounter: public yul::ASTWalker
+{
+public:
+	explicit VariableReferenceCounter(
+		CodeTransformContext& _context,
+		AsmAnalysisInfo const& _assemblyInfo
+	): m_context(_context), m_info(_assemblyInfo)
+	{}
+
+public:
+	void operator()(Identifier const& _identifier);
+	void operator()(FunctionDefinition const&);
+	void operator()(ForLoop const&);
+	void operator()(Block const& _block);
+
+private:
+	void increaseRefIfFound(YulString _variableName);
+
+	CodeTransformContext& m_context;
+	AsmAnalysisInfo const& m_info;
+	Scope* m_scope = nullptr;
+};
+
 class CodeTransform: public boost::static_visitor<>
 {
 public:
@@ -45,6 +86,8 @@ public:
 	CodeTransform(
 		AbstractAssembly& _assembly,
 		AsmAnalysisInfo& _analysisInfo,
+		Block const& _block,
+		bool _allowStackOpt = false,
 		bool _yul = false,
 		bool _evm15 = false,
 		ExternalIdentifierAccess const& _identifierAccess = ExternalIdentifierAccess(),
@@ -52,43 +95,42 @@ public:
 	): CodeTransform(
 		_assembly,
 		_analysisInfo,
+		_block,
+		_allowStackOpt,
 		_yul,
 		_evm15,
 		_identifierAccess,
 		_useNamedLabelsForFunctions,
 		_assembly.stackHeight(),
-		std::make_shared<Context>()
+		nullptr
 	)
 	{
 	}
 
 protected:
-	struct Context
-	{
-		std::map<Scope::Label const*, AbstractAssembly::LabelID> labelIDs;
-		std::map<Scope::Function const*, AbstractAssembly::LabelID> functionEntryIDs;
-		std::map<Scope::Variable const*, int> variableStackHeights;
-	};
+	using Context = CodeTransformContext;
 
 	CodeTransform(
 		AbstractAssembly& _assembly,
 		AsmAnalysisInfo& _analysisInfo,
+		Block const& _block,
+		bool _allowStackOpt,
 		bool _yul,
 		bool _evm15,
 		ExternalIdentifierAccess const& _identifierAccess,
 		bool _useNamedLabelsForFunctions,
 		int _stackAdjustment,
 		std::shared_ptr<Context> _context
-	):
-		m_assembly(_assembly),
-		m_info(_analysisInfo),
-		m_yul(_yul),
-		m_evm15(_evm15),
-		m_useNamedLabelsForFunctions(_useNamedLabelsForFunctions),
-		m_identifierAccess(_identifierAccess),
-		m_stackAdjustment(_stackAdjustment),
-		m_context(_context)
-	{}
+	);
+
+	void decreaseReference(YulString _name, Scope::Variable const& _var);
+	bool unreferenced(Scope::Variable const& _var) const;
+	/// Marks slots of variables that are not used anymore
+	/// and were defined in the current scope for reuse.
+	/// Also POPs unused topmost stack slots.
+	void freeUnusedVariables();
+	/// Marks the stack slot of @a _var to be reused.
+	void deleteVariable(Scope::Variable const& _var);
 
 public:
 	void operator()(Instruction const& _instruction);
@@ -137,6 +179,7 @@ private:
 	AbstractAssembly& m_assembly;
 	AsmAnalysisInfo& m_info;
 	Scope* m_scope = nullptr;
+	bool const m_allowStackOpt = true;
 	bool m_yul = false;
 	bool m_evm15 = false;
 	bool m_useNamedLabelsForFunctions = false;
@@ -147,6 +190,12 @@ private:
 	/// (EVM 1.0 or 1.5).
 	int m_stackAdjustment = 0;
 	std::shared_ptr<Context> m_context;
+
+	/// Set of variables whose reference counter has reached zero,
+	/// and whose stack slot will be marked as unused once we reach
+	/// statement level in the scope where the variable was defined.
+	std::set<Scope::Variable const*> m_variablesScheduledForDeletion;
+	std::set<int> m_unusedStackSlots;
 };
 
 }
diff --git a/libyul/backends/evm/EVMObjectCompiler.cpp b/libyul/backends/evm/EVMObjectCompiler.cpp
index e7e8ad99..13d4b756 100644
--- a/libyul/backends/evm/EVMObjectCompiler.cpp
+++ b/libyul/backends/evm/EVMObjectCompiler.cpp
@@ -27,13 +27,13 @@
 using namespace yul;
 using namespace std;
 
-void EVMObjectCompiler::compile(Object& _object, AbstractAssembly& _assembly, bool _yul, bool _evm15)
+void EVMObjectCompiler::compile(Object& _object, AbstractAssembly& _assembly, bool _yul, bool _evm15, bool _optimize)
 {
 	EVMObjectCompiler compiler(_assembly, _yul, _evm15);
-	compiler.run(_object);
+	compiler.run(_object, _optimize);
 }
 
-void EVMObjectCompiler::run(Object& _object)
+void EVMObjectCompiler::run(Object& _object, bool _optimize)
 {
 	map<YulString, AbstractAssembly::SubID> subIDs;
 
@@ -42,7 +42,7 @@ void EVMObjectCompiler::run(Object& _object)
 		{
 			auto subAssemblyAndID = m_assembly.createSubAssembly();
 			subIDs[subObject->name] = subAssemblyAndID.second;
-			compile(*subObject, *subAssemblyAndID.first, m_yul, m_evm15);
+			compile(*subObject, *subAssemblyAndID.first, m_yul, m_evm15, _optimize);
 		}
 		else
 		{
@@ -51,5 +51,6 @@ void EVMObjectCompiler::run(Object& _object)
 		}
 
 	yulAssert(_object.analysisInfo, "No analysis info.");
-	CodeTransform{m_assembly, *_object.analysisInfo, m_yul, m_evm15}(*_object.code);
+	yulAssert(_object.code, "No code.");
+	CodeTransform{m_assembly, *_object.analysisInfo, *_object.code, _optimize, m_yul, m_evm15}(*_object.code);
 }
diff --git a/libyul/backends/evm/EVMObjectCompiler.h b/libyul/backends/evm/EVMObjectCompiler.h
index c7172e47..826b82a4 100644
--- a/libyul/backends/evm/EVMObjectCompiler.h
+++ b/libyul/backends/evm/EVMObjectCompiler.h
@@ -27,13 +27,13 @@ class AbstractAssembly;
 class EVMObjectCompiler
 {
 public:
-	static void compile(Object& _object, AbstractAssembly& _assembly, bool _yul, bool _evm15);
+	static void compile(Object& _object, AbstractAssembly& _assembly, bool _yul, bool _evm15, bool _optimize);
 private:
 	EVMObjectCompiler(AbstractAssembly& _assembly, bool _yul, bool _evm15):
 		m_assembly(_assembly), m_yul(_yul), m_evm15(_evm15)
 	{}
 
-	void run(Object& _object);
+	void run(Object& _object, bool _optimize);
 
 	AbstractAssembly& m_assembly;
 	bool m_yul = false;