gantz_core 0.4.1

//! Mechanical emission of the IR (`ir.rs`) to Steel.
//!
//! Bodies are built structurally as [`Sexp`]s and parsed exactly once per
//! top-level definition. Only primitive base-engine forms are emitted (`if`,
//! `begin`, `define`, `define-values`, `let`, `set!`) - the VM runs
//! `Engine::new_base()`, which has no `cond`/`and` prelude.
//!
//! Emission contracts (shared with the node fns generated by
//! `codegen/node_fn.rs`):
//!
//! - a node fn returns its single connected output raw, or a `(list ...)` of
//!   its connected outputs (ascending index) when more than one;
//! - a branching node fn returns `(list branch-ix value)` where `value` is
//!   shaped by the selected arm's active output count;
//! - a stateful node fn takes `state` as its trailing arg and returns
//!   `(list output state')`.

use crate::{
    GRAPH_STATE,
    compile::{
        ir::{Arg, Arm, Atom, Body, NodeCall, Step, Subject, Tail, Var},
        names::{join_name, node_fn_name, pair_name, var_name},
    },
    node,
};
use std::fmt;
use steel::{parser::ast::ExprKind, steel_vm::engine::Engine};

/// The sentinel value marking an outlet that did not fire (see
/// [`Atom::Unfired`]).
const UNFIRED: &str = "'%gantz-unfired";

/// A structurally built s-expression, displayed as Steel source.
#[derive(Clone, Debug)]
pub(crate) enum Sexp {
    /// An atom, printed verbatim.
    A(String),
    /// A parenthesised list.
    L(Vec<Sexp>),
}

/// Context fixed across one body's emission.
pub(crate) struct Cx<'a> {
    /// The path of the graph level being emitted, for node fn naming.
    pub path: &'a [node::Id],
}

impl fmt::Display for Sexp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Sexp::A(s) => write!(f, "{s}"),
            Sexp::L(items) => {
                write!(f, "(")?;
                for (i, item) in items.iter().enumerate() {
                    if i > 0 {
                        write!(f, " ")?;
                    }
                    write!(f, "{item}")?;
                }
                write!(f, ")")
            }
        }
    }
}

/// Shorthand for an atom.
pub(crate) fn a(s: impl Into<String>) -> Sexp {
    Sexp::A(s.into())
}

/// Shorthand for a list.
pub(crate) fn l(items: impl IntoIterator<Item = Sexp>) -> Sexp {
    Sexp::L(items.into_iter().collect())
}

/// The unit value `'()`.
pub(crate) fn unit() -> Sexp {
    a("'()")
}

fn atom_sexp(atom: &Atom) -> Sexp {
    match atom {
        Atom::Var(v) => a(var_name(v)),
        Atom::Unit => unit(),
        Atom::Unfired => a(UNFIRED),
    }
}

fn arg_sexp(arg: &Arg) -> Sexp {
    match arg {
        Arg::One(atom) => atom_sexp(atom),
        Arg::List(atoms) => l(std::iter::once(a("list")).chain(atoms.iter().map(atom_sexp))),
    }
}

/// Bind `dst` from `value`: a plain define for one var, or - for several - a
/// temporary bound to the `(list ...)` value, destructured by index.
///
/// `define-values` is deliberately avoided: steel-core 0.7's analysis pass
/// panics on fn bodies mixing several `define-values` with local fn defines
/// (our join points).
fn define_sexps(dst: &[Var], value: Sexp) -> Vec<Sexp> {
    match dst {
        [] => vec![value],
        [v] => vec![l([a("define"), a(var_name(v)), value])],
        _ => {
            let tmp = format!("%vals-{}", var_name(&dst[0]));
            let mut stmts = vec![l([a("define"), a(tmp.clone()), value])];
            for (k, v) in dst.iter().enumerate() {
                stmts.push(l([
                    a("define"),
                    a(var_name(v)),
                    l([a("list-ref"), a(tmp.clone()), a(k.to_string())]),
                ]));
            }
            stmts
        }
    }
}

