oplog - get a workdir tree in a safe manner

crates/gitbutler-oplog/src/oplog.rs

@@ -4,6 +4,7 @@ use gitbutler_branch::{Branch, VirtualBranchesHandle, VirtualBranchesState};
 use gitbutler_diff::{hunks_by_filepath, FileDiff};
 use gitbutler_project::Project;
 use gitbutler_repo::RepositoryExt;
+use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::path::Path;
 use std::str::{from_utf8, FromStr};
@@ -205,21 +206,15 @@ impl OplogExt for Project {
             }
 
             // Get tree id from cache or calculate it
-            let wd_tree_id = wd_trees_cache
-                .entry(commit_id)
-                .or_insert_with(|| tree_from_applied_vbranches(&repo, commit_id).unwrap());
-            let wd_tree = repo.find_tree(wd_tree_id.to_owned())?;
+            let wd_tree = get_workdir_tree(&mut wd_trees_cache, commit_id, &repo)?;
 
             let details = commit
                 .message()
                 .and_then(|msg| SnapshotDetails::from_str(msg).ok());
 
             if let Ok(parent) = commit.parent(0) {
                 // Get tree id from cache or calculate it
-                let parent_wd_tree_id = wd_trees_cache
-                    .entry(parent.id())
-                    .or_insert_with(|| tree_from_applied_vbranches(&repo, parent.id()).unwrap());
-                let parent_tree = repo.find_tree(parent_wd_tree_id.to_owned())?;
+                let parent_tree = get_workdir_tree(&mut wd_trees_cache, parent.id(), &repo)?;
 
                 let mut opts = DiffOptions::new();
                 opts.include_untracked(true);
@@ -317,6 +312,25 @@ impl OplogExt for Project {
         oplog_state.oplog_head()
     }
 }
+
+/// Get a tree of the working dir (applied branches merged)
+fn get_workdir_tree<'a>(
+    wd_trees_cache: &mut HashMap<git2::Oid, git2::Oid>,
+    commit_id: git2::Oid,
+    repo: &'a git2::Repository,
+) -> Result<git2::Tree<'a>, anyhow::Error> {
+    if let Entry::Vacant(e) = wd_trees_cache.entry(commit_id) {
+        if let Ok(wd_tree_id) = tree_from_applied_vbranches(repo, commit_id) {
+            e.insert(wd_tree_id);
+        }
+    }
+    let wd_tree_id = wd_trees_cache.get(&commit_id).ok_or(anyhow!(
+        "Could not get a tree of all applied virtual branches merged"
+    ))?;
+    let wd_tree = repo.find_tree(wd_tree_id.to_owned())?;
+    Ok(wd_tree)
+}
+
 fn prepare_snapshot(ctx: &Project, _shared_access: &WorktreeReadPermission) -> Result<git2::Oid> {
     let worktree_dir = ctx.path.as_path();
     let repo = git2::Repository::open(worktree_dir)?;