Merge pull request #16750 from argotorg/ssa-cfg-slg-populates-spill-set

SSA CFG: Stack Layout generator populates spill set

Author: clonker

libyul/backends/evm/ssa/CodeTransform.cpp

@@ -18,6 +18,7 @@
 
 #include <libyul/backends/evm/ssa/CodeTransform.h>
 
+#include <libyul/backends/evm/ssa/CallGraph.h>
 #include <libyul/backends/evm/ssa/StackLayoutGenerator.h>
 #include <libyul/backends/evm/ssa/StackShuffler.h>
 #include <libyul/backends/evm/ssa/StackUtils.h>
@@ -52,6 +53,7 @@ void CodeTransform::run
 	ControlFlowGraphs const& controlFlowGraphs = _controlFlowLiveness.controlFlowGraphs.get();
 	yulAssert(controlFlowGraphs.functionGraphs.size() == _controlFlowLiveness.cfgLiveness.size());
 	FunctionLabels const functionLabels = registerFunctionLabels(_assembly, controlFlowGraphs);
+	CallGraph const callGraph(controlFlowGraphs);
 
 	for (std::size_t functionIndex = 0; functionIndex < controlFlowGraphs.functionGraphs.size(); ++functionIndex)
 	{
@@ -62,15 +64,16 @@ void CodeTransform::run
 		auto const& liveness = _controlFlowLiveness.cfgLiveness[functionIndex];
 		yulAssert(liveness);
 		auto const graphID = static_cast<ControlFlowGraphs::FunctionGraphID>(functionIndex);
-		auto const& stackLayout = StackLayoutGenerator::generate(*liveness, callSites, graphID);
+		bool const spillingAllowed = !callGraph.isRecursive(graphID);
+		auto const stackLayoutGeneratorResult = StackLayoutGenerator::generate(*liveness, callSites, graphID, spillingAllowed);
 		CodeTransform transform(
 			_assembly,
 			_builtinContext,
 			controlFlowGraphs,
 			functionLabels,
 			callSites,
 			cfg,
-			stackLayout,
+			stackLayoutGeneratorResult.layout,
 			graphID
 		);
 		transform(cfg.entry);

libyul/backends/evm/ssa/StackLayoutGenerator.cpp

@@ -82,25 +82,29 @@ void declareJunk(StackType& _stack, LivenessAnalysis::LivenessData const& _live)
 
 }
 
-SSACFGStackLayout StackLayoutGenerator::generate(
+StackLayoutGenerator::Result StackLayoutGenerator::generate(
 	LivenessAnalysis const& _liveness,
 	CallSites const& _callSites,
-	ControlFlowGraphs::FunctionGraphID const _graphID
+	ControlFlowGraphs::FunctionGraphID const _graphID,
+	bool const _spillingAllowed
 )
 {
-	return StackLayoutGenerator(_liveness, _callSites, _graphID).m_resultLayout;
+	StackLayoutGenerator generator(_liveness, _callSites, _graphID, _spillingAllowed);
+	return Result{std::move(generator.m_resultLayout), std::move(generator.m_spillSet)};
 }
 
 StackLayoutGenerator::StackLayoutGenerator(
 	LivenessAnalysis const& _liveness,
 	CallSites const& _callSites,
-	ControlFlowGraphs::FunctionGraphID const _graphID
+	ControlFlowGraphs::FunctionGraphID const _graphID,
+	bool const _spillingAllowed
 ):
 	m_cfg(_liveness.cfg()),
 	m_liveness(_liveness),
 	m_callSites(_callSites),
 	m_graphID(_graphID),
 	m_hasFunctionReturnLabel(!_liveness.cfg().isMainGraph() && _liveness.cfg().canContinue),
+	m_spillingAllowed(_spillingAllowed),
 	m_junkAdmittingBlocksFinder(std::make_unique<JunkAdmittingBlocksFinder>(_liveness.cfg(), _liveness.topologicalSort())),
 	m_inputStackProposalsPerBlock(m_cfg.numBlocks()),
 	m_resultLayout(m_cfg.numBlocks())
@@ -199,7 +203,8 @@ void StackLayoutGenerator::defineStackIn(SSACFG::BlockId const& _blockId)
 					stack,
 					proposals[i],
 					{},
-					proposals[i].size()
+					proposals[i].size(),
+					&m_spillSet
 				);
 				yulAssert(shuffleResult.status == StackShufflerResult::Status::Admissible);
 				cumulativeCost += stack.callbacks().opGas;
