<commit_before>use std::{
    collections::HashMap,
    error::Error,
    fs::File,
    io::{
        BufRead,
        BufReader,
    },
    path::{
        Path,
        PathBuf,
    },
    str,
    thread,
};

use serde_query::{
    DeserializeQuery,
    Query,
};

use anyhow::{
    Context,
    Result,
};

// use crossbeam_channel::{
//     unbounded,
//     Sender,
// };

use serde_json::Deserializer;

use rkv::{
    backend::{
        BackendEnvironmentBuilder,
        Lmdb,
        LmdbDatabase,
        LmdbEnvironment,
    },
    Manager,
    Rkv,
    SingleStore,
    StoreOptions,
    Value,
};

#[derive(Debug, DeserializeQuery)]
struct DocId {
    #[query(".id")]
    id: String,
}

pub struct Spec<'a> {
    path: &'a Path,
    allow_missing: bool,
    commit_size: usize,
    threads: usize,
}

impl<'a> Spec<'a> {
    pub fn new(path: &Path) -> Spec {
        Spec {
            path,
            allow_missing: false,
            commit_size: 1_000_000usize,
            threads: 4,
        }
    }

    pub fn set_allow_missing(&mut self, allow_missing: bool) -> () {
        self.allow_missing = allow_missing;
    }

    pub fn set_commit_size(&mut self, commit_size: usize) -> () {
        self.commit_size = commit_size;
    }
}

pub fn index(spec: &Spec, data_type: &str, filename: &Path) -> Result<(), Box<dyn Error>> {
    let mut reader = rnc_utils::buf_reader(&filename)?;

    let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
    let mut builder = Rkv::environment_builder::<Lmdb>();
    builder.set_max_dbs(20);
    builder.set_make_dir_if_needed(true);
    builder.set_map_size(200000000000);
    let arc = manager
        .get_or_create_from_builder(spec.path, builder, Rkv::from_builder::<Lmdb>)
        .with_context(|| "Failed to create arc")?;
    let env = arc.read().unwrap();
    let store = env
        .open_single(data_type, StoreOptions::create())
        .with_context(|| format!("Failed to open store {:?}", data_type))?;

    let mut writer = env.write().with_context(|| "Failed to create writer")?;
    let mut buf = String::new();
    let mut count = 0usize;
    loop {
        match reader.read_line(&mut buf)? {
            0 => break,
            _ => {
                let line = buf.replace("\\\\", "\\");
                let value = line.trim_end().to_owned();
                let id: DocId = serde_json::from_str::<Query<DocId>>(&line)?.into();
                let existing = store.get(&writer, &id.id)?;
                let to_add = match existing {
                    None => value,
                    Some(v) => match v {
                        Value::Json(e) => format!("{}{}", value, e),
                        _ => panic!("Invalid existing data, should never happen"),
                    },
                };

                store.put(&mut writer, &id.id, &Value::Json(&to_add))?;
                buf.clear();
                count += 1;
                if count % spec.commit_size == 0 {
                    writer
                        .commit()
                        .with_context(|| format!("Failed to write data from {:?}", filename))?;
                    writer = env.write().with_context(|| "Could not create for new chunk")?;
                }
            },
        }
    }
    writer.commit().with_context(|| format!("Failed to write data from {:?}", filename))?;

    Ok(())
}

// fn path_as_column_name(path: &Path) -> String {
//     path.file_stem().unwrap().to_str().unwrap().to_string().replace("-", "_")
// }

// fn send_file_lines(path: &Path, sender: Sender<(String, DocId, String)>) {
//     let file = File::open(path).unwrap();
//     let data_type = path_as_column_name(&path);
//     let mut reader = BufReader::new(file);
//     let mut buf = String::new();
//     loop {
//         match reader.read_line(&mut buf).unwrap() {
//             0 => break,
//             _ => {
//                 let line = buf.replace("\\\\", "\\");
//                 let id: DocId =
// serde_json::from_str::<Query<DocId>>(&line).unwrap().into();
// sender.send((data_type.to_string(), id, line.to_string())).unwrap();
// buf.clear();             },
//         }
//     }
// }

// pub fn index_files(spec: &Spec, filename: &Path) -> Result<(), Box<dyn Error>> {
//     let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
//     let arc = manager.get_or_create(spec.path, Rkv::new::<Lmdb>)?;
//     let env = arc.read();

//     // let stores = HashMap::<String, Multi>
//     let reader = rnc_utils::buf_reader(&filename)?;
//     let (json_sender, json_reciever) = unbounded();
//     for line in reader.lines() {
//         let line = line?;
//         let filename = line.trim_end();
//         let path = PathBuf::from(filename);
//         let sender = json_sender.clone();
//         let store = env
//             .open_multi(&data_type, StoreOptions::create())
//             .with_context(|| format!("Failed to open store {:?}", data_type))?;
//         stores.insert(name, store);
//         thread::spawn(move || send_file_lines(&path, sender));
//     }

//     let writer = env.write().with_context(|| "Failed to create writer")?;

//     for (data_type, id, data) in json_reciever {
//         let store = stores.get(&data_type)?;
//         store
//             .put(&mut writer, id.id, &Value::Json(data))
//             .with_context(|| format!("Failed to write data for {:}/{:?}", &data_type,
// &id))?;     }

//     Ok(())
// }

pub fn lookup(spec: &Spec, key_file: &Path, output: &Path) -> Result<(), Box<dyn Error>> {
    let mut writer = rnc_utils::buf_writer(&output)
        .with_context(|| format!("Could not open {:?} for writing", &output))?;
    let mut keys = rnc_utils::buf_reader(&key_file)
        .with_context(|| format!("Could not open key file {:?}", &key_file))?;

    // let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
    let mut builder = Rkv::environment_builder::<Lmdb>();
    builder.set_max_dbs(20);
    // let env = arc.read().unwrap();

    let mut manager = Manager::<LmdbEnvironment>::singleton().write().unwrap();
    // let created_arc = manager.get_or_create(spec.path, Rkv::new::<Lmdb>).unwrap();
    let created_arc = manager
        .get_or_create_from_builder(spec.path, builder, Rkv::from_builder::<Lmdb>)
        .with_context(|| "Failed to create arc")?;
    let env = created_arc.read().unwrap();

    let names = env.get_dbs()?.iter().filter_map(|db| db.to_owned()).collect::<Vec<String>>();

    let mut stores: Vec<(String, SingleStore<LmdbDatabase>)> = Vec::new();
    for name in names {
        let store = env
            .open_single(name.as_str(), StoreOptions::default())
            .with_context(|| format!("Failed to open store {:?}", &name))?;
        stores.push((name.to_owned(), store));
    }

    let reader = env.read().with_context(|| "Failed to create writer")?;

    let mut buf = String::new();
    loop {
        match keys.read_line(&mut buf)? {
            0 => break,
            _ => {
                let mut data = HashMap::new();
                let id = buf.trim_end().to_string();
                data.insert("id", serde_json::Value::String(id.to_string()));
                let mut seen = false;
                for (name, store) in &stores {
                    let default = serde_json::Value::Array(Vec::new());
                    let to_update = data.entry(&name).or_insert(default).as_array_mut().unwrap();
                    let current = store
                        .get(&reader, &id)
                        .with_context(|| format!("Failed to get results for {:?}", &id))?;

                    match current {
                        None => (),
                        Some(v) => match v {
                            Value::Json(raw) => {
                                seen = true;
                                let values = Deserializer::from_str(raw).into_iter::<serde_json::Value>();
                                for value in values {
                                    to_update.push(value?);
                                }
                            },
                            _ => panic!("Invalid lookup data"),
                        },
                    }

                    match (seen, spec.allow_missing) {
                        (true, _) => (),
                        (false, true) => {
                            log::info!("No data found for key {} in {}", &id, &name);
                            continue;
                        },
                        (false, false) => {
                            return Err(format!("No data found for key {} in {}", &id, &name).into());
                        },
                    }

                    serde_json::to_writer(&mut writer, &data)?;
                    writeln!(&mut writer)?;
                    buf.clear();
                }
            },
        }
    }

    Ok(())
}
<commit_msg>Correct output<commit_after>use std::{
    collections::HashMap,
    error::Error,
    fs::File,
    io::{
        BufRead,
        BufReader,
    },
    path::{
        Path,
        PathBuf,
    },
    str,
    thread,
};

use serde_query::{
    DeserializeQuery,
    Query,
};

use anyhow::{
    Context,
    Result,
};

// use crossbeam_channel::{
//     unbounded,
//     Sender,
// };

use serde_json::Deserializer;

use rkv::{
    backend::{
        BackendEnvironmentBuilder,
        Lmdb,
        LmdbDatabase,
        LmdbEnvironment,
    },
    Manager,
    Rkv,
    SingleStore,
    StoreOptions,
    Value,
};

#[derive(Debug, DeserializeQuery)]
struct DocId {
    #[query(".id")]
    id: String,
}

pub struct Spec<'a> {
    path: &'a Path,
    allow_missing: bool,
    commit_size: usize,
    threads: usize,
}

impl<'a> Spec<'a> {
    pub fn new(path: &Path) -> Spec {
        Spec {
            path,
            allow_missing: false,
            commit_size: 1_000_000usize,
            threads: 4,
        }
    }

    pub fn set_allow_missing(&mut self, allow_missing: bool) -> () {
        self.allow_missing = allow_missing;
    }

    pub fn set_commit_size(&mut self, commit_size: usize) -> () {
        self.commit_size = commit_size;
    }
}

pub fn index(spec: &Spec, data_type: &str, filename: &Path) -> Result<(), Box<dyn Error>> {
    let mut reader = rnc_utils::buf_reader(&filename)?;

    let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
    let mut builder = Rkv::environment_builder::<Lmdb>();
    builder.set_max_dbs(20);
    builder.set_make_dir_if_needed(true);
    builder.set_map_size(200000000000);
    let arc = manager
        .get_or_create_from_builder(spec.path, builder, Rkv::from_builder::<Lmdb>)
        .with_context(|| "Failed to create arc")?;
    let env = arc.read().unwrap();
    let store = env
        .open_single(data_type, StoreOptions::create())
        .with_context(|| format!("Failed to open store {:?}", data_type))?;

    let mut writer = env.write().with_context(|| "Failed to create writer")?;
    let mut buf = String::new();
    let mut count = 0usize;
    loop {
        match reader.read_line(&mut buf)? {
            0 => break,
            _ => {
                let line = buf.replace("\\\\", "\\");
                let value = line.trim_end().to_owned();
                let id: DocId = serde_json::from_str::<Query<DocId>>(&line)?.into();
                let existing = store.get(&writer, &id.id)?;
                let to_add = match existing {
                    None => value,
                    Some(v) => match v {
                        Value::Json(e) => format!("{}{}", value, e),
                        _ => panic!("Invalid existing data, should never happen"),
                    },
                };

                store.put(&mut writer, &id.id, &Value::Json(&to_add))?;
                buf.clear();
                count += 1;
                if count % spec.commit_size == 0 {
                    writer
                        .commit()
                        .with_context(|| format!("Failed to write data from {:?}", filename))?;
                    writer = env.write().with_context(|| "Could not create for new chunk")?;
                }
            },
        }
    }
    writer.commit().with_context(|| format!("Failed to write data from {:?}", filename))?;

    Ok(())
}

// fn path_as_column_name(path: &Path) -> String {
//     path.file_stem().unwrap().to_str().unwrap().to_string().replace("-", "_")
// }

// fn send_file_lines(path: &Path, sender: Sender<(String, DocId, String)>) {
//     let file = File::open(path).unwrap();
//     let data_type = path_as_column_name(&path);
//     let mut reader = BufReader::new(file);
//     let mut buf = String::new();
//     loop {
//         match reader.read_line(&mut buf).unwrap() {
//             0 => break,
//             _ => {
//                 let line = buf.replace("\\\\", "\\");
//                 let id: DocId =
// serde_json::from_str::<Query<DocId>>(&line).unwrap().into();
// sender.send((data_type.to_string(), id, line.to_string())).unwrap();
// buf.clear();             },
//         }
//     }
// }

// pub fn index_files(spec: &Spec, filename: &Path) -> Result<(), Box<dyn Error>> {
//     let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
//     let arc = manager.get_or_create(spec.path, Rkv::new::<Lmdb>)?;
//     let env = arc.read();

