<commit_before>extern crate osmpbfreader;
use std::collections::HashMap;
use models::*;
use categorize::*;
use std;

// Way as represented in OpenStreetMap
struct Way {
    id: i64,
    nodes: Vec<i64>,
    properties: EdgeProperties,
}

struct Reader {
    nodes: HashMap<i64, Node>,
    ways: Vec<Way>,
}

impl Reader {
    fn new() -> Reader {
        Reader {
            nodes: HashMap::new(),
            ways: Vec::new(),
        }
    }

    fn count_nodes_uses(&mut self) {
        for way in &self.ways {
            for (i, node_id) in way.nodes.iter().enumerate() {
                if let Some(node) = self.nodes.get_mut(node_id) {
                    // Count double extremities nodes
                    if i == 0 || i == way.nodes.len() - 1 {
                        node.uses += 2;
                    } else {
                        node.uses += 1;
                    }
                } else {
                    panic!("Missing node, id: {}", node_id)
                }
            }
        }
    }

    fn split_way(&self, way: &Way) -> Vec<Edge> {
        let mut result = Vec::new();

        let mut source = 0;
        let mut geometry = Vec::new();
        for (i, &node_id) in way.nodes.iter().enumerate() {
            let node = self.nodes[&node_id];
            if i == 0 {
                source = node_id;
                geometry.push(node.coord);
            } else {
                geometry.push(node.coord);

                if node.uses > 1 {
                    result.push(Edge {
                        id: way.id,
                        source: source,
                        target: node_id,
                        geometry: geometry,
                        properties: way.properties,
                    });

                    source = node_id;
                    geometry = vec![node.coord];
                }
            }
        }
        result
    }

    fn read(&mut self, filename: &str) -> Result<(), String> {
        let path = std::path::Path::new(filename);
        let r = try!(std::fs::File::open(&path).map_err(|e| e.to_string()));
        let mut pbf = osmpbfreader::OsmPbfReader::new(r);
        for obj in pbf.iter() {
            match obj {
                osmpbfreader::OsmObj::Node(node) => {
                    let mut n = Node::new();
                    n.id = node.id;
                    n.coord = Coord {
                        lon: node.lon,
                        lat: node.lat,
                    };
                    self.nodes.insert(node.id, n);
                }
                osmpbfreader::OsmObj::Way(way) => {
                    let mut properties = EdgeProperties::new();
                    for (key, val) in way.tags {
                        properties.update(key, val);
                    }
                    properties.normalize();
                    if properties.accessible() {
                        self.ways.push(Way {
                            id: way.id,
                            nodes: way.nodes,
                            properties: properties,
                        })
                    }
                }
                osmpbfreader::OsmObj::Relation(_) => {}
            }
        }
        Ok(())
    }

    fn nodes(self) -> Vec<Node> {
        self.nodes
            .into_iter()
            .map(|(_, node)| node)
            .filter(|node| node.uses > 1)
            .collect()
    }

    fn edges(&self) -> Vec<Edge> {
        self.ways.iter().flat_map(|way| self.split_way(way)).collect()
    }
}

// Read all the nodes and ways of the osm.pbf file
pub fn read(filename: &str) -> Result<(Vec<Node>, Vec<Edge>), String> {
    let mut r = Reader::new();
    try!(r.read(filename));
    r.count_nodes_uses();
    let edges = r.edges();
    Ok((r.nodes(), edges))
}

#[test]
fn test_real_all() {
    let (nodes, ways) = read("src/osm4routing/test_data/minimal.osm.pbf").unwrap();
    assert_eq!(2, nodes.len());
    assert_eq!(1, ways.len());
}

#[test]
fn test_count_nodes() {
    let ways = vec![Way {
                        id: 0,
                        nodes: vec![1, 2, 3],
                        properties: EdgeProperties::new(),
                    }];
    let mut nodes = HashMap::new();
    nodes.insert(1, Node::new());
    nodes.insert(2, Node::new());
    nodes.insert(3, Node::new());
    let mut r = Reader {
        ways: ways,
        nodes: nodes,
    };
    r.count_nodes_uses();
    assert_eq!(2, r.nodes[&1].uses);
    assert_eq!(1, r.nodes[&2].uses);
    assert_eq!(2, r.nodes[&3].uses);

    assert_eq!(2, r.nodes().len());
}

#[test]
fn test_split() {
    let mut nodes = HashMap::new();
    nodes.insert(1, Node::new());
    nodes.insert(2, Node::new());
    nodes.insert(3, Node::new());
    nodes.insert(4, Node::new());
    nodes.insert(5, Node::new());
    let ways = vec![Way {
                        id: 0,
                        nodes: vec![1, 2, 3],
                        properties: EdgeProperties::new(),
                    },
                    Way {
                        id: 0,
                        nodes: vec![4, 5, 2],
                        properties: EdgeProperties::new(),
                    }];
    let mut r = Reader {
        nodes: nodes,
        ways: ways,
    };
    r.count_nodes_uses();
    let edges = r.edges();
    assert_eq!(3, edges.len());
}

#[test]
fn test_wrong_file() {
    let r = read("i hope you have no file name like this one");
    assert!(r.is_err());
}
<commit_msg>Do not read useless nodes<commit_after>extern crate osmpbfreader;
use std::collections::{HashMap, HashSet};
use models::*;
use categorize::*;
use std;

// Way as represented in OpenStreetMap
struct Way {
    id: i64,
    nodes: Vec<i64>,
    properties: EdgeProperties,
}

struct Reader {
    nodes: HashMap<i64, Node>,
    ways: Vec<Way>,
    nodes_to_keep: HashSet<i64>,
}

impl Reader {
    fn new() -> Reader {
        Reader {
            nodes: HashMap::new(),
            ways: Vec::new(),
            nodes_to_keep: HashSet::new(),
        }
    }

    fn count_nodes_uses(&mut self) {
        for way in &self.ways {
            for (i, node_id) in way.nodes.iter().enumerate() {
                if let Some(node) = self.nodes.get_mut(node_id) {
                    // Count double extremities nodes
                    if i == 0 || i == way.nodes.len() - 1 {
                        node.uses += 2;
                    } else {
                        node.uses += 1;
                    }
                } else {
                    panic!("Missing node, id: {}", node_id)
                }
            }
        }
    }

    fn split_way(&self, way: &Way) -> Vec<Edge> {
        let mut result = Vec::new();

        let mut source = 0;
        let mut geometry = Vec::new();
        for (i, &node_id) in way.nodes.iter().enumerate() {
            let node = self.nodes[&node_id];
            if i == 0 {
                source = node_id;
                geometry.push(node.coord);
            } else {
                geometry.push(node.coord);

                if node.uses > 1 {
                    result.push(Edge {
                        id: way.id,
                        source: source,
                        target: node_id,
                        geometry: geometry,
                        properties: way.properties,
                    });

                    source = node_id;
                    geometry = vec![node.coord];
                }
            }
        }
        result
    }

    fn read_ways(&mut self, file: std::fs::File) {
        let mut pbf = osmpbfreader::OsmPbfReader::new(file);
        for obj in pbf.iter() {
            match obj {
                osmpbfreader::OsmObj::Way(way) => {
                    let mut properties = EdgeProperties::new();
                    for (key, val) in way.tags {
                        properties.update(key, val);
                    }
                    properties.normalize();
                    if properties.accessible() {
                        for node in &way.nodes {
                            self.nodes_to_keep.insert(node.clone());
                        }
                        self.ways.push(Way {
                            id: way.id,
                            nodes: way.nodes,
                            properties: properties,
                        });

                    }
                }
                _ => {}
            }
        }
    }

    fn read_nodes(&mut self, file: std::fs::File) {
        let mut pbf = osmpbfreader::OsmPbfReader::new(file);
        self.nodes.reserve(self.nodes_to_keep.len());
        for obj in pbf.iter() {
            match obj {
                osmpbfreader::OsmObj::Node(node) => {
                    if self.nodes_to_keep.contains(&node.id) {
                        self.nodes_to_keep.remove(&node.id);
                        let mut n = Node::new();
                        n.id = node.id;
                        n.coord = Coord {
                            lon: node.lon,
                            lat: node.lat,
                        };
                        self.nodes.insert(node.id, n);
                    }
                }
                _ => {}
            }
        }
    }

    fn nodes(self) -> Vec<Node> {
        self.nodes
            .into_iter()
            .map(|(_, node)| node)
            .filter(|node| node.uses > 1)
            .collect()
    }

    fn edges(&self) -> Vec<Edge> {
        self.ways.iter().flat_map(|way| self.split_way(way)).collect()
    }
}

// Read all the nodes and ways of the osm.pbf file
pub fn read(filename: &str) -> Result<(Vec<Node>, Vec<Edge>), String> {
    let mut r = Reader::new();
    let path = std::path::Path::new(filename);
    let file = try!(std::fs::File::open(&path).map_err(|e| e.to_string()));
    r.read_ways(file);
    let file_nodes = try!(std::fs::File::open(&path).map_err(|e| e.to_string()));
    r.read_nodes(file_nodes);
    r.count_nodes_uses();
    let edges = r.edges();
    Ok((r.nodes(), edges))
}

#[test]
fn test_real_all() {
    let (nodes, ways) = read("src/osm4routing/test_data/minimal.osm.pbf").unwrap();
    assert_eq!(2, nodes.len());
    assert_eq!(1, ways.len());
}

#[test]
fn test_count_nodes() {
    let ways = vec![Way {
                        id: 0,
                        nodes: vec![1, 2, 3],
                        properties: EdgeProperties::new(),
                    }];
    let mut nodes = HashMap::new();
    nodes.insert(1, Node::new());
    nodes.insert(2, Node::new());
    nodes.insert(3, Node::new());
    let mut r = Reader {
        ways: ways,
        nodes: nodes,
        nodes_to_keep: HashSet::new(),
    };
    r.count_nodes_uses();
    assert_eq!(2, r.nodes[&1].uses);
    assert_eq!(1, r.nodes[&2].uses);
    assert_eq!(2, r.nodes[&3].uses);

    assert_eq!(2, r.nodes().len());
}

#[test]
fn test_split() {
    let mut nodes = HashMap::new();
    nodes.insert(1, Node::new());
    nodes.insert(2, Node::new());
    nodes.insert(3, Node::new());
    nodes.insert(4, Node::new());
    nodes.insert(5, Node::new());
    let ways = vec![Way {
                        id: 0,
                        nodes: vec![1, 2, 3],
                        properties: EdgeProperties::new(),
                    },
                    Way {
                        id: 0,
                        nodes: vec![4, 5, 2],
                        properties: EdgeProperties::new(),
                    }];
    let mut r = Reader {
        nodes: nodes,
        ways: ways,
        nodes_to_keep: HashSet::new(),
    };
    r.count_nodes_uses();
    let edges = r.edges();
    assert_eq!(3, edges.len());
}

