Merge pull request #5008 from liangdzou/yul_stack_reuse

Reuse stack slots in Yul codegen

libsolidity/codegen/AsmCodeGen.cpp

@@ -179,13 +179,16 @@ void CodeGenerator::assemble(
 	AsmAnalysisInfo& _analysisInfo,
 	eth::Assembly& _assembly,
 	ExternalIdentifierAccess const& _identifierAccess,
-	bool _useNamedLabelsForFunctions
+	bool _useNamedLabelsForFunctions,
+	bool _optimize
 )
 {
 	EthAssemblyAdapter assemblyAdapter(_assembly);
 	CodeTransform(
 		assemblyAdapter,
 		_analysisInfo,
+		_parsedData,
+		_optimize,
 		false,
 		false,
 		_identifierAccess,

libsolidity/codegen/AsmCodeGen.h

@@ -85,7 +85,8 @@ public:
 		yul::AsmAnalysisInfo& _analysisInfo,
 		dev::eth::Assembly& _assembly,
 		yul::ExternalIdentifierAccess const& _identifierAccess = yul::ExternalIdentifierAccess(),
-		bool _useNamedLabelsForFunctions = false
+		bool _useNamedLabelsForFunctions = false,
+		bool _optimize = false
 	);
 };
 

libsolidity/interface/AssemblyStack.cpp

@@ -122,7 +122,7 @@ void AssemblyStack::optimize(yul::Object& _object)
 	yul::OptimiserSuite::run(*_object.code, *_object.analysisInfo);
 }
 
-MachineAssemblyObject AssemblyStack::assemble(Machine _machine) const
+MachineAssemblyObject AssemblyStack::assemble(Machine _machine, bool _optimize) const
 {
 	solAssert(m_analysisSuccessful, "");
 	solAssert(m_parserResult, "");
@@ -136,7 +136,7 @@ MachineAssemblyObject AssemblyStack::assemble(Machine _machine) const
 		MachineAssemblyObject object;
 		eth::Assembly assembly;
 		EthAssemblyAdapter adapter(assembly);
-		yul::EVMObjectCompiler::compile(*m_parserResult, adapter, m_language == Language::Yul, false);
+		yul::EVMObjectCompiler::compile(*m_parserResult, adapter, m_language == Language::Yul, false, _optimize);
 		object.bytecode = make_shared<eth::LinkerObject>(assembly.assemble());
 		object.assembly = assembly.assemblyString();
 		return object;
@@ -145,7 +145,7 @@ MachineAssemblyObject AssemblyStack::assemble(Machine _machine) const
 	{
 		MachineAssemblyObject object;
 		yul::EVMAssembly assembly(true);
-		yul::EVMObjectCompiler::compile(*m_parserResult, assembly, m_language == Language::Yul, true);
+		yul::EVMObjectCompiler::compile(*m_parserResult, assembly, m_language == Language::Yul, true, _optimize);
 		object.bytecode = make_shared<eth::LinkerObject>(assembly.finalize());
 		/// TODO: fill out text representation
 		return object;

libsolidity/interface/AssemblyStack.h

@@ -73,7 +73,8 @@ public:
 	void optimize();
 
 	/// Run the assembly step (should only be called after parseAndAnalyze).
-	MachineAssemblyObject assemble(Machine _machine) const;
+	/// @param _optimize does not run the optimizer but performs optimized code generation.
+	MachineAssemblyObject assemble(Machine _machine, bool _optimize = false) const;
 
 	/// @returns the errors generated during parsing, analysis (and potentially assembly).
 	langutil::ErrorList const& errors() const { return m_errors; }

libyul/Exceptions.h

@@ -28,6 +28,7 @@ namespace yul
 
 struct YulException: virtual dev::Exception {};
 struct OptimizerException: virtual YulException {};
+struct CodegenException: virtual YulException {};
 struct YulAssertion: virtual YulException {};
 
 /// Assertion that throws an YulAssertion containing the given description if it is not met.

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
 	{

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;
 };
 
 }

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);
 }

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;

libyul/optimiser/ASTWalker.cpp

@@ -93,6 +93,16 @@ void ASTWalker::operator()(Block const& _block)
 	walkVector(_block.statements);
 }
 
