<commit_before>#![cfg(test)]

extern crate compiletest_rs as compiletest;

use api::*;
use rand::{Rng, SeedableRng, XorShiftRng};
use std::path::PathBuf;
use std::error::Error;

fn quick_sort<T: PartialOrd + Send>(v: &mut [T]) {
    if v.len() <= 1 {
        return;
    }

    let mid = partition(v);
    let (lo, hi) = v.split_at_mut(mid);
    join(|| quick_sort(lo), || quick_sort(hi));
}

fn partition<T: PartialOrd + Send>(v: &mut [T]) -> usize {
    let pivot = v.len() - 1;
    let mut i = 0;
    for j in 0..pivot {
        if v[j] <= v[pivot] {
            v.swap(i, j);
            i += 1;
        }
    }
    v.swap(i, pivot);
    i
}

#[test]
fn sort() {
    let mut rng = XorShiftRng::from_seed([0, 1, 2, 3]);
    let mut data: Vec<_> = (0..6 * 1024).map(|_| rng.next_u32()).collect();

    let result = initialize(Configuration::new());

    match result {
        Ok(_) => {
            quick_sort(&mut data);

            let mut sorted_data = data.clone();
            sorted_data.sort();

            assert_eq!(data, sorted_data);
        }
        Err(e) => panic!("expected InitOk, but got {:?}", e),
    }
}

#[test]
fn sort_in_pool() {
    let mut rng = XorShiftRng::from_seed([0, 1, 2, 3]);
    let mut data: Vec<_> = (0..12 * 1024).map(|_| rng.next_u32()).collect();

    let result = ThreadPool::new(Configuration::new());

    match result {
        Ok(pool) => {
            pool.install(|| {
                quick_sort(&mut data);
            });

            let mut sorted_data = data.clone();
            sorted_data.sort();

            assert_eq!(data, sorted_data);
        }
        Err(_) => panic!("expected Ok() but got Err()"),
    }
}

#[test]
fn error_in_pool() {
    let result = ThreadPool::new(Configuration::new().set_num_threads(0));

    match result {
        Ok(_) => panic!("expected Err(), but got Ok()"),
        Err(error) => assert_eq!(error, InitError::NumberOfThreadsZero),
    }
}

