<commit_before>use std::fs::File;
use std::env;
use std::io::BufReader;
use std::io::BufRead;
use std::io::Write;
use std::collections::HashMap;
use std::collections::hash_map::Entry::Occupied;
use std::collections::hash_map::Entry::Vacant;
use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::rc::Rc;

extern crate tools;
use tools::find_source_file;
use tools::file_format::analysis::{read_analysis, read_target, AnalysisKind};
use tools::config;

extern crate rustc_serialize;
use rustc_serialize::json::{Json, ToJson};

#[derive(Debug)]
struct SearchResult {
    lineno: u32,
    bounds: (u32, u32),
    line: Rc<String>,
    context: Rc<String>,
    contextsym: Rc<String>,
}

impl ToJson for SearchResult {
    fn to_json(&self) -> Json {
        let (st, en) = self.bounds;
        let bounds = vec![st, en];

        let mut obj = BTreeMap::new();
        obj.insert("lno".to_string(), self.lineno.to_json());
        obj.insert("bounds".to_string(), bounds.to_json());
        obj.insert("line".to_string(), self.line.to_json());
        obj.insert("context".to_string(), self.context.to_json());
        obj.insert("contextsym".to_string(), self.contextsym.to_json());
        Json::Object(obj)
    }
}

fn split_scopes(id: &str) -> Vec<String> {
    let mut result = Vec::new();
    let mut start = 0;
    let mut argument_nesting = 0;
    for (index, m) in id.match_indices(|c| c == ':' || c == '<' || c == '>') {
        if m == ":" && argument_nesting == 0 {
            if start != index {
                result.push(id[start .. index].to_owned());
                start = index + 1;
            } else {
                start = index + 1;
            }
        } else if m == "<" {
            argument_nesting += 1;
        } else if m == ">" {
            argument_nesting -= 1;
        }
    }
    result.push(id[start ..].to_owned());
    return result;
}

struct StringIntern {
    set: HashMap<Rc<String>, ()>
}

impl StringIntern {
    fn new() -> StringIntern {
        StringIntern {set: HashMap::new()}
    }

    fn add(&mut self, s: String) -> Rc<String> {
        let new_rc = Rc::new(s);
        match self.set.entry(new_rc) {
            Occupied(o) => Rc::clone(&o.key()),
            Vacant(v) => {
                let rval = Rc::clone(&v.key());
                v.insert(());
                rval
            }
        }
    }
}

fn main() {
    let args: Vec<_> = env::args().collect();

    let cfg = config::load(&args[1], false);

    let tree_name = &args[2];
    let tree_config = cfg.trees.get(tree_name).unwrap();

    let filenames_file = &args[3];

    let file_paths: Vec<String> = BufReader::new(File::open(filenames_file).unwrap())
        .lines().map(|x| x.unwrap()).collect();
    let output_file = format!("{}/crossref", tree_config.paths.index_path);
    let jump_file = format!("{}/jumps", tree_config.paths.index_path);
    let id_file = format!("{}/identifiers", tree_config.paths.index_path);

    let mut strings = StringIntern::new();
    let empty_string = strings.add("".to_string());

    let mut table = BTreeMap::new();
    let mut pretty_table = HashMap::new();
    let mut id_table = BTreeMap::new();
    let mut jumps = Vec::new();

    for path in &file_paths {
        print!("File {}\n", path);

        let analysis_fname = format!("{}/analysis/{}", tree_config.paths.index_path, path);
        let analysis = read_analysis(&analysis_fname, &read_target);

        let source_fname = find_source_file(path, &tree_config.paths.files_path, &tree_config.paths.objdir_path);
        let source_file = match File::open(source_fname) {
            Ok(f) => f,
            Err(_) => {
                println!("Unable to open source file");
                continue;
            },
        };
        let reader = BufReader::new(&source_file);
        let lines: Vec<_> = reader.lines().map(
            |l| match l {
                Ok(line) => {
                    let line_cut = line.trim_right();
                    let len = line_cut.len();
                    let line_cut = line_cut.trim_left();
                    let offset = (len - line_cut.len()) as u32;
                    let mut buf = String::new();
                    let mut i = 0;
                    for c in line_cut.chars() {
                        buf.push(c);
                        i += 1;
                        if i > 100 {
                            break;
                        }
                    };
                    (strings.add(buf), offset)
                },
                Err(_) => (Rc::clone(&empty_string), 0)
            }).collect();

        for datum in analysis {
            for piece in datum.data {
                let sym = strings.add(piece.sym.to_owned());
                let t1 = table.entry(Rc::clone(&sym)).or_insert(BTreeMap::new());
                let t2 = t1.entry(piece.kind).or_insert(BTreeMap::new());
                let p: &str = &path;
                let t3 = t2.entry(p).or_insert(Vec::new());
                let lineno = (datum.loc.lineno - 1) as usize;
                if lineno >= lines.len() {
                    print!("Bad line number in file {} (line {})\n", path, lineno);
                    continue;
                }
                let (line, offset) = lines[lineno].clone();
                t3.push(SearchResult {
                    lineno: datum.loc.lineno,
                    bounds: (datum.loc.col_start - offset, datum.loc.col_end - offset),
                    line: line,
                    context: strings.add(piece.context),
                    contextsym: strings.add(piece.contextsym),
                });

                let pretty = strings.add(piece.pretty.to_owned());
                pretty_table.insert(Rc::clone(&sym), Rc::clone(&pretty));

                let ch = piece.sym.chars().nth(0).unwrap();
                if !(ch >= '0' && ch <= '9') && !piece.sym.contains(' ') {
                    let t1 = id_table.entry(pretty).or_insert(BTreeSet::new());
                    t1.insert(sym);
                }
            }
        }
    }

    let mut outputf = File::create(output_file).unwrap();

    for (id, id_data) in table {
        let mut kindmap = BTreeMap::new();
        for (kind, kind_data) in &id_data {
            let mut result = Vec::new();
            for (path, results) in kind_data {
                let mut obj = BTreeMap::new();
                obj.insert("path".to_string(), path.to_json());
                obj.insert("lines".to_string(), results.to_json());
                result.push(Json::Object(obj));
            }
            let kindstr = match *kind {
                AnalysisKind::Use => "Uses",
                AnalysisKind::Def => "Definitions",
                AnalysisKind::Assign => "Assignments",
                AnalysisKind::Decl => "Declarations",
                AnalysisKind::Idl => "IDL",
            };
            kindmap.insert(kindstr.to_string(), Json::Array(result));
        }
        let kindmap = Json::Object(kindmap);

        let _ = outputf.write_all(format!("{}\n{}\n", id, kindmap.to_string()).as_bytes());

        if id_data.contains_key(&AnalysisKind::Def) {
            let defs = id_data.get(&AnalysisKind::Def).unwrap();
            if defs.len() == 1 {
                for (path, results) in defs {
                    if results.len() == 1 {
                        let mut v = Vec::new();
                        v.push(id.to_json());
                        v.push(path.to_json());
                        v.push(results[0].lineno.to_json());
                        let pretty = pretty_table.get(&id).unwrap();
                        v.push(pretty.to_json());
                        jumps.push(Json::Array(v));
                    }
                }
            }
        }
    }

    let mut jumpf = File::create(jump_file).unwrap();
    for jump in jumps {
        let _ = jumpf.write_all((jump.to_string() + "\n").as_bytes());
    }

    let mut idf = File::create(id_file).unwrap();
    for (id, syms) in id_table {
        for sym in syms {
            let components = split_scopes(&id);
            for i in 0..components.len() {
                let sub = &components[i..components.len()];
                let sub = sub.join("::");

                if !sub.is_empty() {
                    let line = format!("{} {}\n", sub, sym);
                    let _ = idf.write_all(line.as_bytes());
                }
            }
        }
    }
}
<commit_msg>Use a more idiomatic way of constructing a buffer in crossref.rs.<commit_after>use std::fs::File;
use std::env;
use std::io::BufReader;
use std::io::BufRead;
use std::io::Write;
use std::collections::HashMap;
use std::collections::hash_map::Entry::Occupied;
use std::collections::hash_map::Entry::Vacant;
use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::rc::Rc;

extern crate tools;
use tools::find_source_file;
use tools::file_format::analysis::{read_analysis, read_target, AnalysisKind};
use tools::config;

extern crate rustc_serialize;
use rustc_serialize::json::{Json, ToJson};

#[derive(Debug)]
struct SearchResult {
    lineno: u32,
    bounds: (u32, u32),
    line: Rc<String>,
    context: Rc<String>,
    contextsym: Rc<String>,
}

impl ToJson for SearchResult {
    fn to_json(&self) -> Json {
        let (st, en) = self.bounds;
        let bounds = vec![st, en];

        let mut obj = BTreeMap::new();
        obj.insert("lno".to_string(), self.lineno.to_json());
        obj.insert("bounds".to_string(), bounds.to_json());
        obj.insert("line".to_string(), self.line.to_json());
        obj.insert("context".to_string(), self.context.to_json());
        obj.insert("contextsym".to_string(), self.contextsym.to_json());
        Json::Object(obj)
    }
}

fn split_scopes(id: &str) -> Vec<String> {
    let mut result = Vec::new();
    let mut start = 0;
    let mut argument_nesting = 0;
    for (index, m) in id.match_indices(|c| c == ':' || c == '<' || c == '>') {
        if m == ":" && argument_nesting == 0 {
            if start != index {
                result.push(id[start .. index].to_owned());
                start = index + 1;
            } else {
                start = index + 1;
            }
        } else if m == "<" {
            argument_nesting += 1;
        } else if m == ">" {
            argument_nesting -= 1;
        }
    }
    result.push(id[start ..].to_owned());
    return result;
}

struct StringIntern {
    set: HashMap<Rc<String>, ()>
}

impl StringIntern {
    fn new() -> StringIntern {
        StringIntern {set: HashMap::new()}
    }

