Skip to main content

ternlang_core/codegen/
tern_asm.rs

1// SPDX-License-Identifier: LGPL-3.0-or-later
2// Ternlang — RFI-IRFOS Ternary Intelligence Stack
3// Copyright (C) 2026 RFI-IRFOS
4//
5// tern_asm.rs — TERN-compatible assembly emitter
6//
7// Produces TERN assembly text from a ternlang `Program` AST.
8//
9// TERN assembly is a RISC-V-inspired balanced ternary assembly language.
10// This emitter generates output that is structurally compatible with the
11// Tern Systems TERN specification (https://github.com/Tern-Computer).
12//
13// Register file: t0–t31 (32 ternary registers, each holds one trit value)
14// Special:       sp (stack pointer), ra (return address), zero (always 0)
15//
16// # Instruction summary
17//
18//   Arithmetic / logic
19//     tadd  rd, rs1, rs2    ; rd = rs1 + rs2  (ternary saturating add)
20//     tsub  rd, rs1, rs2    ; rd = rs1 - rs2
21//     tmul  rd, rs1, rs2    ; rd = rs1 * rs2
22//     tdiv  rd, rs1, rs2    ; rd = rs1 / rs2
23//     tmod  rd, rs1, rs2    ; rd = rs1 mod rs2
24//     tnot  rd, rs1         ; rd = -rs1        (ternary invert / negate)
25//     tcons rd, rs1, rs2    ; rd = consensus(rs1, rs2)  — ternary AND / min
26//     tmax  rd, rs1, rs2    ; rd = max(rs1, rs2)        — ternary OR
27//
28//   Comparison (result is a trit: +1 true, -1 false, 0 hold/equal)
29//     teq   rd, rs1, rs2    ; rd = (rs1 == rs2) ? +1 : -1
30//     tlt   rd, rs1, rs2    ; rd = (rs1 <  rs2) ? +1 : -1
31//     tgt   rd, rs1, rs2    ; rd = (rs1 >  rs2) ? +1 : -1
32//     tle   rd, rs1, rs2    ; rd = (rs1 <= rs2) ? +1 : -1
33//     tge   rd, rs1, rs2    ; rd = (rs1 >= rs2) ? +1 : -1
34//     tne   rd, rs1, rs2    ; rd = (rs1 != rs2) ? +1 : -1
35//
36//   Immediate load
37//     tldi  rd, imm         ; rd = imm  (imm ∈ {-1, 0, 1})
38//     tlii  rd, imm         ; rd = imm  (imm is an integer constant)
39//
40//   Control flow
41//     j     label           ; unconditional jump
42//     bpos  rs, label       ; jump if rs == +1  (affirm branch)
43//     bzero rs, label       ; jump if rs ==  0  (hold branch)
44//     bneg  rs, label       ; jump if rs == -1  (reject branch)
45//     call  label           ; call subroutine (saves ra)
46//     ret                   ; return (jumps to ra)
47//
48//   Memory (ternary word = 1 trit)
49//     tld   rd, offset(rs)  ; load trit from memory[rs + offset]
50//     tst   rs2, offset(rs1); store rs2 → memory[rs1 + offset]
51//
52//   I/O (runtime builtins)
53//     tprint rd             ; print trit value of rd as "affirm"/"hold"/"reject"
54//     tprintint rd          ; print integer value of rd
55
56use crate::ast::*;
57use std::collections::HashMap;
58
59// ── Register allocator ────────────────────────────────────────────────────────
60
61struct RegAlloc {
62    map: HashMap<String, u8>,
63    next: u8,
64}
65
66impl RegAlloc {
67    fn new() -> Self { Self { map: HashMap::new(), next: 2 } } // t0=zero, t1=scratch
68    fn alloc(&mut self, name: &str) -> u8 {
69        if let Some(&r) = self.map.get(name) { return r; }
70        let r = self.next;
71        self.next += 1;
72        self.map.insert(name.to_string(), r);
73        r
74    }
75    fn get(&self, name: &str) -> u8 {
76        *self.map.get(name).unwrap_or(&1) // fall back to t1 scratch
77    }
78    fn scratch(&mut self) -> u8 {
79        let r = self.next;
80        self.next += 1;
81        r
82    }
83    fn snapshot(&self) -> HashMap<String, u8> {
84        self.map.clone()
85    }
86    fn restore(&mut self, snap: HashMap<String, u8>, next: u8) {
87        self.map = snap;
88        self.next = next;
89    }
90}
91
92fn reg(n: u8) -> String {
93    match n {
94        0 => "zero".to_string(),
95        1 => "t1".to_string(),
96        _ => format!("t{}", n),
97    }
98}
99
100// ── Emitter ───────────────────────────────────────────────────────────────────
101
102pub struct TernAsmEmitter {
103    out: Vec<String>,
104    label_counter: usize,
105}
106
107impl TernAsmEmitter {
108    pub fn new() -> Self {
109        Self { out: Vec::new(), label_counter: 0 }
110    }
111
112    fn emit(&mut self, line: &str) {
113        self.out.push(format!("    {}", line));
114    }
115
116    fn emit_label(&mut self, label: &str) {
117        self.out.push(format!("{}:", label));
118    }
119
120    fn fresh_label(&mut self, prefix: &str) -> String {
121        let l = format!(".L_{}_{}", prefix, self.label_counter);
122        self.label_counter += 1;
123        l
124    }
125
126    /// Emit the full program as a TERN ASM string.
127    pub fn emit_program(&mut self, program: &Program) -> String {
128        self.out.push("; Generated by ternlang TERN-ASM emitter — RFI-IRFOS".to_string());
129        self.out.push("; TERN assembly (RISC-V-inspired balanced ternary)".to_string());
130        self.out.push("; Spec compatible with: Tern Systems TERN / BTMC".to_string());
131        self.out.push("".to_string());
132        self.out.push(".section .text".to_string());
133        self.out.push(".global main".to_string());
134        self.out.push("".to_string());
135
136        // Emit all functions
137        for func in &program.functions {
138            let mut ra = RegAlloc::new();
139            self.emit_function(func, &mut ra);
140            self.out.push("".to_string());
141        }
142
143        // Emit agent handlers
144        for agent in &program.agents {
145            for method in &agent.methods {
146                let label = format!("{}__{}", agent.name, method.name);
147                let mut ra = RegAlloc::new();
148                self.emit_function_with_label(&label, method, &mut ra);
149                self.out.push("".to_string());
150            }
151        }
152
153        self.out.join("\n")
154    }
155
156    fn emit_function(&mut self, func: &Function, ra: &mut RegAlloc) {
157        self.emit_function_with_label(&func.name, func, ra);
158    }
159
160    fn emit_function_with_label(&mut self, label: &str, func: &Function, ra: &mut RegAlloc) {
161        self.emit_label(label);
162
163        // Allocate parameter registers
164        for (param_name, _) in &func.params {
165            ra.alloc(param_name);
166        }
167
168        for stmt in &func.body {
169            self.emit_stmt(stmt, ra);
170        }
171
172        // If no explicit return at end, emit ret
173        let has_return = func.body.last().map(|s| matches!(s, Stmt::Return(_))).unwrap_or(false);
174        if !has_return {
175            self.emit("ret");
176        }
177    }
178
179    fn emit_stmt(&mut self, stmt: &Stmt, ra: &mut RegAlloc) {
180        match stmt {
181            Stmt::Let { name, value, .. } => {
182                if let Expr::StructLiteral { fields, .. } = value {
183                    for (f_name, f_val) in fields {
184                        let f_dest = ra.alloc(&format!("{}.{}", name, f_name));
185                        let f_src = self.emit_expr(f_val, ra);
186                        if f_src != f_dest {
187                            self.emit(&format!("tadd  {}, {}, zero   ; struct field init", reg(f_dest), reg(f_src)));
188                        }
189                    }
190                    // The main variable is just a dummy
191                    let dest = ra.alloc(name);
192                    self.emit(&format!("tldi  {}, 0           ; struct root dummy", reg(dest)));
193                } else {
194                    let dest = ra.alloc(name);
195                    let src = self.emit_expr(value, ra);
196                    if src != dest {
197                        self.emit(&format!("tadd  {}, {}, zero   ; {} = {}", reg(dest), reg(src), name, name));
198                    }
199                }
200                // else value was emitted directly into dest
201            }
202
203            Stmt::Set { name, value } => {
204                let dest = ra.get(name);
205                let src = self.emit_expr(value, ra);
206                if src != dest {
207                    self.emit(&format!("tadd  {}, {}, zero   ; {} = expr", reg(dest), reg(src), name));
208                }
209            }
210
211            Stmt::Return(expr) => {
212                let src = self.emit_expr(expr, ra);
213                // Convention: return value in t2
214                if src != 2 {
215                    self.emit(&format!("tadd  t2, {}, zero   ; return value", reg(src)));
216                }
217                self.emit("ret");
218            }
219
220            Stmt::Expr(expr) => {
221                self.emit_expr(expr, ra);
222            }
223
224            Stmt::Block(stmts) => {
225                let snap = ra.snapshot();
226                let next = ra.next;
227                for s in stmts { self.emit_stmt(s, ra); }
228                ra.restore(snap, next);
229            }
230
231            Stmt::IfTernary { condition, on_pos, on_zero, on_neg } => {
232                let cond_reg = self.emit_expr(condition, ra);
233                let lbl_pos  = self.fresh_label("pos");
234                let lbl_zero = self.fresh_label("zero");
235                let lbl_neg  = self.fresh_label("neg");
236                let lbl_end  = self.fresh_label("end");
237
238                self.emit(&format!("bpos  {}, {}", reg(cond_reg), lbl_pos));
239                self.emit(&format!("bzero {}, {}", reg(cond_reg), lbl_zero));
240                self.emit(&format!("j     {}", lbl_neg));
241
242                self.emit_label(&lbl_pos);
243                self.emit_stmt(on_pos, ra);
244                self.emit(&format!("j     {}", lbl_end));
245
246                self.emit_label(&lbl_zero);
247                self.emit_stmt(on_zero, ra);
248                self.emit(&format!("j     {}", lbl_end));
249
250                self.emit_label(&lbl_neg);
251                self.emit_stmt(on_neg, ra);
252
253                self.emit_label(&lbl_end);
254            }
255
256            Stmt::Match { condition, arms } => {
257                let cond_reg = self.emit_expr(condition, ra);
258                let lbl_end  = self.fresh_label("match_end");
259                let mut arm_labels: Vec<(i64, String)> = Vec::new();
260
261                for (val, _) in arms {
262                    arm_labels.push((*val, self.fresh_label(&format!("arm_{}", val))));
263                }
264
265                // branch dispatch
266                for (val, lbl) in &arm_labels {
267                    let tmp = ra.scratch();
268                    self.emit(&format!("tlii  {}, {}", reg(tmp), val));
269                    let cmp = ra.scratch();
270                    self.emit(&format!("teq   {}, {}, {}", reg(cmp), reg(cond_reg), reg(tmp)));
271                    self.emit(&format!("bpos  {}, {}", reg(cmp), lbl));
272                }
273                self.emit(&format!("j     {}", lbl_end));
274
275                for ((_, body_stmt), (_, lbl)) in arms.iter().zip(arm_labels.iter()) {
276                    self.emit_label(lbl);
277                    self.emit_stmt(body_stmt, ra);
278                    self.emit(&format!("j     {}", lbl_end));
279                }
280
281                self.emit_label(&lbl_end);
282            }
283
284            Stmt::WhileTernary { condition, on_pos, on_zero, on_neg } => {
285                let lbl_loop  = self.fresh_label("while");
286                let lbl_pos   = self.fresh_label("wpos");
287                let lbl_zero  = self.fresh_label("wzero");
288                let lbl_neg   = self.fresh_label("wneg");
289                let lbl_end   = self.fresh_label("wend");
290
291                self.emit_label(&lbl_loop);
292                let cond_reg = self.emit_expr(condition, ra);
293
294                self.emit(&format!("bpos  {}, {}", reg(cond_reg), lbl_pos));
295                self.emit(&format!("bzero {}, {}", reg(cond_reg), lbl_zero));
296                self.emit(&format!("j     {}", lbl_neg));
297
298                self.emit_label(&lbl_pos);
299                self.emit_stmt(on_pos, ra);
300                self.emit(&format!("j     {}", lbl_loop));
301
302                self.emit_label(&lbl_zero);
303                self.emit_stmt(on_zero, ra);
304                self.emit(&format!("j     {}", lbl_loop));
305
306                self.emit_label(&lbl_neg);
307                self.emit_stmt(on_neg, ra);
308                // neg branch exits loop
309
310                self.emit_label(&lbl_end);
311            }
312
313            Stmt::Loop { body } => {
314                let lbl_loop = self.fresh_label("loop");
315                let lbl_end  = self.fresh_label("loop_end");
316
317                self.emit_label(&lbl_loop);
318                self.emit_stmt(body, ra);
319                self.emit(&format!("j     {}", lbl_loop));
320                self.emit_label(&lbl_end);
321            }
322
323            Stmt::ForIn { var, iter, body } => {
324                // Simplified: emit a range-like loop if iter is a simple ident
325                let iter_reg  = self.emit_expr(iter, ra);
326                let var_reg   = ra.alloc(var);
327                let idx_reg   = ra.scratch();
328                let lbl_loop  = self.fresh_label("forin");
329                let lbl_end   = self.fresh_label("forin_end");
330
331                self.emit(&format!("tlii  {}, 0          ; for-in idx = 0", reg(idx_reg)));
332                self.emit_label(&lbl_loop);
333                // Load row from tensor iter at idx
334                self.emit(&format!("tld   {}, 0({})      ; load row iter[idx]", reg(var_reg), reg(iter_reg)));
335                self.emit_stmt(body, ra);
336                self.emit(&format!("tadd  {0}, {0}, t1   ; idx++", reg(idx_reg)));
337                self.emit(&format!("j     {}", lbl_loop));
338                self.emit_label(&lbl_end);
339            }
340
341            Stmt::Break    => { self.emit("j     .L_break_TODO   ; break"); }
342            Stmt::Continue => { self.emit("j     .L_continue_TODO ; continue"); }
343
344            Stmt::Send { target, message } => {
345                let t = self.emit_expr(target, ra);
346                let m = self.emit_expr(message, ra);
347                self.emit(&format!("tsend {}, {}           ; send msg to agent", reg(t), reg(m)));
348            }
349
350            Stmt::FieldSet { object, field: _, value } => {
351                let obj_reg = ra.get(object);
352                let val_reg = self.emit_expr(value, ra);
353                self.emit(&format!("tst   {}, 0({})        ; field store", reg(val_reg), reg(obj_reg)));
354            }
355
356            Stmt::IndexSet { object, row, col: _, value } => {
357                let obj_reg = ra.get(object);
358                let row_reg = self.emit_expr(row, ra);
359                let val_reg = self.emit_expr(value, ra);
360                let addr    = ra.scratch();
361                self.emit(&format!("tadd  {}, {}, {}       ; tensor index addr", reg(addr), reg(obj_reg), reg(row_reg)));
362                self.emit(&format!("tst   {}, 0({})        ; tensor store", reg(val_reg), reg(addr)));
363            }
364
365            Stmt::Decorated { stmt, .. } => self.emit_stmt(stmt, ra),
366            Stmt::Use { .. } => {}
367            Stmt::FromImport { .. } => {}
368        }
369    }
370
371    /// Emit an expression, returning the register number containing the result.
372    fn emit_expr(&mut self, expr: &Expr, ra: &mut RegAlloc) -> u8 {
373        match expr {
374            Expr::TritLiteral(v) => {
375                let r = ra.scratch();
376                self.emit(&format!("tldi  {}, {}          ; trit literal", reg(r), v));
377                r
378            }
379
380            Expr::IntLiteral(v) => {
381                let r = ra.scratch();
382                self.emit(&format!("tlii  {}, {}          ; int literal", reg(r), v));
383                r
384            }
385
386            Expr::FloatLiteral(v) => {
387                let r = ra.scratch();
388                self.emit(&format!("tlii  {}, {}          ; float (truncated to int)", reg(r), *v as i64));
389                r
390            }
391
392            Expr::StringLiteral(_) => {
393                let r = ra.scratch();
394                self.emit(&format!("tlii  {}, 0           ; string (addr placeholder)", reg(r)));
395                r
396            }
397
398            Expr::Ident(name) => {
399                ra.get(name)
400            }
401
402            Expr::BinaryOp { op, lhs, rhs } => {
403                let lreg = self.emit_expr(lhs, ra);
404                let rreg = self.emit_expr(rhs, ra);
405                let dest = ra.scratch();
406                let mnemonic = match op {
407                    BinOp::Add          => "tadd",
408                    BinOp::Sub          => "tsub",
409                    BinOp::Mul          => "tmul",
410                    BinOp::Div          => "tdiv",
411                    BinOp::Mod          => "tmod",
412                    BinOp::Equal        => "teq",
413                    BinOp::NotEqual     => "tne",
414                    BinOp::Less         => "tlt",
415                    BinOp::Greater      => "tgt",
416                    BinOp::LessEqual    => "tle",
417                    BinOp::GreaterEqual => "tge",
418                    BinOp::And          => "tcons",
419                    BinOp::Or           => "tmax",
420                };
421                self.emit(&format!("{:<6}{}, {}, {}", mnemonic, reg(dest), reg(lreg), reg(rreg)));
422                dest
423            }
424
425            Expr::UnaryOp { op: UnOp::Neg, expr } => {
426                let src  = self.emit_expr(expr, ra);
427                let dest = ra.scratch();
428                self.emit(&format!("tnot  {}, {}", reg(dest), reg(src)));
429                dest
430            }
431
432            Expr::Call { callee, args } => {
433                // Push args into a0..aN by convention
434                for (i, arg) in args.iter().enumerate() {
435                    let r = self.emit_expr(arg, ra);
436                    // Move into argument registers a0+ (we use t10+ as argument passing regs)
437                    let arg_reg = 10 + i as u8;
438                    if r != arg_reg {
439                        self.emit(&format!("tadd  t{}, {}, zero  ; arg {}", arg_reg, reg(r), i));
440                    }
441                }
442                // Builtin prints
443                match callee.as_str() {
444                    "print" | "debug_print" => {
445                        if !args.is_empty() {
446                            let r = self.emit_expr(&args[0], ra);
447                            self.emit(&format!("tprint {}", reg(r)));
448                        }
449                        return 0; // zero reg
450                    }
451                    _ => {}
452                }
453                self.emit(&format!("call  {}", callee));
454                2 // return value in t2 by convention
455            }
456
457            Expr::Cast { expr, .. } => {
458                self.emit_expr(expr, ra) // transparent pass-through
459            }
460
461            Expr::FieldAccess { object, field: _ } => {
462                let obj_reg = self.emit_expr(object, ra);
463                let dest    = ra.scratch();
464                self.emit(&format!("tld   {}, 0({})       ; field load", reg(dest), reg(obj_reg)));
465                dest
466            }
467
468            Expr::Index { object, row, col: _ } => {
469                let obj_reg = self.emit_expr(object, ra);
470                let row_reg = self.emit_expr(row, ra);
471                let addr    = ra.scratch();
472                let dest    = ra.scratch();
473                self.emit(&format!("tadd  {}, {}, {}      ; tensor index", reg(addr), reg(obj_reg), reg(row_reg)));
474                self.emit(&format!("tld   {}, 0({})       ; tensor load", reg(dest), reg(addr)));
475                dest
476            }
477
478            Expr::TritTensorLiteral(vals) => {
479                let base = ra.scratch();
480                self.emit(&format!("tlii  {}, 0           ; tensor literal (addr)", reg(base)));
481                for (i, v) in vals.iter().enumerate() {
482                    let tmp = ra.scratch();
483                    self.emit(&format!("tldi  {}, {}", reg(tmp), v));
484                    self.emit(&format!("tst   {}, {}({})   ; tensor[{}]", reg(tmp), i, reg(base), i));
485                }
486                base
487            }
488
489            Expr::Spawn { agent_name, .. } => {
490                let r = ra.scratch();
491                self.emit(&format!("tspawn {}, {}         ; spawn agent", reg(r), agent_name));
492                r
493            }
494
495            Expr::Await { target } => {
496                let t = self.emit_expr(target, ra);
497                let r = ra.scratch();
498                self.emit(&format!("tawait {}, {}         ; await agent", reg(r), reg(t)));
499                r
500            }
501
502            Expr::Propagate { expr } => {
503                let src = self.emit_expr(expr, ra);
504                let lbl = self.fresh_label("prop_ok");
505                // If src == -1 (reject), return -1 immediately
506                self.emit(&format!("bneg  {}, .L_prop_ret_{}", reg(src), self.label_counter));
507                self.emit(&format!("j     {}", lbl));
508                self.emit_label(&format!(".L_prop_ret_{}", self.label_counter - 1));
509                self.emit("tldi  t2, -1         ; propagate reject");
510                self.emit("ret");
511                self.emit_label(&lbl);
512                src
513            }
514
515            Expr::NodeId => {
516                let r = ra.scratch();
517                self.emit(&format!("tnodeid {}", reg(r)));
518                r
519            }
520
521            Expr::StructLiteral { .. } => {
522                let r = ra.scratch();
523                self.emit(&format!("tldi  {}, 0           ; struct literal (dummy)", reg(r)));
524                r
525            }
526        }
527    }
528}
529
530// ── Public API ────────────────────────────────────────────────────────────────
531
532/// Emit a ternlang `Program` as TERN-compatible assembly text.
533pub fn emit_tern_asm(program: &Program) -> String {
534    TernAsmEmitter::new().emit_program(program)
535}