InitializeStaticGlobals : reapply support of non-loadable enums with a fix

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,41 @@ private indirect enum GlobalInitValue {
     return false
   }
 
+  private mutating func setEnumCase(to value: InitValue, at path: SmallProjectionPath, index: Int, type: Type) -> Bool {
+    if path.isEmpty, case .enumCaseWithoutPayload(let iea) = value {
+
+      guard case .undefined = self else {
+        // The enum was set twice.
+        return false
+      }
+      assert(index == iea.caseIndex)
+      self = .enumCase(caseIndex: index)
+    } else {
+      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 = Cloner(cloneToGlobal: global, context)
@@ -248,8 +310,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 +337,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 +346,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 +371,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 +385,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)
     }
   }
@@ -361,7 +427,7 @@ private struct InitValueBuilder: AddressDefUseWalker {
       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,13 +442,35 @@ 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:
@@ -477,6 +565,8 @@ private extension Function {
         return false
       case let store as StoreInst:
         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:
@@ -533,3 +623,9 @@ private extension Value {
     return self
   }
 }
+
+private extension EnumCases {
+  func getPayloadType(ofCaseIndex caseIndex: Int) -> Type? {
+    return first(where: { $0.index == caseIndex })!.payload
+  }
+}

test/SILOptimizer/init_static_globals.sil

@@ -37,6 +37,15 @@ struct TwoFields {
   let b: Int32
 }
 
+struct WithOptionalEnum {
+  var e: E?
+}
+
+enum E {
+  case a
+  case b
+}
+
 @_rawLayout(like: Int32)
 struct Raw: ~Copyable {}
 
@@ -168,6 +177,28 @@ sil_global [let] @graw2: $Raw
 // CHECK-LABEL: sil_global [let] @grawArray : $RawArray{{$}}
 sil_global [let] @grawArray: $RawArray
 
+// CHECK-LABEL: sil_global [let] @gIndirectNone : $Optional<TwoFields> = {
+// CHECK-NEXT:    %initval = enum $Optional<TwoFields>, #Optional.none!enumelt
+// CHECK-NEXT:  }
+sil_global [let] @gIndirectNone: $Optional<TwoFields>
+
+// CHECK-LABEL: sil_global [let] @gIndirectSome : $Optional<TwoFields> = {
+// CHECK-NEXT:    %0 = integer_literal $Builtin.Int32, 11
+// CHECK-NEXT:    %1 = struct $Int32 (%0)
+// CHECK-NEXT:    %2 = integer_literal $Builtin.Int32, 10
+// CHECK-NEXT:    %3 = struct $Int32 (%2)
+// CHECK-NEXT:    %4 = struct $TwoFields (%1, %3)
+// CHECK-NEXT:    %initval = enum $Optional<TwoFields>, #Optional.some!enumelt, %4
+// CHECK-NEXT:  }
+sil_global [let] @gIndirectSome: $Optional<TwoFields>
+
+// CHECK-LABEL: sil_global [let] @structWithOptionalEnum : $WithOptionalEnum = {
+// CHECK-NEXT:    %0 = enum $E, #E.a!enumelt
+// CHECK-NEXT:    %1 = enum $Optional<E>, #Optional.some!enumelt, %0
+// CHECK-NEXT:    %initval = struct $WithOptionalEnum (%1)
+// CHECK-NEXT:  }
+sil_global [let] @structWithOptionalEnum : $WithOptionalEnum
+
 sil @unknownfunc : $@convention(thin) () -> ()
 
 // CHECK-LABEL: sil [global_init_once_fn] [ossa] @globalinit_trivialglobal_func :
@@ -579,3 +610,41 @@ bb0(%0 : $Builtin.RawPointer):
   return %21
 }
 
