<commit_before>use memory::Memory;
use std::vec::Vec;
use std::sync::atomic::{AtomicIsize, Ordering};

/// Maps multiple `Memory`s into a single address space.
///
/// The `MappedMemory` allows to unify multiple `Memory`s in one address space.
/// A `Memory` is mounted at a certain address and is from now on responsible
/// for all reads and writes between that starting at that address (the “mount
/// point”) and ending where the mounted memory ends.
/// The read and write calls for the responsible `Memory` will be given an
/// address _relative_ to the mount point.
///
/// ```text
///                   Unmapped
///                <-------->
///       mem_a                   mem_b
/// |--------------|            |------|
///                            NullMemory
///                         |--------------
/// |------------- mapped_mem ----------->
/// |              |            |      |
/// 0            0x100        0x200  0x280
/// ```
///
/// For example: `mapped_mem.get(0x208)` would yield the same value as
/// `mem_b.get(0x008)`.
///
/// # Examples
///
/// ```
/// use osciemu::memory::{Memory, SliceMemory, MappedMemory};
///
/// let mut mm = MappedMemory::new();
/// let m1 = mm.mount(0, Box::new(SliceMemory::from_slice(Box::new([1]))));
/// let m2 = mm.mount(2, Box::new(SliceMemory::from_slice(Box::new([2, 2]))));
/// // Now mm =~ [1, _, 2, 2]
/// assert_eq!(mm.get(0), 1);
/// assert_eq!(mm.get(3), 2);
/// ```
///
/// The `MemoryToken`s can be used to access the `Memory` even if it has been
/// mounted. They work like `std::cell::RefCell`.
///
/// ```
/// use osciemu::memory::{Memory, NullMemory, SliceMemory, MappedMemory};
/// use osciemu::memory::mappedmemory::MemoryToken;
///
/// let mut mm = MappedMemory::new();
/// let m1 = mm.mount(0, Box::new(NullMemory::new()));
/// let m2 = mm.mount(0, Box::new(SliceMemory::from_slice(Box::new([1, 2, 3, 4]))));
/// mm.set(0, 99);
/// assert_eq!(mm.get(0), 99);
/// assert_eq!(mm.borrow(&m2).get(0), 99);
/// ```
///
/// The `MemoryToken`s are also used to unmount a mounted `Memory`.
///
/// ```
/// # use osciemu::memory::{Memory, NullMemory, SliceMemory, MappedMemory};
/// # use osciemu::memory::mappedmemory::MemoryToken;
/// # let mut mm = MappedMemory::new();
/// # let m1 = mm.mount(0, Box::new(NullMemory::new()));
/// # let m2 = mm.mount(0, Box::new(SliceMemory::from_slice(Box::new([1, 2, 3, 4]))));
/// # mm.set(0, 99);
/// // ...
/// mm.unmount(&m2);
/// assert_eq!(mm.get(0), 0);
/// assert_eq!(mm.borrow(&m2).get(0), 99);
/// ```
///
/// # Panics
/// `MappedMemory` panics when an unmapped address is read or written.
pub struct MappedMemory {
    mounted_memories: Vec<Entry>,
    unmounted_memories: Vec<Entry>,
}

static id_counter: AtomicIsize = AtomicIsize::new(0);

struct Entry {
    id: isize,
    start_address: usize,
    size: usize,
    memory: Box<Memory>,
}

#[derive(Clone)]
pub struct MemoryToken {
    id: isize,
}

impl MappedMemory {
    pub fn new() -> MappedMemory {
        MappedMemory {
            mounted_memories: Vec::new(),
            unmounted_memories: Vec::new(),
        }
    }

    /// Mounts a `Memory` at the given address.
    pub fn mount(&mut self, start_address: usize, memory: Box<Memory>) -> MemoryToken {
        let size = memory.size();
        let id = id_counter.fetch_add(1, Ordering::Relaxed);
        let new_entry = Entry {
            id,
            start_address,
            size,
            memory,
        };
        self.mounted_memories.push(new_entry);
        MemoryToken{id}
    }

    /// Unmounts the `Memory` references by the `MemoryToken`. After unmounting,
    /// `MappedMemory` does not hold any references to the `Memory`.
    pub fn unmount(&mut self, token: &MemoryToken) {
        let entry = self.mounted_memories
            .iter()
            .enumerate()
            .find(|&(_, entry)| entry.id == token.id)
            .map(|(idx, _)| idx)
            .map(|idx| {
                self.mounted_memories.remove(idx)
            }).unwrap();
        self.unmounted_memories.push(entry);
    }

    pub fn borrow(&self, token: &MemoryToken) -> &Box<Memory> {
        &self.entry_for_token(token).memory
    }

    pub fn borrow_mut(&mut self, token: &MemoryToken) -> &mut Box<Memory> {
        &mut self.entry_for_token_mut(token).memory
    }

    fn entry_for_token(&self, token: &MemoryToken) -> &Entry {
        let m = self.mounted_memories
            .iter()
            .find(|ref entry| entry.id == token.id);

        if m.is_some() {
            m.unwrap()
        } else {
            self.unmounted_memories
                .iter()
                .find(|ref entry| entry.id == token.id)
                .unwrap()
        }
    }

    fn entry_for_token_mut(&mut self, token: &MemoryToken) -> &mut Entry {
        let m = self.mounted_memories
            .iter_mut()
            .find(|ref mut entry| entry.id == token.id);

        if m.is_some() {
            m.unwrap()
        } else {
            self.unmounted_memories
                .iter_mut()
                .find(|ref mut entry| entry.id == token.id)
                .unwrap()
        }
    }

    fn entry_at_addr(&self, addr: usize) -> Option<&Entry> {
        self.mounted_memories
            .iter()
            .rev()
            .find(|entry| entry.start_address <= addr && entry.start_address + entry.size > addr)
    }

    fn entry_at_addr_mut(&mut self, addr: usize) -> Option<&mut Entry> {
        self.mounted_memories
            .iter_mut()
            .rev()
            .find(|entry| entry.start_address <= addr && entry.start_address + entry.size > addr)
    }
}

impl Memory for MappedMemory {
    fn get(&self, addr: usize) -> u32 {
        self.entry_at_addr(addr)
            .map(|entry| entry.memory.get(addr - entry.start_address))
            .expect("Out of bounds")
    }

    fn set(&mut self, addr: usize, value: u32) {
        self.entry_at_addr_mut(addr)
            .map(|entry| entry.memory.set(addr - entry.start_address, value))
            .expect("Out of bounds")
    }

    fn size(&self) -> usize {
        self.mounted_memories
            .iter()
            .map(|entry| entry.start_address + entry.size)
            .max()
            .unwrap_or(0)
    }
}

#[cfg(test)]
mod tests {
    use memory::{Memory, SliceMemory, NullMemory};

    // #[test]
    // fn memory_at_addr() {
    //     let m1 = SliceMemory::from_slice(Box::new([1]));
    //     let m2 = SliceMemory::from_slice(Box::new([2, 2]));
    //     let mut mm = super::MappedMemory::new();
    //     let m1 = mm.mount(0, &m1);
    //     let m2 = mm.mount(2, &m2);

    //     assert!(mm.memory_at_addr(0)
    //         .map_or(false, |entry| entry.memory.get(0) == 1));
    //     assert!(mm.memory_at_addr(1).is_none());
    //     assert!(mm.memory_at_addr(2)
    //         .map_or(false, |entry| entry.memory.get(0) == 2));
    //     assert!(mm.memory_at_addr(3)
    //         .map_or(false, |entry| entry.memory.get(0) == 2));
    //     assert!(mm.memory_at_addr(4).is_none());
    // }

    // #[test]
    // fn overlapping_mounts() {
    //     let m1 = super::MemoryToken::new(NullMemory::new());
    //     let m2 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([2, 2])));
    //     let m3 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([3])));
    //     let mut mm = super::MappedMemory::new();
    //     mm.mount(0, &m1);
    //     mm.mount(1, &m2);
    //     mm.mount(2, &m3);
    //     assert_eq!(mm.get(0), 0);
    //     assert_eq!(mm.get(1), 2);
    //     assert_eq!(mm.get(2), 3);
    //     assert_eq!(mm.get(3), 0);
    // }

    // #[test]
    // fn get_and_set() {
    //     let m1 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([1])));
    //     let m2 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([2, 2])));
    //     let mut mm = super::MappedMemory::new();

    //     mm.mount(0, &m1);
    //     mm.mount(2, &m2);
    //     assert_eq!(mm.get(0), 1);
    //     assert_eq!(mm.get(2), 2);

    //     mm.set(0, 3);
    //     assert_eq!(mm.get(0), 3);

    //     mm.set(3, 0);
    //     assert_eq!(mm.get(3), 0);
    //     assert_eq!(m1.borrow().get(0), 3);
    //     assert_eq!(m2.borrow().get(1), 0);
    // }

    // #[test]
    // fn size() {
    //     let m1 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([1])));
    //     let m2 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([2, 2])));
    //     let mut mm = super::MappedMemory::new();
    //     assert_eq!(mm.size(), 0);

    //     mm.mount(0, &m1);
    //     assert_eq!(mm.size(), 1);

    //     mm.mount(2, &m2);
    //     assert_eq!(mm.size(), 4);
    // }

    // #[test]
    // fn size_with_overlap() {
    //     let m1 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([1, 1, 1, 1, 1])));
    //     let m2 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([2, 2])));
    //     let m3 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([3, 3, 3])));
    //     let mut mm = super::MappedMemory::new();

    //     mm.mount(0, &m1);
    //     mm.mount(2, &m2);
    //     assert_eq!(mm.size(), 5);

    //     mm.mount(6, &m3);
    //     assert_eq!(mm.size(), 9);
    // }

    // #[test]
    // fn unmount() {
    //     let m1 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([1, 1, 1, 1, 1])));
    //     let m2 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([2, 2, 2, 2, 2])));
    //     let m3 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([3, 3, 3, 3, 3])));
    //     let mut mm = super::MappedMemory::new();
    //     mm.mount(0, &m1);
    //     mm.mount(0, &m2);
    //     mm.mount(0, &m3);

    //     for i in 0..5 {
    //         assert_eq!(mm.get(i), 3);
    //     }
    //     mm.unmount(&m3);
    //     for i in 0..5 {
    //         assert_eq!(mm.get(i), 2);
    //     }
    //     mm.unmount(&m3);
    //     mm.unmount(&m3);
    //     for i in 0..5 {
    //         assert_eq!(mm.get(i), 2);
    //     }
    //     mm.unmount(&m1);
    //     for i in 0..5 {
    //         assert_eq!(mm.get(i), 2);
    //     }
    // }
}
<commit_msg>feat(mappedmemory): Add Entry.contains()<commit_after>use memory::Memory;
use std::vec::Vec;
use std::sync::atomic::{AtomicIsize, Ordering};

/// Maps multiple `Memory`s into a single address space.
///
/// The `MappedMemory` allows to unify multiple `Memory`s in one address space.
/// A `Memory` is mounted at a certain address and is from now on responsible
/// for all reads and writes between that starting at that address (the “mount
/// point”) and ending where the mounted memory ends.
/// The read and write calls for the responsible `Memory` will be given an
/// address _relative_ to the mount point.
///
/// ```text
///                   Unmapped
///                <-------->
///       mem_a                   mem_b
/// |--------------|            |------|
///                            NullMemory
///                         |--------------
/// |------------- mapped_mem ----------->
/// |              |            |      |
/// 0            0x100        0x200  0x280
/// ```
///
/// For example: `mapped_mem.get(0x208)` would yield the same value as
/// `mem_b.get(0x008)`.
///
/// # Examples
///
/// ```
/// use osciemu::memory::{Memory, SliceMemory, MappedMemory};
///
/// let mut mm = MappedMemory::new();
/// let m1 = mm.mount(0, Box::new(SliceMemory::from_slice(Box::new([1]))));
/// let m2 = mm.mount(2, Box::new(SliceMemory::from_slice(Box::new([2, 2]))));
/// // Now mm =~ [1, _, 2, 2]
/// assert_eq!(mm.get(0), 1);
/// assert_eq!(mm.get(3), 2);
/// ```
///
/// The `MemoryToken`s can be used to access the `Memory` even if it has been
/// mounted. They work like `std::cell::RefCell`.
///
/// ```
/// use osciemu::memory::{Memory, NullMemory, SliceMemory, MappedMemory};
/// use osciemu::memory::mappedmemory::MemoryToken;
///
/// let mut mm = MappedMemory::new();
/// let m1 = mm.mount(0, Box::new(NullMemory::new()));
/// let m2 = mm.mount(0, Box::new(SliceMemory::from_slice(Box::new([1, 2, 3, 4]))));
/// mm.set(0, 99);
/// assert_eq!(mm.get(0), 99);
/// assert_eq!(mm.borrow(&m2).get(0), 99);
/// ```
///
/// The `MemoryToken`s are also used to unmount a mounted `Memory`.
///
/// ```
/// # use osciemu::memory::{Memory, NullMemory, SliceMemory, MappedMemory};
/// # use osciemu::memory::mappedmemory::MemoryToken;
/// # let mut mm = MappedMemory::new();
/// # let m1 = mm.mount(0, Box::new(NullMemory::new()));
/// # let m2 = mm.mount(0, Box::new(SliceMemory::from_slice(Box::new([1, 2, 3, 4]))));
/// # mm.set(0, 99);
/// // ...
/// mm.unmount(&m2);
/// assert_eq!(mm.get(0), 0);
/// assert_eq!(mm.borrow(&m2).get(0), 99);
/// ```
///
/// # Panics
/// `MappedMemory` panics when an unmapped address is read or written.
pub struct MappedMemory {
    mounted_memories: Vec<Entry>,
    unmounted_memories: Vec<Entry>,
}

