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    #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
213    use super::*;
214    use crate::aarch64::instdb::InstId;
215    use crate::aarch64::*;
216    use crate::core::buffer::CodeBuffer;
217    use crate::core::operand::OperandCast;
218    use crate::core::target::Environment;
219    use std::vec::Vec;
220
221    fn a64_inst(id: u32, operands: &[crate::core::operand::Operand]) -> Inst {
222        Inst::with_arch_operands(crate::core::arch_traits::Arch::AArch64, id, operands).unwrap()
223    }
224
225    fn build_direct() -> Vec<u8> {
226        let mut buf = CodeBuffer::new(Environment::new(crate::core::arch_traits::Arch::AArch64));
227        let mut asm = Assembler::new(&mut buf);
228        let done = asm.buffer.get_label();
229
230        asm.stp(x29, x30, ptr(sp, 0).pre_offset(-32));
231        asm.mov(x29, sp);
232        asm.add(x0, x1, x2);
233        asm.cmp(x0, imm(0));
234        asm.b_eq(done);
235        asm.sub(x0, x0, imm(1));
236        asm.buffer.bind_label(done);
237        asm.ldp(x29, x30, ptr(sp, 0).post_offset(32));
238        asm.ret(lr);
239
240        buf.finish().unwrap().data().to_vec()
241    }
242
243    fn build_deferred() -> Vec<u8> {
244        let mut buf = CodeBuffer::new(Environment::new(crate::core::arch_traits::Arch::AArch64));
245        let mut builder = Builder::for_arch(crate::core::arch_traits::Arch::AArch64);
246        let done = buf.get_label();
247
248        builder
249            .push_inst(a64_inst(
250                InstId::Stp as u32,
251                &[
252                    *x29.as_operand(),
253                    *x30.as_operand(),
254                    *ptr(sp, 0).pre_offset(-32).as_operand(),
255                ],
256            ))
257            .unwrap();
258        builder
259            .push_inst(a64_inst(
260                InstId::Mov as u32,
261                &[*x29.as_operand(), *sp.as_operand()],
262            ))
263            .unwrap();
264        builder
265            .push_inst(a64_inst(
266                InstId::Add as u32,
267                &[*x0.as_operand(), *x1.as_operand(), *x2.as_operand()],
268            ))
269            .unwrap();
270        builder
271            .push_inst(a64_inst(
272                InstId::Cmp as u32,
273                &[*x0.as_operand(), *imm(0).as_operand()],
274            ))
275            .unwrap();
276        builder
277            .push_inst(a64_inst(
278                InstId::B.with_cc(crate::core::globals::CondCode::EQ),
279                &[*done.as_operand()],
280            ))
281            .unwrap();
282        builder
283            .push_inst(a64_inst(
284                InstId::Sub as u32,
285                &[*x0.as_operand(), *x0.as_operand(), *imm(1).as_operand()],
286            ))
287            .unwrap();
288        builder.push_label(done);
289        builder
290            .push_inst(a64_inst(
291                InstId::Ldp as u32,
292                &[
293                    *x29.as_operand(),
294                    *x30.as_operand(),
295                    *ptr(sp, 0).post_offset(32).as_operand(),
296                ],
297            ))
298            .unwrap();
299        builder
300            .push_inst(a64_inst(InstId::Ret as u32, &[*lr.as_operand()]))
301            .unwrap();
302
303        {
304            let mut asm = Assembler::new(&mut buf);
305            builder.emit_into(&mut asm).unwrap();
306        }
307        buf.finish().unwrap().data().to_vec()
308    }
309
310    #[test]
311    fn replay_matches_direct_assembly() {
312        assert_eq!(build_direct(), build_deferred());
313    }
314}
asmkit/core/builder.rs

asmkit/core/
builder.rs
