<commit_before>#[macro_use]
extern crate abomonation;
extern crate timely;
extern crate differential_dataflow;
extern crate arrayvec;

use std::fs::File;
use std::io::{BufRead, BufReader};
use std::time::Instant;

use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::progress::timestamp::RootTimestamp;
use timely::progress::nested::product::Product;

use differential_dataflow::{Collection, AsCollection};

mod types;
mod queries;

use types::*;

fn main() {

    timely::execute_from_args(std::env::args().skip(4), |worker| {

        let index = worker.index();
        let peers = worker.peers();

        let prefix = ::std::env::args().nth(1).unwrap();;
        let logical_batch = ::std::env::args().nth(2).unwrap().parse::<usize>().unwrap();
        let physical_batch = ::std::env::args().nth(3).unwrap().parse::<usize>().unwrap();
        let query: usize = ::std::env::args().nth(4).unwrap().parse().unwrap();

        let (mut inputs, probes, used) = worker.dataflow::<usize,_,_>(move |scope| {

            // create new inputs to use in workers!
            let (cust_in, cust) = scope.new_input();
            let (line_in, line) = scope.new_input();
            let (nats_in, nats) = scope.new_input();
            let (ords_in, ords) = scope.new_input();
            let (part_in, part) = scope.new_input();
            let (psup_in, psup) = scope.new_input();
            let (regs_in, regs) = scope.new_input();
            let (supp_in, supp) = scope.new_input();

            let mut collections = Collections::new(
                cust.as_collection(),
                line.as_collection(),
                nats.as_collection(),
                ords.as_collection(),
                part.as_collection(),
                psup.as_collection(),
                regs.as_collection(),
                supp.as_collection(),
            );

            let mut probes = Vec::new();

            match query {
                1 => probes.push(queries::query01::query(&mut collections)),
                2 => probes.push(queries::query02::query(&mut collections)),
                3 => probes.push(queries::query03::query(&mut collections)),
                4 => probes.push(queries::query04::query(&mut collections)),
                5 => probes.push(queries::query05::query(&mut collections)),
                6 => probes.push(queries::query06::query(&mut collections)),
                7 => probes.push(queries::query07::query(&mut collections)),
                8 => probes.push(queries::query08::query(&mut collections)),
                9 => probes.push(queries::query09::query(&mut collections)),
                10 => probes.push(queries::query10::query(&mut collections)),
                11 => probes.push(queries::query11::query(&mut collections)),
                12 => probes.push(queries::query12::query(&mut collections)),
                13 => probes.push(queries::query13::query(&mut collections)),
                14 => probes.push(queries::query14::query(&mut collections)),
                15 => probes.push(queries::query15::query(&mut collections)),
                16 => probes.push(queries::query16::query(&mut collections)),
                17 => probes.push(queries::query17::query(&mut collections)),
                18 => probes.push(queries::query18::query(&mut collections)),
                19 => probes.push(queries::query19::query(&mut collections)),
                20 => probes.push(queries::query20::query(&mut collections)),
                21 => probes.push(queries::query21::query(&mut collections)),
                22 => probes.push(queries::query22::query(&mut collections)),
                _ => panic!("query: {:?} unimplemented", query),
            }

            // return the various input handles, and the list of probes.
            ((cust_in, line_in, nats_in, ords_in, part_in, psup_in, regs_in, supp_in), probes, collections.used)
        });

        // customer.tbl lineitem.tbl    nation.tbl  orders.tbl  part.tbl    partsupp.tbl    region.tbl  supplier.tbl
        let mut customers = if used[0] { load::<Customer>(prefix.as_str(), "customer.tbl", index, peers, logical_batch, physical_batch, 0) } else { Vec::new() };
        let mut lineitems = if used[1] { load::<LineItem>(prefix.as_str(), "lineitem.tbl", index, peers, logical_batch, physical_batch, 1) } else { Vec::new() };
        let mut nations = if used[2] { load::<Nation>(prefix.as_str(), "nation.tbl", index, peers, logical_batch, physical_batch, 2) } else { Vec::new() };
        let mut orders = if used[3] { load::<Order>(prefix.as_str(), "orders.tbl", index, peers, logical_batch, physical_batch, 3) } else { Vec::new() };
        let mut parts = if used[4] { load::<Part>(prefix.as_str(), "part.tbl", index, peers, logical_batch, physical_batch, 4) } else { Vec::new() };
        let mut partsupps = if used[5] { load::<PartSupp>(prefix.as_str(), "partsupp.tbl", index, peers, logical_batch, physical_batch, 5) } else { Vec::new() };
        let mut regions = if used[6] { load::<Region>(prefix.as_str(), "region.tbl", index, peers, logical_batch, physical_batch, 6) } else { Vec::new() };
        let mut suppliers = if used[7] { load::<Supplier>(prefix.as_str(), "supplier.tbl", index, peers, logical_batch, physical_batch, 7) } else { Vec::new() };

        let mut tuples = 0usize;
        tuples += customers.iter().map(|x| x.len()).sum::<usize>();
        tuples += lineitems.iter().map(|x| x.len()).sum::<usize>();
        tuples += nations.iter().map(|x| x.len()).sum::<usize>();
        tuples += orders.iter().map(|x| x.len()).sum::<usize>();
        tuples += parts.iter().map(|x| x.len()).sum::<usize>();
        tuples += partsupps.iter().map(|x| x.len()).sum::<usize>();
        tuples += regions.iter().map(|x| x.len()).sum::<usize>();
        tuples += suppliers.iter().map(|x| x.len()).sum::<usize>();

        // Synchronize before starting the timer.
        let next_round = 1;
        inputs.0.advance_to(next_round);
        inputs.1.advance_to(next_round);
        inputs.2.advance_to(next_round);
        inputs.3.advance_to(next_round);
        inputs.4.advance_to(next_round);
        inputs.5.advance_to(next_round);
        inputs.6.advance_to(next_round);
        inputs.7.advance_to(next_round);
        let time = inputs.0.time().clone();
        worker.step_while(|| probes.iter().all(|p| p.less_than(&time)));

        let timer = Instant::now();
        let mut round = 0;
        while customers.len() > 0 || lineitems.len() > 0 || nations.len() > 0 || orders.len() > 0 || parts.len() > 0 || partsupps.len() > 0 || regions.len() > 0 || suppliers.len() > 0 {

            // introduce physical batch of data for each input with remaining data.
            if let Some(mut data) = customers.pop() { inputs.0.send_batch(&mut data); }
            if let Some(mut data) = lineitems.pop() { inputs.1.send_batch(&mut data); }
            if let Some(mut data) = nations.pop() { inputs.2.send_batch(&mut data); }
            if let Some(mut data) = orders.pop() { inputs.3.send_batch(&mut data); }
            if let Some(mut data) = parts.pop() { inputs.4.send_batch(&mut data); }
            if let Some(mut data) = partsupps.pop() { inputs.5.send_batch(&mut data); }
            if let Some(mut data) = regions.pop() { inputs.6.send_batch(&mut data); }
            if let Some(mut data) = suppliers.pop() { inputs.7.send_batch(&mut data); }

            // catch all inputs up to the same (next) round.
            let next_round = 1 + 8 * (round + 1) * physical_batch;
            inputs.0.advance_to(next_round);
            inputs.1.advance_to(next_round);
            inputs.2.advance_to(next_round);
            inputs.3.advance_to(next_round);
            inputs.4.advance_to(next_round);
            inputs.5.advance_to(next_round);
            inputs.6.advance_to(next_round);
            inputs.7.advance_to(next_round);

            let time = inputs.0.time().clone();
            worker.step_while(|| probes.iter().all(|p| p.less_than(&time)));
            round += 1;
        }

        // Finish outstanding work before stopping the timer.
        let next_round = usize::max_value();
        inputs.0.advance_to(next_round);
        inputs.1.advance_to(next_round);
        inputs.2.advance_to(next_round);
        inputs.3.advance_to(next_round);
        inputs.4.advance_to(next_round);
        inputs.5.advance_to(next_round);
        inputs.6.advance_to(next_round);
        inputs.7.advance_to(next_round);
        let time = inputs.0.time().clone();
        worker.step_while(|| probes.iter().all(|p| p.less_than(&time)));

        let query_name = if query < 10 { format!("q0{}", query) } else { format!("q{}", query) };
        let elapsed = timer.elapsed();
        let nanos = elapsed.as_secs() * 1000000000 + elapsed.subsec_nanos() as u64;
        if index == 0 {
            println!("query: {}, elapsed: {:?}, tuples: {:?}, rate: {:?}", query_name, timer.elapsed(), peers * tuples, ((peers * tuples) as f64) / (nanos as f64 / 1000000000.0));
        }

    }).unwrap();
}

pub struct Collections<G: Scope> {
    customers: Collection<G, Customer, isize>,
    lineitems: Collection<G, LineItem, isize>,
    nations: Collection<G, Nation, isize>,
    orders: Collection<G, Order, isize>,
    parts: Collection<G, Part, isize>,
    partsupps: Collection<G, PartSupp, isize>,
    regions: Collection<G, Region, isize>,
    suppliers: Collection<G, Supplier, isize>,
    used: [bool; 8],
}

impl<G: Scope> Collections<G> {
    fn new(
        customers: Collection<G, Customer, isize>,
        lineitems: Collection<G, LineItem, isize>,
        nations: Collection<G, Nation, isize>,
        orders: Collection<G, Order, isize>,
        parts: Collection<G, Part, isize>,
        partsupps: Collection<G, PartSupp, isize>,
        regions: Collection<G, Region, isize>,
        suppliers: Collection<G, Supplier, isize>,
    ) -> Self {

        Collections {
            customers: customers,
            lineitems: lineitems,
            nations: nations,
            orders: orders,
            parts: parts,
            partsupps: partsupps,
            regions: regions,
            suppliers: suppliers,
            used: [false; 8]
        }
    }

