orchid/orchid-extension/src/gen_expr.rs

use std::cell::RefCell;
use std::marker::PhantomData;
use std::mem;
use std::pin::{Pin, pin};
use std::rc::Rc;

use futures::{FutureExt, Stream, StreamExt, stream};
use orchid_base::{
	FmtCtx, FmtUnit, Format, OrcErr, OrcErrv, Pos, Sym, Variants, match_mapping, tl_cache,
};
use substack::Substack;
use task_local::task_local;

use crate::{AtomFactory, AtomicFeatures, Expr, ToExpr, ToExprFuture, api, request, sys_id};

#[derive(Clone, Copy, Debug)]
struct ExprSerializeCx<'a> {
	closures: Substack<'a, u64>,
	lambda_counter: &'a RefCell<u64>,
}

/// Release notifications will not be sent for the slots. Use this with
/// messages that imply ownership transfer
pub async fn serialize(expr: GExpr) -> api::Expression {
	let cx = ExprSerializeCx { closures: Substack::Bottom, lambda_counter: &RefCell::new(0) };
	expr.serialize(cx).await
}

/// Smart object representing AST not-yet-sent to the interpreter. This type can
/// be cloned and persisted, and it must not have unbound arguments. The helper
/// functions in this module let you build trees of [ToExpr] implementors which
/// represent lambdas and their arguments separately, and then convert them into
/// [GExpr] in one pass.
#[derive(Clone, Debug)]
pub struct GExpr {
	/// AST node type
	kind: GExprKind,
	/// Code location associated with the expression for debugging purposes
	pos: Pos,
}
impl GExpr {
	async fn serialize(self, cx: ExprSerializeCx<'_>) -> api::Expression {
		if let GExprKind::Slot(ex) = self.kind {
			let hand = ex.handle();
			mem::drop(ex);
			api::Expression {
				location: api::Location::SlotTarget,
				// an instance is leaked here, we must take ownership of it when we receive this
				kind: api::ExpressionKind::Slot(hand.serialize().await),
			}
		} else {
			api::Expression {
				location: self.pos.to_api(),
				kind: self.kind.serialize(cx).boxed_local().await,
			}
		}
	}
	/// Reassign location information. The typical default is [Pos::Inherit]
	pub fn at(self, pos: Pos) -> Self { GExpr { pos, kind: self.kind } }
	/// Send the expression to the interpreter to be compiled and to become
	/// shareable across extensions
	pub async fn create(self) -> Expr {
		Expr::deserialize(request(api::Create(sys_id(), serialize(self).await)).await).await
	}
}
impl Format for GExpr {
	async fn print<'a>(&'a self, c: &'a (impl FmtCtx + ?Sized + 'a)) -> FmtUnit {
		self.kind.print(c).boxed_local().await
	}
}