#[test]
fn negative_tests_compile_fail() {
    let mode = "compile-fail";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/compile-fail");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
#[cfg(feature = "unstable")]
fn negative_tests_compile_fail_unstable() {
    let mode = "compile-fail";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/compile-fail-unstable");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
fn negative_tests_run_fail() {
    let mode = "run-fail";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/run-fail");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
#[cfg(feature = "unstable")]
fn negative_tests_run_fail_unstable() {
    let mode = "run-fail";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/run-fail-unstable");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
fn positive_test_run_pass() {
    let mode = "run-pass";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/run-pass");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
#[cfg(feature = "unstable")]
fn positive_test_run_pass_unstable() {
    let mode = "run-pass";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/run-pass-unstable");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
fn coerce_to_box_error() {
    // check that coercion succeeds
    let _: Box<Error> = From::from(InitError::NumberOfThreadsZero);
}
<commit_msg>test thread index functions<commit_after>#![cfg(test)]

extern crate compiletest_rs as compiletest;

use api::*;
use rand::{Rng, SeedableRng, XorShiftRng};
use std::path::PathBuf;
use std::error::Error;

fn quick_sort<T: PartialOrd + Send>(v: &mut [T]) {
    if v.len() <= 1 {
        return;
    }

    let mid = partition(v);
    let (lo, hi) = v.split_at_mut(mid);
    join(|| quick_sort(lo), || quick_sort(hi));
}

fn partition<T: PartialOrd + Send>(v: &mut [T]) -> usize {
    let pivot = v.len() - 1;
    let mut i = 0;
    for j in 0..pivot {
        if v[j] <= v[pivot] {
            v.swap(i, j);
            i += 1;
        }
    }
    v.swap(i, pivot);
    i
}

#[test]
fn sort() {
    let mut rng = XorShiftRng::from_seed([0, 1, 2, 3]);
    let mut data: Vec<_> = (0..6 * 1024).map(|_| rng.next_u32()).collect();

    let result = initialize(Configuration::new());

    match result {
        Ok(_) => {
            quick_sort(&mut data);

            let mut sorted_data = data.clone();
            sorted_data.sort();

            assert_eq!(data, sorted_data);
        }
        Err(e) => panic!("expected InitOk, but got {:?}", e),
    }
}

#[test]
fn sort_in_pool() {
    let mut rng = XorShiftRng::from_seed([0, 1, 2, 3]);
    let mut data: Vec<_> = (0..12 * 1024).map(|_| rng.next_u32()).collect();

    let result = ThreadPool::new(Configuration::new());

    match result {
        Ok(pool) => {
            pool.install(|| {
                quick_sort(&mut data);
            });

            let mut sorted_data = data.clone();
            sorted_data.sort();

            assert_eq!(data, sorted_data);
        }
        Err(_) => panic!("expected Ok() but got Err()"),
    }
}

#[test]
fn error_in_pool() {
    let result = ThreadPool::new(Configuration::new().set_num_threads(0));

    match result {
        Ok(_) => panic!("expected Err(), but got Ok()"),
        Err(error) => assert_eq!(error, InitError::NumberOfThreadsZero),
    }
}

#[test]
fn negative_tests_compile_fail() {
    let mode = "compile-fail";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/compile-fail");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
#[cfg(feature = "unstable")]
fn negative_tests_compile_fail_unstable() {
    let mode = "compile-fail";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/compile-fail-unstable");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
fn negative_tests_run_fail() {
    let mode = "run-fail";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/run-fail");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
#[cfg(feature = "unstable")]
fn negative_tests_run_fail_unstable() {
    let mode = "run-fail";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/run-fail-unstable");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
fn positive_test_run_pass() {
    let mode = "run-pass";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/run-pass");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
#[cfg(feature = "unstable")]
fn positive_test_run_pass_unstable() {
    let mode = "run-pass";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/run-pass-unstable");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
fn coerce_to_box_error() {
    // check that coercion succeeds
    let _: Box<Error> = From::from(InitError::NumberOfThreadsZero);
}

#[test]
fn worker_thread_index() {
    let pool = ThreadPool::new(Configuration::new().set_num_threads(22)).unwrap();
    assert_eq!(pool.num_threads(), 22);
    assert_eq!(pool.current_thread_index(), None);
    let index = pool.install(|| pool.current_thread_index().unwrap());
    assert!(index < 22);
}
<|endoftext|>


<commit_before>use std::cmp::*;
use std::os;

pub fn is_word(word:&str) -> bool {
    // Word must either contain no double quotes,
    // or must have a quote on the end
    regex!("^.*\".+\"$|^[^\"]+$").is_match(word)
}

// work around lack of DST
pub fn build_string(ch:char, count:usize) -> String {
    let mut s = String::new();
    let mut i = 0us;
    loop {
        if i == count {
            return s;
        }
        s.push(ch);
        i += 1;
    }
}

pub fn strip_words(line:Vec<String>) -> Vec<String> {
    let mut out = Vec::<String>::new();
    for word in line.iter() {
        out.push(strip_word(word));
    }
    return out;
}

pub fn strip_word(word:&String) -> String {
    let first_removed:String = regex!("\"").splitn(word.as_slice(), 2).collect::<Vec<&str>>().as_slice().concat();
    let second_removed:String = regex!("\"$").splitn(first_removed.as_slice(), 2).collect::<Vec<&str>>().as_slice().concat();
    return second_removed.to_string();
}

pub fn expand_path(path:Path) -> Path {
    if Path::new("~").is_ancestor_of(&path) {
        return match os::homedir() {
            None => Path::new("/"),
            Some(val) => Path::new(val)
        }.join(Path::new(path.as_vec().slice_from(min(path.as_vec().len(), 2))));
    } else {
        return path;
    }
}

pub fn condense_path(path:Path) -> Path {
    let homep = Path::new(match os::homedir() {
            None => return path,
            Some(val) => val
    });
    if homep.is_ancestor_of(&path) {
        match path.path_relative_from(&homep) {
            None => path,
            Some(path) => Path::new("~").join(path)
        }
    } else {
        return path;
    }
}

#[test]
fn build_string_test() {
    assert!(build_string('a', 5) == String::from_str("aaaaa"));
}

#[test]
fn is_word_test() {
    assert!(is_word("hello"));
    assert!(is_word("\"hello world\""));
    assert!(is_word("TEST=\"hello world\""));
    //assert!(is_word("$(test command)"));

    assert!(!is_word("TEST=\"hello"));
    //assert!(!is_word("$(test com"));
}

#[test]
fn strip_word_test() {
    assert!(strip_word(&String::from_str("\"hello world\"")) == String::from_str("hello world"));
    assert!(strip_word(&String::from_str("hello")) == String::from_str("hello"));
}

#[test]
fn strip_words_test() {
    assert!(strip_words(vec![String::from_str("\"hel\"lo world\""), String::from_str("hello")])
            == vec![String::from_str("hel\"lo world"), String::from_str("hello")]);
}

#[test]
fn expand_path_test() {
    // tests require the HOME env set
    let homep = Path::new(os::homedir().unwrap());
    assert!(expand_path(Path::new("~/Documents/scripts/")) == homep.join("Documents/scripts/"));
    assert!(expand_path(Path::new("/etc/wash/")) == Path::new("/etc/wash/"));
}

#[test]
fn condense_path_test() {
    // tests require the HOME env set
    let homep = Path::new(os::homedir().unwrap());
    assert!(condense_path(homep.join("Documents/scripts/")) ==
            Path::new("~/Documents/scripts/"));
    assert!(condense_path(Path::new("/home/")) == Path::new("/home/"));
    assert!(condense_path(Path::new("/etc/wash/")) == Path::new("/etc/wash/"));
}
<commit_msg>Improve word matching<commit_after>use std::cmp::*;
use std::os;

pub fn is_word(word:&str) -> bool {
    // defines the following "container" sequences: (.*) and ".*"
    // words cannot end when either are unclosed
    regex!("^([^ \t\n\r\"()]*|\".*\"|\\(.*\\))*$").is_match(word)
}

// work around lack of DST
pub fn build_string(ch:char, count:usize) -> String {
    let mut s = String::new();
    let mut i = 0us;
    loop {
        if i == count {
            return s;
        }
        s.push(ch);
        i += 1;
    }
}

pub fn strip_words(line:Vec<String>) -> Vec<String> {
    let mut out = Vec::<String>::new();
    for word in line.iter() {
        out.push(strip_word(word));
    }
    return out;
}

pub fn strip_word(word:&String) -> String {
    let first_removed:String = regex!("\"").splitn(word.as_slice(), 2).collect::<Vec<&str>>().as_slice().concat();
    let second_removed:String = regex!("\"$").splitn(first_removed.as_slice(), 2).collect::<Vec<&str>>().as_slice().concat();
    return second_removed.to_string();
}

pub fn expand_path(path:Path) -> Path {
    if Path::new("~").is_ancestor_of(&path) {
        return match os::homedir() {
            None => Path::new("/"),
            Some(val) => Path::new(val)
        }.join(Path::new(path.as_vec().slice_from(min(path.as_vec().len(), 2))));
    } else {
        return path;
    }
}

pub fn condense_path(path:Path) -> Path {
    let homep = Path::new(match os::homedir() {
            None => return path,
            Some(val) => val
    });
    if homep.is_ancestor_of(&path) {
        match path.path_relative_from(&homep) {
            None => path,
            Some(path) => Path::new("~").join(path)
        }
    } else {
        return path;
    }
}

#[test]
fn build_string_test() {
    assert!(build_string('a', 5) == String::from_str("aaaaa"));
}

#[test]
fn is_word_test() {
    assert!(is_word(""));
    assert!(is_word("hello"));
    assert!(is_word("\"hello world\""));
    assert!(is_word("TEST=\"hello world\""));
    assert!(is_word("func(test function)"));
    assert!(is_word("this(is)a\"complex\"series"));

    assert!(!is_word("TEST=\"hello"));
    assert!(!is_word("func(test fun"));
    assert!(!is_word("this(is)a\"complex\"series(with no end"));
}

#[test]
fn strip_word_test() {
    assert!(strip_word(&String::from_str("\"hello world\"")) == String::from_str("hello world"));
    assert!(strip_word(&String::from_str("hello")) == String::from_str("hello"));
}

#[test]
fn strip_words_test() {
    assert!(strip_words(vec![String::from_str("\"hel\"lo world\""), String::from_str("hello")])
            == vec![String::from_str("hel\"lo world"), String::from_str("hello")]);
}

#[test]
fn expand_path_test() {
    // tests require the HOME env set
    let homep = Path::new(os::homedir().unwrap());
    assert!(expand_path(Path::new("~/Documents/scripts/")) == homep.join("Documents/scripts/"));
    assert!(expand_path(Path::new("/etc/wash/")) == Path::new("/etc/wash/"));
}

#[test]
fn condense_path_test() {
    // tests require the HOME env set
    let homep = Path::new(os::homedir().unwrap());
    assert!(condense_path(homep.join("Documents/scripts/")) ==
            Path::new("~/Documents/scripts/"));
    assert!(condense_path(Path::new("/home/")) == Path::new("/home/"));
    assert!(condense_path(Path::new("/etc/wash/")) == Path::new("/etc/wash/"));
}
<|endoftext|>


<commit_before>use crate::ram::cell::*;
use crate::ram::schema::Field;
use crate::ram::types;
use crate::ram::types::OwnedValue;

use std::{
    collections::{HashMap, HashSet}
};

use dovahkiin::types::{key_hash, Type};

enum InstData<'a> {
    Ref(&'a OwnedValue),
    Val(OwnedValue),
}

impl<'a> InstData<'a> {
    fn val_ref(&self) -> &OwnedValue {
        match self {
            InstData::Ref(r) => r,
            InstData::Val(v) => v,
        }
    }
}

pub struct Instruction<'a> {
    data_type: Type,
    val: InstData<'a>,
    offset: usize,
}

pub fn plan_write_field<'a>(
    tail_offset: &mut usize,
    field: &Field,
    value: &'a OwnedValue,
    mut ins: &mut Vec<Instruction<'a>>,
    is_var: bool,
) -> Result<(), WriteError> {
    let mut schema_offset = field.offset.clone();
    let is_field_var = field.is_var();
    let offset = if let Some(ref subs) = field.sub_fields {
        if let OwnedValue::Array(_) = value {
            if !field.is_array {
                return Err(WriteError::DataMismatchSchema(field.clone(), value.clone()));
            }
            if !is_var {
                trace!(
                    "Push jump tailing for map array with {} at {}",
                    tail_offset,
                    schema_offset.unwrap()
                );
                ins.push(Instruction {
                    data_type: Type::U32,
                    val: InstData::Val(OwnedValue::U32(*tail_offset as u32)),
                    offset: schema_offset.unwrap(),
                });
            }
            trace!(
                "Taking tailing offset for array {} at {}",
                field.name,
                tail_offset
            );
            tail_offset
        } else if let OwnedValue::Map(map) = value {
            trace!(
                "Writing map fields with for {} at {}",
                field.name,
                tail_offset
            );
            for sub in subs {
                let val = map.get_by_key_id(sub.name_id);
                plan_write_field(tail_offset, &sub, val, &mut ins, is_var)?;
            }
            return Ok(());
        } else {
            return Err(WriteError::DataMismatchSchema(field.clone(), value.clone()));
        }
    } else if is_field_var {
        // Write position tag for variable sized field
        if !is_var {
            // No need to jump to var region when it is var
            trace!(
                "Push jump tailing inst with {} at {:?}",
                tail_offset,
                schema_offset
            );
            ins.push(Instruction {
                data_type: Type::U32,
                val: InstData::Val(OwnedValue::U32(*tail_offset as u32)),
                offset: schema_offset.unwrap(),
            });
        }
        trace!("Using tailing offset for {} at {}", field.name, tail_offset);
        tail_offset
    } else if !is_var {
        schema_offset.as_mut().expect(&format!(
            "schema should have offset is_var: {}, field var: {}",
            is_var, is_field_var
        ))
    } else {
        tail_offset
    };
    trace!(
        "Plan to write {} at {}, field var {}, in var {}, value {:?}, field {:?}",
        field.name,
        offset,
        is_field_var,
        is_var,
        value,
        field
    );
    if field.nullable {
        let null_bit = match value {
            OwnedValue::Null => true,
            _ => false,
        };
        trace!("Push null bit inst with {} at {}", null_bit, *offset);
        ins.push(Instruction {
            data_type: Type::Bool,
            val: InstData::Val(OwnedValue::Bool(null_bit)),
            offset: *offset,
        });
        *offset += 1;
    }
    if field.is_array {
        if let OwnedValue::Array(array) = value {
            let len = array.len();
            let mut sub_field = field.clone();
            sub_field.is_array = false;
            trace!("Pushing array len inst with {} at {}", len, *offset);
            ins.push(Instruction {
                data_type: types::ARRAY_LEN_TYPE,
                val: InstData::Val(OwnedValue::U32(len as u32)),
                offset: *offset,
            });
            *offset += types::u32_io::type_size();
            for val in array {
                plan_write_field(offset, &sub_field, val, &mut ins, true)?;
            }
        } else if let OwnedValue::PrimArray(ref array) = value {
            let len = array.len();
            let size = array.size();
            // for prim array, just clone it and push into the instruction list with length
            trace!("Pushing prim array len inst with {} at {}", len, *offset);
            ins.push(Instruction {
                data_type: types::ARRAY_LEN_TYPE,
                val: InstData::Val(OwnedValue::U32(len as u32)),
                offset: *offset,
            });
            *offset += types::u32_io::type_size();
            trace!(
                "Pushing prim array ref inst with {:?} at {}",
                value,
                *offset
            );
            ins.push(Instruction {
                data_type: field.data_type,
                val: InstData::Ref(value),
                offset: *offset,
            });
            *offset += size;
        } else {
            return Err(WriteError::DataMismatchSchema(field.clone(), value.clone()));
        }
    } else {
        let is_null = match value {
            OwnedValue::Null => true,
            _ => false,
        };
        if !field.nullable && is_null {
            return Err(WriteError::DataMismatchSchema(field.clone(), value.clone()));
        }
        if !is_null {
            let size = types::get_vsize(field.data_type, &value);
            ins.push(Instruction {
                data_type: field.data_type,
                val: InstData::Ref(value),
                offset: *offset,
            });
            let new_offset = *offset + size;
            trace!(
                "Pushing value ref inst with {:?} at {}, size {}, new offset {}",
                value,
                *offset,
                size,
                new_offset
            );
            *offset = new_offset;
        }
    }
    return Ok(());
}

pub fn plan_write_dynamic_fields<'a>(
    offset: &mut usize,
    field: &Field,
    value: &'a OwnedValue,
    ins: &mut Vec<Instruction<'a>>,
) -> Result<(), WriteError> {
    if let (OwnedValue::Map(data_all), &Some(ref fields)) = (value, &field.sub_fields) {
        let schema_keys: HashSet<u64> = fields.iter().map(|f| f.name_id).collect();
        let dynamic_map: HashMap<_, _> = data_all
            .map
            .iter()
            .filter(|(k, _v)| !schema_keys.contains(k))
            .map(|(k, v)| (k, v))
            .collect();
        let dynamic_names: Vec<_> = data_all
            .fields
            .iter()
            .filter_map(|n| {
                let id = key_hash(&n);
                dynamic_map.get(&id).map(|_| n)
            })
            .collect();
        if !dynamic_map.is_empty() {}
        plan_write_dynamic_map(offset, &dynamic_names, &dynamic_map, ins)?;
    }
    return Ok(());
}

pub const ARRAY_TYPE_MASK: u8 = !(!0 << 1 >> 1); // 1000000...
pub const NULL_PLACEHOLDER: u8 = ARRAY_TYPE_MASK >> 1; // 1000000...

pub fn plan_write_dynamic_map<'a>(
    offset: &mut usize,
    names: &Vec<&String>,
    map: &HashMap<&u64, &'a OwnedValue>,
    ins: &mut Vec<Instruction<'a>>,
) -> Result<(), WriteError> {
    ins.push(Instruction {
        data_type: types::TYPE_CODE_TYPE,
        val: InstData::Val(OwnedValue::U8(Type::Map.id())),
        offset: *offset,
    });
    *offset += types::u8_io::type_size();
    // Write map size
    ins.push(Instruction {
        data_type: types::ARRAY_LEN_TYPE,
        val: InstData::Val(OwnedValue::U32(names.len() as u32)),
        offset: *offset,
    });
    *offset += types::u32_io::type_size();
    for name in names {
        let id = key_hash(name);
        let name_value = OwnedValue::String((*name).to_owned());
        let name_size = types::get_vsize(name_value.base_type(), &name_value);
        ins.push(Instruction {
            data_type: name_value.base_type(),
            val: InstData::Val(name_value),
            offset: *offset,
        });
        *offset += name_size;
        plan_write_dynamic_value(offset, map.get(&id).unwrap(), ins)?;
    }
    Ok(())
}

pub fn plan_write_dynamic_value<'a>(
    offset: &mut usize,
    value: &'a OwnedValue,
    ins: &mut Vec<Instruction<'a>>,
) -> Result<(), WriteError> {
    let base_type = value.base_type();
    match &value {
        &OwnedValue::Array(array) => {
            // Write type id
            ins.push(Instruction {
                data_type: types::TYPE_CODE_TYPE,
                val: InstData::Val(OwnedValue::U8(ARRAY_TYPE_MASK)), // Only put the mask cause we don't know the type
                offset: *offset,
            });
            *offset += types::u8_io::type_size();
            let len = array.len();
            // Write array length
            ins.push(Instruction {
                data_type: types::ARRAY_LEN_TYPE,
                val: InstData::Val(OwnedValue::U32(len as u32)),
                offset: *offset,
            });
            *offset += types::u32_io::type_size();
            for val in array {
                plan_write_dynamic_value(offset, val, ins)?;
            }
        }
        &OwnedValue::PrimArray(array) => {
            // Write type id with array tag
            ins.push(Instruction {
                data_type: types::TYPE_CODE_TYPE,
                val: InstData::Val(OwnedValue::U8(ARRAY_TYPE_MASK | base_type.id())),
                offset: *offset,
            });
            *offset += types::u8_io::type_size();
            let len = array.len();
            ins.push(Instruction {
                data_type: types::ARRAY_LEN_TYPE,
                val: InstData::Val(OwnedValue::U32(len as u32)),
                offset: *offset,
            });
            *offset += types::u32_io::type_size();
            let array_size = array.size();
            ins.push(Instruction {
                data_type: base_type,
                val: InstData::Ref(value),
                offset: *offset,
            });
            *offset += array_size;
        }
        &OwnedValue::Map(map) => plan_write_dynamic_map(
            offset,
            &map.fields.iter().collect(),
            &map.map.iter().collect(),
            ins,
        )?,
        &OwnedValue::Null | OwnedValue::NA => {
            // Write a placeholder because mapping required
            ins.push(Instruction {
                data_type: types::TYPE_CODE_TYPE,
                val: InstData::Val(OwnedValue::U8(NULL_PLACEHOLDER)),
                offset: *offset,
            });
            *offset += types::u8_io::type_size();
        }
        _ => {
            // Primitives
            let ty = value.base_type();
            ins.push(Instruction {
                data_type: types::TYPE_CODE_TYPE,
                val: InstData::Val(OwnedValue::U8(ty.id())),
                offset: *offset,
            });
            *offset += types::u8_io::type_size();
            let value_size = types::get_vsize(ty, &value);
            ins.push(Instruction {
                data_type: ty,
                val: InstData::Ref(value),
                offset: *offset,
            });
            *offset += value_size;
        }
    }
    Ok(())
}

pub fn execute_plan(ptr: usize, instructions: &Vec<Instruction>) {
    for ins in instructions {
        types::set_val(ins.data_type, ins.val.val_ref(), ptr + ins.offset);
    }
}
<commit_msg>cargo fmt<commit_after>use crate::ram::cell::*;
use crate::ram::schema::Field;
use crate::ram::types;
use crate::ram::types::OwnedValue;

use std::collections::{HashMap, HashSet};

use dovahkiin::types::{key_hash, Type};

enum InstData<'a> {
    Ref(&'a OwnedValue),
    Val(OwnedValue),
}

impl<'a> InstData<'a> {
    fn val_ref(&self) -> &OwnedValue {
        match self {
            InstData::Ref(r) => r,
            InstData::Val(v) => v,
        }
    }
}

pub struct Instruction<'a> {
    data_type: Type,
    val: InstData<'a>,
    offset: usize,
}

pub fn plan_write_field<'a>(
    tail_offset: &mut usize,
    field: &Field,
    value: &'a OwnedValue,
    mut ins: &mut Vec<Instruction<'a>>,
    is_var: bool,
) -> Result<(), WriteError> {
    let mut schema_offset = field.offset.clone();
    let is_field_var = field.is_var();
    let offset = if let Some(ref subs) = field.sub_fields {
        if let OwnedValue::Array(_) = value {
            if !field.is_array {
                return Err(WriteError::DataMismatchSchema(field.clone(), value.clone()));
            }
            if !is_var {
                trace!(
                    "Push jump tailing for map array with {} at {}",
                    tail_offset,
                    schema_offset.unwrap()
                );
                ins.push(Instruction {
                    data_type: Type::U32,
                    val: InstData::Val(OwnedValue::U32(*tail_offset as u32)),
                    offset: schema_offset.unwrap(),
                });
            }
            trace!(
                "Taking tailing offset for array {} at {}",
                field.name,
                tail_offset
            );
            tail_offset
        } else if let OwnedValue::Map(map) = value {
            trace!(
                "Writing map fields with for {} at {}",
                field.name,
                tail_offset
            );
            for sub in subs {
                let val = map.get_by_key_id(sub.name_id);
                plan_write_field(tail_offset, &sub, val, &mut ins, is_var)?;
            }
            return Ok(());
        } else {
            return Err(WriteError::DataMismatchSchema(field.clone(), value.clone()));
        }
    } else if is_field_var {
        // Write position tag for variable sized field
        if !is_var {
            // No need to jump to var region when it is var
            trace!(
                "Push jump tailing inst with {} at {:?}",
                tail_offset,
                schema_offset
            );
            ins.push(Instruction {
                data_type: Type::U32,
                val: InstData::Val(OwnedValue::U32(*tail_offset as u32)),
                offset: schema_offset.unwrap(),
            });
        }
        trace!("Using tailing offset for {} at {}", field.name, tail_offset);
        tail_offset
    } else if !is_var {
        schema_offset.as_mut().expect(&format!(
            "schema should have offset is_var: {}, field var: {}",
            is_var, is_field_var
        ))
    } else {
        tail_offset
    };
    trace!(
        "Plan to write {} at {}, field var {}, in var {}, value {:?}, field {:?}",
        field.name,
        offset,
        is_field_var,
        is_var,
        value,
        field
    );
    if field.nullable {
        let null_bit = match value {
            OwnedValue::Null => true,
            _ => false,
        };
        trace!("Push null bit inst with {} at {}", null_bit, *offset);
        ins.push(Instruction {
            data_type: Type::Bool,
            val: InstData::Val(OwnedValue::Bool(null_bit)),
            offset: *offset,
        });
        *offset += 1;
    }
    if field.is_array {
        if let OwnedValue::Array(array) = value {
            let len = array.len();
            let mut sub_field = field.clone();
            sub_field.is_array = false;
            trace!("Pushing array len inst with {} at {}", len, *offset);
            ins.push(Instruction {
                data_type: types::ARRAY_LEN_TYPE,
                val: InstData::Val(OwnedValue::U32(len as u32)),
                offset: *offset,
            });
            *offset += types::u32_io::type_size();
            for val in array {
                plan_write_field(offset, &sub_field, val, &mut ins, true)?;
            }
        } else if let OwnedValue::PrimArray(ref array) = value {
            let len = array.len();
            let size = array.size();
            // for prim array, just clone it and push into the instruction list with length
            trace!("Pushing prim array len inst with {} at {}", len, *offset);
            ins.push(Instruction {
                data_type: types::ARRAY_LEN_TYPE,
                val: InstData::Val(OwnedValue::U32(len as u32)),
                offset: *offset,
            });
            *offset += types::u32_io::type_size();
            trace!(
                "Pushing prim array ref inst with {:?} at {}",
                value,
                *offset
            );
            ins.push(Instruction {
                data_type: field.data_type,
                val: InstData::Ref(value),
                offset: *offset,
            });
            *offset += size;
        } else {
            return Err(WriteError::DataMismatchSchema(field.clone(), value.clone()));
        }
    } else {
        let is_null = match value {
            OwnedValue::Null => true,
            _ => false,
        };
        if !field.nullable && is_null {
            return Err(WriteError::DataMismatchSchema(field.clone(), value.clone()));
        }
        if !is_null {
            let size = types::get_vsize(field.data_type, &value);
            ins.push(Instruction {
                data_type: field.data_type,
                val: InstData::Ref(value),
                offset: *offset,
            });
            let new_offset = *offset + size;
            trace!(
                "Pushing value ref inst with {:?} at {}, size {}, new offset {}",
                value,
                *offset,
                size,
                new_offset
            );
            *offset = new_offset;
        }
    }
    return Ok(());
}

pub fn plan_write_dynamic_fields<'a>(
    offset: &mut usize,
    field: &Field,
    value: &'a OwnedValue,
    ins: &mut Vec<Instruction<'a>>,
) -> Result<(), WriteError> {
    if let (OwnedValue::Map(data_all), &Some(ref fields)) = (value, &field.sub_fields) {
        let schema_keys: HashSet<u64> = fields.iter().map(|f| f.name_id).collect();
        let dynamic_map: HashMap<_, _> = data_all
            .map
            .iter()
            .filter(|(k, _v)| !schema_keys.contains(k))
            .map(|(k, v)| (k, v))
            .collect();
        let dynamic_names: Vec<_> = data_all
            .fields
            .iter()
            .filter_map(|n| {
                let id = key_hash(&n);
                dynamic_map.get(&id).map(|_| n)
            })
            .collect();
        if !dynamic_map.is_empty() {}
        plan_write_dynamic_map(offset, &dynamic_names, &dynamic_map, ins)?;
    }
    return Ok(());
}

pub const ARRAY_TYPE_MASK: u8 = !(!0 << 1 >> 1); // 1000000...
pub const NULL_PLACEHOLDER: u8 = ARRAY_TYPE_MASK >> 1; // 1000000...

pub fn plan_write_dynamic_map<'a>(
    offset: &mut usize,
    names: &Vec<&String>,
    map: &HashMap<&u64, &'a OwnedValue>,
    ins: &mut Vec<Instruction<'a>>,
) -> Result<(), WriteError> {
    ins.push(Instruction {
        data_type: types::TYPE_CODE_TYPE,
        val: InstData::Val(OwnedValue::U8(Type::Map.id())),
        offset: *offset,
    });
    *offset += types::u8_io::type_size();
    // Write map size
    ins.push(Instruction {
        data_type: types::ARRAY_LEN_TYPE,
        val: InstData::Val(OwnedValue::U32(names.len() as u32)),
        offset: *offset,
    });
    *offset += types::u32_io::type_size();
    for name in names {
        let id = key_hash(name);
        let name_value = OwnedValue::String((*name).to_owned());
        let name_size = types::get_vsize(name_value.base_type(), &name_value);
        ins.push(Instruction {
            data_type: name_value.base_type(),
            val: InstData::Val(name_value),
            offset: *offset,
        });
        *offset += name_size;
        plan_write_dynamic_value(offset, map.get(&id).unwrap(), ins)?;
    }
    Ok(())
}

pub fn plan_write_dynamic_value<'a>(
    offset: &mut usize,
    value: &'a OwnedValue,
    ins: &mut Vec<Instruction<'a>>,
) -> Result<(), WriteError> {
    let base_type = value.base_type();
    match &value {
        &OwnedValue::Array(array) => {
            // Write type id
            ins.push(Instruction {
                data_type: types::TYPE_CODE_TYPE,
                val: InstData::Val(OwnedValue::U8(ARRAY_TYPE_MASK)), // Only put the mask cause we don't know the type
                offset: *offset,
            });
            *offset += types::u8_io::type_size();
            let len = array.len();
            // Write array length
            ins.push(Instruction {
                data_type: types::ARRAY_LEN_TYPE,
                val: InstData::Val(OwnedValue::U32(len as u32)),
                offset: *offset,
            });
            *offset += types::u32_io::type_size();
            for val in array {
                plan_write_dynamic_value(offset, val, ins)?;
            }
        }
        &OwnedValue::PrimArray(array) => {
            // Write type id with array tag
            ins.push(Instruction {
                data_type: types::TYPE_CODE_TYPE,
                val: InstData::Val(OwnedValue::U8(ARRAY_TYPE_MASK | base_type.id())),
                offset: *offset,
            });
            *offset += types::u8_io::type_size();
            let len = array.len();
            ins.push(Instruction {
                data_type: types::ARRAY_LEN_TYPE,
                val: InstData::Val(OwnedValue::U32(len as u32)),
                offset: *offset,
            });
            *offset += types::u32_io::type_size();
            let array_size = array.size();
            ins.push(Instruction {
                data_type: base_type,
                val: InstData::Ref(value),
                offset: *offset,
            });
            *offset += array_size;
        }
        &OwnedValue::Map(map) => plan_write_dynamic_map(
            offset,
            &map.fields.iter().collect(),
            &map.map.iter().collect(),
            ins,
        )?,
        &OwnedValue::Null | OwnedValue::NA => {
            // Write a placeholder because mapping required
            ins.push(Instruction {
                data_type: types::TYPE_CODE_TYPE,
                val: InstData::Val(OwnedValue::U8(NULL_PLACEHOLDER)),
                offset: *offset,
            });
            *offset += types::u8_io::type_size();
        }
        _ => {
            // Primitives
            let ty = value.base_type();
            ins.push(Instruction {
                data_type: types::TYPE_CODE_TYPE,
                val: InstData::Val(OwnedValue::U8(ty.id())),
                offset: *offset,
            });
            *offset += types::u8_io::type_size();
            let value_size = types::get_vsize(ty, &value);
            ins.push(Instruction {
                data_type: ty,
                val: InstData::Ref(value),
                offset: *offset,
            });
            *offset += value_size;
        }
    }
    Ok(())
}

pub fn execute_plan(ptr: usize, instructions: &Vec<Instruction>) {
    for ins in instructions {
        types::set_val(ins.data_type, ins.val.val_ref(), ptr + ins.offset);
    }
}
<|endoftext|>


<commit_before>use super::{metadata::{read_metafile,
                       MetaFile},
            PackageIdent,
            PackageTarget};
use crate::error::{Error,
                   Result};
use std::{ffi::OsStr,
          fs,
          io,
          path::{Path,
                 PathBuf},
          str::FromStr};
use tempfile::{Builder,
               TempDir};

pub const INSTALL_TMP_PREFIX: &str = ".hab-pkg-install";

/// Return a directory which can be used as a temp dir during package install/
/// uninstall
///
/// It returns a path which is in the same parent directory as `path`
/// but with TempDir style randomization
pub fn temp_package_directory(path: &Path) -> Result<TempDir> {
    let base = path.parent().ok_or_else(|| {
                                 Error::PackageUnpackFailed(
            "Could not determine parent directory for temporary package directory".to_owned(),
        )
                             })?;
    fs::create_dir_all(base)?;
    let temp_install_prefix =
        path.file_name()
            .and_then(OsStr::to_str)
            .map(|dirname| format!("{}-{}", INSTALL_TMP_PREFIX, dirname))
            .ok_or_else(|| {
                Error::PackageUnpackFailed("Could not generate prefix for temporary package \
                                            directory"
                                                      .to_owned())
            })?;
    Ok(Builder::new().prefix(&temp_install_prefix)
                     .tempdir_in(base)?)
}

/// Returns a list of package structs built from the contents of the given directory.
pub fn all_packages(path: &Path) -> Result<Vec<PackageIdent>> {
    let mut package_list: Vec<PackageIdent> = vec![];
    if fs::metadata(path)?.is_dir() {
        walk_origins(&path, &mut package_list)?;
    }
    Ok(package_list)
}

/// Returns a vector of package idents built from the contents of
/// the given directory, using the given origin to restrict the
/// search.
///
/// The search is restricted by assuming the package directory
/// structure is:
///
///    /base/ORIGIN/NAME/VERSION/RELEASE/
pub fn package_list_for_origin(base_pkg_path: &Path, origin: &str) -> Result<Vec<PackageIdent>> {
    let mut package_list: Vec<PackageIdent> = vec![];
    let mut package_path = PathBuf::from(base_pkg_path);
    package_path.push(&origin);

    if !is_existing_dir(&package_path)? {
        return Ok(package_list);
    };

    walk_names(&origin, &package_path, &mut package_list)?;
    Ok(package_list)
}

/// Returns a vector of package structs built from the contents of
/// the given directory, using the given ident to restrict the
/// search.
///
/// The search is restricted by assuming the package directory
/// structure is:
///
///    /base/ORIGIN/NAME/VERSION/RELEASE/
pub fn package_list_for_ident(base_pkg_path: &Path,
                              ident: &PackageIdent)
                              -> Result<Vec<PackageIdent>> {
    let mut package_list: Vec<PackageIdent> = vec![];
    let mut package_path = PathBuf::from(base_pkg_path);
    package_path.push(&ident.origin);
    package_path.push(&ident.name);

    if !is_existing_dir(&package_path)? {
        return Ok(package_list);
    }

    match (&ident.version, &ident.release) {
        // origin/name
        (None, _) => walk_versions(&ident.origin, &ident.name, &package_path, &mut package_list)?,
        // origin/name/version
        (Some(version), None) => {
            package_path.push(version);
            if !is_existing_dir(&package_path)? {
                return Ok(package_list);
            }
            walk_releases(&ident.origin,
                          &ident.name,
                          &version,
                          &package_path,
                          &mut package_list)?
        }
        // origin/name/version/release
        (Some(version), Some(release)) => {
            package_path.push(version);
            package_path.push(release);
            if !is_existing_dir(&package_path)? {
                return Ok(package_list);
            }

            let active_target = PackageTarget::active_target();
            if let Some(new_ident) = package_ident_from_dir(&ident.origin,
                                                            &ident.name,
                                                            &version,
                                                            active_target,
                                                            &package_path)
            {
                package_list.push(new_ident)
            }
        }
    }
    Ok(package_list)
}

/// Helper function for all_packages. Walks the directory at the given
/// Path for origin directories and builds on the given package list
/// by recursing into name, version, and release directories.
fn walk_origins(path: &Path, packages: &mut Vec<PackageIdent>) -> Result<()> {
    for entry in fs::read_dir(path)? {
        let origin_dir = entry?;
        let origin_path = origin_dir.path();
        if fs::metadata(&origin_path)?.is_dir() {
            let origin = filename_from_entry(&origin_dir);
            walk_names(&origin, &origin_path, packages)?;
        }
    }
    Ok(())
}

/// Helper function for walk_origins. Walks the direcotry at the given
/// Path for name directories and recurses into them to find version
/// and release directories.
fn walk_names(origin: &str, dir: &Path, packages: &mut Vec<PackageIdent>) -> Result<()> {
    for entry in fs::read_dir(dir)? {
        let name_dir = entry?;
        let name_path = name_dir.path();
        if fs::metadata(&name_path)?.is_dir() {
            let name = filename_from_entry(&name_dir);
            walk_versions(&origin, &name, &name_path, packages)?;
        }
    }
    Ok(())
}

/// Helper function for walk_names. Walks the directory at the given
/// Path and recurses into them to find release directories.
fn walk_versions(origin: &str,
                 name: &str,
                 dir: &Path,
                 packages: &mut Vec<PackageIdent>)
                 -> Result<()> {
    for entry in fs::read_dir(dir)? {
        let version_dir = entry?;
        let version_path = version_dir.path();
        if fs::metadata(&version_path)?.is_dir() {
            let version = filename_from_entry(&version_dir);
            walk_releases(origin, name, &version, &version_path, packages)?;
        }
    }
    Ok(())
}

/// Helper function for walk_versions. Walks the directory at the
/// given Path and constructs a Package struct if the directory is a
/// valid package directory. Any resulting packages are pushed onto
/// the given packages vector, assuming the given origin, name, and
/// version.
fn walk_releases(origin: &str,
                 name: &str,
                 version: &str,
                 dir: &Path,
                 packages: &mut Vec<PackageIdent>)
                 -> Result<()> {
    let active_target = PackageTarget::active_target();
    for entry in fs::read_dir(dir)? {
        let release_dir = entry?;
        let release_path = release_dir.path();
        if fs::metadata(&release_path)?.is_dir() {
            if let Some(ident) =
                package_ident_from_dir(origin, name, version, active_target, &release_path)
            {
                packages.push(ident)
            }
        }
    }
    Ok(())
}

/// package_ident_from_dir returns a PackageIdent if the given
/// path contains a valid package for the given active_target.
///
/// Returns None when
///    - The directory is a temporary install directroy
///    - An error occurs reading the package metadata
///    - An error occurs reading the package target
///    - The package target doesn't match the given active target
fn package_ident_from_dir(origin: &str,
                          name: &str,
                          version: &str,
                          active_target: PackageTarget,
                          dir: &Path)
                          -> Option<PackageIdent> {
    let release = dir.file_name().and_then(OsStr::to_str)?;

    if release.starts_with(INSTALL_TMP_PREFIX) {
        debug!("PackageInstall::package_ident_from_dir(): rejected PackageInstall candidate \
                because it matches installation temporary directory prefix: {}",
               dir.display());
        return None;
    }

    let metafile_content = read_metafile(dir, MetaFile::Target);
    // If there is an error reading the target metafile, then skip the candidate
    if let Err(e) = metafile_content {
        debug!("PackageInstall::package_ident_from_dir(): rejected PackageInstall candidate due \
                to error reading TARGET metafile, found={}, reason={:?}",
               dir.display(),
               e,);
        return None;
    }

    // Any errors have been cleared, so unwrap is safe
    let metafile_content = metafile_content.unwrap();
    let install_target = PackageTarget::from_str(&metafile_content);
    // If there is an error parsing the target as a valid PackageTarget, then skip the
    // candidate
    if let Err(e) = install_target {
        debug!("PackageInstall::package_ident_from_dir(): rejected PackageInstall candidate due \
                to error parsing TARGET metafile as a valid PackageTarget, found={}, reason={:?}",
               dir.display(),
               e,);
        return None;
    }
    // Any errors have been cleared, so unwrap is safe
    let install_target = install_target.unwrap();

    // Ensure that the installed package's target matches the active `PackageTarget`,
    // otherwise skip the candidate
    if active_target == install_target {
        Some(PackageIdent::new(origin.to_string(),
                               name.to_string(),
                               Some(version.to_string()),
                               Some(release.to_owned())))
    } else {
        debug!("PackageInstall::package_ident_from_dir(): rejected PackageInstall candidate, \
                found={}, installed_target={}, active_target={}",
               dir.display(),
               install_target,
               active_target,);
        None
    }
}

fn filename_from_entry(entry: &fs::DirEntry) -> String {
    entry.file_name().to_string_lossy().into_owned()
}

fn is_existing_dir(path: &Path) -> Result<bool> {
    match fs::metadata(&path) {
        Err(err) => {
            if err.kind() == io::ErrorKind::NotFound {
                return Ok(false);
            }
            Err(Error::from(err))
        }
        Ok(metadata) => Ok(metadata.is_dir()),
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::package::test_support::testing_package_install;

    use crate::fs;
    use std::fs::File;
    use tempfile::Builder;

    #[test]
    fn empty_dir_gives_empty_list() {
        let package_root = Builder::new().prefix("fs-root").tempdir().unwrap();
        let packages = all_packages(&package_root.path()).unwrap();

        assert_eq!(0, packages.len());
    }

    #[test]
    fn not_a_dir_gives_empty_list() {
        let fs_root = Builder::new().prefix("fs-root").tempdir().unwrap();
        let file_path = fs_root.path().join("not-a-dir");
        let _tmp_file = File::create(&file_path).unwrap();

        let packages = all_packages(&file_path).unwrap();

        assert_eq!(0, packages.len());
    }
    #[test]
    fn can_read_single_package() {
        let fs_root = Builder::new().prefix("fs-root").tempdir().unwrap();
        let package_root = fs::pkg_root_path(Some(fs_root.path()));
        let package_install = testing_package_install("core/redis", fs_root.path());

        let packages = all_packages(&package_root).unwrap();

        assert_eq!(1, packages.len());
        assert_eq!(package_install.ident, packages[0]);
    }

    #[test]
    fn can_read_multiple_packages() {
        let fs_root = Builder::new().prefix("fs-root").tempdir().unwrap();
        let package_root = fs::pkg_root_path(Some(fs_root.path()));
        let expected = vec![testing_package_install("core/redis/1.0.0", fs_root.path()),
                            testing_package_install("test/foobar", fs_root.path()),
                            testing_package_install("core/redis/1.1.0", fs_root.path()),];

        let packages = all_packages(&package_root).unwrap();

        assert_eq!(3, packages.len());
        for p in &expected {
            assert!(packages.contains(&p.ident));
        }
    }

    #[test]
    fn create_temp_package_directory_in_same_parentdir() {
        let p = Builder::new().tempdir().unwrap();
        let temp_dir = temp_package_directory(p.path()).unwrap();
        let temp_path = temp_dir.path();

        assert_eq!(&p.path().parent(), &temp_path.parent());
    }

    #[test]
    fn temp_package_directory_starts_with_prefix() {
        let p = Builder::new().tempdir().unwrap();
        let temp_dir = temp_package_directory(p.path()).unwrap();
        let temp_filename = temp_dir.path().file_name().unwrap().to_str().unwrap();

        assert!(&temp_filename.starts_with(INSTALL_TMP_PREFIX));
    }

    #[test]
    fn temp_package_directory_changes() {
        let p = Builder::new().tempdir().unwrap();
        let temp_dir1 = temp_package_directory(p.path()).unwrap();
        let temp_dir2 = temp_package_directory(p.path()).unwrap();

        assert_ne!(&temp_dir1.path(), &temp_dir2.path());
    }

    #[test]
    fn list_for_origin_skips_non_origin_packages() {
        let fs_root = Builder::new().prefix("fs-root").tempdir().unwrap();
        let package_root = fs::pkg_root_path(Some(fs_root.path()));
        let test_origin = vec![testing_package_install("test/foobar", fs_root.path())];
        let core_origin = vec![testing_package_install("core/redis/1.0.0", fs_root.path()),
                               testing_package_install("core/redis/1.1.0", fs_root.path()),];

        let packages = package_list_for_origin(&package_root, &String::from("core")).unwrap();

        assert_eq!(2, packages.len());
        for p in &core_origin {
            assert!(packages.contains(&p.ident));
        }

        for p in &test_origin {
            assert!(!packages.contains(&p.ident));
        }
    }

    #[test]
    fn list_for_origin_no_packages() {
        let fs_root = Builder::new().prefix("fs-root").tempdir().unwrap();
        let package_root = fs::pkg_root_path(Some(fs_root.path()));

        let packages = package_list_for_origin(&package_root, &String::from("core")).unwrap();

        assert_eq!(0, packages.len());
    }
}
<commit_msg>canonicalize base of temp_package_directory to fix unpacking archives with long paths on windows<commit_after>use super::{metadata::{read_metafile,
                       MetaFile},
            PackageIdent,
            PackageTarget};
use crate::error::{Error,
                   Result};
use std::{ffi::OsStr,
          fs,
          io,
          path::{Path,
                 PathBuf},
          str::FromStr};
use tempfile::{Builder,
               TempDir};

pub const INSTALL_TMP_PREFIX: &str = ".hab-pkg-install";

/// Return a directory which can be used as a temp dir during package install/
/// uninstall
///
/// It returns a path which is in the same parent directory as `path`
/// but with TempDir style randomization
pub fn temp_package_directory(path: &Path) -> Result<TempDir> {
    let base = path.parent().ok_or_else(|| {
                                 Error::PackageUnpackFailed(
            "Could not determine parent directory for temporary package directory".to_owned(),
        )
                             })?;
    fs::create_dir_all(base)?;
    let cannon_base = base.canonicalize()?;
    let temp_install_prefix =
        path.file_name()
            .and_then(OsStr::to_str)
            .map(|dirname| format!("{}-{}", INSTALL_TMP_PREFIX, dirname))
            .ok_or_else(|| {
                Error::PackageUnpackFailed("Could not generate prefix for temporary package \
                                            directory"
                                                      .to_owned())
            })?;
    Ok(Builder::new().prefix(&temp_install_prefix)
                     .tempdir_in(cannon_base)?)
}

/// Returns a list of package structs built from the contents of the given directory.
pub fn all_packages(path: &Path) -> Result<Vec<PackageIdent>> {
    let mut package_list: Vec<PackageIdent> = vec![];
    if fs::metadata(path)?.is_dir() {
        walk_origins(&path, &mut package_list)?;
    }
    Ok(package_list)
}

/// Returns a vector of package idents built from the contents of
/// the given directory, using the given origin to restrict the
/// search.
///
/// The search is restricted by assuming the package directory
/// structure is:
///
///    /base/ORIGIN/NAME/VERSION/RELEASE/
pub fn package_list_for_origin(base_pkg_path: &Path, origin: &str) -> Result<Vec<PackageIdent>> {
    let mut package_list: Vec<PackageIdent> = vec![];
    let mut package_path = PathBuf::from(base_pkg_path);
    package_path.push(&origin);

    if !is_existing_dir(&package_path)? {
        return Ok(package_list);
    };

    walk_names(&origin, &package_path, &mut package_list)?;
    Ok(package_list)
}

/// Returns a vector of package structs built from the contents of
/// the given directory, using the given ident to restrict the
/// search.
///
/// The search is restricted by assuming the package directory
/// structure is:
///
///    /base/ORIGIN/NAME/VERSION/RELEASE/
pub fn package_list_for_ident(base_pkg_path: &Path,
                              ident: &PackageIdent)
                              -> Result<Vec<PackageIdent>> {
    let mut package_list: Vec<PackageIdent> = vec![];
    let mut package_path = PathBuf::from(base_pkg_path);
    package_path.push(&ident.origin);
    package_path.push(&ident.name);

    if !is_existing_dir(&package_path)? {
        return Ok(package_list);
    }

    match (&ident.version, &ident.release) {
        // origin/name
        (None, _) => walk_versions(&ident.origin, &ident.name, &package_path, &mut package_list)?,
        // origin/name/version
        (Some(version), None) => {
            package_path.push(version);
            if !is_existing_dir(&package_path)? {
                return Ok(package_list);
            }
            walk_releases(&ident.origin,
                          &ident.name,
                          &version,
                          &package_path,
                          &mut package_list)?
        }
        // origin/name/version/release
        (Some(version), Some(release)) => {
            package_path.push(version);
            package_path.push(release);
            if !is_existing_dir(&package_path)? {
                return Ok(package_list);
            }

            let active_target = PackageTarget::active_target();
            if let Some(new_ident) = package_ident_from_dir(&ident.origin,
                                                            &ident.name,
                                                            &version,
                                                            active_target,
                                                            &package_path)
            {
                package_list.push(new_ident)
            }
        }
    }
    Ok(package_list)
}

/// Helper function for all_packages. Walks the directory at the given
/// Path for origin directories and builds on the given package list
/// by recursing into name, version, and release directories.
fn walk_origins(path: &Path, packages: &mut Vec<PackageIdent>) -> Result<()> {
    for entry in fs::read_dir(path)? {
        let origin_dir = entry?;
        let origin_path = origin_dir.path();
        if fs::metadata(&origin_path)?.is_dir() {
            let origin = filename_from_entry(&origin_dir);
            walk_names(&origin, &origin_path, packages)?;
        }
    }
    Ok(())
}

/// Helper function for walk_origins. Walks the direcotry at the given
/// Path for name directories and recurses into them to find version
/// and release directories.
fn walk_names(origin: &str, dir: &Path, packages: &mut Vec<PackageIdent>) -> Result<()> {
    for entry in fs::read_dir(dir)? {
        let name_dir = entry?;
        let name_path = name_dir.path();
        if fs::metadata(&name_path)?.is_dir() {
            let name = filename_from_entry(&name_dir);
            walk_versions(&origin, &name, &name_path, packages)?;
        }
    }
    Ok(())
}

/// Helper function for walk_names. Walks the directory at the given
/// Path and recurses into them to find release directories.
fn walk_versions(origin: &str,
                 name: &str,
                 dir: &Path,
                 packages: &mut Vec<PackageIdent>)
                 -> Result<()> {
    for entry in fs::read_dir(dir)? {
        let version_dir = entry?;
        let version_path = version_dir.path();
        if fs::metadata(&version_path)?.is_dir() {
            let version = filename_from_entry(&version_dir);
            walk_releases(origin, name, &version, &version_path, packages)?;
        }
    }
    Ok(())
}

/// Helper function for walk_versions. Walks the directory at the
/// given Path and constructs a Package struct if the directory is a
/// valid package directory. Any resulting packages are pushed onto
/// the given packages vector, assuming the given origin, name, and
/// version.
fn walk_releases(origin: &str,
                 name: &str,
                 version: &str,
                 dir: &Path,
                 packages: &mut Vec<PackageIdent>)
                 -> Result<()> {
    let active_target = PackageTarget::active_target();
    for entry in fs::read_dir(dir)? {
        let release_dir = entry?;
        let release_path = release_dir.path();
        if fs::metadata(&release_path)?.is_dir() {
            if let Some(ident) =
                package_ident_from_dir(origin, name, version, active_target, &release_path)
            {
                packages.push(ident)
            }
        }
    }
    Ok(())
}

/// package_ident_from_dir returns a PackageIdent if the given
/// path contains a valid package for the given active_target.
///
/// Returns None when
///    - The directory is a temporary install directroy
///    - An error occurs reading the package metadata
///    - An error occurs reading the package target
///    - The package target doesn't match the given active target
fn package_ident_from_dir(origin: &str,
                          name: &str,
                          version: &str,
                          active_target: PackageTarget,
                          dir: &Path)
                          -> Option<PackageIdent> {
    let release = dir.file_name().and_then(OsStr::to_str)?;

    if release.starts_with(INSTALL_TMP_PREFIX) {
        debug!("PackageInstall::package_ident_from_dir(): rejected PackageInstall candidate \
                because it matches installation temporary directory prefix: {}",
               dir.display());
        return None;
    }

    let metafile_content = read_metafile(dir, MetaFile::Target);
    // If there is an error reading the target metafile, then skip the candidate
    if let Err(e) = metafile_content {
        debug!("PackageInstall::package_ident_from_dir(): rejected PackageInstall candidate due \
                to error reading TARGET metafile, found={}, reason={:?}",
               dir.display(),
               e,);
        return None;
    }

    // Any errors have been cleared, so unwrap is safe
    let metafile_content = metafile_content.unwrap();
    let install_target = PackageTarget::from_str(&metafile_content);
    // If there is an error parsing the target as a valid PackageTarget, then skip the
    // candidate
    if let Err(e) = install_target {
        debug!("PackageInstall::package_ident_from_dir(): rejected PackageInstall candidate due \
                to error parsing TARGET metafile as a valid PackageTarget, found={}, reason={:?}",
               dir.display(),
               e,);
        return None;
    }
    // Any errors have been cleared, so unwrap is safe
    let install_target = install_target.unwrap();

    // Ensure that the installed package's target matches the active `PackageTarget`,
    // otherwise skip the candidate
    if active_target == install_target {
        Some(PackageIdent::new(origin.to_string(),
                               name.to_string(),
                               Some(version.to_string()),
                               Some(release.to_owned())))
    } else {
        debug!("PackageInstall::package_ident_from_dir(): rejected PackageInstall candidate, \
                found={}, installed_target={}, active_target={}",
               dir.display(),
               install_target,
               active_target,);
        None
    }
}

fn filename_from_entry(entry: &fs::DirEntry) -> String {
    entry.file_name().to_string_lossy().into_owned()
}

fn is_existing_dir(path: &Path) -> Result<bool> {
    match fs::metadata(&path) {
        Err(err) => {
            if err.kind() == io::ErrorKind::NotFound {
                return Ok(false);
            }
            Err(Error::from(err))
        }
        Ok(metadata) => Ok(metadata.is_dir()),
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::package::test_support::testing_package_install;

    use crate::fs;
    use std::fs::File;
    use tempfile::Builder;

    #[test]
    fn empty_dir_gives_empty_list() {
        let package_root = Builder::new().prefix("fs-root").tempdir().unwrap();
        let packages = all_packages(&package_root.path()).unwrap();

        assert_eq!(0, packages.len());
    }

    #[test]
    fn not_a_dir_gives_empty_list() {
        let fs_root = Builder::new().prefix("fs-root").tempdir().unwrap();
        let file_path = fs_root.path().join("not-a-dir");
        let _tmp_file = File::create(&file_path).unwrap();

        let packages = all_packages(&file_path).unwrap();

        assert_eq!(0, packages.len());
    }
    #[test]
    fn can_read_single_package() {
        let fs_root = Builder::new().prefix("fs-root").tempdir().unwrap();
        let package_root = fs::pkg_root_path(Some(fs_root.path()));
        let package_install = testing_package_install("core/redis", fs_root.path());

        let packages = all_packages(&package_root).unwrap();

        assert_eq!(1, packages.len());
        assert_eq!(package_install.ident, packages[0]);
    }

    #[test]
    fn can_read_multiple_packages() {
        let fs_root = Builder::new().prefix("fs-root").tempdir().unwrap();
        let package_root = fs::pkg_root_path(Some(fs_root.path()));
        let expected = vec![testing_package_install("core/redis/1.0.0", fs_root.path()),
                            testing_package_install("test/foobar", fs_root.path()),
                            testing_package_install("core/redis/1.1.0", fs_root.path()),];

        let packages = all_packages(&package_root).unwrap();

        assert_eq!(3, packages.len());
        for p in &expected {
            assert!(packages.contains(&p.ident));
        }
    }

    #[test]
    fn create_temp_package_directory_in_same_parentdir() {
        let p = Builder::new().tempdir().unwrap();
        let temp_dir = temp_package_directory(p.path()).unwrap();
        let temp_path = temp_dir.path();

        assert_eq!(&p.path().parent().unwrap().canonicalize().unwrap(),
                   &temp_path.parent().unwrap());
    }

    #[test]
    fn temp_package_directory_starts_with_prefix() {
        let p = Builder::new().tempdir().unwrap();
        let temp_dir = temp_package_directory(p.path()).unwrap();
        let temp_filename = temp_dir.path().file_name().unwrap().to_str().unwrap();

        assert!(&temp_filename.starts_with(INSTALL_TMP_PREFIX));
    }

    #[test]
    fn temp_package_directory_changes() {
        let p = Builder::new().tempdir().unwrap();
        let temp_dir1 = temp_package_directory(p.path()).unwrap();
        let temp_dir2 = temp_package_directory(p.path()).unwrap();

        assert_ne!(&temp_dir1.path(), &temp_dir2.path());
    }

    #[test]
    fn list_for_origin_skips_non_origin_packages() {
        let fs_root = Builder::new().prefix("fs-root").tempdir().unwrap();
        let package_root = fs::pkg_root_path(Some(fs_root.path()));
        let test_origin = vec![testing_package_install("test/foobar", fs_root.path())];
        let core_origin = vec![testing_package_install("core/redis/1.0.0", fs_root.path()),
                               testing_package_install("core/redis/1.1.0", fs_root.path()),];

        let packages = package_list_for_origin(&package_root, &String::from("core")).unwrap();

        assert_eq!(2, packages.len());
        for p in &core_origin {
            assert!(packages.contains(&p.ident));
        }

        for p in &test_origin {
            assert!(!packages.contains(&p.ident));
        }
    }

    #[test]
    fn list_for_origin_no_packages() {
        let fs_root = Builder::new().prefix("fs-root").tempdir().unwrap();
        let package_root = fs::pkg_root_path(Some(fs_root.path()));

        let packages = package_list_for_origin(&package_root, &String::from("core")).unwrap();

        assert_eq!(0, packages.len());
    }
}
<|endoftext|>


<commit_before>// Copyright 2019 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use darling::{FromDeriveInput, FromMeta};
use heck::SnakeCase;
use proc_macro::TokenStream;
use proc_macro2::{Ident, Span};
use quote::{quote, ToTokens};
use syn::{Attribute, Data, DataEnum, DataStruct, DeriveInput, Path};

use std::convert::TryFrom;

use super::{find_exonum_meta, CratePath};

#[derive(Debug, FromMeta)]
#[darling(default)]
struct ProtobufConvertStructAttrs {
    #[darling(rename = "crate")]
    cr: CratePath,
    pb: Option<Path>,
    serde_pb_convert: bool,
}

impl Default for ProtobufConvertStructAttrs {
    fn default() -> Self {
        Self {
            cr: CratePath::default(),
            pb: None,
            serde_pb_convert: false,
        }
    }
}

impl TryFrom<&[Attribute]> for ProtobufConvertStructAttrs {
    type Error = darling::Error;

    fn try_from(args: &[Attribute]) -> Result<Self, Self::Error> {
        find_exonum_meta(args)
            .map(|meta| Self::from_nested_meta(&meta))
            .unwrap_or_else(|| Ok(Self::default()))
    }
}

#[derive(Debug, FromMeta)]
#[darling(default)]
struct ProtobufConvertEnumAttrs {
    #[darling(rename = "crate")]
    cr: CratePath,
    pb: Option<Path>,
    serde_pb_convert: bool,
    oneof_field: Ident,
}

impl Default for ProtobufConvertEnumAttrs {
    fn default() -> Self {
        Self {
            cr: CratePath::default(),
            pb: None,
            oneof_field: syn::parse_str("message").unwrap(),
            serde_pb_convert: false,
        }
    }
}

impl TryFrom<&[Attribute]> for ProtobufConvertEnumAttrs {
    type Error = darling::Error;

    fn try_from(args: &[Attribute]) -> Result<Self, Self::Error> {
        find_exonum_meta(args)
            .map(|meta| Self::from_nested_meta(&meta))
            .unwrap_or_else(|| Ok(Self::default()))
    }
}

#[derive(Debug)]
struct ProtobufConvertStruct {
    name: Ident,
    fields: Vec<Ident>,
    attrs: ProtobufConvertStructAttrs,
}

impl ProtobufConvertStruct {
    fn from_derive_input(
        name: Ident,
        data: &DataStruct,
        attrs: &[Attribute],
    ) -> Result<Self, darling::Error> {
        let attrs = ProtobufConvertStructAttrs::try_from(attrs)?;
        let fields = data
            .fields
            .iter()
            .map(|f| f.ident.clone().unwrap())
            .collect();

        Ok(Self {
            name,
            attrs,
            fields,
        })
    }
}

impl ToTokens for ProtobufConvertStruct {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        let name = &self.name;
        let pb_name = &self.attrs.pb;
        let from_pb_impl = {
            let getters = self
                .fields
                .iter()
                .map(|i| Ident::new(&format!("get_{}", i), Span::call_site()));
            let fields = self.fields.to_vec();

            quote! {
                let inner = Self {
                    #( #fields: ProtobufConvert::from_pb(pb.#getters().to_owned())?, )*
                };
                Ok(inner)
            }
        };
        let to_pb_impl = {
            let setters = self
                .fields
                .iter()
                .map(|i| Ident::new(&format!("set_{}", i), Span::call_site()));
            let fields = self.fields.to_vec();

            quote! {
                let mut msg = Self::ProtoStruct::default();
                #( msg.#setters(ProtobufConvert::to_pb(&self.#fields).into()); )*
                msg
            }
        };

        let expanded = quote! {
            impl ProtobufConvert for #name {
                type ProtoStruct = #pb_name;

                fn from_pb(pb: Self::ProtoStruct) -> std::result::Result<Self, failure::Error> {
                    #from_pb_impl
                }

                fn to_pb(&self) -> Self::ProtoStruct {
                    #to_pb_impl
                }
            }
        };
        tokens.extend(expanded);
    }
}

#[derive(Debug)]
struct ProtobufConvertEnum {
    name: Ident,
    variants: Vec<Ident>,
    attrs: ProtobufConvertEnumAttrs,
}

impl ProtobufConvertEnum {
    fn from_derive_input(
        name: Ident,
        data: &DataEnum,
        attrs: &[Attribute],
    ) -> Result<Self, darling::Error> {
        let attrs = ProtobufConvertEnumAttrs::try_from(attrs)?;
        let variants = data.variants.iter().map(|v| v.ident.clone()).collect();

        Ok(Self {
            name,
            attrs,
            variants,
        })
    }

    fn impl_protobuf_convert(&self) -> impl ToTokens {
        let pb_oneof_enum = {
            let mut pb = self.attrs.pb.clone().unwrap();
            let oneof = pb.segments.pop().unwrap().value().ident.clone();
            let oneof_enum = Ident::new(
                &format!("{}_oneof_{}", oneof, &self.attrs.oneof_field),
                Span::call_site(),
            );
            quote! { #pb #oneof_enum }
        };
        let name = &self.name;
        let pb_name = &self.attrs.pb;
        let oneof = &self.attrs.oneof_field;

        let from_pb_impl = {
            let match_arms = self.variants.iter().map(|variant| {
                let pb_variant = Ident::new(
                    variant.to_string().to_snake_case().as_ref(),
                    Span::call_site(),
                );
                quote! {
                    Some(#pb_oneof_enum::#pb_variant(pb)) => {
                        #variant::from_pb(pb).map(#name::#variant)
                    }
                }
            });

            quote! {
                match pb.#oneof {
                    #( #match_arms )*
                    None => Err(failure::format_err!("Failed to decode #name from protobuf"))
                }
            }
        };
        let to_pb_impl = {
            let match_arms = self.variants.iter().map(|variant| {
                let pb_variant = variant.to_string().to_snake_case();
                let setter = Ident::new(&format!("set_{}", pb_variant), Span::call_site());
                quote! {
                    #name::#variant(msg) => inner.#setter(msg.to_pb()),
                }
            });

            quote! {
                let mut inner = Self::ProtoStruct::new();
                match self {
                    #( #match_arms )*
                }
                inner
            }
        };

        quote! {
            impl ProtobufConvert for #name {
                type ProtoStruct = #pb_name;

                fn from_pb(mut pb: Self::ProtoStruct) -> std::result::Result<Self, failure::Error> {
                    #from_pb_impl
                }

                fn to_pb(&self) -> Self::ProtoStruct {
                    #to_pb_impl
                }
            }
        }
    }

    fn impl_enum_conversions(&self) -> impl ToTokens {
        let name = &self.name;
        let conversions = self.variants.iter().map(|variant| {
            quote! {
                impl From<#variant> for #name {
                    fn from(variant: #variant) -> Self {
                        #name::#variant(variant)
                    }
                }

                impl std::convert::TryFrom<#name> for #variant {
                    type Error = failure::Error;

                    fn try_from(msg: #name) -> Result<Self, Self::Error> {
                        if let #name::#variant(inner) = msg {
                            Ok(inner)
                        } else {
                            Err(failure::format_err!(
                                "Expected #variant variant, but got {:?}",
                                msg
                            ))
                        }
                    }
                }
            }
        });

        quote! {
            #( #conversions )*
        }
    }
}

