[code-completion] Add completion for keypath dynamic member lookup

Looks into the root type of the keypath to find additional members. This
does not currently map the type of the completion to the subscript's
return type.

rdar://49029126

Author: benlangmuir

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/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"
 
@@ -3563,6 +3564,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();
@@ -3937,8 +3944,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));
@@ -4030,7 +4037,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

@@ -621,6 +621,41 @@ 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) {
+
+  if (!hasDynamicMemberLookupAttribute(baseType))
+    return;
+
+  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);
+    if (auto memberType = rootType->subst(subs))
+      lookupVisibleMemberDeclsImpl(memberType, consumer, dc, LS, reason,
+                                   typeResolver, GSB, visited);
+  }
+}
+
 namespace {
 
 struct FoundDeclTy {
@@ -822,6 +857,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,
@@ -839,6 +886,14 @@ static void lookupVisibleMemberDecls(
   lookupVisibleMemberDeclsImpl(BaseTy, overrideConsumer, CurrDC, LS, Reason,
                                TypeResolver, GSB, Visited);
 
+  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:
 ///