<commit_before>extern crate tempfile;

use std::fmt::{ Display, Formatter, Error };
use std::iter::Iterator;
use std::collections::HashMap;
use std::collections::HashSet;
use std::vec::Vec;
use std::io::prelude::*;
use std::io::BufReader;
use std::fs::File;

use std::clone::Clone;

/**
 * test can be deleted.
 * clauses will house all clauses, key is int so that we can use occurrences
 * ID for simplifications later
 * occurrences tracks which clauses literals are used in
 */
#[derive(Debug)]
pub struct CNF {
    pub nvar        : i32,
    pub nclause     : i32,
    pub clauses     : HashMap<i32, Vec<i32>>,
    pub occurrences : HashMap<i32, Vec<i32>>,
    pub units       : HashSet<i32>
}

impl CNF {
    pub fn new(nvar: i32, nclause: i32) -> CNF {
        CNF {
            nvar        : nvar,
            nclause     : nclause,
            clauses     : HashMap::new(),
            occurrences : HashMap::new(),
            units       : HashSet::new(),
        }
    }

    /// Add a clause, return the ID of the inserted clause
    /// Right now, this isn't public; api is odd as we have an odd representation of literals.
    /// TODO: Mask this with public function?
    fn add_clause(&mut self, clause: Vec<i32>) -> i32 {
        assert!(clause.len() > 0);
        let id: i32 = self.clauses.len() as i32;
        if clause.len() == 1 { self.units.insert(id); }
        for var in &clause {
            let occ = self.occurrences.entry(*var).or_insert(Vec::new());
            occ.push(id);
        }
        self.clauses.insert(id, clause);
        id
    }

    /// Unit propagation: propagate the literal associated with the clause id `clause`, and remove
    /// it from `self.units`.
    pub fn unit_propagate(&mut self, unit: i32) {
        // Remove clauses with lit and remove lit from occurrences.
        let lit = self.clauses.get(&unit).expect("unit clause not found")[0];
        self.propagate(lit);
        // This needs to be true, i.e. unit clause id needs to be in the unit clauses
        assert!(self.units.remove(&unit));
    }

    /// Remove all clauses containing literal `lit` from the CNF, and remove the negation of `lit`
    /// from the remaining clauses.
    /// TODO: unit_propagate and propagate should return true/false
    /// depending on whether the updated CNF problem is satisfiable
    pub fn propagate(&mut self, lit: i32) {
        // Remove clauses with lit and remove lit from occurrences.
        for occ in &self.occurrences.remove(&lit).unwrap() {
            self.clauses.remove(occ);
        }
    }
}

impl Clone for CNF {
    fn clone(&self) -> CNF {
        CNF {
            nvar        : self.nvar,
            nclause     : self.nclause,
            clauses     : self.clauses.clone(),
            occurrences : self.occurrences.clone(),
            units       : self.units.clone()
        }
    }
}

impl Display for CNF {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        let formatted = format!("Nvar: {:?} Nclause: {:?}\nClauses: {:?}\nOccurrrences: {:?}",
                self.nvar, self.nclause,
                self.clauses, self.occurrences);
        write!(f, "{}", formatted)
    }
}

// End CNF

// Begin Assignemnts

pub type Assignment = Vec<bool>;

pub struct PartialAssignment {
    pub assignment : Vec<Option<bool>>,
    pub unassigned : HashSet<i32>
}

impl PartialAssignment {
    pub fn new(n: usize) -> PartialAssignment {
        PartialAssignment {
            assignment : vec!(None; n),
            unassigned : (0..n as i32).collect()
        }
    }

    fn assign(&mut self, v: usize, assn: bool) {
        // TODO: Error check
        self.assignment[v] = Some(assn);
    }

    pub fn assign_literal(&mut self, lit: i32){
        let polarity: bool = lit & 1 == 0;
        let v: usize = (lit >> 1) as usize;
        self.assign(v - 1, polarity);
    }

    fn unassign(&mut self, v: usize) {
        self.assignment[v] = None;
    }

    pub fn unassign_literal(&mut self, lit: i32){
        let v: usize = (lit >> 1) as usize;
        self.unassign(v - 1);
    }
}

// TODO: Implement Display
impl Display for PartialAssignment {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        let formatted = format!("Assignemnts: {:?}\nUnassigned: {:?}",
                self.assignment, self.unassigned);
        write!(f, "{}", formatted)
    }
}

// End Assignments

// Begin Parsing

// TODO: rename to from_dimacs or somehow isolate parsing
fn parse_dimacs(reader: &mut BufReader<File>) -> Result<CNF, &'static str> {
    let mut line_iterator = reader.lines();

    let mut nvar: i32 = -1;
    let mut nclause: i32 = -1;
    // DIMACS file must have a problem statement before other lines.
    // This first loop searches for the problem statement.
    for line in &mut line_iterator {
        let line = line.expect("could not read file");  // Unwrap result
        let words: Vec<&str> = line.split_whitespace().collect();
        if words.len() == 0 { continue; } // Ignore empty lines
        match words[0] {
            "c" => { }
            "p" => { // Problem statement
                // Must have format "p cnf nvar nclause"
                if words.len() != 4 || words[1] != "cnf" {
                    return Err("invalid problem statement");
                }
                nvar = words[2].parse()
                    .expect(&format!("invalid number of variables {}", words[2]));
                nclause = words[3].parse()
                    .expect(&format!("invalid number of clauses {}", words[3]));
                break;
            }
            // TODO: Add words[0] to this error message
            _ => { return Err("unknown statement beginning"); }
        }
    }
    // Then nvar, nclause were never initialized
    if nvar == -1 || nclause == -1 {
       return Err("no problem statement found");
    }
    // TODO: Different errors for different descriptions (hence why I've split up this if
    // statement)
    if nvar == 0 || nclause == 0 {
        return Err("invalid number of literals in problem")
    }

    // Initialize CNF and parse the rest of the file
    let mut cnf = CNF::new(nvar, nclause);
    let mut clauses_read: i32 = 0;
    for line in &mut line_iterator {
        let line = line.unwrap();
        let words: Vec<&str> = line.split_whitespace().collect();
        if words.len() == 0 { continue; }
        match words[0] {
            "c" => { }
            "p" => { return Err("duplicate problem statement"); }
            _   => {
                clauses_read = clauses_read + 1;
                if clauses_read > nclause { return Err("too many clauses in file"); }
                let tokens: Vec<i32> = words.iter()
                    .filter_map(|s| {
                        let n = s.parse::<i32>().unwrap();
                        // FIXME: This ignores zeros not just as line enders but in the formulas
                        // themselves. Split on zeros at the end here.
                        if n == 0 { return None; }
                        // FIXME: This should return an error instead of
                        // panicking, but I can't return an error
                        // in a closure
                        if n > nvar { panic!("variable out of range: {}", n); }
                        Some(if n < 0 { (-n) << 1 | 1 } else { n << 1 })
                    })
                    .collect();
                cnf.add_clause(tokens);
            }
        }
    }
    // Double check that the number of clauses read is equal
    if clauses_read != nclause { return Err("too few clauses in file"); }
    Ok(cnf)
}

pub fn parse_dimacs_file(f: File) -> Result<CNF, &'static str> {
    // Read the file
    let mut reader = BufReader::new(f);

    // TODO: This is definitely not the correct way to handle errors.
    // Should parse_dimacs have options for returning an error AND panicking?
    let line = parse_dimacs(&mut reader);
    line
}

#[cfg(test)]
mod tests {
    extern crate tempfile;
    use std::fs::File;
    use std::io::prelude::*;
    use std::io::SeekFrom;

    use super::parse_dimacs_file;

    macro_rules! create_tempfile {
        ($x: expr) => {{
            let mut tmpfile: File = tempfile::tempfile().unwrap();
            let _ = write!(tmpfile, $x);
            tmpfile.seek(SeekFrom::Start(0)).unwrap();

            tmpfile
        }};
    }

