Merge pull request #71795 from jckarter/yield-read-coroutine-result

jckarter · web-flow · commit 7f9b18432a90 · 2024-02-21T20:40:36.000-08:00
SILGen: Emit the base of a read coroutine as a borrow in a `yield`/`borrow`/noncopyable context.
diff --git a/lib/SILGen/SILGen.h b/lib/SILGen/SILGen.h
@@ -418,12 +418,6 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   /// Emit a global initialization.
   void emitGlobalInitialization(PatternBindingDecl *initializer, unsigned elt);
 
-  /// Should the self argument of the given method always be emitted as
-  /// an r-value (meaning that it can be borrowed only if that is not
-  /// semantically detectable), or it acceptable to emit it as a borrowed
-  /// storage reference?
-  bool shouldEmitSelfAsRValue(FuncDecl *method, CanType selfType);
-
   /// Is the self method of the given nonmutating method passed indirectly?
   bool isNonMutatingSelfIndirect(SILDeclRef method);
 
diff --git a/lib/SILGen/SILGenApply.cpp b/lib/SILGen/SILGenApply.cpp
@@ -5434,22 +5434,6 @@ CallEmission CallEmission::forApplyExpr(SILGenFunction &SGF, ApplyExpr *e) {
   return emission;
 }
 
-bool SILGenModule::shouldEmitSelfAsRValue(FuncDecl *fn, CanType selfType) {
-  if (fn->isStatic())
-    return true;
-
-  switch (fn->getSelfAccessKind()) {
-  case SelfAccessKind::Mutating:
-    return false;
-  case SelfAccessKind::NonMutating:
-  case SelfAccessKind::LegacyConsuming:
-  case SelfAccessKind::Consuming:
-  case SelfAccessKind::Borrowing:
-    return true;
-  }
-  llvm_unreachable("bad self-access kind");
-}
-
 bool SILGenModule::isNonMutatingSelfIndirect(SILDeclRef methodRef) {
   auto method = methodRef.getFuncDecl();
   if (method->isStatic())
diff --git a/lib/SILGen/SILGenFunction.h b/lib/SILGen/SILGenFunction.h
@@ -111,7 +111,7 @@ enum class SGFAccessKind : uint8_t {
   /// The access is a read that would prefer the address of a borrowed value.
   /// This should only be used when it is semantically acceptable to borrow
   /// the value, not just because the caller would benefit from a borrowed
-  /// value.  See shouldEmitSelfAsRValue.
+  /// value.  See shouldEmitSelfAsRValue in SILGenLValue.cpp.
   ///
   /// The caller will be calling emitAddressOfLValue or emitLoadOfLValue
   /// on the l-value.  The latter may be less efficient than an access
@@ -121,7 +121,7 @@ enum class SGFAccessKind : uint8_t {
   /// The access is a read that would prefer a loaded borrowed value.
   /// This should only be used when it is semantically acceptable to borrow
   /// the value, not just because the caller would benefit from a borrowed
-  /// value.  See shouldEmitSelfAsRValue.
+  /// value.  See shouldEmitSelfAsRValue in SILGenLValue.cpp.
   ///
   /// There isn't yet a way to emit the access that takes advantage of this.
   BorrowedObjectRead,
diff --git a/lib/SILGen/SILGenLValue.cpp b/lib/SILGen/SILGenLValue.cpp
@@ -1338,7 +1338,8 @@ static SGFAccessKind getBaseAccessKind(SILGenModule &SGM,
                                        AbstractStorageDecl *member,
                                        SGFAccessKind accessKind,
                                        AccessStrategy strategy,
-                                       CanType baseFormalType);
+                                       CanType baseFormalType,
+                                       bool forBorrowExpr);
 
 namespace {
   /// A helper class for implementing components that involve accessing
@@ -1971,7 +1972,8 @@ namespace {
         if (!base) return LValue();
         auto baseAccessKind =
           getBaseAccessKind(SGF.SGM, Storage, accessKind, strategy,
-                            BaseFormalType);
+                            BaseFormalType,
+                            /*for borrow*/ false);
         return LValue::forValue(baseAccessKind, base, BaseFormalType);
       }();
 
@@ -3049,7 +3051,8 @@ class LLVM_LIBRARY_VISIBILITY SILGenBorrowedBaseVisitor
     auto baseFormalType = getBaseFormalType(e->getBase());
     LValue lv = visit(
         e->getBase(),
-        getBaseAccessKind(SGF.SGM, var, accessKind, strategy, baseFormalType),
+        getBaseAccessKind(SGF.SGM, var, accessKind, strategy, baseFormalType,
+                          /*for borrow*/ true),
         getBaseOptions(options, strategy));
     std::optional<ActorIsolation> actorIso;
     if (e->isImplicitlyAsync())
@@ -3717,14 +3720,50 @@ LValue SILGenLValue::visitDotSyntaxBaseIgnoredExpr(DotSyntaxBaseIgnoredExpr *e,
   return visitRec(e->getRHS(), accessKind, options);
 }
 
+/// Should the self argument of the given method always be emitted as
+/// an r-value (meaning that it can be borrowed only if that is not
+/// semantically detectable), or it acceptable to emit it as a borrowed
+/// storage reference?
+static bool shouldEmitSelfAsRValue(AccessorDecl *fn, CanType selfType,
+                                   bool forBorrowExpr) {
+  if (fn->isStatic())
+    return true;
+
+  switch (fn->getSelfAccessKind()) {
+  case SelfAccessKind::Mutating:
+    return false;
+  case SelfAccessKind::Borrowing:
+  case SelfAccessKind::NonMutating:
+    // If the accessor is a coroutine, we may want to access the projected
+    // value through a borrow of the base. But if it's a regular get/set then
+    // there isn't any real benefit to doing so.
+    if (!fn->isCoroutine()) {
+      return true;
+    }
+    // Normally we'll copy the base to minimize accesses. But if the base
+    // is noncopyable, or we're accessing it in a `borrow` expression, then
+    // we want to keep the access nested on the original base.
+    if (forBorrowExpr || selfType->isNoncopyable()) {
+      return false;
+    }
+    return true;
+
+  case SelfAccessKind::LegacyConsuming:
+  case SelfAccessKind::Consuming:
+    return true;
+  }
+  llvm_unreachable("bad self-access kind");
+}
+
 static SGFAccessKind getBaseAccessKindForAccessor(SILGenModule &SGM,
                                                   AccessorDecl *accessor,
-                                                  CanType baseFormalType) {
+                                                  CanType baseFormalType,
+                                                  bool forBorrowExpr) {
   if (accessor->isMutating())
     return SGFAccessKind::ReadWrite;
 
   auto declRef = SGM.getAccessorDeclRef(accessor, ResilienceExpansion::Minimal);
-  if (SGM.shouldEmitSelfAsRValue(accessor, baseFormalType)) {
+  if (shouldEmitSelfAsRValue(accessor, baseFormalType, forBorrowExpr)) {
     return SGM.isNonMutatingSelfIndirect(declRef)
                ? SGFAccessKind::OwnedAddressRead
                : SGFAccessKind::OwnedObjectRead;
@@ -3748,7 +3787,8 @@ static SGFAccessKind getBaseAccessKind(SILGenModule &SGM,
                                        AbstractStorageDecl *member,
                                        SGFAccessKind accessKind,
                                        AccessStrategy strategy,
-                                       CanType baseFormalType) {
+                                       CanType baseFormalType,
+                                       bool forBorrowExpr) {
   switch (strategy.getKind()) {
   case AccessStrategy::Storage:
     return getBaseAccessKindForStorage(accessKind);
@@ -3757,7 +3797,8 @@ static SGFAccessKind getBaseAccessKind(SILGenModule &SGM,
     assert(accessKind == SGFAccessKind::ReadWrite);
     auto writeBaseKind = getBaseAccessKind(SGM, member, SGFAccessKind::Write,
                                            strategy.getWriteStrategy(),
-                                           baseFormalType);
+                                           baseFormalType,
+                                           /*for borrow*/ false);
 
     // Fast path for the common case that the write will need to mutate
     // the base.
@@ -3767,7 +3808,8 @@ static SGFAccessKind getBaseAccessKind(SILGenModule &SGM,
     auto readBaseKind = getBaseAccessKind(SGM, member,
                                           SGFAccessKind::OwnedAddressRead,
                                           strategy.getReadStrategy(),
-                                          baseFormalType);
+                                          baseFormalType,
+                                          /*for borrow*/ false);
 
     // If they're the same kind, just use that.
     if (readBaseKind == writeBaseKind)
@@ -3786,7 +3828,8 @@ static SGFAccessKind getBaseAccessKind(SILGenModule &SGM,
   case AccessStrategy::DispatchToAccessor:
   case AccessStrategy::DispatchToDistributedThunk: {
     auto accessor = member->getOpaqueAccessor(strategy.getAccessor());
-    return getBaseAccessKindForAccessor(SGM, accessor, baseFormalType);
+    return getBaseAccessKindForAccessor(SGM, accessor, baseFormalType,
+                                        forBorrowExpr);
   }
   }
   llvm_unreachable("bad access strategy");
@@ -3863,7 +3906,8 @@ LValue SILGenLValue::visitMemberRefExpr(MemberRefExpr *e,
   
   LValue lv = visitRec(e->getBase(),
                        getBaseAccessKind(SGF.SGM, var, accessKind, strategy,
-                                         getBaseFormalType(e->getBase())),
+                                         getBaseFormalType(e->getBase()),
+                                         /* for borrow */ false),
                        getBaseOptions(options, strategy));
   assert(lv.isValid());
 
@@ -4069,7 +4113,8 @@ LValue SILGenLValue::visitSubscriptExpr(SubscriptExpr *e,
 
   LValue lv = visitRec(e->getBase(),
                        getBaseAccessKind(SGF.SGM, decl, accessKind, strategy,
-                                         getBaseFormalType(e->getBase())),
+                                         getBaseFormalType(e->getBase()),
+                                         /*for borrow*/ false),
                        getBaseOptions(options, strategy));
   assert(lv.isValid());
 
@@ -4426,7 +4471,8 @@ LValue SILGenFunction::emitPropertyLValue(SILLocation loc, ManagedValue base,
                             F.getResilienceExpansion());
 
   auto baseAccessKind =
-    getBaseAccessKind(SGM, ivar, accessKind, strategy, baseFormalType);
+    getBaseAccessKind(SGM, ivar, accessKind, strategy, baseFormalType,
+                      /*for borrow*/ false);
 
   LValueTypeData baseTypeData =
     getValueTypeData(baseAccessKind, baseFormalType, base.getValue());
@@ -5149,7 +5195,8 @@ RValue SILGenFunction::emitRValueForStorageLoad(
       if (!base) return LValue();
 
       auto baseAccess = getBaseAccessKind(SGM, storage, accessKind,
-                                          strategy, baseFormalType);
+                                          strategy, baseFormalType,
+                                          /*for borrow*/ false);
       return LValue::forValue(baseAccess, base, baseFormalType);
     }();
 
diff --git a/test/SILOptimizer/moveonly_borrowing_switch_yield.swift b/test/SILOptimizer/moveonly_borrowing_switch_yield.swift
@@ -0,0 +1,126 @@
+// RUN: %target-swift-frontend -enable-experimental-feature NoncopyableGenerics -enable-experimental-feature BorrowingSwitch -enable-experimental-feature MoveOnlyPartialConsumption -parse-as-library -O -emit-sil -verify %s
+
+extension List {
+    var peek: Element {
+        _read {
+            // TODO: fix move-only checker induced ownership bug when
+            // we try to switch over `self.head` here.
+            yield head.peek
+        }
+    }
+}
+
+struct MyPointer<Wrapped: ~Copyable>: Copyable {
+    var v: UnsafeMutablePointer<Int>
+
+    static func allocate(capacity: Int) -> Self {
+        fatalError()
+    }
+
+    func initialize(to: consuming Wrapped) {
+    }
+    func deinitialize(count: Int) {
+    }
+    func deallocate() {
+    }
+    func move() -> Wrapped {
+        fatalError()
+    }
+
+    var pointee: Wrapped {
+        _read { fatalError() }
+    }
+}
+
+struct Box<Wrapped: ~Copyable>: ~Copyable {
+    private let _pointer: MyPointer<Wrapped>
+    
+    init(_ element: consuming Wrapped) {
+        _pointer = .allocate(capacity: 1)
+        print("allocatin", _pointer)
+        _pointer.initialize(to: element)
+    }
+        
+    deinit {
+        print("not deallocatin", _pointer)
+        _pointer.deinitialize(count: 1)
+        _pointer.deallocate()
+    }
+    
+    consuming func move() -> Wrapped {
+        let wrapped = _pointer.move()
+        print("deallocatin", _pointer)
+        _pointer.deallocate()
+        discard self
+        return wrapped
+    }
+    
+    var wrapped: Wrapped {
+        _read { yield _pointer.pointee }
+    }
+}
+
+struct List<Element>: ~Copyable {
+    var head: Link<Element> = .empty
+    var bool = false
+}
+
+enum Link<Element>: ~Copyable {
+    case empty
+    case more(Box<Node<Element>>)
+
+    var peek: Element {
+        _read {
+            switch self {
+            case .empty: fatalError()
+            case .more(_borrowing box):
+                yield box.wrapped.element
+            }
+        }
+    }
+}
+
+struct Node<Element>: ~Copyable {
+    let element: Element
+    let next: Link<Element>
+    var bool = true
+}
+
+extension List {
+    mutating func append(_ element: consuming Element) {
+        self = List(head: .more(Box(Node(element: element, next: self.head))))
+    }
+
+    var isEmpty: Bool {
+        switch self.head {
+        case .empty: true
+        default: false
+        }
+    }
+    
+    mutating func pop() -> Element {
+        let h = self.head
+        switch h {
+        case .empty: fatalError()
+        case .more(let box):
+            let node = box.move()
+            self = .init(head: node.next)
+            return node.element
+        }
+    }
+
+}
+
+@main
+struct Main {
+    static func main() {
+        var l: List<Int> = .init()
+        l.append(1)
+        l.append(2)
+        l.append(3)
+        print(l.pop())
+        print(l.pop())
+        print(l.pop())
+        print(l.isEmpty)
+    }
+}
diff --git a/test/SILOptimizer/moveonly_consuming_switch.swift b/test/SILOptimizer/moveonly_consuming_switch.swift
@@ -1,4 +1,7 @@
 // RUN: %target-swift-frontend -emit-sil -enable-experimental-feature NoncopyableGenerics -enable-experimental-feature BorrowingSwitch -enable-experimental-feature MoveOnlyPartialConsumption -verify %s
+
+// TODO: Remove this and just use the real `UnsafeMutablePointer` when
+// noncopyable type support has been upstreamed.
 struct MyPointer<Wrapped: ~Copyable>: Copyable {
     var v: UnsafeMutablePointer<Int>
 
@@ -41,14 +44,6 @@ struct Box<Wrapped: ~Copyable>: ~Copyable {
     }
 }
 
-
-
-
-
-
-
-
-
 enum List<Element: ~Copyable>: ~Copyable {
     case empty
     case node(Element, Box<List>)
