[index] For extensions, relate the symbol reference that gets extended and base references with the extension symbol.

They were getting related with the original extended symbol, which was incorrect.

include/swift/Index/IndexSymbol.h

@@ -45,7 +45,7 @@ inline SymbolPropertySet &operator|=(SymbolPropertySet &SKSet, SymbolProperty SK
 }
 
 struct IndexRelation {
-  const ValueDecl *decl;
+  const Decl *decl;
   SymbolInfo symInfo;
   SymbolRoleSet roles = SymbolRoleSet(0);
 
@@ -55,7 +55,7 @@ struct IndexRelation {
   StringRef USR; // USR may be safely compared by pointer.
   StringRef group;
 
-  IndexRelation(SymbolRoleSet Roles, const ValueDecl *Sym, SymbolInfo SymInfo, StringRef Name, StringRef USR)
+  IndexRelation(SymbolRoleSet Roles, const Decl *Sym, SymbolInfo SymInfo, StringRef Name, StringRef USR)
   : decl(Sym), symInfo(SymInfo), roles(Roles), name(Name), USR(USR) {}
 
   IndexRelation() = default;

lib/Index/Index.cpp

@@ -189,15 +189,21 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
     return false;
   }
 
-  bool addRelation(IndexSymbol &Info, SymbolRoleSet RelationRoles, ValueDecl *D) {
+  bool addRelation(IndexSymbol &Info, SymbolRoleSet RelationRoles, Decl *D) {
     assert(D);
     StringRef Name, USR;
     SymbolInfo SymInfo = getSymbolInfoForDecl(D);
 
     if (SymInfo.Kind == SymbolKind::Unknown)
       return true;
-    if (getNameAndUSR(D, /*ExtD=*/nullptr, Name, USR))
+    if (auto *ExtD = dyn_cast<ExtensionDecl>(D)) {
+      NominalTypeDecl *NTD = ExtD->getExtendedType()->getAnyNominal();
+      if (getNameAndUSR(NTD, ExtD, Name, USR))
         return true;
+    } else {
+      if (getNameAndUSR(cast<ValueDecl>(D), /*ExtD=*/nullptr, Name, USR))
+        return true;
+    }
 
     Info.Relations.push_back(IndexRelation(RelationRoles, D, SymInfo, Name, USR));
     Info.roles |= RelationRoles;
@@ -339,9 +345,9 @@ private:
   bool startEntity(Decl *D, IndexSymbol &Info);
   bool startEntityDecl(ValueDecl *D);
 
-  bool reportRelatedRef(ValueDecl *D, SourceLoc Loc, SymbolRoleSet Relations, ValueDecl *Related);
-  bool reportRelatedTypeRef(const TypeLoc &Ty, SymbolRoleSet Relations, ValueDecl *Related);
-  bool reportInheritedTypeRefs(ArrayRef<TypeLoc> Inherited, ValueDecl *Inheritee);
+  bool reportRelatedRef(ValueDecl *D, SourceLoc Loc, SymbolRoleSet Relations, Decl *Related);
+  bool reportRelatedTypeRef(const TypeLoc &Ty, SymbolRoleSet Relations, Decl *Related);
+  bool reportInheritedTypeRefs(ArrayRef<TypeLoc> Inherited, Decl *Inheritee);
   NominalTypeDecl *getTypeLocAsNominalTypeDecl(const TypeLoc &Ty);
 
   bool reportPseudoGetterDecl(VarDecl *D) {
@@ -609,7 +615,7 @@ bool IndexSwiftASTWalker::startEntityDecl(ValueDecl *D) {
   return startEntity(D, Info);
 }
 
-bool IndexSwiftASTWalker::reportRelatedRef(ValueDecl *D, SourceLoc Loc, SymbolRoleSet Relations, ValueDecl *Related) {
+bool IndexSwiftASTWalker::reportRelatedRef(ValueDecl *D, SourceLoc Loc, SymbolRoleSet Relations, Decl *Related) {
   if (!shouldIndex(D))
     return true;
 
@@ -628,7 +634,7 @@ bool IndexSwiftASTWalker::reportRelatedRef(ValueDecl *D, SourceLoc Loc, SymbolRo
   return !Cancelled;
 }
 
-bool IndexSwiftASTWalker::reportInheritedTypeRefs(ArrayRef<TypeLoc> Inherited, ValueDecl *Inheritee) {
+bool IndexSwiftASTWalker::reportInheritedTypeRefs(ArrayRef<TypeLoc> Inherited, Decl *Inheritee) {
   for (auto Base : Inherited) {
     if(!reportRelatedTypeRef(Base, (SymbolRoleSet) SymbolRole::RelationBaseOf, Inheritee))
       return false;
@@ -636,7 +642,7 @@ bool IndexSwiftASTWalker::reportInheritedTypeRefs(ArrayRef<TypeLoc> Inherited, V
   return true;
 }
 
-bool IndexSwiftASTWalker::reportRelatedTypeRef(const TypeLoc &Ty, SymbolRoleSet Relations, ValueDecl *Related) {
+bool IndexSwiftASTWalker::reportRelatedTypeRef(const TypeLoc &Ty, SymbolRoleSet Relations, Decl *Related) {
 
   if (IdentTypeRepr *T = dyn_cast_or_null<IdentTypeRepr>(Ty.getTypeRepr())) {
     auto Comps = T->getComponentRange();
@@ -718,7 +724,10 @@ IndexSwiftASTWalker::getTypeLocAsNominalTypeDecl(const TypeLoc &Ty) {
 }
 
 bool IndexSwiftASTWalker::reportExtension(ExtensionDecl *D) {
-  SourceLoc Loc = D->getExtendedTypeLoc().getSourceRange().Start;
+  // Use the 'End' token of the range, in case it is a compound name, e.g.
+  //   extension A.B {}
+  // we want the location of 'B' token.
+  SourceLoc Loc = D->getExtendedTypeLoc().getSourceRange().End;
   if (!D->getExtendedType())
     return true;
   NominalTypeDecl *NTD = D->getExtendedType()->getAnyNominal();
@@ -734,11 +743,10 @@ bool IndexSwiftASTWalker::reportExtension(ExtensionDecl *D) {
   if (!startEntity(D, Info))
     return false;
 
-  // FIXME: make extensions their own entity
   if (!reportRelatedRef(NTD, Loc, (SymbolRoleSet)SymbolRole::RelationExtendedBy,
-                        NTD))
+                        D))
       return false;
-  if (!reportInheritedTypeRefs(D->getInherited(), NTD))
+  if (!reportInheritedTypeRefs(D->getInherited(), D))
       return false;
 
   return true;

test/Index/kinds.swift

@@ -123,24 +123,24 @@ protocol AProtocol {
 
 // Extension
 extension AnEnumeration { func extFn() {} }
-// CHECK: [[@LINE-1]]:11 | extension/ext-enum/Swift | AnEnumeration | s:e:s:14swift_ide_test13AnEnumerationO5extFnyyF | Def | rel: 0
+// CHECK: [[@LINE-1]]:11 | extension/ext-enum/Swift | AnEnumeration | [[EXT_AnEnumeration_USR:s:e:s:14swift_ide_test13AnEnumerationO5extFnyyF]] | Def | rel: 0
 // CHECK: [[@LINE-2]]:11 | enum/Swift | AnEnumeration | s:14swift_ide_test13AnEnumerationO | Ref,RelExt | rel: 1
-// CHECK-NEXT: RelExt | AnEnumeration | s:14swift_ide_test13AnEnumerationO
+// CHECK-NEXT: RelExt | AnEnumeration | [[EXT_AnEnumeration_USR]]
 
 extension AStruct { func extFn() {} }
-// CHECK: [[@LINE-1]]:11 | extension/ext-struct/Swift | AStruct | s:e:s:14swift_ide_test7AStructV5extFnyyF | Def | rel: 0
+// CHECK: [[@LINE-1]]:11 | extension/ext-struct/Swift | AStruct | [[EXT_AStruct_USR:s:e:s:14swift_ide_test7AStructV5extFnyyF]] | Def | rel: 0
 // CHECK: [[@LINE-2]]:11 | struct/Swift | AStruct | s:14swift_ide_test7AStructV | Ref,RelExt | rel: 1
-// CHECK-NEXT: RelExt | AStruct | s:14swift_ide_test7AStructV
+// CHECK-NEXT: RelExt | AStruct | [[EXT_AStruct_USR]]
 
 extension AClass { func extFn() {} }
-// CHECK: [[@LINE-1]]:11 | extension/ext-class/Swift | AClass | s:e:s:14swift_ide_test6AClassC5extFnyyF | Def | rel: 0
+// CHECK: [[@LINE-1]]:11 | extension/ext-class/Swift | AClass | [[EXT_AClass_USR:s:e:s:14swift_ide_test6AClassC5extFnyyF]] | Def | rel: 0
 // CHECK: [[@LINE-2]]:11 | class/Swift | AClass | s:14swift_ide_test6AClassC | Ref,RelExt | rel: 1
-// CHECK-NEXT: RelExt | AClass | s:14swift_ide_test6AClassC
+// CHECK-NEXT: RelExt | AClass | [[EXT_AClass_USR]]
 
 extension AProtocol { func extFn() }
-// CHECK: [[@LINE-1]]:11 | extension/ext-protocol/Swift | AProtocol | s:e:s:14swift_ide_test9AProtocolPAAE5extFnyyF | Def | rel: 0
+// CHECK: [[@LINE-1]]:11 | extension/ext-protocol/Swift | AProtocol | [[EXT_AProtocol_USR:s:e:s:14swift_ide_test9AProtocolPAAE5extFnyyF]] | Def | rel: 0
 // CHECK: [[@LINE-2]]:11 | protocol/Swift | AProtocol | s:14swift_ide_test9AProtocolP | Ref,RelExt | rel: 1
-// CHECK-NEXT: RelExt | AProtocol | s:14swift_ide_test9AProtocolP
+// CHECK-NEXT: RelExt | AProtocol | [[EXT_AProtocol_USR]]
 
 // TypeAlias
 typealias SomeAlias = AStruct

test/Index/roles.swift

@@ -134,6 +134,7 @@ class AClass {
 }
 
 protocol AProtocol {
+  // CHECK: [[@LINE-1]]:10 | protocol/Swift | AProtocol | [[AProtocol_USR:.*]] | Def | rel: 0
   func foo() -> Int
   // CHECK: [[@LINE-1]]:8 | instance-method/Swift | foo() | s:FP14swift_ide_test9AProtocol3fooFT_Si | Def,RelChild | rel: 1
   // CHECK-NEXT: RelChild | AProtocol | s:P14swift_ide_test9AProtocol
@@ -157,17 +158,32 @@ class ASubClass : AClass, AProtocol {
 }
 
 // RelationExtendedBy
-// FIXME give extensions their own USR like ObjC?
 extension AClass {
-  // CHECK: [[@LINE-1]]:11 | extension/ext-class/Swift | AClass | s:e:s:FC14swift_ide_test6AClass3barFT_Si | Def | rel: 0
+  // CHECK: [[@LINE-1]]:11 | extension/ext-class/Swift | AClass | [[EXT_ACLASS_USR:.*]] | Def | rel: 0
   // CHECK: [[@LINE-2]]:11 | class/Swift | AClass | s:C14swift_ide_test6AClass | Ref,RelExt | rel: 1
-  // CHECK-NEXT: RelExt | AClass | s:C14swift_ide_test6AClass
+  // CHECK-NEXT: RelExt | AClass | [[EXT_ACLASS_USR]]
 
   func bar() -> Int { return 2 }
   // CHECK: [[@LINE-1]]:8 | instance-method/Swift | bar() | s:FC14swift_ide_test6AClass3barFT_Si | Def,RelChild | rel: 1
   // CHECK-NEXT: RelChild | AClass | s:C14swift_ide_test6AClass
 }
 
+struct OuterS {
+// CHECK: [[@LINE-1]]:8 | struct/Swift | OuterS | [[OUTERS_USR:.*]] | Def | rel: 0
+  struct InnerS {}
+  // CHECK: [[@LINE-1]]:10 | struct/Swift | InnerS | [[INNERS_USR:.*]] | Def,RelChild | rel: 1
+  // CHECK-NEXT: RelChild | OuterS | [[OUTERS_USR]]
+}
+extension OuterS.InnerS : AProtocol {
+  // CHECK: [[@LINE-1]]:18 | extension/ext-struct/Swift | InnerS | [[EXT_INNERS_USR:.*]] | Def | rel: 0
+  // CHECK: [[@LINE-2]]:18 | struct/Swift | InnerS | [[INNERS_USR]] | Ref,RelExt | rel: 1
+  // CHECK-NEXT: RelExt | InnerS | [[EXT_INNERS_USR]]
+  // CHECK: [[@LINE-4]]:27 | protocol/Swift | AProtocol | [[AProtocol_USR]] | Ref,RelBase | rel: 1
+  // CHECK-NEXT: RelBase | InnerS | [[EXT_INNERS_USR]]
+  // CHECK: [[@LINE-6]]:11 | struct/Swift | OuterS | [[OUTERS_USR]] | Ref | rel: 0
+  func foo() {}
+}
+
 var anInstance = AClass(x: 1)
 // CHECK: [[@LINE-1]]:18 | class/Swift | AClass | s:C14swift_ide_test6AClass | Ref | rel: 0
 // CHECK: [[@LINE-2]]:18 | constructor/Swift | init(x:) | s:FC14swift_ide_test6AClasscFT1xSi_S0_ | Ref,Call | rel: 0