Merge pull request #40269 from AnthonyLatsis/assoc-inference-system

AssociatedTypeInference: Initial refactoring of abstract type witness inference

Author: AnthonyLatsis

include/swift/AST/ASTContext.h

@@ -1167,7 +1167,7 @@ class ASTContext final {
   /// conformance itself, along with a bit indicating whether this diagnostic
   /// produces an error.
   struct DelayedConformanceDiag {
-    ValueDecl *Requirement;
+    const ValueDecl *Requirement;
     std::function<void()> Callback;
     bool IsError;
   };

include/swift/AST/Decl.h

@@ -3495,7 +3495,9 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
          SmallVectorImpl<ProtocolConformance *> &conformances) const;
 
   /// Retrieve all of the protocols that this nominal type conforms to.
-  SmallVector<ProtocolDecl *, 2> getAllProtocols() const;
+  ///
+  /// \param sorted Whether to sort the protocols in canonical order.
+  SmallVector<ProtocolDecl *, 2> getAllProtocols(bool sorted = false) const;
 
   /// Retrieve all of the protocol conformances for this nominal type.
   SmallVector<ProtocolConformance *, 2> getAllConformances(

include/swift/Basic/LangOptions.h

@@ -530,6 +530,9 @@ namespace swift {
     RequirementMachineMode RequirementMachineInferredSignatures =
         RequirementMachineMode::Disabled;
 
+    /// Enables dumping type witness systems from associated type inference.
+    bool DumpTypeWitnessSystems = false;
+
     /// Sets the target we are building for and updates platform conditions
     /// to match.
     ///

include/swift/Option/FrontendOptions.td

@@ -352,6 +352,9 @@ def requirement_machine_max_concrete_nesting : Joined<["-"], "requirement-machin
   Flags<[FrontendOption, HelpHidden, DoesNotAffectIncrementalBuild]>,
   HelpText<"Set the maximum concrete type nesting depth before giving up">;
 
+def dump_type_witness_systems : Flag<["-"], "dump-type-witness-systems">,
+  HelpText<"Enables dumping type witness systems from associated type inference">;
+
 def debug_generic_signatures : Flag<["-"], "debug-generic-signatures">,
   HelpText<"Debug generic signatures">;
 

lib/AST/ConformanceLookupTable.cpp

@@ -1061,8 +1061,8 @@ void ConformanceLookupTable::lookupConformances(
 }
 
 void ConformanceLookupTable::getAllProtocols(
-       NominalTypeDecl *nominal,
-       SmallVectorImpl<ProtocolDecl *> &scratch) {
+    NominalTypeDecl *nominal, SmallVectorImpl<ProtocolDecl *> &scratch,
+    bool sorted) {
   // We need to expand all implied conformances to find the complete
   // set of protocols to which this nominal type conforms.
   updateLookupTable(nominal, ConformanceStage::ExpandedImplied);
@@ -1075,7 +1075,9 @@ void ConformanceLookupTable::getAllProtocols(
     scratch.push_back(conformance.first);
   }
 
-  // FIXME: sort the protocols in some canonical order?
+  if (sorted) {
+    llvm::array_pod_sort(scratch.begin(), scratch.end(), TypeDecl::compare);
+  }
 }
 
 int ConformanceLookupTable::compareProtocolConformances(

lib/AST/ConformanceLookupTable.h

@@ -451,9 +451,12 @@ class ConformanceLookupTable : public ASTAllocated<ConformanceLookupTable> {
                           SmallVectorImpl<ConformanceDiagnostic> *diagnostics);
 
   /// Retrieve the complete set of protocols to which this nominal
-  /// type conforms.
+  /// type conforms (if the set contains a protocol, the same is true for any
+  /// inherited protocols).
+  ///
+  /// \param sorted Whether to sort the protocols in canonical order.
   void getAllProtocols(NominalTypeDecl *nominal,
-                       SmallVectorImpl<ProtocolDecl *> &scratch);
+                       SmallVectorImpl<ProtocolDecl *> &scratch, bool sorted);
 
   /// Retrieve the complete set of protocol conformances for this
   /// nominal type.

lib/AST/ProtocolConformance.cpp

@@ -1316,11 +1316,12 @@ bool NominalTypeDecl::lookupConformance(
            conformances);
 }
 
-SmallVector<ProtocolDecl *, 2> NominalTypeDecl::getAllProtocols() const {
+SmallVector<ProtocolDecl *, 2>
+NominalTypeDecl::getAllProtocols(bool sorted) const {
   prepareConformanceTable();
   SmallVector<ProtocolDecl *, 2> result;
-  ConformanceTable->getAllProtocols(const_cast<NominalTypeDecl *>(this),
-                                    result);
+  ConformanceTable->getAllProtocols(const_cast<NominalTypeDecl *>(this), result,
+                                    sorted);
   return result;
 }
 

lib/Frontend/CompilerInvocation.cpp

@@ -959,6 +959,8 @@ static bool ParseLangArgs(LangOptions &Opts, ArgList &Args,
     }
   }
 
+  Opts.DumpTypeWitnessSystems = Args.hasArg(OPT_dump_type_witness_systems);
+
   return HadError || UnsupportedOS || UnsupportedArch;
 }
 

