<commit_before>//*****************************
// Java Codecity data model

extern crate classreader;

use self::classreader::*;
use std::collections::HashMap;
use std::cmp;

#[derive(Debug)]
pub struct Building {
    name:   String,
    width:  u16,
    depth:  u16,
    height: u16,
    // Placement
    pos_w:  u16,
    pos_d:  u16
}

impl Building {
    fn new(class: &Class) -> Building {
        Building {
            name: get_class_name(class).to_string(),
            width: class.methods.len() as u16 + 1,
            depth: class.fields.len() as u16 + 1,
            height: get_total_code_size(&class) as u16 / 10 + 1,
            pos_w: 0,
            pos_d: 0
        }
    }

    fn pos(self: &Building) -> (u16, u16) {
        (self.pos_w, self.pos_d)
    }
}

enum Direction {
    Width,
    Depth
}

#[derive(Debug)]
pub struct Package {
    name:  String,
    width: u16,
    depth: u16,
    packages:  HashMap<String, Package>,
    buildings: Vec<Building>
}

impl Package {
    /// Create a new Package.
    pub fn new(name: &str) -> Package {
        Package {
            name: name.to_string(),
            width: 1,
            depth: 1,
            packages: HashMap::new(),
            buildings: Vec::new()
        }
    }

    fn size(self: &Package) -> (u16, u16) {
        (self.width, self.depth)
    }

    /// Add a class or package to this package
    fn add(&mut self, names: &[&str], class: &Class) {
        if names.len() == 1 {
            self.buildings.push(Building::new(class));
        } else {
            match names.split_first() {
                Some((head, tail)) => {
                    self.packages.entry(head.to_string()).or_insert_with(|| Package::new(head)).add(&tail, class);
                },
                _ => {}
            }
        }
    }

    /// Pack own members for the nice rendering
    fn pack(&mut self) {
        for (_, ch) in &mut self.packages {
            ch.pack();
        }
        let mut free_w = 0;
        let mut free_d = 0;
        let mut cur_w = 1;
        let mut cur_d = 1;
        let mut dir = Direction::Width;
        // TODO Packages
        // Buildings
        self.buildings.sort_by_key(|b| b.width * b.depth);
        self.buildings.reverse();
        for b in &mut self.buildings {
            // println!("Used: {}x{} free: {}x{} placing {}x{}", cur_w, cur_d, free_w, free_d, b.width, b.depth);
            if b.width > free_w || b.depth > free_d {
                // decide the direction for growth
                if self.width <= self.depth {
                    // Grow by width
                    dir = Direction::Width;
                    b.pos_w = self.width;
                    b.pos_d = 1;
                    cur_w = self.width;
                    cur_d = 2 + b.depth;
                    self.width += b.width + 1;
                    self.depth = cmp::max(self.depth, b.depth + 2);
                    free_w = b.width;
                    free_d = self.depth - b.depth - 2;
                } else {
                    // Grow by depth
                    dir = Direction::Depth;
                    b.pos_d = self.depth;
                    b.pos_w = 1;
                    cur_d = self.depth;
                    cur_w = 2 + b.width;
                    self.depth += b.depth + 1;
                    self.width = cmp::max(self.width, b.width + 2);
                    free_d = b.depth;
                    free_w = self.width - b.width - 2;
                }
            } else {
                // Enough room
                match dir {
                    Direction::Width => {
                        b.pos_w = cur_w;
                        b.pos_d = cur_d;
                        free_d -= b.depth + 1;
                        cur_d += b.depth + 1;
                    }
                    Direction::Depth => {
                        b.pos_d = cur_d;
                        b.pos_w = cur_w;
                        free_w -= b.width + 1;
                        cur_w += b.width + 1;
                    }
                }
            }
        }
    }
}

pub fn build_packages(classes: &Vec<Class>) -> Package {
    let mut root_pkg = Package::new("_root_");
    for c in classes {
        let name = get_class_name(&c);
        let names: Vec<&str> = name.split('/').collect();
        root_pkg.add(&names, c);
    }
    return root_pkg;
}

pub fn process_class(class: &Class) {
    assert_eq!(0xCAFEBABE, class.magic);
    println!("class: {}", get_class_name(&class));
    println!("method count: {}", class.methods.len());
    println!("field count:  {}", class.fields.len());
    println!("code size:    {}", get_total_code_size(&class));
    for method in &class.methods {
        println!("\tmethod: {}{}", get_string(&class, method.name_index as usize), get_string(&class, method.descriptor_index as usize));
        println!("\t\tsize:       {}", get_method_size(&method));
        println!("\t\tlines:      {}", get_method_lines(&method));
        println!("\t\tcomplexity: {}", get_method_complexity(&method));
    }

    println!("======================================================");
}

/// Get constant from a pool, correcting for java's 1-based indexes.
fn get_const(class: &Class, i: usize) -> &ConstantPoolInfo {
    &class.constant_pool[i - 1]
}

/// Get string from constant pool
fn get_string(class: &Class, index: usize) -> String {
    match get_const(class, index) {
        &ConstantPoolInfo::Utf8(ref s) => s.clone(),
        _ => "?".to_string()
    }
}

fn get_class_name(class: &Class) -> String {
    match get_const(class, class.this_class as usize) {
        &ConstantPoolInfo::Class(index) => get_string(class, index as usize),
        _ => "?".to_string()
    }
}

fn get_total_code_size(class: &Class) -> usize {
    class.methods.iter().fold(0usize, |sum, method| sum + get_method_size(method))
}

/// Compute method size in byte code instructions
/// If multiple code segments are present in a method (not a javac output) sum it up
fn get_method_size(method: &Method) -> usize {
    method.attributes.iter().fold(0usize, |sum, a| sum + match a {
            &Attribute::Code{ref code, ..} => code.len(),
            _ => 0
        })
}

/// Compute method size in lines of code.
/// If multiple code segments or multiple line number attrbitues are present in a method (not a javac output) sum it up
fn get_method_lines(method: &Method) -> u16 {
    method.attributes.iter().fold(0, |sum, a| sum + match a {
            &Attribute::Code{ref attributes, ..} => attributes.iter().fold(0, |sum, a| sum + match a {
                &Attribute::LineNumberTable(ref line_table) => {
                    let end = get_opt_line_number(line_table.last());
                    let start = get_opt_line_number(line_table.first());
                    if start > end { // TODO: Investigate when it happens
                        0
                    } else {
                        end - start
                    }
                },
                _ => 0
            }),
            _ => 0
        })

}

fn get_opt_line_number(line: Option<&LineNumber>) -> u16 {
    line.map(|x| x.line_number).unwrap_or(0)
}

fn get_method_complexity(method: &Method) -> u16 {
    method.attributes.iter().fold(0, |sum, a| sum + match a {
            &Attribute::Code{ref code, ..} => get_code_complexity(code),
            _ => 0
        })
}

/// Using complexity definition from http://www.leepoint.net/principles_and_practices/complexity/complexity-java-method.html
/// Start with 1 for the method. Add one for each of the following flow-related elements that are found in the method.
/// Returns     Each return that isn't the last statement of a method.
/// Selection   if, else, case, default.
/// Loops   for, while, do-while, break, and continue.
/// Operators   &&, ||, ?, and :
/// Exceptions  catch, finally, throw, or throws clause.
/// Bytecode reference: http://www.angelibrary.com/computer/java_21days/ch32.htm
fn get_code_complexity(code: &Vec<(u32, Instruction)>) -> u16 {
    // println!("{:?}", code);
    // We start from zero to account for always-present final return - it will add the "starting" 1 for us.
    code.iter().fold(0, |sum, instruction| sum + match instruction.1 {
        Instruction::areturn => 1,
        Instruction::athrow => 1,
        Instruction::dreturn => 1,
        Instruction::freturn => 1,
        Instruction::if_acmpeq(..) => 1,
        Instruction::if_acmpne(..) => 1,
        Instruction::if_icmpeq(..) => 1,
        Instruction::if_icmpne(..) => 1,
        Instruction::if_icmplt(..) => 1,
        Instruction::if_icmpge(..) => 1,
        Instruction::if_icmpgt(..) => 1,
        Instruction::if_icmple(..) => 1,
        Instruction::ifeq(..) => 1,
        Instruction::ifne(..) => 1,
        Instruction::iflt(..) => 1,
        Instruction::ifge(..) => 1,
        Instruction::ifgt(..) => 1,
        Instruction::ifle(..) => 1,
        Instruction::ifnonnull(..) => 1,
        Instruction::ifnull(..) => 1,
        Instruction::ireturn => 1,
        Instruction::lookupswitch(_, ref jumps) => jumps.len() as u16, // tricky - lets use the size of the jump table
        Instruction::lreturn => 1,
        Instruction::return_ => 1,
        Instruction::tableswitch(_, _, ref jumps) => jumps.len() as u16, // tricky - lets use the size of the jump table
        // Subroutines in the method. Currently - do not count them towards the complexity. They are used to compile "finally", but I do not see them as branching really. Finally block is executed every time
        // Instruction::jsr(..) => 1,
        // Instruction::jsr_w(..) => 1,
        // Instruction::ret => 1,
        // Instruction::ret_w => 1,
        _ => 0
    })
}

//*****************************
// Unit tests

fn mock(name: &str, w: u16, d: u16) -> Building {
  Building {
    name: name.to_string(),
    width: w,
    depth: d,
    height: 0,
    pos_w: 0,
    pos_d: 0
  }
}

/// A function to print out details about placement
fn debug_test(pkg: &Package) {
    println!("{:?}", pkg);
    for b in &pkg.buildings {
        println!("{:?}", b);
    }
}

#[test]
fn pack_4_singles() {
    let mut pkg = Package::new("_root_");
    for _ in 0..4 {
        pkg.buildings.push(mock("a", 1, 1));
    }
    pkg.pack();
    debug_test(&pkg);
    assert!(pkg.size() == (5, 5));
    // Make sure the last building is in the corner
    assert!(pkg.buildings[3].pos() == (3, 3));
}

#[test]
fn pack_5_singles() {
    let mut pkg = Package::new("_root_");
    for _ in 0..5 {
        pkg.buildings.push(mock("a", 1, 1));
    }
    pkg.pack();
    debug_test(&pkg);
    assert!(pkg.size() == (7, 5));
    assert!(pkg.buildings[4].pos() == (5, 1));
}

#[test]
fn pack_16_singles() {
    let mut pkg = Package::new("_root_");
    for _ in 0..16 {
        pkg.buildings.push(mock("a", 1, 1));
    }
    pkg.pack();
    debug_test(&pkg);
    assert!(pkg.size() == (9, 9));
    // Make sure the last building is in the corner
    assert!(pkg.buildings[15].pos() == (7, 7));
}

#[test]
fn pack_smart_square_width() {
    let mut pkg = Package::new("_root_");
    pkg.buildings.push(mock("w", 3, 1));
    for _ in 0..2 {
        pkg.buildings.push(mock("a", 1, 1));
    }
    pkg.pack();
    debug_test(&pkg);
    assert!(pkg.size() == (5, 5));
    // Make sure the last building is in the corner
    assert!(pkg.buildings[2].pos() == (3, 3));
}

#[test]
fn pack_buildings() {
    let mut pkg = Package::new("_root_");
    pkg.buildings.push(mock("b", 8, 3));
    pkg.buildings.push(mock("c", 4, 2));
    pkg.buildings.push(mock("a", 10, 7));
    pkg.buildings.push(mock("d", 3, 2));
    pkg.buildings.push(mock("e1", 1, 1));
    pkg.buildings.push(mock("e2", 1, 1));
    pkg.buildings.push(mock("e3", 1, 1));
    println!("Starting the test");
    pkg.pack();
    debug_test(&pkg);
    assert!(pkg.size() == (17, 13));
    assert!(pkg.buildings[0].pos() == (1, 1));
    assert!(pkg.buildings[1].pos() == (1, 9));
    assert!(pkg.buildings[2].pos() == (12, 1));
    assert!(pkg.buildings[3].pos() == (12, 4));
    assert!(pkg.buildings[4].pos() == (12, 7));
    assert!(pkg.buildings[5].pos() == (12, 9));
    assert!(pkg.buildings[6].pos() == (12, 11));

}<commit_msg>Refactorign to extract common code<commit_after>//*****************************
// Java Codecity data model

