orchid/orchid-extension/src/tree.rs

use std::cell::RefCell;
use std::num::NonZero;
use std::rc::Rc;

use async_fn_stream::stream;
use dyn_clone::{DynClone, clone_box};
use futures::future::{LocalBoxFuture, join_all};
use futures::{FutureExt, StreamExt};
use hashbrown::HashMap;
use itertools::Itertools;
use orchid_base::{IStr, SrcRange, Sym, TokTree, Token, TokenVariant, is};
use substack::Substack;
use task_local::task_local;
use trait_set::trait_set;

use crate::gen_expr::{GExpr, new_atom, serialize};
use crate::{BorrowedExprStore, Expr, ExprFunc, ExprHandle, Fun, ToExpr, api};

/// Tokens generated by lexers and parsers
///
/// See: [GenTok], [Token], [crate::Lexer], [crate::Parser]
pub type GenTokTree = TokTree<Expr, GExpr>;
/// Tokens generated by lexers and parsers - without location data
///
/// See: [GenTokTree], [Token], [crate::Lexer], [crate::Parser]
pub type GenTok = Token<Expr, GExpr>;

impl TokenVariant<api::Expression> for GExpr {
	type FromApiCtx<'a> = ();
	type ToApiCtx<'a> = ();
	async fn from_api(_: api::Expression, _: &mut Self::FromApiCtx<'_>, _: SrcRange) -> Self {
		panic!("Received new expression from host")
	}
	async fn into_api(self, _: &mut Self::ToApiCtx<'_>) -> api::Expression { serialize(self).await }
}

impl TokenVariant<api::ExprTicket> for Expr {
	type FromApiCtx<'a> = &'a BorrowedExprStore;
	async fn from_api(api: api::ExprTicket, exprs: &mut Self::FromApiCtx<'_>, _: SrcRange) -> Self {
		// SAFETY: receiving trees from sublexers implies borrowing
		Expr::from_handle(ExprHandle::borrowed(api, exprs))
	}
	type ToApiCtx<'a> = ();
	async fn into_api(self, (): &mut Self::ToApiCtx<'_>) -> api::ExprTicket { self.handle().ticket() }
}

/// Embed a literal value in generated syntax.
pub async fn x_tok(x: impl ToExpr) -> GenTok { GenTok::NewExpr(x.to_gen().await) }
/// Embed a reference to a constant in generated syntax. The constant doesn't
/// need to be visible per export rules, and if it doesn't exist, the expression
/// will raise a runtime error
pub async fn ref_tok(path: Sym) -> GenTok { GenTok::NewExpr(path.to_gen().await) }

/// Create a new subtree that is evaluated as-needed, asynchronously, and can
/// use its own path to determine its value
pub fn lazy(
	public: bool,
	name: &str,
	cb: impl AsyncFnOnce(Sym) -> MemKind + Clone + 'static,
) -> Vec<GenMember> {
	vec![GenMember {
		name: name.to_string(),
		kind: LazyMemKind::Lazy(LazyMemberFactory::new(cb)),
		comments: vec![],
		public,
	}]
}
/// A constant node in the module tree
pub fn cnst(public: bool, name: &str, value: impl ToExpr + Clone + 'static) -> Vec<GenMember> {
	lazy(public, name, async |_| MemKind::Const(value.to_gen().await))
}
/// A module in the tree. These can be merged by [merge_trivial]
pub fn module(
	public: bool,
	name: &str,
	mems: impl IntoIterator<Item = Vec<GenMember>>,
) -> Vec<GenMember> {
	let (name, kind) = (name.to_string(), LazyMemKind::Mod(mems.into_iter().flatten().collect()));
	vec![GenMember { name, kind, public, comments: vec![] }]
}
/// A Rust function that is passed to Orchid via [Fun]
pub fn fun<I, O>(public: bool, name: &str, xf: impl ExprFunc<I, O>) -> Vec<GenMember> {
	let fac =
		LazyMemberFactory::new(async move |sym| MemKind::Const(new_atom(Fun::new(sym, xf).await)));
	vec![GenMember { name: name.to_string(), kind: LazyMemKind::Lazy(fac), public, comments: vec![] }]
}
/// A chain of nested modules with names taken from the duble::colon::delimited
/// path ultimately containing the elements
pub fn prefix(path: &str, items: impl IntoIterator<Item = Vec<GenMember>>) -> Vec<GenMember> {
	let mut items = items.into_iter().flatten().collect_vec();
	for step in path.split("::").collect_vec().into_iter().rev() {
		items = module(true, step, [items]);
	}
	items
}
/// Add comments to a set of members
pub fn comments<'a>(
	cmts: impl IntoIterator<Item = &'a str>,
	mut val: Vec<GenMember>,
) -> Vec<GenMember> {
	let cmts = cmts.into_iter().map(|c| c.to_string()).collect_vec();
	for v in val.iter_mut() {
		v.comments.extend(cmts.iter().cloned());
	}
	val
}

/// Trivially merge a gen tree. Behaviours were chosen to make this simple.
///
/// - Comments on imports are discarded
/// - Comments on exports and submodules are combined
/// - Duplicate constants result in an error
/// - A combination of lazy and anything results in an error
pub fn merge_trivial(trees: impl IntoIterator<Item = Vec<GenMember>>) -> Vec<GenMember> {
	let mut all_members = HashMap::<String, (LazyMemKind, Vec<String>)>::new();
	for mem in trees.into_iter().flatten() {
		assert!(mem.public, "Non-trivial merge in {}", mem.name);
		match mem.kind {
			unit @ (LazyMemKind::Const(_) | LazyMemKind::Lazy(_)) => {
				let prev = all_members.insert(mem.name.clone(), (unit, mem.comments.into_iter().collect()));
				assert!(prev.is_none(), "Conflict in trivial tree merge on {}", mem.name);
			},
			LazyMemKind::Mod(members) => match all_members.entry(mem.name.clone()) {
				hashbrown::hash_map::Entry::Vacant(slot) => {
					slot.insert((LazyMemKind::Mod(members), mem.comments.into_iter().collect()));
				},
				hashbrown::hash_map::Entry::Occupied(mut old) => match old.get_mut() {
					(LazyMemKind::Mod(old_items), old_cmts) => {
						let mut swap = vec![];
						std::mem::swap(&mut swap, old_items);
						*old_items = merge_trivial([swap, members]);
						old_cmts.extend(mem.comments);
					},
					_ => panic!("non-trivial merge on {}", mem.name),
				},
			},
		}
	}

	(all_members.into_iter())
		.map(|(name, (kind, comments))| GenMember { comments, kind, name, public: true })
		.collect_vec()
}

trait_set! {
	trait LazyMemberCallback =
		FnOnce(Sym) -> LocalBoxFuture<'static, MemKind> + DynClone
}
pub(crate) struct LazyMemberFactory(Box<dyn LazyMemberCallback>);
impl LazyMemberFactory {
	pub fn new(cb: impl AsyncFnOnce(Sym) -> MemKind + Clone + 'static) -> Self {
		Self(Box::new(|s| cb(s).boxed_local()))
	}
	pub async fn build(self, path: Sym) -> MemKind { (self.0)(path).await }
}
impl Clone for LazyMemberFactory {
	fn clone(&self) -> Self { Self(clone_box(&*self.0)) }
}

pub struct GenMember {
	pub(crate) name: String,
	pub(crate) kind: LazyMemKind,
	pub(crate) public: bool,
	pub(crate) comments: Vec<String>,
}
impl GenMember {
	pub(crate) async fn into_api(self, tia_cx: &mut impl TreeIntoApiCtx) -> api::Member {
		let name = is(&self.name).await;
		let kind = self.kind.into_api(&mut tia_cx.push_path(name.clone())).await;
		let comments = join_all(self.comments.iter().map(async |cmt| is(cmt).await.to_api())).await;
		api::Member { kind, name: name.to_api(), comments, exported: self.public }
	}
}

/// Items in the tree after deferrals have been resolved
pub enum MemKind {
	Const(GExpr),
	Mod(Vec<GenMember>),
}
impl MemKind {
	pub async fn cnst(val: impl ToExpr) -> Self { Self::Const(val.to_gen().await) }
	pub fn module(mems: impl IntoIterator<Item = Vec<GenMember>>) -> Self {
		Self::Mod(mems.into_iter().flatten().collect())
	}
}

pub(crate) enum LazyMemKind {
	Const(GExpr),
	Mod(Vec<GenMember>),
	Lazy(LazyMemberFactory),
}
impl LazyMemKind {
	pub(crate) async fn into_api(self, ctx: &mut impl TreeIntoApiCtx) -> api::MemberKind {
		match self {
			Self::Lazy(lazy) => api::MemberKind::Lazy(add_lazy(ctx, lazy)),
			Self::Const(c) => api::MemberKind::Const(serialize(c).await),
			Self::Mod(members) => api::MemberKind::Module(api::Module {
				members: stream(async |mut cx| {
					for m in members {
						cx.emit(m.into_api(ctx).await).await
					}
				})
				.collect()
				.boxed_local()
				.await,
			}),
		}
	}
}

pub(crate) enum MemberRecord {
	Gen(Vec<IStr>, LazyMemberFactory),
	Res,
}

#[derive(Clone, Default)]
pub(crate) struct LazyMemberStore(Rc<RefCell<HashMap<api::TreeId, MemberRecord>>>);

task_local! {
	static LAZY_MEMBERS: LazyMemberStore;
}

pub(crate) fn with_lazy_member_store<'a>(fut: LocalBoxFuture<'a, ()>) -> LocalBoxFuture<'a, ()> {
	Box::pin(LAZY_MEMBERS.scope(LazyMemberStore::default(), fut))
}

