orchid/orchid-host/src/tree.rs

//! This tree isn't Clone because lazy subtrees are guaranteed to only be loaded
//! once
use std::cell::RefCell;
use std::rc::{Rc, Weak};
use std::slice;

use async_once_cell::OnceCell;
use async_std::sync::RwLock;
use futures::{FutureExt, StreamExt, stream};
use hashbrown::HashMap;
use hashbrown::hash_map::Entry;
use itertools::Itertools;
use orchid_base::clone;
use orchid_base::error::{OrcRes, Reporter, mk_errv};
use orchid_base::interner::Tok;
use orchid_base::location::{CodeGenInfo, Pos};
use orchid_base::name::{NameLike, Sym, VPath};
use orchid_base::reqnot::Requester;

use crate::api;
use crate::ctx::Ctx;
use crate::dealias::{ChildErrorKind, Tree, absolute_path, resolv_glob, walk};
use crate::expr::{Expr, ExprParseCtx, PathSetBuilder};
use crate::parsed::{ItemKind, ParsedMemberKind, ParsedModule};
use crate::system::System;

pub struct RootData {
	pub root: Module,
	pub consts: HashMap<Sym, Expr>,
	pub ctx: Ctx,
}
#[derive(Clone)]
pub struct Root(pub Rc<RwLock<RootData>>);
impl Root {
	#[must_use]
	pub fn new(ctx: Ctx) -> Self {
		Root(Rc::new(RwLock::new(RootData {
			root: Module::default(),
			consts: HashMap::default(),
			ctx,
		})))
	}
	#[must_use]
	pub async fn from_api(api: api::Module, sys: &System) -> Self {
		let mut consts = HashMap::new();
		let mut tfac = TreeFromApiCtx { consts: &mut consts, path: sys.i().i(&[][..]).await, sys };
		let root = Module::from_api(api, &mut tfac).await;
		Root(Rc::new(RwLock::new(RootData { root, consts, ctx: sys.ctx().clone() })))
	}
	pub async fn merge(&self, new: &Root) -> Result<Self, MergeErr> {
		let this = self.0.read().await;
		let that = new.0.read().await;
		let mut consts =
			this.consts.iter().chain(&that.consts).map(|(k, v)| (k.clone(), v.clone())).collect();
		let root = this.root.merge(&that.root, this.ctx.clone(), &mut consts).await?;
		Ok(Self(Rc::new(RwLock::new(RootData { root, consts, ctx: this.ctx.clone() }))))
	}
	#[must_use]
	pub async fn add_parsed(&self, parsed: &ParsedModule, pars_prefix: Sym, rep: &Reporter) -> Self {
		let mut ref_this = self.0.write().await;
		let this = &mut *ref_this;
		let mut deferred_consts = HashMap::new();
		let mut consts = this.consts.clone();
		let mut tfpctx = FromParsedCtx {
			pars_root: parsed,
			deferred_consts: &mut deferred_consts,
			consts: &mut consts,
			pars_prefix: pars_prefix.clone(),
			root: &this.root,
			ctx: &this.ctx,
			rep,
		};
		let mut module = Module::from_parsed(parsed, pars_prefix.clone(), &mut tfpctx).await;
		for step in pars_prefix.iter().rev() {
			let kind = OnceCell::from(MemberKind::Module(module));
			let members = HashMap::from([(
				step.clone(),
				Rc::new(Member { public: true, lazy: RefCell::new(None), kind }),
			)]);
			module = Module { imports: HashMap::new(), members }
		}
		let root = (this.root.merge(&module, this.ctx.clone(), &mut consts).await)
			.expect("Merge conflict between parsed and existing module");
		let new = Root(Rc::new(RwLock::new(RootData { root, consts, ctx: this.ctx.clone() })));
		*this.ctx.root.write().await = new.downgrade();
		for (path, (sys_id, pc_id)) in deferred_consts {
			let sys = this.ctx.system_inst(sys_id).await.expect("System dropped since parsing");
			let api_expr = sys.reqnot().request(api::FetchParsedConst { id: pc_id, sys: sys.id() }).await;
			let mut xp_ctx = ExprParseCtx { ctx: &this.ctx, exprs: sys.ext().exprs() };
			let expr = Expr::from_api(&api_expr, PathSetBuilder::new(), &mut xp_ctx).await;
			new.0.write().await.consts.insert(path, expr);
		}
		new
	}
	pub async fn get_const_value(&self, name: Sym, pos: Pos) -> OrcRes<Expr> {
		let this = &mut *self.0.write().await;
		// shortcut for previously visited
		if let Some(val) = this.consts.get(&name) {
			return Ok(val.clone());
		}
		// load the node, then check if this "walk" call added it to the map
		let ctx = this.ctx.clone();
		let module =
			walk(&this.root, false, name.iter().cloned(), &mut (ctx.clone(), &mut this.consts)).await;
		if let Some(val) = this.consts.get(&name) {
			return Ok(val.clone());
		}
		match module {
			Ok(_) => Err(mk_errv(
				ctx.i.i("module used as constant").await,
				format!("{name} is a module, not a constant"),
				[pos],
			)),
			Err(e) => match e.kind {
				ChildErrorKind::Private => panic!("public_only is false"),
				ChildErrorKind::Constant => panic!("Tree refers to constant not in table"),
				ChildErrorKind::Missing => Err(mk_errv(
					ctx.i.i("Constant does not exist").await,
					format!("{name} does not refer to a constant"),
					[pos],
				)),
			},
		}
	}
	#[must_use]
	pub fn downgrade(&self) -> WeakRoot { WeakRoot(Rc::downgrade(&self.0)) }
}

