cells.rs

#![feature(local_key_cell_methods)]

use std::{cell::RefCell, collections::HashMap, rc::Rc};

/// Traits are separated into RuntimeManaged and ParentManaged because
/// in the COMPONENTS thread_local we need to erase the prop type since we
/// don't know it there (and we don't need props there too).
/// 
/// Luckly for us Rust automatically downcasts a `Rc<RefCell<Component<T>>>` to `Rc<RefCell<dyn RuntimeManaged>>`
/// when we clone the a component to the thread_local, while still keeping the original concrete type 
/// in the parent so it can still pass props of the correct type.

trait RuntimeManaged {
    /// Renders the child components
    fn render(&mut self) {}
    
    /// Runs all effects related to this component
    fn effects(&mut self) -> bool {
        true
    }
}

trait ParentManaged<T> {
    /// Called when re-rendered by the parent.
    /// Props might be stil the same, so its important
    /// to check if they differ and return false to
    /// avoid rendering children again.
    fn update_props(&mut self, props: T) -> bool {
        drop(props);
        true
    }

    /// Called when mounted by the parent.
    fn new(props: T) -> Self where Self: Sized;
}

/// Just a marker trait.
trait Component<T>: RuntimeManaged + ParentManaged<T> {}

/// The wrapper type for all components in the thread_local
type FC<T> = Rc<RefCell<T>>;

/// The reference type used in the owning components
/// 
/// Its optional since a struct field needs to be alocated on 
/// the parent component, but it will just be populated on the `render`
/// method wich will run only on the next microtask cycle
type FCRef<T> = Option<FC<T>>;

////////////////////////////////////////////////////////////////////////////////

thread_local! {
    /// Last component ID generated
    static ID: RefCell<u64> = RefCell::new(0);

    /// The actual live component heap
    static COMPONENTS: RefCell<HashMap<u64, FC<dyn RuntimeManaged>>> = RefCell::new(HashMap::new());
    
    /// The only pourpose of this is to avoid crashes due to nested borrows
    static MICROTASKS: RefCell<Vec<Box<dyn FnOnce()>>> = RefCell::new(Vec::new());

    /// Unused but it could be a fun trick
    /// 
    /// Maybe we can use it to remove the `ref_1` property from the
    /// parent component and make it just pass an macro generated number
    /// to `render` instead
    static CURRENT_COMPONENT_ID: RefCell<Option<u64>> = RefCell::new(None);
}

/// Takes the component ID and runs a function
/// while managing `CURRENT_COMPONENT_ID` so when it
/// exits the thread local has a None value on it
fn with_component<F: FnOnce()>(node_id: u64, f: F) {
    CURRENT_COMPONENT_ID.with_borrow_mut(|value| *value = Some(node_id));
    f();
    CURRENT_COMPONENT_ID.with_borrow_mut(|value| *value = None);
}

/// Runs code in the context of the current component.
///
/// Crashes if there isn't a `with_component` up in the callstack
fn in_component_context<F: FnOnce(u64) -> R, R>(f: F) -> R {
    CURRENT_COMPONENT_ID.with_borrow(|value| {
        if let Some(id) = value {
            return f(*id);
        }

        panic!("Cannot execute fn in in_component_context: no component is being evaluated");
    })
}

/// Calls all deferred functions in the task queue
pub fn drain_microtasks() {
    while MICROTASKS.with(|f| !f.borrow().is_empty()) {
        let mut tasks = MICROTASKS.take();

        for task in tasks.drain(..) {
            task();
        }
    }
}

/// Inserts a task into the deferred task queue
pub fn add_microtask<F: FnOnce() + 'static>(task: F) {
    MICROTASKS.with_borrow_mut(|tasks| {
        tasks.push(Box::new(task));
    });
}

////////////////////////////////////////////////////////////////////////////////

/// Runs the effects of a component and possibly re-renders its children
/// whithout updating the props.
/// 
/// Will be usefull for hooks like `use_state/use_context` which can also trigger re-renders.
fn update_component(node_id: u64) {
    COMPONENTS.with_borrow_mut(|heap| {
        if let Some(existing) = heap.get_mut(&node_id) {
            let microtask_clone = existing.clone();

            add_microtask(move || {
                with_component(node_id, || {
                    if microtask_clone.borrow_mut().effects() {
                        microtask_clone.borrow_mut().render();
                    }
                });
            });
        } else {
            panic!("Component id {} is not allocated", node_id);
        }
    });
}

/// Unmounts the component, in such a way that its drop function will be called
fn unrender<C: Component<T> + 'static, T>(component: &mut (FCRef<C>, u64)) {
    COMPONENTS.with_borrow_mut(|heap| {
        heap.remove(&component.1);
    });

    component.1 = 0;
    
    drop(component.0.take());
}

