lib.rs

pub mod config;
pub mod display;
pub mod loader;
mod package_graph;
pub mod query;
mod specifier;

use std::{
    collections::{HashMap, hash_map::Entry},
    convert::Infallible,
    sync::Arc,
};

use config::{ResolvedTaskConfig, UserConfigFile};
use package_graph::IndexedPackageGraph;
use petgraph::{
    graph::{DefaultIx, DiGraph, EdgeIndex, IndexType, NodeIndex},
    visit::{Control, DfsEvent, depth_first_search},
};
use serde::Serialize;
pub use specifier::TaskSpecifier;
use vec1::smallvec_v1::SmallVec1;
use vite_path::AbsolutePath;
use vite_str::Str;
use vite_workspace::{PackageNodeIndex, WorkspaceRoot};

use crate::display::TaskDisplay;

#[derive(Debug, Clone, Copy, Serialize)]
enum TaskDependencyTypeInner {
    /// The dependency is explicitly declared by user in `dependsOn`.
    Explicit,
    /// The dependency is added due to topological ordering based on package dependencies.
    Topological,
    /// The dependency is explicitly declared by user in `dependsOn` and also added due to topological ordering.
    Both,
}

/// The type of a task dependency, explaining why it's introduced.
#[derive(Debug, Clone, Copy, Serialize)]
#[serde(transparent)]
pub struct TaskDependencyType(TaskDependencyTypeInner);

// It hides `TaskDependencyTypeInner` and only expose `is_explicit`/`is_topological`
// to avoid incorrectly matching only Explicit variant to check if it's explicit.
impl TaskDependencyType {
    pub fn is_explicit(self) -> bool {
        matches!(self.0, TaskDependencyTypeInner::Explicit | TaskDependencyTypeInner::Both)
    }

    pub fn is_topological(self) -> bool {
        matches!(self.0, TaskDependencyTypeInner::Topological | TaskDependencyTypeInner::Both)
    }
}

/// Uniquely identifies a task, by its name and the package where it's defined.
#[derive(Debug, PartialEq, Eq, Hash, Clone, PartialOrd, Ord)]
struct TaskId {
    /// The index of the package where the task is defined.
    pub package_index: PackageNodeIndex,

    /// The name of the script or the entry in `vite.config.*`.
    pub task_name: Str,
}

/// A node in the task graph, representing a task with its resolved configuration.
#[derive(Debug, Serialize)]
pub struct TaskNode {
    /// Printing the task in a human-readable way.
    pub task_display: TaskDisplay,

    /// The resolved configuration of this task.
    ///
    /// This contains information affecting how the task is spawn,
    /// whereas `task_id` is for looking up the task.
    ///
    /// However, it does not contain external factors like additional args from cli and env vars.
    pub resolved_config: ResolvedTaskConfig,
}

impl vite_graph_ser::GetKey for TaskNode {
    type Key<'a> = (&'a AbsolutePath, &'a str);

    fn key(&self) -> Result<Self::Key<'_>, String> {
        Ok((&self.task_display.package_path, &self.task_display.task_name))
    }
}

#[derive(Debug, thiserror::Error)]
pub enum TaskGraphLoadError {
    #[error("Failed to load package graph")]
    PackageGraphLoadError(#[from] vite_workspace::Error),

    #[error("Failed to load task config file for package at {package_path:?}")]
    ConfigLoadError {
        package_path: Arc<AbsolutePath>,
        #[source]
        error: anyhow::Error,
    },

    #[error("Failed to resolve task config for task {task_display}")]
    ResolveConfigError {
        task_display: TaskDisplay,
        #[source]
        error: crate::config::ResolveTaskConfigError,
    },

    #[error("Failed to lookup dependency '{specifier}' for task {task_display}")]
    DependencySpecifierLookupError {
        specifier: Str,
        task_display: TaskDisplay,
        #[source]
        error: SpecifierLookupError,
    },
}

/// Error when looking up a task by its specifier.
///
/// It's generic over `UnknownPackageError`, which is the error type when looking up a task without a package name and without a package origin.
///
/// - When the specifier is from `dependOn` of a known task, `UnknownPackageError` is `Infallible` because the origin package is always known.
/// - When the specifier is from a CLI command, `UnknownPackageError` can be a real error type in case cwd is not in any package.
#[derive(Debug, thiserror::Error, Serialize)]
pub enum SpecifierLookupError<PackageUnknownError = Infallible> {
    #[error("Package '{package_name}' is ambiguous among multiple packages: {package_paths:?}")]
    AmbiguousPackageName { package_name: Str, package_paths: Box<[Arc<AbsolutePath>]> },

