driver: context: change get()/index() argument to TaskRef

With the upcoming changes, we may need to return different TaskDefs
depending on the passed arguments.

3 files changed, 20 insertions, 16 deletions

diff --git a/crates/driver/src/context.rs b/crates/driver/src/context.rs
index ac7162d..25b8e38 100644
--- a/crates/driver/src/context.rs
+++ b/crates/driver/src/context.rs
@@ -178,15 +178,19 @@ impl Context {
 		}
 	}
 
-	pub fn get<'ctx>(&'ctx self, id: &'ctx TaskID) -> Result<&TaskDef> {
+	fn get_with_args<'ctx>(&'ctx self, id: &'ctx TaskID, _args: &TaskArgs) -> Result<&TaskDef> {
 		self.tasks.get(id).ok_or(Error {
 			task: id,
 			kind: ErrorKind::TaskNotFound,
 		})
 	}
 
+	pub fn get<'ctx>(&'ctx self, task: &TaskRef<'ctx>) -> Result<&TaskDef> {
+		self.get_with_args(task.id, task.args.as_ref())
+	}
+
 	fn task_ref<'ctx>(&'ctx self, id: &'ctx TaskID, args: &TaskArgs) -> Result<TaskRef> {
-		let task_def = self.get(id)?;
+		let task_def = self.get_with_args(id, args)?;
 
 		let mut arg_def: HashMap<_, _> = task_def.args.iter().map(|(k, &v)| (k, v)).collect();
 		for (key, arg) in &self.globals {
@@ -322,7 +326,7 @@ impl Context {
 		&'ctx self,
 		task_ref: &TaskRef<'ctx>,
 	) -> Result<Option<TaskRef>> {
-		let task = self.get(task_ref.id)?;
+		let task = self.get(task_ref)?;
 		let inherit = match &task.inherit {
 			Some(inherit) => inherit,
 			None => return Ok(None),
@@ -360,13 +364,13 @@ impl Context {
 		let mut allow_noinherit = true;
 
 		for current in self.inherit_iter(task_ref) {
-			let task_ref = current?;
-			let task = self.get(task_ref.id)?;
+			let current_ref = current?;
+			let task = self.get(&current_ref)?;
 			let entries = task
 				.build_depends
 				.iter()
 				.filter(|dep| allow_noinherit || !dep.noinherit)
-				.map(|dep| self.output_ref(dep, &task_ref.args, true))
+				.map(|dep| self.output_ref(dep, &current_ref.args, true))
 				.collect::<Result<Vec<_>>>()?;
 			ret.extend(entries);
 
@@ -384,13 +388,13 @@ impl Context {
 		let mut allow_noinherit = true;
 
 		for current in self.inherit_iter(task_ref) {
-			let task_ref = current?;
-			let task = self.get(task_ref.id)?;
+			let current_ref = current?;
+			let task = self.get(&current_ref)?;
 			let entries = task
 				.depends
 				.iter()
 				.filter(|dep| allow_noinherit || !dep.noinherit)
-				.map(|dep| self.output_ref(dep, &task_ref.args, false))
+				.map(|dep| self.output_ref(dep, &current_ref.args, false))
 				.collect::<Result<Vec<_>>>()?;
 			ret.extend(entries);
 
@@ -432,10 +436,10 @@ impl Context {
 	}
 }
 
-impl Index<&TaskID> for Context {
+impl<'ctx> Index<&TaskRef<'ctx>> for &'ctx Context {
 	type Output = TaskDef;
 
-	fn index(&self, index: &TaskID) -> &TaskDef {
-		self.tasks.get(index).expect("Invalid TaskID")
+	fn index(&self, index: &TaskRef<'ctx>) -> &TaskDef {
+		self.get(index).expect("Invalid TaskRef")
 	}
 }
diff --git a/crates/driver/src/driver.rs b/crates/driver/src/driver.rs
index bafba1b..ca93b0f 100644
--- a/crates/driver/src/driver.rs
+++ b/crates/driver/src/driver.rs
@@ -40,7 +40,7 @@ impl<'ctx> CompletionState<'ctx> {
 	}
 
 	fn fetch_deps(&self, task: &TaskRef<'ctx>) -> Result<Vec<Dependency>> {
-		let task_def = &self.ctx[task.id];
+		let task_def = &self.ctx[task];
 		task_def
 			.fetch
 			.iter()
@@ -267,7 +267,7 @@ impl<'ctx> Driver<'ctx> {
 	}
 
 	fn spawn_task(&self, task_ref: &TaskRef<'ctx>, runner: &Runner) -> Result<Option<UnixStream>> {
-		let task_def = &self.state.ctx[task_ref.id];
+		let task_def = &self.state.ctx[task_ref];
 		if task_def.action.is_empty() {
 			return Ok(None);
 		}
diff --git a/crates/driver/src/resolve.rs b/crates/driver/src/resolve.rs
index 338ce3f..9ff5708 100644
--- a/crates/driver/src/resolve.rs
+++ b/crates/driver/src/resolve.rs
@@ -133,7 +133,7 @@ where
 		}
 
 		let task = &dep.task;
-		let task_def = match ctx.get(task.id) {
+		let task_def = match ctx.get(task) {
 			Ok(task) => task,
 			Err(err) => return vec![err.into()],
 		};
@@ -240,7 +240,7 @@ impl<'ctx> Resolver<'ctx> {
 			None => (),
 		}
 
-		let task_def = match self.ctx.get(task.id) {
+		let task_def = match self.ctx.get(task) {
 			Ok(task_def) => task_def,
 			Err(err) => return vec![err.into()],
 		};