Merge pull request #70327 from Azoy/inject-get-enum-tag-builtins

[IRGen] Add getEnumTag and injectEnumTag builtins

Author: Azoy

include/swift/AST/Builtins.def

@@ -1048,6 +1048,21 @@ BUILTIN_MISC_OPERATION_WITH_SILGEN(BuildDefaultActorExecutorRef,
 BUILTIN_MISC_OPERATION_WITH_SILGEN(BuildMainActorExecutorRef,
                                    "buildMainActorExecutorRef", "n", Special)
 
+/// getEnumTag: <T>(_: T) -> Builtin.Int32
+///
+/// Given a dynamic generic value, unsafely assume it is an enum type and call
+/// the getEnumTag vwt function.
+BUILTIN_MISC_OPERATION_WITH_SILGEN(GetEnumTag, "getEnumTag", "", Special)
+
+/// injectEnumTag: <T>(_: inout T, _: Builtin.Int32) -> ()
+///
+/// Given a dynamic inout generic value, unsafely assume it is an enum type and
+/// inject the given tag into it.
+///
+/// Note: This assume that either 1. the given tag has no payload or 2. the
+/// tag's payload is already initialized with the given source.
+BUILTIN_MISC_OPERATION_WITH_SILGEN(InjectEnumTag, "injectEnumTag", "", Special)
+
 #undef BUILTIN_MISC_OPERATION_WITH_SILGEN
 
 #undef BUILTIN_MISC_OPERATION

include/swift/SIL/AddressWalker.h

@@ -259,6 +259,8 @@ TransitiveAddressWalker<Impl>::walk(SILValue projectedAddress) && {
         case BuiltinValueKind::GenericXor:
         case BuiltinValueKind::TaskRunInline:
         case BuiltinValueKind::ZeroInitializer:
+        case BuiltinValueKind::GetEnumTag:
+        case BuiltinValueKind::InjectEnumTag:
           callVisitUse(op);
           continue;
         default:

lib/AST/Builtins.cpp

@@ -1978,6 +1978,26 @@ static ValueDecl *getPackLength(ASTContext &ctx, Identifier id) {
   return builder.build(id);
 }
 
+static ValueDecl *getGetEnumTag(ASTContext &ctx, Identifier id) {
+  BuiltinFunctionBuilder builder(ctx, /* genericParamCount */ 1);
+
+  auto paramTy = makeGenericParam();
+  builder.addParameter(paramTy);
+  builder.setResult(makeConcrete(BuiltinIntegerType::get(32, ctx)));
+
+  return builder.build(id);
+}
+
+static ValueDecl *getInjectEnumTag(ASTContext &ctx, Identifier id) {
+  BuiltinFunctionBuilder builder(ctx, /* genericParamCount */ 1);
+
+  builder.addParameter(makeGenericParam(), ParamSpecifier::InOut);
+  builder.addParameter(makeConcrete(BuiltinIntegerType::get(32, ctx)));
+  builder.setResult(makeConcrete(TupleType::getEmpty(ctx)));
+
+  return builder.build(id);
+}
+
 /// An array of the overloaded builtin kinds.
 static const OverloadedBuiltinKind OverloadedBuiltinKinds[] = {
   OverloadedBuiltinKind::None,
@@ -3019,6 +3039,12 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
 
   case BuiltinValueKind::PackLength:
     return getPackLength(Context, Id);
+
+  case BuiltinValueKind::GetEnumTag:
+    return getGetEnumTag(Context, Id);
+
+  case BuiltinValueKind::InjectEnumTag:
+    return getInjectEnumTag(Context, Id);
   }
 
   llvm_unreachable("bad builtin value!");

lib/IRGen/GenBuiltin.cpp

@@ -31,8 +31,10 @@
 #include "GenCast.h"
 #include "GenConcurrency.h"
 #include "GenDistributed.h"
+#include "GenEnum.h"
 #include "GenPointerAuth.h"
 #include "GenIntegerLiteral.h"
+#include "GenOpaque.h"
 #include "IRGenFunction.h"
 #include "IRGenModule.h"
 #include "LoadableTypeInfo.h"
@@ -1503,6 +1505,7 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
     out.add(pointerSrc);
     return;
   }
+
   if (Builtin.ID == BuiltinValueKind::AllocVector) {
     (void)args.claimAll();
     IGF.emitTrap("escaped vector allocation", /*EmitUnreachable=*/true);
@@ -1511,5 +1514,47 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
     IGF.Builder.emitBlock(contBB);
     return;
   }
+
+  if (Builtin.ID == BuiltinValueKind::GetEnumTag) {
+    auto arg = args.claimNext();
+    auto ty = argTypes[0];
+    auto &ti = IGF.getTypeInfo(ty);
+
+    // If the type is just an archetype, then we know nothing about the enum
+    // strategy for it. Just call the vwt function. Otherwise, we know that this
+    // is at least an enum and can optimize away some of the cost of getEnumTag.
+    if (!ty.is<ArchetypeType>()) {
+      assert(ty.getEnumOrBoundGenericEnum() && "expected enum type in "
+             "getEnumTag builtin!");
+
+      auto &strategy = getEnumImplStrategy(IGF.IGM, ty);
+
+      out.add(strategy.emitGetEnumTag(IGF, ty, ti.getAddressForPointer(arg)));
+      return;
+    }
+
+    out.add(emitGetEnumTagCall(IGF, ty, ti.getAddressForPointer(arg)));
+    return;
+  }
+
+  if (Builtin.ID == BuiltinValueKind::InjectEnumTag) {
+    auto input = args.claimNext();
+    auto tag = args.claimNext();
+    auto inputTy = argTypes[0];
+    auto &inputTi = IGF.getTypeInfo(inputTy);
+
+    // In order for us to call 'storeTag' on an enum strategy (when type is not
+    // an archetype), we'd need to be able to map the tag back into an
+    // EnumElementDecl which might be fragile. We don't really care about being
+    // able to optimize this vwt function call anyway because we expect most
+    // use cases to be the truly dynamic case where the compiler has no static
+    // information about the type to be able to optimize it away. Just call the
+    // vwt function.
+
+    emitDestructiveInjectEnumTagCall(IGF, inputTy, tag,
+                                     inputTi.getAddressForPointer(input));
+    return;
+  }
+
   llvm_unreachable("IRGen unimplemented for this builtin!");
 }