    #[error("Package '{package_name}' not found")]
    PackageNameNotFound { package_name: Str },

    #[error("Task '{task_name}' not found in package {package_name}")]
    TaskNameNotFound {
        package_name: Str,
        task_name: Str,
        #[serde(skip)]
        package_index: PackageNodeIndex,
    },

    #[error(
        "Nowhere to look for task '{task_name}' because the package is unknown: {unspecifier_package_error}"
    )]
    PackageUnknown { unspecifier_package_error: PackageUnknownError, task_name: Str },
}

/// newtype of `DefaultIx` for indices in task graphs
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
pub struct TaskIx(DefaultIx);
unsafe impl IndexType for TaskIx {
    fn new(x: usize) -> Self {
        Self(DefaultIx::new(x))
    }

    fn index(&self) -> usize {
        self.0.index()
    }

    fn max() -> Self {
        Self(<DefaultIx as IndexType>::max())
    }
}

pub type TaskNodeIndex = NodeIndex<TaskIx>;
pub type TaskEdgeIndex = EdgeIndex<TaskIx>;

/// Full task graph of a workspace, with necessary HashMaps for quick task lookup
///
/// It's immutable after created. The task nodes contain resolved task configurations and their dependencies.
/// External factors (e.g. additional args from cli, current working directory, environmental variables) are not stored here.
#[derive(Debug)]
pub struct IndexedTaskGraph {
    task_graph: DiGraph<TaskNode, TaskDependencyType, TaskIx>,

    /// Preserve the package graph for two purposes:
    /// - `self.task_graph` refers packages via PackageNodeIndex. To display package names and paths, we need to lookup them in package_graph.
    /// - To find nearest topological tasks when the starting package itself doesn't contain the task with the given name.
    indexed_package_graph: IndexedPackageGraph,

    /// task indices by task id for quick lookup
    node_indices_by_task_id: HashMap<TaskId, TaskNodeIndex>,
}

pub type TaskGraph = DiGraph<TaskNode, TaskDependencyType, TaskIx>;

