<commit_before>//! An example that timelines the specified folder.
//!
//! The timeline action is one of the filesystem tools that the GRR agent
//! exposes. Its description is pretty simple: given a folder, it returns a
//! stream of stat entries for all the files in it (including all subfolders).
//!
//! This example is a tiny wrapper around this action that allows to execute it
//! as a standalone binary without all the RRG setup required. A primary reason
//! for creating it is to compare the efficiency Rust implementation against the
//! existing Python one.

use std::fs::File;
use std::path::{Path, PathBuf};

use rrg::session::Sink;
use structopt::StructOpt;

/// A type for the timelining example command-line arguments.
#[derive(StructOpt)]
#[structopt(name = "timeline", about = "A binary for the timeline action.")]
struct Opts {
    /// A path to the root directory to timeline.
    #[structopt(long="root", name="FILE", default_value="/",
                help="A path to the root directory to timeline.")]
    root: PathBuf,

    /// A path to a file to dump the results into.
    #[structopt(long="output", about="FILE",
                help="A path to a file to dump the results into.")]
    output: PathBuf,
}

/// A session type of the timeline example.
///
/// Every action handler needs a session object. "Real" sessions handle comms
/// with the GRR server, error reporting and alike. However, this is not needed
/// here and the example session simply dumps everything into the specified file
/// encoding data in the gzchunked format.
struct Session {
    /// A file into which the data chunks of the output are dumped.
    output: File,
}

impl Session {

    /// Creates a session object that will write to the given `output` path.
    fn open<P: AsRef<Path>>(output: P) -> Session {
        let output = File::create(output)
            .expect("failed to create the output file");

        Session {
            output: output,
        }
    }
}

impl rrg::session::Session for Session {

    fn reply<R>(&mut self, response: R) -> rrg::session::Result<()>
    where
        R: rrg::action::Response + 'static,
    {
        // We don't care about the chunk ids.
        drop(response);
        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> rrg::session::Result<()>
    where
        R: rrg::action::Response + 'static,
    {
        use std::io::Write as _;
        use byteorder::{BigEndian, WriteBytesExt as _};

        assert_eq!(sink, Sink::TRANSFER_STORE);

        let response = (&response as &dyn std::any::Any)
            .downcast_ref::<rrg::action::timeline::Chunk>()
            .expect("unexpected response type");

        self.output.write_u64::<BigEndian>(response.data.len() as u64)
            .expect("failed to write chunk size tag");
        self.output.write_all(&response.data[..])
            .expect("failed to write chunk data");

        Ok(())
    }
}

fn main() {
    let opts = Opts::from_args();

    let mut session = Session::open(opts.output);

    let request = rrg::action::timeline::Request {
        root: opts.root,
    };

    rrg::action::timeline::handle(&mut session, request)
        .expect("failed to execute the action");
}
<commit_msg>Improve code style in the timeline example.<commit_after>//! An example that timelines the specified folder.
//!
//! The timeline action is one of the filesystem tools that the GRR agent
//! exposes. Its description is pretty simple: given a folder, it returns a
//! stream of stat entries for all the files in it (including all subfolders).
//!
//! This example is a tiny wrapper around this action that allows to execute it
//! as a standalone binary without all the RRG setup required. A primary reason
//! for creating it is to compare the efficiency Rust implementation against the
//! existing Python one.

use std::fs::File;
use std::path::{Path, PathBuf};

use rrg::action::timeline;
use rrg::session::Sink;
use structopt::StructOpt;

/// A type for the timelining example command-line arguments.
#[derive(StructOpt)]
#[structopt(name = "timeline", about = "A binary for the timeline action.")]
struct Opts {
    /// A path to the root directory to timeline.
    #[structopt(long="root", name="FILE", default_value="/",
                help="A path to the root directory to timeline.")]
    root: PathBuf,

    /// A path to a file to dump the results into.
    #[structopt(long="output", about="FILE",
                help="A path to a file to dump the results into.")]
    output: PathBuf,
}

/// A session type of the timeline example.
///
/// Every action handler needs a session object. "Real" sessions handle comms
/// with the GRR server, error reporting and alike. However, this is not needed
/// here and the example session simply dumps everything into the specified file
/// encoding data in the gzchunked format.
struct Session {
    /// A file into which the data chunks of the output are dumped.
    output: File,
}

impl Session {

    /// Creates a session object that will write to the given `output` path.
    fn open<P: AsRef<Path>>(output: P) -> Session {
        let output = File::create(output)
            .expect("failed to create the output file");

        Session {
            output: output,
        }
    }
}

impl rrg::session::Session for Session {

    fn reply<R>(&mut self, response: R) -> rrg::session::Result<()>
    where
        R: rrg::action::Response + 'static,
    {
        // We don't care about the chunk ids.
        drop(response);
        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> rrg::session::Result<()>
    where
        R: rrg::action::Response + 'static,
    {
        use std::io::Write as _;
        use byteorder::{BigEndian, WriteBytesExt as _};

        assert_eq!(sink, Sink::TRANSFER_STORE);

        let response = (&response as &dyn std::any::Any)
            .downcast_ref::<timeline::Chunk>()
            .expect("unexpected response type");

        self.output.write_u64::<BigEndian>(response.data.len() as u64)
            .expect("failed to write chunk size tag");
        self.output.write_all(&response.data[..])
            .expect("failed to write chunk data");

        Ok(())
    }
}

fn main() {
    let opts = Opts::from_args();

    timeline::handle(&mut Session::open(opts.output), timeline::Request {
        root: opts.root,
    }).expect("failed to execute the action");
}
<|endoftext|>


<commit_before>use std::ptr::null;
use libc::{c_char, c_int, c_double, c_void};
use std::ffi::CString;
use std::cell::RefCell;
use utils::{_string, _last_null_pointer_err};
use gdal_sys::{ogr, ogr_enums};
use spatial_ref::{SpatialRef, CoordTransform};

use errors::*;

/// OGR Geometry
pub struct Geometry {
    c_geometry_ref: RefCell<Option<*const c_void>>,
    owned: bool,
}

#[derive(Clone,PartialEq,Debug)]
pub enum WkbType {
    WkbUnknown = ogr::WKB_UNKNOWN as isize,
    WkbPoint = ogr::WKB_POINT as isize,
    WkbLinestring = ogr::WKB_LINESTRING as isize,
    WkbPolygon = ogr::WKB_POLYGON as isize,
    WkbMultipoint = ogr::WKB_MULTIPOINT as isize,
    WkbMultilinestring = ogr::WKB_MULTILINESTRING as isize,
    WkbMultipolygon = ogr::WKB_MULTIPOLYGON as isize,
    WkbGeometrycollection = ogr::WKB_GEOMETRYCOLLECTION as isize,
    WkbLinearring = ogr::WKB_LINEARRING as isize,
}

impl WkbType {
    pub fn from_ogr_type(ogr_type: c_int) -> WkbType {
        match ogr_type {
            ogr::WKB_UNKNOWN => WkbType::WkbUnknown,
            ogr::WKB_POINT => WkbType::WkbPoint,
            ogr::WKB_LINESTRING => WkbType::WkbLinestring,
            ogr::WKB_POLYGON => WkbType::WkbPolygon,
            ogr::WKB_MULTIPOINT => WkbType::WkbMultipoint,
            ogr::WKB_MULTILINESTRING => WkbType::WkbMultilinestring,
            ogr::WKB_MULTIPOLYGON => WkbType::WkbMultipolygon,
            ogr::WKB_GEOMETRYCOLLECTION => WkbType::WkbGeometrycollection,
            ogr::WKB_LINEARRING => WkbType::WkbLinearring,
            _ => WkbType::WkbUnknown
        }
    }
}

impl Geometry {
    pub unsafe fn lazy_feature_geometry() -> Geometry {
        // Geometry objects created with this method map to a Feature's
        // geometry whose memory is managed by the GDAL feature.
        // This object has a tricky lifecycle:
        //
        // * Initially it's created with a null c_geometry
        // * The first time `Feature::geometry` is called, it gets
        //   c_geometry from GDAL and calls `set_c_geometry` with it.
        // * When the Feature is destroyed, this object is also destroyed,
        //   which is good, because that's when c_geometry (which is managed
        //   by the GDAL feature) becomes invalid. Because `self.owned` is
        //   `true`, we don't call `OGR_G_DestroyGeometry`.
        Geometry{c_geometry_ref: RefCell::new(None), owned: false}
    }

    pub fn has_gdal_ptr(&self) -> bool {
        self.c_geometry_ref.borrow().is_some()
    }

    pub unsafe fn set_c_geometry(&self, c_geometry: *const c_void) {
        assert!(! self.has_gdal_ptr());
        assert_eq!(self.owned, false);
        *(self.c_geometry_ref.borrow_mut()) = Some(c_geometry);
    }

    unsafe fn with_c_geometry(c_geom: *const c_void, owned: bool) -> Geometry {
        Geometry{
            c_geometry_ref: RefCell::new(Some(c_geom)),
            owned: owned,
        }
    }

    pub fn empty(wkb_type: c_int) -> Result<Geometry> {
        let c_geom = unsafe { ogr::OGR_G_CreateGeometry(wkb_type) };
        if c_geom.is_null() {
            return Err(_last_null_pointer_err("OGR_G_CreateGeometry").into());
        };
        Ok(unsafe { Geometry::with_c_geometry(c_geom, true) })
    }

    pub fn is_empty(&self) -> bool {
        unsafe { ogr::OGR_G_IsEmpty(self.c_geometry()) }
    }

