asmkit/core/
builder.rs

1//! Deferred instruction builder.
2//!
3//! A [`Builder`] records a sequence of nodes (instructions and label-bind points) so passes
4//! can inspect and mutate them — most importantly a future register-allocation pass — before
5//! machine code is produced. Replaying a builder into an [`InstSink`] (implemented by each
6//! architecture's `Assembler`) emits the exact same bytes as direct assembly: labels and
7//! relocations are recorded at emit time and resolved by `CodeBuffer::finish()` as usual.
8
9use smallvec::SmallVec;
10
11use crate::AsmError;
12
13use super::arch_traits::Arch;
14use super::inst::Inst;
15use super::operand::Label;
16
17/// A node recorded by a [`Builder`].
18#[derive(Clone, Copy, Debug)]
19pub enum Node {
20    /// An instruction with its operands.
21    Inst(Inst),
22    /// Binds a label at this position when replayed.
23    Label(Label),
24}
25
26/// Sink that consumes replayed nodes — implemented by each architecture's `Assembler`.
27pub trait InstSink {
28    /// Target architecture accepted by this sink.
29    fn arch(&self) -> Arch;
30
31    /// Emits one recorded instruction.
32    fn emit_inst(&mut self, inst: &Inst) -> Result<(), AsmError>;
33
34    /// Binds a label at the current position.
35    fn bind_label(&mut self, label: Label) -> Result<(), AsmError>;
36}
37
38/// Records instructions and label-bind points for deferred emission.
39#[derive(Clone, Debug)]
40pub struct Builder {
41    arch: Arch,
42    nodes: SmallVec<[Node; 32]>,
43}
44
45impl Default for Builder {
46    fn default() -> Self {
47        Self::new()
48    }
49}
50
51impl Builder {
52    /// Creates an empty builder.
53    pub fn new() -> Self {
54        Self::for_arch(Arch::HOST)
55    }
56
57    /// Creates a builder that accepts instructions for `arch` only.
58    pub fn for_arch(arch: Arch) -> Self {
59        Self {
60            arch,
61            nodes: SmallVec::new(),
62        }
63    }
64
65    /// Architecture accepted by this builder.
66    pub const fn arch(&self) -> Arch {
67        self.arch
68    }
69
70    /// Returns the number of recorded nodes.
71    pub fn len(&self) -> usize {
72        self.nodes.len()
73    }
74
75    /// Tests whether the builder is empty.
76    pub fn is_empty(&self) -> bool {
77        self.nodes.is_empty()
78    }
79
80    /// Removes all recorded nodes.
81    pub fn clear(&mut self) {
82        self.nodes.clear();
83    }
84
85    /// Records an instruction.
86    pub fn push_inst(&mut self, inst: Inst) -> Result<(), AsmError> {
87        if inst.arch() != self.arch {
88            return Err(AsmError::InvalidArch);
89        }
90        self.nodes.push(Node::Inst(inst));
91        Ok(())
92    }
93
94    /// Records a label-bind point.
95    pub fn push_label(&mut self, label: Label) {
96        self.nodes.push(Node::Label(label));
97    }
98
99    /// Returns the recorded nodes.
100    pub fn nodes(&self) -> &[Node] {
101        &self.nodes
102    }
103
104    /// Replaces an instruction node after checking that it belongs to this builder.
105    pub fn replace_inst(&mut self, index: usize, inst: Inst) -> Result<(), AsmError> {
106        if inst.arch() != self.arch {
107            return Err(AsmError::InvalidArch);
108        }
109        let Some(node) = self.nodes.get_mut(index) else {
110            return Err(AsmError::InvalidArgument);
111        };
112        if !matches!(node, Node::Inst(_)) {
113            return Err(AsmError::InvalidState);
114        }
115        *node = Node::Inst(inst);
116        Ok(())
117    }
118
119    /// Replays all recorded nodes into `sink`, in order.
120    pub fn emit_into<S: InstSink + ?Sized>(&self, sink: &mut S) -> Result<(), AsmError> {
121        if sink.arch() != self.arch {
122            return Err(AsmError::InvalidArch);
123        }
124        for node in self.nodes.iter() {
125            match node {
126                Node::Inst(inst) => sink.emit_inst(inst)?,
127                Node::Label(label) => sink.bind_label(*label)?,
128            }
129        }
130        Ok(())
131    }
132}
133
134#[cfg(test)]
135mod builder_tests {
136    use super::*;
137
138    struct TestSink {
139        arch: Arch,
140        calls: usize,
141        fail: bool,
142    }
143
144    impl InstSink for TestSink {
145        fn arch(&self) -> Arch {
146            self.arch
147        }
148
149        fn emit_inst(&mut self, _: &Inst) -> Result<(), AsmError> {
150            self.calls += 1;
151            if self.fail {
152                Err(AsmError::InvalidOperand)
153            } else {
154                Ok(())
155            }
156        }
157
158        fn bind_label(&mut self, _: Label) -> Result<(), AsmError> {
159            self.calls += 1;
160            Ok(())
161        }
162    }
163
164    #[test]
165    fn wrong_architecture_stops_before_replay() {
166        let mut builder = Builder::for_arch(Arch::AArch64);
167        builder
168            .push_inst(Inst::with_arch_operands(Arch::AArch64, 1, &[]).unwrap())
169            .unwrap();
170        let mut sink = TestSink {
171            arch: Arch::X64,
172            calls: 0,
173            fail: false,
174        };
175
176        assert_eq!(builder.emit_into(&mut sink), Err(AsmError::InvalidArch));
177        assert_eq!(sink.calls, 0);
178    }
179
180    #[test]
181    fn replay_stops_at_the_first_failure() {
182        let mut builder = Builder::for_arch(Arch::X64);
183        let inst = Inst::with_arch_operands(Arch::X64, 1, &[]).unwrap();
184        builder.push_inst(inst).unwrap();
185        builder.push_label(Label::from_id(0));
186        builder.push_inst(inst).unwrap();
187        let mut sink = TestSink {
188            arch: Arch::X64,
189            calls: 0,
190            fail: true,
191        };
192
193        assert_eq!(builder.emit_into(&mut sink), Err(AsmError::InvalidOperand));
194        assert_eq!(sink.calls, 1);
195    }
196
197    #[test]
198    fn replacement_checks_the_instruction_architecture() {
199        let mut builder = Builder::for_arch(Arch::RISCV64);
200        builder
201            .push_inst(Inst::with_arch_operands(Arch::RISCV64, 1, &[]).unwrap())
202            .unwrap();
203        assert_eq!(
204            builder.replace_inst(0, Inst::with_arch_operands(Arch::AArch64, 1, &[]).unwrap()),
205            Err(AsmError::InvalidArch)
206        );
207    }
208}
209
210#[cfg(all(test, feature = "aarch64"))]
211mod tests {
212    use super::*;
213    use crate::aarch64::instdb::InstId;
214    use crate::aarch64::*;
215    use crate::core::buffer::CodeBuffer;
216    use crate::core::operand::OperandCast;
217    use crate::core::target::Environment;
218    use std::vec::Vec;
219
220    fn a64_inst(id: u32, operands: &[crate::core::operand::Operand]) -> Inst {
221        Inst::with_arch_operands(crate::core::arch_traits::Arch::AArch64, id, operands).unwrap()
222    }
223
224    fn build_direct() -> Vec<u8> {
225        let mut buf = CodeBuffer::new(Environment::new(crate::core::arch_traits::Arch::AArch64));
226        let mut asm = Assembler::new(&mut buf);
227        let done = asm.buffer.get_label();
228
229        asm.stp(x29, x30, ptr(sp, 0).pre_offset(-32));
230        asm.mov(x29, sp);
231        asm.add(x0, x1, x2);
232        asm.cmp(x0, imm(0));
233        asm.b_eq(done);
234        asm.sub(x0, x0, imm(1));
235        asm.buffer.bind_label(done);
236        asm.ldp(x29, x30, ptr(sp, 0).post_offset(32));
237        asm.ret(lr);
238
239        buf.finish().unwrap().data().to_vec()
240    }
241
242    fn build_deferred() -> Vec<u8> {
243        let mut buf = CodeBuffer::new(Environment::new(crate::core::arch_traits::Arch::AArch64));
244        let mut builder = Builder::for_arch(crate::core::arch_traits::Arch::AArch64);
245        let done = buf.get_label();
246
247        builder
248            .push_inst(a64_inst(
249                InstId::Stp as u32,
250                &[
251                    *x29.as_operand(),
252                    *x30.as_operand(),
253                    *ptr(sp, 0).pre_offset(-32).as_operand(),
254                ],
255            ))
256            .unwrap();
257        builder
258            .push_inst(a64_inst(
259                InstId::Mov as u32,
260                &[*x29.as_operand(), *sp.as_operand()],
261            ))
262            .unwrap();
263        builder
264            .push_inst(a64_inst(
265                InstId::Add as u32,
266                &[*x0.as_operand(), *x1.as_operand(), *x2.as_operand()],
267            ))
268            .unwrap();
269        builder
270            .push_inst(a64_inst(
271                InstId::Cmp as u32,
272                &[*x0.as_operand(), *imm(0).as_operand()],
273            ))
274            .unwrap();
275        builder
276            .push_inst(a64_inst(
277                InstId::B.with_cc(crate::core::globals::CondCode::EQ),
278                &[*done.as_operand()],
279            ))
280            .unwrap();
281        builder
282            .push_inst(a64_inst(
283                InstId::Sub as u32,
284                &[*x0.as_operand(), *x0.as_operand(), *imm(1).as_operand()],
285            ))
286            .unwrap();
287        builder.push_label(done);
288        builder
289            .push_inst(a64_inst(
290                InstId::Ldp as u32,
291                &[
292                    *x29.as_operand(),
293                    *x30.as_operand(),
294                    *ptr(sp, 0).post_offset(32).as_operand(),
295                ],
296            ))
297            .unwrap();
298        builder
299            .push_inst(a64_inst(InstId::Ret as u32, &[*lr.as_operand()]))
300            .unwrap();
301
302        {
303            let mut asm = Assembler::new(&mut buf);
304            builder.emit_into(&mut asm).unwrap();
305        }
306        buf.finish().unwrap().data().to_vec()
307    }
308
309    #[test]
310    fn replay_matches_direct_assembly() {
311        assert_eq!(build_direct(), build_deferred());
312    }
313}
asmkit/core/builder.rs

asmkit/core/
builder.rs