static id_counter: AtomicIsize = AtomicIsize::new(0);

struct Entry {
    id: isize,
    start_address: usize,
    size: usize,
    memory: Box<Memory>,
}

impl Entry {
    fn contains(&self, addr: usize) -> bool {
        self.start_address <= addr && self.start_address + self.size > addr
    }
}

#[derive(Clone)]
pub struct MemoryToken {
    id: isize,
}

impl MappedMemory {
    pub fn new() -> MappedMemory {
        MappedMemory {
            mounted_memories: Vec::new(),
            unmounted_memories: Vec::new(),
        }
    }

    /// Mounts a `Memory` at the given address.
    pub fn mount(&mut self, start_address: usize, memory: Box<Memory>) -> MemoryToken {
        let size = memory.size();
        let id = id_counter.fetch_add(1, Ordering::Relaxed);
        let new_entry = Entry {
            id,
            start_address,
            size,
            memory,
        };
        self.mounted_memories.push(new_entry);
        MemoryToken{id}
    }

    /// Unmounts the `Memory` references by the `MemoryToken`. After unmounting,
    /// `MappedMemory` does not hold any references to the `Memory`.
    pub fn unmount(&mut self, token: &MemoryToken) {
        let entry = self.mounted_memories
            .iter()
            .enumerate()
            .find(|&(_, entry)| entry.id == token.id)
            .map(|(idx, _)| idx)
            .map(|idx| {
                self.mounted_memories.remove(idx)
            }).unwrap();
        self.unmounted_memories.push(entry);
    }

    pub fn borrow(&self, token: &MemoryToken) -> &Box<Memory> {
        &self.entry_for_token(token).memory
    }

    pub fn borrow_mut(&mut self, token: &MemoryToken) -> &mut Box<Memory> {
        &mut self.entry_for_token_mut(token).memory
    }

    fn entry_for_token(&self, token: &MemoryToken) -> &Entry {
        let m = self.mounted_memories
            .iter()
            .find(|ref entry| entry.id == token.id);

        if m.is_some() {
            m.unwrap()
        } else {
            self.unmounted_memories
                .iter()
                .find(|ref entry| entry.id == token.id)
                .unwrap()
        }
    }

    fn entry_for_token_mut(&mut self, token: &MemoryToken) -> &mut Entry {
        let m = self.mounted_memories
            .iter_mut()
            .find(|ref mut entry| entry.id == token.id);

        if m.is_some() {
            m.unwrap()
        } else {
            self.unmounted_memories
                .iter_mut()
                .find(|ref mut entry| entry.id == token.id)
                .unwrap()
        }
    }

    fn entry_at_addr(&self, addr: usize) -> Option<&Entry> {
        self.mounted_memories
            .iter()
            .rev()
            .find(|entry| entry.contains(addr))
    }

    fn entry_at_addr_mut(&mut self, addr: usize) -> Option<&mut Entry> {
        self.mounted_memories
            .iter_mut()
            .rev()
            .find(|entry| entry.contains(addr))
    }
}

impl Memory for MappedMemory {
    fn get(&self, addr: usize) -> u32 {
        self.entry_at_addr(addr)
            .map(|entry| entry.memory.get(addr - entry.start_address))
            .expect("Out of bounds")
    }

    fn set(&mut self, addr: usize, value: u32) {
        self.entry_at_addr_mut(addr)
            .map(|entry| entry.memory.set(addr - entry.start_address, value))
            .expect("Out of bounds")
    }

    fn size(&self) -> usize {
        self.mounted_memories
            .iter()
            .map(|entry| entry.start_address + entry.size)
            .max()
            .unwrap_or(0)
    }
}

#[cfg(test)]
mod tests {
    use memory::{Memory, SliceMemory, NullMemory};

    // #[test]
    // fn overlapping_mounts() {
    //     let m1 = super::MemoryToken::new(NullMemory::new());
    //     let m2 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([2, 2])));
    //     let m3 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([3])));
    //     let mut mm = super::MappedMemory::new();
    //     mm.mount(0, &m1);
    //     mm.mount(1, &m2);
    //     mm.mount(2, &m3);
    //     assert_eq!(mm.get(0), 0);
    //     assert_eq!(mm.get(1), 2);
    //     assert_eq!(mm.get(2), 3);
    //     assert_eq!(mm.get(3), 0);
    // }

    // #[test]
    // fn get_and_set() {
    //     let m1 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([1])));
    //     let m2 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([2, 2])));
    //     let mut mm = super::MappedMemory::new();

    //     mm.mount(0, &m1);
    //     mm.mount(2, &m2);
    //     assert_eq!(mm.get(0), 1);
    //     assert_eq!(mm.get(2), 2);

    //     mm.set(0, 3);
    //     assert_eq!(mm.get(0), 3);

    //     mm.set(3, 0);
    //     assert_eq!(mm.get(3), 0);
    //     assert_eq!(m1.borrow().get(0), 3);
    //     assert_eq!(m2.borrow().get(1), 0);
    // }

    // #[test]
    // fn size() {
    //     let m1 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([1])));
    //     let m2 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([2, 2])));
    //     let mut mm = super::MappedMemory::new();
    //     assert_eq!(mm.size(), 0);

    //     mm.mount(0, &m1);
    //     assert_eq!(mm.size(), 1);

    //     mm.mount(2, &m2);
    //     assert_eq!(mm.size(), 4);
    // }

    // #[test]
    // fn size_with_overlap() {
    //     let m1 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([1, 1, 1, 1, 1])));
    //     let m2 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([2, 2])));
    //     let m3 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([3, 3, 3])));
    //     let mut mm = super::MappedMemory::new();

    //     mm.mount(0, &m1);
    //     mm.mount(2, &m2);
    //     assert_eq!(mm.size(), 5);

    //     mm.mount(6, &m3);
    //     assert_eq!(mm.size(), 9);
    // }

    // #[test]
    // fn unmount() {
    //     let m1 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([1, 1, 1, 1, 1])));
    //     let m2 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([2, 2, 2, 2, 2])));
    //     let m3 = super::MemoryToken::new(SliceMemory::from_slice(Box::new([3, 3, 3, 3, 3])));
    //     let mut mm = super::MappedMemory::new();
    //     mm.mount(0, &m1);
    //     mm.mount(0, &m2);
    //     mm.mount(0, &m3);

    //     for i in 0..5 {
    //         assert_eq!(mm.get(i), 3);
    //     }
    //     mm.unmount(&m3);
    //     for i in 0..5 {
    //         assert_eq!(mm.get(i), 2);
    //     }
    //     mm.unmount(&m3);
    //     mm.unmount(&m3);
    //     for i in 0..5 {
    //         assert_eq!(mm.get(i), 2);
    //     }
    //     mm.unmount(&m1);
    //     for i in 0..5 {
    //         assert_eq!(mm.get(i), 2);
    //     }
    // }
}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate gimli;
extern crate test;

use gimli::{DebugAbbrev, DebugInfo, LittleEndian};

use std::env;
use std::fs::File;
use std::io::Read;
use std::path::PathBuf;

fn read_section(section: &str) -> Vec<u8> {
    let mut path = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
    path.push("./fixtures/self/");
    path.push(section);

    assert!(path.is_file());
    let mut file = File::open(path).unwrap();

    let mut buf = Vec::new();
    file.read_to_end(&mut buf).unwrap();
    buf
}

#[bench]
fn bench_parsing_debug_abbrev(b: &mut test::Bencher) {
    let debug_info = read_section("debug_info");
    let debug_info = DebugInfo::<LittleEndian>::new(&debug_info);
    let unit = debug_info.units()
        .next()
        .expect("Should have at least one compilation unit")
        .expect("And it should parse OK");

    let debug_abbrev = read_section("debug_abbrev");

    b.iter(|| {
        let debug_abbrev = DebugAbbrev::<LittleEndian>::new(&debug_abbrev);
        test::black_box(unit.abbreviations(debug_abbrev)
            .expect("Should parse abbreviations"));
    });
}

#[bench]
fn bench_parsing_debug_info(b: &mut test::Bencher) {
    let debug_abbrev = read_section("debug_abbrev");
    let debug_abbrev = DebugAbbrev::<LittleEndian>::new(&debug_abbrev);

    let debug_info = read_section("debug_info");

    b.iter(|| {
        let debug_info = DebugInfo::<LittleEndian>::new(&debug_info);

        for unit in debug_info.units() {
            let unit = unit.expect("Should parse compilation unit");
            let abbrevs = unit.abbreviations(debug_abbrev)
                .expect("Should parse abbreviations");

            let mut cursor = unit.entries(&abbrevs);

            loop {
                let entry = cursor.current()
                    .expect("Should have a current entry")
                    .expect("And should parse that entry OK");

                for attr in entry.attrs() {
                    test::black_box(attr.expect("Should parse entry's attribute"));
                }

                if let None = cursor.next_dfs() {
                    break;
                }
            }
        }
    });
}
<commit_msg>Allow running benches manually<commit_after>#![feature(test)]

extern crate gimli;
extern crate test;

use gimli::{DebugAbbrev, DebugInfo, LittleEndian};

use std::env;
use std::fs::File;
use std::io::Read;
use std::path::PathBuf;

fn read_section(section: &str) -> Vec<u8> {
    let mut path = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap_or(".".into()));
    path.push("./fixtures/self/");
    path.push(section);

    assert!(path.is_file());
    let mut file = File::open(path).unwrap();

    let mut buf = Vec::new();
    file.read_to_end(&mut buf).unwrap();
    buf
}

#[bench]
fn bench_parsing_debug_abbrev(b: &mut test::Bencher) {
    let debug_info = read_section("debug_info");
    let debug_info = DebugInfo::<LittleEndian>::new(&debug_info);
    let unit = debug_info.units()
        .next()
        .expect("Should have at least one compilation unit")
        .expect("And it should parse OK");

    let debug_abbrev = read_section("debug_abbrev");

    b.iter(|| {
        let debug_abbrev = DebugAbbrev::<LittleEndian>::new(&debug_abbrev);
        test::black_box(unit.abbreviations(debug_abbrev)
            .expect("Should parse abbreviations"));
    });
}

#[bench]
fn bench_parsing_debug_info(b: &mut test::Bencher) {
    let debug_abbrev = read_section("debug_abbrev");
    let debug_abbrev = DebugAbbrev::<LittleEndian>::new(&debug_abbrev);

    let debug_info = read_section("debug_info");

    b.iter(|| {
        let debug_info = DebugInfo::<LittleEndian>::new(&debug_info);

        for unit in debug_info.units() {
            let unit = unit.expect("Should parse compilation unit");
            let abbrevs = unit.abbreviations(debug_abbrev)
                .expect("Should parse abbreviations");

            let mut cursor = unit.entries(&abbrevs);

            loop {
                let entry = cursor.current()
                    .expect("Should have a current entry")
                    .expect("And should parse that entry OK");

                for attr in entry.attrs() {
                    test::black_box(attr.expect("Should parse entry's attribute"));
                }

                if let None = cursor.next_dfs() {
                    break;
                }
            }
        }
    });
}
<|endoftext|>


