orchid/orchid-extension/src/atom.rs

use std::any::{Any, TypeId, type_name};
use std::collections::HashMap;
use std::fmt::{self, Debug};
use std::future::Future;
use std::io;
use std::marker::PhantomData;
use std::num::NonZeroU32;
use std::ops::Deref;
use std::pin::Pin;
use std::rc::Rc;

use dyn_clone::{DynClone, clone_box};
use futures::future::LocalBoxFuture;
use futures::{AsyncWrite, FutureExt, StreamExt, stream};
use orchid_api_derive::Coding;
use orchid_api_traits::{Coding, Decode, InHierarchy, Request, UnderRoot, enc_vec};
use orchid_base::{
	FmtCtx, FmtUnit, Format, IStr, OrcErrv, Pos, Receipt, ReqHandle, ReqReader, ReqReaderExt, Sym,
	fmt, is, mk_errv, mk_errv_floating, take_first,
};
use trait_set::trait_set;

use crate::gen_expr::GExpr;
use crate::{
	DynSystemCardExt, Expr, ExprData, ExprHandle, ExprKind, OwnedAtom, ToExpr, api, dyn_cted,
	get_obj_store, request, sys_id,
};

/// Every atom managed via this system starts with an ID into the type table
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Coding)]
pub struct AtomTypeId(pub NonZeroU32);

pub trait AtomicVariant {}

/// A value managed by Orchid. The type should also be registered in the
/// [crate::SystemCard] through [AtomicFeatures::ops] which is provided
/// indirectly by either [crate::OwnedAtom] or [crate::ThinAtom]
pub trait Atomic: 'static + Sized {
	/// Either [crate::OwnedVariant] or [crate::ThinVariant] depending on whether
	/// the value implements [crate::OwnedAtom] or [crate::ThinAtom]
	type Variant: AtomicVariant;
	/// Serializable data that gets sent inside the atom to other systems that
	/// depend on this system. Methods on this value are directly accessible
	/// through [TAtom], and this data can also be used for optimized public
	/// functions. The serialized form should have a reasonable length to avoid
	/// overburdening the protocol.
	type Data: Clone + Coding + Sized + 'static;
	/// Register handlers for IPC calls. If this atom implements [Supports], you
	/// should register your implementations here. If this atom doesn't
	/// participate in IPC at all, the default implementation is fine
	fn reg_methods() -> MethodSetBuilder<Self> { MethodSetBuilder::new() }
}

/// Shared interface of all atom types created in this library for use by the
/// library that defines them. This is provided by [Atomic] and either
/// [crate::OwnedAtom] or [crate::ThinAtom]
pub trait AtomicFeatures: Atomic {
	/// Convert a value of this atom inside the defining system into a function
	/// that will perform registrations and serialization
	#[allow(private_interfaces)]
	fn factory(self) -> AtomFactory;
	/// Expose all operations that can be performed on an instance of this type in
	/// an instanceless vtable. This vtable must be registered by the
	/// [crate::System].
	fn ops() -> Box<dyn AtomOps>;
}
pub(crate) trait AtomicFeaturesImpl<Variant: AtomicVariant> {
	fn _factory(self) -> AtomFactory;
	type _Info: AtomOps;
	fn _info() -> Self::_Info;
}
impl<A: Atomic + AtomicFeaturesImpl<A::Variant>> AtomicFeatures for A {
	#[allow(private_interfaces)]
	fn factory(self) -> AtomFactory { self._factory() }
	fn ops() -> Box<dyn AtomOps> { Box::new(Self::_info()) }
}

