math_jit/
library.rs

1//! Management of functions accessible to programs
2
3use std::mem::transmute;
4
5/// Description of a function accessible to the compiled program
6///
7/// Currently only functions of the form `fn(f32, [f32, ..]) -> f32` are
8/// supported.
9pub struct FunctionF32 {
10    /// Name of the function exposed to the program
11    pub name: String,
12    /// Pointer to the function
13    pub(crate) ptr: *const u8,
14    /// Number of arguments of the function
15    pub(crate) param_count: usize,
16}
17
18impl FunctionF32 {
19    /// Constructs a new instance of 1-argument [`FunctionF32`]
20    pub fn new_1(name: String, ptr: extern "C" fn(f32) -> f32) -> Self {
21        FunctionF32 {
22            name,
23            ptr: ptr as *const u8,
24            param_count: 1,
25        }
26    }
27
28    /// Constructs a new instance of 2-argument [`FunctionF32`]
29    pub fn new_2(name: String, ptr: extern "C" fn(f32, f32) -> f32) -> Self {
30        FunctionF32 {
31            name,
32            ptr: ptr as *const u8,
33            param_count: 2,
34        }
35    }
36
37    /// Call this function with 1 argument
38    ///
39    /// Returns None if parameter count is incorrect.
40    pub fn call_1(&self, x: f32) -> Option<f32> {
41        if self.param_count != 1 {
42            None
43        } else {
44            // SAFETY: param_count proves that ptr was instantiated from a
45            // function of this signature.
46            Some(unsafe { transmute::<_, extern "C" fn(f32) -> f32>(self.ptr)(x) })
47        }
48    }
49
50    /// Call this function with 2 arguments
51    ///
52    /// Returns None if parameter count is incorrect.
53    pub fn call_2(&self, x: f32, y: f32) -> Option<f32> {
54        if self.param_count != 2 {
55            None
56        } else {
57            // SAFETY: param_count proves that ptr was instantiated from a
58            // function of this signature.
59            Some(unsafe { transmute::<_, extern "C" fn(f32, f32) -> f32>(self.ptr)(x, y) })
60        }
61    }
62}
63
64/// Library of functions accessible to the program
65pub struct Library {
66    funs: Vec<FunctionF32>,
67}
68
69impl Library {
70    /// Append a new function to the library
71    pub fn insert(&mut self, fun: FunctionF32) {
72        self.funs.push(fun);
73    }
74
75    /// Iterator over the functions
76    pub fn iter(&self) -> impl Iterator<Item = &'_ FunctionF32> {
77        self.funs.iter()
78    }
79}
80
81/// Default implementation of the library
82///
83/// Provider `sin(x)`, `cos(x)` and `pow(a, b)`, as those functions are
84/// accessible to all programs.
85impl Default for Library {
86    fn default() -> Self {
87        extern "C" fn sin_(x: f32) -> f32 {
88            x.sin()
89        }
90
91        extern "C" fn cos_(x: f32) -> f32 {
92            x.cos()
93        }
94
95        extern "C" fn tan_(x: f32) -> f32 {
96            x.tan()
97        }
98
99        extern "C" fn abs_(x: f32) -> f32 {
100            x.abs()
101        }
102
103        extern "C" fn sqrt_(x: f32) -> f32 {
104            x.sqrt()
105        }
106
107        extern "C" fn pow_(a: f32, b: f32) -> f32 {
108            a.powf(b)
109        }
110
111        Library {
112            funs: vec![
113                FunctionF32::new_1("sin".into(), sin_),
114                FunctionF32::new_1("cos".into(), cos_),
115                FunctionF32::new_1("tan".into(), tan_),
116                FunctionF32::new_1("abs".into(), abs_),
117                FunctionF32::new_1("sqrt".into(), sqrt_),
118                FunctionF32::new_2("pow".into(), pow_),
119            ],
120        }
121    }
122}
math_jit/library.rs

math_jit/
library.rs
