Only count variables if only variables are needed.

libyul/backends/evm/EVMCodeTransform.cpp

@@ -668,10 +668,9 @@ bool statementNeedsReturnVariableSetup(Statement const& _statement, vector<Typed
 		holds_alternative<Assignment>(_statement)
 	)
 	{
-		ReferencesCounter referencesCounter{ReferencesCounter::CountWhat::OnlyVariables};
+		map<YulString, size_t> references = VariableReferencesCounter::countReferences(_statement);
-		referencesCounter.visit(_statement);
+		auto isReferenced = [&references](TypedName const& _returnVariable) {
-		auto isReferenced = [&referencesCounter](TypedName const& _returnVariable) {
+			return references.count(_returnVariable.name);
-			return referencesCounter.references().count(_returnVariable.name);
 		};
 		if (ranges::none_of(_returnVariables, isReferenced))
 			return false;

libyul/optimiser/ExpressionJoiner.cpp

@@ -86,7 +86,7 @@ void ExpressionJoiner::visit(Expression& _e)
 
 ExpressionJoiner::ExpressionJoiner(Block& _ast)
 {
-	m_references = ReferencesCounter::countReferences(_ast);
+	m_references = VariableReferencesCounter::countReferences(_ast);
 }
 
 void ExpressionJoiner::handleArguments(vector<Expression>& _arguments)

libyul/optimiser/LoopInvariantCodeMotion.cpp

@@ -71,7 +71,7 @@ bool LoopInvariantCodeMotion::canBePromoted(
 			return false;
 	if (_varDecl.value)
 	{
-		for (auto const& ref: ReferencesCounter::countReferences(*_varDecl.value, ReferencesCounter::OnlyVariables))
+		for (auto const& ref: VariableReferencesCounter::countReferences(*_varDecl.value))
 			if (_varsDefinedInCurrentScope.count(ref.first) || !m_ssaVariables.count(ref.first))
 				return false;
 		SideEffectsCollector sideEffects{m_dialect, *_varDecl.value, &m_functionSideEffects};

libyul/optimiser/NameCollector.cpp

@@ -56,30 +56,62 @@ void ReferencesCounter::operator()(Identifier const& _identifier)
 
 void ReferencesCounter::operator()(FunctionCall const& _funCall)
 {
-	if (m_countWhat == VariablesAndFunctions)
 	++m_references[_funCall.functionName.name];
 	ASTWalker::operator()(_funCall);
 }
 
-map<YulString, size_t> ReferencesCounter::countReferences(Block const& _block, CountWhat _countWhat)
+map<YulString, size_t> ReferencesCounter::countReferences(Block const& _block)
 {
-	ReferencesCounter counter(_countWhat);
+	ReferencesCounter counter;
 	counter(_block);
-	return counter.references();
+	return std::move(counter.m_references);
 }
 
-map<YulString, size_t> ReferencesCounter::countReferences(FunctionDefinition const& _function, CountWhat _countWhat)
+map<YulString, size_t> ReferencesCounter::countReferences(FunctionDefinition const& _function)
 {
-	ReferencesCounter counter(_countWhat);
+	ReferencesCounter counter;
 	counter(_function);
-	return counter.references();
+	return std::move(counter.m_references);
 }
 
-map<YulString, size_t> ReferencesCounter::countReferences(Expression const& _expression, CountWhat _countWhat)
+map<YulString, size_t> ReferencesCounter::countReferences(Expression const& _expression)
 {
-	ReferencesCounter counter(_countWhat);
+	ReferencesCounter counter;
 	counter.visit(_expression);
-	return counter.references();
+	return std::move(counter.m_references);
+}
+
+void VariableReferencesCounter::operator()(Identifier const& _identifier)
+{
+	++m_references[_identifier.name];
+}
+
+map<YulString, size_t> VariableReferencesCounter::countReferences(Block const& _block)
+{
+	VariableReferencesCounter counter;
+	counter(_block);
+	return std::move(counter.m_references);
+}
+
+map<YulString, size_t> VariableReferencesCounter::countReferences(FunctionDefinition const& _function)
+{
+	VariableReferencesCounter counter;
+	counter(_function);
+	return std::move(counter.m_references);
+}
+
+map<YulString, size_t> VariableReferencesCounter::countReferences(Expression const& _expression)
+{
+	VariableReferencesCounter counter;
+	counter.visit(_expression);
+	return std::move(counter.m_references);
+}
+
+map<YulString, size_t> VariableReferencesCounter::countReferences(Statement const& _statement)
+{
+	VariableReferencesCounter counter;
+	counter.visit(_statement);
+	return std::move(counter.m_references);
 }
 
 void AssignmentsSinceContinue::operator()(ForLoop const& _forLoop)

libyul/optimiser/NameCollector.h

@@ -71,23 +71,33 @@ private:
 class ReferencesCounter: public ASTWalker
 {
 public:
-	enum CountWhat { VariablesAndFunctions, OnlyVariables };
-
-	explicit ReferencesCounter(CountWhat _countWhat = VariablesAndFunctions):
-		m_countWhat(_countWhat)
-	{}
-
 	using ASTWalker::operator ();
 	void operator()(Identifier const& _identifier) override;
 	void operator()(FunctionCall const& _funCall) override;
 
-	static std::map<YulString, size_t> countReferences(Block const& _block, CountWhat _countWhat = VariablesAndFunctions);
+	static std::map<YulString, size_t> countReferences(Block const& _block);
-	static std::map<YulString, size_t> countReferences(FunctionDefinition const& _function, CountWhat _countWhat = VariablesAndFunctions);
+	static std::map<YulString, size_t> countReferences(FunctionDefinition const& _function);
-	static std::map<YulString, size_t> countReferences(Expression const& _expression, CountWhat _countWhat = VariablesAndFunctions);
+	static std::map<YulString, size_t> countReferences(Expression const& _expression);
+
+private:
+	std::map<YulString, size_t> m_references;
+};
+
+/**
+ * Specific AST walker that counts all references to all variable declarations.
+ */
+class VariableReferencesCounter: public ASTWalker
+{
+public:
+	using ASTWalker::operator ();
+	void operator()(Identifier const& _identifier) override;
+
+	static std::map<YulString, size_t> countReferences(Block const& _block);
+	static std::map<YulString, size_t> countReferences(FunctionDefinition const& _function);
+	static std::map<YulString, size_t> countReferences(Expression const& _expression);
+	static std::map<YulString, size_t> countReferences(Statement const& _statement);
 
-	std::map<YulString, size_t> const& references() const { return m_references; }
 private:
-	CountWhat m_countWhat = CountWhat::VariablesAndFunctions;
 	std::map<YulString, size_t> m_references;
 };
 

libyul/optimiser/Rematerialiser.cpp

@@ -44,7 +44,7 @@ Rematerialiser::Rematerialiser(
 	bool _onlySelectedVariables
 ):
 	DataFlowAnalyzer(_dialect, MemoryAndStorage::Ignore),
-	m_referenceCounts(ReferencesCounter::countReferences(_ast)),
+	m_referenceCounts(VariableReferencesCounter::countReferences(_ast)),
 	m_varsToAlwaysRematerialize(std::move(_varsToAlwaysRematerialize)),
 	m_onlySelectedVariables(_onlySelectedVariables)
 {
@@ -77,7 +77,9 @@ void Rematerialiser::visit(Expression& _e)
 				{
 					// update reference counts
 					m_referenceCounts[name]--;
-					for (auto const& ref: ReferencesCounter::countReferences(*value->value))
+					for (auto const& ref: VariableReferencesCounter::countReferences(
+						*value->value
+					))
 						m_referenceCounts[ref.first] += ref.second;
 					_e = (ASTCopier{}).translate(*value->value);
 				}

libyul/optimiser/UnusedFunctionParameterPruner.cpp

@@ -43,7 +43,7 @@ using namespace solidity::yul::unusedFunctionsCommon;
 
 void UnusedFunctionParameterPruner::run(OptimiserStepContext& _context, Block& _ast)
 {
-	map<YulString, size_t> references = ReferencesCounter::countReferences(_ast);
+	map<YulString, size_t> references = VariableReferencesCounter::countReferences(_ast);
 	auto used = [&](auto v) -> bool { return references.count(v.name); };
 
 	// Function name and a pair of boolean masks, the first corresponds to parameters and the second