<commit_before>#![cfg_attr(feature = "cargo-clippy", allow(blacklisted_name))]

use std::borrow::Cow;
use std::collections::BTreeSet;

use proc_macro2::{Ident, Span, TokenStream};
use quote::{format_ident, quote, quote_spanned, TokenStreamExt};

struct X;

impl quote::ToTokens for X {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        tokens.append(Ident::new("X", Span::call_site()));
    }
}

#[test]
fn test_quote_impl() {
    let tokens = quote! {
        impl<'a, T: ToTokens> ToTokens for &'a T {
            fn to_tokens(&self, tokens: &mut TokenStream) {
                (**self).to_tokens(tokens)
            }
        }
    };

    let expected = concat!(
        "impl < 'a , T : ToTokens > ToTokens for & 'a T { ",
        "fn to_tokens (& self , tokens : & mut TokenStream) { ",
        "(* * self) . to_tokens (tokens) ",
        "} ",
        "}"
    );

    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_quote_spanned_impl() {
    let span = Span::call_site();
    let tokens = quote_spanned! {span=>
        impl<'a, T: ToTokens> ToTokens for &'a T {
            fn to_tokens(&self, tokens: &mut TokenStream) {
                (**self).to_tokens(tokens)
            }
        }
    };

    let expected = concat!(
        "impl < 'a , T : ToTokens > ToTokens for & 'a T { ",
        "fn to_tokens (& self , tokens : & mut TokenStream) { ",
        "(* * self) . to_tokens (tokens) ",
        "} ",
        "}"
    );

    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_substitution() {
    let x = X;
    let tokens = quote!(#x <#x> (#x) [#x] {#x});

    let expected = "X < X > (X) [X] { X }";

    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_iter() {
    let primes = &[X, X, X, X];

    assert_eq!("X X X X", quote!(#(#primes)*).to_string());

    assert_eq!("X , X , X , X ,", quote!(#(#primes,)*).to_string());

    assert_eq!("X , X , X , X", quote!(#(#primes),*).to_string());
}

#[test]
fn test_advanced() {
    let generics = quote!( <'a, T> );

    let where_clause = quote!( where T: Serialize );

    let field_ty = quote!(String);

    let item_ty = quote!(Cow<'a, str>);

    let path = quote!(SomeTrait::serialize_with);

    let value = quote!(self.x);

    let tokens = quote! {
        struct SerializeWith #generics #where_clause {
            value: &'a #field_ty,
            phantom: ::std::marker::PhantomData<#item_ty>,
        }

        impl #generics ::serde::Serialize for SerializeWith #generics #where_clause {
            fn serialize<S>(&self, s: &mut S) -> Result<(), S::Error>
                where S: ::serde::Serializer
            {
                #path(self.value, s)
            }
        }

        SerializeWith {
            value: #value,
            phantom: ::std::marker::PhantomData::<#item_ty>,
        }
    };

    let expected = concat!(
        "struct SerializeWith < 'a , T > where T : Serialize { ",
        "value : & 'a String , ",
        "phantom : :: std :: marker :: PhantomData < Cow < 'a , str > > , ",
        "} ",
        "impl < 'a , T > :: serde :: Serialize for SerializeWith < 'a , T > where T : Serialize { ",
        "fn serialize < S > (& self , s : & mut S) -> Result < () , S :: Error > ",
        "where S : :: serde :: Serializer ",
        "{ ",
        "SomeTrait :: serialize_with (self . value , s) ",
        "} ",
        "} ",
        "SerializeWith { ",
        "value : self . x , ",
        "phantom : :: std :: marker :: PhantomData :: < Cow < 'a , str > > , ",
        "}"
    );

    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_integer() {
    let ii8 = -1i8;
    let ii16 = -1i16;
    let ii32 = -1i32;
    let ii64 = -1i64;
    let ii128 = -1i128;
    let iisize = -1isize;
    let uu8 = 1u8;
    let uu16 = 1u16;
    let uu32 = 1u32;
    let uu64 = 1u64;
    let uu128 = 1u128;
    let uusize = 1usize;

    let tokens = quote! {
        #ii8 #ii16 #ii32 #ii64 #ii128 #iisize
        #uu8 #uu16 #uu32 #uu64 #uu128 #uusize
    };
    let expected = "- 1i8 - 1i16 - 1i32 - 1i64 - 1i128 - 1isize 1u8 1u16 1u32 1u64 1u128 1usize";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_floating() {
    let e32 = 2.345f32;

    let e64 = 2.345f64;

    let tokens = quote! {
        #e32
        #e64
    };
    let expected = concat!("2.345f32 2.345f64");
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_char() {
    let zero = '\0';
    let pound = '#';
    let quote = '"';
    let apost = '\'';
    let newline = '\n';
    let heart = '\u{2764}';

    let tokens = quote! {
        #zero #pound #quote #apost #newline #heart
    };
    let expected = "'\\u{0}' '#' '\"' '\\'' '\\n' '\u{2764}'";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_str() {
    let s = "\0 a 'b \" c";
    let tokens = quote!(#s);
    let expected = "\"\\u{0} a 'b \\\" c\"";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_string() {
    let s = "\0 a 'b \" c".to_string();
    let tokens = quote!(#s);
    let expected = "\"\\u{0} a 'b \\\" c\"";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_ident() {
    let foo = Ident::new("Foo", Span::call_site());
    let bar = Ident::new(&format!("Bar{}", 7), Span::call_site());
    let tokens = quote!(struct #foo; enum #bar {});
    let expected = "struct Foo ; enum Bar7 { }";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_duplicate() {
    let ch = 'x';

    let tokens = quote!(#ch #ch);

    let expected = "'x' 'x'";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_fancy_repetition() {
    let foo = vec!["a", "b"];
    let bar = vec![true, false];

    let tokens = quote! {
        #(#foo: #bar),*
    };

    let expected = r#""a" : true , "b" : false"#;
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_nested_fancy_repetition() {
    let nested = vec![vec!['a', 'b', 'c'], vec!['x', 'y', 'z']];

    let tokens = quote! {
        #(
            #(#nested)*
        ),*
    };

    let expected = "'a' 'b' 'c' , 'x' 'y' 'z'";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_duplicate_name_repetition() {
    let foo = &["a", "b"];

    let tokens = quote! {
        #(#foo: #foo),*
        #(#foo: #foo),*
    };

    let expected = r#""a" : "a" , "b" : "b" "a" : "a" , "b" : "b""#;
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_duplicate_name_repetition_no_copy() {
    let foo = vec!["a".to_owned(), "b".to_owned()];

    let tokens = quote! {
        #(#foo: #foo),*
    };

    let expected = r#""a" : "a" , "b" : "b""#;
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_btreeset_repetition() {
    let mut set = BTreeSet::new();
    set.insert("a".to_owned());
    set.insert("b".to_owned());

    let tokens = quote! {
        #(#set: #set),*
    };

    let expected = r#""a" : "a" , "b" : "b""#;
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_variable_name_conflict() {
    // The implementation of `#(...),*` uses the variable `_i` but it should be
    // fine, if a little confusing when debugging.
    let _i = vec!['a', 'b'];
    let tokens = quote! { #(#_i),* };
    let expected = "'a' , 'b'";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_nonrep_in_repetition() {
    let rep = vec!["a", "b"];
    let nonrep = "c";

    let tokens = quote! {
        #(#rep #rep : #nonrep #nonrep),*
    };

    let expected = r#""a" "a" : "c" "c" , "b" "b" : "c" "c""#;
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_empty_quote() {
    let tokens = quote!();
    assert_eq!("", tokens.to_string());
}

#[test]
fn test_box_str() {
    let b = "str".to_owned().into_boxed_str();
    let tokens = quote! { #b };
    assert_eq!("\"str\"", tokens.to_string());
}

#[test]
fn test_cow() {
    let owned: Cow<Ident> = Cow::Owned(Ident::new("owned", Span::call_site()));

    let ident = Ident::new("borrowed", Span::call_site());
    let borrowed = Cow::Borrowed(&ident);

    let tokens = quote! { #owned #borrowed };
    assert_eq!("owned borrowed", tokens.to_string());
}

#[test]
fn test_closure() {
    fn field_i(i: usize) -> Ident {
        format_ident!("__field{}", i)
    }

    let fields = (0usize..3)
        .map(field_i as fn(_) -> _)
        .map(|var| quote! { #var });

    let tokens = quote! { #(#fields)* };
    assert_eq!("__field0 __field1 __field2", tokens.to_string());
}

#[test]
fn test_append_tokens() {
    let mut a = quote!(a);
    let b = quote!(b);
    a.append_all(b);
    assert_eq!("a b", a.to_string());
}

#[test]
fn test_format_ident() {
    let id0 = format_ident!("Aa");
    let id1 = format_ident!("Hello{x}", x = id0);
    let id2 = format_ident!("Hello{x}", x = 5usize);
    let id3 = format_ident!("Hello{}_{x}", id0, x = 10usize);
    let id4 = format_ident!("Aa", span = Span::call_site());
    let id5 = format_ident!("Hello{}", Cow::Borrowed("World"));

    assert_eq!(id0, "Aa");
    assert_eq!(id1, "HelloAa");
    assert_eq!(id2, "Hello5");
    assert_eq!(id3, "HelloAa_10");
    assert_eq!(id4, "Aa");
    assert_eq!(id5, "HelloWorld");
}

#[test]
fn test_format_ident_strip_raw() {
    let id = format_ident!("r#struct");
    let my_id = format_ident!("MyId{}", id);
    let raw_my_id = format_ident!("r#MyId{}", id);

    assert_eq!(id, "r#struct");
    assert_eq!(my_id, "MyIdstruct");
    assert_eq!(raw_my_id, "r#MyIdstruct");
}

#[test]
fn test_outer_line_comment() {
    let tokens = quote! {
        /// doc
    };
    let expected = "# [doc = r\" doc\"]";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_inner_line_comment() {
    let tokens = quote! {
        //! doc
    };
    let expected = "# ! [doc = r\" doc\"]";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_outer_block_comment() {
    let tokens = quote! {
        /** doc */
    };
    let expected = "# [doc = r\" doc \"]";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_inner_block_comment() {
    let tokens = quote! {
        /*! doc */
    };
    let expected = "# ! [doc = r\" doc \"]";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_outer_attr() {
    let tokens = quote! {
        #[inline]
    };
    let expected = "# [inline]";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_inner_attr() {
    let tokens = quote! {
        #![no_std]
    };
    let expected = "# ! [no_std]";
    assert_eq!(expected, tokens.to_string());
}

// https://github.com/dtolnay/quote/issues/130
#[test]
fn test_star_after_repetition() {
    let c = vec!['0', '1'];
    let tokens = quote! {
        #(
            f(#c);
        )*
        *out = None;
    };
    let expected = "f ('0') ; f ('1') ; * out = None ;";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_quote_raw_id() {
    let id = quote!(r#raw_id);
    assert_eq!(id.to_string(), "r#raw_id");
}
<commit_msg>Add test of big array repetition<commit_after>#![cfg_attr(feature = "cargo-clippy", allow(blacklisted_name))]

use std::borrow::Cow;
use std::collections::BTreeSet;

use proc_macro2::{Ident, Span, TokenStream};
use quote::{format_ident, quote, quote_spanned, TokenStreamExt};

struct X;

impl quote::ToTokens for X {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        tokens.append(Ident::new("X", Span::call_site()));
    }
}

#[test]
fn test_quote_impl() {
    let tokens = quote! {
        impl<'a, T: ToTokens> ToTokens for &'a T {
            fn to_tokens(&self, tokens: &mut TokenStream) {
                (**self).to_tokens(tokens)
            }
        }
    };

    let expected = concat!(
        "impl < 'a , T : ToTokens > ToTokens for & 'a T { ",
        "fn to_tokens (& self , tokens : & mut TokenStream) { ",
        "(* * self) . to_tokens (tokens) ",
        "} ",
        "}"
    );

    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_quote_spanned_impl() {
    let span = Span::call_site();
    let tokens = quote_spanned! {span=>
        impl<'a, T: ToTokens> ToTokens for &'a T {
            fn to_tokens(&self, tokens: &mut TokenStream) {
                (**self).to_tokens(tokens)
            }
        }
    };

    let expected = concat!(
        "impl < 'a , T : ToTokens > ToTokens for & 'a T { ",
        "fn to_tokens (& self , tokens : & mut TokenStream) { ",
        "(* * self) . to_tokens (tokens) ",
        "} ",
        "}"
    );

    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_substitution() {
    let x = X;
    let tokens = quote!(#x <#x> (#x) [#x] {#x});

    let expected = "X < X > (X) [X] { X }";

    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_iter() {
    let primes = &[X, X, X, X];

    assert_eq!("X X X X", quote!(#(#primes)*).to_string());

    assert_eq!("X , X , X , X ,", quote!(#(#primes,)*).to_string());

    assert_eq!("X , X , X , X", quote!(#(#primes),*).to_string());
}

#[test]
fn test_array() {
    let array: [u8; 40] = [0; 40];
    let _ = quote!(#(#array #array)*);

    let ref_array: &[u8; 40] = &[0; 40];
    let _ = quote!(#(#ref_array #ref_array)*);

    let ref_slice: &[u8] = &[0; 40];
    let _ = quote!(#(#ref_slice #ref_slice)*);
}

#[test]
fn test_advanced() {
    let generics = quote!( <'a, T> );

    let where_clause = quote!( where T: Serialize );

    let field_ty = quote!(String);

    let item_ty = quote!(Cow<'a, str>);

    let path = quote!(SomeTrait::serialize_with);

    let value = quote!(self.x);

    let tokens = quote! {
        struct SerializeWith #generics #where_clause {
            value: &'a #field_ty,
            phantom: ::std::marker::PhantomData<#item_ty>,
        }

        impl #generics ::serde::Serialize for SerializeWith #generics #where_clause {
            fn serialize<S>(&self, s: &mut S) -> Result<(), S::Error>
                where S: ::serde::Serializer
            {
                #path(self.value, s)
            }
        }

        SerializeWith {
            value: #value,
            phantom: ::std::marker::PhantomData::<#item_ty>,
        }
    };

    let expected = concat!(
        "struct SerializeWith < 'a , T > where T : Serialize { ",
        "value : & 'a String , ",
        "phantom : :: std :: marker :: PhantomData < Cow < 'a , str > > , ",
        "} ",
        "impl < 'a , T > :: serde :: Serialize for SerializeWith < 'a , T > where T : Serialize { ",
        "fn serialize < S > (& self , s : & mut S) -> Result < () , S :: Error > ",
        "where S : :: serde :: Serializer ",
        "{ ",
        "SomeTrait :: serialize_with (self . value , s) ",
        "} ",
        "} ",
        "SerializeWith { ",
        "value : self . x , ",
        "phantom : :: std :: marker :: PhantomData :: < Cow < 'a , str > > , ",
        "}"
    );

    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_integer() {
    let ii8 = -1i8;
    let ii16 = -1i16;
    let ii32 = -1i32;
    let ii64 = -1i64;
    let ii128 = -1i128;
    let iisize = -1isize;
    let uu8 = 1u8;
    let uu16 = 1u16;
    let uu32 = 1u32;
    let uu64 = 1u64;
    let uu128 = 1u128;
    let uusize = 1usize;

    let tokens = quote! {
        #ii8 #ii16 #ii32 #ii64 #ii128 #iisize
        #uu8 #uu16 #uu32 #uu64 #uu128 #uusize
    };
    let expected = "- 1i8 - 1i16 - 1i32 - 1i64 - 1i128 - 1isize 1u8 1u16 1u32 1u64 1u128 1usize";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_floating() {
    let e32 = 2.345f32;

    let e64 = 2.345f64;

    let tokens = quote! {
        #e32
        #e64
    };
    let expected = concat!("2.345f32 2.345f64");
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_char() {
    let zero = '\0';
    let pound = '#';
    let quote = '"';
    let apost = '\'';
    let newline = '\n';
    let heart = '\u{2764}';

    let tokens = quote! {
        #zero #pound #quote #apost #newline #heart
    };
    let expected = "'\\u{0}' '#' '\"' '\\'' '\\n' '\u{2764}'";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_str() {
    let s = "\0 a 'b \" c";
    let tokens = quote!(#s);
    let expected = "\"\\u{0} a 'b \\\" c\"";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_string() {
    let s = "\0 a 'b \" c".to_string();
    let tokens = quote!(#s);
    let expected = "\"\\u{0} a 'b \\\" c\"";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_ident() {
    let foo = Ident::new("Foo", Span::call_site());
    let bar = Ident::new(&format!("Bar{}", 7), Span::call_site());
    let tokens = quote!(struct #foo; enum #bar {});
    let expected = "struct Foo ; enum Bar7 { }";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_duplicate() {
    let ch = 'x';

    let tokens = quote!(#ch #ch);

    let expected = "'x' 'x'";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_fancy_repetition() {
    let foo = vec!["a", "b"];
    let bar = vec![true, false];

    let tokens = quote! {
        #(#foo: #bar),*
    };

    let expected = r#""a" : true , "b" : false"#;
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_nested_fancy_repetition() {
    let nested = vec![vec!['a', 'b', 'c'], vec!['x', 'y', 'z']];

    let tokens = quote! {
        #(
            #(#nested)*
        ),*
    };

    let expected = "'a' 'b' 'c' , 'x' 'y' 'z'";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_duplicate_name_repetition() {
    let foo = &["a", "b"];

    let tokens = quote! {
        #(#foo: #foo),*
        #(#foo: #foo),*
    };

    let expected = r#""a" : "a" , "b" : "b" "a" : "a" , "b" : "b""#;
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_duplicate_name_repetition_no_copy() {
    let foo = vec!["a".to_owned(), "b".to_owned()];

    let tokens = quote! {
        #(#foo: #foo),*
    };

    let expected = r#""a" : "a" , "b" : "b""#;
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_btreeset_repetition() {
    let mut set = BTreeSet::new();
    set.insert("a".to_owned());
    set.insert("b".to_owned());

    let tokens = quote! {
        #(#set: #set),*
    };

    let expected = r#""a" : "a" , "b" : "b""#;
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_variable_name_conflict() {
    // The implementation of `#(...),*` uses the variable `_i` but it should be
    // fine, if a little confusing when debugging.
    let _i = vec!['a', 'b'];
    let tokens = quote! { #(#_i),* };
    let expected = "'a' , 'b'";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_nonrep_in_repetition() {
    let rep = vec!["a", "b"];
    let nonrep = "c";

    let tokens = quote! {
        #(#rep #rep : #nonrep #nonrep),*
    };

    let expected = r#""a" "a" : "c" "c" , "b" "b" : "c" "c""#;
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_empty_quote() {
    let tokens = quote!();
    assert_eq!("", tokens.to_string());
}

#[test]
fn test_box_str() {
    let b = "str".to_owned().into_boxed_str();
    let tokens = quote! { #b };
    assert_eq!("\"str\"", tokens.to_string());
}

#[test]
fn test_cow() {
    let owned: Cow<Ident> = Cow::Owned(Ident::new("owned", Span::call_site()));

    let ident = Ident::new("borrowed", Span::call_site());
    let borrowed = Cow::Borrowed(&ident);

    let tokens = quote! { #owned #borrowed };
    assert_eq!("owned borrowed", tokens.to_string());
}

#[test]
fn test_closure() {
    fn field_i(i: usize) -> Ident {
        format_ident!("__field{}", i)
    }

    let fields = (0usize..3)
        .map(field_i as fn(_) -> _)
        .map(|var| quote! { #var });

    let tokens = quote! { #(#fields)* };
    assert_eq!("__field0 __field1 __field2", tokens.to_string());
}

#[test]
fn test_append_tokens() {
    let mut a = quote!(a);
    let b = quote!(b);
    a.append_all(b);
    assert_eq!("a b", a.to_string());
}

#[test]
fn test_format_ident() {
    let id0 = format_ident!("Aa");
    let id1 = format_ident!("Hello{x}", x = id0);
    let id2 = format_ident!("Hello{x}", x = 5usize);
    let id3 = format_ident!("Hello{}_{x}", id0, x = 10usize);
    let id4 = format_ident!("Aa", span = Span::call_site());
    let id5 = format_ident!("Hello{}", Cow::Borrowed("World"));

    assert_eq!(id0, "Aa");
    assert_eq!(id1, "HelloAa");
    assert_eq!(id2, "Hello5");
    assert_eq!(id3, "HelloAa_10");
    assert_eq!(id4, "Aa");
    assert_eq!(id5, "HelloWorld");
}

#[test]
fn test_format_ident_strip_raw() {
    let id = format_ident!("r#struct");
    let my_id = format_ident!("MyId{}", id);
    let raw_my_id = format_ident!("r#MyId{}", id);

    assert_eq!(id, "r#struct");
    assert_eq!(my_id, "MyIdstruct");
    assert_eq!(raw_my_id, "r#MyIdstruct");
}

#[test]
fn test_outer_line_comment() {
    let tokens = quote! {
        /// doc
    };
    let expected = "# [doc = r\" doc\"]";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_inner_line_comment() {
    let tokens = quote! {
        //! doc
    };
    let expected = "# ! [doc = r\" doc\"]";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_outer_block_comment() {
    let tokens = quote! {
        /** doc */
    };
    let expected = "# [doc = r\" doc \"]";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_inner_block_comment() {
    let tokens = quote! {
        /*! doc */
    };
    let expected = "# ! [doc = r\" doc \"]";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_outer_attr() {
    let tokens = quote! {
        #[inline]
    };
    let expected = "# [inline]";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_inner_attr() {
    let tokens = quote! {
        #![no_std]
    };
    let expected = "# ! [no_std]";
    assert_eq!(expected, tokens.to_string());
}

// https://github.com/dtolnay/quote/issues/130
#[test]
fn test_star_after_repetition() {
    let c = vec!['0', '1'];
    let tokens = quote! {
        #(
            f(#c);
        )*
        *out = None;
    };
    let expected = "f ('0') ; f ('1') ; * out = None ;";
    assert_eq!(expected, tokens.to_string());
}

#[test]
fn test_quote_raw_id() {
    let id = quote!(r#raw_id);
    assert_eq!(id.to_string(), "r#raw_id");
}
<|endoftext|>


<commit_before>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use seccomp::{
    Error, SeccompAction, SeccompCmpArgLen as ArgLen, SeccompCmpOp::Eq, SeccompCondition as Cond,
    SeccompRule,
};

#[macro_use]
mod macros;
mod filters;

pub use self::filters::default_filter;
pub use self::filters::get_seccomp_filter;

// See include/uapi/asm-generic/fcntl.h in the kernel code.
const FCNTL_FD_CLOEXEC: u64 = 1;
const FCNTL_F_SETFD: u64 = 2;

// See include/uapi/linux/futex.h in the kernel code.
const FUTEX_WAIT: u64 = 0;
const FUTEX_WAKE: u64 = 1;
const FUTEX_REQUEUE: u64 = 3;
#[cfg(target_env = "gnu")]
const FUTEX_CMP_REQUEUE: u64 = 4;
const FUTEX_PRIVATE_FLAG: u64 = 128;
const FUTEX_WAIT_PRIVATE: u64 = FUTEX_WAIT | FUTEX_PRIVATE_FLAG;
const FUTEX_WAKE_PRIVATE: u64 = FUTEX_WAKE | FUTEX_PRIVATE_FLAG;
const FUTEX_REQUEUE_PRIVATE: u64 = FUTEX_REQUEUE | FUTEX_PRIVATE_FLAG;
#[cfg(target_env = "gnu")]
const FUTEX_CMP_REQUEUE_PRIVATE: u64 = FUTEX_CMP_REQUEUE | FUTEX_PRIVATE_FLAG;

// See include/uapi/asm-generic/ioctls.h in the kernel code.
const TCGETS: u64 = 0x5401;
const TCSETS: u64 = 0x5402;
const TIOCGWINSZ: u64 = 0x5413;
const FIOCLEX: u64 = 0x5451;
const FIONBIO: u64 = 0x5421;

// Hardcoded here instead of getting values from kvm-ioctls, so that filtered values cannot be
// mistakenly or intentionally altered from outside our codebase.

// See include/uapi/linux/if_tun.h in the kernel code.
const KVM_GET_API_VERSION: u64 = 0xae00;
const KVM_CREATE_VM: u64 = 0xae01;
const KVM_CHECK_EXTENSION: u64 = 0xae03;
const KVM_GET_VCPU_MMAP_SIZE: u64 = 0xae04;
const KVM_CREATE_VCPU: u64 = 0xae41;
const KVM_SET_TSS_ADDR: u64 = 0xae47;
const KVM_CREATE_IRQCHIP: u64 = 0xae60;
const KVM_RUN: u64 = 0xae80;
const KVM_SET_MSRS: u64 = 0x4008_ae89;
const KVM_SET_CPUID2: u64 = 0x4008_ae90;
const KVM_SET_USER_MEMORY_REGION: u64 = 0x4020_ae46;
const KVM_IRQFD: u64 = 0x4020_ae76;
const KVM_CREATE_PIT2: u64 = 0x4040_ae77;
const KVM_IOEVENTFD: u64 = 0x4040_ae79;
const KVM_SET_REGS: u64 = 0x4090_ae82;
const KVM_SET_SREGS: u64 = 0x4138_ae84;
const KVM_SET_FPU: u64 = 0x41a0_ae8d;
const KVM_SET_LAPIC: u64 = 0x4400_ae8f;
const KVM_GET_SREGS: u64 = 0x8138_ae83;
const KVM_GET_LAPIC: u64 = 0x8400_ae8e;
const KVM_GET_MSR_INDEX_LIST: u64 = 0xc004_ae02;
const KVM_GET_MSR_FEATURE_INDEX_LIST: u64 = 0xc004_ae0a;
const KVM_GET_SUPPORTED_CPUID: u64 = 0xc008_ae05;
const KVM_GET_IRQCHIP: u64 = 0xc208_ae62;
const KVM_SET_IRQCHIP: u64 = 0x8208_ae63;
const KVM_SET_CLOCK: u64 = 0x4030_ae7b;
const KVM_GET_CLOCK: u64 = 0x8030_ae7c;
const KVM_GET_PIT2: u64 = 0x8070_ae9f;
const KVM_SET_PIT2: u64 = 0x4070_aea0;
const KVM_GET_REGS: u64 = 0x8090_ae81;
const KVM_GET_MSRS: u64 = 0xc008_ae88;
const KVM_GET_CPUID2: u64 = 0xc008_ae91;
const KVM_GET_MP_STATE: u64 = 0x8004_ae98;
const KVM_SET_MP_STATE: u64 = 0x4004_ae99;
const KVM_GET_VCPU_EVENTS: u64 = 0x8040_ae9f;
const KVM_SET_VCPU_EVENTS: u64 = 0x4040_aea0;
const KVM_GET_DEBUGREGS: u64 = 0x8080_aea1;
const KVM_SET_DEBUGREGS: u64 = 0x4080_aea2;
const KVM_GET_XSAVE: u64 = 0x9000_aea4;
const KVM_SET_XSAVE: u64 = 0x5000_aea5;
const KVM_GET_XCRS: u64 = 0x8188_aea6;
const KVM_SET_XCRS: u64 = 0x4188_aea7;

// See include/uapi/linux/if_tun.h in the kernel code.
const TUNSETIFF: u64 = 0x4004_54ca;
const TUNSETOFFLOAD: u64 = 0x4004_54d0;
const TUNSETVNETHDRSZ: u64 = 0x4004_54d8;

fn create_ioctl_seccomp_rule() -> Result<Vec<SeccompRule>, Error> {
    Ok(or![
        and![Cond::new(1, ArgLen::DWORD, Eq, TCSETS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, TCGETS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, TIOCGWINSZ)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_CHECK_EXTENSION,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_CREATE_VM)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_API_VERSION,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_SUPPORTED_CPUID,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_VCPU_MMAP_SIZE,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_CREATE_IRQCHIP,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_CREATE_PIT2)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_CREATE_VCPU)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_IOEVENTFD)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_IRQFD)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_TSS_ADDR,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_USER_MEMORY_REGION,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, FIOCLEX)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, FIONBIO)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, TUNSETIFF)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, TUNSETOFFLOAD)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, TUNSETVNETHDRSZ)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_LAPIC)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_SREGS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_RUN)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_CPUID2)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_FPU)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_LAPIC)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_MSRS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_REGS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_SREGS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_MSR_INDEX_LIST)?],
        and![Cond::new(
            1,
            ArgLen::DWORD,
            Eq,
            KVM_GET_MSR_FEATURE_INDEX_LIST
        )?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_IRQCHIP)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_IRQCHIP)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_CLOCK)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_CLOCK)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_PIT2)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_PIT2)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_REGS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_MSRS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_CPUID2)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_MP_STATE)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_MP_STATE)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_VCPU_EVENTS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_VCPU_EVENTS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_DEBUGREGS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_DEBUGREGS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_XSAVE)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_XSAVE)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_XCRS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_XCRS)?],
    ])
}

#[cfg(test)]
#[cfg(target_env = "musl")]
mod tests {
    use super::*;
    use seccomp::SeccompFilter;
    use std::convert::TryInto;
    use std::thread;

    const EXTRA_SYSCALLS: [i64; 5] = [
        libc::SYS_clone,
        libc::SYS_mprotect,
        libc::SYS_rt_sigprocmask,
        libc::SYS_set_tid_address,
        libc::SYS_sigaltstack,
    ];

    fn add_syscalls_install_filter(mut filter: SeccompFilter) {
        // Test error case: add empty rule array.
        assert!(filter.add_rules(0, vec![],).is_err());
        // Add "Allow" rule for each syscall.
        for syscall in EXTRA_SYSCALLS.iter() {
            assert!(filter
                .add_rules(
                    *syscall,
                    vec![SeccompRule::new(vec![], SeccompAction::Allow)],
                )
                .is_ok());
        }
        assert!(SeccompFilter::apply(filter.try_into().unwrap()).is_ok());
    }

    #[test]
    fn test_basic_seccomp() {
        // Spawn a new thread before running the tests because all tests run
        // in the same thread. Otherwise other tests will fail because of the
        // installed seccomp filters.
        thread::spawn(move || {
            let filter = default_filter().unwrap().allow_all();
            add_syscalls_install_filter(filter);
        })
        .join()
        .unwrap();
    }

    #[test]
    fn test_advanced_seccomp() {
        // Spawn a new thread before running the tests because all tests run
        // in the same thread. Otherwise other tests will fail because of the
        // installed seccomp filters.
        thread::spawn(move || {
            let filter = default_filter().unwrap();
            add_syscalls_install_filter(filter);
        })
        .join()
        .unwrap();
    }
}
<commit_msg>seccomp: allow KVM_GET_DIRTY_LOG ioctl<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use seccomp::{
    Error, SeccompAction, SeccompCmpArgLen as ArgLen, SeccompCmpOp::Eq, SeccompCondition as Cond,
    SeccompRule,
};

#[macro_use]
mod macros;
mod filters;

pub use self::filters::default_filter;
pub use self::filters::get_seccomp_filter;

// See include/uapi/asm-generic/fcntl.h in the kernel code.
const FCNTL_FD_CLOEXEC: u64 = 1;
const FCNTL_F_SETFD: u64 = 2;

// See include/uapi/linux/futex.h in the kernel code.
const FUTEX_WAIT: u64 = 0;
const FUTEX_WAKE: u64 = 1;
const FUTEX_REQUEUE: u64 = 3;
#[cfg(target_env = "gnu")]
const FUTEX_CMP_REQUEUE: u64 = 4;
const FUTEX_PRIVATE_FLAG: u64 = 128;
const FUTEX_WAIT_PRIVATE: u64 = FUTEX_WAIT | FUTEX_PRIVATE_FLAG;
const FUTEX_WAKE_PRIVATE: u64 = FUTEX_WAKE | FUTEX_PRIVATE_FLAG;
const FUTEX_REQUEUE_PRIVATE: u64 = FUTEX_REQUEUE | FUTEX_PRIVATE_FLAG;
#[cfg(target_env = "gnu")]
const FUTEX_CMP_REQUEUE_PRIVATE: u64 = FUTEX_CMP_REQUEUE | FUTEX_PRIVATE_FLAG;

// See include/uapi/asm-generic/ioctls.h in the kernel code.
const TCGETS: u64 = 0x5401;
const TCSETS: u64 = 0x5402;
const TIOCGWINSZ: u64 = 0x5413;
const FIOCLEX: u64 = 0x5451;
const FIONBIO: u64 = 0x5421;

// Hardcoded here instead of getting values from kvm-ioctls, so that filtered values cannot be
// mistakenly or intentionally altered from outside our codebase.

// See include/uapi/linux/if_tun.h in the kernel code.
const KVM_GET_API_VERSION: u64 = 0xae00;
const KVM_CREATE_VM: u64 = 0xae01;
const KVM_CHECK_EXTENSION: u64 = 0xae03;
const KVM_GET_VCPU_MMAP_SIZE: u64 = 0xae04;
const KVM_CREATE_VCPU: u64 = 0xae41;
const KVM_GET_DIRTY_LOG: u64 = 0x4010_ae42;
const KVM_SET_TSS_ADDR: u64 = 0xae47;
const KVM_CREATE_IRQCHIP: u64 = 0xae60;
const KVM_RUN: u64 = 0xae80;
const KVM_SET_MSRS: u64 = 0x4008_ae89;
const KVM_SET_CPUID2: u64 = 0x4008_ae90;
const KVM_SET_USER_MEMORY_REGION: u64 = 0x4020_ae46;
const KVM_IRQFD: u64 = 0x4020_ae76;
const KVM_CREATE_PIT2: u64 = 0x4040_ae77;
const KVM_IOEVENTFD: u64 = 0x4040_ae79;
const KVM_SET_REGS: u64 = 0x4090_ae82;
const KVM_SET_SREGS: u64 = 0x4138_ae84;
const KVM_SET_FPU: u64 = 0x41a0_ae8d;
const KVM_SET_LAPIC: u64 = 0x4400_ae8f;
const KVM_GET_SREGS: u64 = 0x8138_ae83;
const KVM_GET_LAPIC: u64 = 0x8400_ae8e;
const KVM_GET_MSR_INDEX_LIST: u64 = 0xc004_ae02;
const KVM_GET_MSR_FEATURE_INDEX_LIST: u64 = 0xc004_ae0a;
const KVM_GET_SUPPORTED_CPUID: u64 = 0xc008_ae05;
const KVM_GET_IRQCHIP: u64 = 0xc208_ae62;
const KVM_SET_IRQCHIP: u64 = 0x8208_ae63;
const KVM_SET_CLOCK: u64 = 0x4030_ae7b;
const KVM_GET_CLOCK: u64 = 0x8030_ae7c;
const KVM_GET_PIT2: u64 = 0x8070_ae9f;
const KVM_SET_PIT2: u64 = 0x4070_aea0;
const KVM_GET_REGS: u64 = 0x8090_ae81;
const KVM_GET_MSRS: u64 = 0xc008_ae88;
const KVM_GET_CPUID2: u64 = 0xc008_ae91;
const KVM_GET_MP_STATE: u64 = 0x8004_ae98;
const KVM_SET_MP_STATE: u64 = 0x4004_ae99;
const KVM_GET_VCPU_EVENTS: u64 = 0x8040_ae9f;
const KVM_SET_VCPU_EVENTS: u64 = 0x4040_aea0;
const KVM_GET_DEBUGREGS: u64 = 0x8080_aea1;
const KVM_SET_DEBUGREGS: u64 = 0x4080_aea2;
const KVM_GET_XSAVE: u64 = 0x9000_aea4;
const KVM_SET_XSAVE: u64 = 0x5000_aea5;
const KVM_GET_XCRS: u64 = 0x8188_aea6;
const KVM_SET_XCRS: u64 = 0x4188_aea7;

// See include/uapi/linux/if_tun.h in the kernel code.
const TUNSETIFF: u64 = 0x4004_54ca;
const TUNSETOFFLOAD: u64 = 0x4004_54d0;
const TUNSETVNETHDRSZ: u64 = 0x4004_54d8;

fn create_ioctl_seccomp_rule() -> Result<Vec<SeccompRule>, Error> {
    Ok(or![
        and![Cond::new(1, ArgLen::DWORD, Eq, TCSETS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, TCGETS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, TIOCGWINSZ)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_CHECK_EXTENSION,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_CREATE_VM)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_API_VERSION,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_SUPPORTED_CPUID,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_VCPU_MMAP_SIZE,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_CREATE_IRQCHIP,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_CREATE_PIT2)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_CREATE_VCPU)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_DIRTY_LOG)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_IOEVENTFD)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_IRQFD)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_TSS_ADDR,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_USER_MEMORY_REGION,)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, FIOCLEX)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, FIONBIO)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, TUNSETIFF)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, TUNSETOFFLOAD)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, TUNSETVNETHDRSZ)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_LAPIC)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_SREGS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_RUN)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_CPUID2)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_FPU)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_LAPIC)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_MSRS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_REGS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_SREGS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_MSR_INDEX_LIST)?],
        and![Cond::new(
            1,
            ArgLen::DWORD,
            Eq,
            KVM_GET_MSR_FEATURE_INDEX_LIST
        )?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_IRQCHIP)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_IRQCHIP)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_CLOCK)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_CLOCK)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_PIT2)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_PIT2)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_REGS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_MSRS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_CPUID2)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_MP_STATE)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_MP_STATE)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_VCPU_EVENTS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_VCPU_EVENTS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_DEBUGREGS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_DEBUGREGS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_XSAVE)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_XSAVE)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_GET_XCRS)?],
        and![Cond::new(1, ArgLen::DWORD, Eq, KVM_SET_XCRS)?],
    ])
}

#[cfg(test)]
#[cfg(target_env = "musl")]
mod tests {
    use super::*;
    use seccomp::SeccompFilter;
    use std::convert::TryInto;
    use std::thread;

    const EXTRA_SYSCALLS: [i64; 5] = [
        libc::SYS_clone,
        libc::SYS_mprotect,
        libc::SYS_rt_sigprocmask,
        libc::SYS_set_tid_address,
        libc::SYS_sigaltstack,
    ];

    fn add_syscalls_install_filter(mut filter: SeccompFilter) {
        // Test error case: add empty rule array.
        assert!(filter.add_rules(0, vec![],).is_err());
        // Add "Allow" rule for each syscall.
        for syscall in EXTRA_SYSCALLS.iter() {
            assert!(filter
                .add_rules(
                    *syscall,
                    vec![SeccompRule::new(vec![], SeccompAction::Allow)],
                )
                .is_ok());
        }
        assert!(SeccompFilter::apply(filter.try_into().unwrap()).is_ok());
    }

    #[test]
    fn test_basic_seccomp() {
        // Spawn a new thread before running the tests because all tests run
        // in the same thread. Otherwise other tests will fail because of the
        // installed seccomp filters.
        thread::spawn(move || {
            let filter = default_filter().unwrap().allow_all();
            add_syscalls_install_filter(filter);
        })
        .join()
        .unwrap();
    }

    #[test]
    fn test_advanced_seccomp() {
        // Spawn a new thread before running the tests because all tests run
        // in the same thread. Otherwise other tests will fail because of the
        // installed seccomp filters.
        thread::spawn(move || {
            let filter = default_filter().unwrap();
            add_syscalls_install_filter(filter);
        })
        .join()
        .unwrap();
    }
}
<|endoftext|>


<commit_before>//! Reader and writer for the refFlat format.
//!
//! The refFlat format is a transcript-oriented format in which each
//! transcript is denoted in a single line. It is used most prominently
//! by the [picard suite tools](https://broadinstitute.github.io/picard/)
//!
//! A minimum specification of the columns can be found here:
//! https://genome.ucsc.edu/goldenPath/gbdDescriptionsOld.html#RefFlat

use std::io;
use std::fs;
use std::convert::AsRef;
use std::path::Path;
use std::str::FromStr;

use csv;

use feature::FeatureError;
use {Coord, Strand, Transcript, TBuilder, Error};

pub type RefFlatRow = (String, String, String, char, u64, u64, u64, u64, usize, String, String);

pub struct RefFlatRecord {
    pub gene_id: String,
    pub transcript_name: String,
    pub seq_name: String,
    pub strand_char: char,
    pub trx_start: u64,
    pub trx_end: u64,
    pub coding_start: u64,
    pub coding_end: u64,
    pub num_exons: usize,
    pub exon_starts: String,
    pub exon_ends: String,
}

impl RefFlatRecord {
    pub fn new(row: RefFlatRow) -> RefFlatRecord {
        RefFlatRecord {
            gene_id: row.0,
            transcript_name: row.1,
            seq_name: row.2,
            strand_char: row.3,
            trx_start: row.4,
            trx_end: row.5,
            coding_start: row.6,
            coding_end: row.7,
            num_exons: row.8,
            exon_starts: row.9,
            exon_ends: row.10,
        }
    }

    pub fn to_transcript(self) -> Result<Transcript, Error> {

        let exon_coords = self.zip_raw_exon_coords();
        if exon_coords.len() != self.num_exons {
            return Err(Error::RefFlat(
                "number of exon and exon coordinates mismatch"));
        }

        let strand = Strand::from_char(&self.strand_char)
            .map_err(FeatureError::from)
            .map_err(Error::from)?;

        let coding_coord =
            if self.coding_start != self.coding_end {
                Some((self.coding_start, self.coding_end))
            } else {
                None
            };

        TBuilder::new(self.seq_name, self.trx_start, self.trx_end)
            .id(self.transcript_name)
            .strand(strand)
            .coords(exon_coords, coding_coord)
            .coding_incl_stop(true)
            .attribute("gene_id", self.gene_id)
            .build()
            .map_err(Error::from)
    }

    #[inline]
    fn zip_raw_exon_coords(&self) -> Vec<Coord<u64>> {
        let exon_starts = self.exon_starts
            .split(',').filter_map(|item| u64::from_str(item).ok());
        let exon_ends = self.exon_ends
            .split(',').filter_map(|item| u64::from_str(item).ok());
        exon_starts.zip(exon_ends).collect()
    }
}

pub struct Reader<R: io::Read> {
    inner: csv::Reader<R>,
}

impl<R: io::Read> Reader<R> {

    pub fn from_reader(in_reader: R) -> Reader<R> {
        Reader {
            inner: csv::Reader::from_reader(in_reader)
                .delimiter(b'\t')
                .has_headers(false)
        }
    }

    pub fn transcripts<'a>(&'a mut self) -> RefFlatTranscripts<'a, R> {
        RefFlatTranscripts { inner: self.inner.decode() }
    }
}

impl Reader<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::open(path).map(Reader::from_reader)
    }
}

pub struct RefFlatTranscripts<'a, R: 'a> where R: io::Read {
    inner: csv::DecodedRecords<'a, R, RefFlatRow>,
}

impl<'a, R> Iterator for RefFlatTranscripts<'a, R> where R: io::Read {

    type Item = Result<Transcript, Error>;

    fn next(&mut self) -> Option<Result<Transcript, Error>> {
        self.inner.next().map(|row| {
            row.or_else(|err| Err(Error::from(err)))
                .map(RefFlatRecord::new)
                .and_then(|record| record.to_transcript())
        })
    }
}
<commit_msg>Add RefFlatRecord iterator for RefFlatReader<commit_after>//! Reader and writer for the refFlat format.
//!
//! The refFlat format is a transcript-oriented format in which each
//! transcript is denoted in a single line. It is used most prominently
//! by the [picard suite tools](https://broadinstitute.github.io/picard/)
//!
//! A minimum specification of the columns can be found here:
//! https://genome.ucsc.edu/goldenPath/gbdDescriptionsOld.html#RefFlat

use std::io;
use std::fs;
use std::convert::AsRef;
use std::path::Path;
use std::str::FromStr;

use csv;

use feature::FeatureError;
use {Coord, Strand, Transcript, TBuilder, Error};

pub type RefFlatRow = (String, String, String, char, u64, u64, u64, u64, usize, String, String);

pub struct RefFlatRecord {
    pub gene_id: String,
    pub transcript_name: String,
    pub seq_name: String,
    pub strand_char: char,
    pub trx_start: u64,
    pub trx_end: u64,
    pub coding_start: u64,
    pub coding_end: u64,
    pub num_exons: usize,
    pub exon_starts: String,
    pub exon_ends: String,
}

impl RefFlatRecord {
    pub fn new(row: RefFlatRow) -> RefFlatRecord {
        RefFlatRecord {
            gene_id: row.0,
            transcript_name: row.1,
            seq_name: row.2,
            strand_char: row.3,
            trx_start: row.4,
            trx_end: row.5,
            coding_start: row.6,
            coding_end: row.7,
            num_exons: row.8,
            exon_starts: row.9,
            exon_ends: row.10,
        }
    }

    pub fn to_transcript(self) -> Result<Transcript, Error> {
        self.to_tupled_transcript().map(|(_, trx)| trx)
    }

    fn to_tupled_transcript(self) -> Result<(String, Transcript), Error> {

        let exon_coords = self.zip_raw_exon_coords();
        if exon_coords.len() != self.num_exons {
            return Err(Error::RefFlat(
                "number of exon and exon coordinates mismatch"));
        }

        let strand = Strand::from_char(&self.strand_char)
            .map_err(FeatureError::from)
            .map_err(Error::from)?;

        let coding_coord =
            if self.coding_start != self.coding_end {
                Some((self.coding_start, self.coding_end))
            } else {
                None
            };

        let gene_id = self.gene_id;
        TBuilder::new(self.seq_name, self.trx_start, self.trx_end)
            .id(self.transcript_name)
            .strand(strand)
            .coords(exon_coords, coding_coord)
            .coding_incl_stop(true)
            .build()
            .map_err(Error::from)
            .map(|trx| (gene_id, trx))
    }

    #[inline]
    fn zip_raw_exon_coords(&self) -> Vec<Coord<u64>> {
        let exon_starts = self.exon_starts
            .split(',').filter_map(|item| u64::from_str(item).ok());
        let exon_ends = self.exon_ends
            .split(',').filter_map(|item| u64::from_str(item).ok());
        exon_starts.zip(exon_ends).collect()
    }
}

pub struct Reader<R: io::Read> {
    inner: csv::Reader<R>,
}

impl<R: io::Read> Reader<R> {

    pub fn from_reader(in_reader: R) -> Reader<R> {
        Reader {
            inner: csv::Reader::from_reader(in_reader)
                .delimiter(b'\t')
                .has_headers(false)
        }
    }

    pub fn records<'a>(&'a mut self) -> RefFlatRecords<'a, R> {
        RefFlatRecords {
            inner: self.inner.decode()
        }
    }

    pub fn transcripts<'a>(&'a mut self) -> RefFlatTranscripts<'a, R> {
        RefFlatTranscripts {
            inner: self.records()
        }
    }
}

impl Reader<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::open(path).map(Reader::from_reader)
    }
}

pub struct RefFlatRecords<'a, R: 'a> where R: io::Read {
    inner: csv::DecodedRecords<'a, R, RefFlatRow>,
}

impl<'a, R> Iterator for RefFlatRecords<'a, R> where R: io::Read {

    type Item = Result<RefFlatRecord, Error>;

    fn next(&mut self) -> Option<Result<RefFlatRecord, Error>> {
        self.inner.next()
            .map(|row| {
                row.or_else(|err| Err(Error::from(err)))
                    .map(RefFlatRecord::new)
            })
    }
}

pub struct RefFlatTranscripts<'a, R: 'a> where R: io::Read {
    inner: RefFlatRecords<'a, R>,
}

impl<'a, R> Iterator for RefFlatTranscripts<'a, R> where R: io::Read {

    type Item = Result<Transcript, Error>;

    fn next(&mut self) -> Option<Result<Transcript, Error>> {
        self.inner.next()
            .map(|record| record.and_then(|rec| rec.to_transcript()))
    }
}
<|endoftext|>


<commit_before>#[doc = "Rust TNetStrings serialization library."];

export tnetstring;
export str;
export int;
export float;
export bool;
export null;
export map;
export vec;
export to_writer;
export to_bytes;
export to_str;
export from_reader;
export from_bytes;
export from_str;

#[doc = "Represents a tnetstring value."]
enum tnetstring {
    #[doc = "str"]
    str(@[u8]),
    #[doc = "int"]
    int(int),
    #[doc = "floating"]
    float(float),
    #[doc = "boolean"]
    bool(bool),
    #[doc = "null"]
    null,
    #[doc = "map"]
    map(map::hashmap<[u8], tnetstring>),
    #[doc = "list"]
    vec(@[tnetstring]),
}

#[doc = "Serializes a tnetstring value into a io::writer."]
fn to_writer(writer: io::writer, tnetstring: tnetstring) {
    fn write_str(wr: io::writer, s: [u8]) {
        wr.write_str(#fmt("%u:", s.len()));
        wr.write(s);
        wr.write_char(',');
    }

    alt tnetstring {
        str(s) { write_str(writer, *s) }
        int(i) {
            let s = int::str(i);
            writer.write_str(#fmt("%u:%s#", s.len(), s));
        }
        float(f) {
            let s = float::to_str(f, 6u);
            writer.write_str(#fmt("%u:%s^", s.len(), s));
        }
        bool(b) {
            let s = bool::to_str(b);
            writer.write_str(#fmt("%u:%s!", s.len(), s));
        }
        map(m) {
            let buf = io::mem_buffer();
            let wr = io::mem_buffer_writer(buf);
            for m.each { |key, value|
                write_str(wr, key);
                to_writer(wr, value);
            }
            let payload = io::mem_buffer_buf(buf);
            writer.write_str(#fmt("%u:", payload.len()));
            writer.write(payload);
            writer.write_char('}');
        }
        vec(v) {
            let buf = io::mem_buffer();
            let wr = io::mem_buffer_writer(buf);
            for (*v).each {|e| to_writer(wr, e) }
            let payload = io::mem_buffer_buf(buf);
            writer.write_str(#fmt("%u:", payload.len()));
            writer.write(payload);
            writer.write_char(']');
        }
        null {
            writer.write_str("0:~");
        }
    }
}

#[doc = "Serializes a tnetstring value into a byte string."]
fn to_bytes(tnetstring: tnetstring) -> [u8] {
    let buf = io::mem_buffer();
    let wr = io::mem_buffer_writer(buf);
    to_writer(wr, tnetstring);
    io::mem_buffer_buf(buf)
}

#[doc = "Serializes a tnetstring value into a string."]
fn to_str(tnetstring: tnetstring) -> str {
    let buf = io::mem_buffer();
    let wr = io::mem_buffer_writer(buf);
    to_writer(wr, tnetstring);
    io::mem_buffer_str(buf)
}

#[doc = "Deserializes a tnetstring value from an io::reader."]
fn from_reader(reader: io::reader) -> option<tnetstring> {
    assert !reader.eof();

    let mut c = reader.read_byte();
    let mut len = 0u;

    // Note that netstring spec explicitly forbids padding zeros.
    // If the first char is zero, it must be the only char.
    if c < '0' as int || c > '9' as int {
        fail "Not a tnetstring: invalid or missing length prefix";
    } else if c == '0' as int {
        c = reader.read_byte();
    } else {
        loop {
            len = (10u * len) + ((c as uint) - ('0' as uint));

            if reader.eof() {
                fail "Not a tnetstring: invalid or missing length prefix";
            }
            c = reader.read_byte();

            if c < '0' as int || c > '9' as int {
                break;
            }
        }
    }

    // Validate end-of-length-prefix marker.
    if c != ':' as int {
        fail "Not a tnetstring: missing length prefix";
    }

    // Read the data plus terminating type tag.
    let payload = reader.read_bytes(len);

    if payload.len() != len {
        fail "Not a tnetstring: invalid length prefix";
    }

    if reader.eof() {
        fail "Not a tnetstring: missing type tag";
    }

    alt reader.read_byte() as char {
      '#' {
        let s = unsafe { str::unsafe::from_bytes(payload) };
        option::chain(int::from_str(s)) {|v| some(int(v)) }
      }
      '}' { some(map(parse_map(payload))) }
      ']' { some(vec(parse_vec(payload))) }
      '!' {
        let s = unsafe { str::unsafe::from_bytes(payload) };
        option::chain(bool::from_str(s)) {|v| some(bool(v)) }
      }
      '^' {
        let s = unsafe { str::unsafe::from_bytes(payload) };
        option::chain(float::from_str(s)) {|v| some(float(v)) }
      }
      '~' {
        assert payload.len() == 0u;
        some(null)
      }
      ',' { some(str(@payload)) }
      c {
        let s = str::from_char(c);
        fail #fmt("Invalid payload type: %?", s)
      }
    }
}

fn parse_vec(data: [u8]) -> @[tnetstring] {
    if data.len() == 0u { ret @[]; }

    let reader = io::bytes_reader(data);

    let result = dvec();

    let value = from_reader(reader);
    assert option::is_some(value);

    result.push(option::get(value));

    while !reader.eof() {
        let value = from_reader(reader);
        assert option::is_some(value);
        result.push(option::get(value));
    }

    @vec::from_mut(dvec::unwrap(result))
}

fn parse_pair(reader: io::reader) -> (@[u8], tnetstring) {
    let key = from_reader(reader);
    assert option::is_some(key);
    assert !reader.eof();

    alt key {
      some(str(key)) {
        let value = from_reader(reader);
        assert option::is_some(value);

        (key, option::get(value))
      }
      _ { fail "Keys can only be strings." }
    }
}

fn parse_map(data: [u8]) -> map::hashmap<[u8], tnetstring> {
    if data.len() == 0u { ret map::bytes_hash(); }

    let reader = io::bytes_reader(data);
    let result = map::bytes_hash();

    let (key, value) = parse_pair(reader);
    result.insert(copy *key, value);

    while !reader.eof() {
        let (key, value) = parse_pair(reader);
        result.insert(copy *key, value);
    }

    ret result;
}

#[doc = "Deserializes a tnetstring value from a byte string."]
fn from_bytes(data: [u8]) -> (option<tnetstring>, [u8]) {
    let rdr = io::bytes_reader(data);
    let tnetstring = from_reader(rdr);
    (tnetstring, rdr.read_whole_stream())
}

#[doc = "Deserializes a tnetstring value from a string."]
fn from_str(data: str) -> (option<tnetstring>, str) {
    io::with_str_reader(data) {|rdr|
        let tnetstring = from_reader(rdr);
        let bytes = rdr.read_whole_stream();
        (tnetstring, str::from_bytes(bytes))
    }

}

#[doc = "Test the equality between two tnetstring values"]
fn eq(t0: tnetstring, t1: tnetstring) -> bool {
    alt (t0, t1) {
        (str(s0), str(s1)) { s0 == s1 }
        (int(i0), int(i1)) { i0 == i1 }
        (float(f0), float(f1)) { f0 == f1 }
        (bool(b0), bool(b1)) { b0 == b1 }
        (null, null) { true }
        (map(d0), map(d1)) {
          if d0.size() == d1.size() {
              for d0.each() { |k, v|
                  if !d1.contains_key(k) || !eq(d1.get(k), v) {
                      ret false;
                  }
              }
              true
          } else {
              false
          }
        }
        (vec(v0), vec(v1)) { vec::all2(*v0, *v1, eq) }
        _ { false }
    }
}

#[cfg(test)]
mod tests {
    // Tests inspired by https://github.com/rfk/tnetstring.

    fn test(string: str, expected: tnetstring) {
        let (actual, rest) = from_str(string);
        assert option::is_some(actual);
        assert rest == "";

        let actual = option::get(actual);
        assert eq(actual, expected);
        assert to_str(expected) == string;
    }

    #[test]
    fn test_format() {
        test("11:hello world,", str(@str::bytes("hello world")));
        test("0:}", map(map::bytes_hash()));
        test("0:]", vec(@[]));

        let d = map::bytes_hash();
        d.insert(str::bytes("hello"),
                vec(@[
                    int(12345678901),
                    str(@str::bytes("this")),
                    bool(true),
                    null,
                    str(@str::bytes("\x00\x00\x00\x00"))]));

        test("51:5:hello,39:11:12345678901#4:this,4:true!0:~4:\x00\x00\x00" +
            "\x00,]}", map(d));

        test("5:12345#", int(12345));
        test("12:this is cool,", str(@str::bytes("this is cool")));
        test("0:,", str(@str::bytes("")));
        test("0:~", null);
        test("4:true!", bool(true));
        test("5:false!", bool(false));
        test("10:\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00,",
            str(@str::bytes("\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")));
        test("24:5:12345#5:67890#5:xxxxx,]",
            vec(@[
                int(12345),
                int(67890),
                str(@str::bytes("xxxxx"))]));
        test("18:3:0.1^3:0.2^3:0.4^]",
           vec(@[float(0.1), float(0.2), float(0.4)]));
        test("243:238:233:228:223:218:213:208:203:198:193:188:183:178:173:" +
            "168:163:158:153:148:143:138:133:128:123:118:113:108:103:99:95:" +
            "91:87:83:79:75:71:67:63:59:55:51:47:43:39:35:31:27:23:19:15:" +
            "11:hello-there,]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]" +
            "]]]]",
            vec(
                @[vec(@[vec(@[vec(@[vec(@[vec(@[vec(@[vec(@[vec(
                @[vec(@[vec(@[vec(@[vec(@[vec(@[vec(@[vec(@[vec(
                @[vec(@[vec(@[vec(@[vec(@[vec(@[vec(@[vec(@[vec(
                @[vec(@[vec(@[vec(@[vec(@[vec(@[vec(@[vec(@[vec(
                @[vec(@[vec(@[vec(@[vec(@[vec(@[vec(@[vec(@[vec(
                @[vec(@[vec(@[vec(@[vec(@[vec(@[vec(@[vec(@[vec(
                @[vec(@[vec(@[
                    str(@str::bytes("hello-there"))
                ])])])])])])])])])])])])])])])])])])])])])])])])])])])])
                ])])])])])])])])])])])])])])])])])])])])])])]));
    }

    #[test]
    fn test_random() {
        fn randint(rng: rand::rng, a: u32, b: u32) -> u32 {
            (rng.next() % (b - a + 1u32)) + a
        }

        fn get_random_object(rng: rand::rng, depth: u32) -> tnetstring {
            if randint(rng, depth, 10u32) <= 4u32 {
                if randint(rng, 0u32, 1u32) == 0u32 {
                    let n = randint(rng, 0u32, 10u32);
                    vec(@vec::from_fn(n as uint) { |_i|
                        get_random_object(rng, depth + 1u32)
                    })
                } else {
                    let d = map::bytes_hash();

                    let mut i = randint(rng, 0u32, 10u32);
                    while i != 0u32 {
                        let s = rng.gen_bytes(randint(rng, 0u32, 100u32) as uint);
                        d.insert(s, get_random_object(rng, depth + 1u32));
                        i -= 1u32;
                    }
                    map(d)
                }
            } else {
                alt randint(rng, 0u32, 5u32) {
                  0u32 { null }
                  1u32 { bool(true) }
                  2u32 { bool(false) }
                  3u32 {
                    if randint(rng, 0u32, 1u32) == 0u32 {
                        int(rng.next() as int)
                    } else {
                        int(-rng.next() as int)
                    }
                  }
                  4u32 {
                    let mut f = rng.gen_float();

                    // Generate a float that can be exactly converted to
                    // and from a string.
                    loop {
                        alt float::from_str(float::to_str(f, 6u)) {
                          some(f1) {
                            if f == f1 { break; }
                            f = f1;
                          }
                          none { fail }
                        }
                    }

                    if randint(rng, 0u32, 1u32) == 0u32 {
                        float(f)
                    } else {
                        float(-f)
                    }
                  }
                  5u32 {
                    str(@rng.gen_bytes(randint(rng, 0u32, 100u32) as uint))
                  }
                  _ { fail }
                }
            }
        }

        let rng = rand::rng();

        let mut i = 500u;
        while i != 0u {
            let v0 = get_random_object(rng, 0u32);

            alt from_bytes(to_bytes(v0)) {
                (some(v1), rest) if rest == [] {
                    assert eq(v0, v1);
                }
                _ { fail; }
            }
            i -= 1u;
        }
    }
}
<commit_msg>Update to latest rustc syntax<commit_after>#[doc = "Rust TNetStrings serialization library."];

export tnetstring;
export str;
export int;
export float;
export bool;
export null;
export map;
export vec;
export to_writer;
export to_bytes;
export to_str;
export from_reader;
export from_bytes;
export from_str;

#[doc = "Represents a tnetstring value."]
enum tnetstring {
    #[doc = "str"]
    str(@~[u8]),
    #[doc = "int"]
    int(int),
    #[doc = "floating"]
    float(float),
    #[doc = "boolean"]
    bool(bool),
    #[doc = "null"]
    null,
    #[doc = "map"]
    map(map::hashmap<~[u8], tnetstring>),
    #[doc = "list"]
    vec(@~[tnetstring]),
}

#[doc = "Serializes a tnetstring value into a io::writer."]
fn to_writer(writer: io::writer, tnetstring: tnetstring) {
    fn write_str(wr: io::writer, s: &[u8]) {
        wr.write_str(#fmt("%u:", s.len()));
        wr.write(s);
        wr.write_char(',');
    }

    alt tnetstring {
        str(s) { write_str(writer, *s) }
        int(i) {
            let s = int::str(i);
            writer.write_str(#fmt("%u:%s#", s.len(), s));
        }
        float(f) {
            let s = float::to_str(f, 6u);
            writer.write_str(#fmt("%u:%s^", s.len(), s));
        }
        bool(b) {
            let s = bool::to_str(b);
            writer.write_str(#fmt("%u:%s!", s.len(), s));
        }
        map(m) {
            let buf = io::mem_buffer();
            let wr = io::mem_buffer_writer(buf);
            for m.each |key, value| {
                write_str(wr, key);
                to_writer(wr, value);
            }
            let payload = io::mem_buffer_buf(buf);
            writer.write_str(#fmt("%u:", payload.len()));
            writer.write(payload);
            writer.write_char('}');
        }
        vec(v) {
            let buf = io::mem_buffer();
            let wr = io::mem_buffer_writer(buf);
            for v.each |e| { to_writer(wr, e) }
            let payload = io::mem_buffer_buf(buf);
            writer.write_str(#fmt("%u:", payload.len()));
            writer.write(payload);
            writer.write_char(']');
        }
        null {
            writer.write_str("0:~");
        }
    }
}

#[doc = "Serializes a tnetstring value into a byte string."]
fn to_bytes(tnetstring: tnetstring) -> ~[u8] {
    let buf = io::mem_buffer();
    let wr = io::mem_buffer_writer(buf);
    to_writer(wr, tnetstring);
    io::mem_buffer_buf(buf)
}

#[doc = "Serializes a tnetstring value into a string."]
fn to_str(tnetstring: tnetstring) -> str {
    let buf = io::mem_buffer();
    let wr = io::mem_buffer_writer(buf);
    to_writer(wr, tnetstring);
    io::mem_buffer_str(buf)
}

#[doc = "Deserializes a tnetstring value from an io::reader."]
fn from_reader(reader: io::reader) -> option<tnetstring> {
    assert !reader.eof();

    let mut c = reader.read_byte();
    let mut len = 0u;

    // Note that netstring spec explicitly forbids padding zeros.
    // If the first char is zero, it must be the only char.
    if c < '0' as int || c > '9' as int {
        fail "Not a tnetstring: invalid or missing length prefix";
    } else if c == '0' as int {
        c = reader.read_byte();
    } else {
        loop {
            len = (10u * len) + ((c as uint) - ('0' as uint));

            if reader.eof() {
                fail "Not a tnetstring: invalid or missing length prefix";
            }
            c = reader.read_byte();

            if c < '0' as int || c > '9' as int {
                break;
            }
        }
    }

    // Validate end-of-length-prefix marker.
    if c != ':' as int {
        fail "Not a tnetstring: missing length prefix";
    }

    // Read the data plus terminating type tag.
    let payload = reader.read_bytes(len);

    if payload.len() != len {
        fail "Not a tnetstring: invalid length prefix";
    }

    if reader.eof() {
        fail "Not a tnetstring: missing type tag";
    }

    alt reader.read_byte() as char {
      '#' {
        let s = unsafe { str::unsafe::from_bytes(payload) };
        int::from_str(s).map(|v| int(v))
      }
      '}' { some(map(parse_map(payload))) }
      ']' { some(vec(parse_vec(payload))) }
      '!' {
        let s = unsafe { str::unsafe::from_bytes(payload) };
        bool::from_str(s).map(|v| bool(v))
      }
      '^' {
        let s = unsafe { str::unsafe::from_bytes(payload) };
        float::from_str(s).map(|v| float(v))
      }
      '~' {
        assert payload.len() == 0u;
        some(null)
      }
      ',' { some(str(@payload)) }
      c {
        let s = str::from_char(c);
        fail #fmt("Invalid payload type: %?", s)
      }
    }
}

fn parse_vec(data: ~[u8]) -> @~[tnetstring] {
    if data.len() == 0u { ret @~[]; }

    let reader = io::bytes_reader(data);

    let result = dvec();

    let value = from_reader(reader);
    assert value.is_some();

    result.push(value.get());

    while !reader.eof() {
        let value = from_reader(reader);
        assert value.is_some();
        result.push(value.get());
    }

    @vec::from_mut(dvec::unwrap(result))
}

fn parse_pair(reader: io::reader) -> (@~[u8], tnetstring) {
    let key = from_reader(reader);
    assert key.is_some();
    assert !reader.eof();

    alt key {
      some(str(key)) {
        let value = from_reader(reader);
        assert value.is_some();

        (key, value.get())
      }
      _ { fail "Keys can only be strings." }
    }
}

fn parse_map(data: ~[u8]) -> map::hashmap<~[u8], tnetstring> {
    if data.len() == 0u { ret map::bytes_hash(); }

    let reader = io::bytes_reader(data);
    let result = map::bytes_hash();

    let (key, value) = parse_pair(reader);
    result.insert(copy *key, value);

    while !reader.eof() {
        let (key, value) = parse_pair(reader);
        result.insert(copy *key, value);
    }

    ret result;
}

#[doc = "Deserializes a tnetstring value from a byte string."]
fn from_bytes(data: ~[u8]) -> (option<tnetstring>, ~[u8]) {
    let rdr = io::bytes_reader(data);
    let tnetstring = from_reader(rdr);
    (tnetstring, rdr.read_whole_stream())
}

#[doc = "Deserializes a tnetstring value from a string."]
fn from_str(data: str) -> (option<tnetstring>, str) {
    do io::with_str_reader(data) |rdr| {
        let tnetstring = from_reader(rdr);
        let bytes = rdr.read_whole_stream();
        (tnetstring, str::from_bytes(bytes))
    }

}

#[doc = "Test the equality between two tnetstring values"]
fn eq(t0: tnetstring, t1: tnetstring) -> bool {
    alt (t0, t1) {
        (str(s0), str(s1)) { s0 == s1 }
        (int(i0), int(i1)) { i0 == i1 }
        (float(f0), float(f1)) { f0 == f1 }
        (bool(b0), bool(b1)) { b0 == b1 }
        (null, null) { true }
        (map(d0), map(d1)) {
          if d0.size() == d1.size() {
              for d0.each() |k, v| {
                  if !d1.contains_key(k) || !eq(d1.get(k), v) {
                      ret false;
                  }
              }
              true
          } else {
              false
          }
        }
        (vec(v0), vec(v1)) { vec::all2(*v0, *v1, eq) }
        _ { false }
    }
}

#[cfg(test)]
mod tests {
    // Tests inspired by https://github.com/rfk/tnetstring.

    fn test(string: str, expected: tnetstring) {
        let (actual, rest) = from_str(string);
        assert actual.is_some();
        assert rest == "";

        let actual = actual.get();
        assert eq(actual, expected);
        assert to_str(expected) == string;
    }

    #[test]
    fn test_format() {
        test("11:hello world,", str(@str::bytes("hello world")));
        test("0:}", map(map::bytes_hash()));
        test("0:]", vec(@~[]));

        let d = map::bytes_hash();
        d.insert(str::bytes("hello"),
                vec(@~[
                    int(12345678901),
                    str(@str::bytes("this")),
                    bool(true),
                    null,
                    str(@str::bytes("\x00\x00\x00\x00"))]));

        test("51:5:hello,39:11:12345678901#4:this,4:true!0:~4:\x00\x00\x00" +
            "\x00,]}", map(d));

        test("5:12345#", int(12345));
        test("12:this is cool,", str(@str::bytes("this is cool")));
        test("0:,", str(@str::bytes("")));
        test("0:~", null);
        test("4:true!", bool(true));
        test("5:false!", bool(false));
        test("10:\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00,",
            str(@str::bytes("\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")));
        test("24:5:12345#5:67890#5:xxxxx,]",
            vec(@~[
                int(12345),
                int(67890),
                str(@str::bytes("xxxxx"))]));
        test("18:3:0.1^3:0.2^3:0.4^]",
           vec(@~[float(0.1), float(0.2), float(0.4)]));
        test("243:238:233:228:223:218:213:208:203:198:193:188:183:178:173:" +
            "168:163:158:153:148:143:138:133:128:123:118:113:108:103:99:95:" +
            "91:87:83:79:75:71:67:63:59:55:51:47:43:39:35:31:27:23:19:15:" +
            "11:hello-there,]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]" +
            "]]]]",
            vec(
                @~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(
                @~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(
                @~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(
                @~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(
                @~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(
                @~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(@~[vec(
                @~[vec(@~[vec(@~[
                    str(@str::bytes("hello-there"))
                ])])])])])])])])])])])])])])])])])])])])])])])])])])])])
                ])])])])])])])])])])])])])])])])])])])])])])]));
    }

    #[test]
    fn test_random() {
        fn randint(rng: rand::rng, a: u32, b: u32) -> u32 {
            (rng.next() % (b - a + 1u32)) + a
        }

        fn get_random_object(rng: rand::rng, depth: u32) -> tnetstring {
            if randint(rng, depth, 10u32) <= 4u32 {
                if randint(rng, 0u32, 1u32) == 0u32 {
                    let n = randint(rng, 0u32, 10u32);
                    vec(@vec::from_fn(n as uint, |_i|
                        get_random_object(rng, depth + 1u32)
                    ))
                } else {
                    let d = map::bytes_hash();

                    let mut i = randint(rng, 0u32, 10u32);
                    while i != 0u32 {
                        let s = rng.gen_bytes(randint(rng, 0u32, 100u32) as uint);
                        d.insert(s, get_random_object(rng, depth + 1u32));
                        i -= 1u32;
                    }
                    map(d)
                }
            } else {
                alt randint(rng, 0u32, 5u32) {
                  0u32 { null }
                  1u32 { bool(true) }
                  2u32 { bool(false) }
                  3u32 {
                    if randint(rng, 0u32, 1u32) == 0u32 {
                        int(rng.next() as int)
                    } else {
                        int(-rng.next() as int)
                    }
                  }
                  4u32 {
                    let mut f = rng.gen_float();

                    // Generate a float that can be exactly converted to
                    // and from a string.
                    loop {
                        alt float::from_str(float::to_str(f, 6u)) {
                          some(f1) {
                            if f == f1 { break; }
                            f = f1;
                          }
                          none { fail }
                        }
                    }

                    if randint(rng, 0u32, 1u32) == 0u32 {
                        float(f)
                    } else {
                        float(-f)
                    }
                  }
                  5u32 {
                    str(@rng.gen_bytes(randint(rng, 0u32, 100u32) as uint))
                  }
                  _ { fail }
                }
            }
        }

        let rng = rand::rng();

        let mut i = 500u;
        while i != 0u {
            let v0 = get_random_object(rng, 0u32);

            alt from_bytes(to_bytes(v0)) {
                (some(v1), rest) if rest == ~[] {
                    assert eq(v0, v1);
                }
                _ { fail; }
            }
            i -= 1u;
        }
    }
}
<|endoftext|>


