Move function arguments and return values.

2023-10-03 13:03:40 +00:00 · 2021-06-28 18:11:57 +02:00
parent cc0169ad68
commit dfce7b667d
20 changed files with 1037 additions and 85 deletions
@@ -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/filter.hpp>
+#include <range/v3/view/transform.hpp>
+#include <range/v3/view/zip.hpp>
+#include <range/v3/range/conversion.hpp>

 using namespace std;
 using namespace solidity;
@@ -44,7 +50,7 @@ vector<Statement> generateMemoryStore(
 		Identifier{_debugData, memoryStoreFunction->name},
 		{
 			Literal{_debugData, LiteralKind::Number, _mpos, {}},
-			std::move(_value)
+			move(_value)
 		}
 	}});
 	return result;
@@ -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 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,89 +117,180 @@ 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 parameters 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.debugData,
+				*slot,
+				Identifier{param.debugData, 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.debugData,
+				*slot,
+				Literal{returnVariable.debugData, 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 stackParameters = _functionDefinition.parameters | ranges::views::filter(
+			not_fn(m_memoryOffsetTracker)
+		) | ranges::to<TypedNameList>;
+		// Generate new function without return variable and with only the non-moved parameters.
+		YulString newFunctionName = m_context.dispenser.newName(_functionDefinition.name);
+		m_newFunctionDefinitions.emplace_back(FunctionDefinition{
+			_functionDefinition.debugData,
+			newFunctionName,
+			stackParameters,
+			{},
+			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.debugData, move(memoryVariableInits)};
+		_functionDefinition.body.statements.emplace_back(ExpressionStatement{
+			_functionDefinition.debugData,
+			FunctionCall{
+				_functionDefinition.debugData,
+				Identifier{_functionDefinition.debugData, newFunctionName},
+				stackParameters | ranges::views::transform([&](TypedName const& _arg) {
+					return Expression{Identifier{_arg.debugData, _arg.name}};
+				}) | ranges::to<vector<Expression>>
+			}
+		});
+		_functionDefinition.body.statements.emplace_back(Assignment{
+			_functionDefinition.debugData,
+			{Identifier{_functionDefinition.debugData, _functionDefinition.returnVariables.front().name}},
+			make_unique<Expression>(generateMemoryLoad(
+				m_context.dialect,
+				_functionDefinition.debugData,
+				*m_memoryOffsetTracker(_functionDefinition.returnVariables.front().name)
+			))
+		});
+		return;
+	}
+
+	if (!memoryVariableInits.empty())
+		_functionDefinition.body.statements = move(memoryVariableInits) + move(_functionDefinition.body.statements);
+
+	_functionDefinition.returnVariables = _functionDefinition.returnVariables | ranges::views::filter(
+		not_fn(m_memoryOffsetTracker)
+	) | ranges::to<TypedNameList>;
 }

 void StackToMemoryMover::operator()(Block& _block)
 {
-	using OptionalStatements = std::optional<vector<Statement>>;
-	auto rewriteAssignmentOrVariableDeclaration = [&](
+	using OptionalStatements = optional<vector<Statement>>;
+
+	auto rewriteAssignmentOrVariableDeclarationLeftHandSide = [this](
 		auto& _stmt,
-		auto const& _variables
+		auto& _lhsVars
 	) -> OptionalStatements {
 		using StatementType = decay_t<decltype(_stmt)>;
+
+		auto debugData = _stmt.debugData;
+		if (_lhsVars.size() == 1)
+		{
+			if (optional<YulString> offset = m_memoryOffsetTracker(_lhsVars.front().name))
+				return generateMemoryStore(
+					m_context.dialect,
+					debugData,
+					*offset,
+					_stmt.value ? *move(_stmt.value) : Literal{debugData, LiteralKind::Number, "0"_yulstring, {}}
+				);
+			else
+				return {};
+		}
+		vector<optional<YulString>> rhsMemorySlots;
 		if (_stmt.value)
-			visit(*_stmt.value);
-		bool leftHandSideNeedsMoving = util::contains_if(_variables, [&](auto const& var) {
-			return m_memoryOffsetTracker(var.name);
-		});
-		if (!leftHandSideNeedsMoving)
+		{
+			FunctionCall const* functionCall = get_if<FunctionCall>(_stmt.value.get());
+			yulAssert(functionCall, "");
+			if (m_context.dialect.builtin(functionCall->functionName.name))
+				rhsMemorySlots = vector<optional<YulString>>(_lhsVars.size(), nullopt);
+			else
+				rhsMemorySlots =
+					m_functionReturnVariables.at(functionCall->functionName.name) |
+					ranges::views::transform(m_memoryOffsetTracker) |
+					ranges::to<vector<optional<YulString>>>;
+		}
+		else
+			rhsMemorySlots = vector<optional<YulString>>(_lhsVars.size(), nullopt);
+
+		// 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(rhsMemorySlots, [](optional<YulString> const& _slot) { return _slot.has_value(); }) &&
+			!util::contains_if(_lhsVars, m_memoryOffsetTracker)
+		)
 			return {};

-		if (_variables.size() == 1)
-		{
-			optional<YulString> offset = m_memoryOffsetTracker(_variables.front().name);
-			yulAssert(offset, "");
-			return generateMemoryStore(
-				m_context.dialect,
-				_stmt.debugData,
-				*offset,
-				_stmt.value ? *std::move(_stmt.value) : Literal{_stmt.debugData, LiteralKind::Number, "0"_yulstring, {}}
-			);
-		}
-
-		VariableDeclaration tempDecl{_stmt.debugData, {}, 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.debugData, tempVarName, {}});
+		VariableDeclaration tempDecl{debugData, {}, move(_stmt.value)};

-			if (optional<YulString> offset = m_memoryOffsetTracker(var.name))
+		yulAssert(rhsMemorySlots.size() == _lhsVars.size(), "");
+		for (auto&& [lhsVar, rhsSlot]: ranges::views::zip(_lhsVars, rhsMemorySlots))
+		{
+			unique_ptr<Expression> rhs;
+			if (rhsSlot)
+				rhs = make_unique<Expression>(generateMemoryLoad(m_context.dialect, debugData, *rhsSlot));
+			else
+			{
+				YulString tempVarName = m_nameDispenser.newName(lhsVar.name);
+				tempDecl.variables.emplace_back(TypedName{lhsVar.debugData, tempVarName, {}});
+				rhs = make_unique<Expression>(Identifier{debugData, tempVarName});
+			}
+
+			if (optional<YulString> offset = m_memoryOffsetTracker(lhsVar.name))
 				memoryAssignments += generateMemoryStore(
 					m_context.dialect,
 					_stmt.debugData,
 					*offset,
-					Identifier{_stmt.debugData, tempVarName}
+					move(*rhs)
 				);
 			else
 				variableAssignments.emplace_back(StatementType{
-					_stmt.debugData,
-					{move(var)},
-					make_unique<Expression>(Identifier{_stmt.debugData, tempVarName})
+					debugData,
+					{ move(lhsVar) },
+					move(rhs)
 				});
 		}
-		std::vector<Statement> result;
-		result.emplace_back(std::move(tempDecl));
-		std::reverse(memoryAssignments.begin(), memoryAssignments.end());
-		result += std::move(memoryAssignments);
-		std::reverse(variableAssignments.begin(), variableAssignments.end());
-		result += std::move(variableAssignments);
+
+		vector<Statement> result;
+		if (tempDecl.variables.empty())
+			result.emplace_back(ExpressionStatement{debugData, *move(tempDecl.value)});
+		else
+			result.emplace_back(move(tempDecl));
+		reverse(memoryAssignments.begin(), memoryAssignments.end());
+		result += move(memoryAssignments);
+		reverse(variableAssignments.begin(), variableAssignments.end());
+		result += move(variableAssignments);
 		return OptionalStatements{move(result)};
 	};

 	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 = get_if<Assignment>(&_statement))
