driver: add support for specialized tasks

Allow skipping a TaskDef unless it matches specific arg values.

2 files changed, 18 insertions, 9 deletions

diff --git a/crates/driver/src/context.rs b/crates/driver/src/context.rs
index bd8d985..fb0ea96 100644
--- a/crates/driver/src/context.rs
+++ b/crates/driver/src/context.rs
@@ -267,16 +267,22 @@ impl Context {
 		Some(task_def)
 	}
 
-	fn select_task(tasks: &[TaskDef]) -> &TaskDef {
+	fn match_task(task: &TaskDef, args: &TaskArgs) -> bool {
+		task.arg_match
+			.iter()
+			.all(|(key, value)| args.get(key) == Some(value))
+	}
+
+	fn select_task<'ctx>(tasks: &'ctx [TaskDef], args: &TaskArgs) -> Option<&'ctx TaskDef> {
 		tasks
 			.iter()
+			.filter(|task| Self::match_task(task, args))
 			.max_by(|task1, task2| {
 				deb_version::compare_versions(
 					task1.meta.version.as_deref().unwrap_or_default(),
 					task2.meta.version.as_deref().unwrap_or_default(),
 				)
 			})
-			.unwrap()
 	}
 
 	fn get_with_args<'ctx>(&'ctx self, id: &'ctx TaskID, args: &TaskArgs) -> Result<&TaskDef> {
@@ -284,12 +290,13 @@ impl Context {
 			return Ok(rootfs_task);
 		}
 
-		let tasks = self.tasks.get(id).ok_or(Error {
-			task: id,
-			kind: ErrorKind::TaskNotFound,
-		})?;
-
-		Ok(Self::select_task(tasks))
+		self.tasks
+			.get(id)
+			.and_then(|tasks| Self::select_task(tasks, args))
+			.ok_or(Error {
+				task: id,
+				kind: ErrorKind::TaskNotFound,
+			})
 	}
 
 	pub fn get<'ctx>(&'ctx self, task: &TaskRef<'ctx>) -> Result<&TaskDef> {
diff --git a/crates/driver/src/task.rs b/crates/driver/src/task.rs
index 8f3c7d6..51557ae 100644
--- a/crates/driver/src/task.rs
+++ b/crates/driver/src/task.rs
@@ -5,7 +5,7 @@ use serde::Deserialize;
 use common::{string_hash::StringHash, types::TaskID};
 
 use crate::{
-	args::{ArgMapping, ArgType},
+	args::{ArgMapping, ArgType, TaskArgs},
 	recipe,
 };
 
@@ -87,4 +87,6 @@ pub struct TaskDef {
 	pub output: HashMap<String, Output>,
 	#[serde(flatten)]
 	pub action: Action,
+	#[serde(skip)]
+	pub arg_match: TaskArgs,
 }