Merge pull request #77322 from hamishknight/complete-cache

Add cache for `USRBasedTypeContext::typeRelation`

Author: hamishknight

include/swift/IDE/CodeCompletionResultType.h

@@ -207,6 +207,10 @@ class USRBasedTypeContext {
   /// allocated.
   const USRBasedTypeArena &Arena;
 
+  /// A cached set of type relations for this given type context.
+  mutable llvm::DenseMap<const USRBasedType *, CodeCompletionResultTypeRelation>
+      CachedTypeRelations;
+
   SmallVector<ContextualType, 4> ContextualTypes;
 
   /// See \c ExpectedTypeContext::ExpectedAttributeKinds.
@@ -251,12 +255,6 @@ class USRBasedType {
       : USR(USR), Supertypes(Supertypes),
         CustomAttributeKinds(CustomAttributeKinds) {}
 
-  /// Implementation of \c typeRelation. \p VisistedTypes keeps track which
-  /// types have already been visited.
-  CodeCompletionResultTypeRelation
-  typeRelationImpl(const USRBasedType *ResultType, const USRBasedType *VoidType,
-                   SmallPtrSetImpl<const USRBasedType *> &VisitedTypes) const;
-
 public:
   /// A null \c USRBasedType that's represented by an empty USR and has no
   /// supertypes.

lib/IDE/CodeCompletionResultType.cpp

@@ -88,27 +88,37 @@ USRBasedTypeContext::USRBasedTypeContext(const ExpectedTypeContext *TypeContext,
 
 TypeRelation
 USRBasedTypeContext::typeRelation(const USRBasedType *ResultType) const {
-  if (ExpectedCustomAttributeKinds) {
-    return ResultType->getCustomAttributeKinds() & ExpectedCustomAttributeKinds
-               ? TypeRelation::Convertible
-               : TypeRelation::Unrelated;
-  }
-  const USRBasedType *VoidType = Arena.getVoidType();
-  if (ResultType == VoidType) {
-    // Void is not convertible to anything and we don't report Void <-> Void
-    // identical matches (see USRBasedType::typeRelation). So we don't have to
-    // check anything if the result returns Void.
-    return TypeRelation::Unknown;
-  }
+  auto compute = [&]() -> TypeRelation {
+    if (ExpectedCustomAttributeKinds) {
+      return ResultType->getCustomAttributeKinds() &
+                     ExpectedCustomAttributeKinds
+                 ? TypeRelation::Convertible
+                 : TypeRelation::Unrelated;
+    }
+    const USRBasedType *VoidType = Arena.getVoidType();
+    if (ResultType == VoidType) {
+      // Void is not convertible to anything and we don't report Void <-> Void
+      // identical matches (see USRBasedType::typeRelation). So we don't have to
+      // check anything if the result returns Void.
+      return TypeRelation::Unknown;
+    }
 
-  TypeRelation Res = TypeRelation::Unknown;
-  for (auto &ContextualType : ContextualTypes) {
-    Res = std::max(Res, ContextualType.typeRelation(ResultType, VoidType));
-    if (Res == TypeRelation::MAX_VALUE) {
-      return Res; // We can't improve further
+    TypeRelation Res = TypeRelation::Unknown;
+    for (auto &ContextualType : ContextualTypes) {
+      Res = std::max(Res, ContextualType.typeRelation(ResultType, VoidType));
+      if (Res == TypeRelation::MAX_VALUE) {
+        return Res; // We can't improve further
+      }
     }
-  }
-  return Res;
+    return Res;
+  };
+  auto iter = CachedTypeRelations.find(ResultType);
+  if (iter != CachedTypeRelations.end())
+    return iter->second;
+
+  auto relation = compute();
+  CachedTypeRelations.insert({ResultType, relation});
+  return relation;
 }
 
 // MARK: - USRBasedTypeArena
@@ -122,38 +132,6 @@ const USRBasedType *USRBasedTypeArena::getVoidType() const { return VoidType; }
 
 // MARK: - USRBasedType
 
-TypeRelation USRBasedType::typeRelationImpl(
-    const USRBasedType *ResultType, const USRBasedType *VoidType,
-    SmallPtrSetImpl<const USRBasedType *> &VisitedTypes) const {
-
-  // `this` is the contextual type.
-  if (this == VoidType) {
-    // We don't report Void <-> Void matches because that would boost
-    // methods returning Void in e.g.
-    // func foo() { #^COMPLETE^# }
-    // because #^COMPLETE^# is implicitly returned. But that's not very
-    // helpful.
-    return TypeRelation::Unknown;
-  }
-  if (ResultType == this) {
-    return TypeRelation::Convertible;
-  }
-  for (const USRBasedType *Supertype : ResultType->getSupertypes()) {
-    if (!VisitedTypes.insert(Supertype).second) {
-      // Already visited this type.
-      continue;
-    }
-    if (this->typeRelation(Supertype, VoidType) >= TypeRelation::Convertible) {
-      return TypeRelation::Convertible;
-    }
-  }
-  // TypeRelation computation based on USRs is an under-approximation because we
-  // don't take into account generic conversions or retroactive conformance of
-  // library types. Hence, we can't know for sure that ResultType is not
-  // convertible to `this` type and thus can't return Unrelated or Invalid here.
-  return TypeRelation::Unknown;
-}
-
 const USRBasedType *USRBasedType::null(USRBasedTypeArena &Arena) {
   return USRBasedType::fromUSR(/*USR=*/"", /*Supertypes=*/{}, {}, Arena);
 }
@@ -337,8 +315,37 @@ const USRBasedType *USRBasedType::fromType(Type Ty, USRBasedTypeArena &Arena) {
 
 TypeRelation USRBasedType::typeRelation(const USRBasedType *ResultType,
                                         const USRBasedType *VoidType) const {
-  SmallPtrSet<const USRBasedType *, 4> VisitedTypes;
-  return this->typeRelationImpl(ResultType, VoidType, VisitedTypes);
+  // `this` is the contextual type.
+  if (this == VoidType) {
+    // We don't report Void <-> Void matches because that would boost
+    // methods returning Void in e.g.
+    // func foo() { #^COMPLETE^# }
+    // because #^COMPLETE^# is implicitly returned. But that's not very
+    // helpful.
+    return TypeRelation::Unknown;
+  }
+
+  SmallPtrSet<const USRBasedType *, 16> VisitedTypes;
+  SmallVector<const USRBasedType *, 16> Worklist;
+  Worklist.push_back(ResultType);
+  while (!Worklist.empty()) {
+    auto *CurrentType = Worklist.pop_back_val();
+    if (CurrentType == this)
+      return TypeRelation::Convertible;
+
+    for (const USRBasedType *Supertype : CurrentType->getSupertypes()) {
+      if (!VisitedTypes.insert(Supertype).second) {
+        // Already visited this type.
+        continue;
+      }
+      Worklist.push_back(Supertype);
+    }
+  }
+  // TypeRelation computation based on USRs is an under-approximation because we
+  // don't take into account generic conversions or retroactive conformance of
+  // library types. Hence, we can't know for sure that ResultType is not
+  // convertible to `this` type and thus can't return Unrelated or Invalid here.
+  return TypeRelation::Unknown;
 }
 
 // MARK: - USRBasedTypeContext