dijkstra template in rust

dijkstra.rs

use std::ops::{Index, IndexMut, Add};
use std::cmp::Ord;

#[derive(Clone, Copy)]
pub enum NodeStatus {
    Unknown,
    Visiting(usize),
    Visited,
}

impl Default for NodeStatus {
    fn default() -> NodeStatus {
        NodeStatus::Unknown
    }
}

/// Dijkstra algorithm.
/// 
/// Node: type for graph node ID.
/// Edge: type for graph edge, typically containing edge weight.
/// Path: type for a path in the graph, typically containing the path weight.
///  - implements Ord, and must have total order.
///  - implements &Path + Edge -> Path, and must be non-decreasing on Path.
/// Board: an efficitient container mapping Node to NodeStatus.
///  - implements Board[&Node] -> &mut NodeStatus, and must be deterministic.
/// start_node: the starting node for all paths.
/// start_path: the empty path from the starting node to itself.
/// board: the initial Board. Must contain NodeStatus::default() for all Nodes.
/// get_neighbor: given a Node, returns an iterator of all its neighbor Edges 
///               and Nodes.
/// report: receives a result of the shortest Path to a Node, 
///         and returns whether to continue the algorithm.
pub fn dijkstra<
    Node, Edge, Path, Board, GetNeighbor, NeighborIter, ReportResult
> (
    start_node: Node, 
    start_path: Path, 
    mut board: Board, 
    mut get_neighbor: GetNeighbor, 
    mut report: ReportResult, 
)
where 
    // ideally Board should be higher rank type over IndexMut::Output
    // Board: for<'a, N: Default + Clone + Copy> IndexMut<&'a Node, Output = N>,
    Board: for<'a> IndexMut<&'a Node, Output = NodeStatus>,
    Path: Ord,
    for<'a> &'a Path: Add<Edge, Output = Path>,
    GetNeighbor: FnMut(Node) -> NeighborIter,
    NeighborIter: IntoIterator<Item = (Edge, Node)>,
    ReportResult: FnMut(&Node, &Path) -> bool,
{
    let mut heap: Vec<(Path, Node)> = vec![];
    
    board[&start_node] = NodeStatus::Visiting(0);
    heap.push((start_path, start_node));
    
    let swap_heap_element = |
        heap: &mut Vec<(Path, Node)>, 
        board: &mut Board, 
        i: usize, 
        j: usize
    | {
        if i == j {
            return;
        }
        board[&heap[i].1] = NodeStatus::Visiting(j);
        board[&heap[j].1] = NodeStatus::Visiting(i);
        heap.swap(i, j);
    };
    
    let pop_heap = |
        heap: &mut Vec<(Path, Node)>, 
        board: &mut Board
    | -> (Path, Node) {
        swap_heap_element(heap, board, 0, heap.len() - 1);
        let popped = heap.pop().unwrap();
        if !heap.is_empty() {
            let mut cur = 0;
            loop {
                let left = heap.get(cur * 2 + 1).map(|n|&n.0);
                let right = heap.get(cur * 2 + 2).map(|n|&n.0);
                if (left.is_none() || &heap[cur].0 <= left.unwrap()) 
                    && (right.is_none() || &heap[cur].0 <= right.unwrap()) {
                    break;
                }
                let smaller = 
                    if right.is_none() || left.unwrap() < right.unwrap() {
                        cur * 2 + 1
                    } else {
                        cur * 2 + 2
                    };
                swap_heap_element(heap, board, cur, smaller);
                cur = smaller;
            }
        }
        board[&popped.1] = NodeStatus::Visited;
        popped
    };
    
    let new_path = |
        heap: &mut Vec<(Path, Node)>, 
        board: &mut Board, 
        path: Path, 
        node: Node
    | {
        let mut cur;
        match board[&node] {
            NodeStatus::Unknown => {
                cur = heap.len();
                board[&node] = NodeStatus::Visiting(cur);
                heap.push((path, node));
            }
            NodeStatus::Visiting(i) => {
                if heap[i].0 <= path {
                    return;
                }
                cur = i;
                heap[i].0 = path;
            }
            NodeStatus::Visited => return
        }
        while cur != 0 {
            let parent = (cur - 1) / 2;
            if heap[parent].0 <= heap[cur].0 {
                break;
            }
            swap_heap_element(heap, board, parent, cur);
            cur = parent;
        }
    };
    
    while !heap.is_empty() {
        let (path, node) = pop_heap(&mut heap, &mut board);
        if !report(&node, &path) {
            return;
        }
        
        for (edge, new_node) in get_neighbor(node) {
            new_path(&mut heap, &mut board, &path + edge, new_node);
        }
    }
}

/////////////////////////////////////////////////////////////////////////////
// Leetcode 407

enum Node {
    Outside,
    Grid(usize, usize)
}

struct Board {
    width: usize,
    grid: Vec<NodeStatus>,
    outside: NodeStatus,
}

impl Board {
    fn new(width: usize, height: usize) -> Board {
        Board {
            width,
            grid: vec![NodeStatus::default(); width * height],
            outside: NodeStatus::default()
        }
    } 
}

#[derive(Copy, Clone, Default, PartialOrd, Ord, PartialEq, Eq)]
struct Path(i32);

impl<'a> Add<i32> for &'a Path {
    type Output = Path;
    fn add(self, other: i32) -> Path {
        Path(std::cmp::max(self.0, other))
    }
}

impl<'a> Index<&'a Node> for Board {
    type Output = NodeStatus;

    fn index(&self, index: &'a Node) -> &Self::Output {
        match index {
            Node::Outside => &self.outside,
            Node::Grid(x, y) => &self.grid[*x + *y * self.width]
        }
    }
}

impl<'a> IndexMut<&'a Node> for Board {
    fn index_mut(&mut self, index: &'a Node) -> &mut Self::Output {
        match index {
            Node::Outside => &mut self.outside,
            Node::Grid(x, y) => &mut self.grid[*x + *y * self.width]
        }
    }
}

impl Solution {
pub fn trap_rain_water(height_map: Vec<Vec<i32>>) -> i32 {
        let width = height_map.len();
        let height = height_map[0].len();
        let mut result = 0;
        let get_neighbor = |node: Node| {
            let mut list = vec![];
            match node {
                Node::Outside => {
                    for x in 0..width {
                        list.push((height_map[x][0], Node::Grid(x, 0)));
                        if height != 1 {
                            list.push((height_map[x][height - 1], 
                                Node::Grid(x, height - 1)));
                        }
                    }
                    for y in 1..height - 1 {
                        list.push((height_map[0][y], Node::Grid(0, y)));
                        if width != 1 {
                            list.push((height_map[width - 1][y], 
                                Node::Grid(width - 1, y)));
                        }
                    }
                }
                Node::Grid(x, y) => {
                    let mut push = |nx: usize, ny: usize| {
                        list.push((height_map[nx][ny], Node::Grid(nx, ny)));
                    };
                    if x > 0 {
                        push(x - 1, y);
                    }
                    if y > 0 {
                        push(x, y - 1);
                    }
                    if x < width - 1 {
                        push(x + 1, y);
                    }
                    if y < height - 1 {
                        push(x, y + 1);
                    }
                }
            }
            list
        };
        let report = |node: &Node, path: &Path| -> bool {
            if let Node::Grid(x, y) = node {
                result += path.0 - height_map[*x][*y];
            }
            true
        };
        dijkstra(Node::Outside, Path::default(),
            Board::new(width, height), get_neighbor, report
        );
        result
    }
}  

struct Solution;