<commit_before>extern crate klee;

#[test]
fn basic_test() {
  let a : i32 = 0;
  klee::symbol(&a, "a");
  assert_eq!(a, 56);
}

#[test]
fn other_test() {
  let a : i32 = 0;
  let b : i32 = 0;

  klee::symbol(&a, "a");
  klee::symbol(&b, "b");

  if a == 50 {
    if b == 49 {
      assert!("bugger");
    }
  }
}

#[test]
fn another_test() {
  let a = klee::i32("a");
  if a > 60 && a < 90 {
    let b = a + 40;
    assert_eq!(b, 0);

    if b == 12 {
      assert!("This path should be unreachable");
    }

    assert_eq!(b, 101);
    assert_eq!(b, 150000);
  }
}
<commit_msg>Update simple.rs<commit_after>extern crate klee;

#[test]
fn basic_test() {
  let a : i32 = 0;
  klee::symbol(&a, "a");
  assert_eq!(a, 56);
}

#[test]
fn other_test() {
  let a : i32 = 0;
  let b : i32 = 0;

  klee::symbol(&a, "a");
  klee::symbol(&b, "b");

  if a == 50 {
    if b == 49 {
      assert!("bugger");
    }
  }
}

#[test]
fn another_test() {
  let a = klee::i32("a");
  if a > 60 && a < 90 {
    let b = a + 40;
    assert_eq!(b, 0);

    if b == 12 {
      panic!("This path should be unreachable");
    }

    assert_eq!(b, 101);
    assert_eq!(b, 150000);
  }
}
<|endoftext|>


<commit_before>// use std::io;
// use std::io::Write;

use rustbox::{RustBox};
use termbox_sys::tb_change_cell;
use termbox_sys;

use tui::style;

/// A single line added to the widget. May be rendered as multiple lines on the
/// screen.
#[derive(Debug)]
pub struct Line {
    /// Note that this String may not be directly renderable - TODO: explain.
    str       : String,

    /// Number of _visible_ (i.e. excludes color encodings) characters in the
    /// line.
    len_chars : i32,

    /// Visible char indexes (not counting color encodings) of split positions
    /// of the string - when the line doesn't fit into the screen we split it
    /// into multiple lines using these.
    ///
    /// It's important that these are really indices ignoring invisible chars,
    /// as we use difference between two indices in this vector as length of
    /// substrings.
    splits    : Vec<i32>,
}

impl Line {
    pub fn new() -> Line {
        Line {
            str: String::new(),
            len_chars: 0,
            splits: Vec::new(),
        }
    }

    pub fn add_text(&mut self, str : &str) {
        self.str.reserve(str.len());

        let mut iter = str.chars();
        while let Some(char) = iter.next() {
            if char == style::COLOR_PREFIX {
                self.str.push(char);
                // read fg
                self.str.push(iter.next().unwrap());
                self.str.push(iter.next().unwrap());
                if let Some(mb_comma) = iter.next() {
                    self.str.push(mb_comma);
                    if mb_comma == ',' {
                        // read bg
                        self.str.push(iter.next().unwrap());
                        self.str.push(iter.next().unwrap());
                    }
                }
            }

            else if char == style::RESET_PREFIX || char == style::BOLD_PREFIX {
                self.str.push(char);
            }

            // Ignore some chars that break the rendering. These are used by the
            // protocol.
            else if char > '\x07' {
                self.str.push(char);
                if char.is_whitespace() {
                    self.splits.push(self.len_chars);
                }
                self.len_chars += 1;
            }
        }
    }

    pub fn add_char(&mut self, char : char) {
        assert!(char != style::COLOR_PREFIX);
        if char.is_whitespace() {
            self.splits.push(self.len_chars);
        }
        self.str.push(char);
        self.len_chars += 1;
    }

    pub fn len_chars(&self) -> i32 {
        self.len_chars
    }

    /// How many lines does this take when rendered? O(n) where n = number of
    /// split positions in the lines (i.e.  whitespaces).
    pub fn rendered_height(&self, width : i32) -> i32 {
        let mut lines : i32 = 1;
        let mut line_start : i32 = 0;

        for split_idx in 0 .. self.splits.len() {
            let char_idx = *unsafe { self.splits.get_unchecked(split_idx) };
            // writeln!(io::stderr(), "rendered_height: char_idx: {}", char_idx);
            let col = char_idx - line_start;

            // How many more chars can we render in this line?
            let slots_in_line : i32 = width - (col + 1);

            // How many chars do we need to render if until the next split
            // point?
            let chars_until_next_split : i32 =
                // -1 becuase we don't need to render the space or EOL.
                *self.splits.get(split_idx + 1).unwrap_or(&self.len_chars) - 1 - char_idx;

            // writeln!(io::stderr(),
            //          "rendered_height: slots_in_line: {}, chars_until_next_split: {}",
            //          slots_in_line, chars_until_next_split);

            if (chars_until_next_split as i32) > slots_in_line {
                // writeln!(io::stderr(), "splitting at {}", char_idx);
                lines += 1;
                line_start = char_idx + 1;
            }
        }

        lines
    }

    #[inline]
    pub fn draw(&self, rustbox : &RustBox, pos_x : i32, pos_y : i32, width : i32) {
        self.draw_from(rustbox, pos_x, pos_y, 0, width);
    }

    pub fn draw_from(&self, _ : &RustBox, pos_x : i32, pos_y : i32, first_line : i32, width : i32) {
        let mut col = pos_x;
        let mut line = 0;

        let mut next_split_idx : usize = 0;

        let mut char_idx : i32 = 0;

        let mut fg : u16 = 0;
        let mut bg : u16 = 0;

        let mut iter = self.str.chars();
        while let Some(mut char) = iter.next() {
            if char == style::COLOR_PREFIX {
                let fg_1 = to_dec(iter.next().unwrap()) as u16;
                let fg_2 = to_dec(iter.next().unwrap()) as u16;
                // We 'or' here as 'fg' can have 'bold' value
                fg |= fg_1 * 10 + fg_2;

                if let Some(char_) = iter.next() {
                    if char_ == ',' {
                        let bg_1 = to_dec(iter.next().unwrap()) as u16;
                        let bg_2 = to_dec(iter.next().unwrap()) as u16;
                        bg = bg_1 * 10 + bg_2;
                        continue;
                    } else {
                        bg = 0;
                        char = char_;
                    }
                }
            } else if char == style::BOLD_PREFIX {
                fg |= termbox_sys::TB_BOLD;
                continue;
            } else if char == style::RESET_PREFIX {
                fg = 0;
                bg = 0;
                continue;
            }

            if char.is_whitespace() {
                // We may want to move to the next line
                next_split_idx += 1;
                let next_split = self.splits.get(next_split_idx).unwrap_or(&self.len_chars);

                // How many more chars can we render in this line?
                let slots_in_line = width - (col - pos_x);

                // How many chars do we need to render if until the next
                // split point?
                assert!(*next_split > char_idx);
                let chars_until_next_split : i32 = *next_split - char_idx;

                // writeln!(io::stderr(), "chars_until_next_split: {}, slots_in_line: {}",
                //          chars_until_next_split, slots_in_line);

                if (chars_until_next_split as i32) <= slots_in_line {
                    // keep rendering chars
                    if line >= first_line {
                        unsafe { tb_change_cell(col, pos_y + line, char as u32, fg, bg); }
                    }
                    col += 1;
                } else {
                    // need to split here. ignore whitespace char.
                    line += 1;
                    col = pos_x;
                }

                char_idx += 1;
            }

            else {
                // Not possible to split. Need to make sure we don't render out
                // of bounds.
                if col - pos_x < width {
                    if line >= first_line {
                        unsafe { tb_change_cell(col, pos_y + line, char as u32, fg, bg); }
                    }
                    col += 1;
                }

                char_idx += 1;
            }
        }
    }
}

#[inline]
pub fn to_dec(ch : char) -> i8 {
    ((ch as u32) - ('0' as u32)) as i8
}

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

#[cfg(test)]
mod tests {

extern crate test;

use self::test::Bencher;
use std::fs::File;
use std::io::Read;
use super::*;

use tui::style;

#[test]
fn height_test_1() {
    let mut line = Line::new();
    line.add_text("a b c d e");
    assert_eq!(line.rendered_height(1), 5);
    assert_eq!(line.rendered_height(2), 5);
    assert_eq!(line.rendered_height(3), 3);
    assert_eq!(line.rendered_height(4), 3);
    assert_eq!(line.rendered_height(5), 2);
    assert_eq!(line.rendered_height(6), 2);
    assert_eq!(line.rendered_height(7), 2);
    assert_eq!(line.rendered_height(8), 2);
    assert_eq!(line.rendered_height(9), 1);
}

#[test]
fn height_test_2() {
    let mut line = Line::new();
    line.add_text("ab c d e");
    assert_eq!(line.rendered_height(1), 4);
    assert_eq!(line.rendered_height(2), 4);
    assert_eq!(line.rendered_height(3), 3);
    assert_eq!(line.rendered_height(4), 2);
    assert_eq!(line.rendered_height(5), 2);
    assert_eq!(line.rendered_height(6), 2);
    assert_eq!(line.rendered_height(7), 2);
    assert_eq!(line.rendered_height(8), 1);
}

#[test]
fn height_test_3() {
    let mut line = Line::new();
    line.add_text("ab cd e");
    assert_eq!(line.rendered_height(1), 3);
    assert_eq!(line.rendered_height(2), 3);
    assert_eq!(line.rendered_height(3), 3);
    assert_eq!(line.rendered_height(4), 2);
    assert_eq!(line.rendered_height(5), 2);
    assert_eq!(line.rendered_height(6), 2);
    assert_eq!(line.rendered_height(7), 1);
}

#[test]
fn height_test_4() {
    let mut line = Line::new();
    line.add_text("ab cde");
    assert_eq!(line.rendered_height(1), 2);
    assert_eq!(line.rendered_height(2), 2);
    assert_eq!(line.rendered_height(3), 2);
    assert_eq!(line.rendered_height(4), 2);
    assert_eq!(line.rendered_height(5), 2);
    assert_eq!(line.rendered_height(6), 1);
}

#[test]
fn height_test_5() {
    let mut line = Line::new();
    line.add_text("abcde");
    for i in 0 .. 6 {
        assert_eq!(line.rendered_height(i), 1);
    }
}

#[bench]
fn bench_rendered_height(b : &mut Bencher) {

    // 1160 words, 2,237 ns/iter (+/- 150)

    let mut text = String::new();
    {
        let mut file = File::open("test/lipsum.txt").unwrap();
        file.read_to_string(&mut text);
    }

    let mut line = Line::new();
    line.add_text(&text);
    b.iter(|| {
        line.rendered_height(1)
    });
}

} // mod tests
<commit_msg>some warnings in bench<commit_after>// use std::io;
// use std::io::Write;

