reax 0.2.0

use std::sync::atomic::{AtomicUsize, Ordering};
use std::mem;
use std::marker::PhantomData;
use std::io;
use std::fmt;
use std::cell::RefCell;
use fnv::FnvHashMap;
use crate::idset::IdSet;
use crate::handler::SharedDirtiedHandler;

thread_local! {
    static THREAD_SYSTEM: RefCell<Option<ThreadRuntime>> = RefCell::new(None);
}

#[cfg(debug_assertions)]
type NodeLabel = String;

#[cfg(not(debug_assertions))]
type NodeLabel = ();

struct NodeInfo {
    label: NodeLabel,
    used_by: IdSet,
    uses: IdSet,
    dirty: bool,
    on_dirty: Option<SharedDirtiedHandler>,
}

impl NodeInfo {
    fn new() -> Self {
        NodeInfo {
            label: NodeLabel::default(),
            used_by: IdSet::new(),
            uses: IdSet::new(),
            dirty: true, // All nodes are dirty by default.
            on_dirty: None,
        }
    }

    // This is a bit of a tricky function. That is sorta problematic because it
    // is really at the heart of reax. The general idea is that this updates the
    // dirty flag on this node and queues additional nodes to be similarly
    // updated.
    fn propagate_dirty(
        &mut self,
        id: usize,
        cause: usize,
        keep_cause_dirty: bool,
        to_dirty: &mut Vec<usize>,
    ) {
        // Nothing should be done by default. The below cases will decide if
        // downstream nodes actually need to be marked as dirty.
        let mut should_propagate = false;

        // Are we causing the propagation? That is, are we the node that was
        // originally changed/written to? We may encounter the cause several
        // times during propagation due to reference loops.
        if id == cause {
            // We are the cause! We are treated special.

            // Should we keep our existing dirty status?
            if keep_cause_dirty {
                // We will be left as dirty but propagation should continue as
                // per the comment on on_write. This is probably because we act
                // like a `Var` which is always dirty and never transitions
                // *to* dirty but still causes downstream dirtying.
                should_propagate = true;

                // Continuing propagation could be a problem if the node is part
                // of a loop since the dirty status is used to break back-edges
                // in the DFS. However, the `keep_cause_dirty` flag is really
                // only used by Var-like types that don't have upstream vars
                // and so can't be part of loops.
            } else {
                // We are unconditionally cleaned by the write.
                self.dirty = false;
                // Propagation is not needed because downstream vars would
                // already be dirty if they needed our new value.
            }
        } else if !self.dirty {
            // We are a downstream var transitioning to dirty from clean. Make
            // sure we continue propagation so that other downstream vars are
            // aware that the cause is causing sh*t.
            should_propagate = true;
            // Mark us as clean right now so that if propagation encounters this
            // var again we just ignore it.
            self.dirty = true;
            // Also send a notification so that watchers and stuff will fire.
            if let Some(handler) = self.on_dirty.as_mut() {
                handler.on_dirtied(id);
            }
        }

        if should_propagate {
            to_dirty.extend(self.used_by.iter());
        }
    }
}

#[cfg(debug_assertions)]
fn fmt_label(
    var: Option<&NodeInfo>,
    id: usize,
    f: &mut fmt::Formatter,
) -> fmt::Result {
    match var {
        Some(v) if !v.label.is_empty() => fmt::Debug::fmt(&v.label, f),
        _ => fmt::Debug::fmt(&id, f),
    }
}

#[cfg(not(debug_assertions))]
fn fmt_label(
    var: Option<&NodeInfo>,
    id: usize,
    f: &mut fmt::Formatter,
) -> fmt::Result {
    let _ = var.map(|info| info.label); // dead code warning
    fmt::Debug::fmt(&id, f)
}

/// The per-thread reax runtime responsible for tracking changes and
/// dependencies.
#[derive(Default)]
pub struct ThreadRuntime {
    id_buf: Vec<usize>,
    computing_stack: Vec<(usize, bool)>,
    nodes: FnvHashMap<usize, NodeInfo>,
}

