<commit_before>#![crate_type = "dylib"]
#![license = "MIT/ASL2"]

#![feature(plugin_registrar, quote)]

extern crate syntax;
extern crate rustc;

use std::gc::Gc;

use syntax::ast::{
    Ident,
    MetaItem,
    Item,
    Expr,
    MutMutable,
    LitNil,
};
use syntax::ast;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, ItemDecorator};
use syntax::ext::build::AstBuilder;
use syntax::ext::deriving::generic::{
    EnumMatching,
    FieldInfo,
    MethodDef,
    Named,
    StaticFields,
    StaticStruct,
    StaticEnum,
    Struct,
    Substructure,
    TraitDef,
    Unnamed,
    combine_substructure,
};
use syntax::ext::deriving::generic::ty::{
    Borrowed,
    LifetimeBounds,
    Literal,
    Path,
    Ptr,
    Self,
    Tuple,
    borrowed_explicit_self,
};
use syntax::parse::token;

use rustc::plugin::Registry;

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(
        token::intern("deriving_serializable"),
        ItemDecorator(expand_deriving_serializable));

    /*
    reg.register_syntax_extension(
        token::intern("deriving_deserializable"),
        ItemDecorator(expand_deriving_deserializable));
        */
}

fn expand_deriving_serializable(cx: &mut ExtCtxt,
                                sp: Span,
                                mitem: Gc<MetaItem>,
                                item: Gc<Item>,
                                mut push: |Gc<ast::Item>|) {

    let inline = cx.meta_word(sp, token::InternedString::new("inline"));
    let attrs = vec!(cx.attribute(sp, inline));

    let trait_def = TraitDef {
        span: sp,
        attributes: vec!(),
        path: Path::new_(vec!("serde2", "ser", "Serialize"), None,
                         vec!(box Literal(Path::new_local("__S")),
                              box Literal(Path::new_local("__R"))), true),
        additional_bounds: Vec::new(),
        generics: LifetimeBounds {
            lifetimes: Vec::new(),
            bounds: vec!(("__S", None, vec!(Path::new_(
                            vec!("serde2", "ser", "VisitorState"), None,
                            vec!(box Literal(Path::new_local("__R"))), true))),
                         ("__R", None, vec!()))
        },
        methods: vec!(
            MethodDef {
                name: "serialize",
                generics: LifetimeBounds::empty(),
                explicit_self: borrowed_explicit_self(),
                args: vec!(Ptr(box Literal(Path::new_local("__S")),
                            Borrowed(None, MutMutable))),
                ret_ty: Literal(
                    Path::new_local("__R")
                ),
                attributes: attrs,
                combine_substructure: combine_substructure(|a, b, c| {
                    serializable_substructure(a, b, c)
                }),
            })
    };

    trait_def.expand(cx, mitem, item, push)
}

fn serializable_substructure(cx: &ExtCtxt, span: Span, substr: &Substructure) -> Gc<Expr> {
    let serializer = substr.nonself_args[0];

    match *substr.fields {
        Struct(ref fields) => {
            if fields.is_empty() {
                serialize_tuple_struct(cx)
            } else {
                serialize_struct(cx, span, serializer, substr.type_ident, fields)
            }
        }

        EnumMatching(_idx, variant, ref fields) => {
            serialize_enum(cx, span, serializer, substr.type_ident, variant, fields)
        }

        _ => cx.bug("expected Struct or EnumMatching in deriving_serializable")
    }
}

fn serialize_tuple_struct(cx: &ExtCtxt) -> Gc<Expr> {
    // unit structs have no fields and need to return `Ok()`
    quote_expr!(cx, Ok(()))
}

fn serialize_struct(cx: &ExtCtxt,
                    span: Span,
                    serializer: Gc<Expr>,
                    type_ident: Ident,
                    fields: &Vec<FieldInfo>) -> Gc<Expr> {

    let type_name = cx.expr_str(
        span,
        token::get_ident(type_ident));
    let len = fields.len();

    let arms: Vec<ast::Arm> = fields.iter()
        .enumerate()
        .map(|(i, &FieldInfo { name, span, .. })| {
            let first = if i == 0 {
                quote_expr!(cx, true)
            } else {
                quote_expr!(cx, false)
            };

            let name = name.unwrap();
            let name_expr = cx.expr_str(span, token::get_ident(name));

            quote_arm!(cx,
                $i => {
                    self.state += 1;
                    Some(s.visit_map_elt($first, $name_expr, &self.value.$name))
                }
            )
        })
        .collect();

    quote_expr!(cx, {
        struct Visitor<'a> {
            state: uint,
            value: &'a $type_ident,
        }

        impl<
            'a,
            S: ::serde2::VisitorState<R>,
            R
        > ::serde2::Visitor<S, R> for Visitor<'a> {
            #[inline]
            fn visit(&mut self, s: &mut S) -> Option<R> {
                match self.state {
                    $arms
                    _ => None,
                }
            }

            #[inline]
            fn size_hint(&self) -> (uint, Option<uint>) {
                let size = $len - self.state;
                (size, Some(size))
            }
        }

        $serializer.visit_named_map($type_name, Visitor {
            value: self,
            state: 0,
        })
    })
}

fn serialize_enum(cx: &ExtCtxt,
                  span: Span,
                  serializer: Gc<Expr>,
                  type_ident: Ident,
                  variant: &ast::Variant,
                  fields: &Vec<FieldInfo>) -> Gc<Expr> {
    let type_name = cx.expr_str(
        span,
        token::get_ident(type_ident)
    );
    let variant_name = cx.expr_str(
        span,
        token::get_ident(variant.node.name)
    );
    let len = fields.len();

    let stmts: Vec<Gc<ast::Stmt>> = fields.iter()
        .map(|&FieldInfo { self_, span, .. }| {
            quote_stmt!(
                cx,
                try!($serializer.serialize_enum_elt(&$self_))
            )
        })
        .collect();

    quote_expr!(cx, {
        try!($serializer.serialize_enum_start($type_name, $variant_name, $len));
        $stmts
        $serializer.serialize_enum_end()
    })
}

/*
pub fn expand_deriving_deserializable(cx: &mut ExtCtxt,
                                      span: Span,
                                      mitem: Gc<MetaItem>,
                                      item: Gc<Item>,
                                      push: |Gc<Item>|) {
    let trait_def = TraitDef {
        span: span,
        attributes: Vec::new(),
        path: Path::new_(vec!("serde2", "de", "Deserializable"), None,
                         vec!(box Literal(Path::new_local("__D")),
                              box Literal(Path::new_local("__E"))), true),
        additional_bounds: Vec::new(),
        generics: LifetimeBounds {
            lifetimes: Vec::new(),
            bounds: vec!(("__D", None, vec!(Path::new_(
                            vec!("serde2", "de", "Deserializer"), None,
                            vec!(box Literal(Path::new_local("__E"))), true))),
                         ("__E", None, vec!()))
        },
        methods: vec!(
            MethodDef {
                name: "deserialize_token",
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(
                    Ptr(
                        box Literal(Path::new_local("__D")),
                        Borrowed(None, MutMutable)
                    ),
                    Literal(Path::new(vec!("serde2", "de", "Token"))),
                ),
                ret_ty: Literal(
                    Path::new_(
                        vec!("std", "result", "Result"),
                        None,
                        vec!(
                            box Self,
                            box Literal(Path::new_local("__E"))
                        ),
                        true
                    )
                ),
                attributes: Vec::new(),
                combine_substructure: combine_substructure(|a, b, c| {
                    deserializable_substructure(a, b, c)
                }),
            })
    };

    trait_def.expand(cx, mitem, item, push)
}

fn deserializable_substructure(cx: &mut ExtCtxt, span: Span,
                               substr: &Substructure) -> Gc<Expr> {
    let deserializer = substr.nonself_args[0];
    let token = substr.nonself_args[1];

    match *substr.fields {
        StaticStruct(_, ref fields) => {
            deserialize_struct(
                cx,
                span,
                substr.type_ident,
                fields,
                deserializer,
                token)
        }
        StaticEnum(_, ref fields) => {
            deserialize_enum(
                cx,
                span,
                substr.type_ident,
                fields.as_slice(),
                deserializer,
                token)
        }
        _ => cx.bug("expected StaticEnum or StaticStruct in deriving(Deserializable)")
    }
}

fn deserialize_struct(
    cx: &ExtCtxt,
    span: Span,
    type_ident: Ident,
    fields: &StaticFields,
    deserializer: Gc<ast::Expr>,
    token: Gc<ast::Expr>
) -> Gc<ast::Expr> {
    let struct_block = deserialize_struct_from_struct(
        cx,
        span,
        type_ident,
        fields,
        deserializer
    );

    let map_block = deserialize_struct_from_map(
        cx,
        span,
        type_ident,
        fields,
        deserializer
    );

    quote_expr!(
        cx,
        match $token {
            ::serde2::de::StructStart(_, _) => $struct_block,
            ::serde2::de::MapStart(_) => $map_block,
            token => {
                let expected_tokens = [
                    ::serde2::de::StructStartKind,
                    ::serde2::de::MapStartKind,
                ];
                Err($deserializer.syntax_error(token, expected_tokens))
            }
        }
    )
}

fn deserialize_struct_from_struct(
    cx: &ExtCtxt,
    span: Span,
    type_ident: Ident,
    fields: &StaticFields,
    deserializer: Gc<ast::Expr>
) -> Gc<ast::Expr> {
    let expect_struct_field = cx.ident_of("expect_struct_field");

    let call = deserializable_static_fields(
        cx,
        span,
        type_ident,
        fields,
        |cx, span, name| {
            let name = cx.expr_str(span, name);
            quote_expr!(
                cx,
                try!($deserializer.expect_struct_field($name))
            )
        }
    );

    quote_expr!(cx, {
        let result = $call;
        try!($deserializer.expect_struct_end());
        Ok(result)
    })
}

fn deserialize_struct_from_map(
    cx: &ExtCtxt,
    span: Span,
    type_ident: Ident,
    fields: &StaticFields,
    deserializer: Gc<ast::Expr>
) -> Gc<ast::Expr> {
    let fields = match *fields {
        Unnamed(_) => fail!(),
        Named(ref fields) => fields.as_slice(),
    };

    // Declare each field.
    let let_fields: Vec<Gc<ast::Stmt>> = fields.iter()
        .map(|&(name, span)| {
            quote_stmt!(cx, let mut $name = None)
        })
        .collect();

    // Declare key arms.
    let key_arms: Vec<ast::Arm> = fields.iter()
        .map(|&(name, span)| {
            let s = cx.expr_str(span, token::get_ident(name));
            quote_arm!(cx,
                $s => {
                    $name = Some(
                        try!(::serde2::de::Deserializable::deserialize($deserializer))
                    );
                    continue;
                })
        })
        .collect();

    let extract_fields: Vec<Gc<ast::Stmt>> = fields.iter()
        .map(|&(name, span)| {
            let name_str = cx.expr_str(span, token::get_ident(name));
            quote_stmt!(cx,
                let $name = match $name {
                    Some($name) => $name,
                    None => try!($deserializer.missing_field($name_str)),
                };
            )
        })
        .collect();

    let result = cx.expr_struct_ident(
        span,
        type_ident,
        fields.iter()
            .map(|&(name, span)| {
                cx.field_imm(span, name, cx.expr_ident(span, name))
            })
            .collect()
    );

    quote_expr!(cx, {
        $let_fields

        loop {
            let token = match try!($deserializer.expect_token()) {
                ::serde2::de::End => { break; }
                token => token,
            };

            {
                let key = match token {
                    ::serde2::de::Str(s) => s,
                    ::serde2::de::String(ref s) => s.as_slice(),
                    token => {
                        let expected_tokens = [
                            ::serde2::de::StrKind,
                            ::serde2::de::StringKind,
                        ];
                        return Err($deserializer.syntax_error(token, expected_tokens));
                    }
                };

                match key {
                    $key_arms
                    _ => { }
                }
            }

            try!($deserializer.ignore_field(token))
        }

        $extract_fields
        Ok($result)
    })
}

fn deserialize_enum(
    cx: &ExtCtxt,
    span: Span,
    type_ident: Ident,
    fields: &[(Ident, Span, StaticFields)],
    deserializer: Gc<ast::Expr>,
    token: Gc<ast::Expr>
) -> Gc<ast::Expr> {
    let type_name = cx.expr_str(span, token::get_ident(type_ident));

    let variants = fields.iter()
        .map(|&(name, span, _)| {
            cx.expr_str(span, token::get_ident(name))
        })
        .collect();

    let variants = cx.expr_vec(span, variants);

    let arms: Vec<ast::Arm> = fields.iter()
        .enumerate()
        .map(|(i, &(name, span, ref parts))| {
            let call = deserializable_static_fields(
                cx,
                span,
                name,
                parts,
                |cx, span, _| {
                    quote_expr!(cx, try!($deserializer.expect_enum_elt()))
                }
            );

            quote_arm!(cx, $i => $call,)
        })
        .collect();

    quote_expr!(cx, {
        let i = try!($deserializer.expect_enum_start($token, $type_name, $variants));

        let result = match i {
            $arms
            _ => { unreachable!() }
        };

        try!($deserializer.expect_enum_end());

        Ok(result)
    })
}

/// Create a deserializer for a single enum variant/struct:
/// - `outer_pat_ident` is the name of this enum variant/struct
/// - `getarg` should retrieve the `uint`-th field with name `&str`.
fn deserializable_static_fields(
    cx: &ExtCtxt,
    span: Span,
    outer_pat_ident: Ident,
    fields: &StaticFields,
    getarg: |&ExtCtxt, Span, token::InternedString| -> Gc<Expr>
) -> Gc<Expr> {
    match *fields {
        Unnamed(ref fields) => {
            if fields.is_empty() {
                cx.expr_ident(span, outer_pat_ident)
            } else {
                let fields = fields.iter().enumerate().map(|(i, &span)| {
                    getarg(
                        cx,
                        span,
                        token::intern_and_get_ident(format!("_field{}", i).as_slice())
                    )
                }).collect();

                cx.expr_call_ident(span, outer_pat_ident, fields)
            }
        }
        Named(ref fields) => {
            // use the field's span to get nicer error messages.
            let fields = fields.iter().map(|&(name, span)| {
                let arg = getarg(
                    cx,
                    span,
                    token::get_ident(name)
                );
                cx.field_imm(span, name, arg)
            }).collect();

            cx.expr_struct_ident(span, outer_pat_ident, fields)
        }
    }
}
*/
<commit_msg>Fixed serde2<commit_after>#![crate_type = "dylib"]
#![license = "MIT/ASL2"]