    fn add(&mut self, s: String) -> Rc<String> {
        let new_rc = Rc::new(s);
        match self.set.entry(new_rc) {
            Occupied(o) => Rc::clone(&o.key()),
            Vacant(v) => {
                let rval = Rc::clone(&v.key());
                v.insert(());
                rval
            }
        }
    }
}

fn main() {
    let args: Vec<_> = env::args().collect();

    let cfg = config::load(&args[1], false);

    let tree_name = &args[2];
    let tree_config = cfg.trees.get(tree_name).unwrap();

    let filenames_file = &args[3];

    let file_paths: Vec<String> = BufReader::new(File::open(filenames_file).unwrap())
        .lines().map(|x| x.unwrap()).collect();
    let output_file = format!("{}/crossref", tree_config.paths.index_path);
    let jump_file = format!("{}/jumps", tree_config.paths.index_path);
    let id_file = format!("{}/identifiers", tree_config.paths.index_path);

    let mut strings = StringIntern::new();
    let empty_string = strings.add("".to_string());

    let mut table = BTreeMap::new();
    let mut pretty_table = HashMap::new();
    let mut id_table = BTreeMap::new();
    let mut jumps = Vec::new();

    for path in &file_paths {
        print!("File {}\n", path);

        let analysis_fname = format!("{}/analysis/{}", tree_config.paths.index_path, path);
        let analysis = read_analysis(&analysis_fname, &read_target);

        let source_fname = find_source_file(path, &tree_config.paths.files_path, &tree_config.paths.objdir_path);
        let source_file = match File::open(source_fname) {
            Ok(f) => f,
            Err(_) => {
                println!("Unable to open source file");
                continue;
            },
        };
        let reader = BufReader::new(&source_file);
        let lines: Vec<_> = reader.lines().map(
            |l| match l {
                Ok(line) => {
                    let line_cut = line.trim_right();
                    let len = line_cut.len();
                    let line_cut = line_cut.trim_left();
                    let offset = (len - line_cut.len()) as u32;
                    let buf = line_cut.chars().take(100).collect();
                    (strings.add(buf), offset)
                },
                Err(_) => (Rc::clone(&empty_string), 0)
            }).collect();

        for datum in analysis {
            for piece in datum.data {
                let sym = strings.add(piece.sym.to_owned());
                let t1 = table.entry(Rc::clone(&sym)).or_insert(BTreeMap::new());
                let t2 = t1.entry(piece.kind).or_insert(BTreeMap::new());
                let p: &str = &path;
                let t3 = t2.entry(p).or_insert(Vec::new());
                let lineno = (datum.loc.lineno - 1) as usize;
                if lineno >= lines.len() {
                    print!("Bad line number in file {} (line {})\n", path, lineno);
                    continue;
                }
                let (line, offset) = lines[lineno].clone();
                t3.push(SearchResult {
                    lineno: datum.loc.lineno,
                    bounds: (datum.loc.col_start - offset, datum.loc.col_end - offset),
                    line: line,
                    context: strings.add(piece.context),
                    contextsym: strings.add(piece.contextsym),
                });

                let pretty = strings.add(piece.pretty.to_owned());
                pretty_table.insert(Rc::clone(&sym), Rc::clone(&pretty));

                let ch = piece.sym.chars().nth(0).unwrap();
                if !(ch >= '0' && ch <= '9') && !piece.sym.contains(' ') {
                    let t1 = id_table.entry(pretty).or_insert(BTreeSet::new());
                    t1.insert(sym);
                }
            }
        }
    }

    let mut outputf = File::create(output_file).unwrap();

    for (id, id_data) in table {
        let mut kindmap = BTreeMap::new();
        for (kind, kind_data) in &id_data {
            let mut result = Vec::new();
            for (path, results) in kind_data {
                let mut obj = BTreeMap::new();
                obj.insert("path".to_string(), path.to_json());
                obj.insert("lines".to_string(), results.to_json());
                result.push(Json::Object(obj));
            }
            let kindstr = match *kind {
                AnalysisKind::Use => "Uses",
                AnalysisKind::Def => "Definitions",
                AnalysisKind::Assign => "Assignments",
                AnalysisKind::Decl => "Declarations",
                AnalysisKind::Idl => "IDL",
            };
            kindmap.insert(kindstr.to_string(), Json::Array(result));
        }
        let kindmap = Json::Object(kindmap);

        let _ = outputf.write_all(format!("{}\n{}\n", id, kindmap.to_string()).as_bytes());

        if id_data.contains_key(&AnalysisKind::Def) {
            let defs = id_data.get(&AnalysisKind::Def).unwrap();
            if defs.len() == 1 {
                for (path, results) in defs {
                    if results.len() == 1 {
                        let mut v = Vec::new();
                        v.push(id.to_json());
                        v.push(path.to_json());
                        v.push(results[0].lineno.to_json());
                        let pretty = pretty_table.get(&id).unwrap();
                        v.push(pretty.to_json());
                        jumps.push(Json::Array(v));
                    }
                }
            }
        }
    }

    let mut jumpf = File::create(jump_file).unwrap();
    for jump in jumps {
        let _ = jumpf.write_all((jump.to_string() + "\n").as_bytes());
    }

    let mut idf = File::create(id_file).unwrap();
    for (id, syms) in id_table {
        for sym in syms {
            let components = split_scopes(&id);
            for i in 0..components.len() {
                let sub = &components[i..components.len()];
                let sub = sub.join("::");

                if !sub.is_empty() {
                    let line = format!("{} {}\n", sub, sym);
                    let _ = idf.write_all(line.as_bytes());
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::str::FromStr;

fn minmax(seq: & Vec<f32>, adj: bool) -> (f32, f32) {
    let mut min: (usize, &f32) = (0, &seq[0]);
    let mut max: (usize, &f32) = (0, &seq[0]);

    for (i, e) in seq.iter().enumerate() {
        if e < min.1 {
            min = (i, e);
        }
        if e > max.1 {
            max = (i, e);
        }
    }

    if min.1 == max.1 {
        return (-1.0, -1.0); // TODO convert tuple into option
    }
    if min.0 > max.0 {
        let spl;
        let s1: Vec<f32>;
        let s2: Vec<f32>;
        let mm1: (f32, f32);
        let mm2: (f32, f32);

        spl = seq.split_at(max.0 + 1);
        s1 = spl.0.to_vec();
        s2 = spl.1.to_vec();
        if s1.len() == 0 || s2.len() == 0 {
            return (-1.0, -1.0);
        }
        mm1 = minmax(& s1, true);
        mm2 = minmax(& s2, true);
        if (mm1.1 - mm1.0) >= (mm2.1 - mm2.0) {
            return mm1;
        }
        return mm2;
    }
    if adj && min.0 + 1 == max.0 {
        let mut s1: Vec<f32> = seq.clone();
        let mut s2: Vec<f32> = seq.clone();

        let mm1: (f32, f32);
        let mm2: (f32, f32);

        s1.remove(min.0);
        s2.remove(max.0);

        mm1 = minmax(& s1, false);
        mm2 = minmax(& s2, false);

        if (mm1.1 - mm1.0) >= (mm2.1 - mm2.0) {
            return mm1;
        }
        return mm2;
    }

    return (*min.1, *max.1);
}

fn main() {
    let mut input = String::new();
    let market: Vec<f32>;
    let mm: (f32, f32);

    io::stdin().read_line(&mut input);
    market = input.split_whitespace()
        .map(|e| f32::from_str(e).unwrap())
        .collect();

    mm = minmax(&market, true);
    match mm {
        (-1.0, -1.0) => println!("No result"),
        _ => println!("{} {}", mm.0, mm.1),
    }
}
<commit_msg>Refactored from pure floats to options<commit_after>use std::io;
use std::str::FromStr;

fn min_pair(a: & Option<(f32, f32)>, b: & Option<(f32, f32)>) -> Option<(f32, f32)> {
    return match (*a, *b) {
        (None, None) => None,
        (_, None) => *a,
        (None, _) => *b,
        (Some(at), Some(bt)) =>
            if (at.1 - at.0) >= (bt.1 - bt.0) { *a } else { *b }
    }
}

fn minmax(seq: & Vec<f32>, adj: bool) -> Option<(f32, f32)> {
    let mut min: (usize, &f32) = (0, &seq[0]);
    let mut max: (usize, &f32) = (0, &seq[0]);

    for (i, e) in seq.iter().enumerate() {
        if e < min.1 {
            min = (i, e);
        }
        if e > max.1 {
            max = (i, e);
        }
    }

    if min.1 == max.1 {
        return None;
    }
    if min.0 > max.0 {
        let spl;
        let s1: Vec<f32>;
        let s2: Vec<f32>;
        let mm1: Option<(f32, f32)>;
        let mm2: Option<(f32, f32)>;

        spl = seq.split_at(max.0 + 1);
        s1 = spl.0.to_vec();
        s2 = spl.1.to_vec();
        if s1.len() == 0 || s2.len() == 0 {
            return None;
        }
        mm1 = minmax(& s1, true);
        mm2 = minmax(& s2, true);
        return min_pair(&mm1, &mm2);
    }
    if adj && min.0 + 1 == max.0 {
        let mut s1: Vec<f32> = seq.clone();
        let mut s2: Vec<f32> = seq.clone();

        let mm1: Option<(f32, f32)>;
        let mm2: Option<(f32, f32)>;

        s1.remove(min.0);
        s2.remove(max.0);

        mm1 = minmax(& s1, false);
        mm2 = minmax(& s2, false);
        return min_pair(&mm1, &mm2);
    }

    return Some((*min.1, *max.1));
}

fn main() {
    let mut input = String::new();
    let market: Vec<f32>;
    let mm: Option<(f32, f32)>;

    io::stdin().read_line(&mut input);
    market = input.split_whitespace()
        .map(|e| f32::from_str(e).unwrap())
        .collect();

    mm = minmax(&market, true);
    match mm {
        None => println!("No result"),
        Some(v) => println!("{} {}", v.0, v.1),
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

#![deny(warnings)]
// Enable all clippy lints except for many of the pedantic ones. It's a shame this needs to be copied and pasted across crates, but there doesn't appear to be a way to include inner attributes from a common source.
#![deny(
  clippy::all,
  clippy::default_trait_access,
  clippy::expl_impl_clone_on_copy,
  clippy::if_not_else,
  clippy::needless_continue,
  clippy::single_match_else,
  clippy::unseparated_literal_suffix,
  clippy::used_underscore_binding
)]
// It is often more clear to show that nothing is being moved.
#![allow(clippy::match_ref_pats)]
// Subjective style.
#![allow(
  clippy::len_without_is_empty,
  clippy::redundant_field_names,
  clippy::too_many_arguments
)]
// Default isn't as big a deal as people seem to think it is.
#![allow(clippy::new_without_default, clippy::new_ret_no_self)]
// Arc<Mutex> can be more clear than needing to grok Orderings:
#![allow(clippy::mutex_atomic)]

use hex;

use digest::{Digest as DigestTrait, FixedOutput};
use serde::ser::{Serialize, SerializeStruct, Serializer};
use sha2::Sha256;

use std::fmt;
use std::io::{self, Write};

pub const EMPTY_FINGERPRINT: Fingerprint = Fingerprint([
  0xe3, 0xb0, 0xc4, 0x42, 0x98, 0xfc, 0x1c, 0x14, 0x9a, 0xfb, 0xf4, 0xc8, 0x99, 0x6f, 0xb9, 0x24,
  0x27, 0xae, 0x41, 0xe4, 0x64, 0x9b, 0x93, 0x4c, 0xa4, 0x95, 0x99, 0x1b, 0x78, 0x52, 0xb8, 0x55,
]);
pub const EMPTY_DIGEST: Digest = Digest(EMPTY_FINGERPRINT, 0);

const FINGERPRINT_SIZE: usize = 32;

#[derive(Clone, Copy, Eq, Hash, PartialEq, Ord, PartialOrd)]
pub struct Fingerprint(pub [u8; FINGERPRINT_SIZE]);

impl Fingerprint {
  pub fn from_bytes_unsafe(bytes: &[u8]) -> Fingerprint {
    if bytes.len() != FINGERPRINT_SIZE {
      panic!(
        "Input value was not a fingerprint; had length: {}",
        bytes.len()
      );
    }

    let mut fingerprint = [0; FINGERPRINT_SIZE];
    fingerprint.clone_from_slice(&bytes[0..FINGERPRINT_SIZE]);
    Fingerprint(fingerprint)
  }

  pub fn from_hex_string(hex_string: &str) -> Result<Fingerprint, String> {
    <[u8; FINGERPRINT_SIZE] as hex::FromHex>::from_hex(hex_string)
      .map(Fingerprint)
      .map_err(|e| format!("{:?}", e))
  }

  pub fn as_bytes(&self) -> &[u8; FINGERPRINT_SIZE] {
    &self.0
  }

  pub fn to_hex(&self) -> String {
    let mut s = String::new();
    for &byte in &self.0 {
      fmt::Write::write_fmt(&mut s, format_args!("{:02x}", byte)).unwrap();
    }
    s
  }
}

impl fmt::Display for Fingerprint {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "{}", self.to_hex())
  }
}

impl fmt::Debug for Fingerprint {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "Fingerprint<{}>", self.to_hex())
  }
}

