<commit_before>//! An implementation of a buddy allocator.

use core::ptr;
use core::mem;
use core::cmp;

use Allocator;

use types::{UnsafeList, Link, PhantomNode};

pub type FreeBlock = PhantomNode;

#[allow(dead_code)]
pub struct Buddy<'a> {
    min_block_size: usize,
    min_block_order: u32,
    max_order: u32,
    heap_base: *mut u8,
    free_lists: &'a mut [UnsafeList<FreeBlock>],
}

impl<'a> Buddy<'a> {
    pub unsafe fn new(mut r_start: usize, r_size: usize, min_block_size: usize,
                      max_order: u32,
                      free_lists: &'a mut [UnsafeList<FreeBlock>]) -> Self {
        let max_block_size = min_block_size * 2usize.pow(max_order);
        let r_limit = r_start + r_size;

        assert!(min_block_size > ::core::mem::size_of::<FreeBlock>());

        let mut ret = Buddy {
            min_block_size: min_block_size,
            min_block_order: Buddy::log2(min_block_size) as u32,
            max_order: max_order,
            heap_base: r_start as *mut u8,
            free_lists: free_lists,
        };

        // This algorithm only deals with max block at init, so return an empty
        // allocator if the region is not big enough
        if r_size < max_block_size {
            return ret;
        }

        while r_start < r_limit {
            ret.add_block(max_order, r_start as *mut u8);

            r_start += max_block_size;
        }

        ret
    }

    #[allow(exceeding_bitshifts)]
    pub fn next_power_of_two(mut num: usize) -> usize {
        if num == 0 {
            return 1;
        }

        num -= 1;

        num |= num >> 1;
        num |= num >> 2;
        num |= num >> 4;
        num |= num >> 8;
        num |= num >> 16;

        if mem::size_of::<usize>() == mem::size_of::<u64>() {
            num |= num >> 32;
        }

        num + 1
    }

    pub fn log2(mut num: usize) -> usize {
        let mut res = 0;

        loop {
            num >>= 1;

            if num == 0 {
                break;
            }

            res += 1;
        }

        res
    }

    unsafe fn add_block(&mut self, order: u32, start: *mut u8) {
        let link = Link::some(start as *mut FreeBlock);

        self.free_lists[order as usize].push_front(link);
    }

    unsafe fn split_block(&mut self, block: *mut u8, mut order: u32,
                          target_order: u32) {
        while order > target_order {
            order -= 1;

            let buddy_offset = self.get_size_from_order(order);
            let buddy_ptr = block.offset(buddy_offset as isize);

            self.add_block(order, buddy_ptr);
        }
    }

    unsafe fn find_and_pop_buddy(&mut self, ptr: *mut u8,
                                 order: u32) -> *mut u8{
        // Max order blocks are not merged
        if order == self.max_order {
            return ptr::null_mut();
        }

        let size = self.get_size_from_order(order);
        let ptr_from_base = ptr as usize - self.heap_base as usize;
        let buddy_ptr = self.heap_base.offset((ptr_from_base ^ size) as isize);
        let mut cursor = self.free_lists[order as usize].cursor();
        let mut found = false;

        loop {
            match cursor.as_ref() {
                None => break,
                Some(blk) => {
                    if blk as *const FreeBlock as *const u8 == buddy_ptr {
                        found = true;
                        break;
                    }
                }
            }

            cursor.next();
        }

        if found {
            cursor.remove();
            return buddy_ptr
        }

        ptr::null_mut()
    }

    fn get_order_from_size(&self, mut size: usize, _align: usize) -> u32 {
        size = Buddy::next_power_of_two(size);

        size = cmp::max(size, self.min_block_size);

        Buddy::log2(size) as u32 - self.min_block_order
    }

    fn get_size_from_order(&self, order: u32) -> usize {
        self.min_block_size * 2usize.pow(order)
    }
}

impl<'a> Allocator for Buddy<'a> {
    unsafe fn allocate(&mut self, size: usize, align: usize) -> *mut u8 {
        let order = self.get_order_from_size(size, align);

        if order > self.max_order {
            return ptr::null_mut();
        }

        for i in order..(self.max_order + 1) {
            let mut tmp = self.free_lists[i as usize].pop_front();

            if let Some(block) = tmp.as_mut() {
                let ptr = block as *mut FreeBlock as *mut u8;

                if i > order {
                    self.split_block(ptr, i, order);
                }

                return ptr
            }
        }

        ptr::null_mut()
    }

    unsafe fn deallocate(&mut self, mut ptr: *mut u8, old_size: usize,
                         align: usize) {
        let mut order = self.get_order_from_size(old_size, align);

        loop {
            let buddy = self.find_and_pop_buddy(ptr, order);

            if buddy == ptr::null_mut() {
                break;
            }

            ptr = cmp::min(ptr, buddy);

            order += 1;
        }

        self.add_block(order, ptr);
    }
}

#[cfg(test)]
mod test {
    use core::mem;
    use core::ptr;

    use Allocator;

    use super::{Buddy, FreeBlock};

    use types::UnsafeList;

    const HEAP_ALIGN: usize = 4096;
    const HEAP_SIZE: usize = 4096;

    // HEAP_ORDER = 5 & MIN_BLOCK_SIZE = 32 => max alloc 1024
    const HEAP_ORDER: u32 = 5;
    const MIN_BLOCK_SIZE: usize = 32;

    type FreeList = UnsafeList<FreeBlock>;

    extern "C" {
        fn memalign(alignment: usize, size: usize) -> *mut u8;
        fn free(ptr: *mut u8);
    }

    #[test]
    fn test_buddy_next_power_of_two() {
        assert_eq!(Buddy::next_power_of_two(0), 1);
        assert_eq!(Buddy::next_power_of_two(2), 2);
        assert_eq!(Buddy::next_power_of_two(3), 4);
        assert_eq!(Buddy::next_power_of_two(5678), 8192);
        assert_eq!(Buddy::next_power_of_two(8192), 8192);
    }

    #[test]
    fn test_buddy_log2() {
        assert_eq!(Buddy::log2(0), 0);
        assert_eq!(Buddy::log2(1), 0);
        assert_eq!(Buddy::log2(2), 1);
        assert_eq!(Buddy::log2(4), 2);
        assert_eq!(Buddy::log2(8), 3);
        assert_eq!(Buddy::log2(16), 4);
        assert_eq!(Buddy::log2(0x87654321), 31);
    }

    #[test]
    fn test_buddy_alloc_dealloc() {
        unsafe {
            let heap = memalign(HEAP_ALIGN, HEAP_SIZE);
            let mut free_lists: [_; (HEAP_ORDER + 1) as usize];

            free_lists = mem::uninitialized();

            for i in 0..(HEAP_ORDER + 1) {
                free_lists[i as usize] = FreeList::new();
            }

            let max_size = MIN_BLOCK_SIZE * 2usize.pow(HEAP_ORDER);

            let mut alloc = Buddy::new(heap as usize, HEAP_SIZE, MIN_BLOCK_SIZE,
                                       HEAP_ORDER, &mut free_lists[..]);

            // Allocation is too big
            assert_eq!(alloc.allocate(max_size + 1, 1), ptr::null_mut());

            {
                let max_blk = alloc.allocate(max_size, 1);
                let last_blk_offset = ((HEAP_SIZE / max_size) - 1) * max_size;

                // Due to Buddy::new using push front, the first allocated block
                // is gonna be the last pushed
                assert_eq!(max_blk, heap.offset(last_blk_offset as isize));

                alloc.deallocate(max_blk, max_size, 1);
            }

            let blk_32_1 = alloc.allocate(32, 1);
            let blk_32_2 = alloc.allocate(32, 1);

            assert_eq!(blk_32_1.offset(32), blk_32_2);

            let blk_64_1 = alloc.allocate(64, 1);

            assert_eq!(blk_32_1.offset(64), blk_64_1);

            alloc.deallocate(blk_32_1, 32, 1);
            alloc.deallocate(blk_32_2, 32, 1);

            let blk_32_1_1 = alloc.allocate(32, 1);
            let blk_32_2_1 = alloc.allocate(32, 1);

            assert_eq!(blk_32_1_1, blk_32_1);
            assert_eq!(blk_32_2_1, blk_32_2);

            alloc.deallocate(blk_32_2_1, 32, 1);
            alloc.deallocate(blk_32_1_1, 32, 1);

            let blk_64_2 = alloc.allocate(64, 1);

            assert_eq!(blk_64_2, blk_32_1);

            let blk_128 = alloc.allocate(128, 1);

            assert_eq!(blk_128, blk_64_1.offset(64));

            alloc.deallocate(blk_64_1, 64, 1);
            alloc.deallocate(blk_64_2, 64, 1);
            alloc.deallocate(blk_128, 128, 1);

            free(heap);
        }
    }
}
<commit_msg>libheap: Correct minimum block size for the buddy<commit_after>//! An implementation of a buddy allocator.

use core::ptr;
use core::mem;
use core::cmp;

use Allocator;

use types::{UnsafeList, Link, PhantomNode};

pub type FreeBlock = PhantomNode;

#[allow(dead_code)]
pub struct Buddy<'a> {
    min_block_size: usize,
    min_block_order: u32,
    max_order: u32,
    heap_base: *mut u8,
    free_lists: &'a mut [UnsafeList<FreeBlock>],
}

impl<'a> Buddy<'a> {
    pub unsafe fn new(mut r_start: usize, r_size: usize, min_block_size: usize,
                      max_order: u32,
                      free_lists: &'a mut [UnsafeList<FreeBlock>]) -> Self {
        let max_block_size = min_block_size * 2usize.pow(max_order);
        let r_limit = r_start + r_size;

        assert!(min_block_size >= ::core::mem::size_of::<FreeBlock>());

        let mut ret = Buddy {
            min_block_size: min_block_size,
            min_block_order: Buddy::log2(min_block_size) as u32,
            max_order: max_order,
            heap_base: r_start as *mut u8,
            free_lists: free_lists,
        };

        // This algorithm only deals with max block at init, so return an empty
        // allocator if the region is not big enough
        if r_size < max_block_size {
            return ret;
        }

        while r_start < r_limit {
            ret.add_block(max_order, r_start as *mut u8);

            r_start += max_block_size;
        }

        ret
    }