#[test]
fn test_wrong_file() {
    let r = read("i hope you have no file name like this one");
    assert!(r.is_err());
}
<|endoftext|>


<commit_before>// Prelude: Descriptor and its layouts in shader

use prelude::internals::*;
use std;
use vkffi::*;
use render_vk::wrap as vk;

#[derive(Clone, Copy)]
pub enum ShaderStage { Vertex, TessControl, TessEvaluate, Geometry, Fragment }
impl std::convert::Into<VkShaderStageFlags> for ShaderStage
{
	fn into(self) -> VkShaderStageFlags
	{
		match self
		{
			ShaderStage::Vertex => VK_SHADER_STAGE_VERTEX_BIT,
			ShaderStage::TessControl => VK_SHADER_STAGE_TESSELLATION_CONTROL_BIT,
			ShaderStage::TessEvaluate => VK_SHADER_STAGE_TESSELLATION_EVALUATION_BIT,
			ShaderStage::Geometry => VK_SHADER_STAGE_GEOMETRY_BIT,
			ShaderStage::Fragment => VK_SHADER_STAGE_FRAGMENT_BIT
		}
	}
}

#[derive(Clone)]
pub enum Descriptor
{
	Uniform(u32, Vec<ShaderStage>),
	CombinedSampler(u32, Vec<ShaderStage>)
}
pub trait DescriptorInternals
{
	fn count(&self) -> u32;
	fn into_binding(&self, index: u32) -> VkDescriptorSetLayoutBinding;
	fn into_pool_size(&self) -> VkDescriptorPoolSize;
}
impl Descriptor
{
	fn native_type(&self) -> VkDescriptorType
	{
		match self
		{
			&Descriptor::Uniform(_, _) => VkDescriptorType::UniformBuffer,
			&Descriptor::CombinedSampler(_, _) => VkDescriptorType::CombinedImageSampler
		}
	}
	fn stage_mask(&self) -> VkShaderStageFlags
	{
		match self
		{
			&Descriptor::Uniform(_, ref s) => s,
			&Descriptor::CombinedSampler(_, ref s) => s
		}.iter().fold(0, |flag, f| flag | Into::<VkShaderStageFlags>::into(*f))
	}
}
impl DescriptorInternals for Descriptor
{
	fn count(&self) -> u32
	{
		match self
		{
			&Descriptor::Uniform(n, _) => n,
			&Descriptor::CombinedSampler(n, _) => n
		}
	}
	fn into_binding(&self, index: u32) -> VkDescriptorSetLayoutBinding
	{
		VkDescriptorSetLayoutBinding
		{
			binding: index, descriptorType: self.native_type(), descriptorCount: self.count(),
			stageFlags: self.stage_mask(), pImmutableSamplers: std::ptr::null()
		}
	}
	fn into_pool_size(&self) -> VkDescriptorPoolSize
	{
		VkDescriptorPoolSize(self.native_type(), self.count())
	}
}

pub struct DescriptorSetLayout
{
	internal: vk::DescriptorSetLayout,
	structure: Vec<Descriptor>
}
pub trait DescriptorSetLayoutInternals
{
	fn new(dsl: vk::DescriptorSetLayout, structure: &[Descriptor]) -> Self;
	fn descriptors(&self) -> &[Descriptor];
}
impl DescriptorSetLayoutInternals for DescriptorSetLayout
{
	fn new(dsl: vk::DescriptorSetLayout, structure: &[Descriptor]) -> Self
	{
		DescriptorSetLayout { internal: dsl, structure: Vec::from(structure) }
	}
	fn descriptors(&self) -> &[Descriptor] { &self.structure }
}
impl InternalExports<vk::DescriptorSetLayout> for DescriptorSetLayout
{
	fn get_internal(&self) -> &vk::DescriptorSetLayout { &self.internal }
}

pub struct DescriptorSets
{
	pool: vk::DescriptorPool, sets: Vec<VkDescriptorSet>
}
pub trait DescriptorSetsInternals
{
	fn new(pool: vk::DescriptorPool, sets: Vec<VkDescriptorSet>) -> Self;
}
impl DescriptorSetsInternals for DescriptorSets
{
	fn new(pool: vk::DescriptorPool, sets: Vec<VkDescriptorSet>) -> Self
	{
		DescriptorSets { pool: pool, sets: sets }
	}
}
impl std::ops::Index<usize> for DescriptorSets
{
	type Output = VkDescriptorSet;
	fn index<'a>(&'a self, idx: usize) -> &'a Self::Output { &self.sets[idx] }
}
<commit_msg>fix comment<commit_after>// Prelude: Descriptor and its layout in shader

use prelude::internals::*;
use std;
use vkffi::*;
use render_vk::wrap as vk;

#[derive(Clone, Copy)]
pub enum ShaderStage { Vertex, TessControl, TessEvaluate, Geometry, Fragment }
impl std::convert::Into<VkShaderStageFlags> for ShaderStage
{
	fn into(self) -> VkShaderStageFlags
	{
		match self
		{
			ShaderStage::Vertex => VK_SHADER_STAGE_VERTEX_BIT,
			ShaderStage::TessControl => VK_SHADER_STAGE_TESSELLATION_CONTROL_BIT,
			ShaderStage::TessEvaluate => VK_SHADER_STAGE_TESSELLATION_EVALUATION_BIT,
			ShaderStage::Geometry => VK_SHADER_STAGE_GEOMETRY_BIT,
			ShaderStage::Fragment => VK_SHADER_STAGE_FRAGMENT_BIT
		}
	}
}

#[derive(Clone)]
pub enum Descriptor
{
	Uniform(u32, Vec<ShaderStage>),
	CombinedSampler(u32, Vec<ShaderStage>)
}
pub trait DescriptorInternals
{
	fn count(&self) -> u32;
	fn into_binding(&self, index: u32) -> VkDescriptorSetLayoutBinding;
	fn into_pool_size(&self) -> VkDescriptorPoolSize;
}
impl Descriptor
{
	fn native_type(&self) -> VkDescriptorType
	{
		match self
		{
			&Descriptor::Uniform(_, _) => VkDescriptorType::UniformBuffer,
			&Descriptor::CombinedSampler(_, _) => VkDescriptorType::CombinedImageSampler
		}
	}
	fn stage_mask(&self) -> VkShaderStageFlags
	{
		match self
		{
			&Descriptor::Uniform(_, ref s) => s,
			&Descriptor::CombinedSampler(_, ref s) => s
		}.iter().fold(0, |flag, f| flag | Into::<VkShaderStageFlags>::into(*f))
	}
}
impl DescriptorInternals for Descriptor
{
	fn count(&self) -> u32
	{
		match self
		{
			&Descriptor::Uniform(n, _) => n,
			&Descriptor::CombinedSampler(n, _) => n
		}
	}
	fn into_binding(&self, index: u32) -> VkDescriptorSetLayoutBinding
	{
		VkDescriptorSetLayoutBinding
		{
			binding: index, descriptorType: self.native_type(), descriptorCount: self.count(),
			stageFlags: self.stage_mask(), pImmutableSamplers: std::ptr::null()
		}
	}
	fn into_pool_size(&self) -> VkDescriptorPoolSize
	{
		VkDescriptorPoolSize(self.native_type(), self.count())
	}
}

pub struct DescriptorSetLayout
{
	internal: vk::DescriptorSetLayout,
	structure: Vec<Descriptor>
}
pub trait DescriptorSetLayoutInternals
{
	fn new(dsl: vk::DescriptorSetLayout, structure: &[Descriptor]) -> Self;
	fn descriptors(&self) -> &[Descriptor];
}
impl DescriptorSetLayoutInternals for DescriptorSetLayout
{
	fn new(dsl: vk::DescriptorSetLayout, structure: &[Descriptor]) -> Self
	{
		DescriptorSetLayout { internal: dsl, structure: Vec::from(structure) }
	}
	fn descriptors(&self) -> &[Descriptor] { &self.structure }
}
impl InternalExports<vk::DescriptorSetLayout> for DescriptorSetLayout
{
	fn get_internal(&self) -> &vk::DescriptorSetLayout { &self.internal }
}

pub struct DescriptorSets
{
	pool: vk::DescriptorPool, sets: Vec<VkDescriptorSet>
}
pub trait DescriptorSetsInternals
{
	fn new(pool: vk::DescriptorPool, sets: Vec<VkDescriptorSet>) -> Self;
}
impl DescriptorSetsInternals for DescriptorSets
{
	fn new(pool: vk::DescriptorPool, sets: Vec<VkDescriptorSet>) -> Self
	{
		DescriptorSets { pool: pool, sets: sets }
	}
}
impl std::ops::Index<usize> for DescriptorSets
{
	type Output = VkDescriptorSet;
	fn index<'a>(&'a self, idx: usize) -> &'a Self::Output { &self.sets[idx] }
}
<|endoftext|>


<commit_before>use std::path::PathBuf;
use std::fs::File;
use std::error::Error;
use std::fs;
use std::fmt;
use std::io::{Seek, Read, SeekFrom};

#[derive(Debug)]
pub enum FsError {
    FileNotFound,
    InvalidDirectory
}
impl Error for FsError {
    fn description(&self) -> &str {
        match *self {
            FsError::FileNotFound => "the folder does not exist or cannot be read from",
            FsError::InvalidDirectory => "the specified directory is not a valid directory",
        }
    }
}
impl fmt::Display for FsError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FsError::FileNotFound => write!(f, "the folder specified could not be found or read from"),
            FsError::InvalidDirectory => write!(f, "the specified directory is not a valid directory"),
        }
    }
}

#[derive(Debug)]
pub struct FileSystem {
    path: PathBuf,
    mainfile: MainFile
}

#[derive(Debug)]
pub struct MainFile {
    file: Option<File>
}

#[derive(Debug)]
pub struct IndexEntry {
    id: u32,
    size: u32,
    offset: u64
}

impl FileSystem {
    pub fn new(string: &'static str) -> Result<FileSystem, FsError> {
        // Declare some nice variables!!!
        let path = PathBuf::from(string);
        let metadata = fs::metadata(&path);

        // Make sure the folder exists
        if !metadata.is_ok() {
            return Err(FsError::FileNotFound);
        }

        // Make sure it is a directory
        if !metadata.unwrap().is_dir() {
            return Err(FsError::InvalidDirectory);
        }

        // Create mainfile path
        let mut mainfile_path = PathBuf::from(string);
        mainfile_path.push("main_file_cache.dat2");

        let file = File::open(mainfile_path).ok();
        Ok(FileSystem{path: path, mainfile: MainFile{file: file}})
    }

    /// Gets the mainfile, that is, the main_file_cache.dat2 entry in the folder
    /// that holds the actual binary data of the filesystem entries.
    pub fn mainfile(&mut self) -> &mut MainFile {
        &mut self.mainfile
    }
}

