task_inventory.rs

  1//! Project-wide storage of the tasks available, capable of updating itself from the sources set.
  2
  3use std::{
  4    any::TypeId,
  5    path::{Path, PathBuf},
  6    sync::Arc,
  7};
  8
  9use collections::{HashMap, VecDeque};
 10use gpui::{AppContext, Context, Model, ModelContext, Subscription};
 11use itertools::Itertools;
 12use project_core::worktree::WorktreeId;
 13use task::{Task, TaskId, TaskSource};
 14use util::{post_inc, NumericPrefixWithSuffix};
 15
 16/// Inventory tracks available tasks for a given project.
 17pub struct Inventory {
 18    sources: Vec<SourceInInventory>,
 19    last_scheduled_tasks: VecDeque<TaskId>,
 20}
 21
 22struct SourceInInventory {
 23    source: Model<Box<dyn TaskSource>>,
 24    _subscription: Subscription,
 25    type_id: TypeId,
 26    kind: TaskSourceKind,
 27}
 28
 29/// Kind of a source the tasks are fetched from, used to display more source information in the UI.
 30#[derive(Debug, Clone, PartialEq, Eq)]
 31pub enum TaskSourceKind {
 32    /// bash-like commands spawned by users, not associated with any path
 33    UserInput,
 34    /// ~/.config/zed/task.json - like global files with task definitions, applicable to any path
 35    AbsPath(PathBuf),
 36    /// Worktree-specific task definitions, e.g. dynamic tasks from open worktree file, or tasks from the worktree's .zed/task.json
 37    Worktree { id: WorktreeId, abs_path: PathBuf },
 38}
 39
 40impl TaskSourceKind {
 41    fn abs_path(&self) -> Option<&Path> {
 42        match self {
 43            Self::AbsPath(abs_path) | Self::Worktree { abs_path, .. } => Some(abs_path),
 44            Self::UserInput => None,
 45        }
 46    }
 47
 48    fn worktree(&self) -> Option<WorktreeId> {
 49        match self {
 50            Self::Worktree { id, .. } => Some(*id),
 51            _ => None,
 52        }
 53    }
 54}
 55
 56impl Inventory {
 57    pub fn new(cx: &mut AppContext) -> Model<Self> {
 58        cx.new_model(|_| Self {
 59            sources: Vec::new(),
 60            last_scheduled_tasks: VecDeque::new(),
 61        })
 62    }
 63
 64    /// If the task with the same path was not added yet,
 65    /// registers a new tasks source to fetch for available tasks later.
 66    /// Unless a source is removed, ignores future additions for the same path.
 67    pub fn add_source(
 68        &mut self,
 69        kind: TaskSourceKind,
 70        create_source: impl FnOnce(&mut ModelContext<Self>) -> Model<Box<dyn TaskSource>>,
 71        cx: &mut ModelContext<Self>,
 72    ) {
 73        let abs_path = kind.abs_path();
 74        if abs_path.is_some() {
 75            if let Some(a) = self.sources.iter().find(|s| s.kind.abs_path() == abs_path) {
 76                log::debug!("Source for path {abs_path:?} already exists, not adding. Old kind: {OLD_KIND:?}, new kind: {kind:?}", OLD_KIND = a.kind);
 77                return;
 78            }
 79        }
 80
 81        let source = create_source(cx);
 82        let type_id = source.read(cx).type_id();
 83        let source = SourceInInventory {
 84            _subscription: cx.observe(&source, |_, _, cx| {
 85                cx.notify();
 86            }),
 87            source,
 88            type_id,
 89            kind,
 90        };
 91        self.sources.push(source);
 92        cx.notify();
 93    }
 94
 95    /// If present, removes the local static source entry that has the given path,
 96    /// making corresponding task definitions unavailable in the fetch results.
 97    ///
 98    /// Now, entry for this path can be re-added again.
 99    pub fn remove_local_static_source(&mut self, abs_path: &Path) {
100        self.sources.retain(|s| s.kind.abs_path() != Some(abs_path));
101    }
102
103    /// If present, removes the worktree source entry that has the given worktree id,
104    /// making corresponding task definitions unavailable in the fetch results.
105    ///
106    /// Now, entry for this path can be re-added again.
107    pub fn remove_worktree_sources(&mut self, worktree: WorktreeId) {
108        self.sources.retain(|s| s.kind.worktree() != Some(worktree));
109    }
110
111    pub fn source<T: TaskSource>(&self) -> Option<Model<Box<dyn TaskSource>>> {
112        let target_type_id = std::any::TypeId::of::<T>();
113        self.sources.iter().find_map(
114            |SourceInInventory {
115                 type_id, source, ..
116             }| {
117                if &target_type_id == type_id {
118                    Some(source.clone())
119                } else {
120                    None
121                }
122            },
123        )
124    }
125
126    /// Pulls its sources to list runanbles for the path given (up to the source to decide what to return for no path).
127    pub fn list_tasks(
128        &self,
129        path: Option<&Path>,
130        worktree: Option<WorktreeId>,
131        lru: bool,
132        cx: &mut AppContext,
133    ) -> Vec<(TaskSourceKind, Arc<dyn Task>)> {
134        let mut lru_score = 0_u32;
135        let tasks_by_usage = if lru {
136            self.last_scheduled_tasks
137                .iter()
138                .rev()
139                .fold(HashMap::default(), |mut tasks, id| {
140                    tasks.entry(id).or_insert_with(|| post_inc(&mut lru_score));
141                    tasks
142                })
143        } else {
144            HashMap::default()
145        };
146        let not_used_score = post_inc(&mut lru_score);
147        self.sources
148            .iter()
149            .filter(|source| {
150                let source_worktree = source.kind.worktree();
151                worktree.is_none() || source_worktree.is_none() || source_worktree == worktree
152            })
153            .flat_map(|source| {
154                source
155                    .source
156                    .update(cx, |source, cx| source.tasks_for_path(path, cx))
157                    .into_iter()
158                    .map(|task| (&source.kind, task))
159            })
160            .map(|task| {
161                let usages = if lru {
162                    tasks_by_usage
163                        .get(&task.1.id())
164                        .copied()
165                        .unwrap_or(not_used_score)
166                } else {
167                    not_used_score
168                };
169                (task, usages)
170            })
171            .sorted_unstable_by(
172                |((kind_a, task_a), usages_a), ((kind_b, task_b), usages_b)| {
173                    usages_a
174                        .cmp(usages_b)
175                        .then(
176                            kind_a
177                                .worktree()
178                                .is_none()
179                                .cmp(&kind_b.worktree().is_none()),
180                        )
181                        .then(kind_a.worktree().cmp(&kind_b.worktree()))
182                        .then(
183                            kind_a
184                                .abs_path()
185                                .is_none()
186                                .cmp(&kind_b.abs_path().is_none()),
187                        )
188                        .then(kind_a.abs_path().cmp(&kind_b.abs_path()))
189                        .then({
190                            NumericPrefixWithSuffix::from_numeric_prefixed_str(task_a.name())
191                                .cmp(&NumericPrefixWithSuffix::from_numeric_prefixed_str(
192                                    task_b.name(),
193                                ))
194                                .then(task_a.name().cmp(task_b.name()))
195                        })
196                },
197            )
198            .map(|((kind, task), _)| (kind.clone(), task))
199            .collect()
200    }
201
202    /// Returns the last scheduled task, if any of the sources contains one with the matching id.
203    pub fn last_scheduled_task(&self, cx: &mut AppContext) -> Option<Arc<dyn Task>> {
204        self.last_scheduled_tasks.back().and_then(|id| {
205            // TODO straighten the `Path` story to understand what has to be passed here: or it will break in the future.
206            self.list_tasks(None, None, false, cx)
207                .into_iter()
208                .find(|(_, task)| task.id() == id)
209                .map(|(_, task)| task)
210        })
211    }
212
213    /// Registers task "usage" as being scheduled – to be used for LRU sorting when listing all tasks.
214    pub fn task_scheduled(&mut self, id: TaskId) {
215        self.last_scheduled_tasks.push_back(id);
216        if self.last_scheduled_tasks.len() > 5_000 {
217            self.last_scheduled_tasks.pop_front();
218        }
219    }
220}
221
222#[cfg(any(test, feature = "test-support"))]
223pub mod test_inventory {
224    use std::{
225        path::{Path, PathBuf},
226        sync::Arc,
227    };
228
229    use gpui::{AppContext, Context as _, Model, ModelContext, TestAppContext};
230    use project_core::worktree::WorktreeId;
231    use task::{Task, TaskId, TaskSource};
232
233    use crate::Inventory;
234
235    use super::TaskSourceKind;
236
237    #[derive(Debug, Clone, PartialEq, Eq)]
238    pub struct TestTask {
239        pub id: task::TaskId,
240        pub name: String,
241    }
242
243    impl Task for TestTask {
244        fn id(&self) -> &TaskId {
245            &self.id
246        }
247
248        fn name(&self) -> &str {
249            &self.name
250        }
251
252        fn cwd(&self) -> Option<&Path> {
253            None
254        }
255
256        fn exec(&self, _cwd: Option<PathBuf>) -> Option<task::SpawnInTerminal> {
257            None
258        }
259    }
260
261    pub struct StaticTestSource {
262        pub tasks: Vec<TestTask>,
263    }
264
265    impl StaticTestSource {
266        pub fn new(
267            task_names: impl IntoIterator<Item = String>,
268            cx: &mut AppContext,
269        ) -> Model<Box<dyn TaskSource>> {
270            cx.new_model(|_| {
271                Box::new(Self {
272                    tasks: task_names
273                        .into_iter()
274                        .enumerate()
275                        .map(|(i, name)| TestTask {
276                            id: TaskId(format!("task_{i}_{name}")),
277                            name,
278                        })
279                        .collect(),
280                }) as Box<dyn TaskSource>
281            })
282        }
283    }
284
285    impl TaskSource for StaticTestSource {
286        fn tasks_for_path(
287            &mut self,
288            _path: Option<&Path>,
289            _cx: &mut ModelContext<Box<dyn TaskSource>>,
290        ) -> Vec<Arc<dyn Task>> {
291            self.tasks
292                .clone()
293                .into_iter()
294                .map(|task| Arc::new(task) as Arc<dyn Task>)
295                .collect()
296        }
297
298        fn as_any(&mut self) -> &mut dyn std::any::Any {
299            self
300        }
301    }
302
303    pub fn list_task_names(
304        inventory: &Model<Inventory>,
305        path: Option<&Path>,
306        worktree: Option<WorktreeId>,
307        lru: bool,
308        cx: &mut TestAppContext,
309    ) -> Vec<String> {
310        inventory.update(cx, |inventory, cx| {
311            inventory
312                .list_tasks(path, worktree, lru, cx)
313                .into_iter()
314                .map(|(_, task)| task.name().to_string())
315                .collect()
316        })
317    }
318
319    pub fn register_task_used(
320        inventory: &Model<Inventory>,
321        task_name: &str,
322        cx: &mut TestAppContext,
323    ) {
324        inventory.update(cx, |inventory, cx| {
325            let task = inventory
326                .list_tasks(None, None, false, cx)
327                .into_iter()
328                .find(|(_, task)| task.name() == task_name)
329                .unwrap_or_else(|| panic!("Failed to find task with name {task_name}"));
330            inventory.task_scheduled(task.1.id().clone());
331        });
332    }
333
334    pub fn list_tasks(
335        inventory: &Model<Inventory>,
336        path: Option<&Path>,
337        worktree: Option<WorktreeId>,
338        lru: bool,
339        cx: &mut TestAppContext,
340    ) -> Vec<(TaskSourceKind, String)> {
341        inventory.update(cx, |inventory, cx| {
342            inventory
343                .list_tasks(path, worktree, lru, cx)
344                .into_iter()
345                .map(|(source_kind, task)| (source_kind, task.name().to_string()))
346                .collect()
347        })
348    }
349}
350
351#[cfg(test)]
352mod tests {
353    use gpui::TestAppContext;
354
355    use super::test_inventory::*;
356    use super::*;
357
358    #[gpui::test]
359    fn test_task_list_sorting(cx: &mut TestAppContext) {
360        let inventory = cx.update(Inventory::new);
361        let initial_tasks = list_task_names(&inventory, None, None, true, cx);
362        assert!(
363            initial_tasks.is_empty(),
364            "No tasks expected for empty inventory, but got {initial_tasks:?}"
365        );
366        let initial_tasks = list_task_names(&inventory, None, None, false, cx);
367        assert!(
368            initial_tasks.is_empty(),
369            "No tasks expected for empty inventory, but got {initial_tasks:?}"
370        );
371
372        inventory.update(cx, |inventory, cx| {
373            inventory.add_source(
374                TaskSourceKind::UserInput,
375                |cx| StaticTestSource::new(vec!["3_task".to_string()], cx),
376                cx,
377            );
378        });
379        inventory.update(cx, |inventory, cx| {
380            inventory.add_source(
381                TaskSourceKind::UserInput,
382                |cx| {
383                    StaticTestSource::new(
384                        vec![
385                            "1_task".to_string(),
386                            "2_task".to_string(),
387                            "1_a_task".to_string(),
388                        ],
389                        cx,
390                    )
391                },
392                cx,
393            );
394        });
395
396        let expected_initial_state = [
397            "1_a_task".to_string(),
398            "1_task".to_string(),
399            "2_task".to_string(),
400            "3_task".to_string(),
401        ];
402        assert_eq!(
403            list_task_names(&inventory, None, None, false, cx),
404            &expected_initial_state,
405            "Task list without lru sorting, should be sorted alphanumerically"
406        );
407        assert_eq!(
408            list_task_names(&inventory, None, None, true, cx),
409            &expected_initial_state,
410            "Tasks with equal amount of usages should be sorted alphanumerically"
411        );
412
413        register_task_used(&inventory, "2_task", cx);
414        assert_eq!(
415            list_task_names(&inventory, None, None, false, cx),
416            &expected_initial_state,
417            "Task list without lru sorting, should be sorted alphanumerically"
418        );
419        assert_eq!(
420            list_task_names(&inventory, None, None, true, cx),
421            vec![
422                "2_task".to_string(),
423                "1_a_task".to_string(),
424                "1_task".to_string(),
425                "3_task".to_string()
426            ],
427        );
428
429        register_task_used(&inventory, "1_task", cx);
430        register_task_used(&inventory, "1_task", cx);
431        register_task_used(&inventory, "1_task", cx);
432        register_task_used(&inventory, "3_task", cx);
433        assert_eq!(
434            list_task_names(&inventory, None, None, false, cx),
435            &expected_initial_state,
436            "Task list without lru sorting, should be sorted alphanumerically"
437        );
438        assert_eq!(
439            list_task_names(&inventory, None, None, true, cx),
440            vec![
441                "3_task".to_string(),
442                "1_task".to_string(),
443                "2_task".to_string(),
444                "1_a_task".to_string(),
445            ],
446        );
447
448        inventory.update(cx, |inventory, cx| {
449            inventory.add_source(
450                TaskSourceKind::UserInput,
451                |cx| {
452                    StaticTestSource::new(vec!["10_hello".to_string(), "11_hello".to_string()], cx)
453                },
454                cx,
455            );
456        });
457        let expected_updated_state = [
458            "1_a_task".to_string(),
459            "1_task".to_string(),
460            "2_task".to_string(),
461            "3_task".to_string(),
462            "10_hello".to_string(),
463            "11_hello".to_string(),
464        ];
465        assert_eq!(
466            list_task_names(&inventory, None, None, false, cx),
467            &expected_updated_state,
468            "Task list without lru sorting, should be sorted alphanumerically"
469        );
470        assert_eq!(
471            list_task_names(&inventory, None, None, true, cx),
472            vec![
473                "3_task".to_string(),
474                "1_task".to_string(),
475                "2_task".to_string(),
476                "1_a_task".to_string(),
477                "10_hello".to_string(),
478                "11_hello".to_string(),
479            ],
480        );
481
482        register_task_used(&inventory, "11_hello", cx);
483        assert_eq!(
484            list_task_names(&inventory, None, None, false, cx),
485            &expected_updated_state,
486            "Task list without lru sorting, should be sorted alphanumerically"
487        );
488        assert_eq!(
489            list_task_names(&inventory, None, None, true, cx),
490            vec![
491                "11_hello".to_string(),
492                "3_task".to_string(),
493                "1_task".to_string(),
494                "2_task".to_string(),
495                "1_a_task".to_string(),
496                "10_hello".to_string(),
497            ],
498        );
499    }
500
501    #[gpui::test]
502    fn test_inventory_static_task_filters(cx: &mut TestAppContext) {
503        let inventory_with_statics = cx.update(Inventory::new);
504        let common_name = "common_task_name";
505        let path_1 = Path::new("path_1");
506        let path_2 = Path::new("path_2");
507        let worktree_1 = WorktreeId::from_usize(1);
508        let worktree_path_1 = Path::new("worktree_path_1");
509        let worktree_2 = WorktreeId::from_usize(2);
510        let worktree_path_2 = Path::new("worktree_path_2");
511        inventory_with_statics.update(cx, |inventory, cx| {
512            inventory.add_source(
513                TaskSourceKind::UserInput,
514                |cx| {
515                    StaticTestSource::new(
516                        vec!["user_input".to_string(), common_name.to_string()],
517                        cx,
518                    )
519                },
520                cx,
521            );
522            inventory.add_source(
523                TaskSourceKind::AbsPath(path_1.to_path_buf()),
524                |cx| {
525                    StaticTestSource::new(
526                        vec!["static_source_1".to_string(), common_name.to_string()],
527                        cx,
528                    )
529                },
530                cx,
531            );
532            inventory.add_source(
533                TaskSourceKind::AbsPath(path_2.to_path_buf()),
534                |cx| {
535                    StaticTestSource::new(
536                        vec!["static_source_2".to_string(), common_name.to_string()],
537                        cx,
538                    )
539                },
540                cx,
541            );
542            inventory.add_source(
543                TaskSourceKind::Worktree {
544                    id: worktree_1,
545                    abs_path: worktree_path_1.to_path_buf(),
546                },
547                |cx| {
548                    StaticTestSource::new(
549                        vec!["worktree_1".to_string(), common_name.to_string()],
550                        cx,
551                    )
552                },
553                cx,
554            );
555            inventory.add_source(
556                TaskSourceKind::Worktree {
557                    id: worktree_2,
558                    abs_path: worktree_path_2.to_path_buf(),
559                },
560                |cx| {
561                    StaticTestSource::new(
562                        vec!["worktree_2".to_string(), common_name.to_string()],
563                        cx,
564                    )
565                },
566                cx,
567            );
568        });
569
570        let worktree_independent_tasks = vec![
571            (
572                TaskSourceKind::AbsPath(path_1.to_path_buf()),
573                common_name.to_string(),
574            ),
575            (
576                TaskSourceKind::AbsPath(path_1.to_path_buf()),
577                "static_source_1".to_string(),
578            ),
579            (
580                TaskSourceKind::AbsPath(path_2.to_path_buf()),
581                common_name.to_string(),
582            ),
583            (
584                TaskSourceKind::AbsPath(path_2.to_path_buf()),
585                "static_source_2".to_string(),
586            ),
587            (TaskSourceKind::UserInput, common_name.to_string()),
588            (TaskSourceKind::UserInput, "user_input".to_string()),
589        ];
590        let worktree_1_tasks = vec![
591            (
592                TaskSourceKind::Worktree {
593                    id: worktree_1,
594                    abs_path: worktree_path_1.to_path_buf(),
595                },
596                common_name.to_string(),
597            ),
598            (
599                TaskSourceKind::Worktree {
600                    id: worktree_1,
601                    abs_path: worktree_path_1.to_path_buf(),
602                },
603                "worktree_1".to_string(),
604            ),
605        ];
606        let worktree_2_tasks = vec![
607            (
608                TaskSourceKind::Worktree {
609                    id: worktree_2,
610                    abs_path: worktree_path_2.to_path_buf(),
611                },
612                common_name.to_string(),
613            ),
614            (
615                TaskSourceKind::Worktree {
616                    id: worktree_2,
617                    abs_path: worktree_path_2.to_path_buf(),
618                },
619                "worktree_2".to_string(),
620            ),
621        ];
622
623        let all_tasks = worktree_1_tasks
624            .iter()
625            .chain(worktree_2_tasks.iter())
626            // worktree-less tasks come later in the list
627            .chain(worktree_independent_tasks.iter())
628            .cloned()
629            .collect::<Vec<_>>();
630
631        for path in [
632            None,
633            Some(path_1),
634            Some(path_2),
635            Some(worktree_path_1),
636            Some(worktree_path_2),
637        ] {
638            assert_eq!(
639                list_tasks(&inventory_with_statics, path, None, false, cx),
640                all_tasks,
641                "Path {path:?} choice should not adjust static runnables"
642            );
643            assert_eq!(
644                list_tasks(&inventory_with_statics, path, Some(worktree_1), false, cx),
645                worktree_1_tasks
646                    .iter()
647                    .chain(worktree_independent_tasks.iter())
648                    .cloned()
649                    .collect::<Vec<_>>(),
650                "Path {path:?} choice should not adjust static runnables for worktree_1"
651            );
652            assert_eq!(
653                list_tasks(&inventory_with_statics, path, Some(worktree_2), false, cx),
654                worktree_2_tasks
655                    .iter()
656                    .chain(worktree_independent_tasks.iter())
657                    .cloned()
658                    .collect::<Vec<_>>(),
659                "Path {path:?} choice should not adjust static runnables for worktree_2"
660            );
661        }
662    }
663}