extern crate classreader;

use self::classreader::*;
use std::collections::HashMap;
use std::cmp;

#[derive(Debug)]
pub struct Rect {
    // Dimenstions
    width:  u16,
    depth:  u16,
    // Placement
    pos_w:  u16,
    pos_d:  u16
}

impl Rect {
    fn new(w: u16, d: u16) -> Rect {
        Rect {
            width: w,
            depth: d,
            pos_w: 0,
            pos_d: 0
        }
    }

    fn area(self: &Rect) -> u16 {self.width * self.depth}
}

#[derive(Debug)]
pub struct Building {
    name:   String,
    rect:   Rect,
    height: u16,
}

impl Building {
    fn new(class: &Class) -> Building {
        Building {
            name: get_class_name(class).to_string(),
            rect: Rect::new(class.methods.len() as u16 + 1, class.fields.len() as u16 + 1),
            height: get_total_code_size(&class) as u16 / 10 + 1,
        }
    }

    fn pos(self: &Building) -> (u16, u16) {
        (self.rect.pos_w, self.rect.pos_d)
    }
}

enum Direction {
    Width,
    Depth
}

#[derive(Debug)]
pub struct Package {
    name:  String,
    width: u16,
    depth: u16,
    packages:  HashMap<String, Package>,
    buildings: Vec<Building>
}

struct PackState {
    free_w: u16,
    free_d: u16,
    cur_w: u16,
    cur_d: u16,
    dir: Direction
}

impl Package {
    /// Create a new Package.
    pub fn new(name: &str) -> Package {
        Package {
            name: name.to_string(),
            width: 1,
            depth: 1,
            packages: HashMap::new(),
            buildings: Vec::new()
        }
    }

    fn size(self: &Package) -> (u16, u16) {
        (self.width, self.depth)
    }

    /// Add a class or package to this package
    fn add(&mut self, names: &[&str], class: &Class) {
        if names.len() == 1 {
            self.buildings.push(Building::new(class));
        } else {
            match names.split_first() {
                Some((head, tail)) => {
                    self.packages.entry(head.to_string()).or_insert_with(|| Package::new(head)).add(&tail, class);
                },
                _ => {}
            }
        }
    }

    /// Pack own members for the nice rendering
    fn pack(&mut self) {
        for (_, ch) in &mut self.packages {
            ch.pack();
        }
        let mut state = PackState {
            free_w: 0,
            free_d: 0,
            cur_w: 1,
            cur_d: 1,
            dir: Direction::Width
        };

        // TODO Packages
        // Buildings
        self.buildings.sort_by_key(|b| b.rect.area());
        self.buildings.reverse();
        for b in &mut self.buildings {
            let mut r = Rect {.. b.rect};
            self.place_rectangle(&mut state, &mut r);
            b.rect.pos_w = r.pos_w;
            b.rect.pos_d = r.pos_d;
        }
    }

    fn place_rectangle(self: &mut Package, state: &mut PackState, b: &mut Rect) {
        // println!("Used: {}x{} free: {}x{} placing {}x{}", cur_w, cur_d, free_w, free_d, b.width, b.depth);
        if b.width > state.free_w || b.depth > state.free_d {
            // decide the direction for growth
            if self.width <= self.depth {
                // Grow by width
                state.dir = Direction::Width;
                b.pos_w = self.width;
                b.pos_d = 1;
                state.cur_w = self.width;
                state.cur_d = 2 + b.depth;
                self.width += b.width + 1;
                self.depth = cmp::max(self.depth, b.depth + 2);
                state.free_w = b.width;
                state.free_d = self.depth - b.depth - 2;
            } else {
                // Grow by depth
                state.dir = Direction::Depth;
                b.pos_d = self.depth;
                b.pos_w = 1;
                state.cur_d = self.depth;
                state.cur_w = 2 + b.width;
                self.depth += b.depth + 1;
                self.width = cmp::max(self.width, b.width + 2);
                state.free_d = b.depth;
                state.free_w = self.width - b.width - 2;
            }
        } else {
            // Enough room
            match state.dir {
                Direction::Width => {
                    b.pos_w = state.cur_w;
                    b.pos_d = state.cur_d;
                    state.free_d -= b.depth + 1;
                    state.cur_d += b.depth + 1;
                }
                Direction::Depth => {
                    b.pos_d = state.cur_d;
                    b.pos_w = state.cur_w;
                    state.free_w -= b.width + 1;
                    state.cur_w += b.width + 1;
                }
            }
        }
    }
}

pub fn build_packages(classes: &Vec<Class>) -> Package {
    let mut root_pkg = Package::new("_root_");
    for c in classes {
        let name = get_class_name(&c);
        let names: Vec<&str> = name.split('/').collect();
        root_pkg.add(&names, c);
    }
    return root_pkg;
}

pub fn process_class(class: &Class) {
    assert_eq!(0xCAFEBABE, class.magic);
    println!("class: {}", get_class_name(&class));
    println!("method count: {}", class.methods.len());
    println!("field count:  {}", class.fields.len());
    println!("code size:    {}", get_total_code_size(&class));
    for method in &class.methods {
        println!("\tmethod: {}{}", get_string(&class, method.name_index as usize), get_string(&class, method.descriptor_index as usize));
        println!("\t\tsize:       {}", get_method_size(&method));
        println!("\t\tlines:      {}", get_method_lines(&method));
        println!("\t\tcomplexity: {}", get_method_complexity(&method));
    }

    println!("======================================================");
}

/// Get constant from a pool, correcting for java's 1-based indexes.
fn get_const(class: &Class, i: usize) -> &ConstantPoolInfo {
    &class.constant_pool[i - 1]
}

/// Get string from constant pool
fn get_string(class: &Class, index: usize) -> String {
    match get_const(class, index) {
        &ConstantPoolInfo::Utf8(ref s) => s.clone(),
        _ => "?".to_string()
    }
}

fn get_class_name(class: &Class) -> String {
    match get_const(class, class.this_class as usize) {
        &ConstantPoolInfo::Class(index) => get_string(class, index as usize),
        _ => "?".to_string()
    }
}

fn get_total_code_size(class: &Class) -> usize {
    class.methods.iter().fold(0usize, |sum, method| sum + get_method_size(method))
}

/// Compute method size in byte code instructions
/// If multiple code segments are present in a method (not a javac output) sum it up
fn get_method_size(method: &Method) -> usize {
    method.attributes.iter().fold(0usize, |sum, a| sum + match a {
            &Attribute::Code{ref code, ..} => code.len(),
            _ => 0
        })
}

/// Compute method size in lines of code.
/// If multiple code segments or multiple line number attrbitues are present in a method (not a javac output) sum it up
fn get_method_lines(method: &Method) -> u16 {
    method.attributes.iter().fold(0, |sum, a| sum + match a {
            &Attribute::Code{ref attributes, ..} => attributes.iter().fold(0, |sum, a| sum + match a {
                &Attribute::LineNumberTable(ref line_table) => {
                    let end = get_opt_line_number(line_table.last());
                    let start = get_opt_line_number(line_table.first());
                    if start > end { // TODO: Investigate when it happens
                        0
                    } else {
                        end - start
                    }
                },
                _ => 0
            }),
            _ => 0
        })

}

fn get_opt_line_number(line: Option<&LineNumber>) -> u16 {
    line.map(|x| x.line_number).unwrap_or(0)
}

fn get_method_complexity(method: &Method) -> u16 {
    method.attributes.iter().fold(0, |sum, a| sum + match a {
            &Attribute::Code{ref code, ..} => get_code_complexity(code),
            _ => 0
        })
}

/// Using complexity definition from http://www.leepoint.net/principles_and_practices/complexity/complexity-java-method.html
/// Start with 1 for the method. Add one for each of the following flow-related elements that are found in the method.
/// Returns     Each return that isn't the last statement of a method.
/// Selection   if, else, case, default.
/// Loops   for, while, do-while, break, and continue.
/// Operators   &&, ||, ?, and :
/// Exceptions  catch, finally, throw, or throws clause.
/// Bytecode reference: http://www.angelibrary.com/computer/java_21days/ch32.htm
fn get_code_complexity(code: &Vec<(u32, Instruction)>) -> u16 {
    // println!("{:?}", code);
    // We start from zero to account for always-present final return - it will add the "starting" 1 for us.
    code.iter().fold(0, |sum, instruction| sum + match instruction.1 {
        Instruction::areturn => 1,
        Instruction::athrow => 1,
        Instruction::dreturn => 1,
        Instruction::freturn => 1,
        Instruction::if_acmpeq(..) => 1,
        Instruction::if_acmpne(..) => 1,
        Instruction::if_icmpeq(..) => 1,
        Instruction::if_icmpne(..) => 1,
        Instruction::if_icmplt(..) => 1,
        Instruction::if_icmpge(..) => 1,
        Instruction::if_icmpgt(..) => 1,
        Instruction::if_icmple(..) => 1,
        Instruction::ifeq(..) => 1,
        Instruction::ifne(..) => 1,
        Instruction::iflt(..) => 1,
        Instruction::ifge(..) => 1,
        Instruction::ifgt(..) => 1,
        Instruction::ifle(..) => 1,
        Instruction::ifnonnull(..) => 1,
        Instruction::ifnull(..) => 1,
        Instruction::ireturn => 1,
        Instruction::lookupswitch(_, ref jumps) => jumps.len() as u16, // tricky - lets use the size of the jump table
        Instruction::lreturn => 1,
        Instruction::return_ => 1,
        Instruction::tableswitch(_, _, ref jumps) => jumps.len() as u16, // tricky - lets use the size of the jump table
        // Subroutines in the method. Currently - do not count them towards the complexity. They are used to compile "finally", but I do not see them as branching really. Finally block is executed every time
        // Instruction::jsr(..) => 1,
        // Instruction::jsr_w(..) => 1,
        // Instruction::ret => 1,
        // Instruction::ret_w => 1,
        _ => 0
    })
}

//*****************************
// Unit tests

fn mock(name: &str, w: u16, d: u16) -> Building {
  Building {
    name: name.to_string(),
    rect: Rect::new(w, d),
    height: 0
  }
}

/// A function to print out details about placement
fn debug_test(pkg: &Package) {
    println!("{:?}", pkg);
    for b in &pkg.buildings {
        println!("{:?}", b);
    }
}

#[test]
fn pack_4_singles() {
    let mut pkg = Package::new("_root_");
    for _ in 0..4 {
        pkg.buildings.push(mock("a", 1, 1));
    }
    pkg.pack();
    debug_test(&pkg);
    assert!(pkg.size() == (5, 5));
    // Make sure the last building is in the corner
    assert!(pkg.buildings[3].pos() == (3, 3));
}

#[test]
fn pack_5_singles() {
    let mut pkg = Package::new("_root_");
    for _ in 0..5 {
        pkg.buildings.push(mock("a", 1, 1));
    }
    pkg.pack();
    debug_test(&pkg);
    assert!(pkg.size() == (7, 5));
    assert!(pkg.buildings[4].pos() == (5, 1));
}

#[test]
fn pack_16_singles() {
    let mut pkg = Package::new("_root_");
    for _ in 0..16 {
        pkg.buildings.push(mock("a", 1, 1));
    }
    pkg.pack();
    debug_test(&pkg);
    assert!(pkg.size() == (9, 9));
    // Make sure the last building is in the corner
    assert!(pkg.buildings[15].pos() == (7, 7));
}