//     // let stores = HashMap::<String, Multi>
//     let reader = rnc_utils::buf_reader(&filename)?;
//     let (json_sender, json_reciever) = unbounded();
//     for line in reader.lines() {
//         let line = line?;
//         let filename = line.trim_end();
//         let path = PathBuf::from(filename);
//         let sender = json_sender.clone();
//         let store = env
//             .open_multi(&data_type, StoreOptions::create())
//             .with_context(|| format!("Failed to open store {:?}", data_type))?;
//         stores.insert(name, store);
//         thread::spawn(move || send_file_lines(&path, sender));
//     }

//     let writer = env.write().with_context(|| "Failed to create writer")?;

//     for (data_type, id, data) in json_reciever {
//         let store = stores.get(&data_type)?;
//         store
//             .put(&mut writer, id.id, &Value::Json(data))
//             .with_context(|| format!("Failed to write data for {:}/{:?}", &data_type,
// &id))?;     }

//     Ok(())
// }

pub fn lookup(spec: &Spec, key_file: &Path, output: &Path) -> Result<(), Box<dyn Error>> {
    let mut writer = rnc_utils::buf_writer(&output)
        .with_context(|| format!("Could not open {:?} for writing", &output))?;
    let mut keys = rnc_utils::buf_reader(&key_file)
        .with_context(|| format!("Could not open key file {:?}", &key_file))?;

    // let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
    let mut builder = Rkv::environment_builder::<Lmdb>();
    builder.set_max_dbs(20);
    // let env = arc.read().unwrap();

    let mut manager = Manager::<LmdbEnvironment>::singleton().write().unwrap();
    // let created_arc = manager.get_or_create(spec.path, Rkv::new::<Lmdb>).unwrap();
    let created_arc = manager
        .get_or_create_from_builder(spec.path, builder, Rkv::from_builder::<Lmdb>)
        .with_context(|| "Failed to create arc")?;
    let env = created_arc.read().unwrap();

    let names = env.get_dbs()?.iter().filter_map(|db| db.to_owned()).collect::<Vec<String>>();

    let mut stores: Vec<(String, SingleStore<LmdbDatabase>)> = Vec::new();
    for name in names {
        let store = env
            .open_single(name.as_str(), StoreOptions::default())
            .with_context(|| format!("Failed to open store {:?}", &name))?;
        stores.push((name.to_owned(), store));
    }

    let reader = env.read().with_context(|| "Failed to create writer")?;

    let mut buf = String::new();
    loop {
        match keys.read_line(&mut buf)? {
            0 => break,
            _ => {
                let mut data = HashMap::new();
                let id = buf.trim_end().to_string();
                data.insert("id", serde_json::Value::String(id.to_string()));
                let mut seen = false;
                for (name, store) in &stores {
                    let default = serde_json::Value::Array(Vec::new());
                    let to_update = data.entry(&name).or_insert(default).as_array_mut().unwrap();
                    let current = store
                        .get(&reader, &id)
                        .with_context(|| format!("Failed to get results for {:?}", &id))?;

                    match current {
                        None => (),
                        Some(v) => match v {
                            Value::Json(raw) => {
                                seen = true;
                                let values = Deserializer::from_str(raw).into_iter::<serde_json::Value>();
                                for value in values {
                                    to_update.push(value?);
                                }
                            },
                            _ => panic!("Invalid lookup data"),
                        },
                    }

                    match (seen, spec.allow_missing) {
                        (true, _) => (),
                        (false, true) => {
                            log::info!("No data found for key {} in {}", &id, &name);
                            continue;
                        },
                        (false, false) => {
                            return Err(format!("No data found for key {} in {}", &id, &name).into());
                        },
                    }
                }

                serde_json::to_writer(&mut writer, &data)?;
                writeln!(&mut writer)?;
                buf.clear();
            },
        }
    }

    Ok(())
}
<|endoftext|>


<commit_before>use std::{
    collections::HashMap,
    error::Error,
    fs::File,
    io::{
        BufRead,
        BufReader,
    },
    path::{
        Path,
        PathBuf,
    },
    str,
    thread,
};

use serde_query::{
    DeserializeQuery,
    Query,
};

use anyhow::{
    Context,
    Result,
};

// use crossbeam_channel::{
//     unbounded,
//     Sender,
// };

use serde_json::Deserializer;

use rkv::{
    backend::{
        BackendEnvironmentBuilder,
        Lmdb,
        LmdbDatabase,
        LmdbEnvironment,
    },
    Manager,
    Rkv,
    SingleStore,
    StoreOptions,
    Value,
};

#[derive(Debug, DeserializeQuery)]
struct DocId {
    #[query(".id")]
    id: String,
}

pub struct Spec<'a> {
    path: &'a Path,
    allow_missing: bool,
    commit_size: usize,
    threads: usize,
}

impl<'a> Spec<'a> {
    pub fn new(path: &Path) -> Spec {
        Spec {
            path,
            allow_missing: false,
            commit_size: 1_000_000usize,
            threads: 4,
        }
    }

    pub fn set_allow_missing(&mut self, allow_missing: bool) -> () {
        self.allow_missing = allow_missing;
    }

    pub fn set_commit_size(&mut self, commit_size: usize) -> () {
        self.commit_size = commit_size;
    }
}

pub fn index(spec: &Spec, data_type: &str, filename: &Path) -> Result<(), Box<dyn Error>> {
    let mut reader = rnc_utils::buf_reader(&filename)?;

    let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
    let mut builder = Rkv::environment_builder::<Lmdb>();
    builder.set_max_dbs(20);
    builder.set_make_dir_if_needed(true);
    builder.set_map_size(200000000000);
    let arc = manager
        .get_or_create_from_builder(spec.path, builder, Rkv::from_builder::<Lmdb>)
        .with_context(|| "Failed to create arc")?;
    let env = arc.read().unwrap();
    let store = env
        .open_single(data_type, StoreOptions::create())
        .with_context(|| format!("Failed to open store {:?}", data_type))?;

    let mut writer = env.write().with_context(|| "Failed to create writer")?;
    let mut buf = String::new();
    let mut count = 0usize;
    loop {
        match reader.read_line(&mut buf)? {
            0 => break,
            _ => {
                let line = buf.replace("\\\\", "\\");
                let value = line.trim_end().to_owned();
                let id: DocId = serde_json::from_str::<Query<DocId>>(&line)?.into();
                let existing = store.get(&writer, &id.id)?;
                let to_add = match existing {
                    None => value,
                    Some(v) => match v {
                        Value::Json(e) => format!("{}{}", value, e),
                        _ => panic!("Invalid existing data, should never happen"),
                    },
                };

                store.put(&mut writer, &id.id, &Value::Json(&to_add))?;
                buf.clear();
                count += 1;
                if count % spec.commit_size == 0 {
                    println!("Commiting count: {}", count);
                    writer
                        .commit()
                        .with_context(|| format!("Failed to write data from {:?}", filename))?;
                    writer = env.write().with_context(|| "Could not create for new chunk")?;
                }
            },
        }
    }
    writer.commit().with_context(|| format!("Failed to write data from {:?}", filename))?;

    Ok(())
}

// fn path_as_column_name(path: &Path) -> String {
//     path.file_stem().unwrap().to_str().unwrap().to_string().replace("-", "_")
// }

// fn send_file_lines(path: &Path, sender: Sender<(String, DocId, String)>) {
//     let file = File::open(path).unwrap();
//     let data_type = path_as_column_name(&path);
//     let mut reader = BufReader::new(file);
//     let mut buf = String::new();
//     loop {
//         match reader.read_line(&mut buf).unwrap() {
//             0 => break,
//             _ => {
//                 let line = buf.replace("\\\\", "\\");
//                 let id: DocId =
// serde_json::from_str::<Query<DocId>>(&line).unwrap().into();
// sender.send((data_type.to_string(), id, line.to_string())).unwrap();
// buf.clear();             },
//         }
//     }
// }

// pub fn index_files(spec: &Spec, filename: &Path) -> Result<(), Box<dyn Error>> {
//     let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
//     let arc = manager.get_or_create(spec.path, Rkv::new::<Lmdb>)?;
//     let env = arc.read();

//     // let stores = HashMap::<String, Multi>
//     let reader = rnc_utils::buf_reader(&filename)?;
//     let (json_sender, json_reciever) = unbounded();
//     for line in reader.lines() {
//         let line = line?;
//         let filename = line.trim_end();
//         let path = PathBuf::from(filename);
//         let sender = json_sender.clone();
//         let store = env
//             .open_multi(&data_type, StoreOptions::create())
//             .with_context(|| format!("Failed to open store {:?}", data_type))?;
//         stores.insert(name, store);
//         thread::spawn(move || send_file_lines(&path, sender));
//     }

//     let writer = env.write().with_context(|| "Failed to create writer")?;

//     for (data_type, id, data) in json_reciever {
//         let store = stores.get(&data_type)?;
//         store
//             .put(&mut writer, id.id, &Value::Json(data))
//             .with_context(|| format!("Failed to write data for {:}/{:?}", &data_type,
// &id))?;     }

//     Ok(())
// }

pub fn lookup(spec: &Spec, key_file: &Path, output: &Path) -> Result<(), Box<dyn Error>> {
    let mut writer = rnc_utils::buf_writer(&output)
        .with_context(|| format!("Could not open {:?} for writing", &output))?;
    let mut keys = rnc_utils::buf_reader(&key_file)
        .with_context(|| format!("Could not open key file {:?}", &key_file))?;

    // let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
    // let mut builder = Rkv::environment_builder::<Lmdb>();
    // builder.set_max_dbs(20);
    // builder.set_make_dir_if_needed(true);
    // builder.set_map_size(200000000000);
    // println!("Created BUILDER");
    // let arc = manager
    //     .get_or_create_from_builder(spec.path, builder, Rkv::from_builder::<Lmdb>)
    //     .with_context(|| "Failed to create arc")?;
    // let env = arc.read().unwrap();

    let mut manager = Manager::<LmdbEnvironment>::singleton().write().unwrap();
    let created_arc = manager.get_or_create(spec.path, Rkv::new::<Lmdb>).unwrap();
    let env = created_arc.read().unwrap();

    let names = env.get_dbs()?.iter().filter_map(|db| db.to_owned()).collect::<Vec<String>>();

    let mut stores: Vec<(String, SingleStore<LmdbDatabase>)> = Vec::new();
    for name in names {
        let store = env
            .open_single(name.as_str(), StoreOptions::default())
            .with_context(|| format!("Failed to open store {:?}", &name))?;
        stores.push((name.to_owned(), store));
    }

    let reader = env.read().with_context(|| "Failed to create writer")?;

    let mut buf = String::new();
    loop {
        match keys.read_line(&mut buf)? {
            0 => break,
            _ => {
                let mut data = HashMap::new();
                let id = buf.trim_end().to_string();
                data.insert("id", serde_json::Value::String(id.to_string()));
                let mut seen = false;
                for (name, store) in &stores {
                    let default = serde_json::Value::Array(Vec::new());
                    let to_update = data.entry(&name).or_insert(default).as_array_mut().unwrap();
                    let current = store
                        .get(&reader, &id)
                        .with_context(|| format!("Failed to get results for {:?}", &id))?;

                    match current {
                        None => (),
                        Some(v) => match v {
                            Value::Json(raw) => {
                                seen = true;
                                let values = Deserializer::from_str(raw).into_iter::<serde_json::Value>();
                                for value in values {
                                    to_update.push(value?);
                                }
                            },
                            _ => panic!("Invalid lookup data"),
                        },
                    }

                    match (seen, spec.allow_missing) {
                        (true, _) => (),
                        (false, true) => {
                            log::warn!("No data found for key {}", &buf);
                            continue;
                        },
                        (false, false) => {
                            return Err(format!("No data found for key {}", &id).into());
                        },
                    }

                    serde_json::to_writer(&mut writer, &data)?;
                    writeln!(&mut writer)?;
                    buf.clear();
                }
            },
        }
    }

    Ok(())
}
<commit_msg>Fix opening many dbs and logging<commit_after>use std::{
    collections::HashMap,
    error::Error,
    fs::File,
    io::{
        BufRead,
        BufReader,
    },
    path::{
        Path,
        PathBuf,
    },
    str,
    thread,
};

use serde_query::{
    DeserializeQuery,
    Query,
};

use anyhow::{
    Context,
    Result,
};

// use crossbeam_channel::{
//     unbounded,
//     Sender,
// };

use serde_json::Deserializer;

use rkv::{
    backend::{
        BackendEnvironmentBuilder,
        Lmdb,
        LmdbDatabase,
        LmdbEnvironment,
    },
    Manager,
    Rkv,
    SingleStore,
    StoreOptions,
    Value,
};

