ternlang-core 0.3.3

// SPDX-License-Identifier: LGPL-3.0-or-later
// Ternlang — RFI-IRFOS Ternary Intelligence Stack
// Copyright (C) 2026 RFI-IRFOS
//
// tern_asm.rs — TERN-compatible assembly emitter
//
// Produces TERN assembly text from a ternlang `Program` AST.
//
// TERN assembly is a RISC-V-inspired balanced ternary assembly language.
// This emitter generates output that is structurally compatible with the
// Tern Systems TERN specification (https://github.com/Tern-Computer).
//
// Register file: t0–t31 (32 ternary registers, each holds one trit value)
// Special:       sp (stack pointer), ra (return address), zero (always 0)
//
// # Instruction summary
//
//   Arithmetic / logic
//     tadd  rd, rs1, rs2    ; rd = rs1 + rs2  (ternary saturating add)
//     tsub  rd, rs1, rs2    ; rd = rs1 - rs2
//     tmul  rd, rs1, rs2    ; rd = rs1 * rs2
//     tdiv  rd, rs1, rs2    ; rd = rs1 / rs2
//     tmod  rd, rs1, rs2    ; rd = rs1 mod rs2
//     tnot  rd, rs1         ; rd = -rs1        (ternary invert / negate)
//     tcons rd, rs1, rs2    ; rd = consensus(rs1, rs2)  — ternary AND / min
//     tmax  rd, rs1, rs2    ; rd = max(rs1, rs2)        — ternary OR
//
//   Comparison (result is a trit: +1 true, -1 false, 0 hold/equal)
//     teq   rd, rs1, rs2    ; rd = (rs1 == rs2) ? +1 : -1
//     tlt   rd, rs1, rs2    ; rd = (rs1 <  rs2) ? +1 : -1
//     tgt   rd, rs1, rs2    ; rd = (rs1 >  rs2) ? +1 : -1
//     tle   rd, rs1, rs2    ; rd = (rs1 <= rs2) ? +1 : -1
//     tge   rd, rs1, rs2    ; rd = (rs1 >= rs2) ? +1 : -1
//     tne   rd, rs1, rs2    ; rd = (rs1 != rs2) ? +1 : -1
//
//   Immediate load
//     tldi  rd, imm         ; rd = imm  (imm ∈ {-1, 0, 1})
//     tlii  rd, imm         ; rd = imm  (imm is an integer constant)
//
//   Control flow
//     j     label           ; unconditional jump
//     bpos  rs, label       ; jump if rs == +1  (affirm branch)
//     bzero rs, label       ; jump if rs ==  0  (hold branch)
//     bneg  rs, label       ; jump if rs == -1  (reject branch)
//     call  label           ; call subroutine (saves ra)
//     ret                   ; return (jumps to ra)
//
//   Memory (ternary word = 1 trit)
//     tld   rd, offset(rs)  ; load trit from memory[rs + offset]
//     tst   rs2, offset(rs1); store rs2 → memory[rs1 + offset]
//
//   I/O (runtime builtins)
//     tprint rd             ; print trit value of rd as "affirm"/"hold"/"reject"
//     tprintint rd          ; print integer value of rd

use crate::ast::*;
use std::collections::HashMap;

// ── Register allocator ────────────────────────────────────────────────────────

struct RegAlloc {
    map: HashMap<String, u8>,
    next: u8,
}

impl RegAlloc {
    fn new() -> Self { Self { map: HashMap::new(), next: 2 } } // t0=zero, t1=scratch
    fn alloc(&mut self, name: &str) -> u8 {
        if let Some(&r) = self.map.get(name) { return r; }
        let r = self.next;
        self.next += 1;
        self.map.insert(name.to_string(), r);
        r
    }
    fn get(&self, name: &str) -> u8 {
        *self.map.get(name).unwrap_or(&1) // fall back to t1 scratch
    }
    fn scratch(&mut self) -> u8 {
        let r = self.next;
        self.next += 1;
        r
    }
    fn snapshot(&self) -> HashMap<String, u8> {
        self.map.clone()
    }
    fn restore(&mut self, snap: HashMap<String, u8>, next: u8) {
        self.map = snap;
        self.next = next;
    }
}

fn reg(n: u8) -> String {
    match n {
        0 => "zero".to_string(),
        1 => "t1".to_string(),
        _ => format!("t{}", n),
    }
}

// ── Emitter ───────────────────────────────────────────────────────────────────

