1780 lines · 58.3 KB · Rust
RawDownload
//! `ep split-commit` implementation.//!//! This command generates a single evaluation example JSON object from a//! chronologically-ordered unified diff (a "commit").//!//! TODO: Port Python code to generate chronologically-ordered commitsuse crate::FailedHandling;use crate::reorder_patch::{Patch, PatchLine, extract_edits, locate_edited_line};use crate::word_diff::tokenize; /// Find the largest valid UTF-8 char boundary at or before `index` in `s`.fn floor_char_boundary(s: &str, index: usize) -> usize {    if index >= s.len() {        s.len()    } else if s.is_char_boundary(index) {        index    } else {        // Find the nearest valid character boundary at or before index        (0..index)            .rev()            .find(|&i| s.is_char_boundary(i))            .unwrap_or(0)    }}use anyhow::{Context as _, Result};use clap::Args;use edit_prediction::example_spec::ExampleSpec;use rand::Rng;use rand::SeedableRng;use serde::{Deserialize, Serialize};use similar::{DiffTag, TextDiff};use std::collections::BTreeSet;use std::fs;use std::io::{self, Write};use std::path::Path;use std::path::PathBuf; /// `ep split-commit` CLI args.#[derive(Debug, Args, Clone)]pub struct SplitCommitArgs {    /// Split point (float 0.0-1.0 for fraction, or integer for index)    #[arg(long, short = 's')]    pub split_point: Option<String>,     /// Random seed for reproducibility    #[arg(long)]    pub seed: Option<u64>,     /// Pretty-print JSON output    #[arg(long, short = 'p')]    pub pretty: bool,     /// Number of samples to generate per commit (samples random split points)    #[arg(long, short = 'n')]    pub num_samples: Option<usize>,} /// Input format for annotated commits (JSON Lines).#[derive(Debug, Clone, Deserialize)]#[allow(dead_code)]pub struct AnnotatedCommit {    /// Repository path (e.g., "repos/zed")    pub repo: String,    /// Repository URL (e.g., "https://github.com/zed-industries/zed")    pub repo_url: String,    /// Commit SHA    pub commit_sha: String,    /// Chronologically reordered commit diff    pub reordered_commit: String,    /// Original commit diff    pub original_commit: String,    /// Whether diff stats match between original and reordered    pub diff_stats_match: bool,} /// Cursor position in a file.#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]pub struct CursorPosition {    pub file: String,    pub line: usize,    pub column: usize,} impl std::fmt::Display for CursorPosition {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        write!(f, "{}:{}:{}", self.file, self.line, self.column)    }} /// Represents a split commit with source and target patches.#[derive(Debug, Clone)]pub struct SplitCommit {    pub source_patch: String,    pub target_patch: String,} /// Split point specification for evaluation generation.#[derive(Debug, Clone)]pub enum SplitPoint {    /// Fraction of total edits (0.0 to 1.0)    Fraction(f64),    /// Absolute index    Index(usize),} fn parse_split_point(value: &str) -> Option<SplitPoint> {    if value.contains('.') {        value.parse::<f64>().ok().map(SplitPoint::Fraction)    } else {        value.parse::<usize>().ok().map(SplitPoint::Index)    }} /// Entry point for the `ep split-commit` subcommand.////// This runs synchronously and outputs JSON Lines (one output per input line).pub fn run_split_commit(    args: &SplitCommitArgs,    inputs: &[PathBuf],    output_path: Option<&PathBuf>,    failed: FailedHandling,) -> Result<()> {    use std::collections::HashSet;    use std::io::BufRead;     let stdin_path = PathBuf::from("-");    let inputs = if inputs.is_empty() {        std::slice::from_ref(&stdin_path)    } else {        inputs    };     let split_point = args.split_point.as_deref().and_then(parse_split_point);    let mut output_lines = Vec::new();     for input_path in inputs {        let input: Box<dyn BufRead> = if input_path.as_os_str() == "-" {            Box::new(io::BufReader::new(io::stdin()))        } else {            let file = fs::File::open(input_path)                .with_context(|| format!("failed to open input file {}", input_path.display()))?;            Box::new(io::BufReader::new(file))        };         for (line_num, line_result) in input.lines().enumerate() {            let line =                line_result.with_context(|| format!("failed to read line {}", line_num + 1))?;             if line.trim().is_empty() {                continue;            }             let annotated: AnnotatedCommit = serde_json::from_str(&line)                .with_context(|| format!("failed to parse JSON at line {}", line_num + 1))?;             // Generate multiple samples if num_samples is set            if let Some(num_samples) = args.num_samples {                let mut seen_samples: HashSet<String> = HashSet::new();                let base_seed = args.seed.unwrap_or_else(|| rand::random());                 for sample_idx in 0..num_samples {                    let sample_seed = base_seed.wrapping_add(sample_idx as u64);                     let case = match generate_evaluation_example_from_ordered_commit(                        &annotated.reordered_commit,                        &annotated.repo_url,                        &annotated.commit_sha,                        None, // Use random split point for multi-sample mode                        Some(sample_seed),                        Some(sample_idx),                    ) {                        Ok(case) => case,                        Err(e) => {                            let err_msg = format!(                                "failed to generate evaluation example for commit {} at line {} (sample {}): {}",                                annotated.commit_sha,                                line_num + 1,                                sample_idx,                                e                            );                            match failed {                                FailedHandling::Skip | FailedHandling::SkipNoFiles => {                                    eprintln!("{}", err_msg);                                    continue;                                }                                FailedHandling::Keep => {                                    anyhow::bail!(err_msg);                                }                            }                        }                    };                     let json = if args.pretty {                        serde_json::to_string_pretty(&case)                    } else {                        serde_json::to_string(&case)                    }                    .context("failed to serialize evaluation case as JSON")?;                     // Only add unique samples (different split points may produce same result)                    if seen_samples.insert(json.clone()) {                        output_lines.push(json);                    }                }            } else {                let case = match generate_evaluation_example_from_ordered_commit(                    &annotated.reordered_commit,                    &annotated.repo_url,                    &annotated.commit_sha,                    split_point.clone(),                    args.seed,                    None,                ) {                    Ok(case) => case,                    Err(e) => {                        let err_msg = format!(                            "failed to generate evaluation example for commit {} at line {}: {}",                            annotated.commit_sha,                            line_num + 1,                            e                        );                        match failed {                            FailedHandling::Skip | FailedHandling::SkipNoFiles => {                                eprintln!("{}", err_msg);                                continue;                            }                            FailedHandling::Keep => {                                anyhow::bail!(err_msg);                            }                        }                    }                };                 let json = if args.pretty {                    serde_json::to_string_pretty(&case)                } else {                    serde_json::to_string(&case)                }                .context("failed to serialize evaluation case as JSON")?;                 output_lines.push(json);            }        }    }     let output_content = output_lines.join("\n") + if output_lines.is_empty() { "" } else { "\n" };     if let Some(path) = output_path {        fs::write(path, &output_content)            .with_context(|| format!("failed to write output to {}", path.display()))?;    } else {        io::stdout()            .write_all(output_content.as_bytes())            .context("failed to write to stdout")?;    }     Ok(())} /// Main function to generate an evaluation example from an ordered commit.////// # Arguments/// * `commit` - Chronologically ordered unified diff of the commit/// * `repository_url` - URL of the repository/// * `commit_hash` - Hash of the commit/// * `split_point` - Point at which the commit will be split (None for random)/// * `seed` - Optional seed for randomness/// * `sample_num` - Optional sample number for generating unique namespub fn generate_evaluation_example_from_ordered_commit(    commit: &str,    repository_url: &str,    commit_hash: &str,    split_point: Option<SplitPoint>,    seed: Option<u64>,    sample_num: Option<usize>,) -> Result<ExampleSpec> {    let mut rng: Box<dyn rand::RngCore> = match seed {        Some(seed) => Box::new(rand::rngs::StdRng::seed_from_u64(seed)),        None => Box::new(rand::rngs::ThreadRng::default()),    };     // Parse and normalize the commit    let mut patch = Patch::parse_unified_diff(commit);     // Filter header to only keep lines starting with "//"    let header_lines: Vec<&str> = patch        .header        .lines()        .filter(|line| line.starts_with("//"))        .collect();    patch.header = if header_lines.is_empty() {        String::new()    } else {        header_lines.join("\n") + "\n"    };    let commit_normalized = patch.to_string();     // Compute the split point    let stats = patch.stats();    let num_edits = stats.added + stats.removed;     anyhow::ensure!(num_edits != 0, "no edits found in commit");     let split = match split_point {        None => rng.random_range(1..=num_edits),        Some(SplitPoint::Fraction(f)) => {            let v = (f * num_edits as f64).floor() as usize;            v.min(num_edits)        }        Some(SplitPoint::Index(i)) => i.min(num_edits),    };     // Split the commit into source and target patches    let (prefix, suffix) = split_ordered_commit(&commit_normalized, split);     let mut split_commit = SplitCommit {        source_patch: prefix,        target_patch: suffix,    };     // Imitate human edits    let human_edit_seed = rng.random_range(1..=10000u64);    let (src_patch, tgt_patch, cursor_opt) = imitate_human_edits(        &split_commit.source_patch,        &split_commit.target_patch,        human_edit_seed,    );    split_commit.source_patch = src_patch;    split_commit.target_patch = tgt_patch;     // Sample cursor position    let cursor = match cursor_opt {        Some(c) => c,        None => sample_cursor_position(&patch, &split_commit)            .context("failed to sample cursor position")?,    };     // Get cursor excerpt    let cursor_excerpt = get_cursor_excerpt(        &cursor,        &split_commit.source_patch,        &split_commit.target_patch,    )    .context("failed to generate cursor excerpt")?;     // Handle edge case where split_point == 0    if split == 0 {        split_commit.target_patch = String::new();    }     let repo_name = repository_url        .trim_end_matches('/')        .rsplit('/')        .next()        .unwrap_or("unknown");    let short_sha = &commit_hash[..commit_hash.len().min(8)];    let name = match sample_num {        Some(n) => format!("{}-{}-{}", repo_name, short_sha, n),        None => format!("{}-{}", repo_name, short_sha),    };     Ok(ExampleSpec {        name,        repository_url: repository_url.to_string(),        revision: format!("{}~1", commit_hash),        edit_history: split_commit.source_patch.clone(),        cursor_path: Path::new(&cursor.file).into(),        cursor_position: cursor_excerpt,        expected_patches: vec![split_commit.target_patch],        tags: vec![],        reasoning: None,        uncommitted_diff: String::new(),        rejected_patch: None,         telemetry: None,        human_feedback: Vec::new(),        rating: None,    })} /// Split an ordered commit into source and target commits.////// # Arguments/// * `commit` - Ordered commit string/// * `split_pos` - Position to split the commit (number of edited lines)////// # Returns/// A tuple of (source_diff, target_diff)pub fn split_ordered_commit(commit: &str, split_pos: usize) -> (String, String) {    let patch = Patch::parse_unified_diff(commit);    let source_edits: BTreeSet<usize> = (0..split_pos).collect();    let (source, target) = extract_edits(&patch, &source_edits);     let mut source_str = source.to_string();    let target_str = target.to_string();     // Strip last group header from the source (lines starting with "//" at the end)    let source_lines: Vec<&str> = source_str.lines().collect();    let mut end_idx = source_lines.len();    for i in (0..source_lines.len()).rev() {        if source_lines[i].starts_with("//") {            end_idx = i;        } else {            break;        }    }    if end_idx < source_lines.len() {        source_str = source_lines[..end_idx].join("\n");        if !source_str.is_empty() {            source_str.push('\n');        }    }     (source_str, target_str)} /// Calculate the weight for a split position based on the character at that position.////// Higher weights indicate more natural pause points (e.g., after punctuation,/// at identifier boundaries). Lower weights indicate less natural points/// (e.g., mid-identifier).fn position_weight(text: &str, pos: usize) -> u32 {    if pos == 0 || pos > text.len() {        return 1;    }     let chars: Vec<char> = text.chars().collect();    if pos > chars.len() {        return 1;    }     // Get the character just before this position (what we just "typed")    let prev_char = chars[pos - 1];     // High weight: natural pause points (end of statement/argument, opening brackets)    if matches!(prev_char, ',' | ';' | ':' | '(' | '[' | '{') {        return 10;    }     // High weight: closing brackets (finished a group)    if matches!(prev_char, ')' | ']' | '}') {        return 8;    }     // Medium weight: operators and method chains    if matches!(        prev_char,        '.' | '+' | '-' | '*' | '/' | '=' | '<' | '>' | '&' | '|' | '!'    ) {        return 5;    }     // Check if we're at the end of an identifier (word char followed by non-word char)    let is_prev_word_char = prev_char.is_alphanumeric() || prev_char == '_';    let is_next_word_char =        pos < chars.len() && (chars[pos].is_alphanumeric() || chars[pos] == '_');     if is_prev_word_char && !is_next_word_char {        // End of identifier - high weight        return 8;    }     // Whitespace is a natural pause    if prev_char.is_whitespace() {        return 6;    }     // Mid-identifier: low weight (rare autocomplete scenarios)    if is_prev_word_char && is_next_word_char {        return 1;    }     // Default medium-low weight    3} /// Select a weighted random index from a list of weights.////// Returns an index based on the weights, using the provided seed for/// deterministic selection.fn weighted_select(weights: &[u32], seed: u64) -> usize {    if weights.is_empty() {        return 0;    }     let total_weight: u64 = weights.iter().map(|&w| w as u64).sum();    if total_weight == 0 {        // Fallback to uniform selection if all weights are zero        return seed as usize % weights.len();    }     // Use seed to select a value in [0, total_weight)    let target = seed % total_weight;    let mut cumulative: u64 = 0;     for (idx, &weight) in weights.iter().enumerate() {        cumulative += weight as u64;        if target < cumulative {            return idx;        }

Lines 1–500 of 1780.