• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python3
2# Copyright 2021 Google LLC
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16################################################################################
17"""Helper script for upgrading a profraw file to latest version."""
18
19from collections import namedtuple
20import struct
21import subprocess
22import sys
23
24HeaderGeneric = namedtuple('HeaderGeneric', 'magic version')
25HeaderVersion7 = namedtuple(
26    'HeaderVersion7',
27    'BinaryIdsSize DataSize PaddingBytesBeforeCounters CountersSize \
28    PaddingBytesAfterCounters NamesSize CountersDelta NamesDelta ValueKindLast')
29
30PROFRAW_MAGIC = 0xff6c70726f667281
31
32
33def relativize_address(data, offset, databegin, sect_prf_cnts, sect_prf_data):
34  """Turns an absolute offset into a relative one."""
35  value = struct.unpack('Q', data[offset:offset + 8])[0]
36  if sect_prf_cnts <= value < sect_prf_data:
37    # If the value is an address in the right section, make it relative.
38    value = (value - databegin) & 0xffffffffffffffff
39    value = struct.pack('Q', value)
40    for i in range(8):
41      data[offset + i] = value[i]
42
43
44def upgrade(data, sect_prf_cnts, sect_prf_data):
45  """Upgrades profraw data, knowing the sections addresses."""
46  generic_header = HeaderGeneric._make(struct.unpack('QQ', data[:16]))
47  if generic_header.magic != PROFRAW_MAGIC:
48    raise Exception('Bad magic.')
49  if generic_header.version == 5:
50    generic_header = generic_header._replace(version=7)
51    # Upgrade from version 5 to 7 by adding binaryids field.
52    data = struct.pack('QQ', generic_header) + struct.pack('Q', 0) + data[16:]
53  if generic_header.version < 7:
54    raise Exception('Unhandled version.')
55  v7_header = HeaderVersion7._make(struct.unpack('QQQQQQQQQ', data[16:88]))
56
57  if v7_header.BinaryIdsSize % 8 != 0:
58    # Adds padding for binary ids.
59    # cf commit b9f547e8e51182d32f1912f97a3e53f4899ea6be
60    # cf https://reviews.llvm.org/D110365
61    padlen = 8 - (v7_header.BinaryIdsSize % 8)
62    v7_header = v7_header._replace(BinaryIdsSize=v7_header.BinaryIdsSize +
63                                   padlen)
64    data = data[:16] + struct.pack('Q', v7_header.BinaryIdsSize) + data[24:]
65    data = data[:88 + v7_header.BinaryIdsSize] + bytes(
66        padlen) + data[88 + v7_header.BinaryIdsSize:]
67
68  if v7_header.CountersDelta != sect_prf_cnts - sect_prf_data:
69    # Rust linking seems to add an offset...
70    sect_prf_data = v7_header.CountersDelta - sect_prf_cnts + sect_prf_data
71    sect_prf_cnts = v7_header.CountersDelta
72
73  dataref = sect_prf_data
74  relativize_address(data, 64, dataref, sect_prf_cnts, sect_prf_data)
75
76  offset = 88 + v7_header.BinaryIdsSize
77  # This also works for C+Rust binaries compiled with
78  # clang-14/rust-nightly-clang-13.
79  for _ in range(v7_header.DataSize):
80    # 16 is the offset of CounterPtr in ProfrawData structure.
81    relativize_address(data, offset + 16, dataref, sect_prf_cnts, sect_prf_data)
82    # We need this because of CountersDelta -= sizeof(*SrcData);
83    # seen in __llvm_profile_merge_from_buffer.
84    dataref += 44 + 2 * (v7_header.ValueKindLast + 1)
85    # This is the size of one ProfrawData structure.
86    offset += 44 + 2 * (v7_header.ValueKindLast + 1)
87
88  return data
89
90
91def main():
92  """Helper script for upgrading a profraw file to latest version."""
93  if len(sys.argv) != 4:
94    sys.stderr.write('Usage: %s <binary> <profraw> <output>\n' % sys.argv[0])
95    return 1
96
97  # First find llvm profile sections addresses in the elf, quick and dirty.
98  process = subprocess.Popen(['readelf', '-S', sys.argv[1]],
99                             stdout=subprocess.PIPE)
100  output, err = process.communicate()
101  if err:
102    print('readelf failed')
103    return 2
104  for line in iter(output.split(b'\n')):
105    if b'__llvm_prf_cnts' in line:
106      sect_prf_cnts = int(line.split()[4], 16)
107    elif b'__llvm_prf_data' in line:
108      sect_prf_data = int(line.split()[4], 16)
109
110  # Then open and read the input profraw file.
111  with open(sys.argv[2], 'rb') as input_file:
112    profraw_base = bytearray(input_file.read())
113  # Do the upgrade, returning a bytes object.
114  profraw_latest = upgrade(profraw_base, sect_prf_cnts, sect_prf_data)
115  # Write the output to the file given to the command line.
116  with open(sys.argv[3], 'wb') as output_file:
117    output_file.write(profraw_latest)
118
119  return 0
120
121
122if __name__ == '__main__':
123  sys.exit(main())
124