Add new collect_references recipe

crates/brioche-core/src/bake.rs

@@ -18,6 +18,7 @@ use super::{
     Brioche,
 };
 
+mod collect_references;
 mod download;
 mod process;
 mod unarchive;
@@ -492,6 +493,16 @@ async fn run_bake(brioche: &Brioche, recipe: Recipe, meta: &Arc<Meta>) -> anyhow
 
             Ok(Artifact::File(file))
         }
+        Recipe::CollectReferences { recipe } => {
+            let artifact = bake(brioche, *recipe, &scope).await?;
+            let Artifact::Directory(directory) = artifact.value else {
+                anyhow::bail!("tried collecting references for non-directory artifact");
+            };
+
+            let directory = collect_references::bake_collect_references(brioche, directory).await?;
+
+            Ok(Artifact::Directory(directory))
+        }
         Recipe::Sync { recipe } => {
             let result = bake(brioche, *recipe, &scope).await?;
             Ok(result.value)

crates/brioche-core/src/bake/collect_references.rs

@@ -0,0 +1,76 @@
+use crate::{
+    recipe::{Artifact, Directory, WithMeta},
+    Brioche,
+};
+
+pub async fn bake_collect_references(
+    brioche: &Brioche,
+    directory: Directory,
+) -> anyhow::Result<Directory> {
+    let mut resources_dir = Directory::default();
+
+    let mut updated_directory =
+        collect_directory_references(brioche, directory, &mut resources_dir).await?;
+
+    if !resources_dir.is_empty() {
+        let mut root_with_resources_dir = Directory::default();
+        root_with_resources_dir
+            .insert(
+                brioche,
+                b"brioche-resources.d",
+                Some(WithMeta::without_meta(Artifact::Directory(resources_dir))),
+            )
+            .await?;
+
+        updated_directory
+            .merge(&root_with_resources_dir, brioche)
+            .await?;
+    }
+
+    Ok(updated_directory)
+}
+
+async fn collect_references(
+    brioche: &Brioche,
+    artifact: Artifact,
+    resources_dir: &mut Directory,
+) -> anyhow::Result<Artifact> {
+    match artifact {
+        Artifact::File(mut file) => {
+            let file_resources = std::mem::take(&mut file.resources);
+            let file_resources =
+                collect_directory_references(brioche, file_resources, resources_dir).await?;
+            resources_dir.merge(&file_resources, brioche).await?;
+
+            Ok(Artifact::File(file))
+        }
+        Artifact::Symlink { .. } => Ok(artifact),
+        Artifact::Directory(directory) => {
+            let new_directory =
+                collect_directory_references(brioche, directory, resources_dir).await?;
+            Ok(Artifact::Directory(new_directory))
+        }
+    }
+}
+
+async fn collect_directory_references(
+    brioche: &Brioche,
+    directory: Directory,
+    resources_dir: &mut Directory,
+) -> anyhow::Result<Directory> {
+    let mut new_directory = Directory::default();
+
+    for (name, artifact) in directory.entries(brioche).await? {
+        let new_artifact =
+            Box::pin(collect_references(brioche, artifact.value, resources_dir)).await?;
+        new_directory
+            .insert(
+                brioche,
+                &name,
+                Some(WithMeta::new(new_artifact, artifact.meta)),
+            )
+            .await?;
+    }
+
+    Ok(new_directory)
+}

crates/brioche-core/src/recipe.rs

@@ -91,6 +91,9 @@ pub enum Recipe {
         file: Box<WithMeta<Recipe>>,
         executable: Option<bool>,
     },
+    CollectReferences {
+        recipe: Box<WithMeta<Recipe>>,
+    },
     #[serde(rename_all = "camelCase")]
     Proxy(ProxyRecipe),
     #[serde(rename_all = "camelCase")]
@@ -149,6 +152,7 @@ impl Recipe {
             | Recipe::Get { .. }
             | Recipe::Insert { .. }
             | Recipe::SetPermissions { .. }
+            | Recipe::CollectReferences { .. }
             | Recipe::Proxy(_) => false,
         }
     }
@@ -1027,6 +1031,7 @@ impl TryFrom<Recipe> for Artifact {
             | Recipe::Get { .. }
             | Recipe::Insert { .. }
             | Recipe::SetPermissions { .. }
+            | Recipe::CollectReferences { .. }
             | Recipe::Proxy { .. } => Err(RecipeIncomplete),
         }
     }

crates/brioche-core/src/references.rs