    #[allow(exceeding_bitshifts)]
    pub fn next_power_of_two(mut num: usize) -> usize {
        if num == 0 {
            return 1;
        }

        num -= 1;

        num |= num >> 1;
        num |= num >> 2;
        num |= num >> 4;
        num |= num >> 8;
        num |= num >> 16;

        if mem::size_of::<usize>() == mem::size_of::<u64>() {
            num |= num >> 32;
        }

        num + 1
    }

    pub fn log2(mut num: usize) -> usize {
        let mut res = 0;

        loop {
            num >>= 1;

            if num == 0 {
                break;
            }

            res += 1;
        }

        res
    }

    unsafe fn add_block(&mut self, order: u32, start: *mut u8) {
        let link = Link::some(start as *mut FreeBlock);

        self.free_lists[order as usize].push_front(link);
    }

    unsafe fn split_block(&mut self, block: *mut u8, mut order: u32,
                          target_order: u32) {
        while order > target_order {
            order -= 1;

            let buddy_offset = self.get_size_from_order(order);
            let buddy_ptr = block.offset(buddy_offset as isize);

            self.add_block(order, buddy_ptr);
        }
    }

    unsafe fn find_and_pop_buddy(&mut self, ptr: *mut u8,
                                 order: u32) -> *mut u8{
        // Max order blocks are not merged
        if order == self.max_order {
            return ptr::null_mut();
        }

        let size = self.get_size_from_order(order);
        let ptr_from_base = ptr as usize - self.heap_base as usize;
        let buddy_ptr = self.heap_base.offset((ptr_from_base ^ size) as isize);
        let mut cursor = self.free_lists[order as usize].cursor();
        let mut found = false;

        loop {
            match cursor.as_ref() {
                None => break,
                Some(blk) => {
                    if blk as *const FreeBlock as *const u8 == buddy_ptr {
                        found = true;
                        break;
                    }
                }
            }

            cursor.next();
        }

        if found {
            cursor.remove();
            return buddy_ptr
        }

        ptr::null_mut()
    }

    fn get_order_from_size(&self, mut size: usize, _align: usize) -> u32 {
        size = Buddy::next_power_of_two(size);

        size = cmp::max(size, self.min_block_size);

        Buddy::log2(size) as u32 - self.min_block_order
    }

    fn get_size_from_order(&self, order: u32) -> usize {
        self.min_block_size * 2usize.pow(order)
    }
}

impl<'a> Allocator for Buddy<'a> {
    unsafe fn allocate(&mut self, size: usize, align: usize) -> *mut u8 {
        let order = self.get_order_from_size(size, align);

        if order > self.max_order {
            return ptr::null_mut();
        }

        for i in order..(self.max_order + 1) {
            let mut tmp = self.free_lists[i as usize].pop_front();

            if let Some(block) = tmp.as_mut() {
                let ptr = block as *mut FreeBlock as *mut u8;

                if i > order {
                    self.split_block(ptr, i, order);
                }

                return ptr
            }
        }

        ptr::null_mut()
    }

    unsafe fn deallocate(&mut self, mut ptr: *mut u8, old_size: usize,
                         align: usize) {
        let mut order = self.get_order_from_size(old_size, align);

        loop {
            let buddy = self.find_and_pop_buddy(ptr, order);

            if buddy == ptr::null_mut() {
                break;
            }

            ptr = cmp::min(ptr, buddy);

            order += 1;
        }

        self.add_block(order, ptr);
    }
}

#[cfg(test)]
mod test {
    use core::mem;
    use core::ptr;

    use Allocator;

    use super::{Buddy, FreeBlock};

    use types::UnsafeList;

    const HEAP_ALIGN: usize = 4096;
    const HEAP_SIZE: usize = 4096;

    // HEAP_ORDER = 5 & MIN_BLOCK_SIZE = 32 => max alloc 1024
    const HEAP_ORDER: u32 = 5;
    const MIN_BLOCK_SIZE: usize = 32;

    type FreeList = UnsafeList<FreeBlock>;

    extern "C" {
        fn memalign(alignment: usize, size: usize) -> *mut u8;
        fn free(ptr: *mut u8);
    }

    #[test]
    fn test_buddy_next_power_of_two() {
        assert_eq!(Buddy::next_power_of_two(0), 1);
        assert_eq!(Buddy::next_power_of_two(2), 2);
        assert_eq!(Buddy::next_power_of_two(3), 4);
        assert_eq!(Buddy::next_power_of_two(5678), 8192);
        assert_eq!(Buddy::next_power_of_two(8192), 8192);
    }

    #[test]
    fn test_buddy_log2() {
        assert_eq!(Buddy::log2(0), 0);
        assert_eq!(Buddy::log2(1), 0);
        assert_eq!(Buddy::log2(2), 1);
        assert_eq!(Buddy::log2(4), 2);
        assert_eq!(Buddy::log2(8), 3);
        assert_eq!(Buddy::log2(16), 4);
        assert_eq!(Buddy::log2(0x87654321), 31);
    }

    #[test]
    fn test_buddy_alloc_dealloc() {
        unsafe {
            let heap = memalign(HEAP_ALIGN, HEAP_SIZE);
            let mut free_lists: [_; (HEAP_ORDER + 1) as usize];

            free_lists = mem::uninitialized();

            for i in 0..(HEAP_ORDER + 1) {
                free_lists[i as usize] = FreeList::new();
            }

            let max_size = MIN_BLOCK_SIZE * 2usize.pow(HEAP_ORDER);

            let mut alloc = Buddy::new(heap as usize, HEAP_SIZE, MIN_BLOCK_SIZE,
                                       HEAP_ORDER, &mut free_lists[..]);

            // Allocation is too big
            assert_eq!(alloc.allocate(max_size + 1, 1), ptr::null_mut());

            {
                let max_blk = alloc.allocate(max_size, 1);
                let last_blk_offset = ((HEAP_SIZE / max_size) - 1) * max_size;

                // Due to Buddy::new using push front, the first allocated block
                // is gonna be the last pushed
                assert_eq!(max_blk, heap.offset(last_blk_offset as isize));

                alloc.deallocate(max_blk, max_size, 1);
            }

            let blk_32_1 = alloc.allocate(32, 1);
            let blk_32_2 = alloc.allocate(32, 1);

            assert_eq!(blk_32_1.offset(32), blk_32_2);

            let blk_64_1 = alloc.allocate(64, 1);

            assert_eq!(blk_32_1.offset(64), blk_64_1);

            alloc.deallocate(blk_32_1, 32, 1);
            alloc.deallocate(blk_32_2, 32, 1);

            let blk_32_1_1 = alloc.allocate(32, 1);
            let blk_32_2_1 = alloc.allocate(32, 1);

            assert_eq!(blk_32_1_1, blk_32_1);
            assert_eq!(blk_32_2_1, blk_32_2);

            alloc.deallocate(blk_32_2_1, 32, 1);
            alloc.deallocate(blk_32_1_1, 32, 1);

            let blk_64_2 = alloc.allocate(64, 1);

            assert_eq!(blk_64_2, blk_32_1);

            let blk_128 = alloc.allocate(128, 1);

            assert_eq!(blk_128, blk_64_1.offset(64));

            alloc.deallocate(blk_64_1, 64, 1);
            alloc.deallocate(blk_64_2, 64, 1);
            alloc.deallocate(blk_128, 128, 1);

            free(heap);
        }
    }
}
<|endoftext|>


<commit_before>use nom::{
  Consumer, ConsumerState,
  ErrorKind,
  HexDisplay,
  Producer, FileProducer,
  Input, IResult,
  Move, Needed,
};

use metadata;

use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};
use utility::resize_producer;

use std::io;
use std::io::{Error, Result};

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

pub struct Stream {
  info: StreamInfo,
  pub metadata: Vec<Metadata>,
  pub frames: Vec<Frame>,
  state: ParserState,
  consumer_state: ConsumerState<(), ErrorKind, Move>,
}

named!(pub stream_parser <&[u8], Stream>,
  chain!(
    blocks: metadata_parser ~
    frames: many1!(apply!(frame_parser, &blocks.0)),
    move|| {
      let consumed = Move::Consume(0);

      Stream {
        info: blocks.0,
        metadata: blocks.1,
        frames: frames,
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      }
    }
  )
);

impl Stream {
  pub fn new() -> Stream {
    let consumed = Move::Consume(0);

    Stream {
      info: StreamInfo::new(),
      metadata: Vec::new(),
      frames: Vec::new(),
      state: ParserState::Marker,
      consumer_state: ConsumerState::Continue(consumed),
    }
  }

  pub fn info(&self) -> StreamInfo {
    self.info
  }

  pub fn from_file(filename: &str) -> Result<Stream> {
    FileProducer::new(filename, 1024).and_then(|mut producer| {
      let consumed        = Move::Consume(0);
      let mut buffer_size = 1024;
      let mut is_error    = false;
      let mut stream      = Stream {
        info: StreamInfo::new(),
        metadata: Vec::new(),
        frames: Vec::new(),
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      };

      loop {
        match *producer.apply(&mut stream) {
          ConsumerState::Done(_, _)      => break,
          ConsumerState::Continue(await) => {
            let result = resize_producer(&mut producer, &await, buffer_size);

            if let Some(size) = result {
              buffer_size = size;
            }

            continue;
          }
          ConsumerState::Error(_)        => {
            is_error = true;

            break;
          }
        }
      }

      if !is_error {
        Ok(stream)
      } else {
        let error_str = format!("parser: couldn't parse the given file {}",
                                filename);

        Err(Error::new(io::ErrorKind::InvalidData, error_str))
      }
    })
  }

