From 09f28a3b808f6d3b6125ae7eb09328e92a3f2e38 Mon Sep 17 00:00:00 2001
From: Mara Schulke <mara.schulke@helsing.ai>
Date: Tue, 15 Aug 2023 18:32:52 +0200
Subject: [PATCH] Compile locally defined packages

---
 src/generator.rs | 66 ++++++++++++++++++++++++++++++++++--------------
 1 file changed, 47 insertions(+), 19 deletions(-)

diff --git a/src/generator.rs b/src/generator.rs
index 68a27d36..491a3450 100644
--- a/src/generator.rs
+++ b/src/generator.rs
@@ -9,8 +9,8 @@ use serde::{Deserialize, Serialize};
 use tokio::fs;
 
 use crate::{
-    manifest::{Dependency, Manifest},
-    package::PackageStore,
+    manifest::Manifest,
+    package::{PackageId, PackageStore},
 };
 
 /// The directory used for the generated code
@@ -37,6 +37,14 @@ impl fmt::Display for Language {
     }
 }
 
+/// A configuration passed to a generator to compile a package
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
+pub struct GenerationConfig<'i, 'p> {
+    pub package: &'i PackageId,
+    pub location: &'p Path,
+    pub out: &'p Path,
+}
+
 /// Backend used to generate code bindings
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
 pub enum Generator {
@@ -48,25 +56,24 @@ impl Generator {
     pub const TONIC_INCLUDE_FILE: &str = "mod.rs";
 
     /// Run the generator for a dependency and output files into `out`
-    pub async fn run(&self, dependency: &Dependency, out: &Path) -> eyre::Result<()> {
+    pub async fn run(&self, config: GenerationConfig<'_, '_>) -> eyre::Result<()> {
         let protoc = protoc_bin_path().wrap_err("Unable to locate vendored protoc")?;
 
         std::env::set_var("PROTOC", protoc.clone());
 
-        match self {
-            Generator::Tonic => {
-                let out = out.join(dependency.package.as_str());
+        let out = config.out.join(config.package.as_str());
 
-                fs::remove_dir_all(&out).await.ok();
+        fs::remove_dir_all(&out).await.ok();
 
-                fs::create_dir_all(&out)
-                    .await
-                    .wrap_err("Failed to recreate dependency output directory")?;
+        fs::create_dir_all(&out)
+            .await
+            .wrap_err("Failed to recreate dependency output directory")?;
 
-                let package = PackageStore::locate(&dependency.package);
-                let protos = PackageStore::collect(&package).await;
+        match self {
+            Generator::Tonic => {
+                let protos = PackageStore::collect(config.location).await;
 
-                let includes = &[package];
+                let includes = &[&config.location];
 
                 tonic_build::configure()
                     .build_client(true)
@@ -90,8 +97,8 @@ pub async fn generate(language: Language) -> eyre::Result<()> {
     tracing::info!(":: initializing code generator for {language}");
 
     eyre::ensure!(
-        !manifest.dependencies.is_empty(),
-        "At least one dependency is needed to generate code bindings."
+        manifest.package.is_some() || !manifest.dependencies.is_empty(),
+        "Either a local package or at least one dependency is needed to generate code bindings."
     );
 
     // Only tonic is supported right now
@@ -110,15 +117,37 @@ pub async fn generate(language: Language) -> eyre::Result<()> {
         out
     };
 
+    if let Some(ref pkg) = manifest.package {
+        let location = Path::new(PackageStore::PROTO_PATH);
+
+        generator
+            .run(GenerationConfig {
+                package: &pkg.name,
+                location,
+                out: &out,
+            })
+            .await
+            .wrap_err_with(|| format!("Failed to generate bindings for {}", pkg.name))?;
+
+        tracing::info!(":: compiled {} [{}]", pkg.name, location.display());
+    }
+
     for dependency in manifest.dependencies {
+        let location = PackageStore::locate(&dependency.package);
+
         generator
-            .run(&dependency, &out)
+            .run(GenerationConfig {
+                package: &dependency.package,
+                location: &location,
+                out: &out,
+            })
             .await
             .wrap_err_with(|| format!("Failed to generate bindings for {}", dependency.package))?;
 
         tracing::info!(
-            ":: compiled {}",
-            PackageStore::locate(&dependency.package).display()
+            ":: compiled {} [{}]",
+            dependency.package,
+            location.display()
         );
     }
 
@@ -137,7 +166,6 @@ pub async fn generate(language: Language) -> eyre::Result<()> {
 #[macro_export]
 macro_rules! include {
     ($package:expr) => {
-        #[allow(non_snake_case)]
         ::std::include!(concat!(
             env!("CARGO_MANIFEST_DIR"),
             "/proto/build/rust/",