Librarify the plug command

Signed-off-by: Brian H <brian.hardock@fermyon.com>

Author: fibonacci1729

crates/wac-graph/src/lib.rs

@@ -39,6 +39,8 @@
 
 pub(crate) mod encoding;
 mod graph;
+mod plug;
 
 pub use graph::*;
+pub use plug::*;
 pub use wac_types as types;

crates/wac-graph/src/plug.rs

@@ -0,0 +1,86 @@
+use crate::{types::SubtypeChecker, CompositionGraph, PackageId};
+use thiserror::Error;
+
+/// Represents an error that can occur when plugging components together.
+#[derive(Debug, Error)]
+pub enum PlugError {
+    /// The socket component had no matching imports for the plugs that were provided
+    #[error("the socket component had no matching imports for the plugs that were provided")]
+    NoPlugHappened,
+    /// An error occurred when building the composition graph
+    #[error("an error occurred when building the composition graph")]
+    GraphError {
+        /// The underlying graph error
+        #[source]
+        source: anyhow::Error,
+    },
+}
+
+/// Take the exports of the plug components and plug them into the socket component.
+///
+/// Note that the `PackageId`s in `plugs` and `socket` must be registered with the `graph`
+/// before calling this function.
+pub fn plug(
+    graph: &mut CompositionGraph,
+    plugs: Vec<PackageId>,
+    socket: PackageId,
+) -> Result<(), PlugError> {
+    let socket_instantiation = graph.instantiate(socket);
+
+    for plug in plugs {
+        let mut plug_exports = Vec::new();
+        let mut cache = Default::default();
+        let mut checker = SubtypeChecker::new(&mut cache);
+        for (name, plug_ty) in &graph.types()[graph[plug].ty()].exports {
+            if let Some(socket_ty) = graph.types()[graph[socket].ty()].imports.get(name) {
+                if checker
+                    .is_subtype(*plug_ty, graph.types(), *socket_ty, graph.types())
+                    .is_ok()
+                {
+                    plug_exports.push(name.clone());
+                }
+            }
+        }
+
+        // Instantiate the plug component
+        let mut plug_instantiation = None;
+        for plug_name in plug_exports {
+            log::debug!("using export `{plug_name}` for plug");
+            let plug_instantiation =
+                *plug_instantiation.get_or_insert_with(|| graph.instantiate(plug));
+            let export = graph
+                .alias_instance_export(plug_instantiation, &plug_name)
+                .map_err(|err| PlugError::GraphError { source: err.into() })?;
+            graph
+                .set_instantiation_argument(socket_instantiation, &plug_name, export)
+                .map_err(|err| PlugError::GraphError { source: err.into() })?;
+        }
+    }
+
+    // Check we've actually done any plugging.
+    if graph
+        .get_instantiation_arguments(socket_instantiation)
+        .next()
+        .is_none()
+    {
+        return Err(PlugError::NoPlugHappened);
+    }
+
+    // Export all exports from the socket component.
+    for name in graph.types()[graph[socket].ty()]
+        .exports
+        .keys()
+        .cloned()
+        .collect::<Vec<_>>()
+    {
+        let export = graph
+            .alias_instance_export(socket_instantiation, &name)
+            .map_err(|err| PlugError::GraphError { source: err.into() })?;
+
+        graph
+            .export(export, &name)
+            .map_err(|err| PlugError::GraphError { source: err.into() })?;
+    }
+
+    Ok(())
+}

src/commands/plug.rs

