Support: Add YAML I/O support for string maps.

[pcc]

Summary

This will be used to YAMLify parts of the module summary.

[mehdi_amini]

I'm not very familiar with the YAML library, but this seems OK to me.

[mehdi_amini]

One nit: you're using StringMap everywhere and it can be confusing since it is not about the llvm::StringMap

[davide]

llvm/include/llvm/Support/YAMLTraits.h

1550–1557 ↗

(On Diff #82193)

nit: I think the variable names should start with a capital letter.

[pcc]

In D28014#629451, @mehdi_amini wrote:

One nit: you're using StringMap everywhere and it can be confusing since it is not about the llvm::StringMap

The name CustomMapping seems a little better, I'll rename and commit.

llvm/include/llvm/Support/YAMLTraits.h

1550–1557 ↗

(On Diff #82193)

I'm just following the style in the rest of this file.

[majnemer]

llvm/trunk/lib/Support/YAMLTraits.cpp

389–393

Isn't this: return is_contained(ValidKeys, Key);?

[mehdi_amini]

llvm/trunk/lib/Support/YAMLTraits.cpp

389–393

Sweet, I didn't know about llvm::is_contained!

[pcc]

llvm/trunk/lib/Support/YAMLTraits.cpp

389–393

Thanks. I changed the code to use is_contained in r290999.

llvm/trunk/include/llvm/Support/YAMLTraits.h

 /// to/from a list of YAML documents.
 template<typename T>
 struct DocumentListTraits {
   // Must provide:
   // static size_t size(IO &io, T &seq);
   // static T::value_type& element(IO &io, T &seq, size_t index);
 };
 
+/// This class should be specialized by any type that needs to be converted
+/// to/from a YAML mapping in the case where the names of the keys are not known
+/// in advance, e.g. a string map.
+template <typename T>
+struct CustomMappingTraits {
+  // static void inputOne(IO &io, StringRef key, T &elem);
+  // static void output(IO &io, T &elem);
+};
+
 // Only used for better diagnostics of missing traits
 template <typename T>
 struct MissingTrait;
 
 // Test if ScalarEnumerationTraits<T> is defined on type T.
 template <class T>
 struct has_ScalarEnumerationTraits
 {
struct has_SequenceMethodTraits
 
   template <typename U>
   static double test(...);
 
 public:
   static bool const value =  (sizeof(test<SequenceTraits<T>>(nullptr)) == 1);
 };
 
+// Test if CustomMappingTraits<T> is defined on type T.
+template <class T>
+struct has_CustomMappingTraits
+{
+  typedef void (*Signature_input)(IO &io, StringRef key, T &v);
+
+  template <typename U>
+  static char test(SameType<Signature_input, &U::inputOne>*);
+
+  template <typename U>
+  static double test(...);
+
+public:
+  static bool const value =
+      (sizeof(test<CustomMappingTraits<T>>(nullptr)) == 1);
+};
+
 // has_FlowTraits<int> will cause an error with some compilers because
 // it subclasses int.  Using this wrapper only instantiates the
 // real has_FlowTraits only if the template type is a class.
 template <typename T, bool Enabled = std::is_class<T>::value>
 class has_FlowTraits
 {
 public:
    static const bool value = false;

 struct missingTraits
     : public std::integral_constant<bool,
                                     !has_ScalarEnumerationTraits<T>::value &&
                                         !has_ScalarBitSetTraits<T>::value &&
                                         !has_ScalarTraits<T>::value &&
                                         !has_BlockScalarTraits<T>::value &&
                                         !has_MappingTraits<T, Context>::value &&
                                         !has_SequenceTraits<T>::value &&
+                                        !has_CustomMappingTraits<T>::value &&
                                         !has_DocumentListTraits<T>::value> {};
 
 template <typename T, typename Context>
 struct validatedMappingTraits
     : public std::integral_constant<
           bool, has_MappingTraits<T, Context>::value &&
                     has_MappingValidateTraits<T, Context>::value> {};
 
public:
   virtual void postflightFlowElement(void*) = 0;
   virtual void endFlowSequence() = 0;
 
   virtual bool mapTag(StringRef Tag, bool Default=false) = 0;
   virtual void beginMapping() = 0;
   virtual void endMapping() = 0;
   virtual bool preflightKey(const char*, bool, bool, bool &, void *&) = 0;
   virtual void postflightKey(void*) = 0;
+  virtual std::vector<StringRef> keys() = 0;
 
   virtual void beginFlowMapping() = 0;
   virtual void endFlowMapping() = 0;
 
   virtual void beginEnumScalar() = 0;
   virtual bool matchEnumScalar(const char*, bool) = 0;
   virtual bool matchEnumFallback() = 0;
   virtual void endEnumScalar() = 0;
yamlize(IO &io, T &Val, bool, Context &Ctx) {
   } else {
     io.beginMapping();
     detail::doMapping(io, Val, Ctx);
     io.endMapping();
   }
 }
 
 template <typename T>
+typename std::enable_if<has_CustomMappingTraits<T>::value, void>::type
+yamlize(IO &io, T &Val, bool, EmptyContext &Ctx) {
+  if ( io.outputting() ) {
+    io.beginMapping();
+    CustomMappingTraits<T>::output(io, Val);
+    io.endMapping();
+  } else {
+    io.beginMapping();
+    for (StringRef key : io.keys())
+      CustomMappingTraits<T>::inputOne(io, key, Val);
+    io.endMapping();
+  }
+}
+
+template <typename T>
 typename std::enable_if<missingTraits<T, EmptyContext>::value, void>::type
 yamlize(IO &io, T &Val, bool, EmptyContext &Ctx) {
   char missing_yaml_trait_for_type[sizeof(MissingTrait<T>)];
 }
 
 template <typename T, typename Context>
 typename std::enable_if<has_SequenceTraits<T>::value, void>::type
 yamlize(IO &io, T &Seq, bool, Context &Ctx) {

 
 private:
   bool outputting() override;
   bool mapTag(StringRef, bool) override;
   void beginMapping() override;
   void endMapping() override;
   bool preflightKey(const char *, bool, bool, bool &, void *&) override;
   void postflightKey(void *) override;
+  std::vector<StringRef> keys() override;
   void beginFlowMapping() override;
   void endFlowMapping() override;
   unsigned beginSequence() override;
   void endSequence() override;
   bool preflightElement(unsigned index, void *&) override;
   void postflightElement(void *) override;
   unsigned beginFlowSequence() override;
   bool preflightFlowElement(unsigned , void *&) override;
public:
 
     static inline bool classof(const MapHNode *) { return true; }
 
     typedef llvm::StringMap<std::unique_ptr<HNode>> NameToNode;
 
     bool isValidKey(StringRef key);
 
     NameToNode                        Mapping;
-    llvm::SmallVector<const char*, 6> ValidKeys;
+    llvm::SmallVector<std::string, 6> ValidKeys;
   };
 
   class SequenceHNode : public HNode {
     void anchor() override;
 
   public:
     SequenceHNode(Node *n) : HNode(n) { }
 
public:
   ~Output() override;
 
   bool outputting() override;
   bool mapTag(StringRef, bool) override;
   void beginMapping() override;
   void endMapping() override;
   bool preflightKey(const char *key, bool, bool, bool &, void *&) override;
   void postflightKey(void *) override;
+  std::vector<StringRef> keys() override;
   void beginFlowMapping() override;
   void endFlowMapping() override;
   unsigned beginSequence() override;
   void endSequence() override;
   bool preflightElement(unsigned, void *&) override;
   void postflightElement(void *) override;
   unsigned beginFlowSequence() override;
   bool preflightFlowElement(unsigned, void *&) override;

 typename std::enable_if<has_BlockScalarTraits<T>::value, Input &>::type
 operator>>(Input &In, T &Val) {
   EmptyContext Ctx;
   if (In.setCurrentDocument())
     yamlize(In, Val, true, Ctx);
   return In;
 }
 
+// Define non-member operator>> so that Input can stream in a string map.
+template <typename T>
+inline
+typename std::enable_if<has_CustomMappingTraits<T>::value, Input &>::type
+operator>>(Input &In, T &Val) {
+  EmptyContext Ctx;
+  if (In.setCurrentDocument())
+    yamlize(In, Val, true, Ctx);
+  return In;
+}
+
 // Provide better error message about types missing a trait specialization
 template <typename T>
 inline typename std::enable_if<missingTraits<T, EmptyContext>::value,
                                Input &>::type
 operator>>(Input &yin, T &docSeq) {
   char missing_yaml_trait_for_type[sizeof(MissingTrait<T>)];
   return yin;
 }
operator<<(Output &Out, T &Val) {
   if (Out.preflightDocument(0)) {
     yamlize(Out, Val, true, Ctx);
     Out.postflightDocument();
   }
   Out.endDocuments();
   return Out;
 }
 
+// Define non-member operator<< so that Output can stream out a string map.
+template <typename T>
+inline
+typename std::enable_if<has_CustomMappingTraits<T>::value, Output &>::type
+operator<<(Output &Out, T &Val) {
+  EmptyContext Ctx;
+  Out.beginDocuments();
+  if (Out.preflightDocument(0)) {
+    yamlize(Out, Val, true, Ctx);
+    Out.postflightDocument();
+  }
+  Out.endDocuments();
+  return Out;
+}
+
 // Provide better error message about types missing a trait specialization
 template <typename T>
 inline typename std::enable_if<missingTraits<T, EmptyContext>::value,
                                Output &>::type
 operator<<(Output &yout, T &seq) {
   char missing_yaml_trait_for_type[sizeof(MissingTrait<T>)];
   return yout;
 }
 
 template <typename T> struct SequenceTraitsImpl {
   typedef typename T::value_type _type;
   static size_t size(IO &io, T &seq) { return seq.size(); }
   static _type &element(IO &io, T &seq, size_t index) {
     if (index >= seq.size())
       seq.resize(index + 1);
     return seq[index];
   }
 };
 
+/// Implementation of CustomMappingTraits for std::map<std::string, T>.
+template <typename T> struct StdMapStringCustomMappingTraitsImpl {
+  typedef std::map<std::string, T> map_type;
+  static void inputOne(IO &io, StringRef key, map_type &v) {
+    io.mapRequired(key.str().c_str(), v[key]);
+  }
+  static void output(IO &io, map_type &v) {
+    for (auto &p : v)
+      io.mapRequired(p.first.c_str(), p.second);
+  }
+};
+
 } // end namespace yaml
 } // end namespace llvm
 
 /// Utility for declaring that a std::vector of a particular type
 /// should be considered a YAML sequence.
 #define LLVM_YAML_IS_SEQUENCE_VECTOR(_type)                                    \
   namespace llvm {                                                             \
   namespace yaml {                                                             \
#define LLVM_YAML_IS_DOCUMENT_LIST_VECTOR(_type)                               \
   struct DocumentListTraits<SmallVector<_type, N>>                             \
       : public SequenceTraitsImpl<SmallVector<_type, N>> {};                   \
   template <>                                                                  \
   struct DocumentListTraits<std::vector<_type>>                                \
       : public SequenceTraitsImpl<std::vector<_type>> {};                      \
   }                                                                            \
   }
 
+/// Utility for declaring that std::map<std::string, _type> should be considered
+/// a YAML map.
+#define LLVM_YAML_IS_STRING_MAP(_type)                                         \
+  namespace llvm {                                                             \
+  namespace yaml {                                                             \
+  template <>                                                                  \
+  struct CustomMappingTraits<std::map<std::string, _type>>                     \
+      : public StdMapStringCustomMappingTraitsImpl<_type> {};                  \
+  }                                                                            \
+  }
+
 #endif // LLVM_SUPPORT_YAMLTRAITS_H

llvm/trunk/lib/Support/YAMLTraits.cpp

if (EC)
     return;
   // CurrentNode can be null if the document is empty.
   MapHNode *MN = dyn_cast_or_null<MapHNode>(CurrentNode);
   if (MN) {
     MN->ValidKeys.clear();
   }
 }
 