use rustbox::{RustBox};
use termbox_sys::tb_change_cell;
use termbox_sys;

use tui::style;

/// A single line added to the widget. May be rendered as multiple lines on the
/// screen.
#[derive(Debug)]
pub struct Line {
    /// Note that this String may not be directly renderable - TODO: explain.
    str       : String,

    /// Number of _visible_ (i.e. excludes color encodings) characters in the
    /// line.
    len_chars : i32,

    /// Visible char indexes (not counting color encodings) of split positions
    /// of the string - when the line doesn't fit into the screen we split it
    /// into multiple lines using these.
    ///
    /// It's important that these are really indices ignoring invisible chars,
    /// as we use difference between two indices in this vector as length of
    /// substrings.
    splits    : Vec<i32>,
}

impl Line {
    pub fn new() -> Line {
        Line {
            str: String::new(),
            len_chars: 0,
            splits: Vec::new(),
        }
    }

    pub fn add_text(&mut self, str : &str) {
        self.str.reserve(str.len());

        let mut iter = str.chars();
        while let Some(char) = iter.next() {
            if char == style::COLOR_PREFIX {
                self.str.push(char);
                // read fg
                self.str.push(iter.next().unwrap());
                self.str.push(iter.next().unwrap());
                if let Some(mb_comma) = iter.next() {
                    self.str.push(mb_comma);
                    if mb_comma == ',' {
                        // read bg
                        self.str.push(iter.next().unwrap());
                        self.str.push(iter.next().unwrap());
                    }
                }
            }

            else if char == style::RESET_PREFIX || char == style::BOLD_PREFIX {
                self.str.push(char);
            }

            // Ignore some chars that break the rendering. These are used by the
            // protocol.
            else if char > '\x07' {
                self.str.push(char);
                if char.is_whitespace() {
                    self.splits.push(self.len_chars);
                }
                self.len_chars += 1;
            }
        }
    }

    pub fn add_char(&mut self, char : char) {
        assert!(char != style::COLOR_PREFIX);
        if char.is_whitespace() {
            self.splits.push(self.len_chars);
        }
        self.str.push(char);
        self.len_chars += 1;
    }

    pub fn len_chars(&self) -> i32 {
        self.len_chars
    }

    /// How many lines does this take when rendered? O(n) where n = number of
    /// split positions in the lines (i.e.  whitespaces).
    pub fn rendered_height(&self, width : i32) -> i32 {
        let mut lines : i32 = 1;
        let mut line_start : i32 = 0;

        for split_idx in 0 .. self.splits.len() {
            let char_idx = *unsafe { self.splits.get_unchecked(split_idx) };
            // writeln!(io::stderr(), "rendered_height: char_idx: {}", char_idx);
            let col = char_idx - line_start;

            // How many more chars can we render in this line?
            let slots_in_line : i32 = width - (col + 1);

            // How many chars do we need to render if until the next split
            // point?
            let chars_until_next_split : i32 =
                // -1 becuase we don't need to render the space or EOL.
                *self.splits.get(split_idx + 1).unwrap_or(&self.len_chars) - 1 - char_idx;

            // writeln!(io::stderr(),
            //          "rendered_height: slots_in_line: {}, chars_until_next_split: {}",
            //          slots_in_line, chars_until_next_split);

            if (chars_until_next_split as i32) > slots_in_line {
                // writeln!(io::stderr(), "splitting at {}", char_idx);
                lines += 1;
                line_start = char_idx + 1;
            }
        }

        lines
    }

    #[inline]
    pub fn draw(&self, rustbox : &RustBox, pos_x : i32, pos_y : i32, width : i32) {
        self.draw_from(rustbox, pos_x, pos_y, 0, width);
    }

    pub fn draw_from(&self, _ : &RustBox, pos_x : i32, pos_y : i32, first_line : i32, width : i32) {
        let mut col = pos_x;
        let mut line = 0;

        let mut next_split_idx : usize = 0;

        let mut char_idx : i32 = 0;

        let mut fg : u16 = 0;
        let mut bg : u16 = 0;

        let mut iter = self.str.chars();
        while let Some(mut char) = iter.next() {
            if char == style::COLOR_PREFIX {
                let fg_1 = to_dec(iter.next().unwrap()) as u16;
                let fg_2 = to_dec(iter.next().unwrap()) as u16;
                // We 'or' here as 'fg' can have 'bold' value
                fg |= fg_1 * 10 + fg_2;

                if let Some(char_) = iter.next() {
                    if char_ == ',' {
                        let bg_1 = to_dec(iter.next().unwrap()) as u16;
                        let bg_2 = to_dec(iter.next().unwrap()) as u16;
                        bg = bg_1 * 10 + bg_2;
                        continue;
                    } else {
                        bg = 0;
                        char = char_;
                    }
                }
            } else if char == style::BOLD_PREFIX {
                fg |= termbox_sys::TB_BOLD;
                continue;
            } else if char == style::RESET_PREFIX {
                fg = 0;
                bg = 0;
                continue;
            }

            if char.is_whitespace() {
                // We may want to move to the next line
                next_split_idx += 1;
                let next_split = self.splits.get(next_split_idx).unwrap_or(&self.len_chars);

                // How many more chars can we render in this line?
                let slots_in_line = width - (col - pos_x);

                // How many chars do we need to render if until the next
                // split point?
                assert!(*next_split > char_idx);
                let chars_until_next_split : i32 = *next_split - char_idx;

                // writeln!(io::stderr(), "chars_until_next_split: {}, slots_in_line: {}",
                //          chars_until_next_split, slots_in_line);

                if (chars_until_next_split as i32) <= slots_in_line {
                    // keep rendering chars
                    if line >= first_line {
                        unsafe { tb_change_cell(col, pos_y + line, char as u32, fg, bg); }
                    }
                    col += 1;
                } else {
                    // need to split here. ignore whitespace char.
                    line += 1;
                    col = pos_x;
                }

                char_idx += 1;
            }

            else {
                // Not possible to split. Need to make sure we don't render out
                // of bounds.
                if col - pos_x < width {
                    if line >= first_line {
                        unsafe { tb_change_cell(col, pos_y + line, char as u32, fg, bg); }
                    }
                    col += 1;
                }

                char_idx += 1;
            }
        }
    }
}

#[inline]
pub fn to_dec(ch : char) -> i8 {
    ((ch as u32) - ('0' as u32)) as i8
}

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

#[cfg(test)]
mod tests {

extern crate test;

use self::test::Bencher;
use std::fs::File;
use std::io::Read;
use super::*;

#[test]
fn height_test_1() {
    let mut line = Line::new();
    line.add_text("a b c d e");
    assert_eq!(line.rendered_height(1), 5);
    assert_eq!(line.rendered_height(2), 5);
    assert_eq!(line.rendered_height(3), 3);
    assert_eq!(line.rendered_height(4), 3);
    assert_eq!(line.rendered_height(5), 2);
    assert_eq!(line.rendered_height(6), 2);
    assert_eq!(line.rendered_height(7), 2);
    assert_eq!(line.rendered_height(8), 2);
    assert_eq!(line.rendered_height(9), 1);
}

#[test]
fn height_test_2() {
    let mut line = Line::new();
    line.add_text("ab c d e");
    assert_eq!(line.rendered_height(1), 4);
    assert_eq!(line.rendered_height(2), 4);
    assert_eq!(line.rendered_height(3), 3);
    assert_eq!(line.rendered_height(4), 2);
    assert_eq!(line.rendered_height(5), 2);
    assert_eq!(line.rendered_height(6), 2);
    assert_eq!(line.rendered_height(7), 2);
    assert_eq!(line.rendered_height(8), 1);
}

#[test]
fn height_test_3() {
    let mut line = Line::new();
    line.add_text("ab cd e");
    assert_eq!(line.rendered_height(1), 3);
    assert_eq!(line.rendered_height(2), 3);
    assert_eq!(line.rendered_height(3), 3);
    assert_eq!(line.rendered_height(4), 2);
    assert_eq!(line.rendered_height(5), 2);
    assert_eq!(line.rendered_height(6), 2);
    assert_eq!(line.rendered_height(7), 1);
}

#[test]
fn height_test_4() {
    let mut line = Line::new();
    line.add_text("ab cde");
    assert_eq!(line.rendered_height(1), 2);
    assert_eq!(line.rendered_height(2), 2);
    assert_eq!(line.rendered_height(3), 2);
    assert_eq!(line.rendered_height(4), 2);
    assert_eq!(line.rendered_height(5), 2);
    assert_eq!(line.rendered_height(6), 1);
}

#[test]
fn height_test_5() {
    let mut line = Line::new();
    line.add_text("abcde");
    for i in 0 .. 6 {
        assert_eq!(line.rendered_height(i), 1);
    }
}

#[bench]
fn bench_rendered_height(b : &mut Bencher) {

    // 1160 words, 2,237 ns/iter (+/- 150)

    let mut text = String::new();
    {
        let mut file = File::open("test/lipsum.txt").unwrap();
        file.read_to_string(&mut text).unwrap();
    }

    let mut line = Line::new();
    line.add_text(&text);
    b.iter(|| {
        line.rendered_height(1)
    });
}

} // mod tests
<|endoftext|>


<commit_before>// See LICENSE file for copyright and license details.

use std;
use gl;
use gl::types::{GLint, GLchar, GLuint, GLenum};
use cgmath::matrix::Matrix4;
use cgmath::array::Array2;
use error_context;
use core::types::MInt;
use core::misc::read_file;
use visualizer::types::{MatId, MFloat, ColorId, Color4};

pub struct Shader {
    id: GLuint,
}

impl Shader {
    pub fn new(vs_path: &Path, fs_path: &Path) -> Shader {
        set_error_context!("loading shader", vs_path.as_str().unwrap()); // TODO
        compile_program(
            read_file(vs_path).as_slice(),
            read_file(fs_path).as_slice(),
        )
    }

    pub fn activate(&self) {
        verify!(gl::UseProgram(self.id));
    }

    pub fn enable_attr(&self, name: &str, components_count: MInt) {
        let mut attr_id = 0;
        name.with_c_str(|name| {
            unsafe {
                attr_id = verify!(gl::GetAttribLocation(self.id, name));
            }
        });
        verify!(gl::EnableVertexAttribArray(attr_id as GLuint));
        let normalized = gl::FALSE;
        let stride = 0;
        unsafe {
            verify!(gl::VertexAttribPointer(
                attr_id as GLuint,
                components_count,
                gl::FLOAT,
                normalized,
                stride,
                std::ptr::null(),
            ));
        }
    }

