Add value witnesses for single payload enums

include/swift/ABI/ValueWitness.def

@@ -74,6 +74,7 @@
 ///     TYPE_TYPE - a pointer to type metadata
 ///     SIZE_TYPE - size_t
 ///     INT_TYPE - int
+///     UINT_TYPE - unsigned int
 ///     VOID_TYPE - void
 ///   Defaults to VALUE_WITNESS.
 ///   FIXME: The 'copy' witnesses should be using immutable types but aren't.
@@ -179,8 +180,26 @@ FUNCTION_VALUE_WITNESS(initializeBufferWithTakeOfBuffer,
                        MUTABLE_VALUE_TYPE,
                        (MUTABLE_BUFFER_TYPE, MUTABLE_BUFFER_TYPE, TYPE_TYPE))
 
+/// int (*getEnumTagSinglePayload)(const T* enum, UINT_TYPE emptyCases)
+/// Given an instance of valid single payload enum with a payload of this
+/// witness table's type (e.g Optional<ThisType>) , get the tag of the enum.
+FUNCTION_VALUE_WITNESS(getEnumTagSinglePayload,
+                       GetEnumTagSinglePayload,
+                       INT_TYPE,
+                       (IMMUTABLE_VALUE_TYPE, UINT_TYPE, TYPE_TYPE))
+
+/// void (*storeEnumTagSinglePayload)(T* enum, INT_TYPE whichCase,
+///                                   UINT_TYPE emptyCases)
+/// Given uninitialized memory for an instance of a single payload enum with a
+/// payload of this witness table's type (e.g Optional<ThisType>), store the
+/// tag.
+FUNCTION_VALUE_WITNESS(storeEnumTagSinglePayload,
+                       StoreEnumTagSinglePayload,
+                       VOID_TYPE,
+                       (MUTABLE_VALUE_TYPE, INT_TYPE, UINT_TYPE, TYPE_TYPE))
+
 END_VALUE_WITNESS_RANGE(RequiredValueWitnessFunction,
-                        InitializeBufferWithTakeOfBuffer)
+                        StoreEnumTagSinglePayload)
 
 ///   size_t size;
 ///

include/swift/Demangling/ValueWitnessMangling.def

@@ -35,5 +35,7 @@ VALUE_WITNESS(xg, GetExtraInhabitantIndex)
 VALUE_WITNESS(ug, GetEnumTag)
 VALUE_WITNESS(up, DestructiveProjectEnumData)
 VALUE_WITNESS(ui, DestructiveInjectEnumTag)
+VALUE_WITNESS(et, GetEnumTagSinglePayload)
+VALUE_WITNESS(st, StoreEnumTagSinglePayload)
 
 #undef VALUE_WITNESS

include/swift/Runtime/Metadata.h

@@ -357,6 +357,7 @@ namespace value_witness_types {
 #define TYPE_TYPE const Metadata *
 #define SIZE_TYPE size_t
 #define INT_TYPE int
+#define UINT_TYPE unsigned
 #define VOID_TYPE void
 #include "swift/ABI/ValueWitness.def"
 

lib/Demangling/NodePrinter.cpp

@@ -149,6 +149,10 @@ static StringRef toString(ValueWitnessKind k) {
     return "destructiveProjectEnumData";
   case ValueWitnessKind::DestructiveInjectEnumTag:
     return "destructiveInjectEnumTag";
+  case ValueWitnessKind::GetEnumTagSinglePayload:
+    return "getEnumTagSinglePayload";
+  case ValueWitnessKind::StoreEnumTagSinglePayload:
+    return "storeEnumTagSinglePayload";
   }
   printer_unreachable("bad value witness kind");
 }

lib/IRGen/FixedTypeInfo.h

@@ -227,7 +227,16 @@ class FixedTypeInfo : public TypeInfo {
                           llvm::Value *metadata,
                           llvm::Value *vwtable,
                           SILType T) const override {}
