2 files changed, 30 insertions, 2 deletions
diff --git a/libsolidity/ast/AST.cpp b/libsolidity/ast/AST.cpp
index 724a908f..ebc8bd48 100644
--- a/libsolidity/ast/AST.cpp
+++ b/libsolidity/ast/AST.cpp
@@ -84,13 +84,35 @@ SourceUnitAnnotation& SourceUnit::annotation() const
 	return dynamic_cast<SourceUnitAnnotation&>(*m_annotation);
 }
 
-string Declaration::sourceUnitName() const
+set<SourceUnit const*> SourceUnit::referencedSourceUnits(bool _recurse, set<SourceUnit const*> _skipList) const
+{
+	set<SourceUnit const*> sourceUnits;
+	for (ImportDirective const* importDirective: filteredNodes<ImportDirective>(nodes()))
+	{
+		auto const& sourceUnit = importDirective->annotation().sourceUnit;
+		if (!_skipList.count(sourceUnit))
+		{
+			_skipList.insert(sourceUnit);
+			sourceUnits.insert(sourceUnit);
+			if (_recurse)
+				sourceUnits += sourceUnit->referencedSourceUnits(true, _skipList);
+		}
+	}
+	return sourceUnits;
+}
+
+SourceUnit const& Declaration::sourceUnit() const
 {
 	solAssert(!!m_scope, "");
 	ASTNode const* scope = m_scope;
 	while (dynamic_cast<Declaration const*>(scope) && dynamic_cast<Declaration const*>(scope)->m_scope)
 		scope = dynamic_cast<Declaration const*>(scope)->m_scope;
-	return dynamic_cast<SourceUnit const&>(*scope).annotation().path;
+	return dynamic_cast<SourceUnit const&>(*scope);
+}
+
+string Declaration::sourceUnitName() const
+{
+	return sourceUnit().annotation().path;
 }
 
 ImportAnnotation& ImportDirective::annotation() const
diff --git a/libsolidity/ast/AST.h b/libsolidity/ast/AST.h
index 81ddc754..8012bcb4 100644
--- a/libsolidity/ast/AST.h
+++ b/libsolidity/ast/AST.h
@@ -136,6 +136,9 @@ public:
 
 	std::vector<ASTPointer<ASTNode>> nodes() const { return m_nodes; }
 
+	/// @returns a set of referenced SourceUnits. Recursively if @a _recurse is true.
+	std::set<SourceUnit const*> referencedSourceUnits(bool _recurse = false, std::set<SourceUnit const*> _skipList = std::set<SourceUnit const*>()) const;
+
 private:
 	std::vector<ASTPointer<ASTNode>> m_nodes;
 };
@@ -168,6 +171,9 @@ public:
 	ASTNode const* scope() const { return m_scope; }
 	void setScope(ASTNode const* _scope) { m_scope = _scope; }
 
+	/// @returns the source unit this declaration is present in.
+	SourceUnit const& sourceUnit() const;
+
 	/// @returns the source name this declaration is present in.
 	/// Can be combined with annotation().canonicalName to form a globally unique name.
 	std::string sourceUnitName() const;