taxa-core 0.1.0

//! Safe metric-expression DSL → Polars `Expr`. The author-facing AST is the
//! same JSON the Python side used; here it compiles to a typed Polars
//! expression instead of a SQL string — so there is no identifier-quoting or
//! parameter-binding, and no injection surface at all (columns and literals are
//! data, never interpolated text).
//!
//! ```text
//! {"col": "gc_bases"}                       a column reference
//! {"lit": 0}                                a numeric/bool/null/str literal
//! {"op": "/", "args": [a, b]}               + - * /  (left-assoc n-ary)
//! {"op": ">=", "args": [a, b]}              >= > <= < == !=  (binary, → bool)
//! {"op": "and", "args": [a, b, ...]}        and / or  (n-ary bool fold), not (unary)
//! {"fn": "coalesce", "args": [a, b]}        an allow-listed function
//! {"cast": a, "as": "double"}               an allow-listed cast
//! ```
//!
//! Comparison and boolean ops yield a boolean `Expr`, so a single AST node can
//! express a row predicate (e.g. an [`Axis`](crate::manifest::Axis)'s
//! `row_filter`: `{"op": ">=", "args": [{"col": "mcap_usd"}, {"lit": 1e7}]}`).

use std::collections::HashSet;

use polars::prelude::*;
use serde_json::Value as Json;

#[derive(Debug, thiserror::Error)]
#[error("formula error: {0}")]
pub struct FormulaError(pub String);

fn err<T>(msg: impl Into<String>) -> Result<T, FormulaError> {
    Err(FormulaError(msg.into()))
}

const OPS: &[&str] = &["+", "-", "*", "/"];
/// Binary comparison operators — each yields a boolean `Expr`.
const CMP_OPS: &[&str] = &[">=", ">", "<=", "<", "==", "!="];
/// Boolean combinators — `and`/`or` are n-ary folds, `not` is unary.
const BOOL_OPS: &[&str] = &["and", "or", "not"];
/// Unary null predicates — each takes exactly 1 arg and yields a boolean `Expr`.
const NULL_OPS: &[&str] = &["is_null", "is_not_null"];

fn func_arity(name: &str) -> Option<(usize, Option<usize>)> {
    Some(match name {
        "coalesce" => (1, None),
        "nullif" => (2, Some(2)),
        "abs" => (1, Some(1)),
        "round" => (1, Some(2)),
        "least" => (2, None),
        "greatest" => (2, None),
        "ln" => (1, Some(1)),
        "log" => (1, Some(1)),
        "exp" => (1, Some(1)),
        "floor" => (1, Some(1)),
        "ceil" => (1, Some(1)),
        "sqrt" => (1, Some(1)),
        _ => return None,
    })
}