    fn customers(&mut self) -> &Collection<G, Customer, isize> { self.used[0] = true; &self.customers }
    fn lineitems(&mut self) -> &Collection<G, LineItem, isize> { self.used[1] = true; &self.lineitems }
    fn nations(&mut self) -> &Collection<G, Nation, isize> { self.used[2] = true; &self.nations }
    fn orders(&mut self) -> &Collection<G, Order, isize> { self.used[3] = true; &self.orders }
    fn parts(&mut self) -> &Collection<G, Part, isize> { self.used[4] = true; &self.parts }
    fn partsupps(&mut self) -> &Collection<G, PartSupp, isize> { self.used[5] = true; &self.partsupps }
    fn regions(&mut self) -> &Collection<G, Region, isize> { self.used[6] = true; &self.regions }
    fn suppliers(&mut self) -> &Collection<G, Supplier, isize> { self.used[7] = true; &self.suppliers }
}

// Returns a sequence of physical batches of ready-to-go timestamped data. Not clear that `input` can exploit the pre-arrangement yet.
fn load<T>(prefix: &str, name: &str, index: usize, peers: usize, logical_batch: usize, physical_batch: usize, off: usize)
    -> Vec<Vec<(T, Product<RootTimestamp, usize>, isize)>>
where T: for<'a> From<&'a str> {

    let mut result = Vec::new();
    let mut buffer = Vec::new();

    let path = format!("{}{}", prefix, name);

    let items_file = File::open(&path).expect("didn't find items file");
    let mut items_reader =  BufReader::new(items_file);
    let mut count = 0;

    let mut line = String::new();

    while items_reader.read_line(&mut line).unwrap() > 0 {

        if count % peers == index {

            let logical = (8 * count / logical_batch) + off;
            let physical = logical / physical_batch;
            let round = physical / 8;

            while result.len() < round {
                result.push(buffer);
                buffer = Vec::with_capacity(2 + logical_batch * physical_batch / peers);
            }

            let item = T::from(line.as_str());
            buffer.push((item, RootTimestamp::new(logical + 1), 1));
        }

        count += 1;

        line.clear();
    }

    if buffer.len() > 0 {
        result.push(buffer);
    }

    result.reverse();
    result
}
<commit_msg>seal regions, nations<commit_after>#[macro_use]
extern crate abomonation;
extern crate timely;
extern crate differential_dataflow;
extern crate arrayvec;

use std::fs::File;
use std::io::{BufRead, BufReader};
use std::time::Instant;

use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::progress::timestamp::RootTimestamp;
use timely::progress::nested::product::Product;

use differential_dataflow::{Collection, AsCollection};

mod types;
mod queries;

use types::*;

fn main() {

    timely::execute_from_args(std::env::args().skip(4), |worker| {

        let index = worker.index();
        let peers = worker.peers();

        let prefix = ::std::env::args().nth(1).unwrap();;
        let logical_batch = ::std::env::args().nth(2).unwrap().parse::<usize>().unwrap();
        let physical_batch = ::std::env::args().nth(3).unwrap().parse::<usize>().unwrap();
        let query: usize = ::std::env::args().nth(4).unwrap().parse().unwrap();
        let seal: bool = ::std::env::args().any(|x| x == "seal-inputs");

        let (mut inputs, probes, used) = worker.dataflow::<usize,_,_>(move |scope| {

            // create new inputs to use in workers!
            let (cust_in, cust) = scope.new_input();
            let (line_in, line) = scope.new_input();
            let (nats_in, nats) = scope.new_input();
            let (ords_in, ords) = scope.new_input();
            let (part_in, part) = scope.new_input();
            let (psup_in, psup) = scope.new_input();
            let (regs_in, regs) = scope.new_input();
            let (supp_in, supp) = scope.new_input();

            let mut collections = Collections::new(
                cust.as_collection(),
                line.as_collection(),
                nats.as_collection(),
                ords.as_collection(),
                part.as_collection(),
                psup.as_collection(),
                regs.as_collection(),
                supp.as_collection(),
            );

            let mut probes = Vec::new();

            match query {
                // 1 => probes.push(queries::query01::query(&mut collections)),
                // 2 => probes.push(queries::query02::query(&mut collections)),
                // 3 => probes.push(queries::query03::query(&mut collections)),
                // 4 => probes.push(queries::query04::query(&mut collections)),
                // 5 => probes.push(queries::query05::query(&mut collections)),
                // 6 => probes.push(queries::query06::query(&mut collections)),
                // 7 => probes.push(queries::query07::query(&mut collections)),
                // 8 => probes.push(queries::query08::query(&mut collections)),
                9 => probes.push(queries::query09::query(&mut collections)),
                // 10 => probes.push(queries::query10::query(&mut collections)),
                // 11 => probes.push(queries::query11::query(&mut collections)),
                // 12 => probes.push(queries::query12::query(&mut collections)),
                // 13 => probes.push(queries::query13::query(&mut collections)),
                // 14 => probes.push(queries::query14::query(&mut collections)),
                // 15 => probes.push(queries::query15::query(&mut collections)),
                // 16 => probes.push(queries::query16::query(&mut collections)),
                // 17 => probes.push(queries::query17::query(&mut collections)),
                // 18 => probes.push(queries::query18::query(&mut collections)),
                // 19 => probes.push(queries::query19::query(&mut collections)),
                // 20 => probes.push(queries::query20::query(&mut collections)),
                // 21 => probes.push(queries::query21::query(&mut collections)),
                // 22 => probes.push(queries::query22::query(&mut collections)),
                _ => panic!("query: {:?} unimplemented", query),
            }

            // return the various input handles, and the list of probes.
            ((cust_in, line_in, Some(nats_in), ords_in, part_in, psup_in, Some(regs_in), supp_in), probes, collections.used)
        });

        // customer.tbl lineitem.tbl    nation.tbl  orders.tbl  part.tbl    partsupp.tbl    region.tbl  supplier.tbl
        let mut customers = if used[0] { load::<Customer>(prefix.as_str(), "customer.tbl", index, peers, logical_batch, physical_batch, 0) } else { Vec::new() };
        let mut lineitems = if used[1] { load::<LineItem>(prefix.as_str(), "lineitem.tbl", index, peers, logical_batch, physical_batch, 1) } else { Vec::new() };
        let mut nations = if used[2] { load::<Nation>(prefix.as_str(), "nation.tbl", index, peers, logical_batch, physical_batch, 2) } else { Vec::new() };
        let mut orders = if used[3] { load::<Order>(prefix.as_str(), "orders.tbl", index, peers, logical_batch, physical_batch, 3) } else { Vec::new() };
        let mut parts = if used[4] { load::<Part>(prefix.as_str(), "part.tbl", index, peers, logical_batch, physical_batch, 4) } else { Vec::new() };
        let mut partsupps = if used[5] { load::<PartSupp>(prefix.as_str(), "partsupp.tbl", index, peers, logical_batch, physical_batch, 5) } else { Vec::new() };
        let mut regions = if used[6] { load::<Region>(prefix.as_str(), "region.tbl", index, peers, logical_batch, physical_batch, 6) } else { Vec::new() };
        let mut suppliers = if used[7] { load::<Supplier>(prefix.as_str(), "supplier.tbl", index, peers, logical_batch, physical_batch, 7) } else { Vec::new() };

        let mut tuples = 0usize;
        tuples += customers.iter().map(|x| x.len()).sum::<usize>();
        tuples += lineitems.iter().map(|x| x.len()).sum::<usize>();
        tuples += nations.iter().map(|x| x.len()).sum::<usize>();
        tuples += orders.iter().map(|x| x.len()).sum::<usize>();
        tuples += parts.iter().map(|x| x.len()).sum::<usize>();
        tuples += partsupps.iter().map(|x| x.len()).sum::<usize>();
        tuples += regions.iter().map(|x| x.len()).sum::<usize>();
        tuples += suppliers.iter().map(|x| x.len()).sum::<usize>();

        // Synchronize before starting the timer.
        let next_round = 1;
        inputs.0.advance_to(next_round);
        inputs.1.advance_to(next_round);
        inputs.2.as_mut().map(|x| x.advance_to(next_round));
        inputs.3.advance_to(next_round);
        inputs.4.advance_to(next_round);
        inputs.5.advance_to(next_round);
        inputs.6.as_mut().map(|x| x.advance_to(next_round));
        inputs.7.advance_to(next_round);
        let time = inputs.0.time().clone();
        worker.step_while(|| probes.iter().all(|p| p.less_than(&time)));

        let timer = Instant::now();
        let mut round = 0;
        while customers.len() > 0 || lineitems.len() > 0 || nations.len() > 0 || orders.len() > 0 || parts.len() > 0 || partsupps.len() > 0 || regions.len() > 0 || suppliers.len() > 0 {

            // introduce physical batch of data for each input with remaining data.
            if let Some(mut data) = customers.pop() { inputs.0.send_batch(&mut data); }
            if let Some(mut data) = lineitems.pop() { inputs.1.send_batch(&mut data); }
            if let Some(mut data) = nations.pop() { inputs.2.as_mut().map(|x| x.send_batch(&mut data)); } else { if seal { inputs.2 = None; } }
            if let Some(mut data) = orders.pop() { inputs.3.send_batch(&mut data); }
            if let Some(mut data) = parts.pop() { inputs.4.send_batch(&mut data); }
            if let Some(mut data) = partsupps.pop() { inputs.5.send_batch(&mut data); }
            if let Some(mut data) = regions.pop() { inputs.6.as_mut().map(|x| x.send_batch(&mut data)); } else { if seal { inputs.6 = None; } }
            if let Some(mut data) = suppliers.pop() { inputs.7.send_batch(&mut data); }

            // catch all inputs up to the same (next) round.
            let next_round = 1 + 8 * (round + 1) * physical_batch;
            inputs.0.advance_to(next_round);
            inputs.1.advance_to(next_round);
            inputs.2.as_mut().map(|x| x.advance_to(next_round));
            inputs.3.advance_to(next_round);
            inputs.4.advance_to(next_round);
            inputs.5.advance_to(next_round);
            inputs.6.as_mut().map(|x| x.advance_to(next_round));
            inputs.7.advance_to(next_round);

            let time = inputs.0.time().clone();
            worker.step_while(|| probes.iter().all(|p| p.less_than(&time)));
            round += 1;
        }

        // Finish outstanding work before stopping the timer.
        let next_round = usize::max_value();
        inputs.0.advance_to(next_round);
        inputs.1.advance_to(next_round);
        inputs.2.as_mut().map(|x| x.advance_to(next_round));
        inputs.3.advance_to(next_round);
        inputs.4.advance_to(next_round);
        inputs.5.advance_to(next_round);
        inputs.6.as_mut().map(|x| x.advance_to(next_round));
        inputs.7.advance_to(next_round);
        let time = inputs.0.time().clone();
        worker.step_while(|| probes.iter().all(|p| p.less_than(&time)));

        let query_name = if query < 10 { format!("q0{}", query) } else { format!("q{}", query) };
        let elapsed = timer.elapsed();
        let nanos = elapsed.as_secs() * 1000000000 + elapsed.subsec_nanos() as u64;
        if index == 0 {
            println!("query: {}, elapsed: {:?}, tuples: {:?}, rate: {:?}", query_name, timer.elapsed(), peers * tuples, ((peers * tuples) as f64) / (nanos as f64 / 1000000000.0));
        }

    }).unwrap();
}

