From 15fa3ba547bc3ee04af5c32b8f723a97e3feefd8 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov@fb.com>
Date: Tue, 3 Sep 2024 11:27:57 -0700
Subject: [PATCH] [BOLT][YAML] Allow unknown keys in the input (#100824)

This ensures forward compatibility, where old BOLT versions can consume
the profile created by newer versions with extra keys.

Test Plan: added yaml-unknown-keys.test
---
 bolt/lib/Profile/YAMLProfileReader.cpp |  1 +
 bolt/test/X86/yaml-unknown-keys.test   | 50 ++++++++++++++++++++++++++
 bolt/tools/merge-fdata/merge-fdata.cpp |  1 +
 3 files changed, 52 insertions(+)
 create mode 100644 bolt/test/X86/yaml-unknown-keys.test

diff --git a/bolt/lib/Profile/YAMLProfileReader.cpp b/bolt/lib/Profile/YAMLProfileReader.cpp
index 3eca5e972fa5ba..2ab4bf9a4b3e13 100644
--- a/bolt/lib/Profile/YAMLProfileReader.cpp
+++ b/bolt/lib/Profile/YAMLProfileReader.cpp
@@ -373,6 +373,7 @@ Error YAMLProfileReader::preprocessProfile(BinaryContext &BC) {
     return errorCodeToError(EC);
   }
   yaml::Input YamlInput(MB.get()->getBuffer());
+  YamlInput.setAllowUnknownKeys(true);
 
   // Consume YAML file.
   YamlInput >> YamlBP;
diff --git a/bolt/test/X86/yaml-unknown-keys.test b/bolt/test/X86/yaml-unknown-keys.test
new file mode 100644
index 00000000000000..cbcf9a4d304417
--- /dev/null
+++ b/bolt/test/X86/yaml-unknown-keys.test
@@ -0,0 +1,50 @@
+## Test that BOLT gracefully handles a YAML profile with unknown keys.
+
+# REQUIRES: system-linux
+# RUN: split-file %s %t
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %t/main.s -o %t.o
+# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q -nostdlib
+# RUN: llvm-bolt %t.exe -o %t.null --data %t/profile.yaml \
+# RUN:   --profile-ignore-hash -v=1 2>&1 | FileCheck %s
+# CHECK: warning: unknown key 'unknown-header-key'
+# CHECK: warning: unknown key 'unknown_succ_key'
+# CHECK: warning: unknown key 'unknown_block_key'
+# CHECK: warning: unknown key 'unknown_function_key'
+# CHECK: warning: unknown key 'unknown_toplev_key'
+#--- main.s
+  .globl main
+  .type main, %function
+main:
+  .cfi_startproc
+  cmpl	$0x0, %eax
+  jne .LBB0
+.LBB0:
+  retq
+  .cfi_endproc
+.size main, .-main
+#--- profile.yaml
+---
+header:
+  profile-version: 1
+  binary-name:     'yaml-multiple-profiles.test.tmp.exe'
+  binary-build-id: '<unknown>'
+  profile-flags:   [ lbr ]
+  profile-origin:  branch profile reader
+  profile-events:  ''
+  dfs-order:       false
+  unknown-header-key: true
+functions:
+  - name:            'main'
+    fid:             1
+    hash:            0x50BBA3441D436491
+    exec:            1
+    nblocks:         1
+    blocks:
+      - bid:             0
+        insns:           2
+        hash:            0x4D4D8FAF7D4C0000
+        succ:            [ { bid: 1, cnt: 0, unknown_succ_key: 0x10 } ]
+        unknown_block_key: [ ]
+    unknown_function_key: 1
+unknown_toplev_key: ''
+...
diff --git a/bolt/tools/merge-fdata/merge-fdata.cpp b/bolt/tools/merge-fdata/merge-fdata.cpp
index f2ac5ad4492ee5..b640aae808f56e 100644
--- a/bolt/tools/merge-fdata/merge-fdata.cpp
+++ b/bolt/tools/merge-fdata/merge-fdata.cpp
@@ -392,6 +392,7 @@ int main(int argc, char **argv) {
     if (std::error_code EC = MB.getError())
       report_error(InputDataFilename, EC);
     yaml::Input YamlInput(MB.get()->getBuffer());
+    YamlInput.setAllowUnknownKeys(true);
 
     errs() << "Merging data from " << InputDataFilename << "...\n";