impl ThreadRuntime {
    fn with<R>(func: impl FnOnce(&mut ThreadRuntime) -> R) -> R {
        THREAD_SYSTEM.with(|sys| func(sys.borrow_mut()
            .get_or_insert_with(ThreadRuntime::default)))
    }

    /// Replaces the runtime of the current thread with self, returning the old
    /// runtime. All information about existing nodes on this thread will be
    /// lost so be careful.
    pub fn replace_current(self) -> ThreadRuntime {
        ThreadRuntime::with(|sys| mem::replace(sys, self))
    }

    fn set_dependent_impl(&mut self, used_by: usize, uses: usize) {
        if uses == used_by {
            return;
        }

        self.nodes.entry(used_by)
            .or_insert_with(NodeInfo::new)
            .uses
            .insert(uses);

        self.nodes.entry(uses)
            .or_insert_with(NodeInfo::new)
            .used_by
            .insert(used_by);
    }

    /// Write out a simple [graphviz](https://www.graphviz.org/) DSL
    /// representation of the current thread's runtime state; complete with
    /// dependencies, dirty status, and labels. The exact syntax produced may
    /// change at any time so don't rely on it.
    pub fn write_graphviz(mut to: impl io::Write) -> io::Result<()> {
        ThreadRuntime::with(|sys| write!(to, "{:?}", sys))
    }
}

impl fmt::Debug for ThreadRuntime {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("digraph {\n")?;
        for (&id, var) in self.nodes.iter() {
            write!(f, "{} [label=", id)?;
            fmt_label(Some(&var), id, f)?;
            writeln!(f, ", color={}];", match var.dirty {
                true => "red",
                false => "black",
            })?;
        }
        for (&used_by_id, used_by_var) in self.nodes.iter() {
            for uses_id in used_by_var.uses.iter() {
                writeln!(f, "{} -> {};", uses_id, used_by_id)?;
            }
        }
        f.write_str("}\n")
    }
}

/// Offers greater control over how the reax runtime sees the computation of a
/// node.
///
/// Note that handles are isolated across nested computations. For example,
/// ```rust
/// let a = reax::Node::next();
/// let b = reax::Node::next();
/// let c = reax::Node::next();
///
/// a.computation(|handle| {
///     handle.set_reactive_area(false);
///     c.on_read(); // Does not count as dependency of a.
///
///     b.computation(|handle| {
///         assert_eq!(handle.is_reactive_area(), true);
///         c.on_read(); // Does count as a dependency of b.
///     });
///
///     assert_eq!(handle.is_reactive_area(), false);
/// });
///
/// assert_eq!(a.depends_on(&c), false);
/// assert_eq!(b.depends_on(&c), true);
/// ```
pub struct ComputationHandle(usize);

impl ComputationHandle {
    /// Any node accessed after this method is called will count as a dependency
    /// if and only if `should_react` is true.
    pub fn set_reactive_area(&self, should_react: bool) {
        ThreadRuntime::with(|sys| {
            sys.computing_stack
                .get_mut(self.0)
                .expect("computation stack corrupted")
                .1 = should_react;
        })
    }

    /// Checks whether node accesses made around this method call will count as
    /// dependencies.
    pub fn is_reactive_area(&self) -> bool {
        ThreadRuntime::with(|sys| {
            sys.computing_stack
                .get_mut(self.0)
                .expect("computation stack corrupted")
                .1
        })
    }
}

impl Drop for ComputationHandle {
    fn drop(&mut self) {
        ThreadRuntime::with(|sys| {
            assert_eq!(
                sys.computing_stack.len(),
                self.0 + 1,
                "computation stack corrupted",
            );
            sys.computing_stack.pop();
        })
    }
}

static NEXT_ID: AtomicUsize = AtomicUsize::new(0);

