<commit_before>#![feature(zero_one)]

extern crate deeplearn;
extern crate gpuarray as ga;

use std::fs::File;
use std::io::{
    self,
    BufReader,
    Read,
};
use std::num::{Zero, One};
use std::path::Path;
use std::rc::Rc;

use deeplearn::Graph;
use deeplearn::op::{Add, MatMul, Mse, Relu};
use deeplearn::init;
use ga::Array;

fn main() {
    // Training data
    println!("Reading training labels...");
    let train_labels = read_mnist_labels("data/mnist/train-labels-idx1-ubyte", None).unwrap();
    let train_labels_logits: Vec<Array<f32>> = train_labels.iter().cloned().map(|x| one_hot(x, 10)).collect();
    println!("Label count: {}", train_labels.len());

    println!("Reading training images...");
    let (rows, columns, train_images) = read_mnist_images("data/mnist/train-images-idx3-ubyte", None).unwrap();

    // Validation data
    println!("Reading validation labels...");
    let val_labels = read_mnist_labels("data/mnist/t10k-labels-idx1-ubyte", Some(1000)).unwrap();
    let val_labels_logits: Vec<Array<f32>> = val_labels.iter().cloned().map(|x| one_hot(x, 10)).collect();
    println!("Label count: {}", train_labels.len());

    println!("Reading validation images...");
    let (_, _, val_images) = read_mnist_images("data/mnist/t10k-images-idx3-ubyte", Some(1000)).unwrap();

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

    // Setup the graph. It's going to have 2 inputs, 3 hidden nodes in it's 1 hidden layer, and 1
    // output node
    let mut graph = Graph::new(ctx.clone());

    //////////////////////////
    // Layer 1

    // Input.  1 batches of rows*columns inputs each.
    let input = graph.add_variable((1, rows*columns), vec![0.0; rows*columns]);
    // Weights for layer 1: [rows*columns inputs x 300 nodes]
    let l1_w = graph.add_variable((rows*columns, 300), init::Normal(0.001, 0.005));
    // Use matrix multiplication to do a fully connected layer
    let l1_mat_mul = graph.add_node(MatMul(input, l1_w));
    // Grab VarIndex for l1_mat_mul's output
    let l1_mat_mul_out = l1_mat_mul.get(&graph).outputs[0];
    // 300 biases; one for each node in layer 1
    let l1_b = graph.add_variable((1, 300), init::Normal(0.001, 0.005));
    // Here we add the biases to the matrix multiplication output
    let l1_biased = graph.add_node(Add(l1_mat_mul_out, l1_b, 0));
    // Grab VarIndex for l1_biased's output
    let l1_biased_out = l1_biased.get(&graph).outputs[0];
    // Run the biased input*weight sums through an ReLU activation
    let l1_relu = graph.add_node(Relu(l1_biased_out));
    let l1_relu_out = l1_relu.get(&graph).outputs[0];

    //////////////////////////
    // Layer 2

    // Weights for layer 2: [300 inputs x 10 nodes]
    let l2_w = graph.add_variable((300, 10), init::Normal(0.001, 0.005));
    // Fully connected layer 2. Use layer 1's output as layer 2's input
    let l2_mat_mul = graph.add_node(MatMul(l1_relu_out, l2_w));
    // Grab VarIndex for l2_mat_mul's output
    let l2_mat_mul_out = l2_mat_mul.get(&graph).outputs[0];
    // 10 bias for 1 output node
    let l2_b = graph.add_variable((1, 10), init::Normal(0.001, 0.005));
    // Here we add the bias to the matrix multiplication output
    let l2_biased = graph.add_node(Add(l2_mat_mul_out, l2_b, 0));
    // Grab VarIndex for l2_biased's output
    let l2_biased_out = l2_biased.get(&graph).outputs[0];
    // Run the biased input*weight sums through an ReLU activation
    let l2_relu = graph.add_node(Relu(l2_biased_out));
    // Grab VarIndex for l2_relu's output
    let l2_relu_out = l2_relu.get(&graph).outputs[0];

    //////////////////////////
    // Loss

    // Add a variable for the training outputs: [1 batches x 10 output]
    let train_out = graph.add_variable((1, 10), vec![0.0; 10]);
    // Use mean squared error loss function
    let loss = graph.add_node(Mse(l2_relu_out, train_out));
    // Grab the VarIndex for loss's output
    let loss_out = loss.get(&graph).outputs[0];
    // Create a gradient to apply to the loss function
    let loss_d = graph.add_gradient(loss, 0);

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

    // We apply a gradient of -0.001 to the loss function
    let loss_d_cpu = Array::from_vec(vec![1, 10], vec![-0.001; 10]);
    loss_d.get(&graph).set(&ctx, &loss_d_cpu);

    /////////////////////////
    // Train the network
    let mut correct = 0;
    for epoch in 0..60000 {
        // Upload training data
        let train_sample = epoch%train_images.len();
        input.get(&graph).set(&ctx, &train_images[train_sample]);
        train_out.get(&graph).set(&ctx, &train_labels_logits[train_sample]);

        graph.run();
        let out = l2_relu_out.get(&graph).get(&ctx);
        let l1_w_d = graph.get_input_gradient(l1_w).unwrap().get(&graph);
        let l2_w_d = graph.get_input_gradient(l2_w).unwrap().get(&graph);
        let l1_b_d = graph.get_input_gradient(l1_b).unwrap().get(&graph);
        let l2_b_d = graph.get_input_gradient(l2_b).unwrap().get(&graph);
        ga::add(&ctx, &l1_w.get(&graph), -1, &l1_w_d, &l1_w.get(&graph));
        ga::add(&ctx, &l2_w.get(&graph), -1, &l2_w_d, &l2_w.get(&graph));
        ga::add(&ctx, &l1_b.get(&graph), -1, &l1_b_d, &l1_b.get(&graph));
        ga::add(&ctx, &l2_b.get(&graph), -1, &l2_b_d, &l2_b.get(&graph));

        let (mut max_index, mut max_value) = (0, out[&[0, 0]]);
        for (i, val) in out.buffer().iter().enumerate() {
            if *val > max_value {
                max_index = i;
                max_value = *val;
            }
        }

        if max_index == train_labels[train_sample] as usize {
            correct += 1;
        }

        if epoch % 1000 == 999 {
            let out_d = l2_relu.get(&graph)
                               .out_grad[0]
                               .gradient()
                               .get(&graph)
                               .get(&ctx);
            let l = loss_out.get(&graph).get(&ctx);
            println!("===================");
            println!("Epoch: {}", epoch);
            println!("out = {:?}", out);
            println!("out_d = {:?}", out_d);
            println!("loss = {:?}", l);
            println!("Accuracy: {}%", (correct as f32)/(1000 as f32) * 100.0);
            correct = 0;
        }
    }

    /////////////////////////
    // Validate the network
    println!("#######################################");
    println!("Validating");
    let mut correct = 0;
    for epoch in 0..1000 {
        // Upload training data
        let train_sample = epoch%val_images.len();
        input.get(&graph).set(&ctx, &val_images[train_sample]);
        train_out.get(&graph).set(&ctx, &val_labels_logits[train_sample]);

        graph.run();
        let out = l2_relu_out.get(&graph).get(&ctx);

        let (mut max_index, mut max_value) = (0, out[&[0, 0]]);
        for (i, val) in out.buffer().iter().enumerate() {
            if *val > max_value {
                max_index = i;
                max_value = *val;
            }
        }

        if max_index == val_labels[train_sample] as usize {
            correct += 1;
        }

        if epoch % 1000 == 999 {
            let out_d = l2_relu.get(&graph)
                               .out_grad[0]
                               .gradient()
                               .get(&graph)
                               .get(&ctx);
            let l = loss_out.get(&graph).get(&ctx);
            println!("===================");
            println!("Epoch: {}", epoch);
            println!("out = {:?}", out);
            println!("out_d = {:?}", out_d);
            println!("loss = {:?}", l);
            println!("Accuracy: {}%", (correct as f32)/(1000 as f32) * 100.0);
            correct = 0;
        }
    }
}

fn read_mnist_labels<P: AsRef<Path>>(path: P, num_samples: Option<usize>) -> io::Result<Vec<u8>> {
    use std::io::{Error, ErrorKind};

    let ref mut file = BufReader::new(File::open(path).unwrap());

    let magic = u32::from_be(try!(read_u32(file)));
    if magic != 2049 {
        return Err(Error::new(ErrorKind::Other,
                              format!("Invalid magic number. Got expect 2049, got {}",
                                      magic).as_ref()))
    }

    let label_count = u32::from_be(try!(read_u32(file)));

    let mut labels = Vec::with_capacity(label_count as usize);
    for i in 0..label_count {
        if let Some(num_samples) = num_samples {
            if i as usize > num_samples {
                break;
            }
        }
        labels.push(try!(read_u8(file)));
    }

    Ok(labels)
}

fn read_mnist_images<P: AsRef<Path>>(path: P, num_samples: Option<usize>) -> io::Result<(usize, usize, Vec<Array<f32>>)> {
    use std::io::{Error, ErrorKind};

    let ref mut file = BufReader::new(File::open(path).unwrap());

    let magic = u32::from_be(try!(read_u32(file)));
    if magic != 2051 {
        return Err(Error::new(ErrorKind::Other,
                              format!("Invalid magic number. Got expect 2051, got {}",
                                      magic).as_ref()))
    }

    let image_count = u32::from_be(try!(read_u32(file))) as usize;
    let rows = u32::from_be(try!(read_u32(file))) as usize;
    let columns = u32::from_be(try!(read_u32(file))) as usize;

    let mut images = Vec::with_capacity(image_count);
    for i in 0..image_count {
        if let Some(num_samples) = num_samples {
            if i as usize > num_samples {
                break;
            }
        }
        let mut pixel_buf = vec![0u8; rows*columns];
        try!(file.read_exact(pixel_buf.as_mut()));
        let array = Array::from_vec(vec![rows, columns],
                                    pixel_buf.into_iter().map(|x| (x as f32)/255.0).collect());
        images.push(array);
    }

    Ok((rows, columns, images))
}

fn one_hot<N, M>(label: N, classes: N) -> Array<M>
    where usize: From<N>,
          M:     ga::num::Num+Zero+One,
{
    let classes: usize = From::from(classes); // Cast class count to usize
    let label: usize = From::from(label); // Cast label to usize
    let mut buf: Vec<M> = vec![Zero::zero(); classes]; // Create array of zeroes
    buf[label] = One::one(); // Set the one-hot component
    Array::from_vec(vec![1, buf.len()], buf) // Construct the array
}

fn read_u8<T: Read>(reader: &mut T) -> io::Result<u8> {
    use std::mem;

    let mut buf: [u8; 1] = [0];
    reader.read_exact(&mut buf).map(|_| {
        let data: u8 = unsafe { mem::transmute(buf) };
        data
    })
}

fn read_u32<T: Read>(reader: &mut T) -> io::Result<u32> {
    use std::mem;

    let mut buf: [u8; 4] = [0, 0, 0, 0];
    reader.read_exact(&mut buf).map(|_| {
        let data: u32 = unsafe { mem::transmute(buf) };
        data
    })
}
<commit_msg>Don't keep reallocating an output array in mnist example<commit_after>#![feature(zero_one)]

extern crate deeplearn;
extern crate gpuarray as ga;

use std::fs::File;
use std::io::{
    self,
    BufReader,
    Read,
};
use std::num::{Zero, One};
use std::path::Path;
use std::rc::Rc;

