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