path: root/crates/driver/src/context.rs

use std::{
	borrow::Cow,
	collections::{HashMap, HashSet},
	fmt::Display,
	hash::Hash,
	ops::Index,
	rc::Rc,
	result,
};

use lazy_static::lazy_static;
use regex::Regex;

use common::{
	error::{self, Contextualizable},
	types::TaskID,
};

use crate::{args::*, paths, task::*};

#[derive(Debug, Clone, Copy)]
pub enum ErrorKind<'ctx> {
	TaskNotFound,
	InvalidArgument(&'ctx str),
	InvalidArgRef(&'ctx str),
}

#[derive(Debug, Clone, Copy)]
pub struct Error<'ctx> {
	pub task: &'ctx TaskID,
	pub kind: ErrorKind<'ctx>,
}

impl<'ctx> Display for Error<'ctx> {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		let Error { task, kind } = self;
		match kind {
			ErrorKind::TaskNotFound => write!(f, "Task '{}' not found", task),
			ErrorKind::InvalidArgument(arg) => write!(
				f,
				"Invalid or missing argument '{}' for task '{}'",
				arg, task
			),
			ErrorKind::InvalidArgRef(arg) => write!(
				f,
				"Invalid reference for argument '{}' of task '{}'",
				arg, task
			),
		}
	}
}

impl<'ctx> From<Error<'ctx>> for error::Error {
	fn from(err: Error) -> Self {
		error::Error::new(err)
	}
}

pub type Result<'ctx, T> = result::Result<T, Error<'ctx>>;

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct TaskRef<'ctx> {
	pub id: &'ctx TaskID,
	pub args: Rc<TaskArgs>,
}

impl<'ctx> Display for TaskRef<'ctx> {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		if !f.alternate() {
			return self.id.fmt(f);
		}

		let pv_arg = match self.args.get("pv") {
			Some(Arg::String(s)) => Some(s),
			_ => None,
		};
		let host_arg = match self.args.get("host") {
			Some(Arg::Platform(platform)) => Some(platform),
			_ => None,
		};
		let target_arg = match self.args.get("target") {
			Some(Arg::Platform(platform)) => Some(platform),
			_ => None,
		};

		write!(f, "{}", self.id.recipe)?;
		if let Some(pv) = pv_arg {
			write!(f, "-{}", pv)?;
		}
		write!(f, ":{}", self.id.task)?;

		if let Some(host) = host_arg {
			write!(f, "@{}", host.short)?;
		}
		if let Some(target) = target_arg {
			write!(f, "/{}", target.short)?;
		}
		Ok(())
	}
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct OutputRef<'ctx> {
	pub task: TaskRef<'ctx>,
	pub output: &'ctx str,
}

fn platform_relation(args: &TaskArgs, from: &str, to: &str) -> Option<PlatformRelation> {
	let plat_from = match args.get(from)? {
		Arg::Platform(plat) => plat,
		_ => return None,
	};
	let plat_to = match args.get(to)? {
		Arg::Platform(plat) => plat,
		_ => return None,
	};

	let plat_rel = if plat_from == plat_to {
		PlatformRelation {
			is_same: true,
			sysroot: "".to_string(),
			cross_compile: "".to_string(),
		}
	} else {
		PlatformRelation {
			is_same: false,
			sysroot: paths::TASK_SYSROOT.to_string(),
			cross_compile: format!("{}/bin/{}-", plat_from.prefix, plat_to.gnu_triplet),
		}
	};
	Some(plat_rel)
}

#[derive(Debug)]
pub struct Context {
	platforms: HashMap<String, Arg>,
	globals: TaskArgs,
	tasks: HashMap<TaskID, TaskDef>,
}

