Add methods to determine the ownership behavior of patterns.

When matching against a noncopyable value, whether the match operation can
borrow the value in-place or needs to take ownership of it is significant.
This can generally be determined from the kind of pattern being used, except
in the case of expr patterns, where it depends on type-checking the `~=`
operator that was used.

Author: jckarter

include/swift/AST/Ownership.h

@@ -118,14 +118,22 @@ optionalityOf(ReferenceOwnership ownership) {
 }
 
 /// Different kinds of value ownership supported by Swift.
+///
+/// The order of these constants is significant. Ascending order indicates
+/// stricter requirements on the value to be used with the given ownership
+/// kind: any value can be accessed with a shared borrow (so long as there
+/// are no exclusive accesses overlapping). An exclusive `inout` borrow is
+/// only possible for mutable values which the caller already has exclusive
+/// access to or ownership of. Consumption of an owned value requires sole
+/// ownership of that value.
 enum class ValueOwnership : uint8_t {
   /// the context-dependent default ownership (sometimes shared,
   /// sometimes owned)
   Default,
-  /// an 'inout' exclusive, mutating borrow
-  InOut,
   /// a 'borrowing' nonexclusive, usually nonmutating borrow
   Shared,
+  /// an 'inout' exclusive, mutating borrow
+  InOut,
   /// a 'consuming' ownership transfer
   Owned,
 

include/swift/AST/Pattern.h

@@ -230,6 +230,16 @@ class alignas(8) Pattern : public ASTAllocated<Pattern> {
 
   /// Does this pattern have any mutable 'var' bindings?
   bool hasAnyMutableBindings() const;
+  
+  /// Get the ownership behavior of this pattern on the value being matched
+  /// against it.
+  ///
+  /// The pattern must be type-checked for this operation to be valid. If
+  /// \c mostRestrictiveSubpatterns is non-null, the pointed-to vector will be
+  /// populated with references to the subpatterns that cause the pattern to
+  /// have stricter than "shared" ownership behavior for diagnostic purposes.
+  ValueOwnership getOwnership(
+    SmallVectorImpl<Pattern*> *mostRestrictiveSubpatterns = nullptr) const;
 
   static bool classof(const Pattern *P) { return true; }
 
@@ -685,8 +695,10 @@ class ExprPattern : public Pattern {
   DeclContext *DC;
 
   /// A synthesized call to the '~=' operator comparing the match expression
-  /// on the left to the matched value on the right.
-  mutable Expr *MatchExpr = nullptr;
+  /// on the left to the matched value on the right, pairend with a record of the
+  /// ownership of the subject operand.
+  mutable llvm::PointerIntPair<Expr *, 2, ValueOwnership>
+    MatchExprAndOperandOwnership{nullptr, ValueOwnership::Default};
 
   /// An implicit variable used to represent the RHS value of the synthesized
   /// match expression.
@@ -698,6 +710,8 @@ class ExprPattern : public Pattern {
 
   friend class ExprPatternMatchRequest;
 
+  void updateMatchExpr(Expr *matchExpr) const;
+
 public:
   /// Create a new parsed unresolved ExprPattern.
   static ExprPattern *createParsed(ASTContext &ctx, Expr *E, DeclContext *DC);
@@ -711,7 +725,6 @@ class ExprPattern : public Pattern {
 
   Expr *getSubExpr() const { return SubExprAndIsResolved.getPointer(); }
   void setSubExpr(Expr *e) { SubExprAndIsResolved.setPointer(e); }
-
   DeclContext *getDeclContext() const { return DC; }
 
   void setDeclContext(DeclContext *newDC) {
@@ -722,7 +735,17 @@ class ExprPattern : public Pattern {
 
   /// The match expression if it has been computed, \c nullptr otherwise.
   /// Should only be used by the ASTDumper and ASTWalker.
-  Expr *getCachedMatchExpr() const { return MatchExpr; }
+  Expr *getCachedMatchExpr() const {
+    return MatchExprAndOperandOwnership.getPointer();
+  }
+
+  /// Return the ownership of the subject parameter for the `~=` operator being
+  /// used (and thereby, the ownership of the pattern match itself), or
+  /// \c Default if the ownership of the parameter is unresolved.
+  ValueOwnership getCachedMatchOperandOwnership() const {
+    auto ownership = MatchExprAndOperandOwnership.getInt();
+    return ownership;
+  }
 
   /// The match variable if it has been computed, \c nullptr otherwise.
   /// Should only be used by the ASTDumper and ASTWalker.
@@ -731,14 +754,18 @@ class ExprPattern : public Pattern {
   /// A synthesized call to the '~=' operator comparing the match expression
   /// on the left to the matched value on the right.
   Expr *getMatchExpr() const;
+  ValueOwnership getMatchOperandOwnership() const {
+    (void)getMatchExpr();
+    return getCachedMatchOperandOwnership();
+  }
 
   /// An implicit variable used to represent the RHS value of the synthesized
   /// match expression.
   VarDecl *getMatchVar() const;
 
   void setMatchExpr(Expr *e) {
-    assert(MatchExpr && "Should only update an existing MatchExpr");
-    MatchExpr = e;
+    assert(getCachedMatchExpr() && "Should only update an existing MatchExpr");
+    updateMatchExpr(e);
   }
 
   SourceLoc getLoc() const;

lib/AST/ASTDumper.cpp

@@ -928,7 +928,7 @@ namespace {
     using PrintBase::PrintBase;
 
     void printCommon(Pattern *P, const char *Name, StringRef Label) {
-     printHead(Name, PatternColor, Label);
+      printHead(Name, PatternColor, Label);
 
       printFlag(P->isImplicit(), "implicit", ExprModifierColor);
 
@@ -990,6 +990,22 @@ namespace {
     }
     void visitExprPattern(ExprPattern *P, StringRef label) {
       printCommon(P, "pattern_expr", label);
+      switch (P->getCachedMatchOperandOwnership()) {
+      case ValueOwnership::Default:
+        break;
+      case ValueOwnership::Shared:
+        printFieldRaw([](llvm::raw_ostream &os) { os << "borrowing"; },
+                      "ownership");
+        break;
+      case ValueOwnership::InOut:
+        printFieldRaw([](llvm::raw_ostream &os) { os << "mutating"; },
+                      "ownership");
+        break;
+      case ValueOwnership::Owned:
+        printFieldRaw([](llvm::raw_ostream &os) { os << "consuming"; },
+                      "ownership");
+        break;
+      }
       if (auto m = P->getCachedMatchExpr())
         printRec(m);
       else
@@ -2011,6 +2027,22 @@ class PrintStmt : public StmtVisitor<PrintStmt, void, StringRef>,
         printFlag(LabelItem.isDefault(), "default");
 
         if (auto *CasePattern = LabelItem.getPattern()) {
+          switch (CasePattern->getOwnership()) {
+          case ValueOwnership::Default:
+            break;
+          case ValueOwnership::Shared:
+            printFieldRaw([](llvm::raw_ostream &os) { os << "borrowing"; },
+                          "ownership");
+            break;
+          case ValueOwnership::InOut:
+            printFieldRaw([](llvm::raw_ostream &os) { os << "mutating"; },
+                          "ownership");
+            break;
+          case ValueOwnership::Owned:
+            printFieldRaw([](llvm::raw_ostream &os) { os << "consuming"; },
+                          "ownership");
+            break;
+          }
           printRec(CasePattern);
         }
         if (auto *Guard = LabelItem.getGuardExpr()) {

lib/AST/Pattern.cpp

@@ -16,6 +16,7 @@
 
 #include "swift/AST/Pattern.h"
 #include "swift/AST/ASTContext.h"
+#include "swift/AST/ASTVisitor.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/GenericEnvironment.h"
@@ -569,6 +570,49 @@ VarDecl *ExprPattern::getMatchVar() const {
       .getMatchVar();
 }
 
+void ExprPattern::updateMatchExpr(Expr *e) const {
+  class FindMatchOperatorDeclRef: public ASTWalker {
+  public:
+    ValueOwnership Ownership = ValueOwnership::Default;
+  
+    PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
+      // See if this is the reference to the ~= operator used.
+      auto declRef = dyn_cast<DeclRefExpr>(E);
+      if (!declRef) {
+        return Action::Continue(E);
+      }
+      auto decl = declRef->getDecl();
+      auto declName = decl->getName();
+      if (!declName.isOperator()) {
+        return Action::Continue(E);
+      }
+      
+      if (!declName.getBaseIdentifier().is("~=")) {
+        return Action::Continue(E);
+      }
+      
+      // We found a `~=` declref. Get the value ownership from the parameter.
+      auto fnTy = decl->getInterfaceType()->castTo<AnyFunctionType>();
+      if (decl->isStatic()) {
+        fnTy = fnTy->getResult()->castTo<AnyFunctionType>();
+      }
+      // Subject value is the right-hand operand to the operator.
+      assert(fnTy->getParams().size() == 2);
+      Ownership = fnTy->getParams()[1].getValueOwnership();
+      // Operators are always normal functions or methods, so their default
+      // parameter ownership is always borrowing.
+      if (Ownership == ValueOwnership::Default) {
+        Ownership = ValueOwnership::Shared;
+      }
+      return Action::Stop();
+    }
+  };
+  FindMatchOperatorDeclRef walker;
+  e->walk(walker);
+
+  MatchExprAndOperandOwnership = {e, walker.Ownership};
+}
+  
 SourceLoc EnumElementPattern::getStartLoc() const {
   return (ParentType && !ParentType->isImplicit())
              ? ParentType->getSourceRange().Start
@@ -680,3 +724,89 @@ void swift::simple_display(llvm::raw_ostream &out, const Pattern *pattern) {
 SourceLoc swift::extractNearestSourceLoc(const Pattern *pattern) {
   return pattern->getLoc();
 }
+
+ValueOwnership
+Pattern::getOwnership(
+  SmallVectorImpl<Pattern *> *mostRestrictiveSubpatterns) const
+{
+  class GetPatternOwnership: public PatternVisitor<GetPatternOwnership, void> {
+  public:
+    ValueOwnership Ownership = ValueOwnership::Shared;
+    SmallVectorImpl<Pattern *> *RestrictingPatterns = nullptr;
+
+    void increaseOwnership(ValueOwnership newOwnership, Pattern *p) {
+      // If the new ownership is stricter than the current ownership, then
+      // clear the restricting patterns we'd collected and start over with the
+      // new stricter ownership.
+      if (newOwnership > Ownership) {
+        Ownership = newOwnership;
+        if (RestrictingPatterns) {
+          RestrictingPatterns->clear();
+        }
+      }
+      
+      if (RestrictingPatterns
+          && newOwnership == Ownership
+          && Ownership > ValueOwnership::Shared) {
+        RestrictingPatterns->push_back(p);
+      }
+    }
+
+#define USE_SUBPATTERN(Kind) \
+    void visit##Kind##Pattern(Kind##Pattern *pattern) { \
+      return visit(pattern->getSubPattern());            \
+    }
+
+    USE_SUBPATTERN(Paren)
+    USE_SUBPATTERN(Typed)
+    USE_SUBPATTERN(Binding)
+#undef USE_SUBPATTERN
+    void visitTuplePattern(TuplePattern *p) {
+      for (auto &element : p->getElements()) {
+        visit(element.getPattern());
+      }
+    }
+    
+    void visitNamedPattern(NamedPattern *p) {
+      // `var` and `let` bindings consume the matched value.
+      // TODO: borrowing/mutating/consuming parameters
+      increaseOwnership(ValueOwnership::Owned, p);
+    }
+    
+    void visitAnyPattern(AnyPattern *p) {
+      /* no change */
+    }
+    void visitBoolPattern(BoolPattern *p) {
+      /* no change */
+    }
+    
+    void visitIsPattern(IsPattern *p) {
+      // Casting currently always consumes.
+      // TODO: Sometimes maybe it doesn't need to be.
+      increaseOwnership(ValueOwnership::Owned, p);
+    }
+    
+    void visitEnumElementPattern(EnumElementPattern *p) {
+      if (p->hasSubPattern()) {
+        visit(p->getSubPattern());
+      }
+    }
+    
+    void visitOptionalSomePattern(OptionalSomePattern *p) {
+      visit(p->getSubPattern());
+    }
+    
+    void visitExprPattern(ExprPattern *p) {
+      // We can't get the ownership reliably if the pattern hasn't been resolved.
+      if (!p->isResolved()) {
+        return;
+      }
+      increaseOwnership(p->getMatchOperandOwnership(), p);
+    }
+  };
+  
+  GetPatternOwnership visitor;
+  visitor.RestrictingPatterns = mostRestrictiveSubpatterns;
+  visitor.visit(const_cast<Pattern *>(this));
+  return visitor.Ownership;
+}

lib/AST/TypeCheckRequests.cpp

@@ -1029,13 +1029,13 @@ ExprPatternMatchRequest::getCachedResult() const {
   if (!EP->MatchVar)
     return llvm::None;
 
-  return ExprPatternMatchResult(EP->MatchVar, EP->MatchExpr);
+  return ExprPatternMatchResult(EP->MatchVar, EP->getCachedMatchExpr());
 }
 
 void ExprPatternMatchRequest::cacheResult(ExprPatternMatchResult result) const {
   auto *EP = std::get<0>(getStorage());
   EP->MatchVar = result.getMatchVar();
-  EP->MatchExpr = result.getMatchExpr();
+  EP->MatchExprAndOperandOwnership.setPointer(result.getMatchExpr());
 }
 
 //----------------------------------------------------------------------------//

lib/Serialization/ModuleFormat.h

@@ -58,7 +58,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 831; // mark_dependence [nonescaping]
+const uint16_t SWIFTMODULE_VERSION_MINOR = 832; // reorder ValueOwnership
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///

test/Sema/switch-ownership.swift

@@ -0,0 +1,58 @@
+// RUN: %target-swift-frontend -typecheck -dump-ast %s | %FileCheck %s
+
+struct A {
+    static func ~=(pattern: B, subject: A) -> Bool { return true }
+    static func ~=(pattern: C, subject: borrowing A) -> Bool { return true }
+    static func ~=(pattern: D, subject: consuming A) -> Bool { return true }
+}
+struct B { }
+struct C { }
+struct D { }
+struct E { }
+struct F { }
+struct G { }
+
+func ~=(pattern: E, subject: A) -> Bool { return true }
+func ~=(pattern: F, subject: borrowing A) -> Bool { return true }
+func ~=(pattern: G, subject: consuming A) -> Bool { return true }
+
+// CHECK-LABEL: (func_decl{{.*}} "test(value:)"
+func test(value: A) {
+    // CHECK: (switch_stmt
+    switch value {
+    // CHECK: (case_stmt
+    // CHECK:   (case_label_item ownership=borrowing
+    // CHECK:     (pattern_expr type="A" ownership=borrowing
+    case B():
+        break
+    // CHECK: (case_stmt
+    // CHECK:   (case_label_item ownership=borrowing
+    // CHECK:     (pattern_expr type="A" ownership=borrowing
+    case C():
+        break
+    // CHECK: (case_stmt
+    // CHECK:   (case_label_item ownership=consuming
+    // CHECK:     (pattern_expr type="A" ownership=consuming
+    case D():
+        break
+    // CHECK: (case_stmt
+    // CHECK:   (case_label_item ownership=borrowing
+    // CHECK:     (pattern_expr type="A" ownership=borrowing
+    case E():
+        break
+    // CHECK: (case_stmt
+    // CHECK:   (case_label_item ownership=borrowing
+    // CHECK:     (pattern_expr type="A" ownership=borrowing
+    case F():
+        break
+    // CHECK: (case_stmt
+    // CHECK:   (case_label_item ownership=consuming
+    // CHECK:     (pattern_expr type="A" ownership=consuming
+    case G():
+        break
+    // CHECK: (case_stmt
+    // CHECK:   (case_label_item default ownership=borrowing
+    default:
+        break
+    }
+}