/// The call expression for a node fn, state-wrapped when stateful.
fn call_sexp(cx: &Cx, call: &NodeCall) -> Sexp {
    let inputs = node::Conns::try_from_iter(call.args.iter().map(Option::is_some))
        .expect("validated NodeCall arg count exceeds Conns::MAX");
    let node_path: Vec<node::Id> = cx.path.iter().copied().chain([call.node]).collect();
    let fn_name = node_fn_name(&node_path, &inputs, &call.outputs);
    let mut items = vec![a(fn_name)];
    items.extend(call.args.iter().flatten().map(arg_sexp));
    if call.stateful {
        items.push(a("state"));
    }
    let call_expr = Sexp::L(items);
    if call.stateful {
        state_wrapped(call.node, call_expr)
    } else {
        call_expr
    }
}

/// Wrap a stateful node's call: read its state from the graph-state hashmap,
/// run the call, write the returned state back, yield the output.
fn state_wrapped(node: node::Id, call_expr: Sexp) -> Sexp {
    let key = a(format!("'{node}"));
    l([
        a("let"),
        l([l([
            a("state"),
            l([a("hash-ref"), a(GRAPH_STATE), key.clone()]),
        ])]),
        l([
            a("let"),
            l([l([a("results"), call_expr])]),
            l([
                a("let"),
                l([
                    l([a("output"), l([a("car"), a("results")])]),
                    l([a("newstate"), l([a("car"), l([a("cdr"), a("results")])])]),
                ]),
                l([
                    a("set!"),
                    a(GRAPH_STATE),
                    l([a("hash-insert"), a(GRAPH_STATE), key, a("newstate")]),
                ]),
                a("output"),
            ]),
        ]),
    ])
}

/// The statements of a body: its steps followed by exactly one tail
/// expression (the body's value).
pub(crate) fn body_sexps(cx: &Cx, body: &Body) -> Vec<Sexp> {
    let mut stmts = Vec::with_capacity(body.steps.len() + 1);
    for step in &body.steps {
        match step {
            Step::Node { dst, call } => {
                stmts.extend(define_sexps(dst, call_sexp(cx, call)));
            }
            Step::DelayRead { node } => {
                let var = Var::Output {
                    node: *node,
                    output: 0,
                };
                stmts.push(l([
                    a("define"),
                    a(var_name(&var)),
                    l([a("hash-ref"), a(GRAPH_STATE), a(format!("'{node}"))]),
                ]));
            }
            Step::DelayWrite { node, arg } => {
                stmts.push(l([
                    a("set!"),
                    a(GRAPH_STATE),
                    l([
                        a("hash-insert"),
                        a(GRAPH_STATE),
                        a(format!("'{node}")),
                        arg_sexp(arg),
                    ]),
                ]));
            }
            Step::Join(join) => {
                let mut sig = vec![a(join_name(join.id))];
                sig.extend(join.params.iter().map(|p| a(var_name(p))));
                let mut items = vec![a("define"), Sexp::L(sig)];
                items.extend(body_sexps(cx, &join.body));
                stmts.push(Sexp::L(items));
            }
            Step::Branch { subject, dst, arms } => {
                let pair = match subject {
                    Subject::Call(call) => {
                        let pair = pair_name(call.node);
                        stmts.push(l([a("define"), a(pair.clone()), call_sexp(cx, call)]));
                        pair
                    }
                    // Bound by enclosing glue.
                    Subject::PreBound { node } => pair_name(*node),
                };
                stmts.extend(define_sexps(dst, dispatch_sexp(cx, &pair, arms)));
            }
        }
    }
    stmts.push(tail_sexp(&body.tail));
    stmts
}

/// The dispatch expression over a branch's arms: nested `if`s on the pair's
/// branch index, the last arm as the exhaustive fallthrough.
fn dispatch_sexp(cx: &Cx, pair: &str, arms: &[Arm]) -> Sexp {
    let ix_expr = || l([a("list-ref"), a(pair), a("0")]);
    let mut expr = unit();
    for (pos, arm) in arms.iter().enumerate().rev() {
        let arm_expr = arm_sexp(cx, pair, arm);
        // The last arm is the exhaustive fallthrough.
        if pos + 1 == arms.len() {
            expr = arm_expr;
        } else {
            let test = l([a("="), a(arm.ix.to_string()), ix_expr()]);
            expr = l([a("if"), test, arm_expr, expr]);
        }
    }
    expr
}

