dioxus/packages/core/src/scheduler.rs

/*
Welcome to Dioxus's cooperative, priority-based scheduler.

I hope you enjoy your stay.

Some essential reading:
- https://github.com/facebook/react/blob/main/packages/scheduler/src/forks/Scheduler.js#L197-L200
- https://github.com/facebook/react/blob/main/packages/scheduler/src/forks/Scheduler.js#L440
- https://github.com/WICG/is-input-pending
- https://web.dev/rail/
- https://indepth.dev/posts/1008/inside-fiber-in-depth-overview-of-the-new-reconciliation-algorithm-in-react

# What's going on?

Dioxus is a framework for "user experience" - not just "user interfaces." Part of the "experience" is keeping the UI
snappy and "jank free" even under heavy work loads. Dioxus already has the "speed" part figured out - but there's no
point in being "fast" if you can't also be "responsive."

As such, Dioxus can manually decide on what work is most important at any given moment in time. With a properly tuned
priority system, Dioxus can ensure that user interaction is prioritized and committed as soon as possible (sub 100ms).
The controller responsible for this priority management is called the "scheduler" and is responsible for juggling many
different types of work simultaneously.

# How does it work?

Per the RAIL guide, we want to make sure that A) inputs are handled ASAP and B) animations are not blocked.
React-three-fiber is a testament to how amazing this can be - a ThreeJS scene is threaded in between work periods of
React, and the UI still stays snappy!

While it's straightforward to run code ASAP and be as "fast as possible", what's not  _not_ straightforward is how to do
this while not blocking the main thread. The current prevailing thought is to stop working periodically so the browser
has time to paint and run animations. When the browser is finished, we can step in and continue our work.

React-Fiber uses the "Fiber" concept to achieve a pause-resume functionality. This is worth reading up on, but not
necessary to understand what we're doing here. In Dioxus, our DiffMachine is guided by DiffInstructions - essentially
"commands" that guide the Diffing algorithm through the tree. Our "diff_scope" method is async - we can literally pause
our DiffMachine "mid-sentence" (so to speak) by just stopping the poll on the future. The DiffMachine periodically yields
so Rust's async machinery can take over, allowing us to customize when exactly to pause it.

React's "should_yield" method is more complex than ours, and I assume we'll move in that direction as Dioxus matures. For
now, Dioxus just assumes a TimeoutFuture, and selects! on both the Diff algorithm and timeout. If the DiffMachine finishes
before the timeout, then Dioxus will work on any pending work in the interim. If there is no pending work, then the changes
are committed, and coroutines are polled during the idle period. However, if the timeout expires, then the DiffMachine
future is paused and saved (self-referentially).

# Priorty System

So far, we've been able to thread our Dioxus work between animation frames - the main thread is not blocked! But that
doesn't help us _under load_. How do we still stay snappy... even if we're doing a lot of work? Well, that's where
priorities come into play. The goal with priorities is to schedule shorter work as a "high" priority and longer work as
a "lower" priority. That way, we can interrupt long-running low-prioty work with short-running high-priority work.

React's priority system is quite complex.

There are 5 levels of priority and 2 distinctions between UI events (discrete, continuous). I believe React really only
uses 3 priority levels and "idle" priority isn't used... Regardless, there's some batching going on.

For Dioxus, we're going with a 4 tier priorty system:
- Sync: Things that need to be done by the next frame, like TextInput on controlled elements
- High: for events that block all others - clicks, keyboard, and hovers
- Medium: for UI events caused by the user but not directly - scrolls/forms/focus (all other events)
- Low: set_state called asynchronously, and anything generated by suspense

In "Sync" state, we abort our "idle wait" future, and resolve the sync queue immediately and escape. Because we completed
work before the next rAF, any edits can be immediately processed before the frame ends. Generally though, we want to leave
as much time to rAF as possible. "Sync" is currently only used by onInput - we'll leave some docs telling people not to
do anything too arduous from onInput.

For the rest, we defer to the rIC period and work down each queue from high to low.
*/
use crate::heuristics::*;
use crate::innerlude::*;
use futures_channel::mpsc::{UnboundedReceiver, UnboundedSender};
use futures_util::{pin_mut, stream::FuturesUnordered, Future, FutureExt, StreamExt};
use fxhash::FxHashMap;
use fxhash::FxHashSet;
use indexmap::IndexSet;
use slab::Slab;
use std::{
    any::{Any, TypeId},
    cell::{Cell, UnsafeCell},
    collections::{HashSet, VecDeque},
    rc::Rc,
};