<commit_before>use std::io::{BufReader, Cursor, Read};

#[cfg(feature = "enable_logging")]
use log::{debug, info};

fn round_trip(x: &[u8]) {
    let mut compressed: Vec<u8> = Vec::new();
    lzma_rs::xz_compress(&mut std::io::BufReader::new(x), &mut compressed).unwrap();
    #[cfg(feature = "enable_logging")]
    info!("Compressed {} -> {} bytes", x.len(), compressed.len());
    #[cfg(feature = "enable_logging")]
    debug!("Compressed content: {:?}", compressed);
    let mut bf = BufReader::new(compressed.as_slice());
    let mut decomp: Vec<u8> = Vec::new();
    lzma_rs::xz_decompress(&mut bf, &mut decomp).unwrap();
    assert_eq!(decomp, x)
}

fn round_trip_file(filename: &str) {
    let mut x = Vec::new();
    std::fs::File::open(filename)
        .unwrap()
        .read_to_end(&mut x)
        .unwrap();
    round_trip(x.as_slice());
}

#[test]
fn round_trip_basics() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();
    round_trip(b"");
    // Note: we use vec! to avoid storing the slice in the binary
    round_trip(vec![0x00; 1_000_000].as_slice());
    round_trip(vec![0xFF; 1_000_000].as_slice());
}

