<commit_before>use std::fs::File;
use std::io::BufReader;
use std::collections::HashMap;
use std::borrow::Cow;

use zip::read::{ZipFile, ZipArchive};
use zip::result::ZipError;
use quick_xml::{XmlReader, Event, AsStr, Element};

use {DataType, ExcelReader, Range, Cell};
use vba::VbaProject;
use errors::*;

/// A struct representing xml zipped excel file
/// Xlsx, Xlsm, Xlam
pub struct Xlsx {
    zip: ZipArchive<File>,
}

impl Xlsx {
    fn xml_reader<'a>(&'a mut self,
                      path: &str)
                      -> Option<Result<XmlReader<BufReader<ZipFile<'a>>>>> {
        match self.zip.by_name(path) {
            Ok(f) => {
                Some(Ok(XmlReader::from_reader(BufReader::new(f))
                    .with_check(false)
                    .trim_text(false)))
            }
            Err(ZipError::FileNotFound) => None,
            Err(e) => return Some(Err(e.into())),
        }
    }
}

impl ExcelReader for Xlsx {
    fn new(f: File) -> Result<Self> {
        Ok(Xlsx { zip: ZipArchive::new(f)? })
    }

    fn has_vba(&mut self) -> bool {
        self.zip.by_name("xl/vbaProject.bin").is_ok()
    }

    fn vba_project(&mut self) -> Result<Cow<VbaProject>> {
        let mut f = self.zip.by_name("xl/vbaProject.bin")?;
        let len = f.size() as usize;
        VbaProject::new(&mut f, len).map(|v| Cow::Owned(v))
    }

    fn read_shared_strings(&mut self) -> Result<Vec<String>> {
        let mut xml = match self.xml_reader("xl/sharedStrings.xml") {
            None => return Ok(Vec::new()),
            Some(x) => x?,
        };
        let mut strings = Vec::new();
        let mut rich_buffer: Option<String> = None;
        while let Some(res_event) = xml.next() {
            match res_event {
                Ok(Event::Start(ref e)) if e.name() == b"r" => {
                    if let None = rich_buffer {
                        // use a buffer since richtext has multiples <r> and <t> for the same cell
                        rich_buffer = Some(String::new());
                    }
                }
                Ok(Event::End(ref e)) if e.name() == b"si" => {
                    if let Some(s) = rich_buffer {
                        strings.push(s);
                        rich_buffer = None;
                    }
                }
                Ok(Event::Start(ref e)) if e.name() == b"t" => {
                    let value = xml.read_text_unescaped(b"t")?;
                    if let Some(ref mut s) = rich_buffer {
                        s.push_str(&value);
                    } else {
                        strings.push(value);
                    }
                }
                Err(e) => return Err(e.into()),
                _ => (),
            }
        }
        Ok(strings)
    }

    fn read_sheets_names(&mut self,
                         relationships: &HashMap<Vec<u8>, String>)
                         -> Result<Vec<(String, String)>> {
        let xml = match self.xml_reader("xl/workbook.xml") {
            None => return Ok(Vec::new()),
            Some(x) => x?,
        };
        let mut sheets = Vec::new();
        for res_event in xml {
            match res_event {
                Ok(Event::Start(ref e)) if e.name() == b"sheet" => {
                    let mut name = String::new();
                    let mut path = String::new();
                    for a in e.unescaped_attributes() {
                        match a? {
                            (b"name", v) => name = v.as_str()?.to_string(),
                            (b"r:id", v) => {
                                let r = &relationships[&*v][..];
                                // target may have pre-prended "/xl/" or "xl/" path;
                                // strip if present
                                let r = if r.starts_with("/xl/") {
                                    &r[4..]
                                } else if r.starts_with("xl/") {
                                    &r[3..]
                                } else {
                                    r
                                };
                                path = format!("xl/{}", r);
                            }
                            _ => (),
                        }
                    }
                    sheets.push((name, path));
                }
                Err(e) => return Err(e.into()),
                _ => (),
            }
        }
        Ok(sheets)
    }

    fn read_relationships(&mut self) -> Result<HashMap<Vec<u8>, String>> {
        let xml = match self.xml_reader("xl/_rels/workbook.xml.rels") {
            None => return Err("Cannot find relationships file".into()),
            Some(x) => x?,
        };
        let mut relationships = HashMap::new();
        for res_event in xml {
            match res_event {
                Ok(Event::Start(ref e)) if e.name() == b"Relationship" => {
                    let mut id = Vec::new();
                    let mut target = String::new();
                    for a in e.attributes() {
                        match a? {
                            (b"Id", v) => id.extend_from_slice(v),
                            (b"Target", v) => target = v.as_str()?.to_string(),
                            _ => (),
                        }
                    }
                    relationships.insert(id, target);
                }
                Err(e) => return Err(e.into()),
                _ => (),
            }
        }
        Ok(relationships)
    }

    fn read_worksheet_range(&mut self, path: &str, strings: &[String]) -> Result<Range> {
        let mut xml = match self.xml_reader(path) {
            None => return Err(format!("Cannot find {} path", path).into()),
            Some(x) => x?,
        };
        let mut cells = Vec::new();
        'xml: while let Some(res_event) = xml.next() {
            match res_event {
                Err(e) => return Err(e.into()),
                Ok(Event::Start(ref e)) => {
                    match e.name() {
                        b"dimension" => {
                            for a in e.attributes() {
                                if let (b"ref", rdim) = a? {
                                    let (start, end) = get_dimension(rdim)?;
                                    cells.reserve(((end.0 - start.0 + 1) * (end.1 - start.1 + 1)) as
                                                 usize);
                                    continue 'xml;
                                }
                            }
                            return Err(format!("Expecting dimension, got {:?}", e).into());
                        }
                        b"sheetData" => read_sheet_data(&mut xml, strings, &mut cells)?,
                        _ => (),
                    }
                }
                _ => (),
            }
        }
        Ok(Range::from_sparse(cells))
    }
}

/// read sheetData node
fn read_sheet_data(xml: &mut XmlReader<BufReader<ZipFile>>,
                   strings: &[String],
                   cells: &mut Vec<Cell>)
                   -> Result<()> {
    /// read the contents of an <is> cell
    fn read_inline_str(xml: &mut XmlReader<BufReader<ZipFile>>) -> Result<DataType> {
        while let Some(is_evt) = xml.next() {
            match is_evt {
                Err(e) => return Err(e.into()),
                Ok(Event::Start(ref t)) if t.name() == b"t" => {
                    return Ok(DataType::String(xml.read_text_unescaped(b"t")?));
                }
                _ => {}
            }
        }
        Err("unable to read inlineStr".into())
    }
    /// read the contents of a <v> cell
    fn read_value(xml: &mut XmlReader<BufReader<ZipFile>>,
                  strings: &[String],
                  c_element: &Element)
                  -> Result<DataType> {
        let v = xml.read_text_unescaped(b"v")?;
        match get_attribute(c_element, b"t")? {
            Some(b"s") => {
                // shared string
                let idx: usize = v.parse()?;
                Ok(DataType::String(strings[idx].clone()))
            }
            Some(b"b") => {
                // boolean
                Ok(DataType::Bool(v != "0"))
            }
            Some(b"e") => {
                // error
                Ok(DataType::Error(v.parse()?))
            }
            Some(b"d") => {
                // date
                // TODO: create a DataType::Date
                // currently just return as string (ISO 8601)
                Ok(DataType::String(v))
            }
            Some(b"str") => {
                // see http://officeopenxml.com/SScontentOverview.php
                // str - refers to formula cells
                // * <c .. t='v' .. > indicates calculated value (this case)
                // * <c .. t='f' .. > to the formula string (ignored case
                // TODO: Fully support a DataType::Formula representing both Formula string &
                // last calculated value?
                v.parse()
                    .map(DataType::Float)
                    .map_err(Error::from)
            }
            Some(b"n") => {
                // n - number
                v.parse().map(DataType::Float).map_err(Error::from)
            }
            None => {
                v.parse()
                    .map(DataType::Float)
                    .map_err(Error::from)
            }
            Some(b"is") => {
                // this case should be handled in outer loop over cell elements, in which
                // case read_inline_str is called instead. Case included here for completeness.
                return Err("called read_value on a cell of type inlineStr".into());
            }
            Some(t) => return Err(format!("unknown cell 't' attribute={:?}", t).into()),
        }
    }
    /// search through an Element's attributes for the named one
    fn get_attribute<'a>(e: &'a Element, n: &'a [u8]) -> Result<Option<&'a [u8]>> {
        for a in e.attributes() {
            match a {
                Ok((k, v)) if k == n => return Ok(Some(v)),
                Ok(_) => {} // ignore other attributes
                Err(qe) => return Err(qe.into()),
            }
        }

        Ok(None)
    }
    while let Some(res_event) = xml.next() {
        match res_event {
            Err(e) => return Err(e.into()),
            Ok(Event::Start(ref c_element)) if c_element.name() == b"c" => {
                let pos = get_attribute(c_element, b"r")
                    .and_then(|o| o.ok_or_else(|| "Cell missing 'r' attribute tag".into()))
                    .and_then(get_row_column)?;

                loop {
                    match xml.next() {
                        Some(Err(e)) => return Err(e.into()),
                        Some(Ok(Event::Start(ref e))) => {
                            match e.name() {
                                b"is" => {
                                    cells.push(Cell::new(pos, read_inline_str(xml)?));
                                    break;
                                }
                                b"v" => {
                                    cells.push(Cell::new(pos, read_value(xml, strings, c_element)?));
                                    break;
                                }
                                b"f" => {} // ignore f nodes
                                n => return Err(format!("not a v, f, or is node: {:?}", n).into()),
                            }
                        }
                        Some(Ok(Event::End(ref e))) if e.name() == b"c" => {
                            break;
                        }
                        None => return Err("End of xml".into()),
                        _ => (),
                    }
                }
            }
            Ok(Event::End(ref e)) if e.name() == b"sheetData" => return Ok(()),
            _ => (),
        }
    }
    Err("Could not find </sheetData>".into())
}

/// converts a text representation (e.g. "A6:G67") of a dimension into integers
/// - top left (row, column),
/// - bottom right (row, column)
fn get_dimension(dimension: &[u8]) -> Result<((u32, u32), (u32, u32))> {
    let parts: Vec<_> = dimension.split(|c| *c == b':')
        .map(|s| get_row_column(s))
        .collect::<Result<Vec<_>>>()?;

    match parts.len() {
        0 => Err("dimension cannot be empty".into()),
        1 => Ok((parts[0], parts[0])),
        2 => Ok((parts[0], parts[1])),
        len => Err(format!("range dimension has 0 or 1 ':', got {}", len).into()),
    }
}

/// converts a text range name into its position (row, column) (0 based index)
fn get_row_column(range: &[u8]) -> Result<(u32, u32)> {
    let (mut row, mut col) = (0, 0);
    let mut pow = 1;
    let mut readrow = true;
    for c in range.iter().rev() {
        match *c {
            c @ b'0'...b'9' => {
                if readrow {
                    row += ((c - b'0') as u32) * pow;
                    pow *= 10;
                } else {
                    return Err(format!("Numeric character are only allowed \
                        at the end of the range: {:x}",
                                       c)
                        .into());
                }
            }
            c @ b'A'...b'Z' => {
                if readrow {
                    pow = 1;
                    readrow = false;
                }
                col += ((c - b'A') as u32 + 1) * pow;
                pow *= 26;
            }
            c @ b'a'...b'z' => {
                if readrow {
                    pow = 1;
                    readrow = false;
                }
                col += ((c - b'a') as u32 + 1) * pow;
                pow *= 26;
            }
            _ => return Err(format!("Expecting alphanumeric character, got {:x}", c).into()),
        }
    }
    Ok((row - 1, col - 1))
}

