<commit_before>#![cfg(feature = "full")]

#[macro_use]
extern crate quote;
extern crate syn;
extern crate syntex_pos;
extern crate syntex_syntax;
extern crate time;
extern crate walkdir;

use syntex_pos::Span;
use syntex_syntax::ast;
use syntex_syntax::parse::{self, ParseSess, PResult};
use time::PreciseTime;

use std::fs::File;
use std::io::{self, Read, Write};

macro_rules! errorf {
    ($($tt:tt)*) => {
        write!(io::stderr(), $($tt)*).unwrap();
    };
}

#[test]
fn test_round_trip() {
    let mut failed = 0;

    let walk = walkdir::WalkDir::new("tests/cases").sort_by(|a, b| a.cmp(b));
    for entry in walk.into_iter() {
        let entry = entry.unwrap();

        let path = entry.path();
        if path.extension().map(|e| e != "rs").unwrap_or(true) {
            continue;
        }
        errorf!("=== {}: ", path.display());

        let mut file = File::open(path).unwrap();
        let mut content = String::new();
        file.read_to_string(&mut content).unwrap();

        let start = PreciseTime::now();
        let (krate, elapsed) = match syn::parse_crate(&content) {
            Ok(krate) => (krate, start.to(PreciseTime::now())),
            Err(msg) => {
                errorf!("syn failed to parse\n{}\n", msg);
                failed += 1;
                continue;
            }
        };
        let back = quote!(#krate).to_string();

        let sess = ParseSess::new();
        let before = match syntex_parse(content, &sess) {
            Ok(before) => before,
            Err(mut diagnostic) => {
                diagnostic.cancel();
                if diagnostic.message.starts_with("file not found for module") ||
                   diagnostic.message.starts_with("couldn't read") {
                    errorf!("ignore\n");
                    continue;
                } else {
                    panic!(diagnostic.message.clone());
                }
            }
        };
        let after = match syntex_parse(back, &sess) {
            Ok(after) => after,
            Err(mut diagnostic) => {
                errorf!("syntex failed to parse");
                diagnostic.emit();
                failed += 1;
                continue;
            }
        };

        if before == after {
            errorf!("pass in {}ms\n", elapsed.num_milliseconds());
        } else {
            errorf!("FAIL\nbefore: {}\nafter: {}\n",
                    format!("{:?}", before).replace("\n", ""),
                    format!("{:?}", after).replace("\n", ""));
            failed += 1;
        }
    }

    if failed > 0 {
        panic!("{} failures", failed);
    }
}

fn syntex_parse<'a>(content: String, sess: &'a ParseSess) -> PResult<'a, ast::Crate> {
    let name = "test_round_trip".to_string();
    let cfg = Vec::new();
    parse::parse_crate_from_source_str(name, content, cfg, sess).map(respan_crate)
}

fn respan_crate(krate: ast::Crate) -> ast::Crate {
    use std::rc::Rc;
    use syntex_syntax::ast::{Attribute, Expr, ExprKind, Field, FnDecl, FunctionRetTy, ImplItem,
                             ImplItemKind, ItemKind, Mac, MetaItem, MetaItemKind, MethodSig,
                             NestedMetaItem, NestedMetaItemKind, TraitItem, TraitItemKind, TyParam};
    use syntex_syntax::codemap::{self, Spanned};
    use syntex_syntax::fold::{self, Folder};
    use syntex_syntax::parse::token::{Lit, Token, intern};
    use syntex_syntax::ptr::P;
    use syntex_syntax::tokenstream::{Delimited, SequenceRepetition, TokenTree};
    use syntex_syntax::util::move_map::MoveMap;
    use syntex_syntax::util::small_vector::SmallVector;

    struct Respanner;

    impl Respanner {
        fn fold_spanned<T>(&mut self, spanned: Spanned<T>) -> Spanned<T> {
            codemap::respan(self.new_span(spanned.span), spanned.node)
        }

        fn fold_lit(&mut self, l: Lit) -> Lit {
            // Give up on comparing literals inside of macros because there are
            // so many equivalent representations of the same literal; they are
            // tested elsewhere
            match l {
                Lit::Byte(_) => Lit::Byte(intern("")),
                Lit::Char(_) => Lit::Char(intern("")),
                Lit::Integer(_) => Lit::Integer(intern("")),
                Lit::Float(_) => Lit::Float(intern("")),
                Lit::Str_(_) => Lit::Str_(intern("")),
                Lit::ByteStr(_) => Lit::ByteStr(intern("")),
                _ => l,
            }
        }
    }

    impl Folder for Respanner {
        fn new_span(&mut self, _: Span) -> Span {
            syntex_pos::DUMMY_SP
        }

        fn fold_item_kind(&mut self, i: ItemKind) -> ItemKind {
            match i {
                ItemKind::Fn(decl, unsafety, constness, abi, generics, body) => {
                    let generics = self.fold_generics(generics);
                    let decl = self.fold_fn_decl(decl);
                    let body = self.fold_block(body);
                    // default fold_item_kind does not fold this span
                    let constness = self.fold_spanned(constness);
                    ItemKind::Fn(decl, unsafety, constness, abi, generics, body)
                }
                _ => fold::noop_fold_item_kind(i, self),
            }
        }

        fn fold_expr(&mut self, e: P<Expr>) -> P<Expr> {
            e.map(|e| {
                let folded = fold::noop_fold_expr(e, self);
                Expr {
                    node: match folded.node {
                        ExprKind::Lit(l) => {
                            // default fold_expr does not fold lits
                            ExprKind::Lit(l.map(|l| self.fold_spanned(l)))
                        }
                        ExprKind::Binary(op, lhs, rhs) => {
                            // default fold_expr does not fold the op span
                            ExprKind::Binary(self.fold_spanned(op),
                                             self.fold_expr(lhs),
                                             self.fold_expr(rhs))
                        }
                        ExprKind::AssignOp(op, lhs, rhs) => {
                            // default fold_expr does not fold the op span
                            ExprKind::AssignOp(self.fold_spanned(op),
                                               self.fold_expr(lhs),
                                               self.fold_expr(rhs))
                        }
                        other => other,
                    },
                    ..folded
                }
            })
        }

        fn fold_ty_param(&mut self, tp: TyParam) -> TyParam {
            TyParam {
                // default fold_ty_param does not fold the span
                span: self.new_span(tp.span),
                ..fold::noop_fold_ty_param(tp, self)
            }
        }

        fn fold_fn_decl(&mut self, decl: P<FnDecl>) -> P<FnDecl> {
            decl.map(|FnDecl { inputs, output, variadic }| {
                FnDecl {
                    inputs: inputs.move_map(|x| self.fold_arg(x)),
                    output: match output {
                        FunctionRetTy::Ty(ty) => FunctionRetTy::Ty(self.fold_ty(ty)),
                        // default fold_fn_decl does not fold this span
                        FunctionRetTy::Default(span) => FunctionRetTy::Default(self.new_span(span)),
                    },
                    variadic: variadic,
                }
            })
        }

        fn fold_field(&mut self, field: Field) -> Field {
            Field {
                ident: codemap::respan(// default fold_field does not fold this span
                                       self.new_span(field.ident.span),
                                       self.fold_ident(field.ident.node)),
                expr: self.fold_expr(field.expr),
                span: self.new_span(field.span),
            }
        }

        fn fold_trait_item(&mut self, i: TraitItem) -> SmallVector<TraitItem> {
            let noop = fold::noop_fold_trait_item(i, self).expect_one("");
            SmallVector::one(TraitItem {
                node: match noop.node {
                    TraitItemKind::Method(sig, body) => {
                        TraitItemKind::Method(MethodSig {
                                                  constness: self.fold_spanned(sig.constness),
                                                  ..sig
                                              },
                                              body)
                    }
                    node => node,
                },
                ..noop
            })
        }

        fn fold_impl_item(&mut self, i: ImplItem) -> SmallVector<ImplItem> {
            let noop = fold::noop_fold_impl_item(i, self).expect_one("");
            SmallVector::one(ImplItem {
                node: match noop.node {
                    ImplItemKind::Method(sig, body) => {
                        ImplItemKind::Method(MethodSig {
                                                 constness: self.fold_spanned(sig.constness),
                                                 ..sig
                                             },
                                             body)
                    }
                    node => node,
                },
                ..noop
            })
        }

        fn fold_attribute(&mut self, mut at: Attribute) -> Option<Attribute> {
            at.node.id.0 = 0;
            fold::noop_fold_attribute(at, self)
        }

        fn fold_meta_item(&mut self, meta_item: P<MetaItem>) -> P<MetaItem> {
            meta_item.map(|Spanned { node, span }| {
                Spanned {
                    node: match node {
                        MetaItemKind::Word(id) => MetaItemKind::Word(id),
                        MetaItemKind::List(id, mis) => {
                            MetaItemKind::List(id, mis.move_map(|e| self.fold_meta_list_item(e)))
                        }
                        // default fold_meta_item does not fold the value span
                        MetaItemKind::NameValue(id, lit) => {
                            MetaItemKind::NameValue(id, self.fold_spanned(lit))
                        }
                    },
                    span: self.new_span(span),
                }
            })
        }

        fn fold_meta_list_item(&mut self, list_item: NestedMetaItem) -> NestedMetaItem {
            Spanned {
                node: match list_item.node {
                    NestedMetaItemKind::MetaItem(mi) => {
                        NestedMetaItemKind::MetaItem(self.fold_meta_item(mi))
                    }
                    // default fold_meta_list_item does not fold the span
                    NestedMetaItemKind::Literal(lit) => {
                        NestedMetaItemKind::Literal(self.fold_spanned(lit))
                    }
                },
                span: self.new_span(list_item.span),
            }
        }

        fn fold_mac(&mut self, mac: Mac) -> Mac {
            fold::noop_fold_mac(mac, self)
        }

        fn fold_tt(&mut self, tt: &TokenTree) -> TokenTree {
            match *tt {
                TokenTree::Token(span, ref tok) => {
                    TokenTree::Token(self.new_span(span), self.fold_token(tok.clone()))
                }
                TokenTree::Delimited(span, ref delimed) => {
                    TokenTree::Delimited(self.new_span(span),
                                         Rc::new(Delimited {
                                             delim: delimed.delim,
                                             open_span: self.new_span(delimed.open_span),
                                             tts: self.fold_tts(&delimed.tts),
                                             close_span: self.new_span(delimed.close_span),
                                         }))
                }
                TokenTree::Sequence(span, ref seq) => {
                    TokenTree::Sequence(self.new_span(span),
                                        Rc::new(SequenceRepetition {
                                            tts: self.fold_tts(&seq.tts),
                                            separator: seq.separator
                                                .clone()
                                                .map(|tok| self.fold_token(tok)),
                                            ..**seq
                                        }))
                }
            }
        }

        fn fold_token(&mut self, t: Token) -> Token {
            match t {
                // default fold_token does not fold literals
                Token::Literal(lit, repr) => Token::Literal(self.fold_lit(lit), repr),
                Token::Ident(id) => Token::Ident(self.fold_ident(id)),
                Token::Lifetime(id) => Token::Lifetime(self.fold_ident(id)),
                Token::Interpolated(nt) => Token::Interpolated(self.fold_interpolated(nt)),
                Token::SubstNt(ident) => Token::SubstNt(self.fold_ident(ident)),
                Token::MatchNt(name, kind) => {
                    Token::MatchNt(self.fold_ident(name), self.fold_ident(kind))
                }
                _ => t,
            }
        }
    }

    Respanner.fold_crate(krate)
}
<commit_msg>Catch panics from syntex in test<commit_after>#![cfg(feature = "full")]