#[derive(Clone)]
pub(crate) struct EventChannel {
    pub task_counter: Rc<Cell<u64>>,
    pub cur_subtree: Rc<Cell<u32>>,
    pub sender: UnboundedSender<SchedulerMsg>,
    pub schedule_any_immediate: Rc<dyn Fn(ScopeId)>,
    pub submit_task: Rc<dyn Fn(FiberTask) -> TaskHandle>,
    pub get_shared_context: GetSharedContext,
}

pub type GetSharedContext = Rc<dyn Fn(ScopeId, TypeId) -> Option<Rc<dyn Any>>>;

pub enum SchedulerMsg {
    // events from the host
    UiEvent(UserEvent),

    // setstate
    Immediate(ScopeId),

    // tasks
    Task(TaskMsg),
}

pub enum TaskMsg {
    // SubmitTask(FiberTask, u64),
    // SubmitTask(FiberTask, u64),
    ToggleTask(u64),
    PauseTask(u64),
    ResumeTask(u64),
    DropTask(u64),
}

/// The scheduler holds basically everything around "working"
///
/// Each scope has the ability to lightly interact with the scheduler (IE, schedule an update) but ultimately the scheduler calls the components.
///
/// In Dioxus, the scheduler provides 4 priority levels - each with their own "DiffMachine". The DiffMachine state can be saved if the deadline runs
/// out.
///
/// Saved DiffMachine state can be self-referential, so we need to be careful about how we save it. All self-referential data is a link between
/// pending DiffInstructions, Mutations, and their underlying Scope. It's okay for us to be self-referential with this data, provided we don't priority
/// task shift to a higher priority task that needs mutable access to the same scopes.
///
/// We can prevent this safety issue from occurring if we track which scopes are invalidated when starting a new task.
///
///
pub(crate) struct Scheduler {
    /// All mounted components are arena allocated to make additions, removals, and references easy to work with
    /// A generational arena is used to re-use slots of deleted scopes without having to resize the underlying arena.
    ///
    /// This is wrapped in an UnsafeCell because we will need to get mutable access to unique values in unique bump arenas
    /// and rusts's guartnees cannot prove that this is safe. We will need to maintain the safety guarantees manually.
    pub pool: ResourcePool,

    pub heuristics: HeuristicsEngine,

    pub receiver: UnboundedReceiver<SchedulerMsg>,

    // Garbage stored
    pub pending_garbage: FxHashSet<ScopeId>,

    // In-flight futures
    pub async_tasks: FuturesUnordered<FiberTask>,

    // // scheduler stuff
    // pub current_priority: EventPriority,
    pub ui_events: VecDeque<UserEvent>,

    pub pending_immediates: VecDeque<ScopeId>,

    pub pending_tasks: VecDeque<UserEvent>,

    pub batched_events: VecDeque<UserEvent>,

    pub garbage_scopes: HashSet<ScopeId>,

    pub dirty_scopes: IndexSet<ScopeId>,
    pub saved_state: Option<SavedDiffWork<'static>>,
    pub in_progress: bool,
}

