From: Argyrios Kyrtzidis <akyrtzi@gmail.com>
Date: Tue, 21 Jul 2009 00:07:06 +0000 (+0000)
Subject: Change the semantics for Entity.
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=f7cf15ca3c9bee7c0348f549e7a8f0af32b5fa54;p=clang

Change the semantics for Entity.

Entity can now refer to declarations that are not visible outside the translation unit.
It is a wrapper of a pointer union, it's either a Decl* for declarations that don't
"cross" translation units, or an EntityImpl* which is associated with the specific "visible" Decl.

Included is a test case for handling fields across translation units.

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

diff --git a/include/clang/Analysis/CallGraph.h b/include/clang/Analysis/CallGraph.h
index e64c2ca8c7..73072c410e 100644
--- a/include/clang/Analysis/CallGraph.h
+++ b/include/clang/Analysis/CallGraph.h
@@ -24,12 +24,12 @@
 namespace clang {
 
 class CallGraphNode {
-  idx::Entity *F;
+  idx::Entity F;
   typedef std::pair<idx::ASTLocation, CallGraphNode*> CallRecord;
   std::vector<CallRecord> CalledFunctions;
 
 public:
-  CallGraphNode(idx::Entity *f) : F(f) {}
+  CallGraphNode(idx::Entity f) : F(f) {}
 
   typedef std::vector<CallRecord>::iterator iterator;
   typedef std::vector<CallRecord>::const_iterator const_iterator;
@@ -45,14 +45,14 @@ public:
 
   bool hasCallee() const { return begin() != end(); }
 
-  std::string getName(ASTContext &Ctx) { return F->getPrintableName(Ctx); }
+  std::string getName(ASTContext &Ctx) { return F.getPrintableName(Ctx); }
 };
 
 class CallGraph {
   /// Program manages all Entities.
   idx::Program Prog;
 
-  typedef llvm::DenseMap<idx::Entity *, CallGraphNode *> FunctionMapTy;
+  typedef llvm::DenseMap<idx::Entity, CallGraphNode *> FunctionMapTy;
 
   /// FunctionMap owns all CallGraphNodes.
   FunctionMapTy FunctionMap;
@@ -75,7 +75,7 @@ public:
 
   idx::Program &getProgram() { return Prog; }
 
-  CallGraphNode *getOrInsertFunction(idx::Entity *F);
+  CallGraphNode *getOrInsertFunction(idx::Entity F);
 
   void print(llvm::raw_ostream &os);
   void dump();
diff --git a/include/clang/Index/Entity.h b/include/clang/Index/Entity.h
index e1671bbdce..9a5b126d11 100644
--- a/include/clang/Index/Entity.h
+++ b/include/clang/Index/Entity.h
@@ -15,7 +15,9 @@
 #ifndef LLVM_CLANG_INDEX_ENTITY_H
 #define LLVM_CLANG_INDEX_ENTITY_H
 
-#include "llvm/ADT/FoldingSet.h"
+#include "llvm/ADT/PointerUnion.h"
+#include "llvm/ADT/DenseMap.h"
+#include <string>
 
 namespace clang {
   class ASTContext;
@@ -23,23 +25,35 @@ namespace clang {
 
 namespace idx {
   class Program;
+  class EntityImpl;
 
-/// \brief A ASTContext-independent way to refer to declarations that are
-/// visible across translation units.
+/// \brief A ASTContext-independent way to refer to declarations.
 ///
 /// Entity is basically the link for declarations that are semantically the same
 /// in multiple ASTContexts. A client will convert a Decl into an Entity and
 /// later use that Entity to find the "same" Decl into another ASTContext.
+/// Declarations that are semantically the same and visible across translation
+/// units will be associated with the same Entity.
 ///
-/// An Entity may only refer to declarations that can be visible by multiple
-/// translation units, e.g. a static function cannot have an Entity associated
-/// with it.
+/// An Entity may also refer to declarations that cannot be visible across
+/// translation units, e.g. static functions with the same name in multiple
+/// translation units will be associated with different Entities.
 ///
-/// Entities are uniqued so pointer equality can be used (note that the same
-/// Program object should be used when getting Entities).
+/// Entities can be checked for equality but note that the same Program object
+/// should be used when getting Entities.
 ///
-class Entity : public llvm::FoldingSetNode {
+class Entity {
+  /// \brief Stores the Decl directly if it is not visible outside of its own
+  /// translation unit, otherwise it stores the associated EntityImpl.
+  llvm::PointerUnion<Decl *, EntityImpl *> Val;
+
+  explicit Entity(Decl *D) : Val(D) { }
+  explicit Entity(EntityImpl *impl) : Val(impl) { }
+  friend class EntityGetter;
+  
 public:
+  Entity() { }
+
   /// \brief Find the Decl that can be referred to by this entity.
   Decl *getDecl(ASTContext &AST);
 
@@ -48,26 +62,75 @@ public:
 
   /// \brief Get an Entity associated with the given Decl.
   /// \returns Null if an Entity cannot refer to this Decl.
-  static Entity *get(Decl *D, Program &Prog);
+  static Entity get(Decl *D, Program &Prog);
 
-  void Profile(llvm::FoldingSetNodeID &ID) const {
-    Profile(ID, Parent, Id);
-  }
-  static void Profile(llvm::FoldingSetNodeID &ID, Entity *Parent, void *Id) {
-    ID.AddPointer(Parent);
-    ID.AddPointer(Id);
+  /// \brief true if the Entity is not visible outside the trasnlation unit.
+  bool isInternalToTU() const {
+    assert(isValid() && "This Entity is not valid!");
+    return Val.is<Decl *>();
   }
+
+  bool isValid() const { return !Val.isNull(); }
+  bool isInvalid() const { return !isValid(); }
   
-private:
-  Entity *Parent;
-  void *Id;
+  void *getAsOpaquePtr() const { return Val.getOpaqueValue(); }
+  static Entity getFromOpaquePtr(void *Ptr) {
+    Entity Ent;
+    Ent.Val = llvm::PointerUnion<Decl *, EntityImpl *>::getFromOpaqueValue(Ptr);
+    return Ent;
+  }
+
+  friend bool operator==(const Entity &LHS, const Entity &RHS) { 
+    return LHS.getAsOpaquePtr() == RHS.getAsOpaquePtr();
+  }
   
-  Entity(Entity *parent, void *id) : Parent(parent), Id(id) { }
-  friend class EntityGetter;
+  // For use in a std::map.
+  friend bool operator < (const Entity &LHS, const Entity &RHS) { 
+    return LHS.getAsOpaquePtr() < RHS.getAsOpaquePtr();
+  }
+
+  // For use in DenseMap/DenseSet.
+  static Entity getEmptyMarker() {
+    Entity Ent;
+    Ent.Val =
+      llvm::PointerUnion<Decl *, EntityImpl *>::getFromOpaqueValue((void*)-1);
+    return Ent;
+  }
+  static Entity getTombstoneMarker() {
+    Entity Ent;
+    Ent.Val =
+      llvm::PointerUnion<Decl *, EntityImpl *>::getFromOpaqueValue((void*)-2);
+    return Ent;
+  }
 };
   
 } // namespace idx
 
 } // namespace clang
 
+namespace llvm {
+/// Define DenseMapInfo so that Entities can be used as keys in DenseMap and
+/// DenseSets.
+template<>
+struct DenseMapInfo<clang::idx::Entity> {
+  static inline clang::idx::Entity getEmptyKey() {
+    return clang::idx::Entity::getEmptyMarker();
+  }
+
+  static inline clang::idx::Entity getTombstoneKey() {
+    return clang::idx::Entity::getTombstoneMarker();
+  }
+
+  static unsigned getHashValue(clang::idx::Entity);
+
+  static inline bool 
+  isEqual(clang::idx::Entity LHS, clang::idx::Entity RHS) {
+    return LHS == RHS;
+  }
+
+  static inline bool isPod() { return true; }
+};
+
+}  // end namespace llvm
+
 #endif
diff --git a/include/clang/Index/EntityHandler.h b/include/clang/Index/EntityHandler.h
index b5c2cee881..4cde3e7f90 100644
--- a/include/clang/Index/EntityHandler.h
+++ b/include/clang/Index/EntityHandler.h
@@ -23,7 +23,7 @@ namespace idx {
 class EntityHandler {
 public:
   virtual ~EntityHandler();
-  virtual void HandleEntity(Entity *Ent) { }
+  virtual void HandleEntity(Entity Ent);
 };
   
 } // namespace idx
diff --git a/include/clang/Index/IndexProvider.h b/include/clang/Index/IndexProvider.h
index 51fa75f0fe..ab86474f9a 100644
--- a/include/clang/Index/IndexProvider.h
+++ b/include/clang/Index/IndexProvider.h
@@ -26,11 +26,11 @@ namespace idx {
 
 /// \brief Maps Entities to TranslationUnits.
 class IndexProvider {
+public:
   typedef llvm::SmallPtrSet<TranslationUnit *, 4> TUSetTy;
-  typedef std::map<Entity *, TUSetTy> MapTy;
+  typedef std::map<Entity, TUSetTy> MapTy;
   class Indexer;
 
-public:
   explicit IndexProvider(Program &prog) : Prog(prog) { }
 
   Program &getProgram() const { return Prog; }
@@ -40,9 +40,9 @@ public:
 
   typedef TUSetTy::iterator translation_unit_iterator;
 
-  translation_unit_iterator translation_units_begin(Entity *Ent) const;
-  translation_unit_iterator translation_units_end(Entity *Ent) const;
-  bool translation_units_empty(Entity *Ent) const;
+  translation_unit_iterator translation_units_begin(Entity Ent) const;
+  translation_unit_iterator translation_units_end(Entity Ent) const;
+  bool translation_units_empty(Entity Ent) const;
   
 private:
   Program &Prog;
diff --git a/lib/Analysis/CallGraph.cpp b/lib/Analysis/CallGraph.cpp
index cf8f6205ab..2530fc0ad7 100644
--- a/lib/Analysis/CallGraph.cpp
+++ b/lib/Analysis/CallGraph.cpp
@@ -25,12 +25,12 @@ class CGBuilder : public StmtVisitor<CGBuilder> {
   CallGraph &G;
   FunctionDecl *FD;
 
-  Entity *CallerEnt;
+  Entity CallerEnt;
 
   CallGraphNode *CallerNode;
 
 public:
-  CGBuilder(CallGraph &g, FunctionDecl *fd, Entity *E, CallGraphNode *N)
+  CGBuilder(CallGraph &g, FunctionDecl *fd, Entity E, CallGraphNode *N)
     : G(g), FD(fd), CallerEnt(E), CallerNode(N) {}
 
   void VisitStmt(Stmt *S) { VisitChildren(S); }
@@ -47,7 +47,7 @@ public:
 
 void CGBuilder::VisitCallExpr(CallExpr *CE) {
   if (FunctionDecl *CalleeDecl = CE->getDirectCallee()) {
-    Entity *Ent = Entity::get(CalleeDecl, G.getProgram());
+    Entity Ent = Entity::get(CalleeDecl, G.getProgram());
     CallGraphNode *CalleeNode = G.getOrInsertFunction(Ent);
 
     Decl *Parent = ASTLocation::FindImmediateParent(FD, CE);
@@ -75,7 +75,7 @@ void CallGraph::addTU(ASTUnit &AST) {
     if (FunctionDecl *FD = dyn_cast<FunctionDecl>(*I)) {
       if (FD->isThisDeclarationADefinition()) {
         // Set caller's ASTContext.
-        Entity *Ent = Entity::get(FD, Prog);
+        Entity Ent = Entity::get(FD, Prog);
         CallGraphNode *Node = getOrInsertFunction(Ent);
         CallerCtx[Node] = &Ctx;
         
@@ -86,7 +86,7 @@ void CallGraph::addTU(ASTUnit &AST) {
   }
 }
 
-CallGraphNode *CallGraph::getOrInsertFunction(Entity *F) {
+CallGraphNode *CallGraph::getOrInsertFunction(Entity F) {
   CallGraphNode *&Node = FunctionMap[F];
   if (Node)
     return Node;
@@ -98,7 +98,7 @@ void CallGraph::print(llvm::raw_ostream &os) {
   for (iterator I = begin(), E = end(); I != E; ++I) {
     if (I->second->hasCallee()) {
       ASTContext &Ctx = *CallerCtx[I->second];
-      os << "function: " << I->first->getPrintableName(Ctx).c_str() 
+      os << "function: " << I->first.getPrintableName(Ctx).c_str() 
          << " calls:\n";
       for (CallGraphNode::iterator CI = I->second->begin(), 
              CE = I->second->end(); CI != CE; ++CI) {
diff --git a/lib/Index/Entity.cpp b/lib/Index/Entity.cpp
index 7a6c05c416..feed3e4c80 100644
--- a/lib/Index/Entity.cpp
+++ b/lib/Index/Entity.cpp
@@ -12,9 +12,9 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "clang/Index/Entity.h"
-#include "clang/Index/Program.h"
+#include "EntityImpl.h"
 #include "ProgramImpl.h"
+#include "clang/Index/Program.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/DeclVisitor.h"
@@ -33,99 +33,123 @@ namespace clang {
 namespace idx {
 
 /// \brief Gets the Entity associated with a Decl.
-class EntityGetter : public DeclVisitor<EntityGetter, Entity *> {
+class EntityGetter : public DeclVisitor<EntityGetter, Entity> {
   ProgramImpl &Prog;
-  
+
 public:
   EntityGetter(ProgramImpl &prog) : Prog(prog) { }
   
-  Entity *get(Entity *Parent, DeclarationName Name);
-  
-  Entity *VisitNamedDecl(NamedDecl *D);
-  Entity *VisitVarDecl(VarDecl *D);
-  Entity *VisitFunctionDecl(FunctionDecl *D); 
+  Entity VisitNamedDecl(NamedDecl *D);
+  Entity VisitVarDecl(VarDecl *D);
+  Entity VisitFunctionDecl(FunctionDecl *D); 
 };
 
 }
 }
 
-Entity *EntityGetter::get(Entity *Parent, DeclarationName Name) {
+Entity EntityGetter::VisitNamedDecl(NamedDecl *D) {
+  Entity Parent;
+  if (!D->getDeclContext()->isTranslationUnit()) {
+    Parent = Visit(cast<Decl>(D->getDeclContext()));
+    // FIXME: Anonymous structs ?
+    if (Parent.isInvalid())
+      return Entity();
+  }
+  if (Parent.isValid() && Parent.isInternalToTU())
+    return Entity(D);
+
   // FIXME: Only works for DeclarationNames that are identifiers.
 
+  DeclarationName Name = D->getDeclName();
+
   if (!Name.isIdentifier())
-    return 0;
+    return Entity();
 
   IdentifierInfo *II = Name.getAsIdentifierInfo();
   if (!II)
-      return 0;
+      return Entity();
 
-  ProgramImpl::IdEntryTy *Id =
+  EntityImpl::IdEntryTy *Id =
       &Prog.getIdents().GetOrCreateValue(II->getName(),
                                          II->getName() + II->getLength());
-  
+  unsigned IdNS = D->getIdentifierNamespace();
+
   llvm::FoldingSetNodeID ID;
-  Entity::Profile(ID, Parent, Id);
-  
+  EntityImpl::Profile(ID, Parent, Id, IdNS);
+
   ProgramImpl::EntitySetTy &Entities = Prog.getEntities();
   void *InsertPos = 0;
-  if (Entity *Ent = Entities.FindNodeOrInsertPos(ID, InsertPos))
-    return Ent;
-  
-  void *Buf = Prog.Allocate(sizeof(Entity));
-  Entity *New = new (Buf) Entity(Parent, Id);
-  Entities.InsertNode(New, InsertPos);
-  return New;
-}
+  if (EntityImpl *Ent = Entities.FindNodeOrInsertPos(ID, InsertPos))
+    return Entity(Ent);
 
-Entity *EntityGetter::VisitNamedDecl(NamedDecl *D) {
-  // FIXME: Function declarations that are inside functions ?
-  if (!D->getDeclContext()->isFileContext())
-    return 0;
-
-  Entity *Parent = Visit(cast<Decl>(D->getDeclContext()));
-  return get(Parent, D->getDeclName());
+  void *Buf = Prog.Allocate(sizeof(EntityImpl));
+  EntityImpl *New = new (Buf) EntityImpl(Parent, Id, IdNS);
+  Entities.InsertNode(New, InsertPos);
+  
+  return Entity(New);
 }
 
-Entity *EntityGetter::VisitVarDecl(VarDecl *D) {
+Entity EntityGetter::VisitVarDecl(VarDecl *D) {
   // If it's static it cannot be referred to by another translation unit.
   if (D->getStorageClass() == VarDecl::Static)
-    return 0;
+    return Entity(D);
   
   return VisitNamedDecl(D);
 }
 
-Entity *EntityGetter::VisitFunctionDecl(FunctionDecl *D) {
+Entity EntityGetter::VisitFunctionDecl(FunctionDecl *D) {
   // If it's static it cannot be refered to by another translation unit.
   if (D->getStorageClass() == FunctionDecl::Static)
-    return 0;
+    return Entity(D);
   
   return VisitNamedDecl(D);
 }
 
 //===----------------------------------------------------------------------===//
-// Entity Implementation
+// EntityImpl Implementation
 //===----------------------------------------------------------------------===//
 
-/// \brief Find the Decl that can be referred to by this entity.
-Decl *Entity::getDecl(ASTContext &AST) {
+Decl *EntityImpl::getDecl(ASTContext &AST) {
   DeclContext *DC =
-    Parent == 0 ? AST.getTranslationUnitDecl()
-                : cast<DeclContext>(Parent->getDecl(AST));
+    Parent.isInvalid() ? AST.getTranslationUnitDecl()
+                       : cast<DeclContext>(Parent.getDecl(AST));
   if (!DC)
     return 0; // Couldn't get the parent context.
     
-  ProgramImpl::IdEntryTy *Entry = static_cast<ProgramImpl::IdEntryTy *>(Id);
-  IdentifierInfo &II = AST.Idents.get(Entry->getKeyData());
+  IdentifierInfo &II = AST.Idents.get(Id->getKeyData());
   
   DeclContext::lookup_result Res = DC->lookup(DeclarationName(&II));
   for (DeclContext::lookup_iterator I = Res.first, E = Res.second; I!=E; ++I) {
-    if (!isa<TagDecl>(*I))
+    if ((*I)->getIdentifierNamespace() == IdNS)
       return *I;
   }
 
   return 0; // Failed to find a decl using this Entity.
 }
 
+/// \brief Get an Entity associated with the given Decl.
+/// \returns Null if an Entity cannot refer to this Decl.
+Entity EntityImpl::get(Decl *D, ProgramImpl &Prog) {
+  assert(D && "Passed null Decl");
+  return EntityGetter(Prog).Visit(D);
+}
+
+//===----------------------------------------------------------------------===//
+// Entity Implementation
+//===----------------------------------------------------------------------===//
+
+/// \brief Find the Decl that can be referred to by this entity.
+Decl *Entity::getDecl(ASTContext &AST) {
+  if (isInvalid())
+    return 0;
+  
+  if (Decl *D = Val.dyn_cast<Decl *>())
+    // Check that the passed AST is actually the one that this Decl belongs to.
+    return (&D->getASTContext() == &AST) ? D : 0;
+  
+  return Val.get<EntityImpl *>()->getDecl(AST);
+}
+
 std::string Entity::getPrintableName(ASTContext &Ctx) {
   if (const NamedDecl *ND = dyn_cast_or_null<NamedDecl>(getDecl(Ctx))) {
     return ND->getNameAsString();
@@ -135,8 +159,14 @@ std::string Entity::getPrintableName(ASTContext &Ctx) {
 
 /// \brief Get an Entity associated with the given Decl.
 /// \returns Null if an Entity cannot refer to this Decl.
-Entity *Entity::get(Decl *D, Program &Prog) {
-  assert(D && "Passed null Decl");
-  ProgramImpl &Impl = *static_cast<ProgramImpl*>(Prog.Impl);
-  return EntityGetter(Impl).Visit(D);
+Entity Entity::get(Decl *D, Program &Prog) {
+  if (D == 0)
+    return Entity();
+  ProgramImpl &ProgImpl = *static_cast<ProgramImpl*>(Prog.Impl);
+  return EntityImpl::get(D, ProgImpl);
+}
+
+unsigned 
+llvm::DenseMapInfo<Entity>::getHashValue(Entity E) {
+  return DenseMapInfo<void*>::getHashValue(E.getAsOpaquePtr());
 }
diff --git a/lib/Index/EntityImpl.h b/lib/Index/EntityImpl.h
new file mode 100644
index 0000000000..334dcfb4ec
--- /dev/null
+++ b/lib/Index/EntityImpl.h
@@ -0,0 +1,63 @@
+//===--- EntityImpl.h - Internal Entity implementation---------*- C++ -*-=====//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+//  Internal implementation for the Entity class
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_INDEX_ENTITYIMPL_H
+#define LLVM_CLANG_INDEX_ENTITYIMPL_H
+
+#include "clang/Index/Entity.h"
+#include "llvm/ADT/FoldingSet.h"
+#include "llvm/ADT/StringSet.h"
+
+namespace clang {
+
+namespace idx {
+  class ProgramImpl;
+
+class EntityImpl : public llvm::FoldingSetNode {
+public:
+  typedef llvm::StringMapEntry<char> IdEntryTy;
+
+private:
+  Entity Parent;
+  IdEntryTy *Id;
+
+  /// \brief Identifier namespace.
+  unsigned IdNS;
+
+public:
+  EntityImpl(Entity parent, IdEntryTy *id, unsigned idNS)
+    : Parent(parent), Id(id), IdNS(idNS) { }
+
+  /// \brief Find the Decl that can be referred to by this entity.
+  Decl *getDecl(ASTContext &AST);
+
+  /// \brief Get an Entity associated with the given Decl.
+  /// \returns Null if an Entity cannot refer to this Decl.
+  static Entity get(Decl *D, ProgramImpl &Prog);
+
+  void Profile(llvm::FoldingSetNodeID &ID) const {
+    Profile(ID, Parent, Id, IdNS);
+  }
+  static void Profile(llvm::FoldingSetNodeID &ID, Entity Parent, IdEntryTy *Id,
+                      unsigned IdNS) {
+    ID.AddPointer(Parent.getAsOpaquePtr());
+    ID.AddPointer(Id);
+    ID.AddInteger(IdNS);
+  }
+};
+
+} // namespace idx
+
+} // namespace clang
+
+#endif
diff --git a/lib/Index/IndexProvider.cpp b/lib/Index/IndexProvider.cpp
index 074a83d9a6..38317eb9c2 100644
--- a/lib/Index/IndexProvider.cpp
+++ b/lib/Index/IndexProvider.cpp
@@ -13,6 +13,7 @@
 
 #include "clang/Index/IndexProvider.h"
 #include "clang/Index/Program.h"
+#include "clang/Index/Entity.h"
 #include "clang/Index/EntityHandler.h"
 #include "clang/Index/TranslationUnit.h"
 using namespace clang;
@@ -25,7 +26,10 @@ class IndexProvider::Indexer : public EntityHandler {
 public:
   Indexer(TranslationUnit *tu, MapTy &map) : TU(tu), Map(map) { }
 
-  virtual void HandleEntity(Entity *Ent) {
+  virtual void HandleEntity(Entity Ent) {
+    if (Ent.isInternalToTU())
+      return;
+
     MapTy::iterator I = Map.find(Ent);
     if (I != Map.end()) {
       I->second.insert(TU);
@@ -41,25 +45,27 @@ void IndexProvider::IndexAST(TranslationUnit *TU) {
   Prog.FindEntities(TU->getASTContext(), &Idx);
 }
 
+static IndexProvider::TUSetTy EmptySet;
+
 IndexProvider::translation_unit_iterator
-IndexProvider::translation_units_begin(Entity *Ent) const {
+IndexProvider::translation_units_begin(Entity Ent) const {
   MapTy::iterator I = Map.find(Ent);
   if (I == Map.end())
-    return translation_unit_iterator(0);
+    return EmptySet.begin();
   
   return I->second.begin();
 }
 
 IndexProvider::translation_unit_iterator
-IndexProvider::translation_units_end(Entity *Ent) const {
+IndexProvider::translation_units_end(Entity Ent) const {
   MapTy::iterator I = Map.find(Ent);
   if (I == Map.end())
-    return translation_unit_iterator(0);
+    return EmptySet.end();
   
   return I->second.end();
 }
 
-bool IndexProvider::translation_units_empty(Entity *Ent) const {
+bool IndexProvider::translation_units_empty(Entity Ent) const {
   MapTy::iterator I = Map.find(Ent);
   if (I == Map.end())
     return true;
diff --git a/lib/Index/Program.cpp b/lib/Index/Program.cpp
index 807d18c699..73759dc95c 100644
--- a/lib/Index/Program.cpp
+++ b/lib/Index/Program.cpp
@@ -25,6 +25,8 @@ using namespace idx;
 EntityHandler::~EntityHandler() { }
 TranslationUnit::~TranslationUnit() { }
 
+void EntityHandler::HandleEntity(Entity Ent) { }
+
 Program::Program() : Impl(new ProgramImpl()) { }
 
 Program::~Program() {
@@ -34,8 +36,10 @@ Program::~Program() {
 static void FindEntitiesInDC(DeclContext *DC, Program &Prog, EntityHandler *Handler) {
   for (DeclContext::decl_iterator
          I = DC->decls_begin(), E = DC->decls_end(); I != E; ++I) {
-    Entity *Ent = Entity::get(*I, Prog);
-    if (Ent)
+    if (I->getLocation().isInvalid())
+      continue;
+    Entity Ent = Entity::get(*I, Prog);
+    if (Ent.isValid())
       Handler->HandleEntity(Ent);
     if (DeclContext *SubDC = dyn_cast<DeclContext>(*I))
       FindEntitiesInDC(SubDC, Prog, Handler);
diff --git a/lib/Index/ProgramImpl.h b/lib/Index/ProgramImpl.h
index 52f153f1dc..39fc184677 100644
--- a/lib/Index/ProgramImpl.h
+++ b/lib/Index/ProgramImpl.h
@@ -14,8 +14,7 @@
 #ifndef LLVM_CLANG_INDEX_PROGRAMIMPL_H
 #define LLVM_CLANG_INDEX_PROGRAMIMPL_H
 
-#include "clang/Index/Entity.h"
-#include "llvm/ADT/StringSet.h"
+#include "EntityImpl.h"
 
 namespace clang {
 
@@ -24,11 +23,10 @@ namespace idx {
 
 class ProgramImpl {
 public:
-  typedef llvm::FoldingSet<Entity> EntitySetTy;
-  typedef llvm::StringMapEntry<char> IdEntryTy;
-  
+  typedef llvm::FoldingSet<EntityImpl> EntitySetTy;
+
 private:
-  llvm::FoldingSet<Entity> Entities;
+  EntitySetTy Entities;
   llvm::StringSet<> Idents;
   llvm::BumpPtrAllocator BumpAlloc;
 
@@ -38,7 +36,7 @@ private:
 public:
   ProgramImpl() { }
   
-  llvm::FoldingSet<Entity> &getEntities() { return Entities; }
+  EntitySetTy &getEntities() { return Entities; }
   llvm::StringSet<> &getIdents() { return Idents; }
 
   void *Allocate(unsigned Size, unsigned Align = 8) {
diff --git a/test/Index/find-refs.c b/test/Index/find-refs.c
index 5212c24695..9212922ba6 100644
--- a/test/Index/find-refs.c
+++ b/test/Index/find-refs.c
@@ -39,3 +39,8 @@
 // RUN: index-test %t1.ast %t2.ast -point-at %S/t1.c:16:7 -print-refs > %t &&
 // RUN: cat %t | count 1 &&
 // RUN: grep 't1.c:22:3,' %t
+
+// RUN: index-test %t1.ast %t2.ast -point-at %S/foo.h:7:11 -print-refs > %t &&
+// RUN: cat %t | count 2 &&
+// RUN: grep 't1.c:25:3,' %t &&
+// RUN: grep 't2.c:10:3,' %t
diff --git a/test/Index/foo.h b/test/Index/foo.h
index 2e3b403dea..7670c00dfb 100644
--- a/test/Index/foo.h
+++ b/test/Index/foo.h
@@ -2,3 +2,7 @@ extern int global_var;
 
 void foo_func(int param1);
 void bar_func(void);
+
+struct MyStruct {
+  int field_var;
+};
diff --git a/test/Index/t1.c b/test/Index/t1.c
index 17cdb7dc4f..d6ca995604 100644
--- a/test/Index/t1.c
+++ b/test/Index/t1.c
@@ -20,4 +20,7 @@ void field_test(void) {
   struct S1 s1;
   s1.x = 0;
   ((struct S2 *)0)->x = 0;
+  
+  struct MyStruct ms;
+  ms.field_var = 10;
 }
diff --git a/test/Index/t2.c b/test/Index/t2.c
index 8e405cc367..76d5d6ceeb 100644
--- a/test/Index/t2.c
+++ b/test/Index/t2.c
@@ -5,4 +5,7 @@ int global_var = 10;
 void bar_func(void) {
   global_var += 100;
   foo_func(global_var);
+
+  struct MyStruct *ms;
+  ms->field_var = 10;
 }
diff --git a/tools/index-test/index-test.cpp b/tools/index-test/index-test.cpp
index 8fd871d08c..4516e31365 100644
--- a/tools/index-test/index-test.cpp
+++ b/tools/index-test/index-test.cpp
@@ -151,10 +151,8 @@ static void ProcessASTLocation(ASTLocation ASTLoc, IndexProvider &IdxProvider) {
     return;
   }
 
-  Entity *Ent = Entity::get(D, IdxProvider.getProgram());
-  // If there is no Entity associated with this Decl, it means that it's not
-  // visible to other translation units.
-  if (!Ent)
+  Entity Ent = Entity::get(D, IdxProvider.getProgram());
+  if (Ent.isInvalid() || Ent.isInternalToTU())
     return ProcessDecl(D);
 
   // Find the "same" Decl in other translation units and print information.
@@ -162,7 +160,7 @@ static void ProcessASTLocation(ASTLocation ASTLoc, IndexProvider &IdxProvider) {
          I = IdxProvider.translation_units_begin(Ent),
          E = IdxProvider.translation_units_end(Ent); I != E; ++I) {
     TUnit *TU = static_cast<TUnit*>(*I);
-    Decl *OtherD = Ent->getDecl(TU->getASTContext());
+    Decl *OtherD = Ent.getDecl(TU->getASTContext());
     assert(OtherD && "Couldn't resolve Entity");
     ProcessDecl(OtherD);
   }