From def0ebbb3e4a8eeb68d2e29155778eb7ede59188 Mon Sep 17 00:00:00 2001
From: Daniel Kirchner <daniel@ekpyron.org>
Date: Fri, 10 Jul 2020 17:39:23 +0200
Subject: [PATCH] Free variables directly after visiting RHS of Variable
 Declarations during EVMCodeTransform.

---
 Changelog.md                             |  1 +
 libyul/backends/evm/EVMCodeTransform.cpp | 14 +++---
 libyul/backends/evm/EVMCodeTransform.h   |  5 +-
 test/libyul/StackReuseCodegen.cpp        | 63 ++++++++++++++++++++++++
 4 files changed, 75 insertions(+), 8 deletions(-)

diff --git a/Changelog.md b/Changelog.md
index e140ebea5..047b06ddf 100644
--- a/Changelog.md
+++ b/Changelog.md
@@ -6,6 +6,7 @@ Language Features:
 Compiler Features:
  * Code Generator: Evaluate ``keccak256`` of string literals at compile-time.
  * Peephole Optimizer: Remove unnecessary masking of tags.
+ * Yul EVM Code Transform: Free stack slots directly after visiting the right-hand-side of variable declarations instead of at the end of the statement only.
 
 Bugfixes:
  * Type Checker: Fix overload resolution in combination with ``{value: ...}``.
diff --git a/libyul/backends/evm/EVMCodeTransform.cpp b/libyul/backends/evm/EVMCodeTransform.cpp
index 6b0f06e01..041e00380 100644
--- a/libyul/backends/evm/EVMCodeTransform.cpp
+++ b/libyul/backends/evm/EVMCodeTransform.cpp
@@ -141,7 +141,7 @@ bool CodeTransform::unreferenced(Scope::Variable const& _var) const
 	return !m_context->variableReferences.count(&_var) || m_context->variableReferences[&_var] == 0;
 }
 
-void CodeTransform::freeUnusedVariables()
+void CodeTransform::freeUnusedVariables(bool _popUnusedSlotsAtStackTop)
 {
 	if (!m_allowStackOpt)
 		return;
@@ -154,11 +154,12 @@ void CodeTransform::freeUnusedVariables()
 				deleteVariable(var);
 		}
 
-	while (m_unusedStackSlots.count(m_assembly.stackHeight() - 1))
-	{
-		yulAssert(m_unusedStackSlots.erase(m_assembly.stackHeight() - 1), "");
-		m_assembly.appendInstruction(evmasm::Instruction::POP);
-	}
+	if (_popUnusedSlotsAtStackTop)
+		while (m_unusedStackSlots.count(m_assembly.stackHeight() - 1))
+		{
+			yulAssert(m_unusedStackSlots.erase(m_assembly.stackHeight() - 1), "");
+			m_assembly.appendInstruction(evmasm::Instruction::POP);
+		}
 }
 
 void CodeTransform::deleteVariable(Scope::Variable const& _var)
@@ -181,6 +182,7 @@ void CodeTransform::operator()(VariableDeclaration const& _varDecl)
 	{
 		std::visit(*this, *_varDecl.value);
 		expectDeposit(static_cast<int>(numVariables), static_cast<int>(heightAtStart));
+		freeUnusedVariables(false);
 	}
 	else
 	{
diff --git a/libyul/backends/evm/EVMCodeTransform.h b/libyul/backends/evm/EVMCodeTransform.h
index 3c9506b47..1cfa6afd7 100644
--- a/libyul/backends/evm/EVMCodeTransform.h
+++ b/libyul/backends/evm/EVMCodeTransform.h
@@ -161,8 +161,9 @@ protected:
 	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();
+	/// Also POPs unused topmost stack slots,
+	/// unless @a _popUnusedSlotsAtStackTop is set to false.
+	void freeUnusedVariables(bool _popUnusedSlotsAtStackTop = true);
 	/// Marks the stack slot of @a _var to be reused.
 	void deleteVariable(Scope::Variable const& _var);
 
diff --git a/test/libyul/StackReuseCodegen.cpp b/test/libyul/StackReuseCodegen.cpp
index eec882e8c..1b6b91657 100644
--- a/test/libyul/StackReuseCodegen.cpp
+++ b/test/libyul/StackReuseCodegen.cpp
@@ -345,6 +345,69 @@ BOOST_AUTO_TEST_CASE(reuse_slots_function_with_gaps)
 	);
 }
 
+BOOST_AUTO_TEST_CASE(reuse_on_decl_assign_to_last_used)
+{
+	string in = R"({
+		let x := 5
+		let y := x // y should reuse the stack slot of x
+		sstore(y, y)
+	})";
+	BOOST_CHECK_EQUAL(assemble(in),
+		"PUSH1 0x5 "
+		"DUP1 SWAP1 POP "
+		"DUP1 DUP2 SSTORE "
+		"POP "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(reuse_on_decl_assign_to_last_used_expr)
+{
+	string in = R"({
+		let x := 5
+		let y := add(x, 2) // y should reuse the stack slot of x
+		sstore(y, y)
+	})";
+	BOOST_CHECK_EQUAL(assemble(in),
+		"PUSH1 0x5 "
+		"PUSH1 0x2 DUP2 ADD "
+		"SWAP1 POP "
+		"DUP1 DUP2 SSTORE "
+		"POP "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(reuse_on_decl_assign_to_not_last_used)
+{
+	string in = R"({
+		let x := 5
+		let y := x // y should not reuse the stack slot of x, since x is still used below
+		sstore(y, x)
+	})";
+	BOOST_CHECK_EQUAL(assemble(in),
+		"PUSH1 0x5 "
+		"DUP1 "
+		"DUP2 DUP2 SSTORE "
+		"POP POP "
+	);
+}
+
+BOOST_AUTO_TEST_CASE(reuse_on_decl_assign_not_same_scope)
+{
+	string in = R"({
+		let x := 5
+		{
+			let y := x // y should not reuse the stack slot of x, since x is not in the same scope
+			sstore(y, y)
+		}
+	})";
+	BOOST_CHECK_EQUAL(assemble(in),
+		"PUSH1 0x5 "
+		"DUP1 "
+		"DUP1 DUP2 SSTORE "
+		"POP POP "
+	);
+}
+
 
 BOOST_AUTO_TEST_SUITE_END()