<commit_before>use hir::db::HirDatabase;
use ra_syntax::{
    SyntaxKind::{RAW_STRING, STRING},
    TextRange, TextUnit,
};
use rustc_lexer;

use crate::{Assist, AssistCtx, AssistId};

// Assist: make_raw_string
//
// Adds `r#` to a plain string literal.
//
// ```
// fn main() {
//     "Hello,<|> World!";
// }
// ```
// ->
// ```
// fn main() {
//     r#"Hello, World!"#;
// }
// ```
pub(crate) fn make_raw_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let token = ctx.find_token_at_offset(STRING)?;
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    let start_of_inside = usual_string_range.start().to_usize() + 1;
    let end_of_inside = usual_string_range.end().to_usize();
    let inside_str = &text[start_of_inside..end_of_inside];
    let mut unescaped = String::with_capacity(inside_str.len());
    let mut error = Ok(());
    rustc_lexer::unescape::unescape_str(
        inside_str,
        &mut |_, unescaped_char| match unescaped_char {
            Ok(c) => unescaped.push(c),
            Err(_) => error = Err(()),
        },
    );
    if error.is_err() {
        return None;
    }
    ctx.add_action(AssistId("make_raw_string"), "make raw string", |edit| {
        edit.target(token.text_range());
        let max_hash_streak = count_hashes(&unescaped);
        let mut hashes = String::with_capacity(max_hash_streak + 1);
        for _ in 0..hashes.capacity() {
            hashes.push('#');
        }
        edit.replace(token.text_range(), format!("r{}\"{}\"{}", hashes, unescaped, hashes));
    });
    ctx.build()
}

// Assist: make_usual_string
//
// Turns a raw string into a plain string.
//
// ```
// fn main() {
//     r#"Hello,<|> "World!""#;
// }
// ```
// ->
// ```
// fn main() {
//     "Hello, \"World!\"";
// }
// ```
pub(crate) fn make_usual_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let token = ctx.find_token_at_offset(RAW_STRING)?;
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    ctx.add_action(AssistId("make_usual_string"), "make usual string", |edit| {
        edit.target(token.text_range());
        // parse inside string to escape `"`
        let start_of_inside = usual_string_range.start().to_usize() + 1;
        let end_of_inside = usual_string_range.end().to_usize();
        let inside_str = &text[start_of_inside..end_of_inside];
        let escaped = inside_str.escape_default().to_string();
        edit.replace(token.text_range(), format!("\"{}\"", escaped));
    });
    ctx.build()
}

// Assist: add_hash
//
// Adds a hash to a raw string literal.
//
// ```
// fn main() {
//     r#"Hello,<|> World!"#;
// }
// ```
// ->
// ```
// fn main() {
//     r##"Hello, World!"##;
// }
// ```
pub(crate) fn add_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let token = ctx.find_token_at_offset(RAW_STRING)?;
    ctx.add_action(AssistId("add_hash"), "add hash to raw string", |edit| {
        edit.target(token.text_range());
        edit.insert(token.text_range().start() + TextUnit::of_char('r'), "#");
        edit.insert(token.text_range().end(), "#");
    });
    ctx.build()
}

// Assist: remove_hash
//
// Removes a hash from a raw string literal.
//
// ```
// fn main() {
//     r#"Hello,<|> World!"#;
// }
// ```
// ->
// ```
// fn main() {
//     r"Hello, World!";
// }
// ```
pub(crate) fn remove_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let token = ctx.find_token_at_offset(RAW_STRING)?;
    let text = token.text().as_str();
    if text.starts_with("r\"") {
        // no hash to remove
        return None;
    }
    ctx.add_action(AssistId("remove_hash"), "remove hash from raw string", |edit| {
        edit.target(token.text_range());
        let result = &text[2..text.len() - 1];
        let result = if result.starts_with("\"") {
            // no more hash, escape
            let internal_str = &result[1..result.len() - 1];
            format!("\"{}\"", internal_str.escape_default().to_string())
        } else {
            result.to_owned()
        };
        edit.replace(token.text_range(), format!("r{}", result));
    });
    ctx.build()
}

fn count_hashes(s: &str) -> usize {
    let mut max_hash_streak = 0usize;
    for idx in s.match_indices("\"#").map(|(i, _)| i) {
        let (_, sub) = s.split_at(idx + 1);
        let nb_hash = sub.chars().take_while(|c| *c == '#').count();
        if nb_hash > max_hash_streak {
            max_hash_streak = nb_hash;
        }
    }
    max_hash_streak
}

fn find_usual_string_range(s: &str) -> Option<TextRange> {
    Some(TextRange::from_to(
        TextUnit::from(s.find('"')? as u32),
        TextUnit::from(s.rfind('"')? as u32),
    ))
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};

    #[test]
    fn make_raw_string_target() {
        check_assist_target(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r#""random\nstring""#,
        );
    }

    #[test]
    fn make_raw_string_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random