<commit_before>use std::ops::Mul;
use std::sync::Arc;
use std::vec::IntoIter;
use na::{Translate, Cross, Translation, Rotation};
use math::{Scalar, Point, Vect, Isometry};
use utils::data::uid_remap::{UidRemap, FastKey};
use entities::inspection::Repr;
use entities::bounding_volume::{self, AABB};
use queries::geometry::Contact;
use queries::ray::{RayCast, Ray, RayIntersection};
use queries::point::PointQuery;
use narrow_phase::{BasicCollisionDispatcher, ContactSignalHandler, CollisionAlgorithm};
use broad_phase::{BroadPhase, DBVTBroadPhase};
use world::{CollisionObjectsDispatcher, CollisionObject, CollisionGroups};

use na::{Pnt2, Pnt3, Iso2, Iso3};

// FIXME: be generic wrt the BV?
/// Type of the broad phase trait-object used by the collision world.
pub type BroadPhaseObject<P> = Box<BroadPhase<P, AABB<P>, FastKey> + 'static>;

/// A world that handles collision objects.
pub struct CollisionWorld<P, M, T> {
    objects:       UidRemap<CollisionObject<P, M, T>>,
    broad_phase:   BroadPhaseObject<P>,
    narrow_phase:  CollisionObjectsDispatcher<P, M, T>,
    pos_to_update: Vec<(FastKey, M)>,
    timestamp:     usize
    // FIXME: allow modification of the other properties too.
}

impl<P, M, T> CollisionWorld<P, M, T>
    where P: Point,
          P::Vect: Translate<P> + Cross,
          <P::Vect as Cross>::CrossProductType: Vect<Scalar = <P::Vect as Vect>::Scalar> +
                                                Mul<<P::Vect as Vect>::Scalar, Output = <P::Vect as Cross>::CrossProductType>, // FIXME: why do we need this?
          M:  Isometry<P, P::Vect> + Translation<P::Vect> + Rotation<<P::Vect as Cross>::CrossProductType> {
    /// Creates a new collision world.
    // FIXME: use default values for `margin` and `prediction` and allow their modification by the
    // user ?
    pub fn new(margin:     <P::Vect as Vect>::Scalar,
               prediction: <P::Vect as Vect>::Scalar,
               small_uids: bool)
               -> CollisionWorld<P, M, T> {
        let objects = UidRemap::new(small_uids);
        let sdispatcher = Box::new(BasicCollisionDispatcher::new(prediction));
        let broad_phase = Box::new(DBVTBroadPhase::<P, AABB<P>, FastKey>::new(margin, true));
        let narrow_phase = CollisionObjectsDispatcher::new(sdispatcher);

        CollisionWorld {
            objects:       objects,
            broad_phase:   broad_phase,
            narrow_phase:  narrow_phase,
            pos_to_update: Vec::new(),
            timestamp:     0
        }
    }

    /// Adds a collision object to the world.
    pub fn add(&mut self,
               uid: usize,
               position: M,
               shape: Arc<Box<Repr<P, M>>>,
               collision_groups: CollisionGroups,
               data: T) {
        // FIXME: test that we did not add this object already ?

        let mut collision_object = CollisionObject::new(position, shape, collision_groups, data);
        collision_object.timestamp = self.timestamp;
        let aabb = bounding_volume::aabb(&**collision_object.shape, &collision_object.position);
        let fk = self.objects.insert(uid, collision_object).0;
        self.broad_phase.defered_add(fk.uid(), aabb, fk)
    }

    /// Remove a collision object from the world.
    pub fn remove(&mut self, uid: usize) {
        if let Some((fk, _)) = self.objects.remove(uid) {
            self.broad_phase.defered_remove(fk.uid());
        }
    }

    /// Updates the collision world.
    ///
    /// This executes the whole collision detection pipeline: the broad phase first, then the
    /// narrow phase.
    pub fn update(&mut self) {
        self.perform_position_update();
        self.perform_broad_phase();
        self.perform_narrow_phase();
    }

    /// Sets the position the collision object attached to the specified object will have during
    /// the next update.
    pub fn defered_set_position(&mut self, uid: usize, pos: M) {
        if let Some(fk) = self.objects.get_fast_key(uid) {
            self.pos_to_update.push((fk, pos))
        }
    }

    /// Registers a handler for contact start/stop events.
    pub fn register_contact_signal_handler<H>(&mut self, name: &str, handler: H)
        where H: ContactSignalHandler<T> + 'static {
        self.narrow_phase.register_contact_signal_handler(name, Box::new(handler))
    }

    /// Unregisters a handler for contact start/stop events.
    pub fn unregister_contact_signal_handler(&mut self, name: &str) {
        self.narrow_phase.unregister_contact_signal_handler(name)
    }

    /// Executes the position updates.
    pub fn perform_position_update(&mut self) {
        for &(ref fk, ref pos) in self.pos_to_update.iter() {
            if let Some(co) = self.objects.get_fast_mut(fk) {
                co.position = pos.clone();
                co.timestamp = self.timestamp;
                self.broad_phase.defered_set_bounding_volume(fk.uid(), bounding_volume::aabb(&**co.shape, pos));
            }
        }

        self.pos_to_update.clear();
    }

    /// Executes the broad phase of the collision detection pipeline.
    ///
    /// Not that this does not take in account the changes made to the collision updates after the
    /// last `.perform_position_update()` call.
    pub fn perform_broad_phase(&mut self) {
        let bf = &mut self.broad_phase;
        let nf = &mut self.narrow_phase;
        let objs = &self.objects;

        bf.update(
            &mut |b1, b2| CollisionObjectsDispatcher::is_proximity_allowed(objs, b1, b2),
            &mut |b1, b2, started| nf.handle_proximity(objs, b1, b2, started)
        );
    }

    /// Executes the narrow phase of the collision detection pipeline.
    pub fn perform_narrow_phase(&mut self) {
        self.narrow_phase.update(&self.objects, self.timestamp);
        self.timestamp = self.timestamp + 1;
    }

    /// Iterates through all the contact pairs.
    #[inline(always)]
    pub fn contact_pairs<F>(&self, f: F)
          where F: FnMut(&T, &T, &CollisionAlgorithm<P, M>) {
        self.narrow_phase.contact_pairs(&self.objects, f)
    }

    /// Collects every contact detected since the last update.
    #[inline(always)]
    pub fn contacts<F>(&self, f: F)
          where F: FnMut(&T, &T, &Contact<P>) {
        self.narrow_phase.contacts(&self.objects, f)
    }

    /// Computes the interferences between every rigid bodies on this world and a ray.
    #[inline(always)]
    pub fn interferences_with_ray<'a>(&'a self, ray: &'a Ray<P>, groups: &'a CollisionGroups)
        -> InterferencesWithRay<'a, P, M, T> {
        let mut fks = Vec::new();

        self.broad_phase.interferences_with_ray(ray, &mut fks);

        InterferencesWithRay {
            ray:     ray,
            groups:  groups,
            objects: &self.objects,
            idx:     fks.into_iter()
        }
    }

    /// Computes the interferences between every rigid bodies of a given broad phase, and a point.
    #[inline(always)]
    pub fn interferences_with_point<'a>(&'a self, point: &'a P, groups: &'a CollisionGroups)
        -> InterferencesWithPoint<'a, P, M, T> {
        let mut fks = Vec::new();

        self.broad_phase.interferences_with_point(point, &mut fks);

        InterferencesWithPoint {
            point:   point,
            groups:  groups,
            objects: &self.objects,
            idx:     fks.into_iter()
        }
    }

    /// Computes the interferences between every rigid bodies of a given broad phase, and a aabb.
    #[inline(always)]
    pub fn interferences_with_aabb<'a>(&'a self, aabb: &'a AABB<P>, groups: &'a CollisionGroups)
        -> InterferencesWithAABB<'a, P, M, T> {
        let mut fks = Vec::new();

        self.broad_phase.interferences_with_bounding_volume(aabb, &mut fks);

        InterferencesWithAABB {
            groups:  groups,
            objects: &self.objects,
            idx:     fks.into_iter()
        }
    }
}

/// Iterator through all the objects on the world that intersect a specific ray.
pub struct InterferencesWithRay<'a, P: 'a + Point, M: 'a, T: 'a> {
    ray:     &'a Ray<P>,
    objects: &'a UidRemap<CollisionObject<P, M, T>>,
    groups:  &'a CollisionGroups,
    idx:     IntoIter<&'a FastKey>,
}

impl<'a, P, M, T> Iterator for InterferencesWithRay<'a, P, M, T>
    where P: Point,
          M: Isometry<P, P::Vect> + Translation<P::Vect> {
    type Item = (&'a CollisionObject<P, M, T>, RayIntersection<P::Vect>);

    #[inline]
    fn next(&mut self) -> Option<(&'a CollisionObject<P, M, T>, RayIntersection<P::Vect>)> {
        match self.idx.next() {
            Some(id) => {
                let co = &self.objects[*id];

                if co.collision_groups.can_collide_with_groups(self.groups) {
                    let inter = co.shape.toi_and_normal_with_ray(&co.position, self.ray, true);

                    if let Some(inter) = inter {
                        return Some((co, inter))
                    }
                }
                None
            },
            None => None
        }
    }
}

/// Iterator through all the objects on the world that intersect a specific point.
pub struct InterferencesWithPoint<'a, P: 'a, M: 'a, T: 'a> {
    point:   &'a P,
    objects: &'a UidRemap<CollisionObject<P, M, T>>,
    groups:  &'a CollisionGroups,
    idx:     IntoIter<&'a FastKey>,
}

impl<'a, P, M, T> Iterator for InterferencesWithPoint<'a, P, M, T>
    where P: Point,
          M: Isometry<P, P::Vect> + Translation<P::Vect> {
    type Item = &'a CollisionObject<P, M, T>;

    #[inline]
    fn next(&mut self) -> Option<&'a CollisionObject<P, M, T>> {
        match self.idx.next() {
            Some(id) => {
                let co = &self.objects[*id];

                if co.collision_groups.can_collide_with_groups(self.groups) &&
                   co.shape.contains_point(&co.position, self.point) {
                    Some(co)
                }
                else {
                    None
                }
            },
            None => None
        }
    }
}

/// Iterator through all the objects on the world which bounding volume intersects a specific AABB.
pub struct InterferencesWithAABB<'a, P: 'a, M: 'a, T: 'a> {
    objects: &'a UidRemap<CollisionObject<P, M, T>>,
    groups:  &'a CollisionGroups,
    idx:     IntoIter<&'a FastKey>,
}

impl<'a, P, M, T> Iterator for InterferencesWithAABB<'a, P, M, T> {
    type Item = &'a CollisionObject<P, M, T>;

    #[inline]
    fn next(&mut self) -> Option<&'a CollisionObject<P, M, T>> {
        match self.idx.next() {
            Some(id) => {
                let co = &self.objects[*id];

                if co.collision_groups.can_collide_with_groups(self.groups) {
                    Some(co)
                }
                else {
                    None
                }
            },
            None => None
        }
    }
}

/// 2D collision world containing objects of type `T`.
pub type CollisionWorld2<N, T> = CollisionWorld<Pnt2<N>, Iso2<N>, T>;
/// 3D collision world containing objects of type `T`.
pub type CollisionWorld3<N, T> = CollisionWorld<Pnt3<N>, Iso3<N>, T>;
<commit_msg>Fix the new iterators.<commit_after>use std::ops::Mul;
use std::sync::Arc;
use std::vec::IntoIter;
use na::{Translate, Cross, Translation, Rotation};
use math::{Scalar, Point, Vect, Isometry};
use utils::data::uid_remap::{UidRemap, FastKey};
use entities::inspection::Repr;
use entities::bounding_volume::{self, AABB};
use queries::geometry::Contact;
use queries::ray::{RayCast, Ray, RayIntersection};
use queries::point::PointQuery;
use narrow_phase::{BasicCollisionDispatcher, ContactSignalHandler, CollisionAlgorithm};
use broad_phase::{BroadPhase, DBVTBroadPhase};
use world::{CollisionObjectsDispatcher, CollisionObject, CollisionGroups};

use na::{Pnt2, Pnt3, Iso2, Iso3};

// FIXME: be generic wrt the BV?
/// Type of the broad phase trait-object used by the collision world.
pub type BroadPhaseObject<P> = Box<BroadPhase<P, AABB<P>, FastKey> + 'static>;

