<commit_before>use std::cell::Ref;
use std::collections::HashMap;
use std::rc::Rc;

use ga::{self, Array, Tensor, TensorMode};
use rand;

use super::init::Initializer;
use super::op::{OpBuilder, OpDescriptor, Operation};
use super::var_store::{VarIndex, VarStore};

#[derive(Copy, Clone)]
pub enum NodeInput {
    Var(VarIndex),    // Regular input variable
    Recurrent(usize), // Recurrent connection
}

pub struct Node {
    pub inputs: Vec<VarIndex>,
    pub outputs: Vec<VarIndex>,
    pub in_grad: Vec<VarIndex>, // gradients on inputs
    pub out_grad: Vec<OutGrad>, // gradients on outputs
}

pub struct Graph {
    ctx: Rc<ga::Context>,
    nodes: Vec<Node>,
    node_ops: Vec<Box<Operation>>,
    pub var_store: VarStore,
    out_var_map: HashMap<VarIndex, (NodeIndex, usize)>, // Maps output variable to its node and index within node
    // Gradients on variables that are inputs to the graph - they have no corresponding node
    in_var_map: HashMap<VarIndex, usize>,
    learnables: Vec<(VarIndex, GradIndex)>, // Learnable variables
    in_grad: Vec<OutGrad>, // Gradients on variables that are inputs to the graph

    rng: rand::ThreadRng,
}

impl Graph {
    pub fn new(ctx: Rc<ga::Context>) -> Self {
        Graph {
            ctx: ctx,
            nodes: vec![],
            node_ops: vec![],
            var_store: VarStore::new(),
            out_var_map: HashMap::new(),
            in_var_map: HashMap::new(),
            learnables: vec![],
            in_grad: vec![],
            rng: rand::thread_rng(),
        }
    }

    pub fn add_node<T: OpBuilder>(&mut self, op: T) -> NodeIndex {
        let node_index = NodeIndex(self.nodes.len());

        let OpDescriptor { op, inputs: node_inputs, out_shapes } = op.build(&self.ctx, &self.var_store).unwrap();

        // Create output variables
        let mut outputs = vec![];
        for (i, shape) in out_shapes.into_iter().enumerate() {
            let var_index = self.var_store.add(Tensor::new(self.ctx.as_ref(), shape, TensorMode::Mut));
            outputs.push(var_index);
            self.out_var_map.insert(var_index, (node_index, i));
        }
        let mut out_grad = vec![OutGrad::new(); outputs.len()];
        // Set up inputs and gradients on inputs
        let mut inputs = vec![];
        let mut in_grad = vec![];
        for input in node_inputs {
            let (v, gradient) =
                match input {
                    NodeInput::Var(v) => (v, self.add_gradient(v)),
                    NodeInput::Recurrent(out) => {
                        let v = outputs[out];
                        let gradient = Self::create_gradient(&self.ctx, &mut self.var_store, v,
                                                             &mut out_grad[out]);
                        (v, gradient)
                    }
                };
            inputs.push(v);
            in_grad.push(gradient);
        }
        // Create the node
        self.nodes.push(Node { inputs: inputs,
                               outputs: outputs,
                               in_grad: in_grad,
                               out_grad: out_grad });
        // Add the corresponding node op
        self.node_ops.push(Box::new(op));
        node_index
    }

    pub fn add_variable<I: Initializer>(&mut self,
                                        shape: Vec<usize>,
                                        learnable: bool,
                                        init: I) -> VarIndex {
        let a = init.init(&mut self.rng, shape);
        let v = self.var_store.add(Tensor::from_array(&self.ctx, &a, TensorMode::Mut));
        self.in_var_map.insert(v, self.in_grad.len());
        if learnable {
            self.learnables.push((v, GradIndex::InVar(self.in_grad.len())));
        }
        self.in_grad.push(OutGrad::new());
        v
    }

    pub fn get_gradient(&self, v: VarIndex) -> GradIndex {
        match self.out_var_map.get(&v).map(|x| *x) {
            Some((node, out_index)) => {
                // v is the output of a node
                GradIndex::OutVar(node, out_index)
            },
            None => {
                // v is an input to the graph - it has no corresponding node
                let in_grad_index = *self.in_var_map.get(&v)
                                         .expect("Variable is neither input nor output. Nonsense!");
                GradIndex::InVar(in_grad_index)
            },
        }
    }

    pub fn add_gradient(&mut self, v: VarIndex) -> VarIndex {
        match self.out_var_map.get(&v).map(|x| *x) {
            Some((node, out_index)) => {
                // v is the output of a node
                /*self.nodes[node.0].out_grad[out_index]
                                  .fork(&self.ctx, &mut self.var_store, gradient);*/
                Self::create_gradient(&self.ctx, &mut self.var_store, v,
                                      &mut self.nodes[node.0].out_grad[out_index])
            },
            None => {
                // v is an input to the graph - it has no corresponding node
                let in_grad_index = *self.in_var_map.get(&v)
                                         .expect("Variable is neither input nor output. Nonsense!");
                /*self.in_grad[in_grad_index]
                    .fork(&self.ctx, &mut self.var_store, gradient);*/
                Self::create_gradient(&self.ctx, &mut self.var_store, v,
                                      &mut self.in_grad[in_grad_index])
            },
        }
    }

    fn create_gradient(ctx: &ga::Context,
                       var_store: &mut VarStore,
                       v: VarIndex,
                       out_grad: &mut OutGrad)
                       -> VarIndex {
        let shape = var_store.get(v).shape().to_owned();
        let gradient = var_store.add(Tensor::new(ctx, shape, TensorMode::Mut));
        out_grad.fork(ctx, var_store, gradient);
        gradient
    }

    pub fn forward(&mut self) {
        // Forward pass
        //
        // NOTE: We just execute the nodes in order. We can do this because of the way the graph is
        // built. When a user wants to add a node, he/she must also supply the inputs. This means
        // any dependencies must already be added before the node can be added. Therefore, we can
        // assert that all dependents come after their dependencies in the `self.nodes` array.
        for (node, op) in self.nodes.iter().zip(&mut self.node_ops) {
            op.forward(&self.ctx, &self.var_store, node);
        }
    }

    pub fn backward(&mut self) {
        // Backward pass
        // We run through the nodes in reverse order. See note in Graph::forward
        for (node, op) in self.nodes.iter_mut().rev().zip(self.node_ops.iter_mut().rev()) {
            // Sum the gradients on each output if there are multiple gradients
            for out_grad in &node.out_grad {
                out_grad.maybe_sum(self.ctx.as_ref(), &mut self.var_store);
            }
            op.backward(&self.ctx, &mut self.var_store, node);
        }
    }

    pub fn context(&self) -> &ga::Context {
        &self.ctx
    }

    pub fn learnables(&self) -> &[(VarIndex, GradIndex)] {
        &self.learnables
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

#[derive(Clone)]
pub struct OutGrad {
    gradient: Option<VarIndex>, // The gradient or sum of gradients
    gradients: Vec<VarIndex>,
}

impl OutGrad {
    pub fn new() -> Self {
        OutGrad {
            gradient: None,
            gradients: vec![],
        }
    }

    pub fn get(&self) -> VarIndex {
        self.gradient.unwrap()
    }

    pub fn try_get(&self) -> Option<VarIndex> {
        self.gradient
    }

    fn maybe_sum(&self, ctx: &ga::Context, var_store: &mut VarStore) {
        if self.gradients.len() > 0 {
            if let Some(sum) = self.gradient {
                ga::copy_to(ctx, &var_store.get(self.gradients[0]), &var_store.get_mut(sum));
                for grad in &self.gradients[1..] {
                    ga::add(ctx, &var_store.get(sum), -1, &var_store.get(*grad), &var_store.get_mut(sum));
                }
            }
        }
    }

    fn fork(&mut self, ctx: &ga::Context, var_store: &mut VarStore, v: VarIndex) {
        if self.gradients.len() > 0 {
            // There are multiple gradients already, just add the new one to the list
            self.gradients.push(v);
        } else if let Some(gradient) = self.gradient {
            // There is still only one gradient, switch it to a fork
            let shape = {
                let grad = var_store.get(gradient);
                grad.shape().to_vec()
            };
            // Create variable for gradient sum
            self.gradient = Some(var_store.add(Tensor::new(ctx, shape, TensorMode::Mut)));
            self.gradients.push(gradient);
            self.gradients.push(v);
        } else {
            // This is the only gradient so far, so we don't need to sum anything
            self.gradient = Some(v);
        }
    }
}

#[derive(Copy, Clone)]
pub enum GradIndex {
    InVar(usize),
    OutVar(NodeIndex, usize),
}

impl GradIndex {
    pub fn get<'a>(&self, graph: &'a Graph) -> Ref<'a, Tensor<f32>> {
        match *self {
            GradIndex::InVar(in_grad_index) => {
                graph.in_grad[in_grad_index].get().get(graph)
            },
            GradIndex::OutVar(node, out_index) => {
                node.get(graph).out_grad[out_index].get().get(graph)
            },
        }
    }

    pub fn read(&self, g: &Graph, a: &mut Array<f32>) {
        self.get(g).read(g.context(), a);
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

#[derive(Copy, Clone)]
pub struct NodeIndex(usize);

impl NodeIndex {
    pub fn get<'a>(&self, g: &'a Graph) -> &'a Node {
        &g.nodes[self.0]
    }
}

#[test]
fn it_works() {
    use super::op::MatMul;

    let ctx = Rc::new(ga::Context::new());

    // Setup the graph
    let mut graph = Graph::new(ctx.clone());
    let a = graph.add_variable(vec![1, 2], true, vec![1.4, 0.3]);
    let wa = graph.add_variable(vec![2, 3], true, vec![0.5, 0.3, 0.2,
                                                       0.6, 0.7, 0.7]);
    let node = graph.add_node(MatMul(a, wa));
    let node_out = node.get(&graph).outputs[0];
    let node_g = graph.add_gradient(node_out);

    // Send some input data
    let node_g_cpu = ga::Array::from_vec(vec![1, 3], vec![1.0, -1.0, 0.5]);
    node_g.get(&graph).set(&ctx, &node_g_cpu);

    // Run the network
    graph.forward();
    graph.backward();
    let out = node.get(&graph).outputs[0].get(&graph).get(&ctx);
    let wa_d = graph.get_gradient(wa).get(&graph).get(&ctx);
    println!("out = {:?}", out);
    println!("wa_d = {:?}", wa_d);
    assert!(out.buffer() == &[0.88, 0.63, 0.49]);
    assert!(wa_d.buffer() == &[1.4, -1.4, 0.7,
                               0.3, -0.3, 0.15]);
}
<commit_msg>Don't need to borrow vars mutably in OutGrad::maybe_sum<commit_after>use std::cell::Ref;
use std::collections::HashMap;
use std::rc::Rc;

use ga::{self, Array, Tensor, TensorMode};
use rand;

use super::init::Initializer;
use super::op::{OpBuilder, OpDescriptor, Operation};
use super::var_store::{VarIndex, VarStore};

#[derive(Copy, Clone)]
pub enum NodeInput {
    Var(VarIndex),    // Regular input variable
    Recurrent(usize), // Recurrent connection
}

pub struct Node {
    pub inputs: Vec<VarIndex>,
    pub outputs: Vec<VarIndex>,
    pub in_grad: Vec<VarIndex>, // gradients on inputs
    pub out_grad: Vec<OutGrad>, // gradients on outputs
}