string"#;
            }
            "##,
        )
    }

    #[test]
    fn make_raw_string_works_inside_macros() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                format!(<|>"x = {}", 92)
            }
            "#,
            r##"
            fn f() {
                format!(<|>r#"x = {}"#, 92)
            }
            "##,
        )
    }

    #[test]
    fn make_raw_string_hashes_inside_works() {
        check_assist(
            make_raw_string,
            r###"
            fn f() {
                let s = <|>"#random##\nstring";
            }
            "###,
            r####"
            fn f() {
                let s = <|>r#"#random##
string"#;
            }
            "####,
        )
    }

    #[test]
    fn make_raw_string_closing_hashes_inside_works() {
        check_assist(
            make_raw_string,
            r###"
            fn f() {
                let s = <|>"#random\"##\nstring";
            }
            "###,
            r####"
            fn f() {
                let s = <|>r###"#random"##
string"###;
            }
            "####,
        )
    }

    #[test]
    fn make_raw_string_nothing_to_unescape_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_hash_target() {
        check_assist_target(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r#"r"random string""#,
        );
    }

    #[test]
    fn add_hash_works() {
        check_assist(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_more_hash_works() {
        check_assist(
            add_hash,
            r##"
            fn f() {
                let s = <|>r#"random"string"#;
            }
            "##,
            r###"
            fn f() {
                let s = <|>r##"random"string"##;
            }
            "###,
        )
    }

    #[test]
    fn add_hash_not_works() {
        check_assist_not_applicable(
            add_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_target() {
        check_assist_target(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn remove_hash_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        )
    }

    #[test]
    fn remove_hash_with_quote_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn remove_more_hash_works() {
        check_assist(
            remove_hash,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn remove_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_no_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        );
    }

    #[test]
    fn make_usual_string_target() {
        check_assist_target(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn make_usual_string_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_with_quote_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_more_hash_works() {
        check_assist(
            make_usual_string,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>"random string";
            }
            "##,
        )
    }

    #[test]
    fn make_usual_string_not_works() {
        check_assist_not_applicable(
            make_usual_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn count_hashes_test() {
        assert_eq!(0, count_hashes("abc"));
        assert_eq!(0, count_hashes("###"));
        assert_eq!(1, count_hashes("\"#abc"));
        assert_eq!(0, count_hashes("#abc"));
        assert_eq!(2, count_hashes("#ab\"##c"));
        assert_eq!(4, count_hashes("#ab\"##\"####c"));
    }
}
<commit_msg>Fix panic on raw string assist<commit_after>use hir::db::HirDatabase;
use ra_syntax::{
    SyntaxKind::{RAW_STRING, STRING},
    TextRange, TextUnit,
};
use rustc_lexer;

use crate::{Assist, AssistCtx, AssistId};

// Assist: make_raw_string
//
// Adds `r#` to a plain string literal.
//
// ```
// fn main() {
//     "Hello,<|> World!";
// }
// ```
// ->
// ```
// fn main() {
//     r#"Hello, World!"#;
// }
// ```
pub(crate) fn make_raw_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let token = ctx.find_token_at_offset(STRING)?;
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    let start_of_inside = usual_string_range.start().to_usize() + 1;
    let end_of_inside = usual_string_range.end().to_usize();
    let inside_str = &text[start_of_inside..end_of_inside];
    let mut unescaped = String::with_capacity(inside_str.len());
    let mut error = Ok(());
    rustc_lexer::unescape::unescape_str(
        inside_str,
        &mut |_, unescaped_char| match unescaped_char {
            Ok(c) => unescaped.push(c),
            Err(_) => error = Err(()),
        },
    );
    if error.is_err() {
        return None;
    }
    ctx.add_action(AssistId("make_raw_string"), "make raw string", |edit| {
        edit.target(token.text_range());
        let max_hash_streak = count_hashes(&unescaped);
        let mut hashes = String::with_capacity(max_hash_streak + 1);
        for _ in 0..hashes.capacity() {
            hashes.push('#');
        }
        edit.replace(token.text_range(), format!("r{}\"{}\"{}", hashes, unescaped, hashes));
    });
    ctx.build()
}

// Assist: make_usual_string
//
// Turns a raw string into a plain string.
//
// ```
// fn main() {
//     r#"Hello,<|> "World!""#;
// }
// ```
// ->
// ```
// fn main() {
//     "Hello, \"World!\"";
// }
// ```
pub(crate) fn make_usual_string(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let token = ctx.find_token_at_offset(RAW_STRING)?;
    let text = token.text().as_str();
    let usual_string_range = find_usual_string_range(text)?;
    ctx.add_action(AssistId("make_usual_string"), "make usual string", |edit| {
        edit.target(token.text_range());
        // parse inside string to escape `"`
        let start_of_inside = usual_string_range.start().to_usize() + 1;
        let end_of_inside = usual_string_range.end().to_usize();
        let inside_str = &text[start_of_inside..end_of_inside];
        let escaped = inside_str.escape_default().to_string();
        edit.replace(token.text_range(), format!("\"{}\"", escaped));
    });
    ctx.build()
}

// Assist: add_hash
//
// Adds a hash to a raw string literal.
//
// ```
// fn main() {
//     r#"Hello,<|> World!"#;
// }
// ```
// ->
// ```
// fn main() {
//     r##"Hello, World!"##;
// }
// ```
pub(crate) fn add_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let token = ctx.find_token_at_offset(RAW_STRING)?;
    ctx.add_action(AssistId("add_hash"), "add hash to raw string", |edit| {
        edit.target(token.text_range());
        edit.insert(token.text_range().start() + TextUnit::of_char('r'), "#");
        edit.insert(token.text_range().end(), "#");
    });
    ctx.build()
}

// Assist: remove_hash
//
// Removes a hash from a raw string literal.
//
// ```
// fn main() {
//     r#"Hello,<|> World!"#;
// }
// ```
// ->
// ```
// fn main() {
//     r"Hello, World!";
// }
// ```
pub(crate) fn remove_hash(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let token = ctx.find_token_at_offset(RAW_STRING)?;
    let text = token.text().as_str();
    if text.starts_with("r\"") {
        // no hash to remove
        return None;
    }
    ctx.add_action(AssistId("remove_hash"), "remove hash from raw string", |edit| {
        edit.target(token.text_range());
        let result = &text[2..text.len() - 1];
        let result = if result.starts_with("\"") {
            // no more hash, escape
            let internal_str = &result[1..result.len() - 1];
            format!("\"{}\"", internal_str.escape_default().to_string())
        } else {
            result.to_owned()
        };
        edit.replace(token.text_range(), format!("r{}", result));
    });
    ctx.build()
}

fn count_hashes(s: &str) -> usize {
    let mut max_hash_streak = 0usize;
    for idx in s.match_indices("\"#").map(|(i, _)| i) {
        let (_, sub) = s.split_at(idx + 1);
        let nb_hash = sub.chars().take_while(|c| *c == '#').count();
        if nb_hash > max_hash_streak {
            max_hash_streak = nb_hash;
        }
    }
    max_hash_streak
}

fn find_usual_string_range(s: &str) -> Option<TextRange> {
    let left_quote = s.find('"')?;
    let right_quote = s.rfind('"')?;
    if left_quote == right_quote {
        // `s` only contains one quote
        None
    } else {
        Some(TextRange::from_to(
            TextUnit::from(left_quote as u32),
            TextUnit::from(right_quote as u32),
        ))
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};

    #[test]
    fn make_raw_string_target() {
        check_assist_target(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r#""random\nstring""#,
        );
    }

    #[test]
    fn make_raw_string_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random\nstring";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random
string"#;
            }
            "##,
        )
    }

    #[test]
    fn make_raw_string_works_inside_macros() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                format!(<|>"x = {}", 92)
            }
            "#,
            r##"
            fn f() {
                format!(<|>r#"x = {}"#, 92)
            }
            "##,
        )
    }

    #[test]
    fn make_raw_string_hashes_inside_works() {
        check_assist(
            make_raw_string,
            r###"
            fn f() {
                let s = <|>"#random##\nstring";
            }
            "###,
            r####"
            fn f() {
                let s = <|>r#"#random##
string"#;
            }
            "####,
        )
    }

    #[test]
    fn make_raw_string_closing_hashes_inside_works() {
        check_assist(
            make_raw_string,
            r###"
            fn f() {
                let s = <|>"#random\"##\nstring";
            }
            "###,
            r####"
            fn f() {
                let s = <|>r###"#random"##
string"###;
            }
            "####,
        )
    }

    #[test]
    fn make_raw_string_nothing_to_unescape_works() {
        check_assist(
            make_raw_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn make_raw_string_not_works_on_partial_string() {
        check_assist_not_applicable(
            make_raw_string,
            r#"
            fn f() {
                let s = "foo<|>
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_not_works_on_partial_string() {
        check_assist_not_applicable(
            make_usual_string,
            r#"
            fn main() {
                let s = r#"bar<|>
            }
            "#,
        )
    }

    #[test]
    fn add_hash_target() {
        check_assist_target(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r#"r"random string""#,
        );
    }

    #[test]
    fn add_hash_works() {
        check_assist(
            add_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn add_more_hash_works() {
        check_assist(
            add_hash,
            r##"
            fn f() {
                let s = <|>r#"random"string"#;
            }
            "##,
            r###"
            fn f() {
                let s = <|>r##"random"string"##;
            }
            "###,
        )
    }

    #[test]
    fn add_hash_not_works() {
        check_assist_not_applicable(
            add_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_target() {
        check_assist_target(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn remove_hash_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        )
    }

    #[test]
    fn remove_hash_with_quote_works() {
        check_assist(
            remove_hash,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>r"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn remove_more_hash_works() {
        check_assist(
            remove_hash,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
        )
    }

    #[test]
    fn remove_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn remove_hash_no_hash_not_works() {
        check_assist_not_applicable(
            remove_hash,
            r#"
            fn f() {
                let s = <|>r"random string";
            }
            "#,
        );
    }

    #[test]
    fn make_usual_string_target() {
        check_assist_target(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r##"r#"random string"#"##,
        );
    }

    #[test]
    fn make_usual_string_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random string"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_with_quote_works() {
        check_assist(
            make_usual_string,
            r##"
            fn f() {
                let s = <|>r#"random"str"ing"#;
            }
            "##,
            r#"
            fn f() {
                let s = <|>"random\"str\"ing";
            }
            "#,
        )
    }

    #[test]
    fn make_usual_string_more_hash_works() {
        check_assist(
            make_usual_string,
            r###"
            fn f() {
                let s = <|>r##"random string"##;
            }
            "###,
            r##"
            fn f() {
                let s = <|>"random string";
            }
            "##,
        )
    }

    #[test]
    fn make_usual_string_not_works() {
        check_assist_not_applicable(
            make_usual_string,
            r#"
            fn f() {
                let s = <|>"random string";
            }
            "#,
        );
    }

    #[test]
    fn count_hashes_test() {
        assert_eq!(0, count_hashes("abc"));
        assert_eq!(0, count_hashes("###"));
        assert_eq!(1, count_hashes("\"#abc"));
        assert_eq!(0, count_hashes("#abc"));
        assert_eq!(2, count_hashes("#ab\"##c"));
        assert_eq!(4, count_hashes("#ab\"##\"####c"));
    }
}
<|endoftext|>


<commit_before>#![unstable(reason = "not public", issue = "none", feature = "fd")]

#[cfg(test)]
mod tests;

use crate::cmp;
use crate::io::{self, Initializer, IoSlice, IoSliceMut, Read};
use crate::mem;
use crate::sys::cvt;
use crate::sys_common::AsInner;

use libc::{c_int, c_void};

#[derive(Debug)]
#[rustc_layout_scalar_valid_range_start(0)]
// libstd/os/raw/mod.rs assures me that every libstd-supported platform has a
// 32-bit c_int. Below is -2, in two's complement, but that only works out
// because c_int is 32 bits.
#[rustc_layout_scalar_valid_range_end(0xFF_FF_FF_FE)]
pub struct FileDesc {
    fd: c_int,
}

// The maximum read limit on most POSIX-like systems is `SSIZE_MAX`,
// with the man page quoting that if the count of bytes to read is
// greater than `SSIZE_MAX` the result is "unspecified".
//
// On macOS, however, apparently the 64-bit libc is either buggy or
// intentionally showing odd behavior by rejecting any read with a size
// larger than or equal to INT_MAX. To handle both of these the read
// size is capped on both platforms.
#[cfg(target_os = "macos")]
const READ_LIMIT: usize = c_int::MAX as usize - 1;
#[cfg(not(target_os = "macos"))]
const READ_LIMIT: usize = libc::ssize_t::MAX as usize;

#[cfg(any(
    target_os = "dragonfly",
    target_os = "freebsd",
    target_os = "ios",
    target_os = "macos",
    target_os = "netbsd",
    target_os = "openbsd",
))]
const fn max_iov() -> usize {
    libc::IOV_MAX as usize
}

#[cfg(any(target_os = "android", target_os = "emscripten", target_os = "linux"))]
const fn max_iov() -> usize {
    libc::UIO_MAXIOV as usize
}

#[cfg(not(any(
    target_os = "android",
    target_os = "dragonfly",
    target_os = "emscripten",
    target_os = "freebsd",
    target_os = "ios",
    target_os = "linux",
    target_os = "macos",
    target_os = "netbsd",
    target_os = "openbsd",
)))]
const fn max_iov() -> usize {
    16 // The minimum value required by POSIX.
}

impl FileDesc {
    pub fn new(fd: c_int) -> FileDesc {
        assert_ne!(fd, -1);
        // SAFETY: we just asserted that the value is in the valid range and isn't `-1` (the only value bigger than `0xFF_FF_FF_FE` unsigned)
        unsafe { FileDesc { fd } }
    }

    pub fn raw(&self) -> c_int {
        self.fd
    }

    /// Extracts the actual file descriptor without closing it.
    pub fn into_raw(self) -> c_int {
        let fd = self.fd;
        mem::forget(self);
        fd
    }

    pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
        let ret = cvt(unsafe {
            libc::read(self.fd, buf.as_mut_ptr() as *mut c_void, cmp::min(buf.len(), READ_LIMIT))
        })?;
        Ok(ret as usize)
    }

    pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
        let ret = cvt(unsafe {
            libc::readv(
                self.fd,
                bufs.as_ptr() as *const libc::iovec,
                cmp::min(bufs.len(), max_iov()) as c_int,
            )
        })?;
        Ok(ret as usize)
    }

    #[inline]
    pub fn is_read_vectored(&self) -> bool {
        true
    }

    pub fn read_to_end(&self, buf: &mut Vec<u8>) -> io::Result<usize> {
        let mut me = self;
        (&mut me).read_to_end(buf)
    }

    pub fn read_at(&self, buf: &mut [u8], offset: u64) -> io::Result<usize> {
        #[cfg(target_os = "android")]
        use super::android::cvt_pread64;

        #[cfg(not(target_os = "android"))]
        unsafe fn cvt_pread64(
            fd: c_int,
            buf: *mut c_void,
            count: usize,
            offset: i64,
        ) -> io::Result<isize> {
            #[cfg(not(target_os = "linux"))]
            use libc::pread as pread64;
            #[cfg(target_os = "linux")]
            use libc::pread64;
            cvt(pread64(fd, buf, count, offset))
        }

        unsafe {
            cvt_pread64(
                self.fd,
                buf.as_mut_ptr() as *mut c_void,
                cmp::min(buf.len(), READ_LIMIT),
                offset as i64,
            )
            .map(|n| n as usize)
        }
    }

    pub fn write(&self, buf: &[u8]) -> io::Result<usize> {
        let ret = cvt(unsafe {
            libc::write(self.fd, buf.as_ptr() as *const c_void, cmp::min(buf.len(), READ_LIMIT))
        })?;
        Ok(ret as usize)
    }

    pub fn write_vectored(&self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
        let ret = cvt(unsafe {
            libc::writev(
                self.fd,
                bufs.as_ptr() as *const libc::iovec,
                cmp::min(bufs.len(), max_iov()) as c_int,
            )
        })?;
        Ok(ret as usize)
    }

    #[inline]
    pub fn is_write_vectored(&self) -> bool {
        true
    }

    pub fn write_at(&self, buf: &[u8], offset: u64) -> io::Result<usize> {
        #[cfg(target_os = "android")]
        use super::android::cvt_pwrite64;

        #[cfg(not(target_os = "android"))]
        unsafe fn cvt_pwrite64(
            fd: c_int,
            buf: *const c_void,
            count: usize,
            offset: i64,
        ) -> io::Result<isize> {
            #[cfg(not(target_os = "linux"))]
            use libc::pwrite as pwrite64;
            #[cfg(target_os = "linux")]
            use libc::pwrite64;
            cvt(pwrite64(fd, buf, count, offset))
        }

        unsafe {
            cvt_pwrite64(
                self.fd,
                buf.as_ptr() as *const c_void,
                cmp::min(buf.len(), READ_LIMIT),
                offset as i64,
            )
            .map(|n| n as usize)
        }
    }

    #[cfg(target_os = "linux")]
    pub fn get_cloexec(&self) -> io::Result<bool> {
        unsafe { Ok((cvt(libc::fcntl(self.fd, libc::F_GETFD))? & libc::FD_CLOEXEC) != 0) }
    }

    #[cfg(not(any(
        target_env = "newlib",
        target_os = "solaris",
        target_os = "illumos",
        target_os = "emscripten",
        target_os = "fuchsia",
        target_os = "l4re",
        target_os = "linux",
        target_os = "haiku",
        target_os = "redox",
        target_os = "vxworks"
    )))]
    pub fn set_cloexec(&self) -> io::Result<()> {
        unsafe {
            cvt(libc::ioctl(self.fd, libc::FIOCLEX))?;
            Ok(())
        }
    }
    #[cfg(any(
        target_env = "newlib",
        target_os = "solaris",
        target_os = "illumos",
        target_os = "emscripten",
        target_os = "fuchsia",
        target_os = "l4re",
        target_os = "linux",
        target_os = "haiku",
        target_os = "redox",
        target_os = "vxworks"
    ))]
    pub fn set_cloexec(&self) -> io::Result<()> {
        unsafe {
            let previous = cvt(libc::fcntl(self.fd, libc::F_GETFD))?;
            let new = previous | libc::FD_CLOEXEC;
            if new != previous {
                cvt(libc::fcntl(self.fd, libc::F_SETFD, new))?;
            }
            Ok(())
        }
    }

    #[cfg(target_os = "linux")]
    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
        unsafe {
            let v = nonblocking as c_int;
            cvt(libc::ioctl(self.fd, libc::FIONBIO, &v))?;
            Ok(())
        }
    }

    #[cfg(not(target_os = "linux"))]
    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
        unsafe {
            let previous = cvt(libc::fcntl(self.fd, libc::F_GETFL))?;
            let new = if nonblocking {
                previous | libc::O_NONBLOCK
            } else {
                previous & !libc::O_NONBLOCK
            };
            if new != previous {
                cvt(libc::fcntl(self.fd, libc::F_SETFL, new))?;
            }
            Ok(())
        }
    }

    pub fn duplicate(&self) -> io::Result<FileDesc> {
        // We want to atomically duplicate this file descriptor and set the
        // CLOEXEC flag, and currently that's done via F_DUPFD_CLOEXEC. This
        // is a POSIX flag that was added to Linux in 2.6.24.
        let fd = cvt(unsafe { libc::fcntl(self.raw(), libc::F_DUPFD_CLOEXEC, 0) })?;
        Ok(FileDesc::new(fd))
    }
}

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

    #[inline]
    unsafe fn initializer(&self) -> Initializer {
        Initializer::nop()
    }
}