-
+
+  llvm::Value *getEnumTagSinglePayload(IRGenFunction &IGF,
+                                       llvm::Value *numEmptyCases,
+                                       Address enumAddr,
+                                       SILType T) const override;
+
+  void storeEnumTagSinglePayload(IRGenFunction &IGF, llvm::Value *whichCase,
+                                 llvm::Value *numEmptyCases, Address enumAddr,
+                                 SILType T) const override;
+
   static bool classof(const FixedTypeInfo *type) { return true; }
   static bool classof(const TypeInfo *type) { return type->isFixedSize(); }
 };

lib/IRGen/GenEnum.cpp

@@ -1591,19 +1591,17 @@ namespace {
     /// Emit a call into the runtime to get the current enum payload tag.
     /// This returns a tag index in the range
     /// [-ElementsWithPayload..ElementsWithNoPayload-1].
-    llvm::Value *
-    emitGetEnumTag(IRGenFunction &IGF, SILType T, Address enumAddr)
-    const override {
-      auto payloadMetadata = emitPayloadMetadataForLayout(IGF, T);
-      auto numEmptyCases = llvm::ConstantInt::get(IGF.IGM.Int32Ty,
-                                                  ElementsWithNoPayload.size());
-
-      auto opaqueAddr = IGF.Builder.CreateBitCast(enumAddr.getAddress(),
-                                                  IGF.IGM.OpaquePtrTy);
+    llvm::Value *emitGetEnumTag(IRGenFunction &IGF, SILType T,
+                                Address enumAddr) const override {
+      auto numEmptyCases =
+          llvm::ConstantInt::get(IGF.IGM.Int32Ty, ElementsWithNoPayload.size());
 
-      return IGF.Builder.CreateCall(
-                 IGF.IGM.getGetEnumCaseSinglePayloadFn(),
-                 {opaqueAddr, payloadMetadata, numEmptyCases});
+      auto PayloadT = getPayloadType(IGF.IGM, T);
+      auto opaqueAddr = Address(
+          IGF.Builder.CreateBitCast(enumAddr.getAddress(), IGF.IGM.OpaquePtrTy),
+          enumAddr.getAlignment());
+      return emitGetEnumTagSinglePayloadCall(IGF, PayloadT, numEmptyCases,
+                                             opaqueAddr);
     }
 
     /// The payload for a single-payload enum is always placed in front and
@@ -2411,16 +2409,16 @@ namespace {
         }
         return;
       }
+      llvm::Value *opaqueAddr =
+          IGF.Builder.CreateBitCast(dest.getAddress(), IGF.IGM.OpaquePtrTy);
 
-      // Ask the runtime to store the tag.
-      llvm::Value *opaqueAddr = IGF.Builder.CreateBitCast(dest.getAddress(),
-                                                          IGF.IGM.OpaquePtrTy);
-      llvm::Value *metadata = emitPayloadMetadataForLayout(IGF, T);
-      IGF.Builder.CreateCall(IGF.IGM.getStoreEnumTagSinglePayloadFn(),
-                             {opaqueAddr, metadata,
-                              llvm::ConstantInt::getSigned(IGF.IGM.Int32Ty, -1),
-                              llvm::ConstantInt::get(IGF.IGM.Int32Ty,
-                                                 ElementsWithNoPayload.size())});
+      auto PayloadT = getPayloadType(IGF.IGM, T);
+      auto Addr = Address(opaqueAddr, dest.getAlignment());
+      auto *whichCase = llvm::ConstantInt::getSigned(IGF.IGM.Int32Ty, -1);
+      auto *numEmptyCases =
+          llvm::ConstantInt::get(IGF.IGM.Int32Ty, ElementsWithNoPayload.size());
+      emitStoreEnumTagSinglePayloadCall(IGF, PayloadT, whichCase, numEmptyCases,
+                                        Addr);
     }
 
     /// Emit a reassignment sequence from an enum at one address to another.
