earthsci-ast 0.6.0

//! Load-time `enum`-op lowering pass — esm-spec §4.5, §9.3.
//!
//! Walks every expression tree in the file and replaces each
//! `{op: "enum", args: [enum_name, symbol]}` node with a
//! `{op: "const", args: [], value: <integer>}` node, using the file's
//! top-level `enums` block. After this pass, no `enum`-op nodes remain in
//! the in-memory representation — downstream consumers (evaluators,
//! canonicalize, codegen) only ever see resolved integer constants.
//!
//! Mirrors the Python (`registered_functions.lower_enums`) and Julia
//! (`lower_enums!`) passes for cross-binding parity, with the same
//! diagnostic codes (`unknown_enum`, `unknown_enum_symbol`).
//!
//! Operates on the pre-deserialization `serde_json::Value` view so that
//! it covers both top-level components and JSON-typed subsystem blobs in
//! a single walk.

use serde_json::{Map, Value};
use std::collections::HashMap;

const ENUM_OP: &str = "enum";

/// Stable diagnostic codes raised by the enum-lowering pass. Mirrors the
/// codes emitted by the TS / Python / Julia / Go bindings.
pub type EnumLoweringError = crate::diagnostic::DiagnosticError;

use crate::diagnostic::{codes, err};

fn parse_enums_block(
    value: &Value,
) -> Result<HashMap<String, HashMap<String, i64>>, EnumLoweringError> {
    let Some(root) = value.as_object() else {
        return Ok(HashMap::new());
    };
    let Some(enums_value) = root.get("enums") else {
        return Ok(HashMap::new());
    };
    let Some(enums_obj) = enums_value.as_object() else {
        return Err(err(
            codes::INVALID_ENUMS_BLOCK,
            "top-level `enums` must be an object",
        ));
    };
    let mut out = HashMap::new();
    for (enum_name, mapping) in enums_obj {
        let Some(map_obj) = mapping.as_object() else {
            return Err(err(
                codes::INVALID_ENUMS_BLOCK,
                format!("enums.{enum_name}: entry must be an object"),
            ));
        };
        let mut m = HashMap::new();
        for (sym, intval) in map_obj {
            let Some(n) = intval.as_i64() else {
                return Err(err(
                    codes::INVALID_ENUMS_BLOCK,
                    format!("enums.{enum_name}.{sym}: value must be an integer"),
                ));
            };
            m.insert(sym.clone(), n);
        }
        out.insert(enum_name.clone(), m);
    }
    Ok(out)
}

fn lower_enum_node(
    node: &Map<String, Value>,
    enums: &HashMap<String, HashMap<String, i64>>,
) -> Result<Value, EnumLoweringError> {
    let args = node.get("args").and_then(|v| v.as_array()).ok_or_else(|| {
        err(
            codes::ENUM_INVALID_ARGS,
            "`enum` op requires an `args` array of length 2",
        )
    })?;
    if args.len() != 2 {
        return Err(err(
            codes::ENUM_INVALID_ARGS,
            format!(
                "`enum` op expects 2 args (enum_name, symbol_name), got {}",
                args.len()
            ),
        ));
    }
    let enum_name = args[0].as_str().ok_or_else(|| {
        err(
            codes::ENUM_INVALID_ARGS,
            "`enum` op: first arg must be a string (enum name)",
        )
    })?;
    let symbol_name = args[1].as_str().ok_or_else(|| {
        err(
            codes::ENUM_INVALID_ARGS,
            "`enum` op: second arg must be a string (symbol name)",
        )
    })?;
    let mapping = enums.get(enum_name).ok_or_else(|| {
        err(
            codes::UNKNOWN_ENUM,
            format!(
                "unknown_enum: enum `{enum_name}` is not declared in the file's \
                 `enums` block"
            ),
        )
    })?;
    let int_value = mapping.get(symbol_name).ok_or_else(|| {
        err(
            codes::UNKNOWN_ENUM_SYMBOL,
            format!(
                "unknown_enum_symbol: symbol `{symbol_name}` is not declared \
                 under enum `{enum_name}`"
            ),
        )
    })?;

    let mut out = Map::new();
    out.insert("op".to_string(), Value::String("const".to_string()));
    out.insert("args".to_string(), Value::Array(Vec::new()));
    out.insert("value".to_string(), Value::Number((*int_value).into()));
    Ok(Value::Object(out))
}