#![feature(plugin_registrar, quote)]

extern crate syntax;
extern crate rustc;

use syntax::ast::{
    Ident,
    MetaItem,
    Item,
    Expr,
    MutMutable,
    LitNil,
};
use syntax::ast;
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, ItemDecorator};
use syntax::ext::build::AstBuilder;
use syntax::ext::deriving::generic::{
    EnumMatching,
    FieldInfo,
    MethodDef,
    Named,
    StaticFields,
    StaticStruct,
    StaticEnum,
    Struct,
    Substructure,
    TraitDef,
    Unnamed,
    combine_substructure,
};
use syntax::ext::deriving::generic::ty::{
    Borrowed,
    LifetimeBounds,
    Literal,
    Path,
    Ptr,
    Self,
    Tuple,
    borrowed_explicit_self,
};
use syntax::parse::token;
use syntax::ptr::P;

use rustc::plugin::Registry;

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(
        token::intern("deriving_serializable"),
        ItemDecorator(box expand_deriving_serializable));

    /*
    reg.register_syntax_extension(
        token::intern("deriving_deserializable"),
        ItemDecorator(box expand_deriving_deserializable));
        */
}

fn expand_deriving_serializable(cx: &mut ExtCtxt,
                                sp: Span,
                                mitem: &MetaItem,
                                item: &Item,
                                mut push: |P<ast::Item>|) {

    let inline = cx.meta_word(sp, token::InternedString::new("inline"));
    let attrs = vec!(cx.attribute(sp, inline));

    let trait_def = TraitDef {
        span: sp,
        attributes: vec!(),
        path: Path::new_(vec!("serde2", "ser", "Serialize"), None,
                         vec!(box Literal(Path::new_local("__S")),
                              box Literal(Path::new_local("__R"))), true),
        additional_bounds: Vec::new(),
        generics: LifetimeBounds {
            lifetimes: Vec::new(),
            bounds: vec!(("__S", None, vec!(Path::new_(
                            vec!("serde2", "ser", "VisitorState"), None,
                            vec!(box Literal(Path::new_local("__R"))), true))),
                         ("__R", None, vec!()))
        },
        methods: vec!(
            MethodDef {
                name: "serialize",
                generics: LifetimeBounds::empty(),
                explicit_self: borrowed_explicit_self(),
                args: vec!(Ptr(box Literal(Path::new_local("__S")),
                            Borrowed(None, MutMutable))),
                ret_ty: Literal(
                    Path::new_local("__R")
                ),
                attributes: attrs,
                combine_substructure: combine_substructure(|a, b, c| {
                    serializable_substructure(a, b, c)
                }),
            })
    };

    trait_def.expand(cx, mitem, item, push)
}

fn serializable_substructure(cx: &ExtCtxt, span: Span, substr: &Substructure) -> P<Expr> {
    let serializer = substr.nonself_args[0].clone();

    match *substr.fields {
        Struct(ref fields) => {
            if fields.is_empty() {
                serialize_tuple_struct(cx)
            } else {
                serialize_struct(cx, span, serializer, substr.type_ident, fields)
            }
        }

        EnumMatching(_idx, variant, ref fields) => {
            serialize_enum(cx, span, serializer, substr.type_ident, variant, fields)
        }

        _ => cx.bug("expected Struct or EnumMatching in deriving_serializable")
    }
}

fn serialize_tuple_struct(cx: &ExtCtxt) -> P<Expr> {
    // unit structs have no fields and need to return `Ok()`
    quote_expr!(cx, Ok(()))
}

fn serialize_struct(cx: &ExtCtxt,
                    span: Span,
                    serializer: P<Expr>,
                    type_ident: Ident,
                    fields: &Vec<FieldInfo>) -> P<Expr> {

    let type_name = cx.expr_str(
        span,
        token::get_ident(type_ident));
    let len = fields.len();

    let arms: Vec<ast::Arm> = fields.iter()
        .enumerate()
        .map(|(i, &FieldInfo { name, span, .. })| {
            let first = if i == 0 {
                quote_expr!(cx, true)
            } else {
                quote_expr!(cx, false)
            };

            let name = name.unwrap();
            let name_expr = cx.expr_str(span, token::get_ident(name));

            quote_arm!(cx,
                $i => {
                    self.state += 1;
                    Some(s.visit_map_elt($first, $name_expr, &self.value.$name))
                }
            )
        })
        .collect();

    quote_expr!(cx, {
        struct Visitor<'a> {
            state: uint,
            value: &'a $type_ident,
        }

        impl<
            'a,
            S: ::serde2::VisitorState<R>,
            R
        > ::serde2::Visitor<S, R> for Visitor<'a> {
            #[inline]
            fn visit(&mut self, s: &mut S) -> Option<R> {
                match self.state {
                    $arms
                    _ => None,
                }
            }

            #[inline]
            fn size_hint(&self) -> (uint, Option<uint>) {
                let size = $len - self.state;
                (size, Some(size))
            }
        }

        $serializer.visit_named_map($type_name, Visitor {
            value: self,
            state: 0,
        })
    })
}

fn serialize_enum(cx: &ExtCtxt,
                  span: Span,
                  serializer: P<Expr>,
                  type_ident: Ident,
                  variant: &ast::Variant,
                  fields: &Vec<FieldInfo>) -> P<Expr> {
    let type_name = cx.expr_str(
        span,
        token::get_ident(type_ident)
    );
    let variant_name = cx.expr_str(
        span,
        token::get_ident(variant.node.name)
    );
    let len = fields.len();

    let stmts: Vec<P<ast::Stmt>> = fields.iter()
        .map(|&FieldInfo { ref self_, span, .. }| {
            quote_stmt!(
                cx,
                try!($serializer.serialize_enum_elt(&$self_))
            )
        })
        .collect();

    quote_expr!(cx, {
        try!($serializer.serialize_enum_start($type_name, $variant_name, $len));
        $stmts
        $serializer.serialize_enum_end()
    })
}