    #[test]
    #[should_panic(expected = "out of range")]
    fn variable_out_of_range() {
        let tmpfile = create_tempfile!("
            p cnf 2 3
            1 2 0
            4 1 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    #[should_panic(expected = "invalid number")]
    fn invalid_nvar() {
        let tmpfile = create_tempfile!("
            p cnf gd 3
            1 2 0
            4 1 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    #[should_panic(expected = "invalid number")]
    fn invalid_nclause() {
        let tmpfile = create_tempfile!("
            p cnf 2 gdd
            1 2 0
            4 1 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    #[should_panic(expected = "too many")]
    fn too_many_clauses() {
        let tmpfile = create_tempfile!("
            p cnf 5 5
            1 2 0
            2 3 0
            3 4 0
            4 5 0
            1 3 0
            2 4 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    #[should_panic(expected = "too few")]
    fn too_little_clauses() {
        let tmpfile = create_tempfile!("
            p cnf 5 5
            1 2 0
            2 3 0
            3 4 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    fn should_work() {
        let tmpfile = create_tempfile!("
            p cnf 3 3
            1 2 0
            2 3 0
            1 -3 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    fn should_ignore_comments() {
        let tmpfile = create_tempfile!("
            c comment
            p cnf 3 3
            c comment
            1 2 0
            2 3 0
            c comment
            1 -3 0
            c comment
            c comment
            c comment
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    #[should_panic(expected = "duplicate problem")]
    fn duplicate_problem_statement() {
        let tmpfile = create_tempfile!("
            p cnf 3 3
            1 2 0
            p cnf 3 3
            2 3 0
            1 -3 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    #[should_panic(expected = "duplicate problem")]
    fn duplicate_problem_statement_2() {
        let tmpfile = create_tempfile!("
            p cnf 3 3
            p cnf 3 3
            1 2 0
            2 3 0
            1 -3 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    fn adds_unit_clauses() {
        let tmpfile = create_tempfile!("
            p cnf 3 4
            1 0
            2 0
            3 0
            -1 -3 0
        ");
        let cnf = parse_dimacs_file(tmpfile).unwrap();
        assert_eq!(cnf.units.len(), 3);
    }

    // TODO: Make this doc test (with unit propagate); create_tempfile may need to be an importable
    // macro up there
    #[test]
    fn unit_propagate_works() {
        let tmpfile = create_tempfile!("
            p cnf 2 2
            1 0
            1 2 0
        ");
        let mut cnf = parse_dimacs_file(tmpfile).unwrap();
        cnf.unit_propagate(0);
        assert!(cnf.units.is_empty());
    }

    #[test]
    fn unit_propagate_works_2() {
        let tmpfile = create_tempfile!("
            p cnf 4 4
            1 0
            1 2 0
            2 3 0
            4 0
        ");
        let mut cnf = parse_dimacs_file(tmpfile).unwrap();
        cnf.unit_propagate(0);
        // Only one unit clause left
        assert_eq!(cnf.units.len(), 1);
        cnf.unit_propagate(3);
        // No more unit clauses
        assert!(cnf.units.is_empty());
    }

    #[test]
    #[should_panic(expected = "unit clause not found")]
    fn unit_clause_not_found() {
        let tmpfile = create_tempfile!("
            p cnf 2 2
            1 0
            1 2 0
        ");
        let mut cnf = parse_dimacs_file(tmpfile).unwrap();
        cnf.unit_propagate(2);
    }
}
<commit_msg>Add "zeroth" macro to get first element of set<commit_after>extern crate tempfile;

use std::fmt::{ Display, Formatter, Error };
use std::iter::Iterator;
use std::collections::HashMap;
use std::collections::HashSet;
use std::vec::Vec;
use std::io::prelude::*;
use std::io::BufReader;
use std::fs::File;

use std::clone::Clone;

/// Get the (arbitrary) zeroth element of a hashset.
#[marco_export]
macro_rules! zeroth {
    ($hs: expr) => {{
        *$hs.iter().next().unwrap()
    }}
}

/**
 * test can be deleted.
 * clauses will house all clauses, key is int so that we can use occurrences
 * ID for simplifications later
 * occurrences tracks which clauses literals are used in
 */
#[derive(Debug)]
pub struct CNF {
    pub nvar        : i32,
    pub nclause     : i32,
    pub clauses     : HashMap<i32, Vec<i32>>,
    pub occurrences : HashMap<i32, Vec<i32>>,
    pub units       : HashSet<i32>
}

impl CNF {
    pub fn new(nvar: i32, nclause: i32) -> CNF {
        CNF {
            nvar        : nvar,
            nclause     : nclause,
            clauses     : HashMap::new(),
            occurrences : HashMap::new(),
            units       : HashSet::new(),
        }
    }

    /// Add a clause, return the ID of the inserted clause
    /// Right now, this isn't public; api is odd as we have an odd representation of literals.
    /// TODO: Mask this with public function?
    fn add_clause(&mut self, clause: Vec<i32>) -> i32 {
        assert!(clause.len() > 0);
        let id: i32 = self.clauses.len() as i32;
        if clause.len() == 1 { self.units.insert(id); }
        for var in &clause {
            let occ = self.occurrences.entry(*var).or_insert(Vec::new());
            occ.push(id);
        }
        self.clauses.insert(id, clause);
        id
    }

    /// Unit propagation: propagate the literal associated with the clause id `clause`, and remove
    /// it from `self.units`.
    pub fn unit_propagate(&mut self, unit: i32) {
        // Remove clauses with lit and remove lit from occurrences.
        let lit = self.clauses.get(&unit).expect("unit clause not found")[0];
        self.propagate(lit);
        // This needs to be true, i.e. unit clause id needs to be in the unit clauses
        assert!(self.units.remove(&unit));
    }

    /// Remove all clauses containing literal `lit` from the CNF, and remove the negation of `lit`
    /// from the remaining clauses.
    /// TODO: unit_propagate and propagate should return true/false
    /// depending on whether the updated CNF problem is satisfiable
    pub fn propagate(&mut self, lit: i32) {
        // Remove clauses with lit and remove lit from occurrences.
        for occ in &self.occurrences.remove(&lit).unwrap() {
            self.clauses.remove(occ);
        }
    }
}

impl Clone for CNF {
    fn clone(&self) -> CNF {
        CNF {
            nvar        : self.nvar,
            nclause     : self.nclause,
            clauses     : self.clauses.clone(),
            occurrences : self.occurrences.clone(),
            units       : self.units.clone()
        }
    }
}

impl Display for CNF {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        let formatted = format!("Nvar: {:?} Nclause: {:?}\nClauses: {:?}\nOccurrrences: {:?}",
                self.nvar, self.nclause,
                self.clauses, self.occurrences);
        write!(f, "{}", formatted)
    }
}

// End CNF

// Begin Assignemnts

pub type Assignment = Vec<bool>;

pub struct PartialAssignment {
    pub assignment : Vec<Option<bool>>,
    pub unassigned : HashSet<i32>
}

impl PartialAssignment {
    pub fn new(n: usize) -> PartialAssignment {
        PartialAssignment {
            assignment : vec!(None; n),
            unassigned : (0..n as i32).collect()
        }
    }

    fn assign(&mut self, v: usize, assn: bool) {
        // TODO: Error check
        self.assignment[v] = Some(assn);
    }

    pub fn assign_literal(&mut self, lit: i32){
        let polarity: bool = lit & 1 == 0;
        let v: usize = (lit >> 1) as usize;
        self.assign(v - 1, polarity);
    }

    fn unassign(&mut self, v: usize) {
        self.assignment[v] = None;
    }

    pub fn unassign_literal(&mut self, lit: i32){
        let v: usize = (lit >> 1) as usize;
        self.unassign(v - 1);
    }
}

// TODO: Implement Display
impl Display for PartialAssignment {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        let formatted = format!("Assignemnts: {:?}\nUnassigned: {:?}",
                self.assignment, self.unassigned);
        write!(f, "{}", formatted)
    }
}

// End Assignments

// Begin Parsing

// TODO: rename to from_dimacs or somehow isolate parsing
fn parse_dimacs(reader: &mut BufReader<File>) -> Result<CNF, &'static str> {
    let mut line_iterator = reader.lines();

    let mut nvar: i32 = -1;
    let mut nclause: i32 = -1;
    // DIMACS file must have a problem statement before other lines.
    // This first loop searches for the problem statement.
    for line in &mut line_iterator {
        let line = line.expect("could not read file");  // Unwrap result
        let words: Vec<&str> = line.split_whitespace().collect();
        if words.len() == 0 { continue; } // Ignore empty lines
        match words[0] {
            "c" => { }
            "p" => { // Problem statement
                // Must have format "p cnf nvar nclause"
                if words.len() != 4 || words[1] != "cnf" {
                    return Err("invalid problem statement");
                }
                nvar = words[2].parse()
                    .expect(&format!("invalid number of variables {}", words[2]));
                nclause = words[3].parse()
                    .expect(&format!("invalid number of clauses {}", words[3]));
                break;
            }
            // TODO: Add words[0] to this error message
            _ => { return Err("unknown statement beginning"); }
        }
    }
    // Then nvar, nclause were never initialized
    if nvar == -1 || nclause == -1 {
       return Err("no problem statement found");
    }
    // TODO: Different errors for different descriptions (hence why I've split up this if
    // statement)
    if nvar == 0 || nclause == 0 {
        return Err("invalid number of literals in problem")
    }

    // Initialize CNF and parse the rest of the file
    let mut cnf = CNF::new(nvar, nclause);
    let mut clauses_read: i32 = 0;
    for line in &mut line_iterator {
        let line = line.unwrap();
        let words: Vec<&str> = line.split_whitespace().collect();
        if words.len() == 0 { continue; }
        match words[0] {
            "c" => { }
            "p" => { return Err("duplicate problem statement"); }
            _   => {
                clauses_read = clauses_read + 1;
                if clauses_read > nclause { return Err("too many clauses in file"); }
                let tokens: Vec<i32> = words.iter()
                    .filter_map(|s| {
                        let n = s.parse::<i32>().unwrap();
                        // FIXME: This ignores zeros not just as line enders but in the formulas
                        // themselves. Split on zeros at the end here.
                        if n == 0 { return None; }
                        // FIXME: This should return an error instead of
                        // panicking, but I can't return an error
                        // in a closure
                        if n > nvar { panic!("variable out of range: {}", n); }
                        Some(if n < 0 { (-n) << 1 | 1 } else { n << 1 })
                    })
                    .collect();
                cnf.add_clause(tokens);
            }
        }
    }
    // Double check that the number of clauses read is equal
    if clauses_read != nclause { return Err("too few clauses in file"); }
    Ok(cnf)
}

pub fn parse_dimacs_file(f: File) -> Result<CNF, &'static str> {
    // Read the file
    let mut reader = BufReader::new(f);

    // TODO: This is definitely not the correct way to handle errors.
    // Should parse_dimacs have options for returning an error AND panicking?
    let line = parse_dimacs(&mut reader);
    line
}

#[cfg(test)]
mod tests {
    extern crate tempfile;
    use std::fs::File;
    use std::io::prelude::*;
    use std::io::SeekFrom;
    use std::collections::HashSet;

    use super::parse_dimacs_file;
    // use super::zeroth;

    macro_rules! create_tempfile {
        ($x: expr) => {{
            let mut tmpfile: File = tempfile::tempfile().unwrap();
            let _ = write!(tmpfile, $x);
            tmpfile.seek(SeekFrom::Start(0)).unwrap();

            tmpfile
        }};
    }

    #[test]
    #[should_panic(expected = "out of range")]
    fn variable_out_of_range() {
        let tmpfile = create_tempfile!("
            p cnf 2 3
            1 2 0
            4 1 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    #[should_panic(expected = "invalid number")]
    fn invalid_nvar() {
        let tmpfile = create_tempfile!("
            p cnf gd 3
            1 2 0
            4 1 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    #[should_panic(expected = "invalid number")]
    fn invalid_nclause() {
        let tmpfile = create_tempfile!("
            p cnf 2 gdd
            1 2 0
            4 1 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    #[should_panic(expected = "too many")]
    fn too_many_clauses() {
        let tmpfile = create_tempfile!("
            p cnf 5 5
            1 2 0
            2 3 0
            3 4 0
            4 5 0
            1 3 0
            2 4 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    #[should_panic(expected = "too few")]
    fn too_little_clauses() {
        let tmpfile = create_tempfile!("
            p cnf 5 5
            1 2 0
            2 3 0
            3 4 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    fn should_work() {
        let tmpfile = create_tempfile!("
            p cnf 3 3
            1 2 0
            2 3 0
            1 -3 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    fn should_ignore_comments() {
        let tmpfile = create_tempfile!("
            c comment
            p cnf 3 3
            c comment
            1 2 0
            2 3 0
            c comment
            1 -3 0
            c comment
            c comment
            c comment
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    #[should_panic(expected = "duplicate problem")]
    fn duplicate_problem_statement() {
        let tmpfile = create_tempfile!("
            p cnf 3 3
            1 2 0
            p cnf 3 3
            2 3 0
            1 -3 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    #[should_panic(expected = "duplicate problem")]
    fn duplicate_problem_statement_2() {
        let tmpfile = create_tempfile!("
            p cnf 3 3
            p cnf 3 3
            1 2 0
            2 3 0
            1 -3 0
        ");
        let _ = parse_dimacs_file(tmpfile).unwrap();
    }

    #[test]
    fn adds_unit_clauses() {
        let tmpfile = create_tempfile!("
            p cnf 3 4
            1 0
            2 0
            3 0
            -1 -3 0
        ");
        let cnf = parse_dimacs_file(tmpfile).unwrap();
        assert_eq!(cnf.units.len(), 3);
    }

    // TODO: Make this doc test (with unit propagate); create_tempfile may need to be an importable
    // macro up there
    #[test]
    fn unit_propagate_works() {
        let tmpfile = create_tempfile!("
            p cnf 2 2
            1 0
            1 2 0
        ");
        let mut cnf = parse_dimacs_file(tmpfile).unwrap();
        cnf.unit_propagate(0);
        assert!(cnf.units.is_empty());
    }

    #[test]
    fn unit_propagate_works_2() {
        let tmpfile = create_tempfile!("
            p cnf 4 4
            1 0
            1 2 0
            2 3 0
            4 0
        ");
        let mut cnf = parse_dimacs_file(tmpfile).unwrap();
        cnf.unit_propagate(0);
        // Only one unit clause left
        assert_eq!(cnf.units.len(), 1);
        cnf.unit_propagate(3);
        // No more unit clauses
        assert!(cnf.units.is_empty());
    }

    #[test]
    #[should_panic(expected = "unit clause not found")]
    fn unit_clause_not_found() {
        let tmpfile = create_tempfile!("
            p cnf 2 2
            1 0
            1 2 0
        ");
        let mut cnf = parse_dimacs_file(tmpfile).unwrap();
        cnf.unit_propagate(2);
    }

    #[test]
    fn zeroth_works() {
        let mut hs = HashSet::new();
        hs.insert(5);
        assert_eq!(zeroth!(hs), 5);
    }

    #[test]
    fn zeroth_works_2() {
        let mut hs = HashSet::new();
        hs.insert(5);
        hs.insert(535);
        let zth = zeroth!(hs);
        assert!(zth == 5 || zth == 535);
    }
}
<|endoftext|>


<commit_before>// Extracted from the scopeguard crate
use core::ops::{Deref, DerefMut};

pub struct ScopeGuard<T, F>
where
    F: FnMut(&mut T),
{
    dropfn: F,
    value: T,
}

#[cfg_attr(feature = "inline-more", inline)]
pub fn guard<T, F>(value: T, dropfn: F) -> ScopeGuard<T, F>
where
    F: FnMut(&mut T),
{
    ScopeGuard { dropfn, value }
}

impl<T, F> Deref for ScopeGuard<T, F>
where
    F: FnMut(&mut T),
{
    type Target = T;
    #[cfg_attr(feature = "inline-more", inline)]
    fn deref(&self) -> &T {
        &self.value
    }
}

impl<T, F> DerefMut for ScopeGuard<T, F>
where
    F: FnMut(&mut T),
{
    #[cfg_attr(feature = "inline-more", inline)]
    fn deref_mut(&mut self) -> &mut T {
        &mut self.value
    }
}

impl<T, F> Drop for ScopeGuard<T, F>
where
    F: FnMut(&mut T),
{
    #[cfg_attr(feature = "inline-more", inline)]
    fn drop(&mut self) {
        (self.dropfn)(&mut self.value)
    }
}
<commit_msg>Update scopeguard.rs<commit_after>// Extracted from the scopeguard crate
use core::ops::{Deref, DerefMut};

pub struct ScopeGuard<T, F>
where
    F: FnMut(&mut T),
{
    dropfn: F,
    value: T,
}

#[cfg_attr(feature = "inline-more", inline)]
pub fn guard<T, F>(value: T, dropfn: F) -> ScopeGuard<T, F>
where
    F: FnMut(&mut T),
{
    ScopeGuard { dropfn, value }
}

impl<T, F> Deref for ScopeGuard<T, F>
where
    F: FnMut(&mut T),
{
    type Target = T;
    #[cfg_attr(feature = "inline-more", inline)]
    fn deref(&self) -> &T {
        &self.value
    }
}

impl<T, F> DerefMut for ScopeGuard<T, F>
where
    F: FnMut(&mut T),
{
    #[cfg_attr(feature = "inline-more", inline)]
    fn deref_mut(&mut self) -> &mut T {
        &mut self.value
    }
}

impl<T, F> Drop for ScopeGuard<T, F>
where
    F: FnMut(&mut T),
{
    #[inline]
    fn drop(&mut self) {
        (self.dropfn)(&mut self.value)
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::*;
use std::mem;
use std::slice;
use std::thread;

use super::FONT;
use super::Event;
use super::Color;

/// A window
pub struct Window {
    /// The x coordinate of the window
    x: i32,
    /// The y coordinate of the window
    y: i32,
    /// The width of the window
    w: u32,
    /// The height of the window
    h: u32,
    /// The title of the window
    t: String,
    /// The input scheme
    file: File,
    /// Window data
    data: Box<[Color]>,
}

impl Window {
    /// Create a new window
    pub fn new(x: i32, y: i32, w: u32, h: u32, title: &str) -> Option<Box<Self>> {
        match File::open(&format!("orbital:/{}/{}/{}/{}/{}", x, y, w, h, title)) {
            Ok(file) => {
                Some(box Window {
                    x: x,
                    y: y,
                    w: w,
                    h: h,
                    t: title.to_string(),
                    file: file,
                    data: vec![Color::rgb(0, 0, 0); (w * h * 4) as usize].into_boxed_slice(),
                })
            }
            Err(_) => None,
        }
    }

    // TODO: Replace with smarter mechanism, maybe a move event?
    pub fn sync_path(&mut self) {
        if let Ok(path) = self.file.path() {
            // orbital:/x/y/w/h/t
            if let Some(path_str) = path.to_str() {
                let mut parts = path_str.split('/').skip(1);
                if let Some(x) = parts.next() {
                    self.x = x.parse::<i32>().unwrap_or(0);
                }
                if let Some(y) = parts.next() {
                    self.y = y.parse::<i32>().unwrap_or(0);
                }
                if let Some(w) = parts.next() {
                    self.w = w.parse::<u32>().unwrap_or(0);
                }
                if let Some(h) = parts.next() {
                    self.h = h.parse::<u32>().unwrap_or(0);
                }
            }
        }
    }

    /// Get x
    // TODO: Sync with window movements
    pub fn x(&self) -> i32 {
        self.x
    }

    /// Get y
    // TODO: Sync with window movements
    pub fn y(&self) -> i32 {
        self.y
    }

    /// Get width
    pub fn width(&self) -> u32 {
        self.w
    }

    /// Get height
    pub fn height(&self) -> u32 {
        self.h
    }

    /// Get title
    pub fn title(&self) -> String {
        self.t.clone()
    }

    /// Set title
    pub fn set_title(&mut self, _: &str) {
        // TODO
    }

    /// Draw a pixel
    pub fn pixel(&mut self, x: i32, y: i32, color: Color) {
        if x >= 0 && y >= 0 && x < self.w as i32 && y < self.h as i32 {
            let new = color.data;

            let alpha = (new >> 24) & 0xFF;
            if alpha > 0 {
                let old = &mut self.data[y as usize * self.w as usize + x as usize].data;
                if alpha >= 255 {
                    *old = new;
                } else {
                    let n_r = (((new >> 16) & 0xFF) * alpha) >> 8;
                    let n_g = (((new >> 8) & 0xFF) * alpha) >> 8;
                    let n_b = ((new & 0xFF) * alpha) >> 8;

                    let n_alpha = 255 - alpha;
                    let o_a = (((*old >> 24) & 0xFF) * n_alpha) >> 8;
                    let o_r = (((*old >> 16) & 0xFF) * n_alpha) >> 8;
                    let o_g = (((*old >> 8) & 0xFF) * n_alpha) >> 8;
                    let o_b = ((*old & 0xFF) * n_alpha) >> 8;

                    *old = ((o_a << 24) | (o_r << 16) | (o_g << 8) | o_b) + ((alpha << 24) | (n_r << 16) | (n_g << 8) | n_b);
                }
            }
        }
    }

    /// Draw a character, using the loaded font
    pub fn char(&mut self, x: i32, y: i32, c: char, color: Color) {
        let mut offset = (c as usize) * 16;
        for row in 0..16 {
            let row_data;
            if offset < FONT.len() {
                row_data = FONT[offset];
            } else {
                row_data = 0;
            }

            for col in 0..8 {
                let pixel = (row_data >> (7 - col)) & 1;
                if pixel > 0 {
                    self.pixel(x + col as i32, y + row as i32, color);
                }
            }
            offset += 1;
        }
    }

    // TODO move, resize, set_title

    /// Set entire window to a color
    // TODO: Improve speed
    #[allow(unused_variables)]
    pub fn set(&mut self, color: Color) {
        for mut d in self.data.iter_mut() {
            *d = color;
        }
    }

    /// Draw rectangle
    // TODO: Improve speed
    #[allow(unused_variables)]
    pub fn rect(&mut self, start_x: i32, start_y: i32, w: u32, h: u32, color: Color) {
        for y in start_y..start_y + h as i32 {
            for x in start_x..start_x + w as i32 {
                self.pixel(x, y, color);
            }
        }
    }

    /// Display an image
    // TODO: Improve speed
    pub fn image(&mut self, start_x: i32, start_y: i32, w: u32, h: u32, data: &[Color]) {
        let mut i = 0;
        for y in start_y..start_y + h as i32 {
            for x in start_x..start_x + w as i32 {
                if i < data.len() {
                    self.pixel(x, y, data[i])
                }
                i += 1;
            }
        }
    }

    /// Return a iterator over events
    fn events_inner(&mut self, wait: bool) -> EventIter {
        let mut iter = EventIter {
            events: [Event::new(); 128],
            i: 0,
            count: 0,
        };

        'blocking: loop {
            //Should it be cleared? iter.events = [Event::new(); 128];
            match self.file.read(unsafe {
                slice::from_raw_parts_mut(iter.events.as_mut_ptr() as *mut u8, iter.events.len() * mem::size_of::<Event>())
            }){
                Ok(0) => if wait {
                    thread::yield_now();
                } else {
                    break 'blocking;
                },
                Ok(count) => {
                    iter.count = count/mem::size_of::<Event>();
                    break 'blocking;
                },
                Err(_) => break 'blocking,
            }
        }

        iter
    }

    /// Blocking iterator over events
    pub fn events(&mut self) -> EventIter {
        self.events_inner(true)
    }

    /// Nonblocking iterator over events
    pub fn events_no_wait(&mut self) -> EventIter {
        self.events_inner(false)
    }

    /// Flip the window buffer
    pub fn sync(&mut self) -> bool {
        self.file.write(unsafe {
            slice::from_raw_parts(self.data.as_ptr() as *const u8, self.data.len() * mem::size_of::<Color>())
        }).is_ok()
    }
}

/// Event iterator
pub struct EventIter {
    events: [Event; 128],
    i: usize,
    count: usize,
}

impl Iterator for EventIter {
    type Item = Event;
    fn next(&mut self) -> Option<Event> {
        if self.i < self.count {
            if let Some(event) = self.events.get(self.i) {
                self.i += 1;
                Some(*event)
            } else {
                None
            }
        } else {
            None
        }
    }
}
<commit_msg>Add functions to create left, right, top and bottom borders<commit_after>use std::fs::File;
use std::io::*;
use std::mem;
use std::slice;
use std::thread;

use super::FONT;
use super::Event;
use super::Color;

/// A window
pub struct Window {
    /// The x coordinate of the window
    x: i32,
    /// The y coordinate of the window
    y: i32,
    /// The width of the window
    w: u32,
    /// The height of the window
    h: u32,
    /// The title of the window
    t: String,
    /// The input scheme
    file: File,
    /// Window data
    data: Box<[Color]>,
}

impl Window {
    /// Create a new window
    pub fn new(x: i32, y: i32, w: u32, h: u32, title: &str) -> Option<Box<Self>> {
        match File::open(&format!("orbital:/{}/{}/{}/{}/{}", x, y, w, h, title)) {
            Ok(file) => {
                Some(box Window {
                    x: x,
                    y: y,
                    w: w,
                    h: h,
                    t: title.to_string(),
                    file: file,
                    data: vec![Color::rgb(0, 0, 0); (w * h * 4) as usize].into_boxed_slice(),
                })
            }
            Err(_) => None,
        }
    }

    // TODO: Replace with smarter mechanism, maybe a move event?
    pub fn sync_path(&mut self) {
        if let Ok(path) = self.file.path() {
            // orbital:/x/y/w/h/t
            if let Some(path_str) = path.to_str() {
                let mut parts = path_str.split('/').skip(1);
                if let Some(x) = parts.next() {
                    self.x = x.parse::<i32>().unwrap_or(0);
                }
                if let Some(y) = parts.next() {
                    self.y = y.parse::<i32>().unwrap_or(0);
                }
                if let Some(w) = parts.next() {
                    self.w = w.parse::<u32>().unwrap_or(0);
                }
                if let Some(h) = parts.next() {
                    self.h = h.parse::<u32>().unwrap_or(0);
                }
            }
        }
    }

    /// Get x
    // TODO: Sync with window movements
    pub fn x(&self) -> i32 {
        self.x
    }

    /// Get y
    // TODO: Sync with window movements
    pub fn y(&self) -> i32 {
        self.y
    }

    /// Get width
    pub fn width(&self) -> u32 {
        self.w
    }

    /// Get height
    pub fn height(&self) -> u32 {
        self.h
    }

    /// Get title
    pub fn title(&self) -> String {
        self.t.clone()
    }

    /// Set title
    pub fn set_title(&mut self, _: &str) {
        // TODO
    }

    /// Draw a pixel
    pub fn pixel(&mut self, x: i32, y: i32, color: Color) {
        if x >= 0 && y >= 0 && x < self.w as i32 && y < self.h as i32 {
            let new = color.data;

            let alpha = (new >> 24) & 0xFF;
            if alpha > 0 {
                let old = &mut self.data[y as usize * self.w as usize + x as usize].data;
                if alpha >= 255 {
                    *old = new;
                } else {
                    let n_r = (((new >> 16) & 0xFF) * alpha) >> 8;
                    let n_g = (((new >> 8) & 0xFF) * alpha) >> 8;
                    let n_b = ((new & 0xFF) * alpha) >> 8;

                    let n_alpha = 255 - alpha;
                    let o_a = (((*old >> 24) & 0xFF) * n_alpha) >> 8;
                    let o_r = (((*old >> 16) & 0xFF) * n_alpha) >> 8;
                    let o_g = (((*old >> 8) & 0xFF) * n_alpha) >> 8;
                    let o_b = ((*old & 0xFF) * n_alpha) >> 8;

                    *old = ((o_a << 24) | (o_r << 16) | (o_g << 8) | o_b) + ((alpha << 24) | (n_r << 16) | (n_g << 8) | n_b);
                }
            }
        }
    }

    /// Draw a character, using the loaded font
    pub fn char(&mut self, x: i32, y: i32, c: char, color: Color) {
        let mut offset = (c as usize) * 16;
        for row in 0..16 {
            let row_data;
            if offset < FONT.len() {
                row_data = FONT[offset];
            } else {
                row_data = 0;
            }

            for col in 0..8 {
                let pixel = (row_data >> (7 - col)) & 1;
                if pixel > 0 {
                    self.pixel(x + col as i32, y + row as i32, color);
                }
            }
            offset += 1;
        }
    }

    // TODO move, resize, set_title

    /// Set entire window to a color
    // TODO: Improve speed
    #[allow(unused_variables)]
    pub fn set(&mut self, color: Color) {
        for mut d in self.data.iter_mut() {
            *d = color;
        }
    }

    pub fn set_border_left(&mut self, x_limit: i32, color: Color, density: i32) {
        for y in 0..self.height {
            for x in x_limit..(x_limit + density) {
                self.pixel(x, y as i32, color);
            }
        }
    }

    pub fn set_border_right(&mut self, x_limit: i32, color: Color, density: i32) {
        for y in 0..self.height {
            for x in x_limit..(x_limit + density) {
                let new_x : i32 = self.width as i32;
                self.pixel((new_x - x), y as i32, color);
            }
        }
    }

    pub fn set_border_top(&mut self, y_limit: i32, color: Color, density: u32) {
        for x in 0..self.width {
            for y in y_limit..(y_limit + density) {
                self.pixel(x, y, color);
            }
        }
    }

    pub fn set_border_bottom(&mut self, color: Color, density: u32) {
        for x in 0..self.x {
            for y in 0..(y_limit + density) {
                self.pixel(x, self.y - y, color);
            }
        }
    }

    /// Draw rectangle
    // TODO: Improve speed
    #[allow(unused_variables)]
    pub fn rect(&mut self, start_x: i32, start_y: i32, w: u32, h: u32, color: Color) {
        for y in start_y..start_y + h as i32 {
            for x in start_x..start_x + w as i32 {
                self.pixel(x, y, color);
            }
        }
    }

    /// Display an image
    // TODO: Improve speed
    pub fn image(&mut self, start_x: i32, start_y: i32, w: u32, h: u32, data: &[Color]) {
        let mut i = 0;
        for y in start_y..start_y + h as i32 {
            for x in start_x..start_x + w as i32 {
                if i < data.len() {
                    self.pixel(x, y, data[i])
                }
                i += 1;
            }
        }
    }

    /// Return a iterator over events
    fn events_inner(&mut self, wait: bool) -> EventIter {
        let mut iter = EventIter {
            events: [Event::new(); 128],
            i: 0,
            count: 0,
        };

        'blocking: loop {
            //Should it be cleared? iter.events = [Event::new(); 128];
            match self.file.read(unsafe {
                slice::from_raw_parts_mut(iter.events.as_mut_ptr() as *mut u8, iter.events.len() * mem::size_of::<Event>())
            }){
                Ok(0) => if wait {
                    thread::yield_now();
                } else {
                    break 'blocking;
                },
                Ok(count) => {
                    iter.count = count/mem::size_of::<Event>();
                    break 'blocking;
                },
                Err(_) => break 'blocking,
            }
        }

        iter
    }

    /// Blocking iterator over events
    pub fn events(&mut self) -> EventIter {
        self.events_inner(true)
    }

    /// Nonblocking iterator over events
    pub fn events_no_wait(&mut self) -> EventIter {
        self.events_inner(false)
    }

    /// Flip the window buffer
    pub fn sync(&mut self) -> bool {
        self.file.write(unsafe {
            slice::from_raw_parts(self.data.as_ptr() as *const u8, self.data.len() * mem::size_of::<Color>())
        }).is_ok()
    }
}

/// Event iterator
pub struct EventIter {
    events: [Event; 128],
    i: usize,
    count: usize,
}

impl Iterator for EventIter {
    type Item = Event;
    fn next(&mut self) -> Option<Event> {
        if self.i < self.count {
            if let Some(event) = self.events.get(self.i) {
                self.i += 1;
                Some(*event)
            } else {
                None
            }
        } else {
            None
        }
    }
}
<|endoftext|>


<commit_before>// TODO rm allow(unused)
#![allow(unused)]
#![allow(unsafe_code)]

use std::{
    convert::{TryFrom, TryInto},
    fmt::{self, Debug},
    fs::File,
    mem::{transmute, MaybeUninit},
    path::Path,
    sync::{
        atomic::{AtomicU32, Ordering::Acquire},
        Arc,
    },
};

use crossbeam_epoch::pin;

use crate::{
    pagecache::{pread_exact, pwrite_all, MessageKind},
    stack::Stack,
    Error, Lsn, Result,
};

#[cfg(not(feature = "testing"))]
const MIN_SZ: u64 = 64 * 1024;

#[cfg(feature = "testing")]
const MIN_SZ: u64 = 32;

const MIN_TRAILING_ZEROS: u64 = MIN_SZ.trailing_zeros() as u64;

pub type SlabId = u8;
pub type SlabIdx = u32;

/// A unique identifier for a particular slot in the heap
#[allow(clippy::module_name_repetitions)]
#[derive(Clone, Copy, PartialOrd, Ord, Eq, PartialEq, Hash)]
pub struct HeapId {
    pub location: u64,
    pub original_lsn: Lsn,
}

impl Debug for HeapId {
    fn fmt(
        &self,
        f: &mut fmt::Formatter<'_>,
    ) -> std::result::Result<(), fmt::Error> {
        let (slab, idx, original_lsn) = self.decompose();
        f.debug_struct("HeapId")
            .field("slab", &slab)
            .field("idx", &idx)
            .field("original_lsn", &original_lsn)
            .finish()
    }
}

impl HeapId {
    pub fn decompose(&self) -> (SlabId, SlabIdx, Lsn) {
        const IDX_MASK: u64 = (1 << 32) - 1;
        let slab_id =
            u8::try_from((self.location >> 32).trailing_zeros()).unwrap();
        let slab_idx = u32::try_from(self.location & IDX_MASK).unwrap();
        (slab_id, slab_idx, self.original_lsn)
    }

    pub fn compose(
        slab_id: SlabId,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
    ) -> HeapId {
        let slab = 1 << (32 + u64::from(slab_id));
        let heap_id = slab | u64::from(slab_idx);
        HeapId { location: heap_id, original_lsn }
    }

    fn offset(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id) * u64::from(idx)
    }

    fn slab_size(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id)
    }
}

pub(crate) fn slab_size(size: u64) -> u64 {
    slab_id_to_size(size_to_slab_id(size))
}

fn slab_id_to_size(slab_id: u8) -> u64 {
    1 << (MIN_TRAILING_ZEROS + u64::from(slab_id))
}

fn size_to_slab_id(size: u64) -> SlabId {
    // find the power of 2 that is at least 64k
    let normalized_size = std::cmp::max(MIN_SZ, size.next_power_of_two());

    // drop the lowest unused bits
    let rebased_size = normalized_size >> MIN_TRAILING_ZEROS;

    u8::try_from(rebased_size.trailing_zeros()).unwrap()
}

pub(crate) struct Reservation {
    slab_free: Arc<Stack<u32>>,
    completed: bool,
    file: File,
    pub heap_id: HeapId,
}

impl Drop for Reservation {
    fn drop(&mut self) {
        if !self.completed {
            let (slab_id, idx, _) = self.heap_id.decompose();
            self.slab_free.push(idx, &pin());
        }
    }
}

impl Reservation {
    pub fn complete(mut self, data: &[u8]) -> Result<HeapId> {
        log::trace!(
            "Heap::complete({:?}) to offset {} in file {:?}",
            self.heap_id,
            self.heap_id.offset(),
            self.file
        );
        assert_eq!(data.len() as u64, self.heap_id.slab_size());

        pwrite_all(&self.file, data, self.heap_id.offset())?;
        self.file.sync_all()?;

        // if this is not reached due to an IO error,
        // the offset will be returned to the Slab in Drop
        self.completed = true;

        Ok(self.heap_id)
    }

    pub fn abort(self) {
        // actual logic in Drop
    }
}

#[derive(Debug)]
pub(crate) struct Heap {
    // each slab stores
    // items that are double
    // the size of the previous,
    // ranging from 64k in the
    // smallest slab to 2^48 in
    // the last.
    slabs: [Slab; 32],
}

impl Heap {
    pub fn start<P: AsRef<Path>>(p: P) -> Result<Heap> {
        let mut slabs: [MaybeUninit<Slab>; 32] = unsafe { std::mem::zeroed() };

        for slab_id in 0..32 {
            let slab = Slab::start(&p, slab_id)?;
            slabs[slab_id as usize] = MaybeUninit::new(slab);
        }

        Ok(Heap { slabs: unsafe { transmute(slabs) } })
    }

    pub fn gc_unknown_items(
        &self,
        _snapshot: &crate::pagecache::Snapshot,
    ) -> Result<()> {
        //TODO todo!()
        Ok(())
    }

    pub fn read(
        &self,
        heap_id: HeapId,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        log::trace!("Heap::read({:?})", heap_id);
        let (slab_id, slab_idx, original_lsn) = heap_id.decompose();
        self.slabs[slab_id as usize].read(
            slab_idx,
            original_lsn,
            use_compression,
        )
    }

    pub fn free(&self, heap_id: HeapId) {
        log::trace!("Heap::free({:?})", heap_id);
        let (slab_id, slab_idx, _) = heap_id.decompose();
        self.slabs[slab_id as usize].free(slab_idx)
    }

    pub fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        assert!(size < 1 << 48);
        let slab_id = size_to_slab_id(size);
        let ret = self.slabs[slab_id as usize].reserve(size, original_lsn);
        log::trace!("Heap::reserve({}) -> {:?}", size, ret.heap_id);
        ret
    }
}

#[derive(Debug)]
struct Slab {
    file: File,
    slab_id: u8,
    tip: AtomicU32,
    free: Arc<Stack<u32>>,
}

impl Slab {
    pub fn start<P: AsRef<Path>>(directory: P, slab_id: u8) -> Result<Slab> {
        let bs = slab_id_to_size(slab_id);
        let free = Arc::new(Stack::default());

        let mut options = std::fs::OpenOptions::new();
        options.create(true);
        options.read(true);
        options.write(true);

        let file =
            options.open(directory.as_ref().join(format!("{:02}", slab_id)))?;
        let len = file.metadata()?.len();
        let max_idx = len / bs;
        log::trace!(
            "starting heap slab for sizes of {}. tip: {} max idx: {}",
            bs,
            len,
            max_idx
        );
        let tip = AtomicU32::new(u32::try_from(max_idx).unwrap());

        Ok(Slab { file, slab_id, tip, free })
    }

    fn read(
        &self,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        let bs = slab_id_to_size(self.slab_id);
        let offset = u64::from(slab_idx) * bs;

        log::trace!("reading heap slab slot {} at offset {}", slab_idx, offset);

        let mut heap_buf = vec![0; usize::try_from(bs).unwrap()];

        pread_exact(&self.file, &mut heap_buf, offset)?;

        let stored_crc =
            u32::from_le_bytes(heap_buf[1..5].as_ref().try_into().unwrap());

        let mut hasher = crc32fast::Hasher::new();
        hasher.update(&heap_buf[0..1]);
        hasher.update(&heap_buf[5..]);
        let actual_crc = hasher.finalize();

        if actual_crc == stored_crc {
            let actual_lsn = Lsn::from_le_bytes(
                heap_buf[5..13].as_ref().try_into().unwrap(),
            );
            if actual_lsn != original_lsn {
                log::error!(
                    "heap slot lsn {} does not match expected original lsn {}",
                    actual_lsn,
                    original_lsn
                );
                return Err(Error::corruption(None));
            }
            let buf = heap_buf[13..].to_vec();
            let buf = if use_compression {
                crate::pagecache::decompress(buf)
            } else {
                buf
            };
            Ok((MessageKind::from(heap_buf[0]), buf))
        } else {
            log::error!(
                "heap message CRC does not match contents. stored: {} actual: {}",
                stored_crc,
                actual_crc
            );
            Err(Error::corruption(None))
        }
    }

    fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        let idx = if let Some(idx) = self.free.pop(&pin()) {
            log::trace!(
                "reusing heap index {} in slab for sizes of {}",
                idx,
                slab_id_to_size(self.slab_id),
            );
            idx
        } else {
            log::trace!(
                "no free heap slots in slab for sizes of {}",
                slab_id_to_size(self.slab_id),
            );
            self.tip.fetch_add(1, Acquire)
        };

        log::trace!(
            "heap reservation for slot {} in the slab for sizes of {}",
            idx,
            slab_id_to_size(self.slab_id),
        );

        let heap_id = HeapId::compose(self.slab_id, idx, original_lsn);

        Reservation {
            slab_free: self.free.clone(),
            completed: false,
            file: self.file.try_clone().unwrap(),
            heap_id,
        }
    }

    fn free(&self, idx: u32) {
        self.punch_hole(idx);
        self.free.push(idx, &pin());
    }

    fn punch_hole(&self, idx: u32) {
        #[cfg(target_os = "linux")]
        {
            use std::{
                os::unix::io::AsRawFd,
                sync::atomic::{AtomicBool, Ordering::Relaxed},
            };

            use libc::{fallocate, FALLOC_FL_KEEP_SIZE, FALLOC_FL_PUNCH_HOLE};

            static HOLE_PUNCHING_ENABLED: AtomicBool = AtomicBool::new(true);
            const MODE: i32 = FALLOC_FL_KEEP_SIZE | FALLOC_FL_PUNCH_HOLE;

            if HOLE_PUNCHING_ENABLED.load(Relaxed) {
                let bs = i64::try_from(slab_id_to_size(self.slab_id)).unwrap();
                let offset = i64::from(idx) * bs;

                let fd = self.file.as_raw_fd();

                let ret = unsafe {
                    fallocate(
                        fd,
                        MODE,
                        offset.try_into().unwrap(),
                        bs.try_into().unwrap(),
                    )
                };

                if ret != 0 {
                    let err = std::io::Error::last_os_error();
                    log::error!(
                        "failed to punch hole in heap file: {:?}. disabling hole punching",
                        err
                    );
                    HOLE_PUNCHING_ENABLED.store(false, Relaxed);
                }
            }
        }
    }
}
<commit_msg>Use sync_file_range instead of fsync for heap writes on linux<commit_after>// TODO rm allow(unused)
#![allow(unused)]
#![allow(unsafe_code)]

use std::{
    convert::{TryFrom, TryInto},
    fmt::{self, Debug},
    fs::File,
    mem::{transmute, MaybeUninit},
    path::Path,
    sync::{
        atomic::{AtomicU32, Ordering::Acquire},
        Arc,
    },
};

use crossbeam_epoch::pin;

use crate::{
    pagecache::{pread_exact, pwrite_all, MessageKind},
    stack::Stack,
    Error, Lsn, Result,
};

#[cfg(not(feature = "testing"))]
const MIN_SZ: u64 = 64 * 1024;

#[cfg(feature = "testing")]
const MIN_SZ: u64 = 32;

const MIN_TRAILING_ZEROS: u64 = MIN_SZ.trailing_zeros() as u64;

pub type SlabId = u8;
pub type SlabIdx = u32;

/// A unique identifier for a particular slot in the heap
#[allow(clippy::module_name_repetitions)]
#[derive(Clone, Copy, PartialOrd, Ord, Eq, PartialEq, Hash)]
pub struct HeapId {
    pub location: u64,
    pub original_lsn: Lsn,
}

impl Debug for HeapId {
    fn fmt(
        &self,
        f: &mut fmt::Formatter<'_>,
    ) -> std::result::Result<(), fmt::Error> {
        let (slab, idx, original_lsn) = self.decompose();
        f.debug_struct("HeapId")
            .field("slab", &slab)
            .field("idx", &idx)
            .field("original_lsn", &original_lsn)
            .finish()
    }
}

impl HeapId {
    pub fn decompose(&self) -> (SlabId, SlabIdx, Lsn) {
        const IDX_MASK: u64 = (1 << 32) - 1;
        let slab_id =
            u8::try_from((self.location >> 32).trailing_zeros()).unwrap();
        let slab_idx = u32::try_from(self.location & IDX_MASK).unwrap();
        (slab_id, slab_idx, self.original_lsn)
    }

    pub fn compose(
        slab_id: SlabId,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
    ) -> HeapId {
        let slab = 1 << (32 + u64::from(slab_id));
        let heap_id = slab | u64::from(slab_idx);
        HeapId { location: heap_id, original_lsn }
    }

    fn offset(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id) * u64::from(idx)
    }

    fn slab_size(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id)
    }
}

pub(crate) fn slab_size(size: u64) -> u64 {
    slab_id_to_size(size_to_slab_id(size))
}

fn slab_id_to_size(slab_id: u8) -> u64 {
    1 << (MIN_TRAILING_ZEROS + u64::from(slab_id))
}

fn size_to_slab_id(size: u64) -> SlabId {
    // find the power of 2 that is at least 64k
    let normalized_size = std::cmp::max(MIN_SZ, size.next_power_of_two());

    // drop the lowest unused bits
    let rebased_size = normalized_size >> MIN_TRAILING_ZEROS;

    u8::try_from(rebased_size.trailing_zeros()).unwrap()
}

pub(crate) struct Reservation {
    slab_free: Arc<Stack<u32>>,
    completed: bool,
    file: File,
    pub heap_id: HeapId,
}

impl Drop for Reservation {
    fn drop(&mut self) {
        if !self.completed {
            let (slab_id, idx, _) = self.heap_id.decompose();
            self.slab_free.push(idx, &pin());
        }
    }
}

impl Reservation {
    pub fn complete(mut self, data: &[u8]) -> Result<HeapId> {
        log::trace!(
            "Heap::complete({:?}) to offset {} in file {:?}",
            self.heap_id,
            self.heap_id.offset(),
            self.file
        );
        assert_eq!(data.len() as u64, self.heap_id.slab_size());

        // write data
        pwrite_all(&self.file, data, self.heap_id.offset())?;

        // sync data
        #[cfg(target_os = "linux")]
        {
            use std::os::unix::io::AsRawFd;
            let ret = unsafe {
                libc::sync_file_range(
                    self.file.as_raw_fd(),
                    i64::try_from(self.heap_id.offset()).unwrap(),
                    i64::try_from(data.len()).unwrap(),
                    libc::SYNC_FILE_RANGE_WAIT_BEFORE
                        | libc::SYNC_FILE_RANGE_WRITE
                        | libc::SYNC_FILE_RANGE_WAIT_AFTER,
                )
            };
            if ret < 0 {
                let err = std::io::Error::last_os_error();
                if let Some(libc::ENOSYS) = err.raw_os_error() {
                    self.file.sync_all()?;
                } else {
                    return Err(err.into());
                }
            }
        }

        #[cfg(not(target_os = "linux"))]
        self.file.sync_all()?;

        // if this is not reached due to an IO error,
        // the offset will be returned to the Slab in Drop
        self.completed = true;

        Ok(self.heap_id)
    }

    pub fn abort(self) {
        // actual logic in Drop
    }
}

#[derive(Debug)]
pub(crate) struct Heap {
    // each slab stores
    // items that are double
    // the size of the previous,
    // ranging from 64k in the
    // smallest slab to 2^48 in
    // the last.
    slabs: [Slab; 32],
}

impl Heap {
    pub fn start<P: AsRef<Path>>(p: P) -> Result<Heap> {
        let mut slabs: [MaybeUninit<Slab>; 32] = unsafe { std::mem::zeroed() };

        for slab_id in 0..32 {
            let slab = Slab::start(&p, slab_id)?;
            slabs[slab_id as usize] = MaybeUninit::new(slab);
        }

        Ok(Heap { slabs: unsafe { transmute(slabs) } })
    }

    pub fn gc_unknown_items(
        &self,
        _snapshot: &crate::pagecache::Snapshot,
    ) -> Result<()> {
        //TODO todo!()
        Ok(())
    }

    pub fn read(
        &self,
        heap_id: HeapId,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        log::trace!("Heap::read({:?})", heap_id);
        let (slab_id, slab_idx, original_lsn) = heap_id.decompose();
        self.slabs[slab_id as usize].read(
            slab_idx,
            original_lsn,
            use_compression,
        )
    }

    pub fn free(&self, heap_id: HeapId) {
        log::trace!("Heap::free({:?})", heap_id);
        let (slab_id, slab_idx, _) = heap_id.decompose();
        self.slabs[slab_id as usize].free(slab_idx)
    }

    pub fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        assert!(size < 1 << 48);
        let slab_id = size_to_slab_id(size);
        let ret = self.slabs[slab_id as usize].reserve(size, original_lsn);
        log::trace!("Heap::reserve({}) -> {:?}", size, ret.heap_id);
        ret
    }
}

#[derive(Debug)]
struct Slab {
    file: File,
    slab_id: u8,
    tip: AtomicU32,
    free: Arc<Stack<u32>>,
}

impl Slab {
    pub fn start<P: AsRef<Path>>(directory: P, slab_id: u8) -> Result<Slab> {
        let bs = slab_id_to_size(slab_id);
        let free = Arc::new(Stack::default());

        let mut options = std::fs::OpenOptions::new();
        options.create(true);
        options.read(true);
        options.write(true);

        let file =
            options.open(directory.as_ref().join(format!("{:02}", slab_id)))?;
        let len = file.metadata()?.len();
        let max_idx = len / bs;
        log::trace!(
            "starting heap slab for sizes of {}. tip: {} max idx: {}",
            bs,
            len,
            max_idx
        );
        let tip = AtomicU32::new(u32::try_from(max_idx).unwrap());

        Ok(Slab { file, slab_id, tip, free })
    }

    fn read(
        &self,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        let bs = slab_id_to_size(self.slab_id);
        let offset = u64::from(slab_idx) * bs;

        log::trace!("reading heap slab slot {} at offset {}", slab_idx, offset);

        let mut heap_buf = vec![0; usize::try_from(bs).unwrap()];

        pread_exact(&self.file, &mut heap_buf, offset)?;

        let stored_crc =
            u32::from_le_bytes(heap_buf[1..5].as_ref().try_into().unwrap());

        let mut hasher = crc32fast::Hasher::new();
        hasher.update(&heap_buf[0..1]);
        hasher.update(&heap_buf[5..]);
        let actual_crc = hasher.finalize();

        if actual_crc == stored_crc {
            let actual_lsn = Lsn::from_le_bytes(
                heap_buf[5..13].as_ref().try_into().unwrap(),
            );
            if actual_lsn != original_lsn {
                log::error!(
                    "heap slot lsn {} does not match expected original lsn {}",
                    actual_lsn,
                    original_lsn
                );
                return Err(Error::corruption(None));
            }
            let buf = heap_buf[13..].to_vec();
            let buf = if use_compression {
                crate::pagecache::decompress(buf)
            } else {
                buf
            };
            Ok((MessageKind::from(heap_buf[0]), buf))
        } else {
            log::error!(
                "heap message CRC does not match contents. stored: {} actual: {}",
                stored_crc,
                actual_crc
            );
            Err(Error::corruption(None))
        }
    }

    fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        let idx = if let Some(idx) = self.free.pop(&pin()) {
            log::trace!(
                "reusing heap index {} in slab for sizes of {}",
                idx,
                slab_id_to_size(self.slab_id),
            );
            idx
        } else {
            log::trace!(
                "no free heap slots in slab for sizes of {}",
                slab_id_to_size(self.slab_id),
            );
            self.tip.fetch_add(1, Acquire)
        };

        log::trace!(
            "heap reservation for slot {} in the slab for sizes of {}",
            idx,
            slab_id_to_size(self.slab_id),
        );

        let heap_id = HeapId::compose(self.slab_id, idx, original_lsn);

        Reservation {
            slab_free: self.free.clone(),
            completed: false,
            file: self.file.try_clone().unwrap(),
            heap_id,
        }
    }

    fn free(&self, idx: u32) {
        self.punch_hole(idx);
        self.free.push(idx, &pin());
    }

    fn punch_hole(&self, idx: u32) {
        #[cfg(target_os = "linux")]
        {
            use std::{
                os::unix::io::AsRawFd,
                sync::atomic::{AtomicBool, Ordering::Relaxed},
            };

            use libc::{fallocate, FALLOC_FL_KEEP_SIZE, FALLOC_FL_PUNCH_HOLE};

            static HOLE_PUNCHING_ENABLED: AtomicBool = AtomicBool::new(true);
            const MODE: i32 = FALLOC_FL_KEEP_SIZE | FALLOC_FL_PUNCH_HOLE;

            if HOLE_PUNCHING_ENABLED.load(Relaxed) {
                let bs = i64::try_from(slab_id_to_size(self.slab_id)).unwrap();
                let offset = i64::from(idx) * bs;

                let fd = self.file.as_raw_fd();

                let ret = unsafe {
                    fallocate(
                        fd,
                        MODE,
                        offset.try_into().unwrap(),
                        bs.try_into().unwrap(),
                    )
                };

                if ret != 0 {
                    let err = std::io::Error::last_os_error();
                    log::error!(
                        "failed to punch hole in heap file: {:?}. disabling hole punching",
                        err
                    );
                    HOLE_PUNCHING_ENABLED.store(false, Relaxed);
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use boxcars::crc::calc_crc;
use boxcars::*;
use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};

fn bench_crc(c: &mut Criterion) {
    let data = include_bytes!("../assets/replays/good/rumble.replay");
    let mut group = c.benchmark_group("crc_throughput");
    group.throughput(Throughput::Bytes(data.len() as u64));
    group.bench_function("bench_crc", |b| {
        let data = include_bytes!("../assets/replays/good/rumble.replay");
        b.iter(|| {
            black_box(calc_crc(&data[..]));
        })
    });
    group.finish();
}

fn bench_json_serialization(c: &mut Criterion) {
    let data = include_bytes!("../assets/replays/good/3381.replay");
    let json_data_bytes = 19416972_u64;

    let mut group = c.benchmark_group("json_throughput");
    group.sample_size(10);
    group.throughput(Throughput::Bytes(json_data_bytes));
    group.bench_function("bench_json_serialization", |b| {
        let mut bytes = Vec::new();
        let replay = ParserBuilder::new(data)
            .on_error_check_crc()
            .parse()
            .unwrap();
        serde_json::to_writer(&mut bytes, &replay).unwrap();
        assert!(
            json_data_bytes == bytes.len() as u64,
            "update json benchmark with latest throughput size: {}",
            bytes.len()
        );
        unsafe {
            bytes.set_len(0);
        };

        b.iter(|| {
            black_box(serde_json::to_writer(&mut bytes, &replay).unwrap());
            unsafe {
                bytes.set_len(0);
            };
        })
    });
    group.finish();
}

fn bench_parse_crc_body(c: &mut Criterion) {
    let data = include_bytes!("../assets/replays/good/3381.replay");
    let mut group = c.benchmark_group("parse_crc_body");
    group.throughput(Throughput::Bytes(data.len() as u64));
    group.sample_size(20);
    group.bench_function("bench_parse_crc_body", |b| {
        b.iter(|| {
            black_box(
                ParserBuilder::new(data)
                    .always_check_crc()
                    .must_parse_network_data()
                    .parse()
                    .unwrap()
            )
        });
    });
    group.finish();
}

fn bench_parse_no_crc_body(c: &mut Criterion) {
    let data = include_bytes!("../assets/replays/good/3381.replay");
    let mut group = c.benchmark_group("parse_no_crc_body");
    group.throughput(Throughput::Bytes(data.len() as u64));
    group.sample_size(20);
    group.bench_function("bench_parse_no_crc_body", |b| {
        b.iter(|| {
            black_box(
                ParserBuilder::new(data)
                    .on_error_check_crc()
                    .must_parse_network_data()
                    .parse()
                    .unwrap()
            )
        });
    });
    group.finish();
}

fn bench_parse_no_crc_no_body(c: &mut Criterion) {
    // Throughput not included in this benchmark as it is a bit confusing what the number reporting
    // represents. If parsing the header reports based on the entire size of the replay, it will be
    // some very large number (20 GB/s), else it may look deceptively small
    let data = include_bytes!("../assets/replays/good/3381.replay");
    let mut group = c.benchmark_group("parse_no_crc_no_body");
    group.bench_function("bench_parse_no_crc_no_body", |b| {
        b.iter(|| {
            black_box(
                ParserBuilder::new(data)
                    .on_error_check_crc()
                    .never_parse_network_data()
                    .parse()
                    .unwrap()
            )
        });
    });
    group.finish();
}

fn bench_parse_crc_json(c: &mut Criterion) {
    let data = include_bytes!("../assets/replays/good/3381.replay");
    let mut group = c.benchmark_group("parse_crc_json");
    group.throughput(Throughput::Bytes(data.len() as u64));
    group.sample_size(10);
    group.bench_function("bench_parse_crc_json", |b| {
        // allocate a buffer big enough to hold all of the serialized data
        let mut bytes = Vec::with_capacity(2_usize.pow(25));
        b.iter(|| {
            let data = ParserBuilder::new(data).always_check_crc().parse().unwrap();
            black_box(serde_json::to_writer(&mut bytes, &data).unwrap());
            unsafe {
                bytes.set_len(0);
            }
        });
    });
    group.finish();
}

criterion_group!(
    benches,
    bench_crc,
    bench_json_serialization,
    bench_parse_crc_body,
    bench_parse_no_crc_body,
    bench_parse_no_crc_no_body,
    bench_parse_crc_json,
);

criterion_main!(benches);
<commit_msg>Update json serialization benchmark throughput<commit_after>use boxcars::crc::calc_crc;
use boxcars::*;
use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};

fn bench_crc(c: &mut Criterion) {
    let data = include_bytes!("../assets/replays/good/rumble.replay");
    let mut group = c.benchmark_group("crc_throughput");
    group.throughput(Throughput::Bytes(data.len() as u64));
    group.bench_function("bench_crc", |b| {
        let data = include_bytes!("../assets/replays/good/rumble.replay");
        b.iter(|| {
            black_box(calc_crc(&data[..]));
        })
    });
    group.finish();
}

fn bench_json_serialization(c: &mut Criterion) {
    let data = include_bytes!("../assets/replays/good/3381.replay");
    let json_data_bytes = 19483077_u64;

    let mut group = c.benchmark_group("json_throughput");
    group.sample_size(10);
    group.throughput(Throughput::Bytes(json_data_bytes));
    group.bench_function("bench_json_serialization", |b| {
        let mut bytes = Vec::new();
        let replay = ParserBuilder::new(data)
            .on_error_check_crc()
            .parse()
            .unwrap();
        serde_json::to_writer(&mut bytes, &replay).unwrap();
        assert!(
            json_data_bytes == bytes.len() as u64,
            "update json benchmark with latest throughput size: {}",
            bytes.len()
        );
        unsafe {
            bytes.set_len(0);
        };

        b.iter(|| {
            black_box(serde_json::to_writer(&mut bytes, &replay).unwrap());
            unsafe {
                bytes.set_len(0);
            };
        })
    });
    group.finish();
}

fn bench_parse_crc_body(c: &mut Criterion) {
    let data = include_bytes!("../assets/replays/good/3381.replay");
    let mut group = c.benchmark_group("parse_crc_body");
    group.throughput(Throughput::Bytes(data.len() as u64));
    group.sample_size(20);
    group.bench_function("bench_parse_crc_body", |b| {
        b.iter(|| {
            black_box(
                ParserBuilder::new(data)
                    .always_check_crc()
                    .must_parse_network_data()
                    .parse()
                    .unwrap()
            )
        });
    });
    group.finish();
}

fn bench_parse_no_crc_body(c: &mut Criterion) {
    let data = include_bytes!("../assets/replays/good/3381.replay");
    let mut group = c.benchmark_group("parse_no_crc_body");
    group.throughput(Throughput::Bytes(data.len() as u64));
    group.sample_size(20);
    group.bench_function("bench_parse_no_crc_body", |b| {
        b.iter(|| {
            black_box(
                ParserBuilder::new(data)
                    .on_error_check_crc()
                    .must_parse_network_data()
                    .parse()
                    .unwrap()
            )
        });
    });
    group.finish();
}

fn bench_parse_no_crc_no_body(c: &mut Criterion) {
    // Throughput not included in this benchmark as it is a bit confusing what the number reporting
    // represents. If parsing the header reports based on the entire size of the replay, it will be
    // some very large number (20 GB/s), else it may look deceptively small
    let data = include_bytes!("../assets/replays/good/3381.replay");
    let mut group = c.benchmark_group("parse_no_crc_no_body");
    group.bench_function("bench_parse_no_crc_no_body", |b| {
        b.iter(|| {
            black_box(
                ParserBuilder::new(data)
                    .on_error_check_crc()
                    .never_parse_network_data()
                    .parse()
                    .unwrap()
            )
        });
    });
    group.finish();
}

fn bench_parse_crc_json(c: &mut Criterion) {
    let data = include_bytes!("../assets/replays/good/3381.replay");
    let mut group = c.benchmark_group("parse_crc_json");
    group.throughput(Throughput::Bytes(data.len() as u64));
    group.sample_size(10);
    group.bench_function("bench_parse_crc_json", |b| {
        // allocate a buffer big enough to hold all of the serialized data
        let mut bytes = Vec::with_capacity(2_usize.pow(25));
        b.iter(|| {
            let data = ParserBuilder::new(data).always_check_crc().parse().unwrap();
            black_box(serde_json::to_writer(&mut bytes, &data).unwrap());
            unsafe {
                bytes.set_len(0);
            }
        });
    });
    group.finish();
}

criterion_group!(
    benches,
    bench_crc,
    bench_json_serialization,
    bench_parse_crc_body,
    bench_parse_no_crc_body,
    bench_parse_no_crc_no_body,
    bench_parse_crc_json,
);

criterion_main!(benches);
<|endoftext|>


<commit_before>// `cargo test --benches` and `#[feature(test)]` work only in nightly
#![cfg(rustc_nightly)]
#![feature(test)]

use std::sync::Arc;

extern crate bytes;
extern crate futures;
extern crate httpbis;
extern crate test;

use httpbis::*;

use futures::future::Future;
use futures::stream;
use futures::stream::Stream;

use bytes::Bytes;

use test::Bencher;

#[bench]
fn download_megabyte_in_kb_chunks(b: &mut Bencher) {
    b.iter(|| {
        struct Megabyte;

        impl Service for Megabyte {
            fn start_request(&self, _headers: Headers, _req: HttpStreamAfterHeaders) -> Response {
                Response::headers_and_bytes_stream(
                    Headers::ok_200(),
                    stream::iter_ok((0..1024).map(|i| Bytes::from(vec![(i % 0xff) as u8; 1024]))),
                )
            }
        }

        let mut server = ServerBuilder::new_plain();
        server.set_port(0);
        server.service.set_service("/", Arc::new(Megabyte));
        let server = server.build().expect("server");

        let client = Client::new_plain(
            "127.0.0.1",
            server.local_addr().port().unwrap(),
            Default::default(),
        ).expect("client");

        let (header, body) = client
            .start_get("/any", "localhost")
            .0
            .wait()
            .expect("headers");
        assert_eq!(200, header.status());

        let mut s = 0;
        for p in body.wait() {
            match p.expect("part") {
                DataOrTrailers::Data(d, ..) => s += d.len(),
                _ => panic!("unexpected headers"),
            }
        }

        assert_eq!(1024 * 1024, s);
    })
}

#[bench]
fn hundred_small_requests(b: &mut Bencher) {
    b.iter(|| {
        struct My;

        impl Service for My {
            fn start_request(&self, _headers: Headers, _req: HttpStreamAfterHeaders) -> Response {
                Response::found_200_plain_text("hello there")
            }
        }

        let mut server = ServerBuilder::new_plain();
        server.set_port(0);
        server.service.set_service("/", Arc::new(My));
        let server = server.build().expect("server");

        let client = Client::new_plain(
            "127.0.0.1",
            server.local_addr().port().unwrap(),
            Default::default(),
        ).expect("client");

        for _i in 0..100 {
            let (header, body) = client
                .start_get("/any", "localhost")
                .0
                .wait()
                .expect("headers");
            assert_eq!(200, header.status());

            // TODO: check content
            body.collect().wait().expect("body");
        }
    })
}
<commit_msg>Initialize client/server only once in benchmark<commit_after>// `cargo test --benches` and `#[feature(test)]` work only in nightly
#![cfg(rustc_nightly)]
#![feature(test)]

use std::sync::Arc;

extern crate bytes;
extern crate futures;
extern crate httpbis;
extern crate test;

use httpbis::*;

use futures::future::Future;
use futures::stream;
use futures::stream::Stream;

use bytes::Bytes;

use test::Bencher;

#[bench]
fn download_megabyte_in_kb_chunks(b: &mut Bencher) {
    struct Megabyte;

    impl Service for Megabyte {
        fn start_request(&self, _headers: Headers, _req: HttpStreamAfterHeaders) -> Response {
            Response::headers_and_bytes_stream(
                Headers::ok_200(),
                stream::iter_ok((0..1024).map(|i| Bytes::from(vec![(i % 0xff) as u8; 1024]))),
            )
        }
    }

    let mut server = ServerBuilder::new_plain();
    server.set_port(0);
    server.service.set_service("/", Arc::new(Megabyte));
    let server = server.build().expect("server");

    let client = Client::new_plain(
        "127.0.0.1",
        server.local_addr().port().unwrap(),
        Default::default(),
    ).expect("client");

    fn iter(client: &Client) {
        let (header, body) = client
            .start_get("/any", "localhost")
            .0
            .wait()
            .expect("headers");
        assert_eq!(200, header.status());

        let mut s = 0;
        for p in body.wait() {
            match p.expect("part") {
                DataOrTrailers::Data(d, ..) => s += d.len(),
                _ => panic!("unexpected headers"),
            }
        }

        assert_eq!(1024 * 1024, s);
    }

    // Warm-up
    iter(&client);

    b.iter(|| {
        iter(&client)
    })
}

#[bench]
fn small_requests(b: &mut Bencher) {
    struct My;

    impl Service for My {
        fn start_request(&self, _headers: Headers, _req: HttpStreamAfterHeaders) -> Response {
            Response::found_200_plain_text("hello there")
        }
    }

    let mut server = ServerBuilder::new_plain();
    server.set_port(0);
    server.service.set_service("/", Arc::new(My));
    let server = server.build().expect("server");

    let client = Client::new_plain(
        "127.0.0.1",
        server.local_addr().port().unwrap(),
        Default::default(),
    ).expect("client");

    b.iter(|| {
        let (header, body) = client
            .start_get("/any", "localhost")
            .0
            .wait()
            .expect("headers");
        assert_eq!(200, header.status());

        // TODO: check content
        body.collect().wait().expect("body")
    })
}
<|endoftext|>


<commit_before>#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
    }
}
<commit_msg>feat(rpc): Add TCP echo server<commit_after>use std::net::{TcpListener, TcpStream};
use std::io::{Read, Write, Error};
use std::thread;

macro_rules! println_stderr(
    ($($arg:tt)*) => { {
        let r = writeln!(&mut ::std::io::stderr(), $($arg)*);
        r.expect("failed printing to stderr");
    } }
);

///
/// Starts a new echo server that listens for connections in a background thread
/// responds to connections one at a time
/// function returns once echo server has been started and is listening
///
pub fn start_echo_server (port: u16) -> Result<(), Error>{
    let listener = try!(TcpListener::bind(("localhost", port)));
    thread::spawn(move || {
        for stream in listener.incoming() {
            match stream {
                Ok(mut stream) => {
                    let _ = do_echo(&mut stream);
                }
                Err(e) => {
                    println_stderr!("{}", e);
                }
            }
        }
    });
    Ok(())
}

fn do_echo(stream: &mut TcpStream) -> Result<(), Error>{
    let mut buf: String = String::new();
    stream.read_to_string(&mut buf)
          .and_then(|_| {
            stream.write_all(buf.as_bytes())
          })
}

#[cfg(test)]
mod tests {
    use super::start_echo_server;
    use std::net::{TcpStream, Shutdown};
    use std::io::{Read, Write};
    use std::time::Duration;

    #[test]
    fn it_echos() {
        // constants
        let port = 8080;
        let message = "Hello from\n test land!";

        // start the echo server and connect to it
        let _ = start_echo_server(port).unwrap();
        let mut client = TcpStream::connect(("localhost", port)).unwrap();

        // write to the server, send a fin, and wait for the response
        let mut result = String::new();
        let num_bytes_read = client.write_all(message.as_bytes())
        .and_then(|()| client.shutdown(Shutdown::Write))
        .and_then(|()| {
            client.read_to_string(&mut result)
        }).unwrap();

        // check that the response is correct
        assert!(num_bytes_read == message.len());
        assert!(result == message);
    }
}
<|endoftext|>


<commit_before>/* automatically generated by rust-bindgen */

use libc::*;
use sdl2::rwops::ll::SDL_RWops;
use sdl2::version::ll::SDL_version;

pub type MIX_InitFlags = c_uint;
pub static MIX_INIT_FLAC: c_uint = 1;
pub static MIX_INIT_MOD: c_uint = 2;
pub static MIX_INIT_MODPLUG: c_uint = 4;
pub static MIX_INIT_MP3: c_uint = 8;
pub static MIX_INIT_OGG: c_uint = 16;
pub static MIX_INIT_FLUIDSYNTH: c_uint = 32;
pub struct Struct_Mix_Chunk {
    pub allocated: c_int,
    pub abuf: *uint8_t,
    pub alen: uint32_t,
    pub volume: uint8_t,
}
pub type Mix_Chunk = Struct_Mix_Chunk;
pub type Mix_Fading = c_uint;
pub static MIX_NO_FADING: c_uint = 0;
pub static MIX_FADING_OUT: c_uint = 1;
pub static MIX_FADING_IN: c_uint = 2;
pub type Mix_MusicType = c_uint;
pub static MUS_NONE: c_uint = 0;
pub static MUS_CMD: c_uint = 1;
pub static MUS_WAV: c_uint = 2;
pub static MUS_MOD: c_uint = 3;
pub static MUS_MID: c_uint = 4;
pub static MUS_OGG: c_uint = 5;
pub static MUS_MP3: c_uint = 6;
pub static MUS_MP3_MAD: c_uint = 7;
pub static MUS_FLAC: c_uint = 8;
pub static MUS_MODPLUG: c_uint = 9;
pub type Struct__Mix_Music = c_void;
pub type Mix_Music = Struct__Mix_Music;
pub type Mix_EffectFunc_t =
    ::std::option::Option<extern "C" fn
                              (arg1: c_int, arg2: *c_void, arg3: c_int,
                               arg4: *c_void)>;
pub type Mix_EffectDone_t =
    ::std::option::Option<extern "C" fn(arg1: c_int, arg2: *c_void)>;
extern "C" {
    pub fn Mix_Linked_Version() -> *SDL_version;
    pub fn Mix_Init(flags: c_int) -> c_int;
    pub fn Mix_Quit();
    pub fn Mix_OpenAudio(frequency: c_int, format: uint16_t, channels: c_int,
                         chunksize: c_int) -> c_int;
    pub fn Mix_AllocateChannels(numchans: c_int) -> c_int;
    pub fn Mix_QuerySpec(frequency: *c_int, format: *uint16_t,
                         channels: *c_int) -> c_int;
    pub fn Mix_LoadWAV_RW(src: *SDL_RWops, freesrc: c_int) ->
     *Mix_Chunk;
    pub fn Mix_LoadMUS(file: *c_schar) -> *Mix_Music;
    pub fn Mix_LoadMUS_RW(src: *SDL_RWops, freesrc: c_int) ->
     *Mix_Music;
    pub fn Mix_LoadMUSType_RW(src: *SDL_RWops, _type: Mix_MusicType,
                              freesrc: c_int) -> *Mix_Music;
    pub fn Mix_QuickLoad_WAV(mem: *uint8_t) -> *Mix_Chunk;
    pub fn Mix_QuickLoad_RAW(mem: *uint8_t, len: uint32_t) -> *Mix_Chunk;
    pub fn Mix_FreeChunk(chunk: *Mix_Chunk);
    pub fn Mix_FreeMusic(music: *Mix_Music);
    pub fn Mix_GetNumChunkDecoders() -> c_int;
    pub fn Mix_GetChunkDecoder(index: c_int) -> *c_schar;
    pub fn Mix_GetNumMusicDecoders() -> c_int;
    pub fn Mix_GetMusicDecoder(index: c_int) -> *c_schar;
    pub fn Mix_GetMusicType(music: *Mix_Music) -> Mix_MusicType;
    pub fn Mix_SetPostMix(mix_func:
                              ::std::option::Option<extern "C" fn
                                                        (arg1: *c_void,
                                                         arg2: *uint8_t,
                                                         arg3: c_int)>,
                          arg: *c_void);
    pub fn Mix_HookMusic(mix_func:
                             ::std::option::Option<extern "C" fn
                                                       (arg1: *c_void,
                                                        arg2: *uint8_t,
                                                        arg3: c_int)>,
                         arg: *c_void);
    pub fn Mix_HookMusicFinished(music_finished:
                                     ::std::option::Option<extern "C" fn()>);
    pub fn Mix_GetMusicHookData() -> *c_void;
    pub fn Mix_ChannelFinished(channel_finished:
                                   ::std::option::Option<extern "C" fn
                                                             (arg1: c_int)>);
    pub fn Mix_RegisterEffect(chan: c_int, f: Mix_EffectFunc_t,
                              d: Mix_EffectDone_t, arg: *c_void) -> c_int;
    pub fn Mix_UnregisterEffect(channel: c_int, f: Mix_EffectFunc_t) -> c_int;
    pub fn Mix_UnregisterAllEffects(channel: c_int) -> c_int;
    pub fn Mix_SetPanning(channel: c_int, left: uint8_t, right: uint8_t) -> c_int;
    pub fn Mix_SetPosition(channel: c_int, angle: int16_t, distance: uint8_t) ->
     c_int;
    pub fn Mix_SetDistance(channel: c_int, distance: uint8_t) -> c_int;
    pub fn Mix_SetReverseStereo(channel: c_int, flip: c_int) -> c_int;
    pub fn Mix_ReserveChannels(num: c_int) -> c_int;
    pub fn Mix_GroupChannel(which: c_int, tag: c_int) -> c_int;
    pub fn Mix_GroupChannels(from: c_int, to: c_int, tag: c_int) -> c_int;
    pub fn Mix_GroupAvailable(tag: c_int) -> c_int;
    pub fn Mix_GroupCount(tag: c_int) -> c_int;
    pub fn Mix_GroupOldest(tag: c_int) -> c_int;
    pub fn Mix_GroupNewer(tag: c_int) -> c_int;
    pub fn Mix_PlayChannelTimed(channel: c_int, chunk: *Mix_Chunk,
                                loops: c_int, ticks: c_int) -> c_int;
    pub fn Mix_PlayMusic(music: *Mix_Music, loops: c_int) -> c_int;
    pub fn Mix_FadeInMusic(music: *Mix_Music, loops: c_int, ms: c_int) ->
     c_int;
    pub fn Mix_FadeInMusicPos(music: *Mix_Music, loops: c_int, ms: c_int,
                              position: c_double) -> c_int;
    pub fn Mix_FadeInChannelTimed(channel: c_int, chunk: *Mix_Chunk,
                                  loops: c_int, ms: c_int, ticks: c_int) ->
     c_int;
    pub fn Mix_Volume(channel: c_int, volume: c_int) -> c_int;
    pub fn Mix_VolumeChunk(chunk: *Mix_Chunk, volume: c_int) -> c_int;
    pub fn Mix_VolumeMusic(volume: c_int) -> c_int;
    pub fn Mix_HaltChannel(channel: c_int) -> c_int;
    pub fn Mix_HaltGroup(tag: c_int) -> c_int;
    pub fn Mix_HaltMusic() -> c_int;
    pub fn Mix_ExpireChannel(channel: c_int, ticks: c_int) -> c_int;
    pub fn Mix_FadeOutChannel(which: c_int, ms: c_int) -> c_int;
    pub fn Mix_FadeOutGroup(tag: c_int, ms: c_int) -> c_int;
    pub fn Mix_FadeOutMusic(ms: c_int) -> c_int;
    pub fn Mix_FadingMusic() -> Mix_Fading;
    pub fn Mix_FadingChannel(which: c_int) -> Mix_Fading;
    pub fn Mix_Pause(channel: c_int);
    pub fn Mix_Resume(channel: c_int);
    pub fn Mix_Paused(channel: c_int) -> c_int;
    pub fn Mix_PauseMusic();
    pub fn Mix_ResumeMusic();
    pub fn Mix_RewindMusic();
    pub fn Mix_PausedMusic() -> c_int;
    pub fn Mix_SetMusicPosition(position: c_double) -> c_int;
    pub fn Mix_Playing(channel: c_int) -> c_int;
    pub fn Mix_PlayingMusic() -> c_int;
    pub fn Mix_SetMusicCMD(command: *c_schar) -> c_int;
    pub fn Mix_SetSynchroValue(value: c_int) -> c_int;
    pub fn Mix_GetSynchroValue() -> c_int;
    pub fn Mix_SetSoundFonts(paths: *c_schar) -> c_int;
    pub fn Mix_GetSoundFonts() -> *c_schar;
    pub fn Mix_EachSoundFont(function:
                                 ::std::option::Option<extern "C" fn
                                                           (arg1: *c_schar,
                                                            arg2: *c_void)
                                                           -> c_int>,
                             data: *c_void) -> c_int;
    pub fn Mix_GetChunk(channel: c_int) -> *Mix_Chunk;
    pub fn Mix_CloseAudio();
}
<commit_msg>fix libc glob reexports change<commit_after>/* automatically generated by rust-bindgen */

use libc::*;
use sdl2::rwops::ll::SDL_RWops;
use sdl2::version::ll::SDL_version;

pub type MIX_InitFlags = c_uint;
pub static MIX_INIT_FLAC: c_uint = 1;
pub static MIX_INIT_MOD: c_uint = 2;
pub static MIX_INIT_MODPLUG: c_uint = 4;
pub static MIX_INIT_MP3: c_uint = 8;
pub static MIX_INIT_OGG: c_uint = 16;
pub static MIX_INIT_FLUIDSYNTH: c_uint = 32;
pub struct Struct_Mix_Chunk {
    pub allocated: c_int,
    pub abuf: *uint8_t,
    pub alen: uint32_t,
    pub volume: uint8_t,
}
pub type Mix_Chunk = Struct_Mix_Chunk;
pub type Mix_Fading = c_uint;
pub static MIX_NO_FADING: c_uint = 0;
pub static MIX_FADING_OUT: c_uint = 1;
pub static MIX_FADING_IN: c_uint = 2;
pub type Mix_MusicType = c_uint;
pub static MUS_NONE: c_uint = 0;
pub static MUS_CMD: c_uint = 1;
pub static MUS_WAV: c_uint = 2;
pub static MUS_MOD: c_uint = 3;
pub static MUS_MID: c_uint = 4;
pub static MUS_OGG: c_uint = 5;
pub static MUS_MP3: c_uint = 6;
pub static MUS_MP3_MAD: c_uint = 7;
pub static MUS_FLAC: c_uint = 8;
pub static MUS_MODPLUG: c_uint = 9;
pub type Struct__Mix_Music = c_void;
pub type Mix_Music = Struct__Mix_Music;
pub type Mix_EffectFunc_t =
    ::std::option::Option<extern "C" fn
                              (arg1: c_int, arg2: *c_void, arg3: c_int,
                               arg4: *c_void)>;
pub type Mix_EffectDone_t =
    ::std::option::Option<extern "C" fn(arg1: c_int, arg2: *c_void)>;
extern "C" {
    pub fn Mix_Linked_Version() -> *SDL_version;
    pub fn Mix_Init(flags: c_int) -> c_int;
    pub fn Mix_Quit();
    pub fn Mix_OpenAudio(frequency: c_int, format: uint16_t, channels: c_int,
                         chunksize: c_int) -> c_int;
    pub fn Mix_AllocateChannels(numchans: c_int) -> c_int;
    pub fn Mix_QuerySpec(frequency: *c_int, format: *uint16_t,
                         channels: *c_int) -> c_int;
    pub fn Mix_LoadWAV_RW(src: *SDL_RWops, freesrc: c_int) ->
     *Mix_Chunk;
    pub fn Mix_LoadMUS(file: *c_char) -> *Mix_Music;
    pub fn Mix_LoadMUS_RW(src: *SDL_RWops, freesrc: c_int) ->
     *Mix_Music;
    pub fn Mix_LoadMUSType_RW(src: *SDL_RWops, _type: Mix_MusicType,
                              freesrc: c_int) -> *Mix_Music;
    pub fn Mix_QuickLoad_WAV(mem: *uint8_t) -> *Mix_Chunk;
    pub fn Mix_QuickLoad_RAW(mem: *uint8_t, len: uint32_t) -> *Mix_Chunk;
    pub fn Mix_FreeChunk(chunk: *Mix_Chunk);
    pub fn Mix_FreeMusic(music: *Mix_Music);
    pub fn Mix_GetNumChunkDecoders() -> c_int;
    pub fn Mix_GetChunkDecoder(index: c_int) -> *c_char;
    pub fn Mix_GetNumMusicDecoders() -> c_int;
    pub fn Mix_GetMusicDecoder(index: c_int) -> *c_char;
    pub fn Mix_GetMusicType(music: *Mix_Music) -> Mix_MusicType;
    pub fn Mix_SetPostMix(mix_func:
                              ::std::option::Option<extern "C" fn
                                                        (arg1: *c_void,
                                                         arg2: *uint8_t,
                                                         arg3: c_int)>,
                          arg: *c_void);
    pub fn Mix_HookMusic(mix_func:
                             ::std::option::Option<extern "C" fn
                                                       (arg1: *c_void,
                                                        arg2: *uint8_t,
                                                        arg3: c_int)>,
                         arg: *c_void);
    pub fn Mix_HookMusicFinished(music_finished:
                                     ::std::option::Option<extern "C" fn()>);
    pub fn Mix_GetMusicHookData() -> *c_void;
    pub fn Mix_ChannelFinished(channel_finished:
                                   ::std::option::Option<extern "C" fn
                                                             (arg1: c_int)>);
    pub fn Mix_RegisterEffect(chan: c_int, f: Mix_EffectFunc_t,
                              d: Mix_EffectDone_t, arg: *c_void) -> c_int;
    pub fn Mix_UnregisterEffect(channel: c_int, f: Mix_EffectFunc_t) -> c_int;
    pub fn Mix_UnregisterAllEffects(channel: c_int) -> c_int;
    pub fn Mix_SetPanning(channel: c_int, left: uint8_t, right: uint8_t) -> c_int;
    pub fn Mix_SetPosition(channel: c_int, angle: int16_t, distance: uint8_t) ->
     c_int;
    pub fn Mix_SetDistance(channel: c_int, distance: uint8_t) -> c_int;
    pub fn Mix_SetReverseStereo(channel: c_int, flip: c_int) -> c_int;
    pub fn Mix_ReserveChannels(num: c_int) -> c_int;
    pub fn Mix_GroupChannel(which: c_int, tag: c_int) -> c_int;
    pub fn Mix_GroupChannels(from: c_int, to: c_int, tag: c_int) -> c_int;
    pub fn Mix_GroupAvailable(tag: c_int) -> c_int;
    pub fn Mix_GroupCount(tag: c_int) -> c_int;
    pub fn Mix_GroupOldest(tag: c_int) -> c_int;
    pub fn Mix_GroupNewer(tag: c_int) -> c_int;
    pub fn Mix_PlayChannelTimed(channel: c_int, chunk: *Mix_Chunk,
                                loops: c_int, ticks: c_int) -> c_int;
    pub fn Mix_PlayMusic(music: *Mix_Music, loops: c_int) -> c_int;
    pub fn Mix_FadeInMusic(music: *Mix_Music, loops: c_int, ms: c_int) ->
     c_int;
    pub fn Mix_FadeInMusicPos(music: *Mix_Music, loops: c_int, ms: c_int,
                              position: c_double) -> c_int;
    pub fn Mix_FadeInChannelTimed(channel: c_int, chunk: *Mix_Chunk,
                                  loops: c_int, ms: c_int, ticks: c_int) ->
     c_int;
    pub fn Mix_Volume(channel: c_int, volume: c_int) -> c_int;
    pub fn Mix_VolumeChunk(chunk: *Mix_Chunk, volume: c_int) -> c_int;
    pub fn Mix_VolumeMusic(volume: c_int) -> c_int;
    pub fn Mix_HaltChannel(channel: c_int) -> c_int;
    pub fn Mix_HaltGroup(tag: c_int) -> c_int;
    pub fn Mix_HaltMusic() -> c_int;
    pub fn Mix_ExpireChannel(channel: c_int, ticks: c_int) -> c_int;
    pub fn Mix_FadeOutChannel(which: c_int, ms: c_int) -> c_int;
    pub fn Mix_FadeOutGroup(tag: c_int, ms: c_int) -> c_int;
    pub fn Mix_FadeOutMusic(ms: c_int) -> c_int;
    pub fn Mix_FadingMusic() -> Mix_Fading;
    pub fn Mix_FadingChannel(which: c_int) -> Mix_Fading;
    pub fn Mix_Pause(channel: c_int);
    pub fn Mix_Resume(channel: c_int);
    pub fn Mix_Paused(channel: c_int) -> c_int;
    pub fn Mix_PauseMusic();
    pub fn Mix_ResumeMusic();
    pub fn Mix_RewindMusic();
    pub fn Mix_PausedMusic() -> c_int;
    pub fn Mix_SetMusicPosition(position: c_double) -> c_int;
    pub fn Mix_Playing(channel: c_int) -> c_int;
    pub fn Mix_PlayingMusic() -> c_int;
    pub fn Mix_SetMusicCMD(command: *c_char) -> c_int;
    pub fn Mix_SetSynchroValue(value: c_int) -> c_int;
    pub fn Mix_GetSynchroValue() -> c_int;
    pub fn Mix_SetSoundFonts(paths: *c_char) -> c_int;
    pub fn Mix_GetSoundFonts() -> *c_char;
    pub fn Mix_EachSoundFont(function:
                                 ::std::option::Option<extern "C" fn
                                                           (arg1: *c_char,
                                                            arg2: *c_void)
                                                           -> c_int>,
                             data: *c_void) -> c_int;
    pub fn Mix_GetChunk(channel: c_int) -> *Mix_Chunk;
    pub fn Mix_CloseAudio();
}
<|endoftext|>