lib/Sema/TypeCheckProtocol.cpp

@@ -5383,8 +5383,8 @@ void swift::diagnoseConformanceFailure(Type T,
 }
 
 void ConformanceChecker::diagnoseOrDefer(
-       ValueDecl *requirement, bool isError,
-       std::function<void(NormalProtocolConformance *)> fn) {
+    const ValueDecl *requirement, bool isError,
+    std::function<void(NormalProtocolConformance *)> fn) {
   if (isError)
     Conformance->setInvalid();
 

lib/Sema/TypeCheckProtocol.h

@@ -100,61 +100,6 @@ CheckTypeWitnessResult checkTypeWitness(Type type,
                                         const NormalProtocolConformance *Conf,
                                         SubstOptions options = None);
 
-/// Describes the means of inferring an abstract type witness.
-enum class AbstractTypeWitnessKind : uint8_t {
-  /// The type witness was inferred via a same-type-to-concrete constraint
-  /// in a protocol requirement signature.
-  Fixed,
-
-  /// The type witness was inferred via a defaulted associated type.
-  Default,
-
-  /// The type witness was inferred to a generic parameter of the
-  /// conforming type.
-  GenericParam,
-};
-
-/// A type witness inferred without the aid of a specific potential
-/// value witness.
-class AbstractTypeWitness {
-  AbstractTypeWitnessKind Kind;
-  AssociatedTypeDecl *AssocType;
-  Type TheType;
-
-  /// When this is a default type witness, the declaration responsible for it.
-  /// May not necessarilly match \c AssocType.
-  AssociatedTypeDecl *DefaultedAssocType;
-
-  AbstractTypeWitness(AbstractTypeWitnessKind Kind,
-                      AssociatedTypeDecl *AssocType, Type TheType,
-                      AssociatedTypeDecl *DefaultedAssocType)
-      : Kind(Kind), AssocType(AssocType), TheType(TheType),
-        DefaultedAssocType(DefaultedAssocType) {
-    assert(AssocType && TheType);
-  }
-
-public:
-  static AbstractTypeWitness forFixed(AssociatedTypeDecl *assocType, Type type);
-
-  static AbstractTypeWitness forDefault(AssociatedTypeDecl *assocType,
-                                        Type type,
-                                        AssociatedTypeDecl *defaultedAssocType);
-
-  static AbstractTypeWitness forGenericParam(AssociatedTypeDecl *assocType,
-                                             Type type);
-
-public:
-  AbstractTypeWitnessKind getKind() const { return Kind; }
-
-  AssociatedTypeDecl *getAssocType() const { return AssocType; }
-
-  Type getType() const { return TheType; }
-
-  AssociatedTypeDecl *getDefaultedAssocType() const {
-    return DefaultedAssocType;
-  }
-};
-
 /// The set of associated types that have been inferred by matching
 /// the given value witness to its corresponding requirement.
 struct InferredAssociatedTypesByWitness {
@@ -855,9 +800,8 @@ class ConformanceChecker : public WitnessChecker {
   ///
   /// \param fn A function to call to emit the actual diagnostic. If
   /// diagnostics are being deferred,
-  void diagnoseOrDefer(
-         ValueDecl *requirement, bool isError,
-         std::function<void(NormalProtocolConformance *)> fn);
+  void diagnoseOrDefer(const ValueDecl *requirement, bool isError,
+                       std::function<void(NormalProtocolConformance *)> fn);
 
   ArrayRef<MissingWitness> getLocalMissingWitness() {
     return GlobalMissingWitnesses.getArrayRef().
@@ -929,6 +873,136 @@ class ConformanceChecker : public WitnessChecker {
                     llvm::function_ref<bool(AbstractFunctionDecl *)>predicate);
 };
 
+/// A system for recording and probing the intergrity of a type witness solution
+/// for a set of unresolved associated type declarations.
+///
+/// Right now can reason only about abstract type witnesses, i.e., same-type
+/// constraints, default type definitions, and bindings to generic parameters.
+class TypeWitnessSystem final {
+  /// Equivalence classes are used on demand to express equivalences between
+  /// witness candidates and reflect changes to resolved types across their
+  /// members.
+  class EquivalenceClass final {
+    /// The pointer:
+    /// - The resolved type for witness candidates belonging to this equivalence
+    ///   class. The resolved type may be a type parameter, but cannot directly
+    ///   pertain to a name variable in the owning system; instead, witness
+    ///   candidates that should resolve to the same type share an equivalence
+    ///   class.
+    /// The int:
+    /// - A flag indicating whether the resolved type is ambiguous. When set,
+    ///   the resolved type is null.
+    llvm::PointerIntPair<Type, 1, bool> ResolvedTyAndIsAmbiguous;
+
+  public:
+    EquivalenceClass(Type ty) : ResolvedTyAndIsAmbiguous(ty, false) {}
+
+    EquivalenceClass(const EquivalenceClass &) = delete;
+    EquivalenceClass(EquivalenceClass &&) = delete;
+    EquivalenceClass &operator=(const EquivalenceClass &) = delete;
+    EquivalenceClass &operator=(EquivalenceClass &&) = delete;
+
+    Type getResolvedType() const {
+      return ResolvedTyAndIsAmbiguous.getPointer();
+    }
+    void setResolvedType(Type ty);
+
+    bool isAmbiguous() const {
+      return ResolvedTyAndIsAmbiguous.getInt();
+    }
+    void setAmbiguous() {
+      ResolvedTyAndIsAmbiguous = {nullptr, true};
+    }
+  };
+
+  /// A type witness candidate for a name variable.
+  struct TypeWitnessCandidate final {
+    /// The defaulted associated type declaration correlating with this
+    /// candidate, if present.
+    const AssociatedTypeDecl *DefaultedAssocType;
+
+    /// The equivalence class of this candidate.
+    EquivalenceClass *EquivClass;
+  };
+
+  /// The set of equivalence classes in the system.
+  llvm::SmallPtrSet<EquivalenceClass *, 4> EquivalenceClasses;
+
+  /// The mapping from name variables (the names of unresolved associated
+  /// type declarations) to their corresponding type witness candidates.
+  llvm::SmallDenseMap<Identifier, TypeWitnessCandidate, 4> TypeWitnesses;
+
+public:
+  TypeWitnessSystem(ArrayRef<AssociatedTypeDecl *> assocTypes);
+  ~TypeWitnessSystem();
+
+  TypeWitnessSystem(const TypeWitnessSystem &) = delete;
+  TypeWitnessSystem(TypeWitnessSystem &&) = delete;
+  TypeWitnessSystem &operator=(const TypeWitnessSystem &) = delete;
+  TypeWitnessSystem &operator=(TypeWitnessSystem &&) = delete;
+
+  /// Get the resolved type witness for the associated type with the given name.
+  Type getResolvedTypeWitness(Identifier name) const;
+  bool hasResolvedTypeWitness(Identifier name) const;
+
+  /// Get the defaulted associated type relating to the resolved type witness
+  /// for the associated type with the given name, if present.
+  const AssociatedTypeDecl *getDefaultedAssocType(Identifier name) const;
+
+  /// Record a type witness for the given associated type name.
+  ///
+  /// \note This need not lead to the resolution of a type witness, e.g.
+  /// an associated type may be defaulted to another.
+  void addTypeWitness(Identifier name, Type type);
+
+  /// Record a default type witness.
+  ///
+  /// \param defaultedAssocType The specific associated type declaration that
+  /// defines the given default type.
+  ///
+  /// \note This need not lead to the resolution of a type witness.
+  void addDefaultTypeWitness(Type type,
+                             const AssociatedTypeDecl *defaultedAssocType);
+
+  /// Record the given same-type requirement, if regarded of interest to
+  /// the system.
+  ///
+  /// \note This need not lead to the resolution of a type witness.
+  void addSameTypeRequirement(const Requirement &req);
+
+  void dump(llvm::raw_ostream &out,
+            const NormalProtocolConformance *conformance) const;
+
+private:
+  /// Form an equivalence between the given name variables.
+  void addEquivalence(Identifier name1, Identifier name2);
+
+  /// Merge \p equivClass2 into \p equivClass1.
+  ///
+  /// \note This will delete \p equivClass2 after migrating its members to
+  /// \p equivClass1.
+  void mergeEquivalenceClasses(EquivalenceClass *equivClass1,
+                               const EquivalenceClass *equivClass2);
+
+  /// The result of comparing two resolved types targeting a single equivalence
+  /// class, in terms of their relative impact on solving the system.
+  enum class ResolvedTypeComparisonResult {
+    /// The first resolved type is a better choice than the second one.
+    Better,
+
+    /// The first resolved type is an equivalent or worse choice than the
+    /// second one.
+    EquivalentOrWorse,
+
+    /// Both resolved types are concrete and mutually exclusive.
+    Ambiguity
+  };
+
+  /// Compare the given resolved types as targeting a single equivalence class,
+  /// in terms of the their relative impact on solving the system.
+  static ResolvedTypeComparisonResult compareResolvedTypes(Type ty1, Type ty2);
+};
+
 /// Captures the state needed to infer associated types.
 class AssociatedTypeInference {
   /// The type checker we'll need to validate declarations etc.
@@ -956,7 +1030,7 @@ class AssociatedTypeInference {
     typeWitnesses;
 
   /// Information about a failed, defaulted associated type.
-  AssociatedTypeDecl *failedDefaultedAssocType = nullptr;
+  const AssociatedTypeDecl *failedDefaultedAssocType = nullptr;
   Type failedDefaultedWitness;
   CheckTypeWitnessResult failedDefaultedResult;
 
@@ -1005,23 +1079,20 @@ class AssociatedTypeInference {
     ConformanceChecker &checker,
     const llvm::SetVector<AssociatedTypeDecl *> &assocTypes);
 
-  /// Compute a "fixed" type witness for an associated type, e.g.,
-  /// if the refined protocol requires it to be equivalent to some other type.
-  Type computeFixedTypeWitness(AssociatedTypeDecl *assocType);
-
   /// Compute the default type witness from an associated type default,
   /// if there is one.
-  Optional<AbstractTypeWitness>
-  computeDefaultTypeWitness(AssociatedTypeDecl *assocType);
+  Optional<std::pair<AssociatedTypeDecl *, Type>>
+  computeDefaultTypeWitness(AssociatedTypeDecl *assocType) const;
 
   /// Compute the "derived" type witness for an associated type that is
   /// known to the compiler.
   std::pair<Type, TypeDecl *>
   computeDerivedTypeWitness(AssociatedTypeDecl *assocType);
 
-  /// Compute a type witness without using a specific potential witness.
-  Optional<AbstractTypeWitness>
-  computeAbstractTypeWitness(AssociatedTypeDecl *assocType);
+  /// Collect abstract type witnesses and feed them to the given system.
+  void collectAbstractTypeWitnesses(
+      TypeWitnessSystem &system,
+      ArrayRef<AssociatedTypeDecl *> unresolvedAssocTypes) const;
 
   /// Substitute the current type witnesses into the given interface type.
   Type substCurrentTypeWitnesses(Type type);
@@ -1040,14 +1111,12 @@ class AssociatedTypeInference {
   /// requirements of the given constrained extension.
   bool checkConstrainedExtension(ExtensionDecl *ext);
 
-  /// Validate the current tentative solution represented by \p typeWitnesses
-  /// and attempt to resolve abstract type witnesses for associated types that
-  /// could not be inferred otherwise.
+  /// Attempt to infer abstract type witnesses for the given set of associated
+  /// types.
   ///
   /// \returns \c nullptr, or the associated type that failed.
-  AssociatedTypeDecl *
-  completeSolution(ArrayRef<AssociatedTypeDecl *> unresolvedAssocTypes,
-                   unsigned reqDepth);
+  AssociatedTypeDecl *inferAbstractTypeWitnesses(
+      ArrayRef<AssociatedTypeDecl *> unresolvedAssocTypes, unsigned reqDepth);
 
   /// Top-level operation to find solutions for the given unresolved
   /// associated types.