Merge pull request #794 from Demonthos/don't-box-scope-states

do not box scope states

packages/core/Cargo.toml

@@ -37,6 +37,8 @@ log = "0.4.17"
 # Serialize the Edits for use in Webview/Liveview instances
 serde = { version = "1", features = ["derive"], optional = true }
 
+bumpslab = { git = "https://github.com/jkelleyrtp/bumpslab" }
+
 [dev-dependencies]
 tokio = { version = "1", features = ["full"] }
 dioxus = { path = "../dioxus" }

packages/core/src/arena.rs

@@ -92,21 +92,21 @@ impl VirtualDom {
     // Note: This will not remove any ids from the arena
     pub(crate) fn drop_scope(&mut self, id: ScopeId, recursive: bool) {
         self.dirty_scopes.remove(&DirtyScope {
-            height: self.scopes[id.0].height,
+            height: self.scopes[id].height,
             id,
         });
 
         self.ensure_drop_safety(id);
 
         if recursive {
-            if let Some(root) = self.scopes[id.0].as_ref().try_root_node() {
+            if let Some(root) = self.scopes[id].try_root_node() {
                 if let RenderReturn::Ready(node) = unsafe { root.extend_lifetime_ref() } {
                     self.drop_scope_inner(node)
                 }
             }
         }
 
-        let scope = &mut self.scopes[id.0];
+        let scope = &mut self.scopes[id];
 
         // Drop all the hooks once the children are dropped
         // this means we'll drop hooks bottom-up
@@ -119,7 +119,7 @@ impl VirtualDom {
             scope.tasks.remove(task_id);
         }
 
-        self.scopes.remove(id.0);
+        self.scopes.remove(id);
     }
 
     fn drop_scope_inner(&mut self, node: &VNode) {
@@ -140,7 +140,7 @@ impl VirtualDom {
 
     /// Descend through the tree, removing any borrowed props and listeners
     pub(crate) fn ensure_drop_safety(&self, scope_id: ScopeId) {
-        let scope = &self.scopes[scope_id.0];
+        let scope = &self.scopes[scope_id];
 
         // make sure we drop all borrowed props manually to guarantee that their drop implementation is called before we
         // run the hooks (which hold an &mut Reference)

packages/core/src/create.rs

@@ -535,7 +535,7 @@ impl<'b> VirtualDom {
         }
 
         // If running the scope has collected some leaves and *this* component is a boundary, then handle the suspense
-        let boundary = match self.scopes[scope.0].has_context::<Rc<SuspenseContext>>() {
+        let boundary = match self.scopes[scope].has_context::<Rc<SuspenseContext>>() {
             Some(boundary) => boundary,
             _ => return created,
         };
@@ -544,7 +544,7 @@ impl<'b> VirtualDom {
         let new_id = self.next_element(new, parent.template.get().node_paths[idx]);
 
         // Now connect everything to the boundary
-        self.scopes[scope.0].placeholder.set(Some(new_id));
+        self.scopes[scope].placeholder.set(Some(new_id));
 
         // This involves breaking off the mutations to this point, and then creating a new placeholder for the boundary
         // Note that we break off dynamic mutations only - since static mutations aren't rendered immediately
@@ -583,7 +583,7 @@ impl<'b> VirtualDom {
         let new_id = self.next_element(template, template.template.get().node_paths[idx]);
 
         // Set the placeholder of the scope
-        self.scopes[scope.0].placeholder.set(Some(new_id));
+        self.scopes[scope].placeholder.set(Some(new_id));
 
         // Since the placeholder is already in the DOM, we don't create any new nodes
         self.mutations.push(AssignId {

packages/core/src/diff.rs

@@ -15,7 +15,7 @@ use DynamicNode::*;
 
 impl<'b> VirtualDom {
     pub(super) fn diff_scope(&mut self, scope: ScopeId) {
-        let scope_state = &mut self.scopes[scope.0];
+        let scope_state = &mut self.scopes[scope];
 
         self.scope_stack.push(scope);
         unsafe {
@@ -202,7 +202,7 @@ impl<'b> VirtualDom {
         right.scope.set(Some(scope_id));
 
         // copy out the box for both
-        let old = self.scopes[scope_id.0].props.as_ref();
+        let old = self.scopes[scope_id].props.as_ref();
         let new: Box<dyn AnyProps> = right.props.take().unwrap();
         let new: Box<dyn AnyProps> = unsafe { std::mem::transmute(new) };
 
@@ -214,14 +214,14 @@ impl<'b> VirtualDom {
         }
 
         // First, move over the props from the old to the new, dropping old props in the process
-        self.scopes[scope_id.0].props = Some(new);
+        self.scopes[scope_id].props = Some(new);
 
         // Now run the component and diff it
         self.run_scope(scope_id);
         self.diff_scope(scope_id);
 
         self.dirty_scopes.remove(&DirtyScope {
-            height: self.scopes[scope_id.0].height,
+            height: self.scopes[scope_id].height,
             id: scope_id,
         });
     }
@@ -721,7 +721,7 @@ impl<'b> VirtualDom {
 
                     Component(comp) => {
                         let scope = comp.scope.get().unwrap();
-                        match unsafe { self.scopes[scope.0].root_node().extend_lifetime_ref() } {
+                        match unsafe { self.scopes[scope].root_node().extend_lifetime_ref() } {
                             RenderReturn::Ready(node) => self.push_all_real_nodes(node),
                             RenderReturn::Aborted(_node) => todo!(),
                             _ => todo!(),
@@ -923,14 +923,14 @@ impl<'b> VirtualDom {
             .expect("VComponents to always have a scope");
 
         // Remove the component from the dom
-        match unsafe { self.scopes[scope.0].root_node().extend_lifetime_ref() } {
+        match unsafe { self.scopes[scope].root_node().extend_lifetime_ref() } {
             RenderReturn::Ready(t) => self.remove_node(t, gen_muts),
             RenderReturn::Aborted(placeholder) => self.remove_placeholder(placeholder, gen_muts),
             _ => todo!(),
         };
 
         // Restore the props back to the vcomponent in case it gets rendered again
-        let props = self.scopes[scope.0].props.take();
+        let props = self.scopes[scope].props.take();
         *comp.props.borrow_mut() = unsafe { std::mem::transmute(props) };
 
         // Now drop all the resouces
@@ -945,7 +945,7 @@ impl<'b> VirtualDom {
             Some(Placeholder(t)) => t.id.get().unwrap(),
             Some(Component(comp)) => {
                 let scope = comp.scope.get().unwrap();
-                match unsafe { self.scopes[scope.0].root_node().extend_lifetime_ref() } {
+                match unsafe { self.scopes[scope].root_node().extend_lifetime_ref() } {
                     RenderReturn::Ready(t) => self.find_first_element(t),
                     _ => todo!("cannot handle nonstandard nodes"),
                 }
@@ -961,7 +961,7 @@ impl<'b> VirtualDom {
             Some(Placeholder(t)) => t.id.get().unwrap(),
             Some(Component(comp)) => {
                 let scope = comp.scope.get().unwrap();
-                match unsafe { self.scopes[scope.0].root_node().extend_lifetime_ref() } {
+                match unsafe { self.scopes[scope].root_node().extend_lifetime_ref() } {
                     RenderReturn::Ready(t) => self.find_last_element(t),
                     _ => todo!("cannot handle nonstandard nodes"),
                 }

packages/core/src/scheduler/wait.rs

@@ -31,7 +31,7 @@ impl VirtualDom {
         // If the task completes...
         if task.task.borrow_mut().as_mut().poll(&mut cx).is_ready() {
             // Remove it from the scope so we dont try to double drop it when the scope dropes
-            let scope = &self.scopes[task.scope.0];
+            let scope = &self.scopes[task.scope];
             scope.spawned_tasks.borrow_mut().remove(&id);
 
             // Remove it from the scheduler
@@ -40,7 +40,7 @@ impl VirtualDom {
     }
 
     pub(crate) fn acquire_suspense_boundary(&self, id: ScopeId) -> Rc<SuspenseContext> {
-        self.scopes[id.0]
+        self.scopes[id]
             .consume_context::<Rc<SuspenseContext>>()
             .unwrap()
     }
@@ -64,7 +64,7 @@ impl VirtualDom {
         if let Poll::Ready(new_nodes) = as_pinned_mut.poll_unpin(&mut cx) {
             let fiber = self.acquire_suspense_boundary(leaf.scope_id);
 
-            let scope = &self.scopes[scope_id.0];
+            let scope = &self.scopes[scope_id];
             let arena = scope.current_frame();
 
             let ret = arena.bump().alloc(match new_nodes {

packages/core/src/scope_arena.rs

@@ -24,9 +24,9 @@ impl VirtualDom {
         let parent = self.acquire_current_scope_raw();
         let entry = self.scopes.vacant_entry();
         let height = unsafe { parent.map(|f| (*f).height + 1).unwrap_or(0) };
-        let id = ScopeId(entry.key());
+        let id = entry.key();
 
-        entry.insert(Box::new(ScopeState {
+        entry.insert(ScopeState {
             parent,
             id,
             height,
@@ -44,13 +44,13 @@ impl VirtualDom {
             shared_contexts: Default::default(),
             borrowed_props: Default::default(),
             attributes_to_drop: Default::default(),
-        }))
+        })
     }
 
     fn acquire_current_scope_raw(&self) -> Option<*const ScopeState> {
         let id = self.scope_stack.last().copied()?;
-        let scope = self.scopes.get(id.0)?;
-        Some(scope.as_ref())
+        let scope = self.scopes.get(id)?;
+        Some(scope)
     }
 
     pub(crate) fn run_scope(&mut self, scope_id: ScopeId) -> &RenderReturn {
@@ -60,9 +60,9 @@ impl VirtualDom {
         self.ensure_drop_safety(scope_id);
 
         let mut new_nodes = unsafe {
-            self.scopes[scope_id.0].previous_frame().bump_mut().reset();
+            self.scopes[scope_id].previous_frame().bump_mut().reset();
 
-            let scope = &self.scopes[scope_id.0];
+            let scope = &self.scopes[scope_id];
 
             scope.hook_idx.set(0);
 
@@ -127,7 +127,7 @@ impl VirtualDom {
             }
         };
 
-        let scope = &self.scopes[scope_id.0];
+        let scope = &self.scopes[scope_id];
 
         // We write on top of the previous frame and then make it the current by pushing the generation forward
         let frame = scope.previous_frame();

packages/core/src/scopes.rs

@@ -10,12 +10,15 @@ use crate::{
     AnyValue, Attribute, AttributeValue, Element, Event, Properties, TaskId,
 };
 use bumpalo::{boxed::Box as BumpBox, Bump};
+use bumpslab::{BumpSlab, Slot};
 use rustc_hash::{FxHashMap, FxHashSet};
+use slab::{Slab, VacantEntry};
 use std::{
     any::{Any, TypeId},
     cell::{Cell, RefCell},
     fmt::{Arguments, Debug},
     future::Future,
+    ops::{Index, IndexMut},
     rc::Rc,
     sync::Arc,
 };
@@ -63,6 +66,95 @@ impl<'a, T> std::ops::Deref for Scoped<'a, T> {
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug, PartialOrd, Ord)]
 pub struct ScopeId(pub usize);
 
+/// A thin wrapper around a BumpSlab that uses ids to index into the slab.
+pub(crate) struct ScopeSlab {
+    slab: BumpSlab<ScopeState>,
+    // a slab of slots of stable pointers to the ScopeState in the bump slab
+    entries: Slab<Slot<'static, ScopeState>>,
+}
+
+impl Drop for ScopeSlab {
+    fn drop(&mut self) {
+        // Bump slab doesn't drop its contents, so we need to do it manually
+        for slot in self.entries.drain() {
+            self.slab.remove(slot);
+        }
+    }
+}
+
+impl Default for ScopeSlab {
+    fn default() -> Self {
+        Self {
+            slab: BumpSlab::new(),
+            entries: Slab::new(),
+        }
+    }
+}
+
+impl ScopeSlab {
+    pub(crate) fn get(&self, id: ScopeId) -> Option<&ScopeState> {
+        self.entries.get(id.0).map(|slot| unsafe { &*slot.ptr() })
+    }
+
+    pub(crate) fn get_mut(&mut self, id: ScopeId) -> Option<&mut ScopeState> {
+        self.entries
+            .get(id.0)
+            .map(|slot| unsafe { &mut *slot.ptr_mut() })
+    }
+
+    pub(crate) fn vacant_entry(&mut self) -> ScopeSlabEntry {
+        let entry = self.entries.vacant_entry();
+        ScopeSlabEntry {
+            slab: &mut self.slab,
+            entry,
+        }
+    }
+
+    pub(crate) fn remove(&mut self, id: ScopeId) {
+        self.slab.remove(self.entries.remove(id.0));
+    }
+
+    pub(crate) fn contains(&self, id: ScopeId) -> bool {
+        self.entries.contains(id.0)
+    }
+
+    pub(crate) fn iter(&self) -> impl Iterator<Item = &ScopeState> {
+        self.entries.iter().map(|(_, slot)| unsafe { &*slot.ptr() })
+    }
+}
+
+pub(crate) struct ScopeSlabEntry<'a> {
+    slab: &'a mut BumpSlab<ScopeState>,
+    entry: VacantEntry<'a, Slot<'static, ScopeState>>,
+}
+
+impl<'a> ScopeSlabEntry<'a> {
+    pub(crate) fn key(&self) -> ScopeId {
+        ScopeId(self.entry.key())
+    }
+
+    pub(crate) fn insert(self, scope: ScopeState) -> &'a ScopeState {
+        let slot = self.slab.push(scope);
+        // this is safe because the slot is only ever accessed with the lifetime of the borrow of the slab
+        let slot = unsafe { std::mem::transmute(slot) };
+        let entry = self.entry.insert(slot);
+        unsafe { &*entry.ptr() }
+    }
+}
+
+impl Index<ScopeId> for ScopeSlab {
+    type Output = ScopeState;
+    fn index(&self, id: ScopeId) -> &Self::Output {
+        self.get(id).unwrap()
+    }
+}
+
+impl IndexMut<ScopeId> for ScopeSlab {
+    fn index_mut(&mut self, id: ScopeId) -> &mut Self::Output {
+        self.get_mut(id).unwrap()
+    }
+}
+
 /// A component's state separate from its props.
 ///
 /// This struct exists to provide a common interface for all scopes without relying on generics.

packages/core/src/virtual_dom.rs

@@ -5,7 +5,7 @@
 use crate::{
     any_props::VProps,
     arena::{ElementId, ElementRef},
-    innerlude::{DirtyScope, ErrorBoundary, Mutations, Scheduler, SchedulerMsg},
+    innerlude::{DirtyScope, ErrorBoundary, Mutations, Scheduler, SchedulerMsg, ScopeSlab},
     mutations::Mutation,
     nodes::RenderReturn,
     nodes::{Template, TemplateId},
@@ -177,7 +177,7 @@ use std::{any::Any, borrow::BorrowMut, cell::Cell, collections::BTreeSet, future
 pub struct VirtualDom {
     // Maps a template path to a map of byteindexes to templates
     pub(crate) templates: FxHashMap<TemplateId, FxHashMap<usize, Template<'static>>>,
-    pub(crate) scopes: Slab<Box<ScopeState>>,
+    pub(crate) scopes: ScopeSlab,
     pub(crate) dirty_scopes: BTreeSet<DirtyScope>,
     pub(crate) scheduler: Rc<Scheduler>,
 
@@ -258,7 +258,7 @@ impl VirtualDom {
             rx,
             scheduler: Scheduler::new(tx),
             templates: Default::default(),
-            scopes: Slab::default(),
+            scopes: Default::default(),
             elements: Default::default(),
             scope_stack: Vec::new(),
             dirty_scopes: BTreeSet::new(),
@@ -291,14 +291,14 @@ impl VirtualDom {
     ///
     /// This is useful for inserting or removing contexts from a scope, or rendering out its root node
     pub fn get_scope(&self, id: ScopeId) -> Option<&ScopeState> {
-        self.scopes.get(id.0).map(|f| f.as_ref())
+        self.scopes.get(id)
     }
 
     /// Get the single scope at the top of the VirtualDom tree that will always be around
     ///
     /// This scope has a ScopeId of 0 and is the root of the tree
     pub fn base_scope(&self) -> &ScopeState {
-        self.scopes.get(0).unwrap()
+        self.scopes.get(ScopeId(0)).unwrap()
     }
 
     /// Build the virtualdom with a global context inserted into the base scope
@@ -313,7 +313,7 @@ impl VirtualDom {
     ///
     /// Whenever the VirtualDom "works", it will re-render this scope
     pub fn mark_dirty(&mut self, id: ScopeId) {
-        if let Some(scope) = self.scopes.get(id.0) {
+        if let Some(scope) = self.scopes.get(id) {
             let height = scope.height;
             self.dirty_scopes.insert(DirtyScope { height, id });
         }
@@ -324,7 +324,7 @@ impl VirtualDom {
     /// This does not mean the scope is waiting on its own futures, just that the tree that the scope exists in is
     /// currently suspended.
     pub fn is_scope_suspended(&self, id: ScopeId) -> bool {
-        !self.scopes[id.0]
+        !self.scopes[id]
             .consume_context::<Rc<SuspenseContext>>()
             .unwrap()
             .waiting_on
@@ -499,7 +499,7 @@ impl VirtualDom {
     pub fn replace_template(&mut self, template: Template<'static>) {
         self.register_template_first_byte_index(template);
         // iterating a slab is very inefficient, but this is a rare operation that will only happen during development so it's fine
-        for (_, scope) in &self.scopes {
+        for scope in self.scopes.iter() {
             if let Some(RenderReturn::Ready(sync)) = scope.try_root_node() {
                 if sync.template.get().name.rsplit_once(':').unwrap().0
                     == template.name.rsplit_once(':').unwrap().0
@@ -583,7 +583,7 @@ impl VirtualDom {
         loop {
             // first, unload any complete suspense trees
             for finished_fiber in self.finished_fibers.drain(..) {
-                let scope = &self.scopes[finished_fiber.0];
+                let scope = &self.scopes[finished_fiber];
                 let context = scope.has_context::<Rc<SuspenseContext>>().unwrap();
 
                 self.mutations
@@ -607,7 +607,7 @@ impl VirtualDom {
                 self.dirty_scopes.remove(&dirty);
 
                 // If the scope doesn't exist for whatever reason, then we should skip it
-                if !self.scopes.contains(dirty.id.0) {
+                if !self.scopes.contains(dirty.id) {
                     continue;
                 }
 
@@ -626,7 +626,7 @@ impl VirtualDom {
                 // If suspended leaves are present, then we should find the boundary for this scope and attach things
                 // No placeholder necessary since this is a diff
                 if !self.collected_leaves.is_empty() {
-                    let mut boundary = self.scopes[dirty.id.0]
+                    let mut boundary = self.scopes[dirty.id]
                         .consume_context::<Rc<SuspenseContext>>()
                         .unwrap();