tenant.openagents/omega

1// Goal:
2// - Build an index that finds files that are frequently edited in the same git commit
3// - Lookup by path and get a list of related files, sorted by most frequently edited together
4//
5// Path => Path => usize
6//
7// This is a symmetric relationship, so for a => (b, 1), also add b => (a, 1)
8
9use std::collections::HashMap;
10use std::env;
11use std::path::{Path, PathBuf};
12use std::process::ExitCode;
13
14use anyhow::{Context as _, Result, anyhow, bail};
15use util::command::new_command;
16
17pub struct GitLogIndex {
18    index: HashMap<PathBuf, HashMap<PathBuf, usize>>,
19}
20
21impl GitLogIndex {
22    pub fn new() -> Self {
23        Self {
24            index: HashMap::new(),
25        }
26    }
27
28    pub fn add_related(&mut self, path: PathBuf, related: PathBuf) {
29        let count = self
30            .index
31            .entry(path.clone())
32            .or_default()
33            .entry(related.clone())
34            .or_default();
35        *count += 1;
36
37        // add the reverse mapping
38        let reverse_count = self
39            .index
40            .entry(related)
41            .or_default()
42            .entry(path)
43            .or_default();
44        *reverse_count += 1;
45    }
46
47    pub fn get_related(&self, path: &Path, n: usize) -> Vec<PathBuf> {
48        self.get_related_with_counts(path, n)
49            .into_iter()
50            .map(|(path, _)| path)
51            .collect()
52    }
53
54    pub fn get_related_with_counts(&self, path: &Path, n: usize) -> Vec<(PathBuf, usize)> {
55        let Some(counts) = self.index.get(path) else {
56            return Vec::new();
57        };
58
59        let mut related: Vec<_> = counts.iter().collect();
60        related.sort_by(|(left_path, left_count), (right_path, right_count)| {
61            right_count
62                .cmp(left_count)
63                .then_with(|| left_path.cmp(right_path))
64        });
65        related
66            .into_iter()
67            .take(n)
68            .map(|(path, count)| (path.clone(), *count))
69            .collect()
70    }
71}
72
73impl Default for GitLogIndex {
74    fn default() -> Self {
75        Self::new()
76    }
77}
78
79pub async fn build_git_log_index(worktree_dir: &Path) -> Result<GitLogIndex> {
80    let mut index = GitLogIndex::new();
81
82    let output = new_command("git")
83        .arg("log")
84        .arg("-5000")
85        .arg("--pretty=tformat:@@COMMIT %H")
86        .arg("--name-only")
87        .current_dir(worktree_dir)
88        .output()
89        .await
90        .with_context(|| format!("failed to run git log in {}", worktree_dir.display()))?;
91
92    if !output.status.success() {
93        let stderr = String::from_utf8_lossy(&output.stderr);
94        bail!(
95            "git log failed in {} with status {}: {}",
96            worktree_dir.display(),
97            output.status,
98            stderr.trim()
99        );
100    }
101
102    let log = String::from_utf8(output.stdout).context("git log output was not valid UTF-8")?;
103    let parsed = parse_git_log(&log);
104    for files in parsed {
105        for i in 0..files.len() {
106            for j in (i + 1)..files.len() {
107                index.add_related(files[i].clone(), files[j].clone());
108            }
109        }
110    }
111
112    Ok(index)
113}
114
115#[allow(dead_code)]
116fn main() -> ExitCode {
117    match run() {
118        Ok(()) => ExitCode::SUCCESS,
119        Err(error) => {
120            eprintln!("{error:#}");
121            ExitCode::FAILURE
122        }
123    }
124}
125
126#[allow(dead_code)]
127fn run() -> Result<()> {
128    let mut arguments = env::args_os();
129    let program_name = arguments
130        .next()
131        .and_then(|path| PathBuf::from(path).file_name().map(|name| name.to_owned()))
132        .and_then(|name| name.into_string().ok())
133        .unwrap_or_else(|| "git_log_context".to_string());
134
135    let worktree_dir = arguments.next().ok_or_else(|| {
136        print_usage(&program_name);
137        anyhow!("missing worktree path")
138    })?;
139    let query_path = arguments.next().ok_or_else(|| {
140        print_usage(&program_name);
141        anyhow!("missing query path")
142    })?;
143    if arguments.next().is_some() {
144        print_usage(&program_name);
145        bail!("too many arguments");
146    }
147
148    let worktree_dir = PathBuf::from(worktree_dir);
149    let query_path = normalize_query_path(&worktree_dir, &PathBuf::from(query_path));
150    let index = futures::executor::block_on(build_git_log_index(&worktree_dir))?;
151
152    for (path, count) in index.get_related_with_counts(&query_path, 10) {
153        println!("{count}\t{}", path.display());
154    }
155
156    Ok(())
157}
158
159#[allow(dead_code)]
160fn print_usage(program_name: &str) {
161    eprintln!("Usage: {program_name} <worktree-path> <query-path>");
162}
163
164#[allow(dead_code)]
165fn normalize_query_path(worktree_dir: &Path, query_path: &Path) -> PathBuf {
166    if query_path.is_absolute() {
167        query_path
168            .strip_prefix(worktree_dir)
169            .unwrap_or(query_path)
170            .components()
171            .collect()
172    } else {
173        query_path.components().collect()
174    }
175}
176
177fn parse_git_log(log: &str) -> Vec<Vec<PathBuf>> {
178    let mut lines = log.lines().peekable();
179    let mut commits = Vec::new();
180
181    while let Some(line) = lines.next() {
182        if line.starts_with("@@COMMIT ") {
183            // skip blank line
184            lines.next();
185            let mut files = Vec::new();
186            while let Some(next) = lines.peek()
187                && !next.starts_with("@@COMMIT ")
188            {
189                let Some(next) = lines.next() else {
190                    break;
191                };
192                if !next.is_empty() {
193                    files.push(next.into());
194                }
195            }
196            commits.push(files);
197        }
198    }
199
200    commits
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use indoc::indoc;
207
208    #[test]
209    fn test_git_log_index() {
210        let mut index = GitLogIndex::new();
211        index.add_related(PathBuf::from("a"), PathBuf::from("b"));
212        index.add_related(PathBuf::from("a"), PathBuf::from("b"));
213        index.add_related(PathBuf::from("a"), PathBuf::from("c"));
214        index.add_related(PathBuf::from("b"), PathBuf::from("c"));
215
216        let related = index.get_related(&PathBuf::from("a"), 100);
217        assert_eq!(related, vec![PathBuf::from("b"), PathBuf::from("c")]);
218    }
219
220    #[test]
221    fn test_parse_git_log() {
222        let log = indoc! {"
223            @@COMMIT d2e451dd48be67ef8c943e90dabc02e80a6984c9
224
225            crates/edit_prediction/src/edit_prediction.rs
226            crates/edit_prediction_cli/src/format_prompt.rs
227            crates/edit_prediction_cli/src/main.rs
228            crates/edit_prediction_cli/src/predict.rs
229            @@COMMIT d666823f348bd151067464fa31676d28bdb96717
230
231            crates/edit_prediction_cli/src/main.rs
232            crates/edit_prediction_cli/src/predict.rs
233            "};
234        let parsed = parse_git_log(log);
235
236        assert_eq!(parsed.len(), 2);
237        assert_eq!(
238            parsed[0][0],
239            PathBuf::from("crates/edit_prediction/src/edit_prediction.rs")
240        );
241        assert_eq!(
242            parsed[0][1],
243            PathBuf::from("crates/edit_prediction_cli/src/format_prompt.rs")
244        );
245
246        assert_eq!(
247            parsed[1][0],
248            PathBuf::from("crates/edit_prediction_cli/src/main.rs")
249        );
250        assert_eq!(
251            parsed[1][1],
252            PathBuf::from("crates/edit_prediction_cli/src/predict.rs")
253        );
254    }
255}
256