impl AsRef<[u8]> for Fingerprint {
  fn as_ref(&self) -> &[u8] {
    &self.0[..]
  }
}

impl Serialize for Fingerprint {
  fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
  where
    S: Serializer,
  {
    serializer.serialize_str(self.to_hex().as_str())
  }
}

///
/// A Digest is a fingerprint, as well as the size in bytes of the plaintext for which that is the
/// fingerprint.
///
/// It is equivalent to a Bazel Remote Execution Digest, but without the overhead (and awkward API)
/// of needing to create an entire protobuf to pass around the two fields.
///
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Digest(pub Fingerprint, pub usize);

impl Serialize for Digest {
  fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
  where
    S: Serializer,
  {
    let mut obj = serializer.serialize_struct("digest", 2)?;
    obj.serialize_field("fingerprint", &self.0)?;
    obj.serialize_field("size_bytes", &self.1)?;
    obj.end()
  }
}

///
/// A Write instance that fingerprints all data that passes through it.
///
pub struct WriterHasher<W: Write> {
  hasher: Sha256,
  byte_count: usize,
  inner: W,
}

impl<W: Write> WriterHasher<W> {
  pub fn new(inner: W) -> WriterHasher<W> {
    WriterHasher {
      hasher: Sha256::default(),
      byte_count: 0,
      inner: inner,
    }
  }

  ///
  /// Returns the result of fingerprinting this stream, and Drops the stream.
  ///
  pub fn finish(self) -> (Digest, W) {
    (
      Digest(
        Fingerprint::from_bytes_unsafe(&self.hasher.fixed_result()),
        self.byte_count,
      ),
      self.inner,
    )
  }
}

impl<W: Write> Write for WriterHasher<W> {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
    let written = self.inner.write(buf)?;
    // Hash the bytes that were successfully written.
    self.hasher.input(&buf[0..written]);
    self.byte_count += written;
    Ok(written)
  }

  fn flush(&mut self) -> io::Result<()> {
    self.inner.flush()
  }
}

#[cfg(test)]
mod fingerprint_tests {
  use self::serde_test::{assert_ser_tokens, Token};
  use super::Fingerprint;
  use serde_test;

  #[test]
  fn from_bytes_unsafe() {
    assert_eq!(
      Fingerprint::from_bytes_unsafe(&[
        0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab,
        0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab,
        0xab, 0xab,
      ],),
      Fingerprint([0xab; 32])
    );
  }

  #[test]
  fn from_hex_string() {
    assert_eq!(
      Fingerprint::from_hex_string(
        "0123456789abcdefFEDCBA98765432100000000000000000ffFFfFfFFfFfFFff",
      )
      .unwrap(),
      Fingerprint([
        0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef, 0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32,
        0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff,
      ],)
    )
  }

  #[test]
  fn from_hex_string_not_long_enough() {
    Fingerprint::from_hex_string("abcd").expect_err("Want err");
  }

  #[test]
  fn from_hex_string_too_long() {
    Fingerprint::from_hex_string(
      "0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0",
    )
    .expect_err("Want err");
  }

  #[test]
  fn from_hex_string_invalid_chars() {
    Fingerprint::from_hex_string(
      "Q123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF",
    )
    .expect_err("Want err");
  }

  #[test]
  fn to_hex() {
    assert_eq!(
      Fingerprint([
        0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef, 0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32,
        0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff,
      ],)
      .to_hex(),
      "0123456789abcdeffedcba98765432100000000000000000ffffffffffffffff".to_lowercase()
    )
  }

  #[test]
  fn display() {
    let hex = "0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF";
    assert_eq!(
      Fingerprint::from_hex_string(hex).unwrap().to_hex(),
      hex.to_lowercase()
    )
  }

  #[test]
  fn serialize_to_str() {
    let fingerprint = Fingerprint([
      0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef, 0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32,
      0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
      0xff, 0xff,
    ]);
    assert_ser_tokens(
      &fingerprint,
      &[Token::Str(
        "0123456789abcdeffedcba98765432100000000000000000ffffffffffffffff",
      )],
    );
  }
}

#[cfg(test)]
mod digest_tests {
  use self::serde_test::{assert_ser_tokens, Token};
  use super::Digest;
  use super::Fingerprint;
  use serde_test;

  #[test]
  fn serialize_to_str() {
    let digest = Digest(
      Fingerprint::from_hex_string(
        "0123456789abcdeffedcba98765432100000000000000000ffffffffffffffff",
      )
      .unwrap(),
      1,
    );
    assert_ser_tokens(
      &digest,
      &[
        Token::Struct {
          name: "digest",
          len: 2,
        },
        Token::Str("fingerprint"),
        Token::Str("0123456789abcdeffedcba98765432100000000000000000ffffffffffffffff"),
        Token::Str("size_bytes"),
        Token::U64(1),
        Token::StructEnd,
      ],
    );
  }
}

#[cfg(test)]
mod hasher_tests {
  use std;

  #[test]
  fn hashes() {
    let mut src = "meep".as_bytes();

    let dst = Vec::with_capacity(10);
    let mut hasher = super::WriterHasher::new(dst);
    assert_eq!(std::io::copy(&mut src, &mut hasher).unwrap(), 4);
    let want = (
      super::Digest(
        super::Fingerprint::from_hex_string(
          "23e92dfba8fb0c93cfba31ad2962b4e35a47054296d1d375d7f7e13e0185de7a",
        )
        .unwrap(),
        4,
      ),
      "meep".as_bytes().to_vec(),
    );
    assert_eq!(hasher.finish(), want);
  }
}
<commit_msg>Digest can be deserialized as well as serialized (#8497)<commit_after>// Copyright 2017 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

#![deny(warnings)]
// Enable all clippy lints except for many of the pedantic ones. It's a shame this needs to be copied and pasted across crates, but there doesn't appear to be a way to include inner attributes from a common source.
#![deny(
  clippy::all,
  clippy::default_trait_access,
  clippy::expl_impl_clone_on_copy,
  clippy::if_not_else,
  clippy::needless_continue,
  clippy::single_match_else,
  clippy::unseparated_literal_suffix,
  clippy::used_underscore_binding
)]
// It is often more clear to show that nothing is being moved.
#![allow(clippy::match_ref_pats)]
// Subjective style.
#![allow(
  clippy::len_without_is_empty,
  clippy::redundant_field_names,
  clippy::too_many_arguments
)]
// Default isn't as big a deal as people seem to think it is.
#![allow(clippy::new_without_default, clippy::new_ret_no_self)]
// Arc<Mutex> can be more clear than needing to grok Orderings:
#![allow(clippy::mutex_atomic)]

