<commit_before>use log::{debug};

use super::arch::RVA;

const PAGE_SIZE: usize = 0x1000;

#[derive(Debug)]
pub enum Error {
    NotMapped,
}

fn page(rva: RVA) -> usize {
    // depends on PAGE_SIZE
    let v: u64 = rva.into();
    // #yolo
    (v as usize) >> 12
}

fn page_offset(rva: RVA) -> usize {
    // depends on PAGE_SIZE
    let v: u64 = rva.into();
    // #yolo
    (v as usize) & 0xFFF
}

struct Page<T: Default + Copy> {
    elements: [T; PAGE_SIZE],
}

impl<T: Default + Copy> Page<T> {
    fn new(items: &[T]) -> Page<T> {
        let mut page: Page<T> = Default::default();
        page.elements.copy_from_slice(items);
        page
    }
}

impl<T: Default + Copy> Default for Page<T> {
    fn default() -> Self {
        Page {
            elements: [Default::default(); PAGE_SIZE]
        }
    }
}

/// PageMap is a map-like data structure that stores `Copy` elements in pages of 0x1000.
///
/// Its a good choice when representing lots of small elements that are found at
/// contiguous indices. At the moment, indices are `RVA`.
///
/// Lookups should be quick, as they boil down to just a couple dereferences.
pub struct PageMap<T: Default + Copy> {
    pages: Vec<Option<Page<T>>>
}

impl<T: Default + Copy> PageMap<T> {
    pub fn with_capacity(capacity: RVA) -> PageMap<T>{
        let page_count = page(capacity) + 1;
        let mut pages = Vec::with_capacity(page_count);
        pages.resize_with(page_count, || None);

        PageMap {
            pages
        }
    }

    /// error if rva is not in a valid page.
    /// panic due to:
    ///   - rva must be page aligned.
    ///   - must be PAGE_SIZE number of items.
    fn map_page(&mut self, rva: RVA, items: &[T]) -> Result<(), Error> {
        if page_offset(rva) != 0 {
            panic!("invalid map address");
        }
        if items.len() != PAGE_SIZE {
            panic!("invalid map buffer size");
        }
        debug!("mapping page: {}", rva);
        if page(rva) > self.pages.len() {
            return Err(Error::NotMapped);
        }

        self.pages[page(rva)] = Some(Page::new(items));

        Ok(())
    }

    /// map the given items at the given address.
    ///
    /// error if rva or items are not in a valid page.
    /// panic due to:
    ///   - rva must be page aligned.
    ///   - must be multiple of PAGE_SIZE number of items.
    ///
    /// see example under `get`.
    pub fn map(&mut self, rva: RVA, items: &[T]) -> Result<(), Error> {
        for (i, chunk) in items.chunks_exact(PAGE_SIZE).enumerate() {
            self.map_page(rva + i * PAGE_SIZE, chunk)?;
        }
        Ok(())
    }

    /// map the default value (probably zero) at the given address for the given size.
    ///
    /// same error conditions as `map`.
    /// see example under `probe`.
    pub fn map_empty(&mut self, rva: RVA, size: usize) -> Result<(), Error> {
        self.map(rva, &vec![Default::default(); size])
    }

    /// is the given address mapped?
    ///
    /// ```
    /// use lancelot::arch::RVA;
    /// use lancelot::pagemap::PageMap;
    ///
    /// let mut d: PageMap<u32> = PageMap::with_capacity(0x2000.into());
    /// assert_eq!(d.probe(0x0.into()), false);
    /// assert_eq!(d.probe(0x1000.into()), false);
    ///
    /// d.map_empty(0x0.into(), 0x1000).expect("failed to map");
    /// assert_eq!(d.probe(0x0.into()), true);
    /// assert_eq!(d.probe(0x1000.into()), false);
    /// ```
    pub fn probe(&self, rva: RVA) -> bool {
        if page(rva) > self.pages.len() {
            return false;
        }

        return self.pages[page(rva)].is_some()
    }

    /// fetch one item from the given address.
    /// if the address is not mapped, then the result is `None`.
    ///
    /// ```
    /// use lancelot::arch::RVA;
    /// use lancelot::pagemap::PageMap;
    ///
    /// let mut d: PageMap<u32> = PageMap::with_capacity(0x2000.into());
    /// assert_eq!(d.get(0x0.into()), None);
    /// assert_eq!(d.get(0x1000.into()), None);
    ///
    /// d.map(0x1000.into(), &[0x1; 0x1000]).expect("failed to map");
    /// assert_eq!(d.get(0x0.into()), None);
    /// assert_eq!(d.get(0x1000.into()), Some(0x1));
    ///
    /// d.map(0x0.into(), &[0x2; 0x2000]).expect("failed to map");
    ///  assert_eq!(d.get(0x0.into()), Some(0x2));
    ///  assert_eq!(d.get(0x1000.into()), Some(0x2));
    /// ```
    pub fn get(&self, rva: RVA) -> Option<T> {
        if page(rva) > self.pages.len() {
            return None;
        }

        let page = match &self.pages[page(rva)] {
            // page is not mapped
            None => return None,
            // page is mapped
            Some(page) => page,
        };

        Some(page.elements[page_offset(rva)])
    }

    /// fetch one mutable item from the given address.
    ///
    /// ```
    /// use lancelot::arch::RVA;
    /// use lancelot::pagemap::PageMap;
    ///
    /// let mut d: PageMap<u32> = PageMap::with_capacity(0x2000.into());
    /// d.map_empty(0x0.into(), 0x1000).expect("failed to map");
    ///
    /// // address 0x0 starts at 0
    /// assert_eq!(d.get(0x0.into()), Some(0x0));
    ///
    /// // set address 0x0 to 1
    /// let v = d.get_mut(0x0.into()).expect("should be mapped");
    /// *v = 1;
    ///
    /// // address 0x0 is 1
    /// assert_eq!(d.get(0x0.into()), Some(0x1));
    /// ```
    pub fn get_mut(&mut self, rva: RVA) -> Option<&mut T> {
        if page(rva) > self.pages.len() {
            return None;
        }

        let page = match &mut self.pages[page(rva)] {
            // page is not mapped
            None => return None,
            // page is mapped
            Some(page) => page,
        };

        Some(&mut page.elements[page_offset(rva)])
    }

    /// handle the simple slice case: when start and end fall within the same page.
    /// for example, reading a dword from address 0x10.
    ///
    /// ```
    /// use lancelot::arch::RVA;
    /// use lancelot::pagemap::PageMap;
    ///
    /// let mut d: PageMap<u32> = PageMap::with_capacity(0x2000.into());
    /// d.map_empty(0x0.into(), 0x1000).expect("failed to map");
    /// assert_eq!(d.slice(0x0.into(), 0x2.into()).unwrap(), [0x0, 0x0]);
    /// ```
    fn slice_into_simple(&self, start: RVA, buf: &mut [T]) -> Result<(), Error> {
        // precondition: page(start) == page(start + buf.len())

        if page(start) > self.pages.len() {
            return Err(Error::NotMapped);
        }

        let page = match &self.pages[page(start)] {
            // page is not mapped
            None => return Err(Error::NotMapped),
            // page is mapped
            Some(page) => page,
        };

        let end = start + buf.len();
        let elements = &page.elements[page_offset(start)..page_offset(end)];
        buf.copy_from_slice(elements);

        Ok(())
    }

    /// handle the complex slice case: when start and end are on different pages.
    /// for example, reading a dword from address 0xFFE.
    ///
    /// ```
    /// use lancelot::arch::RVA;
    /// use lancelot::pagemap::PageMap;
    ///
    /// let mut d: PageMap<u32> = PageMap::with_capacity(0x5000.into());
    /// d.map_empty(0x1000.into(), 0x3000).expect("failed to map");
    ///
    /// // 0     unmapped
    /// //       unmapped
    /// // 1000  0 0 0 0
    /// //       0 0 0 0
    /// // 2000  0 0 0 0
    /// //       0 0 0 0
    /// // 3000  0 0 0 0
    /// //       0 0 0 0
    /// // 4000  unmapped
    /// //       unmapped
    /// // 5000  unmapped
    ///
    /// assert_eq!(d.slice(0x1FFC.into(), 0x2000.into()).unwrap(), [0x0, 0x0, 0x0, 0x0], "no overlap");
    /// assert_eq!(d.slice(0x1FFD.into(), 0x2001.into()).unwrap(), [0x0, 0x0, 0x0, 0x0], "overlap 1");
    /// assert_eq!(d.slice(0x1FFE.into(), 0x2002.into()).unwrap(), [0x0, 0x0, 0x0, 0x0], "overlap 2");
    /// assert_eq!(d.slice(0x1FFF.into(), 0x2003.into()).unwrap(), [0x0, 0x0, 0x0, 0x0], "overlap 3");
    /// assert_eq!(d.slice(0x2000.into(), 0x2004.into()).unwrap(), [0x0, 0x0, 0x0, 0x0], "overlap 4");
    ///
    /// assert_eq!(d.slice(0x1FFC.into(), 0x3004.into()).unwrap().len(), 0x1008, "4, page, 4");
    ///
    /// assert_eq!(d.slice(0x1FFC.into(), 0x3000.into()).unwrap().len(), 0x1004, "4, page");
    ///
    /// assert_eq!(d.slice(0x2000.into(), 0x3004.into()).unwrap().len(), 0x1004, "page, 4");
    /// ```
    fn slice_into_split(&self, start: RVA, buf: &mut [T]) -> Result<(), Error> {
        let end = start + buf.len();

        if page(end) > self.pages.len() {
            return Err(Error::NotMapped);
        }

        // ensure each page within the requested region is mapped.
        let start_page = page(start);
        let end_page = if page_offset(end) == 0 { page(end) - 1} else { page(end) };
        for page in start_page..end_page {
            if ! self.probe((page * PAGE_SIZE).into()) {
                return Err(Error::NotMapped);
            }
        }

        let mut offset: usize = 0;

        // region one: from `start` to the end of its page
        // region two: any intermediate complete pages
        // region three: from start of final page until `end`

        // one.
        {
            let page = self.pages[page(start)].as_ref().unwrap();
            let elements = &page.elements[page_offset(start)..];
            {
                let dst = &mut buf[offset..offset+elements.len()];
                dst.copy_from_slice(elements);
                offset += elements.len();
            }
        }

        // two.
        if page(start) != page(end) - 1 {
            let start_index = page(start) + 1;
            let end_index = page(end);
            for page_index in start_index..end_index {
                let page = self.pages[page_index].as_ref().unwrap();
                let elements = &page.elements[..];
                {
                    let dst = &mut buf[offset..offset+elements.len()];
                    dst.copy_from_slice(elements);
                    offset += elements.len();
                }
            }
        }

        // three.
        if page_offset(end) != 0x0 {
            let page = self.pages[page(end)].as_ref().unwrap();
            let elements = &page.elements[..page_offset(end)];
            {
                let dst = &mut buf[offset..offset+elements.len()];
                dst.copy_from_slice(elements);
            }
        }

        Ok(())
    }

    /// fetch the items found in the given range, placing them into the given slice.
    /// compared with `slice`, this routine avoids an allocation.
    ///
    /// errors:
    ///   - Error::NotMapped: if any requested address is not mapped
    pub fn slice_into(&self, start: RVA, buf: &mut [T]) -> Result<(), Error> {
        let end = start + buf.len();
        if page(start) == page(end) {
            self.slice_into_simple(start, buf)
        } else {
            self.slice_into_split(start, buf)
        }
    }