/// Return a copy of `node` with every `enum`-op object replaced by its
/// lowered `const` node. A replacement's fields are all leaves, so the
/// visitor's descent into it finds nothing further to lower; on error the
/// partially transformed copy is discarded and the caller's value is left
/// untouched.
fn walk(
    node: &Value,
    enums: &HashMap<String, HashMap<String, i64>>,
) -> Result<Value, EnumLoweringError> {
    let mut out = node.clone();
    crate::json_visit::try_visit_values_mut(&mut out, &mut |v| {
        if let Value::Object(obj) = v
            && obj.get("op").and_then(|w| w.as_str()) == Some(ENUM_OP)
        {
            *v = lower_enum_node(obj, enums)?;
        }
        Ok(())
    })?;
    Ok(out)
}

/// Lower the `enum`-op nodes under `target` against the `enums` block of
/// `document`, the file that wrote `target`. Mutates `target` in place.
///
/// An `enum` op is file-local (esm-spec §9.3): it resolves against the block
/// of the file it is written in. [`lower_enums_raw`] runs once over the root
/// document, so a tree that crosses a file boundary before that pass (a
/// template-library body reaching an importer, §9.7.5) is lowered here, at the
/// edge, while its own file's block is still at hand.
///
/// An op with an argument spelled by a name in `open_names` is left in place:
/// a template parameter substitutes position-blind (§9.6.3 constraint 5), so
/// the call site decides what it spells and the op resolves there. An op whose
/// arguments are not two strings is left for [`lower_enums_raw`], which owns
/// the malformed-op diagnostic.
pub(crate) fn lower_enum_ops_for_file(
    document: &Value,
    target: &mut Value,
    open_names: &std::collections::HashSet<String>,
) -> Result<(), EnumLoweringError> {
    let mut paths: Vec<String> = Vec::new();
    find_enum_paths(target, &mut paths);
    if paths.is_empty() {
        return Ok(());
    }
    let enums = parse_enums_block(document)?;
    crate::json_visit::try_visit_values_mut(target, &mut |v| {
        if let Value::Object(obj) = v
            && obj.get("op").and_then(|w| w.as_str()) == Some(ENUM_OP)
            && let Some([Value::String(name), Value::String(symbol)]) = obj
                .get("args")
                .and_then(|a| a.as_array())
                .map(Vec::as_slice)
            && !open_names.contains(name)
            && !open_names.contains(symbol)
        {
            *v = lower_enum_node(obj, &enums)?;
        }
        Ok(())
    })
}

/// Lower, in place, the `enum` ops of `document`, a document mounted at a
/// §4.7 edge, against its own `enums` block (esm-spec §9.3).
///
/// Called once the document has resolved in its own scope and before its
/// component is spliced into the mounting document, whose block is a different
/// one. `enums` do not merge across a mount, so this is the only block those
/// ops can name.
///
/// A template declaration's body is lowered with that template's `params` left
/// open, as at the import edge ([`lower_enum_ops_for_file`]): an op spelled
/// with a parameter resolves at the call site. Everything else is lowered with
/// no open names.
pub(crate) fn lower_mounted_document_enums(document: &mut Value) -> Result<(), EnumLoweringError> {
    let mut source = Map::new();
    if let Some(enums) = document.get("enums") {
        source.insert("enums".to_string(), enums.clone());
    }
    lower_mounted_enum_ops(&Value::Object(source), document)
}

fn lower_mounted_enum_ops(source: &Value, node: &mut Value) -> Result<(), EnumLoweringError> {
    let no_open_names = std::collections::HashSet::new();
    if node.get("op").and_then(|w| w.as_str()) == Some(ENUM_OP) {
        return lower_enum_ops_for_file(source, node, &no_open_names);
    }
    match node {
        Value::Array(items) => {
            for item in items {
                lower_mounted_enum_ops(source, item)?;
            }
        }
        Value::Object(obj) => {
            for (key, child) in obj.iter_mut() {
                if key == "expression_templates"
                    && let Value::Object(templates) = child
                {
                    for decl in templates.values_mut() {
                        let params: std::collections::HashSet<String> = decl
                            .get("params")
                            .and_then(|p| p.as_array())
                            .map(|ps| {
                                ps.iter()
                                    .filter_map(|p| p.as_str().map(str::to_string))
                                    .collect()
                            })
                            .unwrap_or_default();
                        if let Some(body) = decl.get_mut("body") {
                            lower_enum_ops_for_file(source, body, &params)?;
                        }
                    }
                } else {
                    lower_mounted_enum_ops(source, child)?;
                }
            }
        }
        _ => {}
    }
    Ok(())
}

