Merge pull request #10015 from ethereum/moveFunctionArgumentsToMemoryNew

Stack Limit Evader: Move function arguments and return values to memory (v2)

Author: hrkrshnn

Changelog.md

@@ -4,6 +4,9 @@ Language Features:
 
 
 Compiler Features:
+ * Yul EVM Code Transform: Do not reuse stack slots that immediately become unreachable.
+ * Yul EVM Code Transform: Also pop unused argument slots for functions without return variables (under the same restrictions as for functions with return variables).
+ * Yul Optimizer: Move function arguments and return variables to memory with the experimental Stack Limit Evader (which is not enabled by default).
 
 
 Bugfixes:

libyul/CMakeLists.txt

@@ -129,6 +129,8 @@ add_library(yul
 	optimiser/FullInliner.h
 	optimiser/FunctionCallFinder.cpp
 	optimiser/FunctionCallFinder.h
+	optimiser/FunctionDefinitionCollector.cpp
+	optimiser/FunctionDefinitionCollector.h
 	optimiser/FunctionGrouper.cpp
 	optimiser/FunctionGrouper.h
 	optimiser/FunctionHoister.cpp

libyul/backends/evm/EVMCodeTransform.cpp

@@ -182,16 +182,24 @@ void CodeTransform::operator()(VariableDeclaration const& _varDecl)
 			else
 				m_variablesScheduledForDeletion.insert(&var);
 		}
-		else if (m_unusedStackSlots.empty())
-			atTopOfStack = false;
 		else
 		{
-			auto slot = static_cast<size_t>(*m_unusedStackSlots.begin());
-			m_unusedStackSlots.erase(m_unusedStackSlots.begin());
-			m_context->variableStackHeights[&var] = slot;
-			if (size_t heightDiff = variableHeightDiff(var, varName, true))
-				m_assembly.appendInstruction(evmasm::swapInstruction(static_cast<unsigned>(heightDiff - 1)));
-			m_assembly.appendInstruction(evmasm::Instruction::POP);
+			bool foundUnusedSlot = false;
+			for (auto it = m_unusedStackSlots.begin(); it != m_unusedStackSlots.end(); ++it)
+			{
+				if (m_assembly.stackHeight() - *it > 17)
+					continue;
+				foundUnusedSlot = true;
+				auto slot = static_cast<size_t>(*it);
+				m_unusedStackSlots.erase(it);
+				m_context->variableStackHeights[&var] = slot;
+				if (size_t heightDiff = variableHeightDiff(var, varName, true))
+					m_assembly.appendInstruction(evmasm::swapInstruction(static_cast<unsigned>(heightDiff - 1)));
+				m_assembly.appendInstruction(evmasm::Instruction::POP);
+				break;
+			}
+			if (!foundUnusedSlot)
+				atTopOfStack = false;
 		}
 	}
 }
@@ -404,7 +412,7 @@ void CodeTransform::operator()(FunctionDefinition const& _function)
 				subTransform.deleteVariable(var);
 		}
 
-	if (!m_allowStackOpt || _function.returnVariables.empty())
+	if (!m_allowStackOpt)
 		subTransform.setupReturnVariablesAndFunctionExit();
 
 	subTransform(_function.body);
