<commit_before>//
// Sysinfo
//
// Copyright (c) 2015 Guillaume Gomez
//

use sys::component::Component;
use sys::processor::*;
use sys::process::*;
use std::fs::{File, read_link};
use std::io::Read;
use std::str::FromStr;
use std::path::{Path, PathBuf};
use std::collections::HashMap;
use std::fs;
use libc::{stat, lstat, c_char, sysconf, _SC_CLK_TCK, _SC_PAGESIZE, S_IFLNK, S_IFMT};

pub struct System {
    process_list: HashMap<usize, Process>,
    mem_total: u64,
    mem_free: u64,
    swap_total: u64,
    swap_free: u64,
    processors: Vec<Processor>,
    page_size_kb: u64,
    temperatures: Vec<Component>,
}

impl System {
    pub fn new() -> System {
        let mut s = System {
            process_list: HashMap::new(),
            mem_total: 0,
            mem_free: 0,
            swap_total: 0,
            swap_free: 0,
            processors: Vec::new(),
            page_size_kb: unsafe { sysconf(_SC_PAGESIZE) as u64 / 1024 },
            temperatures: Component::get_components(),
        };
        s.refresh_all();
        s
    }

    pub fn refresh_system(&mut self) {
        let data = get_all_data("/proc/meminfo");
        let lines: Vec<&str> = data.split('\n').collect();

        for component in self.temperatures.iter_mut() {
            component.update();
        }
        for line in lines.iter() {
            match *line {
                l if l.starts_with("MemTotal:") => {
                    let parts: Vec<&str> = line.split(' ').collect();

                    self.mem_total = u64::from_str(parts[parts.len() - 2]).unwrap();
                },
                l if l.starts_with("MemAvailable:") => {
                    let parts: Vec<&str> = line.split(' ').collect();

                    self.mem_free = u64::from_str(parts[parts.len() - 2]).unwrap();
                },
                l if l.starts_with("SwapTotal:") => {
                    let parts: Vec<&str> = line.split(' ').collect();

                    self.swap_total = u64::from_str(parts[parts.len() - 2]).unwrap();
                },
                l if l.starts_with("SwapFree:") => {
                    let parts: Vec<&str> = line.split(' ').collect();

                    self.swap_free = u64::from_str(parts[parts.len() - 2]).unwrap();
                },
                _ => continue,
            }
        }
        let data = get_all_data("/proc/stat");
        let lines: Vec<&str> = data.split('\n').collect();
        let mut i = 0;
        let first = self.processors.len() == 0;
        for line in lines.iter() {
            if !line.starts_with("cpu") {
                break;
            }

            let (parts, _): (Vec<&str>, Vec<&str>) = line.split(' ').partition(|s| s.len() > 0);
            if first {
                self.processors.push(new_processor(parts[0], u64::from_str(parts[1]).unwrap(),
                    u64::from_str(parts[2]).unwrap(),
                    u64::from_str(parts[3]).unwrap(),
                    u64::from_str(parts[4]).unwrap(),
                    u64::from_str(parts[5]).unwrap(),
                    u64::from_str(parts[6]).unwrap(),
                    u64::from_str(parts[7]).unwrap(),
                    u64::from_str(parts[8]).unwrap(),
                    u64::from_str(parts[9]).unwrap(),
                    u64::from_str(parts[10]).unwrap()));
            } else {
                set_processor(self.processors.get_mut(i).unwrap(),
                    u64::from_str(parts[1]).unwrap(),
                    u64::from_str(parts[2]).unwrap(),
                    u64::from_str(parts[3]).unwrap(),
                    u64::from_str(parts[4]).unwrap(),
                    u64::from_str(parts[5]).unwrap(),
                    u64::from_str(parts[6]).unwrap(),
                    u64::from_str(parts[7]).unwrap(),
                    u64::from_str(parts[8]).unwrap(),
                    u64::from_str(parts[9]).unwrap(),
                    u64::from_str(parts[10]).unwrap());
                i += 1;
            }
        }
    }

    pub fn refresh_process(&mut self) {
        match fs::read_dir(&Path::new("/proc")) {
            Ok(d) => {
                for entry in d {
                    if !entry.is_ok() {
                        continue;
                    }
                    let entry = entry.unwrap();
                    let entry = entry.path();

                    if entry.is_dir() {
                        _get_process_data(entry.as_path(), &mut self.process_list, self.page_size_kb);
                    }
                }
                self.clear_procs();
            }
            Err(_) => {}
        }
    }

    fn clear_procs(&mut self) {
        if self.processors.len() > 0 {
            let (new, old) = get_raw_times(&self.processors[0]);
            let total_time = (new - old) as f32;
            let mut to_delete = Vec::new();
            let nb_processors = self.processors.len() as u64 - 1;

            for (pid, proc_) in self.process_list.iter_mut() {
                if has_been_updated(&proc_) == false {
                    to_delete.push(*pid);
                } else {
                    compute_cpu_usage(proc_, nb_processors, total_time);
                }
            }
            for pid in to_delete {
                self.process_list.remove(&pid);
            }
        }
    }

    // COMMON PART
    //
    // Need to be moved into a "common" file to avoid duplication.

    pub fn refresh_all(&mut self) {
        self.refresh_system();
        self.refresh_process();
    }

    pub fn get_process_list<'a>(&'a self) -> &'a HashMap<usize, Process> {
        &self.process_list
    }

    /// Return the process corresponding to the given pid or None if no such process exists.
    pub fn get_process(&self, pid: i64) -> Option<&Process> {
        self.process_list.get(&(pid as usize))
    }

    /// Return a list of process starting with the given name.
    pub fn get_process_by_name(&self, name: &str) -> Vec<&Process> {
        let mut ret = vec!();
        for val in self.process_list.values() {
            if val.name.starts_with(name) {
                ret.push(val);
            }
        }
        ret
    }

    /// The first process in the array is the "main" process
    pub fn get_processor_list<'a>(&'a self) -> &'a [Processor] {
        &self.processors[..]
    }

    pub fn get_total_memory(&self) -> u64 {
        self.mem_total
    }

    pub fn get_free_memory(&self) -> u64 {
        self.mem_free
    }

    pub fn get_used_memory(&self) -> u64 {
        self.mem_total - self.mem_free
    }

    pub fn get_total_swap(&self) -> u64 {
        self.swap_total
    }

    pub fn get_free_swap(&self) -> u64 {
        self.swap_free
    }

    // need to be checked
    pub fn get_used_swap(&self) -> u64 {
        self.swap_total - self.swap_free
    }

    pub fn get_components_list<'a>(&'a self) -> &'a [Component] {
        &self.temperatures[..]
    }
}

fn get_all_data(file_path: &str) -> String {
    let mut file = File::open(file_path).unwrap();
    let mut data = String::new();

    file.read_to_string(&mut data).unwrap();
    data
}

fn update_time_and_memory(entry: &mut Process, parts: &[&str], page_size_kb: u64) {
    //entry.name = parts[1][1..].to_owned();
    //entry.name.pop();
    // we get the rss then we add the vsize
    entry.memory = u64::from_str(parts[23]).unwrap() * page_size_kb +
                   u64::from_str(parts[22]).unwrap() / 1024;
    set_time(entry,
             u64::from_str(parts[13]).unwrap(),
             u64::from_str(parts[14]).unwrap());
}

fn _get_process_data(path: &Path, proc_list: &mut HashMap<usize, Process>, page_size_kb: u64) {
    if !path.exists() || !path.is_dir() {
        return;
    }
    let paths : Vec<&str> = path.as_os_str().to_str().unwrap().split("/").collect();
    let last = paths[paths.len() - 1];
    match i64::from_str(last) {
        Ok(nb) => {
            let mut tmp = PathBuf::from(path);

            tmp.push("stat");
            let data = get_all_data(tmp.to_str().unwrap());

            // The stat file is "interesting" to parse, because spaces cannot
            // be used as delimiters. The second field stores the command name
            // sourrounded by parentheses. Unfortunately, whitespace and
            // parentheses are legal parts of the command, so parsing has to
            // proceed like this: The first field is delimited by the first
            // whitespace, the second field is everything until the last ')'
            // in the entire string. All other fields are delimited by
            // whitespace.

            let mut parts = Vec::new();
            let mut data_it = data.splitn(2, " ");
            parts.push(data_it.next().unwrap());
            // The following loses the ) from the input, but that's ok because
            // we're not using it anyway.
            let mut data_it = data_it.next().unwrap().rsplitn(2, ")");
            let data = data_it.next().unwrap();
            parts.push(data_it.next().unwrap());
            parts.extend(data.split_whitespace());
            if let Some(ref mut entry) = proc_list.get_mut(&(nb as usize)) {
                update_time_and_memory(entry, &parts, page_size_kb);
                return;
            }

            let parent = match i64::from_str(parts[3]).unwrap() {
                0 => None,
                p => Some(p)
            };

            let mut p = Process::new(nb,
                                     parent,
                                     u64::from_str(parts[21]).unwrap() /
                                     unsafe { sysconf(_SC_CLK_TCK) } as u64);

            tmp = PathBuf::from(path);
            tmp.push("status");
            let status_data = get_all_data(tmp.to_str().unwrap());

            // We're only interested in the lines starting with Uid: and Gid:
            // here. From these lines, we're looking at the second entry to get
            // the effective u/gid.

            let f = |h: &str, n: &str| -> Option<i64> {
                if h.starts_with(n) {
                    h.split_whitespace().nth(3).unwrap().parse().ok()
                } else {
                    None
                }
            };
            let mut set_uid = false;
            let mut set_gid = false;
            for line in status_data.lines() {
                if let Some(u) = f(line, "Uid:") {
                    assert!(!set_uid);
                    set_uid = true;
                    p.uid = u;
                }
                if let Some(g) = f(line, "Gid:") {
                    assert!(!set_gid);
                    set_gid = true;
                    p.gid = g;
                }
            }
            assert!(set_uid && set_gid);

            tmp = PathBuf::from(path);
            tmp.push("cmdline");
            p.cmd = copy_from_file(&tmp);
            p.name = p.cmd[0].split("/").last().unwrap().to_owned();
            tmp = PathBuf::from(path);
            tmp.push("environ");
            p.environ = copy_from_file(&tmp);
            tmp = PathBuf::from(path);
            tmp.push("exe");

            let s = read_link(tmp.to_str().unwrap());

            if s.is_ok() {
                p.exe = s.unwrap().to_str().unwrap().to_owned();
            }
            tmp = PathBuf::from(path);
            tmp.push("cwd");
            p.cwd = realpath(Path::new(tmp.to_str().unwrap())).to_str().unwrap().to_owned();
            tmp = PathBuf::from(path);
            tmp.push("root");
            p.root = realpath(Path::new(tmp.to_str().unwrap())).to_str().unwrap().to_owned();

            update_time_and_memory(&mut p, &parts, page_size_kb);
            proc_list.insert(nb as usize, p);
        }
        _ => {}
    }
}