/// One arm's expression: bind the arm's active outputs from the pair's value,
/// then the arm body. Wrapped in `(let () ...)` so the defines form a body.
fn arm_sexp(cx: &Cx, pair: &str, arm: &Arm) -> Sexp {
    let value = || l([a("list-ref"), a(pair), a("1")]);
    let binds = match arm.binds.as_slice() {
        [] => vec![],
        binds => define_sexps(binds, value()),
    };
    let body = body_sexps(cx, &arm.body);
    match (binds.is_empty(), arm.body.steps.is_empty()) {
        // No bindings or local steps: the tail expression alone suffices.
        (true, true) => body.into_iter().next().expect("body yields a tail expr"),
        _ => {
            let mut items = vec![a("let"), l([])];
            items.extend(binds);
            items.extend(body);
            Sexp::L(items)
        }
    }
}

/// The tail expression: the yielded value(s) or a join tail-call.
fn tail_sexp(tail: &Tail) -> Sexp {
    match tail {
        Tail::Ret(atoms) => match atoms.as_slice() {
            [] => unit(),
            [atom] => atom_sexp(atom),
            atoms => l(std::iter::once(a("list")).chain(atoms.iter().map(atom_sexp))),
        },
        Tail::Jump { join, args } => {
            l(std::iter::once(a(join_name(*join))).chain(args.iter().map(atom_sexp)))
        }
    }
}

/// The result expression of a lowered level: the outlet values shaped by
/// activation pattern.
///
/// With fewer than two patterns the level never branches externally: the
/// result is the outlet values directly (raw for one outlet, a `(list ...)`
/// for several, `'()` for none). Otherwise the result is
/// `(list branch-ix value)`: the fired signature (sum of `2^i` for each fired
/// outlet, conditional outlets tested against the unfired sentinel) selects
/// the pattern index, and the value is shaped by that pattern's active
/// outlets - mirroring the `Node::branches` contract of the flow pipeline.
pub(crate) fn level_result_sexp(
    outlets: &[crate::compile::lower::OutletVal],
    patterns: &[node::Conns],
) -> Sexp {
    let atom = |o: &crate::compile::lower::OutletVal| match o.atom {
        Some(ref atom) => atom_sexp(atom),
        None => unit(),
    };
    let shaped = |active: Vec<&crate::compile::lower::OutletVal>| match active.as_slice() {
        [] => unit(),
        [o] => atom(o),
        os => l(std::iter::once(a("list")).chain(os.iter().map(|o| atom(o)))),
    };

    if patterns.len() < 2 {
        return shaped(outlets.iter().collect());
    }

    // The fired signature: constant terms for unconditionally-fired outlets,
    // a sentinel test for conditional ones.
    let mut base: u128 = 0;
    let mut terms: Vec<Sexp> = Vec::new();
    for (i, o) in outlets.iter().enumerate() {
        match (o.atom, o.conditional) {
            (None, _) => {}
            (Some(_), false) => base += 1 << i,
            (Some(atom), true) => terms.push(l([
                a("if"),
                l([a("equal?"), atom_sexp(&atom), a(UNFIRED)]),
                a("0"),
                a((1u128 << i).to_string()),
            ])),
        }
    }
    let sig_expr = match (base, terms.len()) {
        (b, 0) => a(b.to_string()),
        (0, 1) => terms.pop().unwrap(),
        (b, _) => {
            let mut items = vec![a("+"), a(b.to_string())];
            items.extend(terms);
            Sexp::L(items)
        }
    };

    // The unique signature of a pattern's active outlets.
    let signature = |conns: &node::Conns| -> u128 {
        (0..outlets.len())
            .filter(|&i| conns.get(i).unwrap_or(false))
            .map(|i| 1u128 << i)
            .sum()
    };
    let pattern_value = |conns: &node::Conns| -> Sexp {
        let active: Vec<&crate::compile::lower::OutletVal> = outlets
            .iter()
            .enumerate()
            .filter_map(|(i, o)| conns.get(i).unwrap_or(false).then_some(o))
            .collect();
        shaped(active)
    };

    // Nested `if` over patterns; the last is the exhaustive fallthrough.
    let last = patterns.len() - 1;
    let mut expr = l([
        a("list"),
        a(last.to_string()),
        pattern_value(&patterns[last]),
    ]);
    for k in (0..last).rev() {
        expr = l([
            a("if"),
            l([
                a("="),
                a("%gantz-sig"),
                a(signature(&patterns[k]).to_string()),
            ]),
            l([a("list"), a(k.to_string()), pattern_value(&patterns[k])]),
            expr,
        ]);
    }
    l([a("let"), l([l([a("%gantz-sig"), sig_expr])]), expr])
}