#[derive(Clone)]
pub struct WeakRoot(Weak<RwLock<RootData>>);
impl WeakRoot {
	#[must_use]
	pub fn new() -> Self { Self(Weak::new()) }
	#[must_use]
	pub fn upgrade(&self) -> Option<Root> { Some(Root(self.0.upgrade()?)) }
}
impl Default for WeakRoot {
	fn default() -> Self { Self::new() }
}

pub struct TreeFromApiCtx<'a> {
	pub sys: &'a System,
	pub consts: &'a mut HashMap<Sym, Expr>,
	pub path: Tok<Vec<Tok<String>>>,
}
impl<'a> TreeFromApiCtx<'a> {
	#[must_use]
	pub async fn push<'c>(&'c mut self, name: Tok<String>) -> TreeFromApiCtx<'c> {
		let path = self.sys.ctx().i.i(&self.path.iter().cloned().chain([name]).collect_vec()).await;
		TreeFromApiCtx { path, consts: &mut *self.consts, sys: self.sys }
	}
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ResolvedImport {
	pub target: Sym,
	pub pos: Pos,
}

#[derive(Clone, Default)]
pub struct Module {
	pub imports: HashMap<Tok<String>, Result<ResolvedImport, Vec<ResolvedImport>>>,
	pub members: HashMap<Tok<String>, Rc<Member>>,
}
impl Module {
	#[must_use]
	pub async fn from_api(api: api::Module, ctx: &mut TreeFromApiCtx<'_>) -> Self {
		let mut members = HashMap::new();
		for mem in api.members {
			let mem_name = ctx.sys.i().ex(mem.name).await;
			let vname = VPath::new(ctx.path.iter().cloned()).name_with_suffix(mem_name.clone());
			let name = vname.to_sym(ctx.sys.i()).await;
			let (lazy, kind) = match mem.kind {
				api::MemberKind::Lazy(id) =>
					(Some(LazyMemberHandle { id, sys: ctx.sys.id(), path: name.clone() }), None),
				api::MemberKind::Const(val) => {
					let mut expr_ctx = ExprParseCtx { ctx: ctx.sys.ctx(), exprs: ctx.sys.ext().exprs() };
					let expr = Expr::from_api(&val, PathSetBuilder::new(), &mut expr_ctx).await;
					ctx.consts.insert(name.clone(), expr);
					(None, Some(MemberKind::Const))
				},
				api::MemberKind::Module(m) => {
					let m = Self::from_api(m, &mut ctx.push(mem_name.clone()).await).boxed_local().await;
					(None, Some(MemberKind::Module(m)))
				},
			};
			members.insert(
				mem_name.clone(),
				Rc::new(Member {
					public: mem.exported,
					lazy: RefCell::new(lazy),
					kind: kind.map_or_else(OnceCell::new, OnceCell::from),
				}),
			);
		}
		Self { members, imports: HashMap::new() }
	}
	#[must_use]
	async fn from_parsed(parsed: &ParsedModule, path: Sym, ctx: &mut FromParsedCtx<'_>) -> Self {
		let imports_by_name = (parsed.get_imports().into_iter())
			.filter_map(|i| Some((i.name.clone()?, i)))
			.into_group_map();
		let mut glob_imports_by_name = HashMap::<_, Vec<_>>::new();
		for import in parsed.get_imports().into_iter().filter(|i| i.name.is_none()) {
			let pos = import.sr.pos();
			match absolute_path(&path, &import.path, &ctx.ctx.i).await {
				Err(e) => ctx.rep.report(e.err_obj(&ctx.ctx.i, pos, &import.path.to_string()).await),
				Ok(abs_path) => {
					let names_res = match abs_path.strip_prefix(&ctx.pars_prefix[..]) {
						None => {
							let mut tree_ctx = (ctx.ctx.clone(), &mut *ctx.consts);
							resolv_glob(&path, ctx.root, &abs_path, pos, &ctx.ctx.i, &mut tree_ctx).await
						},
						Some(sub_tgt) => {
							let sub_path = (path.strip_prefix(&ctx.pars_prefix[..]))
								.expect("from_parsed called with path outside pars_prefix");
							resolv_glob(sub_path, ctx.pars_root, sub_tgt, pos, &ctx.ctx.i, &mut ()).await
						},
					};
					let abs_path = abs_path.to_sym(&ctx.ctx.i).await;
					match names_res {
						Err(e) => ctx.rep.report(e),
						Ok(names) =>
							for name in names {
								match glob_imports_by_name.entry(name) {
									Entry::Occupied(mut o) => o.get_mut().push((abs_path.clone(), import.sr.clone())),
									Entry::Vacant(v) => {
										v.insert_entry(vec![(abs_path.clone(), import.sr.clone())]);
									},
								}
							},
					}
				},
			}
		}
		let mut imports = HashMap::new();
		if parsed.use_prelude {
			let systems = ctx.ctx.systems.read().await;
			for sys in systems.values().flat_map(|weak| weak.upgrade()) {
				for prelude_item in sys.prelude() {
					imports.insert(
						prelude_item.last_seg(),
						Ok(ResolvedImport {
							target: prelude_item,
							pos: CodeGenInfo::new_details(sys.ctor().name_tok().await, "In prelude", &ctx.ctx.i)
								.await
								.pos(),
						}),
					);
				}
			}
		}
		let conflicting_imports_msg = ctx.ctx.i.i("Conflicting imports").await;
		for (key, values) in imports_by_name {
			if values.len() == 1 {
				let import = values.into_iter().next().unwrap();
				let sr = import.sr.clone();
				let abs_path_res = absolute_path(&path, &import.clone().mspath(), &ctx.ctx.i).await;
				match abs_path_res {
					Err(e) => ctx.rep.report(e.err_obj(&ctx.ctx.i, sr.pos(), &import.to_string()).await),
					Ok(abs_path) => {
						let target = abs_path.to_sym(&ctx.ctx.i).await;
						imports.insert(key, Ok(ResolvedImport { target, pos: sr.pos() }));
					},
				}
			} else {
				for item in values {
					ctx.rep.report(mk_errv(
						conflicting_imports_msg.clone(),
						format!("{key} is imported multiple times from different modules"),
						[item.sr.pos()],
					));
				}
			}
		}
		for (key, values) in glob_imports_by_name {
			if !imports.contains_key(&key) {
				let i = &ctx.ctx.i;
				let values = stream::iter(values)
					.then(|(n, sr)| {
						clone!(key; async move {
							ResolvedImport {
								target: n.to_vname().suffix([key.clone()]).to_sym(i).await,
								pos: sr.pos(),
							}
						})
					})
					.collect::<Vec<_>>()
					.await;
				imports.insert(key, if values.len() == 1 { Ok(values[0].clone()) } else { Err(values) });
			}
		}
		let self_referential_msg = ctx.ctx.i.i("Self-referential import").await;
		for (key, value) in imports.iter() {
			let Ok(import) = value else { continue };
			if import.target.strip_prefix(&path[..]).is_some_and(|t| t.starts_with(slice::from_ref(key)))
			{
				ctx.rep.report(mk_errv(
					self_referential_msg.clone(),
					format!("import {} points to itself or a path within itself", &import.target),
					[import.pos.clone()],
				));
			}
		}
		let mut members = HashMap::new();
		for item in &parsed.items {
			match &item.kind {
				ItemKind::Member(mem) => {
					let path = path.to_vname().suffix([mem.name.clone()]).to_sym(&ctx.ctx.i).await;
					let kind = OnceCell::from(MemberKind::from_parsed(&mem.kind, path.clone(), ctx).await);
					members.insert(
						mem.name.clone(),
						Rc::new(Member { kind, lazy: RefCell::default(), public: mem.exported }),
					);
				},
				ItemKind::Import(_) => (),
			}
		}
		Module { imports, members }
	}
	pub async fn merge(
		&self,
		other: &Module,
		ctx: Ctx,
		consts: &mut HashMap<Sym, Expr>,
	) -> Result<Module, MergeErr> {
		if !self.imports.is_empty() || !other.imports.is_empty() {
			return Err(MergeErr { path: VPath::new([]), kind: MergeErrKind::Imports });
		}
		let mut members = HashMap::new();
		for (key, mem) in &other.members {
			let Some(own) = self.members.get(key) else {
				members.insert(key.clone(), mem.clone());
				continue;
			};
			if own.public != mem.public {
				return Err(MergeErr { path: VPath::new([]), kind: MergeErrKind::Visibility });
			}
			match (own.kind(ctx.clone(), consts).await, mem.kind(ctx.clone(), consts).await) {
				(MemberKind::Module(own_sub), MemberKind::Module(sub)) => {
					match own_sub.merge(sub, ctx.clone(), consts).boxed_local().await {
						Ok(module) => {
							members.insert(
								key.clone(),
								Rc::new(Member {
									lazy: RefCell::new(None),
									public: own.public,
									kind: OnceCell::from(MemberKind::Module(module)),
								}),
							);
						},
						Err(mut e) => {
							e.path = e.path.prefix([key.clone()]);
							return Err(e);
						},
					}
				},
				_ => return Err(MergeErr { path: VPath::new([key.clone()]), kind: MergeErrKind::Const }),
			}
		}
		for (key, mem) in &self.members {
			if let Entry::Vacant(slot) = members.entry(key.clone()) {
				slot.insert(mem.clone());
			}
		}
		Ok(Module { imports: HashMap::new(), members })
	}
}

