Merge pull request swiftlang#79758 from atrick/addressable-deps

Support enforcement of '@'_addressable under -enable-address-dependencies

Author: atrick

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceInsertion.swift

@@ -130,15 +130,16 @@ extension LifetimeDependentApply {
     var info = LifetimeSourceInfo()
     let hasScopedYield = applySite.parameterOperands.contains {
       if let dep = applySite.resultDependence(on: $0) {
-        return dep == .scope
+        return dep.isScoped
       }
       return false
     }
     if hasScopedYield {
-      // for consistency, we you yieldAddress if any yielded value is an address.
+      // for consistency, we use yieldAddress if any yielded value is an address.
       let targetKind = beginApply.yieldedValues.contains(where: { $0.type.isAddress })
         ? TargetKind.yieldAddress : TargetKind.yield
-      info.sources.push(LifetimeSource(targetKind: targetKind, convention: .scope, value: beginApply.token))
+      info.sources.push(LifetimeSource(targetKind: targetKind, convention: .scope(addressable: false),
+                                       value: beginApply.token))
     }
     for operand in applySite.parameterOperands {
       guard let dep = applySite.resultDependence(on: operand) else {
@@ -216,6 +217,11 @@ private extension LifetimeDependentApply.LifetimeSourceInfo {
       // A coroutine creates its own borrow scope, nested within its borrowed operand.
       bases.append(source.value)
     case .result, .inParameter, .inoutParameter:
+      // addressable dependencies directly depend on the incoming address.
+      if context.options.enableAddressDependencies() && source.convention.isAddressable {
+        bases.append(source.value)
+        return
+      }
       // Create a new dependence on the apply's access to the argument.
       for varIntoducer in gatherVariableIntroducers(for: source.value, context) {
         let scope = LifetimeDependence.Scope(base: varIntoducer, context)

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceScopeFixup.swift

@@ -249,7 +249,7 @@ private struct ScopeExtension {
 
 private extension LifetimeDependence.Scope {
   /// The instruction that introduces an extendable scope. This returns a non-nil scope introducer for
-  /// Extendable.nestedScopes.
+  /// ScopeExtension.nestedScopes.
   var extendableBegin: Instruction? {
     switch self {
     case let .access(beginAccess):
@@ -258,6 +258,17 @@ private extension LifetimeDependence.Scope {
       return beginBorrow.value.definingInstruction!
     case let .yield(yieldedValue):
       return yieldedValue.definingInstruction!
+    case let .initialized(initializer):
+      switch initializer {
+      case let .store(initializingStore: store, initialAddress: _):
+        if let sb = store as? StoreBorrowInst {
+          return sb
+        }
+        return nil
+      case .argument, .yield:
+        // TODO: extend indirectly yielded scopes.
+        return nil
+      }
     default:
       return nil
     }
@@ -277,29 +288,31 @@ private extension LifetimeDependence.Scope {
       let accessExtension = gatherAccessExtension(beginAccess: beginAccess, innerScopes: &innerScopes)
       return SingleInlineArray(element: accessExtension)
     case let .borrowed(beginBorrow):
-      let borrowedValue = beginBorrow.baseOperand!.value
-      let enclosingScope = LifetimeDependence.Scope(base: borrowedValue, context)
-      innerScopes.push(self)
-      var innerBorrowScopes = innerScopes
-      innerBorrowScopes.push(enclosingScope)
-      if let extensions = enclosingScope.gatherExtensions(innerScopes: innerBorrowScopes, context) {
-        return extensions
-      }
-      // This is the outermost scope to be extended because gatherExtensions did not find an enclosing scope.
-      return SingleInlineArray(element: getOuterExtension(owner: enclosingScope.parentValue, nestedScopes: innerScopes,
-                                                          context))
+      return gatherBorrowExtension(borrowedValue: beginBorrow.baseOperand!.value, innerScopes: &innerScopes, context)
+
     case let .yield(yieldedValue):
       innerScopes.push(self)
       var extensions = SingleInlineArray<ScopeExtension>()
       let applySite = yieldedValue.definingInstruction as! BeginApplyInst
       for operand in applySite.parameterOperands {
-        guard let dep = applySite.resultDependence(on: operand), dep == .scope else {
+        guard let dep = applySite.resultDependence(on: operand), dep.isScoped else {
           continue
         }
         // Pass a copy of innerScopes without modifying this one.
         extensions.append(contentsOf: gatherOperandExtension(on: operand, innerScopes: innerScopes, context))
       }
       return extensions
+    case let .initialized(initializer):
+      switch initializer {
+      case let .store(initializingStore: store, initialAddress: _):
+        if let sb = store as? StoreBorrowInst {
+          return gatherBorrowExtension(borrowedValue: sb.source, innerScopes: &innerScopes, context)
+        }
+        return nil
+      case .argument, .yield:
+        // TODO: extend indirectly yielded scopes.
+        return nil
+      }
     default:
       return nil
     }
@@ -360,6 +373,23 @@ private extension LifetimeDependence.Scope {
     }
     return ScopeExtension(owner: outerBeginAccess, nestedScopes: innerScopes, dependsOnArg: nil)
   }
+
+  func gatherBorrowExtension(borrowedValue: Value,
+                             innerScopes: inout SingleInlineArray<LifetimeDependence.Scope>,
+                             _ context: FunctionPassContext)
+    -> SingleInlineArray<ScopeExtension> {
+
+    let enclosingScope = LifetimeDependence.Scope(base: borrowedValue, context)
+    innerScopes.push(self)
+    var innerBorrowScopes = innerScopes
+    innerBorrowScopes.push(enclosingScope)
+    if let extensions = enclosingScope.gatherExtensions(innerScopes: innerBorrowScopes, context) {
+      return extensions
+    }
+    // This is the outermost scope to be extended because gatherExtensions did not find an enclosing scope.
+    return SingleInlineArray(element: getOuterExtension(owner: enclosingScope.parentValue, nestedScopes: innerScopes,
+                                                        context))
+  }
 }
 
 /// Compute the range of the a scope owner. Nested scopes must stay within this range.
@@ -584,6 +614,17 @@ private extension LifetimeDependence.Scope {
     case let .yield(yieldedValue):
       let beginApply = yieldedValue.definingInstruction as! BeginApplyInst
       return beginApply.createEnd(builder, context)
+    case let .initialized(initializer):
+      switch initializer {
+      case let .store(initializingStore: store, initialAddress: _):
+        if let sb = store as? StoreBorrowInst {
+          return builder.createEndBorrow(of: sb)
+        }
+        return nil
+      case .argument, .yield:
+        // TODO: extend indirectly yielded scopes.
+        return nil
+      }
     default:
       return nil
     }

SwiftCompilerSources/Sources/Optimizer/PassManager/Options.swift

@@ -32,6 +32,10 @@ struct Options {
     _bridged.enableSimplificationFor(inst.bridged)
   }
 
+  func enableAddressDependencies() -> Bool {
+    _bridged.enableAddressDependencies()
+  }
+
   var enableEmbeddedSwift: Bool {
     _bridged.hasFeature(.Embedded)
   }

SwiftCompilerSources/Sources/Optimizer/Utilities/LifetimeDependenceUtils.swift

@@ -105,10 +105,12 @@ struct LifetimeDependence : CustomStringConvertible {
       case let .initialized(initializer):
         let initialAddress = initializer.initialAddress
         precondition(initialAddress.type.isAddress, "expected an address")
-        precondition(initialAddress is AllocStackInst || initialAddress is FunctionArgument,
+        precondition(initialAddress is AllocStackInst || initialAddress is FunctionArgument
+                       || initialAddress is StoreBorrowInst,
                      "expected storage for a a local 'let'")
         if case let .store(store, _) = initializer {
-          precondition(store is StoringInstruction || store is SourceDestAddrInstruction || store is FullApplySite,
+          precondition(store is StoringInstruction || store is SourceDestAddrInstruction || store is FullApplySite
+                      || store is StoreBorrowInst,
                        "expected a store")
         }
       }
@@ -217,8 +219,8 @@ extension LifetimeDependence.Scope {
   /// Construct a lifetime dependence scope from the base value that other values depend on. This derives the kind of
   /// dependence scope and its parentValue from `base`.
   ///
-  /// The returned Scope must be the only scope for the given 'base' value.  This is generally non-recursive, except
-  /// that finds the single borrow introducer. Use-def walking is handled by a utility such as
+  /// The returned Scope must be the only scope for the given 'base' value. This is generally non-recursive, except
+  /// that it tries to find the single borrow introducer. General use-def walking is handled by a utility such as
   /// VariableIntroducerUseDefWalker, which can handle multiple introducers.
   ///
   /// `base` represents the OSSA lifetime that the dependent value must be used within. If `base` is owned, then it
@@ -291,7 +293,8 @@ extension LifetimeDependence.Scope {
     case let .yield(result):
       self.init(yield: result)
     case .storeBorrow(let sb):
-      self = Self(base: sb.source, context)
+      // Don't follow the stored value in case the dependence requires addressability.
+      self = .initialized(.store(initializingStore: sb, initialAddress: sb))
     }
   }
 
@@ -905,7 +908,7 @@ extension LifetimeDependenceDefUseWalker {
       return leafUse(of: operand)
     }
     if let dep = apply.resultDependence(on: operand),
-       dep == .inherit {
+       !dep.isScoped {
       // Operand is nonescapable and passed as a call argument. If the
       // result inherits its lifetime, then consider any nonescapable
       // result value to be a dependent use.

SwiftCompilerSources/Sources/SIL/FunctionConvention.swift

@@ -237,7 +237,25 @@ extension FunctionConvention {
 
 public enum LifetimeDependenceConvention : CustomStringConvertible {
   case inherit
-  case scope
+  case scope(addressable: Bool)
+
+  public var isScoped: Bool {
+    switch self {
+    case .inherit:
+      return false
+    case .scope:
+      return true
+    }
+  }
+
+  public var isAddressable: Bool {
+    switch self {
+    case .inherit:
+      return false
+    case let .scope(addressable):
+      return addressable
+    }
+  }
 
   public var description: String {
     switch self {
@@ -294,7 +312,8 @@ extension FunctionConvention {
         return .inherit
       }
       if scope {
-        return .scope
+        let addressable = bridged.checkAddressable(bridgedIndex(parameterIndex: index))
+        return .scope(addressable: addressable)
       }
       return nil
     }

include/swift/AST/LifetimeDependence.h

@@ -300,6 +300,11 @@ class LifetimeDependenceInfo {
       && scopeLifetimeParamIndices->contains(index);
   }
 
+  bool checkAddressable(int index) const {
+    return hasAddressableParamIndices()
+      && getAddressableIndices()->contains(index);
+  }
+
   std::string getString() const;
   void Profile(llvm::FoldingSetNodeID &ID) const;
   void getConcatenatedData(SmallVectorImpl<bool> &concatenatedData) const;

include/swift/AST/SILOptions.h

@@ -331,6 +331,10 @@ class SILOptions {
   /// optimizer.
   bool UseAggressiveReg2MemForCodeSize = true;
 
+  /// Enable enforcement of lifetime dependencies on addressable arguments.
+  /// Temporarily used to bootstrap the AddressableParameters feature.
+  bool EnableAddressDependencies = false;
+
   SILOptions() {}
 
   /// Return a hash code of any components from these options that should

include/swift/Option/FrontendOptions.td

@@ -1489,6 +1489,9 @@ def platform_availability_inheritance_map_path
   : Separate<["-"], "platform-availability-inheritance-map-path">, MetaVarName<"<path>">,
     HelpText<"Path of the platform inheritance platform map">;
 
+def enable_address_dependencies : Flag<["-"], "enable-address-dependencies">,
+  HelpText<"Enable enforcement of lifetime dependencies on addressable values.">;
+
 } // end let Flags = [FrontendOption, NoDriverOption, HelpHidden]
 
 def disable_experimental_parser_round_trip : Flag<["-"],

include/swift/SIL/SILBridging.h

@@ -156,6 +156,7 @@ struct BridgedYieldInfoArray {
 struct BridgedLifetimeDependenceInfo {
   swift::IndexSubset *_Nullable inheritLifetimeParamIndices;
   swift::IndexSubset *_Nullable scopeLifetimeParamIndices;
+  swift::IndexSubset *_Nullable addressableParamIndices;
   SwiftUInt targetIndex;
   bool immortal;
 
@@ -164,6 +165,7 @@ struct BridgedLifetimeDependenceInfo {
   BRIDGED_INLINE bool empty() const;
   BRIDGED_INLINE bool checkInherit(SwiftInt index) const;
   BRIDGED_INLINE bool checkScope(SwiftInt index) const;
+  BRIDGED_INLINE bool checkAddressable(SwiftInt index) const;
   BRIDGED_INLINE SwiftInt getTargetIndex() const;
 
   BRIDGED_INLINE BridgedOwnedString getDebugDescription() const;

include/swift/SIL/SILBridgingImpl.h

@@ -146,6 +146,7 @@ BridgedParameterInfo BridgedParameterInfoArray::at(SwiftInt parameterIndex) cons
 BridgedLifetimeDependenceInfo::BridgedLifetimeDependenceInfo(swift::LifetimeDependenceInfo info)
     : inheritLifetimeParamIndices(info.getInheritIndices()),
       scopeLifetimeParamIndices(info.getScopeIndices()),
+      addressableParamIndices(info.getAddressableIndices()),
       targetIndex(info.getTargetIndex()), immortal(info.isImmortal()) {}
 
 SwiftInt BridgedLifetimeDependenceInfoArray::count() const {
@@ -172,6 +173,10 @@ bool BridgedLifetimeDependenceInfo::checkScope(SwiftInt index) const {
          scopeLifetimeParamIndices->contains(index);
 }
 
+bool BridgedLifetimeDependenceInfo::checkAddressable(SwiftInt index) const {
+  return addressableParamIndices && addressableParamIndices->contains(index);
+}
+
 SwiftInt BridgedLifetimeDependenceInfo::getTargetIndex() const {
   return targetIndex;
 }

include/swift/SILOptimizer/OptimizerBridging.h

@@ -382,6 +382,9 @@ struct BridgedPassContext {
   bool enableSimplificationFor(BridgedInstruction inst) const;
   BRIDGED_INLINE bool enableWMORequiredDiagnostics() const;
 
+  // Temporary for AddressableParameters Bootstrapping.
+  BRIDGED_INLINE bool enableAddressDependencies() const;
+
   // Closure specializer
   SWIFT_IMPORT_UNSAFE BridgedFunction ClosureSpecializer_createEmptyFunctionWithSpecializedSignature(BridgedStringRef specializedName,
                                                         const BridgedParameterInfo * _Nullable specializedBridgedParams,

include/swift/SILOptimizer/OptimizerBridgingImpl.h

@@ -573,6 +573,11 @@ bool BridgedPassContext::enableWMORequiredDiagnostics() const {
   return mod->getOptions().EnableWMORequiredDiagnostics;
 }
 
+bool BridgedPassContext::enableAddressDependencies() const {
+  swift::SILModule *mod = invocation->getPassManager()->getModule();
+  return mod->getOptions().EnableAddressDependencies;
+}
+
 static_assert((int)BridgedPassContext::SILStage::Raw == (int)swift::SILStage::Raw);
 static_assert((int)BridgedPassContext::SILStage::Canonical == (int)swift::SILStage::Canonical);
 static_assert((int)BridgedPassContext::SILStage::Lowered == (int)swift::SILStage::Lowered);

lib/DriverTool/sil_opt_main.cpp

@@ -592,6 +592,10 @@ struct SILOptOptions {
       "swift-version",
       llvm::cl::desc(
           "The swift version to assume AST declarations correspond to"));
+
+  llvm::cl::opt<bool> EnableAddressDependencies = llvm::cl::opt<bool>(
+      "enable-address-dependencies",
+      llvm::cl::desc("Enable enforcement of lifetime dependencies on addressable values."));
 };
 
 /// Regular expression corresponding to the value given in one of the
@@ -909,6 +913,8 @@ int sil_opt_main(ArrayRef<const char *> argv, void *MainAddr) {
   SILOpts.EnablePackMetadataStackPromotion =
       options.EnablePackMetadataStackPromotion;
 
+  SILOpts.EnableAddressDependencies = options.EnableAddressDependencies;
+
   if (options.OptModeFlag == OptimizationMode::NotSet) {
     if (options.OptimizationGroup == OptGroup::Diagnostics)
       SILOpts.OptMode = OptimizationMode::NoOptimization;

lib/Frontend/CompilerInvocation.cpp

@@ -3013,6 +3013,9 @@ static bool ParseSILArgs(SILOptions &Opts, ArgList &Args,
     Opts.ShouldFunctionsBePreservedToDebugger &=
         LTOKind.value() == IRGenLLVMLTOKind::None;
 
+  
+  Opts.EnableAddressDependencies = Args.hasArg(OPT_enable_address_dependencies);
+
   return false;
 }