#[test]
fn pack_smart_square_width() {
    let mut pkg = Package::new("_root_");
    pkg.buildings.push(mock("w", 3, 1));
    for _ in 0..2 {
        pkg.buildings.push(mock("a", 1, 1));
    }
    pkg.pack();
    debug_test(&pkg);
    assert!(pkg.size() == (5, 5));
    // Make sure the last building is in the corner
    assert!(pkg.buildings[2].pos() == (3, 3));
}

#[test]
fn pack_buildings() {
    let mut pkg = Package::new("_root_");
    pkg.buildings.push(mock("b", 8, 3));
    pkg.buildings.push(mock("c", 4, 2));
    pkg.buildings.push(mock("a", 10, 7));
    pkg.buildings.push(mock("d", 3, 2));
    pkg.buildings.push(mock("e1", 1, 1));
    pkg.buildings.push(mock("e2", 1, 1));
    pkg.buildings.push(mock("e3", 1, 1));
    println!("Starting the test");
    pkg.pack();
    debug_test(&pkg);
    assert!(pkg.size() == (17, 13));
    assert!(pkg.buildings[0].pos() == (1, 1));
    assert!(pkg.buildings[1].pos() == (1, 9));
    assert!(pkg.buildings[2].pos() == (12, 1));
    assert!(pkg.buildings[3].pos() == (12, 4));
    assert!(pkg.buildings[4].pos() == (12, 7));
    assert!(pkg.buildings[5].pos() == (12, 9));
    assert!(pkg.buildings[6].pos() == (12, 11));

}<|endoftext|>


<commit_before>use std::collections::HashMap;
use errors::RedisError;
use results::RedisResult;
use redis::RedisClient;

pub struct RedisCommand {
    cmd: Vec<u8>,
}

impl<'a> From<&'a mut RedisCommand> for &'a[u8] {
    fn from(command: &mut RedisCommand) -> &[u8] {
        &command.cmd[..]
    }
}

impl RedisCommand {
    pub fn new() -> RedisCommand {
        RedisCommand {
            cmd: vec![],
        }
    }

    pub fn add_cmd<C>(&mut self, command: C) -> &mut RedisCommand where C: ToString {
        self.cmd.extend(command.to_string().into_bytes());
        self
    }

    pub fn add_arg<A>(&mut self, arg: A) -> &mut RedisCommand where A: ToString{
        self.cmd.extend([32].iter().cloned());
        self.cmd.extend(arg.to_string().into_bytes());
        self
    }

    pub fn add_args<A>(&mut self, args: Vec<A>) -> &mut RedisCommand where A: ToString{
        for arg in args {
            self.cmd.extend([32].iter().cloned());
            self.cmd.extend(arg.to_string().into_bytes());
        }
        self
    }

    pub fn add_arg_map(&mut self, args: HashMap<String, String>) -> &mut RedisCommand {
        for (arg, value) in args {
            self.cmd.extend([32].iter().cloned());
            self.cmd.extend(arg.to_string().into_bytes());
            self.cmd.extend([32].iter().cloned());
            self.cmd.extend(value.to_string().into_bytes());
        }
        self
    }

    pub fn add_binary_arg(&mut self, arg: &[u8]) -> &mut RedisCommand {
        self.cmd.extend([32].iter().cloned());
        self.cmd.extend(arg.iter().cloned());
        self
    }

    pub fn end(&mut self) -> &mut RedisCommand {
        self.cmd.extend([13, 10].iter().cloned());
        self
    }
}

macro_rules! generate_command_traits {
    ($(
        fn $func_name:ident$(<$($gen_id:ident: $gen_type:ident),*>)*($($arg_name:ident: $arg_type:ty),*)  {
            $($cmd:ident $bo:expr;)+
        }
    )*)
    =>
    (
        pub trait CommandBuilder {
            $(
                fn $func_name$(<$($gen_id : $gen_type),*>)* (&mut self $(,$arg_name: $arg_type)*) -> &mut RedisCommand;
            )*
        }

        impl CommandBuilder for RedisCommand{
            $(
                fn $func_name$(<$($gen_id : $gen_type),*>)* (&mut self $(,$arg_name: $arg_type)*) -> &mut RedisCommand {
                    $(self.$cmd($bo));*;
                    self.end()
                }
            )*
        }

        pub trait CommandSender {
            $(
                fn $func_name<R: From<RedisResult>, $($($gen_id : $gen_type),*)*> (&mut self $(,$arg_name: $arg_type)*) -> Result<R, RedisError>;
            )*
        }

        impl CommandSender for RedisClient{
            $(
                fn $func_name<R: From<RedisResult>, $($($gen_id : $gen_type),*)*> (&mut self $(,$arg_name: $arg_type)*) -> Result<R, RedisError> {
                    let cmd = &mut RedisCommand::new();
                    cmd.$func_name($($arg_name),*);

                    let res = try!(self.exec_redis_command(cmd));
                    Ok(res.convert::<R>())
                }
            )*
        }
    )
}

generate_command_traits!{
    fn append<K: ToString, V: ToString>(key: K, value: V) {
        add_cmd("APPEND");
        add_arg(key);
        add_arg(value);
    }

    fn auth<P: ToString>(password: P) {
        add_cmd("AUTH");
        add_arg(password);
    }

    fn bgrewriteaof() {
        add_cmd("BGREWRITEAOF");
    }

    fn bgsave() {
        add_cmd("BGSAVE");
    }

    fn bitcount<K: ToString>(key: K) {
        add_cmd("BITCOUNT");
        add_arg(key);
    }

    fn bitcount_range<K: ToString>(key: K, start_range: i64, end_range: i64) {
        add_cmd("BITCOUNT");
        add_arg(key);
        add_arg(start_range);
        add_arg(end_range);
    }

    fn del<K: ToString>(key: K) {
        add_cmd("DEL");
        add_arg(key);
    }

    fn mdel<K: ToString>(keys: Vec<K>){
        add_cmd("DEL");
        add_args(keys);
    }

    fn exists<K: ToString>(key: K) {
        add_cmd("EXISTS");
        add_arg(key);
    }

    fn mexists<K: ToString>(keys: Vec<K>){
        add_cmd("EXISTS");
        add_args(keys);
    }

    fn expire<K: ToString>(key: K, expiry: i64) {
        add_cmd("EXPIRE");
        add_arg(key);
        add_arg(expiry);
    }

    fn expireat<K: ToString>(key: K, timestamp: i64) {
        add_cmd("EXPIREAT");
        add_arg(key);
        add_arg(timestamp);
    }

    fn get<K: ToString>(key: K) {
        add_cmd("GET");
        add_arg(key);
    }

    fn getrange<K: ToString>(key: K, start_range: i64, end_range: i64) {
        add_cmd("GETRANGE");
        add_arg(key);
        add_arg(start_range);
        add_arg(end_range);
    }

    fn hdel<K: ToString, F: ToString>(key: K, field: F) {
        add_cmd("HDEL");
        add_arg(key);
        add_arg(field);
    }

    fn hmdel<K: ToString, V: ToString>(key: K, fields: Vec<V>) {
        add_cmd("HDEL");
        add_arg(key);
        add_args(fields);
    }

    fn hexists<K: ToString, F: ToString>(key: K, field: F) {
        add_cmd("HEXISTS");
        add_arg(key);
        add_arg(field);
    }

    fn hget<K: ToString, F: ToString>(key: K, field: F) {
        add_cmd("HGET");
        add_arg(key);
        add_arg(field);
    }

    fn hgetall<K: ToString>(key: K) {
        add_cmd("HGETALL");
        add_arg(key);
    }

    fn hincrby<K: ToString, F: ToString>(key: K, field: F, increment: i64) {
        add_cmd("HINCRBY");
        add_arg(key);
        add_arg(field);
        add_arg(increment);
    }

    fn hincrbyfloat<K: ToString, F: ToString>(key: K, field: F, increment: f64) {
        add_cmd("HINCRBYBYFLOAT");
        add_arg(key);
        add_arg(field);
        add_arg(increment);
    }

    fn hkeys<K: ToString>(key: K) {
        add_cmd("HKEYS");
        add_arg(key);
    }

    fn hlen<K: ToString>(key: K) {
        add_cmd("HLEN");
        add_arg(key);
    }

    fn hmget<K: ToString>(key: K, fields: Vec<String>) {
        add_cmd("HMGET");
        add_arg(key);
        add_args(fields);
    }

    fn hmset<K: ToString>(key: K, fields: HashMap<String, String>) {
        add_cmd("HMSET");
        add_arg(key);
        add_arg_map(fields);
    }

    fn hset<K: ToString, F: ToString, V: ToString>(key: K, field: F, value: V) {
        add_cmd("HSET");
        add_arg(key);
        add_arg(field);
        add_arg(value);
    }

    fn hstrlen<K: ToString, F: ToString>(key: K, field: F) {
        add_cmd("HSTRLEN");
        add_arg(key);
        add_arg(field);
    }

    fn hsetnx<K: ToString, F: ToString, V: ToString>(key: K, field: F, value: V) {
        add_cmd("HSETNX");
        add_arg(key);
        add_arg(field);
        add_arg(value);
    }

    fn hvals<K: ToString>(key: K) {
        add_cmd("HVALS");
        add_arg(key);
    }

    fn select(db_index: i32){
        add_cmd("SELECT");
        add_arg(db_index);
    }

    fn set<K: ToString, V: ToString>(key: K, value: V) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
    }

    fn set_binary<K: ToString>(key: K, value: &[u8]) {
        add_cmd("SET");
        add_arg(key);
        add_binary_arg(value);
    }

    fn setex<K: ToString, V: ToString>(key: K, value: V, expiry: i64) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("EX");
        add_arg(expiry);
    }

    fn psetex<K: ToString, V: ToString>(key: K, value: V, expiry: i64) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("PX");
        add_arg(expiry);
    }

    fn setnx<K: ToString, V: ToString>(key: K, value: V) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("NX");
    }

    fn setxx<K: ToString, V: ToString>(key: K, value: V) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("XX");
    }

    fn setex_nx<K: ToString, V: ToString>(key: K, value: V, expiry: i64) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("EX");
        add_arg(expiry);
        add_arg("NX");
    }

    fn setex_xx<K: ToString, V: ToString>(key: K, value: V, expiry: i64) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("EX");
        add_arg(expiry);
        add_arg("NX");
    }

    fn psetex_nx<K: ToString, V: ToString>(key: K, value: V, expiry: i64) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("PX");
        add_arg(expiry);
        add_arg("NX");
    }

    fn psetex_xx<K: ToString, V: ToString>(key: K, value: V, expiry: i64) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("PX");
        add_arg(expiry);
        add_arg("XX");
    }

    fn ttl<K: ToString>(key: K) {
        add_cmd("TTL");
        add_arg(key);
    }

    fn zadd<K: ToString, V: ToString>(key: K, score: f64, member: V) {
        add_cmd("ZADD");
        add_arg(key);
        add_arg(score);
        add_arg(member);
    }

    fn zadd_binary<K: ToString>(key: K, score: f64, member: &[u8]) {
        add_cmd("ZADD");
        add_arg(key);
        add_arg(score);
        add_binary_arg(member);
    }

    fn zaddnx<K: ToString, V: ToString>(score: f64, key: K, member: V) {
        add_cmd("ZADD");
        add_arg(key);
        add_arg("NX");
        add_arg(score);
        add_arg(member);
    }

    fn zaddxx<K: ToString, V: ToString>(score: f64, key: K, member: V) {
        add_cmd("ZADD");
        add_arg(key);
        add_arg("XX");
        add_arg(score);
        add_arg(member);
    }

    fn zaddnx_ch<K: ToString, V: ToString>(score: f64, key: K, member: V) {
        add_cmd("ZADD");
        add_arg(key);
        add_arg("NX");
        add_arg("CH");
        add_arg(score);
        add_arg(member);
    }

    fn zaddxx_ch<K: ToString, V: ToString>(score: f64, key: K, member: V) {
        add_cmd("ZADD");
        add_arg(key);
        add_arg("XX");
        add_arg("CH");
        add_arg(score);
        add_arg(member);
    }

    fn zcard<K: ToString>(key: K) {
        add_cmd("ZCARD");
        add_arg(key);
    }

    fn zcount<K: ToString, S: ToString, E: ToString>(key: K, start_range: S, end_range: E) {
        add_cmd("ZCOUNT");
        add_arg(key);
        add_arg(start_range);
        add_arg(end_range);
    }

    fn zincrby<K: ToString, V: ToString>(key: K, increment: f64, member: V) {
        add_cmd("ZINCRBY");
        add_arg(key);
        add_arg(increment);
        add_arg(member);
    }

    fn zrange<K: ToString, S: ToString, E: ToString>(key: K, start_range: S, end_range: E) {
        add_cmd("ZRANGE");
        add_arg(key);
        add_arg(start_range);
        add_arg(end_range);
    }

    fn zrange_with_scores<K: ToString, S: ToString, E: ToString>(key: K, start_range: S, end_range: E) {
        add_cmd("ZRANGE");
        add_arg(key);
        add_arg(start_range);
        add_arg(end_range);
        add_arg("WITHSCORES");
    }

}