use deeplearn::Graph;
use deeplearn::op::{Add, MatMul, Mse, Relu};
use deeplearn::init;
use ga::Array;

fn main() {
    // Training data
    println!("Reading training labels...");
    let train_labels = read_mnist_labels("data/mnist/train-labels-idx1-ubyte", None).unwrap();
    let train_labels_logits: Vec<Array<f32>> = train_labels.iter().cloned().map(|x| one_hot(x, 10)).collect();
    println!("Label count: {}", train_labels.len());

    println!("Reading training images...");
    let (rows, columns, train_images) = read_mnist_images("data/mnist/train-images-idx3-ubyte", None).unwrap();

    // Validation data
    println!("Reading validation labels...");
    let val_labels = read_mnist_labels("data/mnist/t10k-labels-idx1-ubyte", Some(1000)).unwrap();
    let val_labels_logits: Vec<Array<f32>> = val_labels.iter().cloned().map(|x| one_hot(x, 10)).collect();
    println!("Label count: {}", train_labels.len());

    println!("Reading validation images...");
    let (_, _, val_images) = read_mnist_images("data/mnist/t10k-images-idx3-ubyte", Some(1000)).unwrap();

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

    // Setup the graph. It's going to have 2 inputs, 3 hidden nodes in it's 1 hidden layer, and 1
    // output node
    let mut graph = Graph::new(ctx.clone());

    //////////////////////////
    // Layer 1

    // Input.  1 batches of rows*columns inputs each.
    let input = graph.add_variable((1, rows*columns), vec![0.0; rows*columns]);
    // Weights for layer 1: [rows*columns inputs x 300 nodes]
    let l1_w = graph.add_variable((rows*columns, 300), init::Normal(0.001, 0.005));
    // Use matrix multiplication to do a fully connected layer
    let l1_mat_mul = graph.add_node(MatMul(input, l1_w));
    // Grab VarIndex for l1_mat_mul's output
    let l1_mat_mul_out = l1_mat_mul.get(&graph).outputs[0];
    // 300 biases; one for each node in layer 1
    let l1_b = graph.add_variable((1, 300), init::Normal(0.001, 0.005));
    // Here we add the biases to the matrix multiplication output
    let l1_biased = graph.add_node(Add(l1_mat_mul_out, l1_b, 0));
    // Grab VarIndex for l1_biased's output
    let l1_biased_out = l1_biased.get(&graph).outputs[0];
    // Run the biased input*weight sums through an ReLU activation
    let l1_relu = graph.add_node(Relu(l1_biased_out));
    let l1_relu_out = l1_relu.get(&graph).outputs[0];

    //////////////////////////
    // Layer 2

    // Weights for layer 2: [300 inputs x 10 nodes]
    let l2_w = graph.add_variable((300, 10), init::Normal(0.001, 0.005));
    // Fully connected layer 2. Use layer 1's output as layer 2's input
    let l2_mat_mul = graph.add_node(MatMul(l1_relu_out, l2_w));
    // Grab VarIndex for l2_mat_mul's output
    let l2_mat_mul_out = l2_mat_mul.get(&graph).outputs[0];
    // 10 bias for 1 output node
    let l2_b = graph.add_variable((1, 10), init::Normal(0.001, 0.005));
    // Here we add the bias to the matrix multiplication output
    let l2_biased = graph.add_node(Add(l2_mat_mul_out, l2_b, 0));
    // Grab VarIndex for l2_biased's output
    let l2_biased_out = l2_biased.get(&graph).outputs[0];
    // Run the biased input*weight sums through an ReLU activation
    let l2_relu = graph.add_node(Relu(l2_biased_out));
    // Grab VarIndex for l2_relu's output
    let l2_relu_out = l2_relu.get(&graph).outputs[0];

    //////////////////////////
    // Loss

    // Add a variable for the training outputs: [1 batches x 10 output]
    let train_out = graph.add_variable((1, 10), vec![0.0; 10]);
    // Use mean squared error loss function
    let loss = graph.add_node(Mse(l2_relu_out, train_out));
    // Grab the VarIndex for loss's output
    let loss_out = loss.get(&graph).outputs[0];
    // Create a gradient to apply to the loss function
    let loss_d = graph.add_gradient(loss, 0);

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

    // We apply a gradient of -0.001 to the loss function
    let loss_d_cpu = Array::from_vec(vec![1, 10], vec![-0.001; 10]);
    loss_d.get(&graph).set(&ctx, &loss_d_cpu);

    /////////////////////////
    // Train the network
    let mut correct = 0;
    let mut out = Array::new(vec![1, 10], 0.0);
    for epoch in 0..60000 {
        // Upload training data
        let train_sample = epoch%train_images.len();
        input.get(&graph).set(&ctx, &train_images[train_sample]);
        train_out.get(&graph).set(&ctx, &train_labels_logits[train_sample]);

        graph.run();
        l2_relu_out.get(&graph).read(&ctx, &mut out);
        let l1_w_d = graph.get_input_gradient(l1_w).unwrap().get(&graph);
        let l2_w_d = graph.get_input_gradient(l2_w).unwrap().get(&graph);
        let l1_b_d = graph.get_input_gradient(l1_b).unwrap().get(&graph);
        let l2_b_d = graph.get_input_gradient(l2_b).unwrap().get(&graph);
        ga::add(&ctx, &l1_w.get(&graph), -1, &l1_w_d, &l1_w.get(&graph));
        ga::add(&ctx, &l2_w.get(&graph), -1, &l2_w_d, &l2_w.get(&graph));
        ga::add(&ctx, &l1_b.get(&graph), -1, &l1_b_d, &l1_b.get(&graph));
        ga::add(&ctx, &l2_b.get(&graph), -1, &l2_b_d, &l2_b.get(&graph));

        let (mut max_index, mut max_value) = (0, out[&[0, 0]]);
        for (i, val) in out.buffer().iter().enumerate() {
            if *val > max_value {
                max_index = i;
                max_value = *val;
            }
        }

        if max_index == train_labels[train_sample] as usize {
            correct += 1;
        }

        if epoch % 1000 == 999 {
            let out_d = l2_relu.get(&graph)
                               .out_grad[0]
                               .gradient()
                               .get(&graph)
                               .get(&ctx);
            let l = loss_out.get(&graph).get(&ctx);
            println!("===================");
            println!("Epoch: {}", epoch);
            println!("out = {:?}", out);
            println!("out_d = {:?}", out_d);
            println!("loss = {:?}", l);
            println!("Accuracy: {}%", (correct as f32)/(1000 as f32) * 100.0);
            correct = 0;
        }
    }

    /////////////////////////
    // Validate the network
    println!("#######################################");
    println!("Validating");
    let mut correct = 0;
    for epoch in 0..1000 {
        // Upload training data
        let train_sample = epoch%val_images.len();
        input.get(&graph).set(&ctx, &val_images[train_sample]);
        train_out.get(&graph).set(&ctx, &val_labels_logits[train_sample]);

        graph.run();
        let out = l2_relu_out.get(&graph).get(&ctx);

        let (mut max_index, mut max_value) = (0, out[&[0, 0]]);
        for (i, val) in out.buffer().iter().enumerate() {
            if *val > max_value {
                max_index = i;
                max_value = *val;
            }
        }

        if max_index == val_labels[train_sample] as usize {
            correct += 1;
        }

        if epoch % 1000 == 999 {
            let out_d = l2_relu.get(&graph)
                               .out_grad[0]
                               .gradient()
                               .get(&graph)
                               .get(&ctx);
            let l = loss_out.get(&graph).get(&ctx);
            println!("===================");
            println!("Epoch: {}", epoch);
            println!("out = {:?}", out);
            println!("out_d = {:?}", out_d);
            println!("loss = {:?}", l);
            println!("Accuracy: {}%", (correct as f32)/(1000 as f32) * 100.0);
            correct = 0;
        }
    }
}

fn read_mnist_labels<P: AsRef<Path>>(path: P, num_samples: Option<usize>) -> io::Result<Vec<u8>> {
    use std::io::{Error, ErrorKind};

    let ref mut file = BufReader::new(File::open(path).unwrap());

    let magic = u32::from_be(try!(read_u32(file)));
    if magic != 2049 {
        return Err(Error::new(ErrorKind::Other,
                              format!("Invalid magic number. Got expect 2049, got {}",
                                      magic).as_ref()))
    }

    let label_count = u32::from_be(try!(read_u32(file)));

    let mut labels = Vec::with_capacity(label_count as usize);
    for i in 0..label_count {
        if let Some(num_samples) = num_samples {
            if i as usize > num_samples {
                break;
            }
        }
        labels.push(try!(read_u8(file)));
    }

    Ok(labels)
}

fn read_mnist_images<P: AsRef<Path>>(path: P, num_samples: Option<usize>) -> io::Result<(usize, usize, Vec<Array<f32>>)> {
    use std::io::{Error, ErrorKind};

    let ref mut file = BufReader::new(File::open(path).unwrap());

    let magic = u32::from_be(try!(read_u32(file)));
    if magic != 2051 {
        return Err(Error::new(ErrorKind::Other,
                              format!("Invalid magic number. Got expect 2051, got {}",
                                      magic).as_ref()))
    }

    let image_count = u32::from_be(try!(read_u32(file))) as usize;
    let rows = u32::from_be(try!(read_u32(file))) as usize;
    let columns = u32::from_be(try!(read_u32(file))) as usize;

    let mut images = Vec::with_capacity(image_count);
    for i in 0..image_count {
        if let Some(num_samples) = num_samples {
            if i as usize > num_samples {
                break;
            }
        }
        let mut pixel_buf = vec![0u8; rows*columns];
        try!(file.read_exact(pixel_buf.as_mut()));
        let array = Array::from_vec(vec![rows, columns],
                                    pixel_buf.into_iter().map(|x| (x as f32)/255.0).collect());
        images.push(array);
    }

    Ok((rows, columns, images))
}

fn one_hot<N, M>(label: N, classes: N) -> Array<M>
    where usize: From<N>,
          M:     ga::num::Num+Zero+One,
{
    let classes: usize = From::from(classes); // Cast class count to usize
    let label: usize = From::from(label); // Cast label to usize
    let mut buf: Vec<M> = vec![Zero::zero(); classes]; // Create array of zeroes
    buf[label] = One::one(); // Set the one-hot component
    Array::from_vec(vec![1, buf.len()], buf) // Construct the array
}

fn read_u8<T: Read>(reader: &mut T) -> io::Result<u8> {
    use std::mem;

    let mut buf: [u8; 1] = [0];
    reader.read_exact(&mut buf).map(|_| {
        let data: u8 = unsafe { mem::transmute(buf) };
        data
    })
}

fn read_u32<T: Read>(reader: &mut T) -> io::Result<u32> {
    use std::mem;

    let mut buf: [u8; 4] = [0, 0, 0, 0];
    reader.read_exact(&mut buf).map(|_| {
        let data: u32 = unsafe { mem::transmute(buf) };
        data
    })
}
<|endoftext|>


<commit_before>extern crate deeplearn;
extern crate gpuarray as ga;

use std::fs::File;
use std::io::{
    self,
    BufReader,
    Read,
};
use std::path::Path;
use std::rc::Rc;

use deeplearn::{init, layers, util, Graph, Trainer};
use deeplearn::op::Relu;
use ga::Array;

