<commit_before>// Defined note types for GNU systems.

// ABI information.  The descriptor consists of words:
//    word 0: OS descriptor
//    word 1: major version of the ABI
//    word 2: minor version of the ABI
//    word 3: subminor version of the ABI

pub const NT_GNU_ABI_TAG: u32 = 1;
// Old name
pub const ELF_NOTE_ABI: u32 = NT_GNU_ABI_TAG;
// Known OSes.  These values can appear in word 0 of an
//    NT_GNU_ABI_TAG note section entry.
pub const ELF_NOTE_OS_LINUX: u32 = 0;
pub const ELF_NOTE_OS_GNU: u32 = 1;
pub const ELF_NOTE_OS_SOLARIS2: u32 = 2;
pub const ELF_NOTE_OS_FREEBSD: u32 = 3;

// Synthetic hwcap information.  The descriptor begins with two words:
//    word 0: number of entries
//    word 1: bitmask of enabled entries
//    Then follow variable-length entries, one byte followed by a
//    '\0'-terminated hwcap name string.  The byte gives the bit
//    number to test if enabled, (1U << bit) & bitmask.
pub const NT_GNU_HWCAP: u32 = 2;

// Build ID bits as generated by ld --build-id.
//    The descriptor consists of any nonzero number of bytes.
pub const NT_GNU_BUILD_ID: u32 = 3;

// Version note generated by GNU gold containing a version string.
pub const NT_GNU_GOLD_VERSION: u32 = 4;

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "std", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header of a fixed form.
pub struct Nhdr32 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u32,
    /// Length of the note's descriptor
    pub n_descsz: u32,
    /// Type of the note
    pub n_type: u32,
}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "std", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header of a fixed form.
pub struct Nhdr64 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u64,
    /// Length of the note's descriptor.
    pub n_descsz: u64,
    /// Type of the note.
    pub n_type: u64,
}

if_std! {
    use error;
    use container;
    use scroll::{ctx, Pread};

    /// An iterator over ELF binary notes in a note section or segment
    pub struct NoteDataIterator<'a> {
        pub data: &'a [u8],
        pub size: usize,
        pub offset: usize,
        pub ctx: (usize, container::Ctx), // (alignment, ctx)
    }

    impl<'a> Iterator for NoteDataIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            if self.offset >= self.size {
                None
            } else {
                debug!("NoteIterator - {:#x}", self.offset);
                match self.data.gread_with(&mut self.offset, self.ctx) {
                    Ok(res) => Some(Ok(res)),
                    Err(e) => Some(Err(e.into()))
                }
            }
        }
    }

    /// An iterator over ELF binary notes
    pub struct NoteIterator<'a> {
        pub iters: Vec<NoteDataIterator<'a>>,
        pub index: usize,
    }

    impl<'a> Iterator for NoteIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            while self.index < self.iters.len() {
                if let Some(note_result) = self.iters[self.index].next() {
                    return Some(note_result);
                }

                self.index += 1;
            }

            None
        }
    }

    #[derive(Debug)]
    struct NoteHeader {
        n_namesz: usize,
        n_descsz: usize,
        n_type: u32,
    }

    impl From<Nhdr32> for NoteHeader {
        fn from(header: Nhdr32) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type,
            }
        }
    }

    impl From<Nhdr64> for NoteHeader {
        fn from(header: Nhdr64) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type as u32,
            }
        }
    }

    fn align(alignment: usize, offset: &mut usize) {
        let diff = *offset % alignment;
        if diff != 0 {
            *offset += alignment - diff;
        }
    }

    /// A 32/64 bit Note struct, with the name and desc pre-parsed
    #[derive(Debug)]
    pub struct Note<'a> {
        /// The type of this note
        pub n_type: u32,
        /// NUL terminated string, where `namesz` includes the terminator
        pub name: &'a str, // needs padding such that namesz + padding % {wordsize} == 0
        /// arbitrary data of length `descsz`
        pub desc: &'a [u8], // needs padding such that descsz + padding % {wordsize} == 0
    }

    impl<'a> Note<'a> {
        pub fn type_to_str(&self) -> &'static str {
            match self.n_type {
                NT_GNU_ABI_TAG => "NT_GNU_ABI_TAG",
                NT_GNU_HWCAP => "NT_GNU_HWCAP",
                NT_GNU_BUILD_ID => "NT_GNU_BUILD_ID",
                NT_GNU_GOLD_VERSION => "NT_GNU_GOLD_VERSION",
                _ => "NT_UNKNOWN"
            }
        }
    }

    impl<'a> ctx::TryFromCtx<'a, (usize, container::Ctx)> for Note<'a> {
        type Error = error::Error;
        type Size = usize;
        fn try_from_ctx(bytes: &'a [u8], (alignment, ctx): (usize, container::Ctx)) -> Result<(Self, Self::Size), Self::Error> {
            let offset = &mut 0;
            let header: NoteHeader = {
                match alignment {
                    4 => bytes.gread_with::<Nhdr32>(offset, ctx.le)?.into(),
                    // this is a guess; i haven't seen gcc/clang compilers emit 64-bit notes, and i don't have any non gcc/clang compilers
                    8 => bytes.gread_with::<Nhdr64>(offset, ctx.le)?.into(),
                    _ => return Err(error::Error::Malformed(format!("Notes has unimplemented alignment requirement: {:#x}", alignment)))
                }
            };
            debug!("{:?} - {:#x}", header, *offset);
            // -1 because includes \0 terminator
            let name = bytes.gread_with::<&'a str>(offset, ctx::StrCtx::Length(header.n_namesz - 1))?;
            *offset += 1;
            align(alignment, offset);
            debug!("note name {} - {:#x}", name, *offset);
            let desc = bytes.gread_with::<&'a [u8]>(offset, header.n_descsz)?;
            align(alignment, offset);
            debug!("desc {:?} - {:#x}", desc, *offset);
            Ok((Note {
                name,
                desc,
                n_type: header.n_type,
            }, *offset))
        }
    }

    #[cfg(test)]
    mod tests {
        use super::*;

        const NOTE_DATA: [u8; 68] = [0x04, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00,
                                     0x01, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
                                     0x06, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00,
                                     0x04, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00,
                                     0x03, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0xbc, 0xfc, 0x66, 0xcd, 0xc7, 0xd5, 0x14, 0x7b,
                                     0x53, 0xb1, 0x10, 0x11, 0x94, 0x86, 0x8e, 0xf9,
                                     0x4f, 0xe8, 0xdd, 0xdb];

        const CONTEXT: (usize, container::Ctx) = (4, container::Ctx {
            container: container::Container::Big,
            le: ::scroll::Endian::Little,
        });

        fn make_note_iter(start: usize, end: usize) -> NoteDataIterator<'static> {
            NoteDataIterator {
                data: &NOTE_DATA,
                size: end,
                offset: start,
                ctx: CONTEXT,
            }
        }

        #[test]
        fn iter_single_section() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn iter_multiple_sections() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 32), make_note_iter(32, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn skip_empty_sections() {
            let mut notes = NoteIterator {
                iters: vec![
                    make_note_iter(0, 32),
                    make_note_iter(0, 0),
                    make_note_iter(32, 68),
                ],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn ignore_no_sections() {
            let mut notes = NoteIterator { iters: vec![], index: 0 };
            assert!(notes.next().is_none());
        }
    }
}
<commit_msg>Fix tests for rust 1.18.0<commit_after>// Defined note types for GNU systems.

// ABI information.  The descriptor consists of words:
//    word 0: OS descriptor
//    word 1: major version of the ABI
//    word 2: minor version of the ABI
//    word 3: subminor version of the ABI

pub const NT_GNU_ABI_TAG: u32 = 1;
// Old name
pub const ELF_NOTE_ABI: u32 = NT_GNU_ABI_TAG;
// Known OSes.  These values can appear in word 0 of an
//    NT_GNU_ABI_TAG note section entry.
pub const ELF_NOTE_OS_LINUX: u32 = 0;
pub const ELF_NOTE_OS_GNU: u32 = 1;
pub const ELF_NOTE_OS_SOLARIS2: u32 = 2;
pub const ELF_NOTE_OS_FREEBSD: u32 = 3;

// Synthetic hwcap information.  The descriptor begins with two words:
//    word 0: number of entries
//    word 1: bitmask of enabled entries
//    Then follow variable-length entries, one byte followed by a
//    '\0'-terminated hwcap name string.  The byte gives the bit
//    number to test if enabled, (1U << bit) & bitmask.
pub const NT_GNU_HWCAP: u32 = 2;

// Build ID bits as generated by ld --build-id.
//    The descriptor consists of any nonzero number of bytes.
pub const NT_GNU_BUILD_ID: u32 = 3;

// Version note generated by GNU gold containing a version string.
pub const NT_GNU_GOLD_VERSION: u32 = 4;

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "std", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header of a fixed form.
pub struct Nhdr32 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u32,
    /// Length of the note's descriptor
    pub n_descsz: u32,
    /// Type of the note
    pub n_type: u32,
}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "std", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header of a fixed form.
pub struct Nhdr64 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u64,
    /// Length of the note's descriptor.
    pub n_descsz: u64,
    /// Type of the note.
    pub n_type: u64,
}