pub struct Collections<G: Scope> {
    customers: Collection<G, Customer, isize>,
    lineitems: Collection<G, LineItem, isize>,
    nations: Collection<G, Nation, isize>,
    orders: Collection<G, Order, isize>,
    parts: Collection<G, Part, isize>,
    partsupps: Collection<G, PartSupp, isize>,
    regions: Collection<G, Region, isize>,
    suppliers: Collection<G, Supplier, isize>,
    used: [bool; 8],
}

impl<G: Scope> Collections<G> {
    fn new(
        customers: Collection<G, Customer, isize>,
        lineitems: Collection<G, LineItem, isize>,
        nations: Collection<G, Nation, isize>,
        orders: Collection<G, Order, isize>,
        parts: Collection<G, Part, isize>,
        partsupps: Collection<G, PartSupp, isize>,
        regions: Collection<G, Region, isize>,
        suppliers: Collection<G, Supplier, isize>,
    ) -> Self {

        Collections {
            customers: customers,
            lineitems: lineitems,
            nations: nations,
            orders: orders,
            parts: parts,
            partsupps: partsupps,
            regions: regions,
            suppliers: suppliers,
            used: [false; 8]
        }
    }

    fn customers(&mut self) -> &Collection<G, Customer, isize> { self.used[0] = true; &self.customers }
    fn lineitems(&mut self) -> &Collection<G, LineItem, isize> { self.used[1] = true; &self.lineitems }
    fn nations(&mut self) -> &Collection<G, Nation, isize> { self.used[2] = true; &self.nations }
    fn orders(&mut self) -> &Collection<G, Order, isize> { self.used[3] = true; &self.orders }
    fn parts(&mut self) -> &Collection<G, Part, isize> { self.used[4] = true; &self.parts }
    fn partsupps(&mut self) -> &Collection<G, PartSupp, isize> { self.used[5] = true; &self.partsupps }
    fn regions(&mut self) -> &Collection<G, Region, isize> { self.used[6] = true; &self.regions }
    fn suppliers(&mut self) -> &Collection<G, Supplier, isize> { self.used[7] = true; &self.suppliers }
}

// Returns a sequence of physical batches of ready-to-go timestamped data. Not clear that `input` can exploit the pre-arrangement yet.
fn load<T>(prefix: &str, name: &str, index: usize, peers: usize, logical_batch: usize, physical_batch: usize, off: usize)
    -> Vec<Vec<(T, Product<RootTimestamp, usize>, isize)>>
where T: for<'a> From<&'a str> {

    let mut result = Vec::new();
    let mut buffer = Vec::new();

    let path = format!("{}{}", prefix, name);

    let items_file = File::open(&path).expect("didn't find items file");
    let mut items_reader =  BufReader::new(items_file);
    let mut count = 0;

    let mut line = String::new();

    while items_reader.read_line(&mut line).unwrap() > 0 {

        if count % peers == index {

            let logical = (8 * count / logical_batch) + off;
            let physical = logical / physical_batch;
            let round = physical / 8;

            while result.len() < round {
                result.push(buffer);
                buffer = Vec::with_capacity(2 + logical_batch * physical_batch / peers);
            }

            let item = T::from(line.as_str());
            buffer.push((item, RootTimestamp::new(logical + 1), 1));
        }

        count += 1;

        line.clear();
    }

    if buffer.len() > 0 {
        result.push(buffer);
    }

    result.reverse();
    result
}
<|endoftext|>


<commit_before>// Copyright (C) 2019 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use failure::format_err;
use std::{
    fs,
    io::{BufRead, BufReader},
    path::PathBuf,
};

pub(crate) use ffi::is_nand;

pub(crate) fn target_device_from_ubi_volume_name(volume: &str) -> Result<PathBuf, failure::Error> {
    let re = regex::Regex::new(r"^Volume ID:   (\d) \(on ubi(\d)\)$").unwrap();
    let path = fs::read_dir("/dev")?
        .filter(|entry| entry.is_ok())
        .map(|entry| format!("{:?}", entry.unwrap().path()))
        .find(|path| path.starts_with("ubi"))
        .ok_or_else(|| format_err!("Unable to find coorespoing ubi volume"))?;

    let dev_number = path.replace("ubi", "");

    let output = easy_process::run(&format!("ubinfo -d {} -N {}", dev_number, volume))?;
    let line = output
        .stdout
        .lines()
        .next()
        .ok_or_else(|| format_err!("Unable to read first line of ubinfo"))?;

    let re_match = re
        .captures(line)
        .ok_or_else(|| format_err!("Unable to extract any matches for Volume ID"))?;

    Ok(PathBuf::from(format!(
        "/dev/ubi{}_{}",
        dev_number, &re_match[0]
    )))
}

