Introduce forEach for yul ast nodes.

libyul/CMakeLists.txt

@@ -140,8 +140,6 @@ add_library(yul
 	optimiser/FullInliner.h
 	optimiser/FunctionCallFinder.cpp
 	optimiser/FunctionCallFinder.h
-	optimiser/FunctionDefinitionCollector.cpp
-	optimiser/FunctionDefinitionCollector.h
 	optimiser/FunctionGrouper.cpp
 	optimiser/FunctionGrouper.h
 	optimiser/FunctionHoister.cpp

libyul/ControlFlowSideEffectsCollector.cpp

@@ -18,8 +18,6 @@
 
 #include <libyul/ControlFlowSideEffectsCollector.h>
 
-#include <libyul/optimiser/FunctionDefinitionCollector.h>
-
 #include <libyul/AST.h>
 #include <libyul/Dialect.h>
 

libyul/optimiser/ASTWalker.h

@@ -102,4 +102,35 @@ protected:
 	}
 };
 
+namespace detail
+{
+template <
+	typename Node,
+	typename Visitor,
+	typename Base = std::conditional_t<std::is_const_v<Node>, ASTWalker, ASTModifier>
+>
+struct ForEach: Base
+{
+	template<typename Callable>
+	ForEach(Callable&& _visitor): visitor(std::forward<Callable>(_visitor)) {}
+
+	using Base::operator();
+	void operator()(Node& _node) override
+	{
+		visitor(_node);
+		Base::operator()(_node);
+	}
+
+	Visitor visitor;
+};
+}
+
+/// Helper function that traverses the AST and calls the visitor for each
+/// node of a specific type.
+template<typename Node, typename Entry, typename Visitor>
+void forEach(Entry&& _entry, Visitor&& _visitor)
+{
+	detail::ForEach<Node, std::decay_t<Visitor>>{std::forward<Visitor>(_visitor)}(std::forward<Entry>(_entry));
+}
+
 }

libyul/optimiser/DataFlowAnalyzer.cpp

@@ -45,9 +45,9 @@ DataFlowAnalyzer::DataFlowAnalyzer(
 	Dialect const& _dialect,
 	map<YulString, SideEffects> _functionSideEffects
 ):
-m_dialect(_dialect),
+	m_dialect(_dialect),
-m_functionSideEffects(std::move(_functionSideEffects)),
+	m_functionSideEffects(std::move(_functionSideEffects)),
-m_knowledgeBase(_dialect, m_value)
+	m_knowledgeBase(_dialect, m_value)
 {
 	if (auto const* builtin = _dialect.memoryStoreFunction(YulString{}))
 		m_storeFunctionName[static_cast<unsigned>(StoreLoadLocation::Memory)] = builtin->name;
@@ -123,9 +123,7 @@ void DataFlowAnalyzer::operator()(If& _if)
 
 	joinKnowledge(storage, memory);
 
-	Assignments assignments;
+	clearValues(assignedVariableNames(_if.body));
-	assignments(_if.body);
-	clearValues(assignments.names());
 }
 
 void DataFlowAnalyzer::operator()(Switch& _switch)
@@ -140,11 +138,10 @@ void DataFlowAnalyzer::operator()(Switch& _switch)
 		(*this)(_case.body);
 		joinKnowledge(storage, memory);
 
-		Assignments assignments;
+		set<YulString> variables = assignedVariableNames(_case.body);
-		assignments(_case.body);
+		assignedVariables += variables;
-		assignedVariables += assignments.names();
 		// This is a little too destructive, we could retain the old values.
-		clearValues(assignments.names());
+		clearValues(variables);
 		clearKnowledgeIfInvalidated(_case.body);
 	}
 	for (auto& _case: _switch.cases)
@@ -190,10 +187,9 @@ void DataFlowAnalyzer::operator()(ForLoop& _for)
 	AssignmentsSinceContinue assignmentsSinceCont;
 	assignmentsSinceCont(_for.body);
 
-	Assignments assignments;
+	set<YulString> assignedVariables =
-	assignments(_for.body);
+		assignedVariableNames(_for.body) + assignedVariableNames(_for.post);
-	assignments(_for.post);
+	clearValues(assignedVariables);
-	clearValues(assignments.names());
 
 	// break/continue are tricky for storage and thus we almost always clear here.
 	clearKnowledgeIfInvalidated(*_for.condition);
@@ -205,7 +201,7 @@ void DataFlowAnalyzer::operator()(ForLoop& _for)
 	clearValues(assignmentsSinceCont.names());
 	clearKnowledgeIfInvalidated(_for.body);
 	(*this)(_for.post);