if_std! {
    use error;
    use container;
    use scroll::{ctx, Pread};

    /// An iterator over ELF binary notes in a note section or segment
    pub struct NoteDataIterator<'a> {
        pub data: &'a [u8],
        pub size: usize,
        pub offset: usize,
        pub ctx: (usize, container::Ctx), // (alignment, ctx)
    }

    impl<'a> Iterator for NoteDataIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            if self.offset >= self.size {
                None
            } else {
                debug!("NoteIterator - {:#x}", self.offset);
                match self.data.gread_with(&mut self.offset, self.ctx) {
                    Ok(res) => Some(Ok(res)),
                    Err(e) => Some(Err(e.into()))
                }
            }
        }
    }

    /// An iterator over ELF binary notes
    pub struct NoteIterator<'a> {
        pub iters: Vec<NoteDataIterator<'a>>,
        pub index: usize,
    }

    impl<'a> Iterator for NoteIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            while self.index < self.iters.len() {
                if let Some(note_result) = self.iters[self.index].next() {
                    return Some(note_result);
                }

                self.index += 1;
            }

            None
        }
    }

    #[derive(Debug)]
    struct NoteHeader {
        n_namesz: usize,
        n_descsz: usize,
        n_type: u32,
    }

    impl From<Nhdr32> for NoteHeader {
        fn from(header: Nhdr32) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type,
            }
        }
    }

    impl From<Nhdr64> for NoteHeader {
        fn from(header: Nhdr64) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type as u32,
            }
        }
    }

    fn align(alignment: usize, offset: &mut usize) {
        let diff = *offset % alignment;
        if diff != 0 {
            *offset += alignment - diff;
        }
    }

    /// A 32/64 bit Note struct, with the name and desc pre-parsed
    #[derive(Debug)]
    pub struct Note<'a> {
        /// The type of this note
        pub n_type: u32,
        /// NUL terminated string, where `namesz` includes the terminator
        pub name: &'a str, // needs padding such that namesz + padding % {wordsize} == 0
        /// arbitrary data of length `descsz`
        pub desc: &'a [u8], // needs padding such that descsz + padding % {wordsize} == 0
    }

    impl<'a> Note<'a> {
        pub fn type_to_str(&self) -> &'static str {
            match self.n_type {
                NT_GNU_ABI_TAG => "NT_GNU_ABI_TAG",
                NT_GNU_HWCAP => "NT_GNU_HWCAP",
                NT_GNU_BUILD_ID => "NT_GNU_BUILD_ID",
                NT_GNU_GOLD_VERSION => "NT_GNU_GOLD_VERSION",
                _ => "NT_UNKNOWN"
            }
        }
    }

    impl<'a> ctx::TryFromCtx<'a, (usize, container::Ctx)> for Note<'a> {
        type Error = error::Error;
        type Size = usize;
        fn try_from_ctx(bytes: &'a [u8], (alignment, ctx): (usize, container::Ctx)) -> Result<(Self, Self::Size), Self::Error> {
            let offset = &mut 0;
            let header: NoteHeader = {
                match alignment {
                    4 => bytes.gread_with::<Nhdr32>(offset, ctx.le)?.into(),
                    // this is a guess; i haven't seen gcc/clang compilers emit 64-bit notes, and i don't have any non gcc/clang compilers
                    8 => bytes.gread_with::<Nhdr64>(offset, ctx.le)?.into(),
                    _ => return Err(error::Error::Malformed(format!("Notes has unimplemented alignment requirement: {:#x}", alignment)))
                }
            };
            debug!("{:?} - {:#x}", header, *offset);
            // -1 because includes \0 terminator
            let name = bytes.gread_with::<&'a str>(offset, ctx::StrCtx::Length(header.n_namesz - 1))?;
            *offset += 1;
            align(alignment, offset);
            debug!("note name {} - {:#x}", name, *offset);
            let desc = bytes.gread_with::<&'a [u8]>(offset, header.n_descsz)?;
            align(alignment, offset);
            debug!("desc {:?} - {:#x}", desc, *offset);
            Ok((Note {
                name,
                desc,
                n_type: header.n_type,
            }, *offset))
        }
    }

    #[cfg(test)]
    mod tests {
        use super::*;

        static NOTE_DATA: [u8; 68] = [0x04, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00,
                                     0x01, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
                                     0x06, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00,
                                     0x04, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00,
                                     0x03, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0xbc, 0xfc, 0x66, 0xcd, 0xc7, 0xd5, 0x14, 0x7b,
                                     0x53, 0xb1, 0x10, 0x11, 0x94, 0x86, 0x8e, 0xf9,
                                     0x4f, 0xe8, 0xdd, 0xdb];

        static CONTEXT: (usize, container::Ctx) = (4, container::Ctx {
            container: container::Container::Big,
            le: ::scroll::Endian::Little,
        });

        fn make_note_iter(start: usize, end: usize) -> NoteDataIterator<'static> {
            NoteDataIterator {
                data: &NOTE_DATA,
                size: end,
                offset: start,
                ctx: CONTEXT,
            }
        }

        #[test]
        fn iter_single_section() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn iter_multiple_sections() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 32), make_note_iter(32, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn skip_empty_sections() {
            let mut notes = NoteIterator {
                iters: vec![
                    make_note_iter(0, 32),
                    make_note_iter(0, 0),
                    make_note_iter(32, 68),
                ],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn ignore_no_sections() {
            let mut notes = NoteIterator { iters: vec![], index: 0 };
            assert!(notes.next().is_none());
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2013-2014 Simon Sapin.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Abstraction that conditionally compiles either to rust-encoding,
//! or to only support UTF-8.

#[cfg(feature = "query_encoding")] extern crate encoding;

use std::borrow::Cow;

#[cfg(feature = "query_encoding")] use self::encoding::types::{DecoderTrap, EncoderTrap};
#[cfg(feature = "query_encoding")] use self::encoding::label::encoding_from_whatwg_label;
#[cfg(feature = "query_encoding")] pub use self::encoding::types::EncodingRef;

#[cfg(feature = "query_encoding")]
#[deriving(Copy)]
pub struct EncodingOverride {
    /// `None` means UTF-8.
    encoding: Option<EncodingRef>
}

#[cfg(feature = "query_encoding")]
impl EncodingOverride {
    pub fn from_opt_encoding(encoding: Option<EncodingRef>) -> EncodingOverride {
        encoding.map(EncodingOverride::from_encoding).unwrap_or_else(EncodingOverride::utf8)
    }

    pub fn from_encoding(encoding: EncodingRef) -> EncodingOverride {
        EncodingOverride {
            encoding: if encoding.name() == "utf-8" { None } else { Some(encoding) }
        }
    }

    pub fn utf8() -> EncodingOverride {
        EncodingOverride { encoding: None }
    }

    pub fn lookup(label: &[u8]) -> Option<EncodingOverride> {
        ::std::str::from_utf8(label.as_slice())
        .and_then(encoding_from_whatwg_label)
        .map(EncodingOverride::from_encoding)
    }

    pub fn is_utf8(&self) -> bool {
        self.encoding.is_none()
    }

    pub fn decode(&self, input: &[u8]) -> String {
        match self.encoding {
            Some(encoding) => encoding.decode(input, DecoderTrap::Replace).unwrap(),
            None => String::from_utf8_lossy(input).to_string(),
        }
    }

    pub fn encode<'a>(&self, input: &'a str) -> Cow<'a, Vec<u8>, [u8]> {
        match self.encoding {
            Some(encoding) => Cow::Owned(
                encoding.encode(input.as_slice(), EncoderTrap::NcrEscape).unwrap()),
            None => Cow::Borrowed(input.as_bytes()),  // UTF-8
        }
    }
}

#[cfg(not(feature = "query_encoding"))]
#[deriving(Copy)]
pub struct EncodingOverride;

#[cfg(not(feature = "query_encoding"))]
impl EncodingOverride {
    pub fn utf8() -> EncodingOverride {
        EncodingOverride
    }

    pub fn lookup(_label: &[u8]) -> Option<EncodingOverride> {
        None
    }

    pub fn is_utf8(&self) -> bool {
        true
    }

    pub fn decode(&self, input: &[u8]) -> String {
        String::from_utf8_lossy(input).into_owned()
    }

    pub fn encode<'a>(&self, input: &'a str) -> Cow<'a, Vec<u8>, [u8]> {
        Cow::Borrowed(input.as_bytes())
    }
}
<commit_msg>Upgrade to rustc 0.13.0-nightly (96a3c7c6a 2014-12-23 22:21:10 +0000)<commit_after>// Copyright 2013-2014 Simon Sapin.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Abstraction that conditionally compiles either to rust-encoding,
//! or to only support UTF-8.