fn main() {
    // Training data
    println!("Reading training labels...");
    let train_labels = read_mnist_labels("data/mnist/train-labels-idx1-ubyte", None).unwrap();
    let train_labels_logits: Vec<Array<f32>> = train_labels.iter().cloned()
                                                           .map(|x| util::one_hot_rows(x, 10))
                                                           .collect();
    println!("Label count: {}", train_labels.len());

    println!("Reading training images...");
    let (rows, columns, train_images) = read_mnist_images("data/mnist/train-images-idx3-ubyte", None).unwrap();

    // Validation data
    println!("Reading validation labels...");
    let val_labels = read_mnist_labels("data/mnist/t10k-labels-idx1-ubyte", Some(1000)).unwrap();
    println!("Label count: {}", val_labels.len());

    println!("Reading validation images...");
    let (_, _, val_images) = read_mnist_images("data/mnist/t10k-images-idx3-ubyte", Some(1000)).unwrap();

    ////////////////////////////////////////////////////////////////////////////////////////////////
    // Build the graph

    let ctx = Rc::new(ga::Context::new());
    let ref mut graph = Graph::new(ctx.clone());

    let batch_size = 1;

    //////////////////////////
    // Layer 1

    // Input. 1 batch of rows*columns inputs
    let input = graph.add_variable(vec![batch_size, rows*columns], false, vec![0.0; rows*columns]);

    // Biased fully connected layer with 300 neurons
    let (l1_fcb, _, _) = layers::dense_biased(graph, input, 300,
                                              init::Normal(0.001, 0.005),  // Weights initializer
                                              init::Normal(0.001, 0.005)); // Bias initializer
    let l1_out = layers::activation(graph, Relu(l1_fcb));

    //////////////////////////
    // Layer 2

    // Biased fully connected layer with 10 neurons
    let (l2_fcb, _, _) = layers::dense_biased(graph, l1_out, 10,
                                              init::Normal(0.001, 0.005),  // Weights initializer
                                              init::Normal(0.001, 0.005)); // Bias initializer
    let l2_out = layers::activation(graph, Relu(l2_fcb));
    let l2_out_d = graph.get_gradient(l2_out);

    //////////////////////////
    // Loss

    let (loss_out, train_out) = layers::mse(graph, l2_out);
    let loss_d = graph.add_gradient(loss_out); // Create a gradient to apply to the loss function

    ////////////////////////////////////////////////////////////////////////////////////////////////
    // Train and validate the network

    // We apply a gradient of -0.001 to the loss function
    let loss_d_cpu = Array::new(vec![batch_size, 10], -0.001);
    loss_d.write(graph, &loss_d_cpu);

    let mut loss_out_cpu = Array::new(vec![batch_size, 10], 0.0);
    let mut l2_out_cpu = Array::new(vec![batch_size, 10], 0.0);
    let mut l2_out_d_cpu = Array::new(vec![batch_size, 10], 0.0);

    let mut predictions = Array::new(vec![batch_size], 0usize);
    let mut num_correct = 0;

    {
        // Put this in it's own scope so that our train_update closure doesn't hold onto all of our
        // stuff until the end of main()
        let train_update = |graph: &mut Graph, epoch: usize| {
            // Get the output
            l2_out.read(graph, &mut l2_out_cpu);

            // Get the most likely digit (the index of the neuron with the highest output)
            util::argmax_rows(&l2_out_cpu, &mut predictions);
            let prediction = predictions[&[0]];

            // Check if the model was correct
            if prediction == train_labels[epoch] as usize {
                num_correct += 1;
            }

            if epoch % 1000 == 999 {
                l2_out_d.read(graph, &mut l2_out_d_cpu);
                loss_out.read(graph, &mut loss_out_cpu);
                println!("===================");
                println!("Epoch: {}", epoch);
                println!("out = {:?}", l2_out_cpu);
                println!("out_d = {:?}", l2_out_d_cpu);
                println!("loss = {:?}", loss_out_cpu);
                println!("Accuracy: {}%", (num_correct as f32)/(1000 as f32) * 100.0);
                num_correct = 0;
            }
        };

        let trainer = Trainer;
        trainer.train(graph, 60000, train_update,
                      &[(input, &train_images), (train_out, &train_labels_logits)]);
    }

    /////////////////////////
    // Validate the network
    println!("#######################################");
    println!("Validating");
    num_correct = 0;
    for epoch in 0..val_images.len() {
        // Upload training data
        input.write(graph, &val_images[epoch]);

        // Run the graph
        graph.forward();

        // Get the output
        l2_out.read(graph, &mut l2_out_cpu);

        // Get the most likely digit (the index of the neuron with the highest output)
        util::argmax_rows(&l2_out_cpu, &mut predictions);
        let prediction = predictions[&[0]];

        // Check if the model was correct
        if prediction == val_labels[epoch] as usize {
            num_correct += 1;
        }
    }
    println!("Validation Accuracy: {}%", (num_correct as f32)/(val_images.len() as f32) * 100.0);
}

fn read_mnist_labels<P: AsRef<Path>>(path: P, num_samples: Option<usize>) -> io::Result<Vec<u8>> {
    use std::io::{Error, ErrorKind};

    let ref mut file = BufReader::new(File::open(path).unwrap());

    let magic = u32::from_be(try!(read_u32(file)));
    if magic != 2049 {
        return Err(Error::new(ErrorKind::Other,
                              format!("Invalid magic number. Got expect 2049, got {}",
                                      magic).as_ref()))
    }

    let label_count = u32::from_be(try!(read_u32(file)));

    let mut labels = Vec::with_capacity(label_count as usize);
    for i in 0..label_count {
        if let Some(num_samples) = num_samples {
            if i as usize > num_samples {
                break;
            }
        }
        labels.push(try!(read_u8(file)));
    }

    Ok(labels)
}

fn read_mnist_images<P: AsRef<Path>>(path: P, num_samples: Option<usize>)
                                     -> io::Result<(usize, usize, Vec<Array<f32>>)> {
    use std::io::{Error, ErrorKind};

    let ref mut file = BufReader::new(File::open(path).unwrap());

    let magic = u32::from_be(try!(read_u32(file)));
    if magic != 2051 {
        return Err(Error::new(ErrorKind::Other,
                              format!("Invalid magic number. Got expect 2051, got {}",
                                      magic).as_ref()))
    }

    let image_count = u32::from_be(try!(read_u32(file))) as usize;
    let rows = u32::from_be(try!(read_u32(file))) as usize;
    let columns = u32::from_be(try!(read_u32(file))) as usize;

    let mut images = Vec::with_capacity(image_count);
    for i in 0..image_count {
        if let Some(num_samples) = num_samples {
            if i as usize > num_samples {
                break;
            }
        }
        let mut pixel_buf = vec![0u8; rows*columns];
        try!(file.read_exact(pixel_buf.as_mut()));
        let array = Array::from_vec(vec![rows, columns],
                                    pixel_buf.into_iter().map(|x| (x as f32)/255.0).collect());
        images.push(array);
    }

    Ok((rows, columns, images))
}

fn read_u8<T: Read>(reader: &mut T) -> io::Result<u8> {
    use std::mem;

    let mut buf: [u8; 1] = [0];
    reader.read_exact(&mut buf).map(|_| {
        let data: u8 = unsafe { mem::transmute(buf) };
        data
    })
}

fn read_u32<T: Read>(reader: &mut T) -> io::Result<u32> {
    use std::mem;

    let mut buf: [u8; 4] = [0, 0, 0, 0];
    reader.read_exact(&mut buf).map(|_| {
        let data: u32 = unsafe { mem::transmute(buf) };
        data
    })
}
<commit_msg>refactored mnist data reading<commit_after>extern crate deeplearn;
extern crate gpuarray as ga;

use std::fs::File;
use std::io::{
    self,
    BufReader,
    Read,
};
use std::path::Path;
use std::rc::Rc;

use deeplearn::{init, layers, util, Graph, Trainer};
use deeplearn::op::Relu;
use ga::Array;

fn main() {
    // Training data
    println!("Reading training labels...");
    let train_labels = read_mnist_labels("data/mnist/train-labels-idx1-ubyte", None).unwrap();
    let train_labels_logits: Vec<Array<f32>> = train_labels.iter().cloned()
                                                           .map(|x| util::one_hot_rows(x, 10))
                                                           .collect();
    println!("Label count: {}", train_labels.len());

    println!("Reading training images...");
    let (rows, columns, train_images) = read_mnist_images("data/mnist/train-images-idx3-ubyte", None).unwrap();

    // Validation data
    println!("Reading validation labels...");
    let val_labels = read_mnist_labels("data/mnist/t10k-labels-idx1-ubyte", Some(1000)).unwrap();
    println!("Label count: {}", val_labels.len());

    println!("Reading validation images...");
    let (_, _, val_images) = read_mnist_images("data/mnist/t10k-images-idx3-ubyte", Some(1000)).unwrap();

    ////////////////////////////////////////////////////////////////////////////////////////////////
    // Build the graph

    let ctx = Rc::new(ga::Context::new());
    let ref mut graph = Graph::new(ctx.clone());

    let batch_size = 1;

    //////////////////////////
    // Layer 1

    // Input. 1 batch of rows*columns inputs
    let input = graph.add_variable(vec![batch_size, rows*columns], false, vec![0.0; rows*columns]);

    // Biased fully connected layer with 300 neurons
    let (l1_fcb, _, _) = layers::dense_biased(graph, input, 300,
                                              init::Normal(0.001, 0.005),  // Weights initializer
                                              init::Normal(0.001, 0.005)); // Bias initializer
    let l1_out = layers::activation(graph, Relu(l1_fcb));

    //////////////////////////
    // Layer 2

    // Biased fully connected layer with 10 neurons
    let (l2_fcb, _, _) = layers::dense_biased(graph, l1_out, 10,
                                              init::Normal(0.001, 0.005),  // Weights initializer
                                              init::Normal(0.001, 0.005)); // Bias initializer
    let l2_out = layers::activation(graph, Relu(l2_fcb));
    let l2_out_d = graph.get_gradient(l2_out);

    //////////////////////////
    // Loss

    let (loss_out, train_out) = layers::mse(graph, l2_out);
    let loss_d = graph.add_gradient(loss_out); // Create a gradient to apply to the loss function

    ////////////////////////////////////////////////////////////////////////////////////////////////
    // Train and validate the network

    // We apply a gradient of -0.001 to the loss function
    let loss_d_cpu = Array::new(vec![batch_size, 10], -0.001);
    loss_d.write(graph, &loss_d_cpu);

    let mut loss_out_cpu = Array::new(vec![batch_size, 10], 0.0);
    let mut l2_out_cpu = Array::new(vec![batch_size, 10], 0.0);
    let mut l2_out_d_cpu = Array::new(vec![batch_size, 10], 0.0);

    let mut predictions = Array::new(vec![batch_size], 0usize);
    let mut num_correct = 0;

    {
        // Put this in it's own scope so that our train_update closure doesn't hold onto all of our
        // stuff until the end of main()
        let train_update = |graph: &mut Graph, epoch: usize| {
            // Get the output
            l2_out.read(graph, &mut l2_out_cpu);

            // Get the most likely digit (the index of the neuron with the highest output)
            util::argmax_rows(&l2_out_cpu, &mut predictions);
            let prediction = predictions[&[0]];

            // Check if the model was correct
            if prediction == train_labels[epoch] as usize {
                num_correct += 1;
            }

            if epoch % 1000 == 999 {
                l2_out_d.read(graph, &mut l2_out_d_cpu);
                loss_out.read(graph, &mut loss_out_cpu);
                println!("===================");
                println!("Epoch: {}", epoch);
                println!("out = {:?}", l2_out_cpu);
                println!("out_d = {:?}", l2_out_d_cpu);
                println!("loss = {:?}", loss_out_cpu);
                println!("Accuracy: {}%", (num_correct as f32)/(1000 as f32) * 100.0);
                num_correct = 0;
            }
        };

        let trainer = Trainer;
        trainer.train(graph, 60000, train_update,
                      &[(input, &train_images), (train_out, &train_labels_logits)]);
    }

    /////////////////////////
    // Validate the network
    println!("#######################################");
    println!("Validating");
    num_correct = 0;
    for epoch in 0..val_images.len() {
        // Upload training data
        input.write(graph, &val_images[epoch]);

        // Run the graph
        graph.forward();

        // Get the output
        l2_out.read(graph, &mut l2_out_cpu);

        // Get the most likely digit (the index of the neuron with the highest output)
        util::argmax_rows(&l2_out_cpu, &mut predictions);
        let prediction = predictions[&[0]];

        // Check if the model was correct
        if prediction == val_labels[epoch] as usize {
            num_correct += 1;
        }
    }
    println!("Validation Accuracy: {}%", (num_correct as f32)/(val_images.len() as f32) * 100.0);
}

