Draft for moving function arguments to memory.

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>
@@ -66,13 +67,29 @@ struct MemoryOffsetAllocator
 		if (unreachableVariables.count(_function))
 		{
 			yulAssert(!slotAllocations.count(_function), "");
+			set<YulString> arguments;
+			if (functionDefinitions.count(_function))
+			{
+				FunctionDefinition const* functionDefinition = functionDefinitions.at(_function);
+				yulAssert(functionDefinition, "");
+				for (TypedName const& argument: functionDefinition->parameters + functionDefinition->returnVariables)
+					arguments.emplace(argument.name);
+			}
+
 			for (YulString variable: unreachableVariables.at(_function))
-				if (variable.empty())
+				// The variable is empty in case the function has too many arguments, which is handled separately below.
+				if (!variable.empty())
 				{
-					// TODO: Too many function arguments or return parameters.
-				}
-				else
 					slotAllocations[variable] = requiredSlots++;
+					arguments.erase(variable);
+				}
+
+			// TODO: use a clever choice of arguments to move instead.
+			while (arguments.size() > 16)
+			{
+				slotAllocations[*arguments.begin()] = requiredSlots++;
+				arguments.erase(arguments.begin());
+			}
 		}
 
 		return slotsRequiredForFunction[_function] = requiredSlots;
@@ -80,6 +97,7 @@ struct MemoryOffsetAllocator
 
 	map<YulString, set<YulString>> const& unreachableVariables;
 	map<YulString, set<YulString>> const& callGraph;
+	map<YulString, FunctionDefinition const*> const& functionDefinitions;
 
 	map<YulString, uint64_t> slotAllocations{};
 	map<YulString, uint64_t> slotsRequiredForFunction{};
@@ -128,7 +146,9 @@ 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();
 
 	StackToMemoryMover::run(_context, reservedMemory, memoryOffsetAllocator.slotAllocations, requiredSlots, *_object.code);

libyul/optimiser/StackToMemoryMover.cpp

@@ -20,6 +20,7 @@
 #include <libyul/optimiser/ExpressionJoiner.h>
 #include <libyul/optimiser/ForLoopConditionIntoBody.h>
 #include <libyul/optimiser/ForLoopConditionOutOfBody.h>
+#include <libyul/optimiser/FunctionDefinitionCollector.h>
 #include <libyul/backends/evm/EVMDialect.h>
 
 #include <libyul/AsmData.h>
@@ -53,6 +54,22 @@ vector<Statement> generateMemoryStore(
 	}});
 	return result;
 }
