StackToMemoryMover: encapsulate memory offset tracking into its own class.

libyul/optimiser/StackToMemoryMover.cpp

@@ -59,20 +59,17 @@ void StackToMemoryMover::run(
 	Block& _block
 )
 {
-	StackToMemoryMover stackToMemoryMover(_context, _reservedMemory, _memorySlots, _numRequiredSlots);
+	VariableMemoryOffsetTracker memoryOffsetTracker(_reservedMemory, _memorySlots, _numRequiredSlots);
+	StackToMemoryMover stackToMemoryMover(_context, memoryOffsetTracker);
 	stackToMemoryMover(_block);
 }
 
 StackToMemoryMover::StackToMemoryMover(
 	OptimiserStepContext& _context,
-	u256 _reservedMemory,
+	VariableMemoryOffsetTracker const& _memoryOffsetTracker
-	map<YulString, uint64_t> const& _memorySlots,
-	uint64_t _numRequiredSlots
 ):
 m_context(_context),
-m_reservedMemory(std::move(_reservedMemory)),
+m_memoryOffsetTracker(_memoryOffsetTracker),
-m_memorySlots(_memorySlots),
-m_numRequiredSlots(_numRequiredSlots),
 m_nameDispenser(_context.dispenser)
 {
 	auto const* evmDialect = dynamic_cast<EVMDialect const*>(&_context.dialect);
@@ -85,7 +82,7 @@ m_nameDispenser(_context.dispenser)
 void StackToMemoryMover::operator()(FunctionDefinition& _functionDefinition)
 {
 	for (TypedName const& param: _functionDefinition.parameters + _functionDefinition.returnVariables)
-		if (m_memorySlots.count(param.name))
+		if (m_memoryOffsetTracker(param.name))
 		{
 			// TODO: we cannot handle function parameters yet.
 			return;
@@ -98,7 +95,7 @@ void StackToMemoryMover::operator()(Block& _block)
 	using OptionalStatements = std::optional<vector<Statement>>;
 	auto containsVariableNeedingEscalation = [&](auto const& _variables) {
 		return util::contains_if(_variables, [&](auto const& var) {
-			return m_memorySlots.count(var.name);
+			return m_memoryOffsetTracker(var.name);
 		});
 	};
 	auto rewriteAssignmentOrVariableDeclaration = [&](
@@ -107,12 +104,16 @@ void StackToMemoryMover::operator()(Block& _block)
 		std::unique_ptr<Expression> _value
 	) -> std::vector<Statement> {
 		if (_variables.size() == 1)
+		{
+			optional<YulString> offset = m_memoryOffsetTracker(_variables.front().name);
+			yulAssert(offset, "");
 			return generateMemoryStore(
 				m_context.dialect,
 				_loc,
-				memoryOffset(_variables.front().name),
+				*offset,
 				_value ? *std::move(_value) : Literal{_loc, LiteralKind::Number, "0"_yulstring, {}}
 			);
+		}
 
 		VariableDeclaration tempDecl{_loc, {}, std::move(_value)};
 		vector<Statement> memoryAssignments;
@@ -122,11 +123,11 @@ void StackToMemoryMover::operator()(Block& _block)
 			YulString tempVarName = m_nameDispenser.newName(var.name);
 			tempDecl.variables.emplace_back(TypedName{var.location, tempVarName, {}});
 
-			if (m_memorySlots.count(var.name))
+			if (optional<YulString> offset = m_memoryOffsetTracker(var.name))
 				memoryAssignments += generateMemoryStore(
 					m_context.dialect,
 					_loc,
-					memoryOffset(var.name),
+					*offset,
 					Identifier{_loc, tempVarName}
 				);
 			else if constexpr (std::is_same_v<std::decay_t<decltype(var)>, Identifier>)
@@ -185,35 +186,36 @@ void StackToMemoryMover::operator()(Block& _block)
 
 void StackToMemoryMover::visit(Expression& _expression)
 {
-	if (
+	if (Identifier* identifier = std::get_if<Identifier>(&_expression))
-		Identifier* identifier = std::get_if<Identifier>(&_expression);
+		if (optional<YulString> offset = m_memoryOffsetTracker(identifier->name))
-		identifier && m_memorySlots.count(identifier->name)
+		{
-	)
+			BuiltinFunction const* memoryLoadFunction = m_context.dialect.memoryLoadFunction(m_context.dialect.defaultType);
-	{
+			yulAssert(memoryLoadFunction, "");
-		BuiltinFunction const* memoryLoadFunction = m_context.dialect.memoryLoadFunction(m_context.dialect.defaultType);
+			langutil::SourceLocation loc = identifier->location;
-		yulAssert(memoryLoadFunction, "");
+			_expression = FunctionCall{
-		langutil::SourceLocation loc = identifier->location;
+				loc,
-		_expression = FunctionCall{
+				Identifier{loc, memoryLoadFunction->name}, {
-			loc,
+					Literal{
-			Identifier{loc, memoryLoadFunction->name}, {
+						loc,
-				Literal{
+						LiteralKind::Number,
-					loc,
+						*offset,
-					LiteralKind::Number,
+						{}
-					memoryOffset(identifier->name),
+					}
-					{}
 				}
-			}
+			};
-		};
+			return;
+		}
+	ASTModifier::visit(_expression);
+}
+
+optional<YulString> StackToMemoryMover::VariableMemoryOffsetTracker::operator()(YulString _variable) const
+{
+	if (m_memorySlots.count(_variable))
+	{
+		uint64_t slot = m_memorySlots.at(_variable);
+		yulAssert(slot < m_numRequiredSlots, "");
+		return YulString{util::toCompactHexWithPrefix(m_reservedMemory + 32 * (m_numRequiredSlots - slot - 1))};
 	}
 	else
-		ASTModifier::visit(_expression);
+		return std::nullopt;
 }
-
-YulString StackToMemoryMover::memoryOffset(YulString _variable)
-{
-	yulAssert(m_memorySlots.count(_variable), "");
-	uint64_t slot = m_memorySlots.at(_variable);
-	yulAssert(slot < m_numRequiredSlots, "");
-	return YulString{util::toCompactHexWithPrefix(m_reservedMemory + 32 * (m_numRequiredSlots - slot - 1))};
-}
-

libyul/optimiser/StackToMemoryMover.h

@@ -100,20 +100,32 @@ public:
 	void operator()(Block& _block) override;
 	void visit(Expression& _expression) override;
 private:
+	class VariableMemoryOffsetTracker
+	{
+	public:
+		VariableMemoryOffsetTracker(
+			u256 _reservedMemory,
+			std::map<YulString, uint64_t> const& _memorySlots,
+			uint64_t _numRequiredSlots
+		): m_reservedMemory(_reservedMemory), m_memorySlots(_memorySlots), m_numRequiredSlots(_numRequiredSlots)
+		{}
+
+		/// @returns a YulString containing the memory offset to be assigned to @a _variable as number literal
+		/// or std::nullopt if the variable should not be moved.
+		std::optional<YulString> operator()(YulString _variable) const;
+	private:
+		u256 m_reservedMemory;
+		std::map<YulString, uint64_t> const& m_memorySlots;
+		uint64_t m_numRequiredSlots = 0;
+	};
+
 	StackToMemoryMover(
 		OptimiserStepContext& _context,
-		u256 _reservedMemory,
+		VariableMemoryOffsetTracker const& _memoryOffsetTracker
-		std::map<YulString, uint64_t> const& _memorySlots,
-		uint64_t _numRequiredSlots
 	);
 
-	/// @returns a YulString containing the memory offset to be assigned to @a _variable as number literal.
-	YulString memoryOffset(YulString _variable);
-
 	OptimiserStepContext& m_context;
-	u256 m_reservedMemory;
+	VariableMemoryOffsetTracker const& m_memoryOffsetTracker;
-	std::map<YulString, uint64_t> const& m_memorySlots;
-	uint64_t m_numRequiredSlots = 0;
 	NameDispenser& m_nameDispenser;
 };