  fn handle_marker<'a>(&mut self, input: &'a [u8]) -> IResult<&'a [u8], ()> {
    let kind = nom::ErrorKind::Custom(0);

    match tag!(input, "fLaC") {
      IResult::Done(i, _)    => {
        self.state = ParserState::Metadata;

        IResult::Error(Err::Position(kind, i))
      }
      IResult::Error(_)      => IResult::Error(Err::Code(kind)),
      IResult::Incomplete(n) => IResult::Incomplete(n),
    }
  }

  fn handle_metadata(&mut self, input: &[u8]) {
    match metadata::block(input) {
      IResult::Done(i, block) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);
        let is_last  = block.is_last;

        if let metadata::Data::StreamInfo(info) = block.data {
          self.info = info;
        } else {
          self.metadata.push(block);
        }

        if is_last {
          self.state = ParserState::Frame;
        }

        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)       => {
        let kind = ErrorKind::Custom(1);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(s)  => {
        let size = if let Needed::Size(length) = s {
          length
        } else {
          1024
        };
        let needed = Move::Await(Needed::Size(size));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }

  fn handle_frame(&mut self, input: &[u8]) {
    match frame_parser(input, &self.info) {
      IResult::Done(i, frame) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);

        self.frames.push(frame);

        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)       => {
        let kind = ErrorKind::Custom(2);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(s)  => {
        let size = if let Needed::Size(length) = s {
          length
        } else {
          self.info.max_frame_size as usize
        };
        let needed = Move::Await(Needed::Size(size));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }
}

impl<'a> Consumer<&'a [u8], (), ErrorKind, Move> for Stream {
  fn state(&self) -> &ConsumerState<(), ErrorKind, Move> {
    &self.consumer_state
  }

  fn handle(&mut self, input: Input<&'a [u8]>)
            -> &ConsumerState<(), ErrorKind, Move> {
    match input {
      Input::Element(i) | Input::Eof(Some(i)) => {
        match self.state {
          ParserState::Marker   => self.handle_marker(i),
          ParserState::Metadata => self.handle_metadata(i),
          ParserState::Frame    => self.handle_frame(i),
        }
      }
      Input::Empty | Input::Eof(None)         => {
        self.consumer_state = match self.state {
          ParserState::Marker   => ConsumerState::Error(ErrorKind::Custom(0)),
          ParserState::Metadata => ConsumerState::Error(ErrorKind::Custom(1)),
          ParserState::Frame    => ConsumerState::Done(Move::Consume(0), ()),
        };
      }
    }

    &self.consumer_state
  }
}
<commit_msg>Rewrite handle_metadata use with StreamProducer<commit_after>use nom::{
  Consumer, ConsumerState,
  ErrorKind,
  HexDisplay,
  Producer, FileProducer,
  Input, IResult,
  Move, Needed,
};

use metadata;

use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};
use utility::resize_producer;

use std::io;
use std::io::{Error, Result};

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

pub struct Stream {
  info: StreamInfo,
  pub metadata: Vec<Metadata>,
  pub frames: Vec<Frame>,
  state: ParserState,
  consumer_state: ConsumerState<(), ErrorKind, Move>,
}

named!(pub stream_parser <&[u8], Stream>,
  chain!(
    blocks: metadata_parser ~
    frames: many1!(apply!(frame_parser, &blocks.0)),
    move|| {
      let consumed = Move::Consume(0);

      Stream {
        info: blocks.0,
        metadata: blocks.1,
        frames: frames,
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      }
    }
  )
);

impl Stream {
  pub fn new() -> Stream {
    let consumed = Move::Consume(0);

    Stream {
      info: StreamInfo::new(),
      metadata: Vec::new(),
      frames: Vec::new(),
      state: ParserState::Marker,
      consumer_state: ConsumerState::Continue(consumed),
    }
  }

  pub fn info(&self) -> StreamInfo {
    self.info
  }

  pub fn from_file(filename: &str) -> Result<Stream> {
    FileProducer::new(filename, 1024).and_then(|mut producer| {
      let consumed        = Move::Consume(0);
      let mut buffer_size = 1024;
      let mut is_error    = false;
      let mut stream      = Stream {
        info: StreamInfo::new(),
        metadata: Vec::new(),
        frames: Vec::new(),
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      };

      loop {
        match *producer.apply(&mut stream) {
          ConsumerState::Done(_, _)      => break,
          ConsumerState::Continue(await) => {
            let result = resize_producer(&mut producer, &await, buffer_size);

            if let Some(size) = result {
              buffer_size = size;
            }

            continue;
          }
          ConsumerState::Error(_)        => {
            is_error = true;

            break;
          }
        }
      }

      if !is_error {
        Ok(stream)
      } else {
        let error_str = format!("parser: couldn't parse the given file {}",
                                filename);

        Err(Error::new(io::ErrorKind::InvalidData, error_str))
      }
    })
  }

  fn handle_marker<'a>(&mut self, input: &'a [u8]) -> IResult<&'a [u8], ()> {
    let kind = nom::ErrorKind::Custom(0);

    match tag!(input, "fLaC") {
      IResult::Done(i, _)    => {
        self.state = ParserState::Metadata;

        IResult::Error(Err::Position(kind, i))
      }
      IResult::Error(_)      => IResult::Error(Err::Code(kind)),
      IResult::Incomplete(n) => IResult::Incomplete(n),
    }
  }

  fn handle_metadata<'a>(&mut self, input: &'a [u8])
                         -> IResult<&'a [u8], ()> {
    let kind = nom::ErrorKind::Custom(1);

    match metadata_parser(input) {
      IResult::Done(i, block) => {
        let is_last = block.is_last;

        if let metadata::Data::StreamInfo(info) = block.data {
          self.info = info;
        } else {
          self.metadata.push(block);
        }

        if is_last {
          self.state = ParserState::Frame;
        }

        IResult::Error(Err::Position(kind, i))
      }
      IResult::Error(_)      => IResult::Error(Err::Code(kind)),
      IResult::Incomplete(n) => IResult::Incomplete(n),
    }
  }

  fn handle_frame(&mut self, input: &[u8]) {
    match frame_parser(input, &self.info) {
      IResult::Done(i, frame) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);

        self.frames.push(frame);

        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)       => {
        let kind = ErrorKind::Custom(2);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(s)  => {
        let size = if let Needed::Size(length) = s {
          length
        } else {
          self.info.max_frame_size as usize
        };
        let needed = Move::Await(Needed::Size(size));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }
}

impl<'a> Consumer<&'a [u8], (), ErrorKind, Move> for Stream {
  fn state(&self) -> &ConsumerState<(), ErrorKind, Move> {
    &self.consumer_state
  }

  fn handle(&mut self, input: Input<&'a [u8]>)
            -> &ConsumerState<(), ErrorKind, Move> {
    match input {
      Input::Element(i) | Input::Eof(Some(i)) => {
        match self.state {
          ParserState::Marker   => self.handle_marker(i),
          ParserState::Metadata => self.handle_metadata(i),
          ParserState::Frame    => self.handle_frame(i),
        }
      }
      Input::Empty | Input::Eof(None)         => {
        self.consumer_state = match self.state {
          ParserState::Marker   => ConsumerState::Error(ErrorKind::Custom(0)),
          ParserState::Metadata => ConsumerState::Error(ErrorKind::Custom(1)),
          ParserState::Frame    => ConsumerState::Done(Move::Consume(0), ()),
        };
      }
    }

    &self.consumer_state
  }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::ErrorKind::WouldBlock;
#[cfg(any(feature = "ssl", feature = "nativetls"))]
use std::mem::replace;
use std::net::SocketAddr;

use bytes::{Buf, BufMut};
use mio::tcp::TcpStream;
#[cfg(feature = "ssl")]
use openssl::ssl::Error as SslError;
#[cfg(feature = "nativetls")]
use native_tls::{HandshakeError, MidHandshakeTlsStream as MidHandshakeSslStream, TlsStream as SslStream};
use bytes::{Buf, BufMut};

use result::{Error, Kind, Result};

fn map_non_block<T>(res: io::Result<T>) -> io::Result<Option<T>> {
    match res {
        Ok(value) => Ok(Some(value)),
        Err(err) => {
            if let WouldBlock = err.kind() {
                Ok(None)
            } else {
                Err(err)
            }
        }
    }
}

pub trait TryReadBuf: io::Read {
    fn try_read_buf<B: BufMut>(&mut self, buf: &mut B) -> io::Result<Option<usize>>
    where
        Self: Sized,
    {
        // Reads the length of the slice supplied by buf.mut_bytes into the buffer
        // This is not guaranteed to consume an entire datagram or segment.
        // If your protocol is msg based (instead of continuous stream) you should
        // ensure that your buffer is large enough to hold an entire segment (1532 bytes if not jumbo
        // frames)
        let res = map_non_block(self.read(unsafe { buf.bytes_mut() }));

        if let Ok(Some(cnt)) = res {
            unsafe {
                buf.advance_mut(cnt);
            }
        }

        res
    }
}

pub trait TryWriteBuf: io::Write {
    fn try_write_buf<B: Buf>(&mut self, buf: &mut B) -> io::Result<Option<usize>>
    where
        Self: Sized,
    {
        let res = map_non_block(self.write(buf.bytes()));

        if let Ok(Some(cnt)) = res {
            buf.advance(cnt);
        }

        res
    }
}

impl<T: io::Read> TryReadBuf for T {}
impl<T: io::Write> TryWriteBuf for T {}

use self::Stream::*;
pub enum Stream {
    Tcp(TcpStream),
    #[cfg(any(feature = "ssl", feature = "nativetls"))]
    Tls(TlsStream),
}

impl Stream {
    pub fn tcp(stream: TcpStream) -> Stream {
        Tcp(stream)
    }

    #[cfg(any(feature = "ssl", feature = "nativetls"))]
    pub fn tls(stream: MidHandshakeSslStream<TcpStream>) -> Stream {
        Tls(TlsStream::Handshake {
            sock: stream,
            negotiating: false,
        })
    }

    #[cfg(any(feature = "ssl", feature = "nativetls"))]
    pub fn tls_live(stream: SslStream<TcpStream>) -> Stream {
        Tls(TlsStream::Live(stream))
    }

    #[cfg(any(feature = "ssl", feature = "nativetls"))]
    pub fn is_tls(&self) -> bool {
        match *self {
            Tcp(_) => false,
            Tls(_) => true,
        }
    }

    pub fn evented(&self) -> &TcpStream {
        match *self {
            Tcp(ref sock) => sock,
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref inner) => inner.evented(),
        }
    }

    pub fn is_negotiating(&self) -> bool {
        match *self {
            Tcp(_) => false,
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref inner) => inner.is_negotiating(),
        }
    }

