src/gbnf.rs

// SPDX-License-Identifier: GPL-3.0-or-later OR AGPL-3.0-or-later
// Copyright (C) 2025  Red Hat, Inc.

use either::Either;
use rayon::prelude::*;
use std::fs::File;
use std::io::Write;

use crate::Args;
use crate::{FileRange, FileRanges};

fn __quote_char(c: char, in_range: bool) -> String {
    match c {
        '\t' => "\\t".to_string(),
        '\r' => "\\r".to_string(),
        '\n' => "\\n".to_string(),
        '"' | '\\' => format!("\\{}", c),
        '[' | ']' if in_range => format!("\\{}", c),
        _ => c.to_string(),
    }
}

fn quote_char_in_range(c: char) -> String {
    __quote_char(c, true)
}

fn quote_char(c: char) -> String {
    __quote_char(c, false)
}

fn negative(control_line: Option<&String>, separator: &str) -> String {
    let mut rule = String::new();
    let mut rule_terminated = String::new();

    let mut chars_sep = separator.chars().collect::<Vec<char>>();
    chars_sep.push('\n');
    let chars = match control_line {
        Some(s) => {
            let mut chars = s.chars().collect::<Vec<char>>();
            chars.push('\n');
            chars
        }
        None => Vec::new(),
    };

    let nr = chars.len().max(chars_sep.len());
    rule.push_str("( ( (\n");
    for i in 0..nr {
        let mut positive = String::new();
        let mut positive_sep = String::new();
        for (j, (c, s)) in chars.iter().zip(chars_sep.iter()).enumerate().take(i + 1) {
            if i == j {
                let c = if *c != '\n' {
                    &quote_char_in_range(*c)
                } else {
                    ""
                };
                let s = if *s != '\n' {
                    &quote_char_in_range(*s)
                } else {
                    ""
                };
                let positive = if positive.is_empty() {
                    ""
                } else {
                    &format!("\"{}\"", positive)
                };
                let positive_sep = if positive_sep.is_empty() {
                    ""
                } else {
                    &format!("\"{}\"", positive_sep)
                };
                if i > 0 {
                    rule.push_str(&format!(
                        "{} [^\\n{}] |\n {} [^\\n{}]",
                        positive, c, positive_sep, s
                    ));
                    if !c.is_empty() {
                        rule_terminated.push_str(&format!(" | {}\n", positive));
                    }
                    if !s.is_empty() {
                        rule_terminated.push_str(&format!(" | {}\n", positive_sep));
                    }
                } else {
                    rule.push_str(&format!(" [^\\n{}{}]", c, s));
                }
            } else {
                positive.push_str(&quote_char(*c));
                positive_sep.push_str(&quote_char(*s));
            }
        }
        let (chars, mut positive, short_len) = if chars.len() > chars_sep.len() {
            (&chars, positive, chars_sep.len())
        } else {
            (&chars_sep, positive_sep, chars.len())
        };

        for (j, c) in chars.iter().enumerate().take(i + 1) {
            if j < short_len {
                continue;
            }
            if i == j {
                let c = if *c != '\n' {
                    &quote_char_in_range(*c)
                } else {
                    ""
                };
                let positive = if positive.is_empty() {
                    ""
                } else {
                    &format!("\"{}\"", positive)
                };
                rule.push_str(&format!("{} [^\\n{}]", positive, c));
                if !c.is_empty() && !positive.is_empty() {
                    rule_terminated.push_str(&format!(" | {}\n", positive));
                }
            } else {
                positive.push_str(&quote_char(*c));
            }
        }
        if i < nr - 1 {
            rule.push_str(" |\n ");
        }
    }
    rule.push_str(") [^\\n]*\n");
    rule.push_str(&rule_terminated);
    rule.push_str(")? [\\n] )*\n");
    rule
}

macro_rules! filter_unique {
    ($iter:expr, $lines:expr) => {
        $iter.filter(|line| $lines.iter().filter(|x| x == line).count() == 1)
    };
}

