<commit_before>use itertools::Itertools;
use stdx::{format_to, to_lower_snake_case};
use syntax::ast::VisibilityOwner;
use syntax::ast::{self, AstNode, NameOwner};

use crate::{
    utils::{find_impl_block_end, find_struct_impl, generate_impl_text},
    AssistContext, AssistId, AssistKind, Assists,
};

// Assist: generate_enum_is_method
//
// Generate an `is_` method for an enum variant.
//
// ```
// enum Version {
//  Undefined,
//  Minor$0,
//  Major,
// }
// ```
// ->
// ```
// enum Version {
//  Undefined,
//  Minor,
//  Major,
// }
//
// impl Version {
//     /// Returns `true` if the version is [`Minor`].
//     fn is_minor(&self) -> bool {
//         matches!(self, Self::Minor)
//     }
// }
// ```
pub(crate) fn generate_enum_is_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let variant = ctx.find_node_at_offset::<ast::Variant>()?;
    let variant_name = variant.name()?;
    let parent_enum = ast::Adt::Enum(variant.parent_enum());
    let variant_kind = variant_kind(&variant);

    let enum_lowercase_name = to_lower_snake_case(&parent_enum.name()?.to_string());
    let fn_name = format!("is_{}", &to_lower_snake_case(variant_name.text()));

    // Return early if we've found an existing new fn
    let impl_def = find_struct_impl(
        &ctx,
        &parent_enum,
        &fn_name,
    )?;

    let target = variant.syntax().text_range();
    acc.add(
        AssistId("generate_enum_is_method", AssistKind::Generate),
        "Generate an `is_` method for an enum variant",
        target,
        |builder| {
            let mut buf = String::with_capacity(512);

            if impl_def.is_some() {
                buf.push('\n');
            }

            let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
            format_to!(
                buf,
                "    /// Returns `true` if the {} is [`{}`].
    {}fn {}(&self) -> bool {{
        matches!(self, Self::{}{})
    }}",
                enum_lowercase_name,
                variant_name,
                vis,
                fn_name,
                variant_name,
                variant_kind.pattern_suffix(),
            );

            let start_offset = impl_def
                .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf))
                .unwrap_or_else(|| {
                    buf = generate_impl_text(&parent_enum, &buf);
                    parent_enum.syntax().text_range().end()
                });

            builder.insert(start_offset, buf);
        },
    )
}

// Assist: generate_enum_into_method
//
// Generate an `into_` method for an enum variant.
//
// ```
// enum Value {
//  Number(i32),
//  Text(String)$0,
// }
// ```
// ->
// ```
// enum Value {
//  Number(i32),
//  Text(String),
// }
//
// impl Value {
//     fn into_text(self) -> Option<String> {
//         if let Self::Text(v) = self {
//             Some(v)
//         } else {
//             None
//         }
//     }
// }
// ```
pub(crate) fn generate_enum_into_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let variant = ctx.find_node_at_offset::<ast::Variant>()?;
    let variant_name = variant.name()?;
    let parent_enum = ast::Adt::Enum(variant.parent_enum());
    let variant_kind = variant_kind(&variant);

    let fn_name = format!("into_{}", &to_lower_snake_case(variant_name.text()));

    // Return early if we've found an existing new fn
    let impl_def = find_struct_impl(
        &ctx,
        &parent_enum,
        &fn_name,
    )?;

    let field_type = variant_kind.single_field_type()?;
    let (pattern_suffix, bound_name) = variant_kind.binding_pattern()?;

    let target = variant.syntax().text_range();
    acc.add(
        AssistId("generate_enum_into_method", AssistKind::Generate),
        "Generate an `into_` method for an enum variant",
        target,
        |builder| {
            let mut buf = String::with_capacity(512);

            if impl_def.is_some() {
                buf.push('\n');
            }

            let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
            format_to!(
                buf,
                "    {}fn {}(self) -> Option<{}> {{
        if let Self::{}{} = self {{
            Some({})
        }} else {{
            None
        }}
    }}",
                vis,
                fn_name,
                field_type.syntax(),
                variant_name,
                pattern_suffix,
                bound_name,
            );

            let start_offset = impl_def
                .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf))
                .unwrap_or_else(|| {
                    buf = generate_impl_text(&parent_enum, &buf);
                    parent_enum.syntax().text_range().end()
                });

            builder.insert(start_offset, buf);
        },
    )
}

enum VariantKind {
    Unit,
    /// Tuple with a single field
    NewtypeTuple { ty: Option<ast::Type> },
    /// Tuple with 0 or more than 2 fields
    Tuple,
    /// Record with a single field
    NewtypeRecord {
        field_name: Option<ast::Name>,
        field_type: Option<ast::Type>,
    },
    /// Record with 0 or more than 2 fields
    Record,
}

impl VariantKind {
    fn pattern_suffix(&self) -> &'static str {
        match self {
            VariantKind::Unit => "",
            VariantKind::NewtypeTuple { .. } |
            VariantKind::Tuple => "(..)",
            VariantKind::NewtypeRecord { .. } |
            VariantKind::Record => " { .. }",
        }
    }

    fn binding_pattern(&self) -> Option<(String, String)> {
        match self {
            VariantKind::Unit |
            VariantKind::Tuple |
            VariantKind::Record |
            VariantKind::NewtypeRecord { field_name: None, .. } => None,
            VariantKind::NewtypeTuple { .. } => {
                Some(("(v)".to_owned(), "v".to_owned()))
            }
            VariantKind::NewtypeRecord { field_name: Some(name), .. } => {
                Some((
                    format!(" {{ {} }}", name.syntax()),
                    name.syntax().to_string(),
                ))
            }
        }
    }

    fn single_field_type(&self) -> Option<&ast::Type> {
        match self {
            VariantKind::Unit |
            VariantKind::Tuple |
            VariantKind::Record => None,
            VariantKind::NewtypeTuple { ty } => ty.as_ref(),
            VariantKind::NewtypeRecord { field_type, .. } => field_type.as_ref(),
        }
    }
}