#[cfg(feature = "query_encoding")] extern crate encoding;

use std::borrow::Cow;

#[cfg(feature = "query_encoding")] use self::encoding::types::{DecoderTrap, EncoderTrap};
#[cfg(feature = "query_encoding")] use self::encoding::label::encoding_from_whatwg_label;
#[cfg(feature = "query_encoding")] pub use self::encoding::types::EncodingRef;

#[cfg(feature = "query_encoding")]
#[deriving(Copy)]
pub struct EncodingOverride {
    /// `None` means UTF-8.
    encoding: Option<EncodingRef>
}

#[cfg(feature = "query_encoding")]
impl EncodingOverride {
    pub fn from_opt_encoding(encoding: Option<EncodingRef>) -> EncodingOverride {
        encoding.map(EncodingOverride::from_encoding).unwrap_or_else(EncodingOverride::utf8)
    }

    pub fn from_encoding(encoding: EncodingRef) -> EncodingOverride {
        EncodingOverride {
            encoding: if encoding.name() == "utf-8" { None } else { Some(encoding) }
        }
    }

    pub fn utf8() -> EncodingOverride {
        EncodingOverride { encoding: None }
    }

    pub fn lookup(label: &[u8]) -> Option<EncodingOverride> {
        ::std::str::from_utf8(label.as_slice())
        .ok()
        .and_then(encoding_from_whatwg_label)
        .map(EncodingOverride::from_encoding)
    }

    pub fn is_utf8(&self) -> bool {
        self.encoding.is_none()
    }

    pub fn decode(&self, input: &[u8]) -> String {
        match self.encoding {
            Some(encoding) => encoding.decode(input, DecoderTrap::Replace).unwrap(),
            None => String::from_utf8_lossy(input).to_string(),
        }
    }

    pub fn encode<'a>(&self, input: &'a str) -> Cow<'a, Vec<u8>, [u8]> {
        match self.encoding {
            Some(encoding) => Cow::Owned(
                encoding.encode(input.as_slice(), EncoderTrap::NcrEscape).unwrap()),
            None => Cow::Borrowed(input.as_bytes()),  // UTF-8
        }
    }
}

#[cfg(not(feature = "query_encoding"))]
#[deriving(Copy)]
pub struct EncodingOverride;

#[cfg(not(feature = "query_encoding"))]
impl EncodingOverride {
    pub fn utf8() -> EncodingOverride {
        EncodingOverride
    }

    pub fn lookup(_label: &[u8]) -> Option<EncodingOverride> {
        None
    }

    pub fn is_utf8(&self) -> bool {
        true
    }

    pub fn decode(&self, input: &[u8]) -> String {
        String::from_utf8_lossy(input).into_owned()
    }

    pub fn encode<'a>(&self, input: &'a str) -> Cow<'a, Vec<u8>, [u8]> {
        Cow::Borrowed(input.as_bytes())
    }
}
<|endoftext|>


<commit_before>//! The film module provides color types and a render target that the image
//! is written too. Functions are also provided for saving PPM images
//! while I wait to migrate over to the Piston image library due to some
//! compile issues TODO: https://github.com/PistonDevelopers/image

use std::io::{File, Open, Write, BufferedWriter};

pub use self::color::Colorf;
pub use self::render_target::RenderTarget;
pub use self::camera::Camera;

pub mod color;
pub mod render_target;
pub mod camera;

/// Write the sequence of bytes as a PPM image file with the desired name
pub fn write_ppm(name: &str, w: u32, h: u32, img: &[u8]) {
    let file = match File::open_mode(&Path::new(name), Open, Write) {
        Ok(f) => f,
        Err(e) => panic!("Failed to open {}: {}", name, e),
    };
    let mut writer = BufferedWriter::new(file);
    match write!(&mut writer, "P6\n{} {}\n255\n", w, h) {
        Err(e) => panic!("Failed to write pixel data to {}: {}", name, e),
        _ => {},
    }
    match writer.write(img) {
        Err(e) => panic!("Failed to write pixel data to {}: {}", name, e),
        _ => {},
    }
}

<commit_msg>Move to uints for w/h params<commit_after>//! The film module provides color types and a render target that the image
//! is written too. Functions are also provided for saving PPM images
//! while I wait to migrate over to the Piston image library due to some
//! compile issues TODO: https://github.com/PistonDevelopers/image

use std::io::{File, Open, Write, BufferedWriter};

pub use self::color::Colorf;
pub use self::render_target::RenderTarget;
pub use self::camera::Camera;

pub mod color;
pub mod render_target;
pub mod camera;

/// Write the sequence of bytes as a PPM image file with the desired name
pub fn write_ppm(name: &str, w: uint, h: uint, img: &[u8]) {
    let file = match File::open_mode(&Path::new(name), Open, Write) {
        Ok(f) => f,
        Err(e) => panic!("Failed to open {}: {}", name, e),
    };
    let mut writer = BufferedWriter::new(file);
    match write!(&mut writer, "P6\n{} {}\n255\n", w, h) {
        Err(e) => panic!("Failed to write pixel data to {}: {}", name, e),
        _ => {},
    }
    match writer.write(img) {
        Err(e) => panic!("Failed to write pixel data to {}: {}", name, e),
        _ => {},
    }
}

<|endoftext|>


<commit_before>use draw_state::DrawState;
use types::{ self, Matrix2d, Scalar };
use deform::DeformGrid;
use {
    CircleArc,
    Ellipse,
    Image,
    ImageSize,
    Line,
    Polygon,
    Rectangle,
};

/// Implemented by all graphics back-ends.
///
/// Default trait methods can be overridden for better performance or higher
/// quality.
///
/// When drawing, use this trait as generic constraint:
///
/// ```ignore
/// fn draw<G: Graphics>(c: &Context, g: &mut G) { ... }
/// ```
pub trait Graphics: Sized {
    /// The texture type associated with the back-end.
    ///
    /// In generic code, this type is often unknown.
    /// This might lead to more boilerplate code:
    ///
    /// ```ignore
    /// fn draw_texture<G, T>(c: &Context, g: &mut G)
    ///     where G: Graphcis<Texture = T>, T: ImageSize { ... }
    /// ```
    ///
    /// Code written specifically for one back-end can be easier to write.
    /// Later, when the code is done, it can be refactored into generic code.
    type Texture: ImageSize;

    /// Clears background with a color.
    ///
    /// The color should replace the values in the buffer.
    fn clear_color(&mut self, color: types::Color);

    /// Clears stencil buffer with a value, usually 0.
    ///
    /// A stencil buffer contains values that are not visible on the screen.
    /// These values are used to test against the pixel to paint.
    ///
    /// If you are drawing a shape for clipping and forgot to clear the
    /// stencil buffer, then the clipping shape will carry over in next frame
    /// and cause artifacts.
    fn clear_stencil(&mut self, value: u8);

    /// Renders list of 2d triangles.
    fn tri_list<F>(&mut self, draw_state: &DrawState, color: &[f32; 4], f: F)
        where F: FnMut(&mut FnMut(&[f32]));

    /// Renders list of 2d triangles.
    ///
    /// A texture coordinate is assigned per vertex.
    /// The texture coordinates refers to the current texture.
    fn tri_list_uv<F>(
        &mut self,
        draw_state: &DrawState,
        color: &[f32; 4],
        texture: &<Self as Graphics>::Texture,
        f: F
    ) where F: FnMut(&mut FnMut(&[f32], &[f32]));

    /// Draws a rectangle.
    #[inline(always)]
    fn rectangle<R: Into<types::Rectangle>>(
        &mut self,
        r: &Rectangle,
        rectangle: R,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        r.draw_tri(rectangle, draw_state, transform, self);
    }

    /// Draws a polygon.
    #[inline(always)]
    fn polygon(
        &mut self,
        p: &Polygon,
        polygon: types::Polygon,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        p.draw_tri(polygon, draw_state, transform, self);
    }

    /// Draws a tweened polygon using linear interpolation.
    #[inline(always)]
    fn polygon_tween_lerp(
        &mut self,
        p: &Polygon,
        polygons: types::Polygons,
        tween_factor: Scalar,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        p.draw_tween_lerp_tri(polygons, tween_factor,
            draw_state, transform, self);
    }

    /// Draws image.
    #[inline(always)]
    fn image(
        &mut self,
        image: &Image,
        texture: &Self::Texture,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        image.draw_tri(texture, draw_state, transform, self);
    }

    /// Draws ellipse.
    #[inline(always)]
    fn ellipse<R: Into<types::Rectangle>>(
        &mut self,
        e: &Ellipse,
        rectangle: R,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        e.draw_tri(rectangle, draw_state, transform, self);
    }

    /// Draws line.
    #[inline(always)]
    fn line<L: Into<types::Line>>(
        &mut self,
        l: &Line,
        line: L,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        l.draw_tri(line, draw_state, transform, self);
    }