/// A world that handles collision objects.
pub struct CollisionWorld<P, M, T> {
    objects:       UidRemap<CollisionObject<P, M, T>>,
    broad_phase:   BroadPhaseObject<P>,
    narrow_phase:  CollisionObjectsDispatcher<P, M, T>,
    pos_to_update: Vec<(FastKey, M)>,
    timestamp:     usize
    // FIXME: allow modification of the other properties too.
}

impl<P, M, T> CollisionWorld<P, M, T>
    where P: Point,
          P::Vect: Translate<P> + Cross,
          <P::Vect as Cross>::CrossProductType: Vect<Scalar = <P::Vect as Vect>::Scalar> +
                                                Mul<<P::Vect as Vect>::Scalar, Output = <P::Vect as Cross>::CrossProductType>, // FIXME: why do we need this?
          M:  Isometry<P, P::Vect> + Translation<P::Vect> + Rotation<<P::Vect as Cross>::CrossProductType> {
    /// Creates a new collision world.
    // FIXME: use default values for `margin` and `prediction` and allow their modification by the
    // user ?
    pub fn new(margin:     <P::Vect as Vect>::Scalar,
               prediction: <P::Vect as Vect>::Scalar,
               small_uids: bool)
               -> CollisionWorld<P, M, T> {
        let objects = UidRemap::new(small_uids);
        let sdispatcher = Box::new(BasicCollisionDispatcher::new(prediction));
        let broad_phase = Box::new(DBVTBroadPhase::<P, AABB<P>, FastKey>::new(margin, true));
        let narrow_phase = CollisionObjectsDispatcher::new(sdispatcher);

        CollisionWorld {
            objects:       objects,
            broad_phase:   broad_phase,
            narrow_phase:  narrow_phase,
            pos_to_update: Vec::new(),
            timestamp:     0
        }
    }

    /// Adds a collision object to the world.
    pub fn add(&mut self,
               uid: usize,
               position: M,
               shape: Arc<Box<Repr<P, M>>>,
               collision_groups: CollisionGroups,
               data: T) {
        // FIXME: test that we did not add this object already ?

        let mut collision_object = CollisionObject::new(position, shape, collision_groups, data);
        collision_object.timestamp = self.timestamp;
        let aabb = bounding_volume::aabb(&**collision_object.shape, &collision_object.position);
        let fk = self.objects.insert(uid, collision_object).0;
        self.broad_phase.defered_add(fk.uid(), aabb, fk)
    }

    /// Remove a collision object from the world.
    pub fn remove(&mut self, uid: usize) {
        if let Some((fk, _)) = self.objects.remove(uid) {
            self.broad_phase.defered_remove(fk.uid());
        }
    }

    /// Updates the collision world.
    ///
    /// This executes the whole collision detection pipeline: the broad phase first, then the
    /// narrow phase.
    pub fn update(&mut self) {
        self.perform_position_update();
        self.perform_broad_phase();
        self.perform_narrow_phase();
    }

    /// Sets the position the collision object attached to the specified object will have during
    /// the next update.
    pub fn defered_set_position(&mut self, uid: usize, pos: M) {
        if let Some(fk) = self.objects.get_fast_key(uid) {
            self.pos_to_update.push((fk, pos))
        }
    }

    /// Registers a handler for contact start/stop events.
    pub fn register_contact_signal_handler<H>(&mut self, name: &str, handler: H)
        where H: ContactSignalHandler<T> + 'static {
        self.narrow_phase.register_contact_signal_handler(name, Box::new(handler))
    }

    /// Unregisters a handler for contact start/stop events.
    pub fn unregister_contact_signal_handler(&mut self, name: &str) {
        self.narrow_phase.unregister_contact_signal_handler(name)
    }

    /// Executes the position updates.
    pub fn perform_position_update(&mut self) {
        for &(ref fk, ref pos) in self.pos_to_update.iter() {
            if let Some(co) = self.objects.get_fast_mut(fk) {
                co.position = pos.clone();
                co.timestamp = self.timestamp;
                self.broad_phase.defered_set_bounding_volume(fk.uid(), bounding_volume::aabb(&**co.shape, pos));
            }
        }

        self.pos_to_update.clear();
    }

    /// Executes the broad phase of the collision detection pipeline.
    ///
    /// Not that this does not take in account the changes made to the collision updates after the
    /// last `.perform_position_update()` call.
    pub fn perform_broad_phase(&mut self) {
        let bf = &mut self.broad_phase;
        let nf = &mut self.narrow_phase;
        let objs = &self.objects;

        bf.update(
            &mut |b1, b2| CollisionObjectsDispatcher::is_proximity_allowed(objs, b1, b2),
            &mut |b1, b2, started| nf.handle_proximity(objs, b1, b2, started)
        );
    }

    /// Executes the narrow phase of the collision detection pipeline.
    pub fn perform_narrow_phase(&mut self) {
        self.narrow_phase.update(&self.objects, self.timestamp);
        self.timestamp = self.timestamp + 1;
    }

    /// Iterates through all the contact pairs.
    #[inline(always)]
    pub fn contact_pairs<F>(&self, f: F)
          where F: FnMut(&T, &T, &CollisionAlgorithm<P, M>) {
        self.narrow_phase.contact_pairs(&self.objects, f)
    }

    /// Collects every contact detected since the last update.
    #[inline(always)]
    pub fn contacts<F>(&self, f: F)
          where F: FnMut(&T, &T, &Contact<P>) {
        self.narrow_phase.contacts(&self.objects, f)
    }

    /// Computes the interferences between every rigid bodies on this world and a ray.
    #[inline(always)]
    pub fn interferences_with_ray<'a>(&'a self, ray: &'a Ray<P>, groups: &'a CollisionGroups)
        -> InterferencesWithRay<'a, P, M, T> {
        let mut fks = Vec::new();

        self.broad_phase.interferences_with_ray(ray, &mut fks);

        InterferencesWithRay {
            ray:     ray,
            groups:  groups,
            objects: &self.objects,
            idx:     fks.into_iter()
        }
    }

    /// Computes the interferences between every rigid bodies of a given broad phase, and a point.
    #[inline(always)]
    pub fn interferences_with_point<'a>(&'a self, point: &'a P, groups: &'a CollisionGroups)
        -> InterferencesWithPoint<'a, P, M, T> {
        let mut fks = Vec::new();

        self.broad_phase.interferences_with_point(point, &mut fks);

        InterferencesWithPoint {
            point:   point,
            groups:  groups,
            objects: &self.objects,
            idx:     fks.into_iter()
        }
    }

    /// Computes the interferences between every rigid bodies of a given broad phase, and a aabb.
    #[inline(always)]
    pub fn interferences_with_aabb<'a>(&'a self, aabb: &'a AABB<P>, groups: &'a CollisionGroups)
        -> InterferencesWithAABB<'a, P, M, T> {
        let mut fks = Vec::new();

        self.broad_phase.interferences_with_bounding_volume(aabb, &mut fks);

        InterferencesWithAABB {
            groups:  groups,
            objects: &self.objects,
            idx:     fks.into_iter()
        }
    }
}

/// Iterator through all the objects on the world that intersect a specific ray.
pub struct InterferencesWithRay<'a, P: 'a + Point, M: 'a, T: 'a> {
    ray:     &'a Ray<P>,
    objects: &'a UidRemap<CollisionObject<P, M, T>>,
    groups:  &'a CollisionGroups,
    idx:     IntoIter<&'a FastKey>,
}

impl<'a, P, M, T> Iterator for InterferencesWithRay<'a, P, M, T>
    where P: Point,
          M: Isometry<P, P::Vect> + Translation<P::Vect> {
    type Item = (&'a CollisionObject<P, M, T>, RayIntersection<P::Vect>);

    #[inline]
    fn next(&mut self) -> Option<(&'a CollisionObject<P, M, T>, RayIntersection<P::Vect>)> {
        while let Some(id) = self.idx.next() {
            let co = &self.objects[*id];

            if co.collision_groups.can_collide_with_groups(self.groups) {
                let inter = co.shape.toi_and_normal_with_ray(&co.position, self.ray, true);

                if let Some(inter) = inter {
                    return Some((co, inter))
                }
            }
        }

        None
    }
}

/// Iterator through all the objects on the world that intersect a specific point.
pub struct InterferencesWithPoint<'a, P: 'a, M: 'a, T: 'a> {
    point:   &'a P,
    objects: &'a UidRemap<CollisionObject<P, M, T>>,
    groups:  &'a CollisionGroups,
    idx:     IntoIter<&'a FastKey>,
}

impl<'a, P, M, T> Iterator for InterferencesWithPoint<'a, P, M, T>
    where P: Point,
          M: Isometry<P, P::Vect> + Translation<P::Vect> {
    type Item = &'a CollisionObject<P, M, T>;

    #[inline]
    fn next(&mut self) -> Option<&'a CollisionObject<P, M, T>> {
        while let Some(id) = self.idx.next() {
            let co = &self.objects[*id];

            if co.collision_groups.can_collide_with_groups(self.groups) &&
               co.shape.contains_point(&co.position, self.point) {
                return Some(co)
            }
        }

        None
    }
}

/// Iterator through all the objects on the world which bounding volume intersects a specific AABB.
pub struct InterferencesWithAABB<'a, P: 'a, M: 'a, T: 'a> {
    objects: &'a UidRemap<CollisionObject<P, M, T>>,
    groups:  &'a CollisionGroups,
    idx:     IntoIter<&'a FastKey>,
}

impl<'a, P, M, T> Iterator for InterferencesWithAABB<'a, P, M, T> {
    type Item = &'a CollisionObject<P, M, T>;

    #[inline]
    fn next(&mut self) -> Option<&'a CollisionObject<P, M, T>> {
        while let Some(id) = self.idx.next() {
            let co = &self.objects[*id];

            if co.collision_groups.can_collide_with_groups(self.groups) {
                return Some(co)
            }
        }

        None
    }
}

/// 2D collision world containing objects of type `T`.
pub type CollisionWorld2<N, T> = CollisionWorld<Pnt2<N>, Iso2<N>, T>;
/// 3D collision world containing objects of type `T`.
pub type CollisionWorld3<N, T> = CollisionWorld<Pnt3<N>, Iso3<N>, T>;
<|endoftext|>


<commit_before>use std::fmt::Write as _;

use protobuf::reflect::FileDescriptor;
use protobuf::Message;

use crate::gen::code_writer::CodeWriter;
use crate::gen::inside::protobuf_crate_path;
use crate::gen::paths::proto_path_to_fn_file_descriptor;
use crate::gen::rust::snippets::expr_vec_with_capacity;
use crate::gen::scope::FileScope;
use crate::gen::scope::Scope;
use crate::gen::scope::WithScope;
use crate::Customize;

fn escape_byte(s: &mut String, b: u8) {
    if b == b'\n' {
        write!(s, "\\n").unwrap();
    } else if b == b'\r' {
        write!(s, "\\r").unwrap();
    } else if b == b'\t' {
        write!(s, "\\t").unwrap();
    } else if b == b'\\' || b == b'"' {
        write!(s, "\\{}", b as char).unwrap();
    } else if b == b'\0' {
        write!(s, "\\0").unwrap();
    // ASCII printable except space
    } else if b > 0x20 && b < 0x7f {
        write!(s, "{}", b as char).unwrap();
    } else {
        write!(s, "\\x{:02x}", b).unwrap();
    }
}