fn read_mnist_labels<P: AsRef<Path>>(path: P, num_samples: Option<usize>) -> io::Result<Vec<u8>> {
    use std::cmp;
    use std::io::{Error, ErrorKind};

    let ref mut file = BufReader::new(File::open(path).unwrap());

    let magic = u32::from_be(try!(read_u32(file)));
    if magic != 2049 {
        return Err(Error::new(ErrorKind::Other,
                              format!("Invalid magic number. Got expect 2049, got {}",
                                      magic).as_ref()))
    }

    let label_count = u32::from_be(try!(read_u32(file))) as usize;
    let label_count = cmp::min(label_count, num_samples.unwrap_or(label_count));

    let mut labels = Vec::with_capacity(label_count);
    for _ in 0..label_count {
        labels.push(try!(read_u8(file)));
    }

    Ok(labels)
}

fn read_mnist_images<P: AsRef<Path>>(path: P, num_samples: Option<usize>)
                                     -> io::Result<(usize, usize, Vec<Array<f32>>)> {
    use std::cmp;
    use std::io::{Error, ErrorKind};

    let ref mut file = BufReader::new(File::open(path).unwrap());

    let magic = u32::from_be(try!(read_u32(file)));
    if magic != 2051 {
        return Err(Error::new(ErrorKind::Other,
                              format!("Invalid magic number. Got expect 2051, got {}",
                                      magic).as_ref()))
    }

    let image_count = u32::from_be(try!(read_u32(file))) as usize;
    let rows = u32::from_be(try!(read_u32(file))) as usize;
    let columns = u32::from_be(try!(read_u32(file))) as usize;

    let image_count = cmp::min(image_count, num_samples.unwrap_or(image_count));

    let mut images = Vec::with_capacity(image_count);
    for _ in 0..image_count {
        let mut pixel_buf = vec![0u8; rows*columns];
        try!(file.read_exact(pixel_buf.as_mut()));
        let array = Array::from_vec(vec![rows, columns],
                                    pixel_buf.into_iter().map(|x| (x as f32)/255.0).collect());
        images.push(array);
    }

    Ok((rows, columns, images))
}

fn read_u8<T: Read>(reader: &mut T) -> io::Result<u8> {
    use std::mem;

    let mut buf: [u8; 1] = [0];
    reader.read_exact(&mut buf).map(|_| {
        let data: u8 = unsafe { mem::transmute(buf) };
        data
    })
}

fn read_u32<T: Read>(reader: &mut T) -> io::Result<u32> {
    use std::mem;

    let mut buf: [u8; 4] = [0, 0, 0, 0];
    reader.read_exact(&mut buf).map(|_| {
        let data: u32 = unsafe { mem::transmute(buf) };
        data
    })
}
<|endoftext|>


<commit_before>extern crate rier;
#[macro_use] extern crate glium;
extern crate cgmath;
extern crate wavefront_obj;

use std::fs::File;
use std::io::Read;
use wavefront_obj::obj;
use glium::DrawParameters;
use cgmath::{Rad, Point3};

type Mesh = rier::Mesh<Vertex>;

#[derive(Copy, Clone)]
struct Vertex {
    position: [f32; 3],
    normal: [f32; 3],
}

implement_vertex! {Vertex, position, normal}

struct Shader;

type Renderer = rier::Renderer<Shader>;

impl rier::Shader for Shader {
    type Vertex = Vertex;

    fn vertex() -> &'static str {
        r#"
        #version 330 core
        uniform mat4 camera;
        uniform mat4 transform;
        in vec3 position;
        in vec3 normal;
        out vec3 v_normal;
        out vec3 pixel_pos;

        void main() {
            gl_Position = camera * transform * vec4(position, 1.0);
            pixel_pos = vec3(transform * vec4(position, 1.0f));
            v_normal = normalize(mat3(transpose(inverse(transform))) * normal);
        }
        "#
    }

    fn fragment() -> &'static str {
        r#"
        #version 330 core
        uniform vec3 light_pos;
        uniform vec3 view_pos;
        in vec3 v_normal;
        in vec3 pixel_pos;
        out vec4 f_color;

        vec3 light_color = vec3(1.0, 1.0, 1.0);
        vec3 light_dir = normalize(light_pos - pixel_pos);

        vec3 diffuse() {
            float diff = max(dot(v_normal, light_dir), 0.0);
            return diff * light_color;
        }

        vec3 ambient() {
            float strength = 0.1f;
            return strength * light_color;
        }

        vec3 specular() {
            float strength = 0.5f;
            vec3 view_dir = normalize(view_pos - pixel_pos);
            vec3 reflect_dir = reflect(-light_dir, v_normal);
            float spec = pow(max(dot(view_dir, reflect_dir), 0.0), 32);
            return strength * spec * light_color;
        }

        void main() {
            vec3 obj_color = vec3(0.25, 0.7, 0.25);
            vec3 result = (ambient()+diffuse()+specular()) * obj_color;
            f_color = vec4(result, 1.0);
        }
        "#
    }

    fn draw_parameters() -> DrawParameters<'static> {
        use std::default::Default;
        use glium::draw_parameters::{Depth, DepthTest, BackfaceCullingMode};

        let depth = Depth {
            test: DepthTest::IfLess,
            write: true,
            .. Default::default()
        };

        DrawParameters {
            depth: depth,
            backface_culling: BackfaceCullingMode::CullClockwise,
            ..Default::default()
        }
    }
}

struct Model {
    mesh: Mesh,
    transform: rier::Transform,
}

impl Model {
    fn new(renderer: &Renderer) -> Model {
        let mesh = Model::mesh(renderer);

        Model {
            mesh: mesh,
            transform: rier::Transform::new(),
        }
    }

    fn load() -> obj::Object {
        let mut file = File::open("./examples/assets/monkey.obj").unwrap();
        let mut s = String::new();
        file.read_to_string(&mut s).unwrap();
        let mut objects = obj::parse(s).unwrap();
        objects.objects.pop().unwrap()
    }

    fn make_vertex(object: &obj::Object, index: obj::VTNIndex) -> Vertex {
        let (v, _, n) = index;
        let n = n.unwrap();
        let vertex = object.vertices[v];
        let normal = object.normals[n];
        Vertex {
            position: [vertex.x as f32, vertex.y as f32, vertex.z as f32],
            normal  : [normal.x as f32, normal.y as f32, normal.z as f32],
        }
    }

    fn mesh(renderer: &Renderer) -> Mesh {
        let mut vertices: Vec<Vertex> = Vec::new();
        let object = Model::load();
        for geometry in &object.geometry {
            for shape in &geometry.shapes {
                match shape {
                    &obj::Shape::Triangle(a, b, c) => {
                        vertices.push(Model::make_vertex(&object, a));
                        vertices.push(Model::make_vertex(&object, b));
                        vertices.push(Model::make_vertex(&object, c));
                    }
                    _ => unimplemented!(),
                }

            }
        }
        Mesh::new(renderer, &*vertices).unwrap()
    }

    fn render(&self, renderer: &Renderer, camera: &rier::Camera3D, light: Point3<f32>) {
        type Point = [f32; 3];
        let light: Point = light.into();
        let view_pos: Point = camera.eye.into();

        let uniforms = uniform! {
            camera: camera,
            transform: &self.transform,
            light_pos: light,
            view_pos: view_pos,
        };
        renderer.draw(&self.mesh, &uniforms).unwrap();
    }
}

fn main() {
    let gfx = rier::Context::create("Cube", (800, 600)).gfx();
    let renderer = Renderer::new(gfx.clone()).unwrap();
    let mut camera = rier::Camera3D::new(gfx.clone());
    camera.eye = Point3::new(4.0, 3.0, 3.0);
    let mut model = Model::new(&renderer);
    let mut x = 0.0f32;
    let main_loop = rier::Loop::new(move |delta| {
        use rier::main_loop::Return::*;

        camera.update();
        x += delta.subsec_nanos() as f32 / 1000000000.0;
        model.transform.set_rotation(Rad(x), Rad(x), Rad(0.0));
        model.transform.dirty();
        for event in gfx.display.poll_events() {
            match event {
                rier::WindowEvent::Closed => return Exit,
                _ => (),
            }
        }
        gfx.frame(|| {
            model.render(&renderer, &camera, Point3::new(2.2, 1.0, 2.0));
        }).unwrap();
        Next
    });
    main_loop.start();
}
<commit_msg>Breaking fix.<commit_after>extern crate rier;
#[macro_use] extern crate glium;
extern crate cgmath;
extern crate wavefront_obj;

use std::fs::File;
use std::io::Read;
use wavefront_obj::obj;
use glium::DrawParameters;
use cgmath::{Rad, Point3};

type Mesh = rier::Mesh<Vertex>;

#[derive(Copy, Clone)]
struct Vertex {
    position: [f32; 3],
    normal: [f32; 3],
}

implement_vertex! {Vertex, position, normal}

struct Shader;

type Renderer = rier::Renderer<Shader>;

impl rier::Shader for Shader {
    type Vertex = Vertex;

    fn vertex() -> &'static str {
        r#"
        #version 330 core
        uniform mat4 camera;
        uniform mat4 transform;
        in vec3 position;
        in vec3 normal;
        out vec3 v_normal;
        out vec3 pixel_pos;

        void main() {
            gl_Position = camera * transform * vec4(position, 1.0);
            pixel_pos = vec3(transform * vec4(position, 1.0f));
            v_normal = normalize(mat3(transpose(inverse(transform))) * normal);
        }
        "#
    }

    fn fragment() -> &'static str {
        r#"
        #version 330 core
        uniform vec3 light_pos;
        uniform vec3 view_pos;
        in vec3 v_normal;
        in vec3 pixel_pos;
        out vec4 f_color;

        vec3 light_color = vec3(1.0, 1.0, 1.0);
        vec3 light_dir = normalize(light_pos - pixel_pos);

        vec3 diffuse() {
            float diff = max(dot(v_normal, light_dir), 0.0);
            return diff * light_color;
        }

        vec3 ambient() {
            float strength = 0.1f;
            return strength * light_color;
        }

        vec3 specular() {
            float strength = 0.5f;
            vec3 view_dir = normalize(view_pos - pixel_pos);
            vec3 reflect_dir = reflect(-light_dir, v_normal);
            float spec = pow(max(dot(view_dir, reflect_dir), 0.0), 32);
            return strength * spec * light_color;
        }

        void main() {
            vec3 obj_color = vec3(0.25, 0.7, 0.25);
            vec3 result = (ambient()+diffuse()+specular()) * obj_color;
            f_color = vec4(result, 1.0);
        }
        "#
    }

    fn draw_parameters() -> DrawParameters<'static> {
        use std::default::Default;
        use glium::draw_parameters::{Depth, DepthTest, BackfaceCullingMode};

        let depth = Depth {
            test: DepthTest::IfLess,
            write: true,
            .. Default::default()
        };

        DrawParameters {
            depth: depth,
            backface_culling: BackfaceCullingMode::CullClockwise,
            ..Default::default()
        }
    }
}