#[test]
fn test_dimensions() {
    assert_eq!(get_row_column(b"A1").unwrap(), (0, 0));
    assert_eq!(get_row_column(b"C107").unwrap(), (106, 2));
    assert_eq!(get_dimension(b"C2:D35").unwrap(), ((1, 2), (34, 3)));
}
<commit_msg>Parse formula values as Float & fallback to String<commit_after>use std::fs::File;
use std::io::BufReader;
use std::collections::HashMap;
use std::borrow::Cow;

use zip::read::{ZipFile, ZipArchive};
use zip::result::ZipError;
use quick_xml::{XmlReader, Event, AsStr, Element};

use {DataType, ExcelReader, Range, Cell};
use vba::VbaProject;
use errors::*;

/// A struct representing xml zipped excel file
/// Xlsx, Xlsm, Xlam
pub struct Xlsx {
    zip: ZipArchive<File>,
}

impl Xlsx {
    fn xml_reader<'a>(&'a mut self,
                      path: &str)
                      -> Option<Result<XmlReader<BufReader<ZipFile<'a>>>>> {
        match self.zip.by_name(path) {
            Ok(f) => {
                Some(Ok(XmlReader::from_reader(BufReader::new(f))
                    .with_check(false)
                    .trim_text(false)))
            }
            Err(ZipError::FileNotFound) => None,
            Err(e) => return Some(Err(e.into())),
        }
    }
}

impl ExcelReader for Xlsx {
    fn new(f: File) -> Result<Self> {
        Ok(Xlsx { zip: ZipArchive::new(f)? })
    }

    fn has_vba(&mut self) -> bool {
        self.zip.by_name("xl/vbaProject.bin").is_ok()
    }

    fn vba_project(&mut self) -> Result<Cow<VbaProject>> {
        let mut f = self.zip.by_name("xl/vbaProject.bin")?;
        let len = f.size() as usize;
        VbaProject::new(&mut f, len).map(|v| Cow::Owned(v))
    }

    fn read_shared_strings(&mut self) -> Result<Vec<String>> {
        let mut xml = match self.xml_reader("xl/sharedStrings.xml") {
            None => return Ok(Vec::new()),
            Some(x) => x?,
        };
        let mut strings = Vec::new();
        let mut rich_buffer: Option<String> = None;
        while let Some(res_event) = xml.next() {
            match res_event {
                Ok(Event::Start(ref e)) if e.name() == b"r" => {
                    if let None = rich_buffer {
                        // use a buffer since richtext has multiples <r> and <t> for the same cell
                        rich_buffer = Some(String::new());
                    }
                }
                Ok(Event::End(ref e)) if e.name() == b"si" => {
                    if let Some(s) = rich_buffer {
                        strings.push(s);
                        rich_buffer = None;
                    }
                }
                Ok(Event::Start(ref e)) if e.name() == b"t" => {
                    let value = xml.read_text_unescaped(b"t")?;
                    if let Some(ref mut s) = rich_buffer {
                        s.push_str(&value);
                    } else {
                        strings.push(value);
                    }
                }
                Err(e) => return Err(e.into()),
                _ => (),
            }
        }
        Ok(strings)
    }

    fn read_sheets_names(&mut self,
                         relationships: &HashMap<Vec<u8>, String>)
                         -> Result<Vec<(String, String)>> {
        let xml = match self.xml_reader("xl/workbook.xml") {
            None => return Ok(Vec::new()),
            Some(x) => x?,
        };
        let mut sheets = Vec::new();
        for res_event in xml {
            match res_event {
                Ok(Event::Start(ref e)) if e.name() == b"sheet" => {
                    let mut name = String::new();
                    let mut path = String::new();
                    for a in e.unescaped_attributes() {
                        match a? {
                            (b"name", v) => name = v.as_str()?.to_string(),
                            (b"r:id", v) => {
                                let r = &relationships[&*v][..];
                                // target may have pre-prended "/xl/" or "xl/" path;
                                // strip if present
                                let r = if r.starts_with("/xl/") {
                                    &r[4..]
                                } else if r.starts_with("xl/") {
                                    &r[3..]
                                } else {
                                    r
                                };
                                path = format!("xl/{}", r);
                            }
                            _ => (),
                        }
                    }
                    sheets.push((name, path));
                }
                Err(e) => return Err(e.into()),
                _ => (),
            }
        }
        Ok(sheets)
    }

    fn read_relationships(&mut self) -> Result<HashMap<Vec<u8>, String>> {
        let xml = match self.xml_reader("xl/_rels/workbook.xml.rels") {
            None => return Err("Cannot find relationships file".into()),
            Some(x) => x?,
        };
        let mut relationships = HashMap::new();
        for res_event in xml {
            match res_event {
                Ok(Event::Start(ref e)) if e.name() == b"Relationship" => {
                    let mut id = Vec::new();
                    let mut target = String::new();
                    for a in e.attributes() {
                        match a? {
                            (b"Id", v) => id.extend_from_slice(v),
                            (b"Target", v) => target = v.as_str()?.to_string(),
                            _ => (),
                        }
                    }
                    relationships.insert(id, target);
                }
                Err(e) => return Err(e.into()),
                _ => (),
            }
        }
        Ok(relationships)
    }

    fn read_worksheet_range(&mut self, path: &str, strings: &[String]) -> Result<Range> {
        let mut xml = match self.xml_reader(path) {
            None => return Err(format!("Cannot find {} path", path).into()),
            Some(x) => x?,
        };
        let mut cells = Vec::new();
        'xml: while let Some(res_event) = xml.next() {
            match res_event {
                Err(e) => return Err(e.into()),
                Ok(Event::Start(ref e)) => {
                    match e.name() {
                        b"dimension" => {
                            for a in e.attributes() {
                                if let (b"ref", rdim) = a? {
                                    let (start, end) = get_dimension(rdim)?;
                                    cells.reserve(((end.0 - start.0 + 1) * (end.1 - start.1 + 1)) as
                                                 usize);
                                    continue 'xml;
                                }
                            }
                            return Err(format!("Expecting dimension, got {:?}", e).into());
                        }
                        b"sheetData" => read_sheet_data(&mut xml, strings, &mut cells)?,
                        _ => (),
                    }
                }
                _ => (),
            }
        }
        Ok(Range::from_sparse(cells))
    }
}

/// read sheetData node
fn read_sheet_data(xml: &mut XmlReader<BufReader<ZipFile>>,
                   strings: &[String],
                   cells: &mut Vec<Cell>)
                   -> Result<()> {
    /// read the contents of an <is> cell
    fn read_inline_str(xml: &mut XmlReader<BufReader<ZipFile>>) -> Result<DataType> {
        while let Some(is_evt) = xml.next() {
            match is_evt {
                Err(e) => return Err(e.into()),
                Ok(Event::Start(ref t)) if t.name() == b"t" => {
                    return Ok(DataType::String(xml.read_text_unescaped(b"t")?));
                }
                _ => {}
            }
        }
        Err("unable to read inlineStr".into())
    }
    /// read the contents of a <v> cell
    fn read_value(xml: &mut XmlReader<BufReader<ZipFile>>,
                  strings: &[String],
                  c_element: &Element)
                  -> Result<DataType> {
        let v = xml.read_text_unescaped(b"v")?;
        match get_attribute(c_element, b"t")? {
            Some(b"s") => {
                // shared string
                let idx: usize = v.parse()?;
                Ok(DataType::String(strings[idx].clone()))
            }
            Some(b"b") => {
                // boolean
                Ok(DataType::Bool(v != "0"))
            }
            Some(b"e") => {
                // error
                Ok(DataType::Error(v.parse()?))
            }
            Some(b"d") => {
                // date
                // TODO: create a DataType::Date
                // currently just return as string (ISO 8601)
                Ok(DataType::String(v))
            }
            Some(b"str") => {
                // see http://officeopenxml.com/SScontentOverview.php
                // str - refers to formula cells
                // * <c .. t='v' .. > indicates calculated value (this case)
                // * <c .. t='f' .. > to the formula string (ignored case
                // TODO: Fully support a DataType::Formula representing both Formula string &
                // last calculated value?
                //
                // NB: the result of a formula may not be a numeric value (=A3&" "&A4).
                // We do try an initial parse as Float for utility, but fall back to a string
                // representation if that fails
                v.parse()
                    .map(DataType::Float)
                    .map_err(Error::from)
                    .or_else::<Error, _>(|_| Ok(DataType::String(v)))
            }
            Some(b"n") => {
                // n - number
                v.parse().map(DataType::Float).map_err(Error::from)
            }
            None => {
                // If type is not known, we try to parse as Float for utility, but fall back to
                // String if this fails.
                v.parse()
                    .map(DataType::Float)
                    .map_err(Error::from)
                    .or_else::<Error, _>(|_| Ok(DataType::String(v)))
            }
            Some(b"is") => {
                // this case should be handled in outer loop over cell elements, in which
                // case read_inline_str is called instead. Case included here for completeness.
                return Err("called read_value on a cell of type inlineStr".into());
            }
            Some(t) => return Err(format!("unknown cell 't' attribute={:?}", t).into()),
        }
    }
    /// search through an Element's attributes for the named one
    fn get_attribute<'a>(e: &'a Element, n: &'a [u8]) -> Result<Option<&'a [u8]>> {
        for a in e.attributes() {
            match a {
                Ok((k, v)) if k == n => return Ok(Some(v)),
                Ok(_) => {} // ignore other attributes
                Err(qe) => return Err(qe.into()),
            }
        }

        Ok(None)
    }
    while let Some(res_event) = xml.next() {
        match res_event {
            Err(e) => return Err(e.into()),
            Ok(Event::Start(ref c_element)) if c_element.name() == b"c" => {
                let pos = get_attribute(c_element, b"r")
                    .and_then(|o| o.ok_or_else(|| "Cell missing 'r' attribute tag".into()))
                    .and_then(get_row_column)?;

                loop {
                    match xml.next() {
                        Some(Err(e)) => return Err(e.into()),
                        Some(Ok(Event::Start(ref e))) => {
                            match e.name() {
                                b"is" => {
                                    cells.push(Cell::new(pos, read_inline_str(xml)?));
                                    break;
                                }
                                b"v" => {
                                    cells.push(Cell::new(pos, read_value(xml, strings, c_element)?));
                                    break;
                                }
                                b"f" => {} // ignore f nodes
                                n => return Err(format!("not a v, f, or is node: {:?}", n).into()),
                            }
                        }
                        Some(Ok(Event::End(ref e))) if e.name() == b"c" => {
                            break;
                        }
                        None => return Err("End of xml".into()),
                        _ => (),
                    }
                }
            }
            Ok(Event::End(ref e)) if e.name() == b"sheetData" => return Ok(()),
            _ => (),
        }
    }
    Err("Could not find </sheetData>".into())
}

/// converts a text representation (e.g. "A6:G67") of a dimension into integers
/// - top left (row, column),
/// - bottom right (row, column)
fn get_dimension(dimension: &[u8]) -> Result<((u32, u32), (u32, u32))> {
    let parts: Vec<_> = dimension.split(|c| *c == b':')
        .map(|s| get_row_column(s))
        .collect::<Result<Vec<_>>>()?;

    match parts.len() {
        0 => Err("dimension cannot be empty".into()),
        1 => Ok((parts[0], parts[0])),
        2 => Ok((parts[0], parts[1])),
        len => Err(format!("range dimension has 0 or 1 ':', got {}", len).into()),
    }
}