/// A reference to a value of some [Atomic] type. This owns an [Expr]
#[derive(Clone)]
pub struct ForeignAtom {
	pub(crate) expr: Rc<ExprHandle>,
	pub(crate) atom: api::Atom,
	pub(crate) pos: Pos,
}
impl ForeignAtom {
	/// Obtain the position in code of the expression
	pub fn pos(&self) -> Pos { self.pos.clone() }
	/// Obtain the [Expr]
	pub fn ex(self) -> Expr {
		let (handle, pos) = (self.expr.clone(), self.pos.clone());
		let data = ExprData { pos, kind: ExprKind::Atom(ForeignAtom { ..self }) };
		Expr::from_data(handle, data)
	}
	pub(crate) fn new(handle: Rc<ExprHandle>, atom: api::Atom, pos: Pos) -> Self {
		ForeignAtom { atom, expr: handle, pos }
	}
	/// Call an IPC method. If the type does not support the given method type,
	/// this function returns [None]
	pub async fn call<R: Request + UnderRoot<Root: AtomMethod>>(&self, r: R) -> Option<R::Response> {
		let rep = (request(api::Fwd {
			target: self.atom.clone(),
			method: Sym::parse(<R as UnderRoot>::Root::NAME).await.unwrap().tok().to_api(),
			body: enc_vec(&r.into_root()),
		}))
		.await?;
		Some(R::Response::decode_slice(&mut &rep[..]))
	}
	/// Attempt to downcast this value to a concrete atom type
	pub fn downcast<A: Atomic>(self) -> Result<TAtom<A>, NotTypAtom> {
		let mut data = &self.atom.data.0[..];
		let value = AtomTypeId::decode_slice(&mut data);
		let cted = dyn_cted();
		let own_inst = cted.inst();
		let owner_id = self.atom.owner;
		let typ = type_name::<A>();
		let owner = if sys_id() == owner_id {
			own_inst.card()
		} else {
			(cted.deps().find(|s| s.id() == self.atom.owner))
				.ok_or_else(|| NotTypAtom { expr: self.clone().ex(), pos: self.pos(), typ })?
				.get_card()
		};
		let Some(ops) = owner.ops_by_atid(value) else {
			panic!("{value:?} does not refer to an atom in {owner_id:?} when downcasting {typ}");
		};
		if ops.tid() != TypeId::of::<A>() {
			return Err(NotTypAtom { pos: self.pos.clone(), expr: self.ex(), typ });
		}
		let value = A::Data::decode_slice(&mut data);
		Ok(TAtom { value, untyped: self })
	}
}
impl fmt::Display for ForeignAtom {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "Atom::{:?}", self.atom) }
}
impl fmt::Debug for ForeignAtom {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "ForeignAtom({self})") }
}
impl Format for ForeignAtom {
	async fn print<'a>(&'a self, _c: &'a (impl FmtCtx + ?Sized + 'a)) -> FmtUnit {
		FmtUnit::from_api(&request(api::ExtAtomPrint(self.atom.clone())).await)
	}
}
impl ToExpr for ForeignAtom {
	async fn to_expr(self) -> Expr
	where Self: Sized {
		self.ex()
	}
	async fn to_gen(self) -> GExpr { self.ex().to_gen().await }
}

pub struct NotTypAtom {
	pub pos: Pos,
	pub expr: Expr,
	pub typ: &'static str,
}
impl NotTypAtom {
	/// Convert to a generic Orchid error
	pub async fn mk_err(&self) -> OrcErrv {
		mk_errv(
			is("Not the expected type").await,
			format!("The expression {} is not a {}", fmt(&self.expr).await, self.typ),
			[self.pos.clone()],
		)
	}
}
impl Debug for NotTypAtom {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		f.debug_struct("NotTypAtom")
			.field("pos", &self.pos)
			.field("expr", &self.expr)
			.field("typ", &self.typ)
			.finish_non_exhaustive()
	}
}

/// An IPC request associated with an atom. This type should either implement
/// [Request] or be the root of a [orchid_api_derive::Hierarchy] the leaves of
/// which implement [Request].
pub trait AtomMethod: Coding + InHierarchy {
	const NAME: &str;
}

/// A handler for an [AtomMethod] on an [Atomic]. The [AtomMethod] must also be
/// registered in [Atomic::reg_methods]
pub trait Supports<M: AtomMethod>: Atomic {
	fn handle<'a>(
		&self,
		hand: Box<dyn ReqHandle<'a> + '_>,
		req: M,
	) -> impl Future<Output = io::Result<Receipt<'a>>>;
}

trait HandleAtomMethod<A> {
	fn handle<'a, 'b: 'a>(
		&'a self,
		atom: &'a A,
		reader: Box<dyn ReqReader<'b> + 'a>,
	) -> LocalBoxFuture<'a, ()>;
}
struct AtomMethodHandler<M, A>(PhantomData<M>, PhantomData<A>);
impl<M: AtomMethod, A: Supports<M>> HandleAtomMethod<A> for AtomMethodHandler<M, A> {
	fn handle<'a, 'b: 'a>(
		&'a self,
		atom: &'a A,
		mut reader: Box<dyn ReqReader<'b> + 'a>,
	) -> LocalBoxFuture<'a, ()> {
		Box::pin(async {
			let req = reader.read_req::<M>().await.unwrap();
			let _ = Supports::<M>::handle(atom, reader.finish().await, req).await.unwrap();
		})
	}
}