#[test]
fn round_trip_hello() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();
    round_trip(b"Hello world");
}

#[test]
fn round_trip_files() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();
    round_trip_file("tests/files/foo.txt");
}

fn decomp_big_file(compfile: &str, plainfile: &str) {
    let mut expected = Vec::new();
    std::fs::File::open(plainfile)
        .unwrap()
        .read_to_end(&mut expected)
        .unwrap();
    let mut f = BufReader::new(std::fs::File::open(compfile).unwrap());
    let mut decomp: Vec<u8> = Vec::new();
    lzma_rs::xz_decompress(&mut f, &mut decomp).unwrap();
    assert!(decomp == expected)
}

#[test]
fn big_file() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();
    decomp_big_file("tests/files/foo.txt.xz", "tests/files/foo.txt");
    decomp_big_file("tests/files/good-1-lzma2-1.xz", "tests/files/good-1-lzma2-1");
    decomp_big_file("tests/files/good-1-lzma2-2.xz", "tests/files/good-1-lzma2-2");
    decomp_big_file("tests/files/good-1-lzma2-3.xz", "tests/files/good-1-lzma2-3");
    decomp_big_file("tests/files/good-1-lzma2-4.xz", "tests/files/good-1-lzma2-4");
}

#[test]
fn decompress_empty_world() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();
    let mut x: &[u8] = b"\xfd\x37\x7a\x58\x5a\x00\x00\x04\xe6\xd6\xb4\x46\x00\x00\x00\x00\
                         \x1c\xdf\x44\x21\x1f\xb6\xf3\x7d\x01\x00\x00\x00\x00\x04\x59\x5a\
                         ";
    let mut decomp: Vec<u8> = Vec::new();
    lzma_rs::xz_decompress(&mut x, &mut decomp).unwrap();
    assert_eq!(decomp, b"")
}

