contrib/merge-lists/src/main.rs

contrib: add a set of scripts to run pytype in Docker

Having a simple way to run pytype for developers can massively shorten
development cycle. Using the same Docker image and scripts that we use on our
CI guarantees that the result achieved locally will be very similar to (if not
the same as) the output of our CI runners.

Things to note: the Dockerfile needs to do a little dance around user
permissions inside /home/ci-runner/ because:

- on one hand, creating new files on the host (e.g. .pyi files inside .pytype/)
  should use host user's uid and gid

- on the other hand, when we run the image as uid:gid of host user, it needs to
  be able to read/execute files inside the image that are owned by ci-runner

Since local user's uid might be different from ci-runner's uid, we execute this
very broad chmod command inside /home/ci-runner/, but then run the image as the
host user's uid:gid.

There might be a better way to do this.

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
}