impl IndexedTaskGraph {
    /// Load the task graph from a discovered workspace using the provided config loader.
    pub async fn load(
        workspace_root: &WorkspaceRoot,
        config_loader: &dyn loader::UserConfigLoader,
    ) -> Result<Self, TaskGraphLoadError> {
        let mut task_graph = DiGraph::<TaskNode, TaskDependencyType, TaskIx>::default();

        let package_graph = vite_workspace::load_package_graph(workspace_root)?;

        // Record dependency specifiers for each task node to add explicit dependencies later
        let mut task_ids_with_dependency_specifiers: Vec<(TaskId, Arc<[Str]>)> = Vec::new();

        // index tasks by ids
        let mut node_indices_by_task_id: HashMap<TaskId, TaskNodeIndex> =
            HashMap::with_capacity(task_graph.node_count());

        // Load task nodes into `task_graph`
        for package_index in package_graph.node_indices() {
            let package = &package_graph[package_index];
            let package_dir: Arc<AbsolutePath> = workspace_root.path.join(&package.path).into();

            // Collect package.json scripts into a mutable map for draining lookup.
            let mut package_json_scripts: HashMap<&str, &str> = package
                .package_json
                .scripts
                .iter()
                .map(|(name, value)| (name.as_str(), value.as_str()))
                .collect();

            // Load vite.config.* for the package
            let user_config: UserConfigFile =
                config_loader.load_user_config_file(&package_dir).await.map_err(|error| {
                    TaskGraphLoadError::ConfigLoadError { error, package_path: package_dir.clone() }
                })?;

            for (task_name, task_user_config) in user_config.tasks {
                // For each task defined in vite.config.*, look up the corresponding package.json script (if any)
                let package_json_script = package_json_scripts.remove(task_name.as_str());

                let task_id = TaskId { task_name: task_name.clone(), package_index };

                let dependency_specifiers = Arc::clone(&task_user_config.options.depends_on);

                // Resolve the task configuration combining vite.config.* and package.json script
                let resolved_config = ResolvedTaskConfig::resolve(
                    task_user_config,
                    &package_dir,
                    package_json_script,
                )
                .map_err(|err| TaskGraphLoadError::ResolveConfigError {
                    error: err,
                    task_display: TaskDisplay {
                        package_name: package.package_json.name.clone(),
                        task_name: task_name.clone(),
                        package_path: Arc::clone(&package_dir),
                    },
                })?;

                let task_node = TaskNode {
                    task_display: TaskDisplay {
                        package_name: package.package_json.name.clone(),
                        task_name: task_name.clone(),
                        package_path: Arc::clone(&package_dir),
                    },
                    resolved_config,
                };

                let node_index = task_graph.add_node(task_node);
                task_ids_with_dependency_specifiers.push((task_id.clone(), dependency_specifiers));
                node_indices_by_task_id.insert(task_id, node_index);
            }

            // For remaining package.json scripts not defined in vite.config.*, create tasks with default config
            for (script_name, package_json_script) in package_json_scripts.drain() {
                let task_id = TaskId { task_name: Str::from(script_name), package_index };
                let resolved_config = ResolvedTaskConfig::resolve_package_json_script(
                    &package_dir,
                    package_json_script,
                );
                let node_index = task_graph.add_node(TaskNode {
                    task_display: TaskDisplay {
                        package_name: package.package_json.name.clone(),
                        task_name: script_name.into(),
                        package_path: Arc::clone(&package_dir),
                    },
                    resolved_config,
                });
                node_indices_by_task_id.insert(task_id, node_index);
            }
        }

        // Grouping package indices by their package names.
        let mut package_indices_by_name: HashMap<Str, SmallVec1<[PackageNodeIndex; 1]>> =
            HashMap::new();
        for package_index in package_graph.node_indices() {
            let package = &package_graph[package_index];
            match package_indices_by_name.entry(package.package_json.name.clone()) {
                Entry::Vacant(vacant) => {
                    vacant.insert(SmallVec1::new(package_index));
                }
                Entry::Occupied(occupied) => {
                    occupied.into_mut().push(package_index);
                }
            }
        }

        // Construct `Self` with task_graph with all task nodes ready and indexed, but no edges.
        let mut me = Self {
            task_graph,
            indexed_package_graph: IndexedPackageGraph::index(package_graph),
            node_indices_by_task_id,
        };

        // Add explicit dependencies
        for (from_task_id, dependency_specifiers) in task_ids_with_dependency_specifiers {
            let from_node_index = me.node_indices_by_task_id[&from_task_id];
            for specifier in dependency_specifiers.iter().cloned() {
                let to_node_index = me
                    .get_task_index_by_specifier::<Infallible>(
                        TaskSpecifier::parse_raw(&specifier),
                        || Ok(from_task_id.package_index),
                    )
                    .map_err(|error| TaskGraphLoadError::DependencySpecifierLookupError {
                        error,
                        specifier,
                        task_display: me.display_task(from_node_index),
                    })?;
                me.task_graph.update_edge(
                    from_node_index,
                    to_node_index,
                    TaskDependencyType(TaskDependencyTypeInner::Explicit),
                );
            }
        }

        // Add topological dependencies based on package dependencies
        let mut nearest_topological_tasks = Vec::<TaskNodeIndex>::new();
        for (task_id, task_index) in &me.node_indices_by_task_id {
            let task_name = task_id.task_name.as_str();
            let package_index = task_id.package_index;

            // For every task, find nearest tasks with the same name.
            // If there is a task with the same name in a deeper dependency, it will be connected via that nearer dependency's task.
            me.find_nearest_topological_tasks(
                task_name,
                package_index,
                &mut nearest_topological_tasks,
            );
            for nearest_task_index in nearest_topological_tasks.drain(..) {
                if let Some(existing_edge_index) =
                    me.task_graph.find_edge(*task_index, nearest_task_index)
                {
                    // If an edge already exists,
                    let existing_edge = &mut me.task_graph[existing_edge_index];
                    match existing_edge.0 {
                        TaskDependencyTypeInner::Explicit => {
                            // upgrade from Explicit to Both
                            existing_edge.0 = TaskDependencyTypeInner::Both;
                        }
                        TaskDependencyTypeInner::Topological | TaskDependencyTypeInner::Both => {
                            // already has topological dependency, do nothing
                        }
                    }
                } else {
                    // add new topological edge if not exists
                    me.task_graph.add_edge(
                        *task_index,
                        nearest_task_index,
                        TaskDependencyType(TaskDependencyTypeInner::Topological),
                    );
                }
            }
        }
        Ok(me)
    }

