[Interop][SwiftToCxx] Implement enum case switching

include/swift/IRGen/IRABIDetailsProvider.h

@@ -13,8 +13,10 @@
 #ifndef SWIFT_IRGEN_IRABIDETAILSPROVIDER_H
 #define SWIFT_IRGEN_IRABIDETAILSPROVIDER_H
 
+#include "swift/AST/Decl.h"
 #include "swift/AST/Type.h"
 #include "clang/AST/CharUnits.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
@@ -93,6 +95,10 @@ class IRABIDetailsProvider {
   /// access function.
   FunctionABISignature getTypeMetadataAccessFunctionSignature();
 
+  /// Returns EnumElementDecls (enum cases) in their declaration order with
+  /// their tag indices from the given EnumDecl
+  llvm::MapVector<EnumElementDecl *, unsigned> getEnumTagMapping(EnumDecl *ED);
+
 private:
   std::unique_ptr<IRABIDetailsProviderImpl> impl;
 };

lib/IRGen/IRABIDetailsProvider.cpp

@@ -12,6 +12,7 @@
 
 #include "swift/IRGen/IRABIDetailsProvider.h"
 #include "FixedTypeInfo.h"
+#include "GenEnum.h"
 #include "GenType.h"
 #include "IRGen.h"
 #include "IRGenModule.h"
@@ -122,6 +123,19 @@ class IRABIDetailsProviderImpl {
     return {returnTy, {paramTy}};
   }
 
+  llvm::MapVector<EnumElementDecl *, unsigned> getEnumTagMapping(EnumDecl *ED) {
+    llvm::MapVector<EnumElementDecl *, unsigned> elements;
+    auto &enumImplStrat = getEnumImplStrategy(
+        IGM, ED->getDeclaredType()->getCanonicalType());
+
+    for (auto *element : ED->getAllElements()) {
+      auto tagIdx = enumImplStrat.getTagIndex(element);
+      elements.insert({element, tagIdx});
+    }
+
+    return elements;
+  }
+
 private:
   Lowering::TypeConverter typeConverter;
   // Default silOptions are sufficient, as we don't need to generated SIL.
@@ -162,3 +176,8 @@ IRABIDetailsProvider::FunctionABISignature
 IRABIDetailsProvider::getTypeMetadataAccessFunctionSignature() {
   return impl->getTypeMetadataAccessFunctionSignature();
 }
+
+llvm::MapVector<EnumElementDecl *, unsigned>
+IRABIDetailsProvider::getEnumTagMapping(EnumDecl *ED) {
+  return impl->getEnumTagMapping(ED);
+}

lib/PrintAsClang/DeclAndTypePrinter.cpp

@@ -15,6 +15,7 @@
 #include "PrimitiveTypeMapping.h"
 #include "PrintClangFunction.h"
 #include "PrintClangValueType.h"
+#include "SwiftToClangInteropContext.h"
 
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ASTMangler.h"
@@ -31,6 +32,7 @@
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/TypeVisitor.h"
 #include "swift/IDE/CommentConversion.h"
+#include "swift/IRGen/IRABIDetailsProvider.h"
 #include "swift/IRGen/Linking.h"
 #include "swift/Parse/Lexer.h"
 #include "swift/Parse/Parser.h"
@@ -384,30 +386,67 @@ class DeclAndTypePrinter::Implementation
     os << "@end\n";
   }
 