    pub fn uniform_mat4f(&self, mat_id: MatId, mat: &Matrix4<MFloat>) {
        unsafe {
            verify!(gl::UniformMatrix4fv(
                mat_id.id as MInt, 1, gl::FALSE, mat.ptr()));
        }
    }

    pub fn uniform_color(&self, color_id: ColorId, color: Color4) {
        unsafe {
            let data_ptr = std::mem::transmute(&color);
            verify!(gl::Uniform4fv(color_id.id as MInt, 1, data_ptr));
        }
    }

    pub fn get_uniform(&self, name: &str) -> GLuint {
        name.with_c_str(|name| {
            unsafe {
                verify!(gl::GetUniformLocation(self.id, name) as GLuint)
            }
        })
    }
}

impl Drop for Shader {
    fn drop(&mut self) {
        verify!(gl::DeleteProgram(self.id));
    }
}

fn compile_shader(src: &str, shader_type: GLenum) -> GLuint {
    let shader = verify!(gl::CreateShader(shader_type));
    unsafe {
        src.with_c_str(|src| {
            verify!(gl::ShaderSource(shader, 1, &src, std::ptr::null()));
        });
        verify!(gl::CompileShader(shader));
        let mut status = gl::FALSE as GLint;
        verify!(gl::GetShaderiv(shader, gl::COMPILE_STATUS, &mut status));
        if status != (gl::TRUE as GLint) {
            let mut len = 0;
            verify!(gl::GetShaderiv(shader, gl::INFO_LOG_LENGTH, &mut len));
            // subtract 1 to skip the trailing null character
            let mut buf = Vec::from_elem(len as uint - 1, 0u8);
            verify!(gl::GetShaderInfoLog(shader, len, std::ptr::mut_null(),
                buf.as_mut_ptr() as *mut GLchar));
            fail!(
                "compile_shader: {}",
                std::str::raw::from_utf8(buf.as_slice())
            );
        }
    }
    shader
}

fn link_program(vertex_shader: GLuint, fragment_shader: GLuint) -> GLuint {
    let program = verify!(gl::CreateProgram());
    verify!(gl::AttachShader(program, vertex_shader));
    verify!(gl::AttachShader(program, fragment_shader));
    verify!(gl::LinkProgram(program));
    unsafe {
        let mut status = gl::FALSE as GLint;
        verify!(gl::GetProgramiv(program, gl::LINK_STATUS, &mut status));
        if status != (gl::TRUE as GLint) {
            let mut len = 0;
            verify!(gl::GetProgramiv(program, gl::INFO_LOG_LENGTH, &mut len));
            // subtract 1 to skip the trailing null character
            let mut buf = Vec::from_elem(len as uint - 1, 0u8);
            verify!(gl::GetProgramInfoLog(program, len, std::ptr::mut_null(),
                buf.as_mut_ptr() as *mut GLchar));
            fail!(
                "link_program: {}",
                std::str::raw::from_utf8(buf.as_slice())
            );
        }
    }
    program
}

fn compile_program(vertex_shader_src: &str, frag_shader_src: &str) -> Shader {
    let vertex_shader = compile_shader(
        vertex_shader_src, gl::VERTEX_SHADER);
    let fragment_shader = compile_shader(
        frag_shader_src, gl::FRAGMENT_SHADER);
    let program = link_program(vertex_shader, fragment_shader);
    // mark shaders for deletion after program deletion
    verify!(gl::DeleteShader(fragment_shader));
    verify!(gl::DeleteShader(vertex_shader));
    Shader{id: program}
}

// vim: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
<commit_msg>Shader::new: Fixed set_error_context! usage<commit_after>// See LICENSE file for copyright and license details.

use std;
use gl;
use gl::types::{GLint, GLchar, GLuint, GLenum};
use cgmath::matrix::Matrix4;
use cgmath::array::Array2;
use error_context;
use core::types::MInt;
use core::misc::read_file;
use visualizer::types::{MatId, MFloat, ColorId, Color4};

pub struct Shader {
    id: GLuint,
}

impl Shader {
    pub fn new(vs_path: &Path, fs_path: &Path) -> Shader {
        set_error_context!("loading shader (vs)", vs_path.as_str().unwrap());
        set_error_context!("loading shader (fs)", fs_path.as_str().unwrap());
        compile_program(
            read_file(vs_path).as_slice(),
            read_file(fs_path).as_slice(),
        )
    }

    pub fn activate(&self) {
        verify!(gl::UseProgram(self.id));
    }

    pub fn enable_attr(&self, name: &str, components_count: MInt) {
        let mut attr_id = 0;
        name.with_c_str(|name| {
            unsafe {
                attr_id = verify!(gl::GetAttribLocation(self.id, name));
            }
        });
        verify!(gl::EnableVertexAttribArray(attr_id as GLuint));
        let normalized = gl::FALSE;
        let stride = 0;
        unsafe {
            verify!(gl::VertexAttribPointer(
                attr_id as GLuint,
                components_count,
                gl::FLOAT,
                normalized,
                stride,
                std::ptr::null(),
            ));
        }
    }

    pub fn uniform_mat4f(&self, mat_id: MatId, mat: &Matrix4<MFloat>) {
        unsafe {
            verify!(gl::UniformMatrix4fv(
                mat_id.id as MInt, 1, gl::FALSE, mat.ptr()));
        }
    }

    pub fn uniform_color(&self, color_id: ColorId, color: Color4) {
        unsafe {
            let data_ptr = std::mem::transmute(&color);
            verify!(gl::Uniform4fv(color_id.id as MInt, 1, data_ptr));
        }
    }

    pub fn get_uniform(&self, name: &str) -> GLuint {
        name.with_c_str(|name| {
            unsafe {
                verify!(gl::GetUniformLocation(self.id, name) as GLuint)
            }
        })
    }
}

impl Drop for Shader {
    fn drop(&mut self) {
        verify!(gl::DeleteProgram(self.id));
    }
}

fn compile_shader(src: &str, shader_type: GLenum) -> GLuint {
    let shader = verify!(gl::CreateShader(shader_type));
    unsafe {
        src.with_c_str(|src| {
            verify!(gl::ShaderSource(shader, 1, &src, std::ptr::null()));
        });
        verify!(gl::CompileShader(shader));
        let mut status = gl::FALSE as GLint;
        verify!(gl::GetShaderiv(shader, gl::COMPILE_STATUS, &mut status));
        if status != (gl::TRUE as GLint) {
            let mut len = 0;
            verify!(gl::GetShaderiv(shader, gl::INFO_LOG_LENGTH, &mut len));
            // subtract 1 to skip the trailing null character
            let mut buf = Vec::from_elem(len as uint - 1, 0u8);
            verify!(gl::GetShaderInfoLog(shader, len, std::ptr::mut_null(),
                buf.as_mut_ptr() as *mut GLchar));
            fail!(
                "compile_shader: {}",
                std::str::raw::from_utf8(buf.as_slice())
            );
        }
    }
    shader
}

fn link_program(vertex_shader: GLuint, fragment_shader: GLuint) -> GLuint {
    let program = verify!(gl::CreateProgram());
    verify!(gl::AttachShader(program, vertex_shader));
    verify!(gl::AttachShader(program, fragment_shader));
    verify!(gl::LinkProgram(program));
    unsafe {
        let mut status = gl::FALSE as GLint;
        verify!(gl::GetProgramiv(program, gl::LINK_STATUS, &mut status));
        if status != (gl::TRUE as GLint) {
            let mut len = 0;
            verify!(gl::GetProgramiv(program, gl::INFO_LOG_LENGTH, &mut len));
            // subtract 1 to skip the trailing null character
            let mut buf = Vec::from_elem(len as uint - 1, 0u8);
            verify!(gl::GetProgramInfoLog(program, len, std::ptr::mut_null(),
                buf.as_mut_ptr() as *mut GLchar));
            fail!(
                "link_program: {}",
                std::str::raw::from_utf8(buf.as_slice())
            );
        }
    }
    program
}

fn compile_program(vertex_shader_src: &str, frag_shader_src: &str) -> Shader {
    let vertex_shader = compile_shader(
        vertex_shader_src, gl::VERTEX_SHADER);
    let fragment_shader = compile_shader(
        frag_shader_src, gl::FRAGMENT_SHADER);
    let program = link_program(vertex_shader, fragment_shader);
    // mark shaders for deletion after program deletion
    verify!(gl::DeleteShader(fragment_shader));
    verify!(gl::DeleteShader(vertex_shader));
    Shader{id: program}
}

// vim: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
<|endoftext|>


<commit_before>use std::time::SystemTime;

use anyhow::Context;
use aws_sigv4::{
    http_request::{self, SignableRequest, SigningSettings},
    SigningParams,
};
use reqwest::blocking::Client;
use retry::delay::{jitter, Exponential};
use retry::OperationResult;

#[derive(Clone)]
pub struct CloudFront {
    distribution_id: String,
    access_key: String,
    secret_key: String,
}

impl CloudFront {
    pub fn from_environment() -> Option<Self> {
        let distribution_id = dotenv::var("CLOUDFRONT_DISTRIBUTION").ok()?;
        let access_key = dotenv::var("AWS_ACCESS_KEY").expect("missing AWS_ACCESS_KEY");
        let secret_key = dotenv::var("AWS_SECRET_KEY").expect("missing AWS_SECRET_KEY");
        Some(Self {
            distribution_id,
            access_key,
            secret_key,
        })
    }

    /// Invalidate a file on CloudFront
    ///
    /// `path` is the path to the file to invalidate, such as `config.json`, or `re/ge/regex`
    #[instrument(skip(self, client))]
    pub fn invalidate(&self, client: &Client, path: &str) -> anyhow::Result<()> {
        let path = path.trim_start_matches('/');
        let url = format!(
            "https://cloudfront.amazonaws.com/2020-05-31/distribution/{}/invalidation",
            self.distribution_id
        );
        trace!(?url);

        let backoff = Exponential::from_millis(1000).map(jitter).take(2);
        retry::retry_with_index(backoff, |attempt| {
            let now = chrono::offset::Utc::now().timestamp_micros();
            let body = format!(
                r#"
<?xml version="1.0" encoding="UTF-8"?>
<InvalidationBatch xmlns="http://cloudfront.amazonaws.com/doc/2020-05-31/">
    <CallerReference>{now}</CallerReference>
    <Paths>
        <Items>
            <Path>/{path}</Path>
        </Items>
        <Quantity>1</Quantity>
    </Paths>
</InvalidationBatch>
"#
            );
            trace!(?body);

            let request = match http::Request::post(&url)
                .body(&body)
                .context("Failed to construct HTTP request")
            {
                Ok(request) => request,
                Err(error) => return OperationResult::Err(error),
            };

            trace!("Signing invalidation request");
            let request = SignableRequest::from(&request);
            let params = SigningParams::builder()
                .access_key(&self.access_key)
                .secret_key(&self.secret_key)
                .region("us-east-1") // cloudfront is a regionless service, use the default region for signing.
                .service_name("cloudfront")
                .settings(SigningSettings::default())
                .time(SystemTime::now())
                .build()
                .unwrap(); // all required fields are set

            let (mut signature_headers, _) =
                http_request::sign(request, &params).unwrap().into_parts();

            debug!(?attempt, "Sending invalidation request");
            let response = match client
                .post(&url)
                .headers(signature_headers.take_headers().unwrap_or_default())
                .body(body)
                .send()
                .context("Failed to send invalidation request")
            {
                Ok(response) => response,
                Err(error) => return OperationResult::Retry(error),
            };

            let status = response.status();

            let result = match response.error_for_status_ref() {
                Ok(_) => {
                    debug!(?status, "Invalidation request successful");
                    Ok(())
                }
                Err(error) => {
                    let headers = response.headers().clone();
                    let body = response.text();
                    debug!(?status, ?headers, ?body, "Invalidation request failed");

                    Err(error).with_context(|| format!("Failed to invalidate {path}"))
                }
            };

            match result {
                Ok(_) => OperationResult::Ok(()),
                Err(error) if status.is_server_error() => OperationResult::Retry(error),
                Err(error) => OperationResult::Err(error),
            }
        })
        .map_err(|error| error.error)
    }
}
<commit_msg>worker::cloudfront: Increase log level of invalidation request failure<commit_after>use std::time::SystemTime;