+sil [global_init_once_fn] [ossa] @globalinit_indirect_none: $@convention(c) (Builtin.RawPointer) -> () {
+bb0(%0 : $Builtin.RawPointer):
+  alloc_global @gIndirectNone
+  %2 = global_addr @gIndirectNone : $*Optional<TwoFields>
+  inject_enum_addr %2, #Optional.none!enumelt
+  %21 = tuple ()
+  return %21
+}
+
+sil [global_init_once_fn] [ossa] @globalinit_indirect_some : $@convention(c) () -> () {
+bb0:
+  alloc_global @gIndirectSome
+  %1 = global_addr @gIndirectSome : $*Optional<TwoFields>
+  %2 = init_enum_data_addr %1, #Optional.some!enumelt
+  %3 = integer_literal $Builtin.Int32, 10
+  %4 = struct $Int32 (%3)
+  %5 = struct_element_addr %2, #TwoFields.b
+  store %4 to [trivial] %5
+  %7 = integer_literal $Builtin.Int32, 11
+  %8 = struct $Int32 (%7)
+  %9 = struct_element_addr %2, #TwoFields.a
+  store %8 to [trivial] %9
+  %10 = tuple ()
+  return %10 : $()
+}
+
+sil [global_init_once_fn] @globalinit_structWithOptionalEnum : $@convention(c) (Builtin.RawPointer) -> () {
+bb0(%0 : $Builtin.RawPointer):
+  alloc_global @structWithOptionalEnum
+  %2 = global_addr @structWithOptionalEnum : $*WithOptionalEnum
+  %3 = struct_element_addr %2, #WithOptionalEnum.e
+  %4 = init_enum_data_addr %3, #Optional.some!enumelt
+  inject_enum_addr %4, #E.a!enumelt
+  inject_enum_addr %3, #Optional.some!enumelt
+  %7 = tuple ()
+  return %7
+}
+

test/SILOptimizer/static_atomics.swift

@@ -0,0 +1,65 @@
+// RUN: %target-build-swift -parse-as-library -Xfrontend -disable-availability-checking -O %s -module-name=test -emit-sil | %FileCheck %s
+
+// RUN: %empty-directory(%t) 
+// RUN: %target-build-swift -parse-as-library -Xfrontend -disable-availability-checking -O -module-name=test %s -o %t/a.out
+// RUN: %target-run %t/a.out | %FileCheck %s -check-prefix=CHECK-OUTPUT
+
+// REQUIRES: executable_test,swift_stdlib_no_asserts,optimized_stdlib
+// REQUIRES: synchronization
+// UNSUPPORTED: back_deployment_runtime
+
+import Synchronization
+
+
+struct TwoAtomics: ~Copyable {
+  let a: Atomic<UInt8>
+  let b: Atomic<UInt8>
+}
+
+// CHECK-LABEL: sil_global hidden @$s4test7atomic1AA10TwoAtomicsVSgvp : $Optional<TwoAtomics> = {
+var atomic1: TwoAtomics? = nil
+
+// CHECK-LABEL: sil_global hidden @$s4test7atomic2AA10TwoAtomicsVvp : $TwoAtomics = {
+let atomic2: TwoAtomics = TwoAtomics(a: Atomic(29), b: Atomic(30))
+
+// TODO: this is not initialized statically, because missing IRGen support
+var atomic3: TwoAtomics? = TwoAtomics(a: Atomic(27), b: Atomic(28))
+
+// CHECK-LABEL: sil_global hidden @$s4test5mutex15Synchronization5MutexVySiGvp : $Mutex<Int> = {
+let mutex = Mutex<Int>(123)
+
+@main
+struct Main {
+  static func main() {
+
+    precondition(atomic1 == nil)
+
+    // CHECK-OUTPUT: atomic2: 29 30
+    print("atomic2:", atomic2.a.load(ordering: .relaxed), atomic2.b.load(ordering: .relaxed))
+
+    // CHECK-OUTPUT: atomic3: 27 28
+    switch atomic3 {
+    case .some(let a):
+      print("atomic3:", a.a.load(ordering: .relaxed), a.b.load(ordering: .relaxed))
+    case .none:
+      break
+    }
+
+    atomic1 = TwoAtomics(a: Atomic(1), b: Atomic(2))
+
+    // CHECK-OUTPUT: atomic2: 29 30
+    print("atomic2:", atomic2.a.load(ordering: .relaxed), atomic2.b.load(ordering: .relaxed))
+
+
+    mutex.withLock {
+      $0 = $0 + 1
+    }
+
+    // CHECK-OUTPUT: mutex: 124
+    mutex.withLock {
+      print("mutex:", $0)
+    }
+  }
+}
+
+