Merge pull request #83505 from eeckstein/non-loadable-global-inits

Optimizer: support statically initialized `Atomic` and `Mutex` globals

Author: eeckstein

SwiftCompilerSources/Sources/AST/Declarations.swift

@@ -116,7 +116,9 @@ final public class AccessorDecl: FuncDecl {}
 
 final public class MacroDecl: ValueDecl {}
 
-final public class EnumElementDecl: ValueDecl {}
+final public class EnumElementDecl: ValueDecl {
+  public var hasAssociatedValues: Bool { bridged.EnumElementDecl_hasAssociatedValues() }
+}
 
 final public class ExtensionDecl: Decl {}
 

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/InitializeStaticGlobals.swift

@@ -80,8 +80,11 @@ private indirect enum GlobalInitValue {
   // For example, a struct or vector which is initialized by storing its elements.
   case aggregate([GlobalInitValue])
 
-  // An enum with a payload which is not a SIL "constant".
-  case enumCase(caseIndex: Int, payload: GlobalInitValue)
+  // An enum case without a payload of an address-only enum.
+  case enumCase(caseIndex: Int)
+
+  // An enum case which is not a SIL "constant", e.g. because it's address-only
+  case enumCaseWithPayload(caseIndex: Int, payload: GlobalInitValue)
 
   init?(of globalInitFunction: Function, _ context: FunctionPassContext) {
     self = .undefined
@@ -133,20 +136,44 @@ private indirect enum GlobalInitValue {
     self = builder.initValue
   }
 
+  enum InitValue {
+    // The common case
+    case value(Value)
+
+    // For payload-less cases of address-only enums. Such cases are initialized purely with an `inject_enum_addr`,
+    // and we don't have a `Value` which represents the resulting enum(-case).
+    case enumCaseWithoutPayload(InjectEnumAddrInst)
+
+    var parentFunction: Function {
+      switch self {
+        case .value(let value):  return value.parentFunction
+        case .enumCaseWithoutPayload(let iea): return iea.parentFunction
+      }
+    }
+  }
+
   // Sets an element in the constant tree.
   // Returns true if this was successful. One reason for being not successful is if a certain
   // element is set twice, i.e. does not have a single defined value.
-  mutating func setElement(to value: Value, at path: SmallProjectionPath, type: Type) -> Bool {
+  mutating func setElement(to value: InitValue, at path: SmallProjectionPath, type: Type) -> Bool {
     let (kind, index, subPath) = path.pop()
     switch kind {
     case .root:
       guard case .undefined = self else {
         // The element was set twice.
         return false
       }
-      self = .constant(value)
+      switch value {
+      case .value(let value):
+        self = .constant(value)
+      case .enumCaseWithoutPayload:
+        fatalError("should have been handled in the .enumCase of the SmallProjectionPath below")
+      }
       return true
 
+    case .enumCase:
+      return setEnumCase(to: value, at: subPath, index: index, type: type)
+
     case .structField:
       guard let structFields = type.getNominalFields(in: value.parentFunction) else {
         return false
@@ -186,7 +213,7 @@ private indirect enum GlobalInitValue {
   }
 
   private mutating func setField(
-    to value: Value, at path: SmallProjectionPath,
+    to value: InitValue, at path: SmallProjectionPath,
     index: Int, type: Type, numFields: Int
   ) -> Bool {
     if case .undefined = self {
@@ -205,6 +232,43 @@ private indirect enum GlobalInitValue {
     return false
   }
 
+  private mutating func setEnumCase(to value: InitValue, at path: SmallProjectionPath, index: Int, type: Type) -> Bool {
+    switch value {
+
+    case .enumCaseWithoutPayload(let iea):
+      guard case .undefined = self else {
+        // The enum was set twice.
+        return false
+      }
+      assert(index == iea.caseIndex)
+      self = .enumCase(caseIndex: index)
+
+    case .value:
+      guard let payloadType = type.getEnumCases(in: value.parentFunction)!.getPayloadType(ofCaseIndex: index) else {
+        return false
+      }
+      switch self {
+      case .undefined:
+        // It's the first time we set the payload or a sub-field of it.
+        var payload = GlobalInitValue.undefined
+        if !payload.setElement(to: value, at: path, type: payloadType) {
+          return false
+        }
+        self = .enumCaseWithPayload(caseIndex: index, payload: payload)
+      case .enumCaseWithPayload(let existingIndex, var payload) where index == existingIndex:
+        // Some sub-field of the enum-payload was already set.
+        self = .undefined // avoid copy-on-write
+        if !payload.setElement(to: value, at: path, type: payloadType) {
+          return false
+        }
+        self = .enumCaseWithPayload(caseIndex: index, payload: payload)
+      default:
+        return false
+      }
+    }
+    return true
+  }
+
   /// Creates SIL for this global init value in the initializer of the `global`.
   func materialize(into global: GlobalVariable, from function: Function, _ context: FunctionPassContext) {
     var cloner = StaticInitCloner(cloneTo: global, context)
@@ -248,8 +312,11 @@ private indirect enum GlobalInitValue {
       }
       return builder.createVector(type: type, arguments: elementValues)
 
-    case .enumCase(let caseIndex, let payload):
-      let payloadType = type.getEnumCases(in: function)!.first(where: { $0.index == caseIndex })!.payload!
+    case .enumCase(let caseIndex):
+      return builder.createEnum(caseIndex: caseIndex, payload: nil, enumType: type)
+
+    case .enumCaseWithPayload(let caseIndex, let payload):
+      let payloadType = type.getEnumCases(in: function)!.getPayloadType(ofCaseIndex: caseIndex)!
       let payloadValue = payload.materializeRecursively(type: payloadType, &cloner, builder, function)
       return builder.createEnum(caseIndex: caseIndex, payload: payloadValue, enumType: type)
     }
@@ -272,7 +339,7 @@ private indirect enum GlobalInitValue {
     _ context: FunctionPassContext
   ) {
     switch self {
-    case .undefined:
+    case .undefined, .enumCase:
       break
     case .constant(let value):
       if value.containsLoad(context) {
@@ -281,7 +348,7 @@ private indirect enum GlobalInitValue {
           self = .aggregate((value as! Instruction).operands.lazy.map { .constant($0.value) })
           resolveLoadsRecursively(from: &stackValues, &resolvedAllocStacks, context)
         case let ei as EnumInst:
-          self = .enumCase(caseIndex: ei.caseIndex, payload: .constant(ei.payload!))
+          self = .enumCaseWithPayload(caseIndex: ei.caseIndex, payload: .constant(ei.payload!))
           resolveLoadsRecursively(from: &stackValues, &resolvedAllocStacks, context)
         case let li as LoadInst:
           guard let allocStack = li.address as? AllocStackInst,
@@ -306,10 +373,9 @@ private indirect enum GlobalInitValue {
         newFields[i].resolveLoadsRecursively(from: &stackValues, &resolvedAllocStacks, context)
       }
       self = .aggregate(newFields)
-    case .enumCase(let caseIndex, let payload):
-      var newPayload = payload
-      newPayload.resolveLoadsRecursively(from: &stackValues, &resolvedAllocStacks, context)
-      self = .enumCase(caseIndex: caseIndex, payload: newPayload)
+    case .enumCaseWithPayload(let caseIndex, var payload):
+      payload.resolveLoadsRecursively(from: &stackValues, &resolvedAllocStacks, context)
+      self = .enumCaseWithPayload(caseIndex: caseIndex, payload: payload)
     }
   }
 
@@ -321,7 +387,9 @@ private indirect enum GlobalInitValue {
       return value.isValidGlobalInitValue(context)
     case .aggregate(let fields):
       return fields.allSatisfy { $0.isValid(context) }
-    case .enumCase(_, let payload):
+    case .enumCase:
+      return true
+    case .enumCaseWithPayload(_, let payload):
       return payload.isValid(context)
     }
   }
@@ -358,10 +426,10 @@ private struct InitValueBuilder: AddressDefUseWalker {
   mutating func leafUse(address: Operand, path: UnusedWalkingPath) -> WalkResult {
     switch address.instruction {
     case let store as StoreInst:
-      let accessPath = store.destination.constantAccessPath
+      let accessPath = store.destination.lookThroughRawLayoutAddress.constantAccessPath
       switch accessPath.base {
       case .global, .stack:
-        if !initValue.setElement(to: store.source, at: accessPath.projectionPath, type: originalAddress.type) {
+        if !initValue.setElement(to: .value(store.source), at: accessPath.projectionPath, type: originalAddress.type) {
           return .abortWalk
         }
         return .continueWalk
@@ -376,17 +444,49 @@ private struct InitValueBuilder: AddressDefUseWalker {
           return .abortWalk
         }
         // The `nonConstAccessPath` now contains a single `.anyIndexedElement`.
-        if !initValue.setElement(to: store.source, at: nonConstAccessPath.projectionPath, type: originalAddress.type) {
+        if !initValue.setElement(to: .value(store.source), at: nonConstAccessPath.projectionPath, type: originalAddress.type) {
           return .abortWalk
         }
         return .continueWalk
       default:
         fatalError("could not compute access path")
       }
+    case let injectEnum as InjectEnumAddrInst:
+      if injectEnum.element.hasAssociatedValues {
+        if !injectEnum.operand.value.type.isLoadable(in: injectEnum.parentFunction) {
+          // TODO: we don't support non-loadable enum cases with payload yet, because IRGen support is missing.
+          //       e.g. `var global: Atomic<Int>? = Atomic<Int>(0)`
+          // FixedTypeInfo (= used for non-loadable types) is missing the ability to pack a payload into an enum.
+          return .abortWalk
+        }
+        return .continueWalk
+      }
+      let accessPath = injectEnum.enum.getAccessPath(fromInitialPath: SmallProjectionPath(.enumCase,
+                                                                                          index: injectEnum.caseIndex))
+      switch accessPath.base {
+      case .global, .stack:
+        if !initValue.setElement(to: .enumCaseWithoutPayload(injectEnum), at: accessPath.projectionPath, type: originalAddress.type) {
+          return .abortWalk
+        }
+        return .continueWalk
+      default:
+        return .abortWalk
+      }
+
     case is LoadInst, is DeallocStackInst:
       return .continueWalk
-    case let bi as BuiltinInst where bi.id == .PrepareInitialization:
-      return .continueWalk
+    case let bi as BuiltinInst:
+      switch bi.id {
+      case .PrepareInitialization:
+        return .continueWalk
+      case .AddressOfRawLayout:
+        if let addr2Ptr = bi.uses.getSingleUser(ofType: PointerToAddressInst.self) {
+          return walkDownUses(ofAddress: addr2Ptr, path: path)
+        }
+        return .abortWalk
+      default:
+        return .abortWalk
+      }
     default:
       return .abortWalk
     }
@@ -466,14 +566,9 @@ private extension Function {
       case let alloc as AllocGlobalInst where alloc.global == global:
         return false
       case let store as StoreInst:
-        switch store.destination.accessBase {
-        case .global(let g) where g == global:
-          return false
-        case .stack:
-          return false
-        default:
-          return true
-        }
+        return !store.destination.lookThroughRawLayoutAddress.isAddressOfStack(orGlobal: global)
+      case let injectEnum as InjectEnumAddrInst:
+        return !injectEnum.enum.isAddressOfStack(orGlobal: global)
       case let bi as BuiltinInst where bi.id == .PrepareInitialization:
         return false
       default:
@@ -496,3 +591,43 @@ private extension Function {
     context.removeTriviallyDeadInstructionsIgnoringDebugUses(in: self)
   }
 }
+
+private extension Value {
+  func isAddressOfStack(orGlobal global: GlobalVariable) -> Bool {
+    switch accessBase {
+    case .global(let g) where g == global:
+      return true
+    case .stack:
+      return true
+    default:
+      return false
+    }
+  }
+
+  /// Looks through a `@_rawLayout` projection, which "type casts" a raw-layout struct to it's content, which
+  /// must match the like-type of the raw-layout, e.g.
+  /// ```
+  /// @_rawLayout(like: T)
+  /// struct S {}
+  ///
+  ///   %2 = builtin "addressOfRawLayout"<S>(%1 : $*S) : $Builtin.RawPointer
+  ///   %3 = pointer_to_address %2 to $*T
+  /// ```
+  var lookThroughRawLayoutAddress: Value {
+    if let ptr2Addr = self as? PointerToAddressInst,
+       let builtin = ptr2Addr.pointer as? BuiltinInst, builtin.id == .AddressOfRawLayout,
+       let likeType = builtin.arguments[0].type.rawLayoutSubstitutedLikeType,
+       builtin.arguments[0].type.rawLayoutSubstitutedCountType == nil,
+       likeType.canonical == type.canonicalType
+    {
+      return builtin.arguments[0]
+    }
+    return self
+  }
+}
+
+private extension EnumCases {
+  func getPayloadType(ofCaseIndex caseIndex: Int) -> Type? {
+    return first(where: { $0.index == caseIndex })!.payload
+  }
+}

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -391,8 +391,7 @@ extension Instruction {
     case let bi as BuiltinInst:
       switch bi.id {
       case .ZeroInitializer:
-        let type = bi.type.isBuiltinVector ? bi.type.builtinVectorElementType(in: parentFunction) : bi.type
-        return type.isBuiltinInteger || type.isBuiltinFloat
+        return bi.arguments.count == 0
       case .PtrToInt:
         return bi.operands[0].value is StringLiteralInst
       case .IntToPtr:

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -673,6 +673,7 @@ final public class ExtendLifetimeInst : Instruction, UnaryInstruction {}
 final public class InjectEnumAddrInst : Instruction, UnaryInstruction, EnumInstruction {
   public var `enum`: Value { operand.value }
   public var caseIndex: Int { bridged.InjectEnumAddrInst_caseIndex() }
+  public var element: EnumElementDecl { bridged.InjectEnumAddrInst_element().getAs(EnumElementDecl.self) }
 }
 
 //===----------------------------------------------------------------------===//

SwiftCompilerSources/Sources/SIL/Type.swift

@@ -103,6 +103,16 @@ public struct Type : TypeProperties, CustomStringConvertible, NoReflectionChildr
     bridged.isAddressableForDeps(function.bridged)
   }
 
+  /// If this is a raw layout type, returns the substituted like-type.
+  public var rawLayoutSubstitutedLikeType: AST.`Type`? {
+    .init(bridgedOrNil: bridged.getRawLayoutSubstitutedLikeType())
+  }
+
+  /// If this is a raw layout type, returns the substituted count-type.
+  public var rawLayoutSubstitutedCountType: AST.`Type`? {
+    .init(bridgedOrNil: bridged.getRawLayoutSubstitutedCountType())
+  }
+
   //===--------------------------------------------------------------------===//
   //                Properties of lowered `SILFunctionType`s
   //===--------------------------------------------------------------------===//

include/swift/AST/ASTBridging.h

@@ -377,6 +377,7 @@ struct BridgedDeclObj {
   BRIDGED_INLINE bool ProtocolDecl_requiresClass() const;
   BRIDGED_INLINE bool AbstractFunction_isOverridden() const;
   BRIDGED_INLINE bool Destructor_isIsolated() const;
+  BRIDGED_INLINE bool EnumElementDecl_hasAssociatedValues() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedStringRef AccessorDecl_getKindName() const;
 };
 

include/swift/AST/ASTBridgingImpl.h

@@ -219,6 +219,10 @@ bool BridgedDeclObj::Destructor_isIsolated() const {
   return ai.isActorIsolated();
 }
 
+bool BridgedDeclObj::EnumElementDecl_hasAssociatedValues() const {
+  return getAs<swift::EnumElementDecl>()->hasAssociatedValues();
+}
+
 //===----------------------------------------------------------------------===//
 // MARK: BridgedASTNode
 //===----------------------------------------------------------------------===//

include/swift/SIL/SILBridging.h

@@ -263,6 +263,8 @@ struct BridgedType {
   BRIDGED_INLINE bool isExactSuperclassOf(BridgedType t) const;
   BRIDGED_INLINE bool isMarkedAsImmortal() const;
   BRIDGED_INLINE bool isAddressableForDeps(BridgedFunction f) const;
+  BRIDGED_INLINE SWIFT_IMPORT_UNSAFE BridgedASTType getRawLayoutSubstitutedLikeType() const;
+  BRIDGED_INLINE SWIFT_IMPORT_UNSAFE BridgedASTType getRawLayoutSubstitutedCountType() const;
   BRIDGED_INLINE SwiftInt getCaseIdxOfEnumType(BridgedStringRef name) const;
   static BRIDGED_INLINE SwiftInt getNumBoxFields(BridgedCanType boxTy);
   static SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedType getBoxFieldType(BridgedCanType boxTy,
@@ -770,6 +772,7 @@ struct BridgedInstruction {
   BRIDGED_INLINE SwiftInt InitEnumDataAddrInst_caseIndex() const;
   BRIDGED_INLINE SwiftInt UncheckedTakeEnumDataAddrInst_caseIndex() const;
   BRIDGED_INLINE SwiftInt InjectEnumAddrInst_caseIndex() const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedDeclObj InjectEnumAddrInst_element() const;
   BRIDGED_INLINE SwiftInt RefElementAddrInst_fieldIndex() const;
   BRIDGED_INLINE bool RefElementAddrInst_fieldIsLet() const;
   BRIDGED_INLINE bool RefElementAddrInst_isImmutable() const;