impl ToTokens for ProtobufConvertEnum {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        let pb_convert = self.impl_protobuf_convert();
        let conversions = self.impl_enum_conversions();

        let expanded = quote! {
            #pb_convert
            #conversions
        };
        tokens.extend(expanded)
    }
}

#[derive(Debug)]
enum ProtobufConvert {
    Enum(ProtobufConvertEnum),
    Struct(ProtobufConvertStruct),
}

impl FromDeriveInput for ProtobufConvert {
    fn from_derive_input(input: &DeriveInput) -> Result<Self, darling::Error> {
        match &input.data {
            Data::Struct(data) => Ok(ProtobufConvert::Struct(
                ProtobufConvertStruct::from_derive_input(
                    input.ident.clone(),
                    data,
                    input.attrs.as_ref(),
                )?,
            )),
            Data::Enum(data) => Ok(ProtobufConvert::Enum(
                ProtobufConvertEnum::from_derive_input(
                    input.ident.clone(),
                    data,
                    input.attrs.as_ref(),
                )?,
            )),
            _ => Err(darling::Error::unsupported_shape(
                "Only for enums and structs.",
            )),
        }
    }
}

impl ProtobufConvert {
    fn name(&self) -> &Ident {
        match self {
            ProtobufConvert::Enum(inner) => &inner.name,
            ProtobufConvert::Struct(inner) => &inner.name,
        }
    }

