Move function arguments and return values.

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,12 +67,37 @@ struct MemoryOffsetAllocator
 		if (unreachableVariables.count(_function))
 		{
 			yulAssert(!slotAllocations.count(_function), "");
+
+			if (functionDefinitions.count(_function))
+			{
+				FunctionDefinition const* functionDefinition = functionDefinitions.at(_function);
+				yulAssert(functionDefinition, "");
+				size_t totalArgCount = functionDefinition->returnVariables.size() + functionDefinition->parameters.size();
+				size_t slotsNeeded = (totalArgCount > 16) ? totalArgCount - 16 : 0;
+
+				if (slotsNeeded)
+					for (TypedName const& param: functionDefinition->parameters)
+					{
+						slotAllocations[param.name] = requiredSlots++;
+						if (!--slotsNeeded)
+							break;
+					}
+				if (slotsNeeded)
+					for (TypedName const& returnVar: functionDefinition->returnVariables)
+					{
+						slotAllocations[returnVar.name] = requiredSlots++;
+						if (!--slotsNeeded)
+							break;
+					}
+				yulAssert(!slotsNeeded, "");
+			}
+
+			// 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: unreachableVariables.at(_function))
-				if (variable.empty())
-				{
-					// TODO: Too many function arguments or return parameters.
-				}
-				else
+				// 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++;
 		}
 
@@ -80,6 +106,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 +155,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

@@ -15,13 +15,19 @@
 	along with solidity.  If not, see <http://www.gnu.org/licenses/>.
 */
 #include <libyul/optimiser/StackToMemoryMover.h>
+#include <libyul/optimiser/FunctionDefinitionCollector.h>
 #include <libyul/optimiser/NameDispenser.h>
 #include <libyul/backends/evm/EVMDialect.h>
 
 #include <libyul/AST.h>
 
 #include <libsolutil/CommonData.h>
-#include <libsolutil/Visitor.h>
+
+#include <range/v3/algorithm/none_of.hpp>
+#include <range/v3/view/enumerate.hpp>
+#include <range/v3/view/filter.hpp>
+#include <range/v3/view/transform.hpp>
+#include <range/v3/range/conversion.hpp>
 
 using namespace std;
 using namespace solidity;
@@ -77,17 +83,30 @@ void StackToMemoryMover::run(
 )
 {
 	VariableMemoryOffsetTracker memoryOffsetTracker(_reservedMemory, _memorySlots, _numRequiredSlots);
-	StackToMemoryMover stackToMemoryMover(_context, memoryOffsetTracker);
+	StackToMemoryMover stackToMemoryMover(
+		_context,
+		memoryOffsetTracker,
+		util::applyMap(
+			FunctionDefinitionCollector::run(_block),
+			util::mapTuple([](YulString _name, FunctionDefinition const* _funDef) {
+				return std::make_pair(_name, _funDef->returnVariables);
+			}),
+			map<YulString, TypedNameList>{}
+		)
+	);
 	stackToMemoryMover(_block);
+	_block.statements += move(stackToMemoryMover.m_newFunctionDefinitions);
 }
 
 StackToMemoryMover::StackToMemoryMover(
 	OptimiserStepContext& _context,
-	VariableMemoryOffsetTracker const& _memoryOffsetTracker
+	VariableMemoryOffsetTracker const& _memoryOffsetTracker,
+	map<YulString, TypedNameList> _functionReturnVariables
 ):
 m_context(_context),
 m_memoryOffsetTracker(_memoryOffsetTracker),