/// A node handle used to communicate with a [thread's reax
/// runtime](struct.ThreadRuntime.html).
///
/// This is the key primitive used to implement types like `Var` and
/// `ComputedVar`.
///
/// Although raw node ids are unique across threads, this handle does not
/// implement `Sync` or `Send` as the runtime can only receive signals about its
/// nodes from its own thread. Note that when this handle is dropped, the
/// runtime will cleanup all metadata on the node.
#[derive(PartialEq, Eq, Hash)]
#[repr(transparent)]
pub struct Node {
    /// The unique numerical id of this node. Mostly useful for looking up
    /// node info from tables and things. Also useful for evil hacks.
    pub id: usize,
    phantom: PhantomData<*mut ()>,
}

impl Node {
    /// Produces a new, uniquely identified node to be managed by this thread's
    /// runtime. Generally the runtime will not allocate table entries for this
    /// node until it has to.
    pub fn next() -> Self {
        Node {
            id: NEXT_ID.fetch_add(1, Ordering::Relaxed),
            phantom: PhantomData,
        }
    }

    /// Clears this node's dependency list and rebuilds it by watching
    /// accesses made by the given function. The `ComputationHandle` can be used
    /// to more precisely control which accesses are included.
    ///
    /// ```rust
    /// let a = reax::Node::next();
    /// let b = reax::Node::next();
    /// a.computation(|_| b.on_read());
    /// assert!(a.depends_on(&b));
    /// ```
    pub fn computation<T>(&self, func: impl FnOnce(&ComputationHandle) -> T) -> T {
        // Prepare to register the dependencies of this node.
        let handle = ThreadRuntime::with(|sys| {
            if let Some(var) = sys.nodes.get_mut(&self.id) {
                // Use the DFS stack as a temporary buffer to store our
                // dependencies before clearing them. If on_write ever runs at
                // the same time and uses the DFS stack too, `with` will panic.
                sys.id_buf.clear();
                sys.id_buf.extend(var.uses.drain());
                // Remove us as a downstream node of our dependencies.
                for uses in sys.id_buf.drain(..) {
                    if let Some(uses) = sys.nodes.get_mut(&uses) {
                        uses.used_by.remove(self.id);
                    }
                }
            }

            // Setup the computation handle for the new computation.
            let handle = ComputationHandle(sys.computing_stack.len());
            sys.computing_stack.push((self.id, true));
            handle
        });
        // Do the computation.
        func(&handle)
    }

    /// Marks the given node as a direct dependency of this node.
    pub fn add_dependency(&self, on: &Node) {
        ThreadRuntime::with(|sys| sys.set_dependent_impl(self.id, on.id));
    }

    /// Checks if the given node is marked as a direct dependency of this
    /// node.
    pub fn depends_on(&self, on: &Node) -> bool {
        ThreadRuntime::with(|sys| sys.nodes.get(&self.id)
        .map(|var| var.uses.contains(on.id))
        .unwrap_or(false))
    }

    /// Returns the number of direct dependencies that this node uses.
    pub fn num_dependencies(&self) -> usize {
        ThreadRuntime::with(|sys| sys.nodes.get(&self.id)
            .map(|var| var.uses.len())
            .unwrap_or(0))
    }

    /// Returns the number of direct dependents that use this node.
    pub fn num_dependents(&self) -> usize {
        ThreadRuntime::with(|sys| sys.nodes.get(&self.id)
            .map(|var| var.used_by.len())
            .unwrap_or(0))
    }

    /// Returns true if a direct or indirect dependency of this node has been
    /// changed or if the node is not yet registered. Usually this can be
    /// thought of as telling whether or not the node is "outdated". This is
    /// intended to be very cheap to call. Reset by calling
    /// [`on_write`](#method.on_write).
    ///
    /// Note that a [`reax::Var`](struct.Var.html) will always appear as dirty
    /// even though it has no dependencies to make it outdated.
    pub fn is_dirty(&self) -> bool {
        ThreadRuntime::with(|sys| sys.nodes.get(&self.id)
            .map(|var| var.dirty)
            .unwrap_or(true))
    }