pub struct TernAsmEmitter {
    out: Vec<String>,
    label_counter: usize,
}

impl TernAsmEmitter {
    pub fn new() -> Self {
        Self { out: Vec::new(), label_counter: 0 }
    }

    fn emit(&mut self, line: &str) {
        self.out.push(format!("    {}", line));
    }

    fn emit_label(&mut self, label: &str) {
        self.out.push(format!("{}:", label));
    }

    fn fresh_label(&mut self, prefix: &str) -> String {
        let l = format!(".L_{}_{}", prefix, self.label_counter);
        self.label_counter += 1;
        l
    }

    /// Emit the full program as a TERN ASM string.
    pub fn emit_program(&mut self, program: &Program) -> String {
        self.out.push("; Generated by ternlang TERN-ASM emitter — RFI-IRFOS".to_string());
        self.out.push("; TERN assembly (RISC-V-inspired balanced ternary)".to_string());
        self.out.push("; Spec compatible with: Tern Systems TERN / BTMC".to_string());
        self.out.push("".to_string());
        self.out.push(".section .text".to_string());
        self.out.push(".global main".to_string());
        self.out.push("".to_string());

        // Emit all functions
        for func in &program.functions {
            let mut ra = RegAlloc::new();
            self.emit_function(func, &mut ra);
            self.out.push("".to_string());
        }

        // Emit agent handlers
        for agent in &program.agents {
            for method in &agent.methods {
                let label = format!("{}__{}", agent.name, method.name);
                let mut ra = RegAlloc::new();
                self.emit_function_with_label(&label, method, &mut ra);
                self.out.push("".to_string());
            }
        }

        self.out.join("\n")
    }

    fn emit_function(&mut self, func: &Function, ra: &mut RegAlloc) {
        self.emit_function_with_label(&func.name, func, ra);
    }

    fn emit_function_with_label(&mut self, label: &str, func: &Function, ra: &mut RegAlloc) {
        self.emit_label(label);

        // Allocate parameter registers
        for (param_name, _) in &func.params {
            ra.alloc(param_name);
        }

        for stmt in &func.body {
            self.emit_stmt(stmt, ra);
        }

        // If no explicit return at end, emit ret
        let has_return = func.body.last().map(|s| matches!(s, Stmt::Return(_))).unwrap_or(false);
        if !has_return {
            self.emit("ret");
        }
    }