    /// Create a geometry by parsing a
    /// [WKT](https://en.wikipedia.org/wiki/Well-known_text) string.
    pub fn from_wkt(wkt: &str) -> Result<Geometry> {
        let c_wkt = CString::new(wkt)?;
        let mut c_wkt_ptr: *const c_char = c_wkt.as_ptr();
        let mut c_geom: *const c_void = null();
        let rv = unsafe { ogr::OGR_G_CreateFromWkt(&mut c_wkt_ptr, null(), &mut c_geom) };
        if rv != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(rv, "OGR_G_CreateFromWkt").into());
        }
        Ok(unsafe { Geometry::with_c_geometry(c_geom, true) })
    }

    /// Create a rectangular geometry from West, South, East and North values.
    pub fn bbox(w: f64, s: f64, e: f64, n: f64) -> Result<Geometry> {
        Geometry::from_wkt(&format!(
            "POLYGON (({} {}, {} {}, {} {}, {} {}, {} {}))",
            w, n,
            e, n,
            e, s,
            w, s,
            w, n,
        ))
    }

    /// Serialize the geometry as JSON.
    pub fn json(&self) -> Result<String> {
        let c_json = unsafe { ogr::OGR_G_ExportToJson(self.c_geometry()) };
        if c_json.is_null() {
            return Err(_last_null_pointer_err("OGR_G_ExportToJson").into());
        };
        let rv = _string(c_json);
        unsafe { ogr::VSIFree(c_json as *mut c_void) };
        Ok(rv)
    }

    /// Serialize the geometry as WKT.
    pub fn wkt(&self) -> Result<String> {
        let mut c_wkt: *const c_char = null();
        let _err = unsafe { ogr::OGR_G_ExportToWkt(self.c_geometry(), &mut c_wkt) };
        if _err != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(_err, "OGR_G_ExportToWkt").into());
        }
        let wkt = _string(c_wkt);
        unsafe { ogr::OGRFree(c_wkt as *mut c_void) };
        Ok(wkt)
    }

    pub unsafe fn c_geometry(&self) -> *const c_void {
        self.c_geometry_ref.borrow().unwrap()
    }

    pub unsafe fn into_c_geometry(mut self) -> *const c_void {
        assert!(self.owned);
        self.owned = false;
        return self.c_geometry();
    }

    pub fn set_point_2d(&mut self, i: usize, p: (f64, f64)) {
        let (x, y) = p;
        unsafe { ogr::OGR_G_SetPoint_2D(
            self.c_geometry(),
            i as c_int,
            x as c_double,
            y as c_double,
        ) };
    }

    pub fn get_point(&self, i: i32) -> (f64, f64, f64) {
        let mut x: c_double = 0.;
        let mut y: c_double = 0.;
        let mut z: c_double = 0.;
        unsafe { ogr::OGR_G_GetPoint(self.c_geometry(), i, &mut x, &mut y, &mut z) };
        return (x as f64, y as f64, z as f64);
    }

    pub fn get_point_vec(&self) -> Vec<(f64, f64, f64)> {
        let length = unsafe{ ogr::OGR_G_GetPointCount(self.c_geometry()) };
        (0..length).map(|i| self.get_point(i)).collect()
    }

    /// Compute the convex hull of this geometry.
    pub fn convex_hull(&self) -> Result<Geometry> {
        let c_geom = unsafe { ogr::OGR_G_ConvexHull(self.c_geometry()) };
        if c_geom.is_null() {
            return Err(_last_null_pointer_err("OGR_G_ConvexHull").into());
        };
        Ok(unsafe { Geometry::with_c_geometry(c_geom, true) })
    }

    pub fn geometry_type(&self) -> WkbType {
        let ogr_type = unsafe { ogr::OGR_G_GetGeometryType(self.c_geometry()) };
        WkbType::from_ogr_type(ogr_type)
    }

    pub fn geometry_count(&self) -> usize {
        let cnt = unsafe { ogr::OGR_G_GetGeometryCount(self.c_geometry()) };
        cnt as usize
    }

    pub unsafe fn _get_geometry(&self, n: usize) -> Geometry {
        // get the n-th sub-geometry as a non-owned Geometry; don't keep this
        // object for long.
        let c_geom = ogr::OGR_G_GetGeometryRef(self.c_geometry(), n as c_int);
        return Geometry::with_c_geometry(c_geom, false);
    }

    pub fn add_geometry(&mut self, mut sub: Geometry) -> Result<()> {
        assert!(sub.owned);
        sub.owned = false;
        let rv = unsafe { ogr::OGR_G_AddGeometryDirectly(
            self.c_geometry(),
            sub.c_geometry(),
        ) };
        if rv != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(rv, "OGR_G_AddGeometryDirectly").into());
        }
        Ok(())
    }

    // Transform the geometry inplace (when we own the Geometry)
    pub fn transform_inplace(&self, htransform: &CoordTransform) -> Result<()> {
        let rv = unsafe { ogr::OGR_G_Transform(
            self.c_geometry(),
            htransform.to_c_hct()
        ) };
        if rv != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(rv, "OGR_G_Transform").into());
        }
        Ok(())
    }

    // Return a new transformed geometry (when the Geometry is owned by a Feature)
    pub fn transform(&self, htransform: &CoordTransform) -> Result<Geometry> {
        let new_c_geom = unsafe { ogr::OGR_G_Clone(self.c_geometry()) };
        let rv = unsafe { ogr::OGR_G_Transform(new_c_geom, htransform.to_c_hct()) };
        if rv != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(rv, "OGR_G_Transform").into());
        }
        Ok(unsafe { Geometry::with_c_geometry(new_c_geom, true) } )
    }

    pub fn transform_to_inplace(&self, spatial_ref: &SpatialRef) -> Result<()> {
        let rv = unsafe { ogr::OGR_G_TransformTo(
            self.c_geometry(),
            spatial_ref.to_c_hsrs()
        ) };
        if rv != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(rv, "OGR_G_TransformTo").into());
        }
        Ok(())
    }

    pub fn transform_to(&self, spatial_ref: &SpatialRef) -> Result<Geometry> {
        let new_c_geom = unsafe { ogr::OGR_G_Clone(self.c_geometry()) };
        let rv = unsafe { ogr::OGR_G_TransformTo(new_c_geom, spatial_ref.to_c_hsrs()) };
        if rv != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(rv, "OGR_G_TransformTo").into());
        }
        Ok(unsafe { Geometry::with_c_geometry(new_c_geom, true) } )
    }

    pub fn area(&self) -> f64 {
        unsafe { ogr::OGR_G_Area(self.c_geometry()) }
    }

    /// May or may not contain a reference to a SpatialRef: if not, it returns
    /// an `Ok(None)`; if it does, it tries to build a SpatialRef. If that
    /// succeeds, it returns an Ok(Some(SpatialRef)), otherwise, you get the
    /// Err.
    ///
    pub fn spatial_reference(&self) -> Option<SpatialRef> {
        let c_spatial_ref = unsafe { ogr::OGR_G_GetSpatialReference(self.c_geometry()) };

        if c_spatial_ref.is_null() {
            None
        } else {
            match SpatialRef::from_c_obj(c_spatial_ref) {
                Ok(sr) => Some(sr),
                Err(_) => None
            }
        }
    }

    pub fn set_spatial_reference(&mut self, spatial_ref: SpatialRef) {
        unsafe { ogr::OGR_G_AssignSpatialReference(self.c_geometry(), spatial_ref.to_c_hsrs()) };
    }
}

impl Drop for Geometry {
    fn drop(&mut self) {
        if self.owned {
            let c_geometry = self.c_geometry_ref.borrow();
            unsafe { ogr::OGR_G_DestroyGeometry(c_geometry.unwrap() as *mut c_void) };
        }
    }
}

impl Clone for Geometry {
    fn clone(&self) -> Geometry {
        // assert!(self.has_gdal_ptr());
        let c_geometry = self.c_geometry_ref.borrow();
        let new_c_geom = unsafe { ogr::OGR_G_Clone(c_geometry.unwrap())};
        unsafe { Geometry::with_c_geometry(new_c_geom, true) }
    }
}

#[cfg(test)]
mod tests {
    use super::Geometry;
    use spatial_ref::SpatialRef;

    #[test]
    pub fn test_area() {
        let geom = Geometry::empty(::gdal_sys::ogr::WKB_MULTIPOLYGON).unwrap();
        assert_eq!(geom.area(), 0.0);

        let geom = Geometry::from_wkt("POINT(0 0)").unwrap();
        assert_eq!(geom.area(), 0.0);

        let wkt = "POLYGON ((45.0 45.0, 45.0 50.0, 50.0 50.0, 50.0 45.0, 45.0 45.0))";
        let geom = Geometry::from_wkt(wkt).unwrap();
        assert_eq!(geom.area().floor(), 25.0);
    }

    #[test]
    pub fn test_is_empty() {
        let geom = Geometry::empty(::gdal_sys::ogr::WKB_MULTIPOLYGON).unwrap();
        assert!(geom.is_empty());

        let geom = Geometry::from_wkt("POINT(0 0)").unwrap();
        assert!(!geom.is_empty());

        let wkt = "POLYGON ((45.0 45.0, 45.0 50.0, 50.0 50.0, 50.0 45.0, 45.0 45.0))";
        let geom = Geometry::from_wkt(wkt).unwrap();
        assert!(!geom.is_empty());
    }

    #[test]
    pub fn test_spatial_reference() {
        let geom = Geometry::empty(::gdal_sys::ogr::WKB_MULTIPOLYGON).unwrap();
        assert!(geom.spatial_reference().is_none());

        let geom = Geometry::from_wkt("POINT(0 0)").unwrap();
        assert!(geom.spatial_reference().is_none());

        let wkt = "POLYGON ((45.0 45.0, 45.0 50.0, 50.0 50.0, 50.0 45.0, 45.0 45.0))";
        let mut geom = Geometry::from_wkt(wkt).unwrap();
        assert!(geom.spatial_reference().is_none());

        let srs = SpatialRef::from_epsg(4326).unwrap();
        geom.set_spatial_reference(srs);
        assert!(geom.spatial_reference().is_some());
    }
}
<commit_msg>Fix Issue #45<commit_after>use std::ptr::null;
use libc::{c_char, c_int, c_double, c_void};
use std::ffi::CString;
use std::cell::RefCell;
use utils::{_string, _last_null_pointer_err};
use gdal_sys::{ogr, ogr_enums};
use spatial_ref::{SpatialRef, CoordTransform};

use errors::*;

/// OGR Geometry
pub struct Geometry {
    c_geometry_ref: RefCell<Option<*const c_void>>,
    owned: bool,
}

#[derive(Clone,PartialEq,Debug)]
pub enum WkbType {
    WkbUnknown = ogr::WKB_UNKNOWN as isize,
    WkbPoint = ogr::WKB_POINT as isize,
    WkbLinestring = ogr::WKB_LINESTRING as isize,
    WkbPolygon = ogr::WKB_POLYGON as isize,
    WkbMultipoint = ogr::WKB_MULTIPOINT as isize,
    WkbMultilinestring = ogr::WKB_MULTILINESTRING as isize,
    WkbMultipolygon = ogr::WKB_MULTIPOLYGON as isize,
    WkbGeometrycollection = ogr::WKB_GEOMETRYCOLLECTION as isize,
    WkbLinearring = ogr::WKB_LINEARRING as isize,
}