/*
pub fn expand_deriving_deserializable(cx: &mut ExtCtxt,
                                      span: Span,
                                      mitem: &MetaItem,
                                      item: &Item,
                                      push: |P<Item>|) {
    let trait_def = TraitDef {
        span: span,
        attributes: Vec::new(),
        path: Path::new_(vec!("serde2", "de", "Deserializable"), None,
                         vec!(box Literal(Path::new_local("__D")),
                              box Literal(Path::new_local("__E"))), true),
        additional_bounds: Vec::new(),
        generics: LifetimeBounds {
            lifetimes: Vec::new(),
            bounds: vec!(("__D", None, vec!(Path::new_(
                            vec!("serde2", "de", "Deserializer"), None,
                            vec!(box Literal(Path::new_local("__E"))), true))),
                         ("__E", None, vec!()))
        },
        methods: vec!(
            MethodDef {
                name: "deserialize_token",
                generics: LifetimeBounds::empty(),
                explicit_self: None,
                args: vec!(
                    Ptr(
                        box Literal(Path::new_local("__D")),
                        Borrowed(None, MutMutable)
                    ),
                    Literal(Path::new(vec!("serde2", "de", "Token"))),
                ),
                ret_ty: Literal(
                    Path::new_(
                        vec!("std", "result", "Result"),
                        None,
                        vec!(
                            box Self,
                            box Literal(Path::new_local("__E"))
                        ),
                        true
                    )
                ),
                attributes: Vec::new(),
                combine_substructure: combine_substructure(|a, b, c| {
                    deserializable_substructure(a, b, c)
                }),
            })
    };

    trait_def.expand(cx, mitem, item, push)
}

fn deserializable_substructure(cx: &mut ExtCtxt, span: Span,
                               substr: &Substructure) -> P<Expr> {
    let deserializer = substr.nonself_args[0];
    let token = substr.nonself_args[1];

    match *substr.fields {
        StaticStruct(_, ref fields) => {
            deserialize_struct(
                cx,
                span,
                substr.type_ident,
                fields,
                deserializer,
                token)
        }
        StaticEnum(_, ref fields) => {
            deserialize_enum(
                cx,
                span,
                substr.type_ident,
                fields.as_slice(),
                deserializer,
                token)
        }
        _ => cx.bug("expected StaticEnum or StaticStruct in deriving(Deserializable)")
    }
}

fn deserialize_struct(
    cx: &ExtCtxt,
    span: Span,
    type_ident: Ident,
    fields: &StaticFields,
    deserializer: P<ast::Expr>,
    token: P<ast::Expr>
) -> P<ast::Expr> {
    let struct_block = deserialize_struct_from_struct(
        cx,
        span,
        type_ident,
        fields,
        deserializer
    );

    let map_block = deserialize_struct_from_map(
        cx,
        span,
        type_ident,
        fields,
        deserializer
    );

    quote_expr!(
        cx,
        match $token {
            ::serde2::de::StructStart(_, _) => $struct_block,
            ::serde2::de::MapStart(_) => $map_block,
            token => {
                let expected_tokens = [
                    ::serde2::de::StructStartKind,
                    ::serde2::de::MapStartKind,
                ];
                Err($deserializer.syntax_error(token, expected_tokens))
            }
        }
    )
}

fn deserialize_struct_from_struct(
    cx: &ExtCtxt,
    span: Span,
    type_ident: Ident,
    fields: &StaticFields,
    deserializer: P<ast::Expr>
) -> P<ast::Expr> {
    let expect_struct_field = cx.ident_of("expect_struct_field");

    let call = deserializable_static_fields(
        cx,
        span,
        type_ident,
        fields,
        |cx, span, name| {
            let name = cx.expr_str(span, name);
            quote_expr!(
                cx,
                try!($deserializer.expect_struct_field($name))
            )
        }
    );

    quote_expr!(cx, {
        let result = $call;
        try!($deserializer.expect_struct_end());
        Ok(result)
    })
}

fn deserialize_struct_from_map(
    cx: &ExtCtxt,
    span: Span,
    type_ident: Ident,
    fields: &StaticFields,
    deserializer: P<ast::Expr>
) -> P<ast::Expr> {
    let fields = match *fields {
        Unnamed(_) => fail!(),
        Named(ref fields) => fields.as_slice(),
    };

    // Declare each field.
    let let_fields: Vec<P<ast::Stmt>> = fields.iter()
        .map(|&(name, span)| {
            quote_stmt!(cx, let mut $name = None)
        })
        .collect();

    // Declare key arms.
    let key_arms: Vec<ast::Arm> = fields.iter()
        .map(|&(name, span)| {
            let s = cx.expr_str(span, token::get_ident(name));
            quote_arm!(cx,
                $s => {
                    $name = Some(
                        try!(::serde2::de::Deserializable::deserialize($deserializer))
                    );
                    continue;
                })
        })
        .collect();

    let extract_fields: Vec<P<ast::Stmt>> = fields.iter()
        .map(|&(name, span)| {
            let name_str = cx.expr_str(span, token::get_ident(name));
            quote_stmt!(cx,
                let $name = match $name {
                    Some($name) => $name,
                    None => try!($deserializer.missing_field($name_str)),
                };
            )
        })
        .collect();

    let result = cx.expr_struct_ident(
        span,
        type_ident,
        fields.iter()
            .map(|&(name, span)| {
                cx.field_imm(span, name, cx.expr_ident(span, name))
            })
            .collect()
    );

    quote_expr!(cx, {
        $let_fields

        loop {
            let token = match try!($deserializer.expect_token()) {
                ::serde2::de::End => { break; }
                token => token,
            };

            {
                let key = match token {
                    ::serde2::de::Str(s) => s,
                    ::serde2::de::String(ref s) => s.as_slice(),
                    token => {
                        let expected_tokens = [
                            ::serde2::de::StrKind,
                            ::serde2::de::StringKind,
                        ];
                        return Err($deserializer.syntax_error(token, expected_tokens));
                    }
                };

                match key {
                    $key_arms
                    _ => { }
                }
            }

            try!($deserializer.ignore_field(token))
        }

        $extract_fields
        Ok($result)
    })
}

fn deserialize_enum(
    cx: &ExtCtxt,
    span: Span,
    type_ident: Ident,
    fields: &[(Ident, Span, StaticFields)],
    deserializer: P<ast::Expr>,
    token: P<ast::Expr>
) -> P<ast::Expr> {
    let type_name = cx.expr_str(span, token::get_ident(type_ident));

    let variants = fields.iter()
        .map(|&(name, span, _)| {
            cx.expr_str(span, token::get_ident(name))
        })
        .collect();

    let variants = cx.expr_vec(span, variants);

    let arms: Vec<ast::Arm> = fields.iter()
        .enumerate()
        .map(|(i, &(name, span, ref parts))| {
            let call = deserializable_static_fields(
                cx,
                span,
                name,
                parts,
                |cx, span, _| {
                    quote_expr!(cx, try!($deserializer.expect_enum_elt()))
                }
            );

            quote_arm!(cx, $i => $call,)
        })
        .collect();

    quote_expr!(cx, {
        let i = try!($deserializer.expect_enum_start($token, $type_name, $variants));

        let result = match i {
            $arms
            _ => { unreachable!() }
        };

        try!($deserializer.expect_enum_end());

        Ok(result)
    })
}

/// Create a deserializer for a single enum variant/struct:
/// - `outer_pat_ident` is the name of this enum variant/struct
/// - `getarg` should retrieve the `uint`-th field with name `&str`.
fn deserializable_static_fields(
    cx: &ExtCtxt,
    span: Span,
    outer_pat_ident: Ident,
    fields: &StaticFields,
    getarg: |&ExtCtxt, Span, token::InternedString| -> P<Expr>
) -> P<Expr> {
    match *fields {
        Unnamed(ref fields) => {
            if fields.is_empty() {
                cx.expr_ident(span, outer_pat_ident)
            } else {
                let fields = fields.iter().enumerate().map(|(i, &span)| {
                    getarg(
                        cx,
                        span,
                        token::intern_and_get_ident(format!("_field{}", i).as_slice())
                    )
                }).collect();

                cx.expr_call_ident(span, outer_pat_ident, fields)
            }
        }
        Named(ref fields) => {
            // use the field's span to get nicer error messages.
            let fields = fields.iter().map(|&(name, span)| {
                let arg = getarg(
                    cx,
                    span,
                    token::get_ident(name)
                );
                cx.field_imm(span, name, arg)
            }).collect();

            cx.expr_struct_ident(span, outer_pat_ident, fields)
        }
    }
}
*/
<|endoftext|>


<commit_before>//! Representation of migrations

use crate::backend::Backend;
use crate::connection::{BoxableConnection, Connection};
use crate::deserialize::{FromSql, FromSqlRow};
use crate::expression::AsExpression;
use crate::result::QueryResult;
use crate::serialize::ToSql;
use crate::sql_types::Text;
use std::borrow::Cow;
use std::error::Error;
use std::fmt::Display;

/// A specialized result type representing the result of
/// a migration operation
pub type Result<T> = std::result::Result<T, Box<dyn Error + Send + Sync>>;

/// A migration version identifier
///
/// This is used by the migration harness to place migrations
/// in order, therefore two different instances of this type
/// must be sortable
#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Ord, FromSqlRow, AsExpression)]
#[sql_type = "Text"]
pub struct MigrationVersion<'a>(Cow<'a, str>);

impl<'a> MigrationVersion<'a> {
    /// Convert the current migration version into
    /// a owned variant with static life time
    pub fn into_owned(&self) -> MigrationVersion<'static> {
        MigrationVersion(Cow::Owned(self.0.as_ref().to_owned()))
    }
}

impl<'a, DB> FromSql<Text, DB> for MigrationVersion<'a>
where
    String: FromSql<Text, DB>,
    DB: Backend,
{
    fn from_sql(bytes: crate::backend::RawValue<DB>) -> crate::deserialize::Result<Self> {
        let s = String::from_sql(bytes)?;
        Ok(Self(Cow::Owned(s)))
    }
}

impl<'a, DB> ToSql<Text, DB> for MigrationVersion<'a>
where
    Cow<'a, str>: ToSql<Text, DB>,
    DB: Backend,
{
    fn to_sql<W: std::io::Write>(
        &self,
        out: &mut crate::serialize::Output<W, DB>,
    ) -> crate::serialize::Result {
        self.0.to_sql(out)
    }
}

impl<'a> From<String> for MigrationVersion<'a> {
    fn from(s: String) -> Self {
        MigrationVersion(Cow::Owned(s))
    }
}

impl<'a> From<&'a str> for MigrationVersion<'a> {
    fn from(s: &'a str) -> Self {
        MigrationVersion(Cow::Borrowed(s))
    }
}

impl<'a> From<&'a String> for MigrationVersion<'a> {
    fn from(s: &'a String) -> Self {
        MigrationVersion(Cow::Borrowed(s))
    }
}

impl<'a> Display for MigrationVersion<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.0.as_ref())
    }
}

/// Represents the name of a migration
///
/// Users should threat this as `impl Display` type,
/// for implementors of custom migration types
/// this opens the possibility to roll out their own versioning
/// schema.
pub trait MigrationName: Display {
    /// The version corresponding to the current migration name
    fn version(&self) -> MigrationVersion;
}

/// Represents a migration that interacts with diesel
pub trait Migration<DB: Backend> {
    /// Apply this migration
    fn run(&self, conn: &dyn BoxableConnection<DB>) -> Result<()>;

    /// Revert this migration
    fn revert(&self, conn: &dyn BoxableConnection<DB>) -> Result<()>;

    /// Get a the attached metadata for this migration
    fn metadata(&self) -> &dyn MigrationMetadata;

    /// Get the name of the current migration
    ///
    /// The provided name is used by migration harness
    /// to get the version of a migration and to
    /// as something to that is displayed and allows
    /// user to identify a specific migration
    fn name(&self) -> &dyn MigrationName;
}

/// This trait is designed to customize the behaviour
/// of the default migration harness of diesel
///
/// Any new customization option will be added
/// as new function here. Each new function
/// will have a default implementation
/// returning the old a value corresponding
/// to the old uncustomized behaviour
pub trait MigrationMetadata {
    /// Should the current migration be executed in a migration
    /// or not?
    ///
    /// By default this function returns true
    fn run_in_transaction(&self) -> bool {
        true
    }
}

/// A migration source is an entity that can be used
/// to receive a number of migrations from.
pub trait MigrationSource<DB: Backend> {
    /// Get a list of migrations associated with this
    /// migration soucre.
    fn migrations(&self) -> Result<Vec<Box<dyn Migration<DB>>>>;
}

impl<'a, DB: Backend> Migration<DB> for Box<dyn Migration<DB> + 'a> {
    fn run(&self, conn: &dyn BoxableConnection<DB>) -> Result<()> {
        (&**self).run(conn)
    }

    fn revert(&self, conn: &dyn BoxableConnection<DB>) -> Result<()> {
        (&**self).revert(conn)
    }

    fn metadata(&self) -> &dyn MigrationMetadata {
        (&**self).metadata()
    }

