Merge pull request #24072 from benlangmuir/kpdml-ide

Code-completion, indexing, cursor-info, etc. for KeyPath dynamic member lookup

Author: akyrtzi

include/swift/AST/ASTWalker.h

@@ -41,8 +41,10 @@ enum class SemaReferenceKind : uint8_t {
 struct ReferenceMetaData {
   SemaReferenceKind Kind;
   llvm::Optional<AccessKind> AccKind;
-  ReferenceMetaData(SemaReferenceKind Kind, llvm::Optional<AccessKind> AccKind) :
-    Kind(Kind), AccKind(AccKind) {}
+  bool isImplicit = false;
+  ReferenceMetaData(SemaReferenceKind Kind, llvm::Optional<AccessKind> AccKind,
+                    bool isImplicit = false)
+      : Kind(Kind), AccKind(AccKind), isImplicit(isImplicit) {}
 };
 
 /// An abstract class used to traverse an AST.

include/swift/IDE/SourceEntityWalker.h

@@ -113,11 +113,11 @@ class SourceEntityWalker {
   ///
   /// \param D the referenced decl.
   /// \param Range the source range of the source reference.
-  /// \param AccKind whether this is a read, write or read/write access.
+  /// \param Data whether this is a read, write or read/write access, etc.
   /// \param IsOpenBracket this is \c true when the method is called on an
   /// open bracket.
   virtual bool visitSubscriptReference(ValueDecl *D, CharSourceRange Range,
-                                       Optional<AccessKind> AccKind,
+                                       ReferenceMetaData Data,
                                        bool IsOpenBracket);
 
   /// This method is called for each keyword argument in a call expression.

include/swift/IDE/Utils.h

@@ -231,7 +231,7 @@ class CursorInfoResolver : public SourceEntityWalker {
   bool tryResolve(ModuleEntity Mod, SourceLoc Loc);
   bool tryResolve(Stmt *St);
   bool visitSubscriptReference(ValueDecl *D, CharSourceRange Range,
-                               Optional<AccessKind> AccKind,
+                               ReferenceMetaData Data,
                                bool IsOpenBracket) override;
 };
 

include/swift/Sema/IDETypeChecking.h

@@ -227,6 +227,18 @@ namespace swift {
   /// Sometimes for diagnostics we want to work on the original argument list as
   /// written by the user; this performs the reverse transformation.
   OriginalArgumentList getOriginalArgumentList(Expr *expr);
+
+  /// Return true if the specified type or a super-class/super-protocol has the
+  /// @dynamicMemberLookup attribute on it.
+  bool hasDynamicMemberLookupAttribute(Type type);
+
+  /// Returns the root type of the keypath type in a keypath dynamic member
+  /// lookup subscript, or \c None if it cannot be determined.
+  ///
+  /// \param subscript The potential keypath dynamic member lookup subscript.
+  /// \param DC The DeclContext from which the subscript is being referenced.
+  Optional<Type> getRootTypeOfKeypathDynamicMember(SubscriptDecl *subscript,
+                                                   const DeclContext *DC);
 }
 
 #endif

lib/IDE/SourceEntityWalker.cpp

@@ -70,8 +70,7 @@ class SemaAnnotator : public ASTWalker {
   bool passReference(ModuleEntity Mod, std::pair<Identifier, SourceLoc> IdLoc);
 
   bool passSubscriptReference(ValueDecl *D, SourceLoc Loc,
-                              Optional<AccessKind> AccKind,
-                              bool IsOpenBracket);
+                              ReferenceMetaData Data, bool IsOpenBracket);
 
   bool passCallArgNames(Expr *Fn, TupleExpr *TupleE);
 
@@ -265,6 +264,8 @@ std::pair<bool, Expr *> SemaAnnotator::walkToExprPre(Expr *E) {
       !isa<MakeTemporarilyEscapableExpr>(E) &&
       !isa<CollectionUpcastConversionExpr>(E) &&
       !isa<OpaqueValueExpr>(E) &&
+      !isa<SubscriptExpr>(E) &&
+      !isa<KeyPathExpr>(E) &&
       E->isImplicit())
     return { true, E };
 
@@ -326,16 +327,19 @@ std::pair<bool, Expr *> SemaAnnotator::walkToExprPre(Expr *E) {
     if (SE->hasDecl())
       SubscrD = SE->getDecl().getDecl();
 
+    ReferenceMetaData data(SemaReferenceKind::SubscriptRef, OpAccess,
+                           SE->isImplicit());
+
     if (SubscrD) {
-      if (!passSubscriptReference(SubscrD, E->getLoc(), OpAccess, true))
+      if (!passSubscriptReference(SubscrD, E->getLoc(), data, true))
         return { false, nullptr };
     }
 
     if (!SE->getIndex()->walk(*this))
       return { false, nullptr };
 
     if (SubscrD) {
-      if (!passSubscriptReference(SubscrD, E->getEndLoc(), OpAccess, false))
+      if (!passSubscriptReference(SubscrD, E->getEndLoc(), data, false))
         return { false, nullptr };
     }
 
@@ -577,14 +581,14 @@ bool SemaAnnotator::passModulePathElements(
 }
 
 bool SemaAnnotator::passSubscriptReference(ValueDecl *D, SourceLoc Loc,
-                                           Optional<AccessKind> AccKind,
+                                           ReferenceMetaData Data,
                                            bool IsOpenBracket) {
   CharSourceRange Range = Loc.isValid()
                         ? CharSourceRange(Loc, 1)
                         : CharSourceRange();
 
-  bool Continue = SEWalker.visitSubscriptReference(D, Range, AccKind,
-                                                   IsOpenBracket);
+  bool Continue =
+      SEWalker.visitSubscriptReference(D, Range, Data, IsOpenBracket);
   if (!Continue)
     Cancelled = true;
   return Continue;
@@ -727,13 +731,14 @@ bool SourceEntityWalker::visitDeclReference(ValueDecl *D, CharSourceRange Range,
 
 bool SourceEntityWalker::visitSubscriptReference(ValueDecl *D,
                                                  CharSourceRange Range,
-                                                 Optional<AccessKind> AccKind,
+                                                 ReferenceMetaData Data,
                                                  bool IsOpenBracket) {
   // Most of the clients treat subscript reference the same way as a
   // regular reference when called on the open bracket and
   // ignore the closing one.
-  return IsOpenBracket ? visitDeclReference(D, Range, nullptr, nullptr, Type(),
-    ReferenceMetaData(SemaReferenceKind::SubscriptRef, AccKind)) : true;
+  return IsOpenBracket
+             ? visitDeclReference(D, Range, nullptr, nullptr, Type(), Data)
+             : true;
 }
 
 bool SourceEntityWalker::visitCallArgName(Identifier Name,

lib/IDE/SwiftSourceDocInfo.cpp

@@ -118,12 +118,12 @@ bool CursorInfoResolver::tryResolve(Stmt *St) {
   return false;
 }
 
-bool CursorInfoResolver::visitSubscriptReference(ValueDecl *D, CharSourceRange Range,
-                                                 Optional<AccessKind> AccKind,
+bool CursorInfoResolver::visitSubscriptReference(ValueDecl *D,
+                                                 CharSourceRange Range,
+                                                 ReferenceMetaData Data,
                                                  bool IsOpenBracket) {
   // We should treat both open and close brackets equally
-  return visitDeclReference(D, Range, nullptr, nullptr, Type(),
-                    ReferenceMetaData(SemaReferenceKind::SubscriptRef, AccKind));
+  return visitDeclReference(D, Range, nullptr, nullptr, Type(), Data);
 }
 
 ResolvedCursorInfo CursorInfoResolver::resolve(SourceLoc Loc) {
@@ -190,6 +190,8 @@ bool CursorInfoResolver::visitDeclReference(ValueDecl *D,
                                             ReferenceMetaData Data) {
   if (isDone())
     return false;
+  if (Data.isImplicit)
+    return true;
   return !tryResolve(D, CtorTyRef, ExtTyRef, Range.getStart(), /*IsRef=*/true, T);
 }
 
@@ -1423,7 +1425,7 @@ struct RangeResolver::Implementation {
     if (Data.Kind != SemaReferenceKind::DeclRef)
       return;
 
-    if (!isContainedInSelection(CharSourceRange(Start, 0)))
+    if (Data.isImplicit || !isContainedInSelection(CharSourceRange(Start, 0)))
       return;
 
     // If the VD is declared outside of current file, exclude such decl.

lib/Index/Index.cpp

@@ -394,6 +394,10 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
       return true;
 
     IndexSymbol Info;
+
+    if (Data.isImplicit)
+      Info.roles |= (unsigned)SymbolRole::Implicit;
+
     if (CtorTyRef)
       if (!reportRef(CtorTyRef, Loc, Info, Data.AccKind))
         return false;

lib/Migrator/APIDiffMigratorPass.cpp

@@ -484,6 +484,9 @@ struct APIDiffMigratorPass : public ASTMigratorPass, public SourceEntityWalker {
   bool visitDeclReference(ValueDecl *D, CharSourceRange Range,
                           TypeDecl *CtorTyRef, ExtensionDecl *ExtTyRef,
                           Type T, ReferenceMetaData Data) override {
+    if (Data.isImplicit)
+      return true;
+
     for (auto *Item: getRelatedDiffItems(CtorTyRef ? CtorTyRef: D)) {
       std::string RepText;
       if (isSimpleReplacement(Item, isDotMember(Range), RepText)) {
@@ -501,7 +504,7 @@ struct APIDiffMigratorPass : public ASTMigratorPass, public SourceEntityWalker {
     bool visitDeclReference(ValueDecl *D, CharSourceRange Range,
                             TypeDecl *CtorTyRef, ExtensionDecl *ExtTyRef,
                             Type T, ReferenceMetaData Data) override {
-      if (D == Target) {
+      if (D == Target && !Data.isImplicit) {
         Result = Range;
         return false;
       }

lib/Sema/CSSimplify.cpp

@@ -24,6 +24,7 @@
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/Basic/StringExtras.h"
 #include "swift/ClangImporter/ClangModule.h"
+#include "swift/Sema/IDETypeChecking.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/Compiler.h"
 
@@ -3554,6 +3555,12 @@ static bool hasDynamicMemberLookupAttribute(Type type,
   return result;
 }
 
+// for IDETypeChecking
+bool swift::hasDynamicMemberLookupAttribute(Type type) {
+  llvm::DenseMap<CanType, bool> DynamicMemberLookupCache;
+  return ::hasDynamicMemberLookupAttribute(type, DynamicMemberLookupCache);
+}
+
 static bool isKeyPathDynamicMemberLookup(ConstraintLocator *locator) {
   auto path = locator ? locator->getPath() : None;
   return !path.empty() && path.back().isKeyPathDynamicMember();
@@ -3928,8 +3935,8 @@ performMemberLookup(ConstraintKind constraintKind, DeclName memberName,
     // as representing "dynamic lookup" unless it's a direct call
     // to such subscript (in that case label is expected to match).
     if (auto *subscript = dyn_cast<SubscriptDecl>(cand)) {
-      if (hasDynamicMemberLookupAttribute(instanceTy,
-                                          DynamicMemberLookupCache) &&
+      if (::hasDynamicMemberLookupAttribute(instanceTy,
+                                            DynamicMemberLookupCache) &&
           isValidKeyPathDynamicMemberLookup(subscript, TC)) {
         auto info =
             getArgumentLabels(*this, ConstraintLocatorBuilder(memberLocator));
@@ -4021,7 +4028,8 @@ performMemberLookup(ConstraintKind constraintKind, DeclName memberName,
       constraintKind == ConstraintKind::ValueMember &&
       memberName.isSimpleName() && !memberName.isSpecial()) {
     auto name = memberName.getBaseIdentifier();
-    if (hasDynamicMemberLookupAttribute(instanceTy, DynamicMemberLookupCache)) {
+    if (::hasDynamicMemberLookupAttribute(instanceTy,
+                                          DynamicMemberLookupCache)) {
       auto &ctx = getASTContext();
 
       // Recursively look up `subscript(dynamicMember:)` methods in this type.

lib/Sema/LookupVisibleDecls.cpp

@@ -623,6 +623,42 @@ static void lookupVisibleMemberDeclsImpl(
   } while (1);
 }
 
+static void lookupVisibleDynamicMemberLookupDecls(
+    Type baseType, VisibleDeclConsumer &consumer, const DeclContext *dc,
+    LookupState LS, DeclVisibilityKind reason, LazyResolver *typeResolver,
+    GenericSignatureBuilder *GSB, VisitedSet &visited) {
+
+  assert(hasDynamicMemberLookupAttribute(baseType));
+  auto &ctx = dc->getASTContext();
+
+  // Lookup the `subscript(dynamicMember:)` methods in this type.
+  auto subscriptName =
+      DeclName(ctx, DeclBaseName::createSubscript(), ctx.Id_dynamicMember);
+
+  SmallVector<ValueDecl *, 2> subscripts;
+  dc->lookupQualified(baseType, subscriptName, NL_QualifiedDefault,
+                      typeResolver, subscripts);
+
+  for (ValueDecl *VD : subscripts) {
+    auto *subscript = dyn_cast<SubscriptDecl>(VD);
+    if (!subscript)
+      continue;
+
+    auto rootType = getRootTypeOfKeypathDynamicMember(subscript, dc);
+    if (!rootType)
+      continue;
+
+    auto subs =
+        baseType->getMemberSubstitutionMap(dc->getParentModule(), subscript);
+    auto memberType = rootType->subst(subs);
+    if (!memberType || !memberType->mayHaveMembers())
+      continue;
+
+    lookupVisibleMemberDeclsImpl(memberType, consumer, dc, LS, reason,
+                                 typeResolver, GSB, visited);
+  }
+}
+
 namespace {
 
 struct FoundDeclTy {
@@ -824,6 +860,18 @@ class OverrideFilteringConsumer : public VisibleDeclConsumer {
     DeclsToReport.insert(FoundDeclTy(VD, Reason));
   }
 };
+
+struct ShadowedKeyPathMembers : public VisibleDeclConsumer {
+  VisibleDeclConsumer &consumer;
+  llvm::DenseSet<DeclBaseName> &seen;
+  ShadowedKeyPathMembers(VisibleDeclConsumer &consumer, llvm::DenseSet<DeclBaseName> &knownMembers) : consumer(consumer), seen(knownMembers) {}
+  void foundDecl(ValueDecl *VD, DeclVisibilityKind reason) override {
+    // Dynamic lookup members are only visible if they are not shadowed by
+    // non-dynamic members.
+    if (seen.count(VD->getBaseName()) == 0)
+      consumer.foundDecl(VD, reason);
+  }
+};  
 } // end anonymous namespace
 
 /// Enumerate all members in \c BaseTy (including members of extensions,
@@ -841,6 +889,16 @@ static void lookupVisibleMemberDecls(
   lookupVisibleMemberDeclsImpl(BaseTy, overrideConsumer, CurrDC, LS, Reason,
                                TypeResolver, GSB, Visited);
 
+  if (hasDynamicMemberLookupAttribute(BaseTy)) {
+    llvm::DenseSet<DeclBaseName> knownMembers;
+    for (auto &kv : overrideConsumer.FoundDecls) {
+      knownMembers.insert(kv.first);
+    }
+    ShadowedKeyPathMembers dynamicConsumer(overrideConsumer, knownMembers);
+    lookupVisibleDynamicMemberLookupDecls(BaseTy, dynamicConsumer, CurrDC, LS,
+                                          Reason, TypeResolver, GSB, Visited);
+  }
+
   // Report the declarations we found to the real consumer.
   for (const auto &DeclAndReason : overrideConsumer.DeclsToReport)
     Consumer.foundDecl(DeclAndReason.D, DeclAndReason.Reason);

lib/Sema/TypeCheckAttr.cpp

@@ -14,20 +14,21 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "TypeChecker.h"
 #include "MiscDiagnostics.h"
 #include "TypeCheckType.h"
-#include "swift/AST/GenericSignatureBuilder.h"
+#include "TypeChecker.h"
 #include "swift/AST/ASTVisitor.h"
 #include "swift/AST/ClangModuleLoader.h"
 #include "swift/AST/DiagnosticsParse.h"
 #include "swift/AST/GenericEnvironment.h"
+#include "swift/AST/GenericSignatureBuilder.h"
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/NameLookupRequests.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/Types.h"
 #include "swift/Parse/Lexer.h"
+#include "swift/Sema/IDETypeChecking.h"
 #include "llvm/Support/Debug.h"
 
 using namespace swift;
@@ -1032,6 +1033,25 @@ bool swift::isValidKeyPathDynamicMemberLookup(SubscriptDecl *decl,
   return false;
 }
 
+Optional<Type>
+swift::getRootTypeOfKeypathDynamicMember(SubscriptDecl *subscript,
+                                         const DeclContext *DC) {
+  auto &TC = TypeChecker::createForContext(DC->getASTContext());
+
+  if (!isValidKeyPathDynamicMemberLookup(subscript, TC))
+    return None;
+
+  const auto *param = subscript->getIndices()->get(0);
+  auto keyPathType = param->getType()->getAs<BoundGenericType>();
+  if (!keyPathType)
+    return None;
+
+  assert(!keyPathType->getGenericArgs().empty() &&
+         "invalid keypath dynamic member");
+  auto rootType = keyPathType->getGenericArgs()[0];
+  return rootType;
+}
+
 /// The @dynamicMemberLookup attribute is only allowed on types that have at
 /// least one subscript member declared like this:
 ///

test/IDE/annotation.swift

@@ -339,3 +339,30 @@ test_arg_tuple2(p1:0,0)
 test_arg_tuple3(0,p2:0)
 // CHECK: <Func@[[@LINE-7]]:6>test_arg_tuple4</Func>(<Func@[[@LINE-7]]:6#p1>p1</Func>:0,<Func@[[@LINE-7]]:6#p2>p2</Func>:0)
 test_arg_tuple4(p1:0,p2:0)
+
+
+@dynamicMemberLookup
+struct Lens<T> {
+  var obj: T
+  init(_ obj: T) {
+    self.obj = obj
+  }
+  subscript<U>(dynamicMember member: WritableKeyPath<T, U>) -> Lens<U> {
+    get { return Lens<U>(obj[keyPath: member]) }
+    set { obj[keyPath: member] = newValue.obj }
+  }
+}
+struct Point {
+  var x: Int
+  var y: Int
+}
+struct Rectangle {
+  var topLeft: Point
+  var bottomRight: Point
+}
+func testDynamicMemberLookup(r: Lens<Rectangle>) {
+  _ = r.topLeft
+  // CHECK: _ = <Param@[[@LINE-2]]:30>r</Param>.<Var@[[@LINE-5]]:7>topLeft</Var>
+  _ = r.bottomRight.y
+  // CHECK: _ = <Param@[[@LINE-4]]:30>r</Param>.<Var@[[@LINE-6]]:7>bottomRight</Var>.<Var@[[@LINE-10]]:7>y</Var>
+}