#[derive(Debug, DeserializeQuery)]
struct DocId {
    #[query(".id")]
    id: String,
}

pub struct Spec<'a> {
    path: &'a Path,
    allow_missing: bool,
    commit_size: usize,
    threads: usize,
}

impl<'a> Spec<'a> {
    pub fn new(path: &Path) -> Spec {
        Spec {
            path,
            allow_missing: false,
            commit_size: 1_000_000usize,
            threads: 4,
        }
    }

    pub fn set_allow_missing(&mut self, allow_missing: bool) -> () {
        self.allow_missing = allow_missing;
    }

    pub fn set_commit_size(&mut self, commit_size: usize) -> () {
        self.commit_size = commit_size;
    }
}

pub fn index(spec: &Spec, data_type: &str, filename: &Path) -> Result<(), Box<dyn Error>> {
    let mut reader = rnc_utils::buf_reader(&filename)?;

    let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
    let mut builder = Rkv::environment_builder::<Lmdb>();
    builder.set_max_dbs(20);
    builder.set_make_dir_if_needed(true);
    builder.set_map_size(200000000000);
    let arc = manager
        .get_or_create_from_builder(spec.path, builder, Rkv::from_builder::<Lmdb>)
        .with_context(|| "Failed to create arc")?;
    let env = arc.read().unwrap();
    let store = env
        .open_single(data_type, StoreOptions::create())
        .with_context(|| format!("Failed to open store {:?}", data_type))?;

    let mut writer = env.write().with_context(|| "Failed to create writer")?;
    let mut buf = String::new();
    let mut count = 0usize;
    loop {
        match reader.read_line(&mut buf)? {
            0 => break,
            _ => {
                let line = buf.replace("\\\\", "\\");
                let value = line.trim_end().to_owned();
                let id: DocId = serde_json::from_str::<Query<DocId>>(&line)?.into();
                let existing = store.get(&writer, &id.id)?;
                let to_add = match existing {
                    None => value,
                    Some(v) => match v {
                        Value::Json(e) => format!("{}{}", value, e),
                        _ => panic!("Invalid existing data, should never happen"),
                    },
                };

                store.put(&mut writer, &id.id, &Value::Json(&to_add))?;
                buf.clear();
                count += 1;
                if count % spec.commit_size == 0 {
                    writer
                        .commit()
                        .with_context(|| format!("Failed to write data from {:?}", filename))?;
                    writer = env.write().with_context(|| "Could not create for new chunk")?;
                }
            },
        }
    }
    writer.commit().with_context(|| format!("Failed to write data from {:?}", filename))?;

    Ok(())
}

// fn path_as_column_name(path: &Path) -> String {
//     path.file_stem().unwrap().to_str().unwrap().to_string().replace("-", "_")
// }

// fn send_file_lines(path: &Path, sender: Sender<(String, DocId, String)>) {
//     let file = File::open(path).unwrap();
//     let data_type = path_as_column_name(&path);
//     let mut reader = BufReader::new(file);
//     let mut buf = String::new();
//     loop {
//         match reader.read_line(&mut buf).unwrap() {
//             0 => break,
//             _ => {
//                 let line = buf.replace("\\\\", "\\");
//                 let id: DocId =
// serde_json::from_str::<Query<DocId>>(&line).unwrap().into();
// sender.send((data_type.to_string(), id, line.to_string())).unwrap();
// buf.clear();             },
//         }
//     }
// }

// pub fn index_files(spec: &Spec, filename: &Path) -> Result<(), Box<dyn Error>> {
//     let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
//     let arc = manager.get_or_create(spec.path, Rkv::new::<Lmdb>)?;
//     let env = arc.read();

//     // let stores = HashMap::<String, Multi>
//     let reader = rnc_utils::buf_reader(&filename)?;
//     let (json_sender, json_reciever) = unbounded();
//     for line in reader.lines() {
//         let line = line?;
//         let filename = line.trim_end();
//         let path = PathBuf::from(filename);
//         let sender = json_sender.clone();
//         let store = env
//             .open_multi(&data_type, StoreOptions::create())
//             .with_context(|| format!("Failed to open store {:?}", data_type))?;
//         stores.insert(name, store);
//         thread::spawn(move || send_file_lines(&path, sender));
//     }

//     let writer = env.write().with_context(|| "Failed to create writer")?;

//     for (data_type, id, data) in json_reciever {
//         let store = stores.get(&data_type)?;
//         store
//             .put(&mut writer, id.id, &Value::Json(data))
//             .with_context(|| format!("Failed to write data for {:}/{:?}", &data_type,
// &id))?;     }

//     Ok(())
// }

pub fn lookup(spec: &Spec, key_file: &Path, output: &Path) -> Result<(), Box<dyn Error>> {
    let mut writer = rnc_utils::buf_writer(&output)
        .with_context(|| format!("Could not open {:?} for writing", &output))?;
    let mut keys = rnc_utils::buf_reader(&key_file)
        .with_context(|| format!("Could not open key file {:?}", &key_file))?;

    // let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
    let mut builder = Rkv::environment_builder::<Lmdb>();
    builder.set_max_dbs(20);
    // let env = arc.read().unwrap();

    let mut manager = Manager::<LmdbEnvironment>::singleton().write().unwrap();
    // let created_arc = manager.get_or_create(spec.path, Rkv::new::<Lmdb>).unwrap();
    let created_arc = manager
        .get_or_create_from_builder(spec.path, builder, Rkv::from_builder::<Lmdb>)
        .with_context(|| "Failed to create arc")?;
    let env = created_arc.read().unwrap();

    let names = env.get_dbs()?.iter().filter_map(|db| db.to_owned()).collect::<Vec<String>>();

    let mut stores: Vec<(String, SingleStore<LmdbDatabase>)> = Vec::new();
    for name in names {
        let store = env
            .open_single(name.as_str(), StoreOptions::default())
            .with_context(|| format!("Failed to open store {:?}", &name))?;
        stores.push((name.to_owned(), store));
    }

    let reader = env.read().with_context(|| "Failed to create writer")?;

    let mut buf = String::new();
    loop {
        match keys.read_line(&mut buf)? {
            0 => break,
            _ => {
                let mut data = HashMap::new();
                let id = buf.trim_end().to_string();
                data.insert("id", serde_json::Value::String(id.to_string()));
                let mut seen = false;
                for (name, store) in &stores {
                    let default = serde_json::Value::Array(Vec::new());
                    let to_update = data.entry(&name).or_insert(default).as_array_mut().unwrap();
                    let current = store
                        .get(&reader, &id)
                        .with_context(|| format!("Failed to get results for {:?}", &id))?;

                    match current {
                        None => (),
                        Some(v) => match v {
                            Value::Json(raw) => {
                                seen = true;
                                let values = Deserializer::from_str(raw).into_iter::<serde_json::Value>();
                                for value in values {
                                    to_update.push(value?);
                                }
                            },
                            _ => panic!("Invalid lookup data"),
                        },
                    }

                    match (seen, spec.allow_missing) {
                        (true, _) => (),
                        (false, true) => {
                            log::info!("No data found for key {} in {}", &id, &name);
                            continue;
                        },
                        (false, false) => {
                            return Err(format!("No data found for key {} in {}", &id, &name).into());
                        },
                    }

                    serde_json::to_writer(&mut writer, &data)?;
                    writeln!(&mut writer)?;
                    buf.clear();
                }
            },
        }
    }

    Ok(())
}
<|endoftext|>


<commit_before>#[crate_id = "sphere4d"];
#[crate_type = "bin"];
#[warn(non_camel_case_types)];
#[feature(managed_boxes)];

extern mod extra;
extern mod nalgebra;
extern mod ncollide = "ncollide4df64";
extern mod nrays    = "nrays4d";

use extra::arc::Arc;
use std::io::fs::File;
use nalgebra::na::{Iso4, Vec3, Vec4};
use nalgebra::na;
use ncollide::ray::Ray;
use ncollide::geom::{Ball, Box, Cone, Cylinder};
use nrays::scene_node::SceneNode;
use nrays::material::Material;
use nrays::phong_material::PhongMaterial;
use nrays::scene::Scene;
use nrays::light::Light;

fn main() {
    let resolution = Vec3::new(100.0f64, 100.0, 100.0);

    let mut lights = ~[];

    {
        lights.push(Light::new(Vec4::new(0.0f64, 10.0, 10.0, 1.0),
                               0.0,
                               1,
                               Vec3::new(1.0, 1.0, 1.0)));
        lights.push(Light::new(Vec4::new(0.0f64, 10.0, 10.0, 1.0),
                               0.0,
                               1,
                               Vec3::new(1.0, 1.0, 1.0)));
        /*
        lights.push(Light::new(Vec4::new(0.0f64, -10.0, -10.0, 1.0),
                               Vec3::new(1.0, 1.0, 1.0)));
        lights.push(Light::new(Vec4::new(0.0f64, -10.0, 10.0, 1.0),
                               Vec3::new(1.0, 1.0, 1.0)));
        lights.push(Light::new(Vec4::new(0.0f64, 10.0, -10.0, 1.0),
                               Vec3::new(1.0, 1.0, 1.0)));
        lights.push(Light::new(Vec4::new(0.0f64, 10.0, 10.0, 1.0),
                               Vec3::new(1.0, 1.0, 1.0)));
                               */
    }

    let blue = Arc::new(~PhongMaterial::new(
        Vec3::new(0.1, 0.1, 0.1),
        Vec3::new(1.0, 1.0, 1.0),
        Vec3::new(1.0, 1.0, 1.0),
        None,
        None,
        100.0
    ) as ~Material:Freeze+Send);

    let transform: Iso4<f64> = na::one();

    let box_shape = ~Box::new_with_margin(Vec4::new(0.25, 0.25, 0.25, 0.00000001), 0.000000001);
    let ball      = ~Ball::new(0.25);
    let cone      = ~Cone::new_with_margin(0.25, 0.25, 0.0);
    let cylinder  = ~Cylinder::new_with_margin(0.25, 0.25, 0.0);

    let pos  = na::append_translation(&transform, &Vec4::new(0.0, 0.0, 0.0,    4.0));
    let pos2 = na::append_translation(&transform, &Vec4::new(0.75, 0.75, 0.0,  4.0));
    let pos3 = na::append_translation(&transform, &Vec4::new(0.0, 0.75, 0.75,  4.0));
    let pos4 = na::append_translation(&transform, &Vec4::new(0.0, 0.75, -0.75, 4.0));

    let mut nodes = ~[];
    nodes.push(Arc::new(SceneNode::new(blue.clone(), 0.4, 0.2, 1.0, 1.0, pos,  ball, None, true)));
    nodes.push(Arc::new(SceneNode::new(blue.clone(), 0.4, 0.2, 1.0, 1.0, pos2, box_shape, None, true)));
    nodes.push(Arc::new(SceneNode::new(blue.clone(), 0.4, 0.2, 1.0, 1.0, pos3, cone, None, true)));
    nodes.push(Arc::new(SceneNode::new(blue.clone(), 0.4, 0.2, 1.0, 1.0, pos4, cylinder, None, true)));

    let scene = Scene::new(nodes, lights);
    let pixels = scene.render(&resolution, |pt| {
        let x = (pt.x / resolution.x - 0.5) * 2.0;
        let y = (pt.y / resolution.y - 0.5) * 2.0;
        let z = (pt.z / resolution.z - 0.5) * 2.0;
        // Ray::new(eye, na::normalize(&(Vec4::new(pt.x, pt.y, pt.z, 1.0) - eye)))

        Ray::new(Vec4::new(x, y, z, 0.0), Vec4::new(0.0, 0.0, 0.0, 1.0))
    });

    let path = "out.4d";
    let mut file = File::create(&Path::new(path)).expect("Cannot create the file: " + path);
    // let mut file = BufferedWriter::new(file);
    pixels.to_file(&mut file);
}
<commit_msg>Add perspective to the 4d demo.<commit_after>#[crate_id = "sphere4d"];
#[crate_type = "bin"];
#[warn(non_camel_case_types)];
#[feature(managed_boxes)];

extern mod extra;
extern mod nalgebra;
extern mod ncollide = "ncollide4df64";
extern mod nrays    = "nrays4d";

use extra::arc::Arc;
use std::io::fs::File;
use nalgebra::na::{Iso4, Vec3, Vec4};
use nalgebra::na;
use ncollide::ray::Ray;
use ncollide::geom::{Ball, Box, Cone, Cylinder};
use nrays::scene_node::SceneNode;
use nrays::material::Material;
use nrays::phong_material::PhongMaterial;
use nrays::scene::Scene;
use nrays::light::Light;

