orchid/orchid-extension/src/tree.rs

use std::num::NonZero;

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::interner::{Interner, Tok};
use orchid_base::location::SrcRange;
use orchid_base::name::Sym;
use orchid_base::tree::{TokTree, Token, TokenVariant};
use substack::Substack;
use trait_set::trait_set;

use crate::api;
use crate::context::i;
use crate::conv::ToExpr;
use crate::entrypoint::MemberRecord;
use crate::expr::{BorrowedExprStore, Expr, ExprHandle};
use crate::func_atom::{ExprFunc, Fun};
use crate::gen_expr::{GExpr, sym_ref};

pub type GenTokTree = TokTree<Expr, GExpr>;
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,
		_: &Interner,
	) -> Self {
		panic!("Received new expression from host")
	}
	async fn into_api(self, _: &mut Self::ToApiCtx<'_>) -> api::Expression { self.serialize().await }
}

impl TokenVariant<api::ExprTicket> for Expr {
	type FromApiCtx<'a> = &'a BorrowedExprStore;
	async fn from_api(
		api: &api::ExprTicket,
		exprs: &mut Self::FromApiCtx<'_>,
		_: SrcRange,
		_: &Interner,
	) -> 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() }
}

pub async fn x_tok(x: impl ToExpr) -> GenTok { GenTok::NewExpr(x.to_gen().await) }
pub async fn ref_tok(path: Sym) -> GenTok { GenTok::NewExpr(sym_ref(path)) }

pub fn lazy(
	public: bool,
	name: &str,
	cb: impl AsyncFnOnce(Sym) -> MemKind + Clone + 'static,
) -> Vec<GenMember> {
	vec![GenMember {
		name: name.to_string(),
		kind: MemKind::Lazy(LazyMemberFactory::new(cb)),
		comments: vec![],
		public,
	}]
}
pub fn cnst(public: bool, name: &str, value: impl ToExpr + Clone + 'static) -> Vec<GenMember> {
	lazy(public, name, async |_| MemKind::Const(value.to_gen().await))
}
pub fn module(
	public: bool,
	name: &str,
	mems: impl IntoIterator<Item = Vec<GenMember>>,
) -> Vec<GenMember> {
	let (name, kind) = root_mod(name, mems);
	vec![GenMember { name, kind, public, comments: vec![] }]
}
pub fn root_mod(name: &str, mems: impl IntoIterator<Item = Vec<GenMember>>) -> (String, MemKind) {
	let kind = MemKind::Mod { members: mems.into_iter().flatten().collect() };
	(name.to_string(), kind)
}
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(Fun::new(sym, xf).await.to_gen().await)
	});
	vec![GenMember { name: name.to_string(), kind: MemKind::Lazy(fac), public, comments: vec![] }]
}
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
}

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, (MemKind, Vec<String>)>::new();
	for mem in trees.into_iter().flatten() {
		assert!(mem.public, "Non-trivial merge in {}", mem.name);
		match mem.kind {
			unit @ (MemKind::Const(_) | MemKind::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);
			},
			MemKind::Mod { members } => match all_members.entry(mem.name.clone()) {
				hashbrown::hash_map::Entry::Vacant(slot) => {
					slot.insert((MemKind::Mod { members }, mem.comments.into_iter().collect()));
				},
				hashbrown::hash_map::Entry::Occupied(mut old) => match old.get_mut() {
					(MemKind::Mod { members: 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 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 name: String,
	pub kind: MemKind,
	pub public: bool,
	pub comments: Vec<String>,
}
impl GenMember {
	pub async fn into_api(self, tia_cx: &mut impl TreeIntoApiCtx) -> api::Member {
		let name = i().i::<String>(&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| i().i(cmt).await.to_api())).await;
		api::Member { kind, name: name.to_api(), comments, exported: self.public }
	}
}

pub enum MemKind {
	Const(GExpr),
	Mod { members: Vec<GenMember> },
	Lazy(LazyMemberFactory),
}
impl MemKind {
	pub async fn into_api(self, ctx: &mut impl TreeIntoApiCtx) -> api::MemberKind {
		match self {
			Self::Lazy(lazy) => api::MemberKind::Lazy(ctx.with_lazy(lazy)),
			Self::Const(c) => api::MemberKind::Const(c.serialize().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 trait TreeIntoApiCtx {
	fn with_lazy(&mut self, fac: LazyMemberFactory) -> api::TreeId;
	fn push_path(&mut self, seg: Tok<String>) -> impl TreeIntoApiCtx;
}

pub struct TreeIntoApiCtxImpl<'a, 'b> {
	pub basepath: &'a [Tok<String>],
	pub path: Substack<'a, Tok<String>>,
	pub lazy_members: &'b mut HashMap<api::TreeId, MemberRecord>,
}

impl TreeIntoApiCtx for TreeIntoApiCtxImpl<'_, '_> {
	fn push_path(&mut self, seg: Tok<String>) -> impl TreeIntoApiCtx {
		TreeIntoApiCtxImpl {
			lazy_members: self.lazy_members,
			basepath: self.basepath,
			path: self.path.push(seg),
		}
	}
	fn with_lazy(&mut self, fac: LazyMemberFactory) -> api::TreeId {
		let id = api::TreeId(NonZero::new((self.lazy_members.len() + 2) as u64).unwrap());
		let path = self.basepath.iter().cloned().chain(self.path.unreverse()).collect_vec();
		self.lazy_members.insert(id, MemberRecord::Gen(path, fac));
		id
	}
}