#[test]
fn decompress_hello_world() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();
    let mut x: &[u8] = b"\xfd\x37\x7a\x58\x5a\x00\x00\x04\xe6\xd6\xb4\x46\x02\x00\x21\x01\
                         \x16\x00\x00\x00\x74\x2f\xe5\xa3\x01\x00\x0b\x48\x65\x6c\x6c\x6f\
                         \x20\x77\x6f\x72\x6c\x64\x0a\x00\xca\xec\x49\x05\x66\x3f\x67\x98\
                         \x00\x01\x24\x0c\xa6\x18\xd8\xd8\x1f\xb6\xf3\x7d\x01\x00\x00\x00\
                         \x00\x04\x59\x5a";
    let mut decomp: Vec<u8> = Vec::new();
    lzma_rs::xz_decompress(&mut x, &mut decomp).unwrap();
    assert_eq!(decomp, b"Hello world\x0a")
}

#[test]
fn test_xz_block_check_crc32() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();

    decomp_big_file(
        "tests/files/block-check-crc32.txt.xz",
        "tests/files/block-check-crc32.txt",
    );
}

#[test]
fn test_xz_block_check_crc32_invalid() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();

    let testcase = "tests/files/block-check-crc32.txt.xz";
    let mut corrupted = {
        let mut buf = vec![];
        std::fs::File::open(testcase)
            .unwrap()
            .read_to_end(&mut buf)
            .unwrap();
        // Mangle the "Block Check" field.
        buf[0x54] = 0x67;
        buf[0x55] = 0x45;
        buf[0x56] = 0x23;
        buf[0x57] = 0x01;
        BufReader::new(Cursor::new(buf))
    };
    let mut decomp = vec![];

    let err_msg = lzma_rs::xz_decompress(&mut corrupted, &mut decomp)
        .unwrap_err()
        .to_string();
    assert_eq!(
        err_msg,
        "xz error: Invalid footer CRC32: expected 0x01234567 but got 0x8b0d303e"
    )
}
<commit_msg>Format according to cargo's rules<commit_after>use std::io::{BufReader, Cursor, Read};

