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

use core::{
    marker::PhantomData,
    mem::transmute,
    ops::Deref,
    sync::atomic::{AtomicUsize, Ordering},
};

/// Creates a function defenition with explicitly defined offset from module or signature.
/// ```
/// # use memflex::ResolveBy;
/// fn get_address_in_module<const N: usize>(_: ResolveBy<N>) -> usize {
///     todo!()
/// }
///
/// memflex::function! {
///     extern fn CALLE(i32, i32) -> i32 = "file.exe"#0x1122;
///
///     // Use `get_address_in_module` function to get the address duh
///     extern fn OTHER_FN(f32, bool) -> u64 = (get_address_in_module)"file.exe"#0x1122;
///
///     // Resolve by signature
///     extern fn SIG_FN() -> u64 = "file.exe"%"4A 5B 3C AA 14";
/// }
/// ```
#[macro_export]
macro_rules! function {
    (
        $(
            $(extern $($abi:literal)?)? fn $fname:ident( $($atype:ty),* ) $(-> $ret:ty)? = $( ($resolver:ident) )? $modname:literal $sep:tt $offset:expr;
        )*
    ) => {
        $(
            static $fname: $crate::Function< $(extern $($abi)?)? fn($($atype),*) $(-> $ret)?> = $crate::Function::new(
                || unsafe {
                    ($crate::__resolver!( $($resolver)? ))( $crate::__resolve_by!($sep $modname, $offset ) )
                }
            );
        )*
    };
}

/// Function that resolve its address on the first access.
pub struct Function<F> {
    address: AtomicUsize,
    init: fn() -> usize,
    _ph: PhantomData<F>,
}

impl<F> Function<F> {
    #[doc(hidden)]
    pub const fn new(init: fn() -> usize) -> Self {
        Self {
            address: AtomicUsize::new(0),
            init,
            _ph: PhantomData,
        }
    }

    /// Returns function address
    pub fn address(&self) -> usize {
        let mut address = self.address.load(Ordering::Acquire);
        if address == 0 {
            address = (self.init)();
            self.address.store(address, Ordering::Release);
        }

        address
    }

    /// Force inits function
    pub fn force(&self) {
        _ = self.address();
    }
}

impl<F> Deref for Function<F> {
    type Target = F;

    fn deref(&self) -> &Self::Target {
        self.force();
        unsafe { transmute(self) }
    }
}