5 files changed, 28 insertions, 13 deletions

diff --git a/src/shared/cplusplus/ASTVisitor.cpp b/src/shared/cplusplus/ASTVisitor.cpp
index 6f331515a0..9281efdf2b 100644
--- a/src/shared/cplusplus/ASTVisitor.cpp
+++ b/src/shared/cplusplus/ASTVisitor.cpp
@@ -53,8 +53,8 @@
 
 using namespace CPlusPlus;
 
-ASTVisitor::ASTVisitor(Control *control)
-    : _control(control)
+ASTVisitor::ASTVisitor(TranslationUnit *translationUnit)
+    : _translationUnit(translationUnit)
 { }
 
 ASTVisitor::~ASTVisitor()
@@ -64,10 +64,18 @@ void ASTVisitor::accept(AST *ast)
 { AST::accept(ast, this); }
 
 Control *ASTVisitor::control() const
-{ return _control; }
+{
+    if (_translationUnit)
+        return _translationUnit->control();
+
+    return 0;
+}
 
 TranslationUnit *ASTVisitor::translationUnit() const
-{ return _control->translationUnit(); }
+{ return _translationUnit; }
+
+void ASTVisitor::setTranslationUnit(TranslationUnit *translationUnit)
+{ _translationUnit = translationUnit; }
 
 unsigned ASTVisitor::tokenCount() const
 { return translationUnit()->tokenCount(); }
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
diff --git a/src/shared/cplusplus/Semantic.cpp b/src/shared/cplusplus/Semantic.cpp
index 8baf10d4b0..23e2fa2c95 100644
--- a/src/shared/cplusplus/Semantic.cpp
+++ b/src/shared/cplusplus/Semantic.cpp
@@ -64,9 +64,10 @@ using namespace CPlusPlus;
 class Semantic::Data
 {
 public:
-    Data(Semantic *semantic, Control *control)
+    Data(Semantic *semantic, TranslationUnit *translationUnit)
         : semantic(semantic),
-          control(control),
+          translationUnit(translationUnit),
+          control(translationUnit->control()),
           skipFunctionBodies(false),
           visibility(Symbol::Public),
           ojbcVisibility(Symbol::Protected),
@@ -90,6 +91,7 @@ public:
     }
 
     Semantic *semantic;
+    TranslationUnit *translationUnit;
     Control *control;
     bool skipFunctionBodies;
     int visibility;
@@ -103,9 +105,9 @@ public:
     CheckName *checkName;
 };
 
-Semantic::Semantic(Control *control)
+Semantic::Semantic(TranslationUnit *translationUnit)
 {
-    d = new Data(this, control);
+    d = new Data(this, translationUnit);
     d->checkSpecifier = new CheckSpecifier(this);
     d->checkDeclaration = new CheckDeclaration(this);
     d->checkDeclarator = new CheckDeclarator(this);
@@ -117,6 +119,9 @@ Semantic::Semantic(Control *control)
 Semantic::~Semantic()
 { delete d; }
 
+TranslationUnit *Semantic::translationUnit() const
+{ return d->translationUnit; }
+
 Control *Semantic::control() const
 { return d->control; }
 
diff --git a/src/shared/cplusplus/Semantic.h b/src/shared/cplusplus/Semantic.h
index 391872d847..4a2ebfc2fb 100644
--- a/src/shared/cplusplus/Semantic.h
+++ b/src/shared/cplusplus/Semantic.h
@@ -61,9 +61,10 @@ class CPLUSPLUS_EXPORT Semantic
     void operator =(const Semantic &other);
 
 public:
-    Semantic(Control *control);
+    Semantic(TranslationUnit *translationUnit);
     virtual ~Semantic();
 
+    TranslationUnit *translationUnit() const;    
     Control *control() const;
 
     FullySpecifiedType check(SpecifierListAST *specifier, Scope *scope);
diff --git a/src/shared/cplusplus/SemanticCheck.cpp b/src/shared/cplusplus/SemanticCheck.cpp
index 1feffd93aa..b28c01636b 100644
--- a/src/shared/cplusplus/SemanticCheck.cpp
+++ b/src/shared/cplusplus/SemanticCheck.cpp
@@ -52,7 +52,7 @@
 using namespace CPlusPlus;
 
 SemanticCheck::SemanticCheck(Semantic *semantic)
-    : ASTVisitor(semantic->control()),
+    : ASTVisitor(semantic->translationUnit()),
       _semantic(semantic)
 { }