    /// Draws circle arc.
    #[inline(always)]
    fn circle_arc<R: Into<types::Rectangle>>(
        &mut self,
        c: &CircleArc,
        rectangle: R,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        c.draw_tri(rectangle, draw_state, transform, self);
    }

    /// Draws deformed image.
    #[inline(always)]
    fn deform_image(
        &mut self,
        d: &DeformGrid,
        texture: &Self::Texture,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        d.draw_image_tri(texture, draw_state, transform, self);
    }
}
<commit_msg>Improved docs for `Graphics::tri_list`<commit_after>use draw_state::DrawState;
use types::{ self, Matrix2d, Scalar };
use deform::DeformGrid;
use {
    CircleArc,
    Ellipse,
    Image,
    ImageSize,
    Line,
    Polygon,
    Rectangle,
};

/// Implemented by all graphics back-ends.
///
/// Default trait methods can be overridden for better performance or higher
/// quality.
///
/// When drawing, use this trait as generic constraint:
///
/// ```ignore
/// fn draw<G: Graphics>(c: &Context, g: &mut G) { ... }
/// ```
pub trait Graphics: Sized {
    /// The texture type associated with the back-end.
    ///
    /// In generic code, this type is often unknown.
    /// This might lead to more boilerplate code:
    ///
    /// ```ignore
    /// fn draw_texture<G, T>(c: &Context, g: &mut G)
    ///     where G: Graphcis<Texture = T>, T: ImageSize { ... }
    /// ```
    ///
    /// Code written specifically for one back-end can be easier to write.
    /// Later, when the code is done, it can be refactored into generic code.
    type Texture: ImageSize;

    /// Clears background with a color.
    ///
    /// The color should replace the values in the buffer.
    fn clear_color(&mut self, color: types::Color);

    /// Clears stencil buffer with a value, usually 0.
    ///
    /// A stencil buffer contains values that are not visible on the screen.
    /// These values are used to test against the pixel to paint.
    ///
    /// If you are drawing a shape for clipping and forgot to clear the
    /// stencil buffer, then the clipping shape will carry over in next frame
    /// and cause artifacts.
    fn clear_stencil(&mut self, value: u8);

    /// Renders list of 2d triangles using a solid color.
    ///
    /// All vertices share the same color.
    ///
    /// The back-end calls the closure with a closure to receive vertices.
    /// First, the back-end sets up shaders and such to prepare.
    /// Then it calls the closure, which calls back with chunks of vertices.
    /// The number of vertices per chunk never exceeds
    /// `BACK_END_MAX_VERTEX_COUNT`.
    /// Vertex positions are encoded `[x0, y0, x1, y1, ...]`.
    fn tri_list<F>(&mut self, draw_state: &DrawState, color: &[f32; 4], f: F)
        where F: FnMut(&mut FnMut(&[f32]));

    /// Renders list of 2d triangles.
    ///
    /// A texture coordinate is assigned per vertex.
    /// The texture coordinates refers to the current texture.
    fn tri_list_uv<F>(
        &mut self,
        draw_state: &DrawState,
        color: &[f32; 4],
        texture: &<Self as Graphics>::Texture,
        f: F
    ) where F: FnMut(&mut FnMut(&[f32], &[f32]));

    /// Draws a rectangle.
    #[inline(always)]
    fn rectangle<R: Into<types::Rectangle>>(
        &mut self,
        r: &Rectangle,
        rectangle: R,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        r.draw_tri(rectangle, draw_state, transform, self);
    }

    /// Draws a polygon.
    #[inline(always)]
    fn polygon(
        &mut self,
        p: &Polygon,
        polygon: types::Polygon,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        p.draw_tri(polygon, draw_state, transform, self);
    }

    /// Draws a tweened polygon using linear interpolation.
    #[inline(always)]
    fn polygon_tween_lerp(
        &mut self,
        p: &Polygon,
        polygons: types::Polygons,
        tween_factor: Scalar,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        p.draw_tween_lerp_tri(polygons, tween_factor,
            draw_state, transform, self);
    }

    /// Draws image.
    #[inline(always)]
    fn image(
        &mut self,
        image: &Image,
        texture: &Self::Texture,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        image.draw_tri(texture, draw_state, transform, self);
    }

    /// Draws ellipse.
    #[inline(always)]
    fn ellipse<R: Into<types::Rectangle>>(
        &mut self,
        e: &Ellipse,
        rectangle: R,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        e.draw_tri(rectangle, draw_state, transform, self);
    }

    /// Draws line.
    #[inline(always)]
    fn line<L: Into<types::Line>>(
        &mut self,
        l: &Line,
        line: L,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        l.draw_tri(line, draw_state, transform, self);
    }

    /// Draws circle arc.
    #[inline(always)]
    fn circle_arc<R: Into<types::Rectangle>>(
        &mut self,
        c: &CircleArc,
        rectangle: R,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        c.draw_tri(rectangle, draw_state, transform, self);
    }

    /// Draws deformed image.
    #[inline(always)]
    fn deform_image(
        &mut self,
        d: &DeformGrid,
        texture: &Self::Texture,
        draw_state: &DrawState,
        transform: Matrix2d
    ) {
        d.draw_image_tri(texture, draw_state, transform, self);
    }
}
<|endoftext|>


<commit_before>pub mod twodim;
pub mod nhood;

use traits::EvolutionState;

pub struct EmptyState;

impl EvolutionState for EmptyState {
    fn update(&mut self) {}
}
<commit_msg>Note about EmptyState.<commit_after>pub mod twodim;
pub mod nhood;

use traits::EvolutionState;

/// Dummy evolution state to be used with
/// cellular automata where this concept is
/// not applicable.
pub struct EmptyState;

impl EvolutionState for EmptyState {
    fn update(&mut self) {}
}
<|endoftext|>


<commit_before>use alloc::rc::Rc;
use spin::RwLock;
use core::hash::{Hash, Hasher};
// Compiler warns about a field not exported in core.
#[allow(deprecated)]
use core::hash::SipHasher;
use collections::vec::Vec;
use core::cmp;

enum HashMapMember<K: Hash + cmp::PartialEq, V> {
    None,
    Tombstone,
    Bucket(Bucket<K, V>),
}

struct Bucket<K: Hash + cmp::PartialEq, V> {
    key: K,
    value: Rc<V>,
}

pub struct HashMap<K: Hash + cmp::PartialEq, V> {
    count: usize,
    table: RwLock<Vec<HashMapMember<K, V>>>,
}

impl<K: Hash + cmp::PartialEq, V> Bucket<K, V> {
    fn new(key: K, value: V) -> Self {
        Bucket {
            key: key,
            value: Rc::new(value),
        }
    }
}

impl<K: Hash + cmp::PartialEq, V> HashMap<K, V> {
    pub fn new(size: usize) -> Self {
        let mut table = vec![];
        for _ in 0..size {
            table.push(HashMapMember::None);
        }
        HashMap {
            count: 0,
            table: RwLock::new(table),
        }
    }

    // Compiler warns about a field not exported in core.
    #[allow(deprecated)]
    fn get_index(&self, key: &K) -> usize {
        let mut hasher = SipHasher::new();
        key.hash(&mut hasher);
        (hasher.finish() as usize) % self.table.read().len()
    }

    pub fn remove(&mut self, key: K) {
        let index = self.get_index(&key);
        let mut table = self.table.write();
        let (begin, end) = table.split_at_mut(index);
        for entry in end.iter_mut().chain(begin) {
            match entry {
                &mut HashMapMember::Tombstone => continue,
                &mut HashMapMember::None => return,
                &mut HashMapMember::Bucket(ref mut b) if b.key != key => continue,
                loc @ &mut HashMapMember::Bucket(_) => {
                    *loc = HashMapMember::Tombstone;
                    return;
                }
            }
        }
    }

    pub fn insert(&mut self, key: K, value: V) {
        let index = self.get_index(&key);
        let mut table = self.table.write();
        let (begin, end) = table.split_at_mut(index);
        for entry in end.iter_mut().chain(begin) {
            match entry {
                &mut HashMapMember::Tombstone => continue,
                loc @ &mut HashMapMember::None => {
                    *loc = HashMapMember::Bucket(Bucket::new(key, value));
                    return;
                }
                &mut HashMapMember::Bucket(ref mut b) => {
                    if b.key == key {
                        b.value = Rc::new(value);
                        self.count += 1;
                        return;
                    } else {
                        continue;
                    }
                }
            }
        }
    }

    pub fn contains(&self, key: &K) -> bool {
        let index = self.get_index(key);
        let table = self.table.read();
        let (begin, end) = table.split_at(index);
        for entry in end.iter().chain(begin) {
            match *entry {
                HashMapMember::Tombstone => continue,
                HashMapMember::None => return false,
                HashMapMember::Bucket(ref b) => {
                    if &b.key == key {
                        return true;
                    } else {
                        continue;
                    }
                }
            }
        }
        false
    }