/// Compile an AST node to a Polars expression. `columns` is the source's real
/// column allow-list — an unknown column is an error.
pub fn compile_formula(node: &Json, columns: &HashSet<String>) -> Result<Expr, FormulaError> {
    let obj = match node.as_object() {
        Some(o) => o,
        None => return err(format!("expected an AST node (object), got {node}")),
    };

    if let Some(c) = obj.get("col") {
        let name = c
            .as_str()
            .ok_or_else(|| FormulaError("col must be a string".into()))?;
        if !columns.contains(name) {
            return err(format!("unknown column {name:?} (not in source schema)"));
        }
        return Ok(col(name));
    }

    if let Some(lit_node) = obj.get("lit") {
        return lit_expr(lit_node);
    }

    if let Some(op) = obj.get("op") {
        let op = op.as_str().unwrap_or("");
        // Arithmetic: left-assoc n-ary fold (≥2 args).
        if OPS.contains(&op) {
            let args = compile_args(obj, columns)?;
            let mut it = args.into_iter();
            let Some(mut acc) = it.next() else {
                return err(format!("operator {op:?} needs ≥2 args"));
            };
            let mut n = 1;
            for a in it {
                acc = match op {
                    "+" => acc + a,
                    "-" => acc - a,
                    "*" => acc * a,
                    _ => acc / a,
                };
                n += 1;
            }
            if n < 2 {
                return err(format!("operator {op:?} needs ≥2 args"));
            }
            return Ok(acc);
        }
        // Comparison: binary, yields a boolean Expr (the predicate building block).
        if CMP_OPS.contains(&op) {
            let args = compile_args(obj, columns)?;
            let mut it = args.into_iter();
            let (Some(a), Some(b), None) = (it.next(), it.next(), it.next()) else {
                return err(format!("comparison {op:?} needs exactly 2 args"));
            };
            return Ok(match op {
                ">=" => a.gt_eq(b),
                ">" => a.gt(b),
                "<=" => a.lt_eq(b),
                "<" => a.lt(b),
                "==" => a.eq(b),
                _ => a.neq(b),
            });
        }
        // Boolean combinators: and/or n-ary fold, not unary.
        if BOOL_OPS.contains(&op) {
            let args = compile_args(obj, columns)?;
            if op == "not" {
                let mut it = args.into_iter();
                let (Some(a), None) = (it.next(), it.next()) else {
                    return err("`not` needs exactly 1 arg");
                };
                return Ok(a.not());
            }
            let mut it = args.into_iter();
            let Some(mut acc) = it.next() else {
                return err(format!("operator {op:?} needs ≥2 args"));
            };
            let mut n = 1;
            for a in it {
                acc = if op == "and" { acc.and(a) } else { acc.or(a) };
                n += 1;
            }
            if n < 2 {
                return err(format!("operator {op:?} needs ≥2 args"));
            }
            return Ok(acc);
        }
        // Null predicates: unary, yield a boolean Expr (`col.is_null()` /
        // `col.is_not_null()`). Lets a row_filter drop nulls directly instead of a
        // `coalesce(col,'sentinel') != 'sentinel'` workaround.
        if NULL_OPS.contains(&op) {
            let args = compile_args(obj, columns)?;
            let mut it = args.into_iter();
            let (Some(a), None) = (it.next(), it.next()) else {
                return err(format!("{op:?} needs exactly 1 arg"));
            };
            return Ok(if op == "is_null" {
                a.is_null()
            } else {
                a.is_not_null()
            });
        }
        return err(format!("operator {op:?} not allowed"));
    }

    if let Some(f) = obj.get("fn") {
        let f = f.as_str().unwrap_or("");
        let (lo, hi) = match func_arity(f) {
            Some(a) => a,
            None => return err(format!("function {f:?} not allowed")),
        };
        let raw = args_of(obj);
        if raw.len() < lo || hi.map(|h| raw.len() > h).unwrap_or(false) {
            return err(format!("function {f:?} arity {} out of range", raw.len()));
        }
        // round's optional 2nd arg is a decimal-places count, which Polars needs
        // as a literal (not an expr) — parse it from the AST rather than silently
        // ignoring it.
        if f == "round" {
            let decimals: u32 = match raw.get(1) {
                None => 0,
                Some(p) => p.get("lit").and_then(|v| v.as_u64()).ok_or_else(|| {
                    FormulaError("round's 2nd arg must be a literal integer".into())
                })? as u32,
            };
            let first = raw
                .first()
                .ok_or_else(|| FormulaError("round needs ≥1 arg".into()))?;
            let inner = compile_formula(first, columns)?;
            return Ok(inner.round(decimals, RoundMode::HalfToEven));
        }
        let args = raw
            .iter()
            .map(|n| compile_formula(n, columns))
            .collect::<Result<Vec<Expr>, FormulaError>>()?;
        return apply_fn(f, args);
    }

    if let Some(inner) = obj.get("cast") {
        let t = obj
            .get("as")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_lowercase();
        let dt = cast_dtype(&t)?;
        let e = compile_formula(inner, columns)?;
        return Ok(e.cast(dt));
    }

    let mut keys: Vec<&String> = obj.keys().collect();
    keys.sort();
    err(format!("unrecognized AST node: {keys:?}"))
}