-	clearValues(assignments.names());
+	clearValues(assignedVariables);
 	clearKnowledgeIfInvalidated(*_for.condition);
 	clearKnowledgeIfInvalidated(_for.post);
 	clearKnowledgeIfInvalidated(_for.body);

libyul/optimiser/FunctionDefinitionCollector.cpp

@@ -1,36 +0,0 @@
-/*
-	This file is part of solidity.
-
-	solidity is free software: you can redistribute it and/or modify
-	it under the terms of the GNU General Public License as published by
-	the Free Software Foundation, either version 3 of the License, or
-	(at your option) any later version.
-
-	solidity is distributed in the hope that it will be useful,
-	but WITHOUT ANY WARRANTY; without even the implied warranty of
-	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-	GNU General Public License for more details.
-
-	You should have received a copy of the GNU General Public License
-	along with solidity.  If not, see <http://www.gnu.org/licenses/>.
-*/
-
-#include <libyul/optimiser/FunctionDefinitionCollector.h>
-#include <libyul/AST.h>
-
-using namespace std;
-using namespace solidity;
-using namespace solidity::yul;
-
-map<YulString, FunctionDefinition const*> FunctionDefinitionCollector::run(Block const& _block)
-{
-	FunctionDefinitionCollector functionDefinitionCollector;
-	functionDefinitionCollector(_block);
-	return functionDefinitionCollector.m_functionDefinitions;
-}
-
-void FunctionDefinitionCollector::operator()(FunctionDefinition const& _functionDefinition)
-{
-	m_functionDefinitions[_functionDefinition.name] = &_functionDefinition;
-	ASTWalker::operator()(_functionDefinition);
-}

libyul/optimiser/FunctionDefinitionCollector.h

@@ -1,44 +0,0 @@
-/*
-	This file is part of solidity.
-
-	solidity is free software: you can redistribute it and/or modify
-	it under the terms of the GNU General Public License as published by
-	the Free Software Foundation, either version 3 of the License, or
-	(at your option) any later version.
-
-	solidity is distributed in the hope that it will be useful,
-	but WITHOUT ANY WARRANTY; without even the implied warranty of
-	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-	GNU General Public License for more details.
-
-	You should have received a copy of the GNU General Public License
-	along with solidity.  If not, see <http://www.gnu.org/licenses/>.
-*/
-/**
- * AST walker that finds all function definitions and stores them into a map indexed by the function names.
- */
-#pragma once
-
-#include <libyul/optimiser/ASTWalker.h>
-
-#include <map>
-
-namespace solidity::yul
-{
-
-/**
- * AST walker that finds all function definitions and stores them into a map indexed by the function names.
- *
- * Prerequisite: Disambiguator
- */
-class FunctionDefinitionCollector: ASTWalker
-{
-public:
-	static std::map<YulString, FunctionDefinition const*> run(Block const& _block);
-private:
-	using ASTWalker::operator();
-	void operator()(FunctionDefinition const& _functionDefinition) override;
-	std::map<YulString, FunctionDefinition const*> m_functionDefinitions;
-};
-
-}

libyul/optimiser/NameCollector.cpp

@@ -78,13 +78,6 @@ map<YulString, size_t> ReferencesCounter::countReferences(Expression const& _exp
 	return counter.references();
 }
 
-void Assignments::operator()(Assignment const& _assignment)
-{
-	for (auto const& var: _assignment.variableNames)
-		m_names.emplace(var.name);
-}
-
-
 void AssignmentsSinceContinue::operator()(ForLoop const& _forLoop)
 {
 	m_forLoopDepth++;
@@ -109,3 +102,22 @@ void AssignmentsSinceContinue::operator()(FunctionDefinition const&)
 {
 	yulAssert(false, "");
 }
+
+std::set<YulString> solidity::yul::assignedVariableNames(Block const& _code)
+{
+	std::set<YulString> names;
+	forEach<Assignment const>(_code, [&](Assignment const& _assignment) {
+		for (auto const& var: _assignment.variableNames)
+			names.emplace(var.name);
+	});
+	return names;
+}
+
+map<YulString, FunctionDefinition const*> solidity::yul::allFunctionDefinitions(Block const& _block)
+{
+	std::map<YulString, FunctionDefinition const*> result;
+	forEach<FunctionDefinition const>(_block, [&](FunctionDefinition const& _function) {
+		result[_function.name] = &_function;
+	});
+	return result;
+}

libyul/optimiser/NameCollector.h

@@ -91,20 +91,6 @@ private:
 	std::map<YulString, size_t> m_references;
 };
 