#[allow(unused_must_use)]
fn copy_from_file(entry: &Path) -> Vec<String> {
    match File::open(entry.to_str().unwrap()) {
        Ok(mut f) => {
            let mut d = String::new();

            f.read_to_string(&mut d);
            d.split('\0').map(|x| x.to_owned()).collect()
        },
        Err(_) => Vec::new()
    }
}

fn realpath(original: &Path) -> PathBuf {
    let ori = Path::new(original.to_str().unwrap());

    // Right now lstat on windows doesn't work quite well
    if cfg!(windows) {
        return PathBuf::from(ori);
    }
    let result = PathBuf::from(original);
    let mut buf: stat = unsafe { ::std::mem::uninitialized() };
    let res = unsafe { lstat(result.to_str().unwrap().as_ptr() as *const c_char, &mut buf as *mut stat) };
    if res < 0 || (buf.st_mode & S_IFMT) != S_IFLNK {
        PathBuf::new()
    } else {
        match fs::read_link(&result) {
            Ok(f) => f,
        Err(_) => PathBuf::new(),
        }
    }
}
<commit_msg>Report the effective uid field, not saved<commit_after>//
// Sysinfo
//
// Copyright (c) 2015 Guillaume Gomez
//

use sys::component::Component;
use sys::processor::*;
use sys::process::*;
use std::fs::{File, read_link};
use std::io::Read;
use std::str::FromStr;
use std::path::{Path, PathBuf};
use std::collections::HashMap;
use std::fs;
use libc::{stat, lstat, c_char, sysconf, _SC_CLK_TCK, _SC_PAGESIZE, S_IFLNK, S_IFMT};

pub struct System {
    process_list: HashMap<usize, Process>,
    mem_total: u64,
    mem_free: u64,
    swap_total: u64,
    swap_free: u64,
    processors: Vec<Processor>,
    page_size_kb: u64,
    temperatures: Vec<Component>,
}

impl System {
    pub fn new() -> System {
        let mut s = System {
            process_list: HashMap::new(),
            mem_total: 0,
            mem_free: 0,
            swap_total: 0,
            swap_free: 0,
            processors: Vec::new(),
            page_size_kb: unsafe { sysconf(_SC_PAGESIZE) as u64 / 1024 },
            temperatures: Component::get_components(),
        };
        s.refresh_all();
        s
    }

    pub fn refresh_system(&mut self) {
        let data = get_all_data("/proc/meminfo");
        let lines: Vec<&str> = data.split('\n').collect();

        for component in self.temperatures.iter_mut() {
            component.update();
        }
        for line in lines.iter() {
            match *line {
                l if l.starts_with("MemTotal:") => {
                    let parts: Vec<&str> = line.split(' ').collect();

                    self.mem_total = u64::from_str(parts[parts.len() - 2]).unwrap();
                },
                l if l.starts_with("MemAvailable:") => {
                    let parts: Vec<&str> = line.split(' ').collect();

                    self.mem_free = u64::from_str(parts[parts.len() - 2]).unwrap();
                },
                l if l.starts_with("SwapTotal:") => {
                    let parts: Vec<&str> = line.split(' ').collect();

                    self.swap_total = u64::from_str(parts[parts.len() - 2]).unwrap();
                },
                l if l.starts_with("SwapFree:") => {
                    let parts: Vec<&str> = line.split(' ').collect();

                    self.swap_free = u64::from_str(parts[parts.len() - 2]).unwrap();
                },
                _ => continue,
            }
        }
        let data = get_all_data("/proc/stat");
        let lines: Vec<&str> = data.split('\n').collect();
        let mut i = 0;
        let first = self.processors.len() == 0;
        for line in lines.iter() {
            if !line.starts_with("cpu") {
                break;
            }

            let (parts, _): (Vec<&str>, Vec<&str>) = line.split(' ').partition(|s| s.len() > 0);
            if first {
                self.processors.push(new_processor(parts[0], u64::from_str(parts[1]).unwrap(),
                    u64::from_str(parts[2]).unwrap(),
                    u64::from_str(parts[3]).unwrap(),
                    u64::from_str(parts[4]).unwrap(),
                    u64::from_str(parts[5]).unwrap(),
                    u64::from_str(parts[6]).unwrap(),
                    u64::from_str(parts[7]).unwrap(),
                    u64::from_str(parts[8]).unwrap(),
                    u64::from_str(parts[9]).unwrap(),
                    u64::from_str(parts[10]).unwrap()));
            } else {
                set_processor(self.processors.get_mut(i).unwrap(),
                    u64::from_str(parts[1]).unwrap(),
                    u64::from_str(parts[2]).unwrap(),
                    u64::from_str(parts[3]).unwrap(),
                    u64::from_str(parts[4]).unwrap(),
                    u64::from_str(parts[5]).unwrap(),
                    u64::from_str(parts[6]).unwrap(),
                    u64::from_str(parts[7]).unwrap(),
                    u64::from_str(parts[8]).unwrap(),
                    u64::from_str(parts[9]).unwrap(),
                    u64::from_str(parts[10]).unwrap());
                i += 1;
            }
        }
    }

    pub fn refresh_process(&mut self) {
        match fs::read_dir(&Path::new("/proc")) {
            Ok(d) => {
                for entry in d {
                    if !entry.is_ok() {
                        continue;
                    }
                    let entry = entry.unwrap();
                    let entry = entry.path();

                    if entry.is_dir() {
                        _get_process_data(entry.as_path(), &mut self.process_list, self.page_size_kb);
                    }
                }
                self.clear_procs();
            }
            Err(_) => {}
        }
    }

    fn clear_procs(&mut self) {
        if self.processors.len() > 0 {
            let (new, old) = get_raw_times(&self.processors[0]);
            let total_time = (new - old) as f32;
            let mut to_delete = Vec::new();
            let nb_processors = self.processors.len() as u64 - 1;

            for (pid, proc_) in self.process_list.iter_mut() {
                if has_been_updated(&proc_) == false {
                    to_delete.push(*pid);
                } else {
                    compute_cpu_usage(proc_, nb_processors, total_time);
                }
            }
            for pid in to_delete {
                self.process_list.remove(&pid);
            }
        }
    }

    // COMMON PART
    //
    // Need to be moved into a "common" file to avoid duplication.

    pub fn refresh_all(&mut self) {
        self.refresh_system();
        self.refresh_process();
    }

    pub fn get_process_list<'a>(&'a self) -> &'a HashMap<usize, Process> {
        &self.process_list
    }

    /// Return the process corresponding to the given pid or None if no such process exists.
    pub fn get_process(&self, pid: i64) -> Option<&Process> {
        self.process_list.get(&(pid as usize))
    }

    /// Return a list of process starting with the given name.
    pub fn get_process_by_name(&self, name: &str) -> Vec<&Process> {
        let mut ret = vec!();
        for val in self.process_list.values() {
            if val.name.starts_with(name) {
                ret.push(val);
            }
        }
        ret
    }

    /// The first process in the array is the "main" process
    pub fn get_processor_list<'a>(&'a self) -> &'a [Processor] {
        &self.processors[..]
    }

    pub fn get_total_memory(&self) -> u64 {
        self.mem_total
    }

    pub fn get_free_memory(&self) -> u64 {
        self.mem_free
    }

    pub fn get_used_memory(&self) -> u64 {
        self.mem_total - self.mem_free
    }

    pub fn get_total_swap(&self) -> u64 {
        self.swap_total
    }

    pub fn get_free_swap(&self) -> u64 {
        self.swap_free
    }

    // need to be checked
    pub fn get_used_swap(&self) -> u64 {
        self.swap_total - self.swap_free
    }

    pub fn get_components_list<'a>(&'a self) -> &'a [Component] {
        &self.temperatures[..]
    }
}

fn get_all_data(file_path: &str) -> String {
    let mut file = File::open(file_path).unwrap();
    let mut data = String::new();

    file.read_to_string(&mut data).unwrap();
    data
}

fn update_time_and_memory(entry: &mut Process, parts: &[&str], page_size_kb: u64) {
    //entry.name = parts[1][1..].to_owned();
    //entry.name.pop();
    // we get the rss then we add the vsize
    entry.memory = u64::from_str(parts[23]).unwrap() * page_size_kb +
                   u64::from_str(parts[22]).unwrap() / 1024;
    set_time(entry,
             u64::from_str(parts[13]).unwrap(),
             u64::from_str(parts[14]).unwrap());
}

fn _get_process_data(path: &Path, proc_list: &mut HashMap<usize, Process>, page_size_kb: u64) {
    if !path.exists() || !path.is_dir() {
        return;
    }
    let paths : Vec<&str> = path.as_os_str().to_str().unwrap().split("/").collect();
    let last = paths[paths.len() - 1];
    match i64::from_str(last) {
        Ok(nb) => {
            let mut tmp = PathBuf::from(path);

            tmp.push("stat");
            let data = get_all_data(tmp.to_str().unwrap());

            // The stat file is "interesting" to parse, because spaces cannot
            // be used as delimiters. The second field stores the command name
            // sourrounded by parentheses. Unfortunately, whitespace and
            // parentheses are legal parts of the command, so parsing has to
            // proceed like this: The first field is delimited by the first
            // whitespace, the second field is everything until the last ')'
            // in the entire string. All other fields are delimited by
            // whitespace.

            let mut parts = Vec::new();
            let mut data_it = data.splitn(2, " ");
            parts.push(data_it.next().unwrap());
            // The following loses the ) from the input, but that's ok because
            // we're not using it anyway.
            let mut data_it = data_it.next().unwrap().rsplitn(2, ")");
            let data = data_it.next().unwrap();
            parts.push(data_it.next().unwrap());
            parts.extend(data.split_whitespace());
            if let Some(ref mut entry) = proc_list.get_mut(&(nb as usize)) {
                update_time_and_memory(entry, &parts, page_size_kb);
                return;
            }

            let parent = match i64::from_str(parts[3]).unwrap() {
                0 => None,
                p => Some(p)
            };

            let mut p = Process::new(nb,
                                     parent,
                                     u64::from_str(parts[21]).unwrap() /
                                     unsafe { sysconf(_SC_CLK_TCK) } as u64);

            tmp = PathBuf::from(path);
            tmp.push("status");
            let status_data = get_all_data(tmp.to_str().unwrap());

            // We're only interested in the lines starting with Uid: and Gid:
            // here. From these lines, we're looking at the second entry to get
            // the effective u/gid.

            let f = |h: &str, n: &str| -> Option<i64> {
                if h.starts_with(n) {
                    h.split_whitespace().nth(2).unwrap().parse().ok()
                } else {
                    None
                }
            };
            let mut set_uid = false;
            let mut set_gid = false;
            for line in status_data.lines() {
                if let Some(u) = f(line, "Uid:") {
                    assert!(!set_uid);
                    set_uid = true;
                    p.uid = u;
                }
                if let Some(g) = f(line, "Gid:") {
                    assert!(!set_gid);
                    set_gid = true;
                    p.gid = g;
                }
            }
            assert!(set_uid && set_gid);

            tmp = PathBuf::from(path);
            tmp.push("cmdline");
            p.cmd = copy_from_file(&tmp);
            p.name = p.cmd[0].split("/").last().unwrap().to_owned();
            tmp = PathBuf::from(path);
            tmp.push("environ");
            p.environ = copy_from_file(&tmp);
            tmp = PathBuf::from(path);
            tmp.push("exe");

            let s = read_link(tmp.to_str().unwrap());

            if s.is_ok() {
                p.exe = s.unwrap().to_str().unwrap().to_owned();
            }
            tmp = PathBuf::from(path);
            tmp.push("cwd");
            p.cwd = realpath(Path::new(tmp.to_str().unwrap())).to_str().unwrap().to_owned();
            tmp = PathBuf::from(path);
            tmp.push("root");
            p.root = realpath(Path::new(tmp.to_str().unwrap())).to_str().unwrap().to_owned();

            update_time_and_memory(&mut p, &parts, page_size_kb);
            proc_list.insert(nb as usize, p);
        }
        _ => {}
    }
}