+				return rewriteAssignmentOrVariableDeclarationLeftHandSide(*assignment, assignment->variableNames);
+			else if (auto* varDecl = get_if<VariableDeclaration>(&_statement))
+				return rewriteAssignmentOrVariableDeclarationLeftHandSide(*varDecl, varDecl->variables);
+			return {};
 		}
 	);
 }
@@ -188,7 +298,7 @@ void StackToMemoryMover::operator()(Block& _block)
 void StackToMemoryMover::visit(Expression& _expression)
 {
 	ASTModifier::visit(_expression);
-	if (Identifier* identifier = std::get_if<Identifier>(&_expression))
+	if (Identifier* identifier = get_if<Identifier>(&_expression))
 		if (optional<YulString> offset = m_memoryOffsetTracker(identifier->name))
 			_expression = generateMemoryLoad(m_context.dialect, identifier->debugData, *offset);
 }
@@ -202,5 +312,15 @@ optional<YulString> StackToMemoryMover::VariableMemoryOffsetTracker::operator()(
 		return YulString{util::toCompactHexWithPrefix(m_reservedMemory + 32 * (m_numRequiredSlots - slot - 1))};
 	}
 	else
-		return std::nullopt;
+		return nullopt;
+}
+
+optional<YulString> StackToMemoryMover::VariableMemoryOffsetTracker::operator()(TypedName const& _variable) const
+{
+	return (*this)(_variable.name);
+}
+
+optional<YulString> StackToMemoryMover::VariableMemoryOffsetTracker::operator()(Identifier const& _variable) const
+{
+	return (*this)(_variable.name);
 }