ai-harness/crates/harness-tools/src/edit/replacers.rs

//! Verbatim port of opencode's replacer chain (`packages/opencode/src/tool/edit.ts`).
//! Each replacer returns candidate matches in the same order the TS generator yields them;
//! `replace()` walks replacers in a fixed sequence and takes the first candidate that
//! resolves unambiguously (or throws — the disproportionate-match guard is a hard stop,
//! not a "try the next candidate").

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum EditError {
    #[error("No changes to apply: oldString and newString are identical.")]
    Identical,
    #[error("oldString cannot be empty when editing an existing file. Provide the exact text to replace, or use write for an intentional full-file replacement.")]
    EmptyOldString,
    #[error("Could not find oldString in the file. It must match exactly, including whitespace, indentation, and line endings.")]
    NotFound,
    #[error("Found multiple matches for oldString. Provide more surrounding context to make the match unique.")]
    MultipleMatches,
    #[error("Refusing replacement because the matched span is much larger than oldString. Re-read the file and provide the full exact oldString for the intended replacement.")]
    DisproportionateMatch,
}

const SINGLE_CANDIDATE_SIMILARITY_THRESHOLD: f64 = 0.65;
const MULTIPLE_CANDIDATES_SIMILARITY_THRESHOLD: f64 = 0.65;

fn levenshtein(a: &str, b: &str) -> usize {
    let a: Vec<char> = a.chars().collect();
    let b: Vec<char> = b.chars().collect();
    if a.is_empty() || b.is_empty() {
        return a.len().max(b.len());
    }
    let mut matrix = vec![vec![0usize; b.len() + 1]; a.len() + 1];
    for (i, row) in matrix.iter_mut().enumerate() {
        row[0] = i;
    }
    for (j, cell) in matrix[0].iter_mut().enumerate() {
        *cell = j;
    }
    for i in 1..=a.len() {
        for j in 1..=b.len() {
            let cost = if a[i - 1] == b[j - 1] { 0 } else { 1 };
            matrix[i][j] = (matrix[i - 1][j] + 1)
                .min(matrix[i][j - 1] + 1)
                .min(matrix[i - 1][j - 1] + cost);
        }
    }
    matrix[a.len()][b.len()]
}

fn simple(_content: &str, find: &str) -> Vec<String> {
    vec![find.to_string()]
}

fn line_trimmed(content: &str, find: &str) -> Vec<String> {
    let original_lines: Vec<&str> = content.split('\n').collect();
    let mut search_lines: Vec<&str> = find.split('\n').collect();
    if search_lines.last() == Some(&"") {
        search_lines.pop();
    }
    let mut out = Vec::new();
    if search_lines.is_empty() || original_lines.len() < search_lines.len() {
        return out;
    }

    for i in 0..=(original_lines.len() - search_lines.len()) {
        let matches =
            (0..search_lines.len()).all(|j| original_lines[i + j].trim() == search_lines[j].trim());
        if matches {
            let mut match_start = 0usize;
            for line in &original_lines[..i] {
                match_start += line.len() + 1;
            }
            let mut match_end = match_start;
            for k in 0..search_lines.len() {
                match_end += original_lines[i + k].len();
                if k < search_lines.len() - 1 {
                    match_end += 1;
                }
            }
            out.push(content[match_start..match_end].to_string());
        }
    }
    out
}

fn extract_lines(original_lines: &[&str], content: &str, start: usize, end: usize) -> String {
    let mut match_start = 0usize;
    for line in &original_lines[..start] {
        match_start += line.len() + 1;
    }
    let mut match_end = match_start;
    for (k, line) in original_lines.iter().enumerate().take(end + 1).skip(start) {
        match_end += line.len();
        if k < end {
            match_end += 1;
        }
    }
    content[match_start..match_end].to_string()
}