fn variant_kind(variant: &ast::Variant) -> VariantKind {
    match variant.kind() {
        ast::StructKind::Record(record) => {
            if let Some((single_field,)) = record.fields().collect_tuple() {
                let field_name = single_field.name();
                let field_type = single_field.ty();
                VariantKind::NewtypeRecord { field_name, field_type }
            } else {
                VariantKind::Record
            }
        }
        ast::StructKind::Tuple(tuple) => {
            if let Some((single_field,)) = tuple.fields().collect_tuple() {
                let ty = single_field.ty();
                VariantKind::NewtypeTuple { ty }
            } else {
                VariantKind::Tuple
            }
        }
        ast::StructKind::Unit => VariantKind::Unit,
    }
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn test_generate_enum_is_from_variant() {
        check_assist(
            generate_enum_is_method,
            r#"
enum Variant {
    Undefined,
    Minor$0,
    Major,
}"#,
            r#"enum Variant {
    Undefined,
    Minor,
    Major,
}

impl Variant {
    /// Returns `true` if the variant is [`Minor`].
    fn is_minor(&self) -> bool {
        matches!(self, Self::Minor)
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_is_already_implemented() {
        check_assist_not_applicable(
            generate_enum_is_method,
            r#"
enum Variant {
    Undefined,
    Minor$0,
    Major,
}

impl Variant {
    fn is_minor(&self) -> bool {
        matches!(self, Self::Minor)
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_is_from_tuple_variant() {
        check_assist(
            generate_enum_is_method,
            r#"
enum Variant {
    Undefined,
    Minor(u32)$0,
    Major,
}"#,
            r#"enum Variant {
    Undefined,
    Minor(u32),
    Major,
}

impl Variant {
    /// Returns `true` if the variant is [`Minor`].
    fn is_minor(&self) -> bool {
        matches!(self, Self::Minor(..))
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_is_from_record_variant() {
        check_assist(
            generate_enum_is_method,
            r#"
enum Variant {
    Undefined,
    Minor { foo: i32 }$0,
    Major,
}"#,
            r#"enum Variant {
    Undefined,
    Minor { foo: i32 },
    Major,
}

impl Variant {
    /// Returns `true` if the variant is [`Minor`].
    fn is_minor(&self) -> bool {
        matches!(self, Self::Minor { .. })
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_is_from_variant_with_one_variant() {
        check_assist(
            generate_enum_is_method,
            r#"enum Variant { Undefi$0ned }"#,
            r#"
enum Variant { Undefined }

impl Variant {
    /// Returns `true` if the variant is [`Undefined`].
    fn is_undefined(&self) -> bool {
        matches!(self, Self::Undefined)
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_is_from_variant_with_visibility_marker() {
        check_assist(
            generate_enum_is_method,
            r#"
pub(crate) enum Variant {
    Undefined,
    Minor$0,
    Major,
}"#,
            r#"pub(crate) enum Variant {
    Undefined,
    Minor,
    Major,
}

impl Variant {
    /// Returns `true` if the variant is [`Minor`].
    pub(crate) fn is_minor(&self) -> bool {
        matches!(self, Self::Minor)
    }
}"#,
        );
    }

    #[test]
    fn test_multiple_generate_enum_is_from_variant() {
        check_assist(
            generate_enum_is_method,
            r#"
enum Variant {
    Undefined,
    Minor,
    Major$0,
}

impl Variant {
    /// Returns `true` if the variant is [`Minor`].
    fn is_minor(&self) -> bool {
        matches!(self, Self::Minor)
    }
}"#,
            r#"enum Variant {
    Undefined,
    Minor,
    Major,
}

impl Variant {
    /// Returns `true` if the variant is [`Minor`].
    fn is_minor(&self) -> bool {
        matches!(self, Self::Minor)
    }

    /// Returns `true` if the variant is [`Major`].
    fn is_major(&self) -> bool {
        matches!(self, Self::Major)
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_into_tuple_variant() {
        check_assist(
            generate_enum_into_method,
            r#"
enum Value {
    Number(i32),
    Text(String)$0,
}"#,
            r#"enum Value {
    Number(i32),
    Text(String),
}

impl Value {
    fn into_text(self) -> Option<String> {
        if let Self::Text(v) = self {
            Some(v)
        } else {
            None
        }
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_into_already_implemented() {
        check_assist_not_applicable(
            generate_enum_into_method,
            r#"enum Value {
    Number(i32),
    Text(String)$0,
}

impl Value {
    fn into_text(self) -> Option<String> {
        if let Self::Text(v) = self {
            Some(v)
        } else {
            None
        }
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_into_unit_variant() {
        check_assist_not_applicable(
            generate_enum_into_method,
            r#"enum Value {
    Number(i32),
    Text(String),
    Unit$0,
}"#,
        );
    }

    #[test]
    fn test_generate_enum_into_record_with_multiple_fields() {
        check_assist_not_applicable(
            generate_enum_into_method,
            r#"enum Value {
    Number(i32),
    Text(String),
    Both { first: i32, second: String }$0,
}"#,
        );
    }

    #[test]
    fn test_generate_enum_into_tuple_with_multiple_fields() {
        check_assist_not_applicable(
            generate_enum_into_method,
            r#"enum Value {
    Number(i32),
    Text(String, String)$0,
}"#,
        );
    }

    #[test]
    fn test_generate_enum_into_record_variant() {
        check_assist(
            generate_enum_into_method,
            r#"enum Value {
    Number(i32),
    Text { text: String }$0,
}"#,
            r#"enum Value {
    Number(i32),
    Text { text: String },
}

impl Value {
    fn into_text(self) -> Option<String> {
        if let Self::Text { text } = self {
            Some(text)
        } else {
            None
        }
    }
}"#,
        );
    }
}
<commit_msg>deduplicate some<commit_after>use itertools::Itertools;
use stdx::to_lower_snake_case;
use syntax::ast::VisibilityOwner;
use syntax::ast::{self, AstNode, NameOwner};

use crate::{AssistContext, AssistId, AssistKind, Assists, assist_context::AssistBuilder, utils::{find_impl_block_end, find_struct_impl, generate_impl_text}};

// Assist: generate_enum_is_method
//
// Generate an `is_` method for an enum variant.
//
// ```
// enum Version {
//  Undefined,
//  Minor$0,
//  Major,
// }
// ```
// ->
// ```
// enum Version {
//  Undefined,
//  Minor,
//  Major,
// }
//
// impl Version {
//     /// Returns `true` if the version is [`Minor`].
//     fn is_minor(&self) -> bool {
//         matches!(self, Self::Minor)
//     }
// }
// ```
pub(crate) fn generate_enum_is_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let variant = ctx.find_node_at_offset::<ast::Variant>()?;
    let variant_name = variant.name()?;
    let parent_enum = ast::Adt::Enum(variant.parent_enum());
    let variant_kind = variant_kind(&variant);

    let enum_lowercase_name = to_lower_snake_case(&parent_enum.name()?.to_string());
    let fn_name = format!("is_{}", &to_lower_snake_case(variant_name.text()));

    // Return early if we've found an existing new fn
    let impl_def = find_struct_impl(
        &ctx,
        &parent_enum,
        &fn_name,
    )?;

    let target = variant.syntax().text_range();
    acc.add(
        AssistId("generate_enum_is_method", AssistKind::Generate),
        "Generate an `is_` method for an enum variant",
        target,
        |builder| {
            let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
            let method = format!(
                "    /// Returns `true` if the {} is [`{}`].
    {}fn {}(&self) -> bool {{
        matches!(self, Self::{}{})
    }}",
                enum_lowercase_name,
                variant_name,
                vis,
                fn_name,
                variant_name,
                variant_kind.pattern_suffix(),
            );

            add_method_to_adt(builder, &parent_enum, impl_def, &method);
        },
    )
}

// Assist: generate_enum_into_method
//
// Generate an `into_` method for an enum variant.
//
// ```
// enum Value {
//  Number(i32),
//  Text(String)$0,
// }
// ```
// ->
// ```
// enum Value {
//  Number(i32),
//  Text(String),
// }
//
// impl Value {
//     fn into_text(self) -> Option<String> {
//         if let Self::Text(v) = self {
//             Some(v)
//         } else {
//             None
//         }
//     }
// }
// ```
pub(crate) fn generate_enum_into_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let variant = ctx.find_node_at_offset::<ast::Variant>()?;
    let variant_name = variant.name()?;
    let parent_enum = ast::Adt::Enum(variant.parent_enum());
    let variant_kind = variant_kind(&variant);

    let fn_name = format!("into_{}", &to_lower_snake_case(variant_name.text()));

    // Return early if we've found an existing new fn
    let impl_def = find_struct_impl(
        &ctx,
        &parent_enum,
        &fn_name,
    )?;

    let field_type = variant_kind.single_field_type()?;
    let (pattern_suffix, bound_name) = variant_kind.binding_pattern()?;

    let target = variant.syntax().text_range();
    acc.add(
        AssistId("generate_enum_into_method", AssistKind::Generate),
        "Generate an `into_` method for an enum variant",
        target,
        |builder| {
            let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
            let method = format!(
                "    {}fn {}(self) -> Option<{}> {{
        if let Self::{}{} = self {{
            Some({})
        }} else {{
            None
        }}
    }}",
                vis,
                fn_name,
                field_type.syntax(),
                variant_name,
                pattern_suffix,
                bound_name,
            );

            add_method_to_adt(builder, &parent_enum, impl_def, &method);
        },
    )
}

fn add_method_to_adt(
    builder: &mut AssistBuilder,
    adt: &ast::Adt,
    impl_def: Option<ast::Impl>,
    method: &str,
) {
    let mut buf = String::with_capacity(method.len() + 2);
    if impl_def.is_some() {
        buf.push('\n');
    }
    buf.push_str(method);

    let start_offset = impl_def
        .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf))
        .unwrap_or_else(|| {
            buf = generate_impl_text(&adt, &buf);
            adt.syntax().text_range().end()
        });

    builder.insert(start_offset, buf);
}

enum VariantKind {
    Unit,
    /// Tuple with a single field
    NewtypeTuple { ty: Option<ast::Type> },
    /// Tuple with 0 or more than 2 fields
    Tuple,
    /// Record with a single field
    NewtypeRecord {
        field_name: Option<ast::Name>,
        field_type: Option<ast::Type>,
    },
    /// Record with 0 or more than 2 fields
    Record,
}

impl VariantKind {
    fn pattern_suffix(&self) -> &'static str {
        match self {
            VariantKind::Unit => "",
            VariantKind::NewtypeTuple { .. } |
            VariantKind::Tuple => "(..)",
            VariantKind::NewtypeRecord { .. } |
            VariantKind::Record => " { .. }",
        }
    }

    fn binding_pattern(&self) -> Option<(String, String)> {
        match self {
            VariantKind::Unit |
            VariantKind::Tuple |
            VariantKind::Record |
            VariantKind::NewtypeRecord { field_name: None, .. } => None,
            VariantKind::NewtypeTuple { .. } => {
                Some(("(v)".to_owned(), "v".to_owned()))
            }
            VariantKind::NewtypeRecord { field_name: Some(name), .. } => {
                Some((
                    format!(" {{ {} }}", name.syntax()),
                    name.syntax().to_string(),
                ))
            }
        }
    }

    fn single_field_type(&self) -> Option<&ast::Type> {
        match self {
            VariantKind::Unit |
            VariantKind::Tuple |
            VariantKind::Record => None,
            VariantKind::NewtypeTuple { ty } => ty.as_ref(),
            VariantKind::NewtypeRecord { field_type, .. } => field_type.as_ref(),
        }
    }
}

fn variant_kind(variant: &ast::Variant) -> VariantKind {
    match variant.kind() {
        ast::StructKind::Record(record) => {
            if let Some((single_field,)) = record.fields().collect_tuple() {
                let field_name = single_field.name();
                let field_type = single_field.ty();
                VariantKind::NewtypeRecord { field_name, field_type }
            } else {
                VariantKind::Record
            }
        }
        ast::StructKind::Tuple(tuple) => {
            if let Some((single_field,)) = tuple.fields().collect_tuple() {
                let ty = single_field.ty();
                VariantKind::NewtypeTuple { ty }
            } else {
                VariantKind::Tuple
            }
        }
        ast::StructKind::Unit => VariantKind::Unit,
    }
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn test_generate_enum_is_from_variant() {
        check_assist(
            generate_enum_is_method,
            r#"
enum Variant {
    Undefined,
    Minor$0,
    Major,
}"#,
            r#"enum Variant {
    Undefined,
    Minor,
    Major,
}

impl Variant {
    /// Returns `true` if the variant is [`Minor`].
    fn is_minor(&self) -> bool {
        matches!(self, Self::Minor)
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_is_already_implemented() {
        check_assist_not_applicable(
            generate_enum_is_method,
            r#"
enum Variant {
    Undefined,
    Minor$0,
    Major,
}

impl Variant {
    fn is_minor(&self) -> bool {
        matches!(self, Self::Minor)
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_is_from_tuple_variant() {
        check_assist(
            generate_enum_is_method,
            r#"
enum Variant {
    Undefined,
    Minor(u32)$0,
    Major,
}"#,
            r#"enum Variant {
    Undefined,
    Minor(u32),
    Major,
}

impl Variant {
    /// Returns `true` if the variant is [`Minor`].
    fn is_minor(&self) -> bool {
        matches!(self, Self::Minor(..))
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_is_from_record_variant() {
        check_assist(
            generate_enum_is_method,
            r#"
enum Variant {
    Undefined,
    Minor { foo: i32 }$0,
    Major,
}"#,
            r#"enum Variant {
    Undefined,
    Minor { foo: i32 },
    Major,
}

impl Variant {
    /// Returns `true` if the variant is [`Minor`].
    fn is_minor(&self) -> bool {
        matches!(self, Self::Minor { .. })
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_is_from_variant_with_one_variant() {
        check_assist(
            generate_enum_is_method,
            r#"enum Variant { Undefi$0ned }"#,
            r#"
enum Variant { Undefined }

impl Variant {
    /// Returns `true` if the variant is [`Undefined`].
    fn is_undefined(&self) -> bool {
        matches!(self, Self::Undefined)
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_is_from_variant_with_visibility_marker() {
        check_assist(
            generate_enum_is_method,
            r#"
pub(crate) enum Variant {
    Undefined,
    Minor$0,
    Major,
}"#,
            r#"pub(crate) enum Variant {
    Undefined,
    Minor,
    Major,
}

impl Variant {
    /// Returns `true` if the variant is [`Minor`].
    pub(crate) fn is_minor(&self) -> bool {
        matches!(self, Self::Minor)
    }
}"#,
        );
    }

    #[test]
    fn test_multiple_generate_enum_is_from_variant() {
        check_assist(
            generate_enum_is_method,
            r#"
enum Variant {
    Undefined,
    Minor,
    Major$0,
}

impl Variant {
    /// Returns `true` if the variant is [`Minor`].
    fn is_minor(&self) -> bool {
        matches!(self, Self::Minor)
    }
}"#,
            r#"enum Variant {
    Undefined,
    Minor,
    Major,
}

impl Variant {
    /// Returns `true` if the variant is [`Minor`].
    fn is_minor(&self) -> bool {
        matches!(self, Self::Minor)
    }

    /// Returns `true` if the variant is [`Major`].
    fn is_major(&self) -> bool {
        matches!(self, Self::Major)
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_into_tuple_variant() {
        check_assist(
            generate_enum_into_method,
            r#"
enum Value {
    Number(i32),
    Text(String)$0,
}"#,
            r#"enum Value {
    Number(i32),
    Text(String),
}

impl Value {
    fn into_text(self) -> Option<String> {
        if let Self::Text(v) = self {
            Some(v)
        } else {
            None
        }
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_into_already_implemented() {
        check_assist_not_applicable(
            generate_enum_into_method,
            r#"enum Value {
    Number(i32),
    Text(String)$0,
}

impl Value {
    fn into_text(self) -> Option<String> {
        if let Self::Text(v) = self {
            Some(v)
        } else {
            None
        }
    }
}"#,
        );
    }

    #[test]
    fn test_generate_enum_into_unit_variant() {
        check_assist_not_applicable(
            generate_enum_into_method,
            r#"enum Value {
    Number(i32),
    Text(String),
    Unit$0,
}"#,
        );
    }

    #[test]
    fn test_generate_enum_into_record_with_multiple_fields() {
        check_assist_not_applicable(
            generate_enum_into_method,
            r#"enum Value {
    Number(i32),
    Text(String),
    Both { first: i32, second: String }$0,
}"#,
        );
    }

    #[test]
    fn test_generate_enum_into_tuple_with_multiple_fields() {
        check_assist_not_applicable(
            generate_enum_into_method,
            r#"enum Value {
    Number(i32),
    Text(String, String)$0,
}"#,
        );
    }

    #[test]
    fn test_generate_enum_into_record_variant() {
        check_assist(
            generate_enum_into_method,
            r#"enum Value {
    Number(i32),
    Text { text: String }$0,
}"#,
            r#"enum Value {
    Number(i32),
    Text { text: String },
}

impl Value {
    fn into_text(self) -> Option<String> {
        if let Self::Text { text } = self {
            Some(text)
        } else {
            None
        }
    }
}"#,
        );
    }
}
<|endoftext|>


<commit_before>#![cfg(target_os = "linux")]

mod common;

use std::fs;
use std::mem::size_of;
use std::path::Path;
use std::thread;
use std::time::Duration;

use rand::Rng;

use sled::{Config, ConfigBuilder, Db};

const CYCLE: usize = 256;
const BATCH_SIZE: u32 = 8;

/// Verifies that the keys in the tree are correctly recovered.
/// Panics if they are incorrect.
/// Returns the key that should be resumed at, and the current cycle value.
fn verify(tree: &sled::Tree) -> (u32, u32) {
    // key 0 should always be the highest value, as that's where we increment
    // at some point, it might go down by one
    // it should never return, or go down again after that
    let mut iter = tree.iter();
    let highest = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return (0, 0),
    };

    let highest_vec = u32_to_vec(highest);

    // find how far we got
    let mut contiguous: u32 = 0;
    let mut lowest = 0;
    for res in iter {
        let (_k, v) = res.unwrap();
        if &v[..4] == &highest_vec[..4] {
            contiguous += 1;
        } else {
            let expected = if highest == 0 {
                CYCLE as u32 - 1
            } else {
                (highest - 1) % CYCLE as u32
            };
            let actual = slice_to_u32(&*v);
            assert_eq!(expected, actual);
            lowest = actual;
            break;
        }
    }

    // ensure nothing changes after this point
    let low_beginning = u32_to_vec(contiguous + 1);

    for res in tree.range(&*low_beginning..) {
        let (k, v): (sled::IVec, _) = res.unwrap();
        assert_eq!(
            slice_to_u32(&*v),
            lowest,
            "expected key {} to have value {}, instead it had value {} in db: {:?}",
            slice_to_u32(&*k),
            lowest,
            slice_to_u32(&*v),
            tree
        );
    }

    (contiguous, highest)
}

/// Verifies that the keys in the tree are correctly recovered (i.e., equal).
/// Panics if they are incorrect.
fn verify_batches(tree: &sled::Tree) -> u32 {
    let mut iter = tree.iter();
    let first_value = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return 0,
    };
    for key in 0..BATCH_SIZE {
        let res = tree.get(u32_to_vec(key));
        let option = res.unwrap();
        let v = match option {
            Some(v) => v,
            None => panic!(
                "expected key {} to have a value, instead it was missing in db: {:?}",
                key,
                tree
            ),
        };
        let value = slice_to_u32(&*v);
        assert_eq!(
            first_value,
            value,
            "expected key {} to have value {}, instead it had value {} in db: {:?}",
            key,
            first_value,
            value,
            tree
        );
    }

    first_value
}

fn u32_to_vec(u: u32) -> Vec<u8> {
    let buf: [u8; size_of::<u32>()] = u.to_be_bytes();
    buf.to_vec()
}

fn slice_to_u32(b: &[u8]) -> u32 {
    let mut buf = [0u8; size_of::<u32>()];
    buf.copy_from_slice(&b[..size_of::<u32>()]);

    u32::from_be_bytes(buf)
}

fn spawn_killah() {
    thread::spawn(|| {
        let runtime = rand::thread_rng().gen_range(0, 200);
        thread::sleep(Duration::from_millis(runtime));
        unsafe {
            libc::raise(9);
        }
    });
}

fn run(config: Config) {
    common::setup_logger();

    let crash_during_initialization = rand::thread_rng().gen_bool(0.1);

    if crash_during_initialization {
        spawn_killah();
    }

    let tree = Db::start(config).unwrap();

    if !crash_during_initialization {
        spawn_killah();
    }

    let (key, highest) = verify(&tree);

    let mut hu = ((highest as usize) * CYCLE) + key as usize;
    assert_eq!(hu % CYCLE, key as usize);
    assert_eq!(hu / CYCLE, highest as usize);

    loop {
        hu += 1;

        if hu / CYCLE >= CYCLE {
            hu = 0;
        }

        let key = u32_to_vec((hu % CYCLE) as u32);

        let mut value = u32_to_vec((hu / CYCLE) as u32);
        let additional_len = rand::thread_rng().gen_range(0, 1000);
        value.append(&mut vec![0u8; additional_len]);

        tree.insert(&key, value).unwrap();
    }
}

fn run_batches(config: Config) {
    fn do_batch(i: u32, tree: &sled::Tree) {
        let mut rng = rand::thread_rng();
        let base_value = u32_to_vec(i);

        let mut batch = sled::Batch::default();
        if rng.gen_bool(0.1) {
            for key in 0..BATCH_SIZE {
                batch.remove(u32_to_vec(key));
            }
        } else {
            for key in 0..BATCH_SIZE {
                let mut value = base_value.clone();
                let additional_len = rng.gen_range(0, 1000);
                value.append(&mut vec![0u8; additional_len]);

                batch.insert(u32_to_vec(key), value);
            }
        }
        tree.apply_batch(batch).unwrap();
    }

    common::setup_logger();
    let crash_during_initialization = rand::thread_rng().gen_bool(0.1);

    if crash_during_initialization {
        spawn_killah();
    }

    let tree = sled::Db::start(config).unwrap();

    let mut i = verify_batches(&tree);
    i += 1;
    do_batch(i, &tree);

    if !crash_during_initialization {
        spawn_killah();
    }

    loop {
        i += 1;
        do_batch(i, &tree);
    }
}

fn run_without_snapshot(dir: &str) {
    let mut config_builder = ConfigBuilder::new()
        .page_consolidation_threshold(10)
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(1 << 56);

    config_builder.io_buf_size = 1024;

    let config = config_builder.build();

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_with_snapshot(dir: &str) {
    let mut config_builder = ConfigBuilder::new()
        .page_consolidation_threshold(10)
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(5000);

    config_builder.io_buf_size = 1024;

    let config = config_builder.build();

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_batches_without_snapshot(dir: &str) {
    let mut config_builder = ConfigBuilder::new()
        .page_consolidation_threshold(10)
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(1 << 56);

    config_builder.io_buf_size = 1024;

    let config = config_builder.build();

    match thread::spawn(|| run_batches(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_batches_with_snapshot(dir: &str) {
    let mut config_builder = ConfigBuilder::new()
        .page_consolidation_threshold(10)
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(5000);

    config_builder.io_buf_size = 1024;

    let config = config_builder.build();

    match thread::spawn(|| run_batches(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

#[test]
fn test_crash_recovery_with_runtime_snapshot() {
    let dir = "test_crashes_with_snapshot";
    cleanup(dir);
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_with_snapshot(dir)
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup(dir);
                panic!("child exited abnormally");
            }
        }
    }
    cleanup(dir);
}

#[test]
fn test_crash_recovery_no_runtime_snapshot() {
    let dir = "test_crashes";
    cleanup(dir);
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_without_snapshot(dir)
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup(dir);
                panic!("child exited abnormally");
            }
        }
    }
    cleanup(dir);
}

#[test]
#[ignore]
fn test_crash_batches_with_runtime_snapshot() {
    let dir = "test_batches_with_snapshot";
    cleanup(dir);
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_batches_with_snapshot(dir)
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup(dir);
                panic!("child exited abnormally");
            }
        }
    }
    cleanup(dir);
}

#[test]
#[ignore]
fn test_crash_batches_no_runtime_snapshot() {
    let dir = "test_batches";
    cleanup(dir);
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_batches_without_snapshot(dir)
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup(dir);
                panic!("child exited abnormally");
            }
        }
    }
    cleanup(dir);
}

fn cleanup(dir: &str) {
    let dir = Path::new(dir);
    if dir.exists() {
        fs::remove_dir_all(dir).unwrap();
    }
}
<commit_msg>Remove page_consolidation_threshold from crash tests<commit_after>#![cfg(target_os = "linux")]

mod common;

use std::fs;
use std::mem::size_of;
use std::path::Path;
use std::thread;
use std::time::Duration;

use rand::Rng;

use sled::{Config, ConfigBuilder, Db};

const CYCLE: usize = 256;
const BATCH_SIZE: u32 = 8;

/// Verifies that the keys in the tree are correctly recovered.
/// Panics if they are incorrect.
/// Returns the key that should be resumed at, and the current cycle value.
fn verify(tree: &sled::Tree) -> (u32, u32) {
    // key 0 should always be the highest value, as that's where we increment
    // at some point, it might go down by one
    // it should never return, or go down again after that
    let mut iter = tree.iter();
    let highest = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return (0, 0),
    };

    let highest_vec = u32_to_vec(highest);

    // find how far we got
    let mut contiguous: u32 = 0;
    let mut lowest = 0;
    for res in iter {
        let (_k, v) = res.unwrap();
        if &v[..4] == &highest_vec[..4] {
            contiguous += 1;
        } else {
            let expected = if highest == 0 {
                CYCLE as u32 - 1
            } else {
                (highest - 1) % CYCLE as u32
            };
            let actual = slice_to_u32(&*v);
            assert_eq!(expected, actual);
            lowest = actual;
            break;
        }
    }

    // ensure nothing changes after this point
    let low_beginning = u32_to_vec(contiguous + 1);

    for res in tree.range(&*low_beginning..) {
        let (k, v): (sled::IVec, _) = res.unwrap();
        assert_eq!(
            slice_to_u32(&*v),
            lowest,
            "expected key {} to have value {}, instead it had value {} in db: {:?}",
            slice_to_u32(&*k),
            lowest,
            slice_to_u32(&*v),
            tree
        );
    }

    (contiguous, highest)
}

/// Verifies that the keys in the tree are correctly recovered (i.e., equal).
/// Panics if they are incorrect.
fn verify_batches(tree: &sled::Tree) -> u32 {
    let mut iter = tree.iter();
    let first_value = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return 0,
    };
    for key in 0..BATCH_SIZE {
        let res = tree.get(u32_to_vec(key));
        let option = res.unwrap();
        let v = match option {
            Some(v) => v,
            None => panic!(
                "expected key {} to have a value, instead it was missing in db: {:?}",
                key,
                tree
            ),
        };
        let value = slice_to_u32(&*v);
        assert_eq!(
            first_value,
            value,
            "expected key {} to have value {}, instead it had value {} in db: {:?}",
            key,
            first_value,
            value,
            tree
        );
    }

    first_value
}

fn u32_to_vec(u: u32) -> Vec<u8> {
    let buf: [u8; size_of::<u32>()] = u.to_be_bytes();
    buf.to_vec()
}

fn slice_to_u32(b: &[u8]) -> u32 {
    let mut buf = [0u8; size_of::<u32>()];
    buf.copy_from_slice(&b[..size_of::<u32>()]);

    u32::from_be_bytes(buf)
}

fn spawn_killah() {
    thread::spawn(|| {
        let runtime = rand::thread_rng().gen_range(0, 200);
        thread::sleep(Duration::from_millis(runtime));
        unsafe {
            libc::raise(9);
        }
    });
}

fn run(config: Config) {
    common::setup_logger();

    let crash_during_initialization = rand::thread_rng().gen_bool(0.1);

    if crash_during_initialization {
        spawn_killah();
    }

    let tree = Db::start(config).unwrap();

    if !crash_during_initialization {
        spawn_killah();
    }

    let (key, highest) = verify(&tree);

    let mut hu = ((highest as usize) * CYCLE) + key as usize;
    assert_eq!(hu % CYCLE, key as usize);
    assert_eq!(hu / CYCLE, highest as usize);

    loop {
        hu += 1;

        if hu / CYCLE >= CYCLE {
            hu = 0;
        }

        let key = u32_to_vec((hu % CYCLE) as u32);

        let mut value = u32_to_vec((hu / CYCLE) as u32);
        let additional_len = rand::thread_rng().gen_range(0, 1000);
        value.append(&mut vec![0u8; additional_len]);

        tree.insert(&key, value).unwrap();
    }
}

fn run_batches(config: Config) {
    fn do_batch(i: u32, tree: &sled::Tree) {
        let mut rng = rand::thread_rng();
        let base_value = u32_to_vec(i);

        let mut batch = sled::Batch::default();
        if rng.gen_bool(0.1) {
            for key in 0..BATCH_SIZE {
                batch.remove(u32_to_vec(key));
            }
        } else {
            for key in 0..BATCH_SIZE {
                let mut value = base_value.clone();
                let additional_len = rng.gen_range(0, 1000);
                value.append(&mut vec![0u8; additional_len]);

                batch.insert(u32_to_vec(key), value);
            }
        }
        tree.apply_batch(batch).unwrap();
    }

    common::setup_logger();
    let crash_during_initialization = rand::thread_rng().gen_bool(0.1);

    if crash_during_initialization {
        spawn_killah();
    }

    let tree = sled::Db::start(config).unwrap();

    let mut i = verify_batches(&tree);
    i += 1;
    do_batch(i, &tree);

    if !crash_during_initialization {
        spawn_killah();
    }

    loop {
        i += 1;
        do_batch(i, &tree);
    }
}

fn run_without_snapshot(dir: &str) {
    let mut config_builder = ConfigBuilder::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(1 << 56);

    config_builder.io_buf_size = 1024;

    let config = config_builder.build();

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_with_snapshot(dir: &str) {
    let mut config_builder = ConfigBuilder::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(5000);

    config_builder.io_buf_size = 1024;

    let config = config_builder.build();

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_batches_without_snapshot(dir: &str) {
    let mut config_builder = ConfigBuilder::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(1 << 56);

    config_builder.io_buf_size = 1024;

    let config = config_builder.build();

    match thread::spawn(|| run_batches(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_batches_with_snapshot(dir: &str) {
    let mut config_builder = ConfigBuilder::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(5000);

    config_builder.io_buf_size = 1024;

    let config = config_builder.build();

    match thread::spawn(|| run_batches(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

#[test]
fn test_crash_recovery_with_runtime_snapshot() {
    let dir = "test_crashes_with_snapshot";
    cleanup(dir);
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_with_snapshot(dir)
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup(dir);
                panic!("child exited abnormally");
            }
        }
    }
    cleanup(dir);
}

#[test]
fn test_crash_recovery_no_runtime_snapshot() {
    let dir = "test_crashes";
    cleanup(dir);
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_without_snapshot(dir)
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup(dir);
                panic!("child exited abnormally");
            }
        }
    }
    cleanup(dir);
}

#[test]
#[ignore]
fn test_crash_batches_with_runtime_snapshot() {
    let dir = "test_batches_with_snapshot";
    cleanup(dir);
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_batches_with_snapshot(dir)
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup(dir);
                panic!("child exited abnormally");
            }
        }
    }
    cleanup(dir);
}

#[test]
#[ignore]
fn test_crash_batches_no_runtime_snapshot() {
    let dir = "test_batches";
    cleanup(dir);
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_batches_without_snapshot(dir)
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup(dir);
                panic!("child exited abnormally");
            }
        }
    }
    cleanup(dir);
}

fn cleanup(dir: &str) {
    let dir = Path::new(dir);
    if dir.exists() {
        fs::remove_dir_all(dir).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::f64;
use error::StatsError;

pub trait Statistics {
    /// Returns the minimum value in the data
    ///
    /// # Rermarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn min(&self) -> f64;

    /// Returns the maximum value in the data
    ///
    /// # Rermarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn max(&self) -> f64;

    /// Returns the minimum absolute value in the data
    ///
    /// # Rermarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn abs_min(&self) -> f64;

    /// Returns the maximum absolute value in the data
    ///
    /// # Rermarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn abs_max(&self) -> f64;

    /// Evaluates the sample mean, an estimate of the population
    /// mean.
    ///
    /// # Remarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn mean(&self) -> f64;

    /// Evaluates the geometric mean of the data
    ///
    /// # Remarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn geometric_mean(&self) -> f64;

    /// Evaluates the harmonic mean of the data
    ///
    /// # Remarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn harmonic_mean(&self) -> f64;

    /// Estimates the unbiased population variance from the provided samples
    ///
    /// # Remarks
    ///
    /// On a dataset of size `N`, `N-1` is used as a normalizer (Bessel's correction).
    ///
    /// Returns `f64::NAN` if data has less than two entries or if any entry is `f64::NAN`
    fn variance(&self) -> f64;

    /// Evaluates the population variance from a full population.
    ///
    /// # Remarks
    ///
    /// On a dataset of size `N`, `N` is used as a normalizer and would thus
    /// be biased if applied to a subset
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn population_variance(&self) -> f64;

    /// Estimates the unbiased population standard deviation from the provided samples
    ///
    /// # Remarks
    ///
    /// On a dataset of size `N`, `N-1` is used as a normalizer (Bessel's correction).
    ///
    /// Returns `f64::NAN` if data has less than two entries or if any entry is `f64::NAN`
    fn std_dev(&self) -> f64;

    /// Evaluates the population standard deviation from a full population.
    ///
    /// # Remarks
    ///
    /// On a dataset of size `N`, `N` is used as a normalizer and would thus
    /// be biased if applied to a subset
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn population_std_dev(&self) -> f64;

    /// Estimates the unbiased population covariance between the two provided samples
    ///
    /// # Remarks
    ///
    /// On a dataset of size `N`, `N-1` is used as a normalizer (Bessel's correction).
    ///
    /// Returns `f64::NAN` if data has less than two entries or if any entry is `f64::NAN`
    ///
    /// # Panics
    ///
    /// If the two sample containers do not contain the same number of elements
    fn covariance(&self, other: &Self) -> f64;

    /// Evaluates the population covariance between the two provider populations
    ///
    /// # Remarks
    ///
    /// On a dataset of size `N`, `N` is used as a normalizer and would thus be
    /// biased if applied to a subset
    ///
    /// Returns `f64::NAN` if data is empty or any entry is `f64::NAN`
    ///
    /// # Panics
    ///
    /// If the two sample containers do not contain the same number of elements
    fn population_covariance(&self, other: &Self) -> f64;

    /// Estimates the quadratic mean (Root Mean Square) of the data
    ///
    /// # Remarks
    ///
    /// Returns `f64::NAN` if data is empty or any entry is `f64::NAN`
    fn quadratic_mean(&self) -> f64;

    /// Returns the order statistic `(order 1..N)` from the data
    ///
    /// # Remarks
    ///
    /// No sorting is assumed. Order must be one-based (between `1` and `N` inclusive)
    /// Returns `f64::NAN` if order is outside the viable range or data is empty.
    fn order_statistic(&mut self, order: usize) -> f64;
}

impl Statistics for [f64] {
    fn min(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        self.iter().fold(f64::INFINITY, |acc, &x| if x < acc || x.is_nan() { x } else { acc })
    }

    fn max(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        self.iter().fold(f64::NEG_INFINITY, |acc, &x| if x > acc || x.is_nan() { x } else { acc })
    }

    fn abs_min(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        self.iter()
            .map(|x| x.abs())
            .fold(f64::INFINITY, |acc, x| if x < acc || x.is_nan() { x } else { acc })
    }

    fn abs_max(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        self.iter()
            .map(|x| x.abs())
            .fold(f64::NEG_INFINITY, |acc, x| if x > acc || x.is_nan() { x } else { acc })
    }

    fn mean(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        let mut m = 0.0;
        self.iter()
            .fold(0.0, |acc, &x| {
                m += 1.0;
                acc + (x - acc) / m
            })
    }

    fn geometric_mean(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        (self.iter().fold(0.0, |acc, &x| acc + x.ln()) / self.len() as f64).exp()
    }

    fn harmonic_mean(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        self.len() as f64 / self.iter().fold(0.0, |acc, &x| acc + 1.0 / x)
    }

    fn variance(&self) -> f64 {
        if self.len() <= 1 {
            return f64::NAN;
        }

        unsafe {
            let mut var = 0.0;
            let mut t = *self.get_unchecked(0);
            for i in 1..self.len() {
                t += *self.get_unchecked(i);
                let diff = (i as f64 + 1.0) * *self.get_unchecked(i) - t;
                var += (diff * diff) / ((i + 1) * i) as f64;
            }
            var / (self.len() - 1) as f64
        }
    }

    fn population_variance(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        unsafe {
            let mut var = 0.0;
            let mut t = *self.get_unchecked(0);
            for i in 1..self.len() {
                t += *self.get_unchecked(i);
                let diff = (i as f64 + 1.0) * *self.get_unchecked(i) - t;
                var += (diff * diff) / ((i + 1) * i) as f64
            }
            var / self.len() as f64
        }
    }

    fn std_dev(&self) -> f64 {
        self.variance().sqrt()
    }

    fn population_std_dev(&self) -> f64 {
        self.population_variance().sqrt()
    }

    fn covariance(&self, other: &[f64]) -> f64 {
        let n1 = self.len();
        let n2 = other.len();
        assert!(n1 == n2, format!("{}", StatsError::VectorsSameLength));
        if n1 <= 1 {
            return f64::NAN;
        }

        let mean1 = self.mean();
        let mean2 = other.mean();
        self.iter()
            .zip(other.iter())
            .fold(0.0, |acc, x| acc + (x.0 - mean1) * (x.1 - mean2)) / (n1 - 1) as f64
    }

    fn population_covariance(&self, other: &[f64]) -> f64 {
        let n1 = self.len();
        let n2 = other.len();
        assert!(n1 == n2, format!("{}", StatsError::VectorsSameLength));
        if n1 == 0 {
            return f64::NAN;
        }

        let mean1 = self.mean();
        let mean2 = other.mean();
        self.iter()
            .zip(other.iter())
            .fold(0.0, |acc, x| acc + (x.0 - mean1) * (x.1 - mean2)) / n1 as f64
    }

    fn quadratic_mean(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        let mut m = 0.0;
        self.iter()
            .fold(0.0, |acc, &x| {
                m += 1.0;
                acc + (x * x - acc) / m
            })
    }

    /// Returns the order statistic `(order 1..N)` from the data
    ///
    /// # Remarks
    ///
    /// No sorting is assumed. Order must be one-based (between `1` and `N` inclusive).
    /// Returns `f64::NAN` if order is outside the viable range or data is empty.
    ///
    /// **NOTE:** This method works inplace for arrays and may cause the array to be reordered
    fn order_statistic(&mut self, order: usize) -> f64 {
        let n = self.len();
        match order {
            1 => self.min(),
            _ if order == n => self.max(),
            _ if order < 1 || order > n => f64::NAN,
            _ => select_inplace(self, order - 1)
        }
    }
}

// Selection algorithm from Numerical Recipes
// See: https://en.wikipedia.org/wiki/Selection_algorithm
fn select_inplace(arr: &mut [f64], rank: usize) -> f64 {
    if rank == 0 {
        return arr.min();
    }
    if rank > arr.len() - 1 {
        return arr.max();
    }

    unsafe {
        let mut low = 0;
        let mut high = arr.len() - 1;
        loop {
            if high <= low + 1 {
                if high == low + 1 && *arr.get_unchecked(high) < *arr.get_unchecked(low) {
                    arr.swap(low, high)
                }
                return *arr.get_unchecked(rank);
            }

            let middle = (low + high) >> 1;
            arr.swap(middle, low + 1);

            if *arr.get_unchecked(low) > *arr.get_unchecked(high) {
                arr.swap(low, high);
            }
            if *arr.get_unchecked(low + 1) > *arr.get_unchecked(high) {
                arr.swap(low + 1, high);
            }
            if *arr.get_unchecked(low) > *arr.get_unchecked(low + 1) {
                arr.swap(low, low + 1);
            }

            let mut begin = low + 1;
            let mut end = high;
            let pivot = *arr.get_unchecked(begin);
            loop {
                loop {
                    begin += 1;
                    if *arr.get_unchecked(begin) < pivot {
                        break;
                    }
                }
                loop {
                    end -= 1;
                    if *arr.get_unchecked(end) > pivot {
                        break;
                    }
                }
                if end < begin {
                    break;
                }
                arr.swap(begin, end);
            }

            arr[low + 1] = *arr.get_unchecked(end);
            arr[end] = pivot;

            if end >= rank {
                high = end - 1;
            }
            if end <= rank {
                low = begin;
            }
        }
    }
}
<commit_msg>implement median and quantile for array statistics<commit_after>use std::f64;
use error::StatsError;

pub trait Statistics {
    /// Returns the minimum value in the data
    ///
    /// # Rermarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn min(&self) -> f64;

    /// Returns the maximum value in the data
    ///
    /// # Rermarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn max(&self) -> f64;

    /// Returns the minimum absolute value in the data
    ///
    /// # Rermarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn abs_min(&self) -> f64;

    /// Returns the maximum absolute value in the data
    ///
    /// # Rermarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn abs_max(&self) -> f64;

    /// Evaluates the sample mean, an estimate of the population
    /// mean.
    ///
    /// # Remarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn mean(&self) -> f64;

    /// Evaluates the geometric mean of the data
    ///
    /// # Remarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn geometric_mean(&self) -> f64;

    /// Evaluates the harmonic mean of the data
    ///
    /// # Remarks
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn harmonic_mean(&self) -> f64;

    /// Estimates the unbiased population variance from the provided samples
    ///
    /// # Remarks
    ///
    /// On a dataset of size `N`, `N-1` is used as a normalizer (Bessel's correction).
    ///
    /// Returns `f64::NAN` if data has less than two entries or if any entry is `f64::NAN`
    fn variance(&self) -> f64;

    /// Evaluates the population variance from a full population.
    ///
    /// # Remarks
    ///
    /// On a dataset of size `N`, `N` is used as a normalizer and would thus
    /// be biased if applied to a subset
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn population_variance(&self) -> f64;

    /// Estimates the unbiased population standard deviation from the provided samples
    ///
    /// # Remarks
    ///
    /// On a dataset of size `N`, `N-1` is used as a normalizer (Bessel's correction).
    ///
    /// Returns `f64::NAN` if data has less than two entries or if any entry is `f64::NAN`
    fn std_dev(&self) -> f64;

    /// Evaluates the population standard deviation from a full population.
    ///
    /// # Remarks
    ///
    /// On a dataset of size `N`, `N` is used as a normalizer and would thus
    /// be biased if applied to a subset
    ///
    /// Returns `f64::NAN` if data is empty or an entry is `f64::NAN`
    fn population_std_dev(&self) -> f64;

    /// Estimates the unbiased population covariance between the two provided samples
    ///
    /// # Remarks
    ///
    /// On a dataset of size `N`, `N-1` is used as a normalizer (Bessel's correction).
    ///
    /// Returns `f64::NAN` if data has less than two entries or if any entry is `f64::NAN`
    ///
    /// # Panics
    ///
    /// If the two sample containers do not contain the same number of elements
    fn covariance(&self, other: &Self) -> f64;

    /// Evaluates the population covariance between the two provider populations
    ///
    /// # Remarks
    ///
    /// On a dataset of size `N`, `N` is used as a normalizer and would thus be
    /// biased if applied to a subset
    ///
    /// Returns `f64::NAN` if data is empty or any entry is `f64::NAN`
    ///
    /// # Panics
    ///
    /// If the two sample containers do not contain the same number of elements
    fn population_covariance(&self, other: &Self) -> f64;

    /// Estimates the quadratic mean (Root Mean Square) of the data
    ///
    /// # Remarks
    ///
    /// Returns `f64::NAN` if data is empty or any entry is `f64::NAN`
    fn quadratic_mean(&self) -> f64;

    /// Returns the order statistic `(order 1..N)` from the data
    ///
    /// # Remarks
    ///
    /// No sorting is assumed. Order must be one-based (between `1` and `N` inclusive)
    /// Returns `f64::NAN` if order is outside the viable range or data is empty.
    fn order_statistic(&mut self, order: usize) -> f64;

    /// Returns the median value from the data
    ///
    /// # Remarks
    ///
    /// Returns `f64::NAN` if data is empty
    fn median(&mut self) -> f64;

    /// Estimates the tau-th quantile from the data. The tau-th quantile
    /// is the data value where the cumulative distribution function crosses tau.
    ///
    /// # Remarks
    ///
    /// No sorting is assumed. Tau must be between `0` and `1` inclusive.
    /// Returns `f64::NAN` if data is empty or tau is outside the inclusive range.
    fn quantile(&mut self, tau: f64) -> f64;
}

impl Statistics for [f64] {
    fn min(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        self.iter().fold(f64::INFINITY, |acc, &x| if x < acc || x.is_nan() { x } else { acc })
    }

    fn max(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        self.iter().fold(f64::NEG_INFINITY, |acc, &x| if x > acc || x.is_nan() { x } else { acc })
    }

    fn abs_min(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        self.iter()
            .map(|x| x.abs())
            .fold(f64::INFINITY, |acc, x| if x < acc || x.is_nan() { x } else { acc })
    }

    fn abs_max(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        self.iter()
            .map(|x| x.abs())
            .fold(f64::NEG_INFINITY, |acc, x| if x > acc || x.is_nan() { x } else { acc })
    }

    fn mean(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        let mut m = 0.0;
        self.iter()
            .fold(0.0, |acc, &x| {
                m += 1.0;
                acc + (x - acc) / m
            })
    }

    fn geometric_mean(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        (self.iter().fold(0.0, |acc, &x| acc + x.ln()) / self.len() as f64).exp()
    }

    fn harmonic_mean(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        self.len() as f64 / self.iter().fold(0.0, |acc, &x| acc + 1.0 / x)
    }

    fn variance(&self) -> f64 {
        if self.len() <= 1 {
            return f64::NAN;
        }

        unsafe {
            let mut var = 0.0;
            let mut t = *self.get_unchecked(0);
            for i in 1..self.len() {
                t += *self.get_unchecked(i);
                let diff = (i as f64 + 1.0) * *self.get_unchecked(i) - t;
                var += (diff * diff) / ((i + 1) * i) as f64;
            }
            var / (self.len() - 1) as f64
        }
    }

    fn population_variance(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        unsafe {
            let mut var = 0.0;
            let mut t = *self.get_unchecked(0);
            for i in 1..self.len() {
                t += *self.get_unchecked(i);
                let diff = (i as f64 + 1.0) * *self.get_unchecked(i) - t;
                var += (diff * diff) / ((i + 1) * i) as f64
            }
            var / self.len() as f64
        }
    }

    fn std_dev(&self) -> f64 {
        self.variance().sqrt()
    }

    fn population_std_dev(&self) -> f64 {
        self.population_variance().sqrt()
    }

    fn covariance(&self, other: &[f64]) -> f64 {
        let n1 = self.len();
        let n2 = other.len();
        assert!(n1 == n2, format!("{}", StatsError::VectorsSameLength));
        if n1 <= 1 {
            return f64::NAN;
        }

        let mean1 = self.mean();
        let mean2 = other.mean();
        self.iter()
            .zip(other.iter())
            .fold(0.0, |acc, x| acc + (x.0 - mean1) * (x.1 - mean2)) / (n1 - 1) as f64
    }

    fn population_covariance(&self, other: &[f64]) -> f64 {
        let n1 = self.len();
        let n2 = other.len();
        assert!(n1 == n2, format!("{}", StatsError::VectorsSameLength));
        if n1 == 0 {
            return f64::NAN;
        }

        let mean1 = self.mean();
        let mean2 = other.mean();
        self.iter()
            .zip(other.iter())
            .fold(0.0, |acc, x| acc + (x.0 - mean1) * (x.1 - mean2)) / n1 as f64
    }

    fn quadratic_mean(&self) -> f64 {
        if self.len() == 0 {
            return f64::NAN;
        }

        let mut m = 0.0;
        self.iter()
            .fold(0.0, |acc, &x| {
                m += 1.0;
                acc + (x * x - acc) / m
            })
    }

    /// Returns the order statistic `(order 1..N)` from the data
    ///
    /// # Remarks
    ///
    /// No sorting is assumed. Order must be one-based (between `1` and `N` inclusive).
    /// Returns `f64::NAN` if order is outside the viable range or data is empty.
    ///
    /// **NOTE:** This method works inplace for arrays and may cause the array to be reordered
    fn order_statistic(&mut self, order: usize) -> f64 {
        let n = self.len();
        match order {
            1 => self.min(),
            _ if order == n => self.max(),
            _ if order < 1 || order > n => f64::NAN,
            _ => select_inplace(self, order - 1)
        }
    }

    /// Returns the median value from the data
    ///
    /// # Remarks
    ///
    /// Returns `f64::NAN` if data is empty
    ///
    /// **NOTE:** This method works inplace for arrays and may cause the array to be reordered
    fn median(&mut self) -> f64 {
        let k = self.len() / 2;
        if self.len() % 2 != 0 {
            select_inplace(self, k)
        } else {
            (select_inplace(self, k.saturating_sub(1)) + select_inplace(self, k)) / 2.0
        }
    }

    /// Estimates the tau-th quantile from the data. The tau-th quantile
    /// is the data value where the cumulative distribution function crosses tau.
    ///
    /// # Remarks
    ///
    /// No sorting is assumed. Tau must be between `0` and `1` inclusive.
    /// Returns `f64::NAN` if data is empty or tau is outside the inclusive range.
    ///
    /// **NOTE:** This method works inplace for arrays and may cause the array to be reordered
    fn quantile(&mut self, tau: f64) -> f64 {
        if tau < 0.0 || tau > 1.0 || self.len() == 0 {
            return f64::NAN;
        }

        let h = (self.len() as f64 + 1.0 / 3.0) * tau + 1.0 / 3.0;
        let hf = h as i64;

        if hf <= 0 || tau == 0.0 {
            return self.min();
        }
        if hf >= self.len() as i64 || tau == 1.0 {
            return self.max();
        }

        let a = select_inplace(self, (hf as usize).saturating_sub(1));
        let b = select_inplace(self, hf as usize);
        a + (h - hf as f64) * (b - a)
    }
}

// Selection algorithm from Numerical Recipes
// See: https://en.wikipedia.org/wiki/Selection_algorithm
fn select_inplace(arr: &mut [f64], rank: usize) -> f64 {
    if rank == 0 {
        return arr.min();
    }
    if rank > arr.len() - 1 {
        return arr.max();
    }

    unsafe {
        let mut low = 0;
        let mut high = arr.len() - 1;
        loop {
            if high <= low + 1 {
                if high == low + 1 && *arr.get_unchecked(high) < *arr.get_unchecked(low) {
                    arr.swap(low, high)
                }
                return *arr.get_unchecked(rank);
            }

            let middle = (low + high) >> 1;
            arr.swap(middle, low + 1);

            if *arr.get_unchecked(low) > *arr.get_unchecked(high) {
                arr.swap(low, high);
            }
            if *arr.get_unchecked(low + 1) > *arr.get_unchecked(high) {
                arr.swap(low + 1, high);
            }
            if *arr.get_unchecked(low) > *arr.get_unchecked(low + 1) {
                arr.swap(low, low + 1);
            }

            let mut begin = low + 1;
            let mut end = high;
            let pivot = *arr.get_unchecked(begin);
            loop {
                loop {
                    begin += 1;
                    if *arr.get_unchecked(begin) < pivot {
                        break;
                    }
                }
                loop {
                    end -= 1;
                    if *arr.get_unchecked(end) > pivot {
                        break;
                    }
                }
                if end < begin {
                    break;
                }
                arr.swap(begin, end);
            }

            arr[low + 1] = *arr.get_unchecked(end);
            arr[end] = pivot;

            if end >= rank {
                high = end - 1;
            }
            if end <= rank {
                low = begin;
            }
        }
    }
}
<|endoftext|>


<commit_before>use mio::Poll;
use mio::Token;
use net2::TcpBuilder;
use net2::TcpStreamExt;
use std::io::Read;
use std::io::Write;
use std::io;
use std::net::SocketAddr;
use std::net::ToSocketAddrs;
use std::net;
use std::os::unix::io::AsRawFd;
use std::result::Result;

use stream::utils::*;

/// A buffered TCP stream that registers itself for readiness events when the buffer is not empty.
/// Use `Write` instance for writing to the buffer. Use `Read` instance when the stream is ready
/// for reading. Use `send` when the stream is ready for writing.
pub struct TcpStream<'poll> {
    inner: net::TcpStream,
    poll: &'poll Poll,
    out_buf: Vec<u8>,
}

#[derive(Debug)]
pub enum TcpError {
    IoError(io::Error),
    CantResolveAddr,
}

impl<'poll> TcpStream<'poll> {
    pub fn new(
        poll: &'poll Poll,
        serv_addr: &str,
        serv_port: u16,
    ) -> Result<TcpStream<'poll>, TcpError> {
        let mut addr_iter = (serv_addr, serv_port)
            .to_socket_addrs()
            .map_err(TcpError::IoError)?;
        let addr = addr_iter.next().ok_or(TcpError::CantResolveAddr)?;
        let stream = {
            match addr {
                SocketAddr::V4(_) =>
                    TcpBuilder::new_v4().unwrap().to_tcp_stream().unwrap(),
                SocketAddr::V6(_) =>
                    TcpBuilder::new_v6().unwrap().to_tcp_stream().unwrap(),
            }
        };
        stream.set_nonblocking(true).unwrap();
        // This will fail with EINPROGRESS
        let _ = stream.connect(addr);
        register_for_r(poll, stream.as_raw_fd());
        Ok(TcpStream {
            inner: stream,
            poll: poll,
            out_buf: Vec::with_capacity(1024),
        })
    }

    /// Call when the stream is ready for writing.
    pub fn write_ready(&mut self) -> io::Result<()> {
        let to_send = self.out_buf.len();
        match self.inner.write(&self.out_buf) {
            Ok(bytes_sent) => {
                self.out_buf.drain(0..bytes_sent);
                let register = if bytes_sent == to_send {
                    reregister_for_r
                } else {
                    reregister_for_rw
                };
                register(&self.poll, self.inner.as_raw_fd());
                Ok(())
            }
            Err(err) => {
                reregister_for_rw(&self.poll, self.inner.as_raw_fd());
                if err.kind() == io::ErrorKind::WouldBlock {
                    Ok(())
                } else {
                    Err(err)
                }
            }
        }
    }

    pub fn get_tok(&self) -> Token {
        Token(self.inner.as_raw_fd() as usize)
    }
}

// Drop just deregisters the socket
impl<'poll> Drop for TcpStream<'poll> {
    fn drop(&mut self) {
        deregister(&self.poll, self.inner.as_raw_fd());
    }
}

impl<'poll> Read for TcpStream<'poll> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.inner.read(buf)
    }
}

impl<'poll> Write for TcpStream<'poll> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        // TODO inefficient when the socket is already ready for writing
        self.out_buf.extend(buf);
        reregister_for_rw(&self.poll, self.inner.as_raw_fd());
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        self.write_ready()
    }
}
<commit_msg>Remove some unwraps<commit_after>use mio::Poll;
use mio::Token;
use net2::TcpBuilder;
use net2::TcpStreamExt;
use std::io::Read;
use std::io::Write;
use std::io;
use std::net::SocketAddr;
use std::net::ToSocketAddrs;
use std::net;
use std::os::unix::io::AsRawFd;
use std::result::Result;

use stream::utils::*;

/// A buffered TCP stream that registers itself for readiness events when the buffer is not empty.
/// Use `Write` instance for writing to the buffer. Use `Read` instance when the stream is ready
/// for reading. Use `send` when the stream is ready for writing.
pub struct TcpStream<'poll> {
    inner: net::TcpStream,
    poll: &'poll Poll,
    out_buf: Vec<u8>,
}

#[derive(Debug)]
pub enum TcpError {
    IoError(io::Error),
    CantResolveAddr,
}

impl<'poll> TcpStream<'poll> {
    pub fn new(
        poll: &'poll Poll,
        serv_addr: &str,
        serv_port: u16,
    ) -> Result<TcpStream<'poll>, TcpError> {
        let mut addr_iter = (serv_addr, serv_port)
            .to_socket_addrs()
            .map_err(TcpError::IoError)?;
        let addr = addr_iter.next().ok_or(TcpError::CantResolveAddr)?;
        let stream = {
            match addr {
                SocketAddr::V4(_) =>
                    TcpBuilder::new_v4()
                        .map_err(TcpError::IoError)?
                        .to_tcp_stream()
                        .map_err(TcpError::IoError)?,
                SocketAddr::V6(_) =>
                    TcpBuilder::new_v6()
                        .map_err(TcpError::IoError)?
                        .to_tcp_stream()
                        .map_err(TcpError::IoError)?,
            }
        };
        stream.set_nonblocking(true).map_err(TcpError::IoError)?;
        // This will fail with EINPROGRESS
        let _ = stream.connect(addr);
        register_for_r(poll, stream.as_raw_fd());
        Ok(TcpStream {
            inner: stream,
            poll: poll,
            out_buf: Vec::with_capacity(1024),
        })
    }

    /// Call when the stream is ready for writing.
    pub fn write_ready(&mut self) -> io::Result<()> {
        let to_send = self.out_buf.len();
        match self.inner.write(&self.out_buf) {
            Ok(bytes_sent) => {
                self.out_buf.drain(0..bytes_sent);
                let register = if bytes_sent == to_send {
                    reregister_for_r
                } else {
                    reregister_for_rw
                };
                register(&self.poll, self.inner.as_raw_fd());
                Ok(())
            }
            Err(err) => {
                reregister_for_rw(&self.poll, self.inner.as_raw_fd());
                if err.kind() == io::ErrorKind::WouldBlock {
                    Ok(())
                } else {
                    Err(err)
                }
            }
        }
    }

    pub fn get_tok(&self) -> Token {
        Token(self.inner.as_raw_fd() as usize)
    }
}

// Drop just deregisters the socket
impl<'poll> Drop for TcpStream<'poll> {
    fn drop(&mut self) {
        deregister(&self.poll, self.inner.as_raw_fd());
    }
}

impl<'poll> Read for TcpStream<'poll> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.inner.read(buf)
    }
}

impl<'poll> Write for TcpStream<'poll> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        // TODO inefficient when the socket is already ready for writing
        self.out_buf.extend(buf);
        reregister_for_rw(&self.poll, self.inner.as_raw_fd());
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        self.write_ready()
    }
}
<|endoftext|>


<commit_before>// Portions of this file are Copyright 2014 The Rust Project Developers.
// See http://rust-lang.org/COPYRIGHT.

use libc;
use core::mem;
use errno::{SysError, SysResult};

pub use libc::consts::os::posix88::{
    SIGHUP,   // 1
    SIGINT,   // 2
    SIGQUIT,  // 3
    SIGILL,   // 4
    SIGABRT,  // 6
    SIGFPE,   // 8
    SIGKILL,  // 9
    SIGSEGV,  // 11
    SIGPIPE,  // 13
    SIGALRM,  // 14
    SIGTERM,  // 15
};

pub use self::signal::{
    SIGTRAP,
    SIGIOT,
    SIGBUS,
    SIGSYS,
    SIGURG,
    SIGSTOP,
    SIGTSTP,
    SIGCONT,
    SIGCHLD,
    SIGTTIN,
    SIGTTOU,
    SIGIO,
    SIGXCPU,
    SIGXFSZ,
    SIGVTALRM,
    SIGPROF,
    SIGWINCH,
    SIGUSR1,
    SIGUSR2,
};

pub use self::signal::SockFlag;
pub use self::signal::sigset_t;

// This doesn't always exist, but when it does, it's 7
pub const SIGEMT: libc::c_int = 7;

#[cfg(any(all(target_os = "linux",
              any(target_arch = "x86",
                  target_arch = "x86_64",
                  target_arch = "arm")),
          target_os = "android"))]
pub mod signal {
    use libc;

    bitflags!(
        flags SockFlag: libc::c_ulong {
            const SA_NOCLDSTOP = 0x00000001,
            const SA_NOCLDWAIT = 0x00000002,
            const SA_NODEFER   = 0x40000000,
            const SA_ONSTACK   = 0x08000000,
            const SA_RESETHAND = 0x80000000,
            const SA_RESTART   = 0x10000000,
            const SA_SIGINFO   = 0x00000004,
        }
    );

    pub const SIGTRAP:      libc::c_int = 5;
    pub const SIGIOT:       libc::c_int = 6;
    pub const SIGBUS:       libc::c_int = 7;
    pub const SIGUSR1:      libc::c_int = 10;
    pub const SIGUSR2:      libc::c_int = 12;
    pub const SIGSTKFLT:    libc::c_int = 16;
    pub const SIGCHLD:      libc::c_int = 17;
    pub const SIGCONT:      libc::c_int = 18;
    pub const SIGSTOP:      libc::c_int = 19;
    pub const SIGTSTP:      libc::c_int = 20;
    pub const SIGTTIN:      libc::c_int = 21;
    pub const SIGTTOU:      libc::c_int = 22;
    pub const SIGURG:       libc::c_int = 23;
    pub const SIGXCPU:      libc::c_int = 24;
    pub const SIGXFSZ:      libc::c_int = 25;
    pub const SIGVTALRM:    libc::c_int = 26;
    pub const SIGPROF:      libc::c_int = 27;
    pub const SIGWINCH:     libc::c_int = 28;
    pub const SIGIO:        libc::c_int = 29;
    pub const SIGPOLL:      libc::c_int = 29;
    pub const SIGPWR:       libc::c_int = 30;
    pub const SIGSYS:       libc::c_int = 31;
    pub const SIGUNUSED:    libc::c_int = 31;

    // This definition is not as accurate as it could be, {pid, uid, status} is
    // actually a giant union. Currently we're only interested in these fields,
    // however.
    #[repr(C)]
    #[derive(Copy)]
    pub struct siginfo {
        si_signo: libc::c_int,
        si_errno: libc::c_int,
        si_code: libc::c_int,
        pub pid: libc::pid_t,
        pub uid: libc::uid_t,
        pub status: libc::c_int,
    }

    #[repr(C)]
    #[allow(missing_copy_implementations)]
    pub struct sigaction {
        pub sa_handler: extern fn(libc::c_int),
        pub sa_mask: sigset_t,
        pub sa_flags: SockFlag,
        sa_restorer: *mut libc::c_void,
    }

    #[repr(C)]
    #[cfg(target_word_size = "32")]
    #[derive(Copy)]
    pub struct sigset_t {
        __val: [libc::c_ulong; 32],
    }

    #[repr(C)]
    #[cfg(target_word_size = "64")]
    #[derive(Copy)]
    pub struct sigset_t {
        __val: [libc::c_ulong; 16],
    }
}

#[cfg(all(target_os = "linux",
          any(target_arch = "mips", target_arch = "mipsel")))]