impl Scheduler {
    pub(crate) fn new(
        sender: UnboundedSender<SchedulerMsg>,
        receiver: UnboundedReceiver<SchedulerMsg>,
    ) -> Self {
        /*
        Preallocate 2000 elements and 100 scopes to avoid dynamic allocation.
        Perhaps this should be configurable from some external config?
        */
        let components = Rc::new(UnsafeCell::new(Slab::with_capacity(100)));
        let raw_elements = Rc::new(UnsafeCell::new(Slab::with_capacity(2000)));

        let heuristics = HeuristicsEngine::new();

        let task_counter = Rc::new(Cell::new(0));
        let cur_subtree = Rc::new(Cell::new(0));

        let channel = EventChannel {
            cur_subtree,
            task_counter: task_counter.clone(),
            sender: sender.clone(),
            schedule_any_immediate: {
                let sender = sender.clone();
                Rc::new(move |id| sender.unbounded_send(SchedulerMsg::Immediate(id)).unwrap())
            },
            // todo: we want to get the futures out of the scheduler message
            // the scheduler message should be send/sync
            submit_task: {
                Rc::new(move |fiber_task| {
                    let task_id = task_counter.get();
                    task_counter.set(task_id + 1);

                    todo!();
                    // sender
                    //     .unbounded_send(SchedulerMsg::Task(TaskMsg::SubmitTask(
                    //         fiber_task, task_id,
                    //     )))
                    //     .unwrap();
                    TaskHandle {
                        our_id: task_id,
                        sender: sender.clone(),
                    }
                })
            },
            get_shared_context: {
                let components = components.clone();
                Rc::new(move |id, ty| {
                    let components = unsafe { &*components.get() };
                    let mut search: Option<&Scope> = components.get(id.0);
                    while let Some(inner) = search.take() {
                        if let Some(shared) = inner.shared_contexts.borrow().get(&ty) {
                            return Some(shared.clone());
                        } else {
                            search = inner.parent_idx.map(|id| components.get(id.0)).flatten();
                        }
                    }
                    None
                })
            },
        };

        let pool = ResourcePool {
            components,
            raw_elements,
            channel,
        };

        let async_tasks = FuturesUnordered::new();

        // push a task that would never resolve - prevents us from immediately aborting the scheduler
        async_tasks.push(Box::pin(async {
            std::future::pending::<()>().await;
            ScopeId(0)
        }) as FiberTask);

        let saved_state = SavedDiffWork {
            mutations: Mutations::new(),
            stack: DiffStack::new(),
            seen_scopes: Default::default(),
        };

        Self {
            pool,

            receiver,

            async_tasks,

            pending_garbage: FxHashSet::default(),

            heuristics,

            ui_events: VecDeque::new(),

            pending_immediates: VecDeque::new(),

            pending_tasks: VecDeque::new(),

            batched_events: VecDeque::new(),

            garbage_scopes: HashSet::new(),

            dirty_scopes: Default::default(),
            saved_state: Some(saved_state),
            in_progress: false,
        }
    }

    // returns true if the event is discrete
    pub fn handle_ui_event(&mut self, event: UserEvent) -> bool {
        let (discrete, priority) = event_meta(&event);

        if let Some(scope) = self.pool.get_scope_mut(event.scope) {
            if let Some(element) = event.mounted_dom_id {
                // TODO: bubble properly here
                scope.call_listener(event, element);

                while let Ok(Some(dirty_scope)) = self.receiver.try_next() {
                    //
                    //     self.add_dirty_scope(dirty_scope, trigger.priority)
                }
            }
        }

        // use EventPriority::*;

        // match priority {
        //     Immediate => todo!(),
        //     High => todo!(),
        //     Medium => todo!(),
        //     Low => todo!(),
        // }

        discrete
    }

    fn prepare_work(&mut self) {
        // while let Some(trigger) = self.ui_events.pop_back() {
        //     if let Some(scope) = self.pool.get_scope_mut(trigger.scope) {}
        // }
    }

    // nothing to do, no events on channels, no work
    pub fn has_any_work(&self) -> bool {
        !(self.dirty_scopes.is_empty() && self.ui_events.is_empty())
    }

    /// re-balance the work lanes, ensuring high-priority work properly bumps away low priority work
    fn balance_lanes(&mut self) {}

    fn save_work(&mut self, lane: SavedDiffWork) {
        let saved: SavedDiffWork<'static> = unsafe { std::mem::transmute(lane) };
        self.saved_state = Some(saved);
    }