pub(crate) fn target_device_from_mtd_name(name: &str) -> Result<PathBuf, failure::Error> {
    let re =
        regex::Regex::new(r#"^(?P<dev>mtd\d): ([[:xdigit:]]+) ([[:xdigit:]]+) "(?P<name>.*)"$"#)
            .unwrap();
    let proc = fs::File::open("/proc/mtd")?;

    BufReader::new(proc)
        .lines()
        .filter_map(Result::ok)
        .find_map(|line| {
            re.captures(&line).and_then(|re_match| {
                let re_dev = re_match.name("dev").unwrap().as_str();
                let re_name = re_match.name("name").unwrap().as_str();
                if re_name == name {
                    Some(PathBuf::from(format!("/dev/{}", re_dev)))
                } else {
                    None
                }
            })
        })
        .ok_or_else(|| format_err!("Unable to find match for mtd device: {}", name))
}

mod ffi {
    use nix::ioctl_read;
    use std::{mem::MaybeUninit, os::unix::io::AsRawFd, path::Path};

    // From https://github.com/torvalds/linux/blob/master/include/uapi/mtd/mtd-abi.h
    const MTD_NANDFLASH: u8 = 4;
    const MTD_MLCNANDFLASH: u8 = 8;
    const MEMGETINFO: u8 = b'M';
    const MEMGETINFO_MODE: u8 = 1;

    #[repr(C)]
    pub struct mtd_info_user {
        kind: u8,
        flags: u32,
        size: u32,
        erasesize: u32,
        writesize: u32,
        oobsize: u32,
        padding: u64,
    }

    ioctl_read!(mtd_get_info, MEMGETINFO, MEMGETINFO_MODE, mtd_info_user);

    pub fn is_nand(device: &Path) -> Result<bool, failure::Error> {
        let device = std::fs::File::open(device)?;
        let info = unsafe {
            let mut info = MaybeUninit::<mtd_info_user>::uninit();
            mtd_get_info(device.as_raw_fd(), info.as_mut_ptr())?;
            info.assume_init()
        };

        Ok(info.kind == MTD_NANDFLASH || info.kind == MTD_MLCNANDFLASH)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use lazy_static::lazy_static;
    use pretty_assertions::assert_eq;
    use std::sync::{Arc, Mutex};

    pub(crate) struct FakeMtd {
        pub(crate) devices: Vec<PathBuf>,
        pub(crate) kind: MtdKind,
    }

    pub(crate) enum MtdKind {
        Nand,
        Nor,
    }

    impl FakeMtd {
        pub(crate) fn new(names: &[&str], kind: MtdKind) -> Result<FakeMtd, failure::Error> {
            match kind {
                MtdKind::Nand => {
                    easy_process::run("modprobe nandsim second_id_byte=0x36")?;
                    easy_process::run("mtdpart del /dev/mtd0 1")?;
                }
                MtdKind::Nor => {
                    easy_process::run("modprobe mtdram total_size=1000 erase_size=10")?;
                }
            }
            let total_size = 1000;

            // FakeMtd created here so if any subsequent command fails the drop will still
            // be called to cleanup mtd devices
            let mut mtd = FakeMtd {
                devices: vec![],
                kind,
            };
            let size = total_size / names.len();
            names.iter().enumerate().try_for_each(|(i, n)| {
                easy_process::run(&format!(
                    "mtdpart add /dev/mtd0 {} {} {}",
                    n,
                    i * size,
                    size
                ))
                .map(|_| {
                    mtd.devices
                        .push(PathBuf::from(format!("/dev/mtd{}", i + 1)))
                })
            })?;

            Ok(mtd)
        }
    }

    impl Drop for FakeMtd {
        fn drop(&mut self) {
            let module = match self.kind {
                MtdKind::Nand => "nandsim",
                MtdKind::Nor => "mtdram",
            };

            // Sleep time for nandsim to sync and avoid errors
            std::thread::sleep(std::time::Duration::from_millis(500));

            if let Err(e) = easy_process::run(&format!("rmmod {}", module)) {
                eprintln!("Failed to cleanup FakeMtd, Error: {}", e);
            }
        }
    }

    // Used to serialize access to MTD devices
    lazy_static! {
        pub static ref SERIALIZE: Arc<Mutex<()>> = Arc::new(Mutex::default());
    }

    #[test]
    #[ignore]
    fn device_from_mtd_name_nor() -> Result<(), failure::Error> {
        let _lock = SERIALIZE.lock();
        let dev_names = vec!["system0", "system1"];

        let mtd = FakeMtd::new(&dev_names, MtdKind::Nor).unwrap();

        assert_eq!(
            dev_names
                .into_iter()
                .map(target_device_from_mtd_name)
                .map(Result::unwrap)
                .collect::<Vec<_>>(),
            mtd.devices,
        );
        assert!(target_device_from_mtd_name("some_inexistent_device").is_err());

        Ok(())
    }

    #[test]
    #[ignore]
    fn device_from_mtd_name_nand() -> Result<(), failure::Error> {
        let _lock = SERIALIZE.lock();
        let dev_names = vec!["system0", "system1"];

        let mtd = FakeMtd::new(&dev_names, MtdKind::Nand).unwrap();

        assert_eq!(
            dev_names
                .into_iter()
                .map(target_device_from_mtd_name)
                .map(Result::unwrap)
                .collect::<Vec<_>>(),
            mtd.devices,
        );
        assert!(target_device_from_mtd_name("some_inexistent_device").is_err());

        Ok(())
    }

    #[test]
    #[ignore]
    fn test_is_nand() -> Result<(), failure::Error> {
        let _lock = SERIALIZE.lock();
        let dev_names = vec!["system0"];

        {
            let mtd = FakeMtd::new(&dev_names, MtdKind::Nand).unwrap();
            assert_eq!(is_nand(&mtd.devices[0]).unwrap(), true);
        }
        {
            let mtd = FakeMtd::new(&dev_names, MtdKind::Nor).unwrap();
            assert_eq!(is_nand(&mtd.devices[0]).unwrap(), false);
        }

        Ok(())
    }
}
<commit_msg>updatehub: target_device_from_ubi_volume_name fix and test<commit_after>// Copyright (C) 2019 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use failure::format_err;
use std::{
    fs,
    io::{BufRead, BufReader},
    path::PathBuf,
};

pub(crate) use ffi::is_nand;

pub(crate) fn target_device_from_ubi_volume_name(volume: &str) -> Result<PathBuf, failure::Error> {
    let re = regex::Regex::new(r"^Volume ID:   (?P<volume>\d+) \(on ubi(\d+)\)$").unwrap();
    walkdir::WalkDir::new("/dev")
        .min_depth(1)
        .into_iter()
        .filter_entry(|p| {
            p.file_name()
                .to_str()
                .map(|n| n.starts_with("ubi") && !n.contains('_'))
                .unwrap_or(false)
        })
        .filter_map(Result::ok)
        .find_map(|entry| {
            let path = entry.path();
            let output =
                easy_process::run(&format!("ubinfo {} -N {}", path.display(), volume)).ok()?;

            let line = output.stdout.lines().next()?;
            let re_match = re.captures(line)?;

            Some(PathBuf::from(format!(
                "{}_{}",
                path.display(),
                &re_match.name("volume").unwrap().as_str()
            )))
        })
        .ok_or_else(|| format_err!("Unable to find Ubi Volume"))
}

pub(crate) fn target_device_from_mtd_name(name: &str) -> Result<PathBuf, failure::Error> {
    let re =
        regex::Regex::new(r#"^(?P<dev>mtd\d): ([[:xdigit:]]+) ([[:xdigit:]]+) "(?P<name>.*)"$"#)
            .unwrap();
    let proc = fs::File::open("/proc/mtd")?;

    BufReader::new(proc)
        .lines()
        .filter_map(Result::ok)
        .find_map(|line| {
            re.captures(&line).and_then(|re_match| {
                let re_dev = re_match.name("dev").unwrap().as_str();
                let re_name = re_match.name("name").unwrap().as_str();
                if re_name == name {
                    Some(PathBuf::from(format!("/dev/{}", re_dev)))
                } else {
                    None
                }
            })
        })
        .ok_or_else(|| format_err!("Unable to find match for mtd device: {}", name))
}

mod ffi {
    use nix::ioctl_read;
    use std::{mem::MaybeUninit, os::unix::io::AsRawFd, path::Path};

    // From https://github.com/torvalds/linux/blob/master/include/uapi/mtd/mtd-abi.h
    const MTD_NANDFLASH: u8 = 4;
    const MTD_MLCNANDFLASH: u8 = 8;
    const MEMGETINFO: u8 = b'M';
    const MEMGETINFO_MODE: u8 = 1;

    #[repr(C)]
    pub struct mtd_info_user {
        kind: u8,
        flags: u32,
        size: u32,
        erasesize: u32,
        writesize: u32,
        oobsize: u32,
        padding: u64,
    }

    ioctl_read!(mtd_get_info, MEMGETINFO, MEMGETINFO_MODE, mtd_info_user);

    pub fn is_nand(device: &Path) -> Result<bool, failure::Error> {
        let device = std::fs::File::open(device)?;
        let info = unsafe {
            let mut info = MaybeUninit::<mtd_info_user>::uninit();
            mtd_get_info(device.as_raw_fd(), info.as_mut_ptr())?;
            info.assume_init()
        };

        Ok(info.kind == MTD_NANDFLASH || info.kind == MTD_MLCNANDFLASH)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use lazy_static::lazy_static;
    use pretty_assertions::assert_eq;
    use std::sync::{Arc, Mutex};

    pub(crate) struct FakeUbi {
        pub(crate) mtd: FakeMtd,
    }

    impl FakeUbi {
        pub(crate) fn new(names: &[&str], kind: MtdKind) -> Result<FakeUbi, failure::Error> {
            let mtd = FakeMtd::new(&["system"], kind)?;
            easy_process::run("modprobe ubi mtd=0")?;

            // Ubi created here so if anything fails the Drop will still be executed
            let ubi = FakeUbi { mtd };

            for name in names {
                easy_process::run(&format!("ubimkvol /dev/ubi0 -N {} -s 1MiB", name))?;
            }

            Ok(ubi)
        }
    }

    impl Drop for FakeUbi {
        fn drop(&mut self) {
            if let Err(e) = easy_process::run(&format!("rmmod ubi")) {
                eprintln!("Failed to cleanup FakeUbi, Error: {}", e);
            }
        }
    }

    pub(crate) struct FakeMtd {
        pub(crate) devices: Vec<PathBuf>,
        pub(crate) kind: MtdKind,
    }

    pub(crate) enum MtdKind {
        Nand,
        Nor,
    }

    impl FakeMtd {
        pub(crate) fn new(names: &[&str], kind: MtdKind) -> Result<FakeMtd, failure::Error> {
            match kind {
                MtdKind::Nand => {
                    easy_process::run("modprobe nandsim second_id_byte=0x36")?;
                    easy_process::run("mtdpart del /dev/mtd0 1")?;
                }
                MtdKind::Nor => {
                    easy_process::run("modprobe mtdram total_size=20000 erase_size=10")?;
                }
            }
            let total_size = 20000;

            // FakeMtd created here so if any subsequent command fails the drop will still
            // be called to cleanup mtd devices
            let mut mtd = FakeMtd {
                devices: vec![],
                kind,
            };
            let size = total_size / names.len();
            names.iter().enumerate().try_for_each(|(i, n)| {
                easy_process::run(&format!(
                    "mtdpart add /dev/mtd0 {} {} {}",
                    n,
                    i * size,
                    size
                ))
                .map(|_| {
                    mtd.devices
                        .push(PathBuf::from(format!("/dev/mtd{}", i + 1)))
                })
            })?;

            Ok(mtd)
        }
    }

    impl Drop for FakeMtd {
        fn drop(&mut self) {
            let module = match self.kind {
                MtdKind::Nand => "nandsim",
                MtdKind::Nor => "mtdram",
            };

            // Sleep time for nandsim to sync and avoid errors
            std::thread::sleep(std::time::Duration::from_millis(500));

            if let Err(e) = easy_process::run(&format!("rmmod {}", module)) {
                eprintln!("Failed to cleanup FakeMtd, Error: {}", e);
            }
        }
    }

    // Used to serialize access to MTD devices
    lazy_static! {
        pub static ref SERIALIZE: Arc<Mutex<()>> = Arc::new(Mutex::default());
    }

    #[test]
    #[ignore]
    fn device_from_mtd_name_nor() {
        let _lock = SERIALIZE.lock();
        let dev_names = vec!["system0", "system1"];

        let mtd = FakeMtd::new(&dev_names, MtdKind::Nor).unwrap();

        assert_eq!(
            dev_names
                .into_iter()
                .map(target_device_from_mtd_name)
                .map(Result::unwrap)
                .collect::<Vec<_>>(),
            mtd.devices,
        );
        assert!(target_device_from_mtd_name("some_inexistent_device").is_err());
    }

    #[test]
    #[ignore]
    fn device_from_mtd_name_nand() {
        let _lock = SERIALIZE.lock();
        let dev_names = vec!["system0", "system1"];

        let mtd = FakeMtd::new(&dev_names, MtdKind::Nand).unwrap();

        assert_eq!(
            dev_names
                .into_iter()
                .map(target_device_from_mtd_name)
                .map(Result::unwrap)
                .collect::<Vec<_>>(),
            mtd.devices,
        );
        assert!(target_device_from_mtd_name("some_inexistent_device").is_err());
    }

    #[test]
    #[ignore]
    fn test_is_nand() {
        let _lock = SERIALIZE.lock();
        let dev_names = vec!["system0"];

        {
            let mtd = FakeMtd::new(&dev_names, MtdKind::Nand).unwrap();
            assert_eq!(is_nand(&mtd.devices[0]).unwrap(), true);
        }
        {
            let mtd = FakeMtd::new(&dev_names, MtdKind::Nor).unwrap();
            assert_eq!(is_nand(&mtd.devices[0]).unwrap(), false);
        }
    }

    #[test]
    #[ignore]
    fn device_from_ubi_volume_name() {
        let _lock = SERIALIZE.lock();
        let volume_names = vec!["some_ui_volume", "another_ubi_volume"];

        let _ubi = FakeUbi::new(&volume_names, MtdKind::Nor).unwrap();
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[1]).unwrap(),
            PathBuf::from("/dev/ubi0_1")
        );
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[0]).unwrap(),
            PathBuf::from("/dev/ubi0_0")
        );
    }

    #[test]
    #[ignore]
    fn device_from_ubi_volume_name_multiple_volumes() {
        let _lock = SERIALIZE.lock();
        let volume_names = vec![
            "volume0", "volume1", "volume2", "volume3", "volume4", "volume5", "volume6", "volume7",
            "volume8", "volume9", "volume10", "volume11", "volume12", "volume13",
        ];

        let _ubi = FakeUbi::new(&volume_names, MtdKind::Nor).unwrap();
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[8]).unwrap(),
            PathBuf::from("/dev/ubi0_8")
        );
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[12]).unwrap(),
            PathBuf::from("/dev/ubi0_12")
        );
    }
}
<|endoftext|>


