From 89a31730fc40bbc3849ab12c8205b45ea63138b1 Mon Sep 17 00:00:00 2001
From: Joao Matos <joao@tritao.eu>
Date: Tue, 21 Feb 2017 21:29:44 +0000
Subject: [PATCH] Added a visitor overload for translation units.

---
 src/AST/ASTVisitor.cs                     | 8 ++++++++
 src/AST/CppTypePrinter.cs                 | 5 +++++
 src/AST/Declaration.cs                    | 1 +
 src/AST/TranslationUnit.cs                | 7 ++++++-
 src/Generator.Tests/AST/TestAST.cs        | 5 +++++
 src/Generator/Generators/CodeGenerator.cs | 5 +++++
 src/Generator/Generators/TypePrinter.cs   | 5 +++++
 src/Generator/Passes/Pass.cs              | 2 +-
 8 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/src/AST/ASTVisitor.cs b/src/AST/ASTVisitor.cs
index 337c51e2..e463e5c9 100644
--- a/src/AST/ASTVisitor.cs
+++ b/src/AST/ASTVisitor.cs
@@ -302,6 +302,14 @@ namespace CppSharp.AST
             return !AlreadyVisited(decl);
         }
 
+        public virtual bool VisitTranslationUnit(TranslationUnit unit)
+        {
+            if (!VisitDeclaration(unit))
+                return false;
+
+            return VisitDeclarationContext(unit);
+        }
+
         public virtual bool VisitClassDecl(Class @class)
         {
             if (!VisitDeclarationContext(@class))
diff --git a/src/AST/CppTypePrinter.cs b/src/AST/CppTypePrinter.cs
index 087ec0fe..b9bf3bf8 100644
--- a/src/AST/CppTypePrinter.cs
+++ b/src/AST/CppTypePrinter.cs
@@ -316,6 +316,11 @@ namespace CppSharp.AST
             return GetDeclName(decl, PrintScopeKind);
         }
 
+        public string VisitTranslationUnit(TranslationUnit unit)
+        {
+            return VisitDeclaration(unit);
+        }
+
         public virtual string VisitClassDecl(Class @class)
         {
             return VisitDeclaration(@class);
diff --git a/src/AST/Declaration.cs b/src/AST/Declaration.cs
index 3d5b160b..d54aab5a 100644
--- a/src/AST/Declaration.cs
+++ b/src/AST/Declaration.cs
@@ -367,6 +367,7 @@ namespace CppSharp.AST
     public interface IDeclVisitor<out T>
     {
         T VisitDeclaration(Declaration decl);
+        T VisitTranslationUnit(TranslationUnit unit);
         T VisitClassDecl(Class @class);
         T VisitFieldDecl(Field field);
         T VisitFunctionDecl(Function function);
diff --git a/src/AST/TranslationUnit.cs b/src/AST/TranslationUnit.cs
index 31a760f9..992d6db2 100644
--- a/src/AST/TranslationUnit.cs
+++ b/src/AST/TranslationUnit.cs
@@ -6,7 +6,7 @@ using System.IO;
 namespace CppSharp.AST
 {
     /// <summary>
-    /// Represents a parsed C++ unit.
+    /// Represents a source code unit.
     /// </summary>
     [DebuggerDisplay("File = {FileName}, Ignored = {Ignore}")]
     public class TranslationUnit : Namespace
@@ -84,5 +84,10 @@ namespace CppSharp.AST
         }
 
         private WeakReference module;
+
+        public override T Visit<T>(IDeclVisitor<T> visitor)
+        {
+            return visitor.VisitTranslationUnit(this);
+        }
     }
 }
\ No newline at end of file
diff --git a/src/Generator.Tests/AST/TestAST.cs b/src/Generator.Tests/AST/TestAST.cs
index 74d30429..7a98ccd2 100644
--- a/src/Generator.Tests/AST/TestAST.cs
+++ b/src/Generator.Tests/AST/TestAST.cs
@@ -82,6 +82,11 @@ namespace CppSharp.Generator.Tests.AST
                 throw new System.NotImplementedException();
             }
 
+            public bool VisitTranslationUnit(TranslationUnit unit)
+            {
+                throw new System.NotImplementedException();
+            }
+
             public bool VisitClassDecl(Class @class)
             {
                 throw new System.NotImplementedException();
diff --git a/src/Generator/Generators/CodeGenerator.cs b/src/Generator/Generators/CodeGenerator.cs
index 606b6fa9..49672ad6 100644
--- a/src/Generator/Generators/CodeGenerator.cs
+++ b/src/Generator/Generators/CodeGenerator.cs
@@ -56,6 +56,11 @@ namespace CppSharp.Generators
             throw new NotImplementedException();
         }
 
+        public virtual bool VisitTranslationUnit(TranslationUnit unit)
+        {
+            return VisitDeclContext(unit);
+        }
+
         public virtual bool VisitDeclContext(DeclarationContext context)
         {
             foreach (var decl in context.Declarations)
diff --git a/src/Generator/Generators/TypePrinter.cs b/src/Generator/Generators/TypePrinter.cs
index e91cb39b..17457573 100644
--- a/src/Generator/Generators/TypePrinter.cs
+++ b/src/Generator/Generators/TypePrinter.cs
@@ -258,6 +258,11 @@ namespace CppSharp.Generators
             throw new NotImplementedException();
         }
 
+        public virtual TypePrinterResult VisitTranslationUnit(TranslationUnit unit)
+        {
+            throw new NotImplementedException();
+        }
+
         public virtual TypePrinterResult VisitTypeAliasDecl(TypeAlias typeAlias)
         {
             throw new NotImplementedException();
diff --git a/src/Generator/Passes/Pass.cs b/src/Generator/Passes/Pass.cs
index 3f2e185e..c2872f10 100644
--- a/src/Generator/Passes/Pass.cs
+++ b/src/Generator/Passes/Pass.cs
@@ -26,7 +26,7 @@ namespace CppSharp.Passes
             return true;
         }
 
-        public virtual bool VisitTranslationUnit(TranslationUnit unit)
+        public override bool VisitTranslationUnit(TranslationUnit unit)
         {
             if (!unit.IsValid || unit.Ignore)
                 return false;