    /// fetch the items found in the given range.
    ///
    /// errors:
    ///   - Error::NotMapped: if any requested address is not mapped
    ///
    /// panic if:
    ///   - start > end
    ///
    /// ```
    /// use lancelot::arch::RVA;
    /// use lancelot::pagemap::PageMap;
    ///
    /// let mut d: PageMap<u32> = PageMap::with_capacity(0x2000.into());
    /// d.map_empty(0x0.into(), 0x1000).expect("failed to map");
    /// assert_eq!(d.slice(0x0.into(), 0x2.into()).unwrap(), [0x0, 0x0]);
    /// ```
    pub fn slice(&self, start: RVA, end: RVA) -> Result<Vec<T>, Error> {
        if start > end {
            panic!("start > end");
        }

        let mut ret = vec![Default::default(); (end - start).into()];
        self.slice_into(start, &mut ret)?;

        Ok(ret)
    }
}
<commit_msg>pagemap: implement mapzx<commit_after>use log::{debug};

use super::arch::RVA;

const PAGE_SIZE: usize = 0x1000;

#[derive(Debug)]
pub enum Error {
    NotMapped,
}

fn page(rva: RVA) -> usize {
    // depends on PAGE_SIZE
    let v: u64 = rva.into();
    // #yolo
    (v as usize) >> 12
}

fn page_offset(rva: RVA) -> usize {
    // depends on PAGE_SIZE
    let v: u64 = rva.into();
    // #yolo
    (v as usize) & 0xFFF
}

struct Page<T: Default + Copy> {
    elements: [T; PAGE_SIZE],
}

impl<T: Default + Copy> Page<T> {
    fn new(items: &[T]) -> Page<T> {
        let mut page: Page<T> = Default::default();
        page.elements.copy_from_slice(items);
        page
    }
}

impl<T: Default + Copy> Default for Page<T> {
    fn default() -> Self {
        Page {
            elements: [Default::default(); PAGE_SIZE]
        }
    }
}

/// PageMap is a map-like data structure that stores `Copy` elements in pages of 0x1000.
///
/// Its a good choice when representing lots of small elements that are found at
/// contiguous indices. At the moment, indices are `RVA`.
///
/// Lookups should be quick, as they boil down to just a couple dereferences.
pub struct PageMap<T: Default + Copy> {
    pages: Vec<Option<Page<T>>>
}

impl<T: Default + Copy> PageMap<T> {
    pub fn with_capacity(capacity: RVA) -> PageMap<T>{
        let page_count = page(capacity) + 1;
        let mut pages = Vec::with_capacity(page_count);
        pages.resize_with(page_count, || None);

        PageMap {
            pages
        }
    }

    /// error if rva is not in a valid page.
    /// panic due to:
    ///   - rva must be page aligned.
    ///   - must be PAGE_SIZE number of items.
    fn map_page(&mut self, rva: RVA, items: &[T]) -> Result<(), Error> {
        if page_offset(rva) != 0 {
            panic!("invalid map address");
        }
        if items.len() != PAGE_SIZE {
            panic!("invalid map buffer size");
        }
        debug!("mapping page: {}", rva);
        if page(rva) > self.pages.len() {
            return Err(Error::NotMapped);
        }

        self.pages[page(rva)] = Some(Page::new(items));

        Ok(())
    }

    /// map the given items at the given address.
    ///
    /// error if rva or items are not in a valid page.
    /// panic due to:
    ///   - rva must be page aligned.
    ///   - must be multiple of PAGE_SIZE number of items.
    ///
    /// see example under `get`.
    pub fn map(&mut self, rva: RVA, items: &[T]) -> Result<(), Error> {
        for (i, chunk) in items.chunks_exact(PAGE_SIZE).enumerate() {
            self.map_page(rva + i * PAGE_SIZE, chunk)?;
        }
        Ok(())
    }

    /// map the default value (probably zero) at the given address for the given size.
    ///
    /// same error conditions as `map`.
    /// see example under `probe`.
    pub fn map_empty(&mut self, rva: RVA, size: usize) -> Result<(), Error> {
        self.map(rva, &vec![Default::default(); size])
    }

    /// map the given items at the given address, padding with the default value until the next page.
    /// (map zero-extend).
    ///
    /// same error conditions as `map`.
    ///
    /// ```
    /// use lancelot::arch::RVA;
    /// use lancelot::pagemap::PageMap;
    ///
    /// let mut d: PageMap<u32> = PageMap::with_capacity(0x2000.into());
    /// assert_eq!(d.get(0x0.into()), None);
    /// assert_eq!(d.get(0x1.into()), None);
    ///
    /// d.mapzx(0x0.into(), &[0x1, ]).expect("failed to map");
    /// assert_eq!(d.get(0x0.into()), Some(0x1));
    /// assert_eq!(d.get(0x0.into()), Some(0x0));
    /// ```
    pub fn mapzx(&mut self, rva: RVA, items: &[T]) -> Result<(), Error> {
        let empty_count = PAGE_SIZE - page_offset(items.len().into());
        let mut padded_items = Vec::with_capacity(items.len() + empty_count);
        padded_items.extend(items);
        padded_items.extend(&vec![Default::default(); empty_count]);
        self.map(rva, &padded_items)
    }

    /// is the given address mapped?
    ///
    /// ```
    /// use lancelot::arch::RVA;
    /// use lancelot::pagemap::PageMap;
    ///
    /// let mut d: PageMap<u32> = PageMap::with_capacity(0x2000.into());
    /// assert_eq!(d.probe(0x0.into()), false);
    /// assert_eq!(d.probe(0x1000.into()), false);
    ///
    /// d.map_empty(0x0.into(), 0x1000).expect("failed to map");
    /// assert_eq!(d.probe(0x0.into()), true);
    /// assert_eq!(d.probe(0x1000.into()), false);
    /// ```
    pub fn probe(&self, rva: RVA) -> bool {
        if page(rva) > self.pages.len() {
            return false;
        }

        return self.pages[page(rva)].is_some()
    }

    /// fetch one item from the given address.
    /// if the address is not mapped, then the result is `None`.
    ///
    /// ```
    /// use lancelot::arch::RVA;
    /// use lancelot::pagemap::PageMap;
    ///
    /// let mut d: PageMap<u32> = PageMap::with_capacity(0x2000.into());
    /// assert_eq!(d.get(0x0.into()), None);
    /// assert_eq!(d.get(0x1000.into()), None);
    ///
    /// d.map(0x1000.into(), &[0x1; 0x1000]).expect("failed to map");
    /// assert_eq!(d.get(0x0.into()), None);
    /// assert_eq!(d.get(0x1000.into()), Some(0x1));
    ///
    /// d.map(0x0.into(), &[0x2; 0x2000]).expect("failed to map");
    ///  assert_eq!(d.get(0x0.into()), Some(0x2));
    ///  assert_eq!(d.get(0x1000.into()), Some(0x2));
    /// ```
    pub fn get(&self, rva: RVA) -> Option<T> {
        if page(rva) > self.pages.len() {
            return None;
        }

        let page = match &self.pages[page(rva)] {
            // page is not mapped
            None => return None,
            // page is mapped
            Some(page) => page,
        };

        Some(page.elements[page_offset(rva)])
    }

    /// fetch one mutable item from the given address.
    ///
    /// ```
    /// use lancelot::arch::RVA;
    /// use lancelot::pagemap::PageMap;
    ///
    /// let mut d: PageMap<u32> = PageMap::with_capacity(0x2000.into());
    /// d.map_empty(0x0.into(), 0x1000).expect("failed to map");
    ///
    /// // address 0x0 starts at 0
    /// assert_eq!(d.get(0x0.into()), Some(0x0));
    ///
    /// // set address 0x0 to 1
    /// let v = d.get_mut(0x0.into()).expect("should be mapped");
    /// *v = 1;
    ///
    /// // address 0x0 is 1
    /// assert_eq!(d.get(0x0.into()), Some(0x1));
    /// ```
    pub fn get_mut(&mut self, rva: RVA) -> Option<&mut T> {
        if page(rva) > self.pages.len() {
            return None;
        }

        let page = match &mut self.pages[page(rva)] {
            // page is not mapped
            None => return None,
            // page is mapped
            Some(page) => page,
        };

        Some(&mut page.elements[page_offset(rva)])
    }

    /// handle the simple slice case: when start and end fall within the same page.
    /// for example, reading a dword from address 0x10.
    ///
    /// ```
    /// use lancelot::arch::RVA;
    /// use lancelot::pagemap::PageMap;
    ///
    /// let mut d: PageMap<u32> = PageMap::with_capacity(0x2000.into());
    /// d.map_empty(0x0.into(), 0x1000).expect("failed to map");
    /// assert_eq!(d.slice(0x0.into(), 0x2.into()).unwrap(), [0x0, 0x0]);
    /// ```
    fn slice_into_simple(&self, start: RVA, buf: &mut [T]) -> Result<(), Error> {
        // precondition: page(start) == page(start + buf.len())

        if page(start) > self.pages.len() {
            return Err(Error::NotMapped);
        }

        let page = match &self.pages[page(start)] {
            // page is not mapped
            None => return Err(Error::NotMapped),
            // page is mapped
            Some(page) => page,
        };

        let end = start + buf.len();
        let elements = &page.elements[page_offset(start)..page_offset(end)];
        buf.copy_from_slice(elements);

        Ok(())
    }

    /// handle the complex slice case: when start and end are on different pages.
    /// for example, reading a dword from address 0xFFE.
    ///
    /// ```
    /// use lancelot::arch::RVA;
    /// use lancelot::pagemap::PageMap;
    ///
    /// let mut d: PageMap<u32> = PageMap::with_capacity(0x5000.into());
    /// d.map_empty(0x1000.into(), 0x3000).expect("failed to map");
    ///
    /// // 0     unmapped
    /// //       unmapped
    /// // 1000  0 0 0 0
    /// //       0 0 0 0
    /// // 2000  0 0 0 0
    /// //       0 0 0 0
    /// // 3000  0 0 0 0
    /// //       0 0 0 0
    /// // 4000  unmapped
    /// //       unmapped
    /// // 5000  unmapped
    ///
    /// assert_eq!(d.slice(0x1FFC.into(), 0x2000.into()).unwrap(), [0x0, 0x0, 0x0, 0x0], "no overlap");
    /// assert_eq!(d.slice(0x1FFD.into(), 0x2001.into()).unwrap(), [0x0, 0x0, 0x0, 0x0], "overlap 1");
    /// assert_eq!(d.slice(0x1FFE.into(), 0x2002.into()).unwrap(), [0x0, 0x0, 0x0, 0x0], "overlap 2");
    /// assert_eq!(d.slice(0x1FFF.into(), 0x2003.into()).unwrap(), [0x0, 0x0, 0x0, 0x0], "overlap 3");
    /// assert_eq!(d.slice(0x2000.into(), 0x2004.into()).unwrap(), [0x0, 0x0, 0x0, 0x0], "overlap 4");
    ///
    /// assert_eq!(d.slice(0x1FFC.into(), 0x3004.into()).unwrap().len(), 0x1008, "4, page, 4");
    ///
    /// assert_eq!(d.slice(0x1FFC.into(), 0x3000.into()).unwrap().len(), 0x1004, "4, page");
    ///
    /// assert_eq!(d.slice(0x2000.into(), 0x3004.into()).unwrap().len(), 0x1004, "page, 4");
    /// ```
    fn slice_into_split(&self, start: RVA, buf: &mut [T]) -> Result<(), Error> {
        let end = start + buf.len();

        if page(end) > self.pages.len() {
            return Err(Error::NotMapped);
        }

        // ensure each page within the requested region is mapped.
        let start_page = page(start);
        let end_page = if page_offset(end) == 0 { page(end) - 1} else { page(end) };
        for page in start_page..end_page {
            if ! self.probe((page * PAGE_SIZE).into()) {
                return Err(Error::NotMapped);
            }
        }

        let mut offset: usize = 0;

        // region one: from `start` to the end of its page
        // region two: any intermediate complete pages
        // region three: from start of final page until `end`

        // one.
        {
            let page = self.pages[page(start)].as_ref().unwrap();
            let elements = &page.elements[page_offset(start)..];
            {
                let dst = &mut buf[offset..offset+elements.len()];
                dst.copy_from_slice(elements);
                offset += elements.len();
            }
        }

        // two.
        if page(start) != page(end) - 1 {
            let start_index = page(start) + 1;
            let end_index = page(end);
            for page_index in start_index..end_index {
                let page = self.pages[page_index].as_ref().unwrap();
                let elements = &page.elements[..];
                {
                    let dst = &mut buf[offset..offset+elements.len()];
                    dst.copy_from_slice(elements);
                    offset += elements.len();
                }
            }
        }

        // three.
        if page_offset(end) != 0x0 {
            let page = self.pages[page(end)].as_ref().unwrap();
            let elements = &page.elements[..page_offset(end)];
            {
                let dst = &mut buf[offset..offset+elements.len()];
                dst.copy_from_slice(elements);
            }
        }

        Ok(())
    }