fn main() {
    let resolution = Vec3::new(100.0f64, 100.0, 100.0);

    let mut lights = ~[];

    {
        lights.push(Light::new(Vec4::new(10.0f64, -10.0, 10.0, 1.0),
                               0.0,
                               1,
                               Vec3::new(2.0, 2.0, 2.0)));
        lights.push(Light::new(Vec4::new(-10.0f64, -10.0, 10.0, 1.0),
                               0.0,
                               1,
                               Vec3::new(2.0, 2.0, 2.0)));
        lights.push(Light::new(Vec4::new(10.0f64, 10.0, 10.0, 1.0),
                               0.0,
                               1,
                               Vec3::new(2.0, 2.0, 2.0)));
        lights.push(Light::new(Vec4::new(-10.0f64, 10.0, 10.0, 1.0),
                               0.0,
                               1,
                               Vec3::new(2.0, 2.0, 2.0)));
        // lights.push(Light::new(Vec4::new(0.0f64, 10.0, 10.0, 1.0),
        //                        0.0,
        //                        1,
        //                        Vec3::new(1.0, 1.0, 1.0)));
    }

    let blue = Arc::new(~PhongMaterial::new(
        Vec3::new(1.0, 1.0, 1.0),
        Vec3::new(1.0, 1.0, 1.0),
        Vec3::new(1.0, 1.0, 1.0),
        None,
        None,
        60.0
    ) as ~Material:Freeze+Send);

    let transform: Iso4<f64> = na::one();

    let box_shape = ~Box::new_with_margin(Vec4::new(0.25, 0.25, 0.25, 0.25), 0.0);
    let ball      = ~Ball::new(0.25);
    let cone      = ~Cone::new_with_margin(0.25, 0.25, 0.0);
    let cylinder  = ~Cylinder::new_with_margin(0.25, 0.25, 0.0);

    let pos  = na::append_translation(&transform, &Vec4::new(0.0, 0.0, 0.0,    1.5));
    let pos2 = na::append_translation(&transform, &Vec4::new(0.75, 0.75, 0.0,  1.5));
    let pos3 = na::append_translation(&transform, &Vec4::new(0.0, 0.75, 0.75,  1.5));
    let pos4 = na::append_translation(&transform, &Vec4::new(0.0, 0.75, -0.75, 1.5));

    let mut nodes = ~[];
    nodes.push(Arc::new(SceneNode::new(blue.clone(), 0.0, 0.0, 1.0, 1.0, pos,  ball, None, true)));
    nodes.push(Arc::new(SceneNode::new(blue.clone(), 0.0, 0.0, 1.0, 1.0, pos2, box_shape, None, true)));
    nodes.push(Arc::new(SceneNode::new(blue.clone(), 0.0, 0.0, 1.0, 1.0, pos3, cone, None, true)));
    nodes.push(Arc::new(SceneNode::new(blue.clone(), 0.0, 0.0, 1.0, 1.0, pos4, cylinder, None, true)));

    let scene = Scene::new(nodes, lights);
    let pixels = scene.render(&resolution, |pt| {
        let x = (pt.x / resolution.x - 0.5) * 2.0;
        let y = (pt.y / resolution.y - 0.5) * 2.0;
        let z = (pt.z / resolution.z - 0.5) * 2.0;
        let eye = Vec4::new(0.0, 0.0, 0.0, 0.0);
        Ray::new(eye, na::normalize(&(Vec4::new(x, y, z, 1.0) - eye)))

        // Ray::new(Vec4::new(x, y, z, 0.0), Vec4::new(0.0, 0.0, 0.0, 1.0))
    });

    let path = "out.4d";
    let mut file = File::create(&Path::new(path)).expect("Cannot create the file: " + path);
    // let mut file = BufferedWriter::new(file);
    pixels.to_file(&mut file);
}
<|endoftext|>


<commit_before>use aes::BLOCK_SIZE;
use errors::*;
use helper::ceil_quotient;
use std;

use prefix_suffix_oracles::Oracle12;
use prefix_suffix_oracles::{DeterministicOracle, Oracle};

fn uses_padding<T: Oracle>(oracle: &T) -> Result<bool, Error> {
    Ok((oracle.encrypt(&[0])?.len() - oracle.encrypt(&[])?.len()) % BLOCK_SIZE == 0)
}

pub fn prefix_plus_suffix_length<T: Oracle>(oracle: &T) -> Result<usize, Error> {
    let initial = oracle.encrypt(&[])?.len();
    if !uses_padding(oracle)? {
        return Ok(initial);
    }

    let input = [0; BLOCK_SIZE];
    if let Some(index) = (1..=BLOCK_SIZE).find(|&i| {
        if let Ok(ciphertext) = oracle.encrypt(&input[..i]) {
            initial != ciphertext.len()
        } else {
            false
        }
    }) {
        Ok(initial - index)
    } else {
        bail!(
            "length of oracle output did not change, something is wrong with the provided oracle"
        );
    }
}

// For an oracle prepending prefix and appending suffix to its input, this function returns
// prefix.len()/BLOCK_SIZE, that is the number of blocks fully occupied by the prefix.
//
// To determine this number, we pass two different cleartexts to the oracle and count the number
// of identical blocks at the start of the corresponding ciphertexts.

fn full_prefix_blocks_count<T: DeterministicOracle>(oracle: &T) -> Result<usize, Error> {
    if let Some(result) = oracle
        .encrypt(&[0])?
        .chunks(BLOCK_SIZE)
        .zip(oracle.encrypt(&[1])?.chunks(BLOCK_SIZE))
        .position(|(x, y)| x != y)
    {
        Ok(result)
    } else {
        bail!("no differing blocks found, something is wrong with the provided oracle");
    }
}

// We look at the first block C not fully occupied by the prefix and fill it with a
// constant block B, say consisting of 0's.
// This part of the cleartext in oracle looks as follows:
//
//                               <------ B ------->
// prefix[?] prefix[?] prefix[?] 0 0 ... 0 || 0 0 0 suffix[0] suffix[1]
// <---------- C ------------------------>
//
// We then successively reduce the length of B until the ciphertext of C
// changes. This happens as soon as the cleartext in oracle looks as follows:
//
//                               <-- B -->
// prefix[?] prefix[?] prefix[?] 0 0 ... 0 suffix[0] || suffix[1]
// <---------- C ---------------------------------->
//
// This gives us the length of the prefix in C.
//
// We need to do this with two different constants because suffix[0] might accidentally
// coincide with the constant we have chosen.

pub fn prefix_length<T: DeterministicOracle>(oracle: &T) -> Result<usize, Error> {
    let n = full_prefix_blocks_count(oracle)?;
    let helper = |k: u8| -> Result<usize, Error> {
        let constant_block = vec![k; BLOCK_SIZE];

        let mut prev = oracle.encrypt(&constant_block)?;

        for i in 0..BLOCK_SIZE {
            let cur = oracle.encrypt(&constant_block[i + 1..])?;
            if prev.chunks(BLOCK_SIZE).nth(n) != cur.chunks(BLOCK_SIZE).nth(n) {
                return Ok(i);
            }
            prev = cur;
        }
        Ok(BLOCK_SIZE)
    };

    Ok(n * BLOCK_SIZE + std::cmp::min(helper(0)?, helper(1)?))
}

pub fn suffix_length<T: DeterministicOracle>(oracle: &T) -> Result<usize, Error> {
    Ok(prefix_plus_suffix_length(oracle)? - prefix_length(oracle)?)
}

pub fn decrypt_suffix<T: DeterministicOracle>(oracle: &T) -> Result<Vec<u8>, Error> {
    // The following input is chosen in such a way that the cleartext in oracle looks as follows:
    //
    //            input start      input end
    //                ↓                ↓
    // <-- prefix --> 0 ... 0 || 0 ... 0 suffix[0] || suffix[1] ...
    //                ↑          ↑
    //            prefix_len  prefix_blocks*BLOCK_SIZE
    //
    // The resulting ciphertext is compared to oracle([input, u]). The u yielding a match is
    // equal to suffix[0].

    let prefix_len = prefix_length(oracle)?;
    let (prefix_blocks, prefix_padding) = ceil_quotient(prefix_len, BLOCK_SIZE);
    let suffix_len = suffix_length(oracle)?;

    let mut suffix = Vec::with_capacity(suffix_len);

    let mut input = vec![0; prefix_padding + BLOCK_SIZE - 1];
    let reference_ciphertexts = (0..BLOCK_SIZE)
        .map(|left_shift| oracle.encrypt(&input[left_shift..]))
        .collect::<Result<Vec<Vec<u8>>, Error>>()?;

    for i in 0..suffix_len {
        let block_index = prefix_blocks + i / BLOCK_SIZE;
        let left_shift = i % BLOCK_SIZE;
        for u in 0u8..=255 {
            input.push(u);
            if reference_ciphertexts[left_shift][block_index * BLOCK_SIZE..(block_index + 1) * BLOCK_SIZE]
                == oracle.encrypt(&input[left_shift..])?
                    [block_index * BLOCK_SIZE..(block_index + 1) * BLOCK_SIZE]
            {
                suffix.push(u);
                break;
            }
            input.pop();
        }
    }
    Ok(suffix)
}

pub fn run() -> Result<(), Error> {
    let oracle = Oracle12::new()?;
    oracle.verify_suffix(&decrypt_suffix(&oracle)?)
}
<commit_msg>Minor improvement<commit_after>use aes::BLOCK_SIZE;
use errors::*;
use helper::ceil_quotient;
use std;

use prefix_suffix_oracles::Oracle12;
use prefix_suffix_oracles::{DeterministicOracle, Oracle};

fn uses_padding<T: Oracle>(oracle: &T) -> Result<bool, Error> {
    Ok((oracle.encrypt(&[0])?.len() - oracle.encrypt(&[])?.len()) % BLOCK_SIZE == 0)
}

pub fn prefix_plus_suffix_length<T: Oracle>(oracle: &T) -> Result<usize, Error> {
    let initial = oracle.encrypt(&[])?.len();
    if !uses_padding(oracle)? {
        return Ok(initial);
    }

    let input = [0; BLOCK_SIZE];
    if let Some(index) = (1..=BLOCK_SIZE).find(|&i| {
        if let Ok(ciphertext) = oracle.encrypt(&input[..i]) {
            initial != ciphertext.len()
        } else {
            false
        }
    }) {
        Ok(initial - index)
    } else {
        bail!(
            "length of oracle output did not change, something is wrong with the provided oracle"
        );
    }
}

// For an oracle prepending prefix and appending suffix to its input, this function returns
// prefix.len()/BLOCK_SIZE, that is the number of blocks fully occupied by the prefix.
//
// To determine this number, we pass two different cleartexts to the oracle and count the number
// of identical blocks at the start of the corresponding ciphertexts.

fn full_prefix_blocks_count<T: DeterministicOracle>(oracle: &T) -> Result<usize, Error> {
    if let Some(result) = oracle
        .encrypt(&[0])?
        .chunks(BLOCK_SIZE)
        .zip(oracle.encrypt(&[1])?.chunks(BLOCK_SIZE))
        .position(|(x, y)| x != y)
    {
        Ok(result)
    } else {
        bail!("no differing blocks found, something is wrong with the provided oracle");
    }
}

// We look at the first block C not fully occupied by the prefix and fill it with a
// constant block B, say consisting of 0's.
// This part of the cleartext in oracle looks as follows:
//
//                               <------ B ------->
// prefix[?] prefix[?] prefix[?] 0 0 ... 0 || 0 0 0 suffix[0] suffix[1]
// <---------- C ------------------------>
//
// We then successively reduce the length of B until the ciphertext of C
// changes. This happens as soon as the cleartext in oracle looks as follows:
//
//                               <-- B -->
// prefix[?] prefix[?] prefix[?] 0 0 ... 0 suffix[0] || suffix[1]
// <---------- C ---------------------------------->
//
// This gives us the length of the prefix in C.
//
// We need to do this with two different constants because suffix[0] might accidentally
// coincide with the constant we have chosen.