    fn cr(&self) -> &CratePath {
        match self {
            ProtobufConvert::Enum(inner) => &inner.attrs.cr,
            ProtobufConvert::Struct(inner) => &inner.attrs.cr,
        }
    }

    fn serde_needed(&self) -> bool {
        match self {
            ProtobufConvert::Enum(inner) => inner.attrs.serde_pb_convert,
            ProtobufConvert::Struct(inner) => inner.attrs.serde_pb_convert,
        }
    }

    fn implement_merkledb_traits(&self) -> impl ToTokens {
        let name = self.name();
        let cr = self.cr();

        quote! {
            impl exonum_merkledb::ObjectHash for #name {
                fn object_hash(&self) -> #cr::crypto::Hash {
                    use exonum_merkledb::BinaryValue;
                    let v = self.to_bytes();
                    #cr::crypto::hash(&v)
                }
            }

            // This trait assumes that we work with trusted data so we can unwrap here.
            impl exonum_merkledb::BinaryValue for #name {
                fn to_bytes(&self) -> Vec<u8> {
                    self.to_pb().write_to_bytes().expect(
                        concat!("Failed to serialize in BinaryValue for ", stringify!(#name))
                    )
                }

                fn from_bytes(value: std::borrow::Cow<[u8]>) -> Result<Self, failure::Error> {
                    let mut block = <Self as ProtobufConvert>::ProtoStruct::new();
                    block.merge_from_bytes(value.as_ref())?;
                    ProtobufConvert::from_pb(block)
                }
            }
        }
    }

    fn implement_serde_protobuf_convert(&self) -> impl ToTokens {
        let name = self.name();
        quote! {
            impl serde::Serialize for #name {
                fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
                where
                    S: serde::Serializer,
                {
                    self.to_pb().serialize(serializer)
                }
            }

            impl<'de> serde::Deserialize<'de> for #name {
                fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
                where
                    D: serde::Deserializer<'de>,
                {
                    let pb = <#name as ProtobufConvert>::ProtoStruct::deserialize(deserializer)?;
                    ProtobufConvert::from_pb(pb).map_err(serde::de::Error::custom)
                }
            }
        }
    }

    fn implement_protobuf_convert(&self) -> impl ToTokens {
        match self {
            ProtobufConvert::Enum(data) => quote! { #data },
            ProtobufConvert::Struct(data) => quote! { #data },
        }
    }

    fn implement_any_conversion(&self) -> impl ToTokens {
        let name = self.name();
        let cr = self.cr();

        quote! {
            impl From<#name> for #cr::proto::Any {
                fn from(v: #name) -> Self {
                    Self::new(v)
                }
            }

            impl std::convert::TryFrom<#cr::proto::Any> for #name {
                type Error = failure::Error;

                fn try_from(v: #cr::proto::Any) -> Result<Self, Self::Error> {
                    v.try_into()
                }
            }
        }
    }
}