<commit_msg>add the transaction commands<commit_after>use std::collections::HashMap;
use errors::RedisError;
use results::RedisResult;
use redis::RedisClient;

pub struct RedisCommand {
    cmd: Vec<u8>,
}

impl<'a> From<&'a mut RedisCommand> for &'a[u8] {
    fn from(command: &mut RedisCommand) -> &[u8] {
        &command.cmd[..]
    }
}

impl RedisCommand {
    pub fn new() -> RedisCommand {
        RedisCommand {
            cmd: vec![],
        }
    }

    pub fn add_cmd<C>(&mut self, command: C) -> &mut RedisCommand where C: ToString {
        self.cmd.extend(command.to_string().into_bytes());
        self
    }

    pub fn add_arg<A>(&mut self, arg: A) -> &mut RedisCommand where A: ToString{
        self.cmd.extend([32].iter().cloned());
        self.cmd.extend(arg.to_string().into_bytes());
        self
    }

    pub fn add_args<A>(&mut self, args: Vec<A>) -> &mut RedisCommand where A: ToString{
        for arg in args {
            self.cmd.extend([32].iter().cloned());
            self.cmd.extend(arg.to_string().into_bytes());
        }
        self
    }

    pub fn add_arg_map(&mut self, args: HashMap<String, String>) -> &mut RedisCommand {
        for (arg, value) in args {
            self.cmd.extend([32].iter().cloned());
            self.cmd.extend(arg.to_string().into_bytes());
            self.cmd.extend([32].iter().cloned());
            self.cmd.extend(value.to_string().into_bytes());
        }
        self
    }

    pub fn add_binary_arg(&mut self, arg: &[u8]) -> &mut RedisCommand {
        self.cmd.extend([32].iter().cloned());
        self.cmd.extend(arg.iter().cloned());
        self
    }

    pub fn end(&mut self) -> &mut RedisCommand {
        self.cmd.extend([13, 10].iter().cloned());
        self
    }
}

macro_rules! generate_command_traits {
    ($(
        fn $func_name:ident$(<$($gen_id:ident: $gen_type:ident),*>)*($($arg_name:ident: $arg_type:ty),*)  {
            $($cmd:ident $bo:expr;)+
        }
    )*)
    =>
    (
        pub trait CommandBuilder {
            $(
                fn $func_name$(<$($gen_id : $gen_type),*>)* (&mut self $(,$arg_name: $arg_type)*) -> &mut RedisCommand;
            )*
        }

        impl CommandBuilder for RedisCommand{
            $(
                fn $func_name$(<$($gen_id : $gen_type),*>)* (&mut self $(,$arg_name: $arg_type)*) -> &mut RedisCommand {
                    $(self.$cmd($bo));*;
                    self.end()
                }
            )*
        }

        pub trait CommandSender {
            $(
                fn $func_name<R: From<RedisResult>, $($($gen_id : $gen_type),*)*> (&mut self $(,$arg_name: $arg_type)*) -> Result<R, RedisError>;
            )*
        }

        impl CommandSender for RedisClient{
            $(
                fn $func_name<R: From<RedisResult>, $($($gen_id : $gen_type),*)*> (&mut self $(,$arg_name: $arg_type)*) -> Result<R, RedisError> {
                    let cmd = &mut RedisCommand::new();
                    cmd.$func_name($($arg_name),*);

                    let res = try!(self.exec_redis_command(cmd));
                    Ok(res.convert::<R>())
                }
            )*
        }
    )
}

generate_command_traits!{
    fn append<K: ToString, V: ToString>(key: K, value: V) {
        add_cmd("APPEND");
        add_arg(key);
        add_arg(value);
    }

    fn auth<P: ToString>(password: P) {
        add_cmd("AUTH");
        add_arg(password);
    }

    fn bgrewriteaof() {
        add_cmd("BGREWRITEAOF");
    }

    fn bgsave() {
        add_cmd("BGSAVE");
    }

    fn bitcount<K: ToString>(key: K) {
        add_cmd("BITCOUNT");
        add_arg(key);
    }

    fn bitcount_range<K: ToString>(key: K, start_range: i64, end_range: i64) {
        add_cmd("BITCOUNT");
        add_arg(key);
        add_arg(start_range);
        add_arg(end_range);
    }

    fn del<K: ToString>(key: K) {
        add_cmd("DEL");
        add_arg(key);
    }

    fn mdel<K: ToString>(keys: Vec<K>){
        add_cmd("DEL");
        add_args(keys);
    }

    fn discard() {
        add_cmd("DISCARD");
    }

    fn exec() {
        add_cmd("EXEC");
    }

    fn exists<K: ToString>(key: K) {
        add_cmd("EXISTS");
        add_arg(key);
    }

    fn mexists<K: ToString>(keys: Vec<K>){
        add_cmd("EXISTS");
        add_args(keys);
    }

    fn expire<K: ToString>(key: K, expiry: i64) {
        add_cmd("EXPIRE");
        add_arg(key);
        add_arg(expiry);
    }

    fn expireat<K: ToString>(key: K, timestamp: i64) {
        add_cmd("EXPIREAT");
        add_arg(key);
        add_arg(timestamp);
    }

    fn get<K: ToString>(key: K) {
        add_cmd("GET");
        add_arg(key);
    }

    fn getrange<K: ToString>(key: K, start_range: i64, end_range: i64) {
        add_cmd("GETRANGE");
        add_arg(key);
        add_arg(start_range);
        add_arg(end_range);
    }

    fn hdel<K: ToString, F: ToString>(key: K, field: F) {
        add_cmd("HDEL");
        add_arg(key);
        add_arg(field);
    }

    fn hmdel<K: ToString, V: ToString>(key: K, fields: Vec<V>) {
        add_cmd("HDEL");
        add_arg(key);
        add_args(fields);
    }

    fn hexists<K: ToString, F: ToString>(key: K, field: F) {
        add_cmd("HEXISTS");
        add_arg(key);
        add_arg(field);
    }

    fn hget<K: ToString, F: ToString>(key: K, field: F) {
        add_cmd("HGET");
        add_arg(key);
        add_arg(field);
    }

    fn hgetall<K: ToString>(key: K) {
        add_cmd("HGETALL");
        add_arg(key);
    }

    fn hincrby<K: ToString, F: ToString>(key: K, field: F, increment: i64) {
        add_cmd("HINCRBY");
        add_arg(key);
        add_arg(field);
        add_arg(increment);
    }

    fn hincrbyfloat<K: ToString, F: ToString>(key: K, field: F, increment: f64) {
        add_cmd("HINCRBYBYFLOAT");
        add_arg(key);
        add_arg(field);
        add_arg(increment);
    }

    fn hkeys<K: ToString>(key: K) {
        add_cmd("HKEYS");
        add_arg(key);
    }

    fn hlen<K: ToString>(key: K) {
        add_cmd("HLEN");
        add_arg(key);
    }

    fn hmget<K: ToString>(key: K, fields: Vec<String>) {
        add_cmd("HMGET");
        add_arg(key);
        add_args(fields);
    }

    fn hmset<K: ToString>(key: K, fields: HashMap<String, String>) {
        add_cmd("HMSET");
        add_arg(key);
        add_arg_map(fields);
    }

    fn hset<K: ToString, F: ToString, V: ToString>(key: K, field: F, value: V) {
        add_cmd("HSET");
        add_arg(key);
        add_arg(field);
        add_arg(value);
    }

    fn hstrlen<K: ToString, F: ToString>(key: K, field: F) {
        add_cmd("HSTRLEN");
        add_arg(key);
        add_arg(field);
    }

    fn hsetnx<K: ToString, F: ToString, V: ToString>(key: K, field: F, value: V) {
        add_cmd("HSETNX");
        add_arg(key);
        add_arg(field);
        add_arg(value);
    }

    fn hvals<K: ToString>(key: K) {
        add_cmd("HVALS");
        add_arg(key);
    }

    fn multi() {
        add_cmd("MULTI");
    }

    fn select(db_index: i32){
        add_cmd("SELECT");
        add_arg(db_index);
    }

    fn set<K: ToString, V: ToString>(key: K, value: V) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
    }

    fn set_binary<K: ToString>(key: K, value: &[u8]) {
        add_cmd("SET");
        add_arg(key);
        add_binary_arg(value);
    }

    fn setex<K: ToString, V: ToString>(key: K, value: V, expiry: i64) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("EX");
        add_arg(expiry);
    }

    fn psetex<K: ToString, V: ToString>(key: K, value: V, expiry: i64) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("PX");
        add_arg(expiry);
    }

    fn setnx<K: ToString, V: ToString>(key: K, value: V) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("NX");
    }

    fn setxx<K: ToString, V: ToString>(key: K, value: V) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("XX");
    }

    fn setex_nx<K: ToString, V: ToString>(key: K, value: V, expiry: i64) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("EX");
        add_arg(expiry);
        add_arg("NX");
    }

    fn setex_xx<K: ToString, V: ToString>(key: K, value: V, expiry: i64) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("EX");
        add_arg(expiry);
        add_arg("NX");
    }

    fn psetex_nx<K: ToString, V: ToString>(key: K, value: V, expiry: i64) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("PX");
        add_arg(expiry);
        add_arg("NX");
    }

    fn psetex_xx<K: ToString, V: ToString>(key: K, value: V, expiry: i64) {
        add_cmd("SET");
        add_arg(key);
        add_arg(value);
        add_arg("PX");
        add_arg(expiry);
        add_arg("XX");
    }

    fn ttl<K: ToString>(key: K) {
        add_cmd("TTL");
        add_arg(key);
    }

    fn unwatch() {
        add_cmd("UNWATCH");
    }

    fn watch<K: ToString>(key: K) {
        add_cmd("WATCH");
        add_arg(key);
    }

    fn mwatch<K: ToString>(keys: Vec<K>) {
        add_cmd("WATCH");
        add_args(keys);
    }

    fn zadd<K: ToString, V: ToString>(key: K, score: f64, member: V) {
        add_cmd("ZADD");
        add_arg(key);
        add_arg(score);
        add_arg(member);
    }

    fn zadd_binary<K: ToString>(key: K, score: f64, member: &[u8]) {
        add_cmd("ZADD");
        add_arg(key);
        add_arg(score);
        add_binary_arg(member);
    }

    fn zaddnx<K: ToString, V: ToString>(score: f64, key: K, member: V) {
        add_cmd("ZADD");
        add_arg(key);
        add_arg("NX");
        add_arg(score);
        add_arg(member);
    }

    fn zaddxx<K: ToString, V: ToString>(score: f64, key: K, member: V) {
        add_cmd("ZADD");
        add_arg(key);
        add_arg("XX");
        add_arg(score);
        add_arg(member);
    }

    fn zaddnx_ch<K: ToString, V: ToString>(score: f64, key: K, member: V) {
        add_cmd("ZADD");
        add_arg(key);
        add_arg("NX");
        add_arg("CH");
        add_arg(score);
        add_arg(member);
    }

    fn zaddxx_ch<K: ToString, V: ToString>(score: f64, key: K, member: V) {
        add_cmd("ZADD");
        add_arg(key);
        add_arg("XX");
        add_arg("CH");
        add_arg(score);
        add_arg(member);
    }

    fn zcard<K: ToString>(key: K) {
        add_cmd("ZCARD");
        add_arg(key);
    }

    fn zcount<K: ToString, S: ToString, E: ToString>(key: K, start_range: S, end_range: E) {
        add_cmd("ZCOUNT");
        add_arg(key);
        add_arg(start_range);
        add_arg(end_range);
    }

    fn zincrby<K: ToString, V: ToString>(key: K, increment: f64, member: V) {
        add_cmd("ZINCRBY");
        add_arg(key);
        add_arg(increment);
        add_arg(member);
    }

    fn zrange<K: ToString, S: ToString, E: ToString>(key: K, start_range: S, end_range: E) {
        add_cmd("ZRANGE");
        add_arg(key);
        add_arg(start_range);
        add_arg(end_range);
    }

    fn zrange_with_scores<K: ToString, S: ToString, E: ToString>(key: K, start_range: S, end_range: E) {
        add_cmd("ZRANGE");
        add_arg(key);
        add_arg(start_range);
        add_arg(end_range);
        add_arg("WITHSCORES");
    }

}

