Merge pull request swiftlang#35607 from xedin/type-vars-accepts-set

[AST] Adjust `TypeBase::getTypeVariables` to accept a set

Author: xedin

include/swift/AST/Types.h

@@ -632,7 +632,7 @@ class alignas(1 << TypeAlignInBits) TypeBase {
   ///
   /// \param typeVariables This vector is populated with the set of
   /// type variables referenced by this type.
-  void getTypeVariables(SmallVectorImpl<TypeVariableType *> &typeVariables);
+  void getTypeVariables(SmallPtrSetImpl<TypeVariableType *> &typeVariables);
 
   /// Determine whether this type is a type parameter, which is either a
   /// GenericTypeParamType or a DependentMemberType.

include/swift/Sema/Constraint.h

@@ -381,43 +381,45 @@ class Constraint final : public llvm::ilist_node<Constraint>,
   void *operator new(size_t) = delete;
 
   Constraint(ConstraintKind kind, ArrayRef<Constraint *> constraints,
-             ConstraintLocator *locator, ArrayRef<TypeVariableType *> typeVars);
+             ConstraintLocator *locator,
+             SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
   /// Construct a new constraint.
   Constraint(ConstraintKind kind, Type first, Type second,
              ConstraintLocator *locator,
-             ArrayRef<TypeVariableType *> typeVars);
+             SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
   /// Construct a new constraint.
   Constraint(ConstraintKind kind, Type first, Type second, Type third,
              ConstraintLocator *locator,
-             ArrayRef<TypeVariableType *> typeVars);
+             SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
   /// Construct a new member constraint.
   Constraint(ConstraintKind kind, Type first, Type second, DeclNameRef member,
              DeclContext *useDC, FunctionRefKind functionRefKind,
              ConstraintLocator *locator,
-             ArrayRef<TypeVariableType *> typeVars);
+             SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
   /// Construct a new value witness constraint.
   Constraint(ConstraintKind kind, Type first, Type second,
              ValueDecl *requirement, DeclContext *useDC,
              FunctionRefKind functionRefKind, ConstraintLocator *locator,
-             ArrayRef<TypeVariableType *> typeVars);
+             SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
   /// Construct a new overload-binding constraint, which might have a fix.
   Constraint(Type type, OverloadChoice choice, DeclContext *useDC,
              ConstraintFix *fix, ConstraintLocator *locator,
-             ArrayRef<TypeVariableType *> typeVars);
+             SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
   /// Construct a restricted constraint.
   Constraint(ConstraintKind kind, ConversionRestrictionKind restriction,
              Type first, Type second, ConstraintLocator *locator,
-             ArrayRef<TypeVariableType *> typeVars);
-  
+             SmallPtrSetImpl<TypeVariableType *> &typeVars);
+
   /// Construct a relational constraint with a fix.
   Constraint(ConstraintKind kind, ConstraintFix *fix, Type first, Type second,
-             ConstraintLocator *locator, ArrayRef<TypeVariableType *> typeVars);
+             ConstraintLocator *locator,
+             SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
   /// Retrieve the type variables buffer, for internal mutation.
   MutableArrayRef<TypeVariableType *> getTypeVariablesBuffer() {

lib/AST/Type.cpp

@@ -445,13 +445,13 @@ Type TypeBase::addCurriedSelfType(const DeclContext *dc) {
   return FunctionType::get({selfParam}, type);
 }
 
-void
-TypeBase::getTypeVariables(SmallVectorImpl<TypeVariableType *> &typeVariables) {
+void TypeBase::getTypeVariables(
+    SmallPtrSetImpl<TypeVariableType *> &typeVariables) {
   // If we know we don't have any type variables, we're done.
   if (hasTypeVariable()) {
     auto addTypeVariables = [&](Type type) -> bool {
       if (auto tv = dyn_cast<TypeVariableType>(type.getPointer())) {
-        typeVariables.push_back(tv);
+        typeVariables.insert(tv);
       }
 
       return false;

lib/Sema/CSBindings.cpp

@@ -830,9 +830,9 @@ PotentialBindings ConstraintSystem::inferBindingsFor(TypeVariableType *typeVar,
 /// \returns the type to bind to, if the binding is okay.
 static Optional<Type> checkTypeOfBinding(TypeVariableType *typeVar, Type type) {
   // If the type references the type variable, don't permit the binding.
-  SmallVector<TypeVariableType *, 4> referencedTypeVars;
+  SmallPtrSet<TypeVariableType *, 4> referencedTypeVars;
   type->getTypeVariables(referencedTypeVars);
-  if (count(referencedTypeVars, typeVar))
+  if (referencedTypeVars.count(typeVar))
     return None;
 
   // If type variable is not allowed to bind to `lvalue`,
@@ -940,21 +940,9 @@ PotentialBindings::inferFromRelational(Constraint *constraint) {
   if (type->getWithoutSpecifierType()
           ->lookThroughAllOptionalTypes()
           ->is<DependentMemberType>()) {
-    SmallVector<TypeVariableType *, 4> referencedVars;
-    type->getTypeVariables(referencedVars);
-
-    bool containsSelf = false;
-    for (auto *var : referencedVars) {
-      // Add all type variables encountered in the type except
-      // to the current type variable.
-      if (var != TypeVar) {
-        AdjacentVars.insert(var);
-        continue;
-      }
-
-      containsSelf = true;
-    }
+    type->getTypeVariables(AdjacentVars);
 
+    bool containsSelf = AdjacentVars.erase(TypeVar);
     // If inferred type doesn't contain the current type variable,
     // let's mark bindings as delayed until dependent member type
     // is resolved.
@@ -979,11 +967,8 @@ PotentialBindings::inferFromRelational(Constraint *constraint) {
   // Check whether we can perform this binding.
   if (auto boundType = checkTypeOfBinding(TypeVar, type)) {
     type = *boundType;
-    if (type->hasTypeVariable()) {
-      SmallVector<TypeVariableType *, 4> referencedVars;
-      type->getTypeVariables(referencedVars);
-      AdjacentVars.insert(referencedVars.begin(), referencedVars.end());
-    }
+    if (type->hasTypeVariable())
+      type->getTypeVariables(AdjacentVars);
   } else {
     auto *bindingTypeVar = type->getRValueType()->getAs<TypeVariableType>();
 

lib/Sema/CSGen.cpp

@@ -2471,7 +2471,7 @@ namespace {
       // the type variable for the pattern `i`.
       struct CollectVarRefs : public ASTWalker {
         ConstraintSystem &cs;
-        llvm::SmallVector<TypeVariableType *, 4> varRefs;
+        llvm::SmallPtrSet<TypeVariableType *, 4> varRefs;
         bool hasErrorExprs = false;
 
         CollectVarRefs(ConstraintSystem &cs) : cs(cs) { }
@@ -2532,10 +2532,12 @@ namespace {
       if (!inferredType || inferredType->hasError())
         return Type();
 
-      CS.addUnsolvedConstraint(
-          Constraint::create(CS, ConstraintKind::DefaultClosureType,
-                             closureType, inferredType, locator,
-                             collectVarRefs.varRefs));
+      SmallVector<TypeVariableType *, 4> referencedVars{
+          collectVarRefs.varRefs.begin(), collectVarRefs.varRefs.end()};
+
+      CS.addUnsolvedConstraint(Constraint::create(
+          CS, ConstraintKind::DefaultClosureType, closureType, inferredType,
+          locator, referencedVars));
 
       CS.setClosureType(closure, inferredType);
       return closureType;

lib/Sema/CSSimplify.cpp

@@ -10574,13 +10574,16 @@ ConstraintSystem::addArgumentConversionConstraintImpl(
   if (auto *argTypeVar = first->getAs<TypeVariableType>()) {
     if (argTypeVar->getImpl().isClosureType()) {
       // Extract any type variables present in the parameter's result builder.
-      SmallVector<TypeVariableType *, 4> typeVars;
+      SmallPtrSet<TypeVariableType *, 4> typeVars;
       if (auto builderTy = getOpenedResultBuilderTypeFor(*this, locator))
         builderTy->getTypeVariables(typeVars);
 
+      SmallVector<TypeVariableType *, 4> referencedVars{typeVars.begin(),
+                                                        typeVars.end()};
+
       auto *loc = getConstraintLocator(locator);
       addUnsolvedConstraint(
-          Constraint::create(*this, kind, first, second, loc, typeVars));
+          Constraint::create(*this, kind, first, second, loc, referencedVars));
       return SolutionKind::Solved;
     }
   }