Merge pull request #66768 from Azoy/pack-count-builtin

[AST] Add Builtin.packLength

Author: Azoy

include/swift/AST/Builtins.def

@@ -538,6 +538,11 @@ BUILTIN_SIL_OPERATION(WithUnsafeThrowingContinuation, "withUnsafeThrowingContinu
 /// Force the current task to be rescheduled on the specified actor.
 BUILTIN_SIL_OPERATION(HopToActor, "hopToActor", None)
 
+/// packLength: <each T>(_: repeat each T) -> Int
+///
+/// Returns the number of items in a pack.
+BUILTIN_SIL_OPERATION(PackLength, "packLength", Special)
+
 #undef BUILTIN_SIL_OPERATION
 
 // BUILTIN_RUNTIME_CALL - A call into a runtime function.

lib/AST/Builtins.cpp

@@ -235,21 +235,25 @@ static const char * const GenericParamNames[] = {
 };
 
 static GenericTypeParamDecl*
-createGenericParam(ASTContext &ctx, const char *name, unsigned index) {
+createGenericParam(ASTContext &ctx, const char *name, unsigned index,
+                   bool isParameterPack = false) {
   ModuleDecl *M = ctx.TheBuiltinModule;
   Identifier ident = ctx.getIdentifier(name);
   return GenericTypeParamDecl::createImplicit(
-      &M->getMainFile(FileUnitKind::Builtin), ident, /*depth*/ 0, index);
+      &M->getMainFile(FileUnitKind::Builtin), ident, /*depth*/ 0, index,
+                      isParameterPack);
 }
 
 /// Create a generic parameter list with multiple generic parameters.
 static GenericParamList *getGenericParams(ASTContext &ctx,
-                                          unsigned numParameters) {
+                                          unsigned numParameters,
+                                          bool areParameterPacks = false) {
   assert(numParameters <= std::size(GenericParamNames));
 
   SmallVector<GenericTypeParamDecl *, 2> genericParams;
   for (unsigned i = 0; i != numParameters; ++i)
-    genericParams.push_back(createGenericParam(ctx, GenericParamNames[i], i));
+    genericParams.push_back(createGenericParam(ctx, GenericParamNames[i], i,
+                                               areParameterPacks));
 
   auto paramList = GenericParamList::create(ctx, SourceLoc(), genericParams,
                                             SourceLoc());
@@ -678,9 +682,11 @@ namespace {
 
   public:
     BuiltinFunctionBuilder(ASTContext &ctx, unsigned numGenericParams = 1,
-                           bool wantsAdditionalAnyObjectRequirement = false)
+                           bool wantsAdditionalAnyObjectRequirement = false,
+                           bool areParametersPacks = false)
         : Context(ctx) {
-      TheGenericParamList = getGenericParams(ctx, numGenericParams);
+      TheGenericParamList = getGenericParams(ctx, numGenericParams,
+                                             areParametersPacks);
       if (wantsAdditionalAnyObjectRequirement) {
         Requirement req(RequirementKind::Conformance,
                         TheGenericParamList->getParams()[0]->getInterfaceType(),
@@ -750,7 +756,7 @@ namespace {
       unsigned Index;
       Type build(BuiltinFunctionBuilder &builder) const {
         return builder.TheGenericParamList->getParams()[Index]
-            ->getDeclaredInterfaceType();
+                    ->getDeclaredInterfaceType();
       }
     };
     struct LambdaGenerator {
@@ -767,6 +773,16 @@ namespace {
         return MetatypeType::get(Object.build(builder), Repr);
       }
     };
+    template <class T>
+    struct PackExpansionGenerator {
+      T Object;
+      Type build(BuiltinFunctionBuilder &builder) const {
+        auto patternTy = Object.build(builder);
+        SmallVector<Type, 2> packs;
+        patternTy->getTypeParameterPacks(packs);
+        return PackExpansionType::get(patternTy, packs[0]);
+      }
+    };
   };
 } // end anonymous namespace
 
@@ -814,6 +830,12 @@ makeMetatype(const T &object,
   return { object, repr };
 }
 
+template <class T>
+static BuiltinFunctionBuilder::PackExpansionGenerator<T>
+makePackExpansion(const T &object) {
+  return { object };
+}
+
 /// Create a function with type <T> T -> ().
 static ValueDecl *getRefCountingOperation(ASTContext &ctx, Identifier id) {
   return getBuiltinFunction(ctx, id, _thin,
@@ -1934,6 +1956,18 @@ static ValueDecl *getHopToActor(ASTContext &ctx, Identifier id) {
   return builder.build(id);
 }
 
+static ValueDecl *getPackLength(ASTContext &ctx, Identifier id) {
+  BuiltinFunctionBuilder builder(ctx, /* genericParamCount */ 1,
+                                 /* anyObject */ false,
+                                 /* areParametersPack */ true);
+
+  auto paramTy = makeMetatype(makeTuple(makePackExpansion(makeGenericParam())));
+  builder.addParameter(paramTy);
+  builder.setResult(makeConcrete(BuiltinIntegerType::getWordType(ctx)));
+
+  return builder.build(id);
+}
+
 /// An array of the overloaded builtin kinds.
 static const OverloadedBuiltinKind OverloadedBuiltinKinds[] = {
   OverloadedBuiltinKind::None,
@@ -2975,6 +3009,9 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
 
   case BuiltinValueKind::AutoDiffAllocateSubcontextWithType:
     return getAutoDiffAllocateSubcontext(Context, Id);
+
+  case BuiltinValueKind::PackLength:
+    return getPackLength(Context, Id);
   }
 
   llvm_unreachable("bad builtin value!");

lib/SILGen/SILGenBuiltin.cpp

@@ -1708,6 +1708,19 @@ static ManagedValue emitBuiltinHopToActor(SILGenFunction &SGF, SILLocation loc,
   return ManagedValue::forObjectRValueWithoutOwnership(SGF.emitEmptyTuple(loc));
 }
 
+static ManagedValue emitBuiltinPackLength(SILGenFunction &SGF, SILLocation loc,
+                                          SubstitutionMap subs,
+                                          ArrayRef<ManagedValue> args,
+                                          SGFContext C) {
+  auto argTy = args[0].getType().getASTType();
+  auto tupleTy = CanTupleType(argTy->getMetatypeInstanceType()
+      ->castTo<TupleType>());
+  auto packTy = tupleTy.getInducedPackType();
+
+  return ManagedValue::forObjectRValueWithoutOwnership(
+    SGF.B.createPackLength(loc, packTy));
+}
+
 static ManagedValue emitBuiltinAutoDiffCreateLinearMapContextWithType(
     SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
     ArrayRef<ManagedValue> args, SGFContext C) {

test/IRGen/builtin_pack_length.swift

@@ -0,0 +1,40 @@
+// RUN: %target-swift-frontend -module-name builtins -enable-builtin-module -Xllvm -sil-disable-pass=target-constant-folding -disable-access-control -primary-file %s -emit-ir -o - -disable-objc-attr-requires-foundation-module -disable-availability-checking -O | %FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-%target-runtime
+
+// REQUIRES: CPU=x86_64 || CPU=arm64 || CPU=arm64e
+
+import Builtin
+
+// CHECK: define {{.*}} @"$s8builtins9packCountyBwxxQpRvzlF"(ptr {{.*}} %0, i64 {{.*}} [[PACK_COUNT:%.*]], ptr {{.*}} %"each T")
+// CHECK: ret i64 [[PACK_COUNT]]
+func packCount<each T>(_: repeat each T) -> Builtin.Word {
+  Builtin.packLength((repeat each T).self)
+}
+
+// CHECK: define {{.*}} @"$s8builtins12packCountUseBwyF"()
+// CHECK-NEXT: entry:
+// CHECK-NEXT: ret i64 4
+func packCountUse() -> Builtin.Word {
+  packCount(123, 321, 456, 654)
+}
+
+// CHECK: define {{.*}} @"$s8builtins18weirdPackCountUse0yBwxxQpRvzlF"(ptr {{.*}} %0, i64 [[PACK_COUNT:%.*]], ptr {{.*}} %"each T")
+// CHECK: [[ADD:%.*]] = add i64 [[PACK_COUNT]], 2
+// CHECK-NEXT: ret i64 [[ADD]]
+func weirdPackCountUse0<each T>(_ x: repeat each T) -> Builtin.Word {
+  Builtin.packLength((String, repeat each T, Int).self)
+}
+
+// CHECK: define {{.*}} @"$s8builtins18weirdPackCountUse1yBwxxQpRvzlF"(ptr {{.*}} %0, i64 %1, ptr {{.*}} %"each T")
+// CHECK: ret i64 3
+func weirdPackCountUse1<each T>(_ x: repeat each T) -> Builtin.Word {
+  Builtin.packLength((String, (repeat each T), Int).self)
+}
+
+struct Pack<each T> {
+// CHECK: define {{.*}} @"$s8builtins4PackV5countBwvgZ"(i64 returned [[PACK_COUNT:%.*]], ptr nocapture readnone %"each T")
+// CHECK-NEXT: entry:
+// CHECK-NEXT: ret i64 [[PACK_COUNT]]
+  static var count: Builtin.Word {
+    Builtin.packLength((repeat each T).self)
+  }
+}

test/SILGen/builtins.swift

@@ -876,3 +876,12 @@ func assumeTrue(_ x: Builtin.Int1) {
 func assumeAlignment(_ p: Builtin.RawPointer, _ x: Builtin.Word) {
   Builtin.assumeAlignment(p, x)
 }
+
+// CHECK-LABEL: sil hidden [ossa] @$s8builtins9packCountyBwxxQpRvzlF : $@convention(thin) <each T> (@pack_guaranteed Pack{repeat each T}) -> Builtin.Word {
+// CHECK: bb0(%0 : $*Pack{repeat each T}):
+// CHECK:   [[META:%.*]] = metatype $@thin (repeat each T).Type
+// CHECK:   [[PACK_LENGTH:%.*]] = pack_length $Pack{repeat each T}
+// CHECK:   return [[PACK_LENGTH]] : $Builtin.Word
+func packCount<each T>(_ x: repeat each T) -> Builtin.Word {
+  Builtin.packLength((repeat each T).self)
+}