-m_nameDispenser(_context.dispenser)
+m_nameDispenser(_context.dispenser),
+m_functionReturnVariables(move(_functionReturnVariables))
 {
 	auto const* evmDialect = dynamic_cast<EVMDialect const*>(&_context.dialect);
 	yulAssert(
@@ -98,68 +117,160 @@ m_nameDispenser(_context.dispenser)
 
 void StackToMemoryMover::operator()(FunctionDefinition& _functionDefinition)
 {
-	for (TypedName const& param: _functionDefinition.parameters + _functionDefinition.returnVariables)
-		if (m_memoryOffsetTracker(param.name))
-		{
-			// TODO: we cannot handle function parameters yet.
-			return;
-		}
+	// It is important to first visit the function body, so that it doesn't replace the memory inits for
+	// variable arguments we might generate below.
 	ASTModifier::operator()(_functionDefinition);
+
+	vector<Statement> memoryVariableInits;
+
+	// All function arguments with a memory slot are moved at the beginning of the function body.
+	for (TypedName const& param: _functionDefinition.parameters)
+		if (auto slot = m_memoryOffsetTracker(param.name))
+			memoryVariableInits += generateMemoryStore(
+				m_context.dialect,
+				param.location,
+				*slot,
+				Identifier{param.location, param.name}
+			);
+
+	// All memory return variables have to be initialized to zero in memory.
+	for (TypedName const& returnVariable: _functionDefinition.returnVariables)
+		if (auto slot = m_memoryOffsetTracker(returnVariable.name))
+			memoryVariableInits += generateMemoryStore(
+				m_context.dialect,
+				returnVariable.location,
+				*slot,
+				Literal{returnVariable.location, LiteralKind::Number, "0"_yulstring, {}}
+			);
+
+	// Special case of a function with a single return argument that needs to move to memory.
+	if (_functionDefinition.returnVariables.size() == 1 && m_memoryOffsetTracker(_functionDefinition.returnVariables.front().name))
+	{
+		TypedNameList stackArguments = _functionDefinition.parameters | ranges::views::filter([&](TypedName const& _arg){
+			return !m_memoryOffsetTracker(_arg.name);
+		}) | ranges::to<TypedNameList>;
+		// Generate new function without return argument and with only the non-moved arguments.
+		YulString newFunctionName = m_context.dispenser.newName(_functionDefinition.name);
+		m_newFunctionDefinitions.emplace_back(FunctionDefinition{
+			_functionDefinition.location,
+			newFunctionName,
+			stackArguments,
+			{},
+			move(_functionDefinition.body)
+		});
+		// Replace original function by a call to the new function and an assignment to the return variable from memory.
+		_functionDefinition.body = Block{_functionDefinition.location, move(memoryVariableInits)};
+		_functionDefinition.body.statements.emplace_back(ExpressionStatement{
+			_functionDefinition.location,
+			FunctionCall{
+				_functionDefinition.location,
+				Identifier{_functionDefinition.location, newFunctionName},
+				stackArguments | ranges::views::transform([&](TypedName const& _arg) {
+					return Expression{Identifier{
+						_functionDefinition.location,
+						_arg.name
+					}};
+				}) | ranges::to<vector<Expression>>
+			}
+		});
+		_functionDefinition.body.statements.emplace_back(Assignment{
+			_functionDefinition.location,
+			{Identifier{_functionDefinition.location, _functionDefinition.returnVariables.front().name}},
+			make_unique<Expression>(generateMemoryLoad(
+				m_context.dialect,
+				_functionDefinition.location,
+				*m_memoryOffsetTracker(_functionDefinition.returnVariables.front().name)
+			))
+		});
+		return;
+	}
+
+	if (!memoryVariableInits.empty())
+	{
+		memoryVariableInits += move(_functionDefinition.body.statements);
+		_functionDefinition.body.statements = move(memoryVariableInits);
+	}
+
+	_functionDefinition.returnVariables = _functionDefinition.returnVariables | ranges::views::filter([&](TypedName const& _name){
+		return !m_memoryOffsetTracker(_name.name);
+	}) | ranges::to<TypedNameList>;
 }
 
 void StackToMemoryMover::operator()(Block& _block)
 {
 	using OptionalStatements = std::optional<vector<Statement>>;
-	auto rewriteAssignmentOrVariableDeclaration = [&](
+
+	auto rewriteAssignmentOrVariableDeclarationLeftHandSide = [this](
 		auto& _stmt,
-		auto const& _variables
+		auto& _lhsVars
 	) -> OptionalStatements {
 		using StatementType = decay_t<decltype(_stmt)>;
-		if (_stmt.value)
-			visit(*_stmt.value);
-		bool leftHandSideNeedsMoving = util::contains_if(_variables, [&](auto const& var) {
-			return m_memoryOffsetTracker(var.name);
-		});
-		if (!leftHandSideNeedsMoving)
+
+		langutil::SourceLocation  loc = _stmt.location;
+		if (_lhsVars.size() == 1)
+		{
+			optional<YulString> offset = m_memoryOffsetTracker(_lhsVars.front().name);
+			if (offset)
+				return generateMemoryStore(
+					m_context.dialect,
+					loc,
+					*offset,
+					_stmt.value ? *std::move(_stmt.value) : Literal{loc, LiteralKind::Number, "0"_yulstring, {}}
+				);
+			else
+				return {};
+		}
+		FunctionCall const* functionCall = get_if<FunctionCall>(_stmt.value.get());
+		yulAssert(functionCall, "");
+		auto rhsSlots = m_functionReturnVariables.at(functionCall->functionName.name) |
+			ranges::views::transform([&](auto const& _var) {
+				return m_memoryOffsetTracker(_var.name);
+			}) | ranges::to<vector<optional<YulString>>>;
+
+		// Nothing to do, if the right-hand-side remains entirely on the stack and
+		// none of the variables in the left-hand-side are moved.
+		if (
+			ranges::none_of(rhsSlots, [](optional<YulString> const& _slot) { return _slot.has_value(); }) &&
+			!util::contains_if(_lhsVars, [&](auto const& var) { return m_memoryOffsetTracker(var.name); })
+		)
 			return {};
 
-		langutil::SourceLocation loc = _stmt.location;
-
-		if (_variables.size() == 1)
-		{
-			optional<YulString> offset = m_memoryOffsetTracker(_variables.front().name);
-			yulAssert(offset, "");
-			return generateMemoryStore(
-				m_context.dialect,
-				loc,
-				*offset,
-				_stmt.value ? *std::move(_stmt.value) : Literal{loc, LiteralKind::Number, "0"_yulstring, {}}
-			);
-		}
-
-		VariableDeclaration tempDecl{loc, {}, std::move(_stmt.value)};
 		vector<Statement> memoryAssignments;
 		vector<Statement> variableAssignments;
-		for (auto& var: _variables)
-		{
-			YulString tempVarName = m_nameDispenser.newName(var.name);
-			tempDecl.variables.emplace_back(TypedName{var.location, tempVarName, {}});
+		VariableDeclaration tempDecl{loc, {}, std::move(_stmt.value)};
 
-			if (optional<YulString> offset = m_memoryOffsetTracker(var.name))
+		yulAssert(rhsSlots.size() == _lhsVars.size(), "");
+		for (auto&& [lhsVar, rhsSlot]: ranges::views::zip(_lhsVars, rhsSlots))
+		{
+			unique_ptr<Expression> rhs;
+			if (rhsSlot)
+				rhs = make_unique<Expression>(generateMemoryLoad(m_context.dialect, loc, *rhsSlot));
+			else
+			{
+				YulString tempVarName = m_nameDispenser.newName(lhsVar.name);
+				tempDecl.variables.emplace_back(TypedName{lhsVar.location, tempVarName, {}});
+				rhs = make_unique<Expression>(Identifier{loc, tempVarName});
+			}
+
+			if (optional<YulString> offset = m_memoryOffsetTracker(lhsVar.name))
 				memoryAssignments += generateMemoryStore(
 					m_context.dialect,
 					loc,
 					*offset,
-					Identifier{loc, tempVarName}
+					move(*rhs)
 				);
 			else
 				variableAssignments.emplace_back(StatementType{
-					loc, {move(var)},
-					make_unique<Expression>(Identifier{loc, tempVarName})
+					loc, { std::move(lhsVar) },
+					move(rhs)
 				});
 		}
+
 		std::vector<Statement> result;
-		result.emplace_back(std::move(tempDecl));
+		if (tempDecl.variables.empty())
+			result.emplace_back(ExpressionStatement{loc, *move(tempDecl.value)});
+		else
+			result.emplace_back(std::move(tempDecl));
 		std::reverse(memoryAssignments.begin(), memoryAssignments.end());
 		result += std::move(memoryAssignments);
 		std::reverse(variableAssignments.begin(), variableAssignments.end());
@@ -169,19 +280,14 @@ void StackToMemoryMover::operator()(Block& _block)
 
 	util::iterateReplacing(
 		_block.statements,
-		[&](Statement& _statement)
+		[&](Statement& _statement) -> OptionalStatements
 		{
-			return std::visit(util::GenericVisitor{
-				[&](Assignment& _assignment) -> OptionalStatements
-				{
-					return rewriteAssignmentOrVariableDeclaration(_assignment, _assignment.variableNames);
-				},
-				[&](VariableDeclaration& _varDecl) -> OptionalStatements
-				{
-					return rewriteAssignmentOrVariableDeclaration(_varDecl, _varDecl.variables);
-				},
-				[&](auto& _stmt) -> OptionalStatements { (*this)(_stmt); return {}; }
-			}, _statement);
+			visit(_statement);
+			if (auto* assignment = std::get_if<Assignment>(&_statement))
+				return rewriteAssignmentOrVariableDeclarationLeftHandSide(*assignment, assignment->variableNames);
+			else if (auto* varDecl = std::get_if<VariableDeclaration>(&_statement))
+				return rewriteAssignmentOrVariableDeclarationLeftHandSide(*varDecl, varDecl->variables);
+			return {};
 		}
 	);
 }

libyul/optimiser/StackToMemoryMover.h

@@ -22,6 +22,9 @@
 
 #include <libyul/optimiser/ASTWalker.h>
 #include <libyul/optimiser/OptimiserStep.h>
+#include <libyul/ASTForward.h>
+
+#include <liblangutil/SourceLocation.h>
 #include <libsolutil/Common.h>
 
 namespace solidity::yul
@@ -52,6 +55,16 @@ namespace solidity::yul
  *     let c := _3
  *     let a := _1
  *
+ * In case f has return parameters that are moved to memory, fewer variables are returned and the return values read
+ * from memory instead. Assume the third return parameter of f (i.e. c) has to be moved to memory:
+ *     let a, b, c, d := f()
+ *   then it is replaced by
+ *     let _1, _2, _4 := f()
+ *     mstore(<memory offset for d>, _4)
+ *     mstore(<memory offset for b>, _2)
+ *     let c := mload(<memory offset of third return parameter of f>)
+ *     let a := _1
+ *
  * Assignments to single variables are replaced by mstore's:
  *   If a is in the map, replace
  *     a := expr
@@ -70,8 +83,44 @@ namespace solidity::yul
  *
  * Replace all references to a variable ``a`` in the map 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).
+ * Function arguments are moved at the beginning of a function body:
+ *   If a1 is in the map, replace
+ *     function f(a1, a2, ..., a17)
+ *     {
+ *       ...
+ *       sstore(a1, a17)
+ *     }
+ *   by
+ *     function f(a1, a2, ..., a17)
+ *     {
+ *       mstore(<memory offset for a1>, a1)
+ *       ...
+ *       sstore(mload(<memory offset for a1>, a17)
+ *     }
+ * This relies on the code transform popping arguments that are no longer used, if they are on the stack top.
+ *
+ * Functions with only one return argument that has to be moved are encapsulated in a wrapper function as follows:
+ *   Suppose b and r need to be moved in:
+ *     function f(a, b) -> r
+ *     {
+ *       ...body of f...
+ *       r := b
+ *       ...body of f continued...
+ *     }
+ *   then replace by:
+ *     function f(a, b) -> r
+ *     {
+ *       mstore(<memory offset of b>, b)
+ *       mstore(<memory offset of r>, 0)
+ *       f_1(a)
+ *       r := mload(<memory offset of r>)
+ *     }
+ *     function f_1(a)
+ *     {
+ *       ...body of f...
+ *       mstore(<memory offset of r>, mload(<memory offset of b>))
+ *       ...body of f continued...
+ *     }
  *
  * Prerequisite: Disambiguator, ForLoopInitRewriter, FunctionHoister.
  */
@@ -98,6 +147,7 @@ public:
 
 	void operator()(FunctionDefinition& _functionDefinition) override;
 	void operator()(Block& _block) override;
+	using ASTModifier::visit;
 	void visit(Expression& _expression) override;
 private:
 	class VariableMemoryOffsetTracker
@@ -118,15 +168,34 @@ private:
 		std::map<YulString, uint64_t> const& m_memorySlots;
 		uint64_t m_numRequiredSlots = 0;
 	};
+	struct FunctionMoveInfo
+	{
+		std::vector<std::optional<YulString>> returnVariableSlots;
+	};
 
 	StackToMemoryMover(
 		OptimiserStepContext& _context,
-		VariableMemoryOffsetTracker const& _memoryOffsetTracker
+		VariableMemoryOffsetTracker const& _memoryOffsetTracker,
+		std::map<YulString, std::vector<TypedName>> _functionReturnVariables
+	);
+
+	template<typename StatementType, typename VariableType>
+	std::optional<std::vector<Statement>> rewriteAssignmentOrVariableDeclarationToFunctionCalls(
+		StatementType& _stmt,
+		std::vector<VariableType>& _variables
+	);
+
+	template<typename StatementType, typename VariableType>
+	std::optional<std::vector<Statement>> rewriteAssignmentOrVariableDeclarationLeftHandSide(
+		StatementType& _stmt,
+		std::vector<VariableType>& _variables
 	);
 
 	OptimiserStepContext& m_context;
 	VariableMemoryOffsetTracker const& m_memoryOffsetTracker;
 	NameDispenser& m_nameDispenser;
+	std::map<YulString, std::vector<TypedName>> m_functionReturnVariables;
+	std::list<Statement> m_newFunctionDefinitions;
 };
 
 }