RequirementMachine: Generalize ConcreteTypeWitness to TypeWitness

slavapestov · slavapestov · commit 4c83fe851d35 · 2021-12-17T10:04:52.000-05:00
diff --git a/lib/AST/RequirementMachine/PropertyUnification.cpp b/lib/AST/RequirementMachine/PropertyUnification.cpp
@@ -636,48 +636,53 @@ void PropertyMap::concretizeTypeWitnessInConformance(
   }
 }
 
-RewriteSystem::ConcreteTypeWitness::ConcreteTypeWitness(
-    Symbol concreteConformance,
-    Symbol assocType,
-    Symbol concreteType)
-  : ConcreteConformance(concreteConformance),
-    AssocType(assocType),
-    ConcreteType(concreteType) {
-  assert(concreteConformance.getKind() == Symbol::Kind::ConcreteConformance);
-  assert(assocType.getKind() == Symbol::Kind::AssociatedType);
-  assert(concreteType.getKind() == Symbol::Kind::ConcreteType);
-  assert(assocType.getProtocols().size() == 1);
-  assert(assocType.getProtocols()[0] == concreteConformance.getProtocol());
+RewriteSystem::TypeWitness::TypeWitness(
+    Term lhs, llvm::PointerUnion<Symbol, Term> rhs)
+  : LHS(lhs), RHS(rhs) {
+  assert(LHS.size() >= 2);
+  assert(getConcreteConformance().getKind() ==
+         Symbol::Kind::ConcreteConformance);
+  assert(getAssocType().getKind() == Symbol::Kind::AssociatedType);
+  if (RHS.is<Symbol>())
+    assert(RHS.get<Symbol>().getKind() == Symbol::Kind::ConcreteType);
+  assert(getAssocType().getProtocols().size() == 1);
+  assert(getAssocType().getProtocols()[0] ==
+         getConcreteConformance().getProtocol());
 }
 
 bool swift::rewriting::operator==(
-    const RewriteSystem::ConcreteTypeWitness &lhs,
-    const RewriteSystem::ConcreteTypeWitness &rhs) {
-  return (lhs.ConcreteConformance == rhs.ConcreteConformance &&
-          lhs.AssocType == rhs.AssocType &&
-          lhs.ConcreteType == rhs.ConcreteType);
+    const RewriteSystem::TypeWitness &lhs,
+    const RewriteSystem::TypeWitness &rhs) {
+  return (lhs.LHS == rhs.LHS &&
+          lhs.RHS == rhs.RHS);
 }
 
-unsigned RewriteSystem::recordConcreteTypeWitness(
-    RewriteSystem::ConcreteTypeWitness witness) {
-  auto key = std::make_pair(witness.ConcreteConformance,
-                            witness.AssocType);
-  unsigned index = ConcreteTypeWitnesses.size();
-  auto inserted = ConcreteTypeWitnessMap.insert(std::make_pair(key, index));
+void RewriteSystem::TypeWitness::dump(llvm::raw_ostream &out) const {
+  out << "Subject type: " << LHS << "\n";
+  if (RHS.is<Symbol>())
+    out << "Concrete type witness: " << RHS.get<Symbol>() << "\n";
+  else
+    out << "Abstract type witness: " << RHS.get<Term>() << "\n";
+}
+
+unsigned RewriteSystem::recordTypeWitness(
+    RewriteSystem::TypeWitness witness) {
+  unsigned index = TypeWitnesses.size();
+  auto inserted = TypeWitnessMap.insert(std::make_pair(witness.LHS, index));
 
   if (!inserted.second) {
     index = inserted.first->second;
   } else {
-    ConcreteTypeWitnesses.push_back(witness);
+    TypeWitnesses.push_back(witness);
   }
 
-  assert(ConcreteTypeWitnesses[index] == witness);
+  assert(TypeWitnesses[index] == witness);
   return index;
 }
 
-const RewriteSystem::ConcreteTypeWitness &
-RewriteSystem::getConcreteTypeWitness(unsigned index) const {
-  return ConcreteTypeWitnesses[index];
+const RewriteSystem::TypeWitness &
+RewriteSystem::getTypeWitness(unsigned index) const {
+  return TypeWitnesses[index];
 }
 
 /// Given the key of a property bag known to have \p concreteType,
@@ -737,13 +742,9 @@ MutableTerm PropertyMap::computeConstraintTermForTypeWitness(
   auto typeWitnessSymbol =
       Symbol::forConcreteType(typeWitnessSchema, result, Context);
 
-  auto concreteConformanceSymbol = *(subjectType.end() - 2);
-  auto assocTypeSymbol = *(subjectType.end() - 1);
-
-  RewriteSystem::ConcreteTypeWitness witness(concreteConformanceSymbol,
-                                             assocTypeSymbol,
-                                             typeWitnessSymbol);
-  unsigned witnessID = System.recordConcreteTypeWitness(witness);
+  RewriteSystem::TypeWitness witness(Term::get(subjectType, Context),
+                                     typeWitnessSymbol);
+  unsigned witnessID = System.recordTypeWitness(witness);
 
   // Simplify the substitution terms in the type witness symbol.
   RewritePath substPath;
@@ -764,7 +765,7 @@ MutableTerm PropertyMap::computeConstraintTermForTypeWitness(
 
     // Add a rule T.[concrete: C : P] => T.[concrete: C : P].[P:X].
     MutableTerm result(key);
-    result.add(concreteConformanceSymbol);
+    result.add(witness.getConcreteConformance());
 
     // ([concrete: C : P] => [concrete: C : P].[P:X].[concrete: C])
     path.add(RewriteStep::forSameTypeWitness(
diff --git a/lib/AST/RequirementMachine/RewriteLoop.cpp b/lib/AST/RequirementMachine/RewriteLoop.cpp
@@ -436,13 +436,11 @@ void RewritePathEvaluator::applyConcreteTypeWitness(const RewriteStep &step,
   checkA();
   auto &term = A.back();
 
-  const auto &witness = system.getConcreteTypeWitness(step.RuleID);
+  const auto &witness = system.getTypeWitness(step.RuleID);
   auto fail = [&]() {
     llvm::errs() << "Bad concrete type witness term:\n";
     llvm::errs() << term << "\n\n";
-    llvm::errs() << "Conformance: " << witness.ConcreteConformance << "\n";
-    llvm::errs() << "Assoc type: " << witness.AssocType << "\n";
-    llvm::errs() << "Concrete type: " << witness.ConcreteType << "\n";
+    witness.dump(llvm::errs());
     llvm::errs() << "End offset: " << step.EndOffset << "\n";
     llvm::errs() << "Inverse: " << step.Inverse << "\n";
     abort();
@@ -453,9 +451,9 @@ void RewritePathEvaluator::applyConcreteTypeWitness(const RewriteStep &step,
     //
     //    U.[concrete: C : P].[P:X].[concrete: C.X].V
     if (term.size() <= step.EndOffset + 3 ||
-        *(term.end() - step.EndOffset - 3) != witness.ConcreteConformance ||
-        *(term.end() - step.EndOffset - 2) != witness.AssocType ||
-        *(term.end() - step.EndOffset - 1) != witness.ConcreteType) {
+        *(term.end() - step.EndOffset - 3) != witness.getConcreteConformance() ||
+        *(term.end() - step.EndOffset - 2) != witness.getAssocType() ||
+        *(term.end() - step.EndOffset - 1) != witness.getConcreteType()) {
       fail();
     }
 
@@ -471,16 +469,16 @@ void RewritePathEvaluator::applyConcreteTypeWitness(const RewriteStep &step,
     //
     //    U.[concrete: C : P].[P:X].V
     if (term.size() <= step.EndOffset + 2 ||
-        *(term.end() - step.EndOffset - 2) != witness.ConcreteConformance ||
-        *(term.end() - step.EndOffset - 1) != witness.AssocType) {
+        *(term.end() - step.EndOffset - 2) != witness.getConcreteConformance() ||
+        *(term.end() - step.EndOffset - 1) != witness.getAssocType()) {
       fail();
     }
 
     // Get the subterm U.[concrete: C : P].[P:X].
     MutableTerm newTerm(term.begin(), term.end() - step.EndOffset);
 
     // Add the symbol [concrete: C.X].
-    newTerm.add(witness.ConcreteType);
+    newTerm.add(witness.getConcreteType());
 
     // Add the subterm V, to get
     //
@@ -496,35 +494,35 @@ void RewritePathEvaluator::applySameTypeWitness(const RewriteStep &step,
   checkA();
   auto &term = A.back();
 
-  const auto &witness = system.getConcreteTypeWitness(step.RuleID);
+  const auto &witness = system.getTypeWitness(step.RuleID);
   auto fail = [&]() {
     llvm::errs() << "Bad same-type witness term:\n";
     llvm::errs() << term << "\n\n";
-    llvm::errs() << "Conformance: " << witness.ConcreteConformance << "\n";
-    llvm::errs() << "Assoc type: " << witness.AssocType << "\n";
-    llvm::errs() << "Concrete type: " << witness.ConcreteType << "\n";
+    witness.dump(llvm::errs());
     abort();
   };
 
+  auto concreteConformanceSymbol = witness.getConcreteConformance();
+
 #ifndef NDEBUG
-  if (witness.ConcreteType.getConcreteType() !=
-      witness.ConcreteConformance.getConcreteType()) {
+  if (witness.getConcreteType().getConcreteType() !=
+      concreteConformanceSymbol.getConcreteType()) {
     fail();
   }
 #endif
 
   auto witnessConcreteType = Symbol::forConcreteType(
-      witness.ConcreteConformance.getConcreteType(),
-      witness.ConcreteConformance.getSubstitutions(),
+      concreteConformanceSymbol.getConcreteType(),
+      concreteConformanceSymbol.getSubstitutions(),
       system.getRewriteContext());
 
   if (!step.Inverse) {
     // Make sure the term takes the following form, where |V| == EndOffset:
     //
     //    U.[concrete: C : P].[P:X].[concrete: C].V
     if (term.size() <= step.EndOffset + 3 ||
-        *(term.end() - step.EndOffset - 3) != witness.ConcreteConformance ||
-        *(term.end() - step.EndOffset - 2) != witness.AssocType ||
+        *(term.end() - step.EndOffset - 3) != concreteConformanceSymbol ||
+        *(term.end() - step.EndOffset - 2) != witness.getAssocType() ||
         *(term.end() - step.EndOffset - 1) != witnessConcreteType) {
       fail();
     }
@@ -541,15 +539,15 @@ void RewritePathEvaluator::applySameTypeWitness(const RewriteStep &step,
     //
     //    U.[concrete: C : P].V
     if (term.size() <= step.EndOffset + 1 ||
-        *(term.end() - step.EndOffset - 1) != witness.ConcreteConformance) {
+        *(term.end() - step.EndOffset - 1) != concreteConformanceSymbol) {
       fail();
     }
 
     // Get the subterm U.[concrete: C : P].
     MutableTerm newTerm(term.begin(), term.end() - step.EndOffset);
 
     // Add the symbol [P:X].
-    newTerm.add(witness.AssocType);
+    newTerm.add(witness.getAssocType());
 
     // Add the symbol [concrete: C].
     newTerm.add(witnessConcreteType);
diff --git a/lib/AST/RequirementMachine/RewriteLoop.h b/lib/AST/RequirementMachine/RewriteLoop.h
@@ -118,8 +118,8 @@ struct RewriteStep {
     /// If inverted: the concrete type symbol [concrete: C.X] is introduced.
     ///
     /// The RuleID field is repurposed to store the result of calling
-    /// RewriteSystem::recordConcreteTypeWitness(). This index is then
-    /// passed in to RewriteSystem::getConcreteTypeWitness() when applying
+    /// RewriteSystem::recordTypeWitness(). This index is then passed in
+    /// to RewriteSystem::getTypeWitness() when applying
     /// the step.
     ConcreteTypeWitness,
 
@@ -129,9 +129,7 @@ struct RewriteStep {
     ///
     /// If inverted: the associated type symbol [P:X] is introduced.
     ///
-    /// The RuleID field is repurposed to store the result of calling
-    /// RewriteSystem::recordConcreteTypeWitness(). This index is then
-    /// passed in to RewriteSystem::getConcreteTypeWitness() when applying
+    /// The RuleID field is a TypeWitness ID as above.
     /// the step.
     SameTypeWitness,
   };
diff --git a/lib/AST/RequirementMachine/RewriteSystem.h b/lib/AST/RequirementMachine/RewriteSystem.h
@@ -14,6 +14,7 @@
 #define SWIFT_REWRITESYSTEM_H
 
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/PointerUnion.h"
 
 #include "Debug.h"
 #include "RewriteLoop.h"
@@ -324,28 +325,68 @@ class RewriteSystem final {
   //////////////////////////////////////////////////////////////////////////////
 
 public:
-  struct ConcreteTypeWitness {
-    Symbol ConcreteConformance;
-    Symbol AssocType;
-    Symbol ConcreteType;
+  /// A type witness has a subject type, stored in LHS, which takes the form:
+  ///
+  /// T.[concrete: C : P].[P:X]
+  ///
+  /// For some concrete type C, protocol P and associated type X.
+  ///
+  /// The type witness of X in the conformance C : P is either a concrete type,
+  /// or an abstract type parameter.
+  ///
+  /// If it is a concrete type, then RHS stores the concrete type symbol.
+  ///
+  /// If it is an abstract type parameter, then RHS stores the type term.
+  ///
+  /// Think of these as rewrite rules which are lazily created, but always
+  /// "there" -- they encode information about concrete conformances, which
+  /// are solved outside of the requirement machine itself.
+  ///
+  /// We don't want to eagerly pull in all concrete conformances and walk
+  /// them recursively introducing rewrite rules.
+  ///
+  /// The RewriteStep::{Concrete,Same,Abstract}TypeWitness rewrite step kinds
+  /// reference TypeWitnesses via their RuleID field.
+  ///
+  /// Type witnesses are recorded lazily in property map construction, in
+  /// PropertyMap::computeConstraintTermForTypeWitness().
+  struct TypeWitness {
+    Term LHS;
+    llvm::PointerUnion<Symbol, Term> RHS;
+
+    TypeWitness(Term lhs, llvm::PointerUnion<Symbol, Term> rhs);
+
+    friend bool operator==(const TypeWitness &lhs,
+                           const TypeWitness &rhs);
+
+    Symbol getConcreteConformance() const {
+      return *(LHS.end() - 2);
+    }
+
+    Symbol getAssocType() const {
+      return *(LHS.end() - 1);
+    }
+
+    Symbol getConcreteType() const {
+      return RHS.get<Symbol>();
+    }
 
-    ConcreteTypeWitness(Symbol concreteConformance,
-                        Symbol assocType,
-                        Symbol concreteType);
+    Term getAbstractType() const {
+      return RHS.get<Term>();
+    }
 
-    friend bool operator==(const ConcreteTypeWitness &lhs,
-                           const ConcreteTypeWitness &rhs);
+    void dump(llvm::raw_ostream &out) const;
   };
 
 private:
   /// Cache for concrete type witnesses. The value in the map is an index
   /// into the vector.
-  llvm::DenseMap<std::pair<Symbol, Symbol>, unsigned> ConcreteTypeWitnessMap;
-  std::vector<ConcreteTypeWitness> ConcreteTypeWitnesses;
+  llvm::DenseMap<Term, unsigned> TypeWitnessMap;
+  std::vector<TypeWitness> TypeWitnesses;
 
 public:
-  unsigned recordConcreteTypeWitness(ConcreteTypeWitness witness);
-  const ConcreteTypeWitness &getConcreteTypeWitness(unsigned index) const;
+  unsigned recordTypeWitness(TypeWitness witness);
+  const TypeWitness &getTypeWitness(unsigned index) const;
 
   //////////////////////////////////////////////////////////////////////////////
   ///
