Simple sparse graph structure and djikstra's algorithm in rust (rustc 0.11.0-pre (e55f64f 2014-06-09 01:11:58 -0700))

graph.rs

use collections::slice;
use std::cell::Cell;

#[deriving(Show, PartialEq, Eq, PartialOrd, Ord)]
pub struct NodeId(pub int);
#[deriving(Show, PartialEq, Eq)]
pub struct EdgeId(int);

#[deriving(Show)]
pub struct Node<N, E> {
    pub id: NodeId,
    pub val: Cell<N>,
    pub visited: Cell<bool>,
    edges_: Vec<EdgeId>,
}

#[deriving(Show)]
pub struct Edge<N, E> {
    pub id: EdgeId,
    pub val: Cell<E>,
    left: NodeId,
    right: NodeId,
}

impl<N, E> Edge<N, E> {
    fn other(&self, id: NodeId) -> NodeId {
        if self.left == id {
            self.right
        } else if self.right == id {
            self.left
        } else {
            fail!("Node id not found in edge");
        }
    }
}

pub struct EdgeIterator<'a, N, E> {
    graph: &'a Graph<N,E>,
    id: NodeId,
    iter: slice::Items<'a, EdgeId>,
}

impl<'a, N: Copy, E: Copy> Iterator<(&'a Edge<N,E>, &'a Node<N,E>)> for EdgeIterator<'a, N, E> {
    fn next(&mut self) -> Option<(&'a Edge<N, E>, &'a Node<N, E>)> {
        self.iter.next().map(|next_id| {
            let edge = self.graph.get_edge(*next_id);
            let node = self.graph.get_node(edge.other(self.id));
            (edge, node)
        })
    }
}

#[deriving(Show)]
pub struct Graph<N, E> {
    edges: Vec<Edge<N, E>>,
    nodes: Vec<Node<N, E>>,
}

impl<N: Copy, E: Copy> Graph<N, E> {
    pub fn new() -> Graph<N, E> {
        Graph {
            edges: Vec::new(),
            nodes: Vec::new()
        }
    }

    pub fn add_node(&mut self, node: N) -> NodeId {
        let next_id = NodeId(self.nodes.len() as int);
        self.nodes.push( Node {
            id: next_id,
            val: Cell::new(node),
            visited: Cell::new(false),
            edges_: vec!(),
        });
        next_id
    }

    pub fn mark_unvisited(&mut self) {
        for node in self.nodes.iter() {
            node.visited.set(false);
        }
    }

    pub fn edges<'a> (&'a self, id: NodeId) -> EdgeIterator<'a, N, E> {
        EdgeIterator {
            graph: self,
            id: id,
            iter: self.get_node(id).edges_.iter(),
        }
    }

    pub fn get_node<'a> (&'a self, id: NodeId) -> &'a Node<N, E> {
        let NodeId(index) = id;
        self.nodes.get(index as uint)
    }

    pub fn get_edge<'a> (&'a self, id: EdgeId) -> &'a Edge<N, E> {
        let EdgeId(index) = id;
        self.edges.get(index as uint)
    }

    pub fn get_mut_node<'a> (&'a mut self, id: NodeId) -> &'a mut Node<N, E> {
        let NodeId(index) = id;
        &mut self.nodes.as_mut_slice()[index as uint]
    }

    pub fn add_edge(&mut self, edge: E, left: NodeId, right: NodeId) -> EdgeId {
        let next_id = EdgeId(self.edges.len() as int);
        self.edges.push( Edge {
            id: next_id,
            val: Cell::new(edge),
            left: left,
            right: right,
        } );
        self.get_mut_node(left).edges_.push(next_id);
        self.get_mut_node(right).edges_.push(next_id);
        next_id
    }
}

routing.rs

extern crate collections;

use std::os;
use graph::{NodeId, Graph};
use collections::priority_queue;

mod graph;

#[deriving(Show)]
struct DijkstraNode {
    cost: int,
    previous: NodeId,
}

type DijkstraGraph = Graph<Option<DijkstraNode>, int>;

fn read_graph (from: &mut Reader) -> std::io::IoResult<(DijkstraGraph, NodeId, NodeId)> {
    let string = try!(from.read_to_str());
    let mut g = Graph::new();
    //println!("{}", string);
    let mut in_iter = string.as_slice().
        split('\n').flat_map(|x| x.split(',') ).flat_map(|x| x.split(' ') );
    let matrix_size = from_str::<int>(in_iter.next().unwrap()).unwrap();
    for _ in range(0, matrix_size) {
        g.add_node(None);
    }
    for y in range(0, matrix_size) {
        for x in range(0, matrix_size) {
            let cost = from_str::<int>(in_iter.next().unwrap()).unwrap();
            if cost != -1 {
                g.add_edge(cost, NodeId(x), NodeId(y));
            }
        }
    }
    fn alpha_index(string: Option<&str>) -> NodeId {
        NodeId(string.map(|x| x.char_at(0) as int).unwrap() - 'A' as int)
    }
    let start = alpha_index(in_iter.next());
    let end = alpha_index(in_iter.next());
    Ok((g, start, end))
}

fn dijkstra(graph: &mut DijkstraGraph, start: NodeId, end: NodeId) -> Vec<NodeId> {
    let mut unfinished_nodes = priority_queue::PriorityQueue::from_vec(vec!((0, start)));
    graph.mark_unvisited();
    loop {
        let (current_cost, current_node) = match unfinished_nodes.pop() {
            Some((c, n)) => (-c, graph.get_node(n)), None => break
        };
        if current_node.visited.get() {
            continue;
        }
        for (edge, node) in graph.edges(current_node.id) {
            if node.visited.get() {
                continue;
            }
            let mut val = node.val.get();
            let cost;
            let new_cost = current_cost + edge.val.get();
            if val.is_none() || val.unwrap().cost > new_cost {
                val = Some(DijkstraNode{cost: new_cost, previous: current_node.id});
                node.val.set(val);
            }
            cost = val.unwrap().cost;
            unfinished_nodes.push((-cost, node.id));
            let best_cost = unfinished_nodes.top().unwrap();
            if node.id == end && -best_cost.val0() >= cost {
                break;
            }
        }
        current_node.visited.set(true);
    }
    let mut path : Vec<NodeId> = vec!();
    let mut n = graph.get_node(end);
    if n.val.get().is_none() {
        fail!("Failed to find end node");
    }
    loop {
        path.push(n.id);
        n = match n.val.get() {
            Some(n) => graph.get_node(n.previous),
            None => {
                path.reverse();
                return path;
            }
        };
    }
}

fn path_to_str(path: &[NodeId]) -> String {
    let mut path_str = String::new();
    for x in path.iter() {
        let NodeId(x) = *x;
        path_str.push_char(std::char::from_u32(x as u32 + 'A' as u32).unwrap());
    }
    path_str
}

fn main() {
    let args = os::args();
    let mut file;
    if args.len() != 2 {
        fail!(format!("usage: {} [cost matrix file]", args.get(0)));
    } else {
        file = std::io::File::open(&Path::new(args.get(1).as_slice()));
    }
    let (mut gr, start, end) = read_graph(&mut file).unwrap();
    let path = dijkstra(&mut gr, start, end);
    println!("{}", path_to_str(path.as_slice()));
}