impl WkbType {
    pub fn from_ogr_type(ogr_type: c_int) -> WkbType {
        match ogr_type {
            ogr::WKB_UNKNOWN => WkbType::WkbUnknown,
            ogr::WKB_POINT => WkbType::WkbPoint,
            ogr::WKB_LINESTRING => WkbType::WkbLinestring,
            ogr::WKB_POLYGON => WkbType::WkbPolygon,
            ogr::WKB_MULTIPOINT => WkbType::WkbMultipoint,
            ogr::WKB_MULTILINESTRING => WkbType::WkbMultilinestring,
            ogr::WKB_MULTIPOLYGON => WkbType::WkbMultipolygon,
            ogr::WKB_GEOMETRYCOLLECTION => WkbType::WkbGeometrycollection,
            ogr::WKB_LINEARRING => WkbType::WkbLinearring,
            _ => WkbType::WkbUnknown
        }
    }
}

impl Geometry {
    pub unsafe fn lazy_feature_geometry() -> Geometry {
        // Geometry objects created with this method map to a Feature's
        // geometry whose memory is managed by the GDAL feature.
        // This object has a tricky lifecycle:
        //
        // * Initially it's created with a null c_geometry
        // * The first time `Feature::geometry` is called, it gets
        //   c_geometry from GDAL and calls `set_c_geometry` with it.
        // * When the Feature is destroyed, this object is also destroyed,
        //   which is good, because that's when c_geometry (which is managed
        //   by the GDAL feature) becomes invalid. Because `self.owned` is
        //   `true`, we don't call `OGR_G_DestroyGeometry`.
        Geometry{c_geometry_ref: RefCell::new(None), owned: false}
    }

    pub fn has_gdal_ptr(&self) -> bool {
        self.c_geometry_ref.borrow().is_some()
    }

    pub unsafe fn set_c_geometry(&self, c_geometry: *const c_void) {
        assert!(! self.has_gdal_ptr());
        assert_eq!(self.owned, false);
        *(self.c_geometry_ref.borrow_mut()) = Some(c_geometry);
    }

    unsafe fn with_c_geometry(c_geom: *const c_void, owned: bool) -> Geometry {
        Geometry{
            c_geometry_ref: RefCell::new(Some(c_geom)),
            owned: owned,
        }
    }

    pub fn empty(wkb_type: c_int) -> Result<Geometry> {
        let c_geom = unsafe { ogr::OGR_G_CreateGeometry(wkb_type) };
        if c_geom.is_null() {
            return Err(_last_null_pointer_err("OGR_G_CreateGeometry").into());
        };
        Ok(unsafe { Geometry::with_c_geometry(c_geom, true) })
    }

    pub fn is_empty(&self) -> bool {
        unsafe { ogr::OGR_G_IsEmpty(self.c_geometry()) }
    }

    /// Create a geometry by parsing a
    /// [WKT](https://en.wikipedia.org/wiki/Well-known_text) string.
    pub fn from_wkt(wkt: &str) -> Result<Geometry> {
        let c_wkt = CString::new(wkt)?;
        let mut c_wkt_ptr: *const c_char = c_wkt.as_ptr();
        let mut c_geom: *const c_void = null();
        let rv = unsafe { ogr::OGR_G_CreateFromWkt(&mut c_wkt_ptr, null(), &mut c_geom) };
        if rv != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(rv, "OGR_G_CreateFromWkt").into());
        }
        Ok(unsafe { Geometry::with_c_geometry(c_geom, true) })
    }

    /// Create a rectangular geometry from West, South, East and North values.
    pub fn bbox(w: f64, s: f64, e: f64, n: f64) -> Result<Geometry> {
        Geometry::from_wkt(&format!(
            "POLYGON (({} {}, {} {}, {} {}, {} {}, {} {}))",
            w, n,
            e, n,
            e, s,
            w, s,
            w, n,
        ))
    }

    /// Serialize the geometry as JSON.
    pub fn json(&self) -> Result<String> {
        let c_json = unsafe { ogr::OGR_G_ExportToJson(self.c_geometry()) };
        if c_json.is_null() {
            return Err(_last_null_pointer_err("OGR_G_ExportToJson").into());
        };
        let rv = _string(c_json);
        unsafe { ogr::VSIFree(c_json as *mut c_void) };
        Ok(rv)
    }

    /// Serialize the geometry as WKT.
    pub fn wkt(&self) -> Result<String> {
        let mut c_wkt: *const c_char = null();
        let _err = unsafe { ogr::OGR_G_ExportToWkt(self.c_geometry(), &mut c_wkt) };
        if _err != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(_err, "OGR_G_ExportToWkt").into());
        }
        let wkt = _string(c_wkt);
        unsafe { ogr::OGRFree(c_wkt as *mut c_void) };
        Ok(wkt)
    }

    pub unsafe fn c_geometry(&self) -> *const c_void {
        self.c_geometry_ref.borrow().unwrap()
    }

    pub unsafe fn into_c_geometry(mut self) -> *const c_void {
        assert!(self.owned);
        self.owned = false;
        return self.c_geometry();
    }

    pub fn set_point_2d(&mut self, i: usize, p: (f64, f64)) {
        let (x, y) = p;
        unsafe { ogr::OGR_G_SetPoint_2D(
            self.c_geometry(),
            i as c_int,
            x as c_double,
            y as c_double,
        ) };
    }

    pub fn get_point(&self, i: i32) -> (f64, f64, f64) {
        let mut x: c_double = 0.;
        let mut y: c_double = 0.;
        let mut z: c_double = 0.;
        unsafe { ogr::OGR_G_GetPoint(self.c_geometry(), i, &mut x, &mut y, &mut z) };
        return (x as f64, y as f64, z as f64);
    }

    pub fn get_point_vec(&self) -> Vec<(f64, f64, f64)> {
        let length = unsafe{ ogr::OGR_G_GetPointCount(self.c_geometry()) };
        (0..length).map(|i| self.get_point(i)).collect()
    }

    /// Compute the convex hull of this geometry.
    pub fn convex_hull(&self) -> Result<Geometry> {
        let c_geom = unsafe { ogr::OGR_G_ConvexHull(self.c_geometry()) };
        if c_geom.is_null() {
            return Err(_last_null_pointer_err("OGR_G_ConvexHull").into());
        };
        Ok(unsafe { Geometry::with_c_geometry(c_geom, true) })
    }

    pub fn geometry_type(&self) -> WkbType {
        let ogr_type = unsafe { ogr::OGR_G_GetGeometryType(self.c_geometry()) };
        WkbType::from_ogr_type(ogr_type)
    }

    pub fn geometry_count(&self) -> usize {
        let cnt = unsafe { ogr::OGR_G_GetGeometryCount(self.c_geometry()) };
        cnt as usize
    }

    pub unsafe fn _get_geometry(&self, n: usize) -> Geometry {
        // get the n-th sub-geometry as a non-owned Geometry; don't keep this
        // object for long.
        let c_geom = ogr::OGR_G_GetGeometryRef(self.c_geometry(), n as c_int);
        return Geometry::with_c_geometry(c_geom, false);
    }

    pub fn add_geometry(&mut self, mut sub: Geometry) -> Result<()> {
        assert!(sub.owned);
        sub.owned = false;
        let rv = unsafe { ogr::OGR_G_AddGeometryDirectly(
            self.c_geometry(),
            sub.c_geometry(),
        ) };
        if rv != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(rv, "OGR_G_AddGeometryDirectly").into());
        }
        Ok(())
    }

    // Transform the geometry inplace (when we own the Geometry)
    pub fn transform_inplace(&self, htransform: &CoordTransform) -> Result<()> {
        let rv = unsafe { ogr::OGR_G_Transform(
            self.c_geometry(),
            htransform.to_c_hct()
        ) };
        if rv != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(rv, "OGR_G_Transform").into());
        }
        Ok(())
    }

    // Return a new transformed geometry (when the Geometry is owned by a Feature)
    pub fn transform(&self, htransform: &CoordTransform) -> Result<Geometry> {
        let new_c_geom = unsafe { ogr::OGR_G_Clone(self.c_geometry()) };
        let rv = unsafe { ogr::OGR_G_Transform(new_c_geom, htransform.to_c_hct()) };
        if rv != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(rv, "OGR_G_Transform").into());
        }
        Ok(unsafe { Geometry::with_c_geometry(new_c_geom, true) } )
    }

    pub fn transform_to_inplace(&self, spatial_ref: &SpatialRef) -> Result<()> {
        let rv = unsafe { ogr::OGR_G_TransformTo(
            self.c_geometry(),
            spatial_ref.to_c_hsrs()
        ) };
        if rv != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(rv, "OGR_G_TransformTo").into());
        }
        Ok(())
    }

    pub fn transform_to(&self, spatial_ref: &SpatialRef) -> Result<Geometry> {
        let new_c_geom = unsafe { ogr::OGR_G_Clone(self.c_geometry()) };
        let rv = unsafe { ogr::OGR_G_TransformTo(new_c_geom, spatial_ref.to_c_hsrs()) };
        if rv != ogr_enums::OGRErr::OGRERR_NONE {
            return Err(ErrorKind::OgrError(rv, "OGR_G_TransformTo").into());
        }
        Ok(unsafe { Geometry::with_c_geometry(new_c_geom, true) } )
    }

    pub fn area(&self) -> f64 {
        unsafe { ogr::OGR_G_Area(self.c_geometry()) }
    }

    /// May or may not contain a reference to a SpatialRef: if not, it returns
    /// an `Ok(None)`; if it does, it tries to build a SpatialRef. If that
    /// succeeds, it returns an Ok(Some(SpatialRef)), otherwise, you get the
    /// Err.
    ///
    pub fn spatial_reference(&self) -> Option<SpatialRef> {
        let c_spatial_ref = unsafe { ogr::OGR_G_GetSpatialReference(self.c_geometry()) };

        if c_spatial_ref.is_null() {
            None
        } else {
            match SpatialRef::from_c_obj(c_spatial_ref) {
                Ok(sr) => Some(sr),
                Err(_) => None
            }
        }
    }

    pub fn set_spatial_reference(&mut self, spatial_ref: &SpatialRef) {
        unsafe { ogr::OGR_G_AssignSpatialReference(self.c_geometry(), (*spatial_ref).to_c_hsrs()) };
    }
}