    fn name(&self) -> &dyn MigrationName {
        (&**self).name()
    }
}

impl<'a, DB: Backend> Migration<DB> for &'a dyn Migration<DB> {
    fn run(&self, conn: &dyn BoxableConnection<DB>) -> Result<()> {
        (&**self).run(conn)
    }

    fn revert(&self, conn: &dyn BoxableConnection<DB>) -> Result<()> {
        (&**self).revert(conn)
    }

    fn metadata(&self) -> &dyn MigrationMetadata {
        (&**self).metadata()
    }

    fn name(&self) -> &dyn MigrationName {
        (&**self).name()
    }
}

/// Create table statement for the `__diesel_schema_migrations` used
/// used by the postgresql, sqlite and mysql backend
pub const CREATE_MIGRATIONS_TABLE: &str = include_str!("setup_migration_table.sql");

/// A trait indicating that a connection could be used to manage migrations
///
/// Only custom backend implementations need to think about this trait
pub trait MigrationConnection: Connection {
    /// Setup the following table:
    ///
    /// ```rust
    /// diesel::table! {
    ///      __diesel_schema_migrations(version) {
    ///          version -> Text,
    ///          /// defaults to `CURRENT_TIMESTAMP`
    ///          run_on -> Timestamp,
    ///      }
    /// }
    /// ```
    fn setup(&self) -> QueryResult<usize>;
}

#[cfg(feature = "postgres")]
impl MigrationConnection for crate::pg::PgConnection {
    fn setup(&self) -> QueryResult<usize> {
        use crate::RunQueryDsl;
        crate::sql_query(CREATE_MIGRATIONS_TABLE).execute(self)
    }
}

#[cfg(feature = "mysql")]
impl MigrationConnection for crate::mysql::MysqlConnection {
    fn setup(&self) -> QueryResult<usize> {
        use crate::RunQueryDsl;
        crate::sql_query(CREATE_MIGRATIONS_TABLE).execute(self)
    }
}

#[cfg(feature = "sqlite")]
impl MigrationConnection for crate::sqlite::SqliteConnection {
    fn setup(&self) -> QueryResult<usize> {
        use crate::RunQueryDsl;
        crate::sql_query(CREATE_MIGRATIONS_TABLE).execute(self)
    }
}
<commit_msg>Update diesel/src/migration/mod.rs<commit_after>//! Representation of migrations

use crate::backend::Backend;
use crate::connection::{BoxableConnection, Connection};
use crate::deserialize::{FromSql, FromSqlRow};
use crate::expression::AsExpression;
use crate::result::QueryResult;
use crate::serialize::ToSql;
use crate::sql_types::Text;
use std::borrow::Cow;
use std::error::Error;
use std::fmt::Display;

/// A specialized result type representing the result of
/// a migration operation
pub type Result<T> = std::result::Result<T, Box<dyn Error + Send + Sync>>;

/// A migration version identifier
///
/// This is used by the migration harness to place migrations
/// in order, therefore two different instances of this type
/// must be sortable
#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Ord, FromSqlRow, AsExpression)]
#[sql_type = "Text"]
pub struct MigrationVersion<'a>(Cow<'a, str>);

impl<'a> MigrationVersion<'a> {
    /// Convert the current migration version into
    /// a owned variant with static life time
    pub fn into_owned(&self) -> MigrationVersion<'static> {
        MigrationVersion(Cow::Owned(self.0.as_ref().to_owned()))
    }
}

impl<'a, DB> FromSql<Text, DB> for MigrationVersion<'a>
where
    String: FromSql<Text, DB>,
    DB: Backend,
{
    fn from_sql(bytes: crate::backend::RawValue<DB>) -> crate::deserialize::Result<Self> {
        let s = String::from_sql(bytes)?;
        Ok(Self(Cow::Owned(s)))
    }
}

impl<'a, DB> ToSql<Text, DB> for MigrationVersion<'a>
where
    Cow<'a, str>: ToSql<Text, DB>,
    DB: Backend,
{
    fn to_sql<W: std::io::Write>(
        &self,
        out: &mut crate::serialize::Output<W, DB>,
    ) -> crate::serialize::Result {
        self.0.to_sql(out)
    }
}

impl<'a> From<String> for MigrationVersion<'a> {
    fn from(s: String) -> Self {
        MigrationVersion(Cow::Owned(s))
    }
}

impl<'a> From<&'a str> for MigrationVersion<'a> {
    fn from(s: &'a str) -> Self {
        MigrationVersion(Cow::Borrowed(s))
    }
}

impl<'a> From<&'a String> for MigrationVersion<'a> {
    fn from(s: &'a String) -> Self {
        MigrationVersion(Cow::Borrowed(s))
    }
}

impl<'a> Display for MigrationVersion<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.0.as_ref())
    }
}

/// Represents the name of a migration
///
/// Users should threat this as `impl Display` type,
/// for implementors of custom migration types
/// this opens the possibility to roll out their own versioning
/// schema.
pub trait MigrationName: Display {
    /// The version corresponding to the current migration name
    fn version(&self) -> MigrationVersion;
}

/// Represents a migration that interacts with diesel
pub trait Migration<DB: Backend> {
    /// Apply this migration
    fn run(&self, conn: &dyn BoxableConnection<DB>) -> Result<()>;

    /// Revert this migration
    fn revert(&self, conn: &dyn BoxableConnection<DB>) -> Result<()>;

    /// Get a the attached metadata for this migration
    fn metadata(&self) -> &dyn MigrationMetadata;

    /// Get the name of the current migration
    ///
    /// The provided name is used by migration harness
    /// to get the version of a migration and to
    /// as something to that is displayed and allows
    /// user to identify a specific migration
    fn name(&self) -> &dyn MigrationName;
}

/// This trait is designed to customize the behaviour
/// of the default migration harness of diesel
///
/// Any new customization option will be added
/// as new function here. Each new function
/// will have a default implementation
/// returning the old a value corresponding
/// to the old uncustomized behaviour
pub trait MigrationMetadata {
    /// Whether the current migration be executed in a migration
    /// or not
    ///
    /// By default this function returns true
    fn run_in_transaction(&self) -> bool {
        true
    }
}

/// A migration source is an entity that can be used
/// to receive a number of migrations from.
pub trait MigrationSource<DB: Backend> {
    /// Get a list of migrations associated with this
    /// migration soucre.
    fn migrations(&self) -> Result<Vec<Box<dyn Migration<DB>>>>;
}

impl<'a, DB: Backend> Migration<DB> for Box<dyn Migration<DB> + 'a> {
    fn run(&self, conn: &dyn BoxableConnection<DB>) -> Result<()> {
        (&**self).run(conn)
    }

    fn revert(&self, conn: &dyn BoxableConnection<DB>) -> Result<()> {
        (&**self).revert(conn)
    }

    fn metadata(&self) -> &dyn MigrationMetadata {
        (&**self).metadata()
    }

    fn name(&self) -> &dyn MigrationName {
        (&**self).name()
    }
}

impl<'a, DB: Backend> Migration<DB> for &'a dyn Migration<DB> {
    fn run(&self, conn: &dyn BoxableConnection<DB>) -> Result<()> {
        (&**self).run(conn)
    }

    fn revert(&self, conn: &dyn BoxableConnection<DB>) -> Result<()> {
        (&**self).revert(conn)
    }

    fn metadata(&self) -> &dyn MigrationMetadata {
        (&**self).metadata()
    }

    fn name(&self) -> &dyn MigrationName {
        (&**self).name()
    }
}

/// Create table statement for the `__diesel_schema_migrations` used
/// used by the postgresql, sqlite and mysql backend
pub const CREATE_MIGRATIONS_TABLE: &str = include_str!("setup_migration_table.sql");

/// A trait indicating that a connection could be used to manage migrations
///
/// Only custom backend implementations need to think about this trait
pub trait MigrationConnection: Connection {
    /// Setup the following table:
    ///
    /// ```rust
    /// diesel::table! {
    ///      __diesel_schema_migrations(version) {
    ///          version -> Text,
    ///          /// defaults to `CURRENT_TIMESTAMP`
    ///          run_on -> Timestamp,
    ///      }
    /// }
    /// ```
    fn setup(&self) -> QueryResult<usize>;
}

#[cfg(feature = "postgres")]
impl MigrationConnection for crate::pg::PgConnection {
    fn setup(&self) -> QueryResult<usize> {
        use crate::RunQueryDsl;
        crate::sql_query(CREATE_MIGRATIONS_TABLE).execute(self)
    }
}

#[cfg(feature = "mysql")]
impl MigrationConnection for crate::mysql::MysqlConnection {
    fn setup(&self) -> QueryResult<usize> {
        use crate::RunQueryDsl;
        crate::sql_query(CREATE_MIGRATIONS_TABLE).execute(self)
    }
}

