Introduce TaskIDRef type

A borrowed version of a TaskID, allowing to avoid allocations in some
places. Tasks are now stored in a two-level map in Context.

4 files changed, 76 insertions, 37 deletions

diff --git a/crates/rebel-common/src/types.rs b/crates/rebel-common/src/types.rs
index 2a06275..d3beb70 100644
--- a/crates/rebel-common/src/types.rs
+++ b/crates/rebel-common/src/types.rs
@@ -13,8 +13,29 @@ pub struct TaskID {
 	pub task: String,
 }
 
+impl TaskID {
+	pub fn as_ref(&self) -> TaskIDRef<'_> {
+		TaskIDRef {
+			recipe: &self.recipe,
+			task: &self.task,
+		}
+	}
+}
+
 impl Display for TaskID {
 	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		self.as_ref().fmt(f)
+	}
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct TaskIDRef<'a> {
+	pub recipe: &'a str,
+	pub task: &'a str,
+}
+
+impl<'a> Display for TaskIDRef<'a> {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 		write!(f, "{}::{}", self.recipe, self.task)
 	}
 }
diff --git a/crates/rebel/src/context.rs b/crates/rebel/src/context.rs
index fa24399..bf16546 100644
--- a/crates/rebel/src/context.rs
+++ b/crates/rebel/src/context.rs
@@ -12,7 +12,7 @@ use std::{
 use rebel_common::{
 	error::{self, Contextualizable},
 	string_hash::ArchiveHash,
-	types::TaskID,
+	types::TaskIDRef,
 };
 use rebel_parse::{self as parse, TaskFlags};
 
@@ -32,7 +32,7 @@ pub enum ErrorKind<'a> {
 
 #[derive(Debug, Clone, Copy)]
 pub struct Error<'a> {
-	pub task: &'a TaskID,
+	pub task: TaskIDRef<'a>,
 	pub kind: ErrorKind<'a>,
 }
 
@@ -65,7 +65,7 @@ pub type Result<'a, T> = result::Result<T, Error<'a>>;
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct TaskRef<'ctx> {
-	pub id: &'ctx TaskID,
+	pub id: TaskIDRef<'ctx>,
 	pub args: Rc<TaskArgs>,
 }
 