struct Model {
    mesh: Mesh,
    transform: rier::Transform,
}

impl Model {
    fn new(renderer: &Renderer) -> Model {
        let mesh = Model::mesh(renderer);

        Model {
            mesh: mesh,
            transform: rier::Transform::new(),
        }
    }

    fn load() -> obj::Object {
        let mut file = File::open("./examples/assets/monkey.obj").unwrap();
        let mut s = String::new();
        file.read_to_string(&mut s).unwrap();
        let mut objects = obj::parse(s).unwrap();
        objects.objects.pop().unwrap()
    }

    fn make_vertex(object: &obj::Object, index: obj::VTNIndex) -> Vertex {
        let (v, _, n) = index;
        let n = n.unwrap();
        let vertex = object.vertices[v];
        let normal = object.normals[n];
        Vertex {
            position: [vertex.x as f32, vertex.y as f32, vertex.z as f32],
            normal  : [normal.x as f32, normal.y as f32, normal.z as f32],
        }
    }

    fn mesh(renderer: &Renderer) -> Mesh {
        let mut vertices: Vec<Vertex> = Vec::new();
        let object = Model::load();
        for geometry in &object.geometry {
            for shape in &geometry.shapes {
                match shape.primitive {
                    obj::Primitive::Triangle(a, b, c) => {
                        vertices.push(Model::make_vertex(&object, a));
                        vertices.push(Model::make_vertex(&object, b));
                        vertices.push(Model::make_vertex(&object, c));
                    }
                    _ => unimplemented!(),
                }

            }
        }
        Mesh::new(renderer, &*vertices).unwrap()
    }

    fn render(&self, renderer: &Renderer, camera: &rier::Camera3D, light: Point3<f32>) {
        type Point = [f32; 3];
        let light: Point = light.into();
        let view_pos: Point = camera.eye.into();

        let uniforms = uniform! {
            camera: camera,
            transform: &self.transform,
            light_pos: light,
            view_pos: view_pos,
        };
        renderer.draw(&self.mesh, &uniforms).unwrap();
    }
}

fn main() {
    let gfx = rier::Context::create("Cube", (800, 600)).gfx();
    let renderer = Renderer::new(gfx.clone()).unwrap();
    let mut camera = rier::Camera3D::new(gfx.clone());
    camera.eye = Point3::new(4.0, 3.0, 3.0);
    let mut model = Model::new(&renderer);
    let mut x = 0.0f32;
    let main_loop = rier::Loop::new(move |delta| {
        use rier::main_loop::Return::*;

        camera.update();
        x += delta.subsec_nanos() as f32 / 1000000000.0;
        model.transform.set_rotation(Rad(x), Rad(x), Rad(0.0));
        model.transform.dirty();
        for event in gfx.display.poll_events() {
            match event {
                rier::WindowEvent::Closed => return Exit,
                _ => (),
            }
        }
        gfx.frame(|| {
            model.render(&renderer, &camera, Point3::new(2.2, 1.0, 2.0));
        }).unwrap();
        Next
    });
    main_loop.start();
}
<|endoftext|>


<commit_before>extern crate pretty;

use pretty::{
    Doc
};

#[deriving(Clone)]
#[deriving(Show)]
pub struct Tree<'a> {
    node:String,
    subtrees:&'a[Tree<'a>]
}

impl<'a> Tree<'a> {
    pub fn new(node:&str, subtrees:&'a[Tree<'a>]) -> Tree<'a> {
        Tree {
            node: String::from_str(node),
            subtrees: subtrees
        }
    }

    fn pretty_trees(trees:&'a[Tree<'a>]) -> Doc {
        match trees {
            [] => panic!(),
            [ref t] => t.pretty(),
            [ref t, ref ts..] => {
                t.pretty().append(
                    Doc::text(",")
                ).append(
                    Doc::newline()
                ).append(
                    Tree::pretty_trees(*ts)
                )
            }
        }
    }

    fn pretty_bracket(ts:&'a[Tree<'a>]) -> Doc {
        match ts {
            [] => Doc::nil(),
            ts => {
                Doc::text("[").append(
                    Doc::newline().append(
                        Tree::pretty_trees(ts)
                    ).nest(2)
                ).append(
                    Doc::newline()
                ).append(
                    Doc::text("]")
                )
            }
        }
    }

    pub fn pretty(&self) -> Doc {
        Doc::text(
            self.node.clone()
        ).append(
            Tree::pretty_bracket(
                self.subtrees
            )
        ).group()
    }
}

#[allow(dead_code)]
pub fn main() {
    let bbbbbbs =
        [ Tree::new("ccc", [])
        , Tree::new("dd", [])
        ];
    let ffffs =
        [ Tree::new("gg", [])
        , Tree::new("hhh", [])
        , Tree::new("ii", [])
        ];

    let aaas =
        [ Tree::new("bbbbbb", bbbbbbs)
        , Tree::new("eee", [])
        , Tree::new("ffff", ffffs)
        ];

    let example = Tree::new("aaaa", aaas);

    print!("{}", example.pretty().render(70));
}
<commit_msg>change print! to println!<commit_after>extern crate pretty;

use pretty::{
    Doc
};

#[deriving(Clone)]
#[deriving(Show)]
pub struct Tree<'a> {
    node:String,
    subtrees:&'a[Tree<'a>]
}

impl<'a> Tree<'a> {
    pub fn new(node:&str, subtrees:&'a[Tree<'a>]) -> Tree<'a> {
        Tree {
            node: String::from_str(node),
            subtrees: subtrees
        }
    }

    fn pretty_trees(trees:&'a[Tree<'a>]) -> Doc {
        match trees {
            [] => panic!(),
            [ref t] => t.pretty(),
            [ref t, ref ts..] => {
                t.pretty().append(
                    Doc::text(",")
                ).append(
                    Doc::newline()
                ).append(
                    Tree::pretty_trees(*ts)
                )
            }
        }
    }

    fn pretty_bracket(ts:&'a[Tree<'a>]) -> Doc {
        match ts {
            [] => Doc::nil(),
            ts => {
                Doc::text("[").append(
                    Doc::newline().append(
                        Tree::pretty_trees(ts)
                    ).nest(2)
                ).append(
                    Doc::newline()
                ).append(
                    Doc::text("]")
                )
            }
        }
    }

    pub fn pretty(&self) -> Doc {
        Doc::text(
            self.node.clone()
        ).append(
            Tree::pretty_bracket(
                self.subtrees
            )
        ).group()
    }
}

#[allow(dead_code)]
pub fn main() {
    let bbbbbbs =
        [ Tree::new("ccc", [])
        , Tree::new("dd", [])
        ];
    let ffffs =
        [ Tree::new("gg", [])
        , Tree::new("hhh", [])
        , Tree::new("ii", [])
        ];

    let aaas =
        [ Tree::new("bbbbbb", bbbbbbs)
        , Tree::new("eee", [])
        , Tree::new("ffff", ffffs)
        ];

    let example = Tree::new("aaaa", aaas);

    println!("{}", example.pretty().render(70));
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;
use std::collections::{HashMap, HashSet};
use std::ops::Range;
use std::rc::Rc;
use std::time::{Duration, Instant};

use core::interning::InternedString;
use core::{Dependency, PackageId, PackageIdSpec, Registry, Summary};
use util::errors::CargoResult;
use util::Config;

use im_rc;

pub struct ResolverProgress {
    ticks: u16,
    start: Instant,
    time_to_print: Duration,
    printed: bool,
    deps_time: Duration,
}

impl ResolverProgress {
    pub fn new() -> ResolverProgress {
        ResolverProgress {
            ticks: 0,
            start: Instant::now(),
            time_to_print: Duration::from_millis(500),
            printed: false,
            deps_time: Duration::new(0, 0),
        }
    }
    pub fn shell_status(&mut self, config: Option<&Config>) -> CargoResult<()> {
        // If we spend a lot of time here (we shouldn't in most cases) then give
        // a bit of a visual indicator as to what we're doing. Only enable this
        // when stderr is a tty (a human is likely to be watching) to ensure we
        // get deterministic output otherwise when observed by tools.
        //
        // Also note that we hit this loop a lot, so it's fairly performance
        // sensitive. As a result try to defer a possibly expensive operation
        // like `Instant::now` by only checking every N iterations of this loop
        // to amortize the cost of the current time lookup.
        self.ticks += 1;
        if let Some(config) = config {
            if config.shell().is_err_tty()
                && !self.printed
                && self.ticks % 1000 == 0
                && self.start.elapsed() - self.deps_time > self.time_to_print
            {
                self.printed = true;
                config.shell().status("Resolving", "dependency graph...")?;
            }
        }
        // The largest test in our suite takes less then 5000 ticks
        // with all the algorithm improvements.
        // If any of them are removed then it takes more than I am willing to measure.
        // So lets fail the test fast if we have ben running for two long.
        debug_assert!(self.ticks < 50_000);
        // The largest test in our suite takes less then 30 sec
        // with all the improvements to how fast a tick can go.
        // If any of them are removed then it takes more than I am willing to measure.
        // So lets fail the test fast if we have ben running for two long.
        if cfg!(debug_assertions) && (self.ticks % 1000 == 0) {
            assert!(self.start.elapsed() - self.deps_time < Duration::from_secs(90));
        }
        Ok(())
    }
    pub fn elapsed(&mut self, dur: Duration) {
        self.deps_time += dur;
    }
}

pub struct RegistryQueryer<'a> {
    pub registry: &'a mut (Registry + 'a),
    replacements: &'a [(PackageIdSpec, Dependency)],
    try_to_use: &'a HashSet<&'a PackageId>,
    cache: HashMap<Dependency, Rc<Vec<Candidate>>>,
    // If set the list of dependency candidates will be sorted by minimal
    // versions first. That allows `cargo update -Z minimal-versions` which will
    // specify minimum dependency versions to be used.
    minimal_versions: bool,
}

impl<'a> RegistryQueryer<'a> {
    pub fn new(
        registry: &'a mut Registry,
        replacements: &'a [(PackageIdSpec, Dependency)],
        try_to_use: &'a HashSet<&'a PackageId>,
        minimal_versions: bool,
    ) -> Self {
        RegistryQueryer {
            registry,
            replacements,
            cache: HashMap::new(),
            try_to_use,
            minimal_versions,
        }
    }