    /// fetch the items found in the given range, placing them into the given slice.
    /// compared with `slice`, this routine avoids an allocation.
    ///
    /// errors:
    ///   - Error::NotMapped: if any requested address is not mapped
    pub fn slice_into(&self, start: RVA, buf: &mut [T]) -> Result<(), Error> {
        let end = start + buf.len();
        if page(start) == page(end) {
            self.slice_into_simple(start, buf)
        } else {
            self.slice_into_split(start, buf)
        }
    }

    /// fetch the items found in the given range.
    ///
    /// errors:
    ///   - Error::NotMapped: if any requested address is not mapped
    ///
    /// panic if:
    ///   - start > end
    ///
    /// ```
    /// use lancelot::arch::RVA;
    /// use lancelot::pagemap::PageMap;
    ///
    /// let mut d: PageMap<u32> = PageMap::with_capacity(0x2000.into());
    /// d.map_empty(0x0.into(), 0x1000).expect("failed to map");
    /// assert_eq!(d.slice(0x0.into(), 0x2.into()).unwrap(), [0x0, 0x0]);
    /// ```
    pub fn slice(&self, start: RVA, end: RVA) -> Result<Vec<T>, Error> {
        if start > end {
            panic!("start > end");
        }

        let mut ret = vec![Default::default(); (end - start).into()];
        self.slice_into(start, &mut ret)?;

        Ok(ret)
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

/// Types for a Cuckoo proof of work and its encapsulation as a fully usable
/// proof of work within a block header.
use std::cmp::max;
use std::ops::{Add, Div, Mul, Sub};
use std::{fmt, iter};

use rand::{thread_rng, Rng};
use serde::{de, Deserialize, Deserializer, Serialize, Serializer};

use consensus::SECOND_POW_SIZESHIFT;
use core::hash::Hashed;
use global;
use ser::{self, Readable, Reader, Writeable, Writer};

/// The difficulty is defined as the maximum target divided by the block hash.
#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord)]
pub struct Difficulty {
	num: u64,
}

impl Difficulty {
	/// Difficulty of zero, which is invalid (no target can be
	/// calculated from it) but very useful as a start for additions.
	pub fn zero() -> Difficulty {
		Difficulty { num: 0 }
	}

	/// Difficulty of one, which is the minimum difficulty
	/// (when the hash equals the max target)
	pub fn one() -> Difficulty {
		Difficulty { num: 1 }
	}

	/// Convert a `u32` into a `Difficulty`
	pub fn from_num(num: u64) -> Difficulty {
		// can't have difficulty lower than 1
		Difficulty { num: max(num, 1) }
	}

	/// Computes the difficulty from a hash. Divides the maximum target by the
	/// provided hash and applies the Cuckoo sizeshift adjustment factor (see
	/// https://lists.launchpad.net/mimblewimble/msg00494.html).
	pub fn from_proof_adjusted(proof: &Proof) -> Difficulty {
		let max_target = <u64>::max_value();
		let target = proof.hash().to_u64();
		let shift = proof.cuckoo_sizeshift;

		// Adjust the difficulty based on a 2^(N-M)*(N-1) factor, with M being
		// the minimum sizeshift and N the provided sizeshift
		let adjust_factor = (1 << (shift - global::ref_sizeshift()) as u64) * (shift as u64 - 1);
		let difficulty = (max_target / target) * adjust_factor;

		Difficulty::from_num(difficulty)
	}

	/// Same as `from_proof_adjusted` but instead of an adjustment based on
	/// cycle size, scales based on a provided factor. Used by dual PoW system
	/// to scale one PoW against the other.
	pub fn from_proof_scaled(proof: &Proof, scaling: u64) -> Difficulty {
		let max_target = <u64>::max_value();
		let target = proof.hash().to_u64();

		// Scaling between 2 proof of work algos
		Difficulty::from_num((max_target / scaling) / target)
	}

	/// Converts the difficulty into a u64
	pub fn to_num(&self) -> u64 {
		self.num
	}
}

impl fmt::Display for Difficulty {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		write!(f, "{}", self.num)
	}
}

impl Add<Difficulty> for Difficulty {
	type Output = Difficulty;
	fn add(self, other: Difficulty) -> Difficulty {
		Difficulty {
			num: self.num + other.num,
		}
	}
}

impl Sub<Difficulty> for Difficulty {
	type Output = Difficulty;
	fn sub(self, other: Difficulty) -> Difficulty {
		Difficulty {
			num: self.num - other.num,
		}
	}
}

impl Mul<Difficulty> for Difficulty {
	type Output = Difficulty;
	fn mul(self, other: Difficulty) -> Difficulty {
		Difficulty {
			num: self.num * other.num,
		}
	}
}

impl Div<Difficulty> for Difficulty {
	type Output = Difficulty;
	fn div(self, other: Difficulty) -> Difficulty {
		Difficulty {
			num: self.num / other.num,
		}
	}
}

impl Writeable for Difficulty {
	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), ser::Error> {
		writer.write_u64(self.num)
	}
}

impl Readable for Difficulty {
	fn read(reader: &mut Reader) -> Result<Difficulty, ser::Error> {
		let data = reader.read_u64()?;
		Ok(Difficulty { num: data })
	}
}

impl Serialize for Difficulty {
	fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
	where
		S: Serializer,
	{
		serializer.serialize_u64(self.num)
	}
}

impl<'de> Deserialize<'de> for Difficulty {
	fn deserialize<D>(deserializer: D) -> Result<Difficulty, D::Error>
	where
		D: Deserializer<'de>,
	{
		deserializer.deserialize_u64(DiffVisitor)
	}
}

struct DiffVisitor;

impl<'de> de::Visitor<'de> for DiffVisitor {
	type Value = Difficulty;

	fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
		formatter.write_str("a difficulty")
	}

	fn visit_str<E>(self, s: &str) -> Result<Self::Value, E>
	where
		E: de::Error,
	{
		let num_in = s.parse::<u64>();
		if num_in.is_err() {
			return Err(de::Error::invalid_value(
				de::Unexpected::Str(s),
				&"a value number",
			));
		};
		Ok(Difficulty {
			num: num_in.unwrap(),
		})
	}

	fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E>
	where
		E: de::Error,
	{
		Ok(Difficulty { num: value })
	}
}

/// Block header information pertaining to the proof of work
#[derive(Clone, Debug, PartialEq)]
pub struct ProofOfWork {
	/// Total accumulated difficulty since genesis block
	pub total_difficulty: Difficulty,
	/// Difficulty scaling factor between the different proofs of work
	pub scaling_difficulty: u64,
	/// Nonce increment used to mine this block.
	pub nonce: u64,
	/// Proof of work data.
	pub proof: Proof,
}

impl Default for ProofOfWork {
	fn default() -> ProofOfWork {
		let proof_size = global::proofsize();
		ProofOfWork {
			total_difficulty: Difficulty::one(),
			scaling_difficulty: 1,
			nonce: 0,
			proof: Proof::zero(proof_size),
		}
	}
}

impl ProofOfWork {
	/// Read implementation, can't define as trait impl as we need a version
	pub fn read(ver: u16, reader: &mut Reader) -> Result<ProofOfWork, ser::Error> {
		let (total_difficulty, scaling_difficulty) = if ver == 1 {
			// read earlier in the header on older versions
			(Difficulty::one(), 1)
		} else {
			(Difficulty::read(reader)?, reader.read_u64()?)
		};
		let nonce = reader.read_u64()?;
		let proof = Proof::read(reader)?;
		Ok(ProofOfWork {
			total_difficulty,
			scaling_difficulty,
			nonce,
			proof,
		})
	}

	/// Write implementation, can't define as trait impl as we need a version
	pub fn write<W: Writer>(&self, ver: u16, writer: &mut W) -> Result<(), ser::Error> {
		if writer.serialization_mode() != ser::SerializationMode::Hash {
			self.write_pre_pow(ver, writer)?;
			writer.write_u64(self.nonce)?;
		}

		self.proof.write(writer)?;
		Ok(())
	}

	/// Write the pre-hash portion of the header
	pub fn write_pre_pow<W: Writer>(&self, ver: u16, writer: &mut W) -> Result<(), ser::Error> {
		if ver > 1 {
			ser_multiwrite!(
				writer,
				[write_u64, self.total_difficulty.to_num()],
				[write_u64, self.scaling_difficulty]
			);
		}
		Ok(())
	}

	/// Maximum difficulty this proof of work can achieve
	pub fn to_difficulty(&self) -> Difficulty {
		// 2 proof of works, Cuckoo29 (for now) and Cuckoo30+, which are scaled
		// differently (scaling not controlled for now)
		if self.proof.cuckoo_sizeshift == SECOND_POW_SIZESHIFT {
			Difficulty::from_proof_scaled(&self.proof, self.scaling_difficulty)
		} else {
			Difficulty::from_proof_adjusted(&self.proof)
		}
	}

	/// The shift used for the cuckoo cycle size on this proof
	pub fn cuckoo_sizeshift(&self) -> u8 {
		self.proof.cuckoo_sizeshift
	}
}

/// A Cuckoo Cycle proof of work, consisting of the shift to get the graph
/// size (i.e. 31 for Cuckoo31 with a 2^31 or 1<<31 graph size) and the nonces
/// of the graph solution. While being expressed as u64 for simplicity, each
/// nonce is strictly less than half the cycle size (i.e. <2^30 for Cuckoo 31).
///
/// The hash of the `Proof` is the hash of its packed nonces when serializing
/// them at their exact bit size. The resulting bit sequence is padded to be
/// byte-aligned.
///
#[derive(Clone, PartialOrd, PartialEq)]
pub struct Proof {
	/// Power of 2 used for the size of the cuckoo graph
	pub cuckoo_sizeshift: u8,
	/// The nonces
	pub nonces: Vec<u64>,
}

impl fmt::Debug for Proof {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		write!(f, "Cuckoo{}(", self.cuckoo_sizeshift)?;
		for (i, val) in self.nonces[..].iter().enumerate() {
			write!(f, "{:x}", val)?;
			if i < self.nonces.len() - 1 {
				write!(f, " ")?;
			}
		}
		write!(f, ")")
	}
}

impl Eq for Proof {}

impl Proof {
	/// Builds a proof with provided nonces at default sizeshift
	pub fn new(mut in_nonces: Vec<u64>) -> Proof {
		in_nonces.sort();
		Proof {
			cuckoo_sizeshift: global::min_sizeshift(),
			nonces: in_nonces,
		}
	}

	/// Builds a proof with all bytes zeroed out
	pub fn zero(proof_size: usize) -> Proof {
		Proof {
			cuckoo_sizeshift: global::min_sizeshift(),
			nonces: vec![0; proof_size],
		}
	}

	/// Builds a proof with random POW data,
	/// needed so that tests that ignore POW
	/// don't fail due to duplicate hashes
	pub fn random(proof_size: usize) -> Proof {
		let sizeshift = global::min_sizeshift();
		let nonce_mask = (1 << (sizeshift - 1)) - 1;
		let mut rng = thread_rng();
		// force the random num to be within sizeshift bits
		let mut v: Vec<u64> = iter::repeat(())
			.map(|()| (rng.gen::<u32>() & nonce_mask) as u64)
			.take(proof_size)
			.collect();
		v.sort();
		Proof {
			cuckoo_sizeshift: global::min_sizeshift(),
			nonces: v,
		}
	}