+FunctionCall generateMemoryLoad(Dialect const& _dialect, langutil::SourceLocation const& _loc, YulString _mpos)
+{
+	BuiltinFunction const* memoryLoadFunction = _dialect.memoryLoadFunction(_dialect.defaultType);
+	yulAssert(memoryLoadFunction, "");
+	return FunctionCall{
+		_loc,
+		Identifier{_loc, memoryLoadFunction->name}, {
+			Literal{
+				_loc,
+				LiteralKind::Number,
+				_mpos,
+				{}
+			}
+		}
+	};
+}
 }
 
 void StackToMemoryMover::run(
@@ -79,38 +96,22 @@ void StackToMemoryMover::run(
 	ForLoopConditionIntoBody::run(_context, _block);
 	ExpressionSplitter::run(_context, _block);
 
+	FunctionCallRewriter functionCallRewriter(_context, memoryOffsetTracker, FunctionDefinitionCollector::run(_block));
+	functionCallRewriter(_block);
+
 	VariableDeclarationAndAssignmentMover declarationAndAssignmentMover(_context, memoryOffsetTracker);
 	declarationAndAssignmentMover(_block);
 
 	ExpressionJoiner::runUntilStabilized(_context, _block);
 	ForLoopConditionOutOfBody::run(_context, _block);
 
+	FunctionDefinitionRewriter functionDefinitionRewriter(_context, memoryOffsetTracker);
+	functionDefinitionRewriter(_block);
+
 	IdentifierMover identifierMover(_context, memoryOffsetTracker);
 	identifierMover(_block);
 }
 
-void StackToMemoryMover::VariableDeclarationAndAssignmentMover::operator()(FunctionDefinition& _functionDefinition)
-{
-	for (TypedName const& param: _functionDefinition.parameters + _functionDefinition.returnVariables)
-		if (m_memoryOffsetTracker(param.name))
-		{
-			// TODO: we cannot handle function parameters yet.
-			return;
-		}
-	ASTModifier::operator()(_functionDefinition);
-}
-
-void StackToMemoryMover::IdentifierMover::operator()(FunctionDefinition& _functionDefinition)
-{
-	for (TypedName const& param: _functionDefinition.parameters + _functionDefinition.returnVariables)
-		if (m_memoryOffsetTracker(param.name))
-		{
-			// TODO: we cannot handle function parameters yet.
-			return;
-		}
-	ASTModifier::operator()(_functionDefinition);
-}
-
 void StackToMemoryMover::VariableDeclarationAndAssignmentMover::operator()(Block& _block)
 {
 	using OptionalStatements = std::optional<vector<Statement>>;
@@ -198,20 +199,7 @@ void StackToMemoryMover::IdentifierMover::visit(Expression& _expression)
 	if (Identifier* identifier = std::get_if<Identifier>(&_expression))
 		if (optional<YulString> offset = m_memoryOffsetTracker(identifier->name))
 		{
-			BuiltinFunction const* memoryLoadFunction = m_context.dialect.memoryLoadFunction(m_context.dialect.defaultType);
+			_expression = generateMemoryLoad(m_context.dialect, identifier->location, *offset);
-			yulAssert(memoryLoadFunction, "");
-			langutil::SourceLocation loc = identifier->location;
-			_expression = FunctionCall{
-				loc,
-				Identifier{loc, memoryLoadFunction->name}, {
-					Literal{
-						loc,
-						LiteralKind::Number,
-						*offset,
-						{}
-					}
-				}
-			};
 			return;
 		}
 	ASTModifier::visit(_expression);
@@ -228,3 +216,150 @@ optional<YulString> StackToMemoryMover::VariableMemoryOffsetTracker::operator()(
 	else
 		return std::nullopt;
 }
+
+void StackToMemoryMover::FunctionDefinitionRewriter::operator()(FunctionDefinition& _functionDefinition)
+{
+	vector<Statement> returnVariableMemoryInits;
+	TypedNameList parameters;
+	for (TypedName& parameter: _functionDefinition.parameters)
+		if (!m_memoryOffsetTracker(parameter.name))
+			parameters.emplace_back(std::move(parameter));
+	_functionDefinition.parameters = std::move(parameters);
+	TypedNameList returnVariables;
+	for (TypedName& returnVariable: _functionDefinition.returnVariables)
+		if (auto slot = m_memoryOffsetTracker(returnVariable.name))
+			returnVariableMemoryInits += generateMemoryStore(
+				m_context.dialect,
+				returnVariable.location,
+				*slot,
+				Literal{returnVariable.location, LiteralKind::Number, "0"_yulstring, {}}
+			);
+		else
+			returnVariables.emplace_back(std::move(returnVariable));
+	_functionDefinition.returnVariables = std::move(returnVariables);
+
+	if (!returnVariableMemoryInits.empty())
+	{
+		returnVariableMemoryInits += std::move(_functionDefinition.body.statements);
+		_functionDefinition.body.statements = std::move(returnVariableMemoryInits);
+	}
+}
+
+StackToMemoryMover::FunctionCallRewriter::FunctionCallRewriter(
+	OptimiserStepContext& _context,
+	VariableMemoryOffsetTracker const& _memoryOffsetTracker,
+	std::map<YulString, FunctionDefinition const*> const& _functionDefinitions
+): m_context(_context), m_memoryOffsetTracker(_memoryOffsetTracker)
+{
+	for (auto [functionName, functionDefinition]: _functionDefinitions)
+	{
+		FunctionArguments& functionArguments = m_functionArguments[functionName];
+		for (TypedName const& parameter: functionDefinition->parameters)
+			functionArguments.parameters.push_back(parameter.name);
+		for (TypedName const& returnVariable: functionDefinition->returnVariables)
+			functionArguments.returnVariables.push_back(returnVariable.name);
+	}
+}
+
+void StackToMemoryMover::FunctionCallRewriter::operator()(FunctionCall& _functionCall)
+{
+	if (!m_functionArguments.count(_functionCall.functionName.name))
+	{
+		// If it is not in the list of user-defined functions, it should be a builtin.
+		yulAssert(m_context.dialect.builtin(_functionCall.functionName.name), "");
+		ASTModifier::operator()(_functionCall);
+		return;
+	}
+
+	yulAssert(m_slotsForCurrentReturns.empty(), "");
+
+	FunctionArguments const& functionArguments = m_functionArguments.at(_functionCall.functionName.name);
+
+	yulAssert(_functionCall.arguments.size() == functionArguments.parameters.size(), "");
+
+	vector<Expression> arguments;
+	for (size_t i = 0; i < functionArguments.parameters.size(); ++i)
+	{
+		yulAssert(
+			holds_alternative<Identifier>(_functionCall.arguments[i]) ||
+			holds_alternative<Literal>(_functionCall.arguments[i]),
+			"Expected fully split expressions."
+		);
+
+		if (auto slot = m_memoryOffsetTracker(functionArguments.parameters[i]))
+		{
+			auto loc = locationOf(_functionCall.arguments[i]);
+			m_statementsToPrefix += generateMemoryStore(
+				m_context.dialect,
+				loc,
+				*slot,
+				std::move(_functionCall.arguments[i])
+			);
+		}
+		else
+			arguments.emplace_back(_functionCall.arguments[i]);
+	}
+
+	for (YulString returnVariable: functionArguments.returnVariables)
+		if (auto slot = m_memoryOffsetTracker(returnVariable))
+			m_slotsForCurrentReturns.emplace_back(*slot);
+		else
+			m_slotsForCurrentReturns.emplace_back(std::nullopt);
+
+	_functionCall.arguments = std::move(arguments);
+}
+
+void StackToMemoryMover::FunctionCallRewriter::operator()(Block& _block)
+{
+	util::iterateReplacing(_block.statements, [&](Statement& _statement) -> std::optional<vector<Statement>> {
+		yulAssert(m_statementsToPrefix.empty(), "");
+		vector<Statement> statementsToSuffix;
+		auto handleVariableDeclarationOrAssignment = [&](auto& _stmt, auto& _variableList) {
+			visit(*_stmt.value);
+			if (!m_slotsForCurrentReturns.empty())
+			{
+				std::decay_t<decltype(_variableList)> variableList;
+				yulAssert(_variableList.size() == m_slotsForCurrentReturns.size(), "");
+				for (size_t i = 0; i < m_slotsForCurrentReturns.size(); ++i)
+					if (auto slot = m_slotsForCurrentReturns[i])
+					{
+						auto loc = _variableList[i].location;
+						statementsToSuffix.emplace_back(std::decay_t<decltype(_stmt)>{
+							loc,
+							{std::move(_variableList[i])},
+							std::make_unique<Expression>(generateMemoryLoad(
+								m_context.dialect,
+								loc,
+								*slot
+							))
+						});
+					}
+					else
+						variableList.emplace_back(move(_variableList[i]));
+				if (variableList.empty())
+					_statement = ExpressionStatement{_stmt.location, std::move(*_stmt.value)};
+				else
+					_variableList = move(variableList);
+				m_slotsForCurrentReturns.clear();
+			}
+		};
+		if (Assignment* _assignment = std::get_if<Assignment>(&_statement))
+			handleVariableDeclarationOrAssignment(*_assignment, _assignment->variableNames);
+		else if (VariableDeclaration* _variableDeclaration = std::get_if<VariableDeclaration>(&_statement))
+		{
+			if (_variableDeclaration->value)
+				handleVariableDeclarationOrAssignment(*_variableDeclaration, _variableDeclaration->variables);
+		}
+		else
+			visit(_statement);
+
+		if (m_statementsToPrefix.empty() && statementsToSuffix.empty())
+			return {};
+		vector<Statement> result;
+		result.swap(m_statementsToPrefix);
+		result.emplace_back(std::move(_statement));
+		result += std::move(statementsToSuffix);
+		statementsToSuffix.clear();
+		return result;
+	});
+}

libyul/optimiser/StackToMemoryMover.h

@@ -56,7 +56,7 @@ namespace solidity::yul
  * and then all references to a variable ``a`` in the map are replaced by ``mload(<memory offset for a>)``.
  *
  * If a visited function has arguments or return parameters that are contained in the map,
- * the entire function is skipped (no local variables in the function will be moved at all).
+ * TODO: document exactly what happens.
  *
  * Prerequisite: Disambiguator, ForLoopInitRewriter, FunctionHoister.
  */
@@ -108,7 +108,6 @@ private:
 		{
 		}
 		using ASTModifier::operator();
-		void operator()(FunctionDefinition& _functionDefinition) override;
 		void operator()(Block& _block) override;
 	private:
 		OptimiserStepContext& m_context;
@@ -124,12 +123,50 @@ private:
 		{
 		}
 		using ASTModifier::operator();
-		void operator()(FunctionDefinition& _functionDefinition) override;
 		void visit(Expression& _expression) override;
 	private:
 		OptimiserStepContext& m_context;
 		VariableMemoryOffsetTracker const& m_memoryOffsetTracker;
 	};
+	class FunctionDefinitionRewriter: ASTModifier
+	{
+	public:
+		FunctionDefinitionRewriter(
+			OptimiserStepContext& _context,
+			VariableMemoryOffsetTracker const& _memoryOffsetTracker
+		): m_context(_context), m_memoryOffsetTracker(_memoryOffsetTracker)
+		{
+		}
+		using ASTModifier::operator();
+		void operator()(FunctionDefinition& _functionDefinition) override;
+	private:
+		OptimiserStepContext& m_context;
+		VariableMemoryOffsetTracker const& m_memoryOffsetTracker;
+	};
+	class FunctionCallRewriter: ASTModifier
+	{
+	public:
+		FunctionCallRewriter(
+			OptimiserStepContext& _context,
+			VariableMemoryOffsetTracker const& _memoryOffsetTracker,
+			std::map<YulString, FunctionDefinition const*> const& _functionDefinitions
+		);
+		using ASTModifier::operator();
+		using ASTModifier::visit;
+		void operator()(FunctionCall& _functionCall) override;
+		void operator()(Block& _block) override;
+	private:
+		OptimiserStepContext& m_context;
+		VariableMemoryOffsetTracker const& m_memoryOffsetTracker;
+		struct FunctionArguments
+		{
+			std::vector<YulString> parameters;
+			std::vector<YulString> returnVariables;
+		};
+		std::map<YulString, FunctionArguments> m_functionArguments;
+		std::vector<Statement> m_statementsToPrefix;
+		std::vector<std::optional<YulString>> m_slotsForCurrentReturns;
+	};
 };
 
 }