orchid/orchid-extension/src/atom_owned.rs

use std::any::{Any, TypeId, type_name};
use std::borrow::Cow;
use std::cell::RefCell;
use std::future::Future;
use std::marker::PhantomData;
use std::num::NonZero;
use std::ops::Deref;
use std::pin::Pin;
use std::rc::Rc;

use async_once_cell::OnceCell;
use dyn_clone::{DynClone, clone_box};
use futures::future::{LocalBoxFuture, ready};
use futures::{AsyncRead, AsyncWrite, FutureExt};
use futures_locks::{RwLock, RwLockReadGuard};
use itertools::Itertools;
use memo_map::MemoMap;
use never::Never;
use orchid_api_traits::{Decode, Encode, enc_vec};
use orchid_base::error::OrcRes;
use orchid_base::format::{FmtCtx, FmtCtxImpl, FmtUnit, take_first};
use orchid_base::name::Sym;
use task_local::task_local;

use crate::api;
use crate::atom::{
	AtomCard, AtomCtx, AtomDynfo, AtomFactory, Atomic, AtomicFeaturesImpl, AtomicVariant, MethodSet,
	MethodSetBuilder, TAtom, err_not_callable, err_not_command, get_info,
};
use crate::expr::Expr;
use crate::gen_expr::{GExpr, bot};
use crate::system::{cted, sys_id};

pub struct OwnedVariant;
impl AtomicVariant for OwnedVariant {}
impl<A: OwnedAtom + Atomic<Variant = OwnedVariant>> AtomicFeaturesImpl<OwnedVariant> for A {
	fn _factory(self) -> AtomFactory {
		AtomFactory::new(async move || {
			let obj_store = get_obj_store();
			let atom_id = {
				let mut id = obj_store.next_id.borrow_mut();
				*id += 1;
				api::AtomId(NonZero::new(*id + 1).unwrap())
			};
			let (typ_id, _) = get_info::<A>(cted().inst().card());
			let mut data = enc_vec(&typ_id);
			self.encode(Pin::<&mut Vec<u8>>::new(&mut data)).await;
			obj_store.objects.read().await.insert(atom_id, Box::new(self));
			api::Atom { drop: Some(atom_id), data: api::AtomData(data), owner: sys_id() }
		})
	}
	fn _info() -> Self::_Info { OwnedAtomDynfo { msbuild: A::reg_reqs(), ms: OnceCell::new() } }
	type _Info = OwnedAtomDynfo<A>;
}

/// While an atom read guard is held, no atom can be removed.
pub(crate) struct AtomReadGuard<'a> {
	id: api::AtomId,
	_lock: PhantomData<&'a ()>,
	guard: RwLockReadGuard<MemoMap<api::AtomId, Box<dyn DynOwnedAtom>>>,
}
impl<'a> AtomReadGuard<'a> {
	async fn new(id: api::AtomId) -> Self {
		let guard = get_obj_store().objects.read().await;
		if guard.get(&id).is_none() {
			panic!("Received invalid atom ID: {id:?}");
		}
		Self { id, guard, _lock: PhantomData }
	}
}
impl Deref for AtomReadGuard<'_> {
	type Target = dyn DynOwnedAtom;
	fn deref(&self) -> &Self::Target { &**self.guard.get(&self.id).unwrap() }
}

/// Remove an atom from the store
pub(crate) async fn take_atom(id: api::AtomId) -> Box<dyn DynOwnedAtom> {
	let mut g = get_obj_store().objects.write().await;
	g.remove(&id).unwrap_or_else(|| panic!("Received invalid atom ID: {}", id.0))
}