	/// Returns the proof size
	pub fn proof_size(&self) -> usize {
		self.nonces.len()
	}
}

impl Readable for Proof {
	fn read(reader: &mut Reader) -> Result<Proof, ser::Error> {
		let cuckoo_sizeshift = reader.read_u8()?;
		if cuckoo_sizeshift == 0 || cuckoo_sizeshift > 64 {
			return Err(ser::Error::CorruptedData);
		}

		let mut nonces = Vec::with_capacity(global::proofsize());
		let nonce_bits = cuckoo_sizeshift as usize - 1;
		let bytes_len = BitVec::bytes_len(nonce_bits * global::proofsize());
		let bits = reader.read_fixed_bytes(bytes_len)?;
		let bitvec = BitVec { bits };
		for n in 0..global::proofsize() {
			let mut nonce = 0;
			for bit in 0..nonce_bits {
				if bitvec.bit_at(n * nonce_bits + (bit as usize)) {
					nonce |= 1 << bit;
				}
			}
			nonces.push(nonce);
		}
		Ok(Proof {
			cuckoo_sizeshift,
			nonces,
		})
	}
}

impl Writeable for Proof {
	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), ser::Error> {
		if writer.serialization_mode() != ser::SerializationMode::Hash {
			writer.write_u8(self.cuckoo_sizeshift)?;
		}

		let nonce_bits = self.cuckoo_sizeshift as usize - 1;
		let mut bitvec = BitVec::new(nonce_bits * global::proofsize());
		for (n, nonce) in self.nonces.iter().enumerate() {
			for bit in 0..nonce_bits {
				if nonce & (1 << bit) != 0 {
					bitvec.set_bit_at(n * nonce_bits + (bit as usize))
				}
			}
		}
		writer.write_fixed_bytes(&bitvec.bits)?;
		Ok(())
	}
}

// TODO this could likely be optimized by writing whole bytes (or even words)
// in the `BitVec` at once, dealing with the truncation, instead of bits by bits
struct BitVec {
	bits: Vec<u8>,
}

impl BitVec {
	/// Number of bytes required to store the provided number of bits
	fn bytes_len(bits_len: usize) -> usize {
		(bits_len + 7) / 8
	}

	fn new(bits_len: usize) -> BitVec {
		BitVec {
			bits: vec![0; BitVec::bytes_len(bits_len)],
		}
	}

	fn set_bit_at(&mut self, pos: usize) {
		self.bits[pos / 8] |= 1 << (pos % 8) as u8;
	}

	fn bit_at(&self, pos: usize) -> bool {
		self.bits[pos / 8] & (1 << (pos % 8) as u8) != 0
	}
}
<commit_msg>Scaling multiplier fix and slightly clearer structure<commit_after>// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

/// Types for a Cuckoo proof of work and its encapsulation as a fully usable
/// proof of work within a block header.
use std::cmp::max;
use std::ops::{Add, Div, Mul, Sub};
use std::{fmt, iter};

use rand::{thread_rng, Rng};
use serde::{de, Deserialize, Deserializer, Serialize, Serializer};

use consensus::SECOND_POW_SIZESHIFT;
use core::hash::Hashed;
use global;
use ser::{self, Readable, Reader, Writeable, Writer};

/// The difficulty is defined as the maximum target divided by the block hash.
#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord)]
pub struct Difficulty {
	num: u64,
}

impl Difficulty {
	/// Difficulty of zero, which is invalid (no target can be
	/// calculated from it) but very useful as a start for additions.
	pub fn zero() -> Difficulty {
		Difficulty { num: 0 }
	}

	/// Difficulty of one, which is the minimum difficulty
	/// (when the hash equals the max target)
	pub fn one() -> Difficulty {
		Difficulty { num: 1 }
	}

	/// Convert a `u32` into a `Difficulty`
	pub fn from_num(num: u64) -> Difficulty {
		// can't have difficulty lower than 1
		Difficulty { num: max(num, 1) }
	}

	/// Computes the difficulty from a hash. Divides the maximum target by the
	/// provided hash and applies the Cuckoo sizeshift adjustment factor (see
	/// https://lists.launchpad.net/mimblewimble/msg00494.html).
	pub fn from_proof_adjusted(proof: &Proof) -> Difficulty {
		// Adjust the difficulty based on a 2^(N-M)*(N-1) factor, with M being
		// the minimum sizeshift and N the provided sizeshift
		let shift = proof.cuckoo_sizeshift;
		let adjust_factor = (1 << (shift - global::ref_sizeshift()) as u64) * (shift as u64 - 1);

		Difficulty::from_num(Difficulty::raw_difficulty(proof) * adjust_factor)
	}

	/// Same as `from_proof_adjusted` but instead of an adjustment based on
	/// cycle size, scales based on a provided factor. Used by dual PoW system
	/// to scale one PoW against the other.
	pub fn from_proof_scaled(proof: &Proof, scaling: u64) -> Difficulty {
		// Scaling between 2 proof of work algos
		Difficulty::from_num(Difficulty::raw_difficulty(proof) * scaling)
	}

	fn raw_difficulty(proof: &Proof) -> u64 {
		<u64>::max_value() / proof.hash().to_u64()
	}

	/// Converts the difficulty into a u64
	pub fn to_num(&self) -> u64 {
		self.num
	}
}

impl fmt::Display for Difficulty {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		write!(f, "{}", self.num)
	}
}

impl Add<Difficulty> for Difficulty {
	type Output = Difficulty;
	fn add(self, other: Difficulty) -> Difficulty {
		Difficulty {
			num: self.num + other.num,
		}
	}
}

impl Sub<Difficulty> for Difficulty {
	type Output = Difficulty;
	fn sub(self, other: Difficulty) -> Difficulty {
		Difficulty {
			num: self.num - other.num,
		}
	}
}

impl Mul<Difficulty> for Difficulty {
	type Output = Difficulty;
	fn mul(self, other: Difficulty) -> Difficulty {
		Difficulty {
			num: self.num * other.num,
		}
	}
}

impl Div<Difficulty> for Difficulty {
	type Output = Difficulty;
	fn div(self, other: Difficulty) -> Difficulty {
		Difficulty {
			num: self.num / other.num,
		}
	}
}

impl Writeable for Difficulty {
	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), ser::Error> {
		writer.write_u64(self.num)
	}
}

impl Readable for Difficulty {
	fn read(reader: &mut Reader) -> Result<Difficulty, ser::Error> {
		let data = reader.read_u64()?;
		Ok(Difficulty { num: data })
	}
}

impl Serialize for Difficulty {
	fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
	where
		S: Serializer,
	{
		serializer.serialize_u64(self.num)
	}
}

impl<'de> Deserialize<'de> for Difficulty {
	fn deserialize<D>(deserializer: D) -> Result<Difficulty, D::Error>
	where
		D: Deserializer<'de>,
	{
		deserializer.deserialize_u64(DiffVisitor)
	}
}

struct DiffVisitor;

impl<'de> de::Visitor<'de> for DiffVisitor {
	type Value = Difficulty;

	fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
		formatter.write_str("a difficulty")
	}

	fn visit_str<E>(self, s: &str) -> Result<Self::Value, E>
	where
		E: de::Error,
	{
		let num_in = s.parse::<u64>();
		if num_in.is_err() {
			return Err(de::Error::invalid_value(
				de::Unexpected::Str(s),
				&"a value number",
			));
		};
		Ok(Difficulty {
			num: num_in.unwrap(),
		})
	}

	fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E>
	where
		E: de::Error,
	{
		Ok(Difficulty { num: value })
	}
}

/// Block header information pertaining to the proof of work
#[derive(Clone, Debug, PartialEq)]
pub struct ProofOfWork {
	/// Total accumulated difficulty since genesis block
	pub total_difficulty: Difficulty,
	/// Difficulty scaling factor between the different proofs of work
	pub scaling_difficulty: u64,
	/// Nonce increment used to mine this block.
	pub nonce: u64,
	/// Proof of work data.
	pub proof: Proof,
}

impl Default for ProofOfWork {
	fn default() -> ProofOfWork {
		let proof_size = global::proofsize();
		ProofOfWork {
			total_difficulty: Difficulty::one(),
			scaling_difficulty: 1,
			nonce: 0,
			proof: Proof::zero(proof_size),
		}
	}
}

impl ProofOfWork {
	/// Read implementation, can't define as trait impl as we need a version
	pub fn read(ver: u16, reader: &mut Reader) -> Result<ProofOfWork, ser::Error> {
		let (total_difficulty, scaling_difficulty) = if ver == 1 {
			// read earlier in the header on older versions
			(Difficulty::one(), 1)
		} else {
			(Difficulty::read(reader)?, reader.read_u64()?)
		};
		let nonce = reader.read_u64()?;
		let proof = Proof::read(reader)?;
		Ok(ProofOfWork {
			total_difficulty,
			scaling_difficulty,
			nonce,
			proof,
		})
	}

	/// Write implementation, can't define as trait impl as we need a version
	pub fn write<W: Writer>(&self, ver: u16, writer: &mut W) -> Result<(), ser::Error> {
		if writer.serialization_mode() != ser::SerializationMode::Hash {
			self.write_pre_pow(ver, writer)?;
			writer.write_u64(self.nonce)?;
		}

		self.proof.write(writer)?;
		Ok(())
	}

	/// Write the pre-hash portion of the header
	pub fn write_pre_pow<W: Writer>(&self, ver: u16, writer: &mut W) -> Result<(), ser::Error> {
		if ver > 1 {
			ser_multiwrite!(
				writer,
				[write_u64, self.total_difficulty.to_num()],
				[write_u64, self.scaling_difficulty]
			);
		}
		Ok(())
	}

	/// Maximum difficulty this proof of work can achieve
	pub fn to_difficulty(&self) -> Difficulty {
		// 2 proof of works, Cuckoo29 (for now) and Cuckoo30+, which are scaled
		// differently (scaling not controlled for now)
		if self.proof.cuckoo_sizeshift == SECOND_POW_SIZESHIFT {
			Difficulty::from_proof_scaled(&self.proof, self.scaling_difficulty)
		} else {
			Difficulty::from_proof_adjusted(&self.proof)
		}
	}

	/// The shift used for the cuckoo cycle size on this proof
	pub fn cuckoo_sizeshift(&self) -> u8 {
		self.proof.cuckoo_sizeshift
	}
}

/// A Cuckoo Cycle proof of work, consisting of the shift to get the graph
/// size (i.e. 31 for Cuckoo31 with a 2^31 or 1<<31 graph size) and the nonces
/// of the graph solution. While being expressed as u64 for simplicity, each
/// nonce is strictly less than half the cycle size (i.e. <2^30 for Cuckoo 31).
///
/// The hash of the `Proof` is the hash of its packed nonces when serializing
/// them at their exact bit size. The resulting bit sequence is padded to be
/// byte-aligned.
///
#[derive(Clone, PartialOrd, PartialEq)]
pub struct Proof {
	/// Power of 2 used for the size of the cuckoo graph
	pub cuckoo_sizeshift: u8,
	/// The nonces
	pub nonces: Vec<u64>,
}

impl fmt::Debug for Proof {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		write!(f, "Cuckoo{}(", self.cuckoo_sizeshift)?;
		for (i, val) in self.nonces[..].iter().enumerate() {
			write!(f, "{:x}", val)?;
			if i < self.nonces.len() - 1 {
				write!(f, " ")?;
			}
		}
		write!(f, ")")
	}
}

impl Eq for Proof {}

impl Proof {
	/// Builds a proof with provided nonces at default sizeshift
	pub fn new(mut in_nonces: Vec<u64>) -> Proof {
		in_nonces.sort();
		Proof {
			cuckoo_sizeshift: global::min_sizeshift(),
			nonces: in_nonces,
		}
	}

