feat: retain live page mounts and update shared UI references across aliases

src/ui_bridge.rs

@@ -1,27 +1,29 @@
 //! Tessera-owned UI payloads using Aster's existing embedding interface.
 use std::{
     any::Any,
-    collections::HashMap,
-    sync::{Arc, Mutex, OnceLock},
+    collections::{HashMap, HashSet},
+    sync::{Arc, Mutex},
 };
 
 use anyhow::{Result, bail};
 use endor_core::{EndorReal, ObjectPayload, ObjectRef, Value, Vm, VmError};
 
-use crate::document::parse_color;
+use crate::{
+    document::parse_color,
+    live_ui::{LiveKind, LiveNode},
+};
 
 pub const MAX_UI_NODES: usize = 4096;
 pub const MAX_UI_DEPTH: usize = 64;
 pub const MAX_PAGES: usize = 64;
 pub const MAX_DIMENSION: f64 = 1_000_000.0;
 
-#[derive(Debug)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct UiNode {
     pub kind: UiKind,
-    depth: usize,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone, PartialEq)]
 pub enum UiKind {
     Text {
         value: String,
@@ -50,7 +52,7 @@ pub enum UiKind {
 #[derive(Debug)]
 struct NodePayload(UiSlot);
 
-type UiSlot = Arc<OnceLock<Arc<UiNode>>>;
+type UiSlot = Arc<Mutex<Option<Arc<LiveNode>>>>;
 
 impl ObjectPayload for NodePayload {
     fn as_any(&self) -> &dyn Any {
@@ -59,63 +61,84 @@ impl ObjectPayload for NodePayload {
     fn kind_name(&self) -> &'static str {
         "tessera-ui-node"
     }
+    fn assign_from(&self, source: &dyn ObjectPayload) -> Result<bool, VmError> {
+        let Some(source) = source.as_any().downcast_ref::<Self>() else {
+            return Ok(false);
+        };
+        let old = self
+            .0
+            .lock()
+            .map_err(|_| VmError::InvalidArgument("UI slot is poisoned"))?
+            .clone()
+            .ok_or(VmError::InvalidArgument("UI reference is uninitialized"))?;
+        let new = source
+            .0
+            .lock()
+            .map_err(|_| VmError::InvalidArgument("UI slot is poisoned"))?
+            .clone()
+            .ok_or(VmError::InvalidArgument("UI source is uninitialized"))?;
+        old.replace(new).map_err(|_| {
+            VmError::InvalidArgument("UI assignment creates a cycle or exceeds graph limits")
+        })?;
+        Ok(true)
+    }
 }
 
 #[derive(Debug)]
 struct BuildSession {
     // Borrowed handles only. Candidate owns every pool root until extraction
-    // and execution have finished. Native child edges contain Arc<UiNode>,
+    // and execution have finished. Native child edges contain Arc<LiveNode>,
     // never VM handles, so payload-owned GC edges remain empty and immutable.
     pool: Vec<(ObjectRef, UiSlot)>,
     indexes: HashMap<ObjectRef, usize>,
-    used: usize,
+    free: Vec<usize>,
     pages: Vec<ObjectRef>,
-    pages_used: usize,
-    rendered: Option<(usize, Arc<UiNode>)>,
+    initialized_pages: Vec<bool>,
+    page_functions: Vec<(ObjectRef, ObjectRef)>,
+    native_roots: HashSet<usize>,
+    mounts: Vec<Option<Arc<LiveNode>>>,
+    active: Option<usize>,
+    pending: Option<(usize, Option<Arc<LiveNode>>)>,
+    last_snapshot: Option<Arc<UiNode>>,
 }
 
 impl BuildSession {
     fn page(&mut self) -> Result<Value, VmError> {
-        let page = self
-            .pages
-            .get(self.pages_used)
+        let index = self
+            .initialized_pages
+            .iter()
+            .position(|used| !*used)
             .ok_or(VmError::InvalidArgument(
-                "UI build exceeds the 64 page budget",
+                "UI session exceeds 64 retained pages",
             ))?;
-        self.pages_used += 1;
-        Ok(Value::Object(*page))
+        self.initialized_pages[index] = true;
+        Ok(Value::Object(self.pages[index]))
     }
-    fn child(&self, value: &Value) -> Result<Arc<UiNode>, VmError> {
+    fn child(&self, value: &Value) -> Result<Arc<LiveNode>, VmError> {
         let Value::Object(reference) = value else {
             return Err(VmError::InvalidArgument("UI children must be UI nodes"));
         };
-        self.indexes
-            .get(reference)
-            .and_then(|index| self.pool[*index].1.get())
-            .cloned()
-            .ok_or(VmError::InvalidArgument(
-                "UI children must be initialized nodes from this build",
-            ))
+        let index = self.indexes.get(reference).ok_or(VmError::InvalidArgument(
+            "UI child belongs to another session",
+        ))?;
+        self.pool[*index]
+            .1
+            .lock()
+            .map_err(|_| VmError::InvalidArgument("UI slot is poisoned"))?
+            .clone()
+            .ok_or(VmError::InvalidArgument("UI child is uninitialized"))
     }
-
-    fn emit(&mut self, kind: UiKind) -> Result<Value, VmError> {
-        let depth = match &kind {
-            UiKind::Column { children, .. } | UiKind::Row { children, .. } => {
-                children.iter().map(|node| node.depth).max().unwrap_or(0) + 1
-            }
-            _ => 1,
-        };
-        if depth > MAX_UI_DEPTH {
-            return Err(VmError::InvalidArgument("UI tree exceeds 64 levels"));
-        }
-        let Some((reference, slot)) = self.pool.get(self.used) else {
-            return Err(VmError::InvalidArgument(
-                "UI build exceeds the 4096 node budget",
-            ));
-        };
-        slot.set(Arc::new(UiNode { kind, depth }))
-            .map_err(|_| VmError::InvalidArgument("UI node initialized twice"))?;
-        self.used += 1;
+    fn emit(&mut self, kind: LiveKind) -> Result<Value, VmError> {
+        let node = LiveNode::new(kind);
+        node.snapshot()
+            .map_err(|_| VmError::InvalidArgument("UI graph exceeds depth or node limits"))?;
+        let index = self.free.pop().ok_or(VmError::InvalidArgument(
+            "UI session exceeds 4096 retained nodes",
+        ))?;
+        let (reference, slot) = &self.pool[index];
+        *slot
+            .lock()
+            .map_err(|_| VmError::InvalidArgument("UI slot is poisoned"))? = Some(node);
         Ok(Value::Object(*reference))
     }
 }
@@ -201,18 +224,18 @@ impl ObjectPayload for UiConstructor {
                 } else {
                     1
                 };
-                UiKind::Text {
+                LiveKind::Leaf(UiKind::Text {
                     value: value.clone(),
                     size,
                     color: color(arguments.get(2), "#242422")?,
-                }
+                })
             }
-            Constructor::Rect => UiKind::Rect {
+            Constructor::Rect => LiveKind::Leaf(UiKind::Rect {
                 width: dimension(&arguments[0])?,
                 height: dimension(&arguments[1])?,
                 radius: arguments.get(2).map(dimension).transpose()?.unwrap_or(0.0),
                 color: color(arguments.get(3), "#DCE8FF")?,
-            },
+            }),
             Constructor::Column | Constructor::Row => {
                 let padding = dimension(&arguments[0])?;
                 let gap = dimension(&arguments[1])?;
@@ -221,20 +244,20 @@ impl ObjectPayload for UiConstructor {
                     .map(|value| session.child(value))
                     .collect::<Result<Vec<_>, _>>()?;
                 if matches!(constructor, Constructor::Column) {
-                    UiKind::Column {
+                    LiveKind::Column {
                         padding,
                         gap,
                         children,
                     }
                 } else {
-                    UiKind::Row {
+                    LiveKind::Row {
                         padding,
                         gap,
                         children,
                     }
                 }
             }
-            Constructor::Spacer => UiKind::Spacer(dimension(&arguments[0])?),
+            Constructor::Spacer => LiveKind::Leaf(UiKind::Spacer(dimension(&arguments[0])?)),
             Constructor::Page => return session.page(),
         };
         session.emit(kind)
@@ -282,11 +305,45 @@ impl ObjectPayload for PageRender {
             .session
             .lock()
             .map_err(|_| VmError::InvalidArgument("UI build session is poisoned"))?;
-        if self.page >= session.pages_used {
+        if !session.initialized_pages[self.page] {
             return Err(VmError::InvalidArgument("page is not initialized"));
         }
         let root = session.child(&arguments[0])?;
-        session.rendered = Some((self.page, root));
+        session.pending = Some((self.page, Some(root)));
+        Ok(Value::Nil)
+    }
+}
+
+#[derive(Debug)]
+struct PageUnmount {
+    page: usize,
+    session: Arc<Mutex<BuildSession>>,
+}
+impl ObjectPayload for PageUnmount {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+    fn kind_name(&self) -> &'static str {
+        "tessera-page-unmount"
+    }
+    fn call(&self, arguments: &[Value]) -> Result<Value, VmError> {
+        if !arguments.is_empty() {
+            return Err(VmError::InvalidArgument("page.unmount takes no arguments"));
+        }
+        let mut session = self
+            .session
+            .lock()
+            .map_err(|_| VmError::InvalidArgument("UI session is poisoned"))?;
+        if session.active == Some(self.page)
+            || session
+                .pending
+                .as_ref()
+                .is_some_and(|(page, _)| *page == self.page)
+        {
+            session.pending = Some((self.page, None));
+        } else {
+            session.mounts[self.page] = None;
+        }
         Ok(Value::Nil)
     }
 }
@@ -298,20 +355,117 @@ pub struct UiBridge {
 }
 
 impl UiBridge {
-    /// Drain this input's render request; retained UI objects remain in the session pool.
-    pub fn take_rendered(&self) -> Result<Option<Arc<UiNode>>> {
+    /// Native mount/child edges retain explicit Aster roots. Drop those roots
+    /// when the native edge disappears, then recycle only pool-only VM slots.
+    pub fn reclaim(&self, vm: &mut Vm, roots: &mut Vec<ObjectRef>) -> Result<()> {
         let mut session = self
             .session
             .lock()
-            .map_err(|_| anyhow::anyhow!("UI build session is poisoned"))?;
-        Ok(session.rendered.take().map(|(_, root)| root))
+            .map_err(|_| anyhow::anyhow!("UI session is poisoned"))?;
+        loop {
+            let mut freed = false;
+            for index in 0..session.pool.len() {
+                let handle = session.pool[index].0;
+                let referenced = session.pool[index]
+                    .1
+                    .lock()
+                    .map_err(|_| anyhow::anyhow!("UI slot is poisoned"))?
+                    .as_ref()
+                    .is_some_and(|node| Arc::strong_count(node) > 1);
+                if referenced && !session.native_roots.contains(&index) {
+                    vm.retain(handle)?;
+                    roots.push(handle);
+                    session.native_roots.insert(index);
+                } else if !referenced && session.native_roots.remove(&index) {
+                    let at = roots
+                        .iter()
+                        .rposition(|root| *root == handle)
+                        .expect("native edge owns a root");
+                    roots.remove(at);
+                    vm.release(handle)?;
+                }
+                if !referenced && vm.ref_count(handle)? == 1 {
+                    let removed = session.pool[index]
+                        .1
+                        .lock()
+                        .map_err(|_| anyhow::anyhow!("UI slot is poisoned"))?
+                        .take()
+                        .is_some();
+                    if removed {
+                        session.free.push(index);
+                        freed = true;
+                    }
+                }
+            }
+            if !freed {
+                break;
+            }
+        }
+        for index in 0..session.pages.len() {
+            let (render, unmount) = session.page_functions[index];
+            if session.initialized_pages[index]
+                && session.active != Some(index)
+                && session.mounts[index].is_none()
+                && session
+                    .pending
+                    .as_ref()
+                    .is_none_or(|(page, _)| *page != index)
+                && vm.ref_count(session.pages[index])? == 1
+                && vm.ref_count(render)? == 2
+                && vm.ref_count(unmount)? == 2
+            {
+                session.initialized_pages[index] = false;
+            }
+        }
+        Ok(())
+    }
+    pub fn discard_pending(&self) -> Result<()> {
+        self.session
+            .lock()
+            .map_err(|_| anyhow::anyhow!("UI session is poisoned"))?
+            .pending = None;
+        Ok(())
+    }
+    /// Commit a mount only when layout/publication can succeed. Return None for
+    /// unchanged pages, Some(None) for an unmounted page, or an updated snapshot.
+    pub fn prepare_frame(&self) -> Result<Option<Option<Arc<UiNode>>>> {
+        let session = self
+            .session
+            .lock()
+            .map_err(|_| anyhow::anyhow!("UI session is poisoned"))?;
+        let root = match &session.pending {
+            Some((_, root)) => root.as_ref(),
+            None => session
+                .active
+                .and_then(|index| session.mounts[index].as_ref()),
+        };
+        let snapshot = root.map(|root| root.snapshot()).transpose()?;
+        let changed = session.pending.is_some() || snapshot != session.last_snapshot;
+        Ok(changed.then_some(snapshot))
+    }
+    pub fn commit_frame(&self, snapshot: Option<Arc<UiNode>>) -> Result<()> {
+        let mut session = self
+            .session
+            .lock()
+            .map_err(|_| anyhow::anyhow!("UI session is poisoned"))?;
+        if let Some((page, root)) = session.pending.take() {
+            session.active = root.as_ref().map(|_| page);
+            session.mounts[page] = root;
+        }
+        session.last_snapshot = snapshot;
+        Ok(())
     }
     pub fn rendered(&self) -> Result<Option<Arc<UiNode>>> {
         let session = self
             .session
             .lock()
-            .map_err(|_| anyhow::anyhow!("UI build session is poisoned"))?;
-        Ok(session.rendered.as_ref().map(|(_, root)| root.clone()))
+            .map_err(|_| anyhow::anyhow!("UI session is poisoned"))?;
+        session
+            .pending
+            .as_ref()
+            .and_then(|(_, root)| root.as_ref())
+            .map(|root| root.snapshot())
+            .transpose()
     }
 }
 
@@ -352,7 +506,7 @@ pub fn install(vm: &mut Vm, roots: &mut Vec<ObjectRef>) -> Result<UiBridge> {
     let mut pool = Vec::with_capacity(MAX_UI_NODES);
     let mut indexes = HashMap::with_capacity(MAX_UI_NODES);
     for index in 0..MAX_UI_NODES {
-        let slot = Arc::new(OnceLock::new());
+        let slot = Arc::new(Mutex::new(None));
         let reference = vm.allocate(NodePayload(slot.clone()))?;
         roots.push(reference);
         indexes.insert(reference, index);
@@ -361,10 +515,15 @@ pub fn install(vm: &mut Vm, roots: &mut Vec<ObjectRef>) -> Result<UiBridge> {
     let session = Arc::new(Mutex::new(BuildSession {
         pool,
         indexes,
-        used: 0,
+        free: (0..MAX_UI_NODES).rev().collect(),
         pages: Vec::with_capacity(MAX_PAGES),
-        pages_used: 0,
-        rendered: None,
+        initialized_pages: vec![false; MAX_PAGES],
+        page_functions: Vec::with_capacity(MAX_PAGES),
+        native_roots: HashSet::new(),
+        mounts: vec![None; MAX_PAGES],
+        active: None,
+        pending: None,
+        last_snapshot: None,
     }));
     for page in 0..MAX_PAGES {
         let reference = vm.allocate(PagePayload)?;
@@ -375,6 +534,18 @@ pub fn install(vm: &mut Vm, roots: &mut Vec<ObjectRef>) -> Result<UiBridge> {
         })?;
         roots.push(render);
         vm.set_attr(reference, "render", Value::Object(render))?;
+        let unmount = vm.allocate(PageUnmount {
+            page,
+            session: session.clone(),
+        })?;
+        roots.push(unmount);
+        vm.set_attr(reference, "unmount", Value::Object(unmount))?;
+        session
+            .lock()
+            .map_err(|_| anyhow::anyhow!("UI session is poisoned"))?
+            .page_functions
+            .push((render, unmount));
+
         session
             .lock()
             .map_err(|_| anyhow::anyhow!("UI build session is poisoned"))?
@@ -405,10 +576,11 @@ pub fn snapshot(vm: &Vm, result: &Value) -> Result<Arc<UiNode>> {
     let Value::Object(reference) = result else {
         bail!("main must return a UI node");
     };
-    Ok(vm
-        .payload::<NodePayload>(*reference, "Tessera UI node")?
+    vm.payload::<NodePayload>(*reference, "Tessera UI node")?
         .0
-        .get()
+        .lock()
+        .map_err(|_| anyhow::anyhow!("UI slot is poisoned"))?
+        .as_ref()
         .ok_or_else(|| anyhow::anyhow!("main returned an uninitialized UI node"))?
-        .clone())
+        .snapshot()
 }