shark/fish/src/bot/node.rs

//! Graph data structures used by `Bot` in its search algorithm.

use mino::matrix::{Mat, MatBuf};
use mino::srs::{Piece, PieceType, Queue};

use crate::bot::Arena;
use crate::find::find_locations;

/// Represents a node in the search tree. A node basically just consists of a board state
/// (incl. queue) and some extra metadata relating it to previous nodes in the tree.
pub(crate) struct Node {
    matrix: *const Mat,
    queue: RawQueue,
    edge: Option<Edge>,
    rating: i32,
    // currently there is no need to store a node's children, but maybe this could change
    // in the future.
}

// Reallocates the matrix into the arena.
fn copy_matrix<'a>(arena: &'a Arena, matrix: &Mat) -> &'a Mat {
    Mat::new(arena.alloc_slice_copy(&matrix[..]))
}

// Reallocates the queue into the arena.
fn copy_queue<'a>(arena: &'a Arena, queue: Queue<'_>) -> Queue<'a> {
    Queue {
        hold: queue.hold,
        next: arena.alloc_slice_copy(&queue.next),
    }
}

impl Node {
    /// Allocate a root node using the given arena and initial configuration. The initial
    /// matrix and queue are also allocated onto the arena, so you do not need to worry
    /// about their lifetimes when managing the lifetime of the root.
    pub fn alloc_root<'a>(arena: &'a Arena, matrix: &Mat, queue: Queue<'_>) -> &'a Self {
        let matrix = copy_matrix(arena, matrix);
        let queue = copy_queue(arena, queue);
        Node::alloc(arena, matrix, queue, i32::MIN, None)
    }

    // `matrix` and `queue` must be allocated inside `arena`
    fn alloc<'a>(
        arena: &'a Arena,
        matrix: &'a Mat,
        queue: Queue<'a>,
        rating: i32,
        edge: Option<Edge>,
    ) -> &'a Self {
        let matrix = matrix as *const Mat;
        let queue = RawQueue::from(queue);
        arena.alloc_with(|| Self {
            matrix,
            queue,
            edge,
            rating,
        })
    }

    pub fn matrix(&self) -> &Mat {
        unsafe { &*self.matrix }
    }

    pub fn queue(&self) -> Queue<'_> {
        unsafe { self.queue.as_queue() }
    }

    pub fn is_better(&self, other: &Node) -> bool {
        self.rating > other.rating
    }

    pub fn is_terminal(&self) -> bool {
        // TODO: additional terminal-node conditions e.g. clears last row of garbage
        self.queue().is_empty()
    }

    /// Get the initial placement made after the root node which eventually arrives at
    /// this node.
    pub fn root_placement(&self) -> Option<Piece> {
        let mut root_placement = None;
        let mut parent = Some(self);
        while let Some(node) = parent.take() {
            parent = node.edge.as_ref().map(|e| {
                root_placement = Some(e.placement);
                e.parent()
            });
        }
        root_placement
    }

    /// Expands this node, allocating the children into the given arena.
    // `self` must be allocated inside `arena`
    pub fn expand<'a, E>(
        &'a self,
        arena: &'a Arena,
        evaluate: E,
    ) -> impl Iterator<Item = &'a Node> + 'a
    where
        E: Fn(&Mat, Queue<'_>) -> i32 + 'a,
    {
        let placements = self.queue().reachable().flat_map(|ty| {
            let locs = find_locations(self.matrix(), ty);
            locs.map(move |loc| Piece { ty, loc })
        });

        let mut matrix = MatBuf::new();

        placements.map(move |placement| {
            matrix.copy_from(self.matrix());
            placement.cells().fill(&mut matrix);
            matrix.clear_lines();
            // TODO: the above call returns useful information about if this placement is
            // a combo, does it clear the bottom row of garbage. this should be used for
            // prioritizing nodes

            let parent = RawNodePtr::from(self);
            let edge = Edge { placement, parent };
            let suc_matrix = copy_matrix(arena, &matrix);
            let suc_queue = self.queue().remove(placement.ty);

            // TODO: transposition table

            let rating = evaluate(suc_matrix, suc_queue);

            Node::alloc(arena, suc_matrix, suc_queue, rating, Some(edge))
        })
    }
}

impl core::fmt::Debug for Node {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "Node(rating={}", self.rating)?;
        if let Some(pc) = self.root_placement() {
            write!(f, ", root_placement={}", pc.ty)?;
        }
        if !self.queue().is_empty() {
            write!(f, ", queue={}", self.queue())?;
        }
        write!(f, ")")
    }
}

/// Represents an edge in the graph, pointing from a node to its parent. Particularly,
/// contains the placement made in order to arrive at the child from the parent.
struct Edge {
    placement: Piece,
    parent: RawNodePtr,
}

impl Edge {
    fn parent(&self) -> &Node {
        unsafe { self.parent.as_node() }
    }
}

/// Wraps a raw pointer to a `Node`, requiring you to manage the lifetime yourself.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
#[repr(transparent)]
pub(crate) struct RawNodePtr(*const Node);

impl RawNodePtr {
    pub unsafe fn as_node<'a>(self) -> &'a Node {
        &*self.0
    }
}

impl From<&Node> for RawNodePtr {
    fn from(node: &Node) -> Self {
        Self(node)
    }
}

/// Wraps the raw components of a `Queue`, requiring you to manage the lifetime yourself.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
struct RawQueue {
    hold: Option<PieceType>,
    len: u16, // u16 to save space esp. considering padding
    next: *const PieceType,
}

impl RawQueue {
    pub unsafe fn as_queue<'a>(self) -> Queue<'a> {
        let hold = self.hold;
        let next = core::slice::from_raw_parts(self.next, self.len as usize);
        Queue { hold, next }
    }
}

impl From<Queue<'_>> for RawQueue {
    fn from(queue: Queue<'_>) -> Self {
        Self {
            hold: queue.hold,
            len: queue.next.len() as u16,
            next: queue.next.as_ptr(),
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use mino::mat;

    #[test]
    fn test_copy_matrix() {
        let arena = Arena::new();
        let mat0 = mat! {
            "..xxx..x.x";
            "xxxxxx.xxx";
        };
        let mat1 = copy_matrix(&arena, mat0);
        assert_eq!(mat0, mat1);
    }

    #[test]
    fn test_copy_queue() {
        use PieceType::*;
        let arena = Arena::new();
        let q0 = Queue::new(None, &[I, L, J, O]);
        let q1 = copy_queue(&arena, q0);
        assert_eq!(q0, q1);
    }

    #[test]
    fn test_sizeof_raw_queue() {
        assert_eq!(
            core::mem::size_of::<RawQueue>(),
            core::mem::size_of::<(u16, u16, *const ())>(),
        );
    }

    #[test]
    fn test_raw_queue_roundtrip() {
        use PieceType::*;
        let q0 = Queue::new(None, &[I, L, J, O]);
        let rq0 = RawQueue::from(q0);
        let q1 = unsafe { rq0.as_queue() };
        assert_eq!(q1, q0);

        let q0 = Queue::new(None, &[]);
        let rq0 = RawQueue::from(q0);
        let q1 = unsafe { rq0.as_queue() };
        assert_eq!(q1, q0);
    }
}