@@ -128,6 +128,7 @@ pub fn referenced_blobs(recipe: &Recipe) -> Vec<BlobHash> {
         | Recipe::Insert { .. }
         | Recipe::SetPermissions { .. }
         | Recipe::Proxy(_)
+        | Recipe::CollectReferences { .. }
         | Recipe::Sync { .. } => vec![],
     }
 }
@@ -258,6 +259,7 @@ pub fn referenced_recipes(recipe: &Recipe) -> Vec<RecipeHash> {
             executable: _,
         } => referenced_recipes(file),
         Recipe::Proxy(proxy) => vec![proxy.recipe],
+        Recipe::CollectReferences { recipe } => referenced_recipes(recipe),
         Recipe::Sync { recipe } => referenced_recipes(recipe),
     }
 }

crates/brioche-core/tests/bake_collect_references.rs

@@ -0,0 +1,114 @@
+use brioche_core::recipe::{Recipe, WithMeta};
+
+mod brioche_test;
+
+#[tokio::test]
+async fn test_bake_collect_references() -> anyhow::Result<()> {
+    let (brioche, _context) = brioche_test::brioche_test().await;
+
+    let foo_blob = brioche_test::blob(&brioche, b"foo").await;
+    let bar_blob = brioche_test::blob(&brioche, b"bar").await;
+    let resource_a_blob = brioche_test::blob(&brioche, b"resource a").await;
+    let resource_b_blob = brioche_test::blob(&brioche, b"resource b").await;
+    let resource_c_blob = brioche_test::blob(&brioche, b"resource c").await;
+
+    let resource_a = brioche_test::file(resource_a_blob, false);
+    let resource_b = brioche_test::file(resource_b_blob, false);
+    let resource_c = brioche_test::file_with_resources(
+        resource_c_blob,
+        false,
+        brioche_test::dir_value(&brioche, [("foo/a.txt", resource_a.clone())]).await,
+    );
+
+    let fizz = brioche_test::dir(
+        &brioche,
+        [(
+            "file.txt",
+            brioche_test::file_with_resources(
+                foo_blob,
+                false,
+                brioche_test::dir_value(&brioche, [("foo/b.txt", resource_b.clone())]).await,
+            ),
+        )],
+    )
+    .await;
+
+    let buzz = brioche_test::file_with_resources(
+        brioche_test::blob(&brioche, b"bar").await,
+        false,
+        brioche_test::dir_value(&brioche, [("foo/c.txt", resource_c.clone())]).await,
+    );
+
+    let dir = brioche_test::dir(
+        &brioche,
+        [("fizz", fizz.clone()), ("buzz.txt", buzz.clone())],
+    )
+    .await;
+    let recipe = Recipe::CollectReferences {
+        recipe: Box::new(WithMeta::without_meta(dir.clone().into())),
+    };
+
+    let expected_output = brioche_test::dir(
+        &brioche,
+        [
+            (
+                "fizz",
+                brioche_test::dir(
+                    &brioche,
+                    [("file.txt", brioche_test::file(foo_blob, false))],
+                )
+                .await,
+            ),
+            ("buzz.txt", brioche_test::file(bar_blob, false)),
+            (
+                "brioche-resources.d",
+                brioche_test::dir(
+                    &brioche,
+                    [
+                        ("foo/a.txt", brioche_test::file(resource_a_blob, false)),
+                        ("foo/b.txt", brioche_test::file(resource_b_blob, false)),
+                        ("foo/c.txt", brioche_test::file(resource_c_blob, false)),
+                    ],
+                )
+                .await,
+            ),
+        ],
+    )
+    .await;
+
+    let output = brioche_test::bake_without_meta(&brioche, recipe).await?;
+
+    assert_eq!(output, expected_output);
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_bake_collect_references_no_references() -> anyhow::Result<()> {
+    let (brioche, _context) = brioche_test::brioche_test().await;
+
+    let dir = brioche_test::dir(
+        &brioche,
+        [
+            (
+                "foo.txt",
+                brioche_test::file(brioche_test::blob(&brioche, "foo!").await, false),
+            ),
+            (
+                "bar/baz.txt",
+                brioche_test::file(brioche_test::blob(&brioche, "baz!").await, false),
+            ),
+        ],
+    )
+    .await;
+
+    let recipe = Recipe::CollectReferences {
+        recipe: Box::new(WithMeta::without_meta(dir.clone().into())),
+    };
+
+    let output = brioche_test::bake_without_meta(&brioche, recipe).await?;
+
+    assert_eq!(output, dir);
+
+    Ok(())
+}