Add for case support and tests

simanerush · simanerush · commit 0e1d3f17a401 · 2023-12-02T15:27:28.000-08:00
diff --git a/lib/SILGen/SILGenFunction.h b/lib/SILGen/SILGenFunction.h
@@ -2724,26 +2724,24 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   ///
   ///   This function will be called within a cleanups scope and with
   ///   InnermostPackExpansion set up properly for the context.
-  void emitDynamicPackLoop(SILLocation loc,
-                           CanPackType formalPackType,
-                           unsigned componentIndex,
-                           SILValue startingAfterIndexWithinComponent,
-                           SILValue limitWithinComponent,
-                           GenericEnvironment *openedElementEnv,
-                           bool reverse,
-                        llvm::function_ref<void(SILValue indexWithinComponent,
-                                                SILValue packExpansionIndex,
-                                                SILValue packIndex)> emitBody);
+  void emitDynamicPackLoop(
+      SILLocation loc, CanPackType formalPackType, unsigned componentIndex,
+      SILValue startingAfterIndexWithinComponent, SILValue limitWithinComponent,
+      GenericEnvironment *openedElementEnv, bool reverse,
+      llvm::function_ref<void(SILValue indexWithinComponent,
+                              SILValue packExpansionIndex, SILValue packIndex)>
+          emitBody,
+      SILBasicBlock *loopLatch = nullptr);
 
   /// A convenience version of dynamic pack loop that visits an entire
   /// pack expansion component in forward order.
-  void emitDynamicPackLoop(SILLocation loc,
-                           CanPackType formalPackType,
-                           unsigned componentIndex,
-                           GenericEnvironment *openedElementEnv,
-                        llvm::function_ref<void(SILValue indexWithinComponent,
-                                                SILValue packExpansionIndex,
-                                                SILValue packIndex)> emitBody);
+  void emitDynamicPackLoop(
+      SILLocation loc, CanPackType formalPackType, unsigned componentIndex,
+      GenericEnvironment *openedElementEnv,
+      llvm::function_ref<void(SILValue indexWithinComponent,
+                              SILValue packExpansionIndex, SILValue packIndex)>
+          emitBody,
+      SILBasicBlock *loopLatch = nullptr);
 
   /// Emit a transform on each element of a pack-expansion component
   /// of a pack, write the result into a pack-expansion component of
diff --git a/lib/SILGen/SILGenPack.cpp b/lib/SILGen/SILGenPack.cpp
@@ -662,28 +662,27 @@ void SILGenFunction::projectTupleElementsToPack(SILLocation loc,
   });
 }
 
-void SILGenFunction::emitDynamicPackLoop(SILLocation loc,
-                                         CanPackType formalPackType,
-                                         unsigned componentIndex,
-                                         GenericEnvironment *openedElementEnv,
-                      llvm::function_ref<void(SILValue indexWithinComponent,
-                                              SILValue packExpansionIndex,
-                                              SILValue packIndex)> emitBody) {
+void SILGenFunction::emitDynamicPackLoop(
+    SILLocation loc, CanPackType formalPackType, unsigned componentIndex,
+    GenericEnvironment *openedElementEnv,
+    llvm::function_ref<void(SILValue indexWithinComponent,
+                            SILValue packExpansionIndex, SILValue packIndex)>
+        emitBody,
+    SILBasicBlock *loopLatch) {
   return emitDynamicPackLoop(loc, formalPackType, componentIndex,
                              /*startAfter*/ SILValue(), /*limit*/ SILValue(),
-                             openedElementEnv, /*reverse*/false, emitBody);
+                             openedElementEnv, /*reverse*/ false, emitBody,
+                             loopLatch);
 }
 