#[allow(unused_must_use)]
fn copy_from_file(entry: &Path) -> Vec<String> {
    match File::open(entry.to_str().unwrap()) {
        Ok(mut f) => {
            let mut d = String::new();

            f.read_to_string(&mut d);
            d.split('\0').map(|x| x.to_owned()).collect()
        },
        Err(_) => Vec::new()
    }
}

fn realpath(original: &Path) -> PathBuf {
    let ori = Path::new(original.to_str().unwrap());

    // Right now lstat on windows doesn't work quite well
    if cfg!(windows) {
        return PathBuf::from(ori);
    }
    let result = PathBuf::from(original);
    let mut buf: stat = unsafe { ::std::mem::uninitialized() };
    let res = unsafe { lstat(result.to_str().unwrap().as_ptr() as *const c_char, &mut buf as *mut stat) };
    if res < 0 || (buf.st_mode & S_IFMT) != S_IFLNK {
        PathBuf::new()
    } else {
        match fs::read_link(&result) {
            Ok(f) => f,
        Err(_) => PathBuf::new(),
        }
    }
}
<|endoftext|>


<commit_before>use std::ffi::CStr;
use std::marker::PhantomData;

use llvm_sys::prelude::*;
use llvm_sys::core::*;

use owned::{Owned, DropInPlace};

use llvm::{Context, BasicBlock, Value};

pub struct Builder<'cid: 'context, 'context> {
    _context: PhantomData<&'context Context<'cid>>
}

impl<'cid, 'context> DropInPlace for Builder<'cid, 'context> {
    unsafe fn drop_in_place(&mut self) {
        LLVMDisposeBuilder(self.as_raw());
    }
}

impl<'cid, 'context> Builder<'cid, 'context> {
    pub fn new(context: &'context Context<'cid>) -> Owned<Builder<'cid, 'context>> {
        unsafe {
            Owned::from_raw(
                LLVMCreateBuilderInContext(context.as_raw()) as *mut Builder
            )
        }
    }

    pub fn position_at_end<'fid, 'block, 'builder>(&'builder mut self, block: &'block mut BasicBlock<'cid, 'fid>) -> &'builder mut PositionedBuilder<'cid, 'context, 'fid, 'block> {
        unsafe {
            LLVMPositionBuilderAtEnd(self.as_raw(), block.as_raw());
            &mut *(self as *mut Builder as *mut PositionedBuilder)
        }
    }

    pub fn as_raw(&self) -> LLVMBuilderRef {
        self as *const Builder as *mut Builder as LLVMBuilderRef
    }
}

pub struct PositionedBuilder<'cid: 'context, 'context: 'block, 'fid: 'block, 'block> {
    _block: PhantomData<&'block mut BasicBlock<'cid, 'fid>>,
    _builder: PhantomData<Builder<'cid, 'context>>
}

