<commit_before>use syn;
use quote;
use proc_macro2::Span;

use field::*;
use model::*;
use util::*;

pub fn derive(item: syn::DeriveInput) -> Result<quote::Tokens, Diagnostic> {
    let model = Model::from_item(&item)?;

    if model.fields().is_empty() {
        return Err(Span::call_site()
            .error("Cannot derive Insertable for unit structs")
            .help(format!(
                "Use `insert_into({}::table).default_values()` if you want `DEFAULT VALUES`",
                model.table_name()
            )));
    }

    let table_name = model.table_name();
    let struct_name = item.ident;

    let (impl_generics, ty_generics, where_clause) = item.generics.split_for_impl();
    let mut ref_generics = item.generics.clone();
    ref_generics.params.push(parse_quote!('insert));
    let (ref_generics, ..) = ref_generics.split_for_impl();

    let (field_ty, field_assign): (Vec<_>, Vec<_>) = model
        .fields()
        .iter()
        .map(|f| (field_ty(f, table_name), field_expr(f, table_name)))
        .unzip();

    Ok(wrap_in_dummy_mod(
        model.dummy_mod_name("insertable"),
        quote! {
            use self::diesel::insertable::Insertable;
            use self::diesel::query_builder::UndecoratedInsertRecord;
            use self::diesel::prelude::*;

            impl #ref_generics Insertable<#table_name::table>
                for &'insert #struct_name #ty_generics
            #where_clause
            {
                type Values = <(#(#field_ty,)*) as Insertable<#table_name::table>>::Values;

                fn values(self) -> Self::Values {
                    (#(#field_assign,)*).values()
                }
            }

            impl #impl_generics UndecoratedInsertRecord<#table_name::table>
                for #struct_name #ty_generics
            #where_clause
            {
            }
        },
    ))
}

fn field_ty(field: &Field, table_name: syn::Ident) -> syn::Type {
    if field.has_flag("embed") {
        let field_ty = &field.ty;
        parse_quote!(&'insert #field_ty)
    } else {
        let inner_ty = inner_of_option_ty(&field.ty);
        let column_name = field.column_name();
        parse_quote!(
            std::option::Option<diesel::dsl::Eq<
                #table_name::#column_name,
                &'insert #inner_ty,
            >>
        )
    }
}

fn field_expr(field: &Field, table_name: syn::Ident) -> syn::Expr {
    let field_access = field.name.access();
    if field.has_flag("embed") {
        parse_quote!(&self#field_access)
    } else {
        let column_name = field.column_name();
        let column: syn::Expr = parse_quote!(#table_name::#column_name);
        if is_option_ty(&field.ty) {
            parse_quote!(self#field_access.as_ref().map(|x| #column.eq(x)))
        } else {
            parse_quote!(std::option::Option::Some(#column.eq(&self#field_access)))
        }
    }
}
<commit_msg>Derive insertable also on the type, not only on the reference<commit_after>use syn;
use quote;
use proc_macro2::Span;

use field::*;
use model::*;
use util::*;

pub fn derive(item: syn::DeriveInput) -> Result<quote::Tokens, Diagnostic> {
    let model = Model::from_item(&item)?;

    if model.fields().is_empty() {
        return Err(Span::call_site()
            .error("Cannot derive Insertable for unit structs")
            .help(format!(
                "Use `insert_into({}::table).default_values()` if you want `DEFAULT VALUES`",
                model.table_name()
            )));
    }

    let table_name = model.table_name();
    let struct_name = item.ident;

    let (_, ty_generics, where_clause) = item.generics.split_for_impl();
    let mut impl_generics = item.generics.clone();
    impl_generics.params.push(parse_quote!('insert));
    let (impl_generics, ..) = impl_generics.split_for_impl();

    let (direct_field_ty, direct_field_assign): (Vec<_>, Vec<_>) = model
        .fields()
        .iter()
        .map(|f| {
            (
                (field_ty(f, table_name, None)),
                (field_expr(f, table_name, None)),
            )
        })
        .unzip();

    let (ref_field_ty, ref_field_assign): (Vec<_>, Vec<_>) = model
        .fields()
        .iter()
        .map(|f| {
            (
                (field_ty(f, table_name, Some(quote!(&'insert)))),
                (field_expr(f, table_name, Some(quote!(&)))),
            )
        })
        .unzip();

    Ok(wrap_in_dummy_mod(
        model.dummy_mod_name("insertable"),
        quote! {
            use self::diesel::insertable::Insertable;
            use self::diesel::query_builder::UndecoratedInsertRecord;
            use self::diesel::prelude::*;

            impl #impl_generics Insertable<#table_name::table> for #struct_name #ty_generics
                #where_clause
            {
                type Values = <(#(#direct_field_ty,)*) as Insertable<#table_name::table>>::Values;

                fn values(self) -> Self::Values {
                    (#(#direct_field_assign,)*).values()
                }
            }

            impl #impl_generics Insertable<#table_name::table>
                for &'insert #struct_name #ty_generics
            #where_clause
            {
                type Values = <(#(#ref_field_ty,)*) as Insertable<#table_name::table>>::Values;

                fn values(self) -> Self::Values {
                    (#(#ref_field_assign,)*).values()
                }
            }

            impl #impl_generics UndecoratedInsertRecord<#table_name::table>
                for #struct_name #ty_generics
            #where_clause
            {
            }
        },
    ))
}

fn field_ty(field: &Field, table_name: syn::Ident, lifetime: Option<quote::Tokens>) -> syn::Type {
    if field.has_flag("embed") {
        let field_ty = &field.ty;
        parse_quote!(#lifetime #field_ty)
    } else {
        let inner_ty = inner_of_option_ty(&field.ty);
        let column_name = field.column_name();
        parse_quote!(
            std::option::Option<diesel::dsl::Eq<
                #table_name::#column_name,
                #lifetime #inner_ty,
            >>
        )
    }
}

fn field_expr(field: &Field, table_name: syn::Ident, lifetime: Option<quote::Tokens>) -> syn::Expr {
    let field_access = field.name.access();
    if field.has_flag("embed") {
        parse_quote!(#lifetime self#field_access)
    } else {
        let column_name = field.column_name();
        let column: syn::Expr = parse_quote!(#table_name::#column_name);
        if is_option_ty(&field.ty) {
            if lifetime.is_some() {
                parse_quote!(self#field_access.as_ref().map(|x| #column.eq(x)))
            } else {
                parse_quote!(self#field_access.map(|x| #column.eq(x)))
            }
        } else {
            parse_quote!(std::option::Option::Some(#column.eq(#lifetime self#field_access)))
        }
    }
}
<|endoftext|>


<commit_before>use nom::{IResult,le_i32,le_u8};
use nom::Err;
use nom::IResult::*;
use directory_header::*;
use helpers::*;
use raw_bsp::*;
use std::str::from_utf8;

named! {
    parse_entity_value <EntityValue>,
    alt!(
        chain!(
                       tag!(b"*")    ~
            reference: parse_str_int ,
            || {
                EntityValue::ModelRef(reference)
            }
        ) |
        chain!(
                       tag!(b"t")    ~
            reference: parse_str_int ,
            || {
                EntityValue::TargetRef(reference)
            }
        ) |
        chain!(
            first: parse_str_int        ~
                   mandatory_whitespace ~
            secnd: parse_str_int        ~
                   mandatory_whitespace ~
            third: parse_str_int        ~
                   whitespace ,
            || { EntityValue::IVec3([first, secnd, third]) }
        ) |
        chain!(
            first: parse_str_float      ~
                   mandatory_whitespace  ~
            secnd: parse_str_float      ~
                   mandatory_whitespace  ~
            third: parse_str_float      ~
                   whitespace  ,
            || { EntityValue::Vec3([first, secnd, third]) }
        ) |
        chain!(
            first: parse_str_int        ~
                   mandatory_whitespace ~
            secnd: parse_str_int        ~
                   whitespace ,
            || { EntityValue::IVec2([first, secnd]) }
        ) |
        chain!(
            first: parse_str_float      ~
                   mandatory_whitespace ~
            secnd: parse_str_float      ~
                   whitespace ,
            || { EntityValue::Vec2([first, secnd]) }
        ) |
        map!(
            parse_str_int,
            EntityValue::Int
        ) |
        map!(
            parse_str_float,
            EntityValue::Float
        ) |
        map!(
            take_s_until!("\""),
            EntityValue::Text
        )
    )
}

fn parse_entity(i: &[u8]) -> IResult<&[u8], Entity> {
    use nom::{Err, ErrorKind};

    let (rest, keyvals) = itry!(
        chain!(i,
                    whitespace ~
                    char!('{')  ~
                    whitespace  ~
            params: many0!(
                chain!(
                           char!('\"')        ~
                    name:  take_s_until!("\"")  ~
                           char!('\"')        ~
                           whitespace         ~
                           char!('\"')        ~
                    value: parse_entity_value ~
                           char!('\"')        ~
                           whitespace         ,
                    || { (name, value) }
                )
            )                   ~
                    char!('}')  ,
            || {
                params
            }
        )
    );

    fn byte_arrays_and_val_to_entity(
        v: Vec<(String, EntityValue)>
    ) -> Entity {
        Entity {
            parameters: v.into_iter().collect()
        }
    }

    Done(
        rest,
        byte_arrays_and_val_to_entity(keyvals)
    )
}

named! {
    parse_texture<Texture>,
    chain!(
        name:     take_s!(64) ~
        flags:    le_i32      ~
        contents: le_i32      ,
        || {
            Texture {
                name: name,
                flags: flags,
                contents: contents,
            }
        }
    )
}

named! {
    parse_plane<Plane>,
    chain!(
        normal:   parse_vec3 ~
        distance: le_f32     ,
        || {
            Plane {
                normal: normal,
                distance: distance,
            }
        }
    )
}

named! {
    parse_raw_node<RawNode>,
    chain!(
        plane:   le_i32       ~
        children: pair!(
            le_i32,
            le_i32
        )                     ~
        min:     parse_ivec3  ~
        max:     parse_ivec3  ,
        || {
            RawNode {
                plane: plane,
                children: children,
                min: min,
                max: max,
            }
        }
    )
}

named! {
    parse_vertex<Vertex>,
    chain!(
        position:        parse_vec3     ~
        surface_coords:  parse_vec2     ~
        lightmap_coords: parse_vec2     ~
        normal:          parse_vec3     ~
        color:           take_exact!(4) ,
        || {
            Vertex {
                position: position,
                surface_coords: surface_coords,
                lightmap_coords: lightmap_coords,
                normal: normal,
                color: color,
            }
        }
    )
}

named! {
    parse_raw_leaf<RawLeaf>,
    chain!(
        visdata_cluster:  le_i32      ~
        areaportal_area:  le_i32      ~
        min:              parse_ivec3 ~
        max:              parse_ivec3 ~
        first_leaf_face:  le_i32      ~
        num_leaf_faces:   le_i32      ~
        first_leaf_brush: le_i32      ~
        num_leaf_brushes: le_i32      ,
        || {
            RawLeaf {
                visdata_cluster: visdata_cluster,
                areaportal_area: areaportal_area,
                min: min,
                max: max,
                first_leaf_face: first_leaf_face,
                num_leaf_faces: num_leaf_faces,
                first_leaf_brush: first_leaf_brush,
                num_leaf_brushes: num_leaf_brushes,
            }
        }
    )
}

named! {
    parse_raw_leaf_face<RawLeafFace>,
    chain!(
        index: le_i32 ,
        || {
            RawLeafFace {
                index: index,
            }
        }
    )
}

named! {
    parse_raw_leaf_brush<RawLeafBrush>,
    chain!(
        index: le_i32 ,
        || {
            RawLeafBrush {
                index: index,
            }
        }
    )
}

named! {
    parse_raw_model<RawModel>,
    chain!(
        min:         parse_vec3 ~
        max:         parse_vec3 ~
        first_face:  le_i32     ~
        num_faces:   le_i32     ~
        first_brush: le_i32     ~
        num_brushes: le_i32     ,
        || {
            RawModel {
                min: min,
                max: max,
                first_face: first_face,
                num_faces: num_faces,
                first_brush: first_brush,
                num_brushes: num_brushes,
            }
        }
    )
}

named! {
    parse_raw_brush<RawBrush>,
    chain!(
        first_brush_side: le_i32 ~
        num_brush_sides:  le_i32 ~
        texture_index:    le_i32 ,
        || {
            RawBrush {
                first_brush_side: first_brush_side,
                num_brush_sides: num_brush_sides,
                texture_index: texture_index,
            }
        }
    )
}

named! {
    parse_raw_brush_side<RawBrushSide>,
    chain!(
        plane_index:   le_i32 ~
        texture_index: le_i32 ,
        || {
            RawBrushSide {
                plane_index: plane_index,
                texture_index: texture_index,
            }
        }
    )
}

named! {
    parse_raw_mesh_vertex<RawMeshVertex>,
    chain!(
        offset: le_i32 ,
        || {
            RawMeshVertex {
                offset: offset,
            }
        }
    )
}

named! {
    parse_raw_effect<RawEffect>,
    chain!(
        name:        take_s!(64) ~
        brush_index: le_i32      ~
                     take!(4)    ,
        || {
            RawEffect {
                name: name,
                brush_index: brush_index,
            }
        }
    )
}

fn parse_face_type(i: &[u8]) -> IResult<&[u8], FaceType> {
    use nom::ErrorKind;
    let (rest, t) = itry!(le_i32(i));
    match t {
        1 => Done(rest, FaceType::Polygon),
        2 => Done(rest, FaceType::Patch),
        3 => Done(rest, FaceType::Mesh),
        4 => Done(rest, FaceType::Billboard),
        _ => Error(Err::Code(ErrorKind::Custom(0))),
    }
}

named! {
    parse_raw_face<RawFace>,
    chain!(
        texture_index:     le_i32          ~
        effect_index:      le_i32          ~
        face_type:         parse_face_type ~
        first_vertex:      le_i32          ~
        num_vertexes:      le_i32          ~
        first_mesh_vertex: le_i32          ~
        num_mesh_vertices: le_i32          ~
        lightmap_index:    le_i32          ~
        lightmap_start:    parse_ivec2     ~
        lightmap_size:     parse_ivec2     ~
        lightmap_origin:   parse_vec3      ~
        lightmap_vecs:     pair!(
            parse_vec3,
            parse_vec3
        )                                  ~
        normal:            parse_vec3      ~
        size:              parse_ivec2     ,
        || {
            RawFace {
                texture_index: texture_index,
                effect_index: effect_index,
                face_type: face_type,
                first_vertex: first_vertex,
                num_vertexes: num_vertexes,
                first_mesh_vertex: first_mesh_vertex,
                num_mesh_vertices: num_mesh_vertices,
                lightmap_index: lightmap_index,
                lightmap_start: lightmap_start,
                lightmap_size: lightmap_size,
                lightmap_origin: lightmap_origin,
                lightmap_vecs: lightmap_vecs,
                normal: normal,
                size: size,
            }
        }
    )
}

fn parse_visibility_data(i: &[u8]) -> IResult<&[u8], RawVisibilityData> {
    let (rest, (num_vectors, sizeof_vector)) = itry!(
        chain! {i,
            num_vectors:   le_i32 ~
            sizeof_vector: le_i32 ,
            || { (num_vectors, sizeof_vector) }
        }
    );

    let size = num_vectors * sizeof_vector;
    let (rest, raw) = itry!(take!(rest, size));

    Done(rest, RawVisibilityData {
        num_vectors: num_vectors,
        sizeof_vector: sizeof_vector,
        raw_bytes: raw.into_iter().map(|&a| a).collect::<Vec<_>>(),
    })
}

fn parse_lightmap(i: &[u8]) -> IResult<&[u8], Lightmap> {
    let mut out = [[[0u8, 0u8, 0u8]; 128]; 128];
    let mut rest = i;
    for x in 0..128 {
        for y in 0..128 {
            let (a, b) = itry!(take_exact!(rest, 3));
            rest = a;
            out[y][x] = b;
        }
    }
    Done(rest, Lightmap { colors: out })
}

named! {
    parse_rotation_direction<RotationDirection>,
    chain!(
        phi:   le_u8 ~
        theta: le_u8 ,
        || {
            RotationDirection {
                phi: phi,
                theta: theta,
            }
        }
    )
}

named! {
    parse_light_volume<LightVolume>,
    chain!(
        ambient:     take_exact!(3)           ~
        directional: take_exact!(3)           ~
        direction:   parse_rotation_direction ,
        || {
            LightVolume {
                ambient: ambient,
                directional: directional,
                direction: direction,
            }
        }
    )
}

pub fn parse_raw_bsp(i: &[u8]) -> IResult<&[u8], RawBsp> {
    let (_, header) = itry!(directory_header(i));
    let (_, entities) = many1_from_header!(i,
        header.entities,
        parse_entity,
        Entity
    );
    let (_, textures) = get_from_header!(i,
        header.textures,
        parse_texture,
        Texture,
        72
    );
    let (_, planes) = get_from_header!(i,
        header.planes,
        parse_plane,
        Plane
    );
    let (_, nodes) = get_from_header!(i,
        header.nodes,
        parse_raw_node,
        RawNode
    );
    let (_, leaves) = get_from_header!(i,
        header.leaves,
        parse_raw_leaf,
        RawLeaf
    );
    let (_, leaf_faces) = get_from_header!(i,
        header.leaf_faces,
        parse_raw_leaf_face,
        RawLeafFace
    );
    let (_, leaf_brushes) = get_from_header!(i,
        header.leaf_brushes,
        parse_raw_leaf_brush,
        RawLeafBrush
    );
    let (_, models) = get_from_header!(i,
        header.models,
        parse_raw_model,
        RawModel
    );
    let (_, brushes) = get_from_header!(i,
        header.brushes,
        parse_raw_brush,
        RawBrush
    );
    let (_, brush_sides) = get_from_header!(i,
        header.brush_sides,
        parse_raw_brush_side,
        RawBrushSide
    );
    let (_, vertices) = get_from_header!(i,
        header.vertices,
        parse_vertex,
        Vertex
    );
    let (_, mesh_vertices) = get_from_header!(i,
        header.mesh_vertices,
        parse_raw_mesh_vertex,
        RawMeshVertex
    );
    let (_, effects) = get_from_header!(i,
        header.effects,
        parse_raw_effect,
        RawEffect,
        72
    );
    let (_, faces) = get_from_header!(i,
        header.faces,
        parse_raw_face,
        RawFace
    );
    let (_, light_maps) = get_from_header!(i,
        header.light_maps,
        parse_lightmap,
        Lightmap
    );
    let (_, light_volumes) = get_from_header!(i,
        header.light_volumes,
        parse_light_volume,
        LightVolume
    );
    let (_, visibility_data) = many1_from_header!(i,
        header.visibility_data,
        parse_visibility_data,
        RawVisibilityData
    );

    Done(&[], RawBsp {
        header: header,
        entities: entities,
        textures: textures,
        planes: planes,
        nodes: nodes,
        leaves: leaves,
        leaf_faces: leaf_faces,
        leaf_brushes: leaf_brushes,
        models: models,
        brushes: brushes,
        brush_sides: brush_sides,
        vertices: vertices,
        mesh_vertices: mesh_vertices,
        effects: effects,
        faces: faces,
        light_maps: light_maps,
        light_volumes: light_volumes,
        visibility_data: visibility_data,
    })
}
<commit_msg>Formatting<commit_after>use nom::{IResult,le_i32,le_u8};
use nom::Err;
use nom::IResult::*;
use directory_header::*;
use helpers::*;
use raw_bsp::*;
use std::str::from_utf8;

named! {
    parse_entity_value <EntityValue>,
    alt!(
        chain!(
                       tag!(b"*")    ~
            reference: parse_str_int ,
            || { EntityValue::ModelRef(reference) }
        ) |
        chain!(
                       tag!(b"t")    ~
            reference: parse_str_int ,
            || { EntityValue::TargetRef(reference) }
        ) |
        chain!(
            first: parse_str_int        ~
                   mandatory_whitespace ~
            secnd: parse_str_int        ~
                   mandatory_whitespace ~
            third: parse_str_int        ~
                   whitespace ,
            || { EntityValue::IVec3([first, secnd, third]) }
        ) |
        chain!(
            first: parse_str_float      ~
                   mandatory_whitespace ~
            secnd: parse_str_float      ~
                   mandatory_whitespace ~
            third: parse_str_float      ~
                   whitespace  ,
            || { EntityValue::Vec3([first, secnd, third]) }
        ) |
        chain!(
            first: parse_str_int        ~
                   mandatory_whitespace ~
            secnd: parse_str_int        ~
                   whitespace ,
            || { EntityValue::IVec2([first, secnd]) }
        ) |
        chain!(
            first: parse_str_float      ~
                   mandatory_whitespace ~
            secnd: parse_str_float      ~
                   whitespace ,
            || { EntityValue::Vec2([first, secnd]) }
        ) |
        map!(parse_str_int, EntityValue::Int) |
        map!(parse_str_float, EntityValue::Float) |
        map!(take_s_until!("\""), EntityValue::Text)
    )
}

fn parse_entity(i: &[u8]) -> IResult<&[u8], Entity> {
    use nom::{Err, ErrorKind};

    let (rest, keyvals) = itry!(
        chain!(i,
                    whitespace ~
                    char!('{')  ~
                    whitespace  ~
            params: many0!(
                chain!(
                           char!('\"')        ~
                    name:  take_s_until!("\"")  ~
                           char!('\"')        ~
                           whitespace         ~
                           char!('\"')        ~
                    value: parse_entity_value ~
                           char!('\"')        ~
                           whitespace         ,
                    || { (name, value) }
                )
            )                   ~
                    char!('}')  ,
            || {
                params
            }
        )
    );

    fn byte_arrays_and_val_to_entity(
        v: Vec<(String, EntityValue)>
    ) -> Entity {
        Entity {
            parameters: v.into_iter().collect()
        }
    }

    Done(
        rest,
        byte_arrays_and_val_to_entity(keyvals)
    )
}

named! {
    parse_texture<Texture>,
    chain!(
        name:     take_s!(64) ~
        flags:    le_i32      ~
        contents: le_i32      ,
        || {
            Texture {
                name: name,
                flags: flags,
                contents: contents,
            }
        }
    )
}

named! {
    parse_plane<Plane>,
    chain!(
        normal:   parse_vec3 ~
        distance: le_f32     ,
        || {
            Plane {
                normal: normal,
                distance: distance,
            }
        }
    )
}

named! {
    parse_raw_node<RawNode>,
    chain!(
        plane:   le_i32       ~
        children: pair!(
            le_i32,
            le_i32
        )                     ~
        min:     parse_ivec3  ~
        max:     parse_ivec3  ,
        || {
            RawNode {
                plane: plane,
                children: children,
                min: min,
                max: max,
            }
        }
    )
}

named! {
    parse_vertex<Vertex>,
    chain!(
        position:        parse_vec3     ~
        surface_coords:  parse_vec2     ~
        lightmap_coords: parse_vec2     ~
        normal:          parse_vec3     ~
        color:           take_exact!(4) ,
        || {
            Vertex {
                position: position,
                surface_coords: surface_coords,
                lightmap_coords: lightmap_coords,
                normal: normal,
                color: color,
            }
        }
    )
}

named! {
    parse_raw_leaf<RawLeaf>,
    chain!(
        visdata_cluster:  le_i32      ~
        areaportal_area:  le_i32      ~
        min:              parse_ivec3 ~
        max:              parse_ivec3 ~
        first_leaf_face:  le_i32      ~
        num_leaf_faces:   le_i32      ~
        first_leaf_brush: le_i32      ~
        num_leaf_brushes: le_i32      ,
        || {
            RawLeaf {
                visdata_cluster: visdata_cluster,
                areaportal_area: areaportal_area,
                min: min,
                max: max,
                first_leaf_face: first_leaf_face,
                num_leaf_faces: num_leaf_faces,
                first_leaf_brush: first_leaf_brush,
                num_leaf_brushes: num_leaf_brushes,
            }
        }
    )
}

named! {
    parse_raw_leaf_face<RawLeafFace>,
    chain!(
        index: le_i32 ,
        || {
            RawLeafFace {
                index: index,
            }
        }
    )
}

named! {
    parse_raw_leaf_brush<RawLeafBrush>,
    chain!(
        index: le_i32 ,
        || {
            RawLeafBrush {
                index: index,
            }
        }
    )
}

named! {
    parse_raw_model<RawModel>,
    chain!(
        min:         parse_vec3 ~
        max:         parse_vec3 ~
        first_face:  le_i32     ~
        num_faces:   le_i32     ~
        first_brush: le_i32     ~
        num_brushes: le_i32     ,
        || {
            RawModel {
                min: min,
                max: max,
                first_face: first_face,
                num_faces: num_faces,
                first_brush: first_brush,
                num_brushes: num_brushes,
            }
        }
    )
}

named! {
    parse_raw_brush<RawBrush>,
    chain!(
        first_brush_side: le_i32 ~
        num_brush_sides:  le_i32 ~
        texture_index:    le_i32 ,
        || {
            RawBrush {
                first_brush_side: first_brush_side,
                num_brush_sides: num_brush_sides,
                texture_index: texture_index,
            }
        }
    )
}

named! {
    parse_raw_brush_side<RawBrushSide>,
    chain!(
        plane_index:   le_i32 ~
        texture_index: le_i32 ,
        || {
            RawBrushSide {
                plane_index: plane_index,
                texture_index: texture_index,
            }
        }
    )
}

named! {
    parse_raw_mesh_vertex<RawMeshVertex>,
    chain!(
        offset: le_i32 ,
        || {
            RawMeshVertex {
                offset: offset,
            }
        }
    )
}

named! {
    parse_raw_effect<RawEffect>,
    chain!(
        name:        take_s!(64) ~
        brush_index: le_i32      ~
                     take!(4)    ,
        || {
            RawEffect {
                name: name,
                brush_index: brush_index,
            }
        }
    )
}

fn parse_face_type(i: &[u8]) -> IResult<&[u8], FaceType> {
    use nom::ErrorKind;
    let (rest, t) = itry!(le_i32(i));
    match t {
        1 => Done(rest, FaceType::Polygon),
        2 => Done(rest, FaceType::Patch),
        3 => Done(rest, FaceType::Mesh),
        4 => Done(rest, FaceType::Billboard),
        _ => Error(Err::Code(ErrorKind::Custom(0))),
    }
}

named! {
    parse_raw_face<RawFace>,
    chain!(
        texture_index:     le_i32          ~
        effect_index:      le_i32          ~
        face_type:         parse_face_type ~
        first_vertex:      le_i32          ~
        num_vertexes:      le_i32          ~
        first_mesh_vertex: le_i32          ~
        num_mesh_vertices: le_i32          ~
        lightmap_index:    le_i32          ~
        lightmap_start:    parse_ivec2     ~
        lightmap_size:     parse_ivec2     ~
        lightmap_origin:   parse_vec3      ~
        lightmap_vecs:     pair!(
            parse_vec3,
            parse_vec3
        )                                  ~
        normal:            parse_vec3      ~
        size:              parse_ivec2     ,
        || {
            RawFace {
                texture_index: texture_index,
                effect_index: effect_index,
                face_type: face_type,
                first_vertex: first_vertex,
                num_vertexes: num_vertexes,
                first_mesh_vertex: first_mesh_vertex,
                num_mesh_vertices: num_mesh_vertices,
                lightmap_index: lightmap_index,
                lightmap_start: lightmap_start,
                lightmap_size: lightmap_size,
                lightmap_origin: lightmap_origin,
                lightmap_vecs: lightmap_vecs,
                normal: normal,
                size: size,
            }
        }
    )
}

fn parse_visibility_data(i: &[u8]) -> IResult<&[u8], RawVisibilityData> {
    let (rest, (num_vectors, sizeof_vector)) = itry!(
        chain! {i,
            num_vectors:   le_i32 ~
            sizeof_vector: le_i32 ,
            || { (num_vectors, sizeof_vector) }
        }
    );

    let size = num_vectors * sizeof_vector;
    let (rest, raw) = itry!(take!(rest, size));

    Done(rest, RawVisibilityData {
        num_vectors: num_vectors,
        sizeof_vector: sizeof_vector,
        raw_bytes: raw.into_iter().map(|&a| a).collect::<Vec<_>>(),
    })
}

fn parse_lightmap(i: &[u8]) -> IResult<&[u8], Lightmap> {
    let mut out = [[[0u8, 0u8, 0u8]; 128]; 128];
    let mut rest = i;
    for x in 0..128 {
        for y in 0..128 {
            let (a, b) = itry!(take_exact!(rest, 3));
            rest = a;
            out[y][x] = b;
        }
    }
    Done(rest, Lightmap { colors: out })
}

named! {
    parse_rotation_direction<RotationDirection>,
    chain!(
        phi:   le_u8 ~
        theta: le_u8 ,
        || {
            RotationDirection {
                phi: phi,
                theta: theta,
            }
        }
    )
}

named! {
    parse_light_volume<LightVolume>,
    chain!(
        ambient:     take_exact!(3)           ~
        directional: take_exact!(3)           ~
        direction:   parse_rotation_direction ,
        || {
            LightVolume {
                ambient: ambient,
                directional: directional,
                direction: direction,
            }
        }
    )
}

pub fn parse_raw_bsp(i: &[u8]) -> IResult<&[u8], RawBsp> {
    let (_, header) = itry!(directory_header(i));
    let (_, entities) = many1_from_header!(i,
        header.entities,
        parse_entity,
        Entity
    );
    let (_, textures) = get_from_header!(i,
        header.textures,
        parse_texture,
        Texture,
        72
    );
    let (_, planes) = get_from_header!(i,
        header.planes,
        parse_plane,
        Plane
    );
    let (_, nodes) = get_from_header!(i,
        header.nodes,
        parse_raw_node,
        RawNode
    );
    let (_, leaves) = get_from_header!(i,
        header.leaves,
        parse_raw_leaf,
        RawLeaf
    );
    let (_, leaf_faces) = get_from_header!(i,
        header.leaf_faces,
        parse_raw_leaf_face,
        RawLeafFace
    );
    let (_, leaf_brushes) = get_from_header!(i,
        header.leaf_brushes,
        parse_raw_leaf_brush,
        RawLeafBrush
    );
    let (_, models) = get_from_header!(i,
        header.models,
        parse_raw_model,
        RawModel
    );
    let (_, brushes) = get_from_header!(i,
        header.brushes,
        parse_raw_brush,
        RawBrush
    );
    let (_, brush_sides) = get_from_header!(i,
        header.brush_sides,
        parse_raw_brush_side,
        RawBrushSide
    );
    let (_, vertices) = get_from_header!(i,
        header.vertices,
        parse_vertex,
        Vertex
    );
    let (_, mesh_vertices) = get_from_header!(i,
        header.mesh_vertices,
        parse_raw_mesh_vertex,
        RawMeshVertex
    );
    let (_, effects) = get_from_header!(i,
        header.effects,
        parse_raw_effect,
        RawEffect,
        72
    );
    let (_, faces) = get_from_header!(i,
        header.faces,
        parse_raw_face,
        RawFace
    );
    let (_, light_maps) = get_from_header!(i,
        header.light_maps,
        parse_lightmap,
        Lightmap
    );
    let (_, light_volumes) = get_from_header!(i,
        header.light_volumes,
        parse_light_volume,
        LightVolume
    );
    let (_, visibility_data) = many1_from_header!(i,
        header.visibility_data,
        parse_visibility_data,
        RawVisibilityData
    );

    Done(&[], RawBsp {
        header: header,
        entities: entities,
        textures: textures,
        planes: planes,
        nodes: nodes,
        leaves: leaves,
        leaf_faces: leaf_faces,
        leaf_brushes: leaf_brushes,
        models: models,
        brushes: brushes,
        brush_sides: brush_sides,
        vertices: vertices,
        mesh_vertices: mesh_vertices,
        effects: effects,
        faces: faces,
        light_maps: light_maps,
        light_volumes: light_volumes,
        visibility_data: visibility_data,
    })
}
<|endoftext|>


<commit_before>#![feature(globs)]
use super::events::*;
use super::net::*;
use super::net::RaftRpc;
use super::server::RaftServerState;
use super::server::{RaftStateTransition, NextState, Continue};
use super::server::{RaftNextState, RaftCandidate, RaftFollower};
use std::vec::Vec;
use uuid::{Uuid, UuidVersion, Version4Random};

// Rust doesn't support arbitrary methods being defined in another
// file... except with traits.
//
// I don't know how this copes with small private helper functions.
// Hopefully you won't need to declare them in the trait before you
// can define them in the impl Leader for RaftServerState.

pub trait Leader {
    fn leader_setup(&mut self) -> RaftStateTransition;
    fn leader_teardown(&mut self);
    fn leader_heartbeat(&mut self) -> RaftStateTransition;

    fn leader_append_entries_req(&mut self, req: AppendEntriesReq, chan: Sender<AppendEntriesRes>) -> RaftStateTransition;
    fn leader_append_entries_res(&mut self, res: AppendEntriesRes) -> RaftStateTransition;

    fn leader_vote_req(&mut self, req: VoteReq, chan: Sender<VoteRes>) -> RaftStateTransition;
    fn leader_vote_res(&mut self, res: VoteRes) -> RaftStateTransition;
}

impl Leader for RaftServerState {
    fn leader_setup(&mut self) -> RaftStateTransition {
        //send initial heartbeat to all followers
        self.leader_heartbeat();
        Continue
    }

    fn leader_teardown(&mut self) {
        self.new_state(RaftFollower);
    }

    fn leader_heartbeat(&mut self) -> RaftStateTransition {
        let aer = AppendEntriesReq{
                                term: self.log.term,
                                //TODO get more accurate term info from election or system var
                                prev_log_idx: self.log.idx,
                                prev_log_term: self.log.term,
                                commit_idx: self.log.idx,
                                leader_id: self.id,
                                uuid: Uuid::new(Version4Random).unwrap(),
                                entries: Vec::new()};

        self.peers.msg_all_peers(RpcARQ(aer));

        Continue
    }

    fn leader_append_entries_req(&mut self, req: AppendEntriesReq, chan: Sender<AppendEntriesRes>) -> RaftStateTransition {

        //From Raft spec:
        //If RPC request or response contains term T > currentTerm: set currentTerm = T,
        //convert to follower
        for entry in req.entries.iter() {
            if (entry.term > self.log.term) {
                self.log.term = entry.term;
                self.new_state(RaftFollower);
            }
        }

        //send entry to followers
        self.peers.msg_all_peers(RpcARQ(req));
        //add all peers to peers_to_confirm list
        self.peers_to_confirm = self.peers.get_peer_ids();

        Continue
    }

    fn leader_append_entries_res(&mut self, res: AppendEntriesRes) -> RaftStateTransition {

        if (res.term > self.log.term) {
            self.log.term = res.term;
            self.new_state(RaftFollower);
        }

        //listen for a conform (majority of peers) of the log
        if (self.peers_to_confirm.contains(&res.id) &&
            !self.peers_have_confirmed.contains(&res.id)) {
            self.peers_have_confirmed.push(res.id);
            if (self.peers_have_confirmed.len() >= self.peers_to_confirm.len()) {
                //conform reached
                //pass the log into application state machine
                //pass the log to client using RaftAppMsg<AppMsg>
            }
        }

        Continue
    }

    fn leader_vote_req(&mut self, req: VoteReq, chan: Sender<VoteRes>) -> RaftStateTransition {
        Continue
    }

    fn leader_vote_res(&mut self, res: VoteRes) -> RaftStateTransition  {
        Continue
    }
}

<commit_msg>More on logic.<commit_after>#![feature(globs)]
use super::events::*;
use super::net::*;
use super::net::RaftRpc;
use super::server::RaftServerState;
use super::server::{RaftStateTransition, NextState, Continue};
use super::server::{RaftNextState, RaftCandidate, RaftFollower};
use std::vec::Vec;
use uuid::{Uuid, UuidVersion, Version4Random};

// Rust doesn't support arbitrary methods being defined in another
// file... except with traits.
//
// I don't know how this copes with small private helper functions.
// Hopefully you won't need to declare them in the trait before you
// can define them in the impl Leader for RaftServerState.

pub trait Leader {
    fn leader_setup(&mut self) -> RaftStateTransition;
    fn leader_teardown(&mut self);
    fn leader_heartbeat(&mut self) -> RaftStateTransition;

    fn leader_append_entries_req(&mut self, req: AppendEntriesReq, chan: Sender<AppendEntriesRes>) -> RaftStateTransition;
    fn leader_append_entries_res(&mut self, res: AppendEntriesRes) -> RaftStateTransition;

    fn leader_vote_req(&mut self, req: VoteReq, chan: Sender<VoteRes>) -> RaftStateTransition;
    fn leader_vote_res(&mut self, res: VoteRes) -> RaftStateTransition;
}

impl Leader for RaftServerState {
    fn leader_setup(&mut self) -> RaftStateTransition {
        //send initial heartbeat to all followers
        self.leader_heartbeat();
        Continue
    }

    fn leader_teardown(&mut self) {
        self.new_state(RaftFollower);
    }

    fn leader_heartbeat(&mut self) -> RaftStateTransition {
        let aer = AppendEntriesReq{
                                term: self.log.term,
                                //TODO get more accurate term info from election or system var
                                prev_log_idx: self.log.idx,
                                prev_log_term: self.log.term,
                                commit_idx: self.log.idx,
                                leader_id: self.id,
                                uuid: Uuid::new(Version4Random).unwrap(),
                                entries: Vec::new()};

        self.peers.msg_all_peers(RpcARQ(aer));

        Continue
    }

    fn leader_append_entries_req(&mut self, req: AppendEntriesReq, chan: Sender<AppendEntriesRes>) -> RaftStateTransition {
        if req.commit_idx > self.last_applied {
            self.last_applied += 1;
            // TODO:
            // self.to_app_sm.send(log(lastApplied))
        }

        //Raft paper:
        //If RPC request or response contains term T > currentTerm: set currentTerm = T,
        //convert to follower
        for entry in req.entries.iter() {
            if (entry.term > self.log.term) {
                self.log.term = entry.term;
                self.new_state(RaftFollower);
            }
        }

        //send entry to followers
        self.peers.msg_all_peers(RpcARQ(req));
        //add all peers to peers_to_confirm list
        self.peers_to_confirm = self.peers.get_peer_ids();

        Continue
    }

    fn leader_append_entries_res(&mut self, res: AppendEntriesRes) -> RaftStateTransition {

        if (res.term > self.log.term) {
            self.log.term = res.term;
            self.new_state(RaftFollower);
        }

        //listen for a conform (majority of peers) of the log
        if (self.peers_to_confirm.contains(&res.id) &&
            !self.peers_have_confirmed.contains(&res.id)) {
            self.peers_have_confirmed.push(res.id);
            if (self.peers_have_confirmed.len() >= self.peers_to_confirm.len()) {
                //conform reached
                //pass the log into application state machine
                //pass the log to client using RaftAppMsg<AppMsg>
            }
        }

        Continue
    }

    fn leader_vote_req(&mut self, req: VoteReq, chan: Sender<VoteRes>) -> RaftStateTransition {
        Continue
    }

    fn leader_vote_res(&mut self, res: VoteRes) -> RaftStateTransition  {
        Continue
    }
}

<|endoftext|>


<commit_before>//! Server Responses
//!
//! These are responses sent by a `hyper::Server` to clients, after
//! receiving a request.
use std::any::{Any, TypeId};
use std::marker::PhantomData;
use std::mem;
use std::io::{self, Write};
use std::ptr;

use time::now_utc;

use header;
use http::h1::{CR, LF, LINE_ENDING, HttpWriter};
use http::h1::HttpWriter::{ThroughWriter, ChunkedWriter, SizedWriter, EmptyWriter};
use status;
use net::{Fresh, Streaming};
use version;

/// The outgoing half for a Tcp connection, created by a `Server` and given to a `Handler`.
///
/// The default `StatusCode` for a `Response` is `200 OK`.
///
/// There is a `Drop` implementation for `Response` that will automatically
/// write the head and flush the body, if the handler has not already done so,
/// so that the server doesn't accidentally leave dangling requests.
#[derive(Debug)]
pub struct Response<'a, W: Any = Fresh> {
    /// The HTTP version of this response.
    pub version: version::HttpVersion,
    // Stream the Response is writing to, not accessible through UnwrittenResponse
    body: HttpWriter<&'a mut (Write + 'a)>,
    // The status code for the request.
    status: status::StatusCode,
    // The outgoing headers on this response.
    headers: &'a mut header::Headers,

    _writing: PhantomData<W>
}

impl<'a, W: Any> Response<'a, W> {
    /// The status of this response.
    #[inline]
    pub fn status(&self) -> status::StatusCode { self.status }

    /// The headers of this response.
    #[inline]
    pub fn headers(&self) -> &header::Headers { &*self.headers }

    /// Construct a Response from its constituent parts.
    #[inline]
    pub fn construct(version: version::HttpVersion,
                     body: HttpWriter<&'a mut (Write + 'a)>,
                     status: status::StatusCode,
                     headers: &'a mut header::Headers) -> Response<'a, Fresh> {
        Response {
            status: status,
            version: version,
            body: body,
            headers: headers,
            _writing: PhantomData,
        }
    }

    /// Deconstruct this Response into its constituent parts.
    #[inline]
    pub fn deconstruct(self) -> (version::HttpVersion, HttpWriter<&'a mut (Write + 'a)>,
                                 status::StatusCode, &'a mut header::Headers) {
        unsafe {
            let parts = (
                self.version,
                ptr::read(&self.body),
                self.status,
                ptr::read(&self.headers)
            );
            mem::forget(self);
            parts
        }
    }

    fn write_head(&mut self) -> io::Result<Body> {
        debug!("writing head: {:?} {:?}", self.version, self.status);
        try!(write!(&mut self.body, "{} {}{}{}", self.version, self.status,
            CR as char, LF as char));

        if !self.headers.has::<header::Date>() {
            self.headers.set(header::Date(header::HttpDate(now_utc())));
        }

        let body_type = match self.status {
            status::StatusCode::NoContent | status::StatusCode::NotModified => Body::Empty,
            c if c.class() == status::StatusClass::Informational => Body::Empty,
            _ => if let Some(cl) = self.headers.get::<header::ContentLength>() {
                Body::Sized(**cl)
            } else {
                Body::Chunked
            }
        };

        // can't do in match above, thanks borrowck
        if body_type == Body::Chunked {
            let encodings = match self.headers.get_mut::<header::TransferEncoding>() {
                Some(&mut header::TransferEncoding(ref mut encodings)) => {
                    //TODO: check if chunked is already in encodings. use HashSet?
                    encodings.push(header::Encoding::Chunked);
                    false
                },
                None => true
            };

            if encodings {
                self.headers.set::<header::TransferEncoding>(
                    header::TransferEncoding(vec![header::Encoding::Chunked]))
            }
        }

        debug!("headers [\n{:?}]", self.headers);
        try!(write!(&mut self.body, "{}", self.headers));
        try!(write!(&mut self.body, "{}", LINE_ENDING));

        Ok(body_type)
    }
}

impl<'a> Response<'a, Fresh> {
    /// Creates a new Response that can be used to write to a network stream.
    #[inline]
    pub fn new(stream: &'a mut (Write + 'a), headers: &'a mut header::Headers) ->
            Response<'a, Fresh> {
        Response {
            status: status::StatusCode::Ok,
            version: version::HttpVersion::Http11,
            headers: headers,
            body: ThroughWriter(stream),
            _writing: PhantomData,
        }
    }

    /// Writes the body and ends the response.
    ///
    /// This is a shortcut method for when you have a response with a fixed
    /// size, and would only need a single `write` call normally.
    ///
    /// # Example
    ///
    /// ```
    /// # use hyper::server::Response;
    /// fn handler(res: Response) {
    ///     res.send(b"Hello World!").unwrap();
    /// }
    /// ```
    ///
    /// The above is the same, but shorter, than the longer:
    ///
    /// ```
    /// # use hyper::server::Response;
    /// use std::io::Write;
    /// use hyper::header::ContentLength;
    /// fn handler(mut res: Response) {
    ///     let body = b"Hello World!";
    ///     res.headers_mut().set(ContentLength(body.len() as u64));
    ///     let mut res = res.start().unwrap();
    ///     res.write_all(body).unwrap();
    /// }
    /// ```
    #[inline]
    pub fn send(mut self, body: &[u8]) -> io::Result<()> {
        self.headers.set(header::ContentLength(body.len() as u64));
        let mut stream = try!(self.start());
        try!(stream.write_all(body));
        stream.end()
    }

    /// Consume this Response<Fresh>, writing the Headers and Status and
    /// creating a Response<Streaming>
    pub fn start(mut self) -> io::Result<Response<'a, Streaming>> {
        let body_type = try!(self.write_head());
        let (version, body, status, headers) = self.deconstruct();
        let stream = match body_type {
            Body::Chunked => ChunkedWriter(body.into_inner()),
            Body::Sized(len) => SizedWriter(body.into_inner(), len),
            Body::Empty => EmptyWriter(body.into_inner()),
        };

        // "copy" to change the phantom type
        Ok(Response {
            version: version,
            body: stream,
            status: status,
            headers: headers,
            _writing: PhantomData,
        })
    }
    /// Get a mutable reference to the status.
    #[inline]
    pub fn status_mut(&mut self) -> &mut status::StatusCode { &mut self.status }

    /// Get a mutable reference to the Headers.
    #[inline]
    pub fn headers_mut(&mut self) -> &mut header::Headers { self.headers }
}

impl<'a> Response<'a, Streaming> {
    /// Flushes all writing of a response to the client.
    #[inline]
    pub fn end(self) -> io::Result<()> {
        trace!("ending");
        let (_, body, _, _) = self.deconstruct();
        try!(body.end());
        Ok(())
    }
}

impl<'a> Write for Response<'a, Streaming> {
    #[inline]
    fn write(&mut self, msg: &[u8]) -> io::Result<usize> {
        debug!("write {:?} bytes", msg.len());
        self.body.write(msg)
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        self.body.flush()
    }
}

#[derive(PartialEq)]
enum Body {
    Chunked,
    Sized(u64),
    Empty,
}

impl<'a, T: Any> Drop for Response<'a, T> {
    fn drop(&mut self) {
        if TypeId::of::<T>() == TypeId::of::<Fresh>() {
            let mut body = match self.write_head() {
                Ok(Body::Chunked) => ChunkedWriter(self.body.get_mut()),
                Ok(Body::Sized(len)) => SizedWriter(self.body.get_mut(), len),
                Ok(Body::Empty) => EmptyWriter(self.body.get_mut()),
                Err(e) => {
                    debug!("error dropping request: {:?}", e);
                    return;
                }
            };
            end(&mut body);
        } else {
            end(&mut self.body);
        };

        #[inline]
        fn end<W: Write>(w: &mut W) {
            match w.write(&[]) {
                Ok(_) => match w.flush() {
                    Ok(_) => debug!("drop successful"),
                    Err(e) => debug!("error dropping request: {:?}", e)
                },
                Err(e) => debug!("error dropping request: {:?}", e)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use header::Headers;
    use mock::MockStream;
    use super::Response;

    macro_rules! lines {
        ($s:ident = $($line:pat),+) => ({
            let s = String::from_utf8($s.write).unwrap();
            let mut lines = s.split_terminator("\r\n");

            $(
                match lines.next() {
                    Some($line) => (),
                    other => panic!("line mismatch: {:?} != {:?}", other, stringify!($line))
                }
            )+

            assert_eq!(lines.next(), None);
        })
    }

    #[test]
    fn test_fresh_start() {
        let mut headers = Headers::new();
        let mut stream = MockStream::new();
        {
            let res = Response::new(&mut stream, &mut headers);
            res.start().unwrap().deconstruct();
        }

        lines! { stream =
            "HTTP/1.1 200 OK",
            _date,
            _transfer_encoding,
            ""
        }
    }

    #[test]
    fn test_streaming_end() {
        let mut headers = Headers::new();
        let mut stream = MockStream::new();
        {
            let res = Response::new(&mut stream, &mut headers);
            res.start().unwrap().end().unwrap();
        }

        lines! { stream =
            "HTTP/1.1 200 OK",
            _date,
            _transfer_encoding,
            "",
            "0",
            "" // empty zero body
        }
    }

    #[test]
    fn test_fresh_drop() {
        use status::StatusCode;
        let mut headers = Headers::new();
        let mut stream = MockStream::new();
        {
            let mut res = Response::new(&mut stream, &mut headers);
            *res.status_mut() = StatusCode::NotFound;
        }

        lines! { stream =
            "HTTP/1.1 404 Not Found",
            _date,
            _transfer_encoding,
            "",
            "0",
            "" // empty zero body
        }
    }

    #[test]
    fn test_streaming_drop() {
        use std::io::Write;
        use status::StatusCode;
        let mut headers = Headers::new();
        let mut stream = MockStream::new();
        {
            let mut res = Response::new(&mut stream, &mut headers);
            *res.status_mut() = StatusCode::NotFound;
            let mut stream = res.start().unwrap();
            stream.write_all(b"foo").unwrap();
        }

        lines! { stream =
            "HTTP/1.1 404 Not Found",
            _date,
            _transfer_encoding,
            "",
            "3",
            "foo",
            "0",
            "" // empty zero body
        }
    }

    #[test]
    fn test_no_content() {
        use std::io::Write;
        use status::StatusCode;
        let mut headers = Headers::new();
        let mut stream = MockStream::new();
        {
            let mut res = Response::new(&mut stream, &mut headers);
            *res.status_mut() = StatusCode::NoContent;
            res.start().unwrap();
        }

        lines! { stream =
            "HTTP/1.1 204 No Content",
            _date,
            ""
        }
    }
}
<commit_msg>fix(response): respond with a 500 if a handler panics<commit_after>//! Server Responses
//!
//! These are responses sent by a `hyper::Server` to clients, after
//! receiving a request.
use std::any::{Any, TypeId};
use std::marker::PhantomData;
use std::mem;
use std::io::{self, Write};
use std::ptr;
use std::thread;

use time::now_utc;

use header;
use http::h1::{CR, LF, LINE_ENDING, HttpWriter};
use http::h1::HttpWriter::{ThroughWriter, ChunkedWriter, SizedWriter, EmptyWriter};
use status;
use net::{Fresh, Streaming};
use version;

/// The outgoing half for a Tcp connection, created by a `Server` and given to a `Handler`.
///
/// The default `StatusCode` for a `Response` is `200 OK`.
///
/// There is a `Drop` implementation for `Response` that will automatically
/// write the head and flush the body, if the handler has not already done so,
/// so that the server doesn't accidentally leave dangling requests.
#[derive(Debug)]
pub struct Response<'a, W: Any = Fresh> {
    /// The HTTP version of this response.
    pub version: version::HttpVersion,
    // Stream the Response is writing to, not accessible through UnwrittenResponse
    body: HttpWriter<&'a mut (Write + 'a)>,
    // The status code for the request.
    status: status::StatusCode,
    // The outgoing headers on this response.
    headers: &'a mut header::Headers,

    _writing: PhantomData<W>
}

impl<'a, W: Any> Response<'a, W> {
    /// The status of this response.
    #[inline]
    pub fn status(&self) -> status::StatusCode { self.status }

    /// The headers of this response.
    #[inline]
    pub fn headers(&self) -> &header::Headers { &*self.headers }

    /// Construct a Response from its constituent parts.
    #[inline]
    pub fn construct(version: version::HttpVersion,
                     body: HttpWriter<&'a mut (Write + 'a)>,
                     status: status::StatusCode,
                     headers: &'a mut header::Headers) -> Response<'a, Fresh> {
        Response {
            status: status,
            version: version,
            body: body,
            headers: headers,
            _writing: PhantomData,
        }
    }

    /// Deconstruct this Response into its constituent parts.
    #[inline]
    pub fn deconstruct(self) -> (version::HttpVersion, HttpWriter<&'a mut (Write + 'a)>,
                                 status::StatusCode, &'a mut header::Headers) {
        unsafe {
            let parts = (
                self.version,
                ptr::read(&self.body),
                self.status,
                ptr::read(&self.headers)
            );
            mem::forget(self);
            parts
        }
    }

    fn write_head(&mut self) -> io::Result<Body> {
        debug!("writing head: {:?} {:?}", self.version, self.status);
        try!(write!(&mut self.body, "{} {}{}{}", self.version, self.status,
            CR as char, LF as char));

        if !self.headers.has::<header::Date>() {
            self.headers.set(header::Date(header::HttpDate(now_utc())));
        }

        let body_type = match self.status {
            status::StatusCode::NoContent | status::StatusCode::NotModified => Body::Empty,
            c if c.class() == status::StatusClass::Informational => Body::Empty,
            _ => if let Some(cl) = self.headers.get::<header::ContentLength>() {
                Body::Sized(**cl)
            } else {
                Body::Chunked
            }
        };

        // can't do in match above, thanks borrowck
        if body_type == Body::Chunked {
            let encodings = match self.headers.get_mut::<header::TransferEncoding>() {
                Some(&mut header::TransferEncoding(ref mut encodings)) => {
                    //TODO: check if chunked is already in encodings. use HashSet?
                    encodings.push(header::Encoding::Chunked);
                    false
                },
                None => true
            };

            if encodings {
                self.headers.set::<header::TransferEncoding>(
                    header::TransferEncoding(vec![header::Encoding::Chunked]))
            }
        }

        debug!("headers [\n{:?}]", self.headers);
        try!(write!(&mut self.body, "{}", self.headers));
        try!(write!(&mut self.body, "{}", LINE_ENDING));

        Ok(body_type)
    }
}

impl<'a> Response<'a, Fresh> {
    /// Creates a new Response that can be used to write to a network stream.
    #[inline]
    pub fn new(stream: &'a mut (Write + 'a), headers: &'a mut header::Headers) ->
            Response<'a, Fresh> {
        Response {
            status: status::StatusCode::Ok,
            version: version::HttpVersion::Http11,
            headers: headers,
            body: ThroughWriter(stream),
            _writing: PhantomData,
        }
    }

    /// Writes the body and ends the response.
    ///
    /// This is a shortcut method for when you have a response with a fixed
    /// size, and would only need a single `write` call normally.
    ///
    /// # Example
    ///
    /// ```
    /// # use hyper::server::Response;
    /// fn handler(res: Response) {
    ///     res.send(b"Hello World!").unwrap();
    /// }
    /// ```
    ///
    /// The above is the same, but shorter, than the longer:
    ///
    /// ```
    /// # use hyper::server::Response;
    /// use std::io::Write;
    /// use hyper::header::ContentLength;
    /// fn handler(mut res: Response) {
    ///     let body = b"Hello World!";
    ///     res.headers_mut().set(ContentLength(body.len() as u64));
    ///     let mut res = res.start().unwrap();
    ///     res.write_all(body).unwrap();
    /// }
    /// ```
    #[inline]
    pub fn send(mut self, body: &[u8]) -> io::Result<()> {
        self.headers.set(header::ContentLength(body.len() as u64));
        let mut stream = try!(self.start());
        try!(stream.write_all(body));
        stream.end()
    }

    /// Consume this Response<Fresh>, writing the Headers and Status and
    /// creating a Response<Streaming>
    pub fn start(mut self) -> io::Result<Response<'a, Streaming>> {
        let body_type = try!(self.write_head());
        let (version, body, status, headers) = self.deconstruct();
        let stream = match body_type {
            Body::Chunked => ChunkedWriter(body.into_inner()),
            Body::Sized(len) => SizedWriter(body.into_inner(), len),
            Body::Empty => EmptyWriter(body.into_inner()),
        };

        // "copy" to change the phantom type
        Ok(Response {
            version: version,
            body: stream,
            status: status,
            headers: headers,
            _writing: PhantomData,
        })
    }
    /// Get a mutable reference to the status.
    #[inline]
    pub fn status_mut(&mut self) -> &mut status::StatusCode { &mut self.status }

    /// Get a mutable reference to the Headers.
    #[inline]
    pub fn headers_mut(&mut self) -> &mut header::Headers { self.headers }
}

impl<'a> Response<'a, Streaming> {
    /// Flushes all writing of a response to the client.
    #[inline]
    pub fn end(self) -> io::Result<()> {
        trace!("ending");
        let (_, body, _, _) = self.deconstruct();
        try!(body.end());
        Ok(())
    }
}

impl<'a> Write for Response<'a, Streaming> {
    #[inline]
    fn write(&mut self, msg: &[u8]) -> io::Result<usize> {
        debug!("write {:?} bytes", msg.len());
        self.body.write(msg)
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        self.body.flush()
    }
}

#[derive(PartialEq)]
enum Body {
    Chunked,
    Sized(u64),
    Empty,
}

impl<'a, T: Any> Drop for Response<'a, T> {
    fn drop(&mut self) {
        if TypeId::of::<T>() == TypeId::of::<Fresh>() {
            if thread::panicking() {
                self.status = status::StatusCode::InternalServerError;
            }

            let mut body = match self.write_head() {
                Ok(Body::Chunked) => ChunkedWriter(self.body.get_mut()),
                Ok(Body::Sized(len)) => SizedWriter(self.body.get_mut(), len),
                Ok(Body::Empty) => EmptyWriter(self.body.get_mut()),
                Err(e) => {
                    debug!("error dropping request: {:?}", e);
                    return;
                }
            };
            end(&mut body);
        } else {
            end(&mut self.body);
        };

        #[inline]
        fn end<W: Write>(w: &mut W) {
            match w.write(&[]) {
                Ok(_) => match w.flush() {
                    Ok(_) => debug!("drop successful"),
                    Err(e) => debug!("error dropping request: {:?}", e)
                },
                Err(e) => debug!("error dropping request: {:?}", e)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use header::Headers;
    use mock::MockStream;
    use super::Response;

    macro_rules! lines {
        ($s:ident = $($line:pat),+) => ({
            let s = String::from_utf8($s.write).unwrap();
            let mut lines = s.split_terminator("\r\n");

            $(
                match lines.next() {
                    Some($line) => (),
                    other => panic!("line mismatch: {:?} != {:?}", other, stringify!($line))
                }
            )+

            assert_eq!(lines.next(), None);
        })
    }

    #[test]
    fn test_fresh_start() {
        let mut headers = Headers::new();
        let mut stream = MockStream::new();
        {
            let res = Response::new(&mut stream, &mut headers);
            res.start().unwrap().deconstruct();
        }

        lines! { stream =
            "HTTP/1.1 200 OK",
            _date,
            _transfer_encoding,
            ""
        }
    }

    #[test]
    fn test_streaming_end() {
        let mut headers = Headers::new();
        let mut stream = MockStream::new();
        {
            let res = Response::new(&mut stream, &mut headers);
            res.start().unwrap().end().unwrap();
        }

        lines! { stream =
            "HTTP/1.1 200 OK",
            _date,
            _transfer_encoding,
            "",
            "0",
            "" // empty zero body
        }
    }

    #[test]
    fn test_fresh_drop() {
        use status::StatusCode;
        let mut headers = Headers::new();
        let mut stream = MockStream::new();
        {
            let mut res = Response::new(&mut stream, &mut headers);
            *res.status_mut() = StatusCode::NotFound;
        }

        lines! { stream =
            "HTTP/1.1 404 Not Found",
            _date,
            _transfer_encoding,
            "",
            "0",
            "" // empty zero body
        }
    }

    #[test]
    fn test_fresh_drop_panicing() {
        use std::thread;
        use std::sync::{Arc, Mutex};

        use status::StatusCode;

        let stream = MockStream::new();
        let stream = Arc::new(Mutex::new(stream));
        let inner_stream = stream.clone();
        let join_handle = thread::spawn(move || {
            let mut headers = Headers::new();
            let mut stream = inner_stream.lock().unwrap();
            let mut res = Response::new(&mut *stream, &mut headers);
            *res.status_mut() = StatusCode::NotFound;

            panic!("inside")
        });

        assert!(join_handle.join().is_err());

        let stream = match stream.lock() {
            Err(poisoned) => poisoned.into_inner().clone(),
            Ok(_) => unreachable!()
        };

        lines! { stream =
            "HTTP/1.1 500 Internal Server Error",
            _date,
            _transfer_encoding,
            "",
            "0",
            "" // empty zero body
        }
    }

    #[test]
    fn test_streaming_drop() {
        use std::io::Write;
        use status::StatusCode;
        let mut headers = Headers::new();
        let mut stream = MockStream::new();
        {
            let mut res = Response::new(&mut stream, &mut headers);
            *res.status_mut() = StatusCode::NotFound;
            let mut stream = res.start().unwrap();
            stream.write_all(b"foo").unwrap();
        }

        lines! { stream =
            "HTTP/1.1 404 Not Found",
            _date,
            _transfer_encoding,
            "",
            "3",
            "foo",
            "0",
            "" // empty zero body
        }
    }

    #[test]
    fn test_no_content() {
        use std::io::Write;
        use status::StatusCode;
        let mut headers = Headers::new();
        let mut stream = MockStream::new();
        {
            let mut res = Response::new(&mut stream, &mut headers);
            *res.status_mut() = StatusCode::NoContent;
            res.start().unwrap();
        }

        lines! { stream =
            "HTTP/1.1 204 No Content",
            _date,
            ""
        }
    }
}
<|endoftext|>


<commit_before>use primal_bit::BitVec;
use std::cmp;

use crate::wheel;

pub mod primes;
mod presieve;

/// A heavily optimised prime sieve.
///
/// This is a streaming segmented sieve, meaning it sieves numbers in
/// intervals, extracting whatever it needs and discarding it. See
/// `Sieve` for a wrapper that caches the information to allow for
/// repeated queries, at the cost of *O(limit)* memory use.
///
/// This uses *O(sqrt(limit))* memory, and is designed to be as
/// cache-friendly as possible. `StreamingSieve` should be used for
/// one-off calls, or simple linear iteration.
///
/// The design is *heavily* inspired/adopted from Kim Walisch's
/// [primesieve](http://primesieve.org/), and has similar speed
/// (around 5-20% slower).
///
/// # Examples
///
/// ```rust
/// let count = primal::StreamingSieve::prime_pi(123456);
/// println!("𝜋(123456) = {}", count);
/// ```
#[derive(Debug)]
pub struct StreamingSieve {
    small: Option<crate::Sieve>,
    sieve: BitVec,
    primes: Vec<wheel::State<wheel::Wheel210>>,
    small_primes: Vec<wheel::State<wheel::Wheel30>>,
    large_primes: Vec<wheel::State<wheel::Wheel210>>,
    presieve: presieve::Presieve,

    low: usize,
    current: usize,
    limit: usize,
}

const CACHE: usize = 32 << 10;
const SEG_ELEMS: usize = 8 * CACHE;
const SEG_LEN: usize = SEG_ELEMS * wheel::BYTE_MODULO / wheel::BYTE_SIZE;

fn isqrt(x: usize) -> usize {
    (x as f64).sqrt() as usize
}

impl StreamingSieve {
    /// Create a new instance of the streaming sieve that will
    /// correctly progressively filter primes up to `limit`.
    pub(crate) fn new(limit: usize) -> StreamingSieve {
        let low = 0;

        let elems = cmp::min(wheel::bits_for(limit), SEG_ELEMS);
        let presieve = presieve::Presieve::new(elems);
        let current = presieve.smallest_unincluded_prime();

        let small = if limit < current * current {
            None
        } else {
            Some(crate::Sieve::new(isqrt(limit) + 1))
        };

        StreamingSieve {
            small,
            sieve: BitVec::from_elem(elems, true),
            primes: vec![],
            small_primes: vec![],
            large_primes: vec![],
            presieve,

            low,
            current,
            limit
        }
    }
    fn split_index(&self, idx: usize) -> (usize, usize) {
        let len = SEG_ELEMS;
        (idx / len,idx % len)
    }
    fn index_for(&self, n: usize) -> (bool, usize, usize) {
        let (b, idx) = wheel::bit_index(n);
        let (base, tweak) = self.split_index(idx);
        (b, base, tweak)
    }

    /// Count the number of primes upto and including `n`, that is, 𝜋,
    /// the [prime counting
    /// function](https://en.wikipedia.org/wiki/Prime-counting_function).
    ///
    /// # Examples
    ///
    /// ```rust
    /// assert_eq!(primal::StreamingSieve::prime_pi(10), 4);
    /// // the endpoint is included
    /// assert_eq!(primal::StreamingSieve::prime_pi(11), 5);
    ///
    /// assert_eq!(primal::StreamingSieve::prime_pi(100), 25);
    /// assert_eq!(primal::StreamingSieve::prime_pi(1000), 168);
    /// ```
    pub fn prime_pi(n: usize) -> usize {
        match n {
            0..=1 => 0,
            2 => 1,
            3..=4 => 2,
            5..=6 => 3,
            7..=10 => 4,
            _ => {
                let mut sieve = StreamingSieve::new(n);
                let (includes, base, tweak) = sieve.index_for(n);
                let mut count = match wheel::BYTE_MODULO {
                    30 => 3,
                    _ => unimplemented!()
                };

                for _ in 0..base {
                    let (_, bitv) = sieve.next().unwrap();
                    count += bitv.count_ones();
                }

                let (_, last) = sieve.next().unwrap();
                count += last.count_ones_before(tweak + includes as usize);
                count
            }
        }
    }

    /// Compute *pn*, the `n` prime number, 1-indexed
    /// (i.e. *p1* = 2, *p2* = 3).
    ///
    /// # Panics
    ///
    /// `n` must be larger than 0 and less than the total number of
    /// primes in this sieve (that is,
    /// `self.prime_pi(self.upper_bound())`).
    ///
    /// # Example
    ///
    /// ```rust
    /// assert_eq!(primal::StreamingSieve::nth_prime(1_000), 7919);
    /// ```
    pub fn nth_prime(n: usize) -> usize {
        assert!(n > 0);
        match n {
            1 => 2,
            2 => 3,
            3 => 5,
            _ => {
                let mut bit_n = n - 3;
                let (_, hi) = primal_estimate::nth_prime(n as u64);
                let mut sieve = StreamingSieve::new(hi as usize);
                while let Some((low, bits)) = sieve.next() {
                    let count = bits.count_ones();
                    if count >= bit_n {
                        let bit_idx = bits.find_nth_bit(bit_n - 1).unwrap();
                        return low + wheel::from_bit_index(bit_idx)
                    }

                    bit_n -= count
                }
                unreachable!()
            }
        }
    }

    fn add_sieving_prime(&mut self, p: usize, low: usize) {
        if p <= CACHE / 2 {
            self.small_primes.push(wheel::State::new(wheel::Wheel30, p, low));
        } else {
            let elem = wheel::State::new(wheel::Wheel210, p, low);
            if p < CACHE * 5 / 2 {
                self.primes.push(elem)
            } else {
                self.large_primes.push(elem)
            }
        }
    }

    fn find_new_sieving_primes(&mut self, low: usize, high: usize) {
        if let Some(small) = self.small.take() {
            for p in small.primes_from(self.current) {
                if p * p >= high {
                    self.current = p;
                    break
                }
                self.add_sieving_prime(p, low);
            }
            self.small = Some(small);
        }
    }

    fn small_primes_sieve<W: wheel::Wheel>(sieve: &mut BitVec,
                                           small_primes: &mut [wheel::State<W>]) {
        let bytes = sieve.as_bytes_mut();
        for wi in small_primes {
            wi.sieve_hardcoded(bytes);
        }
    }

    fn direct_sieve(&mut self) {
        let bytes = self.sieve.as_bytes_mut();
        let mut chunks = self.primes.chunks_exact_mut(3);

        while let Some([wi1, wi2, wi3]) = chunks.next() {
            wi1.sieve_triple(wi2, wi3, bytes);
        }

        for wi in chunks.into_remainder() {
            wi.sieve(bytes);
        }
    }

    fn large_primes_sieve(&mut self) {
        let bytes = self.sieve.as_bytes_mut();
        let mut chunks = self.large_primes.chunks_exact_mut(2);

        while let Some([wi1, wi2]) = chunks.next() {
            wi1.sieve_pair(wi2, bytes);
        }

        for wi in chunks.into_remainder() {
            wi.sieve(bytes);
        }
    }

    /// Extract the next chunk of filtered primes, the return value is
    /// `Some((low, v))` or `None` if the sieve has reached the limit.
    ///
    /// The vector stores bits for each odd number starting at `low`.
    /// Bit `n` of `v` is set if and only if `low + 2 * n + 1` is
    /// prime.
    ///
    /// NB. the prime 2 is not included in any of these sieves and so
    /// needs special handling.
    pub(crate) fn next(&mut self) -> Option<(usize, &BitVec)> {
        if self.low >= self.limit {
            return None
        }

        let low = self.low;
        self.low = self.low.saturating_add(SEG_LEN);
        let high = cmp::min(low.saturating_add(SEG_LEN - 1), self.limit);

        self.find_new_sieving_primes(low, high);

        self.presieve.apply(&mut self.sieve, low);
        StreamingSieve::small_primes_sieve(&mut self.sieve, &mut self.small_primes);
        self.direct_sieve();
        self.large_primes_sieve();

        if low == 0 {
            // 1 is not prime.
            self.sieve.set(0, false);
            self.presieve.mark_small_primes(&mut self.sieve);
        }

        Some((low, &self.sieve))
    }
}

#[cfg(test)]
mod tests {
    use crate::Sieve;
    use primal_slowsieve::Primes;
    use crate::wheel;
    use super::StreamingSieve;
    fn gcd(x: usize, y: usize) -> usize {
        if y == 0 { x }
        else { gcd(y, x % y) }
    }
    fn coprime_to(x: usize) -> Vec<usize> {
        (1..x).filter(|&n| gcd(n, x) == 1).collect()
    }
    #[test]
    fn test() {
        let coprime = coprime_to(wheel::BYTE_MODULO);
        const LIMIT: usize = 2_000_000;
        let mut sieve = StreamingSieve::new(LIMIT);
        let primes = ::primal_slowsieve::Primes::sieve(LIMIT);

        let mut base = 0;
        let mut index = 0;

        while let Some((_low, next)) = sieve.next() {
            for val in next {
                let i = wheel::BYTE_MODULO * base + coprime[index];
                if i >= LIMIT { break }
                assert!(primes.is_prime(i) == val,
                        "failed for {} (is prime = {})", i, primes.is_prime(i));

                index += 1;
                if index == wheel::BYTE_SIZE {
                    index = 0;
                    base += 1
                }
            }
        }
    }
    #[test]
    fn prime_pi() {
        let (limit, mult) = if cfg!(feature = "slow_tests") {
            (2_000_000, 19_998)
        } else {
            (200_000, 1_998)
        };
        let real = Primes::sieve(limit);

        for i in (0..20).chain((0..100).map(|n| n * mult + 1)) {
            let val = StreamingSieve::prime_pi(i);
            let true_ = real.primes().take_while(|p| *p <= i).count();
            assert!(val == true_, "failed for {}, true {}, computed {}",
                    i, true_, val)
        }
    }

    #[test]
    fn nth_prime() {
        let primes = Sieve::new(2_000_000);

        for (i, p) in primes.primes_from(0).enumerate() {
            let n = i + 1;
            if n < 2000 || n % 1000 == 0 {
                assert_eq!(StreamingSieve::nth_prime(n), p);
            }
        }
    }

    // These are designed to specifically test the medium sized and
    // large prime sieving.
    #[test]
    fn prime_pi_huge() {
        #[cfg(all(feature = "slow_tests", target_pointer_width = "64"))]
        const LIMIT_RESULT: (usize, usize) = (10_000_000_000, 455_052_511);
        #[cfg(all(feature = "slow_tests", target_pointer_width = "32"))]
        const LIMIT_RESULT: (usize, usize) = (4_294_000_000, 203_236_859);
        #[cfg(not(feature = "slow_tests"))]
        const LIMIT_RESULT: (usize, usize) = (500_000_000, 26_355_867);

        assert_eq!(StreamingSieve::prime_pi(LIMIT_RESULT.0), LIMIT_RESULT.1);
    }

    #[test]
    fn nth_prime_huge() {
        #[cfg(all(feature = "slow_tests", target_pointer_width = "64"))]
        const LIMIT_RESULT: (usize, usize) = (455_052_512, 10_000_000_019);
        #[cfg(all(feature = "slow_tests", target_pointer_width = "32"))]
        const LIMIT_RESULT: (usize, usize) = (203_236_860, 4_294_000_079);
        #[cfg(not(feature = "slow_tests"))]
        const LIMIT_RESULT: (usize, usize) = (26_355_868, 500_000_003);

        assert_eq!(StreamingSieve::nth_prime(LIMIT_RESULT.0), LIMIT_RESULT.1);
    }
}
<commit_msg>Fix the limit in find_new_sieving_primes<commit_after>use primal_bit::BitVec;
use std::cmp;

use crate::wheel;

pub mod primes;
mod presieve;

/// A heavily optimised prime sieve.
///
/// This is a streaming segmented sieve, meaning it sieves numbers in
/// intervals, extracting whatever it needs and discarding it. See
/// `Sieve` for a wrapper that caches the information to allow for
/// repeated queries, at the cost of *O(limit)* memory use.
///
/// This uses *O(sqrt(limit))* memory, and is designed to be as
/// cache-friendly as possible. `StreamingSieve` should be used for
/// one-off calls, or simple linear iteration.
///
/// The design is *heavily* inspired/adopted from Kim Walisch's
/// [primesieve](http://primesieve.org/), and has similar speed
/// (around 5-20% slower).
///
/// # Examples
///
/// ```rust
/// let count = primal::StreamingSieve::prime_pi(123456);
/// println!("𝜋(123456) = {}", count);
/// ```
#[derive(Debug)]
pub struct StreamingSieve {
    small: Option<crate::Sieve>,
    sieve: BitVec,
    primes: Vec<wheel::State<wheel::Wheel210>>,
    small_primes: Vec<wheel::State<wheel::Wheel30>>,
    large_primes: Vec<wheel::State<wheel::Wheel210>>,
    presieve: presieve::Presieve,

    low: usize,
    current: usize,
    limit: usize,
}

const CACHE: usize = 32 << 10;
const SEG_ELEMS: usize = 8 * CACHE;
const SEG_LEN: usize = SEG_ELEMS * wheel::BYTE_MODULO / wheel::BYTE_SIZE;

fn isqrt(x: usize) -> usize {
    (x as f64).sqrt() as usize
}

impl StreamingSieve {
    /// Create a new instance of the streaming sieve that will
    /// correctly progressively filter primes up to `limit`.
    pub(crate) fn new(limit: usize) -> StreamingSieve {
        let low = 0;

        let elems = cmp::min(wheel::bits_for(limit), SEG_ELEMS);
        let presieve = presieve::Presieve::new(elems);
        let current = presieve.smallest_unincluded_prime();

        let small = if limit < current * current {
            None
        } else {
            Some(crate::Sieve::new(isqrt(limit) + 1))
        };

        StreamingSieve {
            small,
            sieve: BitVec::from_elem(elems, true),
            primes: vec![],
            small_primes: vec![],
            large_primes: vec![],
            presieve,

            low,
            current,
            limit
        }
    }
    fn split_index(&self, idx: usize) -> (usize, usize) {
        let len = SEG_ELEMS;
        (idx / len,idx % len)
    }
    fn index_for(&self, n: usize) -> (bool, usize, usize) {
        let (b, idx) = wheel::bit_index(n);
        let (base, tweak) = self.split_index(idx);
        (b, base, tweak)
    }

    /// Count the number of primes upto and including `n`, that is, 𝜋,
    /// the [prime counting
    /// function](https://en.wikipedia.org/wiki/Prime-counting_function).
    ///
    /// # Examples
    ///
    /// ```rust
    /// assert_eq!(primal::StreamingSieve::prime_pi(10), 4);
    /// // the endpoint is included
    /// assert_eq!(primal::StreamingSieve::prime_pi(11), 5);
    ///
    /// assert_eq!(primal::StreamingSieve::prime_pi(100), 25);
    /// assert_eq!(primal::StreamingSieve::prime_pi(1000), 168);
    /// ```
    pub fn prime_pi(n: usize) -> usize {
        match n {
            0..=1 => 0,
            2 => 1,
            3..=4 => 2,
            5..=6 => 3,
            7..=10 => 4,
            _ => {
                let mut sieve = StreamingSieve::new(n);
                let (includes, base, tweak) = sieve.index_for(n);
                let mut count = match wheel::BYTE_MODULO {
                    30 => 3,
                    _ => unimplemented!()
                };

                for _ in 0..base {
                    let (_, bitv) = sieve.next().unwrap();
                    count += bitv.count_ones();
                }

                let (_, last) = sieve.next().unwrap();
                count += last.count_ones_before(tweak + includes as usize);
                count
            }
        }
    }

    /// Compute *pn*, the `n` prime number, 1-indexed
    /// (i.e. *p1* = 2, *p2* = 3).
    ///
    /// # Panics
    ///
    /// `n` must be larger than 0 and less than the total number of
    /// primes in this sieve (that is,
    /// `self.prime_pi(self.upper_bound())`).
    ///
    /// # Example
    ///
    /// ```rust
    /// assert_eq!(primal::StreamingSieve::nth_prime(1_000), 7919);
    /// ```
    pub fn nth_prime(n: usize) -> usize {
        assert!(n > 0);
        match n {
            1 => 2,
            2 => 3,
            3 => 5,
            _ => {
                let mut bit_n = n - 3;
                let (_, hi) = primal_estimate::nth_prime(n as u64);
                let mut sieve = StreamingSieve::new(hi as usize);
                while let Some((low, bits)) = sieve.next() {
                    let count = bits.count_ones();
                    if count >= bit_n {
                        let bit_idx = bits.find_nth_bit(bit_n - 1).unwrap();
                        return low + wheel::from_bit_index(bit_idx)
                    }

                    bit_n -= count
                }
                unreachable!()
            }
        }
    }

    fn add_sieving_prime(&mut self, p: usize, low: usize) {
        if p <= CACHE / 2 {
            self.small_primes.push(wheel::State::new(wheel::Wheel30, p, low));
        } else {
            let elem = wheel::State::new(wheel::Wheel210, p, low);
            if p < CACHE * 5 / 2 {
                self.primes.push(elem)
            } else {
                self.large_primes.push(elem)
            }
        }
    }

    fn find_new_sieving_primes(&mut self, low: usize, high: usize) {
        if let Some(small) = self.small.take() {
            for p in small.primes_from(self.current) {
                if p * p > high {
                    self.current = p;
                    break
                }
                self.add_sieving_prime(p, low);
            }
            self.small = Some(small);
        }
    }

    fn small_primes_sieve<W: wheel::Wheel>(sieve: &mut BitVec,
                                           small_primes: &mut [wheel::State<W>]) {
        let bytes = sieve.as_bytes_mut();
        for wi in small_primes {
            wi.sieve_hardcoded(bytes);
        }
    }

    fn direct_sieve(&mut self) {
        let bytes = self.sieve.as_bytes_mut();
        let mut chunks = self.primes.chunks_exact_mut(3);

        while let Some([wi1, wi2, wi3]) = chunks.next() {
            wi1.sieve_triple(wi2, wi3, bytes);
        }

        for wi in chunks.into_remainder() {
            wi.sieve(bytes);
        }
    }

    fn large_primes_sieve(&mut self) {
        let bytes = self.sieve.as_bytes_mut();
        let mut chunks = self.large_primes.chunks_exact_mut(2);

        while let Some([wi1, wi2]) = chunks.next() {
            wi1.sieve_pair(wi2, bytes);
        }

        for wi in chunks.into_remainder() {
            wi.sieve(bytes);
        }
    }

    /// Extract the next chunk of filtered primes, the return value is
    /// `Some((low, v))` or `None` if the sieve has reached the limit.
    ///
    /// The vector stores bits for each odd number starting at `low`.
    /// Bit `n` of `v` is set if and only if `low + 2 * n + 1` is
    /// prime.
    ///
    /// NB. the prime 2 is not included in any of these sieves and so
    /// needs special handling.
    pub(crate) fn next(&mut self) -> Option<(usize, &BitVec)> {
        if self.low >= self.limit {
            return None
        }

        let low = self.low;
        self.low = self.low.saturating_add(SEG_LEN);
        let high = cmp::min(low.saturating_add(SEG_LEN - 1), self.limit);

        self.find_new_sieving_primes(low, high);

        self.presieve.apply(&mut self.sieve, low);
        StreamingSieve::small_primes_sieve(&mut self.sieve, &mut self.small_primes);
        self.direct_sieve();
        self.large_primes_sieve();

        if low == 0 {
            // 1 is not prime.
            self.sieve.set(0, false);
            self.presieve.mark_small_primes(&mut self.sieve);
        }

        Some((low, &self.sieve))
    }
}

#[cfg(test)]
mod tests {
    use crate::Sieve;
    use primal_slowsieve::Primes;
    use crate::wheel;
    use super::StreamingSieve;
    fn gcd(x: usize, y: usize) -> usize {
        if y == 0 { x }
        else { gcd(y, x % y) }
    }
    fn coprime_to(x: usize) -> Vec<usize> {
        (1..x).filter(|&n| gcd(n, x) == 1).collect()
    }
    #[test]
    fn test() {
        let coprime = coprime_to(wheel::BYTE_MODULO);
        const LIMIT: usize = 2_000_000;
        let mut sieve = StreamingSieve::new(LIMIT);
        let primes = ::primal_slowsieve::Primes::sieve(LIMIT);

        let mut base = 0;
        let mut index = 0;

        while let Some((_low, next)) = sieve.next() {
            for val in next {
                let i = wheel::BYTE_MODULO * base + coprime[index];
                if i >= LIMIT { break }
                assert!(primes.is_prime(i) == val,
                        "failed for {} (is prime = {})", i, primes.is_prime(i));

                index += 1;
                if index == wheel::BYTE_SIZE {
                    index = 0;
                    base += 1
                }
            }
        }
    }
    #[test]
    fn prime_pi() {
        let (limit, mult) = if cfg!(feature = "slow_tests") {
            (2_000_000, 19_998)
        } else {
            (200_000, 1_998)
        };
        let real = Primes::sieve(limit);

        for i in (0..20).chain((0..100).map(|n| n * mult + 1)) {
            let val = StreamingSieve::prime_pi(i);
            let true_ = real.primes().take_while(|p| *p <= i).count();
            assert!(val == true_, "failed for {}, true {}, computed {}",
                    i, true_, val)
        }
    }

    #[test]
    fn prime_pi_issue_48() {
        // Expected values independently confirmed on Wolfram|Alpha.
        let limits = [49, 121, 289, 961, 11_047, 32_611, 230_907, 455_166_135];
        let expected = [15, 30, 61, 162, 1_338, 3_501, 20_513, 24_112_077];
        for (&limit, &expected) in limits.iter().zip(&expected) {
            let val = StreamingSieve::prime_pi(limit);
            assert_eq!(val, expected, "failed for limit {}", limit);
        }
    }

    #[test]
    fn nth_prime() {
        let primes = Sieve::new(2_000_000);

        for (i, p) in primes.primes_from(0).enumerate() {
            let n = i + 1;
            if n < 2000 || n % 1000 == 0 {
                assert_eq!(StreamingSieve::nth_prime(n), p);
            }
        }
    }

    // These are designed to specifically test the medium sized and
    // large prime sieving.
    #[test]
    fn prime_pi_huge() {
        #[cfg(all(feature = "slow_tests", target_pointer_width = "64"))]
        const LIMIT_RESULT: (usize, usize) = (10_000_000_000, 455_052_511);
        #[cfg(all(feature = "slow_tests", target_pointer_width = "32"))]
        const LIMIT_RESULT: (usize, usize) = (4_294_000_000, 203_236_859);
        #[cfg(not(feature = "slow_tests"))]
        const LIMIT_RESULT: (usize, usize) = (500_000_000, 26_355_867);

        assert_eq!(StreamingSieve::prime_pi(LIMIT_RESULT.0), LIMIT_RESULT.1);
    }

    #[test]
    fn nth_prime_huge() {
        #[cfg(all(feature = "slow_tests", target_pointer_width = "64"))]
        const LIMIT_RESULT: (usize, usize) = (455_052_512, 10_000_000_019);
        #[cfg(all(feature = "slow_tests", target_pointer_width = "32"))]
        const LIMIT_RESULT: (usize, usize) = (203_236_860, 4_294_000_079);
        #[cfg(not(feature = "slow_tests"))]
        const LIMIT_RESULT: (usize, usize) = (26_355_868, 500_000_003);

        assert_eq!(StreamingSieve::nth_prime(LIMIT_RESULT.0), LIMIT_RESULT.1);
    }
}
<|endoftext|>