[yaml] Add the ability to pass in context at the map level

[zturner]

Summary

When doing mapRequired or mapOptional it is useful to be able to pass context to the operation. This has been partially addressed by having a void *context in the IO instance, but this is inflexible when you context is constructed on the fly, for example, and needs to be passed to only one specific mapping.

The solution implemented here adds an additional overload of mapRequired and mapOptional, taking a parameterized type. If you invoke this overload, you are expected to specialize a MappingTraits with not only the default implementation which takes a n IO& and a T &, but an additional overload that takes an IO&, T &, and Context &. This object is then passed into the map operation. The original no context overload is still required

[zturner]

It occurred to me that a limitation of my original patch is that it would have been impossible to define a mapping function that accepts a context if a specialization of MappingTraits had already been defined for T out of your control, since they all reused the same exact MappingTraits<T> specialization. To fix this mapping with context now requires an entirely different class. Instead of reusing MappingTraits<T>::mapping() it relies on MappingContextTraits<T, Context>::mapping(). Now one can have both a MappingContextTraits<T, Context> as well as a MappingTraits<T>, and if the non-context mapping functionality is identical, one could even define the context mapping traits like this:

template<>
struct MappingContextTraits<MyClass, MyContext> : public MappingTraits<MyClass> {
  static void mapping(IO &io, MyClass &Obj, MyContext &Context) {
    MappingTraits<MyClass>::mapping(io, Obj);
    // Do something with the context.
  }
};

[zturner]

Not sure why I added lhames, must have gotten confused. + people who are hopefully more appropriate.

[zturner]

Would someone mind taking a look at this? I've got 2-3 patches which can't go in until this one does.

[zturner]

Chandler,

AFAICT you are the only listed owner of Support. Would you mind reviewing this (or finding someone else to review it)?

[chandlerc]

I guess I can look. I was hoping someone with more familiarity with the code would, but unless Alex shows back up, that seems unlikely. And I understand you need to get unblocked.

I find the design pattern here pretty deeply confusing. But I don't think that it is at all related to your patch, but to my lack of familiarity with YAMLIO overall. Your patch does seem to fit the existing design patterns.

Since you might be able to, have you checked if MSVC 2013 is happy with the SFINAE trait you're using? I've had trouble getting those past it in the past.

Also a minor comment on the test below.

unittests/Support/YAMLIOTest.cpp

2305–2317

Shouldn't these be outside the llvm::yaml namespace and in the test's anonymous namespace?

Above, only the template specializations are in the llvm::yaml namespace which makes more sense to me.

[zturner]

I moved the structs out of the yaml namespace as you suggested. Tested on VS 2013 and everything works fine. I'll upload a new patch if you want, but the only change was simply moving the classes in the unit tests to the global namespace.

And yes, I agree with you it uses some confusing design patterns. I can try to explain it in more detail if you like, let me know if that would be helpful.

[chandlerc]

LGTM with the previous changes and another change below.

include/llvm/Support/YAMLTraits.h

765

Is this in some 'internal' or 'detail' namespace that I don't see? If not, I think it probably should be. Seems like too generic of a name.

Also, probably should be doMapping to be consistent with naming conventions (not that this file is terribly consistent with the rest of LLVM's conventions).

include/llvm/Support/YAMLTraits.h

 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/YAMLParser.h"
 #include "llvm/Support/raw_ostream.h"
 #include <system_error>
 
 namespace llvm {
 namespace yaml {
 
+struct EmptyContext {};
+
 /// This class should be specialized by any type that needs to be converted
 /// to/from a YAML mapping.  For example:
 ///
 ///     struct MappingTraits<MyStruct> {
 ///       static void mapping(IO &io, MyStruct &s) {
 ///         io.mapRequired("name", s.name);
 ///         io.mapRequired("size", s.size);
 ///         io.mapOptional("age",  s.age);
 ///       }
 ///     };
 template<class T>
 struct MappingTraits {
   // Must provide:
   // static void mapping(IO &io, T &fields);
   // Optionally may provide:
   // static StringRef validate(IO &io, T &fields);
   //
   // The optional flow flag will cause generated YAML to use a flow mapping
   // (e.g. { a: 0, b: 1 }):
   // static const bool flow = true;
 };
 
+/// This class is similar to MappingTraits<T> but allows you to pass in
+/// additional context for each map operation.  For example:
+///
+///     struct MappingContextTraits<MyStruct, MyContext> {
+///       static void mapping(IO &io, MyStruct &s, MyContext &c) {
+///         io.mapRequired("name", s.name);
+///         io.mapRequired("size", s.size);
+///         io.mapOptional("age",  s.age);
+///         ++c.TimesMapped;
+///       }
+///     };
+template <class T, class Context> struct MappingContextTraits {
+  // Must provide:
+  // static void mapping(IO &io, T &fields, Context &Ctx);
+  // Optionally may provide:
+  // static StringRef validate(IO &io, T &fields, Context &Ctx);
+  //
+  // The optional flow flag will cause generated YAML to use a flow mapping
+  // (e.g. { a: 0, b: 1 }):
+  // static const bool flow = true;
+};
+
 /// This class should be specialized by any integral type that converts
 /// to/from a YAML scalar where there is a one-to-one mapping between
 /// in-memory values and a string in YAML.  For example:
 ///
 ///     struct ScalarEnumerationTraits<Colors> {
 ///         static void enumeration(IO &io, Colors &value) {
 ///           io.enumCase(value, "red",   cRed);
 ///           io.enumCase(value, "blue",  cBlue);
struct has_BlockScalarTraits
   template <typename U>
   static double test(...);
 
 public:
   static bool const value =
       (sizeof(test<BlockScalarTraits<T>>(nullptr, nullptr)) == 1);
 };
 
+// Test if MappingContextTraits<T> is defined on type T.
+template <class T, class Context> struct has_MappingTraits {
+  typedef void (*Signature_mapping)(class IO &, T &, Context &);
+
+  template <typename U>
+  static char test(SameType<Signature_mapping, &U::mapping>*);
+
+  template <typename U>
+  static double test(...);
+
+public:
+  static bool const value =
+      (sizeof(test<MappingContextTraits<T, Context>>(nullptr)) == 1);
+};
+
 // Test if MappingTraits<T> is defined on type T.
-template <class T>
-struct has_MappingTraits
-{
-  typedef void (*Signature_mapping)(class IO&, T&);
+template <class T> struct has_MappingTraits<T, EmptyContext> {
+  typedef void (*Signature_mapping)(class IO &, T &);
 
   template <typename U>
-  static char test(SameType<Signature_mapping, &U::mapping>*);
+  static char test(SameType<Signature_mapping, &U::mapping> *);
 
+  template <typename U> static double test(...);
+
+public:
+  static bool const value = (sizeof(test<MappingTraits<T>>(nullptr)) == 1);
+};
+
+// Test if MappingContextTraits<T>::validate() is defined on type T.
+template <class T, class Context> struct has_MappingValidateTraits {
+  typedef StringRef (*Signature_validate)(class IO &, T &, Context &);
+
+  template <typename U>
+  static char test(SameType<Signature_validate, &U::validate>*);
+
   template <typename U>
   static double test(...);
 
 public:
-  static bool const value = (sizeof(test<MappingTraits<T> >(nullptr)) == 1);
+  static bool const value =
+      (sizeof(test<MappingContextTraits<T, Context>>(nullptr)) == 1);
 };
 
 // Test if MappingTraits<T>::validate() is defined on type T.
-template <class T>
-struct has_MappingValidateTraits
-{
-  typedef StringRef (*Signature_validate)(class IO&, T&);
+template <class T> struct has_MappingValidateTraits<T, EmptyContext> {
+  typedef StringRef (*Signature_validate)(class IO &, T &);
 
   template <typename U>
-  static char test(SameType<Signature_validate, &U::validate>*);
+  static char test(SameType<Signature_validate, &U::validate> *);
 
-  template <typename U>
-  static double test(...);
+  template <typename U> static double test(...);
 
 public:
-  static bool const value = (sizeof(test<MappingTraits<T> >(nullptr)) == 1);
+  static bool const value = (sizeof(test<MappingTraits<T>>(nullptr)) == 1);
 };
 
 // Test if SequenceTraits<T> is defined on type T.
 template <class T>
 struct has_SequenceMethodTraits
 {
   typedef size_t (*Signature_size)(class IO&, T&);
 
inline bool needsQuotes(StringRef S) {
   if (isBool(S))
     return true;
   if (isNumeric(S))
     return true;
 
   return false;
 }
 
-template<typename T>
-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>::value
-                                      && !has_SequenceTraits<T>::value
-                                      && !has_DocumentListTraits<T>::value >  {};
-
-template<typename T>
-struct validatedMappingTraits : public std::integral_constant<bool,
-                                       has_MappingTraits<T>::value
-                                    && has_MappingValidateTraits<T>::value> {};
-
-template<typename T>
-struct unvalidatedMappingTraits : public std::integral_constant<bool,
-                                        has_MappingTraits<T>::value
-                                    && !has_MappingValidateTraits<T>::value> {};
+template <typename T, typename Context>
+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_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> {};
+
+template <typename T, typename Context>
+struct unvalidatedMappingTraits
+    : public std::integral_constant<
+          bool, has_MappingTraits<T, Context>::value &&
+                    !has_MappingValidateTraits<T, Context>::value> {};
+
 // Base class for Input and Output.
 class IO {
 public:
 
   IO(void *Ctxt=nullptr);
   virtual ~IO();
 
   virtual bool outputting() = 0;
void maskedBitSetCase(T &Val, const char *Str, uint32_t ConstVal,
                         uint32_t Mask) {
     if (bitSetMatch(Str, outputting() && (Val & Mask) == ConstVal))
       Val = Val | ConstVal;
   }
 
   void *getContext();
   void setContext(void *);
 
-  template <typename T>
-  void mapRequired(const char* Key, T& Val) {
-    this->processKey(Key, Val, true);
+  template <typename T> void mapRequired(const char *Key, T &Val) {
+    EmptyContext Ctx;
+    this->processKey(Key, Val, true, Ctx);
+  }
+  template <typename T, typename Context>
+  void mapRequired(const char *Key, T &Val, Context &Ctx) {
+    this->processKey(Key, Val, true, Ctx);
+  }
+
+  template <typename T> void mapOptional(const char *Key, T &Val) {
+    EmptyContext Ctx;
+    mapOptionalWithContext(Key, Val, Ctx);
   }
 
   template <typename T>
-  typename std::enable_if<has_SequenceTraits<T>::value,void>::type
-  mapOptional(const char* Key, T& Val) {
+  void mapOptional(const char *Key, T &Val, const T &Default) {
+    EmptyContext Ctx;
+    mapOptionalWithContext(Key, Val, Default, Ctx);
+  }
+
+  template <typename T, typename Context>
+  typename std::enable_if<has_SequenceTraits<T>::value, void>::type
+  mapOptionalWithContext(const char *Key, T &Val, Context &Ctx) {
     // omit key/value instead of outputting empty sequence
-    if ( this->canElideEmptySequence() && !(Val.begin() != Val.end()) )
+    if (this->canElideEmptySequence() && !(Val.begin() != Val.end()))
       return;
-    this->processKey(Key, Val, false);
+    this->processKey(Key, Val, false, Ctx);
   }
 
-  template <typename T>
-  void mapOptional(const char* Key, Optional<T> &Val) {
-    processKeyWithDefault(Key, Val, Optional<T>(), /*Required=*/false);
+  template <typename T, typename Context>
+  void mapOptionalWithContext(const char *Key, Optional<T> &Val, Context &Ctx) {
+    this->processKeyWithDefault(Key, Val, Optional<T>(), /*Required=*/false,
+                                Ctx);
   }
 
-  template <typename T>
-  typename std::enable_if<!has_SequenceTraits<T>::value,void>::type
-  mapOptional(const char* Key, T& Val) {
-    this->processKey(Key, Val, false);
+  template <typename T, typename Context>
+  typename std::enable_if<!has_SequenceTraits<T>::value, void>::type
+  mapOptionalWithContext(const char *Key, T &Val, Context &Ctx) {
+    this->processKey(Key, Val, false, Ctx);
   }
 
-  template <typename T>
-  void mapOptional(const char* Key, T& Val, const T& Default) {
-    this->processKeyWithDefault(Key, Val, Default, false);
+  template <typename T, typename Context>
+  void mapOptionalWithContext(const char *Key, T &Val, const T &Default,
+                              Context &Ctx) {
+    this->processKeyWithDefault(Key, Val, Default, false, Ctx);
   }
 
 private:
-  template <typename T>
+  template <typename T, typename Context>
   void processKeyWithDefault(const char *Key, Optional<T> &Val,
-                             const Optional<T> &DefaultValue, bool Required) {
+                             const Optional<T> &DefaultValue, bool Required,
+                             Context &Ctx) {
     assert(DefaultValue.hasValue() == false &&
            "Optional<T> shouldn't have a value!");
     void *SaveInfo;
     bool UseDefault;
     const bool sameAsDefault = outputting() && !Val.hasValue();
     if (!outputting() && !Val.hasValue())
       Val = T();
     if (this->preflightKey(Key, Required, sameAsDefault, UseDefault,
                            SaveInfo)) {
-      yamlize(*this, Val.getValue(), Required);
+      yamlize(*this, Val.getValue(), Required, Ctx);
       this->postflightKey(SaveInfo);
     } else {
       if (UseDefault)
         Val = DefaultValue;
     }
   }
 
-  template <typename T>
-  void processKeyWithDefault(const char *Key, T &Val, const T& DefaultValue,
-                                                                bool Required) {
+  template <typename T, typename Context>
+  void processKeyWithDefault(const char *Key, T &Val, const T &DefaultValue,
+                             bool Required, Context &Ctx) {
     void *SaveInfo;
     bool UseDefault;
     const bool sameAsDefault = outputting() && Val == DefaultValue;
     if ( this->preflightKey(Key, Required, sameAsDefault, UseDefault,
                                                                   SaveInfo) ) {
-      yamlize(*this, Val, Required);
+      yamlize(*this, Val, Required, Ctx);
       this->postflightKey(SaveInfo);
     }
     else {
       if ( UseDefault )
         Val = DefaultValue;
     }
   }
 
-  template <typename T>
-  void processKey(const char *Key, T &Val, bool Required) {
+  template <typename T, typename Context>
+  void processKey(const char *Key, T &Val, bool Required, Context &Ctx) {
     void *SaveInfo;
     bool UseDefault;
     if ( this->preflightKey(Key, Required, false, UseDefault, SaveInfo) ) {
-      yamlize(*this, Val, Required);
+      yamlize(*this, Val, Required, Ctx);
       this->postflightKey(SaveInfo);
     }
   }
 
 private:
   void  *Ctxt;
 };
 
-template<typename T>
-typename std::enable_if<has_ScalarEnumerationTraits<T>::value,void>::type
-yamlize(IO &io, T &Val, bool) {
+template <typename T>
+typename std::enable_if<has_ScalarEnumerationTraits<T>::value, void>::type
+yamlize(IO &io, T &Val, bool, EmptyContext &Ctx) {
   io.beginEnumScalar();
   ScalarEnumerationTraits<T>::enumeration(io, Val);
   io.endEnumScalar();
 }
 
-template<typename T>
-typename std::enable_if<has_ScalarBitSetTraits<T>::value,void>::type
-yamlize(IO &io, T &Val, bool) {
+template <typename T>
+typename std::enable_if<has_ScalarBitSetTraits<T>::value, void>::type
+yamlize(IO &io, T &Val, bool, EmptyContext &Ctx) {
   bool DoClear;
   if ( io.beginBitSetScalar(DoClear) ) {
     if ( DoClear )
       Val = static_cast<T>(0);
     ScalarBitSetTraits<T>::bitset(io, Val);
     io.endBitSetScalar();
   }
 }
 
-template<typename T>
-typename std::enable_if<has_ScalarTraits<T>::value,void>::type
-yamlize(IO &io, T &Val, bool) {
+template <typename T>
+typename std::enable_if<has_ScalarTraits<T>::value, void>::type
+yamlize(IO &io, T &Val, bool, EmptyContext &Ctx) {
   if ( io.outputting() ) {
     std::string Storage;
     llvm::raw_string_ostream Buffer(Storage);
     ScalarTraits<T>::output(Val, io.getContext(), Buffer);
     StringRef Str = Buffer.str();
     io.scalarString(Str, ScalarTraits<T>::mustQuote(Str));
   }
   else {
     StringRef Str;
     io.scalarString(Str, ScalarTraits<T>::mustQuote(Str));
     StringRef Result = ScalarTraits<T>::input(Str, io.getContext(), Val);
     if ( !Result.empty() ) {
       io.setError(llvm::Twine(Result));
     }
   }
 }
 
 template <typename T>
 typename std::enable_if<has_BlockScalarTraits<T>::value, void>::type
-yamlize(IO &YamlIO, T &Val, bool) {
+yamlize(IO &YamlIO, T &Val, bool, EmptyContext &Ctx) {
   if (YamlIO.outputting()) {
     std::string Storage;
     llvm::raw_string_ostream Buffer(Storage);
     BlockScalarTraits<T>::output(Val, YamlIO.getContext(), Buffer);
     StringRef Str = Buffer.str();
     YamlIO.blockScalarString(Str);
   } else {
     StringRef Str;
     YamlIO.blockScalarString(Str);
     StringRef Result =
         BlockScalarTraits<T>::input(Str, YamlIO.getContext(), Val);
     if (!Result.empty())
       YamlIO.setError(llvm::Twine(Result));
   }
 }
 
-template<typename T>
-typename std::enable_if<validatedMappingTraits<T>::value, void>::type
-yamlize(IO &io, T &Val, bool) {
+template <typename T, typename Context>
+void do_mapping(IO &io, T &Val, Context &Ctx) {
+  MappingContextTraits<T, Context>::mapping(io, Val, Ctx);
+}
+
+template <typename T> void do_mapping(IO &io, T &Val, EmptyContext &Ctx) {
+  MappingTraits<T>::mapping(io, Val);
+}
+
+template <typename T, typename Context>
+typename std::enable_if<validatedMappingTraits<T, Context>::value, void>::type
+yamlize(IO &io, T &Val, bool, Context &Ctx) {
   if (has_FlowTraits<MappingTraits<T>>::value)
     io.beginFlowMapping();
   else
     io.beginMapping();
   if (io.outputting()) {
     StringRef Err = MappingTraits<T>::validate(io, Val);
     if (!Err.empty()) {
       llvm::errs() << Err << "\n";
       assert(Err.empty() && "invalid struct trying to be written as yaml");
     }
   }
-  MappingTraits<T>::mapping(io, Val);
+  do_mapping(io, Val, Ctx);
   if (!io.outputting()) {
     StringRef Err = MappingTraits<T>::validate(io, Val);
     if (!Err.empty())
       io.setError(Err);
   }
   if (has_FlowTraits<MappingTraits<T>>::value)
     io.endFlowMapping();
   else
     io.endMapping();
 }
 
-template<typename T>
-typename std::enable_if<unvalidatedMappingTraits<T>::value, void>::type
-yamlize(IO &io, T &Val, bool) {
+template <typename T, typename Context>
+typename std::enable_if<unvalidatedMappingTraits<T, Context>::value, void>::type
+yamlize(IO &io, T &Val, bool, Context &Ctx) {
   if (has_FlowTraits<MappingTraits<T>>::value) {
     io.beginFlowMapping();
-    MappingTraits<T>::mapping(io, Val);
+    do_mapping(io, Val, Ctx);
     io.endFlowMapping();
   } else {
     io.beginMapping();
-    MappingTraits<T>::mapping(io, Val);
+    do_mapping(io, Val, Ctx);
     io.endMapping();
   }
 }
 
-template<typename T>
-typename std::enable_if<missingTraits<T>::value, void>::type
-yamlize(IO &io, T &Val, bool) {
+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 std::enable_if<has_SequenceTraits<T>::value,void>::type
-yamlize(IO &io, T &Seq, bool) {
+template <typename T, typename Context>
+typename std::enable_if<has_SequenceTraits<T>::value, void>::type
+yamlize(IO &io, T &Seq, bool, Context &Ctx) {
   if ( has_FlowTraits< SequenceTraits<T> >::value ) {
     unsigned incnt = io.beginFlowSequence();
     unsigned count = io.outputting() ? SequenceTraits<T>::size(io, Seq) : incnt;
     for(unsigned i=0; i < count; ++i) {
       void *SaveInfo;
       if ( io.preflightFlowElement(i, SaveInfo) ) {
-        yamlize(io, SequenceTraits<T>::element(io, Seq, i), true);
+        yamlize(io, SequenceTraits<T>::element(io, Seq, i), true, Ctx);
         io.postflightFlowElement(SaveInfo);
       }
     }
     io.endFlowSequence();
   }
   else {
     unsigned incnt = io.beginSequence();
     unsigned count = io.outputting() ? SequenceTraits<T>::size(io, Seq) : incnt;
     for(unsigned i=0; i < count; ++i) {
       void *SaveInfo;
       if ( io.preflightElement(i, SaveInfo) ) {
-        yamlize(io, SequenceTraits<T>::element(io, Seq, i), true);
+        yamlize(io, SequenceTraits<T>::element(io, Seq, i), true, Ctx);
         io.postflightElement(SaveInfo);
       }
     }
     io.endSequence();
   }
 }
 
 template<>

 };
 
 // Define non-member operator>> so that Input can stream in a document list.
 template <typename T>
 inline
 typename std::enable_if<has_DocumentListTraits<T>::value, Input &>::type
 operator>>(Input &yin, T &docList) {
   int i = 0;
+  EmptyContext Ctx;
   while ( yin.setCurrentDocument() ) {
-    yamlize(yin, DocumentListTraits<T>::element(yin, docList, i), true);
+    yamlize(yin, DocumentListTraits<T>::element(yin, docList, i), true, Ctx);
     if ( yin.error() )
       return yin;
     yin.nextDocument();
     ++i;
   }
   return yin;
 }
 
 // Define non-member operator>> so that Input can stream in a map as a document.
 template <typename T>
-inline
-typename std::enable_if<has_MappingTraits<T>::value, Input &>::type
+inline typename std::enable_if<has_MappingTraits<T, EmptyContext>::value,
+                               Input &>::type
 operator>>(Input &yin, T &docMap) {
+  EmptyContext Ctx;
   yin.setCurrentDocument();
-  yamlize(yin, docMap, true);
+  yamlize(yin, docMap, true, Ctx);
   return yin;
 }
 
 // Define non-member operator>> so that Input can stream in a sequence as
 // a document.
 template <typename T>
 inline
 typename std::enable_if<has_SequenceTraits<T>::value, Input &>::type
 operator>>(Input &yin, T &docSeq) {
+  EmptyContext Ctx;
   if (yin.setCurrentDocument())
-    yamlize(yin, docSeq, true);
+    yamlize(yin, docSeq, true, Ctx);
   return yin;
 }
 
 // Define non-member operator>> so that Input can stream in a block scalar.
 template <typename T>
 inline
 typename std::enable_if<has_BlockScalarTraits<T>::value, Input &>::type
 operator>>(Input &In, T &Val) {
+  EmptyContext Ctx;
   if (In.setCurrentDocument())
-    yamlize(In, Val, true);
+    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>::value, Input &>::type
+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;
 }
 
 // Define non-member operator<< so that Output can stream out document list.
 template <typename T>
 inline
 typename std::enable_if<has_DocumentListTraits<T>::value, Output &>::type
 operator<<(Output &yout, T &docList) {
+  EmptyContext Ctx;
   yout.beginDocuments();
   const size_t count = DocumentListTraits<T>::size(yout, docList);
   for(size_t i=0; i < count; ++i) {
     if ( yout.preflightDocument(i) ) {
-      yamlize(yout, DocumentListTraits<T>::element(yout, docList, i), true);
+      yamlize(yout, DocumentListTraits<T>::element(yout, docList, i), true,
+              Ctx);
       yout.postflightDocument();
     }
   }
   yout.endDocuments();
   return yout;
 }
 
 // Define non-member operator<< so that Output can stream out a map.
 template <typename T>
-inline
-typename std::enable_if<has_MappingTraits<T>::value, Output &>::type
+inline typename std::enable_if<has_MappingTraits<T, EmptyContext>::value,
+                               Output &>::type
 operator<<(Output &yout, T &map) {
+  EmptyContext Ctx;
   yout.beginDocuments();
   if ( yout.preflightDocument(0) ) {
-    yamlize(yout, map, true);
+    yamlize(yout, map, true, Ctx);
     yout.postflightDocument();
   }
   yout.endDocuments();
   return yout;
 }
 
 // Define non-member operator<< so that Output can stream out a sequence.
 template <typename T>
 inline
 typename std::enable_if<has_SequenceTraits<T>::value, Output &>::type
 operator<<(Output &yout, T &seq) {
+  EmptyContext Ctx;
   yout.beginDocuments();
   if ( yout.preflightDocument(0) ) {
-    yamlize(yout, seq, true);
+    yamlize(yout, seq, true, Ctx);
     yout.postflightDocument();
   }
   yout.endDocuments();
   return yout;
 }
 
 // Define non-member operator<< so that Output can stream out a block scalar.
 template <typename T>
 inline
 typename std::enable_if<has_BlockScalarTraits<T>::value, Output &>::type
 operator<<(Output &Out, T &Val) {
+  EmptyContext Ctx;
   Out.beginDocuments();
   if (Out.preflightDocument(0)) {
-    yamlize(Out, Val, true);
+    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>::value, Output &>::type
+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(); }

unittests/Support/YAMLIOTest.cpp

Seq.emplace_back("This is str3");
     ostr.flush();
     EXPECT_EQ(out,
               "---\n"
               "[ This is str1, This is str2, This is str3 ]\n"
               "...\n");
     out.clear();
   }
 }
+
+namespace llvm {
+namespace yaml {
+struct MappingContext {
+  int A = 0;
+};
+struct SimpleMap {
+  int B = 0;
+  int C = 0;
+};
+
+struct NestedMap {
+  NestedMap(MappingContext &Context) : Context(Context) {}
+  SimpleMap Simple;
+  MappingContext &Context;
+};
+
+template <> struct MappingContextTraits<SimpleMap, MappingContext> {
+  static void mapping(IO &io, SimpleMap &sm, MappingContext &Context) {
+    io.mapRequired("B", sm.B);
+    io.mapRequired("C", sm.C);
+    ++Context.A;
+    io.mapRequired("Context", Context.A);
+  }
+};
+
+template <> struct MappingTraits<NestedMap> {
+  static void mapping(IO &io, NestedMap &nm) {
+    io.mapRequired("Simple", nm.Simple, nm.Context);
+  }
+};
+}
+}
+
+TEST(YAMLIO, TestMapWithContext) {
+  llvm::yaml::MappingContext Context;
+  llvm::yaml::NestedMap Nested(Context);
+  std::string out;
+  llvm::raw_string_ostream ostr(out);
+
+  Output yout(ostr, nullptr, 15);
+
+  yout << Nested;
+  ostr.flush();
+  EXPECT_EQ(1, Context.A);
+  EXPECT_EQ("---\n"
+            "Simple:          \n"
+            "  B:               0\n"
+            "  C:               0\n"
+            "  Context:         1\n"
+            "...\n",
+            out);
+
+  out.clear();
+
+  Nested.Simple.B = 2;
+  Nested.Simple.C = 3;
+  yout << Nested;
+  ostr.flush();
+  EXPECT_EQ(2, Context.A);
+  EXPECT_EQ("---\n"
+            "Simple:          \n"
+            "  B:               2\n"
+            "  C:               3\n"
+            "  Context:         2\n"
+            "...\n",
+            out);
+  out.clear();
+}