#[cfg(feature = "enable_logging")]
use log::{debug, info};

fn round_trip(x: &[u8]) {
    let mut compressed: Vec<u8> = Vec::new();
    lzma_rs::xz_compress(&mut std::io::BufReader::new(x), &mut compressed).unwrap();
    #[cfg(feature = "enable_logging")]
    info!("Compressed {} -> {} bytes", x.len(), compressed.len());
    #[cfg(feature = "enable_logging")]
    debug!("Compressed content: {:?}", compressed);
    let mut bf = BufReader::new(compressed.as_slice());
    let mut decomp: Vec<u8> = Vec::new();
    lzma_rs::xz_decompress(&mut bf, &mut decomp).unwrap();
    assert_eq!(decomp, x)
}

fn round_trip_file(filename: &str) {
    let mut x = Vec::new();
    std::fs::File::open(filename)
        .unwrap()
        .read_to_end(&mut x)
        .unwrap();
    round_trip(x.as_slice());
}

#[test]
fn round_trip_basics() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();
    round_trip(b"");
    // Note: we use vec! to avoid storing the slice in the binary
    round_trip(vec![0x00; 1_000_000].as_slice());
    round_trip(vec![0xFF; 1_000_000].as_slice());
}

#[test]
fn round_trip_hello() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();
    round_trip(b"Hello world");
}

