orchid/orchid-extension/src/reflection.rs

use std::cell::{OnceCell, RefCell};
use std::rc::Rc;

use futures::FutureExt;
use futures::future::LocalBoxFuture;
use futures::lock::Mutex;
use hashbrown::HashMap;
use memo_map::MemoMap;
use orchid_base::interner::{IStr, es, iv};
use orchid_base::name::{NameLike, VPath};
use task_local::task_local;

use crate::api;
use crate::entrypoint::request;
use crate::system::sys_id;

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

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

#[derive(Debug)]
pub struct ReflModData {
	inferred: Mutex<bool>,
	path: VPath,
	members: MemoMap<IStr, ReflMem>,
}

#[derive(Clone, Debug)]
pub struct ReflMod(Rc<ReflModData>);
impl ReflMod {
	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 path_tok = iv(&self.0.path[..]).await;
		let reply = match request(api::LsModule(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 = es(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().await;
					let kind = match v.kind {
						api::MemberInfoKind::Constant => ReflMemKind::Const,
						api::MemberInfoKind::Module =>
							ReflMemKind::Mod(default_module(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 (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(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.clone());
					Ok(const_mem)
				},
				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
			},
		}
	}
}

task_local! {
	static REFL_ROOTS: RefCell<HashMap<api::SysId, ReflMod>>
}

#[derive(Clone, Debug)]
pub struct InvalidPathError {
	keep_ancestry: bool,
}

fn default_module(path: VPath) -> ReflMod {
	ReflMod(Rc::new(ReflModData { 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,
	}))
}

pub fn refl() -> ReflMod {
	REFL_ROOTS.with(|tbl| {
		tbl.borrow_mut().entry(sys_id()).or_insert_with(|| default_module(VPath::new([]))).clone()
	})
}

pub fn with_refl_roots<'a>(fut: LocalBoxFuture<'a, ()>) -> LocalBoxFuture<'a, ()> {
	Box::pin(REFL_ROOTS.scope(RefCell::default(), fut))
}