Correctly halts

orchid-host/src/expr_store.rs

@@ -2,6 +2,7 @@ use std::cell::RefCell;
 use std::fmt;
 use std::rc::Rc;
 
+use bound::Bound;
 use hashbrown::HashMap;
 use hashbrown::hash_map::Entry;
 
@@ -12,15 +13,29 @@ use crate::expr::Expr;
 pub struct ExprStoreData {
 	exprs: RefCell<HashMap<api::ExprTicket, (u32, Expr)>>,
 	parent: Option<ExprStore>,
+	tracking_parent: bool,
 }
 #[derive(Clone, Default)]
 pub struct ExprStore(Rc<ExprStoreData>);
 impl ExprStore {
+	/// If tracking_parent is false, get_expr can fall back to the parent if none
+	/// is found here.
+	///
+	/// If tracking_parent is true, get_expr can still fall back to the parent,
+	/// but operations on the parent can access the child exprs too until this
+	/// store is dropped.
 	#[must_use]
-	pub fn derive(&self) -> Self {
-		Self(Rc::new(ExprStoreData { exprs: RefCell::default(), parent: Some(self.clone()) }))
+	pub fn derive(&self, tracking_parent: bool) -> Self {
+		Self(Rc::new(ExprStoreData {
+			exprs: RefCell::default(),
+			parent: Some(self.clone()),
+			tracking_parent,
+		}))
 	}
 	pub fn give_expr(&self, expr: Expr) {
+		if self.0.tracking_parent {
+			self.0.parent.as_ref().unwrap().give_expr(expr.clone());
+		}
 		match self.0.exprs.borrow_mut().entry(expr.id()) {
 			Entry::Occupied(mut oe) => oe.get_mut().0 += 1,
 			Entry::Vacant(v) => {
@@ -29,8 +44,11 @@ impl ExprStore {
 		}
 	}
 	pub fn take_expr(&self, ticket: api::ExprTicket) -> Option<Expr> {
+		if self.0.tracking_parent {
+			self.0.parent.as_ref().unwrap().take_expr(ticket);
+		}
 		match self.0.exprs.borrow_mut().entry(ticket) {
-			Entry::Vacant(_) => None,
+			Entry::Vacant(_) => panic!("Attempted to double-take expression"),
 			Entry::Occupied(oe) if oe.get().0 == 1 => Some(oe.remove().1),
 			Entry::Occupied(mut oe) => {
 				oe.get_mut().0 -= 1;
@@ -43,6 +61,11 @@ impl ExprStore {
 		(self.0.exprs.borrow().get(&ticket).map(|(_, expr)| expr.clone()))
 			.or_else(|| self.0.parent.as_ref()?.get_expr(ticket))
 	}
+	pub fn iter(&self) -> impl Iterator<Item = (u32, Expr)> {
+		let r = Bound::new(self.clone(), |this| this.0.exprs.borrow());
+		let mut iter = Bound::new(r, |r| r.values());
+		std::iter::from_fn(move || iter.wrapped_mut().next().cloned())
+	}
 }
 impl fmt::Display for ExprStore {
 	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -51,3 +74,19 @@ impl fmt::Display for ExprStore {
 		write!(f, "Store holding {rc} refs to {} exprs", r.len())
 	}
 }
+impl Drop for ExprStore {
+	fn drop(&mut self) {
+		if 1 < Rc::strong_count(&self.0) {
+			return;
+		}
+		if !self.0.tracking_parent {
+			return;
+		}
+		let parent = self.0.parent.as_ref().unwrap();
+		for (id, (count, _)) in self.0.exprs.borrow().iter() {
+			for _ in 0..*count {
+				parent.take_expr(*id);
+			}
+		}
+	}
+}