/// AST nodes recognized by the interpreter
#[derive(Clone, Debug)]
pub enum GExprKind {
	/// Function call
	Call(Box<GExpr>, Box<GExpr>),
	/// Lambda expression. Argument must be the same for slot
	Lambda(Box<GExpr>),
	/// Slot for a lambda argument
	Arg(u64),
	/// The second expression is only valid after the first one had already been
	/// fully normalized. The main use case is the pattern `Lambda(0, Seq(0,
	/// Call(foo, 0)))` where foo is an atom that attempts to downcast its
	/// argument.
	Seq(Box<GExpr>, Box<GExpr>),
	/// A reference to a constant from the shared constant tree. It is best to
	/// mark the system that provides named constants as a dependency, but this is
	/// not required
	Const(Sym),
	/// A newly created atom. Since at this point the atom needs to be registered
	/// inside the extension but doesn't yet have an [api::ExprTicket], atoms need
	/// their own [api::Atom::drop] if they have an identity
	#[allow(private_interfaces)]
	NewAtom(AtomFactory),
	/// An expression previously registered or coming from outside the extension
	Slot(Expr),
	/// A runtime error
	Bottom(OrcErrv),
}
impl GExprKind {
	pub fn at(self, pos: Pos) -> GExpr { GExpr { kind: self, pos } }
	async fn serialize(self, cx: ExprSerializeCx<'_>) -> api::ExpressionKind {
		match_mapping!(self, Self => api::ExpressionKind {
			Call(
				f => Box::new(f.serialize(cx).await),
				x => Box::new(x.serialize(cx).await)
			),
			Seq(
				a => Box::new(a.serialize(cx).await),
				b => Box::new(b.serialize(cx).await)
			),
			Const(name.to_api()),
			Bottom(err.to_api()),
			NewAtom(fac.clone().build().await),
		} {
			Self::Slot(_) => panic!("processed elsewhere"),
			Self::Lambda(body) => {
				let id: u64;
				{
					let mut g = cx.lambda_counter.borrow_mut();
					id = *g;
					*g += 1;
				};
				let cx = ExprSerializeCx {
					lambda_counter: cx.lambda_counter,
					closures: cx.closures.push(id)
				};
				api::ExpressionKind::Lambda(id,
					Box::new(body.serialize(cx).await)
				)
			},
			Self::Arg(arg) => {
				api::ExpressionKind::Arg(*cx.closures.iter().nth(arg as usize).expect("Unbound arg"))
			},
		})
	}
}
impl Format for GExprKind {
	async fn print<'a>(&'a self, c: &'a (impl FmtCtx + ?Sized + 'a)) -> FmtUnit {
		match self {
			GExprKind::Call(f, x) =>
				tl_cache!(Rc<Variants>: Rc::new(Variants::default().bounded("{0} ({1})")))
					.units([f.print(c).await, x.print(c).await]),
			GExprKind::Lambda(body) =>
				tl_cache!(Rc<Variants>: Rc::new(Variants::default().bounded("\\{1}")))
					.units([body.print(c).await]),
			GExprKind::Arg(arg) => arg.to_string().into(),
			GExprKind::Seq(a, b) =>
				tl_cache!(Rc<Variants>: Rc::new(Variants::default().bounded("[{0}] {1}")))
					.units([a.print(c).await, b.print(c).await]),
			GExprKind::Const(sym) => sym.to_string().into(),
			GExprKind::NewAtom(atom_factory) => atom_factory.to_string().into(),
			GExprKind::Slot(expr) =>
				tl_cache!(Rc<Variants>: Rc::new(Variants::default().bounded("{{{0}}}")))
					.units([expr.print(c).await]),
			GExprKind::Bottom(orc_errv) => orc_errv.to_string().into(),
		}
	}
}

pub fn inherit(kind: GExprKind) -> GExpr { GExpr { pos: Pos::Inherit, kind } }

task_local! {
	pub static CLOSURE_DEPTH: u64;
}

impl ToExpr for Sym {
	async fn to_expr(self) -> Expr
	where Self: Sized {
		self.to_gen().await.create().await
	}
	async fn to_gen(self) -> GExpr { inherit(GExprKind::Const(self)) }
}
/// Creates an expression from a new atom that we own.
pub fn new_atom<A: AtomicFeatures>(atom: A) -> GExpr { inherit(GExprKind::NewAtom(atom.factory())) }

pub fn slot(expr: Expr) -> GExpr { GExpr { pos: Pos::SlotTarget, kind: GExprKind::Slot(expr) } }

/// An expression which is only valid if a number of dependencies had already
/// been normalized
pub fn seq(
	deps: impl IntoGExprStream,
	val: impl ToExpr,
) -> ToExprFuture<impl Future<Output = GExpr>> {
	ToExprFuture(async {
		async fn recur(mut ops: Pin<&mut impl Stream<Item = GExpr>>) -> Option<GExpr> {
			let op = ops.next().await?.to_gen().await;
			Some(match recur(ops).boxed_local().await {
				None => op,
				Some(rec) => inherit(GExprKind::Seq(Box::new(op), Box::new(rec))),
			})
		}
		recur(pin!(deps.into_gexpr_stream().chain(stream::iter([val.to_gen().await]))))
			.await
			.expect("Empty list provided to seq!")
	})
}

#[derive(Debug, Clone, Copy)]
pub enum ArgState {
	Building,
	Serializing { depth: u64 },
	Ready,
}