use hex;

use digest::{Digest as DigestTrait, FixedOutput};
use serde::ser::{Serialize, SerializeStruct, Serializer};
use serde::{Deserialize, Deserializer};
use sha2::Sha256;

use serde::de::{MapAccess, Visitor};
use serde::export::fmt::Error;
use serde::export::Formatter;
use std::fmt;
use std::io::{self, Write};

pub const EMPTY_FINGERPRINT: Fingerprint = Fingerprint([
  0xe3, 0xb0, 0xc4, 0x42, 0x98, 0xfc, 0x1c, 0x14, 0x9a, 0xfb, 0xf4, 0xc8, 0x99, 0x6f, 0xb9, 0x24,
  0x27, 0xae, 0x41, 0xe4, 0x64, 0x9b, 0x93, 0x4c, 0xa4, 0x95, 0x99, 0x1b, 0x78, 0x52, 0xb8, 0x55,
]);
pub const EMPTY_DIGEST: Digest = Digest(EMPTY_FINGERPRINT, 0);

const FINGERPRINT_SIZE: usize = 32;

#[derive(Clone, Copy, Eq, Hash, PartialEq, Ord, PartialOrd)]
pub struct Fingerprint(pub [u8; FINGERPRINT_SIZE]);

impl Fingerprint {
  pub fn from_bytes_unsafe(bytes: &[u8]) -> Fingerprint {
    if bytes.len() != FINGERPRINT_SIZE {
      panic!(
        "Input value was not a fingerprint; had length: {}",
        bytes.len()
      );
    }

    let mut fingerprint = [0; FINGERPRINT_SIZE];
    fingerprint.clone_from_slice(&bytes[0..FINGERPRINT_SIZE]);
    Fingerprint(fingerprint)
  }

  pub fn from_hex_string(hex_string: &str) -> Result<Fingerprint, String> {
    <[u8; FINGERPRINT_SIZE] as hex::FromHex>::from_hex(hex_string)
      .map(Fingerprint)
      .map_err(|e| format!("{:?}", e))
  }

  pub fn as_bytes(&self) -> &[u8; FINGERPRINT_SIZE] {
    &self.0
  }

  pub fn to_hex(&self) -> String {
    let mut s = String::new();
    for &byte in &self.0 {
      fmt::Write::write_fmt(&mut s, format_args!("{:02x}", byte)).unwrap();
    }
    s
  }
}

impl fmt::Display for Fingerprint {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "{}", self.to_hex())
  }
}

impl fmt::Debug for Fingerprint {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "Fingerprint<{}>", self.to_hex())
  }
}

impl AsRef<[u8]> for Fingerprint {
  fn as_ref(&self) -> &[u8] {
    &self.0[..]
  }
}

impl Serialize for Fingerprint {
  fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
  where
    S: Serializer,
  {
    serializer.serialize_str(self.to_hex().as_str())
  }
}

impl<'de> Deserialize<'de> for Fingerprint {
  fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
  where
    D: Deserializer<'de>,
  {
    struct FingerprintVisitor;

    impl<'de> Visitor<'de> for FingerprintVisitor {
      type Value = Fingerprint;

      fn expecting(&self, formatter: &mut Formatter) -> Result<(), Error> {
        formatter.write_str("struct Fingerprint")
      }

      fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
      where
        E: serde::de::Error,
      {
        Fingerprint::from_hex_string(v).map_err(|err| {
          serde::de::Error::invalid_value(
            serde::de::Unexpected::Str(&format!("{:?}: {}", v, err)),
            &format!("A hex representation of a {} byte value", FINGERPRINT_SIZE).as_str(),
          )
        })
      }
    }

    deserializer.deserialize_string(FingerprintVisitor)
  }
}

///
/// A Digest is a fingerprint, as well as the size in bytes of the plaintext for which that is the
/// fingerprint.
///
/// It is equivalent to a Bazel Remote Execution Digest, but without the overhead (and awkward API)
/// of needing to create an entire protobuf to pass around the two fields.
///
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Digest(pub Fingerprint, pub usize);

impl Serialize for Digest {
  fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
  where
    S: Serializer,
  {
    let mut obj = serializer.serialize_struct("digest", 2)?;
    obj.serialize_field("fingerprint", &self.0)?;
    obj.serialize_field("size_bytes", &self.1)?;
    obj.end()
  }
}

#[derive(Deserialize)]
#[serde(field_identifier, rename_all = "snake_case")]
enum Field {
  Fingerprint,
  SizeBytes,
}

impl<'de> Deserialize<'de> for Digest {
  fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
  where
    D: Deserializer<'de>,
  {
    struct DigestVisitor;

    impl<'de> Visitor<'de> for DigestVisitor {
      type Value = Digest;

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

      fn visit_map<V>(self, mut map: V) -> Result<Digest, V::Error>
      where
        V: MapAccess<'de>,
      {
        use serde::de;

        let mut fingerprint = None;
        let mut size_bytes = None;
        while let Some(key) = map.next_key()? {
          match key {
            Field::Fingerprint => {
              if fingerprint.is_some() {
                return Err(de::Error::duplicate_field("fingerprint"));
              }
              fingerprint = Some(map.next_value()?);
            }
            Field::SizeBytes => {
              if size_bytes.is_some() {
                return Err(de::Error::duplicate_field("size_bytes"));
              }
              size_bytes = Some(map.next_value()?);
            }
          }
        }
        let fingerprint = fingerprint.ok_or_else(|| de::Error::missing_field("fingerprint"))?;
        let size_bytes = size_bytes.ok_or_else(|| de::Error::missing_field("size_bytes"))?;
        Ok(Digest(fingerprint, size_bytes))
      }
    }

    const FIELDS: &[&str] = &["fingerprint", "size_bytes"];
    deserializer.deserialize_struct("digest", FIELDS, DigestVisitor)
  }
}

///
/// A Write instance that fingerprints all data that passes through it.
///
pub struct WriterHasher<W: Write> {
  hasher: Sha256,
  byte_count: usize,
  inner: W,
}

impl<W: Write> WriterHasher<W> {
  pub fn new(inner: W) -> WriterHasher<W> {
    WriterHasher {
      hasher: Sha256::default(),
      byte_count: 0,
      inner: inner,
    }
  }

  ///
  /// Returns the result of fingerprinting this stream, and Drops the stream.
  ///
  pub fn finish(self) -> (Digest, W) {
    (
      Digest(
        Fingerprint::from_bytes_unsafe(&self.hasher.fixed_result()),
        self.byte_count,
      ),
      self.inner,
    )
  }
}

impl<W: Write> Write for WriterHasher<W> {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
    let written = self.inner.write(buf)?;
    // Hash the bytes that were successfully written.
    self.hasher.input(&buf[0..written]);
    self.byte_count += written;
    Ok(written)
  }

  fn flush(&mut self) -> io::Result<()> {
    self.inner.flush()
  }
}

#[cfg(test)]
mod fingerprint_tests {
  use self::serde_test::{assert_ser_tokens, Token};
  use super::Fingerprint;
  use serde_test;

  #[test]
  fn from_bytes_unsafe() {
    assert_eq!(
      Fingerprint::from_bytes_unsafe(&[
        0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab,
        0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab, 0xab,
        0xab, 0xab,
      ],),
      Fingerprint([0xab; 32])
    );
  }

  #[test]
  fn from_hex_string() {
    assert_eq!(
      Fingerprint::from_hex_string(
        "0123456789abcdefFEDCBA98765432100000000000000000ffFFfFfFFfFfFFff",
      )
      .unwrap(),
      Fingerprint([
        0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef, 0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32,
        0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff,
      ],)
    )
  }

  #[test]
  fn from_hex_string_not_long_enough() {
    Fingerprint::from_hex_string("abcd").expect_err("Want err");
  }

  #[test]
  fn from_hex_string_too_long() {
    Fingerprint::from_hex_string(
      "0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0",
    )
    .expect_err("Want err");
  }

  #[test]
  fn from_hex_string_invalid_chars() {
    Fingerprint::from_hex_string(
      "Q123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF",
    )
    .expect_err("Want err");
  }

  #[test]
  fn to_hex() {
    assert_eq!(
      Fingerprint([
        0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef, 0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32,
        0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff,
      ],)
      .to_hex(),
      "0123456789abcdeffedcba98765432100000000000000000ffffffffffffffff".to_lowercase()
    )
  }

  #[test]
  fn display() {
    let hex = "0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF";
    assert_eq!(
      Fingerprint::from_hex_string(hex).unwrap().to_hex(),
      hex.to_lowercase()
    )
  }

  #[test]
  fn serialize_to_str() {
    let fingerprint = Fingerprint([
      0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef, 0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32,
      0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
      0xff, 0xff,
    ]);
    assert_ser_tokens(
      &fingerprint,
      &[Token::Str(
        "0123456789abcdeffedcba98765432100000000000000000ffffffffffffffff",
      )],
    );
  }
}

#[cfg(test)]
mod digest_tests {
  use self::serde_test::{assert_tokens, Token};
  use super::Digest;
  use super::Fingerprint;
  use serde_test;

  #[test]
  fn serialize_and_deserialize() {
    let digest = Digest(
      Fingerprint::from_hex_string(
        "0123456789abcdeffedcba98765432100000000000000000ffffffffffffffff",
      )
      .unwrap(),
      1,
    );
    assert_tokens(
      &digest,
      &[
        Token::Struct {
          name: "digest",
          len: 2,
        },
        Token::Str("fingerprint"),
        Token::Str("0123456789abcdeffedcba98765432100000000000000000ffffffffffffffff"),
        Token::Str("size_bytes"),
        Token::U64(1),
        Token::StructEnd,
      ],
    );
  }
}

#[cfg(test)]
mod hasher_tests {
  use std;