<commit_before>use petgraph::{Graph, Undirected};
use petgraph::graph;
use petgraph::graph::IndexType;
use rand;
use rand::distributions::{IndependentSample, Range};

#[derive(Hash, Eq, PartialEq, Debug, Copy, Clone, Ord, PartialOrd)]
pub struct GridCoordinate {
    pub x: isize,
    pub y: isize,
}
impl GridCoordinate {
    pub fn new(x: isize, y: isize) -> GridCoordinate {
        GridCoordinate { x: x, y: y }
    }
}

pub struct SquareGrid<GridIndexType: IndexType> {
    graph: Graph<(), (), Undirected, GridIndexType>,
    dimension_size: GridIndexType,
}

enum GridDirection {
    North,
    South,
    East,
    West,
}

impl<GridIndexType: IndexType> SquareGrid<GridIndexType> {
    pub fn new(dimension_size: GridIndexType) -> SquareGrid<GridIndexType> {

        let dim_size = dimension_size.index();
        let cells_count = dim_size * dim_size;
        let nodes_count_hint = cells_count;
        let edges_count_hint = 4 * cells_count - 4 * dim_size; // Probably overkill, but don't want any capacity panics

        let mut grid = SquareGrid {
            graph: Graph::with_capacity(nodes_count_hint, edges_count_hint),
            dimension_size: dimension_size,
        };
        for _ in 0..cells_count {
            let _ = grid.graph.add_node(());
        }

        grid
    }

    pub fn random_cell(&self) -> GridCoordinate {
        let range_end_exclusive = self.dimension_size.index() * self.dimension_size.index();
        let random_pos = Range::new(0, range_end_exclusive);
        let mut rng = rand::thread_rng();
        let index = random_pos.ind_sample(&mut rng);
        index_to_grid_coordinate(self.dimension_size.index(), index)
    }

    /// Link two cells
    ///
    /// Todo - only allow links between adjacent cells? If `b` not in `g.neighbours(a)`.
    ///      - better to change the API to take an index and GridDirection
    ///
    /// Panics if a cell does not exist.
    pub fn link(&mut self, a: GridCoordinate, b: GridCoordinate) {
        if a != b {
            let a_index = self.grid_coordinate_graph_index(&a);
            let b_index = self.grid_coordinate_graph_index(&b);
            let _ = self.graph.update_edge(a_index, b_index, ());
        }
    }

    /// Unlink two cells, if a link exists between them.
    pub fn unlink(&mut self, a: GridCoordinate, b: GridCoordinate) {
        let a_index = self.grid_coordinate_graph_index(&a);
        let b_index = self.grid_coordinate_graph_index(&b);
        if let Some(edge_index) = self.graph.find_edge(a_index, b_index) {
            // This will invalidate the last edge index in the graph, which is fine as we
            // are not storing them for any reason.
            self.graph.remove_edge(edge_index);
        }
    }

    /// Cell nodes that are linked to a particular node by a passage.
    pub fn links(&self, coord: GridCoordinate) -> Vec<GridCoordinate> {
        self.graph
            .edges(self.grid_coordinate_graph_index(&coord))
            .map(|index_edge_data_pair| {
                let grid_node_index = index_edge_data_pair.0.clone();
                index_to_grid_coordinate(self.dimension_size.index(), grid_node_index.index())
            })
            .collect()
    }

    /// Cell nodes that are to the North, South, East or West of a particular node, but not
    /// necessarily linked by a passage.
    pub fn neighbours(&self, coord: GridCoordinate) -> Vec<GridCoordinate> {

        vec![offset_coordinate(&coord, GridDirection::North),
             offset_coordinate(&coord, GridDirection::South),
             offset_coordinate(&coord, GridDirection::East),
             offset_coordinate(&coord, GridDirection::West)]
            .into_iter()
            .filter(|adjacent_coord| self.is_valid_coordinate(adjacent_coord))
            .collect()
    }

    /// Are two cells in the grid linked?
    pub fn is_linked(&self, a: GridCoordinate, b: GridCoordinate) -> bool {
        let a_index = self.grid_coordinate_graph_index(&a);
        let b_index = self.grid_coordinate_graph_index(&b);
        self.graph.find_edge(a_index, b_index).is_some()
    }

    pub fn iter(&self) -> CellIter {
        let dim_size = self.dimension_size.index();
        CellIter {
            current_cell_number: 0,
            dimension_size: dim_size,
            cells_count: dim_size * dim_size,
        }
    }

    pub fn iter_row(&self) -> BatchIter {
        BatchIter {
            iter_type: BatchIterType::Row,
            current_index: 0,
            dimension_size: self.dimension_size.index(),
        }
    }

    pub fn iter_column(&self) -> BatchIter {
        BatchIter {
            iter_type: BatchIterType::Column,
            current_index: 0,
            dimension_size: self.dimension_size.index(),
        }
    }

    fn is_neighbour(&self, a: GridCoordinate, b: GridCoordinate) -> bool {
        self.neighbours(a).iter().any(|&coord| coord == b)
    }

    fn is_valid_coordinate(&self, coord: &GridCoordinate) -> bool {
        let (x, y) = (coord.x, coord.y);
        let dim_size = self.dimension_size.index() as isize;
        if x < 0 || y < 0 || x >= dim_size || y >= dim_size {
            return false;
        }
        true
    }

    fn grid_coordinate_graph_index(&self, coord: &GridCoordinate) -> graph::NodeIndex<GridIndexType> {
        let grid_index_raw = ((coord.y * self.dimension_size.index() as isize) + coord.x) as usize;
        graph::NodeIndex::<GridIndexType>::new(grid_index_raw)
    }
}

pub struct CellIter {
    current_cell_number: usize,
    dimension_size: usize,
    cells_count: usize,
}
impl Iterator for CellIter {
    type Item = GridCoordinate;
    fn next(&mut self) -> Option<Self::Item> {
        if self.current_cell_number < self.cells_count {
            let coord = index_to_grid_coordinate(self.dimension_size, self.current_cell_number);
            self.current_cell_number += 1;
            Some(coord)
        } else {
            None
        }
    }
}

enum BatchIterType {
    Row,
    Column,
}
pub struct BatchIter {
    iter_type: BatchIterType,
    current_index: usize,
    dimension_size: usize,
}
impl Iterator for BatchIter {
    type Item = Vec<GridCoordinate>;
    fn next(&mut self) -> Option<Self::Item> {
        if self.current_index < self.dimension_size {
            let coords = (0..self.dimension_size)
                             .into_iter()
                             .map(|i| {
                                 if let BatchIterType::Row = self.iter_type {
                                     GridCoordinate::new(i as isize, self.current_index as isize)
                                 } else {
                                     GridCoordinate::new(self.current_index as isize, i as isize)
                                 }
                             })
                             .collect();
            self.current_index += 1;
            Some(coords)
        } else {
            None
        }
    }
}

fn index_to_grid_coordinate(dimension_size: usize, one_dimensional_index: usize) -> GridCoordinate {
    let y = one_dimensional_index / dimension_size;
    let x = one_dimensional_index - (y * dimension_size);
    GridCoordinate {
        x: x as isize,
        y: y as isize,
    }
}

fn offset_coordinate(coord: &GridCoordinate, dir: GridDirection) -> GridCoordinate {
    let (x, y) = (coord.x, coord.y);
    match dir {
        GridDirection::North => GridCoordinate { x: x, y: y - 1 },
        GridDirection::South => GridCoordinate { x: x, y: y + 1 },
        GridDirection::East => GridCoordinate { x: x + 1, y: y },
        GridDirection::West => GridCoordinate { x: x - 1, y: y },
    }
}

#[cfg(test)]
mod test {

    use super::*;
    use itertools::Itertools; // a trait

    type SmallGrid = SquareGrid<u8>;

    #[test]
    fn neighbour_cells() {
        let g = SmallGrid::new(10);

        let check_expected_neighbours = |coord, vec_expected_neighbours: Vec<GridCoordinate>| {
            let node_indices = g.neighbours(coord)
                                .into_iter()
                                .sorted();
            let expected_indices = vec_expected_neighbours.into_iter()
                                                          .sorted();
            assert_eq!(node_indices, expected_indices);
        };
        let gc = |x, y| GridCoordinate::new(x, y);

        // corners
        check_expected_neighbours(gc(0, 0), vec![gc(1, 0), gc(0, 1)]);
        check_expected_neighbours(gc(9, 0), vec![gc(8, 0), gc(9, 1)]);
        check_expected_neighbours(gc(0, 9), vec![gc(0, 8), gc(1, 9)]);
        check_expected_neighbours(gc(9, 9), vec![gc(9, 8), gc(8, 9)]);

        // side element examples
        check_expected_neighbours(gc(1, 0), vec![gc(0, 0), gc(1, 1), gc(2, 0)]);
        check_expected_neighbours(gc(0, 1), vec![gc(0, 0), gc(0, 2), gc(1, 1)]);
        check_expected_neighbours(gc(0, 8), vec![gc(1, 8), gc(0, 7), gc(0, 9)]);
        check_expected_neighbours(gc(9, 8), vec![gc(9, 7), gc(9, 9), gc(8, 8)]);

        // Some place with 4 neighbours inside the grid
        check_expected_neighbours(gc(1, 1), vec![gc(0, 1), gc(1, 0), gc(2, 1), gc(1, 2)]);
    }

