feat: refactor and cleanup Adg

Signed-off-by: Andrew Lilley Brinker <[email protected]>

Author: alilleybrinker

omnibor/src/adg.rs

@@ -1,41 +1,6 @@
 //! Operations and types for handling Artifact Dependency Graphs (ADGs).
-#![allow(unused)]
 
-use crate::{
-    error::InputManifestError,
-    hash_algorithm::HashAlgorithm,
-    storage::{Match, Storage},
-    ArtifactId,
-};
-use petgraph::{graph::NodeIndex, Graph};
+pub(crate) mod adg_internal;
+mod adg_node;
 
-// Convenience type for our internal dependency graph.
-pub(crate) type DepGraphInner<H> = Graph<AdgNode<H>, ()>;
-
-/// The Artifact Dependency Graph.
-#[derive(Debug)]
-pub struct DepGraph<H>
-where
-    H: HashAlgorithm,
-{
-    graph: DepGraphInner<H>,
-}
-
-impl<H> DepGraph<H>
-where
-    H: HashAlgorithm,
-{
-    pub(crate) fn from_graph(graph: DepGraphInner<H>) -> Self {
-        Self { graph }
-    }
-}
-
-/// A node in the Artifact Dependency Graph.
-#[derive(Debug, PartialEq, Eq, Clone, Copy, PartialOrd, Ord)]
-pub struct AdgNode<H>
-where
-    H: HashAlgorithm,
-{
-    /// The Artifact ID of the node's artifact.
-    pub id: ArtifactId<H>,
-}
+pub use adg_node::AdgNode;

omnibor/src/adg/adg_internal.rs

@@ -0,0 +1,74 @@
+#![allow(unused)]
+
+use crate::{
+    adg::AdgNode,
+    error::ArtifactIdError,
+    hash_algorithm::HashAlgorithm,
+    storage::{Match, Storage},
+    ArtifactId, Identify,
+};
+use petgraph::{algo::astar, graph::NodeIndex, Graph};
+use std::ops::Not;
+
+// Convenience type for our internal dependency graph.
+pub(crate) type AdgInner<H> = Graph<AdgNode<H>, ()>;
+
+/// The Artifact Dependency Graph.
+#[derive(Debug)]
+pub struct Adg<H>
+where
+    H: HashAlgorithm,
+{
+    /// The internal graph type.
+    graph: AdgInner<H>,
+    /// The target artifact the ADG is describing.
+    target: NodeIndex,
+}
+
+impl<H> Adg<H>
+where
+    H: HashAlgorithm,
+{
+    // Actual construction is handled by the storage.
+    pub(crate) fn new(graph: AdgInner<H>, target: NodeIndex) -> Self {
+        Self { graph, target }
+    }
+
+    /// Get the target artifact for this ADG.
+    pub fn target(&self) -> ArtifactId<H> {
+        self.graph[self.target].artifact_id()
+    }
+
+    /// Get the path to a dependency, if present in the graph.
+    pub fn dependency_path(&self, dep: ArtifactId<H>) -> Result<Vec<&AdgNode<H>>, ArtifactIdError> {
+        let dep = dep.identify()?;
+
+        let search = astar(
+            &self.graph,
+            // Start from the target.
+            self.target,
+            // End when we find the dependency.
+            |n| self.graph[n].artifact_id() == dep,
+            // Each edge has weight 1.
+            |_| 1,
+            // We don't do estimation.
+            |_| 0,
+        );
+
+        match search {
+            None => Ok(Vec::new()),
+            Some((_, path)) => Ok(path.iter().map(|idx| &self.graph[*idx]).collect()),
+        }
+    }
+
+    /// Check if the ADG's target artifact depends on a specific artifact.
+    pub fn depends_on(&self, dep: ArtifactId<H>) -> Result<bool, ArtifactIdError> {
+        let dep = dep.identify()?;
+
+        // Rather than doing a graph algo, check if the dep is one of the nodes.
+        Ok(self
+            .graph
+            .node_indices()
+            .any(|idx| self.graph[idx].artifact_id() == dep))
+    }
+}

omnibor/src/adg/adg_node.rs

