uiua_parser 0.19.0-dev.4

//! Splitting primitive names

use std::{collections::HashMap, fmt, sync::LazyLock};

use enum_iterator::Sequence;

use crate::{Complex, SysOp, ast::NumWord};

use super::Primitive;

static ALIASES: LazyLock<HashMap<Primitive, &[&str]>> = LazyLock::new(|| {
    [
        (Primitive::Identity, &["id"] as &[_]),
        (Primitive::Gap, &["ga"]),
        (Primitive::Pop, &["po"]),
        (Primitive::Fix, &["fx"]),
        (Primitive::Box, &["bx"]),
        (Primitive::MemberOf, &["elem"]),
        (Primitive::IndexIn, &["idx"]),
        (Primitive::Switch, &["sw"]),
        (Primitive::Stencil, &["st"]),
        (Primitive::Floor, &["flr", "flor"]),
        (Primitive::Range, &["ran"]),
        (Primitive::Partition, &["par"]),
        (Primitive::Deshape, &["flat", "fla"]),
        (Primitive::Ne, &["ne", "neq"]),
        (Primitive::Eq, &["eq"]),
        (Primitive::Lt, &["lt", "less"]),
        (Primitive::Le, &["le", "leq", "lte"]),
        (Primitive::Gt, &["gt", "greater"]),
        (Primitive::Ge, &["ge", "geq", "gte"]),
        (Primitive::Utf8, &["utf"]),
        (Primitive::First, &["fst"]),
        (Primitive::Last, &["lst"]),
        (Primitive::Slf, &["slf"]),
        (Primitive::Select, &["sel"]),
        (Primitive::Parse, &["prs"]),
        (Primitive::Rand, &["rnd"]),
        (Primitive::Voxels, &["project"]),
        (Primitive::Repr, &["rpr"]),
        (Primitive::Sign, &["sgn"]),
        (Primitive::Backward, &["bw"]),
        (Primitive::Mul, &["mul", "mult"]),
        (Primitive::Multivector, &["mv", "multi"]),
        (Primitive::OuterProduct, &["wedge", "wed"]),
    ]
    .into()
});

/// A numeric syntax component that can be parsed along with primitive names
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Sequence)]
pub enum NumComponent {
    /// The real complex unit
    R,
    /// The imaginary complex unit
    I,
}

/// A parsed primitive component
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Sequence)]
pub enum PrimComponent {
    /// A primitive
    Prim(Primitive),
    /// A numeric component
    Num(NumComponent),
    /// Subscript 0
    Sub0,
    /// Subscript 1
    Sub1,
    /// Subscript 2
    Sub2,
    /// Subscript 4
    Sub4,
    /// Subscript i
    SubI,
    /// Subscript ⌟
    SubRight,
    /// (
    OpenParen,
    /// )
    CloseParen,
    /// ε
    Epsilon,
    /// ∞
    Infinity,
    /// ℝ
    Real,
    /// ℤ
    Int,
    /// ℕ
    Nat,
    /// 𝔹
    Bool,
}

impl From<Primitive> for PrimComponent {
    fn from(prim: Primitive) -> Self {
        PrimComponent::Prim(prim)
    }
}

impl From<NumComponent> for PrimComponent {
    fn from(num: NumComponent) -> Self {
        PrimComponent::Num(num)
    }
}

impl PrimComponent {
    /// Get the name of this component
    pub fn name(&self) -> &'static str {
        match self {
            PrimComponent::Prim(prim) => prim.name(),
            PrimComponent::Num(num) => num.name(),
            PrimComponent::Sub0 => "₀",
            PrimComponent::Sub1 => "₁",
            PrimComponent::Sub2 => "₂",
            PrimComponent::Sub4 => "₄",
            PrimComponent::SubI => "ᵢ",
            PrimComponent::SubRight => "⌟",
            PrimComponent::OpenParen => "(",
            PrimComponent::CloseParen => ")",
            PrimComponent::Epsilon => "ε",
            PrimComponent::Infinity => "∞",
            PrimComponent::Real => "ℝ",
            PrimComponent::Int => "ℤ",
            PrimComponent::Nat => "ℕ",
            PrimComponent::Bool => "𝔹",
        }
    }
    /// Try to parse a component from a name prefix
    pub fn from_format_name(name: &str) -> Option<Self> {
        (Primitive::from_format_name(name).map(Into::into))
            .or_else(|| NumComponent::from_format_name(name).map(Into::into))
    }
}

