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.

1 files changed, 39 insertions, 24 deletions

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,
 		})
 	}