@@ -0,0 +1,44 @@
+#![allow(unused)]
+
+use crate::{
+    error::{ArtifactIdError, InputManifestError},
+    hash_algorithm::HashAlgorithm,
+    storage::{Match, Storage},
+    ArtifactId, Identify,
+};
+use petgraph::{algo::astar, graph::NodeIndex, Graph};
+
+/// A node in the Artifact Dependency Graph.
+#[derive(Debug, PartialEq, Eq, Copy, PartialOrd, Ord)]
+pub struct AdgNode<H>
+where
+    H: HashAlgorithm,
+{
+    /// The Artifact ID of the node's artifact.
+    artifact_id: ArtifactId<H>,
+}
+
+impl<H> AdgNode<H>
+where
+    H: HashAlgorithm,
+{
+    pub(crate) fn new(artifact_id: ArtifactId<H>) -> Self {
+        AdgNode { artifact_id }
+    }
+
+    /// Get the Artifact ID for the node.
+    pub fn artifact_id(&self) -> ArtifactId<H> {
+        self.artifact_id
+    }
+}
+
+impl<H> Clone for AdgNode<H>
+where
+    H: HashAlgorithm,
+{
+    fn clone(&self) -> Self {
+        AdgNode {
+            artifact_id: self.artifact_id,
+        }
+    }
+}

omnibor/src/lib.rs

@@ -517,6 +517,7 @@ pub mod hash_algorithm;
 pub mod hash_provider;
 pub mod storage;
 
+pub use crate::adg::adg_internal::Adg;
 pub use crate::artifact_id::ArtifactId;
 pub use crate::artifact_id::Identify;
 pub use crate::artifact_id::IdentifyAsync;

omnibor/src/storage.rs

@@ -23,12 +23,12 @@ pub use crate::storage::in_memory_storage::InMemoryStorage;
 pub use crate::storage::query::Match;
 
 use crate::{
-    adg::{AdgNode, DepGraph, DepGraphInner},
+    adg::{adg_internal::AdgInner, AdgNode},
     artifact_id::ArtifactId,
     error::InputManifestError,
     hash_algorithm::HashAlgorithm,
     input_manifest::InputManifest,
-    Identify,
+    Adg, Identify,
 };
 
 /// Represents the interface for storing and querying manifests.
@@ -78,13 +78,13 @@ pub trait Storage<H: HashAlgorithm> {
         I: Identify<H>;
 
     /// Get the Artifact Dependency Graph of the target.
-    fn get_adg<I>(&self, target_aid: I) -> Result<DepGraph<H>, InputManifestError>
+    fn get_adg<I>(&self, target_aid: I) -> Result<Adg<H>, InputManifestError>
     where
         I: Identify<H>,
     {
         let target_aid = target_aid.identify()?;
-        let graph = build_graph(target_aid, self)?;
-        Ok(DepGraph::from_graph(graph))
+        let (graph, root) = build_graph(target_aid, self)?;
+        Ok(Adg::new(graph, root))
     }
 }
 
@@ -137,29 +137,29 @@ impl<H: HashAlgorithm, S: Storage<H>> Storage<H> for &mut S {
 fn build_graph<H, S>(
     target_aid: ArtifactId<H>,
     storage: &S,
-) -> Result<DepGraphInner<H>, InputManifestError>
+) -> Result<(AdgInner<H>, NodeIndex), InputManifestError>
 where
     H: HashAlgorithm,
     S: Storage<H> + ?Sized,
 {
     let mut graph = Graph::new();
-    populate_graph(&mut graph, target_aid, None, storage)?;
-    Ok(graph)
+    let root = populate_graph(&mut graph, target_aid, None, storage)?;
+    Ok((graph, root))
 }
 
 /// Build one layer of the ADG, recursing down for further layers.
 fn populate_graph<H, S>(
-    graph: &mut DepGraphInner<H>,
+    graph: &mut AdgInner<H>,
     target_aid: ArtifactId<H>,
     parent_idx: Option<NodeIndex>,
     storage: &S,
-) -> Result<(), InputManifestError>
+) -> Result<NodeIndex, InputManifestError>
 where
     H: HashAlgorithm,
     S: Storage<H> + ?Sized,
 {
     // Add current node to the graph.
-    let self_idx = graph.add_node(AdgNode { id: target_aid });
+    let self_idx = graph.add_node(AdgNode::new(target_aid));
 
     // If there's a parent, add an edge from the parent to the current node.
     if let Some(parent_idx) = parent_idx {
@@ -169,9 +169,9 @@ where
     // If there's a manifest for current node, recurse for each child.
     if let Some(manifest) = storage.get_manifest(Match::target(target_aid))? {
         for input in manifest.inputs() {
-            populate_graph(graph, input.artifact(), Some(self_idx), storage)?;
+            let _ = populate_graph(graph, input.artifact(), Some(self_idx), storage)?;
         }
     }
 
-    Ok(())
+    Ok(self_idx)
 }