    #[test]
    fn random_cell() {
        let g = SmallGrid::new(4);
        let cells_count = 4 * 4;
        for _ in 0..1000 {
            let coord = g.random_cell();
            assert!(coord.x >= 0 && coord.x < cells_count);
            assert!(coord.y >= 0 && coord.y < cells_count);
        }
    }

    #[test]
    fn cell_iter() {
        let g = SmallGrid::new(2);
        assert_eq!(g.iter().collect::<Vec<GridCoordinate>>(),
                   vec![GridCoordinate::new(0, 0),
                        GridCoordinate::new(1, 0),
                        GridCoordinate::new(0, 1),
                        GridCoordinate::new(1, 1)]);
    }

    #[test]
    fn row_iter() {
        let g = SmallGrid::new(2);
        assert_eq!(g.iter_row().collect::<Vec<Vec<GridCoordinate>>>(),
                   vec![vec![GridCoordinate::new(0, 0), GridCoordinate::new(1, 0)],
                        vec![GridCoordinate::new(0, 1), GridCoordinate::new(1, 1)]]);
    }

    #[test]
    fn column_iter() {
        let g = SmallGrid::new(2);
        assert_eq!(g.iter_column().collect::<Vec<Vec<GridCoordinate>>>(),
                   vec![vec![GridCoordinate::new(0, 0), GridCoordinate::new(0, 1)],
                        vec![GridCoordinate::new(1, 0), GridCoordinate::new(1, 1)]]);
    }

    #[test]
    fn linking_cells() {
        let mut g = SmallGrid::new(4);
        let a = GridCoordinate::new(0, 1);
        let b = GridCoordinate::new(0, 2);
        let c = GridCoordinate::new(0, 3);

        // I'd rather use a closure, but it needs to borrow the graph immutably until
        // it goes out of scope and it's ugly to pass the grid into each function call
        macro_rules! links_sorted {
            ($x:expr) => (g.links($x).into_iter().sorted())
        }
        // Testing that the order of the arguments to `is_linked` does not matter
        macro_rules! bi_check_linked {
            ($x:expr, $y:expr) => (g.is_linked($x, $y) && g.is_linked($y, $x))
        }

        // a, b and c start with no links
        assert!(!bi_check_linked!(a, b));
        assert!(!bi_check_linked!(a, c));
        assert!(!bi_check_linked!(b, c));
        assert_eq!(links_sorted!(a), vec![]);
        assert_eq!(links_sorted!(b), vec![]);
        assert_eq!(links_sorted!(c), vec![]);

        g.link(a, b);
        // a - b linked bi-directionally
        assert!(bi_check_linked!(a, b));
        assert_eq!(links_sorted!(a), vec![b]);
        assert_eq!(links_sorted!(b), vec![a]);

        g.link(b, c);
        // a - b still linked bi-directionally after linking b - c
        // b linked to a & c bi-directionally
        // c linked to b bi-directionally
        assert!(bi_check_linked!(a, b));
        assert!(bi_check_linked!(b, c));
        assert!(!bi_check_linked!(a, c));
        assert_eq!(links_sorted!(a), vec![b]);
        assert_eq!(links_sorted!(b), vec![a, c]);
        assert_eq!(links_sorted!(c), vec![b]);

        // a - b still linked bi-directionally after updating exist link
        assert!(bi_check_linked!(a, b));
        assert!(bi_check_linked!(b, c));
        assert_eq!(links_sorted!(a), vec![b]);
        assert_eq!(links_sorted!(b), vec![a, c]);

        // a - b unlinked
        // b still linked to c bi-directionally
        g.unlink(a, b);
        assert!(!bi_check_linked!(a, b));
        assert!(bi_check_linked!(b, c));
        assert_eq!(links_sorted!(a), vec![]);
        assert_eq!(links_sorted!(b), vec![c]);
        assert_eq!(links_sorted!(c), vec![b]);

        // a, b and c start all unlinked again
        g.unlink(b, c);
        assert!(!bi_check_linked!(a, b));
        assert!(!bi_check_linked!(a, c));
        assert!(!bi_check_linked!(b, c));
        assert_eq!(links_sorted!(a), vec![]);
        assert_eq!(links_sorted!(b), vec![]);
        assert_eq!(links_sorted!(c), vec![]);
    }

    #[test]
    fn no_self_linked_cycles() {
        let mut g = SmallGrid::new(4);
        let a = GridCoordinate::new(0, 0);
        g.link(a, a);
        assert_eq!(g.links(a), vec![]);
    }

    #[test]
    fn no_parallel_duplicated_linked_cells() {
        let mut g = SmallGrid::new(4);
        let a = GridCoordinate::new(0, 0);
        let b = GridCoordinate::new(0, 1);
        g.link(a, b);
        g.link(a, b);
        assert_eq!(g.links(a), vec![b]);
        assert_eq!(g.links(b), vec![a]);

        g.unlink(a, b);
        assert_eq!(g.links(a), vec![]);
        assert_eq!(g.links(b), vec![]);
    }
}
<commit_msg>Update syntax for constructing offset GridCoordinates.<commit_after>use petgraph::{Graph, Undirected};
use petgraph::graph;
use petgraph::graph::IndexType;
use rand;
use rand::distributions::{IndependentSample, Range};

#[derive(Hash, Eq, PartialEq, Debug, Copy, Clone, Ord, PartialOrd)]
pub struct GridCoordinate {
    pub x: isize,
    pub y: isize,
}
impl GridCoordinate {
    pub fn new(x: isize, y: isize) -> GridCoordinate {
        GridCoordinate { x: x, y: y }
    }
}

pub struct SquareGrid<GridIndexType: IndexType> {
    graph: Graph<(), (), Undirected, GridIndexType>,
    dimension_size: GridIndexType,
}

enum GridDirection {
    North,
    South,
    East,
    West,
}

impl<GridIndexType: IndexType> SquareGrid<GridIndexType> {
    pub fn new(dimension_size: GridIndexType) -> SquareGrid<GridIndexType> {

        let dim_size = dimension_size.index();
        let cells_count = dim_size * dim_size;
        let nodes_count_hint = cells_count;
        let edges_count_hint = 4 * cells_count - 4 * dim_size; // Probably overkill, but don't want any capacity panics

        let mut grid = SquareGrid {
            graph: Graph::with_capacity(nodes_count_hint, edges_count_hint),
            dimension_size: dimension_size,
        };
        for _ in 0..cells_count {
            let _ = grid.graph.add_node(());
        }

        grid
    }

    pub fn random_cell(&self) -> GridCoordinate {
        let range_end_exclusive = self.dimension_size.index() * self.dimension_size.index();
        let random_pos = Range::new(0, range_end_exclusive);
        let mut rng = rand::thread_rng();
        let index = random_pos.ind_sample(&mut rng);
        index_to_grid_coordinate(self.dimension_size.index(), index)
    }

    /// Link two cells
    ///
    /// Todo - only allow links between adjacent cells? If `b` not in `g.neighbours(a)`.
    ///      - better to change the API to take an index and GridDirection
    ///
    /// Panics if a cell does not exist.
    pub fn link(&mut self, a: GridCoordinate, b: GridCoordinate) {
        if a != b {
            let a_index = self.grid_coordinate_graph_index(&a);
            let b_index = self.grid_coordinate_graph_index(&b);
            let _ = self.graph.update_edge(a_index, b_index, ());
        }
    }

    /// Unlink two cells, if a link exists between them.
    pub fn unlink(&mut self, a: GridCoordinate, b: GridCoordinate) {
        let a_index = self.grid_coordinate_graph_index(&a);
        let b_index = self.grid_coordinate_graph_index(&b);
        if let Some(edge_index) = self.graph.find_edge(a_index, b_index) {
            // This will invalidate the last edge index in the graph, which is fine as we
            // are not storing them for any reason.
            self.graph.remove_edge(edge_index);
        }
    }

    /// Cell nodes that are linked to a particular node by a passage.
    pub fn links(&self, coord: GridCoordinate) -> Vec<GridCoordinate> {
        self.graph
            .edges(self.grid_coordinate_graph_index(&coord))
            .map(|index_edge_data_pair| {
                let grid_node_index = index_edge_data_pair.0.clone();
                index_to_grid_coordinate(self.dimension_size.index(), grid_node_index.index())
            })
            .collect()
    }

    /// Cell nodes that are to the North, South, East or West of a particular node, but not
    /// necessarily linked by a passage.
    pub fn neighbours(&self, coord: GridCoordinate) -> Vec<GridCoordinate> {

        vec![offset_coordinate(&coord, GridDirection::North),
             offset_coordinate(&coord, GridDirection::South),
             offset_coordinate(&coord, GridDirection::East),
             offset_coordinate(&coord, GridDirection::West)]
            .into_iter()
            .filter(|adjacent_coord| self.is_valid_coordinate(adjacent_coord))
            .collect()
    }

    /// Are two cells in the grid linked?
    pub fn is_linked(&self, a: GridCoordinate, b: GridCoordinate) -> bool {
        let a_index = self.grid_coordinate_graph_index(&a);
        let b_index = self.grid_coordinate_graph_index(&b);
        self.graph.find_edge(a_index, b_index).is_some()
    }

    pub fn iter(&self) -> CellIter {
        let dim_size = self.dimension_size.index();
        CellIter {
            current_cell_number: 0,
            dimension_size: dim_size,
            cells_count: dim_size * dim_size,
        }
    }

    pub fn iter_row(&self) -> BatchIter {
        BatchIter {
            iter_type: BatchIterType::Row,
            current_index: 0,
            dimension_size: self.dimension_size.index(),
        }
    }