@@ -594,6 +602,7 @@ void CodeTransform::visitExpression(Expression const& _expression)
 
 void CodeTransform::setupReturnVariablesAndFunctionExit()
 {
+	yulAssert(isInsideFunction(), "");
 	yulAssert(!returnVariablesAndFunctionExitAreSetup(), "");
 	yulAssert(m_scope, "");
 
@@ -656,7 +665,8 @@ void CodeTransform::visitStatements(vector<Statement> const& _statements)
 	{
 		freeUnusedVariables();
 		if (
-			!m_delayedReturnVariables.empty() &&
+			isInsideFunction() &&
+			!returnVariablesAndFunctionExitAreSetup() &&
 			statementNeedsReturnVariableSetup(statement, m_delayedReturnVariables)
 		)
 			setupReturnVariablesAndFunctionExit();

libyul/backends/evm/EVMCodeTransform.h

@@ -181,6 +181,10 @@ class CodeTransform
 	{
 		return m_functionExitStackHeight.has_value();
 	}
+	bool isInsideFunction() const
+	{
+		return m_functionExitLabel.has_value();
+	}
 
 	AbstractAssembly& m_assembly;
 	AsmAnalysisInfo& m_info;

libyul/optimiser/FunctionDefinitionCollector.cpp

@@ -0,0 +1,36 @@
+/*
+	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/>.
+*/
+
+#include <libyul/optimiser/FunctionDefinitionCollector.h>
+#include <libyul/AST.h>
+
+using namespace std;
+using namespace solidity;
+using namespace solidity::yul;
+
+map<YulString, FunctionDefinition const*> FunctionDefinitionCollector::run(Block& _block)
+{
+	FunctionDefinitionCollector functionDefinitionCollector;
+	functionDefinitionCollector(_block);
+	return functionDefinitionCollector.m_functionDefinitions;
+}
+
+void FunctionDefinitionCollector::operator()(FunctionDefinition const& _functionDefinition)
+{
+	m_functionDefinitions[_functionDefinition.name] = &_functionDefinition;
+	ASTWalker::operator()(_functionDefinition);
+}

libyul/optimiser/FunctionDefinitionCollector.h

@@ -0,0 +1,44 @@
+/*
+	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/>.
+*/
+/**
+ * AST walker that finds all function definitions and stores them into a map indexed by the function names.
+ */
+#pragma once
+
+#include <libyul/optimiser/ASTWalker.h>
+
+#include <map>
+
+namespace solidity::yul
+{
+
+/**
+ * AST walker that finds all function definitions and stores them into a map indexed by the function names.
+ *
+ * Prerequisite: Disambiguator
+ */
+class FunctionDefinitionCollector: ASTWalker
+{
+public:
+	static std::map<YulString, FunctionDefinition const*> run(Block& _block);
+private:
+	using ASTWalker::operator();
+	void operator()(FunctionDefinition const& _functionDefinition) override;
+	std::map<YulString, FunctionDefinition const*> m_functionDefinitions;
+};
+
+}

libyul/optimiser/StackLimitEvader.cpp

@@ -18,6 +18,7 @@
 #include <libyul/optimiser/StackLimitEvader.h>
 #include <libyul/optimiser/CallGraphGenerator.h>
 #include <libyul/optimiser/FunctionCallFinder.h>
+#include <libyul/optimiser/FunctionDefinitionCollector.h>
 #include <libyul/optimiser/NameDispenser.h>
 #include <libyul/optimiser/StackToMemoryMover.h>
 #include <libyul/backends/evm/EVMDialect.h>
@@ -28,6 +29,9 @@
 #include <libsolutil/Algorithms.h>
 #include <libsolutil/CommonData.h>
 
+#include <range/v3/view/concat.hpp>
+#include <range/v3/view/take.hpp>
+
 using namespace std;
 using namespace solidity;
 using namespace solidity::yul;
@@ -62,25 +66,42 @@ struct MemoryOffsetAllocator
 			for (YulString child: callGraph.at(_function))
 				requiredSlots = std::max(run(child), requiredSlots);
 
-		if (unreachableVariables.count(_function))
+		if (auto const* unreachables = util::valueOrNullptr(unreachableVariables, _function))
 		{
-			yulAssert(!slotAllocations.count(_function), "");
-			for (YulString variable: unreachableVariables.at(_function))
-				if (variable.empty())
-				{
-					// TODO: Too many function arguments or return parameters.
-				}
-				else
+			if (FunctionDefinition const* functionDefinition = util::valueOrDefault(functionDefinitions, _function, nullptr, util::allow_copy))
+				if (
+					size_t totalArgCount = functionDefinition->returnVariables.size() + functionDefinition->parameters.size();
+					totalArgCount > 16
+				)
+					for (TypedName const& var: ranges::concat_view(
+						functionDefinition->parameters,
+						functionDefinition->returnVariables
+					) | ranges::views::take(totalArgCount - 16))
+						slotAllocations[var.name] = requiredSlots++;
+
+			// Assign slots for all variables that become unreachable in the function body, if the above did not
+			// assign a slot for them already.
+			for (YulString variable: *unreachables)
+				// The empty case is a function with too many arguments or return values,
+				// which was already handled above.
+				if (!variable.empty() && !slotAllocations.count(variable))
 					slotAllocations[variable] = requiredSlots++;
 		}
 
 		return slotsRequiredForFunction[_function] = requiredSlots;
 	}
 
+	/// Maps function names to the set of unreachable variables in that function.
+	/// An empty variable name means that the function has too many arguments or return variables.
 	map<YulString, set<YulString>> const& unreachableVariables;
+	/// The graph of immediate function calls of all functions.
 	map<YulString, set<YulString>> const& callGraph;
+	/// Maps the name of each user-defined function to its definition.
+	map<YulString, FunctionDefinition const*> const& functionDefinitions;
 
+	/// Maps variable names to the memory slot the respective variable is assigned.
 	map<YulString, uint64_t> slotAllocations{};
+	/// Maps function names to the number of memory slots the respective function requires.
 	map<YulString, uint64_t> slotsRequiredForFunction{};
 };
 
@@ -116,6 +137,8 @@ void StackLimitEvader::run(
 
 	// Make sure all calls to ``memoryguard`` we found have the same value as argument (otherwise, abort).
 	u256 reservedMemory = literalArgumentValue(*memoryGuardCalls.front());
+	yulAssert(reservedMemory < u256(1) << 32 - 1, "");
+
 	for (FunctionCall const* memoryGuardCall: memoryGuardCalls)
 		if (reservedMemory != literalArgumentValue(*memoryGuardCall))
 			return;
@@ -127,12 +150,14 @@ void StackLimitEvader::run(
 		if (_unreachableVariables.count(function))
 			return;
 
-	MemoryOffsetAllocator memoryOffsetAllocator{_unreachableVariables, callGraph.functionCalls};
+	map<YulString, FunctionDefinition const*> functionDefinitions = FunctionDefinitionCollector::run(*_object.code);
+
+	MemoryOffsetAllocator memoryOffsetAllocator{_unreachableVariables, callGraph.functionCalls, functionDefinitions};
 	uint64_t requiredSlots = memoryOffsetAllocator.run();
+	yulAssert(requiredSlots < (uint64_t(1) << 32) - 1, "");
 
 	StackToMemoryMover::run(_context, reservedMemory, memoryOffsetAllocator.slotAllocations, requiredSlots, *_object.code);
 
-	yulAssert(requiredSlots < std::numeric_limits<uint64_t>::max() / 32, "");
 	reservedMemory += 32 * requiredSlots;
 	for (FunctionCall* memoryGuardCall: FunctionCallFinder::run(*_object.code, "memoryguard"_yulstring))
 	{