Skip to main content

ternlang_core/codegen/
betbc.rs

1use crate::ast::*;
2use crate::vm::bet::pack_trits;
3use crate::trit::Trit;
4
5pub struct BytecodeEmitter {
6    code: Vec<u8>,
7    symbols: std::collections::HashMap<String, u8>,
8    func_addrs: std::collections::HashMap<String, u16>,
9    break_patches: Vec<usize>, // addresses to patch when a loop ends
10    next_reg: u8,
11    /// Struct layouts: struct_name → ordered field names
12    struct_layouts: std::collections::HashMap<String, Vec<String>>,
13    /// Agent type IDs: agent_name → type_id (u16 index)
14    agent_type_ids: std::collections::HashMap<String, u16>,
15    /// Agent handler addresses emitted during emit_program (type_id → addr)
16    agent_handlers: Vec<(u16, u16)>,
17}
18
19impl BytecodeEmitter {
20    pub fn new() -> Self {
21        Self {
22            code: Vec::new(),
23            symbols: std::collections::HashMap::new(),
24            func_addrs: std::collections::HashMap::new(),
25            break_patches: Vec::new(),
26            next_reg: 0,
27            struct_layouts: std::collections::HashMap::new(),
28            agent_type_ids: std::collections::HashMap::new(),
29            agent_handlers: Vec::new(),
30        }
31    }
32
33    /// After emit_program, call this to wire agent handler addresses into a VM.
34    pub fn register_agents(&self, vm: &mut crate::vm::BetVm) {
35        for &(type_id, addr) in &self.agent_handlers {
36            vm.register_agent_type(type_id, addr as usize);
37        }
38    }
39
40    pub fn emit_program(&mut self, program: &Program) {
41        // Register struct layouts so field-access codegen knows field order.
42        for s in &program.structs {
43            let field_names: Vec<String> = s.fields.iter().map(|(n, _)| n.clone()).collect();
44            self.struct_layouts.insert(s.name.clone(), field_names);
45        }
46        // Register agent type IDs before emitting bodies.
47        for (idx, agent) in program.agents.iter().enumerate() {
48            self.agent_type_ids.insert(agent.name.clone(), idx as u16);
49        }
50
51        // Two-pass: first emit a TJMP over all function/agent bodies.
52        let entry_jmp_patch = self.code.len() + 1;
53        self.code.push(0x0b); // TJMP — skip over function bodies
54        self.code.extend_from_slice(&[0u8, 0u8]);
55
56        // Emit agent handler methods (the `handle` fn of each agent).
57        for agent in &program.agents {
58            let type_id = self.agent_type_ids[&agent.name];
59            // The first method named "handle" is the entry point.
60            // All methods are emitted; the first one becomes the handler addr.
61            let mut handler_addr: Option<u16> = None;
62            for method in &agent.methods {
63                let addr = self.code.len() as u16;
64                if handler_addr.is_none() {
65                    handler_addr = Some(addr);
66                }
67                self.emit_function(method);
68                // Also register under the fully-qualified name "AgentName::method"
69                let fq = format!("{}::{}", agent.name, method.name);
70                self.func_addrs.insert(fq, addr);
71            }
72            if let Some(addr) = handler_addr {
73                self.agent_handlers.push((type_id, addr));
74            }
75        }
76
77        // Emit regular function bodies.
78        for func in &program.functions {
79            self.emit_function(func);
80        }
81
82        // Patch entry jump to land after all bodies.
83        let after_funcs = self.code.len() as u16;
84        self.patch_u16(entry_jmp_patch, after_funcs);
85    }
86
87    pub fn emit_function(&mut self, func: &Function) {
88        // Record address of this function's first instruction.
89        let func_addr = self.code.len() as u16;
90        self.func_addrs.insert(func.name.clone(), func_addr);
91
92        for stmt in &func.body {
93            self.emit_stmt(stmt);
94        }
95        // Emit TRET at end of every function body.
96        self.code.push(0x11); // TRET
97    }
98
99    pub fn emit_stmt(&mut self, stmt: &Stmt) {
100        match stmt {
101            Stmt::Let { name, ty, value } => {
102                match ty {
103                    Type::TritTensor { dims } => {
104                        let size: usize = dims.iter().product();
105                        self.code.push(0x0f); // TALLOC
106                        self.code.extend_from_slice(&(size as u16).to_le_bytes());
107                        let reg = self.next_reg;
108                        self.symbols.insert(name.clone(), reg);
109                        self.next_reg += 1;
110                        self.code.push(0x08); // TSTORE
111                        self.code.push(reg);
112                    }
113                    Type::Named(struct_name) => {
114                        // Allocate one register per field, zero-initialised.
115                        // Fields stored as "<instance>.<field>" in the symbol table.
116                        let fields = self.struct_layouts.get(struct_name)
117                            .cloned()
118                            .unwrap_or_default();
119                        // Record base register under the instance name too (not strictly needed).
120                        let base_reg = self.next_reg;
121                        self.symbols.insert(name.clone(), base_reg);
122                        for field in &fields {
123                            let reg = self.next_reg;
124                            self.next_reg += 1;
125                            self.symbols.insert(format!("{}.{}", name, field), reg);
126                            // Zero-initialise each field
127                            self.code.push(0x01); // TPUSH hold
128                            self.code.extend(crate::vm::bet::pack_trits(&[crate::trit::Trit::Tend]));
129                            self.code.push(0x08); // TSTORE
130                            self.code.push(reg);
131                        }
132                        // If no fields, still emit the base placeholder
133                        if fields.is_empty() {
134                            self.next_reg += 1;
135                            self.code.push(0x01);
136                            self.code.extend(crate::vm::bet::pack_trits(&[crate::trit::Trit::Tend]));
137                            self.code.push(0x08);
138                            self.code.push(base_reg);
139                        }
140                    }
141                    _ => {
142                        self.emit_expr(value);
143                        let reg = self.next_reg;
144                        self.symbols.insert(name.clone(), reg);
145                        self.next_reg += 1;
146                        self.code.push(0x08); // TSTORE
147                        self.code.push(reg);
148                    }
149                }
150            }
151            Stmt::FieldSet { object, field, value } => {
152                // Resolve the mangled register name for this field.
153                let key = format!("{}.{}", object, field);
154                self.emit_expr(value);
155                if let Some(&reg) = self.symbols.get(&key) {
156                    self.code.push(0x08); // TSTORE
157                    self.code.push(reg);
158                }
159                // Unknown field — emit nothing (will be a runtime no-op).
160            }
161            Stmt::IndexSet { object, row, col, value } => {
162                if let Some(&reg) = self.symbols.get(object) {
163                    self.code.push(0x09); self.code.push(reg); // TLOAD tensor ref
164                    self.emit_expr(row);
165                    self.emit_expr(col);
166                    self.emit_expr(value);
167                    self.code.push(0x23); // TSET
168                }
169            }
170            Stmt::IfTernary { condition, on_pos, on_zero, on_neg } => {
171                self.emit_expr(condition);
172                
173                // Jump to POS branch
174                self.code.push(0x0a); // TDUP
175                let jmp_pos_patch = self.code.len() + 1;
176                self.code.push(0x05); // TJMP_POS
177                self.code.extend_from_slice(&[0, 0]);
178
179                // Jump to ZERO branch
180                self.code.push(0x0a); // TDUP
181                let jmp_zero_patch = self.code.len() + 1;
182                self.code.push(0x06); // TJMP_ZERO
183                self.code.extend_from_slice(&[0, 0]);
184
185                // Fallthrough to NEG branch
186                self.code.push(0x0c); // TPOP
187                self.emit_stmt(on_neg);
188                let end_jmp_neg_patch = self.code.len() + 1;
189                self.code.push(0x0b); // TJMP
190                self.code.extend_from_slice(&[0, 0]);
191
192                // POS Branch
193                let pos_addr = self.code.len() as u16;
194                self.patch_u16(jmp_pos_patch, pos_addr);
195                self.code.push(0x0c); // TPOP
196                self.emit_stmt(on_pos);
197                let end_jmp_pos_patch = self.code.len() + 1;
198                self.code.push(0x0b); // TJMP
199                self.code.extend_from_slice(&[0, 0]);
200
201                // ZERO Branch
202                let zero_addr = self.code.len() as u16;
203                self.patch_u16(jmp_zero_patch, zero_addr);
204                self.code.push(0x0c); // TPOP
205                self.emit_stmt(on_zero);
206                
207                // End Label
208                let end_addr = self.code.len() as u16;
209                self.patch_u16(end_jmp_neg_patch, end_addr);
210                self.patch_u16(end_jmp_pos_patch, end_addr);
211            }
212            Stmt::Match { condition, arms } => {
213                self.emit_expr(condition);
214                
215                let mut patches = Vec::new();
216                let mut end_patches = Vec::new();
217
218                for (val, _stmt) in arms {
219                    self.code.push(0x0a); // TDUP
220                    let patch_pos = self.code.len() + 1;
221                    match val {
222                        1  => self.code.push(0x05), // TJMP_POS
223                        0  => self.code.push(0x06), // TJMP_ZERO
224                        -1 => self.code.push(0x07), // TJMP_NEG
225                        _  => unreachable!(),
226                    }
227                    self.code.extend_from_slice(&[0, 0]);
228                    patches.push((patch_pos, *val));
229                }
230
231                self.code.push(0x0c); // TPOP (If no match found)
232                let end_jmp_no_match = self.code.len() + 1;
233                self.code.push(0x0b); // TJMP
234                self.code.extend_from_slice(&[0, 0]);
235
236                for (patch_pos, val) in patches {
237                    let addr = self.code.len() as u16;
238                    self.patch_u16(patch_pos, addr);
239                    self.code.push(0x0c); // TPOP
240                    
241                    // Find the stmt for this val
242                    let stmt = arms.iter().find(|(v, _)| *v == val).unwrap().1.clone();
243                    self.emit_stmt(&stmt);
244                    
245                    let end_patch = self.code.len() + 1;
246                    self.code.push(0x0b); // TJMP
247                    self.code.extend_from_slice(&[0, 0]);
248                    end_patches.push(end_patch);
249                }
250
251                let end_addr = self.code.len() as u16;
252                self.patch_u16(end_jmp_no_match, end_addr);
253                for p in end_patches {
254                    self.patch_u16(p, end_addr);
255                }
256            }
257            // for <var> in <tensor_ref> { body }
258            // Iterates over each trit in a tensor sequentially.
259            Stmt::ForIn { var, iter, body } => {
260                // Emit the iterable expression — expects a TensorRef on stack
261                self.emit_expr(iter);
262                let iter_reg = self.next_reg;
263                self.symbols.insert(format!("__iter_{}", var), iter_reg);
264                self.next_reg += 1;
265                self.code.push(0x08); self.code.push(iter_reg); // TSTORE iter ref
266
267                // Index register
268                let idx_reg = self.next_reg;
269                self.next_reg += 1;
270                // TSHAPE → push (rows, cols); store cols for bound check
271                self.code.push(0x09); self.code.push(iter_reg); // TLOAD tensor ref
272                self.code.push(0x24); // TSHAPE → rows, cols
273                let bound_reg = self.next_reg; self.next_reg += 1;
274                self.code.push(0x08); self.code.push(bound_reg); // TSTORE cols (bound)
275                // discard rows
276                self.code.push(0x0c); // TPOP rows
277
278                // Initialise index to 0 (hold trit — used as int here)
279                self.code.push(0x01);
280                self.code.extend(pack_trits(&[Trit::Tend]));
281                self.code.push(0x08); self.code.push(idx_reg); // TSTORE idx=0
282
283                let loop_top = self.code.len() as u16;
284
285
286
287                // Load current element: TIDX(tensor, 0, idx) — simplified 1D walk
288                self.code.push(0x09); self.code.push(iter_reg); // TLOAD tensor
289                self.code.push(0x01); self.code.extend(pack_trits(&[Trit::Tend])); // row 0
290                self.code.push(0x09); self.code.push(idx_reg); // TLOAD idx
291                self.code.push(0x22); // TIDX → trit element
292                let var_reg = self.next_reg; self.next_reg += 1;
293                self.symbols.insert(var.clone(), var_reg);
294                self.code.push(0x08); self.code.push(var_reg); // TSTORE var
295
296                // Emit body
297                self.emit_stmt(body);
298
299                // Unconditional jump back to loop top
300                let jmp_back = self.code.len() + 1;
301                self.code.push(0x0b);
302                self.code.extend_from_slice(&[0, 0]);
303                self.patch_u16(jmp_back, loop_top);
304            }
305
306            // loop { body } — infinite loop, exited by break
307            Stmt::Loop { body } => {
308                let loop_top = self.code.len() as u16;
309
310                // Track break patch sites
311                let pre_break_count = self.break_patches.len();
312                self.emit_stmt(body);
313                // Jump back to top
314                let jmp_back = self.code.len() + 1;
315                self.code.push(0x0b);
316                self.code.extend_from_slice(&[0, 0]);
317                self.patch_u16(jmp_back, loop_top);
318                // Collect break patches, then apply (avoids double borrow)
319                let after_loop = self.code.len() as u16;
320                let patches: Vec<usize> = self.break_patches.drain(pre_break_count..).collect();
321                for patch in patches {
322                    self.patch_u16(patch, after_loop);
323                }
324            }
325
326            Stmt::Break => {
327                let patch = self.code.len() + 1;
328                self.code.push(0x0b); // TJMP (address patched by enclosing loop)
329                self.code.extend_from_slice(&[0, 0]);
330                self.break_patches.push(patch);
331            }
332
333            Stmt::Continue => {
334                // Continue is a TJMP to loop top — needs enclosing loop context.
335                // Emit as no-op for now (safe: loop naturally continues).
336            }
337
338            Stmt::Use { .. } => {
339                // Module resolution not yet implemented — no-op at codegen level.
340            }
341            Stmt::Send { target, message } => {
342                // Push AgentRef, then message, then emit TSEND (0x31).
343                self.emit_expr(target);
344                self.emit_expr(message);
345                self.code.push(0x31); // TSEND
346            }
347
348            Stmt::WhileTernary { condition, on_pos, on_zero, on_neg } => {
349                let loop_top = self.code.len() as u16;
350                self.emit_expr(condition);
351                // Ternary branch: pos → on_pos body, zero → on_zero body, neg → break
352                self.code.push(0x0a); // TDUP
353                let jmp_pos_patch = self.code.len() + 1;
354                self.code.push(0x05); self.code.extend_from_slice(&[0, 0]); // TJMP_POS
355                self.code.push(0x0a); // TDUP
356                let jmp_zero_patch = self.code.len() + 1;
357                self.code.push(0x06); self.code.extend_from_slice(&[0, 0]); // TJMP_ZERO
358                // Neg branch: exit loop
359                self.code.push(0x0c); // TPOP
360                self.emit_stmt(on_neg);
361                let exit_patch = self.code.len() + 1;
362                self.code.push(0x0b); self.code.extend_from_slice(&[0, 0]); // TJMP exit
363
364                // Pos branch
365                let pos_addr = self.code.len() as u16;
366                self.patch_u16(jmp_pos_patch, pos_addr);
367                self.code.push(0x0c); // TPOP
368                self.emit_stmt(on_pos);
369                let back_pos = self.code.len() + 1;
370                self.code.push(0x0b); self.code.extend_from_slice(&[0, 0]);
371                self.patch_u16(back_pos, loop_top);
372
373                // Zero branch
374                let zero_addr = self.code.len() as u16;
375                self.patch_u16(jmp_zero_patch, zero_addr);
376                self.code.push(0x0c); // TPOP
377                self.emit_stmt(on_zero);
378                let back_zero = self.code.len() + 1;
379                self.code.push(0x0b); self.code.extend_from_slice(&[0, 0]);
380                self.patch_u16(back_zero, loop_top);
381
382                // Exit label
383                let exit_addr = self.code.len() as u16;
384                self.patch_u16(exit_patch, exit_addr);
385            }
386
387            Stmt::Return(expr) => {
388                self.emit_expr(expr);
389                self.code.push(0x11); // TRET
390            }
391            Stmt::Block(stmts) => {
392                for stmt in stmts {
393                    self.emit_stmt(stmt);
394                }
395            }
396            Stmt::Decorated { directive, stmt } => {
397                if directive == "sparseskip" {
398                    let opcode = 0x21u8; // TSPARSE_MATMUL — always emit sparse opcode
399
400                    // Case 1: @sparseskip on a bare expression: matmul(a, b);
401                    if let Stmt::Expr(inner_expr) = stmt.as_ref() {
402                        if let Expr::Call { callee, args } = inner_expr {
403                            if callee == "matmul" && args.len() == 2 {
404                                self.emit_expr(&args[0]);
405                                self.emit_expr(&args[1]);
406                                self.code.push(opcode); // TSPARSE_MATMUL or TMATMUL
407                                return;
408                            }
409                        }
410                    }
411                    // Case 2: @sparseskip on a let binding: let c: trittensor = matmul(a, b);
412                    if let Stmt::Let { name, value, .. } = stmt.as_ref() {
413                        if let Expr::Call { callee, args } = value {
414                            if callee == "matmul" && args.len() == 2 {
415                                self.emit_expr(&args[0]);
416                                self.emit_expr(&args[1]);
417                                self.code.push(opcode); // TSPARSE_MATMUL or TMATMUL
418                                
419                                if opcode == 0x21 {
420                                    // TSPARSE_MATMUL pushes TensorRef then Int(skipped_count)
421                                    self.code.push(0x0c); // TPOP — discard skipped_count
422                                }
423
424                                let reg = self.next_reg;
425                                self.symbols.insert(name.clone(), reg);
426                                self.next_reg += 1;
427                                self.code.push(0x08); // TSTORE tensor ref into register
428                                self.code.push(reg);
429                                return;
430                            }
431                        }
432                    }
433                }
434                // Fallthrough: emit the inner statement unchanged
435                self.emit_stmt(stmt);
436            }
437            _ => {}
438        }
439    }
440
441    fn emit_expr(&mut self, expr: &Expr) {
442        match expr {
443            Expr::TritLiteral(val) => {
444                self.code.push(0x01); // TPUSH
445                let trit = Trit::from(*val);
446                self.code.extend(pack_trits(&[trit]));
447            }
448            Expr::Ident(name) => {
449                if let Some(&reg) = self.symbols.get(name) {
450                    self.code.push(0x09); // TLOAD
451                    self.code.push(reg);
452                }
453            }
454            Expr::BinaryOp { op, lhs, rhs } => {
455                self.emit_expr(lhs);
456                self.emit_expr(rhs);
457                match op {
458                    BinOp::Add      => self.code.push(0x02), // TADD
459                    BinOp::Mul      => self.code.push(0x03), // TMUL
460                    BinOp::Sub      => { self.code.push(0x04); self.code.push(0x02); } // TNEG rhs, TADD
461                    BinOp::Equal    => self.code.push(0x16), // TEQ
462                    BinOp::NotEqual => { self.code.push(0x16); self.code.push(0x04); } // TEQ then TNEG
463                    BinOp::And      => self.code.push(0x03), // TMUL (ternary AND = multiply)
464                    BinOp::Or       => self.code.push(0x0e), // TCONS (ternary OR = consensus)
465                    BinOp::Less     => self.code.push(0x14), // TLESS  (a < b → affirm/tend/reject)
466                    BinOp::Greater  => self.code.push(0x15), // TGREATER (a > b → affirm/tend/reject)
467                }
468            }
469            Expr::UnaryOp { op, expr } => {
470                self.emit_expr(expr);
471                match op {
472                    UnOp::Neg => self.code.push(0x04),
473                }
474            }
475            Expr::Call { callee, args } => {
476                for arg in args {
477                    self.emit_expr(arg);
478                }
479                match callee.as_str() {
480                    "consensus" => {
481                        if args.len() == 2 {
482                            self.code.push(0x0e); // TCONS
483                        }
484                    }
485                    "invert" => {
486                        if args.len() == 1 {
487                            self.code.push(0x04); // TNEG
488                        }
489                    }
490                    "truth" => {
491                        self.code.push(0x01); // TPUSH
492                        self.code.extend(pack_trits(&[Trit::Affirm]));
493                    }
494                    "hold" => {
495                        self.code.push(0x01); // TPUSH
496                        self.code.extend(pack_trits(&[Trit::Tend]));
497                    }
498                    "conflict" => {
499                        self.code.push(0x01); // TPUSH
500                        self.code.extend(pack_trits(&[Trit::Reject]));
501                    }
502                    "matmul" => {
503                        if args.len() == 2 {
504                            self.code.push(0x20); // TMATMUL (dense)
505                        }
506                    }
507                    "sparsity" => {
508                        if args.len() == 1 {
509                            self.code.push(0x25); // TSPARSITY
510                        }
511                    }
512                    "shape" => {
513                        if args.len() == 1 {
514                            self.code.push(0x24); // TSHAPE
515                        }
516                    }
517                    _ => {
518                        // User-defined function call — emit TCALL if address known
519                        if let Some(&addr) = self.func_addrs.get(callee) {
520                            self.code.push(0x10); // TCALL
521                            self.code.extend_from_slice(&addr.to_le_bytes());
522                        } else {
523                            // Forward reference: emit TCALL with placeholder, needs second pass
524                            // For now emit hold as safe default
525                            self.code.push(0x01); // TPUSH hold
526                            self.code.extend(pack_trits(&[Trit::Tend]));
527                        }
528                    }
529                }
530            }
531            Expr::FieldAccess { object, field } => {
532                // Resolve to mangled register: "<object_name>.<field>"
533                // Only handles single-level ident.field — nested access falls back to hold.
534                if let Expr::Ident(obj_name) = object.as_ref() {
535                    let key = format!("{}.{}", obj_name, field);
536                    if let Some(&reg) = self.symbols.get(&key) {
537                        self.code.push(0x09); // TLOAD
538                        self.code.push(reg);
539                        return;
540                    }
541                }
542                // Fallback: push hold
543                self.code.push(0x01);
544                self.code.extend(pack_trits(&[Trit::Tend]));
545            }
546            Expr::Index { object, row, col } => {
547                self.emit_expr(object);
548                self.emit_expr(row);
549                self.emit_expr(col);
550                self.code.push(0x22); // TIDX
551            }
552            Expr::Propagate { expr } => {
553                // Evaluate inner expression → stack: [val]
554                self.emit_expr(expr);
555                // TDUP → [val, dup]
556                self.code.push(0x0a);
557                // TJMP_NEG to propagate path — consumes dup; if -1 jumps, else [val] remains
558                let neg_patch = self.code.len() + 1;
559                self.code.push(0x07); // TJMP_NEG
560                self.code.extend_from_slice(&[0u8, 0u8]);
561                // Not -1: skip over the early return
562                let skip_patch = self.code.len() + 1;
563                self.code.push(0x0b); // TJMP
564                self.code.extend_from_slice(&[0u8, 0u8]);
565                // propagate path: val=-1 is still on stack — TRET returns it
566                let prop_addr = self.code.len() as u16;
567                self.patch_u16(neg_patch, prop_addr);
568                self.code.push(0x11); // TRET
569                // skip label: continue with val on stack
570                let skip_addr = self.code.len() as u16;
571                self.patch_u16(skip_patch, skip_addr);
572            }
573            Expr::Cast { expr, .. } => {
574                // cast() is a no-op at the BET level — trits are already in canonical form.
575                // Emit the inner expression; the type annotation guides the type checker only.
576                self.emit_expr(expr);
577            }
578            Expr::Spawn { agent_name, node_addr } => {
579                if let Some(addr) = node_addr {
580                    // Remote spawn: push addr string, then TREMOTE_SPAWN(0x33)
581                    self.emit_expr(&Expr::StringLiteral(addr.clone()));
582                    if let Some(&type_id) = self.agent_type_ids.get(agent_name) {
583                        self.code.push(0x33); // TREMOTE_SPAWN
584                        self.code.extend_from_slice(&type_id.to_le_bytes());
585                    } else {
586                        self.code.push(0x01);
587                        self.code.extend(pack_trits(&[Trit::Tend]));
588                    }
589                } else if let Some(&type_id) = self.agent_type_ids.get(agent_name) {
590                    // Local spawn
591                    self.code.push(0x30); // TSPAWN
592                    self.code.extend_from_slice(&type_id.to_le_bytes());
593                } else {
594                    // Unknown agent — push hold as fallback
595                    self.code.push(0x01);
596                    self.code.extend(pack_trits(&[Trit::Tend]));
597                }
598            }
599            Expr::StringLiteral(_s) => {
600                // For v0.1: we don't have a TPUSH_STRING opcode.
601                // Instead, we hack it by passing strings out-of-band or
602                // just ignoring them in the BET bytecode for now.
603                // Actually, for remote spawn to work, we need to pass the address.
604                // Let's assume the VM can handle a raw string in the value stack if pushed via a hook.
605                // For now, emit a placeholder.
606                self.code.push(0x01);
607                self.code.extend(pack_trits(&[Trit::Tend]));
608            }
609            Expr::NodeId => {
610                self.code.push(0x12); // TNODEID
611            }
612            Expr::Await { target } => {
613                // Emit the AgentRef expression, then TAWAIT (0x32).
614                // TAWAIT pops the AgentRef, pops its mailbox front, calls handler, pushes result.
615                self.emit_expr(target);
616                self.code.push(0x32); // TAWAIT
617            }
618            _ => {}
619        }
620    }
621
622    /// Emit a TCALL to a named function.  Call this after `emit_program()` to
623    /// create an entry point that executes a specific function (typically `main`).
624    /// The TJMP in `emit_program` already points past all function bodies, so
625    /// code appended here is what actually runs at startup.
626    ///
627    /// The function's return value will be on the stack when the VM halts.
628    pub fn emit_entry_call(&mut self, func_name: &str) {
629        if let Some(&addr) = self.func_addrs.get(func_name) {
630            self.code.push(0x10); // TCALL — push return addr, jump to func
631            self.code.extend_from_slice(&addr.to_le_bytes());
632        }
633    }
634
635    pub fn finalize(mut self) -> Vec<u8> {
636        self.code.push(0x00); // THALT
637        self.code
638    }
639
640    fn patch_u16(&mut self, pos: usize, val: u16) {
641        let bytes = val.to_le_bytes();
642        self.code[pos] = bytes[0];
643        self.code[pos + 1] = bytes[1];
644    }
645}
646
647#[cfg(test)]
648mod tests {
649    use super::*;
650    use crate::parser::Parser;
651    use crate::vm::{BetVm, Value};
652
653    #[test]
654    fn test_compile_and_run_simple() {
655        let input = "let x: trit = 1; let y: trit = -x; return y;";
656        let mut parser = Parser::new(input);
657        let mut emitter = BytecodeEmitter::new();
658        
659        // Parse and emit statements
660        while let Ok(stmt) = parser.parse_stmt() {
661            emitter.emit_stmt(&stmt);
662        }
663        
664        let code = emitter.finalize();
665        let mut vm = BetVm::new(code);
666        vm.run().unwrap();
667        
668        // Final 'y' should be in register 1
669        assert_eq!(vm.get_register(1), Value::Trit(Trit::Reject));
670    }
671
672    #[test]
673    fn test_sparseskip_emits_tsparse_matmul() {
674        // @sparseskip on a let binding with matmul rhs should emit TSPARSE_MATMUL (0x21)
675        // and the result tensor ref should be stored in a register
676        let input = "let a: trittensor<2 x 2>; let b: trittensor<2 x 2>; @sparseskip let c: trittensor<2 x 2> = matmul(a, b);";
677        let mut parser = Parser::new(input);
678        let mut emitter = BytecodeEmitter::new();
679
680        while let Ok(stmt) = parser.parse_stmt() {
681            emitter.emit_stmt(&stmt);
682        }
683
684        let code = emitter.finalize();
685        // Verify TSPARSE_MATMUL (0x21) appears in the bytecode
686        assert!(code.contains(&0x21), "Expected TSPARSE_MATMUL (0x21) in bytecode");
687        // Verify dense TMATMUL (0x20) does NOT appear (we used sparseskip)
688        assert!(!code.contains(&0x20), "Expected no dense TMATMUL (0x20) when @sparseskip used");
689
690        // Run it — both tensors are zero-initialized, result should be TensorRef in reg2
691        let mut vm = BetVm::new(code);
692        vm.run().unwrap();
693        assert!(matches!(vm.get_register(2), Value::TensorRef(_)));
694    }
695
696    #[test]
697    fn test_compile_match() {
698        let input = "let x: trit = 1; match x { 1 => { let y: trit = -1; } 0 => { let y: trit = 0; } -1 => { let y: trit = 1; } }";
699        let mut parser = Parser::new(input);
700        let mut emitter = BytecodeEmitter::new();
701        
702        while let Ok(stmt) = parser.parse_stmt() {
703            emitter.emit_stmt(&stmt);
704        }
705        
706        let code = emitter.finalize();
707        let mut vm = BetVm::new(code);
708        vm.run().unwrap();
709        
710        // 'x' is 1, so 'y' in the first branch should be -1.
711        assert_eq!(vm.get_register(1), Value::Trit(Trit::Reject));
712    }
713}