Skip to main content

vm/compiler/
ir.rs

1use std::collections::HashMap;
2
3use crate::ValueType;
4use crate::builtins::default_host_callable;
5
6use super::ParseError;
7
8pub type LocalSlot = u16;
9
10#[derive(Clone, Debug, PartialEq, Eq)]
11pub enum TypeSchema {
12    Unknown,
13    Null,
14    Int,
15    Float,
16    Number,
17    Bool,
18    String,
19    Bytes,
20    Optional(Box<TypeSchema>),
21    GenericParam(String),
22    Named(String, Vec<TypeSchema>),
23    Array(Box<TypeSchema>),
24    ArrayTuple(Vec<TypeSchema>),
25    ArrayTupleRest {
26        prefix: Vec<TypeSchema>,
27        rest: Box<TypeSchema>,
28    },
29    Map(Box<TypeSchema>),
30    Object(HashMap<String, TypeSchema>),
31    Callable {
32        params: Vec<TypeSchema>,
33        result: Box<TypeSchema>,
34    },
35}
36
37impl TypeSchema {
38    pub(crate) fn is_optional(&self) -> bool {
39        matches!(self, TypeSchema::Optional(_))
40    }
41
42    pub(crate) fn clone_inner_if_optional(&self) -> TypeSchema {
43        match self {
44            TypeSchema::Optional(inner) => inner.as_ref().clone(),
45            other => other.clone(),
46        }
47    }
48
49    pub(crate) fn split_optional(&self) -> (TypeSchema, bool) {
50        match self {
51            TypeSchema::Optional(inner) => (inner.as_ref().clone(), true),
52            other => (other.clone(), false),
53        }
54    }
55
56    pub(crate) fn coarse_value_type(&self) -> ValueType {
57        match self {
58            TypeSchema::Unknown
59            | TypeSchema::GenericParam(_)
60            | TypeSchema::Number
61            | TypeSchema::Callable { .. } => ValueType::Unknown,
62            TypeSchema::Null => ValueType::Null,
63            TypeSchema::Int => ValueType::Int,
64            TypeSchema::Float => ValueType::Float,
65            TypeSchema::Bool => ValueType::Bool,
66            TypeSchema::String => ValueType::String,
67            TypeSchema::Bytes => ValueType::Bytes,
68            TypeSchema::Optional(inner) => inner.coarse_value_type(),
69            TypeSchema::Named(_, _) | TypeSchema::Map(_) | TypeSchema::Object(_) => ValueType::Map,
70            TypeSchema::Array(_)
71            | TypeSchema::ArrayTuple(_)
72            | TypeSchema::ArrayTupleRest { .. } => ValueType::Array,
73        }
74    }
75
76    pub(crate) fn array_prefix_and_rest(&self) -> Option<(&[TypeSchema], Option<&TypeSchema>)> {
77        match self {
78            TypeSchema::Array(element) => Some((&[], Some(element.as_ref()))),
79            TypeSchema::ArrayTuple(items) => Some((items.as_slice(), None)),
80            TypeSchema::ArrayTupleRest { prefix, rest } => {
81                Some((prefix.as_slice(), Some(rest.as_ref())))
82            }
83            _ => None,
84        }
85    }
86
87    pub(crate) fn array_item_schema_at(&self, index: usize) -> Option<TypeSchema> {
88        let (prefix, rest) = self.array_prefix_and_rest()?;
89        prefix.get(index).cloned().or_else(|| rest.cloned())
90    }
91
92    pub(crate) fn collapsed_array_item_schema(&self) -> Option<TypeSchema> {
93        let (prefix, rest) = self.array_prefix_and_rest()?;
94        let mut items = prefix.iter();
95        let Some(first) = items.next().cloned().or_else(|| rest.cloned()) else {
96            return Some(TypeSchema::Unknown);
97        };
98        if items.all(|schema| schema == &first) && rest.is_none_or(|schema| schema == &first) {
99            Some(first)
100        } else {
101            Some(TypeSchema::Unknown)
102        }
103    }
104}
105
106#[derive(Clone, Debug, PartialEq, Eq)]
107pub struct FunctionParam {
108    pub name: String,
109    pub schema: Option<TypeSchema>,
110}
111
112#[derive(Clone, Debug, PartialEq, Eq)]
113pub struct StructDecl {
114    pub name: String,
115    pub type_params: Vec<String>,
116    pub body_schema: TypeSchema,
117}
118
119fn known_host_accepts_arity(name: &str, arity: u8) -> bool {
120    #[cfg(feature = "edge-abi")]
121    if let Some(function) = edge_abi::function_by_name(name) {
122        return function.param_types.len() == usize::from(arity);
123    }
124    default_host_callable(name).is_some_and(|callable| {
125        let required = callable
126            .signature
127            .params
128            .iter()
129            .take_while(|param| !param.optional)
130            .count();
131        required <= usize::from(arity) && usize::from(arity) <= callable.signature.params.len()
132    })
133}
134
135/// Shared frontend-independent program representation that all source
136/// frontends lower into before bytecode emission.
137#[derive(Clone, Debug)]
138pub struct ClosureExpr {
139    pub param_slots: Vec<LocalSlot>,
140    pub capture_copies: Vec<(LocalSlot, LocalSlot)>,
141    pub body: Box<Expr>,
142}
143
144#[derive(Clone, Debug)]
145pub enum MatchPattern {
146    Int(i64),
147    String(String),
148    Bytes(Vec<u8>),
149    Null,
150    None,
151    SomeBinding(LocalSlot),
152    Type(MatchTypePattern),
153}
154
155impl MatchPattern {
156    pub(crate) fn binding_slot(&self) -> Option<LocalSlot> {
157        match self {
158            MatchPattern::SomeBinding(slot) => Some(*slot),
159            _ => None,
160        }
161    }
162
163    pub(crate) fn requires_optional_value(&self) -> bool {
164        matches!(self, MatchPattern::None | MatchPattern::SomeBinding(_))
165    }
166}
167
168#[derive(Clone, Debug)]
169pub enum MatchTypePattern {
170    Int,
171    Float,
172    Number,
173    Bool,
174    String,
175    Bytes,
176    Array,
177    Map,
178}
179
180#[derive(Clone, Debug)]
181pub enum Expr {
182    Null,
183    Int(i64),
184    Float(f64),
185    Bool(bool),
186    String(String),
187    Bytes(Vec<u8>),
188    FunctionRef(u16),
189    OptionalGet {
190        container: Box<Expr>,
191        key: Box<Expr>,
192        container_slot: LocalSlot,
193        key_slot: LocalSlot,
194    },
195    OptionUnwrapOr {
196        value: Box<Expr>,
197        value_slot: LocalSlot,
198        fallback: Box<Expr>,
199    },
200    Call(u16, Vec<TypeSchema>, Vec<Expr>),
201    LocalCall(LocalSlot, Vec<TypeSchema>, Vec<Expr>),
202    Closure(ClosureExpr),
203    ClosureCall(ClosureExpr, Vec<Expr>),
204    Add(Box<Expr>, Box<Expr>),
205    Sub(Box<Expr>, Box<Expr>),
206    Mul(Box<Expr>, Box<Expr>),
207    Div(Box<Expr>, Box<Expr>),
208    Mod(Box<Expr>, Box<Expr>),
209    Neg(Box<Expr>),
210    Not(Box<Expr>),
211    And(Box<Expr>, Box<Expr>),
212    Or(Box<Expr>, Box<Expr>),
213    Eq(Box<Expr>, Box<Expr>),
214    Lt(Box<Expr>, Box<Expr>),
215    Gt(Box<Expr>, Box<Expr>),
216    Var(LocalSlot),
217    MoveVar(LocalSlot),
218    MoveField {
219        root: LocalSlot,
220        key: String,
221    },
222    MoveIndex {
223        root: LocalSlot,
224        index: i64,
225    },
226    IfElse {
227        condition: Box<Expr>,
228        then_expr: Box<Expr>,
229        else_expr: Box<Expr>,
230    },
231    Match {
232        value_slot: LocalSlot,
233        result_slot: LocalSlot,
234        value: Box<Expr>,
235        arms: Vec<(MatchPattern, Expr)>,
236        default: Box<Expr>,
237    },
238    ToOwned(Box<Expr>),
239    Borrow(Box<Expr>),
240    BorrowMut(Box<Expr>),
241    Block {
242        stmts: Vec<Stmt>,
243        expr: Box<Expr>,
244    },
245}
246
247#[derive(Clone, Debug)]
248pub enum AssignmentKind {
249    Set,
250    Add,
251    Increment,
252}
253
254impl AssignmentKind {
255    pub(crate) fn requires_numeric_operands(&self) -> bool {
256        matches!(self, Self::Add | Self::Increment)
257    }
258
259    pub(crate) fn diagnostic_label(&self) -> &'static str {
260        match self {
261            Self::Set => "'=' assignment",
262            Self::Add => "'+=' assignment",
263            Self::Increment => "'++' increment",
264        }
265    }
266}
267
268#[derive(Clone, Debug)]
269pub enum Stmt {
270    Noop {
271        line: u32,
272    },
273    Let {
274        index: LocalSlot,
275        declared_schema: Option<TypeSchema>,
276        expr: Expr,
277        line: u32,
278    },
279    Assign {
280        kind: AssignmentKind,
281        index: LocalSlot,
282        expr: Expr,
283        line: u32,
284    },
285    ClosureLet {
286        line: u32,
287        closure: ClosureExpr,
288    },
289    FuncDecl {
290        name: String,
291        index: u16,
292        arity: u8,
293        args: Vec<String>,
294        exported: bool,
295        has_impl: bool,
296        line: u32,
297    },
298    Expr {
299        expr: Expr,
300        line: u32,
301    },
302    IfElse {
303        condition: Expr,
304        then_branch: Vec<Stmt>,
305        else_branch: Vec<Stmt>,
306        line: u32,
307    },
308    For {
309        init: Box<Stmt>,
310        condition: Expr,
311        post: Box<Stmt>,
312        body: Vec<Stmt>,
313        line: u32,
314    },
315    While {
316        condition: Expr,
317        body: Vec<Stmt>,
318        line: u32,
319    },
320    Break {
321        line: u32,
322    },
323    Continue {
324        line: u32,
325    },
326    /// Explicit compile-time drop: null-out the local slot and trigger the
327    /// runtime drop-contract for whatever value was previously stored there.
328    Drop {
329        index: LocalSlot,
330        line: u32,
331    },
332}
333
334#[derive(Clone, Debug, PartialEq, Eq)]
335pub struct FunctionDecl {
336    pub name: String,
337    pub arity: u8,
338    pub index: u16,
339    pub args: Vec<String>,
340    pub arg_schemas: Vec<Option<TypeSchema>>,
341    pub return_schema: Option<TypeSchema>,
342    pub type_params: Vec<String>,
343    pub exported: bool,
344    pub return_type: ValueType,
345}
346
347#[derive(Clone, Debug)]
348pub struct FunctionImpl {
349    pub param_slots: Vec<LocalSlot>,
350    pub capture_copies: Vec<(LocalSlot, LocalSlot)>,
351    pub body_stmts: Vec<Stmt>,
352    pub body_expr: Expr,
353    pub body_expr_line: u32,
354}
355
356#[derive(Clone, Debug)]
357pub struct FrontendIr {
358    pub stmts: Vec<Stmt>,
359    pub locals: usize,
360    pub local_bindings: Vec<(String, LocalSlot)>,
361    pub struct_schemas: HashMap<String, StructDecl>,
362    pub unknown_type_spans: Vec<crate::compiler::source_map::Span>,
363    pub functions: Vec<FunctionDecl>,
364    pub function_impls: HashMap<u16, FunctionImpl>,
365    pub stmt_sources: Vec<Option<String>>,
366    pub function_sources: HashMap<u16, String>,
367}
368
369pub struct LocalIrBuilder {
370    locals: HashMap<String, LocalSlot>,
371    next_local: LocalSlot,
372    functions: Vec<FunctionDecl>,
373    function_meta: HashMap<String, (u16, Option<u8>)>,
374}
375
376impl Default for LocalIrBuilder {
377    fn default() -> Self {
378        Self::new()
379    }
380}
381
382impl LocalIrBuilder {
383    pub fn new() -> Self {
384        Self {
385            locals: HashMap::new(),
386            next_local: 0,
387            functions: Vec::new(),
388            function_meta: HashMap::new(),
389        }
390    }
391
392    pub fn lower_local(&mut self, name: &str, expr: Expr, line: u32) -> Result<Stmt, ParseError> {
393        let index = self.alloc_local_named(name)?;
394        Ok(Stmt::Let {
395            index,
396            declared_schema: None,
397            expr,
398            line,
399        })
400    }
401
402    pub fn lower_assign(&self, name: &str, expr: Expr, line: u32) -> Result<Stmt, ParseError> {
403        let Some(index) = self.locals.get(name).copied() else {
404            return Err(ParseError {
405                span: None,
406                code: None,
407                line: line as usize,
408                message: format!("unknown local '{name}'"),
409            });
410        };
411        Ok(Stmt::Assign {
412            kind: AssignmentKind::Set,
413            index,
414            expr,
415            line,
416        })
417    }
418
419    pub fn resolve_local_expr(&self, name: &str) -> Option<Expr> {
420        self.locals.get(name).copied().map(Expr::Var)
421    }
422
423    pub fn has_declared_function(&self, name: &str) -> bool {
424        self.function_meta.contains_key(name)
425    }
426
427    pub fn function_index(&self, name: &str) -> Option<u16> {
428        self.function_meta.get(name).map(|(index, _)| *index)
429    }
430
431    pub fn declare_function(&mut self, name: &str, arity: Option<u8>) -> Result<(), ParseError> {
432        if let Some((index, existing_arity)) = self.function_meta.get(name).copied() {
433            match (existing_arity, arity) {
434                (Some(expected), Some(actual))
435                    if expected != actual && !known_host_accepts_arity(name, actual) =>
436                {
437                    return Err(ParseError {
438                        span: None,
439                        code: None,
440                        line: 1,
441                        message: format!(
442                            "function '{name}' declared with conflicting arity {expected} vs {actual}"
443                        ),
444                    });
445                }
446                (None, Some(actual)) => {
447                    if let Some(function) = self.functions.get_mut(index as usize) {
448                        function.arity = actual;
449                        function.args = (0..actual).map(|slot| format!("arg{slot}")).collect();
450                    }
451                    self.function_meta
452                        .insert(name.to_string(), (index, Some(actual)));
453                }
454                _ => {}
455            }
456            return Ok(());
457        }
458
459        let index = u16::try_from(self.functions.len()).map_err(|_| ParseError {
460            span: None,
461            code: None,
462            line: 1,
463            message: "too many declared functions".to_string(),
464        })?;
465        let effective_arity = arity.unwrap_or(0);
466        self.functions.push(FunctionDecl {
467            name: name.to_string(),
468            arity: effective_arity,
469            index,
470            args: (0..effective_arity)
471                .map(|slot| format!("arg{slot}"))
472                .collect(),
473            arg_schemas: vec![None; usize::from(effective_arity)],
474            return_schema: None,
475            type_params: Vec::new(),
476            exported: false,
477            return_type: ValueType::Unknown,
478        });
479        self.function_meta.insert(name.to_string(), (index, arity));
480        Ok(())
481    }
482
483    pub fn resolve_call_expr(&mut self, name: &str, args: Vec<Expr>) -> Option<Expr> {
484        if let Some(local_index) = self.locals.get(name).copied() {
485            return Some(Expr::LocalCall(local_index, Vec::new(), args));
486        }
487        let (func_index, declared_arity) = self.function_meta.get(name).copied()?;
488        let call_arity = u8::try_from(args.len()).ok()?;
489        match declared_arity {
490            Some(expected)
491                if expected != call_arity && !known_host_accepts_arity(name, call_arity) =>
492            {
493                return None;
494            }
495            Some(_) => {}
496            None => {
497                if let Some(function) = self.functions.get_mut(func_index as usize) {
498                    function.arity = call_arity;
499                    function.args = (0..call_arity).map(|slot| format!("arg{slot}")).collect();
500                }
501                self.function_meta
502                    .insert(name.to_string(), (func_index, Some(call_arity)));
503            }
504        }
505        Some(Expr::Call(func_index, Vec::new(), args))
506    }
507
508    pub fn finish(self, stmts: Vec<Stmt>) -> FrontendIr {
509        let mut local_bindings = self
510            .locals
511            .into_iter()
512            .collect::<Vec<(String, LocalSlot)>>();
513        local_bindings.sort_by_key(|(_, index)| *index);
514        FrontendIr {
515            stmts,
516            locals: self.next_local as usize,
517            local_bindings,
518            struct_schemas: HashMap::new(),
519            unknown_type_spans: Vec::new(),
520            functions: self.functions,
521            function_impls: HashMap::new(),
522            stmt_sources: Vec::new(),
523            function_sources: HashMap::new(),
524        }
525    }
526
527    pub fn alloc_local_named(&mut self, name: &str) -> Result<LocalSlot, ParseError> {
528        if let Some(index) = self.locals.get(name).copied() {
529            return Ok(index);
530        }
531        let index = self.alloc_local()?;
532        self.locals.insert(name.to_string(), index);
533        Ok(index)
534    }
535
536    fn alloc_local(&mut self) -> Result<LocalSlot, ParseError> {
537        let index = self.next_local;
538        self.next_local = self.next_local.checked_add(1).ok_or(ParseError {
539            span: None,
540            code: None,
541            line: 1,
542            message: "local index overflow".to_string(),
543        })?;
544        Ok(index)
545    }
546}