  #[test]
  fn hashes() {
    let mut src = "meep".as_bytes();

    let dst = Vec::with_capacity(10);
    let mut hasher = super::WriterHasher::new(dst);
    assert_eq!(std::io::copy(&mut src, &mut hasher).unwrap(), 4);
    let want = (
      super::Digest(
        super::Fingerprint::from_hex_string(
          "23e92dfba8fb0c93cfba31ad2962b4e35a47054296d1d375d7f7e13e0185de7a",
        )
        .unwrap(),
        4,
      ),
      "meep".as_bytes().to_vec(),
    );
    assert_eq!(hasher.finish(), want);
  }
}
<|endoftext|>


<commit_before>// ignore-windows failing on win32 bot
// ignore-android: FIXME(#10381)
// compile-flags:-g
// min-gdb-version: 7.7
// min-lldb-version: 310

// === GDB TESTS ===================================================================================

// gdb-command: run

// gdb-command: print empty_string
// gdb-check:$1 = ""

// gdb-command: print empty_str
// gdb-check:$2 = ""

// === LLDB TESTS ==================================================================================

// lldb-command: run

// lldb-command: fr v empty_string
// lldb-check:[...]empty_string = ""

// lldb-command: fr v empty_str
// lldb-check:[...]empty_str = ""

fn main() {
    let empty_string = String::new();

    let empty_str = "";

    zzz(); // #break
}

fn zzz() {}
<commit_msg>Rollup merge of #69223 - spunit262:empty-string-gdb-version, r=Dylan-DPC<commit_after>// ignore-windows failing on win32 bot
// ignore-android: FIXME(#10381)
// compile-flags:-g
// min-gdb-version: 7.7
// ignore-gdb-version: 7.11.90 - 8.0.9
// min-lldb-version: 310

// === GDB TESTS ===================================================================================

// gdb-command: run

// gdb-command: print empty_string
// gdb-check:$1 = ""

// gdb-command: print empty_str
// gdb-check:$2 = ""

// === LLDB TESTS ==================================================================================

// lldb-command: run

// lldb-command: fr v empty_string
// lldb-check:[...]empty_string = ""

// lldb-command: fr v empty_str
// lldb-check:[...]empty_str = ""

fn main() {
    let empty_string = String::new();

    let empty_str = "";

    zzz(); // #break
}

fn zzz() {}
<|endoftext|>


<commit_before>use pyo3;
use pyo3::types::{PyBytes};
use pyo3::prelude::*;
use failure::{Error};

use lancelot::workspace;
use lancelot::arch::{RVA};