/// converts a text range name into its position (row, column) (0 based index)
fn get_row_column(range: &[u8]) -> Result<(u32, u32)> {
    let (mut row, mut col) = (0, 0);
    let mut pow = 1;
    let mut readrow = true;
    for c in range.iter().rev() {
        match *c {
            c @ b'0'...b'9' => {
                if readrow {
                    row += ((c - b'0') as u32) * pow;
                    pow *= 10;
                } else {
                    return Err(format!("Numeric character are only allowed \
                        at the end of the range: {:x}",
                                       c)
                        .into());
                }
            }
            c @ b'A'...b'Z' => {
                if readrow {
                    pow = 1;
                    readrow = false;
                }
                col += ((c - b'A') as u32 + 1) * pow;
                pow *= 26;
            }
            c @ b'a'...b'z' => {
                if readrow {
                    pow = 1;
                    readrow = false;
                }
                col += ((c - b'a') as u32 + 1) * pow;
                pow *= 26;
            }
            _ => return Err(format!("Expecting alphanumeric character, got {:x}", c).into()),
        }
    }
    Ok((row - 1, col - 1))
}

#[test]
fn test_dimensions() {
    assert_eq!(get_row_column(b"A1").unwrap(), (0, 0));
    assert_eq!(get_row_column(b"C107").unwrap(), (106, 2));
    assert_eq!(get_dimension(b"C2:D35").unwrap(), ((1, 2), (34, 3)));
}
<|endoftext|>


<commit_before>use criterion::{criterion_group, criterion_main, Criterion};

fn fetch_benchmark(c: &mut Criterion) {
    // 0:  48 c7 c0 01 00 00 00    mov    rax,0x1
    let mut emu = lancelot::test::emu_from_shellcode64(&b"\x48\xC7\xC0\x01\x00\x00\x00"[..]);
    emu.reg.rip = 0x0;

    c.bench_function("fetch", |b| {
        b.iter(|| {
            emu.fetch().unwrap();
        })
    });
}

fn insn_benchmark(c: &mut Criterion) {
    c.bench_function("mov rax, 0x1", |b| {
        // 0:  48 c7 c0 01 00 00 00    mov    rax,0x1
        let mut emu = lancelot::test::emu_from_shellcode64(&b"\x48\xC7\xC0\x01\x00\x00\x00"[..]);
        b.iter(|| {
            emu.reg.rip = 0x0;
            emu.step().unwrap();
        })
    });

    c.bench_function("push 0x1", |b| {
        // 0:  6a 01                   push   0x1
        let mut emu = lancelot::test::emu_from_shellcode64(&b"\x6A\x01"[..]);
        b.iter(|| {
            emu.reg.rip = 0x0;
            emu.reg.rsp = 0x6000;
            emu.step().unwrap();
        })
    });
}

criterion_group!(fetch, fetch_benchmark);
criterion_group!(insn, insn_benchmark);
criterion_main!(fetch, insn);
<commit_msg>emu: bench: add some more aspects to benchmark<commit_after>use criterion::{criterion_group, criterion_main, Criterion};

fn fetch_benchmark(c: &mut Criterion) {
    // 0:  48 c7 c0 01 00 00 00    mov    rax,0x1
    let mut emu = lancelot::test::emu_from_shellcode64(&b"\x48\xC7\xC0\x01\x00\x00\x00"[..]);
    emu.reg.rip = 0x0;

    c.bench_function("fetch", |b| {
        b.iter(|| {
            emu.mem.read_u128(criterion::black_box(0x0)).unwrap();
        })
    });

    c.bench_function("fetch and decode", |b| {
        b.iter(|| {
            emu.fetch().unwrap();
        })
    });
}

fn insn_benchmark(c: &mut Criterion) {
    c.bench_function("mov rax, 0x1", |b| {
        // 0:  48 c7 c0 01 00 00 00    mov    rax,0x1
        let mut emu = lancelot::test::emu_from_shellcode64(&b"\x48\xC7\xC0\x01\x00\x00\x00"[..]);
        b.iter(|| {
            emu.reg.rip = 0x0;
            emu.step().unwrap();
        })
    });

    c.bench_function("push 0x1", |b| {
        // 0:  6a 01                   push   0x1
        let mut emu = lancelot::test::emu_from_shellcode64(&b"\x6A\x01"[..]);
        b.iter(|| {
            emu.reg.rip = 0x0;
            emu.reg.rsp = 0x6000;
            emu.step().unwrap();
        })
    });

    c.bench_function("sub rax, rbx", |b| {
        // 0:  48 29 d8                sub    rax,rbx
        let mut emu = lancelot::test::emu_from_shellcode64(&b"\x48\x29\xD8"[..]);
        b.iter(|| {
            emu.reg.rip = 0x0;
            emu.reg.rax = 0x1;
            emu.reg.rbx = 0x1;
            emu.step().unwrap();
        })
    });

    c.bench_function("add rax, rbx", |b| {
        // 0:  48 01 d8                add    rax,rbx
        let mut emu = lancelot::test::emu_from_shellcode64(&b"\x48\x01\xD8"[..]);
        b.iter(|| {
            emu.reg.rip = 0x0;
            emu.reg.rax = 0x1;
            emu.reg.rbx = 0x1;
            emu.step().unwrap();
        })
    });
}

criterion_group!(fetch, fetch_benchmark);
criterion_group!(insn, insn_benchmark);
criterion_main!(fetch, insn);
<|endoftext|>


<commit_before><commit_msg>Working on connect to multiple addresses<commit_after><|endoftext|>


<commit_before>/*******************************************************************************
 *
 * kit/kernel/kernel.rs
 * - main kernel entry point and top level management
 *
 * vim:ts=4:sw=4:et:tw=80:ft=rust
 *
 * Copyright (C) 2015, Devyn Cairns
 * Redistribution of this file is permitted under the terms of the simplified
 * BSD license. See LICENSE for more information.
 *
 * Based on OSDev Bare Bones tutorial
 * http://wiki.osdev.org/Bare_Bones
 *
 ******************************************************************************/

#![crate_type="lib"]
#![feature(core)]
#![feature(asm)]
#![feature(no_std)]
#![feature(lang_items)]
#![no_std]

#[macro_use]
extern crate core;

use core::prelude::*;
use core::fmt::Write;

use terminal::color;
use terminal::Terminal;

mod terminal;

#[no_mangle]
pub extern fn kernel_main() -> ! {

    terminal::initialize();
    terminal::set_color(color::RED, color::WHITE);

    terminal::write("+ Hello, I'm Kit.\n");

    panic!("2 + 2 = {}", 2 + 2);
}

#[lang = "stack_exhausted"]
extern fn stack_exhausted() {
}

#[lang = "eh_personality"]
extern fn eh_personality() {
}

#[lang = "panic_fmt"]
#[allow(unused_must_use)]
extern fn panic_fmt(fmt: core::fmt::Arguments,
                    file: &'static str,
                    line: usize) -> ! {

    terminal::set_color(color::WHITE, color::RED);

    write!(&mut Terminal, "\nKernel panic in {}:{}:\n  {}\n\n", file, line, fmt);

    unsafe {
        asm!("cli");

        loop {
            asm!("hlt");
        }
    }
}
<commit_msg>[kernel] more panic!() testing<commit_after>/*******************************************************************************
 *
 * kit/kernel/kernel.rs
 * - main kernel entry point and top level management
 *
 * vim:ts=4:sw=4:et:tw=80:ft=rust
 *
 * Copyright (C) 2015, Devyn Cairns
 * Redistribution of this file is permitted under the terms of the simplified
 * BSD license. See LICENSE for more information.
 *
 * Based on OSDev Bare Bones tutorial
 * http://wiki.osdev.org/Bare_Bones
 *
 ******************************************************************************/

#![crate_type="lib"]
#![feature(core)]
#![feature(asm)]
#![feature(no_std)]
#![feature(lang_items)]
#![no_std]

#[macro_use]
extern crate core;

use core::prelude::*;
use core::fmt::Write;

use terminal::color;
use terminal::Terminal;

mod terminal;

#[no_mangle]
pub extern fn kernel_main() -> ! {

    terminal::initialize();
    terminal::set_color(color::RED, color::WHITE);

    terminal::write("+ Hello, I'm Kit.\n");

    let result: Result<(), &str> = Err("foo");

    result.unwrap();

    unreachable!();
}

#[lang = "stack_exhausted"]
extern fn stack_exhausted() {
}

#[lang = "eh_personality"]
extern fn eh_personality() {
}

#[lang = "panic_fmt"]
#[allow(unused_must_use)]
extern fn panic_fmt(fmt: core::fmt::Arguments,
                    file: &'static str,
                    line: usize) -> ! {

    terminal::set_color(color::WHITE, color::RED);

    write!(&mut Terminal, "\nKernel panic in {}:{}:\n  {}\n\n", file, line, fmt);

    unsafe {
        asm!("cli");

        loop {
            asm!("hlt");
        }
    }
}
<|endoftext|>


<commit_before>use failure::{Error, Fail};
use strum_macros::Display;
use log::{info};
use bitflags::{bitflags};

use super::arch::{Arch, RVA, VA};
use super::pagemap::{PageMap};
use super::loaders::pe::PELoader;
use super::loaders::sc::ShellcodeLoader;
use super::analysis::{Analyzer};

#[derive(Debug, Fail)]
pub enum LoaderError {
    #[fail(display = "The given buffer is not supported (arch/plat/file format)")]
    NotSupported,
    #[fail(display = "The given buffer uses a bitness incompatible with the architecture")]
    MismatchedBitness,
}

#[derive(Display, Clone, Copy)]
pub enum FileFormat {
    Raw, // shellcode
    PE,
}

#[derive(Display, Clone, Copy)]
pub enum Platform {
    Windows,
}

bitflags! {
    pub struct Permissions: u8 {
        const R = 0b00000001;
        const W = 0b00000010;
        const X = 0b00000100;
        const RW = Self::R.bits | Self::W.bits;
        const RX =  Self::R.bits | Self::X.bits;
        const WX =  Self::W.bits | Self::X.bits;
        const RWX =  Self::R.bits | Self::W.bits | Self::X.bits;
    }
}

#[derive(Debug)]
pub struct Section {
    pub addr: RVA,
    pub size: u32,
    pub perms: Permissions,
    pub name: String,
}

impl Section {
    pub fn contains(self: &Section, rva: RVA) -> bool {
        if rva < self.addr {
            return false;
        }

        let max = self.addr + self.size as i32;  // danger
        if rva >= max {
            return false;
        }

        true
    }

    pub fn is_executable(&self) -> bool {
        self.perms.intersects(Permissions::X)
    }

    pub fn end(&self) -> RVA {
        self.addr + RVA::from(self.size as i64)
    }
}

pub struct LoadedModule {
    pub base_address: VA,
    pub sections: Vec<Section>,
    pub address_space: PageMap<u8>,
}

impl LoadedModule {
    pub fn max_address(&self) -> RVA {
        self.sections
            .iter()
            .map(|sec| sec.end())
            .max()
            .unwrap()  // danger: assume there's at least one section.
    }
}

pub trait Loader {
    /// Fetch the number of bits for a pointer in this architecture.
    fn get_arch(&self) -> Arch;
    fn get_plat(&self) -> Platform;
    fn get_file_format(&self) -> FileFormat;

    fn get_name(&self) -> String {
        return format!(
            "{}/{}/{}",
            self.get_plat(),
            self.get_arch(),
            self.get_file_format()
        );
    }

    /// Returns True if this Loader knows how to load the given bytes.
    fn taste(&self, buf: &[u8]) -> bool;

    /// Load the given bytes into a Module and suggest the appropriate Analyzers.
    ///
    /// While the loader is parsing a file, it should determine what
    ///  the most appropriate analyzers are, e.g. a PE loader may inspect the headers
    ///  to determine if there is Control Flow Guard metadata that can be analyzed.
    fn load(&self, buf: &[u8]) -> Result<(LoadedModule, Vec<Box<dyn Analyzer>>), Error>;
}