    /// Queries the `registry` to return a list of candidates for `dep`.
    ///
    /// This method is the location where overrides are taken into account. If
    /// any candidates are returned which match an override then the override is
    /// applied by performing a second query for what the override should
    /// return.
    pub fn query(&mut self, dep: &Dependency) -> CargoResult<Rc<Vec<Candidate>>> {
        if let Some(out) = self.cache.get(dep).cloned() {
            return Ok(out);
        }

        let mut ret = Vec::new();
        self.registry.query(
            dep,
            &mut |s| {
                ret.push(Candidate {
                    summary: s,
                    replace: None,
                });
            },
            false,
        )?;
        for candidate in ret.iter_mut() {
            let summary = &candidate.summary;

            let mut potential_matches = self
                .replacements
                .iter()
                .filter(|&&(ref spec, _)| spec.matches(summary.package_id()));

            let &(ref spec, ref dep) = match potential_matches.next() {
                None => continue,
                Some(replacement) => replacement,
            };
            debug!(
                "found an override for {} {}",
                dep.package_name(),
                dep.version_req()
            );

            let mut summaries = self.registry.query_vec(dep, false)?.into_iter();
            let s = summaries.next().ok_or_else(|| {
                format_err!(
                    "no matching package for override `{}` found\n\
                     location searched: {}\n\
                     version required: {}",
                    spec,
                    dep.source_id(),
                    dep.version_req()
                )
            })?;
            let summaries = summaries.collect::<Vec<_>>();
            if !summaries.is_empty() {
                let bullets = summaries
                    .iter()
                    .map(|s| format!("  * {}", s.package_id()))
                    .collect::<Vec<_>>();
                bail!(
                    "the replacement specification `{}` matched \
                     multiple packages:\n  * {}\n{}",
                    spec,
                    s.package_id(),
                    bullets.join("\n")
                );
            }

            // The dependency should be hard-coded to have the same name and an
            // exact version requirement, so both of these assertions should
            // never fail.
            assert_eq!(s.version(), summary.version());
            assert_eq!(s.name(), summary.name());

            let replace = if s.source_id() == summary.source_id() {
                debug!("Preventing\n{:?}\nfrom replacing\n{:?}", summary, s);
                None
            } else {
                Some(s)
            };
            let matched_spec = spec.clone();

            // Make sure no duplicates
            if let Some(&(ref spec, _)) = potential_matches.next() {
                bail!(
                    "overlapping replacement specifications found:\n\n  \
                     * {}\n  * {}\n\nboth specifications match: {}",
                    matched_spec,
                    spec,
                    summary.package_id()
                );
            }

            for dep in summary.dependencies() {
                debug!("\t{} => {}", dep.package_name(), dep.version_req());
            }

            candidate.replace = replace;
        }

        // When we attempt versions for a package we'll want to do so in a
        // sorted fashion to pick the "best candidates" first. Currently we try
        // prioritized summaries (those in `try_to_use`) and failing that we
        // list everything from the maximum version to the lowest version.
        ret.sort_unstable_by(|a, b| {
            let a_in_previous = self.try_to_use.contains(a.summary.package_id());
            let b_in_previous = self.try_to_use.contains(b.summary.package_id());
            let previous_cmp = a_in_previous.cmp(&b_in_previous).reverse();
            match previous_cmp {
                Ordering::Equal => {
                    let cmp = a.summary.version().cmp(b.summary.version());
                    if self.minimal_versions {
                        // Lower version ordered first.
                        cmp
                    } else {
                        // Higher version ordered first.
                        cmp.reverse()
                    }
                }
                _ => previous_cmp,
            }
        });

        let out = Rc::new(ret);

        self.cache.insert(dep.clone(), out.clone());

        Ok(out)
    }
}

#[derive(Clone, Copy)]
pub enum Method<'a> {
    Everything, // equivalent to Required { dev_deps: true, all_features: true, .. }
    Required {
        dev_deps: bool,
        features: &'a [InternedString],
        all_features: bool,
        uses_default_features: bool,
    },
}

impl<'r> Method<'r> {
    pub fn split_features(features: &[String]) -> Vec<InternedString> {
        features
            .iter()
            .flat_map(|s| s.split_whitespace())
            .flat_map(|s| s.split(','))
            .filter(|s| !s.is_empty())
            .map(|s| InternedString::new(s))
            .collect::<Vec<InternedString>>()
    }
}

#[derive(Clone)]
pub struct Candidate {
    pub summary: Summary,
    pub replace: Option<Summary>,
}

#[derive(Clone)]
pub struct DepsFrame {
    pub parent: Summary,
    pub just_for_error_messages: bool,
    pub remaining_siblings: RcVecIter<DepInfo>,
}

impl DepsFrame {
    /// Returns the least number of candidates that any of this frame's siblings
    /// has.
    ///
    /// The `remaining_siblings` array is already sorted with the smallest
    /// number of candidates at the front, so we just return the number of
    /// candidates in that entry.
    fn min_candidates(&self) -> usize {
        self.remaining_siblings
            .peek()
            .map(|(_, (_, candidates, _))| candidates.len())
            .unwrap_or(0)
    }

    pub fn flatten(&self) -> impl Iterator<Item = (&PackageId, Dependency)> {
        self.remaining_siblings
            .clone()
            .map(move |(_, (d, _, _))| (self.parent.package_id(), d))
    }
}

impl PartialEq for DepsFrame {
    fn eq(&self, other: &DepsFrame) -> bool {
        self.just_for_error_messages == other.just_for_error_messages
            && self.min_candidates() == other.min_candidates()
    }
}

impl Eq for DepsFrame {}

impl PartialOrd for DepsFrame {
    fn partial_cmp(&self, other: &DepsFrame) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for DepsFrame {
    fn cmp(&self, other: &DepsFrame) -> Ordering {
        self.just_for_error_messages
            .cmp(&other.just_for_error_messages)
            .reverse()
            .then_with(|| self.min_candidates().cmp(&other.min_candidates()))
    }
}

/// Note that a `OrdSet` is used for the remaining dependencies that need
/// activation. This set is sorted by how many candidates each dependency has.
///
/// This helps us get through super constrained portions of the dependency
/// graph quickly and hopefully lock down what later larger dependencies can
/// use (those with more candidates).
#[derive(Clone)]
pub struct RemainingDeps(usize, im_rc::OrdSet<(DepsFrame, usize)>);

impl RemainingDeps {
    pub fn new() -> RemainingDeps {
        RemainingDeps(0, im_rc::OrdSet::new())
    }
    pub fn push(&mut self, x: DepsFrame) {
        self.1.insert((x, self.0));
        self.0 += 1;
    }
    pub fn pop_most_constrained(&mut self) -> Option<(bool, (Summary, (usize, DepInfo)))> {
        while let Some((mut deps_frame, i)) = self.1.remove_min() {
            let just_here_for_the_error_messages = deps_frame.just_for_error_messages;

            // Figure out what our next dependency to activate is, and if nothing is
            // listed then we're entirely done with this frame (yay!) and we can
            // move on to the next frame.
            if let Some(sibling) = deps_frame.remaining_siblings.next() {
                let parent = Summary::clone(&deps_frame.parent);
                self.1.insert((deps_frame, i));
                return Some((just_here_for_the_error_messages, (parent, sibling)));
            }
        }
        None
    }
    pub fn iter(&mut self) -> impl Iterator<Item = (&PackageId, Dependency)> {
        self.1.iter().flat_map(|(other, _)| other.flatten())
    }
}

// Information about the dependencies for a crate, a tuple of:
//
// (dependency info, candidates, features activated)
pub type DepInfo = (Dependency, Rc<Vec<Candidate>>, Rc<Vec<InternedString>>);

/// All possible reasons that a package might fail to activate.
///
/// We maintain a list of conflicts for error reporting as well as backtracking
/// purposes. Each reason here is why candidates may be rejected or why we may
/// fail to resolve a dependency.
#[derive(Debug, Clone, PartialOrd, Ord, PartialEq, Eq)]
pub enum ConflictReason {
    /// There was a semver conflict, for example we tried to activate a package
    /// 1.0.2 but 1.1.0 was already activated (aka a compatible semver version
    /// is already activated)
    Semver,

    /// The `links` key is being violated. For example one crate in the
    /// dependency graph has `links = "foo"` but this crate also had that, and
    /// we're only allowed one per dependency graph.
    Links(InternedString),

    /// A dependency listed features that weren't actually available on the
    /// candidate. For example we tried to activate feature `foo` but the
    /// candidate we're activating didn't actually have the feature `foo`.
    MissingFeatures(String),
}

impl ConflictReason {
    pub fn is_links(&self) -> bool {
        if let ConflictReason::Links(_) = *self {
            return true;
        }
        false
    }

    pub fn is_missing_features(&self) -> bool {
        if let ConflictReason::MissingFeatures(_) = *self {
            return true;
        }
        false
    }
}

pub struct RcVecIter<T> {
    vec: Rc<Vec<T>>,
    rest: Range<usize>,
}

impl<T> RcVecIter<T> {
    pub fn new(vec: Rc<Vec<T>>) -> RcVecIter<T> {
        RcVecIter {
            rest: 0..vec.len(),
            vec,
        }
    }

    fn peek(&self) -> Option<(usize, &T)> {
        self.rest
            .clone()
            .next()
            .and_then(|i| self.vec.get(i).map(|val| (i, &*val)))
    }
}

// Not derived to avoid `T: Clone`
impl<T> Clone for RcVecIter<T> {
    fn clone(&self) -> RcVecIter<T> {
        RcVecIter {
            vec: self.vec.clone(),
            rest: self.rest.clone(),
        }
    }
}

impl<T> Iterator for RcVecIter<T>
where
    T: Clone,
{
    type Item = (usize, T);

    fn next(&mut self) -> Option<Self::Item> {
        self.rest
            .next()
            .and_then(|i| self.vec.get(i).map(|val| (i, val.clone())))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        // rest is a std::ops::Range, which is an ExactSizeIterator.
        self.rest.size_hint()
    }
}

impl<T: Clone> ExactSizeIterator for RcVecIter<T> {}

pub struct RcList<T> {
    pub head: Option<Rc<(T, RcList<T>)>>,
}

impl<T> RcList<T> {
    pub fn new() -> RcList<T> {
        RcList { head: None }
    }

    pub fn push(&mut self, data: T) {
        let node = Rc::new((
            data,
            RcList {
                head: self.head.take(),
            },
        ));
        self.head = Some(node);
    }
}

// Not derived to avoid `T: Clone`
impl<T> Clone for RcList<T> {
    fn clone(&self) -> RcList<T> {
        RcList {
            head: self.head.clone(),
        }
    }
}

// Avoid stack overflows on drop by turning recursion into a loop
impl<T> Drop for RcList<T> {
    fn drop(&mut self) {
        let mut cur = self.head.take();
        while let Some(head) = cur {
            match Rc::try_unwrap(head) {
                Ok((_data, mut next)) => cur = next.head.take(),
                Err(_) => break,
            }
        }
    }
}

pub enum GraphNode {
    Add(PackageId),
    Link(PackageId, PackageId, Dependency),
}
<commit_msg>better error for benchmarking<commit_after>use std::cmp::Ordering;
use std::collections::{HashMap, HashSet};
use std::ops::Range;
use std::rc::Rc;
use std::time::{Duration, Instant};

use core::interning::InternedString;
use core::{Dependency, PackageId, PackageIdSpec, Registry, Summary};
use util::errors::CargoResult;
use util::Config;

use im_rc;

pub struct ResolverProgress {
    ticks: u16,
    start: Instant,
    time_to_print: Duration,
    printed: bool,
    deps_time: Duration,
}

impl ResolverProgress {
    pub fn new() -> ResolverProgress {
        ResolverProgress {
            ticks: 0,
            start: Instant::now(),
            time_to_print: Duration::from_millis(500),
            printed: false,
            deps_time: Duration::new(0, 0),
        }
    }
    pub fn shell_status(&mut self, config: Option<&Config>) -> CargoResult<()> {
        // If we spend a lot of time here (we shouldn't in most cases) then give
        // a bit of a visual indicator as to what we're doing. Only enable this
        // when stderr is a tty (a human is likely to be watching) to ensure we
        // get deterministic output otherwise when observed by tools.
        //
        // Also note that we hit this loop a lot, so it's fairly performance
        // sensitive. As a result try to defer a possibly expensive operation
        // like `Instant::now` by only checking every N iterations of this loop
        // to amortize the cost of the current time lookup.
        self.ticks += 1;
        if let Some(config) = config {
            if config.shell().is_err_tty()
                && !self.printed
                && self.ticks % 1000 == 0
                && self.start.elapsed() - self.deps_time > self.time_to_print
            {
                self.printed = true;
                config.shell().status("Resolving", "dependency graph...")?;
            }
        }
        // The largest test in our suite takes less then 5000 ticks
        // with all the algorithm improvements.
        // If any of them are removed then it takes more than I am willing to measure.
        // So lets fail the test fast if we have ben running for two long.
        debug_assert!(
            self.ticks < 50_000,
            "got to 50_000 ticks in {:?}",
            self.start.elapsed()
        );
        // The largest test in our suite takes less then 30 sec
        // with all the improvements to how fast a tick can go.
        // If any of them are removed then it takes more than I am willing to measure.
        // So lets fail the test fast if we have ben running for two long.
        if cfg!(debug_assertions) && (self.ticks % 1000 == 0) {
            assert!(self.start.elapsed() - self.deps_time < Duration::from_secs(90));
        }
        Ok(())
    }
    pub fn elapsed(&mut self, dur: Duration) {
        self.deps_time += dur;
    }
}

pub struct RegistryQueryer<'a> {
    pub registry: &'a mut (Registry + 'a),
    replacements: &'a [(PackageIdSpec, Dependency)],
    try_to_use: &'a HashSet<&'a PackageId>,
    cache: HashMap<Dependency, Rc<Vec<Candidate>>>,
    // If set the list of dependency candidates will be sorted by minimal
    // versions first. That allows `cargo update -Z minimal-versions` which will
    // specify minimum dependency versions to be used.
    minimal_versions: bool,
}