<|endoftext|>


<commit_before>use token::Token;
use token::Token::*;

struct CompileState {
    output: String,
    loopnum: i32,
    ifnum: i32,
    outbuffsize: i32,
}

fn compile_iter(state: &mut CompileState, tokens: Vec<Token>) {
    let mut outbuffpos = 0;
    for token in tokens {
        match token {
            Add(offset, value) => {
                let dest = if offset == 0 {
                    "%r12".to_string()
                } else {
                    format!("{}(%rbx)", (offset * 8))
                };
                if value == 1 && dest == "%r12" {
                    state.output.push_str("    inc %r12\n");
                } else if value >= 1 {
                    state.output.push_str(&format!("    addq ${}, {}\n", value, dest));
                } else if value == -1 && dest == "%r12" {
                    state.output.push_str("    dec %r12\n");
                } else if value <= -1 {
                    state.output.push_str(&format!("    subq ${}, {}\n", -value, dest));
                }
            }
            MulCopy(src_idx, dest_idx, mul) => {
                let mut src = if src_idx == 0 {
                    "%r12".to_string()
                } else {
                    format!("{}(%rbx)", (src_idx * 8))
                };
                let dest = if dest_idx == 0 {
                    "%r12".to_string()
                } else {
                    format!("{}(%rbx)", (dest_idx * 8))
                };

                if mul != -1 && mul != 1 {
                    state.output.push_str(&format!(concat!("    movq {}, %rax\n",
                                                           "    movq ${}, %rdx\n",
                                                           "    mulq %rdx\n"),
                                                   src,
                                                   mul.abs()));
                    src = "%rax".to_string();
                } else if src != "%r12" && dest != "%r12" {
                    // x86 cannot move memory to memory
                    state.output.push_str(&format!("    movq {}, %rax\n", src));
                    src = "%rax".to_string();
                }

                if mul > 0 {
                    state.output.push_str(&format!("    addq {}, {}\n", src, dest));
                } else {
                    state.output.push_str(&format!("    subq {}, {}\n", src, dest));
                }
            }
            Set(offset, value) => {
                if offset == 0 && value == 0 {
                    state.output.push_str("    xor %r12, %r12\n");
                } else if offset == 0 {
                    state.output.push_str(&format!("    movq ${}, %r12\n", value));
                } else {
                    state.output.push_str(&format!("    movq ${}, {}(%rbx)\n", value, offset * 8));
                }
            }
            Move(offset) => {
                if offset != 0 {
                    state.output.push_str("    movq %r12, (%rbx)\n");
                    if offset > 0 {
                        state.output.push_str(&format!("    addq ${}, %rbx\n", offset * 8));
                    } else {
                        state.output.push_str(&format!("    subq ${}, %rbx\n", -offset * 8));
                    }
                    state.output.push_str("    movq (%rbx), %r12\n");
                }
            }
            Loop(content) => {
                state.loopnum += 1;
                let curloop = state.loopnum;
                state.output.push_str(&format!(concat!("    jmp endloop{}\n", "    loop{}:\n"),
                                               curloop,
                                               curloop));

                compile_iter(state, content);

                state.output.push_str(&format!(concat!("    endloop{}:\n",
                                                       "    test %r12, %r12\n",
                                                       "    jnz loop{}\n"),
                                               curloop,
                                               curloop))
            }
            If(offset, content) => {
                state.ifnum += 1;
                let curif = state.ifnum;
                if offset == 0 {
                    state.output.push_str("    test %r12, %r12\n");
                } else {
                    state.output.push_str(&format!("    cmpq $0, {}(%rbx)\n", offset * 8));
                }
                state.output.push_str(&format!("    jz endif{}\n", curif));

                compile_iter(state, content);

                state.output.push_str(&format!("    endif{}:\n", curif))
            }
            Scan(offset) => {
                // Slighly more optimal than normal loop and move
                state.loopnum += 1;
                state.output.push_str(&format!(concat!("    movq %r12, (%rbx)\n",
                                                       "    jmp endloop{}\n",
                                                       "    loop{}:\n"),
                                               state.loopnum,
                                               state.loopnum));
                if offset > 0 {
                    state.output.push_str(&format!("    addq ${}, %rbx\n", offset * 8));
                } else {
                    state.output.push_str(&format!("    subq ${}, %rbx\n", -offset * 8));
                }
                state.output.push_str(&format!(concat!("    endloop{}:\n",
                                                       "    cmp $0, (%rbx)\n",
                                                       "    jnz loop{}\n",
                                                       "    movq (%rbx), %r12\n"),
                                               state.loopnum,
                                               state.loopnum));
            }
            Input => {
                state.output.push_str(concat!("\n    xor %rax, %rax\n",
                                              "    xor %rdi, %rdi\n",
                                              "    movq %rbx, %rsi\n",
                                              "    movq $1, %rdx\n",
                                              "    syscall\n",
                                              "    movq (%rbx), %r12\n\n"))
            }
            LoadOut(offset, add) => {
                let outaddr = format!("(strbuff+{})", outbuffpos);
                if offset == 0 {
                    state.output.push_str(&format!("    movq %r12, {}\n", outaddr));
                } else {
                    state.output.push_str(&format!("    movq {}(%rbx), %rax\n", offset * 8));
                    state.output.push_str(&format!("    movq %rax, {}\n", outaddr));
                }
                if add > 0 {
                    state.output.push_str(&format!("    addb ${}, {}\n", add, outaddr));
                } else if add < 0 {
                    state.output.push_str(&format!("    subb ${}, {}\n", -add, outaddr));
                }
                outbuffpos += 1;
            }
            LoadOutSet(value) => {
                let outaddr = format!("(strbuff+{})", outbuffpos);
                state.output.push_str(&format!("    movq ${}, {}\n", value, outaddr));
                outbuffpos += 1;
            }
            Output => {
                state.output.push_str(&format!(concat!("    movq $1, %rax\n",
                                                       "    movq $1, %rdi\n",
                                                       "    movq $strbuff, %rsi\n",
                                                       "    movq ${}, %rdx\n",
                                                       "    syscall\n\n"),
                                               outbuffpos));

                if state.outbuffsize < outbuffpos + 8 {
                    state.outbuffsize = outbuffpos + 8;
                }
                outbuffpos = 0;
            }
        }
    }
}

pub fn compile(tokens: Vec<Token>, tape_size: i32) -> String {
    let mut state = CompileState {
        output: String::new(),
        loopnum: 0,
        ifnum: 0,
        outbuffsize: 0,
    };

    compile_iter(&mut state, tokens);

    // Exit syscall
    state.output.push_str(concat!("\n    movq $60, %rax\n",
                                  "    movq $0, %rdi\n",
                                  "    syscall\n"));

    format!(concat!(".section .bss\n",
                    "    .lcomm strbuff, {outbuffsize}\n",
                    "    .lcomm mem, {}\n",
                    "    .set startidx, mem + {}\n",
                    ".section .text\n",
                    ".global _start\n",
                    "_start:\n",
                    "    xor %r12, %r12\n",
                    "    movq $startidx, %rbx\n\n{}"),
            tape_size,
            tape_size / 2,
            state.output,
            outbuffsize = state.outbuffsize)
}
<commit_msg>Combine some code in compile()<commit_after>use token::Token;
use token::Token::*;

struct CompileState {
    output: String,
    loopnum: i32,
    ifnum: i32,
    outbuffsize: i32,
}