impl AsInner<c_int> for FileDesc {
    fn as_inner(&self) -> &c_int {
        &self.fd
    }
}

impl Drop for FileDesc {
    fn drop(&mut self) {
        // Note that errors are ignored when closing a file descriptor. The
        // reason for this is that if an error occurs we don't actually know if
        // the file descriptor was closed or not, and if we retried (for
        // something like EINTR), we might close another valid file descriptor
        // opened after we closed ours.
        let _ = unsafe { libc::close(self.fd) };
    }
}
<commit_msg>Check that c_int is i32 in FileDesc::new.<commit_after>#![unstable(reason = "not public", issue = "none", feature = "fd")]

#[cfg(test)]
mod tests;

use crate::cmp;
use crate::io::{self, Initializer, IoSlice, IoSliceMut, Read};
use crate::mem;
use crate::sys::cvt;
use crate::sys_common::AsInner;

use libc::{c_int, c_void};

#[derive(Debug)]
#[rustc_layout_scalar_valid_range_start(0)]
// libstd/os/raw/mod.rs assures me that every libstd-supported platform has a
// 32-bit c_int. Below is -2, in two's complement, but that only works out
// because c_int is 32 bits.
#[rustc_layout_scalar_valid_range_end(0xFF_FF_FF_FE)]
pub struct FileDesc {
    fd: c_int,
}

// The maximum read limit on most POSIX-like systems is `SSIZE_MAX`,
// with the man page quoting that if the count of bytes to read is
// greater than `SSIZE_MAX` the result is "unspecified".
//
// On macOS, however, apparently the 64-bit libc is either buggy or
// intentionally showing odd behavior by rejecting any read with a size
// larger than or equal to INT_MAX. To handle both of these the read
// size is capped on both platforms.
#[cfg(target_os = "macos")]
const READ_LIMIT: usize = c_int::MAX as usize - 1;
#[cfg(not(target_os = "macos"))]
const READ_LIMIT: usize = libc::ssize_t::MAX as usize;

#[cfg(any(
    target_os = "dragonfly",
    target_os = "freebsd",
    target_os = "ios",
    target_os = "macos",
    target_os = "netbsd",
    target_os = "openbsd",
))]
const fn max_iov() -> usize {
    libc::IOV_MAX as usize
}

#[cfg(any(target_os = "android", target_os = "emscripten", target_os = "linux"))]
const fn max_iov() -> usize {
    libc::UIO_MAXIOV as usize
}

#[cfg(not(any(
    target_os = "android",
    target_os = "dragonfly",
    target_os = "emscripten",
    target_os = "freebsd",
    target_os = "ios",
    target_os = "linux",
    target_os = "macos",
    target_os = "netbsd",
    target_os = "openbsd",
)))]
const fn max_iov() -> usize {
    16 // The minimum value required by POSIX.
}

impl FileDesc {
    pub fn new(fd: c_int) -> FileDesc {
        assert_ne!(fd, -1i32);
        // SAFETY: we just asserted that the value is in the valid range and isn't `-1` (the only value bigger than `0xFF_FF_FF_FE` unsigned)
        unsafe { FileDesc { fd } }
    }

    pub fn raw(&self) -> c_int {
        self.fd
    }

    /// Extracts the actual file descriptor without closing it.
    pub fn into_raw(self) -> c_int {
        let fd = self.fd;
        mem::forget(self);
        fd
    }

    pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
        let ret = cvt(unsafe {
            libc::read(self.fd, buf.as_mut_ptr() as *mut c_void, cmp::min(buf.len(), READ_LIMIT))
        })?;
        Ok(ret as usize)
    }

    pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
        let ret = cvt(unsafe {
            libc::readv(
                self.fd,
                bufs.as_ptr() as *const libc::iovec,
                cmp::min(bufs.len(), max_iov()) as c_int,
            )
        })?;
        Ok(ret as usize)
    }

    #[inline]
    pub fn is_read_vectored(&self) -> bool {
        true
    }

    pub fn read_to_end(&self, buf: &mut Vec<u8>) -> io::Result<usize> {
        let mut me = self;
        (&mut me).read_to_end(buf)
    }

    pub fn read_at(&self, buf: &mut [u8], offset: u64) -> io::Result<usize> {
        #[cfg(target_os = "android")]
        use super::android::cvt_pread64;

        #[cfg(not(target_os = "android"))]
        unsafe fn cvt_pread64(
            fd: c_int,
            buf: *mut c_void,
            count: usize,
            offset: i64,
        ) -> io::Result<isize> {
            #[cfg(not(target_os = "linux"))]
            use libc::pread as pread64;
            #[cfg(target_os = "linux")]
            use libc::pread64;
            cvt(pread64(fd, buf, count, offset))
        }

        unsafe {
            cvt_pread64(
                self.fd,
                buf.as_mut_ptr() as *mut c_void,
                cmp::min(buf.len(), READ_LIMIT),
                offset as i64,
            )
            .map(|n| n as usize)
        }
    }

    pub fn write(&self, buf: &[u8]) -> io::Result<usize> {
        let ret = cvt(unsafe {
            libc::write(self.fd, buf.as_ptr() as *const c_void, cmp::min(buf.len(), READ_LIMIT))
        })?;
        Ok(ret as usize)
    }

    pub fn write_vectored(&self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
        let ret = cvt(unsafe {
            libc::writev(
                self.fd,
                bufs.as_ptr() as *const libc::iovec,
                cmp::min(bufs.len(), max_iov()) as c_int,
            )
        })?;
        Ok(ret as usize)
    }

    #[inline]
    pub fn is_write_vectored(&self) -> bool {
        true
    }

    pub fn write_at(&self, buf: &[u8], offset: u64) -> io::Result<usize> {
        #[cfg(target_os = "android")]
        use super::android::cvt_pwrite64;

        #[cfg(not(target_os = "android"))]
        unsafe fn cvt_pwrite64(
            fd: c_int,
            buf: *const c_void,
            count: usize,
            offset: i64,
        ) -> io::Result<isize> {
            #[cfg(not(target_os = "linux"))]
            use libc::pwrite as pwrite64;
            #[cfg(target_os = "linux")]
            use libc::pwrite64;
            cvt(pwrite64(fd, buf, count, offset))
        }

        unsafe {
            cvt_pwrite64(
                self.fd,
                buf.as_ptr() as *const c_void,
                cmp::min(buf.len(), READ_LIMIT),
                offset as i64,
            )
            .map(|n| n as usize)
        }
    }

    #[cfg(target_os = "linux")]
    pub fn get_cloexec(&self) -> io::Result<bool> {
        unsafe { Ok((cvt(libc::fcntl(self.fd, libc::F_GETFD))? & libc::FD_CLOEXEC) != 0) }
    }

    #[cfg(not(any(
        target_env = "newlib",
        target_os = "solaris",
        target_os = "illumos",
        target_os = "emscripten",
        target_os = "fuchsia",
        target_os = "l4re",
        target_os = "linux",
        target_os = "haiku",
        target_os = "redox",
        target_os = "vxworks"
    )))]
    pub fn set_cloexec(&self) -> io::Result<()> {
        unsafe {
            cvt(libc::ioctl(self.fd, libc::FIOCLEX))?;
            Ok(())
        }
    }
    #[cfg(any(
        target_env = "newlib",
        target_os = "solaris",
        target_os = "illumos",
        target_os = "emscripten",
        target_os = "fuchsia",
        target_os = "l4re",
        target_os = "linux",
        target_os = "haiku",
        target_os = "redox",
        target_os = "vxworks"
    ))]
    pub fn set_cloexec(&self) -> io::Result<()> {
        unsafe {
            let previous = cvt(libc::fcntl(self.fd, libc::F_GETFD))?;
            let new = previous | libc::FD_CLOEXEC;
            if new != previous {
                cvt(libc::fcntl(self.fd, libc::F_SETFD, new))?;
            }
            Ok(())
        }
    }

    #[cfg(target_os = "linux")]
    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
        unsafe {
            let v = nonblocking as c_int;
            cvt(libc::ioctl(self.fd, libc::FIONBIO, &v))?;
            Ok(())
        }
    }

    #[cfg(not(target_os = "linux"))]
    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
        unsafe {
            let previous = cvt(libc::fcntl(self.fd, libc::F_GETFL))?;
            let new = if nonblocking {
                previous | libc::O_NONBLOCK
            } else {
                previous & !libc::O_NONBLOCK
            };
            if new != previous {
                cvt(libc::fcntl(self.fd, libc::F_SETFL, new))?;
            }
            Ok(())
        }
    }

    pub fn duplicate(&self) -> io::Result<FileDesc> {
        // We want to atomically duplicate this file descriptor and set the
        // CLOEXEC flag, and currently that's done via F_DUPFD_CLOEXEC. This
        // is a POSIX flag that was added to Linux in 2.6.24.
        let fd = cvt(unsafe { libc::fcntl(self.raw(), libc::F_DUPFD_CLOEXEC, 0) })?;
        Ok(FileDesc::new(fd))
    }
}

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

    #[inline]
    unsafe fn initializer(&self) -> Initializer {
        Initializer::nop()
    }
}