lib/SIL/IR/OperandOwnership.cpp

@@ -879,6 +879,8 @@ BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, PoundAssert)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, GlobalStringTablePointer)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, TypePtrAuthDiscriminator)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, TargetOSVersionAtLeast)
+BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, GetEnumTag)
+BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, InjectEnumTag)
 OperandOwnership OperandOwnershipBuiltinClassifier::visitCopy(BuiltinInst *bi,
                                                               StringRef) {
   if (bi->getFunction()->getConventions().useLoweredAddresses()) {

lib/SIL/IR/ValueOwnership.cpp

@@ -618,6 +618,8 @@ CONSTANT_OWNERSHIP_BUILTIN(None, CreateTaskGroupWithFlags)
 CONSTANT_OWNERSHIP_BUILTIN(None, DestroyTaskGroup)
 CONSTANT_OWNERSHIP_BUILTIN(None, TaskRunInline)
 CONSTANT_OWNERSHIP_BUILTIN(None, Copy)
+CONSTANT_OWNERSHIP_BUILTIN(None, GetEnumTag)
+CONSTANT_OWNERSHIP_BUILTIN(None, InjectEnumTag)
 
 #undef CONSTANT_OWNERSHIP_BUILTIN
 

lib/SIL/Utils/MemAccessUtils.cpp

@@ -2647,6 +2647,12 @@ static void visitBuiltinAddress(BuiltinInst *builtin,
       }
       return;
 
+    // These builtins take a generic 'T' as their operand.
+    case BuiltinValueKind::GetEnumTag:
+    case BuiltinValueKind::InjectEnumTag:
+      visitor(&builtin->getAllOperands()[0]);
+      return;
+
     // Arrays: (T.Type, Builtin.RawPointer, Builtin.RawPointer,
     // Builtin.Word)
     case BuiltinValueKind::CopyArray:

lib/SIL/Verifier/SILVerifier.cpp

@@ -559,6 +559,11 @@ struct ImmutableAddressUseVerifier {
           if (isPolymorphicBuiltin(*builtinKind)) {
             break;
           }
+
+          // Get enum tag borrows its operand address value.
+          if (builtinKind == BuiltinValueKind::GetEnumTag) {
+            return false;
+          }
         }
 
         // Otherwise this is a builtin that we are not expecting to see, so bail

lib/SILGen/SILGenBuiltin.cpp

@@ -1817,6 +1817,34 @@ static ManagedValue emitBuiltinBuildMainActorExecutorRef(
                               BuiltinValueKind::BuildMainActorExecutorRef);
 }
 
+static ManagedValue emitBuiltinGetEnumTag(SILGenFunction &SGF, SILLocation loc,
+                                          SubstitutionMap subs,
+                                          ArrayRef<ManagedValue> args,
+                                          SGFContext C) {
+  auto &ctx = SGF.getASTContext();
+
+  auto bi = SGF.B.createBuiltin(
+    loc, ctx.getIdentifier(getBuiltinName(BuiltinValueKind::GetEnumTag)),
+    SILType::getBuiltinIntegerType(32, ctx), subs,
+    { args[0].getValue() });
+
+  return ManagedValue::forObjectRValueWithoutOwnership(bi);
+}
+
+static ManagedValue emitBuiltinInjectEnumTag(SILGenFunction &SGF, SILLocation loc,
+                                             SubstitutionMap subs,
+                                             ArrayRef<ManagedValue> args,
+                                             SGFContext C) {
+  auto &ctx = SGF.getASTContext();
+
+  auto bi = SGF.B.createBuiltin(
+    loc, ctx.getIdentifier(getBuiltinName(BuiltinValueKind::InjectEnumTag)),
+    SILType::getEmptyTupleType(ctx), subs,
+    { args[0].getValue(), args[1].getValue() });
+
+  return ManagedValue::forObjectRValueWithoutOwnership(bi);
+}
+
 llvm::Optional<SpecializedEmitter>
 SpecializedEmitter::forDecl(SILGenModule &SGM, SILDeclRef function) {
   // Only consider standalone declarations in the Builtin module.

lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp

@@ -157,6 +157,8 @@ static bool isBarrier(SILInstruction *inst) {
     case BuiltinValueKind::StackDealloc:
     case BuiltinValueKind::AllocVector:
     case BuiltinValueKind::AssumeAlignment:
+    case BuiltinValueKind::GetEnumTag:
+    case BuiltinValueKind::InjectEnumTag:
       return false;
 
     // Handle some rare builtins that may be sensitive to object lifetime

test/IRGen/builtins.swift

@@ -827,5 +827,19 @@ func convertTaskToJob(_ task: Builtin.NativeObject) -> Builtin.Job {
   return Builtin.convertTaskToJob(task)
 }
 
+// CHECK-LABEL: define {{.*}} swiftcc i32 @"$s8builtins10getEnumTagys6UInt32VxlF"(ptr {{.*}} %0, ptr %T)
+// CHECK: %GetEnumTag = load ptr, ptr {{%.*}}
+// CHECK: [[TAG:%.*]] = call i32 %GetEnumTag(ptr {{.*}} %0, ptr %T)
+// CHECK: ret i32 [[TAG]]
+func getEnumTag<T>(_ x: T) -> UInt32 {
+  UInt32(Builtin.getEnumTag(x))
+}
+
+// CHECK-LABEL: define {{.*}} swiftcc void @"$s8builtins13injectEnumTag_3tagyxz_s6UInt32VtlF"(ptr %0, i32 %1, ptr %T)
+// CHECK: %DestructiveInjectEnumTag = load ptr, ptr {{%.*}}
+// CHECK: call void %DestructiveInjectEnumTag(ptr {{.*}} %0, i32 %1, ptr %T)
+func injectEnumTag<T>(_ x: inout T, tag: UInt32) {
+  Builtin.injectEnumTag(&x, tag._value)
+}
 
 // CHECK: ![[R]] = !{i64 0, i64 9223372036854775807}

test/SILGen/builtins.swift

@@ -882,3 +882,22 @@ func assumeAlignment(_ p: Builtin.RawPointer, _ x: Builtin.Word) {
 func packCount<each T>(_ x: repeat each T) -> Builtin.Word {
   Builtin.packLength((repeat each T).self)
 }
+
+// CHECK-LABEL: sil hidden [ossa] @$s8builtins10getEnumTagyBi32_xlF : $@convention(thin) <T> (@in_guaranteed T) -> Builtin.Int32 {
+// CHECK: bb0([[INPUT:%.*]] : $*T):
+// CHECK-NOT: copy_addr
+// CHECK:   [[TAG:%.*]] = builtin "getEnumTag"<T>([[INPUT]] : $*T)
+// CHECK:   return [[TAG]] : $Builtin.Int32
+func getEnumTag<T>(_ x: T) -> Builtin.Int32 {
+  Builtin.getEnumTag(x)
+}
+
+// CHECK-LABEL: sil hidden [ossa] @$s8builtins13injectEnumTag_3tagyxz_Bi32_tlF : $@convention(thin) <T> (@inout T, Builtin.Int32) -> () {
+// CHECK: bb0([[INPUT:%.*]] : $*T, [[TAG:%.*]] : $Builtin.Int32):
+// CHECK-NOT: copy_addr
+// CHECK:   [[ACCESS:%.*]] = begin_access [modify] [unknown] [[INPUT]] : $*T
+// CHECK:   builtin "injectEnumTag"<T>([[ACCESS]] : $*T, [[TAG]] : $Builtin.Int32)
+// CHECK:   end_access [[ACCESS]]
+func injectEnumTag<T>(_ x: inout T, tag: Builtin.Int32) {
+  Builtin.injectEnumTag(&x, tag)
+}