impl ToTokens for ProtobufConvert {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        let mod_name = Ident::new(
            &format!("pb_convert_impl_{}", self.name()),
            Span::call_site(),
        );
        let cr = self.cr();

        let protobuf_convert = self.implement_protobuf_convert();
        let merkledb_traits = self.implement_merkledb_traits();
        let any_conversion = self.implement_any_conversion();
        let serde_traits = if self.serde_needed() {
            let serde = self.implement_serde_protobuf_convert();
            quote! { #serde }
        } else {
            quote! {}
        };

        let expanded = quote! {
            mod #mod_name {
                use super::*;

                use protobuf::Message as _ProtobufMessage;
                use #cr::proto::ProtobufConvert;

                #protobuf_convert
                #merkledb_traits
                #serde_traits
                #any_conversion
            }
        };
        tokens.extend(expanded)
    }
}

pub fn implement_protobuf_convert(input: TokenStream) -> TokenStream {
    let input = ProtobufConvert::from_derive_input(&syn::parse(input).unwrap())
        .unwrap_or_else(|e| panic!("ProtobufConvert: {}", e));
    let tokens = quote! {#input};
    tokens.into()
}
<commit_msg>Improve ProtobufConvert for enum variants (#1465)<commit_after>// Copyright 2019 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use darling::{FromDeriveInput, FromMeta};
use heck::SnakeCase;
use proc_macro::TokenStream;
use proc_macro2::{Ident, Span};
use quote::{quote, ToTokens};
use syn::{Attribute, Data, DataEnum, DataStruct, DeriveInput, Fields, Path, Type, Variant};

use std::convert::TryFrom;

use super::{find_exonum_meta, CratePath};

#[derive(Debug, FromMeta)]
#[darling(default)]
struct ProtobufConvertStructAttrs {
    #[darling(rename = "crate")]
    cr: CratePath,
    pb: Option<Path>,
    serde_pb_convert: bool,
}

impl Default for ProtobufConvertStructAttrs {
    fn default() -> Self {
        Self {
            cr: CratePath::default(),
            pb: None,
            serde_pb_convert: false,
        }
    }
}

impl TryFrom<&[Attribute]> for ProtobufConvertStructAttrs {
    type Error = darling::Error;

    fn try_from(args: &[Attribute]) -> Result<Self, Self::Error> {
        find_exonum_meta(args)
            .map(|meta| Self::from_nested_meta(&meta))
            .unwrap_or_else(|| Ok(Self::default()))
    }
}

#[derive(Debug, FromMeta)]
#[darling(default)]
struct ProtobufConvertEnumAttrs {
    #[darling(rename = "crate")]
    cr: CratePath,
    pb: Option<Path>,
    serde_pb_convert: bool,
    oneof_field: Ident,
}

impl Default for ProtobufConvertEnumAttrs {
    fn default() -> Self {
        Self {
            cr: CratePath::default(),
            pb: None,
            oneof_field: syn::parse_str("message").unwrap(),
            serde_pb_convert: false,
        }
    }
}

impl TryFrom<&[Attribute]> for ProtobufConvertEnumAttrs {
    type Error = darling::Error;

    fn try_from(args: &[Attribute]) -> Result<Self, Self::Error> {
        find_exonum_meta(args)
            .map(|meta| Self::from_nested_meta(&meta))
            .unwrap_or_else(|| Ok(Self::default()))
    }
}

#[derive(Debug)]
struct ProtobufConvertStruct {
    name: Ident,
    fields: Vec<Ident>,
    attrs: ProtobufConvertStructAttrs,
}

impl ProtobufConvertStruct {
    fn from_derive_input(
        name: Ident,
        data: &DataStruct,
        attrs: &[Attribute],
    ) -> Result<Self, darling::Error> {
        let attrs = ProtobufConvertStructAttrs::try_from(attrs)?;
        let fields = data
            .fields
            .iter()
            .map(|f| f.ident.clone().unwrap())
            .collect();

        Ok(Self {
            name,
            attrs,
            fields,
        })
    }
}

impl ToTokens for ProtobufConvertStruct {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        let name = &self.name;
        let pb_name = &self.attrs.pb;
        let from_pb_impl = {
            let getters = self
                .fields
                .iter()
                .map(|i| Ident::new(&format!("get_{}", i), Span::call_site()));
            let fields = self.fields.to_vec();

            quote! {
                let inner = Self {
                    #( #fields: ProtobufConvert::from_pb(pb.#getters().to_owned())?, )*
                };
                Ok(inner)
            }
        };
        let to_pb_impl = {
            let setters = self
                .fields
                .iter()
                .map(|i| Ident::new(&format!("set_{}", i), Span::call_site()));
            let fields = self.fields.to_vec();

            quote! {
                let mut msg = Self::ProtoStruct::default();
                #( msg.#setters(ProtobufConvert::to_pb(&self.#fields).into()); )*
                msg
            }
        };

        let expanded = quote! {
            impl ProtobufConvert for #name {
                type ProtoStruct = #pb_name;

                fn from_pb(pb: Self::ProtoStruct) -> std::result::Result<Self, failure::Error> {
                    #from_pb_impl
                }

                fn to_pb(&self) -> Self::ProtoStruct {
                    #to_pb_impl
                }
            }
        };
        tokens.extend(expanded);
    }
}