#[macro_use]
extern crate quote;
extern crate syn;
extern crate syntex_pos;
extern crate syntex_syntax;
extern crate time;
extern crate walkdir;

use syntex_pos::Span;
use syntex_syntax::ast;
use syntex_syntax::parse::{self, ParseSess, PResult};
use time::PreciseTime;

use std::fs::File;
use std::io::{self, Read, Write};
use std::panic;

macro_rules! errorf {
    ($($tt:tt)*) => {
        write!(io::stderr(), $($tt)*).unwrap();
    };
}

#[test]
fn test_round_trip() {
    let mut failed = 0;

    let walk = walkdir::WalkDir::new("tests/cases").sort_by(|a, b| a.cmp(b));
    for entry in walk.into_iter() {
        let entry = entry.unwrap();

        let path = entry.path();
        if path.extension().map(|e| e != "rs").unwrap_or(true) {
            continue;
        }
        errorf!("=== {}: ", path.display());

        let mut file = File::open(path).unwrap();
        let mut content = String::new();
        file.read_to_string(&mut content).unwrap();

        let start = PreciseTime::now();
        let (krate, elapsed) = match syn::parse_crate(&content) {
            Ok(krate) => (krate, start.to(PreciseTime::now())),
            Err(msg) => {
                errorf!("syn failed to parse\n{}\n", msg);
                failed += 1;
                continue;
            }
        };
        let back = quote!(#krate).to_string();

        let equal = panic::catch_unwind(|| {
            let sess = ParseSess::new();
            let before = match syntex_parse(content, &sess) {
                Ok(before) => before,
                Err(mut diagnostic) => {
                    diagnostic.cancel();
                    if diagnostic.message.starts_with("file not found for module") {
                        errorf!("ignore\n");
                        return true;
                    } else {
                        errorf!("FAIL: {}\n", diagnostic.message);
                        return false;
                    }
                }
            };
            let after = match syntex_parse(back, &sess) {
                Ok(after) => after,
                Err(mut diagnostic) => {
                    errorf!("syntex failed to parse");
                    diagnostic.emit();
                    return false;
                }
            };

            if before == after {
                errorf!("pass in {}ms\n", elapsed.num_milliseconds());
                true
            } else {
                errorf!("FAIL\nbefore: {}\nafter: {}\n",
                        format!("{:?}", before).replace("\n", ""),
                        format!("{:?}", after).replace("\n", ""));
                false
            }
        });
        match equal {
            Err(_) => errorf!("syntex panic\n"),
            Ok(true) => {},
            Ok(false) => failed += 1,
        }
    }

    if failed > 0 {
        panic!("{} failures", failed);
    }
}

fn syntex_parse<'a>(content: String, sess: &'a ParseSess) -> PResult<'a, ast::Crate> {
    let name = "test_round_trip".to_string();
    let cfg = Vec::new();
    parse::parse_crate_from_source_str(name, content, cfg, sess).map(respan_crate)
}

fn respan_crate(krate: ast::Crate) -> ast::Crate {
    use std::rc::Rc;
    use syntex_syntax::ast::{Attribute, Expr, ExprKind, Field, FnDecl, FunctionRetTy, ImplItem,
                             ImplItemKind, ItemKind, Mac, MetaItem, MetaItemKind, MethodSig,
                             NestedMetaItem, NestedMetaItemKind, TraitItem, TraitItemKind, TyParam};
    use syntex_syntax::codemap::{self, Spanned};
    use syntex_syntax::fold::{self, Folder};
    use syntex_syntax::parse::token::{Lit, Token, intern};
    use syntex_syntax::ptr::P;
    use syntex_syntax::tokenstream::{Delimited, SequenceRepetition, TokenTree};
    use syntex_syntax::util::move_map::MoveMap;
    use syntex_syntax::util::small_vector::SmallVector;

    struct Respanner;

    impl Respanner {
        fn fold_spanned<T>(&mut self, spanned: Spanned<T>) -> Spanned<T> {
            codemap::respan(self.new_span(spanned.span), spanned.node)
        }

        fn fold_lit(&mut self, l: Lit) -> Lit {
            // Give up on comparing literals inside of macros because there are
            // so many equivalent representations of the same literal; they are
            // tested elsewhere
            match l {
                Lit::Byte(_) => Lit::Byte(intern("")),
                Lit::Char(_) => Lit::Char(intern("")),
                Lit::Integer(_) => Lit::Integer(intern("")),
                Lit::Float(_) => Lit::Float(intern("")),
                Lit::Str_(_) => Lit::Str_(intern("")),
                Lit::ByteStr(_) => Lit::ByteStr(intern("")),
                _ => l,
            }
        }
    }

    impl Folder for Respanner {
        fn new_span(&mut self, _: Span) -> Span {
            syntex_pos::DUMMY_SP
        }

        fn fold_item_kind(&mut self, i: ItemKind) -> ItemKind {
            match i {
                ItemKind::Fn(decl, unsafety, constness, abi, generics, body) => {
                    let generics = self.fold_generics(generics);
                    let decl = self.fold_fn_decl(decl);
                    let body = self.fold_block(body);
                    // default fold_item_kind does not fold this span
                    let constness = self.fold_spanned(constness);
                    ItemKind::Fn(decl, unsafety, constness, abi, generics, body)
                }
                _ => fold::noop_fold_item_kind(i, self),
            }
        }

        fn fold_expr(&mut self, e: P<Expr>) -> P<Expr> {
            e.map(|e| {
                let folded = fold::noop_fold_expr(e, self);
                Expr {
                    node: match folded.node {
                        ExprKind::Lit(l) => {
                            // default fold_expr does not fold lits
                            ExprKind::Lit(l.map(|l| self.fold_spanned(l)))
                        }
                        ExprKind::Binary(op, lhs, rhs) => {
                            // default fold_expr does not fold the op span
                            ExprKind::Binary(self.fold_spanned(op),
                                             self.fold_expr(lhs),
                                             self.fold_expr(rhs))
                        }
                        ExprKind::AssignOp(op, lhs, rhs) => {
                            // default fold_expr does not fold the op span
                            ExprKind::AssignOp(self.fold_spanned(op),
                                               self.fold_expr(lhs),
                                               self.fold_expr(rhs))
                        }
                        other => other,
                    },
                    ..folded
                }
            })
        }

        fn fold_ty_param(&mut self, tp: TyParam) -> TyParam {
            TyParam {
                // default fold_ty_param does not fold the span
                span: self.new_span(tp.span),
                ..fold::noop_fold_ty_param(tp, self)
            }
        }

        fn fold_fn_decl(&mut self, decl: P<FnDecl>) -> P<FnDecl> {
            decl.map(|FnDecl { inputs, output, variadic }| {
                FnDecl {
                    inputs: inputs.move_map(|x| self.fold_arg(x)),
                    output: match output {
                        FunctionRetTy::Ty(ty) => FunctionRetTy::Ty(self.fold_ty(ty)),
                        // default fold_fn_decl does not fold this span
                        FunctionRetTy::Default(span) => FunctionRetTy::Default(self.new_span(span)),
                    },
                    variadic: variadic,
                }
            })
        }

        fn fold_field(&mut self, field: Field) -> Field {
            Field {
                ident: codemap::respan(// default fold_field does not fold this span
                                       self.new_span(field.ident.span),
                                       self.fold_ident(field.ident.node)),
                expr: self.fold_expr(field.expr),
                span: self.new_span(field.span),
            }
        }

        fn fold_trait_item(&mut self, i: TraitItem) -> SmallVector<TraitItem> {
            let noop = fold::noop_fold_trait_item(i, self).expect_one("");
            SmallVector::one(TraitItem {
                node: match noop.node {
                    TraitItemKind::Method(sig, body) => {
                        TraitItemKind::Method(MethodSig {
                                                  constness: self.fold_spanned(sig.constness),
                                                  ..sig
                                              },
                                              body)
                    }
                    node => node,
                },
                ..noop
            })
        }

        fn fold_impl_item(&mut self, i: ImplItem) -> SmallVector<ImplItem> {
            let noop = fold::noop_fold_impl_item(i, self).expect_one("");
            SmallVector::one(ImplItem {
                node: match noop.node {
                    ImplItemKind::Method(sig, body) => {
                        ImplItemKind::Method(MethodSig {
                                                 constness: self.fold_spanned(sig.constness),
                                                 ..sig
                                             },
                                             body)
                    }
                    node => node,
                },
                ..noop
            })
        }

        fn fold_attribute(&mut self, mut at: Attribute) -> Option<Attribute> {
            at.node.id.0 = 0;
            fold::noop_fold_attribute(at, self)
        }

        fn fold_meta_item(&mut self, meta_item: P<MetaItem>) -> P<MetaItem> {
            meta_item.map(|Spanned { node, span }| {
                Spanned {
                    node: match node {
                        MetaItemKind::Word(id) => MetaItemKind::Word(id),
                        MetaItemKind::List(id, mis) => {
                            MetaItemKind::List(id, mis.move_map(|e| self.fold_meta_list_item(e)))
                        }
                        // default fold_meta_item does not fold the value span
                        MetaItemKind::NameValue(id, lit) => {
                            MetaItemKind::NameValue(id, self.fold_spanned(lit))
                        }
                    },
                    span: self.new_span(span),
                }
            })
        }

        fn fold_meta_list_item(&mut self, list_item: NestedMetaItem) -> NestedMetaItem {
            Spanned {
                node: match list_item.node {
                    NestedMetaItemKind::MetaItem(mi) => {
                        NestedMetaItemKind::MetaItem(self.fold_meta_item(mi))
                    }
                    // default fold_meta_list_item does not fold the span
                    NestedMetaItemKind::Literal(lit) => {
                        NestedMetaItemKind::Literal(self.fold_spanned(lit))
                    }
                },
                span: self.new_span(list_item.span),
            }
        }

        fn fold_mac(&mut self, mac: Mac) -> Mac {
            fold::noop_fold_mac(mac, self)
        }

        fn fold_tt(&mut self, tt: &TokenTree) -> TokenTree {
            match *tt {
                TokenTree::Token(span, ref tok) => {
                    TokenTree::Token(self.new_span(span), self.fold_token(tok.clone()))
                }
                TokenTree::Delimited(span, ref delimed) => {
                    TokenTree::Delimited(self.new_span(span),
                                         Rc::new(Delimited {
                                             delim: delimed.delim,
                                             open_span: self.new_span(delimed.open_span),
                                             tts: self.fold_tts(&delimed.tts),
                                             close_span: self.new_span(delimed.close_span),
                                         }))
                }
                TokenTree::Sequence(span, ref seq) => {
                    TokenTree::Sequence(self.new_span(span),
                                        Rc::new(SequenceRepetition {
                                            tts: self.fold_tts(&seq.tts),
                                            separator: seq.separator
                                                .clone()
                                                .map(|tok| self.fold_token(tok)),
                                            ..**seq
                                        }))
                }
            }
        }

        fn fold_token(&mut self, t: Token) -> Token {
            match t {
                // default fold_token does not fold literals
                Token::Literal(lit, repr) => Token::Literal(self.fold_lit(lit), repr),
                Token::Ident(id) => Token::Ident(self.fold_ident(id)),
                Token::Lifetime(id) => Token::Lifetime(self.fold_ident(id)),
                Token::Interpolated(nt) => Token::Interpolated(self.fold_interpolated(nt)),
                Token::SubstNt(ident) => Token::SubstNt(self.fold_ident(ident)),
                Token::MatchNt(name, kind) => {
                    Token::MatchNt(self.fold_ident(name), self.fold_ident(kind))
                }
                _ => t,
            }
        }
    }

    Respanner.fold_crate(krate)
}
<|endoftext|>