pub fn prefix_length<T: DeterministicOracle>(oracle: &T) -> Result<usize, Error> {
    let offset = full_prefix_blocks_count(oracle)? * BLOCK_SIZE;
    let helper = |k: u8| -> Result<usize, Error> {
        let constant_block = vec![k; BLOCK_SIZE];
        let initial = &oracle.encrypt(&constant_block)?[offset..(offset + BLOCK_SIZE)];
        for i in 0..BLOCK_SIZE {
            let cur = oracle.encrypt(&constant_block[i + 1..])?;
            if cur.len() < offset + BLOCK_SIZE || initial != &cur[offset..(offset + BLOCK_SIZE)] {
                return Ok(i);
            }
        }
        Ok(BLOCK_SIZE)
    };

    Ok(offset + std::cmp::min(helper(0)?, helper(1)?))
}

pub fn suffix_length<T: DeterministicOracle>(oracle: &T) -> Result<usize, Error> {
    Ok(prefix_plus_suffix_length(oracle)? - prefix_length(oracle)?)
}

pub fn decrypt_suffix<T: DeterministicOracle>(oracle: &T) -> Result<Vec<u8>, Error> {
    // The following input is chosen in such a way that the cleartext in oracle looks as follows:
    //
    //            input start      input end
    //                ↓                ↓
    // <-- prefix --> 0 ... 0 || 0 ... 0 suffix[0] || suffix[1] ...
    //                ↑          ↑
    //            prefix_len  prefix_blocks*BLOCK_SIZE
    //
    // The resulting ciphertext is compared to oracle([input, u]). The u yielding a match is
    // equal to suffix[0].

    let prefix_len = prefix_length(oracle)?;
    let (prefix_blocks, prefix_padding) = ceil_quotient(prefix_len, BLOCK_SIZE);
    let suffix_len = suffix_length(oracle)?;

    let mut suffix = Vec::with_capacity(suffix_len);

    let mut input = vec![0; prefix_padding + BLOCK_SIZE - 1];
    let reference_ciphertexts = (0..BLOCK_SIZE)
        .map(|left_shift| oracle.encrypt(&input[left_shift..]))
        .collect::<Result<Vec<Vec<u8>>, Error>>()?;

    for i in 0..suffix_len {
        let block_index = prefix_blocks + i / BLOCK_SIZE;
        let left_shift = i % BLOCK_SIZE;
        for u in 0u8..=255 {
            input.push(u);
            if reference_ciphertexts[left_shift]
                [block_index * BLOCK_SIZE..(block_index + 1) * BLOCK_SIZE]
                == oracle.encrypt(&input[left_shift..])?
                    [block_index * BLOCK_SIZE..(block_index + 1) * BLOCK_SIZE]
            {
                suffix.push(u);
                break;
            }
            input.pop();
        }
    }
    Ok(suffix)
}

pub fn run() -> Result<(), Error> {
    let oracle = Oracle12::new()?;
    oracle.verify_suffix(&decrypt_suffix(&oracle)?)
}
<|endoftext|>


<commit_before>use std::fmt::Debug;
use std::intrinsics::type_name;
use super::Matcher;

pub struct Nothing;
pub struct Something;

impl<T: Debug> Matcher<Option<T>> for Nothing {
    fn matches(&self, lhs: &Option<T>) -> bool {
        lhs.is_none()
    }

    fn fail_msg(&self, lhs: &Option<T>) -> String {
        format!("expected {:?} to be None", lhs)
    }

    fn negated_fail_msg(&self, lhs: &Option<T>) -> String {
        format!("expected {:?} to be Some<{}>", lhs, unsafe { type_name::<T>() })
    }
}

impl<T: Debug> Matcher<Option<T>> for Something {
    fn matches(&self, lhs: &Option<T>) -> bool {
        lhs.is_some()
    }

    fn fail_msg(&self, lhs: &Option<T>) -> String {
        format!("expected {:?} to be Some<{}>", lhs, unsafe { type_name::<T>() })
    }

    fn negated_fail_msg(&self, lhs: &Option<T>) -> String {
        format!("expected {:?} to be None", lhs)
    }
}

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

    #[test]
    fn test_none_matches() {
        let n: Option<String> = None;
        expect(n).is(none());
    }

    #[test]
    #[should_panic(expected="expected Some(\"Hello\") to be None")]
    fn test_none_fails_with_message() {
        let n = Some("Hello");
        expect(n).is(none());
    }

    #[test]
    #[should_panic(expected="expected None to be Some<i32>")]
    fn test_negated_none_fails_with_message() {
        let n: Option<i32> = None;
        expect(n).is_not(none());
    }

    #[test]
    fn test_some_matches() {
        let n = Some(32);
        expect(n).is(some());
    }

    #[test]
    #[should_panic(expected="expected None to be Some<i32>")]
    fn test_some_fails_with_message() {
        let n: Option<i32> = None;
        expect(n).is(some());
    }

    #[test]
    #[should_panic(expected="expected Some(32) to be None")]
    fn test_negated_some_fails_with_message() {
        let n = Some(32);
        expect(n).is_not(some());
    }
}
<commit_msg>Drop option dependency on intrinsics<commit_after>use std::fmt::Debug;
use super::Matcher;

pub struct Nothing;
pub struct Something;

impl<T: Debug> Matcher<Option<T>> for Nothing {
    fn matches(&self, lhs: &Option<T>) -> bool {
        lhs.is_none()
    }

    fn fail_msg(&self, lhs: &Option<T>) -> String {
        format!("expected {:?} to be None", lhs)
    }

    fn negated_fail_msg(&self, lhs: &Option<T>) -> String {
        format!("expected {:?} to be Some", lhs)
    }
}

impl<T: Debug> Matcher<Option<T>> for Something {
    fn matches(&self, lhs: &Option<T>) -> bool {
        lhs.is_some()
    }

    fn fail_msg(&self, lhs: &Option<T>) -> String {
        format!("expected {:?} to be Some", lhs)
    }

    fn negated_fail_msg(&self, lhs: &Option<T>) -> String {
        format!("expected {:?} to be None", lhs)
    }
}

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

    #[test]
    fn test_none_matches() {
        let n: Option<String> = None;
        expect(n).is(none());
    }

    #[test]
    #[should_panic(expected="expected Some(\"Hello\") to be None")]
    fn test_none_fails_with_message() {
        let n = Some("Hello");
        expect(n).is(none());
    }

    #[test]
    #[should_panic(expected="expected None to be Some<i32>")]
    fn test_negated_none_fails_with_message() {
        let n: Option<i32> = None;
        expect(n).is_not(none());
    }

    #[test]
    fn test_some_matches() {
        let n = Some(32);
        expect(n).is(some());
    }

    #[test]
    #[should_panic(expected="expected None to be Some<i32>")]
    fn test_some_fails_with_message() {
        let n: Option<i32> = None;
        expect(n).is(some());
    }

    #[test]
    #[should_panic(expected="expected Some(32) to be None")]
    fn test_negated_some_fails_with_message() {
        let n = Some(32);
        expect(n).is_not(some());
    }
}
<|endoftext|>


<commit_before>use crypto::{ symmetriccipher, buffer, aes, blockmodes };
use crypto::buffer::{ ReadBuffer, WriteBuffer, BufferResult };

pub fn decrypt_aes_128_ecb_no_padding(encrypted_data: Vec<u8>, key: Vec<u8>) -> Result<Vec<u8>, symmetriccipher::SymmetricCipherError> {
    // https://github.com/DaGenix/rust-crypto/blob/master/examples/symmetriccipher.rs
    let mut decryptor = aes::ecb_decryptor(
        aes::KeySize::KeySize128,
        &key[..],
        blockmodes::NoPadding
    );

    let mut final_result = Vec::<u8>::new();
    let mut read_buffer = buffer::RefReadBuffer::new(&encrypted_data[..]);
    let mut buffer = [0; 4096];
    let mut write_buffer = buffer::RefWriteBuffer::new(&mut buffer);

    loop {
        let result = try!(decryptor.decrypt(&mut read_buffer, &mut write_buffer, true));
        final_result.extend(write_buffer.take_read_buffer().take_remaining().iter().map(|&i| i));
        match result {
            BufferResult::BufferUnderflow => break,
            BufferResult::BufferOverflow => { }
        }
    }

    Ok(final_result)
}<commit_msg>Added aes test<commit_after>use crypto::{ symmetriccipher, buffer, aes, blockmodes };
use crypto::buffer::{ ReadBuffer, WriteBuffer, BufferResult };

pub fn decrypt_aes_128_ecb_no_padding(encrypted_data: Vec<u8>, key: Vec<u8>) -> Result<Vec<u8>, symmetriccipher::SymmetricCipherError> {
    // https://github.com/DaGenix/rust-crypto/blob/master/examples/symmetriccipher.rs
    let mut decryptor = aes::ecb_decryptor(
        aes::KeySize::KeySize128,
        &key[..],
        blockmodes::NoPadding
    );

    let mut final_result = Vec::<u8>::new();
    let mut read_buffer = buffer::RefReadBuffer::new(&encrypted_data[..]);
    let mut buffer = [0; 4096];
    let mut write_buffer = buffer::RefWriteBuffer::new(&mut buffer);

    loop {
        let result = try!(decryptor.decrypt(&mut read_buffer, &mut write_buffer, true));
        final_result.extend(write_buffer.take_read_buffer().take_remaining().iter().map(|&i| i));
        match result {
            BufferResult::BufferUnderflow => break,
            BufferResult::BufferOverflow => { }
        }
    }

    Ok(final_result)
}