    pub fn clear_negotiating(&mut self) -> Result<()> {
        match *self {
            Tcp(_) => Err(Error::new(
                Kind::Internal,
                "Attempted to clear negotiating flag on non ssl connection.",
            )),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref mut inner) => inner.clear_negotiating(),
        }
    }

    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
        match *self {
            Tcp(ref sock) => sock.peer_addr(),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref inner) => inner.peer_addr(),
        }
    }

    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        match *self {
            Tcp(ref sock) => sock.local_addr(),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref inner) => inner.local_addr(),
        }
    }
}

impl io::Read for Stream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match *self {
            Tcp(ref mut sock) => sock.read(buf),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(TlsStream::Live(ref mut sock)) => sock.read(buf),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref mut tls_stream) => {
                trace!("Attempting to read ssl handshake.");
                match replace(tls_stream, TlsStream::Upgrading) {
                    TlsStream::Live(_) | TlsStream::Upgrading => unreachable!(),
                    TlsStream::Handshake {
                        sock,
                        mut negotiating,
                    } => match sock.handshake() {
                        Ok(mut sock) => {
                            trace!("Completed SSL Handshake");
                            let res = sock.read(buf);
                            *tls_stream = TlsStream::Live(sock);
                            res
                        }
                        #[cfg(feature = "ssl")]
                        Err(HandshakeError::SetupFailure(err)) => {
                            Err(io::Error::new(io::ErrorKind::Other, err))
                        }
                        #[cfg(feature = "ssl")]
                        Err(HandshakeError::Failure(mid))
                        | Err(HandshakeError::Interrupted(mid)) => {
                            let err = match *mid.error() {
                                SslError::WantWrite(_) => {
                                    negotiating = true;
                                    Err(io::Error::new(
                                        io::ErrorKind::WouldBlock,
                                        "SSL wants writing",
                                    ))
                                }
                                SslError::WantRead(_) => Err(io::Error::new(
                                    io::ErrorKind::WouldBlock,
                                    "SSL wants reading",
                                )),
                                SslError::Stream(ref e) => Err(From::from(e.kind())),
                                ref err => {
                                    Err(io::Error::new(io::ErrorKind::Other, format!("{}", err)))
                                }
                            };
                            *tls_stream = TlsStream::Handshake {
                                sock: mid,
                                negotiating: negotiating,
                            };
                            err
                        }
                        #[cfg(feature = "nativetls")]
                        Err(HandshakeError::Interrupted(mid)) => {
                            negotiating = true;
                            *tls_stream = TlsStream::Handshake {
                                sock: mid,
                                negotiating: negotiating,
                            };
                            Err(io::Error::new(
                                io::ErrorKind::WouldBlock,
                                "SSL would block",
                            ))
                        }
                        #[cfg(feature = "nativetls")]
                        Err(HandshakeError::Failure(err)) => {
                            Err(io::Error::new(
                                io::ErrorKind::Other,
                                format!("{}", err),
                            ))
                        }
                    },
                }
            }
        }
    }
}

impl io::Write for Stream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match *self {
            Tcp(ref mut sock) => sock.write(buf),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(TlsStream::Live(ref mut sock)) => sock.write(buf),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref mut tls_stream) => {
                trace!("Attempting to write ssl handshake.");
                match replace(tls_stream, TlsStream::Upgrading) {
                    TlsStream::Live(_) | TlsStream::Upgrading => unreachable!(),
                    TlsStream::Handshake {
                        sock,
                        mut negotiating,
                    } => {
                        match sock.handshake() {
                            Ok(mut sock) => {
                                trace!("Completed SSL Handshake");
                                let res = sock.write(buf);
                                *tls_stream = TlsStream::Live(sock);
                                res
                            }
                            #[cfg(feature = "ssl")]
                            Err(HandshakeError::SetupFailure(err)) => {
                                Err(io::Error::new(io::ErrorKind::Other, err))
                            }
                            #[cfg(feature = "ssl")]
                            Err(HandshakeError::Failure(mid))
                            | Err(HandshakeError::Interrupted(mid)) => {
                                let err = match *mid.error() {
                                    SslError::WantRead(_) => {
                                        negotiating = true;
                                        Err(io::Error::new(
                                            io::ErrorKind::WouldBlock,
                                            "SSL wants reading",
                                        ))
                                    }
                                    SslError::WantWrite(_) => {
                                        negotiating = false;
                                        Err(io::Error::new(
                                            io::ErrorKind::WouldBlock,
                                            "SSL wants writing",
                                        ))
                                    }
                                    SslError::Stream(ref e) => {
                                        negotiating = false;
                                        Err(From::from(e.kind()))
                                    }
                                    ref err => {
                                        negotiating = false;
                                        Err(io::Error::new(
                                            io::ErrorKind::Other,
                                            format!("{}", err),
                                        ))
                                    }
                                };
                                *tls_stream = TlsStream::Handshake {
                                    sock: mid,
                                    negotiating: negotiating,
                                };
                                err
                            }
                            #[cfg(feature = "nativetls")]
                            Err(HandshakeError::Interrupted(mid)) => {
                                negotiating = true;
                                *tls_stream = TlsStream::Handshake {
                                    sock: mid,
                                    negotiating: negotiating,
                                };
                                Err(io::Error::new(
                                    io::ErrorKind::WouldBlock,
                                    "SSL would block",
                                ))
                            }
                            #[cfg(feature = "nativetls")]
                            Err(HandshakeError::Failure(err)) => {
                                Err(io::Error::new(
                                    io::ErrorKind::Other,
                                    format!("{}", err),
                                ))
                            }
                        }
                    }
                }
            }
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match *self {
            Tcp(ref mut sock) => sock.flush(),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(TlsStream::Live(ref mut sock)) => sock.flush(),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(TlsStream::Handshake { ref mut sock, .. }) => sock.get_mut().flush(),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(TlsStream::Upgrading) => panic!("Tried to access actively upgrading TlsStream"),
        }
    }
}

#[cfg(any(feature = "ssl", feature = "nativetls"))]
pub enum TlsStream {
    Live(SslStream<TcpStream>),
    Handshake {
        sock: MidHandshakeSslStream<TcpStream>,
        negotiating: bool,
    },
    Upgrading,
}

#[cfg(any(feature = "ssl", feature = "nativetls"))]
impl TlsStream {
    pub fn evented(&self) -> &TcpStream {
        match *self {
            TlsStream::Live(ref sock) => sock.get_ref(),
            TlsStream::Handshake { ref sock, .. } => sock.get_ref(),
            TlsStream::Upgrading => panic!("Tried to access actively upgrading TlsStream"),
        }
    }

    pub fn is_negotiating(&self) -> bool {
        match *self {
            TlsStream::Live(_) => false,
            TlsStream::Handshake {
                sock: _,
                negotiating,
            } => negotiating,
            TlsStream::Upgrading => panic!("Tried to access actively upgrading TlsStream"),
        }
    }

    pub fn clear_negotiating(&mut self) -> Result<()> {
        match *self {
            TlsStream::Live(_) => Err(Error::new(
                Kind::Internal,
                "Attempted to clear negotiating flag on live ssl connection.",
            )),
            TlsStream::Handshake {
                sock: _,
                ref mut negotiating,
            } => Ok(*negotiating = false),
            TlsStream::Upgrading => panic!("Tried to access actively upgrading TlsStream"),
        }
    }

    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
        match *self {
            TlsStream::Live(ref sock) => sock.get_ref().peer_addr(),
            TlsStream::Handshake { ref sock, .. } => sock.get_ref().peer_addr(),
            TlsStream::Upgrading => panic!("Tried to access actively upgrading TlsStream"),
        }
    }

    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        match *self {
            TlsStream::Live(ref sock) => sock.get_ref().local_addr(),
            TlsStream::Handshake { ref sock, .. } => sock.get_ref().local_addr(),
            TlsStream::Upgrading => panic!("Tried to access actively upgrading TlsStream"),
        }
    }
}
<commit_msg>Update stream.rs<commit_after>use std::io;
use std::io::ErrorKind::WouldBlock;
#[cfg(any(feature = "ssl", feature = "nativetls"))]
use std::mem::replace;
use std::net::SocketAddr;

use bytes::{Buf, BufMut};
use mio::tcp::TcpStream;
#[cfg(feature = "ssl")]
use openssl::ssl::Error as SslError;
#[cfg(feature = "nativetls")]
use native_tls::{HandshakeError, MidHandshakeTlsStream as MidHandshakeSslStream, TlsStream as SslStream};

use result::{Error, Kind, Result};

fn map_non_block<T>(res: io::Result<T>) -> io::Result<Option<T>> {
    match res {
        Ok(value) => Ok(Some(value)),
        Err(err) => {
            if let WouldBlock = err.kind() {
                Ok(None)
            } else {
                Err(err)
            }
        }
    }
}

pub trait TryReadBuf: io::Read {
    fn try_read_buf<B: BufMut>(&mut self, buf: &mut B) -> io::Result<Option<usize>>
    where
        Self: Sized,
    {
        // Reads the length of the slice supplied by buf.mut_bytes into the buffer
        // This is not guaranteed to consume an entire datagram or segment.
        // If your protocol is msg based (instead of continuous stream) you should
        // ensure that your buffer is large enough to hold an entire segment (1532 bytes if not jumbo
        // frames)
        let res = map_non_block(self.read(unsafe { buf.bytes_mut() }));

        if let Ok(Some(cnt)) = res {
            unsafe {
                buf.advance_mut(cnt);
            }
        }

        res
    }
}

pub trait TryWriteBuf: io::Write {
    fn try_write_buf<B: Buf>(&mut self, buf: &mut B) -> io::Result<Option<usize>>
    where
        Self: Sized,
    {
        let res = map_non_block(self.write(buf.bytes()));

        if let Ok(Some(cnt)) = res {
            buf.advance(cnt);
        }

        res
    }
}

impl<T: io::Read> TryReadBuf for T {}
impl<T: io::Write> TryWriteBuf for T {}

use self::Stream::*;
pub enum Stream {
    Tcp(TcpStream),
    #[cfg(any(feature = "ssl", feature = "nativetls"))]
    Tls(TlsStream),
}

impl Stream {
    pub fn tcp(stream: TcpStream) -> Stream {
        Tcp(stream)
    }

    #[cfg(any(feature = "ssl", feature = "nativetls"))]
    pub fn tls(stream: MidHandshakeSslStream<TcpStream>) -> Stream {
        Tls(TlsStream::Handshake {
            sock: stream,
            negotiating: false,
        })
    }