fn write_file_descriptor(
    file_descriptor: &FileDescriptor,
    customize: &Customize,
    w: &mut CodeWriter,
) {
    w.write_line("/// `FileDescriptor` object which allows dynamic access to files");
    w.pub_fn(
        &format!(
            "file_descriptor() -> {}::reflect::FileDescriptor",
            protobuf_crate_path(customize)
        ),
        |w| {
            w.lazy_static(
                "file_descriptor_lazy",
                &format!(
                    "{}::reflect::GeneratedFileDescriptor",
                    protobuf_crate_path(customize)
                ),
                &format!("{}", protobuf_crate_path(customize)),
            );
            w.block(
                "let file_descriptor = file_descriptor_lazy.get(|| {",
                "});",
                |w| {
                    let deps = &file_descriptor.proto().dependency;
                    w.write_line(&format!(
                        "let mut deps = {};",
                        expr_vec_with_capacity(&format!("{}", deps.len()))
                    ));
                    for f in deps {
                        w.write_line(&format!(
                            "deps.push({}());",
                            proto_path_to_fn_file_descriptor(f, customize)
                        ));
                    }

                    let scope = FileScope { file_descriptor };

                    let messages = scope.find_messages_except_map();
                    w.write_line(&format!(
                        "let mut messages = {};",
                        expr_vec_with_capacity(&format!("{}", messages.len()))
                    ));
                    for m in &messages {
                        w.write_line(&format!(
                            "messages.push({}::generated_message_descriptor_data());",
                            m.rust_name_to_file(),
                        ));
                    }

                    let enums = scope.find_enums();
                    w.write_line(&format!(
                        "let mut enums = {};",
                        expr_vec_with_capacity(&format!("{}", enums.len()))
                    ));
                    for e in &enums {
                        w.write_line(&format!(
                            "enums.push({}::generated_enum_descriptor_data());",
                            e.rust_name_to_file(),
                        ));
                    }

                    w.write_line(&format!(
                        "{}::reflect::GeneratedFileDescriptor::new_generated(",
                        protobuf_crate_path(&customize),
                    ));
                    w.indented(|w| {
                        w.write_line(&format!("file_descriptor_proto(),"));
                        w.write_line(&format!("deps,"));
                        w.write_line(&format!("messages,"));
                        w.write_line(&format!("enums,"));
                    });
                    w.write_line(")");
                },
            );
            w.write_line(&format!(
                "{}::reflect::FileDescriptor::new_generated_2(file_descriptor)",
                protobuf_crate_path(&customize),
            ));
        },
    );
}

pub(crate) fn write_file_descriptor_data(
    file: &FileDescriptor,
    customize: &Customize,
    w: &mut CodeWriter,
) {
    let fdp_bytes = file.proto().write_to_bytes().unwrap();
    w.write_line("static file_descriptor_proto_data: &'static [u8] = b\"\\");
    w.indented(|w| {
        const MAX_LINE_LEN: usize = 72;

        let mut s = String::new();
        for &b in &fdp_bytes {
            let prev_len = s.len();
            escape_byte(&mut s, b);
            let truncate = s.len() > MAX_LINE_LEN;
            if truncate {
                s.truncate(prev_len);
            }
            if truncate || s.len() == MAX_LINE_LEN {
                write!(s, "\\").unwrap();
                w.write_line(&s);
                s.clear();
            }
            if truncate {
                escape_byte(&mut s, b);
            }
        }
        if !s.is_empty() {
            write!(s, "\\").unwrap();
            w.write_line(&s);
            s.clear();
        }
    });
    w.write_line("\";");
    w.write_line("");
    w.write_line("/// `FileDescriptorProto` object which was a source for this generated file");
    w.def_fn(
        &format!(
            "file_descriptor_proto() -> &'static {protobuf_crate}::descriptor::FileDescriptorProto",
            protobuf_crate=protobuf_crate_path(customize)
        ),
        |w| {
            w.lazy_static_decl_get(
                "file_descriptor_proto_lazy",
                &format!(
                    "{protobuf_crate}::descriptor::FileDescriptorProto",
                    protobuf_crate=protobuf_crate_path(customize)
                ),
                &format!("{}", protobuf_crate_path(customize)),
                |w| {
                    w.write_line(&format!(
                        "{protobuf_crate}::Message::parse_from_bytes(file_descriptor_proto_data).unwrap()",
                        protobuf_crate=protobuf_crate_path(customize)
                    ));
                },
            );
        },
    );
    w.write_line("");
    write_file_descriptor(file, &customize, w);
}

/// Code to generate call `module::file_descriptor()`.
pub(crate) fn file_descriptor_call_expr(scope: &Scope) -> String {
    format!(
        "{}()",
        scope
            .rust_path_to_file()
            .to_reverse()
            .append_ident("file_descriptor".into())
    )
}
<commit_msg>Modernize codebase<commit_after>use std::fmt::Write as _;

use protobuf::reflect::FileDescriptor;
use protobuf::Message;

use crate::gen::code_writer::CodeWriter;
use crate::gen::inside::protobuf_crate_path;
use crate::gen::paths::proto_path_to_fn_file_descriptor;
use crate::gen::rust::snippets::expr_vec_with_capacity;
use crate::gen::scope::FileScope;
use crate::gen::scope::Scope;
use crate::gen::scope::WithScope;
use crate::Customize;

fn escape_byte(s: &mut String, b: u8) {
    if b == b'\n' {
        write!(s, "\\n").unwrap();
    } else if b == b'\r' {
        write!(s, "\\r").unwrap();
    } else if b == b'\t' {
        write!(s, "\\t").unwrap();
    } else if b == b'\\' || b == b'"' {
        write!(s, "\\{}", b as char).unwrap();
    } else if b == b'\0' {
        write!(s, "\\0").unwrap();
    // ASCII printable except space
    } else if b > 0x20 && b < 0x7f {
        write!(s, "{}", b as char).unwrap();
    } else {
        write!(s, "\\x{:02x}", b).unwrap();
    }
}

fn write_file_descriptor(
    file_descriptor: &FileDescriptor,
    customize: &Customize,
    w: &mut CodeWriter,
) {
    w.write_line("/// `FileDescriptor` object which allows dynamic access to files");
    w.pub_fn(
        &format!(
            "file_descriptor() -> {protobuf_crate}::reflect::FileDescriptor",
            protobuf_crate = protobuf_crate_path(customize)
        ),
        |w| {
            w.lazy_static(
                "file_descriptor_lazy",
                &format!(
                    "{protobuf_crate}::reflect::GeneratedFileDescriptor",
                    protobuf_crate = protobuf_crate_path(customize)
                ),
                &format!("{}", protobuf_crate_path(customize)),
            );
            w.block(
                "let file_descriptor = file_descriptor_lazy.get(|| {",
                "});",
                |w| {
                    let deps = &file_descriptor.proto().dependency;
                    w.write_line(&format!(
                        "let mut deps = {vec_with_capacity};",
                        vec_with_capacity = expr_vec_with_capacity(&format!("{}", deps.len()))
                    ));
                    for f in deps {
                        w.write_line(&format!(
                            "deps.push({}());",
                            proto_path_to_fn_file_descriptor(f, customize)
                        ));
                    }

                    let scope = FileScope { file_descriptor };

                    let messages = scope.find_messages_except_map();
                    w.write_line(&format!(
                        "let mut messages = {vec_with_capacity};",
                        vec_with_capacity = expr_vec_with_capacity(&format!("{}", messages.len()))
                    ));
                    for m in &messages {
                        w.write_line(&format!(
                            "messages.push({}::generated_message_descriptor_data());",
                            m.rust_name_to_file(),
                        ));
                    }

                    let enums = scope.find_enums();
                    w.write_line(&format!(
                        "let mut enums = {};",
                        expr_vec_with_capacity(&format!("{}", enums.len()))
                    ));
                    for e in &enums {
                        w.write_line(&format!(
                            "enums.push({}::generated_enum_descriptor_data());",
                            e.rust_name_to_file(),
                        ));
                    }

                    w.write_line(&format!(
                        "{}::reflect::GeneratedFileDescriptor::new_generated(",
                        protobuf_crate_path(&customize),
                    ));
                    w.indented(|w| {
                        w.write_line(&format!("file_descriptor_proto(),"));
                        w.write_line(&format!("deps,"));
                        w.write_line(&format!("messages,"));
                        w.write_line(&format!("enums,"));
                    });
                    w.write_line(")");
                },
            );
            w.write_line(&format!(
                "{}::reflect::FileDescriptor::new_generated_2(file_descriptor)",
                protobuf_crate_path(&customize),
            ));
        },
    );
}

pub(crate) fn write_file_descriptor_data(
    file: &FileDescriptor,
    customize: &Customize,
    w: &mut CodeWriter,
) {
    let fdp_bytes = file.proto().write_to_bytes().unwrap();
    w.write_line("static file_descriptor_proto_data: &'static [u8] = b\"\\");
    w.indented(|w| {
        const MAX_LINE_LEN: usize = 72;

        let mut s = String::new();
        for &b in &fdp_bytes {
            let prev_len = s.len();
            escape_byte(&mut s, b);
            let truncate = s.len() > MAX_LINE_LEN;
            if truncate {
                s.truncate(prev_len);
            }
            if truncate || s.len() == MAX_LINE_LEN {
                write!(s, "\\").unwrap();
                w.write_line(&s);
                s.clear();
            }
            if truncate {
                escape_byte(&mut s, b);
            }
        }
        if !s.is_empty() {
            write!(s, "\\").unwrap();
            w.write_line(&s);
            s.clear();
        }
    });
    w.write_line("\";");
    w.write_line("");
    w.write_line("/// `FileDescriptorProto` object which was a source for this generated file");
    w.def_fn(
        &format!(
            "file_descriptor_proto() -> &'static {protobuf_crate}::descriptor::FileDescriptorProto",
            protobuf_crate=protobuf_crate_path(customize)
        ),
        |w| {
            w.lazy_static_decl_get(
                "file_descriptor_proto_lazy",
                &format!(
                    "{protobuf_crate}::descriptor::FileDescriptorProto",
                    protobuf_crate=protobuf_crate_path(customize)
                ),
                &format!("{}", protobuf_crate_path(customize)),
                |w| {
                    w.write_line(&format!(
                        "{protobuf_crate}::Message::parse_from_bytes(file_descriptor_proto_data).unwrap()",
                        protobuf_crate=protobuf_crate_path(customize)
                    ));
                },
            );
        },
    );
    w.write_line("");
    write_file_descriptor(file, &customize, w);
}

/// Code to generate call `module::file_descriptor()`.
pub(crate) fn file_descriptor_call_expr(scope: &Scope) -> String {
    format!(
        "{}()",
        scope
            .rust_path_to_file()
            .to_reverse()
            .append_ident("file_descriptor".into())
    )
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    code_cache::module_cache::ModuleCache,
    data_cache::BlockDataCache,
    process_txn::execute::ExecutedTransaction,
    txn_executor::{TransactionExecutor, LIBRA_SYSTEM_MODULE},
};
use lazy_static::lazy_static;
use libra_types::{
    block_metadata::BlockMetadata,
    identifier::Identifier,
    transaction::TransactionOutput,
    vm_error::{StatusCode, VMStatus},
};
use vm::{
    gas_schedule::{CostTable, GasAlgebra, GasUnits},
    transaction_metadata::TransactionMetadata,
};
use vm_runtime_types::value::Value;

lazy_static! {
    static ref BLOCK_PROLOGUE: Identifier = Identifier::new("block_prologue").unwrap();
}

pub(crate) fn process_block_metadata<'alloc, P>(
    block_metadata: BlockMetadata,
    module_cache: P,
    data_cache: &mut BlockDataCache<'_>,
) -> TransactionOutput
where
    P: ModuleCache<'alloc>,
{
    // TODO: How should we setup the metadata here? A couple of thoughts here:
    // 1. We might make the txn_data to be poisoned so that reading anything will result in a panic.
    // 2. The most important consideration is figuring out the sender address.  Having a notion of a
    //    "null address" (probably 0x0...0) that is prohibited from containing modules or resources
    //    might be useful here.
    // 3. We set the max gas to a big number just to get rid of the potential out of gas error.
    let mut txn_data = TransactionMetadata::default();

    txn_data.max_gas_amount = GasUnits::new(std::u64::MAX);
    // TODO: We might need a non zero cost table here so that we can at least bound the execution
    //       time by a reasonable amount.
    let gas_schedule = CostTable::zero();

    let mut txn_executor =
        TransactionExecutor::new(&module_cache, &gas_schedule, data_cache, txn_data);
    let result = if let Ok((id, timestamp, previous_vote, proposer)) = block_metadata.into_inner() {
        let args = vec![
            Value::u64(timestamp),
            Value::byte_array(id),
            Value::byte_array(previous_vote),
            Value::address(proposer),
        ];
        txn_executor.execute_function(&LIBRA_SYSTEM_MODULE, &BLOCK_PROLOGUE, args)
    } else {
        Err(VMStatus::new(StatusCode::MALFORMED))
    };
    match result {
        Ok(_) => txn_executor.transaction_cleanup(vec![]),
        Err(err) => ExecutedTransaction::discard_error_output(err),
    }
}
<commit_msg>[vm] Fix caching rule for block prologue<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    code_cache::module_cache::ModuleCache,
    data_cache::BlockDataCache,
    process_txn::execute::ExecutedTransaction,
    txn_executor::{TransactionExecutor, LIBRA_SYSTEM_MODULE},
};
use lazy_static::lazy_static;
use libra_types::{
    block_metadata::BlockMetadata,
    identifier::Identifier,
    transaction::TransactionOutput,
    vm_error::{StatusCode, VMStatus},
};
use vm::{
    gas_schedule::{CostTable, GasAlgebra, GasUnits},
    transaction_metadata::TransactionMetadata,
};
use vm_runtime_types::value::Value;

