regast-core 0.1.0

use regast_syntax::{AstKind, ClassSet, NodeId, NodeMap, Pattern, RepKind};

use crate::{Disambiguation, IrId, IrPool, Prov};

const MAX_REPEAT: u32 = 1_000;

#[derive(Clone, Debug)]
pub struct Lowering {
    pub root: IrId,
    pub origins: NodeMap<Vec<IrId>>,
}

impl Lowering {
    #[must_use]
    pub fn ir_for(&self, node: NodeId) -> Option<IrId> {
        self.origins.get(node)?.last().copied()
    }
}

pub fn lower(pattern: &Pattern, pool: &mut IrPool, disambiguation: Disambiguation) -> Lowering {
    let mut origins: NodeMap<Vec<IrId>> = NodeMap::new(pattern);
    for cursor in pattern.root().walk().collect::<Vec<_>>().into_iter().rev() {
        let node = cursor.id();
        let prov = Prov::Ast(node);
        let root = match cursor.kind() {
            AstKind::Empty => pool.one(),
            AstKind::Anchor { kind } => pool.anchor(*kind, prov),
            AstKind::Literal { c, .. } => {
                let set = pool.class_set(ClassSet::singleton(*c));
                pool.class(set, prov)
            }
            AstKind::Dot => {
                let set = pool.class_set(ClassSet::singleton('\n').negate());
                pool.class(set, prov)
            }
            AstKind::Class { set, .. } => {
                let set = pool.class_set(set.clone());
                pool.class(set, prov)
            }
            AstKind::Group { inner, .. } => origins.get(*inner).expect("inner lowered")[0],
            AstKind::Alt { arms } => fold_alt(
                &origins,
                pool,
                arms,
                node,
                disambiguation == Disambiguation::Greedy,
            ),
            AstKind::Concat { parts } => fold_seq(&origins, pool, parts, node),
            AstKind::Repeat {
                inner,
                kind,
                greedy,
            } => {
                let inner = origins.get(*inner).expect("inner lowered")[0];
                let greedy = disambiguation == Disambiguation::Posix || *greedy;
                lower_repeat(pool, inner, *kind, greedy, node)
            }
        };
        origins.insert(node, vec![root]);
    }
    Lowering {
        root: origins.get(pattern.root_id()).expect("root lowered")[0],
        origins,
    }
}

fn fold_alt(
    origins: &NodeMap<Vec<IrId>>,
    pool: &mut IrPool,
    arms: &[NodeId],
    node: NodeId,
    ordered: bool,
) -> IrId {
    let mut iter = arms.iter().rev();
    let mut result = origins
        .get(*iter.next().expect("at least two arms"))
        .unwrap()[0];
    for arm in iter {
        let left = origins.get(*arm).unwrap()[0];
        result = pool.alt(left, result, Prov::Ast(node), ordered);
    }
    result
}

fn fold_seq(
    origins: &NodeMap<Vec<IrId>>,
    pool: &mut IrPool,
    parts: &[NodeId],
    node: NodeId,
) -> IrId {
    let mut iter = parts.iter().rev();
    let mut result = origins
        .get(*iter.next().expect("at least two parts"))
        .unwrap()[0];
    for part in iter {
        let left = origins.get(*part).unwrap()[0];
        result = pool.seq(left, result, Prov::Ast(node));
    }
    result
}

fn lower_repeat(pool: &mut IrPool, inner: IrId, kind: RepKind, greedy: bool, node: NodeId) -> IrId {
    let prov = Prov::Ast(node);
    match kind {
        RepKind::ZeroOrOne => optional(pool, inner, greedy, node),
        RepKind::ZeroOrMore => pool.star(inner, prov, greedy),
        RepKind::OneOrMore => {
            let star = pool.star(inner, Prov::Synth(node), greedy);
            pool.seq(inner, star, prov)
        }
        RepKind::Range { min, max } => {
            assert!(min <= MAX_REPEAT, "repeat lower bound exceeds limit");
            assert!(
                max.is_none_or(|max| max <= MAX_REPEAT),
                "repeat upper bound exceeds limit"
            );
            lower_range(pool, inner, min, max, greedy, node)
        }
    }
}

fn lower_range(
    pool: &mut IrPool,
    inner: IrId,
    min: u32,
    max: Option<u32>,
    greedy: bool,
    node: NodeId,
) -> IrId {
    let (mut suffix, has_suffix) = match max {
        None => (pool.star(inner, Prov::Synth(node), greedy), true),
        Some(max) => {
            let mut optional_suffix = pool.one();
            for _ in 0..max - min {
                let seq = if optional_suffix == pool.one() {
                    inner
                } else {
                    pool.seq(inner, optional_suffix, Prov::Synth(node))
                };
                optional_suffix = optional(pool, seq, greedy, node);
            }
            (optional_suffix, max > min)
        }
    };
    if min == 0 {
        return suffix;
    }
    if !has_suffix {
        suffix = inner;
        for _ in 1..min {
            suffix = pool.seq(inner, suffix, Prov::Synth(node));
        }
        return suffix;
    }
    for _ in 0..min {
        suffix = pool.seq(inner, suffix, Prov::Synth(node));
    }
    suffix
}

fn optional(pool: &mut IrPool, inner: IrId, greedy: bool, node: NodeId) -> IrId {
    if greedy {
        pool.alt(inner, pool.one(), Prov::Ast(node), false)
    } else {
        pool.alt(pool.one(), inner, Prov::Ast(node), true)
    }
}

#[cfg(test)]
mod tests {
    use regast_syntax::Pattern;

    use super::*;

    #[test]
    fn lowers_every_repeat_form() {
        for pattern in ["a?", "a*", "a+", "a{2}", "a{1,3}", "a{2,}", "a??"] {
            let pattern = Pattern::parse(pattern).unwrap();
            let mut pool = IrPool::new();
            let lowered = lower(&pattern, &mut pool, Disambiguation::Posix);
            assert!(lowered.root.index() < u32::try_from(pool.len()).unwrap());
        }
    }
}