fn compile_iter(state: &mut CompileState, tokens: Vec<Token>) {
    let mut outbuffpos = 0;
    for token in tokens {
        match token {
            Add(offset, value) => {
                let dest = if offset == 0 {
                    "%r12".to_string()
                } else {
                    format!("{}(%rbx)", (offset * 8))
                };
                if value == 1 && dest == "%r12" {
                    state.output.push_str("    inc %r12\n");
                } else if value >= 1 {
                    state.output.push_str(&format!("    addq ${}, {}\n", value, dest));
                } else if value == -1 && dest == "%r12" {
                    state.output.push_str("    dec %r12\n");
                } else if value <= -1 {
                    state.output.push_str(&format!("    subq ${}, {}\n", -value, dest));
                }
            }
            MulCopy(src_idx, dest_idx, mul) => {
                let mut src = if src_idx == 0 {
                    "%r12".to_string()
                } else {
                    format!("{}(%rbx)", (src_idx * 8))
                };
                let dest = if dest_idx == 0 {
                    "%r12".to_string()
                } else {
                    format!("{}(%rbx)", (dest_idx * 8))
                };

                if mul != -1 && mul != 1 {
                    state.output.push_str(&format!(concat!("    movq {}, %rax\n",
                                                           "    movq ${}, %rdx\n",
                                                           "    mulq %rdx\n"),
                                                   src,
                                                   mul.abs()));
                    src = "%rax".to_string();
                } else if src != "%r12" && dest != "%r12" {
                    // x86 cannot move memory to memory
                    state.output.push_str(&format!("    movq {}, %rax\n", src));
                    src = "%rax".to_string();
                }

                if mul > 0 {
                    state.output.push_str(&format!("    addq {}, {}\n", src, dest));
                } else {
                    state.output.push_str(&format!("    subq {}, {}\n", src, dest));
                }
            }
            Set(offset, value) => {
                if offset == 0 && value == 0 {
                    state.output.push_str("    xor %r12, %r12\n");
                } else if offset == 0 {
                    state.output.push_str(&format!("    movq ${}, %r12\n", value));
                } else {
                    state.output.push_str(&format!("    movq ${}, {}(%rbx)\n", value, offset * 8));
                }
            }
            Move(offset) => {
                if offset != 0 {
                    state.output.push_str("    movq %r12, (%rbx)\n");
                    if offset > 0 {
                        state.output.push_str(&format!("    addq ${}, %rbx\n", offset * 8));
                    } else {
                        state.output.push_str(&format!("    subq ${}, %rbx\n", -offset * 8));
                    }
                    state.output.push_str("    movq (%rbx), %r12\n");
                }
            }
            Loop(content) => {
                state.loopnum += 1;
                let curloop = state.loopnum;
                state.output.push_str(&format!(concat!("    jmp endloop{}\n", "    loop{}:\n"),
                                               curloop,
                                               curloop));

                compile_iter(state, content);

                state.output.push_str(&format!(concat!("    endloop{}:\n",
                                                       "    test %r12, %r12\n",
                                                       "    jnz loop{}\n"),
                                               curloop,
                                               curloop))
            }
            If(offset, content) => {
                state.ifnum += 1;
                let curif = state.ifnum;
                if offset == 0 {
                    state.output.push_str("    test %r12, %r12\n");
                } else {
                    state.output.push_str(&format!("    cmpq $0, {}(%rbx)\n", offset * 8));
                }
                state.output.push_str(&format!("    jz endif{}\n", curif));

                compile_iter(state, content);

                state.output.push_str(&format!("    endif{}:\n", curif))
            }
            Scan(offset) => {
                // Slighly more optimal than normal loop and move
                state.loopnum += 1;
                state.output.push_str(&format!(concat!("    movq %r12, (%rbx)\n",
                                                       "    jmp endloop{}\n",
                                                       "    loop{}:\n"),
                                               state.loopnum,
                                               state.loopnum));
                if offset > 0 {
                    state.output.push_str(&format!("    addq ${}, %rbx\n", offset * 8));
                } else {
                    state.output.push_str(&format!("    subq ${}, %rbx\n", -offset * 8));
                }
                state.output.push_str(&format!(concat!("    endloop{}:\n",
                                                       "    cmp $0, (%rbx)\n",
                                                       "    jnz loop{}\n",
                                                       "    movq (%rbx), %r12\n"),
                                               state.loopnum,
                                               state.loopnum));
            }
            Input => {
                state.output.push_str(concat!("\n    xor %rax, %rax\n",
                                              "    xor %rdi, %rdi\n",
                                              "    movq %rbx, %rsi\n",
                                              "    movq $1, %rdx\n",
                                              "    syscall\n",
                                              "    movq (%rbx), %r12\n\n"))
            }
            LoadOut(offset, add) => {
                let outaddr = format!("(strbuff+{})", outbuffpos);
                if offset == 0 {
                    state.output.push_str(&format!("    movq %r12, {}\n", outaddr));
                } else {
                    state.output.push_str(&format!("    movq {}(%rbx), %rax\n", offset * 8));
                    state.output.push_str(&format!("    movq %rax, {}\n", outaddr));
                }
                if add > 0 {
                    state.output.push_str(&format!("    addb ${}, {}\n", add, outaddr));
                } else if add < 0 {
                    state.output.push_str(&format!("    subb ${}, {}\n", -add, outaddr));
                }
                outbuffpos += 1;
            }
            LoadOutSet(value) => {
                let outaddr = format!("(strbuff+{})", outbuffpos);
                state.output.push_str(&format!("    movq ${}, {}\n", value, outaddr));
                outbuffpos += 1;
            }
            Output => {
                state.output.push_str(&format!(concat!("    movq $1, %rax\n",
                                                       "    movq $1, %rdi\n",
                                                       "    movq $strbuff, %rsi\n",
                                                       "    movq ${}, %rdx\n",
                                                       "    syscall\n\n"),
                                               outbuffpos));

                if state.outbuffsize < outbuffpos + 8 {
                    state.outbuffsize = outbuffpos + 8;
                }
                outbuffpos = 0;
            }
        }
    }
}

pub fn compile(tokens: Vec<Token>, tape_size: i32) -> String {
    let mut state = CompileState {
        output: String::new(),
        loopnum: 0,
        ifnum: 0,
        outbuffsize: 0,
    };

    compile_iter(&mut state, tokens);

    format!(concat!(".section .bss\n",
                    "    .lcomm strbuff, {outbuffsize}\n",
                    "    .lcomm mem, {}\n",
                    "    .set startidx, mem + {}\n",
                    ".section .text\n",
                    ".global _start\n",
                    "_start:\n",
                    "    xor %r12, %r12\n",
                    "    movq $startidx, %rbx\n\n",
                    "{}\n",
                    // Exit syscall
                    "    movq $60, %rax\n",
                    "    movq $0, %rdi\n",
                    "    syscall\n"),
            tape_size,
            tape_size / 2,
            state.output,
            outbuffsize = state.outbuffsize)
}
<|endoftext|>


<commit_before>// partially (c) 2016 Google Inc. (Lewin Bormann, lewinb@google.com)
//
// See project root for licensing information.
//
use crate::types::TokenInfo;

use std::collections::HashMap;
use std::io;
use std::path::{Path, PathBuf};
use std::sync::Mutex;

use serde::{Deserialize, Serialize};

// The storage layer allows retrieving tokens for scopes that have been
// previously granted tokens. One wrinkle is that a token granted for a set
// of scopes X is also valid for any subset of X's scopes. So when retrieving a
// token for a set of scopes provided by the caller it's beneficial to compare
// that set to all previously stored tokens to see if it is a subset of any
// existing set. To do this efficiently we store a bloom filter along with each
// token that represents the set of scopes the token is associated with. The
// bloom filter allows for efficiently skipping any entries that are
// definitively not a superset.
// The current implementation uses a 64bit bloom filter with 4 hash functions.

/// ScopeHash is a hash value derived from a list of scopes. The hash value
/// represents a fingerprint of the set of scopes *independent* of the ordering.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
struct ScopeHash(u64);

/// ScopeFilter represents a filter for a set of scopes. It can definitively
/// prove that a given list of scopes is not a subset of another.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
struct ScopeFilter(u64);

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum FilterResponse {
    Maybe,
    No,
}

impl ScopeFilter {
    /// Determine if this ScopeFilter could be a subset of the provided filter.
    fn is_subset_of(self, filter: ScopeFilter) -> FilterResponse {
        if self.0 & filter.0 == self.0 {
            FilterResponse::Maybe
        } else {
            FilterResponse::No
        }
    }
}

#[derive(Debug)]
pub(crate) struct ScopeSet<'a, T> {
    hash: ScopeHash,
    filter: ScopeFilter,
    scopes: &'a [T],
}

// Implement Clone manually. Auto derive fails to work correctly because we want
// Clone to be implemented regardless of whether T is Clone or not.
impl<'a, T> Clone for ScopeSet<'a, T> {
    fn clone(&self) -> Self {
        ScopeSet {
            hash: self.hash,
            filter: self.filter,
            scopes: self.scopes,
        }
    }
}
impl<'a, T> Copy for ScopeSet<'a, T> {}

impl<'a, T> ScopeSet<'a, T>
where
    T: AsRef<str>,
{
    // implement an inherent from method even though From is implemented. This
    // is because passing an array ref like &[&str; 1] (&["foo"]) will be auto
    // deref'd to a slice on function boundaries, but it will not implement the
    // From trait. This inherent method just serves to auto deref from array
    // refs to slices and proxy to the From impl.
    pub fn from(scopes: &'a [T]) -> Self {
        let (hash, filter) = scopes.iter().fold(
            (ScopeHash(0), ScopeFilter(0)),
            |(mut scope_hash, mut scope_filter), scope| {
                let h = seahash::hash(scope.as_ref().as_bytes());

                // Use the first 4 6-bit chunks of the seahash as the 4 hash values
                // in the bloom filter.
                for i in 0..4 {
                    // h is a hash derived value in the range 0..64
                    let h = (h >> (6 * i)) & 0b11_1111;
                    scope_filter.0 |= 1 << h;
                }

                // xor the hashes together to get an order independent fingerprint.
                scope_hash.0 ^= h;
                (scope_hash, scope_filter)
            },
        );
        ScopeSet {
            hash,
            filter,
            scopes,
        }
    }
}

pub(crate) enum Storage {
    Memory { tokens: Mutex<JSONTokens> },
    Disk(DiskStorage),
}

impl Storage {
    pub(crate) async fn set<T>(
        &self,
        scopes: ScopeSet<'_, T>,
        token: TokenInfo,
    ) -> Result<(), io::Error>
    where
        T: AsRef<str>,
    {
        match self {
            Storage::Memory { tokens } => tokens.lock().unwrap().set(scopes, token),
            Storage::Disk(disk_storage) => disk_storage.set(scopes, token).await,
        }
    }

    pub(crate) fn get<T>(&self, scopes: ScopeSet<T>) -> Option<TokenInfo>
    where
        T: AsRef<str>,
    {
        match self {
            Storage::Memory { tokens } => tokens.lock().unwrap().get(scopes),
            Storage::Disk(disk_storage) => disk_storage.get(scopes),
        }
    }
}

/// A single stored token.

#[derive(Debug, Clone)]
struct JSONToken {
    scopes: Vec<String>,
    token: TokenInfo,
    hash: ScopeHash,
    filter: ScopeFilter,
}

impl<'de> Deserialize<'de> for JSONToken {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(Deserialize)]
        struct RawJSONToken {
            scopes: Vec<String>,
            token: TokenInfo,
        }
        let RawJSONToken { scopes, token } = RawJSONToken::deserialize(deserializer)?;
        let ScopeSet { hash, filter, .. } = ScopeSet::from(&scopes);
        Ok(JSONToken {
            scopes,
            token,
            hash,
            filter,
        })
    }
}

impl Serialize for JSONToken {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        #[derive(Serialize)]
        struct RawJSONToken<'a> {
            scopes: &'a [String],
            token: &'a TokenInfo,
        }
        RawJSONToken {
            scopes: &self.scopes,
            token: &self.token,
        }
        .serialize(serializer)
    }
}

/// List of tokens in a JSON object
#[derive(Debug, Clone)]
pub(crate) struct JSONTokens {
    token_map: HashMap<ScopeHash, JSONToken>,
}

impl Serialize for JSONTokens {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.collect_seq(self.token_map.values())
    }
}

impl<'de> Deserialize<'de> for JSONTokens {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        struct V;
        impl<'de> serde::de::Visitor<'de> for V {
            type Value = JSONTokens;

            // Format a message stating what data this Visitor expects to receive.
            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
                formatter.write_str("a sequence of JSONToken's")
            }

            fn visit_seq<M>(self, mut access: M) -> Result<Self::Value, M::Error>
            where
                M: serde::de::SeqAccess<'de>,
            {
                let mut token_map = HashMap::with_capacity(access.size_hint().unwrap_or(0));
                while let Some(json_token) = access.next_element::<JSONToken>()? {
                    token_map.insert(json_token.hash, json_token);
                }
                Ok(JSONTokens { token_map })
            }
        }

        // Instantiate our Visitor and ask the Deserializer to drive
        // it over the input data.
        deserializer.deserialize_seq(V)
    }
}

impl JSONTokens {
    pub(crate) fn new() -> Self {
        JSONTokens {
            token_map: HashMap::new(),
        }
    }

    async fn load_from_file(filename: &Path) -> Result<Self, io::Error> {
        let contents = tokio::fs::read(filename).await?;
        serde_json::from_slice(&contents).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
    }

    fn get<T>(
        &self,
        ScopeSet {
            hash,
            filter,
            scopes,
        }: ScopeSet<T>,
    ) -> Option<TokenInfo>
    where
        T: AsRef<str>,
    {
        if let Some(json_token) = self.token_map.get(&hash) {
            return Some(json_token.token.clone());
        }

        let requested_scopes_are_subset_of = |other_scopes: &[String]| {
            scopes
                .iter()
                .all(|s| other_scopes.iter().any(|t| t.as_str() == s.as_ref()))
        };
        // No exact match for the scopes provided. Search for any tokens that
        // exist for a superset of the scopes requested.
        self.token_map
            .values()
            .filter(|json_token| filter.is_subset_of(json_token.filter) == FilterResponse::Maybe)
            .find(|v: &&JSONToken| requested_scopes_are_subset_of(&v.scopes))
            .map(|t: &JSONToken| t.token.clone())
    }