//-----------------------------------------------------------------------------

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

    #[test]
    fn decrypt_aes_128_ecb_no_padding_i_problem_7() {
        // http://cryptopals.com/sets/1/challenges/7/
        use byte_conversion::*;
        let encrypted_bytes = base64_to_bytes("CRIwqt4+szDbqkNY+I0qbDe3LQz0wiw0SuxBQtAM5TDdMbjCMD/venUDW9BL
PEXODbk6a48oMbAY6DDZsuLbc0uR9cp9hQ0QQGATyyCESq2NSsvhx5zKlLtz
dsnfK5ED5srKjK7Fz4Q38/ttd+stL/9WnDzlJvAo7WBsjI5YJc2gmAYayNfm
CW2lhZE/ZLG0CBD2aPw0W417QYb4cAIOW92jYRiJ4PTsBBHDe8o4JwqaUac6
rqdi833kbyAOV/Y2RMbN0oDb9Rq8uRHvbrqQJaJieaswEtMkgUt3P5Ttgeh7
J+hE6TR0uHot8WzHyAKNbUWHoi/5zcRCUipvVOYLoBZXlNu4qnwoCZRSBgvC
wTdz3Cbsp/P2wXB8tiz6l9rL2bLhBt13Qxyhhu0H0+JKj6soSeX5ZD1Rpilp
9ncR1tHW8+uurQKyXN4xKeGjaKLOejr2xDIw+aWF7GszU4qJhXBnXTIUUNUf
RlwEpS6FZcsMzemQF30ezSJHfpW7DVHzwiLyeiTJRKoVUwo43PXupnJXDmUy
sCa2nQz/iEwyor6kPekLv1csm1Pa2LZmbA9Ujzz8zb/gFXtQqBAN4zA8/wt0
VfoOsEZwcsaLOWUPtF/Ry3VhlKwXE7gGH/bbShAIKQqMqqUkEucZ3HPHAVp7
ZCn3Ox6+c5QJ3Uv8V7L7SprofPFN6F+kfDM4zAc59do5twgDoClCbxxG0L19
TBGHiYP3CygeY1HLMrX6KqypJfFJW5O9wNIF0qfOC2lWFgwayOwq41xdFSCW
0/EBSc7cJw3N06WThrW5LimAOt5L9c7Ik4YIxu0K9JZwAxfcU4ShYu6euYmW
LP98+qvRnIrXkePugS9TSOJOHzKUoOcb1/KYd9NZFHEcp58Df6rXFiz9DSq8
0rR5Kfs+M+Vuq5Z6zY98/SP0A6URIr9NFu+Cs9/gf+q4TRwsOzRMjMQzJL8f
7TXPEHH2+qEcpDKz/5pE0cvrgHr63XKu4XbzLCOBz0DoFAw3vkuxGwJq4Cpx
kt+eCtxSKUzNtXMn/mbPqPl4NZNJ8yzMqTFSODS4bYTBaN/uQYcOAF3NBYFd
5x9TzIAoW6ai13a8h/s9i5FlVRJDe2cetQhArrIVBquF0L0mUXMWNPFKkaQE
BsxpMCYh7pp7YlyCNode12k5jY1/lc8jQLQJ+EJHdCdM5t3emRzkPgND4a7O
NhoIkUUS2R1oEV1toDj9iDzGVFwOvWyt4GzA9XdxT333JU/n8m+N6hs23MBc
Z086kp9rJGVxZ5f80jRz3ZcjU6zWjR9ucRyjbsuVn1t4EJEm6A7KaHm13m0v
wN/O4KYTiiY3aO3siayjNrrNBpn1OeLv9UUneLSCdxcUqjRvOrdA5NYv25Hb
4wkFCIhC/Y2ze/kNyis6FrXtStcjKC1w9Kg8O25VXB1Fmpu+4nzpbNdJ9LXa
hF7wjOPXN6dixVKpzwTYjEFDSMaMhaTOTCaqJig97624wv79URbCgsyzwaC7
YXRtbTstbFuEFBee3uW7B3xXw72mymM2BS2uPQ5NIwmacbhta8aCRQEGqIZ0
78YrrOlZIjar3lbTCo5o6nbbDq9bvilirWG/SgWINuc3pWl5CscRcgQQNp7o
LBgrSkQkv9AjZYcvisnr89TxjoxBO0Y93jgp4T14LnVwWQVx3l3d6S1wlsci
dVeaM24E/JtS8k9XAvgSoKCjyiqsawBMzScXCIRCk6nqX8ZaJU3rZ0LeOMTU
w6MC4dC+aY9SrCvNQub19mBdtJUwOBOqGdfd5IoqQkaL6DfOkmpnsCs5PuLb
GZBVhah5L87IY7r6TB1V7KboXH8PZIYc1zlemMZGU0o7+etxZWHgpdeX6JbJ
Is3ilAzYqw/Hz65no7eUxcDg1aOaxemuPqnYRGhW6PvjZbwAtfQPlofhB0jT
Ht5bRlzF17rn9q/6wzlc1ssp2xmeFzXoxffpELABV6+yj3gfQ/bxIB9NWjdZ
K08RX9rjm9CcBlRQeTZrD67SYQWqRpT5t7zcVDnx1s7ZffLBWm/vXLfPzMaQ
YEJ4EfoduSutjshXvR+VQRPs2TWcF7OsaE4csedKUGFuo9DYfFIHFDNg+1Py
rlWJ0J/X0PduAuCZ+uQSsM/ex/vfXp6Z39ngq4exUXoPtAIqafrDMd8SuAty
EZhyY9V9Lp2qNQDbl6JI39bDz+6pDmjJ2jlnpMCezRK89cG11IqiUWvIPxHj
oiT1guH1uk4sQ2Pc1J4zjJNsZgoJDcPBbfss4kAqUJvQyFbzWshhtVeAv3dm
gwUENIhNK/erjpgw2BIRayzYw001jAIF5c7rYg38o6x3YdAtU3d3QpuwG5xD
fODxzfL3yEKQr48C/KqxI87uGwyg6H5gc2AcLU9JYt5QoDFoC7PFxcE3RVqc
7/Um9Js9X9UyriEjftWt86/tEyG7F9tWGxGNEZo3MOydwX/7jtwoxQE5ybFj
WndqLp8DV3naLQsh/Fz8JnTYHvOR72vuiw/x5D5PFuXV0aSVvmw5Wnb09q/B
owS14WzoHH6ekaWbh78xlypn/L/M+nIIEX1Ol3TaVOqIxvXZ2sjm86xRz0Ed
oHFfupSekdBULCqptxpFpBshZFvauUH8Ez7wA7wjL65GVlZ0f74U7MJVu9Sw
sZdgsLmnsQvr5n2ojNNBEv+qKG2wpUYTmWRaRc5EClUNfhzh8iDdHIsl6edO
ewORRrNiBay1NCzlfz1cj6VlYYQUM9bDEyqrwO400XQNpoFOxo4fxUdd+AHm
CBhHbyCR81/C6LQTG2JQBvjykG4pmoqnYPxDyeiCEG+JFHmP1IL+jggdjWhL
WQatslrWxuESEl3PEsrAkMF7gt0dBLgnWsc1cmzntG1rlXVi/Hs2TAU3RxEm
MSWDFubSivLWSqZj/XfGWwVpP6fsnsfxpY3d3h/fTxDu7U8GddaFRQhJ+0ZO
dx6nRJUW3u6xnhH3mYVRk88EMtpEpKrSIWfXphgDUPZ0f4agRzehkn9vtzCm
NjFnQb0/shnqTh4Mo/8oommbsBTUKPYS7/1oQCi12QABjJDt+LyUan+4iwvC
i0k0IUIHvk21381vC0ixYDZxzY64+xx/RNID+iplgzq9PDZgjc8L7jMg+2+m
rxPS56e71m5E2zufZ4d+nFjIg+dHD/ShNPzVpXizRVUERztLuak8Asah3/yv
wOrH1mKEMMGC1/6qfvZUgFLJH5V0Ep0n2K/Fbs0VljENIN8cjkCKdG8aBnef
EhITdV7CVjXcivQ6efkbOQCfkfcwWpaBFC8tD/zebXFE+JshW16D4EWXMnSm
/9HcGwHvtlAj04rwrZ5tRvAgf1IR83kqqiTvqfENcj7ddCFwtNZrQK7EJhgB
5Tr1tBFcb9InPRtS3KYteYHl3HWR9t8E2YGE8IGrS1sQibxaK/C0kKbqIrKp
npwtoOLsZPNbPw6K2jpko9NeZAx7PYFmamR4D50KtzgELQcaEsi5aCztMg7f
p1mK6ijyMKIRKwNKIYHagRRVLNgQLg/WTKzGVbWwq6kQaQyArwQCUXo4uRty
zGMaKbTG4dns1OFB1g7NCiPb6s1lv0/lHFAF6HwoYV/FPSL/pirxyDSBb/FR
RA3PIfmvGfMUGFVWlyS7+O73l5oIJHxuaJrR4EenzAu4Avpa5d+VuiYbM10a
LaVegVPvFn4pCP4U/Nbbw4OTCFX2HKmWEiVBB0O3J9xwXWpxN1Vr5CDi75Fq
NhxYCjgSJzWOUD34Y1dAfcj57VINmQVEWyc8Tch8vg9MnHGCOfOjRqp0VGyA
S15AVD2QS1V6fhRimJSVyT6QuGb8tKRsl2N+a2Xze36vgMhw7XK7zh//jC2H");

        let expected = readable_text_to_bytes("I\'m back and I\'m ringin\' the bell \nA rockin\' on the mike while the fly girls yell \nIn ecstasy in the back of me \nWell that\'s my DJ Deshay cuttin\' all them Z\'s \nHittin\' hard and the girlies goin\' crazy \nVanilla\'s on the mike, man I\'m not lazy. \n\nI\'m lettin\' my drug kick in \nIt controls my mouth and I begin \nTo just let it flow, let my concepts go \nMy posse\'s to the side yellin\', Go Vanilla Go! \n\nSmooth \'cause that\'s the way I will be \nAnd if you don\'t give a damn, then \nWhy you starin\' at me \nSo get off \'cause I control the stage \nThere\'s no dissin\' allowed \nI\'m in my own phase \nThe girlies sa y they love me and that is ok \nAnd I can dance better than any kid n\' play \n\nStage 2 -- Yea the one ya\' wanna listen to \nIt\'s off my head so let the beat play through \nSo I can funk it up and make it sound good \n1-2-3 Yo -- Knock on some wood \nFor good luck, I like my rhymes atrocious \nSupercalafragilisticexpialidocious \nI\'m an effect and that you can bet \nI can take a fly girl and make her wet. \n\nI\'m like Samson -- Samson to Delilah \nThere\'s no denyin\', You can try to hang \nBut you\'ll keep tryin\' to get my style \nOver and over, practice makes perfect \nBut not if you\'re a loafer. \n\nYou\'ll get nowhere, no place, no time, no girls \nSoon -- Oh my God, homebody, you probably eat \nSpaghetti with a spoon! Come on and say it! \n\nVIP. Vanilla Ice yep, yep, I\'m comin\' hard like a rhino \nIntoxicating so you stagger like a wino \nSo punks stop trying and girl stop cryin\' \nVanilla Ice is sellin\' and you people are buyin\' \n\'Cause why the freaks are jockin\' like Crazy Glue \nMovin\' and groovin\' trying to sing along \nAll through the ghetto groovin\' this here song \nNow you\'re amazed by the VIP posse. \n\nSteppin\' so hard like a German Nazi \nStartled by the bases hittin\' ground \nThere\'s no trippin\' on mine, I\'m just gettin\' down \nSparkamatic, I\'m hangin\' tight like a fanatic \nYou trapped me once and I thought that \nYou might have it \nSo step down and lend me your ear \n\'89 in my time! You, \'90 is my year. \n\nYou\'re weakenin\' fast, YO! and I can tell it \nYour body\'s gettin\' hot, so, so I can smell it \nSo don\'t be mad and don\'t be sad \n\'Cause the lyrics belong to ICE, You can call me Dad \nYou\'re pitchin\' a fit, so step back and endure \nLet the witch doctor, Ice, do the dance to cure \nSo come up close and don\'t be square \nYou wanna battle me -- Anytime, anywhere \n\nYou thought that I was weak, Boy, you\'re dead wrong \nSo come on, everybody and sing this song \n\nSay -- Play that funky music Say, go white boy, go white boy go \nplay that funky music Go white boy, go white boy, go \nLay down and boogie and play that funky music till you die. \n\nPlay that funky music Come on, Come on, let me hear \nPlay that funky music white boy you say it, say it \nPlay that funky music A little louder now \nPlay that funky music, white boy Come on, Come on, Come on \nPlay that funky music \n\u{4}\u{4}\u{4}\u{4}");

        let key = readable_text_to_bytes("YELLOW SUBMARINE");
        assert!(key.len() == 16);

        let decrypted_data_ = decrypt_aes_128_ecb_no_padding(encrypted_bytes, key);
        let decrypted_data = decrypted_data_.ok().unwrap();

        assert_eq!(expected, decrypted_data);
    }
}<|endoftext|>


<commit_before>//! Definition of the `Result` (immediately finished) combinator

use core::result;

use {Future, Poll, Async};

/// A future representing a value that is immediately ready.
///
/// Created by the `result` function.
#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
// TODO: rename this to `Result` on the next major version
pub struct FutureResult<T, E> {
    inner: Option<result::Result<T, E>>,
}

/// Creates a new "leaf future" which will resolve with the given result.
///
/// The returned future represents a computation which is finshed immediately.
/// This can be useful with the `finished` and `failed` base future types to
/// convert an immediate value to a future to interoperate elsewhere.
///
/// # Examples
///
/// ```
/// use futures::future::*;
///
/// let future_of_1 = result::<u32, u32>(Ok(1));
/// let future_of_err_2 = result::<u32, u32>(Err(2));
/// ```
pub fn result<T, E>(r: result::Result<T, E>) -> FutureResult<T, E> {
    FutureResult { inner: Some(r) }
}

/// Creates a "leaf future" from an immediate value of a finished and
/// successful computation.
///
/// The returned future is similar to `result` where it will immediately run a
/// scheduled callback with the provided value.
///
/// # Examples
///
/// ```
/// use futures::future::*;
///
/// let future_of_1 = ok::<u32, u32>(1);
/// ```
pub fn ok<T, E>(t: T) -> FutureResult<T, E> {
    result(Ok(t))
}

/// Creates a "leaf future" from an immediate value of a failed computation.
///
/// The returned future is similar to `result` where it will immediately run a
/// scheduled callback with the provided value.
///
/// # Examples
///
/// ```
/// use futures::future::*;
///
/// let future_of_err_1 = err::<u32, u32>(1);
/// ```
pub fn err<T, E>(e: E) -> FutureResult<T, E> {
    result(Err(e))
}

impl<T, E> Future for FutureResult<T, E> {
    type Item = T;
    type Error = E;

    fn poll(&mut self) -> Poll<T, E> {
        self.inner.take().expect("cannot poll Result twice").map(Async::Ready)
    }
}
<commit_msg>Make FutureResult Clone.<commit_after>//! Definition of the `Result` (immediately finished) combinator

use core::result;

use {Future, Poll, Async};

/// A future representing a value that is immediately ready.
///
/// Created by the `result` function.
#[derive(Debug, Clone)]
#[must_use = "futures do nothing unless polled"]
// TODO: rename this to `Result` on the next major version
pub struct FutureResult<T, E> {
    inner: Option<result::Result<T, E>>,
}