pub mod signal {
    use libc;

    bitflags!(
        flags SockFlag: libc::c_uint {
            const SA_NOCLDSTOP = 0x00000001,
            const SA_NOCLDWAIT = 0x00001000,
            const SA_NODEFER   = 0x40000000,
            const SA_ONSTACK   = 0x08000000,
            const SA_RESETHAND = 0x80000000,
            const SA_RESTART   = 0x10000000,
            const SA_SIGINFO   = 0x00000008,
        }
    );

    pub const SIGTRAP:      libc::c_int = 5;
    pub const SIGIOT:       libc::c_int = 6;
    pub const SIGBUS:       libc::c_int = 10;
    pub const SIGSYS:       libc::c_int = 12;
    pub const SIGUSR1:      libc::c_int = 16;
    pub const SIGUSR2:      libc::c_int = 17;
    pub const SIGCHLD:      libc::c_int = 18;
    pub const SIGCLD:       libc::c_int = 18;
    pub const SIGPWR:       libc::c_int = 19;
    pub const SIGWINCH:     libc::c_int = 20;
    pub const SIGURG:       libc::c_int = 21;
    pub const SIGIO:        libc::c_int = 22;
    pub const SIGPOLL:      libc::c_int = 22;
    pub const SIGSTOP:      libc::c_int = 23;
    pub const SIGTSTP:      libc::c_int = 24;
    pub const SIGCONT:      libc::c_int = 25;
    pub const SIGTTIN:      libc::c_int = 26;
    pub const SIGTTOU:      libc::c_int = 27;
    pub const SIGVTALRM:    libc::c_int = 28;
    pub const SIGPROF:      libc::c_int = 29;
    pub const SIGXCPU:      libc::c_int = 30;
    pub const SIGFSZ:       libc::c_int = 31;