impl MainFile {
    /// Checks if the file exists.
    pub fn exists(&self) -> bool {
        self.file.is_some()
    }

    /// Gets the backing file, if existant. Returns a new instance with a fresh seek pointer.
    pub fn file(&mut self) -> Option<&mut File> {
        self.file.as_mut()
    }

    /// Calculates the number of data blocks in the mainfile (if existant). This is done by
    /// taking the file size and dividing that by 520 (rouding up), because each block
    /// takes up 520 bytes of data.
    pub fn num_blocks(&self) -> Option<u64> {
        match self.file {
            Some(ref x) => Some((x.metadata().unwrap().len() + 519u64) / 520u64),
            None => None
        }
    }

    /// Reads a block of data, specified by the block id. The data is read at 520 * block_id
    /// and is exactly 520 bytes big. It is not guaranteed all 520 bytes are occupied if the
    /// block is the last one, thus possible to be trimmed.
    pub fn read_block(&mut self, block: u32) -> Option<[u8; 520]> {
        // Do we have a valid file?
        if self.file.is_none() {
            return None;
        }

        let mut data: [u8; 520] = [0; 520];
        let file = self.file().unwrap();

        // Seek to the right position and read the data
        file.seek(SeekFrom::Start(block as u64 * 520u64));
        file.read(&mut data);

        return Some(data);
    }
}
<commit_msg>Add getters for index entry fields<commit_after>use std::path::PathBuf;
use std::fs::File;
use std::error::Error;
use std::fs;
use std::fmt;
use std::io::{Seek, Read, SeekFrom};

#[derive(Debug)]
pub enum FsError {
    FileNotFound,
    InvalidDirectory
}
impl Error for FsError {
    fn description(&self) -> &str {
        match *self {
            FsError::FileNotFound => "the folder does not exist or cannot be read from",
            FsError::InvalidDirectory => "the specified directory is not a valid directory",
        }
    }
}
impl fmt::Display for FsError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FsError::FileNotFound => write!(f, "the folder specified could not be found or read from"),
            FsError::InvalidDirectory => write!(f, "the specified directory is not a valid directory"),
        }
    }
}

#[derive(Debug)]
pub struct FileSystem {
    path: PathBuf,
    mainfile: MainFile
}

#[derive(Debug)]
pub struct MainFile {
    file: Option<File>
}

#[derive(Debug)]
pub struct IndexEntry {
    id: u32,
    size: u32,
    offset: u64
}

impl IndexEntry {
    pub fn id(&self) -> u32 {
        self.id
    }
    pub fn size(&self) -> u32 {
        self.size
    }
    pub fn offset(&self) -> u64 {
        self.offset
    }
}

impl FileSystem {
    pub fn new(string: &'static str) -> Result<FileSystem, FsError> {
        // Declare some nice variables!!!
        let path = PathBuf::from(string);
        let metadata = fs::metadata(&path);

        // Make sure the folder exists
        if !metadata.is_ok() {
            return Err(FsError::FileNotFound);
        }

        // Make sure it is a directory
        if !metadata.unwrap().is_dir() {
            return Err(FsError::InvalidDirectory);
        }

        // Create mainfile path
        let mut mainfile_path = PathBuf::from(string);
        mainfile_path.push("main_file_cache.dat2");

        let file = File::open(mainfile_path).ok();
        Ok(FileSystem{path: path, mainfile: MainFile{file: file}})
    }

    /// Gets the mainfile, that is, the main_file_cache.dat2 entry in the folder
    /// that holds the actual binary data of the filesystem entries.
    pub fn mainfile(&mut self) -> &mut MainFile {
        &mut self.mainfile
    }
}

impl MainFile {
    /// Checks if the file exists.
    pub fn exists(&self) -> bool {
        self.file.is_some()
    }

    /// Gets the backing file, if existant. Returns a new instance with a fresh seek pointer.
    pub fn file(&mut self) -> Option<&mut File> {
        self.file.as_mut()
    }

    /// Calculates the number of data blocks in the mainfile (if existant). This is done by
    /// taking the file size and dividing that by 520 (rouding up), because each block
    /// takes up 520 bytes of data.
    pub fn num_blocks(&self) -> Option<u64> {
        match self.file {
            Some(ref x) => Some((x.metadata().unwrap().len() + 519u64) / 520u64),
            None => None
        }
    }

    /// Reads a block of data, specified by the block id. The data is read at 520 * block_id
    /// and is exactly 520 bytes big. It is not guaranteed all 520 bytes are occupied if the
    /// block is the last one, thus possible to be trimmed.
    pub fn read_block(&mut self, block: u32) -> Option<[u8; 520]> {
        // Do we have a valid file?
        if self.file.is_none() {
            return None;
        }

        let mut data: [u8; 520] = [0; 520];
        let file = self.file().unwrap();

        // Seek to the right position and read the data
        file.seek(SeekFrom::Start(block as u64 * 520u64));
        file.read(&mut data);

        return Some(data);
    }
}
<|endoftext|>


<commit_before><commit_msg>sozuctl: refactor table creation of queried applications<commit_after><|endoftext|>


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

use crate::{
    format_module_id,
    test_reporter::{FailureReason, TestFailure, TestResults, TestRunInfo, TestStatistics},
};
use anyhow::Result;
use bytecode_interpreter::{
    concrete::{settings::InterpreterSettings, value::GlobalState},
    shared::bridge::{adapt_move_vm_change_set, adapt_move_vm_result},
    StacklessBytecodeInterpreter,
};
use colored::*;
use move_binary_format::{errors::VMResult, file_format::CompiledModule};
use move_core_types::{
    effects::ChangeSet,
    gas_schedule::{CostTable, GasAlgebra, GasCost, GasUnits},
    identifier::IdentStr,
    value::serialize_values,
    vm_status::StatusCode,
};
use move_lang::{
    shared::Flags,
    unit_test::{ExpectedFailure, ModuleTestPlan, TestCase, TestPlan},
};
use move_model::{model::GlobalEnv, run_model_builder_with_compilation_flags};
use move_vm_runtime::{logging::NoContextLog, move_vm::MoveVM};
use move_vm_test_utils::InMemoryStorage;
use move_vm_types::gas_schedule::{zero_cost_schedule, GasStatus};
use rayon::prelude::*;
use resource_viewer::MoveValueAnnotator;
use std::{io::Write, marker::Send, sync::Mutex, time::Instant};

/// Test state common to all tests
#[derive(Debug)]
pub struct SharedTestingConfig {
    save_storage_state_on_failure: bool,
    execution_bound: u64,
    cost_table: CostTable,
    starting_storage_state: InMemoryStorage,
    source_files: Vec<String>,
    check_stackless_vm: bool,
    verbose: bool,
}

#[derive(Debug)]
pub struct TestRunner {
    num_threads: usize,
    testing_config: SharedTestingConfig,
    tests: TestPlan,
}

/// A gas schedule where every instruction has a cost of "1". This is used to bound execution of a
/// test to a certain number of ticks.
fn unit_cost_table() -> CostTable {
    let mut cost_schedule = zero_cost_schedule();
    cost_schedule.instruction_table.iter_mut().for_each(|cost| {
        *cost = GasCost::new(1, 1);
    });
    cost_schedule.native_table.iter_mut().for_each(|cost| {
        *cost = GasCost::new(1, 1);
    });
    cost_schedule
}

/// Setup storage state with the set of modules that will be needed for all tests
fn setup_test_storage<'a>(
    modules: impl Iterator<Item = &'a CompiledModule>,
) -> Result<InMemoryStorage> {
    let mut storage = InMemoryStorage::new();
    for module in modules {
        let module_id = module.self_id();
        let mut module_bytes = Vec::new();
        module.serialize(&mut module_bytes)?;
        storage.publish_or_overwrite_module(module_id, module_bytes);
    }

    Ok(storage)
}

/// Print the updates to storage represented by `cs` in the context of the starting storage state
/// `storage`.
fn print_resources(cs: &ChangeSet, storage: &InMemoryStorage) -> Result<String> {
    use std::fmt::Write;
    let mut buf = String::new();
    let annotator = MoveValueAnnotator::new(storage);
    for (account_addr, account_state) in cs.accounts() {
        writeln!(&mut buf, "0x{}:", account_addr.short_str_lossless())?;

        for (tag, resource_opt) in account_state.resources() {
            if let Some(resource) = resource_opt {
                writeln!(
                    &mut buf,
                    "\t{}",
                    format!("=> {}", annotator.view_resource(tag, resource)?).replace("\n", "\n\t")
                )?;
            }
        }
    }

    Ok(buf)
}

impl TestRunner {
    pub fn new(
        execution_bound: u64,
        num_threads: usize,
        check_stackless_vm: bool,
        verbose: bool,
        save_storage_state_on_failure: bool,
        tests: TestPlan,
    ) -> Result<Self> {
        let source_files = tests
            .files
            .keys()
            .map(|filepath| filepath.to_string())
            .collect();
        let modules = tests.module_info.values().map(|info| &info.0);
        let starting_storage_state = setup_test_storage(modules)?;
        Ok(Self {
            testing_config: SharedTestingConfig {
                save_storage_state_on_failure,
                starting_storage_state,
                execution_bound,
                cost_table: unit_cost_table(),
                source_files,
                check_stackless_vm,
                verbose,
            },
            num_threads,
            tests,
        })
    }

    pub fn run<W: Write + Send>(self, writer: &Mutex<W>) -> Result<TestResults> {
        rayon::ThreadPoolBuilder::new()
            .num_threads(self.num_threads)
            .build()
            .unwrap()
            .install(|| {
                let final_statistics = self
                    .tests
                    .module_tests
                    .par_iter()
                    .map(|(_, test_plan)| self.testing_config.exec_module_tests(test_plan, writer))
                    .reduce(TestStatistics::new, |acc, stats| acc.combine(stats));

                Ok(TestResults::new(final_statistics, self.tests))
            })
    }

    pub fn filter(&mut self, test_name_slice: &str) {
        for (module_id, module_test) in self.tests.module_tests.iter_mut() {
            if module_id.name().as_str().contains(test_name_slice) {
                continue;
            } else {
                let tests = std::mem::take(&mut module_test.tests);
                module_test.tests = tests
                    .into_iter()
                    .filter(|(test_name, _)| test_name.as_str().contains(test_name_slice))
                    .collect();
            }
        }
    }
}