impl AsInner<c_int> for FileDesc {
    fn as_inner(&self) -> &c_int {
        &self.fd
    }
}

impl Drop for FileDesc {
    fn drop(&mut self) {
        // Note that errors are ignored when closing a file descriptor. The
        // reason for this is that if an error occurs we don't actually know if
        // the file descriptor was closed or not, and if we retried (for
        // something like EINTR), we might close another valid file descriptor
        // opened after we closed ours.
        let _ = unsafe { libc::close(self.fd) };
    }
}
<|endoftext|>


<commit_before>//! Estimate upper and lower bounds for the *n*-th prime, and π(*n*),
//! the number of primes less than or equal to *n*.
//!
//! This is designed to be used via the `primal` crate.

#[allow(dead_code)]
mod tables;

/// Returns estimated bounds for π(*n*), the number of primes less
/// than or equal to `n`.
///
/// That is, if (*a*, *b*) = `estimate_prime_pi(n)`, *a* ≤ π(*n*) ≤
/// *b*. The bounds used are proved in [1] and [2, Théorème 1.10],
/// and are summarised in [2, pp. 14–15].
///
/// [1]: Barkley Rosser. "Explicit Bounds for Some Functions of Prime
/// Numbers". American Journal of Mathematics 63 (1):
/// 211–232. 1941. doi:[10.2307/2371291](http://dx.doi.org/10.2307/2371291).
///
///  [2]: Dusart, Pierre. ["Autour de la fonction qui compte le nombre
/// de nombres premiers."][pdf] PhD diss., Université de Limoges,
/// 1998.
///
/// [pdf]: http://www.unilim.fr/laco/theses/1998/T1998_01.html
pub fn prime_pi(n: u64) -> (u64, u64) {
    if n < tables::SMALL_PRIME_PI.len() as u64 {
        let x = tables::SMALL_PRIME_PI[n as usize] as u64;
        (x, x)
    } else {
        let n_ = n as f64;
        let lg = n_.ln();
        let inv_lg = 1.0 / lg;
        let n_lg = n_ * inv_lg;

        // numbers refer to parts of theorem 1.10 of [2].
        let lo = if n >= 32299 {
            // 6.
            n_lg * (1.0 + inv_lg * (1.0 + 1.8 * inv_lg))
        } else if n >= 5393 {
            // 5.
            n_ / (lg - 1.0)
        } else if n >= 599 {
            // 1.
            n_lg * (1.0 + inv_lg)
        } else {
            // [1]
            n_ / (lg + 2.0)
        };

        let hi = if n >= 13_220_000_000 {
            // 3.
            n_lg * (1.0 + 1.0992 * inv_lg)
        } else if n >= 355991 {
            // 7.
            n_lg * (1.0 + inv_lg * (1.0 + 2.51 * inv_lg))
        } else if n >= 60184 {
            // 4.
            n_ / (lg - 1.1)
        } else {
            // 2.
            n_lg * (1.0 + 1.2762 * inv_lg)
        };

        (lo as u64, hi as u64)
    }
}

/// Gives estimated bounds for *pn*, the `n`th prime number,
/// 1-indexed (i.e. *p1* = 2, *p2* = 3).
///
/// That is, if (a,b) = `estimate_nth_prime(n)`, *a* ≤
/// *pn* ≤ *b*. The bounds used are proved in [1] and [2,
/// Théorèmes 1.6–1.8], and are summarised in [2, pp. 14–15].
///
/// [1]: Massias, Jean-Pierre; Robin, Guy. ["Bornes effectives pour
/// certaines fonctions concernant les nombres
/// premiers."](http://eudml.org/doc/247826) Journal de théorie des
/// nombres de Bordeaux 8.1 (1996): 215-242.
///
/// [2]: Dusart, Pierre. ["Autour de la fonction qui compte le nombre
/// de nombres premiers."][pdf] PhD diss., Université de Limoges, 1998.
///
/// [pdf]: http://www.unilim.fr/laco/theses/1998/T1998_01.html
pub fn nth_prime(n: u64) -> (u64, u64) {
    if n == 0 {
        (0, 0)
    } else if n <= tables::SMALL_PRIMES.len() as u64 {
        // table is 0-indexed, n is 1-indexed, need to adjust.
        let x = tables::SMALL_PRIMES[n as usize - 1] as u64;
        (x, x)
    } else {
        let n_ = n as f64;
        let lg = n_.ln();
        let lglg = lg.ln();

        let lo = lg + lglg - 1.0 + if n > 3 {//13196 {
            // [2] Theorem 1.6
            (lglg - 2.1) / lg
        } else {
            // [1] Theorem A (ii)
            0.0
        };

        let hi = lg + lglg + if n >= 39017 {
            // [2] Theorem 1.8
            -0.9484
        } else if n >= 27076 {
            // [2] Theorem 1.7
            -1.0 + (lglg - 1.8) / lg
        } else if n >= 15985 {
            // [1] Theorem A (v)
            -0.9427
        } else if n >= 13 {
            // [1] Theorem A (v)
            -1.0 + 1.8 * lglg / lg
        } else {
            // [1] Theorem A (iv)
            0.0
        };
        ((n_ * lo) as u64, (n_ * hi) as u64)
    }
}

#[cfg(test)]
mod tests {
    extern crate primal;
    use self::primal::Sieve;

    #[test]
    fn prime_pi() {
        fn check(n: u64, pi: u64) {
            let (lo, hi) = super::prime_pi(n);
            assert!(lo <= pi && pi <= hi,
                    "found failing estimate at {}, should satisfy: {} <= {} <= {}",
                    n, lo, pi, hi)
        }
        let primes = Sieve::new(1_000_000);

        let mut last = 0;
        for (i, p) in primes.primes_from(0).enumerate() {
            for j in last..p {
                check(j as u64, i as u64);
            }
            last = p;
        }

        let sporadic = [
            (1, 4),
            (2, 25),
            (3, 168),
            (4, 1229),
            (5, 9592),
            (6, 78498),
            (7, 664579),
            (8, 5761455),
            (9, 50847534),
            (10, 455052511),
            (11, 4118054813),
            (12, 37607912018),
            (13, 346065536839),
            (14, 3204941750802),
            (15, 29844570422669),
            (16, 279238341033925),
            (17, 2623557157654233),
            ];
        for &(exponent, real) in sporadic.iter() {
            let n = 10u64.pow(exponent);
            check(n, real);
        }
    }

    #[test]
    fn nth_prime() {
        fn check(n: u64, p: u64) {
            let (lo, hi) = super::nth_prime(n);
            assert!(lo <= p && p <= hi,
                    "found failing estimate at {}, should satisfy: {} <= {} <= {}",
                    n, lo, p, hi);
        }
        let sieve = Sieve::new(1_000_000);

        for (i, p) in sieve.primes_from(0).enumerate() {
            let n = i as u64 + 1;
            check(n, p as u64);
        }

        let sporadic = [
            (0, 2),
            (1, 29),
            (2, 541),
            (3, 7919),
            (4, 104729),
            (5, 1299709),
            (6, 15485863),
            (7, 179424673),
            (8, 2038074743),
            (9, 22801763489),
            (10, 252097800623),
            (11, 2760727302517),
            (12, 29996224275833),
            (13, 323780508946331),
            (14, 3475385758524527),
            (15, 37124508045065437),
            ];

        for &(exponent, nth_prime) in sporadic.iter() {
            let n = 10u64.pow(exponent);
            check(n, nth_prime);
        }
    }
}
<commit_msg>estimate: add doc examples.<commit_after>//! Estimate upper and lower bounds for the *n*-th prime, and π(*n*),
//! the number of primes less than or equal to *n*.
//!
//! This is designed to be used via the `primal` crate.

#[allow(dead_code)]
mod tables;

/// Returns estimated bounds for π(*n*), the number of primes less
/// than or equal to `n`.
///
/// That is, if (*a*, *b*) = `estimate_prime_pi(n)`, *a* ≤ π(*n*) ≤
/// *b*. The bounds used are proved in [1] and [2, Théorème 1.10],
/// and are summarised in [2, pp. 14–15].
///
/// [1]: Barkley Rosser. "Explicit Bounds for Some Functions of Prime
/// Numbers". American Journal of Mathematics 63 (1):
/// 211–232. 1941. doi:[10.2307/2371291](http://dx.doi.org/10.2307/2371291).
///
///  [2]: Dusart, Pierre. ["Autour de la fonction qui compte le nombre
/// de nombres premiers."][pdf] PhD diss., Université de Limoges,
/// 1998.
///
/// [pdf]: http://www.unilim.fr/laco/theses/1998/T1998_01.html
///
/// # Examples
///
/// ```rust
/// # extern crate primal;
/// // the number of primes below 1e9
/// let count_billion = 50_847_534;
///
/// let (low, high) = primal::estimate_prime_pi(1_000_000_000);
/// // check the bounds are satisfied
/// assert!(low <= count_billion && count_billion <= high);
/// ```
pub fn prime_pi(n: u64) -> (u64, u64) {
    if n < tables::SMALL_PRIME_PI.len() as u64 {
        let x = tables::SMALL_PRIME_PI[n as usize] as u64;
        (x, x)
    } else {
        let n_ = n as f64;
        let lg = n_.ln();
        let inv_lg = 1.0 / lg;
        let n_lg = n_ * inv_lg;

        // numbers refer to parts of theorem 1.10 of [2].
        let lo = if n >= 32299 {
            // 6.
            n_lg * (1.0 + inv_lg * (1.0 + 1.8 * inv_lg))
        } else if n >= 5393 {
            // 5.
            n_ / (lg - 1.0)
        } else if n >= 599 {
            // 1.
            n_lg * (1.0 + inv_lg)
        } else {
            // [1]
            n_ / (lg + 2.0)
        };

        let hi = if n >= 13_220_000_000 {
            // 3.
            n_lg * (1.0 + 1.0992 * inv_lg)
        } else if n >= 355991 {
            // 7.
            n_lg * (1.0 + inv_lg * (1.0 + 2.51 * inv_lg))
        } else if n >= 60184 {
            // 4.
            n_ / (lg - 1.1)
        } else {
            // 2.
            n_lg * (1.0 + 1.2762 * inv_lg)
        };

        (lo as u64, hi as u64)
    }
}