#[derive(Debug)]
struct ParsedVariant {
    name: Ident,
    field_name: Path,
}

impl TryFrom<&Variant> for ParsedVariant {
    type Error = darling::Error;

    fn try_from(value: &Variant) -> Result<Self, Self::Error> {
        let name = value.ident.clone();
        let field_name = match &value.fields {
            Fields::Unnamed(fields) => {
                if fields.unnamed.len() != 1 {
                    return Err(darling::Error::unsupported_shape(
                        "Too many fields in the enum variant",
                    ));
                }

                match &fields.unnamed.first().unwrap().ty {
                    Type::Path(type_path) => Ok(type_path.path.clone()),
                    _ => Err(darling::Error::unsupported_shape(
                        "Only variants in form Foo(Bar) are supported.",
                    )),
                }
            }
            _ => Err(darling::Error::unsupported_shape(
                "Only variants in form Foo(Bar) are supported.",
            )),
        }?;

        Ok(Self { name, field_name })
    }
}

#[derive(Debug)]
struct ProtobufConvertEnum {
    name: Ident,
    variants: Vec<ParsedVariant>,
    attrs: ProtobufConvertEnumAttrs,
}

impl ProtobufConvertEnum {
    fn from_derive_input(
        name: Ident,
        data: &DataEnum,
        attrs: &[Attribute],
    ) -> Result<Self, darling::Error> {
        let attrs = ProtobufConvertEnumAttrs::try_from(attrs)?;
        let variants = data
            .variants
            .iter()
            .map(ParsedVariant::try_from)
            .collect::<Result<Vec<_>, _>>()?;

        Ok(Self {
            name,
            attrs,
            variants,
        })
    }