    /// Find the nearest tasks with the given name starting from the given package node index.
    /// This method only considers the package graph topology. It doesn't rely on existing topological edges of
    /// the task graph, because the starting package itself may not contain a task with the given name
    ///
    /// The task with the given name in the starting package won't be included in the result even if there's one.
    ///
    /// This performs a BFS on the package graph starting from `starting_from`,
    /// and collects the first found tasks with the given name in each branch.
    ///
    /// For example, if the package graph is A -> B -> C and A -> D -> C,
    /// and we are looking for task "build" starting from A,
    ///
    /// - No matter A contains "build" or not, it's not included in the result.
    /// - If B and D both have a "build" task, both will be returned, but C's "build" task will not be returned
    ///   because it's not the nearest in either branch.
    /// - If B or D doesn't have a "build" task, then C's "build" task will be returned.
    fn find_nearest_topological_tasks(
        &self,
        task_name: &str,
        starting_from: PackageNodeIndex,
        out: &mut Vec<TaskNodeIndex>,
    ) {
        // DFS the package graph starting from `starting_from`,
        depth_first_search(
            self.indexed_package_graph.package_graph(),
            Some(starting_from),
            |event| {
                let DfsEvent::TreeEdge(_, dependency_package_index) = event else {
                    return Control::<()>::Continue;
                };

                if let Some(dependency_task_index) = self.node_indices_by_task_id.get(&TaskId {
                    package_index: dependency_package_index,
                    task_name: task_name.into(),
                }) {
                    // Encountered a package containing the task with the same name
                    // collect the task index
                    out.push(*dependency_task_index);

                    // And stop searching further down this branch
                    Control::Prune
                } else {
                    Control::Continue
                }
            },
        );
    }

    /// Lookup the node index of a task by a specifier.
    ///
    /// The specifier can be either 'packageName#taskName' or just 'taskName' (in which case the task in the origin package is looked up).
    fn get_task_index_by_specifier<PackageUnknownError>(
        &self,
        specifier: TaskSpecifier,
        get_package_origin: impl FnOnce() -> Result<PackageNodeIndex, PackageUnknownError>,
    ) -> Result<TaskNodeIndex, SpecifierLookupError<PackageUnknownError>> {
        let package_index = if let Some(package_name) = specifier.package_name {
            // Lookup package path by the package name from '#'
            let Some(package_indices) =
                self.indexed_package_graph.get_package_indices_by_name(&package_name)
            else {
                return Err(SpecifierLookupError::PackageNameNotFound {
                    package_name: package_name.into(),
                });
            };
            if package_indices.len() > 1 {
                return Err(SpecifierLookupError::AmbiguousPackageName {
                    package_name: package_name.into(),
                    package_paths: package_indices
                        .iter()
                        .map(|package_index| {
                            Arc::clone(
                                &self.indexed_package_graph.package_graph()[*package_index]
                                    .absolute_path,
                            )
                        })
                        .collect(),
                });
            };
            *package_indices.first()
        } else {
            // No '#', so the specifier only contains task name, look up in the origin path package
            get_package_origin().map_err(|err| SpecifierLookupError::PackageUnknown {
                unspecifier_package_error: err,
                task_name: specifier.task_name.clone(),
            })?
        };
        let task_id_to_lookup = TaskId { task_name: specifier.task_name, package_index };
        let Some(node_index) = self.node_indices_by_task_id.get(&task_id_to_lookup) else {
            return Err(SpecifierLookupError::TaskNameNotFound {
                package_name: self.indexed_package_graph.package_graph()[package_index]
                    .package_json
                    .name
                    .clone(),
                task_name: task_id_to_lookup.task_name.clone(),
                package_index,
            });
        };
        Ok(*node_index)
    }

    pub fn task_graph(&self) -> &TaskGraph {
        &self.task_graph
    }

    pub fn get_package_name(&self, package_index: PackageNodeIndex) -> &str {
        self.indexed_package_graph.package_graph()[package_index].package_json.name.as_str()
    }

    pub fn get_package_path(&self, package_index: PackageNodeIndex) -> &Arc<AbsolutePath> {
        &self.indexed_package_graph.package_graph()[package_index].absolute_path
    }

    pub fn get_package_path_for_task(&self, task_index: TaskNodeIndex) -> &Arc<AbsolutePath> {
        &self.task_graph[task_index].task_display.package_path
    }
}