use anyhow::Context;
use aws_sigv4::{
    http_request::{self, SignableRequest, SigningSettings},
    SigningParams,
};
use reqwest::blocking::Client;
use retry::delay::{jitter, Exponential};
use retry::OperationResult;

#[derive(Clone)]
pub struct CloudFront {
    distribution_id: String,
    access_key: String,
    secret_key: String,
}

impl CloudFront {
    pub fn from_environment() -> Option<Self> {
        let distribution_id = dotenv::var("CLOUDFRONT_DISTRIBUTION").ok()?;
        let access_key = dotenv::var("AWS_ACCESS_KEY").expect("missing AWS_ACCESS_KEY");
        let secret_key = dotenv::var("AWS_SECRET_KEY").expect("missing AWS_SECRET_KEY");
        Some(Self {
            distribution_id,
            access_key,
            secret_key,
        })
    }

    /// Invalidate a file on CloudFront
    ///
    /// `path` is the path to the file to invalidate, such as `config.json`, or `re/ge/regex`
    #[instrument(skip(self, client))]
    pub fn invalidate(&self, client: &Client, path: &str) -> anyhow::Result<()> {
        let path = path.trim_start_matches('/');
        let url = format!(
            "https://cloudfront.amazonaws.com/2020-05-31/distribution/{}/invalidation",
            self.distribution_id
        );
        trace!(?url);

        let backoff = Exponential::from_millis(1000).map(jitter).take(2);
        retry::retry_with_index(backoff, |attempt| {
            let now = chrono::offset::Utc::now().timestamp_micros();
            let body = format!(
                r#"
<?xml version="1.0" encoding="UTF-8"?>
<InvalidationBatch xmlns="http://cloudfront.amazonaws.com/doc/2020-05-31/">
    <CallerReference>{now}</CallerReference>
    <Paths>
        <Items>
            <Path>/{path}</Path>
        </Items>
        <Quantity>1</Quantity>
    </Paths>
</InvalidationBatch>
"#
            );
            trace!(?body);

            let request = match http::Request::post(&url)
                .body(&body)
                .context("Failed to construct HTTP request")
            {
                Ok(request) => request,
                Err(error) => return OperationResult::Err(error),
            };

            trace!("Signing invalidation request");
            let request = SignableRequest::from(&request);
            let params = SigningParams::builder()
                .access_key(&self.access_key)
                .secret_key(&self.secret_key)
                .region("us-east-1") // cloudfront is a regionless service, use the default region for signing.
                .service_name("cloudfront")
                .settings(SigningSettings::default())
                .time(SystemTime::now())
                .build()
                .unwrap(); // all required fields are set

            let (mut signature_headers, _) =
                http_request::sign(request, &params).unwrap().into_parts();

            debug!(?attempt, "Sending invalidation request");
            let response = match client
                .post(&url)
                .headers(signature_headers.take_headers().unwrap_or_default())
                .body(body)
                .send()
                .context("Failed to send invalidation request")
            {
                Ok(response) => response,
                Err(error) => return OperationResult::Retry(error),
            };

            let status = response.status();

            let result = match response.error_for_status_ref() {
                Ok(_) => {
                    debug!(?status, "Invalidation request successful");
                    Ok(())
                }
                Err(error) => {
                    let headers = response.headers().clone();
                    let body = response.text();
                    warn!(?status, ?headers, ?body, "Invalidation request failed");

                    Err(error).with_context(|| format!("Failed to invalidate {path}"))
                }
            };

            match result {
                Ok(_) => OperationResult::Ok(()),
                Err(error) if status.is_server_error() => OperationResult::Retry(error),
                Err(error) => OperationResult::Err(error),
            }
        })
        .map_err(|error| error.error)
    }
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;
use std::{slice, mem};

use libc::{size_t, c_int, c_uint};

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

#[derive(Debug)]
struct Node {
    weight: size_t,
    leaf_counts: Vec<c_int>,
}

#[derive(Debug)]
struct Leaf {
    pub weight: size_t,
    pub index: size_t,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug)]
struct List {
    lookahead1: Node,
    lookahead2: Node,
    next_leaf_index: size_t,
}

#[no_mangle]
#[allow(non_snake_case)]
pub extern fn ZopfliLengthLimitedCodeLengths(frequencies: *const size_t, n: c_int, maxbits: c_int, bitlengths: *mut c_uint) -> c_int {
    let freqs = unsafe { slice::from_raw_parts(frequencies, n as usize) };
    let result = length_limited_code_lengths(freqs, maxbits);

    for (i, res) in result.into_iter().enumerate() {
        unsafe {
            *bitlengths.offset(i as isize) = res as c_uint;
        }
    }

    return 0;
}

pub fn length_limited_code_lengths(frequencies: &[size_t], maxbits: c_int) -> Vec<size_t> {
    let mut leaves = vec![];

    // Count used symbols and place them in the leaves.
    for (i, &freq) in frequencies.iter().enumerate() {
        if freq != 0 {
            leaves.push(Leaf { weight: freq, index: i });
        }
    }

    // Short circuit some special cases
    if leaves.is_empty() {
        // There are no non-zero frequencies.
        return vec![0; frequencies.len()];
    }
    if leaves.len() == 1 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        return result;
    }
    if leaves.len() == 2 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        result[leaves[1].index] = 1;
        return result;
    }

    // Sort the leaves from least frequent to most frequent.
    // Add index into the same variable for stable sorting.
    for leaf in leaves.iter_mut() {
        leaf.weight = (leaf.weight << 9) | leaf.index;
    }
    leaves.sort();
    for leaf in leaves.iter_mut() {
        leaf.weight >>= 9;
    }

    let mut lists = Vec::with_capacity(maxbits as usize);
    for _ in 0..maxbits {
        lists.push(List {
            lookahead1: Node { weight: leaves[0].weight, leaf_counts: vec![1] },
            lookahead2: Node { weight: leaves[1].weight, leaf_counts: vec![2] },
            next_leaf_index: 2,
        });
    }

    // In the last list, 2 * numsymbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = 2 * leaves.len() - 4;
    for _ in 0..num_boundary_pm_runs {
        boundary_pm_toplevel(&mut lists[..], &leaves);
    }

    let n = frequencies.len();
    let mut result = vec![0; n];

    let mut a = lists.pop().unwrap().lookahead2.leaf_counts.into_iter().rev().peekable();

    let mut bitlength_value = 1;
    while let Some(leaf_count) = a.next() {
        let next_count = *a.peek().unwrap_or(&0);
        for i in next_count..leaf_count {
            result[leaves[i as usize].index as usize] = bitlength_value;
        }
        bitlength_value += 1;
    }

    result
}

fn lowest_list(lists: &mut [List], leaves: &Vec<Leaf>) {
    // We're in the lowest list, just add another leaf to the lookaheads
    // There will always be more leaves to be added on level 0 so this is safe.
    let mut current_list = lists.get_mut(0).unwrap();
    let ref next_leaf = leaves[current_list.next_leaf_index];
    current_list.lookahead2.weight = next_leaf.weight;

    current_list.lookahead2.leaf_counts[0] = current_list.lookahead1.leaf_counts.last().unwrap() + 1;
    current_list.next_leaf_index += 1;
}

fn next_leaf(lists: &mut [List], leaves: &Vec<Leaf>, current_list_index: usize) {
    let mut current_list = lists.get_mut(current_list_index).unwrap();

    // The next leaf goes next; counting itself makes the leaf_count increase by one.
    current_list.lookahead2.weight = leaves[current_list.next_leaf_index].weight;
    current_list.lookahead2.leaf_counts.clear();
    current_list.lookahead2.leaf_counts.extend(current_list.lookahead1.leaf_counts.iter());
    let last_index = current_list.lookahead2.leaf_counts.len() - 1;
    current_list.lookahead2.leaf_counts[last_index] += 1;
    current_list.next_leaf_index += 1;
}

fn next_tree(weight_sum: size_t, lists: &mut [List], leaves: &Vec<Leaf>, current_list_index: usize) {
    let previous_list_leaf_counts = lists[current_list_index - 1].lookahead2.leaf_counts.clone();
    {
        let ref mut current_list = lists[current_list_index];

        // Make a tree from the lookaheads from the previous list; that goes next.
        // This is not a leaf node, so the leaf count stays the same.
        current_list.lookahead2.weight = weight_sum;
        current_list.lookahead2.leaf_counts.clear();
        current_list.lookahead2.leaf_counts.extend(previous_list_leaf_counts);
        current_list.lookahead2.leaf_counts.push(*current_list.lookahead1.leaf_counts.last().unwrap());
    }

    // The previous list needs two new lookahead nodes.
    boundary_pm(lists, leaves, current_list_index - 1);
    boundary_pm(lists, leaves, current_list_index - 1);
}

fn boundary_pm_toplevel(lists: &mut [List], leaves: &Vec<Leaf>) {
    let last_index = lists.len() - 1;
    boundary_pm(lists, leaves, last_index);
}