impl<'a> RegistryQueryer<'a> {
    pub fn new(
        registry: &'a mut Registry,
        replacements: &'a [(PackageIdSpec, Dependency)],
        try_to_use: &'a HashSet<&'a PackageId>,
        minimal_versions: bool,
    ) -> Self {
        RegistryQueryer {
            registry,
            replacements,
            cache: HashMap::new(),
            try_to_use,
            minimal_versions,
        }
    }

    /// Queries the `registry` to return a list of candidates for `dep`.
    ///
    /// This method is the location where overrides are taken into account. If
    /// any candidates are returned which match an override then the override is
    /// applied by performing a second query for what the override should
    /// return.
    pub fn query(&mut self, dep: &Dependency) -> CargoResult<Rc<Vec<Candidate>>> {
        if let Some(out) = self.cache.get(dep).cloned() {
            return Ok(out);
        }

        let mut ret = Vec::new();
        self.registry.query(
            dep,
            &mut |s| {
                ret.push(Candidate {
                    summary: s,
                    replace: None,
                });
            },
            false,
        )?;
        for candidate in ret.iter_mut() {
            let summary = &candidate.summary;

            let mut potential_matches = self
                .replacements
                .iter()
                .filter(|&&(ref spec, _)| spec.matches(summary.package_id()));

            let &(ref spec, ref dep) = match potential_matches.next() {
                None => continue,
                Some(replacement) => replacement,
            };
            debug!(
                "found an override for {} {}",
                dep.package_name(),
                dep.version_req()
            );

            let mut summaries = self.registry.query_vec(dep, false)?.into_iter();
            let s = summaries.next().ok_or_else(|| {
                format_err!(
                    "no matching package for override `{}` found\n\
                     location searched: {}\n\
                     version required: {}",
                    spec,
                    dep.source_id(),
                    dep.version_req()
                )
            })?;
            let summaries = summaries.collect::<Vec<_>>();
            if !summaries.is_empty() {
                let bullets = summaries
                    .iter()
                    .map(|s| format!("  * {}", s.package_id()))
                    .collect::<Vec<_>>();
                bail!(
                    "the replacement specification `{}` matched \
                     multiple packages:\n  * {}\n{}",
                    spec,
                    s.package_id(),
                    bullets.join("\n")
                );
            }

            // The dependency should be hard-coded to have the same name and an
            // exact version requirement, so both of these assertions should
            // never fail.
            assert_eq!(s.version(), summary.version());
            assert_eq!(s.name(), summary.name());

            let replace = if s.source_id() == summary.source_id() {
                debug!("Preventing\n{:?}\nfrom replacing\n{:?}", summary, s);
                None
            } else {
                Some(s)
            };
            let matched_spec = spec.clone();

            // Make sure no duplicates
            if let Some(&(ref spec, _)) = potential_matches.next() {
                bail!(
                    "overlapping replacement specifications found:\n\n  \
                     * {}\n  * {}\n\nboth specifications match: {}",
                    matched_spec,
                    spec,
                    summary.package_id()
                );
            }

            for dep in summary.dependencies() {
                debug!("\t{} => {}", dep.package_name(), dep.version_req());
            }

            candidate.replace = replace;
        }

        // When we attempt versions for a package we'll want to do so in a
        // sorted fashion to pick the "best candidates" first. Currently we try
        // prioritized summaries (those in `try_to_use`) and failing that we
        // list everything from the maximum version to the lowest version.
        ret.sort_unstable_by(|a, b| {
            let a_in_previous = self.try_to_use.contains(a.summary.package_id());
            let b_in_previous = self.try_to_use.contains(b.summary.package_id());
            let previous_cmp = a_in_previous.cmp(&b_in_previous).reverse();
            match previous_cmp {
                Ordering::Equal => {
                    let cmp = a.summary.version().cmp(b.summary.version());
                    if self.minimal_versions {
                        // Lower version ordered first.
                        cmp
                    } else {
                        // Higher version ordered first.
                        cmp.reverse()
                    }
                }
                _ => previous_cmp,
            }
        });

        let out = Rc::new(ret);

        self.cache.insert(dep.clone(), out.clone());

        Ok(out)
    }
}

#[derive(Clone, Copy)]
pub enum Method<'a> {
    Everything, // equivalent to Required { dev_deps: true, all_features: true, .. }
    Required {
        dev_deps: bool,
        features: &'a [InternedString],
        all_features: bool,
        uses_default_features: bool,
    },
}

impl<'r> Method<'r> {
    pub fn split_features(features: &[String]) -> Vec<InternedString> {
        features
            .iter()
            .flat_map(|s| s.split_whitespace())
            .flat_map(|s| s.split(','))
            .filter(|s| !s.is_empty())
            .map(|s| InternedString::new(s))
            .collect::<Vec<InternedString>>()
    }
}

#[derive(Clone)]
pub struct Candidate {
    pub summary: Summary,
    pub replace: Option<Summary>,
}

#[derive(Clone)]
pub struct DepsFrame {
    pub parent: Summary,
    pub just_for_error_messages: bool,
    pub remaining_siblings: RcVecIter<DepInfo>,
}

impl DepsFrame {
    /// Returns the least number of candidates that any of this frame's siblings
    /// has.
    ///
    /// The `remaining_siblings` array is already sorted with the smallest
    /// number of candidates at the front, so we just return the number of
    /// candidates in that entry.
    fn min_candidates(&self) -> usize {
        self.remaining_siblings
            .peek()
            .map(|(_, (_, candidates, _))| candidates.len())
            .unwrap_or(0)
    }

    pub fn flatten(&self) -> impl Iterator<Item = (&PackageId, Dependency)> {
        self.remaining_siblings
            .clone()
            .map(move |(_, (d, _, _))| (self.parent.package_id(), d))
    }
}

impl PartialEq for DepsFrame {
    fn eq(&self, other: &DepsFrame) -> bool {
        self.just_for_error_messages == other.just_for_error_messages
            && self.min_candidates() == other.min_candidates()
    }
}

impl Eq for DepsFrame {}

impl PartialOrd for DepsFrame {
    fn partial_cmp(&self, other: &DepsFrame) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for DepsFrame {
    fn cmp(&self, other: &DepsFrame) -> Ordering {
        self.just_for_error_messages
            .cmp(&other.just_for_error_messages)
            .reverse()
            .then_with(|| self.min_candidates().cmp(&other.min_candidates()))
    }
}

/// Note that a `OrdSet` is used for the remaining dependencies that need
/// activation. This set is sorted by how many candidates each dependency has.
///
/// This helps us get through super constrained portions of the dependency
/// graph quickly and hopefully lock down what later larger dependencies can
/// use (those with more candidates).
#[derive(Clone)]
pub struct RemainingDeps(usize, im_rc::OrdSet<(DepsFrame, usize)>);

impl RemainingDeps {
    pub fn new() -> RemainingDeps {
        RemainingDeps(0, im_rc::OrdSet::new())
    }
    pub fn push(&mut self, x: DepsFrame) {
        self.1.insert((x, self.0));
        self.0 += 1;
    }
    pub fn pop_most_constrained(&mut self) -> Option<(bool, (Summary, (usize, DepInfo)))> {
        while let Some((mut deps_frame, i)) = self.1.remove_min() {
            let just_here_for_the_error_messages = deps_frame.just_for_error_messages;

            // Figure out what our next dependency to activate is, and if nothing is
            // listed then we're entirely done with this frame (yay!) and we can
            // move on to the next frame.
            if let Some(sibling) = deps_frame.remaining_siblings.next() {
                let parent = Summary::clone(&deps_frame.parent);
                self.1.insert((deps_frame, i));
                return Some((just_here_for_the_error_messages, (parent, sibling)));
            }
        }
        None
    }
    pub fn iter(&mut self) -> impl Iterator<Item = (&PackageId, Dependency)> {
        self.1.iter().flat_map(|(other, _)| other.flatten())
    }
}

// Information about the dependencies for a crate, a tuple of:
//
// (dependency info, candidates, features activated)
pub type DepInfo = (Dependency, Rc<Vec<Candidate>>, Rc<Vec<InternedString>>);

/// All possible reasons that a package might fail to activate.
///
/// We maintain a list of conflicts for error reporting as well as backtracking
/// purposes. Each reason here is why candidates may be rejected or why we may
/// fail to resolve a dependency.
#[derive(Debug, Clone, PartialOrd, Ord, PartialEq, Eq)]
pub enum ConflictReason {
    /// There was a semver conflict, for example we tried to activate a package
    /// 1.0.2 but 1.1.0 was already activated (aka a compatible semver version
    /// is already activated)
    Semver,

    /// The `links` key is being violated. For example one crate in the
    /// dependency graph has `links = "foo"` but this crate also had that, and
    /// we're only allowed one per dependency graph.
    Links(InternedString),

    /// A dependency listed features that weren't actually available on the
    /// candidate. For example we tried to activate feature `foo` but the
    /// candidate we're activating didn't actually have the feature `foo`.
    MissingFeatures(String),
}

impl ConflictReason {
    pub fn is_links(&self) -> bool {
        if let ConflictReason::Links(_) = *self {
            return true;
        }
        false
    }

    pub fn is_missing_features(&self) -> bool {
        if let ConflictReason::MissingFeatures(_) = *self {
            return true;
        }
        false
    }
}

pub struct RcVecIter<T> {
    vec: Rc<Vec<T>>,
    rest: Range<usize>,
}

impl<T> RcVecIter<T> {
    pub fn new(vec: Rc<Vec<T>>) -> RcVecIter<T> {
        RcVecIter {
            rest: 0..vec.len(),
            vec,
        }
    }

