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 itertools::Itertools;
use orchid_base::error::{OrcErrv, Reporter, mk_err};
use orchid_base::format::{FmtCtx, FmtUnit, Format, Variants, fmt};
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::conv::ToExpr;
use orchid_extension::expr::Expr;
use orchid_extension::gen_expr::{GExpr, arg, bot, call, lambda, sym_ref};
use orchid_extension::system::SysCtx;
use substack::Substack;

#[derive(Clone)]
pub struct LowerCtx<'a> {
	pub sys: SysCtx,
	pub rep: &'a Reporter,
}

#[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 async fn lower(&self, ctx: LowerCtx<'_>, args: Substack<'_, Sym>) -> GExpr {
		let expr = match self.tok() {
			MacTok::Bottom(e) => bot(e.clone()),
			MacTok::Lambda(arg, body) => {
				let MacTok::Name(name) = &*arg.tok else {
					let err = mk_err(
						ctx.sys.i().i("Syntax error after macros").await,
						"This token ends up as a binding, consider replacing it with a name",
						[arg.pos()],
					);
					ctx.rep.report(err.clone());
					return bot(err);
				};
				lambda(args.len() as u64, lower_v(body, ctx, args.push(name.clone())).await)
			},
			MacTok::Name(name) => match args.iter().enumerate().find(|(_, n)| *n == name) {
				None => sym_ref(name.clone()),
				Some((i, _)) => arg((args.len() - i) as u64),
			},
			MacTok::Ph(ph) => {
				let err = mk_err(
					ctx.sys.i().i("Placeholder in value").await,
					format!("Placeholder {ph} is only supported in macro patterns"),
					[self.pos()],
				);
				ctx.rep.report(err.clone());
				return bot(err);
			},
			MacTok::S(Paren::Round, body) => call(lower_v(body, ctx, args).await),
			MacTok::S(..) => {
				let err = mk_err(
					ctx.sys.i().i("[] or {} after macros").await,
					format!("{} didn't match any macro", fmt(self, ctx.sys.i()).await),
					[self.pos()],
				);
				ctx.rep.report(err.clone());
				return bot(err);
			},
			MacTok::Slot => panic!("Uninstantiated template should never be exposed"),
			MacTok::Value(v) => v.clone().to_expr().await,
		};
		expr.at(self.pos())
	}
}
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 {
		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
	}
}

pub async fn lower_v(v: &[MacTree], ctx: LowerCtx<'_>, args: Substack<'_, Sym>) -> Vec<GExpr> {
	join_all(v.iter().map(|t| t.lower(ctx.clone(), args.clone())).collect::<Vec<_>>()).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 build_glossary(&self) -> HashSet<Sym> {
		match self {
			MacTok::Bottom(_) | MacTok::Ph(_) | MacTok::Slot | MacTok::Value(_) => HashSet::new(),
			MacTok::Name(sym) => HashSet::from([sym.clone()]),
			MacTok::S(_, body) => body.iter().flat_map(|mt| &*mt.glossary).cloned().collect(),
			MacTok::Lambda(arg, body) =>
				body.iter().chain([arg]).flat_map(|mt| &*mt.glossary).cloned().collect(),
		}
	}
	pub fn at(self, pos: impl Into<Pos>) -> MacTree {
		MacTree { pos: pos.into(), glossary: Rc::new(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) => 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: IStr,
	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(_) => return src.clone(),
			MacTok::Slot | MacTok::Ph(_) | MacTok::Bottom(_) => return src.clone(),
			MacTok::S(p, body) => MacTok::S(*p, map_mactree_v(body, changed, map)),
		},
	};
	if *changed { tok.at(src.pos()) } 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
}

pub fn glossary_v(src: &[MacTree]) -> impl Iterator<Item = Sym> {
	src.iter().flat_map(|mt| mt.glossary()).cloned()
}