pub struct Graph {
    ctx: Rc<ga::Context>,
    nodes: Vec<Node>,
    node_ops: Vec<Box<Operation>>,
    pub var_store: VarStore,
    out_var_map: HashMap<VarIndex, (NodeIndex, usize)>, // Maps output variable to its node and index within node
    // Gradients on variables that are inputs to the graph - they have no corresponding node
    in_var_map: HashMap<VarIndex, usize>,
    learnables: Vec<(VarIndex, GradIndex)>, // Learnable variables
    in_grad: Vec<OutGrad>, // Gradients on variables that are inputs to the graph

    rng: rand::ThreadRng,
}

impl Graph {
    pub fn new(ctx: Rc<ga::Context>) -> Self {
        Graph {
            ctx: ctx,
            nodes: vec![],
            node_ops: vec![],
            var_store: VarStore::new(),
            out_var_map: HashMap::new(),
            in_var_map: HashMap::new(),
            learnables: vec![],
            in_grad: vec![],
            rng: rand::thread_rng(),
        }
    }

    pub fn add_node<T: OpBuilder>(&mut self, op: T) -> NodeIndex {
        let node_index = NodeIndex(self.nodes.len());

        let OpDescriptor { op, inputs: node_inputs, out_shapes } = op.build(&self.ctx, &self.var_store).unwrap();

        // Create output variables
        let mut outputs = vec![];
        for (i, shape) in out_shapes.into_iter().enumerate() {
            let var_index = self.var_store.add(Tensor::new(self.ctx.as_ref(), shape, TensorMode::Mut));
            outputs.push(var_index);
            self.out_var_map.insert(var_index, (node_index, i));
        }
        let mut out_grad = vec![OutGrad::new(); outputs.len()];
        // Set up inputs and gradients on inputs
        let mut inputs = vec![];
        let mut in_grad = vec![];
        for input in node_inputs {
            let (v, gradient) =
                match input {
                    NodeInput::Var(v) => (v, self.add_gradient(v)),
                    NodeInput::Recurrent(out) => {
                        let v = outputs[out];
                        let gradient = Self::create_gradient(&self.ctx, &mut self.var_store, v,
                                                             &mut out_grad[out]);
                        (v, gradient)
                    }
                };
            inputs.push(v);
            in_grad.push(gradient);
        }
        // Create the node
        self.nodes.push(Node { inputs: inputs,
                               outputs: outputs,
                               in_grad: in_grad,
                               out_grad: out_grad });
        // Add the corresponding node op
        self.node_ops.push(Box::new(op));
        node_index
    }

    pub fn add_variable<I: Initializer>(&mut self,
                                        shape: Vec<usize>,
                                        learnable: bool,
                                        init: I) -> VarIndex {
        let a = init.init(&mut self.rng, shape);
        let v = self.var_store.add(Tensor::from_array(&self.ctx, &a, TensorMode::Mut));
        self.in_var_map.insert(v, self.in_grad.len());
        if learnable {
            self.learnables.push((v, GradIndex::InVar(self.in_grad.len())));
        }
        self.in_grad.push(OutGrad::new());
        v
    }

    pub fn get_gradient(&self, v: VarIndex) -> GradIndex {
        match self.out_var_map.get(&v).map(|x| *x) {
            Some((node, out_index)) => {
                // v is the output of a node
                GradIndex::OutVar(node, out_index)
            },
            None => {
                // v is an input to the graph - it has no corresponding node
                let in_grad_index = *self.in_var_map.get(&v)
                                         .expect("Variable is neither input nor output. Nonsense!");
                GradIndex::InVar(in_grad_index)
            },
        }
    }

    pub fn add_gradient(&mut self, v: VarIndex) -> VarIndex {
        match self.out_var_map.get(&v).map(|x| *x) {
            Some((node, out_index)) => {
                // v is the output of a node
                /*self.nodes[node.0].out_grad[out_index]
                                  .fork(&self.ctx, &mut self.var_store, gradient);*/
                Self::create_gradient(&self.ctx, &mut self.var_store, v,
                                      &mut self.nodes[node.0].out_grad[out_index])
            },
            None => {
                // v is an input to the graph - it has no corresponding node
                let in_grad_index = *self.in_var_map.get(&v)
                                         .expect("Variable is neither input nor output. Nonsense!");
                /*self.in_grad[in_grad_index]
                    .fork(&self.ctx, &mut self.var_store, gradient);*/
                Self::create_gradient(&self.ctx, &mut self.var_store, v,
                                      &mut self.in_grad[in_grad_index])
            },
        }
    }

    fn create_gradient(ctx: &ga::Context,
                       var_store: &mut VarStore,
                       v: VarIndex,
                       out_grad: &mut OutGrad)
                       -> VarIndex {
        let shape = var_store.get(v).shape().to_owned();
        let gradient = var_store.add(Tensor::new(ctx, shape, TensorMode::Mut));
        out_grad.fork(ctx, var_store, gradient);
        gradient
    }

    pub fn forward(&mut self) {
        // Forward pass
        //
        // NOTE: We just execute the nodes in order. We can do this because of the way the graph is
        // built. When a user wants to add a node, he/she must also supply the inputs. This means
        // any dependencies must already be added before the node can be added. Therefore, we can
        // assert that all dependents come after their dependencies in the `self.nodes` array.
        for (node, op) in self.nodes.iter().zip(&mut self.node_ops) {
            op.forward(&self.ctx, &self.var_store, node);
        }
    }

    pub fn backward(&mut self) {
        // Backward pass
        // We run through the nodes in reverse order. See note in Graph::forward
        for (node, op) in self.nodes.iter_mut().rev().zip(self.node_ops.iter_mut().rev()) {
            // Sum the gradients on each output if there are multiple gradients
            for out_grad in &node.out_grad {
                out_grad.maybe_sum(self.ctx.as_ref(), &mut self.var_store);
            }
            op.backward(&self.ctx, &mut self.var_store, node);
        }
    }

    pub fn context(&self) -> &ga::Context {
        &self.ctx
    }

    pub fn learnables(&self) -> &[(VarIndex, GradIndex)] {
        &self.learnables
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

#[derive(Clone)]
pub struct OutGrad {
    gradient: Option<VarIndex>, // The gradient or sum of gradients
    gradients: Vec<VarIndex>,
}

impl OutGrad {
    pub fn new() -> Self {
        OutGrad {
            gradient: None,
            gradients: vec![],
        }
    }

    pub fn get(&self) -> VarIndex {
        self.gradient.unwrap()
    }

    pub fn try_get(&self) -> Option<VarIndex> {
        self.gradient
    }

    fn maybe_sum(&self, ctx: &ga::Context, var_store: &mut VarStore) {
        if self.gradients.len() > 0 {
            if let Some(sum) = self.gradient {
                ga::copy_to(ctx, &var_store.get(self.gradients[0]), &var_store.get(sum));
                for grad in &self.gradients[1..] {
                    ga::add(ctx, &var_store.get(sum), -1, &var_store.get(*grad), &var_store.get(sum));
                }
            }
        }
    }

    fn fork(&mut self, ctx: &ga::Context, var_store: &mut VarStore, v: VarIndex) {
        if self.gradients.len() > 0 {
            // There are multiple gradients already, just add the new one to the list
            self.gradients.push(v);
        } else if let Some(gradient) = self.gradient {
            // There is still only one gradient, switch it to a fork
            let shape = {
                let grad = var_store.get(gradient);
                grad.shape().to_vec()
            };
            // Create variable for gradient sum
            self.gradient = Some(var_store.add(Tensor::new(ctx, shape, TensorMode::Mut)));
            self.gradients.push(gradient);
            self.gradients.push(v);
        } else {
            // This is the only gradient so far, so we don't need to sum anything
            self.gradient = Some(v);
        }
    }
}

#[derive(Copy, Clone)]
pub enum GradIndex {
    InVar(usize),
    OutVar(NodeIndex, usize),
}

impl GradIndex {
    pub fn get<'a>(&self, graph: &'a Graph) -> Ref<'a, Tensor<f32>> {
        match *self {
            GradIndex::InVar(in_grad_index) => {
                graph.in_grad[in_grad_index].get().get(graph)
            },
            GradIndex::OutVar(node, out_index) => {
                node.get(graph).out_grad[out_index].get().get(graph)
            },
        }
    }

    pub fn read(&self, g: &Graph, a: &mut Array<f32>) {
        self.get(g).read(g.context(), a);
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

#[derive(Copy, Clone)]
pub struct NodeIndex(usize);

impl NodeIndex {
    pub fn get<'a>(&self, g: &'a Graph) -> &'a Node {
        &g.nodes[self.0]
    }
}

#[test]
fn it_works() {
    use super::op::MatMul;

    let ctx = Rc::new(ga::Context::new());

    // Setup the graph
    let mut graph = Graph::new(ctx.clone());
    let a = graph.add_variable(vec![1, 2], true, vec![1.4, 0.3]);
    let wa = graph.add_variable(vec![2, 3], true, vec![0.5, 0.3, 0.2,
                                                       0.6, 0.7, 0.7]);
    let node = graph.add_node(MatMul(a, wa));
    let node_out = node.get(&graph).outputs[0];
    let node_g = graph.add_gradient(node_out);

    // Send some input data
    let node_g_cpu = ga::Array::from_vec(vec![1, 3], vec![1.0, -1.0, 0.5]);
    node_g.get(&graph).set(&ctx, &node_g_cpu);

    // Run the network
    graph.forward();
    graph.backward();
    let out = node.get(&graph).outputs[0].get(&graph).get(&ctx);
    let wa_d = graph.get_gradient(wa).get(&graph).get(&ctx);
    println!("out = {:?}", out);
    println!("wa_d = {:?}", wa_d);
    assert!(out.buffer() == &[0.88, 0.63, 0.49]);
    assert!(wa_d.buffer() == &[1.4, -1.4, 0.7,
                               0.3, -0.3, 0.15]);
}
<|endoftext|>


<commit_before>use {HashSet};
use std::collections::{HashMap, RingBuf, Deque};

// call this Unigraph instead?
pub trait Graph<T> {
    /// Insert a new node, returning its index
    fn add_node(&mut self, val: T) -> NodeIndex;

    /// Insert a new edge, returning an error if one of the indices is invalid
    /// and a result otherwise
    fn add_edge(&mut self, i: NodeIndex, j: NodeIndex) -> GraphResult<bool>;

    /// Return an iterator over the out-neighbors for a given node
    fn adj(&self, i: NodeIndex) -> NodeIndices;

    /// Return the number of nodes in the graph
    fn num_nodes(&self) -> uint;

    fn node_indices(&self) -> NodeIndices;

    fn dfs(&self) -> Digraph<NodeIndex> {
        let mut tree = Digraph::new();

        if self.num_nodes() == 0 {
            return tree;
        }

        let mut visited = HashSet::new();
        let mut discovered = vec!();

        discovered.push((0, None));
        loop {
            match discovered.pop() {
                None => break,
                Some((ind, parent)) => {
                    tree.add_node(ind);
                    if parent.is_some() {
                        tree.add_edge(parent.unwrap(), ind);
                    }
                    visited.insert(ind);

                    for i in self.adj(ind) {
                        if !visited.contains(i) {
                            discovered.push((*i, Some(ind)));
                        }
                    }
                }
            }
        }
        return tree;

    }
}

struct GraphError {
    error_type: GraphErrorType,
}

impl GraphError {
    fn invalid_index(ind: NodeIndex) -> GraphError {
        GraphError { error_type: InvalidNodeIndex(ind) }
    }
}

type GraphResult<T> = Result<T, GraphError>;

enum GraphErrorType {
    InvalidNodeIndex(NodeIndex),
}

pub type NodeIndex = uint;
type NodeIndexSet = HashSet<NodeIndex>;

