Merge pull request #79978 from jckarter/addressable-representation-silgen-2

SILGen: Emit an addressable representation for immutable bindings on demand. [take 2]

Author: jckarter

include/swift/AST/Types.h

@@ -4331,9 +4331,7 @@ inline bool isConsumedParameterInCaller(ParameterConvention conv) {
   return isConsumedParameter<false>(conv);
 }
 
-/// Returns true if conv is a guaranteed parameter. This may look unnecessary
-/// but this will allow code to generalize to handle Indirect_Guaranteed
-/// parameters when they are added.
+/// Returns true if conv is a guaranteed parameter.
 template <bool InCallee>
 bool isGuaranteedParameter(ParameterConvention conv) {
   switch (conv) {

lib/SIL/IR/AbstractionPattern.cpp

@@ -1679,9 +1679,7 @@ AbstractionPattern::isFunctionParamAddressable(TypeConverter &TC,
     auto type = getType();
 
     if (type->isTypeParameter() || type->is<ArchetypeType>()) {
-      // If the function abstraction pattern is completely opaque, assume we
-      // may need to preserve the address for dependencies.
-      return true;    
+      return false;
     }
 
     auto fnTy = cast<AnyFunctionType>(getType());

lib/SILGen/SILGenApply.cpp

@@ -3462,11 +3462,9 @@ SILGenFunction::tryEmitAddressableParameterAsAddress(ArgumentSource &&arg,
       expr = le->getSubExpr();
     }
     if (auto dre = dyn_cast<DeclRefExpr>(expr)) {
-      if (auto param = dyn_cast<ParamDecl>(dre->getDecl())) {
-        if (VarLocs.count(param)
-            && VarLocs[param].addressable
-            && param->getValueOwnership() == ownership) {
-          auto addr = VarLocs[param].value;
+      if (auto param = dyn_cast<VarDecl>(dre->getDecl())) {
+        if (auto addr = getLocalVariableAddressableBuffer(param, expr,
+                                                          ownership)) {
           return ManagedValue::forBorrowedAddressRValue(addr);
         }
       }
@@ -3577,12 +3575,20 @@ class ArgEmitter {
   void emit(ArgumentSource &&arg, AbstractionPattern origParamType,
             bool isAddressable,
             std::optional<AnyFunctionType::Param> origParam = std::nullopt) {
-    if (isAddressable && origParam) {
+    if (isAddressable) {
       // If the function takes an addressable parameter, and its argument is
       // a reference to an addressable declaration with compatible ownership,
       // forward the address along in-place.
+      ValueOwnership paramOwnership;
+      if (isGuaranteedParameterInCaller(ParamInfos.front().getConvention())) {
+        paramOwnership = ValueOwnership::Shared;
+      } else if (isMutatingParameter(ParamInfos.front().getConvention())) {
+        paramOwnership = ValueOwnership::InOut;
+      } else {
+        paramOwnership = ValueOwnership::Owned;      
+      }
       if (auto addr = SGF.tryEmitAddressableParameterAsAddress(std::move(arg),
-                                             origParam->getValueOwnership())) {
+                                                              paramOwnership)) {
         claimNextParameters(1);
         Args.push_back(addr);
         return;

lib/SILGen/SILGenBuilder.cpp

@@ -1081,6 +1081,18 @@ class EndAccessCleanup final : public Cleanup {
 };
 }
 
+SILValue
+SILGenBuilder::emitBeginAccess(SILLocation loc,
+                               SILValue address,
+                               SILAccessKind kind,
+                               SILAccessEnforcement enforcement) {
+  auto access = createBeginAccess(loc, address,
+                                  kind, enforcement,
+                                  /*no nested conflict*/ false, false);
+  SGF.Cleanups.pushCleanup<EndAccessCleanup>(access);
+  return access;
+}
+
 ManagedValue
 SILGenBuilder::createOpaqueBorrowBeginAccess(SILLocation loc,
                                              ManagedValue address) {

lib/SILGen/SILGenBuilder.h

@@ -477,6 +477,10 @@ class SILGenBuilder : public SILBuilder {
                                     ManagedValue base,
                                     MarkDependenceKind dependencekind);
 
+  SILValue emitBeginAccess(SILLocation loc, SILValue address,
+                           SILAccessKind kind,
+                           SILAccessEnforcement enforcement);
+
   ManagedValue createOpaqueBorrowBeginAccess(SILLocation loc,
                                              ManagedValue address);
   ManagedValue createOpaqueConsumeBeginAccess(SILLocation loc,

lib/SILGen/SILGenConstructor.cpp

@@ -1173,7 +1173,8 @@ void SILGenFunction::emitClassConstructorInitializer(ConstructorDecl *ctor) {
             MarkUnresolvedNonCopyableValueInst::CheckKind::
                 ConsumableAndAssignable);
       }
-      VarLocs[selfDecl] = VarLoc::get(selfArg.getValue());
+      VarLocs[selfDecl] = VarLoc(selfArg.getValue(),
+                                 SILAccessEnforcement::Static);
     }
   }
 
@@ -1696,7 +1697,7 @@ void SILGenFunction::emitIVarInitializer(SILDeclRef ivarInitializer) {
   selfArg = B.createMarkUninitialized(selfDecl, selfArg,
                                       MarkUninitializedInst::RootSelf);
   assert(selfTy.hasReferenceSemantics() && "can't emit a value type ctor here");
-  VarLocs[selfDecl] = VarLoc::get(selfArg);
+  VarLocs[selfDecl] = VarLoc(selfArg, SILAccessEnforcement::Unknown);
 
   auto cleanupLoc = CleanupLocation(loc);
   prepareEpilog(cd, std::nullopt, std::nullopt, cleanupLoc);
@@ -1727,11 +1728,11 @@ void SILGenFunction::emitInitAccessor(AccessorDecl *accessor) {
     loc.markAutoGenerated();
 
     SILValue argValue = F.begin()->createFunctionArgument(type, arg);
-    VarLocs[arg] =
-        markUninitialized
-            ? VarLoc::get(B.createMarkUninitializedOut(loc, argValue))
-            : VarLoc::get(argValue);
-
+    if (markUninitialized) {
+      argValue = B.createMarkUninitializedOut(loc, argValue);
+    }
+    
+    VarLocs[arg] = VarLoc(argValue, SILAccessEnforcement::Static);
     InitAccessorArgumentMappings[property] = arg;
   };
 

lib/SILGen/SILGenDecl.cpp

@@ -472,7 +472,7 @@ class DestroyLocalVariable : public Cleanup {
     Var->print(llvm::errs());
     llvm::errs() << "\n";
     if (isActive()) {
-      auto loc = SGF.VarLocs[Var];
+      auto &loc = SGF.VarLocs[Var];
       assert((loc.box || loc.value) && "One of box or value should be set");
       if (loc.box) {
         llvm::errs() << "Box: " << loc.box << "\n";
@@ -664,7 +664,8 @@ class LocalVariableInitialization : public SingleBufferInitialization {
     /// decl to.
     assert(SGF.VarLocs.count(decl) == 0 && "Already emitted the local?");
 
-    SGF.VarLocs[decl] = SILGenFunction::VarLoc::get(Addr, Box);
+    SGF.VarLocs[decl] = SILGenFunction::VarLoc(Addr,
+      SILAccessEnforcement::Dynamic, Box);
 
     SingleBufferInitialization::finishInitialization(SGF);
     assert(!DidFinish &&
@@ -677,6 +678,54 @@ class LocalVariableInitialization : public SingleBufferInitialization {
 } // end anonymous namespace
 
 namespace {
+
+static void deallocateAddressable(SILGenFunction &SGF,
+                SILLocation l,
+                const SILGenFunction::VarLoc::AddressableBuffer::State &state) {
+  SGF.B.createEndBorrow(l, state.storeBorrow);
+  SGF.B.createDeallocStack(l, state.allocStack);
+  if (state.reabstraction) {
+    SGF.B.createDestroyValue(l, state.reabstraction);
+  }
+}
+
+/// Cleanup to deallocate the addressable buffer for a parameter or let
+/// binding.
+class DeallocateLocalVariableAddressableBuffer : public Cleanup {
+  ValueDecl *vd;
+public:
+  DeallocateLocalVariableAddressableBuffer(ValueDecl *vd) : vd(vd) {}
+
+  void emit(SILGenFunction &SGF, CleanupLocation l,
+            ForUnwind_t forUnwind) override {
+    auto found = SGF.VarLocs.find(vd);
+    if (found == SGF.VarLocs.end()) {
+      return;
+    }
+    auto &loc = found->second;
+    
+    if (auto &state = loc.addressableBuffer.state) {
+      // The addressable buffer was forced, so clean it up now.
+      deallocateAddressable(SGF, l, *state);
+    } else {
+      // Remember this insert location in case we need to force the addressable
+      // buffer later.
+      SILInstruction *marker = SGF.B.createTuple(l, {});
+      loc.addressableBuffer.cleanupPoints.emplace_back(marker);
+    }
+  }
+
+  void dump(SILGenFunction &SGF) const override {
+#ifndef NDEBUG
+    llvm::errs() << "DeallocateLocalVariableAddressableBuffer\n"
+                 << "State:" << getState() << "\n"
+                 << "Decl: ";
+    vd->print(llvm::errs());
+    llvm::errs() << "\n";
+#endif
+  }
+};
+
 /// Initialize a writeback buffer that receives the value of a 'let'
 /// declaration.
 class LetValueInitialization : public Initialization {
@@ -755,7 +804,8 @@ class LetValueInitialization : public Initialization {
       if (isUninitialized)
         address = SGF.B.createMarkUninitializedVar(vd, address);
       DestroyCleanup = SGF.enterDormantTemporaryCleanup(address, *lowering);
-      SGF.VarLocs[vd] = SILGenFunction::VarLoc::get(address);
+      SGF.VarLocs[vd] = SILGenFunction::VarLoc(address,
+                                               SILAccessEnforcement::Unknown);
     }
     // Push a cleanup to destroy the let declaration.  This has to be
     // inactive until the variable is initialized: if control flow exits the
@@ -766,6 +816,10 @@ class LetValueInitialization : public Initialization {
     SGF.Cleanups.pushCleanupInState<DestroyLocalVariable>(
       CleanupState::Dormant, vd);
     DestroyCleanup = SGF.Cleanups.getTopCleanup();
+
+    // If the binding has an addressable buffer forced, it should be cleaned
+    // up here.
+    SGF.enterLocalVariableAddressableBufferScope(vd);
   }
 
   ~LetValueInitialization() override {
@@ -883,7 +937,8 @@ class LetValueInitialization : public Initialization {
     if (SGF.getASTContext().SILOpts.supportsLexicalLifetimes(SGF.getModule()))
       value = getValueForLexicalLifetimeBinding(SGF, loc, value, wasPlusOne);
 
-    SGF.VarLocs[vd] = SILGenFunction::VarLoc::get(value);
+    SGF.VarLocs[vd] = SILGenFunction::VarLoc(value,
+                                             SILAccessEnforcement::Unknown);
 
     // Emit a debug_value[_addr] instruction to record the start of this value's
     // lifetime, if permitted to do so.
@@ -1463,7 +1518,7 @@ SILGenFunction::emitInitializationForVarDecl(VarDecl *vd, bool forceImmutable,
     assert(SILDebugClient && "Debugger client doesn't support SIL");
     SILValue SV = SILDebugClient->emitLValueForVariable(vd, B);
 
-    VarLocs[vd] = SILGenFunction::VarLoc::get(SV);
+    VarLocs[vd] = VarLoc(SV, SILAccessEnforcement::Dynamic);
     return InitializationPtr(new KnownAddressInitialization(SV));
   }
 
@@ -1494,7 +1549,7 @@ SILGenFunction::emitInitializationForVarDecl(VarDecl *vd, bool forceImmutable,
     if (isUninitialized)
       addr = B.createMarkUninitializedVar(loc, addr);
 
-    VarLocs[vd] = SILGenFunction::VarLoc::get(addr);
+    VarLocs[vd] = VarLoc(addr, SILAccessEnforcement::Dynamic);
     Result = InitializationPtr(new KnownAddressInitialization(addr));
   } else {
     std::optional<MarkUninitializedInst::Kind> uninitKind;
@@ -2309,11 +2364,9 @@ void SILGenFunction::destroyLocalVariable(SILLocation silLoc, VarDecl *vd) {
     }
   };
 
-  auto loc = VarLocs[vd];
-
   // For a heap variable, the box is responsible for the value. We just need
   // to give up our retain count on it.
-  if (auto boxValue = loc.box) {
+  if (auto boxValue = VarLocs[vd].box) {
     if (!getASTContext().SILOpts.supportsLexicalLifetimes(getModule())) {
       emitDestroy(boxValue);
       return;
@@ -2329,7 +2382,7 @@ void SILGenFunction::destroyLocalVariable(SILLocation silLoc, VarDecl *vd) {
 
   // For 'let' bindings, we emit a release_value or destroy_addr, depending on
   // whether we have an address or not.
-  SILValue Val = loc.value;
+  SILValue Val = VarLocs[vd].value;
 
   if (Val->getType().isAddress()) {
     B.createDestroyAddr(silLoc, Val);
@@ -2406,6 +2459,108 @@ void SILGenFunction::destroyLocalVariable(SILLocation silLoc, VarDecl *vd) {
   llvm_unreachable("unhandled case");
 }
 
+void
+SILGenFunction::enterLocalVariableAddressableBufferScope(VarDecl *decl) {
+  Cleanups.pushCleanup<DeallocateLocalVariableAddressableBuffer>(decl);
+}
+
+SILValue
+SILGenFunction::getLocalVariableAddressableBuffer(VarDecl *decl,
+                                                  SILLocation curLoc,
+                                                  ValueOwnership ownership) {
+  auto foundVarLoc = VarLocs.find(decl);
+  if (foundVarLoc == VarLocs.end()) {
+    return SILValue();
+  }
+  
+  auto value = foundVarLoc->second.value;
+  auto access = foundVarLoc->second.access;
+  auto *state = foundVarLoc->second.addressableBuffer.state.get();
+  
+  SILType fullyAbstractedTy = getLoweredType(AbstractionPattern::getOpaque(),
+                                     decl->getTypeInContext()->getRValueType());
+  
+  // Check whether the bound value is inherently suitable for addressability.
+  // It must already be in memory and fully abstracted.
+  if (value->getType().isAddress()
+      && fullyAbstractedTy.getASTType() == value->getType().getASTType()) {
+    SILValue address = value;
+    // Begin an access if the address is mutable.
+    if (access != SILAccessEnforcement::Unknown) {
+      address = B.emitBeginAccess(curLoc, address,
+        ownership == ValueOwnership::InOut ? SILAccessKind::Modify
+                                           : SILAccessKind::Read,
+        access);
+    }
+    return address;
+  }
+  
+  // We can't retroactively introduce a reabstracted representation for a
+  // mutable binding (since we would now have two mutable memory locations
+  // representing the same value).
+  if (access != SILAccessEnforcement::Unknown) {
+    return SILValue();
+  }
+  
+  assert(ownership == ValueOwnership::Shared);
+  
+  // Check whether the in-memory representation has already been forced.
+  if (state) {
+    return state->storeBorrow;
+  }
+  
+  // Otherwise, force the addressable representation.
+  SILValue reabstraction, allocStack, storeBorrow;
+  {
+    SavedInsertionPointRAII save(B);
+    B.setInsertionPoint(value->getNextInstruction());
+    auto declarationLoc = value->getDefiningInsertionPoint()->getLoc();
+    
+    // Reabstract if necessary.
+    auto newValue = value;
+    reabstraction = SILValue();
+    if (newValue->getType().getASTType() != fullyAbstractedTy.getASTType()){
+      auto reabstracted = emitSubstToOrigValue(curLoc,
+        ManagedValue::forBorrowedRValue(value),
+        AbstractionPattern::getOpaque(),
+        decl->getTypeInContext()->getCanonicalType(),
+        SGFContext());
+      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);
+  }
+  
+  // Record the addressable representation.
+  auto &addressableBuffer = VarLocs[decl].addressableBuffer;
+  addressableBuffer.state
+    = std::make_unique<VarLoc::AddressableBuffer::State>(reabstraction,
+                                                         allocStack,
+                                                         storeBorrow);
+  auto *newState = addressableBuffer.state.get();
+
+  // Emit cleanups on any paths where we previously would have cleaned up
+  // the addressable representation if it had been forced earlier.
+  decltype(addressableBuffer.cleanupPoints) cleanupPoints;
+  cleanupPoints.swap(addressableBuffer.cleanupPoints);
+  
+  for (SILInstruction *cleanupPoint : cleanupPoints) {
+    SavedInsertionPointRAII insertCleanup(B, cleanupPoint);
+    deallocateAddressable(*this, cleanupPoint->getLoc(), *newState);
+    cleanupPoint->eraseFromParent();
+  }
+  
+  return storeBorrow;
+}
+
 void BlackHoleInitialization::performPackExpansionInitialization(
                                         SILGenFunction &SGF,
                                         SILLocation loc,
@@ -2437,3 +2592,9 @@ void BlackHoleInitialization::copyOrInitValueInto(SILGenFunction &SGF, SILLocati
   value = SGF.B.createMoveValue(loc, value);
   SGF.B.createIgnoredUse(loc, value.getValue());
 }
+
+SILGenFunction::VarLoc::AddressableBuffer::~AddressableBuffer() {
+  for (auto cleanupPoint : cleanupPoints) {
+    cleanupPoint->eraseFromParent();
+  }
+}

lib/SILGen/SILGenFunction.cpp

@@ -770,7 +770,7 @@ void SILGenFunction::emitCaptures(SILLocation loc,
       }
     };
 
-    auto Entry = found->second;
+    auto &Entry = found->second;
     auto val = Entry.value;
 
     switch (SGM.Types.getDeclCaptureKind(capture, expansion)) {