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(crate) 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(test)]
223mod tests {
224    use std::path::PathBuf;
225
226    use gpui::TestAppContext;
227
228    use super::*;
229
230    #[gpui::test]
231    fn test_task_list_sorting(cx: &mut TestAppContext) {
232        let inventory = cx.update(Inventory::new);
233        let initial_tasks = list_task_names(&inventory, None, None, true, cx);
234        assert!(
235            initial_tasks.is_empty(),
236            "No tasks expected for empty inventory, but got {initial_tasks:?}"
237        );
238        let initial_tasks = list_task_names(&inventory, None, None, false, cx);
239        assert!(
240            initial_tasks.is_empty(),
241            "No tasks expected for empty inventory, but got {initial_tasks:?}"
242        );
243
244        inventory.update(cx, |inventory, cx| {
245            inventory.add_source(
246                TaskSourceKind::UserInput,
247                |cx| StaticTestSource::new(vec!["3_task".to_string()], cx),
248                cx,
249            );
250        });
251        inventory.update(cx, |inventory, cx| {
252            inventory.add_source(
253                TaskSourceKind::UserInput,
254                |cx| {
255                    StaticTestSource::new(
256                        vec![
257                            "1_task".to_string(),
258                            "2_task".to_string(),
259                            "1_a_task".to_string(),
260                        ],
261                        cx,
262                    )
263                },
264                cx,
265            );
266        });
267
268        let expected_initial_state = [
269            "1_a_task".to_string(),
270            "1_task".to_string(),
271            "2_task".to_string(),
272            "3_task".to_string(),
273        ];
274        assert_eq!(
275            list_task_names(&inventory, None, None, false, cx),
276            &expected_initial_state,
277            "Task list without lru sorting, should be sorted alphanumerically"
278        );
279        assert_eq!(
280            list_task_names(&inventory, None, None, true, cx),
281            &expected_initial_state,
282            "Tasks with equal amount of usages should be sorted alphanumerically"
283        );
284
285        register_task_used(&inventory, "2_task", cx);
286        assert_eq!(
287            list_task_names(&inventory, None, None, false, cx),
288            &expected_initial_state,
289            "Task list without lru sorting, should be sorted alphanumerically"
290        );
291        assert_eq!(
292            list_task_names(&inventory, None, None, true, cx),
293            vec![
294                "2_task".to_string(),
295                "1_a_task".to_string(),
296                "1_task".to_string(),
297                "3_task".to_string()
298            ],
299        );
300
301        register_task_used(&inventory, "1_task", cx);
302        register_task_used(&inventory, "1_task", cx);
303        register_task_used(&inventory, "1_task", cx);
304        register_task_used(&inventory, "3_task", cx);
305        assert_eq!(
306            list_task_names(&inventory, None, None, false, cx),
307            &expected_initial_state,
308            "Task list without lru sorting, should be sorted alphanumerically"
309        );
310        assert_eq!(
311            list_task_names(&inventory, None, None, true, cx),
312            vec![
313                "3_task".to_string(),
314                "1_task".to_string(),
315                "2_task".to_string(),
316                "1_a_task".to_string(),
317            ],
318        );
319
320        inventory.update(cx, |inventory, cx| {
321            inventory.add_source(
322                TaskSourceKind::UserInput,
323                |cx| {
324                    StaticTestSource::new(vec!["10_hello".to_string(), "11_hello".to_string()], cx)
325                },
326                cx,
327            );
328        });
329        let expected_updated_state = [
330            "1_a_task".to_string(),
331            "1_task".to_string(),
332            "2_task".to_string(),
333            "3_task".to_string(),
334            "10_hello".to_string(),
335            "11_hello".to_string(),
336        ];
337        assert_eq!(
338            list_task_names(&inventory, None, None, false, cx),
339            &expected_updated_state,
340            "Task list without lru sorting, should be sorted alphanumerically"
341        );
342        assert_eq!(
343            list_task_names(&inventory, None, None, true, cx),
344            vec![
345                "3_task".to_string(),
346                "1_task".to_string(),
347                "2_task".to_string(),
348                "1_a_task".to_string(),
349                "10_hello".to_string(),
350                "11_hello".to_string(),
351            ],
352        );
353
354        register_task_used(&inventory, "11_hello", cx);
355        assert_eq!(
356            list_task_names(&inventory, None, None, false, cx),
357            &expected_updated_state,
358            "Task list without lru sorting, should be sorted alphanumerically"
359        );
360        assert_eq!(
361            list_task_names(&inventory, None, None, true, cx),
362            vec![
363                "11_hello".to_string(),
364                "3_task".to_string(),
365                "1_task".to_string(),
366                "2_task".to_string(),
367                "1_a_task".to_string(),
368                "10_hello".to_string(),
369            ],
370        );
371    }
372
373    #[gpui::test]
374    fn test_inventory_static_task_filters(cx: &mut TestAppContext) {
375        let inventory_with_statics = cx.update(Inventory::new);
376        let common_name = "common_task_name";
377        let path_1 = Path::new("path_1");
378        let path_2 = Path::new("path_2");
379        let worktree_1 = WorktreeId::from_usize(1);
380        let worktree_path_1 = Path::new("worktree_path_1");
381        let worktree_2 = WorktreeId::from_usize(2);
382        let worktree_path_2 = Path::new("worktree_path_2");
383        inventory_with_statics.update(cx, |inventory, cx| {
384            inventory.add_source(
385                TaskSourceKind::UserInput,
386                |cx| {
387                    StaticTestSource::new(
388                        vec!["user_input".to_string(), common_name.to_string()],
389                        cx,
390                    )
391                },
392                cx,
393            );
394            inventory.add_source(
395                TaskSourceKind::AbsPath(path_1.to_path_buf()),
396                |cx| {
397                    StaticTestSource::new(
398                        vec!["static_source_1".to_string(), common_name.to_string()],
399                        cx,
400                    )
401                },
402                cx,
403            );
404            inventory.add_source(
405                TaskSourceKind::AbsPath(path_2.to_path_buf()),
406                |cx| {
407                    StaticTestSource::new(
408                        vec!["static_source_2".to_string(), common_name.to_string()],
409                        cx,
410                    )
411                },
412                cx,
413            );
414            inventory.add_source(
415                TaskSourceKind::Worktree {
416                    id: worktree_1,
417                    abs_path: worktree_path_1.to_path_buf(),
418                },
419                |cx| {
420                    StaticTestSource::new(
421                        vec!["worktree_1".to_string(), common_name.to_string()],
422                        cx,
423                    )
424                },
425                cx,
426            );
427            inventory.add_source(
428                TaskSourceKind::Worktree {
429                    id: worktree_2,
430                    abs_path: worktree_path_2.to_path_buf(),
431                },
432                |cx| {
433                    StaticTestSource::new(
434                        vec!["worktree_2".to_string(), common_name.to_string()],
435                        cx,
436                    )
437                },
438                cx,
439            );
440        });
441
442        let worktree_independent_tasks = vec![
443            (
444                TaskSourceKind::AbsPath(path_1.to_path_buf()),
445                common_name.to_string(),
446            ),
447            (
448                TaskSourceKind::AbsPath(path_1.to_path_buf()),
449                "static_source_1".to_string(),
450            ),
451            (
452                TaskSourceKind::AbsPath(path_2.to_path_buf()),
453                common_name.to_string(),
454            ),
455            (
456                TaskSourceKind::AbsPath(path_2.to_path_buf()),
457                "static_source_2".to_string(),
458            ),
459            (TaskSourceKind::UserInput, common_name.to_string()),
460            (TaskSourceKind::UserInput, "user_input".to_string()),
461        ];
462        let worktree_1_tasks = vec![
463            (
464                TaskSourceKind::Worktree {
465                    id: worktree_1,
466                    abs_path: worktree_path_1.to_path_buf(),
467                },
468                common_name.to_string(),
469            ),
470            (
471                TaskSourceKind::Worktree {
472                    id: worktree_1,
473                    abs_path: worktree_path_1.to_path_buf(),
474                },
475                "worktree_1".to_string(),
476            ),
477        ];
478        let worktree_2_tasks = vec![
479            (
480                TaskSourceKind::Worktree {
481                    id: worktree_2,
482                    abs_path: worktree_path_2.to_path_buf(),
483                },
484                common_name.to_string(),
485            ),
486            (
487                TaskSourceKind::Worktree {
488                    id: worktree_2,
489                    abs_path: worktree_path_2.to_path_buf(),
490                },
491                "worktree_2".to_string(),
492            ),
493        ];
494
495        let all_tasks = worktree_1_tasks
496            .iter()
497            .chain(worktree_2_tasks.iter())
498            // worktree-less tasks come later in the list
499            .chain(worktree_independent_tasks.iter())
500            .cloned()
501            .collect::<Vec<_>>();
502
503        for path in [
504            None,
505            Some(path_1),
506            Some(path_2),
507            Some(worktree_path_1),
508            Some(worktree_path_2),
509        ] {
510            assert_eq!(
511                list_tasks(&inventory_with_statics, path, None, false, cx),
512                all_tasks,
513                "Path {path:?} choice should not adjust static runnables"
514            );
515            assert_eq!(
516                list_tasks(&inventory_with_statics, path, Some(worktree_1), false, cx),
517                worktree_1_tasks
518                    .iter()
519                    .chain(worktree_independent_tasks.iter())
520                    .cloned()
521                    .collect::<Vec<_>>(),
522                "Path {path:?} choice should not adjust static runnables for worktree_1"
523            );
524            assert_eq!(
525                list_tasks(&inventory_with_statics, path, Some(worktree_2), false, cx),
526                worktree_2_tasks
527                    .iter()
528                    .chain(worktree_independent_tasks.iter())
529                    .cloned()
530                    .collect::<Vec<_>>(),
531                "Path {path:?} choice should not adjust static runnables for worktree_2"
532            );
533        }
534    }
535
536    #[derive(Debug, Clone, PartialEq, Eq)]
537    struct TestTask {
538        id: TaskId,
539        name: String,
540    }
541
542    impl Task for TestTask {
543        fn id(&self) -> &TaskId {
544            &self.id
545        }
546
547        fn name(&self) -> &str {
548            &self.name
549        }
550
551        fn cwd(&self) -> Option<&Path> {
552            None
553        }
554
555        fn exec(&self, _cwd: Option<PathBuf>) -> Option<task::SpawnInTerminal> {
556            None
557        }
558    }
559
560    struct StaticTestSource {
561        tasks: Vec<TestTask>,
562    }
563
564    impl StaticTestSource {
565        fn new(
566            task_names: impl IntoIterator<Item = String>,
567            cx: &mut AppContext,
568        ) -> Model<Box<dyn TaskSource>> {
569            cx.new_model(|_| {
570                Box::new(Self {
571                    tasks: task_names
572                        .into_iter()
573                        .enumerate()
574                        .map(|(i, name)| TestTask {
575                            id: TaskId(format!("task_{i}_{name}")),
576                            name,
577                        })
578                        .collect(),
579                }) as Box<dyn TaskSource>
580            })
581        }
582    }
583
584    impl TaskSource for StaticTestSource {
585        fn tasks_for_path(
586            &mut self,
587            // static task source does not depend on path input
588            _: Option<&Path>,
589            _cx: &mut ModelContext<Box<dyn TaskSource>>,
590        ) -> Vec<Arc<dyn Task>> {
591            self.tasks
592                .clone()
593                .into_iter()
594                .map(|task| Arc::new(task) as Arc<dyn Task>)
595                .collect()
596        }
597
598        fn as_any(&mut self) -> &mut dyn std::any::Any {
599            self
600        }
601    }
602
603    fn list_task_names(
604        inventory: &Model<Inventory>,
605        path: Option<&Path>,
606        worktree: Option<WorktreeId>,
607        lru: bool,
608        cx: &mut TestAppContext,
609    ) -> Vec<String> {
610        inventory.update(cx, |inventory, cx| {
611            inventory
612                .list_tasks(path, worktree, lru, cx)
613                .into_iter()
614                .map(|(_, task)| task.name().to_string())
615                .collect()
616        })
617    }
618
619    fn list_tasks(
620        inventory: &Model<Inventory>,
621        path: Option<&Path>,
622        worktree: Option<WorktreeId>,
623        lru: bool,
624        cx: &mut TestAppContext,
625    ) -> Vec<(TaskSourceKind, String)> {
626        inventory.update(cx, |inventory, cx| {
627            inventory
628                .list_tasks(path, worktree, lru, cx)
629                .into_iter()
630                .map(|(source_kind, task)| (source_kind, task.name().to_string()))
631                .collect()
632        })
633    }
634
635    fn register_task_used(inventory: &Model<Inventory>, task_name: &str, cx: &mut TestAppContext) {
636        inventory.update(cx, |inventory, cx| {
637            let (_, task) = inventory
638                .list_tasks(None, None, false, cx)
639                .into_iter()
640                .find(|(_, task)| task.name() == task_name)
641                .unwrap_or_else(|| panic!("Failed to find task with name {task_name}"));
642            inventory.task_scheduled(task.id().clone());
643        });
644    }
645}