    fn emit_stmt(&mut self, stmt: &Stmt, ra: &mut RegAlloc) {
        match stmt {
            Stmt::Let { name, value, .. } => {
                if let Expr::StructLiteral { fields, .. } = value {
                    for (f_name, f_val) in fields {
                        let f_dest = ra.alloc(&format!("{}.{}", name, f_name));
                        let f_src = self.emit_expr(f_val, ra);
                        if f_src != f_dest {
                            self.emit(&format!("tadd  {}, {}, zero   ; struct field init", reg(f_dest), reg(f_src)));
                        }
                    }
                    // The main variable is just a dummy
                    let dest = ra.alloc(name);
                    self.emit(&format!("tldi  {}, 0           ; struct root dummy", reg(dest)));
                } else {
                    let dest = ra.alloc(name);
                    let src = self.emit_expr(value, ra);
                    if src != dest {
                        self.emit(&format!("tadd  {}, {}, zero   ; {} = {}", reg(dest), reg(src), name, name));
                    }
                }
                // else value was emitted directly into dest
            }

            Stmt::Set { name, value } => {
                let dest = ra.get(name);
                let src = self.emit_expr(value, ra);
                if src != dest {
                    self.emit(&format!("tadd  {}, {}, zero   ; {} = expr", reg(dest), reg(src), name));
                }
            }

            Stmt::Return(expr) => {
                let src = self.emit_expr(expr, ra);
                // Convention: return value in t2
                if src != 2 {
                    self.emit(&format!("tadd  t2, {}, zero   ; return value", reg(src)));
                }
                self.emit("ret");
            }

            Stmt::Expr(expr) => {
                self.emit_expr(expr, ra);
            }

            Stmt::Block(stmts) => {
                let snap = ra.snapshot();
                let next = ra.next;
                for s in stmts { self.emit_stmt(s, ra); }
                ra.restore(snap, next);
            }

            Stmt::IfTernary { condition, on_pos, on_zero, on_neg } => {
                let cond_reg = self.emit_expr(condition, ra);
                let lbl_pos  = self.fresh_label("pos");
                let lbl_zero = self.fresh_label("zero");
                let lbl_neg  = self.fresh_label("neg");
                let lbl_end  = self.fresh_label("end");

                self.emit(&format!("bpos  {}, {}", reg(cond_reg), lbl_pos));
                self.emit(&format!("bzero {}, {}", reg(cond_reg), lbl_zero));
                self.emit(&format!("j     {}", lbl_neg));

                self.emit_label(&lbl_pos);
                self.emit_stmt(on_pos, ra);
                self.emit(&format!("j     {}", lbl_end));

                self.emit_label(&lbl_zero);
                self.emit_stmt(on_zero, ra);
                self.emit(&format!("j     {}", lbl_end));

                self.emit_label(&lbl_neg);
                self.emit_stmt(on_neg, ra);

                self.emit_label(&lbl_end);
            }

            Stmt::Match { condition, arms } => {
                let cond_reg = self.emit_expr(condition, ra);
                let lbl_end  = self.fresh_label("match_end");
                let mut arm_labels: Vec<(i64, String)> = Vec::new();

                for (pattern, _) in arms {
                    let val = match pattern {
                        Pattern::Int(v) => *v,
                        Pattern::Trit(t) => *t as i64,
                        Pattern::Float(f) => *f as i64,
                    };
                    arm_labels.push((val, self.fresh_label(&format!("arm_{}", val))));
                }

                // branch dispatch
                for (val, lbl) in &arm_labels {
                    let tmp = ra.scratch();
                    self.emit(&format!("tlii  {}, {}", reg(tmp), val));
                    let cmp = ra.scratch();
                    self.emit(&format!("teq   {}, {}, {}", reg(cmp), reg(cond_reg), reg(tmp)));
                    self.emit(&format!("bpos  {}, {}", reg(cmp), lbl));
                }
                self.emit(&format!("j     {}", lbl_end));

                for ((_, body_stmt), (_, lbl)) in arms.iter().zip(arm_labels.iter()) {
                    self.emit_label(lbl);
                    self.emit_stmt(body_stmt, ra);
                    self.emit(&format!("j     {}", lbl_end));
                }

                self.emit_label(&lbl_end);
            }

            Stmt::WhileTernary { condition, on_pos, on_zero, on_neg } => {
                let lbl_loop  = self.fresh_label("while");
                let lbl_pos   = self.fresh_label("wpos");
                let lbl_zero  = self.fresh_label("wzero");
                let lbl_neg   = self.fresh_label("wneg");
                let lbl_end   = self.fresh_label("wend");

                self.emit_label(&lbl_loop);
                let cond_reg = self.emit_expr(condition, ra);

                self.emit(&format!("bpos  {}, {}", reg(cond_reg), lbl_pos));
                self.emit(&format!("bzero {}, {}", reg(cond_reg), lbl_zero));
                self.emit(&format!("j     {}", lbl_neg));

                self.emit_label(&lbl_pos);
                self.emit_stmt(on_pos, ra);
                self.emit(&format!("j     {}", lbl_loop));

                self.emit_label(&lbl_zero);
                self.emit_stmt(on_zero, ra);
                self.emit(&format!("j     {}", lbl_loop));

                self.emit_label(&lbl_neg);
                self.emit_stmt(on_neg, ra);
                // neg branch exits loop

                self.emit_label(&lbl_end);
            }

            Stmt::Loop { body } => {
                let lbl_loop = self.fresh_label("loop");
                let lbl_end  = self.fresh_label("loop_end");

                self.emit_label(&lbl_loop);
                self.emit_stmt(body, ra);
                self.emit(&format!("j     {}", lbl_loop));
                self.emit_label(&lbl_end);
            }

            Stmt::ForIn { var, iter, body } => {
                // Simplified: emit a range-like loop if iter is a simple ident
                let iter_reg  = self.emit_expr(iter, ra);
                let var_reg   = ra.alloc(var);
                let idx_reg   = ra.scratch();
                let lbl_loop  = self.fresh_label("forin");
                let lbl_end   = self.fresh_label("forin_end");

                self.emit(&format!("tlii  {}, 0          ; for-in idx = 0", reg(idx_reg)));
                self.emit_label(&lbl_loop);
                // Load row from tensor iter at idx
                self.emit(&format!("tld   {}, 0({})      ; load row iter[idx]", reg(var_reg), reg(iter_reg)));
                self.emit_stmt(body, ra);
                self.emit(&format!("tadd  {0}, {0}, t1   ; idx++", reg(idx_reg)));
                self.emit(&format!("j     {}", lbl_loop));
                self.emit_label(&lbl_end);
            }

            Stmt::Break    => { self.emit("j     .L_break_arch_def ; break"); }
            Stmt::Continue => { self.emit("j     .L_continue_arch_def ; continue"); }

            Stmt::Send { target, message } => {
                let t = self.emit_expr(target, ra);
                let m = self.emit_expr(message, ra);
                self.emit(&format!("tsend {}, {}           ; send msg to agent", reg(t), reg(m)));
            }

            Stmt::FieldSet { object, field: _, value } => {
                let obj_reg = ra.get(object);
                let val_reg = self.emit_expr(value, ra);
                self.emit(&format!("tst   {}, 0({})        ; field store", reg(val_reg), reg(obj_reg)));
            }

            Stmt::IndexSet { object, row, col: _, value } => {
                let obj_reg = ra.get(object);
                let row_reg = self.emit_expr(row, ra);
                let val_reg = self.emit_expr(value, ra);
                let addr    = ra.scratch();
                self.emit(&format!("tadd  {}, {}, {}       ; tensor index addr", reg(addr), reg(obj_reg), reg(row_reg)));
                self.emit(&format!("tst   {}, 0({})        ; tensor store", reg(val_reg), reg(addr)));
            }

            Stmt::Decorated { stmt, .. } => self.emit_stmt(stmt, ra),
            Stmt::Use { .. } => {}
            Stmt::FromImport { .. } => {}
        }
    }