    pub fn get(&self, key: &K) -> Option<Rc<V>> {
        let index = self.get_index(key);
        let mut table = self.table.write();
        let (begin, end) = table.split_at_mut(index);
        for entry in end.iter_mut().chain(begin) {
            match *entry {
                HashMapMember::Tombstone => continue,
                HashMapMember::None => return None,
                HashMapMember::Bucket(ref b) => {
                    if &b.key == key {
                        return Some(b.value.clone());
                    } else {
                        continue;
                    }
                }
            }
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use alloc::rc::Rc;
    use super::HashMap;

    #[test]
    fn test_new() {
        let ht = HashMap::<i32, i32>::new(10);
        let length = ht.table.read().len();
        assert_eq!(length, 10);

        let ht1 = HashMap::<&str, &str>::new(100);
        let length = ht1.table.read().len();
        assert_eq!(length, 100);

        #[allow(unused)]
        struct Custom<'a> {
            int_field: i32,
            str_field: &'a str,
        }

        let ht2 = HashMap::<i32, Custom>::new(1000);
        let length = ht2.table.read().len();
        assert_eq!(length, 1000);
    }

    #[test]
    fn test_basic_operations() {
        let mut ht = HashMap::<i32, i32>::new(10);
        let k = 42;
        ht.insert(k, 43);
        assert!(ht.contains(&k));
        assert_eq!(ht.get(&k), Some(Rc::new(43)));

        ht.remove(42);
        assert!(!ht.contains(&k));
        assert_eq!(ht.get(&k), None);

        ht.insert(42, 44);
        assert!(ht.contains(&k));
        assert_eq!(ht.get(&k), Some(Rc::new(44)));

        ht.insert(42, 45);
        assert!(ht.contains(&k));
        assert_eq!(ht.get(&k), Some(Rc::new(45)));
    }
}
<commit_msg>Automatically rebalance hash table<commit_after>use alloc::rc::Rc;
use spin::RwLock;
use core::hash::{Hash, Hasher};
// Compiler warns about a field not exported in core.
#[allow(deprecated)]
use core::hash::SipHasher;
use collections::vec::Vec;
use core::mem;
use core::cmp;

enum HashMapMember<K: Hash + cmp::PartialEq, V> {
    None,
    Tombstone,
    Bucket(Bucket<K, V>),
}

struct Bucket<K: Hash + cmp::PartialEq, V> {
    key: K,
    value: Rc<V>,
}

pub struct HashMap<K: Hash + cmp::PartialEq, V> {
    count: usize,
    table: RwLock<Vec<HashMapMember<K, V>>>,
    rebalance_factor: f32,
}

impl<K: Hash + cmp::PartialEq, V> Bucket<K, V> {
    fn new(key: K, value: Rc<V>) -> Self {
        Bucket {
            key: key,
            value: value,
        }
    }
}

impl<K: Hash + cmp::PartialEq, V> HashMap<K, V> {
    pub fn new(size: usize) -> Self {
        let mut table = vec![];
        for _ in 0..size {
            table.push(HashMapMember::None);
        }
        HashMap {
            count: 0,
            table: RwLock::new(table),
            rebalance_factor: 0.75,
        }
    }

    fn calculate_load(&self) -> f32 {
        self.count as f32 / self.table.read().len() as f32
    }

    fn rebalance(&mut self) {
        let mut table = self.table.write();
        let mut temp = HashMap::new(table.len() * 2);
        for entry in table.drain(1..) {
            match entry {
                HashMapMember::Bucket(b) => temp.insert_rc(b.key, b.value),
                _ => continue,
            }
        }
        mem::replace(&mut *table, temp.table.into_inner());
    }

    // Compiler warns about a field not exported in core.
    #[allow(deprecated)]
    fn get_index(&self, key: &K) -> usize {
        let mut hasher = SipHasher::new();
        key.hash(&mut hasher);
        (hasher.finish() as usize) % self.table.read().len()
    }

    pub fn remove(&mut self, key: K) {
        let index = self.get_index(&key);
        let mut table = self.table.write();
        let (begin, end) = table.split_at_mut(index);
        for entry in end.iter_mut().chain(begin) {
            match entry {
                &mut HashMapMember::Tombstone => continue,
                &mut HashMapMember::None => return,
                &mut HashMapMember::Bucket(ref mut b) if b.key != key => continue,
                loc @ &mut HashMapMember::Bucket(_) => {
                    *loc = HashMapMember::Tombstone;
                    return;
                }
            }
        }
    }

    pub fn insert_rc(&mut self, key: K, value: Rc<V>) {
        self.count += 1;
        if self.calculate_load() >= self.rebalance_factor {
            self.rebalance();
        }
        let index = self.get_index(&key);
        let mut table = self.table.write();
        let (begin, end) = table.split_at_mut(index);
        for entry in end.iter_mut().chain(begin) {
            match entry {
                &mut HashMapMember::Tombstone => continue,
                loc @ &mut HashMapMember::None => {
                    *loc = HashMapMember::Bucket(Bucket::new(key, value));
                    return;
                }
                &mut HashMapMember::Bucket(ref mut b) => {
                    if b.key == key {
                        b.value = value;
                        return;
                    } else {
                        continue;
                    }
                }
            }
        }
    }

    pub fn insert(&mut self, key: K, value: V) {
        self.insert_rc(key, Rc::new(value))
    }

    pub fn contains(&self, key: &K) -> bool {
        let index = self.get_index(key);
        let table = self.table.read();
        let (begin, end) = table.split_at(index);
        for entry in end.iter().chain(begin) {
            match *entry {
                HashMapMember::Tombstone => continue,
                HashMapMember::None => return false,
                HashMapMember::Bucket(ref b) => {
                    if &b.key == key {
                        return true;
                    } else {
                        continue;
                    }
                }
            }
        }
        false
    }

    pub fn get(&self, key: &K) -> Option<Rc<V>> {
        let index = self.get_index(key);
        let mut table = self.table.write();
        let (begin, end) = table.split_at_mut(index);
        for entry in end.iter_mut().chain(begin) {
            match *entry {
                HashMapMember::Tombstone => continue,
                HashMapMember::None => return None,
                HashMapMember::Bucket(ref b) => {
                    if &b.key == key {
                        return Some(b.value.clone());
                    } else {
                        continue;
                    }
                }
            }
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use alloc::rc::Rc;
    use super::HashMap;

    #[test]
    fn test_new() {
        let ht = HashMap::<i32, i32>::new(10);
        let length = ht.table.read().len();
        assert_eq!(length, 10);

        let ht1 = HashMap::<&str, &str>::new(100);
        let length = ht1.table.read().len();
        assert_eq!(length, 100);

        #[allow(unused)]
        struct Custom<'a> {
            int_field: i32,
            str_field: &'a str,
        }

        let ht2 = HashMap::<i32, Custom>::new(1000);
        let length = ht2.table.read().len();
        assert_eq!(length, 1000);
    }

    #[test]
    fn test_basic_operations() {
        let mut ht = HashMap::<i32, i32>::new(10);
        let k = 42;
        ht.insert(k, 43);
        assert!(ht.contains(&k));
        assert_eq!(ht.get(&k), Some(Rc::new(43)));

        ht.remove(42);
        assert!(!ht.contains(&k));
        assert_eq!(ht.get(&k), None);

        ht.insert(42, 44);
        assert!(ht.contains(&k));
        assert_eq!(ht.get(&k), Some(Rc::new(44)));

        ht.insert(42, 45);
        assert!(ht.contains(&k));
        assert_eq!(ht.get(&k), Some(Rc::new(45)));
    }

    #[test]
    fn test_rebalance() {

        let initial_size = 4;
        let mut ht = HashMap::<usize, usize>::new(initial_size);
        assert_eq!(ht.count, 0);

        ht.insert(0, 0);
        assert_eq!(ht.count, 1);
        assert_eq!(ht.table.read().len(), 4);

        ht.insert(1, 1);
        assert_eq!(ht.count, 2);
        assert_eq!(ht.table.read().len(), 4);

        ht.insert(2, 2);
        assert_eq!(ht.count, 3);
        assert_eq!(ht.table.read().len(), 8);

        ht.insert(3, 3);
        assert_eq!(ht.count, 4);
        assert!(ht.calculate_load() < ht.rebalance_factor);
        assert_eq!(ht.table.read().len(), 8);

        ht.insert(4, 4);
        assert_eq!(ht.count, 5);
        assert!(ht.calculate_load() < ht.rebalance_factor);
        assert_eq!(ht.table.read().len(), 8);

        ht.insert(5, 5);
        assert_eq!(ht.count, 6);
        assert!(ht.calculate_load() < ht.rebalance_factor);
        assert_eq!(ht.table.read().len(), 16);
    }
}
<|endoftext|>


<commit_before>use alloc::rc::Rc;
use spin::RwLock;
use core::hash::{Hash, Hasher};
// Compiler warns about a field not exported in core.
#[allow(deprecated)]
use core::hash::SipHasher;
use collections::vec::Vec;
use core::cmp;

enum HashMapMember<K: Hash + cmp::PartialEq, V> {
    None,
    Tombstone,
    Bucket(Bucket<K, V>),
}

struct Bucket<K: Hash + cmp::PartialEq, V> {
    key: K,
    value: Rc<V>,
}

pub struct HashMap<K: Hash + cmp::PartialEq, V> {
    count: usize,
    table: RwLock<Vec<HashMapMember<K, V>>>,
}

impl<K: Hash + cmp::PartialEq, V> Bucket<K, V> {
    fn new(key: K, value: V) -> Self {
        Bucket {
            key: key,
            value: Rc::new(value),
        }
    }
}

impl<K: Hash + cmp::PartialEq, V> HashMap<K, V> {
    pub fn new(size: usize) -> Self {
        let mut table = vec![];
        for _ in 0..size {
            table.push(HashMapMember::None);
        }
        HashMap {
            count: 0,
            table: RwLock::new(table),
        }
    }