+std::vector<StringRef> Input::keys() {
+  MapHNode *MN = dyn_cast<MapHNode>(CurrentNode);
+  std::vector<StringRef> Ret;
+  if (!MN) {
+    setError(CurrentNode, "not a mapping");
+    return Ret;
+  }
+  for (auto &P : MN->Mapping)
+    Ret.push_back(P.first());
+  return Ret;
+}
+
 bool Input::preflightKey(const char *Key, bool Required, bool, bool &UseDefault,
                          void *&SaveInfo) {
   UseDefault = false;
   if (EC)
     return false;
 
   // CurrentNode is null for empty documents, which is an error in case required
   // nodes are present.
if (ScalarNode *SN = dyn_cast<ScalarNode>(N)) {
     return llvm::make_unique<EmptyHNode>(N);
   } else {
     setError(N, "unknown node kind");
     return nullptr;
   }
 }
 
 bool Input::MapHNode::isValidKey(StringRef Key) {
-  for (const char *K : ValidKeys) {
-    if (Key.equals(K))
+  for (std::string &K : ValidKeys) {
+    if (Key == K)
       return true;
   }
   return false;
 }
 
 void Input::setError(const Twine &Message) {
   this->setError(CurrentNode, Message);
 }
 
 bool Input::canElideEmptySequence() {
   return false;
bool Output::mapTag(StringRef Tag, bool Use) {
   }
   return Use;
 }
 
 void Output::endMapping() {
   StateStack.pop_back();
 }
 
+std::vector<StringRef> Output::keys() {
+  report_fatal_error("invalid call");
+}
+
 bool Output::preflightKey(const char *Key, bool Required, bool SameAsDefault,
                           bool &UseDefault, void *&) {
   UseDefault = false;
   if (Required || !SameAsDefault) {
     auto State = StateStack.back();
     if (State == inFlowMapFirstKey || State == inFlowMapOtherKey) {
       flowKey(Key);
     } else {

llvm/trunk/unittests/Support/YAMLIOTest.cpp

EXPECT_EQ("---\n"
             "  B:               2\n"
             "  C:               3\n"
             "  Context:         2\n"
             "...\n",
             out);
   out.clear();
 }
 
+LLVM_YAML_IS_STRING_MAP(int)
+
+TEST(YAMLIO, TestCustomMapping) {
+  std::map<std::string, int> x;
+  x["foo"] = 1;
+  x["bar"] = 2;
+
+  std::string out;
+  llvm::raw_string_ostream ostr(out);
+  Output xout(ostr, nullptr, 0);
+
+  xout << x;
+  ostr.flush();
+  EXPECT_EQ("---\n"
+            "bar:             2\n"
+            "foo:             1\n"
+            "...\n",
+            out);
+
+  Input yin(out);
+  std::map<std::string, int> y;
+  yin >> y;
+  EXPECT_EQ(2ul, y.size());
+  EXPECT_EQ(1, y["foo"]);
+  EXPECT_EQ(2, y["bar"]);
+}
+
+LLVM_YAML_IS_STRING_MAP(FooBar)
+
+TEST(YAMLIO, TestCustomMappingStruct) {
+  std::map<std::string, FooBar> x;
+  x["foo"].foo = 1;
+  x["foo"].bar = 2;
+  x["bar"].foo = 3;
+  x["bar"].bar = 4;
+
+  std::string out;
+  llvm::raw_string_ostream ostr(out);
+  Output xout(ostr, nullptr, 0);
+
+  xout << x;
+  ostr.flush();
+  EXPECT_EQ("---\n"
+            "bar:             \n"
+            "  foo:             3\n"
+            "  bar:             4\n"
+            "foo:             \n"
+            "  foo:             1\n"
+            "  bar:             2\n"
+            "...\n",
+            out);
+
+  Input yin(out);
+  std::map<std::string, FooBar> y;
+  yin >> y;
+  EXPECT_EQ(2ul, y.size());
+  EXPECT_EQ(1, y["foo"].foo);
+  EXPECT_EQ(2, y["foo"].bar);
+  EXPECT_EQ(3, y["bar"].foo);
+  EXPECT_EQ(4, y["bar"].bar);
+}
+
 TEST(YAMLIO, InvalidInput) {
   // polluting 1 value in the sequence
   Input yin("---\n- foo:  3\n  bar:  5\n1\n- foo:  3\n  bar:  5\n...\n");
   std::vector<FooBar> Data;
   yin >> Data;
   EXPECT_TRUE((bool)yin.error());
 }