contrib/merge-lists/src/main.rs

compare-disco: support for `file` nodes specification

This leverage the `nodefromfile` feature in core. This make it possible for
callers to no longer pay the subset computation cost (and to make sure the
subset is the right one, even when the base repository is different)

use clap::{ArgGroup, Parser};
use itertools::Itertools;
use regex::bytes::Regex;
use similar::ChangeTag;
use std::cmp::{max, min, Ordering};
use std::collections::HashSet;
use std::ffi::OsString;
use std::ops::Range;
use std::path::PathBuf;

fn find_unchanged_ranges(
    old_bytes: &[u8],
    new_bytes: &[u8],
) -> Vec<(Range<usize>, Range<usize>)> {
    let diff = similar::TextDiff::configure()
        .algorithm(similar::Algorithm::Patience)
        .diff_lines(old_bytes, new_bytes);
    let mut new_unchanged_ranges = vec![];
    let mut old_index = 0;
    let mut new_index = 0;
    for diff in diff.iter_all_changes() {
        match diff.tag() {
            ChangeTag::Equal => {
                new_unchanged_ranges.push((
                    old_index..old_index + diff.value().len(),
                    new_index..new_index + diff.value().len(),
                ));
                old_index += diff.value().len();
                new_index += diff.value().len();
            }
            ChangeTag::Delete => {
                old_index += diff.value().len();
            }
            ChangeTag::Insert => {
                new_index += diff.value().len();
            }
        }
    }
    new_unchanged_ranges
}

/// Returns a list of all the lines in the input (including trailing newlines),
/// but only if they all match the regex and they are sorted.
fn get_lines<'input>(
    input: &'input [u8],
    regex: &Regex,
) -> Option<Vec<&'input [u8]>> {
    let lines = input.split_inclusive(|x| *x == b'\n').collect_vec();
    let mut previous_line = "".as_bytes();
    for line in &lines {
        if *line < previous_line {
            return None;
        }
        if !regex.is_match(line) {
            return None;
        }
        previous_line = line;
    }
    Some(lines)
}

fn resolve_conflict(
    base_slice: &[u8],
    local_slice: &[u8],
    other_slice: &[u8],
    regex: &Regex,
) -> Option<Vec<u8>> {
    let base_lines = get_lines(base_slice, regex)?;
    let local_lines = get_lines(local_slice, regex)?;
    let other_lines = get_lines(other_slice, regex)?;
    let base_lines_set: HashSet<_> = base_lines.iter().copied().collect();
    let local_lines_set: HashSet<_> = local_lines.iter().copied().collect();
    let other_lines_set: HashSet<_> = other_lines.iter().copied().collect();
    let mut result = local_lines_set;
    for to_add in other_lines_set.difference(&base_lines_set) {
        result.insert(to_add);
    }
    for to_remove in base_lines_set.difference(&other_lines_set) {
        result.remove(to_remove);
    }
    Some(result.into_iter().sorted().collect_vec().concat())
}

fn resolve(
    base_bytes: &[u8],
    local_bytes: &[u8],
    other_bytes: &[u8],
    regex: &Regex,
) -> (Vec<u8>, Vec<u8>, Vec<u8>) {
    // Find unchanged ranges between the base and the two sides. We do that by
    // initially considering the whole base unchanged. Then we compare each
    // side with the base and intersect the unchanged ranges we find with
    // what we had before.
    let unchanged_ranges = vec![UnchangedRange {
        base_range: 0..base_bytes.len(),
        offsets: vec![],
    }];
    let unchanged_ranges = intersect_regions(
        unchanged_ranges,
        &find_unchanged_ranges(base_bytes, local_bytes),
    );
    let mut unchanged_ranges = intersect_regions(
        unchanged_ranges,
        &find_unchanged_ranges(base_bytes, other_bytes),
    );
    // Add an empty UnchangedRange at the end to make it easier to find change
    // ranges. That way there's a changed range before each UnchangedRange.
    unchanged_ranges.push(UnchangedRange {
        base_range: base_bytes.len()..base_bytes.len(),
        offsets: vec![
            local_bytes.len().wrapping_sub(base_bytes.len()) as isize,
            other_bytes.len().wrapping_sub(base_bytes.len()) as isize,
        ],
    });

    let mut new_base_bytes: Vec<u8> = vec![];
    let mut new_local_bytes: Vec<u8> = vec![];
    let mut new_other_bytes: Vec<u8> = vec![];
    let mut previous = UnchangedRange {
        base_range: 0..0,
        offsets: vec![0, 0],
    };
    for current in unchanged_ranges {
        let base_slice =
            &base_bytes[previous.base_range.end..current.base_range.start];
        let local_slice = &local_bytes[previous.end(0)..current.start(0)];
        let other_slice = &other_bytes[previous.end(1)..current.start(1)];
        if let Some(resolution) =
            resolve_conflict(base_slice, local_slice, other_slice, regex)
        {
            new_base_bytes.extend(&resolution);
            new_local_bytes.extend(&resolution);
            new_other_bytes.extend(&resolution);
        } else {
            new_base_bytes.extend(base_slice);
            new_local_bytes.extend(local_slice);
            new_other_bytes.extend(other_slice);
        }
        new_base_bytes.extend(&base_bytes[current.base_range.clone()]);
        new_local_bytes.extend(&local_bytes[current.start(0)..current.end(0)]);
        new_other_bytes.extend(&other_bytes[current.start(1)..current.end(1)]);
        previous = current;
    }

    (new_base_bytes, new_local_bytes, new_other_bytes)
}