/// Gives estimated bounds for *pn*, the `n`th prime number,
/// 1-indexed (i.e. *p1* = 2, *p2* = 3).
///
/// That is, if (a,b) = `estimate_nth_prime(n)`, *a* ≤
/// *pn* ≤ *b*. The bounds used are proved in [1] and [2,
/// Théorèmes 1.6–1.8], and are summarised in [2, pp. 14–15].
///
/// [1]: Massias, Jean-Pierre; Robin, Guy. ["Bornes effectives pour
/// certaines fonctions concernant les nombres
/// premiers."](http://eudml.org/doc/247826) Journal de théorie des
/// nombres de Bordeaux 8.1 (1996): 215-242.
///
/// [2]: Dusart, Pierre. ["Autour de la fonction qui compte le nombre
/// de nombres premiers."][pdf] PhD diss., Université de Limoges, 1998.
///
/// [pdf]: http://www.unilim.fr/laco/theses/1998/T1998_01.html
///
/// # Examples
///
/// ```rust
/// # extern crate primal;
/// // the 1e9-th prime
/// let billionth = 22_801_763_489;
///
/// let (low, high) = primal::estimate_nth_prime(1_000_000_000);
/// // check the bounds are satisfied
/// assert!(low <= billionth && billionth <= high);
/// ```
pub fn nth_prime(n: u64) -> (u64, u64) {
    if n == 0 {
        (0, 0)
    } else if n <= tables::SMALL_PRIMES.len() as u64 {
        // table is 0-indexed, n is 1-indexed, need to adjust.
        let x = tables::SMALL_PRIMES[n as usize - 1] as u64;
        (x, x)
    } else {
        let n_ = n as f64;
        let lg = n_.ln();
        let lglg = lg.ln();

        let lo = lg + lglg - 1.0 + if n > 3 {//13196 {
            // [2] Theorem 1.6
            (lglg - 2.1) / lg
        } else {
            // [1] Theorem A (ii)
            0.0
        };

        let hi = lg + lglg + if n >= 39017 {
            // [2] Theorem 1.8
            -0.9484
        } else if n >= 27076 {
            // [2] Theorem 1.7
            -1.0 + (lglg - 1.8) / lg
        } else if n >= 15985 {
            // [1] Theorem A (v)
            -0.9427
        } else if n >= 13 {
            // [1] Theorem A (v)
            -1.0 + 1.8 * lglg / lg
        } else {
            // [1] Theorem A (iv)
            0.0
        };
        ((n_ * lo) as u64, (n_ * hi) as u64)
    }
}

#[cfg(test)]
mod tests {
    extern crate primal;
    use self::primal::Sieve;

    #[test]
    fn prime_pi() {
        fn check(n: u64, pi: u64) {
            let (lo, hi) = super::prime_pi(n);
            assert!(lo <= pi && pi <= hi,
                    "found failing estimate at {}, should satisfy: {} <= {} <= {}",
                    n, lo, pi, hi)
        }
        let primes = Sieve::new(1_000_000);

        let mut last = 0;
        for (i, p) in primes.primes_from(0).enumerate() {
            for j in last..p {
                check(j as u64, i as u64);
            }
            last = p;
        }

        let sporadic = [
            (1, 4),
            (2, 25),
            (3, 168),
            (4, 1229),
            (5, 9592),
            (6, 78498),
            (7, 664579),
            (8, 5761455),
            (9, 50847534),
            (10, 455052511),
            (11, 4118054813),
            (12, 37607912018),
            (13, 346065536839),
            (14, 3204941750802),
            (15, 29844570422669),
            (16, 279238341033925),
            (17, 2623557157654233),
            ];
        for &(exponent, real) in sporadic.iter() {
            let n = 10u64.pow(exponent);
            check(n, real);
        }
    }