#[test]
fn round_trip_files() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();
    round_trip_file("tests/files/foo.txt");
}

fn decomp_big_file(compfile: &str, plainfile: &str) {
    let mut expected = Vec::new();
    std::fs::File::open(plainfile)
        .unwrap()
        .read_to_end(&mut expected)
        .unwrap();
    let mut f = BufReader::new(std::fs::File::open(compfile).unwrap());
    let mut decomp: Vec<u8> = Vec::new();
    lzma_rs::xz_decompress(&mut f, &mut decomp).unwrap();
    assert!(decomp == expected)
}

#[test]
fn big_file() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();
    decomp_big_file("tests/files/foo.txt.xz", "tests/files/foo.txt");
    decomp_big_file(
        "tests/files/good-1-lzma2-1.xz",
        "tests/files/good-1-lzma2-1",
    );
    decomp_big_file(
        "tests/files/good-1-lzma2-2.xz",
        "tests/files/good-1-lzma2-2",
    );
    decomp_big_file(
        "tests/files/good-1-lzma2-3.xz",
        "tests/files/good-1-lzma2-3",
    );
    decomp_big_file(
        "tests/files/good-1-lzma2-4.xz",
        "tests/files/good-1-lzma2-4",
    );
}

#[test]
fn decompress_empty_world() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();
    let mut x: &[u8] = b"\xfd\x37\x7a\x58\x5a\x00\x00\x04\xe6\xd6\xb4\x46\x00\x00\x00\x00\
                         \x1c\xdf\x44\x21\x1f\xb6\xf3\x7d\x01\x00\x00\x00\x00\x04\x59\x5a\
                         ";
    let mut decomp: Vec<u8> = Vec::new();
    lzma_rs::xz_decompress(&mut x, &mut decomp).unwrap();
    assert_eq!(decomp, b"")
}