impl NumComponent {
    /// Get the name of this numeric component
    pub fn name(&self) -> &'static str {
        use NumComponent::*;
        match self {
            R => "r",
            I => "i",
        }
    }
    /// Get the number value
    pub fn value(&self) -> NumWord {
        use NumComponent::*;
        match self {
            R => Complex::ONE.into(),
            I => Complex::I.into(),
        }
    }
    /// Try to parse a numeric component from a name prefix
    pub fn from_format_name(name: &str) -> Option<Self> {
        use NumComponent::*;
        Some(match name {
            "r" => R,
            "i" => I,
            _ => return None,
        })
    }
}

impl fmt::Display for NumComponent {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.name())
    }
}

impl fmt::Display for PrimComponent {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            PrimComponent::Prim(prim) => prim.fmt(f),
            PrimComponent::Num(num) => num.fmt(f),
            _ => self.name().fmt(f),
        }
    }
}

impl Primitive {
    /// Get the short aliases for this primitive
    pub fn aliases(&self) -> &'static [&'static str] {
        ALIASES.get(self).copied().unwrap_or_default()
    }
    /// Try to parse a primitive from a name prefix
    pub fn from_format_name(name: &str) -> Option<Self> {
        if name.chars().any(char::is_uppercase) {
            return None;
        }
        if name.len() < 2 {
            return None;
        }
        static REVERSE_ALIASES: LazyLock<HashMap<&'static str, Primitive>> = LazyLock::new(|| {
            ALIASES
                .iter()
                .flat_map(|(prim, aliases)| aliases.iter().map(|&s| (s, *prim)))
                .collect()
        });
        if let Some(prim) = REVERSE_ALIASES.get(name) {
            return Some(*prim);
        }
        if let Some(prim) = Primitive::non_deprecated().find(|p| p.name() == name) {
            return Some(prim);
        }
        if let Some(prim) = Primitive::all().find(|p| p.glyph().is_none() && p.name() == name) {
            return Some(prim);
        }
        if let Some(prim) = SysOp::ALL.iter().find(|s| s.name() == name) {
            return Some(Primitive::Sys(*prim));
        }
        if name.len() < 3 || name == "int" {
            return None;
        }
        let mut matching = Primitive::non_deprecated()
            .filter(|p| p.glyph().is_some() && p.name().starts_with(name));
        let res = matching.next()?;
        let exact_match = res.name() == name;
        (exact_match || matching.next().is_none()).then_some(res)
    }
    /// The list of strings where each character maps to an entire primitive
    pub fn multi_aliases() -> &'static [(&'static str, &'static [(PrimComponent, &'static str)])] {
        use Primitive::*;
        macro_rules! alias {
            (PrimComponent::$comp:ident) => {
                PrimComponent::$comp
            };
            ($prim:ident) => {
                PrimComponent::Prim($prim)
            };
            ($(($($s:ident)?, $($prim:tt)*)),+ $(,)*) => {
                (
                    concat!($(stringify!($($s)?)),*),
                    &[$((alias!($($prim)*), stringify!($($s)?))),*]
                )
            };
        }
        &[
            alias!((a, Assert), (w, With), (m, Match)),
            alias!((e, Evert), (d, Div), (o, On), (r, Range)),
            alias!((d, Div), (o, On), (r, Range)),
            alias!((p, Partition), (i, Identity), (b, By), (n, Ne)),
            alias!((p, Partition), (b, Box), (b, By), (n, Ne)),
            alias!((p, Partition), (p, Parse), (b, By), (n, Ne)),
            alias!(
                (p, Partition),
                (i, Identity),
                (n, Not),
                (b, By),
                (m, MemberOf)
            ),
            alias!((p, Partition), (b, Box), (n, Not), (b, By), (m, MemberOf)),
            alias!((p, Partition), (p, Parse), (n, Not), (b, By), (m, MemberOf)),
            alias!((k, Keep), (b, By), (n, Ne)),
            alias!((f, Fork), (t, Take), (d, Drop)),
            alias!((f, Fork), (d, Drop), (t, Take)),
            alias!((d, Drop), (n, Neg), (o, On), (d, Drop)),
            alias!(
                (a, Anti),
                (d, Drop),
                (n, Neg),
                (o, On),
                (a, Anti),
                (d, Drop)
            ),
            alias!((p, Dip), (e, Pop), (r, Under), (f, Now)),
            alias!((s, Un), (et, By)),
            alias!((wr, Sub), (en, By), (ch, Not)),
            alias!((w, Sub), (r, By), (e, Not)),
            alias!(
                (u, Un),
                (n, PrimComponent::OpenParen),
                (wr, Sub),
                (en, By),
                (ch, Not),
                (, PrimComponent::CloseParen),
            ),
            alias!(
                (u, Un),
                (n, PrimComponent::OpenParen),
                (w, Sub),
                (r, By),
                (e, Not),
                (, PrimComponent::CloseParen),
            ),
            alias!((s, Un), (huf, Sort)),
            alias!((sh, Un), (uffle, Sort)),
            alias!((sel, Select), (first, First)),
            alias!((l, Un), (og, Exp)),
            alias!((wi, Stencil), (n, Identity)),
            alias!((du, On), (p, Identity)),
            alias!((r, Pop), (e, Un), (al, Complex)),
            alias!((d, Fork), (u, Reach), (n, Both), (e, Join)),
            alias!((kor, Keep), (k, PrimComponent::Sub2)),
            alias!((eac, Rows), (h, PrimComponent::Sub0)),
            alias!((in, Add), (c, PrimComponent::Sub1)),
            alias!((de, Sub), (c, PrimComponent::Sub1)),
            alias!((eps, PrimComponent::Epsilon)),
            alias!((epsil, PrimComponent::Epsilon)),
            alias!((epsilon, PrimComponent::Epsilon)),
            alias!((r, Reduce), (a, Content), (ze, Join)),
            alias!((r, Reduce), (a, Content), (z, Join)),
            alias!((li, Validate), (st, PrimComponent::Infinity)),
            alias!(
                (s, Validate),
                (t, PrimComponent::Sub1),
                (r, PrimComponent::Infinity)
            ),
            alias!((nu, Validate), (m, PrimComponent::Real)),
            alias!((in, Validate), (t, PrimComponent::Int)),
            alias!((na, Validate), (t, PrimComponent::Nat)),
            alias!((bo, Validate), (ol, PrimComponent::Bool)),
            // alias!((co, Multivector), (mp, PrimComponent::SubI)),
            // alias!((com, Multivector), (pl, PrimComponent::SubI)),
            // alias!((comp, Multivector), (plex, PrimComponent::SubI)),
            alias!((co, Neg), (nj, PrimComponent::SubRight)),
            alias!((conju, Neg), (gate, PrimComponent::SubRight)),
            alias!((du, Neg), (al, PrimComponent::Sub4)),
            alias!((r, Under), (e, Both), (g, Neg), (, PrimComponent::Sub4), (r, OuterProduct)),
            alias!(
                (re, Under),
                (g, Both),
                (r, Neg),
                (e, PrimComponent::Sub4),
                (s, OuterProduct)
            ),
            alias!(
                (re, Under),
                (gr, Both),
                (es, Neg),
                (si, PrimComponent::Sub4),
                (ve, OuterProduct)
            ),
            alias!(
                (s, Mul),
                (a, Fork),
                (n, Neg),
                (d, PrimComponent::SubRight),
                (w, Backward),
                (i, Mul)
            ),
            alias!(
                (sa, Mul),
                (b, Fork),
                (d, Neg),
                (w, PrimComponent::SubRight),
                (i, Backward),
                (ch, Mul)
            ),
        ]
    }
    /// Look up a multi-alias from [`Self::multi_aliases`]
    pub fn get_multi_alias(name: &str) -> Option<&'static [(PrimComponent, &'static str)]> {
        Self::multi_aliases()
            .iter()
            .find(|(alias, _)| *alias == name)
            .map(|(_, aliases)| *aliases)
    }
}