    fn impl_protobuf_convert(&self) -> impl ToTokens {
        let pb_oneof_enum = {
            let mut pb = self.attrs.pb.clone().unwrap();
            let oneof = pb.segments.pop().unwrap().value().ident.clone();
            let oneof_enum = Ident::new(
                &format!("{}_oneof_{}", oneof, &self.attrs.oneof_field),
                Span::call_site(),
            );
            quote! { #pb #oneof_enum }
        };
        let name = &self.name;
        let pb_name = &self.attrs.pb;
        let oneof = &self.attrs.oneof_field;

        let from_pb_impl = {
            let match_arms = self.variants.iter().map(|variant| {
                let pb_variant = Ident::new(
                    variant.name.to_string().to_snake_case().as_ref(),
                    Span::call_site(),
                );
                let variant_name = &variant.name;
                let field_name = &variant.field_name;

                quote! {
                    Some(#pb_oneof_enum::#pb_variant(pb)) => {
                        #field_name::from_pb(pb).map(#name::#variant_name)
                    }
                }
            });

            quote! {
                match pb.#oneof {
                    #( #match_arms )*
                    None => Err(failure::format_err!("Failed to decode #name from protobuf"))
                }
            }
        };
        let to_pb_impl = {
            let match_arms = self.variants.iter().map(|variant| {
                let pb_variant = variant.name.to_string().to_snake_case();
                let variant_name = &variant.name;

                let setter = Ident::new(&format!("set_{}", pb_variant), Span::call_site());
                quote! {
                    #name::#variant_name(msg) => inner.#setter(msg.to_pb()),
                }
            });

            quote! {
                let mut inner = Self::ProtoStruct::new();
                match self {
                    #( #match_arms )*
                }
                inner
            }
        };

        quote! {
            impl ProtobufConvert for #name {
                type ProtoStruct = #pb_name;

                fn from_pb(mut pb: Self::ProtoStruct) -> std::result::Result<Self, failure::Error> {
                    #from_pb_impl
                }

                fn to_pb(&self) -> Self::ProtoStruct {
                    #to_pb_impl
                }
            }
        }
    }

    fn impl_enum_conversions(&self) -> impl ToTokens {
        let name = &self.name;
        let conversions = self.variants.iter().map(|variant| {
            let variant_name = &variant.name;
            let field_name = &variant.field_name;

            quote! {
                impl From<#field_name> for #name {
                    fn from(variant: #field_name) -> Self {
                        #name::#variant_name(variant)
                    }
                }

                impl std::convert::TryFrom<#name> for #field_name {
                    type Error = failure::Error;

                    fn try_from(msg: #name) -> Result<Self, Self::Error> {
                        if let #name::#variant_name(inner) = msg {
                            Ok(inner)
                        } else {
                            Err(failure::format_err!(
                                "Expected #variant_name variant, but got {:?}",
                                msg
                            ))
                        }
                    }
                }
            }
        });

        quote! {
            #( #conversions )*
        }
    }
}