impl SharedTestingConfig {
    fn execute_via_move_vm(
        &self,
        test_plan: &ModuleTestPlan,
        function_name: &str,
        test_info: &TestCase,
    ) -> (VMResult<ChangeSet>, VMResult<Vec<Vec<u8>>>, TestRunInfo) {
        let move_vm = MoveVM::new();
        let mut session = move_vm.new_session(&self.starting_storage_state);
        let mut gas_meter = GasStatus::new(&self.cost_table, GasUnits::new(self.execution_bound));
        // TODO: collect VM logs if the verbose flag (i.e, `self.verbose`) is set
        let log_context = NoContextLog::new();

        let now = Instant::now();
        let return_result = session.execute_function(
            &test_plan.module_id,
            &IdentStr::new(function_name).unwrap(),
            vec![], // no ty args, at least for now
            serialize_values(test_info.arguments.iter()),
            &mut gas_meter,
            &log_context,
        );
        let test_run_info = TestRunInfo::new(
            function_name.to_string(),
            now.elapsed(),
            self.execution_bound - gas_meter.remaining_gas().get(),
        );
        (
            session.finish().map(|(cs, _)| cs),
            return_result,
            test_run_info,
        )
    }

    fn execute_via_stackless_vm(
        &self,
        env: &GlobalEnv,
        test_plan: &ModuleTestPlan,
        function_name: &str,
        test_info: &TestCase,
    ) -> (VMResult<ChangeSet>, VMResult<Vec<Vec<u8>>>, TestRunInfo) {
        let now = Instant::now();

        // NOTE: as of now, `self.starting_storage_state` contains modules only and no resources.
        // The modules are captured by `env: &GlobalEnv` and the default GlobalState captures the
        // empty-resource state.
        let interpreter =
            StacklessBytecodeInterpreter::new(env, None, InterpreterSettings::default());
        let (return_result, change_set, _) = interpreter.interpret(
            &test_plan.module_id,
            &IdentStr::new(function_name).unwrap(),
            &[], // no ty args, at least for now
            &test_info.arguments,
            &GlobalState::default(),
        );

        let test_run_info = TestRunInfo::new(
            function_name.to_string(),
            now.elapsed(),
            // NOTE (mengxu) instruction counting on stackless VM might not be very useful because
            // gas is not charged against stackless VM instruction.
            0,
        );
        (Ok(change_set), return_result, test_run_info)
    }

    fn exec_module_tests<W: Write>(
        &self,
        test_plan: &ModuleTestPlan,
        writer: &Mutex<W>,
    ) -> TestStatistics {
        let mut stats = TestStatistics::new();
        let pass = |fn_name: &str| {
            writeln!(
                writer.lock().unwrap(),
                "[ {}    ] {}::{}",
                "PASS".bold().bright_green(),
                format_module_id(&test_plan.module_id),
                fn_name
            )
            .unwrap()
        };
        let fail = |fn_name: &str| {
            writeln!(
                writer.lock().unwrap(),
                "[ {}    ] {}::{}",
                "FAIL".bold().bright_red(),
                format_module_id(&test_plan.module_id),
                fn_name,
            )
            .unwrap()
        };
        let timeout = |fn_name: &str| {
            writeln!(
                writer.lock().unwrap(),
                "[ {} ] {}::{}",
                "TIMEOUT".bold().bright_yellow(),
                format_module_id(&test_plan.module_id),
                fn_name,
            )
            .unwrap();
        };

        let stackless_model = if self.check_stackless_vm {
            let model =
                run_model_builder_with_compilation_flags(&self.source_files, &[], Flags::testing())
                    .unwrap_or_else(|e| panic!("Unable to build stackless bytecode: {}", e));
            Some(model)
        } else {
            None
        };

        for (function_name, test_info) in &test_plan.tests {
            let (cs_result, exec_result, test_run_info) =
                self.execute_via_move_vm(test_plan, function_name, test_info);
            if self.check_stackless_vm {
                let (stackless_vm_change_set, stackless_vm_result, _) = self
                    .execute_via_stackless_vm(
                        stackless_model.as_ref().unwrap(),
                        test_plan,
                        function_name,
                        test_info,
                    );
                let move_vm_result = adapt_move_vm_result(exec_result.clone());
                let move_vm_change_set =
                    adapt_move_vm_change_set(cs_result.clone(), &self.starting_storage_state);
                if stackless_vm_result != move_vm_result
                    || stackless_vm_change_set != move_vm_change_set
                {
                    fail(function_name);
                    stats.test_failure(
                        TestFailure::new(
                            FailureReason::mismatch(
                                move_vm_result,
                                move_vm_change_set,
                                stackless_vm_result,
                                stackless_vm_change_set,
                            ),
                            test_run_info,
                            None,
                            None,
                        ),
                        &test_plan,
                    );
                    continue;
                }
            };

            let save_session_state = || {
                if self.save_storage_state_on_failure {
                    cs_result.ok().and_then(|changeset| {
                        print_resources(&changeset, &self.starting_storage_state).ok()
                    })
                } else {
                    None
                }
            };
            match exec_result {
                Err(err) => match (test_info.expected_failure.as_ref(), err.sub_status()) {
                    // Ran out of ticks, report a test timeout and log a test failure
                    _ if err.major_status() == StatusCode::OUT_OF_GAS => {
                        timeout(function_name);
                        stats.test_failure(
                            TestFailure::new(
                                FailureReason::timeout(),
                                test_run_info,
                                Some(err),
                                save_session_state(),
                            ),
                            &test_plan,
                        )
                    }
                    // Expected the test to not abort, but it aborted with `code`
                    (None, Some(code)) => {
                        fail(function_name);
                        stats.test_failure(
                            TestFailure::new(
                                FailureReason::aborted(code),
                                test_run_info,
                                Some(err),
                                save_session_state(),
                            ),
                            &test_plan,
                        )
                    }
                    // Expected the test the abort with a specific `code`, and it did abort with
                    // that abort code
                    (Some(ExpectedFailure::ExpectedWithCode(code)), Some(other_code))
                        if err.major_status() == StatusCode::ABORTED && *code == other_code =>
                    {
                        pass(function_name);
                        stats.test_success(test_run_info, &test_plan);
                    }
                    // Expected the test to abort with a specific `code` but it aborted with a
                    // different `other_code`
                    (Some(ExpectedFailure::ExpectedWithCode(code)), Some(other_code)) => {
                        fail(function_name);
                        stats.test_failure(
                            TestFailure::new(
                                FailureReason::wrong_abort(*code, other_code),
                                test_run_info,
                                Some(err),
                                save_session_state(),
                            ),
                            &test_plan,
                        )
                    }
                    // Expected the test to abort and it aborted, but we don't need to check the code
                    (Some(ExpectedFailure::Expected), Some(_)) => {
                        pass(function_name);
                        stats.test_success(test_run_info, &test_plan);
                    }
                    // Expected the test to abort and it aborted with internal error
                    (Some(ExpectedFailure::Expected), None)
                        if err.major_status() != StatusCode::EXECUTED =>
                    {
                        pass(function_name);
                        stats.test_success(test_run_info, &test_plan);
                    }
                    // Unexpected return status from the VM, signal that we hit an unknown error.
                    (_, None) => {
                        fail(function_name);
                        stats.test_failure(
                            TestFailure::new(
                                FailureReason::unknown(),
                                test_run_info,
                                Some(err),
                                save_session_state(),
                            ),
                            &test_plan,
                        )
                    }
                },
                Ok(_) => {
                    // Expected the test to fail, but it executed
                    if test_info.expected_failure.is_some() {
                        fail(function_name);
                        stats.test_failure(
                            TestFailure::new(
                                FailureReason::no_abort(),
                                test_run_info,
                                None,
                                save_session_state(),
                            ),
                            &test_plan,
                        )
                    } else {
                        // Expected the test to execute fully and it did
                        pass(function_name);
                        stats.test_success(test_run_info, &test_plan);
                    }
                }
            }
        }

        stats
    }
}
<commit_msg>[move-unit-test] bring back the verbose mode<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    format_module_id,
    test_reporter::{FailureReason, TestFailure, TestResults, TestRunInfo, TestStatistics},
};
use anyhow::Result;
use bytecode_interpreter::{
    concrete::{settings::InterpreterSettings, value::GlobalState},
    shared::bridge::{adapt_move_vm_change_set, adapt_move_vm_result},
    StacklessBytecodeInterpreter,
};
use colored::*;
use move_binary_format::{errors::VMResult, file_format::CompiledModule};
use move_core_types::{
    effects::ChangeSet,
    gas_schedule::{CostTable, GasAlgebra, GasCost, GasUnits},
    identifier::IdentStr,
    value::serialize_values,
    vm_status::StatusCode,
};
use move_lang::{
    shared::Flags,
    unit_test::{ExpectedFailure, ModuleTestPlan, TestCase, TestPlan},
};
use move_model::{model::GlobalEnv, run_model_builder_with_compilation_flags};
use move_vm_runtime::{logging::NoContextLog, move_vm::MoveVM};
use move_vm_test_utils::InMemoryStorage;
use move_vm_types::gas_schedule::{zero_cost_schedule, GasStatus};
use rayon::prelude::*;
use resource_viewer::MoveValueAnnotator;
use std::{io::Write, marker::Send, sync::Mutex, time::Instant};

/// Test state common to all tests
#[derive(Debug)]
pub struct SharedTestingConfig {
    save_storage_state_on_failure: bool,
    execution_bound: u64,
    cost_table: CostTable,
    starting_storage_state: InMemoryStorage,
    source_files: Vec<String>,
    check_stackless_vm: bool,
    verbose: bool,
}

#[derive(Debug)]
pub struct TestRunner {
    num_threads: usize,
    testing_config: SharedTestingConfig,
    tests: TestPlan,
}

/// A gas schedule where every instruction has a cost of "1". This is used to bound execution of a
/// test to a certain number of ticks.
fn unit_cost_table() -> CostTable {
    let mut cost_schedule = zero_cost_schedule();
    cost_schedule.instruction_table.iter_mut().for_each(|cost| {
        *cost = GasCost::new(1, 1);
    });
    cost_schedule.native_table.iter_mut().for_each(|cost| {
        *cost = GasCost::new(1, 1);
    });
    cost_schedule
}

/// Setup storage state with the set of modules that will be needed for all tests
fn setup_test_storage<'a>(
    modules: impl Iterator<Item = &'a CompiledModule>,
) -> Result<InMemoryStorage> {
    let mut storage = InMemoryStorage::new();
    for module in modules {
        let module_id = module.self_id();
        let mut module_bytes = Vec::new();
        module.serialize(&mut module_bytes)?;
        storage.publish_or_overwrite_module(module_id, module_bytes);
    }

    Ok(storage)
}

/// Print the updates to storage represented by `cs` in the context of the starting storage state
/// `storage`.
fn print_resources(cs: &ChangeSet, storage: &InMemoryStorage) -> Result<String> {
    use std::fmt::Write;
    let mut buf = String::new();
    let annotator = MoveValueAnnotator::new(storage);
    for (account_addr, account_state) in cs.accounts() {
        writeln!(&mut buf, "0x{}:", account_addr.short_str_lossless())?;

        for (tag, resource_opt) in account_state.resources() {
            if let Some(resource) = resource_opt {
                writeln!(
                    &mut buf,
                    "\t{}",
                    format!("=> {}", annotator.view_resource(tag, resource)?).replace("\n", "\n\t")
                )?;
            }
        }
    }

    Ok(buf)
}

