YAMLIO: Improve endian type support

Summary:
Now that endian types support enumerations (D59141), the existing yaml
support for them is somewhat insufficient. The current solution was to
define the ScalarTraits class for these types, which always forwards to
the ScalarTraits of the underlying type. However, the enum types will
usually have ScalarEnumerationTraits of ScalarBitsetTraits.

In this patch I add the two extra Traits types to the endian types. In
order to properly SFINAE-ize them, I've also added an extra "Enable"
template argument to the Traits template classes.

Reviewers: zturner, sammccall

Subscribers: kristina, Bigcheese, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D59289

llvm-svn: 356269

2 files changed, 128 insertions, 12 deletions

diff --git a/llvm/include/llvm/Support/YAMLTraits.h b/llvm/include/llvm/Support/YAMLTraits.h
index b80c7426d9b..ec508ecb14f 100644
--- a/llvm/include/llvm/Support/YAMLTraits.h
+++ b/llvm/include/llvm/Support/YAMLTraits.h
@@ -101,8 +101,7 @@ template <class T, class Context> struct MappingContextTraits {
 ///           io.enumCase(value, "green", cGreen);
 ///         }
 ///       };
-template<typename T>
-struct ScalarEnumerationTraits {
+template <typename T, typename Enable = void> struct ScalarEnumerationTraits {
   // Must provide:
   // static void enumeration(IO &io, T &value);
 };
@@ -118,8 +117,7 @@ struct ScalarEnumerationTraits {
 ///          io.bitSetCase(value, "round", flagRound);
 ///        }
 ///      };
-template<typename T>
-struct ScalarBitSetTraits {
+template <typename T, typename Enable = void> struct ScalarBitSetTraits {
   // Must provide:
   // static void bitset(IO &io, T &value);
 };
@@ -145,8 +143,7 @@ enum class QuotingType { None, Single, Double };
 ///      }
 ///      static QuotingType mustQuote(StringRef) { return QuotingType::Single; }
 ///    };
-template<typename T>
-struct ScalarTraits {
+template <typename T, typename Enable = void> struct ScalarTraits {
   // Must provide:
   //
   // Function to write the value as a string:
@@ -980,7 +977,7 @@ yamlize(IO &io, T &Val, bool, EmptyContext &Ctx) {
   bool DoClear;
   if ( io.beginBitSetScalar(DoClear) ) {
     if ( DoClear )
-      Val = static_cast<T>(0);
+      Val = T();
     ScalarBitSetTraits<T>::bitset(io, Val);
     io.endBitSetScalar();
   }
@@ -1245,12 +1242,14 @@ struct ScalarTraits<double> {
   static QuotingType mustQuote(StringRef) { return QuotingType::None; }
 };
 
-// For endian types, we just use the existing ScalarTraits for the underlying
-// type.  This way endian aware types are supported whenever a ScalarTraits
-// is defined for the underlying type.
+// For endian types, we use existing scalar Traits class for the underlying
+// type.  This way endian aware types are supported whenever the traits are
+// defined for the underlying type.
 template <typename value_type, support::endianness endian, size_t alignment>
-struct ScalarTraits<support::detail::packed_endian_specific_integral<
-    value_type, endian, alignment>> {
+struct ScalarTraits<
+    support::detail::packed_endian_specific_integral<value_type, endian,
+                                                     alignment>,
+    typename std::enable_if<has_ScalarTraits<value_type>::value>::type> {
   using endian_type =
       support::detail::packed_endian_specific_integral<value_type, endian,
                                                        alignment>;
@@ -1271,6 +1270,38 @@ struct ScalarTraits<support::detail::packed_endian_specific_integral<
   }
 };
 
+template <typename value_type, support::endianness endian, size_t alignment>
+struct ScalarEnumerationTraits<
+    support::detail::packed_endian_specific_integral<value_type, endian,
+                                                     alignment>,
+    typename std::enable_if<
+        has_ScalarEnumerationTraits<value_type>::value>::type> {
+  using endian_type =
+      support::detail::packed_endian_specific_integral<value_type, endian,
+                                                       alignment>;
+
+  static void enumeration(IO &io, endian_type &E) {
+    value_type V = E;
+    ScalarEnumerationTraits<value_type>::enumeration(io, V);
+    E = V;
+  }
+};
+
+template <typename value_type, support::endianness endian, size_t alignment>
+struct ScalarBitSetTraits<
+    support::detail::packed_endian_specific_integral<value_type, endian,
+                                                     alignment>,
+    typename std::enable_if<has_ScalarBitSetTraits<value_type>::value>::type> {
+  using endian_type =
+      support::detail::packed_endian_specific_integral<value_type, endian,
+                                                       alignment>;
+  static void bitset(IO &io, endian_type &E) {
+    value_type V = E;
+    ScalarBitSetTraits<value_type>::bitset(io, V);
+    E = V;
+  }
+};
+
 // Utility for use within MappingTraits<>::mapping() method
 // to [de]normalize an object for use with YAML conversion.
 template <typename TNorm, typename TFinal>
diff --git a/llvm/unittests/Support/YAMLIOTest.cpp b/llvm/unittests/Support/YAMLIOTest.cpp
index 941a765a3c1..26858015bcb 100644
--- a/llvm/unittests/Support/YAMLIOTest.cpp
+++ b/llvm/unittests/Support/YAMLIOTest.cpp
@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/BitmaskEnum.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
@@ -579,6 +580,90 @@ TEST(YAMLIO, TestReadWriteEndianTypes) {
   }
 }
 
+enum class Enum : uint16_t { One, Two };
+enum class BitsetEnum : uint16_t {
+  ZeroOne = 0x01,
+  OneZero = 0x10,
+  LLVM_MARK_AS_BITMASK_ENUM(/*LargestValue*/ OneZero),
+};
+LLVM_ENABLE_BITMASK_ENUMS_IN_NAMESPACE();
+struct EndianEnums {
+  llvm::support::little_t<Enum> LittleEnum;
+  llvm::support::big_t<Enum> BigEnum;
+  llvm::support::little_t<BitsetEnum> LittleBitset;
+  llvm::support::big_t<BitsetEnum> BigBitset;
+};
+namespace llvm {
+namespace yaml {
+template <> struct ScalarEnumerationTraits<Enum> {
+  static void enumeration(IO &io, Enum &E) {
+    io.enumCase(E, "One", Enum::One);
+    io.enumCase(E, "Two", Enum::Two);
+  }
+};
+
+template <> struct ScalarBitSetTraits<BitsetEnum> {
+  static void bitset(IO &io, BitsetEnum &E) {
+    io.bitSetCase(E, "ZeroOne", BitsetEnum::ZeroOne);
+    io.bitSetCase(E, "OneZero", BitsetEnum::OneZero);
+  }
+};
+
+template <> struct MappingTraits<EndianEnums> {
+  static void mapping(IO &io, EndianEnums &EE) {
+    io.mapRequired("LittleEnum", EE.LittleEnum);
+    io.mapRequired("BigEnum", EE.BigEnum);
+    io.mapRequired("LittleBitset", EE.LittleBitset);
+    io.mapRequired("BigBitset", EE.BigBitset);
+  }
+};
+} // namespace yaml
+} // namespace llvm
+
+TEST(YAMLIO, TestReadEndianEnums) {
+  EndianEnums map;
+  Input yin("---\n"
+            "LittleEnum:   One\n"
+            "BigEnum:      Two\n"
+            "LittleBitset: [ ZeroOne ]\n"
+            "BigBitset:    [ ZeroOne, OneZero ]\n"
+            "...\n");
+  yin >> map;
+
+  EXPECT_FALSE(yin.error());
+  EXPECT_EQ(Enum::One, map.LittleEnum);
+  EXPECT_EQ(Enum::Two, map.BigEnum);
+  EXPECT_EQ(BitsetEnum::ZeroOne, map.LittleBitset);
+  EXPECT_EQ(BitsetEnum::ZeroOne | BitsetEnum::OneZero, map.BigBitset);
+}
+
+TEST(YAMLIO, TestReadWriteEndianEnums) {
+  std::string intermediate;
+  {
+    EndianEnums map;
+    map.LittleEnum = Enum::Two;
+    map.BigEnum = Enum::One;
+    map.LittleBitset = BitsetEnum::OneZero | BitsetEnum::ZeroOne;
+    map.BigBitset = BitsetEnum::OneZero;
+
+    llvm::raw_string_ostream ostr(intermediate);
+    Output yout(ostr);
+    yout << map;
+  }
+
+  {
+    Input yin(intermediate);
+    EndianEnums map;
+    yin >> map;
+
+    EXPECT_FALSE(yin.error());
+    EXPECT_EQ(Enum::Two, map.LittleEnum);
+    EXPECT_EQ(Enum::One, map.BigEnum);
+    EXPECT_EQ(BitsetEnum::OneZero | BitsetEnum::ZeroOne, map.LittleBitset);
+    EXPECT_EQ(BitsetEnum::OneZero, map.BigBitset);
+  }
+}
+
 struct StringTypes {
   llvm::StringRef str1;
   llvm::StringRef str2;