#[test]
fn decompress_hello_world() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();
    let mut x: &[u8] = b"\xfd\x37\x7a\x58\x5a\x00\x00\x04\xe6\xd6\xb4\x46\x02\x00\x21\x01\
                         \x16\x00\x00\x00\x74\x2f\xe5\xa3\x01\x00\x0b\x48\x65\x6c\x6c\x6f\
                         \x20\x77\x6f\x72\x6c\x64\x0a\x00\xca\xec\x49\x05\x66\x3f\x67\x98\
                         \x00\x01\x24\x0c\xa6\x18\xd8\xd8\x1f\xb6\xf3\x7d\x01\x00\x00\x00\
                         \x00\x04\x59\x5a";
    let mut decomp: Vec<u8> = Vec::new();
    lzma_rs::xz_decompress(&mut x, &mut decomp).unwrap();
    assert_eq!(decomp, b"Hello world\x0a")
}

#[test]
fn test_xz_block_check_crc32() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();

    decomp_big_file(
        "tests/files/block-check-crc32.txt.xz",
        "tests/files/block-check-crc32.txt",
    );
}

#[test]
fn test_xz_block_check_crc32_invalid() {
    #[cfg(feature = "enable_logging")]
    let _ = env_logger::try_init();

    let testcase = "tests/files/block-check-crc32.txt.xz";
    let mut corrupted = {
        let mut buf = vec![];
        std::fs::File::open(testcase)
            .unwrap()
            .read_to_end(&mut buf)
            .unwrap();
        // Mangle the "Block Check" field.
        buf[0x54] = 0x67;
        buf[0x55] = 0x45;
        buf[0x56] = 0x23;
        buf[0x57] = 0x01;
        BufReader::new(Cursor::new(buf))
    };
    let mut decomp = vec![];

    let err_msg = lzma_rs::xz_decompress(&mut corrupted, &mut decomp)
        .unwrap_err()
        .to_string();
    assert_eq!(
        err_msg,
        "xz error: Invalid footer CRC32: expected 0x01234567 but got 0x8b0d303e"
    )
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate coreutils;
extern crate extra;

use std::env;
use std::fs;
use std::path::Path;
use std::io::{stdout, stderr, StdoutLock, Stderr, Write};
use std::process::exit;
use coreutils::{ArgParser, to_human_readable_string};
use extra::option::OptionalExt;

const MAN_PAGE: &'static str = /* @MANSTART{du} */ r#"
NAME
    du - list directory content with sizes

SYNOPSIS
    du [ -h | --help ][FILE]...

DESCRIPTION
    List the name and size of the FILE(s), or the current directory

OPTIONS
    -h
    --human-readable
        human readable output
    --help
        display this help and exit
"#; /* @MANEND */

fn list_entry(path: &str, name: &str, parser: &ArgParser, stdout: &mut StdoutLock, stderr: &mut Stderr) {
    let metadata = fs::metadata(path).try(stderr);
    let size = metadata.len();

    if parser.found(&'h') || parser.found("human-readable") {
        stdout.write(to_human_readable_string(size).as_bytes()).try(stderr);
    } else {
        stdout.write(((size + 1023) / 1024).to_string().as_bytes()).try(stderr);
    }

    stdout.write(b"\t").try(stderr);
    stdout.write(name.as_bytes()).try(stderr);
    stdout.write(b"\n").try(stderr);
}

fn list_dir(path: &str, parser: &ArgParser, stdout: &mut StdoutLock, stderr: &mut Stderr) {
    if fs::metadata(path).try(stderr).is_dir() {
        let read_dir = Path::new(path).read_dir().try(stderr);

        let mut entries = vec![];
        for dir in read_dir {
            let dir = match dir {
                Ok(x) => x,
                Err(_) => continue,
            };
            let mut file_name = dir.file_name().to_string_lossy().into_owned();
            if dir.file_type().try(stderr).is_dir() {
                file_name.push('/');
            }
            entries.push(file_name);
        }

        entries.sort();

        for entry in entries.iter() {
            let mut entry_path = path.to_owned();
            if !entry_path.ends_with('/') {
                entry_path.push('/');
            }
            entry_path.push_str(&entry);

            list_entry(&entry_path, &entry, parser, stdout, stderr);
        }
    } else {
        list_entry(path, path, parser, stdout, stderr);
    }
}

fn main() {
    let stdout = stdout();
    let mut stdout = stdout.lock();
    let mut stderr = stderr();
    let mut parser = ArgParser::new(2)
        .add_flag("h", "human-readable")
        .add_flag("", "help");
    parser.parse(env::args());

    if parser.found("help") {
        stdout.write(MAN_PAGE.as_bytes()).try(&mut stderr);
        stdout.flush().try(&mut stderr);
        exit(0);
    }

    if parser.args.is_empty() {
        list_dir(".", &parser, &mut stdout, &mut stderr);
    } else {
        for dir in parser.args.iter() {
            list_dir(&dir, &parser, &mut stdout, &mut stderr);
        }
    }
}
<commit_msg>du adeheres to new ArgParser interface<commit_after>#![deny(warnings)]

extern crate coreutils;
extern crate extra;

use std::env;
use std::fs;
use std::path::Path;
use std::io::{stdout, stderr, StdoutLock, Stderr, Write};
use std::process::exit;
use coreutils::{ArgParser, to_human_readable_string};
use extra::option::OptionalExt;

const MAN_PAGE: &'static str = /* @MANSTART{du} */ r#"
NAME
    du - list directory content with sizes

SYNOPSIS
    du [ -h | --help ][FILE]...

DESCRIPTION
    List the name and size of the FILE(s), or the current directory

OPTIONS
    -h
    --human-readable
        human readable output
    --help
        display this help and exit
"#; /* @MANEND */

fn list_entry(path: &str, name: &str, parser: &ArgParser, stdout: &mut StdoutLock, stderr: &mut Stderr) {
    let metadata = fs::metadata(path).try(stderr);
    let size = metadata.len();

    if parser.found("human-readable") {
        stdout.write(to_human_readable_string(size).as_bytes()).try(stderr);
    } else {
        stdout.write(((size + 1023) / 1024).to_string().as_bytes()).try(stderr);
    }

    stdout.write(b"\t").try(stderr);
    stdout.write(name.as_bytes()).try(stderr);
    stdout.write(b"\n").try(stderr);
}

fn list_dir(path: &str, parser: &ArgParser, stdout: &mut StdoutLock, stderr: &mut Stderr) {
    if fs::metadata(path).try(stderr).is_dir() {
        let read_dir = Path::new(path).read_dir().try(stderr);

        let mut entries = vec![];
        for dir in read_dir {
            let dir = match dir {
                Ok(x) => x,
                Err(_) => continue,
            };
            let mut file_name = dir.file_name().to_string_lossy().into_owned();
            if dir.file_type().try(stderr).is_dir() {
                file_name.push('/');
            }
            entries.push(file_name);
        }

        entries.sort();

        for entry in entries.iter() {
            let mut entry_path = path.to_owned();
            if !entry_path.ends_with('/') {
                entry_path.push('/');
            }
            entry_path.push_str(&entry);

            list_entry(&entry_path, &entry, parser, stdout, stderr);
        }
    } else {
        list_entry(path, path, parser, stdout, stderr);
    }
}

fn main() {
    let stdout = stdout();
    let mut stdout = stdout.lock();
    let mut stderr = stderr();
    let mut parser = ArgParser::new(2)
        .add_flag(&["h", "human-readable"])
        .add_flag(&["help"]);
    parser.parse(env::args());

    if parser.found("help") {
        stdout.write(MAN_PAGE.as_bytes()).try(&mut stderr);
        stdout.flush().try(&mut stderr);
        exit(0);
    }

    if parser.args.is_empty() {
        list_dir(".", &parser, &mut stdout, &mut stderr);
    } else {
        for dir in parser.args.iter() {
            list_dir(&dir, &parser, &mut stdout, &mut stderr);
        }
    }
}
<|endoftext|>