lazy_static! {
    static ref BLOCK_PROLOGUE: Identifier = Identifier::new("block_prologue").unwrap();
}

pub(crate) fn process_block_metadata<'alloc, P>(
    block_metadata: BlockMetadata,
    module_cache: P,
    data_cache: &mut BlockDataCache<'_>,
) -> TransactionOutput
where
    P: ModuleCache<'alloc>,
{
    // TODO: How should we setup the metadata here? A couple of thoughts here:
    // 1. We might make the txn_data to be poisoned so that reading anything will result in a panic.
    // 2. The most important consideration is figuring out the sender address.  Having a notion of a
    //    "null address" (probably 0x0...0) that is prohibited from containing modules or resources
    //    might be useful here.
    // 3. We set the max gas to a big number just to get rid of the potential out of gas error.
    let mut txn_data = TransactionMetadata::default();

    txn_data.max_gas_amount = GasUnits::new(std::u64::MAX);
    // TODO: We might need a non zero cost table here so that we can at least bound the execution
    //       time by a reasonable amount.
    let gas_schedule = CostTable::zero();

    let mut txn_executor =
        TransactionExecutor::new(&module_cache, &gas_schedule, data_cache, txn_data);
    let result = if let Ok((id, timestamp, previous_vote, proposer)) = block_metadata.into_inner() {
        let args = vec![
            Value::u64(timestamp),
            Value::byte_array(id),
            Value::byte_array(previous_vote),
            Value::address(proposer),
        ];
        txn_executor.execute_function(&LIBRA_SYSTEM_MODULE, &BLOCK_PROLOGUE, args)
    } else {
        Err(VMStatus::new(StatusCode::MALFORMED))
    };
    match result.and_then(|_| txn_executor.make_write_set(vec![], Ok(()))) {
        Ok(output) => {
            data_cache.push_write_set(output.write_set());
            output
        }
        Err(err) => ExecutedTransaction::discard_error_output(err),
    }
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::ops::Not;

use crate::requests::*;
use crate::types::*;

/// Use this method to export chat invite links.
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
#[must_use = "requests do nothing unless sent"]
pub struct ExportChatInviteLink {
    chat_id: ChatRef,
}

impl Request for ExportChatInviteLink {
    type Type = JsonRequestType<Self>;
    type Response = JsonIdResponse<String>;

    fn serialize(&self) -> Result<HttpRequest, Error> {
        Self::Type::serialize(RequestUrl::method("exportChatInviteLink"), self)
    }
}

impl ExportChatInviteLink {
    pub fn new<C>(chat: C) -> Self
    where
        C: ToChatRef,
    {
        ExportChatInviteLink {
            chat_id: chat.to_chat_ref(),
        }
    }
}

/// Export chat invite link.
pub trait CanExportChatInviteLink {
    fn export_invite_link(&self) -> ExportChatInviteLink;
}

impl<C> CanExportChatInviteLink for C
where
    C: ToChatRef,
{
    fn export_invite_link(&self) -> ExportChatInviteLink {
        ExportChatInviteLink::new(self)
    }
}
<commit_msg>Remove unused imports<commit_after>use crate::requests::*;
use crate::types::*;

/// Use this method to export chat invite links.
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
#[must_use = "requests do nothing unless sent"]
pub struct ExportChatInviteLink {
    chat_id: ChatRef,
}

impl Request for ExportChatInviteLink {
    type Type = JsonRequestType<Self>;
    type Response = JsonIdResponse<String>;

    fn serialize(&self) -> Result<HttpRequest, Error> {
        Self::Type::serialize(RequestUrl::method("exportChatInviteLink"), self)
    }
}

impl ExportChatInviteLink {
    pub fn new<C>(chat: C) -> Self
    where
        C: ToChatRef,
    {
        ExportChatInviteLink {
            chat_id: chat.to_chat_ref(),
        }
    }
}

/// Export chat invite link.
pub trait CanExportChatInviteLink {
    fn export_invite_link(&self) -> ExportChatInviteLink;
}

impl<C> CanExportChatInviteLink for C
where
    C: ToChatRef,
{
    fn export_invite_link(&self) -> ExportChatInviteLink {
        ExportChatInviteLink::new(self)
    }
}
<|endoftext|>


<commit_before>use crate::convert::*;
use core::hash::{Hasher};

/// A `Hasher` for hashing an arbitrary stream of bytes.
///
/// Instances of [`AHasher`] represent state that is updated while hashing data.
///
/// Each method updates the internal state based on the new data provided. Once
/// all of the data has been provided, the resulting hash can be obtained by calling
/// `finish()`
///
/// [Clone] is also provided in case you wish to calculate hashes for two different items that
/// start with the same data.
///
#[derive(Debug, Clone)]
pub struct AHasher {
    buffer: [u64; 2],
    key: u128,
}

impl AHasher {
    /// Creates a new hasher keyed to the provided keys.
    /// # Example
    ///
    /// ```
    /// use std::hash::Hasher;
    /// use ahash::AHasher;
    ///
    /// let mut hasher = AHasher::new_with_keys(123, 456);
    ///
    /// hasher.write_u32(1989);
    /// hasher.write_u8(11);
    /// hasher.write_u8(9);
    /// hasher.write(b"Huh?");
    ///
    /// println!("Hash is {:x}!", hasher.finish());
    /// ```
    #[inline]
    pub fn new_with_keys(key0: u64, key1: u64) -> Self {
        Self { buffer: [key0, key1], key: [key1, key0].convert()  }
    }

    #[cfg(test)]
    pub(crate) fn test_with_keys(key1: u64, key2: u64) -> AHasher {
        use crate::scramble_keys;
        let (k1, k2) = scramble_keys(key1, key2);
        AHasher { buffer: [k1, k2], key: [k2, k1].convert() }
    }
}

#[inline(never)]
#[no_mangle]
fn hash_test_aes(input: &[u8]) -> u64 {
    let mut a = AHasher::new_with_keys(67, 87);
    a.write(input);
    a.finish()
}

/// Provides methods to hash all of the primitive types.
impl Hasher for AHasher {
    #[inline]
    fn write_u8(&mut self, i: u8) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u16(&mut self, i: u16) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u32(&mut self, i: u32) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u128(&mut self, i: u128) {
        self.buffer = aeshashx2(self.buffer.convert(), i, self.key).convert();
    }

    #[inline]
    fn write_usize(&mut self, i: usize) {
        self.write_u64(i as u64);
    }

    #[inline]
    fn write_u64(&mut self, i: u64) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write(&mut self, input: &[u8]) {
        let mut data = input;
        let length = data.len() as u64;
        //This will be scrambled by the first AES round in any branch.
        self.buffer[1] = self.buffer[1].wrapping_add(length);
        //A 'binary search' on sizes reduces the number of comparisons.
        if data.len() <= 8 {
            if data.len() >= 2 {
                if data.len() >= 4 {
                    //len 4-8
                    self.buffer = aeshash(self.buffer.convert(), data.read_u32().0 as u128).convert();
                    self.buffer = aeshash(self.buffer.convert(), data.read_last_u32() as u128).convert();
                } else {
                    //len 2-3
                    self.buffer = aeshash(self.buffer.convert(), data.read_u16().0 as u128).convert();
                    self.buffer = aeshash(self.buffer.convert(), data[data.len() - 1] as u128).convert();
                }
            } else {
                let value;
                if data.len() > 0 {
                    value = data[0]; //len 1
                } else {
                    value = 0;
                }
                self.buffer = aeshash(self.buffer.convert(), value as u128).convert();
            }
            self.buffer = aeshash(self.buffer.convert(), self.key).convert();
        } else {
            if data.len() > 32 {
                if data.len() > 64 {
                    let (_, tail) = data.split_at(data.len() - 32);
                    let mut par_block: u128 = self.buffer.convert();
                    while data.len() > 32 {
                        let (b1, rest) = data.read_u128();
                        self.buffer = aeshashx2(self.buffer.convert(), b1, self.key).convert();
                        data = rest;
                        let (b2, rest) = data.read_u128();
                        par_block = aeshashx2(par_block, b2, self.key);
                        data = rest;
                    }
                    let (b1, rest) = tail.read_u128();
                    self.buffer = aeshashx2(self.buffer.convert(), b1, self.key).convert();
                    let (b2, _) = rest.read_u128();
                    par_block = aeshashx2(par_block, b2, self.key);
                    self.buffer = aeshashx2(self.buffer.convert(), par_block, self.key).convert();
                } else {
                    //len 33-64
                    let (head, _) = data.split_at(32);
                    let (_, tail) = data.split_at(data.len() - 32);
                    self.buffer = aeshashx2(self.buffer.convert(), head.read_u128().0, self.key).convert();
                    self.buffer = aeshashx2(self.buffer.convert(), head.read_last_u128(), self.key).convert();
                    self.buffer = aeshashx2(self.buffer.convert(), tail.read_u128().0, self.key).convert();
                    self.buffer = aeshashx2(self.buffer.convert(), tail.read_last_u128(), self.key).convert();
                }
            } else {
                if data.len() > 16 {
                    //len 17-32
                    self.buffer = aeshashx2(self.buffer.convert(), data.read_u128().0, self.key).convert();
                    self.buffer = aeshashx2(self.buffer.convert(), data.read_last_u128(), self.key).convert();
                } else {
                    //len 9-16
                    self.buffer = aeshash(self.buffer.convert(), data.read_u64().0 as u128).convert();
                    self.buffer = aeshashx2(self.buffer.convert(), data.read_last_u64() as u128, self.key).convert();
                }
            }
        }
    }
    #[inline]
    fn finish(&self) -> u64 {
        let result: [u64; 2] = aeshash(self.buffer.convert(), self.key).convert();
        result[0] //.wrapping_add(result[1])
    }
}

#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes"))]
#[inline(always)]
fn aeshash(value: u128, xor: u128) -> u128 {
    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use core::arch::x86_64::*;
    use core::mem::transmute;
    unsafe {
        let value = transmute(value);
        transmute(_mm_aesdec_si128(value, transmute(xor)))
    }
}
#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes"))]
#[inline(always)]
fn aeshashx2(value: u128, k1: u128, k2: u128) -> u128 {
    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use core::arch::x86_64::*;
    use core::mem::transmute;
    unsafe {
        let value = transmute(value);
        let value = _mm_aesdec_si128(value, transmute(k1));
        transmute(_mm_aesdec_si128(value, transmute(k2)))
    }
}

#[cfg(test)]
mod tests {
    use crate::aes_hash::*;
    use crate::convert::Convert;
    use std::collections::HashMap;
    use std::hash::BuildHasherDefault;

    #[cfg(feature = "compile-time-rng")]
    #[test]
    fn test_builder() {
        let mut map = HashMap::<u32, u64, BuildHasherDefault<AHasher>>::default();
        map.insert(1, 3);
    }

    #[cfg(feature = "compile-time-rng")]
    #[test]
    fn test_default() {
        let hasher_a = AHasher::default();
        assert_ne!(0, hasher_a.buffer[0]);
        assert_ne!(0, hasher_a.buffer[1]);
        assert_ne!(hasher_a.buffer[0], hasher_a.buffer[1]);
        let hasher_b = AHasher::default();
        assert_eq!(hasher_a.buffer[0], hasher_b.buffer[0]);
        assert_eq!(hasher_a.buffer[1], hasher_b.buffer[1]);
    }

    #[test]
    fn test_hash() {
        let mut result: [u64; 2] = [0x6c62272e07bb0142, 0x62b821756295c58d];
        let value: [u64; 2] = [1 << 32, 0xFEDCBA9876543210];
        result = aeshash(value.convert(), result.convert()).convert();
        result = aeshash(result.convert(), result.convert()).convert();
        let mut result2: [u64; 2] = [0x6c62272e07bb0142, 0x62b821756295c58d];
        let value2: [u64; 2] = [1, 0xFEDCBA9876543210];
        result2 = aeshash(value2.convert(), result2.convert()).convert();
        result2 = aeshash(result2.convert(), result.convert()).convert();
        let result: [u8; 16] = result.convert();
        let result2: [u8; 16] = result2.convert();
        assert_ne!(hex::encode(result), hex::encode(result2));
    }

    #[test]
    fn test_conversion() {
        let input: &[u8] = "dddddddd".as_bytes();
        let bytes: u64 = as_array!(input, 8).convert();
        assert_eq!(bytes, 0x6464646464646464);
    }
}
<commit_msg>Simplify loading for short strings<commit_after>use crate::convert::*;
use core::hash::Hasher;

/// A `Hasher` for hashing an arbitrary stream of bytes.
///
/// Instances of [`AHasher`] represent state that is updated while hashing data.
///
/// Each method updates the internal state based on the new data provided. Once
/// all of the data has been provided, the resulting hash can be obtained by calling
/// `finish()`
///
/// [Clone] is also provided in case you wish to calculate hashes for two different items that
/// start with the same data.
///
#[derive(Debug, Clone)]
pub struct AHasher {
    buffer: [u64; 2],
    key: u128,
}

impl AHasher {
    /// Creates a new hasher keyed to the provided keys.
    /// # Example
    ///
    /// ```
    /// use std::hash::Hasher;
    /// use ahash::AHasher;
    ///
    /// let mut hasher = AHasher::new_with_keys(123, 456);
    ///
    /// hasher.write_u32(1989);
    /// hasher.write_u8(11);
    /// hasher.write_u8(9);
    /// hasher.write(b"Huh?");
    ///
    /// println!("Hash is {:x}!", hasher.finish());
    /// ```
    #[inline]
    pub fn new_with_keys(key0: u64, key1: u64) -> Self {
        Self {
            buffer: [key0, key1],
            key: [key1, key0].convert(),
        }
    }

    #[cfg(test)]
    pub(crate) fn test_with_keys(key1: u64, key2: u64) -> AHasher {
        use crate::scramble_keys;
        let (k1, k2) = scramble_keys(key1, key2);
        AHasher {
            buffer: [k1, k2],
            key: [k2, k1].convert(),
        }
    }
}

#[inline(never)]
#[no_mangle]
fn hash_test_aes(input: &[u8]) -> u64 {
    let mut a = AHasher::new_with_keys(67, 87);
    a.write(input);
    a.finish()
}

/// Provides methods to hash all of the primitive types.
impl Hasher for AHasher {
    #[inline]
    fn write_u8(&mut self, i: u8) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u16(&mut self, i: u16) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u32(&mut self, i: u32) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u128(&mut self, i: u128) {
        self.buffer = aeshashx2(self.buffer.convert(), i, self.key).convert();
    }

    #[inline]
    fn write_usize(&mut self, i: usize) {
        self.write_u64(i as u64);
    }

    #[inline]
    fn write_u64(&mut self, i: u64) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write(&mut self, input: &[u8]) {
        let mut data = input;
        let length = data.len() as u64;
        //This will be scrambled by the first AES round in any branch.
        self.buffer[1] = self.buffer[1].wrapping_add(length);
        //A 'binary search' on sizes reduces the number of comparisons.
        if data.len() <= 8 {
            let value: [u64; 2] = if data.len() >= 2 {
                if data.len() >= 4 {
                    //len 4-8
                    [data.read_u32().0 as u64, data.read_last_u32() as u64]
                } else {
                    //len 2-3
                    [data.read_u16().0 as u64, data[data.len() - 1] as u64]
                }
            } else {
                if data.len() > 0 {
                    [data[0] as u64, 0]
                } else {
                    [0, 0]
                }
            };
            self.buffer = aeshashx2(self.buffer.convert(), value.convert(), self.key).convert();
        } else {
            if data.len() > 32 {
                if data.len() > 64 {
                    let (_, tail) = data.split_at(data.len() - 32);
                    let mut par_block: u128 = self.buffer.convert();
                    while data.len() > 32 {
                        let (b1, rest) = data.read_u128();
                        self.buffer = aeshashx2(self.buffer.convert(), b1, self.key).convert();
                        data = rest;
                        let (b2, rest) = data.read_u128();
                        par_block = aeshashx2(par_block, b2, self.key);
                        data = rest;
                    }
                    let (b1, rest) = tail.read_u128();
                    self.buffer = aeshashx2(self.buffer.convert(), b1, self.key).convert();
                    let (b2, _) = rest.read_u128();
                    par_block = aeshashx2(par_block, b2, self.key);
                    self.buffer = aeshashx2(self.buffer.convert(), par_block, self.key).convert();
                } else {
                    //len 33-64
                    let (head, _) = data.split_at(32);
                    let (_, tail) = data.split_at(data.len() - 32);
                    self.buffer = aeshashx2(self.buffer.convert(), head.read_u128().0, self.key).convert();
                    self.buffer = aeshashx2(self.buffer.convert(), head.read_last_u128(), self.key).convert();
                    self.buffer = aeshashx2(self.buffer.convert(), tail.read_u128().0, self.key).convert();
                    self.buffer = aeshashx2(self.buffer.convert(), tail.read_last_u128(), self.key).convert();
                }
            } else {
                if data.len() > 16 {
                    //len 17-32
                    self.buffer = aeshashx2(self.buffer.convert(), data.read_u128().0, self.key).convert();
                    self.buffer = aeshashx2(self.buffer.convert(), data.read_last_u128(), self.key).convert();
                } else {
                    //len 9-16
                    let value: [u64; 2] = [data.read_u64().0, data.read_last_u64()];
                    self.buffer = aeshashx2(self.buffer.convert(), value.convert(), self.key).convert();
                }
            }
        }
    }
    #[inline]
    fn finish(&self) -> u64 {
        let result: [u64; 2] = aeshash(self.buffer.convert(), self.key).convert();
        result[0] //.wrapping_add(result[1])
    }
}

#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes"))]
#[inline(always)]
fn aeshash(value: u128, xor: u128) -> u128 {
    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use core::arch::x86_64::*;
    use core::mem::transmute;
    unsafe {
        let value = transmute(value);
        transmute(_mm_aesdec_si128(value, transmute(xor)))
    }
}
#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes"))]
#[inline(always)]
fn aeshashx2(value: u128, k1: u128, k2: u128) -> u128 {
    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use core::arch::x86_64::*;
    use core::mem::transmute;
    unsafe {
        let value = transmute(value);
        let value = _mm_aesdec_si128(value, transmute(k1));
        transmute(_mm_aesdec_si128(value, transmute(k2)))
    }
}

#[cfg(test)]
mod tests {
    use crate::aes_hash::*;
    use crate::convert::Convert;
    use std::collections::HashMap;
    use std::hash::BuildHasherDefault;

    #[cfg(feature = "compile-time-rng")]
    #[test]
    fn test_builder() {
        let mut map = HashMap::<u32, u64, BuildHasherDefault<AHasher>>::default();
        map.insert(1, 3);
    }

    #[cfg(feature = "compile-time-rng")]
    #[test]
    fn test_default() {
        let hasher_a = AHasher::default();
        assert_ne!(0, hasher_a.buffer[0]);
        assert_ne!(0, hasher_a.buffer[1]);
        assert_ne!(hasher_a.buffer[0], hasher_a.buffer[1]);
        let hasher_b = AHasher::default();
        assert_eq!(hasher_a.buffer[0], hasher_b.buffer[0]);
        assert_eq!(hasher_a.buffer[1], hasher_b.buffer[1]);
    }

    #[test]
    fn test_hash() {
        let mut result: [u64; 2] = [0x6c62272e07bb0142, 0x62b821756295c58d];
        let value: [u64; 2] = [1 << 32, 0xFEDCBA9876543210];
        result = aeshash(value.convert(), result.convert()).convert();
        result = aeshash(result.convert(), result.convert()).convert();
        let mut result2: [u64; 2] = [0x6c62272e07bb0142, 0x62b821756295c58d];
        let value2: [u64; 2] = [1, 0xFEDCBA9876543210];
        result2 = aeshash(value2.convert(), result2.convert()).convert();
        result2 = aeshash(result2.convert(), result.convert()).convert();
        let result: [u8; 16] = result.convert();
        let result2: [u8; 16] = result2.convert();
        assert_ne!(hex::encode(result), hex::encode(result2));
    }

    #[test]
    fn test_conversion() {
        let input: &[u8] = "dddddddd".as_bytes();
        let bytes: u64 = as_array!(input, 8).convert();
        assert_eq!(bytes, 0x6464646464646464);
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::collections::VecDeque;
use num::{ToPrimitive};

use log::{warn, debug};
use failure::{Error, Fail};
use zydis::gen::*;

use super::arch::{Arch};
use super::loader::{LoadedModule, Section};
use super::flowmeta::FlowMeta;
use super::workspace::{Workspace};

#[derive(Debug, Fail)]
pub enum AnalysisError {
    #[fail(display = "foo")]
    Foo,
    #[fail(display = "foo")]
    InvalidInstruction
}

pub enum AnalysisCommand<A: Arch> {
    MakeInsn(A::RVA),
}

pub struct XrefAnalysis<A: Arch> {
    // TODO: use FNV because the keys are small.
    // TODO: use SmallVec(1) for `.from` values,
    // TODO: use SmallVec(X) for `.to` values,

    // dst rva -> src rva
    to: HashMap<A::RVA, A::RVA>,
    // src rva -> dst rva
    from: HashMap<A::RVA, A::RVA>,
}

pub struct FlowAnalysis<A: Arch> {
    // one entry for each section in the module.
    // if executable, then one FlowMeta for each address in the section.
    // that is, Vec<FlowMeta>.len() == Section.buf.len()
    // TODO: order these entries so that the most common sections are first (`.code`?)
    meta: Vec<Vec<FlowMeta>>,
    xrefs: XrefAnalysis<A>,
}

pub struct Analysis<A: Arch> {
    queue: VecDeque<AnalysisCommand<A>>,

    flow: FlowAnalysis<A>,
    // datameta
    // symbols
    // functions
}

impl<A: Arch> Analysis<A> {
    pub fn new(module: &LoadedModule<A>) -> Analysis<A> {
        let flow_meta: Vec<Vec<FlowMeta>> = module.sections.iter()
            .map(|section| -> Vec<FlowMeta> {
                if section.is_executable() {
                    vec![FlowMeta::zero(); section.buf.len()]
                } else {
                    vec![]
                }
            })
            .collect();

        Analysis {
            queue: VecDeque::new(),
            flow: FlowAnalysis {
                meta: flow_meta,
                xrefs: XrefAnalysis {
                    to: HashMap::new(),
                    from: HashMap::new(),
                }
            }
        }
    }

}

// here we've logically split off the analysis portion of workspace.
// this should keep file sizes smaller, and hopefully easier to understand.
impl<A: Arch + 'static> Workspace<A> {
    pub fn make_insn(&mut self, rva: A::RVA) -> Result<(), Error> {
        self.analysis.queue.push_back(AnalysisCommand::MakeInsn(rva));
        Ok(())
    }

    fn get_meta(&self, rva: A::RVA) -> Option<FlowMeta> {
        self.module.sections
            .iter()
            .enumerate()
            .filter(|(_, section)| section.contains(rva))
            .nth(0)
            .and_then(|(i, section): (usize, &Section<A>)| -> Option<FlowMeta> {
                // rva is guaranteed to be within this section,
                // so we can do an unchecked subtract here.
                let offset = rva - section.addr;
                A::RVA::to_usize(&offset)
                    .and_then(|offset| {
                        Some(self.analysis.flow.meta[i][offset].clone())
                    })
            })
    }

    /// Does the given instruction have a fallthrough flow?
    ///
    /// ```
    /// use lancelot::test;
    ///
    /// // JMP $+0;
    /// let insn = test::get_shellcode32_workspace(b"\xEB\xFE").read_insn(0x0).unwrap();
    /// assert_eq!(ws::does_insn_fallthrough(&insn), false);
    ///
    /// // PUSH 0x11
    /// let ws = test::get_shellcode32_workspace(b"\x6A\x11").read_insn(0x0).unwrap();
    /// assert_eq!(ws::does_insn_fallthrough(&insn), true);
    /// ```
    pub fn does_insn_fallthrough(insn: &ZydisDecodedInstruction) -> bool {
        match insn.mnemonic as i32 {
            ZYDIS_MNEMONIC_JMP => false,
            ZYDIS_MNEMONIC_RET => false,
            ZYDIS_MNEMONIC_IRET => false,
            ZYDIS_MNEMONIC_IRETD => false,
            ZYDIS_MNEMONIC_IRETQ => false,
            // TODO: call may not fallthrough if function is noret.
            // will need another pass to clean this up.
            ZYDIS_MNEMONIC_CALL => true,
            _ => true,
        }
    }

    fn handle_make_insn(&mut self, rva: A::RVA) -> Result<Vec<AnalysisCommand<A>>, Error> {
        let mut ret = vec![];

        // TODO: 0. probe address

        // 1. ensure instruction doesn't already exist
        //
        // if we get a result here, then there's not yet an instruction at the rva.
        // otherwise, we will have returned early, and there's no work to be done.
        //
        // now, we might worry about lots of extra allocations for the Vec if the insn already exists.
        //  but, its not a problem: Vec only allocates when there's a non-zero element in it!
        // so, its mostly ok to spam `make_insn`.
        let mut meta = match self.get_meta(rva) {
            None => {
                // this might happen if:
                //   - the instruction is in a non-executable section
                //   - the memory is not mapped
                warn!("invalid instruction: no flow meta: {:x}", rva);
                return Ok(vec![]);
            },
            Some(meta) => {
                if meta.is_insn() {
                    debug!("duplicate instruction: {:x}", rva);
                    return Ok(vec![]);
                } else {
                    // this is the happy path:
                    //   - this is a valid place for an instruction
                    //   - and it doesn't yet exist
                    meta
                }
            }
        };

        let insn = match self.read_insn(rva) {
            Err(e) => {
                warn!("invalid instruction: {:}: {:x}", e, rva);
                return Ok(vec![]);
            },
            Ok(insn) => insn,
        };

        // 2. compute instruction len
        let length = insn.length;

        // 3. compute fallthrough
        let does_fallthrough = Workspace::<A>::does_insn_fallthrough(&insn);

        // 4. compute flow ref
        // 5. update flowmeta
        Ok(ret)
    }

    pub fn analyze(&mut self) -> Result<(), Error> {
        while let Some(cmd) = self.analysis.queue.pop_front() {
            let cmds = match cmd {
                AnalysisCommand::MakeInsn(rva) => self.handle_make_insn(rva)?,
            };
            self.analysis.queue.extend(cmds);
        }

        Ok(())
    }
}
<commit_msg>analysis: begin to sketch out insn flow analysis<commit_after>use std::collections::HashMap;
use std::collections::VecDeque;
use num::{ToPrimitive};

use log::{warn, debug};
use failure::{Error, Fail};
use zydis::gen::*;

use super::arch::{Arch};
use super::xref::{Xref, XrefType};
use super::loader::{LoadedModule, Section};
use super::flowmeta::FlowMeta;
use super::workspace::{Workspace};

#[derive(Debug, Fail)]
pub enum AnalysisError {
    #[fail(display = "foo")]
    Foo,
    #[fail(display = "foo")]
    InvalidInstruction
}

pub enum AnalysisCommand<A: Arch> {
    MakeInsn(A::RVA),
}

pub struct XrefAnalysis<A: Arch> {
    // TODO: use FNV because the keys are small.
    // TODO: use SmallVec(1) for `.from` values,
    // TODO: use SmallVec(X) for `.to` values,

    // dst rva -> src rva
    to: HashMap<A::RVA, A::RVA>,
    // src rva -> dst rva
    from: HashMap<A::RVA, A::RVA>,
}

pub struct FlowAnalysis<A: Arch> {
    // one entry for each section in the module.
    // if executable, then one FlowMeta for each address in the section.
    // that is, Vec<FlowMeta>.len() == Section.buf.len()
    // TODO: order these entries so that the most common sections are first (`.code`?)
    meta: Vec<Vec<FlowMeta>>,
    xrefs: XrefAnalysis<A>,
}

pub struct Analysis<A: Arch> {
    queue: VecDeque<AnalysisCommand<A>>,

    flow: FlowAnalysis<A>,
    // datameta
    // symbols
    // functions
}

impl<A: Arch> Analysis<A> {
    pub fn new(module: &LoadedModule<A>) -> Analysis<A> {
        let flow_meta: Vec<Vec<FlowMeta>> = module.sections.iter()
            .map(|section| -> Vec<FlowMeta> {
                if section.is_executable() {
                    vec![FlowMeta::zero(); section.buf.len()]
                } else {
                    vec![]
                }
            })
            .collect();

        Analysis {
            queue: VecDeque::new(),
            flow: FlowAnalysis {
                meta: flow_meta,
                xrefs: XrefAnalysis {
                    to: HashMap::new(),
                    from: HashMap::new(),
                }
            }
        }
    }

}

// here we've logically split off the analysis portion of workspace.
// this should keep file sizes smaller, and hopefully easier to understand.
impl<A: Arch + 'static> Workspace<A> {
    pub fn make_insn(&mut self, rva: A::RVA) -> Result<(), Error> {
        self.analysis.queue.push_back(AnalysisCommand::MakeInsn(rva));
        Ok(())
    }

    fn get_meta(&self, rva: A::RVA) -> Option<FlowMeta> {
        self.module.sections
            .iter()
            .enumerate()
            .filter(|(_, section)| section.contains(rva))
            .nth(0)
            .and_then(|(i, section): (usize, &Section<A>)| -> Option<FlowMeta> {
                // rva is guaranteed to be within this section,
                // so we can do an unchecked subtract here.
                let offset = rva - section.addr;
                A::RVA::to_usize(&offset)
                    .and_then(|offset| {
                        Some(self.analysis.flow.meta[i][offset].clone())
                    })
            })
    }

    /// Does the given instruction have a fallthrough flow?
    ///
    /// ```
    /// use lancelot::test;
    ///
    /// // JMP $+0;
    /// let insn = test::get_shellcode32_workspace(b"\xEB\xFE").read_insn(0x0).unwrap();
    /// assert_eq!(ws::does_insn_fallthrough(&insn), false);
    ///
    /// // PUSH 0x11
    /// let ws = test::get_shellcode32_workspace(b"\x6A\x11").read_insn(0x0).unwrap();
    /// assert_eq!(ws::does_insn_fallthrough(&insn), true);
    /// ```
    pub fn does_insn_fallthrough(insn: &ZydisDecodedInstruction) -> bool {
        match insn.mnemonic as i32 {
            ZYDIS_MNEMONIC_JMP => false,
            ZYDIS_MNEMONIC_RET => false,
            ZYDIS_MNEMONIC_IRET => false,
            ZYDIS_MNEMONIC_IRETD => false,
            ZYDIS_MNEMONIC_IRETQ => false,
            // TODO: call may not fallthrough if function is noret.
            // will need another pass to clean this up.
            ZYDIS_MNEMONIC_CALL => true,
            _ => true,
        }
    }

    fn get_call_insn_flow(&self, rva: A::RVA, insn: &ZydisDecodedInstruction) -> Result<Vec<Xref<A>>, Error> {
        // TODO
        Ok(vec![])
    }
    fn get_jmp_insn_flow(&self, rva: A::RVA, insn: &ZydisDecodedInstruction) -> Result<Vec<Xref<A>>, Error> {
        // TODO
        Ok(vec![])
    }
    fn get_ret_insn_flow(&self, rva: A::RVA, insn: &ZydisDecodedInstruction) -> Result<Vec<Xref<A>>, Error> {
        // TODO
        Ok(vec![])
    }
    fn get_cjmp_insn_flow(&self, rva: A::RVA, insn: &ZydisDecodedInstruction) -> Result<Vec<Xref<A>>, Error> {
        // TODO
        Ok(vec![])
    }
    fn get_cmov_insn_flow(&self, rva: A::RVA, insn: &ZydisDecodedInstruction) -> Result<Vec<Xref<A>>, Error> {
        // TODO
        Ok(vec![])
    }

    fn get_insn_flow(&self, rva: A::RVA, insn: &ZydisDecodedInstruction) -> Result<Vec<Xref<A>>, Error> {
        match insn.mnemonic as i32 {
            ZYDIS_MNEMONIC_CALL => self.get_call_insn_flow(rva, insn),

            ZYDIS_MNEMONIC_JMP => self.get_jmp_insn_flow(rva, insn),

            ZYDIS_MNEMONIC_RET
            | ZYDIS_MNEMONIC_IRET
            | ZYDIS_MNEMONIC_IRETD
            | ZYDIS_MNEMONIC_IRETQ => self.get_ret_insn_flow(rva, insn),

            ZYDIS_MNEMONIC_JB
            | ZYDIS_MNEMONIC_JBE
            | ZYDIS_MNEMONIC_JCXZ
            | ZYDIS_MNEMONIC_JECXZ
            | ZYDIS_MNEMONIC_JKNZD
            | ZYDIS_MNEMONIC_JKZD
            | ZYDIS_MNEMONIC_JL
            | ZYDIS_MNEMONIC_JLE
            | ZYDIS_MNEMONIC_JNB
            | ZYDIS_MNEMONIC_JNBE
            | ZYDIS_MNEMONIC_JNL
            | ZYDIS_MNEMONIC_JNLE
            | ZYDIS_MNEMONIC_JNO
            | ZYDIS_MNEMONIC_JNP
            | ZYDIS_MNEMONIC_JNS
            | ZYDIS_MNEMONIC_JNZ
            | ZYDIS_MNEMONIC_JO
            | ZYDIS_MNEMONIC_JP
            | ZYDIS_MNEMONIC_JRCXZ
            | ZYDIS_MNEMONIC_JS
            | ZYDIS_MNEMONIC_JZ => self.get_cjmp_insn_flow(rva, insn),

            ZYDIS_MNEMONIC_CMOVB
            | ZYDIS_MNEMONIC_CMOVBE
            | ZYDIS_MNEMONIC_CMOVL
            | ZYDIS_MNEMONIC_CMOVLE
            | ZYDIS_MNEMONIC_CMOVNB
            | ZYDIS_MNEMONIC_CMOVNBE
            | ZYDIS_MNEMONIC_CMOVNL
            | ZYDIS_MNEMONIC_CMOVNLE
            | ZYDIS_MNEMONIC_CMOVNO
            | ZYDIS_MNEMONIC_CMOVNP
            | ZYDIS_MNEMONIC_CMOVNS
            | ZYDIS_MNEMONIC_CMOVNZ
            | ZYDIS_MNEMONIC_CMOVO
            | ZYDIS_MNEMONIC_CMOVP
            | ZYDIS_MNEMONIC_CMOVS
            | ZYDIS_MNEMONIC_CMOVZ => self.get_cmov_insn_flow(rva, insn),

            // TODO: syscall, sysexit, sysret, vmcall, vmmcall

            _ => Ok(vec![]),
        }
    }

    fn handle_make_insn(&mut self, rva: A::RVA) -> Result<Vec<AnalysisCommand<A>>, Error> {
        let mut ret = vec![];

        // TODO: 0. probe address

        // 1. ensure instruction doesn't already exist
        //
        // if we get a result here, then there's not yet an instruction at the rva.
        // otherwise, we will have returned early, and there's no work to be done.
        //
        // now, we might worry about lots of extra allocations for the Vec if the insn already exists.
        //  but, its not a problem: Vec only allocates when there's a non-zero element in it!
        // so, its mostly ok to spam `make_insn`.
        let mut meta = match self.get_meta(rva) {
            None => {
                // this might happen if:
                //   - the instruction is in a non-executable section
                //   - the memory is not mapped
                warn!("invalid instruction: no flow meta: {:x}", rva);
                return Ok(vec![]);
            },
            Some(meta) => {
                if meta.is_insn() {
                    debug!("duplicate instruction: {:x}", rva);
                    return Ok(vec![]);
                } else {
                    // this is the happy path:
                    //   - this is a valid place for an instruction
                    //   - and it doesn't yet exist
                    meta
                }
            }
        };

        let insn = match self.read_insn(rva) {
            Err(e) => {
                warn!("invalid instruction: {:}: {:x}", e, rva);
                return Ok(vec![]);
            },
            Ok(insn) => insn,
        };

        // 2. compute instruction len
        let length = insn.length;

        // 3. compute fallthrough
        let does_fallthrough = Workspace::<A>::does_insn_fallthrough(&insn);

        // 4. compute flow ref
        let flows = self.get_insn_flow(rva, &insn)?;

        // 5. update flowmeta
        Ok(ret)
    }

    pub fn analyze(&mut self) -> Result<(), Error> {
        while let Some(cmd) = self.analysis.queue.pop_front() {
            let cmds = match cmd {
                AnalysisCommand::MakeInsn(rva) => self.handle_make_insn(rva)?,
            };
            self.analysis.queue.extend(cmds);
        }

        Ok(())
    }
}
<|endoftext|>