impl Drop for Geometry {
    fn drop(&mut self) {
        if self.owned {
            let c_geometry = self.c_geometry_ref.borrow();
            unsafe { ogr::OGR_G_DestroyGeometry(c_geometry.unwrap() as *mut c_void) };
        }
    }
}

impl Clone for Geometry {
    fn clone(&self) -> Geometry {
        // assert!(self.has_gdal_ptr());
        let c_geometry = self.c_geometry_ref.borrow();
        let new_c_geom = unsafe { ogr::OGR_G_Clone(c_geometry.unwrap())};
        unsafe { Geometry::with_c_geometry(new_c_geom, true) }
    }
}

#[cfg(test)]
mod tests {
    use super::Geometry;
    use spatial_ref::SpatialRef;

    #[test]
    pub fn test_area() {
        let geom = Geometry::empty(::gdal_sys::ogr::WKB_MULTIPOLYGON).unwrap();
        assert_eq!(geom.area(), 0.0);

        let geom = Geometry::from_wkt("POINT(0 0)").unwrap();
        assert_eq!(geom.area(), 0.0);

        let wkt = "POLYGON ((45.0 45.0, 45.0 50.0, 50.0 50.0, 50.0 45.0, 45.0 45.0))";
        let geom = Geometry::from_wkt(wkt).unwrap();
        assert_eq!(geom.area().floor(), 25.0);
    }

    #[test]
    pub fn test_is_empty() {
        let geom = Geometry::empty(::gdal_sys::ogr::WKB_MULTIPOLYGON).unwrap();
        assert!(geom.is_empty());

        let geom = Geometry::from_wkt("POINT(0 0)").unwrap();
        assert!(!geom.is_empty());

        let wkt = "POLYGON ((45.0 45.0, 45.0 50.0, 50.0 50.0, 50.0 45.0, 45.0 45.0))";
        let geom = Geometry::from_wkt(wkt).unwrap();
        assert!(!geom.is_empty());
    }

    #[test]
    pub fn test_spatial_reference() {
        let geom = Geometry::empty(::gdal_sys::ogr::WKB_MULTIPOLYGON).unwrap();
        assert!(geom.spatial_reference().is_none());

        let geom = Geometry::from_wkt("POINT(0 0)").unwrap();
        assert!(geom.spatial_reference().is_none());

        let wkt = "POLYGON ((45.0 45.0, 45.0 50.0, 50.0 50.0, 50.0 45.0, 45.0 45.0))";
        let mut geom = Geometry::from_wkt(wkt).unwrap();
        assert!(geom.spatial_reference().is_none());

        let srs = SpatialRef::from_epsg(4326).unwrap();
        geom.set_spatial_reference(&srs);
        assert!(geom.spatial_reference().is_some());
    }
}
<|endoftext|>


<commit_before>use rational::Rational64;
use data::params::CodecParams;

#[derive(Clone,Debug,PartialEq)]
pub struct Stream {
    pub id: usize,
    pub index: usize,
    pub params : CodecParams,
    pub start: Option<u64>,
    pub duration: Option<u64>,
    pub timebase : Rational64,
//  seek_index : SeekIndex
}

pub struct StreamGroup<'a> {
    pub id: usize,
    pub start: u64,
    pub end: u64,
    pub streams: &'a [Stream],
}
<commit_msg>format: Add a simplified accessor<commit_after>use rational::Rational64;
use data::params::CodecParams;

#[derive(Clone,Debug,PartialEq)]
pub struct Stream {
    pub id: usize,
    pub index: usize,
    pub params : CodecParams,
    pub start: Option<u64>,
    pub duration: Option<u64>,
    pub timebase : Rational64,
//  seek_index : SeekIndex
}

impl Stream {
    pub fn get_extradata<'a>(&'a self) -> Option<&'a [u8]> {
        self.params.extradata.as_ref().map(|e| e.as_slice())
    }
}

pub struct StreamGroup<'a> {
    pub id: usize,
    pub start: u64,
    pub end: u64,
    pub streams: &'a [Stream],
}
<|endoftext|>


<commit_before>//! Runtime lifecycle related functionality.

use util::bigint::{U256, M256};
use util::gas::Gas;
use vm::errors::{RequireError, MachineError};
use vm::commit::AccountState;
use vm::Memory;
use super::{Machine, MachineStatus};
use super::util::copy_into_memory_apply;
use super::cost::code_deposit_gas;

/// # Lifecycle of a Machine
///
/// When a new non-invoked transaction is created, `initialize_call`
/// or `initialize_create` should be called. After this, the machine
/// can be stepped as normal. When the machine meets a CALL/CALLCODE
/// or CREATE instruction, a sub-machine will be created. This
/// submachine should first call `invoke_call` or
/// `invoke_create`. After the submachine is finished, it should call
/// `apply_sub`. When the non-invoked transaction is finished, it
/// should first call `code_deposit` if it is a contract creation
/// transaction. After that, it should call `finalize`.

impl<M: Memory + Default> Machine<M> {
    /// Initialize a MessageCall transaction.
    pub fn initialize_call(&mut self, preclaimed_value: U256) {
        self.state.account_state.premark_exists(self.state.context.address);
        self.state.account_state.decrease_balance(self.state.context.caller, preclaimed_value);
        self.state.account_state.decrease_balance(self.state.context.caller, self.state.context.value);
        self.state.account_state.increase_balance(self.state.context.address, self.state.context.value);
    }

    /// Initialize the runtime as a call from a CALL or CALLCODE opcode.
    pub fn invoke_call(&mut self) {
        self.state.account_state.premark_exists(self.state.context.address);
        self.state.account_state.decrease_balance(self.state.context.caller, self.state.context.value);
        self.state.account_state.increase_balance(self.state.context.address, self.state.context.value);
    }

    /// Initialize a ContractCreation transaction.
    pub fn initialize_create(&mut self, preclaimed_value: U256) -> Result<(), RequireError> {
        self.state.account_state.require(self.state.context.address)?;

        self.state.account_state.premark_exists(self.state.context.address);
        self.state.account_state.decrease_balance(self.state.context.caller, preclaimed_value);
        self.state.account_state.decrease_balance(self.state.context.caller, self.state.context.value);
        self.state.account_state.create(self.state.context.address, self.state.context.value).unwrap();

        Ok(())
    }

    /// Initialize the runtime as a call from a CREATE opcode.
    pub fn invoke_create(&mut self) -> Result<(), RequireError> {
        self.state.account_state.require(self.state.context.address)?;

        self.state.account_state.premark_exists(self.state.context.address);
        self.state.account_state.decrease_balance(self.state.context.caller, self.state.context.value);
        self.state.account_state.create(self.state.context.address, self.state.context.value).unwrap();

        Ok(())
    }

    /// Deposit code for a ContractCreation transaction or a CREATE opcode.
    pub fn code_deposit(&mut self) {
        match self.status() {
            MachineStatus::ExitedOk | MachineStatus::ExitedErr(_) => (),
            _ => panic!(),
        }

        let deposit_cost = code_deposit_gas(self.state.out.len());
        if deposit_cost > self.state.available_gas() {
            if !self.state.patch.force_code_deposit {
                self.status = MachineStatus::ExitedErr(MachineError::EmptyGas);
            } else {
                self.state.account_state.code_deposit(self.state.context.address, &[]);
            }
        } else {
            self.state.used_gas = self.state.used_gas + deposit_cost;
            self.state.account_state.code_deposit(self.state.context.address,
                                                  self.state.out.as_slice());
        }
    }

    /// Finalize a transaction. This should not be used when invoked
    /// by an opcode.
    pub fn finalize(&mut self, real_used_gas: Gas, preclaimed_value: U256, fresh_account_state: &AccountState) -> Result<(), RequireError> {
        self.state.account_state.require(self.state.context.address)?;

        match self.status() {
            MachineStatus::ExitedOk => {
                // Requires removed accounts to exist.
                for address in &self.state.removed {
                    self.state.account_state.require(*address)?;
                }
            },
            MachineStatus::ExitedErr(_) => {
                // If exited with error, reset all changes.
                self.state.account_state = fresh_account_state.clone();
                self.state.account_state.decrease_balance(self.state.context.caller, preclaimed_value);
                self.state.logs = Vec::new();
                self.state.removed = Vec::new();
            },
            _ => panic!(),
        }

        let gas_dec = real_used_gas * self.state.context.gas_price;
        self.state.account_state.increase_balance(self.state.context.caller, preclaimed_value);
        self.state.account_state.decrease_balance(self.state.context.caller, gas_dec.into());

        for address in &self.state.removed {
            self.state.account_state.remove(*address).unwrap();
        }

        match self.status() {
            MachineStatus::ExitedOk => Ok(()),
            MachineStatus::ExitedErr(_) => Ok(()),
            _ => panic!(),
        }
    }

    /// Apply a sub runtime into the current runtime. This sub runtime
    /// should have been created by the current runtime's `derive`
    /// function. Depending whether the current runtime is invoking a
    /// ContractCreation or MessageCall instruction, it will apply
    /// various states back.
    pub fn apply_sub(&mut self, sub: Machine<M>) {
        use std::mem::swap;
        let mut status = MachineStatus::Running;
        swap(&mut status, &mut self.status);
        match status {
            MachineStatus::InvokeCreate(_) => {
                self.apply_create(sub);
            },
            MachineStatus::InvokeCall(_, (out_start, out_len)) => {
                self.apply_call(sub, out_start, out_len);
            },
            _ => panic!(),
        }
    }

    fn apply_create(&mut self, mut sub: Machine<M>) {
        if self.state.available_gas() < sub.state.used_gas {
            panic!();
        }

        sub.code_deposit();

        match sub.status() {
            MachineStatus::ExitedOk => {
                let sub_total_used_gas = sub.state.total_used_gas();

                self.state.account_state = sub.state.account_state;
                self.state.blockhash_state = sub.state.blockhash_state;
                self.state.logs = sub.state.logs;
                self.state.removed = sub.state.removed;
                self.state.used_gas = self.state.used_gas + sub_total_used_gas;
                self.state.refunded_gas = self.state.refunded_gas + sub.state.refunded_gas;

            },
            MachineStatus::ExitedErr(_) => {
                self.state.used_gas = self.state.used_gas + sub.state.context.gas_limit;
                self.state.stack.pop().unwrap();
                self.state.stack.push(M256::zero()).unwrap();
            },
            _ => panic!(),
        }
    }