-/**
- * Specific AST walker that finds all variables that are assigned to.
- */
-class Assignments: public ASTWalker
-{
-public:
-	using ASTWalker::operator ();
-	void operator()(Assignment const& _assignment) override;
-
-	std::set<YulString> const& names() const { return m_names; }
-private:
-	std::set<YulString> m_names;
-};
-
 /**
  * Collects all names from a given continue statement on onwards.
  *
@@ -130,4 +116,12 @@ private:
 	std::set<YulString> m_names;
 };
 
+/// @returns the names of all variables that are assigned to inside @a _code.
+/// (ignores variable declarations)
+std::set<YulString> assignedVariableNames(Block const& _code);
+
+/// @returns all function definitions anywhere in the AST.
+/// Requires disambiguated source.
+std::map<YulString, FunctionDefinition const*> allFunctionDefinitions(Block const& _block);
+
 }

libyul/optimiser/SSATransform.cpp

@@ -196,12 +196,7 @@ void IntroduceControlFlowSSA::operator()(ForLoop& _for)
 {
 	yulAssert(_for.pre.statements.empty(), "For loop init rewriter not run.");
 
-	Assignments assignments;
+	for (auto const& var: assignedVariableNames(_for.body) + assignedVariableNames(_for.post))
-	assignments(_for.body);
-	assignments(_for.post);
-
-
-	for (auto const& var: assignments.names())
 		if (m_variablesInScope.count(var))
 			m_variablesToReassign.insert(var);
 
@@ -359,11 +354,7 @@ void PropagateValues::operator()(ForLoop& _for)
 {
 	yulAssert(_for.pre.statements.empty(), "For loop init rewriter not run.");
 
-	Assignments assignments;
+	for (auto const& var: assignedVariableNames(_for.body) + assignedVariableNames(_for.post))
-	assignments(_for.body);
-	assignments(_for.post);
-
-	for (auto const& var: assignments.names())
 		m_currentVariableValues.erase(var);
 
 	visit(*_for.condition);
@@ -389,11 +380,10 @@ void PropagateValues::operator()(Block& _block)
 void SSATransform::run(OptimiserStepContext& _context, Block& _ast)
 {
 	TypeInfo typeInfo(_context.dialect, _ast);
-	Assignments assignments;
+	set<YulString> assignedVariables = assignedVariableNames(_ast);
-	assignments(_ast);
+	IntroduceSSA{_context.dispenser, assignedVariables, typeInfo}(_ast);
-	IntroduceSSA{_context.dispenser, assignments.names(), typeInfo}(_ast);
+	IntroduceControlFlowSSA{_context.dispenser, assignedVariables, typeInfo}(_ast);
-	IntroduceControlFlowSSA{_context.dispenser, assignments.names(), typeInfo}(_ast);
+	PropagateValues{assignedVariables}(_ast);
-	PropagateValues{assignments.names()}(_ast);
 }
 
 

libyul/optimiser/StackLimitEvader.cpp

@@ -18,8 +18,8 @@
 #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/NameCollector.h>
 #include <libyul/optimiser/StackToMemoryMover.h>
 #include <libyul/backends/evm/ControlFlowGraphBuilder.h>
 #include <libyul/backends/evm/EVMDialect.h>
@@ -193,7 +193,7 @@ void StackLimitEvader::run(
 		if (_unreachableVariables.count(function))
 			return;
 
-	map<YulString, FunctionDefinition const*> functionDefinitions = FunctionDefinitionCollector::run(*_object.code);
+	map<YulString, FunctionDefinition const*> functionDefinitions = allFunctionDefinitions(*_object.code);
 
 	MemoryOffsetAllocator memoryOffsetAllocator{_unreachableVariables, callGraph.functionCalls, functionDefinitions};
 	uint64_t requiredSlots = memoryOffsetAllocator.run();

libyul/optimiser/StackToMemoryMover.cpp

@@ -15,7 +15,7 @@
 	along with solidity.  If not, see <http://www.gnu.org/licenses/>.
 */
 #include <libyul/optimiser/StackToMemoryMover.h>
-#include <libyul/optimiser/FunctionDefinitionCollector.h>
+#include <libyul/optimiser/NameCollector.h>
 #include <libyul/optimiser/NameDispenser.h>
 #include <libyul/backends/evm/EVMDialect.h>
 
@@ -87,7 +87,7 @@ void StackToMemoryMover::run(
 		_context,
 		memoryOffsetTracker,
 		util::applyMap(
-			FunctionDefinitionCollector::run(_block),
+			allFunctionDefinitions(_block),
 			util::mapTuple([](YulString _name, FunctionDefinition const* _funDef) {
 				return make_pair(_name, _funDef->returnVariables);
 			}),