pub struct OwnedAtomDynfo<T: OwnedAtom> {
	msbuild: MethodSetBuilder<T>,
	ms: OnceCell<MethodSet<T>>,
}
impl<T: OwnedAtom> AtomDynfo for OwnedAtomDynfo<T> {
	fn tid(&self) -> TypeId { TypeId::of::<T>() }
	fn name(&self) -> &'static str { type_name::<T>() }
	fn decode<'a>(&'a self, AtomCtx(data, ..): AtomCtx<'a>) -> LocalBoxFuture<'a, Box<dyn Any>> {
		Box::pin(async {
			Box::new(<T as AtomCard>::Data::decode_slice(&mut &data[..])) as Box<dyn Any>
		})
	}
	fn call(&self, AtomCtx(_, id): AtomCtx, arg: Expr) -> LocalBoxFuture<'_, GExpr> {
		Box::pin(async move { take_atom(id.unwrap()).await.dyn_call(arg).await })
	}
	fn call_ref<'a>(&'a self, AtomCtx(_, id): AtomCtx<'a>, arg: Expr) -> LocalBoxFuture<'a, GExpr> {
		Box::pin(async move { AtomReadGuard::new(id.unwrap()).await.dyn_call_ref(arg).await })
	}
	fn print(&self, AtomCtx(_, id): AtomCtx<'_>) -> LocalBoxFuture<'_, FmtUnit> {
		Box::pin(async move { AtomReadGuard::new(id.unwrap()).await.dyn_print().await })
	}
	fn handle_req<'a, 'b: 'a, 'c: 'a>(
		&'a self,
		AtomCtx(_, id): AtomCtx,
		key: Sym,
		req: Pin<&'b mut dyn AsyncRead>,
		rep: Pin<&'c mut dyn AsyncWrite>,
	) -> LocalBoxFuture<'a, bool> {
		Box::pin(async move {
			let a = AtomReadGuard::new(id.unwrap()).await;
			let ms = self.ms.get_or_init(self.msbuild.pack()).await;
			ms.dispatch(a.as_any_ref().downcast_ref().unwrap(), key, req, rep).await
		})
	}
	fn command<'a>(
		&'a self,
		AtomCtx(_, id): AtomCtx<'a>,
	) -> LocalBoxFuture<'a, OrcRes<Option<GExpr>>> {
		Box::pin(async move { take_atom(id.unwrap()).await.dyn_command().await })
	}
	fn drop(&self, AtomCtx(_, id): AtomCtx) -> LocalBoxFuture<'_, ()> {
		Box::pin(async move { take_atom(id.unwrap()).await.dyn_free().await })
	}
	fn serialize<'a, 'b: 'a>(
		&'a self,
		AtomCtx(_, id): AtomCtx<'a>,
		mut write: Pin<&'b mut dyn AsyncWrite>,
	) -> LocalBoxFuture<'a, Option<Vec<Expr>>> {
		Box::pin(async move {
			let id = id.unwrap();
			id.encode(write.as_mut()).await.unwrap();
			AtomReadGuard::new(id).await.dyn_serialize(write).await
		})
	}
	fn deserialize<'a>(&'a self, data: &'a [u8], refs: &'a [Expr]) -> LocalBoxFuture<'a, api::Atom> {
		Box::pin(async move {
			let refs = T::Refs::from_iter(refs.iter().cloned());
			let obj = T::deserialize(DeserCtxImpl(data), refs).await;
			obj._factory().build().await
		})
	}
}

pub trait DeserializeCtx: Sized {
	fn read<T: Decode>(&mut self) -> impl Future<Output = T>;
	fn is_empty(&self) -> bool;
	fn assert_empty(&self) { assert!(self.is_empty(), "Bytes found after decoding") }
	fn decode<T: Decode>(&mut self) -> impl Future<Output = T> {
		async {
			let t = self.read().await;
			self.assert_empty();
			t
		}
	}
}

struct DeserCtxImpl<'a>(&'a [u8]);
impl DeserializeCtx for DeserCtxImpl<'_> {
	async fn read<T: Decode>(&mut self) -> T { T::decode(Pin::new(&mut self.0)).await.unwrap() }
	fn is_empty(&self) -> bool { self.0.is_empty() }
}

pub trait RefSet {
	fn from_iter<I: Iterator<Item = Expr> + ExactSizeIterator>(refs: I) -> Self;
	fn to_vec(self) -> Vec<Expr>;
}

static E_NON_SER: &str = "Never is a stand-in refset for non-serializable atoms";

