[memprof] Take Schema into account in PortableMemInfoBlock::serializedSize

llvm/include/llvm/ProfileData/MemProf.h

@@ -138,11 +138,22 @@ struct PortableMemInfoBlock {
     return !operator==(Other);
   }
 
-  static constexpr size_t serializedSize() {
+  static size_t serializedSize(const MemProfSchema &Schema) {
     size_t Result = 0;
-#define MIBEntryDef(NameTag, Name, Type) Result += sizeof(Type);
+
+    for (const Meta Id : Schema) {
+      switch (Id) {
+#define MIBEntryDef(NameTag, Name, Type)                                       \
+  case Meta::Name: {                                                           \
+    Result += sizeof(Type);                                                    \
+  } break;
 #include "llvm/ProfileData/MIBEntryDef.inc"
 #undef MIBEntryDef
+      default:
+        llvm_unreachable("Unknown meta type id, invalid input?");
+      }
+    }
+
     return Result;
   }
 
@@ -292,7 +303,8 @@ struct IndexedAllocationInfo {
       : CallStack(CS.begin(), CS.end()), CSId(CSId), Info(MB) {}
 
   // Returns the size in bytes when this allocation info struct is serialized.
-  size_t serializedSize(IndexedVersion Version) const;
+  size_t serializedSize(const MemProfSchema &Schema,
+                        IndexedVersion Version) const;
 
   bool operator==(const IndexedAllocationInfo &Other) const {
     if (Other.Info != Info)
@@ -367,7 +379,8 @@ struct IndexedMemProfRecord {
     CallSites.append(Other.CallSites);
   }
 
-  size_t serializedSize(IndexedVersion Version) const;
+  size_t serializedSize(const MemProfSchema &Schema,
+                        IndexedVersion Version) const;
 
   bool operator==(const IndexedMemProfRecord &Other) const {
     if (Other.AllocSites != AllocSites)
@@ -535,7 +548,7 @@ class RecordWriterTrait {
     endian::Writer LE(Out, llvm::endianness::little);
     offset_type N = sizeof(K);
     LE.write<offset_type>(N);
-    offset_type M = V.serializedSize(Version);
+    offset_type M = V.serializedSize(*Schema, Version);
     LE.write<offset_type>(M);
     return std::make_pair(N, M);
   }

llvm/lib/ProfileData/MemProf.cpp

@@ -10,42 +10,46 @@
 
 namespace llvm {
 namespace memprof {
-static size_t serializedSizeV0(const IndexedAllocationInfo &IAI) {
+static size_t serializedSizeV0(const IndexedAllocationInfo &IAI,
+                               const MemProfSchema &Schema) {
   size_t Size = 0;
   // The number of frames to serialize.
   Size += sizeof(uint64_t);
   // The callstack frame ids.
   Size += sizeof(FrameId) * IAI.CallStack.size();
   // The size of the payload.
-  Size += PortableMemInfoBlock::serializedSize();
+  Size += PortableMemInfoBlock::serializedSize(Schema);
   return Size;
 }
 
-static size_t serializedSizeV2(const IndexedAllocationInfo &IAI) {
+static size_t serializedSizeV2(const IndexedAllocationInfo &IAI,
+                               const MemProfSchema &Schema) {
   size_t Size = 0;
   // The CallStackId
   Size += sizeof(CallStackId);
   // The size of the payload.
-  Size += PortableMemInfoBlock::serializedSize();
+  Size += PortableMemInfoBlock::serializedSize(Schema);
   return Size;
 }
 
-size_t IndexedAllocationInfo::serializedSize(IndexedVersion Version) const {
+size_t IndexedAllocationInfo::serializedSize(const MemProfSchema &Schema,
+                                             IndexedVersion Version) const {
   switch (Version) {
   case Version0:
   case Version1:
-    return serializedSizeV0(*this);
+    return serializedSizeV0(*this, Schema);
   case Version2:
-    return serializedSizeV2(*this);
+    return serializedSizeV2(*this, Schema);
   }
   llvm_unreachable("unsupported MemProf version");
 }
 
-static size_t serializedSizeV0(const IndexedMemProfRecord &Record) {
+static size_t serializedSizeV0(const IndexedMemProfRecord &Record,
+                               const MemProfSchema &Schema) {
   // The number of alloc sites to serialize.
   size_t Result = sizeof(uint64_t);
   for (const IndexedAllocationInfo &N : Record.AllocSites)
-    Result += N.serializedSize(Version0);
+    Result += N.serializedSize(Schema, Version0);
 
   // The number of callsites we have information for.
   Result += sizeof(uint64_t);
@@ -57,11 +61,12 @@ static size_t serializedSizeV0(const IndexedMemProfRecord &Record) {
   return Result;
 }
 
-static size_t serializedSizeV2(const IndexedMemProfRecord &Record) {
+static size_t serializedSizeV2(const IndexedMemProfRecord &Record,
+                               const MemProfSchema &Schema) {
   // The number of alloc sites to serialize.
   size_t Result = sizeof(uint64_t);
   for (const IndexedAllocationInfo &N : Record.AllocSites)
-    Result += N.serializedSize(Version2);
+    Result += N.serializedSize(Schema, Version2);
 
   // The number of callsites we have information for.
   Result += sizeof(uint64_t);
@@ -70,13 +75,14 @@ static size_t serializedSizeV2(const IndexedMemProfRecord &Record) {
   return Result;
 }
 
-size_t IndexedMemProfRecord::serializedSize(IndexedVersion Version) const {
+size_t IndexedMemProfRecord::serializedSize(const MemProfSchema &Schema,
+                                            IndexedVersion Version) const {
   switch (Version) {
   case Version0:
   case Version1:
-    return serializedSizeV0(*this);
+    return serializedSizeV0(*this, Schema);
   case Version2:
-    return serializedSizeV2(*this);
+    return serializedSizeV2(*this, Schema);
   }
   llvm_unreachable("unsupported MemProf version");
 }
@@ -156,7 +162,7 @@ static IndexedMemProfRecord deserializeV0(const MemProfSchema &Schema,
     }
     Node.CSId = hashCallStack(Node.CallStack);
     Node.Info.deserialize(Schema, Ptr);
-    Ptr += PortableMemInfoBlock::serializedSize();
+    Ptr += PortableMemInfoBlock::serializedSize(Schema);
     Record.AllocSites.push_back(Node);
   }
 
@@ -193,7 +199,7 @@ static IndexedMemProfRecord deserializeV2(const MemProfSchema &Schema,
     IndexedAllocationInfo Node;
     Node.CSId = endian::readNext<CallStackId, llvm::endianness::little>(Ptr);
     Node.Info.deserialize(Schema, Ptr);
-    Ptr += PortableMemInfoBlock::serializedSize();
+    Ptr += PortableMemInfoBlock::serializedSize(Schema);
     Record.AllocSites.push_back(Node);
   }