use crate::properties::RootProps;
use crate::root_wrapper::RootScopeWrapper;
use crate::{
arena::ElementId,
innerlude::{NoOpMutations, SchedulerMsg, ScopeOrder, ScopeState, VProps, WriteMutations},
runtime::{Runtime, RuntimeGuard},
scopes::ScopeId,
ComponentFunction, Element, Mutations,
};
use crate::{innerlude::Work, scopes::LastRenderedNode};
use crate::{Task, VComponent};
use futures_util::StreamExt;
use slab::Slab;
use std::collections::BTreeSet;
use std::{any::Any, rc::Rc};
use tracing::instrument;
pub struct VirtualDom {
pub(crate) scopes: Slab<ScopeState>,
pub(crate) dirty_scopes: BTreeSet<ScopeOrder>,
pub(crate) runtime: Rc<Runtime>,
pub(crate) resolved_scopes: Vec<ScopeId>,
rx: futures_channel::mpsc::UnboundedReceiver<SchedulerMsg>,
}
impl VirtualDom {
pub fn new(app: fn() -> Element) -> Self {
Self::new_with_props(app, ())
}
pub fn new_with_props<P: Clone + 'static, M: 'static>(
root: impl ComponentFunction<P, M>,
root_props: P,
) -> Self {
let render_fn = root.fn_ptr();
let props = VProps::new(root, |_, _| true, root_props, "Root");
Self::new_with_component(VComponent {
name: "root",
render_fn,
props: Box::new(props),
})
}
pub fn prebuilt(app: fn() -> Element) -> Self {
let mut dom = Self::new(app);
dom.rebuild_in_place();
dom
}
#[instrument(skip(root), level = "trace", name = "VirtualDom::new")]
pub(crate) fn new_with_component(root: VComponent) -> Self {
let (tx, rx) = futures_channel::mpsc::unbounded();
let mut dom = Self {
rx,
runtime: Runtime::new(tx),
scopes: Default::default(),
dirty_scopes: Default::default(),
resolved_scopes: Default::default(),
};
let root = VProps::new(
RootScopeWrapper,
|_, _| true,
RootProps(root),
"RootWrapper",
);
dom.new_scope(Box::new(root), "app");
#[cfg(debug_assertions)]
dom.register_subsecond_handler();
dom
}
pub fn get_scope(&self, id: ScopeId) -> Option<&ScopeState> {
self.scopes.get(id.0)
}
pub fn base_scope(&self) -> &ScopeState {
self.get_scope(ScopeId::ROOT).unwrap()
}
#[instrument(skip(self, f), level = "trace", name = "VirtualDom::in_runtime")]
pub fn in_runtime<O>(&self, f: impl FnOnce() -> O) -> O {
let _runtime = RuntimeGuard::new(self.runtime.clone());
f()
}
pub fn in_scope<T>(&self, scope: ScopeId, f: impl FnOnce() -> T) -> T {
self.runtime.in_scope(scope, f)
}
pub fn with_root_context<T: Clone + 'static>(self, context: T) -> Self {
self.base_scope().state().provide_context(context);
self
}
pub fn provide_root_context<T: Clone + 'static>(&self, context: T) {
self.base_scope().state().provide_context(context);
}
pub fn insert_any_root_context(&mut self, context: Box<dyn Any>) {
self.base_scope().state().provide_any_context(context);
}
pub fn mark_all_dirty(&mut self) {
let mut orders = vec![];
for (_idx, scope) in self.scopes.iter() {
orders.push(ScopeOrder::new(scope.state().height(), scope.id()));
}
for order in orders {
self.queue_scope(order);
}
}
pub fn mark_dirty(&mut self, id: ScopeId) {
let Some(scope) = self.runtime.try_get_state(id) else {
return;
};
tracing::event!(tracing::Level::TRACE, "Marking scope {:?} as dirty", id);
let order = ScopeOrder::new(scope.height(), id);
drop(scope);
self.queue_scope(order);
}
fn mark_task_dirty(&mut self, task: Task) {
let Some(scope) = self.runtime.task_scope(task) else {
return;
};
let Some(scope) = self.runtime.try_get_state(scope) else {
return;
};
tracing::event!(
tracing::Level::TRACE,
"Marking task {:?} (spawned in {:?}) as dirty",
task,
scope.id,
);
let order = ScopeOrder::new(scope.height(), scope.id);
drop(scope);
self.queue_task(task, order);
}
#[instrument(skip(self), level = "trace", name = "VirtualDom::wait_for_work")]
pub async fn wait_for_work(&mut self) {
loop {
self.process_events();
if self.has_dirty_scopes() {
return;
}
let _runtime = RuntimeGuard::new(self.runtime.clone());
self.wait_for_event().await;
}
}
#[instrument(skip(self), level = "trace", name = "VirtualDom::wait_for_event")]
async fn wait_for_event(&mut self) {
match self.rx.next().await.expect("channel should never close") {
SchedulerMsg::Immediate(id) => self.mark_dirty(id),
SchedulerMsg::TaskNotified(id) => {
self.mark_task_dirty(Task::from_id(id));
}
SchedulerMsg::EffectQueued => {}
SchedulerMsg::AllDirty => self.mark_all_dirty(),
};
}
fn queue_events(&mut self) {
while let Ok(msg) = self.rx.try_recv() {
match msg {
SchedulerMsg::Immediate(id) => self.mark_dirty(id),
SchedulerMsg::TaskNotified(task) => self.mark_task_dirty(Task::from_id(task)),
SchedulerMsg::EffectQueued => {}
SchedulerMsg::AllDirty => self.mark_all_dirty(),
}
}
}
#[instrument(skip(self), level = "trace", name = "VirtualDom::process_events")]
pub fn process_events(&mut self) {
self.queue_events();
if self.has_dirty_scopes() {
return;
}
self.poll_tasks()
}
#[instrument(skip(self), level = "trace", name = "VirtualDom::poll_tasks")]
fn poll_tasks(&mut self) {
let _runtime = RuntimeGuard::new(self.runtime.clone());
while !self.runtime.dirty_tasks.borrow().is_empty()
|| !self.runtime.pending_effects.borrow().is_empty()
{
while let Some(task) = self.pop_task() {
let _ = self.runtime.handle_task_wakeup(task);
self.queue_events();
if self.has_dirty_scopes() {
return;
}
}
while let Some(effect) = self.pop_effect() {
effect.run();
self.queue_events();
if self.has_dirty_scopes() {
return;
}
}
}
}
pub fn rebuild_in_place(&mut self) {
self.rebuild(&mut NoOpMutations);
}
pub fn rebuild_to_vec(&mut self) -> Mutations {
let mut mutations = Mutations::default();
self.rebuild(&mut mutations);
mutations
}
#[instrument(skip(self, to), level = "trace", name = "VirtualDom::rebuild")]
pub fn rebuild(&mut self, to: &mut impl WriteMutations) {
let _runtime = RuntimeGuard::new(self.runtime.clone());
let new_nodes = self
.runtime
.clone()
.while_rendering(|| self.run_scope(ScopeId::ROOT));
let new_nodes = LastRenderedNode::new(new_nodes);
self.scopes[ScopeId::ROOT.0].last_rendered_node = Some(new_nodes.clone());
let m = self.create_scope(Some(to), ScopeId::ROOT, new_nodes, None);
to.append_children(ElementId(0), m);
}
#[instrument(skip(self, to), level = "trace", name = "VirtualDom::render_immediate")]
pub fn render_immediate(&mut self, to: &mut impl WriteMutations) {
self.process_events();
let _runtime = RuntimeGuard::new(self.runtime.clone());
while let Some(work) = self.pop_work() {
match work {
Work::PollTask(task) => {
_ = self.runtime.handle_task_wakeup(task);
self.queue_events();
}
Work::RerunScope(scope) => {
self.runtime.clone().while_rendering(|| {
self.run_and_diff_scope(Some(to), scope.id);
});
}
}
}
self.runtime.finish_render();
}
pub fn render_immediate_to_vec(&mut self) -> Mutations {
let mut mutations = Mutations::default();
self.render_immediate(&mut mutations);
mutations
}
#[instrument(skip(self), level = "trace", name = "VirtualDom::wait_for_suspense")]
pub async fn wait_for_suspense(&mut self) {
loop {
self.queue_events();
if !self.suspended_tasks_remaining() && !self.has_dirty_scopes() {
break;
}
self.wait_for_suspense_work().await;
self.render_suspense_immediate().await;
}
}
pub fn suspended_tasks_remaining(&self) -> bool {
self.runtime.suspended_tasks.get() > 0
}
pub async fn wait_for_suspense_work(&mut self) {
loop {
self.queue_events();
if self.has_dirty_scopes() {
break;
}
{
let _runtime = RuntimeGuard::new(self.runtime.clone());
let mut tasks_polled = 0;
while let Some(task) = self.pop_task() {
if self.runtime.task_runs_during_suspense(task) {
let _ = self.runtime.handle_task_wakeup(task);
self.queue_events();
if self.has_dirty_scopes() {
return;
}
}
tasks_polled += 1;
if tasks_polled > 32 {
yield_now().await;
tasks_polled = 0;
}
}
}
self.wait_for_event().await;
}
}
pub async fn render_suspense_immediate(&mut self) -> Vec<ScopeId> {
self.queue_events();
let _runtime = RuntimeGuard::new(self.runtime.clone());
let mut work_done = 0;
while let Some(work) = self.pop_work() {
match work {
Work::PollTask(task) => {
if self.runtime.task_runs_during_suspense(task) {
let _ = self.runtime.handle_task_wakeup(task);
}
}
Work::RerunScope(scope) => {
let scope_id: ScopeId = scope.id;
let run_scope = self
.runtime
.try_get_state(scope.id)
.filter(|scope| scope.should_run_during_suspense())
.is_some();
if run_scope {
self.runtime.clone().while_rendering(|| {
self.run_and_diff_scope(None::<&mut NoOpMutations>, scope_id);
});
tracing::trace!("Ran scope {:?} during suspense", scope_id);
} else {
tracing::warn!(
"Scope {:?} was marked as dirty, but will not rerun during suspense. Only nodes that are under a suspense boundary rerun during suspense",
scope_id
);
}
}
}
self.queue_events();
work_done += 1;
if work_done > 32 {
yield_now().await;
work_done = 0;
}
}
self.resolved_scopes
.sort_by_key(|&id| self.runtime.get_state(id).height);
std::mem::take(&mut self.resolved_scopes)
}
pub fn runtime(&self) -> Rc<Runtime> {
self.runtime.clone()
}
#[deprecated = "Use [VirtualDom::runtime().handle_event] instead"]
pub fn handle_event(&self, name: &str, event: Rc<dyn Any>, element: ElementId, bubbling: bool) {
let event = crate::Event::new(event, bubbling);
self.runtime().handle_event(name, event, element);
}
#[cfg(debug_assertions)]
fn register_subsecond_handler(&self) {
let sender = self.runtime().sender.clone();
subsecond::register_handler(std::sync::Arc::new(move || {
_ = sender.unbounded_send(SchedulerMsg::AllDirty);
}));
}
}
impl Drop for VirtualDom {
fn drop(&mut self) {
let mut scopes = self.scopes.drain().collect::<Vec<_>>();
scopes.sort_by_key(|scope| scope.state().height);
for scope in scopes.into_iter().rev() {
drop(scope);
}
self.runtime.pending_effects.borrow_mut().clear();
self.runtime.tasks.borrow_mut().clear();
self.runtime.mounts.borrow_mut().clear();
}
}
async fn yield_now() {
let mut yielded = false;
std::future::poll_fn::<(), _>(move |cx| {
if !yielded {
cx.waker().wake_by_ref();
yielded = true;
std::task::Poll::Pending
} else {
std::task::Poll::Ready(())
}
})
.await;
}