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 (pattern, _) in arms {
262                    let val = match pattern {
263                        Pattern::Int(v) => *v,
264                        Pattern::Trit(t) => *t as i64,
265                        Pattern::Float(f) => *f as i64,
266                    };
267                    arm_labels.push((val, self.fresh_label(&format!("arm_{}", val))));
268                }
269
270                // branch dispatch
271                for (val, lbl) in &arm_labels {
272                    let tmp = ra.scratch();
273                    self.emit(&format!("tlii  {}, {}", reg(tmp), val));
274                    let cmp = ra.scratch();
275                    self.emit(&format!("teq   {}, {}, {}", reg(cmp), reg(cond_reg), reg(tmp)));
276                    self.emit(&format!("bpos  {}, {}", reg(cmp), lbl));
277                }
278                self.emit(&format!("j     {}", lbl_end));
279
280                for ((_, body_stmt), (_, lbl)) in arms.iter().zip(arm_labels.iter()) {
281                    self.emit_label(lbl);
282                    self.emit_stmt(body_stmt, ra);
283                    self.emit(&format!("j     {}", lbl_end));
284                }
285
286                self.emit_label(&lbl_end);
287            }
288
289            Stmt::WhileTernary { condition, on_pos, on_zero, on_neg } => {
290                let lbl_loop  = self.fresh_label("while");
291                let lbl_pos   = self.fresh_label("wpos");
292                let lbl_zero  = self.fresh_label("wzero");
293                let lbl_neg   = self.fresh_label("wneg");
294                let lbl_end   = self.fresh_label("wend");
295
296                self.emit_label(&lbl_loop);
297                let cond_reg = self.emit_expr(condition, ra);
298
299                self.emit(&format!("bpos  {}, {}", reg(cond_reg), lbl_pos));
300                self.emit(&format!("bzero {}, {}", reg(cond_reg), lbl_zero));
301                self.emit(&format!("j     {}", lbl_neg));
302
303                self.emit_label(&lbl_pos);
304                self.emit_stmt(on_pos, ra);
305                self.emit(&format!("j     {}", lbl_loop));
306
307                self.emit_label(&lbl_zero);
308                self.emit_stmt(on_zero, ra);
309                self.emit(&format!("j     {}", lbl_loop));
310
311                self.emit_label(&lbl_neg);
312                self.emit_stmt(on_neg, ra);
313                // neg branch exits loop
314
315                self.emit_label(&lbl_end);
316            }
317
318            Stmt::Loop { body } => {
319                let lbl_loop = self.fresh_label("loop");
320                let lbl_end  = self.fresh_label("loop_end");
321
322                self.emit_label(&lbl_loop);
323                self.emit_stmt(body, ra);
324                self.emit(&format!("j     {}", lbl_loop));
325                self.emit_label(&lbl_end);
326            }
327
328            Stmt::ForIn { var, iter, body } => {
329                // Simplified: emit a range-like loop if iter is a simple ident
330                let iter_reg  = self.emit_expr(iter, ra);
331                let var_reg   = ra.alloc(var);
332                let idx_reg   = ra.scratch();
333                let lbl_loop  = self.fresh_label("forin");
334                let lbl_end   = self.fresh_label("forin_end");
335
336                self.emit(&format!("tlii  {}, 0          ; for-in idx = 0", reg(idx_reg)));
337                self.emit_label(&lbl_loop);
338                // Load row from tensor iter at idx
339                self.emit(&format!("tld   {}, 0({})      ; load row iter[idx]", reg(var_reg), reg(iter_reg)));
340                self.emit_stmt(body, ra);
341                self.emit(&format!("tadd  {0}, {0}, t1   ; idx++", reg(idx_reg)));
342                self.emit(&format!("j     {}", lbl_loop));
343                self.emit_label(&lbl_end);
344            }
345
346            Stmt::Break    => { self.emit("j     .L_break_arch_def ; break"); }
347            Stmt::Continue => { self.emit("j     .L_continue_arch_def ; continue"); }
348
349            Stmt::Send { target, message } => {
350                let t = self.emit_expr(target, ra);
351                let m = self.emit_expr(message, ra);
352                self.emit(&format!("tsend {}, {}           ; send msg to agent", reg(t), reg(m)));
353            }
354
355            Stmt::FieldSet { object, field: _, value } => {
356                let obj_reg = ra.get(object);
357                let val_reg = self.emit_expr(value, ra);
358                self.emit(&format!("tst   {}, 0({})        ; field store", reg(val_reg), reg(obj_reg)));
359            }
360
361            Stmt::IndexSet { object, row, col: _, value } => {
362                let obj_reg = ra.get(object);
363                let row_reg = self.emit_expr(row, ra);
364                let val_reg = self.emit_expr(value, ra);
365                let addr    = ra.scratch();
366                self.emit(&format!("tadd  {}, {}, {}       ; tensor index addr", reg(addr), reg(obj_reg), reg(row_reg)));
367                self.emit(&format!("tst   {}, 0({})        ; tensor store", reg(val_reg), reg(addr)));
368            }
369
370            Stmt::Decorated { stmt, .. } => self.emit_stmt(stmt, ra),
371            Stmt::Use { .. } => {}
372            Stmt::FromImport { .. } => {}
373        }
374    }
375
376    /// Emit an expression, returning the register number containing the result.
377    fn emit_expr(&mut self, expr: &Expr, ra: &mut RegAlloc) -> u8 {
378        match expr {
379            Expr::TritLiteral(v) => {
380                let r = ra.scratch();
381                self.emit(&format!("tldi  {}, {}          ; trit literal", reg(r), v));
382                r
383            }
384
385            Expr::IntLiteral(v) => {
386                let r = ra.scratch();
387                self.emit(&format!("tlii  {}, {}          ; int literal", reg(r), v));
388                r
389            }
390
391            Expr::FloatLiteral(v) => {
392                let r = ra.scratch();
393                self.emit(&format!("tlii  {}, {}          ; float (truncated to int)", reg(r), *v as i64));
394                r
395            }
396
397            Expr::StringLiteral(_) => {
398                let r = ra.scratch();
399                self.emit(&format!("tlii  {}, 0           ; string (addr architecture defined)", reg(r)));
400                r
401            }
402
403            Expr::Ident(name) => {
404                ra.get(name)
405            }
406
407            Expr::BinaryOp { op, lhs, rhs } => {
408                let lreg = self.emit_expr(lhs, ra);
409                let rreg = self.emit_expr(rhs, ra);
410                let dest = ra.scratch();
411                let mnemonic = match op {
412                    BinOp::Add          => "tadd",
413                    BinOp::Sub          => "tsub",
414                    BinOp::Mul          => "tmul",
415                    BinOp::Div          => "tdiv",
416                    BinOp::Mod          => "tmod",
417                    BinOp::Equal        => "teq",
418                    BinOp::NotEqual     => "tne",
419                    BinOp::Less         => "tlt",
420                    BinOp::Greater      => "tgt",
421                    BinOp::LessEqual    => "tle",
422                    BinOp::GreaterEqual => "tge",
423                    BinOp::And          => "tcons",
424                    BinOp::Or           => "tmax",
425                };
426                self.emit(&format!("{:<6}{}, {}, {}", mnemonic, reg(dest), reg(lreg), reg(rreg)));
427                dest
428            }
429
430            Expr::UnaryOp { op: UnOp::Neg, expr } => {
431                let src  = self.emit_expr(expr, ra);
432                let dest = ra.scratch();
433                self.emit(&format!("tnot  {}, {}", reg(dest), reg(src)));
434                dest
435            }
436
437            Expr::Call { callee, args } => {
438                // Push args into a0..aN by convention
439                for (i, arg) in args.iter().enumerate() {
440                    let r = self.emit_expr(arg, ra);
441                    // Move into argument registers a0+ (we use t10+ as argument passing regs)
442                    let arg_reg = 10 + i as u8;
443                    if r != arg_reg {
444                        self.emit(&format!("tadd  t{}, {}, zero  ; arg {}", arg_reg, reg(r), i));
445                    }
446                }
447                // Builtin prints
448                match callee.as_str() {
449                    "print" | "debug_print" => {
450                        if !args.is_empty() {
451                            let r = self.emit_expr(&args[0], ra);
452                            self.emit(&format!("tprint {}", reg(r)));
453                        }
454                        return 0; // zero reg
455                    }
456                    "opent" => { // opent(path, mode) -> handle
457                        if args.len() == 2 {
458                            let r_path = self.emit_expr(&args[0], ra);
459                            let r_mode = self.emit_expr(&args[1], ra);
460                            self.emit(&format!("tpush {}", reg(r_path)));
461                            self.emit(&format!("tpush {}", reg(r_mode)));
462                            self.emit("topent");
463                            self.emit("tpop t2"); // return handle in t2
464                        }
465                        return 2;
466                    }
467                    "readt" => { // readt(handle) -> trit
468                        if !args.is_empty() {
469                            let r_handle = self.emit_expr(&args[0], ra);
470                            self.emit(&format!("tpush {}", reg(r_handle)));
471                            self.emit("treadt");
472                            self.emit("tpop t2");
473                        }
474                        return 2;
475                    }
476                    "writet" => { // writet(handle, trit) -> void
477                        if args.len() == 2 {
478                            let r_handle = self.emit_expr(&args[0], ra);
479                            let r_trit = self.emit_expr(&args[1], ra);
480                            self.emit(&format!("tpush {}", reg(r_handle)));
481                            self.emit(&format!("tpush {}", reg(r_trit)));
482                            self.emit("twritet");
483                        }
484                        return 0;
485                    }
486                    _ => {}
487                }
488                self.emit(&format!("call  {}", callee));
489                2 // return value in t2 by convention
490            }
491
492            Expr::Cast { expr, .. } => {
493                self.emit_expr(expr, ra) // transparent pass-through
494            }
495
496            Expr::FieldAccess { object, field: _ } => {
497                let obj_reg = self.emit_expr(object, ra);
498                let dest    = ra.scratch();
499                self.emit(&format!("tld   {}, 0({})       ; field load", reg(dest), reg(obj_reg)));
500                dest
501            }
502
503            Expr::Index { object, row, col: _ } => {
504                let obj_reg = self.emit_expr(object, ra);
505                let row_reg = self.emit_expr(row, ra);
506                let addr    = ra.scratch();
507                let dest    = ra.scratch();
508                self.emit(&format!("tadd  {}, {}, {}      ; tensor index", reg(addr), reg(obj_reg), reg(row_reg)));
509                self.emit(&format!("tld   {}, 0({})       ; tensor load", reg(dest), reg(addr)));
510                dest
511            }
512
513            Expr::TritTensorLiteral(vals) => {
514                let base = ra.scratch();
515                self.emit(&format!("tlii  {}, 0           ; tensor literal (addr)", reg(base)));
516                for (i, v) in vals.iter().enumerate() {
517                    let tmp = ra.scratch();
518                    self.emit(&format!("tldi  {}, {}", reg(tmp), v));
519                    self.emit(&format!("tst   {}, {}({})   ; tensor[{}]", reg(tmp), i, reg(base), i));
520                }
521                base
522            }
523
524            Expr::Spawn { agent_name, .. } => {
525                let r = ra.scratch();
526                self.emit(&format!("tspawn {}, {}         ; spawn agent", reg(r), agent_name));
527                r
528            }
529
530            Expr::Await { target } => {
531                let t = self.emit_expr(target, ra);
532                let r = ra.scratch();
533                self.emit(&format!("tawait {}, {}         ; await agent", reg(r), reg(t)));
534                r
535            }
536
537            Expr::Propagate { expr } => {
538                let src = self.emit_expr(expr, ra);
539                let lbl = self.fresh_label("prop_ok");
540                // If src == -1 (reject), return -1 immediately
541                self.emit(&format!("bneg  {}, .L_prop_ret_{}", reg(src), self.label_counter));
542                self.emit(&format!("j     {}", lbl));
543                self.emit_label(&format!(".L_prop_ret_{}", self.label_counter - 1));
544                self.emit("tldi  t2, -1         ; propagate reject");
545                self.emit("ret");
546                self.emit_label(&lbl);
547                src
548            }
549
550            Expr::NodeId => {
551                let r = ra.scratch();
552                self.emit(&format!("tnodeid {}", reg(r)));
553                r
554            }
555
556            Expr::StructLiteral { .. } => {
557                let r = ra.scratch();
558                self.emit(&format!("tldi  {}, 0           ; struct literal (dummy)", reg(r)));
559                r
560            }
561        }
562    }
563}
564
565// ── Public API ────────────────────────────────────────────────────────────────
566
567/// Emit a ternlang `Program` as TERN-compatible assembly text.
568pub fn emit_tern_asm(program: &Program) -> String {
569    TernAsmEmitter::new().emit_program(program)
570}