#[cfg(feature = "sqlite")]
impl MigrationConnection for crate::sqlite::SqliteConnection {
    fn setup(&self) -> QueryResult<usize> {
        use crate::RunQueryDsl;
        crate::sql_query(CREATE_MIGRATIONS_TABLE).execute(self)
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::sync::mpsc::{Sender, Receiver};
use std::sync::mpsc::TryRecvError;
use std::sync::{Arc, Mutex};
use std::old_io::timer;
use std::time::Duration;

use messages::{MarketMessages, MarketHistory, ActorMessages, TransactionRequest};
use messages::ActorMessages::{StockRequest, MoneyRequest, CommitTransaction, AbortTransaction, History, Time, ReceiveActivityCount};
use messages::MarketMessages::{BuyRequest, Commit, Cancel, RegisterActor, SellRequest};
use actor::Actor;

pub fn start_dummy_actor_2(actor_id: usize, existing_markets: HashMap<usize, Sender<MarketMessages>>, actor_tx: Sender<ActorMessages>, actor_rx: Receiver<ActorMessages>) {
  println!("Starting Dummy_Actor_2 {}", actor_id);
  let mut actor = Actor { id: actor_id,
                          money: 100,
                          stocks: HashMap::new(),
                          pending_money: 0,
                          pending_stock: (0, 0),
                          markets: existing_markets,
                          history: Arc::new(Mutex::new(MarketHistory {history: HashMap::new(), stocks: vec![]}))};

  //Register the actor with every market
  for (_, market_tx) in actor.markets.iter() {
    market_tx.send(RegisterActor(actor.id, actor_tx.clone())).unwrap();
  }

  let mut current_time: usize = 0;
  let mut max_time: usize = 0;

  //Number of stocks available to buy
  let numStocks = actor.stocks.len();
  let mut init_history = false;
  let mut stop_flag = false;
  let mut toSellPrices = HashMap::new();
  let mut stock_id_incr = 0;
  let mut activeBuyRequests = HashMap::new();
  let mut activeSellRequests = HashMap::new();
  let mut backupSellRequests = HashMap::new();

  loop {
    if stop_flag {
      timer::sleep(Duration::milliseconds(1000));
      continue;
    }

    //Logic
    let mark_clone = actor.markets.clone();
    let stock_clone = actor.stocks.clone();

    //buying and selling decisions
    ////////////////////////////////////////////////////////////////////
    if current_time > 3 * (max_time / 4) {
      for transaction_id in activeBuyRequests.iter() {

      }
      for transaction_id in activeSellRequests.iter() {

      }
    }

    if init_history {
      //Get variables for the actor's stocks
      let local_stocks;
      {
        local_stocks = actor.history.lock().unwrap().stocks.clone();
      }

      //Iterate through the actor's stocks
      for stock in local_stocks.iter() {
        match actor.stocks.get(stock) {
          //If the actor has some of a stock
          Some(count) => {

            //And he has not yet sent out a sell request
            if toSellPrices.contains_key(stock) {
              //Get the price he should sell it at (remove from HashMap)
              let sell_price = toSellPrices.remove(stock);
              match sell_price {
                Some(price) => {
                  for (_, market_tx) in actor.markets.iter() {
                    //Send out a sell request to sell it
                    let t = TransactionRequest{actor_id: actor.id, transaction_id: stock_id_incr, stock_id: *stock, price: price, quantity: 1};
                    market_tx.send(SellRequest(t)).unwrap();
                    activeSellRequests.insert(stock_id_incr, *stock);
                    stock_id_incr = stock_id_incr + 1;
                    backupSellRequests.insert(*stock, price / 2);
                  }
                },
                None => {}
              }
            }

            if current_time > 3 * (max_time / 4) {
              match backupSellRequests.remove(stock) {
                Some(price) => {
                  for (_, market_tx) in actor.markets.iter() {
                    let t = TransactionRequest{actor_id: actor.id, transaction_id: stock_id_incr, stock_id: *stock, price: 0, quantity: 1};
                    market_tx.send(SellRequest(t)).unwrap();
                    stock_id_incr = stock_id_incr + 1;
                  }
                },
                None => {}
              }
            }
          }
          //If the actor has none of a stock
          None => {
            if current_time < 3 * (max_time / 4) {
              if !toSellPrices.contains_key(stock) {
                //Make the price he should buy it at the most recently bought price
                let buy_price = actor.history.lock().unwrap().last_sold_price(*stock);
                match buy_price {
                  //If the stock was last bought at a price
                  Some(price) => {
                    //If the actor can afford to buy it
                    if actor.money > price {
                      for (_, market_tx) in actor.markets.iter() {
                        let t = TransactionRequest{actor_id: actor.id, transaction_id: stock_id_incr, stock_id: *stock, price: price, quantity: 1};
                        market_tx.send(BuyRequest(t)).unwrap();
                        activeBuyRequests.insert(stock_id_incr, *stock);
                        stock_id_incr = stock_id_incr + 1;
                        toSellPrices.insert(*stock, price * 2);
                      }
                    }
                  },
                  //If the stock has not been bought yet
                  None => {}
                }
              }
            }
          }
        }
      }

    }
    /////////////////////////////////////////////////////////////////////

    match actor_rx.try_recv() {
      Ok(message) => {
          match message {
            StockRequest(stock_request) => {
              //println!("Started Stock Request in dummy_actor_2 {}", actor.id);
              let market_tx;
              let tx_text = mark_clone.get(&stock_request.market_id);
              match tx_text {
                Some(tx) => {
                  market_tx = tx;
                },
                None => {
                  return; //HOW DID WE GET A MISSING MARKET?
                }
              }
              if has_pending_transaction(&actor) {
                market_tx.send(Cancel(actor.id)).unwrap();
              }
              else {
                //if we have the stock. set it aside.
                let stock_id = stock_request.stock_id;
                let quantity = stock_request.quantity;
                let stock = stock_clone.get(&stock_id);
                match stock {
                  Some(owned_quantity) => {
                    if *owned_quantity >= quantity {
                      remove_stock(&mut actor, (stock_id, quantity));
                      actor.pending_stock = (stock_id, quantity);
                      market_tx.send(Commit(actor.id)).unwrap();
                    }
                    else {
                      market_tx.send(Cancel(actor.id)).unwrap();
                    }
                  },
                  None => {
                    market_tx.send(Cancel(actor.id)).unwrap();
                  }
                }
              }
              //print_status(&actor);
              },
            MoneyRequest(money_request) => {
              let market_tx;
              let tx_text = actor.markets.get(&money_request.market_id);
              match tx_text {
                Some(tx) => {
                  market_tx = tx;
                },
                None => {
                  return; //HOW DID WE GET A MISSING MARKET?
                }
              }
              if has_pending_transaction(&actor) {
                market_tx.send(Cancel(actor.id)).unwrap();
              }
              else {
                //if we have the money. set it aside.
                if actor.money >= money_request.amount {
                  actor.money = actor.money - money_request.amount;
                  actor.pending_money = money_request.amount;
                  market_tx.send(Commit(actor.id)).unwrap();
                }
                else {
                  market_tx.send(Cancel(actor.id)).unwrap();
                }
              }
              //print_status(&actor);
              },
            CommitTransaction(commit_transaction_request) => {
              //if we have money pending, then look up the stock id and add that quantity purchased.
              //remove the pending money
              if actor.pending_money > 0 {
                let price_per_unit = commit_transaction_request.price;
                let units = commit_transaction_request.quantity;
                let leftover_money = actor.pending_money - price_per_unit * units;

                //make a function for adding stock.
                add_stock(&mut actor, (commit_transaction_request.stock_id, units));
                actor.money = actor.money + leftover_money;

                actor.pending_money = 0;

                //Remove this stock from the active buy requests
                activeBuyRequests.remove(&(commit_transaction_request.transaction_id));
              }

              //if we have stock pending, look up the quantity purchased and add the money.
              //remove the pending stock
              if actor.pending_stock.1 > 0 {
                let money = commit_transaction_request.price * commit_transaction_request.quantity;
                let restore_stock = (commit_transaction_request.stock_id, actor.pending_stock.1 - commit_transaction_request.quantity);
                if restore_stock.1 > 0 {
                  add_stock(&mut actor, restore_stock);
                }
                actor.money = actor.money + money;
                actor.pending_stock = (0,0);

                activeSellRequests.remove(&(commit_transaction_request.transaction_id));
                backupSellRequests.remove(&(commit_transaction_request.stock_id));
              }
              //println!("After a committed transaction in actor {} ", actor.id);
              //print_status(&actor);
              },
            AbortTransaction => {
              //move pending stock back into stocks.
              if actor.pending_stock.1 != 0 {
                let pending_stock_clone = actor.pending_stock.clone();
                add_stock(&mut actor, pending_stock_clone);
                //now that we have moved it. Clear out the pending stock.
                actor.pending_stock = (0,0); //setting the quantity to zero clears it.
              }
              for (stock_id, quantity) in actor.stocks.iter() {
                //println!("After aborting the transaction, actor {} now has StockId: {} Quantity: {}", actor.id, *stock_id, *quantity);
              }

              //move pending money back into money.
              if actor.pending_money > 0 {
                actor.money = actor.money + actor.pending_money;
                actor.pending_money = 0;
              }
              //println!("After aborting the transaction, actor {} now has {} money.", actor.id, actor.money);
            },
            History(history) => {
              //println!("Actor {} received history {}", actor.id, *(history.lock().unwrap()));
              actor.history = history;
              init_history = true},
            Time(current, max) => {
                current_time = current;
                max_time = max;
            },
            ReceiveActivityCount(_,_,_) => {},
            Stop => {
              print_status(&actor);
              stop_flag = true;
            }
          }
        },
      Err(TryRecvError::Empty) => {timer::sleep(Duration::milliseconds(1));},
      Err(TryRecvError::Disconnected) => {println!("ERROR: Actor {} disconnected", actor.id);}
    }
  }
}

fn add_stock(actor: &mut Actor, stock_to_add: (usize, usize)) {
  let stock_clone = actor.stocks.clone();
  let held_stock = stock_clone.get(&stock_to_add.0);
  match held_stock {
    Some(stock_count) => {
        //we have some stock. We need to add to our reserve.
        actor.stocks.insert(stock_to_add.0, stock_to_add.1 + *stock_count);
      },
    None => {
      //we don't have any stock left. Just add it back.
      actor.stocks.insert(stock_to_add.0, stock_to_add.1);
    }
  }
}

fn remove_stock(actor: &mut Actor, stock_to_remove: (usize, usize)) {
  let stock_clone = actor.stocks.clone();
  let held_stock = stock_clone.get(&stock_to_remove.0);
  match held_stock {
    Some(stock_count) => {
        //we have some stock. We need to add to our reserve.
        actor.stocks.insert(stock_to_remove.0, *stock_count - stock_to_remove.1);
      },
    None => {} //TODO Should we error handle here?
  }
}

fn print_status(actor: &Actor) {
  for (stock_id, quantity) in (*actor).stocks.iter() {
    println!("Dummy Actor_2 {} now has StockId: {} Quantity: {}", (*actor).id, *stock_id, *quantity);
  }
  println!("Dummy Actor_2 {} has {} money.", (*actor).id, (*actor).money);
}

fn has_pending_transaction(actor: &Actor) -> bool {
  (*actor).pending_money > 0 || (*actor).pending_stock.1 > 0
}
<commit_msg>Removed warnings from Dummy Actor 2<commit_after>use std::collections::HashMap;
use std::sync::mpsc::{Sender, Receiver};
use std::sync::mpsc::TryRecvError;
use std::sync::{Arc, Mutex};
use std::old_io::timer;
use std::time::Duration;

use messages::{MarketMessages, MarketHistory, ActorMessages, TransactionRequest};
use messages::ActorMessages::{StockRequest, MoneyRequest, CommitTransaction, AbortTransaction, History, Time, ReceiveActivityCount, Stop};
use messages::MarketMessages::{BuyRequest, Commit, Cancel, RegisterActor, SellRequest};
use actor::Actor;

pub fn start_dummy_actor_2(actor_id: usize, existing_markets: HashMap<usize, Sender<MarketMessages>>, actor_tx: Sender<ActorMessages>, actor_rx: Receiver<ActorMessages>) {
  println!("Starting Dummy_Actor_2 {}", actor_id);
  let mut actor = Actor { id: actor_id,
                          money: 100,
                          stocks: HashMap::new(),
                          pending_money: 0,
                          pending_stock: (0, 0),
                          markets: existing_markets,
                          history: Arc::new(Mutex::new(MarketHistory {history: HashMap::new(), stocks: vec![]}))};

  //Register the actor with every market
  for (_, market_tx) in actor.markets.iter() {
    market_tx.send(RegisterActor(actor.id, actor_tx.clone())).unwrap();
  }

  let mut current_time: usize = 0;
  let mut max_time: usize = 0;

  //Number of stocks available to buy
  let mut init_history = false;
  let mut stop_flag = false;
  let mut to_sell_prices = HashMap::new();
  let mut stock_id_incr = 0;
  let mut active_buy_requests = HashMap::new();
  let mut active_sell_requests = HashMap::new();
  let mut backup_sell_requests = HashMap::new();

  loop {
    if stop_flag {
      timer::sleep(Duration::milliseconds(1000));
      continue;
    }

    //Logic
    let mark_clone = actor.markets.clone();
    let stock_clone = actor.stocks.clone();

    //buying and selling decisions
    ////////////////////////////////////////////////////////////////////

    if init_history {
      //Get variables for the actor's stocks
      let local_stocks;
      {
        local_stocks = actor.history.lock().unwrap().stocks.clone();
      }

      //Iterate through the actor's stocks
      for stock in local_stocks.iter() {
        match actor.stocks.get(stock) {
          //If the actor has some of a stock
          Some(_) => {

            //And he has not yet sent out a sell request
            if to_sell_prices.contains_key(stock) {
              //Get the price he should sell it at (remove from HashMap)
              let sell_price = to_sell_prices.remove(stock);
              match sell_price {
                Some(price) => {
                  for (_, market_tx) in actor.markets.iter() {
                    //Send out a sell request to sell it
                    let t = TransactionRequest{actor_id: actor.id, transaction_id: stock_id_incr, stock_id: *stock, price: price, quantity: 1};
                    market_tx.send(SellRequest(t)).unwrap();
                    active_sell_requests.insert(stock_id_incr, *stock);
                    stock_id_incr = stock_id_incr + 1;
                    backup_sell_requests.insert(*stock, price / 2);
                  }
                },
                None => {}
              }
            }

            if current_time > 3 * (max_time / 4) {
              match backup_sell_requests.remove(stock) {
                Some(_) => {
                  for (_, market_tx) in actor.markets.iter() {
                    let t = TransactionRequest{actor_id: actor.id, transaction_id: stock_id_incr, stock_id: *stock, price: 0, quantity: 1};
                    market_tx.send(SellRequest(t)).unwrap();
                    stock_id_incr = stock_id_incr + 1;
                  }
                },
                None => {}
              }
            }
          }
          //If the actor has none of a stock
          None => {
            if current_time < 3 * (max_time / 4) {
              if !to_sell_prices.contains_key(stock) {
                //Make the price he should buy it at the most recently bought price
                let buy_price = actor.history.lock().unwrap().last_sold_price(*stock);
                match buy_price {
                  //If the stock was last bought at a price
                  Some(price) => {
                    //If the actor can afford to buy it
                    if actor.money > price {
                      for (_, market_tx) in actor.markets.iter() {
                        let t = TransactionRequest{actor_id: actor.id, transaction_id: stock_id_incr, stock_id: *stock, price: price, quantity: 1};
                        market_tx.send(BuyRequest(t)).unwrap();
                        active_buy_requests.insert(stock_id_incr, *stock);
                        stock_id_incr = stock_id_incr + 1;
                        to_sell_prices.insert(*stock, price * 2);
                      }
                    }
                  },
                  //If the stock has not been bought yet
                  None => {}
                }
              }
            }
          }
        }
      }

    }
    /////////////////////////////////////////////////////////////////////

    match actor_rx.try_recv() {
      Ok(message) => {
          match message {
            StockRequest(stock_request) => {
              //println!("Started Stock Request in dummy_actor_2 {}", actor.id);
              let market_tx;
              let tx_text = mark_clone.get(&stock_request.market_id);
              match tx_text {
                Some(tx) => {
                  market_tx = tx;
                },
                None => {
                  return; //HOW DID WE GET A MISSING MARKET?
                }
              }
              if has_pending_transaction(&actor) {
                market_tx.send(Cancel(actor.id)).unwrap();
              }
              else {
                //if we have the stock. set it aside.
                let stock_id = stock_request.stock_id;
                let quantity = stock_request.quantity;
                let stock = stock_clone.get(&stock_id);
                match stock {
                  Some(owned_quantity) => {
                    if *owned_quantity >= quantity {
                      remove_stock(&mut actor, (stock_id, quantity));
                      actor.pending_stock = (stock_id, quantity);
                      market_tx.send(Commit(actor.id)).unwrap();
                    }
                    else {
                      market_tx.send(Cancel(actor.id)).unwrap();
                    }
                  },
                  None => {
                    market_tx.send(Cancel(actor.id)).unwrap();
                  }
                }
              }
              //print_status(&actor);
              },
            MoneyRequest(money_request) => {
              let market_tx;
              let tx_text = actor.markets.get(&money_request.market_id);
              match tx_text {
                Some(tx) => {
                  market_tx = tx;
                },
                None => {
                  return; //HOW DID WE GET A MISSING MARKET?
                }
              }
              if has_pending_transaction(&actor) {
                market_tx.send(Cancel(actor.id)).unwrap();
              }
              else {
                //if we have the money. set it aside.
                if actor.money >= money_request.amount {
                  actor.money = actor.money - money_request.amount;
                  actor.pending_money = money_request.amount;
                  market_tx.send(Commit(actor.id)).unwrap();
                }
                else {
                  market_tx.send(Cancel(actor.id)).unwrap();
                }
              }
              //print_status(&actor);
              },
            CommitTransaction(commit_transaction_request) => {
              //if we have money pending, then look up the stock id and add that quantity purchased.
              //remove the pending money
              if actor.pending_money > 0 {
                let price_per_unit = commit_transaction_request.price;
                let units = commit_transaction_request.quantity;
                let leftover_money = actor.pending_money - price_per_unit * units;

                //make a function for adding stock.
                add_stock(&mut actor, (commit_transaction_request.stock_id, units));
                actor.money = actor.money + leftover_money;

                actor.pending_money = 0;

                //Remove this stock from the active buy requests
                active_buy_requests.remove(&(commit_transaction_request.transaction_id));
              }

              //if we have stock pending, look up the quantity purchased and add the money.
              //remove the pending stock
              if actor.pending_stock.1 > 0 {
                let money = commit_transaction_request.price * commit_transaction_request.quantity;
                let restore_stock = (commit_transaction_request.stock_id, actor.pending_stock.1 - commit_transaction_request.quantity);
                if restore_stock.1 > 0 {
                  add_stock(&mut actor, restore_stock);
                }
                actor.money = actor.money + money;
                actor.pending_stock = (0,0);

                active_sell_requests.remove(&(commit_transaction_request.transaction_id));
                backup_sell_requests.remove(&(commit_transaction_request.stock_id));
              }
              //println!("After a committed transaction in actor {} ", actor.id);
              //print_status(&actor);
              },
            AbortTransaction => {
              //move pending stock back into stocks.
              if actor.pending_stock.1 != 0 {
                let pending_stock_clone = actor.pending_stock.clone();
                add_stock(&mut actor, pending_stock_clone);
                //now that we have moved it. Clear out the pending stock.
                actor.pending_stock = (0,0); //setting the quantity to zero clears it.
              }
              //move pending money back into money.
              if actor.pending_money > 0 {
                actor.money = actor.money + actor.pending_money;
                actor.pending_money = 0;
              }
              //println!("After aborting the transaction, actor {} now has {} money.", actor.id, actor.money);
            },
            History(history) => {
              //println!("Actor {} received history {}", actor.id, *(history.lock().unwrap()));
              actor.history = history;
              init_history = true},
            Time(current, max) => {
                current_time = current;
                max_time = max;
            },
            ReceiveActivityCount(_,_,_) => {},
            Stop => {
              print_status(&actor);
              stop_flag = true;
            }
          }
        },
      Err(TryRecvError::Empty) => {timer::sleep(Duration::milliseconds(1));},
      Err(TryRecvError::Disconnected) => {println!("ERROR: Actor {} disconnected", actor.id);}
    }
  }
}

fn add_stock(actor: &mut Actor, stock_to_add: (usize, usize)) {
  let stock_clone = actor.stocks.clone();
  let held_stock = stock_clone.get(&stock_to_add.0);
  match held_stock {
    Some(stock_count) => {
        //we have some stock. We need to add to our reserve.
        actor.stocks.insert(stock_to_add.0, stock_to_add.1 + *stock_count);
      },
    None => {
      //we don't have any stock left. Just add it back.
      actor.stocks.insert(stock_to_add.0, stock_to_add.1);
    }
  }
}

fn remove_stock(actor: &mut Actor, stock_to_remove: (usize, usize)) {
  let stock_clone = actor.stocks.clone();
  let held_stock = stock_clone.get(&stock_to_remove.0);
  match held_stock {
    Some(stock_count) => {
        //we have some stock. We need to add to our reserve.
        actor.stocks.insert(stock_to_remove.0, *stock_count - stock_to_remove.1);
      },
    None => {} //TODO Should we error handle here?
  }
}

fn print_status(actor: &Actor) {
  for (stock_id, quantity) in (*actor).stocks.iter() {
    println!("Dummy Actor_2 {} now has StockId: {} Quantity: {}", (*actor).id, *stock_id, *quantity);
  }
  println!("Dummy Actor_2 {} has {} money.", (*actor).id, (*actor).money);
}

fn has_pending_transaction(actor: &Actor) -> bool {
  (*actor).pending_money > 0 || (*actor).pending_stock.1 > 0
}
<|endoftext|>


<commit_before>use decoders::*;
use decoders::tiff::*;
use decoders::basics::*;
use std::mem::transmute;
extern crate itertools;
use self::itertools::Itertools;
use std::f32;
use std::cmp;

#[derive(Debug, Clone)]
pub struct ArwDecoder<'a> {
  buffer: &'a [u8],
  rawloader: &'a RawLoader,
  tiff: TiffIFD<'a>,
}

impl<'a> ArwDecoder<'a> {
  pub fn new(buf: &'a [u8], tiff: TiffIFD<'a>, rawloader: &'a RawLoader) -> ArwDecoder<'a> {
    ArwDecoder {
      buffer: buf,
      tiff: tiff,
      rawloader: rawloader,
    }
  }
}

impl<'a> Decoder for ArwDecoder<'a> {
  fn identify(&self) -> Result<&Camera, String> {
    let make = fetch_tag!(self.tiff, Tag::Make, "ARW: Couldn't find Make").get_str();
    let model = fetch_tag!(self.tiff, Tag::Model, "ARW: Couldn't find Model").get_str();
    self.rawloader.check_supported(make, model)
  }