impl ToTokens for ProtobufConvertEnum {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        let pb_convert = self.impl_protobuf_convert();
        let conversions = self.impl_enum_conversions();

        let expanded = quote! {
            #pb_convert
            #conversions
        };
        tokens.extend(expanded)
    }
}

#[derive(Debug)]
enum ProtobufConvert {
    Enum(ProtobufConvertEnum),
    Struct(ProtobufConvertStruct),
}

impl FromDeriveInput for ProtobufConvert {
    fn from_derive_input(input: &DeriveInput) -> Result<Self, darling::Error> {
        match &input.data {
            Data::Struct(data) => Ok(ProtobufConvert::Struct(
                ProtobufConvertStruct::from_derive_input(
                    input.ident.clone(),
                    data,
                    input.attrs.as_ref(),
                )?,
            )),
            Data::Enum(data) => Ok(ProtobufConvert::Enum(
                ProtobufConvertEnum::from_derive_input(
                    input.ident.clone(),
                    data,
                    input.attrs.as_ref(),
                )?,
            )),
            _ => Err(darling::Error::unsupported_shape(
                "Only for enums and structs.",
            )),
        }
    }
}

impl ProtobufConvert {
    fn name(&self) -> &Ident {
        match self {
            ProtobufConvert::Enum(inner) => &inner.name,
            ProtobufConvert::Struct(inner) => &inner.name,
        }
    }

    fn cr(&self) -> &CratePath {
        match self {
            ProtobufConvert::Enum(inner) => &inner.attrs.cr,
            ProtobufConvert::Struct(inner) => &inner.attrs.cr,
        }
    }

    fn serde_needed(&self) -> bool {
        match self {
            ProtobufConvert::Enum(inner) => inner.attrs.serde_pb_convert,
            ProtobufConvert::Struct(inner) => inner.attrs.serde_pb_convert,
        }
    }

    fn implement_merkledb_traits(&self) -> impl ToTokens {
        let name = self.name();
        let cr = self.cr();

        quote! {
            impl exonum_merkledb::ObjectHash for #name {
                fn object_hash(&self) -> #cr::crypto::Hash {
                    use exonum_merkledb::BinaryValue;
                    let v = self.to_bytes();
                    #cr::crypto::hash(&v)
                }
            }

            // This trait assumes that we work with trusted data so we can unwrap here.
            impl exonum_merkledb::BinaryValue for #name {
                fn to_bytes(&self) -> Vec<u8> {
                    self.to_pb().write_to_bytes().expect(
                        concat!("Failed to serialize in BinaryValue for ", stringify!(#name))
                    )
                }

                fn from_bytes(value: std::borrow::Cow<[u8]>) -> Result<Self, failure::Error> {
                    let mut block = <Self as ProtobufConvert>::ProtoStruct::new();
                    block.merge_from_bytes(value.as_ref())?;
                    ProtobufConvert::from_pb(block)
                }
            }
        }
    }

    fn implement_serde_protobuf_convert(&self) -> impl ToTokens {
        let name = self.name();
        quote! {
            impl serde::Serialize for #name {
                fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
                where
                    S: serde::Serializer,
                {
                    self.to_pb().serialize(serializer)
                }
            }

            impl<'de> serde::Deserialize<'de> for #name {
                fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
                where
                    D: serde::Deserializer<'de>,
                {
                    let pb = <#name as ProtobufConvert>::ProtoStruct::deserialize(deserializer)?;
                    ProtobufConvert::from_pb(pb).map_err(serde::de::Error::custom)
                }
            }
        }
    }

    fn implement_protobuf_convert(&self) -> impl ToTokens {
        match self {
            ProtobufConvert::Enum(data) => quote! { #data },
            ProtobufConvert::Struct(data) => quote! { #data },
        }
    }

    fn implement_any_conversion(&self) -> impl ToTokens {
        let name = self.name();
        let cr = self.cr();

        quote! {
            impl From<#name> for #cr::proto::Any {
                fn from(v: #name) -> Self {
                    Self::new(v)
                }
            }

            impl std::convert::TryFrom<#cr::proto::Any> for #name {
                type Error = failure::Error;

                fn try_from(v: #cr::proto::Any) -> Result<Self, Self::Error> {
                    v.try_into()
                }
            }
        }
    }
}

impl ToTokens for ProtobufConvert {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        let mod_name = Ident::new(
            &format!("pb_convert_impl_{}", self.name()),
            Span::call_site(),
        );
        let cr = self.cr();

        let protobuf_convert = self.implement_protobuf_convert();
        let merkledb_traits = self.implement_merkledb_traits();
        let any_conversion = self.implement_any_conversion();
        let serde_traits = if self.serde_needed() {
            let serde = self.implement_serde_protobuf_convert();
            quote! { #serde }
        } else {
            quote! {}
        };

        let expanded = quote! {
            mod #mod_name {
                use super::*;

                use protobuf::Message as _ProtobufMessage;
                use #cr::proto::ProtobufConvert;

                #protobuf_convert
                #merkledb_traits
                #serde_traits
                #any_conversion
            }
        };
        tokens.extend(expanded)
    }
}

pub fn implement_protobuf_convert(input: TokenStream) -> TokenStream {
    let input = ProtobufConvert::from_derive_input(&syn::parse(input).unwrap())
        .unwrap_or_else(|e| panic!("ProtobufConvert: {}", e));
    let tokens = quote! {#input};
    tokens.into()
}
<|endoftext|>