/// A collection of [Supports] impls for an [Atomic]. If a [Supports]
/// impl is not added to the method set, it will not be recognized. Note that
/// the [Supports] implementors must be registered, which are not necessarily
/// the same as the [Request] implementors
pub struct MethodSetBuilder<A: Atomic> {
	handlers: Vec<(&'static str, Rc<dyn HandleAtomMethod<A>>)>,
}
impl<A: Atomic> MethodSetBuilder<A> {
	pub fn new() -> Self { Self { handlers: vec![] } }

	/// Add an [AtomMethod]
	pub fn handle<M: AtomMethod>(mut self) -> Self
	where A: Supports<M> {
		assert!(!M::NAME.is_empty(), "AtomMethod::NAME cannoot be empty");
		self.handlers.push((M::NAME, Rc::new(AtomMethodHandler::<M, A>(PhantomData, PhantomData))));
		self
	}

	pub(crate) async fn pack(&self) -> MethodSet<A> {
		MethodSet {
			handlers: stream::iter(self.handlers.iter())
				.then(async |(k, v)| (Sym::parse(k).await.unwrap(), v.clone()))
				.collect()
				.await,
		}
	}
}

pub(crate) struct MethodSet<A: Atomic> {
	handlers: HashMap<Sym, Rc<dyn HandleAtomMethod<A>>>,
}
impl<A: Atomic> MethodSet<A> {
	pub(crate) async fn dispatch<'a>(
		&self,
		atom: &'_ A,
		key: Sym,
		req: Box<dyn ReqReader<'a> + 'a>,
	) -> bool {
		match self.handlers.get(&key) {
			None => false,
			Some(handler) => {
				handler.handle(atom, req).await;
				true
			},
		}
	}
}

impl<A: Atomic> Default for MethodSetBuilder<A> {
	fn default() -> Self { Self::new() }
}

/// A handle to a value defined by this or another system. This owns an [Expr]
#[derive(Clone)]
pub struct TAtom<A: Atomic> {
	pub untyped: ForeignAtom,
	pub value: A::Data,
}
impl<A: Atomic> TAtom<A> {
	/// Obtain the underlying [Expr]
	pub fn ex(&self) -> Expr { self.untyped.clone().ex() }
	/// Obtain the position in code associated with the atom
	pub fn pos(&self) -> Pos { self.untyped.pos() }
	/// Produce from an [ExprHandle] directly
	pub async fn downcast(expr: Rc<ExprHandle>) -> Result<Self, NotTypAtom> {
		match Expr::from_handle(expr).atom().await {
			Err(expr) =>
				Err(NotTypAtom { pos: expr.data().await.pos.clone(), expr, typ: type_name::<A>() }),
			Ok(atm) => atm.downcast(),
		}
	}
	/// Find the instance associated with a [TAtom] that we own
	///
	/// # Panics
	///
	/// if we don't actually own this atom
	pub async fn own(&self) -> A
	where A: OwnedAtom {
		let g = get_obj_store().objects.read().await;
		let atom_id = self.untyped.atom.drop.expect("Owned atoms always have a drop ID");
		let dyn_atom =
			g.get(&atom_id).expect("Atom ID invalid; atom type probably not owned by this crate");
		dyn_atom.as_any_ref().downcast_ref().cloned().expect("The ID should imply a type as well")
	}
	/// Call an IPC method on the value. Since we know the type, unlike
	/// [ForeignAtom::call], we can ensure that the callee recognizes this method
	pub async fn call<R: Request + UnderRoot<Root: AtomMethod>>(&self, req: R) -> R::Response
	where A: Supports<<R as UnderRoot>::Root> {
		R::Response::decode_slice(
			&mut &(request(api::Fwd {
				target: self.untyped.atom.clone(),
				method: Sym::parse(<R as UnderRoot>::Root::NAME).await.unwrap().tok().to_api(),
				body: enc_vec(&req.into_root()),
			}))
			.await
			.unwrap()[..],
		)
	}
}
impl<A: AtomicFeatures> Deref for TAtom<A> {
	type Target = A::Data;
	fn deref(&self) -> &Self::Target { &self.value }
}
impl<A: AtomicFeatures> ToExpr for TAtom<A> {
	async fn to_gen(self) -> GExpr { self.untyped.to_gen().await }
}
impl<A: AtomicFeatures> Format for TAtom<A> {
	async fn print<'a>(&'a self, c: &'a (impl FmtCtx + ?Sized + 'a)) -> FmtUnit {
		self.untyped.print(c).await
	}
}

pub(crate) struct AtomCtx<'a>(pub &'a [u8], pub Option<api::AtomId>);