    // This definition is not as accurate as it could be, {pid, uid, status} is
    // actually a giant union. Currently we're only interested in these fields,
    // however.
    #[repr(C)]
    pub struct siginfo {
        si_signo: libc::c_int,
        si_code: libc::c_int,
        si_errno: libc::c_int,
        pub pid: libc::pid_t,
        pub uid: libc::uid_t,
        pub status: libc::c_int,
    }

    #[repr(C)]
    pub struct sigaction {
        pub sa_flags: SockFlag,
        pub sa_handler: extern fn(libc::c_int),
        pub sa_mask: sigset_t,
        sa_restorer: *mut libc::c_void,
        sa_resv: [libc::c_int; 1],
    }

    #[repr(C)]
    pub struct sigset_t {
        __val: [libc::c_ulong; 32],
    }
}

#[cfg(any(target_os = "macos",
          target_os = "ios",
          target_os = "freebsd",
          target_os = "dragonfly"))]
pub mod signal {
    use libc;

    bitflags!(
        flags SockFlag: libc::c_int {
            const SA_NOCLDSTOP = 0x0008,
            const SA_NOCLDWAIT = 0x0020,
            const SA_NODEFER   = 0x0010,
            const SA_ONSTACK   = 0x0001,
            const SA_RESETHAND = 0x0004,
            const SA_RESTART   = 0x0002,
            const SA_SIGINFO   = 0x0040,
        }
    );