    fn set<T>(
        &mut self,
        ScopeSet {
            hash,
            filter,
            scopes,
        }: ScopeSet<T>,
        token: TokenInfo,
    ) -> Result<(), io::Error>
    where
        T: AsRef<str>,
    {
        use std::collections::hash_map::Entry;
        match self.token_map.entry(hash) {
            Entry::Occupied(mut entry) => {
                entry.get_mut().token = token;
            }
            Entry::Vacant(entry) => {
                let json_token = JSONToken {
                    scopes: scopes.iter().map(|x| x.as_ref().to_owned()).collect(),
                    token,
                    hash,
                    filter,
                };
                entry.insert(json_token.clone());
            }
        }
        Ok(())
    }
}

pub(crate) struct DiskStorage {
    tokens: Mutex<JSONTokens>,
    filename: PathBuf,
}

impl DiskStorage {
    pub(crate) async fn new(filename: PathBuf) -> Result<Self, io::Error> {
        let tokens = match JSONTokens::load_from_file(&filename).await {
            Ok(tokens) => tokens,
            Err(e) if e.kind() == io::ErrorKind::NotFound => JSONTokens::new(),
            Err(e) => return Err(e),
        };

        Ok(DiskStorage {
            tokens: Mutex::new(tokens),
            filename,
        })
    }

    pub(crate) async fn set<T>(
        &self,
        scopes: ScopeSet<'_, T>,
        token: TokenInfo,
    ) -> Result<(), io::Error>
    where
        T: AsRef<str>,
    {
        let json = {
            use std::ops::Deref;
            let mut lock = self.tokens.lock().unwrap();
            lock.set(scopes, token)?;
            serde_json::to_string(lock.deref())
                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
        };
        tokio::fs::write(self.filename.clone(), json).await
    }

    pub(crate) fn get<T>(&self, scopes: ScopeSet<T>) -> Option<TokenInfo>
    where
        T: AsRef<str>,
    {
        self.tokens.lock().unwrap().get(scopes)
    }
}

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

    #[test]
    fn test_scope_filter() {
        let foo = ScopeSet::from(&["foo"]).filter;
        let bar = ScopeSet::from(&["bar"]).filter;
        let foobar = ScopeSet::from(&["foo", "bar"]).filter;

        // foo and bar are both subsets of foobar. This condition should hold no
        // matter what changes are made to the bloom filter implementation.
        assert!(foo.is_subset_of(foobar) == FilterResponse::Maybe);
        assert!(bar.is_subset_of(foobar) == FilterResponse::Maybe);

        // These conditions hold under the current bloom filter implementation
        // because "foo" and "bar" don't collide, but if the bloom filter
        // implementations change it could be valid for them to return Maybe.
        assert!(foo.is_subset_of(bar) == FilterResponse::No);
        assert!(bar.is_subset_of(foo) == FilterResponse::No);
        assert!(foobar.is_subset_of(foo) == FilterResponse::No);
        assert!(foobar.is_subset_of(bar) == FilterResponse::No);
    }

    #[tokio::test]
    async fn test_disk_storage() {
        let new_token = |access_token: &str| TokenInfo {
            access_token: access_token.to_owned(),
            refresh_token: None,
            expires_at: None,
        };
        let scope_set = ScopeSet::from(&["myscope"]);
        let tempdir = tempfile::tempdir().unwrap();
        {
            let storage = DiskStorage::new(tempdir.path().join("tokenstorage.json"))
                .await
                .unwrap();
            assert!(storage.get(scope_set).is_none());
            storage
                .set(scope_set, new_token("my_access_token"))
                .await
                .unwrap();
            assert_eq!(storage.get(scope_set), Some(new_token("my_access_token")));
        }
        {
            // Create a new DiskStorage instance and verify the tokens were read from disk correctly.
            let storage = DiskStorage::new(tempdir.path().join("tokenstorage.json"))
                .await
                .unwrap();
            assert_eq!(storage.get(scope_set), Some(new_token("my_access_token")));
        }
    }
}
<commit_msg>Create the token file with more secure permissions on unix.<commit_after>// partially (c) 2016 Google Inc. (Lewin Bormann, lewinb@google.com)
//
// See project root for licensing information.
//
use crate::types::TokenInfo;

use std::collections::HashMap;
use std::io;
use std::path::{Path, PathBuf};
use std::sync::Mutex;

use serde::{Deserialize, Serialize};

// The storage layer allows retrieving tokens for scopes that have been
// previously granted tokens. One wrinkle is that a token granted for a set
// of scopes X is also valid for any subset of X's scopes. So when retrieving a
// token for a set of scopes provided by the caller it's beneficial to compare
// that set to all previously stored tokens to see if it is a subset of any
// existing set. To do this efficiently we store a bloom filter along with each
// token that represents the set of scopes the token is associated with. The
// bloom filter allows for efficiently skipping any entries that are
// definitively not a superset.
// The current implementation uses a 64bit bloom filter with 4 hash functions.

/// ScopeHash is a hash value derived from a list of scopes. The hash value
/// represents a fingerprint of the set of scopes *independent* of the ordering.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
struct ScopeHash(u64);

/// ScopeFilter represents a filter for a set of scopes. It can definitively
/// prove that a given list of scopes is not a subset of another.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
struct ScopeFilter(u64);

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum FilterResponse {
    Maybe,
    No,
}

impl ScopeFilter {
    /// Determine if this ScopeFilter could be a subset of the provided filter.
    fn is_subset_of(self, filter: ScopeFilter) -> FilterResponse {
        if self.0 & filter.0 == self.0 {
            FilterResponse::Maybe
        } else {
            FilterResponse::No
        }
    }
}

#[derive(Debug)]
pub(crate) struct ScopeSet<'a, T> {
    hash: ScopeHash,
    filter: ScopeFilter,
    scopes: &'a [T],
}

// Implement Clone manually. Auto derive fails to work correctly because we want
// Clone to be implemented regardless of whether T is Clone or not.
impl<'a, T> Clone for ScopeSet<'a, T> {
    fn clone(&self) -> Self {
        ScopeSet {
            hash: self.hash,
            filter: self.filter,
            scopes: self.scopes,
        }
    }
}
impl<'a, T> Copy for ScopeSet<'a, T> {}

impl<'a, T> ScopeSet<'a, T>
where
    T: AsRef<str>,
{
    // implement an inherent from method even though From is implemented. This
    // is because passing an array ref like &[&str; 1] (&["foo"]) will be auto
    // deref'd to a slice on function boundaries, but it will not implement the
    // From trait. This inherent method just serves to auto deref from array
    // refs to slices and proxy to the From impl.
    pub fn from(scopes: &'a [T]) -> Self {
        let (hash, filter) = scopes.iter().fold(
            (ScopeHash(0), ScopeFilter(0)),
            |(mut scope_hash, mut scope_filter), scope| {
                let h = seahash::hash(scope.as_ref().as_bytes());

                // Use the first 4 6-bit chunks of the seahash as the 4 hash values
                // in the bloom filter.
                for i in 0..4 {
                    // h is a hash derived value in the range 0..64
                    let h = (h >> (6 * i)) & 0b11_1111;
                    scope_filter.0 |= 1 << h;
                }

                // xor the hashes together to get an order independent fingerprint.
                scope_hash.0 ^= h;
                (scope_hash, scope_filter)
            },
        );
        ScopeSet {
            hash,
            filter,
            scopes,
        }
    }
}

pub(crate) enum Storage {
    Memory { tokens: Mutex<JSONTokens> },
    Disk(DiskStorage),
}

impl Storage {
    pub(crate) async fn set<T>(
        &self,
        scopes: ScopeSet<'_, T>,
        token: TokenInfo,
    ) -> Result<(), io::Error>
    where
        T: AsRef<str>,
    {
        match self {
            Storage::Memory { tokens } => tokens.lock().unwrap().set(scopes, token),
            Storage::Disk(disk_storage) => disk_storage.set(scopes, token).await,
        }
    }

    pub(crate) fn get<T>(&self, scopes: ScopeSet<T>) -> Option<TokenInfo>
    where
        T: AsRef<str>,
    {
        match self {
            Storage::Memory { tokens } => tokens.lock().unwrap().get(scopes),
            Storage::Disk(disk_storage) => disk_storage.get(scopes),
        }
    }
}

/// A single stored token.

#[derive(Debug, Clone)]
struct JSONToken {
    scopes: Vec<String>,
    token: TokenInfo,
    hash: ScopeHash,
    filter: ScopeFilter,
}

impl<'de> Deserialize<'de> for JSONToken {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        #[derive(Deserialize)]
        struct RawJSONToken {
            scopes: Vec<String>,
            token: TokenInfo,
        }
        let RawJSONToken { scopes, token } = RawJSONToken::deserialize(deserializer)?;
        let ScopeSet { hash, filter, .. } = ScopeSet::from(&scopes);
        Ok(JSONToken {
            scopes,
            token,
            hash,
            filter,
        })
    }
}

impl Serialize for JSONToken {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        #[derive(Serialize)]
        struct RawJSONToken<'a> {
            scopes: &'a [String],
            token: &'a TokenInfo,
        }
        RawJSONToken {
            scopes: &self.scopes,
            token: &self.token,
        }
        .serialize(serializer)
    }
}

/// List of tokens in a JSON object
#[derive(Debug, Clone)]
pub(crate) struct JSONTokens {
    token_map: HashMap<ScopeHash, JSONToken>,
}

impl Serialize for JSONTokens {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.collect_seq(self.token_map.values())
    }
}

impl<'de> Deserialize<'de> for JSONTokens {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        struct V;
        impl<'de> serde::de::Visitor<'de> for V {
            type Value = JSONTokens;

            // Format a message stating what data this Visitor expects to receive.
            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
                formatter.write_str("a sequence of JSONToken's")
            }

            fn visit_seq<M>(self, mut access: M) -> Result<Self::Value, M::Error>
            where
                M: serde::de::SeqAccess<'de>,
            {
                let mut token_map = HashMap::with_capacity(access.size_hint().unwrap_or(0));
                while let Some(json_token) = access.next_element::<JSONToken>()? {
                    token_map.insert(json_token.hash, json_token);
                }
                Ok(JSONTokens { token_map })
            }
        }

        // Instantiate our Visitor and ask the Deserializer to drive
        // it over the input data.
        deserializer.deserialize_seq(V)
    }
}

impl JSONTokens {
    pub(crate) fn new() -> Self {
        JSONTokens {
            token_map: HashMap::new(),
        }
    }

    async fn load_from_file(filename: &Path) -> Result<Self, io::Error> {
        let contents = tokio::fs::read(filename).await?;
        serde_json::from_slice(&contents).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
    }

    fn get<T>(
        &self,
        ScopeSet {
            hash,
            filter,
            scopes,
        }: ScopeSet<T>,
    ) -> Option<TokenInfo>
    where
        T: AsRef<str>,
    {
        if let Some(json_token) = self.token_map.get(&hash) {
            return Some(json_token.token.clone());
        }

        let requested_scopes_are_subset_of = |other_scopes: &[String]| {
            scopes
                .iter()
                .all(|s| other_scopes.iter().any(|t| t.as_str() == s.as_ref()))
        };
        // No exact match for the scopes provided. Search for any tokens that
        // exist for a superset of the scopes requested.
        self.token_map
            .values()
            .filter(|json_token| filter.is_subset_of(json_token.filter) == FilterResponse::Maybe)
            .find(|v: &&JSONToken| requested_scopes_are_subset_of(&v.scopes))
            .map(|t: &JSONToken| t.token.clone())
    }

    fn set<T>(
        &mut self,
        ScopeSet {
            hash,
            filter,
            scopes,
        }: ScopeSet<T>,
        token: TokenInfo,
    ) -> Result<(), io::Error>
    where
        T: AsRef<str>,
    {
        use std::collections::hash_map::Entry;
        match self.token_map.entry(hash) {
            Entry::Occupied(mut entry) => {
                entry.get_mut().token = token;
            }
            Entry::Vacant(entry) => {
                let json_token = JSONToken {
                    scopes: scopes.iter().map(|x| x.as_ref().to_owned()).collect(),
                    token,
                    hash,
                    filter,
                };
                entry.insert(json_token.clone());
            }
        }
        Ok(())
    }
}

