RequirementMachine: PropertyMap can use Terms instead of MutableTerms as keys

slavapestov · slavapestov · commit 9738d4eef6ed · 2021-08-06T14:17:20.000-04:00
Also, we don't have to sort rules in term order before adding them to
the map; a bucket sort by term length is sufficient.
diff --git a/lib/AST/RequirementMachine/PropertyMap.cpp b/lib/AST/RequirementMachine/PropertyMap.cpp
@@ -672,14 +672,10 @@ PropertyMap::lookUpProperties(const MutableTerm &key) const {
 ///
 /// This must be called in monotonically non-decreasing key order.
 PropertyBag *
-PropertyMap::getOrCreateProperties(const MutableTerm &key) {
-  assert(!key.empty());
-
-  if (!Entries.empty()) {
-    const auto &lastEntry = Entries.back();
-    if (lastEntry->getKey() == key)
-      return lastEntry;
-  }
+PropertyMap::getOrCreateProperties(Term key) {
+  auto next = Trie.find(key.rbegin(), key.rend());
+  if (next && (*next)->getKey() == key)
+    return *next;
 
   auto *props = new PropertyBag(key);
 
@@ -704,12 +700,9 @@ PropertyMap::getOrCreateProperties(const MutableTerm &key) {
   //
   // Since 'A' has no proper suffix with additional properties, the next
   // property bag of 'A' is nullptr.
-  if (auto *next = lookUpProperties(key))
-    props->copyPropertiesFrom(next, Context);
+  if (next)
+    props->copyPropertiesFrom(*next, Context);
 
-  // Here is where we make the assumption that if the new key is not equal to
-  // the key of the last entry, then it is larger than any key already in the
-  // map.
   Entries.push_back(props);
   auto oldProps = Trie.insert(key.rbegin(), key.rend(), props);
   if (oldProps) {
@@ -738,7 +731,7 @@ void PropertyMap::clear() {
 /// Record a protocol conformance, layout or superclass constraint on the given
 /// key. Must be called in monotonically non-decreasing key order.
 void PropertyMap::addProperty(
-    const MutableTerm &key, Symbol property,
+    Term key, Symbol property,
     SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules) {
   assert(property.isProperty());
   auto *props = getOrCreateProperties(key);
@@ -763,12 +756,8 @@ void PropertyMap::computeConcreteTypeInDomainMap() {
     auto concreteTypeKey = std::make_pair(concreteType, domain);
 
     auto found = ConcreteTypeInDomainMap.find(concreteTypeKey);
-    if (found != ConcreteTypeInDomainMap.end()) {
-      const auto &otherTerm = found->second;
-      assert(props->Key.compare(otherTerm, Protos) > 0 &&
-             "Out-of-order keys?");
+    if (found != ConcreteTypeInDomainMap.end())
       continue;
-    }
 
     auto inserted = ConcreteTypeInDomainMap.insert(
         std::make_pair(concreteTypeKey, props->Key));
@@ -857,7 +846,7 @@ void PropertyMap::concretizeNestedTypesFromConcreteParents(
 ///    T.[P:B] => U.V
 ///
 void PropertyMap::concretizeNestedTypesFromConcreteParent(
-    const MutableTerm &key, RequirementKind requirementKind,
+    Term key, RequirementKind requirementKind,
     CanType concreteType, ArrayRef<Term> substitutions,
     ArrayRef<const ProtocolDecl *> conformsTo,
     llvm::TinyPtrVector<ProtocolConformance *> &conformances,
@@ -920,7 +909,7 @@ void PropertyMap::concretizeNestedTypesFromConcreteParent(
                      << " of " << concreteType << " is " << typeWitness << "\n";
       }
 
-      MutableTerm subjectType = key;
+      MutableTerm subjectType(key);
       subjectType.add(Symbol::forAssociatedType(proto, assocType->getName(),
                                                 Context));
 
@@ -936,7 +925,7 @@ void PropertyMap::concretizeNestedTypesFromConcreteParent(
         }
 
         // Add a rule T.[P:A] => T.
-        constraintType = key;
+        constraintType = MutableTerm(key);
       } else {
         constraintType = computeConstraintTermForTypeWitness(
             key, concreteType, typeWitness, subjectType,
@@ -977,7 +966,7 @@ void PropertyMap::concretizeNestedTypesFromConcreteParent(
 ///
 ///        T.[P:A] => V
 MutableTerm PropertyMap::computeConstraintTermForTypeWitness(
-    const MutableTerm &key, CanType concreteType, CanType typeWitness,
+    Term key, CanType concreteType, CanType typeWitness,
     const MutableTerm &subjectType, ArrayRef<Term> substitutions) const {
   if (!typeWitness->hasTypeParameter()) {
     // Check if we have a shorter representative we can use.
@@ -986,12 +975,13 @@ MutableTerm PropertyMap::computeConstraintTermForTypeWitness(
 
     auto found = ConcreteTypeInDomainMap.find(concreteTypeKey);
     if (found != ConcreteTypeInDomainMap.end()) {
-      if (found->second != subjectType) {
+      MutableTerm result(found->second);
+      if (result != subjectType) {
         if (DebugConcretizeNestedTypes) {
           llvm::dbgs() << "^^ Type witness can re-use property bag of "
                        << found->second << "\n";
         }
-        return found->second;
+        return result;
       }
     }
   }
@@ -1048,7 +1038,11 @@ RewriteSystem::buildPropertyMap(PropertyMap &map,
                                 unsigned maxDepth) {
   map.clear();
 
-  std::vector<std::pair<MutableTerm, Symbol>> properties;
+  // PropertyMap::addRule() requires that shorter rules are added
+  // before longer rules, so that it can perform lookups on suffixes and call
+  // PropertyBag::copyPropertiesFrom(). However, we don't have to perform a
+  // full sort by term order here; a bucket sort by term length suffices.
+  SmallVector<std::vector<std::pair<Term, Symbol>>, 4> properties;
 
   for (const auto &rule : Rules) {
     if (rule.isDeleted())
@@ -1068,31 +1062,19 @@ RewriteSystem::buildPropertyMap(PropertyMap &map,
     if (!std::equal(rhs.begin(), rhs.end(), lhs.begin()))
       continue;
 
-    MutableTerm key(rhs);
-
-#ifndef NDEBUG
-    assert(!simplify(key) &&
-           "Right hand side of a property rule should already be reduced");
-#endif
-
-    properties.emplace_back(key, property);
+    if (rhs.size() >= properties.size())
+      properties.resize(rhs.size() + 1);
+    properties[rhs.size()].emplace_back(rhs, property);
   }
 
-  // PropertyMap::addRule() requires that shorter rules are added
-  // before longer rules, so that it can perform lookups on suffixes and call
-  // PropertyBag::copyPropertiesFrom().
-  std::sort(properties.begin(), properties.end(),
-            [&](const std::pair<MutableTerm, Symbol> &lhs,
-                const std::pair<MutableTerm, Symbol> &rhs) -> bool {
-              return lhs.first.compare(rhs.first, Protos) < 0;
-            });
-
   // Merging multiple superclass or concrete type rules can induce new rules
   // to unify concrete type constructor arguments.
   SmallVector<std::pair<MutableTerm, MutableTerm>, 3> inducedRules;
 
-  for (auto pair : properties) {
-    map.addProperty(pair.first, pair.second, inducedRules);
+  for (const auto &bucket : properties) {
+    for (auto pair : bucket) {
+      map.addProperty(pair.first, pair.second, inducedRules);
+    }
   }
 
   // We collect terms with fully concrete types so that we can re-use them
diff --git a/lib/AST/RequirementMachine/PropertyMap.h b/lib/AST/RequirementMachine/PropertyMap.h
@@ -50,7 +50,7 @@ class PropertyBag {
   friend class PropertyMap;
 
   /// The fully reduced term whose properties are recorded in this property bag.
-  MutableTerm Key;
+  Term Key;
 
   /// All protocols this type conforms to.
   llvm::TinyPtrVector<const ProtocolDecl *> ConformsTo;
@@ -72,7 +72,7 @@ class PropertyBag {
   /// ConformsTo list.
   llvm::TinyPtrVector<ProtocolConformance *> ConcreteConformances;
 
-  explicit PropertyBag(const MutableTerm &key) : Key(key) {}
+  explicit PropertyBag(Term key) : Key(key) {}
 
   void addProperty(Symbol property,
                    RewriteContext &ctx,
@@ -87,7 +87,7 @@ class PropertyBag {
   PropertyBag &operator=(PropertyBag &&) = delete;
 
 public:
-  const MutableTerm &getKey() const { return Key; }
+  Term getKey() const { return Key; }
   void dump(llvm::raw_ostream &out) const;
 
   bool hasSuperclassBound() const {
@@ -142,13 +142,13 @@ class PropertyMap {
   Trie<PropertyBag *, MatchKind::Longest> Trie;
 
   using ConcreteTypeInDomain = std::pair<CanType, ArrayRef<const ProtocolDecl *>>;
-  llvm::DenseMap<ConcreteTypeInDomain, MutableTerm> ConcreteTypeInDomainMap;
+  llvm::DenseMap<ConcreteTypeInDomain, Term> ConcreteTypeInDomainMap;
 
   const ProtocolGraph &Protos;
   unsigned DebugConcreteUnification : 1;
   unsigned DebugConcretizeNestedTypes : 1;
 
-  PropertyBag *getOrCreateProperties(const MutableTerm &key);
+  PropertyBag *getOrCreateProperties(Term key);
 
   PropertyMap(const PropertyMap &) = delete;
   PropertyMap(PropertyMap &&) = delete;
@@ -170,7 +170,7 @@ class PropertyMap {
   void dump(llvm::raw_ostream &out) const;
 
   void clear();
-  void addProperty(const MutableTerm &key, Symbol property,
+  void addProperty(Term key, Symbol property,
                    SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules);
 
   void computeConcreteTypeInDomainMap();
@@ -179,14 +179,14 @@ class PropertyMap {
 
 private:
   void concretizeNestedTypesFromConcreteParent(
-                   const MutableTerm &key, RequirementKind requirementKind,
+                   Term key, RequirementKind requirementKind,
                    CanType concreteType, ArrayRef<Term> substitutions,
                    ArrayRef<const ProtocolDecl *> conformsTo,
                    llvm::TinyPtrVector<ProtocolConformance *> &conformances,
                    SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules) const;
 
   MutableTerm computeConstraintTermForTypeWitness(
-      const MutableTerm &key, CanType concreteType, CanType typeWitness,
+      Term key, CanType concreteType, CanType typeWitness,
       const MutableTerm &subjectType, ArrayRef<Term> substitutions) const;
 };
 
diff --git a/test/Generics/rdar79564324.swift b/test/Generics/rdar79564324.swift
@@ -35,6 +35,6 @@ public func test<T : P>(_ t: T) where T == T.A {
 // CHECK-NEXT: }
 // CHECK-NEXT: Property map: {
 // CHECK-NEXT:   [P:A] => { conforms_to: [P] }
-// CHECK-NEXT:   τ_0_0 => { conforms_to: [P] }
 // CHECK-NEXT:   τ_0_1 => { conforms_to: [P] }
+// CHECK-NEXT:   τ_0_0 => { conforms_to: [P] }
 // CHECK-NEXT: }
diff --git a/test/Generics/unify_associated_types.swift b/test/Generics/unify_associated_types.swift
@@ -29,10 +29,10 @@ struct MergeTest<G : P1a & P2a> {}
 // CHECK: - [P1a&P2a:T].[P1:X] => [P1a&P2a:T].[P1&P2:X]
 // CHECK: }
 // CHECK: Property map: {
-// CHECK:   [P1a&P2a:T] => { conforms_to: [P1 P2] }
 // CHECK:   [P1a:T] => { conforms_to: [P1] }
 // CHECK:   [P2a:T] => { conforms_to: [P2] }
 // CHECK:   τ_0_0 => { conforms_to: [P1a P2a] }
+// CHECK:   [P1a&P2a:T] => { conforms_to: [P1 P2] }
 // CHECK: }
 // CHECK: }
 
diff --git a/test/Generics/unify_superclass_types_2.swift b/test/Generics/unify_superclass_types_2.swift
@@ -42,6 +42,6 @@ func unifySuperclassTest<T : P1 & P2>(_: T) {
 // CHECK-NEXT:   [P2:X] => { layout: _NativeClass superclass: [superclass: Generic<τ_0_0, String, τ_0_1> with <[P2:A2], [P2:B2]>] }
 // CHECK-NEXT:   τ_0_0 => { conforms_to: [P1 P2] }
 // CHECK-NEXT:   τ_0_0.[P1&P2:X] => { layout: _NativeClass superclass: [superclass: Generic<Int, τ_0_0, τ_0_1> with <τ_0_0.[P1:A1], τ_0_0.[P1:B1]>] }
-// CHECK-NEXT:   τ_0_0.[P1:A1] => { concrete_type: [concrete: String] }
 // CHECK-NEXT:   τ_0_0.[P2:A2] => { concrete_type: [concrete: Int] }
+// CHECK-NEXT:   τ_0_0.[P1:A1] => { concrete_type: [concrete: String] }
 // CHECK-NEXT: }
diff --git a/test/Generics/unify_superclass_types_3.swift b/test/Generics/unify_superclass_types_3.swift
@@ -44,6 +44,6 @@ func unifySuperclassTest<T : P1 & P2>(_: T) {
 // CHECK-NEXT:   [P2:X] => { layout: _NativeClass superclass: [superclass: Generic<τ_0_0, String, τ_0_1> with <[P2:A2], [P2:B2]>] }
 // CHECK-NEXT:   τ_0_0 => { conforms_to: [P1 P2] }
 // CHECK-NEXT:   τ_0_0.[P1&P2:X] => { layout: _NativeClass superclass: [superclass: Derived<τ_0_0, τ_0_1> with <τ_0_0.[P1:A1], τ_0_0.[P1:B1]>] }
-// CHECK-NEXT:   τ_0_0.[P1:A1] => { concrete_type: [concrete: String] }
 // CHECK-NEXT:   τ_0_0.[P2:A2] => { concrete_type: [concrete: Int] }
+// CHECK-NEXT:   τ_0_0.[P1:A1] => { concrete_type: [concrete: String] }
 // CHECK-NEXT: }