pub fn default_loaders() -> Vec<Box<dyn Loader>> {
    // we might like these to come from a lazy_static global,
    //  however, then these have to be Sync.
    // I'm not sure if that's a good idea yet.
    let mut loaders: Vec<Box<dyn Loader>> = vec![];
    // the order here matters!
    // the default `load` routine will pick the first matching loader,
    //  so the earlier entries here have higher precedence.

    loaders.push(Box::new(PELoader::new(Arch::X32)));
    loaders.push(Box::new(PELoader::new(Arch::X64)));
    loaders.push(Box::new(ShellcodeLoader::new(Platform::Windows, Arch::X32)));
    loaders.push(Box::new(ShellcodeLoader::new(Platform::Windows, Arch::X64)));

    loaders
}

/// Find the loaders that support loading the given sample.
///
/// The result is an iterator so that the caller can use the first
///  matching result without waiting for all Loaders to taste the bytes.
///
/// Loaders are tasted in the order defined in `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::arch::*;
/// use lancelot::loader::*;
///
/// match taste(b"\xEB\xFE").nth(0) {
///   Some(loader) => assert_eq!(loader.get_name(), "Windows/x32/Raw"),
///   None => panic!("no matching loaders"),
/// };
/// ```
pub fn taste(buf: &[u8]) -> impl Iterator<Item = Box<dyn Loader>> {
    default_loaders()
        .into_iter()
        .filter(move |loader| loader.taste(buf))
}

/// Load the given sample using the first matching loader from `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::arch::*;
/// use lancelot::loader::*;
///
/// load(b"\xEB\xFE")
///   .map(|(loader, module, analyzers)| {
///     assert_eq!(loader.get_name(),       "Windows/x32/Raw");
///     assert_eq!(module.base_address,     VA(0x0));
///     assert_eq!(module.sections[0].name, "raw");
///   })
///   .map_err(|e| panic!(e));
/// ```
pub fn load(buf: &[u8]) -> Result<(Box<dyn Loader>, LoadedModule, Vec<Box<dyn Analyzer>>), Error> {
    match taste(buf).nth(0) {
        Some(loader) => {
            info!("auto-detected loader: {}", loader.get_name());
            loader.load(buf).map(|(module, analyzers)| (loader, module, analyzers))
        },
        None => Err(LoaderError::NotSupported.into()),
    }
}
<commit_msg>loader: fix new compile error on 1.40<commit_after>use failure::{Error, Fail};
use strum_macros::Display;
use log::{info};
use bitflags::{bitflags};

use super::arch::{Arch, RVA, VA};
use super::pagemap::{PageMap};
use super::loaders::pe::PELoader;
use super::loaders::sc::ShellcodeLoader;
use super::analysis::{Analyzer};

#[derive(Debug, Fail)]
pub enum LoaderError {
    #[fail(display = "The given buffer is not supported (arch/plat/file format)")]
    NotSupported,
    #[fail(display = "The given buffer uses a bitness incompatible with the architecture")]
    MismatchedBitness,
}

#[derive(Display, Clone, Copy)]
pub enum FileFormat {
    Raw, // shellcode
    PE,
}

#[derive(Display, Clone, Copy)]
pub enum Platform {
    Windows,
}

bitflags! {
    pub struct Permissions: u8 {
        const R = 0b00000001;
        const W = 0b00000010;
        const X = 0b00000100;
        const RW = Self::R.bits | Self::W.bits;
        const RX =  Self::R.bits | Self::X.bits;
        const WX =  Self::W.bits | Self::X.bits;
        const RWX =  Self::R.bits | Self::W.bits | Self::X.bits;
    }
}

#[derive(Debug)]
pub struct Section {
    pub addr: RVA,
    pub size: u32,
    pub perms: Permissions,
    pub name: String,
}

impl Section {
    pub fn contains(self: &Section, rva: RVA) -> bool {
        if rva < self.addr {
            return false;
        }

        let max = self.addr + self.size as i32;  // danger
        if rva >= max {
            return false;
        }

        true
    }

    pub fn is_executable(&self) -> bool {
        self.perms.intersects(Permissions::X)
    }

    pub fn end(&self) -> RVA {
        self.addr + RVA::from(self.size as i64)
    }
}

pub struct LoadedModule {
    pub base_address: VA,
    pub sections: Vec<Section>,
    pub address_space: PageMap<u8>,
}

impl LoadedModule {
    pub fn max_address(&self) -> RVA {
        self.sections
            .iter()
            .map(|sec| sec.end())
            .max()
            .unwrap()  // danger: assume there's at least one section.
    }
}

pub trait Loader {
    /// Fetch the number of bits for a pointer in this architecture.
    fn get_arch(&self) -> Arch;
    fn get_plat(&self) -> Platform;
    fn get_file_format(&self) -> FileFormat;

    fn get_name(&self) -> String {
        return format!(
            "{}/{}/{}",
            self.get_plat(),
            self.get_arch(),
            self.get_file_format()
        );
    }

    /// Returns True if this Loader knows how to load the given bytes.
    fn taste(&self, buf: &[u8]) -> bool;

    /// Load the given bytes into a Module and suggest the appropriate Analyzers.
    ///
    /// While the loader is parsing a file, it should determine what
    ///  the most appropriate analyzers are, e.g. a PE loader may inspect the headers
    ///  to determine if there is Control Flow Guard metadata that can be analyzed.
    fn load(&self, buf: &[u8]) -> Result<(LoadedModule, Vec<Box<dyn Analyzer>>), Error>;
}

pub fn default_loaders() -> Vec<Box<dyn Loader>> {
    // we might like these to come from a lazy_static global,
    //  however, then these have to be Sync.
    // I'm not sure if that's a good idea yet.
    let mut loaders: Vec<Box<dyn Loader>> = vec![];
    // the order here matters!
    // the default `load` routine will pick the first matching loader,
    //  so the earlier entries here have higher precedence.

    loaders.push(Box::new(PELoader::new(Arch::X32)));
    loaders.push(Box::new(PELoader::new(Arch::X64)));
    loaders.push(Box::new(ShellcodeLoader::new(Platform::Windows, Arch::X32)));
    loaders.push(Box::new(ShellcodeLoader::new(Platform::Windows, Arch::X64)));

    loaders
}

/// Find the loaders that support loading the given sample.
///
/// The result is an iterator so that the caller can use the first
///  matching result without waiting for all Loaders to taste the bytes.
///
/// Loaders are tasted in the order defined in `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::arch::*;
/// use lancelot::loader::*;
///
/// match taste(b"\xEB\xFE").nth(0) {
///   Some(loader) => assert_eq!(loader.get_name(), "Windows/x32/Raw"),
///   None => panic!("no matching loaders"),
/// };
/// ```
pub fn taste(buf: &[u8]) -> impl Iterator<Item = Box<dyn Loader>> + '_ {
    default_loaders()
        .into_iter()
        .filter(move |loader| loader.taste(buf))
}

/// Load the given sample using the first matching loader from `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::arch::*;
/// use lancelot::loader::*;
///
/// load(b"\xEB\xFE")
///   .map(|(loader, module, analyzers)| {
///     assert_eq!(loader.get_name(),       "Windows/x32/Raw");
///     assert_eq!(module.base_address,     VA(0x0));
///     assert_eq!(module.sections[0].name, "raw");
///   })
///   .map_err(|e| panic!(e));
/// ```
pub fn load(buf: &[u8]) -> Result<(Box<dyn Loader>, LoadedModule, Vec<Box<dyn Analyzer>>), Error> {
    match taste(buf).nth(0) {
        Some(loader) => {
            info!("auto-detected loader: {}", loader.get_name());
            loader.load(buf).map(|(module, analyzers)| (loader, module, analyzers))
        },
        None => Err(LoaderError::NotSupported.into()),
    }
}
<|endoftext|>


<commit_before>use crate::transform::MirPass;
use crate::util::patch::MirPatch;
use rustc_index::bit_set::BitSet;
use rustc_middle::mir::*;
use rustc_middle::ty::TyCtxt;
use rustc_target::spec::PanicStrategy;

/// A pass that removes noop landing pads and replaces jumps to them with
/// `None`. This is important because otherwise LLVM generates terrible
/// code for these.
pub struct RemoveNoopLandingPads;

pub fn remove_noop_landing_pads<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
    if tcx.sess.panic_strategy() == PanicStrategy::Abort {
        return;
    }
    debug!("remove_noop_landing_pads({:?})", body);

    RemoveNoopLandingPads.remove_nop_landing_pads(body)
}

impl<'tcx> MirPass<'tcx> for RemoveNoopLandingPads {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        remove_noop_landing_pads(tcx, body);
    }
}

impl RemoveNoopLandingPads {
    fn is_nop_landing_pad(
        &self,
        bb: BasicBlock,
        body: &Body<'_>,
        nop_landing_pads: &BitSet<BasicBlock>,
    ) -> bool {
        for stmt in &body[bb].statements {
            match &stmt.kind {
                StatementKind::FakeRead(..)
                | StatementKind::StorageLive(_)
                | StatementKind::StorageDead(_)
                | StatementKind::AscribeUserType(..)
                | StatementKind::Coverage(..)
                | StatementKind::Nop => {
                    // These are all nops in a landing pad
                }

                StatementKind::Assign(box (place, Rvalue::Use(_))) => {
                    if place.as_local().is_some() {
                        // Writing to a local (e.g., a drop flag) does not
                        // turn a landing pad to a non-nop
                    } else {
                        return false;
                    }
                }

                StatementKind::Assign { .. }
                | StatementKind::SetDiscriminant { .. }
                | StatementKind::LlvmInlineAsm { .. }
                | StatementKind::Retag { .. } => {
                    return false;
                }
            }
        }

        let terminator = body[bb].terminator();
        match terminator.kind {
            TerminatorKind::Goto { .. }
            | TerminatorKind::Resume
            | TerminatorKind::SwitchInt { .. }
            | TerminatorKind::FalseEdge { .. }
            | TerminatorKind::FalseUnwind { .. } => {
                terminator.successors().all(|&succ| nop_landing_pads.contains(succ))
            }
            TerminatorKind::GeneratorDrop
            | TerminatorKind::Yield { .. }
            | TerminatorKind::Return
            | TerminatorKind::Abort
            | TerminatorKind::Unreachable
            | TerminatorKind::Call { .. }
            | TerminatorKind::Assert { .. }
            | TerminatorKind::DropAndReplace { .. }
            | TerminatorKind::Drop { .. }
            | TerminatorKind::InlineAsm { .. } => false,
        }
    }

    fn remove_nop_landing_pads(&self, body: &mut Body<'_>) {
        // make sure there's a single resume block
        let resume_block = {
            let patch = MirPatch::new(body);
            let resume_block = patch.resume_block();
            patch.apply(body);
            resume_block
        };
        debug!("remove_noop_landing_pads: resume block is {:?}", resume_block);

        let mut jumps_folded = 0;
        let mut landing_pads_removed = 0;
        let mut nop_landing_pads = BitSet::new_empty(body.basic_blocks().len());

        // This is a post-order traversal, so that if A post-dominates B
        // then A will be visited before B.
        let postorder: Vec<_> = traversal::postorder(body).map(|(bb, _)| bb).collect();
        for bb in postorder {
            debug!("  processing {:?}", bb);
            if let Some(unwind) = body[bb].terminator_mut().unwind_mut() {
                if let Some(unwind_bb) = *unwind {
                    if nop_landing_pads.contains(unwind_bb) {
                        debug!("    removing noop landing pad");
                        landing_pads_removed += 1;
                        *unwind = None;
                    }
                }
            }

            for target in body[bb].terminator_mut().successors_mut() {
                if *target != resume_block && nop_landing_pads.contains(*target) {
                    debug!("    folding noop jump to {:?} to resume block", target);
                    *target = resume_block;
                    jumps_folded += 1;
                }
            }

            let is_nop_landing_pad = self.is_nop_landing_pad(bb, body, &nop_landing_pads);
            if is_nop_landing_pad {
                nop_landing_pads.insert(bb);
            }
            debug!("    is_nop_landing_pad({:?}) = {}", bb, is_nop_landing_pad);
        }

        debug!("removed {:?} jumps and {:?} landing pads", jumps_folded, landing_pads_removed);
    }
}
<commit_msg>Recognize discriminant reads as no-ops in RemoveNoopLandingPads<commit_after>use crate::transform::MirPass;
use crate::util::patch::MirPatch;
use rustc_index::bit_set::BitSet;
use rustc_middle::mir::*;
use rustc_middle::ty::TyCtxt;
use rustc_target::spec::PanicStrategy;