fn boundary_pm(lists: &mut [List], leaves: &Vec<Leaf>, current_list_index: usize) {
    let next_leaf_index = lists[current_list_index].next_leaf_index;

    if current_list_index == 0 && next_leaf_index == leaves.len() {
        // We've added all the leaves to the lowest list, so we're done here
        return;
    }

    mem::swap(&mut lists[current_list_index].lookahead1, &mut lists[current_list_index].lookahead2);

    if current_list_index == 0 {
        lowest_list(lists, leaves);
    } else {
        // We're at a list other than the lowest list.
        let weight_sum = {
            let previous_list = lists.get(current_list_index - 1).unwrap();
            previous_list.lookahead1.weight + previous_list.lookahead2.weight
        };

        if next_leaf_index < leaves.len() && weight_sum > leaves[next_leaf_index].weight {
            next_leaf(lists, leaves, current_list_index);
        } else {
            next_tree(weight_sum, lists, leaves, current_list_index);
        }
    }
}

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<commit_msg>Access raw pointer to get around borrow checker<commit_after>use std::cmp::Ordering;
use std::{slice, mem};

use libc::{size_t, c_int, c_uint};

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

#[derive(Debug)]
struct Node {
    weight: size_t,
    leaf_counts: Vec<c_int>,
}

#[derive(Debug)]
struct Leaf {
    pub weight: size_t,
    pub index: size_t,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug)]
struct List {
    lookahead1: Node,
    lookahead2: Node,
    next_leaf_index: size_t,
}

#[no_mangle]
#[allow(non_snake_case)]
pub extern fn ZopfliLengthLimitedCodeLengths(frequencies: *const size_t, n: c_int, maxbits: c_int, bitlengths: *mut c_uint) -> c_int {
    let freqs = unsafe { slice::from_raw_parts(frequencies, n as usize) };
    let result = length_limited_code_lengths(freqs, maxbits);

    for (i, res) in result.into_iter().enumerate() {
        unsafe {
            *bitlengths.offset(i as isize) = res as c_uint;
        }
    }

    return 0;
}

pub fn length_limited_code_lengths(frequencies: &[size_t], maxbits: c_int) -> Vec<size_t> {
    let mut leaves = vec![];

    // Count used symbols and place them in the leaves.
    for (i, &freq) in frequencies.iter().enumerate() {
        if freq != 0 {
            leaves.push(Leaf { weight: freq, index: i });
        }
    }

    // Short circuit some special cases
    if leaves.is_empty() {
        // There are no non-zero frequencies.
        return vec![0; frequencies.len()];
    }
    if leaves.len() == 1 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        return result;
    }
    if leaves.len() == 2 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        result[leaves[1].index] = 1;
        return result;
    }

    // Sort the leaves from least frequent to most frequent.
    // Add index into the same variable for stable sorting.
    for leaf in leaves.iter_mut() {
        leaf.weight = (leaf.weight << 9) | leaf.index;
    }
    leaves.sort();
    for leaf in leaves.iter_mut() {
        leaf.weight >>= 9;
    }

    let mut lists = Vec::with_capacity(maxbits as usize);
    for _ in 0..maxbits {
        lists.push(List {
            lookahead1: Node { weight: leaves[0].weight, leaf_counts: vec![1] },
            lookahead2: Node { weight: leaves[1].weight, leaf_counts: vec![2] },
            next_leaf_index: 2,
        });
    }

    // In the last list, 2 * numsymbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = 2 * leaves.len() - 4;
    for _ in 0..num_boundary_pm_runs {
        boundary_pm_toplevel(&mut lists[..], &leaves);
    }

    let n = frequencies.len();
    let mut result = vec![0; n];

    let mut a = lists.pop().unwrap().lookahead2.leaf_counts.into_iter().rev().peekable();

    let mut bitlength_value = 1;
    while let Some(leaf_count) = a.next() {
        let next_count = *a.peek().unwrap_or(&0);
        for i in next_count..leaf_count {
            result[leaves[i as usize].index as usize] = bitlength_value;
        }
        bitlength_value += 1;
    }

    result
}

fn lowest_list(lists: &mut [List], leaves: &Vec<Leaf>) {
    // We're in the lowest list, just add another leaf to the lookaheads
    // There will always be more leaves to be added on level 0 so this is safe.
    let mut current_list = lists.get_mut(0).unwrap();
    let ref next_leaf = leaves[current_list.next_leaf_index];
    current_list.lookahead2.weight = next_leaf.weight;

    current_list.lookahead2.leaf_counts[0] = current_list.lookahead1.leaf_counts.last().unwrap() + 1;
    current_list.next_leaf_index += 1;
}

fn next_leaf(lists: &mut [List], leaves: &Vec<Leaf>, current_list_index: usize) {
    let mut current_list = lists.get_mut(current_list_index).unwrap();

    // The next leaf goes next; counting itself makes the leaf_count increase by one.
    current_list.lookahead2.weight = leaves[current_list.next_leaf_index].weight;
    current_list.lookahead2.leaf_counts.clear();
    current_list.lookahead2.leaf_counts.extend(current_list.lookahead1.leaf_counts.iter());
    let last_index = current_list.lookahead2.leaf_counts.len() - 1;
    current_list.lookahead2.leaf_counts[last_index] += 1;
    current_list.next_leaf_index += 1;
}

fn next_tree(weight_sum: size_t, lists: &mut [List], leaves: &Vec<Leaf>, current_list_index: usize) {
    let num_leaf_counts = lists[current_list_index - 1].lookahead2.leaf_counts.len();
    let previous_list_leaf_counts = lists[current_list_index - 1].lookahead2.leaf_counts.as_ptr();
    {
        let ref mut current_list = lists[current_list_index];

        // Make a tree from the lookaheads from the previous list; that goes next.
        // This is not a leaf node, so the leaf count stays the same.
        current_list.lookahead2.weight = weight_sum;
        current_list.lookahead2.leaf_counts.clear();
        for i in 0..num_leaf_counts {
            current_list.lookahead2.leaf_counts.push(unsafe {
                *previous_list_leaf_counts.offset(i as isize)
            });
        }
        current_list.lookahead2.leaf_counts.push(*current_list.lookahead1.leaf_counts.last().unwrap());
    }

    // The previous list needs two new lookahead nodes.
    boundary_pm(lists, leaves, current_list_index - 1);
    boundary_pm(lists, leaves, current_list_index - 1);
}

fn boundary_pm_toplevel(lists: &mut [List], leaves: &Vec<Leaf>) {
    let last_index = lists.len() - 1;
    boundary_pm(lists, leaves, last_index);
}

fn boundary_pm(lists: &mut [List], leaves: &Vec<Leaf>, current_list_index: usize) {
    let next_leaf_index = lists[current_list_index].next_leaf_index;

    if current_list_index == 0 && next_leaf_index == leaves.len() {
        // We've added all the leaves to the lowest list, so we're done here
        return;
    }

    mem::swap(&mut lists[current_list_index].lookahead1, &mut lists[current_list_index].lookahead2);

    if current_list_index == 0 {
        lowest_list(lists, leaves);
    } else {
        // We're at a list other than the lowest list.
        let weight_sum = {
            let previous_list = lists.get(current_list_index - 1).unwrap();
            previous_list.lookahead1.weight + previous_list.lookahead2.weight
        };

        if next_leaf_index < leaves.len() && weight_sum > leaves[next_leaf_index].weight {
            next_leaf(lists, leaves, current_list_index);
        } else {
            next_tree(weight_sum, lists, leaves, current_list_index);
        }
    }
}

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;
use std::{slice, mem};

use libc::{size_t, c_int, c_uint};

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

#[derive(Debug)]
struct Node {
    weight: size_t,
    leaf_counts: Vec<c_int>,
}

#[derive(Debug)]
struct Leaf {
    pub weight: size_t,
    pub index: size_t,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug)]
struct List {
    lookahead1: Node,
    lookahead2: Node,
    next_leaf_index: size_t,
}

#[no_mangle]
#[allow(non_snake_case)]
pub extern fn ZopfliLengthLimitedCodeLengths(frequencies: *const size_t, n: c_int, maxbits: c_int, bitlengths: *mut c_uint) -> c_int {
    let freqs = unsafe { slice::from_raw_parts(frequencies, n as usize) };
    let result = length_limited_code_lengths(freqs, maxbits);

    for (i, res) in result.into_iter().enumerate() {
        unsafe {
            *bitlengths.offset(i as isize) = res as c_uint;
        }
    }

    return 0;
}

pub fn length_limited_code_lengths(frequencies: &[size_t], maxbits: c_int) -> Vec<size_t> {
    let mut leaves = vec![];

    // Count used symbols and place them in the leaves.
    for (i, &freq) in frequencies.iter().enumerate() {
        if freq != 0 {
            leaves.push(Leaf { weight: freq, index: i });
        }
    }

    // Short circuit some special cases
    if leaves.is_empty() {
        // There are no non-zero frequencies.
        return vec![0; frequencies.len()];
    }
    if leaves.len() == 1 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        return result;
    }
    if leaves.len() == 2 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        result[leaves[1].index] = 1;
        return result;
    }

    // Sort the leaves from least frequent to most frequent.
    // Add index into the same variable for stable sorting.
    for leaf in leaves.iter_mut() {
        leaf.weight = (leaf.weight << 9) | leaf.index;
    }
    leaves.sort();
    for leaf in leaves.iter_mut() {
        leaf.weight >>= 9;
    }

    let mut lists = Vec::with_capacity(maxbits as usize);
    for _ in 0..maxbits {
        lists.push(List {
            lookahead1: Node { weight: leaves[0].weight, leaf_counts: vec![1] },
            lookahead2: Node { weight: leaves[1].weight, leaf_counts: vec![2] },
            next_leaf_index: 2,
        });
    }

    // In the last list, 2 * numsymbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = 2 * leaves.len() - 4;
    for _ in 0..num_boundary_pm_runs {
        lists = boundary_pm(lists, &leaves);
    }

    let n = frequencies.len();
    let mut result = vec![0; n];

    let mut a = lists.pop().unwrap().lookahead2.leaf_counts.into_iter().rev().peekable();

    let mut bitlength_value = 1;
    while let Some(leaf_count) = a.next() {
        let next_count = *a.peek().unwrap_or(&0);
        for i in next_count..leaf_count {
            result[leaves[i as usize].index as usize] = bitlength_value;
        }
        bitlength_value += 1;
    }

    result
}

fn lowest_list(mut current_list: List, mut lists: Vec<List>, leaves: &Vec<Leaf>) -> Vec<List> {
    // We're in the lowest list, just add another leaf to the lookaheads
    // There will always be more leaves to be added on level 0 so this is safe.
    let ref next_leaf = leaves[current_list.next_leaf_index];
    current_list.lookahead2.weight = next_leaf.weight;

    current_list.lookahead2.leaf_counts[0] = current_list.lookahead1.leaf_counts.last().unwrap() + 1;
    current_list.next_leaf_index += 1;
    lists.push(current_list);
    lists
}