/// Every `{"col": "..."}` referenced anywhere in a formula AST. Used by the
/// `dims_from` loader to know which snapshot columns an axis `row_filter` needs
/// joined onto the narrow series frame (the filter is applied AFTER the join, so
/// its inputs must be present). Lenient: a malformed node simply contributes no
/// columns — validation happens later in `compile_formula`.
pub fn referenced_columns(node: &Json) -> HashSet<String> {
    let mut out = HashSet::new();
    collect_columns(node, &mut out);
    out
}

fn collect_columns(node: &Json, out: &mut HashSet<String>) {
    match node {
        Json::Object(obj) => {
            if let Some(Json::String(name)) = obj.get("col") {
                out.insert(name.clone());
            }
            for (_, v) in obj {
                collect_columns(v, out);
            }
        }
        Json::Array(items) => {
            for v in items {
                collect_columns(v, out);
            }
        }
        _ => {}
    }
}

fn compile_args(
    obj: &serde_json::Map<String, Json>,
    columns: &HashSet<String>,
) -> Result<Vec<Expr>, FormulaError> {
    obj.get("args")
        .and_then(|a| a.as_array())
        .map(|a| a.iter().map(|n| compile_formula(n, columns)).collect())
        .unwrap_or_else(|| Ok(vec![]))
}

/// Raw (uncompiled) AST args — `round` reads its literal places arg from these.
fn args_of(obj: &serde_json::Map<String, Json>) -> Vec<Json> {
    obj.get("args")
        .and_then(|a| a.as_array())
        .cloned()
        .unwrap_or_default()
}

fn apply_fn(name: &str, mut args: Vec<Expr>) -> Result<Expr, FormulaError> {
    // Pull the sole argument of a unary function (arity was checked by the caller;
    // the error path keeps this total rather than relying on that invariant).
    fn one(a: Vec<Expr>) -> Result<Expr, FormulaError> {
        a.into_iter()
            .next()
            .ok_or_else(|| FormulaError("function needs exactly 1 arg".into()))
    }
    Ok(match name {
        "coalesce" => coalesce(&args),
        "nullif" => {
            let two = || FormulaError("`nullif` needs exactly 2 args".into());
            let b = args.pop().ok_or_else(two)?;
            let a = args.pop().ok_or_else(two)?;
            when(a.clone().eq(b)).then(lit(NULL)).otherwise(a)
        }
        "abs" => one(args)?.abs(),
        // "round" is handled in compile_formula (it needs its literal 2nd arg).
        "least" => fold_minmax(args, true)?,
        "greatest" => fold_minmax(args, false)?,
        "ln" => one(args)?.log(lit(std::f64::consts::E)),
        "log" => one(args)?.log(lit(10.0)),
        "exp" => one(args)?.exp(),
        "floor" => one(args)?.floor(),
        "ceil" => one(args)?.ceil(),
        "sqrt" => one(args)?.sqrt(),
        other => return err(format!("function {other:?} not allowed")),
    })
}

/// least/greatest as a pairwise `when` fold (avoids polars' ambiguous
/// `min_horizontal`/`max_horizontal` glob re-exports).
fn fold_minmax(args: Vec<Expr>, least: bool) -> Result<Expr, FormulaError> {
    let mut it = args.into_iter();
    let mut acc = it
        .next()
        .ok_or_else(|| FormulaError("`least`/`greatest` needs ≥1 arg".into()))?;
    for a in it {
        let pick_a = if least {
            a.clone().lt(acc.clone())
        } else {
            a.clone().gt(acc.clone())
        };
        acc = when(pick_a).then(a).otherwise(acc);
    }
    Ok(acc)
}