impl TestRunner {
    pub fn new(
        execution_bound: u64,
        num_threads: usize,
        check_stackless_vm: bool,
        verbose: bool,
        save_storage_state_on_failure: bool,
        tests: TestPlan,
    ) -> Result<Self> {
        let source_files = tests
            .files
            .keys()
            .map(|filepath| filepath.to_string())
            .collect();
        let modules = tests.module_info.values().map(|info| &info.0);
        let starting_storage_state = setup_test_storage(modules)?;
        Ok(Self {
            testing_config: SharedTestingConfig {
                save_storage_state_on_failure,
                starting_storage_state,
                execution_bound,
                cost_table: unit_cost_table(),
                source_files,
                check_stackless_vm,
                verbose,
            },
            num_threads,
            tests,
        })
    }

    pub fn run<W: Write + Send>(self, writer: &Mutex<W>) -> Result<TestResults> {
        rayon::ThreadPoolBuilder::new()
            .num_threads(self.num_threads)
            .build()
            .unwrap()
            .install(|| {
                let final_statistics = self
                    .tests
                    .module_tests
                    .par_iter()
                    .map(|(_, test_plan)| self.testing_config.exec_module_tests(test_plan, writer))
                    .reduce(TestStatistics::new, |acc, stats| acc.combine(stats));

                Ok(TestResults::new(final_statistics, self.tests))
            })
    }

    pub fn filter(&mut self, test_name_slice: &str) {
        for (module_id, module_test) in self.tests.module_tests.iter_mut() {
            if module_id.name().as_str().contains(test_name_slice) {
                continue;
            } else {
                let tests = std::mem::take(&mut module_test.tests);
                module_test.tests = tests
                    .into_iter()
                    .filter(|(test_name, _)| test_name.as_str().contains(test_name_slice))
                    .collect();
            }
        }
    }
}

impl SharedTestingConfig {
    fn execute_via_move_vm(
        &self,
        test_plan: &ModuleTestPlan,
        function_name: &str,
        test_info: &TestCase,
    ) -> (VMResult<ChangeSet>, VMResult<Vec<Vec<u8>>>, TestRunInfo) {
        let move_vm = MoveVM::new();
        let mut session = move_vm.new_session(&self.starting_storage_state);
        let mut gas_meter = GasStatus::new(&self.cost_table, GasUnits::new(self.execution_bound));
        // TODO: collect VM logs if the verbose flag (i.e, `self.verbose`) is set
        let log_context = NoContextLog::new();

        let now = Instant::now();
        let return_result = session.execute_function(
            &test_plan.module_id,
            &IdentStr::new(function_name).unwrap(),
            vec![], // no ty args, at least for now
            serialize_values(test_info.arguments.iter()),
            &mut gas_meter,
            &log_context,
        );
        let test_run_info = TestRunInfo::new(
            function_name.to_string(),
            now.elapsed(),
            self.execution_bound - gas_meter.remaining_gas().get(),
        );
        (
            session.finish().map(|(cs, _)| cs),
            return_result,
            test_run_info,
        )
    }

    fn execute_via_stackless_vm(
        &self,
        env: &GlobalEnv,
        test_plan: &ModuleTestPlan,
        function_name: &str,
        test_info: &TestCase,
    ) -> (VMResult<ChangeSet>, VMResult<Vec<Vec<u8>>>, TestRunInfo) {
        let now = Instant::now();

        let settings = if self.verbose {
            InterpreterSettings::verbose_default()
        } else {
            InterpreterSettings::default()
        };
        let interpreter = StacklessBytecodeInterpreter::new(env, None, settings);

        // NOTE: as of now, `self.starting_storage_state` contains modules only and no resources.
        // The modules are captured by `env: &GlobalEnv` and the default GlobalState captures the
        // empty-resource state.
        let global_state = GlobalState::default();
        let (return_result, change_set, _) = interpreter.interpret(
            &test_plan.module_id,
            &IdentStr::new(function_name).unwrap(),
            &[], // no ty args, at least for now
            &test_info.arguments,
            &global_state,
        );

        let test_run_info = TestRunInfo::new(
            function_name.to_string(),
            now.elapsed(),
            // NOTE (mengxu) instruction counting on stackless VM might not be very useful because
            // gas is not charged against stackless VM instruction.
            0,
        );
        (Ok(change_set), return_result, test_run_info)
    }

    fn exec_module_tests<W: Write>(
        &self,
        test_plan: &ModuleTestPlan,
        writer: &Mutex<W>,
    ) -> TestStatistics {
        let mut stats = TestStatistics::new();
        let pass = |fn_name: &str| {
            writeln!(
                writer.lock().unwrap(),
                "[ {}    ] {}::{}",
                "PASS".bold().bright_green(),
                format_module_id(&test_plan.module_id),
                fn_name
            )
            .unwrap()
        };
        let fail = |fn_name: &str| {
            writeln!(
                writer.lock().unwrap(),
                "[ {}    ] {}::{}",
                "FAIL".bold().bright_red(),
                format_module_id(&test_plan.module_id),
                fn_name,
            )
            .unwrap()
        };
        let timeout = |fn_name: &str| {
            writeln!(
                writer.lock().unwrap(),
                "[ {} ] {}::{}",
                "TIMEOUT".bold().bright_yellow(),
                format_module_id(&test_plan.module_id),
                fn_name,
            )
            .unwrap();
        };

        let stackless_model = if self.check_stackless_vm {
            let model =
                run_model_builder_with_compilation_flags(&self.source_files, &[], Flags::testing())
                    .unwrap_or_else(|e| panic!("Unable to build stackless bytecode: {}", e));
            Some(model)
        } else {
            None
        };

        for (function_name, test_info) in &test_plan.tests {
            let (cs_result, exec_result, test_run_info) =
                self.execute_via_move_vm(test_plan, function_name, test_info);
            if self.check_stackless_vm {
                let (stackless_vm_change_set, stackless_vm_result, _) = self
                    .execute_via_stackless_vm(
                        stackless_model.as_ref().unwrap(),
                        test_plan,
                        function_name,
                        test_info,
                    );
                let move_vm_result = adapt_move_vm_result(exec_result.clone());
                let move_vm_change_set =
                    adapt_move_vm_change_set(cs_result.clone(), &self.starting_storage_state);
                if stackless_vm_result != move_vm_result
                    || stackless_vm_change_set != move_vm_change_set
                {
                    fail(function_name);
                    stats.test_failure(
                        TestFailure::new(
                            FailureReason::mismatch(
                                move_vm_result,
                                move_vm_change_set,
                                stackless_vm_result,
                                stackless_vm_change_set,
                            ),
                            test_run_info,
                            None,
                            None,
                        ),
                        &test_plan,
                    );
                    continue;
                }
            };

            let save_session_state = || {
                if self.save_storage_state_on_failure {
                    cs_result.ok().and_then(|changeset| {
                        print_resources(&changeset, &self.starting_storage_state).ok()
                    })
                } else {
                    None
                }
            };
            match exec_result {
                Err(err) => match (test_info.expected_failure.as_ref(), err.sub_status()) {
                    // Ran out of ticks, report a test timeout and log a test failure
                    _ if err.major_status() == StatusCode::OUT_OF_GAS => {
                        timeout(function_name);
                        stats.test_failure(
                            TestFailure::new(
                                FailureReason::timeout(),
                                test_run_info,
                                Some(err),
                                save_session_state(),
                            ),
                            &test_plan,
                        )
                    }
                    // Expected the test to not abort, but it aborted with `code`
                    (None, Some(code)) => {
                        fail(function_name);
                        stats.test_failure(
                            TestFailure::new(
                                FailureReason::aborted(code),
                                test_run_info,
                                Some(err),
                                save_session_state(),
                            ),
                            &test_plan,
                        )
                    }
                    // Expected the test the abort with a specific `code`, and it did abort with
                    // that abort code
                    (Some(ExpectedFailure::ExpectedWithCode(code)), Some(other_code))
                        if err.major_status() == StatusCode::ABORTED && *code == other_code =>
                    {
                        pass(function_name);
                        stats.test_success(test_run_info, &test_plan);
                    }
                    // Expected the test to abort with a specific `code` but it aborted with a
                    // different `other_code`
                    (Some(ExpectedFailure::ExpectedWithCode(code)), Some(other_code)) => {
                        fail(function_name);
                        stats.test_failure(
                            TestFailure::new(
                                FailureReason::wrong_abort(*code, other_code),
                                test_run_info,
                                Some(err),
                                save_session_state(),
                            ),
                            &test_plan,
                        )
                    }
                    // Expected the test to abort and it aborted, but we don't need to check the code
                    (Some(ExpectedFailure::Expected), Some(_)) => {
                        pass(function_name);
                        stats.test_success(test_run_info, &test_plan);
                    }
                    // Expected the test to abort and it aborted with internal error
                    (Some(ExpectedFailure::Expected), None)
                        if err.major_status() != StatusCode::EXECUTED =>
                    {
                        pass(function_name);
                        stats.test_success(test_run_info, &test_plan);
                    }
                    // Unexpected return status from the VM, signal that we hit an unknown error.
                    (_, None) => {
                        fail(function_name);
                        stats.test_failure(
                            TestFailure::new(
                                FailureReason::unknown(),
                                test_run_info,
                                Some(err),
                                save_session_state(),
                            ),
                            &test_plan,
                        )
                    }
                },
                Ok(_) => {
                    // Expected the test to fail, but it executed
                    if test_info.expected_failure.is_some() {
                        fail(function_name);
                        stats.test_failure(
                            TestFailure::new(
                                FailureReason::no_abort(),
                                test_run_info,
                                None,
                                save_session_state(),
                            ),
                            &test_plan,
                        )
                    } else {
                        // Expected the test to execute fully and it did
                        pass(function_name);
                        stats.test_success(test_run_info, &test_plan);
                    }
                }
            }
        }

        stats
    }
}
<|endoftext|>


<commit_before>use futures::{Poll, Future, Async, BoxFuture};
use tokio_service::Service;
use store::Store;
use {LogPos, Error};
use std::net::SocketAddr;
use std::io;
use std::mem;
use replica::{ReplicaRequest, ReplicaResponse, LogEntry};

enum ReplicatorState<S: Store, D>
    where D: DownstreamService
{
    Idle,
    Fetching(S::FetchFut),
    Forwarding(D::Future),
}

pub trait DownstreamService {
    type Future: Future<Item = ReplicaResponse, Error = io::Error>;
}