fn block_anchor(content: &str, find: &str) -> Vec<String> {
    let original_lines: Vec<&str> = content.split('\n').collect();
    let mut search_lines: Vec<&str> = find.split('\n').collect();
    if search_lines.len() < 3 {
        return vec![];
    }
    if search_lines.last() == Some(&"") {
        search_lines.pop();
    }

    let first_line_search = search_lines[0].trim();
    let last_line_search = search_lines[search_lines.len() - 1].trim();
    let search_block_size = search_lines.len();
    let max_line_delta = ((search_block_size as f64 * 0.25).floor() as i64).max(1);

    let mut candidates: Vec<(usize, usize)> = Vec::new();
    for i in 0..original_lines.len() {
        if original_lines[i].trim() != first_line_search {
            continue;
        }
        let mut j = i + 2;
        while j < original_lines.len() {
            if original_lines[j].trim() == last_line_search {
                let actual_block_size = (j - i + 1) as i64;
                if (actual_block_size - search_block_size as i64).abs() <= max_line_delta {
                    candidates.push((i, j));
                }
                break;
            }
            j += 1;
        }
    }
    if candidates.is_empty() {
        return vec![];
    }

    let middle_similarity = |start: usize, end: usize, early_exit: bool| -> f64 {
        let actual_block_size = end - start + 1;
        let lines_to_check = (search_block_size as i64 - 2).min(actual_block_size as i64 - 2);
        if lines_to_check <= 0 {
            return 1.0;
        }
        let lines_to_check = lines_to_check as usize;
        let upper = (search_block_size - 1).min(actual_block_size - 1);
        let mut similarity = 0.0f64;
        for j in 1..upper {
            let original_line = original_lines[start + j].trim();
            let search_line = search_lines[j].trim();
            let max_len = original_line.len().max(search_line.len());
            if max_len == 0 {
                continue;
            }
            let distance = levenshtein(original_line, search_line);
            similarity += (1.0 - distance as f64 / max_len as f64) / lines_to_check as f64;
            if early_exit && similarity >= SINGLE_CANDIDATE_SIMILARITY_THRESHOLD {
                break;
            }
        }
        similarity
    };

    if candidates.len() == 1 {
        let (start, end) = candidates[0];
        let similarity = middle_similarity(start, end, true);
        if similarity >= SINGLE_CANDIDATE_SIMILARITY_THRESHOLD {
            return vec![extract_lines(&original_lines, content, start, end)];
        }
        return vec![];
    }

    let mut best: Option<(usize, usize)> = None;
    let mut max_similarity = -1.0f64;
    for &(start, end) in &candidates {
        let actual_block_size = end - start + 1;
        let lines_to_check = (search_block_size as i64 - 2).min(actual_block_size as i64 - 2);
        let similarity = if lines_to_check <= 0 {
            1.0
        } else {
            let upper = (search_block_size - 1).min(actual_block_size - 1);
            let mut sum = 0.0;
            for j in 1..upper {
                let original_line = original_lines[start + j].trim();
                let search_line = search_lines[j].trim();
                let max_len = original_line.len().max(search_line.len());
                if max_len == 0 {
                    continue;
                }
                sum += 1.0 - levenshtein(original_line, search_line) as f64 / max_len as f64;
            }
            sum / lines_to_check as f64
        };
        if similarity > max_similarity {
            max_similarity = similarity;
            best = Some((start, end));
        }
    }
    if max_similarity >= MULTIPLE_CANDIDATES_SIMILARITY_THRESHOLD {
        if let Some((start, end)) = best {
            return vec![extract_lines(&original_lines, content, start, end)];
        }
    }
    vec![]
}

fn normalize_whitespace(text: &str) -> String {
    let mut out = String::new();
    let mut in_ws = false;
    for c in text.chars() {
        if c.is_whitespace() {
            if !in_ws {
                out.push(' ');
                in_ws = true;
            }
        } else {
            out.push(c);
            in_ws = false;
        }
    }
    out.trim().to_string()
}

/// Finds `words` in `line` separated by runs of whitespace, returning the exact
/// (non-normalized) substring — a hand-rolled stand-in for the TS version's regex built
/// from escaped words joined by `\s+`.
fn find_flexible_whitespace_match(line: &str, words: &[&str]) -> Option<String> {
    if words.is_empty() {
        return None;
    }
    let chars: Vec<char> = line.chars().collect();
    let n = chars.len();
    'outer: for start in 0..n {
        let mut pos = start;
        for (wi, word) in words.iter().enumerate() {
            let wchars: Vec<char> = word.chars().collect();
            if pos + wchars.len() > n || chars[pos..pos + wchars.len()] != wchars[..] {
                continue 'outer;
            }
            pos += wchars.len();
            if wi < words.len() - 1 {
                let ws_start = pos;
                while pos < n && chars[pos].is_whitespace() {
                    pos += 1;
                }
                if pos == ws_start {
                    continue 'outer;
                }
            }
        }
        return Some(chars[start..pos].iter().collect());
    }
    None
}