+  void visitEnumDeclCxx(EnumDecl *ED) {
+    // FIXME: Print enum's availability
+    ClangValueTypePrinter printer(os, owningPrinter.prologueOS,
+                                  owningPrinter.typeMapping,
+                                  owningPrinter.interopContext);
+    printer.printValueTypeDecl(ED, /*bodyPrinter=*/[&]() {
+      ClangSyntaxPrinter syntaxPrinter(os);
+      auto elementTagMapping =
+          owningPrinter.interopContext.getIrABIDetails().getEnumTagMapping(ED);
+      // Sort cases based on their assigned tag indices
+      llvm::stable_sort(elementTagMapping, [](const auto &p1, const auto &p2) {
+        return p1.second < p2.second;
+      });
+
+      if (elementTagMapping.empty()) {
+        os << "\n";
+        return;
+      }
+
+      os << "  enum class cases {\n";
+      for (const auto &pair : elementTagMapping) {
+        os << "    ";
+        syntaxPrinter.printIdentifier(pair.first->getNameStr());
+        os << ",\n";
+      }
+      os << "  };\n"; // enum class cases' closing bracket
+
+      // Printing operator cases()
+      os << "  inline operator cases() const {\n";
+      os << "    switch (_getEnumTag()) {\n";
+      for (const auto &pair : elementTagMapping) {
+        os << "      case " << pair.second << ": return cases::";
+        syntaxPrinter.printIdentifier(pair.first->getNameStr());
+        os << ";\n";
+      }
+      // TODO: change to Swift's fatalError when it's available in C++
+      os << "      default: abort();\n";
+      os << "    }\n"; // switch's closing bracket
+      os << "  }\n";   // operator cases()'s closing bracket
+
+      // Printing predicates
+      for (const auto &pair : elementTagMapping) {
+        os << "  inline bool is";
+        auto name = pair.first->getNameStr().str();
+        name[0] = std::toupper(name[0]);
+        os << name << "() const {\n";
+        os << "    return *this == cases::";
+        syntaxPrinter.printIdentifier(pair.first->getNameStr());
+        os << ";\n  }\n";
+      }
+      os << "\n";
+    });
+    os << outOfLineDefinitions;
+    outOfLineDefinitions.clear();
+  }
+
   void visitEnumDecl(EnumDecl *ED) {
     printDocumentationComment(ED);
 
     if (outputLang == OutputLanguageMode::Cxx) {
-      // FIXME: Print enum's availability
-      ClangValueTypePrinter printer(os, owningPrinter.prologueOS,
-                                    owningPrinter.typeMapping,
-                                    owningPrinter.interopContext);
-      printer.printValueTypeDecl(ED, /*bodyPrinter=*/[&]() {
-        ClangSyntaxPrinter syntaxPrinter(os);
-        os << "  enum class cases {";
-        llvm::interleaveComma(
-            ED->getAllCases(), os, [&](const EnumCaseDecl *caseDecl) {
-              llvm::interleaveComma(caseDecl->getElements(), os,
-                                    [&](const EnumElementDecl *elementDecl) {
-                                      os << "\n    ";
-                                      syntaxPrinter.printIdentifier(
-                                          elementDecl->getNameStr());
-                                    });
-            });
-        os << "\n  };\n";
-      });
-      os << outOfLineDefinitions;
-      outOfLineDefinitions.clear();
+      visitEnumDeclCxx(ED);
       return;
     }
 

lib/PrintAsClang/PrintClangValueType.cpp

@@ -202,7 +202,25 @@ void ClangValueTypePrinter::printValueTypeDecl(
     os << ".getOpaquePointer()";
   os << "; }\n";
   os << "\n";
-
+  // Print out helper function for getting enum tag for enum type
+  if (isa<EnumDecl>(typeDecl)) {
+    // FIXME: (tongjie) return type should be unsigned
+    os << "  inline int _getEnumTag() const {\n";
+    os << "    auto metadata = " << cxx_synthesis::getCxxImplNamespaceName()
+       << "::";
+    printer.printSwiftTypeMetadataAccessFunctionCall(typeMetadataFuncName);
+    os << ";\n";
+    os << "    auto *vwTable = ";
+    printer.printValueWitnessTableAccessFromTypeMetadata("metadata");
+    os << ";\n";
+    os << "    const auto *enumVWTable = reinterpret_cast<";
+    ClangSyntaxPrinter(os).printSwiftImplQualifier();
+    os << "EnumValueWitnessTable";
+    os << " *>(vwTable);\n";
+    os << "    return enumVWTable->getEnumTag(_getOpaquePointer(), "
+          "metadata._0);\n";
+    os << "  }\n";
+  }
   // Print out the storage for the value type.
   os << "  ";
   if (isOpaqueLayout) {

lib/PrintAsClang/PrintSwiftToClangCoreScaffold.cpp

@@ -62,11 +62,12 @@ static void printKnownType(
   printKnownStruct(typeMapping, os, name, typeRecord);
 }
 
-static void printValueWitnessTableFunctionType(raw_ostream &os, StringRef name,
-                                               StringRef returnType,
-                                               std::string paramTypes,
-                                               uint16_t ptrauthDisc) {
-  os << "using ValueWitness" << name << "Ty = " << returnType
+static void
+printValueWitnessTableFunctionType(raw_ostream &os, StringRef prefix,
+                                   StringRef name, StringRef returnType,
+                                   std::string paramTypes,
+                                   uint16_t ptrauthDisc) {
+  os << "using " << prefix << name << "Ty = " << returnType
      << "(* __ptrauth_swift_value_witness_function_pointer(" << ptrauthDisc
      << "))(" << paramTypes << ");\n";
 }
@@ -87,7 +88,7 @@ static void printValueWitnessTable(raw_ostream &os) {
   membersOS << "  " << type << " " << #lowerId << ";\n";
 #define FUNCTION_VALUE_WITNESS(lowerId, upperId, returnType, paramTypes)       \
   printValueWitnessTableFunctionType(                                          \
-      os, #upperId, returnType, makeParams paramTypes,                         \
+      os, "ValueWitness", #upperId, returnType, makeParams paramTypes,         \
       SpecialPointerAuthDiscriminators::upperId);                              \
   membersOS << "  ValueWitness" << #upperId << "Ty _Nonnull " << #lowerId      \
             << ";\n";
@@ -102,7 +103,32 @@ static void printValueWitnessTable(raw_ostream &os) {
 #define VOID_TYPE "void"
 #include "swift/ABI/ValueWitness.def"
 
-  os << "\nstruct ValueWitnessTable {\n" << membersOS.str() << "};\n";
+  os << "\nstruct ValueWitnessTable {\n" << membersOS.str() << "};\n\n";
+  membersOS.str().clear();
+
+#define WANT_ONLY_ENUM_VALUE_WITNESSES
+#define DATA_VALUE_WITNESS(lowerId, upperId, type)                             \
+  membersOS << "  " << type << " " << #lowerId << ";\n";
+#define FUNCTION_VALUE_WITNESS(lowerId, upperId, returnType, paramTypes)       \
+  printValueWitnessTableFunctionType(                                          \
+      os, "EnumValueWitness", #upperId, returnType, makeParams paramTypes,     \
+      SpecialPointerAuthDiscriminators::upperId);                              \
+  membersOS << "  EnumValueWitness" << #upperId << "Ty _Nonnull " << #lowerId  \
+            << ";\n";
+#define MUTABLE_VALUE_TYPE "void * _Nonnull"
+#define IMMUTABLE_VALUE_TYPE "const void * _Nonnull"
+#define MUTABLE_BUFFER_TYPE "void * _Nonnull"
+#define IMMUTABLE_BUFFER_TYPE "const void * _Nonnull"
+#define TYPE_TYPE "void * _Nonnull"
+#define SIZE_TYPE "size_t"
+#define INT_TYPE "int"
+#define UINT_TYPE "unsigned"
+#define VOID_TYPE "void"
+#include "swift/ABI/ValueWitness.def"
+
+  os << "\nstruct EnumValueWitnessTable {\n"
+     << "  ValueWitnessTable vwTable;\n"
+     << membersOS.str() << "};\n\n";
 }
 
 static void printTypeMetadataResponseType(SwiftToClangInteropContext &ctx,

test/Inputs/clang-importer-sdk/usr/include/stdlib.h

@@ -16,4 +16,6 @@ void free(void *);
 ldiv_t ldiv(long int, long int);
 lldiv_t lldiv(long long int, long long int);
 
+_Noreturn void abort(void);
+
 #endif // SDK_STDLIB_H

test/Interop/SwiftToCxx/core/swift-impl-defs-in-cxx.swift

@@ -47,6 +47,18 @@
 // CHECK-NEXT:  unsigned extraInhabitantCount;
 // CHECK-NEXT: };
 // CHECK-EMPTY:
+// CHECK-NEXT: using EnumValueWitnessGetEnumTagTy = int(* __ptrauth_swift_value_witness_function_pointer(41909))(const void * _Nonnull, void * _Nonnull);
+// CHECK-NEXT: using EnumValueWitnessDestructiveProjectEnumDataTy = void(* __ptrauth_swift_value_witness_function_pointer(1053))(void * _Nonnull, void * _Nonnull);
+// CHECK-NEXT: using EnumValueWitnessDestructiveInjectEnumTagTy = void(* __ptrauth_swift_value_witness_function_pointer(45796))(void * _Nonnull, unsigned, void * _Nonnull);
+// CHECK-EMPTY:
+// CHECK-NEXT: struct EnumValueWitnessTable {
+// CHECK-NEXT:   ValueWitnessTable vwTable;
+// CHECK-NEXT:   EnumValueWitnessGetEnumTagTy _Nonnull getEnumTag;
+// CHECK-NEXT:   EnumValueWitnessDestructiveProjectEnumDataTy _Nonnull destructiveProjectEnumData;
+// CHECK-NEXT:   EnumValueWitnessDestructiveInjectEnumTagTy _Nonnull destructiveInjectEnumTag;
+// CHECK-NEXT: };
+// CHECK-EMPTY:
+// CHECK-EMPTY:
 // CHECK-NEXT: #ifdef __cplusplus
 // CHECK-NEXT: }
 // CHECK-NEXT: #endif