impl<D> DownstreamService for D
    where D: Service<Request = (SocketAddr, ReplicaRequest),
                     Response = ReplicaResponse,
                     Error = io::Error,
                     Future = BoxFuture<ReplicaResponse, io::Error>>
{
    type Future = D::Future;
}

pub struct Replicator<S: Store, D: DownstreamService> {
    store: S,
    downstream: D,
    last_seen_seq: LogPos,
    downstream_addr: Option<SocketAddr>,
    state: ReplicatorState<S, D>,
}

impl<S: Store, D: DownstreamService> Replicator<S, D> {
    pub fn new(store: S, downstream: D) -> Self {
        Replicator {
            store: store,
            downstream: downstream,
            last_seen_seq: LogPos::zero(),
            downstream_addr: None,
            state: ReplicatorState::Idle,
        }
    }
}

impl<S: Store, D: DownstreamService> Drop for Replicator<S, D> {
    fn drop(&mut self) {
        debug!("Drop Replicator");
    }
}

impl<S: Store, D> Future for Replicator<S, D>
    where S: Store,
          D: Service<Request = (SocketAddr, ReplicaRequest),
                     Response = ReplicaResponse,
                     Error = io::Error,
                     Future = BoxFuture<ReplicaResponse, io::Error>>
{
    type Item = ();
    type Error = Error;
    fn poll(&mut self) -> Poll<(), Error> {
        debug!("Replicator#poll");
        loop {
            match mem::replace(&mut self.state, ReplicatorState::Idle) {
                ReplicatorState::Idle => {
                    debug!("Idle, fetching after: {:?}", self.last_seen_seq);
                    let fetch_f = self.store.fetch_next(self.last_seen_seq);
                    self.state = ReplicatorState::Fetching(fetch_f);
                }
                // Now, what have we learnt from this, children? Surely it's
                // that because we're forwarding LogEntries at least, we
                // should make it easy to get a LogEntry into/out of the
                // store, rather than having to futz around with serialisation
                // rubbish.
                ReplicatorState::Fetching(mut fetch_f) => {
                    debug!("Fetching");
                    match try!(fetch_f.poll()) {
                        Async::NotReady => {
                            self.state = ReplicatorState::Fetching(fetch_f);
                            return Ok(Async::NotReady);
                        }
                        Async::Ready((off, val)) => {
                            debug!("Fetched: {:?}", (&off, &val));
                            if let &LogEntry::Config(ref conf) = &val {
                                debug!("logged Config message: {:?}", conf);
                                // FIXME: Well, this is blatantly wrong.
                                self.downstream_addr = Some(conf.head);
                            };

                            if let Some(addr) = self.downstream_addr {
                                debug!("Forward {:?} to: {:?}", off, addr);
                                // XXX: Symmetry with ReplClient?
                                let req = ReplicaRequest::AppendLogEntry {
                                    assumed_offset: self.last_seen_seq,
                                    entry_offset: off,
                                    datum: val,
                                };
                                let f = self.downstream.call((addr, req));
                                self.state = ReplicatorState::Forwarding(f);
                            } else {
                                debug!("No downstream at {:?}", off);
                            }
                            self.last_seen_seq = off;
                        }
                    }
                }
                ReplicatorState::Forwarding(mut f) => {
                    match try!(f.poll()) {
                        Async::NotReady => {
                            self.state = ReplicatorState::Forwarding(f);
                            return Ok(Async::NotReady);
                        }
                        Async::Ready(ReplicaResponse::Done(pos)) => {
                            debug!("woo! {:?}", pos);
                        }
                        Async::Ready(ReplicaResponse::BadSequence(pos)) => {
                            panic!("ARGH! {:?}", pos);
                        }
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod test {
    use futures::{self, Future, BoxFuture, Async};
    use futures::stream::Stream;
    use service::Service;
    use tokio::channel;
    use replica::LogPos;
    use store::{RamStore, Store};
    use replica::{ReplicaRequest, ReplicaResponse, LogEntry, HostConfig};
    use std::net::SocketAddr;
    use super::*;
    use tokio::reactor::Core;
    use std::time::Duration;
    use std::io;
    use tokio_timer;
    use env_logger;

    struct MyMagicalDownstream {
        send: channel::Sender<((SocketAddr, ReplicaRequest), futures::Complete<ReplicaResponse>)>,
    }

    impl Service for MyMagicalDownstream {
        type Request = (SocketAddr, ReplicaRequest);
        type Response = ReplicaResponse;
        type Error = io::Error;
        type Future = BoxFuture<ReplicaResponse, io::Error>;
        fn poll_ready(&self) -> Async<()> {
            Async::Ready(())
        }
        fn call(&self, req: Self::Request) -> Self::Future {
            let (c, p) = futures::oneshot();
            match self.send.send((req, c)) {
                Ok(()) => (),
                Err(e) => return futures::failed(e).boxed(),
            };
            p.map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "broken pipe"))
                .boxed()
        }
    }

    #[test]
    // #[ignore] // WIP
    fn should_start_replicating_to_downstream_on_committed_config_message() {
        env_logger::init().unwrap_or(());
        let store = RamStore::new();
        let timer = tokio_timer::wheel().tick_duration(Duration::from_millis(1)).build();

        let mut core = Core::new().expect("core::new");
        let (tx, rx) = channel::channel(&core.handle()).expect("channel");

        let downstream = MyMagicalDownstream { send: tx };

        let replica = Replicator::new(store.clone(), downstream);

        core.handle().spawn(replica.map_err(|e| panic!("Replicator failed!: {:?}", e)));
        debug!("Spawned replica task");

        let off = LogPos::zero();

        let head_addr = "1.2.3.4:5".parse().expect("parse");

        let config: HostConfig<::std::net::SocketAddr> = HostConfig {
            head: head_addr,
            tail: "1.2.3.6:7".parse().expect("parse"),
        };

        debug!("append config message");
        core.run(store.append_entry(off, off.next(), LogEntry::Config(config.clone())))
            .expect("append");

        let f = timer.timeout(rx.into_future().map_err(|(e, _)| e),
                              Duration::from_millis(1000));
        debug!("wait for stub message");
        let (data, _rx) = core.run(f).expect("receive downstream");
        let ((addr, msg), _completer) = data.expect("Some message");
        println!("Some: {:?}", (&addr, &msg));
        // Expect that the target at `head_addr` receives a set of replication messages.
        let (_assumed, _off, entry) = match msg {
            ReplicaRequest::AppendLogEntry { assumed_offset, entry_offset, datum } => {
                (assumed_offset, entry_offset, datum)
            }
        };
        assert_eq!((addr, entry), (head_addr, LogEntry::Config(config)));
    }
}
<commit_msg>Show that the replicator will send the next message too.<commit_after>use futures::{Poll, Future, Async, BoxFuture};
use tokio_service::Service;
use store::Store;
use {LogPos, Error};
use std::net::SocketAddr;
use std::io;
use std::mem;
use replica::{ReplicaRequest, ReplicaResponse, LogEntry};

enum ReplicatorState<S: Store, D>
    where D: DownstreamService
{
    Idle,
    Fetching(S::FetchFut),
    Forwarding(D::Future),
}

pub trait DownstreamService {
    type Future: Future<Item = ReplicaResponse, Error = io::Error>;
}

impl<D> DownstreamService for D
    where D: Service<Request = (SocketAddr, ReplicaRequest),
                     Response = ReplicaResponse,
                     Error = io::Error,
                     Future = BoxFuture<ReplicaResponse, io::Error>>
{
    type Future = D::Future;
}

pub struct Replicator<S: Store, D: DownstreamService> {
    store: S,
    downstream: D,
    last_seen_seq: LogPos,
    downstream_addr: Option<SocketAddr>,
    state: ReplicatorState<S, D>,
}

impl<S: Store, D: DownstreamService> Replicator<S, D> {
    pub fn new(store: S, downstream: D) -> Self {
        Replicator {
            store: store,
            downstream: downstream,
            last_seen_seq: LogPos::zero(),
            downstream_addr: None,
            state: ReplicatorState::Idle,
        }
    }
}

impl<S: Store, D: DownstreamService> Drop for Replicator<S, D> {
    fn drop(&mut self) {
        debug!("Drop Replicator");
    }
}

impl<S: Store, D> Future for Replicator<S, D>
    where S: Store,
          D: Service<Request = (SocketAddr, ReplicaRequest),
                     Response = ReplicaResponse,
                     Error = io::Error,
                     Future = BoxFuture<ReplicaResponse, io::Error>>
{
    type Item = ();
    type Error = Error;
    fn poll(&mut self) -> Poll<(), Error> {
        debug!("Replicator#poll");
        loop {
            match mem::replace(&mut self.state, ReplicatorState::Idle) {
                ReplicatorState::Idle => {
                    debug!("Idle, fetching after: {:?}", self.last_seen_seq);
                    let fetch_f = self.store.fetch_next(self.last_seen_seq);
                    self.state = ReplicatorState::Fetching(fetch_f);
                }
                // Now, what have we learnt from this, children? Surely it's
                // that because we're forwarding LogEntries at least, we
                // should make it easy to get a LogEntry into/out of the
                // store, rather than having to futz around with serialisation
                // rubbish.
                ReplicatorState::Fetching(mut fetch_f) => {
                    debug!("Fetching");
                    match try!(fetch_f.poll()) {
                        Async::NotReady => {
                            self.state = ReplicatorState::Fetching(fetch_f);
                            return Ok(Async::NotReady);
                        }
                        Async::Ready((off, val)) => {
                            debug!("Fetched: {:?}", (&off, &val));
                            if let &LogEntry::Config(ref conf) = &val {
                                debug!("logged Config message: {:?}", conf);
                                // FIXME: Well, this is blatantly wrong.
                                self.downstream_addr = Some(conf.head);
                            };

                            if let Some(addr) = self.downstream_addr {
                                debug!("Forward {:?} to: {:?}", off, addr);
                                // XXX: Symmetry with ReplClient?
                                let req = ReplicaRequest::AppendLogEntry {
                                    assumed_offset: self.last_seen_seq,
                                    entry_offset: off,
                                    datum: val,
                                };
                                let f = self.downstream.call((addr, req));
                                self.state = ReplicatorState::Forwarding(f);
                            } else {
                                debug!("No downstream at {:?}", off);
                            }
                            self.last_seen_seq = off;
                        }
                    }
                }
                ReplicatorState::Forwarding(mut f) => {
                    match try!(f.poll()) {
                        Async::NotReady => {
                            self.state = ReplicatorState::Forwarding(f);
                            return Ok(Async::NotReady);
                        }
                        Async::Ready(ReplicaResponse::Done(pos)) => {
                            debug!("woo! {:?}", pos);
                        }
                        Async::Ready(ReplicaResponse::BadSequence(pos)) => {
                            panic!("ARGH! {:?}", pos);
                        }
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod test {
    use futures::{self, Future, BoxFuture, Async};
    use futures::stream::Stream;
    use service::Service;
    use tokio::channel;
    use replica::LogPos;
    use store::{RamStore, Store};
    use replica::{ReplicaRequest, ReplicaResponse, LogEntry, HostConfig};
    use std::net::SocketAddr;
    use super::*;
    use tokio::reactor::Core;
    use std::time::Duration;
    use std::io;
    use tokio_timer;
    use env_logger;

    struct MyMagicalDownstream {
        send: channel::Sender<((SocketAddr, ReplicaRequest), futures::Complete<ReplicaResponse>)>,
    }

    impl Service for MyMagicalDownstream {
        type Request = (SocketAddr, ReplicaRequest);
        type Response = ReplicaResponse;
        type Error = io::Error;
        type Future = BoxFuture<ReplicaResponse, io::Error>;
        fn poll_ready(&self) -> Async<()> {
            Async::Ready(())
        }
        fn call(&self, req: Self::Request) -> Self::Future {
            let (c, p) = futures::oneshot();
            match self.send.send((req, c)) {
                Ok(()) => (),
                Err(e) => return futures::failed(e).boxed(),
            };
            p.map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "broken pipe"))
                .boxed()
        }
    }

    #[test]
    fn should_start_replicating_to_downstream_on_committed_config_message() {
        env_logger::init().unwrap_or(());
        let store = RamStore::new();
        let timer = tokio_timer::wheel().tick_duration(Duration::from_millis(1)).build();

        let mut core = Core::new().expect("core::new");
        let (tx, rx) = channel::channel(&core.handle()).expect("channel");

        let downstream = MyMagicalDownstream { send: tx };

        let replica = Replicator::new(store.clone(), downstream);

        core.handle().spawn(replica.map_err(|e| panic!("Replicator failed!: {:?}", e)));
        debug!("Spawned replica task");

        let off = LogPos::zero();

        let head_addr = "1.2.3.4:5".parse().expect("parse");

        let config: HostConfig<::std::net::SocketAddr> = HostConfig {
            head: head_addr,
            tail: "1.2.3.6:7".parse().expect("parse"),
        };

        debug!("append config message");
        core.run(store.append_entry(off, off.next(), LogEntry::Config(config.clone())))
            .expect("append");

        let f = timer.timeout(rx.into_future().map_err(|(e, _)| e),
                              Duration::from_millis(1000));
        debug!("wait for stub message");
        let (data, _rx) = core.run(f).expect("receive downstream");
        let ((addr, msg), _completer) = data.expect("Some message");
        println!("Some: {:?}", (&addr, &msg));
        let (_assumed, _off, entry) = match msg {
            ReplicaRequest::AppendLogEntry { assumed_offset, entry_offset, datum } => {
                (assumed_offset, entry_offset, datum)
            }
        };
        assert_eq!((addr, entry), (head_addr, LogEntry::Config(config)));
    }

    #[test]
    fn should_send_next_log_entry_on_okay_from_downstream() {
        env_logger::init().unwrap_or(());
        let store = RamStore::new();
        let timer = tokio_timer::wheel().tick_duration(Duration::from_millis(1)).build();

        let mut core = Core::new().expect("core::new");
        let (tx, rx) = channel::channel(&core.handle()).expect("channel");

        let downstream = MyMagicalDownstream { send: tx };

        let replica = Replicator::new(store.clone(), downstream);

        core.handle().spawn(replica.map_err(|e| panic!("Replicator failed!: {:?}", e)));
        debug!("Spawned replica task");

        let head_addr = "1.2.3.4:5".parse().expect("parse");

        let config: HostConfig<::std::net::SocketAddr> = HostConfig {
            head: head_addr,
            tail: "1.2.3.6:7".parse().expect("parse"),
        };

        {
            let off = LogPos::zero();
            let next = off.next();
            debug!("append config message {:?} -> {:?}", off, next);
            core.run(store.append_entry(off, next, LogEntry::Config(config.clone())))
                .expect("append");

            let off = next;
            let next = off.next();
            debug!("append data message {:?} -> {:?}", off, next);
            core.run(store.append_entry(off, next, LogEntry::Data(b"Hello world!".to_vec())))
                .expect("append");
        }

        let f = timer.timeout(rx.into_future().map_err(|(e, _)| e),
                              Duration::from_millis(1000));
        debug!("wait for stub message");
        let (data, rx) = core.run(f).expect("receive downstream");
        let ((addr, msg), completer) = data.expect("Some message");
        println!("Some: {:?}", (&addr, &msg));
        // Expect that the target at `head_addr` receives a set of replication messages.
        let (_assumed0, off0, _entry0) = match msg {
            ReplicaRequest::AppendLogEntry { assumed_offset, entry_offset, datum } => {
                (assumed_offset, entry_offset, datum)
            }
        };

        completer.complete(ReplicaResponse::Done(off0));

        let f = timer.timeout(rx.into_future().map_err(|(e, _)| e),
                              Duration::from_millis(1000));
        debug!("wait for stub message");
        let (data, _rx) = core.run(f).expect("receive downstream");
        let ((addr, msg), _completer) = data.expect("Some message");
        println!("Some: {:?}", (&addr, &msg));
        // Expect that the target at `head_addr` receives a set of replication messages.
        let (assumed1, off1, entry1) = match msg {
            ReplicaRequest::AppendLogEntry { assumed_offset, entry_offset, datum } => {
                (assumed_offset, entry_offset, datum)
            }
        };

        assert_eq!((addr, assumed1, entry1),
                   (head_addr, off0, LogEntry::Data(b"Hello world!".to_vec())));
        assert!(off1 > assumed1);
    }

    #[test]
    #[ignore]
    fn should_restart_replication_on_badsequence_from_downstream() {}

}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate aho_corasick;
extern crate test;

use std::iter;

use aho_corasick::{AcAutomaton, Transitions};
use test::Bencher;

const HAYSTACK_RANDOM: &'static str = include_str!("random.txt");

fn bench_aut_no_match<T: Transitions>(
    b: &mut Bencher,
    aut: AcAutomaton<T>,
    haystack: &str,
) {
    b.bytes = haystack.len() as u64;
    b.iter(|| assert!(aut.find(haystack).next().is_none()));
}

fn bench_naive_no_match<S>(b: &mut Bencher, needles: Vec<S>, haystack: &str)
        where S: Into<String> {
    b.bytes = haystack.len() as u64;
    let needles: Vec<String> = needles.into_iter().map(Into::into).collect();
    b.iter(|| assert!(!naive_find(&needles, haystack)));
}

fn haystack_same(letter: char) -> String {
    iter::repeat(letter).take(10000).collect()
}

macro_rules! aut_benches {
    ($prefix:ident, $aut:expr) => { mod $prefix {
use aho_corasick::AcAutomaton;
use test::Bencher;

use super::{HAYSTACK_RANDOM, bench_aut_no_match, haystack_same};

#[bench]
fn ac_one_byte(b: &mut Bencher) {
    let aut = $aut(vec!["a"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_one_prefix_byte_no_match(b: &mut Bencher) {
    let aut = $aut(vec!["zbc"]);
    bench_aut_no_match(b, aut, &haystack_same('y'));
}

#[bench]
fn ac_one_prefix_byte_every_match(b: &mut Bencher) {
    // We lose the benefit of `memchr` because the first byte matches
    // in every position in the haystack.
    let aut = $aut(vec!["zbc"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_one_prefix_byte_random(b: &mut Bencher) {
    let aut = $aut(vec!["zbc\x00"]);
    bench_aut_no_match(b, aut, HAYSTACK_RANDOM);
}

#[bench]
fn ac_two_bytes(b: &mut Bencher) {
    let aut = $aut(vec!["a", "b"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_two_diff_prefix(b: &mut Bencher) {
    let aut = $aut(vec!["abcdef", "bmnopq"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_two_one_prefix_byte_every_match(b: &mut Bencher) {
    let aut = $aut(vec!["zbcdef", "zmnopq"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_two_one_prefix_byte_no_match(b: &mut Bencher) {
    let aut = $aut(vec!["zbcdef", "zmnopq"]);
    bench_aut_no_match(b, aut, &haystack_same('y'));
}

#[bench]
fn ac_two_one_prefix_byte_random(b: &mut Bencher) {
    let aut = $aut(vec!["zbcdef\x00", "zmnopq\x00"]);
    bench_aut_no_match(b, aut, HAYSTACK_RANDOM);
}

#[bench]
fn ac_ten_bytes(b: &mut Bencher) {
    let aut = $aut(vec!["a", "b", "c", "d", "e",
                        "f", "g", "h", "i", "j"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_ten_diff_prefix(b: &mut Bencher) {
    let aut = $aut(vec!["abcdef", "bbcdef", "cbcdef", "dbcdef",
                        "ebcdef", "fbcdef", "gbcdef", "hbcdef",
                        "ibcdef", "jbcdef"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_ten_one_prefix_byte_every_match(b: &mut Bencher) {
    let aut = $aut(vec!["zacdef", "zbcdef", "zccdef", "zdcdef",
                        "zecdef", "zfcdef", "zgcdef", "zhcdef",
                        "zicdef", "zjcdef"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_ten_one_prefix_byte_no_match(b: &mut Bencher) {
    let aut = $aut(vec!["zacdef", "zbcdef", "zccdef", "zdcdef",
                        "zecdef", "zfcdef", "zgcdef", "zhcdef",
                        "zicdef", "zjcdef"]);
    bench_aut_no_match(b, aut, &haystack_same('y'));
}

#[bench]
fn ac_ten_one_prefix_byte_random(b: &mut Bencher) {
    let aut = $aut(vec!["zacdef\x00", "zbcdef\x00", "zccdef\x00",
                        "zdcdef\x00", "zecdef\x00", "zfcdef\x00",
                        "zgcdef\x00", "zhcdef\x00", "zicdef\x00",
                        "zjcdef\x00"]);
    bench_aut_no_match(b, aut, HAYSTACK_RANDOM);
}
}}}

aut_benches!(dense, AcAutomaton::new);
aut_benches!(sparse, AcAutomaton::<Sparse>::with_transitions);

// A naive multi-pattern search.
// We use this to benchmark *throughput*, so it should never match anything.
fn naive_find(needles: &[String], haystack: &str) -> bool {
    for hi in 0..haystack.len() {
        let rest = &haystack.as_bytes()[hi..];
        for needle in needles {
            let needle = needle.as_bytes();
            if needle.len() > rest.len() {
                continue;
            }
            if needle == &rest[..needle.len()] {
                // should never happen in throughput benchmarks.
                return true;
            }
        }
    }
    false
}

#[bench]
fn naive_one_byte(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["a"], &haystack_same('z'));
}

#[bench]
fn naive_one_prefix_byte_no_match(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["zbc"], &haystack_same('y'));
}

#[bench]
fn naive_one_prefix_byte_every_match(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["zbc"], &haystack_same('z'));
}

#[bench]
fn naive_one_prefix_byte_random(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["zbc\x00"], HAYSTACK_RANDOM);
}

#[bench]
fn naive_two_bytes(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["a", "b"], &haystack_same('z'));
}

#[bench]
fn naive_two_diff_prefix(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["abcdef", "bmnopq"], &haystack_same('z'));
}

#[bench]
fn naive_two_one_prefix_byte_every_match(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["zbcdef", "zmnopq"], &haystack_same('z'));
}

#[bench]
fn naive_two_one_prefix_byte_no_match(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["zbcdef", "zmnopq"], &haystack_same('y'));
}

#[bench]
fn naive_two_one_prefix_byte_random(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["zbcdef\x00", "zmnopq\x00"], HAYSTACK_RANDOM);
}

#[bench]
fn naive_ten_bytes(b: &mut Bencher) {
    let needles = vec!["a", "b", "c", "d", "e",
                       "f", "g", "h", "i", "j"];
    bench_naive_no_match(b, needles, &haystack_same('z'));
}

#[bench]
fn naive_ten_diff_prefix(b: &mut Bencher) {
    let needles = vec!["abcdef", "bbcdef", "cbcdef", "dbcdef",
                       "ebcdef", "fbcdef", "gbcdef", "hbcdef",
                       "ibcdef", "jbcdef"];
    bench_naive_no_match(b, needles, &haystack_same('z'));
}

#[bench]
fn naive_ten_one_prefix_byte_every_match(b: &mut Bencher) {
    let needles = vec!["zacdef", "zbcdef", "zccdef", "zdcdef",
                       "zecdef", "zfcdef", "zgcdef", "zhcdef",
                       "zicdef", "zjcdef"];
    bench_naive_no_match(b, needles, &haystack_same('z'));
}

#[bench]
fn naive_ten_one_prefix_byte_no_match(b: &mut Bencher) {
    let needles = vec!["zacdef", "zbcdef", "zccdef", "zdcdef",
                       "zecdef", "zfcdef", "zgcdef", "zhcdef",
                       "zicdef", "zjcdef"];
    bench_naive_no_match(b, needles, &haystack_same('y'));
}

#[bench]
fn naive_ten_one_prefix_byte_random(b: &mut Bencher) {
    let needles = vec!["zacdef\x00", "zbcdef\x00", "zccdef\x00",
                       "zdcdef\x00", "zecdef\x00", "zfcdef\x00",
                       "zgcdef\x00", "zhcdef\x00", "zicdef\x00",
                       "zjcdef\x00"];
    bench_naive_no_match(b, needles, HAYSTACK_RANDOM);
}
<commit_msg>fix benchmarks<commit_after>#![feature(test)]

extern crate aho_corasick;
extern crate test;

use std::iter;

use aho_corasick::{AcAutomaton, Transitions};
use test::Bencher;

const HAYSTACK_RANDOM: &'static str = include_str!("random.txt");

fn bench_aut_no_match<T: Transitions>(
    b: &mut Bencher,
    aut: AcAutomaton<T>,
    haystack: &str,
) {
    b.bytes = haystack.len() as u64;
    b.iter(|| assert!(aut.find(haystack).next().is_none()));
}

fn bench_naive_no_match<S>(b: &mut Bencher, needles: Vec<S>, haystack: &str)
        where S: Into<String> {
    b.bytes = haystack.len() as u64;
    let needles: Vec<String> = needles.into_iter().map(Into::into).collect();
    b.iter(|| assert!(!naive_find(&needles, haystack)));
}

fn haystack_same(letter: char) -> String {
    iter::repeat(letter).take(10000).collect()
}

macro_rules! aut_benches {
    ($prefix:ident, $aut:expr) => { mod $prefix {
#![allow(unused_imports)]
use aho_corasick::{AcAutomaton, Sparse};
use test::Bencher;

use super::{HAYSTACK_RANDOM, bench_aut_no_match, haystack_same};

#[bench]
fn ac_one_byte(b: &mut Bencher) {
    let aut = $aut(vec!["a"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_one_prefix_byte_no_match(b: &mut Bencher) {
    let aut = $aut(vec!["zbc"]);
    bench_aut_no_match(b, aut, &haystack_same('y'));
}

#[bench]
fn ac_one_prefix_byte_every_match(b: &mut Bencher) {
    // We lose the benefit of `memchr` because the first byte matches
    // in every position in the haystack.
    let aut = $aut(vec!["zbc"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_one_prefix_byte_random(b: &mut Bencher) {
    let aut = $aut(vec!["zbc\x00"]);
    bench_aut_no_match(b, aut, HAYSTACK_RANDOM);
}

#[bench]
fn ac_two_bytes(b: &mut Bencher) {
    let aut = $aut(vec!["a", "b"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_two_diff_prefix(b: &mut Bencher) {
    let aut = $aut(vec!["abcdef", "bmnopq"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_two_one_prefix_byte_every_match(b: &mut Bencher) {
    let aut = $aut(vec!["zbcdef", "zmnopq"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_two_one_prefix_byte_no_match(b: &mut Bencher) {
    let aut = $aut(vec!["zbcdef", "zmnopq"]);
    bench_aut_no_match(b, aut, &haystack_same('y'));
}

#[bench]
fn ac_two_one_prefix_byte_random(b: &mut Bencher) {
    let aut = $aut(vec!["zbcdef\x00", "zmnopq\x00"]);
    bench_aut_no_match(b, aut, HAYSTACK_RANDOM);
}

#[bench]
fn ac_ten_bytes(b: &mut Bencher) {
    let aut = $aut(vec!["a", "b", "c", "d", "e",
                        "f", "g", "h", "i", "j"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_ten_diff_prefix(b: &mut Bencher) {
    let aut = $aut(vec!["abcdef", "bbcdef", "cbcdef", "dbcdef",
                        "ebcdef", "fbcdef", "gbcdef", "hbcdef",
                        "ibcdef", "jbcdef"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_ten_one_prefix_byte_every_match(b: &mut Bencher) {
    let aut = $aut(vec!["zacdef", "zbcdef", "zccdef", "zdcdef",
                        "zecdef", "zfcdef", "zgcdef", "zhcdef",
                        "zicdef", "zjcdef"]);
    bench_aut_no_match(b, aut, &haystack_same('z'));
}

#[bench]
fn ac_ten_one_prefix_byte_no_match(b: &mut Bencher) {
    let aut = $aut(vec!["zacdef", "zbcdef", "zccdef", "zdcdef",
                        "zecdef", "zfcdef", "zgcdef", "zhcdef",
                        "zicdef", "zjcdef"]);
    bench_aut_no_match(b, aut, &haystack_same('y'));
}

#[bench]
fn ac_ten_one_prefix_byte_random(b: &mut Bencher) {
    let aut = $aut(vec!["zacdef\x00", "zbcdef\x00", "zccdef\x00",
                        "zdcdef\x00", "zecdef\x00", "zfcdef\x00",
                        "zgcdef\x00", "zhcdef\x00", "zicdef\x00",
                        "zjcdef\x00"]);
    bench_aut_no_match(b, aut, HAYSTACK_RANDOM);
}
}}}

aut_benches!(dense, AcAutomaton::new);
aut_benches!(sparse, AcAutomaton::<Sparse>::with_transitions);

// A naive multi-pattern search.
// We use this to benchmark *throughput*, so it should never match anything.
fn naive_find(needles: &[String], haystack: &str) -> bool {
    for hi in 0..haystack.len() {
        let rest = &haystack.as_bytes()[hi..];
        for needle in needles {
            let needle = needle.as_bytes();
            if needle.len() > rest.len() {
                continue;
            }
            if needle == &rest[..needle.len()] {
                // should never happen in throughput benchmarks.
                return true;
            }
        }
    }
    false
}

#[bench]
fn naive_one_byte(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["a"], &haystack_same('z'));
}

#[bench]
fn naive_one_prefix_byte_no_match(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["zbc"], &haystack_same('y'));
}

#[bench]
fn naive_one_prefix_byte_every_match(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["zbc"], &haystack_same('z'));
}

#[bench]
fn naive_one_prefix_byte_random(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["zbc\x00"], HAYSTACK_RANDOM);
}

#[bench]
fn naive_two_bytes(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["a", "b"], &haystack_same('z'));
}

#[bench]
fn naive_two_diff_prefix(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["abcdef", "bmnopq"], &haystack_same('z'));
}

#[bench]
fn naive_two_one_prefix_byte_every_match(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["zbcdef", "zmnopq"], &haystack_same('z'));
}

#[bench]
fn naive_two_one_prefix_byte_no_match(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["zbcdef", "zmnopq"], &haystack_same('y'));
}

#[bench]
fn naive_two_one_prefix_byte_random(b: &mut Bencher) {
    bench_naive_no_match(b, vec!["zbcdef\x00", "zmnopq\x00"], HAYSTACK_RANDOM);
}

#[bench]
fn naive_ten_bytes(b: &mut Bencher) {
    let needles = vec!["a", "b", "c", "d", "e",
                       "f", "g", "h", "i", "j"];
    bench_naive_no_match(b, needles, &haystack_same('z'));
}

#[bench]
fn naive_ten_diff_prefix(b: &mut Bencher) {
    let needles = vec!["abcdef", "bbcdef", "cbcdef", "dbcdef",
                       "ebcdef", "fbcdef", "gbcdef", "hbcdef",
                       "ibcdef", "jbcdef"];
    bench_naive_no_match(b, needles, &haystack_same('z'));
}

#[bench]
fn naive_ten_one_prefix_byte_every_match(b: &mut Bencher) {
    let needles = vec!["zacdef", "zbcdef", "zccdef", "zdcdef",
                       "zecdef", "zfcdef", "zgcdef", "zhcdef",
                       "zicdef", "zjcdef"];
    bench_naive_no_match(b, needles, &haystack_same('z'));
}

#[bench]
fn naive_ten_one_prefix_byte_no_match(b: &mut Bencher) {
    let needles = vec!["zacdef", "zbcdef", "zccdef", "zdcdef",
                       "zecdef", "zfcdef", "zgcdef", "zhcdef",
                       "zicdef", "zjcdef"];
    bench_naive_no_match(b, needles, &haystack_same('y'));
}

#[bench]
fn naive_ten_one_prefix_byte_random(b: &mut Bencher) {
    let needles = vec!["zacdef\x00", "zbcdef\x00", "zccdef\x00",
                       "zdcdef\x00", "zecdef\x00", "zfcdef\x00",
                       "zgcdef\x00", "zhcdef\x00", "zicdef\x00",
                       "zjcdef\x00"];
    bench_naive_no_match(b, needles, HAYSTACK_RANDOM);
}
<|endoftext|>