1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

use std::ffi::CString;
use std::mem;
use std::pin::Pin;

use detour::RawDetour;
use dynasmrt::DynasmApi;
use dynasmrt::DynasmLabelApi;
use dynasmrt::ExecutableBuffer;
use dynasmrt::{dynasm, x64::Assembler};
use once_cell::sync::OnceCell;
use tracy_client::sys::___tracy_alloc_srcloc;
use windows::Win32::System::Threading::TlsAlloc;

use super::FrameworkError;
use super::FrameworkGlobal;

pub mod runtime;

pub struct Profiler {
    tls_index: u32,
    tracy: tracy_client::Client,
}

type WindowsFn = unsafe extern "C" fn();

#[repr(C)]
#[derive(Debug)]
pub struct ZoneData {
    name: u64,
    addr: Option<unsafe extern "C" fn()>,
    profiler_entry_thunk: unsafe extern "C" fn(),
    profiler_exit_thunk: unsafe extern "C" fn(),
    prologue: Option<unsafe extern "C" fn()>,
    tls_index: u32,
}

#[derive(Debug)]
pub struct ProfiledFunction {
    code_buffer: ExecutableBuffer,
    zone: Pin<Box<ZoneData>>,
    detour: RawDetour,
    name_cstr: CString,
}

#[cfg(test)]
mod test {
    use faithe::offset_of;

    use super::*;
    #[test]
    fn zone_data_layout() {
        assert_eq!(0x0, offset_of!(ZoneData, name));
        assert_eq!(0x08, offset_of!(ZoneData, addr));
        assert_eq!(0x10, offset_of!(ZoneData, profiler_entry_thunk));
        assert_eq!(0x18, offset_of!(ZoneData, profiler_exit_thunk));
        assert_eq!(0x20, offset_of!(ZoneData, prologue));
        assert_eq!(0x28, offset_of!(ZoneData, tls_index));
        assert_eq!(0x30, mem::size_of::<ZoneData>());
    }
}

impl FrameworkGlobal for Profiler {
    fn cell() -> &'static OnceCell<Self> {
        static INSTANCE: OnceCell<Profiler> = OnceCell::new();
        &INSTANCE
    }

    fn create() -> Result<Self, super::FrameworkError> {
        let tracy = tracy_client::Client::start();

        Ok(Profiler {
            tls_index: unsafe { TlsAlloc() },
            tracy,
        })
    }
}

impl Profiler {
    pub fn tracy(&self) -> &tracy_client::Client {
        &self.tracy
    }

    pub fn install_at(
        &self,
        addr: *const (),
        name: &'static str,
    ) -> Result<ProfiledFunction, FrameworkError> {
        let name_cstr = CString::new(name).expect("given name is not a valid ASCII string"); // TODO: should be a FrameworkError

        let source_loc = unsafe {
            ___tracy_alloc_srcloc(
                1,
                name_cstr.as_ptr(),
                name_cstr.to_bytes().len(),
                std::ptr::null(),
                0,
            )
        };

        let mut zone = Box::pin(ZoneData {
            name: source_loc,
            addr: None,
            profiler_entry_thunk: runtime::profiler_entry,
            profiler_exit_thunk: runtime::profiler_exit,
            prologue: None,
            tls_index: self.tls_index,
        });

        let mut ops = Assembler::new().expect("unable to create assembler");
        let data_ptr = &*zone as *const ZoneData;

        let prelude_offset = ops.offset();
        dynasm!(ops
            ; -> prelude:
            ; push rbx
            ; mov rbx, QWORD data_ptr as _
            ; jmp QWORD [rbx + 16]
            ; int3
        );

        let prologue_offset = ops.offset();
        dynasm!(ops
            ; -> prologue:
            ; mov rbx, QWORD data_ptr as _
            ; jmp QWORD [rbx + 24]
            ; int3
        );

        let code_buffer = ops
            .finalize()
            .expect("unable to create profiler prelude/prologue buffer");

        let prelude: WindowsFn = unsafe { mem::transmute(code_buffer.ptr(prelude_offset)) };
        let prologue: WindowsFn = unsafe { mem::transmute(code_buffer.ptr(prologue_offset)) };

        let detour = unsafe { RawDetour::new(addr as *const _, prelude as *const _)? };

        unsafe { detour.enable()? };

        zone.addr = Some(unsafe { mem::transmute::<*const (), WindowsFn>(detour.trampoline()) });
        zone.prologue = Some(prologue);

        Ok(ProfiledFunction {
            code_buffer,
            detour,
            name_cstr,
            zone,
        })
    }
}