orchid/orchid-std/src/macros/mactree.rs

use std::borrow::Cow;
use std::fmt::Display;
use std::rc::Rc;

use futures::FutureExt;
use futures::future::join_all;
use itertools::Itertools;
use orchid_api::Paren;
use orchid_base::error::OrcErrv;
use orchid_base::format::{FmtCtx, FmtUnit, Format, Variants};
use orchid_base::interner::Tok;
use orchid_base::location::Pos;
use orchid_base::name::Sym;
use orchid_base::tl_cache;
use orchid_base::tree::indent;
use orchid_extension::atom::Atomic;
use orchid_extension::atom_owned::{OwnedAtom, OwnedVariant};
use orchid_extension::expr::Expr;

#[derive(Debug, Clone)]
pub struct MacTree {
	pub pos: Pos,
	pub tok: Rc<MacTok>,
}
impl MacTree {
	pub fn tok(&self) -> &MacTok { &*self.tok }
	pub fn pos(&self) -> Pos { self.pos.clone() }
}
impl Atomic for MacTree {
	type Data = ();
	type Variant = OwnedVariant;
}
impl OwnedAtom for MacTree {
	type Refs = ();

	async fn val(&self) -> Cow<'_, Self::Data> { Cow::Owned(()) }
	async fn print<'a>(&'a self, c: &'a (impl FmtCtx + ?Sized + 'a)) -> FmtUnit {
		self.tok.print(c).await
	}
}

#[derive(Debug, Clone)]
pub enum MacTok {
	S(Paren, Vec<MacTree>),
	Name(Sym),
	/// Only permitted in arguments to `instantiate_tpl`
	Slot,
	Value(Expr),
	Lambda(MacTree, Vec<MacTree>),
	/// Only permitted in "pattern" values produced by macro blocks, which are
	/// never accessed as variables by usercode
	Ph(Ph),
	Bottom(OrcErrv),
}
impl MacTok {
	pub fn at(self, pos: impl Into<Pos>) -> MacTree {
		MacTree { pos: pos.into(), tok: Rc::new(self) }
	}
}
impl Format for MacTok {
	async fn print<'a>(&'a self, c: &'a (impl FmtCtx + ?Sized + 'a)) -> FmtUnit {
		match self {
			Self::Value(v) => v.print(c).await,
			Self::Lambda(arg, b) => FmtUnit::new(
				tl_cache!(Rc<Variants>: Rc::new(Variants::default()
					.unbounded("\\{0b}.{1l}")
					.bounded("(\\{0b}.{1b})"))),
				[arg.print(c).boxed_local().await, mtreev_fmt(b, c).await],
			),
			Self::Name(n) => format!("{n}").into(),
			Self::Ph(ph) => format!("{ph}").into(),
			Self::S(p, body) => FmtUnit::new(
				match *p {
					Paren::Round => Rc::new(Variants::default().bounded("({0b})")),
					Paren::Curly => Rc::new(Variants::default().bounded("{{0b}}")),
					Paren::Square => Rc::new(Variants::default().bounded("[{0b}]")),
				},
				[mtreev_fmt(body, c).await],
			),
			Self::Slot => "SLOT".into(),
			Self::Bottom(err) if err.len() == 1 => format!("Bottom({}) ", err.one().unwrap()).into(),
			Self::Bottom(err) => format!("Botttom(\n{}) ", indent(&err.to_string())).into(),
		}
	}
}

pub async fn mtreev_fmt<'b>(
	v: impl IntoIterator<Item = &'b MacTree>,
	c: &(impl FmtCtx + ?Sized),
) -> FmtUnit {
	FmtUnit::sequence(" ", None, join_all(v.into_iter().map(|t| t.print(c))).await)
}

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub struct Ph {
	pub name: Tok<String>,
	pub kind: PhKind,
}
impl Display for Ph {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		match self.kind {
			PhKind::Scalar => write!(f, "${}", self.name),
			PhKind::Vector { at_least_one: false, priority: 0 } => write!(f, "..${}", self.name),
			PhKind::Vector { at_least_one: true, priority: 0 } => write!(f, "...${}", self.name),
			PhKind::Vector { at_least_one: false, priority } => write!(f, "..${}:{priority}", self.name),
			PhKind::Vector { at_least_one: true, priority } => write!(f, "...${}:{priority}", self.name),
		}
	}
}

#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub enum PhKind {
	Scalar,
	Vector { at_least_one: bool, priority: u8 },
}

pub fn map_mactree<F: FnMut(MacTree) -> Option<MacTree>>(
	src: &MacTree,
	changed: &mut bool,
	map: &mut F,
) -> MacTree {
	let tok = match map(src.clone()) {
		Some(new_tok) => {
			*changed = true;
			return new_tok;
		},
		None => match &*src.tok {
			MacTok::Lambda(arg, body) => MacTok::Lambda(
				ro(changed, |changed| map_mactree(arg, changed, map)),
				map_mactree_v(body, changed, map),
			),
			MacTok::Name(_) | MacTok::Value(_) | MacTok::Slot | MacTok::Ph(_) | MacTok::Bottom(_) =>
				return src.clone(),
			MacTok::S(p, body) => MacTok::S(*p, map_mactree_v(body, changed, map)),
		},
	};
	if *changed { MacTree { pos: src.pos.clone(), tok: Rc::new(tok) } } else { src.clone() }
}
pub fn map_mactree_v<F: FnMut(MacTree) -> Option<MacTree>>(
	src: &[MacTree],
	changed: &mut bool,
	map: &mut F,
) -> Vec<MacTree> {
	src.iter().map(|tree| ro(changed, |changed| map_mactree(tree, changed, map))).collect_vec()
}

/// reverse "or". Inside, the flag is always false, but raising it will raise
/// the outside flag too.
fn ro<T>(flag: &mut bool, cb: impl FnOnce(&mut bool) -> T) -> T {
	let mut new_flag = false;
	let val = cb(&mut new_flag);
	*flag |= new_flag;
	val
}