/// A pass that removes noop landing pads and replaces jumps to them with
/// `None`. This is important because otherwise LLVM generates terrible
/// code for these.
pub struct RemoveNoopLandingPads;

pub fn remove_noop_landing_pads<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
    if tcx.sess.panic_strategy() == PanicStrategy::Abort {
        return;
    }
    debug!("remove_noop_landing_pads({:?})", body);

    RemoveNoopLandingPads.remove_nop_landing_pads(body)
}

impl<'tcx> MirPass<'tcx> for RemoveNoopLandingPads {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        remove_noop_landing_pads(tcx, body);
    }
}

impl RemoveNoopLandingPads {
    fn is_nop_landing_pad(
        &self,
        bb: BasicBlock,
        body: &Body<'_>,
        nop_landing_pads: &BitSet<BasicBlock>,
    ) -> bool {
        for stmt in &body[bb].statements {
            match &stmt.kind {
                StatementKind::FakeRead(..)
                | StatementKind::StorageLive(_)
                | StatementKind::StorageDead(_)
                | StatementKind::AscribeUserType(..)
                | StatementKind::Coverage(..)
                | StatementKind::Nop => {
                    // These are all nops in a landing pad
                }

                StatementKind::Assign(box (place, Rvalue::Use(_) | Rvalue::Discriminant(_))) => {
                    if place.as_local().is_some() {
                        // Writing to a local (e.g., a drop flag) does not
                        // turn a landing pad to a non-nop
                    } else {
                        return false;
                    }
                }

                StatementKind::Assign { .. }
                | StatementKind::SetDiscriminant { .. }
                | StatementKind::LlvmInlineAsm { .. }
                | StatementKind::Retag { .. } => {
                    return false;
                }
            }
        }

        let terminator = body[bb].terminator();
        match terminator.kind {
            TerminatorKind::Goto { .. }
            | TerminatorKind::Resume
            | TerminatorKind::SwitchInt { .. }
            | TerminatorKind::FalseEdge { .. }
            | TerminatorKind::FalseUnwind { .. } => {
                terminator.successors().all(|&succ| nop_landing_pads.contains(succ))
            }
            TerminatorKind::GeneratorDrop
            | TerminatorKind::Yield { .. }
            | TerminatorKind::Return
            | TerminatorKind::Abort
            | TerminatorKind::Unreachable
            | TerminatorKind::Call { .. }
            | TerminatorKind::Assert { .. }
            | TerminatorKind::DropAndReplace { .. }
            | TerminatorKind::Drop { .. }
            | TerminatorKind::InlineAsm { .. } => false,
        }
    }

    fn remove_nop_landing_pads(&self, body: &mut Body<'_>) {
        // make sure there's a single resume block
        let resume_block = {
            let patch = MirPatch::new(body);
            let resume_block = patch.resume_block();
            patch.apply(body);
            resume_block
        };
        debug!("remove_noop_landing_pads: resume block is {:?}", resume_block);

        let mut jumps_folded = 0;
        let mut landing_pads_removed = 0;
        let mut nop_landing_pads = BitSet::new_empty(body.basic_blocks().len());

        // This is a post-order traversal, so that if A post-dominates B
        // then A will be visited before B.
        let postorder: Vec<_> = traversal::postorder(body).map(|(bb, _)| bb).collect();
        for bb in postorder {
            debug!("  processing {:?}", bb);
            if let Some(unwind) = body[bb].terminator_mut().unwind_mut() {
                if let Some(unwind_bb) = *unwind {
                    if nop_landing_pads.contains(unwind_bb) {
                        debug!("    removing noop landing pad");
                        landing_pads_removed += 1;
                        *unwind = None;
                    }
                }
            }

            for target in body[bb].terminator_mut().successors_mut() {
                if *target != resume_block && nop_landing_pads.contains(*target) {
                    debug!("    folding noop jump to {:?} to resume block", target);
                    *target = resume_block;
                    jumps_folded += 1;
                }
            }

            let is_nop_landing_pad = self.is_nop_landing_pad(bb, body, &nop_landing_pads);
            if is_nop_landing_pad {
                nop_landing_pads.insert(bb);
            }
            debug!("    is_nop_landing_pad({:?}) = {}", bb, is_nop_landing_pad);
        }

        debug!("removed {:?} jumps and {:?} landing pads", jumps_folded, landing_pads_removed);
    }
}
<|endoftext|>


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

#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]
#![cfg_attr(not(feature="clippy"), allow(unknown_lints))]

extern crate env_logger;

#[cfg(any(feature = "amqp0-build-primitives", feature = "amqp0-pregen-primitives"))]
extern crate amqp0_codegen as codegen;
#[cfg(any(feature = "amqp0-build-primitives", feature = "amqp0-pregen-primitives"))]
extern crate amqp0_specs as specs;

fn main() {
    amqp0::build();
}

#[cfg(not(any(feature = "amqp0-build-primitives", feature = "amqp0-pregen-primitives")))]
mod amqp0 {
    pub fn build() {
        println!("Skipping build (neither amqp0-build-primitives nor amqp0-pregen-primitives specified)");
    }
}

#[cfg(any(feature = "amqp0-build-primitives", feature = "amqp0-pregen-primitives"))]
mod amqp0 {
    use std::env;
    use std::fs;
    use std::path::{Path, PathBuf};

    use env_logger;

    use codegen::{self, Builder, CodeWriter};
    use codegen::primalgen::{SpecsModuleWriter, SpecModuleWriter};
    use specs::specs as amqp0_specs;

    struct PrimitivesBuilder;

    const BUILDER_CRATES: &'static [&'static str] = &["amqp0-codegen"];
    const BUILDER_REBUILD: &'static [&'static str] = &["amqp0-build-primitives"];
    const BUILDER_PREGEN: &'static [&'static str] = &["amqp0-build-primitives"];

    impl Builder for PrimitivesBuilder {
        fn name(&self) -> &str { "amqp0-primitives" }
        fn crates(&self) -> &[&str] { BUILDER_CRATES }
        fn rebuild_features(&self) -> &[&str] { BUILDER_REBUILD }
        fn pregeneration_features(&self) -> &[&str] { BUILDER_PREGEN }
    }

    pub fn build() {
        env_logger::init().unwrap();
        println!("Building primitives from amqpspec");

        let out_path = env::var_os("OUT_DIR").map(PathBuf::from).expect("Error: OUT_DIR not set");
        let specs = amqp0_specs();
        let mut paths: Vec<PathBuf> = Vec::with_capacity(1 + specs.len());

        // mod.rs
        let specs_writer = SpecsModuleWriter::from_spec_slice(&specs[..]);
        let writer = CodeWriter::new(PrimitivesBuilder, specs_writer);
        println!("Generated mod.rs");
        let mod_path = out_path.join("mod.rs");
        writer.write_to_path(&mod_path).expect("Failed to write amqp0.rs");
        paths.push(mod_path);

        // {name}{minor}_{revision}.rs
        for spec in specs {
            let spec_writer = SpecModuleWriter::from_spec(spec);
            let filename = format!("{}.rs", spec_writer.mod_name());
            let path = out_path.join(&filename);
            let writer = CodeWriter::new(PrimitivesBuilder, spec_writer);
            println!("Generating {}", filename);
            writer.write_to_path(&path).expect("Failed to write spec module");
            paths.push(path);
        }

        if cfg!(feature = "rustfmt") {
            codegen::format_files(paths.clone());
        }

        if cfg!(feature = "amqp0-pregen-primitives") {
            let pregen_dir = Path::new("pregen");
            for path in paths {
                let suffix = path.strip_prefix(&out_path).unwrap();
                let dst = pregen_dir.join(suffix);
                println!("Saving to {} to {}", path.display(), dst.display());
                fs::copy(&path, dst).unwrap();
            }
        }
    }
}<commit_msg>[primitives] Make build script conditionals more simple.<commit_after>// Copyright 2016 Alexander Reece
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]
#![cfg_attr(not(feature="clippy"), allow(unknown_lints))]

extern crate env_logger;

#[cfg(feature = "amqp0-codegen")]
extern crate amqp0_codegen as codegen;
#[cfg(feature = "amqp0-specs")]
extern crate amqp0_specs as specs;

fn main() {
    amqp0::build();
}

#[cfg(not(feature = "amqp0-specs"))]
mod amqp0 {
    pub fn build() {
        println!("Skipping build (neither amqp0-build-primitives nor amqp0-pregen-primitives specified)");
    }
}

#[cfg(feature = "amqp0-specs")]
mod amqp0 {
    use std::env;
    use std::fs;
    use std::path::{Path, PathBuf};

    use env_logger;

    use codegen::{self, Builder, CodeWriter};
    use codegen::primalgen::{SpecsModuleWriter, SpecModuleWriter};
    use specs::specs as amqp0_specs;

    struct PrimitivesBuilder;

    const BUILDER_CRATES: &'static [&'static str] = &["amqp0-codegen"];
    const BUILDER_REBUILD: &'static [&'static str] = &["amqp0-build-primitives"];
    const BUILDER_PREGEN: &'static [&'static str] = &["amqp0-build-primitives"];

    impl Builder for PrimitivesBuilder {
        fn name(&self) -> &str { "amqp0-primitives" }
        fn crates(&self) -> &[&str] { BUILDER_CRATES }
        fn rebuild_features(&self) -> &[&str] { BUILDER_REBUILD }
        fn pregeneration_features(&self) -> &[&str] { BUILDER_PREGEN }
    }

    pub fn build() {
        env_logger::init().unwrap();
        println!("Building primitives from amqpspec");

        let out_path = env::var_os("OUT_DIR").map(PathBuf::from).expect("Error: OUT_DIR not set");
        let specs = amqp0_specs();
        let mut paths: Vec<PathBuf> = Vec::with_capacity(1 + specs.len());

        // mod.rs
        let specs_writer = SpecsModuleWriter::from_spec_slice(&specs[..]);
        let writer = CodeWriter::new(PrimitivesBuilder, specs_writer);
        println!("Generated mod.rs");
        let mod_path = out_path.join("mod.rs");
        writer.write_to_path(&mod_path).expect("Failed to write amqp0.rs");
        paths.push(mod_path);

        // {name}{minor}_{revision}.rs
        for spec in specs {
            let spec_writer = SpecModuleWriter::from_spec(spec);
            let filename = format!("{}.rs", spec_writer.mod_name());
            let path = out_path.join(&filename);
            let writer = CodeWriter::new(PrimitivesBuilder, spec_writer);
            println!("Generating {}", filename);
            writer.write_to_path(&path).expect("Failed to write spec module");
            paths.push(path);
        }

        if cfg!(feature = "rustfmt") {
            codegen::format_files(paths.clone());
        }

        if cfg!(feature = "amqp0-pregen-primitives") {
            let pregen_dir = Path::new("pregen");
            for path in paths {
                let suffix = path.strip_prefix(&out_path).unwrap();
                let dst = pregen_dir.join(suffix);
                println!("Saving to {} to {}", path.display(), dst.display());
                fs::copy(&path, dst).unwrap();
            }
        }
    }
}<|endoftext|>