pub fn find_distant_control_line<'a>(
    lines: &'a [String],
    control_lines: &'a [String],
    after: bool,
) -> Option<&'a String> {
    let unique_control_lines: Vec<_> = filter_unique!(control_lines.iter(), lines).collect();
    if unique_control_lines.len() == 1 {
        return Some(unique_control_lines[0]);
    }
    let different_control_lines: Vec<String> = {
        let mut seen_lines = std::collections::HashSet::new();
        control_lines
            .iter()
            .filter(|line| seen_lines.insert(line.to_string()))
            .map(|s| s.to_string())
            .collect()
    };
    let mut max_dist = 0.0;
    let mut max_line: Option<&String> = None;
    let mut seen = std::collections::HashMap::new();
    let (unique_control_lines_iter, boundary_line) = if !after {
        (
            Either::Left(unique_control_lines.iter().rev()),
            lines.first(),
        )
    } else {
        (Either::Right(unique_control_lines.iter()), lines.last())
    };
    for control_line in unique_control_lines_iter {
        let mut control_min_dist: Option<f64> = None;
        for line in &different_control_lines {
            if **control_line == *line {
                continue;
            }
            let key = if control_line < &line {
                (*control_line, line)
            } else {
                (line, *control_line)
            };
            let dist = seen.get(&key);
            let dist = match dist {
                Some(dist) => *dist,
                None => {
                    let dist =
                        textdistance::nstr::damerau_levenshtein_restricted(control_line, line);
                    //println!("{}\n{}\n{}\n", dist, control_line, line);
                    seen.insert(key, dist);
                    dist
                }
            };
            if control_min_dist.is_none_or(|d| dist < d) {
                control_min_dist = Some(dist);
            }
        }
        if (control_min_dist.is_some_and(|d| d > max_dist)
            && control_line != &boundary_line.unwrap())
            || max_line.is_none()
        {
            max_dist = control_min_dist.unwrap_or(0.0);
            max_line = Some(control_line);
        }
    }
    if max_line.is_none() {
        assert!(unique_control_lines.is_empty());
    }
    max_line
}

pub fn find_long_control_line<'a>(
    lines: &'a [String],
    control_lines: &'a [String],
    after: bool,
) -> Option<&'a String> {
    let mut max_len = 0;
    let mut max_line: Option<&String> = None;
    let unique_control_lines_iter = if !after {
        Either::Left(filter_unique!(control_lines.iter(), lines).rev())
    } else {
        Either::Right(filter_unique!(control_lines.iter(), lines))
    };
    for line in unique_control_lines_iter {
        if line.len() > max_len || max_line.is_none() {
            max_len = line.len();
            max_line = Some(line);
        }
    }
    max_line
}

pub fn find_control_line<'a>(
    lines: &'a [String],
    control_lines: &'a [String],
    after: bool,
) -> Option<&'a String> {
    let control_line = if false {
        find_long_control_line(lines, control_lines, after)
    } else {
        find_distant_control_line(lines, control_lines, after)
    };
    if let Some(found) = control_line {
        assert!(control_lines.contains(found));
    } else if !after {
        // use the last dup as "start" control line
        return control_lines.last();
    }
    control_line
}

fn push_control_line(file_rule: &mut String, control_line: &str) {
    file_rule.push_str(" \"");
    file_rule.push_str(&control_line.chars().map(quote_char).collect::<String>());
    file_rule.push_str("\" [\\n] \n");
}

fn find_control_lines_parallel(
    file_ranges: &mut FileRanges,
    args: &Args,
) -> Result<(), anyhow::Error> {
    let num_cpus = if args.threads == 0 {
        num_cpus::get()
    } else {
        args.threads.try_into().unwrap()
    };

    let before_context = args.context.max(args.before_context) as usize;
    let after_context = args.context.max(args.after_context) as usize;

    // Define a closure to process a single range
    let process_range = |range: &FileRange| -> (Option<String>, Option<String>, usize, usize) {
        let lines = &range.lines;
        let before_context = range.before_context.saturating_sub(before_context);
        let after_context = range.after_context.saturating_sub(after_context);

        let control_lines_before = &lines[..before_context];
        let control_lines_after = &lines[lines.len() - after_context..];

        let control_line_before = find_control_line(lines, control_lines_before, false);
        let control_line_after = find_control_line(lines, control_lines_after, true);

        let control_line_before_index = if let Some(control_line) = control_line_before {
            control_lines_before
                .iter()
                .rposition(|line| line == control_line)
                .unwrap()
        } else {
            control_lines_before.len()
        };

        let control_line_after_index = if let Some(control_line) = control_line_after {
            control_lines_after.len()
                - control_lines_after
                    .iter()
                    .position(|line| line == control_line)
                    .unwrap()
                - 1
        } else {
            control_lines_after.len()
        };

        (
            control_line_before.cloned(),
            control_line_after.cloned(),
            control_line_before_index,
            control_line_after_index,
        )
    };

    // Process all ranges sequentially
    let all_ranges: Vec<_> = file_ranges.hash.values().flatten().collect();

    // Process ranges in parallel if num_threads > 1, otherwise process sequentially
    let results: Vec<_> = if num_cpus > 1 {
        // Create a thread pool with limited parallelism
        let pool = rayon::ThreadPoolBuilder::new()
            .num_threads(num_cpus)
            .build()
            .unwrap();

        // Process all ranges in parallel using the configured thread pool
        pool.install(|| {
            all_ranges
                .par_iter()
                .cloned()
                .map(process_range)
                .collect::<Vec<_>>()
        })
    } else {
        all_ranges
            .into_iter()
            .map(process_range)
            .collect::<Vec<_>>()
    };

    // Update ranges using stable iteration
    let mut result_index = 0;
    for ranges in file_ranges.hash.values_mut() {
        for range in ranges.iter_mut() {
            let result: &(Option<String>, Option<String>, usize, usize) = &results[result_index];
            let (
                control_line_before,
                control_line_after,
                control_line_before_index,
                control_line_after_index,
            ) = result;
            range.control_line_before = control_line_before.clone();
            range.control_line_after = control_line_after.clone();
            range.control_line_before_index = *control_line_before_index;
            range.control_line_after_index = *control_line_after_index;
            result_index += 1;
        }
    }

    Ok(())
}

