1 files changed, 4 insertions, 3 deletions
diff --git a/src/shared/cplusplus/ASTVisitor.h b/src/shared/cplusplus/ASTVisitor.h
index e1429fa2f6..0a42b6f4d4 100644
--- a/src/shared/cplusplus/ASTVisitor.h
+++ b/src/shared/cplusplus/ASTVisitor.h
@@ -61,12 +61,13 @@ class CPLUSPLUS_EXPORT ASTVisitor
     void operator =(const ASTVisitor &other);
 
 public:
-    ASTVisitor(Control *control);
+    ASTVisitor(TranslationUnit *unit);
     virtual ~ASTVisitor();
 
-    Control *control() const;
     TranslationUnit *translationUnit() const;
+    void setTranslationUnit(TranslationUnit *translationUnit);
 
+    Control *control() const;
     unsigned tokenCount() const;
     const Token &tokenAt(unsigned index) const;
     int tokenKind(unsigned index) const;
@@ -352,7 +353,7 @@ public:
     virtual void endVisit(ObjCSynchronizedStatementAST *) { }
 
 private:
-    Control *_control;
+    TranslationUnit *_translationUnit;
 };
 
 } // end of namespace CPlusPlus