    #[cfg(any(feature = "ssl", feature = "nativetls"))]
    pub fn tls_live(stream: SslStream<TcpStream>) -> Stream {
        Tls(TlsStream::Live(stream))
    }

    #[cfg(any(feature = "ssl", feature = "nativetls"))]
    pub fn is_tls(&self) -> bool {
        match *self {
            Tcp(_) => false,
            Tls(_) => true,
        }
    }

    pub fn evented(&self) -> &TcpStream {
        match *self {
            Tcp(ref sock) => sock,
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref inner) => inner.evented(),
        }
    }

    pub fn is_negotiating(&self) -> bool {
        match *self {
            Tcp(_) => false,
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref inner) => inner.is_negotiating(),
        }
    }

    pub fn clear_negotiating(&mut self) -> Result<()> {
        match *self {
            Tcp(_) => Err(Error::new(
                Kind::Internal,
                "Attempted to clear negotiating flag on non ssl connection.",
            )),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref mut inner) => inner.clear_negotiating(),
        }
    }

    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
        match *self {
            Tcp(ref sock) => sock.peer_addr(),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref inner) => inner.peer_addr(),
        }
    }

    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        match *self {
            Tcp(ref sock) => sock.local_addr(),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref inner) => inner.local_addr(),
        }
    }
}

impl io::Read for Stream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match *self {
            Tcp(ref mut sock) => sock.read(buf),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(TlsStream::Live(ref mut sock)) => sock.read(buf),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref mut tls_stream) => {
                trace!("Attempting to read ssl handshake.");
                match replace(tls_stream, TlsStream::Upgrading) {
                    TlsStream::Live(_) | TlsStream::Upgrading => unreachable!(),
                    TlsStream::Handshake {
                        sock,
                        mut negotiating,
                    } => match sock.handshake() {
                        Ok(mut sock) => {
                            trace!("Completed SSL Handshake");
                            let res = sock.read(buf);
                            *tls_stream = TlsStream::Live(sock);
                            res
                        }
                        #[cfg(feature = "ssl")]
                        Err(HandshakeError::SetupFailure(err)) => {
                            Err(io::Error::new(io::ErrorKind::Other, err))
                        }
                        #[cfg(feature = "ssl")]
                        Err(HandshakeError::Failure(mid))
                        | Err(HandshakeError::Interrupted(mid)) => {
                            let err = match *mid.error() {
                                SslError::WantWrite(_) => {
                                    negotiating = true;
                                    Err(io::Error::new(
                                        io::ErrorKind::WouldBlock,
                                        "SSL wants writing",
                                    ))
                                }
                                SslError::WantRead(_) => Err(io::Error::new(
                                    io::ErrorKind::WouldBlock,
                                    "SSL wants reading",
                                )),
                                SslError::Stream(ref e) => Err(From::from(e.kind())),
                                ref err => {
                                    Err(io::Error::new(io::ErrorKind::Other, format!("{}", err)))
                                }
                            };
                            *tls_stream = TlsStream::Handshake {
                                sock: mid,
                                negotiating: negotiating,
                            };
                            err
                        }
                        #[cfg(feature = "nativetls")]
                        Err(HandshakeError::Interrupted(mid)) => {
                            negotiating = true;
                            *tls_stream = TlsStream::Handshake {
                                sock: mid,
                                negotiating: negotiating,
                            };
                            Err(io::Error::new(
                                io::ErrorKind::WouldBlock,
                                "SSL would block",
                            ))
                        }
                        #[cfg(feature = "nativetls")]
                        Err(HandshakeError::Failure(err)) => {
                            Err(io::Error::new(
                                io::ErrorKind::Other,
                                format!("{}", err),
                            ))
                        }
                    },
                }
            }
        }
    }
}

impl io::Write for Stream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match *self {
            Tcp(ref mut sock) => sock.write(buf),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(TlsStream::Live(ref mut sock)) => sock.write(buf),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(ref mut tls_stream) => {
                trace!("Attempting to write ssl handshake.");
                match replace(tls_stream, TlsStream::Upgrading) {
                    TlsStream::Live(_) | TlsStream::Upgrading => unreachable!(),
                    TlsStream::Handshake {
                        sock,
                        mut negotiating,
                    } => {
                        match sock.handshake() {
                            Ok(mut sock) => {
                                trace!("Completed SSL Handshake");
                                let res = sock.write(buf);
                                *tls_stream = TlsStream::Live(sock);
                                res
                            }
                            #[cfg(feature = "ssl")]
                            Err(HandshakeError::SetupFailure(err)) => {
                                Err(io::Error::new(io::ErrorKind::Other, err))
                            }
                            #[cfg(feature = "ssl")]
                            Err(HandshakeError::Failure(mid))
                            | Err(HandshakeError::Interrupted(mid)) => {
                                let err = match *mid.error() {
                                    SslError::WantRead(_) => {
                                        negotiating = true;
                                        Err(io::Error::new(
                                            io::ErrorKind::WouldBlock,
                                            "SSL wants reading",
                                        ))
                                    }
                                    SslError::WantWrite(_) => {
                                        negotiating = false;
                                        Err(io::Error::new(
                                            io::ErrorKind::WouldBlock,
                                            "SSL wants writing",
                                        ))
                                    }
                                    SslError::Stream(ref e) => {
                                        negotiating = false;
                                        Err(From::from(e.kind()))
                                    }
                                    ref err => {
                                        negotiating = false;
                                        Err(io::Error::new(
                                            io::ErrorKind::Other,
                                            format!("{}", err),
                                        ))
                                    }
                                };
                                *tls_stream = TlsStream::Handshake {
                                    sock: mid,
                                    negotiating: negotiating,
                                };
                                err
                            }
                            #[cfg(feature = "nativetls")]
                            Err(HandshakeError::Interrupted(mid)) => {
                                negotiating = true;
                                *tls_stream = TlsStream::Handshake {
                                    sock: mid,
                                    negotiating: negotiating,
                                };
                                Err(io::Error::new(
                                    io::ErrorKind::WouldBlock,
                                    "SSL would block",
                                ))
                            }
                            #[cfg(feature = "nativetls")]
                            Err(HandshakeError::Failure(err)) => {
                                Err(io::Error::new(
                                    io::ErrorKind::Other,
                                    format!("{}", err),
                                ))
                            }
                        }
                    }
                }
            }
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match *self {
            Tcp(ref mut sock) => sock.flush(),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(TlsStream::Live(ref mut sock)) => sock.flush(),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(TlsStream::Handshake { ref mut sock, .. }) => sock.get_mut().flush(),
            #[cfg(any(feature = "ssl", feature = "nativetls"))]
            Tls(TlsStream::Upgrading) => panic!("Tried to access actively upgrading TlsStream"),
        }
    }
}

#[cfg(any(feature = "ssl", feature = "nativetls"))]
pub enum TlsStream {
    Live(SslStream<TcpStream>),
    Handshake {
        sock: MidHandshakeSslStream<TcpStream>,
        negotiating: bool,
    },
    Upgrading,
}

#[cfg(any(feature = "ssl", feature = "nativetls"))]
impl TlsStream {
    pub fn evented(&self) -> &TcpStream {
        match *self {
            TlsStream::Live(ref sock) => sock.get_ref(),
            TlsStream::Handshake { ref sock, .. } => sock.get_ref(),
            TlsStream::Upgrading => panic!("Tried to access actively upgrading TlsStream"),
        }
    }

    pub fn is_negotiating(&self) -> bool {
        match *self {
            TlsStream::Live(_) => false,
            TlsStream::Handshake {
                sock: _,
                negotiating,
            } => negotiating,
            TlsStream::Upgrading => panic!("Tried to access actively upgrading TlsStream"),
        }
    }

    pub fn clear_negotiating(&mut self) -> Result<()> {
        match *self {
            TlsStream::Live(_) => Err(Error::new(
                Kind::Internal,
                "Attempted to clear negotiating flag on live ssl connection.",
            )),
            TlsStream::Handshake {
                sock: _,
                ref mut negotiating,
            } => Ok(*negotiating = false),
            TlsStream::Upgrading => panic!("Tried to access actively upgrading TlsStream"),
        }
    }

    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
        match *self {
            TlsStream::Live(ref sock) => sock.get_ref().peer_addr(),
            TlsStream::Handshake { ref sock, .. } => sock.get_ref().peer_addr(),
            TlsStream::Upgrading => panic!("Tried to access actively upgrading TlsStream"),
        }
    }

    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        match *self {
            TlsStream::Live(ref sock) => sock.get_ref().local_addr(),
            TlsStream::Handshake { ref sock, .. } => sock.get_ref().local_addr(),
            TlsStream::Upgrading => panic!("Tried to access actively upgrading TlsStream"),
        }
    }
}
<|endoftext|>


<commit_before>use nom::{
  Consumer, ConsumerState,
  ErrorKind,
  HexDisplay,
  Producer, FileProducer,
  Input, IResult,
  Move, Needed,
};

use metadata;

use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};
use utility::resize_producer;

use std::io;
use std::io::{Error, Result};

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

pub struct Stream {
  info: StreamInfo,
  pub metadata: Vec<Metadata>,
  pub frames: Vec<Frame>,
  state: ParserState,
  consumer_state: ConsumerState<(), ErrorKind, Move>,
}

named!(pub stream_parser <&[u8], Stream>,
  chain!(
    blocks: metadata_parser ~
    frames: many1!(apply!(frame_parser, &blocks.0)),
    move|| {
      let consumed = Move::Consume(0);

      Stream {
        info: blocks.0,
        metadata: blocks.1,
        frames: frames,
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      }
    }
  )
);

impl Stream {
  pub fn new() -> Stream {
    let consumed = Move::Consume(0);

    Stream {
      info: StreamInfo::new(),
      metadata: Vec::new(),
      frames: Vec::new(),
      state: ParserState::Marker,
      consumer_state: ConsumerState::Continue(consumed),
    }
  }

  pub fn info(&self) -> StreamInfo {
    self.info
  }