    #[test]
    fn nth_prime() {
        fn check(n: u64, p: u64) {
            let (lo, hi) = super::nth_prime(n);
            assert!(lo <= p && p <= hi,
                    "found failing estimate at {}, should satisfy: {} <= {} <= {}",
                    n, lo, p, hi);
        }
        let sieve = Sieve::new(1_000_000);

        for (i, p) in sieve.primes_from(0).enumerate() {
            let n = i as u64 + 1;
            check(n, p as u64);
        }

        let sporadic = [
            (0, 2),
            (1, 29),
            (2, 541),
            (3, 7919),
            (4, 104729),
            (5, 1299709),
            (6, 15485863),
            (7, 179424673),
            (8, 2038074743),
            (9, 22801763489),
            (10, 252097800623),
            (11, 2760727302517),
            (12, 29996224275833),
            (13, 323780508946331),
            (14, 3475385758524527),
            (15, 37124508045065437),
            ];

        for &(exponent, nth_prime) in sporadic.iter() {
            let n = 10u64.pow(exponent);
            check(n, nth_prime);
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::io::BufReader;
use std::io::Result;
use std::io::{Error, ErrorKind};
use std::net::{TcpListener, TcpStream};
use std::sync::mpsc::sync_channel;
use std::sync::mpsc::{Receiver, SyncSender};
use std::sync::Arc;
use std::sync::Mutex;
use std::thread;

use birch::client::ClientConnection;
use birch::proto::RawMessage;
use birch::socket::IrcReader;
use birch::socket::{IrcSocket, IrcSocketConfig, IrcWriter};

#[derive(Clone)]
struct IrcChannel(SyncSender<RawMessage>);

impl IrcChannel {
    fn new() -> (Self, Receiver<RawMessage>) {
        let (sender, receiver) = sync_channel(100);
        (IrcChannel(sender), receiver)
    }
}

impl IrcWriter for IrcChannel {
    fn write_message(&mut self, msg: &RawMessage) -> Result<()> {
        if let Err(err) = self.0.send(msg.clone()) {
            println!("Failed to write to client: {:?}", err);
            return Err(Error::new(ErrorKind::Other, "other end disconnected"));
        }

        Ok(())
    }
}

struct ClientFanoutManager {
    id_counter: usize,
    clients: HashMap<usize, SyncSender<RawMessage>>,
}

impl ClientFanoutManager {
    fn new() -> Self {
        Self {
            id_counter: 0,
            clients: HashMap::new(),
        }
    }

    fn next_id(&mut self) -> usize {
        let id = self.id_counter;
        self.id_counter += 1;
        id
    }

    fn add_client(&mut self, sender: SyncSender<RawMessage>) -> usize {
        let id = self.next_id();
        self.clients.insert(id, sender);
        id
    }

    fn remove_client(&mut self, id: usize) -> Option<()> {
        self.clients.remove(&id).map(|_| ())
    }

    fn write_message(&mut self, msg: &RawMessage) {
        let mut dead_clients = vec![];
        for (id, client) in self.clients.iter() {
            if let Err(err) = client.send(msg.clone()) {
                println!("Failed to write to client: {:?}", err);
                dead_clients.push(*id);
            }
        }
        // Prune anything we failed to write to
        for id in dead_clients.into_iter() {
            self.remove_client(id);
        }
    }
}

struct ClientManager {
    fanout_manager: Arc<Mutex<ClientFanoutManager>>,
    fanout: IrcChannel,
}

impl ClientManager {
    fn new() -> Self {
        let (sender, receiver) = IrcChannel::new();
        let manager = Arc::new(Mutex::new(ClientFanoutManager::new()));

        let fanout = Arc::clone(&manager);
        thread::spawn(move || {
            for msg in receiver {
                fanout.lock().expect("mutex poisoned").write_message(&msg);
            }
        });

        Self {
            fanout_manager: manager,
            fanout: sender,
        }
    }

    fn fanout_sender(&self) -> IrcChannel {
        self.fanout.clone()
    }

    fn accept_client_connection(&mut self, stream: TcpStream) {
        let (sender, receiver) = sync_channel(100);

        // TODO: write up disconnect logic
        let _client_id = self
            .fanout_manager
            .lock()
            .unwrap()
            .add_client(sender.clone());

        let mut writer = stream.try_clone().expect("clone stream");
        // Writer thread
        thread::spawn(move || {
            for msg in receiver {
                if let Err(err) = writer.write_message(&msg) {
                    println!("failed to write to client: {:?}", err);
                    break;
                }
            }
        });

        // Reader thread
        thread::spawn(move || {
            let mut writer = IrcChannel(sender);
            let mut client = ClientConnection::new(&mut writer);
            let mut reader = BufReader::new(stream);

            loop {
                let result = reader
                    .read_message()
                    .and_then(|msg| client.handle_message(&msg));

                if result.is_err() {
                    break;
                }
            }
        });
    }
}

fn main() -> Result<()> {
    let mut client_manager = ClientManager::new();

    let config = IrcSocketConfig {
        nick: "ep",
        addr: "irc.freenode.net:6667",
        max_retries: Some(3),
    };

    let mut fanout_sender = client_manager.fanout_sender();
    let mut sock = IrcSocket::new(config);
    thread::spawn(move || {
        // TODO: surface any errors here + shutdown hook
        sock.start_loop(&mut fanout_sender)
    });

    // TODO: hook this up to config
    let listener = TcpListener::bind("127.0.0.1:9123").unwrap();
    for stream in listener.incoming() {
        if let Ok(stream) = stream {
            println!("Client connected");
            client_manager.accept_client_connection(stream);
        }
    }

    Ok(())
}
<commit_msg>Switch to crossbeam channels for select!{}<commit_after>#[macro_use]
extern crate crossbeam;

use std::collections::HashMap;
use std::io::BufReader;
use std::io::Result;
use std::io::{Error, ErrorKind};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::sync::Arc;
use std::sync::Mutex;
use std::thread;
use std::time::Duration;

use crossbeam::crossbeam_channel::{tick, unbounded, Receiver, Sender};

use birch::client::ClientConnection;
use birch::proto::RawMessage;
use birch::socket::IrcReader;
use birch::socket::{IrcSocket, IrcSocketConfig, IrcWriter};

#[derive(Clone)]
struct IrcChannel(Sender<RawMessage>);

impl IrcChannel {
    fn new() -> (Self, Receiver<RawMessage>) {
        let (sender, receiver) = unbounded();
        (IrcChannel(sender), receiver)
    }
}

impl IrcWriter for IrcChannel {
    fn write_message(&mut self, msg: &RawMessage) -> Result<()> {
        if let Err(err) = self.0.send(msg.clone()) {
            println!("Failed to write to client: {:?}", err);
            return Err(Error::new(ErrorKind::Other, "other end disconnected"));
        }
        Ok(())
    }
}

struct ClientFanoutManager {
    id_counter: usize,
    clients: HashMap<usize, Sender<RawMessage>>,
}

impl ClientFanoutManager {
    fn new() -> Self {
        Self {
            id_counter: 0,
            clients: HashMap::new(),
        }
    }

    fn next_id(&mut self) -> usize {
        let id = self.id_counter;
        self.id_counter += 1;
        id
    }

    fn add_client(&mut self, sender: Sender<RawMessage>) -> usize {
        let id = self.next_id();
        self.clients.insert(id, sender);
        id
    }

    fn remove_client(&mut self, id: usize) -> Option<()> {
        self.clients.remove(&id).map(|_| ())
    }

    fn write_message(&mut self, msg: &RawMessage) {
        let mut dead_clients = vec![];
        for (id, client) in self.clients.iter() {
            if let Err(err) = client.send(msg.clone()) {
                println!("Failed to write to client: {:?}", err);
                dead_clients.push(*id);
            }
        }
        // Prune anything we failed to write to
        for id in dead_clients.into_iter() {
            self.remove_client(id);
        }
    }
}

struct ClientManager {
    fanout_manager: Arc<Mutex<ClientFanoutManager>>,
    fanout: IrcChannel,
}

impl ClientManager {
    fn new() -> Self {
        let (sender, receiver) = IrcChannel::new();
        let manager = Arc::new(Mutex::new(ClientFanoutManager::new()));

        let fanout = Arc::clone(&manager);
        thread::spawn(move || {
            for msg in receiver {
                fanout.lock().expect("mutex poisoned").write_message(&msg);
            }
        });

        Self {
            fanout_manager: manager,
            fanout: sender,
        }
    }

    fn fanout_sender(&self) -> IrcChannel {
        self.fanout.clone()
    }

    fn accept_client_connection(&mut self, stream: TcpStream) {
        let to_client = unbounded();
        let from_client = unbounded();

        // TODO: write up disconnect logic
        let _client_id = self
            .fanout_manager
            .lock()
            .unwrap()
            .add_client(to_client.0.clone());

        let reader = stream.try_clone().expect("clone stream");
        let mut writer = stream.try_clone().expect("clone stream");

        // Read thread - handle input from client
        let client_sender = from_client.0.clone();
        thread::spawn(move || {
            reader
                .set_read_timeout(Some(Duration::from_secs(120)))
                .expect("set_read_timeout call failed");

            let mut reader = BufReader::new(reader);

            loop {
                match reader.read_message() {
                    Ok(msg) => {
                        if let Err(err) = client_sender.send(msg) {
                            println!("Failed to send: {:?}", err);
                            break;
                        }
                    }
                    Err(err) => {
                        println!("Failed to read client message: {:?}", err);
                        break;
                    }
                }
            }
        });

        let mut channel = IrcChannel(to_client.0.clone());
        let mut client = ClientConnection::new(&mut channel);
        let ping_ticker = tick(Duration::from_secs(120));

        loop {
            select! {
                recv(to_client.1) -> msg => match msg {
                    Ok(msg) => if let Err(err) = writer.write_message(&msg) {
                        println!("failed to write to client: {:?}", err);
                        break;
                    },
                    Err(err) => {
                        println!("receive failed: {:?}", err);
                        break;
                    }
                },

                recv(from_client.1) -> msg => match msg {
                    Ok(msg) => if let Err(err) = client.handle_message(&msg) {
                        println!("failed to handle client message {:?}", err);
                        break;
                    },
                    Err(err) => {
                        println!("receive failed: {:?}", err);
                        break;
                    }
                },

                recv(ping_ticker) -> _ => if let Err(err) = client.ping() {
                    println!("Ping failed, disconnecting client: {:?}", err);
                    break;
                }
            }
        }

        stream
            .shutdown(Shutdown::Both)
            .expect("shutdown call failed");
    }
}

fn main() -> Result<()> {
    let mut client_manager = ClientManager::new();

    let config = IrcSocketConfig {
        nick: "ep",
        addr: "irc.freenode.net:6667",
        max_retries: Some(3),
    };

    let mut fanout_sender = client_manager.fanout_sender();
    let mut sock = IrcSocket::new(config);
    thread::spawn(move || {
        // TODO: surface any errors here + shutdown hook
        sock.start_loop(&mut fanout_sender)
    });

    // TODO: hook this up to config
    let listener = TcpListener::bind("127.0.0.1:9123").unwrap();
    for stream in listener.incoming() {
        if let Ok(stream) = stream {
            println!("Client connected");
            client_manager.accept_client_connection(stream);
        }
    }

    Ok(())
}
<|endoftext|>


<commit_before>use std::io::prelude::*;
use std::net::{TcpStream, ToSocketAddrs};

use std::io::{BufReader, BufWriter};

use ber::{self, common};
use ber::types::ASNType;
use Result;

pub struct LDAP
{
    stream: TcpStream,
    msgid: i32,
    enc: ber::Encoder<BufWriter<TcpStream>>,
    dec: ber::Decoder<BufReader<TcpStream>>,
}

impl LDAP
{
    pub fn connect<A: ToSocketAddrs>(addr: A) -> Result<LDAP>
    {
        let stream = try!(TcpStream::connect(addr));
        let enc = ber::Encoder::from_writer(try!(stream.try_clone()));
        let dec = ber::Decoder::from_reader(try!(stream.try_clone()));

        Ok(LDAP
        {
            stream: stream,
            enc: enc,
            dec: dec,
            msgid: 0,
        })
    }

    pub fn send(&mut self, mut tag: common::Tag) -> Result<()>
    {
        let messageid = self.msgid.into_ber_universal();

        let tagvec = vec![messageid, tag];

        try!(self.enc.encode(&tagvec.into_ber_universal()));

        self.enc.flush();

        Ok(())
    }

    pub fn read(&mut self) -> Result<common::Tag>
    {
        let tag = try!(self.dec.decode());

        println!("{:?}", &tag);

        Ok(tag)
    }
}
<commit_msg>We don't need that.<commit_after><|endoftext|>


<commit_before>use std::{
    io::Write,
    net::{TcpListener, TcpStream},
};

use dicom::{
    core::{DataElement, PrimitiveValue, VR},
    dicom_value,
    dictionary_std::tags,
    object::{InMemDicomObject, StandardDataDictionary},
};
use dicom_ul::{
    pdu::reader::read_pdu,
    pdu::{
        writer::write_pdu, PDataValueType, PresentationContextProposed, PresentationContextResult,
    },
};
use structopt::StructOpt;

/// DICOM C-STORE SCP
#[derive(Debug, StructOpt)]
struct App {
    /// verbose mode
    #[structopt(short = "v", long = "verbose")]
    verbose: bool,
    /// enforce max pdu length
    #[structopt(short = "s", long = "strict")]
    strict: bool,
    /// max pdu length
    #[structopt(short = "m", long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,

    /// Which port to listen on
    #[structopt(short, default_value = "11111")]
    port: u16,
}

fn run(
    scu_stream: &mut TcpStream,
    strict: bool,
    verbose: bool,
    max_pdu_length: u32,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut buffer: Vec<u8> = Vec::with_capacity(max_pdu_length as usize);
    let mut instance_buffer: Vec<u8> = Vec::with_capacity(1024 * 1024);
    let mut pcid = 1;
    let mut msgid = 1;
    let mut sop_class_uid = "".to_string();
    let mut sop_instance_uid = "".to_string();
    loop {
        match read_pdu(scu_stream, max_pdu_length, strict) {
            Ok(mut pdu) => {
                if verbose {
                    println!("scu ----> scp: {}", pdu.short_description());
                }
                match pdu {
                    dicom_ul::Pdu::AssociationRQ {
                        protocol_version,
                        calling_ae_title,
                        called_ae_title,
                        application_context_name,
                        presentation_contexts,
                        user_variables,
                    } => {
                        buffer.clear();
                        let PresentationContextProposed {
                            id,
                            abstract_syntax: _,
                            transfer_syntaxes,
                        } = &presentation_contexts[0];
                        let presentation_context_result = PresentationContextResult {
                            id: *id,
                            reason: dicom_ul::pdu::PresentationContextResultReason::Acceptance,
                            transfer_syntax: transfer_syntaxes[0].clone(),
                        };
                        pcid = *id;
                        let response = dicom_ul::Pdu::AssociationAC {
                            protocol_version,
                            calling_ae_title,
                            called_ae_title,
                            application_context_name,
                            presentation_contexts: vec![presentation_context_result],
                            user_variables,
                        };
                        write_pdu(&mut buffer, &response).unwrap();
                        if verbose {
                            println!("scu <---- scp: {}", response.short_description());
                        }
                        scu_stream.write_all(&buffer).unwrap();
                    }
                    dicom_ul::Pdu::PData { ref mut data } => {
                        if data[0].value_type == PDataValueType::Data && !data[0].is_last {
                            instance_buffer.append(&mut data[0].data);
                        } else if data[0].value_type == PDataValueType::Command && data[0].is_last {
                            // commands are always in implict VR LE
                            let ts =
                                dicom::transfer_syntax::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased();
                            let data_value = &data[0];
                            let v = &data_value.data;

                            let obj = InMemDicomObject::read_dataset_with_ts(v.as_slice(), &ts)?;
                            msgid = obj.element(tags::MESSAGE_ID)?.to_int()?;
                            sop_class_uid = obj
                                .element(tags::AFFECTED_SOP_CLASS_UID)?
                                .to_str()?
                                .to_string();
                            sop_instance_uid = obj
                                .element(tags::AFFECTED_SOP_INSTANCE_UID)?
                                .to_str()?
                                .to_string();
                            instance_buffer.clear();
                        } else if data[0].value_type == PDataValueType::Data && data[0].is_last {
                            instance_buffer.append(&mut data[0].data);
                            use std::fs;

                            fs::write(sop_instance_uid.trim_end_matches('\0'), &instance_buffer)
                                .expect("Unable to write file");

                            // send C-STORE-RSP object
                            // commands are always in implict VR LE
                            let ts =
                                dicom::transfer_syntax::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased();

                            let obj =
                                create_cstore_response(msgid, &sop_class_uid, &sop_instance_uid);

                            let mut data = Vec::new();

                            obj.write_dataset_with_ts(&mut data, &ts)?;

                            let pdu_response = dicom_ul::Pdu::PData {
                                data: vec![dicom_ul::pdu::PDataValue {
                                    presentation_context_id: pcid,
                                    value_type: PDataValueType::Command,
                                    is_last: true,
                                    data,
                                }],
                            };
                            buffer.clear();
                            write_pdu(&mut buffer, &pdu_response).unwrap();
                            scu_stream.write_all(&buffer).unwrap();
                        }
                    }
                    dicom_ul::Pdu::ReleaseRQ => {
                        buffer.clear();
                        write_pdu(&mut buffer, &dicom_ul::Pdu::ReleaseRP).unwrap();
                        scu_stream.write_all(&buffer).unwrap();
                    }
                    _ => {}
                }
            }
            Err(dicom_ul::pdu::reader::Error::NoPduAvailable { .. }) => {
                break;
            }
            Err(_err) => {
                break;
            }
        }
    }
    Ok(())
}

fn create_cstore_response(
    message_id: u16,
    sop_class_uid: &str,
    sop_instance_uid: &str,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::create_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            8 + sop_class_uid.len() as i32
                + 8
                + 2
                + 8
                + 2
                + 8
                + 2
                + 8
                + 2
                + sop_instance_uid.len() as i32,
        ),
    ));