fn whitespace_normalized(content: &str, find: &str) -> Vec<String> {
    let mut out = Vec::new();
    let normalized_find = normalize_whitespace(find);

    let lines: Vec<&str> = content.split('\n').collect();
    for line in &lines {
        if normalize_whitespace(line) == normalized_find {
            out.push(line.to_string());
        } else if normalize_whitespace(line).contains(&normalized_find) {
            let words: Vec<&str> = find.split_whitespace().collect();
            if let Some(m) = find_flexible_whitespace_match(line, &words) {
                out.push(m);
            }
        }
    }

    let find_lines: Vec<&str> = find.split('\n').collect();
    if find_lines.len() > 1 && lines.len() >= find_lines.len() {
        for i in 0..=(lines.len() - find_lines.len()) {
            let block = lines[i..i + find_lines.len()].join("\n");
            if normalize_whitespace(&block) == normalized_find {
                out.push(block);
            }
        }
    }
    out
}

fn remove_indentation(text: &str) -> String {
    let lines: Vec<&str> = text.split('\n').collect();
    let non_empty: Vec<&&str> = lines.iter().filter(|l| !l.trim().is_empty()).collect();
    if non_empty.is_empty() {
        return text.to_string();
    }
    let min_indent = non_empty
        .iter()
        .map(|l| l.len() - l.trim_start().len())
        .min()
        .unwrap_or(0);
    lines
        .iter()
        .map(|l| {
            if l.trim().is_empty() {
                l.to_string()
            } else {
                l[min_indent.min(l.len())..].to_string()
            }
        })
        .collect::<Vec<_>>()
        .join("\n")
}

fn indentation_flexible(content: &str, find: &str) -> Vec<String> {
    let normalized_find = remove_indentation(find);
    let content_lines: Vec<&str> = content.split('\n').collect();
    let find_lines: Vec<&str> = find.split('\n').collect();
    let mut out = Vec::new();
    if content_lines.len() < find_lines.len() {
        return out;
    }
    for i in 0..=(content_lines.len() - find_lines.len()) {
        let block = content_lines[i..i + find_lines.len()].join("\n");
        if remove_indentation(&block) == normalized_find {
            out.push(block);
        }
    }
    out
}

fn unescape_string(s: &str) -> String {
    let mut out = String::new();
    let mut chars = s.chars().peekable();
    while let Some(c) = chars.next() {
        if c == '\\' {
            match chars.peek().copied() {
                Some('n') => {
                    out.push('\n');
                    chars.next();
                }
                Some('t') => {
                    out.push('\t');
                    chars.next();
                }
                Some('r') => {
                    out.push('\r');
                    chars.next();
                }
                Some('\'') => {
                    out.push('\'');
                    chars.next();
                }
                Some('"') => {
                    out.push('"');
                    chars.next();
                }
                Some('`') => {
                    out.push('`');
                    chars.next();
                }
                Some('\\') => {
                    out.push('\\');
                    chars.next();
                }
                Some('\n') => {
                    out.push('\n');
                    chars.next();
                }
                Some('$') => {
                    out.push('$');
                    chars.next();
                }
                _ => out.push(c),
            }
        } else {
            out.push(c);
        }
    }
    out
}

fn escape_normalized(content: &str, find: &str) -> Vec<String> {
    let mut out = Vec::new();
    let unescaped_find = unescape_string(find);
    if content.contains(&unescaped_find) {
        out.push(unescaped_find.clone());
    }
    let lines: Vec<&str> = content.split('\n').collect();
    let find_lines: Vec<&str> = unescaped_find.split('\n').collect();
    if lines.len() >= find_lines.len() {
        for i in 0..=(lines.len() - find_lines.len()) {
            let block = lines[i..i + find_lines.len()].join("\n");
            if unescape_string(&block) == unescaped_find {
                out.push(block);
            }
        }
    }
    out
}

fn multi_occurrence(content: &str, find: &str) -> Vec<String> {
    let mut out = Vec::new();
    if find.is_empty() {
        return out;
    }
    let mut start = 0usize;
    while let Some(idx) = content[start..].find(find) {
        out.push(find.to_string());
        start += idx + find.len();
    }
    out
}

