1*6db8c59fSFangrui Song // Check that we can patch and unpatch specific function ids.
2*6db8c59fSFangrui Song //
3*6db8c59fSFangrui Song // RUN: %clangxx_xray -std=c++11 %s -o %t
4*6db8c59fSFangrui Song // RUN: XRAY_OPTIONS="patch_premain=false xray_naive_log=false" %run %t | FileCheck %s
5*6db8c59fSFangrui Song 
6*6db8c59fSFangrui Song // UNSUPPORTED: target-is-mips64,target-is-mips64el
7*6db8c59fSFangrui Song 
8*6db8c59fSFangrui Song #include "xray/xray_interface.h"
9*6db8c59fSFangrui Song 
10*6db8c59fSFangrui Song #include <set>
11*6db8c59fSFangrui Song #include <cstdio>
12*6db8c59fSFangrui Song #include <cassert>
13*6db8c59fSFangrui Song 
14*6db8c59fSFangrui Song std::set<int32_t> function_ids;
15*6db8c59fSFangrui Song 
16*6db8c59fSFangrui Song [[clang::xray_never_instrument]] void coverage_handler(int32_t fid,
17*6db8c59fSFangrui Song                                                        XRayEntryType) {
18*6db8c59fSFangrui Song   thread_local bool patching = false;
19*6db8c59fSFangrui Song   if (patching) return;
20*6db8c59fSFangrui Song   patching = true;
21*6db8c59fSFangrui Song   function_ids.insert(fid);
22*6db8c59fSFangrui Song   __xray_unpatch_function(fid);
23*6db8c59fSFangrui Song   patching = false;
24*6db8c59fSFangrui Song }
25*6db8c59fSFangrui Song 
26*6db8c59fSFangrui Song [[clang::xray_always_instrument]] void baz() {
27*6db8c59fSFangrui Song   // do nothing!
28*6db8c59fSFangrui Song }
29*6db8c59fSFangrui Song 
30*6db8c59fSFangrui Song [[clang::xray_always_instrument]] void bar() {
31*6db8c59fSFangrui Song   baz();
32*6db8c59fSFangrui Song }
33*6db8c59fSFangrui Song 
34*6db8c59fSFangrui Song [[clang::xray_always_instrument]] void foo() {
35*6db8c59fSFangrui Song   bar();
36*6db8c59fSFangrui Song }
37*6db8c59fSFangrui Song 
38*6db8c59fSFangrui Song [[clang::xray_always_instrument]] int main(int argc, char *argv[]) {
39*6db8c59fSFangrui Song   __xray_set_handler(coverage_handler);
40*6db8c59fSFangrui Song   assert(__xray_patch() == XRayPatchingStatus::SUCCESS);
41*6db8c59fSFangrui Song   foo();
42*6db8c59fSFangrui Song   assert(__xray_unpatch() == XRayPatchingStatus::SUCCESS);
43*6db8c59fSFangrui Song 
44*6db8c59fSFangrui Song   // print out the function_ids.
45*6db8c59fSFangrui Song   printf("first pass.\n");
46*6db8c59fSFangrui Song   for (const auto id : function_ids)
47*6db8c59fSFangrui Song     printf("patched: %d\n", id);
48*6db8c59fSFangrui Song 
49*6db8c59fSFangrui Song   // CHECK-LABEL: first pass.
50*6db8c59fSFangrui Song   // CHECK-DAG: patched: [[F1:.*]]
51*6db8c59fSFangrui Song   // CHECK-DAG: patched: [[F2:.*]]
52*6db8c59fSFangrui Song   // CHECK-DAG: patched: [[F3:.*]]
53*6db8c59fSFangrui Song 
54*6db8c59fSFangrui Song   // make a copy of the function_ids, then patch them later.
55*6db8c59fSFangrui Song   auto called_fns = function_ids;
56*6db8c59fSFangrui Song 
57*6db8c59fSFangrui Song   // clear the function_ids.
58*6db8c59fSFangrui Song   function_ids.clear();
59*6db8c59fSFangrui Song 
60*6db8c59fSFangrui Song   // patch the functions we've called before.
61*6db8c59fSFangrui Song   for (const auto id : called_fns)
62*6db8c59fSFangrui Song     assert(__xray_patch_function(id) == XRayPatchingStatus::SUCCESS);
63*6db8c59fSFangrui Song 
64*6db8c59fSFangrui Song   // then call them again.
65*6db8c59fSFangrui Song   foo();
66*6db8c59fSFangrui Song   assert(__xray_unpatch() == XRayPatchingStatus::SUCCESS);
67*6db8c59fSFangrui Song 
68*6db8c59fSFangrui Song   // confirm that we've seen the same functions again.
69*6db8c59fSFangrui Song   printf("second pass.\n");
70*6db8c59fSFangrui Song   for (const auto id : function_ids)
71*6db8c59fSFangrui Song     printf("patched: %d\n", id);
72*6db8c59fSFangrui Song   // CHECK-LABEL: second pass.
73*6db8c59fSFangrui Song   // CHECK-DAG: patched: [[F1]]
74*6db8c59fSFangrui Song   // CHECK-DAG: patched: [[F2]]
75*6db8c59fSFangrui Song   // CHECK-DAG: patched: [[F3]]
76*6db8c59fSFangrui Song 
77*6db8c59fSFangrui Song   // Now we want to make sure that if we unpatch one, that we're only going to
78*6db8c59fSFangrui Song   // see two calls of the coverage_handler.
79*6db8c59fSFangrui Song   function_ids.clear();
80*6db8c59fSFangrui Song   assert(__xray_patch() == XRayPatchingStatus::SUCCESS);
81*6db8c59fSFangrui Song   assert(__xray_unpatch_function(1) == XRayPatchingStatus::SUCCESS);
82*6db8c59fSFangrui Song   foo();
83*6db8c59fSFangrui Song   assert(__xray_unpatch() == XRayPatchingStatus::SUCCESS);
84*6db8c59fSFangrui Song 
85*6db8c59fSFangrui Song   // confirm that we don't see function id one called anymore.
86*6db8c59fSFangrui Song   printf("missing 1.\n");
87*6db8c59fSFangrui Song   for (const auto id : function_ids)
88*6db8c59fSFangrui Song     printf("patched: %d\n", id);
89*6db8c59fSFangrui Song   // CHECK-LABEL: missing 1.
90*6db8c59fSFangrui Song   // CHECK-NOT: patched: 1
91*6db8c59fSFangrui Song }
92