    pub const SIGTRAP:      libc::c_int = 5;
    pub const SIGIOT:       libc::c_int = 6;
    pub const SIGBUS:       libc::c_int = 10;
    pub const SIGSYS:       libc::c_int = 12;
    pub const SIGURG:       libc::c_int = 16;
    pub const SIGSTOP:      libc::c_int = 17;
    pub const SIGTSTP:      libc::c_int = 18;
    pub const SIGCONT:      libc::c_int = 19;
    pub const SIGCHLD:      libc::c_int = 20;
    pub const SIGTTIN:      libc::c_int = 21;
    pub const SIGTTOU:      libc::c_int = 22;
    pub const SIGIO:        libc::c_int = 23;
    pub const SIGXCPU:      libc::c_int = 24;
    pub const SIGXFSZ:      libc::c_int = 25;
    pub const SIGVTALRM:    libc::c_int = 26;
    pub const SIGPROF:      libc::c_int = 27;
    pub const SIGWINCH:     libc::c_int = 28;
    pub const SIGINFO:      libc::c_int = 29;
    pub const SIGUSR1:      libc::c_int = 30;
    pub const SIGUSR2:      libc::c_int = 31;

    #[cfg(any(target_os = "macos", target_os = "ios"))]
    pub type sigset_t = u32;
    #[cfg(any(target_os = "freebsd", target_os = "dragonfly"))]
    #[repr(C)]
    pub struct sigset_t {
        bits: [u32; 4],
    }