fn next_leaf(mut current_list: List, previous_list: List, mut lists: Vec<List>, leaves: &Vec<Leaf>) -> Vec<List> {
    // The next leaf goes next; counting itself makes the leaf_count increase by one.
    current_list.lookahead2.weight = leaves[current_list.next_leaf_index].weight;
    current_list.lookahead2.leaf_counts.clear();
    current_list.lookahead2.leaf_counts.extend(current_list.lookahead1.leaf_counts.iter());
    let last_index = current_list.lookahead2.leaf_counts.len() - 1;
    current_list.lookahead2.leaf_counts[last_index] += 1;
    current_list.next_leaf_index += 1;
    lists.push(previous_list);
    lists.push(current_list);
    lists
}

fn next_tree(mut current_list: List, previous_list: List, weight_sum: size_t, mut lists: Vec<List>, leaves: &Vec<Leaf>) -> Vec<List> {
    // Make a tree from the lookaheads from the previous list; that goes next.
    // This is not a leaf node, so the leaf count stays the same.

    current_list.lookahead2.weight = weight_sum;
    current_list.lookahead2.leaf_counts.clear();
    current_list.lookahead2.leaf_counts.extend(previous_list.lookahead2.leaf_counts.iter());
    current_list.lookahead2.leaf_counts.push(*current_list.lookahead1.leaf_counts.last().unwrap());

    // The previous list needs two new lookahead nodes.
    lists.push(previous_list);
    lists = boundary_pm(lists, leaves);
    lists = boundary_pm(lists, leaves);
    lists.push(current_list);

    lists
}

fn boundary_pm(mut lists: Vec<List>, leaves: &Vec<Leaf>) -> Vec<List> {
    let mut current_list = lists.pop().unwrap();
    if lists.is_empty() && current_list.next_leaf_index == leaves.len() {
        // We've added all the leaves to the lowest list, so we're done here
        lists.push(current_list);
        return lists;
    }

    mem::swap(&mut current_list.lookahead1, &mut current_list.lookahead2);

    if lists.is_empty() {
        lists = lowest_list(current_list, lists, leaves);
    } else {
        // We're at a list other than the lowest list.
        let previous_list = lists.pop().unwrap();
        let weight_sum = previous_list.lookahead1.weight + previous_list.lookahead2.weight;

        if current_list.next_leaf_index < leaves.len() && weight_sum > leaves[current_list.next_leaf_index].weight {
            lists = next_leaf(current_list, previous_list, lists, leaves);
        } else {
            lists = next_tree(current_list, previous_list, weight_sum, lists, leaves);
        }
    }

    lists
}

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<commit_msg>Pass mutable references of lists around instead of being functional<commit_after>use std::cmp::Ordering;
use std::{slice, mem};

use libc::{size_t, c_int, c_uint};

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

#[derive(Debug)]
struct Node {
    weight: size_t,
    leaf_counts: Vec<c_int>,
}

#[derive(Debug)]
struct Leaf {
    pub weight: size_t,
    pub index: size_t,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug)]
struct List {
    lookahead1: Node,
    lookahead2: Node,
    next_leaf_index: size_t,
}

#[no_mangle]
#[allow(non_snake_case)]
pub extern fn ZopfliLengthLimitedCodeLengths(frequencies: *const size_t, n: c_int, maxbits: c_int, bitlengths: *mut c_uint) -> c_int {
    let freqs = unsafe { slice::from_raw_parts(frequencies, n as usize) };
    let result = length_limited_code_lengths(freqs, maxbits);

    for (i, res) in result.into_iter().enumerate() {
        unsafe {
            *bitlengths.offset(i as isize) = res as c_uint;
        }
    }

    return 0;
}

pub fn length_limited_code_lengths(frequencies: &[size_t], maxbits: c_int) -> Vec<size_t> {
    let mut leaves = vec![];

    // Count used symbols and place them in the leaves.
    for (i, &freq) in frequencies.iter().enumerate() {
        if freq != 0 {
            leaves.push(Leaf { weight: freq, index: i });
        }
    }

    // Short circuit some special cases
    if leaves.is_empty() {
        // There are no non-zero frequencies.
        return vec![0; frequencies.len()];
    }
    if leaves.len() == 1 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        return result;
    }
    if leaves.len() == 2 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        result[leaves[1].index] = 1;
        return result;
    }

    // Sort the leaves from least frequent to most frequent.
    // Add index into the same variable for stable sorting.
    for leaf in leaves.iter_mut() {
        leaf.weight = (leaf.weight << 9) | leaf.index;
    }
    leaves.sort();
    for leaf in leaves.iter_mut() {
        leaf.weight >>= 9;
    }

    let mut lists = Vec::with_capacity(maxbits as usize);
    for _ in 0..maxbits {
        lists.push(List {
            lookahead1: Node { weight: leaves[0].weight, leaf_counts: vec![1] },
            lookahead2: Node { weight: leaves[1].weight, leaf_counts: vec![2] },
            next_leaf_index: 2,
        });
    }

    // In the last list, 2 * numsymbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = 2 * leaves.len() - 4;
    for _ in 0..num_boundary_pm_runs {
        boundary_pm(&mut lists, &leaves);
    }

    let n = frequencies.len();
    let mut result = vec![0; n];

    let mut a = lists.pop().unwrap().lookahead2.leaf_counts.into_iter().rev().peekable();

    let mut bitlength_value = 1;
    while let Some(leaf_count) = a.next() {
        let next_count = *a.peek().unwrap_or(&0);
        for i in next_count..leaf_count {
            result[leaves[i as usize].index as usize] = bitlength_value;
        }
        bitlength_value += 1;
    }

    result
}

fn lowest_list(mut current_list: List, lists: &mut Vec<List>, leaves: &Vec<Leaf>) {
    // We're in the lowest list, just add another leaf to the lookaheads
    // There will always be more leaves to be added on level 0 so this is safe.
    let ref next_leaf = leaves[current_list.next_leaf_index];
    current_list.lookahead2.weight = next_leaf.weight;

    current_list.lookahead2.leaf_counts[0] = current_list.lookahead1.leaf_counts.last().unwrap() + 1;
    current_list.next_leaf_index += 1;
    lists.push(current_list);
}

fn next_leaf(mut current_list: List, previous_list: List, lists: &mut Vec<List>, leaves: &Vec<Leaf>) {
    // The next leaf goes next; counting itself makes the leaf_count increase by one.
    current_list.lookahead2.weight = leaves[current_list.next_leaf_index].weight;
    current_list.lookahead2.leaf_counts.clear();
    current_list.lookahead2.leaf_counts.extend(current_list.lookahead1.leaf_counts.iter());
    let last_index = current_list.lookahead2.leaf_counts.len() - 1;
    current_list.lookahead2.leaf_counts[last_index] += 1;
    current_list.next_leaf_index += 1;
    lists.push(previous_list);
    lists.push(current_list);
}

fn next_tree(mut current_list: List, previous_list: List, weight_sum: size_t, lists: &mut Vec<List>, leaves: &Vec<Leaf>) {
    // Make a tree from the lookaheads from the previous list; that goes next.
    // This is not a leaf node, so the leaf count stays the same.

    current_list.lookahead2.weight = weight_sum;
    current_list.lookahead2.leaf_counts.clear();
    current_list.lookahead2.leaf_counts.extend(previous_list.lookahead2.leaf_counts.iter());
    current_list.lookahead2.leaf_counts.push(*current_list.lookahead1.leaf_counts.last().unwrap());

    // The previous list needs two new lookahead nodes.
    lists.push(previous_list);
    boundary_pm(lists, leaves);
    boundary_pm(lists, leaves);
    lists.push(current_list);
}

fn boundary_pm(lists: &mut Vec<List>, leaves: &Vec<Leaf>) {
    let mut current_list = lists.pop().unwrap();
    if lists.is_empty() && current_list.next_leaf_index == leaves.len() {
        // We've added all the leaves to the lowest list, so we're done here
        lists.push(current_list);
        return;
    }

    mem::swap(&mut current_list.lookahead1, &mut current_list.lookahead2);

    if lists.is_empty() {
        lowest_list(current_list, lists, leaves);
    } else {
        // We're at a list other than the lowest list.
        let previous_list = lists.pop().unwrap();
        let weight_sum = previous_list.lookahead1.weight + previous_list.lookahead2.weight;

        if current_list.next_leaf_index < leaves.len() && weight_sum > leaves[current_list.next_leaf_index].weight {
            next_leaf(current_list, previous_list, lists, leaves);
        } else {
            next_tree(current_list, previous_list, weight_sum, lists, leaves);
        }
    }
}

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<|endoftext|>


<commit_before>use std::os::unix::process::CommandExt;
use std::process::Command;

const SYSTEMD_RUN: &'static str = "/usr/bin/systemd-run";
const FIREJAIL: &'static str = "/usr/bin/firejail";

const FIREFOX: &'static str = "/usr/bin/firefox";
const FIREFOX_PROFILE: &'static str = "/etc/firejail/firefox.profile";
const FIREFOX_MOZILLA_SFW_PROFILE: &'static str = "default-1473025815439";

const FIREFOX_MEMORY_HIGH: i32 = 4420;
const FIREFOX_MEMORY_MAX: i32 = 5200;

const KEIRA: &'static str = "/home/kalvin/Downloads/.firefox-nightly/firefox";
const KEIRA_PROFILE: &'static str = "/etc/firejail/firefox-nightly.profile";

const TERM: &'static str = "/usr/bin/alacritty";
const TERM_PROFILE: &'static str = "/etc/firejail/x-terminal-emulator.profile";

const THUNDERBIRD: &'static str = "/home/kalvin/Downloads/.thunderbird-beta/thunderbird";
const THUNDERBIRD_PROFILE: &'static str = "/etc/firejail/thunderbird-beta.profile";

#[derive(Debug)]
struct CgroupedFirejailedCommandOptions<'a> {
    bin_path: &'a str,
    memory_high: Option<i32>,
    memory_max: Option<i32>,
    firejail_profile: Option<&'a str>,
    implicit_extra_args: Option<&'a [&'a str]>,
    argv_remainder: Vec<String>,
}

fn cgrouped_firejail_command(options: CgroupedFirejailedCommandOptions) -> Command {
    let mut command = Command::new(SYSTEMD_RUN);
    command.args(&["--user", "--scope"]);
    command.args(&["-p", "MemorySwapMax=0"]);

    if let Some(memory_high) = options.memory_high {
        let memory_high_owned_string = format!("MemoryHigh={}M", memory_high);
        command.args(&["-p", memory_high_owned_string.as_str()]);
    }
    if let Some(memory_max) = options.memory_max {
        let memory_max_owned_string = format!("MemoryMax={}M", memory_max);
        command.args(&["-p", memory_max_owned_string.as_str()]);
    }

    if let Some(firejail_profile) = options.firejail_profile {
        command.arg(FIREJAIL);
        command.arg("--ignore=seccomp");
        command.arg(format!("--profile={}", firejail_profile));
    }

    command.arg(options.bin_path);

    if let Some(implicit_extra_args) = options.implicit_extra_args {
        command.args(implicit_extra_args);
    }

    command.args(options.argv_remainder);

    command
}