    fn apply_call(&mut self, sub: Machine<M>, out_start: M256, out_len: M256) {
        if self.state.available_gas() < sub.state.used_gas {
            panic!();
        }

        match sub.status() {
            MachineStatus::ExitedOk => {
                let sub_total_used_gas = sub.state.total_used_gas();

                self.state.account_state = sub.state.account_state;
                self.state.blockhash_state = sub.state.blockhash_state;
                self.state.logs = sub.state.logs;
                self.state.removed = sub.state.removed;
                self.state.used_gas = self.state.used_gas + sub_total_used_gas;
                self.state.refunded_gas = self.state.refunded_gas + sub.state.refunded_gas;
                copy_into_memory_apply(&mut self.state.memory, sub.state.out.as_slice(),
                                       out_start, out_len);
            },
            MachineStatus::ExitedErr(_) => {
                self.state.used_gas = self.state.used_gas + sub.state.context.gas_limit;
                self.state.stack.pop().unwrap();
                self.state.stack.push(M256::zero()).unwrap();
            },
            _ => panic!(),
        }
    }
}
<commit_msg>Apply miner rewards on transaction finalization<commit_after>//! Runtime lifecycle related functionality.

use util::bigint::{U256, M256};
use util::gas::Gas;
use vm::errors::{RequireError, MachineError};
use vm::commit::AccountState;
use vm::Memory;
use super::{Machine, MachineStatus};
use super::util::copy_into_memory_apply;
use super::cost::code_deposit_gas;

/// # Lifecycle of a Machine
///
/// When a new non-invoked transaction is created, `initialize_call`
/// or `initialize_create` should be called. After this, the machine
/// can be stepped as normal. When the machine meets a CALL/CALLCODE
/// or CREATE instruction, a sub-machine will be created. This
/// submachine should first call `invoke_call` or
/// `invoke_create`. After the submachine is finished, it should call
/// `apply_sub`. When the non-invoked transaction is finished, it
/// should first call `code_deposit` if it is a contract creation
/// transaction. After that, it should call `finalize`.

impl<M: Memory + Default> Machine<M> {
    /// Initialize a MessageCall transaction.
    pub fn initialize_call(&mut self, preclaimed_value: U256) {
        self.state.account_state.premark_exists(self.state.context.address);
        self.state.account_state.decrease_balance(self.state.context.caller, preclaimed_value);
        self.state.account_state.decrease_balance(self.state.context.caller, self.state.context.value);
        self.state.account_state.increase_balance(self.state.context.address, self.state.context.value);
    }

    /// Initialize the runtime as a call from a CALL or CALLCODE opcode.
    pub fn invoke_call(&mut self) {
        self.state.account_state.premark_exists(self.state.context.address);
        self.state.account_state.decrease_balance(self.state.context.caller, self.state.context.value);
        self.state.account_state.increase_balance(self.state.context.address, self.state.context.value);
    }

    /// Initialize a ContractCreation transaction.
    pub fn initialize_create(&mut self, preclaimed_value: U256) -> Result<(), RequireError> {
        self.state.account_state.require(self.state.context.address)?;

        self.state.account_state.premark_exists(self.state.context.address);
        self.state.account_state.decrease_balance(self.state.context.caller, preclaimed_value);
        self.state.account_state.decrease_balance(self.state.context.caller, self.state.context.value);
        self.state.account_state.create(self.state.context.address, self.state.context.value).unwrap();

        Ok(())
    }

    /// Initialize the runtime as a call from a CREATE opcode.
    pub fn invoke_create(&mut self) -> Result<(), RequireError> {
        self.state.account_state.require(self.state.context.address)?;

        self.state.account_state.premark_exists(self.state.context.address);
        self.state.account_state.decrease_balance(self.state.context.caller, self.state.context.value);
        self.state.account_state.create(self.state.context.address, self.state.context.value).unwrap();

        Ok(())
    }

    /// Deposit code for a ContractCreation transaction or a CREATE opcode.
    pub fn code_deposit(&mut self) {
        match self.status() {
            MachineStatus::ExitedOk | MachineStatus::ExitedErr(_) => (),
            _ => panic!(),
        }

        let deposit_cost = code_deposit_gas(self.state.out.len());
        if deposit_cost > self.state.available_gas() {
            if !self.state.patch.force_code_deposit {
                self.status = MachineStatus::ExitedErr(MachineError::EmptyGas);
            } else {
                self.state.account_state.code_deposit(self.state.context.address, &[]);
            }
        } else {
            self.state.used_gas = self.state.used_gas + deposit_cost;
            self.state.account_state.code_deposit(self.state.context.address,
                                                  self.state.out.as_slice());
        }
    }

    /// Finalize a transaction. This should not be used when invoked
    /// by an opcode.
    pub fn finalize(&mut self, real_used_gas: Gas, preclaimed_value: U256, fresh_account_state: &AccountState) -> Result<(), RequireError> {
        self.state.account_state.require(self.state.context.address)?;

        match self.status() {
            MachineStatus::ExitedOk => {
                // Requires removed accounts to exist.
                for address in &self.state.removed {
                    self.state.account_state.require(*address)?;
                }
            },
            MachineStatus::ExitedErr(_) => {
                // If exited with error, reset all changes.
                self.state.account_state = fresh_account_state.clone();
                self.state.account_state.decrease_balance(self.state.context.caller, preclaimed_value);
                self.state.logs = Vec::new();
                self.state.removed = Vec::new();
            },
            _ => panic!(),
        }

        let gas_dec = real_used_gas * self.state.context.gas_price;
        self.state.account_state.increase_balance(self.state.context.caller, preclaimed_value);
        self.state.account_state.decrease_balance(self.state.context.caller, gas_dec.into());

        // Apply miner rewards
        self.state.account_state.increase_balance(self.state.block.coinbase, gas_dec.into());

        for address in &self.state.removed {
            self.state.account_state.remove(*address).unwrap();
        }

        match self.status() {
            MachineStatus::ExitedOk => Ok(()),
            MachineStatus::ExitedErr(_) => Ok(()),
            _ => panic!(),
        }
    }

    /// Apply a sub runtime into the current runtime. This sub runtime
    /// should have been created by the current runtime's `derive`
    /// function. Depending whether the current runtime is invoking a
    /// ContractCreation or MessageCall instruction, it will apply
    /// various states back.
    pub fn apply_sub(&mut self, sub: Machine<M>) {
        use std::mem::swap;
        let mut status = MachineStatus::Running;
        swap(&mut status, &mut self.status);
        match status {
            MachineStatus::InvokeCreate(_) => {
                self.apply_create(sub);
            },
            MachineStatus::InvokeCall(_, (out_start, out_len)) => {
                self.apply_call(sub, out_start, out_len);
            },
            _ => panic!(),
        }
    }

    fn apply_create(&mut self, mut sub: Machine<M>) {
        if self.state.available_gas() < sub.state.used_gas {
            panic!();
        }

        sub.code_deposit();

        match sub.status() {
            MachineStatus::ExitedOk => {
                let sub_total_used_gas = sub.state.total_used_gas();

                self.state.account_state = sub.state.account_state;
                self.state.blockhash_state = sub.state.blockhash_state;
                self.state.logs = sub.state.logs;
                self.state.removed = sub.state.removed;
                self.state.used_gas = self.state.used_gas + sub_total_used_gas;
                self.state.refunded_gas = self.state.refunded_gas + sub.state.refunded_gas;

            },
            MachineStatus::ExitedErr(_) => {
                self.state.used_gas = self.state.used_gas + sub.state.context.gas_limit;
                self.state.stack.pop().unwrap();
                self.state.stack.push(M256::zero()).unwrap();
            },
            _ => panic!(),
        }
    }

    fn apply_call(&mut self, sub: Machine<M>, out_start: M256, out_len: M256) {
        if self.state.available_gas() < sub.state.used_gas {
            panic!();
        }

        match sub.status() {
            MachineStatus::ExitedOk => {
                let sub_total_used_gas = sub.state.total_used_gas();

                self.state.account_state = sub.state.account_state;
                self.state.blockhash_state = sub.state.blockhash_state;
                self.state.logs = sub.state.logs;
                self.state.removed = sub.state.removed;
                self.state.used_gas = self.state.used_gas + sub_total_used_gas;
                self.state.refunded_gas = self.state.refunded_gas + sub.state.refunded_gas;
                copy_into_memory_apply(&mut self.state.memory, sub.state.out.as_slice(),
                                       out_start, out_len);
            },
            MachineStatus::ExitedErr(_) => {
                self.state.used_gas = self.state.used_gas + sub.state.context.gas_limit;
                self.state.stack.pop().unwrap();
                self.state.stack.push(M256::zero()).unwrap();
            },
            _ => panic!(),
        }
    }
}
<|endoftext|>


<commit_before>use std::vec::Vec;
use super::defines::*;

//TODO: convert to macro with usage
//format!(indent!(5, "format:{}"), 6)
pub fn tabs(num: usize) -> String {
    format!("{:1$}", "", TAB_SIZE * num)
}

pub fn format_block(prefix: &str, suffix: &str, body: &[String]) -> Vec<String> {
    let mut v = Vec::new();
    if !prefix.is_empty() { v.push(prefix.into()); }
    for s in body.iter() {
        let s = format!("{}{}", TAB, s);
        v.push(s);
    }
    if !suffix.is_empty() { v.push(suffix.into()); }
    v
}

pub fn format_block_smart(prefix: &str, suffix: &str, body: &[String],
                          outer_separator: &str, inner_separator: &str) -> Vec<String> {
    format_block_smart_width(prefix, suffix, body, outer_separator, inner_separator, MAX_TEXT_WIDTH)
}

pub fn format_block_smart_width(prefix: &str, suffix: &str, body: &[String], outer_separator: &str,
                                inner_separator: &str, max_width: usize) -> Vec<String> {
    let outer_len = prefix.len() + suffix.len() + 2 * outer_separator.len();
    let mut inner_len = inner_separator.len() * (body.len() - 1);
    //TODO: change to sum()
    for s in body {
        inner_len = inner_len + s.len();
    }
    if (outer_len + inner_len) > max_width {
        format_block(prefix, suffix, body)
    } else {
        let mut s = format!("{}{}", prefix, outer_separator);
        let mut first = true;
        for s_ in body {
            if first {
                first = false;
                s = s + s_;
            } else {
                s = s + inner_separator + s_;
            }
        }
        s = s + outer_separator + suffix;
        vec![s]
    }
}