/// Creates a new "leaf future" which will resolve with the given result.
///
/// The returned future represents a computation which is finshed immediately.
/// This can be useful with the `finished` and `failed` base future types to
/// convert an immediate value to a future to interoperate elsewhere.
///
/// # Examples
///
/// ```
/// use futures::future::*;
///
/// let future_of_1 = result::<u32, u32>(Ok(1));
/// let future_of_err_2 = result::<u32, u32>(Err(2));
/// ```
pub fn result<T, E>(r: result::Result<T, E>) -> FutureResult<T, E> {
    FutureResult { inner: Some(r) }
}

/// Creates a "leaf future" from an immediate value of a finished and
/// successful computation.
///
/// The returned future is similar to `result` where it will immediately run a
/// scheduled callback with the provided value.
///
/// # Examples
///
/// ```
/// use futures::future::*;
///
/// let future_of_1 = ok::<u32, u32>(1);
/// ```
pub fn ok<T, E>(t: T) -> FutureResult<T, E> {
    result(Ok(t))
}

/// Creates a "leaf future" from an immediate value of a failed computation.
///
/// The returned future is similar to `result` where it will immediately run a
/// scheduled callback with the provided value.
///
/// # Examples
///
/// ```
/// use futures::future::*;
///
/// let future_of_err_1 = err::<u32, u32>(1);
/// ```
pub fn err<T, E>(e: E) -> FutureResult<T, E> {
    result(Err(e))
}

impl<T, E> Future for FutureResult<T, E> {
    type Item = T;
    type Error = E;

    fn poll(&mut self) -> Poll<T, E> {
        self.inner.take().expect("cannot poll Result twice").map(Async::Ready)
    }
}
<|endoftext|>


<commit_before>use chrono::NaiveDateTime;
use diesel::*;

use models::Crate;
use schema::*;
use views::EncodableCategory;

#[derive(Clone, Identifiable, Queryable, QueryableByName, Debug)]
#[table_name = "categories"]
pub struct Category {
    pub id: i32,
    pub category: String,
    pub slug: String,
    pub description: String,
    pub crates_cnt: i32,
    pub created_at: NaiveDateTime,
}

#[derive(Associations, Insertable, Identifiable, Debug, Clone, Copy)]
#[belongs_to(Category)]
#[belongs_to(Crate)]
#[table_name = "crates_categories"]
#[primary_key(crate_id, category_id)]
pub struct CrateCategory {
    crate_id: i32,
    category_id: i32,
}

impl Category {
    pub fn encodable(self) -> EncodableCategory {
        let Category {
            crates_cnt,
            category,
            slug,
            description,
            created_at,
            ..
        } = self;
        EncodableCategory {
            id: slug.clone(),
            slug: slug.clone(),
            description: description.clone(),
            created_at,
            crates_cnt,
            category: category.rsplit("::").collect::<Vec<_>>()[0].to_string(),
        }
    }

    pub fn update_crate<'a>(
        conn: &PgConnection,
        krate: &Crate,
        slugs: &[&'a str],
    ) -> QueryResult<Vec<&'a str>> {
        use diesel::dsl::any;

        conn.transaction(|| {
            let categories = categories::table
                .filter(categories::slug.eq(any(slugs)))
                .load::<Category>(conn)?;
            let invalid_categories = slugs
                .iter()
                .cloned()
                .filter(|s| !categories.iter().any(|c| c.slug == *s))
                .collect();
            let crate_categories = categories
                .iter()
                .map(|c| CrateCategory {
                    category_id: c.id,
                    crate_id: krate.id,
                })
                .collect::<Vec<_>>();

            delete(CrateCategory::belonging_to(krate)).execute(conn)?;
            insert_into(crates_categories::table)
                .values(&crate_categories)
                .execute(conn)?;
            Ok(invalid_categories)
        })
    }

    pub fn count_toplevel(conn: &PgConnection) -> QueryResult<i64> {
        use self::categories::dsl::*;

        categories
            .filter(category.not_like("%::%"))
            .count()
            .get_result(conn)
    }

    pub fn toplevel(
        conn: &PgConnection,
        sort: &str,
        limit: i64,
        offset: i64,
    ) -> QueryResult<Vec<Category>> {
        use diesel::dsl::*;
        use diesel::select;

        let sort_sql = match sort {
            "crates" => "ORDER BY crates_cnt DESC",
            _ => "ORDER BY category ASC",
        };

        // Collect all the top-level categories and sum up the crates_cnt of
        // the crates in all subcategories
        select(sql::<categories::SqlType>(&format!(
            "c.id, c.category, c.slug, c.description,
                sum(c2.crates_cnt)::int as crates_cnt, c.created_at
             FROM categories as c
             INNER JOIN categories c2 ON split_part(c2.slug, '::', 1) = c.slug
             WHERE split_part(c.slug, '::', 1) = c.slug
             GROUP BY c.id
             {} LIMIT {} OFFSET {}",
            sort_sql, limit, offset
        ))).load(conn)
    }

    pub fn subcategories(&self, conn: &PgConnection) -> QueryResult<Vec<Category>> {
        use diesel::sql_types::Text;

        sql_query(
            "SELECT c.id, c.category, c.slug, c.description, \
             COALESCE (( \
             SELECT sum(c2.crates_cnt)::int \
             FROM categories as c2 \
             WHERE c2.slug = c.slug \
             OR c2.slug LIKE c.slug || '::%' \
             ), 0) as crates_cnt, c.created_at \
             FROM categories as c \
             WHERE c.category ILIKE $1 || '::%' \
             AND c.category NOT ILIKE $1 || '::%::%'",
        ).bind::<Text, _>(&self.category)
            .load(conn)
    }

    /// Gathers the parent categories from the top-level Category to the direct parent of this Category.
    /// Returns categories as a Vector in order of traversal, not including this Category.
    /// The intention is to be able to have slugs or parent categories arrayed in order, to
    /// offer the frontend, for examples, slugs to create links to each parent category in turn.
    pub fn parent_categories(&self, conn: &PgConnection) -> QueryResult<Vec<Category>> {
        use diesel::expression::dsl::*;
        use diesel::types::Text;

        sql::<categories::SqlType>(
            "SELECT c.id, c.category, c.slug, c.description, \
            COALESCE(( \
            SELECT sum(c2.crates_cnt)::int from categories c2 \
            WHERE path <@ subltree(c.path, 0, 2)), 0) as crates_cnt, c.created_at \
            FROM categories c \
            WHERE c.path @> (select path from categories where slug = $1) \
            AND c.slug <> $1
            ORDER BY c.path",
        ).bind::<Text, _>(&self.slug)
            .load(conn)
    }
}

#[derive(Insertable, AsChangeset, Default, Debug)]
#[table_name = "categories"]
pub struct NewCategory<'a> {
    pub category: &'a str,
    pub slug: &'a str,
}

impl<'a> NewCategory<'a> {
    /// Inserts the category into the database, or updates an existing one.
    pub fn create_or_update(&self, conn: &PgConnection) -> QueryResult<Category> {
        use schema::categories::dsl::*;

        insert_into(categories)
            .values(self)
            .on_conflict(slug)
            .do_update()
            .set(self)
            .get_result(conn)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use diesel::connection::SimpleConnection;
    use dotenv::dotenv;
    use std::env;

    fn pg_connection() -> PgConnection {
        let _ = dotenv();
        let database_url =
            env::var("TEST_DATABASE_URL").expect("TEST_DATABASE_URL must be set to run tests");
        let conn = PgConnection::establish(&database_url).unwrap();
        // These tests deadlock if run concurrently
        conn.batch_execute("BEGIN; LOCK categories IN ACCESS EXCLUSIVE MODE")
            .unwrap();
        conn
    }

    #[test]
    fn category_toplevel_excludes_subcategories() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug) VALUES
            ('Cat 2', 'cat2'), ('Cat 1', 'cat1'), ('Cat 1::sub', 'cat1::sub')
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "", 10, 0)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec!["Cat 1".to_string(), "Cat 2".to_string()];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_toplevel_orders_by_crates_cnt_when_sort_given() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug, crates_cnt) VALUES
            ('Cat 1', 'cat1', 0), ('Cat 2', 'cat2', 2), ('Cat 3', 'cat3', 1)
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "crates", 10, 0)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec![
            "Cat 2".to_string(),
            "Cat 3".to_string(),
            "Cat 1".to_string(),
        ];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_toplevel_applies_limit_and_offset() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug) VALUES
            ('Cat 1', 'cat1'), ('Cat 2', 'cat2')
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "", 1, 0)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec!["Cat 1".to_string()];
        assert_eq!(expected, categories);

        let categories = Category::toplevel(&conn, "", 1, 1)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec!["Cat 2".to_string()];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_toplevel_includes_subcategories_in_crate_cnt() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug, crates_cnt) VALUES
            ('Cat 1', 'cat1', 1), ('Cat 1::sub', 'cat1::sub', 2),
            ('Cat 2', 'cat2', 3), ('Cat 2::Sub 1', 'cat2::sub1', 4),
            ('Cat 2::Sub 2', 'cat2::sub2', 5), ('Cat 3', 'cat3', 6)
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "crates", 10, 0)
            .unwrap()
            .into_iter()
            .map(|c| (c.category, c.crates_cnt))
            .collect::<Vec<_>>();
        let expected = vec![
            ("Cat 2".to_string(), 12),
            ("Cat 3".to_string(), 6),
            ("Cat 1".to_string(), 3),
        ];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_toplevel_applies_limit_and_offset_after_grouping() {
        let conn = pg_connection();
        conn.batch_execute(
            "INSERT INTO categories (category, slug, crates_cnt) VALUES
            ('Cat 1', 'cat1', 1), ('Cat 1::sub', 'cat1::sub', 2),
            ('Cat 2', 'cat2', 3), ('Cat 2::Sub 1', 'cat2::sub1', 4),
            ('Cat 2::Sub 2', 'cat2::sub2', 5), ('Cat 3', 'cat3', 6)
            ",
        ).unwrap();

        let categories = Category::toplevel(&conn, "crates", 2, 0)
            .unwrap()
            .into_iter()
            .map(|c| (c.category, c.crates_cnt))
            .collect::<Vec<_>>();
        let expected = vec![("Cat 2".to_string(), 12), ("Cat 3".to_string(), 6)];
        assert_eq!(expected, categories);

        let categories = Category::toplevel(&conn, "crates", 2, 1)
            .unwrap()
            .into_iter()
            .map(|c| (c.category, c.crates_cnt))
            .collect::<Vec<_>>();
        let expected = vec![("Cat 3".to_string(), 6), ("Cat 1".to_string(), 3)];
        assert_eq!(expected, categories);
    }

    #[test]
    fn category_parent_categories_includes_path_to_node_with_count() {
        use self::categories::dsl::*;
        let conn = pg_connection();
        insert_into(categories)
            .values(&vec![
                (category.eq("Cat 1"), slug.eq("cat1"), crates_cnt.eq(1)),
                (category.eq("Cat 1::sub1"), slug.eq("cat1::sub1"), crates_cnt.eq(2)),
                (category.eq("Cat 1::sub2"), slug.eq("cat1::sub2"), crates_cnt.eq(2)),
                (category.eq("Cat 1::sub1::subsub1"), slug.eq("cat1::sub1::subsub1"), crates_cnt.eq(2)),
                (category.eq("Cat 2"), slug.eq("cat2"), crates_cnt.eq(3)),
                (category.eq("Cat 2::Sub 1"), slug.eq("cat2::sub1"), crates_cnt.eq(4)),
                (category.eq("Cat 2::Sub 2"), slug.eq("cat2::sub2"), crates_cnt.eq(5)),
                (category.eq("Cat 3"), slug.eq("cat3"), crates_cnt.eq(200))
            ])
            .execute(&conn)
            .unwrap();

        let cat = categories
            .filter(slug.eq("cat1::sub1"))
            .first::<Category>(&conn)
            .unwrap();
        let subcats = cat.subcategories(&conn).unwrap();
        let parents = cat.parent_categories(&conn).unwrap();

        assert_eq!(parents.len(), 1);
        assert_eq!(parents[0].slug, "cat1");
        assert_eq!(parents[0].crates_cnt, 7);
        assert_eq!(subcats.len(), 1);
        assert_eq!(subcats[0].slug, "cat1::sub1::subsub1");
    }
}
<commit_msg>Fix up the other tests too<commit_after>use chrono::NaiveDateTime;
use diesel::*;

use models::Crate;
use schema::*;
use views::EncodableCategory;