-void SILGenFunction::emitDynamicPackLoop(SILLocation loc,
-                                         CanPackType formalPackType,
-                                         unsigned componentIndex,
-                                         SILValue startingAfterIndexInComponent,
-                                         SILValue limitWithinComponent,
-                                         GenericEnvironment *openedElementEnv,
-                                         bool reverse,
-                      llvm::function_ref<void(SILValue indexWithinComponent,
-                                              SILValue packExpansionIndex,
-                                              SILValue packIndex)> emitBody) {
+void SILGenFunction::emitDynamicPackLoop(
+    SILLocation loc, CanPackType formalPackType, unsigned componentIndex,
+    SILValue startingAfterIndexInComponent, SILValue limitWithinComponent,
+    GenericEnvironment *openedElementEnv, bool reverse,
+    llvm::function_ref<void(SILValue indexWithinComponent,
+                            SILValue packExpansionIndex, SILValue packIndex)>
+        emitBody,
+    SILBasicBlock *loopLatch) {
   assert(isa<PackExpansionType>(formalPackType.getElementType(componentIndex)));
   assert((!startingAfterIndexInComponent || !reverse) &&
          "cannot reverse with a starting index");
@@ -764,6 +763,7 @@ void SILGenFunction::emitDynamicPackLoop(SILLocation loc,
   // the incoming index - 1 if reverse)
   SILValue curIndex = incomingIndex;
   if (reverse) {
+    assert(!loopLatch && "Only forward iteration supported with loop latch");
     curIndex = B.createBuiltinBinaryFunction(loc, "sub", wordTy, wordTy,
                                              { incomingIndex, one });
   }
@@ -791,6 +791,13 @@ void SILGenFunction::emitDynamicPackLoop(SILLocation loc,
   {
     FullExpr scope(Cleanups, CleanupLocation(loc));
     emitBody(curIndex, packExpansionIndex, packIndex);
+    if (loopLatch) {
+      B.createBranch(loc, loopLatch);
+    }
+  }
+
+  if (loopLatch) {
+    B.emitBlock(loopLatch);
   }
 
   // The index to pass to the loop condition block (the current index + 1
diff --git a/lib/SILGen/SILGenStmt.cpp b/lib/SILGen/SILGenStmt.cpp
@@ -1255,9 +1255,7 @@ void StmtEmitter::visitForEachStmt(ForEachStmt *S) {
         PackType::get(SGF.getASTContext(), expansion->getType())
             ->getCanonicalType());
 
-    // Create a new basic block and jump into it.
     JumpDest loopDest = createJumpDest(S->getBody());
-    SGF.B.emitBlock(loopDest.getBlock(), S);
 
     SGF.emitDynamicPackLoop(
         SILLocation(expansion), formalPackType, 0,
@@ -1271,7 +1269,8 @@ void StmtEmitter::visitForEachStmt(ForEachStmt *S) {
           SGF.emitExprInto(expansion->getPatternExpr(), letValueInit.get());
           visit(S->getBody());
           return;
-        });
+        },
+        loopDest.getBlock());
 
     return;
   }
diff --git a/test/SILGen/pack_iteration.swift b/test/SILGen/pack_iteration.swift
@@ -30,6 +30,10 @@ enum E<T> {
 // CHECK: [[COND:%.*]] = builtin "cmp_eq_Word"([[IDX3]] : $Builtin.Word, [[PACK_LENGTH]] : $Builtin.Word) : $Builtin.Int1
 // CHECK: cond_br [[COND]], [[NONE_BB:bb[0-9]+]], [[SOME_BB:bb[0-9]+]]
 //
+// CHECK: [[NONE_BB]]:
+// CHECK:  [[FUNC_END_FUNC:%.*]] = function_ref @funcEnd : $@convention(thin) () -> ()
+// CHECK: apply [[FUNC_END_FUNC]]() : $@convention(thin) () -> ()
+//
 // CHECK: [[SOME_BB]]:
 // CHECK: [[DYN_PACK_IDX:%.*]] = dynamic_pack_index [[IDX3]] of $Pack{repeat each Element}
 // CHECK: open_pack_element [[DYN_PACK_IDX]] of <each Element> at <Pack{repeat each Element}>, shape $each Element, uuid "[[UUID:.*]]"
@@ -43,9 +47,6 @@ enum E<T> {
 // CHECK: [[IDX4:%.*]] = builtin "add_Word"([[IDX3]] : $Builtin.Word, [[IDX2]] : $Builtin.Word) : $Builtin.Word
 // CHECK: br [[LOOP_DEST]]([[IDX4]] : $Builtin.Word)
 //
-// CHECK: [[NONE_BB]]:
-// CHECK:  [[FUNC_END_FUNC:%.*]] = function_ref @funcEnd : $@convention(thin) () -> ()
-// CHECK: apply [[FUNC_END_FUNC]]() : $@convention(thin) () -> ()
 // CHECK: } // end sil function '$s14pack_iteration14iterateTrivial4overyxxQp_tRvzlF'
 func iterateTrivial<each Element>(over elements: repeat each Element) {
   for element in repeat each elements {
@@ -54,32 +55,89 @@ func iterateTrivial<each Element>(over elements: repeat each Element) {
   funcEnd()
 }
 
-// TODO: Write this test
+// CHECK-LABEL: sil hidden [ossa] @$s14pack_iteration11equalTuples3lhs3rhsSbxxQp_t_xxQp_ttRvzSQRzlF : $@convention(thin) <each Element where repeat each Element : Equatable> (@pack_guaranteed Pack{repeat each Element}, @pack_guaranteed Pack{repeat each Element}) -> Bool {
+// CHECK: bb6:
+// CHECK: [[STACK1:%.*]] = alloc_stack $(repeat each Element)
+// CHECK: [[STACK2:%.*]] = alloc_stack $(repeat each Element)
+// CHECK: [[IDX1:%.*]] = integer_literal $Builtin.Word, 0
+// CHECK: [[IDX2:%.*]] = integer_literal $Builtin.Word, 1
+// CHECK: [[PACK_LENGTH:%.*]] = pack_length $Pack{repeat each Element}
+// CHECK: br [[LOOP_DEST:bb[0-9]+]]([[IDX1]] : $Builtin.Word)
+//
+// CHECK: [[LOOP_DEST]]([[IDX3:%.*]] : $Builtin.Word):
+// CHECK: [[COND:%.*]] = builtin "cmp_eq_Word"([[IDX3]] : $Builtin.Word, [[PACK_LENGTH]] : $Builtin.Word) : $Builtin.Int1
+// CHECK: cond_br [[COND]], [[NONE_BB:bb[0-9]+]], [[SOME_BB:bb[0-9]+]]
+//
+// CHECK: [[SOME_BB]]:
+// CHECK: [[DYN_PACK_IDX:%.*]] = dynamic_pack_index [[IDX3]] of $Pack{repeat (each Element, each Element)}
+// CHECK: [[OPEN_PACK_ELT:%.*]] = open_pack_element [[DYN_PACK_IDX]] of <each Element where repeat each Element : Equatable> at <Pack{repeat each Element}>, shape $each Element, uuid "[[UUID:.*]]"
+// CHECK: [[STACK_LEFT:%.*]] = alloc_stack [lexical] $@pack_element("[[UUID]]") each Element, let, name "left"
+// CHECK: [[STACK_RIGHT:%.*]] = alloc_stack [lexical] $@pack_element("[[UUID]]") each Element, let, name "right"
+// CHECK: tuple_pack_element_addr [[DYN_PACK_IDX]] of [[STACK1]] : $*(repeat each Element) as $*@pack_element("[[UUID]]") each Element
+// CHECK: tuple_pack_element_addr [[DYN_PACK_IDX]] of [[STACK2]] : $*(repeat each Element) as $*@pack_element("[[UUID]]") each Element
+// CHECK: [[METATYPE:%.*]] = metatype $@thick (@pack_element("[[UUID]]") each Element).Type
+// CHECK: [[WITNESS_METHOD:%.*]] = witness_method $@pack_element("[[UUID]]") each Element, #Equatable."==" : <Self where Self : Equatable> (Self.Type) -> (Self, Self) -> Bool, [[OPEN_PACK_ELT]] : $Builtin.SILToken : $@convention(witness_method: Equatable) <τ_0_0 where τ_0_0 : Equatable> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0, @thick τ_0_0.Type) -> Bool
+// CHECK: apply [[WITNESS_METHOD]]<@pack_element("[[UUID]]") each Element>([[STACK_LEFT]], [[STACK_RIGHT]], [[METATYPE]]) : $@convention(witness_method: Equatable) <τ_0_0 where τ_0_0 : Equatable> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_0, @thick τ_0_0.Type) -> Bool
+//
+// CHECK: } // end sil function '$s14pack_iteration11equalTuples3lhs3rhsSbxxQp_t_xxQp_ttRvzSQRzlF'
 func equalTuples<each Element: Equatable>(lhs: (repeat each Element), rhs: (repeat each Element)) -> Bool {
-
-//  %12 = dynamic_pack_index %9 of $Pack{repeat (each Element, each Element)} // users: %19, %17, %14, %13
-//  %13 = open_pack_element %12 of <each Element where repeat each Element : Equatable> at <Pack{repeat each Element}>, shape $each Element, uuid "E53D635E-3D89-11EE-82A2-7AABAFDC7DCA" // users: %19, %17, %14
-//  %14 = tuple_pack_element_addr %12 of %4 : $*(repeat (each Element, each Element)) as $*(@pack_element("E53D635E-3D89-11EE-82A2-7AABAFDC7DCA") each Element, @pack_element("E53D635E-3D89-11EE-82A2-7AABAFDC7DCA") each Element) // users: %16, %15
-//  %15 = tuple_element_addr %14 : $*(@pack_element("E53D635E-3D89-11EE-82A2-7AABAFDC7DCA") each Element, @pack_element("E53D635E-3D89-11EE-82A2-7AABAFDC7DCA") each Element), 0 // user: %18
-//  %16 = tuple_element_addr %14 : $*(@pack_element("E53D635E-3D89-11EE-82A2-7AABAFDC7DCA") each Element, @pack_element("E53D635E-3D89-11EE-82A2-7AABAFDC7DCA") each Element), 1 // user: %20
-//  %17 = pack_element_get %12 of %0 : $*Pack{repeat each Element} as $*@pack_element("E53D635E-3D89-11EE-82A2-7AABAFDC7DCA") each Element // user: %18
-//  copy_addr %17 to [init] %15 : $*@pack_element("E53D635E-3D89-11EE-82A2-7AABAFDC7DCA") each Element // id: %18
-//  %19 = pack_element_get %12 of %1 : $*Pack{repeat each Element} as $*@pack_element("E53D635E-3D89-11EE-82A2-7AABAFDC7DCA") each Element // user: %20
-//  copy_addr %19 to [init] %16 : $*@pack_element("E53D635E-3D89-11EE-82A2-7AABAFDC7DCA") each Element // id: %20
-//  %21 = builtin "add_Word"(%9 : $Builtin.Word, %6 : $Builtin.Word) : $Builtin.Word // user: %22
-//  br bb1(%21 : $Builtin.Word)
-
+  
   for (left, right) in repeat (each lhs, each rhs) {
     guard left == right else { return false }
   }
+  
   return true
 }
 
-// TODO: Write this test
+// CHECK-LABEL: sil hidden [ossa] @$s14pack_iteration19iteratePatternMatch4overyAA1EOyxGxQp_tRvzlF : $@convention(thin) <each Element> (@pack_guaranteed Pack{repeat E<each Element>}) -> () {
+// CHECK: bb0([[PACK:%.*]] : $*Pack{repeat E<each Element>}):
+// CHECK: [[IDX1:%.*]] = integer_literal $Builtin.Word, 0
+// CHECK: [[IDX2:%.*]] = integer_literal $Builtin.Word, 1
+// CHECK: [[PACK_LENGTH:%.*]] = pack_length $Pack{repeat each Element}
+// CHECK: br [[LOOP_DEST:bb[0-9]+]]([[IDX1]] : $Builtin.Word)
+//
+// CHECK: [[LOOP_DEST]]([[IDX3:%.*]] : $Builtin.Word):
+// CHECK: [[COND:%.*]] = builtin "cmp_eq_Word"([[IDX3]] : $Builtin.Word, [[PACK_LENGTH]] : $Builtin.Word) : $Builtin.Int1
+// CHECK: cond_br [[COND]], [[NONE_BB:bb[0-9]+]], [[SOME_BB:bb[0-9]+]]
+//
+// CHECK: [[SOME_BB]]:
+// CHECK: [[DYN_PACK_IDX:%.*]] = dynamic_pack_index [[IDX3]] of $Pack{repeat E<each Element>}
+// CHECK: open_pack_element [[DYN_PACK_IDX]] of <each Element> at <Pack{repeat each Element}>, shape $each Element, uuid "[[UUID:.*]]"
+// CHECK: [[STACK:%.*]] = alloc_stack [lexical] $@pack_element("[[UUID]]") each Element, let, name "value"
+// CHECK: [[PACK_ELT_GET:%.*]] = pack_element_get [[DYN_PACK_IDX]] of [[PACK]] : $*Pack{repeat E<each Element>} as $*E<@pack_element("[[UUID]]") each Element>
+// CHECK: [[ENUM_STACK:%.*]] = alloc_stack $E<@pack_element("[[UUID]]") each Element>
+// CHECK: copy_addr [[PACK_ELT_GET]] to [init] [[ENUM_STACK]] : $*E<@pack_element("[[UUID]]") each Element>
+// CHECK: switch_enum_addr [[ENUM_STACK]] : $*E<@pack_element("[[UUID]]") each Element>, case #E.one!enumelt: [[ENUM_MATCH_BB:bb[0-9]+]], case #E.two!enumelt: [[CONTINUE_BB:bb[0-9]+]]
+//
+// CHECK: [[CONTINUE_BB]]:
+// CHECK: destroy_addr [[ENUM_STACK]]
+// CHECK: dealloc_stack [[ENUM_STACK]]
+// CHECK: dealloc_stack [[STACK]]
+// CHECK:  br [[LATCH_BB:bb[0-9]+]]
+//
+// CHECK: [[ENUM_MATCH_BB]]:
+// CHECK: [[ENUM_DATA_ADDR:%.*]] = unchecked_take_enum_data_addr %13 : $*E<@pack_element("[[UUID]]") each Element>, #E.one!enumelt
+// CHECK: copy_addr [take] [[ENUM_DATA_ADDR]] to [init] [[STACK]]
+// CHECK: [[LOOP_END_FUNC:%.*]] = function_ref @loopBodyEnd : $@convention(thin) () -> ()
+// CHECK: apply [[LOOP_END_FUNC]]() : $@convention(thin) () -> ()
+// CHECK: dealloc_stack [[ENUM_STACK]]
+// CHECK: destroy_addr [[STACK]]
+// CHECK: dealloc_stack [[STACK]]
+// CHECK:  br [[LATCH:bb[0-9]+]]
+//
+// CHECK: [[NONE_BB]]:
+// CHECK: [[FUNC_END_FUNC:%.*]] = function_ref @funcEnd : $@convention(thin) () -> ()
+// CHECK: apply [[FUNC_END_FUNC]]() : $@convention(thin) () -> ()
+//
+// CHECK: [[LATCH_BB]]:
+// CHECK: [[ADD_WORD:%.*]] = builtin "add_Word"([[IDX3]] : $Builtin.Word, [[IDX2]] : $Builtin.Word) : $Builtin.Word
+// CHECK: br [[LOOP_DEST]]([[ADD_WORD]] : $Builtin.Word)
+// CHECK: } // end sil function '$s14pack_iteration19iteratePatternMatch4overyAA1EOyxGxQp_tRvzlF'
 func iteratePatternMatch<each Element>(over element: repeat E<each Element>) {
   for case .one(let value) in repeat each element {
-    print(value)
+    loopBodyEnd()
   }
+  funcEnd()
 }
 
 