pub fn comment_block(body: &[String]) -> Vec<String> {
    body.iter().map(|s| format!("//{}", s)).collect()
}

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

    #[test]
    fn test_tabs() {
        assert_eq!(tabs(0), "");
        assert_eq!(tabs(1), format!("{}", TAB));
        assert_eq!(tabs(2), format!("{0}{0}", TAB));
    }

    #[test]
    fn test_format_block() {
        let body = vec!["0 => 1,".into(), "1 => 0,".into()];
        let actual = format_block("match a {", "}", &body);
        let expected = [
            "match a {",
            "    0 => 1,",
            "    1 => 0,",
            "}"
        ];
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_format_block_smart_width_one_line_outer_separator() {
        let body = vec!["f()".into()];
        let actual = format_block_smart_width("unsafe {", "}", &body, " ", "", 14);
        let expected = ["unsafe { f() }"];
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_format_block_smart_width_many_lines_outer_separator() {
        let body = vec!["f()".into()];
        let actual = format_block_smart_width("unsafe {", "}", &body, " ", "", 13);
        let expected = ["unsafe {", "    f()", "}"];
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_format_block_smart_one_line_inner_separator() {
        let body = vec!["a: &str".into(), "b: &str".into()];
        let actual = format_block_smart("f(", ")", &body, "", ", ");
        let expected = ["f(a: &str, b: &str)"];
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_comment_block() {
        let body = vec!["f(a,".into(), "  b)".into()];
        let actual = comment_block(&body);
        let expected = ["//f(a,", "//  b)"];
        assert_eq!(actual, expected);
    }
}
<commit_msg>Add format_block_one_line<commit_after>use std::vec::Vec;
use super::defines::*;

//TODO: convert to macro with usage
//format!(indent!(5, "format:{}"), 6)
pub fn tabs(num: usize) -> String {
    format!("{:1$}", "", TAB_SIZE * num)
}

pub fn format_block(prefix: &str, suffix: &str, body: &[String]) -> Vec<String> {
    let mut v = Vec::new();
    if !prefix.is_empty() { v.push(prefix.into()); }
    for s in body.iter() {
        let s = format!("{}{}", TAB, s);
        v.push(s);
    }
    if !suffix.is_empty() { v.push(suffix.into()); }
    v
}

pub fn format_block_one_line(prefix: &str, suffix: &str, body: &[String],
                             outer_separator: &str, inner_separator: &str) -> String {
    let mut s = format!("{}{}", prefix, outer_separator);
    let mut first = true;
    for s_ in body {
        if first {
            first = false;
            s = s + s_;
        } else {
            s = s + inner_separator + s_;
        }
    }
    s + outer_separator + suffix
}

pub fn format_block_smart(prefix: &str, suffix: &str, body: &[String],
                          outer_separator: &str, inner_separator: &str) -> Vec<String> {
    format_block_smart_width(prefix, suffix, body, outer_separator, inner_separator, MAX_TEXT_WIDTH)
}

pub fn format_block_smart_width(prefix: &str, suffix: &str, body: &[String], outer_separator: &str,
                                inner_separator: &str, max_width: usize) -> Vec<String> {
    let outer_len = prefix.len() + suffix.len() + 2 * outer_separator.len();
    let mut inner_len = inner_separator.len() * (body.len() - 1);
    //TODO: change to sum()
    for s in body {
        inner_len = inner_len + s.len();
    }
    if (outer_len + inner_len) > max_width {
        format_block(prefix, suffix, body)
    } else {
        let s = format_block_one_line(prefix, suffix, body, outer_separator, inner_separator);
        vec![s]
    }
}

pub fn comment_block(body: &[String]) -> Vec<String> {
    body.iter().map(|s| format!("//{}", s)).collect()
}

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

    #[test]
    fn test_tabs() {
        assert_eq!(tabs(0), "");
        assert_eq!(tabs(1), format!("{}", TAB));
        assert_eq!(tabs(2), format!("{0}{0}", TAB));
    }

    #[test]
    fn test_format_block() {
        let body = vec!["0 => 1,".into(), "1 => 0,".into()];
        let actual = format_block("match a {", "}", &body);
        let expected = [
            "match a {",
            "    0 => 1,",
            "    1 => 0,",
            "}"
        ];
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_format_block_smart_width_one_line_outer_separator() {
        let body = vec!["f()".into()];
        let actual = format_block_smart_width("unsafe {", "}", &body, " ", "", 14);
        let expected = ["unsafe { f() }"];
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_format_block_smart_width_many_lines_outer_separator() {
        let body = vec!["f()".into()];
        let actual = format_block_smart_width("unsafe {", "}", &body, " ", "", 13);
        let expected = ["unsafe {", "    f()", "}"];
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_format_block_smart_one_line_inner_separator() {
        let body = vec!["a: &str".into(), "b: &str".into()];
        let actual = format_block_smart("f(", ")", &body, "", ", ");
        let expected = ["f(a: &str, b: &str)"];
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_comment_block() {
        let body = vec!["f(a,".into(), "  b)".into()];
        let actual = comment_block(&body);
        let expected = ["//f(a,", "//  b)"];
        assert_eq!(actual, expected);
    }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::{io, io::prelude::*};
use std::hash::Hash;
use std::fmt::Debug;
use std::collections::HashMap;

type Recipe<T> = HashMap<T, u64>;
type RecipeBook<T> = HashMap<T, (Recipe<T>, u64)>;

fn get_input() -> io::Result<String> {
  let mut file = File::open("14.txt")?;
  let mut contents = String::new();
  file.read_to_string(&mut contents)?;

  Ok(contents)
}

fn parse_input(input: &str) -> RecipeBook<&str> {
  fn parse_ingredient(input: &str) -> Option<(&str, u64)> {
    let mut tokens = input.trim().split(" ");

    match (tokens.next(), tokens.next()) {
      (Some(amount), Some(id)) => {
        amount.parse::<u64>().ok()
        .map(|amount| (id, amount))
      },
      _ => None
    }
  }

  input.lines()
  .map(|line| line.split(" => "))
  .filter_map(|mut tokens| match (tokens.next(), tokens.next()) {
    (Some(ingredients), Some(outcome)) => {
      Some((ingredients.split(", "), outcome))
    },
    _ => None
  })
  .map(|(ingredients, outcome)| (ingredients.filter_map(parse_ingredient), parse_ingredient(outcome)))
  .filter_map(|(ingredients, outcome)| outcome.and_then(|outcome| {
    let ingredients = ingredients.collect::<Recipe<_>>();

    if ingredients.len() == 0 {
      return None;
    }

    Some((outcome.0, (ingredients, outcome.1)))
  }))
  .collect()
}

fn get_base_ingredient_amount<T>(recipes: &RecipeBook<T>, ingredient: T, amount: u64, residues: &mut Recipe<T>) -> u64
where T: Eq + Clone + Copy + Debug + Hash {
  let mut amount = amount;

  if let Some(residue) = residues.get_mut(&ingredient) {
    if *residue >= amount {
      *residue -= amount;
      amount = 0;
    } else {
      amount -= *residue;
      *residue = 0;
    }
  }

  if amount == 0 {
    return 0;
  }

  match recipes.get(&ingredient) {
    None => amount,
    Some((subrecipe, yield_amount)) => {
      let times = (amount as f64 / *yield_amount as f64).ceil() as u64;
      let result = subrecipe.iter()
        .map(|(&ingredient, &amount)| {
          get_base_ingredient_amount(recipes, ingredient, amount * times, residues)
        })
        .sum();

      let new_residue = residues.get(&ingredient).cloned().unwrap_or(0) + *yield_amount * times - amount;
      residues.insert(ingredient, new_residue);

      result
    }
  }
}

fn main() {
  let input = get_input().unwrap();
  let recipes = parse_input(&input);

  let one_fuel_ore = get_base_ingredient_amount(&recipes, "FUEL", 1, &mut Recipe::new());
  println!("Part 1: {}", one_fuel_ore);

  let max_ore = 1000000000000u64;
  let mut max_fuel_approx = max_ore / one_fuel_ore;
  let mut cooldown = false;

  loop {
    let ore = get_base_ingredient_amount(&recipes, "FUEL", max_fuel_approx, &mut Recipe::new());

    if ore < max_ore {
      if cooldown {
        break;
      }

      max_fuel_approx += (max_ore - ore) / one_fuel_ore + 1;
    } else {
      cooldown = true;
      max_fuel_approx -= 1;
    }
  }

  println!("Part 2: {}", max_fuel_approx);
}
<commit_msg>Simplify trait constraints<commit_after>use std::fs::File;
use std::{io, io::prelude::*};
use std::hash::Hash;
use std::collections::HashMap;

type Recipe<T> = HashMap<T, u64>;
type RecipeBook<T> = HashMap<T, (Recipe<T>, u64)>;

fn get_input() -> io::Result<String> {
  let mut file = File::open("14.txt")?;
  let mut contents = String::new();
  file.read_to_string(&mut contents)?;

  Ok(contents)
}

fn parse_input(input: &str) -> RecipeBook<&str> {
  fn parse_ingredient(input: &str) -> Option<(&str, u64)> {
    let mut tokens = input.trim().split(" ");

    match (tokens.next(), tokens.next()) {
      (Some(amount), Some(id)) => {
        amount.parse::<u64>().ok()
        .map(|amount| (id, amount))
      },
      _ => None
    }
  }

  input.lines()
  .map(|line| line.split(" => "))
  .filter_map(|mut tokens| match (tokens.next(), tokens.next()) {
    (Some(ingredients), Some(outcome)) => {
      Some((ingredients.split(", "), outcome))
    },
    _ => None
  })
  .map(|(ingredients, outcome)| (ingredients.filter_map(parse_ingredient), parse_ingredient(outcome)))
  .filter_map(|(ingredients, outcome)| outcome.and_then(|(product, amount)| {
    let ingredients = ingredients.collect::<Recipe<_>>();

    if ingredients.len() == 0 {
      return None;
    }

    Some((product, (ingredients, amount)))
  }))
  .collect()
}

fn get_base_ingredient_amount<'a, T>(
  recipes: &'a RecipeBook<T>,
  ingredient: &'a T,
  amount: u64,
  residues: &mut HashMap<&'a T, u64>
) -> u64
where
  T: Eq + Hash
{
  let mut amount = amount;

  if let Some(residue) = residues.get_mut(ingredient) {
    if *residue >= amount {
      *residue -= amount;
      amount = 0;
    } else {
      amount -= *residue;
      *residue = 0;
    }
  }

  if amount == 0 {
    return 0;
  }

  match recipes.get(&ingredient) {
    None => amount,
    Some((subrecipe, yield_amount)) => {
      let times = (amount as f64 / *yield_amount as f64).ceil() as u64;
      let result = subrecipe.iter()
        .map(|(ingredient, &amount)| {
          get_base_ingredient_amount(recipes, ingredient, amount * times, residues)
        })
        .sum();

      let new_residue = residues.get(ingredient).cloned().unwrap_or(0) + *yield_amount * times - amount;
      residues.insert(ingredient, new_residue);

      result
    }
  }
}

fn main() {
  let input = get_input().unwrap();
  let recipes = parse_input(&input);
  let get_ore_for_fuel = |n| get_base_ingredient_amount(&recipes, &"FUEL", n, &mut HashMap::new());

  let one_fuel_ore = get_ore_for_fuel(1);
  println!("Part 1: {}", one_fuel_ore);

  let max_ore = 1000000000000u64;
  let mut max_fuel_approx = max_ore / one_fuel_ore;
  let mut cooldown = false;

  loop {
    let ore = get_ore_for_fuel(max_fuel_approx);

    if ore < max_ore {
      if cooldown {
        break;
      }

      max_fuel_approx += (max_ore - ore) / one_fuel_ore + 1;
    } else {
      cooldown = true;
      max_fuel_approx -= 1;
    }
  }

  println!("Part 2: {}", max_fuel_approx);
}
<|endoftext|>


<commit_before>fn mod_mul_(a: u64, b: u64, m: u64) -> u64 {
    (u128::from(a) * u128::from(b) % u128::from(m)) as u64
}

fn mod_mul(a: u64, b: u64, m: u64) -> u64 {
    match a.checked_mul(b) {
        Some(r) => if r >= m { r % m } else { r },
        None => mod_mul_(a, b, m),
    }
}

fn mod_sqr(a: u64, m: u64) -> u64 {
    if a < (1 << 32) {
        let r = a * a;
        if r >= m {
            r % m
        } else {
            r
        }
    } else {
        mod_mul_(a, a, m)
    }
}

fn mod_exp(mut x: u64, mut d: u64, n: u64) -> u64 {
    let mut ret: u64 = 1;
    while d != 0 {
        if d % 2 == 1 {
            ret = mod_mul(ret, x, n)
        }
        d /= 2;
        x = mod_sqr(x, n);
    }
    ret
}

/// Test if `n` is prime, using the deterministic version of the
/// Miller-Rabin test.
///
/// Doing a lot of primality tests with numbers strictly below some
/// upper bound will be faster using the `is_prime` method of a
/// `Sieve` instance.
///
/// # Examples
///
/// ```rust
/// assert_eq!(primal::is_prime(1), false);
/// assert_eq!(primal::is_prime(2), true);
/// assert_eq!(primal::is_prime(3), true);
/// assert_eq!(primal::is_prime(4), false);
/// assert_eq!(primal::is_prime(5), true);
///
/// assert_eq!(primal::is_prime(22_801_763_487), false);
/// assert_eq!(primal::is_prime(22_801_763_489), true);
/// assert_eq!(primal::is_prime(22_801_763_491), false);
/// ```
pub fn miller_rabin(n: u64) -> bool {
    const HINT: &[u64] = &[2];

    // we have a strict upper bound, so we can just use the witness
    // table of Pomerance, Selfridge & Wagstaff and Jeaschke to be as
    // efficient as possible, without having to fall back to
    // randomness.
    const WITNESSES: &[(u64, &[u64])] =
        &[(2_046, HINT),
          (1_373_652, &[2, 3]),
          (9_080_190, &[31, 73]),
          (25_326_000, &[2, 3, 5]),
          (4_759_123_140, &[2, 7, 61]),
          (1_112_004_669_632, &[2, 13, 23, 1662803]),
          (2_152_302_898_746, &[2, 3, 5, 7, 11]),
          (3_474_749_660_382, &[2, 3, 5, 7, 11, 13]),
          (341_550_071_728_320, &[2, 3, 5, 7, 11, 13, 17]),
          (0xFFFF_FFFF_FFFF_FFFF, &[2, 3, 5, 7, 11, 13, 17, 19, 23])
         ];

    if n % 2 == 0 { return n == 2 }
    if n == 1 { return false }

    let mut d = n - 1;
    let mut s = 0;
    while d % 2 == 0 { d /= 2; s += 1 }

    let witnesses =
        WITNESSES.iter().find(|&&(hi, _)| hi >= n)
            .map(|&(_, wtnss)| wtnss).unwrap();
    'next_witness: for &a in witnesses.iter() {
        let mut power = mod_exp(a, d, n);
        assert!(power < n);
        if power == 1 || power == n - 1 { continue 'next_witness }

        for _r in 0..s {
            power = mod_sqr(power, n);
            assert!(power < n);
            if power == 1 { return false }
            if power == n - 1 {
                continue 'next_witness
            }
        }
        return false
    }

    true
}

#[cfg(test)]
mod tests {
    use primal::Sieve;

    #[test]
    fn mod_mul() {
        assert_eq!(super::mod_mul(1 << 63, 1 << 32, 3), 2);
        assert_eq!(super::mod_mul(1 << 31, 1 << 31, (1 << 32) - 7), 3221225479);
        assert_eq!(super::mod_mul(1 << 32, 1 << 32, (1 << 32) - 7), 49);
        assert_eq!(super::mod_mul(1 << 32, 1 << 32, (1 << 32) + 7), 49);
        assert_eq!(super::mod_mul(1 << 63, 1 << 32, (1 << 32) + 7), 2_147_483_480);
        assert_eq!(super::mod_mul(1 << 63, 1 << 32, (1 << 63) + 7), 9_223_372_006_790_004_743);
        assert_eq!(super::mod_mul(1 << 32, 1 << 32, !0), 1);
    }

    #[test]
    fn miller_rabin() {
        const LIMIT: usize = 1_000_000;
        let sieve = Sieve::new(LIMIT);
        for x in 0..LIMIT {
            let s = sieve.is_prime(x);
            let mr = super::miller_rabin(x as u64);

            assert!(s == mr, "miller_rabin {} mismatches sieve {} for {}",
                    mr, s, x)
        }
    }
    #[test]
    fn miller_rabin_large() {
        let tests = &[
            (4_294_967_311, true),
            (4_294_967_291, true),
            (4_294_967_291 * 4_294_967_291, false),
            (!0, false),
            ];
        for &(n, is_prime) in tests {
            assert!(super::miller_rabin(n) == is_prime,
                    "mismatch for {} (should be {})", n, is_prime);
        }
    }
}
<commit_msg>primal-check: add an additional witness check<commit_after>fn mod_mul_(a: u64, b: u64, m: u64) -> u64 {
    (u128::from(a) * u128::from(b) % u128::from(m)) as u64
}

fn mod_mul(a: u64, b: u64, m: u64) -> u64 {
    match a.checked_mul(b) {
        Some(r) => if r >= m { r % m } else { r },
        None => mod_mul_(a, b, m),
    }
}

fn mod_sqr(a: u64, m: u64) -> u64 {
    if a < (1 << 32) {
        let r = a * a;
        if r >= m {
            r % m
        } else {
            r
        }
    } else {
        mod_mul_(a, a, m)
    }
}

fn mod_exp(mut x: u64, mut d: u64, n: u64) -> u64 {
    let mut ret: u64 = 1;
    while d != 0 {
        if d % 2 == 1 {
            ret = mod_mul(ret, x, n)
        }
        d /= 2;
        x = mod_sqr(x, n);
    }
    ret
}

/// Test if `n` is prime, using the deterministic version of the
/// Miller-Rabin test.
///
/// Doing a lot of primality tests with numbers strictly below some
/// upper bound will be faster using the `is_prime` method of a
/// `Sieve` instance.
///
/// # Examples
///
/// ```rust
/// assert_eq!(primal::is_prime(1), false);
/// assert_eq!(primal::is_prime(2), true);
/// assert_eq!(primal::is_prime(3), true);
/// assert_eq!(primal::is_prime(4), false);
/// assert_eq!(primal::is_prime(5), true);
///
/// assert_eq!(primal::is_prime(22_801_763_487), false);
/// assert_eq!(primal::is_prime(22_801_763_489), true);
/// assert_eq!(primal::is_prime(22_801_763_491), false);
/// ```
pub fn miller_rabin(n: u64) -> bool {
    const HINT: &[u64] = &[2];

    // we have a strict upper bound, so we can just use the witness
    // table of Pomerance, Selfridge & Wagstaff and Jeaschke to be as
    // efficient as possible, without having to fall back to
    // randomness. Additional limits from Feitsma and Galway complete
    // the entire range of `u64`. See also:
    // https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test#Testing_against_small_sets_of_bases
    const WITNESSES: &[(u64, &[u64])] = &[
        (2_046, HINT),
        (1_373_652, &[2, 3]),
        (9_080_190, &[31, 73]),
        (25_326_000, &[2, 3, 5]),
        (4_759_123_140, &[2, 7, 61]),
        (1_112_004_669_632, &[2, 13, 23, 1662803]),
        (2_152_302_898_746, &[2, 3, 5, 7, 11]),
        (3_474_749_660_382, &[2, 3, 5, 7, 11, 13]),
        (341_550_071_728_320, &[2, 3, 5, 7, 11, 13, 17]),
        (3_825_123_056_546_413_050, &[2, 3, 5, 7, 11, 13, 17, 19, 23]),
        (std::u64::MAX, &[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37]),
    ];

    if n % 2 == 0 { return n == 2 }
    if n == 1 { return false }

    let mut d = n - 1;
    let mut s = 0;
    while d % 2 == 0 { d /= 2; s += 1 }

    let witnesses =
        WITNESSES.iter().find(|&&(hi, _)| hi >= n)
            .map(|&(_, wtnss)| wtnss).unwrap();
    'next_witness: for &a in witnesses.iter() {
        let mut power = mod_exp(a, d, n);
        assert!(power < n);
        if power == 1 || power == n - 1 { continue 'next_witness }

        for _r in 0..s {
            power = mod_sqr(power, n);
            assert!(power < n);
            if power == 1 { return false }
            if power == n - 1 {
                continue 'next_witness
            }
        }
        return false
    }

    true
}

#[cfg(test)]
mod tests {
    use primal::Sieve;

    #[test]
    fn mod_mul() {
        assert_eq!(super::mod_mul(1 << 63, 1 << 32, 3), 2);
        assert_eq!(super::mod_mul(1 << 31, 1 << 31, (1 << 32) - 7), 3221225479);
        assert_eq!(super::mod_mul(1 << 32, 1 << 32, (1 << 32) - 7), 49);
        assert_eq!(super::mod_mul(1 << 32, 1 << 32, (1 << 32) + 7), 49);
        assert_eq!(super::mod_mul(1 << 63, 1 << 32, (1 << 32) + 7), 2_147_483_480);
        assert_eq!(super::mod_mul(1 << 63, 1 << 32, (1 << 63) + 7), 9_223_372_006_790_004_743);
        assert_eq!(super::mod_mul(1 << 32, 1 << 32, !0), 1);
    }

    #[test]
    fn miller_rabin() {
        const LIMIT: usize = 1_000_000;
        let sieve = Sieve::new(LIMIT);
        for x in 0..LIMIT {
            let s = sieve.is_prime(x);
            let mr = super::miller_rabin(x as u64);

            assert!(s == mr, "miller_rabin {} mismatches sieve {} for {}",
                    mr, s, x)
        }
    }
    #[test]
    fn miller_rabin_large() {
        let tests = &[
            (4_294_967_311, true),
            (4_294_967_291, true),
            (4_294_967_291 * 4_294_967_291, false),
            (!0, false),
            ];
        for &(n, is_prime) in tests {
            assert!(super::miller_rabin(n) == is_prime,
                    "mismatch for {} (should be {})", n, is_prime);
        }
    }

    #[test]
    fn oeis_a014233() {
        // https://oeis.org/A014233
        const A014233: [u64; 9] = [
            2047,
            1373653,
            25326001,
            3215031751,
            2152302898747,
            3474749660383,
            341550071728321,
            341550071728321,
            3825123056546413051,
            // 3825123056546413051,
            // 3825123056546413051,
            // 318665857834031151167461,
            // 3317044064679887385961981,
        ];
        for &n in &A014233 {
            assert!(!super::miller_rabin(n), "{} is composite!", n);
        }
    }
}
<|endoftext|>


<commit_before>/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Maciej Dziardziel <fiedzia@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE file
 * that was distributed with this source code.
 */

#![feature(process_id)]

extern crate libc;
extern crate time;

use libc::{c_int, pid_t};
use std::fmt;
use std::io;
use std::process::Child;
use std::sync::{Arc, Condvar, Mutex};
use std::thread;
use time::{Duration, get_time};

// This is basically sys::unix::process::ExitStatus
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum ExitStatus {
    Code(i32),
    Signal(i32),
}

impl ExitStatus {
    fn from_status(status: c_int) -> ExitStatus {
        if status & 0x7F != 0 { // WIFSIGNALED(status)
            ExitStatus::Signal(status & 0x7F)
        } else {
            ExitStatus::Code(status & 0xFF00 >> 8)
        }
    }

    pub fn success(&self) -> bool {
        match *self {
            ExitStatus::Code(code) => code == 0,
            _ => false,
        }
    }

    pub fn code(&self) -> Option<i32> {
        match *self {
            ExitStatus::Code(code) => Some(code),
            _ => None,
        }
    }

    pub fn signal(&self) -> Option<i32> {
        match *self {
            ExitStatus::Signal(code) => Some(code),
            _ => None,
        }
    }
}

impl fmt::Display for ExitStatus {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ExitStatus::Code(code) => write!(f, "exit code: {}", code),
            ExitStatus::Signal(code) => write!(f, "exit code: {}", code),
        }
    }
}

/// Missing methods for Child objects
pub trait ChildExt {
    /// Send a signal to a Child process.
    fn send_signal(&mut self, signal: usize) -> io::Result<()>;

    /// Wait for a process to finish or return after the specified duration.
    fn wait_or_timeout(&mut self, timeout: f64) -> io::Result<Option<ExitStatus>>;
}

impl ChildExt for Child {
    fn send_signal(&mut self, signal: usize) -> io::Result<()> {
        if unsafe { libc::funcs::posix88::signal::kill(self.id() as pid_t,
                                                       signal as i32) } != 0 {
            Err(io::Error::last_os_error())
        } else {
            Ok(())
        }
    }

    fn wait_or_timeout(&mut self, timeout: f64) -> io::Result<Option<ExitStatus>> {
        // The result will be written to that Option, protected by a Mutex
        // Then the Condvar will be signaled
        let state = Arc::new((
            Mutex::new(Option::None::<io::Result<ExitStatus>>),
            Condvar::new(),
        ));

        // Start the waiting thread
        let state_th = state.clone();
        let pid_th = self.id();
        thread::spawn(move || {
            let &(ref lock_th, ref cvar_th) = &*state_th;
            // Child::wait() would need a &mut to self, can't use that...
            // use waitpid() directly, with our own ExitStatus
            let mut status: c_int = 0;
            let r = unsafe { libc::waitpid(pid_th as i32, &mut status, 0) };
            // Fill the Option and notify on the Condvar
            let mut exitstatus_th = lock_th.lock().unwrap();
            if r != pid_th as c_int {
                *exitstatus_th = Some(Err(io::Error::last_os_error()));
            } else {
                let s = ExitStatus::from_status(status);
                *exitstatus_th = Some(Ok(s));
            }
            cvar_th.notify_one();
        });

        // Main thread waits
        let &(ref lock, ref cvar) = &*state;
        let mut exitstatus = lock.lock().unwrap();
        // Condvar::wait_timeout_ms() can wake too soon, in this case wait again
        let target = get_time() +
            Duration::seconds(timeout as i64) +
            Duration::nanoseconds((timeout * 1.0e-6) as i64);
        while exitstatus.is_none() {
            let now = get_time();
            if now >= target {
                return Ok(None)
            }
            let ms = (target - get_time()).num_milliseconds() as u32;
            exitstatus = cvar.wait_timeout_ms(exitstatus, ms).unwrap().0;
        }

        // Turn Option<Result<ExitStatus>> into Result<Option<ExitStatus>>
        match exitstatus.take() {
            Some(r) => match r {
                Ok(s) => Ok(Some(s)),
                Err(e) => Err(e),
            },
            None => panic!(),
        }
    }
}
<commit_msg>Remove unused attribute.<commit_after>/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Maciej Dziardziel <fiedzia@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE file
 * that was distributed with this source code.
 */

extern crate libc;
extern crate time;

use libc::{c_int, pid_t};
use std::fmt;
use std::io;
use std::process::Child;
use std::sync::{Arc, Condvar, Mutex};
use std::thread;
use time::{Duration, get_time};

// This is basically sys::unix::process::ExitStatus
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum ExitStatus {
    Code(i32),
    Signal(i32),
}

impl ExitStatus {
    fn from_status(status: c_int) -> ExitStatus {
        if status & 0x7F != 0 { // WIFSIGNALED(status)
            ExitStatus::Signal(status & 0x7F)
        } else {
            ExitStatus::Code(status & 0xFF00 >> 8)
        }
    }

    pub fn success(&self) -> bool {
        match *self {
            ExitStatus::Code(code) => code == 0,
            _ => false,
        }
    }

    pub fn code(&self) -> Option<i32> {
        match *self {
            ExitStatus::Code(code) => Some(code),
            _ => None,
        }
    }

    pub fn signal(&self) -> Option<i32> {
        match *self {
            ExitStatus::Signal(code) => Some(code),
            _ => None,
        }
    }
}

impl fmt::Display for ExitStatus {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ExitStatus::Code(code) => write!(f, "exit code: {}", code),
            ExitStatus::Signal(code) => write!(f, "exit code: {}", code),
        }
    }
}

/// Missing methods for Child objects
pub trait ChildExt {
    /// Send a signal to a Child process.
    fn send_signal(&mut self, signal: usize) -> io::Result<()>;

    /// Wait for a process to finish or return after the specified duration.
    fn wait_or_timeout(&mut self, timeout: f64) -> io::Result<Option<ExitStatus>>;
}

impl ChildExt for Child {
    fn send_signal(&mut self, signal: usize) -> io::Result<()> {
        if unsafe { libc::funcs::posix88::signal::kill(self.id() as pid_t,
                                                       signal as i32) } != 0 {
            Err(io::Error::last_os_error())
        } else {
            Ok(())
        }
    }

    fn wait_or_timeout(&mut self, timeout: f64) -> io::Result<Option<ExitStatus>> {
        // The result will be written to that Option, protected by a Mutex
        // Then the Condvar will be signaled
        let state = Arc::new((
            Mutex::new(Option::None::<io::Result<ExitStatus>>),
            Condvar::new(),
        ));

        // Start the waiting thread
        let state_th = state.clone();
        let pid_th = self.id();
        thread::spawn(move || {
            let &(ref lock_th, ref cvar_th) = &*state_th;
            // Child::wait() would need a &mut to self, can't use that...
            // use waitpid() directly, with our own ExitStatus
            let mut status: c_int = 0;
            let r = unsafe { libc::waitpid(pid_th as i32, &mut status, 0) };
            // Fill the Option and notify on the Condvar
            let mut exitstatus_th = lock_th.lock().unwrap();
            if r != pid_th as c_int {
                *exitstatus_th = Some(Err(io::Error::last_os_error()));
            } else {
                let s = ExitStatus::from_status(status);
                *exitstatus_th = Some(Ok(s));
            }
            cvar_th.notify_one();
        });

        // Main thread waits
        let &(ref lock, ref cvar) = &*state;
        let mut exitstatus = lock.lock().unwrap();
        // Condvar::wait_timeout_ms() can wake too soon, in this case wait again
        let target = get_time() +
            Duration::seconds(timeout as i64) +
            Duration::nanoseconds((timeout * 1.0e-6) as i64);
        while exitstatus.is_none() {
            let now = get_time();
            if now >= target {
                return Ok(None)
            }
            let ms = (target - get_time()).num_milliseconds() as u32;
            exitstatus = cvar.wait_timeout_ms(exitstatus, ms).unwrap().0;
        }

        // Turn Option<Result<ExitStatus>> into Result<Option<ExitStatus>>
        match exitstatus.take() {
            Some(r) => match r {
                Ok(s) => Ok(Some(s)),
                Err(e) => Err(e),
            },
            None => panic!(),
        }
    }
}
<|endoftext|>