  pub fn from_file(filename: &str) -> Result<Stream> {
    FileProducer::new(filename, 1024).and_then(|mut producer| {
      let consumed        = Move::Consume(0);
      let mut buffer_size = 1024;
      let mut is_error    = false;
      let mut stream      = Stream {
        info: StreamInfo::new(),
        metadata: Vec::new(),
        frames: Vec::new(),
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      };

      loop {
        match *producer.apply(&mut stream) {
          ConsumerState::Done(_, _)      => break,
          ConsumerState::Continue(await) => {
            let result = resize_producer(&mut producer, &await, buffer_size);

            if let Some(size) = result {
              buffer_size = size;
            }

            continue;
          }
          ConsumerState::Error(_)        => {
            is_error = true;

            break;
          }
        }
      }

      if !is_error {
        Ok(stream)
      } else {
        let error_str = format!("parser: couldn't parse the given file {}",
                                filename);

        Err(Error::new(io::ErrorKind::InvalidData, error_str))
      }
    })
  }

  fn handle_marker<'a>(&mut self, input: &'a [u8]) -> IResult<&'a [u8], ()> {
    let kind = nom::ErrorKind::Custom(0);

    match tag!(input, "fLaC") {
      IResult::Done(i, _)    => {
        self.state = ParserState::Metadata;

        IResult::Error(Err::Position(kind, i))
      }
      IResult::Error(_)      => IResult::Error(Err::Code(kind)),
      IResult::Incomplete(n) => IResult::Incomplete(n),
    }
  }

  fn handle_metadata<'a>(&mut self, input: &'a [u8])
                         -> IResult<&'a [u8], ()> {
    let kind = nom::ErrorKind::Custom(1);

    match metadata_parser(input) {
      IResult::Done(i, block) => {
        let is_last = block.is_last;

        if let metadata::Data::StreamInfo(info) = block.data {
          self.info = info;
        } else {
          self.metadata.push(block);
        }

        if is_last {
          self.state = ParserState::Frame;
        }

        IResult::Error(Err::Position(kind, i))
      }
      IResult::Error(_)      => IResult::Error(Err::Code(kind)),
      IResult::Incomplete(n) => IResult::Incomplete(n),
    }
  }

  fn handle_frame<'a>(&mut self, input: &'a [u8]) -> IResult<&'a [u8], ()> {
    let kind = nom::ErrorKind::Custom(2);

    match frame_parser(input, &self.info) {
      IResult::Done(i, frame) => {
        self.frames.push(frame);

        IResult::Error(Err::Position(kind, i))
      }
      IResult::Error(_)      => IResult::Error(Err::Code(kind)),
      IResult::Incomplete(n) => IResult::Incomplete(n),
    }
  }

  pub fn handle<S: StreamProducer>(&mut self, stream: &mut S)
                                   -> Result<(), ErrorKind> {
    stream.parse(|input| {
      match self.state {
        ParserState::Marker   => self.handle_marker(input),
        ParserState::Metadata => self.handle_metadata(input),
        ParserState::Frame    => self.handle_frame(input),
      }
    })
  }
}
<commit_msg>Remove `consumer_state` from `Stream`<commit_after>use nom::{
  Consumer, ConsumerState,
  ErrorKind,
  HexDisplay,
  Producer, FileProducer,
  Input, IResult,
  Move, Needed,
};

use metadata;

use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};
use utility::resize_producer;

use std::io;
use std::io::{Error, Result};

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

pub struct Stream {
  info: StreamInfo,
  pub metadata: Vec<Metadata>,
  pub frames: Vec<Frame>,
  state: ParserState,
}

named!(pub stream_parser <&[u8], Stream>,
  chain!(
    blocks: metadata_parser ~
    frames: many1!(apply!(frame_parser, &blocks.0)),
    move|| {
      let consumed = Move::Consume(0);

      Stream {
        info: blocks.0,
        metadata: blocks.1,
        frames: frames,
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      }
    }
  )
);

impl Stream {
  pub fn new() -> Stream {
    Stream {
      info: StreamInfo::new(),
      metadata: Vec::new(),
      frames: Vec::new(),
      state: ParserState::Marker,
    }
  }

  pub fn info(&self) -> StreamInfo {
    self.info
  }

  pub fn from_file(filename: &str) -> Result<Stream> {
    FileProducer::new(filename, 1024).and_then(|mut producer| {
      let consumed        = Move::Consume(0);
      let mut buffer_size = 1024;
      let mut is_error    = false;
      let mut stream      = Stream {
        info: StreamInfo::new(),
        metadata: Vec::new(),
        frames: Vec::new(),
        state: ParserState::Marker,
      };

      loop {
        match *producer.apply(&mut stream) {
          ConsumerState::Done(_, _)      => break,
          ConsumerState::Continue(await) => {
            let result = resize_producer(&mut producer, &await, buffer_size);

            if let Some(size) = result {
              buffer_size = size;
            }

            continue;
          }
          ConsumerState::Error(_)        => {
            is_error = true;

            break;
          }
        }
      }

      if !is_error {
        Ok(stream)
      } else {
        let error_str = format!("parser: couldn't parse the given file {}",
                                filename);

        Err(Error::new(io::ErrorKind::InvalidData, error_str))
      }
    })
  }

  fn handle_marker<'a>(&mut self, input: &'a [u8]) -> IResult<&'a [u8], ()> {
    let kind = nom::ErrorKind::Custom(0);

    match tag!(input, "fLaC") {
      IResult::Done(i, _)    => {
        self.state = ParserState::Metadata;

        IResult::Error(Err::Position(kind, i))
      }
      IResult::Error(_)      => IResult::Error(Err::Code(kind)),
      IResult::Incomplete(n) => IResult::Incomplete(n),
    }
  }

  fn handle_metadata<'a>(&mut self, input: &'a [u8])
                         -> IResult<&'a [u8], ()> {
    let kind = nom::ErrorKind::Custom(1);

    match metadata_parser(input) {
      IResult::Done(i, block) => {
        let is_last = block.is_last;

        if let metadata::Data::StreamInfo(info) = block.data {
          self.info = info;
        } else {
          self.metadata.push(block);
        }

        if is_last {
          self.state = ParserState::Frame;
        }

        IResult::Error(Err::Position(kind, i))
      }
      IResult::Error(_)      => IResult::Error(Err::Code(kind)),
      IResult::Incomplete(n) => IResult::Incomplete(n),
    }
  }

  fn handle_frame<'a>(&mut self, input: &'a [u8]) -> IResult<&'a [u8], ()> {
    let kind = nom::ErrorKind::Custom(2);

    match frame_parser(input, &self.info) {
      IResult::Done(i, frame) => {
        self.frames.push(frame);

        IResult::Error(Err::Position(kind, i))
      }
      IResult::Error(_)      => IResult::Error(Err::Code(kind)),
      IResult::Incomplete(n) => IResult::Incomplete(n),
    }
  }

  pub fn handle<S: StreamProducer>(&mut self, stream: &mut S)
                                   -> Result<(), ErrorKind> {
    stream.parse(|input| {
      match self.state {
        ParserState::Marker   => self.handle_marker(input),
        ParserState::Metadata => self.handle_metadata(input),
        ParserState::Frame    => self.handle_frame(input),
      }
    })
  }
}
<|endoftext|>


<commit_before>use std::os::raw::{c_int, c_double, c_void};

use ffi;
use base::*;

extern "C" fn wrapper<W>(raw_out: *mut ffi::SoundIoOutStream, min: c_int, max: c_int)
    where W: Fn(OutStream, i32, i32),
{
    let out = OutStream::new(raw_out);
    let cb_ptr = unsafe { (*out.stream).userdata as *const Box<W> };
    let cb: &W = unsafe { &*cb_ptr };
    cb(out, min, max);
}

struct OutStreamCallbacks
{
    write: Option<Box<Fn(OutStream, i32, i32)>>,
    underflow: Option<Box<Fn(OutStream)>>,
    error: Option<Box<Fn(OutStream, ffi::SioError)>>,
}
impl Default for OutStreamCallbacks
{
    fn default() -> Self {
        OutStreamCallbacks {
            write: None,
            underflow: None,
            error: None,
        }
    }
}

pub struct OutStream
{
    // TODO: make this private again
    pub stream: *mut ffi::SoundIoOutStream,
    callbacks: OutStreamCallbacks,
}
impl OutStream
{
    pub fn new(raw_stream: *mut ffi::SoundIoOutStream) -> Self {
        OutStream {
            stream: raw_stream,
            callbacks: OutStreamCallbacks::default(),
        }
    }

    pub fn open(&self) -> Option<ffi::SioError> {
        match unsafe { ffi::soundio_outstream_open(self.stream) } {
            ffi::SioError::None => None,
            err @ _ => Some(err),
        }
    }

    pub fn start(&self) -> Option<ffi::SioError> {
        match unsafe { ffi::soundio_outstream_start(self.stream) } {
            ffi::SioError::None => None,
            err @ _ => Some(err),
        }
    }

    pub fn register_write_callback<W>(&mut self, callback: Box<W>)
        where W: Fn(OutStream, i32, i32)
    {
        let userdata = &*callback as *const W as *mut c_void;
        unsafe {
            (*self.stream).userdata = userdata;
            (*self.stream).write_callback = Some(wrapper::<W>);
        }
    }

    pub fn begin_write(&self,
                       areas: *mut *mut ffi::SoundIoChannelArea,
                       frame_count: *mut c_int)
                       -> Option<ffi::SioError> {
        match unsafe { ffi::soundio_outstream_begin_write(self.stream, areas, frame_count) } {
            ffi::SioError::None => None,
            err @ _ => Some(err),
        }
    }

    pub fn end_write(&self) -> Option<ffi::SioError> {
        match unsafe { ffi::soundio_outstream_end_write(self.stream) } {
            ffi::SioError::None => None,
            err @ _ => Some(err),
        }
    }

    pub fn clear_buffer(&self) -> Option<ffi::SioError> {
        match unsafe { ffi::soundio_outstream_clear_buffer(self.stream) } {
            ffi::SioError::None => None,
            err @ _ => Some(err),
        }
    }

    pub fn pause(&self, pause: bool) -> Option<ffi::SioError> {
        let pause_c_bool = match pause {
            true => 1u8,
            false => 0u8,
        };
        match unsafe { ffi::soundio_outstream_pause(self.stream, pause_c_bool) } {
            ffi::SioError::None => None,
            err @ _ => Some(err),
        }
    }

