fuel-dispatch 0.20.0

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Parse the FKC authoring DSL (`extent(role, axis)`, `const(N)`, `param(N)`,
//! `add|sub|mul|div(a, b)`) into the positional `shape_expr::Dim` AST. Role names are
//! resolved to positional operand indices via the probe combo's canonical order
//! (§6.4-0009 wire form is positional). Returns `None` on any malformed / unknown-role
//! input (the caller maps `None` → skip; never a false reject).

use crate::fkc::return_check::ProbeComboRef;
use crate::fkc::shape_expr::{Dim, LAST};

/// Position of `role` in the combo's canonical order.
fn role_pos(combo: ProbeComboRef, role: &str) -> Option<u8> {
    combo
        .iter()
        .position(|(r, _, _)| r == role)
        .and_then(|p| u8::try_from(p).ok())
}

/// Split `"a, b"` (the two args of a binary node) at the top-level comma (depth 0).
fn split_top_comma(s: &str) -> Option<(&str, &str)> {
    let mut depth = 0i32;
    for (i, ch) in s.char_indices() {
        match ch {
            '(' => depth += 1,
            ')' => depth -= 1,
            ',' if depth == 0 => return Some((s[..i].trim(), s[i + 1..].trim())),
            _ => {}
        }
    }
    None
}

/// Strip a `head( ... )` wrapper, returning the inner text.
fn inner<'a>(s: &'a str, head: &str) -> Option<&'a str> {
    s.trim()
        .strip_prefix(head)?
        .strip_suffix(')')
        .map(str::trim)
}

/// Parse a whole-shape `matmul(role_a, role_b)` rule into the two operand
/// shapes (concrete `i64` extents), resolved against the probe combo in
/// `(lhs, rhs)` order. `None` on a non-`matmul(...)` rule, a malformed arg
/// list, or an unknown role — the caller maps `None` → skip (never a false
/// reject). §6.20-0008 role-woven `matmul` kind: the two named roles ARE the
/// M/N/K-bearing operands; the caller derives the contraction shape via
/// [`crate::fkc::shape_expr::matmul_shape`] (from operand shapes + the M/N/K
/// role structure — NEVER from `entry.decompose`, per the §4 guardrail).
pub fn parse_matmul_operands(rule: &str, combo: ProbeComboRef) -> Option<(Vec<i64>, Vec<i64>)> {
    let args = inner(rule.trim(), "matmul(")?;
    let (role_a, role_b) = split_top_comma(args)?;
    let shape_of = |role: &str| {
        combo
            .iter()
            .find(|(r, _, _)| r == role)
            .map(|(_, s, _)| s.dims().iter().map(|&d| d as i64).collect::<Vec<i64>>())
    };
    Some((shape_of(role_a)?, shape_of(role_b)?))
}

pub fn parse_dim(rule: &str, combo: ProbeComboRef) -> Option<Dim> {
    let rule = rule.trim();
    if let Some(args) = inner(rule, "extent(") {
        let (role, axis) = split_top_comma(args)?;
        let operand = role_pos(combo, role)?;
        let axis = if axis == "last" {
            LAST
        } else {
            axis.parse::<u8>().ok()?
        };
        return Some(Dim::Extent { operand, axis });
    }
    if let Some(n) = inner(rule, "const(") {
        return Some(Dim::Const(n.parse::<i64>().ok()?));
    }
    if let Some(f) = inner(rule, "param(") {
        return Some(Dim::Param(f.parse::<u8>().ok()?));
    }
    for (head, ctor) in [("add(", 0u8), ("sub(", 1), ("mul(", 2), ("div(", 3)] {
        if let Some(args) = inner(rule, head) {
            let (a, b) = split_top_comma(args)?;
            let (da, db) = (
                Box::new(parse_dim(a, combo)?),
                Box::new(parse_dim(b, combo)?),
            );
            return Some(match ctor {
                0 => Dim::Add(da, db),
                1 => Dim::Sub(da, db),
                2 => Dim::Mul(da, db),
                _ => Dim::Div(da, db),
            });
        }
    }
    None
}