<commit_before>#![feature(alloc_system)]

extern crate alloc_system;

extern crate libloading;
extern crate capnp;

pub mod agent;

pub mod scheduler;

pub mod ports;
pub mod result;
<commit_msg>Question_mark<commit_after>#![feature(question_mark)]
#![feature(alloc_system)]

extern crate alloc_system;

extern crate libloading;
extern crate capnp;

pub mod agent;

pub mod scheduler;

pub mod ports;
pub mod result;
<|endoftext|>


<commit_before>//! Tools for working with OSC. [**sender()**](./fn.sender.html) creates an OSC sender,
//! [**receiver(port)**](./fn.receiver.html) creates an OSC receiver.

pub use rosc;

// Re-export rosc items.
//
// Remove `Osc` prefix as items are already namespaced via a module, e.g. `OscMessage` becomes
// `nannou_osc::Message`.
pub use self::recv::Receiver;
#[doc(inline)]
pub use self::rosc::{
    decoder, encoder, OscBundle as Bundle, OscColor as Color, OscError as Error,
    OscMessage as Message, OscMidiMessage as MidiMessage, OscType as Type,
};
pub use self::send::Sender;

use std;
use std::net::{Ipv4Addr, SocketAddr};

pub mod recv;
pub mod send;

/// Indicates that a `Sender` is not currently connected to a target address, and that the target
/// address will have to be supplied manually when sending packets.
pub struct Unconnected;

/// Indicates that a `Sender` is currently connected to a specific target address.
pub struct Connected {
    addr: SocketAddr,
}

/// An *OSC packet* can contain an OSC message or a bundle of nested packets which is called an
/// *OSC bundle*.
#[derive(Clone, Debug, PartialEq)]
pub enum Packet {
    Message(Message),
    Bundle(Bundle),
}

impl From<Message> for Packet {
    fn from(msg: Message) -> Self {
        Packet::Message(msg)
    }
}

impl From<Bundle> for Packet {
    fn from(bundle: Bundle) -> Self {
        Packet::Bundle(bundle)
    }
}

impl From<rosc::OscPacket> for Packet {
    fn from(packet: rosc::OscPacket) -> Self {
        match packet {
            rosc::OscPacket::Message(msg) => msg.into(),
            rosc::OscPacket::Bundle(bundle) => bundle.into(),
        }
    }
}

impl<A> From<(A, Vec<Type>)> for Packet
where
    A: Into<String>,
{
    fn from((addr, args): (A, Vec<Type>)) -> Self {
        msg(addr, args).into()
    }
}

impl Into<rosc::OscPacket> for Packet {
    fn into(self) -> rosc::OscPacket {
        match self {
            Packet::Message(msg) => rosc::OscPacket::Message(msg),
            Packet::Bundle(bundle) => rosc::OscPacket::Bundle(bundle),
        }
    }
}

impl Packet {
    /// A recursive function that unfolds the packet into the end of the given buffer of messages.
    pub fn unfold(self, msgs: &mut Vec<Message>) {
        match self {
            Packet::Message(msg) => msgs.push(msg),
            Packet::Bundle(bundle) => {
                for packet in bundle.content {
                    Packet::unfold(packet.into(), msgs);
                }
            }
        }
    }

    /// Convert the `Packet` into a Vec containing all `Message`s contained within.
    ///
    /// This uses the `unfold` method internally.
    pub fn into_msgs(self) -> Vec<Message> {
        let mut msgs = vec![];
        self.unfold(&mut msgs);
        msgs
    }
}

/// The default local IP address.
pub fn default_ipv4_addr() -> Ipv4Addr {
    Ipv4Addr::new(0, 0, 0, 0)
}

/// A simple wrapper around the most commonly used `Receiver` constructor.
pub fn receiver(port: u16) -> Result<Receiver, std::io::Error> {
    Receiver::bind(port)
}

/// A simple wrapper around the most commonly used `Sender` constructor.
pub fn sender() -> Result<Sender, std::io::Error> {
    Sender::bind()
}

/// A simplified constructor for an OSC `Message`.
pub fn msg<A>(addr: A, args: Vec<Type>) -> Message
where
    A: Into<String>,
{
    let addr = addr.into();
    let args = Some(args);
    Message { addr, args }
}

/// Decodes the given slice of `bytes` into a `Packet`.
///
/// Returns an `Error` if the slice does not contain a valid OSC packet.
pub fn decode(bytes: &[u8]) -> Result<Packet, Error> {
    rosc::decoder::decode(bytes).map(|p| p.into())
}

/// Encodes the given `Packet` into a `Vec` of bytes.
///
/// Returns an `Error` if the packet is invalid.
pub fn encode(packet: Packet) -> Result<Vec<u8>, Error> {
    let rosc_packet = packet.into();
    rosc::encoder::encode(&rosc_packet)
}

/// Errors that might occur whilst attempting to send or receive an OSC packet.
#[derive(Debug)]
pub enum CommunicationError {
    Io(std::io::Error),
    Osc(Error),
    Poisoned,
}

impl From<std::io::Error> for CommunicationError {
    fn from(err: std::io::Error) -> Self {
        CommunicationError::Io(err)
    }
}

impl From<Error> for CommunicationError {
    fn from(err: Error) -> Self {
        CommunicationError::Osc(err)
    }
}

impl<T> From<std::sync::PoisonError<T>> for CommunicationError {
    fn from(_: std::sync::PoisonError<T>) -> Self {
        CommunicationError::Poisoned
    }
}

impl std::error::Error for CommunicationError {
    fn description(&self) -> &str {
        match *self {
            CommunicationError::Io(ref err) => std::error::Error::description(err),
            // TODO: Error isn't implemented for OscError - should fix this upstream.
            CommunicationError::Osc(ref _err) => "Failed to decode the OSC packet",
            // CommunicationError::Osc(ref err) => std::error::Error::description(err),
            CommunicationError::Poisoned => "The inner buffer's mutex was poisoned",
        }
    }
    fn cause(&self) -> Option<&std::error::Error> {
        match *self {
            CommunicationError::Io(ref err) => Some(err),
            // TODO: Error isn't implemented for OscError - should fix this upstream.
            CommunicationError::Osc(ref _err) => None,
            // CommunicationError::Osc(ref err) => Some(err),
            _ => None,
        }
    }
}

impl std::fmt::Display for CommunicationError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", std::error::Error::description(self))
    }
}
<commit_msg>Address warnings in nannou_osc<commit_after>//! Tools for working with OSC. [**sender()**](./fn.sender.html) creates an OSC sender,
//! [**receiver(port)**](./fn.receiver.html) creates an OSC receiver.

pub use rosc;

// Re-export rosc items.
//
// Remove `Osc` prefix as items are already namespaced via a module, e.g. `OscMessage` becomes
// `nannou_osc::Message`.
pub use self::recv::Receiver;
#[doc(inline)]
pub use self::rosc::{
    decoder, encoder, OscBundle as Bundle, OscColor as Color, OscError as Error,
    OscMessage as Message, OscMidiMessage as MidiMessage, OscType as Type,
};
pub use self::send::Sender;

use std;
use std::net::{Ipv4Addr, SocketAddr};

pub mod recv;
pub mod send;

/// Indicates that a `Sender` is not currently connected to a target address, and that the target
/// address will have to be supplied manually when sending packets.
pub struct Unconnected;

/// Indicates that a `Sender` is currently connected to a specific target address.
pub struct Connected {
    addr: SocketAddr,
}

/// An *OSC packet* can contain an OSC message or a bundle of nested packets which is called an
/// *OSC bundle*.
#[derive(Clone, Debug, PartialEq)]
pub enum Packet {
    Message(Message),
    Bundle(Bundle),
}

impl From<Message> for Packet {
    fn from(msg: Message) -> Self {
        Packet::Message(msg)
    }
}

impl From<Bundle> for Packet {
    fn from(bundle: Bundle) -> Self {
        Packet::Bundle(bundle)
    }
}

impl From<rosc::OscPacket> for Packet {
    fn from(packet: rosc::OscPacket) -> Self {
        match packet {
            rosc::OscPacket::Message(msg) => msg.into(),
            rosc::OscPacket::Bundle(bundle) => bundle.into(),
        }
    }
}

impl<A> From<(A, Vec<Type>)> for Packet
where
    A: Into<String>,
{
    fn from((addr, args): (A, Vec<Type>)) -> Self {
        msg(addr, args).into()
    }
}

impl Into<rosc::OscPacket> for Packet {
    fn into(self) -> rosc::OscPacket {
        match self {
            Packet::Message(msg) => rosc::OscPacket::Message(msg),
            Packet::Bundle(bundle) => rosc::OscPacket::Bundle(bundle),
        }
    }
}

impl Packet {
    /// A recursive function that unfolds the packet into the end of the given buffer of messages.
    pub fn unfold(self, msgs: &mut Vec<Message>) {
        match self {
            Packet::Message(msg) => msgs.push(msg),
            Packet::Bundle(bundle) => {
                for packet in bundle.content {
                    Packet::unfold(packet.into(), msgs);
                }
            }
        }
    }

    /// Convert the `Packet` into a Vec containing all `Message`s contained within.
    ///
    /// This uses the `unfold` method internally.
    pub fn into_msgs(self) -> Vec<Message> {
        let mut msgs = vec![];
        self.unfold(&mut msgs);
        msgs
    }
}

/// The default local IP address.
pub fn default_ipv4_addr() -> Ipv4Addr {
    Ipv4Addr::new(0, 0, 0, 0)
}

/// A simple wrapper around the most commonly used `Receiver` constructor.
pub fn receiver(port: u16) -> Result<Receiver, std::io::Error> {
    Receiver::bind(port)
}

/// A simple wrapper around the most commonly used `Sender` constructor.
pub fn sender() -> Result<Sender, std::io::Error> {
    Sender::bind()
}

/// A simplified constructor for an OSC `Message`.
pub fn msg<A>(addr: A, args: Vec<Type>) -> Message
where
    A: Into<String>,
{
    let addr = addr.into();
    let args = Some(args);
    Message { addr, args }
}

/// Decodes the given slice of `bytes` into a `Packet`.
///
/// Returns an `Error` if the slice does not contain a valid OSC packet.
pub fn decode(bytes: &[u8]) -> Result<Packet, Error> {
    rosc::decoder::decode(bytes).map(|p| p.into())
}

/// Encodes the given `Packet` into a `Vec` of bytes.
///
/// Returns an `Error` if the packet is invalid.
pub fn encode(packet: Packet) -> Result<Vec<u8>, Error> {
    let rosc_packet = packet.into();
    rosc::encoder::encode(&rosc_packet)
}

/// Errors that might occur whilst attempting to send or receive an OSC packet.
#[derive(Debug)]
pub enum CommunicationError {
    Io(std::io::Error),
    Osc(Error),
    Poisoned,
}

impl From<std::io::Error> for CommunicationError {
    fn from(err: std::io::Error) -> Self {
        CommunicationError::Io(err)
    }
}

impl From<Error> for CommunicationError {
    fn from(err: Error) -> Self {
        CommunicationError::Osc(err)
    }
}

impl<T> From<std::sync::PoisonError<T>> for CommunicationError {
    fn from(_: std::sync::PoisonError<T>) -> Self {
        CommunicationError::Poisoned
    }
}

impl std::error::Error for CommunicationError {
    fn cause(&self) -> Option<&dyn std::error::Error> {
        match *self {
            CommunicationError::Io(ref err) => Some(err),
            // TODO: Error isn't implemented for OscError - should fix this upstream.
            CommunicationError::Osc(ref _err) => None,
            // CommunicationError::Osc(ref err) => Some(err),
            _ => None,
        }
    }
}