fn add_lazy(cx: &impl TreeIntoApiCtx, fac: LazyMemberFactory) -> api::TreeId {
	LAZY_MEMBERS.with(|lazy_members| {
		let mut g = lazy_members.0.borrow_mut();
		let id = api::TreeId(NonZero::new((g.len() + 2) as u64).unwrap());
		let path = cx.path().collect_vec();
		g.insert(id, MemberRecord::Gen(path, fac));
		id
	})
}

pub(crate) async fn get_lazy(id: api::TreeId) -> (Sym, LazyMemKind) {
	let (path, cb) =
		LAZY_MEMBERS.with(|tbl| match tbl.0.borrow_mut().insert(id, MemberRecord::Res) {
			None => panic!("Tree for ID not found"),
			Some(MemberRecord::Res) => panic!("This tree has already been transmitted"),
			Some(MemberRecord::Gen(path, cb)) => (path, cb),
		});
	let path = Sym::new(path).await.unwrap();
	(path.clone(), match cb.build(path).await {
		MemKind::Const(c) => LazyMemKind::Const(c),
		MemKind::Mod(m) => LazyMemKind::Mod(m),
	})
}

pub(crate) trait TreeIntoApiCtx {
	fn push_path(&mut self, seg: IStr) -> impl TreeIntoApiCtx;
	fn path(&self) -> impl Iterator<Item = IStr>;
}

pub(crate) struct TreeIntoApiCtxImpl<'a> {
	pub basepath: &'a [IStr],
	pub path: Substack<'a, IStr>,
}

impl TreeIntoApiCtx for TreeIntoApiCtxImpl<'_> {
	fn push_path(&mut self, seg: IStr) -> impl TreeIntoApiCtx {
		TreeIntoApiCtxImpl { basepath: self.basepath, path: self.path.push(seg) }
	}
	fn path(&self) -> impl Iterator<Item = IStr> {
		self.basepath.iter().cloned().chain(self.path.unreverse())
	}
}