@@ -6,8 +6,8 @@ use std::{
 use anyhow::{bail, Context as _, Error, Result};
 use clap::Args;
 use std::str::FromStr;
-use wac_graph::{CompositionGraph, EncodeOptions, NodeId, PackageId};
-use wac_types::{Package, SubtypeChecker};
+use wac_graph::{CompositionGraph, EncodeOptions};
+use wac_types::Package;
 
 #[cfg(feature = "registry")]
 use warg_client::FileSystemClient;
@@ -151,7 +151,6 @@ impl PlugCommand {
 
         let socket = Package::from_bytes("socket", None, socket, graph.types_mut())?;
         let socket = graph.register_package(socket)?;
-        let socket_instantiation = graph.instantiate(socket);
 
         // Collect the plugs by their names
         let mut plugs_by_name = std::collections::HashMap::<_, Vec<_>>::new();
@@ -169,6 +168,8 @@ impl PlugCommand {
             plugs_by_name.entry(name).or_default().push(plug);
         }
 
+        let mut plug_packages = Vec::new();
+
         // Plug each plug into the socket.
         for (name, plug_refs) in plugs_by_name {
             for (i, plug_ref) in plug_refs.iter().enumerate() {
@@ -221,37 +222,23 @@ impl PlugCommand {
                     use core::fmt::Write;
                     write!(&mut name, "{i}").unwrap();
                 }
-                plug_into_socket(&name, &path, socket, socket_instantiation, &mut graph)?;
+
+                let plug_package = Package::from_file(&name, None, &path, graph.types_mut())?;
+                let package_id = graph.register_package(plug_package)?;
+                plug_packages.push(package_id);
             }
         }
 
-        // Check we've actually done any plugging.
-        if graph
-            .get_instantiation_arguments(socket_instantiation)
-            .next()
-            .is_none()
-        {
-            bail!("the socket component had no matching imports for the plugs that were provided")
-        }
+        wac_graph::plug(&mut graph, plug_packages, socket)?;
 
-        // Export all exports from the socket component.
-        for name in graph.types()[graph[socket].ty()]
-            .exports
-            .keys()
-            .cloned()
-            .collect::<Vec<_>>()
-        {
-            let export = graph.alias_instance_export(socket_instantiation, &name)?;
-            graph.export(export, &name)?;
-        }
+        let mut bytes = graph.encode(EncodeOptions::default())?;
 
         let binary_output_to_terminal =
             !self.wat && self.output.is_none() && std::io::stdout().is_terminal();
         if binary_output_to_terminal {
             bail!("cannot print binary wasm output to a terminal; pass the `-t` flag to print the text format instead");
         }
 
-        let mut bytes = graph.encode(EncodeOptions::default())?;
         if self.wat {
             bytes = wasmprinter::print_bytes(&bytes)
                 .context("failed to convert binary wasm output to text")?
@@ -278,39 +265,3 @@ impl PlugCommand {
         Ok(())
     }
 }
-
-/// Take the exports of the plug component and plug them into the socket component.
-fn plug_into_socket(
-    name: &str,
-    plug_path: &std::path::Path,
-    socket: PackageId,
-    socket_instantiation: NodeId,
-    graph: &mut CompositionGraph,
-) -> Result<(), anyhow::Error> {
-    let plug = Package::from_file(name, None, plug_path, graph.types_mut())?;
-    let plug = graph.register_package(plug)?;
-
-    let mut plugs = Vec::new();
-    let mut cache = Default::default();
-    let mut checker = SubtypeChecker::new(&mut cache);
-    for (name, plug_ty) in &graph.types()[graph[plug].ty()].exports {
-        if let Some(socket_ty) = graph.types()[graph[socket].ty()].imports.get(name) {
-            if checker
-                .is_subtype(*plug_ty, graph.types(), *socket_ty, graph.types())
-                .is_ok()
-            {
-                plugs.push(name.clone());
-            }
-        }
-    }
-
-    // Instantiate the plug component
-    let mut plug_instantiation = None;
-    for plug_name in plugs {
-        log::debug!("using export `{plug_name}` for plug");
-        let plug_instantiation = *plug_instantiation.get_or_insert_with(|| graph.instantiate(plug));
-        let export = graph.alias_instance_export(plug_instantiation, &plug_name)?;
-        graph.set_instantiation_argument(socket_instantiation, &plug_name, export)?;
-    }
-    Ok(())
-}