    pub fn get_latency(&self) -> Result<f64, ffi::SioError> {
        let mut latency = 0.0f64;
        match unsafe {
            ffi::soundio_outstream_get_latency(self.stream, &mut latency as *mut c_double)
        } {
            ffi::SioError::None => Ok(latency),
            err @ _ => Err(err),
        }

    }

    pub fn current_format(&self) -> Result<ffi::SioFormat, ffi::SioError> {
        match unsafe { (*self.stream).format } {
            ffi::SioFormat::Invalid => Err(ffi::SioError::Invalid),
            fmt @ _ => Ok(fmt),
        }
    }

    pub fn get_layout(&self) -> ChannelLayout {
        ChannelLayout::new(unsafe { &(*self.stream).layout })
    }

    pub fn get_sample_rate(&self) -> i32 {
        unsafe { (*self.stream).sample_rate as i32 }
    }

    pub fn get_device(&self) -> Device {
        let dev = Device::new(unsafe { (*self.stream).device });
        dev.inc_ref();
        dev
    }

    pub fn destroy(&self) {
        unsafe { ffi::soundio_outstream_destroy(self.stream) }
    }
}
impl Drop for OutStream
{
    fn drop(&mut self) {
        // TODO: call destroy manually.
        // OutStream will get dropped each time a new
        // struct is created from the same *mut pointer.
    }
}
<commit_msg>Implement callback struct<commit_after>use std::os::raw::{c_int, c_double, c_void};

use ffi;
use base::*;

extern "C" fn write_callback_wrapper<W>(raw_out: *mut ffi::SoundIoOutStream,
                                        min: c_int,
                                        max: c_int)
    where W: Fn(OutStream, i32, i32)
{
    let out = OutStream::new(raw_out);
    let callbacks_ptr = unsafe { (*out.stream).userdata as *const Box<OutStreamCallbacks> };
    let callbacks: &Box<OutStreamCallbacks> = unsafe { &*callbacks_ptr };
    println!("match");
    match callbacks.write {
        Some(ref f) => f(out, min as i32, max as i32),
        None => println!("No registered callback found!"),
    }
    // let cb_ptr = unsafe { (*out.stream).userdata as *const Box<W> };
    // let cb: &W = unsafe { &*cb_ptr };
    // cb(out, min, max);
}

struct OutStreamCallbacks<'a> {
    write: Option<Box<Fn(OutStream, i32, i32) + 'a>>,
    underflow: Option<Box<Fn(OutStream) + 'a>>,
    error: Option<Box<Fn(OutStream, ffi::SioError) + 'a>>,
}
impl<'a> Default for OutStreamCallbacks<'a> {
    fn default() -> Self {
        OutStreamCallbacks {
            write: None,
            underflow: None,
            error: None,
        }
    }
}
impl<'a> Drop for OutStreamCallbacks<'a> {
    fn drop(&mut self) {
    }
}

pub struct OutStream<'a> {
    // TODO: make this private again
    pub stream: *mut ffi::SoundIoOutStream,
    callbacks: Box<OutStreamCallbacks<'a>>,
}
impl<'a> OutStream<'a> {
    pub fn new(raw_stream: *mut ffi::SoundIoOutStream) -> Self {
        let callbacks = Box::new(OutStreamCallbacks::default());
        OutStream {
            stream: raw_stream,
            callbacks: callbacks,
        }
    }

    pub fn open(&self) -> Option<ffi::SioError> {
        match unsafe { ffi::soundio_outstream_open(self.stream) } {
            ffi::SioError::None => None,
            err @ _ => Some(err),
        }
    }

    pub fn start(&self) -> Option<ffi::SioError> {
        match unsafe { ffi::soundio_outstream_start(self.stream) } {
            ffi::SioError::None => None,
            err @ _ => Some(err),
        }
    }

    pub fn register_write_callback<W>(&mut self, callback: Box<W>)
        where W: Fn(OutStream, i32, i32) + 'a
    {
        // stored box reference to callback closure
        self.callbacks.write = Some(callback);
        unsafe {
            // register wrapper for write_callback
            (*self.stream).write_callback = Some(write_callback_wrapper::<W>);
            // store reference to callbacks struct in userdata pointer
            (*self.stream).userdata =
                &self.callbacks as *const Box<OutStreamCallbacks> as *mut c_void
        }
    }

    pub fn register_underflow_callback<U>(&mut self, callback: Box<U>)
        where U: Fn(OutStream) + 'a
    {
        self.callbacks.underflow = Some(callback)
    }

    pub fn register_error_callback<E>(&mut self, callback: Box<E>)
        where E: Fn(OutStream, ffi::SioError) + 'a
    {
        self.callbacks.error = Some(callback)
    }

    pub fn begin_write(&self,
                       areas: *mut *mut ffi::SoundIoChannelArea,
                       frame_count: *mut c_int)
                       -> Option<ffi::SioError> {
        match unsafe { ffi::soundio_outstream_begin_write(self.stream, areas, frame_count) } {
            ffi::SioError::None => None,
            err @ _ => Some(err),
        }
    }

    pub fn end_write(&self) -> Option<ffi::SioError> {
        match unsafe { ffi::soundio_outstream_end_write(self.stream) } {
            ffi::SioError::None => None,
            err @ _ => Some(err),
        }
    }

    pub fn clear_buffer(&self) -> Option<ffi::SioError> {
        match unsafe { ffi::soundio_outstream_clear_buffer(self.stream) } {
            ffi::SioError::None => None,
            err @ _ => Some(err),
        }
    }

    pub fn pause(&self, pause: bool) -> Option<ffi::SioError> {
        let pause_c_bool = match pause {
            true => 1u8,
            false => 0u8,
        };
        match unsafe { ffi::soundio_outstream_pause(self.stream, pause_c_bool) } {
            ffi::SioError::None => None,
            err @ _ => Some(err),
        }
    }

    pub fn get_latency(&self) -> Result<f64, ffi::SioError> {
        let mut latency = 0.0f64;
        match unsafe {
            ffi::soundio_outstream_get_latency(self.stream, &mut latency as *mut c_double)
        } {
            ffi::SioError::None => Ok(latency),
            err @ _ => Err(err),
        }

    }

    pub fn current_format(&self) -> Result<ffi::SioFormat, ffi::SioError> {
        match unsafe { (*self.stream).format } {
            ffi::SioFormat::Invalid => Err(ffi::SioError::Invalid),
            fmt @ _ => Ok(fmt),
        }
    }

    pub fn get_layout(&self) -> ChannelLayout {
        ChannelLayout::new(unsafe { &(*self.stream).layout })
    }

    pub fn get_sample_rate(&self) -> i32 {
        unsafe { (*self.stream).sample_rate as i32 }
    }

    pub fn get_device(&self) -> Device {
        let dev = Device::new(unsafe { (*self.stream).device });
        dev.inc_ref();
        dev
    }

    pub fn destroy(&self) {
        unsafe { ffi::soundio_outstream_destroy(self.stream) }
    }
}
impl<'a> Drop for OutStream<'a> {
    fn drop(&mut self) {
        // TODO: call destroy manually.
        // OutStream will get dropped each time a new
        // struct is created from the same *mut pointer.
    }
}
<|endoftext|>


<commit_before>use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

pub struct Stream {
  pub info: StreamInfo,
  pub metadata: Vec<Metadata>,
  pub frames: Vec<Frame>,
}

named!(pub stream_parser <&[u8], Stream>,
  chain!(
    blocks: metadata_parser ~
    frames: many1!(apply!(frame_parser, &blocks.0)),
    move|| {
      Stream {
        info: blocks.0,
        metadata: blocks.1,
        frames: frames,
      }
    }
  )
);
<commit_msg>Add the two different states to `Stream`<commit_after>use nom::{ConsumerState, Move};

use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

pub struct Stream {
  pub info: StreamInfo,
  pub metadata: Vec<Metadata>,
  pub frames: Vec<Frame>,
  state: ParserState,
  consumer_state: ConsumerState<(), ErrorKind, Move>,
}

named!(pub stream_parser <&[u8], Stream>,
  chain!(
    blocks: metadata_parser ~
    frames: many1!(apply!(frame_parser, &blocks.0)),
    move|| {
      let consumed = Move::Consume(0);

      Stream {
        info: blocks.0,
        metadata: blocks.1,
        frames: frames,
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      }
    }
  )
);
<|endoftext|>


<commit_before>//! A byte-offset based string table.
//! Commonly used in ELF binaries, Unix archives, and even PE binaries.

use core::ops::Index;
use core::slice;
use core::str;
use core::fmt;
use scroll::{self, ctx, Pread};
#[cfg(feature = "std")]
use error;

/// A common string table format which is indexed by byte offsets (and not
/// member index). Constructed using [`parse`](#method.parse)
/// with your choice of delimiter. Please be careful.
pub struct Strtab<'a> {
    bytes: &'a[u8],
    delim: ctx::StrCtx,
}

#[inline(always)]
fn get_str(offset: usize, bytes: &[u8], delim: ctx::StrCtx) -> scroll::Result<&str> {
    bytes.pread_with::<&str>(offset, delim)
}

impl<'a> Strtab<'a> {
    /// Construct a new strtab with `bytes` as the backing string table, using `delim` as the delimiter between entries
    pub fn new (bytes: &'a [u8], delim: u8) -> Self {
        Strtab { delim: ctx::StrCtx::from(delim), bytes: bytes }
    }
    /// Construct a strtab from a `ptr`, and a `size`, using `delim` as the delimiter
    pub unsafe fn from_raw(ptr: *const u8, size: usize, delim: u8) -> Strtab<'a> {
        Strtab { delim: ctx::StrCtx::from(delim), bytes: slice::from_raw_parts(ptr, size) }
    }
    #[cfg(feature = "std")]
    /// Parses a strtab from `bytes` at `offset` with `len` size as the backing string table, using `delim` as the delimiter
    pub fn parse(bytes: &'a [u8], offset: usize, len: usize, delim: u8) -> error::Result<Strtab<'a>> {
        let bytes: &'a [u8] = bytes.pread_slice(offset, len)?;
        Ok(Strtab { bytes: bytes, delim: ctx::StrCtx::from(delim) })
    }
    #[cfg(feature = "std")]
    /// Converts the string table to a vector, with the original `delim` used to separate the strings
    pub fn to_vec(self) -> error::Result<Vec<String>> {
        let len = self.bytes.len();
        let mut strings = Vec::with_capacity(len);
        let mut i = 0;
        while i < len {
            let string = self.get(i)?;
            i = i + string.len() + 1;
            strings.push(string.to_string());
        }
        Ok(strings)
    }
    /// Safely gets a str reference from the backing bytes starting at byte `offset`
    pub fn get(&self, offset: usize) -> error::Result<&'a str> {
        Ok(get_str(offset, self.bytes, self.delim)?)
    }
}

