Skip to main content

wazabin_qcode_parser/
ast.rs

1#[derive(Clone, Debug, PartialEq, Eq)]
2pub struct SourcePosition {
3    pub offset: usize,
4    pub line: usize,
5    pub column: usize,
6}
7
8#[derive(Clone, Debug, PartialEq, Eq)]
9pub struct SourceSpan {
10    pub start: SourcePosition,
11    pub end: SourcePosition,
12}
13
14impl SourceSpan {
15    pub fn contains_offset(&self, offset: usize) -> bool {
16        self.start.offset <= offset && offset < self.end.offset
17    }
18}
19
20#[derive(Clone, Debug)]
21pub enum Atom {
22    /// `{name}` — captures a Rust variable from the surrounding scope.
23    External(String),
24    /// `%name` — references an SSA instruction result.
25    Ssa(String),
26    /// `@name` — references a block parameter.
27    BlockParam(String),
28    /// bare `name` — references a varnode (valid only in pointer positions).
29    Varnode(String),
30    /// `&name` — takes the address of a varnode.
31    AddressOf(String),
32    Int(u64),
33    /// `true` / `false` — a byte-stored `bool` constant.
34    Bool(bool),
35}
36
37#[derive(Clone, Debug)]
38pub struct TypedAtom {
39    pub size_bytes: Option<usize>,
40    pub atom: Atom,
41    pub span: SourceSpan,
42}
43
44/// A direct function reference in textual QCode.
45#[derive(Clone, Debug, PartialEq, Eq)]
46pub enum Callee {
47    Named(String),
48    /// An unresolved pass-local function slot, written `<minted:N>`.
49    Minted(u32),
50}
51
52#[derive(Clone, Debug)]
53pub struct TupleField {
54    pub name: Option<String>,
55    pub value: TypedAtom,
56}
57
58#[derive(Clone, Debug)]
59pub enum ExtractField {
60    Name(String),
61    Index(u64),
62}
63
64/// The field selector of a `gep(...)` — by field name or by raw byte offset.
65#[derive(Clone, Debug)]
66pub enum GepField {
67    Name(String),
68    Offset(u64),
69}
70
71/// The declared type of a struct field: either a scalar of `n` bytes, or a
72/// pointer to a named (nominal) struct, written `Foo*`.
73#[derive(Clone, Debug)]
74pub enum StructFieldType {
75    Int(usize),
76    StructPtr(String),
77}
78
79/// One field of a `type Foo { ... }` declaration. A field named `_` is padding:
80/// it advances the running offset by its byte size without naming a slot.
81#[derive(Clone, Debug)]
82pub struct StructFieldDecl {
83    pub name: String,
84    pub ty: StructFieldType,
85}
86
87impl StructFieldDecl {
88    pub fn is_padding(&self) -> bool {
89        self.name == "_"
90    }
91}
92
93/// A nominal struct definition: `type Foo { a: 4, _: 5, b: 2 }`. Field offsets
94/// are the running byte sum (padding included); the struct `size` is the total.
95#[derive(Clone, Debug)]
96pub struct StructDecl {
97    pub name: String,
98    pub fields: Vec<StructFieldDecl>,
99    pub span: SourceSpan,
100}
101
102#[derive(Clone, Debug)]
103pub enum ExprNode {
104    Atom(TypedAtom),
105    Unop {
106        op: String,
107        src: TypedAtom,
108    },
109    Binary {
110        lhs: TypedAtom,
111        op: String,
112        rhs: TypedAtom,
113    },
114    Cast {
115        op: CastOp,
116        size_bytes: usize,
117        src: TypedAtom,
118    },
119    /// `load(space:size, ptr)` — read `size` bytes from address `ptr` in the
120    /// named `space`.
121    Load {
122        space: String,
123        size_bytes: usize,
124        ptr: TypedAtom,
125    },
126    /// `store(space:size, ptr <- value)` — write `value` (`size` bytes) to
127    /// address `ptr` in the named `space`.
128    Store {
129        space: String,
130        size_bytes: usize,
131        ptr: TypedAtom,
132        src: TypedAtom,
133    },
134    FuncCall {
135        op: String,
136        args: Vec<TypedAtom>,
137    },
138    /// A pure intrinsic call, e.g. `$rol(%x, %k)`. `name` excludes the `$`.
139    Intrinsic {
140        name: String,
141        args: Vec<TypedAtom>,
142    },
143    /// `apply lambda(args...)` — value-level application of a pure lambda.
144    Apply {
145        target: Callee,
146        args: Vec<TypedAtom>,
147    },
148    /// `body <$> src` / `(body c0 c1) <$> src` — an element-wise `map` over the
149    /// array `src`. `body` is a named function or an unresolved minted callee;
150    /// `captures` are the loop-invariant operands the body closes over.
151    Map {
152        body: Callee,
153        src: TypedAtom,
154        captures: Vec<TypedAtom>,
155    },
156    /// `scanl @body init src` / `scanl (@body c0 c1) init src` — a left-scan over
157    /// the array `src`. Named bodies are stored without the `@`; minted bodies
158    /// use their explicit placeholder. `init` is the initial accumulator;
159    /// `captures` are loop-invariant operands.
160    Scan {
161        body: Callee,
162        init: TypedAtom,
163        src: TypedAtom,
164        captures: Vec<TypedAtom>,
165    },
166    /// `pack(a=x, b=y)` — build an aggregate value from its named fields.
167    Tuple {
168        fields: Vec<TupleField>,
169    },
170    /// `extract(agg.field)` — project a field out of an aggregate value.
171    Extract {
172        agg: TypedAtom,
173        field: ExtractField,
174    },
175    /// `gep(base.field)` — compute the address of a struct field (typed, named
176    /// pointer arithmetic; no memory access).
177    Gep {
178        base: TypedAtom,
179        field: GepField,
180    },
181    /// `src[start:end]` — extract the byte range `[start, end)` of `src`. A
182    /// missing `start` defaults to 0; a missing `end` defaults to `src`'s width.
183    Range {
184        src: TypedAtom,
185        start: Option<u64>,
186        end: Option<u64>,
187    },
188}
189
190#[derive(Clone, Copy, Debug)]
191pub enum CastOp {
192    Zext,
193    Sext,
194    IntToFloat,
195    FloatToFloat,
196    Trunc,
197}
198
199#[derive(Clone, Debug)]
200pub struct BlockParamDecl {
201    pub name: String,
202    pub size_bytes: Option<usize>,
203}
204
205/// A branch target or label declaration — either a named label or a block address.
206#[derive(Clone, Debug)]
207pub enum Label {
208    /// A named label such as `<entry>` or `<done @v1 @v2>`. Generates a `BlockId` binding.
209    Named {
210        name: String,
211        /// Block parameters declared on this label (e.g. `@v1`, `@v2:i64`).
212        /// Non-empty only when this `Label` appears inside a `LabelDecl`.
213        params: Vec<BlockParamDecl>,
214        span: SourceSpan,
215    },
216    /// A numeric address such as `<0x1001>`. Sets the block's address; no binding generated.
217    Address { value: u64, span: SourceSpan },
218}
219
220impl Label {
221    pub fn span(&self) -> &SourceSpan {
222        match self {
223            Self::Named { span, .. } | Self::Address { span, .. } => span,
224        }
225    }
226
227    pub fn name(&self) -> Option<&str> {
228        match self {
229            Self::Named { name, .. } => Some(name),
230            Self::Address { .. } => None,
231        }
232    }
233}
234
235/// Per-parameter arguments on a branch edge: `(param_name, value)` in
236/// declaration order. Empty unless the edge was written as
237/// `goto <block @v1=e1 @v2=e2>`.
238pub type BlockArgs = Vec<(String, TypedAtom)>;
239
240/// One `switch` arm: `(case value, target block, arguments on that edge)`.
241pub type SwitchCase = (u64, Label, BlockArgs);
242
243#[derive(Clone, Debug)]
244pub enum Statement {
245    LocalDecl {
246        name: String,
247        name_span: SourceSpan,
248        display_name: String,
249        size_bytes: usize,
250        span: SourceSpan,
251    },
252    Assign {
253        name: String,
254        name_span: SourceSpan,
255        expr: ExprNode,
256        /// A `Foo*` struct-pointer type declared on the assignment, if any. When
257        /// present, the result value is retyped to that struct pointer (used to
258        /// seed struct typing in tests). A plain `iN`/`fN` declared type is not
259        /// recorded here — it only drives size coercion of the rhs.
260        decl_struct_ptr: Option<String>,
261        span: SourceSpan,
262    },
263    Expr(ExprNode),
264    LabelDecl {
265        label: Label,
266        span: SourceSpan,
267    },
268    Branch {
269        target: Label,
270        /// Per-parameter arguments: `(param_name, value)` in declaration order.
271        /// Non-empty when the branch was written as `goto <block @v1=e1 @v2=e2>`.
272        args: BlockArgs,
273        span: SourceSpan,
274    },
275    BranchInd {
276        ptr: TypedAtom,
277        /// Resolved jump targets, from a `// -> <a>, <b>` edge hint. The indirect
278        /// jump's own syntax encodes no successors, so without this the block has
279        /// no out-edges.
280        targets: Vec<Label>,
281        span: SourceSpan,
282    },
283    /// Multi-way dispatch: `switch %idx { 0x0 => <a>, default => <d> }`.
284    Switch {
285        scrutinee: TypedAtom,
286        /// `(case value, target, per-parameter arguments)` in written order.
287        cases: Vec<SwitchCase>,
288        /// The `default => <block>` arm, when written.
289        default: Option<(Label, BlockArgs)>,
290        span: SourceSpan,
291    },
292    CBranch {
293        condition: TypedAtom,
294        target: Label,
295        target_args: BlockArgs,
296        fallthrough: Label,
297        fallthrough_args: BlockArgs,
298        span: SourceSpan,
299    },
300    Call {
301        target: Callee,
302        /// `true` for `tailcall fn ...`, which has no return-edge hint.
303        tail: bool,
304        /// Arguments passed to the callee, one per inferred callee input, in
305        /// order. The name in each pair is the callee's parameter name as
306        /// printed (`@r0`, `@arg1`, …); it is decorative and discarded on
307        /// lowering, where only the positional atoms matter.
308        args: BlockArgs,
309        /// The call's return (fall-through) block(s), from a `// -> <ret>` edge
310        /// hint. A `call` terminates its block; this records where control resumes
311        /// after the callee returns. Empty for a non-returning call.
312        targets: Vec<Label>,
313        span: SourceSpan,
314    },
315    CallInd {
316        ptr: TypedAtom,
317        /// Positional arguments passed to the indirect callee.
318        args: Vec<TypedAtom>,
319        /// The call's return (fall-through) block(s), from a `// -> <ret>` edge
320        /// hint. See [`Statement::Call`].
321        targets: Vec<Label>,
322        span: SourceSpan,
323    },
324    Return {
325        ptr: TypedAtom,
326        value: Option<TypedAtom>,
327        span: SourceSpan,
328    },
329    ReturnValue {
330        value: TypedAtom,
331        span: SourceSpan,
332    },
333    /// Bytes that did not decode to a valid instruction; a terminator with no
334    /// successors and no operands.
335    BadInsn {
336        span: SourceSpan,
337    },
338    Assert {
339        condition: TypedAtom,
340        span: SourceSpan,
341    },
342    /// A comment attached to this statement, written as `# text` on the preceding line.
343    Commented {
344        comment: String,
345        inner: Box<Statement>,
346    },
347}
348
349impl Statement {
350    /// Strips any wrapping `Commented` variant and returns the inner statement.
351    pub fn inner(&self) -> &Statement {
352        match self {
353            Self::Commented { inner, .. } => inner.inner(),
354            other => other,
355        }
356    }
357}
358
359/// A function declaration (`fn name: <entry> stmts...`).
360#[derive(Clone, Debug)]
361pub struct FnDecl {
362    pub kind: FnKind,
363    pub name: String,
364    pub name_span: SourceSpan,
365    pub span: SourceSpan,
366    pub statements: Vec<Statement>,
367}
368
369#[derive(Clone, Copy, Debug, PartialEq, Eq)]
370pub enum FnKind {
371    Machine,
372    Lambda,
373}
374
375/// Top-level program representation. Any leading `type` declarations are
376/// collected into `structs`; `kind` is the statement or function body.
377#[derive(Clone, Debug)]
378pub struct Program {
379    pub structs: Vec<StructDecl>,
380    pub kind: ProgramKind,
381}
382
383#[derive(Clone, Debug)]
384pub enum ProgramKind {
385    Statements(Vec<Statement>),
386    Functions {
387        varnodes: Vec<Statement>,
388        fns: Vec<FnDecl>,
389    },
390}