	/// Builds a proof with all bytes zeroed out
	pub fn zero(proof_size: usize) -> Proof {
		Proof {
			cuckoo_sizeshift: global::min_sizeshift(),
			nonces: vec![0; proof_size],
		}
	}

	/// Builds a proof with random POW data,
	/// needed so that tests that ignore POW
	/// don't fail due to duplicate hashes
	pub fn random(proof_size: usize) -> Proof {
		let sizeshift = global::min_sizeshift();
		let nonce_mask = (1 << (sizeshift - 1)) - 1;
		let mut rng = thread_rng();
		// force the random num to be within sizeshift bits
		let mut v: Vec<u64> = iter::repeat(())
			.map(|()| (rng.gen::<u32>() & nonce_mask) as u64)
			.take(proof_size)
			.collect();
		v.sort();
		Proof {
			cuckoo_sizeshift: global::min_sizeshift(),
			nonces: v,
		}
	}

	/// Returns the proof size
	pub fn proof_size(&self) -> usize {
		self.nonces.len()
	}
}

impl Readable for Proof {
	fn read(reader: &mut Reader) -> Result<Proof, ser::Error> {
		let cuckoo_sizeshift = reader.read_u8()?;
		if cuckoo_sizeshift == 0 || cuckoo_sizeshift > 64 {
			return Err(ser::Error::CorruptedData);
		}

		let mut nonces = Vec::with_capacity(global::proofsize());
		let nonce_bits = cuckoo_sizeshift as usize - 1;
		let bytes_len = BitVec::bytes_len(nonce_bits * global::proofsize());
		let bits = reader.read_fixed_bytes(bytes_len)?;
		let bitvec = BitVec { bits };
		for n in 0..global::proofsize() {
			let mut nonce = 0;
			for bit in 0..nonce_bits {
				if bitvec.bit_at(n * nonce_bits + (bit as usize)) {
					nonce |= 1 << bit;
				}
			}
			nonces.push(nonce);
		}
		Ok(Proof {
			cuckoo_sizeshift,
			nonces,
		})
	}
}

impl Writeable for Proof {
	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), ser::Error> {
		if writer.serialization_mode() != ser::SerializationMode::Hash {
			writer.write_u8(self.cuckoo_sizeshift)?;
		}

		let nonce_bits = self.cuckoo_sizeshift as usize - 1;
		let mut bitvec = BitVec::new(nonce_bits * global::proofsize());
		for (n, nonce) in self.nonces.iter().enumerate() {
			for bit in 0..nonce_bits {
				if nonce & (1 << bit) != 0 {
					bitvec.set_bit_at(n * nonce_bits + (bit as usize))
				}
			}
		}
		writer.write_fixed_bytes(&bitvec.bits)?;
		Ok(())
	}
}

// TODO this could likely be optimized by writing whole bytes (or even words)
// in the `BitVec` at once, dealing with the truncation, instead of bits by bits
struct BitVec {
	bits: Vec<u8>,
}

impl BitVec {
	/// Number of bytes required to store the provided number of bits
	fn bytes_len(bits_len: usize) -> usize {
		(bits_len + 7) / 8
	}

	fn new(bits_len: usize) -> BitVec {
		BitVec {
			bits: vec![0; BitVec::bytes_len(bits_len)],
		}
	}

	fn set_bit_at(&mut self, pos: usize) {
		self.bits[pos / 8] |= 1 << (pos % 8) as u8;
	}

	fn bit_at(&self, pos: usize) -> bool {
		self.bits[pos / 8] & (1 << (pos % 8) as u8) != 0
	}
}
<|endoftext|>


<commit_before>use flowgger::config::Config;
use flowgger::decoder::Decoder;
use flowgger::encoder::Encoder;
use flate2::FlateReadExt;
use std::io::{stderr, Read, Write};
use std::net::UdpSocket;
use std::str;
use std::sync::mpsc::SyncSender;
use super::Input;

const DEFAULT_LISTEN: &'static str = "0.0.0.0:514";
const MAX_UDP_PACKET_SIZE: usize = 65527;
const MAX_COMPRESSION_RATIO: usize = 5;

pub struct UdpInput {
    listen: String,
}

impl UdpInput {
    pub fn new(config: &Config) -> UdpInput {
        let listen = config.lookup("input.listen")
            .map_or(DEFAULT_LISTEN,
                    |x| x.as_str().expect("input.listen must be an ip:port string"))
            .to_owned();
        UdpInput { listen: listen }
    }
}

impl Input for UdpInput {
    fn accept(&self,
              tx: SyncSender<Vec<u8>>,
              decoder: Box<Decoder + Send>,
              encoder: Box<Encoder + Send>) {
        let socket = UdpSocket::bind(&self.listen as &str)
            .expect(&format!("Unable to listen to {}", self.listen));
        let tx = tx.clone();
        let (decoder, encoder): (Box<Decoder>, Box<Encoder>) = (decoder.clone_boxed(),
                                                                encoder.clone_boxed());
        let mut buf = [0; MAX_UDP_PACKET_SIZE];
        loop {
            let (length, _src) = match socket.recv_from(&mut buf) {
                Ok(res) => res,
                Err(_) => continue,
            };
            let line = &buf[..length];
            if let Err(e) = handle_record_maybe_compressed(&line, &tx, &decoder, &encoder) {
                let _ = writeln!(stderr(), "{}", e);
            }
        }
    }
}

fn handle_record_maybe_compressed(line: &[u8],
                                  tx: &SyncSender<Vec<u8>>,
                                  decoder: &Box<Decoder>,
                                  encoder: &Box<Encoder>)
                                  -> Result<(), &'static str> {
    if line.len() > 2 &&
       (line[0] == 0x78 && (line[1] == 0x01 || line[1] == 0x9c || line[1] == 0xda)) {
        let mut decompressed = Vec::with_capacity(MAX_UDP_PACKET_SIZE * MAX_COMPRESSION_RATIO);
        match line.zlib_decode().read_to_end(&mut decompressed) {
            Ok(_) => handle_record(&decompressed, tx, decoder, encoder),
            Err(_) => return Err("Corrupted compressed (zlib) record"),
        }
    } else if line.len() > 3 && (line[0] == 0x1f && line[1] == 0x8b && line[2] == 0x08) {
        let mut decompressed = Vec::with_capacity(MAX_UDP_PACKET_SIZE * MAX_COMPRESSION_RATIO);
        match line.gz_decode().and_then(|mut x| x.read_to_end(&mut decompressed)) {
            Ok(_) => handle_record(&decompressed, tx, decoder, encoder),
            Err(_) => return Err("Corrupted compressed (gzip) record"),
        }
    } else {
        handle_record(line, tx, decoder, encoder)
    }
}

fn handle_record(line: &[u8],
                 tx: &SyncSender<Vec<u8>>,
                 decoder: &Box<Decoder>,
                 encoder: &Box<Encoder>)
                 -> Result<(), &'static str> {
    let line = match str::from_utf8(line) {
        Err(_) => return Err("Invalid UTF-8 input"),
        Ok(line) => line,
    };
    let decoded = try!(decoder.decode(line));
    let reencoded = try!(encoder.encode(decoded));
    tx.send(reencoded).unwrap();
    Ok(())
}
<commit_msg>Raise the lower bounds for compressed UDP packets<commit_after>use flowgger::config::Config;
use flowgger::decoder::Decoder;
use flowgger::encoder::Encoder;
use flate2::FlateReadExt;
use std::io::{stderr, Read, Write};
use std::net::UdpSocket;
use std::str;
use std::sync::mpsc::SyncSender;
use super::Input;

const DEFAULT_LISTEN: &'static str = "0.0.0.0:514";
const MAX_UDP_PACKET_SIZE: usize = 65527;
const MAX_COMPRESSION_RATIO: usize = 5;

pub struct UdpInput {
    listen: String,
}

impl UdpInput {
    pub fn new(config: &Config) -> UdpInput {
        let listen = config.lookup("input.listen")
            .map_or(DEFAULT_LISTEN,
                    |x| x.as_str().expect("input.listen must be an ip:port string"))
            .to_owned();
        UdpInput { listen: listen }
    }
}

impl Input for UdpInput {
    fn accept(&self,
              tx: SyncSender<Vec<u8>>,
              decoder: Box<Decoder + Send>,
              encoder: Box<Encoder + Send>) {
        let socket = UdpSocket::bind(&self.listen as &str)
            .expect(&format!("Unable to listen to {}", self.listen));
        let tx = tx.clone();
        let (decoder, encoder): (Box<Decoder>, Box<Encoder>) = (decoder.clone_boxed(),
                                                                encoder.clone_boxed());
        let mut buf = [0; MAX_UDP_PACKET_SIZE];
        loop {
            let (length, _src) = match socket.recv_from(&mut buf) {
                Ok(res) => res,
                Err(_) => continue,
            };
            let line = &buf[..length];
            if let Err(e) = handle_record_maybe_compressed(&line, &tx, &decoder, &encoder) {
                let _ = writeln!(stderr(), "{}", e);
            }
        }
    }
}

fn handle_record_maybe_compressed(line: &[u8],
                                  tx: &SyncSender<Vec<u8>>,
                                  decoder: &Box<Decoder>,
                                  encoder: &Box<Encoder>)
                                  -> Result<(), &'static str> {
    if line.len() >= 8 &&
       (line[0] == 0x78 && (line[1] == 0x01 || line[1] == 0x9c || line[1] == 0xda)) {
        let mut decompressed = Vec::with_capacity(MAX_UDP_PACKET_SIZE * MAX_COMPRESSION_RATIO);
        match line.zlib_decode().read_to_end(&mut decompressed) {
            Ok(_) => handle_record(&decompressed, tx, decoder, encoder),
            Err(_) => return Err("Corrupted compressed (zlib) record"),
        }
    } else if line.len() >= 24 && (line[0] == 0x1f && line[1] == 0x8b && line[2] == 0x08) {
        let mut decompressed = Vec::with_capacity(MAX_UDP_PACKET_SIZE * MAX_COMPRESSION_RATIO);
        match line.gz_decode().and_then(|mut x| x.read_to_end(&mut decompressed)) {
            Ok(_) => handle_record(&decompressed, tx, decoder, encoder),
            Err(_) => return Err("Corrupted compressed (gzip) record"),
        }
    } else {
        handle_record(line, tx, decoder, encoder)
    }
}

fn handle_record(line: &[u8],
                 tx: &SyncSender<Vec<u8>>,
                 decoder: &Box<Decoder>,
                 encoder: &Box<Encoder>)
                 -> Result<(), &'static str> {
    let line = match str::from_utf8(line) {
        Err(_) => return Err("Invalid UTF-8 input"),
        Ok(line) => line,
    };
    let decoded = try!(decoder.decode(line));
    let reencoded = try!(encoder.encode(decoded));
    tx.send(reencoded).unwrap();
    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2014 The Gfx-rs Developers.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt;
use std::from_str::FromStr;
use std::gc::Gc;
use syntax::{ast, ext};
use syntax::ext::build::AstBuilder;
use syntax::ext::deriving::generic;
use syntax::{attr, codemap};
use syntax::parse::token;

/// A component modifier.
#[deriving(PartialEq)]
enum Modifier {
    /// Corresponds to the `#[normalized]` attribute.
    ///
    /// Normalizes the component at runtime. Unsigned integers are normalized to
    /// `[0, 1]`. Signed integers are normalized to `[-1, 1]`.
    Normalized,
    /// Corresponds to the `#[as_float]` attribute.
    ///
    /// Casts the component to a float precision floating-point number at runtime.
    AsFloat,
    /// Corresponds to the `#[as_double]` attribute.
    ///
    /// Specifies a high-precision float.
    AsDouble,
}

impl fmt::Show for Modifier {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Normalized => write!(f, "normalized"),
            AsFloat => write!(f, "as_float"),
            AsDouble => write!(f, "as_double"),
        }
    }
}

