Skip to main content

patchouly_core/
stencils.rs

1use crate::relocation::{PatchInfo, Relocation};
2
3/// Describes the way to pass a variable in and out of a stencil
4pub struct StencilFamily<
5    const IN: usize,
6    const OUT: usize,
7    const MAX_REGS: usize,
8    const HOLES: usize,
9    const JUMPS: usize,
10> {
11    /// Pool of stencil relocation data, possibly shared by multiple stencils
12    ///
13    /// The relocation data for each stencil is sorted by relocation offset.
14    pub relocation_data: &'static [Relocation],
15    pub stencils: &'static [Stencil<IN, OUT, HOLES, JUMPS>],
16}
17
18#[derive(Clone, Copy, Default, Debug)]
19#[repr(C)]
20pub struct Stencil<const IN: usize, const OUT: usize, const HOLES: usize, const JUMPS: usize> {
21    pub code_index: u32,
22    pub code_len: u16,
23    /// Points to the first relocation in [StencilFamily::relocation_data]
24    ///
25    /// The end of the relocations are marked by a
26    /// `Relocation { encoding: Invalid }` (all zeroes) entry.
27    pub relocation_index: u16,
28}
29pub type UntypedStencil = Stencil<0, 0, 0, 0>;
30
31pub struct SelectedStencil<
32    const IN: usize,
33    const OUT: usize,
34    const HOLES: usize,
35    const JUMPS: usize,
36> {
37    pub wide: bool,
38    pub stencil: &'static Stencil<IN, OUT, HOLES, JUMPS>,
39}
40impl<const IN: usize, const OUT: usize, const HOLES: usize, const JUMPS: usize>
41    SelectedStencil<IN, OUT, HOLES, JUMPS>
42{
43    fn new(wide: bool, stencil: &'static Stencil<IN, OUT, HOLES, JUMPS>) -> Option<Self> {
44        if stencil.code_len == 0 {
45            None
46        } else {
47            Some(Self { wide, stencil })
48        }
49    }
50}
51
52#[derive(Clone, Copy, Debug, PartialEq, Eq)]
53pub enum Location {
54    Stack(u16),
55    Register(u16),
56}
57impl Location {
58    #[doc(hidden)]
59    pub fn from_bits(bits: u16) -> Self {
60        match bits {
61            0 => Location::Stack(0),
62            _ => Location::Register(bits - 1),
63        }
64    }
65
66    #[doc(hidden)]
67    pub fn into_bits(&self) -> u16 {
68        match self {
69            Location::Stack(_) => 0,
70            Location::Register(i) => i + 1,
71        }
72    }
73}
74
75impl<
76    const IN: usize,
77    const OUT: usize,
78    const MAX_REGS: usize,
79    const HOLES: usize,
80    const JUMPS: usize,
81> StencilFamily<IN, OUT, MAX_REGS, HOLES, JUMPS>
82{
83    pub fn inputs(&self) -> usize {
84        IN
85    }
86    pub fn outputs(&self) -> usize {
87        OUT
88    }
89    pub fn max_regs(&self) -> usize {
90        MAX_REGS
91    }
92    pub fn holes(&self) -> usize {
93        HOLES
94    }
95    pub fn jumps(&self) -> usize {
96        JUMPS
97    }
98
99    pub fn select(
100        &self,
101        inputs: &[Location; IN],
102        outputs: &[Location; OUT],
103        holes: &[usize; HOLES],
104    ) -> Option<SelectedStencil<IN, OUT, HOLES, JUMPS>> {
105        let index = io_to_index(inputs, outputs, MAX_REGS, false);
106        let stencil = &self.stencils[index];
107        for reloc in stencil.relocations(self) {
108            if reloc.is_invalid() {
109                return SelectedStencil::new(false, stencil);
110            }
111            if let Some(PatchInfo::Hole(i)) = reloc.patch_info()
112                && !reloc.supports_value(holes[i as usize])
113            {
114                break;
115            }
116        }
117        SelectedStencil::new(
118            true,
119            &self.stencils[index + stencils_len(IN, OUT, MAX_REGS)],
120        )
121    }
122}
123
124#[doc(hidden)]
125pub fn io_to_index(
126    inputs: &[Location],
127    outputs: &[Location],
128    max_regs: usize,
129    wide: bool,
130) -> usize {
131    let mut i = 0;
132    for var in inputs.iter().chain(outputs.iter()) {
133        i = max_regs * i
134            + match var {
135                Location::Stack(_) => 0,
136                Location::Register(i) => *i as usize + 1,
137            };
138    }
139    i + if wide {
140        stencils_len(inputs.len(), outputs.len(), max_regs)
141    } else {
142        0
143    }
144}
145
146#[doc(hidden)]
147pub fn index_to_io_lossy(
148    index: usize,
149    max_regs: usize,
150    inputs: &mut [Location],
151    outputs: &mut [Location],
152) -> bool {
153    fn process_index(mut index: usize, max_regs: usize, slots: &mut [Location]) -> usize {
154        for v in slots.iter_mut().rev() {
155            let reg = index % max_regs;
156            index = (index - reg) / max_regs;
157            *v = if reg == 0 {
158                Location::Stack(0)
159            } else {
160                Location::Register(reg as u16 - 1)
161            };
162        }
163        index
164    }
165
166    let index = process_index(index, max_regs, outputs);
167    let index = process_index(index, max_regs, inputs);
168    assert!(index == 0 || index == 1);
169    index == 1
170}
171
172#[doc(hidden)]
173pub const fn stencils_len(inputs: usize, outputs: usize, max_regs: usize) -> usize {
174    max_regs.pow(inputs as u32 + outputs as u32)
175}
176
177impl<const IN: usize, const OUT: usize, const HOLES: usize, const JUMPS: usize>
178    Stencil<IN, OUT, HOLES, JUMPS>
179{
180    #[doc(hidden)]
181    pub const fn from_bits(bits: u64) -> Self {
182        Stencil {
183            code_index: (bits & 0xFFFFFFFF) as u32,
184            code_len: ((bits >> 32) & 0xFFFF) as u16,
185            relocation_index: ((bits >> 48) & 0xFFFF) as u16,
186        }
187    }
188
189    #[doc(hidden)]
190    pub const fn into_bits(self) -> u64 {
191        (self.code_index as u64)
192            | ((self.code_len as u64) << 32)
193            | ((self.relocation_index as u64) << 48)
194    }
195
196    pub fn code<'a>(&self, store: &'a [u8]) -> &'a [u8] {
197        &store[self.code_index as usize..self.code_index as usize + self.code_len as usize]
198    }
199
200    pub fn relocations<'a, const MAX_REGS: usize>(
201        &self,
202        store: &'a StencilFamily<IN, OUT, MAX_REGS, HOLES, JUMPS>,
203    ) -> &'a [Relocation] {
204        &store.relocation_data[self.relocation_index as usize..]
205    }
206
207    pub fn untyped(&self) -> UntypedStencil {
208        UntypedStencil::from_bits(self.into_bits())
209    }
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215
216    type TestStencil = Stencil<0, 0, 0, 0>;
217
218    #[test]
219    fn test_stencil_bits() {
220        let mut i = 1u64;
221        for _ in 0..10000 {
222            i = i.wrapping_mul(31);
223            assert_eq!(TestStencil::from_bits(i).into_bits(), i);
224        }
225    }
226
227    #[test]
228    fn test_io_to_index() {
229        let mut i = 1usize;
230        let len = stencils_len(4, 4, 10);
231        for _ in 0..10000 {
232            i = i.wrapping_mul(31);
233            let mut inputs = [Location::Stack(0); 4];
234            let mut outputs = [Location::Stack(0); 4];
235            index_to_io_lossy(i % len, 10, &mut inputs, &mut outputs);
236            assert_eq!(io_to_index(&inputs, &outputs, 10, false), i % len);
237        }
238    }
239}