    // Compiler warns about a field not exported in core.
    #[allow(deprecated)]
    fn get_index(&self, key: &K) -> usize {
        let mut hasher = SipHasher::new();
        key.hash(&mut hasher);
        (hasher.finish() as usize) % self.table.read().len()
    }

    pub fn remove(&mut self, key: K) {
        let index = self.get_index(&key);
        let mut table = self.table.write();
        let (begin, end) = table.split_at_mut(index);
        for entry in end.iter_mut().chain(begin) {
            match entry {
                &mut HashMapMember::Tombstone => continue,
                &mut HashMapMember::None => return,
                &mut HashMapMember::Bucket(ref mut b) if b.key != key => continue,
                loc @ &mut HashMapMember::Bucket(_) => {
                    *loc = HashMapMember::Tombstone;
                    return;
                }
            }
        }
    }

    pub fn insert(&mut self, key: K, value: V) {
        let index = self.get_index(&key);
        let mut table = self.table.write();
        let (begin, end) = table.split_at_mut(index);
        for entry in end.iter_mut().chain(begin) {
            match entry {
                &mut HashMapMember::Tombstone => continue,
                loc @ &mut HashMapMember::None => {
                    *loc = HashMapMember::Bucket(Bucket::new(key, value));
                    return;
                }
                &mut HashMapMember::Bucket(ref mut b) => {
                    if b.key == key {
                        b.value = Rc::new(value);
                        self.count += 1;
                        return;
                    } else {
                        continue;
                    }
                }
            }
        }
    }

    pub fn contains(&self, key: &K) -> bool {
        let index = self.get_index(&key);
        let table = self.table.read();
        let (begin, end) = table.split_at(index);
        for entry in end.iter().chain(begin) {
            match entry {
                &HashMapMember::Tombstone => continue,
                &HashMapMember::None => return false,
                &HashMapMember::Bucket(ref b) => {
                    if &b.key == key {
                        return true;
                    } else {
                        continue;
                    }
                }
            }
        }
        false
    }

    pub fn get(&self, key: &K) -> Option<Rc<V>> {
        let index = self.get_index(&key);
        let mut table = self.table.write();
        let (begin, end) = table.split_at_mut(index);
        for entry in end.iter_mut().chain(begin) {
            match entry {
                &mut HashMapMember::Tombstone => continue,
                &mut HashMapMember::None => return None,
                &mut HashMapMember::Bucket(ref b) => {
                    if &b.key == key {
                        return Some(b.value.clone());
                    } else {
                        continue;
                    }
                }
            }
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use alloc::rc::Rc;
    use super::HashMap;

    #[test]
    fn test_new() {
        let ht = HashMap::<i32, i32>::new(10);
        let length = ht.table.read().len();
        assert_eq!(length, 10);

        let ht1 = HashMap::<&str, &str>::new(100);
        let length = ht1.table.read().len();
        assert_eq!(length, 100);

        #[allow(unused)]
        struct Custom<'a> {
            int_field: i32,
            str_field: &'a str,
        }

        let ht2 = HashMap::<i32, Custom>::new(1000);
        let length = ht2.table.read().len();
        assert_eq!(length, 1000);
    }

    #[test]
    fn test_basic_operations() {
        let mut ht = HashMap::<i32, i32>::new(10);
        let k = 42;
        ht.insert(k, 43);
        assert!(ht.contains(&k));
        assert_eq!(ht.get(&k), Some(Rc::new(43)));

        ht.remove(42);
        assert!(!ht.contains(&k));
        assert_eq!(ht.get(&k), None);

        ht.insert(42, 44);
        assert!(ht.contains(&k));
        assert_eq!(ht.get(&k), Some(Rc::new(44)));

        ht.insert(42, 45);
        assert!(ht.contains(&k));
        assert_eq!(ht.get(&k), Some(Rc::new(45)));
    }
}
<commit_msg>Clippy lints<commit_after>use alloc::rc::Rc;
use spin::RwLock;
use core::hash::{Hash, Hasher};
// Compiler warns about a field not exported in core.
#[allow(deprecated)]
use core::hash::SipHasher;
use collections::vec::Vec;
use core::cmp;

enum HashMapMember<K: Hash + cmp::PartialEq, V> {
    None,
    Tombstone,
    Bucket(Bucket<K, V>),
}

struct Bucket<K: Hash + cmp::PartialEq, V> {
    key: K,
    value: Rc<V>,
}

pub struct HashMap<K: Hash + cmp::PartialEq, V> {
    count: usize,
    table: RwLock<Vec<HashMapMember<K, V>>>,
}

impl<K: Hash + cmp::PartialEq, V> Bucket<K, V> {
    fn new(key: K, value: V) -> Self {
        Bucket {
            key: key,
            value: Rc::new(value),
        }
    }
}

impl<K: Hash + cmp::PartialEq, V> HashMap<K, V> {

    pub fn new(size: usize) -> Self {
        let mut table = vec![];
        for _ in 0..size {
            table.push(HashMapMember::None);
        }
        HashMap {
            count: 0,
            table: RwLock::new(table),
        }
    }

    // Compiler warns about a field not exported in core.
    #[allow(deprecated)]
    fn get_index(&self, key: &K) -> usize {
        let mut hasher = SipHasher::new();
        key.hash(&mut hasher);
        (hasher.finish() as usize) % self.table.read().len()
    }

    pub fn remove(&mut self, key: K) {
        let index = self.get_index(&key);
        let mut table = self.table.write();
        let (begin, end) = table.split_at_mut(index);
        for entry in end.iter_mut().chain(begin) {
            match entry {
                &mut HashMapMember::Tombstone => continue,
                &mut HashMapMember::None => return,
                &mut HashMapMember::Bucket(ref mut b) if b.key != key => continue,
                loc @ &mut HashMapMember::Bucket(_) => {
                    *loc = HashMapMember::Tombstone;
                    return;
                }
            }
        }
    }

    pub fn insert(&mut self, key: K, value: V) {
        let index = self.get_index(&key);
        let mut table = self.table.write();
        let (begin, end) = table.split_at_mut(index);
        for entry in end.iter_mut().chain(begin) {
            match entry {
                &mut HashMapMember::Tombstone => continue,
                loc @ &mut HashMapMember::None => {
                    *loc = HashMapMember::Bucket(Bucket::new(key, value));
                    return;
                }
                &mut HashMapMember::Bucket(ref mut b) => {
                    if b.key == key {
                        b.value = Rc::new(value);
                        self.count += 1;
                        return;
                    } else {
                        continue;
                    }
                }
            }
        }
    }

    pub fn contains(&self, key: &K) -> bool {
        let index = self.get_index(key);
        let table = self.table.read();
        let (begin, end) = table.split_at(index);
        for entry in end.iter().chain(begin) {
            match *entry {
                HashMapMember::Tombstone => continue,
                HashMapMember::None => return false,
                HashMapMember::Bucket(ref b) => {
                    if &b.key == key {
                        return true;
                    } else {
                        continue;
                    }
                }
            }
        }
        false
    }