    // This structure has more fields, but we're not all that interested in
    // them.
    #[repr(C)]
    #[derive(Copy)]
    pub struct siginfo {
        pub si_signo: libc::c_int,
        pub si_errno: libc::c_int,
        pub si_code: libc::c_int,
        pub pid: libc::pid_t,
        pub uid: libc::uid_t,
        pub status: libc::c_int,
    }

    #[cfg(any(target_os = "macos", target_os = "ios"))]
    #[repr(C)]
    #[allow(missing_copy_implementations)]
    pub struct sigaction {
        pub sa_handler: extern fn(libc::c_int),
        sa_tramp: *mut libc::c_void,
        pub sa_mask: sigset_t,
        pub sa_flags: SockFlag,
    }

    #[cfg(any(target_os = "freebsd", target_os = "dragonfly"))]
    #[repr(C)]
    pub struct sigaction {
        pub sa_handler: extern fn(libc::c_int),
        pub sa_flags: SockFlag,
        pub sa_mask: sigset_t,
    }

}

mod ffi {
    use libc;
    use super::signal::{sigaction, sigset_t};

    #[allow(improper_ctypes)]
    extern {
        pub fn sigaction(signum: libc::c_int,
                         act: *const sigaction,
                         oldact: *mut sigaction) -> libc::c_int;

        pub fn sigaddset(set: *mut sigset_t, signum: libc::c_int) -> libc::c_int;
        pub fn sigdelset(set: *mut sigset_t, signum: libc::c_int) -> libc::c_int;
        pub fn sigemptyset(set: *mut sigset_t) -> libc::c_int;

        pub fn kill(pid: libc::pid_t, signum: libc::c_int) -> libc::c_int;
    }
}