    pub fn iter_column(&self) -> BatchIter {
        BatchIter {
            iter_type: BatchIterType::Column,
            current_index: 0,
            dimension_size: self.dimension_size.index(),
        }
    }

    fn is_neighbour(&self, a: GridCoordinate, b: GridCoordinate) -> bool {
        self.neighbours(a).iter().any(|&coord| coord == b)
    }

    fn is_valid_coordinate(&self, coord: &GridCoordinate) -> bool {
        let (x, y) = (coord.x, coord.y);
        let dim_size = self.dimension_size.index() as isize;
        if x < 0 || y < 0 || x >= dim_size || y >= dim_size {
            return false;
        }
        true
    }

    fn grid_coordinate_graph_index(&self, coord: &GridCoordinate) -> graph::NodeIndex<GridIndexType> {
        let grid_index_raw = ((coord.y * self.dimension_size.index() as isize) + coord.x) as usize;
        graph::NodeIndex::<GridIndexType>::new(grid_index_raw)
    }
}

pub struct CellIter {
    current_cell_number: usize,
    dimension_size: usize,
    cells_count: usize,
}
impl Iterator for CellIter {
    type Item = GridCoordinate;
    fn next(&mut self) -> Option<Self::Item> {
        if self.current_cell_number < self.cells_count {
            let coord = index_to_grid_coordinate(self.dimension_size, self.current_cell_number);
            self.current_cell_number += 1;
            Some(coord)
        } else {
            None
        }
    }
}

enum BatchIterType {
    Row,
    Column,
}
pub struct BatchIter {
    iter_type: BatchIterType,
    current_index: usize,
    dimension_size: usize,
}
impl Iterator for BatchIter {
    type Item = Vec<GridCoordinate>;
    fn next(&mut self) -> Option<Self::Item> {
        if self.current_index < self.dimension_size {
            let coords = (0..self.dimension_size)
                             .into_iter()
                             .map(|i| {
                                 if let BatchIterType::Row = self.iter_type {
                                     GridCoordinate::new(i as isize, self.current_index as isize)
                                 } else {
                                     GridCoordinate::new(self.current_index as isize, i as isize)
                                 }
                             })
                             .collect();
            self.current_index += 1;
            Some(coords)
        } else {
            None
        }
    }
}

fn index_to_grid_coordinate(dimension_size: usize, one_dimensional_index: usize) -> GridCoordinate {
    let y = one_dimensional_index / dimension_size;
    let x = one_dimensional_index - (y * dimension_size);
    GridCoordinate {
        x: x as isize,
        y: y as isize,
    }
}

fn offset_coordinate(coord: &GridCoordinate, dir: GridDirection) -> GridCoordinate {
    let (x, y) = (coord.x, coord.y);
    match dir {
        GridDirection::North => GridCoordinate { y: y - 1, ..*coord },
        GridDirection::South => GridCoordinate { y: y + 1, ..*coord },
        GridDirection::East => GridCoordinate { x: x + 1, ..*coord },
        GridDirection::West => GridCoordinate { x: x - 1, ..*coord },
    }
}

#[cfg(test)]
mod test {

    use super::*;
    use itertools::Itertools; // a trait

    type SmallGrid = SquareGrid<u8>;

    #[test]
    fn neighbour_cells() {
        let g = SmallGrid::new(10);

        let check_expected_neighbours = |coord, vec_expected_neighbours: Vec<GridCoordinate>| {
            let node_indices = g.neighbours(coord)
                                .into_iter()
                                .sorted();
            let expected_indices = vec_expected_neighbours.into_iter()
                                                          .sorted();
            assert_eq!(node_indices, expected_indices);
        };
        let gc = |x, y| GridCoordinate::new(x, y);

        // corners
        check_expected_neighbours(gc(0, 0), vec![gc(1, 0), gc(0, 1)]);
        check_expected_neighbours(gc(9, 0), vec![gc(8, 0), gc(9, 1)]);
        check_expected_neighbours(gc(0, 9), vec![gc(0, 8), gc(1, 9)]);
        check_expected_neighbours(gc(9, 9), vec![gc(9, 8), gc(8, 9)]);

        // side element examples
        check_expected_neighbours(gc(1, 0), vec![gc(0, 0), gc(1, 1), gc(2, 0)]);
        check_expected_neighbours(gc(0, 1), vec![gc(0, 0), gc(0, 2), gc(1, 1)]);
        check_expected_neighbours(gc(0, 8), vec![gc(1, 8), gc(0, 7), gc(0, 9)]);
        check_expected_neighbours(gc(9, 8), vec![gc(9, 7), gc(9, 9), gc(8, 8)]);

        // Some place with 4 neighbours inside the grid
        check_expected_neighbours(gc(1, 1), vec![gc(0, 1), gc(1, 0), gc(2, 1), gc(1, 2)]);
    }

    #[test]
    fn random_cell() {
        let g = SmallGrid::new(4);
        let cells_count = 4 * 4;
        for _ in 0..1000 {
            let coord = g.random_cell();
            assert!(coord.x >= 0 && coord.x < cells_count);
            assert!(coord.y >= 0 && coord.y < cells_count);
        }
    }

    #[test]
    fn cell_iter() {
        let g = SmallGrid::new(2);
        assert_eq!(g.iter().collect::<Vec<GridCoordinate>>(),
                   vec![GridCoordinate::new(0, 0),
                        GridCoordinate::new(1, 0),
                        GridCoordinate::new(0, 1),
                        GridCoordinate::new(1, 1)]);
    }

    #[test]
    fn row_iter() {
        let g = SmallGrid::new(2);
        assert_eq!(g.iter_row().collect::<Vec<Vec<GridCoordinate>>>(),
                   vec![vec![GridCoordinate::new(0, 0), GridCoordinate::new(1, 0)],
                        vec![GridCoordinate::new(0, 1), GridCoordinate::new(1, 1)]]);
    }

    #[test]
    fn column_iter() {
        let g = SmallGrid::new(2);
        assert_eq!(g.iter_column().collect::<Vec<Vec<GridCoordinate>>>(),
                   vec![vec![GridCoordinate::new(0, 0), GridCoordinate::new(0, 1)],
                        vec![GridCoordinate::new(1, 0), GridCoordinate::new(1, 1)]]);
    }

    #[test]
    fn linking_cells() {
        let mut g = SmallGrid::new(4);
        let a = GridCoordinate::new(0, 1);
        let b = GridCoordinate::new(0, 2);
        let c = GridCoordinate::new(0, 3);

        // I'd rather use a closure, but it needs to borrow the graph immutably until
        // it goes out of scope and it's ugly to pass the grid into each function call
        macro_rules! links_sorted {
            ($x:expr) => (g.links($x).into_iter().sorted())
        }
        // Testing that the order of the arguments to `is_linked` does not matter
        macro_rules! bi_check_linked {
            ($x:expr, $y:expr) => (g.is_linked($x, $y) && g.is_linked($y, $x))
        }

        // a, b and c start with no links
        assert!(!bi_check_linked!(a, b));
        assert!(!bi_check_linked!(a, c));
        assert!(!bi_check_linked!(b, c));
        assert_eq!(links_sorted!(a), vec![]);
        assert_eq!(links_sorted!(b), vec![]);
        assert_eq!(links_sorted!(c), vec![]);

        g.link(a, b);
        // a - b linked bi-directionally
        assert!(bi_check_linked!(a, b));
        assert_eq!(links_sorted!(a), vec![b]);
        assert_eq!(links_sorted!(b), vec![a]);

        g.link(b, c);
        // a - b still linked bi-directionally after linking b - c
        // b linked to a & c bi-directionally
        // c linked to b bi-directionally
        assert!(bi_check_linked!(a, b));
        assert!(bi_check_linked!(b, c));
        assert!(!bi_check_linked!(a, c));
        assert_eq!(links_sorted!(a), vec![b]);
        assert_eq!(links_sorted!(b), vec![a, c]);
        assert_eq!(links_sorted!(c), vec![b]);

        // a - b still linked bi-directionally after updating exist link
        assert!(bi_check_linked!(a, b));
        assert!(bi_check_linked!(b, c));
        assert_eq!(links_sorted!(a), vec![b]);
        assert_eq!(links_sorted!(b), vec![a, c]);

        // a - b unlinked
        // b still linked to c bi-directionally
        g.unlink(a, b);
        assert!(!bi_check_linked!(a, b));
        assert!(bi_check_linked!(b, c));
        assert_eq!(links_sorted!(a), vec![]);
        assert_eq!(links_sorted!(b), vec![c]);
        assert_eq!(links_sorted!(c), vec![b]);

        // a, b and c start all unlinked again
        g.unlink(b, c);
        assert!(!bi_check_linked!(a, b));
        assert!(!bi_check_linked!(a, c));
        assert!(!bi_check_linked!(b, c));
        assert_eq!(links_sorted!(a), vec![]);
        assert_eq!(links_sorted!(b), vec![]);
        assert_eq!(links_sorted!(c), vec![]);
    }

    #[test]
    fn no_self_linked_cycles() {
        let mut g = SmallGrid::new(4);
        let a = GridCoordinate::new(0, 0);
        g.link(a, a);
        assert_eq!(g.links(a), vec![]);
    }

    #[test]
    fn no_parallel_duplicated_linked_cells() {
        let mut g = SmallGrid::new(4);
        let a = GridCoordinate::new(0, 0);
        let b = GridCoordinate::new(0, 1);
        g.link(a, b);
        g.link(a, b);
        assert_eq!(g.links(a), vec![b]);
        assert_eq!(g.links(b), vec![a]);

        g.unlink(a, b);
        assert_eq!(g.links(a), vec![]);
        assert_eq!(g.links(b), vec![]);
    }
}
<|endoftext|>


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

/// A shader module.
///
/// A shader module is a container that associates some shading code to several identifiers.
struct ShaderModule {
  symbols: HashMap<Identifier, ShadingCode>
}

