RequirementMachine: Adding shape abstractions

Adding abstractions to check terms for shape symbol and remove the shape
symbol from the end of the sequence of symbols, rather than manually
manipulating the end() sequence externally.

Author: kathygray-pl

lib/AST/RequirementMachine/GenericSignatureQueries.cpp

@@ -755,16 +755,16 @@ RequirementMachine::getReducedShapeTerm(Type type) const {
   // Get the term T', which is the reduced shape of T.
   if (term.size() != 2 ||
       term[0].getKind() != Symbol::Kind::GenericParam ||
-      term[1].getKind() != Symbol::Kind::Shape) {
+      !term.hasShape()) {
     ABORT([&](auto &out) {
       out << "Invalid reduced shape\n";
       out << "Type: " << type << "\n";
       out << "Term: " << term;
     });
   }
 
-  MutableTerm reducedTerm(term.begin(), term.end() - 1);
-  return reducedTerm;
+  term.removeShape();
+  return term;
 }
 
 Type RequirementMachine::getReducedShape(Type type,

lib/AST/RequirementMachine/InterfaceType.cpp

@@ -424,7 +424,7 @@ RewriteContext::getRelativeTermForType(CanType typeWitness,
   // Get the substitution S corresponding to τ_0_n.
   unsigned index = getGenericParamIndex(typeWitness->getRootGenericParam());
   result = MutableTerm(substitutions[index]);
-  ASSERT(result.back().getKind() != Symbol::Kind::Shape);
+  ASSERT(!result.hasShape());
 
   // If the substitution is a term consisting of a single protocol symbol
   // [P], save P for later.
@@ -485,7 +485,7 @@ Type PropertyMap::getTypeFromSubstitutionSchema(
           auto substitution = substitutions[index];
 
           bool isShapePosition = (pos == TypePosition::Shape);
-          bool isShapeTerm = (substitution.back() == Symbol::forShape(Context));
+          bool isShapeTerm = substitution.hasShape();
           if (isShapePosition != isShapeTerm) {
             ABORT([&](auto &out) {
               out << "Shape vs. type mixup\n\n";
@@ -504,8 +504,8 @@ Type PropertyMap::getTypeFromSubstitutionSchema(
           // Undo the thing where the count type of a PackExpansionType
           // becomes a shape term.
           if (isShapeTerm) {
-            MutableTerm mutTerm(substitution.begin(), substitution.end() - 1);
-            substitution = Term::get(mutTerm, Context);
+            MutableTerm noShape = substitution.termWithoutShape();
+            substitution = Term::get(noShape, Context);
           }
 
           // Prepend the prefix of the lookup key to the substitution.

lib/AST/RequirementMachine/RequirementBuilder.cpp

@@ -257,11 +257,10 @@ void RequirementBuilder::addRequirementRules(ArrayRef<unsigned> rules) {
 
     ASSERT(rule.getLHS().back().getKind() != Symbol::Kind::Protocol);
 
-    if (constraintTerm.back().getKind() == Symbol::Kind::Shape) {
-      ASSERT(rule.getRHS().back().getKind() == Symbol::Kind::Shape);
+    if (constraintTerm.hasShape()) {
+      ASSERT(rule.getRHS().hasShape());
       // Strip off the shape symbol from the constraint term.
-      constraintTerm = MutableTerm(constraintTerm.begin(),
-                                   constraintTerm.end() - 1);
+      constraintTerm.removeShape();
     }
 
     if (constraintTerm.front().getKind() == Symbol::Kind::PackElement) {
@@ -332,10 +331,10 @@ void RequirementBuilder::processConnectedComponents() {
   for (auto &pair : Components) {
     MutableTerm subjectTerm(pair.first);
     RequirementKind kind;
-    if (subjectTerm.back().getKind() == Symbol::Kind::Shape) {
+    if (subjectTerm.hasShape()) {
       kind = RequirementKind::SameShape;
       // Strip off the shape symbol from the subject term.
-      subjectTerm = MutableTerm(subjectTerm.begin(), subjectTerm.end() - 1);
+      subjectTerm.removeShape();
     } else {
       kind = RequirementKind::SameType;
       if (subjectTerm.front().getKind() == Symbol::Kind::PackElement) {

lib/AST/RequirementMachine/RewriteSystem.cpp

@@ -634,7 +634,7 @@ void RewriteSystem::verifyRewriteRules(ValidityPolicy policy) const {
 
     if (rhs.size() == 1 && rhs[0].getKind() == Symbol::Kind::Shape) {
       // We can have a rule like T.[shape] => [shape].
-      ASSERT_RULE(lhs.back().getKind() == Symbol::Kind::Shape);
+      ASSERT_RULE(lhs.hasShape());
     } else {
       // Otherwise, LHS and RHS must have the same domain.
       auto lhsDomain = lhs.getRootProtocol();

lib/AST/RequirementMachine/Term.cpp

@@ -68,6 +68,17 @@ Symbol Term::back() const {
   return Ptr->getElements().back();
 }
 
+bool Term::hasShape() const {
+  return back().getKind() == Symbol::Kind::Shape;
+}
+
+MutableTerm Term::termWithoutShape() const {
+  if (hasShape())
+    return MutableTerm(begin(), end() - 1);
+  else
+    return MutableTerm(begin(), end());
+}
+
 Symbol Term::operator[](size_t index) const {
   return Ptr->getElements()[index];
 }
@@ -224,6 +235,15 @@ std::optional<int> MutableTerm::compare(const MutableTerm &other,
   return compareImpl(begin(), end(), other.begin(), other.end(), ctx);
 }
 
+bool MutableTerm::hasShape() const {
+  return back().getKind() == Symbol::Kind::Shape;
+}
+
+void MutableTerm::removeShape() {
+  if (hasShape())
+    Symbols.pop_back();
+}
+
 /// Replace the subterm in the range [from,to) of this term with \p rhs.
 void MutableTerm::rewriteSubTerm(Symbol *from, Symbol *to, Term rhs) {
   auto oldSize = size();

lib/AST/RequirementMachine/Term.h

@@ -59,6 +59,9 @@ class Term final {
 
   Symbol back() const;
 
+  bool hasShape() const;
+  MutableTerm termWithoutShape() const;
+
   Symbol operator[](size_t index) const;
 
   /// Returns an opaque pointer that uniquely identifies this term.
@@ -184,6 +187,9 @@ class MutableTerm final {
     return Symbols.back();
   }
 
+  bool hasShape() const;
+  void removeShape();
+
   Symbol operator[](size_t index) const {
     return Symbols[index];
   }