impl<'cid, 'context, 'fid, 'block> PositionedBuilder<'cid, 'context, 'fid, 'block> {
    pub fn br(&mut self, target: &BasicBlock<'cid, 'fid>) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildBr(self.as_raw(), target.as_raw()) as *const Value)
        }
    }

    pub fn or(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildOr(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn add(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildAdd(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn and(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildAnd(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn mul(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildMul(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn neg(&mut self, value: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildNeg(self.as_raw(), value.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn not(&mut self, value: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildNot(self.as_raw(), value.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn ret(&mut self, value: &Value<'cid, 'fid>) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildRet(self.as_raw(), value.as_raw()) as *const Value)
        }
    }

    pub fn as_raw(&self) -> LLVMBuilderRef {
        self as *const PositionedBuilder as *mut PositionedBuilder as LLVMBuilderRef
    }
}
<commit_msg>Add a bunch of new instructions that can be built<commit_after>use std::ffi::CStr;
use std::marker::PhantomData;

use llvm_sys::prelude::*;
use llvm_sys::core::*;

use owned::{Owned, DropInPlace};

use llvm::{Context, BasicBlock, Value, Type};

pub struct Builder<'cid: 'context, 'context> {
    _context: PhantomData<&'context Context<'cid>>
}

impl<'cid, 'context> DropInPlace for Builder<'cid, 'context> {
    unsafe fn drop_in_place(&mut self) {
        LLVMDisposeBuilder(self.as_raw());
    }
}

impl<'cid, 'context> Builder<'cid, 'context> {
    pub fn new(context: &'context Context<'cid>) -> Owned<Builder<'cid, 'context>> {
        unsafe {
            Owned::from_raw(
                LLVMCreateBuilderInContext(context.as_raw()) as *mut Builder
            )
        }
    }

    pub fn position_at_end<'fid, 'block, 'builder>(&'builder mut self, block: &'block mut BasicBlock<'cid, 'fid>) -> &'builder mut PositionedBuilder<'cid, 'context, 'fid, 'block> {
        unsafe {
            LLVMPositionBuilderAtEnd(self.as_raw(), block.as_raw());
            &mut *(self as *mut Builder as *mut PositionedBuilder)
        }
    }

    pub fn as_raw(&self) -> LLVMBuilderRef {
        self as *const Builder as *mut Builder as LLVMBuilderRef
    }
}

pub struct PositionedBuilder<'cid: 'context, 'context: 'block, 'fid: 'block, 'block> {
    _block: PhantomData<&'block mut BasicBlock<'cid, 'fid>>,
    _builder: PhantomData<Builder<'cid, 'context>>
}

impl<'cid, 'context, 'fid, 'block> PositionedBuilder<'cid, 'context, 'fid, 'block> {
    pub fn br(&mut self, target: &BasicBlock<'cid, 'fid>) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildBr(self.as_raw(), target.as_raw()) as *const Value)
        }
    }

    pub fn cond_br(&mut self, cond: &Value<'cid, 'fid>, then_block: &BasicBlock<'cid, 'fid>, else_block: &BasicBlock<'cid, 'fid>) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildCondBr(self.as_raw(), cond.as_raw(), then_block.as_raw(), else_block.as_raw()) as *const Value)
        }
    }

    pub fn add(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildAdd(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn fadd(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildFAdd(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn sub(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildSub(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn fsub(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildFSub(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn mul(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildMul(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn fmul(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildFMul(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn udiv(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildUDiv(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn sdiv(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildSDiv(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn fdiv(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildFDiv(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn urem(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildURem(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn srem(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildSRem(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn frem(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildFRem(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn shl(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildShl(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn lshr(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildLShr(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn ashr(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildAShr(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn and(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildAnd(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn or(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildOr(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn xor(&mut self, lhs: &Value<'cid, 'fid>, rhs: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildXor(self.as_raw(), lhs.as_raw(), rhs.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn neg(&mut self, value: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildNeg(self.as_raw(), value.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn fneg(&mut self, value: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildFNeg(self.as_raw(), value.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn not(&mut self, value: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildNot(self.as_raw(), value.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn alloca(&mut self, ty: &Type<'cid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildAlloca(self.as_raw(), ty.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn array_alloca(&mut self, ty: &Type<'cid>, len: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildArrayAlloca(self.as_raw(), ty.as_raw(), len.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn load(&mut self, ptr: &Value<'cid, 'fid>, name: &CStr) -> &'block Value<'cid, 'fid> {
        unsafe {
            &*(LLVMBuildLoad(self.as_raw(), ptr.as_raw(), name.as_ptr()) as *const Value)
        }
    }

    pub fn store(&mut self, value: &Value<'cid, 'fid>, ptr: &Value<'cid, 'fid>) {
        unsafe {
            LLVMBuildStore(self.as_raw(), value.as_raw(), ptr.as_raw());
        }
    }

    pub fn ret(&mut self, value: &Value<'cid, 'fid>) {
        unsafe {
            LLVMBuildRet(self.as_raw(), value.as_raw());
        }
    }

    pub fn ret_void(&mut self) {
        unsafe {
            LLVMBuildRetVoid(self.as_raw());
        }
    }

    pub fn as_raw(&self) -> LLVMBuilderRef {
        self as *const PositionedBuilder as *mut PositionedBuilder as LLVMBuilderRef
    }
}
<|endoftext|>


<commit_before>use subframe;
use subframe::{Subframe, MAX_FIXED_ORDER, MAX_LPC_ORDER};

use std::ptr;

pub fn fixed_restore_signal(order: usize,
                            residual: &[i32],
                            output: &mut [i32]) {
  debug_assert!(order <= MAX_FIXED_ORDER);

  let polynomial = [ &[][..]
                   , &[1][..]
                   , &[-1, 2][..]
                   , &[1, -3, 3][..]
                   , &[-1, 4, -6, 4][..]
                   ];

  let coefficients = polynomial[order];

  for i in 0..residual.len() {
    let offset     = i + order;
    let prediction = coefficients.iter()
                      .zip(&output[i..offset])
                      .fold(0, |result, (coefficient, signal)|
                            result + coefficient * signal);

    output[offset] = residual[i] + prediction;
  }
}

pub fn lpc_restore_signal(quantization_level: i8,
                          coefficients: &[i32],
                          residual: &[i32],
                          output: &mut [i32]) {
  let order = coefficients.len();

  debug_assert!(order <= MAX_LPC_ORDER);

  for i in 0..residual.len() {
    let offset     = i + order;
    let prediction = coefficients.iter().rev()
                       .zip(&output[i..offset])
                       .fold(0, |result, (coefficient, signal)|
                             result + coefficient * signal);

    output[offset] = residual[i] + (prediction >> quantization_level);
  }
}

pub fn decode(subframe: &Subframe, output: &mut [i32]) {
  match subframe.data {
    subframe::Data::Constant(constant)     => {
      for i in 0..output.len() {
        output[i] = constant
      }
    }
    subframe::Data::Verbatim(ref verbatim) => {
      let length = verbatim.len();

      unsafe {
        ptr::copy(verbatim.as_ptr(), output.as_mut_ptr(), length)
      }
    }
    subframe::Data::Fixed(ref fixed)       => {
      let order = fixed.order as usize;

      for i in 0..order {
        output[i] = fixed.warmup[i];
      }

      fixed_restore_signal(order, &fixed.residual, output);
    }
    subframe::Data::LPC(ref lpc)           => {
      let order        = lpc.order as usize;
      let coefficients = &lpc.qlp_coefficients[0..order];

      for i in 0..order {
        output[i] = lpc.warmup[i];
      }

      lpc_restore_signal(lpc.quantization_level, coefficients, &lpc.residual,
                         output);
    }
  }

  if subframe.wasted_bits > 0 {
    for value in output {
      *value <<= subframe.wasted_bits;
    }
  }
}

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

  use subframe;
  use subframe::{

  #[test]
  fn test_fixed_restore_signal() {
    let residuals   = [ &[-19, -16, 17, -23, -7, 16, -16, -5, 3
                         , -8, -13, -15, -1][..]
                      , &[-6513][..]
                      ];
    let mut outputs = [ &mut [-729, -722, -667, 0, 0, 0, 0, 0, 0
                             , 0, 0, 0, 0, 0, 0, 0][..]
                      , &mut [21877, 27482, 0][..]
                      ];

    fixed_restore_signal(3, &residuals[0], &mut outputs[0]);
    fixed_restore_signal(2, &residuals[1], &mut outputs[1]);

    assert_eq!(&outputs[0], &[-729, -722, -667, -583, -486, -359, -225, -91
                             , 59, 209, 354, 497, 630, 740, 812, 845]);
    assert_eq!(&outputs[1], &[21877, 27482, 26574]);
  }

  #[test]
  fn test_lpc_restore_signal() {
    let residuals    = [ &[-2, -23, 14, 6, 3, -4, 12, -2, 10][..]
                       , &[3157][..]
                       ];
    let coefficients = [ &[1042, -399, -75, -269, 121, 166, -75][..]
                       , &[1757, -1199, 879, -836, 555, -255, 119][..]
                       ];
    let mut outputs = [ &mut [-796, -547, -285, -32, 199, 443, 670
                             , 0, 0, 0, 0, 0, 0, 0, 0, 0][..]
                      , &mut [-21363, -21951, -22649, -24364, -27297, -26870
                             ,-30017, 0][..]
                      ];

    lpc_restore_signal(9, &coefficients[0], &residuals[0], &mut outputs[0]);
    lpc_restore_signal(10, &coefficients[1], &residuals[1], &mut outputs[1]);

    assert_eq!(&outputs[0], &[-796, -547, -285, -32, 199, 443, 670, 875
                             , 1046, 1208, 1343, 1454, 1541, 1616, 1663
                             , 1701]);
    assert_eq!(&outputs[1], &[-21363, -21951, -22649, -24364, -27297, -26870
                             , -30017, -29718]);
  }

  #[test]
  fn test_decode() {
    let mut output = [0; 16];

    let constant = Subframe {
      data: subframe::Data::Constant(4),
      wasted_bits: 0,
    };

    let verbatim = Subframe {
      data: subframe::Data::Verbatim(vec![16, -3, 55, 49, -32, 6, 40, -90, 1
                                         ,0, 77, -12, 84, 10, -112, 136]),
      wasted_bits: 0,
    };

    decode(&constant, &mut output);
    assert_eq!(&output, &[4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4]);

    decode(&verbatim, &mut output);
    assert_eq!(&output, &[16, -3, 55, 49, -32, 6, 40 , -90, 1, 0, 77, -12, 84
                         ,10 , -112, 136]);
  }
}
<commit_msg>Add test for `Fixed` subframe data for decode<commit_after>use subframe;
use subframe::{Subframe, MAX_FIXED_ORDER, MAX_LPC_ORDER};

use std::ptr;

pub fn fixed_restore_signal(order: usize,
                            residual: &[i32],
                            output: &mut [i32]) {
  debug_assert!(order <= MAX_FIXED_ORDER);

  let polynomial = [ &[][..]
                   , &[1][..]
                   , &[-1, 2][..]
                   , &[1, -3, 3][..]
                   , &[-1, 4, -6, 4][..]
                   ];

  let coefficients = polynomial[order];

  for i in 0..residual.len() {
    let offset     = i + order;
    let prediction = coefficients.iter()
                      .zip(&output[i..offset])
                      .fold(0, |result, (coefficient, signal)|
                            result + coefficient * signal);

    output[offset] = residual[i] + prediction;
  }
}

pub fn lpc_restore_signal(quantization_level: i8,
                          coefficients: &[i32],
                          residual: &[i32],
                          output: &mut [i32]) {
  let order = coefficients.len();

  debug_assert!(order <= MAX_LPC_ORDER);

  for i in 0..residual.len() {
    let offset     = i + order;
    let prediction = coefficients.iter().rev()
                       .zip(&output[i..offset])
                       .fold(0, |result, (coefficient, signal)|
                             result + coefficient * signal);

    output[offset] = residual[i] + (prediction >> quantization_level);
  }
}

pub fn decode(subframe: &Subframe, output: &mut [i32]) {
  match subframe.data {
    subframe::Data::Constant(constant)     => {
      for i in 0..output.len() {
        output[i] = constant
      }
    }
    subframe::Data::Verbatim(ref verbatim) => {
      let length = verbatim.len();

      unsafe {
        ptr::copy(verbatim.as_ptr(), output.as_mut_ptr(), length)
      }
    }
    subframe::Data::Fixed(ref fixed)       => {
      let order = fixed.order as usize;

      for i in 0..order {
        output[i] = fixed.warmup[i];
      }

      fixed_restore_signal(order, &fixed.residual, output);
    }
    subframe::Data::LPC(ref lpc)           => {
      let order        = lpc.order as usize;
      let coefficients = &lpc.qlp_coefficients[0..order];

      for i in 0..order {
        output[i] = lpc.warmup[i];
      }

      lpc_restore_signal(lpc.quantization_level, coefficients, &lpc.residual,
                         output);
    }
  }

  if subframe.wasted_bits > 0 {
    for value in output {
      *value <<= subframe.wasted_bits;
    }
  }
}

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

  use subframe;
  use subframe::{
    Subframe, Fixed, LPC,
    EntropyCodingMethod, CodingMethod, PartitionedRice,
    PartitionedRiceContents,
  };

  #[test]
  fn test_fixed_restore_signal() {
    let residuals   = [ &[-19, -16, 17, -23, -7, 16, -16, -5, 3
                         , -8, -13, -15, -1][..]
                      , &[-6513][..]
                      ];
    let mut outputs = [ &mut [-729, -722, -667, 0, 0, 0, 0, 0, 0
                             , 0, 0, 0, 0, 0, 0, 0][..]
                      , &mut [21877, 27482, 0][..]
                      ];

    fixed_restore_signal(3, &residuals[0], &mut outputs[0]);
    fixed_restore_signal(2, &residuals[1], &mut outputs[1]);

    assert_eq!(&outputs[0], &[-729, -722, -667, -583, -486, -359, -225, -91
                             , 59, 209, 354, 497, 630, 740, 812, 845]);
    assert_eq!(&outputs[1], &[21877, 27482, 26574]);
  }

  #[test]
  fn test_lpc_restore_signal() {
    let residuals    = [ &[-2, -23, 14, 6, 3, -4, 12, -2, 10][..]
                       , &[3157][..]
                       ];
    let coefficients = [ &[1042, -399, -75, -269, 121, 166, -75][..]
                       , &[1757, -1199, 879, -836, 555, -255, 119][..]
                       ];
    let mut outputs = [ &mut [-796, -547, -285, -32, 199, 443, 670
                             , 0, 0, 0, 0, 0, 0, 0, 0, 0][..]
                      , &mut [-21363, -21951, -22649, -24364, -27297, -26870
                             ,-30017, 0][..]
                      ];

    lpc_restore_signal(9, &coefficients[0], &residuals[0], &mut outputs[0]);
    lpc_restore_signal(10, &coefficients[1], &residuals[1], &mut outputs[1]);

    assert_eq!(&outputs[0], &[-796, -547, -285, -32, 199, 443, 670, 875
                             , 1046, 1208, 1343, 1454, 1541, 1616, 1663
                             , 1701]);
    assert_eq!(&outputs[1], &[-21363, -21951, -22649, -24364, -27297, -26870
                             , -30017, -29718]);
  }

  #[test]
  fn test_decode() {
    let mut output = [0; 16];

    let constant = Subframe {
      data: subframe::Data::Constant(4),
      wasted_bits: 0,
    };

    let verbatim = Subframe {
      data: subframe::Data::Verbatim(vec![16, -3, 55, 49, -32, 6, 40, -90, 1
                                         ,0, 77, -12, 84, 10, -112, 136]),
      wasted_bits: 0,
    };

    let fixed = Subframe {
      data: subframe::Data::Fixed(Fixed {
        entropy_coding_method: EntropyCodingMethod {
          method_type: CodingMethod::PartitionedRice,
          data: PartitionedRice {
            order: 0,
            contents: PartitionedRiceContents {
              parameters: vec![],
              raw_bits: vec![],
            },
          },
        },
        order: 3,
        warmup: [-729, -722, -667, 0],
        residual: vec![-19, -16, 17, -23, -7, 16, -16, -5, 3 , -8, -13, -15
                      ,-1],
      }),
      wasted_bits: 0,
    };

    decode(&constant, &mut output);
    assert_eq!(&output, &[4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4]);

    decode(&verbatim, &mut output);
    assert_eq!(&output, &[16, -3, 55, 49, -32, 6, 40 , -90, 1, 0, 77, -12, 84
                         ,10 , -112, 136]);

    decode(&fixed, &mut output);
    assert_eq!(&output, &[-729, -722, -667, -583, -486, -359, -225, -91, 59
                         ,209, 354, 497, 630, 740, 812, 845]);
  }
}
<|endoftext|>


<commit_before>use subframe;
use subframe::{Subframe, MAX_FIXED_ORDER, MAX_LPC_ORDER};

use std::ptr;

pub fn fixed_restore_signal(order: usize,
                            residual: &[i32],
                            output: &mut [i32]) {
  debug_assert!(order <= MAX_FIXED_ORDER);

  let polynomial = [ &[][..]
                   , &[1][..]
                   , &[-1, 2][..]
                   , &[1, -3, 3][..]
                   , &[-1, 4, -6, 4][..]
                   ];

  let coefficients = polynomial[order];

  for i in 0..residual.len() {
    let offset     = i + order;
    let prediction = coefficients.iter()
                      .zip(&output[i..offset])
                      .fold(0, |result, (coefficient, signal)|
                            result + coefficient * signal);

    output[offset] = residual[i] + prediction;
  }
}

pub fn lpc_restore_signal(quantization_level: i8,
                          coefficients: &[i32],
                          residual: &[i32],
                          output: &mut [i32]) {
  let order = coefficients.len();

  debug_assert!(order <= MAX_LPC_ORDER);

  for i in 0..residual.len() {
    let offset     = i + order;
    let prediction = coefficients.iter().rev()
                       .zip(&output[i..offset])
                       .fold(0, |result, (coefficient, signal)|
                             result + coefficient * signal);

    output[offset] = residual[i] + (prediction >> quantization_level);
  }
}

/// Decodes the current subframe.
///
/// * `Constant` - fills the length of `output` with the constant value
///   within the subframe.
/// * `Verbatim` - copies the data within the verbatim subframe over to
///   `output`.
/// * `Fixed` - restore the signal of the fixed linear prediction and put
///   the result into `output`.
/// * `LPC` - restore the signal of the finite impulse response linear
///   prediction and put the result into `output`.
pub fn decode(subframe: &Subframe, output: &mut [i32]) {
  match subframe.data {
    subframe::Data::Constant(constant)     => {
      for i in 0..output.len() {
        output[i] = constant
      }
    }
    subframe::Data::Verbatim(ref verbatim) => {
      let length = verbatim.len();

      unsafe {
        ptr::copy(verbatim.as_ptr(), output.as_mut_ptr(), length)
      }
    }
    subframe::Data::Fixed(ref fixed)       => {
      let order = fixed.order as usize;

      for i in 0..order {
        output[i] = fixed.warmup[i];
      }

      fixed_restore_signal(order, &fixed.residual, output);
    }
    subframe::Data::LPC(ref lpc)           => {
      let order        = lpc.order as usize;
      let coefficients = &lpc.qlp_coefficients[0..order];

      for i in 0..order {
        output[i] = lpc.warmup[i];
      }

      lpc_restore_signal(lpc.quantization_level, coefficients, &lpc.residual,
                         output);
    }
  }

  if subframe.wasted_bits > 0 {
    for value in output {
      *value <<= subframe.wasted_bits;
    }
  }
}

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

  use subframe;
  use subframe::{
    Subframe, Fixed, LPC,
    EntropyCodingMethod, CodingMethod, PartitionedRice,
    PartitionedRiceContents,
  };

  #[test]
  fn test_fixed_restore_signal() {
    let residuals   = [ &[-19, -16, 17, -23, -7, 16, -16, -5, 3
                         , -8, -13, -15, -1][..]
                      , &[-6513][..]
                      ];
    let mut outputs = [ &mut [-729, -722, -667, 0, 0, 0, 0, 0, 0
                             , 0, 0, 0, 0, 0, 0, 0][..]
                      , &mut [21877, 27482, 0][..]
                      ];

    fixed_restore_signal(3, &residuals[0], &mut outputs[0]);
    fixed_restore_signal(2, &residuals[1], &mut outputs[1]);

    assert_eq!(&outputs[0], &[-729, -722, -667, -583, -486, -359, -225, -91
                             , 59, 209, 354, 497, 630, 740, 812, 845]);
    assert_eq!(&outputs[1], &[21877, 27482, 26574]);
  }

  #[test]
  fn test_lpc_restore_signal() {
    let residuals    = [ &[-2, -23, 14, 6, 3, -4, 12, -2, 10][..]
                       , &[3157][..]
                       ];
    let coefficients = [ &[1042, -399, -75, -269, 121, 166, -75][..]
                       , &[1757, -1199, 879, -836, 555, -255, 119][..]
                       ];
    let mut outputs = [ &mut [-796, -547, -285, -32, 199, 443, 670
                             , 0, 0, 0, 0, 0, 0, 0, 0, 0][..]
                      , &mut [-21363, -21951, -22649, -24364, -27297, -26870
                             ,-30017, 0][..]
                      ];

    lpc_restore_signal(9, &coefficients[0], &residuals[0], &mut outputs[0]);
    lpc_restore_signal(10, &coefficients[1], &residuals[1], &mut outputs[1]);

    assert_eq!(&outputs[0], &[-796, -547, -285, -32, 199, 443, 670, 875
                             , 1046, 1208, 1343, 1454, 1541, 1616, 1663
                             , 1701]);
    assert_eq!(&outputs[1], &[-21363, -21951, -22649, -24364, -27297, -26870
                             , -30017, -29718]);
  }

  #[test]
  fn test_decode() {
    let mut output = [0; 16];

    let constant = Subframe {
      data: subframe::Data::Constant(4),
      wasted_bits: 0,
    };

    let verbatim = Subframe {
      data: subframe::Data::Verbatim(vec![16, -3, 55, 49, -32, 6, 40, -90, 1
                                         ,0, 77, -12, 84, 10, -112, 136]),
      wasted_bits: 0,
    };

    let fixed = Subframe {
      data: subframe::Data::Fixed(Fixed {
        entropy_coding_method: EntropyCodingMethod {
          method_type: CodingMethod::PartitionedRice,
          data: PartitionedRice {
            order: 0,
            contents: PartitionedRiceContents {
              parameters: vec![],
              raw_bits: vec![],
            },
          },
        },
        order: 3,
        warmup: [-729, -722, -667, 0],
        residual: vec![-19, -16, 17, -23, -7, 16, -16, -5, 3 , -8, -13, -15
                      ,-1],
      }),
      wasted_bits: 0,
    };

    let lpc = Subframe {
      data: subframe::Data::LPC(LPC {
        entropy_coding_method: EntropyCodingMethod {
          method_type: CodingMethod::PartitionedRice,
          data: PartitionedRice {
            order: 0,
            contents: PartitionedRiceContents {
              parameters: vec![],
              raw_bits: vec![],
            },
          },
        },
        order: 7,
        qlp_coeff_precision: 0,
        quantization_level: 9,
        qlp_coefficients: [1042, -399, -75, -269, 121, 166, -75, 0, 0, 0, 0
                          ,0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
                          ,0, 0, 0, 0],
        warmup: [-796, -547, -285, -32, 199, 443, 670, 0, 0, 0, 0, 0, 0, 0, 0
                ,0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        residual: vec![-2, -23, 14, 6, 3, -4, 12, -2, 10],
      }),
      wasted_bits: 0,
    };

    decode(&constant, &mut output);
    assert_eq!(&output, &[4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4]);

    decode(&verbatim, &mut output);
    assert_eq!(&output, &[16, -3, 55, 49, -32, 6, 40 , -90, 1, 0, 77, -12, 84
                         ,10 , -112, 136]);

    decode(&fixed, &mut output);
    assert_eq!(&output, &[-729, -722, -667, -583, -486, -359, -225, -91, 59
                         ,209, 354, 497, 630, 740, 812, 845]);

    decode(&lpc, &mut output);
    assert_eq!(&output, &[-796, -547, -285, -32, 199, 443, 670, 875, 1046
                         ,1208, 1343, 1454, 1541, 1616, 1663, 1701]);
  }

  #[test]
  fn test_wasted_bit_decode() {
    let mut output = [0; 4];

    let constant = Subframe {
      data: subframe::Data::Constant(1),
      wasted_bits: 10,
    };

    decode(&constant, &mut output);
    assert_eq!(&output, &[1024, 1024, 1024, 1024]);
  }
}
<commit_msg>Add explanation for `fixed_restore_signal`<commit_after>use subframe;
use subframe::{Subframe, MAX_FIXED_ORDER, MAX_LPC_ORDER};

use std::ptr;

// Restore the original signal from a fixed linear prediction.
//
// Signal restoration is accomplished by summing up the residual and the
// predictor. With the linear prediction being "fixed", the polynomials will
// remain the same depending on the order value. And the max order is
// `MAX_FIXED_ORDER`, which is 4.
//
// This function also assumes that `output` already has the warm up values
// from the `Fixed` subframe in it.
pub fn fixed_restore_signal(order: usize,
                            residual: &[i32],
                            output: &mut [i32]) {
  debug_assert!(order <= MAX_FIXED_ORDER);

  let polynomial = [ &[][..]
                   , &[1][..]
                   , &[-1, 2][..]
                   , &[1, -3, 3][..]
                   , &[-1, 4, -6, 4][..]
                   ];

  let coefficients = polynomial[order];

  for i in 0..residual.len() {
    let offset     = i + order;
    let prediction = coefficients.iter()
                      .zip(&output[i..offset])
                      .fold(0, |result, (coefficient, signal)|
                            result + coefficient * signal);

    output[offset] = residual[i] + prediction;
  }
}

pub fn lpc_restore_signal(quantization_level: i8,
                          coefficients: &[i32],
                          residual: &[i32],
                          output: &mut [i32]) {
  let order = coefficients.len();

  debug_assert!(order <= MAX_LPC_ORDER);

  for i in 0..residual.len() {
    let offset     = i + order;
    let prediction = coefficients.iter().rev()
                       .zip(&output[i..offset])
                       .fold(0, |result, (coefficient, signal)|
                             result + coefficient * signal);

    output[offset] = residual[i] + (prediction >> quantization_level);
  }
}

/// Decodes the current subframe.
///
/// * `Constant` - fills the length of `output` with the constant value
///   within the subframe.
/// * `Verbatim` - copies the data within the verbatim subframe over to
///   `output`.
/// * `Fixed` - restore the signal of the fixed linear prediction and put
///   the result into `output`.
/// * `LPC` - restore the signal of the finite impulse response linear
///   prediction and put the result into `output`.
pub fn decode(subframe: &Subframe, output: &mut [i32]) {
  match subframe.data {
    subframe::Data::Constant(constant)     => {
      for i in 0..output.len() {
        output[i] = constant
      }
    }
    subframe::Data::Verbatim(ref verbatim) => {
      let length = verbatim.len();

      unsafe {
        ptr::copy(verbatim.as_ptr(), output.as_mut_ptr(), length)
      }
    }
    subframe::Data::Fixed(ref fixed)       => {
      let order = fixed.order as usize;

      for i in 0..order {
        output[i] = fixed.warmup[i];
      }

      fixed_restore_signal(order, &fixed.residual, output);
    }
    subframe::Data::LPC(ref lpc)           => {
      let order        = lpc.order as usize;
      let coefficients = &lpc.qlp_coefficients[0..order];

      for i in 0..order {
        output[i] = lpc.warmup[i];
      }

      lpc_restore_signal(lpc.quantization_level, coefficients, &lpc.residual,
                         output);
    }
  }

  if subframe.wasted_bits > 0 {
    for value in output {
      *value <<= subframe.wasted_bits;
    }
  }
}

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

  use subframe;
  use subframe::{
    Subframe, Fixed, LPC,
    EntropyCodingMethod, CodingMethod, PartitionedRice,
    PartitionedRiceContents,
  };

  #[test]
  fn test_fixed_restore_signal() {
    let residuals   = [ &[-19, -16, 17, -23, -7, 16, -16, -5, 3
                         , -8, -13, -15, -1][..]
                      , &[-6513][..]
                      ];
    let mut outputs = [ &mut [-729, -722, -667, 0, 0, 0, 0, 0, 0
                             , 0, 0, 0, 0, 0, 0, 0][..]
                      , &mut [21877, 27482, 0][..]
                      ];

    fixed_restore_signal(3, &residuals[0], &mut outputs[0]);
    fixed_restore_signal(2, &residuals[1], &mut outputs[1]);

    assert_eq!(&outputs[0], &[-729, -722, -667, -583, -486, -359, -225, -91
                             , 59, 209, 354, 497, 630, 740, 812, 845]);
    assert_eq!(&outputs[1], &[21877, 27482, 26574]);
  }

  #[test]
  fn test_lpc_restore_signal() {
    let residuals    = [ &[-2, -23, 14, 6, 3, -4, 12, -2, 10][..]
                       , &[3157][..]
                       ];
    let coefficients = [ &[1042, -399, -75, -269, 121, 166, -75][..]
                       , &[1757, -1199, 879, -836, 555, -255, 119][..]
                       ];
    let mut outputs = [ &mut [-796, -547, -285, -32, 199, 443, 670
                             , 0, 0, 0, 0, 0, 0, 0, 0, 0][..]
                      , &mut [-21363, -21951, -22649, -24364, -27297, -26870
                             ,-30017, 0][..]
                      ];

    lpc_restore_signal(9, &coefficients[0], &residuals[0], &mut outputs[0]);
    lpc_restore_signal(10, &coefficients[1], &residuals[1], &mut outputs[1]);

    assert_eq!(&outputs[0], &[-796, -547, -285, -32, 199, 443, 670, 875
                             , 1046, 1208, 1343, 1454, 1541, 1616, 1663
                             , 1701]);
    assert_eq!(&outputs[1], &[-21363, -21951, -22649, -24364, -27297, -26870
                             , -30017, -29718]);
  }

  #[test]
  fn test_decode() {
    let mut output = [0; 16];

    let constant = Subframe {
      data: subframe::Data::Constant(4),
      wasted_bits: 0,
    };

    let verbatim = Subframe {
      data: subframe::Data::Verbatim(vec![16, -3, 55, 49, -32, 6, 40, -90, 1
                                         ,0, 77, -12, 84, 10, -112, 136]),
      wasted_bits: 0,
    };

    let fixed = Subframe {
      data: subframe::Data::Fixed(Fixed {
        entropy_coding_method: EntropyCodingMethod {
          method_type: CodingMethod::PartitionedRice,
          data: PartitionedRice {
            order: 0,
            contents: PartitionedRiceContents {
              parameters: vec![],
              raw_bits: vec![],
            },
          },
        },
        order: 3,
        warmup: [-729, -722, -667, 0],
        residual: vec![-19, -16, 17, -23, -7, 16, -16, -5, 3 , -8, -13, -15
                      ,-1],
      }),
      wasted_bits: 0,
    };

    let lpc = Subframe {
      data: subframe::Data::LPC(LPC {
        entropy_coding_method: EntropyCodingMethod {
          method_type: CodingMethod::PartitionedRice,
          data: PartitionedRice {
            order: 0,
            contents: PartitionedRiceContents {
              parameters: vec![],
              raw_bits: vec![],
            },
          },
        },
        order: 7,
        qlp_coeff_precision: 0,
        quantization_level: 9,
        qlp_coefficients: [1042, -399, -75, -269, 121, 166, -75, 0, 0, 0, 0
                          ,0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
                          ,0, 0, 0, 0],
        warmup: [-796, -547, -285, -32, 199, 443, 670, 0, 0, 0, 0, 0, 0, 0, 0
                ,0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        residual: vec![-2, -23, 14, 6, 3, -4, 12, -2, 10],
      }),
      wasted_bits: 0,
    };

    decode(&constant, &mut output);
    assert_eq!(&output, &[4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4]);

    decode(&verbatim, &mut output);
    assert_eq!(&output, &[16, -3, 55, 49, -32, 6, 40 , -90, 1, 0, 77, -12, 84
                         ,10 , -112, 136]);

    decode(&fixed, &mut output);
    assert_eq!(&output, &[-729, -722, -667, -583, -486, -359, -225, -91, 59
                         ,209, 354, 497, 630, 740, 812, 845]);

    decode(&lpc, &mut output);
    assert_eq!(&output, &[-796, -547, -285, -32, 199, 443, 670, 875, 1046
                         ,1208, 1343, 1454, 1541, 1616, 1663, 1701]);
  }

  #[test]
  fn test_wasted_bit_decode() {
    let mut output = [0; 4];

    let constant = Subframe {
      data: subframe::Data::Constant(1),
      wasted_bits: 10,
    };

    decode(&constant, &mut output);
    assert_eq!(&output, &[1024, 1024, 1024, 1024]);
  }
}
<|endoftext|>


<commit_before>#![feature(phase)]

extern crate regex;
#[phase(plugin)] extern crate regex_macros;

use std::io::{File, IoResult, IoError, OtherIoError};
use std::collections::HashMap;
use std::default::Default;
use std::os;

pub struct Dotenv {
	env: HashMap<Vec<u8>, Vec<u8>>
}

fn merge_env_hashes(
	a: HashMap<Vec<u8>, Vec<u8>>,
	b: HashMap<Vec<u8>, Vec<u8>>,
) -> HashMap<Vec<u8>, Vec<u8>> {
	let mut result = a.clone();
	result.extend(b.into_iter());
	result
}

fn parse_env(data: Vec<u8>) -> HashMap<Vec<u8>, Vec<u8>> {
	let line_splitter = regex!(r"[\r\n]+");
	let env_format = regex!(
		r"^(?P<key>[A-Za-z_][A-Za-z0-9_]*)\s*=\s*(?P<value>.+)$"
	);

	let data_string = String::from_utf8_lossy(data.as_slice());
	let data_slice = data_string.as_slice();

	line_splitter.split(data_slice).filter_map(|line| {
		env_format.captures(line).and_then(|captures| {
			let key = captures.name("key");
			let value = captures.name("value");

			// regex captures return empty string on failure
			if key == "" || value == "" {
				None
			} else {
				Some((
					String::from_str(key).into_bytes(),
					String::from_str(value).into_bytes()
				))
			}
		})
	}).collect()
}

impl Dotenv {
	pub fn dotenv() -> Dotenv {
		Dotenv::from_filename(".env").ok().unwrap_or_default()
	}

	pub fn from_str(s: &str) -> Dotenv {
		Dotenv::from_bytes(String::from_str(s).into_bytes())
	}

	pub fn from_bytes(bytes: Vec<u8>) -> Dotenv {
		let dotenv_hash = parse_env(bytes);
		let env_hash = os::env_as_bytes().into_iter().collect();
		Dotenv{env: merge_env_hashes(env_hash, dotenv_hash)}
	}

	pub fn from_file(mut file: File) -> IoResult<Dotenv> {
		file.read_to_end().map(|file_contents|
			Dotenv::from_bytes(file_contents)
		)
	}

	pub fn from_filename(n: &str) -> IoResult<Dotenv> {
		os::self_exe_path().as_mut().map(|path| {
			path.push(n);
			Dotenv::from_path(path)
		}).unwrap_or(Err(IoError{
			kind: OtherIoError,
			desc: "Could not fetch the path of this executable",
			detail: None
		}))
	}

	pub fn from_path(path: &Path) -> IoResult<Dotenv> {
		File::open(path).and_then(|file| {
			Dotenv::from_file(file)
		})
	}

	pub fn getenv(&self, n: &str) -> Option<String> {
		let n_bytes = &String::from_str(n).into_bytes();
		self.env.find(n_bytes).and_then(|bytes| {
			String::from_utf8(bytes.clone()).ok()
		})
	}

	pub fn getenv_as_bytes(&self, n: &str) -> Option<Vec<u8>> {
		let n_bytes = &String::from_str(n).into_bytes();
		self.env.find(n_bytes).map(|bytes| {
			bytes.clone()})
		}

	pub fn env(&self) -> Vec<(String, String)> {
		self.env.iter().filter_map(|(key, value)| {
			String::from_utf8(key.clone()).ok().and_then(|key_string| {
				String::from_utf8(value.clone()).ok().map(|value_string| {
					(key_string.clone(), value_string)
				})
			})
		}).collect()
	}

	pub fn env_as_bytes(&self) -> Vec<(Vec<u8>, Vec<u8>)> {
		self.env.iter().map(|(key, value)| {
			(key.clone(), value.clone())
		}).collect()
	}
}

impl Default for Dotenv {
	fn default() -> Dotenv {
		Dotenv::from_bytes(Default::default())
	}
}

#[test]
fn merge_env_hashes_test() {
	let hash_a: HashMap<Vec<u8>, Vec<u8>> = vec![
		(
			String::from_str("bacon").into_bytes(),
			String::from_str("bad").into_bytes()
		), (
			String::from_str("cabbage").into_bytes(),
			String::from_str("neutral").into_bytes()
		)
	].into_iter().collect();

	let hash_b: HashMap<Vec<u8>, Vec<u8>> = vec![
		(
			String::from_str("bacon").into_bytes(),
			String::from_str("good").into_bytes()
		)
	].into_iter().collect();

	let expected_hash: HashMap<Vec<u8>, Vec<u8>> = vec![
		(
			String::from_str("bacon").into_bytes(),
			String::from_str("good").into_bytes()
		), (
			String::from_str("cabbage").into_bytes(),
			String::from_str("neutral").into_bytes()
		)
	].into_iter().collect();

	let actual_hash = merge_env_hashes(hash_a, hash_b);

	assert_eq!(expected_hash, actual_hash)
}

#[test]
fn parse_env_test_basic() {
	let input = String::from_str(concat!(
		"bacon=good", "\n",
		"cabbage=neutral", "\n"
	)).into_bytes();

	let expected_hash: HashMap<Vec<u8>, Vec<u8>> = vec![
		(
			String::from_str("bacon").into_bytes(),
			String::from_str("good").into_bytes()
		), (
			String::from_str("cabbage").into_bytes(),
			String::from_str("neutral").into_bytes()
		)
	].into_iter().collect();

	let actual_hash = parse_env(input);

	assert_eq!(expected_hash, actual_hash);
}

#[test]
fn parse_env_test_edges() {
	let input = String::from_str(concat!(
		"chili=spicy=EVIL", "\n",       // equals on value, uppercase on value
		"to_MAH_to=okay", "\n",         // uppercase and underscores on key
		"bacon   =  good", "\n\r\n\r",  // spaces around equal, newlines
		"cabbage=not evil"              // space on value, no final newline
	)).into_bytes();

	let expected_hash: HashMap<Vec<u8>, Vec<u8>> = vec![
		(
			String::from_str("chili").into_bytes(),
			String::from_str("spicy=EVIL").into_bytes()
		), (
			String::from_str("to_MAH_to").into_bytes(),
			String::from_str("okay").into_bytes()
		), (
			String::from_str("bacon").into_bytes(),
			String::from_str("good").into_bytes()
		), (
			String::from_str("cabbage").into_bytes(),
			String::from_str("not evil").into_bytes()
		)
	].into_iter().collect();

	let actual_hash = parse_env(input);

	assert_eq!(expected_hash, actual_hash);
}
<commit_msg>Big rewrite: simpler API, more testable<commit_after>#![feature(phase)]

extern crate regex;
#[phase(plugin)] extern crate regex_macros;

use std::io::{BufferedReader, File, IoError, OtherIoError};
use std::os::{setenv, self_exe_path};

#[deriving(Show, Clone)]
pub struct ParseError {
	line: String
}

#[deriving(Show, Clone)]
pub enum DotenvError {
	Parse(ParseError),
	Io(IoError)
}

fn parse_line(line: String) -> Result<Option<(String, String)>, ParseError> {
	let line_regex = regex!(concat!(r"^(\s*(",
		r"#.*|", // A comment, or...
		r"(export\s+)?", // ...(optionally preceded by "export")...
		r"(?P<key>[A-Za-z_][A-Za-z0-9_]*)", // ...a key,...
		r"\s*=\s*", // ...then an equal sign,...
		r"(?P<value>.+?)", // ...and then its corresponding value.
	r")\s*)$"));

	line_regex.captures(line.as_slice()).map_or(
		Err(ParseError{line: line.clone()}),
		|captures| {
			let key = captures.name("key");
			let value = captures.name("value");

			if key == "" || value == "" {
				// If there's no key and value, but capturing did not fail,
				// then this means we're dealing with a comment.
				Ok(None)
			} else {
				Ok(Some((key.to_string(), value.to_string())))
			}
		}
	)
}

fn from_line_iter<T: Iterator<String>>(lines: T) -> Result<(), ParseError> {
	let mut parsed_lines = lines.map(parse_line);
	let failure = parsed_lines.find(|line| line.is_err());

	if failure.is_some() {
		return Err(failure.unwrap().err().unwrap());
	}

	for (key, value) in parsed_lines.filter_map(|line| {
		line.clone().unwrap()
	}) {
		setenv(key.as_slice(), value);
	}

	Ok(())
}

fn from_file(file: File) -> Result<(), DotenvError> {
	let mut reader = BufferedReader::new(file);
	let lines = reader.lines();

	from_line_iter(lines.filter_map(|result| {
		result.ok()
	})).map_err(|err| {
		Parse(err)
	})
}

pub fn from_path(path: &Path) -> Result<(), DotenvError> {
	match File::open(path) {
		Ok(file) => from_file(file),
		Err(err) => Err(Io(err))
	}
}

pub fn from_filename(filename: &str) -> Result<(), DotenvError> {
	self_exe_path().as_mut().map(|path| {
		path.push(filename);
		from_path(path)
	}).unwrap_or(Err(Io(IoError{
		kind: OtherIoError,
		desc: "Could not fetch the path of this executable",
		detail: None
	})))
}

pub fn dotenv() -> Result<(), DotenvError> {
	from_filename(".env")
}

#[test]
fn test_parse_line_env() {
	let input_iter = vec![
		"THIS_IS_KEY=hi this is value",
		"   many_spaces  =   wow a  maze   ",
		"export   SHELL_LOVER=1"
	].into_iter().map(|input| input.to_string());
	let actual_iter = input_iter.map(|input| parse_line(input));

	let expected_iter = vec![
		("THIS_IS_KEY", "hi this is value"),
		("many_spaces", "wow a  maze"),
		("SHELL_LOVER", "1")
	].into_iter().map(|(key, value)| (key.to_string(), value.to_string()));

	for (expected, actual) in expected_iter.zip(actual_iter) {
		assert!(actual.is_ok());
		assert!(actual.clone().ok().unwrap().is_some())
		assert_eq!(expected, actual.ok().unwrap().unwrap())
	}
}

#[test]
fn test_parse_line_comment() {
	let input_iter = vec![
		"# foo=bar",
		"    #    "
	].into_iter().map(|input| input.to_string());
	let mut actual_iter = input_iter.map(|input| parse_line(input));

	for actual in actual_iter {
		assert!(actual.is_ok());
		assert!(actual.ok().unwrap().is_none());
	}
}

#[test]
fn test_parse_line_invalid() {
	let input_iter = vec![
		"  invalid    ",
		"very bacon = yes indeed",
		"key=",
		"=value"
	].into_iter().map(|input| input.to_string());
	let mut actual_iter = input_iter.map(|input| parse_line(input));

	for actual in actual_iter {
		assert!(actual.is_err());
	}
}

#[test]
fn test_from_line_iter_valid() {
	let input = vec![
		"test_env_one=1",
		"# a comment",
		"test_env_two=2"
	].into_iter().map(|line| line.to_string());
	let actual = from_line_iter(input);

	assert!(actual.is_ok())
}

#[test]
fn test_from_line_iter_invalid() {
	let input = vec![
		"test_env_one=1",
		"# a comment",
		"not valid"
	].into_iter().map(|line| line.to_string());
	let actual = from_line_iter(input);

	assert!(actual.is_err())
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::errors::*;
use determinator::Utf8Paths0;
use guppy::{graph::PackageGraph, MetadataCommand};
use indoc::formatdoc;
use log::{debug, info};
use once_cell::sync::OnceCell;
use std::{
    borrow::Cow,
    ffi::{OsStr, OsString},
    fmt,
    path::{Path, PathBuf},
    process::{Command, Stdio},
};

/// Support for source control operations through running Git commands.
///
/// This assumes that the underlying Git repository doesn't change in the middle of an operation,
/// and caches data as a result. If mutation operations are added, the caches would need to be
/// invalidated.
#[derive(Clone, Debug)]
pub struct GitCli {
    root: &'static Path,
    // Caches.
    tracked_files: OnceCell<Utf8Paths0>,
}

impl GitCli {
    /// Creates a new instance of the Git CLI.
    pub fn new(root: &'static Path) -> Result<Self> {
        let git_cli = Self {
            root,
            tracked_files: OnceCell::new(),
        };
        git_cli.validate()?;
        Ok(git_cli)
    }

    /// Returns the files tracked by Git in this working copy.
    ///
    /// The return value can be iterated on to get a list of paths.
    pub fn tracked_files(&self) -> Result<&Utf8Paths0> {
        self.tracked_files.get_or_try_init(|| {
            // TODO: abstract out SCM and command-running functionality.
            let output = self
                .git_command()
                // The -z causes files to not be quoted, and to be separated by \0.
                .args(&["ls-files", "-z"])
                .output()
                .map_err(|err| SystemError::io("running git ls-files", err))?;
            if !output.status.success() {
                return Err(SystemError::Exec {
                    cmd: "git ls-files",
                    status: output.status,
                });
            }

            Utf8Paths0::from_bytes(output.stdout)
                .map_err(|(path, err)| SystemError::NonUtf8Path { path, err })
        })
    }

    /// Returns the merge base of the current commit (`HEAD`) with the specified commit.
    pub fn merge_base(&self, commit_ref: &str) -> Result<GitHash> {
        let output = self
            .git_command()
            .args(&["merge-base", "HEAD", commit_ref])
            .output()
            .map_err(|err| {
                SystemError::io(format!("running git merge-base HEAD {}", commit_ref), err)
            })?;
        if !output.status.success() {
            return Err(SystemError::Exec {
                cmd: "git merge-base",
                status: output.status,
            });
        }

        // The output is a hex-encoded hash followed by a newline.
        let stdout = &output.stdout[..(output.stdout.len() - 1)];
        GitHash::from_hex(stdout)
    }

    /// Returns the files changed between the given commits, or the current directory if the new
    /// commit isn't specified.
    ///
    /// For more about the diff filter, see `man git-diff`'s help for `--diff-filter`.
    pub fn files_changed_between<'a>(
        &self,
        old: impl Into<Cow<'a, OsStr>>,
        new: impl Into<Option<Cow<'a, OsStr>>>,
        // TODO: make this more well-typed/express more of the diff model in Rust
        diff_filter: Option<&str>,
    ) -> Result<Utf8Paths0> {
        let mut command = self.git_command();
        command.args(&["diff", "-z", "--name-only"]);
        if let Some(diff_filter) = diff_filter {
            command.arg(format!("--diff-filter={}", diff_filter));
        }
        command.arg(old.into());
        if let Some(new) = new.into() {
            command.arg(new);
        }

        let output = command
            .output()
            .map_err(|err| SystemError::io("running git diff", err))?;
        if !output.status.success() {
            return Err(SystemError::Exec {
                cmd: "git diff",
                status: output.status,
            });
        }

        Utf8Paths0::from_bytes(output.stdout)
            .map_err(|(path, err)| SystemError::NonUtf8Path { path, err })
    }

    /// Returns a package graph for the given commit, using a scratch repo if necessary.
    pub fn package_graph_at(&self, commit_ref: &GitHash) -> Result<PackageGraph> {
        // Create or initialize the scratch worktree.
        let scratch = self.get_or_init_scratch(commit_ref)?;

        // Compute the package graph for the scratch worktree.
        MetadataCommand::new()
            .current_dir(scratch)
            .build_graph()
            .map_err(|err| SystemError::guppy("building package graph", err))
    }

    // ---
    // Helper methods
    // ---

    fn validate(&self) -> Result<()> {
        // Check that the project root and the Git root match.
        let output = self
            .git_command()
            .args(&["rev-parse", "--show-toplevel"])
            .stderr(Stdio::inherit())
            .output()
            .map_err(|err| SystemError::io("running git rev-parse --show-toplevel", err))?;
        if !output.status.success() {
            let msg = formatdoc!(
                "unable to find a git repo at {}
                (hint: did you download an archive from GitHub? x requires a git clone)",
                self.root.display()
            );
            return Err(SystemError::git_root(msg));
        }

        let mut git_root_bytes = output.stdout;
        // Pop the newline off the git root bytes.
        git_root_bytes.pop();
        let git_root = match String::from_utf8(git_root_bytes) {
            Ok(git_root) => git_root,
            Err(_) => {
                return Err(SystemError::git_root(
                    "git rev-parse --show-toplevel returned a non-Unicode path",
                ));
            }
        };
        if self.root != Path::new(&git_root) {
            let msg = formatdoc!(
                "git root expected to be at {}, but actually found at {}
                (hint: did you download an archive from GitHub? x requires a git clone)",
                self.root.display(),
                git_root,
            );
            return Err(SystemError::git_root(msg));
        }
        Ok(())
    }

    // TODO: abstract out command running and error handling
    fn git_command(&self) -> Command {
        // TODO: add support for the GIT environment variable?
        let mut command = Command::new("git");
        command.current_dir(self.root).stderr(Stdio::inherit());
        command
    }

    /// Gets the scratch worktree if it exists, or initializes it if it doesn't.
    ///
    /// The scratch worktree is meant to be persistent across invocations of `x`. This is done for
    /// performance reasons.
    fn get_or_init_scratch(&self, hash: &GitHash) -> Result<PathBuf> {
        let mut scratch_dir = self.root.join("target");
        scratch_dir.extend(&["x-scratch", "tree"]);

        if scratch_dir.is_dir() {
            debug!(
                "Using existing scratch worktree at {}",
                scratch_dir.display()
            );

            // TODO: check if the directory is actually a Git worktree.

            // Check out the given hash in the scratch worktree.
            let output = self
                .git_command()
                .current_dir(&scratch_dir)
                // TODO: also git clean?
                .args(&["reset", &format!("{:x}", hash), "--hard"])
                .output()
                .map_err(|err| SystemError::io("running git checkout in scratch tree", err))?;
            if !output.status.success() {
                return Err(SystemError::Exec {
                    cmd: "git checkout",
                    status: output.status,
                });
            }
        } else {
            // Try creating a scratch worktree at that location.
            info!("Setting up scratch worktree in {}", scratch_dir.display());
            let output = self
                .git_command()
                .args(&["worktree", "add"])
                .arg(&scratch_dir)
                .args(&[&format!("{:x}", hash), "--detach"])
                .output()
                .map_err(|err| SystemError::io("running git worktree add", err))?;
            if !output.status.success() {
                return Err(SystemError::Exec {
                    cmd: "git worktree add",
                    status: output.status,
                });
            }
        }

        // TODO: some sort of cross-process locking may be necessary in the future. Don't worry
        // about it for now.
        Ok(scratch_dir)
    }
}

/// A Git hash.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct GitHash([u8; 20]);

impl GitHash {
    /// Creates a new Git hash from a hex-encoded string.
    pub fn from_hex(hex: impl AsRef<[u8]>) -> Result<Self> {
        let hex = hex.as_ref();
        Ok(GitHash(hex::FromHex::from_hex(hex).map_err(|err| {
            SystemError::from_hex(format!("parsing a Git hash: {:?}", hex), err)
        })?))
    }
}

impl<'a, 'b> From<&'a GitHash> for Cow<'b, OsStr> {
    fn from(git_hash: &'a GitHash) -> Cow<'b, OsStr> {
        OsString::from(format!("{:x}", git_hash)).into()
    }
}

impl fmt::LowerHex for GitHash {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", hex::encode(&self.0))
    }
}
<commit_msg>x: check if scratch dir is git repo<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::errors::*;
use determinator::Utf8Paths0;
use guppy::{graph::PackageGraph, MetadataCommand};
use indoc::formatdoc;
use log::{debug, info};
use once_cell::sync::OnceCell;
use std::{
    borrow::Cow,
    ffi::{OsStr, OsString},
    fmt,
    path::{Path, PathBuf},
    process::{Command, Stdio},
};

/// Support for source control operations through running Git commands.
///
/// This assumes that the underlying Git repository doesn't change in the middle of an operation,
/// and caches data as a result. If mutation operations are added, the caches would need to be
/// invalidated.
#[derive(Clone, Debug)]
pub struct GitCli {
    root: &'static Path,
    // Caches.
    tracked_files: OnceCell<Utf8Paths0>,
}

impl GitCli {
    /// Creates a new instance of the Git CLI.
    pub fn new(root: &'static Path) -> Result<Self> {
        let git_cli = Self {
            root,
            tracked_files: OnceCell::new(),
        };
        git_cli.validate()?;
        Ok(git_cli)
    }

    /// Returns the files tracked by Git in this working copy.
    ///
    /// The return value can be iterated on to get a list of paths.
    pub fn tracked_files(&self) -> Result<&Utf8Paths0> {
        self.tracked_files.get_or_try_init(|| {
            // TODO: abstract out SCM and command-running functionality.
            let output = self
                .git_command()
                // The -z causes files to not be quoted, and to be separated by \0.
                .args(&["ls-files", "-z"])
                .output()
                .map_err(|err| SystemError::io("running git ls-files", err))?;
            if !output.status.success() {
                return Err(SystemError::Exec {
                    cmd: "git ls-files",
                    status: output.status,
                });
            }

            Utf8Paths0::from_bytes(output.stdout)
                .map_err(|(path, err)| SystemError::NonUtf8Path { path, err })
        })
    }

    /// Returns the merge base of the current commit (`HEAD`) with the specified commit.
    pub fn merge_base(&self, commit_ref: &str) -> Result<GitHash> {
        let output = self
            .git_command()
            .args(&["merge-base", "HEAD", commit_ref])
            .output()
            .map_err(|err| {
                SystemError::io(format!("running git merge-base HEAD {}", commit_ref), err)
            })?;
        if !output.status.success() {
            return Err(SystemError::Exec {
                cmd: "git merge-base",
                status: output.status,
            });
        }

        // The output is a hex-encoded hash followed by a newline.
        let stdout = &output.stdout[..(output.stdout.len() - 1)];
        GitHash::from_hex(stdout)
    }

    /// Returns the files changed between the given commits, or the current directory if the new
    /// commit isn't specified.
    ///
    /// For more about the diff filter, see `man git-diff`'s help for `--diff-filter`.
    pub fn files_changed_between<'a>(
        &self,
        old: impl Into<Cow<'a, OsStr>>,
        new: impl Into<Option<Cow<'a, OsStr>>>,
        // TODO: make this more well-typed/express more of the diff model in Rust
        diff_filter: Option<&str>,
    ) -> Result<Utf8Paths0> {
        let mut command = self.git_command();
        command.args(&["diff", "-z", "--name-only"]);
        if let Some(diff_filter) = diff_filter {
            command.arg(format!("--diff-filter={}", diff_filter));
        }
        command.arg(old.into());
        if let Some(new) = new.into() {
            command.arg(new);
        }

        let output = command
            .output()
            .map_err(|err| SystemError::io("running git diff", err))?;
        if !output.status.success() {
            return Err(SystemError::Exec {
                cmd: "git diff",
                status: output.status,
            });
        }

        Utf8Paths0::from_bytes(output.stdout)
            .map_err(|(path, err)| SystemError::NonUtf8Path { path, err })
    }

    /// Returns a package graph for the given commit, using a scratch repo if necessary.
    pub fn package_graph_at(&self, commit_ref: &GitHash) -> Result<PackageGraph> {
        // Create or initialize the scratch worktree.
        let scratch = self.get_or_init_scratch(commit_ref)?;

        // Compute the package graph for the scratch worktree.
        MetadataCommand::new()
            .current_dir(scratch)
            .build_graph()
            .map_err(|err| SystemError::guppy("building package graph", err))
    }

    // ---
    // Helper methods
    // ---

    fn validate(&self) -> Result<()> {
        // Check that the project root and the Git root match.
        let output = self
            .git_command()
            .args(&["rev-parse", "--show-toplevel"])
            .stderr(Stdio::inherit())
            .output()
            .map_err(|err| SystemError::io("running git rev-parse --show-toplevel", err))?;
        if !output.status.success() {
            let msg = formatdoc!(
                "unable to find a git repo at {}
                (hint: did you download an archive from GitHub? x requires a git clone)",
                self.root.display()
            );
            return Err(SystemError::git_root(msg));
        }

        let mut git_root_bytes = output.stdout;
        // Pop the newline off the git root bytes.
        git_root_bytes.pop();
        let git_root = match String::from_utf8(git_root_bytes) {
            Ok(git_root) => git_root,
            Err(_) => {
                return Err(SystemError::git_root(
                    "git rev-parse --show-toplevel returned a non-Unicode path",
                ));
            }
        };
        if self.root != Path::new(&git_root) {
            let msg = formatdoc!(
                "git root expected to be at {}, but actually found at {}
                (hint: did you download an archive from GitHub? x requires a git clone)",
                self.root.display(),
                git_root,
            );
            return Err(SystemError::git_root(msg));
        }
        Ok(())
    }

    // TODO: abstract out command running and error handling
    fn git_command(&self) -> Command {
        // TODO: add support for the GIT environment variable?
        let mut command = Command::new("git");
        command.current_dir(self.root).stderr(Stdio::inherit());
        command
    }

    /// Gets the scratch worktree if it exists, or initializes it if it doesn't.
    ///
    /// The scratch worktree is meant to be persistent across invocations of `x`. This is done for
    /// performance reasons.
    fn get_or_init_scratch(&self, hash: &GitHash) -> Result<PathBuf> {
        let mut scratch_dir = self.root.join("target");
        scratch_dir.extend(&["x-scratch", "tree"]);

        if scratch_dir.is_dir() && self.is_git_repo(&scratch_dir)? {
            debug!(
                "Using existing scratch worktree at {}",
                scratch_dir.display()
            );

            // Check out the given hash in the scratch worktree.
            let output = self
                .git_command()
                .current_dir(&scratch_dir)
                // TODO: also git clean?
                .args(&["reset", &format!("{:x}", hash), "--hard"])
                .output()
                .map_err(|err| SystemError::io("running git checkout in scratch tree", err))?;
            if !output.status.success() {
                return Err(SystemError::Exec {
                    cmd: "git checkout",
                    status: output.status,
                });
            }
        } else {
            if scratch_dir.is_dir() {
                std::fs::remove_dir_all(&scratch_dir)
                    .map_err(|err| SystemError::io("cleaning old scratch_dir", err))?;
            }

            // Try creating a scratch worktree at that location.
            info!("Setting up scratch worktree in {}", scratch_dir.display());
            let output = self
                .git_command()
                .args(&["worktree", "add"])
                .arg(&scratch_dir)
                .args(&[&format!("{:x}", hash), "--detach"])
                .output()
                .map_err(|err| SystemError::io("running git worktree add", err))?;
            if !output.status.success() {
                return Err(SystemError::Exec {
                    cmd: "git worktree add",
                    status: output.status,
                });
            }
        }

        // TODO: some sort of cross-process locking may be necessary in the future. Don't worry
        // about it for now.
        Ok(scratch_dir)
    }

    pub fn is_git_repo(&self, dir: &Path) -> Result<bool> {
        let output = self
            .git_command()
            .current_dir(dir)
            .args(&["rev-parse", "--git-dir"])
            .output()
            .map_err(|err| SystemError::io("checking if a directory is a git repo", err))?;

        Ok(output.status.success())
    }
}

/// A Git hash.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct GitHash([u8; 20]);

impl GitHash {
    /// Creates a new Git hash from a hex-encoded string.
    pub fn from_hex(hex: impl AsRef<[u8]>) -> Result<Self> {
        let hex = hex.as_ref();
        Ok(GitHash(hex::FromHex::from_hex(hex).map_err(|err| {
            SystemError::from_hex(format!("parsing a Git hash: {:?}", hex), err)
        })?))
    }
}

impl<'a, 'b> From<&'a GitHash> for Cow<'b, OsStr> {
    fn from(git_hash: &'a GitHash) -> Cow<'b, OsStr> {
        OsString::from(format!("{:x}", git_hash)).into()
    }
}

impl fmt::LowerHex for GitHash {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", hex::encode(&self.0))
    }
}
<|endoftext|>