#[derive(Debug)]
pub struct MergeErr {
	pub path: VPath,
	pub kind: MergeErrKind,
}
#[derive(Debug)]
pub enum MergeErrKind {
	Imports,
	Visibility,
	Const,
}

pub struct FromParsedCtx<'a> {
	pars_prefix: Sym,
	pars_root: &'a ParsedModule,
	root: &'a Module,
	rep: &'a Reporter,
	ctx: &'a Ctx,
	consts: &'a mut HashMap<Sym, Expr>,
	deferred_consts: &'a mut HashMap<Sym, (api::SysId, api::ParsedConstId)>,
}

impl Tree for Module {
	type Ctx<'a> = (Ctx, &'a mut HashMap<Sym, Expr>);
	async fn child(
		&self,
		key: Tok<String>,
		public_only: bool,
		(ctx, consts): &mut Self::Ctx<'_>,
	) -> crate::dealias::ChildResult<'_, Self> {
		let Some(member) = self.members.get(&key) else {
			return Err(ChildErrorKind::Missing);
		};
		if public_only && !member.public {
			return Err(ChildErrorKind::Private);
		}
		match &member.kind(ctx.clone(), consts).await {
			MemberKind::Module(m) => Ok(m),
			MemberKind::Const => Err(ChildErrorKind::Constant),
		}
	}
	fn children(&self, public_only: bool) -> hashbrown::HashSet<Tok<String>> {
		self.members.iter().filter(|(_, v)| !public_only || v.public).map(|(k, _)| k.clone()).collect()
	}
}