#[derive(Clone, Identifiable, Queryable, QueryableByName, Debug)]
#[table_name = "categories"]
pub struct Category {
    pub id: i32,
    pub category: String,
    pub slug: String,
    pub description: String,
    pub crates_cnt: i32,
    pub created_at: NaiveDateTime,
}

#[derive(Associations, Insertable, Identifiable, Debug, Clone, Copy)]
#[belongs_to(Category)]
#[belongs_to(Crate)]
#[table_name = "crates_categories"]
#[primary_key(crate_id, category_id)]
pub struct CrateCategory {
    crate_id: i32,
    category_id: i32,
}

impl Category {
    pub fn encodable(self) -> EncodableCategory {
        let Category {
            crates_cnt,
            category,
            slug,
            description,
            created_at,
            ..
        } = self;
        EncodableCategory {
            id: slug.clone(),
            slug: slug.clone(),
            description: description.clone(),
            created_at,
            crates_cnt,
            category: category.rsplit("::").collect::<Vec<_>>()[0].to_string(),
        }
    }

    pub fn update_crate<'a>(
        conn: &PgConnection,
        krate: &Crate,
        slugs: &[&'a str],
    ) -> QueryResult<Vec<&'a str>> {
        use diesel::dsl::any;

        conn.transaction(|| {
            let categories = categories::table
                .filter(categories::slug.eq(any(slugs)))
                .load::<Category>(conn)?;
            let invalid_categories = slugs
                .iter()
                .cloned()
                .filter(|s| !categories.iter().any(|c| c.slug == *s))
                .collect();
            let crate_categories = categories
                .iter()
                .map(|c| CrateCategory {
                    category_id: c.id,
                    crate_id: krate.id,
                })
                .collect::<Vec<_>>();

            delete(CrateCategory::belonging_to(krate)).execute(conn)?;
            insert_into(crates_categories::table)
                .values(&crate_categories)
                .execute(conn)?;
            Ok(invalid_categories)
        })
    }

    pub fn count_toplevel(conn: &PgConnection) -> QueryResult<i64> {
        use self::categories::dsl::*;

        categories
            .filter(category.not_like("%::%"))
            .count()
            .get_result(conn)
    }

    pub fn toplevel(
        conn: &PgConnection,
        sort: &str,
        limit: i64,
        offset: i64,
    ) -> QueryResult<Vec<Category>> {
        use diesel::dsl::*;
        use diesel::select;

        let sort_sql = match sort {
            "crates" => "ORDER BY crates_cnt DESC",
            _ => "ORDER BY category ASC",
        };

        // Collect all the top-level categories and sum up the crates_cnt of
        // the crates in all subcategories
        select(sql::<categories::SqlType>(&format!(
            "c.id, c.category, c.slug, c.description,
                sum(c2.crates_cnt)::int as crates_cnt, c.created_at
             FROM categories as c
             INNER JOIN categories c2 ON split_part(c2.slug, '::', 1) = c.slug
             WHERE split_part(c.slug, '::', 1) = c.slug
             GROUP BY c.id
             {} LIMIT {} OFFSET {}",
            sort_sql, limit, offset
        ))).load(conn)
    }

    pub fn subcategories(&self, conn: &PgConnection) -> QueryResult<Vec<Category>> {
        use diesel::sql_types::Text;

        sql_query(
            "SELECT c.id, c.category, c.slug, c.description, \
             COALESCE (( \
             SELECT sum(c2.crates_cnt)::int \
             FROM categories as c2 \
             WHERE c2.slug = c.slug \
             OR c2.slug LIKE c.slug || '::%' \
             ), 0) as crates_cnt, c.created_at \
             FROM categories as c \
             WHERE c.category ILIKE $1 || '::%' \
             AND c.category NOT ILIKE $1 || '::%::%'",
        ).bind::<Text, _>(&self.category)
            .load(conn)
    }

    /// Gathers the parent categories from the top-level Category to the direct parent of this Category.
    /// Returns categories as a Vector in order of traversal, not including this Category.
    /// The intention is to be able to have slugs or parent categories arrayed in order, to
    /// offer the frontend, for examples, slugs to create links to each parent category in turn.
    pub fn parent_categories(&self, conn: &PgConnection) -> QueryResult<Vec<Category>> {
        use diesel::expression::dsl::*;
        use diesel::types::Text;

        sql::<categories::SqlType>(
            "SELECT c.id, c.category, c.slug, c.description, \
            COALESCE(( \
            SELECT sum(c2.crates_cnt)::int from categories c2 \
            WHERE path <@ subltree(c.path, 0, 2)), 0) as crates_cnt, c.created_at \
            FROM categories c \
            WHERE c.path @> (select path from categories where slug = $1) \
            AND c.slug <> $1
            ORDER BY c.path",
        ).bind::<Text, _>(&self.slug)
            .load(conn)
    }
}

#[derive(Insertable, AsChangeset, Default, Debug)]
#[table_name = "categories"]
pub struct NewCategory<'a> {
    pub category: &'a str,
    pub slug: &'a str,
}

impl<'a> NewCategory<'a> {
    /// Inserts the category into the database, or updates an existing one.
    pub fn create_or_update(&self, conn: &PgConnection) -> QueryResult<Category> {
        use schema::categories::dsl::*;

        insert_into(categories)
            .values(self)
            .on_conflict(slug)
            .do_update()
            .set(self)
            .get_result(conn)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use diesel::connection::SimpleConnection;
    use dotenv::dotenv;
    use std::env;

    fn pg_connection() -> PgConnection {
        let _ = dotenv();
        let database_url =
            env::var("TEST_DATABASE_URL").expect("TEST_DATABASE_URL must be set to run tests");
        let conn = PgConnection::establish(&database_url).unwrap();
        // These tests deadlock if run concurrently
        conn.batch_execute("BEGIN; LOCK categories IN ACCESS EXCLUSIVE MODE")
            .unwrap();
        conn
    }

    #[test]
    fn category_toplevel_excludes_subcategories() {
        use self::categories::dsl::*;
        let conn = pg_connection();
        insert_into(categories)
            .values(&vec![
                (category.eq("Cat 2"), slug.eq("cat2")),
                (category.eq("Cat 1"), slug.eq("cat1")),
                (category.eq("Cat 1::sub"), slug.eq("cat1::sub"))
            ])
            .execute(&conn)
            .unwrap();

        let cats = Category::toplevel(&conn, "", 10, 0)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec!["Cat 1".to_string(), "Cat 2".to_string()];
        assert_eq!(expected, cats);
    }

    #[test]
    fn category_toplevel_orders_by_crates_cnt_when_sort_given() {
        use self::categories::dsl::*;
        let conn = pg_connection();
        insert_into(categories)
            .values(&vec![
                (category.eq("Cat 1"), slug.eq("cat1"), crates_cnt.eq(0)),
                (category.eq("Cat 2"), slug.eq("cat2"), crates_cnt.eq(2)),
                (category.eq("Cat 3"), slug.eq("cat3"), crates_cnt.eq(1))
            ])
            .execute(&conn)
            .unwrap();

        let cats = Category::toplevel(&conn, "crates", 10, 0)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec![
            "Cat 2".to_string(),
            "Cat 3".to_string(),
            "Cat 1".to_string(),
        ];
        assert_eq!(expected, cats);
    }

    #[test]
    fn category_toplevel_applies_limit_and_offset() {
        use self::categories::dsl::*;
        let conn = pg_connection();
        insert_into(categories)
            .values(&vec![
                (category.eq("Cat 1"), slug.eq("cat1")),
                (category.eq("Cat 2"), slug.eq("cat2"))
            ])
            .execute(&conn)
            .unwrap();

        let cats = Category::toplevel(&conn, "", 1, 0)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec!["Cat 1".to_string()];
        assert_eq!(expected, cats);

        let cats = Category::toplevel(&conn, "", 1, 1)
            .unwrap()
            .into_iter()
            .map(|c| c.category)
            .collect::<Vec<_>>();
        let expected = vec!["Cat 2".to_string()];
        assert_eq!(expected, cats);
    }

    #[test]
    fn category_toplevel_includes_subcategories_in_crate_cnt() {
        use self::categories::dsl::*;
        let conn = pg_connection();
        insert_into(categories)
            .values(&vec![
                (category.eq("Cat 1"), slug.eq("cat1"), crates_cnt.eq(1)),
                (category.eq("Cat 1::sub"), slug.eq("cat1::sub"), crates_cnt.eq(2)),
                (category.eq("Cat 2"), slug.eq("cat2"), crates_cnt.eq(3)),
                (category.eq("Cat 2::Sub 1"), slug.eq("cat2::sub1"), crates_cnt.eq(4)),
                (category.eq("Cat 2::Sub 2"), slug.eq("cat2::sub2"), crates_cnt.eq(5)),
                (category.eq("Cat 3"), slug.eq("cat3"), crates_cnt.eq(6))
            ])
            .execute(&conn)
            .unwrap();

        let cats = Category::toplevel(&conn, "crates", 10, 0)
            .unwrap()
            .into_iter()
            .map(|c| (c.category, c.crates_cnt))
            .collect::<Vec<_>>();
        let expected = vec![
            ("Cat 2".to_string(), 12),
            ("Cat 3".to_string(), 6),
            ("Cat 1".to_string(), 3),
        ];
        assert_eq!(expected, cats);
    }

    #[test]
    fn category_toplevel_applies_limit_and_offset_after_grouping() {
        use self::categories::dsl::*;
        let conn = pg_connection();
        insert_into(categories)
            .values(&vec![
                (category.eq("Cat 1"), slug.eq("cat1"), crates_cnt.eq(1)),
                (category.eq("Cat 1::sub"), slug.eq("cat1::sub"), crates_cnt.eq(2)),
                (category.eq("Cat 2"), slug.eq("cat2"), crates_cnt.eq(3)),
                (category.eq("Cat 2::Sub 1"), slug.eq("cat2::sub1"), crates_cnt.eq(4)),
                (category.eq("Cat 2::Sub 2"), slug.eq("cat2::sub2"), crates_cnt.eq(5)),
                (category.eq("Cat 3"), slug.eq("cat3"), crates_cnt.eq(6))
            ])
            .execute(&conn)
            .unwrap();

        let cats = Category::toplevel(&conn, "crates", 2, 0)
            .unwrap()
            .into_iter()
            .map(|c| (c.category, c.crates_cnt))
            .collect::<Vec<_>>();
        let expected = vec![("Cat 2".to_string(), 12), ("Cat 3".to_string(), 6)];
        assert_eq!(expected, cats);

        let cats = Category::toplevel(&conn, "crates", 2, 1)
            .unwrap()
            .into_iter()
            .map(|c| (c.category, c.crates_cnt))
            .collect::<Vec<_>>();
        let expected = vec![("Cat 3".to_string(), 6), ("Cat 1".to_string(), 3)];
        assert_eq!(expected, cats);
    }

    #[test]
    fn category_parent_categories_includes_path_to_node_with_count() {
        use self::categories::dsl::*;
        let conn = pg_connection();
        insert_into(categories)
            .values(&vec![
                (category.eq("Cat 1"), slug.eq("cat1"), crates_cnt.eq(1)),
                (category.eq("Cat 1::sub1"), slug.eq("cat1::sub1"), crates_cnt.eq(2)),
                (category.eq("Cat 1::sub2"), slug.eq("cat1::sub2"), crates_cnt.eq(2)),
                (category.eq("Cat 1::sub1::subsub1"), slug.eq("cat1::sub1::subsub1"), crates_cnt.eq(2)),
                (category.eq("Cat 2"), slug.eq("cat2"), crates_cnt.eq(3)),
                (category.eq("Cat 2::Sub 1"), slug.eq("cat2::sub1"), crates_cnt.eq(4)),
                (category.eq("Cat 2::Sub 2"), slug.eq("cat2::sub2"), crates_cnt.eq(5)),
                (category.eq("Cat 3"), slug.eq("cat3"), crates_cnt.eq(200))
            ])
            .execute(&conn)
            .unwrap();

        let cat = categories
            .filter(slug.eq("cat1::sub1"))
            .first::<Category>(&conn)
            .unwrap();
        let subcats = cat.subcategories(&conn).unwrap();
        let parents = cat.parent_categories(&conn).unwrap();

        assert_eq!(parents.len(), 1);
        assert_eq!(parents[0].slug, "cat1");
        assert_eq!(parents[0].crates_cnt, 7);
        assert_eq!(subcats.len(), 1);
        assert_eq!(subcats[0].slug, "cat1::sub1::subsub1");
    }
}
<|endoftext|>