impl RefSet for Never {
	fn from_iter<I>(_: I) -> Self { panic!("{E_NON_SER}") }
	fn to_vec(self) -> Vec<Expr> { panic!("{E_NON_SER}") }
}

impl RefSet for () {
	fn to_vec(self) -> Vec<Expr> { Vec::new() }
	fn from_iter<I: Iterator<Item = Expr> + ExactSizeIterator>(refs: I) -> Self {
		assert_eq!(refs.len(), 0, "Expected no refs")
	}
}

impl RefSet for Vec<Expr> {
	fn from_iter<I: Iterator<Item = Expr> + ExactSizeIterator>(refs: I) -> Self { refs.collect_vec() }
	fn to_vec(self) -> Vec<Expr> { self }
}

impl<const N: usize> RefSet for [Expr; N] {
	fn to_vec(self) -> Vec<Expr> { self.into_iter().collect_vec() }
	fn from_iter<I: Iterator<Item = Expr> + ExactSizeIterator>(refs: I) -> Self {
		assert_eq!(refs.len(), N, "Wrong number of refs provided");
		refs.collect_vec().try_into().unwrap_or_else(|_: Vec<_>| unreachable!())
	}
}

/// Atoms that have a [Drop]
pub trait OwnedAtom: Atomic<Variant = OwnedVariant> + Any + Clone + 'static {
	/// If serializable, the collection that best stores subexpression references
	/// for this atom.
	///
	/// - `()` for no subexppressions,
	/// - `[Expr; N]` for a static number of subexpressions
	/// - `Vec<Expr>` for a variable number of subexpressions
	/// - `Never` if not serializable
	///
	/// If this isn't `Never`, you must override the default, panicking
	/// `serialize` and `deserialize` implementation
	type Refs: RefSet;
	fn val(&self) -> impl Future<Output = Cow<'_, Self::Data>>;
	#[allow(unused_variables)]
	fn call_ref(&self, arg: Expr) -> impl Future<Output = GExpr> {
		async move { bot(err_not_callable().await) }
	}
	fn call(self, arg: Expr) -> impl Future<Output = GExpr> {
		async {
			let gcl = self.call_ref(arg).await;
			self.free().await;
			gcl
		}
	}
	#[allow(unused_variables)]
	fn command(self) -> impl Future<Output = OrcRes<Option<GExpr>>> {
		async move { Err(err_not_command().await) }
	}
	#[allow(unused_variables)]
	fn free(self) -> impl Future<Output = ()> { async {} }
	#[allow(unused_variables)]
	fn print_atom<'a>(&'a self, c: &'a (impl FmtCtx + ?Sized + 'a)) -> impl Future<Output = FmtUnit> {
		async { format!("OwnedAtom({})", type_name::<Self>()).into() }
	}
	#[allow(unused_variables)]
	fn serialize(
		&self,
		write: Pin<&mut (impl AsyncWrite + ?Sized)>,
	) -> impl Future<Output = Self::Refs> {
		assert_serializable::<Self>();
		async { panic!("Either implement serialize or set Refs to Never for {}", type_name::<Self>()) }
	}
	#[allow(unused_variables)]
	fn deserialize(dctx: impl DeserializeCtx, refs: Self::Refs) -> impl Future<Output = Self> {
		assert_serializable::<Self>();
		async {
			panic!("Either implement deserialize or set Refs to Never for {}", type_name::<Self>())
		}
	}
}

fn assert_serializable<T: OwnedAtom>() {
	static MSG: &str = "The extension scaffold is broken, Never Refs should prevent serialization";
	assert_ne!(TypeId::of::<T::Refs>(), TypeId::of::<Never>(), "{MSG}");
}