#[pymodule]
fn pylancelot(_py: Python, m: &PyModule) -> PyResult<()> {
    #[pyfn(m, "from_bytes")]
    /// from_bytes(filename, buf, /)
    /// --
    ///
    /// Create a workspace from the given bytes.
    ///
    /// Args:
    ///   filename (str): the source filename
    ///   buf (bytes): the bytes containing a PE, shellcode, etc.
    ///
    /// Raises:
    ///   ValueError: if failed to create the workspace. TODO: more specific.
    ///
    fn from_bytes(_py: Python, filename: String, buf: &PyBytes) -> PyResult<PyWorkspace> {
        let ws = match workspace::Workspace::from_bytes(&filename, buf.as_bytes()).load() {
            Err(_) => return Err(pyo3::exceptions::ValueError::py_err("failed to create workspace")),
            Ok(ws) => ws,
        };
        Ok(PyWorkspace{ws})
    }

    #[pyclass]
    pub struct PyWorkspace {
        ws: workspace::Workspace
    }

    #[pyclass]
    pub struct PySection {
        #[pyo3(get)]
        pub addr: i64,
        #[pyo3(get)]
        pub length: u64,
        #[pyo3(get)]
        pub perms: u8,
        #[pyo3(get)]
        pub name: String,
    }

    m.add("PERM_NONE", lancelot::loader::Permissions::empty().bits())?;
    m.add("PERM_R", lancelot::loader::Permissions::R.bits())?;
    m.add("PERM_W", lancelot::loader::Permissions::W.bits())?;
    m.add("PERM_X", lancelot::loader::Permissions::X.bits())?;
    m.add("PERM_RW", lancelot::loader::Permissions::RW.bits())?;
    m.add("PERM_RX", lancelot::loader::Permissions::RX.bits())?;
    m.add("PERM_RWX", lancelot::loader::Permissions::RWX.bits())?;

    #[pyclass]
    pub struct PyXref {
        #[pyo3(get)]
        pub src: i64,
        #[pyo3(get)]
        pub dst: i64,
        /// one of XREF_* constants
        #[pyo3(get)]
        pub typ: u8,
    }

    // keep in sync with pylancelot::PyWorkspace::translate_xref
    m.add("XREF_FALLTHROUGH", 1)?;
    m.add("XREF_CALL", 2)?;
    m.add("XREF_UNCONDITIONAL_JUMP", 3)?;
    m.add("XREF_CONDITIONAL_JUMP", 4)?;
    m.add("XREF_CONDITIONAL_MOVE", 5)?;

    #[pymethods]
    impl PyWorkspace {
        #[getter]
        /// filename(self, /)
        /// --
        ///
        /// Fetch the filename.
        /// ```
        pub fn filename(&self) -> PyResult<String> {
            Ok(self.ws.filename.clone())
        }

        #[getter]
        /// loader(self, /)
        /// --
        ///
        /// Fetch the name of the loader used to create the workspace.
        pub fn loader(&self) -> PyResult<String> {
            Ok(self.ws.loader.get_name())
        }

        #[getter]
        /// base_address(self, /)
        /// --
        ///
        /// Fetch the base address to which the module was loaded.
        pub fn base_address(&self) -> PyResult<u64> {
            Ok(self.ws.module.base_address.into())
        }

        #[getter]
        pub fn sections(&self) -> PyResult<Vec<PySection>> {
            Ok(self.ws.module.sections.iter()
                .map(|section| PySection{
                    addr: section.addr.into(),
                    length: section.buf.len() as u64,
                    perms: section.perms.bits(),
                    name: section.name.clone(),
                })
                .collect())
        }

        #[getter]
        pub fn functions(&self) -> PyResult<Vec<i64>> {
            Ok(self.ws.get_functions().map(|&rva| rva.into()).collect())
        }

        pub fn get_xrefs_from(&self, rva: i64) -> PyResult<Vec<PyXref>> {
            // doesn't currently contain the fallthrough xref,
            // because in get_xrefs_to, its more difficult to compute.
            let mut xrefs = match self.ws.analysis.flow.xrefs.from.get(&RVA::from(rva)) {
                Some(xrefs) => {
                    xrefs.iter().map(|x| PyWorkspace::translate_xref(x)).collect()
                },
                None => vec![],
            };

            if let Some(meta) = self.ws.get_meta(RVA::from(rva)) {
                if meta.does_fallthrough() {
                    // we have to do some extra legwork here to correctly handle long insns
                    let length = match meta.get_insn_length() {
                        Ok(length) => length,
                        Err(lancelot::flowmeta::Error::LongInstruction) => {
                            match self.ws.read_insn(RVA::from(rva)) {
                                Ok(insn) => insn.length,
                                Err(_) => 0,
                            }
                        },
                        Err(lancelot::flowmeta::Error::NotAnInstruction) => 0,
                    };

                    if length > 0 {
                        xrefs.push(PyXref {
                            src: rva,
                            dst: rva + length as i64,
                            typ: 1,
                        });
                    }
                }
            }

            Ok(xrefs)
        }

        pub fn get_xrefs_to(&self, rva: i64) -> PyResult<Vec<PyXref>> {
            Ok(match self.ws.analysis.flow.xrefs.to.get(&RVA::from(rva)) {
                Some(xrefs) => {
                    xrefs.iter().map(|x| PyWorkspace::translate_xref(x)).collect()
                },
                None => vec![],
            })
        }

        /// probe(self, rva, length=1, /)
        /// --
        ///
        /// Is the given address mapped?
        pub fn probe(&self, rva: i64, length: Option<usize>) -> PyResult<bool> {
            match length {
                Some(length) => Ok(self.ws.probe(RVA::from(rva), length)),
                None =>  Ok(self.ws.probe(RVA::from(rva), 1)),
            }
        }

        /// read_bytes(self, rva, length, /)
        /// --
        ///
        /// Read bytes from the given memory address.
        ///
        /// raises:
        ///   - LookupError: if the address is not mapped, or the length overruns.
        pub fn read_bytes<'p>(&self, py: Python<'p>, rva: i64, length: usize) -> PyResult<&'p pyo3::types::PyBytes> {
            match self.ws.read_bytes(RVA::from(rva), length) {
                Ok(buf) => Ok(pyo3::types::PyBytes::new(py, buf)),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_u8(self, rva, /)
        /// --
        ///
        pub fn read_u8(&self, rva: i64) -> PyResult<u8> {
            match self.ws.read_u8(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_u16(self, rva, /)
        /// --
        ///
        pub fn read_u16(&self, rva: i64) -> PyResult<u16> {
            match self.ws.read_u16(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_u32(self, rva, /)
        /// --
        ///
        pub fn read_u32(&self, rva: i64) -> PyResult<u32> {
            match self.ws.read_u32(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_u64(self, rva, /)
        /// --
        ///
        pub fn read_u64(&self, rva: i64) -> PyResult<u64> {
            match self.ws.read_u64(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_rva(self, rva, /)
        /// --
        ///
        /// Read the integer relative virtual address at the given address.
        /// The size of the RVA is dependent upon the bitness of the current workspace.
        pub fn read_rva(&self, rva: i64) -> PyResult<i64> {
            match self.ws.read_rva(RVA::from(rva)) {
                Ok(b) => Ok(b.into()),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_va(self, rva, /)
        /// --
        ///
        /// Read the integer virtual address at the given address.
        /// The size of the VA is dependent upon the bitness of the current workspace.
        pub fn read_va(&self, rva: i64) -> PyResult<u64> {
            match self.ws.read_va(RVA::from(rva)) {
                Ok(b) => Ok(b.into()),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }
    }

    impl PyWorkspace {
        /// translate from a failure::Error returned by a Workspace
        /// while reading an element from memory into a pyo3/Python error
        /// corresponding to out-of-bounds-access.
        fn translate_buffer_error(e: Error) -> PyErr {
            match e.downcast::<lancelot::workspace::WorkspaceError>() {
                // we have to use the documentation for the associated routine,
                // such as `read_u8`, to determine what to inspect here.
                Ok(lancelot::workspace::WorkspaceError::InvalidAddress) => {
                    pyo3::exceptions::LookupError::py_err("invalid address")
                },
                Ok(lancelot::workspace::WorkspaceError::BufferOverrun) => {
                    pyo3::exceptions::LookupError::py_err("buffer overrun")
                }
                Ok(_) => {
                    // this should never be hit
                    pyo3::exceptions::RuntimeError::py_err("unexpected error")
                },
                Err(_) => {
                    // this should never be hit
                    pyo3::exceptions::RuntimeError::py_err("unexpected error")
                },
            }
        }

        fn translate_xref(x: &lancelot::xref::Xref) -> PyXref {
            PyXref {
                src: x.src.into(),
                dst: x.dst.into(),
                typ: match x.typ {
                    lancelot::xref::XrefType::Fallthrough => 1,
                    lancelot::xref::XrefType::Call => 2,
                    lancelot::xref::XrefType::UnconditionalJump => 3,
                    lancelot::xref::XrefType::ConditionalJump => 4,
                    lancelot::xref::XrefType::ConditionalMove => 5,
                }
            }
        }
    }

    #[pyproto]
    impl pyo3::class::basic::PyObjectProtocol for PyWorkspace {
        fn __str__(&self) -> PyResult<String> {
            PyWorkspace::__repr__(self)
        }

        fn __repr__(&self) -> PyResult<String> {
            Ok(format!("PyWorkspace(filename: {} loader: {})",
                self.ws.filename.clone(),
                self.ws.loader.get_name(),
            ))
        }
    }

    #[pyproto]
    impl pyo3::class::basic::PyObjectProtocol for PySection {
        fn __str__(&self) -> PyResult<String> {
            PySection::__repr__(self)
        }

        fn __repr__(&self) -> PyResult<String> {
            Ok(format!("PySection(addr: {:#x} length: {:#x} perms: {:#x} name: {})",
                self.addr,
                self.length,
                self.perms,
                self.name,
            ))
        }
    }

    #[pyproto]
    impl pyo3::class::basic::PyObjectProtocol for PyXref {
        fn __str__(&self) -> PyResult<String> {
            PyXref::__repr__(self)
        }

        fn __repr__(&self) -> PyResult<String> {
            Ok(format!("PyXref(src: {:#x} dst: {:#x} type: {})",
                self.src,
                self.dst,
                match self.typ {
                    1 => "XREF_FALLTHROUGH",
                    2 => "XREF_CALL",
                    3 => "XREF_UNCONDITIONAL_JUMP",
                    4 => "XREF_CONDITIONAL_JUMP",
                    5 => "XREF_CONDITIONAL_MOVE",
                    _ => unreachable!(),
                },
            ))
        }
    }

    Ok(())
}

<commit_msg>pylancelot: add doc<commit_after>use pyo3;
use pyo3::types::{PyBytes};
use pyo3::prelude::*;
use failure::{Error};

use lancelot::workspace;
use lancelot::arch::{RVA};

#[pymodule]
fn pylancelot(_py: Python, m: &PyModule) -> PyResult<()> {
    #[pyfn(m, "from_bytes")]
    /// from_bytes(filename, buf, /)
    /// --
    ///
    /// Create a workspace from the given bytes.
    ///
    /// Args:
    ///   filename (str): the source filename
    ///   buf (bytes): the bytes containing a PE, shellcode, etc.
    ///
    /// Raises:
    ///   ValueError: if failed to create the workspace. TODO: more specific.
    ///
    fn from_bytes(_py: Python, filename: String, buf: &PyBytes) -> PyResult<PyWorkspace> {
        let ws = match workspace::Workspace::from_bytes(&filename, buf.as_bytes()).load() {
            Err(_) => return Err(pyo3::exceptions::ValueError::py_err("failed to create workspace")),
            Ok(ws) => ws,
        };
        Ok(PyWorkspace{ws})
    }

    #[pyclass]
    pub struct PyWorkspace {
        ws: workspace::Workspace
    }

    #[pyclass]
    pub struct PySection {
        #[pyo3(get)]
        pub addr: i64,
        #[pyo3(get)]
        pub length: u64,
        #[pyo3(get)]
        pub perms: u8,
        #[pyo3(get)]
        pub name: String,
    }

    m.add("PERM_NONE", lancelot::loader::Permissions::empty().bits())?;
    m.add("PERM_R", lancelot::loader::Permissions::R.bits())?;
    m.add("PERM_W", lancelot::loader::Permissions::W.bits())?;
    m.add("PERM_X", lancelot::loader::Permissions::X.bits())?;
    m.add("PERM_RW", lancelot::loader::Permissions::RW.bits())?;
    m.add("PERM_RX", lancelot::loader::Permissions::RX.bits())?;
    m.add("PERM_RWX", lancelot::loader::Permissions::RWX.bits())?;

    #[pyclass]
    pub struct PyXref {
        #[pyo3(get)]
        pub src: i64,
        #[pyo3(get)]
        pub dst: i64,
        /// one of XREF_* constants
        #[pyo3(get)]
        pub typ: u8,
    }

    // keep in sync with pylancelot::PyWorkspace::translate_xref
    m.add("XREF_FALLTHROUGH", 1)?;
    m.add("XREF_CALL", 2)?;
    m.add("XREF_UNCONDITIONAL_JUMP", 3)?;
    m.add("XREF_CONDITIONAL_JUMP", 4)?;
    m.add("XREF_CONDITIONAL_MOVE", 5)?;

    #[pymethods]
    impl PyWorkspace {
        #[getter]
        /// filename(self, /)
        /// --
        ///
        /// Fetch the filename.
        /// ```
        pub fn filename(&self) -> PyResult<String> {
            Ok(self.ws.filename.clone())
        }

        #[getter]
        /// loader(self, /)
        /// --
        ///
        /// Fetch the name of the loader used to create the workspace.
        pub fn loader(&self) -> PyResult<String> {
            Ok(self.ws.loader.get_name())
        }

        #[getter]
        /// base_address(self, /)
        /// --
        ///
        /// Fetch the base address to which the module was loaded.
        pub fn base_address(&self) -> PyResult<u64> {
            Ok(self.ws.module.base_address.into())
        }

        #[getter]
        /// sections(self, /)
        /// --
        ///
        /// Fetch the sections loaded from the module.
        pub fn sections(&self) -> PyResult<Vec<PySection>> {
            Ok(self.ws.module.sections.iter()
                .map(|section| PySection{
                    addr: section.addr.into(),
                    length: section.buf.len() as u64,
                    perms: section.perms.bits(),
                    name: section.name.clone(),
                })
                .collect())
        }

        #[getter]
        /// functions(self, /)
        /// --
        ///
        /// Fetch the addresses of functions discovered during analysis.
        pub fn functions(&self) -> PyResult<Vec<i64>> {
            Ok(self.ws.get_functions().map(|&rva| rva.into()).collect())
        }

        /// get_xrefs_from(self, rva, /)
        /// --
        ///
        /// Fetch the xrefs flowing from the given address.
        pub fn get_xrefs_from(&self, rva: i64) -> PyResult<Vec<PyXref>> {
            let mut xrefs = match self.ws.analysis.flow.xrefs.from.get(&RVA::from(rva)) {
                Some(xrefs) => {
                    xrefs.iter().map(|x| PyWorkspace::translate_xref(x)).collect()
                },
                None => vec![],
            };

            if let Some(meta) = self.ws.get_meta(RVA::from(rva)) {
                if meta.does_fallthrough() {
                    // we have to do some extra legwork here to correctly handle long insns
                    let length = match meta.get_insn_length() {
                        Ok(length) => length,
                        Err(lancelot::flowmeta::Error::LongInstruction) => {
                            match self.ws.read_insn(RVA::from(rva)) {
                                Ok(insn) => insn.length,
                                Err(_) => 0,
                            }
                        },
                        Err(lancelot::flowmeta::Error::NotAnInstruction) => 0,
                    };

                    if length > 0 {
                        xrefs.push(PyXref {
                            src: rva,
                            dst: rva + length as i64,
                            typ: 1,  // XREF_TYPE_FALLTHROUGH
                        });
                    }
                }
            }

            Ok(xrefs)
        }

        pub fn get_xrefs_to(&self, rva: i64) -> PyResult<Vec<PyXref>> {
            Ok(match self.ws.analysis.flow.xrefs.to.get(&RVA::from(rva)) {
                Some(xrefs) => {
                    xrefs.iter().map(|x| PyWorkspace::translate_xref(x)).collect()
                },
                None => vec![],
            })
        }

        /// probe(self, rva, length=1, /)
        /// --
        ///
        /// Is the given address mapped?
        pub fn probe(&self, rva: i64, length: Option<usize>) -> PyResult<bool> {
            match length {
                Some(length) => Ok(self.ws.probe(RVA::from(rva), length)),
                None =>  Ok(self.ws.probe(RVA::from(rva), 1)),
            }
        }

        /// read_bytes(self, rva, length, /)
        /// --
        ///
        /// Read bytes from the given memory address.
        ///
        /// raises:
        ///   - LookupError: if the address is not mapped, or the length overruns.
        pub fn read_bytes<'p>(&self, py: Python<'p>, rva: i64, length: usize) -> PyResult<&'p pyo3::types::PyBytes> {
            match self.ws.read_bytes(RVA::from(rva), length) {
                Ok(buf) => Ok(pyo3::types::PyBytes::new(py, buf)),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_u8(self, rva, /)
        /// --
        ///
        pub fn read_u8(&self, rva: i64) -> PyResult<u8> {
            match self.ws.read_u8(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_u16(self, rva, /)
        /// --
        ///
        pub fn read_u16(&self, rva: i64) -> PyResult<u16> {
            match self.ws.read_u16(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_u32(self, rva, /)
        /// --
        ///
        pub fn read_u32(&self, rva: i64) -> PyResult<u32> {
            match self.ws.read_u32(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_u64(self, rva, /)
        /// --
        ///
        pub fn read_u64(&self, rva: i64) -> PyResult<u64> {
            match self.ws.read_u64(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_rva(self, rva, /)
        /// --
        ///
        /// Read the integer relative virtual address at the given address.
        /// The size of the RVA is dependent upon the bitness of the current workspace.
        pub fn read_rva(&self, rva: i64) -> PyResult<i64> {
            match self.ws.read_rva(RVA::from(rva)) {
                Ok(b) => Ok(b.into()),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_va(self, rva, /)
        /// --
        ///
        /// Read the integer virtual address at the given address.
        /// The size of the VA is dependent upon the bitness of the current workspace.
        pub fn read_va(&self, rva: i64) -> PyResult<u64> {
            match self.ws.read_va(RVA::from(rva)) {
                Ok(b) => Ok(b.into()),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }
    }

    impl PyWorkspace {
        /// translate from a failure::Error returned by a Workspace
        /// while reading an element from memory into a pyo3/Python error
        /// corresponding to out-of-bounds-access.
        fn translate_buffer_error(e: Error) -> PyErr {
            match e.downcast::<lancelot::workspace::WorkspaceError>() {
                // we have to use the documentation for the associated routine,
                // such as `read_u8`, to determine what to inspect here.
                Ok(lancelot::workspace::WorkspaceError::InvalidAddress) => {
                    pyo3::exceptions::LookupError::py_err("invalid address")
                },
                Ok(lancelot::workspace::WorkspaceError::BufferOverrun) => {
                    pyo3::exceptions::LookupError::py_err("buffer overrun")
                }
                Ok(_) => {
                    // this should never be hit
                    pyo3::exceptions::RuntimeError::py_err("unexpected error")
                },
                Err(_) => {
                    // this should never be hit
                    pyo3::exceptions::RuntimeError::py_err("unexpected error")
                },
            }
        }

        fn translate_xref(x: &lancelot::xref::Xref) -> PyXref {
            PyXref {
                src: x.src.into(),
                dst: x.dst.into(),
                typ: match x.typ {
                    lancelot::xref::XrefType::Fallthrough => 1,
                    lancelot::xref::XrefType::Call => 2,
                    lancelot::xref::XrefType::UnconditionalJump => 3,
                    lancelot::xref::XrefType::ConditionalJump => 4,
                    lancelot::xref::XrefType::ConditionalMove => 5,
                }
            }
        }
    }

    #[pyproto]
    impl pyo3::class::basic::PyObjectProtocol for PyWorkspace {
        fn __str__(&self) -> PyResult<String> {
            PyWorkspace::__repr__(self)
        }

        fn __repr__(&self) -> PyResult<String> {
            Ok(format!("PyWorkspace(filename: {} loader: {})",
                self.ws.filename.clone(),
                self.ws.loader.get_name(),
            ))
        }
    }

    #[pyproto]
    impl pyo3::class::basic::PyObjectProtocol for PySection {
        fn __str__(&self) -> PyResult<String> {
            PySection::__repr__(self)
        }

        fn __repr__(&self) -> PyResult<String> {
            Ok(format!("PySection(addr: {:#x} length: {:#x} perms: {:#x} name: {})",
                self.addr,
                self.length,
                self.perms,
                self.name,
            ))
        }
    }

    #[pyproto]
    impl pyo3::class::basic::PyObjectProtocol for PyXref {
        fn __str__(&self) -> PyResult<String> {
            PyXref::__repr__(self)
        }

        fn __repr__(&self) -> PyResult<String> {
            Ok(format!("PyXref(src: {:#x} dst: {:#x} type: {})",
                self.src,
                self.dst,
                match self.typ {
                    1 => "XREF_FALLTHROUGH",
                    2 => "XREF_CALL",
                    3 => "XREF_UNCONDITIONAL_JUMP",
                    4 => "XREF_CONDITIONAL_JUMP",
                    5 => "XREF_CONDITIONAL_MOVE",
                    _ => unreachable!(),
                },
            ))
        }
    }

    Ok(())
}

<|endoftext|>


<commit_before><commit_msg>feature(parse-collada): better support parsing normal data<commit_after><|endoftext|>


<commit_before>// rustfmt-format_strings: false

fn main() {
    let expected = "; ModuleID = \'foo\'

; Function Attrs: nounwind
declare void @llvm.memset.p0i8.i32(i8* nocapture, i8, i32, i32, i1) #0

declare i32 @write(i32, i8*, i32)

declare i32 @putchar(i32)

declare i32 @getchar()

define i32 @main() {
entry:
  ret i32 0
}

attributes #0 = { nounwind }
";
}
<commit_msg>removed comment with default: rustfmt-format_strings<commit_after>fn main() {
    let expected = "; ModuleID = \'foo\'

; Function Attrs: nounwind
declare void @llvm.memset.p0i8.i32(i8* nocapture, i8, i32, i32, i1) #0

declare i32 @write(i32, i8*, i32)

declare i32 @putchar(i32)

declare i32 @getchar()

define i32 @main() {
entry:
  ret i32 0
}

attributes #0 = { nounwind }
";
}
<|endoftext|>


<commit_before><commit_msg>note source of test values<commit_after><|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use bincode::deserialize;
use byteorder::{LittleEndian, ReadBytesExt};
use clap;
use std::any::TypeId;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};
use std::{mem, process};
use webrender::WEBRENDER_RECORDING_HEADER;
use webrender::api::{ApiMsg, DocumentMsg};
use wrench::{Wrench, WrenchThing};

#[derive(Clone)]
enum Item {
    Message(ApiMsg),
    Data(Vec<u8>),
}

pub struct BinaryFrameReader {
    file: File,
    eof: bool,
    frame_offsets: Vec<u64>,

    skip_uploads: bool,
    replay_api: bool,
    play_through: bool,

    frame_data: Vec<Item>,
    frame_num: u32,
    frame_read: bool,
}

impl BinaryFrameReader {
    pub fn new(file_path: &Path) -> BinaryFrameReader {
        let mut file = File::open(&file_path).expect("Can't open recording file");
        let header = file.read_u64::<LittleEndian>().unwrap();
        assert_eq!(
            header,
            WEBRENDER_RECORDING_HEADER,
            "Binary recording is missing recording header!"
        );

        let apimsg_type_id = unsafe {
            assert_eq!(mem::size_of::<TypeId>(), mem::size_of::<u64>());
            mem::transmute::<TypeId, u64>(TypeId::of::<ApiMsg>())
        };

        let written_apimsg_type_id = file.read_u64::<LittleEndian>().unwrap();
        if written_apimsg_type_id != apimsg_type_id {
            println!(
                "Warning: binary file ApiMsg type mismatch: expected 0x{:x}, found 0x{:x}",
                apimsg_type_id,
                written_apimsg_type_id
            );
        }

        BinaryFrameReader {
            file,
            eof: false,
            frame_offsets: vec![],

            skip_uploads: false,
            replay_api: false,
            play_through: false,

            frame_data: vec![],
            frame_num: 0,
            frame_read: false,
        }
    }

    pub fn new_from_args(args: &clap::ArgMatches) -> BinaryFrameReader {
        let bin_file = args.value_of("INPUT").map(|s| PathBuf::from(s)).unwrap();
        let mut r = BinaryFrameReader::new(&bin_file);
        r.skip_uploads = args.is_present("skip-uploads");
        r.replay_api = args.is_present("api");
        r.play_through = args.is_present("play");
        r
    }

    // FIXME I don't think we can skip uploads without also skipping
    // payload (I think? Unused payload ranges may also be ignored.). But
    // either way we'd need to track image updates and deletions -- if we
    // delete an image, we can't go back to a previous frame.
    //
    // We could probably introduce a mode where going backwards replays all
    // frames up until that frame, so that we know we can be accurate.
    fn should_skip_upload_msg(&self, msg: &ApiMsg) -> bool {
        if !self.skip_uploads {
            return false;
        }

        match *msg {
            ApiMsg::UpdateResources(..) => true,
            _ => false,
        }
    }

    // a frame exists if we either haven't hit eof yet, or if
    // we have, then if we've seen its offset.
    fn frame_exists(&self, frame: u32) -> bool {
        !self.eof || (frame as usize) < self.frame_offsets.len()
    }
}

impl WrenchThing for BinaryFrameReader {
    fn do_frame(&mut self, wrench: &mut Wrench) -> u32 {
        // save where the frame begins as we read through the file
        if self.frame_num as usize >= self.frame_offsets.len() {
            self.frame_num = self.frame_offsets.len() as u32;
            let pos = self.file.seek(SeekFrom::Current(0)).unwrap();
            println!("Frame {} offset: {}", self.frame_offsets.len(), pos);
            self.frame_offsets.push(pos);
        }

        let first_time = !self.frame_read;
        if first_time {
            let offset = self.frame_offsets[self.frame_num as usize];
            self.file.seek(SeekFrom::Start(offset)).unwrap();

            wrench.set_title(&format!("frame {}", self.frame_num));

            self.frame_data.clear();
            let mut found_frame_marker = false;
            let mut found_display_list = false;
            let mut found_pipeline = false;
            while let Ok(mut len) = self.file.read_u32::<LittleEndian>() {
                if len > 0 {
                    let mut buffer = vec![0; len as usize];
                    self.file.read_exact(&mut buffer).unwrap();
                    let msg = deserialize(&buffer).unwrap();
                    // In order to detect the first valid frame, we
                    // need to find:
                    // (a) SetRootPipeline
                    // (b) SetDisplayList
                    // (c) GenerateFrame that occurs *after* (a) and (b)
                    match msg {
                        ApiMsg::UpdateDocument(_, DocumentMsg::GenerateFrame(..)) => {
                            found_frame_marker = true;
                        }
                        ApiMsg::UpdateDocument(_, DocumentMsg::SetDisplayList { .. }) => {
                            found_frame_marker = false;
                            found_display_list = true;
                        }
                        ApiMsg::UpdateDocument(_, DocumentMsg::SetRootPipeline(..)) => {
                            found_frame_marker = false;
                            found_pipeline = true;
                        }
                        _ => {}
                    }
                    self.frame_data.push(Item::Message(msg));
                    // Frames are marked by the GenerateFrame message.
                    // On the first frame, we additionally need to find at least
                    // a SetDisplayList and a SetRootPipeline.
                    // After the first frame, any GenerateFrame message marks a new
                    // frame being rendered.
                    if found_frame_marker && (self.frame_num > 0 || (found_display_list && found_pipeline)) {
                        break;
                    }
                } else {
                    len = self.file.read_u32::<LittleEndian>().unwrap();
                    let mut buffer = vec![0; len as usize];
                    self.file.read_exact(&mut buffer).unwrap();
                    self.frame_data.push(Item::Data(buffer));
                }
            }

            if self.eof == false &&
                self.file.seek(SeekFrom::Current(0)).unwrap() == self.file.metadata().unwrap().len()
            {
                self.eof = true;
            }

            self.frame_read = true;
        }

        if first_time || self.replay_api {
            wrench.begin_frame();
            let frame_items = self.frame_data.clone();
            for item in frame_items {
                match item {
                    Item::Message(msg) => if !self.should_skip_upload_msg(&msg) {
                        wrench.api.send_message(msg);
                    },
                    Item::Data(buf) => {
                        wrench.api.send_payload(&buf);
                    }
                }
            }
        } else if self.play_through {
            if !self.frame_exists(self.frame_num + 1) {
                process::exit(0);
            }
            self.next_frame();
            self.do_frame(wrench);
        } else {
            wrench.refresh();
        }

        self.frame_num
    }

    // note that we don't loop here; we could, but that'll require
    // some tracking to avoid reuploading resources every run.  We
    // sort of try in should_skip_upload_msg, but this needs work.
    fn next_frame(&mut self) {
        if self.frame_exists(self.frame_num + 1) {
            self.frame_num += 1;
            self.frame_read = false;
        }
    }

    fn prev_frame(&mut self) {
        if self.frame_num > 0 {
            self.frame_num -= 1;
            self.frame_read = false;
        }
    }
}
<commit_msg>Auto merge of #2209 - salty-horse:ignore_delete_document, r=glennw<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use bincode::deserialize;
use byteorder::{LittleEndian, ReadBytesExt};
use clap;
use std::any::TypeId;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};
use std::{mem, process};
use webrender::WEBRENDER_RECORDING_HEADER;
use webrender::api::{ApiMsg, DocumentMsg};
use wrench::{Wrench, WrenchThing};

#[derive(Clone)]
enum Item {
    Message(ApiMsg),
    Data(Vec<u8>),
}

pub struct BinaryFrameReader {
    file: File,
    eof: bool,
    frame_offsets: Vec<u64>,

    skip_uploads: bool,
    replay_api: bool,
    play_through: bool,

    frame_data: Vec<Item>,
    frame_num: u32,
    frame_read: bool,
}

impl BinaryFrameReader {
    pub fn new(file_path: &Path) -> BinaryFrameReader {
        let mut file = File::open(&file_path).expect("Can't open recording file");
        let header = file.read_u64::<LittleEndian>().unwrap();
        assert_eq!(
            header,
            WEBRENDER_RECORDING_HEADER,
            "Binary recording is missing recording header!"
        );

        let apimsg_type_id = unsafe {
            assert_eq!(mem::size_of::<TypeId>(), mem::size_of::<u64>());
            mem::transmute::<TypeId, u64>(TypeId::of::<ApiMsg>())
        };

        let written_apimsg_type_id = file.read_u64::<LittleEndian>().unwrap();
        if written_apimsg_type_id != apimsg_type_id {
            println!(
                "Warning: binary file ApiMsg type mismatch: expected 0x{:x}, found 0x{:x}",
                apimsg_type_id,
                written_apimsg_type_id
            );
        }

        BinaryFrameReader {
            file,
            eof: false,
            frame_offsets: vec![],

            skip_uploads: false,
            replay_api: false,
            play_through: false,

            frame_data: vec![],
            frame_num: 0,
            frame_read: false,
        }
    }

    pub fn new_from_args(args: &clap::ArgMatches) -> BinaryFrameReader {
        let bin_file = args.value_of("INPUT").map(|s| PathBuf::from(s)).unwrap();
        let mut r = BinaryFrameReader::new(&bin_file);
        r.skip_uploads = args.is_present("skip-uploads");
        r.replay_api = args.is_present("api");
        r.play_through = args.is_present("play");
        r
    }

    // FIXME I don't think we can skip uploads without also skipping
    // payload (I think? Unused payload ranges may also be ignored.). But
    // either way we'd need to track image updates and deletions -- if we
    // delete an image, we can't go back to a previous frame.
    //
    // We could probably introduce a mode where going backwards replays all
    // frames up until that frame, so that we know we can be accurate.
    fn should_skip_upload_msg(&self, msg: &ApiMsg) -> bool {
        if !self.skip_uploads {
            return false;
        }

        match *msg {
            ApiMsg::UpdateResources(..) => true,
            _ => false,
        }
    }

    // a frame exists if we either haven't hit eof yet, or if
    // we have, then if we've seen its offset.
    fn frame_exists(&self, frame: u32) -> bool {
        !self.eof || (frame as usize) < self.frame_offsets.len()
    }
}

impl WrenchThing for BinaryFrameReader {
    fn do_frame(&mut self, wrench: &mut Wrench) -> u32 {
        // save where the frame begins as we read through the file
        if self.frame_num as usize >= self.frame_offsets.len() {
            self.frame_num = self.frame_offsets.len() as u32;
            let pos = self.file.seek(SeekFrom::Current(0)).unwrap();
            println!("Frame {} offset: {}", self.frame_offsets.len(), pos);
            self.frame_offsets.push(pos);
        }

        let first_time = !self.frame_read;
        if first_time {
            let offset = self.frame_offsets[self.frame_num as usize];
            self.file.seek(SeekFrom::Start(offset)).unwrap();

            wrench.set_title(&format!("frame {}", self.frame_num));

            self.frame_data.clear();
            let mut found_frame_marker = false;
            let mut found_display_list = false;
            let mut found_pipeline = false;
            while let Ok(mut len) = self.file.read_u32::<LittleEndian>() {
                if len > 0 {
                    let mut buffer = vec![0; len as usize];
                    self.file.read_exact(&mut buffer).unwrap();
                    let msg = deserialize(&buffer).unwrap();
                    let mut store_message = true;
                    // In order to detect the first valid frame, we
                    // need to find:
                    // (a) SetRootPipeline
                    // (b) SetDisplayList
                    // (c) GenerateFrame that occurs *after* (a) and (b)
                    match msg {
                        ApiMsg::UpdateDocument(_, DocumentMsg::GenerateFrame(..)) => {
                            found_frame_marker = true;
                        }
                        ApiMsg::UpdateDocument(_, DocumentMsg::SetDisplayList { .. }) => {
                            found_frame_marker = false;
                            found_display_list = true;
                        }
                        ApiMsg::UpdateDocument(_, DocumentMsg::SetRootPipeline(..)) => {
                            found_frame_marker = false;
                            found_pipeline = true;
                        }
                        // Wrench depends on the document always existing
                        ApiMsg::DeleteDocument(_) => {
                            store_message = false;
                        }
                        _ => {}
                    }
                    if store_message {
                        self.frame_data.push(Item::Message(msg));
                    }
                    // Frames are marked by the GenerateFrame message.
                    // On the first frame, we additionally need to find at least
                    // a SetDisplayList and a SetRootPipeline.
                    // After the first frame, any GenerateFrame message marks a new
                    // frame being rendered.
                    if found_frame_marker && (self.frame_num > 0 || (found_display_list && found_pipeline)) {
                        break;
                    }
                } else {
                    len = self.file.read_u32::<LittleEndian>().unwrap();
                    let mut buffer = vec![0; len as usize];
                    self.file.read_exact(&mut buffer).unwrap();
                    self.frame_data.push(Item::Data(buffer));
                }
            }

            if self.eof == false &&
                self.file.seek(SeekFrom::Current(0)).unwrap() == self.file.metadata().unwrap().len()
            {
                self.eof = true;
            }

            self.frame_read = true;
        }

        if first_time || self.replay_api {
            wrench.begin_frame();
            let frame_items = self.frame_data.clone();
            for item in frame_items {
                match item {
                    Item::Message(msg) => if !self.should_skip_upload_msg(&msg) {
                        wrench.api.send_message(msg);
                    },
                    Item::Data(buf) => {
                        wrench.api.send_payload(&buf);
                    }
                }
            }
        } else if self.play_through {
            if !self.frame_exists(self.frame_num + 1) {
                process::exit(0);
            }
            self.next_frame();
            self.do_frame(wrench);
        } else {
            wrench.refresh();
        }

        self.frame_num
    }

    // note that we don't loop here; we could, but that'll require
    // some tracking to avoid reuploading resources every run.  We
    // sort of try in should_skip_upload_msg, but this needs work.
    fn next_frame(&mut self) {
        if self.frame_exists(self.frame_num + 1) {
            self.frame_num += 1;
            self.frame_read = false;
        }
    }

    fn prev_frame(&mut self) {
        if self.frame_num > 0 {
            self.frame_num -= 1;
            self.frame_read = false;
        }
    }
}
<|endoftext|>