/// A vtable-like type that collects operations defined by an [Atomic] without
/// associating with an instance of that type. This must be registered in
/// [crate::SystemCard]
#[allow(private_interfaces)]
pub trait AtomOps: 'static {
	fn tid(&self) -> TypeId;
	fn name(&self) -> &'static str;
	fn decode<'a>(&'a self, ctx: AtomCtx<'a>) -> LocalBoxFuture<'a, Box<dyn Any>>;
	fn call<'a>(&'a self, ctx: AtomCtx<'a>, arg: Expr) -> LocalBoxFuture<'a, GExpr>;
	fn call_ref<'a>(&'a self, ctx: AtomCtx<'a>, arg: Expr) -> LocalBoxFuture<'a, GExpr>;
	fn print<'a>(&'a self, ctx: AtomCtx<'a>) -> LocalBoxFuture<'a, FmtUnit>;
	fn handle_req_ref<'a>(
		&'a self,
		ctx: AtomCtx<'a>,
		key: Sym,
		req: Box<dyn ReqReader<'a> + 'a>,
	) -> LocalBoxFuture<'a, bool>;
	fn serialize<'a, 'b: 'a>(
		&'a self,
		ctx: AtomCtx<'a>,
		write: Pin<&'b mut dyn AsyncWrite>,
	) -> LocalBoxFuture<'a, Option<Vec<Expr>>>;
	fn deserialize<'a>(
		&'a self,
		data: &'a [u8],
		refs: &'a [Expr],
	) -> LocalBoxFuture<'a, api::LocalAtom>;
	fn drop<'a>(&'a self, ctx: AtomCtx<'a>) -> LocalBoxFuture<'a, ()>;
}

trait_set! {
	pub trait AtomFactoryFn = FnOnce() -> LocalBoxFuture<'static, api::LocalAtom> + DynClone;
}
pub(crate) struct AtomFactory(Box<dyn AtomFactoryFn>, String);
impl AtomFactory {
	pub fn new(name: String, f: impl AsyncFnOnce() -> api::LocalAtom + Clone + 'static) -> Self {
		Self(Box::new(|| f().boxed_local()), name)
	}
	pub async fn build(self) -> api::LocalAtom { (self.0)().await }
}
impl Clone for AtomFactory {
	fn clone(&self) -> Self { AtomFactory(clone_box(&*self.0), self.1.clone()) }
}
impl fmt::Debug for AtomFactory {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "AtomFactory<{}>", self.1) }
}
impl fmt::Display for AtomFactory {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "{self:?}") }
}
impl Format for AtomFactory {
	async fn print<'a>(&'a self, _c: &'a (impl FmtCtx + ?Sized + 'a)) -> FmtUnit {
		self.to_string().into()
	}
}

/// Error produced when an atom can not be applied to a value as a function
pub async fn err_not_callable(unit: &FmtUnit) -> OrcErrv {
	mk_errv_floating(
		is("This atom is not callable").await,
		format!("Attempted to apply {} as function", take_first(unit, false)),
	)
}

/// Error produced when an atom can not be the final value of the program
pub async fn err_not_command(unit: &FmtUnit) -> OrcErrv {
	mk_errv_floating(
		is("This atom is not a command").await,
		format!("Settled on {} which is an inactionable value", take_first(unit, false)),
	)
}

pub(crate) async fn err_exit_success_msg() -> IStr { is("Early successful exit").await }
pub(crate) async fn err_exit_failure_msg() -> IStr { is("Early failure exit").await }

/// Sentinel error returnable from [crate::OwnedAtom::command] or
/// [crate::ThinAtom::command] to indicate that the program should exit with a
/// success
pub async fn err_exit_success() -> OrcErrv {
	mk_errv_floating(
		err_exit_success_msg().await,
		"Sentinel error indicating that the program should exit with a success.",
	)
}

/// Sentinel error returnable from [crate::OwnedAtom::command] or
/// [crate::ThinAtom::command] to indicate that the program should exit with a
/// failure
pub async fn err_exit_failure() -> OrcErrv {
	mk_errv_floating(
		err_exit_failure_msg().await,
		"Sentinel error indicating that the program should exit with a failure \
		but without raising an error.",
	)
}

/// Read the type ID prefix from an atom, return type information and the rest
/// of the data
pub(crate) fn resolve_atom_type(atom: &api::Atom) -> (Box<dyn AtomOps>, AtomTypeId, &[u8]) {
	let mut data = &atom.data.0[..];
	let atid = AtomTypeId::decode_slice(&mut data);
	let atom_record = dyn_cted().inst().card().ops_by_atid(atid).expect("Unrecognized atom type ID");
	(atom_record, atid, data)
}