/// Argument bound by an enclosing [lam] or [dyn_lambda]
#[derive(Debug, Clone, Copy)]
pub struct GenArg<'a>(*const RefCell<ArgState>, PhantomData<&'a ()>);
impl ToExpr for GenArg<'_> {
	async fn to_gen(self) -> GExpr {
		// SAFETY: Created from a Rc that lives as long as the lifetime arg, see [lam]
		let state = unsafe { self.0.as_ref().unwrap() };
		match (*state.borrow(), CLOSURE_DEPTH.try_with(|r| *r)) {
			(ArgState::Serializing { .. }, Err(_)) =>
				panic!("Lambda should have cleared up argstate alongside CLOSURE_DEPTH"),
			(ArgState::Serializing { depth }, Ok(total)) => inherit(GExprKind::Arg(total - depth)),
			(ArgState::Building, _) =>
				panic!("Argument serialized before lambda. Likely an over-eager ToExpr impl"),
			(ArgState::Ready, _) =>
				unreachable!("The arg should never be available this long, the GenArg is a convenience"),
		}
	}
}

/// A lambda expression.
pub fn lam<'a>(
	cb: impl for<'b> AsyncFnOnce(GenArg<'b>) -> GExpr + 'a,
) -> ToExprFuture<impl Future<Output = GExpr> + 'a> {
	let state = Rc::new(RefCell::new(ArgState::Building));
	ToExprFuture(async move {
		let rank = CLOSURE_DEPTH.try_with(|r| *r + 1).unwrap_or(0);
		match *state.borrow_mut() {
			ref mut state @ ArgState::Building => *state = ArgState::Serializing { depth: rank },
			ArgState::Serializing { .. } => panic!("Lambda serialized twice, found interrupted"),
			ArgState::Ready => panic!("Lambda serialized twice"),
		}
		let gen_arg = GenArg(Rc::as_ptr(&state), PhantomData);
		let ret = CLOSURE_DEPTH.scope(rank, async { cb(gen_arg).await.to_gen().await }).await;
		mem::drop(state);
		inherit(GExprKind::Lambda(Box::new(ret)))
	})
}

/// one or more items that are convertible to expressions. In practice, a
/// [ToExpr], [Vec<GExpr>], or a tuple of types that all implement [ToExpr]. For
/// compilation performance, the tuple's arity may not be more than 6
pub trait IntoGExprStream {
	/// Convert each item to an expression and return them
	fn into_gexpr_stream(self) -> impl Stream<Item = GExpr>;
}
impl<T: ToExpr> IntoGExprStream for T {
	fn into_gexpr_stream(self) -> impl Stream<Item = GExpr> { (self,).into_gexpr_stream() }
}
impl IntoGExprStream for Vec<GExpr> {
	fn into_gexpr_stream(self) -> impl Stream<Item = GExpr> { stream::iter(self) }
}

mod tuple_impls {
	use futures::{Stream, StreamExt, stream};

	use super::IntoGExprStream;
	use crate::conv::ToExpr;
	use crate::gen_expr::GExpr;

	macro_rules! tuple_impl {
		($($T:ident)*) => {
			pastey::paste!{
				impl<$($T: ToExpr),*> IntoGExprStream for ($($T,)*) {
					fn into_gexpr_stream(self) -> impl Stream<Item = GExpr> {
						let ($([< $T:snake >],)*) = self;
						stream::once(async { stream::iter([$([< $T:snake >].to_gen().await,)*]) }).flatten()
					}
				}
			}
		};
	}

	tuple_impl!();
	tuple_impl!(A);
	tuple_impl!(A B);
	tuple_impl!(A B C);
	tuple_impl!(A B C D);
	tuple_impl!(A B C D E);
	tuple_impl!(A B C D E F);
}

/// Call a (curried) function
pub fn call(
	f: impl ToExpr,
	argv: impl IntoGExprStream,
) -> ToExprFuture<impl Future<Output = GExpr>> {
	ToExprFuture(async {
		(argv.into_gexpr_stream())
			.fold(f.to_gen().await, async |f, x| inherit(GExprKind::Call(Box::new(f), Box::new(x))))
			.await
	})
}

/// Call a function on a dynamic number of arguments
pub fn call_v(
	f: impl ToExpr,
	argv: impl IntoIterator<Item: ToExpr>,
) -> ToExprFuture<impl Future<Output = GExpr>> {
	ToExprFuture(async {
		stream::iter(argv)
			.fold(f.to_gen().await, async |f, x| {
				inherit(GExprKind::Call(Box::new(f), Box::new(x.to_gen().await)))
			})
			.await
	})
}

/// A runtime error
pub fn bot(ev: impl IntoIterator<Item = OrcErr>) -> GExpr {
	inherit(GExprKind::Bottom(OrcErrv::new(ev).unwrap()))
}