Merge pull request #80083 from jckarter/addressable-buffer-begin-scope

SILGen: Adjust allocation point for addressable representations.

Author: jckarter

lib/SILGen/SILGenDecl.cpp

@@ -818,7 +818,7 @@ class LetValueInitialization : public Initialization {
     DestroyCleanup = SGF.Cleanups.getTopCleanup();
 
     // If the binding has an addressable buffer forced, it should be cleaned
-    // up here.
+    // up at this scope.
     SGF.enterLocalVariableAddressableBufferScope(vd);
   }
 
@@ -2461,6 +2461,8 @@ void SILGenFunction::destroyLocalVariable(SILLocation silLoc, VarDecl *vd) {
 
 void
 SILGenFunction::enterLocalVariableAddressableBufferScope(VarDecl *decl) {
+  auto marker = B.createTuple(decl, {});
+  AddressableBuffers[decl] = marker;
   Cleanups.pushCleanup<DeallocateLocalVariableAddressableBuffer>(decl);
 }
 
@@ -2513,7 +2515,28 @@ SILGenFunction::getLocalVariableAddressableBuffer(VarDecl *decl,
   SILValue reabstraction, allocStack, storeBorrow;
   {
     SavedInsertionPointRAII save(B);
-    B.setInsertionPoint(value->getNextInstruction());
+    auto insertPoint = AddressableBuffers[decl].insertPoint;
+    B.setInsertionPoint(insertPoint);
+    auto allocStackTy = fullyAbstractedTy;
+    if (value->getType().isMoveOnlyWrapped()) {
+      allocStackTy = allocStackTy.addingMoveOnlyWrapper();
+    }
+    allocStack = B.createAllocStack(decl,
+                                allocStackTy,
+                                std::nullopt,
+                                DoesNotHaveDynamicLifetime,
+                                IsNotLexical,
+                                IsNotFromVarDecl,
+                                DoesNotUseMoveableValueDebugInfo,
+                                /*skipVarDeclAssert*/ true);
+  }
+  {
+    SavedInsertionPointRAII save(B);
+    if (isa<ParamDecl>(decl)) {
+      B.setInsertionPoint(allocStack->getNextInstruction());
+    } else {
+      B.setInsertionPoint(value->getNextInstruction());
+    }
     auto declarationLoc = value->getDefiningInsertionPoint()->getLoc();
 
     // Reabstract if necessary.
@@ -2528,14 +2551,6 @@ SILGenFunction::getLocalVariableAddressableBuffer(VarDecl *decl,
       reabstraction = reabstracted.forward(*this);
       newValue = reabstraction;
     }
-    // TODO: reabstract
-    allocStack = B.createAllocStack(declarationLoc, newValue->getType(),
-                                    std::nullopt,
-                                    DoesNotHaveDynamicLifetime,
-                                    IsNotLexical,
-                                    IsNotFromVarDecl,
-                                    DoesNotUseMoveableValueDebugInfo,
-                                    /*skipVarDeclAssert*/ true);
     storeBorrow = B.createStoreBorrow(declarationLoc, newValue, allocStack);
   }
 

lib/SILGen/SILGenFunction.h

@@ -535,6 +535,37 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   /// a local variable.
   llvm::DenseMap<ValueDecl*, VarLoc> VarLocs;
 
+  // Represents an addressable buffer that has been allocated but not yet used.
+  struct PreparedAddressableBuffer {
+    SILInstruction *insertPoint = nullptr;
+    
+    PreparedAddressableBuffer() = default;
+    
+    PreparedAddressableBuffer(SILInstruction *insertPoint)
+      : insertPoint(insertPoint)
+    {}
+    
+    PreparedAddressableBuffer(PreparedAddressableBuffer &&other)
+      : insertPoint(other.insertPoint)
+    {
+      other.insertPoint = nullptr;
+    }
+    
+    PreparedAddressableBuffer &operator=(PreparedAddressableBuffer &&other) {
+      insertPoint = other.insertPoint;
+      other.insertPoint = nullptr;
+      return *this;
+    }
+    
+    ~PreparedAddressableBuffer() {
+      if (insertPoint) {
+        // Remove the insertion point if it went unused.
+        insertPoint->eraseFromParent();
+      }
+    }
+  };
+  llvm::DenseMap<VarDecl *, PreparedAddressableBuffer> AddressableBuffers;
+  
   /// Establish the scope for the addressable buffer that might be allocated
   /// for a local variable binding.
   ///

test/SILGen/addressable_representation.swift

@@ -134,9 +134,9 @@ func test4(_ c: CO1, _ c1: Bool, _ c2: Bool) {
 
 // CHECK-LABEL: sil {{.*}} @$s{{.*}}5test5
 func test5(_ f: @escaping () -> ()) {
+    // CHECK: [[ADDRESSABLE:%.*]] = alloc_stack
     // CHECK: [[THUNKED:%.*]] = partial_apply
     // CHECK: [[CONV:%.*]] = convert_function [[THUNKED]]
-    // CHECK: [[ADDRESSABLE:%.*]] = alloc_stack
     // CHECK: [[ADDRESSABLE_BORROW:%.*]] = store_borrow [[CONV]] to [[ADDRESSABLE]]
 
     // CHECK: function_ref @$s{{.*}}22addressableFunctionArg
@@ -181,6 +181,30 @@ func test7(_ c: CO1) {
   print(v)
 }
 
+// CHECK-LABEL: sil {{.*}} @$s{{.*}}5test8
+func test8() {
+  guard #available(Span 0.1, *) else { return }
+
+  var s = "A long string that is absolutely not smol at all."
+  let u = Array(s.utf8)
+
+  // CHECK: [[C_ADDRESSABLE:%.*]] = alloc_stack $CO2<String>
+  // CHECK: [[C_ADDRESSABLE_BORROW:%.*]] = store_borrow {{%.*}} to [[C_ADDRESSABLE]]
+  let c = CO2(consume s)
+  s = ""
+  // CHECK: [[GET_STORAGE:%.*]] = function_ref @$s{{.*}}7storage{{.*}}Gvg
+  // CHECK: apply [[GET_STORAGE]]<String>([[C_ADDRESSABLE_BORROW]]
+  let span = c.storage
+
+  _ = span.count == 1
+
+  // CHECK: function_ref @$s{{.*}}4utf8
+  let v = Array(span[0].utf8)
+  _ = u == v
+  // CHECK: end_borrow [[C_ADDRESSABLE_BORROW]]
+}
+
+
 func addressableFunctionArg(_ f: @_addressable @escaping () -> ()) {}
 
 @available(Span 0.1, *)
@@ -203,3 +227,23 @@ struct CO1 {
     }
 }
 
+@available(Span 0.1, *)
+@_addressableForDependencies
+struct CO2<T> {
+    var s: T
+
+    init(_ s: T) { self.s = s }
+
+    var storage: Span<T> {
+        @lifetime(borrow self)
+        borrowing get {
+            fatalError()
+        }
+    }
+
+    @lifetime(borrow self)
+    func getStorage() -> Span<T> {
+        fatalError()
+    }
+}
+