/// Try to parse multiple primitives from the concatenation of their name prefixes
pub fn split_name(name: &str) -> Option<Vec<(PrimComponent, &str)>> {
    let mut indices: Vec<usize> = name.char_indices().map(|(i, _)| i).collect();
    indices.push(name.len());
    // Forward parsing
    let mut prims = Vec::new();
    let mut start = 0;
    'outer: loop {
        if start == indices.len() {
            return Some(prims);
        }
        let start_index = indices[start];
        for len in (1..=indices.len() - start).rev() {
            let end_index = indices.get(start + len).copied().unwrap_or(name.len());
            if start_index == end_index {
                continue;
            }
            let sub_name = &name[start_index..end_index];
            // Normal primitive matching
            if let Some(p) = PrimComponent::from_format_name(sub_name)
                && (sub_name.len() > 1 || p.name().len() == 1)
            {
                prims.push((p, sub_name));
                start += len;
                continue 'outer;
            }
            // Greek
            if sub_name.chars().count() == 1 {
                for prim in [Primitive::Eta, Primitive::Pi, Primitive::Tau] {
                    if sub_name.chars().next() == prim.glyph() {
                        prims.push((prim.into(), sub_name));
                        start += len;
                        continue 'outer;
                    }
                }
            }
            if sub_name.len() == 1 {
                continue;
            }
            // 1-letter planet notation
            let unforked = sub_name.strip_prefix('f').unwrap_or(sub_name);
            if unforked
                .strip_suffix(['i', 'p', 'f'])
                .unwrap_or(unforked)
                .chars()
                .all(|c| "gd".contains(c))
            {
                for (i, c) in sub_name.char_indices() {
                    let prim = match c {
                        'f' if i == 0 => Primitive::Fork,
                        'f' => Primitive::Fix,
                        'g' => Primitive::Gap,
                        'd' => Primitive::Dip,
                        'i' => Primitive::Identity,
                        'p' => Primitive::Pop,
                        _ => unreachable!(),
                    };
                    prims.push((prim.into(), &sub_name[i..i + 1]))
                }
                start += len;
                continue 'outer;
            }
            // Aliases
            if let Some(ps) = Primitive::get_multi_alias(sub_name) {
                prims.extend(ps.iter().map(|(p, s)| (*p, *s)));
                start += len;
                continue 'outer;
            }
        }
        break;
    }
    // Backward parsing
    prims.clear();
    let mut end = indices.len() - 1;
    'outer: loop {
        if end == 0 {
            prims.reverse();
            return Some(prims);
        }
        let end_index = indices[end];
        for len in (1..=end).rev() {
            let start_index = indices.get(end - len).copied().unwrap_or(0);
            let sub_name = &name[start_index..end_index];
            // Normal primitive matching
            if let Some(p) = Primitive::from_format_name(sub_name)
                && (sub_name.len() > 1 || p.name().len() == 1)
            {
                prims.push((p.into(), sub_name));
                end -= len;
                continue 'outer;
            }
            // Greek
            if sub_name.chars().count() == 1 {
                for prim in [Primitive::Eta, Primitive::Pi, Primitive::Tau] {
                    if sub_name.chars().next() == prim.glyph() {
                        prims.push((prim.into(), sub_name));
                        end -= len;
                        continue 'outer;
                    }
                }
            }
            if sub_name.len() == 1 {
                continue;
            }
            // 1-letter planet notation
            if sub_name
                .strip_prefix('f')
                .unwrap_or(sub_name)
                .strip_suffix(['i', 'p'])
                .unwrap_or(sub_name)
                .chars()
                .all(|c| "gd".contains(c))
                && sub_name != "fi"
            {
                for (i, c) in sub_name.char_indices().rev() {
                    let prim = match c {
                        'f' => Primitive::Fork,
                        'g' => Primitive::Gap,
                        'd' => Primitive::Dip,
                        'i' => Primitive::Identity,
                        'p' => Primitive::Pop,
                        _ => unreachable!(),
                    };
                    prims.push((prim.into(), &sub_name[i..i + 1]))
                }
                end -= len;
                continue 'outer;
            }
            // Dip fix
            if sub_name
                .strip_suffix('f')
                .unwrap_or(sub_name)
                .chars()
                .all(|c| c == 'd')
            {
                for (i, c) in sub_name.char_indices().rev() {
                    let prim = match c {
                        'd' => Primitive::Dip,
                        'f' => Primitive::Fix,
                        _ => unreachable!(),
                    };
                    prims.push((prim.into(), &sub_name[i..i + 1]))
                }
                end -= len;
                continue 'outer;
            }
            // Aliases
            if let Some(ps) = Primitive::get_multi_alias(sub_name) {
                prims.extend(ps.iter().rev().map(|(p, s)| (*p, *s)));
                end -= len;
                continue 'outer;
            }
            // Greek
            if sub_name.chars().count() == 1 {
                for prim in [Primitive::Eta, Primitive::Pi, Primitive::Tau] {
                    if sub_name.chars().next() == prim.glyph() {
                        prims.push((prim.into(), sub_name));
                        end -= len;
                        continue 'outer;
                    }
                }
            }
        }
        break;
    }
    None
}