    unsafe fn load_work(&mut self) -> SavedDiffWork<'static> {
        self.saved_state.take().unwrap().extend()
    }

    pub fn handle_channel_msg(&mut self, msg: SchedulerMsg) {
        match msg {
            //
            SchedulerMsg::Task(msg) => todo!(),

            SchedulerMsg::Immediate(_) => todo!(),

            SchedulerMsg::UiEvent(event) => {
                //

                let (discrete, priority) = event_meta(&event);

                if let Some(scope) = self.pool.get_scope_mut(event.scope) {
                    if let Some(element) = event.mounted_dom_id {
                        // TODO: bubble properly here
                        scope.call_listener(event, element);

                        while let Ok(Some(dirty_scope)) = self.receiver.try_next() {
                            //
                            //     self.add_dirty_scope(dirty_scope, trigger.priority)
                        }
                    }
                }

                discrete;
            }
        }
    }

    /// Load the current lane, and work on it, periodically checking in if the deadline has been reached.
    ///
    /// Returns true if the lane is finished before the deadline could be met.
    pub fn work_on_current_lane(
        &mut self,
        deadline_reached: impl FnMut() -> bool,
        mutations: &mut Vec<Mutations>,
    ) -> bool {
        // Work through the current subtree, and commit the results when it finishes
        // When the deadline expires, give back the work
        let saved_state = unsafe { self.load_work() };

        // We have to split away some parts of ourself - current lane is borrowed mutably
        let shared = self.pool.clone();
        let mut machine = unsafe { saved_state.promote(&shared) };

        let mut ran_scopes = FxHashSet::default();

        if machine.stack.is_empty() {
            let shared = self.pool.clone();

            self.dirty_scopes.sort_by(|a, b| {
                let h1 = shared.get_scope(*a).unwrap().height;
                let h2 = shared.get_scope(*b).unwrap().height;
                h1.cmp(&h2)
            });

            if let Some(scopeid) = self.dirty_scopes.pop() {
                log::info!("handlng dirty scope {:#?}", scopeid);
                if !ran_scopes.contains(&scopeid) {
                    ran_scopes.insert(scopeid);

                    let mut component = self.pool.get_scope_mut(scopeid).unwrap();
                    if component.run_scope(&self.pool) {
                        let (old, new) = (component.frames.wip_head(), component.frames.fin_head());
                        // let (old, new) = (component.frames.wip_head(), component.frames.fin_head());
                        machine.stack.scope_stack.push(scopeid);
                        machine.stack.push(DiffInstruction::Diff { new, old });
                    }
                }
            }
        }

        let work_completed = machine.work(deadline_reached);

        log::debug!("Working finished? {:?}", work_completed);

        // log::debug!("raw edits {:?}", machine.mutations.edits);

        let mut machine: DiffMachine<'static> = unsafe { std::mem::transmute(machine) };
        // let mut saved = machine.save();

        if work_completed {
            for node in machine.seen_scopes.drain() {
                // self.dirty_scopes.clear();
                // self.ui_events.clear();
                self.dirty_scopes.remove(&node);
                // self.dirty_scopes.remove(&node);
            }

            let mut new_mutations = Mutations::new();

            for edit in machine.mutations.edits.drain(..) {
                new_mutations.edits.push(edit);
            }

            // for edit in saved.edits.drain(..) {
            //     new_mutations.edits.push(edit);
            // }

            // std::mem::swap(&mut new_mutations, &mut saved.mutations);

            mutations.push(new_mutations);

            // log::debug!("saved edits {:?}", mutations);

            let mut saved = machine.save();
            self.save_work(saved);
            true

            // self.save_work(saved);
            // false
        } else {
            false
        }
    }

    /// The primary workhorse of the VirtualDOM.
    ///
    /// Uses some fairly complex logic to schedule what work should be produced.
    ///
    /// Returns a list of successful mutations.
    pub fn work_with_deadline<'a>(
        &'a mut self,
        mut deadline: impl FnMut() -> bool,
    ) -> Vec<Mutations<'a>> {
        /*
        Strategy:
        - When called, check for any UI events that might've been received since the last frame.
        - Dump all UI events into a "pending discrete" queue and a "pending continuous" queue.

        - If there are any pending discrete events, then elevate our priorty level. If our priority level is already "high,"
            then we need to finish the high priority work first. If the current work is "low" then analyze what scopes
            will be invalidated by this new work. If this interferes with any in-flight medium or low work, then we need
            to bump the other work out of the way, or choose to process it so we don't have any conflicts.
            'static components have a leg up here since their work can be re-used among multiple scopes.
            "High priority" is only for blocking! Should only be used on "clicks"

        - If there are no pending discrete events, then check for continuous events. These can be completely batched

        - we batch completely until we run into a discrete event
        - all continuous events are batched together
        - so D C C C C C would be two separate events - D and C. IE onclick and onscroll
        - D C C C C C C D C C C D would be D C D C D in 5 distinct phases.

        - !listener bubbling is not currently implemented properly and will need to be implemented somehow in the future
            - we need to keep track of element parents to be able to traverse properly


        Open questions:
        - what if we get two clicks from the component during the same slice?
            - should we batch?
            - react says no - they are continuous
            - but if we received both - then we don't need to diff, do we? run as many as we can and then finally diff?
        */
        let mut committed_mutations = Vec::<Mutations<'static>>::new();

        while self.has_any_work() {
            // switch our priority, pop off any work
            while let Some(event) = self.ui_events.pop_front() {
                if let Some(scope) = self.pool.get_scope_mut(event.scope) {
                    if let Some(element) = event.mounted_dom_id {
                        log::info!("Calling listener {:?}, {:?}", event.scope, element);

                        // TODO: bubble properly here
                        scope.call_listener(event, element);

                        while let Ok(Some(dirty_scope)) = self.receiver.try_next() {
                            match dirty_scope {
                                SchedulerMsg::Immediate(im) => {
                                    self.dirty_scopes.insert(im);
                                }
                                SchedulerMsg::UiEvent(e) => self.ui_events.push_back(e),
                                SchedulerMsg::Task(_) => todo!(),
                            }
                        }
                    }
                }
            }

            let work_complete = self.work_on_current_lane(&mut deadline, &mut committed_mutations);

            if !work_complete {
                return committed_mutations;
            }
        }

        log::debug!("work with deadline completed: {:#?}", committed_mutations);

        committed_mutations
    }

    /// Work the scheduler down, not polling any ongoing tasks.
    ///
    /// Will use the standard priority-based scheduling, batching, etc, but just won't interact with the async reactor.
    pub fn work_sync<'a>(&'a mut self) -> Vec<Mutations<'a>> {
        let mut committed_mutations = Vec::new();

        while let Ok(Some(msg)) = self.receiver.try_next() {
            self.handle_channel_msg(msg);
        }

        if !self.has_any_work() {
            return committed_mutations;
        }

        while self.has_any_work() {
            self.prepare_work();
            self.work_on_current_lane(|| false, &mut committed_mutations);
        }

        committed_mutations
    }

    /// Restart the entire VirtualDOM from scratch, wiping away any old state and components.
    ///
    /// Typically used to kickstart the VirtualDOM after initialization.
    pub fn rebuild(&mut self, base_scope: ScopeId) -> Mutations {
        let mut shared = self.pool.clone();
        let mut diff_machine = DiffMachine::new(Mutations::new(), &mut shared);

        // TODO: drain any in-flight work
        let cur_component = self
            .pool
            .get_scope_mut(base_scope)
            .expect("The base scope should never be moved");

        // We run the component. If it succeeds, then we can diff it and add the changes to the dom.
        if cur_component.run_scope(&self.pool) {
            diff_machine
                .stack
                .create_node(cur_component.frames.fin_head(), MountType::Append);

            diff_machine.stack.scope_stack.push(base_scope);

            diff_machine.work(|| false);
        } else {
            // todo: should this be a hard error?
            log::warn!(
                "Component failed to run succesfully during rebuild.
                This does not result in a failed rebuild, but indicates a logic failure within your app."
            );
        }

        unsafe { std::mem::transmute(diff_machine.mutations) }
    }

    pub fn hard_diff(&mut self, base_scope: ScopeId) -> Mutations {
        let cur_component = self
            .pool
            .get_scope_mut(base_scope)
            .expect("The base scope should never be moved");

        if cur_component.run_scope(&self.pool) {
            let mut diff_machine = DiffMachine::new(Mutations::new(), &mut self.pool);
            diff_machine.cfg.force_diff = true;
            diff_machine.diff_scope(base_scope);
            diff_machine.mutations
        } else {
            Mutations::new()
        }
    }
}