#[derive(Copy)]
pub struct SigSet {
    sigset: sigset_t
}

pub type SigNum = libc::c_int;

impl SigSet {
    pub fn empty() -> SigSet {
        let mut sigset = unsafe { mem::uninitialized::<sigset_t>() };
        let _ = unsafe { ffi::sigemptyset(&mut sigset as *mut sigset_t) };

        SigSet { sigset: sigset }
    }

    pub fn add(&mut self, signum: SigNum) -> SysResult<()> {
        let res = unsafe { ffi::sigaddset(&mut self.sigset as *mut sigset_t, signum) };

        if res < 0 {
            return Err(SysError::last());
        }

        Ok(())
    }

    pub fn remove(&mut self, signum: SigNum) -> SysResult<()> {
        let res = unsafe { ffi::sigdelset(&mut self.sigset as *mut sigset_t, signum) };

        if res < 0 {
            return Err(SysError::last());
        }

        Ok(())
    }
}

type sigaction_t = self::signal::sigaction;

pub struct SigAction {
    sigaction: sigaction_t
}

impl SigAction {
    pub fn new(handler: extern fn(libc::c_int), flags: SockFlag, mask: SigSet) -> SigAction {
        let mut s = unsafe { mem::uninitialized::<sigaction_t>() };
        s.sa_handler = handler;
        s.sa_flags = flags;
        s.sa_mask = mask.sigset;

        SigAction { sigaction: s }
    }
}

pub fn sigaction(signum: SigNum, sigaction: &SigAction) -> SysResult<SigAction> {
    let mut oldact = unsafe { mem::uninitialized::<sigaction_t>() };

    let res = unsafe {
        ffi::sigaction(signum, &sigaction.sigaction as *const sigaction_t, &mut oldact as *mut sigaction_t)
    };

    if res < 0 {
        return Err(SysError::last());
    }

    Ok(SigAction { sigaction: oldact })
}

pub fn kill(pid: libc::pid_t, signum: SigNum) -> SysResult<()> {
    let res = unsafe { ffi::kill(pid, signum) };

    if res < 0 {
        return Err(SysError::last());
    }

    Ok(())
}
<commit_msg>Fix a bad cfg on linux<commit_after>// Portions of this file are Copyright 2014 The Rust Project Developers.
// See http://rust-lang.org/COPYRIGHT.

use libc;
use core::mem;
use errno::{SysError, SysResult};

pub use libc::consts::os::posix88::{
    SIGHUP,   // 1
    SIGINT,   // 2
    SIGQUIT,  // 3
    SIGILL,   // 4
    SIGABRT,  // 6
    SIGFPE,   // 8
    SIGKILL,  // 9
    SIGSEGV,  // 11
    SIGPIPE,  // 13
    SIGALRM,  // 14
    SIGTERM,  // 15
};

pub use self::signal::{
    SIGTRAP,
    SIGIOT,
    SIGBUS,
    SIGSYS,
    SIGURG,
    SIGSTOP,
    SIGTSTP,
    SIGCONT,
    SIGCHLD,
    SIGTTIN,
    SIGTTOU,
    SIGIO,
    SIGXCPU,
    SIGXFSZ,
    SIGVTALRM,
    SIGPROF,
    SIGWINCH,
    SIGUSR1,
    SIGUSR2,
};

pub use self::signal::SockFlag;
pub use self::signal::sigset_t;

// This doesn't always exist, but when it does, it's 7
pub const SIGEMT: libc::c_int = 7;

#[cfg(any(all(target_os = "linux",
              any(target_arch = "x86",
                  target_arch = "x86_64",
                  target_arch = "arm")),
          target_os = "android"))]
pub mod signal {
    use libc;

    bitflags!(
        flags SockFlag: libc::c_ulong {
            const SA_NOCLDSTOP = 0x00000001,
            const SA_NOCLDWAIT = 0x00000002,
            const SA_NODEFER   = 0x40000000,
            const SA_ONSTACK   = 0x08000000,
            const SA_RESETHAND = 0x80000000,
            const SA_RESTART   = 0x10000000,
            const SA_SIGINFO   = 0x00000004,
        }
    );