pub struct Member {
	pub public: bool,
	pub lazy: RefCell<Option<LazyMemberHandle>>,
	pub kind: OnceCell<MemberKind>,
}
impl Member {
	#[must_use]
	pub async fn kind<'a>(&'a self, ctx: Ctx, consts: &mut HashMap<Sym, Expr>) -> &'a MemberKind {
		(self.kind.get_or_init(async {
			let handle = self.lazy.borrow_mut().take().expect("If kind is uninit, lazy must be Some");
			handle.run(ctx, consts).await
		}))
		.await
	}
}

pub enum MemberKind {
	Const,
	Module(Module),
}
impl MemberKind {
	#[must_use]
	async fn from_parsed(parsed: &ParsedMemberKind, path: Sym, ctx: &mut FromParsedCtx<'_>) -> Self {
		match parsed {
			ParsedMemberKind::DeferredConst(id, sys) => {
				ctx.deferred_consts.insert(path, (sys.id(), *id));
				MemberKind::Const
			},
			ParsedMemberKind::Mod(m) =>
				MemberKind::Module(Module::from_parsed(m, path, ctx).boxed_local().await),
		}
	}
}

pub struct LazyMemberHandle {
	id: api::TreeId,
	sys: api::SysId,
	path: Sym,
}
impl LazyMemberHandle {
	#[must_use]
	pub async fn run(self, ctx: Ctx, consts: &mut HashMap<Sym, Expr>) -> MemberKind {
		let sys = ctx.system_inst(self.sys).await.expect("Missing system for lazy member");
		match sys.get_tree(self.id).await {
			api::MemberKind::Const(c) => {
				let mut pctx = ExprParseCtx { ctx: &ctx, exprs: sys.ext().exprs() };
				consts.insert(self.path, Expr::from_api(&c, PathSetBuilder::new(), &mut pctx).await);
				MemberKind::Const
			},
			api::MemberKind::Module(m) => MemberKind::Module(
				Module::from_api(m, &mut TreeFromApiCtx { sys: &sys, consts, path: self.path.tok() }).await,
			),
			api::MemberKind::Lazy(id) => Self { id, ..self }.run(ctx, consts).boxed_local().await,
		}
	}
	#[must_use]
	pub async fn into_member(self, public: bool) -> Member {
		Member { public, kind: OnceCell::new(), lazy: RefCell::new(Some(self)) }
	}
}