id: make IdDb::new use non-legacy Context

Author: jonathantanmy2

crates/but/src/id.rs

@@ -1,63 +1,74 @@
 use std::{collections::HashMap, fmt::Display};
 
-use bstr::ByteSlice;
+use bstr::{BString, ByteSlice};
 use but_core::ref_metadata::StackId;
 use but_ctx::Context;
 use but_hunk_assignment::HunkAssignment;
 use gitbutler_command_context::CommandContext;
 
+fn branch_names(ctx: &Context) -> anyhow::Result<Vec<BString>> {
+    let guard = ctx.shared_worktree_access();
+    let meta = ctx.meta(guard.read_permission())?;
+    let head_info = but_workspace::head_info(&ctx.repo, &meta, Default::default())?;
+    let mut branch_names: Vec<BString> = Vec::new();
+    for stack in head_info.stacks {
+        for segment in stack.segments {
+            if let Some(ref_info) = segment.ref_info {
+                branch_names.push(ref_info.ref_name.shorten().to_owned());
+            }
+        }
+    }
+    Ok(branch_names)
+}
+
 pub struct IdDb {
     branch_name_to_cli_id: HashMap<String, CliId>,
     unassigned: CliId,
 }
 
 impl IdDb {
-    pub fn new(ctx: &CommandContext) -> anyhow::Result<Self> {
+    pub fn new(ctx: &Context) -> anyhow::Result<Self> {
         let mut max_zero_count = 1; // Ensure at least two "0" in ID.
-        let stacks = crate::utils::commits::stacks(ctx)?;
+        let branch_names = branch_names(ctx)?;
         let mut pairs_to_count: HashMap<u16, u8> = HashMap::new();
         fn u8_pair_to_u16(two: [u8; 2]) -> u16 {
             two[0] as u16 * 256 + two[1] as u16
         }
-        for stack in &stacks {
-            for head in &stack.heads {
-                for pair in head.name.windows(2) {
-                    let pair: [u8; 2] = pair.try_into()?;
-                    if !pair[0].is_ascii_alphanumeric() || !pair[1].is_ascii_alphanumeric() {
-                        continue;
-                    }
-                    let could_collide_with_commits =
-                        pair[0].is_ascii_hexdigit() && pair[1].is_ascii_hexdigit();
-                    if could_collide_with_commits {
-                        continue;
-                    }
-                    let u16pair = u8_pair_to_u16(pair);
-                    pairs_to_count
-                        .entry(u16pair)
-                        .and_modify(|count| *count = count.saturating_add(1))
-                        .or_insert(1);
+        for branch_name in &branch_names {
+            for pair in branch_name.windows(2) {
+                let pair: [u8; 2] = pair.try_into()?;
+                if !pair[0].is_ascii_alphanumeric() || !pair[1].is_ascii_alphanumeric() {
+                    continue;
                 }
-                for field in head.name.fields_with(|c| c != '0') {
-                    max_zero_count = std::cmp::max(field.len(), max_zero_count);
+                let could_collide_with_commits =
+                    pair[0].is_ascii_hexdigit() && pair[1].is_ascii_hexdigit();
+                if could_collide_with_commits {
+                    continue;
                 }
+                let u16pair = u8_pair_to_u16(pair);
+                pairs_to_count
+                    .entry(u16pair)
+                    .and_modify(|count| *count = count.saturating_add(1))
+                    .or_insert(1);
+            }
+            for field in branch_name.fields_with(|c| c != '0') {
+                max_zero_count = std::cmp::max(field.len(), max_zero_count);
             }
         }
 
         let mut branch_name_to_cli_id: HashMap<String, CliId> = HashMap::new();
-        for stack in stacks {
-            'head: for head in &stack.heads {
-                // Find first non-conflicting pair and use it as CliId.
-                for pair in head.name.windows(2) {
-                    let pair: [u8; 2] = pair.try_into()?;
-                    let u16pair = u8_pair_to_u16(pair);
-                    if let Some(1) = pairs_to_count.get(&u16pair) {
-                        let name = head.name.to_string();
-                        let id = str::from_utf8(&pair)
-                            .expect("if we stored it, it's ascii-alphanum")
-                            .to_owned();
-                        branch_name_to_cli_id.insert(name.clone(), CliId::Branch { name, id });
-                        continue 'head;
-                    }
+        'branch_name: for branch_name in &branch_names {
+            // Find first non-conflicting pair and use it as CliId.
+            for pair in branch_name.windows(2) {
+                let pair: [u8; 2] = pair.try_into()?;
+                let u16pair = u8_pair_to_u16(pair);
+                if let Some(1) = pairs_to_count.get(&u16pair) {
+                    let name = branch_name.to_string();
+                    let id = str::from_utf8(&pair)
+                        .expect("if we stored it, it's ascii-alphanum")
+                        .to_owned();
+                    branch_name_to_cli_id.insert(name.clone(), CliId::Branch { name, id });
+                    continue 'branch_name;
                 }
             }
         }
@@ -69,21 +80,15 @@ impl IdDb {
         })
     }
 
-    fn find_branches_by_name(
-        &mut self,
-        ctx: &CommandContext,
-        name: &str,
-    ) -> anyhow::Result<Vec<CliId>> {
-        let stacks = crate::utils::commits::stacks(ctx)?;
+    fn find_branches_by_name(&mut self, ctx: &Context, name: &str) -> anyhow::Result<Vec<CliId>> {
+        let branch_names = branch_names(ctx)?;
         let mut matches = Vec::new();
 
-        for stack in stacks {
-            for head in &stack.heads {
-                let branch_name = head.name.to_string();
-                // Exact match or partial match
-                if branch_name == name || branch_name.contains(name) {
-                    matches.push(self.branch(&branch_name).clone())
-                }
+        for branch_name in branch_names {
+            let branch_name = branch_name.to_string();
+            // Exact match or partial match
+            if branch_name == name || branch_name.contains(name) {
+                matches.push(self.branch(&branch_name).clone())
             }
         }
 
@@ -204,12 +209,12 @@ impl CliId {
         }
 
         // TODO: make callers of this function pass IdDb instead
-        let mut id_db = IdDb::new(&ctx.legacy_ctx()?)?;
+        let mut id_db = IdDb::new(ctx)?;
 
         let mut matches = Vec::new();
 
         // First, try exact branch name match
-        if let Ok(branch_matches) = id_db.find_branches_by_name(&ctx.legacy_ctx()?, s) {
+        if let Ok(branch_matches) = id_db.find_branches_by_name(ctx, s) {
             matches.extend(branch_matches);
         }
 
@@ -230,7 +235,7 @@ impl CliId {
                 .into_iter()
                 .filter(|id| id.matches_prefix(s))
                 .for_each(|id| cli_matches.push(id));
-            crate::status::all_branches(&ctx.legacy_ctx()?)?
+            crate::status::all_branches(ctx)?
                 .into_iter()
                 .filter(|id| id.matches_prefix(s))
                 .for_each(|id| cli_matches.push(id));
@@ -253,7 +258,7 @@ impl CliId {
                 .into_iter()
                 .filter(|id| id.matches(s))
                 .for_each(|id| cli_matches.push(id));
-            crate::status::all_branches(&ctx.legacy_ctx()?)?
+            crate::status::all_branches(ctx)?
                 .into_iter()
                 .filter(|id| id.matches(s))
                 .for_each(|id| cli_matches.push(id));

crates/but/src/status/mod.rs

@@ -104,7 +104,7 @@ pub(crate) async fn worktree(
 
     let unassigned = assignment::filter_by_stack_id(assignments_by_file.values(), &None);
     stack_details.push((None, (None, unassigned)));
-    let mut id_db = IdDb::new(&legacy_ctx)?;
+    let mut id_db = IdDb::new(ctx)?;
 
     // For JSON output, we'll need the original StackDetails to avoid redundant conversions
     let mut original_stack_details: Vec<(Option<gitbutler_stack::StackId>, Option<StackDetails>)> =
@@ -236,7 +236,7 @@ pub(crate) async fn worktree(
             verbose,
             review,
             &mut stack_mark,
-            &mut legacy_ctx,
+            ctx,
             i == stack_details_len - 1,
             i == 0,
             &review_map,
@@ -373,7 +373,7 @@ pub fn print_group(
     verbose: bool,
     show_url: bool,
     stack_mark: &mut Option<ColoredString>,
-    ctx: &mut CommandContext,
+    ctx: &Context,
     _last: bool,
     first: bool,
     review_map: &std::collections::HashMap<String, Vec<but_forge::ForgeReview>>,
@@ -436,7 +436,7 @@ pub fn print_group(
                 let dot = "●".yellow();
                 print_commit(
                     commit.id,
-                    created_at_of_commit(ctx, commit.id)?,
+                    created_at_of_commit(&ctx.legacy_ctx()?, commit.id)?,
                     commit.message.to_string(),
                     commit.author.name.clone(),
                     dot,
@@ -453,7 +453,8 @@ pub fn print_group(
             for cli_commit in &branch.commits {
                 let commit = &cli_commit;
                 let marked =
-                    crate::mark::commit_marked(ctx, commit.id.to_string()).unwrap_or_default();
+                    crate::mark::commit_marked(&mut ctx.legacy_ctx()?, commit.id.to_string())
+                        .unwrap_or_default();
                 let dot = match commit.state {
                     but_workspace::ui::CommitState::LocalOnly => "●".normal(),
                     but_workspace::ui::CommitState::LocalAndRemote(object_id) => {
@@ -467,7 +468,7 @@ pub fn print_group(
                 };
                 print_commit(
                     commit.id,
-                    created_at_of_commit(ctx, commit.id)?,
+                    created_at_of_commit(&ctx.legacy_ctx()?, commit.id)?,
                     commit.message.to_string(),
                     commit.author.name.clone(),
                     dot,
@@ -543,9 +544,9 @@ pub(crate) fn all_files(ctx: &mut CommandContext) -> anyhow::Result<Vec<CliId>>
     Ok(out)
 }
 
-pub(crate) fn all_branches(ctx: &CommandContext) -> anyhow::Result<Vec<CliId>> {
+pub(crate) fn all_branches(ctx: &Context) -> anyhow::Result<Vec<CliId>> {
     let mut id_db = IdDb::new(ctx)?;
-    let stacks = crate::utils::commits::stacks(ctx)?;
+    let stacks = crate::utils::commits::stacks(&ctx.legacy_ctx()?)?;
     let mut branches = Vec::new();
     for stack in stacks {
         for head in stack.heads {