#[deriving(Show, Clone)]
struct Node<T> {
    data: T,
    index: NodeIndex,
}

/// Undirected graph. Allows loops.
#[deriving(Show)]
pub struct UnGraph<T> {
    nodes: Vec<Node<T>>,
    adj: Vec<NodeIndexSet>,
    num_nodes: uint,
}

/// Directed graph. Allows loops.
#[deriving(Show)]
pub struct Digraph<T> {
    nodes: Vec<Node<T>>,
    in_adj: Vec<NodeIndexSet>,
    out_adj: Vec<NodeIndexSet>,
    num_nodes: uint,
}

impl<T> UnGraph<T> {
    pub fn new() -> UnGraph<T> {
        UnGraph { nodes: vec!(), adj: vec!(), num_nodes: 0 }
    }

    pub fn degree(&mut self, ind: NodeIndex) -> Result<uint, GraphError> {
        if ind >= self.num_nodes {
            Err(GraphError::invalid_index(ind))
        } else {
            Ok(self.adj[ind].len())
        }
    }
}
impl<T: Clone> UnGraph<T> {
    /// Returns a new graph induced by a set of node indices
    pub fn induced_subgraph(&self, nodes: &NodeIndexSet) -> UnGraph<T> {
        let mut new = UnGraph::new();
        let mut ind_map = HashMap::new();

        for (i, ind) in nodes.iter().enumerate() {
            ind_map.insert(ind, i);
            new.add_node(self.nodes[*ind].data.clone());
        }

        // Here we're assuming that NodeIndex = uint. Not sure how to easily do otherwise
        for i in range(0, self.num_nodes) {
            for j in self.adj[i].iter() {
                new.add_edge(i, *j);
            }
        }

        new
    }
}

impl<T> Graph<T> for UnGraph<T> {
    fn add_node(&mut self, val: T) -> NodeIndex {
        let ind = self.num_nodes;
        self.nodes.push(Node { data: val, index: ind });
        self.adj.push(HashSet::new());
        self.num_nodes += 1;
        ind
    }

    fn add_edge(&mut self, i: NodeIndex, j: NodeIndex) -> GraphResult<bool> {
        if i >= self.num_nodes {
            Err(GraphError::invalid_index(i))
        } else if j >= self.num_nodes {
            Err(GraphError::invalid_index(j))
        } else {
            if self.adj[i].contains(&j) {
                Ok(false)
            } else {
                self.adj.get_mut(i).insert(j);
                self.adj.get_mut(j).insert(i);
                Ok(true)
            }
        }
    }

    fn adj(&self, i: NodeIndex) -> NodeIndices {
        NodeIndices::from_set(&self.adj[i])
    }

    fn num_nodes(&self) -> uint {
        self.num_nodes
    }

    fn node_indices(&self) -> NodeIndices {
        NodeIndices::from_node_vec(&self.nodes)
    }
}

impl<T> Digraph<T> {
    pub fn new() -> Digraph<T> {
        Digraph { nodes: vec!(), in_adj: vec!(), out_adj: vec!(), num_nodes: 0 }
    }

    pub fn out_degree(&mut self, ind: NodeIndex) -> Result<uint, GraphError> {
        if ind >= self.num_nodes {
            Err(GraphError::invalid_index(ind))
        } else {
            Ok(self.out_adj[ind].len())
        }
    }

    pub fn in_degree(&mut self, ind: NodeIndex) -> Result<uint, GraphError> {
        if ind >= self.num_nodes {
            Err(GraphError::invalid_index(ind))
        } else {
            Ok(self.in_adj[ind].len())
        }
    }
}

impl<T: Clone> Digraph<T> {
    /// Returns a new graph induced by a set of node indices
    pub fn induced_subgraph(&self, nodes: &NodeIndexSet) -> Digraph<T> {
        let mut new = Digraph::new();
        let mut ind_map = HashMap::new();

        for (i, ind) in nodes.iter().enumerate() {
            ind_map.insert(ind, i);
            new.add_node(self.nodes[*ind].data.clone());
        }

        // Here we're assuming that NodeIndex = uint. Not sure how to easily do otherwise
        for i in range(0, self.num_nodes) {
            for j in self.in_adj[i].iter() {
                new.add_edge(*j, i);
            }

            for j in self.out_adj[i].iter() {
                new.add_edge(i, *j);
            }
        }

        new
    }
}

impl<T> Graph<T> for Digraph<T> {
    fn add_node(&mut self, val: T) -> NodeIndex {
        let ind = self.num_nodes;
        self.nodes.push(Node { data: val, index: ind });
        self.in_adj.push(HashSet::new());
        self.out_adj.push(HashSet::new());
        self.num_nodes += 1;
        ind
    }

    fn add_edge(&mut self, i: NodeIndex, j: NodeIndex) -> GraphResult<bool> {
        if i >= self.num_nodes {
            Err(GraphError::invalid_index(i))
        } else if j >= self.num_nodes {
            Err(GraphError::invalid_index(j))
        } else {
            if self.out_adj[i].contains(&j) {
                Ok(false)
            } else {
                self.out_adj.get_mut(i).insert(j);
                self.in_adj.get_mut(j).insert(i);
                Ok(true)
            }
        }
    }

    fn adj(&self, i: NodeIndex) -> NodeIndices {
        NodeIndices::from_set(&self.out_adj[i])
    }

    fn num_nodes(&self) -> uint {
        self.num_nodes
    }

    fn node_indices(&self) -> NodeIndices {
        NodeIndices::from_node_vec(&self.nodes)
    }
}

struct NodeIndices {
    indices: Vec<NodeIndex>,
    curr: uint,
}

impl Iterator<NodeIndex> for NodeIndices {
    fn next(&mut self) -> Option<NodeIndex> {
        if self.curr < self.indices.len() {
            self.curr += 1;
            Some(self.indices[self.curr - 1])
        } else {
            None
        }
    }
}

impl NodeIndices {
    fn from_set(set: &NodeIndexSet) -> NodeIndices {
        let mut vec = vec!();
        for &i in set.iter() {
            vec.push(i);
        }
        NodeIndices { indices: vec, curr: 0 }
    }

    fn from_node_vec<T>(vec: &Vec<Node<T>>) -> NodeIndices {
        let mut new = vec!();
        for n in vec.iter() {
            new.push(n.index);
        }

        NodeIndices { indices: new, curr: 0 }
    }
}
<commit_msg>Change UnGraph, Digraph to use hashmaps for nodes/adj lists instead of vectors.<commit_after>use {HashSet};
use std::collections::{HashMap, RingBuf, Deque};

// call this Unigraph instead?
pub trait Graph<T> {
    /// Insert a new node, returning its index
    fn add_node(&mut self, val: T) -> NodeIndex;

    /// Insert a new edge, returning an error if one of the indices is invalid
    /// and a result otherwise
    fn add_edge(&mut self, i: NodeIndex, j: NodeIndex) -> GraphResult<bool>;

    /// Return an iterator over the out-neighbors for a given node
    fn adj(&self, i: NodeIndex) -> NodeIndices;

    /// Return the number of nodes in the graph
    fn num_nodes(&self) -> uint;

    fn node_indices(&self) -> NodeIndices;

    fn dfs(&self) -> Digraph<NodeIndex> {
        let mut tree = Digraph::new();

        if self.num_nodes() == 0 {
            return tree;
        }

        let mut visited = HashSet::new();
        let mut discovered = vec!();

        discovered.push((0, None));
        loop {
            match discovered.pop() {
                None => break,
                Some((ind, parent)) => {
                    tree.add_node(ind);
                    if parent.is_some() {
                        tree.add_edge(parent.unwrap(), ind);
                    }
                    visited.insert(ind);

                    for i in self.adj(ind) {
                        if !visited.contains(i) {
                            discovered.push((*i, Some(ind)));
                        }
                    }
                }
            }
        }
        return tree;

    }
}

struct GraphError {
    error_type: GraphErrorType,
}

impl GraphError {
    fn invalid_index(ind: NodeIndex) -> GraphError {
        GraphError { error_type: InvalidNodeIndex(ind) }
    }
}

type GraphResult<T> = Result<T, GraphError>;

enum GraphErrorType {
    InvalidNodeIndex(NodeIndex),
}

pub type NodeIndex = uint;
type NodeIndexSet = HashSet<NodeIndex>;

#[deriving(Show, Clone)]
struct Node<T> {
    data: T,
    index: NodeIndex,
}

/// Undirected graph. Allows loops.
#[deriving(Show)]
pub struct UnGraph<T> {
    nodes: HashMap<NodeIndex, Node<T>>,
    adj: HashMap<NodeIndex, NodeIndexSet>,
    num_nodes: uint,
}

/// Directed graph. Allows loops.
#[deriving(Show)]
pub struct Digraph<T> {
    nodes: HashMap<NodeIndex, Node<T>>,
    in_adj: HashMap<NodeIndex, NodeIndexSet>,
    out_adj: HashMap<NodeIndex, NodeIndexSet>,
    num_nodes: uint,
}

impl<T> UnGraph<T> {
    pub fn new() -> UnGraph<T> {
        UnGraph { nodes: HashMap::new(), adj: HashMap::new(), num_nodes: 0 }
    }

    pub fn degree(&mut self, ind: NodeIndex) -> Result<uint, GraphError> {
        if !self.adj.contains_key(&ind) {
            Err(GraphError::invalid_index(ind))
        } else {
            Ok( self.get_adj(ind).len() )
        }
    }

    pub fn get_node<'a>(&'a self, i: NodeIndex) -> &'a Node<T> {
        self.nodes.find(&i).unwrap()
    }

    pub fn get_adj<'a>(&'a self, i: NodeIndex) -> &'a NodeIndexSet {
        self.adj.find(&i).unwrap()
    }
}
impl<T: Clone> UnGraph<T> {
    /// Returns a new graph induced by a set of node indices
    pub fn induced_subgraph(&self, nodes: &NodeIndexSet) -> UnGraph<T> {
        let mut new = UnGraph::new();
        let mut ind_map = HashMap::new();

        for (i, ind) in nodes.iter().enumerate() {
            ind_map.insert(ind, i);
            new.add_node(self.get_node(*ind).data.clone());
        }

        // Here we're assuming that NodeIndex = uint. Not sure how to easily do otherwise
        for i in range(0, self.num_nodes) {
            for j in self.get_adj(i).iter() {
                new.add_edge(i, *j);
            }
        }

        new
    }
}

impl<T> Graph<T> for UnGraph<T> {
    fn add_node(&mut self, val: T) -> NodeIndex {
        let ind = self.num_nodes;
        self.nodes.insert(ind, Node { data: val, index: ind });
        self.adj.insert(ind, HashSet::new());
        self.num_nodes += 1;
        ind
    }

    fn add_edge(&mut self, i: NodeIndex, j: NodeIndex) -> GraphResult<bool> {
        if !self.nodes.contains_key(&i) {
            Err(GraphError::invalid_index(i))
        } else if !self.nodes.contains_key(&j) {
            Err(GraphError::invalid_index(j))
        } else {
            if self.get_adj(i).contains(&j) {
                Ok(false)
            } else {
                self.adj.find_mut(&i).unwrap().insert(j);
                self.adj.find_mut(&j).unwrap().insert(i);
                Ok(true)
            }
        }
    }

    fn adj(&self, i: NodeIndex) -> NodeIndices {
        FromIterator::from_iter(self.get_adj(i).iter().map(|&x| x))
    }

    fn num_nodes(&self) -> uint {
        self.num_nodes
    }

    fn node_indices(&self) -> NodeIndices {
        FromIterator::from_iter(self.nodes.keys().map(|&x| x))
    }
}