pub trait DynOwnedAtom: DynClone + 'static {
	fn atom_tid(&self) -> TypeId;
	fn as_any_ref(&self) -> &dyn Any;
	fn encode<'a>(&'a self, buffer: Pin<&'a mut dyn AsyncWrite>) -> LocalBoxFuture<'a, ()>;
	fn dyn_call_ref(&self, arg: Expr) -> LocalBoxFuture<'_, GExpr>;
	fn dyn_call(self: Box<Self>, arg: Expr) -> LocalBoxFuture<'static, GExpr>;
	fn dyn_command(self: Box<Self>) -> LocalBoxFuture<'static, OrcRes<Option<GExpr>>>;
	fn dyn_free(self: Box<Self>) -> LocalBoxFuture<'static, ()>;
	fn dyn_print(&self) -> LocalBoxFuture<'_, FmtUnit>;
	fn dyn_serialize<'a>(
		&'a self,
		sink: Pin<&'a mut dyn AsyncWrite>,
	) -> LocalBoxFuture<'a, Option<Vec<Expr>>>;
}
impl<T: OwnedAtom> DynOwnedAtom for T {
	fn atom_tid(&self) -> TypeId { TypeId::of::<T>() }
	fn as_any_ref(&self) -> &dyn Any { self }
	fn encode<'a>(&'a self, buffer: Pin<&'a mut dyn AsyncWrite>) -> LocalBoxFuture<'a, ()> {
		async { self.val().await.as_ref().encode(buffer).await.unwrap() }.boxed_local()
	}
	fn dyn_call_ref(&self, arg: Expr) -> LocalBoxFuture<'_, GExpr> {
		self.call_ref(arg).boxed_local()
	}
	fn dyn_call(self: Box<Self>, arg: Expr) -> LocalBoxFuture<'static, GExpr> {
		self.call(arg).boxed_local()
	}
	fn dyn_command(self: Box<Self>) -> LocalBoxFuture<'static, OrcRes<Option<GExpr>>> {
		self.command().boxed_local()
	}
	fn dyn_free(self: Box<Self>) -> LocalBoxFuture<'static, ()> { self.free().boxed_local() }
	fn dyn_print(&self) -> LocalBoxFuture<'_, FmtUnit> {
		async move { self.print_atom(&FmtCtxImpl::default()).await }.boxed_local()
	}
	fn dyn_serialize<'a>(
		&'a self,
		sink: Pin<&'a mut dyn AsyncWrite>,
	) -> LocalBoxFuture<'a, Option<Vec<Expr>>> {
		match TypeId::of::<Never>() == TypeId::of::<<Self as OwnedAtom>::Refs>() {
			true => ready(None).boxed_local(),
			false => async { Some(self.serialize(sink).await.to_vec()) }.boxed_local(),
		}
	}
}

#[derive(Default)]
pub(crate) struct ObjStore {
	pub(crate) next_id: RefCell<u64>,
	pub(crate) objects: RwLock<MemoMap<api::AtomId, Box<dyn DynOwnedAtom>>>,
}

task_local! {
	static OBJ_STORE: Rc<ObjStore>;
}

pub(crate) fn with_obj_store<'a>(fut: LocalBoxFuture<'a, ()>) -> LocalBoxFuture<'a, ()> {
	Box::pin(OBJ_STORE.scope(Rc::new(ObjStore::default()), fut))
}

pub(crate) fn get_obj_store() -> Rc<ObjStore> {
	OBJ_STORE.try_with(|store| store.clone()).expect("Owned atom store not initialized")
}

pub async fn own<A: OwnedAtom>(typ: &TAtom<A>) -> A {
	let g = get_obj_store().objects.read().await;
	let atom_id = typ.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")
}

pub async fn debug_print_obj_store(show_atoms: bool) {
	let store = get_obj_store();
	let keys = store.objects.read().await.keys().cloned().collect_vec();
	let mut message = "Atoms in store:".to_string();
	if !show_atoms {
		message += &keys.iter().map(|k| format!(" {:?}", k)).join("");
	} else {
		for k in keys {
			let g = store.objects.read().await;
			let Some(atom) = g.get(&k) else {
				message += &format!("\n{k:?} has since been deleted");
				continue;
			};
			let atom = clone_box(&**atom);
			std::mem::drop(g);
			message += &format!("\n{k:?} -> {}", take_first(&atom.dyn_print().await, true));
		}
	}
	eprintln!("{message}")
}