orchid/orchid-extension/src/expr.rs

use std::cell::RefCell;
use std::fmt;
use std::hash::Hash;
use std::rc::Rc;

use async_once_cell::OnceCell;
use derive_destructure::destructure;
use hashbrown::HashSet;
use orchid_base::error::OrcErrv;
use orchid_base::format::{FmtCtx, FmtUnit, Format};
use orchid_base::location::Pos;
use orchid_base::reqnot::Requester;

use crate::api;
use crate::atom::ForeignAtom;
use crate::context::{ctx, i};
use crate::gen_expr::{GExpr, GExprKind};

pub struct BorrowedExprStore(RefCell<Option<HashSet<Rc<ExprHandle>>>>);
impl BorrowedExprStore {
	pub(crate) fn new() -> Self { Self(RefCell::new(Some(HashSet::new()))) }
	pub async fn dispose(self) {
		let elements = self.0.borrow_mut().take().unwrap();
		for handle in elements {
			handle.on_borrow_expire().await
		}
	}
}
impl Drop for BorrowedExprStore {
	fn drop(&mut self) {
		if self.0.borrow().is_some() {
			panic!("This should always be explicitly disposed")
		}
	}
}

#[derive(destructure, PartialEq, Eq, Hash)]
pub struct ExprHandle(api::ExprTicket);
impl ExprHandle {
	/// Do not signal to take ownership of the expr. Instead, the
	/// [BorrowedExprStore] signifies the lifetime of the borrow, and when it is
	/// freed, it signals to take ownership of any exprs that ended up outliving
	/// it. It is used to receive exprs sent via [ExprHandle::ticket] as an
	/// optimization over [ExprHandle::from_ticket]
	pub fn borrowed(tk: api::ExprTicket, store: &BorrowedExprStore) -> Rc<Self> {
		let this = Rc::new(Self(tk));
		store.0.borrow_mut().as_mut().unwrap().insert(this.clone());
		this
	}
	/// This function takes over the loose reference pre-created via
	/// [ExprHandle::serialize] in the sender. It must therefore pair up with a
	/// corresponding call to that function.
	pub fn deserialize(tk: api::ExprTicket) -> Rc<Self> { Rc::new(Self(tk)) }
	/// This function takes ownership of a borrowed expr sent via
	/// [ExprHandle::ticket] and signals immediately to record that ownership. It
	/// is used in place of [ExprHandle::borrowed] when it's impractical to
	/// determine how long the borrow will live.
	///
	/// # Safety
	///
	/// You need to ensure that the [api::Acquire] sent by this function arrives
	/// before the borrow expires, so you still need a borrow delimited by some
	/// message you will send in the future.
	pub async fn from_ticket(tk: api::ExprTicket) -> Rc<Self> {
		let store = BorrowedExprStore::new();
		let expr = Self::borrowed(tk, &store);
		store.dispose().await;
		expr
	}
	/// The raw ticket used in messages. If you want to transfer ownership via the
	/// ticket, you should use [ExprHandle::serialize]. Only send this if you want
	/// to lend the expr, and you expect the receiver to use
	/// [ExprHandle::borrowed] or [ExprHandle::from_ticket]
	pub fn ticket(&self) -> api::ExprTicket { self.0 }
	async fn send_acq(&self) { ctx().reqnot().notify(api::Acquire(ctx().sys_id(), self.0)).await }
	/// If this is the last one reference, do nothing, otherwise send an Acquire
	pub async fn on_borrow_expire(self: Rc<Self>) { self.serialize().await; }
	/// Drop the handle and get the ticket without a release notification.
	/// Use this with messages that imply ownership transfer. This function is
	/// safe because abusing it is a memory leak.
	pub async fn serialize(self: Rc<Self>) -> api::ExprTicket {
		match Rc::try_unwrap(self) {
			Err(rc) => {
				rc.send_acq().await;
				rc.0
			},
			Ok(hand) => hand.destructure().0,
		}
	}
}
impl fmt::Debug for ExprHandle {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "ExprHandle({})", self.0.0) }
}
impl Drop for ExprHandle {
	fn drop(&mut self) {
		let notif = api::Release(ctx().sys_id(), self.0);
		ctx().spawn(async move { ctx().reqnot().clone().notify(notif).await })
	}
}

#[derive(Clone, Debug, destructure)]
pub struct Expr {
	handle: Rc<ExprHandle>,
	data: Rc<OnceCell<ExprData>>,
}
impl Expr {
	pub fn from_handle(handle: Rc<ExprHandle>) -> Self { Self { handle, data: Rc::default() } }
	pub fn from_data(handle: Rc<ExprHandle>, d: ExprData) -> Self {
		Self { handle, data: Rc::new(OnceCell::from(d)) }
	}
	/// Creates an instance without incrementing the reference count. This is
	/// only safe to be called on a reference created with an [Expr::serialize]
	/// call which created the loose reference it can take ownership of.
	pub async fn deserialize(tk: api::ExprTicket) -> Self {
		Self::from_handle(ExprHandle::deserialize(tk))
	}
	pub async fn data(&self) -> &ExprData {
		(self.data.get_or_init(async {
			let details = ctx().reqnot().request(api::Inspect { target: self.handle.ticket() }).await;
			let pos = Pos::from_api(&details.location, &i()).await;
			let kind = match details.kind {
				api::InspectedKind::Atom(a) =>
					ExprKind::Atom(ForeignAtom::new(self.handle.clone(), a, pos.clone())),
				api::InspectedKind::Bottom(b) => ExprKind::Bottom(OrcErrv::from_api(&b, &i()).await),
				api::InspectedKind::Opaque => ExprKind::Opaque,
			};
			ExprData { pos, kind }
		}))
		.await
	}
	pub async fn atom(self) -> Result<ForeignAtom, Self> {
		match self.data().await {
			ExprData { kind: ExprKind::Atom(atom), .. } => Ok(atom.clone()),
			_ => Err(self),
		}
	}
	pub fn handle(&self) -> Rc<ExprHandle> { self.handle.clone() }

	pub fn slot(&self) -> GExpr {
		GExpr { pos: Pos::SlotTarget, kind: GExprKind::Slot(self.clone()) }
	}
	/// Increments the refcount to ensure that the ticket remains valid even if
	/// the handle is freed. To avoid a leak, [Expr::deserialize] must eventually
	/// be called.
	pub async fn serialize(self) -> api::ExprTicket { self.handle.serialize().await }
}
impl Format for Expr {
	async fn print<'a>(&'a self, _c: &'a (impl FmtCtx + ?Sized + 'a)) -> FmtUnit {
		match &self.data().await.kind {
			ExprKind::Opaque => "OPAQUE".to_string().into(),
			ExprKind::Bottom(b) => format!("Bottom({b})").into(),
			ExprKind::Atom(a) =>
				FmtUnit::from_api(&ctx().reqnot().request(api::ExtAtomPrint(a.atom.clone())).await),
		}
	}
}

#[derive(Clone, Debug)]
pub struct ExprData {
	pub pos: Pos,
	pub kind: ExprKind,
}

#[derive(Clone, Debug)]
pub enum ExprKind {
	Atom(ForeignAtom),
	Bottom(OrcErrv),
	Opaque,
}