/// Assemble a top-level fn definition from glue statements.
pub(crate) fn fn_def(name: &str, params: &[String], stmts: Vec<Sexp>) -> ExprKind {
    let mut sig = vec![a(name)];
    sig.extend(params.iter().map(|p| a(p.clone())));
    let mut items = vec![a("define"), Sexp::L(sig)];
    items.extend(stmts);
    parse_one(&Sexp::L(items).to_string())
}

/// Parse one emitted top-level definition into the Steel AST.
pub(crate) fn parse_one(src: &str) -> ExprKind {
    Engine::emit_ast(src)
        .unwrap_or_else(|e| panic!("emitted Steel failed to parse: {e}\n{src}"))
        .into_iter()
        .next()
        .expect("emitted Steel parsed to no expression")
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::compile::ir::{self, Join};
    use steel::SteelVal;

    fn out(node: node::Id, output: usize) -> Var {
        Var::Output { node, output }
    }

    fn call(node: node::Id, args: Vec<Option<Arg>>, outputs: &str) -> NodeCall {
        NodeCall {
            node,
            args,
            outputs: outputs.parse().unwrap(),
            stateful: false,
        }
    }

    /// Emit `body` as a no-arg fn alongside the given node fn definitions,
    /// run it in a base VM, and return the result.
    fn run(node_fns: &[&str], body: &Body) -> SteelVal {
        ir::validate(body, 1, &[]).unwrap();
        let cx = Cx { path: &[] };
        let mut src = node_fns.join(" ");
        // A test driver without state threading: just the body statements.
        let mut items = vec![a("define"), l([a("test-fn")])];
        items.extend(body_sexps(&cx, body));
        src.push_str(&format!(" {}", Sexp::L(items)));
        let mut vm = Engine::new_base();
        vm.run(src).unwrap();
        vm.run("(test-fn)".to_string())
            .unwrap()
            .into_iter()
            .next()
            .unwrap()
    }

    /// A linear chain: constant -> double, returning the result.
    #[test]
    fn linear_chain() {
        let body = Body {
            steps: vec![
                Step::Node {
                    dst: vec![out(0, 0)],
                    call: call(0, vec![], "1"),
                },
                Step::Node {
                    dst: vec![out(1, 0)],
                    call: call(1, vec![Some(Arg::One(Atom::Var(out(0, 0))))], "1"),
                },
            ],
            tail: Tail::Ret(vec![Atom::Var(out(1, 0))]),
        };
        let result = run(
            &[
                "(define (node-fn-0-o1) 21)",
                "(define (node-fn-1-i1-o1 input0) (* 2 input0))",
            ],
            &body,
        );
        assert_eq!(result, SteelVal::IntV(42));
    }

    /// A multi-output node destructures via define-values.
    #[test]
    fn multi_output_destructure() {
        let body = Body {
            steps: vec![
                Step::Node {
                    dst: vec![out(0, 0), out(0, 1)],
                    call: call(0, vec![], "11"),
                },
                Step::Node {
                    dst: vec![out(1, 0)],
                    call: call(
                        1,
                        vec![
                            Some(Arg::One(Atom::Var(out(0, 0)))),
                            Some(Arg::One(Atom::Var(out(0, 1)))),
                        ],
                        "1",
                    ),
                },
            ],
            tail: Tail::Ret(vec![Atom::Var(out(1, 0))]),
        };
        let result = run(
            &[
                "(define (node-fn-0-o11) (list 40 2))",
                "(define (node-fn-1-i11-o1 input0 input1) (+ input0 input1))",
            ],
            &body,
        );
        assert_eq!(result, SteelVal::IntV(42));
    }

    /// Multiple sources into one input pass as a `(list ...)` arg.
    #[test]
    fn list_arg() {
        let body = Body {
            steps: vec![
                Step::Node {
                    dst: vec![out(0, 0)],
                    call: call(0, vec![], "1"),
                },
                Step::Node {
                    dst: vec![out(1, 0)],
                    call: call(
                        1,
                        vec![Some(Arg::List(vec![
                            Atom::Var(out(0, 0)),
                            Atom::Var(out(0, 0)),
                        ]))],
                        "1",
                    ),
                },
            ],
            tail: Tail::Ret(vec![Atom::Var(out(1, 0))]),
        };
        let result = run(
            &[
                "(define (node-fn-0-o1) 21)",
                "(define (node-fn-1-i1-o1 input0) (+ (list-ref input0 0) (list-ref input0 1)))",
            ],
            &body,
        );
        assert_eq!(result, SteelVal::IntV(42));
    }

    /// A branch whose arms jump to a join; the join's param carries the
    /// arm-varying value and the branch exports the join's result.
    #[test]
    fn branch_join_export() {
        let param = Var::Input { node: 2, input: 0 };
        let export = out(2, 0);
        let arm = |ix: usize, output: usize| Arm {
            ix,
            binds: vec![out(1, output)],
            body: Body {
                steps: vec![],
                tail: Tail::Jump {
                    join: 2,
                    args: vec![Atom::Var(out(1, output))],
                },
            },
        };
        let body = Body {
            steps: vec![
                Step::Node {
                    dst: vec![out(0, 0)],
                    call: call(0, vec![], "1"),
                },
                Step::Join(Join {
                    id: 2,
                    params: vec![param],
                    rec: false,
                    body: Body {
                        steps: vec![Step::Node {
                            dst: vec![out(2, 0)],
                            call: call(2, vec![Some(Arg::One(Atom::Var(param)))], "1"),
                        }],
                        tail: Tail::Ret(vec![Atom::Var(out(2, 0))]),
                    },
                }),
                Step::Branch {
                    subject: Subject::Call(call(
                        1,
                        vec![Some(Arg::One(Atom::Var(out(0, 0))))],
                        "11",
                    )),
                    dst: vec![export],
                    arms: vec![arm(0, 0), arm(1, 1)],
                },
            ],
            tail: Tail::Ret(vec![Atom::Var(export)]),
        };
        let node_fns = [
            // Selects arm 1, yielding 20.
            "(define (node-fn-0-o1) 20)",
            // Arm 0 iff input < 10, arm 1 otherwise; value passes through.
            "(define (node-fn-1-i1-o11 input0)
               (if (< input0 10) (list 0 input0) (list 1 input0)))",
            "(define (node-fn-2-i1-o1 input0) (+ input0 1))",
        ];
        let result = run(&node_fns, &body);
        assert_eq!(result, SteelVal::IntV(21));
    }

    /// A rec join expresses a loop: iterate until the count reaches zero.
    /// (The shape Phase 3's iterate-until-branch lowering will produce.)
    #[test]
    fn rec_join_loop() {
        let acc = Var::Input { node: 1, input: 0 };
        let n = Var::Input { node: 1, input: 1 };
        let body = Body {
            steps: vec![Step::Join(Join {
                id: 1,
                params: vec![acc, n],
                rec: true,
                body: Body {
                    steps: vec![Step::Branch {
                        subject: Subject::Call(call(
                            1,
                            vec![Some(Arg::One(Atom::Var(acc))), Some(Arg::One(Atom::Var(n)))],
                            "11",
                        )),
                        dst: vec![out(1, 0)],
                        arms: vec![
                            // Continue: jump back with updated acc/n.
                            Arm {
                                ix: 0,
                                binds: vec![out(9, 0), out(9, 1)],
                                body: Body {
                                    steps: vec![],
                                    tail: Tail::Jump {
                                        join: 1,
                                        args: vec![Atom::Var(out(9, 0)), Atom::Var(out(9, 1))],
                                    },
                                },
                            },
                            // Exit: yield the accumulator.
                            Arm {
                                ix: 1,
                                binds: vec![out(9, 0)],
                                body: Body {
                                    steps: vec![],
                                    tail: Tail::Ret(vec![Atom::Var(out(9, 0))]),
                                },
                            },
                        ],
                    }],
                    tail: Tail::Ret(vec![Atom::Var(out(1, 0))]),
                },
            })],
            tail: Tail::Jump {
                join: 1,
                args: vec![Atom::Unit, Atom::Unit],
            },
        };
        let node_fns = ["(define (node-fn-1-i11-o11 input0 input1)
            (let ((acc (if (number? input0) input0 0))
                  (n (if (number? input1) input1 100000)))
              (if (= n 0)
                  (list 1 acc)
                  (list 0 (list (+ acc 1) (- n 1))))))"];
        let result = run(&node_fns, &body);
        assert_eq!(result, SteelVal::IntV(100000));
    }
}