    /// If this node is clean (or if `keep_dirty` is given), calling this method
    /// marks all clean downstream nodes as dirty. Otherwise, it marks this node
    /// as clean and leaves all other nodes unchanged. Any nodes transitioned
    /// from clean to dirty by this call will send notifications to their
    /// respective channels if present. This method is idempotent and generally
    /// fast to call repeatedly.
    ///
    /// The `keep_dirty` flag is kinda subtle and mostly used for types like
    /// [`reax::Var`](struct.Var.html) which are supposed to remain dirty
    /// through writes but cause the dirty status to propagate to downstream
    /// nodes. In particular, this function will hang if `keep_dirty` is used
    /// for a node which depends on itself since the dirty flag will not be used
    /// to break the reference loop.
    ///
    /// Lazily computed nodes should *not* pass `keep_dirty` since writes to
    /// them are assumed to bring them up-to-date (that is, clean them).
    pub fn on_write(&self, keep_dirty: bool) {
        ThreadRuntime::with(|sys| {
            // As a safety measure, clear the old DFS stack.
            sys.id_buf.clear();

            // Push the cause onto the stack to begin the DFS.
            sys.id_buf.push(self.id);

            // Make sure that the cause has an entry in our nodes table if
            // its dirty status needs to be changed.
            if !keep_dirty {
               sys.nodes.entry(self.id).or_insert_with(NodeInfo::new);
            }

            // Run a DFS through the dependency graph with propagate_dirty doing
            // the heavy lifting.
            while let Some(dep_id) = sys.id_buf.pop() {
                if let Some(var) = sys.nodes.get_mut(&dep_id) {
                    var.propagate_dirty(
                        dep_id, // downstream var id
                        self.id, // cause id
                        keep_dirty, // should the cause's dirty status be unchanged?
                        &mut sys.id_buf, // DFS stack
                    );
                }
            }
        });
    }

    /// When a computation is in progress, this method generally causes self to
    /// be marked as a dependency of the actively re-computing node. See
    /// [`computation`](#method.computation).
    pub fn on_read(&self) {
        ThreadRuntime::with(|sys| {
            if let Some(&(used_by, true)) = sys.computing_stack.last() {
                sys.set_dependent_impl(used_by, self.id);
            }
        });
    }

    /// Sets the object which will handle this node transitioning from clean to
    /// dirty. This is useful for quickly updating "outdated" nodes and for
    /// firing watchers. Users should be aware that dependencies of a dirtied
    /// node are not yet updated (and may panic if accessed) at the time that
    /// the handler fires. Note that a node can only be connected to one handler
    /// at a time so usage of this may interfere with utilities like
    /// [`EagerCompute::install`](struct.EagerCompute.html#method.install).
    pub fn send_dirtied_signal_to(&self, handler: SharedDirtiedHandler) {
        ThreadRuntime::with(|sys| {
            sys.nodes.entry(self.id)
                .or_insert_with(NodeInfo::new)
                .on_dirty = Some(handler);
        });
    }

    /// Sets the name of this node used by the runtime for debug output.
    /// This is a no-op in release mode.
    #[cfg(debug_assertions)]
    pub fn set_label(&self, text: impl fmt::Display) {
        use fmt::Write;

        ThreadRuntime::with(|sys| {
            let label = &mut sys.nodes
                .entry(self.id)
                .or_insert_with(NodeInfo::new)
                .label;
            write!(label, "{}", text).unwrap();
        })
    }

    /// Sets the name of this node used by the runtime for debug output.
    /// This is a no-op in release mode.
    #[cfg(not(debug_assertions))]
    pub fn set_label(&self, _text: impl fmt::Display) { }

