[opt] Re-work broadenSingleElementStores to use projections. (#32318)

* Fixes loadable edge case for address-only types.
* Re-work, generalize, and simplify by using projections.
* Support `-enable-cxx-interop` in sil-opt.

Author: zoecarver

include/swift/SIL/Projection.h

@@ -462,7 +462,7 @@ class Projection {
   static NullablePtr<SingleValueInstruction>
   createAggFromFirstLevelProjections(SILBuilder &B, SILLocation Loc,
                                      SILType BaseType,
-                                     llvm::SmallVectorImpl<SILValue> &Values);
+                                     ArrayRef<SILValue> Values);
 
   void print(raw_ostream &os, SILType baseType) const;
 private:

lib/SIL/Utils/Projection.cpp

@@ -803,10 +803,9 @@ Projection::operator<(const Projection &Other) const {
 }
 
 NullablePtr<SingleValueInstruction>
-Projection::
-createAggFromFirstLevelProjections(SILBuilder &B, SILLocation Loc,
-                                   SILType BaseType,
-                                   llvm::SmallVectorImpl<SILValue> &Values) {
+Projection::createAggFromFirstLevelProjections(
+    SILBuilder &B, SILLocation Loc, SILType BaseType,
+    ArrayRef<SILValue> Values) {
   if (BaseType.getStructOrBoundGenericStruct()) {
     return B.createStruct(Loc, BaseType, Values);
   }

lib/SILOptimizer/Utils/CanonicalizeInstruction.cpp

@@ -315,68 +315,65 @@ splitAggregateLoad(LoadInst *loadInst, CanonicalizeInstruction &pass) {
 // (store (struct_element_addr %base) object)
 //   ->
 // (store %base (struct object))
+//
+// TODO: supporting enums here would be very easy. The main thing is adding
+// support in `createAggFromFirstLevelProjections`.
+// Note: we will not be able to support tuples because we cannot have a
+// single-element tuple.
 static SILBasicBlock::iterator
 broadenSingleElementStores(StoreInst *storeInst,
                            CanonicalizeInstruction &pass) {
   // Keep track of the next iterator after any newly added or to-be-deleted
   // instructions. This must be valid regardless of whether the pass immediately
   // deletes the instructions or simply records them for later deletion.
   auto nextII = std::next(storeInst->getIterator());
-
-  // Bail if the store's destination is not a struct_element_addr.
-  auto *sea = dyn_cast<StructElementAddrInst>(storeInst->getDest());
-  if (!sea)
-    return nextII;
-
   auto *f = storeInst->getFunction();
 
-  // Continue up the struct_element_addr chain, as long as each struct only has
-  // a single property, creating StoreInsts along the way.
-  SILBuilderWithScope builder(storeInst);
-
-  SILValue result = storeInst->getSrc();
-  SILValue baseAddr = sea->getOperand();
-  SILValue storeAddr;
-  while (true) {
+  ProjectionPath projections(storeInst->getDest()->getType());
+  SILValue op = storeInst->getDest();
+  while (isa<StructElementAddrInst>(op)) {
+    auto *inst = cast<SingleValueInstruction>(op);
+    SILValue baseAddr = inst->getOperand(0);
     SILType baseAddrType = baseAddr->getType();
-
-    // If our aggregate has unreferenced storage then we can never prove if it
-    // actually has a single field.
-    if (baseAddrType.aggregateHasUnreferenceableStorage())
-      break;
-
     auto *decl = baseAddrType.getStructOrBoundGenericStruct();
     assert(
       !decl->isResilient(f->getModule().getSwiftModule(),
                          f->getResilienceExpansion()) &&
         "This code assumes resilient structs can not have fragile fields. If "
         "this assert is hit, this has been changed. Please update this code.");
 
-    // NOTE: If this is ever changed to support enums, we must check for address
-    // only types here. For structs we do not have to check since a single
-    // element struct with a loadable element can never be address only. We
-    // additionally do not have to worry about our input value being address
-    // only since we are storing into it.
-    if (decl->getStoredProperties().size() != 1)
+    // Bail if the store's destination is not a struct_element_addr or if the
+    // store's destination (%base) is not a loadable type. If %base is not a
+    // loadable type, we can't create it as a struct later on.
+    // If our aggregate has unreferenced storage then we can never prove if it
+    // actually has a single field.
+    if (!baseAddrType.isLoadable(*f) ||
+        baseAddrType.aggregateHasUnreferenceableStorage() ||
+        decl->getStoredProperties().size() != 1)
       break;
 
-    // Update the store location now that we know it is safe.
-    storeAddr = baseAddr;
+    projections.push_back(Projection(inst));
+    op = baseAddr;
+  }
 
-    // Otherwise, create the struct.
-    result = builder.createStruct(storeInst->getLoc(),
-                                  baseAddrType.getObjectType(), result);
+  // If we couldn't create a projection, bail.
+  if (projections.empty())
+    return nextII;
 
-    // See if we have another struct_element_addr we can strip off. If we don't
-    // then this as much as we can promote.
-    sea = dyn_cast<StructElementAddrInst>(sea->getOperand());
-    if (!sea)
-      break;
-    baseAddr = sea->getOperand();
+  // Now work our way back up. At this point we know all operations we are going
+  // to do succeed (cast<SingleValueInst>, createAggFromFirstLevelProjections,
+  // etc.) so we can omit null checks. We should not bail at this point (we
+  // could create a double consume, or worse).
+  SILBuilderWithScope builder(storeInst);
+  SILValue result = storeInst->getSrc();
+  SILValue storeAddr = storeInst->getDest();
+  for (Projection proj : llvm::reverse(projections)) {
+    storeAddr = cast<SingleValueInstruction>(storeAddr)->getOperand(0);
+    result = proj.createAggFromFirstLevelProjections(
+                     builder, storeInst->getLoc(),
+                     storeAddr->getType().getObjectType(), {result})
+                 .get();
   }
-  // If we failed to create any structs, bail.
-  if (result == storeInst->getSrc())
-    return nextII;
 
   // Store the new struct-wrapped value into the final base address.
   builder.createStore(storeInst->getLoc(), result, storeAddr,

test/SILOptimizer/Inputs/foo.h

@@ -0,0 +1,18 @@
+#ifndef VALIDATION_TEST_SILOPTIMIZER_FOO_H
+#define VALIDATION_TEST_SILOPTIMIZER_FOO_H
+
+struct Foo {
+  int x;
+  ~Foo() {}
+};
+
+struct Loadable {
+  int x;
+};
+
+struct Bar {
+  Loadable y;
+  ~Bar() {}
+};
+
+#endif // VALIDATION_TEST_SILOPTIMIZER_FOO_H

test/SILOptimizer/Inputs/module.modulemap

@@ -0,0 +1,3 @@
+module Foo {
+  header "foo.h"
+}

test/SILOptimizer/dont_broaden_cxx_address_only.sil

@@ -0,0 +1,44 @@
+// RUN: %target-sil-opt -silgen-cleanup %s -I %S/Inputs -enable-sil-verify-all -enable-cxx-interop | %FileCheck %s
+
+sil_stage canonical
+
+import Builtin
+import Swift
+import SwiftShims
+import Foo
+
+// Make sure we don't try to create a struct here. Foo is not loadable, even
+// though it's only property is.
+// CHECK-LABEL: @test_foo
+// CHECK: bb0
+// CHECK-NEXT: [[E:%.*]] = struct_element_addr
+// CHECK-NEXT: store %1 to [trivial] [[E]]
+// CHECK-NEXT: tuple
+// CHECK-NEXT: return
+// CHECK-LABEL: end sil function 'test_foo'
+sil shared [transparent] [serializable] [ossa] @test_foo : $@convention(method) (Int32, @thin Foo.Type) -> @out Foo {
+bb0(%0 : $*Foo, %1 : $Int32, %2 : $@thin Foo.Type):
+  %3 = struct_element_addr %0 : $*Foo, #Foo.x
+  store %1 to [trivial] %3 : $*Int32
+  %5 = tuple ()
+  return %5 : $()
+}
+
+// Make sure we create a struct for the first (loadable) type but not the second
+// type (Bar).
+// CHECK-LABEL: @test_bar
+// CHECK: bb0
+// CHECK-NEXT: [[E:%.*]] = struct_element_addr
+// CHECK-NEXT: [[AGG:%.*]] = struct $Loadable (%1 : $Int32)
+// CHECK-NEXT: store [[AGG]] to [trivial] [[E]]
+// CHECK-NEXT: tuple
+// CHECK-NEXT: return
+// CHECK-LABEL: end sil function 'test_bar'
+sil shared [transparent] [serializable] [ossa] @test_bar : $@convention(method) (Int32, @thin Bar.Type) -> @out Bar {
+bb0(%0 : $*Bar, %1 : $Int32, %2 : $@thin Bar.Type):
+  %3 = struct_element_addr %0 : $*Bar, #Bar.y
+  %3a = struct_element_addr %3 : $*Loadable, #Loadable.x
+  store %1 to [trivial] %3a : $*Int32
+  %5 = tuple ()
+  return %5 : $()
+}

tools/sil-opt/SILOpt.cpp

@@ -270,6 +270,11 @@ static cl::opt<std::string> RemarksFormat(
     cl::desc("The format used for serializing remarks (default: YAML)"),
     cl::value_desc("format"), cl::init("yaml"));
 
+static llvm::cl::opt<bool>
+    EnableCxxInterop("enable-cxx-interop",
+                     llvm::cl::desc("Enable C++ interop."),
+                     llvm::cl::init(false));
+
 static void runCommandLineSelectedPasses(SILModule *Module,
                                          irgen::IRGenModule *IRGenMod) {
   auto &opts = Module->getOptions();
@@ -348,6 +353,8 @@ int main(int argc, char **argv) {
   Invocation.getLangOptions().EnableExperimentalDifferentiableProgramming =
       EnableExperimentalDifferentiableProgramming;
 
+  Invocation.getLangOptions().EnableCXXInterop = EnableCxxInterop;
+
   Invocation.getDiagnosticOptions().VerifyMode =
       VerifyMode ? DiagnosticOptions::Verify : DiagnosticOptions::NoVerify;