impl<'a> fmt::Debug for Strtab<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "delim: {:?} {:?}", self.delim, str::from_utf8(&self.bytes))
    }
}

impl<'a> Default for Strtab<'a> {
    fn default() -> Strtab<'a> {
        Strtab { bytes: &[], delim: ctx::StrCtx::default() }
    }
}

impl<'a> Index<usize> for Strtab<'a> {
    type Output = str;
    /// Gets str reference at starting at byte `offset`.
    /// **NB**: this will panic if the underlying bytes are not valid utf8, or the offset is invalid
    fn index(&self, offset: usize) -> &Self::Output {
        get_str(offset, &self.bytes, self.delim).unwrap()
    }
}

#[test]
fn as_vec_no_final_null() {
    let bytes = b"\0printf\0memmove\0busta";
    let strtab = unsafe { Strtab::from_raw(bytes.as_ptr(), bytes.len(), 0x0) };
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 4);
    assert_eq!(vec, vec!["", "printf", "memmove", "busta"]);
}

#[test]
fn as_vec_no_first_null_no_final_null() {
    let bytes = b"printf\0memmove\0busta";
    let strtab = unsafe { Strtab::from_raw(bytes.as_ptr(), bytes.len(), 0x0) };
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 3);
    assert_eq!(vec, vec!["printf", "memmove", "busta"]);
}

#[test]
fn to_vec_final_null() {
    let bytes = b"\0printf\0memmove\0busta\0";
    let strtab = unsafe { Strtab::from_raw(bytes.as_ptr(), bytes.len(), 0x0) };
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 4);
    assert_eq!(vec, vec!["", "printf", "memmove", "busta"]);
}

#[test]
fn to_vec_newline_delim() {
    let bytes = b"\nprintf\nmemmove\nbusta\n";
    let strtab = unsafe { Strtab::from_raw(bytes.as_ptr(), bytes.len(), '\n' as u8) };
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 4);
    assert_eq!(vec, vec!["", "printf", "memmove", "busta"]);
}
<commit_msg>strtab: use scroll error for Strtab::get so we can compile on nostd<commit_after>//! A byte-offset based string table.
//! Commonly used in ELF binaries, Unix archives, and even PE binaries.

use core::ops::Index;
use core::slice;
use core::str;
use core::fmt;
use scroll::{self, ctx, Pread};
#[cfg(feature = "std")]
use error;

/// A common string table format which is indexed by byte offsets (and not
/// member index). Constructed using [`parse`](#method.parse)
/// with your choice of delimiter. Please be careful.
pub struct Strtab<'a> {
    bytes: &'a[u8],
    delim: ctx::StrCtx,
}

#[inline(always)]
fn get_str(offset: usize, bytes: &[u8], delim: ctx::StrCtx) -> scroll::Result<&str> {
    bytes.pread_with::<&str>(offset, delim)
}

impl<'a> Strtab<'a> {
    /// Construct a new strtab with `bytes` as the backing string table, using `delim` as the delimiter between entries
    pub fn new (bytes: &'a [u8], delim: u8) -> Self {
        Strtab { delim: ctx::StrCtx::from(delim), bytes: bytes }
    }
    /// Construct a strtab from a `ptr`, and a `size`, using `delim` as the delimiter
    pub unsafe fn from_raw(ptr: *const u8, size: usize, delim: u8) -> Strtab<'a> {
        Strtab { delim: ctx::StrCtx::from(delim), bytes: slice::from_raw_parts(ptr, size) }
    }
    #[cfg(feature = "std")]
    /// Parses a strtab from `bytes` at `offset` with `len` size as the backing string table, using `delim` as the delimiter
    pub fn parse(bytes: &'a [u8], offset: usize, len: usize, delim: u8) -> error::Result<Strtab<'a>> {
        let bytes: &'a [u8] = bytes.pread_slice(offset, len)?;
        Ok(Strtab { bytes: bytes, delim: ctx::StrCtx::from(delim) })
    }
    #[cfg(feature = "std")]
    /// Converts the string table to a vector, with the original `delim` used to separate the strings
    pub fn to_vec(self) -> error::Result<Vec<String>> {
        let len = self.bytes.len();
        let mut strings = Vec::with_capacity(len);
        let mut i = 0;
        while i < len {
            let string = self.get(i)?;
            i = i + string.len() + 1;
            strings.push(string.to_string());
        }
        Ok(strings)
    }
    /// Safely parses and gets a str reference from the backing bytes starting at byte `offset`
    pub fn get(&self, offset: usize) -> scroll::Result<&'a str> {
        get_str(offset, self.bytes, self.delim)
    }
}

impl<'a> fmt::Debug for Strtab<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "delim: {:?} {:?}", self.delim, str::from_utf8(&self.bytes))
    }
}

impl<'a> Default for Strtab<'a> {
    fn default() -> Strtab<'a> {
        Strtab { bytes: &[], delim: ctx::StrCtx::default() }
    }
}

impl<'a> Index<usize> for Strtab<'a> {
    type Output = str;
    /// Gets str reference at starting at byte `offset`.
    /// **NB**: this will panic if the underlying bytes are not valid utf8, or the offset is invalid
    fn index(&self, offset: usize) -> &Self::Output {
        get_str(offset, &self.bytes, self.delim).unwrap()
    }
}

#[test]
fn as_vec_no_final_null() {
    let bytes = b"\0printf\0memmove\0busta";
    let strtab = unsafe { Strtab::from_raw(bytes.as_ptr(), bytes.len(), 0x0) };
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 4);
    assert_eq!(vec, vec!["", "printf", "memmove", "busta"]);
}

#[test]
fn as_vec_no_first_null_no_final_null() {
    let bytes = b"printf\0memmove\0busta";
    let strtab = unsafe { Strtab::from_raw(bytes.as_ptr(), bytes.len(), 0x0) };
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 3);
    assert_eq!(vec, vec!["printf", "memmove", "busta"]);
}

#[test]
fn to_vec_final_null() {
    let bytes = b"\0printf\0memmove\0busta\0";
    let strtab = unsafe { Strtab::from_raw(bytes.as_ptr(), bytes.len(), 0x0) };
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 4);
    assert_eq!(vec, vec!["", "printf", "memmove", "busta"]);
}

#[test]
fn to_vec_newline_delim() {
    let bytes = b"\nprintf\nmemmove\nbusta\n";
    let strtab = unsafe { Strtab::from_raw(bytes.as_ptr(), bytes.len(), '\n' as u8) };
    let vec = strtab.to_vec().unwrap();
    assert_eq!(vec.len(), 4);
    assert_eq!(vec, vec!["", "printf", "memmove", "busta"]);
}
<|endoftext|>


<commit_before>use util::{bitfields, bits::Bits, primitive_enum};

bitfields! {
    /// 4000208h - IME - Interrupt Master Enable Register (R/W)
    ///   Bit   Expl.
    ///   0     Disable all interrupts         (0=Disable All, 1=See [`InterruptEnable`])
    ///   1-31  Not used
    pub struct InterruptMasterEnable: u16 {
        [0] disable_all, set_disable_all: bool,
    }
}

bitfields! {
    pub struct InterruptEnable: u16 {}
}

impl InterruptEnable {
    pub fn enabled(&self, interrupt: Interrupt) -> bool {
        self.value.is_bit_set(u8::from(interrupt) as u32)
    }

    pub fn set_enabled(&self, interrupt: Interrupt, enabled: bool) {
        self.value.replace_bit(u8::from(interrupt) as u32, enabled);
    }
}

bitfields! {
    pub struct InterruptReqAck: u16 {}
}

impl InterruptReqAck {
    /// Interrupts can be acknowledged by writing a one to one of the IRQ bits.
    pub fn set_acknowledge(&mut self, value: u16) {
        self.value &= !value;
    }
}

primitive_enum! {
    pub enum Interrupt: u8 {
        VBlank = 0,
        HBlank,
        VCounterMatch,
        Timer0Overflow,
        Timer1Overflow,
        Timer2Overflow,
        Timer3Overflow,
        SerialCommunication,
        DMA0,
        DMA1,
        DMA2,
        DMA3,
        Keypad,
        GamePak,
    }
}
<commit_msg>Correct naming of master enable bit.<commit_after>use util::{bitfields, bits::Bits, primitive_enum};

bitfields! {
    /// 4000208h - IME - Interrupt Master Enable Register (R/W)
    ///   Bit   Expl.
    ///   0     Disable all interrupts         (0=Disable All, 1=See [`InterruptEnable`])
    ///   1-31  Not used
    pub struct InterruptMasterEnable: u16 {
        [0] enabled, set_enabled: bool,
    }
}

bitfields! {
    pub struct InterruptEnable: u16 {}
}

impl InterruptEnable {
    pub fn enabled(&self, interrupt: Interrupt) -> bool {
        self.value.is_bit_set(u8::from(interrupt) as u32)
    }

    pub fn set_enabled(&self, interrupt: Interrupt, enabled: bool) {
        self.value.replace_bit(u8::from(interrupt) as u32, enabled);
    }
}

bitfields! {
    pub struct InterruptReqAck: u16 {}
}

impl InterruptReqAck {
    /// Interrupts can be acknowledged by writing a one to one of the IRQ bits.
    pub fn set_acknowledge(&mut self, value: u16) {
        self.value &= !value;
    }
}

primitive_enum! {
    pub enum Interrupt: u8 {
        VBlank = 0,
        HBlank,
        VCounterMatch,
        Timer0Overflow,
        Timer1Overflow,
        Timer2Overflow,
        Timer3Overflow,
        SerialCommunication,
        DMA0,
        DMA1,
        DMA2,
        DMA3,
        Keypad,
        GamePak,
    }
}
<|endoftext|>