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