  fn image(&self) -> Result<Image,String> {
    let camera = try!(self.identify());
    let data = self.tiff.find_ifds_with_tag(Tag::StripOffsets);
    if data.len() == 0 {
      if camera.model == "DSLR-A100" {
        return self.image_a100(camera)
      } else { // try decoding as SRF
        return self.image_srf(camera)
      }
    }
    let raw = data[0];
    let width = fetch_tag!(raw, Tag::ImageWidth, "ARW: Couldn't find width").get_u16(0) as u32;
    let mut height = fetch_tag!(raw, Tag::ImageLength, "ARW: Couldn't find height").get_u16(0) as u32;
    let offset = fetch_tag!(raw, Tag::StripOffsets, "ARW: Couldn't find offset").get_u32(0) as usize;
    let count = fetch_tag!(raw, Tag::StripByteCounts, "ARW: Couldn't find byte count").get_u32(0) as usize;
    let compression = fetch_tag!(raw, Tag::Compression, "ARW: Couldn't find Compression").get_u16(0);
    let bps = if camera.bps != 0 {
      camera.bps
    } else {
      fetch_tag!(raw, Tag::BitsPerSample, "ARW: Couldn't find bps").get_u16(0) as u32
    };
    let src = &self.buffer[offset .. self.buffer.len()];

    let image = match compression {
      1 => {
        if camera.model == "DSC-R1" {
          decode_14be_unpacked(src, width as usize, height as usize)
        } else {
          decode_16le(src, width as usize, height as usize)
        }
      }
      32767 => {
        if ((width*height*bps) as usize) != count*8 {
          height += 8;
          ArwDecoder::decode_arw1(src, width as usize, height as usize)
        } else {
          match bps {
            8 => {let curve = try!(ArwDecoder::get_curve(raw)); ArwDecoder::decode_arw2(src, width as usize, height as usize, &curve)},
            12 => decode_12le(src, width as usize, height as usize),
            _ => return Err(format!("ARW2: Don't know how to decode images with {} bps", bps)),
          }
        }
      },
      _ => return Err(format!("ARW: Don't know how to decode type {}", compression).to_string()),
    };

    ok_image(camera, width, height, try!(self.get_wb()), image)
  }
}

impl<'a> ArwDecoder<'a> {
  fn image_a100(&self, camera: &Camera) -> Result<Image,String> {
    // We've caught the elusive A100 in the wild, a transitional format
    // between the simple sanity of the MRW custom format and the wordly
    // wonderfullness of the Tiff-based ARW format, let's shoot from the hip
    let data = self.tiff.find_ifds_with_tag(Tag::SubIFDs);
    if data.len() == 0 {
      return Err("ARW: Couldn't find the data IFD!".to_string())
    }
    let raw = data[0];
    let width = 3881;
    let height = 2608;
    let offset = fetch_tag!(raw, Tag::SubIFDs, "ARW: Couldn't find offset").get_u32(0) as usize;

    let src = &self.buffer[offset .. self.buffer.len()];
    let image = ArwDecoder::decode_arw1(src, width as usize, height as usize);

    // Get the WB the MRW way
    let priv_offset = fetch_tag!(self.tiff, Tag::DNGPrivateArea, "ARW: Couldn't find private offset").get_u32(0);
    let buf = &self.buffer[priv_offset as usize..];
    let mut currpos: usize = 8;
    let mut wb_vals: [u16;4] = [0;4];
    // At most we read 20 bytes from currpos so check we don't step outside that
    while currpos+20 < buf.len() {
      let tag: u32 = BEu32(buf,currpos);
      let len: u32 = LEu32(buf,currpos+4);
      if tag == 0x574247 { // WBG
        for i in 0..4 {
          wb_vals[i] = LEu16(buf, currpos+12+i*2);
        }
        break;
      }
      currpos += (len+8) as usize;
    }
    let wb_coeffs = [wb_vals[0] as f32, wb_vals[1] as f32, wb_vals[3] as f32, f32::NAN];

    ok_image(camera, width, height, wb_coeffs, image)
  }