impl FromStr for Modifier {
    fn from_str(src: &str) -> Option<Modifier> {
        match src {
            "normalized" => Some(Normalized),
            "as_float" => Some(AsFloat),
            "as_double" => Some(AsDouble),
            _ => None,
        }
    }
}

/// Scan through the field's attributes and extract a relevant modifier. If
/// multiple modifier attributes are found, use the first modifier and emit a
/// warning.
fn find_modifier(cx: &mut ext::base::ExtCtxt, span: codemap::Span,
                 attributes: &[ast::Attribute]) -> Option<Modifier> {
    attributes.iter().fold(None, |modifier, attribute| {
        match attribute.node.value.node {
            ast::MetaWord(ref word) => {
                from_str(word.get()).and_then(|new_modifier| {
                    attr::mark_used(attribute);
                    modifier.map_or(Some(new_modifier), |modifier| {
                        cx.span_warn(span, format!(
                            "Extra attribute modifier detected: `#[{}]` - \
                            ignoring in favour of `#[{}]`.", new_modifier, modifier
                        ).as_slice());
                        None
                    })
                }).or(modifier)
            },
            _ => modifier,
        }
    })
}

/// Find a `gfx::attrib::Type` that describes the given type identifier.
fn decode_type(cx: &mut ext::base::ExtCtxt, span: codemap::Span,
               ty_ident: &ast::Ident, modifier: Option<Modifier>) -> Gc<ast::Expr> {
    let ty_str = ty_ident.name.as_str();
    match ty_str {
        "f32" | "f64" => {
            let kind = cx.ident_of(match modifier {
                None | Some(AsFloat) => "FloatDefault",
                Some(AsDouble) => "FloatPrecision",
                Some(Normalized) => {
                    cx.span_warn(span, format!(
                        "Incompatible float modifier attribute: `#[{}]`", modifier
                    ).as_slice());
                    ""
                }
            });
            let sub_type = cx.ident_of(format!("F{}", ty_str.slice_from(1)).as_slice());
            quote_expr!(cx, gfx::attrib::Float(gfx::attrib::$kind,
                                               gfx::attrib::$sub_type))
        },
        "u8" | "u16" | "u32" | "u64" |
        "i8" | "i16" | "i32" | "i64" => {
            let sign = cx.ident_of({
                if ty_str.starts_with("i") { "Signed" } else { "Unsigned" }
            });
            let kind = cx.ident_of(match modifier {
                None => "IntRaw",
                Some(Normalized) => "IntNormalized",
                Some(AsFloat) => "IntAsFloat",
                Some(AsDouble) => {
                    cx.span_warn(span, format!(
                        "Incompatible int modifier attribute: `#[{}]`", modifier
                    ).as_slice());
                    ""
                }
            });
            let sub_type = cx.ident_of(format!("U{}", ty_str.slice_from(1)).as_slice());
            quote_expr!(cx, gfx::attrib::Int(gfx::attrib::$kind,
                                             gfx::attrib::$sub_type,
                                             gfx::attrib::$sign))
        },
        "uint" | "int" => {
            cx.span_err(span, format!("Pointer-sized integer components are \
                                      not supported, but found: `{}`. Use an \
                                      integer component with an explicit size \
                                      instead.", ty_str).as_slice());
            cx.expr_lit(span, ast::LitNil)
        },
        ty_str => {
            cx.span_err(span, format!("Unrecognized component type: `{}`",
                                      ty_str).as_slice());
            cx.expr_lit(span, ast::LitNil)
        },
    }
}

fn decode_count_and_type(cx: &mut ext::base::ExtCtxt, span: codemap::Span,
                         field: &ast::StructField) -> (Gc<ast::Expr>, Gc<ast::Expr>) {
    let modifier = find_modifier(cx, span, field.node.attrs.as_slice());
    match field.node.ty.node {
        ast::TyPath(ref p, _, _) => (
            cx.expr_lit(span, ast::LitIntUnsuffixed(1)),
            decode_type(cx, span, &p.segments[0].identifier, modifier),
        ),
        ast::TyFixedLengthVec(pty, expr) => (expr, match pty.node {
            ast::TyPath(ref p, _, _) => {
                decode_type(cx, span, &p.segments[0].identifier, modifier)
            },
            _ => {
                cx.span_err(span, format!("Unsupported fixed vector sub-type: \
                                          `{}`",pty.node).as_slice());
                cx.expr_lit(span, ast::LitNil)
            },
        }),
        _ => {
            cx.span_err(span, format!("Unsupported attribute type: `{}`",
                                      field.node.ty.node).as_slice());
            (cx.expr_lit(span, ast::LitNil), cx.expr_lit(span, ast::LitNil))
        },
    }
}

/// Generates the the method body for `gfx::VertexFormat::generate`.
fn method_body(cx: &mut ext::base::ExtCtxt, span: codemap::Span,
                   substr: &generic::Substructure) -> Gc<ast::Expr> {
    match *substr.fields {
        generic::StaticStruct(ref definition, generic::Named(ref fields)) => {
            let attribute_pushes = definition.fields.iter().zip(fields.iter())
                .map(|(def, &(ident, _))| {
                    let struct_ident = substr.type_ident;
                    let buffer_expr = substr.nonself_args[1];
                    let (count_expr, type_expr) = decode_count_and_type(cx, span, def);
                    let ident_str = token::get_ident(ident);
                    let ident_str = ident_str.get();
                    super::ugh(cx, |cx| quote_expr!(cx, {
                        attributes.push(gfx::Attribute {
                            buffer: $buffer_expr,
                            elem_count: $count_expr,
                            elem_type: $type_expr,
                            offset: unsafe {
                                &(*(0u as *const $struct_ident)).$ident as *const _ as gfx::attrib::Offset
                            },
                            stride: { use std::mem; mem::size_of::<$struct_ident>() as gfx::attrib::Stride },
                            name: $ident_str.to_string(),
                        });
                    }))
                }).collect::<Vec<Gc<ast::Expr>>>();
            let capacity = fields.len();
            super::ugh(cx, |cx| quote_expr!(cx, {
                let mut attributes = Vec::with_capacity($capacity);
                $attribute_pushes;
                attributes
            }))
        },
        _ => {
            cx.span_err(span, "Unable to implement `gfx::VertexFormat::generate` \
                              on a non-structure");
            cx.expr_lit(span, ast::LitNil)
        }
    }
}

/// Derive a `gfx::VertexFormat` implementation for the `struct`
pub fn expand(context: &mut ext::base::ExtCtxt, span: codemap::Span,
              meta_item: Gc<ast::MetaItem>, item: Gc<ast::Item>,
              push: |Gc<ast::Item>|) {
    // `impl gfx::VertexFormat for $item`
    generic::TraitDef {
        span: span,
        attributes: Vec::new(),
        path: generic::ty::Path {
            path: vec!["gfx", "VertexFormat"],
            lifetime: None,
            params: Vec::new(),
            global: true,
        },
        additional_bounds: Vec::new(),
        generics: generic::ty::LifetimeBounds::empty(),
        methods: vec![
            // `fn generate(Option<Self>, gfx::BufferHandle) -> Vec<gfx::Attribute>`
            generic::MethodDef {
                name: "generate",
                generics: generic::ty::LifetimeBounds::empty(),
                explicit_self: None,
                args: vec![
                    generic::ty::Literal(generic::ty::Path {
                        path: vec!["Option"],
                        lifetime: None,
                        params: vec![box generic::ty::Self],
                        global: false,
                    }),
                    generic::ty::Literal(generic::ty::Path::new(
                        vec!["gfx", "BufferHandle"]
                    )),
                ],
                ret_ty: generic::ty::Literal(
                    generic::ty::Path {
                        path: vec!["Vec"],
                        lifetime: None,
                        params: vec![
                            box generic::ty::Literal(generic::ty::Path::new(
                                vec!["gfx", "Attribute"])),
                        ],
                        global: false,
                    },
                ),
                attributes: Vec::new(),
                // generate the method body
                combine_substructure: generic::combine_substructure(method_body),
            },
        ],
    }.expand(context, meta_item, item, push);
}
<commit_msg>Change LitIntUnsuffixed to LitInt to reflect rustc change https://github.com/rust-lang/rust/commit/0dc215741b34236c310e409c437600ba0165c97c<commit_after>// Copyright 2014 The Gfx-rs Developers.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt;
use std::from_str::FromStr;
use std::gc::Gc;
use syntax::{ast, ext};
use syntax::ext::build::AstBuilder;
use syntax::ext::deriving::generic;
use syntax::{attr, codemap};
use syntax::parse::token;

/// A component modifier.
#[deriving(PartialEq)]
enum Modifier {
    /// Corresponds to the `#[normalized]` attribute.
    ///
    /// Normalizes the component at runtime. Unsigned integers are normalized to
    /// `[0, 1]`. Signed integers are normalized to `[-1, 1]`.
    Normalized,
    /// Corresponds to the `#[as_float]` attribute.
    ///
    /// Casts the component to a float precision floating-point number at runtime.
    AsFloat,
    /// Corresponds to the `#[as_double]` attribute.
    ///
    /// Specifies a high-precision float.
    AsDouble,
}

impl fmt::Show for Modifier {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Normalized => write!(f, "normalized"),
            AsFloat => write!(f, "as_float"),
            AsDouble => write!(f, "as_double"),
        }
    }
}

impl FromStr for Modifier {
    fn from_str(src: &str) -> Option<Modifier> {
        match src {
            "normalized" => Some(Normalized),
            "as_float" => Some(AsFloat),
            "as_double" => Some(AsDouble),
            _ => None,
        }
    }
}

/// Scan through the field's attributes and extract a relevant modifier. If
/// multiple modifier attributes are found, use the first modifier and emit a
/// warning.
fn find_modifier(cx: &mut ext::base::ExtCtxt, span: codemap::Span,
                 attributes: &[ast::Attribute]) -> Option<Modifier> {
    attributes.iter().fold(None, |modifier, attribute| {
        match attribute.node.value.node {
            ast::MetaWord(ref word) => {
                from_str(word.get()).and_then(|new_modifier| {
                    attr::mark_used(attribute);
                    modifier.map_or(Some(new_modifier), |modifier| {
                        cx.span_warn(span, format!(
                            "Extra attribute modifier detected: `#[{}]` - \
                            ignoring in favour of `#[{}]`.", new_modifier, modifier
                        ).as_slice());
                        None
                    })
                }).or(modifier)
            },
            _ => modifier,
        }
    })
}

/// Find a `gfx::attrib::Type` that describes the given type identifier.
fn decode_type(cx: &mut ext::base::ExtCtxt, span: codemap::Span,
               ty_ident: &ast::Ident, modifier: Option<Modifier>) -> Gc<ast::Expr> {
    let ty_str = ty_ident.name.as_str();
    match ty_str {
        "f32" | "f64" => {
            let kind = cx.ident_of(match modifier {
                None | Some(AsFloat) => "FloatDefault",
                Some(AsDouble) => "FloatPrecision",
                Some(Normalized) => {
                    cx.span_warn(span, format!(
                        "Incompatible float modifier attribute: `#[{}]`", modifier
                    ).as_slice());
                    ""
                }
            });
            let sub_type = cx.ident_of(format!("F{}", ty_str.slice_from(1)).as_slice());
            quote_expr!(cx, gfx::attrib::Float(gfx::attrib::$kind,
                                               gfx::attrib::$sub_type))
        },
        "u8" | "u16" | "u32" | "u64" |
        "i8" | "i16" | "i32" | "i64" => {
            let sign = cx.ident_of({
                if ty_str.starts_with("i") { "Signed" } else { "Unsigned" }
            });
            let kind = cx.ident_of(match modifier {
                None => "IntRaw",
                Some(Normalized) => "IntNormalized",
                Some(AsFloat) => "IntAsFloat",
                Some(AsDouble) => {
                    cx.span_warn(span, format!(
                        "Incompatible int modifier attribute: `#[{}]`", modifier
                    ).as_slice());
                    ""
                }
            });
            let sub_type = cx.ident_of(format!("U{}", ty_str.slice_from(1)).as_slice());
            quote_expr!(cx, gfx::attrib::Int(gfx::attrib::$kind,
                                             gfx::attrib::$sub_type,
                                             gfx::attrib::$sign))
        },
        "uint" | "int" => {
            cx.span_err(span, format!("Pointer-sized integer components are \
                                      not supported, but found: `{}`. Use an \
                                      integer component with an explicit size \
                                      instead.", ty_str).as_slice());
            cx.expr_lit(span, ast::LitNil)
        },
        ty_str => {
            cx.span_err(span, format!("Unrecognized component type: `{}`",
                                      ty_str).as_slice());
            cx.expr_lit(span, ast::LitNil)
        },
    }
}