impl Context {
	pub fn new(tasks: HashMap<TaskID, TaskDef>) -> Self {
		let platforms: HashMap<_, _> = [
			arg(
				"build",
				Platform {
					short: "build".to_string(),
					gnu_triplet: "x86_64-linux-gnu".to_string(),
					karch: "x86_64".to_string(),
					prefix: "/opt/toolchain".to_string(),
				},
			),
			arg(
				"aarch64",
				Platform {
					short: "aarch64".to_string(),
					gnu_triplet: "aarch64-linux-gnu".to_string(),
					karch: "arm64".to_string(),
					prefix: "/usr".to_string(),
				},
			),
		]
		.into_iter()
		.collect();

		let globals = TaskArgs::from_iter([
			("build".to_string(), platforms["build"].clone()),
			arg("workdir", paths::TASK_WORKDIR.to_string()),
			arg("dldir", paths::TASK_DLDIR.to_string()),
			arg("destdir", paths::TASK_DESTDIR.to_string()),
			arg("sysroot", paths::TASK_SYSROOT.to_string()),
		]);

		Context {
			platforms,
			globals,
			tasks,
		}
	}

	pub fn get<'ctx>(&'ctx self, id: &'ctx TaskID) -> Result<&TaskDef> {
		self.tasks.get(id).ok_or(Error {
			task: id,
			kind: ErrorKind::TaskNotFound,
		})
	}

	fn task_ref<'ctx>(&'ctx self, id: &'ctx TaskID, args: &TaskArgs) -> Result<TaskRef> {
		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 {
			// TODO: Handle conflicts between explicit args and globals
			arg_def.insert(key, ArgType::from(arg));
		}

		let mut new_args = TaskArgs::default();

		for (key, typ) in arg_def {
			if let Some(arg) = args.get(key) {
				if ArgType::from(arg) == typ {
					new_args.set(key, Some(arg));
					continue;
				}
			}
			return Err(Error {
				task: id,
				kind: ErrorKind::InvalidArgument(key),
			});
		}

		let build_to_host = platform_relation(&new_args, "build", "host");
		let host_to_target = platform_relation(&new_args, "host", "target");
		let build_to_target = platform_relation(&new_args, "build", "target");

		let cross_compile = build_to_host
			.as_ref()
			.map(|build_to_host| build_to_host.cross_compile.clone());

		new_args.set("build_to_host", build_to_host);
		new_args.set("host_to_target", host_to_target);
		new_args.set("build_to_target", build_to_target);

		new_args.set("cross_compile", cross_compile);

		new_args.set("pn", Some(task_def.meta.name.clone()));
		new_args.set("pv", task_def.meta.version.clone());

		Ok(TaskRef {
			id,
			args: Rc::new(new_args),
		})
	}

	pub fn parse<'ctx>(&'ctx self, s: &str) -> error::Result<TaskRef> {
		lazy_static! {
			static ref RE: Regex = Regex::new(
				r"^(?P<recipe>[[:word:]-]+):(?P<task>[[:word:]-]+)(?:@(?P<host>[[:word:]-]+))?(?:/(?P<target>[[:word:]-]+))?$",
			).unwrap();
		}

		let cap = RE.captures(s).context("Invalid task syntax")?;

		let recipe = cap["recipe"].to_string();
		let task = cap["task"].to_string();

		let id = TaskID { recipe, task };
		let (ctx_id, _) = self
			.tasks
			.get_key_value(&id)
			.with_context(|| format!("Task {}:{} not found", id.recipe, id.task))?;

		let mut args = self.globals.clone();

		if let Some(host) = cap.name("host") {
			let plat = self
				.platforms
				.get(host.as_str())
				.with_context(|| format!("Platform '{}' not found", host.as_str()))?;
			args.set("host", Some(plat));
			args.set("target", Some(plat));
		}
		if let Some(target) = cap.name("target") {
			let plat = self
				.platforms
				.get(target.as_str())
				.with_context(|| format!("Platform '{}' not found", target.as_str()))?;
			args.set("target", Some(plat));
		}

		self.task_ref(ctx_id, &args)
			.with_context(|| format!("Failed to instantiate task {}:{}", id.recipe, id.task))
	}

	fn map_args<'ctx, 'args>(
		task: &'ctx TaskID,
		mapping: &'ctx ArgMapping,
		args: &'args TaskArgs,
		build_dep: bool,
	) -> Result<'ctx, Cow<'args, TaskArgs>> {
		if mapping.0.is_empty() && !build_dep {
			return Ok(Cow::Borrowed(args));
		}

		let mut ret = args.clone();

		if build_dep {
			ret.set("host", args.get("build"));
			ret.set("target", args.get("host"));
		}

		for (to, from) in &mapping.0 {
			let value = args.get(from).ok_or(Error {
				task,
				kind: ErrorKind::InvalidArgRef(to),
			})?;
			ret.set(to, Some(value.clone()));
		}

		Ok(Cow::Owned(ret))
	}

	fn inherit_ref<'ctx>(&'ctx self, dep: &'ctx InheritDep, 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())
	}

	pub fn output_ref<'ctx>(
		&'ctx self,
		dep: &'ctx OutputDep,
		args: &TaskArgs,
		build_dep: bool,
	) -> Result<OutputRef<'ctx>> {
		let mapped_args = Context::map_args(&dep.dep.id, &dep.dep.args, args, build_dep)?;
		Ok(OutputRef {
			task: self.task_ref(&dep.dep.id, mapped_args.as_ref())?,
			output: &dep.output,
		})
	}

	pub fn get_inherit_depend<'ctx>(
		&'ctx self,
		task_ref: &TaskRef<'ctx>,
	) -> Result<Option<TaskRef>> {
		let task = self.get(task_ref.id)?;
		let inherit = match &task.inherit {
			Some(inherit) => inherit,
			None => return Ok(None),
		};
		Some(self.inherit_ref(inherit, &task_ref.args)).transpose()
	}

	fn inherit_iter<'ctx>(
		&'ctx self,
		task_ref: &TaskRef<'ctx>,
	) -> impl Iterator<Item = Result<TaskRef>> {
		struct Iter<'ctx>(&'ctx Context, Option<Result<'ctx, TaskRef<'ctx>>>);

		impl<'ctx> Iterator for Iter<'ctx> {
			type Item = Result<'ctx, TaskRef<'ctx>>;

			fn next(&mut self) -> Option<Self::Item> {
				let task_ref = match self.1.take()? {
					Ok(task_ref) => task_ref,
					Err(err) => return Some(Err(err)),
				};
				self.1 = self.0.get_inherit_depend(&task_ref).transpose();
				Some(Ok(task_ref))
			}
		}

		Iter(self, Some(Ok(task_ref.clone())))
	}

	pub fn get_build_depends<'ctx>(
		&'ctx self,
		task_ref: &TaskRef<'ctx>,
	) -> Result<HashSet<OutputRef>> {
		let mut ret = HashSet::new();
		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 entries = task
				.build_depends
				.iter()
				.filter(|dep| allow_noinherit || !dep.noinherit)
				.map(|dep| self.output_ref(dep, &task_ref.args, true))
				.collect::<Result<Vec<_>>>()?;
			ret.extend(entries);

			allow_noinherit = false;
		}

		Ok(ret)
	}

	pub fn get_host_depends<'ctx>(
		&'ctx self,
		task_ref: &TaskRef<'ctx>,
	) -> Result<HashSet<OutputRef>> {
		let mut ret = HashSet::new();
		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 entries = task
				.depends
				.iter()
				.filter(|dep| allow_noinherit || !dep.noinherit)
				.map(|dep| self.output_ref(dep, &task_ref.args, false))
				.collect::<Result<Vec<_>>>()?;
			ret.extend(entries);

			allow_noinherit = false;
		}

		Ok(ret)
	}

	pub fn in_rootfs<'ctx>(&'ctx self, output: &OutputRef<'ctx>) -> bool {
		let build = self.globals.get("build").unwrap();
		if output.task.args.get("host") != Some(build) {
			return false;
		}
		if let Some(target) = output.task.args.get("target") {
			if target != build {
				return false;
			}
		}

		// TODO: Do not hardcode this
		match (
			output.task.id.recipe.as_str(),
			output.task.id.task.as_str(),
			output.output,
		) {
			("gmp", "install", "default") => true,
			("mpfr", "install", "default") => true,
			("mpc", "install", "default") => true,
			("zlib", "install", "default") => true,
			("binutils", "install", "default") => true,
			("gcc", "install", "default") => true,
			("libgcc", "install", "default") => true,
			("gcc-libs", "install", "default") => true,
			("linux-uapi-headers", "install", "default") => true,
			("glibc", "install", "default") => true,
			_ => false,
		}
	}
}

impl Index<&TaskID> for Context {
	type Output = TaskDef;

	fn index(&self, index: &TaskID) -> &TaskDef {
		self.tasks.get(index).expect("Invalid TaskID")
	}
}