/// Record the path of every `enum`-op node under `view`, in pre-order. The
/// walk also descends into a matched node's own fields (they hold only the
/// two name strings in a well-formed node, so nothing further matches there).
fn find_enum_paths(view: &Value, hits: &mut Vec<String>) {
    crate::json_visit::visit_values(view, &mut |path, v| {
        if let Value::Object(obj) = v
            && crate::json_visit::small_get(obj, "op").and_then(|w| w.as_str()) == Some(ENUM_OP)
        {
            hits.push(path.to_string());
        }
    });
}

/// Lower every `enum`-op node in `file` to a `{op: "const", value: <int>}`
/// node using the document's top-level `enums` block, returning a NEW
/// document. `file` is left untouched.
///
/// The canonical entry point (API_SPEC.md §8 items 14 and 15): TYPED, because
/// every other binding takes a typed document, and PURE, because TypeScript
/// and Python are pure and the mutating variants take a `_mut` suffix
/// (Julia's `!`). The in-place raw-JSON pass the loader itself runs is
/// [`lower_enums_raw`].
///
/// Pre-condition: the input has been schema-validated and any
/// `apply_expression_template` ops have already been expanded.
///
/// Errors with an [`EnumLoweringError`] carrying one of the diagnostic codes
/// listed on [`lower_enums_raw`].
pub fn lower_enums(file: &crate::EsmFile) -> Result<crate::EsmFile, EnumLoweringError> {
    let mut value = serde_json::to_value(file).map_err(|e| {
        err(
            codes::INVALID_ENUMS_BLOCK,
            format!("could not render the document as JSON for enum lowering: {e}"),
        )
    })?;
    lower_enums_raw(&mut value)?;
    serde_json::from_value(value).map_err(|e| {
        err(
            codes::ENUM_LOWERING_RESIDUAL,
            format!("lowered document did not deserialize back into an EsmFile: {e}"),
        )
    })
}

/// [`lower_enums`], applied in place. The `_mut` spelling API_SPEC.md §8 item
/// 15 assigns to a mutating variant (Julia spells the same thing
/// `lower_enums!`).
pub fn lower_enums_mut(file: &mut crate::EsmFile) -> Result<(), EnumLoweringError> {
    *file = lower_enums(file)?;
    Ok(())
}

/// Lower every `enum`-op node in `value` to a `{op: "const", value: <int>}`
/// node using the file's top-level `enums` block. Mutates `value` in place.
///
/// The raw-`serde_json::Value` extension seam (API_SPEC.md §8 item 14): this
/// is what the loader runs, BEFORE the document is coerced to an
/// [`crate::EsmFile`], so that JSON-typed subsystem blobs are covered by the
/// same walk. Callers holding a typed document want [`lower_enums`].
///
/// Pre-condition: the input has been schema-validated and any
/// `apply_expression_template` ops have already been expanded.
///
/// Errors with diagnostic codes:
/// - `unknown_enum` — `enum` op references an undeclared enum name
/// - `unknown_enum_symbol` — `enum` op references an undeclared symbol
/// - `enum_invalid_args` — `enum` op has a non-positional or wrong-arity body
/// - `invalid_enums_block` — top-level `enums` block is malformed
pub fn lower_enums_raw(value: &mut Value) -> Result<(), EnumLoweringError> {
    lower_enums_raw_scanned(value, true)
}

/// Whether `value` is an `enum`-op node, as [`lower_enums_raw`] finds them.
pub(crate) fn is_enum_node(value: &Value) -> bool {
    matches!(value, Value::Object(obj)
        if crate::json_visit::small_get(obj, "op").and_then(|w| w.as_str()) == Some(ENUM_OP))
}

