wmaker/wutil-rs/src/tree.rs

//! N-ary tree structure.
//!
//! ## Rust rewrite notes
//!
//! This is only used in one place (`util/wmmenugen.c`), and it should probably
//! be moved out of wutil-rs (if not deleted entirely) once we migrate that
//! utility.
//!
//! The FFI functions provided here assume that they will be used as they are in
//! `wmmenugen.c`. (For example, the C interface may allow for some callback
//! parameters to be null, but they aren't in practice, so the Rust layer
//! doesn't check for that.)
//!
//! The original C library had a larger surface area and tracked parent
//! pointers, but the Rust version only has the API used by wmmenugen. Case in
//! point: `WMTreeNode` originally had an optional destructor function pointer
//! for freeing its `data` field, but wmmenugen never actually deleted a tree
//! node, so the destructor was never called. This Rust rewrite doesn't track a
//! `data` destructor and doesn't even provide a way to delete tree nodes.
//!
//! If a generic tree really becomes necessary, [`Tree`] may be adapted, but a
//! different design is probably warranted. (At the very least, `Tree` should be
//! generic over the type of `data` so that we're not using `c_void`. It is also
//! probably better to index into an owned `Vec` or use an external
//! arena. `GhostCell` or one of its ilk may even make sense.)

use std::ffi::c_void;

pub struct Tree {
    depth: u32,
    children: Vec<Box<Tree>>,
    data: *mut c_void,
}

pub mod ffi {
    use super::Tree;

    use std::{
        ffi::{c_int, c_void},
        ptr,
    };

    /// Creates the root of a new tree with the data `data`.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMCreateTreeNode(data: *mut c_void) -> *mut Tree {
        Box::leak(Box::new(Tree {
            depth: 0,
            children: Vec::new(),
            data,
        }))
    }

    /// Creates a tree node as the `index`th child of `parent`, with the data
    /// `item`. If `index` is negative, the node is added as the last child of
    /// `parent`.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMInsertItemInTree(
        parent: *mut Tree,
        index: c_int,
        item: *mut c_void,
    ) -> *mut Tree {
        if parent.is_null() {
            return ptr::null_mut();
        }
        let parent = unsafe { &mut *parent };

        let child = Tree {
            depth: parent.depth + 1,
            children: Vec::new(),
            data: item,
        };

        if index < 0 {
            parent.children.push(Box::new(child));
            parent.children.last_mut().unwrap().as_mut() as *mut _
        } else {
            let index = index as usize;
            parent.children.insert(index, Box::new(child));
            parent.children[index].as_mut() as *mut _
        }
    }

    /// Inserts `tree` as the `index`th child of `parent`. If `index` is
    /// negative, `tree` is added as the last child of `parent`.
    ///
    /// This eagerly updates the depth of the entire subtree rooted at `tree`,
    /// so it has a O(N) cost rather than O(1).
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMInsertNodeInTree(
        parent: *mut Tree,
        index: c_int,
        tree: *mut Tree,
    ) -> *mut Tree {
        if parent.is_null() {
            return ptr::null_mut();
        }
        let parent = unsafe { &mut *parent };

        if tree.is_null() {
            return ptr::null_mut();
        }

        let mut stack = vec![(unsafe { &mut *tree }, parent.depth + 1)];
        while let Some((tree, depth)) = stack.pop() {
            tree.depth = depth;
            for child in &mut tree.children {
                stack.push((child, depth + 1));
            }
        }

        if index < 0 {
            parent.children.push(unsafe { Box::from_raw(tree) });
            tree
        } else {
            let index = index as usize;
            parent
                .children
                .insert(index, unsafe { Box::from_raw(tree) });
            tree
        }
    }

    /// Returns the data field of `tree`.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMGetDataForTreeNode(tree: *mut Tree) -> *mut c_void {
        if tree.is_null() {
            return ptr::null_mut();
        }
        unsafe { (*tree).data }
    }

    /// Returns the depth of `tree` (0 if `tree` is a root, 1 if it is an
    /// immediate child of the root, etc.). This is an `O(1)` operation.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMGetTreeNodeDepth(tree: *mut Tree) -> c_int {
        if tree.is_null() {
            return 0;
        }
        unsafe { (*tree).depth as c_int }
    }

    /// Recursively sorts the children of each node of the subtree rooted at
    /// `tree`, according to `comparer`.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMSortTree(
        tree: *mut Tree,
        comparer: unsafe extern "C" fn(*const Tree, *const Tree) -> c_int,
    ) {
        use std::cmp::Ordering;

        if tree.is_null() {
            return;
        }

        let comparer = |a: &Box<Tree>, b: &Box<Tree>| {
            let a = a.as_ref() as *const Tree as *const _;
            let b = b.as_ref() as *const Tree as *const _;
            match unsafe { comparer(a, b) }.signum() {
                -1 => Ordering::Less,
                0 => Ordering::Equal,
                1 => Ordering::Greater,
                _ => unreachable!(),
            }
        };

        let mut stack = vec![unsafe { &mut *tree }];

        while let Some(tree) = stack.pop() {
            tree.children.sort_by(comparer);
            stack.extend(tree.children.iter_mut().map(|c| c.as_mut()));
        }
    }

    /// Returns the first tree node in the subtree rooted at `tree` that matches
    /// `match_p`, or null if none is found.  If `match_p` is `None`, returns
    /// the first node whose `data` field is equal to `cdata`. Search is
    /// recursive, up to `limit` depth (which must be non-negative).
    ///
    /// ## Rust rewrite notes
    ///
    /// This was originally a DFS with a `depthFirst` parameter which actually
    /// controlled whether tree traversal was pre- or post-order (and not
    /// whether the search was depth-first or breadth-first, which the name
    /// might seem to suggest). Since this was only ever called with a non-zero
    /// `depthFirst`, the Rust version is hard-coded as a pre-order DFS.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMFindInTreeWithDepthLimit(
        tree: *mut Tree,
        match_p: Option<unsafe extern "C" fn(*const Tree, *const c_void) -> c_int>,
        cdata: *mut c_void,
        limit: c_int,
    ) -> *mut Tree {
        if tree.is_null() {
            return ptr::null_mut();
        }
        let safe_tree = unsafe { &mut *tree };

        let match_p = |t: &Tree| -> bool {
            if let Some(p) = match_p {
                (unsafe { (p)(t, cdata) }) != 0
            } else {
                t.data == cdata
            }
        };

        let mut stack = vec![(safe_tree, 0)];

        while let Some((tree, depth)) = stack.pop() {
            if depth > limit {
                continue;
            }

            if match_p(tree) {
                return tree as *mut Tree;
            }

            for child in tree.children.iter_mut() {
                stack.push((child.as_mut(), depth + 1));
            }
        }

        ptr::null_mut()
    }

    /// Traverses `tree` recursively, invoking `walk` on each tree node with
    /// `cdata` passed in to provide a side channel.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMTreeWalk(
        tree: *mut Tree,
        walk: unsafe extern "C" fn(*mut Tree, *mut c_void),
        cdata: *mut c_void,
    ) {
        if tree.is_null() {
            return;
        }
        let tree = unsafe { &mut *tree };

        let walk_fn = |t: &mut Tree| unsafe {
            walk(t as *mut Tree, cdata);
        };

        let mut stack = vec![tree];

        while let Some(tree) = stack.pop() {
            walk_fn(tree);

            for child in &mut tree.children {
                stack.push(child);
            }
        }
    }
}