  fn image_srf(&self, camera: &Camera) -> Result<Image,String> {
    let data = self.tiff.find_ifds_with_tag(Tag::ImageWidth);
    if data.len() == 0 {
      return Err("ARW: Couldn't find the data IFD!".to_string())
    }
    let raw = data[0];

    let width = fetch_tag!(raw, Tag::ImageWidth, "SRF: Couldn't find width").get_u16(0) as u32;
    let height = fetch_tag!(raw, Tag::ImageLength, "SRF: Couldn't find height").get_u16(0) as u32;
    let len = (width*height*2) as usize;

    // Constants taken from dcraw
    let off: usize = 862144;
    let key_off: usize = 200896;
    let head_off: usize = 164600;

    // Replicate the dcraw contortions to get the "decryption" key
    let offset = (self.buffer[key_off as usize] as usize)*4;
    let first_key = BEu32(self.buffer, key_off+offset);
    let head = ArwDecoder::sony_decrypt(self.buffer, head_off, 40, first_key);
    let second_key = LEu32(&head, 22);

    // "Decrypt" the whole image buffer
    let image_data = ArwDecoder::sony_decrypt(self.buffer, off, len, second_key);
    let image = decode_16be(&image_data, width as usize, height as usize);

    ok_image(camera, width, height, [f32::NAN,f32::NAN,f32::NAN,f32::NAN], image)
  }

  fn decode_arw1(buf: &[u8], width: usize, height: usize) -> Vec<u16> {
    let mut pump = BitPump::new(buf, PumpOrder::MSB);
    let mut out: Vec<u16> = vec![0; (width*height) as usize];

    let mut sum: i32 = 0;
    for x in 0..width {
      let col = width-1-x;
      let mut row = 0;
      while row <= height {
        if row == height {
          row = 1;
        }

        let mut len: u32 = 4 - pump.get_bits(2);
        if len == 3 && pump.get_bits(1) != 0 {
          len = 0;
        } else if len == 4 {
          let zeros = pump.peek_bits(13).leading_zeros() - 19;
          len += zeros;
          pump.get_bits(cmp::min(13, zeros+1));
        }
        let diff: i32 = pump.get_ibits(len);
        sum += diff;
        if len > 0 && (diff & (1 << (len - 1))) == 0 {
          sum -= (1 << len) - 1;
        }
        out[row*width+col] = sum as u16;
        row += 2
      }
    }
    out
  }

  fn decode_arw2(buf: &[u8], width: usize, height: usize, curve: &LookupTable) -> Vec<u16> {
    decode_threaded(width, height, &(|out: &mut [u16], start, width, height| {
      let mut pump = BitPump::new(&buf[(start*width)..], PumpOrder::LSB);

      for row in 0..height {
        // Process 16 pixels at a time in interleaved fashion
        let mut col = 0;
        let mut random = pump.peek_bits(24);
        while col < (width-30) {
          let max = pump.get_bits(11);
          let min = pump.get_bits(11);
          let delta = max-min;
          // Calculate the size of the data shift needed by how large the delta is
          // A delta with 11 bits requires a shift of 4, 10 bits of 3, etc
          let delta_shift: u32 = cmp::max(0, (32-(delta.leading_zeros() as i32)) - 7) as u32;
          let imax = pump.get_bits(4);
          let imin = pump.get_bits(4);

          for i in 0..16 {
            let val = if i == imax {
              max
            } else if i == imin {
              min
            } else {
              cmp::min(0x7ff,(pump.get_bits(7) << delta_shift) + min)
            };
            out[row*width+col+(i*2) as usize] = curve.dither(val as u16, &mut random);
          }
          col += if (col & 1) != 0 {31} else {1};  // Skip to next 16 pixels
        }
      }
    }))
  }

  fn get_wb(&self) -> Result<[f32;4], String> {
    let priv_offset = fetch_tag!(self.tiff, Tag::DNGPrivateArea, "ARW: Couldn't find private offset").get_u32(0);
    let priv_tiff = TiffIFD::new(self.buffer, priv_offset as usize, 0, 0, LITTLE_ENDIAN).unwrap();
    let sony_offset = fetch_tag!(priv_tiff, Tag::SonyOffset, "ARW: Couldn't find sony offset").get_u32(0) as usize;
    let sony_length = fetch_tag!(priv_tiff, Tag::SonyLength, "ARW: Couldn't find sony length").get_u32(0) as usize;
    let sony_key = fetch_tag!(priv_tiff, Tag::SonyKey, "ARW: Couldn't find sony key").get_u32(0);
    let decrypted_buf = ArwDecoder::sony_decrypt(self.buffer, sony_offset, sony_length, sony_key);
    let decrypted_tiff = TiffIFD::new(&decrypted_buf, 0, sony_offset, 0, LITTLE_ENDIAN).unwrap();
    let grgb_levels = decrypted_tiff.find_entry(Tag::SonyGRBG);
    let rggb_levels = decrypted_tiff.find_entry(Tag::SonyRGGB);
    if grgb_levels.is_some() {
      let levels = grgb_levels.unwrap();
      Ok([levels.get_u16(1) as f32, levels.get_u16(0) as f32, levels.get_u16(2) as f32, f32::NAN])
    } else if rggb_levels.is_some() {
      let levels = rggb_levels.unwrap();
      Ok([levels.get_u16(0) as f32, levels.get_u16(1) as f32, levels.get_u16(3) as f32, f32::NAN])
    } else {
      Err("ARW: Couldn't find GRGB or RGGB levels".to_string())
    }
  }

  fn get_curve(raw: &TiffIFD) -> Result<LookupTable, String> {
    let centry = fetch_tag!(raw, Tag::SonyCurve, "ARW: Couldn't find sony curve");
    let mut curve: [u32;6] = [ 0, 0, 0, 0, 0, 4095 ];

    for i in 0..4 {
      curve[i+1] = ((centry.get_u16(i) >> 2) & 0xfff) as u32;
    }

    let mut out = vec![0 as u16; (curve[5]+1) as usize];
    for i in 0..5 {
      for j in (curve[i]+1)..(curve[i+1]+1) {
        out[j as usize] = out[(j-1) as usize] + (1<<i);
      }
    }

    Ok(LookupTable::new(&out))
  }

  fn sony_decrypt(buf: &[u8], offset: usize, length: usize, key: u32) -> Vec<u8>{
    let mut pad: [u32; 128] = [0 as u32; 128];
    let mut mkey = key;
    // Make sure we always have space for the final bytes even if the buffer
    // isn't a multiple of 4
    let mut out = vec![0;length+4];
    // Initialize the decryption pad from the key
    for p in 0..4 {
      mkey = mkey.wrapping_mul(48828125).wrapping_add(1);
      pad[p] = mkey;
    }
    pad[3] = pad[3] << 1 | (pad[0]^pad[2]) >> 31;
    for p in 4..127 {
      pad[p] = (pad[p-4]^pad[p-2]) << 1 | (pad[p-3]^pad[p-1]) >> 31;
    }
    for p in 0..127 {
      pad[p] = u32::from_be(pad[p]);
    }

    for i in (0..length).step(4) {
      let p = i/4 + 127;
      pad[p & 127] = pad[(p+1) & 127] ^ pad[(p+1+64) & 127];
      let output = LEu32(buf, offset+i) ^ pad[p & 127];
      let bytes: [u8; 4] = unsafe { transmute(output.to_le()) };
      out[i]   = bytes[0];
      out[i+1] = bytes[1];
      out[i+2] = bytes[2];
      out[i+3] = bytes[3];
    }
    out
  }
}
<commit_msg>ARW: Correctly shift values before applying curve<commit_after>use decoders::*;
use decoders::tiff::*;
use decoders::basics::*;
use std::mem::transmute;
extern crate itertools;
use self::itertools::Itertools;
use std::f32;
use std::cmp;

#[derive(Debug, Clone)]
pub struct ArwDecoder<'a> {
  buffer: &'a [u8],
  rawloader: &'a RawLoader,
  tiff: TiffIFD<'a>,
}

impl<'a> ArwDecoder<'a> {
  pub fn new(buf: &'a [u8], tiff: TiffIFD<'a>, rawloader: &'a RawLoader) -> ArwDecoder<'a> {
    ArwDecoder {
      buffer: buf,
      tiff: tiff,
      rawloader: rawloader,
    }
  }
}

impl<'a> Decoder for ArwDecoder<'a> {
  fn identify(&self) -> Result<&Camera, String> {
    let make = fetch_tag!(self.tiff, Tag::Make, "ARW: Couldn't find Make").get_str();
    let model = fetch_tag!(self.tiff, Tag::Model, "ARW: Couldn't find Model").get_str();
    self.rawloader.check_supported(make, model)
  }