fn trimmed_boundary(content: &str, find: &str) -> Vec<String> {
    let trimmed_find = find.trim();
    let mut out = Vec::new();
    if trimmed_find == find {
        return out;
    }
    if content.contains(trimmed_find) {
        out.push(trimmed_find.to_string());
    }
    let lines: Vec<&str> = content.split('\n').collect();
    let find_lines: Vec<&str> = find.split('\n').collect();
    if lines.len() >= find_lines.len() {
        for i in 0..=(lines.len() - find_lines.len()) {
            let block = lines[i..i + find_lines.len()].join("\n");
            if block.trim() == trimmed_find {
                out.push(block);
            }
        }
    }
    out
}

fn context_aware(content: &str, find: &str) -> Vec<String> {
    let mut find_lines: Vec<&str> = find.split('\n').collect();
    if find_lines.len() < 3 {
        return vec![];
    }
    if find_lines.last() == Some(&"") {
        find_lines.pop();
    }
    let content_lines: Vec<&str> = content.split('\n').collect();
    let first_line = find_lines[0].trim();
    let last_line = find_lines[find_lines.len() - 1].trim();
    let mut out = Vec::new();

    for i in 0..content_lines.len() {
        if content_lines[i].trim() != first_line {
            continue;
        }
        let mut j = i + 2;
        while j < content_lines.len() {
            if content_lines[j].trim() == last_line {
                let block_lines = &content_lines[i..=j];
                if block_lines.len() == find_lines.len() {
                    let mut matching = 0usize;
                    let mut total_non_empty = 0usize;
                    for k in 1..block_lines.len() - 1 {
                        let block_line = block_lines[k].trim();
                        let find_line = find_lines[k].trim();
                        if !block_line.is_empty() || !find_line.is_empty() {
                            total_non_empty += 1;
                            if block_line == find_line {
                                matching += 1;
                            }
                        }
                    }
                    if total_non_empty == 0 || (matching as f64 / total_non_empty as f64) >= 0.5 {
                        out.push(block_lines.join("\n"));
                    }
                }
                break;
            }
            j += 1;
        }
    }
    out
}

type ReplacerFn = fn(&str, &str) -> Vec<String>;

const REPLACER_CHAIN: &[ReplacerFn] = &[
    simple,
    line_trimmed,
    block_anchor,
    whitespace_normalized,
    indentation_flexible,
    escape_normalized,
    trimmed_boundary,
    context_aware,
    multi_occurrence,
];

fn is_disproportionate_match(search: &str, old_string: &str) -> bool {
    let old_lines = old_string.split('\n').count() as i64;
    let search_lines = search.split('\n').count() as i64;
    if search_lines >= (old_lines + 3).max(old_lines * 2) {
        return true;
    }
    if old_lines == 1 {
        return false;
    }
    let old_trimmed_len = old_string.trim().len() as i64;
    search.trim().len() as i64 > (old_trimmed_len + 500).max(old_trimmed_len * 4)
}