@@ -258,15 +263,20 @@ void StackLayoutGenerator::visitBlock(SSACFG::BlockId const& _blockId)
 
 		StackSlotLiveness const opLiveOutSlots = toStackSlotLiveness(m_cfg, opLiveOutWithoutOutputs);
 		{
-			StackData const target = findOptimalTarget(
+			auto [target, plannedSpillSet] = findOptimalTarget(
 				stack.data(),
 				requiredStackTop,
 				opLiveOutSlots,
 				junkCanBeAdded,
-				m_hasFunctionReturnLabel
+				m_hasFunctionReturnLabel,
+				m_spillSet,
+				m_spillingAllowed
 			);
-			auto const shuffleResult = StackShuffler<StackType::Callbacks>::shuffle(stack, target);
+			auto const spillCountBefore = m_spillSet.numSpilled();
+			m_spillSet = std::move(plannedSpillSet);
+			auto const shuffleResult = shuffleWithSpillDiscovery(currentStackData, target, m_spillSet);
 			yulAssert(shuffleResult.status == StackShufflerResult::Status::Admissible);
+			yulAssert(m_spillingAllowed || m_spillSet.numSpilled() == spillCountBefore, "Spilling not allowed, stack too deep.");
 		}
 
 		blockLayout.operationIn.push_back(currentStackData);
@@ -293,15 +303,20 @@ void StackLayoutGenerator::visitBlock(SSACFG::BlockId const& _blockId)
 				{
 					auto const condition = Slot::makeValue(m_cfg, _cJump.condition);
 					StackSlotLiveness const blockLiveOutSlots = toStackSlotLiveness(m_cfg, blockLiveOut);
-					StackData const target = findOptimalTarget(
+					auto [target, plannedSpillSet] = findOptimalTarget(
 						stack.data(),
 						{condition},
 						blockLiveOutSlots,
 						false,
-						m_hasFunctionReturnLabel
+						m_hasFunctionReturnLabel,
+						m_spillSet,
+						m_spillingAllowed
 					);
-					auto const shuffleResult = StackShuffler<StackType::Callbacks>::shuffle(stack, target);
+					auto const spillCountBefore = m_spillSet.numSpilled();
+					m_spillSet = std::move(plannedSpillSet);
+					auto const shuffleResult = shuffleWithSpillDiscovery(currentStackData, target, m_spillSet);
 					yulAssert(shuffleResult.status == StackShufflerResult::Status::Admissible);
+					yulAssert(m_spillingAllowed || m_spillSet.numSpilled() == spillCountBefore, "Spilling not allowed, stack too deep.");
 				}
 
 				yulAssert(!stack.empty() && stack.top().isValue() && stack.top().value() == _cJump.condition);
@@ -322,8 +337,10 @@ void StackLayoutGenerator::visitBlock(SSACFG::BlockId const& _blockId)
 				// in case there are return values, let's bring the function return label to the top
 				StackData returnStack = _functionReturn.returnValues | ranges::views::transform([this](InstId const _id) { return StackSlot::makeValue(m_cfg, _id); }) | ranges::to<std::vector>;
 				returnStack.push_back(StackSlot::makeFunctionReturnLabel(m_graphID));
-				auto const shuffleResult = StackShuffler<StackType::Callbacks>::shuffle(stack, returnStack);
+				auto const spillCountBefore = m_spillSet.numSpilled();
+				auto const shuffleResult = shuffleWithSpillDiscovery(currentStackData, returnStack, m_spillSet);
 				yulAssert(shuffleResult.status == StackShufflerResult::Status::Admissible);
+				yulAssert(m_spillingAllowed || m_spillSet.numSpilled() == spillCountBefore, "Spilling not allowed, stack too deep.");
 				blockLayout.exitIn = currentStackData;
 			},
 			[&](SSACFG::BasicBlock::Jump const& _jump) {

libyul/backends/evm/ssa/StackLayoutGenerator.h

@@ -18,6 +18,8 @@
 
 #pragma once
 
+#include <libyul/backends/evm/ssa/spill/SpillSet.h>
+
 #include <libyul/backends/evm/ssa/Stack.h>
 #include <libyul/backends/evm/ssa/StackLayout.h>
 
@@ -32,17 +34,27 @@ class StackLayoutGenerator
 {
 public:
 	using Slot = StackSlot;
-	static SSACFGStackLayout generate(
+
+	/// the per-block stack layout plus the set of values the layout generator decided to spill to memory
+	struct Result
+	{
+		SSACFGStackLayout layout;
+		spill::SpillSet spillSet;
+	};
+
+	static Result generate(
 		LivenessAnalysis const& _liveness,
 		CallSites const& _callSites,
-		ControlFlowGraphs::FunctionGraphID _graphID
+		ControlFlowGraphs::FunctionGraphID _graphID,
+		bool _spillingAllowed
 	);
 
 private:
 	explicit StackLayoutGenerator(
 		LivenessAnalysis const& _liveness,
 		CallSites const& _callSites,
-		ControlFlowGraphs::FunctionGraphID _graphID
+		ControlFlowGraphs::FunctionGraphID _graphID,
+		bool _spillingAllowed
 	);
 
 	void defineStackIn(SSACFG::BlockId const& _blockId);
@@ -53,13 +65,15 @@ class StackLayoutGenerator
 	CallSites const& m_callSites;
 	ControlFlowGraphs::FunctionGraphID m_graphID;
 	bool m_hasFunctionReturnLabel;
+	bool m_spillingAllowed;
 
 	std::unique_ptr<JunkAdmittingBlocksFinder> m_junkAdmittingBlocksFinder;
 	// Per-edge stack proposals for defineStackIn.
 	// m_inputStackProposalsPerBlock[blockId] contains (predecessorId, edgeStack) pairs
 	// representing the stack state flowing from each predecessor into the block.
 	std::vector<std::vector<std::pair<SSACFG::BlockId, StackData>>> m_inputStackProposalsPerBlock;
 	SSACFGStackLayout m_resultLayout;
+	spill::SpillSet m_spillSet;
 };
 
 }

libyul/backends/evm/ssa/StackShuffler.h

@@ -951,4 +951,49 @@ class StackShuffler
 	}
 };
 
+[[nodiscard]] inline StackShufflerResult shuffleWithSpillDiscovery(
+	StackData& _data,
+	StackData const& _args,
+	StackSlotLiveness const& _liveOut,
+	std::size_t const _targetStackSize,
+	spill::SpillSet& _spilledVariables
+)
+{
+	StackData const initialData = _data;
+	StackShufflerResult result;
+	do
+	{
+		_data = initialData;
+		Stack<> stack(_data, {});
+		result = StackShuffler<NoOpStackManipulationCallbacks>::shuffle(stack, _args, _liveOut, _targetStackSize, &_spilledVariables);
+		switch (result.status)
+		{
+		case StackShufflerResult::Status::Continue:
+			yulAssert(false);
+		case StackShufflerResult::Status::Admissible:
+			break;
+		case StackShufflerResult::Status::StackTooDeep:
+		{
+			yulAssert(result.culprit.isValue() && !result.culprit.isLiteralValue());
+			yulAssert(!_spilledVariables.isSpilled(result.culprit.value()));
+			_spilledVariables.add(result.culprit.value());
+			break;
+		}
+		case StackShufflerResult::Status::MaxIterationsReached:
+			break;
+		}
+	}
+	while (result.status == StackShufflerResult::Status::StackTooDeep);
+	return result;
+}
+
+[[nodiscard]] inline StackShufflerResult shuffleWithSpillDiscovery(
+	StackData& _data,
+	StackData const& _target,
+	spill::SpillSet& _spilledVariables)
+{
+	return shuffleWithSpillDiscovery(_data, _target, {}, _target.size(), _spilledVariables);
+}
+
+
 }

libyul/backends/evm/ssa/StackUtils.cpp

@@ -54,13 +54,19 @@ void GasAccumulatingCallbacks::push(StackSlot const& _slot)
 		auto const op = cfg.evmDialect.evmVersion().hasPush0() ? evmasm::Instruction::PUSH0 : evmasm::Instruction::CODESIZE;
 		opGas += evmasm::GasMeter::runGas(op, cfg.evmDialect.evmVersion());
 	}
-	else
+	else if (_slot.isFunctionCallReturnLabel())
 	{
-		yulAssert(_slot.isFunctionCallReturnLabel(), "we can only push literals, junk, and function call return labels");
 		// this is a jump dest, we don't really know yet how big it is going to be, just assume that it fits into
 		// a 2-byte number
 		opGas += evmasm::GasMeter::runGas(evmasm::Instruction::PUSH2, cfg.evmDialect.evmVersion());
 	}
+	else
+	{
+		// Spilled SSA value
+		yulAssert(_slot.isValue(), "unexpected slot kind in GasAccumulatingCallbacks::push");
+		opGas += evmasm::GasMeter::runGas(evmasm::Instruction::PUSH32, cfg.evmDialect.evmVersion());
+		opGas += evmasm::GasMeter::runGas(evmasm::Instruction::MLOAD, cfg.evmDialect.evmVersion());
+	}
 }
 
 void GasAccumulatingCallbacks::pop()
@@ -80,13 +86,15 @@ StackData solidity::yul::ssa::stackPreImage(SSACFG const& _cfg, StackData _stack
 	return _stack;
 }
 
-StackData solidity::yul::ssa::findOptimalTarget
+OptimalTarget solidity::yul::ssa::findOptimalTarget
 (
 	StackData const& _stackData,
 	StackData const& _targetArgs,
 	StackSlotLiveness const& _targetLiveOut,
 	bool const _canIntroduceJunk,
-	bool const _hasFunctionReturnLabel
+	bool const _hasFunctionReturnLabel,
+	spill::SpillSet const& _spillSet,
+	bool const _spillingAllowed
 )
 {
 	std::size_t const minSize = _targetLiveOut.size() + _targetArgs.size() + (_hasFunctionReturnLabel ? 1 : 0);
@@ -105,10 +113,11 @@ StackData solidity::yul::ssa::findOptimalTarget
 
 	StackData data;
 	data.reserve(startSize + maxUpwardExpansion);
+	spill::SpillSet spillSet;
 	auto const evaluateCost = [&](std::size_t const _targetSize) -> std::size_t
 	{
 		StackShufflerResult result;
-		spill::SpillSet spillSet;
+		spillSet = _spillSet;
 		OpsCountingCallbacks callbacks;
 		do
 		{
@@ -142,6 +151,7 @@ StackData solidity::yul::ssa::findOptimalTarget
 
 	std::size_t bestCost = evaluateCost(startSize);
 	StackData bestData = data;
+	spill::SpillSet bestSpillSet = spillSet;
 
 	// On non-reverting paths, only search downward from pivot to avoid growing the stack.
 	// On reverting paths, search in both directions since stack cleanup doesn't matter.
@@ -158,6 +168,7 @@ StackData solidity::yul::ssa::findOptimalTarget
 			{
 				bestCost = cost;
 				bestData = data;
+				bestSpillSet = spillSet;
 				consecutiveIncreases = 0;
 			}
 			else if (++consecutiveIncreases >= stopAfter)
@@ -175,14 +186,16 @@ StackData solidity::yul::ssa::findOptimalTarget
 			{
 				bestCost = cost;
 				bestData = data;
+				bestSpillSet = spillSet;
 				consecutiveIncreases = 0;
 			}
 			else if (++consecutiveIncreases >= stopAfter)
 				break;
 		}
 	}
 
-	return bestData;
+	yulAssert(_spillingAllowed || bestSpillSet.numSpilled() == _spillSet.numSpilled(), "Spilling not allowed, stack too deep.");
+	return OptimalTarget{std::move(bestData), std::move(bestSpillSet)};
 }
 
 CallSites solidity::yul::ssa::gatherCallSites(SSACFG const& _cfg)

libyul/backends/evm/ssa/StackUtils.h

@@ -18,6 +18,8 @@
 
 #pragma once
 
+#include <libyul/backends/evm/ssa/spill/SpillSet.h>
+
 #include <libyul/backends/evm/ssa/PhiInverse.h>
 #include <libyul/backends/evm/ssa/Stack.h>
 #include <libyul/backends/evm/ssa/StackSlotLiveness.h>
@@ -63,15 +65,26 @@ struct GasAccumulatingCallbacks
 /// Transform stack data by replacing all its phi variables with their respective preimages.
 StackData stackPreImage(SSACFG const& _cfg, StackData _stack, PhiInverse const& _phiInverse);
 
+struct OptimalTarget
+{
+	StackData data;
+	spill::SpillSet spillSet;
+};
+
 /// Searches for the cheapest target stack size for shuffling _stackData towards
-/// (_targetArgs on top, _targetLiveOut anywhere in the tail), then realizes that size and
-/// returns the resolved post-shuffle StackData
-StackData findOptimalTarget(
+/// (_targetArgs on top, _targetLiveOut anywhere in the tail). Realizes that size and returns the resolved post-shuffle
+/// StackData together with the spill set that was needed to reach it.
+/// Callers must use the returned spill set for any subsequent shuffle of the real stack to this target, otherwise
+/// live values that were planned-spilled here will silently disappear from the real stack while
+/// staying outside the caller's spill set.
+OptimalTarget findOptimalTarget(
 	StackData const& _stackData,
 	StackData const& _targetArgs,
 	StackSlotLiveness const& _targetLiveOut,
 	bool _canIntroduceJunk,
-	bool _hasFunctionReturnLabel
+	bool _hasFunctionReturnLabel,
+	spill::SpillSet const& _spillSet,
+	bool _spillingAllowed
 );
 
 CallSites gatherCallSites(SSACFG const& _cfg);