fn decode_count_and_type(cx: &mut ext::base::ExtCtxt, span: codemap::Span,
                         field: &ast::StructField) -> (Gc<ast::Expr>, Gc<ast::Expr>) {
    let modifier = find_modifier(cx, span, field.node.attrs.as_slice());
    match field.node.ty.node {
        ast::TyPath(ref p, _, _) => (
            cx.expr_lit(span, ast::LitInt(1, ast::UnsuffixedIntLit(ast::Plus))),
            decode_type(cx, span, &p.segments[0].identifier, modifier),
        ),
        ast::TyFixedLengthVec(pty, expr) => (expr, match pty.node {
            ast::TyPath(ref p, _, _) => {
                decode_type(cx, span, &p.segments[0].identifier, modifier)
            },
            _ => {
                cx.span_err(span, format!("Unsupported fixed vector sub-type: \
                                          `{}`",pty.node).as_slice());
                cx.expr_lit(span, ast::LitNil)
            },
        }),
        _ => {
            cx.span_err(span, format!("Unsupported attribute type: `{}`",
                                      field.node.ty.node).as_slice());
            (cx.expr_lit(span, ast::LitNil), cx.expr_lit(span, ast::LitNil))
        },
    }
}

/// Generates the the method body for `gfx::VertexFormat::generate`.
fn method_body(cx: &mut ext::base::ExtCtxt, span: codemap::Span,
                   substr: &generic::Substructure) -> Gc<ast::Expr> {
    match *substr.fields {
        generic::StaticStruct(ref definition, generic::Named(ref fields)) => {
            let attribute_pushes = definition.fields.iter().zip(fields.iter())
                .map(|(def, &(ident, _))| {
                    let struct_ident = substr.type_ident;
                    let buffer_expr = substr.nonself_args[1];
                    let (count_expr, type_expr) = decode_count_and_type(cx, span, def);
                    let ident_str = token::get_ident(ident);
                    let ident_str = ident_str.get();
                    super::ugh(cx, |cx| quote_expr!(cx, {
                        attributes.push(gfx::Attribute {
                            buffer: $buffer_expr,
                            elem_count: $count_expr,
                            elem_type: $type_expr,
                            offset: unsafe {
                                &(*(0u as *const $struct_ident)).$ident as *const _ as gfx::attrib::Offset
                            },
                            stride: { use std::mem; mem::size_of::<$struct_ident>() as gfx::attrib::Stride },
                            name: $ident_str.to_string(),
                        });
                    }))
                }).collect::<Vec<Gc<ast::Expr>>>();
            let capacity = fields.len();
            super::ugh(cx, |cx| quote_expr!(cx, {
                let mut attributes = Vec::with_capacity($capacity);
                $attribute_pushes;
                attributes
            }))
        },
        _ => {
            cx.span_err(span, "Unable to implement `gfx::VertexFormat::generate` \
                              on a non-structure");
            cx.expr_lit(span, ast::LitNil)
        }
    }
}

/// Derive a `gfx::VertexFormat` implementation for the `struct`
pub fn expand(context: &mut ext::base::ExtCtxt, span: codemap::Span,
              meta_item: Gc<ast::MetaItem>, item: Gc<ast::Item>,
              push: |Gc<ast::Item>|) {
    // `impl gfx::VertexFormat for $item`
    generic::TraitDef {
        span: span,
        attributes: Vec::new(),
        path: generic::ty::Path {
            path: vec!["gfx", "VertexFormat"],
            lifetime: None,
            params: Vec::new(),
            global: true,
        },
        additional_bounds: Vec::new(),
        generics: generic::ty::LifetimeBounds::empty(),
        methods: vec![
            // `fn generate(Option<Self>, gfx::BufferHandle) -> Vec<gfx::Attribute>`
            generic::MethodDef {
                name: "generate",
                generics: generic::ty::LifetimeBounds::empty(),
                explicit_self: None,
                args: vec![
                    generic::ty::Literal(generic::ty::Path {
                        path: vec!["Option"],
                        lifetime: None,
                        params: vec![box generic::ty::Self],
                        global: false,
                    }),
                    generic::ty::Literal(generic::ty::Path::new(
                        vec!["gfx", "BufferHandle"]
                    )),
                ],
                ret_ty: generic::ty::Literal(
                    generic::ty::Path {
                        path: vec!["Vec"],
                        lifetime: None,
                        params: vec![
                            box generic::ty::Literal(generic::ty::Path::new(
                                vec!["gfx", "Attribute"])),
                        ],
                        global: false,
                    },
                ),
                attributes: Vec::new(),
                // generate the method body
                combine_substructure: generic::combine_substructure(method_body),
            },
        ],
    }.expand(context, meta_item, item, push);
}
<|endoftext|>


<commit_before>// Copyright 2015 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

use {CommandResult, Host, ProviderFactory, Providers, Result};
use error::Error;
use regex::Regex;
use std::{fs, process, str};
use target::bin_resolver::BinResolver;
use telemetry::{FsMount, Netif, NetifIPv4, NetifIPv6, NetifStatus};

pub fn default_provider(host: &mut Host, providers: Vec<Providers>) -> Result<Providers> {
    for p in providers {
        let provider = ProviderFactory::create(host, Some(p));

        if provider.is_ok() {
            return Ok(provider.unwrap().get_providers());
        }
    }

    Err(Error::Generic("No package providers are available".to_string()))
}

pub fn command_exec(cmd: &str) -> Result<CommandResult> {
    let output = try!(process::Command::new("sh").arg("-c").arg(cmd).output());

    Ok(CommandResult {
        exit_code: output.status.code().unwrap(),
        stdout: str::from_utf8(&output.stdout).unwrap().trim().to_string(),
        stderr: str::from_utf8(&output.stderr).unwrap().trim().to_string(),
    })
}

pub fn file_is_file(path: &str) -> Result<bool> {
    let meta = fs::metadata(path);
    Ok(meta.is_err() || meta.unwrap().is_file())
}

pub fn file_exists(path: &str) -> Result<bool> {
    Ok(fs::metadata(path).is_ok())
}

pub fn file_delete(path: &str) -> Result<()> {
    try!(fs::remove_file(path));
    Ok(())
}

pub fn file_mv(path: &str, new_path: &str) -> Result<()> {
    try!(fs::rename(path, new_path))
}

pub fn file_copy(path: &str, new_path: &str) -> Result<()> {
    try!(fs::copy(path, new_path));
    Ok(())
}

pub fn file_set_owner(path: &str, user: &str, group: &str) -> Result<()> {
    let user_group = format!("{}:{}", user, group);
    let args: Vec<&str> = vec![&user_group, path];
    let output = process::Command::new(&try!(BinResolver::resolve("chown"))).args(&args).output().unwrap();

    if !output.status.success() {
        return Err(Error::Generic(format!("Could not chown file with error: {}", str::from_utf8(&output.stderr).unwrap())));
    }

    Ok(())
}

pub fn file_stat<'a>(path: &'a str, args: Vec<&'a str>) -> Result<String> {
    let mut args = args;
    args.push(path);
    let output = process::Command::new(&try!(BinResolver::resolve("stat"))).args(&args).output().unwrap();

    if !output.status.success() {
        return Err(Error::Generic(format!("Could not stat file with error: {}", str::from_utf8(&output.stderr).unwrap())));
    }

    Ok(try!(str::from_utf8(&output.stdout)).trim().to_string())
}

pub fn file_set_mode(path: &str, mode: u16) -> Result<()> {
    let mode_s: &str = &mode.to_string();
    let output = process::Command::new(&try!(BinResolver::resolve("chmod"))).args(&vec![mode_s, path]).output().unwrap();

    if !output.status.success() {
        return Err(Error::Generic(format!("Could not chmod file with error: {}", str::from_utf8(&output.stderr).unwrap())));
    }

    Ok(())
}

pub fn hostname() -> Result<String> {
    let output = try!(process::Command::new(&try!(BinResolver::resolve("hostname"))).arg("-f").output());

    if output.status.success() == false {
        return Err(Error::Generic(format!("Could not determine hostname with error: {}", str::from_utf8(&output.stderr).unwrap())));
    }

    Ok(try!(str::from_utf8(&output.stdout)).trim().to_string())
}

pub enum FsFieldOrder {
    Filesystem,
    Size,
    Used,
    Available,
    Capacity,
    Mount,
    Blank,
}

pub fn fs() -> Result<Vec<FsMount>> {
    self::parse_fs(vec![
        self::FsFieldOrder::Filesystem,
        self::FsFieldOrder::Size,
        self::FsFieldOrder::Used,
        self::FsFieldOrder::Available,
        self::FsFieldOrder::Capacity,
        self::FsFieldOrder::Mount,
    ])
}

pub fn parse_fs(fields: Vec<FsFieldOrder>) -> Result<Vec<FsMount>> {
    let mount_out = try!(process::Command::new(&try!(BinResolver::resolve("df"))).arg("-Pk").output());
    let mount = try!(String::from_utf8(mount_out.stdout));

    let mut pattern = "(?m)^".to_string();

    for field in fields {
        match field {
            FsFieldOrder::Filesystem => pattern.push_str("(?P<fs>.+?)"),
            FsFieldOrder::Size => pattern.push_str("(?P<size>[0-9]+)"),
            FsFieldOrder::Used => pattern.push_str("(?P<used>[0-9]+)"),
            FsFieldOrder::Available => pattern.push_str("(?P<available>[0-9]+)"),
            FsFieldOrder::Capacity => pattern.push_str("(?P<capacity>[0-9]{1,3})%"),
            FsFieldOrder::Mount => pattern.push_str("(?P<mount>/.*)"),
            FsFieldOrder::Blank => pattern.push_str(r"[^\s]+"),
        }

        pattern.push_str(r"[\s]*");
    }

    pattern.push_str("$");

    let regex = Regex::new(&pattern).unwrap();
    let mut fs = vec!();

    let lines: Vec<&str> = mount.lines().collect();
    for line in lines {
        if let Some(cap) = regex.captures(line) {
            fs.push(FsMount {
                filesystem: cap.name("fs").unwrap().to_string(),
                mountpoint: cap.name("mount").unwrap().to_string(),
                size: try!(cap.name("size").unwrap().parse::<u64>()),
                used: try!(cap.name("used").unwrap().parse::<u64>()),
                available: try!(cap.name("available").unwrap().parse::<u64>()),
                capacity: try!(cap.name("capacity").unwrap().parse::<f32>())/100.0,
                // inodes_used: try!(cap.name("iused").unwrap().parse::<u64>()),
                // inodes_available: try!(cap.name("iavailable").unwrap().parse::<u64>()),
                // inodes_capacity: try!(cap.name("icapacity").unwrap().parse::<f32>())/100.0,
            });
        }
    };

    Ok(fs)
}