fn simple_firejail_command(target: &str, args: Vec<String>) -> Command {
    let mut command = Command::new(FIREJAIL);

    command.arg(format!("--profile=/etc/firejail/{}.profile", target));
    command.arg(format!("/usr/bin/{}", target));
    command.args(args);

    command
}

fn init_command() -> Command {
    let mut args_less_the_first = std::env::args();
    // Pops the name of this executable.
    args_less_the_first.next();

    // Pops the name of the unner target.
    let target = args_less_the_first.next().unwrap();

    // Copies and collects the remaining argv.
    let args: Vec<String> = args_less_the_first.collect();

    match target.as_str() {
        "ff" => {
            return cgrouped_firejail_command(CgroupedFirejailedCommandOptions {
                bin_path: FIREFOX,
                memory_high: Some(FIREFOX_MEMORY_HIGH),
                memory_max: Some(FIREFOX_MEMORY_MAX),
                firejail_profile: Some(FIREFOX_PROFILE),
                implicit_extra_args: Some(&["-P", FIREFOX_MOZILLA_SFW_PROFILE]),
                argv_remainder: args,
            });
        }
        "keira" => {
            return cgrouped_firejail_command(CgroupedFirejailedCommandOptions {
                bin_path: KEIRA,
                memory_high: Some(FIREFOX_MEMORY_HIGH),
                memory_max: Some(FIREFOX_MEMORY_MAX),
                firejail_profile: Some(KEIRA_PROFILE),
                implicit_extra_args: Some(&["-P", "nightly"]),
                argv_remainder: args,
            });
        }
        "npv" => return simple_firejail_command("mpv", args),
        "t" => {
            return cgrouped_firejail_command(CgroupedFirejailedCommandOptions {
                bin_path: TERM,
                memory_high: None,
                memory_max: None,
                firejail_profile: Some(TERM_PROFILE),
                implicit_extra_args: None,
                argv_remainder: args,
            })
        }
        "read" => {
            let formatted_args: Vec<String> = vec![format!("about:reader?url={}", args[0])];
            return cgrouped_firejail_command(CgroupedFirejailedCommandOptions {
                bin_path: FIREFOX,
                memory_high: Some(FIREFOX_MEMORY_HIGH),
                memory_max: Some(FIREFOX_MEMORY_MAX),
                firejail_profile: Some(FIREFOX_PROFILE),
                implicit_extra_args: Some(&["-P", FIREFOX_MOZILLA_SFW_PROFILE]),
                argv_remainder: formatted_args,
            });
        }
        "tbb" => {
            return cgrouped_firejail_command(CgroupedFirejailedCommandOptions {
                bin_path: THUNDERBIRD,
                memory_high: None,
                memory_max: None,
                firejail_profile: Some(THUNDERBIRD_PROFILE),
                implicit_extra_args: None,
                argv_remainder: args,
            });
        }
        "vlk" => return simple_firejail_command("vlc", args),
        "z" => return simple_firejail_command("zathura", args),
        _ => panic!(format!("no handler for ``{}''", target)),
    }
}

fn main() {
    let mut command = init_command();
    println!("{:#?}", command);
    command.exec();
}
<commit_msg>unner: disable seccomp in VLC<commit_after>use std::os::unix::process::CommandExt;
use std::process::Command;

const SYSTEMD_RUN: &'static str = "/usr/bin/systemd-run";
const FIREJAIL: &'static str = "/usr/bin/firejail";

const FIREFOX: &'static str = "/usr/bin/firefox";
const FIREFOX_PROFILE: &'static str = "/etc/firejail/firefox.profile";
const FIREFOX_MOZILLA_SFW_PROFILE: &'static str = "default-1473025815439";

const FIREFOX_MEMORY_HIGH: i32 = 4420;
const FIREFOX_MEMORY_MAX: i32 = 5200;

const KEIRA: &'static str = "/home/kalvin/Downloads/.firefox-nightly/firefox";
const KEIRA_PROFILE: &'static str = "/etc/firejail/firefox-nightly.profile";

const TERM: &'static str = "/usr/bin/alacritty";
const TERM_PROFILE: &'static str = "/etc/firejail/x-terminal-emulator.profile";

const THUNDERBIRD: &'static str = "/home/kalvin/Downloads/.thunderbird-beta/thunderbird";
const THUNDERBIRD_PROFILE: &'static str = "/etc/firejail/thunderbird-beta.profile";

#[derive(Debug)]
struct CgroupedFirejailedCommandOptions<'a> {
    bin_path: &'a str,
    memory_high: Option<i32>,
    memory_max: Option<i32>,
    firejail_profile: Option<&'a str>,
    implicit_extra_args: Option<&'a [&'a str]>,
    argv_remainder: Vec<String>,
}

fn cgrouped_firejail_command(options: CgroupedFirejailedCommandOptions) -> Command {
    let mut command = Command::new(SYSTEMD_RUN);
    command.args(&["--user", "--scope"]);
    command.args(&["-p", "MemorySwapMax=0"]);

    if let Some(memory_high) = options.memory_high {
        let memory_high_owned_string = format!("MemoryHigh={}M", memory_high);
        command.args(&["-p", memory_high_owned_string.as_str()]);
    }
    if let Some(memory_max) = options.memory_max {
        let memory_max_owned_string = format!("MemoryMax={}M", memory_max);
        command.args(&["-p", memory_max_owned_string.as_str()]);
    }

    if let Some(firejail_profile) = options.firejail_profile {
        command.arg(FIREJAIL);
        command.arg("--ignore=seccomp");
        command.arg(format!("--profile={}", firejail_profile));
    }

    command.arg(options.bin_path);

    if let Some(implicit_extra_args) = options.implicit_extra_args {
        command.args(implicit_extra_args);
    }

    command.args(options.argv_remainder);

    command
}

fn simple_firejail_command(target: &str, args: Vec<String>) -> Command {
    let mut command = Command::new(FIREJAIL);

    command.arg(format!("--profile=/etc/firejail/{}.profile", target));
    command.arg(format!("/usr/bin/{}", target));
    command.args(args);

    command
}

fn init_command() -> Command {
    let mut args_less_the_first = std::env::args();
    // Pops the name of this executable.
    args_less_the_first.next();

    // Pops the name of the unner target.
    let target = args_less_the_first.next().unwrap();

    // Copies and collects the remaining argv.
    let args: Vec<String> = args_less_the_first.collect();

    match target.as_str() {
        "ff" => {
            return cgrouped_firejail_command(CgroupedFirejailedCommandOptions {
                bin_path: FIREFOX,
                memory_high: Some(FIREFOX_MEMORY_HIGH),
                memory_max: Some(FIREFOX_MEMORY_MAX),
                firejail_profile: Some(FIREFOX_PROFILE),
                implicit_extra_args: Some(&["-P", FIREFOX_MOZILLA_SFW_PROFILE]),
                argv_remainder: args,
            });
        }
        "keira" => {
            return cgrouped_firejail_command(CgroupedFirejailedCommandOptions {
                bin_path: KEIRA,
                memory_high: Some(FIREFOX_MEMORY_HIGH),
                memory_max: Some(FIREFOX_MEMORY_MAX),
                firejail_profile: Some(KEIRA_PROFILE),
                implicit_extra_args: Some(&["-P", "nightly"]),
                argv_remainder: args,
            });
        }
        "npv" => return simple_firejail_command("mpv", args),
        "t" => {
            return cgrouped_firejail_command(CgroupedFirejailedCommandOptions {
                bin_path: TERM,
                memory_high: None,
                memory_max: None,
                firejail_profile: Some(TERM_PROFILE),
                implicit_extra_args: None,
                argv_remainder: args,
            })
        }
        "read" => {
            let formatted_args: Vec<String> = vec![format!("about:reader?url={}", args[0])];
            return cgrouped_firejail_command(CgroupedFirejailedCommandOptions {
                bin_path: FIREFOX,
                memory_high: Some(FIREFOX_MEMORY_HIGH),
                memory_max: Some(FIREFOX_MEMORY_MAX),
                firejail_profile: Some(FIREFOX_PROFILE),
                implicit_extra_args: Some(&["-P", FIREFOX_MOZILLA_SFW_PROFILE]),
                argv_remainder: formatted_args,
            });
        }
        "tbb" => {
            return cgrouped_firejail_command(CgroupedFirejailedCommandOptions {
                bin_path: THUNDERBIRD,
                memory_high: None,
                memory_max: None,
                firejail_profile: Some(THUNDERBIRD_PROFILE),
                implicit_extra_args: None,
                argv_remainder: args,
            });
        }
        "vlk" => {
            return cgrouped_firejail_command(CgroupedFirejailedCommandOptions {
                bin_path: "/usr/bin/vlc",
                memory_high: None,
                memory_max: None,
                firejail_profile: Some("/etc/firejail/vlc.profile"),
                implicit_extra_args: None,
                argv_remainder: args,
            })
        }
        "z" => return simple_firejail_command("zathura", args),
        _ => panic!(format!("no handler for ``{}''", target)),
    }
}

fn main() {
    let mut command = init_command();
    println!("{:#?}", command);
    command.exec();
}
<|endoftext|>


<commit_before>//! Resource Records

use std::fmt;
use super::{Composer, ComposeError, ComposeResult, DName, GenericRecordData,
            ParsedDName, ParsedRecordData, Parser, ParseResult, RecordData};
use ::iana::{Class, Rtype};

//------------ Record --------------------------------------------------------

#[derive(Clone, Debug)]
pub struct Record<N: DName, D: RecordData> {
    name: N,
    class: Class,
    ttl: u32,
    data: D
}

/// # Creation and Element Access
///
impl<N: DName, D: RecordData> Record<N, D> {
    /// Creates a new record from its parts.
    pub fn new(name: N, class: Class, ttl: u32, data: D) -> Self {
        Record{name: name, class: class, ttl: ttl, data: data}
    }

    /// Returns a reference to the domain name.
    pub fn name(&self) -> &N {
        &self.name
    }

    /// Returns a mutable reference to the domain name.
    pub fn name_mut(&mut self) -> &mut N {
        &mut self.name
    }

    /// Returns the record type.
    pub fn rtype(&self) -> Rtype {
        self.data.rtype()
    }

    /// Returns the record class.
    pub fn class(&self) -> Class {
        self.class
    }

    /// Sets the record’s class.
    pub fn set_class(&mut self, class: Class) {
        self.class = class
    }