@@ -2625,7 +2623,6 @@ namespace {
                   EnumElementDecl *Case) const override {
       if (TIK < Fixed) {
         // If the enum isn't fixed-layout, get the runtime to do this for us.
-        llvm::Value *payload = emitPayloadMetadataForLayout(IGF, T);
         llvm::Value *caseIndex;
         if (Case == getPayloadElement()) {
           caseIndex = llvm::ConstantInt::getSigned(IGF.IGM.Int32Ty, -1);
@@ -2645,10 +2642,10 @@ namespace {
         llvm::Value *opaqueAddr
           = IGF.Builder.CreateBitCast(enumAddr.getAddress(),
                                       IGF.IGM.OpaquePtrTy);
-
-        IGF.Builder.CreateCall(IGF.IGM.getStoreEnumTagSinglePayloadFn(),
-                               {opaqueAddr, payload, caseIndex, numEmptyCases});
-
+        auto PayloadT = getPayloadType(IGF.IGM, T);
+        auto Addr = Address(opaqueAddr, enumAddr.getAlignment());
+        emitStoreEnumTagSinglePayloadCall(IGF, PayloadT, caseIndex,
+                                          numEmptyCases, Addr);
         return;
       }
 
@@ -2675,20 +2672,18 @@ namespace {
 
     /// Constructs an enum value using a tag index in the range
     /// [-ElementsWithPayload..ElementsWithNoPayload-1].
-    void emitStoreTag(IRGenFunction &IGF,
-                      SILType T,
-                      Address enumAddr,
+    void emitStoreTag(IRGenFunction &IGF, SILType T, Address enumAddr,
                       llvm::Value *tag) const override {
-      llvm::Value *payload = emitPayloadMetadataForLayout(IGF, T);
-      llvm::Value *numEmptyCases = llvm::ConstantInt::get(IGF.IGM.Int32Ty,
-                                                ElementsWithNoPayload.size());
-
+      auto PayloadT = getPayloadType(IGF.IGM, T);
       llvm::Value *opaqueAddr
         = IGF.Builder.CreateBitCast(enumAddr.getAddress(),
-                                    IGF.IGM.OpaquePtrTy);
+                                      IGF.IGM.OpaquePtrTy);
 
-      IGF.Builder.CreateCall(IGF.IGM.getStoreEnumTagSinglePayloadFn(),
-                             {opaqueAddr, payload, tag, numEmptyCases});
+      llvm::Value *numEmptyCases = llvm::ConstantInt::get(IGF.IGM.Int32Ty,
+                                                ElementsWithNoPayload.size());
+      auto Addr = Address(opaqueAddr, enumAddr.getAlignment());
+      emitStoreEnumTagSinglePayloadCall(IGF, PayloadT, tag, numEmptyCases,
+                                        Addr);
     }
 
     void initializeMetadata(IRGenFunction &IGF,

lib/IRGen/GenOpaque.cpp

@@ -23,6 +23,7 @@
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/IR/DerivedTypes.h"
+#include "swift/AST/IRGenOptions.h"
 #include "swift/IRGen/ValueWitness.h"
 
 #include "Callee.h"
@@ -141,13 +142,38 @@ static llvm::Type *createWitnessType(IRGenModule &IGM, ValueWitness index) {
     return llvm::FunctionType::get(voidTy, args, /*isVarArg*/ false);
   }
 
+  /// int (*getEnumTagSinglePayload)(const T* enum, UINT_TYPE emptyCases,
+  ///                                M *self)
+  case ValueWitness::GetEnumTagSinglePayload: {
+    llvm::Type *ptrTy = IGM.OpaquePtrTy;
+    llvm::Type *indexTy = IGM.Int32Ty;
+    llvm::Type *metaTy = IGM.TypeMetadataPtrTy;
+
+    llvm::Type *args[] = { ptrTy, indexTy, metaTy };
+    return llvm::FunctionType::get(indexTy, args, false);
+  }
+
+  /// void (*storeEnumTagSinglePayload)(T* enum, INT_TYPE whichCase,
+  ///                                   UINT_TYPE emptyCases,
+  ///                                   M *self)
+  case ValueWitness::StoreEnumTagSinglePayload: {
+    llvm::Type *voidTy = IGM.VoidTy;
+    llvm::Type *ptrTy = IGM.OpaquePtrTy;
+    llvm::Type *indexTy = IGM.Int32Ty;
+    llvm::Type *metaTy = IGM.TypeMetadataPtrTy;
+
+    llvm::Type *args[] = { ptrTy, indexTy, indexTy, metaTy };
+    return llvm::FunctionType::get(voidTy, args, false);
+  }
+
   case ValueWitness::Size:
   case ValueWitness::Flags:
   case ValueWitness::Stride:
   case ValueWitness::ExtraInhabitantFlags:
     // Non-function witnesses all have type size_t.
     return IGM.SizeTy;
   }
+
   llvm_unreachable("bad value witness!");
 }
 
@@ -173,8 +199,15 @@ static llvm::AttributeList getValueWitnessAttrs(IRGenModule &IGM,
   case ValueWitness::GetEnumTag:
   case ValueWitness::GetExtraInhabitantIndex:
   case ValueWitness::StoreExtraInhabitant:
+  case ValueWitness::StoreEnumTagSinglePayload:
     return attrs.addAttribute(ctx, 1, llvm::Attribute::NoAlias);
 
+  case ValueWitness::GetEnumTagSinglePayload:
+    return attrs
+        .addAttribute(ctx, llvm::AttributeList::FunctionIndex,
+                      llvm::Attribute::ReadOnly)
+        .addAttribute(ctx, 1, llvm::Attribute::NoAlias);
+
   // These have two arguments and they don't alias each other.
   case ValueWitness::AssignWithTake:
   case ValueWitness::InitializeBufferWithCopyOfBuffer:
@@ -245,6 +278,10 @@ static StringRef getValueWitnessLabel(ValueWitness index) {
     return "destructiveProjectEnumData";
   case ValueWitness::DestructiveInjectEnumTag:
     return "destructiveInjectEnumTag";
+  case ValueWitness::GetEnumTagSinglePayload:
+    return "getEnumTagSinglePayload";
+  case ValueWitness::StoreEnumTagSinglePayload:
+    return "storeEnumTagSinglePayload";
   }
   llvm_unreachable("bad value witness index");
 }
@@ -635,6 +672,107 @@ llvm::Value *irgen::emitStoreExtraInhabitantCall(IRGenFunction &IGF,
   return call;
 }
 
+/// Emit a trampoline to call the getEnumTagSinglePayload witness. API:
+/// INT_TYPE (const T* enum, UINT_TYPE emptyCases, M *self)
+static llvm::Constant *
+getGetEnumTagSinglePayloadTrampolineFn(IRGenModule &IGM) {
+
+  llvm::Type *argTys[] = {IGM.OpaquePtrTy, IGM.Int32Ty, IGM.TypeMetadataPtrTy};
+
+  llvm::SmallString<40> fnName("__swift_getEnumTagSinglePayload");
+
+  auto func = IGM.getOrCreateHelperFunction(
+      fnName, IGM.Int32Ty, argTys,
+      [&](IRGenFunction &IGF) {
+        auto it = IGF.CurFn->arg_begin();
+        auto *enumAddr = &*(it++);
+        auto *numEmptyCases = &*(it++);
+        auto *metadata = &*(it++);
+        auto &Builder = IGF.Builder;
+        auto witnessFunc = emitLoadOfValueWitnessFunctionFromMetadata(
+            IGF, metadata, ValueWitness::GetEnumTagSinglePayload);
+        auto *result = Builder.CreateCall(witnessFunc,
+                                          {enumAddr, numEmptyCases, metadata});
+        Builder.CreateRet(result);
+      },
+      true /*noinline*/);
+
+  // This function is readonly.
+  cast<llvm::Function>(func)->addFnAttr(llvm::Attribute::ReadOnly);
+  return func;
+}
+
+/// Emit a trampoline to call the storeEnumTagSinglePayload witness. API:
+/// VOID_TYPE (const T* enum, INT_TYPE whichCase, UINT_TYPE emptyCases,
+///            M *self)
+static llvm::Constant *
+getStoreEnumTagSinglePayloadTrampolineFn(IRGenModule &IGM) {
+
+  llvm::Type *argTys[] = {IGM.OpaquePtrTy, IGM.Int32Ty, IGM.Int32Ty,
+                          IGM.TypeMetadataPtrTy};
+
+  llvm::SmallString<40> fnName("__swift_storeEnumTagSinglePayload");
+
+  return IGM.getOrCreateHelperFunction(
+      fnName, IGM.VoidTy, argTys,
+      [&](IRGenFunction &IGF) {
+        auto it = IGF.CurFn->arg_begin();
+        auto *enumAddr = &*(it++);
+        auto *whichCase = &*(it++);
+        auto *numEmptyCases = &*(it++);
+        auto *metadata = &*(it++);
+        auto &Builder = IGF.Builder;
+        auto witnessFunc = emitLoadOfValueWitnessFunctionFromMetadata(
+            IGF, metadata, ValueWitness::StoreEnumTagSinglePayload);
+        Builder.CreateCall(witnessFunc,
+                           {enumAddr, whichCase, numEmptyCases, metadata});
+        Builder.CreateRetVoid();
+      },
+      true /*noinline*/);
+}
+
+llvm::Value *irgen::emitGetEnumTagSinglePayloadCall(IRGenFunction &IGF,
+                                                    SILType T,
+                                                    llvm::Value *numEmptyCases,
+                                                    Address destObject) {
+  if (!IGF.IGM.getOptions().OptimizeForSize) {
+    llvm::Value *metadata;
+    auto fn = IGF.emitValueWitnessFunctionRef(
+        T, metadata, ValueWitness::GetEnumTagSinglePayload);
+    llvm::CallInst *call = IGF.Builder.CreateCall(
+        fn, {destObject.getAddress(), numEmptyCases, metadata});
+    return call;
+  }
+  auto *metadata = IGF.emitTypeMetadataRefForLayout(T);
+  auto *func = getGetEnumTagSinglePayloadTrampolineFn(IGF.IGM);
+  auto *result = IGF.Builder.CreateCall(
+      func,
+      {IGF.Builder.CreateBitCast(destObject.getAddress(), IGF.IGM.OpaquePtrTy),
+       numEmptyCases, metadata});
+  return result;
+}
+
+llvm::Value *irgen::emitStoreEnumTagSinglePayloadCall(
+    IRGenFunction &IGF, SILType T, llvm::Value *whichCase,
+    llvm::Value *numEmptyCases, Address destObject) {
+  if (!IGF.IGM.getOptions().OptimizeForSize) {
+    llvm::Value *metadata;
+    auto fn = IGF.emitValueWitnessFunctionRef(
+        T, metadata, ValueWitness::StoreEnumTagSinglePayload);
+    llvm::CallInst *call = IGF.Builder.CreateCall(
+        fn, {destObject.getAddress(), whichCase, numEmptyCases, metadata});
+    return call;
+  }
+
+  auto *metadata = IGF.emitTypeMetadataRefForLayout(T);
+  auto *func = getStoreEnumTagSinglePayloadTrampolineFn(IGF.IGM);
+  auto *result = IGF.Builder.CreateCall(
+      func,
+      {IGF.Builder.CreateBitCast(destObject.getAddress(), IGF.IGM.OpaquePtrTy),
+       whichCase, numEmptyCases, metadata});
+  return result;
+}
+
 /// Emit a call to the 'getEnumTag' operation.
 llvm::Value *irgen::emitGetEnumTagCall(IRGenFunction &IGF,
                                        SILType T,

lib/IRGen/GenOpaque.h

@@ -180,6 +180,17 @@ namespace irgen {
                                             llvm::Value *index,
                                             Address destObject);
 
+  /// Emit a call to the 'getEnumTagSinglePayload' operation.
+  llvm::Value *emitGetEnumTagSinglePayloadCall(IRGenFunction &IGF, SILType T,
+                                               llvm::Value *numEmptyCases,
+                                               Address destObject);
+
+  /// Emit a call to the 'storeEnumTagSinglePayload' operation.
+  llvm::Value *emitStoreEnumTagSinglePayloadCall(IRGenFunction &IGF, SILType T,
+                                                 llvm::Value *whichCase,
+                                                 llvm::Value *numEmptyCases,
+                                                 Address destObject);
+
   /// Emit a call to the 'getEnumTag' operation.
   llvm::Value *emitGetEnumTagCall(IRGenFunction &IGF,
                                   SILType T,