    /// Emit an expression, returning the register number containing the result.
    fn emit_expr(&mut self, expr: &Expr, ra: &mut RegAlloc) -> u8 {
        match expr {
            Expr::TritLiteral(v) => {
                let r = ra.scratch();
                self.emit(&format!("tldi  {}, {}          ; trit literal", reg(r), v));
                r
            }

            Expr::IntLiteral(v) => {
                let r = ra.scratch();
                self.emit(&format!("tlii  {}, {}          ; int literal", reg(r), v));
                r
            }

            Expr::FloatLiteral(v) => {
                let r = ra.scratch();
                self.emit(&format!("tlii  {}, {}          ; float (truncated to int)", reg(r), *v as i64));
                r
            }

            Expr::StringLiteral(_) => {
                let r = ra.scratch();
                self.emit(&format!("tlii  {}, 0           ; string (addr architecture defined)", reg(r)));
                r
            }

            Expr::Ident(name) => {
                ra.get(name)
            }

            Expr::BinaryOp { op, lhs, rhs } => {
                let lreg = self.emit_expr(lhs, ra);
                let rreg = self.emit_expr(rhs, ra);
                let dest = ra.scratch();
                let mnemonic = match op {
                    BinOp::Add          => "tadd",
                    BinOp::Sub          => "tsub",
                    BinOp::Mul          => "tmul",
                    BinOp::Div          => "tdiv",
                    BinOp::Mod          => "tmod",
                    BinOp::Equal        => "teq",
                    BinOp::NotEqual     => "tne",
                    BinOp::Less         => "tlt",
                    BinOp::Greater      => "tgt",
                    BinOp::LessEqual    => "tle",
                    BinOp::GreaterEqual => "tge",
                    BinOp::And          => "tcons",
                    BinOp::Or           => "tmax",
                };
                self.emit(&format!("{:<6}{}, {}, {}", mnemonic, reg(dest), reg(lreg), reg(rreg)));
                dest
            }

            Expr::UnaryOp { op: UnOp::Neg, expr } => {
                let src  = self.emit_expr(expr, ra);
                let dest = ra.scratch();
                self.emit(&format!("tnot  {}, {}", reg(dest), reg(src)));
                dest
            }

            Expr::Call { callee, args } => {
                // Push args into a0..aN by convention
                for (i, arg) in args.iter().enumerate() {
                    let r = self.emit_expr(arg, ra);
                    // Move into argument registers a0+ (we use t10+ as argument passing regs)
                    let arg_reg = 10 + i as u8;
                    if r != arg_reg {
                        self.emit(&format!("tadd  t{}, {}, zero  ; arg {}", arg_reg, reg(r), i));
                    }
                }
                // Builtin prints
                match callee.as_str() {
                    "print" | "debug_print" => {
                        if !args.is_empty() {
                            let r = self.emit_expr(&args[0], ra);
                            self.emit(&format!("tprint {}", reg(r)));
                        }
                        return 0; // zero reg
                    }
                    "opent" => { // opent(path, mode) -> handle
                        if args.len() == 2 {
                            let r_path = self.emit_expr(&args[0], ra);
                            let r_mode = self.emit_expr(&args[1], ra);
                            self.emit(&format!("tpush {}", reg(r_path)));
                            self.emit(&format!("tpush {}", reg(r_mode)));
                            self.emit("topent");
                            self.emit("tpop t2"); // return handle in t2
                        }
                        return 2;
                    }
                    "readt" => { // readt(handle) -> trit
                        if !args.is_empty() {
                            let r_handle = self.emit_expr(&args[0], ra);
                            self.emit(&format!("tpush {}", reg(r_handle)));
                            self.emit("treadt");
                            self.emit("tpop t2");
                        }
                        return 2;
                    }
                    "writet" => { // writet(handle, trit) -> void
                        if args.len() == 2 {
                            let r_handle = self.emit_expr(&args[0], ra);
                            let r_trit = self.emit_expr(&args[1], ra);
                            self.emit(&format!("tpush {}", reg(r_handle)));
                            self.emit(&format!("tpush {}", reg(r_trit)));
                            self.emit("twritet");
                        }
                        return 0;
                    }
                    _ => {}
                }
                self.emit(&format!("call  {}", callee));
                2 // return value in t2 by convention
            }

            Expr::Cast { expr, .. } => {
                self.emit_expr(expr, ra) // transparent pass-through
            }

            Expr::FieldAccess { object, field: _ } => {
                let obj_reg = self.emit_expr(object, ra);
                let dest    = ra.scratch();
                self.emit(&format!("tld   {}, 0({})       ; field load", reg(dest), reg(obj_reg)));
                dest
            }

            Expr::Index { object, row, col: _ } => {
                let obj_reg = self.emit_expr(object, ra);
                let row_reg = self.emit_expr(row, ra);
                let addr    = ra.scratch();
                let dest    = ra.scratch();
                self.emit(&format!("tadd  {}, {}, {}      ; tensor index", reg(addr), reg(obj_reg), reg(row_reg)));
                self.emit(&format!("tld   {}, 0({})       ; tensor load", reg(dest), reg(addr)));
                dest
            }

            Expr::TritTensorLiteral(vals) => {
                let base = ra.scratch();
                self.emit(&format!("tlii  {}, 0           ; tensor literal (addr)", reg(base)));
                for (i, v) in vals.iter().enumerate() {
                    let tmp = ra.scratch();
                    self.emit(&format!("tldi  {}, {}", reg(tmp), v));
                    self.emit(&format!("tst   {}, {}({})   ; tensor[{}]", reg(tmp), i, reg(base), i));
                }
                base
            }

            Expr::Spawn { agent_name, .. } => {
                let r = ra.scratch();
                self.emit(&format!("tspawn {}, {}         ; spawn agent", reg(r), agent_name));
                r
            }

            Expr::Await { target } => {
                let t = self.emit_expr(target, ra);
                let r = ra.scratch();
                self.emit(&format!("tawait {}, {}         ; await agent", reg(r), reg(t)));
                r
            }

            Expr::Propagate { expr } => {
                let src = self.emit_expr(expr, ra);
                let lbl = self.fresh_label("prop_ok");
                // If src == -1 (reject), return -1 immediately
                self.emit(&format!("bneg  {}, .L_prop_ret_{}", reg(src), self.label_counter));
                self.emit(&format!("j     {}", lbl));
                self.emit_label(&format!(".L_prop_ret_{}", self.label_counter - 1));
                self.emit("tldi  t2, -1         ; propagate reject");
                self.emit("ret");
                self.emit_label(&lbl);
                src
            }

            Expr::NodeId => {
                let r = ra.scratch();
                self.emit(&format!("tnodeid {}", reg(r)));
                r
            }

            Expr::StructLiteral { .. } => {
                let r = ra.scratch();
                self.emit(&format!("tldi  {}, 0           ; struct literal (dummy)", reg(r)));
                r
            }
        }
    }
}

// ── Public API ────────────────────────────────────────────────────────────────

/// Emit a ternlang `Program` as TERN-compatible assembly text.
pub fn emit_tern_asm(program: &Program) -> String {
    TernAsmEmitter::new().emit_program(program)
}