[memprof] Accept a function name in YAML

llvm/include/llvm/ProfileData/MemProf.h

@@ -11,6 +11,7 @@
 #include "llvm/Support/BLAKE3.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/EndianStream.h"
+#include "llvm/Support/Format.h"
 #include "llvm/Support/HashBuilder.h"
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/Support/raw_ostream.h"
@@ -491,11 +492,15 @@ struct MemProfRecord {
   }
 };
 
+// A "typedef" for GUID.  See ScalarTraits<memprof::GUIDHex64> for how a GUID is
+// serialized and deserialized in YAML.
+LLVM_YAML_STRONG_TYPEDEF(uint64_t, GUIDHex64)
+
 // Helper struct for AllMemProfData.  In YAML, we treat the GUID and the fields
 // within MemProfRecord at the same level as if the GUID were part of
 // MemProfRecord.
 struct GUIDMemProfRecordPair {
-  GlobalValue::GUID GUID;
+  GUIDHex64 GUID;
   MemProfRecord Record;
 };
 
@@ -1166,6 +1171,27 @@ template <typename FrameIdTy> class CallStackRadixTreeBuilder {
 } // namespace memprof
 
 namespace yaml {
+template <> struct ScalarTraits<memprof::GUIDHex64> {
+  static void output(const memprof::GUIDHex64 &Val, void *, raw_ostream &Out) {
+    // Print GUID as a 16-digit hexadecimal number.
+    Out << format("0x%016" PRIx64, (uint64_t)Val);
+  }
+  static StringRef input(StringRef Scalar, void *, memprof::GUIDHex64 &Val) {
+    uint64_t Num;
+    if (Scalar.starts_with_insensitive("0x")) {
+      // Accept hexadecimal numbers starting with 0x or 0X.
+      if (Scalar.getAsInteger(0, Num))
+        return "invalid hex64 number";
+      Val = Num;
+    } else {
+      // Otherwise, treat the input as a string containing a function name.
+      Val = memprof::IndexedMemProfRecord::getGUID(Scalar);
+    }
+    return StringRef();
+  }
+  static QuotingType mustQuote(StringRef) { return QuotingType::None; }
+};
+
 template <> struct MappingTraits<memprof::Frame> {
   static void mapping(IO &Io, memprof::Frame &F) {
     Io.mapRequired("Function", F.Function);

llvm/test/tools/llvm-profdata/memprof-yaml.test

@@ -8,7 +8,7 @@
 ;--- memprof-in.yaml
 ---
 HeapProfileRecords:
-  - GUID:            16045690981402826360
+  - GUID:            0xdeadbeef12345678
     AllocSites:
       - Callstack:
           - { Function: 100, LineOffset: 11, Column: 10, IsInlineFrame: true }

llvm/unittests/ProfileData/MemProfTest.cpp

@@ -35,6 +35,7 @@ using ::llvm::StringRef;
 using ::llvm::object::SectionedAddress;
 using ::llvm::symbolize::SymbolizableModule;
 using ::testing::ElementsAre;
+using ::testing::IsEmpty;
 using ::testing::Pair;
 using ::testing::Return;
 using ::testing::SizeIs;
@@ -760,6 +761,43 @@ TEST(MemProf, YAMLParser) {
               ElementsAre(hashCallStack(CS3), hashCallStack(CS4)));
 }
 
+// Verify that the YAML parser accepts a GUID expressed as a function name.
+TEST(MemProf, YAMLParserGUID) {
+  StringRef YAMLData = R"YAML(
+---
+HeapProfileRecords:
+- GUID: _Z3fooi
+  AllocSites:
+  - Callstack:
+    - {Function: 0x100, LineOffset: 11, Column: 10, IsInlineFrame: true}
+    MemInfoBlock: {}
+  CallSites: []
+)YAML";
+
+  YAMLMemProfReader YAMLReader;
+  YAMLReader.parse(YAMLData);
+  IndexedMemProfData MemProfData = YAMLReader.takeMemProfData();
+
+  Frame F1(0x100, 11, 10, true);
+
+  llvm::SmallVector<FrameId> CS1 = {F1.hash()};
+
+  // Verify the entire contents of MemProfData.Frames.
+  EXPECT_THAT(MemProfData.Frames, UnorderedElementsAre(Pair(F1.hash(), F1)));
+
+  // Verify the entire contents of MemProfData.Frames.
+  EXPECT_THAT(MemProfData.CallStacks,
+              UnorderedElementsAre(Pair(hashCallStack(CS1), CS1)));
+
+  // Verify the entire contents of MemProfData.Records.
+  ASSERT_THAT(MemProfData.Records, SizeIs(1));
+  const auto &[GUID, Record] = *MemProfData.Records.begin();
+  EXPECT_EQ(GUID, IndexedMemProfRecord::getGUID("_Z3fooi"));
+  ASSERT_THAT(Record.AllocSites, SizeIs(1));
+  EXPECT_EQ(Record.AllocSites[0].CSId, hashCallStack(CS1));
+  EXPECT_THAT(Record.CallSiteIds, IsEmpty());
+}
+
 template <typename T> std::string serializeInYAML(T &Val) {
   std::string Out;
   llvm::raw_string_ostream OS(Out);