distro/src/plan.rs

use crate::recipe::{PackageRecipe, RecipeSet, SourceRecipe, ToolRecipe};

use petgraph::{
    Direction,
    graph::{DiGraph, NodeIndex},
};
use smallvec::{SmallVec, smallvec};
use std::{
    cmp::Reverse,
    collections::{BinaryHeap, HashMap, HashSet},
};
use thiserror::Error;

#[derive(Debug, Error)]
pub enum PlanError {
    #[error("missing source recipe '{0}'")]
    MissingSource(String),
    #[error("missing tool recipe '{0}'")]
    MissingTool(String),
    #[error("missing package recipe '{0}'")]
    MissingPackage(String),
    #[error("plan cycle detected")]
    CycleDetected,
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum PlanKey {
    SourceFetch(String),
    SourcePatch(String),
    SourcePrepare(String),

    ToolConfigure(String),
    ToolBuild(String),
    ToolInstall(String),

    PkgConfigure(String),
    PkgBuild(String),
    PkgInstall(String),
    PkgPackage(String),
}

impl PlanKey {
    fn weight(&self) -> i8 {
        match self {
            PlanKey::SourceFetch(_) => 0,
            PlanKey::SourcePatch(_) => 1,
            PlanKey::SourcePrepare(_) => 2,
            PlanKey::ToolConfigure(_) => 3,
            PlanKey::ToolBuild(_) => 4,
            PlanKey::ToolInstall(_) => 5,
            PlanKey::PkgConfigure(_) => 6,
            PlanKey::PkgBuild(_) => 7,
            PlanKey::PkgInstall(_) => 8,
            PlanKey::PkgPackage(_) => 9,
        }
    }

    fn dependencies(&self, recipes: &RecipeSet) -> Result<SmallVec<[PlanKey; 8]>, PlanError> {
        match self {
            PlanKey::SourceFetch(_) => Ok(smallvec![]),
            PlanKey::SourcePatch(recipe) => Ok(smallvec![PlanKey::SourceFetch(recipe.clone())]),
            PlanKey::SourcePrepare(recipe) => Ok(smallvec![PlanKey::SourcePatch(recipe.clone())]),
            PlanKey::ToolConfigure(recipe) => {
                let recipe = recipes
                    .tool(recipe)
                    .ok_or(PlanError::MissingTool(recipe.clone()))?;

                let source_deps = recipe.sources.iter().map(|name| {
                    recipes
                        .source(name)
                        .map(|_| PlanKey::SourcePrepare(name.to_string()))
                        .ok_or(PlanError::MissingSource(name.to_string()))
                });

                let tool_deps = recipe.tools_wanted.iter().map(|name| {
                    recipes
                        .tool(name)
                        .map(|_| PlanKey::ToolInstall(name.to_string()))
                        .ok_or(PlanError::MissingTool(name.to_string()))
                });

                let pkg_deps = recipe.pkgs_wanted.iter().map(|name| {
                    recipes
                        .package(name)
                        .map(|_| PlanKey::PkgPackage(name.to_string()))
                        .ok_or(PlanError::MissingPackage(name.to_string()))
                });

                source_deps.chain(tool_deps).chain(pkg_deps).collect()
            }
            PlanKey::ToolBuild(recipe) => Ok(smallvec![PlanKey::ToolConfigure(recipe.clone())]),
            PlanKey::ToolInstall(recipe) => Ok(smallvec![PlanKey::ToolBuild(recipe.clone())]),
            PlanKey::PkgConfigure(recipe) => {
                let recipe = recipes
                    .package(recipe)
                    .ok_or(PlanError::MissingPackage(recipe.clone()))?;

                let source_deps = recipe.sources.iter().map(|name| {
                    recipes
                        .source(name)
                        .map(|_| PlanKey::SourcePrepare(name.to_string()))
                        .ok_or(PlanError::MissingSource(name.to_string()))
                });

                let tool_deps = recipe.tools_wanted.iter().map(|name| {
                    recipes
                        .tool(name)
                        .map(|_| PlanKey::ToolInstall(name.to_string()))
                        .ok_or(PlanError::MissingTool(name.to_string()))
                });

                let pkg_deps = recipe.pkgs_wanted.iter().map(|name| {
                    recipes
                        .package(name)
                        .map(|_| PlanKey::PkgPackage(name.to_string()))
                        .ok_or(PlanError::MissingPackage(name.to_string()))
                });

                source_deps.chain(tool_deps).chain(pkg_deps).collect()
            }
            PlanKey::PkgBuild(recipe) => Ok(smallvec![PlanKey::PkgConfigure(recipe.clone())]),
            PlanKey::PkgInstall(recipe) => Ok(smallvec![PlanKey::PkgBuild(recipe.clone())]),
            PlanKey::PkgPackage(recipe) => Ok(smallvec![PlanKey::PkgInstall(recipe.clone())]),
        }
    }

    fn is_active(&self) -> bool {
        true
    }
}

pub struct Plan<'a> {
    recipes: &'a RecipeSet<'a>,
    wanted: HashSet<PlanKey>,
}

impl<'a> Plan<'a> {
    pub fn new(recipes: &'a RecipeSet) -> Self {
        Self {
            recipes,
            wanted: HashSet::new(),
        }
    }

    pub fn add_wanted(&mut self, key: PlanKey) {
        self.wanted.insert(key);
    }

    pub fn steps(&self) -> Result<Vec<PlanKey>, PlanError> {
        let mut stack: Vec<_> = self.wanted.iter().cloned().collect();
        let mut graph: DiGraph<_, ()> = DiGraph::new();
        let mut nodes = HashMap::new();

        while let Some(node) = stack.pop() {
            let node_idx = match nodes.get(&node) {
                Some(&idx) => idx,
                None => {
                    let idx = graph.add_node(node.clone());
                    nodes.insert(node.clone(), idx);
                    idx
                }
            };

            for dep in node.dependencies(self.recipes)? {
                let dep_idx = match nodes.get(&dep) {
                    Some(&idx) => idx,
                    None => {
                        let idx = graph.add_node(dep.clone());
                        nodes.insert(dep.clone(), idx);
                        stack.push(dep);
                        idx
                    }
                };

                graph.update_edge(dep_idx, node_idx, ());
            }
        }

        // petgraph::algo::toposort(&graph, None)
        //     .and_then(|nodes| {
        //         Ok(nodes
        //             .iter()
        //             .map(|&k| graph[k])
        //             .filter(|node| node.is_active())
        //             .collect())
        //     })
        //     .map_err(|_| PlanError::CycleDetected)

        let mut in_degree: HashMap<NodeIndex, usize> = graph
            .node_indices()
            .map(|i| (i, graph.neighbors_directed(i, Direction::Incoming).count()))
            .collect();

        let mut heap: BinaryHeap<(Reverse<i8>, NodeIndex)> = in_degree
            .iter()
            .filter(|&(_, d)| *d == 0)
            .map(|(&i, _)| (Reverse(graph[i].weight()), i))
            .collect();

        let mut result = Vec::with_capacity(graph.node_count());

        while let Some((_, idx)) = heap.pop() {
            result.push(graph[idx].clone());

            for neighbor in graph.neighbors_directed(idx, Direction::Outgoing) {
                let d = in_degree.get_mut(&neighbor).unwrap();
                *d -= 1;
                if *d == 0 {
                    heap.push((Reverse(graph[neighbor].weight()), neighbor));
                }
            }
        }

        if result.len() != graph.node_count() {
            Err(PlanError::CycleDetected)
        } else {
            result.retain(|node| node.is_active());

            Ok(result)
        }
    }
}