impl<T> Digraph<T> {
    pub fn new() -> Digraph<T> {
        Digraph { nodes: HashMap::new(),
                  in_adj: HashMap::new(),
                  out_adj: HashMap::new(),
                  num_nodes: 0 }
    }

    pub fn out_degree(&mut self, ind: NodeIndex) -> Result<uint, GraphError> {
        if !self.nodes.contains_key(&ind) {
            Err(GraphError::invalid_index(ind))
        } else {
            Ok(self.get_out_adj(ind).len())
        }
    }

    pub fn in_degree(&mut self, ind: NodeIndex) -> Result<uint, GraphError> {
        if ind >= self.num_nodes {
            Err(GraphError::invalid_index(ind))
        } else {
            Ok(self.get_in_adj(ind).len())
        }
    }

    pub fn get_node<'a>(&'a self, i: NodeIndex) -> &'a Node<T> {
        self.nodes.find(&i).unwrap()
    }

    pub fn get_in_adj<'a>(&'a self, i: NodeIndex) -> &'a NodeIndexSet {
        self.in_adj.find(&i).unwrap()
    }

    pub fn get_out_adj<'a>(&'a self, i: NodeIndex) -> &'a NodeIndexSet {
        self.out_adj.find(&i).unwrap()
    }
}

impl<T: Clone> Digraph<T> {
    /// Returns a new graph induced by a set of node indices
    pub fn induced_subgraph(&self, nodes: &NodeIndexSet) -> Digraph<T> {
        let mut new = Digraph::new();
        let mut ind_map = HashMap::new();

        for (i, ind) in nodes.iter().enumerate() {
            ind_map.insert(ind, i);
            new.add_node(self.get_node(*ind).data.clone());
        }

        // Here we're assuming that NodeIndex = uint. Not sure how to easily do otherwise
        for i in range(0, self.num_nodes) {
            for j in self.get_in_adj(i).iter() {
                new.add_edge(*j, i);
            }

            for j in self.get_out_adj(i).iter() {
                new.add_edge(i, *j);
            }
        }

        new
    }
}

impl<T> Graph<T> for Digraph<T> {
    fn add_node(&mut self, val: T) -> NodeIndex {
        let ind = self.num_nodes;
        self.nodes.insert(ind, Node { data: val, index: ind });
        self.in_adj.insert(ind, HashSet::new());
        self.out_adj.insert(ind, HashSet::new());
        self.num_nodes += 1;
        ind
    }

    fn add_edge(&mut self, i: NodeIndex, j: NodeIndex) -> GraphResult<bool> {
        if !self.nodes.contains_key(&i) {
            Err(GraphError::invalid_index(i))
        } else if !self.nodes.contains_key(&j) {
            Err(GraphError::invalid_index(j))
        } else {
            if self.get_out_adj(i).contains(&j) {
                Ok(false)
            } else {
                self.out_adj.find_mut(&i).unwrap().insert(j);
                self.in_adj.find_mut(&j).unwrap().insert(i);
                Ok(true)
            }
        }
    }

    fn adj(&self, i: NodeIndex) -> NodeIndices {
        FromIterator::from_iter(self.get_out_adj(i).iter().map(|&x| x))
    }

    fn num_nodes(&self) -> uint {
        self.num_nodes
    }

    fn node_indices(&self) -> NodeIndices {
        FromIterator::from_iter(self.nodes.keys().map(|&x| x))
    }
}

struct NodeIndices {
    indices: Vec<NodeIndex>,
    curr: uint,
}

impl Iterator<NodeIndex> for NodeIndices {
    fn next(&mut self) -> Option<NodeIndex> {
        if self.curr < self.indices.len() {
            self.curr += 1;
            Some(self.indices[self.curr - 1])
        } else {
            None
        }
    }
}

impl FromIterator<NodeIndex> for NodeIndices {
    fn from_iter<T: Iterator<NodeIndex>>(mut it: T) -> NodeIndices {
        let mut vec = vec!();
        for i in it {
            vec.push(i);
        }
        NodeIndices { indices: vec, curr: 0 }
    }

}
<|endoftext|>


<commit_before>//! This is initial draft implementation of HTTP/1.x protocol. It's based
//! on `rotor::transports::greedy_stream` which is not very DDoS-safe. Also
//! this implementation has hard-coded limits on certain request parameters.
//! Otherwise implementation should be good enough for general consumption
//!
//!
use std::io::Write;
use std::error::Error;
use std::marker::PhantomData;
use std::mem::replace;

use time::now_utc;
use rotor::transports::stream::{Transport, Protocol};
use rotor::transports::StreamSocket;
use rotor::buffer_util::find_substr;
use rotor::async::Async;
use hyper::version::HttpVersion;
use hyper::method::Method;
use hyper::header::{Headers, Date, HttpDate, ContentLength};
use hyper::uri::RequestUri;
use hyper::status::StatusCode;
use httparse;

/// Note httparse requires we preallocate array of this size so be wise
pub const MAX_HEADERS_NUM: usize = 256;
/// This one is not preallocated, but too large buffer is of limited use
/// because of previous parameter.
pub const MAX_HEADERS_SIZE: usize = 16384;
/// This is not "enough for everyone" but we probably need some limit anyway.
/// Note that underlying `netbuf` impl limited to a little less than 4GiB
pub const MAX_BODY_SIZE: usize = 104_856_700;

pub trait Handler<C> {
    /// Dispatched when request arrives.
    ///
    /// We don't support POST body yet, so this is only one callback, but will
    /// probably be split into many in future
    fn request(_request: Request, _response: &mut ResponseBuilder,
               _ctx: &mut C)
    {
        // The 404 or BadRequest for all requests
    }
}

#[derive(Debug)]
pub enum ResponseFsm {
    Head {
        status: StatusCode,
        version: HttpVersion,
        headers: Headers,
    },
    // TODO(tailhook) WritingFixed { bytes_remaining: u64 }
    // TODO(tailhook) WritingChunked
    End,
}

pub struct ResponseBuilder<'a, 'b: 'a>{
    state: ResponseFsm,
    transport: &'a mut Transport<'b>,
}

pub enum Client<C, H:Handler<C>> {
    Initial,
    ReadHeaders,  // TODO(tailhook) 100 Expect?
    Processing(H, PhantomData<*const C>),
    // ReadFixedSize(Request, usize),  // TODO
    // ReadChunked(Request, usize),
    KeepAlive,
}

#[derive(Debug)]
pub struct Request {
    pub version: HttpVersion,
    pub method: Method,
    pub uri: RequestUri,
    pub headers: Headers,
    pub body: Vec<u8>,
}

fn parse_headers(transport: &mut Transport)
    -> Result<Option<Request>, Box<Error+Send+Sync>>
{
    use hyper::version::HttpVersion::*;

    let mut buf = transport.input();
    let headers_end = match find_substr(&buf[..], b"\r\n\r\n") {
        Some(x) => x,
        None => { return Ok(None); }
    };
    let mut headers = [httparse::EMPTY_HEADER; MAX_HEADERS_NUM];
    let req = {
        let mut raw = httparse::Request::new(&mut headers);
        match raw.parse(&buf[..]) {
            Ok(httparse::Status::Complete(x)) => {
                assert!(x == headers_end+4);
            }
            Ok(_) => unreachable!(),
            Err(_) => {
                return Err(From::from("Header syntax mismatch"));
            }
        }
        Request {
            version: if raw.version.unwrap() == 1 { Http11 } else { Http10 },
            method: try!(raw.method.unwrap().parse()),
            uri: try!(raw.path.unwrap().parse()),
            headers: try!(Headers::from_raw(raw.headers)),
            body: Vec::new(),
        }
    };
    buf.consume(headers_end+4);
    Ok(Some(req))
}

impl<C, H: Handler<C>> Client<C, H> {
    fn handle_request(&self, req: Request, transport: &mut Transport,
                        ctx: &mut C)
    {
        let mut bld = ResponseBuilder {
            state: ResponseFsm::Head {
                status: if req.method == Method::Get
                    { StatusCode::NotFound } else { StatusCode::BadRequest },
                version: req.version,
                headers: Headers::new(),
            },
            transport: transport,
        };
        <H as Handler<C>>::request(req, &mut bld, ctx);
        bld.default_body();
    }
}

impl<C, H: Handler<C>> Protocol<C> for Client<C, H> {
    fn accepted<S: StreamSocket>(_conn: &mut S, _context: &mut C)
        -> Option<Self>
    {
        Some(Client::Initial)
    }
    fn data_received(self, transport: &mut Transport, ctx: &mut C)
        -> Async<Self, ()>
    {
        use self::Client::*;
        match self {
            // We keep Initial and KeepAlive states separate from
            // ReadHeaders just for debugging and for potentially different
            // timeouts in future
            Initial|ReadHeaders|KeepAlive => {
                match parse_headers(transport) {
                    Err(e) => {
                        debug!("Error parsing HTTP headers: {}", e);
                        Async::Stop
                    }
                    Ok(None) => {
                        Async::Continue(ReadHeaders, ())
                    }
                    Ok(Some(req)) => {
                        self.handle_request(req, transport, ctx);
                        Async::Continue(KeepAlive, ())
                    }
                }
            }
            Processing(_, _) => {
                unimplemented!();
            }
        }
    }
}

impl<'a, 'b> ResponseBuilder<'a, 'b> {

    /// Write request body in a single go
    pub fn put_body<B:AsRef<[u8]>>(&mut self, body: B) {
        use self::ResponseFsm::*;
        match replace(&mut self.state, End) {
            Head { status, version, mut headers } => {
                let body = body.as_ref();
                let out = self.transport.output();
                write!(out, "{} {}\r\n", version, status).unwrap();
                if !headers.has::<Date>() {
                    headers.set(Date(HttpDate(now_utc())));
                }
                headers.set(ContentLength(body.len() as u64));
                write!(out, "{}\r\n", headers).unwrap();
                out.extend(body);
            }
            state => {
                panic!("Too late to send body in state: {:?}", state);
            }
        }
    }

    pub fn set_status(&mut self, new_status: StatusCode) {
        use self::ResponseFsm::*;
        match self.state {
            Head { ref mut status, .. } => {
                *status = new_status;
            }
            ref state => {
                panic!("Too late to set status in state: {:?}", state);
            }
        }
    }

    pub fn default_body(&mut self) {
        use self::ResponseFsm::*;
        match self.state {
            Head { status, .. } => {
                // TODO(tailhook) maybe assert on 200 Ok status?
                self.put_body(&format!("{}", status));
                self.state = End;
            }
            End => {}
        }
    }
}
<commit_msg>Support reading fixed length bodys.<commit_after>//! This is initial draft implementation of HTTP/1.x protocol. It's based
//! on `rotor::transports::greedy_stream` which is not very DDoS-safe. Also
//! this implementation has hard-coded limits on certain request parameters.
//! Otherwise implementation should be good enough for general consumption
//!
//!
use std::cmp;
use std::io::Write;
use std::error::Error;
use std::marker::PhantomData;
use std::mem::replace;

use time::now_utc;
use rotor::transports::stream::{Transport, Protocol};
use rotor::transports::StreamSocket;
use rotor::buffer_util::find_substr;
use rotor::async::Async;
use hyper::version::HttpVersion;
use hyper::method::Method;
use hyper::header::{Headers, Date, HttpDate, ContentLength};
use hyper::uri::RequestUri;
use hyper::status::StatusCode;
use httparse;

