Merge pull request #70039 from eeckstein/fix-let-property-lowering

LetPropertyLowering: handle `store_borrow` and re-borrows

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LetPropertyLowering.swift

@@ -135,8 +135,8 @@ private func constructLetInitRegion(
   // root-class initializer).
   initRegion.insert(markUninitialized)
 
-  var beginBorrows = Stack<BeginBorrowInst>(context)
-  defer { beginBorrows.deinitialize() }
+  var borrows = Stack<BorrowIntroducingInstruction>(context)
+  defer { borrows.deinitialize() }
 
   for inst in markUninitialized.parentFunction.instructions {
     switch inst {
@@ -161,8 +161,13 @@ private func constructLetInitRegion(
       // Include let-field partial de-initializations in the region.
       initRegion.insert(inst)
 
-    case let beginBorrow as BeginBorrowInst:
-      beginBorrows.append(beginBorrow)
+    case let beginBorrow as BeginBorrowInst
+         where beginBorrow.borrowedValue.referenceRoot == markUninitialized:
+      borrows.append(beginBorrow)
+
+    case let storeBorrow as StoreBorrowInst
+         where storeBorrow.source.referenceRoot == markUninitialized:
+      borrows.append(storeBorrow)
 
     default:
       break
@@ -171,8 +176,8 @@ private func constructLetInitRegion(
 
   // Extend the region to whole borrow scopes to avoid that we insert an `end_init_let_ref` in the
   // middle of a borrow scope.
-  for beginBorrow in beginBorrows where initRegion.contains(beginBorrow) {
-    initRegion.insert(contentsOf: beginBorrow.endBorrows)
+  for borrow in borrows where initRegion.contains(borrow) {
+    initRegion.insert(borrowScopeOf: borrow, context)
   }
 }
 

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -532,3 +532,25 @@ extension GlobalVariable {
     }
   }
 }
+
+extension InstructionRange {
+  /// Adds the instruction range of a borrow-scope by transitively visiting all (potential) re-borrows.
+  mutating func insert(borrowScopeOf borrow: BorrowIntroducingInstruction, _ context: some Context) {
+    var worklist = ValueWorklist(context)
+    defer { worklist.deinitialize() }
+
+    worklist.pushIfNotVisited(borrow)
+    while let value = worklist.pop() {
+      for use in value.uses {
+        switch use.instruction {
+        case let endBorrow as EndBorrowInst:
+          self.insert(endBorrow)
+        case let branch as BranchInst:
+          worklist.pushIfNotVisited(branch.getArgument(for: use))
+        default:
+          break
+        }
+      }
+    }
+  }
+}

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -234,7 +234,7 @@ extension UnaryInstruction {
 final public class UnimplementedInstruction : Instruction {
 }
 
-public protocol StoringInstruction : AnyObject {
+public protocol StoringInstruction : Instruction {
   var operands: OperandArray { get }
 }
 
@@ -270,6 +270,8 @@ final public class StoreInst : Instruction, StoringInstruction {
 final public class StoreWeakInst : Instruction, StoringInstruction { }
 final public class StoreUnownedInst : Instruction, StoringInstruction { }
 
+final public class StoreBorrowInst : SingleValueInstruction, StoringInstruction, BorrowIntroducingInstruction { }
+
 final public class AssignInst : Instruction, StoringInstruction {
   // must match with enum class swift::AssignOwnershipQualifier
   public enum AssignOwnership: Int {
@@ -495,6 +497,9 @@ extension LoadInstruction {
   public var address: Value { operand.value }
 }
 
+/// Instructions, beginning a borrow-scope which must be ended by `end_borrow`.
+public protocol BorrowIntroducingInstruction : SingleValueInstruction {}
+
 final public class LoadInst : SingleValueInstruction, LoadInstruction {
   // must match with enum class LoadOwnershipQualifier
   public enum LoadOwnership: Int {
@@ -507,7 +512,7 @@ final public class LoadInst : SingleValueInstruction, LoadInstruction {
 
 final public class LoadWeakInst : SingleValueInstruction, LoadInstruction {}
 final public class LoadUnownedInst : SingleValueInstruction, LoadInstruction {}
-final public class LoadBorrowInst : SingleValueInstruction, LoadInstruction {}
+final public class LoadBorrowInst : SingleValueInstruction, LoadInstruction, BorrowIntroducingInstruction {}
 
 final public class BuiltinInst : SingleValueInstruction {
   public typealias ID = BridgedInstruction.BuiltinValueKind
@@ -872,16 +877,9 @@ extension BeginAccessInst : ScopedInstruction {
   }
 }
 
-final public class BeginBorrowInst : SingleValueInstruction, UnaryInstruction {
+final public class BeginBorrowInst : SingleValueInstruction, UnaryInstruction, BorrowIntroducingInstruction {
   public var borrowedValue: Value { operand.value }
 
-  public typealias EndBorrowSequence = LazyMapSequence<LazyFilterSequence<LazyMapSequence<UseList, EndBorrowInst?>>,
-                                                       EndBorrowInst>
-
-  public var endBorrows: EndBorrowSequence {
-    uses.lazy.compactMap({ $0.instruction as? EndBorrowInst })
-  }
-
   public var isLexical: Bool { bridged.BeginBorrow_isLexical() }
   public var hasPointerEscape: Bool { bridged.BeginBorrow_hasPointerEscape() }
 }

SwiftCompilerSources/Sources/SIL/Registration.swift

@@ -41,6 +41,7 @@ public func registerSILClasses() {
   register(StoreInst.self)
   register(StoreWeakInst.self)
   register(StoreUnownedInst.self)
+  register(StoreBorrowInst.self)
   register(AssignInst.self)
   register(CopyAddrInst.self)
   register(EndAccessInst.self)

test/SILOptimizer/let-property-lowering.sil

@@ -254,3 +254,59 @@ bb3:
   return %2 : $C
 }
 
+// CHECK-LABEL: sil [ossa] @test_store_borrow :
+// CHECK:       bb1:
+// CHECK-NEXT:    end_borrow
+// CHECK-NEXT:    end_init_let_ref %2
+// CHECK:       bb2:
+// CHECK:         end_borrow {{%[0-9]}} : $*C
+// CHECK:         [[B:%.*]] = begin_borrow %2
+// CHECK:         end_borrow [[B]] : $C
+// CHECK-NEXT:    end_init_let_ref %2
+// CHECK:       } // end sil function 'test_store_borrow'
+sil [ossa] @test_store_borrow : $@convention(thin) (@owned C, Int) -> @owned C {
+bb0(%0 : @owned $C, %1 : $Int):
+  %2 = mark_uninitialized [rootself] %0 : $C
+  %3 = begin_borrow %2 : $C
+  %4 = ref_element_addr %3 : $C, #C.a
+  store %1 to [trivial] %4 : $*Int
+  end_borrow %3 : $C
+  %10 = alloc_stack $C
+  %11 = store_borrow %2 to %10 : $*C
+  cond_br undef, bb1, bb2
+bb1:
+  end_borrow %11 : $*C
+  dealloc_stack %10 : $*C
+  return %2 : $C
+bb2:
+  end_borrow %11 : $*C
+  %17 = begin_borrow %2 : $C
+  %18 = ref_element_addr %17 : $C, #C.a
+  %19 = begin_access [deinit] [static] %18 : $*Int
+  destroy_addr %19 : $*Int
+  end_access %19 : $*Int
+  end_borrow %17 : $C
+  unreachable
+}
+
+// CHECK-LABEL: sil [ossa] @test_reborrow :
+// CHECK:       bb1([[A:%.*]] : @reborrow @guaranteed $C):
+// CHECK:         end_borrow [[A]]
+// CHECK-NEXT:    end_init_let_ref %2
+// CHECK:       } // end sil function 'test_reborrow'
+sil [ossa] @test_reborrow : $@convention(thin) (@owned C, Int) -> @owned C {
+bb0(%0 : @owned $C, %1 : $Int):
+  %2 = mark_uninitialized [rootself] %0 : $C
+  %3 = begin_borrow %2 : $C
+  %4 = ref_element_addr %3 : $C, #C.a
+  store %1 to [trivial] %4 : $*Int
+  br bb1(%3 : $C)
+bb1(%3a : @reborrow @guaranteed $C):
+  end_borrow %3a : $C
+  %7 = begin_borrow %2 : $C
+  %8 = ref_element_addr %7 : $C, #C.a
+  %9 = load [trivial] %8 : $*Int
+  end_borrow %7 : $C
+  return %2 : $C
+}
+