/// Mounts a component, and schedules it's effects to be executed and it's children 
/// to also be rendered in the next microtask cycle.
/// 
/// If the component has already been mounted it will instead call update_props on it.
fn render<C: Component<T> + 'static, T: 'static>(component: &mut (FCRef<C>, u64), props: T) {
    if let Some(existing) = component.0.as_ref() {
        if existing.borrow_mut().update_props(props) {
            let microtask_clone = existing.clone();

            let node_id = component.1;

            add_microtask(move || {
                with_component(node_id, || {
                    if microtask_clone.borrow_mut().effects() {
                        microtask_clone.borrow_mut().render();
                    }
                });
            });
        }
    } else {
        let new = Rc::new(RefCell::new(C::new(props)));

        let node_id = ID.with_borrow_mut(|id| {
            *id += 1;
            *id
        });

        component.0 = Some(new.clone());
        component.1 = node_id;

        let microtask_clone = new.clone();

        let heap_clone = new.clone();

        COMPONENTS.with_borrow_mut(|components| components.insert(node_id, heap_clone));

        add_microtask(move || {
            with_component(node_id, || {
                if microtask_clone.borrow_mut().effects() {
                    microtask_clone.borrow_mut().render();
                }
            });
        });
    }
}

////////////////////////////////////////////////////////////////////////////////

fn use_context<T: Default>() -> T {
    in_component_context(|_| {
        // TODO
        Default::default()
    })
}

////////////////////////////////////////////////////////////////////////////////

#[derive(Default, Debug)]
struct Example(u64);

////////////////////////////////////////////////////////////////////////////////


/// Printer component example
/// 
/// Recevies a string as a prop and prints it 
/// only if it changes
struct Printer {
    prop: String,
}

impl Drop for Printer {
    fn drop(&mut self) {
        println!("Printer getting dropped from memory");
        // also drop context subscriptions
    }
}

impl RuntimeManaged for Printer {
    fn effects(&mut self) -> bool {
        let example = use_context::<Example>();
        println!("Printer::effects, prop: '{}' context: {:?} ", self.prop, example);
        false
    }
}

impl ParentManaged<String> for Printer {
    fn new(prop: String) -> Self {
        println!("A printer is getting mounted");
        Self { prop }
    }

    fn update_props(&mut self, prop: String) -> bool {
        if self.prop != prop {
            self.prop = prop;
            return true;
        }

        false
    }
}

impl Component<String> for Printer {}

////////////////////////////////////////////////////////////////////////

/// The parent component mounts a Printer component
/// 
/// It as a len prop that is a number
/// It will repeat a `"X"` string by that number and 
/// pass it to the printer component
/// 
/// If len is zero it unmounts the Printer component
/// while keeping itself mounted. It will re-mount Printer
/// if len ever comes back to being a non-zero value
struct Parent {
    len: u64,
    ref_1: (FCRef<Printer>, u64),
}

impl RuntimeManaged for Parent {
    /// fn render(&mut self) {
    ///    rsx! {
    ///       if self.len >= 1 {
    ///          <Printer prop={String::from("X").repeat(self.len as usize)} />
    ///       }
    ///    }
    /// }
    ///
    fn render(&mut self) {
        if self.len >= 1 {
            render(&mut self.ref_1, String::from("X").repeat(self.len as usize));
        } else {
            unrender(&mut self.ref_1);
        }
    }

    fn effects(&mut self) -> bool {
        true
    }
}

impl ParentManaged<u64> for Parent {
    fn new(len: u64) -> Self {
        println!("A Parent is getting mounted");
        Self {
            ref_1: (None, 0),
            len,
        }
    }
    fn update_props(&mut self, props: u64) -> bool {
        if self.len != props {
            self.len = props;
            // Only re-render children if props change
            return true;
        }

        false
    }
}

impl Drop for Parent {
    fn drop(&mut self) {
        // Drop must unmount any components Self mounts
        println!("Parent getting dropped from memory");
        unrender(&mut self.ref_1);
    }
}

impl Component<u64> for Parent {}

fn main() {
    let mut main = (None as FCRef<Parent>, 0);

    // First render, pass 3 as prop to main
    render(&mut main, 3);
    drain_microtasks();
    // No-op, just API usage to remove warnings
    update_component(main.1);

    // These re-renders shoudn't print anything
    render(&mut main, 3);
    drain_microtasks();
    render(&mut main, 3);
    drain_microtasks();
    
    // This should print since the props changed
    render(&mut main, 5);
    drain_microtasks();

    // This should print since the props changed again
    render(&mut main, 3);
    drain_microtasks();

    // But not now
    render(&mut main, 3);
    drain_microtasks();
    render(&mut main, 3);
    drain_microtasks();

    // This should make the Parent component Unmount the Printer component
    render(&mut main, 0);
    drain_microtasks();

    // The next re-renders shoudn't do anything
    render(&mut main, 0);
    drain_microtasks();
    render(&mut main, 0);
    drain_microtasks();

    // This should mount and render the Printer
    render(&mut main, 1);
    drain_microtasks();

    // Demonstration of safe runtime cleanup
    unrender(&mut main);
    drain_microtasks();

    // Last log to confirm that Printer and Parent
    // are droped before main ends since they could
    // still be leaked on thread_locals 
    println!("End main")
}