/// Note httparse requires we preallocate array of this size so be wise
pub const MAX_HEADERS_NUM: usize = 256;
/// This one is not preallocated, but too large buffer is of limited use
/// because of previous parameter.
pub const MAX_HEADERS_SIZE: usize = 16384;
/// This is not "enough for everyone" but we probably need some limit anyway.
/// Note that underlying `netbuf` impl limited to a little less than 4GiB
pub const MAX_BODY_SIZE: usize = 104_856_700;

pub trait Handler<C> {
    /// Dispatched when request arrives.
    ///
    /// We don't support POST body yet, so this is only one callback, but will
    /// probably be split into many in future
    fn request(_request: Request, _response: &mut ResponseBuilder,
               _ctx: &mut C)
    {
        // The 404 or BadRequest for all requests
    }
}

#[derive(Debug)]
pub enum ResponseFsm {
    Head {
        status: StatusCode,
        version: HttpVersion,
        headers: Headers,
    },
    // TODO(tailhook) WritingFixed { bytes_remaining: u64 }
    // TODO(tailhook) WritingChunked
    End,
}

pub struct ResponseBuilder<'a, 'b: 'a>{
    state: ResponseFsm,
    transport: &'a mut Transport<'b>,
}

impl <'a, 'b>ResponseBuilder<'a, 'b> {
    pub fn new(req: &Request, transport: &'a mut Transport<'b>)
        -> ResponseBuilder<'a, 'b>
    {
        ResponseBuilder {
            state: ResponseFsm::Head {
                status: if req.method == Method::Get {
                    StatusCode::NotFound
                } else {
                    StatusCode::BadRequest
                },
                version: req.version,
                headers: Headers::new(),
            },
            transport: transport,
        }
    }
}

/// A connection with a client.
///
/// The `Initial`, `KeepAlive` and `ReadHeaders` states are kept separate for
/// debugging and different timeouts in future eventuallly.
pub enum Client<C, H: Handler<C>> {
    /// The initial state of a connection.
    Initial,
    /// The state after some headers have been read.
    ReadHeaders, // TODO(tailhook) 100 Expect?
    /// Not yet supported.
    Processing(H, PhantomData<*const C>),
    /// Reading a request body with a fixed size.
    ///
    /// The `usize` gives the number of remaining bytes.
    ReadFixedSize(Request, usize),
    // TODO ReadChunked(Request, usize),
    /// A connection in idle state.
    KeepAlive,
}

#[derive(Debug)]
pub struct Request {
    pub version: HttpVersion,
    pub method: Method,
    pub uri: RequestUri,
    pub headers: Headers,
    pub body: Vec<u8>,
}

impl Request {
    fn content_length(&self) -> Option<usize> {
        if let Some(length_header) = self.headers.get::<ContentLength>() {
            return Some(**length_header as usize)
        }
        None
    }
}

fn parse_headers(transport: &mut Transport)
    -> Result<Option<Request>, Box<Error+Send+Sync>>
{
    use hyper::version::HttpVersion::*;

    let mut buf = transport.input();
    let headers_end = match find_substr(&buf[..], b"\r\n\r\n") {
        Some(x) => x,
        None => { return Ok(None); }
    };
    let mut headers = [httparse::EMPTY_HEADER; MAX_HEADERS_NUM];
    let req = {
        let mut raw = httparse::Request::new(&mut headers);
        match raw.parse(&buf[..]) {
            Ok(httparse::Status::Complete(x)) => {
                assert!(x == headers_end+4);
            }
            Ok(_) => unreachable!(),
            Err(_) => {
                return Err(From::from("Header syntax mismatch"));
            }
        }
        Request {
            version: if raw.version.unwrap() == 1 { Http11 } else { Http10 },
            method: try!(raw.method.unwrap().parse()),
            uri: try!(raw.path.unwrap().parse()),
            headers: try!(Headers::from_raw(raw.headers)),
            body: Vec::new(),
        }
    };
    buf.consume(headers_end+4);
    Ok(Some(req))
}

impl<C, H: Handler<C>> Protocol<C> for Client<C, H> {
    fn accepted<S: StreamSocket>(_conn: &mut S, _context: &mut C)
        -> Option<Self>
    {
        Some(Client::Initial)
    }
    fn data_received(self, transport: &mut Transport, ctx: &mut C)
        -> Async<Self, ()>
    {
        use self::Client::*;
        match self {
            Initial | ReadHeaders | KeepAlive => {
                match parse_headers(transport) {
                    Err(_) => Async::Stop,
                    Ok(None) => Async::Continue(ReadHeaders, ()),
                    Ok(Some(req)) => {
                        if let Some(length) = req.content_length() {
                            Async::Continue(ReadFixedSize(req, length), ())
                        } else {
                            let mut bld = ResponseBuilder::new(&req,
                                                               transport);
                            <H as Handler<C>>::request(req, &mut bld, ctx);
                            bld.default_body();
                            Async::Continue(Client::KeepAlive, ())
                        }
                    }
                }
            }
            ReadFixedSize(mut req, mut length) => {
                {
                    let mut buf = transport.input();
                    let read_length = cmp::min(length, buf.len());
                    length -= read_length;
                    req.body.extend(&buf[..read_length]);
                    buf.consume(read_length);
                }
                if length != 0 {
                    Async::Continue(ReadFixedSize(req, length), ())
                } else {
                    let mut bld = ResponseBuilder::new(&req, transport);
                    <H as Handler<C>>::request(req, &mut bld, ctx);
                    bld.default_body();
                    Async::Continue(Client::KeepAlive, ())
                }
            }
            _ => unimplemented!()
        }
    }
}

impl<'a, 'b> ResponseBuilder<'a, 'b> {

    /// Write request body in a single go
    pub fn put_body<B:AsRef<[u8]>>(&mut self, body: B) {
        use self::ResponseFsm::*;
        match replace(&mut self.state, End) {
            Head { status, version, mut headers } => {
                let body = body.as_ref();
                let out = self.transport.output();
                write!(out, "{} {}\r\n", version, status).unwrap();
                if !headers.has::<Date>() {
                    headers.set(Date(HttpDate(now_utc())));
                }
                headers.set(ContentLength(body.len() as u64));
                write!(out, "{}\r\n", headers).unwrap();
                out.extend(body);
            }
            state => {
                panic!("Too late to send body in state: {:?}", state);
            }
        }
    }

    pub fn set_status(&mut self, new_status: StatusCode) {
        use self::ResponseFsm::*;
        match self.state {
            Head { ref mut status, .. } => {
                *status = new_status;
            }
            ref state => {
                panic!("Too late to set status in state: {:?}", state);
            }
        }
    }

    pub fn default_body(&mut self) {
        use self::ResponseFsm::*;
        match self.state {
            Head { status, .. } => {
                // TODO(tailhook) maybe assert on 200 Ok status?
                self.put_body(&format!("{}", status));
                self.state = End;
            }
            End => {}
        }
    }
}
<|endoftext|>


<commit_before>use diesel::associations::HasTable;
use diesel::pg::Pg;
use diesel::prelude::*;
use serde_json;
use std::collections::HashMap;

use models::Crate;
use schema::badges;
use views::EncodableBadge;

#[derive(Debug, Queryable, Associations)]
#[belongs_to(Crate)]
#[table_name = "badges"]
pub struct CrateBadge {
    pub crate_id: i32,
    pub badge: Badge,
}

#[derive(Debug, PartialEq, Clone, Deserialize, Serialize)]
#[serde(rename_all = "kebab-case", tag = "badge_type", content = "attributes")]
pub enum Badge {
    TravisCi {
        repository: String,
        branch: Option<String>,
    },
    Appveyor {
        repository: String,
        id: Option<String>,
        branch: Option<String>,
        project_name: Option<String>,
        service: Option<String>,
    },
    #[serde(rename = "gitlab")]
    GitLab {
        repository: String,
        branch: Option<String>,
    },
    CircleCi {
        repository: String,
        branch: Option<String>,
    },
    IsItMaintainedIssueResolution {
        repository: String,
    },
    IsItMaintainedOpenIssues {
        repository: String,
    },
    Codecov {
        repository: String,
        branch: Option<String>,
        service: Option<String>,
    },
    Coveralls {
        repository: String,
        branch: Option<String>,
        service: Option<String>,
    },
    Maintenance {
        status: MaintenanceStatus,
    },
}

#[derive(Debug, PartialEq, Clone, Copy, Deserialize, Serialize)]
#[serde(rename_all = "kebab-case")]
pub enum MaintenanceStatus {
    ActivelyDeveloped,
    PassivelyMaintained,
    AsIs,
    None,
    Experimental,
    LookingForMaintainer,
    Deprecated,
}

// FIXME: Derive this in Diesel 1.3.
impl HasTable for CrateBadge {
    type Table = badges::table;

    fn table() -> Self::Table {
        badges::table
    }
}

impl Queryable<badges::SqlType, Pg> for Badge {
    type Row = (i32, String, serde_json::Value);

    fn build((_, badge_type, attributes): Self::Row) -> Self {
        let json = json!({"badge_type": badge_type, "attributes": attributes});
        serde_json::from_value(json).expect("Invalid CI badge in the database")
    }
}

impl Badge {
    pub fn encodable(self) -> EncodableBadge {
        serde_json::from_value(serde_json::to_value(self).unwrap()).unwrap()
    }

    pub fn update_crate<'a>(
        conn: &PgConnection,
        krate: &Crate,
        badges: Option<&'a HashMap<String, HashMap<String, String>>>,
    ) -> QueryResult<Vec<&'a str>> {
        use diesel::{delete, insert_into};

        let mut invalid_badges = vec![];
        let mut new_badges = vec![];

        if let Some(badges) = badges {
            for (k, v) in badges {
                let attributes_json = serde_json::to_value(v).unwrap();

                let json = json!({"badge_type": k, "attributes": attributes_json});
                if serde_json::from_value::<Badge>(json).is_ok() {
                    new_badges.push((
                        badges::crate_id.eq(krate.id),
                        badges::badge_type.eq(k),
                        badges::attributes.eq(attributes_json),
                    ));
                } else {
                    invalid_badges.push(&**k);
                }
            }
        }

        conn.transaction(|| {
            delete(badges::table)
                .filter(badges::crate_id.eq(krate.id))
                .execute(conn)?;
            insert_into(badges::table)
                .values(&new_badges)
                .execute(conn)?;
            Ok(invalid_badges)
        })
    }
}
<commit_msg>Add some docs to `CrateBadge`<commit_after>use diesel::associations::HasTable;
use diesel::pg::Pg;
use diesel::prelude::*;
use serde_json;
use std::collections::HashMap;

use models::Crate;
use schema::badges;
use views::EncodableBadge;

/// A combination of a `Badge` and a crate ID.
///
/// We don't typically care about the crate ID when dealing with badges.
/// However, when we are eager loading all badges for a group of crates, we need
/// the crate ID to group badges by their owner.
#[derive(Debug, Queryable, Associations)]
#[belongs_to(Crate)]
#[table_name = "badges"]
pub struct CrateBadge {
    pub crate_id: i32,
    pub badge: Badge,
}

#[derive(Debug, PartialEq, Clone, Deserialize, Serialize)]
#[serde(rename_all = "kebab-case", tag = "badge_type", content = "attributes")]
pub enum Badge {
    TravisCi {
        repository: String,
        branch: Option<String>,
    },
    Appveyor {
        repository: String,
        id: Option<String>,
        branch: Option<String>,
        project_name: Option<String>,
        service: Option<String>,
    },
    #[serde(rename = "gitlab")]
    GitLab {
        repository: String,
        branch: Option<String>,
    },
    CircleCi {
        repository: String,
        branch: Option<String>,
    },
    IsItMaintainedIssueResolution {
        repository: String,
    },
    IsItMaintainedOpenIssues {
        repository: String,
    },
    Codecov {
        repository: String,
        branch: Option<String>,
        service: Option<String>,
    },
    Coveralls {
        repository: String,
        branch: Option<String>,
        service: Option<String>,
    },
    Maintenance {
        status: MaintenanceStatus,
    },
}