pub(crate) struct DiskStorage {
    tokens: Mutex<JSONTokens>,
    filename: PathBuf,
}

impl DiskStorage {
    pub(crate) async fn new(filename: PathBuf) -> Result<Self, io::Error> {
        let tokens = match JSONTokens::load_from_file(&filename).await {
            Ok(tokens) => tokens,
            Err(e) if e.kind() == io::ErrorKind::NotFound => JSONTokens::new(),
            Err(e) => return Err(e),
        };

        Ok(DiskStorage {
            tokens: Mutex::new(tokens),
            filename,
        })
    }

    pub(crate) async fn set<T>(
        &self,
        scopes: ScopeSet<'_, T>,
        token: TokenInfo,
    ) -> Result<(), io::Error>
    where
        T: AsRef<str>,
    {
        use tokio::io::AsyncWriteExt;
        let json = {
            use std::ops::Deref;
            let mut lock = self.tokens.lock().unwrap();
            lock.set(scopes, token)?;
            serde_json::to_string(lock.deref())
                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
        };
        let mut f = open_writeable_file(&self.filename).await?;
        f.write_all(json.as_bytes()).await?;
        Ok(())
    }

    pub(crate) fn get<T>(&self, scopes: ScopeSet<T>) -> Option<TokenInfo>
    where
        T: AsRef<str>,
    {
        self.tokens.lock().unwrap().get(scopes)
    }
}

#[cfg(unix)]
async fn open_writeable_file(
    filename: impl AsRef<Path>,
) -> Result<tokio::fs::File, tokio::io::Error> {
    // Ensure if the file is created it's only readable and writable by the
    // current user.
    use std::os::unix::fs::OpenOptionsExt;
    let opts: tokio::fs::OpenOptions = {
        let mut opts = std::fs::OpenOptions::new();
        opts.write(true).create(true).truncate(true).mode(0o600);
        opts.into()
    };
    opts.open(filename).await
}

#[cfg(not(unix))]
async fn open_writeable_file(
    filename: impl AsRef<Path>,
) -> Result<tokio::fs::File, tokio::io::Error> {
    // I don't have knowledge of windows or other platforms to know how to
    // create a file that's only readable by the current user.
    tokio::fs::File::create(filename).await
}

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

    #[test]
    fn test_scope_filter() {
        let foo = ScopeSet::from(&["foo"]).filter;
        let bar = ScopeSet::from(&["bar"]).filter;
        let foobar = ScopeSet::from(&["foo", "bar"]).filter;

        // foo and bar are both subsets of foobar. This condition should hold no
        // matter what changes are made to the bloom filter implementation.
        assert!(foo.is_subset_of(foobar) == FilterResponse::Maybe);
        assert!(bar.is_subset_of(foobar) == FilterResponse::Maybe);

        // These conditions hold under the current bloom filter implementation
        // because "foo" and "bar" don't collide, but if the bloom filter
        // implementations change it could be valid for them to return Maybe.
        assert!(foo.is_subset_of(bar) == FilterResponse::No);
        assert!(bar.is_subset_of(foo) == FilterResponse::No);
        assert!(foobar.is_subset_of(foo) == FilterResponse::No);
        assert!(foobar.is_subset_of(bar) == FilterResponse::No);
    }

    #[tokio::test]
    async fn test_disk_storage() {
        let new_token = |access_token: &str| TokenInfo {
            access_token: access_token.to_owned(),
            refresh_token: None,
            expires_at: None,
        };
        let scope_set = ScopeSet::from(&["myscope"]);
        let tempdir = tempfile::tempdir().unwrap();
        {
            let storage = DiskStorage::new(tempdir.path().join("tokenstorage.json"))
                .await
                .unwrap();
            assert!(storage.get(scope_set).is_none());
            storage
                .set(scope_set, new_token("my_access_token"))
                .await
                .unwrap();
            assert_eq!(storage.get(scope_set), Some(new_token("my_access_token")));
        }
        {
            // Create a new DiskStorage instance and verify the tokens were read from disk correctly.
            let storage = DiskStorage::new(tempdir.path().join("tokenstorage.json"))
                .await
                .unwrap();
            assert_eq!(storage.get(scope_set), Some(new_token("my_access_token")));
        }
    }
}
<|endoftext|>


<commit_before>/* Copyright 2017 David Li.

 * This file is part of BigRedRPG.

 * BigRedRPG is free software: you can redistribute it and/or modify it
 * under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.

 * BigRedRPG is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.

 * You should have received a copy of the GNU Affero General Public License
 * along with BigRedRPG.  If not, see <http://www.gnu.org/licenses/>.
 */

use std::collections::HashMap;

use location::LocationId;
use nature::Nature;

#[derive(Clone,Copy,Debug,Eq,PartialEq)]
pub enum StatusKind {
    Stress,
    Boredom,
    Exhaustion,
    Grades,
}

#[derive(Clone)]
pub struct StatusModifier {
    pub description: String,
    pub modifiers: Vec<(StatusKind, i32)>,
}

pub struct Status {
    base: u32,
    range: (u32, u32),
    current: i32,
}

impl Status {
    pub fn new(base: u32) -> Status {
        if base > 100 {
            panic!("Base status value outside range (0, 100): {}", base);
        }
        Status {
            base: base,
            range: (0, 100),
            current: base as i32,
        }
    }
}

pub struct Student {
    pub id: u32,
    pub name: String,
    pub major: String,
    pub dorm: LocationId,

    stress: Status,
    boredom: Status,
    exhaustion: Status,
    grades: Status,

    modifiers: Vec<StatusModifier>,
    natures: HashMap<String, Nature>,
}

impl Student {
    pub fn new<S: Into<String>>(id: u32, name: S, major: S, dorm: LocationId) -> Student {
        Student {
            id: id,
            name: name.into(),
            major: major.into(),
            dorm: dorm,

            stress: Status::new(0),
            boredom: Status::new(0),
            exhaustion: Status::new(0),
            grades: Status::new(100),

            modifiers: Vec::new(),
            natures: HashMap::new(),
        }
    }

    pub fn apply_modifier(&mut self, modifier: StatusModifier) {
        for &(status, amount) in modifier.modifiers.iter() {
            let mut stat = match status {
                StatusKind::Stress => &mut self.stress,
                StatusKind::Boredom => &mut self.boredom,
                StatusKind::Exhaustion => &mut self.exhaustion,
                StatusKind::Grades => &mut self.grades,
            };
            stat.current += amount;
        }

        self.modifiers.push(modifier);
    }

    pub fn has_nature<'a, S: Into<&'a str>>(&self, nature: S) -> bool {
        self.natures.contains_key(nature.into())
    }
}
<commit_msg>refactor(Student): add StudentId<commit_after>/* Copyright 2017 David Li.

 * This file is part of BigRedRPG.

 * BigRedRPG is free software: you can redistribute it and/or modify it
 * under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.

 * BigRedRPG is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.

 * You should have received a copy of the GNU Affero General Public License
 * along with BigRedRPG.  If not, see <http://www.gnu.org/licenses/>.
 */

use std::collections::HashMap;

use location::LocationId;
use nature::Nature;

pub type StudentId = u32;

#[derive(Clone,Copy,Debug,Eq,PartialEq)]
pub enum StatusKind {
    Stress,
    Boredom,
    Exhaustion,
    Grades,
}

#[derive(Clone)]
pub struct StatusModifier {
    pub description: String,
    pub modifiers: Vec<(StatusKind, i32)>,
}

pub struct Status {
    base: u32,
    range: (u32, u32),
    current: i32,
}

impl Status {
    pub fn new(base: u32) -> Status {
        if base > 100 {
            panic!("Base status value outside range (0, 100): {}", base);
        }
        Status {
            base: base,
            range: (0, 100),
            current: base as i32,
        }
    }
}

pub struct Student {
    pub id: StudentId,
    pub name: String,
    pub major: String,
    pub dorm: LocationId,

    stress: Status,
    boredom: Status,
    exhaustion: Status,
    grades: Status,

    modifiers: Vec<StatusModifier>,
    natures: HashMap<String, Nature>,
}

impl Student {
    pub fn new<S: Into<String>>(id: StudentId, name: S, major: S, dorm: LocationId) -> Student {
        Student {
            id: id,
            name: name.into(),
            major: major.into(),
            dorm: dorm,

            stress: Status::new(0),
            boredom: Status::new(0),
            exhaustion: Status::new(0),
            grades: Status::new(100),

            modifiers: Vec::new(),
            natures: HashMap::new(),
        }
    }

    pub fn apply_modifier(&mut self, modifier: StatusModifier) {
        for &(status, amount) in modifier.modifiers.iter() {
            let mut stat = match status {
                StatusKind::Stress => &mut self.stress,
                StatusKind::Boredom => &mut self.boredom,
                StatusKind::Exhaustion => &mut self.exhaustion,
                StatusKind::Grades => &mut self.grades,
            };
            stat.current += amount;
        }

        self.modifiers.push(modifier);
    }

    pub fn has_nature<'a, S: Into<&'a str>>(&self, nature: S) -> bool {
        self.natures.contains_key(nature.into())
    }
}
<|endoftext|>


<commit_before>use luminance::{Dim2, Flat, Sampler};
use luminance_gl::gl33::Texture;
use image::{self, ImageResult};
use std::path::Path;

pub use luminance::RGBA32F;

pub type TextureImage<F> = Texture<Flat, Dim2, F>;

/// Load an RGBA texture from an image at a path.
pub fn load_rgba_texture<P>(path: P, sampler: &Sampler) -> ImageResult<TextureImage<RGBA32F>> where P: AsRef<Path> {
  let image = try!(image::open(path)).to_rgba();
  let dim = image.dimensions();
  let raw: Vec<f32> = image.into_raw().into_iter().map(|x| x as f32 / 255.).collect();

  let tex = Texture::new(dim, 0, sampler);
  tex.upload_raw(false, &raw);

  Ok(tex)
}

/// Save an RGBA image on disk.
pub fn save_rgba_texture<P>(texture: &TextureImage<RGBA32F>, path: P) where P: AsRef<Path> {
  let texels = texture.get_raw_texels();
  let (w, h) = texture.size;
  let mut output = Vec::with_capacity((w * h) as usize);

  for texel in &texels {
    output.push((texel * 255.) as u8);
  }

  let _ = image::save_buffer(path, &output, w, h, image::ColorType::RGBA(8));
}
<commit_msg>Support for sRGB textures.<commit_after>use luminance::{Dim2, Flat, Sampler};
use luminance_gl::gl33::Texture;
use image::{self, ImageResult};
use std::path::Path;

pub use luminance::RGBA32F;

pub type TextureImage<F> = Texture<Flat, Dim2, F>;

/// Load an RGBA texture from an image at a path.
pub fn load_rgba_texture<P>(path: P, sampler: &Sampler, linear: bool) -> ImageResult<TextureImage<RGBA32F>> where P: AsRef<Path> {
  let image = try!(image::open(path)).to_rgba();
  let dim = image.dimensions();
  let raw: Vec<f32> = image.into_raw().into_iter().map(|x| {
    let y = x as f32 / 255.;

    if linear {
      y
    } else {
      y.powf(1. / 2.2)
    }
  }).collect();

  let tex = Texture::new(dim, 0, sampler);
  tex.upload_raw(false, &raw);

  Ok(tex)
}

/// Save an RGBA image on disk.
pub fn save_rgba_texture<P>(texture: &TextureImage<RGBA32F>, path: P) where P: AsRef<Path> {
  let texels = texture.get_raw_texels();
  let (w, h) = texture.size;
  let mut output = Vec::with_capacity((w * h) as usize);

  for texel in &texels {
    output.push((texel * 255.) as u8);
  }

  let _ = image::save_buffer(path, &output, w, h, image::ColorType::RGBA(8));
}
<|endoftext|>