    // service
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        dicom_value!(Str, sop_class_uid),
    ));
    // command
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        dicom_value!(U16, [0x8001]),
    ));
    // message ID being responded to
    obj.put(DataElement::new(
        tags::MESSAGE_ID_BEING_RESPONDED_TO,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));
    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0101]),
    ));
    // status https://dicom.nema.org/dicom/2013/output/chtml/part07/chapter_C.html
    obj.put(DataElement::new(
        tags::STATUS,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));
    // SOPInstanceUID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_INSTANCE_UID,
        VR::UI,
        dicom_value!(Str, sop_instance_uid),
    ));

    obj
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let App {
        verbose,
        strict,
        port,
        max_pdu_length,
    } = App::from_args();

    let listen_addr = format!("0.0.0.0:{}", port);
    let listener = TcpListener::bind(&listen_addr).unwrap();
    if verbose {
        println!("listening on: {}", listen_addr);
    }

    for mut stream in listener.incoming() {
        match stream {
            Ok(ref mut scu_stream) => {
                if let Err(e) = run(scu_stream, strict, verbose, max_pdu_length) {
                    eprintln!("[ERROR] {}", e);
                }
            }
            Err(e) => {
                eprintln!("[ERROR] {}", e);
            }
        }
    }

    Ok(())
}
<commit_msg>comments<commit_after>use std::{
    io::Write,
    net::{TcpListener, TcpStream},
};

use dicom::{
    core::{DataElement, PrimitiveValue, VR},
    dicom_value,
    dictionary_std::tags,
    object::{InMemDicomObject, StandardDataDictionary},
};
use dicom_ul::{
    pdu::reader::read_pdu,
    pdu::{
        writer::write_pdu, PDataValueType, PresentationContextProposed, PresentationContextResult,
    },
};
use structopt::StructOpt;

/// DICOM C-STORE SCP
#[derive(Debug, StructOpt)]
struct App {
    /// verbose mode
    #[structopt(short = "v", long = "verbose")]
    verbose: bool,
    /// enforce max pdu length
    #[structopt(short = "s", long = "strict")]
    strict: bool,
    /// max pdu length
    #[structopt(short = "m", long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,

    /// Which port to listen on
    #[structopt(short, default_value = "11111")]
    port: u16,
}

fn run(
    scu_stream: &mut TcpStream,
    strict: bool,
    verbose: bool,
    max_pdu_length: u32,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut buffer: Vec<u8> = Vec::with_capacity(max_pdu_length as usize);
    let mut instance_buffer: Vec<u8> = Vec::with_capacity(1024 * 1024);
    let mut pcid = 1;
    let mut msgid = 1;
    let mut sop_class_uid = "".to_string();
    let mut sop_instance_uid = "".to_string();
    loop {
        match read_pdu(scu_stream, max_pdu_length, strict) {
            Ok(mut pdu) => {
                if verbose {
                    println!("scu ----> scp: {}", pdu.short_description());
                }
                match pdu {
                    dicom_ul::Pdu::AssociationRQ {
                        protocol_version,
                        calling_ae_title,
                        called_ae_title,
                        application_context_name,
                        presentation_contexts,
                        user_variables,
                    } => {
                        buffer.clear();
                        let PresentationContextProposed {
                            id,
                            abstract_syntax: _,
                            transfer_syntaxes,
                        } = &presentation_contexts[0];
                        let presentation_context_result = PresentationContextResult {
                            id: *id,
                            reason: dicom_ul::pdu::PresentationContextResultReason::Acceptance,
                            // accept the first proposed transfer syntax
                            transfer_syntax: transfer_syntaxes[0].clone(),
                        };
                        pcid = *id;

                        // copying most variables for now, should be set to application specific values
                        let response = dicom_ul::Pdu::AssociationAC {
                            protocol_version,
                            calling_ae_title,
                            called_ae_title,
                            application_context_name,
                            presentation_contexts: vec![presentation_context_result],
                            user_variables,
                        };
                        write_pdu(&mut buffer, &response).unwrap();
                        if verbose {
                            println!("scu <---- scp: {}", response.short_description());
                        }
                        scu_stream.write_all(&buffer).unwrap();
                    }
                    dicom_ul::Pdu::PData { ref mut data } => {
                        if data[0].value_type == PDataValueType::Data && !data[0].is_last {
                            instance_buffer.append(&mut data[0].data);
                        } else if data[0].value_type == PDataValueType::Command && data[0].is_last {
                            // commands are always in implict VR LE
                            let ts =
                                dicom::transfer_syntax::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased();
                            let data_value = &data[0];
                            let v = &data_value.data;

                            let obj = InMemDicomObject::read_dataset_with_ts(v.as_slice(), &ts)?;
                            msgid = obj.element(tags::MESSAGE_ID)?.to_int()?;
                            sop_class_uid = obj
                                .element(tags::AFFECTED_SOP_CLASS_UID)?
                                .to_str()?
                                .to_string();
                            sop_instance_uid = obj
                                .element(tags::AFFECTED_SOP_INSTANCE_UID)?
                                .to_str()?
                                .to_string();
                            instance_buffer.clear();
                        } else if data[0].value_type == PDataValueType::Data && data[0].is_last {
                            instance_buffer.append(&mut data[0].data);
                            use std::fs;

                            // write the files to the current directory with their SOPInstanceUID as filenames
                            fs::write(sop_instance_uid.trim_end_matches('\0'), &instance_buffer)
                                .expect("Unable to write file");

                            // send C-STORE-RSP object
                            // commands are always in implict VR LE
                            let ts =
                                dicom::transfer_syntax::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased();

                            let obj =
                                create_cstore_response(msgid, &sop_class_uid, &sop_instance_uid);

                            let mut data = Vec::new();

                            obj.write_dataset_with_ts(&mut data, &ts)?;

                            let pdu_response = dicom_ul::Pdu::PData {
                                data: vec![dicom_ul::pdu::PDataValue {
                                    presentation_context_id: pcid,
                                    value_type: PDataValueType::Command,
                                    is_last: true,
                                    data,
                                }],
                            };
                            buffer.clear();
                            write_pdu(&mut buffer, &pdu_response).unwrap();
                            scu_stream.write_all(&buffer).unwrap();
                        }
                    }
                    dicom_ul::Pdu::ReleaseRQ => {
                        buffer.clear();
                        write_pdu(&mut buffer, &dicom_ul::Pdu::ReleaseRP).unwrap();
                        scu_stream.write_all(&buffer).unwrap();
                    }
                    _ => {}
                }
            }
            Err(dicom_ul::pdu::reader::Error::NoPduAvailable { .. }) => {
                break;
            }
            Err(_err) => {
                break;
            }
        }
    }
    Ok(())
}

fn create_cstore_response(
    message_id: u16,
    sop_class_uid: &str,
    sop_instance_uid: &str,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::create_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            8 + sop_class_uid.len() as i32
                + 8
                + 2
                + 8
                + 2
                + 8
                + 2
                + 8
                + 2
                + sop_instance_uid.len() as i32,
        ),
    ));

    // service
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        dicom_value!(Str, sop_class_uid),
    ));
    // command
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        dicom_value!(U16, [0x8001]),
    ));
    // message ID being responded to
    obj.put(DataElement::new(
        tags::MESSAGE_ID_BEING_RESPONDED_TO,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));
    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0101]),
    ));
    // status https://dicom.nema.org/dicom/2013/output/chtml/part07/chapter_C.html
    obj.put(DataElement::new(
        tags::STATUS,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));
    // SOPInstanceUID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_INSTANCE_UID,
        VR::UI,
        dicom_value!(Str, sop_instance_uid),
    ));

    obj
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let App {
        verbose,
        strict,
        port,
        max_pdu_length,
    } = App::from_args();

    let listen_addr = format!("0.0.0.0:{}", port);
    let listener = TcpListener::bind(&listen_addr).unwrap();
    if verbose {
        println!("listening on: {}", listen_addr);
    }

    for mut stream in listener.incoming() {
        match stream {
            Ok(ref mut scu_stream) => {
                if let Err(e) = run(scu_stream, strict, verbose, max_pdu_length) {
                    eprintln!("[ERROR] {}", e);
                }
            }
            Err(e) => {
                eprintln!("[ERROR] {}", e);
            }
        }
    }

    Ok(())
}
<|endoftext|>


<commit_before>pub use self::fs::FileAbstraction;

// TODO:
// This whole thing can be written better with a trait based mechanism that is embedded into the
// store. However it would mean rewriting most things to be generic which can be a pain in the ass.

#[cfg(test)]
mod fs {
    use error::StoreError as SE;
    use std::io::Cursor;
    use std::path::PathBuf;

    use std::collections::HashMap;
    use std::sync::Mutex;

    lazy_static! {
        static ref MAP: Mutex<HashMap<PathBuf, Cursor<Vec<u8>>>> = {
            Mutex::new(HashMap::new())
        };
    }

    /// `FileAbstraction` type, this is the Test version!
    ///
    /// A lazy file is either absent, but a path to it is available, or it is present.
    #[derive(Debug)]
    pub enum FileAbstraction {
        Absent(PathBuf),
    }

    impl FileAbstraction {

        /**
         * Get the mutable file behind a FileAbstraction object
         */
        pub fn get_file_content(&mut self) -> Result<Cursor<Vec<u8>>, SE> {
            debug!("Getting lazy file: {:?}", self);
            match *self {
                FileAbstraction::Absent(ref f) => {
                    let map = MAP.lock().unwrap();
                    return Ok(map.get(f).unwrap().clone());
                },
            };
        }

        pub fn write_file_content(&mut self, buf: &[u8]) -> Result<(), SE> {
            match *self {
                FileAbstraction::Absent(ref f) => {
                    let mut map = MAP.lock().unwrap();
                    if let Some(ref mut cur) = map.get_mut(f) {
                        let mut vec = cur.get_mut();
                        vec.clear();
                        vec.extend_from_slice(buf);
                        return Ok(());
                    }
                    let vec = Vec::from(buf);
                    map.insert(f.clone(), Cursor::new(vec));
                    return Ok(());
                },
            };
        }

        pub fn remove_file(_: &PathBuf) -> Result<(), SE> {
            Ok(())
        }