#[derive(Debug, PartialEq, Clone, Copy, Deserialize, Serialize)]
#[serde(rename_all = "kebab-case")]
pub enum MaintenanceStatus {
    ActivelyDeveloped,
    PassivelyMaintained,
    AsIs,
    None,
    Experimental,
    LookingForMaintainer,
    Deprecated,
}

// FIXME: Derive this in Diesel 1.3.
impl HasTable for CrateBadge {
    type Table = badges::table;

    fn table() -> Self::Table {
        badges::table
    }
}

impl Queryable<badges::SqlType, Pg> for Badge {
    type Row = (i32, String, serde_json::Value);

    fn build((_, badge_type, attributes): Self::Row) -> Self {
        let json = json!({"badge_type": badge_type, "attributes": attributes});
        serde_json::from_value(json).expect("Invalid CI badge in the database")
    }
}

impl Badge {
    pub fn encodable(self) -> EncodableBadge {
        serde_json::from_value(serde_json::to_value(self).unwrap()).unwrap()
    }

    pub fn update_crate<'a>(
        conn: &PgConnection,
        krate: &Crate,
        badges: Option<&'a HashMap<String, HashMap<String, String>>>,
    ) -> QueryResult<Vec<&'a str>> {
        use diesel::{delete, insert_into};

        let mut invalid_badges = vec![];
        let mut new_badges = vec![];

        if let Some(badges) = badges {
            for (k, v) in badges {
                let attributes_json = serde_json::to_value(v).unwrap();

                let json = json!({"badge_type": k, "attributes": attributes_json});
                if serde_json::from_value::<Badge>(json).is_ok() {
                    new_badges.push((
                        badges::crate_id.eq(krate.id),
                        badges::badge_type.eq(k),
                        badges::attributes.eq(attributes_json),
                    ));
                } else {
                    invalid_badges.push(&**k);
                }
            }
        }

        conn.transaction(|| {
            delete(badges::table)
                .filter(badges::crate_id.eq(krate.id))
                .execute(conn)?;
            insert_into(badges::table)
                .values(&new_badges)
                .execute(conn)?;
            Ok(invalid_badges)
        })
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]

//! The `ParallelIterator` module makes it easy to write parallel
//! programs using an iterator-style interface. To get access to all
//! the methods you want, the easiest is to write `use
//! rayon::par_iter::*;` at the top of your module, which will import
//! the various traits and methods you need.
//!
//! The submodules of this module mostly just contain implementaton
//! details of little interest to an end-user.

use std::f64;
use std::ops::Fn;
use self::collect::collect_into;
use self::enumerate::Enumerate;
use self::filter::Filter;
use self::filter_map::FilterMap;
use self::map::Map;
use self::reduce::{reduce, ReduceOp, SumOp, MulOp, MinOp, MaxOp, ReduceWithOp,
                   SUM, MUL, MIN, MAX};
use self::state::ParallelIteratorState;
use self::weight::Weight;
use self::zip::ZipIter;

pub mod collect;
pub mod enumerate;
pub mod filter;
pub mod filter_map;
pub mod len;
pub mod for_each;
pub mod reduce;
pub mod slice;
pub mod slice_mut;
pub mod state;
pub mod map;
pub mod weight;
pub mod zip;
pub mod range;

#[cfg(test)]
mod test;

pub trait IntoParallelIterator {
    type Iter: ParallelIterator<Item=Self::Item>;
    type Item: Send;

    fn into_par_iter(self) -> Self::Iter;
}

pub trait IntoParallelRefIterator<'data> {
    type Iter: ParallelIterator<Item=&'data Self::Item>;
    type Item: Sync + 'data;

    fn par_iter(&'data self) -> Self::Iter;
}

pub trait IntoParallelRefMutIterator<'data> {
    type Iter: ParallelIterator<Item=&'data mut Self::Item>;
    type Item: Send + 'data;

    fn par_iter_mut(&'data mut self) -> Self::Iter;
}

/// The `ParallelIterator` interface.
pub trait ParallelIterator {
    type Item: Send;
    type Shared: Sync;
    type State: ParallelIteratorState<Shared=Self::Shared, Item=Self::Item> + Send;

    fn state(self) -> (Self::Shared, Self::State);

    /// Indicates the relative "weight" of producing each item in this
    /// parallel iterator. A higher weight will cause finer-grained
    /// parallel subtasks. 1.0 indicates something very cheap and
    /// uniform, like copying a value out of an array, or computing `x
    /// + 1`. If your tasks are either very expensive, or very
    /// unpredictable, you are better off with higher values. See also
    /// `weight_max`, which is a convenient shorthand to force the
    /// finest grained parallel execution posible. Tuning this value
    /// should not affect correctness but can improve (or hurt)
    /// performance.
    fn weight(self, scale: f64) -> Weight<Self>
        where Self: Sized
    {
        Weight::new(self, scale)
    }

    /// Shorthand for `self.weight(f64::INFINITY)`. This forces the
    /// smallest granularity of parallel execution, which makes sense
    /// when your parallel tasks are (potentially) very expensive to
    /// execute.
    fn weight_max(self) -> Weight<Self>
        where Self: Sized
    {
        self.weight(f64::INFINITY)
    }

    /// Executes `OP` on each item produced by the iterator, in parallel.
    fn for_each<OP>(self, op: OP)
        where OP: Fn(Self::Item) + Sync, Self: Sized
    {
        for_each::for_each(self, &op)
    }

    /// Applies `map_op` to each item of his iterator, producing a new
    /// iterator with the results.
    fn map<MAP_OP,R>(self, map_op: MAP_OP) -> Map<Self, MAP_OP>
        where MAP_OP: Fn(Self::Item) -> R, Self: Sized
    {
        Map::new(self, map_op)
    }

    /// Applies `map_op` to each item of his iterator, producing a new
    /// iterator with the results.
    fn filter<FILTER_OP>(self, filter_op: FILTER_OP) -> Filter<Self, FILTER_OP>
        where FILTER_OP: Fn(&Self::Item) -> bool, Self: Sized
    {
        Filter::new(self, filter_op)
    }

    /// Applies `map_op` to each item of his iterator, producing a new
    /// iterator with the results.
    fn filter_map<FILTER_OP,R>(self, filter_op: FILTER_OP) -> FilterMap<Self, FILTER_OP>
        where FILTER_OP: Fn(Self::Item) -> Option<R>, Self: Sized
    {
        FilterMap::new(self, filter_op)
    }

    /// Reduces the items in the iterator into one item using `op`.
    /// See also `sum`, `mul`, `min`, etc, which are slightly more
    /// efficient. Returns `None` if the iterator is empty.
    ///
    /// Note: unlike in a sequential iterator, the order in which `op`
    /// will be applied to reduce the result is not specified. So `op`
    /// should be commutative and associative or else the results will
    /// be non-deterministic.
    fn reduce_with<OP>(self, op: OP) -> Option<Self::Item>
        where Self: Sized, OP: Fn(Self::Item, Self::Item) -> Self::Item + Sync,
    {
        reduce(self.map(Some), &ReduceWithOp::new(op))
    }

    /// Sums up the items in the iterator.
    ///
    /// Note that the order in items will be reduced is not specified,
    /// so if the `+` operator is not truly commutative and
    /// associative (as is the case for floating point numbers), then
    /// the results are not fully deterministic.
    fn sum(self) -> Self::Item
        where Self: Sized, SumOp: ReduceOp<Self::Item>
    {
        reduce(self, SUM)
    }

    /// Multiplies all the items in the iterator.
    ///
    /// Note that the order in items will be reduced is not specified,
    /// so if the `*` operator is not truly commutative and
    /// associative (as is the case for floating point numbers), then
    /// the results are not fully deterministic.
    fn mul(self) -> Self::Item
        where Self: Sized, MulOp: ReduceOp<Self::Item>
    {
        reduce(self, MUL)
    }

    /// Computes the minimum of all the items in the iterator.
    ///
    /// Note that the order in items will be reduced is not specified,
    /// so if the `Ord` impl is not truly commutative and associative
    /// (as is the case for floating point numbers), then the results
    /// are not deterministic.
    fn min(self) -> Self::Item
        where Self: Sized, MinOp: ReduceOp<Self::Item>
    {
        reduce(self, MIN)
    }

    /// Computes the maximum of all the items in the iterator.
    ///
    /// Note that the order in items will be reduced is not specified,
    /// so if the `Ord` impl is not truly commutative and associative
    /// (as is the case for floating point numbers), then the results
    /// are not deterministic.
    fn max(self) -> Self::Item
        where Self: Sized, MaxOp: ReduceOp<Self::Item>
    {
        reduce(self, MAX)
    }

    /// Reduces the items using the given "reduce operator". You may
    /// prefer `reduce_with` for a simpler interface.
    ///
    /// Note that the order in items will be reduced is not specified,
    /// so if the `reduce_op` impl is not truly commutative and
    /// associative, then the results are not deterministic.
    fn reduce<REDUCE_OP>(self, reduce_op: &REDUCE_OP) -> Self::Item
        where Self: Sized, REDUCE_OP: ReduceOp<Self::Item>
    {
        reduce(self, reduce_op)
    }
}

impl<T: ParallelIterator> IntoParallelIterator for T {
    type Iter = T;
    type Item = T::Item;

    fn into_par_iter(self) -> T {
        self
    }
}

/// A trait for parallel iterators items where the precise number of
/// items is not known, but we can at least give an upper-bound. These
/// sorts of iterators result from filtering.
///
/// # Safety note
///
/// This trait is declared as **unsafe to implement**, but it is
/// perfectly safe to **use**. It is unsafe to implement because other
/// code relies on the fact that the estimated length is an upper
/// bound in order to guarantee safety invariants.
pub unsafe trait BoundedParallelIterator: ParallelIterator {
}

/// A trait for parallel iterators items where the precise number of
/// items is known. This occurs when e.g. iterating over a
/// vector. Knowing precisely how many items will be produced is very
/// useful.
///
/// # Safety note
///
/// This trait is declared as **unsafe to implement**, but it is
/// perfectly safe to **use**. It is unsafe to implement because other
/// code relies on the fact that the estimated length from an
/// `ExactParallelIterator` is precisely correct in order to guarantee
/// safety invariants.
pub unsafe trait ExactParallelIterator: BoundedParallelIterator {
    /// Iterate over tuples `(A, B)`, where the items `A` are from
    /// this iterator and `B` are from the iterator given as argument.
    /// Like the `zip` method on ordinary iterators, if the two
    /// iterators are of unequal length, you only get the items they
    /// have in common.
    fn zip<ZIP_OP>(self, zip_op: ZIP_OP) -> ZipIter<Self, ZIP_OP::Iter>
        where Self: Sized, ZIP_OP: IntoParallelIterator, ZIP_OP::Iter: ExactParallelIterator
    {
        ZipIter::new(self, zip_op.into_par_iter())
    }

    /// Yields an index along with each item.
    fn enumerate(self) -> Enumerate<Self>
        where Self: Sized
    {
        Enumerate::new(self)
    }

    /// Collects the results of the iterator into the specified
    /// vector. The vector is always truncated before execution
    /// begins. If possible, reusing the vector across calls can lead
    /// to better performance since it reuses the same backing buffer.
    fn collect_into(self, target: &mut Vec<Self::Item>)
        where Self: Sized
    {
        collect_into(self, target);
    }
}