    fn peek(&self) -> Option<(usize, &T)> {
        self.rest
            .clone()
            .next()
            .and_then(|i| self.vec.get(i).map(|val| (i, &*val)))
    }
}

// Not derived to avoid `T: Clone`
impl<T> Clone for RcVecIter<T> {
    fn clone(&self) -> RcVecIter<T> {
        RcVecIter {
            vec: self.vec.clone(),
            rest: self.rest.clone(),
        }
    }
}

impl<T> Iterator for RcVecIter<T>
where
    T: Clone,
{
    type Item = (usize, T);

    fn next(&mut self) -> Option<Self::Item> {
        self.rest
            .next()
            .and_then(|i| self.vec.get(i).map(|val| (i, val.clone())))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        // rest is a std::ops::Range, which is an ExactSizeIterator.
        self.rest.size_hint()
    }
}

impl<T: Clone> ExactSizeIterator for RcVecIter<T> {}

pub struct RcList<T> {
    pub head: Option<Rc<(T, RcList<T>)>>,
}

impl<T> RcList<T> {
    pub fn new() -> RcList<T> {
        RcList { head: None }
    }

    pub fn push(&mut self, data: T) {
        let node = Rc::new((
            data,
            RcList {
                head: self.head.take(),
            },
        ));
        self.head = Some(node);
    }
}

// Not derived to avoid `T: Clone`
impl<T> Clone for RcList<T> {
    fn clone(&self) -> RcList<T> {
        RcList {
            head: self.head.clone(),
        }
    }
}

// Avoid stack overflows on drop by turning recursion into a loop
impl<T> Drop for RcList<T> {
    fn drop(&mut self) {
        let mut cur = self.head.take();
        while let Some(head) = cur {
            match Rc::try_unwrap(head) {
                Ok((_data, mut next)) => cur = next.head.take(),
                Err(_) => break,
            }
        }
    }
}

pub enum GraphNode {
    Add(PackageId),
    Link(PackageId, PackageId, Dependency),
}
<|endoftext|>


<commit_before>// pest. The Elegant Parser
// Copyright (c) 2018 Dragoș Tiselice
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

/// Implementation of a `Stack` which maintains an log of `StackOp`s in order to rewind the stack
/// to a previous state.
#[derive(Debug)]
pub struct Stack<T: Clone> {
    ops: Vec<StackOp<T>>,
    cache: Vec<T>,
    snapshots: Vec<usize>
}

impl<T: Clone> Stack<T> {
    /// Creates a new `Stack`.
    pub fn new() -> Self {
        Stack {
            ops: vec![],
            cache: vec![],
            snapshots: vec![]
        }
    }

    /// Returns `true` if the stack is currently empty.
    #[allow(dead_code)]
    pub fn is_empty(&self) -> bool {
        self.cache.is_empty()
    }

    /// Returns the top-most `&T` in the `Stack`.
    pub fn peek(&self) -> Option<&T> {
        self.cache.last()
    }

    /// Pushes a `T` onto the `Stack`.
    pub fn push(&mut self, elem: T) {
        self.ops.push(StackOp::Push(elem.clone()));
        self.cache.push(elem);
    }

    /// Pops the top-most `T` from the `Stack`.
    pub fn pop(&mut self) -> Option<T> {
        let popped = self.cache.pop();
        if let Some(ref val) = popped {
            self.ops.push(StackOp::Pop(val.clone()));
        }
        popped
    }

    /// Returns an iterator to the current state of the cache in fifo order.
    pub fn iter(&self) -> impl Iterator<Item = &T> {
        self.cache.iter().rev()
    }

    /// Takes a snapshot of the current `Stack`.
    pub fn snapshot(&mut self) {
        let ops_index = self.ops.len();

        if ops_index > self.most_recent_snap() {
            self.snapshots.push(ops_index);
        }
    }

    /// Rewinds the `Stack` to the most recent `snapshot()`. If no `snapshot()` has been taken, this
    /// function return the stack to its initial state.
    pub fn restore(&mut self) {
        match self.snapshots.pop() {
            Some(ops_index) => {
                self.rewind_to(ops_index);
                self.ops.truncate(ops_index);
            },
            None => {
                self.cache.clear();
                self.ops.clear();
            }
        }
    }

    // Rewind the stack to a particular index
    fn rewind_to(&mut self, index: usize) {
        let ops_to_rewind = &self.ops[index..];
        for op in ops_to_rewind.iter().rev() {
            match *op {
                StackOp::Push(_) => {
                    self.cache.pop();
                }
                StackOp::Pop(ref elem) => {
                    self.cache.push(elem.clone());
                }
            }
        }
    }

    fn most_recent_snap(&self) -> usize {
        *self.snapshots.last().unwrap_or(&0)
    }
}

#[derive(Debug)]
enum StackOp<T> {
    Push(T),
    Pop(T)
}

#[cfg(test)]
mod test {
    use super::Stack;

    #[test]
    fn snapshot_with_empty() {
        let mut stack = Stack::new();

        stack.snapshot();
        // []
        assert!(stack.is_empty());
        // [0]
        stack.push(0);
        stack.restore();
        assert!(stack.is_empty());
    }

    #[test]
    fn iter() {
        let mut stack = Stack::new();

        stack.push(0);
        stack.push(1);
        stack.push(2);

        assert_eq!(stack.iter().collect::<Vec<&i32>>(), vec![&2, &1, &0]);
    }

    #[test]
    fn stack_ops() {
        let mut stack = Stack::new();

        // []
        assert!(stack.is_empty());
        assert_eq!(stack.peek(), None);
        assert_eq!(stack.pop(), None);

        // [0]
        stack.push(0);
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&0));

        // [0, 1]
        stack.push(1);
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&1));

        // [0]
        assert_eq!(stack.pop(), Some(1));
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&0));

        // [0, 2]
        stack.push(2);
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&2));

        // [0, 2, 3]
        stack.push(3);
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&3));

        // Take a snapshot of the current stack
        // [0, 2, 3]
        stack.snapshot();

        // [0, 2]
        assert_eq!(stack.pop(), Some(3));
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&2));

        // Take a snapshot of the current stack
        // [0, 2]
        stack.snapshot();

        // [0]
        assert_eq!(stack.pop(), Some(2));
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&0));

        // []
        assert_eq!(stack.pop(), Some(0));
        assert!(stack.is_empty());

        // Test backtracking
        // [0, 2]
        stack.restore();
        assert_eq!(stack.pop(), Some(2));
        assert_eq!(stack.pop(), Some(0));
        assert_eq!(stack.pop(), None);

        // Test backtracking
        // [0, 2, 3]
        stack.restore();
        assert_eq!(stack.pop(), Some(3));
        assert_eq!(stack.pop(), Some(2));
        assert_eq!(stack.pop(), Some(0));
        assert_eq!(stack.pop(), None);
    }
}
<commit_msg>Fixed nested Stack snapshotting.<commit_after>// pest. The Elegant Parser
// Copyright (c) 2018 Dragoș Tiselice
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

/// Implementation of a `Stack` which maintains an log of `StackOp`s in order to rewind the stack
/// to a previous state.
#[derive(Debug)]
pub struct Stack<T: Clone> {
    ops: Vec<StackOp<T>>,
    cache: Vec<T>,
    snapshots: Vec<usize>
}

impl<T: Clone> Stack<T> {
    /// Creates a new `Stack`.
    pub fn new() -> Self {
        Stack {
            ops: vec![],
            cache: vec![],
            snapshots: vec![]
        }
    }

    /// Returns `true` if the stack is currently empty.
    #[allow(dead_code)]
    pub fn is_empty(&self) -> bool {
        self.cache.is_empty()
    }

    /// Returns the top-most `&T` in the `Stack`.
    pub fn peek(&self) -> Option<&T> {
        self.cache.last()
    }

    /// Pushes a `T` onto the `Stack`.
    pub fn push(&mut self, elem: T) {
        self.ops.push(StackOp::Push(elem.clone()));
        self.cache.push(elem);
    }

    /// Pops the top-most `T` from the `Stack`.
    pub fn pop(&mut self) -> Option<T> {
        let popped = self.cache.pop();
        if let Some(ref val) = popped {
            self.ops.push(StackOp::Pop(val.clone()));
        }
        popped
    }

    /// Returns an iterator to the current state of the cache in fifo order.
    pub fn iter(&self) -> impl Iterator<Item = &T> {
        self.cache.iter().rev()
    }

    /// Takes a snapshot of the current `Stack`.
    pub fn snapshot(&mut self) {
        self.snapshots.push(self.ops.len());
    }

    /// Rewinds the `Stack` to the most recent `snapshot()`. If no `snapshot()` has been taken, this
    /// function return the stack to its initial state.
    pub fn restore(&mut self) {
        match self.snapshots.pop() {
            Some(ops_index) => {
                self.rewind_to(ops_index);
                self.ops.truncate(ops_index);
            },
            None => {
                self.cache.clear();
                self.ops.clear();
            }
        }
    }

    // Rewind the stack to a particular index
    fn rewind_to(&mut self, index: usize) {
        let ops_to_rewind = &self.ops[index..];
        for op in ops_to_rewind.iter().rev() {
            match *op {
                StackOp::Push(_) => {
                    self.cache.pop();
                }
                StackOp::Pop(ref elem) => {
                    self.cache.push(elem.clone());
                }
            }
        }
    }
}

#[derive(Debug)]
enum StackOp<T> {
    Push(T),
    Pop(T)
}

#[cfg(test)]
mod test {
    use super::Stack;

    #[test]
    fn snapshot_with_empty() {
        let mut stack = Stack::new();

        stack.snapshot();
        // []
        assert!(stack.is_empty());
        // [0]
        stack.push(0);
        stack.restore();
        assert!(stack.is_empty());
    }

    #[test]
    fn snapshot_twice() {
        let mut stack = Stack::new();

        stack.push(0);

        stack.snapshot();
        stack.snapshot();
        stack.restore();
        stack.restore();

        assert_eq!(stack.iter().collect::<Vec<&i32>>(), vec![&0]);
    }

    #[test]
    fn iter() {
        let mut stack = Stack::new();

        stack.push(0);
        stack.push(1);
        stack.push(2);

        assert_eq!(stack.iter().collect::<Vec<&i32>>(), vec![&2, &1, &0]);
    }

    #[test]
    fn stack_ops() {
        let mut stack = Stack::new();

        // []
        assert!(stack.is_empty());
        assert_eq!(stack.peek(), None);
        assert_eq!(stack.pop(), None);

        // [0]
        stack.push(0);
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&0));

        // [0, 1]
        stack.push(1);
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&1));

        // [0]
        assert_eq!(stack.pop(), Some(1));
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&0));

        // [0, 2]
        stack.push(2);
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&2));

        // [0, 2, 3]
        stack.push(3);
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&3));

        // Take a snapshot of the current stack
        // [0, 2, 3]
        stack.snapshot();

        // [0, 2]
        assert_eq!(stack.pop(), Some(3));
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&2));

        // Take a snapshot of the current stack
        // [0, 2]
        stack.snapshot();

        // [0]
        assert_eq!(stack.pop(), Some(2));
        assert!(!stack.is_empty());
        assert_eq!(stack.peek(), Some(&0));

        // []
        assert_eq!(stack.pop(), Some(0));
        assert!(stack.is_empty());

        // Test backtracking
        // [0, 2]
        stack.restore();
        assert_eq!(stack.pop(), Some(2));
        assert_eq!(stack.pop(), Some(0));
        assert_eq!(stack.pop(), None);

        // Test backtracking
        // [0, 2, 3]
        stack.restore();
        assert_eq!(stack.pop(), Some(3));
        assert_eq!(stack.pop(), Some(2));
        assert_eq!(stack.pop(), Some(0));
        assert_eq!(stack.pop(), None);
    }
}
<|endoftext|>