Stash a CXXUnit pointer into each cursor. This allows us to simplify

the interface to clang_visitChildren() by eliminating the
CXTranslationUnit pointer.


git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@94051 91177308-0d34-0410-b5e6-96231b3b80d8

5 files changed, 89 insertions, 110 deletions

diff --git a/include/clang-c/Index.h b/include/clang-c/Index.h
index 397e274b8d..110e49904b 100644
--- a/include/clang-c/Index.h
+++ b/include/clang-c/Index.h
@@ -303,7 +303,7 @@ CINDEX_LINKAGE void clang_disposeString(CXString string);
  *   TU = clang_createTranslationUnit(Idx, "IndexTest.pch");
  *
  *   // This will load all the symbols from 'IndexTest.pch'
- *   clang_visitChildren(TU, clang_getTranslationUnitCursor(TU), 
+ *   clang_visitChildren(clang_getTranslationUnitCursor(TU), 
  *                       TranslationUnitVisitor, 0);
  *   clang_disposeTranslationUnit(TU);
  *
@@ -311,8 +311,8 @@ CINDEX_LINKAGE void clang_disposeString(CXString string);
  *   // from 'IndexTest.pch'.
  *   char *args[] = { "-Xclang", "-include-pch=IndexTest.pch", 0 };
  *   TU = clang_createTranslationUnitFromSourceFile(Idx, "IndexTest.c", 2, args);
- *   clang_loadTranslationUnit(TU, clang_getTranslationUnitCursor(TU),
- *                             TranslationUnitVisitor, 0);
+ *   clang_visitChildren(clang_getTranslationUnitCursor(TU),
+ *                       TranslationUnitVisitor, 0);
  *   clang_disposeTranslationUnit(TU);
  *
  * This process of creating the 'pch', loading it separately, and using it (via
@@ -451,8 +451,7 @@ typedef enum CXChildVisitResult (*CXCursorVisitor)(CXCursor cursor,
  * \returns a non-zero value if the traversal was terminated
  * prematurely by the visitor returning \c CXChildVisit_Break.
  */
-CINDEX_LINKAGE unsigned clang_visitChildren(CXTranslationUnit tu,
-                                            CXCursor parent, 
+CINDEX_LINKAGE unsigned clang_visitChildren(CXCursor parent, 
                                             CXCursorVisitor visitor,
                                             CXClientData client_data);
 
diff --git a/tools/CIndex/CIndex.cpp b/tools/CIndex/CIndex.cpp
index 267fcf7336..a1c1230701 100644
--- a/tools/CIndex/CIndex.cpp
+++ b/tools/CIndex/CIndex.cpp
@@ -141,6 +141,7 @@ namespace {
   
 // Cursor visitor.
 class CursorVisitor : public DeclVisitor<CursorVisitor, bool> {
+  ASTUnit *TU;
   CXCursor Parent;
   CXCursorVisitor Visitor;
   CXClientData ClientData;
@@ -153,9 +154,9 @@ class CursorVisitor : public DeclVisitor<CursorVisitor, bool> {
   using DeclVisitor<CursorVisitor, bool>::Visit;
   
 public:
-  CursorVisitor(CXCursorVisitor Visitor, CXClientData ClientData, 
+  CursorVisitor(ASTUnit *TU, CXCursorVisitor Visitor, CXClientData ClientData, 
                 unsigned MaxPCHLevel)
-    : Visitor(Visitor), ClientData(ClientData), MaxPCHLevel(MaxPCHLevel)
+    : TU(TU), Visitor(Visitor), ClientData(ClientData), MaxPCHLevel(MaxPCHLevel)
   {
     Parent.kind = CXCursor_NoDeclFound;
     Parent.data[0] = 0;
@@ -242,12 +243,12 @@ bool CursorVisitor::VisitChildren(CXCursor Cursor) {
   }
   
   if (clang_isTranslationUnit(Cursor.kind)) {
-    ASTUnit *CXXUnit = static_cast<ASTUnit *>(Cursor.data[0]);
+    ASTUnit *CXXUnit = getCursorASTUnit(Cursor);
     if (!CXXUnit->isMainFileAST() && CXXUnit->getOnlyLocalDecls()) {
       const std::vector<Decl*> &TLDs = CXXUnit->getTopLevelDecls();
       for (std::vector<Decl*>::const_iterator it = TLDs.begin(),
            ie = TLDs.end(); it != ie; ++it) {
-        if (Visit(MakeCXCursor(*it)))
+        if (Visit(MakeCXCursor(*it, CXXUnit)))
           return true;
       }
     } else {
@@ -271,7 +272,7 @@ bool CursorVisitor::VisitTranslationUnitDecl(TranslationUnitDecl *D) {
 bool CursorVisitor::VisitDeclContext(DeclContext *DC) {
   for (DeclContext::decl_iterator
        I = DC->decls_begin(), E = DC->decls_end(); I != E; ++I) {
-    if (Visit(MakeCXCursor(*I)))
+    if (Visit(MakeCXCursor(*I, TU)))
       return true;
   }
   
@@ -296,13 +297,14 @@ bool CursorVisitor::VisitFunctionDecl(FunctionDecl *ND) {
 }
 
 bool CursorVisitor::VisitObjCCategoryDecl(ObjCCategoryDecl *ND) {
-  if (Visit(MakeCursorObjCClassRef(ND->getClassInterface(), ND->getLocation())))
+  if (Visit(MakeCursorObjCClassRef(ND->getClassInterface(), ND->getLocation(),
+                                   TU)))
     return true;
   
   ObjCCategoryDecl::protocol_loc_iterator PL = ND->protocol_loc_begin();
   for (ObjCCategoryDecl::protocol_iterator I = ND->protocol_begin(),
          E = ND->protocol_end(); I != E; ++I, ++PL)
-    if (Visit(MakeCursorObjCProtocolRef(*I, *PL)))
+    if (Visit(MakeCursorObjCProtocolRef(*I, *PL, TU)))
       return true;
   
   return VisitDeclContext(ND);
@@ -312,13 +314,14 @@ bool CursorVisitor::VisitObjCInterfaceDecl(ObjCInterfaceDecl *D) {
   // Issue callbacks for super class.
   if (D->getSuperClass() &&
       Visit(MakeCursorObjCSuperClassRef(D->getSuperClass(),
-                                        D->getSuperClassLoc())))
+                                        D->getSuperClassLoc(), 
+                                        TU)))
     return true;
   
   ObjCInterfaceDecl::protocol_loc_iterator PL = D->protocol_loc_begin();
   for (ObjCInterfaceDecl::protocol_iterator I = D->protocol_begin(),
          E = D->protocol_end(); I != E; ++I, ++PL)
-    if (Visit(MakeCursorObjCProtocolRef(*I, *PL)))
+    if (Visit(MakeCursorObjCProtocolRef(*I, *PL, TU)))
       return true;
   
   return VisitDeclContext(D);
@@ -336,7 +339,7 @@ bool CursorVisitor::VisitObjCProtocolDecl(ObjCProtocolDecl *PID) {
   ObjCProtocolDecl::protocol_loc_iterator PL = PID->protocol_loc_begin();
   for (ObjCProtocolDecl::protocol_iterator I = PID->protocol_begin(),
          E = PID->protocol_end(); I != E; ++I, ++PL)
-    if (Visit(MakeCursorObjCProtocolRef(*I, *PL)))
+    if (Visit(MakeCursorObjCProtocolRef(*I, *PL, TU)))
       return true;
   
   return VisitDeclContext(PID);
@@ -513,7 +516,7 @@ CXString clang_getTranslationUnitSpelling(CXTranslationUnit CTUnit) {
 }
 
 CXCursor clang_getTranslationUnitCursor(CXTranslationUnit TU) {
-  CXCursor Result = { CXCursor_TranslationUnit, { TU, 0, 0 } };
+  CXCursor Result = { CXCursor_TranslationUnit, { 0, 0, TU } };
   return Result;
 }
 
@@ -641,11 +644,10 @@ static Decl *getDeclFromExpr(Stmt *E) {
 
 extern "C" {
   
-unsigned clang_visitChildren(CXTranslationUnit tu,
-                             CXCursor parent, 
+unsigned clang_visitChildren(CXCursor parent, 
                              CXCursorVisitor visitor,
                              CXClientData client_data) {
-  ASTUnit *CXXUnit = static_cast<ASTUnit *>(tu);
+  ASTUnit *CXXUnit = getCursorASTUnit(parent);
 
   unsigned PCHLevel = Decl::MaxPCHLevel;
   
@@ -658,7 +660,7 @@ unsigned clang_visitChildren(CXTranslationUnit tu,
       ++PCHLevel;
   }
   
-  CursorVisitor CursorVis(visitor, client_data, PCHLevel);
+  CursorVisitor CursorVis(CXXUnit, visitor, client_data, PCHLevel);
   return CursorVis.VisitChildren(parent);
 }
 
@@ -686,7 +688,7 @@ static CXString getDeclSpelling(Decl *D) {
 CXString clang_getCursorSpelling(CXCursor C) {
   assert(getCursorDecl(C) && "CXCursor has null decl");
   if (clang_isTranslationUnit(C.kind))
-    return clang_getTranslationUnitSpelling(C.data[0]);
+    return clang_getTranslationUnitSpelling(C.data[2]);
 
   if (clang_isReference(C.kind)) {
     switch (C.kind) {
@@ -787,15 +789,15 @@ CXCursor clang_getCursor(CXTranslationUnit CTUnit, const char *source_name,
   Stmt *Stm = ALoc.dyn_AsStmt();
   if (Dcl) {
     if (Stm)
-      return MakeCXCursor(Stm, Dcl);
+      return MakeCXCursor(Stm, Dcl, CXXUnit);
 
     if (ALoc.isNamedRef()) {
       if (ObjCInterfaceDecl *Class = dyn_cast<ObjCInterfaceDecl>(Dcl))
-        return MakeCursorObjCClassRef(Class, ALoc.AsNamedRef().Loc);
+        return MakeCursorObjCClassRef(Class, ALoc.AsNamedRef().Loc, CXXUnit);
       if (ObjCProtocolDecl *Proto = dyn_cast<ObjCProtocolDecl>(Dcl))
-        return MakeCursorObjCProtocolRef(Proto, ALoc.AsNamedRef().Loc);
+        return MakeCursorObjCProtocolRef(Proto, ALoc.AsNamedRef().Loc, CXXUnit);
     }
-    return MakeCXCursor(Dcl);
+    return MakeCXCursor(Dcl, CXXUnit);
   }
   return MakeCXCursorInvalid(CXCursor_NoDeclFound);
 }
@@ -933,13 +935,17 @@ CXSourceRange clang_getCursorExtent(CXCursor C) {
 }
 
 CXCursor clang_getCursorReferenced(CXCursor C) {
+  if (clang_isInvalid(C.kind))
+    return clang_getNullCursor();
+  
+  ASTUnit *CXXUnit = getCursorASTUnit(C);
   if (clang_isDeclaration(C.kind))
     return C;
   
   if (clang_isExpression(C.kind)) {
     Decl *D = getDeclFromExpr(getCursorExpr(C));
     if (D)
-      return MakeCXCursor(D);
+      return MakeCXCursor(D, CXXUnit);
     return clang_getNullCursor();
   }
 
@@ -948,13 +954,13 @@ CXCursor clang_getCursorReferenced(CXCursor C) {
   
   switch (C.kind) {
     case CXCursor_ObjCSuperClassRef:
-      return MakeCXCursor(getCursorObjCSuperClassRef(C).first);
+      return MakeCXCursor(getCursorObjCSuperClassRef(C).first, CXXUnit);
       
     case CXCursor_ObjCProtocolRef: {       
-      return MakeCXCursor(getCursorObjCProtocolRef(C).first);
+      return MakeCXCursor(getCursorObjCProtocolRef(C).first, CXXUnit);
       
     case CXCursor_ObjCClassRef:      
-      return MakeCXCursor(getCursorObjCClassRef(C).first);
+      return MakeCXCursor(getCursorObjCClassRef(C).first, CXXUnit);
       
     default:
       // We would prefer to enumerate all non-reference cursor kinds here.
@@ -967,6 +973,11 @@ CXCursor clang_getCursorReferenced(CXCursor C) {
 }
 
 CXCursor clang_getCursorDefinition(CXCursor C) {
+  if (clang_isInvalid(C.kind))
+    return clang_getNullCursor();
+  
+  ASTUnit *CXXUnit = getCursorASTUnit(C);
+  
   bool WasReference = false;
   if (clang_isReference(C.kind) || clang_isExpression(C.kind)) {
     C = clang_getCursorReferenced(C);
@@ -1016,10 +1027,11 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
     break;
 
   case Decl::UsingDirective:
-    return MakeCXCursor(cast<UsingDirectiveDecl>(D)->getNominatedNamespace());
+    return MakeCXCursor(cast<UsingDirectiveDecl>(D)->getNominatedNamespace(),
+                        CXXUnit);
 
   case Decl::NamespaceAlias:
-    return MakeCXCursor(cast<NamespaceAliasDecl>(D)->getNamespace());
+    return MakeCXCursor(cast<NamespaceAliasDecl>(D)->getNamespace(), CXXUnit);
 
   case Decl::Enum:
   case Decl::Record:
@@ -1027,7 +1039,7 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
   case Decl::ClassTemplateSpecialization:
   case Decl::ClassTemplatePartialSpecialization:
     if (TagDecl *Def = cast<TagDecl>(D)->getDefinition(D->getASTContext()))
-      return MakeCXCursor(Def);
+      return MakeCXCursor(Def, CXXUnit);
     return clang_getNullCursor();
 
   case Decl::Function:
@@ -1037,7 +1049,7 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
   case Decl::CXXConversion: {
     const FunctionDecl *Def = 0;
     if (cast<FunctionDecl>(D)->getBody(Def))
-      return MakeCXCursor(const_cast<FunctionDecl *>(Def));
+      return MakeCXCursor(const_cast<FunctionDecl *>(Def), CXXUnit);
     return clang_getNullCursor();
   }
 
@@ -1047,7 +1059,7 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
     // Variables with initializers have definitions.
     const VarDecl *Def = 0;
     if (Var->getDefinition(Def))
-      return MakeCXCursor(const_cast<VarDecl *>(Def));
+      return MakeCXCursor(const_cast<VarDecl *>(Def), CXXUnit);
 
     // extern and private_extern variables are not definitions.
     if (Var->hasExternalStorage())
@@ -1064,7 +1076,7 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
   case Decl::FunctionTemplate: {
     const FunctionDecl *Def = 0;
     if (cast<FunctionTemplateDecl>(D)->getTemplatedDecl()->getBody(Def))
-      return MakeCXCursor(Def->getDescribedFunctionTemplate());
+      return MakeCXCursor(Def->getDescribedFunctionTemplate(), CXXUnit);
     return clang_getNullCursor();
   }
    
@@ -1072,7 +1084,8 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
     if (RecordDecl *Def = cast<ClassTemplateDecl>(D)->getTemplatedDecl()
                                           ->getDefinition(D->getASTContext()))
       return MakeCXCursor(
-                         cast<CXXRecordDecl>(Def)->getDescribedClassTemplate());
+                         cast<CXXRecordDecl>(Def)->getDescribedClassTemplate(), 
+                          CXXUnit);
     return clang_getNullCursor();
   }
 
@@ -1087,7 +1100,8 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
         return clang_getNullCursor();
       }
 
-      Def = clang_getCursorDefinition(MakeCXCursor((*S)->getTargetDecl()));
+      Def = clang_getCursorDefinition(MakeCXCursor((*S)->getTargetDecl(), 
+                                                   CXXUnit));
     }
 
     return Def;
@@ -1095,7 +1109,8 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
 
   case Decl::UsingShadow:
     return clang_getCursorDefinition(
-                       MakeCXCursor(cast<UsingShadowDecl>(D)->getTargetDecl()));
+                       MakeCXCursor(cast<UsingShadowDecl>(D)->getTargetDecl(), 
+                                    CXXUnit));
 
   case Decl::ObjCMethod: {
     ObjCMethodDecl *Method = cast<ObjCMethodDecl>(D);
@@ -1111,7 +1126,7 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
         if (ObjCMethodDecl *Def = ClassImpl->getMethod(Method->getSelector(),
                                                   Method->isInstanceMethod()))
           if (Def->isThisDeclarationADefinition())
-            return MakeCXCursor(Def);
+            return MakeCXCursor(Def, CXXUnit);
 
     return clang_getNullCursor();
   }
@@ -1119,7 +1134,7 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
   case Decl::ObjCCategory:
     if (ObjCCategoryImplDecl *Impl
                                = cast<ObjCCategoryDecl>(D)->getImplementation())
-      return MakeCXCursor(Impl);
+      return MakeCXCursor(Impl, CXXUnit);
     return clang_getNullCursor();
 
   case Decl::ObjCProtocol:
@@ -1138,7 +1153,7 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
         return C;
     } else if (ObjCImplementationDecl *Impl
                               = cast<ObjCInterfaceDecl>(D)->getImplementation())
-      return MakeCXCursor(Impl);
+      return MakeCXCursor(Impl, CXXUnit);
     return clang_getNullCursor();
   
   case Decl::ObjCProperty:
@@ -1150,7 +1165,7 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
     if (ObjCInterfaceDecl *Class
           = cast<ObjCCompatibleAliasDecl>(D)->getClassInterface())
       if (!Class->isForwardDecl())
-        return MakeCXCursor(Class);
+        return MakeCXCursor(Class, CXXUnit);
     
     return clang_getNullCursor();
 
@@ -1158,7 +1173,8 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
     ObjCForwardProtocolDecl *Forward = cast<ObjCForwardProtocolDecl>(D);
     if (Forward->protocol_size() == 1)
       return clang_getCursorDefinition(
-                                     MakeCXCursor(*Forward->protocol_begin()));
+                                     MakeCXCursor(*Forward->protocol_begin(), 
+                                                  CXXUnit));
 
     // FIXME: Cannot return multiple definitions.
     return clang_getNullCursor();
@@ -1169,7 +1185,7 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
     if (Class->size() == 1) {
       ObjCInterfaceDecl *IFace = Class->begin()->getInterface();
       if (!IFace->isForwardDecl())
-        return MakeCXCursor(IFace);
+        return MakeCXCursor(IFace, CXXUnit);
       return clang_getNullCursor();
     }
 
@@ -1179,12 +1195,12 @@ CXCursor clang_getCursorDefinition(CXCursor C) {
 
   case Decl::Friend:
     if (NamedDecl *Friend = cast<FriendDecl>(D)->getFriendDecl())
-      return clang_getCursorDefinition(MakeCXCursor(Friend));
+      return clang_getCursorDefinition(MakeCXCursor(Friend, CXXUnit));
     return clang_getNullCursor();
 
   case Decl::FriendTemplate:
     if (NamedDecl *Friend = cast<FriendTemplateDecl>(D)->getFriendDecl())
-      return clang_getCursorDefinition(MakeCXCursor(Friend));
+      return clang_getCursorDefinition(MakeCXCursor(Friend, CXXUnit));
     return clang_getNullCursor();
   }
 
diff --git a/tools/CIndex/CXCursor.cpp b/tools/CIndex/CXCursor.cpp
index 7c8b33d1ae..29bd4e2bbc 100644
--- a/tools/CIndex/CXCursor.cpp
+++ b/tools/CIndex/CXCursor.cpp
@@ -71,12 +71,12 @@ static CXCursorKind GetCursorKind(Decl *D) {
   return CXCursor_NotImplemented;  
 }
 
-CXCursor cxcursor::MakeCXCursor(Decl *D) {
-  CXCursor C = { GetCursorKind(D), { D, 0, 0 } };
+CXCursor cxcursor::MakeCXCursor(Decl *D, ASTUnit *TU) {
+  CXCursor C = { GetCursorKind(D), { D, 0, TU } };
   return C;
 }
 
-CXCursor cxcursor::MakeCXCursor(Stmt *S, Decl *Parent) {
+CXCursor cxcursor::MakeCXCursor(Stmt *S, Decl *Parent, ASTUnit *TU) {
   CXCursorKind K = CXCursor_NotImplemented;
   
   switch (S->getStmtClass()) {
@@ -207,14 +207,15 @@ CXCursor cxcursor::MakeCXCursor(Stmt *S, Decl *Parent) {
     break;
   }
   
-  CXCursor C = { K, { Parent, S, 0 } };
+  CXCursor C = { K, { Parent, S, TU } };
   return C;
 }
 
 CXCursor cxcursor::MakeCursorObjCSuperClassRef(ObjCInterfaceDecl *Super, 
-                                         SourceLocation Loc) {
+                                               SourceLocation Loc, 
+                                               ASTUnit *TU) {
   void *RawLoc = reinterpret_cast<void *>(Loc.getRawEncoding());
-  CXCursor C = { CXCursor_ObjCSuperClassRef, { Super, RawLoc, 0 } };
+  CXCursor C = { CXCursor_ObjCSuperClassRef, { Super, RawLoc, TU } };
   return C;    
 }
 
@@ -227,9 +228,10 @@ cxcursor::getCursorObjCSuperClassRef(CXCursor C) {
 }
 
 CXCursor cxcursor::MakeCursorObjCProtocolRef(ObjCProtocolDecl *Super, 
-                                             SourceLocation Loc) {
+                                             SourceLocation Loc, 
+                                             ASTUnit *TU) {
   void *RawLoc = reinterpret_cast<void *>(Loc.getRawEncoding());
-  CXCursor C = { CXCursor_ObjCProtocolRef, { Super, RawLoc, 0 } };
+  CXCursor C = { CXCursor_ObjCProtocolRef, { Super, RawLoc, TU } };
   return C;    
 }
 
@@ -242,9 +244,10 @@ cxcursor::getCursorObjCProtocolRef(CXCursor C) {
 }
 
 CXCursor cxcursor::MakeCursorObjCClassRef(ObjCInterfaceDecl *Class, 
-                                         SourceLocation Loc) {
+                                          SourceLocation Loc, 
+                                          ASTUnit *TU) {
   void *RawLoc = reinterpret_cast<void *>(Loc.getRawEncoding());
-  CXCursor C = { CXCursor_ObjCClassRef, { Class, RawLoc, 0 } };
+  CXCursor C = { CXCursor_ObjCClassRef, { Class, RawLoc, TU } };
   return C;    
 }
 
@@ -274,55 +277,11 @@ Stmt *cxcursor::getCursorStmt(CXCursor Cursor) {
 }
 
 ASTContext &cxcursor::getCursorContext(CXCursor Cursor) {
-  switch (Cursor.kind) {
-  case CXCursor_TypedefDecl:
-  case CXCursor_StructDecl:
-  case CXCursor_UnionDecl:
-  case CXCursor_ClassDecl:
-  case CXCursor_EnumDecl:
-  case CXCursor_FieldDecl:
-  case CXCursor_EnumConstantDecl:
-  case CXCursor_FunctionDecl:
-  case CXCursor_VarDecl:
-  case CXCursor_ParmDecl:
-  case CXCursor_ObjCInterfaceDecl:
-  case CXCursor_ObjCCategoryDecl:
-  case CXCursor_ObjCProtocolDecl:
-  case CXCursor_ObjCPropertyDecl:
-  case CXCursor_ObjCIvarDecl:
-  case CXCursor_ObjCInstanceMethodDecl:
-  case CXCursor_ObjCClassMethodDecl:
-  case CXCursor_ObjCImplementationDecl:
-  case CXCursor_ObjCCategoryImplDecl:
-  case CXCursor_UnexposedDecl:
-    return static_cast<Decl *>(Cursor.data[0])->getASTContext();
-
-  case CXCursor_ObjCSuperClassRef:
-  case CXCursor_ObjCProtocolRef:
-  case CXCursor_ObjCClassRef:
-    return static_cast<Decl *>(Cursor.data[0])->getASTContext();
-    
-  case CXCursor_InvalidFile:
-  case CXCursor_NoDeclFound:
-  case CXCursor_NotImplemented:
-    llvm_unreachable("No context in an invalid cursor");
-    break;
-
-  case CXCursor_UnexposedExpr:
-  case CXCursor_DeclRefExpr:
-  case CXCursor_MemberRefExpr:
-  case CXCursor_CallExpr:
-  case CXCursor_ObjCMessageExpr:
-  case CXCursor_UnexposedStmt:
-    return static_cast<Decl *>(Cursor.data[0])->getASTContext();
+  return getCursorASTUnit(Cursor)->getASTContext();
+}
 
-  case CXCursor_TranslationUnit: {
-    ASTUnit *CXXUnit = static_cast<ASTUnit *>(Cursor.data[0]);
-    return CXXUnit->getASTContext();
-  }
-  }
-  
-  llvm_unreachable("No context available");
+ASTUnit *cxcursor::getCursorASTUnit(CXCursor Cursor) {
+  return static_cast<ASTUnit *>(Cursor.data[2]);
 }
 
 bool cxcursor::operator==(CXCursor X, CXCursor Y) {
diff --git a/tools/CIndex/CXCursor.h b/tools/CIndex/CXCursor.h
index 9c438d9e69..5f81c8a173 100644
--- a/tools/CIndex/CXCursor.h
+++ b/tools/CIndex/CXCursor.h
@@ -21,6 +21,7 @@
 namespace clang {
 
 class ASTContext;
+class ASTUnit;
 class Decl;
 class Expr;
 class NamedDecl;
@@ -31,12 +32,13 @@ class Stmt;
 namespace cxcursor {
   
 CXCursor MakeCXCursorInvalid(CXCursorKind K);
-CXCursor MakeCXCursor(clang::Stmt *S, clang::Decl *Parent);
-CXCursor MakeCXCursor(clang::Decl *D);
+CXCursor MakeCXCursor(clang::Stmt *S, clang::Decl *Parent, ASTUnit *TU);
+CXCursor MakeCXCursor(clang::Decl *D, ASTUnit *TU);
 
 /// \brief Create an Objective-C superclass reference at the given location.
 CXCursor MakeCursorObjCSuperClassRef(ObjCInterfaceDecl *Super, 
-                                     SourceLocation Loc);
+                                     SourceLocation Loc, 
+                                     ASTUnit *TU);
 
 /// \brief Unpack an ObjCSuperClassRef cursor into the interface it references
 /// and optionally the location where the reference occurred.
@@ -44,7 +46,8 @@ std::pair<ObjCInterfaceDecl *, SourceLocation>
   getCursorObjCSuperClassRef(CXCursor C);
 
 /// \brief Create an Objective-C protocol reference at the given location.
-CXCursor MakeCursorObjCProtocolRef(ObjCProtocolDecl *Proto, SourceLocation Loc);
+CXCursor MakeCursorObjCProtocolRef(ObjCProtocolDecl *Proto, SourceLocation Loc, 
+                                   ASTUnit *TU);
 
 /// \brief Unpack an ObjCProtocolRef cursor into the protocol it references
 /// and optionally the location where the reference occurred.
@@ -52,7 +55,8 @@ std::pair<ObjCProtocolDecl *, SourceLocation>
   getCursorObjCProtocolRef(CXCursor C);
 
 /// \brief Create an Objective-C class reference at the given location.
-CXCursor MakeCursorObjCClassRef(ObjCInterfaceDecl *Class, SourceLocation Loc);
+CXCursor MakeCursorObjCClassRef(ObjCInterfaceDecl *Class, SourceLocation Loc, 
+                                ASTUnit *TU);
 
 /// \brief Unpack an ObjCClassRef cursor into the class it references
 /// and optionally the location where the reference occurred.
@@ -63,6 +67,7 @@ Decl *getCursorDecl(CXCursor Cursor);
 Expr *getCursorExpr(CXCursor Cursor);
 Stmt *getCursorStmt(CXCursor Cursor);
 ASTContext &getCursorContext(CXCursor Cursor);
+ASTUnit *getCursorASTUnit(CXCursor Cursor);
   
 bool operator==(CXCursor X, CXCursor Y);
   
diff --git a/tools/c-index-test/c-index-test.c b/tools/c-index-test/c-index-test.c
index 185c059737..473773e080 100644
--- a/tools/c-index-test/c-index-test.c
+++ b/tools/c-index-test/c-index-test.c
@@ -230,7 +230,7 @@ static int perform_test_load(CXIndex Idx, CXTranslationUnit TU,
   
   Data.TU = TU;
   Data.Filter = ck;
-  clang_visitChildren(TU, clang_getTranslationUnitCursor(TU), Visitor, &Data);
+  clang_visitChildren(clang_getTranslationUnitCursor(TU), Visitor, &Data);
   clang_disposeTranslationUnit(TU);
   return 0;
 }