Add option to recurse referencedSourceUnits

2 files changed, 10 insertions, 5 deletions

diff --git a/libsolidity/ast/AST.cpp b/libsolidity/ast/AST.cpp
index c5fb35f6..a126e6ed 100644
--- a/libsolidity/ast/AST.cpp
+++ b/libsolidity/ast/AST.cpp
@@ -84,12 +84,17 @@ SourceUnitAnnotation& SourceUnit::annotation() const
 	return dynamic_cast<SourceUnitAnnotation&>(*m_annotation);
 }
 
-vector<SourceUnit const*> SourceUnit::referencedSourceUnits() const
+set<SourceUnit const*> SourceUnit::referencedSourceUnits(bool _recurse) const
 {
-	vector<SourceUnit const*> sourceUnits;
+	set<SourceUnit const*> sourceUnits;
 	for (ImportDirective const* importDirective: filteredNodes<ImportDirective>(nodes()))
 	{
-		sourceUnits.push_back(importDirective->annotation().sourceUnit);
+		sourceUnits.insert(importDirective->annotation().sourceUnit);
+		if (_recurse)
+		{
+			set<SourceUnit const*> referencedSourceUnits = importDirective->annotation().sourceUnit->referencedSourceUnits(true);
+			sourceUnits.insert(referencedSourceUnits.begin(), referencedSourceUnits.end());
+		}
 	}
 	return sourceUnits;
 }
diff --git a/libsolidity/ast/AST.h b/libsolidity/ast/AST.h
index e2117613..2522867b 100644
--- a/libsolidity/ast/AST.h
+++ b/libsolidity/ast/AST.h
@@ -136,8 +136,8 @@ public:
 
 	std::vector<ASTPointer<ASTNode>> nodes() const { return m_nodes; }
 
-	/// @returns a vector of referenced SourceUnits.
-	std::vector<SourceUnit const*> referencedSourceUnits() const;
+	/// @returns a set of referenced SourceUnits. Recursively if @a _recurse is true.
+	std::set<SourceUnit const*> referencedSourceUnits(bool _recurse = false) const;
 
 private:
 	std::vector<ASTPointer<ASTNode>> m_nodes;