+void ASTWalker::visit(Statement const& _st)
+{
+	boost::apply_visitor(*this, _st);
+}
+
+void ASTWalker::visit(Expression const& _e)
+{
+	boost::apply_visitor(*this, _e);
+}
+
 void ASTModifier::operator()(FunctionalInstruction& _instr)
 {
 	walkVector(_instr.arguments | boost::adaptors::reversed);
@@ -155,3 +165,13 @@ void ASTModifier::operator()(Block& _block)
 {
 	walkVector(_block.statements);
 }
+
+void ASTModifier::visit(Statement& _st)
+{
+	boost::apply_visitor(*this, _st);
+}
+
+void ASTModifier::visit(Expression& _e)
+{
+	boost::apply_visitor(*this, _e);
+}

libyul/optimiser/ASTWalker.h

@@ -58,14 +58,8 @@ public:
 	virtual void operator()(ForLoop const&);
 	virtual void operator()(Block const& _block);
 
-	virtual void visit(Statement const& _st)
-	{
-		boost::apply_visitor(*this, _st);
-	}
-	virtual void visit(Expression const& _e)
-	{
-		boost::apply_visitor(*this, _e);
-	}
+	virtual void visit(Statement const& _st);
+	virtual void visit(Expression const& _e);
 
 protected:
 	template <class T>
@@ -99,14 +93,8 @@ public:
 	virtual void operator()(ForLoop&);
 	virtual void operator()(Block& _block);
 
-	virtual void visit(Statement& _st)
-	{
-		boost::apply_visitor(*this, _st);
-	}
-	virtual void visit(Expression& _e)
-	{
-		boost::apply_visitor(*this, _e);
-	}
+	virtual void visit(Statement& _st);
+	virtual void visit(Expression& _e);
 
 protected:
 	template <class T>

test/libyul/StackReuseCodegen.cpp

@@ -0,0 +1,353 @@
+/*
+    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/>.
+*/
+/**
+ * Unit tests for stack-reusing code generator.
+ */
+
+#include <test/Options.h>
+
+#include <libsolidity/interface/AssemblyStack.h>
+#include <libevmasm/Instruction.h>
+
+using namespace std;
+
+namespace dev
+{
+namespace yul
+{
+namespace test
+{
+
+namespace
+{
+string assemble(string const& _input)
+{
+	solidity::AssemblyStack asmStack;
+	BOOST_REQUIRE_MESSAGE(asmStack.parseAndAnalyze("", _input), "Source did not parse: " + _input);
+	return solidity::disassemble(asmStack.assemble(solidity::AssemblyStack::Machine::EVM, true).bytecode->bytecode);
+}
+}
+
+BOOST_AUTO_TEST_SUITE(StackReuseCodegen)
+
+BOOST_AUTO_TEST_CASE(smoke_test)
+{
+	string out = assemble("{}");
+	BOOST_CHECK_EQUAL(out, "");
+}
+
+BOOST_AUTO_TEST_CASE(single_var)
+{
+	string out = assemble("{ let x }");
+	BOOST_CHECK_EQUAL(out, "PUSH1 0x0 POP ");
+}
+
+BOOST_AUTO_TEST_CASE(single_var_assigned)
+{
+	string out = assemble("{ let x := 1 }");
+	BOOST_CHECK_EQUAL(out, "PUSH1 0x1 POP ");
+}
+
+BOOST_AUTO_TEST_CASE(single_var_assigned_plus_code)
+{
+	string out = assemble("{ let x := 1 mstore(3, 4) }");
+	BOOST_CHECK_EQUAL(out, "PUSH1 0x1 POP PUSH1 0x4 PUSH1 0x3 MSTORE ");
+}
+
+BOOST_AUTO_TEST_CASE(single_var_assigned_plus_code_and_reused)
+{
+	string out = assemble("{ let x := 1 mstore(3, 4) pop(mload(x)) }");
+	BOOST_CHECK_EQUAL(out, "PUSH1 0x1 PUSH1 0x4 PUSH1 0x3 MSTORE DUP1 MLOAD POP POP ");
+}
+
+BOOST_AUTO_TEST_CASE(multi_reuse_single_slot)
+{
+	string out = assemble("{ let x := 1 x := 6 let y := 2 y := 4 }");
+	BOOST_CHECK_EQUAL(out, "PUSH1 0x1 PUSH1 0x6 SWAP1 POP POP PUSH1 0x2 PUSH1 0x4 SWAP1 POP POP ");
+}
+
+BOOST_AUTO_TEST_CASE(multi_reuse_single_slot_nested)
+{
+	string out = assemble("{ let x := 1 x := 6 { let y := 2 y := 4 } }");
+	BOOST_CHECK_EQUAL(out, "PUSH1 0x1 PUSH1 0x6 SWAP1 POP POP PUSH1 0x2 PUSH1 0x4 SWAP1 POP POP ");
+}
+
+BOOST_AUTO_TEST_CASE(multi_reuse_same_variable_name)
+{
+	string out = assemble("{ let z := mload(0) { let x := 1 x := 6 z := x } { let x := 2 z := x x := 4 } }");
+	BOOST_CHECK_EQUAL(out,
+		"PUSH1 0x0 MLOAD "
+		"PUSH1 0x1 PUSH1 0x6 SWAP1 POP DUP1 SWAP2 POP POP "
+		"PUSH1 0x2 DUP1 SWAP2 POP PUSH1 0x4 SWAP1 POP POP "
+		"POP "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(last_use_in_nested_block)
+{
+	string out = assemble("{ let z := 0 { pop(z) } let x := 1 }");
+	BOOST_CHECK_EQUAL(out, "PUSH1 0x0 DUP1 POP POP PUSH1 0x1 POP ");
+}
+
+BOOST_AUTO_TEST_CASE(if_)
+{
+	// z is only removed after the if (after the jumpdest)
+	string out = assemble("{ let z := mload(0) if z { let x := z } let t := 3 }");
+	BOOST_CHECK_EQUAL(out, "PUSH1 0x0 MLOAD DUP1 ISZERO PUSH1 0xA JUMPI DUP1 POP JUMPDEST POP PUSH1 0x3 POP ");
+}
+
+BOOST_AUTO_TEST_CASE(switch_)
+{
+	string out = assemble("{ let z := 0 switch z case 0 { let x := 2 let y := 3 } default { z := 3 } let t := 9 }");
+	BOOST_CHECK_EQUAL(out,
+		"PUSH1 0x0 DUP1 "
+		"PUSH1 0x0 DUP2 EQ PUSH1 0x11 JUMPI "
+		"PUSH1 0x3 SWAP2 POP PUSH1 0x18 JUMP "
+		"JUMPDEST PUSH1 0x2 POP PUSH1 0x3 POP "
+		"JUMPDEST POP POP " // This is where z and its copy (switch condition) can be removed.
+		"PUSH1 0x9 POP "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(reuse_slots)
+{
+	// x and y should reuse the slots of b and d
+	string out = assemble("{ let a, b, c, d let x := 2 let y := 3 mstore(x, a) mstore(y, c) }");
+	BOOST_CHECK_EQUAL(out,
+		"PUSH1 0x0 PUSH1 0x0 PUSH1 0x0 PUSH1 0x0 "
+		"POP " // d is removed right away
+		"PUSH1 0x2 SWAP2 POP " // x is stored at b's slot
+		"PUSH1 0x3 DUP4 DUP4 MSTORE "
+		"DUP2 DUP2 MSTORE "
+		"POP POP POP POP "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(for_1)
+{
+	// Special scoping rules, but can remove z early
+	string out = assemble("{ for { let z := 0 } 1 { } { let x := 3 } let t := 2 }");
+	BOOST_CHECK_EQUAL(out,
+		"PUSH1 0x0 POP "
+		"JUMPDEST PUSH1 0x1 ISZERO PUSH1 0x11 JUMPI "
+		"PUSH1 0x3 POP JUMPDEST PUSH1 0x3 JUMP "
+		"JUMPDEST PUSH1 0x2 POP "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(for_2)
+{
+	// Special scoping rules, cannot remove z until after the loop!
+	string out = assemble("{ for { let z := 0 } 1 { } { z := 8 let x := 3 } let t := 2 }");
+	BOOST_CHECK_EQUAL(out,
+		"PUSH1 0x0 "
+		"JUMPDEST PUSH1 0x1 ISZERO PUSH1 0x14 JUMPI "
+		"PUSH1 0x8 SWAP1 POP "
+		"PUSH1 0x3 POP "
+		"JUMPDEST PUSH1 0x2 JUMP "
+		"JUMPDEST POP " // z is removed
+		"PUSH1 0x2 POP "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(function_trivial)
+{
+	string in = R"({
+		function f() { }
+	})";
+	BOOST_CHECK_EQUAL(assemble(in),
+		"PUSH1 0x5 JUMP JUMPDEST JUMP JUMPDEST "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(function_retparam)
+{
+	string in = R"({
+		function f() -> x, y { }
+	})";
+	BOOST_CHECK_EQUAL(assemble(in),
+		"PUSH1 0xB JUMP "
+		"JUMPDEST PUSH1 0x0 PUSH1 0x0 SWAP1 SWAP2 JUMP "
+		"JUMPDEST "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(function_params)
+{
+	string in = R"({
+		function f(a, b) { }
+	})";
+	BOOST_CHECK_EQUAL(assemble(in), "PUSH1 0x7 JUMP JUMPDEST POP POP JUMP JUMPDEST ");
+}
+
+BOOST_AUTO_TEST_CASE(function_params_and_retparams)
+{
+	string in = R"({
+		function f(a, b, c, d) -> x, y { }
+	})";
+	// This does not re-use the parameters for the return parameters
+	// We do not expect parameters to be fully unused, so the stack
+	// layout for a function is still fixed, even though parameters
+	// can be re-used.
+	BOOST_CHECK_EQUAL(assemble(in),
+		"PUSH1 0x10 JUMP JUMPDEST PUSH1 0x0 PUSH1 0x0 SWAP5 POP SWAP5 SWAP3 POP POP POP JUMP JUMPDEST "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(function_params_and_retparams_partly_unused)
+{
+	string in = R"({
+		function f(a, b, c, d) -> x, y { b := 3 let s := 9 y := 2 mstore(s, y) }
+	})";
+	BOOST_CHECK_EQUAL(assemble(in),
+		"PUSH1 0x1E JUMP "
+		"JUMPDEST PUSH1 0x0 PUSH1 0x0 "
+		"PUSH1 0x3 SWAP4 POP "
+		"PUSH1 0x9 PUSH1 0x2 SWAP2 POP "
+		"DUP2 DUP2 MSTORE "
+		"POP SWAP5 POP SWAP5 SWAP3 POP POP POP JUMP "
+		"JUMPDEST "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(function_with_body_embedded)
+{
+	string in = R"({
+		let b := 3
+		function f(a, r) -> t {
+			// r could be removed right away, but a cannot - this is not implemented, though
+			let x := a a := 3 t := a
+		}
+		b := 7
+	})";
+	BOOST_CHECK_EQUAL(assemble(in),
+		"PUSH1 0x3 PUSH1 "
+		"0x16 JUMP "
+		"JUMPDEST PUSH1 0x0 " // start of f, initialize t
+		"DUP2 POP " // let x := a
+		"PUSH1 0x3 SWAP2 POP "
+		"DUP2 SWAP1 POP "
+		"SWAP3 SWAP2 POP POP JUMP "
+		"JUMPDEST PUSH1 0x7 SWAP1 "
+		"POP POP "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(function_call)
+{
+	string in = R"({
+		let b := f(1, 2)
+		function f(a, r) -> t { }
+		b := f(3, 4)
+	})";
+	BOOST_CHECK_EQUAL(assemble(in),
+		"PUSH1 0x9 PUSH1 0x2 PUSH1 0x1 PUSH1 0xD JUMP "
+		"JUMPDEST PUSH1 0x15 JUMP " // jump over f
+		"JUMPDEST PUSH1 0x0 SWAP3 SWAP2 POP POP JUMP " // f
+		"JUMPDEST PUSH1 0x1F PUSH1 0x4 PUSH1 0x3 PUSH1 0xD JUMP "
+		"JUMPDEST SWAP1 POP POP "
+	);
+}
+
+
+BOOST_AUTO_TEST_CASE(functions_multi_return)
+{
+	string in = R"({
+		function f(a, b) -> t { }
+		function g() -> r, s { }
+		let x := f(1, 2)
+		x := f(3, 4)
+		let y, z := g()
+		y, z := g()
+		let unused := 7
+	})";
+	BOOST_CHECK_EQUAL(assemble(in),
+		"PUSH1 0xB JUMP "
+		"JUMPDEST PUSH1 0x0 SWAP3 SWAP2 POP POP JUMP " // f
+		"JUMPDEST PUSH1 0x17 JUMP "
+		"JUMPDEST PUSH1 0x0 PUSH1 0x0 SWAP1 SWAP2 JUMP " // g
+		"JUMPDEST PUSH1 0x21 PUSH1 0x2 PUSH1 0x1 PUSH1 0x3 JUMP " // f(1, 2)
+		"JUMPDEST PUSH1 0x2B PUSH1 0x4 PUSH1 0x3 PUSH1 0x3 JUMP " // f(3, 4)
+		"JUMPDEST SWAP1 POP " // assignment to x
+		"POP " // remove x
+		"PUSH1 0x34 PUSH1 0xF JUMP " // g()
+		"JUMPDEST PUSH1 0x3A PUSH1 0xF JUMP " // g()
+		"JUMPDEST SWAP2 POP SWAP2 POP " // assignments
+		"POP POP " // removal of y and z
+		"PUSH1 0x7 POP "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(reuse_slots_function)
+{
+	string in = R"({
+		function f() -> x, y, z, t {}
+		let a, b, c, d := f() let x1 := 2 let y1 := 3 mstore(x1, a) mstore(y1, c)
+	})";
+	BOOST_CHECK_EQUAL(assemble(in),
+		"PUSH1 0x11 JUMP "
+		"JUMPDEST PUSH1 0x0 PUSH1 0x0 PUSH1 0x0 PUSH1 0x0 SWAP1 SWAP2 SWAP3 SWAP4 JUMP "
+		"JUMPDEST PUSH1 0x17 PUSH1 0x3 JUMP "
+		// Stack: a b c d
+		"JUMPDEST POP " // d is unused
+		// Stack: a b c
+		"PUSH1 0x2 SWAP2 POP " // x1 reuses b's slot
+		"PUSH1 0x3 "
+		// Stack: a x1 c y1
+		"DUP4 DUP4 MSTORE "
+		"DUP2 DUP2 MSTORE "
+		"POP POP POP POP "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(reuse_slots_function_with_gaps)
+{
+	string in = R"({
+		// Only x3 is actually used, the slots of
+		// x1 and x2 will be reused right away.
+		let x1 := 5 let x2 := 6 let x3 := 7
+		mstore(x1, x2)
+		function f() -> x, y, z, t {}
+		let a, b, c, d := f() mstore(x3, a) mstore(c, d)
+	})";
+	BOOST_CHECK_EQUAL(assemble(in),
+		"PUSH1 0x5 PUSH1 0x6 PUSH1 0x7 "
+		"DUP2 DUP4 MSTORE "
+		"PUSH1 0x1A JUMP " // jump across function
+		"JUMPDEST PUSH1 0x0 PUSH1 0x0 PUSH1 0x0 PUSH1 0x0 SWAP1 SWAP2 SWAP3 SWAP4 JUMP "
+		"JUMPDEST PUSH1 0x20 PUSH1 0xC JUMP "
+		// stack: x1 x2 x3 a b c d
+		"JUMPDEST SWAP6 POP " // move d into x1
+		// stack: d x2 x3 a b c
+		"SWAP4 POP "
+		// stack: d c x3 a b
+		"POP "
+		// stack: d c x3 a
+		"DUP1 DUP3 MSTORE "
+		"POP POP "
+		// stack: d c
+		"DUP2 DUP2 MSTORE "
+		"POP POP "
+	);
+}
+
+
+BOOST_AUTO_TEST_SUITE_END()
+
+}
+}
+}