  fn image(&self) -> Result<Image,String> {
    let camera = try!(self.identify());
    let data = self.tiff.find_ifds_with_tag(Tag::StripOffsets);
    if data.len() == 0 {
      if camera.model == "DSLR-A100" {
        return self.image_a100(camera)
      } else { // try decoding as SRF
        return self.image_srf(camera)
      }
    }
    let raw = data[0];
    let width = fetch_tag!(raw, Tag::ImageWidth, "ARW: Couldn't find width").get_u16(0) as u32;
    let mut height = fetch_tag!(raw, Tag::ImageLength, "ARW: Couldn't find height").get_u16(0) as u32;
    let offset = fetch_tag!(raw, Tag::StripOffsets, "ARW: Couldn't find offset").get_u32(0) as usize;
    let count = fetch_tag!(raw, Tag::StripByteCounts, "ARW: Couldn't find byte count").get_u32(0) as usize;
    let compression = fetch_tag!(raw, Tag::Compression, "ARW: Couldn't find Compression").get_u16(0);
    let bps = if camera.bps != 0 {
      camera.bps
    } else {
      fetch_tag!(raw, Tag::BitsPerSample, "ARW: Couldn't find bps").get_u16(0) as u32
    };
    let src = &self.buffer[offset .. self.buffer.len()];

    let image = match compression {
      1 => {
        if camera.model == "DSC-R1" {
          decode_14be_unpacked(src, width as usize, height as usize)
        } else {
          decode_16le(src, width as usize, height as usize)
        }
      }
      32767 => {
        if ((width*height*bps) as usize) != count*8 {
          height += 8;
          ArwDecoder::decode_arw1(src, width as usize, height as usize)
        } else {
          match bps {
            8 => {let curve = try!(ArwDecoder::get_curve(raw)); ArwDecoder::decode_arw2(src, width as usize, height as usize, &curve)},
            12 => decode_12le(src, width as usize, height as usize),
            _ => return Err(format!("ARW2: Don't know how to decode images with {} bps", bps)),
          }
        }
      },
      _ => return Err(format!("ARW: Don't know how to decode type {}", compression).to_string()),
    };

    ok_image(camera, width, height, try!(self.get_wb()), image)
  }
}

impl<'a> ArwDecoder<'a> {
  fn image_a100(&self, camera: &Camera) -> Result<Image,String> {
    // We've caught the elusive A100 in the wild, a transitional format
    // between the simple sanity of the MRW custom format and the wordly
    // wonderfullness of the Tiff-based ARW format, let's shoot from the hip
    let data = self.tiff.find_ifds_with_tag(Tag::SubIFDs);
    if data.len() == 0 {
      return Err("ARW: Couldn't find the data IFD!".to_string())
    }
    let raw = data[0];
    let width = 3881;
    let height = 2608;
    let offset = fetch_tag!(raw, Tag::SubIFDs, "ARW: Couldn't find offset").get_u32(0) as usize;

    let src = &self.buffer[offset .. self.buffer.len()];
    let image = ArwDecoder::decode_arw1(src, width as usize, height as usize);

    // Get the WB the MRW way
    let priv_offset = fetch_tag!(self.tiff, Tag::DNGPrivateArea, "ARW: Couldn't find private offset").get_u32(0);
    let buf = &self.buffer[priv_offset as usize..];
    let mut currpos: usize = 8;
    let mut wb_vals: [u16;4] = [0;4];
    // At most we read 20 bytes from currpos so check we don't step outside that
    while currpos+20 < buf.len() {
      let tag: u32 = BEu32(buf,currpos);
      let len: u32 = LEu32(buf,currpos+4);
      if tag == 0x574247 { // WBG
        for i in 0..4 {
          wb_vals[i] = LEu16(buf, currpos+12+i*2);
        }
        break;
      }
      currpos += (len+8) as usize;
    }
    let wb_coeffs = [wb_vals[0] as f32, wb_vals[1] as f32, wb_vals[3] as f32, f32::NAN];

    ok_image(camera, width, height, wb_coeffs, image)
  }

  fn image_srf(&self, camera: &Camera) -> Result<Image,String> {
    let data = self.tiff.find_ifds_with_tag(Tag::ImageWidth);
    if data.len() == 0 {
      return Err("ARW: Couldn't find the data IFD!".to_string())
    }
    let raw = data[0];

    let width = fetch_tag!(raw, Tag::ImageWidth, "SRF: Couldn't find width").get_u16(0) as u32;
    let height = fetch_tag!(raw, Tag::ImageLength, "SRF: Couldn't find height").get_u16(0) as u32;
    let len = (width*height*2) as usize;

    // Constants taken from dcraw
    let off: usize = 862144;
    let key_off: usize = 200896;
    let head_off: usize = 164600;

    // Replicate the dcraw contortions to get the "decryption" key
    let offset = (self.buffer[key_off as usize] as usize)*4;
    let first_key = BEu32(self.buffer, key_off+offset);
    let head = ArwDecoder::sony_decrypt(self.buffer, head_off, 40, first_key);
    let second_key = LEu32(&head, 22);

    // "Decrypt" the whole image buffer
    let image_data = ArwDecoder::sony_decrypt(self.buffer, off, len, second_key);
    let image = decode_16be(&image_data, width as usize, height as usize);

    ok_image(camera, width, height, [f32::NAN,f32::NAN,f32::NAN,f32::NAN], image)
  }

  fn decode_arw1(buf: &[u8], width: usize, height: usize) -> Vec<u16> {
    let mut pump = BitPump::new(buf, PumpOrder::MSB);
    let mut out: Vec<u16> = vec![0; (width*height) as usize];

    let mut sum: i32 = 0;
    for x in 0..width {
      let col = width-1-x;
      let mut row = 0;
      while row <= height {
        if row == height {
          row = 1;
        }

        let mut len: u32 = 4 - pump.get_bits(2);
        if len == 3 && pump.get_bits(1) != 0 {
          len = 0;
        } else if len == 4 {
          let zeros = pump.peek_bits(13).leading_zeros() - 19;
          len += zeros;
          pump.get_bits(cmp::min(13, zeros+1));
        }
        let diff: i32 = pump.get_ibits(len);
        sum += diff;
        if len > 0 && (diff & (1 << (len - 1))) == 0 {
          sum -= (1 << len) - 1;
        }
        out[row*width+col] = sum as u16;
        row += 2
      }
    }
    out
  }

  fn decode_arw2(buf: &[u8], width: usize, height: usize, curve: &LookupTable) -> Vec<u16> {
    decode_threaded(width, height, &(|out: &mut [u16], start, width, height| {
      let mut pump = BitPump::new(&buf[(start*width)..], PumpOrder::LSB);

      for row in 0..height {
        // Process 16 pixels at a time in interleaved fashion
        let mut col = 0;
        let mut random = pump.peek_bits(24);
        while col < (width-30) {
          let max = pump.get_bits(11);
          let min = pump.get_bits(11);
          let delta = max-min;
          // Calculate the size of the data shift needed by how large the delta is
          // A delta with 11 bits requires a shift of 4, 10 bits of 3, etc
          let delta_shift: u32 = cmp::max(0, (32-(delta.leading_zeros() as i32)) - 7) as u32;
          let imax = pump.get_bits(4);
          let imin = pump.get_bits(4);

          for i in 0..16 {
            let val = if i == imax {
              max
            } else if i == imin {
              min
            } else {
              cmp::min(0x7ff,(pump.get_bits(7) << delta_shift) + min)
            };
            out[row*width+col+(i*2) as usize] = curve.dither((val<<1) as u16, &mut random);
          }
          col += if (col & 1) != 0 {31} else {1};  // Skip to next 16 pixels
        }
      }
    }))
  }

  fn get_wb(&self) -> Result<[f32;4], String> {
    let priv_offset = fetch_tag!(self.tiff, Tag::DNGPrivateArea, "ARW: Couldn't find private offset").get_u32(0);
    let priv_tiff = TiffIFD::new(self.buffer, priv_offset as usize, 0, 0, LITTLE_ENDIAN).unwrap();
    let sony_offset = fetch_tag!(priv_tiff, Tag::SonyOffset, "ARW: Couldn't find sony offset").get_u32(0) as usize;
    let sony_length = fetch_tag!(priv_tiff, Tag::SonyLength, "ARW: Couldn't find sony length").get_u32(0) as usize;
    let sony_key = fetch_tag!(priv_tiff, Tag::SonyKey, "ARW: Couldn't find sony key").get_u32(0);
    let decrypted_buf = ArwDecoder::sony_decrypt(self.buffer, sony_offset, sony_length, sony_key);
    let decrypted_tiff = TiffIFD::new(&decrypted_buf, 0, sony_offset, 0, LITTLE_ENDIAN).unwrap();
    let grgb_levels = decrypted_tiff.find_entry(Tag::SonyGRBG);
    let rggb_levels = decrypted_tiff.find_entry(Tag::SonyRGGB);
    if grgb_levels.is_some() {
      let levels = grgb_levels.unwrap();
      Ok([levels.get_u16(1) as f32, levels.get_u16(0) as f32, levels.get_u16(2) as f32, f32::NAN])
    } else if rggb_levels.is_some() {
      let levels = rggb_levels.unwrap();
      Ok([levels.get_u16(0) as f32, levels.get_u16(1) as f32, levels.get_u16(3) as f32, f32::NAN])
    } else {
      Err("ARW: Couldn't find GRGB or RGGB levels".to_string())
    }
  }

  fn get_curve(raw: &TiffIFD) -> Result<LookupTable, String> {
    let centry = fetch_tag!(raw, Tag::SonyCurve, "ARW: Couldn't find sony curve");
    let mut curve: [u32;6] = [ 0, 0, 0, 0, 0, 4095 ];

    for i in 0..4 {
      curve[i+1] = ((centry.get_u16(i) >> 2) & 0xfff) as u32;
    }

    let mut out = vec![0 as u16; (curve[5]+1) as usize];
    for i in 0..5 {
      for j in (curve[i]+1)..(curve[i+1]+1) {
        out[j as usize] = out[(j-1) as usize] + (1<<i);
      }
    }

    Ok(LookupTable::new(&out))
  }

  fn sony_decrypt(buf: &[u8], offset: usize, length: usize, key: u32) -> Vec<u8>{
    let mut pad: [u32; 128] = [0 as u32; 128];
    let mut mkey = key;
    // Make sure we always have space for the final bytes even if the buffer
    // isn't a multiple of 4
    let mut out = vec![0;length+4];
    // Initialize the decryption pad from the key
    for p in 0..4 {
      mkey = mkey.wrapping_mul(48828125).wrapping_add(1);
      pad[p] = mkey;
    }
    pad[3] = pad[3] << 1 | (pad[0]^pad[2]) >> 31;
    for p in 4..127 {
      pad[p] = (pad[p-4]^pad[p-2]) << 1 | (pad[p-3]^pad[p-1]) >> 31;
    }
    for p in 0..127 {
      pad[p] = u32::from_be(pad[p]);
    }

    for i in (0..length).step(4) {
      let p = i/4 + 127;
      pad[p & 127] = pad[(p+1) & 127] ^ pad[(p+1+64) & 127];
      let output = LEu32(buf, offset+i) ^ pad[p & 127];
      let bytes: [u8; 4] = unsafe { transmute(output.to_le()) };
      out[i]   = bytes[0];
      out[i+1] = bytes[1];
      out[i+2] = bytes[2];
      out[i+3] = bytes[3];
    }
    out
  }
}
<|endoftext|>