    /// Returns the record’s time-to-live.
    pub fn ttl(&self) -> u32 {
        self.ttl
    }

    /// Sets the record’s time-to-live.
    pub fn set_ttl(&mut self, ttl: u32) {
        self.ttl = ttl
    }

    /// Return a reference to the record data.
    pub fn data(&self) -> &D {
        &self.data
    }

    /// Returns a mutable reference to the record data.
    pub fn data_mut(&mut self) -> &mut D {
        &mut self.data
    }

    /// Converts the record into its record data.
    pub fn into_data(self) -> D {
        self.data
    }
}

/// # Parsing
///
impl<'a, D: ParsedRecordData<'a>> Record<ParsedDName<'a>, D> {
    pub fn parse(parser: &mut Parser<'a>) -> ParseResult<Option<Self>> {
        let name = try!(ParsedDName::parse(parser));
        let rtype = try!(Rtype::parse(parser));
        let class = try!(Class::parse(parser));
        let ttl = try!(parser.parse_u32());
        let rdlen = try!(parser.parse_u16()) as usize;
        try!(parser.set_limit(rdlen));
        let data = try!(D::parse(rtype, parser));
        if data.is_none() {
            try!(parser.skip(rdlen));
        }
        parser.remove_limit();
        Ok(data.map(|data| Record::new(name, class, ttl, data)))
    }
}

impl<'a> Record<ParsedDName<'a>, GenericRecordData<'a>> {
    pub fn parse_generic(parser: &mut Parser<'a>) -> ParseResult<Self> {
        Self::parse(parser).map(Option::unwrap)
    }
}

/// # Composing
///
impl<N: DName, D: RecordData> Record<N, D> {
    pub fn compose<C: AsMut<Composer>>(&self, mut composer: C)
                                       -> ComposeResult<()> {
        try!(self.name.compose(composer.as_mut()));
        try!(self.data.rtype().compose(composer.as_mut()));
        try!(self.class.compose(composer.as_mut()));
        try!(composer.as_mut().compose_u32(self.ttl));
        let pos = composer.as_mut().pos();
        try!(composer.as_mut().compose_u16(0));
        try!(self.data.compose(composer.as_mut()));
        let delta = composer.as_mut().delta(pos) - 2;
        if delta > (::std::u16::MAX as usize) {
            return Err(ComposeError::Overflow)
        }
        composer.as_mut().update_u16(pos, delta as u16);
        Ok(())
    }
}

//--- From

impl<N: DName, D: RecordData> From<(N, Class, u32, D)> for Record<N, D> {
    fn from(x: (N, Class, u32, D)) -> Self {
        Record::new(x.0, x.1, x.2, x.3)
    }
}

impl<N: DName, D: RecordData> From<(N, u32, D)> for Record<N, D> {
    fn from(x: (N, u32, D)) -> Self {
        Record::new(x.0, Class::In, x.1, x.2)
    }
}

//--- Display

impl<N, D> fmt::Display for Record<N, D>
     where N: DName + fmt::Display,
           D: RecordData + fmt::Display {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}\t{}\t{}\t{}\t{}",
               self.name, self.ttl, self.class, self.data.rtype(),
               self.data)
    }
}

//------------ GenericRecord -------------------------------------------------

pub type GenericRecord<'a> = Record<ParsedDName<'a>, GenericRecordData<'a>>;

<commit_msg>Add documentation to bits::record.<commit_after>//! Resource Records
//!
//! This module defines the [`Record`] type that represents DNS resource
//! records. It also provides a type alias for parsed records with generic
//! data through [`GenericRecord`].
//!
//! [`Record`]: struct.Record.html
//! [`GenericRecord`]: type.GenericRecord.html

use std::fmt;
use super::{Composer, ComposeError, ComposeResult, DName, GenericRecordData,
            ParsedDName, ParsedRecordData, Parser, ParseResult, RecordData};
use ::iana::{Class, Rtype};

//------------ Record --------------------------------------------------------

/// A DNS resource record.
///
/// All information available through the DNS is stored in resource records.
/// They have a three part key of a domain name, resource record type, and
/// class. Data is arranged in a tree which is navigated using the domain
/// name. Each node in the tree carries a label, starting with the root
/// label as the top-most node. The tree is traversed by stepping through the
/// name from right to left, finding a child node carring the label of each
/// step.
///
/// The record type describes the kind of data the record holds, such as IP
/// addresses. The class, finally, describes which sort of network the
/// information is for since DNS was originally intended to be used for
/// networks other than the Internet as well. In practice, the only relevant
/// class is IN, the Internet. Note that each class has its own tree of nodes.
///
/// The payload of a resource record is its data. Its purpose, meaning, and
/// format is determined by the record type. For each unique three-part key
/// there can be multiple resource records. All these records for the same
/// key are called *resource record sets,* most often shortened to ‘RRset.’
///
/// There is one more piece of data: the TTL or time to live. This value
/// says how long a record remains valid before it should be refreshed from
/// its original source, given in seconds. The TTL is used to add caching
/// facilities to the DNS.
///
/// Values of the `Record` type represent one single resource record. Since
/// there are currently more than eighty record types—see [`Rtype`] for a
/// complete list—, the type is generic over a trait for record data. This
/// trait holds both the record type value and the record data as they are
/// inseparably entwined.
///
/// Since records contain domain names, the `Record` type is additionally
/// generic over a domain name type.
///
/// There is three ways to create a record value. First, you can make one
/// yourself using the [`new()`] function. In will neatly take care of all
/// the generics through type inference. Secondly, you can parse a record
/// from an existing message. [`Message`] and its friends provide a way to
/// do that; see there for all the details. Finally, you can scan a record
/// from master data (aka zonefiles). See the [domain::master] module for
/// that.
///
/// Records can be place into DNS messages through the value chain starting
/// with [`MessageBuilder`]. In order to make adding records easier, `Record`
/// implements the `From` trait for two kinds of tuples: A four-tuple of
/// name, class, time-to-live value, and record data and a triple leaving
/// out the class and assuming it to `Class::In`.
///
/// [`new()´]: #method.new
/// [`Message`]: ../message/struct.Message.html
/// [`MessageBuilder`]: ../message_builder/struct.MessageBuilder.html
/// [`Rtype`]: ../../iana/enum.Rtype.html
/// [domain::master]: ../../master/index.html
#[derive(Clone, Debug)]
pub struct Record<N: DName, D: RecordData> {
    name: N,
    class: Class,
    ttl: u32,
    data: D
}

/// # Creation and Element Access
///
impl<N: DName, D: RecordData> Record<N, D> {
    /// Creates a new record from its parts.
    pub fn new(name: N, class: Class, ttl: u32, data: D) -> Self {
        Record{name: name, class: class, ttl: ttl, data: data}
    }

    /// Returns a reference to the domain name.
    ///
    /// The domain name, sometimes called the *owner* of the record,
    /// specifies the node in the DNS tree this record belongs to.
    pub fn name(&self) -> &N {
        &self.name
    }

    /// Returns the record type.
    pub fn rtype(&self) -> Rtype {
        self.data.rtype()
    }

    /// Returns the record class.
    pub fn class(&self) -> Class {
        self.class
    }

    /// Sets the record’s class.
    pub fn set_class(&mut self, class: Class) {
        self.class = class
    }

    /// Returns the record’s time-to-live.
    pub fn ttl(&self) -> u32 {
        self.ttl
    }

    /// Sets the record’s time-to-live.
    pub fn set_ttl(&mut self, ttl: u32) {
        self.ttl = ttl
    }

    /// Return a reference to the record data.
    pub fn data(&self) -> &D {
        &self.data
    }

    /// Returns a mutable reference to the record data.
    pub fn data_mut(&mut self) -> &mut D {
        &mut self.data
    }

    /// Trades the record for its record data.
    pub fn into_data(self) -> D {
        self.data
    }
}

/// # Parsing
///
impl<'a, D: ParsedRecordData<'a>> Record<ParsedDName<'a>, D> {
    /// Parses a record from a parser.
    ///
    /// This function is only available for records that use a parsed domain
    /// name and a record data type that, too, can be parsed.
    pub fn parse(parser: &mut Parser<'a>) -> ParseResult<Option<Self>> {
        let name = try!(ParsedDName::parse(parser));
        let rtype = try!(Rtype::parse(parser));
        let class = try!(Class::parse(parser));
        let ttl = try!(parser.parse_u32());
        let rdlen = try!(parser.parse_u16()) as usize;
        try!(parser.set_limit(rdlen));
        let data = try!(D::parse(rtype, parser));
        if data.is_none() {
            try!(parser.skip(rdlen));
        }
        parser.remove_limit();
        Ok(data.map(|data| Record::new(name, class, ttl, data)))
    }
}

impl<'a> Record<ParsedDName<'a>, GenericRecordData<'a>> {
    /// Parses a record with generic data from a parser.
    pub fn parse_generic(parser: &mut Parser<'a>) -> ParseResult<Self> {
        Self::parse(parser).map(Option::unwrap)
    }
}

/// # Composing
///
impl<N: DName, D: RecordData> Record<N, D> {
    /// Appends the record’s wire-format representation to a composer.
    pub fn compose<C: AsMut<Composer>>(&self, mut composer: C)
                                       -> ComposeResult<()> {
        try!(self.name.compose(composer.as_mut()));
        try!(self.data.rtype().compose(composer.as_mut()));
        try!(self.class.compose(composer.as_mut()));
        try!(composer.as_mut().compose_u32(self.ttl));
        let pos = composer.as_mut().pos();
        try!(composer.as_mut().compose_u16(0));
        try!(self.data.compose(composer.as_mut()));
        let delta = composer.as_mut().delta(pos) - 2;
        if delta > (::std::u16::MAX as usize) {
            return Err(ComposeError::Overflow)
        }
        composer.as_mut().update_u16(pos, delta as u16);
        Ok(())
    }
}

//--- From

impl<N: DName, D: RecordData> From<(N, Class, u32, D)> for Record<N, D> {
    fn from(x: (N, Class, u32, D)) -> Self {
        Record::new(x.0, x.1, x.2, x.3)
    }
}

impl<N: DName, D: RecordData> From<(N, u32, D)> for Record<N, D> {
    fn from(x: (N, u32, D)) -> Self {
        Record::new(x.0, Class::In, x.1, x.2)
    }
}

//--- Display

impl<N, D> fmt::Display for Record<N, D>
     where N: DName + fmt::Display,
           D: RecordData + fmt::Display {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}\t{}\t{}\t{}\t{}",
               self.name, self.ttl, self.class, self.data.rtype(),
               self.data)
    }
}

//------------ GenericRecord -------------------------------------------------

/// A record with generic record data.
pub type GenericRecord<'a> = Record<ParsedDName<'a>, GenericRecordData<'a>>;

<|endoftext|>