/// Applies the replacer chain in order, returning the new content on the first
/// unambiguous match. Mirrors opencode's `replace()` exactly, including the
/// disproportionate-match guard short-circuiting immediately (not skipping ahead).
pub fn replace(
    content: &str,
    old_string: &str,
    new_string: &str,
    replace_all: bool,
) -> Result<String, EditError> {
    if old_string == new_string {
        return Err(EditError::Identical);
    }
    if old_string.is_empty() {
        return Err(EditError::EmptyOldString);
    }

    let mut not_found = true;
    for replacer in REPLACER_CHAIN {
        for search in replacer(content, old_string) {
            let index = match content.find(&search) {
                Some(i) => i,
                None => continue,
            };
            not_found = false;
            if is_disproportionate_match(&search, old_string) {
                return Err(EditError::DisproportionateMatch);
            }
            if replace_all {
                return Ok(content.replace(&search, new_string));
            }
            let last_index = content.rfind(&search).unwrap();
            if index != last_index {
                continue;
            }
            return Ok(format!(
                "{}{}{}",
                &content[..index],
                new_string,
                &content[index + search.len()..]
            ));
        }
    }

    if not_found {
        Err(EditError::NotFound)
    } else {
        Err(EditError::MultipleMatches)
    }
}

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

    #[test]
    fn replaces_exact_text_in_existing_content() {
        let result = replace("old content here", "old content", "new content", false).unwrap();
        assert_eq!(result, "new content here");
    }

    #[test]
    fn rejects_identical_old_and_new() {
        assert_eq!(
            replace("content", "same", "same", false),
            Err(EditError::Identical)
        );
    }

    #[test]
    fn rejects_empty_old_string() {
        assert_eq!(
            replace("content", "", "x", false),
            Err(EditError::EmptyOldString)
        );
    }

    #[test]
    fn reports_not_found() {
        let err = replace("actual content", "not in file", "replacement", false).unwrap_err();
        assert_eq!(err, EditError::NotFound);
    }

    #[test]
    fn rejects_loose_block_anchor_matches() {
        let original = [
            "function configure() {",
            "  keepImportantState()",
            "  removeAllUserData()",
            "  archiveBackups()",
            "  auditLog()",
            "}",
        ]
        .join("\n");
        let old = ["function configure() {", "  const enabled = true", "}"].join("\n");
        let new = ["function configure() {", "  const enabled = false", "}"].join("\n");
        assert_eq!(
            replace(&original, &old, &new, false),
            Err(EditError::NotFound)
        );
    }

    #[test]
    fn rejects_block_anchor_matches_with_unrelated_middle_content() {
        let original = ["function configure() {", "  removeAllUserData()", "}"].join("\n");
        let old = ["function configure() {", "  const enabled = true", "}"].join("\n");
        let new = ["function configure() {", "  const enabled = false", "}"].join("\n");
        assert_eq!(
            replace(&original, &old, &new, false),
            Err(EditError::NotFound)
        );
    }

    #[test]
    fn replaces_all_occurrences_with_replace_all() {
        let result = replace("foo bar foo baz foo", "foo", "qux", true).unwrap();
        assert_eq!(result, "qux bar qux baz qux");
    }

    #[test]
    fn rejects_ambiguous_single_occurrence_request() {
        let err = replace("foo bar foo", "foo", "qux", false).unwrap_err();
        assert_eq!(err, EditError::MultipleMatches);
    }

    #[test]
    fn handles_multiline_replacements() {
        let result = replace(
            "line1\nline2\nline3",
            "line2",
            "new line 2\nextra line",
            false,
        )
        .unwrap();
        assert_eq!(result, "line1\nnew line 2\nextra line\nline3");
    }

    #[test]
    fn line_trimmed_replacer_matches_despite_surrounding_whitespace_changes() {
        let content = "  line one  \n  line two  \n";
        let matches = line_trimmed(content, "line one\nline two");
        assert_eq!(matches, vec!["  line one  \n  line two  ".to_string()]);
    }

    #[test]
    fn whitespace_normalized_matches_reformatted_line() {
        let content = "const x   =    1;\n";
        let matches = whitespace_normalized(content, "const x = 1;");
        assert!(matches.iter().any(|m| m == "const x   =    1;"));
    }

    #[test]
    fn indentation_flexible_matches_reindented_block() {
        let content = "    if true {\n        foo()\n    }\n";
        let find = "if true {\n    foo()\n}";
        let matches = indentation_flexible(content, find);
        assert_eq!(matches.len(), 1);
    }

    #[test]
    fn escape_normalized_matches_escaped_newline_in_find() {
        let content = "line one\nline two\n";
        let matches = escape_normalized(content, "line one\\nline two");
        assert!(matches.contains(&"line one\nline two".to_string()));
    }

    #[test]
    fn trimmed_boundary_matches_padded_find() {
        let content = "exact text here\n";
        let matches = trimmed_boundary(content, "  exact text here  ");
        assert!(matches.contains(&"exact text here".to_string()));
    }

    #[test]
    fn disproportionate_match_flags_far_larger_line_count() {
        let old = "one line";
        let search = "a\nb\nc\nd"; // 4 lines >= max(1+3, 1*2) = 4
        assert!(is_disproportionate_match(search, old));
    }

    #[test]
    fn disproportionate_match_flags_far_larger_char_span_same_line_count() {
        let old = "short";
        let search = "x".repeat(600); // > max(5+500, 5*4) = 505, still 1 line like `old`
        assert!(!is_disproportionate_match(&search, old)); // old_lines == 1 short-circuits to false

        let old_multiline = "short\nline";
        let search_multiline = format!("{}\nline", "x".repeat(600));
        assert!(is_disproportionate_match(&search_multiline, old_multiline));
    }

    #[test]
    fn disproportionate_match_allows_reasonably_sized_matches() {
        assert!(!is_disproportionate_match(
            "old content",
            "old content here"
        ));
    }
}