From 157cbd41f460fa1abe72479b468b04a0fc847d55 Mon Sep 17 00:00:00 2001
From: Philippe Antoine <p.antoine@catenacyber.fr>
Date: Fri, 2 Aug 2024 21:50:22 +0200
Subject: [PATCH] infra/coverage: update profraw to version 9

---
 .../base-images/base-runner/profraw_update.py | 65 +++++++++++++------
 1 file changed, 46 insertions(+), 19 deletions(-)

diff --git a/infra/base-images/base-runner/profraw_update.py b/infra/base-images/base-runner/profraw_update.py
index bd2bf41cc3eb9..c8053a0886162 100755
--- a/infra/base-images/base-runner/profraw_update.py
+++ b/infra/base-images/base-runner/profraw_update.py
@@ -22,10 +22,11 @@
 import sys
 
 HeaderGeneric = namedtuple('HeaderGeneric', 'magic version')
-HeaderVersion7 = namedtuple(
-    'HeaderVersion7',
+HeaderVersion9 = namedtuple(
+    'HeaderVersion9',
     'BinaryIdsSize DataSize PaddingBytesBeforeCounters CountersSize \
-    PaddingBytesAfterCounters NamesSize CountersDelta NamesDelta ValueKindLast')
+    PaddingBytesAfterCounters NumBitmapBytes PaddingBytesAfterBitmapBytes NamesSize CountersDelta BitmapDelta NamesDelta ValueKindLast'
+)
 
 PROFRAW_MAGIC = 0xff6c70726f667281
 
@@ -57,39 +58,65 @@ def upgrade(data, sect_prf_cnts, sect_prf_data):
     # cf https://reviews.llvm.org/D111123
     generic_header = generic_header._replace(version=8)
     data = data[:8] + struct.pack('Q', generic_header.version) + data[16:]
-  v7_header = HeaderVersion7._make(struct.unpack('QQQQQQQQQ', data[16:88]))
-
-  if v7_header.BinaryIdsSize % 8 != 0:
+  was8 = False
+  if generic_header.version == 8:
+    was8 = True
+    # see https://reviews.llvm.org/D138846
+    generic_header = generic_header._replace(version=9)
+    # Upgrade from version 8 to 9 by adding NumBitmapBytes, PaddingBytesAfterBitmapBytes and BitmapDelta fields.
+    data = data[:8] + struct.pack(
+        'Q', generic_header.version) + data[16:56] + struct.pack(
+            'QQ', 0, 0) + data[56:72] + struct.pack('Q', 0) + data[72:]
+
+  v9_header = HeaderVersion9._make(struct.unpack('QQQQQQQQQQQQ', data[16:112]))
+
+  if v9_header.BinaryIdsSize % 8 != 0:
     # Adds padding for binary ids.
     # cf commit b9f547e8e51182d32f1912f97a3e53f4899ea6be
     # cf https://reviews.llvm.org/D110365
-    padlen = 8 - (v7_header.BinaryIdsSize % 8)
-    v7_header = v7_header._replace(BinaryIdsSize=v7_header.BinaryIdsSize +
+    padlen = 8 - (v9_header.BinaryIdsSize % 8)
+    v7_header = v9_header._replace(BinaryIdsSize=v9_header.BinaryIdsSize +
                                    padlen)
-    data = data[:16] + struct.pack('Q', v7_header.BinaryIdsSize) + data[24:]
-    data = data[:88 + v7_header.BinaryIdsSize] + bytes(
-        padlen) + data[88 + v7_header.BinaryIdsSize:]
-
-  if v7_header.CountersDelta != (sect_prf_cnts -
+    data = data[:16] + struct.pack('Q', v9_header.BinaryIdsSize) + data[24:]
+    data = data[:112 + v9_header.BinaryIdsSize] + bytes(
+        padlen) + data[112 + v9_header.BinaryIdsSize:]
+
+  if was8:
+    offset = 112 + v9_header.BinaryIdsSize
+    for d in range(v9_header.DataSize):
+      # add BitmapPtr and aligned u32(NumBitmapBytes)
+      data = data[:offset + 3 * 8] + struct.pack(
+          'Q', 0) + data[offset + 3 * 8:offset + 6 * 8] + struct.pack(
+              'Q', 0) + data[offset + 6 * 8:]
+      value = struct.unpack('Q',
+                            data[offset + 2 * 8:offset + 3 * 8])[0] - 16 * d
+      data = data[:offset + 2 * 8] + struct.pack('Q',
+                                                 value) + data[offset + 3 * 8:]
+      offset += 8 * 8
+
+  if v9_header.CountersDelta != (sect_prf_cnts -
                                  sect_prf_data) & 0xffffffffffffffff:
     # Rust linking seems to add an offset...
-    sect_prf_data = v7_header.CountersDelta - sect_prf_cnts + sect_prf_data
-    sect_prf_cnts = v7_header.CountersDelta
+    sect_prf_data = v9_header.CountersDelta - sect_prf_cnts + sect_prf_data
+    sect_prf_cnts = v9_header.CountersDelta
 
   dataref = sect_prf_data
   relativize_address(data, 64, dataref, sect_prf_cnts, sect_prf_data)
 
-  offset = 88 + v7_header.BinaryIdsSize
+  offset = 112 + v9_header.BinaryIdsSize
   # This also works for C+Rust binaries compiled with
   # clang-14/rust-nightly-clang-13.
-  for _ in range(v7_header.DataSize):
+  for _ in range(v9_header.DataSize):
     # 16 is the offset of CounterPtr in ProfrawData structure.
     relativize_address(data, offset + 16, dataref, sect_prf_cnts, sect_prf_data)
     # We need this because of CountersDelta -= sizeof(*SrcData);
     # seen in __llvm_profile_merge_from_buffer.
-    dataref += 44 + 2 * (v7_header.ValueKindLast + 1)
+    dataref += 44 + 2 * (v9_header.ValueKindLast + 1)
+    if was8:
+      #profraw9 added RelativeBitmapPtr and NumBitmapBytes (8+4 rounded up to 16)
+      dataref -= 16
     # This is the size of one ProfrawData structure.
-    offset += 44 + 2 * (v7_header.ValueKindLast + 1)
+    offset += 44 + 2 * (v9_header.ValueKindLast + 1)
 
   return data