IDE+Evaluator: refactor IDE type-checking utilities for subscript decl into requests

We used to have a function getRootAndResultTypeOfKeypathDynamicMember to return
both the root and result type of a subscript. This patch splits the function
into two functions returning root type and result type respectively. It also refactors
the implementation into the evaluator model.

Author: nkcsgexi

include/swift/AST/ASTTypeIDZone.def

@@ -20,6 +20,7 @@ SWIFT_TYPEID_NAMED(ValueDecl *, ValueDecl)
 SWIFT_TYPEID_NAMED(ProtocolDecl *, ProtocolDecl)
 SWIFT_TYPEID_NAMED(Decl *, Decl)
 SWIFT_TYPEID(Type)
+SWIFT_TYPEID(TypePair)
 SWIFT_TYPEID(PropertyWrapperBackingPropertyInfo)
 SWIFT_TYPEID(PropertyWrapperTypeInfo)
 SWIFT_TYPEID_NAMED(Optional<PropertyWrapperMutability>, PropertyWrapperMutability)

include/swift/AST/ASTTypeIDs.h

@@ -29,6 +29,7 @@ class Type;
 class VarDecl;
 class TypeAliasDecl;
 class Type;
+struct TypePair;
 
 #define SWIFT_AST_TYPEID_ZONE 1
 

include/swift/Sema/IDETypeChecking.h

