ruwex 0.1.0

Fast Rust rewrite of wikiextractor: extract and clean text from Wikimedia XML dumps
Documentation
//! Brace and bracket matching, ported from wikiextractor's
//! `findMatchingBraces` and `findBalanced` — including their handling of
//! ambiguous and unbalanced input, which downstream behavior depends on.

use std::sync::LazyLock;

use regex::{Match, Regex};

static OPEN_BRACES_2: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\{{2,}").unwrap());
static OPEN_BRACES_3: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\{{3,}").unwrap());
static NEXT_BRACES: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\{{2,}|\}{2,}").unwrap());
static OPEN_ANY: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\{{2,}|\[{2,}").unwrap());
static NEXT_ANY: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"\{{2,}|\}{2,}|\[{2,}|\]{2,}").unwrap());

/// Finds spans of matching `{{..}}` / `{{{..}}}` (and, with `ldelim == 0`,
/// also `[[..]]`) expressions. `ldelim` is the minimum number of opening
/// braces (wikiextractor uses 2 for template expansion, 0 for parameter
/// splitting).
///
/// Stack entries are opening-run lengths; negative values track `[[` runs.
pub fn find_matching_braces(text: &str, ldelim: usize) -> Vec<(usize, usize)> {
    let (re_open, re_next): (&Regex, &Regex) = match ldelim {
        0 => (&OPEN_ANY, &NEXT_ANY),
        2 => (&OPEN_BRACES_2, &NEXT_BRACES),
        3 => (&OPEN_BRACES_3, &NEXT_BRACES),
        _ => unreachable!("ldelim is 0 (parameters), 2 (templates) or 3 (tplargs)"),
    };

    let mut result = Vec::new();
    let mut cur = 0;
    'outer: loop {
        let Some(m1) = re_open.find_at(text, cur) else {
            return result;
        };
        let mut stack: Vec<i64> = Vec::new();
        let first_len = (m1.end() - m1.start()) as i64;
        if text.as_bytes()[m1.start()] == b'{' {
            stack.push(first_len);
        } else {
            stack.push(-first_len);
        }
        let mut end = m1.end();
        loop {
            let Some(m2) = re_next.find_at(text, end) else {
                return result; // unbalanced
            };
            end = m2.end();
            let mut lmatch = (m2.end() - m2.start()) as i64;
            match text.as_bytes()[m2.start()] {
                b'{' => stack.push(lmatch),
                b'}' => {
                    while let Some(open_count) = stack.pop() {
                        if open_count == 0 {
                            continue; // illegal unmatched [[
                        }
                        if lmatch >= open_count {
                            lmatch -= open_count;
                            if lmatch <= 1 {
                                break; // either close or stray }
                            }
                        } else {
                            // put back unmatched
                            stack.push(open_count - lmatch);
                            break;
                        }
                    }
                    if stack.is_empty() {
                        result.push((m1.start(), (end as i64 - lmatch) as usize));
                        cur = end;
                        continue 'outer;
                    } else if stack.len() == 1 && 0 < stack[0] && stack[0] < ldelim as i64 {
                        // ambiguous {{{{{ }}} }}
                        result.push((m1.start() + stack[0] as usize, end));
                        cur = end;
                        continue 'outer;
                    }
                }
                b'[' => stack.push(-lmatch),
                b']' => {
                    while stack.last().is_some_and(|&top| top < 0) {
                        let open_count = -stack.pop().expect("checked non-empty");
                        if lmatch >= open_count {
                            lmatch -= open_count;
                            if lmatch <= 1 {
                                break; // either close or stray ]
                            }
                        } else {
                            // put back unmatched (negative)
                            stack.push(lmatch - open_count);
                            break;
                        }
                    }
                    if stack.is_empty() {
                        result.push((m1.start(), (end as i64 - lmatch) as usize));
                        cur = end;
                        continue 'outer;
                    }
                    // unmatched ]] are discarded; scanning continues
                }
                _ => unreachable!("regex only matches braces and brackets"),
            }
        }
    }
}

/// Finds spans of balanced `open`…`close` pairs (e.g. `[[`…`]]`), skipping
/// stray closers outside any pair.
pub fn find_balanced(text: &str, open: &str, close: &str) -> Vec<(usize, usize)> {
    let mut result = Vec::new();
    let mut depth = 0usize;
    let mut start = 0usize;
    let mut cur = 0usize;
    loop {
        let next = if depth == 0 {
            // only opening delimiters start a new expression
            text[cur..].find(open).map(|i| (cur + i, true))
        } else {
            let o = text[cur..].find(open).map(|i| cur + i);
            let c = text[cur..].find(close).map(|i| cur + i);
            match (o, c) {
                (Some(a), Some(b)) if a < b => Some((a, true)),
                (_, Some(b)) => Some((b, false)),
                (Some(a), None) => Some((a, true)),
                (None, None) => None,
            }
        };
        let Some((pos, is_open)) = next else {
            return result;
        };
        if depth == 0 {
            start = pos;
        }
        if is_open {
            depth += 1;
            cur = pos + open.len();
        } else {
            depth -= 1;
            cur = pos + close.len();
            if depth == 0 {
                result.push((start, cur));
            }
        }
    }
}

/// Convenience wrapper matching a regex starting exactly at `pos`
/// (Python's `pattern.match(text, pos)`).
pub fn match_at<'t>(re: &Regex, text: &'t str, pos: usize) -> Option<Match<'t>> {
    re.find_at(text, pos).filter(|m| m.start() == pos)
}

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

    fn spans(text: &str, ldelim: usize) -> Vec<&str> {
        find_matching_braces(text, ldelim)
            .into_iter()
            .map(|(s, e)| &text[s..e])
            .collect()
    }

    #[test]
    fn matches_simple_and_nested_templates() {
        assert_eq!(
            spans("a {{x}} b {{y|{{z}}}} c", 2),
            vec!["{{x}}", "{{y|{{z}}}}"]
        );
    }

    #[test]
    fn matches_triple_braces_as_one_span() {
        assert_eq!(spans("{{{1|default}}}", 2), vec!["{{{1|default}}}"]);
    }

    #[test]
    fn leaves_stray_closers_out_of_span() {
        // {{ }}}} -> span is {{ }} with trailing }} left as text
        assert_eq!(spans("{{ x }}}}", 2), vec!["{{ x }}"]);
    }

    #[test]
    fn unbalanced_open_is_dropped_silently() {
        assert_eq!(spans("a {{x b c", 2), Vec::<&str>::new());
    }

    #[test]
    fn brackets_tracked_in_zero_delim_mode() {
        assert_eq!(
            spans("a [[link|x]] and {{t|p}}", 0),
            vec!["[[link|x]]", "{{t|p}}"]
        );
    }

    #[test]
    fn balanced_finds_nested_links() {
        let text = "see [[File:x|cap [[inner]] end]] tail [[plain]]";
        let found: Vec<&str> = find_balanced(text, "[[", "]]")
            .into_iter()
            .map(|(s, e)| &text[s..e])
            .collect();
        assert_eq!(found, vec!["[[File:x|cap [[inner]] end]]", "[[plain]]"]);
    }

    #[test]
    fn balanced_skips_stray_closers() {
        let text = "a ]] b [[x]] c";
        let found = find_balanced(text, "[[", "]]");
        assert_eq!(found, vec![(7, 12)]);
    }
}