<commit_before>/*
 * Cymbalum, Molecular Simulation in Rust
 * Copyright (C) 2015 Guillaume Fraux
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/
*/

//! Saving properties of the simulation to some stream (file or stdout mainly)

extern crate chemharp;

use std::io::prelude::*;
use std::io;
use std::fs::File;
use std::path::Path;

use ::units;
use ::universe::Universe;

use ::universe::chemharp::universe_to_frame;

/// The `Output` trait define the interface for all the quantities outputed by
/// the simulation during the run. An Output can be a text or a binary data
/// file, an image, a text log, …
pub trait Output {
    /// Function called once at the beggining of the simulation, which allow
    /// for some setup of the output if needed.
    fn setup(&mut self, _: &Universe) {}

    /// Write the output from the universe.
    fn write(&mut self, universe: &Universe);

    /// Function called once at the end of the simulation.
    fn finish(&mut self, _: &Universe) {}
}

/******************************************************************************/
/// The `TrajectoryOutput` allow to write the trajectory of the system to a
/// file, using any format supported by the
/// [Chemharp](http://chemharp.readthedocs.org/en/latest/formats.html) library.
pub struct TrajectoryOutput {
    file: chemharp::Trajectory,
}

impl TrajectoryOutput {
    /// Create a new `TrajectoryOutput` writing to `filename`. The file is
    /// replaced if it already exists.
    pub fn new<'a, S>(filename: S) -> Result<TrajectoryOutput, chemharp::Error> where S: Into<&'a str> {
        Ok(TrajectoryOutput{
            file: try!(chemharp::Trajectory::create(filename.into()))
        })
    }
}

impl Output for TrajectoryOutput {
    fn write(&mut self, universe: &Universe) {
        let frame = match universe_to_frame(universe) {
            Ok(val) => val,
            Err(err) => {
                error!("Error in Chemharp runtime while converting data: {}", err.message());
                panic!();
            }
        };

        match self.file.write(&frame) {
            Ok(()) => {},
            Err(err) => {
                error!("Error in Chemharp runtime while writing file: {}", err.message());
                panic!();
            }
        }
    }
}

/******************************************************************************/
/// The `CellOutput` write all the components of a cell to a file . The columns
/// in the file contains the following values: `A B C α β γ`.
pub struct CellOutput {
    file: File,
    path: String
}

impl CellOutput {
    /// Create a new `CellOutput` writing to `filename`. The file is replaced if
    /// it already exists.
    pub fn new<P: AsRef<Path>>(filename: P) -> Result<CellOutput, io::Error> {
        let path = String::from(filename.as_ref().to_str().unwrap());
        Ok(CellOutput{
            file: try!(File::create(filename)),
            path: path,
        })
    }
}