<commit_msg>Add `Sized` bound to `ParallelIterator`; it was not really practical to use it as an object anyhow, since you would have had to specify your `ParallelIteratorState` precisely (and that trait is not object safe).<commit_after>#![allow(dead_code)]

//! The `ParallelIterator` module makes it easy to write parallel
//! programs using an iterator-style interface. To get access to all
//! the methods you want, the easiest is to write `use
//! rayon::par_iter::*;` at the top of your module, which will import
//! the various traits and methods you need.
//!
//! The submodules of this module mostly just contain implementaton
//! details of little interest to an end-user.

use std::f64;
use std::ops::Fn;
use self::collect::collect_into;
use self::enumerate::Enumerate;
use self::filter::Filter;
use self::filter_map::FilterMap;
use self::map::Map;
use self::reduce::{reduce, ReduceOp, SumOp, MulOp, MinOp, MaxOp, ReduceWithOp,
                   SUM, MUL, MIN, MAX};
use self::state::ParallelIteratorState;
use self::weight::Weight;
use self::zip::ZipIter;

pub mod collect;
pub mod enumerate;
pub mod filter;
pub mod filter_map;
pub mod len;
pub mod for_each;
pub mod reduce;
pub mod slice;
pub mod slice_mut;
pub mod state;
pub mod map;
pub mod weight;
pub mod zip;
pub mod range;

#[cfg(test)]
mod test;

pub trait IntoParallelIterator {
    type Iter: ParallelIterator<Item=Self::Item>;
    type Item: Send;

    fn into_par_iter(self) -> Self::Iter;
}

pub trait IntoParallelRefIterator<'data> {
    type Iter: ParallelIterator<Item=&'data Self::Item>;
    type Item: Sync + 'data;

    fn par_iter(&'data self) -> Self::Iter;
}

pub trait IntoParallelRefMutIterator<'data> {
    type Iter: ParallelIterator<Item=&'data mut Self::Item>;
    type Item: Send + 'data;

    fn par_iter_mut(&'data mut self) -> Self::Iter;
}

/// The `ParallelIterator` interface.
pub trait ParallelIterator: Sized {
    type Item: Send;
    type Shared: Sync;
    type State: ParallelIteratorState<Shared=Self::Shared, Item=Self::Item> + Send;

    fn state(self) -> (Self::Shared, Self::State);

    /// Indicates the relative "weight" of producing each item in this
    /// parallel iterator. A higher weight will cause finer-grained
    /// parallel subtasks. 1.0 indicates something very cheap and
    /// uniform, like copying a value out of an array, or computing `x
    /// + 1`. If your tasks are either very expensive, or very
    /// unpredictable, you are better off with higher values. See also
    /// `weight_max`, which is a convenient shorthand to force the
    /// finest grained parallel execution posible. Tuning this value
    /// should not affect correctness but can improve (or hurt)
    /// performance.
    fn weight(self, scale: f64) -> Weight<Self> {
        Weight::new(self, scale)
    }

    /// Shorthand for `self.weight(f64::INFINITY)`. This forces the
    /// smallest granularity of parallel execution, which makes sense
    /// when your parallel tasks are (potentially) very expensive to
    /// execute.
    fn weight_max(self) -> Weight<Self> {
        self.weight(f64::INFINITY)
    }

    /// Executes `OP` on each item produced by the iterator, in parallel.
    fn for_each<OP>(self, op: OP)
        where OP: Fn(Self::Item) + Sync
    {
        for_each::for_each(self, &op)
    }

    /// Applies `map_op` to each item of his iterator, producing a new
    /// iterator with the results.
    fn map<MAP_OP,R>(self, map_op: MAP_OP) -> Map<Self, MAP_OP>
        where MAP_OP: Fn(Self::Item) -> R
    {
        Map::new(self, map_op)
    }

    /// Applies `map_op` to each item of his iterator, producing a new
    /// iterator with the results.
    fn filter<FILTER_OP>(self, filter_op: FILTER_OP) -> Filter<Self, FILTER_OP>
        where FILTER_OP: Fn(&Self::Item) -> bool
    {
        Filter::new(self, filter_op)
    }

    /// Applies `map_op` to each item of his iterator, producing a new
    /// iterator with the results.
    fn filter_map<FILTER_OP,R>(self, filter_op: FILTER_OP) -> FilterMap<Self, FILTER_OP>
        where FILTER_OP: Fn(Self::Item) -> Option<R>
    {
        FilterMap::new(self, filter_op)
    }

    /// Reduces the items in the iterator into one item using `op`.
    /// See also `sum`, `mul`, `min`, etc, which are slightly more
    /// efficient. Returns `None` if the iterator is empty.
    ///
    /// Note: unlike in a sequential iterator, the order in which `op`
    /// will be applied to reduce the result is not specified. So `op`
    /// should be commutative and associative or else the results will
    /// be non-deterministic.
    fn reduce_with<OP>(self, op: OP) -> Option<Self::Item>
        where OP: Fn(Self::Item, Self::Item) -> Self::Item + Sync,
    {
        reduce(self.map(Some), &ReduceWithOp::new(op))
    }

    /// Sums up the items in the iterator.
    ///
    /// Note that the order in items will be reduced is not specified,
    /// so if the `+` operator is not truly commutative and
    /// associative (as is the case for floating point numbers), then
    /// the results are not fully deterministic.
    fn sum(self) -> Self::Item
        where SumOp: ReduceOp<Self::Item>
    {
        reduce(self, SUM)
    }

    /// Multiplies all the items in the iterator.
    ///
    /// Note that the order in items will be reduced is not specified,
    /// so if the `*` operator is not truly commutative and
    /// associative (as is the case for floating point numbers), then
    /// the results are not fully deterministic.
    fn mul(self) -> Self::Item
        where MulOp: ReduceOp<Self::Item>
    {
        reduce(self, MUL)
    }

    /// Computes the minimum of all the items in the iterator.
    ///
    /// Note that the order in items will be reduced is not specified,
    /// so if the `Ord` impl is not truly commutative and associative
    /// (as is the case for floating point numbers), then the results
    /// are not deterministic.
    fn min(self) -> Self::Item
        where MinOp: ReduceOp<Self::Item>
    {
        reduce(self, MIN)
    }

    /// Computes the maximum of all the items in the iterator.
    ///
    /// Note that the order in items will be reduced is not specified,
    /// so if the `Ord` impl is not truly commutative and associative
    /// (as is the case for floating point numbers), then the results
    /// are not deterministic.
    fn max(self) -> Self::Item
        where MaxOp: ReduceOp<Self::Item>
    {
        reduce(self, MAX)
    }

    /// Reduces the items using the given "reduce operator". You may
    /// prefer `reduce_with` for a simpler interface.
    ///
    /// Note that the order in items will be reduced is not specified,
    /// so if the `reduce_op` impl is not truly commutative and
    /// associative, then the results are not deterministic.
    fn reduce<REDUCE_OP>(self, reduce_op: &REDUCE_OP) -> Self::Item
        where REDUCE_OP: ReduceOp<Self::Item>
    {
        reduce(self, reduce_op)
    }
}

impl<T: ParallelIterator> IntoParallelIterator for T {
    type Iter = T;
    type Item = T::Item;

    fn into_par_iter(self) -> T {
        self
    }
}

/// A trait for parallel iterators items where the precise number of
/// items is not known, but we can at least give an upper-bound. These
/// sorts of iterators result from filtering.
///
/// # Safety note
///
/// This trait is declared as **unsafe to implement**, but it is
/// perfectly safe to **use**. It is unsafe to implement because other
/// code relies on the fact that the estimated length is an upper
/// bound in order to guarantee safety invariants.
pub unsafe trait BoundedParallelIterator: ParallelIterator {
}

/// A trait for parallel iterators items where the precise number of
/// items is known. This occurs when e.g. iterating over a
/// vector. Knowing precisely how many items will be produced is very
/// useful.
///
/// # Safety note
///
/// This trait is declared as **unsafe to implement**, but it is
/// perfectly safe to **use**. It is unsafe to implement because other
/// code relies on the fact that the estimated length from an
/// `ExactParallelIterator` is precisely correct in order to guarantee
/// safety invariants.
pub unsafe trait ExactParallelIterator: BoundedParallelIterator {
    /// Iterate over tuples `(A, B)`, where the items `A` are from
    /// this iterator and `B` are from the iterator given as argument.
    /// Like the `zip` method on ordinary iterators, if the two
    /// iterators are of unequal length, you only get the items they
    /// have in common.
    fn zip<ZIP_OP>(self, zip_op: ZIP_OP) -> ZipIter<Self, ZIP_OP::Iter>
        where ZIP_OP: IntoParallelIterator, ZIP_OP::Iter: ExactParallelIterator
    {
        ZipIter::new(self, zip_op.into_par_iter())
    }

    /// Yields an index along with each item.
    fn enumerate(self) -> Enumerate<Self> {
        Enumerate::new(self)
    }

    /// Collects the results of the iterator into the specified
    /// vector. The vector is always truncated before execution
    /// begins. If possible, reusing the vector across calls can lead
    /// to better performance since it reuses the same backing buffer.
    fn collect_into(self, target: &mut Vec<Self::Item>) {
        collect_into(self, target);
    }
}

<|endoftext|>


<commit_before>use std::collections::{BTreeMap, HashMap, HashSet};

use mapping::{Mapping, FieldMapping};
use document::Document;

#[derive(Debug)]
pub struct Index {
    pub mappings: HashMap<String, Mapping>,
    pub docs: BTreeMap<u64, Document>,
    pub aliases: HashSet<String>,
    next_doc_num: u64,
    doc_id_map: HashMap<String, u64>,
}

impl Index {
    pub fn new() -> Index {
        Index {
            mappings: HashMap::new(),
            docs: BTreeMap::new(),
            aliases: HashSet::new(),
            next_doc_num: 1,
            doc_id_map: HashMap::new(),
        }
    }

    pub fn get_mapping_by_name(&self, name: &str) -> Option<&Mapping> {
        self.mappings.get(name)
    }

    pub fn get_field_mapping_by_name(&self, name: &str) -> Option<&FieldMapping> {
        for mapping in self.mappings.values() {
            if let Some(ref field_mapping) = mapping.fields.get(name) {
                return Some(field_mapping);
            }
        }

        None
    }

    pub fn get_document_by_id(&self, id: &str) -> Option<&Document> {
        match self.doc_id_map.get(id) {
            Some(doc_num) => self.docs.get(doc_num),
            None => None,
        }
    }

    pub fn contains_document_id(&self, id: &str) -> bool {
        self.doc_id_map.contains_key(id)
    }

    pub fn remove_document_by_id(&mut self, id: &str) -> bool {
        match self.doc_id_map.remove(id) {
            Some(doc_num) => {
                self.docs.remove(&doc_num);

                true
            }
            None => false
        }
    }

    pub fn insert_or_update_document(&mut self, doc: Document) {
        let doc_num = self.next_doc_num;
        self.next_doc_num += 1;

        self.doc_id_map.insert(doc.id.clone(), doc_num);
        self.docs.insert(doc_num, doc);
    }

    pub fn initialise(&mut self) {}
}
<commit_msg>Added basic inverted index<commit_after>use std::collections::{BTreeMap, HashMap, HashSet};

use term::Term;
use mapping::{Mapping, FieldMapping};
use document::Document;

#[derive(Debug)]
pub struct Index {
    pub mappings: HashMap<String, Mapping>,
    pub docs: BTreeMap<u64, Document>,
    pub index: BTreeMap<Term, BTreeMap<String, Vec<(u64, u32)>>>,
    pub aliases: HashSet<String>,
    next_doc_num: u64,
    doc_id_map: HashMap<String, u64>,
}