        pub fn copy(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            let mut map = MAP.lock().unwrap();
            let a = map.get(from).unwrap().clone();
            map.insert(to.clone(), a);
            Ok(())
        }

        pub fn rename(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            let mut map = MAP.lock().unwrap();
            let a = map.get(from).unwrap().clone();
            map.insert(to.clone(), a);
            Ok(())
        }

        pub fn create_dir_all(_: &PathBuf) -> Result<(), SE> {
            Ok(())
        }
    }
}

#[cfg(not(test))]
mod fs {
    use error::{MapErrInto, StoreError as SE, StoreErrorKind as SEK};
    use std::io::{Seek, SeekFrom, Read};
    use std::path::{Path, PathBuf};
    use std::fs::{File, OpenOptions, create_dir_all, remove_file, copy, rename};

    /// `FileAbstraction` type
    ///
    /// A lazy file is either absent, but a path to it is available, or it is present.
    #[derive(Debug)]
    pub enum FileAbstraction {
        Absent(PathBuf),
        File(File, PathBuf)
    }

    fn open_file<A: AsRef<Path>>(p: A) -> ::std::io::Result<File> {
        OpenOptions::new().write(true).read(true).open(p)
    }

    fn create_file<A: AsRef<Path>>(p: A) -> ::std::io::Result<File> {
        if let Some(parent) = p.as_ref().parent() {
            debug!("Implicitely creating directory: {:?}", parent);
            if let Err(e) = create_dir_all(parent) {
                return Err(e);
            }
        }
        OpenOptions::new().write(true).read(true).create(true).open(p)
    }

    impl FileAbstraction {

        /**
         * Get the content behind this file
         */
        pub fn get_file_content(&mut self) -> Result<&mut Read, SE> {
            debug!("Getting lazy file: {:?}", self);
            let (file, path) = match *self {
                FileAbstraction::File(ref mut f, _) => return {
                    // We seek to the beginning of the file since we expect each
                    // access to the file to be in a different context
                    try!(f.seek(SeekFrom::Start(0))
                        .map_err_into(SEK::FileNotSeeked));
                    Ok(f)
                },
                FileAbstraction::Absent(ref p) => (try!(open_file(p).map_err_into(SEK::FileNotFound)),
                                            p.clone()),
            };
            *self = FileAbstraction::File(file, path);
            if let FileAbstraction::File(ref mut f, _) = *self {
                return Ok(f);
            }
            unreachable!()
        }

        /**
         * Write the content of this file
         */
        pub fn write_file_content(&mut self, buf: &[u8]) -> Result<(), SE> {
            use std::io::Write;
            let (file, path) = match *self {
                FileAbstraction::File(ref mut f, _) => return {
                    // We seek to the beginning of the file since we expect each
                    // access to the file to be in a different context
                    try!(f.seek(SeekFrom::Start(0))
                        .map_err_into(SEK::FileNotCreated));
                    f.write_all(buf).map_err_into(SEK::FileNotWritten)
                },
                FileAbstraction::Absent(ref p) => (try!(create_file(p).map_err_into(SEK::FileNotCreated)),
                                            p.clone()),
            };
            *self = FileAbstraction::File(file, path);
            if let FileAbstraction::File(ref mut f, _) = *self {
                return f.write_all(buf).map_err_into(SEK::FileNotWritten);
            }
            unreachable!();
        }

        pub fn remove_file(path: &PathBuf) -> Result<(), SE> {
            remove_file(path).map_err_into(SEK::FileNotRemoved)
        }

        pub fn copy(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            copy(from, to).map_err_into(SEK::FileNotCopied).map(|_| ())
        }

        pub fn rename(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            rename(from, to).map_err_into(SEK::FileNotRenamed)
        }

        pub fn create_dir_all(path: &PathBuf) -> Result<(), SE> {
            create_dir_all(path).map_err_into(SEK::DirNotCreated)
        }
    }
}

#[cfg(test)]
mod test {
    use super::FileAbstraction;
    use std::io::Read;
    use std::path::PathBuf;

    #[test]
    fn lazy_file() {
        let mut path = PathBuf::from("/tests");
        path.set_file_name("test1");
        let mut lf = FileAbstraction::Absent(path);
        lf.write_file_content(b"Hello World").unwrap();
        let mut bah = Vec::new();
        lf.get_file_content().unwrap().read_to_end(&mut bah).unwrap();
        assert_eq!(bah, b"Hello World");
    }

}
<commit_msg>Fix test backend implementation<commit_after>pub use self::fs::FileAbstraction;

// TODO:
// This whole thing can be written better with a trait based mechanism that is embedded into the
// store. However it would mean rewriting most things to be generic which can be a pain in the ass.

#[cfg(test)]
mod fs {
    use error::StoreError as SE;
    use error::StoreErrorKind as SEK;
    use std::io::Cursor;
    use std::path::PathBuf;

    use libimagerror::into::IntoError;

    use std::collections::HashMap;
    use std::sync::Mutex;

    lazy_static! {
        static ref MAP: Mutex<HashMap<PathBuf, Cursor<Vec<u8>>>> = {
            Mutex::new(HashMap::new())
        };
    }

    /// `FileAbstraction` type, this is the Test version!
    ///
    /// A lazy file is either absent, but a path to it is available, or it is present.
    #[derive(Debug)]
    pub enum FileAbstraction {
        Absent(PathBuf),
    }

    impl FileAbstraction {

        /**
         * Get the mutable file behind a FileAbstraction object
         */
        pub fn get_file_content(&mut self) -> Result<Cursor<Vec<u8>>, SE> {
            debug!("Getting lazy file: {:?}", self);
            match *self {
                FileAbstraction::Absent(ref f) => {
                    let map = MAP.lock().unwrap();
                    return map.get(f).cloned().ok_or(SEK::FileNotFound.into_error());
                },
            };
        }

        pub fn write_file_content(&mut self, buf: &[u8]) -> Result<(), SE> {
            match *self {
                FileAbstraction::Absent(ref f) => {
                    let mut map = MAP.lock().unwrap();
                    if let Some(ref mut cur) = map.get_mut(f) {
                        let mut vec = cur.get_mut();
                        vec.clear();
                        vec.extend_from_slice(buf);
                        return Ok(());
                    }
                    let vec = Vec::from(buf);
                    map.insert(f.clone(), Cursor::new(vec));
                    return Ok(());
                },
            };
        }

        pub fn remove_file(_: &PathBuf) -> Result<(), SE> {
            Ok(())
        }

        pub fn copy(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            let mut map = MAP.lock().unwrap();
            let a = map.get(from).unwrap().clone();
            map.insert(to.clone(), a);
            Ok(())
        }

        pub fn rename(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            let mut map = MAP.lock().unwrap();
            let a = map.get(from).unwrap().clone();
            map.insert(to.clone(), a);
            Ok(())
        }

        pub fn create_dir_all(_: &PathBuf) -> Result<(), SE> {
            Ok(())
        }
    }
}

#[cfg(not(test))]
mod fs {
    use error::{MapErrInto, StoreError as SE, StoreErrorKind as SEK};
    use std::io::{Seek, SeekFrom, Read};
    use std::path::{Path, PathBuf};
    use std::fs::{File, OpenOptions, create_dir_all, remove_file, copy, rename};

    /// `FileAbstraction` type
    ///
    /// A lazy file is either absent, but a path to it is available, or it is present.
    #[derive(Debug)]
    pub enum FileAbstraction {
        Absent(PathBuf),
        File(File, PathBuf)
    }

    fn open_file<A: AsRef<Path>>(p: A) -> ::std::io::Result<File> {
        OpenOptions::new().write(true).read(true).open(p)
    }

    fn create_file<A: AsRef<Path>>(p: A) -> ::std::io::Result<File> {
        if let Some(parent) = p.as_ref().parent() {
            debug!("Implicitely creating directory: {:?}", parent);
            if let Err(e) = create_dir_all(parent) {
                return Err(e);
            }
        }
        OpenOptions::new().write(true).read(true).create(true).open(p)
    }

    impl FileAbstraction {

        /**
         * Get the content behind this file
         */
        pub fn get_file_content(&mut self) -> Result<&mut Read, SE> {
            debug!("Getting lazy file: {:?}", self);
            let (file, path) = match *self {
                FileAbstraction::File(ref mut f, _) => return {
                    // We seek to the beginning of the file since we expect each
                    // access to the file to be in a different context
                    try!(f.seek(SeekFrom::Start(0))
                        .map_err_into(SEK::FileNotSeeked));
                    Ok(f)
                },
                FileAbstraction::Absent(ref p) => (try!(open_file(p).map_err_into(SEK::FileNotFound)),
                                            p.clone()),
            };
            *self = FileAbstraction::File(file, path);
            if let FileAbstraction::File(ref mut f, _) = *self {
                return Ok(f);
            }
            unreachable!()
        }

        /**
         * Write the content of this file
         */
        pub fn write_file_content(&mut self, buf: &[u8]) -> Result<(), SE> {
            use std::io::Write;
            let (file, path) = match *self {
                FileAbstraction::File(ref mut f, _) => return {
                    // We seek to the beginning of the file since we expect each
                    // access to the file to be in a different context
                    try!(f.seek(SeekFrom::Start(0))
                        .map_err_into(SEK::FileNotCreated));
                    f.write_all(buf).map_err_into(SEK::FileNotWritten)
                },
                FileAbstraction::Absent(ref p) => (try!(create_file(p).map_err_into(SEK::FileNotCreated)),
                                            p.clone()),
            };
            *self = FileAbstraction::File(file, path);
            if let FileAbstraction::File(ref mut f, _) = *self {
                return f.write_all(buf).map_err_into(SEK::FileNotWritten);
            }
            unreachable!();
        }

        pub fn remove_file(path: &PathBuf) -> Result<(), SE> {
            remove_file(path).map_err_into(SEK::FileNotRemoved)
        }

        pub fn copy(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            copy(from, to).map_err_into(SEK::FileNotCopied).map(|_| ())
        }

        pub fn rename(from: &PathBuf, to: &PathBuf) -> Result<(), SE> {
            rename(from, to).map_err_into(SEK::FileNotRenamed)
        }

        pub fn create_dir_all(path: &PathBuf) -> Result<(), SE> {
            create_dir_all(path).map_err_into(SEK::DirNotCreated)
        }
    }
}

#[cfg(test)]
mod test {
    use super::FileAbstraction;
    use std::io::Read;
    use std::path::PathBuf;

    #[test]
    fn lazy_file() {
        let mut path = PathBuf::from("/tests");
        path.set_file_name("test1");
        let mut lf = FileAbstraction::Absent(path);
        lf.write_file_content(b"Hello World").unwrap();
        let mut bah = Vec::new();
        lf.get_file_content().unwrap().read_to_end(&mut bah).unwrap();
        assert_eq!(bah, b"Hello World");
    }

}
<|endoftext|>