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 hashbrown::HashSet;
use orchid_api_derive::Coding;
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::{Paren, indent};
use orchid_extension::atom::Atomic;
use orchid_extension::atom_owned::{OwnedAtom, OwnedVariant};
use orchid_extension::expr::Expr;

fn union_rc_sets(seq: impl IntoIterator<Item = Rc<HashSet<Sym>>>) -> Rc<HashSet<Sym>> {
	let mut acc = Rc::<HashSet<Sym>>::default();
	for right in seq {
		if acc.is_empty() {
			acc = right;
			continue;
		}
		if right.is_empty() {
			continue;
		}
		acc = match (Rc::try_unwrap(acc), Rc::try_unwrap(right)) {
			(Ok(mut left), Ok(right)) => {
				left.extend(right);
				Rc::new(left)
			},
			(Ok(mut owned), Err(borrowed)) | (Err(borrowed), Ok(mut owned)) => {
				owned.extend(borrowed.iter().cloned());
				Rc::new(owned)
			},
			(Err(left), Err(right)) => Rc::new(left.union(&right).cloned().collect()),
		}
	}
	acc
}

#[derive(Debug, Clone)]
pub struct MacTreeSeq {
	pub items: Rc<Vec<MacTree>>,
	pub top_glossary: Rc<HashSet<Sym>>,
	pub glossary: Rc<HashSet<Sym>>,
}
impl MacTreeSeq {
	pub fn new(i: impl IntoIterator<Item = MacTree>) -> Self {
		let mut items = Vec::new();
		let mut top_glossary = HashSet::new();
		let mut glossary = HashSet::new();
		for item in i {
			glossary.extend(item.glossary().iter().cloned());
			if let MacTok::Name(n) = item.tok() {
				top_glossary.insert(n.clone());
			}
			items.push(item);
		}
		Self { items: Rc::new(items), top_glossary: Rc::new(top_glossary), glossary: Rc::new(glossary) }
	}
	pub fn map<F: FnMut(MacTree) -> Option<MacTree>>(&self, changed: &mut bool, map: &mut F) -> Self {
		Self::new(self.items.iter().map(|tree| ro(changed, |changed| tree.map(changed, map))))
	}
	pub fn glossary(&self) -> &HashSet<Sym> { &self.glossary }
	pub fn concat(self, other: Self) -> Self {
		if self.items.is_empty() {
			return other;
		} else if other.items.is_empty() {
			return self;
		}
		let items = match (Rc::try_unwrap(self.items), Rc::try_unwrap(other.items)) {
			(Ok(mut left), Ok(mut right)) => {
				left.append(&mut right);
				left
			},
			(Ok(mut left), Err(right)) => {
				left.extend_from_slice(&right[..]);
				left
			},
			(Err(left), Ok(mut right)) => {
				right.splice(0..0, left.iter().cloned());
				right
			},
			(Err(left), Err(right)) => left.iter().chain(&right[..]).cloned().collect(),
		};
		Self {
			items: Rc::new(items),
			top_glossary: union_rc_sets([self.top_glossary, other.top_glossary]),
			glossary: union_rc_sets([self.glossary, other.glossary]),
		}
	}
}
impl Format for MacTreeSeq {
	async fn print<'a>(&'a self, c: &'a (impl FmtCtx + ?Sized + 'a)) -> FmtUnit {
		mtreev_fmt(&self.items[..], c).await
	}
}

#[derive(Debug, Clone)]
pub struct MacTree {
	pub pos: Pos,
	pub tok: Rc<MacTok>,
	pub glossary: Rc<HashSet<Sym>>,
}
impl MacTree {
	pub fn tok(&self) -> &MacTok { &self.tok }
	pub fn pos(&self) -> Pos { self.pos.clone() }
	pub fn glossary(&self) -> &HashSet<Sym> { &self.glossary }
	pub fn map<F: FnMut(Self) -> Option<Self>>(&self, changed: &mut bool, map: &mut F) -> Self {
		let tok = match map(self.clone()) {
			Some(new_tok) => {
				*changed = true;
				return new_tok;
			},
			None => match &*self.tok {
				MacTok::Lambda(arg, body) =>
					MacTok::Lambda(ro(changed, |changed| arg.map(changed, map)), body.map(changed, map)),
				MacTok::Name(_) | MacTok::Value(_) => return self.clone(),
				MacTok::Slot | MacTok::Ph(_) | MacTok::Bottom(_) => return self.clone(),
				MacTok::S(p, body) => MacTok::S(*p, body.map(changed, map)),
			},
		};
		if *changed { tok.at(self.pos()) } else { self.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_atom<'a>(&'a self, c: &'a (impl FmtCtx + ?Sized + 'a)) -> FmtUnit {
		tl_cache!(Rc<Variants>: Rc::new(Variants::default().bounded("'{0}")))
			.units([self.tok.print(c).await])
	}
}
impl Format for MacTree {
	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, MacTreeSeq),
	Name(Sym),
	/// Only permitted in arguments to `instantiate_tpl`
	Slot,
	Value(Expr),
	Lambda(MacTree, MacTreeSeq),
	/// 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 build_glossary(&self) -> Rc<HashSet<Sym>> {
		match self {
			MacTok::Bottom(_) | MacTok::Ph(_) | MacTok::Slot | MacTok::Value(_) => Rc::default(),
			MacTok::Name(sym) => Rc::new(HashSet::from([sym.clone()])),
			MacTok::S(_, body) => union_rc_sets(body.items.iter().map(|mt| mt.glossary.clone())),
			MacTok::Lambda(arg, body) =>
				union_rc_sets(body.items.iter().chain([arg]).map(|mt| mt.glossary.clone())),
		}
	}
	pub fn at(self, pos: impl Into<Pos>) -> MacTree {
		MacTree { pos: pos.into(), glossary: self.build_glossary(), 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) => tl_cache!(Rc<Variants>: Rc::new(Variants::default()
					.unbounded("\\{0} {1l}")
					.bounded("(\\{0} {1b})")))
			.units([arg.print(c).boxed_local().await, b.print(c).await]),
			Self::Name(n) => format!("{n}").into(),
			Self::Ph(ph) => format!("{ph}").into(),
			Self::S(p, body) => match *p {
				Paren::Round => tl_cache!(Rc<Variants>: Rc::new(Variants::default().bounded("({0b})"))),
				Paren::Curly => tl_cache!(Rc<Variants>: Rc::new(Variants::default().bounded("{{0b}}"))),
				Paren::Square => tl_cache!(Rc<Variants>: Rc::new(Variants::default().bounded("[{0b}]"))),
			}
			.units([body.print(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, Coding)]
pub enum PhKind {
	Scalar,
	Vector { at_least_one: bool, priority: u8 },
}

/// 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
}