pub fn generate_gbnf_file(
    file: &mut File,
    file_ranges: &mut FileRanges,
    args: &Args,
) -> Result<(), anyhow::Error> {
    find_control_lines_parallel(file_ranges, args)?;

    let mut file_rules = Vec::new();

    for (filename, ranges) in &mut file_ranges.hash {
        assert!(!ranges.is_empty());
        let mut file_rule = String::new();

        file_rule.push_str("prefix \"");
        file_rule.push_str(filename);
        file_rule.push_str("\" \"\\n\"\n");

        for range in ranges {
            // Generate rules for each snippet
            let control_line_before = range.control_line_before.as_ref();
            let control_line_after = range.control_line_after.as_ref();

            if let Some(control_line) = control_line_before {
                push_control_line(&mut file_rule, control_line);
            }
            let index = range.control_line_before_index;
            range.start += index;
            range.lines.drain(..index);

            let index = range.control_line_after_index;
            range.end -= index;
            range.lines.drain(range.lines.len() - index..);

            assert!(range.end - range.start == range.lines.len() - 1);

            file_rule.push_str(&negative(control_line_after, &args.context_separator));
            if let Some(control_line) = control_line_after {
                file_rule.push_str(" ( ");
                push_control_line(&mut file_rule, control_line);
                file_rule.push_str(" )?\n");
            }
            file_rule.push_str(" separator\n");
        }
        file_rule.push_str(" [\\n]?");

        file_rules.push(file_rule);
    }

    let mut gbnf_content = String::new();
    gbnf_content.push_str("root ::= \"```\\n\" ( ");
    gbnf_content.push_str(&file_rules.join(" |\n"));
    if false {
        gbnf_content.push_str(&format!(
            " ){{1,{}}} \"```\"\n",
            file_ranges.filenames.len()
        ));
    } else {
        /*
         * The DoS fixes in llama.cpp
         * 990e4d96980d0b016a2b07049cc9031642fb9903 rightfully had to
         * enforce some limits for repetitions, while unlimited
         * repetitions aren't bound.
         */
        gbnf_content.push_str(" )+ \"```\"\n");
    }

    gbnf_content.push_str("prefix ::= \"");
    gbnf_content.push_str(
        &args
            .filename_prefix
            .chars()
            .map(quote_char)
            .collect::<String>(),
    );
    gbnf_content.push_str("\"\n");

    gbnf_content.push_str("separator ::= \"");
    gbnf_content.push_str(
        &args
            .context_separator
            .chars()
            .map(quote_char)
            .collect::<String>(),
    );
    gbnf_content.push_str("\\n\"");

    writeln!(file, "{}", gbnf_content)?;
    file.flush()?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs::Permissions;
    use std::os::unix::fs::PermissionsExt;

    #[test]
    fn test_gbnf_generation() {
        let mut temp_file = tempfile::Builder::new()
            .permissions(Permissions::from_mode(0o700))
            .prefix("rg-edit-")
            .suffix(&format!(
                "{}{}",
                crate::RG_EDIT_EXTENSION,
                crate::GBNF_EXTENSION
            ))
            .tempfile()
            .unwrap();

        // Create a mock Args
        let args = Args {
            context_separator: "-- DO NOT DELETE THIS SEPARATOR".to_string(),
            filename_prefix: "-- ".to_string(),
            ..Default::default()
        };

        // This test would require more complex setup to test actual FileRanges
        // For now, we just ensure the function compiles and doesn't panic
        let mut file_ranges = FileRanges::new("", "", "").unwrap();
        assert!(generate_gbnf_file(temp_file.as_file_mut(), &mut file_ranges, &args).is_ok());
    }
}

// Local Variables:
// rust-format-on-save: t
// End: