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