impl std::fmt::Display for CommunicationError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match *self {
            CommunicationError::Io(ref err) => std::fmt::Display::fmt(err, f),
            // TODO: Error isn't implemented for OscError - should fix this upstream.
            CommunicationError::Osc(ref _err) => write!(f, "Failed to decode the OSC packet"),
            CommunicationError::Poisoned => write!(f, "The inner buffer's mutex was poisoned"),
        }
    }
}
<|endoftext|>


<commit_before><commit_msg>Convert more code to result<commit_after><|endoftext|>


<commit_before>use std::sync::Arc;

use wayland_backend::server::{
    ClientData, ClientId, GlobalHandler, GlobalId, Handle, ObjectData, ObjectId,
};

use crate::{
    dispatch::{DelegateDispatch, DelegateDispatchBase},
    Client, DataInit, Dispatch, DisplayHandle, New, Resource,
};

pub(crate) struct GlobalData<I: Resource, D: GlobalDispatch<I>> {
    pub(crate) data: <D as GlobalDispatch<I>>::GlobalData,
}

impl<I: Resource + 'static, D: GlobalDispatch<I> + 'static> GlobalHandler<D> for GlobalData<I, D> {
    fn can_view(&self, id: ClientId, data: &Arc<dyn ClientData<D>>, _: GlobalId) -> bool {
        let client = Client { id, data: data.clone().into_any_arc() };
        <D as GlobalDispatch<I>>::can_view(client, &self.data)
    }

    fn bind(
        self: Arc<Self>,
        handle: &mut Handle<D>,
        data: &mut D,
        client_id: ClientId,
        _: GlobalId,
        object_id: ObjectId,
    ) -> Arc<dyn ObjectData<D>> {
        let mut handle = DisplayHandle::from(handle);
        let client = Client::from_id(&mut handle, client_id).expect("Dead client in bind ?!");
        let resource = <I as Resource>::from_id(&mut handle, object_id)
            .expect("Wrong object_id in GlobalHandler ?!");

        let mut new_data = None;

        data.bind(
            &mut handle,
            &client,
            New::wrap(resource),
            &self.data,
            &mut DataInit { store: &mut new_data },
        );

        match new_data {
            Some(data) => data,
            None => panic!(
                "Bind callback for interface {} did not init new instance.",
                I::interface().name
            ),
        }
    }
}

/// A trait which provides an implementation for handling advertisement of a global to clients with some type
/// of associated user data.
pub trait GlobalDispatch<I: Resource>: Dispatch<I> {
    /// Data associated with the global.
    type GlobalData: Send + Sync + 'static;

    /// Called when a client has bound this global.
    ///
    /// The return value of this function should contain user data to associate the object created by the
    /// client.
    fn bind(
        &mut self,
        handle: &mut DisplayHandle<'_>,
        client: &Client,
        resource: New<I>,
        global_data: &Self::GlobalData,
        data_init: &mut DataInit<'_, Self>,
    );

    /// Checks if the global should be advertised to some client.
    ///
    /// The implementation of this function determines whether a client may see and bind some global. If this
    /// function returns false, the client will not be told the global exists and attempts to bind the global
    /// will raise a protocol error.
    ///
    /// One use of this function is implementing privileged protocols such as XWayland keyboard grabbing
    /// which must only be used by XWayland.
    ///
    /// The default implementation allows all clients to see the global.
    fn can_view(_client: Client, _global_data: &Self::GlobalData) -> bool {
        true
    }
}

/*
 * Dispatch delegation helpers
 */

/// The base trait used to define a delegate type to handle some global.
pub trait DelegateGlobalDispatchBase<I: Resource>: DelegateDispatchBase<I> {
    /// Data associated with the global.
    type GlobalData: Send + Sync + 'static;
}

/// A trait which defines a delegate type to handle some type of global.
///
/// This trait is useful for building modular handlers for globals.
pub trait DelegateGlobalDispatch<
    I: Resource,
    D: GlobalDispatch<I, GlobalData = <Self as DelegateGlobalDispatchBase<I>>::GlobalData>
        + Dispatch<I, UserData = Self::UserData>,
>: Sized + DelegateGlobalDispatchBase<I> + DelegateDispatch<I, D>
{
    /// Called when a client has bound this global.
    ///
    /// The return value of this function should contain user data to associate the object created by the
    /// client.
    fn bind(
        state: &mut D,
        handle: &mut DisplayHandle<'_>,
        client: &Client,
        resource: New<I>,
        global_data: &Self::GlobalData,
        data_init: &mut DataInit<'_, D>,
    );

    /// Checks if the global should be advertised to some client.
    ///
    /// The implementation of this function determines whether a client may see and bind some global. If this
    /// function returns false, the client will not be told the global exists and attempts to bind the global
    /// will raise a protocol error.
    ///
    /// One use of this function is implementing privileged protocols such as XWayland keyboard grabbing
    /// which must only be used by XWayland.
    ///
    /// The default implementation allows all clients to see the global.
    fn can_view(_client: Client, _global_data: &Self::GlobalData) -> bool {
        true
    }
}