@@ -234,10 +234,10 @@ namespace swift {
   /// 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<std::pair<Type, Type>>
-  getRootAndResultTypeOfKeypathDynamicMember(SubscriptDecl *subscript,
-                                             const DeclContext *DC);
+  Type getRootTypeOfKeypathDynamicMember(SubscriptDecl *subscript);
+
+  Type getResultTypeOfKeypathDynamicMember(SubscriptDecl *subscript);
+
   /// Collect all the protocol requirements that a given declaration can
   ///   provide default implementations for. VD is a declaration in extension
   ///   declaration. Scratch is the buffer to collect those protocol

include/swift/Sema/IDETypeCheckingRequestIDZone.def

@@ -16,3 +16,5 @@
 //===----------------------------------------------------------------------===//
 SWIFT_TYPEID(IsDeclApplicableRequest)
 SWIFT_TYPEID(TypeRelationCheckRequest)
+SWIFT_TYPEID(RootAndResultTypeOfKeypathDynamicMemberRequest)
+SWIFT_TYPEID(RootTypeOfKeypathDynamicMemberRequest)

include/swift/Sema/IDETypeCheckingRequests.h

@@ -92,30 +92,57 @@ enum class TypeRelation: uint8_t {
   ConvertTo,
 };
 
+struct TypePair {
+  Type FirstTy;
+  Type SecondTy;
+  TypePair(Type FirstTy, Type SecondTy): FirstTy(FirstTy), SecondTy(SecondTy) {}
+  TypePair(): TypePair(Type(), Type()) {}
+  friend llvm::hash_code hash_value(const TypePair &TI) {
+    return hash_combine(hash_value(TI.FirstTy.getPointer()),
+                        hash_value(TI.SecondTy.getPointer()));
+  }
+
+  friend bool operator==(const TypePair &lhs,
+                         const TypePair &rhs) {
+    return lhs.FirstTy.getPointer() == rhs.FirstTy.getPointer() &&
+      lhs.SecondTy.getPointer() == rhs.SecondTy.getPointer();
+  }
+
+  friend bool operator!=(const TypePair &lhs,
+                         const TypePair &rhs) {
+    return !(lhs == rhs);
+  }
+
+  friend void simple_display(llvm::raw_ostream &out,
+                             const TypePair &owner) {
+    out << "<";
+    simple_display(out, owner.FirstTy);
+    out << ", ";
+    simple_display(out, owner.SecondTy);
+    out << ">";
+  }
+};
+
 struct TypeRelationCheckInput {
   DeclContext *DC;
-  Type FirstType;
-  Type SecondType;
+  TypePair Pair;
   TypeRelation Relation;
   bool OpenArchetypes;
 
   TypeRelationCheckInput(DeclContext *DC, Type FirstType, Type SecondType,
                          TypeRelation Relation, bool OpenArchetypes = true):
-    DC(DC), FirstType(FirstType), SecondType(SecondType), Relation(Relation),
+    DC(DC), Pair(FirstType, SecondType), Relation(Relation),
     OpenArchetypes(OpenArchetypes) {}
 
   friend llvm::hash_code hash_value(const TypeRelationCheckInput &TI) {
-    return hash_combine(hash_value(TI.FirstType.getPointer()),
-                        hash_value(TI.SecondType.getPointer()),
+    return hash_combine(hash_value(TI.Pair),
                         hash_value(TI.Relation),
                         hash_value(TI.OpenArchetypes));
   }
 
   friend bool operator==(const TypeRelationCheckInput &lhs,
                          const TypeRelationCheckInput &rhs) {
-    return lhs.FirstType.getPointer() == rhs.FirstType.getPointer() &&
-      lhs.SecondType.getPointer() == rhs.SecondType.getPointer() &&
-      lhs.Relation == rhs.Relation &&
+    return lhs.Pair == rhs.Pair && lhs.Relation == rhs.Relation &&
       lhs.OpenArchetypes == rhs.OpenArchetypes;
   }
 
@@ -127,7 +154,7 @@ struct TypeRelationCheckInput {
   friend void simple_display(llvm::raw_ostream &out,
                                const TypeRelationCheckInput &owner) {
     out << "Check if ";
-    simple_display(out, owner.FirstType);
+    simple_display(out, owner.Pair);
     out << " is ";
     switch(owner.Relation) {
 #define CASE(NAME) case TypeRelation::NAME: out << #NAME << " "; break;
@@ -136,7 +163,6 @@ struct TypeRelationCheckInput {
     CASE(ConvertTo)
 #undef CASE
     }
-    simple_display(out, owner.SecondType);
   }
 };
 
@@ -161,6 +187,53 @@ class TypeRelationCheckRequest:
   SourceLoc getNearestLoc() const { return SourceLoc(); };
 };
 
+//----------------------------------------------------------------------------//
+// RootAndResultTypeOfKeypathDynamicMemberRequest
+//----------------------------------------------------------------------------//
+class RootAndResultTypeOfKeypathDynamicMemberRequest:
+    public SimpleRequest<RootAndResultTypeOfKeypathDynamicMemberRequest,
+                         TypePair(SubscriptDecl*),
+                         CacheKind::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  llvm::Expected<TypePair> evaluate(Evaluator &evaluator, SubscriptDecl* SD) const;
+
+public:
+  // Caching
+  bool isCached() const { return true; }
+  // Source location
+  SourceLoc getNearestLoc() const { return SourceLoc(); };
+};
+
+class RootTypeOfKeypathDynamicMemberRequest:
+    public SimpleRequest<RootTypeOfKeypathDynamicMemberRequest,
+                         Type(SubscriptDecl*),
+                         /*Cached in the request above*/CacheKind::Uncached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  llvm::Expected<Type> evaluate(Evaluator &evaluator, SubscriptDecl* SD) const {
+    return evaluateOrDefault(SD->getASTContext().evaluator,
+      RootAndResultTypeOfKeypathDynamicMemberRequest{SD}, TypePair()).
+        FirstTy;
+  }
+
+public:
+  // Caching
+  bool isCached() const { return true; }
+  // Source location
+  SourceLoc getNearestLoc() const { return SourceLoc(); };
+};
+
 /// The zone number for the IDE.
 #define SWIFT_IDE_TYPE_CHECK_REQUESTS_TYPEID_ZONE 97
 #define SWIFT_TYPEID_ZONE SWIFT_IDE_TYPE_CHECK_REQUESTS_TYPEID_ZONE

lib/IDE/CodeCompletion.cpp

@@ -1994,18 +1994,14 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
       //     τ_1_0(U) => U }
       auto subs = keyPathInfo.baseType->getMemberSubstitutions(SD);
 
-      // Extract the root and result type of the KeyPath type in the parameter.
-      // i.e. 'T' and 'U'
-      auto rootAndResult =
-          getRootAndResultTypeOfKeypathDynamicMember(SD, CurrDeclContext);
-
       // If the keyPath result type has type parameters, that might affect the
       // subscript result type.
-      auto keyPathResultTy = rootAndResult->second->mapTypeOutOfContext();
+      auto keyPathResultTy = getResultTypeOfKeypathDynamicMember(SD)->
+        mapTypeOutOfContext();
       if (keyPathResultTy->hasTypeParameter()) {
-        auto keyPathRootTy =
-            rootAndResult->first.subst(QueryTypeSubstitutionMap{subs},
-                                       LookUpConformanceInModule(CurrModule));
+        auto keyPathRootTy = getRootTypeOfKeypathDynamicMember(SD).
+          subst(QueryTypeSubstitutionMap{subs},
+                LookUpConformanceInModule(CurrModule));
 
         // The result type of the VD.
         // i.e. 'Circle.center' => 'Point'.

lib/IDE/IDETypeChecking.cpp

@@ -771,3 +771,14 @@ bool swift::isConvertibleTo(Type T1, Type T2, bool openArchetypes,
     TypeRelationCheckRequest(TypeRelationCheckInput(&DC, T1, T2,
       TypeRelation::ConvertTo, openArchetypes)), false);
 }
+
+Type swift::getRootTypeOfKeypathDynamicMember(SubscriptDecl *SD) {
+  return evaluateOrDefault(SD->getASTContext().evaluator,
+    RootTypeOfKeypathDynamicMemberRequest{SD}, Type());
+}
+
+Type swift::getResultTypeOfKeypathDynamicMember(SubscriptDecl *SD) {
+  return evaluateOrDefault(SD->getASTContext().evaluator,
+    RootAndResultTypeOfKeypathDynamicMemberRequest{SD}, TypePair()).
+      SecondTy;
+}

lib/Sema/IDETypeCheckingRequests.cpp

@@ -105,7 +105,7 @@ TypeRelationCheckRequest::evaluate(Evaluator &evaluator,
   Optional<constraints::ConstraintKind> CKind;
   switch (Owner.Relation) {
   case TypeRelation::EqualTo:
-    return Owner.FirstType->isEqual(Owner.SecondType);
+    return Owner.Pair.FirstTy->isEqual(Owner.Pair.SecondTy);
   case TypeRelation::PossiblyEqualTo:
     CKind = constraints::ConstraintKind::Bind;
     break;
@@ -114,6 +114,24 @@ TypeRelationCheckRequest::evaluate(Evaluator &evaluator,
     break;
   }
   assert(CKind.hasValue());
-  return canSatisfy(Owner.FirstType, Owner.SecondType, Owner.OpenArchetypes,
+  return canSatisfy(Owner.Pair.FirstTy, Owner.Pair.SecondTy, Owner.OpenArchetypes,
                     *CKind, Owner.DC);
 }
+
+llvm::Expected<TypePair>
+RootAndResultTypeOfKeypathDynamicMemberRequest::evaluate(Evaluator &evaluator,
+                                              SubscriptDecl *subscript) const {
+  auto &TC = TypeChecker::createForContext(subscript->getASTContext());
+
+  if (!isValidKeyPathDynamicMemberLookup(subscript, TC))
+    return TypePair();
+
+  const auto *param = subscript->getIndices()->get(0);
+  auto keyPathType = param->getType()->getAs<BoundGenericType>();
+  if (!keyPathType)
+    return TypePair();
+  auto genericArgs = keyPathType->getGenericArgs();
+  assert(!genericArgs.empty() && genericArgs.size() == 2 &&
+         "invalid keypath dynamic member");
+  return TypePair(genericArgs[0], genericArgs[1]);
+}

lib/Sema/LookupVisibleDecls.cpp

@@ -977,11 +977,10 @@ static void lookupVisibleDynamicMemberLookupDecls(
     if (!subscript)
       continue;
 
-    auto rootAndResult =
-        getRootAndResultTypeOfKeypathDynamicMember(subscript, dc);
-    if (!rootAndResult)
+    auto rootType = evaluateOrDefault(subscript->getASTContext().evaluator,
+      RootTypeOfKeypathDynamicMemberRequest{subscript}, Type());
+    if (rootType.isNull())
       continue;
-    auto rootType = rootAndResult->first;
 
     auto subs =
         baseType->getMemberSubstitutionMap(dc->getParentModule(), subscript);

lib/Sema/TypeCheckAttr.cpp

@@ -972,24 +972,6 @@ bool swift::isValidKeyPathDynamicMemberLookup(SubscriptDecl *decl,
   return false;
 }
 
-Optional<std::pair<Type, Type>>
-swift::getRootAndResultTypeOfKeypathDynamicMember(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;
-  auto genericArgs = keyPathType->getGenericArgs();
-  assert(!genericArgs.empty() && genericArgs.size() == 2 &&
-         "invalid keypath dynamic member");
-  return std::pair<Type, Type>{genericArgs[0], genericArgs[1]};
-}
-
 /// The @dynamicMemberLookup attribute is only allowed on types that have at
 /// least one subscript member declared like this:
 ///