    pub const SIGTRAP:      libc::c_int = 5;
    pub const SIGIOT:       libc::c_int = 6;
    pub const SIGBUS:       libc::c_int = 7;
    pub const SIGUSR1:      libc::c_int = 10;
    pub const SIGUSR2:      libc::c_int = 12;
    pub const SIGSTKFLT:    libc::c_int = 16;
    pub const SIGCHLD:      libc::c_int = 17;
    pub const SIGCONT:      libc::c_int = 18;
    pub const SIGSTOP:      libc::c_int = 19;
    pub const SIGTSTP:      libc::c_int = 20;
    pub const SIGTTIN:      libc::c_int = 21;
    pub const SIGTTOU:      libc::c_int = 22;
    pub const SIGURG:       libc::c_int = 23;
    pub const SIGXCPU:      libc::c_int = 24;
    pub const SIGXFSZ:      libc::c_int = 25;
    pub const SIGVTALRM:    libc::c_int = 26;
    pub const SIGPROF:      libc::c_int = 27;
    pub const SIGWINCH:     libc::c_int = 28;
    pub const SIGIO:        libc::c_int = 29;
    pub const SIGPOLL:      libc::c_int = 29;
    pub const SIGPWR:       libc::c_int = 30;
    pub const SIGSYS:       libc::c_int = 31;
    pub const SIGUNUSED:    libc::c_int = 31;

    // This definition is not as accurate as it could be, {pid, uid, status} is
    // actually a giant union. Currently we're only interested in these fields,
    // however.
    #[repr(C)]
    #[derive(Copy)]
    pub struct siginfo {
        si_signo: libc::c_int,
        si_errno: libc::c_int,
        si_code: libc::c_int,
        pub pid: libc::pid_t,
        pub uid: libc::uid_t,
        pub status: libc::c_int,
    }

    #[repr(C)]
    #[allow(missing_copy_implementations)]
    pub struct sigaction {
        pub sa_handler: extern fn(libc::c_int),
        pub sa_mask: sigset_t,
        pub sa_flags: SockFlag,
        sa_restorer: *mut libc::c_void,
    }

    #[repr(C)]
    #[cfg(target_pointer_width = "32")]
    #[derive(Copy)]
    pub struct sigset_t {
        __val: [libc::c_ulong; 32],
    }

    #[repr(C)]
    #[cfg(target_pointer_width = "64")]
    #[derive(Copy)]
    pub struct sigset_t {
        __val: [libc::c_ulong; 16],
    }
}

#[cfg(all(target_os = "linux",
          any(target_arch = "mips", target_arch = "mipsel")))]
pub mod signal {
    use libc;

    bitflags!(
        flags SockFlag: libc::c_uint {
            const SA_NOCLDSTOP = 0x00000001,
            const SA_NOCLDWAIT = 0x00001000,
            const SA_NODEFER   = 0x40000000,
            const SA_ONSTACK   = 0x08000000,
            const SA_RESETHAND = 0x80000000,
            const SA_RESTART   = 0x10000000,
            const SA_SIGINFO   = 0x00000008,
        }
    );

    pub const SIGTRAP:      libc::c_int = 5;
    pub const SIGIOT:       libc::c_int = 6;
    pub const SIGBUS:       libc::c_int = 10;
    pub const SIGSYS:       libc::c_int = 12;
    pub const SIGUSR1:      libc::c_int = 16;
    pub const SIGUSR2:      libc::c_int = 17;
    pub const SIGCHLD:      libc::c_int = 18;
    pub const SIGCLD:       libc::c_int = 18;
    pub const SIGPWR:       libc::c_int = 19;
    pub const SIGWINCH:     libc::c_int = 20;
    pub const SIGURG:       libc::c_int = 21;
    pub const SIGIO:        libc::c_int = 22;
    pub const SIGPOLL:      libc::c_int = 22;
    pub const SIGSTOP:      libc::c_int = 23;
    pub const SIGTSTP:      libc::c_int = 24;
    pub const SIGCONT:      libc::c_int = 25;
    pub const SIGTTIN:      libc::c_int = 26;
    pub const SIGTTOU:      libc::c_int = 27;
    pub const SIGVTALRM:    libc::c_int = 28;
    pub const SIGPROF:      libc::c_int = 29;
    pub const SIGXCPU:      libc::c_int = 30;
    pub const SIGFSZ:       libc::c_int = 31;

    // This definition is not as accurate as it could be, {pid, uid, status} is
    // actually a giant union. Currently we're only interested in these fields,
    // however.
    #[repr(C)]
    pub struct siginfo {
        si_signo: libc::c_int,
        si_code: libc::c_int,
        si_errno: libc::c_int,
        pub pid: libc::pid_t,
        pub uid: libc::uid_t,
        pub status: libc::c_int,
    }

    #[repr(C)]
    pub struct sigaction {
        pub sa_flags: SockFlag,
        pub sa_handler: extern fn(libc::c_int),
        pub sa_mask: sigset_t,
        sa_restorer: *mut libc::c_void,
        sa_resv: [libc::c_int; 1],
    }

    #[repr(C)]
    pub struct sigset_t {
        __val: [libc::c_ulong; 32],
    }
}

#[cfg(any(target_os = "macos",
          target_os = "ios",
          target_os = "freebsd",
          target_os = "dragonfly"))]
pub mod signal {
    use libc;

    bitflags!(
        flags SockFlag: libc::c_int {
            const SA_NOCLDSTOP = 0x0008,
            const SA_NOCLDWAIT = 0x0020,
            const SA_NODEFER   = 0x0010,
            const SA_ONSTACK   = 0x0001,
            const SA_RESETHAND = 0x0004,
            const SA_RESTART   = 0x0002,
            const SA_SIGINFO   = 0x0040,
        }
    );

    pub const SIGTRAP:      libc::c_int = 5;
    pub const SIGIOT:       libc::c_int = 6;
    pub const SIGBUS:       libc::c_int = 10;
    pub const SIGSYS:       libc::c_int = 12;
    pub const SIGURG:       libc::c_int = 16;
    pub const SIGSTOP:      libc::c_int = 17;
    pub const SIGTSTP:      libc::c_int = 18;
    pub const SIGCONT:      libc::c_int = 19;
    pub const SIGCHLD:      libc::c_int = 20;
    pub const SIGTTIN:      libc::c_int = 21;
    pub const SIGTTOU:      libc::c_int = 22;
    pub const SIGIO:        libc::c_int = 23;
    pub const SIGXCPU:      libc::c_int = 24;
    pub const SIGXFSZ:      libc::c_int = 25;
    pub const SIGVTALRM:    libc::c_int = 26;
    pub const SIGPROF:      libc::c_int = 27;
    pub const SIGWINCH:     libc::c_int = 28;
    pub const SIGINFO:      libc::c_int = 29;
    pub const SIGUSR1:      libc::c_int = 30;
    pub const SIGUSR2:      libc::c_int = 31;

    #[cfg(any(target_os = "macos", target_os = "ios"))]
    pub type sigset_t = u32;
    #[cfg(any(target_os = "freebsd", target_os = "dragonfly"))]
    #[repr(C)]
    pub struct sigset_t {
        bits: [u32; 4],
    }

    // This structure has more fields, but we're not all that interested in
    // them.
    #[repr(C)]
    #[derive(Copy)]
    pub struct siginfo {
        pub si_signo: libc::c_int,
        pub si_errno: libc::c_int,
        pub si_code: libc::c_int,
        pub pid: libc::pid_t,
        pub uid: libc::uid_t,
        pub status: libc::c_int,
    }

    #[cfg(any(target_os = "macos", target_os = "ios"))]
    #[repr(C)]
    #[allow(missing_copy_implementations)]
    pub struct sigaction {
        pub sa_handler: extern fn(libc::c_int),
        sa_tramp: *mut libc::c_void,
        pub sa_mask: sigset_t,
        pub sa_flags: SockFlag,
    }

    #[cfg(any(target_os = "freebsd", target_os = "dragonfly"))]
    #[repr(C)]
    pub struct sigaction {
        pub sa_handler: extern fn(libc::c_int),
        pub sa_flags: SockFlag,
        pub sa_mask: sigset_t,
    }

}

mod ffi {
    use libc;
    use super::signal::{sigaction, sigset_t};

    #[allow(improper_ctypes)]
    extern {
        pub fn sigaction(signum: libc::c_int,
                         act: *const sigaction,
                         oldact: *mut sigaction) -> libc::c_int;

        pub fn sigaddset(set: *mut sigset_t, signum: libc::c_int) -> libc::c_int;
        pub fn sigdelset(set: *mut sigset_t, signum: libc::c_int) -> libc::c_int;
        pub fn sigemptyset(set: *mut sigset_t) -> libc::c_int;

        pub fn kill(pid: libc::pid_t, signum: libc::c_int) -> libc::c_int;
    }
}

#[derive(Copy)]
pub struct SigSet {
    sigset: sigset_t
}

pub type SigNum = libc::c_int;

impl SigSet {
    pub fn empty() -> SigSet {
        let mut sigset = unsafe { mem::uninitialized::<sigset_t>() };
        let _ = unsafe { ffi::sigemptyset(&mut sigset as *mut sigset_t) };

        SigSet { sigset: sigset }
    }

    pub fn add(&mut self, signum: SigNum) -> SysResult<()> {
        let res = unsafe { ffi::sigaddset(&mut self.sigset as *mut sigset_t, signum) };

        if res < 0 {
            return Err(SysError::last());
        }

        Ok(())
    }

    pub fn remove(&mut self, signum: SigNum) -> SysResult<()> {
        let res = unsafe { ffi::sigdelset(&mut self.sigset as *mut sigset_t, signum) };

        if res < 0 {
            return Err(SysError::last());
        }

        Ok(())
    }
}

type sigaction_t = self::signal::sigaction;

pub struct SigAction {
    sigaction: sigaction_t
}

impl SigAction {
    pub fn new(handler: extern fn(libc::c_int), flags: SockFlag, mask: SigSet) -> SigAction {
        let mut s = unsafe { mem::uninitialized::<sigaction_t>() };
        s.sa_handler = handler;
        s.sa_flags = flags;
        s.sa_mask = mask.sigset;

        SigAction { sigaction: s }
    }
}

pub fn sigaction(signum: SigNum, sigaction: &SigAction) -> SysResult<SigAction> {
    let mut oldact = unsafe { mem::uninitialized::<sigaction_t>() };

    let res = unsafe {
        ffi::sigaction(signum, &sigaction.sigaction as *const sigaction_t, &mut oldact as *mut sigaction_t)
    };

    if res < 0 {
        return Err(SysError::last());
    }

    Ok(SigAction { sigaction: oldact })
}

pub fn kill(pid: libc::pid_t, signum: SigNum) -> SysResult<()> {
    let res = unsafe { ffi::kill(pid, signum) };

    if res < 0 {
        return Err(SysError::last());
    }

    Ok(())
}
<|endoftext|>


<commit_before>pub mod a {
    pub mod series {
        pub mod of {
            pub fn nested_modules() {}
        }
    }
}

enum TrafficLight {
    Red,
    Yellow,
    Green
}

use a::series::of;
use TrafficLight::{Red, Green};

fn main() {
    of::nested_modules();
    let red = Red;
    let yellow = TrafficLight::Yellow;
    let green = Green;
}

<commit_msg>Example of glob import with Enum<commit_after>pub mod a {
    pub mod series {
        pub mod of {
            pub fn nested_modules() {}
        }
    }
}

enum TrafficLight {
    Red,
    Yellow,
    Green
}

use a::series::of;
use TrafficLight::{Red, Green};
// Example of glob import
// use TrafficLight::*;

fn main() {
    of::nested_modules();
    let red = Red;
    let yellow = TrafficLight::Yellow;
    let green = Green;
}

<|endoftext|>