/// A tool that performs a 3-way merge, resolving conflicts in sorted lists and
/// leaving other conflicts unchanged. This is useful with Mercurial's support
/// for partial merge tools (configured in `[partial-merge-tools]`).
#[derive(Parser, Debug)]
#[clap(version, about, long_about = None)]
#[clap(group(ArgGroup::new("match").required(true).args(&["pattern", "python-imports"])))]
struct Args {
    /// Path to the file's content in the "local" side
    local: OsString,

    /// Path to the file's content in the base
    base: OsString,

    /// Path to the file's content in the "other" side
    other: OsString,

    /// Regular expression to use
    #[clap(long, short)]
    pattern: Option<String>,

    /// Use built-in regular expression for Python imports
    #[clap(long)]
    python_imports: bool,
}

fn get_regex(args: &Args) -> Regex {
    let pattern = if args.python_imports {
        r"import \w+(\.\w+)*( +#.*)?\n|from (\w+(\.\w+)* import \w+( as \w+)?(, \w+( as \w+)?)*( +#.*)?)"
    } else if let Some(pattern) = &args.pattern {
        pattern
    } else {
        ".*"
    };
    let pattern = format!(r"{}\r?\n?", pattern);
    regex::bytes::Regex::new(&pattern).unwrap()
}

fn main() {
    let args: Args = Args::parse();

    let base_path = PathBuf::from(&args.base);
    let local_path = PathBuf::from(&args.local);
    let other_path = PathBuf::from(&args.other);

    let base_bytes = std::fs::read(&base_path).unwrap();
    let local_bytes = std::fs::read(&local_path).unwrap();
    let other_bytes = std::fs::read(&other_path).unwrap();

    let regex = get_regex(&args);
    let (new_base_bytes, new_local_bytes, new_other_bytes) =
        resolve(&base_bytes, &local_bytes, &other_bytes, &regex);

    // Write out the result if anything changed
    if new_base_bytes != base_bytes {
        std::fs::write(&base_path, new_base_bytes).unwrap();
    }
    if new_local_bytes != local_bytes {
        std::fs::write(&local_path, new_local_bytes).unwrap();
    }
    if new_other_bytes != other_bytes {
        std::fs::write(&other_path, new_other_bytes).unwrap();
    }
}

fn checked_add(base: usize, offset: isize) -> usize {
    if offset < 0 {
        base.checked_sub(offset.checked_abs().unwrap() as usize)
            .unwrap()
    } else {
        base.checked_add(offset as usize).unwrap()
    }
}

// The remainder of the file is copied from
// https://github.com/martinvonz/jj/blob/main/lib/src/diff.rs

#[derive(Clone, PartialEq, Eq, Debug)]
struct UnchangedRange {
    base_range: Range<usize>,
    offsets: Vec<isize>,
}

impl UnchangedRange {
    fn start(&self, side: usize) -> usize {
        checked_add(self.base_range.start, self.offsets[side])
    }

    fn end(&self, side: usize) -> usize {
        checked_add(self.base_range.end, self.offsets[side])
    }
}

impl PartialOrd for UnchangedRange {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for UnchangedRange {
    fn cmp(&self, other: &Self) -> Ordering {
        self.base_range
            .start
            .cmp(&other.base_range.start)
            .then_with(|| self.base_range.end.cmp(&other.base_range.end))
    }
}

/// Takes the current regions and intersects it with the new unchanged ranges
/// from a 2-way diff. The result is a map of unchanged regions with one more
/// offset in the map's values.
fn intersect_regions(
    current_ranges: Vec<UnchangedRange>,
    new_unchanged_ranges: &[(Range<usize>, Range<usize>)],
) -> Vec<UnchangedRange> {
    let mut result = vec![];
    let mut current_ranges_iter = current_ranges.into_iter().peekable();
    for (new_base_range, other_range) in new_unchanged_ranges.iter() {
        assert_eq!(new_base_range.len(), other_range.len());
        while let Some(UnchangedRange {
            base_range,
            offsets,
        }) = current_ranges_iter.peek()
        {
            // No need to look further if we're past the new range.
            if base_range.start >= new_base_range.end {
                break;
            }
            // Discard any current unchanged regions that don't match between
            // the base and the new input.
            if base_range.end <= new_base_range.start {
                current_ranges_iter.next();
                continue;
            }
            let new_start = max(base_range.start, new_base_range.start);
            let new_end = min(base_range.end, new_base_range.end);
            let mut new_offsets = offsets.clone();
            new_offsets
                .push(other_range.start.wrapping_sub(new_base_range.start)
                    as isize);
            result.push(UnchangedRange {
                base_range: new_start..new_end,
                offsets: new_offsets,
            });
            if base_range.end >= new_base_range.end {
                // Break without consuming the item; there may be other new
                // ranges that overlap with it.
                break;
            }
            current_ranges_iter.next();
        }
    }
    result
}