    /// Creates a node handle from a raw numerical id. Although this can't
    /// be used to cause undefined behavior, it is a huge footgun. In
    /// particular:
    /// - Dropping the returned node will clear all relevant runtime
    ///   metadata available on the current thread.
    /// - Runtime metadata including dependencies, dependents, labels, and dirty
    ///   status will not carry over thread boundaries. If you want to share a
    ///   node across threads you should find a way to sync it.
    pub fn from_id(id: usize) -> mem::ManuallyDrop<Node> {
        mem::ManuallyDrop::new(Node {
            id,
            phantom: PhantomData,
        })
    }

    /// Creates a copy of this node handle with all the same gotchas as
    /// [`from_raw`](#method.from_raw).
    pub fn clone_id(&self) -> mem::ManuallyDrop<Node> {
        Node::from_id(self.id)
    }
}

#[cfg(debug_assertions)]
impl fmt::Debug for Node {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        ThreadRuntime::with(|sys| fmt_label(
            sys.nodes.get(&self.id),
            self.id,
            f,
        ))
    }
}

#[cfg(not(debug_assertions))]
impl fmt::Debug for Node {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt_label(None, self.id, f)
    }
}

impl Drop for Node {
    fn drop(&mut self) {
        let id = self.id;
        ThreadRuntime::with(|sys| {
            if let Some(var) = sys.nodes.remove(&id) {
                for uses in var.uses.iter() {
                    if let Some(uses) = sys.nodes.get_mut(&uses) {
                        uses.used_by.remove(id);
                    }
                }
                for used_by in var.used_by.iter() {
                    if let Some(used_by) = sys.nodes.get_mut(&used_by) {
                        used_by.uses.remove(id);
                    }
                }
            }
        });
    }
}

#[no_mangle]
extern "C" fn reax_node_create() -> usize {
    mem::ManuallyDrop::new(Node::next()).id
}

#[no_mangle]
extern "C" fn reax_node_destroy(node: usize) {
    mem::ManuallyDrop::into_inner(Node::from_id(node));
}

#[no_mangle]
extern "C" fn reax_node_is_dirty(node: usize) -> bool {
    Node::from_id(node).is_dirty()
}

#[no_mangle]
extern "C" fn reax_node_num_dependents(node: usize) -> usize {
    Node::from_id(node).num_dependents()
}

#[no_mangle]
extern "C" fn reax_node_num_dependencies(node: usize) -> usize {
    Node::from_id(node).num_dependencies()
}

#[no_mangle]
extern "C" fn reax_node_add_dependency(node: usize, on: usize) {
    Node::from_id(node).add_dependency(&*Node::from_id(on))
}

#[no_mangle]
extern "C" fn reax_node_on_write(node: usize, keep_dirty: bool) {
    Node::from_id(node).on_write(keep_dirty)
}

#[no_mangle]
extern "C" fn reax_node_on_read(node: usize) {
    Node::from_id(node).on_read()
}

#[no_mangle]
unsafe extern "C" fn reax_write_graphviz(ptr: *mut u8, capacity: usize) -> usize {
    let buf = std::slice::from_raw_parts_mut(ptr, capacity);
    let mut writer = io::Cursor::new(buf);
    match ThreadRuntime::write_graphviz(&mut writer) {
        Ok(_) => writer.position() as usize,
        Err(_) => 0,
    }
}

#[no_mangle]
extern "C" fn reax_computation_push(node: usize) {
    ThreadRuntime::with(|rt| {
        rt.computing_stack.push((node, true));
    })
}

#[no_mangle]
extern "C" fn reax_computation_pop() {
    ThreadRuntime::with(|rt| {
        rt.computing_stack.pop();
    })
}

#[no_mangle]
extern "C" fn reax_computation_set_reactive(reactive: bool) {
    ThreadRuntime::with(|rt| {
        if let Some((_, ref mut r)) = rt.computing_stack.last_mut() {
            *r = reactive;
        }
    })
}

#[no_mangle]
extern "C" fn reax_computation_is_reactive() -> bool {
    ThreadRuntime::with(|rt| {
        rt.computing_stack.last().map(|&(_, r)| r).unwrap_or(false)
    })
}