Merge pull request #72144 from nate-chandler/partial-consumption/20240306/1

nate-chandler · web-flow · commit 8ddbc8bf0e29 · 2024-03-07T10:51:55.000-08:00
[NoncopyablePartialConsumption] Allow consume.
diff --git a/include/swift/AST/DiagnosticsSema.def b/include/swift/AST/DiagnosticsSema.def
@@ -7741,6 +7741,16 @@ ERROR(moveOnly_not_allowed_here,none,
       "'@_moveOnly' attribute is only valid on structs or enums", ())
 ERROR(consume_expression_not_passed_lvalue,none,
       "'consume' can only be applied to a local binding ('let', 'var', or parameter)", ())
+ERROR(consume_expression_partial_copyable,none,
+      "'consume' can only be used to partially consume storage of a noncopyable type", ())
+ERROR(consume_expression_non_storage,none,
+      "'consume' can only be used to partially consume storage", ())
+NOTE(note_consume_expression_non_storage_call,none,
+      "non-storage produced by this call", ())
+NOTE(note_consume_expression_non_storage_subscript,none,
+      "non-storage produced by this subscript", ())
+NOTE(note_consume_expression_non_storage_property,none,
+      "non-storage produced by this computed property", ())
 ERROR(borrow_expression_not_passed_lvalue,none,
       "'borrow' can only be applied to a local binding ('let', 'var', or parameter)", ())
 ERROR(copy_expression_not_passed_lvalue,none,
diff --git a/lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerTester.cpp b/lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerTester.cpp
@@ -94,7 +94,8 @@ class MoveOnlyAddressCheckerTesterPass : public SILFunctionTransform {
     llvm::SmallSetVector<MarkUnresolvedNonCopyableValueInst *, 32>
         moveIntroducersToProcess;
     searchForCandidateAddressMarkUnresolvedNonCopyableValueInsts(
-        fn, moveIntroducersToProcess, diagnosticEmitter);
+        fn, getAnalysis<PostOrderAnalysis>(), moveIntroducersToProcess,
+        diagnosticEmitter);
 
     LLVM_DEBUG(llvm::dbgs()
                << "Emitting diagnostic when checking for mark must check inst: "
diff --git a/lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp b/lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp
@@ -479,14 +479,14 @@ static bool isCopyableValue(SILValue value) {
 
 void swift::siloptimizer::
     searchForCandidateAddressMarkUnresolvedNonCopyableValueInsts(
-        SILFunction *fn,
+        SILFunction *fn, PostOrderAnalysis *poa,
         llvm::SmallSetVector<MarkUnresolvedNonCopyableValueInst *, 32>
             &moveIntroducersToProcess,
         DiagnosticEmitter &diagnosticEmitter) {
-  for (auto &block : *fn) {
-    for (auto ii = block.begin(), ie = block.end(); ii != ie;) {
-      auto *mmci = dyn_cast<MarkUnresolvedNonCopyableValueInst>(&*ii);
-      ++ii;
+  auto *po = poa->get(fn);
+  for (auto *block : po->getPostOrder()) {
+    for (auto &ii : llvm::make_range(block->rbegin(), block->rend())) {
+      auto *mmci = dyn_cast<MarkUnresolvedNonCopyableValueInst>(&ii);
 
       if (!mmci || !mmci->hasMoveCheckerKind() || !mmci->getType().isAddress())
         continue;
@@ -3877,6 +3877,9 @@ bool MoveOnlyAddressChecker::check(
       // If we fail the address check in some way, set the diagnose!
       diagnosticEmitter.emitCheckerDoesntUnderstandDiagnostic(markedValue);
     }
+
+    markedValue->replaceAllUsesWith(markedValue->getOperand());
+    markedValue->eraseFromParent();
   }
 
   if (DumpSILBeforeRemovingMarkUnresolvedNonCopyableValueInst) {
diff --git a/lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.h b/lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.h
@@ -17,6 +17,8 @@
 
 namespace swift {
 
+class PostOrderAnalysis;
+
 namespace siloptimizer {
 
 class DiagnosticEmitter;
@@ -26,7 +28,7 @@ class DiagnosticEmitter;
 /// NOTE: To see if we emitted a diagnostic, use \p
 /// diagnosticEmitter.getDiagnosticCount().
 void searchForCandidateAddressMarkUnresolvedNonCopyableValueInsts(
-    SILFunction *fn,
+    SILFunction *fn, PostOrderAnalysis *poa,
     llvm::SmallSetVector<MarkUnresolvedNonCopyableValueInst *, 32>
         &moveIntroducersToProcess,
     DiagnosticEmitter &diagnosticEmitter);
diff --git a/lib/SILOptimizer/Mandatory/MoveOnlyChecker.cpp b/lib/SILOptimizer/Mandatory/MoveOnlyChecker.cpp
@@ -186,7 +186,7 @@ void MoveOnlyChecker::checkAddresses() {
   llvm::SmallSetVector<MarkUnresolvedNonCopyableValueInst *, 32>
       moveIntroducersToProcess;
   searchForCandidateAddressMarkUnresolvedNonCopyableValueInsts(
-      fn, moveIntroducersToProcess, diagnosticEmitter);
+      fn, poa, moveIntroducersToProcess, diagnosticEmitter);
 
   LLVM_DEBUG(
       llvm::dbgs()
diff --git a/lib/SILOptimizer/Mandatory/MoveOnlyTempAllocationFromLetTester.cpp b/lib/SILOptimizer/Mandatory/MoveOnlyTempAllocationFromLetTester.cpp
@@ -51,7 +51,8 @@ struct MoveOnlyTempAllocationFromLetTester : SILFunctionTransform {
 
     unsigned diagCount = diagnosticEmitter.getDiagnosticCount();
     searchForCandidateAddressMarkUnresolvedNonCopyableValueInsts(
-        getFunction(), moveIntroducersToProcess, diagnosticEmitter);
+        getFunction(), getAnalysis<PostOrderAnalysis>(),
+        moveIntroducersToProcess, diagnosticEmitter);
 
     // Return early if we emitted a diagnostic.
     if (diagCount != diagnosticEmitter.getDiagnosticCount())
diff --git a/lib/Sema/MiscDiagnostics.cpp b/lib/Sema/MiscDiagnostics.cpp
@@ -420,12 +420,89 @@ static void diagSyntacticUseRestrictions(const Expr *E, const DeclContext *DC,
     }
 
     void checkConsumeExpr(ConsumeExpr *consumeExpr) {
-      auto *subExpr = consumeExpr->getSubExpr();
-      if (auto *li = dyn_cast<LoadExpr>(subExpr))
-        subExpr = li->getSubExpr();
-      if (!isa<DeclRefExpr>(subExpr)) {
+      auto partialConsumptionEnabled =
+          Ctx.LangOpts.hasFeature(Feature::MoveOnlyPartialConsumption);
+
+      bool noncopyable = false;
+      bool partial = false;
+      Expr *current = consumeExpr->getSubExpr();
+      while (current) {
+        if (current->getType()->getCanonicalType()->isNoncopyable()) {
+          noncopyable = true;
+        }
+        if (auto *dre = dyn_cast<DeclRefExpr>(current)) {
+          if (partial & !noncopyable) {
+            Ctx.Diags.diagnose(consumeExpr->getLoc(),
+                               diag::consume_expression_partial_copyable);
+            return;
+          }
+          // The chain of member_ref_exprs and load_exprs terminates at a
+          // declref_expr.  This is legal.
+          return;
+        }
+        // Look through loads.
+        if (auto *le = dyn_cast<LoadExpr>(current)) {
+          current = le->getSubExpr();
+          continue;
+        }
+        auto *mre = dyn_cast<MemberRefExpr>(current);
+        if (mre && partialConsumptionEnabled) {
+          auto *vd = dyn_cast<VarDecl>(mre->getMember().getDecl());
+          if (!vd) {
+            Ctx.Diags.diagnose(consumeExpr->getLoc(),
+                               diag::consume_expression_non_storage);
+            return;
+          }
+          partial = true;
+          AccessStrategy strategy = vd->getAccessStrategy(
+              mre->getAccessSemantics(), AccessKind::Read,
+              DC->getParentModule(), ResilienceExpansion::Minimal);
+          if (strategy.getKind() != AccessStrategy::Storage) {
+            if (noncopyable) {
+              Ctx.Diags.diagnose(consumeExpr->getLoc(),
+                                 diag::consume_expression_non_storage);
+              Ctx.Diags.diagnose(
+                  mre->getLoc(),
+                  diag::note_consume_expression_non_storage_property);
+            } else {
+              Ctx.Diags.diagnose(consumeExpr->getLoc(),
+                                 diag::consume_expression_partial_copyable);
+            }
+            return;
+          }
+          current = mre->getBase();
+          continue;
+        }
+        auto *ce = dyn_cast<CallExpr>(current);
+        if (ce && partialConsumptionEnabled) {
+          if (noncopyable) {
+            Ctx.Diags.diagnose(consumeExpr->getLoc(),
+                               diag::consume_expression_non_storage);
+            Ctx.Diags.diagnose(ce->getLoc(),
+                               diag::note_consume_expression_non_storage_call);
+          } else {
+            Ctx.Diags.diagnose(consumeExpr->getLoc(),
+                               diag::consume_expression_partial_copyable);
+          }
+          return;
+        }
+        auto *se = dyn_cast<SubscriptExpr>(current);
+        if (se && partialConsumptionEnabled) {
+          if (noncopyable) {
+            Ctx.Diags.diagnose(consumeExpr->getLoc(),
+                               diag::consume_expression_non_storage);
+            Ctx.Diags.diagnose(
+                se->getLoc(),
+                diag::note_consume_expression_non_storage_subscript);
+          } else {
+            Ctx.Diags.diagnose(consumeExpr->getLoc(),
+                               diag::consume_expression_partial_copyable);
+          }
+          return;
+        }
         Ctx.Diags.diagnose(consumeExpr->getLoc(),
                            diag::consume_expression_not_passed_lvalue);
+        return;
       }
     }
 
diff --git a/test/Interop/Cxx/operators/move-only/move-only-synthesized-property-typecheck.swift b/test/Interop/Cxx/operators/move-only/move-only-synthesized-property-typecheck.swift
@@ -30,8 +30,8 @@ func testNonCopyableHolderConstDerefPointee() {
     holder.pointee.mutMethod(1) // expected-error {{cannot use mutating member on immutable value: 'pointee' is a get-only property}}
     holder.pointee.x = 2 // expected-error {{cannot assign to property: 'pointee' is a get-only property}}
 #else
-    consumingNC(holder.pointee) // CHECK: [[@LINE]]:{{.*}}: error:
-    let consumeVal = holder.pointee // CHECK: [[@LINE]]:{{.*}}: error:
+    consumingNC(holder.pointee) // CHECK-DAG: [[@LINE]]:{{.*}}: error:
+    let consumeVal = holder.pointee // CHECK-DAG: [[@LINE]]:{{.*}}: error:
 #endif
 }
 
diff --git a/test/Interpreter/move_expr_moveonly_partial_consumption.swift b/test/Interpreter/move_expr_moveonly_partial_consumption.swift
@@ -0,0 +1,115 @@
+// RUN: %target-run-simple-swift(-parse-as-library -Xfrontend -sil-verify-all -enable-experimental-feature MoveOnlyPartialConsumption) | %FileCheck %s
+// RUN: %target-run-simple-swift(-parse-as-library -O -Xfrontend -sil-verify-all -enable-experimental-feature MoveOnlyPartialConsumption) | %FileCheck %s
+
+// REQUIRES: executable_test
+
+@main struct App { static func main() {
+  test1((1,2,3,4))
+}}
+
+func barrier() { print("barrier") }
+
+struct Ur<T> : ~Copyable {
+  var t: T
+  var name: String
+  init(_ t: T, named name: String) {
+    self.t = t
+    self.name = name
+    print("hi", name)
+  }
+  deinit {
+    print("bye", name)
+  }
+}
+func take<T>(_ u: consuming Ur<T>) {}
+
+struct Pair<T> : ~Copyable {
+  var u1: Ur<T>
+  var u2: Ur<T>
+  init(_ t1: T, _ t2: T, named name: String) {
+    u1 = .init(t1, named: "\(name).u1")
+    u2 = .init(t2, named: "\(name).u2")
+  }
+}
+func take<T>(_ u: consuming Pair<T>) {}
+
+struct Quad<T> : ~Copyable {
+  var p1: Pair<T>
+  var p2: Pair<T>
+  init(_ t1: T, _ t2: T, _ t3: T, _ t4: T, named name: String) {
+    p1 = .init(t1, t2, named: "\(name).p1")
+    p2 = .init(t3, t4, named: "\(name).p2")
+  }
+}
+func take<T>(_ u: consuming Quad<T>) {}
+
+func test1<T>(_ ts: (T, T, T, T)) {
+do {
+  let q1 = Quad<T>(ts.0, ts.1, ts.2, ts.3, named: "\(#function).q1")
+  barrier()
+  // CHECK: barrier
+  // CHECK: bye test1(_:).q1.p1.u1
+  // CHECK: bye test1(_:).q1.p1.u2
+  // CHECK: bye test1(_:).q1.p2.u1
+  // CHECK: bye test1(_:).q1.p2.u2
+}
+
+  let q2 = Quad<T>(ts.0, ts.1, ts.2, ts.3, named: "\(#function).q2")
+  take(q2.p2.u2)
+  take(q2.p2.u1)
+  barrier()
+  take(q2.p1.u2)
+  take(q2.p1.u1)
+  // CHECK: bye test1(_:).q2.p2.u2
+  // CHECK: bye test1(_:).q2.p2.u1
+  // CHECK: barrier
+  // CHECK: bye test1(_:).q2.p1.u2
+  // CHECK: bye test1(_:).q2.p1.u1
+
+  let q3 = Quad<T>(ts.0, ts.1, ts.2, ts.3, named: "\(#function).q3")
+  _ = consume q3.p2.u2
+  _ = consume q3.p2.u1
+  _ = consume q3.p1.u2
+  barrier()
+  _ = consume q3.p1.u1
+  // CHECK: bye test1(_:).q3.p2.u2
+  // CHECK: bye test1(_:).q3.p2.u1
+  // CHECK: bye test1(_:).q3.p1.u2
+  // CHECK: barrier
+  // CHECK: bye test1(_:).q3.p1.u1
+
+  let q4 = Quad<T>(ts.0, ts.1, ts.2, ts.3, named: "\(#function).q4")
+  _ = consume q4.p1.u1
+  barrier()
+  _ = consume q4.p2.u1
+  _ = consume q4.p1.u2
+  _ = consume q4.p2.u2
+  // CHECK: bye test1(_:).q4.p1.u1
+  // CHECK: barrier
+  // CHECK: bye test1(_:).q4.p2.u1
+  // CHECK: bye test1(_:).q4.p1.u2
+  // CHECK: bye test1(_:).q4.p2.u2
+
+do {
+  let q5 = Quad<T>(ts.0, ts.1, ts.2, ts.3, named: "\(#function).q5")
+  take(q5.p1.u1)
+  _ = consume q5.p2.u1
+  _ = consume q5.p1.u2
+  take(q5.p2.u2)
+  // CHECK: bye test1(_:).q5.p1.u1
+  // CHECK: bye test1(_:).q5.p2.u1
+  // CHECK: bye test1(_:).q5.p1.u2
+  // CHECK: bye test1(_:).q5.p2.u2
+}
+do {
+  let q6 = Quad<T>(ts.0, ts.1, ts.2, ts.3, named: "\(#function).q6")
+  if Bool.random() {
+    take(q6.p1.u1)
+    _ = consume q6.p2.u1
+  } else {
+    _ = consume q6.p1.u2
+    take(q6.p2.u2)
+  }
+}
+}
+
diff --git a/test/Sema/move_expr_moveonly_partial_consumption.swift b/test/Sema/move_expr_moveonly_partial_consumption.swift
