Add convenience methods and trait implementations to TaskArgs

3 files changed, 60 insertions, 28 deletions

diff --git a/src/args.rs b/src/args.rs
index 1b58dd4..cd4a360 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -1,4 +1,9 @@
-use std::{collections::HashMap, hash, rc::Rc};
+use std::{
+	collections::{hash_map, HashMap},
+	hash,
+	iter::FromIterator,
+	rc::Rc,
+};
 
 use enum_kinds::EnumKind;
 use serde::{Deserialize, Serialize};
@@ -29,7 +34,41 @@ impl From<String> for Arg {
 }
 
 #[derive(Clone, Debug, Serialize, PartialEq, Eq, Default)]
-pub struct TaskArgs(pub HashMap<String, Rc<Arg>>);
+pub struct TaskArgs(HashMap<String, Rc<Arg>>);
+
+impl TaskArgs {
+	pub fn get(&self, key: &str) -> Option<&Rc<Arg>> {
+		self.0.get(key)
+	}
+
+	pub fn set(&mut self, key: &str, value: Option<Rc<Arg>>) {
+		if let Some(v) = value {
+			self.0.insert(key.to_string(), v);
+		} else {
+			self.0.remove(key);
+		}
+	}
+
+	pub fn iter(&self) -> hash_map::Iter<String, Rc<Arg>> {
+		self.into_iter()
+	}
+}
+
+impl FromIterator<(String, Rc<Arg>)> for TaskArgs {
+	fn from_iter<T: IntoIterator<Item = (String, Rc<Arg>)>>(iter: T) -> Self {
+		TaskArgs(HashMap::from_iter(iter))
+	}
+}
+
+impl<'a> IntoIterator for &'a TaskArgs {
+	type Item = (&'a String, &'a Rc<Arg>);
+
+	type IntoIter = hash_map::Iter<'a, String, Rc<Arg>>;
+
+	fn into_iter(self) -> Self::IntoIter {
+		self.0.iter()
+	}
+}
 
 impl hash::Hash for TaskArgs {
 	fn hash<H: hash::Hasher>(&self, _state: &mut H) {
diff --git a/src/context.rs b/src/context.rs
index 0d72669..13e3c8a 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -3,6 +3,7 @@ use std::{
 	collections::{HashMap, HashSet},
 	fmt::Display,
 	hash::Hash,
+	iter::FromIterator,
 	ops::Index,
 	rc::Rc,
 	result,
@@ -101,21 +102,16 @@ impl Context {
 			prefix: "/usr".to_string(),
 			cross_compile: "/opt/toolchain/bin/aarch64-linux-gnu-".to_string(),
 		};
-		let globals = TaskArgs(
-			IntoIterator::into_iter([
-				arg("build", build_platform),
-				arg("workdir", paths::abs(paths::TASK_WORKDIR)),
-				arg("dldir", paths::abs(paths::TASK_DLDIR)),
-				arg("destdir", paths::abs(paths::TASK_DESTDIR)),
-				arg("sysroot", paths::abs(paths::TASK_SYSROOT)),
-			])
-			.collect(),
-		);
+		let globals = TaskArgs::from_iter([
+			arg("build", build_platform),
+			arg("workdir", paths::abs(paths::TASK_WORKDIR)),
+			arg("dldir", paths::abs(paths::TASK_DLDIR)),
+			arg("destdir", paths::abs(paths::TASK_DESTDIR)),
+			arg("sysroot", paths::abs(paths::TASK_SYSROOT)),
+		]);
 
 		let mut default_args = globals.clone();
-		default_args
-			.0
-			.insert("host".to_string(), host_platform.into());
+		default_args.set("host", Some(host_platform.into()));
 
 		Context {
 			tasks,
@@ -135,7 +131,7 @@ impl Context {
 		let task_def = self.get(id)?;
 
 		let mut arg_def: HashMap<_, _> = task_def.args.iter().map(|(k, &v)| (k, v)).collect();
-		for (key, arg) in &self.globals.0 {
+		for (key, arg) in &self.globals {
 			// TODO: Handle conflicts between explicit args and globals
 			arg_def.insert(key, ArgType::from(arg.as_ref()));
 		}
@@ -143,9 +139,9 @@ impl Context {
 		let mut new_args = TaskArgs::default();
 
 		for (key, typ) in arg_def {
-			if let Some(arg) = args.0.get(key) {
+			if let Some(arg) = args.get(key) {
 				if ArgType::from(arg.as_ref()) == typ {
-					new_args.0.insert(key.clone(), arg.clone());
+					new_args.set(key, Some(arg.clone()));
 					continue;
 				}
 			}
@@ -184,18 +180,16 @@ impl Context {
 		let mut ret = args.clone();
 
 		if build_dep {
-			ret.0.insert("host".to_string(), args.0["build"].clone());
-			if let Some(host) = args.0.get("host") {
-				ret.0.insert("target".to_string(), host.clone());
-			}
+			ret.set("host", args.get("build").map(Rc::clone));
+			ret.set("target", args.get("host").map(Rc::clone));
 		}
 
 		for (to, from) in &mapping.0 {
-			let value = args.0.get(from).ok_or_else(|| Error {
+			let value = args.get(from).ok_or_else(|| Error {
 				task,
 				kind: ErrorKind::InvalidArgRef(to),
 			})?;
-			ret.0.insert(to.clone(), value.clone());
+			ret.set(to, Some(value.clone()));
 		}
 
 		Ok(Cow::Owned(ret))
@@ -302,11 +296,11 @@ impl Context {
 	}
 
 	pub fn in_rootfs<'ctx>(&'ctx self, output: &OutputRef<'ctx>) -> bool {
-		let build = &self.globals.0["build"];
-		if output.task.args.0.get("host") != Some(build) {
+		let build = self.globals.get("build").unwrap();
+		if output.task.args.get("host") != Some(build) {
 			return false;
 		}
-		if let Some(target) = output.task.args.0.get("target") {
+		if let Some(target) = output.task.args.get("target") {
 			if target != build {
 				return false;
 			}
diff --git a/src/executor.rs b/src/executor.rs
index da651c0..8e87fe1 100644
--- a/src/executor.rs
+++ b/src/executor.rs
@@ -205,7 +205,6 @@ impl<'ctx> Executor<'ctx> {
 			id: task_ref.id.clone(),
 			args: task_ref
 				.args
-				.0
 				.iter()
 				.map(|(k, v)| (k.clone(), serde_json::to_value(v).unwrap()))
 				.collect(),