    pub fn get(&self, key: &K) -> Option<Rc<V>> {
        let index = self.get_index(key);
        let mut table = self.table.write();
        let (begin, end) = table.split_at_mut(index);
        for entry in end.iter_mut().chain(begin) {
            match *entry {
                HashMapMember::Tombstone => continue,
                HashMapMember::None => return None,
                HashMapMember::Bucket(ref b) => {
                    if &b.key == key {
                        return Some(b.value.clone());
                    } else {
                        continue;
                    }
                }
            }
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use alloc::rc::Rc;
    use super::HashMap;

    #[test]
    fn test_new() {
        let ht = HashMap::<i32, i32>::new(10);
        let length = ht.table.read().len();
        assert_eq!(length, 10);

        let ht1 = HashMap::<&str, &str>::new(100);
        let length = ht1.table.read().len();
        assert_eq!(length, 100);

        #[allow(unused)]
        struct Custom<'a> {
            int_field: i32,
            str_field: &'a str,
        }

        let ht2 = HashMap::<i32, Custom>::new(1000);
        let length = ht2.table.read().len();
        assert_eq!(length, 1000);
    }

    #[test]
    fn test_basic_operations() {
        let mut ht = HashMap::<i32, i32>::new(10);
        let k = 42;
        ht.insert(k, 43);
        assert!(ht.contains(&k));
        assert_eq!(ht.get(&k), Some(Rc::new(43)));

        ht.remove(42);
        assert!(!ht.contains(&k));
        assert_eq!(ht.get(&k), None);

        ht.insert(42, 44);
        assert!(ht.contains(&k));
        assert_eq!(ht.get(&k), Some(Rc::new(44)));

        ht.insert(42, 45);
        assert!(ht.contains(&k));
        assert_eq!(ht.get(&k), Some(Rc::new(45)));
    }
}
<|endoftext|>


<commit_before>use ::libc::{self, c_int, O_EXCL, O_RDWR, O_CREAT};
use ::libc::types::os::arch::posix01::stat as stat_t;
use std::io;
use std::os::unix::io::{RawFd, FromRawFd};
use std::fs::{self, File, OpenOptions};
use std::path::Path;
use std::ffi::CString;
use ::util::{tmpname, cstr};
use super::unix_common::O_CLOEXEC;
pub use super::unix_common::create;

unsafe fn stat(fd: RawFd) -> io::Result<stat_t> {
    let mut meta: stat_t = ::std::mem::zeroed();
    if libc::fstat(fd, &mut meta as *mut stat_t) != 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(meta)
    }
}

// Helper for ensuring that the temporary file gets deleted.
struct DeleteGuard<'a>(&'a Path);

impl<'a> Drop for DeleteGuard<'a> {
    fn drop(&mut self) {
        let _ = fs::remove_file(self.0);
    }
}

pub fn create_shared(dir: &Path, count: usize) -> io::Result<Vec<File>> {
    let mut opts = OpenOptions::new();
    opts.read(true).write(true).create(false);

    if count == 0 {
        return Ok(vec![]);
    }
    'outer: loop {
        let tmp_path = dir.join(&tmpname());
        return match unsafe {
            libc::open(try!(cstr(&tmp_path)).as_ptr(), O_CLOEXEC | O_EXCL | O_RDWR | O_CREAT, 0o600)
        } {
            -1 => {
                let err = io::Error::last_os_error();
                if err.kind() == io::ErrorKind::AlreadyExists {
                    continue;
                } else {
                    Err(err)
                }
            },
            fd => unsafe {
                let first = FromRawFd::from_raw_fd(fd);
                let dg = DeleteGuard(&tmp_path);

                let target_meta = try!(stat(fd));
                let mut files: Vec<File> = try!((1..count).map(|_| opts.open(&tmp_path)).collect());
                for file in &files {
                    let meta = try!(stat(file.as_raw_fd()));
                    if meta.st_dev != target_meta.st_dev ||
                       meta.st_ino != target_meta.st_ino ||
                       // Even if the device information get's reused, the owner should actually be
                       // sufficient.
                       meta.st_uid != target_meta.st_uid ||
                       meta.st_gid != target_meta.st_gid {

                        // Error? Panic? If we hit this, we're likely under attack (or a hardware
                        // glitch/reconfiguration?.
                        continue 'outer;
                    }

                }
                files.push(first);
                Ok(files)
            },
        }
    }
}
<commit_msg>Fix warnings on unix.<commit_after>use ::libc::{self, O_EXCL, O_RDWR, O_CREAT};
use ::libc::types::os::arch::posix01::stat as stat_t;
use std::io;
use std::os::unix::io::{RawFd, FromRawFd, AsRawFd};
use std::fs::{self, File, OpenOptions};
use std::path::Path;
use ::util::{tmpname, cstr};
use super::unix_common::O_CLOEXEC;
pub use super::unix_common::create;

unsafe fn stat(fd: RawFd) -> io::Result<stat_t> {
    let mut meta: stat_t = ::std::mem::zeroed();
    if libc::fstat(fd, &mut meta as *mut stat_t) != 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(meta)
    }
}

// Helper for ensuring that the temporary file gets deleted.
struct DeleteGuard<'a>(&'a Path);

impl<'a> Drop for DeleteGuard<'a> {
    fn drop(&mut self) {
        let _ = fs::remove_file(self.0);
    }
}

pub fn create_shared(dir: &Path, count: usize) -> io::Result<Vec<File>> {
    let mut opts = OpenOptions::new();
    opts.read(true).write(true).create(false);

    if count == 0 {
        return Ok(vec![]);
    }
    'outer: loop {
        let tmp_path = dir.join(&tmpname());
        return match unsafe {
            libc::open(try!(cstr(&tmp_path)).as_ptr(), O_CLOEXEC | O_EXCL | O_RDWR | O_CREAT, 0o600)
        } {
            -1 => {
                let err = io::Error::last_os_error();
                if err.kind() == io::ErrorKind::AlreadyExists {
                    continue;
                } else {
                    Err(err)
                }
            },
            fd => unsafe {
                let first = FromRawFd::from_raw_fd(fd);
                let _dg = DeleteGuard(&tmp_path);

                let target_meta = try!(stat(fd));
                let mut files: Vec<File> = try!((1..count).map(|_| opts.open(&tmp_path)).collect());
                for file in &files {
                    let meta = try!(stat(file.as_raw_fd()));
                    if meta.st_dev != target_meta.st_dev ||
                       meta.st_ino != target_meta.st_ino ||
                       // Even if the device information get's reused, the owner should actually be
                       // sufficient.
                       meta.st_uid != target_meta.st_uid ||
                       meta.st_gid != target_meta.st_gid {

                        // Error? Panic? If we hit this, we're likely under attack (or a hardware
                        // glitch/reconfiguration?.
                        continue 'outer;
                    }

                }
                files.push(first);
                Ok(files)
            },
        }
    }
}
<|endoftext|>


<commit_before>use std::{
    io::{self, ErrorKind},
    net::SocketAddr,
    sync::Arc,
    time::Duration,
};

use byte_string::ByteStr;
use kcp::{Error as KcpError, KcpResult};
use log::{debug, error, trace};
use tokio::{
    net::{ToSocketAddrs, UdpSocket},
    sync::mpsc,
    task::JoinHandle,
    time,
};

use crate::{config::KcpConfig, session::KcpSessionManager, stream::KcpStream};

pub struct KcpListener {
    udp: Arc<UdpSocket>,
    accept_rx: mpsc::Receiver<(KcpStream, SocketAddr)>,
    task_watcher: JoinHandle<()>,
}

impl Drop for KcpListener {
    fn drop(&mut self) {
        self.task_watcher.abort();
    }
}

impl KcpListener {
    pub async fn bind<A: ToSocketAddrs>(config: KcpConfig, addr: A) -> KcpResult<KcpListener> {
        let udp = UdpSocket::bind(addr).await?;
        let udp = Arc::new(udp);
        let server_udp = udp.clone();

        let (accept_tx, accept_rx) = mpsc::channel(1024 /* backlogs */);
        let task_watcher = tokio::spawn(async move {
            let (close_tx, mut close_rx) = mpsc::channel(64);

            let mut sessions = KcpSessionManager::new();
            let mut packet_buffer = [0u8; 65536];
            loop {
                tokio::select! {
                    peer_addr = close_rx.recv() => {
                        let peer_addr = peer_addr.expect("close_tx closed unexpectly");
                        sessions.close_peer(peer_addr);
                        trace!("session peer_addr: {} removed", peer_addr);
                    }

                    recv_res = udp.recv_from(&mut packet_buffer) => {
                        match recv_res {
                            Err(err) => {
                                error!("udp.recv_from failed, error: {}", err);
                                time::sleep(Duration::from_secs(1)).await;
                            }
                            Ok((n, peer_addr)) => {
                                let packet = &mut packet_buffer[..n];

                                log::trace!("received peer: {}, {:?}", peer_addr, ByteStr::new(packet));

                                let mut conv = kcp::get_conv(packet);
                                if conv == 0 {
                                    // Allocate a conv for client.
                                    conv = sessions.alloc_conv();
                                    debug!("allocate {} conv for peer: {}", conv, peer_addr);

                                    kcp::set_conv(packet, conv);
                                }

                                let sn = kcp::get_sn(packet);

                                let session = match sessions.get_or_create(&config, conv, sn, &udp, peer_addr, &close_tx).await {
                                    Ok((s, created)) => {
                                        if created {
                                            // Created a new session, constructed a new accepted client
                                            let stream = KcpStream::with_session(s.clone());
                                            if let Err(..) = accept_tx.try_send((stream, peer_addr)) {
                                                debug!("failed to create accepted stream due to channel failure");

                                                // remove it from session
                                                sessions.close_peer(peer_addr);
                                                continue;
                                            }
                                        } else {
                                            let session_conv = s.conv().await;
                                            if session_conv != conv {
                                                debug!("received peer: {} with conv: {} not match with session conv: {}",
                                                       peer_addr,
                                                       conv,
                                                       session_conv);
                                                continue;
                                            }
                                        }

                                        s
                                    },
                                    Err(err) => {
                                        error!("failed to create session, error: {}, peer: {}, conv: {}", err, peer_addr, conv);
                                        continue;
                                    }
                                };

                                // let mut kcp = session.kcp_socket().lock().await;
                                // if let Err(err) = kcp.input(packet) {
                                //     error!("kcp.input failed, peer: {}, conv: {}, error: {}, packet: {:?}", peer_addr, conv, err, ByteStr::new(packet));
                                // }
                                session.input(packet).await;
                            }
                        }
                    }
                }
            }
        });

        Ok(KcpListener {
            udp: server_udp,
            accept_rx,
            task_watcher,
        })
    }

