shark/fish/src/bot/node.rs

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

use core::ops::Range;

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

use crate::bot::trans::TransTable;
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 struct Node {
    matrix: *const Mat,
    queue: RawQueue,
    edge: Option<Edge>,
    rating: (bool, 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);
        let rating = (false, i32::MIN);
        Node::alloc(arena, matrix, queue, rating, None)
    }

    // `matrix` and `queue` must be allocated inside `arena`
    fn alloc<'a>(
        arena: &'a Arena,
        matrix: &'a Mat,
        queue: Queue<'a>,
        rating: (bool, 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 rating(&self) -> (bool, i32) {
        self.rating
    }

    /// 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,
        trans: &'a mut TransTable,
        mut evaluate: E,
    ) -> impl Iterator<Item = &'a Node> + 'a
    where
        E: FnMut(&Mat, Queue<'_>, &Range<i16>) -> (bool, i32) + 'a,
    {
        let mut matrix = MatBuf::new();

        let placements = self.queue().reachable().flat_map(|ty| {
            let locs = find_locations(self.matrix(), ty);
            locs.map(move |loc| Piece { ty, loc })
        });

        let children = placements.map(move |placement| {
            // compute new board state from placement
            matrix.copy_from(self.matrix());
            placement.cells().fill(&mut matrix);
            let cleared = matrix.clear_lines();
            let queue = self.queue().remove(placement.ty);

            // allocate new node unless this state already exists
            trans.get_or_insert(&matrix, queue, || {
                Node::alloc(
                    arena,
                    copy_matrix(arena, &matrix),
                    // `queue` is already allocated on the arena so don't need to copy it
                    queue,
                    evaluate(&matrix, queue, &cleared),
                    Some(Edge {
                        placement,
                        parent: self.into(),
                    }),
                )
            })
        });

        // exclude duplicate children to avoid redundant search sub-trees
        children.filter_map(|child| {
            match child {
                Ok(new) => Some(new),
                Err(dup) => {
                    // TODO: reparent dup so that it has a more preferable initial path.
                    tracing::trace!("ignoring duplicate child {dup:?}");
                    None
                }
            }
        })
    }
}

impl core::fmt::Debug for Node {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "Node {{ ")?;
        match self.rating {
            (false, h) => write!(f, "rating: {}", h),
            (true, n) => write!(f, "solution: {}", -n),
        }?;
        if let Some(pc) = self.root_placement() {
            write!(f, ", root_placement: {:?}", pc)?;
        }
        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 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);
    }
}