#[macro_export]
macro_rules! delegate_global_dispatch {
    (@impl $dispatch_from:ident $(< $( $lt:tt $( : $clt:tt $(+ $dlt:tt )* )? ),+ >)? : $interface: ty => $dispatch_to: ty) => {
        impl$(< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $crate::GlobalDispatch<$interface> for $dispatch_from$(< $( $lt ),+ >)? {
            type GlobalData = <$dispatch_to as $crate::DelegateGlobalDispatchBase<$interface>>::GlobalData;

            fn bind(
                &mut self,
                dhandle: &mut $crate::DisplayHandle<'_>,
                client: &$crate::Client,
                resource: $crate::New<$interface>,
                global_data: &Self::GlobalData,
                data_init: &mut $crate::DataInit<'_, Self>,
            ) {
                <$dispatch_to as $crate::DelegateGlobalDispatch<$interface, Self>>::bind(self, dhandle, client, resource, global_data, data_init)
            }

            fn can_view(client: $crate::Client, global_data: &Self::GlobalData) -> bool {
                <$dispatch_to as $crate::DelegateGlobalDispatch<$interface, Self>>::can_view(client, global_data)
            }
        }
    };
    ($impl:tt : [$($interface: ty),*] => $dispatch_to: ty) => {
        $(
            $crate::delegate_global_dispatch!(@impl $impl : $interface => $dispatch_to);
        )*
    };
}
<commit_msg>server/global: Add smoke test compiling a delegate_global_dispatch invocation<commit_after>use std::sync::Arc;

use wayland_backend::server::{
    ClientData, ClientId, GlobalHandler, GlobalId, Handle, ObjectData, ObjectId,
};

use crate::{
    dispatch::{DelegateDispatch, DelegateDispatchBase},
    Client, DataInit, Dispatch, DisplayHandle, New, Resource,
};

pub(crate) struct GlobalData<I: Resource, D: GlobalDispatch<I>> {
    pub(crate) data: <D as GlobalDispatch<I>>::GlobalData,
}

impl<I: Resource + 'static, D: GlobalDispatch<I> + 'static> GlobalHandler<D> for GlobalData<I, D> {
    fn can_view(&self, id: ClientId, data: &Arc<dyn ClientData<D>>, _: GlobalId) -> bool {
        let client = Client { id, data: data.clone().into_any_arc() };
        <D as GlobalDispatch<I>>::can_view(client, &self.data)
    }

    fn bind(
        self: Arc<Self>,
        handle: &mut Handle<D>,
        data: &mut D,
        client_id: ClientId,
        _: GlobalId,
        object_id: ObjectId,
    ) -> Arc<dyn ObjectData<D>> {
        let mut handle = DisplayHandle::from(handle);
        let client = Client::from_id(&mut handle, client_id).expect("Dead client in bind ?!");
        let resource = <I as Resource>::from_id(&mut handle, object_id)
            .expect("Wrong object_id in GlobalHandler ?!");

        let mut new_data = None;

        data.bind(
            &mut handle,
            &client,
            New::wrap(resource),
            &self.data,
            &mut DataInit { store: &mut new_data },
        );

        match new_data {
            Some(data) => data,
            None => panic!(
                "Bind callback for interface {} did not init new instance.",
                I::interface().name
            ),
        }
    }
}

/// A trait which provides an implementation for handling advertisement of a global to clients with some type
/// of associated user data.
pub trait GlobalDispatch<I: Resource>: Dispatch<I> {
    /// Data associated with the global.
    type GlobalData: Send + Sync + 'static;

    /// Called when a client has bound this global.
    ///
    /// The return value of this function should contain user data to associate the object created by the
    /// client.
    fn bind(
        &mut self,
        handle: &mut DisplayHandle<'_>,
        client: &Client,
        resource: New<I>,
        global_data: &Self::GlobalData,
        data_init: &mut DataInit<'_, Self>,
    );

    /// Checks if the global should be advertised to some client.
    ///
    /// The implementation of this function determines whether a client may see and bind some global. If this
    /// function returns false, the client will not be told the global exists and attempts to bind the global
    /// will raise a protocol error.
    ///
    /// One use of this function is implementing privileged protocols such as XWayland keyboard grabbing
    /// which must only be used by XWayland.
    ///
    /// The default implementation allows all clients to see the global.
    fn can_view(_client: Client, _global_data: &Self::GlobalData) -> bool {
        true
    }
}

/*
 * Dispatch delegation helpers
 */

/// The base trait used to define a delegate type to handle some global.
pub trait DelegateGlobalDispatchBase<I: Resource>: DelegateDispatchBase<I> {
    /// Data associated with the global.
    type GlobalData: Send + Sync + 'static;
}

/// A trait which defines a delegate type to handle some type of global.
///
/// This trait is useful for building modular handlers for globals.
pub trait DelegateGlobalDispatch<
    I: Resource,
    D: GlobalDispatch<I, GlobalData = <Self as DelegateGlobalDispatchBase<I>>::GlobalData>
        + Dispatch<I, UserData = Self::UserData>,
>: Sized + DelegateGlobalDispatchBase<I> + DelegateDispatch<I, D>
{
    /// Called when a client has bound this global.
    ///
    /// The return value of this function should contain user data to associate the object created by the
    /// client.
    fn bind(
        state: &mut D,
        handle: &mut DisplayHandle<'_>,
        client: &Client,
        resource: New<I>,
        global_data: &Self::GlobalData,
        data_init: &mut DataInit<'_, D>,
    );

    /// Checks if the global should be advertised to some client.
    ///
    /// The implementation of this function determines whether a client may see and bind some global. If this
    /// function returns false, the client will not be told the global exists and attempts to bind the global
    /// will raise a protocol error.
    ///
    /// One use of this function is implementing privileged protocols such as XWayland keyboard grabbing
    /// which must only be used by XWayland.
    ///
    /// The default implementation allows all clients to see the global.
    fn can_view(_client: Client, _global_data: &Self::GlobalData) -> bool {
        true
    }
}

#[macro_export]
macro_rules! delegate_global_dispatch {
    (@impl $dispatch_from:ident $(< $( $lt:tt $( : $clt:tt $(+ $dlt:tt )* )? ),+ >)? : $interface: ty => $dispatch_to: ty) => {
        impl$(< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $crate::GlobalDispatch<$interface> for $dispatch_from$(< $( $lt ),+ >)? {
            type GlobalData = <$dispatch_to as $crate::DelegateGlobalDispatchBase<$interface>>::GlobalData;

            fn bind(
                &mut self,
                dhandle: &mut $crate::DisplayHandle<'_>,
                client: &$crate::Client,
                resource: $crate::New<$interface>,
                global_data: &Self::GlobalData,
                data_init: &mut $crate::DataInit<'_, Self>,
            ) {
                <$dispatch_to as $crate::DelegateGlobalDispatch<$interface, Self>>::bind(self, dhandle, client, resource, global_data, data_init)
            }

            fn can_view(client: $crate::Client, global_data: &Self::GlobalData) -> bool {
                <$dispatch_to as $crate::DelegateGlobalDispatch<$interface, Self>>::can_view(client, global_data)
            }
        }
    };
    ($impl:tt : [$($interface: ty),*] => $dispatch_to: ty) => {
        $(
            $crate::delegate_global_dispatch!(@impl $impl : $interface => $dispatch_to);
        )*
    };
}

#[cfg(test)]
mod tests {
    #[test]
    fn smoke_test_dispatch_global_dispatch() {
        use crate::{
            delegate_dispatch, protocol::wl_output, Client, DataInit, DelegateDispatch,
            DelegateDispatchBase, DelegateGlobalDispatch, DelegateGlobalDispatchBase, Dispatch,
            DisplayHandle, GlobalDispatch, New,
        };

        struct DelegateToMe;

        impl DelegateDispatchBase<wl_output::WlOutput> for DelegateToMe {
            type UserData = ();
        }
        impl<D> DelegateDispatch<wl_output::WlOutput, D> for DelegateToMe
        where
            D: Dispatch<wl_output::WlOutput, UserData = Self::UserData> + AsMut<DelegateToMe>,
        {
            fn request(
                _state: &mut D,
                _client: &Client,
                _resource: &wl_output::WlOutput,
                _request: wl_output::Request,
                _data: &Self::UserData,
                _dhandle: &mut DisplayHandle,
                _data_init: &mut DataInit<'_, D>,
            ) {
            }
        }
        impl DelegateGlobalDispatchBase<wl_output::WlOutput> for DelegateToMe {
            type GlobalData = ();
        }
        impl<D> DelegateGlobalDispatch<wl_output::WlOutput, D> for DelegateToMe
        where
            D: GlobalDispatch<wl_output::WlOutput, GlobalData = Self::GlobalData>,
            D: Dispatch<wl_output::WlOutput, UserData = ()>,
            D: AsMut<DelegateToMe>,
        {
            fn bind(
                _state: &mut D,
                _handle: &mut DisplayHandle<'_>,
                _client: &Client,
                _resource: New<wl_output::WlOutput>,
                _global_data: &Self::GlobalData,
                _data_init: &mut DataInit<'_, D>,
            ) {
            }
        }

        struct ExampleApp {
            delegate: DelegateToMe,
        }

        delegate_dispatch!(ExampleApp: [wl_output::WlOutput] => DelegateToMe);
        delegate_global_dispatch!(ExampleApp: [wl_output::WlOutput] => DelegateToMe);

        impl AsMut<DelegateToMe> for ExampleApp {
            fn as_mut(&mut self) -> &mut DelegateToMe {
                &mut self.delegate
            }
        }
    }
}
<|endoftext|>


<commit_before>#![feature(iter_arith)]
use std::io;
use std::io::prelude::*;

struct House { lights: Vec<Vec<bool>>, }
impl House {
    fn step(&self) -> House {
        let lights = (0..self.lights.len()).map(|i| {
            (0..self.lights.len()).map(|j| {
                let light = *self.lights.get(i).unwrap().get(j).unwrap();
                let neighbors = self.neighbors(i, j);
                let neighbors_on = neighbors.iter().filter(|b| **b).count();
                if light {
                    if neighbors_on == 2 || neighbors_on == 3 {
                        true
                    } else {
                        false
                    }
                } else {
                    if neighbors_on == 3 { true } else { false }
                }
            }).collect()
        }).collect();
        House { lights: lights, }
    }

    fn neighbors(&self, x: usize, y: usize) -> Vec<bool> {
        let indices =
            (-1..2).fold(vec![], |mut list, i| {
                let points = (-1..2).filter(|j| {
                    let x_i = x as i32 + i;
                    let y_j = y as i32 + j;
                    (i != 0 || *j != 0) && self.in_bounds(x_i, y_j)
                }).map(|j| (i, j));
                list.extend(points);
                list
            });
        let neighbors: Vec<bool> = indices.iter().map(|p| {
            let x_i = x as i32 + p.0;
            let y_j = y as i32 + p.1;
            *self.lights.get(x_i as usize).unwrap().get(y_j as usize).unwrap()
        }).collect();
        neighbors
    }

    fn lights_on(&self) -> u32 {
        (0..self.lights.len()).map(|i| {
            let on_in_row: u32 = (0..self.lights.len()).map(|j| {
                if *self.lights.get(i).unwrap().get(j).unwrap() {
                    1
                } else {
                    0
                }
            }).sum();
            on_in_row
        }).sum()
    }

    fn in_bounds(&self, x: i32, y: i32) -> bool {
        x >= 0 && y >= 0 &&
            x < self.lights.len() as i32 && y < self.lights.len() as i32
    }
}

fn parse(line: &String) -> Vec<bool> {
    line.chars().map(|c| c == '#').collect()
}

fn main() {
    let stdin = io::stdin();
    let lights: Vec<Vec<bool>> =
        stdin
        .lock()
        .lines()
        .map(|line| parse(&line.unwrap()))
        .collect();
    let house = House { lights: lights, };
    let last = (0..100).fold(house, |h, _| h.step());
    println!("{}", last.lights_on());
}
<commit_msg>day 18, part 2<commit_after>#![feature(iter_arith)]
use std::io;
use std::io::prelude::*;

struct House { lights: Vec<Vec<bool>>, }
impl House {
    fn corners_on(&self) -> House {
        let lights = (0..self.lights.len()).map(|i| {
            (0..self.lights.len()).map(|j| {
                if self.is_corner(i, j) {
                    true
                } else {
                    let light = *self.lights.get(i).unwrap().get(j).unwrap();
                    light
                }
            }).collect()
        }).collect();
        House { lights: lights, }
    }

    fn is_corner(&self, x: usize, y: usize) -> bool {
        (x == 0 || x == self.lights.len() - 1) &&
            (y == 0 || y == self.lights.len() - 1)
    }

    fn step(&self) -> House {
        let lights = (0..self.lights.len()).map(|i| {
            (0..self.lights.len()).map(|j| {
                let light = *self.lights.get(i).unwrap().get(j).unwrap();
                let neighbors = self.neighbors(i, j);
                let neighbors_on = neighbors.iter().filter(|b| **b).count();
                if light {
                    if neighbors_on == 2 || neighbors_on == 3 {
                        true
                    } else {
                        false
                    }
                } else {
                    if neighbors_on == 3 { true } else { false }
                }
            }).collect()
        }).collect();
        House { lights: lights, }
    }

    fn neighbors(&self, x: usize, y: usize) -> Vec<bool> {
        let indices =
            (-1..2).fold(vec![], |mut list, i| {
                let points = (-1..2).filter(|j| {
                    let x_i = x as i32 + i;
                    let y_j = y as i32 + j;
                    (i != 0 || *j != 0) && self.in_bounds(x_i, y_j)
                }).map(|j| (i, j));
                list.extend(points);
                list
            });
        let neighbors: Vec<bool> = indices.iter().map(|p| {
            let x_i = x as i32 + p.0;
            let y_j = y as i32 + p.1;
            *self.lights.get(x_i as usize).unwrap().get(y_j as usize).unwrap()
        }).collect();
        neighbors
    }

    fn lights_on(&self) -> u32 {
        (0..self.lights.len()).map(|i| {
            let on_in_row: u32 = (0..self.lights.len()).map(|j| {
                if *self.lights.get(i).unwrap().get(j).unwrap() {
                    1
                } else {
                    0
                }
            }).sum();
            on_in_row
        }).sum()
    }

    fn in_bounds(&self, x: i32, y: i32) -> bool {
        x >= 0 && y >= 0 &&
            x < self.lights.len() as i32 && y < self.lights.len() as i32
    }
}

fn parse(line: &String) -> Vec<bool> {
    line.chars().map(|c| c == '#').collect()
}

fn main() {
    let stdin = io::stdin();
    let lights: Vec<Vec<bool>> =
        stdin
        .lock()
        .lines()
        .map(|line| parse(&line.unwrap()))
        .collect();
    let house = House { lights: lights, };
    let last = (0..100).fold(house.corners_on(), |h, _| h.step().corners_on());
    println!("{}", last.lights_on());
}
<|endoftext|>


<commit_before>extern crate onmf;
use onmf::helpers::{ToImage, Normalize};
use onmf::testimage_generator;
use onmf::testimage_generator::{magnify};

extern crate nalgebra;
use self::nalgebra::{DMat};

extern crate rand;
use rand::{StdRng, SeedableRng};

fn main() {
    let mag_factor = 10;

    let seed: &[_] = &[1, 2, 3, 4];
    let mut rng: StdRng = SeedableRng::from_seed(seed);

    let steps = 6;
    let per_step = 1000;
    let nsamples = steps * per_step;
    let nobserved = 10 * 10;
    let nhidden = 20 + steps * 2;

    let mut data = DMat::<f64>::new_zeros(nsamples, nobserved);

    // write on testimage into each row of data
    for (step, i, factor) in testimage_generator::testimages::<f64, _>(per_step, &mut rng) {
        let row = step * i;
        for (col, val) in factor.as_vec().iter().enumerate() {
            data[(row, col)] = val.clone();
        }
    }

    let mut nmf = onmf::OrthogonalNMF::<f64>::init_randomly(
        nhidden, nobserved, nsamples);

    for _ in 0..200 { nmf.iterate(&data) }

    // take a snapshot
    // read testimage out of each row of nmf.hidden
    for irow in 0..nhidden {
        let mut column = Vec::<f64>::new();
        for icol in 0..nobserved {
            column.push(nmf.hidden[(irow, icol)]);
        }
        let image = DMat::<f64>::from_col_vec(10, 10, &column[..]);
        // println!("image {}", irow);
        // println!("{:?}", image);
        magnify(image.normalize(), mag_factor).save_to_png(&format!("orthogonal-nmf-hidden-snapshot-{}.png", irow)[..]).unwrap();
    }
}
<commit_msg>write out testimages on each iteration to observe progress<commit_after>extern crate onmf;
use onmf::helpers::{ToImage, Normalize};
use onmf::testimage_generator;
use onmf::testimage_generator::{magnify};

extern crate nalgebra;
use self::nalgebra::{DMat};

extern crate rand;
use rand::{StdRng, SeedableRng};

fn main() {
    let mag_factor = 10;

    let seed: &[_] = &[1, 2, 3, 4];
    let mut rng: StdRng = SeedableRng::from_seed(seed);

    let steps = 6;
    let per_step = 1000;
    let nsamples = steps * per_step;
    let nobserved = 10 * 10;
    let nhidden = 20 + steps * 2;

    let mut data = DMat::<f64>::new_zeros(nsamples, nobserved);

    // write one testimage into each row of data
    for (step, i, factor) in testimage_generator::testimages::<f64, _>(per_step, &mut rng) {
        let row = step * i;
        for (col, val) in factor.as_vec().iter().enumerate() {
            data[(row, col)] = val.clone();
        }
    }

    let mut nmf = onmf::OrthogonalNMF::<f64>::init_randomly(
        nhidden, nobserved, nsamples);

    for _ in 0..1000 {
        nmf.iterate(&data);

        // read testimage out of each row of nmf.hidden
        for irow in 0..nhidden {
            let mut column = Vec::<f64>::new();
            for icol in 0..nobserved {
                column.push(nmf.hidden[(irow, icol)]);
            }
            let image = DMat::<f64>::from_col_vec(10, 10, &column[..]);
            // println!("image {}", irow);
            // println!("{:?}", image);
            magnify(image.normalize(), mag_factor).save_to_png(&format!("orthogonal-nmf-hidden-snapshot-{}.png", irow)[..]).unwrap();
        }
    }
}
<|endoftext|>