    pub async fn accept(&mut self) -> KcpResult<(KcpStream, SocketAddr)> {
        match self.accept_rx.recv().await {
            Some(s) => Ok(s),
            None => Err(KcpError::IoError(io::Error::new(
                ErrorKind::Other,
                "accept channel closed unexpectly",
            ))),
        }
    }

    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.udp.local_addr()
    }
}

#[cfg(test)]
mod test {
    use super::KcpListener;
    use crate::{config::KcpConfig, stream::KcpStream};
    use futures::future;

    #[tokio::test]
    async fn multi_echo() {
        let _ = env_logger::try_init();

        let mut listener = KcpListener::bind(KcpConfig::default(), "127.0.0.1:0").await.unwrap();
        let server_addr = listener.local_addr().unwrap();

        tokio::spawn(async move {
            loop {
                let (mut stream, _) = listener.accept().await.unwrap();

                tokio::spawn(async move {
                    let mut buffer = [0u8; 8192];
                    while let Ok(n) = stream.recv(&mut buffer).await {
                        let data = &buffer[..n];

                        let mut sent = 0;
                        while sent < n {
                            let sn = stream.send(&data[sent..]).await.unwrap();
                            sent += sn;
                        }
                    }
                });
            }
        });

        let mut vfut = Vec::new();

        for _ in 0..10 {
            vfut.push(async move {
                let mut stream = KcpStream::connect(&KcpConfig::default(), server_addr).await.unwrap();

                for _ in 0..5 {
                    const SEND_BUFFER: &[u8] = b"HELLO WORLD";
                    assert_eq!(SEND_BUFFER.len(), stream.send(SEND_BUFFER).await.unwrap());

                    let mut buffer = [0u8; 1024];
                    let n = stream.recv(&mut buffer).await.unwrap();
                    assert_eq!(SEND_BUFFER, &buffer[..n]);
                }
            });
        }

        future::join_all(vfut).await;
    }
}
<commit_msg>flushing stream write IO in test<commit_after>use std::{
    io::{self, ErrorKind},
    net::SocketAddr,
    sync::Arc,
    time::Duration,
};

use byte_string::ByteStr;
use kcp::{Error as KcpError, KcpResult};
use log::{debug, error, trace};
use tokio::{
    net::{ToSocketAddrs, UdpSocket},
    sync::mpsc,
    task::JoinHandle,
    time,
};

use crate::{config::KcpConfig, session::KcpSessionManager, stream::KcpStream};

pub struct KcpListener {
    udp: Arc<UdpSocket>,
    accept_rx: mpsc::Receiver<(KcpStream, SocketAddr)>,
    task_watcher: JoinHandle<()>,
}

impl Drop for KcpListener {
    fn drop(&mut self) {
        self.task_watcher.abort();
    }
}

impl KcpListener {
    pub async fn bind<A: ToSocketAddrs>(config: KcpConfig, addr: A) -> KcpResult<KcpListener> {
        let udp = UdpSocket::bind(addr).await?;
        let udp = Arc::new(udp);
        let server_udp = udp.clone();

        let (accept_tx, accept_rx) = mpsc::channel(1024 /* backlogs */);
        let task_watcher = tokio::spawn(async move {
            let (close_tx, mut close_rx) = mpsc::channel(64);

            let mut sessions = KcpSessionManager::new();
            let mut packet_buffer = [0u8; 65536];
            loop {
                tokio::select! {
                    peer_addr = close_rx.recv() => {
                        let peer_addr = peer_addr.expect("close_tx closed unexpectly");
                        sessions.close_peer(peer_addr);
                        trace!("session peer_addr: {} removed", peer_addr);
                    }

                    recv_res = udp.recv_from(&mut packet_buffer) => {
                        match recv_res {
                            Err(err) => {
                                error!("udp.recv_from failed, error: {}", err);
                                time::sleep(Duration::from_secs(1)).await;
                            }
                            Ok((n, peer_addr)) => {
                                let packet = &mut packet_buffer[..n];

                                log::trace!("received peer: {}, {:?}", peer_addr, ByteStr::new(packet));

                                let mut conv = kcp::get_conv(packet);
                                if conv == 0 {
                                    // Allocate a conv for client.
                                    conv = sessions.alloc_conv();
                                    debug!("allocate {} conv for peer: {}", conv, peer_addr);

                                    kcp::set_conv(packet, conv);
                                }

                                let sn = kcp::get_sn(packet);

                                let session = match sessions.get_or_create(&config, conv, sn, &udp, peer_addr, &close_tx).await {
                                    Ok((s, created)) => {
                                        if created {
                                            // Created a new session, constructed a new accepted client
                                            let stream = KcpStream::with_session(s.clone());
                                            if let Err(..) = accept_tx.try_send((stream, peer_addr)) {
                                                debug!("failed to create accepted stream due to channel failure");

                                                // remove it from session
                                                sessions.close_peer(peer_addr);
                                                continue;
                                            }
                                        } else {
                                            let session_conv = s.conv().await;
                                            if session_conv != conv {
                                                debug!("received peer: {} with conv: {} not match with session conv: {}",
                                                       peer_addr,
                                                       conv,
                                                       session_conv);
                                                continue;
                                            }
                                        }

                                        s
                                    },
                                    Err(err) => {
                                        error!("failed to create session, error: {}, peer: {}, conv: {}", err, peer_addr, conv);
                                        continue;
                                    }
                                };

                                // let mut kcp = session.kcp_socket().lock().await;
                                // if let Err(err) = kcp.input(packet) {
                                //     error!("kcp.input failed, peer: {}, conv: {}, error: {}, packet: {:?}", peer_addr, conv, err, ByteStr::new(packet));
                                // }
                                session.input(packet).await;
                            }
                        }
                    }
                }
            }
        });

        Ok(KcpListener {
            udp: server_udp,
            accept_rx,
            task_watcher,
        })
    }

    pub async fn accept(&mut self) -> KcpResult<(KcpStream, SocketAddr)> {
        match self.accept_rx.recv().await {
            Some(s) => Ok(s),
            None => Err(KcpError::IoError(io::Error::new(
                ErrorKind::Other,
                "accept channel closed unexpectly",
            ))),
        }
    }

    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.udp.local_addr()
    }
}

#[cfg(test)]
mod test {
    use super::KcpListener;
    use crate::{config::KcpConfig, stream::KcpStream};
    use futures::future;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    #[tokio::test]
    async fn multi_echo() {
        let _ = env_logger::try_init();

        let config = KcpConfig::default();

        let mut listener = KcpListener::bind(config.clone(), "127.0.0.1:0").await.unwrap();
        let server_addr = listener.local_addr().unwrap();

        tokio::spawn(async move {
            loop {
                let (mut stream, _) = listener.accept().await.unwrap();

                tokio::spawn(async move {
                    let mut buffer = [0u8; 8192];
                    while let Ok(n) = stream.read(&mut buffer).await {
                        if n == 0 {
                            break;
                        }

                        let data = &buffer[..n];
                        stream.write_all(data).await.unwrap();
                        stream.flush().await.unwrap();
                    }
                });
            }
        });

        let mut vfut = Vec::new();

        for _ in 0..100 {
            vfut.push(async move {
                let mut stream = KcpStream::connect(&config, server_addr).await.unwrap();

                for _ in 0..20 {
                    const SEND_BUFFER: &[u8] = b"HELLO WORLD";
                    stream.write_all(SEND_BUFFER).await.unwrap();
                    stream.flush().await.unwrap();

                    let mut buffer = [0u8; 1024];
                    let n = stream.recv(&mut buffer).await.unwrap();
                    assert_eq!(SEND_BUFFER, &buffer[..n]);
                }
            });
        }

        future::join_all(vfut).await;
    }
}
<|endoftext|>