fn cast_dtype(t: &str) -> Result<DataType, FormulaError> {
    Ok(match t {
        "double" | "decimal" => DataType::Float64,
        "bigint" => DataType::Int64,
        "integer" => DataType::Int32,
        "varchar" => DataType::String,
        "boolean" => DataType::Boolean,
        other => return err(format!("cast type {other:?} not allowed")),
    })
}

fn lit_expr(v: &Json) -> Result<Expr, FormulaError> {
    Ok(match v {
        Json::Null => lit(NULL),
        Json::Bool(b) => lit(*b),
        Json::Number(n) => {
            if let Some(i) = n.as_i64() {
                lit(i)
            } else {
                lit(n.as_f64().unwrap_or(0.0))
            }
        }
        Json::String(s) => lit(s.clone()),
        other => return err(format!("unsupported literal {other}")),
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    fn cols() -> HashSet<String> {
        ["stars", "forks"].iter().map(|s| s.to_string()).collect()
    }

    #[test]
    fn validates_columns_and_funcs() {
        let c = cols();
        assert!(compile_formula(&json!({"col": "stars"}), &c).is_ok());
        assert!(compile_formula(
            &json!({"op": "/", "args": [{"col": "forks"}, {"col": "stars"}]}),
            &c
        )
        .is_ok());
        assert!(compile_formula(
            &json!({"fn": "coalesce", "args": [{"col": "stars"}, {"lit": 0}]}),
            &c
        )
        .is_ok());
        assert!(compile_formula(&json!({"cast": {"col": "stars"}, "as": "double"}), &c).is_ok());
        // rejects
        assert!(compile_formula(&json!({"col": "evil; DROP TABLE"}), &c).is_err());
        assert!(compile_formula(&json!({"fn": "system", "args": []}), &c).is_err());
        assert!(compile_formula(
            &json!({"op": "%", "args": [{"col": "stars"}, {"lit": 2}]}),
            &c
        )
        .is_err());
        assert!(compile_formula(&json!({"cast": {"col": "stars"}, "as": "evil"}), &c).is_err());
    }

    #[test]
    fn malformed_arity_errors_never_panics() {
        // Every arity boundary that used to `.unwrap()` an iterator must now return
        // an Err for the off-by-one input — exercising the panic-free paths.
        let c = cols();
        let bad = [
            json!({"op": "+", "args": [{"col": "stars"}]}), // arithmetic needs ≥2
            json!({"op": ">=", "args": [{"col": "stars"}]}), // cmp needs exactly 2
            json!({"op": ">=", "args": [{"col": "stars"}, {"lit": 1}, {"lit": 2}]}),
            json!({"op": "and", "args": [{"col": "stars"}]}), // and/or need ≥2
            json!({"op": "not", "args": [{"col": "stars"}, {"col": "forks"}]}), // not is unary
            json!({"op": "is_null", "args": [{"col": "stars"}, {"col": "forks"}]}),
            json!({"fn": "abs", "args": []}),   // unary fn needs 1
            json!({"fn": "least", "args": []}), // least/greatest need ≥1 (arity ≥2 here)
            json!({"fn": "round", "args": []}), // round needs ≥1
            json!({"fn": "nullif", "args": [{"col": "stars"}]}),
        ];
        for node in bad {
            assert!(
                compile_formula(&node, &c).is_err(),
                "expected Err (not a panic) for {node}"
            );
        }
    }

    #[test]
    fn comparison_and_boolean_predicates_compile() {
        let c = cols();
        // a single comparison node → a boolean predicate Expr
        assert!(compile_formula(
            &json!({"op": ">=", "args": [{"col": "stars"}, {"lit": 1e7}]}),
            &c
        )
        .is_ok());
        assert!(compile_formula(
            &json!({"op": "<", "args": [{"col": "forks"}, {"lit": 5}]}),
            &c
        )
        .is_ok());
        assert!(compile_formula(
            &json!({"op": "==", "args": [{"col": "stars"}, {"col": "forks"}]}),
            &c
        )
        .is_ok());
        // boolean combinators
        assert!(compile_formula(
            &json!({"op": "and", "args": [
                {"op": ">=", "args": [{"col": "stars"}, {"lit": 100}]},
                {"op": "<",  "args": [{"col": "forks"}, {"lit": 1000}]}
            ]}),
            &c
        )
        .is_ok());
        assert!(compile_formula(
            &json!({"op": "not", "args": [{"op": ">", "args": [{"col": "stars"}, {"lit": 0}]}]}),
            &c
        )
        .is_ok());
        // arity errors
        assert!(compile_formula(&json!({"op": ">=", "args": [{"col": "stars"}]}), &c).is_err());
        assert!(compile_formula(
            &json!({"op": "not", "args": [{"col": "stars"}, {"col": "forks"}]}),
            &c
        )
        .is_err());
        // still rejects unknown ops
        assert!(compile_formula(
            &json!({"op": "%", "args": [{"col": "stars"}, {"lit": 2}]}),
            &c
        )
        .is_err());
    }

    #[test]
    fn referenced_columns_walks_the_ast() {
        // A row-filter formula referencing two snapshot columns nested in
        // comparison + boolean ops. `referenced_columns` must surface BOTH so the
        // `dims_from` loader joins them onto the narrow series frame.
        let ast = json!({"op": "and", "args": [
            {"op": ">=", "args": [{"col": "mcap_usd"}, {"lit": 1e7}]},
            {"op": "<",  "args": [{"col": "founding_year"}, {"lit": 2000}]}
        ]});
        let got = referenced_columns(&ast);
        let mut got: Vec<String> = got.into_iter().collect();
        got.sort();
        assert_eq!(
            got,
            vec!["founding_year".to_string(), "mcap_usd".to_string()]
        );

        // A leaf col node, and a node with no cols, are handled.
        assert_eq!(
            referenced_columns(&json!({"col": "x"})),
            ["x".to_string()].into_iter().collect()
        );
        assert!(referenced_columns(&json!({"lit": 1})).is_empty());
    }

    #[test]
    fn null_predicates_compile_and_filter() {
        let c = cols();
        // Both unary null ops compile to a boolean predicate Expr.
        assert!(compile_formula(&json!({"op": "is_null", "args": [{"col": "stars"}]}), &c).is_ok());
        assert!(compile_formula(
            &json!({"op": "is_not_null", "args": [{"col": "stars"}]}),
            &c
        )
        .is_ok());
        // arity errors
        assert!(compile_formula(
            &json!({"op": "is_null", "args": [{"col": "stars"}, {"col": "forks"}]}),
            &c
        )
        .is_err());
        assert!(compile_formula(&json!({"op": "is_not_null", "args": []}), &c).is_err());

        // `is_not_null` actually filters null rows out of a frame.
        let df = df!["stars" => &[Some(1i64), None, Some(3)]].unwrap();
        let pred = compile_formula(
            &json!({"op": "is_not_null", "args": [{"col": "stars"}]}),
            &c,
        )
        .unwrap();
        let kept = df
            .lazy()
            .filter(pred)
            .collect()
            .unwrap()
            .column("stars")
            .unwrap()
            .i64()
            .unwrap()
            .into_no_null_iter()
            .collect::<Vec<i64>>();
        assert_eq!(kept, vec![1, 3]);
    }

    #[test]
    fn round_arg_handling() {
        let c = cols();
        assert!(compile_formula(&json!({"fn": "round", "args": [{"col": "stars"}]}), &c).is_ok());
        // literal places arg is honored
        assert!(compile_formula(
            &json!({"fn": "round", "args": [{"col": "stars"}, {"lit": 2}]}),
            &c
        )
        .is_ok());
        // non-literal places arg is rejected (rather than silently ignored)
        assert!(compile_formula(
            &json!({"fn": "round", "args": [{"col": "stars"}, {"col": "forks"}]}),
            &c
        )
        .is_err());
    }
}