orchid/orchid-extension/src/reflection.rs

use std::cell::OnceCell;
use std::rc::Rc;

use futures::FutureExt;
use futures::lock::Mutex;
use memo_map::MemoMap;
use orchid_base::interner::Tok;
use orchid_base::name::{NameLike, VPath};
use orchid_base::reqnot::Requester;

use crate::api;
use crate::system::{SysCtx, SysCtxEntry, WeakSysCtx};

pub struct ReflMemData {
	// None for inferred steps
	public: OnceCell<bool>,
	kind: ReflMemKind,
}
#[derive(Clone)]
pub struct ReflMem(Rc<ReflMemData>);
impl ReflMem {
	pub fn kind(&self) -> ReflMemKind { self.0.kind.clone() }
}

#[derive(Clone)]
pub enum ReflMemKind {
	Const,
	Mod(ReflMod),
}

pub struct ReflModData {
	inferred: Mutex<bool>,
	path: VPath,
	ctx: WeakSysCtx,
	members: MemoMap<IStr, ReflMem>,
}

#[derive(Clone)]
pub struct ReflMod(Rc<ReflModData>);
impl ReflMod {
	fn ctx(&self) -> SysCtx {
		self.0.ctx.upgrade().expect("ReflectedModule accessed after context drop")
	}
	pub fn path(&self) -> &[IStr] { &self.0.path[..] }
	pub fn is_root(&self) -> bool { self.0.path.is_empty() }
	async fn try_populate(&self) -> Result<(), api::LsModuleError> {
		let ctx = self.ctx();
		let path_tok = ctx.i().i(&self.0.path[..]).await;
		let reply = match ctx.reqnot().request(api::LsModule(ctx.sys_id(), path_tok.to_api())).await {
			Err(api::LsModuleError::TreeUnavailable) =>
				panic!("Reflected tree accessed outside an interpreter call. This extension is faulty."),
			Err(err) => return Err(err),
			Ok(details) => details,
		};
		for (k, v) in reply.members {
			let k = ctx.i().ex(k).await;
			let mem = match self.0.members.get(&k) {
				Some(mem) => mem,
				None => {
					let path = self.0.path.clone().name_with_suffix(k.clone()).to_sym(ctx.i()).await;
					let kind = match v.kind {
						api::MemberInfoKind::Constant => ReflMemKind::Const,
						api::MemberInfoKind::Module =>
							ReflMemKind::Mod(default_module(&ctx, VPath::new(path.segs()))),
					};
					self.0.members.get_or_insert(&k, || default_member(self.is_root(), kind))
				},
			};
			let _ = mem.0.public.set(v.public);
		}
		Ok(())
	}
	pub async fn get_child(&self, key: &IStr) -> Option<ReflMem> {
		let inferred_g = self.0.inferred.lock().await;
		if let Some(mem) = self.0.members.get(key) {
			return Some(mem.clone());
		}
		if !*inferred_g {
			return None;
		}
		match self.try_populate().await {
			Err(api::LsModuleError::InvalidPath) =>
				panic!("Path became invalid since module was created"),
			Err(api::LsModuleError::IsConstant) =>
				panic!("Path previously contained a module but now contains a constant"),
			Err(api::LsModuleError::TreeUnavailable) => unreachable!(),
			Ok(()) => (),
		}
		self.0.members.get(key).cloned()
	}
	pub async fn get_by_path(&self, path: &[IStr]) -> Result<ReflMem, InvalidPathError> {
		let ctx = self.ctx();
		let (next, tail) = path.split_first().expect("Attempted to walk by empty path");
		let inferred_g = self.0.inferred.lock().await;
		if let Some(next) = self.0.members.get(next) {
			return if tail.is_empty() {
				Ok(next.clone())
			} else {
				match next.kind() {
					ReflMemKind::Const => Err(InvalidPathError { keep_ancestry: true }),
					ReflMemKind::Mod(m) => m.get_by_path(tail).boxed_local().await,
				}
			};
		}
		if !*inferred_g {
			return Err(InvalidPathError { keep_ancestry: true });
		}
		let candidate = default_module(&ctx, self.0.path.clone().suffix([next.clone()]));
		if tail.is_empty() {
			return match candidate.try_populate().await {
				Ok(()) => {
					let tgt_mem = default_member(self.is_root(), ReflMemKind::Mod(candidate));
					self.0.members.insert(next.clone(), tgt_mem.clone());
					Ok(tgt_mem)
				},
				Err(api::LsModuleError::InvalidPath) => Err(InvalidPathError { keep_ancestry: false }),
				Err(api::LsModuleError::IsConstant) => {
					let const_mem = default_member(self.is_root(), ReflMemKind::Const);
					self.0.members.insert(next.clone(), const_mem);
					Err(InvalidPathError { keep_ancestry: true })
				},
				Err(api::LsModuleError::TreeUnavailable) => unreachable!(),
			};
		}
		match candidate.get_by_path(tail).boxed_local().await {
			e @ Err(InvalidPathError { keep_ancestry: false }) => e,
			res @ Err(InvalidPathError { keep_ancestry: true }) | res @ Ok(_) => {
				let tgt_mem = default_member(self.is_root(), ReflMemKind::Mod(candidate));
				self.0.members.insert(next.clone(), tgt_mem);
				res
			},
		}
	}
}

struct ReflRoot(ReflMod);
impl SysCtxEntry for ReflRoot {}

pub struct InvalidPathError {
	keep_ancestry: bool,
}

fn default_module(ctx: &SysCtx, path: VPath) -> ReflMod {
	ReflMod(Rc::new(ReflModData {
		ctx: ctx.downgrade(),
		inferred: Mutex::new(true),
		path,
		members: MemoMap::new(),
	}))
}

fn default_member(is_root: bool, kind: ReflMemKind) -> ReflMem {
	ReflMem(Rc::new(ReflMemData {
		public: if is_root { true.into() } else { OnceCell::new() },
		kind,
	}))
}

fn get_root(ctx: &SysCtx) -> &ReflRoot {
	ctx.get_or_insert(|| ReflRoot(default_module(ctx, VPath::new([]))))
}

pub fn refl(ctx: &SysCtx) -> ReflMod { get_root(ctx).0.clone() }