/// [`lower_enums_raw`] for a caller whose own walk of `value` already found
/// whether it has any `enum`-op node ([`is_enum_node`]): `false` skips the
/// search, leaving only the `enums` block check.
pub(crate) fn lower_enums_raw_scanned(
    value: &mut Value,
    may_have_enum_ops: bool,
) -> Result<(), EnumLoweringError> {
    let enums = parse_enums_block(value)?;
    if !may_have_enum_ops {
        return Ok(());
    }

    // Fast path: no enum-op nodes anywhere; skip the rebuild.
    let mut paths: Vec<String> = Vec::new();
    find_enum_paths(value, &mut paths);
    if paths.is_empty() {
        return Ok(());
    }

    let walked = walk(value, &enums)?;
    *value = walked;

    // Defensive: assert no enum nodes remain.
    let mut leftover: Vec<String> = Vec::new();
    find_enum_paths(value, &mut leftover);
    if !leftover.is_empty() {
        return Err(err(
            codes::ENUM_LOWERING_RESIDUAL,
            format!(
                "enum-op nodes remain after lowering at: {}",
                leftover.join(", ")
            ),
        ));
    }
    Ok(())
}

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

    fn lowered(input: Value) -> Value {
        let mut v = input;
        lower_enums_raw(&mut v).expect("lowering should succeed");
        v
    }

    #[test]
    fn basic_enum_replaced_with_const() {
        let input = json!({
            "enums": {
                "season": {"winter": 1, "spring": 2, "summer": 3, "autumn": 4}
            },
            "models": {
                "M": {
                    "variables": {
                        "s": {
                            "type": "unknown"}
                    },
                    "equations": [
                {"lhs": "s", "rhs": {
                                "op": "enum",
                                "args": ["season", "summer"]
                            }}]
                }
            }
        });
        let out = lowered(input);
        let expr = crate::classification::observed_definition_json(&out["models"]["M"], "s")
            .expect("s defining equation");
        assert_eq!(expr["op"], "const");
        assert_eq!(expr["value"], 3);
        assert_eq!(expr["args"], json!([]));
    }

    #[test]
    fn nested_enum_in_index_op() {
        let input = json!({
            "enums": {
                "season": {"summer": 3},
                "land_use": {"forest": 5}
            },
            "models": {
                "M": {
                    "variables": {
                        "r": {
                            "type": "unknown"}
                    },
                    "equations": [
                {"lhs": "r", "rhs": {
                                "op": "index",
                                "args": [
                                    {"op": "const", "args": [], "value": [[1, 2, 3]]},
                                    {"op": "enum", "args": ["season", "summer"]},
                                    {"op": "enum", "args": ["land_use", "forest"]}
                                ]
                            }}]
                }
            }
        });
        let out = lowered(input);
        let expr = crate::classification::observed_definition_json(&out["models"]["M"], "r")
            .expect("r defining equation");
        assert_eq!(expr["op"], "index");
        assert_eq!(expr["args"][1]["op"], "const");
        assert_eq!(expr["args"][1]["value"], 3);
        assert_eq!(expr["args"][2]["op"], "const");
        assert_eq!(expr["args"][2]["value"], 5);
    }

    #[test]
    fn unknown_enum_rejected() {
        let mut v = json!({
            "enums": {"season": {"summer": 3}},
            "models": {"M": {
                "variables": {"x": {
                    "type": "unknown"}},
                "equations": [
                {"lhs": "x", "rhs": {"op": "enum", "args": ["weekday", "monday"]}}]
            }}
        });
        let err = lower_enums_raw(&mut v).unwrap_err();
        assert_eq!(err.code, "unknown_enum");
        assert!(err.message.contains("weekday"));
    }

    #[test]
    fn unknown_enum_symbol_rejected() {
        let mut v = json!({
            "enums": {"season": {"summer": 3}},
            "models": {"M": {
                "variables": {"x": {
                    "type": "unknown"}},
                "equations": [
                {"lhs": "x", "rhs": {"op": "enum", "args": ["season", "winter"]}}]
            }}
        });
        let err = lower_enums_raw(&mut v).unwrap_err();
        assert_eq!(err.code, "unknown_enum_symbol");
        assert!(err.message.contains("winter"));
        assert!(err.message.contains("season"));
    }

    #[test]
    fn missing_enums_block_rejects_enum_op() {
        let mut v = json!({
            "models": {"M": {
                "variables": {"x": {
                    "type": "unknown"}},
                "equations": [
                {"lhs": "x", "rhs": {"op": "enum", "args": ["season", "summer"]}}]
            }}
        });
        let err = lower_enums_raw(&mut v).unwrap_err();
        assert_eq!(err.code, "unknown_enum");
    }

    #[test]
    fn wrong_arity_rejected() {
        let mut v = json!({
            "enums": {"season": {"summer": 3}},
            "models": {"M": {
                "variables": {"x": {
                    "type": "unknown"}},
                "equations": [
                {"lhs": "x", "rhs": {"op": "enum", "args": ["season"]}}]
            }}
        });
        let err = lower_enums_raw(&mut v).unwrap_err();
        assert_eq!(err.code, "enum_invalid_args");
    }

    #[test]
    fn no_enums_no_change() {
        let input = json!({
            "models": {"M": {
                "variables": {"x": {"type": "unknown", "default": 0.0}},
                "equations": [{
                    "lhs": {"op": "D", "args": ["x"], "wrt": "t"},
                    "rhs": {"op": "*", "args": [0.1, "x"]}
                }]
            }}
        });
        let out = lowered(input.clone());
        assert_eq!(out, input);
    }

    #[test]
    fn enum_op_in_reaction_rate() {
        let input = json!({
            "enums": {"phase": {"day": 1, "night": 2}},
            "reaction_systems": {
                "RS": {
                    "species": [],
                    "reactions": [{
                        "id": "r1",
                        "substrates": {},
                        "products": {},
                        "rate": {
                            "op": "*",
                            "args": [1.0, {"op": "enum", "args": ["phase", "day"]}]
                        }
                    }]
                }
            }
        });
        let out = lowered(input);
        let rate = &out["reaction_systems"]["RS"]["reactions"][0]["rate"];
        assert_eq!(rate["args"][1]["op"], "const");
        assert_eq!(rate["args"][1]["value"], 1);
    }

    #[test]
    fn enum_op_in_subsystem_blob() {
        // Subsystems land as raw JSON in the Rust types, so the JSON walk
        // must descend through them too.
        let input = json!({
            "enums": {"season": {"summer": 3}},
            "models": {"Outer": {
                "variables": {},
                "equations": [],
                "subsystems": {
                    "Inner": {
                        "variables": {"y": {
                            "type": "unknown"}},
                        "equations": [
                {"lhs": "y", "rhs": {"op": "enum", "args": ["season", "summer"]}}]
                    }
                }
            }}
        });
        let out = lowered(input);
        let expr = crate::classification::observed_definition_json(
            &out["models"]["Outer"]["subsystems"]["Inner"],
            "y",
        )
        .expect("y defining equation");
        assert_eq!(expr["op"], "const");
        assert_eq!(expr["value"], 3);
    }

    /// API_SPEC.md §8 items 14 + 15: the canonical `lower_enums` is TYPED and
    /// PURE — it takes an `EsmFile`, returns a new one, and leaves the input
    /// alone. `lower_enums_mut` is the in-place variant.
    #[test]
    fn typed_lower_enums_is_pure_and_agrees_with_the_raw_pass() {
        let doc = json!({
            "esm": "1.0.0",
            "metadata": {"name": "enum_doc"},
            "enums": {"season": {"winter": 1, "summer": 3}},
            "models": {
                "M": {
                    "variables": {"s": {"type": "unknown"}},
                    "equations": [
                        {"lhs": "s", "rhs": {"op": "enum", "args": ["season", "summer"]}}
                    ]
                }
            }
        });
        let file: crate::EsmFile = serde_json::from_value(doc.clone()).expect("typed document");

        let lowered_file = lower_enums(&file).expect("typed lowering succeeds");

        // Pure: the input still carries the enum op.
        let before = serde_json::to_value(&file).expect("re-render");
        assert_eq!(before["models"]["M"]["equations"][0]["rhs"]["op"], "enum");

        let after = serde_json::to_value(&lowered_file).expect("re-render");
        assert_eq!(after["models"]["M"]["equations"][0]["rhs"]["op"], "const");
        assert_eq!(after["models"]["M"]["equations"][0]["rhs"]["value"], 3);

        // Agrees with the raw seam.
        assert_eq!(after, lowered(doc));

        // And `_mut` is the same transform applied in place.
        let mut mutated = file.clone();
        lower_enums_mut(&mut mutated).expect("in-place lowering succeeds");
        assert_eq!(serde_json::to_value(&mutated).expect("re-render"), after);
    }

    /// The typed form raises `EnumLoweringError` with the cross-binding code.
    #[test]
    fn typed_lower_enums_raises_enum_lowering_error() {
        let file: crate::EsmFile = serde_json::from_value(json!({
            "esm": "1.0.0",
            "metadata": {"name": "bad_enum"},
            "enums": {"season": {"winter": 1}},
            "models": {
                "M": {
                    "variables": {"s": {"type": "unknown"}},
                    "equations": [
                        {"lhs": "s", "rhs": {"op": "enum", "args": ["season", "nope"]}}
                    ]
                }
            }
        }))
        .expect("typed document");
        let e: EnumLoweringError = lower_enums(&file).unwrap_err();
        assert_eq!(e.code, codes::UNKNOWN_ENUM_SYMBOL);
    }
}