/// Spectra Shading Language AST.
#[derive(Clone, Debug, Eq, PartialEq)]
enum SSL {
  /// An `export list_of_identifiers_` statement.
  Export(ExportList),
  /// A `from module use list of identifiers` statement.
  FromUse(ImportList),
  /// A `pipeline { list_of_pipeline_attributes }` statement.
  Pipeline(PipelineStatement),
  /// A yield statement, valid in geometry shaders.
  Yield(GeometryYieldExpression),
}

/// A module.
type Module = String;
/// An identifier.
type Identifier = String;
/// Some opaque shading code.
type ShadingCode = String;
/// An expression.
type Expression = String;

/// An export non-empty list.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ExportList {
  pub export_list: HashSet<Identifier>
}

/// An import non-empty list.
#[derive(Clone, Debug, Eq, PartialEq)]
struct ImportList {
  module: Module,
  list: HashSet<Identifier>
}

/// A pipeline statement.
#[derive(Clone, Debug, Eq, PartialEq)]
struct PipelineStatement {
  attributes: Vec<PipelineAttribute>
}

/// Attributes that can be set in a pipeline.
#[derive(Clone, Debug, Eq, PartialEq)]
enum PipelineAttribute {
  /// Maximum vertices that the geometry shader can output.
  GeometryShaderMaxVertices(u32),
  /// Number of times the geometry shader must be invoked.
  GeometryShaderInvokations(u32)
}

/// Expressions that can be yielded in a geometry shader.
#[derive(Clone, Debug, Eq, PartialEq)]
enum GeometryYieldExpression {
  /// Yield a primitive.
  YieldPrimitive,
  /// Yield a primitive’s vertex (fold vertex).
  YieldFoldVertex(Expression)
}

/// Error that can occur when parsing SSL code.
#[derive(Clone, Debug, Eq, PartialEq)]
enum ParseError {
  ExpressionError(String)
}
<commit_msg>Added ModulePath. #24<commit_after>use std::collections::{HashMap, HashSet};

/// A shader module.
///
/// A shader module is a container that associates some shading code to several identifiers.
struct ShaderModule {
  symbols: HashMap<Identifier, ShadingCode>
}

/// Spectra Shading Language AST.
#[derive(Clone, Debug, Eq, PartialEq)]
enum SSL {
  /// An `export list_of_identifiers_` statement.
  Export(ExportList),
  /// A `from module use list of identifiers` statement.
  FromUse(ImportList),
  /// A `pipeline { list_of_pipeline_attributes }` statement.
  Pipeline(PipelineStatement),
  /// A yield statement, valid in geometry shaders.
  Yield(GeometryYieldExpression),
}

/// A module.
type ModuleName = String;
/// An identifier.
type Identifier = String;
/// Some opaque shading code.
type ShadingCode = String;
/// An expression.
type Expression = String;

/// An export non-empty list.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ExportList {
  pub export_list: HashSet<ModulePath>
}

/// An import non-empty list.
#[derive(Clone, Debug, Eq, PartialEq)]
struct ImportList {
  module: ModuleName,
  list: HashSet<ModulePath>
}

/// A module path is a list of module(s), representing a hierarchy.
#[derive(Clone, Debug, Hash, Eq, PartialEq)]
pub struct ModulePath {
  hierarchy: Vec<ModuleName>
}

/// A pipeline statement.
#[derive(Clone, Debug, Eq, PartialEq)]
struct PipelineStatement {
  attributes: Vec<PipelineAttribute>
}

/// Attributes that can be set in a pipeline.
#[derive(Clone, Debug, Eq, PartialEq)]
enum PipelineAttribute {
  /// Maximum vertices that the geometry shader can output.
  GeometryShaderMaxVertices(u32),
  /// Number of times the geometry shader must be invoked.
  GeometryShaderInvokations(u32)
}

/// Expressions that can be yielded in a geometry shader.
#[derive(Clone, Debug, Eq, PartialEq)]
enum GeometryYieldExpression {
  /// Yield a primitive.
  YieldPrimitive,
  /// Yield a primitive’s vertex (fold vertex).
  YieldFoldVertex(Expression)
}

/// Error that can occur when parsing SSL code.
#[derive(Clone, Debug, Eq, PartialEq)]
enum ParseError {
  ExpressionError(String)
}
<|endoftext|>


<commit_before>use storage::Storage;

use aws::aws::common::credentials::DefaultCredentialsProviderSync;
use aws::aws::common::region::Region;
use aws::aws::s3::bucket::CreateBucketRequest;
use aws::aws::s3::endpoint::{Endpoint, Signature};
use aws::aws::s3::object::{GetObjectRequest, PutObjectRequest};
use aws::aws::s3::s3client::S3Client;
use hyper::client::Client;
use url::Url;

pub struct S3Storage {
    client: S3Client<DefaultCredentialsProviderSync, Client>,
}

impl S3Storage {
    pub fn new(url: Url) -> Self {
        let provider = DefaultCredentialsProviderSync::new(None).unwrap();
        let endpoint = Endpoint::new(
            Region::EuWest1,
            Signature::V4,
            Some(url),
            None,
            None,
            Some(false),
        );
        let client = S3Client::new(provider, endpoint);

        let storage = S3Storage { client: client };
        storage.init();
        storage
    }

    fn init(&self) {
        if !self.client
            .list_buckets()
            .unwrap()
            .buckets
            .iter()
            .any(|ref bucket| bucket.name == "raw")
        {
            let mut req = CreateBucketRequest::default();
            req.bucket = "raw".to_string();
            self.client.create_bucket(&req).unwrap();
        }
        if !self.client
            .list_buckets()
            .unwrap()
            .buckets
            .iter()
            .any(|ref bucket| bucket.name == "image")
        {
            let mut req = CreateBucketRequest::default();
            req.bucket = "image".to_string();
            self.client.create_bucket(&req).unwrap();
        }
    }
}

impl Storage for S3Storage {
    fn store(&self, bucket: &String, name: &String, data: &[u8]) {
        let mut put_obj = PutObjectRequest::default();
        put_obj.bucket = bucket.to_owned();
        put_obj.key = name.to_string();
        put_obj.body = Some(data);
        self.client.put_object(&put_obj, None).unwrap();
    }

    fn get(&self, bucket: &String, name: &String) -> Vec<u8> {
        let mut req = GetObjectRequest::default();
        req.bucket = bucket.to_owned();
        req.key = name.to_owned();
        let resp = self.client.get_object(&req, None).unwrap();
        if resp.is_body {
            resp.body
        } else {
            resp.body_buffer
        }
    }
}

#[cfg(test)]
mod tests {
    use dotenv::dotenv;

    use super::*;

    #[test]
    fn store_and_get_text() {
        dotenv().unwrap();

        let s3_storage = S3Storage::new("http://127.0.0.1:9100".parse().unwrap());
        s3_storage.store(&"raw".to_string(), &"text.txt".to_string(), b"hello");
        let result = s3_storage.get(&"raw".to_string(), &"text.txt".to_string());
        assert_eq!(b"hello", result.as_slice());
    }
}
<commit_msg>Create sidecar bucket if it does not exist<commit_after>use storage::Storage;

use aws::aws::common::credentials::DefaultCredentialsProviderSync;
use aws::aws::common::region::Region;
use aws::aws::s3::bucket::CreateBucketRequest;
use aws::aws::s3::endpoint::{Endpoint, Signature};
use aws::aws::s3::object::{GetObjectRequest, PutObjectRequest};
use aws::aws::s3::s3client::S3Client;
use hyper::client::Client;
use url::Url;

pub struct S3Storage {
    client: S3Client<DefaultCredentialsProviderSync, Client>,
}

impl S3Storage {
    pub fn new(url: Url) -> Self {
        let provider = DefaultCredentialsProviderSync::new(None).unwrap();
        let endpoint = Endpoint::new(
            Region::EuWest1,
            Signature::V4,
            Some(url),
            None,
            None,
            Some(false),
        );
        let client = S3Client::new(provider, endpoint);

        let storage = S3Storage { client: client };
        storage.init();
        storage
    }

    fn init(&self) {
        if !self.client
            .list_buckets()
            .unwrap()
            .buckets
            .iter()
            .any(|ref bucket| bucket.name == "raw")
        {
            let mut req = CreateBucketRequest::default();
            req.bucket = "raw".to_string();
            self.client.create_bucket(&req).unwrap();
        }
        if !self.client
            .list_buckets()
            .unwrap()
            .buckets
            .iter()
            .any(|ref bucket| bucket.name == "sidecar")
        {
            let mut req = CreateBucketRequest::default();
            req.bucket = "sidecar".to_string();
            self.client.create_bucket(&req).unwrap();
        }
        if !self.client
            .list_buckets()
            .unwrap()
            .buckets
            .iter()
            .any(|ref bucket| bucket.name == "image")
        {
            let mut req = CreateBucketRequest::default();
            req.bucket = "image".to_string();
            self.client.create_bucket(&req).unwrap();
        }
    }
}

impl Storage for S3Storage {
    fn store(&self, bucket: &String, name: &String, data: &[u8]) {
        let mut put_obj = PutObjectRequest::default();
        put_obj.bucket = bucket.to_owned();
        put_obj.key = name.to_string();
        put_obj.body = Some(data);
        self.client.put_object(&put_obj, None).unwrap();
    }

    fn get(&self, bucket: &String, name: &String) -> Vec<u8> {
        let mut req = GetObjectRequest::default();
        req.bucket = bucket.to_owned();
        req.key = name.to_owned();
        let resp = self.client.get_object(&req, None).unwrap();
        if resp.is_body {
            resp.body
        } else {
            resp.body_buffer
        }
    }
}

#[cfg(test)]
mod tests {
    use dotenv::dotenv;

    use super::*;

    #[test]
    fn store_and_get_text() {
        dotenv().unwrap();

        let s3_storage = S3Storage::new("http://127.0.0.1:9100".parse().unwrap());
        s3_storage.store(&"raw".to_string(), &"text.txt".to_string(), b"hello");
        let result = s3_storage.get(&"raw".to_string(), &"text.txt".to_string());
        assert_eq!(b"hello", result.as_slice());
    }
}
<|endoftext|>