impl Index {
    pub fn new() -> Index {
        Index {
            mappings: HashMap::new(),
            docs: BTreeMap::new(),
            index: BTreeMap::new(),
            aliases: HashSet::new(),
            next_doc_num: 1,
            doc_id_map: HashMap::new(),
        }
    }

    pub fn get_mapping_by_name(&self, name: &str) -> Option<&Mapping> {
        self.mappings.get(name)
    }

    pub fn get_field_mapping_by_name(&self, name: &str) -> Option<&FieldMapping> {
        for mapping in self.mappings.values() {
            if let Some(ref field_mapping) = mapping.fields.get(name) {
                return Some(field_mapping);
            }
        }

        None
    }

    pub fn get_document_by_id(&self, id: &str) -> Option<&Document> {
        match self.doc_id_map.get(id) {
            Some(doc_num) => self.docs.get(doc_num),
            None => None,
        }
    }

    pub fn contains_document_id(&self, id: &str) -> bool {
        self.doc_id_map.contains_key(id)
    }

    pub fn remove_document_by_id(&mut self, id: &str) -> bool {
        match self.doc_id_map.remove(id) {
            Some(doc_num) => {
                self.docs.remove(&doc_num);

                true
            }
            None => false
        }
    }

    pub fn insert_or_update_document(&mut self, doc: Document) {
        let doc_num = self.next_doc_num;
        self.next_doc_num += 1;

        // Put field contents in inverted index
        for (field_name, tokens) in doc.fields.iter() {
            let mut position: u32 = 1;

            for token in tokens.iter() {
                if !self.index.contains_key(&token.term) {
                    self.index.insert(token.term.clone(), BTreeMap::new());
                }

                let mut index_fields = self.index.get_mut(&token.term).unwrap();

                if !index_fields.contains_key(field_name) {
                    index_fields.insert(field_name.clone(), Vec::new());
                }

                let mut postings_list = index_fields.get_mut(field_name).unwrap();
                postings_list.push((doc_num, token.position));
            }
        }

        self.doc_id_map.insert(doc.id.clone(), doc_num);
        self.docs.insert(doc_num, doc);
    }

    pub fn initialise(&mut self) {}
}
<|endoftext|>


<commit_before>use fragment::matching;
use ExclusionPattern;
use walkdir::{DirEntry, Error, WalkDir};
use std::path::PathBuf;
use std::clone::Clone;

pub struct Index {
    path: PathBuf,
    entries: Vec<IndexedPath>,
}

#[derive(Debug, Eq, Ord, PartialEq, PartialOrd)]
pub struct IndexedPath(PathBuf);

impl ToString for IndexedPath {
    fn to_string(&self) -> String {
        self.0.to_string_lossy().into_owned()
    }
}

impl Clone for IndexedPath {
    fn clone(&self) -> Self {
        IndexedPath(self.0.clone())
    }

    fn clone_from(&mut self, source: &Self) {
        self.0 = source.0.clone()
    }
}

impl Index {
    pub fn new(path: PathBuf) -> Index {
        Index {
            path: path,
            entries: Vec::new()
        }
    }

    /// Finds all files inside and beneath the index path
    /// and adds them to the index entries vector.
    pub fn populate(&mut self, exclusions: Option<Vec<ExclusionPattern>>, case_sensitive: bool) {
        // The entries listed by read_dir include the root index path; we want
        // relative paths, so we get this length so that we can strip it.
        let prefix_length = match self.path.to_str() {
            Some(path) => path.len() + 1,
            None => return,
        };

        // Start indexing at the specified path.
        let filtered_entries = WalkDir::new(&self.path).into_iter().filter_entry(|entry| {
            if let Some(ref exclusions) = exclusions {
                !exclusions.iter().any(|exclusion| {
                    exclusion.matches(entry.path().to_string_lossy().as_ref())
                })
            } else {
                true
            }
        });

        for entry in filtered_entries {
            relative_entry_path(entry, prefix_length).map(|entry_path| {
                self.entries.push(
                    if case_sensitive {
                        IndexedPath(PathBuf::from(entry_path))
                    } else {
                        IndexedPath(PathBuf::from(entry_path.to_lowercase()))
                    }
                );
            });
        }
    }

    pub fn find(&self, term: &str, limit: usize) -> Vec<PathBuf> {
        matching::find(term, &self.entries, limit, true)
            .into_iter()
            .map(|r| {
                let IndexedPath(other_thing) = r.clone();
                other_thing
            })
            .collect()
    }
}

/// Transforms a DirEntry object into an optional relative path string,
/// returning None if any errors occur or if the entry is not a file.
fn relative_entry_path(entry: Result<DirEntry, Error>, prefix_length: usize) -> Option<String> {
    entry.ok().and_then(|e| {
        // Limit path look-ups to files.
        e.path().metadata().ok().and_then(|metadata| {
            if metadata.is_file() {
                // Map the absolute path to a relative one.
                e.path().to_str().map(|path| path[prefix_length..].to_string())
            } else {
                None
            }
        })
    })
}

#[cfg(test)]
mod tests {
    extern crate fragment;

    use super::{Index, IndexedPath, ExclusionPattern};
    use std::path::PathBuf;

    #[test]
    fn populate_adds_all_files_to_entries() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath(PathBuf::from("directory/Capitalized_file")),
                                    IndexedPath(PathBuf::from("directory/nested_file")),
                                    IndexedPath(PathBuf::from("root_file"))];
        index.populate(None, true);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_respects_exclusions() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath(PathBuf::from("root_file"))];
        index.populate(Some(vec![ExclusionPattern::new("**/directory").unwrap()]), true);

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_lowercases_entries_when_case_sensitive_is_false() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath(PathBuf::from("directory/capitalized_file")),
                                    IndexedPath(PathBuf::from("directory/nested_file")),
                                    IndexedPath(PathBuf::from("root_file"))];
        index.populate(None, false);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_lowercases_entries_when_case_sensitive_is_true() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath(PathBuf::from("directory/Capitalized_file")),
                                    IndexedPath(PathBuf::from("directory/nested_file")),
                                    IndexedPath(PathBuf::from("root_file"))];
        index.populate(None, true);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn find_defers_to_matching_module() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        index.populate(None, true);
        let term = "root";
        let limit = 5;

        // Get a string version of the results (PathBuf doesn't implement the display trait).
        let results: Vec<String> = index.find(term, limit)
                                        .iter()
                                        .map(|r| r.to_string_lossy().into_owned())
                                        .collect();

        assert_eq!(results, vec!["root_file".to_string()]);
    }
}
<commit_msg>Add a test ensuring that find always opts for case sensitivity<commit_after>use fragment::matching;
use ExclusionPattern;
use walkdir::{DirEntry, Error, WalkDir};
use std::path::PathBuf;
use std::clone::Clone;

pub struct Index {
    path: PathBuf,
    entries: Vec<IndexedPath>,
}

#[derive(Debug, Eq, Ord, PartialEq, PartialOrd)]
pub struct IndexedPath(PathBuf);

impl ToString for IndexedPath {
    fn to_string(&self) -> String {
        self.0.to_string_lossy().into_owned()
    }
}

impl Clone for IndexedPath {
    fn clone(&self) -> Self {
        IndexedPath(self.0.clone())
    }

    fn clone_from(&mut self, source: &Self) {
        self.0 = source.0.clone()
    }
}

impl Index {
    pub fn new(path: PathBuf) -> Index {
        Index {
            path: path,
            entries: Vec::new()
        }
    }

    /// Finds all files inside and beneath the index path
    /// and adds them to the index entries vector.
    pub fn populate(&mut self, exclusions: Option<Vec<ExclusionPattern>>, case_sensitive: bool) {
        // The entries listed by read_dir include the root index path; we want
        // relative paths, so we get this length so that we can strip it.
        let prefix_length = match self.path.to_str() {
            Some(path) => path.len() + 1,
            None => return,
        };

        // Start indexing at the specified path.
        let filtered_entries = WalkDir::new(&self.path).into_iter().filter_entry(|entry| {
            if let Some(ref exclusions) = exclusions {
                !exclusions.iter().any(|exclusion| {
                    exclusion.matches(entry.path().to_string_lossy().as_ref())
                })
            } else {
                true
            }
        });

        for entry in filtered_entries {
            relative_entry_path(entry, prefix_length).map(|entry_path| {
                self.entries.push(
                    if case_sensitive {
                        IndexedPath(PathBuf::from(entry_path))
                    } else {
                        IndexedPath(PathBuf::from(entry_path.to_lowercase()))
                    }
                );
            });
        }
    }

    pub fn find(&self, term: &str, limit: usize) -> Vec<PathBuf> {
        matching::find(term, &self.entries, limit, true)
            .into_iter()
            .map(|r| {
                let IndexedPath(other_thing) = r.clone();
                other_thing
            })
            .collect()
    }
}

/// Transforms a DirEntry object into an optional relative path string,
/// returning None if any errors occur or if the entry is not a file.
fn relative_entry_path(entry: Result<DirEntry, Error>, prefix_length: usize) -> Option<String> {
    entry.ok().and_then(|e| {
        // Limit path look-ups to files.
        e.path().metadata().ok().and_then(|metadata| {
            if metadata.is_file() {
                // Map the absolute path to a relative one.
                e.path().to_str().map(|path| path[prefix_length..].to_string())
            } else {
                None
            }
        })
    })
}

#[cfg(test)]
mod tests {
    extern crate fragment;

    use super::{Index, IndexedPath, ExclusionPattern};
    use std::path::PathBuf;

    #[test]
    fn populate_adds_all_files_to_entries() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath(PathBuf::from("directory/Capitalized_file")),
                                    IndexedPath(PathBuf::from("directory/nested_file")),
                                    IndexedPath(PathBuf::from("root_file"))];
        index.populate(None, true);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_respects_exclusions() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath(PathBuf::from("root_file"))];
        index.populate(Some(vec![ExclusionPattern::new("**/directory").unwrap()]), true);

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_lowercases_entries_when_case_sensitive_is_false() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath(PathBuf::from("directory/capitalized_file")),
                                    IndexedPath(PathBuf::from("directory/nested_file")),
                                    IndexedPath(PathBuf::from("root_file"))];
        index.populate(None, false);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn populate_lowercases_entries_when_case_sensitive_is_true() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        let expected_entries = vec![IndexedPath(PathBuf::from("directory/Capitalized_file")),
                                    IndexedPath(PathBuf::from("directory/nested_file")),
                                    IndexedPath(PathBuf::from("root_file"))];
        index.populate(None, true);
        index.entries.sort();

        assert_eq!(index.entries, expected_entries);
    }

    #[test]
    fn find_defers_to_matching_module() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        index.populate(None, true);
        let term = "root";
        let limit = 5;

        // Get a string version of the results (PathBuf doesn't implement the display trait).
        let results: Vec<String> = index.find(term, limit)
                                        .iter()
                                        .map(|r| r.to_string_lossy().into_owned())
                                        .collect();

        assert_eq!(results, vec!["root_file".to_string()]);
    }

    #[test]
    fn find_always_uses_case_sensitive_queries() {
        let path = PathBuf::from("tests/sample");
        let mut index = Index::new(path);
        index.populate(None, true);
        let term = "Root";
        let limit = 5;

        // Get a string version of the results (PathBuf doesn't implement the display trait).
        let results: Vec<String> = index.find(term, limit)
                                        .iter()
                                        .map(|r| r.to_string_lossy().into_owned())
                                        .collect();

        assert!(results.is_empty());
    }
}
<|endoftext|>