Merge pull request #64272 from hborla/emit-materialize-pack-expr

[SILGen] Emit `MaterializePackExpr`s.

Author: hborla

lib/SILGen/SILGenApply.cpp

@@ -3793,8 +3793,6 @@ class ArgEmitter {
     auto expansionExpr =
       cast<PackExpansionExpr>(expr->getSemanticsProvidingExpr());
 
-    SGF.prepareToEmitPackExpansionExpr(expansionExpr);
-
     // TODO: try to borrow the existing elements if we can do that within
     // the limitations of the SIL representation.
 

lib/SILGen/SILGenExpr.cpp

@@ -1534,8 +1534,6 @@ RValueEmitter::visitPackExpansionExpr(PackExpansionExpr *E,
   assert(init && init->canPerformPackExpansionInitialization() &&
          "cannot emit a PackExpansionExpr without an appropriate context");
 
-  SGF.prepareToEmitPackExpansionExpr(E);
-
   auto type = E->getType()->getCanonicalType();
   assert(isa<PackExpansionType>(type));
   auto formalPackType = CanPackType::get(SGF.getASTContext(), {type});
@@ -1570,7 +1568,8 @@ RValueEmitter::visitPackElementExpr(PackElementExpr *E, SGFContext C) {
 
 RValue
 RValueEmitter::visitMaterializePackExpr(MaterializePackExpr *E, SGFContext C) {
-  llvm_unreachable("not implemented for MaterializePackExpr");
+  // Always emitted through `visitPackElementExpr`.
+  llvm_unreachable("materialized pack outside of PackElementExpr");
 }
 
 RValue RValueEmitter::visitArchetypeToSuperExpr(ArchetypeToSuperExpr *E,
@@ -6202,8 +6201,6 @@ RValue SILGenFunction::emitPlusZeroRValue(Expr *E) {
 
 static void emitIgnoredPackExpansion(SILGenFunction &SGF,
                                      PackExpansionExpr *E) {
-  SGF.prepareToEmitPackExpansionExpr(E);
-
   auto expansionType =
     cast<PackExpansionType>(E->getType()->getCanonicalType());
   auto formalPackType = CanPackType::get(SGF.getASTContext(), expansionType);

lib/SILGen/SILGenFunction.h

@@ -560,8 +560,17 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   SILValue ExpectedExecutor;
 
   struct ActivePackExpansion {
-    SILValue ExpansionIndex;
     GenericEnvironment *OpenedElementEnv;
+    SILValue ExpansionIndex;
+
+    /// Mapping from temporary pack expressions to their values. These
+    /// are evaluated once, with their elements projected in a dynamic
+    /// pack loop.
+    llvm::SmallDenseMap<MaterializePackExpr *, SILValue>
+      MaterializedPacks;
+
+    ActivePackExpansion(GenericEnvironment *OpenedElementEnv)
+        : OpenedElementEnv(OpenedElementEnv) {}
   };
 
   /// The innermost active pack expansion.
@@ -1819,7 +1828,7 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   /// places that will need to support any sort of future feature
   /// where e.g. certain `each` operands need to be evaluated once
   /// for the entire expansion.
-  void prepareToEmitPackExpansionExpr(PackExpansionExpr *E) {}
+  void prepareToEmitPackExpansionExpr(PackExpansionExpr *E);
 
   //
   // Helpers for emitting ApplyExpr chains.

lib/SILGen/SILGenLValue.cpp

@@ -3342,6 +3342,24 @@ LValue SILGenLValue::visitPackElementExpr(PackElementExpr *e,
     }
   }
 
+  if (auto packExpr = dyn_cast<MaterializePackExpr>(refExpr)) {
+    auto *activeExpansion = SGF.getInnermostPackExpansion();
+    assert(activeExpansion->MaterializedPacks.count(packExpr) &&
+           "didn't materialize pack before dynamic pack loop emission");
+
+    auto elementTy =
+      SGF.getLoweredType(substFormalType).getAddressType();
+    auto tupleAddr = activeExpansion->MaterializedPacks.find(packExpr);
+    auto packIndex = activeExpansion->ExpansionIndex;
+    auto elementAddr =
+      SGF.B.createTuplePackElementAddr(e, packIndex, tupleAddr->second, elementTy);
+    return LValue::forAddress(accessKind,
+                              ManagedValue::forLValue(elementAddr),
+                              /*access enforcement*/ None,
+                              AbstractionPattern(substFormalType),
+                              substFormalType);
+  }
+
   SGF.SGM.diagnose(refExpr, diag::not_implemented,
                    "emission of 'each' for this kind of expression");
   auto loweredTy = SGF.getLoweredType(substFormalType).getAddressType();

lib/SILGen/SILGenPack.cpp

@@ -13,6 +13,7 @@
 #include "Initialization.h"
 #include "Scope.h"
 #include "SILGenFunction.h"
+#include "swift/AST/ASTWalker.h"
 #include "swift/AST/GenericEnvironment.h"
 
 using namespace swift;
@@ -165,8 +166,52 @@ class PartialDestroyRemainingTupleCleanup : public Cleanup {
 #endif
   }
 };
+
+/// An ASTWalker to emit tuple values in `MaterializePackExpr` nodes.
+///
+/// Materialized packs are emitted inside a pack expansion context before
+/// entering the dynamic pack loop so that the values are only evaluated
+/// once, rather than at each pack element iteration.
+struct MaterializePackEmitter : public ASTWalker {
+  SILGenFunction &SGF;
+
+  MaterializePackEmitter(SILGenFunction &SGF) : SGF(SGF) {}
+
+  ASTWalker::PreWalkResult<Expr *> walkToExprPre(Expr *expr) override {
+    using Action = ASTWalker::Action;
+
+    // Don't walk into nested pack expansions.
+    if (isa<PackExpansionExpr>(expr))
+      return Action::SkipChildren(expr);
+
+    if (auto *packExpr = dyn_cast<MaterializePackExpr>(expr)) {
+      auto *fromExpr = packExpr->getFromExpr();
+      assert(fromExpr->getType()->is<TupleType>());
+
+      auto &lowering = SGF.getTypeLowering(fromExpr->getType());
+      auto loweredTy = lowering.getLoweredType();
+      auto tupleAddr = SGF.emitTemporaryAllocation(fromExpr, loweredTy);
+      auto init = SGF.useBufferAsTemporary(tupleAddr, lowering);
+      SGF.emitExprInto(fromExpr, init.get());
+
+      // Write the tuple value to a side table in the active pack expansion
+      // to be projected later within the dynamic pack loop.
+      auto *activeExpansion = SGF.getInnermostPackExpansion();
+      activeExpansion->MaterializedPacks[packExpr] = tupleAddr;
+    }
+
+    return Action::Continue(expr);
+  }
+};
+
 } // end anonymous namespace
 
+void
+SILGenFunction::prepareToEmitPackExpansionExpr(PackExpansionExpr *E) {
+  MaterializePackEmitter tempPackEmission(*this);
+  E->getPatternExpr()->walk(tempPackEmission);
+}
+
 CleanupHandle SILGenFunction::enterDeallocPackCleanup(SILValue temp) {
   assert(temp->getType().isAddress() &&  "dealloc must have an address type");
   assert(temp->getType().is<SILPackType>());
@@ -390,6 +435,17 @@ void SILGenFunction::emitDynamicPackLoop(SILLocation loc,
          "cannot reverse with a starting index");
   ASTContext &ctx = SGM.getASTContext();
 
+  // Save and restore the innermost pack expansion.
+  ActivePackExpansion activeExpansionRecord = {
+    openedElementEnv
+  };
+
+  llvm::SaveAndRestore<ActivePackExpansion*>
+    packExpansionScope(InnermostPackExpansion, &activeExpansionRecord);
+
+  if (auto *expansion = loc.getAsASTNode<PackExpansionExpr>())
+    prepareToEmitPackExpansionExpr(expansion);
+
   auto wordTy = SILType::getBuiltinWordType(ctx);
   auto boolTy = SILType::getBuiltinIntegerType(1, ctx);
 
@@ -461,6 +517,7 @@ void SILGenFunction::emitDynamicPackLoop(SILLocation loc,
   // Construct the dynamic pack index into the component.
   SILValue packExpansionIndex =
     B.createDynamicPackIndex(loc, curIndex, formalDynamicPackType);
+  getInnermostPackExpansion()->ExpansionIndex = packExpansionIndex;
 
   // If there's an opened element environment, open it here.
   if (openedElementEnv) {
@@ -478,14 +535,6 @@ void SILGenFunction::emitDynamicPackLoop(SILLocation loc,
   // Emit the loop body in a scope as a convenience, since it's necessary
   // to avoid dominance problems anyway.
   {
-    // Save and restore the innermost pack expansion.
-    ActivePackExpansion activeExpansionRecord = {
-      packExpansionIndex,
-      openedElementEnv
-    };
-    llvm::SaveAndRestore<ActivePackExpansion*>
-      packExpansionScope(InnermostPackExpansion, &activeExpansionRecord);
-
     FullExpr scope(Cleanups, CleanupLocation(loc));
     emitBody(curIndex, packExpansionIndex, packIndex);
   }

lib/SILOptimizer/Mandatory/PMOMemoryUseCollector.cpp

@@ -289,7 +289,8 @@ bool ElementUseCollector::collectUses(SILValue Pointer) {
     if (auto *si = dyn_cast<StoreInst>(User)) {
       if (UI->getOperandNumber() == StoreInst::Dest) {
         if (auto tupleType = PointeeType.getAs<TupleType>()) {
-          if (!tupleType->isEqual(Module.getASTContext().TheEmptyTupleType)) {
+          if (!tupleType->isEqual(Module.getASTContext().TheEmptyTupleType) &&
+              !tupleType->containsPackExpansionType()) {
             UsesToScalarize.push_back(User);
             continue;
           }
@@ -325,7 +326,8 @@ bool ElementUseCollector::collectUses(SILValue Pointer) {
       // If this is a copy of a tuple, we should scalarize it so that we don't
       // have an access that crosses elements.
       if (auto tupleType = PointeeType.getAs<TupleType>()) {
-        if (!tupleType->isEqual(Module.getASTContext().TheEmptyTupleType)) {
+        if (!tupleType->isEqual(Module.getASTContext().TheEmptyTupleType) &&
+            !tupleType->containsPackExpansionType()) {
           UsesToScalarize.push_back(CAI);
           continue;
         }

test/Interpreter/variadic_generic_tuples.swift

@@ -57,4 +57,19 @@ tuples.test("makeTuple3") {
   expectEqual("(Swift.Int, Swift.Float)", _typeName(makeTuple3(t: Int.self, u: Float.self)))
 }
 
-runAllTests()
+func makeTuple<each Element>(
+  _ element: repeat each Element
+) -> (repeat each Element) {
+  return (repeat each element)
+}
+
+func expandTupleElements<each T: Equatable>(_ value: repeat each T) {
+  let values = makeTuple(repeat each value)
+  _ = (repeat expectEqual(each value, each values.element))
+}
+
+tuples.test("expandTuple") {
+  expandTupleElements(1, "hello", true)
+}
+
+runAllTests()

test/SILGen/variadic-generic-tuples.swift

@@ -112,3 +112,56 @@ public struct Container<each T> {
 // CHECK-NEXT:    return [[RET]] : $()
 
 // CHECK-LABEL: sil {{.*}}@$s4main9ContainerV7storageAA6StoredVyxGxQp_tvM
+
+struct Wrapper<Value> {
+  let value: Value
+}
+
+// CHECK-LABEL: @$s4main17wrapTupleElementsyAA7WrapperVyxGxQp_txxQpRvzlF : $@convention(thin) <each T> (@pack_guaranteed Pack{repeat each T}) -> @pack_out Pack{repeat Wrapper<each T>}
+func wrapTupleElements<each T>(_ value: repeat each T) -> (repeat Wrapper<each T>) {
+  // CHECK: [[RETURN_VAL:%.*]] : $*Pack{repeat Wrapper<each T>}
+
+  // CHECK: [[VAR:%.*]] = alloc_stack [lexical] $(repeat each T)
+  let values = (repeat each value)
+
+  // Create a temporary for the 'values' in 'each values.element'
+  // CHECK: bb3:
+  // CHECK-NEXT: [[TEMP:%.*]] = alloc_stack $(repeat each T)
+  // CHECK-NEXT: copy_addr [[VAR]] to [init] [[TEMP]] : $*(repeat each T)
+
+  // Integer values for dynamic pack loop
+  // CHECK-NEXT: [[ZERO:%.*]] = integer_literal $Builtin.Word, 0
+  // CHECK-NEXT: [[ONE:%.*]] = integer_literal $Builtin.Word, 1
+  // CHECK-NEXT: [[PACK_LEN:%.*]] = pack_length $Pack{repeat each T}
+  // CHECK-NEXT: br bb4([[ZERO]] : $Builtin.Word)
+
+  // Loop condition
+  // CHECK: bb4([[INDEX:%.*]] : $Builtin.Word)
+  // CHECK-NEXT: [[INDEX_EQ_LEN:%.*]] = builtin "cmp_eq_Word"([[INDEX]] : $Builtin.Word, [[PACK_LEN]] : $Builtin.Word) : $Builtin.Int1
+  // CHECK-NEXT: cond_br [[INDEX_EQ_LEN]], bb6, bb5
+
+  // Loop body
+  // CHECK: bb5:
+  // CHECK-NEXT: [[CUR_INDEX:%.*]] = dynamic_pack_index [[INDEX]] of $Pack{repeat Wrapper<each T>}
+  // CHECK-NEXT: open_pack_element [[CUR_INDEX]] of <each T> at <Pack{repeat each T}>, shape $T, uuid [[UUID:".*"]]
+  // CHECK-NEXT: [[RETURN_VAL_ELT_ADDR:%.*]] = pack_element_get [[CUR_INDEX]] of [[RETURN_VAL]] : $*Pack{repeat Wrapper<each T>} as $*Wrapper<@pack_element([[UUID]]) T>
+  // CHECK-NEXT: [[METATYPE:%.*]] = metatype $@thin Wrapper<@pack_element([[UUID]]) T>.Type
+  // CHECK-NEXT: [[TUPLE_ELT_ADDR:%.*]] = tuple_pack_element_addr [[CUR_INDEX]] of [[TEMP]] : $*(repeat each T) as $*@pack_element([[UUID]]) T
+  // CHECK-NEXT: [[INIT_ARG:%.*]] = alloc_stack $@pack_element([[UUID]]) T
+  // CHECK-NEXT: copy_addr [[TUPLE_ELT_ADDR]] to [init] [[INIT_ARG]] : $*@pack_element([[UUID]]) T
+  // function_ref Wrapper.init(value:)
+  // CHECK: [[INIT:%.*]] = function_ref @$s4main7WrapperV5valueACyxGx_tcfC : $@convention(method) <τ_0_0> (@in τ_0_0, @thin Wrapper<τ_0_0>.Type) -> @out Wrapper<τ_0_0>
+  // CHECK-NEXT: apply [[INIT]]<@pack_element([[UUID]]) T>([[RETURN_VAL_ELT_ADDR]], [[INIT_ARG]], [[METATYPE]]) : $@convention(method) <τ_0_0> (@in τ_0_0, @thin Wrapper<τ_0_0>.Type) -> @out Wrapper<τ_0_0>
+  // CHECK-NEXT: dealloc_stack [[INIT_ARG]] : $*@pack_element([[UUID]]) T
+  // CHECK-NEXT: [[NEXT_INDEX:%.*]] = builtin "add_Word"([[INDEX]] : $Builtin.Word, [[ONE]] : $Builtin.Word) : $Builtin.Word
+  // CHECK-NEXT: br bb4([[NEXT_INDEX]] : $Builtin.Word)
+
+  return (repeat Wrapper(value: each values.element))
+
+  // CHECK: destroy_addr [[TEMP]] : $*(repeat each T)
+  // CHECK: dealloc_stack [[TEMP]] : $*(repeat each T)
+  // CHECK: destroy_addr [[VAR]] : $*(repeat each T)
+  // CHECK: dealloc_stack [[VAR]] : $*(repeat each T)
+  // CHECK-NEXT:    [[RET:%.*]] = tuple ()
+  // CHECK-NEXT:    return [[RET]] : $()
+}