1 files changed, 16 insertions, 12 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")
 	}
 }