pub fn parse_net(if_pattern: &str, kv_pattern: &str, ipv4_pattern: &str, ipv6_pattern: &str) -> Result<Vec<Netif>> {
    let ifconfig_out = try!(process::Command::new(&try!(BinResolver::resolve("ifconfig"))).output());
    let ifconfig = try!(str::from_utf8(&ifconfig_out.stdout));

    // Rust Regex doesn't support negative lookahead, so we are
    // forced to pass this line by line.
    // Match interfaces:
    // (?m:^([a-z]+[0-9]+):((?:(?!^[a-z]+[0-9]+:).)+))
    // Match options:
    // (?m:^\s*([a-z0-9]+)(?:\s|:\s|=)(.+))
    let if_regex = Regex::new(if_pattern).unwrap();

    let mut net = vec!();

    for cap in if_regex.captures_iter(ifconfig) {
        net.push(try!(parse_netif(cap.name("if").unwrap(), cap.name("content").unwrap(), kv_pattern, ipv4_pattern, ipv6_pattern)));
    }

    Ok(net)
}

fn parse_netif(iface: &str, content: &str, kv_pattern: &str, ipv4_pattern: &str, ipv6_pattern: &str) -> Result<Netif> {
    let mut netif = Netif {
        interface: iface.to_string(),
        mac: None,
        inet: None,
        inet6: None,
        status: None,
    };

    let kv_regex = Regex::new(kv_pattern).unwrap();
    let ipv4_regex = Regex::new(ipv4_pattern).unwrap();
    let ipv6_regex = Regex::new(ipv6_pattern).unwrap();

    let lines: Vec<&str> = content.lines().collect();
    for line in lines {
        if let Some(kv_capture) = kv_regex.captures(line) {
            let value = kv_capture.name("value").unwrap().trim();

            match kv_capture.name("key").unwrap() {
                "ether" | "HWaddr" => netif.mac = Some(value.to_string()),
                "inet" => {
                    if let Some(ipv4_capture) = ipv4_regex.captures(value) {
                        netif.inet = Some(NetifIPv4 {
                            address: ipv4_capture.name("ip").unwrap().to_string(),
                            netmask: ipv4_capture.name("mask").unwrap().to_string(),
                        });
                    }
                },
                "inet6" => {
                    if let Some(ipv6_capture) = ipv6_regex.captures(value) {
                        netif.inet6 = Some(NetifIPv6 {
                            address: ipv6_capture.name("ip").unwrap().to_string(),
                            prefixlen: try!(ipv6_capture.name("prefix").unwrap().parse::<u8>()),
                            scopeid: if ipv6_capture.name("scope").is_some() {
                                Some(ipv6_capture.name("scope").unwrap().to_string())
                            } else {
                                None
                            },
                        });
                    }
                },
                "status" => netif.status = match value {
                    "active" => Some(NetifStatus::Active),
                    "inactive" => Some(NetifStatus::Inactive),
                    _ => unreachable!(),
                },
                _ => (),
            }
        }
    }

    Ok(netif)
}

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

    #[test]
    fn test_hostname() {
        // XXX Not a proper test. Requires mocking.
        assert!(hostname().is_ok());
    }
}
<commit_msg>Fixed bug with file_mv return value<commit_after>// Copyright 2015 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

use {CommandResult, Host, ProviderFactory, Providers, Result};
use error::Error;
use regex::Regex;
use std::{fs, process, str};
use target::bin_resolver::BinResolver;
use telemetry::{FsMount, Netif, NetifIPv4, NetifIPv6, NetifStatus};

pub fn default_provider(host: &mut Host, providers: Vec<Providers>) -> Result<Providers> {
    for p in providers {
        let provider = ProviderFactory::create(host, Some(p));

        if provider.is_ok() {
            return Ok(provider.unwrap().get_providers());
        }
    }

    Err(Error::Generic("No package providers are available".to_string()))
}

pub fn command_exec(cmd: &str) -> Result<CommandResult> {
    let output = try!(process::Command::new("sh").arg("-c").arg(cmd).output());

    Ok(CommandResult {
        exit_code: output.status.code().unwrap(),
        stdout: str::from_utf8(&output.stdout).unwrap().trim().to_string(),
        stderr: str::from_utf8(&output.stderr).unwrap().trim().to_string(),
    })
}

pub fn file_is_file(path: &str) -> Result<bool> {
    let meta = fs::metadata(path);
    Ok(meta.is_err() || meta.unwrap().is_file())
}

pub fn file_exists(path: &str) -> Result<bool> {
    Ok(fs::metadata(path).is_ok())
}

pub fn file_delete(path: &str) -> Result<()> {
    try!(fs::remove_file(path));
    Ok(())
}

pub fn file_mv(path: &str, new_path: &str) -> Result<()> {
    Ok(try!(fs::rename(path, new_path)))
}

pub fn file_copy(path: &str, new_path: &str) -> Result<()> {
    try!(fs::copy(path, new_path));
    Ok(())
}

pub fn file_set_owner(path: &str, user: &str, group: &str) -> Result<()> {
    let user_group = format!("{}:{}", user, group);
    let args: Vec<&str> = vec![&user_group, path];
    let output = process::Command::new(&try!(BinResolver::resolve("chown"))).args(&args).output().unwrap();

    if !output.status.success() {
        return Err(Error::Generic(format!("Could not chown file with error: {}", str::from_utf8(&output.stderr).unwrap())));
    }

    Ok(())
}

pub fn file_stat<'a>(path: &'a str, args: Vec<&'a str>) -> Result<String> {
    let mut args = args;
    args.push(path);
    let output = process::Command::new(&try!(BinResolver::resolve("stat"))).args(&args).output().unwrap();

    if !output.status.success() {
        return Err(Error::Generic(format!("Could not stat file with error: {}", str::from_utf8(&output.stderr).unwrap())));
    }

    Ok(try!(str::from_utf8(&output.stdout)).trim().to_string())
}

pub fn file_set_mode(path: &str, mode: u16) -> Result<()> {
    let mode_s: &str = &mode.to_string();
    let output = process::Command::new(&try!(BinResolver::resolve("chmod"))).args(&vec![mode_s, path]).output().unwrap();

    if !output.status.success() {
        return Err(Error::Generic(format!("Could not chmod file with error: {}", str::from_utf8(&output.stderr).unwrap())));
    }

    Ok(())
}

pub fn hostname() -> Result<String> {
    let output = try!(process::Command::new(&try!(BinResolver::resolve("hostname"))).arg("-f").output());

    if output.status.success() == false {
        return Err(Error::Generic(format!("Could not determine hostname with error: {}", str::from_utf8(&output.stderr).unwrap())));
    }

    Ok(try!(str::from_utf8(&output.stdout)).trim().to_string())
}

pub enum FsFieldOrder {
    Filesystem,
    Size,
    Used,
    Available,
    Capacity,
    Mount,
    Blank,
}

pub fn fs() -> Result<Vec<FsMount>> {
    self::parse_fs(vec![
        self::FsFieldOrder::Filesystem,
        self::FsFieldOrder::Size,
        self::FsFieldOrder::Used,
        self::FsFieldOrder::Available,
        self::FsFieldOrder::Capacity,
        self::FsFieldOrder::Mount,
    ])
}

pub fn parse_fs(fields: Vec<FsFieldOrder>) -> Result<Vec<FsMount>> {
    let mount_out = try!(process::Command::new(&try!(BinResolver::resolve("df"))).arg("-Pk").output());
    let mount = try!(String::from_utf8(mount_out.stdout));

    let mut pattern = "(?m)^".to_string();

    for field in fields {
        match field {
            FsFieldOrder::Filesystem => pattern.push_str("(?P<fs>.+?)"),
            FsFieldOrder::Size => pattern.push_str("(?P<size>[0-9]+)"),
            FsFieldOrder::Used => pattern.push_str("(?P<used>[0-9]+)"),
            FsFieldOrder::Available => pattern.push_str("(?P<available>[0-9]+)"),
            FsFieldOrder::Capacity => pattern.push_str("(?P<capacity>[0-9]{1,3})%"),
            FsFieldOrder::Mount => pattern.push_str("(?P<mount>/.*)"),
            FsFieldOrder::Blank => pattern.push_str(r"[^\s]+"),
        }

        pattern.push_str(r"[\s]*");
    }

    pattern.push_str("$");

    let regex = Regex::new(&pattern).unwrap();
    let mut fs = vec!();

    let lines: Vec<&str> = mount.lines().collect();
    for line in lines {
        if let Some(cap) = regex.captures(line) {
            fs.push(FsMount {
                filesystem: cap.name("fs").unwrap().to_string(),
                mountpoint: cap.name("mount").unwrap().to_string(),
                size: try!(cap.name("size").unwrap().parse::<u64>()),
                used: try!(cap.name("used").unwrap().parse::<u64>()),
                available: try!(cap.name("available").unwrap().parse::<u64>()),
                capacity: try!(cap.name("capacity").unwrap().parse::<f32>())/100.0,
                // inodes_used: try!(cap.name("iused").unwrap().parse::<u64>()),
                // inodes_available: try!(cap.name("iavailable").unwrap().parse::<u64>()),
                // inodes_capacity: try!(cap.name("icapacity").unwrap().parse::<f32>())/100.0,
            });
        }
    };

    Ok(fs)
}

pub fn parse_net(if_pattern: &str, kv_pattern: &str, ipv4_pattern: &str, ipv6_pattern: &str) -> Result<Vec<Netif>> {
    let ifconfig_out = try!(process::Command::new(&try!(BinResolver::resolve("ifconfig"))).output());
    let ifconfig = try!(str::from_utf8(&ifconfig_out.stdout));

    // Rust Regex doesn't support negative lookahead, so we are
    // forced to pass this line by line.
    // Match interfaces:
    // (?m:^([a-z]+[0-9]+):((?:(?!^[a-z]+[0-9]+:).)+))
    // Match options:
    // (?m:^\s*([a-z0-9]+)(?:\s|:\s|=)(.+))
    let if_regex = Regex::new(if_pattern).unwrap();

    let mut net = vec!();

    for cap in if_regex.captures_iter(ifconfig) {
        net.push(try!(parse_netif(cap.name("if").unwrap(), cap.name("content").unwrap(), kv_pattern, ipv4_pattern, ipv6_pattern)));
    }

    Ok(net)
}

fn parse_netif(iface: &str, content: &str, kv_pattern: &str, ipv4_pattern: &str, ipv6_pattern: &str) -> Result<Netif> {
    let mut netif = Netif {
        interface: iface.to_string(),
        mac: None,
        inet: None,
        inet6: None,
        status: None,
    };

    let kv_regex = Regex::new(kv_pattern).unwrap();
    let ipv4_regex = Regex::new(ipv4_pattern).unwrap();
    let ipv6_regex = Regex::new(ipv6_pattern).unwrap();

    let lines: Vec<&str> = content.lines().collect();
    for line in lines {
        if let Some(kv_capture) = kv_regex.captures(line) {
            let value = kv_capture.name("value").unwrap().trim();

            match kv_capture.name("key").unwrap() {
                "ether" | "HWaddr" => netif.mac = Some(value.to_string()),
                "inet" => {
                    if let Some(ipv4_capture) = ipv4_regex.captures(value) {
                        netif.inet = Some(NetifIPv4 {
                            address: ipv4_capture.name("ip").unwrap().to_string(),
                            netmask: ipv4_capture.name("mask").unwrap().to_string(),
                        });
                    }
                },
                "inet6" => {
                    if let Some(ipv6_capture) = ipv6_regex.captures(value) {
                        netif.inet6 = Some(NetifIPv6 {
                            address: ipv6_capture.name("ip").unwrap().to_string(),
                            prefixlen: try!(ipv6_capture.name("prefix").unwrap().parse::<u8>()),
                            scopeid: if ipv6_capture.name("scope").is_some() {
                                Some(ipv6_capture.name("scope").unwrap().to_string())
                            } else {
                                None
                            },
                        });
                    }
                },
                "status" => netif.status = match value {
                    "active" => Some(NetifStatus::Active),
                    "inactive" => Some(NetifStatus::Inactive),
                    _ => unreachable!(),
                },
                _ => (),
            }
        }
    }

    Ok(netif)
}

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

    #[test]
    fn test_hostname() {
        // XXX Not a proper test. Requires mocking.
        assert!(hostname().is_ok());
    }
}
<|endoftext|>