@@ -144,12 +144,15 @@ fn platform_relation(args: &TaskArgs, from: &str, to: &str) -> Option<PlatformRe
 pub struct Context {
 	platforms: HashMap<String, Arg>,
 	globals: TaskArgs,
-	tasks: HashMap<TaskID, Vec<TaskDef>>,
+	tasks: HashMap<String, HashMap<String, Vec<TaskDef>>>,
 	rootfs: (ArchiveHash, String),
 }
 
 impl Context {
-	pub fn new(mut tasks: HashMap<TaskID, Vec<TaskDef>>, pins: Pins) -> error::Result<Self> {
+	pub fn new(
+		mut tasks: HashMap<String, HashMap<String, Vec<TaskDef>>>,
+		pins: Pins,
+	) -> error::Result<Self> {
 		let platforms: HashMap<_, _> = [
 			arg(
 				"build",
@@ -212,7 +215,7 @@ impl Context {
 	}
 
 	fn add_rootfs_tasks(
-		tasks: &mut HashMap<TaskID, Vec<TaskDef>>,
+		tasks: &mut HashMap<String, HashMap<String, Vec<TaskDef>>>,
 		provides: Vec<pin::Provides>,
 		globals: &TaskArgs,
 	) -> error::Result<()> {
@@ -255,10 +258,9 @@ impl Context {
 			task_def.priority = i32::MAX;
 
 			tasks
-				.entry(TaskID {
-					recipe: recipe.to_string(),
-					task: task.to_string(),
-				})
+				.entry(recipe)
+				.or_default()
+				.entry(task)
 				.or_default()
 				.push(task_def);
 		}
@@ -293,9 +295,12 @@ impl Context {
 			.max_by(|task1, task2| Self::compare_tasks(task1, task2))
 	}
 
-	fn get_with_args<'a>(&self, id: &'a TaskID, args: &TaskArgs) -> Result<'a, &TaskDef> {
-		self.tasks
-			.get(id)
+	fn get_by_ref(&self, id: TaskIDRef) -> Option<&[TaskDef]> {
+		Some(self.tasks.get(id.recipe)?.get(id.task)?)
+	}
+
+	fn get_with_args<'a>(&self, id: TaskIDRef<'a>, args: &TaskArgs) -> Result<'a, &TaskDef> {
+		self.get_by_ref(id)
 			.and_then(|tasks| Self::select_task(tasks, args))
 			.ok_or(Error {
 				task: id,
@@ -307,7 +312,7 @@ impl Context {
 		self.get_with_args(task.id, task.args.as_ref())
 	}
 
-	fn task_ref<'ctx>(&'ctx self, id: &'ctx TaskID, args: &TaskArgs) -> Result<TaskRef> {
+	fn task_ref<'ctx>(&'ctx self, id: TaskIDRef<'ctx>, args: &TaskArgs) -> Result<TaskRef> {
 		let task_def = self.get_with_args(id, args)?;
 
 		let mut arg_def: HashMap<_, _> = task_def.args.iter().map(|(k, &v)| (k, v)).collect();
@@ -362,14 +367,22 @@ impl Context {
 			.ok()
 			.context("Invalid task syntax")?;
 
-		let recipe = parsed.id.recipe.to_string();
-		let task = parsed.id.task.to_string();
+		let id = TaskIDRef {
+			recipe: parsed.id.recipe,
+			task: parsed.id.task,
+		};
 
-		let id = TaskID { recipe, task };
-		let (ctx_id, _) = self
+		let (ctx_recipe, recipe_tasks) = self
 			.tasks
-			.get_key_value(&id)
+			.get_key_value(id.recipe)
+			.with_context(|| format!("Task {} not found", id))?;
+		let (ctx_task, _) = recipe_tasks
+			.get_key_value(id.task)
 			.with_context(|| format!("Task {} not found", id))?;
+		let ctx_id = TaskIDRef {
+			recipe: ctx_recipe,
+			task: ctx_task,
+		};
 
 		let mut args = self.globals.clone();
 
@@ -397,7 +410,7 @@ impl Context {
 	}
 
 	fn map_args<'ctx, 'args>(
-		task: &'ctx TaskID,
+		task: TaskIDRef<'ctx>,
 		mapping: &'ctx ArgMapping,
 		args: &'args TaskArgs,
 		build_dep: bool,
@@ -425,8 +438,9 @@ impl Context {
 	}
 
 	fn parent_ref<'ctx>(&'ctx self, dep: &'ctx ParentDep, args: &TaskArgs) -> Result<TaskRef> {
-		let mapped_args = Context::map_args(&dep.dep.id, &dep.dep.args, args, false)?;
-		self.task_ref(&dep.dep.id, mapped_args.as_ref())
+		let id = dep.dep.id.as_ref();
+		let mapped_args = Context::map_args(id, &dep.dep.args, args, false)?;
+		self.task_ref(id, mapped_args.as_ref())
 	}
 
 	pub fn output_ref<'ctx>(
@@ -435,9 +449,10 @@ impl Context {
 		args: &TaskArgs,
 		build_dep: bool,
 	) -> Result<OutputRef<'ctx>> {
-		let mapped_args = Context::map_args(&dep.dep.id, &dep.dep.args, args, build_dep)?;
+		let id = dep.dep.id.as_ref();
+		let mapped_args = Context::map_args(id, &dep.dep.args, args, build_dep)?;
 		Ok(OutputRef {
-			task: self.task_ref(&dep.dep.id, mapped_args.as_ref())?,
+			task: self.task_ref(id, mapped_args.as_ref())?,
 			output: &dep.output,
 		})
 	}
diff --git a/crates/rebel/src/recipe.rs b/crates/rebel/src/recipe.rs
index fd5367a..24b4bff 100644
--- a/crates/rebel/src/recipe.rs
+++ b/crates/rebel/src/recipe.rs
@@ -74,24 +74,25 @@ fn recipe_name(path: &Path) -> Option<&str> {
 }
 
 fn handle_recipe_tasks(
-	tasks: &mut HashMap<TaskID, Vec<TaskDef>>,
+	tasks: &mut HashMap<String, HashMap<String, Vec<TaskDef>>>,
 	recipe_tasks: HashMap<String, TaskDef>,
 	meta: &TaskMeta,
 ) {
+	let task_map = match tasks.get_mut(&meta.recipe) {
+		Some(task_map) => task_map,
+		None => tasks.entry(meta.recipe.clone()).or_default(),
+	};
+
 	for (label, mut task) in recipe_tasks {
-		let task_id = TaskID {
-			recipe: meta.recipe.clone(),
-			task: label,
-		};
 		task.meta = meta.clone();
-		tasks.entry(task_id).or_default().push(task);
+		task_map.entry(label).or_default().push(task);
 	}
 }
 
 fn read_recipe_tasks(
 	path: &Path,
 	basename: &str,
-	tasks: &mut HashMap<TaskID, Vec<TaskDef>>,
+	tasks: &mut HashMap<String, HashMap<String, Vec<TaskDef>>>,
 ) -> Result<RecipeMeta> {
 	let recipe_def = CURRENT_RECIPE.set(basename, || read_yaml::<Recipe>(path))?;
 
@@ -120,7 +121,7 @@ fn read_subrecipe_tasks(
 	basename: &str,
 	recipename: &str,
 	recipe_meta: &RecipeMeta,
-	tasks: &mut HashMap<TaskID, Vec<TaskDef>>,
+	tasks: &mut HashMap<String, HashMap<String, Vec<TaskDef>>>,
 ) -> Result<()> {
 	let recipe = format!("{basename}/{recipename}");
 	let recipe_def = CURRENT_RECIPE.set(&recipe, || read_yaml::<Subrecipe>(path))?;
@@ -140,8 +141,10 @@ fn read_subrecipe_tasks(
 	Ok(())
 }
 
-pub fn read_recipes<P: AsRef<Path>>(path: P) -> Result<HashMap<TaskID, Vec<TaskDef>>> {
-	let mut tasks = HashMap::<TaskID, Vec<TaskDef>>::new();
+pub fn read_recipes<P: AsRef<Path>>(
+	path: P,
+) -> Result<HashMap<String, HashMap<String, Vec<TaskDef>>>> {
+	let mut tasks = HashMap::<String, HashMap<String, Vec<TaskDef>>>::new();
 	let mut recipe_metas = HashMap::<String, RecipeMeta>::new();
 
 	for entry in WalkDir::new(path)
diff --git a/crates/rebel/src/resolve.rs b/crates/rebel/src/resolve.rs
index 896903f..27a6d5c 100644
--- a/crates/rebel/src/resolve.rs
+++ b/crates/rebel/src/resolve.rs
@@ -2,7 +2,7 @@ use std::collections::{HashMap, HashSet};
 use std::fmt;
 use std::rc::Rc;
 
-use rebel_common::types::TaskID;
+use rebel_common::types::TaskIDRef;
 
 use crate::args::TaskArgs;
 use crate::context::{self, Context, OutputRef, TaskRef};
@@ -38,8 +38,8 @@ impl<'ctx> From<&TaskRef<'ctx>> for DepChain<'ctx> {
 	}
 }
 
-impl<'ctx> From<&'ctx TaskID> for DepChain<'ctx> {
-	fn from(id: &'ctx TaskID) -> Self {
+impl<'ctx> From<TaskIDRef<'ctx>> for DepChain<'ctx> {
+	fn from(id: TaskIDRef<'ctx>) -> Self {
 		TaskRef {
 			id,
 			args: Rc::new(TaskArgs::default()),