Fix --sync missing some recipes (#122)

* Update `get_recipes` to take a hashmap to insert into

* Fix sync missing some recipes by improving `recipe_references`

* Update failing test in `sync_to_registry.rs`

crates/brioche-core/src/recipe.rs

@@ -167,9 +167,8 @@ impl Recipe {
 pub async fn get_recipes(
     brioche: &Brioche,
     recipe_hashes: impl IntoIterator<Item = RecipeHash>,
-) -> anyhow::Result<HashMap<RecipeHash, Recipe>> {
-    let mut recipes = HashMap::new();
-
+    recipes: &mut HashMap<RecipeHash, Recipe>,
+) -> anyhow::Result<()> {
     let cached_recipes = brioche.cached_recipes.read().await;
     let mut uncached_recipes = HashSet::new();
 
@@ -189,7 +188,7 @@ pub async fn get_recipes(
 
     // Return early if we have no uncached recipess to fetch
     if uncached_recipes.is_empty() {
-        return Ok(recipes);
+        return Ok(());
     }
 
     let mut db_conn = brioche.db_conn.lock().await;
@@ -251,11 +250,12 @@ pub async fn get_recipes(
         anyhow::bail!("recipes not found: {uncached_recipes:?}");
     }
 
-    Ok(recipes)
+    Ok(())
 }
 
 pub async fn get_recipe(brioche: &Brioche, recipe_hash: RecipeHash) -> anyhow::Result<Recipe> {
-    let mut recipes = get_recipes(brioche, [recipe_hash]).await?;
+    let mut recipes = HashMap::new();
+    get_recipes(brioche, [recipe_hash], &mut recipes).await?;
     let recipe = recipes
         .remove(&recipe_hash)
         .expect("recipe not returned in collection");
@@ -719,7 +719,8 @@ impl Directory {
     }
 
     pub async fn entries(&self, brioche: &Brioche) -> anyhow::Result<BTreeMap<BString, Artifact>> {
-        let entry_recipes = get_recipes(brioche, self.entries.values().copied()).await?;
+        let mut entry_recipes = HashMap::new();
+        get_recipes(brioche, self.entries.values().copied(), &mut entry_recipes).await?;
 
         let entries = self
             .entries

crates/brioche-core/src/references.rs

@@ -17,6 +17,21 @@ use crate::{
     Brioche,
 };
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum ReferencedRecipe {
+    Recipe(Recipe),
+    RecipeHash(RecipeHash),
+}
+
+impl ReferencedRecipe {
+    pub fn hash(&self) -> RecipeHash {
+        match self {
+            ReferencedRecipe::Recipe(recipe) => recipe.hash(),
+            ReferencedRecipe::RecipeHash(hash) => *hash,
+        }
+    }
+}
+
 #[derive(Debug, Default, Clone, serde::Serialize, serde::Deserialize)]
 pub struct RecipeReferences {
     pub blobs: HashSet<BlobHash>,
@@ -26,18 +41,32 @@ pub struct RecipeReferences {
 pub async fn recipe_references(
     brioche: &Brioche,
     references: &mut RecipeReferences,
-    recipes: impl IntoIterator<Item = RecipeHash>,
+    recipes: impl IntoIterator<Item = ReferencedRecipe>,
 ) -> anyhow::Result<()> {
     let mut unvisited = VecDeque::from_iter(recipes);
 
     loop {
-        unvisited.retain(|recipe| !references.recipes.contains_key(recipe));
+        unvisited.retain(|recipe| !references.recipes.contains_key(&recipe.hash()));
 
         if unvisited.is_empty() {
             break;
         }
 
-        let recipes = crate::recipe::get_recipes(brioche, unvisited.drain(..)).await?;
+        let mut recipe_hashes = vec![];
+        let mut recipes = HashMap::new();
+
+        for referenced in unvisited.drain(..) {
+            match referenced {
+                ReferencedRecipe::Recipe(recipe) => {
+                    recipes.insert(recipe.hash(), recipe);
+                }
+                ReferencedRecipe::RecipeHash(hash) => {
+                    recipe_hashes.push(hash);
+                }
+            }
+        }
+
+        crate::recipe::get_recipes(brioche, recipe_hashes, &mut recipes).await?;
 
         for recipe in recipes.values() {
             unvisited.extend(referenced_recipes(recipe));
@@ -64,7 +93,7 @@ pub async fn project_references(
     project_hashes: impl IntoIterator<Item = ProjectHash>,
 ) -> anyhow::Result<()> {
     let mut unvisited = VecDeque::from_iter(project_hashes);
-    let mut new_recipes = HashSet::<RecipeHash>::new();
+    let mut new_recipes = HashSet::<ReferencedRecipe>::new();
 
     loop {
         unvisited.retain(|project_hash| !references.projects.contains_key(project_hash));
@@ -101,7 +130,7 @@ pub async fn project_references(
                     .as_ref()
                     .with_context(|| format!("static not loaded for module {module_path}"))?;
                 let static_recipe_hash = static_.output_recipe_hash(output)?;
-                new_recipes.insert(static_recipe_hash);
+                new_recipes.insert(ReferencedRecipe::RecipeHash(static_recipe_hash));
             }
         }
     }
@@ -135,17 +164,24 @@ pub fn referenced_blobs(recipe: &Recipe) -> Vec<BlobHash> {
     }
 }
 
-pub fn referenced_recipes(recipe: &Recipe) -> Vec<RecipeHash> {
+pub fn referenced_recipes(recipe: &Recipe) -> Vec<ReferencedRecipe> {
     match recipe {
         Recipe::File {
             resources,
             content_blob: _,
             executable: _,
-        } => referenced_recipes(resources),
-        Recipe::Directory(directory) => directory.entry_hashes().values().copied().collect(),
+        } => vec![ReferencedRecipe::Recipe(resources.value.clone())],
+        Recipe::Directory(directory) => directory
+            .entry_hashes()
+            .values()
+            .copied()
+            .map(ReferencedRecipe::RecipeHash)
+            .collect(),
         Recipe::Symlink { .. } => vec![],
         Recipe::Download(_) => vec![],
-        Recipe::Unarchive(unarchive) => referenced_recipes(&unarchive.file),
+        Recipe::Unarchive(unarchive) => {
+            vec![ReferencedRecipe::Recipe(unarchive.file.value.clone())]
+        }
         Recipe::Process(process) => {
             let ProcessRecipe {
                 command,
@@ -163,26 +199,28 @@ pub fn referenced_recipes(recipe: &Recipe) -> Vec<RecipeHash> {
 
             templates
                 .flat_map(|template| &template.components)
-                .flat_map(|component| match component {
-                    ProcessTemplateComponent::Input { recipe } => referenced_recipes(recipe),
+                .filter_map(|component| match component {
+                    ProcessTemplateComponent::Input { recipe } => {
+                        Some(ReferencedRecipe::Recipe(recipe.value.clone()))
+                    }
                     ProcessTemplateComponent::Literal { .. }
                     | ProcessTemplateComponent::OutputPath
                     | ProcessTemplateComponent::ResourceDir
                     | ProcessTemplateComponent::InputResourceDirs
                     | ProcessTemplateComponent::HomeDir
                     | ProcessTemplateComponent::WorkDir
-                    | ProcessTemplateComponent::TempDir => vec![],
+                    | ProcessTemplateComponent::TempDir => None,
                 })
                 .chain(
                     dependencies
                         .iter()
-                        .flat_map(|dep| referenced_recipes(&dep.value)),
+                        .map(|recipe| ReferencedRecipe::Recipe(recipe.value.clone())),
                 )
-                .chain(referenced_recipes(work_dir))
+                .chain([ReferencedRecipe::Recipe(work_dir.value.clone())])
                 .chain(
                     output_scaffold
-                        .iter()
-                        .flat_map(|recipe| referenced_recipes(recipe)),
+                        .as_ref()
+                        .map(|recipe| ReferencedRecipe::Recipe(recipe.value.clone())),
                 )
                 .collect()
         }
@@ -198,71 +236,79 @@ pub fn referenced_recipes(recipe: &Recipe) -> Vec<RecipeHash> {
                 networking: _,
             } = process;
 
-            let work_dir = Recipe::from(work_dir.clone());
-            let output_scaffold = output_scaffold
-                .as_ref()
-                .map(|artifact| Recipe::from((**artifact).clone()));
-
             let templates = [command].into_iter().chain(args).chain(env.values());
 
             templates
                 .flat_map(|template| &template.components)
-                .flat_map(|component| match component {
+                .filter_map(|component| match component {
                     CompleteProcessTemplateComponent::Input { artifact } => {
-                        let recipe = Recipe::from(artifact.value.clone());
-                        referenced_recipes(&recipe)
+                        Some(ReferencedRecipe::Recipe(artifact.value.clone().into()))
                     }
                     CompleteProcessTemplateComponent::Literal { .. }
                     | CompleteProcessTemplateComponent::OutputPath
                     | CompleteProcessTemplateComponent::ResourceDir
                     | CompleteProcessTemplateComponent::InputResourceDirs
                     | CompleteProcessTemplateComponent::HomeDir
                     | CompleteProcessTemplateComponent::WorkDir
-                    | CompleteProcessTemplateComponent::TempDir => vec![],
+                    | CompleteProcessTemplateComponent::TempDir => None,
                 })
-                .chain(referenced_recipes(&work_dir))
-                .chain(output_scaffold.iter().flat_map(referenced_recipes))
+                .chain([ReferencedRecipe::Recipe(Recipe::Directory(
+                    work_dir.clone(),
+                ))])
+                .chain(
+                    output_scaffold
+                        .as_deref()
+                        .map(|recipe| ReferencedRecipe::Recipe(recipe.clone().into())),
+                )
                 .collect()
         }
         Recipe::CreateFile {
             content: _,
             executable: _,
             resources,
-        } => referenced_recipes(resources),
+        } => vec![ReferencedRecipe::Recipe(resources.value.clone())],
         Recipe::CreateDirectory(directory) => directory
             .entries
             .values()
-            .flat_map(|entry| referenced_recipes(entry))
+            .map(|entry| ReferencedRecipe::Recipe(entry.value.clone()))
             .collect(),
-        Recipe::Cast { recipe, to: _ } => referenced_recipes(recipe),
+        Recipe::Cast { recipe, to: _ } => vec![ReferencedRecipe::Recipe(recipe.value.clone())],
         Recipe::Merge { directories } => directories
             .iter()
-            .flat_map(|dir| referenced_recipes(dir))
+            .map(|dir| ReferencedRecipe::Recipe(dir.value.clone()))
             .collect(),
         Recipe::Peel {
             directory,
             depth: _,
-        } => referenced_recipes(directory),
-        Recipe::Get { directory, path: _ } => referenced_recipes(directory),
+        } => vec![ReferencedRecipe::Recipe(directory.value.clone())],
+        Recipe::Get { directory, path: _ } => {
+            vec![ReferencedRecipe::Recipe(directory.value.clone())]
+        }
         Recipe::Insert {
             directory,
             path: _,
             recipe,
-        } => referenced_recipes(directory)
+        } => [ReferencedRecipe::Recipe(directory.value.clone())]
             .into_iter()
-            .chain(recipe.iter().flat_map(|recipe| referenced_recipes(recipe)))
+            .chain(
+                recipe
+                    .as_ref()
+                    .map(|recipe| ReferencedRecipe::Recipe(recipe.value.clone())),
+            )
             .collect(),
         Recipe::Glob {
             directory,
             patterns: _,
-        } => referenced_recipes(directory),
+        } => vec![ReferencedRecipe::Recipe(directory.value.clone())],
         Recipe::SetPermissions {
             file,
             executable: _,
-        } => referenced_recipes(file),
-        Recipe::Proxy(proxy) => vec![proxy.recipe],
-        Recipe::CollectReferences { recipe } => referenced_recipes(recipe),
-        Recipe::Sync { recipe } => referenced_recipes(recipe),
+        } => vec![ReferencedRecipe::Recipe(file.value.clone())],
+        Recipe::Proxy(proxy) => vec![ReferencedRecipe::RecipeHash(proxy.recipe)],
+        Recipe::CollectReferences { recipe } => {
+            vec![ReferencedRecipe::Recipe(recipe.value.clone())]
+        }
+        Recipe::Sync { recipe } => vec![ReferencedRecipe::Recipe(recipe.value.clone())],
     }
 }
 

crates/brioche-core/src/registry.rs

@@ -12,6 +12,7 @@ use crate::{
     blob::BlobHash,
     project::{Project, ProjectHash},
     recipe::{Artifact, Recipe, RecipeHash},
+    references::ReferencedRecipe,
     Brioche,
 };
 
@@ -502,7 +503,14 @@ pub async fn fetch_recipes_deep(
 
         for recipe in &new_recipes {
             let referenced_recipes = crate::references::referenced_recipes(recipe);
-            pending_recipes.extend(referenced_recipes);
+            let referenced_recipe_hashes =
+                referenced_recipes
+                    .into_iter()
+                    .filter_map(|reference| match reference {
+                        ReferencedRecipe::RecipeHash(hash) => Some(hash),
+                        ReferencedRecipe::Recipe(_) => None,
+                    });
+            pending_recipes.extend(referenced_recipe_hashes);
         }
 
         checked_recipes.extend(new_recipes.iter().map(|recipe| recipe.hash()));

crates/brioche-core/src/sync.rs

@@ -4,7 +4,7 @@ use human_repr::HumanDuration;
 
 use crate::{
     project::ProjectHash,
-    references::{ProjectReferences, RecipeReferences},
+    references::{ProjectReferences, RecipeReferences, ReferencedRecipe},
     Brioche,
 };
 
@@ -55,9 +55,12 @@ pub async fn sync_bakes(
 
     let mut sync_references = RecipeReferences::default();
 
-    let recipe_hashes = bakes
-        .iter()
-        .flat_map(|(input, output)| [input.hash(), output.hash()]);
+    let recipe_hashes = bakes.iter().flat_map(|(input, output)| {
+        [
+            ReferencedRecipe::Recipe(input.clone()),
+            ReferencedRecipe::Recipe(output.clone().into()),
+        ]
+    });
     crate::references::recipe_references(brioche, &mut sync_references, recipe_hashes).await?;
 
     let num_recipe_refs = sync_references.recipes.len();

crates/brioche-core/tests/recipe.rs

@@ -114,9 +114,13 @@ async fn test_recipe_save_multiple() -> anyhow::Result<()> {
         brioche_core::recipe::save_recipes(&brioche, [file_2.clone(), file_3.clone()]).await?;
     assert_eq!(new_artifacts, 1);
 
-    let results =
-        brioche_core::recipe::get_recipes(&brioche, [file_1.hash(), file_2.hash(), file_3.hash()])
-            .await?;
+    let mut results = HashMap::new();
+    brioche_core::recipe::get_recipes(
+        &brioche,
+        [file_1.hash(), file_2.hash(), file_3.hash()],
+        &mut results,
+    )
+    .await?;
 
     assert_eq!(
         results,

crates/brioche-core/tests/sync_to_registry.rs

@@ -22,8 +22,11 @@ async fn test_sync_to_registry_process_and_complete_process() -> anyhow::Result<
         Recipe::CompleteProcess(complete_process_recipe.clone()).hash();
 
     // Create a mocked output for the complete_process recipe
+    let output_resource_dir = brioche_test_support::empty_dir_value();
+    let output_resource_dir_hash = Recipe::Directory(output_resource_dir.clone()).hash();
     let dummy_blob = brioche_test_support::blob(&brioche, "dummy value").await;
-    let mocked_output = brioche_test_support::file(dummy_blob, false);
+    let mocked_output =
+        brioche_test_support::file_with_resources(dummy_blob, false, output_resource_dir);
     brioche_test_support::mock_bake(
         &brioche,
         &Recipe::CompleteProcess(complete_process_recipe.clone()),
@@ -84,6 +87,7 @@ async fn test_sync_to_registry_process_and_complete_process() -> anyhow::Result<
                 process_recipe_hash,
                 complete_process_recipe_hash,
                 mocked_output.hash(),
+                output_resource_dir_hash,
             ])?)
             .create(),
     );