impl Output for CellOutput {
    fn setup(&mut self, _: &Universe) {
        if let Err(e) = writeln!(&mut self.file, "# Unit cell of the simulation") {
            // Do panic in early time
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
        if let Err(e) = writeln!(&mut self.file, "# Step A/Å B/Å C/Å α/deg β/deg γ/deg") {
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
    }

    fn write(&mut self, universe: &Universe) {
        let cell = universe.cell();
        if let Err(e) = writeln!(
            &mut self.file, "{} {} {} {} {} {} {}",
            universe.step(), cell.a(), cell.b(), cell.c(), cell.alpha(), cell.beta(), cell.gamma()
        ) {
            // Do not panic during the simulation
            error!("Could not write to file '{}': {:?}", self.path, e);
        }
    }
}

/******************************************************************************/
/// The `EnergyOutput` write the energy of the system to a text file, organized
/// as: `PotentialEnergy     KineticEnergy     TotalEnergy`.
pub struct EnergyOutput {
    file: File,
}

impl EnergyOutput {
    /// Create a new `EnergyOutput` writing to `filename`. The file is replaced
    /// if it already exists.
    pub fn new<P: AsRef<Path>>(filename: P) -> Result<EnergyOutput, io::Error> {
        Ok(EnergyOutput{
            file: try!(File::create(filename))
        })
    }
}

impl Output for EnergyOutput {
    fn setup(&mut self, _: &Universe) {
        writeln!(&mut self.file, "# Energy of the simulation (kJ/mol)").unwrap();
        writeln!(&mut self.file, "# Step Potential Kinetic Total").unwrap();
    }

    fn write(&mut self, universe: &Universe) {
        let potential = units::to(universe.potential_energy(), "kJ/mol").unwrap();
        let kinetic = units::to(universe.kinetic_energy(), "kJ/mol").unwrap();
        let total = units::to(universe.total_energy(), "kJ/mol").unwrap();
        writeln!(&mut self.file, "{} {} {} {}", universe.step(), potential, kinetic, total).unwrap();
    }
}

/******************************************************************************/
/// The `PropertiesOutput` write various physical properties of the universe to
/// a file. These properties are:
///
/// - volume of the unit cell;
/// - instant temperature;
/// - instant pressure;
pub struct PropertiesOutput {
    file: File,
    path: String,
}

impl PropertiesOutput {
    /// Create a new `PropertiesOutput` writing to `filename`. The file is replaced
    /// if it already exists.
    pub fn new<P: AsRef<Path>>(filename: P) -> Result<PropertiesOutput, io::Error> {
        let path = String::from(filename.as_ref().to_str().unwrap());
        Ok(PropertiesOutput{
            file: try!(File::create(filename)),
            path: path,
        })
    }
}

impl Output for PropertiesOutput {
    fn setup(&mut self, _: &Universe) {
        if let Err(e) = writeln!(&mut self.file, "# Physical properties of the simulation") {
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
        if let Err(e) = writeln!(&mut self.file, "# Step Volume/A^3 Temperature/K Pressure/bar") {
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
    }

    fn write(&mut self, universe: &Universe) {
        let V = units::to(universe.volume(), "A^3").unwrap();
        let T = units::to(universe.temperature(), "K").unwrap();
        let P = units::to(universe.pressure(), "bar").unwrap();
        if let Err(e) = writeln!(&mut self.file, "{} {} {} {}", universe.step(), V, T, P) {
            error!("Could not write to file '{}': {:?}", self.path, e);
        }
    }
}

#[cfg(test)]
mod tests {
    use std::io::prelude::*;
    use std::fs;

    use super::*;
    use ::universe::*;
    use ::types::*;
    use ::potentials::*;
    use ::units;

    fn testing_universe() -> Universe {
        let mut universe = Universe::from_cell(UnitCell::cubic(10.0));;

        universe.add_particle(Particle::new("F"));
        universe[0].set_position(Vector3D::new(0.0, 0.0, 0.0));

        universe.add_particle(Particle::new("F"));
        universe[1].set_position(Vector3D::new(1.3, 0.0, 0.0));

        universe.add_pair_interaction("F", "F",
            Harmonic{k: units::from(300.0, "kJ/mol/A^2").unwrap(), x0: units::from(1.2, "A").unwrap()});
        return universe;
    }

    fn check_file_content(filename: &str, content: &str) {
        let mut file = fs::File::open(filename).unwrap();
        let mut buffer = String::new();
        file.read_to_string(&mut buffer).unwrap();

        assert_eq!(buffer, content);
    }

    #[test]
    fn trajectory() {
        let filename = "testing-trajectory-output.xyz";
        let universe = testing_universe();
        {
            let mut out = TrajectoryOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "2\nWritten by Chemharp\nF 0 0 0\nF 1.3 0 0\n");
        fs::remove_file(filename).unwrap();
    }

    #[test]
    fn energy() {
        let filename = "testing-energy-output.dat";
        let universe = testing_universe();
        {
            let mut out = EnergyOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "# Energy of the simulation (kJ/mol)\n# Step Potential Kinetic Total\n0 1.5000000000000027 0 1.5000000000000027\n");
        fs::remove_file(filename).unwrap();
    }

    #[test]
    fn cell() {
        let filename = "testing-cell-output.dat";
        let universe = testing_universe();
        {
            let mut out = CellOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "# Unit cell of the simulation\n# Step A/Å B/Å C/Å α/deg β/deg γ/deg\n0 10 10 10 90 90 90\n");
        fs::remove_file(filename).unwrap();
    }

    #[test]
    fn properties() {
        let filename = "testing-properties-output.dat";
        let universe = testing_universe();
        {
            let mut out = PropertiesOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "# Physical properties of the simulation\n# Step Volume/A^3 Temperature/K Pressure/bar\n0 1000 0 -431.7400836223091\n");
        fs::remove_file(filename).unwrap();
    }
}
<commit_msg>Better errors handling in EnergyOutput<commit_after>/*
 * Cymbalum, Molecular Simulation in Rust
 * Copyright (C) 2015 Guillaume Fraux
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/
*/

//! Saving properties of the simulation to some stream (file or stdout mainly)

extern crate chemharp;

use std::io::prelude::*;
use std::io;
use std::fs::File;
use std::path::Path;

use ::units;
use ::universe::Universe;

use ::universe::chemharp::universe_to_frame;

/// The `Output` trait define the interface for all the quantities outputed by
/// the simulation during the run. An Output can be a text or a binary data
/// file, an image, a text log, …
pub trait Output {
    /// Function called once at the beggining of the simulation, which allow
    /// for some setup of the output if needed.
    fn setup(&mut self, _: &Universe) {}

    /// Write the output from the universe.
    fn write(&mut self, universe: &Universe);

    /// Function called once at the end of the simulation.
    fn finish(&mut self, _: &Universe) {}
}

/******************************************************************************/
/// The `TrajectoryOutput` allow to write the trajectory of the system to a
/// file, using any format supported by the
/// [Chemharp](http://chemharp.readthedocs.org/en/latest/formats.html) library.
pub struct TrajectoryOutput {
    file: chemharp::Trajectory,
}

impl TrajectoryOutput {
    /// Create a new `TrajectoryOutput` writing to `filename`. The file is
    /// replaced if it already exists.
    pub fn new<'a, S>(filename: S) -> Result<TrajectoryOutput, chemharp::Error> where S: Into<&'a str> {
        Ok(TrajectoryOutput{
            file: try!(chemharp::Trajectory::create(filename.into()))
        })
    }
}

impl Output for TrajectoryOutput {
    fn write(&mut self, universe: &Universe) {
        let frame = match universe_to_frame(universe) {
            Ok(val) => val,
            Err(err) => {
                error!("Error in Chemharp runtime while converting data: {}", err.message());
                panic!();
            }
        };

        match self.file.write(&frame) {
            Ok(()) => {},
            Err(err) => {
                error!("Error in Chemharp runtime while writing file: {}", err.message());
                panic!();
            }
        }
    }
}

/******************************************************************************/
/// The `CellOutput` write all the components of a cell to a file . The columns
/// in the file contains the following values: `A B C α β γ`.
pub struct CellOutput {
    file: File,
    path: String
}

impl CellOutput {
    /// Create a new `CellOutput` writing to `filename`. The file is replaced if
    /// it already exists.
    pub fn new<P: AsRef<Path>>(filename: P) -> Result<CellOutput, io::Error> {
        let path = String::from(filename.as_ref().to_str().unwrap());
        Ok(CellOutput{
            file: try!(File::create(filename)),
            path: path,
        })
    }
}

impl Output for CellOutput {
    fn setup(&mut self, _: &Universe) {
        if let Err(e) = writeln!(&mut self.file, "# Unit cell of the simulation") {
            // Do panic in early time
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
        if let Err(e) = writeln!(&mut self.file, "# Step A/Å B/Å C/Å α/deg β/deg γ/deg") {
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
    }

    fn write(&mut self, universe: &Universe) {
        let cell = universe.cell();
        if let Err(e) = writeln!(
            &mut self.file, "{} {} {} {} {} {} {}",
            universe.step(), cell.a(), cell.b(), cell.c(), cell.alpha(), cell.beta(), cell.gamma()
        ) {
            // Do not panic during the simulation
            error!("Could not write to file '{}': {:?}", self.path, e);
        }
    }
}

/******************************************************************************/
/// The `EnergyOutput` write the energy of the system to a text file, organized
/// as: `PotentialEnergy     KineticEnergy     TotalEnergy`.
pub struct EnergyOutput {
    file: File,
    path: String
}

impl EnergyOutput {
    /// Create a new `EnergyOutput` writing to `filename`. The file is replaced
    /// if it already exists.
    pub fn new<P: AsRef<Path>>(filename: P) -> Result<EnergyOutput, io::Error> {
        let path = String::from(filename.as_ref().to_str().unwrap());
        Ok(EnergyOutput{
            file: try!(File::create(filename)),
            path: path,
        })
    }
}

impl Output for EnergyOutput {
    fn setup(&mut self, _: &Universe) {
        if let Err(e) = writeln!(&mut self.file, "# Energy of the simulation (kJ/mol)") {
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
        if let Err(e) = writeln!(&mut self.file, "# Step Potential Kinetic Total") {
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
    }

    fn write(&mut self, universe: &Universe) {
        let potential = units::to(universe.potential_energy(), "kJ/mol").unwrap();
        let kinetic = units::to(universe.kinetic_energy(), "kJ/mol").unwrap();
        let total = units::to(universe.total_energy(), "kJ/mol").unwrap();
        if let Err(e) = writeln!(&mut self.file, "{} {} {} {}", universe.step(), potential, kinetic, total) {
            error!("Could not write to file '{}': {:?}", self.path, e);
        }
    }
}

/******************************************************************************/
/// The `PropertiesOutput` write various physical properties of the universe to
/// a file. These properties are:
///
/// - volume of the unit cell;
/// - instant temperature;
/// - instant pressure;
pub struct PropertiesOutput {
    file: File,
    path: String,
}

impl PropertiesOutput {
    /// Create a new `PropertiesOutput` writing to `filename`. The file is replaced
    /// if it already exists.
    pub fn new<P: AsRef<Path>>(filename: P) -> Result<PropertiesOutput, io::Error> {
        let path = String::from(filename.as_ref().to_str().unwrap());
        Ok(PropertiesOutput{
            file: try!(File::create(filename)),
            path: path,
        })
    }
}

impl Output for PropertiesOutput {
    fn setup(&mut self, _: &Universe) {
        if let Err(e) = writeln!(&mut self.file, "# Physical properties of the simulation") {
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
        if let Err(e) = writeln!(&mut self.file, "# Step Volume/A^3 Temperature/K Pressure/bar") {
            panic!("Could not write to file '{}': {:?}", self.path, e);
        }
    }

    fn write(&mut self, universe: &Universe) {
        let V = units::to(universe.volume(), "A^3").unwrap();
        let T = units::to(universe.temperature(), "K").unwrap();
        let P = units::to(universe.pressure(), "bar").unwrap();
        if let Err(e) = writeln!(&mut self.file, "{} {} {} {}", universe.step(), V, T, P) {
            error!("Could not write to file '{}': {:?}", self.path, e);
        }
    }
}

#[cfg(test)]
mod tests {
    use std::io::prelude::*;
    use std::fs;

    use super::*;
    use ::universe::*;
    use ::types::*;
    use ::potentials::*;
    use ::units;

    fn testing_universe() -> Universe {
        let mut universe = Universe::from_cell(UnitCell::cubic(10.0));;

        universe.add_particle(Particle::new("F"));
        universe[0].set_position(Vector3D::new(0.0, 0.0, 0.0));

        universe.add_particle(Particle::new("F"));
        universe[1].set_position(Vector3D::new(1.3, 0.0, 0.0));

        universe.add_pair_interaction("F", "F",
            Harmonic{k: units::from(300.0, "kJ/mol/A^2").unwrap(), x0: units::from(1.2, "A").unwrap()});
        return universe;
    }

    fn check_file_content(filename: &str, content: &str) {
        let mut file = fs::File::open(filename).unwrap();
        let mut buffer = String::new();
        file.read_to_string(&mut buffer).unwrap();

        assert_eq!(buffer, content);
    }

    #[test]
    fn trajectory() {
        let filename = "testing-trajectory-output.xyz";
        let universe = testing_universe();
        {
            let mut out = TrajectoryOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "2\nWritten by Chemharp\nF 0 0 0\nF 1.3 0 0\n");
        fs::remove_file(filename).unwrap();
    }

    #[test]
    fn energy() {
        let filename = "testing-energy-output.dat";
        let universe = testing_universe();
        {
            let mut out = EnergyOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "# Energy of the simulation (kJ/mol)\n# Step Potential Kinetic Total\n0 1.5000000000000027 0 1.5000000000000027\n");
        fs::remove_file(filename).unwrap();
    }

    #[test]
    fn cell() {
        let filename = "testing-cell-output.dat";
        let universe = testing_universe();
        {
            let mut out = CellOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "# Unit cell of the simulation\n# Step A/Å B/Å C/Å α/deg β/deg γ/deg\n0 10 10 10 90 90 90\n");
        fs::remove_file(filename).unwrap();
    }

    #[test]
    fn properties() {
        let filename = "testing-properties-output.dat";
        let universe = testing_universe();
        {
            let mut out = PropertiesOutput::new(filename).unwrap();
            out.setup(&universe);
            out.write(&universe);
            out.finish(&universe);
        }

        check_file_content(filename, "# Physical properties of the simulation\n# Step Volume/A^3 Temperature/K Pressure/bar\n0 1000 0 -431.7400836223091\n");
        fs::remove_file(filename).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::sync::{Mutex, MutexGuard};
use std::hash::{Hash, SipHasher, Hasher};
use std::fmt;

use storage::{Key, RefKey, Value};
use storage::engine::{Engine, Modify};
use proto::mvccpb::{MetaLock, MetaItem};
use util::codec::bytes;
use super::meta::Meta;
use super::codec;
use super::{Error, Result};

const SHARD_LOCK_SIZE: usize = 256;

fn shard_index(key: RefKey) -> usize {
    let mut s = SipHasher::new();
    key.hash(&mut s);
    (s.finish() as usize) % SHARD_LOCK_SIZE
}

#[derive(Debug)]
pub struct MvccStore {
    engine: Box<Engine>,
    shard_locks: Vec<Mutex<()>>,
}

impl MvccStore {
    pub fn new(engine: Box<Engine>) -> MvccStore {
        let mut locks = Vec::<_>::with_capacity(SHARD_LOCK_SIZE);
        for _ in 0..SHARD_LOCK_SIZE {
            locks.push(Mutex::new(()));
        }
        MvccStore {
            engine: engine,
            shard_locks: locks,
        }
    }

    pub fn start_row_transaction(&self, key: Key) -> Result<RowTxn> {
        let index = shard_index(&key);
        let guard = self.shard_locks[index].lock().unwrap();
        RowTxn::new(self.engine.as_ref(), key, guard)
    }

    pub fn get(&self, key: Key, ts: u64) -> Result<Option<Value>> {
        let txn = try!(self.start_row_transaction(key));
        txn.get(ts)
    }

    pub fn prewrite(&self, key: Key, args: Prewrite, primary: Key, ts: u64) -> Result<()> {
        let mut txn = try!(self.start_row_transaction(key));
        try!(txn.prewrite(args, primary, ts));
        try!(txn.write());
        Ok(())
    }

    pub fn commit(&self, key: Key, start_ts: u64, commit_ts: u64) -> Result<()> {
        let mut txn = try!(self.start_row_transaction(key));
        try!(txn.commit(start_ts, commit_ts));
        try!(txn.write());
        Ok(())
    }

    pub fn rollback(&self, key: Key, start_ts: u64) -> Result<()> {
        let mut txn = try!(self.start_row_transaction(key));
        try!(txn.rollback(start_ts));
        try!(txn.write());
        Ok(())
    }
}

#[allow(dead_code)]
pub struct RowTxn<'a> {
    engine: &'a Engine,
    row_key: Key,
    row_lock: MutexGuard<'a, ()>,
    meta_key: Key,
    meta: Meta,
    writes: Vec<Modify>,
}

impl<'a> fmt::Debug for RowTxn<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "txn: {:?} - {:?}", self.row_key, self.engine)
    }
}

impl<'a> RowTxn<'a> {
    fn new(engine: &'a Engine, row_key: Key, lock: MutexGuard<'a, ()>) -> Result<RowTxn<'a>> {
        let meta_key = bytes::encode_bytes(&row_key);
        let meta = match try!(engine.get(&meta_key)) {
            Some(x) => try!(Meta::parse(&x)),
            None => Meta::new(),
        };
        Ok(RowTxn {
            engine: engine,
            row_key: row_key,
            row_lock: lock,
            meta_key: meta_key,
            meta: meta,
            writes: vec![],
        })
    }

    pub fn write(&mut self) -> Result<()> {
        let batch = self.writes.drain(..).collect();
        try!(self.engine.write(batch));
        Ok(())
    }

    pub fn get(&self, ts: u64) -> Result<Option<Value>> {
        if let Some(lock) = self.meta.get_lock() {
            if lock.get_start_ts() <= ts {
                return Err(Error::KeyIsLocked {
                    primary: lock.get_primary_key().to_vec(),
                    ts: lock.get_start_ts(),
                });
            }
        }
        match self.meta.iter_items().find(|x| x.get_commit_ts() <= ts) {
            Some(x) => {
                let data_key = codec::encode_key(&self.row_key, x.get_start_ts());
                Ok(try!(self.engine.get(&data_key)))
            }
            None => Ok(None),
        }
    }

    pub fn prewrite(&mut self, args: Prewrite, primary: Key, ts: u64) -> Result<()> {
        if let Some(lock) = self.meta.get_lock() {
            return Err(Error::KeyIsLocked {
                primary: lock.get_primary_key().to_vec(),
                ts: lock.get_start_ts(),
            });
        }
        if self.meta.iter_items().any(|x| x.get_commit_ts() >= ts) {
            return Err(Error::WriteConflict);
        }

        let mut lock = MetaLock::new();
        lock.set_primary_key(primary);
        lock.set_start_ts(ts);
        lock.set_read_only(match args {
            Prewrite::Lock => true,
            _ => false,
        });
        self.meta.set_lock(lock);
        let modify = Modify::Put((self.meta_key.clone(), self.meta_bytes()));
        self.writes.push(modify);

        if let Prewrite::Put(value) = args {
            let value_key = codec::encode_key(&self.row_key, ts);
            self.writes.push(Modify::Put((value_key, value)));
        }
        Ok(())
    }

    pub fn commit(&mut self, start_ts: u64, commit_ts: u64) -> Result<()> {
        match self.meta.get_lock() {
            Some(lock) if lock.get_start_ts() == start_ts => {}
            _ => return Err(Error::TxnAbortedWhileWorking),
        }
        let meta_item = {
            let lock = self.meta.get_lock().unwrap();
            let mut item = MetaItem::new();
            item.set_start_ts(lock.get_start_ts());
            item.set_commit_ts(commit_ts);
            item
        };
        self.meta.push_item(meta_item);
        self.meta.clear_lock();
        let modify = Modify::Put((self.meta_key.clone(), self.meta_bytes()));
        self.writes.push(modify);
        Ok(())
    }

    #[allow(dead_code)]
    pub fn rollback(&mut self, start_ts: u64) -> Result<()> {
        match self.meta.get_lock() {
            Some(lock) if lock.get_start_ts() == start_ts => {
                if !lock.get_read_only() {
                    let value_key = codec::encode_key(&self.row_key, lock.get_start_ts());
                    self.writes.push(Modify::Delete(value_key));
                }
            }
            _ => return Err(Error::TxnAbortedWhileWorking),
        }
        self.meta.clear_lock();
        let modify = Modify::Put((self.meta_key.clone(), self.meta_bytes()));
        self.writes.push(modify);
        Ok(())
    }

    fn meta_bytes(&self) -> Value {
        let mut data = vec![];
        self.meta.write_to(&mut data);
        data
    }
}

#[derive(Debug)]
pub enum Prewrite {
    Put(Value),
    Delete,
    Lock,
}

#[cfg(test)]
mod tests {
    use super::{MvccStore, Prewrite};
    use storage::engine::{self, Engine, Dsn};

    #[test]
    fn test_mvcc_store_read() {
        let engine = engine::new_engine(Dsn::Memory).unwrap();
        let store = MvccStore::new(engine);

        must_get_none(&store, b"x", 1);

        must_prewrite_put(&store, b"x", b"x5", b"x", 5);
        must_get_none(&store, b"x", 3);
        must_get_err(&store, b"x", 7);

        must_commit(&store, b"x", 5, 10);
        must_get_none(&store, b"x", 3);
        must_get_none(&store, b"x", 7);
        must_get(&store, b"x", 13, b"x5");

        must_prewrite_delete(&store, b"x", b"x", 15);
        must_commit(&store, b"x", 15, 20);
        must_get_none(&store, b"x", 3);
        must_get_none(&store, b"x", 7);
        must_get(&store, b"x", 13, b"x5");
        must_get(&store, b"x", 17, b"x5");
        must_get_none(&store, b"x", 23);
    }

    #[test]
    fn test_mvcc_store_write() {
        let engine = engine::new_engine(Dsn::Memory).unwrap();
        let store = MvccStore::new(engine);

        must_prewrite_put(&store, b"x", b"x5", b"x", 5);
        must_prewrite_lock_err(&store, b"x", b"x", 6);
        must_commit(&store, b"x", 5, 10);
        must_prewrite_lock_err(&store, b"x", b"x", 6);
        must_prewrite_lock(&store, b"x", b"x", 12);
    }

    #[test]
    fn test_mvcc_store_rollback() {
        let engine = engine::new_engine(Dsn::Memory).unwrap();
        let store = MvccStore::new(engine);

        must_prewrite_put(&store, b"x", b"x5", b"x", 5);
        must_rollback(&store, b"x", 5);
        must_prewrite_lock(&store, b"x", b"x", 6);
    }

    fn must_get(store: &MvccStore, key: &[u8], ts: u64, expect: &[u8]) {
        assert_eq!(store.get(key.to_vec(), ts).unwrap().unwrap(), expect);
    }

    fn must_get_none(store: &MvccStore, key: &[u8], ts: u64) {
        assert!(store.get(key.to_vec(), ts).unwrap().is_none());
    }

    fn must_get_err(store: &MvccStore, key: &[u8], ts: u64) {
        assert!(store.get(key.to_vec(), ts).is_err());
    }

    fn must_prewrite_put(store: &MvccStore, key: &[u8], value: &[u8], pk: &[u8], ts: u64) {
        store.prewrite(key.to_vec(), Prewrite::Put(value.to_vec()), pk.to_vec(), ts).unwrap();
    }

    fn must_prewrite_delete(store: &MvccStore, key: &[u8], pk: &[u8], ts: u64) {
        store.prewrite(key.to_vec(), Prewrite::Delete, pk.to_vec(), ts).unwrap();
    }

    fn must_prewrite_lock(store: &MvccStore, key: &[u8], pk: &[u8], ts: u64) {
        store.prewrite(key.to_vec(), Prewrite::Lock, pk.to_vec(), ts).unwrap();
    }

    fn must_prewrite_lock_err(store: &MvccStore, key: &[u8], pk: &[u8], ts: u64) {
        assert!(store.prewrite(key.to_vec(), Prewrite::Lock, pk.to_vec(), ts).is_err());
    }

    fn must_commit(store: &MvccStore, key: &[u8], start_ts: u64, commit_ts: u64) {
        store.commit(key.to_vec(), start_ts, commit_ts).unwrap();
    }

    fn must_rollback(store: &MvccStore, key: &[u8], start_ts: u64) {
        store.rollback(key.to_vec(), start_ts).unwrap();
    }
}
<commit_msg>storage: update commit logic<commit_after>use std::sync::{Mutex, MutexGuard};
use std::hash::{Hash, SipHasher, Hasher};
use std::fmt;

use storage::{Key, RefKey, Value};
use storage::engine::{Engine, Modify};
use proto::mvccpb::{MetaLock, MetaItem};
use util::codec::bytes;
use super::meta::Meta;
use super::codec;
use super::{Error, Result};

const SHARD_LOCK_SIZE: usize = 256;

fn shard_index(key: RefKey) -> usize {
    let mut s = SipHasher::new();
    key.hash(&mut s);
    (s.finish() as usize) % SHARD_LOCK_SIZE
}

#[derive(Debug)]
pub struct MvccStore {
    engine: Box<Engine>,
    shard_locks: Vec<Mutex<()>>,
}

impl MvccStore {
    pub fn new(engine: Box<Engine>) -> MvccStore {
        let mut locks = Vec::<_>::with_capacity(SHARD_LOCK_SIZE);
        for _ in 0..SHARD_LOCK_SIZE {
            locks.push(Mutex::new(()));
        }
        MvccStore {
            engine: engine,
            shard_locks: locks,
        }
    }

    pub fn start_row_transaction(&self, key: Key) -> Result<RowTxn> {
        let index = shard_index(&key);
        let guard = self.shard_locks[index].lock().unwrap();
        RowTxn::new(self.engine.as_ref(), key, guard)
    }

    pub fn get(&self, key: Key, ts: u64) -> Result<Option<Value>> {
        let txn = try!(self.start_row_transaction(key));
        txn.get(ts)
    }

    pub fn prewrite(&self, key: Key, args: Prewrite, primary: Key, ts: u64) -> Result<()> {
        let mut txn = try!(self.start_row_transaction(key));
        try!(txn.prewrite(args, primary, ts));
        try!(txn.write());
        Ok(())
    }

    pub fn commit(&self, key: Key, start_ts: u64, commit_ts: u64) -> Result<()> {
        let mut txn = try!(self.start_row_transaction(key));
        try!(txn.commit(start_ts, commit_ts));
        try!(txn.write());
        Ok(())
    }

    pub fn rollback(&self, key: Key, start_ts: u64) -> Result<()> {
        let mut txn = try!(self.start_row_transaction(key));
        try!(txn.rollback(start_ts));
        try!(txn.write());
        Ok(())
    }
}

#[allow(dead_code)]
pub struct RowTxn<'a> {
    engine: &'a Engine,
    row_key: Key,
    row_lock: MutexGuard<'a, ()>,
    meta_key: Key,
    meta: Meta,
    writes: Vec<Modify>,
}

impl<'a> fmt::Debug for RowTxn<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "txn: {:?} - {:?}", self.row_key, self.engine)
    }
}

impl<'a> RowTxn<'a> {
    fn new(engine: &'a Engine, row_key: Key, lock: MutexGuard<'a, ()>) -> Result<RowTxn<'a>> {
        let meta_key = bytes::encode_bytes(&row_key);
        let meta = match try!(engine.get(&meta_key)) {
            Some(x) => try!(Meta::parse(&x)),
            None => Meta::new(),
        };
        Ok(RowTxn {
            engine: engine,
            row_key: row_key,
            row_lock: lock,
            meta_key: meta_key,
            meta: meta,
            writes: vec![],
        })
    }

    pub fn write(&mut self) -> Result<()> {
        let batch = self.writes.drain(..).collect();
        try!(self.engine.write(batch));
        Ok(())
    }

    pub fn get(&self, ts: u64) -> Result<Option<Value>> {
        if let Some(lock) = self.meta.get_lock() {
            if lock.get_start_ts() <= ts {
                return Err(Error::KeyIsLocked {
                    primary: lock.get_primary_key().to_vec(),
                    ts: lock.get_start_ts(),
                });
            }
        }
        match self.meta.iter_items().find(|x| x.get_commit_ts() <= ts) {
            Some(x) => {
                let data_key = codec::encode_key(&self.row_key, x.get_start_ts());
                Ok(try!(self.engine.get(&data_key)))
            }
            None => Ok(None),
        }
    }

    pub fn prewrite(&mut self, args: Prewrite, primary: Key, ts: u64) -> Result<()> {
        if let Some(lock) = self.meta.get_lock() {
            return Err(Error::KeyIsLocked {
                primary: lock.get_primary_key().to_vec(),
                ts: lock.get_start_ts(),
            });
        }
        if self.meta.iter_items().any(|x| x.get_commit_ts() >= ts) {
            return Err(Error::WriteConflict);
        }

        let mut lock = MetaLock::new();
        lock.set_primary_key(primary);
        lock.set_start_ts(ts);
        lock.set_read_only(match args {
            Prewrite::Lock => true,
            _ => false,
        });
        self.meta.set_lock(lock);
        let modify = Modify::Put((self.meta_key.clone(), self.meta_bytes()));
        self.writes.push(modify);

        if let Prewrite::Put(value) = args {
            let value_key = codec::encode_key(&self.row_key, ts);
            self.writes.push(Modify::Put((value_key, value)));
        }
        Ok(())
    }

    pub fn commit(&mut self, start_ts: u64, commit_ts: u64) -> Result<()> {
        let read_only = match self.meta.get_lock() {
            Some(lock) if lock.get_start_ts() == start_ts => lock.get_read_only(),
            _ => return Err(Error::TxnAbortedWhileWorking),
        };
        if !read_only {
            let mut item = MetaItem::new();
            item.set_start_ts(start_ts);
            item.set_commit_ts(commit_ts);
            self.meta.push_item(item);
        }
        self.meta.clear_lock();
        let modify = Modify::Put((self.meta_key.clone(), self.meta_bytes()));
        self.writes.push(modify);
        Ok(())
    }

    #[allow(dead_code)]
    pub fn rollback(&mut self, start_ts: u64) -> Result<()> {
        match self.meta.get_lock() {
            Some(lock) if lock.get_start_ts() == start_ts => {
                if !lock.get_read_only() {
                    let value_key = codec::encode_key(&self.row_key, lock.get_start_ts());
                    self.writes.push(Modify::Delete(value_key));
                }
            }
            _ => return Err(Error::TxnAbortedWhileWorking),
        }
        self.meta.clear_lock();
        let modify = Modify::Put((self.meta_key.clone(), self.meta_bytes()));
        self.writes.push(modify);
        Ok(())
    }

    fn meta_bytes(&self) -> Value {
        let mut data = vec![];
        self.meta.write_to(&mut data);
        data
    }
}

#[derive(Debug)]
pub enum Prewrite {
    Put(Value),
    Delete,
    Lock,
}

#[cfg(test)]
mod tests {
    use super::{MvccStore, Prewrite};
    use storage::engine::{self, Engine, Dsn};

    #[test]
    fn test_mvcc_store_read() {
        let engine = engine::new_engine(Dsn::Memory).unwrap();
        let store = MvccStore::new(engine);

        must_get_none(&store, b"x", 1);

        must_prewrite_put(&store, b"x", b"x5", b"x", 5);
        must_get_none(&store, b"x", 3);
        must_get_err(&store, b"x", 7);

        must_commit(&store, b"x", 5, 10);
        must_get_none(&store, b"x", 3);
        must_get_none(&store, b"x", 7);
        must_get(&store, b"x", 13, b"x5");

        must_prewrite_delete(&store, b"x", b"x", 15);
        must_commit(&store, b"x", 15, 20);
        must_get_none(&store, b"x", 3);
        must_get_none(&store, b"x", 7);
        must_get(&store, b"x", 13, b"x5");
        must_get(&store, b"x", 17, b"x5");
        must_get_none(&store, b"x", 23);
    }

    #[test]
    fn test_mvcc_store_write() {
        let engine = engine::new_engine(Dsn::Memory).unwrap();
        let store = MvccStore::new(engine);

        must_prewrite_put(&store, b"x", b"x5", b"x", 5);
        must_prewrite_lock_err(&store, b"x", b"x", 6);
        must_commit(&store, b"x", 5, 10);
        must_prewrite_lock_err(&store, b"x", b"x", 6);
        must_prewrite_lock(&store, b"x", b"x", 12);
    }

    #[test]
    fn test_mvcc_store_rollback() {
        let engine = engine::new_engine(Dsn::Memory).unwrap();
        let store = MvccStore::new(engine);

        must_prewrite_put(&store, b"x", b"x5", b"x", 5);
        must_rollback(&store, b"x", 5);
        must_prewrite_lock(&store, b"x", b"x", 6);
    }

    fn must_get(store: &MvccStore, key: &[u8], ts: u64, expect: &[u8]) {
        assert_eq!(store.get(key.to_vec(), ts).unwrap().unwrap(), expect);
    }

    fn must_get_none(store: &MvccStore, key: &[u8], ts: u64) {
        assert!(store.get(key.to_vec(), ts).unwrap().is_none());
    }

    fn must_get_err(store: &MvccStore, key: &[u8], ts: u64) {
        assert!(store.get(key.to_vec(), ts).is_err());
    }

    fn must_prewrite_put(store: &MvccStore, key: &[u8], value: &[u8], pk: &[u8], ts: u64) {
        store.prewrite(key.to_vec(), Prewrite::Put(value.to_vec()), pk.to_vec(), ts).unwrap();
    }

    fn must_prewrite_delete(store: &MvccStore, key: &[u8], pk: &[u8], ts: u64) {
        store.prewrite(key.to_vec(), Prewrite::Delete, pk.to_vec(), ts).unwrap();
    }

    fn must_prewrite_lock(store: &MvccStore, key: &[u8], pk: &[u8], ts: u64) {
        store.prewrite(key.to_vec(), Prewrite::Lock, pk.to_vec(), ts).unwrap();
    }

    fn must_prewrite_lock_err(store: &MvccStore, key: &[u8], pk: &[u8], ts: u64) {
        assert!(store.prewrite(key.to_vec(), Prewrite::Lock, pk.to_vec(), ts).is_err());
    }

    fn must_commit(store: &MvccStore, key: &[u8], start_ts: u64, commit_ts: u64) {
        store.commit(key.to_vec(), start_ts, commit_ts).unwrap();
    }

    fn must_rollback(store: &MvccStore, key: &[u8], start_ts: u64) {
        store.rollback(key.to_vec(), start_ts).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Condvar, Mutex};
use std::usize;

use sleep::Sleep;

/// We define various kinds of latches, which are all a primitive signaling
/// mechanism. A latch starts as false. Eventually someone calls `set()` and
/// it becomes true. You can test if it has been set by calling `probe()`.
///
/// Some kinds of latches, but not all, support a `wait()` operation
/// that will wait until the latch is set, blocking efficiently. That
/// is not part of the trait since it is not possibly to do with all
/// latches.
///
/// The intention is that `set()` is called once, but `probe()` may be
/// called any number of times. Once `probe()` returns true, the memory
/// effects that occurred before `set()` become visible.
///
/// It'd probably be better to refactor the API into two paired types,
/// but that's a bit of work, and this is not a public API.
///
/// ## Memory ordering
///
/// Latches need to guarantee two things:
///
/// - Once `probe()` returns true, all memory effects from the `set()`
///   are visible (in other words, the set should synchronize-with
///   the probe).
/// - Once `set()` occurs, the next `probe()` *will* observe it.  This
///   typically requires a seq-cst ordering. See [the "tickle-then-get-sleepy" scenario in the sleep
///   README](/src/sleep/README.md#tickle-then-get-sleepy) for details.
pub(super) trait Latch: LatchProbe {
    /// Set the latch, signalling others.
    fn set(&self);
}

pub(super) trait LatchProbe {
    /// Test if the latch is set.
    fn probe(&self) -> bool;
}

/// Spin latches are the simplest, most efficient kind, but they do
/// not support a `wait()` operation. They just have a boolean flag
/// that becomes true when `set()` is called.
pub(super) struct SpinLatch {
    b: AtomicBool,
}

impl SpinLatch {
    #[inline]
    pub(super) fn new() -> SpinLatch {
        SpinLatch {
            b: AtomicBool::new(false),
        }
    }
}

impl LatchProbe for SpinLatch {
    #[inline]
    fn probe(&self) -> bool {
        self.b.load(Ordering::SeqCst)
    }
}

impl Latch for SpinLatch {
    #[inline]
    fn set(&self) {
        self.b.store(true, Ordering::SeqCst);
    }
}

/// A Latch starts as false and eventually becomes true. You can block
/// until it becomes true.
pub(super) struct LockLatch {
    m: Mutex<bool>,
    v: Condvar,
}

impl LockLatch {
    #[inline]
    pub(super) fn new() -> LockLatch {
        LockLatch {
            m: Mutex::new(false),
            v: Condvar::new(),
        }
    }

    /// Resets this lock latch so it can be reused again.
    pub(super) fn reset(&self) {
        *self.m.lock().unwrap() = false;
    }

    /// Block until latch is set.
    pub(super) fn wait(&self) {
        let mut guard = self.m.lock().unwrap();
        while !*guard {
            guard = self.v.wait(guard).unwrap();
        }
    }
}

impl LatchProbe for LockLatch {
    #[inline]
    fn probe(&self) -> bool {
        // Not particularly efficient, but we don't really use this operation
        let guard = self.m.lock().unwrap();
        *guard
    }
}

impl Latch for LockLatch {
    #[inline]
    fn set(&self) {
        let mut guard = self.m.lock().unwrap();
        *guard = true;
        self.v.notify_all();
    }
}

/// Counting latches are used to implement scopes. They track a
/// counter. Unlike other latches, calling `set()` does not
/// necessarily make the latch be considered `set()`; instead, it just
/// decrements the counter. The latch is only "set" (in the sense that
/// `probe()` returns true) once the counter reaches zero.
#[derive(Debug)]
pub(super) struct CountLatch {
    counter: AtomicUsize,
}

impl CountLatch {
    #[inline]
    pub(super) fn new() -> CountLatch {
        CountLatch {
            counter: AtomicUsize::new(1),
        }
    }

    #[inline]
    pub(super) fn increment(&self) {
        debug_assert!(!self.probe());
        self.counter.fetch_add(1, Ordering::Relaxed);
    }
}

impl LatchProbe for CountLatch {
    #[inline]
    fn probe(&self) -> bool {
        // Need to acquire any memory reads before latch was set:
        self.counter.load(Ordering::SeqCst) == 0
    }
}

impl Latch for CountLatch {
    /// Set the latch to true, releasing all threads who are waiting.
    #[inline]
    fn set(&self) {
        self.counter.fetch_sub(1, Ordering::SeqCst);
    }
}

/// A tickling latch wraps another latch type, and will also awaken a thread
/// pool when it is set.  This is useful for jobs injected between thread pools,
/// so the source pool can continue processing its own work while waiting.
pub(super) struct TickleLatch<'a, L: Latch> {
    inner: L,
    sleep: &'a Sleep,
}

impl<'a, L: Latch> TickleLatch<'a, L> {
    #[inline]
    pub(super) fn new(latch: L, sleep: &'a Sleep) -> Self {
        TickleLatch {
            inner: latch,
            sleep,
        }
    }
}

impl<'a, L: Latch> LatchProbe for TickleLatch<'a, L> {
    #[inline]
    fn probe(&self) -> bool {
        self.inner.probe()
    }
}

impl<'a, L: Latch> Latch for TickleLatch<'a, L> {
    #[inline]
    fn set(&self) {
        self.inner.set();
        self.sleep.tickle(usize::MAX);
    }
}

impl<'a, L> LatchProbe for &'a L where L: LatchProbe {
    fn probe(&self) -> bool {
        unimplemented!()
    }
}

impl<'a, L> Latch for &'a L where L: Latch {
    fn set(&self) {
        unimplemented!()
    }
}
<commit_msg>Actually impl `LatchProbe` and `Latch` for &L where L: Latch.<commit_after>use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Condvar, Mutex};
use std::usize;

use sleep::Sleep;

/// We define various kinds of latches, which are all a primitive signaling
/// mechanism. A latch starts as false. Eventually someone calls `set()` and
/// it becomes true. You can test if it has been set by calling `probe()`.
///
/// Some kinds of latches, but not all, support a `wait()` operation
/// that will wait until the latch is set, blocking efficiently. That
/// is not part of the trait since it is not possibly to do with all
/// latches.
///
/// The intention is that `set()` is called once, but `probe()` may be
/// called any number of times. Once `probe()` returns true, the memory
/// effects that occurred before `set()` become visible.
///
/// It'd probably be better to refactor the API into two paired types,
/// but that's a bit of work, and this is not a public API.
///
/// ## Memory ordering
///
/// Latches need to guarantee two things:
///
/// - Once `probe()` returns true, all memory effects from the `set()`
///   are visible (in other words, the set should synchronize-with
///   the probe).
/// - Once `set()` occurs, the next `probe()` *will* observe it.  This
///   typically requires a seq-cst ordering. See [the "tickle-then-get-sleepy" scenario in the sleep
///   README](/src/sleep/README.md#tickle-then-get-sleepy) for details.
pub(super) trait Latch: LatchProbe {
    /// Set the latch, signalling others.
    fn set(&self);
}

pub(super) trait LatchProbe {
    /// Test if the latch is set.
    fn probe(&self) -> bool;
}

/// Spin latches are the simplest, most efficient kind, but they do
/// not support a `wait()` operation. They just have a boolean flag
/// that becomes true when `set()` is called.
pub(super) struct SpinLatch {
    b: AtomicBool,
}

impl SpinLatch {
    #[inline]
    pub(super) fn new() -> SpinLatch {
        SpinLatch {
            b: AtomicBool::new(false),
        }
    }
}

impl LatchProbe for SpinLatch {
    #[inline]
    fn probe(&self) -> bool {
        self.b.load(Ordering::SeqCst)
    }
}

impl Latch for SpinLatch {
    #[inline]
    fn set(&self) {
        self.b.store(true, Ordering::SeqCst);
    }
}

/// A Latch starts as false and eventually becomes true. You can block
/// until it becomes true.
pub(super) struct LockLatch {
    m: Mutex<bool>,
    v: Condvar,
}

impl LockLatch {
    #[inline]
    pub(super) fn new() -> LockLatch {
        LockLatch {
            m: Mutex::new(false),
            v: Condvar::new(),
        }
    }

    /// Resets this lock latch so it can be reused again.
    pub(super) fn reset(&self) {
        *self.m.lock().unwrap() = false;
    }

    /// Block until latch is set.
    pub(super) fn wait(&self) {
        let mut guard = self.m.lock().unwrap();
        while !*guard {
            guard = self.v.wait(guard).unwrap();
        }
    }
}

impl LatchProbe for LockLatch {
    #[inline]
    fn probe(&self) -> bool {
        // Not particularly efficient, but we don't really use this operation
        let guard = self.m.lock().unwrap();
        *guard
    }
}

impl Latch for LockLatch {
    #[inline]
    fn set(&self) {
        let mut guard = self.m.lock().unwrap();
        *guard = true;
        self.v.notify_all();
    }
}

/// Counting latches are used to implement scopes. They track a
/// counter. Unlike other latches, calling `set()` does not
/// necessarily make the latch be considered `set()`; instead, it just
/// decrements the counter. The latch is only "set" (in the sense that
/// `probe()` returns true) once the counter reaches zero.
#[derive(Debug)]
pub(super) struct CountLatch {
    counter: AtomicUsize,
}

impl CountLatch {
    #[inline]
    pub(super) fn new() -> CountLatch {
        CountLatch {
            counter: AtomicUsize::new(1),
        }
    }

    #[inline]
    pub(super) fn increment(&self) {
        debug_assert!(!self.probe());
        self.counter.fetch_add(1, Ordering::Relaxed);
    }
}

impl LatchProbe for CountLatch {
    #[inline]
    fn probe(&self) -> bool {
        // Need to acquire any memory reads before latch was set:
        self.counter.load(Ordering::SeqCst) == 0
    }
}

impl Latch for CountLatch {
    /// Set the latch to true, releasing all threads who are waiting.
    #[inline]
    fn set(&self) {
        self.counter.fetch_sub(1, Ordering::SeqCst);
    }
}

/// A tickling latch wraps another latch type, and will also awaken a thread
/// pool when it is set.  This is useful for jobs injected between thread pools,
/// so the source pool can continue processing its own work while waiting.
pub(super) struct TickleLatch<'a, L: Latch> {
    inner: L,
    sleep: &'a Sleep,
}

impl<'a, L: Latch> TickleLatch<'a, L> {
    #[inline]
    pub(super) fn new(latch: L, sleep: &'a Sleep) -> Self {
        TickleLatch {
            inner: latch,
            sleep,
        }
    }
}

impl<'a, L: Latch> LatchProbe for TickleLatch<'a, L> {
    #[inline]
    fn probe(&self) -> bool {
        self.inner.probe()
    }
}

impl<'a, L: Latch> Latch for TickleLatch<'a, L> {
    #[inline]
    fn set(&self) {
        self.inner.set();
        self.sleep.tickle(usize::MAX);
    }
}

impl<'a, L> LatchProbe for &'a L where L: LatchProbe {
    fn probe(&self) -> bool {
        L::probe(&self)
    }
}

impl<'a, L> Latch for &'a L where L: Latch {
    fn set(&self) {
        L::set(&self);
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate const_cstr;
extern crate dlopen;
extern crate rand;
extern crate test_server;

use dlopen::raw::Library;
use std::{env, str};
use std::u16::MAX;
use test_server::http::StatusCode;

#[cfg(target_os = "windows")]
const LIB_PATH: &str = "../../target/debug/ctehxk2.dll";
#[cfg(target_os = "linux")]
const LIB_PATH: &str = "../../target/debug/libctehxk2.so";
#[cfg(target_os = "macos")]
const LIB_PATH: &str = "../../target/debug/libctehxk2.dylib";

#[test]
fn has_ct_api_functions() {
    let lib = Library::open(LIB_PATH).expect("Could not open library");

    let init: unsafe extern "system" fn(u16, u16) -> i8 =
        unsafe { lib.symbol_cstr(const_cstr!("CT_init").as_cstr()) }.unwrap();

    let data: unsafe extern "system" fn(
        u16,
        *mut u8,
        *mut u8,
        u16,
        *const u8,
        *mut u16,
        *mut u8,
    ) -> i8 = unsafe { lib.symbol_cstr(const_cstr!("CT_data").as_cstr()) }.unwrap();

    let close: unsafe extern "system" fn(u16) -> i8 =
        unsafe { lib.symbol_cstr(const_cstr!("CT_close").as_cstr()) }.unwrap();

    let ctn = rand::random::<u16>();
    let pn = rand::random::<u16>();
    let mut dad = rand::random::<u8>();
    let mut sad = rand::random::<u8>();

    let commands: [u8; 1] = [rand::random::<u8>(); 1];
    let commands_ptr: *const u8 = &commands[0];
    let lenc: u16 = commands.len() as u16;

    let mut response: [u8; MAX as usize] = [rand::random::<u8>(); MAX as usize];
    let response_ptr: *mut u8 = &mut response[0];
    let mut lenr: u16 = response.len() as u16;

    let server = test_server::serve(Some("127.0.0.1:65432"));
    env::set_var("K2_BASE_URL", "http://127.0.0.1:65432");

    server.reply().status(StatusCode::OK).body("0");
    assert_eq!(0, unsafe { init(ctn, pn) });

    server
        .reply()
        .status(StatusCode::OK)
        .body("{\"dad\":1,\"sad\":1,\"lenr\":5,\"response\":\"AQIDBAU=\",\"responseCode\":0}");
    assert_eq!(0, unsafe {
        data(
            ctn,
            &mut dad,
            &mut sad,
            lenc,
            commands_ptr,
            &mut lenr,
            response_ptr,
        )
    });

    server.reply().status(StatusCode::OK).body("0");
    assert_eq!(0, unsafe { close(ctn) });
}
<commit_msg>cargo fmt<commit_after>#[macro_use]
extern crate const_cstr;
extern crate dlopen;
extern crate rand;
extern crate test_server;

use dlopen::raw::Library;
use std::u16::MAX;
use std::{env, str};
use test_server::http::StatusCode;

#[cfg(target_os = "windows")]
const LIB_PATH: &str = "../../target/debug/ctehxk2.dll";
#[cfg(target_os = "linux")]
const LIB_PATH: &str = "../../target/debug/libctehxk2.so";
#[cfg(target_os = "macos")]
const LIB_PATH: &str = "../../target/debug/libctehxk2.dylib";

#[test]
fn has_ct_api_functions() {
    let lib = Library::open(LIB_PATH).expect("Could not open library");

    let init: unsafe extern "system" fn(u16, u16) -> i8 =
        unsafe { lib.symbol_cstr(const_cstr!("CT_init").as_cstr()) }.unwrap();

    let data: unsafe extern "system" fn(
        u16,
        *mut u8,
        *mut u8,
        u16,
        *const u8,
        *mut u16,
        *mut u8,
    ) -> i8 = unsafe { lib.symbol_cstr(const_cstr!("CT_data").as_cstr()) }.unwrap();

    let close: unsafe extern "system" fn(u16) -> i8 =
        unsafe { lib.symbol_cstr(const_cstr!("CT_close").as_cstr()) }.unwrap();

    let ctn = rand::random::<u16>();
    let pn = rand::random::<u16>();
    let mut dad = rand::random::<u8>();
    let mut sad = rand::random::<u8>();

    let commands: [u8; 1] = [rand::random::<u8>(); 1];
    let commands_ptr: *const u8 = &commands[0];
    let lenc: u16 = commands.len() as u16;

    let mut response: [u8; MAX as usize] = [rand::random::<u8>(); MAX as usize];
    let response_ptr: *mut u8 = &mut response[0];
    let mut lenr: u16 = response.len() as u16;

    let server = test_server::serve(Some("127.0.0.1:65432"));
    env::set_var("K2_BASE_URL", "http://127.0.0.1:65432");

    server.reply().status(StatusCode::OK).body("0");
    assert_eq!(0, unsafe { init(ctn, pn) });

    server
        .reply()
        .status(StatusCode::OK)
        .body("{\"dad\":1,\"sad\":1,\"lenr\":5,\"response\":\"AQIDBAU=\",\"responseCode\":0}");
    assert_eq!(0, unsafe {
        data(
            ctn,
            &mut dad,
            &mut sad,
            lenc,
            commands_ptr,
            &mut lenr,
            response_ptr,
        )
    });

    server.reply().status(StatusCode::OK).body("0");
    assert_eq!(0, unsafe { close(ctn) });
}
<|endoftext|>