<commit_before>//! Generator of Rust-Qt crates.
//!
//! See [README](https://github.com/rust-qt/cpp_to_rust/tree/master/qt_generator/qt_generator)
//! for more information.

use crate::config::GlobalConfig;
use crate::processor;
use crate::workspace::Workspace;
use flexi_logger::{Duplicate, LevelFilter, LogSpecification, Logger};
use itertools::Itertools;
use log::{error, info};
use ritual_common::errors::{bail, err_msg, Result};
use ritual_common::file_utils::canonicalize;
use ritual_common::file_utils::path_to_str;
use std::path::PathBuf;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
/// Generates rust_qt crates using ritual.
/// See [ritual](https://github.com/rust-qt/ritual) for more details.
pub struct Options {
    #[structopt(parse(from_os_str))]
    /// Directory for output and temporary files
    pub workspace: PathBuf,
    #[structopt(long = "local-paths")]
    /// Write local paths to `ritual` crates in generated `Cargo.toml`
    pub local_paths: Option<bool>,
    #[structopt(short = "c", long = "crates", required = true)]
    /// Crates to process (e.g. `qt_core`)
    pub crates: Vec<String>,
    #[structopt(short = "o", long = "operations", required = true)]
    /// Operations to perform
    pub operations: Vec<String>,
}

pub fn run_from_args(config: GlobalConfig) -> Result<()> {
    run(Options::from_args(), config)
}

pub fn run(options: Options, mut config: GlobalConfig) -> Result<()> {
    let workspace_path = canonicalize(options.workspace)?;

    let mut workspace = Workspace::new(workspace_path.clone())?;

    Logger::with(LogSpecification::default(LevelFilter::Trace).build())
        .log_to_file()
        .directory(path_to_str(&workspace.log_path())?)
        .print_message()
        .duplicate_to_stderr(Duplicate::Info)
        .start()
        .unwrap_or_else(|e| panic!("Logger initialization failed: {}", e));

    info!("Workspace: {}", workspace_path.display());

    if let Some(local_paths) = options.local_paths {
        workspace.set_write_dependencies_local_paths(local_paths)?;
    }

    let mut was_any_action = false;

    let final_crates = if options.crates.iter().any(|x| *x == "all") {
        let all = config.all_crate_names();
        if all.is_empty() {
            bail!("\"all\" is not supported as crate name specifier");
        }
        all.to_vec()
    } else {
        options.crates.clone()
    };

    let operations = options
        .operations
        .iter()
        .map(|s| s.to_lowercase())
        .collect_vec();

    if operations.is_empty() {
        error!("No action requested. Run \"qt_generator --help\".");
        return Ok(());
    }

    for crate_name in &final_crates {
        let create_config = config
            .create_config_hook()
            .ok_or_else(|| err_msg("create_config_hook is missing"))?;

        let config = create_config(&crate_name)?;
        was_any_action = true;
        processor::process(&mut workspace, &config, &operations)?;
    }

    //workspace.save_data()?;
    if was_any_action {
        info!("ritual finished");
    } else {
        error!("No action requested. Run \"qt_generator --help\".");
    }
    Ok(())
}
<commit_msg>create workspace if it doesn't exist<commit_after>//! Generator of Rust-Qt crates.
//!
//! See [README](https://github.com/rust-qt/cpp_to_rust/tree/master/qt_generator/qt_generator)
//! for more information.

use crate::config::GlobalConfig;
use crate::processor;
use crate::workspace::Workspace;
use flexi_logger::{Duplicate, LevelFilter, LogSpecification, Logger};
use itertools::Itertools;
use log::{error, info};
use ritual_common::errors::{bail, err_msg, Result};
use ritual_common::file_utils::path_to_str;
use ritual_common::file_utils::{canonicalize, create_dir};
use std::path::PathBuf;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
/// Generates rust_qt crates using ritual.
/// See [ritual](https://github.com/rust-qt/ritual) for more details.
pub struct Options {
    #[structopt(parse(from_os_str))]
    /// Directory for output and temporary files
    pub workspace: PathBuf,
    #[structopt(long = "local-paths")]
    /// Write local paths to `ritual` crates in generated `Cargo.toml`
    pub local_paths: Option<bool>,
    #[structopt(short = "c", long = "crates", required = true)]
    /// Crates to process (e.g. `qt_core`)
    pub crates: Vec<String>,
    #[structopt(short = "o", long = "operations", required = true)]
    /// Operations to perform
    pub operations: Vec<String>,
}

pub fn run_from_args(config: GlobalConfig) -> Result<()> {
    run(Options::from_args(), config)
}

pub fn run(options: Options, mut config: GlobalConfig) -> Result<()> {
    if !options.workspace.exists() {
        create_dir(&options.workspace)?;
    }
    let workspace_path = canonicalize(options.workspace)?;

    let mut workspace = Workspace::new(workspace_path.clone())?;

    Logger::with(LogSpecification::default(LevelFilter::Trace).build())
        .log_to_file()
        .directory(path_to_str(&workspace.log_path())?)
        .print_message()
        .duplicate_to_stderr(Duplicate::Info)
        .start()
        .unwrap_or_else(|e| panic!("Logger initialization failed: {}", e));

    info!("Workspace: {}", workspace_path.display());

    if let Some(local_paths) = options.local_paths {
        workspace.set_write_dependencies_local_paths(local_paths)?;
    }

    let mut was_any_action = false;

    let final_crates = if options.crates.iter().any(|x| *x == "all") {
        let all = config.all_crate_names();
        if all.is_empty() {
            bail!("\"all\" is not supported as crate name specifier");
        }
        all.to_vec()
    } else {
        options.crates.clone()
    };

    let operations = options
        .operations
        .iter()
        .map(|s| s.to_lowercase())
        .collect_vec();

    if operations.is_empty() {
        error!("No action requested. Run \"qt_generator --help\".");
        return Ok(());
    }

    for crate_name in &final_crates {
        let create_config = config
            .create_config_hook()
            .ok_or_else(|| err_msg("create_config_hook is missing"))?;

        let config = create_config(&crate_name)?;
        was_any_action = true;
        processor::process(&mut workspace, &config, &operations)?;
    }

    //workspace.save_data()?;
    if was_any_action {
        info!("ritual finished");
    } else {
        error!("No action requested. Run \"qt_generator --help\".");
    }
    Ok(())
}
<|endoftext|>


<commit_before>use std::slice;
use std::libc::{c_int, c_char};
use std::str::raw;
use std::os::getenv;
use std::path::Path;
use gdal::driver::Driver;
use gdal::register_drivers;
use gdal::geom::Point;

#[link(name="gdal")]
extern {
    fn GDALCreateCopy(
            hDriver: *(),
            pszFilename: *c_char,
            hSrcDS: *(),
            bStrict: c_int,
            papszOptions: **c_char,
            pfnProgres: *(),
            pProgressData: *()
        ) -> *();
    fn GDALOpen(pszFilename: *c_char, eAccess: c_int) -> *();
    fn GDALClose(hDS: *());
    fn GDALGetDatasetDriver(hDataset: *()) -> *();
    fn GDALGetRasterXSize(hDataset: *()) -> c_int;
    fn GDALGetRasterYSize(hDataset: *()) -> c_int;
    fn GDALGetRasterCount(hDataset: *()) -> c_int;
    fn GDALGetProjectionRef(hDS: *()) -> *c_char;
    fn GDALSetProjection(hDS: *(), pszProjection: *c_char) -> c_int;
    fn GDALGetRasterBand(hDS: *(), nBandId: c_int) -> *();
    fn GDALRasterIO(
            hBand: *(),
            eRWFlag: c_int,
            nXOff: c_int,
            nYOff: c_int,
            nXSize: c_int,
            nYSize: c_int,
            pData: *(),
            nBufXSize: c_int,
            nBufYSize: c_int,
            GDALDataType: c_int,
            nPixelSpace: c_int,
            nLineSpace: c_int
        ) -> c_int;
}

static GA_ReadOnly:  c_int = 0;
static GA_Update:    c_int = 1;

static GDT_Unknown:  c_int = 0;
static GDT_Byte:     c_int = 1;
static GDT_UInt16:   c_int = 2;
static GDT_Int16:    c_int = 3;
static GDT_UInt32:   c_int = 4;
static GDT_Int32:    c_int = 5;
static GDT_Float32:  c_int = 6;
static GDT_Float64:  c_int = 7;
static GDT_CInt16:   c_int = 8;
static GDT_CInt32:   c_int = 9;
static GDT_CFloat32: c_int = 10;
static GDT_CFloat64: c_int = 11;

static GF_Read:      c_int = 0;
static GF_Write:     c_int = 1;

pub struct Dataset {
    c_dataset: *(),
}

impl Drop for Dataset {
    fn drop(&mut self) {
        unsafe { GDALClose(self.c_dataset); }
    }
}

impl Dataset {
    pub fn get_raster_size(&self) -> (int, int) {
        let size_x = unsafe { GDALGetRasterXSize(self.c_dataset) } as int;
        let size_y = unsafe { GDALGetRasterYSize(self.c_dataset) } as int;
        return (size_x, size_y);
    }

    pub fn get_driver(&self) -> Driver {
        let c_driver = unsafe { GDALGetDatasetDriver(self.c_dataset) };
        return Driver{c_driver: c_driver};
    }

    pub fn get_raster_count(&self) -> int {
        return unsafe { GDALGetRasterCount(self.c_dataset) } as int;
    }

    pub fn get_projection(&self) -> ~str {
        unsafe {
            let rv = GDALGetProjectionRef(self.c_dataset);
            return raw::from_c_str(rv);
        }
    }

    pub fn set_projection(&self, projection: &str) {
        projection.with_c_str(|c_projection| {
            unsafe { GDALSetProjection(self.c_dataset, c_projection) };
        });
    }

    pub fn create_copy(
        &self,
        driver: Driver,
        filename: &str
    ) -> Option<Dataset> {
        use std::ptr::null;
        let c_dataset = filename.with_c_str(|c_filename| {
            unsafe {
                return GDALCreateCopy(
                    driver.c_driver,
                    c_filename,
                    self.c_dataset,
                    0,
                    null(),
                    null(),
                    null()
                )
            }
        });
        return match c_dataset.is_null() {
            true  => None,
            false => Some(Dataset{c_dataset: c_dataset}),
        };
    }

    pub fn read_raster(
        &self,
        band_index: int,
        window: Point<int>,
        window_size: Point<uint>,
        buffer_size: Point<uint>
    ) -> ByteBuffer {
        let buffer_size_bytes = buffer_size.x * buffer_size.y;
        let mut data: ~[u8] = slice::with_capacity(buffer_size_bytes);
        for _ in range(0, buffer_size_bytes) { data.push(0u8); } // TODO zero fill
        unsafe {
            let c_band = GDALGetRasterBand(self.c_dataset, band_index as c_int);
            let rv = GDALRasterIO(
                c_band,
                GF_Read,
                window.x as c_int,
                window.y as c_int,
                window_size.x as c_int,
                window_size.y as c_int,
                data.as_mut_ptr() as *(),
                buffer_size.x as c_int,
                buffer_size.y as c_int,
                GDT_Byte,
                0,
                0
            ) as int;
            assert!(rv == 0);
        };
        return ByteBuffer{
            size: buffer_size,
            data: data,
        };
    }

    pub fn write_raster(
        &self,
        band_index: int,
        window: Point<int>,
        window_size: Point<uint>,
        buffer: ByteBuffer
    ) {
        assert_eq!(buffer.data.len(), buffer.size.x * buffer.size.y);
        unsafe {
            let c_band = GDALGetRasterBand(self.c_dataset, band_index as c_int);
            let rv = GDALRasterIO(
                c_band,
                GF_Write,
                window.x as c_int,
                window.y as c_int,
                window_size.x as c_int,
                window_size.y as c_int,
                buffer.data.as_ptr() as *(),
                buffer.size.x as c_int,
                buffer.size.y as c_int,
                GDT_Byte,
                0,
                0
            ) as int;
            assert!(rv == 0);
        };
    }
}

pub fn open(path: &Path) -> Option<Dataset> {
    register_drivers();
    let filename = path.as_str().unwrap();
    let c_dataset = filename.with_c_str(|c_filename| {
        return unsafe { GDALOpen(c_filename, GA_ReadOnly) };
    });
    return match c_dataset.is_null() {
        true  => None,
        false => Some(Dataset{c_dataset: c_dataset}),
    };
}

struct ByteBuffer {
    size: Point<uint>,
    data: ~[u8],
}

fn fixture_path(name: &str) -> Path {
    let envvar = "RUSTILES_TEST_FIXTURES";
    let fixtures = match getenv(envvar) {
        Some(p) => Path::new(p),
        None => fail!("Environment variable {} not set", envvar)
    };
    let rv = fixtures.join(name);
    return rv;
}

#[test]
fn test_open() {
    let dataset = open(&fixture_path("tinymarble.jpeg"));
    assert!(dataset.is_some());

    let missing_dataset = open(&fixture_path("no_such_file.jpeg"));
    assert!(missing_dataset.is_none());
}

#[test]
fn test_get_raster_size() {
    let dataset = open(&fixture_path("tinymarble.jpeg")).unwrap();
    let (size_x, size_y) = dataset.get_raster_size();
    assert_eq!(size_x, 100);
    assert_eq!(size_y, 50);
}

#[test]
fn test_get_raster_count() {
    let dataset = open(&fixture_path("tinymarble.jpeg")).unwrap();
    let count = dataset.get_raster_count();
    assert_eq!(count, 3);
}

#[test]
fn test_get_projection() {
    let dataset = open(&fixture_path("tinymarble.jpeg")).unwrap();
    //dataset.set_projection("WGS84");
    let projection = dataset.get_projection();
    assert_eq!(projection.slice(0, 16), "GEOGCS[\"WGS 84\",");
}

#[test]
fn test_read_raster() {
    let dataset = open(&fixture_path("tinymarble.jpeg")).unwrap();
    let rv = dataset.read_raster(1, Point(20, 30), Point(10, 10), Point(3, 5));
    assert_eq!(rv.size.x, 3);
    assert_eq!(rv.size.y, 5);
    assert_eq!(rv.data, ~[13, 3, 18, 6, 9, 1, 2, 9, 4, 6, 11, 4, 6, 2, 9]);
}

#[test]
fn test_write_raster() {
    use gdal::driver::get_driver;
    let driver = get_driver("MEM").unwrap();
    let dataset = driver.create("", 20, 10, 1).unwrap();

    // create a 2x1 raster
    let raster = ByteBuffer{size: Point(2, 1), data: ~[50u8, 20u8]};

    // epand it to fill the image (20x10)
    dataset.write_raster(1, Point(0, 0), Point(20, 10), raster);

    // read a pixel from the left side
    let left = dataset.read_raster(1, Point(5, 5), Point(1, 1), Point(1, 1));
    assert_eq!(left.data[0], 50u8);

    // read a pixel from the right side
    let right = dataset.read_raster(1, Point(15, 5), Point(1, 1), Point(1, 1));
    assert_eq!(right.data[0], 20u8);
}

#[test]
fn test_get_dataset_driver() {
    let dataset = open(&fixture_path("tinymarble.jpeg")).unwrap();
    let driver = dataset.get_driver();
    assert_eq!(driver.get_short_name(), ~"JPEG");
    assert_eq!(driver.get_long_name(), ~"JPEG JFIF");
}

#[test]
fn test_create() {
    use gdal::driver::get_driver;
    let driver = get_driver("MEM").unwrap();
    let dataset = driver.create("", 10, 20, 3).unwrap();
    assert_eq!(dataset.get_raster_size(), (10, 20));
    assert_eq!(dataset.get_raster_count(), 3);
    assert_eq!(dataset.get_driver().get_short_name(), ~"MEM");
}

#[test]
fn test_create_copy() {
    use gdal::driver::get_driver;
    let driver = get_driver("MEM").unwrap();
    let dataset = open(&fixture_path("tinymarble.jpeg")).unwrap();
    let copy = dataset.create_copy(driver, "").unwrap();
    assert_eq!(copy.get_raster_size(), (100, 50));
    assert_eq!(copy.get_raster_count(), 3);
}
<commit_msg>wrap GDALSetGeoTransform<commit_after>use std::slice;
use std::libc::{c_int, c_char, c_double};
use std::str::raw;
use std::os::getenv;
use std::path::Path;
use gdal::driver::Driver;
use gdal::register_drivers;
use gdal::geom::Point;

#[link(name="gdal")]
extern {
    fn GDALCreateCopy(
            hDriver: *(),
            pszFilename: *c_char,
            hSrcDS: *(),
            bStrict: c_int,
            papszOptions: **c_char,
            pfnProgres: *(),
            pProgressData: *()
        ) -> *();
    fn GDALOpen(pszFilename: *c_char, eAccess: c_int) -> *();
    fn GDALClose(hDS: *());
    fn GDALGetDatasetDriver(hDataset: *()) -> *();
    fn GDALGetRasterXSize(hDataset: *()) -> c_int;
    fn GDALGetRasterYSize(hDataset: *()) -> c_int;
    fn GDALGetRasterCount(hDataset: *()) -> c_int;
    fn GDALGetProjectionRef(hDS: *()) -> *c_char;
    fn GDALSetProjection(hDS: *(), pszProjection: *c_char) -> c_int;
    fn GDALSetGeoTransform(hDS: *(), padfTransform: *c_double) -> c_int;
    fn GDALGetGeoTransform(hDS: *(), padfTransform: *mut c_double) -> c_int;
    fn GDALGetRasterBand(hDS: *(), nBandId: c_int) -> *();
    fn GDALRasterIO(
            hBand: *(),
            eRWFlag: c_int,
            nXOff: c_int,
            nYOff: c_int,
            nXSize: c_int,
            nYSize: c_int,
            pData: *(),
            nBufXSize: c_int,
            nBufYSize: c_int,
            GDALDataType: c_int,
            nPixelSpace: c_int,
            nLineSpace: c_int
        ) -> c_int;
}

static GA_ReadOnly:  c_int = 0;
static GA_Update:    c_int = 1;

static GDT_Unknown:  c_int = 0;
static GDT_Byte:     c_int = 1;
static GDT_UInt16:   c_int = 2;
static GDT_Int16:    c_int = 3;
static GDT_UInt32:   c_int = 4;
static GDT_Int32:    c_int = 5;
static GDT_Float32:  c_int = 6;
static GDT_Float64:  c_int = 7;
static GDT_CInt16:   c_int = 8;
static GDT_CInt32:   c_int = 9;
static GDT_CFloat32: c_int = 10;
static GDT_CFloat64: c_int = 11;

static GF_Read:      c_int = 0;
static GF_Write:     c_int = 1;

pub struct Dataset {
    c_dataset: *(),
}

impl Drop for Dataset {
    fn drop(&mut self) {
        unsafe { GDALClose(self.c_dataset); }
    }
}

impl Dataset {
    pub fn get_raster_size(&self) -> (int, int) {
        let size_x = unsafe { GDALGetRasterXSize(self.c_dataset) } as int;
        let size_y = unsafe { GDALGetRasterYSize(self.c_dataset) } as int;
        return (size_x, size_y);
    }

    pub fn get_driver(&self) -> Driver {
        let c_driver = unsafe { GDALGetDatasetDriver(self.c_dataset) };
        return Driver{c_driver: c_driver};
    }

    pub fn get_raster_count(&self) -> int {
        return unsafe { GDALGetRasterCount(self.c_dataset) } as int;
    }

    pub fn get_projection(&self) -> ~str {
        unsafe {
            let rv = GDALGetProjectionRef(self.c_dataset);
            return raw::from_c_str(rv);
        }
    }

    pub fn set_projection(&self, projection: &str) {
        projection.with_c_str(|c_projection| {
            unsafe { GDALSetProjection(self.c_dataset, c_projection) };
        });
    }

    pub fn set_geo_transform(&self, tr: (f64, f64, f64, f64, f64, f64)) {
        let (tr_0, tr_1, tr_2, tr_3, tr_4, tr_5) = tr;
        let tr_vec: ~[c_double] = ~[tr_0, tr_1, tr_2, tr_3, tr_4, tr_5];

        let rv = unsafe {
            GDALSetGeoTransform(self.c_dataset, tr_vec.as_ptr())
        } as int;
        assert!(rv == 0);
    }

    pub fn get_geo_transform(&self) -> (f64, f64, f64, f64, f64, f64) {
        let mut tr: ~[c_double] = slice::with_capacity(6);
        for _ in range(0, 6) { tr.push(0.0); }
        let rv = unsafe {
            GDALGetGeoTransform(
                self.c_dataset,
                tr.as_mut_ptr()
            )
        } as int;
        assert!(rv == 0);
        return (tr[0], tr[1], tr[2], tr[3], tr[4], tr[5]);
    }

    pub fn create_copy(
        &self,
        driver: Driver,
        filename: &str
    ) -> Option<Dataset> {
        use std::ptr::null;
        let c_dataset = filename.with_c_str(|c_filename| {
            unsafe {
                return GDALCreateCopy(
                    driver.c_driver,
                    c_filename,
                    self.c_dataset,
                    0,
                    null(),
                    null(),
                    null()
                )
            }
        });
        return match c_dataset.is_null() {
            true  => None,
            false => Some(Dataset{c_dataset: c_dataset}),
        };
    }

    pub fn read_raster(
        &self,
        band_index: int,
        window: Point<int>,
        window_size: Point<uint>,
        buffer_size: Point<uint>
    ) -> ByteBuffer {
        let buffer_size_bytes = buffer_size.x * buffer_size.y;
        let mut data: ~[u8] = slice::with_capacity(buffer_size_bytes);
        for _ in range(0, buffer_size_bytes) { data.push(0u8); } // TODO zero fill
        unsafe {
            let c_band = GDALGetRasterBand(self.c_dataset, band_index as c_int);
            let rv = GDALRasterIO(
                c_band,
                GF_Read,
                window.x as c_int,
                window.y as c_int,
                window_size.x as c_int,
                window_size.y as c_int,
                data.as_mut_ptr() as *(),
                buffer_size.x as c_int,
                buffer_size.y as c_int,
                GDT_Byte,
                0,
                0
            ) as int;
            assert!(rv == 0);
        };
        return ByteBuffer{
            size: buffer_size,
            data: data,
        };
    }

    pub fn write_raster(
        &self,
        band_index: int,
        window: Point<int>,
        window_size: Point<uint>,
        buffer: ByteBuffer
    ) {
        assert_eq!(buffer.data.len(), buffer.size.x * buffer.size.y);
        unsafe {
            let c_band = GDALGetRasterBand(self.c_dataset, band_index as c_int);
            let rv = GDALRasterIO(
                c_band,
                GF_Write,
                window.x as c_int,
                window.y as c_int,
                window_size.x as c_int,
                window_size.y as c_int,
                buffer.data.as_ptr() as *(),
                buffer.size.x as c_int,
                buffer.size.y as c_int,
                GDT_Byte,
                0,
                0
            ) as int;
            assert!(rv == 0);
        };
    }
}

pub fn open(path: &Path) -> Option<Dataset> {
    register_drivers();
    let filename = path.as_str().unwrap();
    let c_dataset = filename.with_c_str(|c_filename| {
        return unsafe { GDALOpen(c_filename, GA_ReadOnly) };
    });
    return match c_dataset.is_null() {
        true  => None,
        false => Some(Dataset{c_dataset: c_dataset}),
    };
}

struct ByteBuffer {
    size: Point<uint>,
    data: ~[u8],
}

fn fixture_path(name: &str) -> Path {
    let envvar = "RUSTILES_TEST_FIXTURES";
    let fixtures = match getenv(envvar) {
        Some(p) => Path::new(p),
        None => fail!("Environment variable {} not set", envvar)
    };
    let rv = fixtures.join(name);
    return rv;
}

#[test]
fn test_open() {
    let dataset = open(&fixture_path("tinymarble.jpeg"));
    assert!(dataset.is_some());

    let missing_dataset = open(&fixture_path("no_such_file.jpeg"));
    assert!(missing_dataset.is_none());
}

#[test]
fn test_get_raster_size() {
    let dataset = open(&fixture_path("tinymarble.jpeg")).unwrap();
    let (size_x, size_y) = dataset.get_raster_size();
    assert_eq!(size_x, 100);
    assert_eq!(size_y, 50);
}

#[test]
fn test_get_raster_count() {
    let dataset = open(&fixture_path("tinymarble.jpeg")).unwrap();
    let count = dataset.get_raster_count();
    assert_eq!(count, 3);
}

#[test]
fn test_get_projection() {
    let dataset = open(&fixture_path("tinymarble.jpeg")).unwrap();
    //dataset.set_projection("WGS84");
    let projection = dataset.get_projection();
    assert_eq!(projection.slice(0, 16), "GEOGCS[\"WGS 84\",");
}

#[test]
fn test_read_raster() {
    let dataset = open(&fixture_path("tinymarble.jpeg")).unwrap();
    let rv = dataset.read_raster(1, Point(20, 30), Point(10, 10), Point(3, 5));
    assert_eq!(rv.size.x, 3);
    assert_eq!(rv.size.y, 5);
    assert_eq!(rv.data, ~[13, 3, 18, 6, 9, 1, 2, 9, 4, 6, 11, 4, 6, 2, 9]);
}

#[test]
fn test_write_raster() {
    use gdal::driver::get_driver;
    let driver = get_driver("MEM").unwrap();
    let dataset = driver.create("", 20, 10, 1).unwrap();

    // create a 2x1 raster
    let raster = ByteBuffer{size: Point(2, 1), data: ~[50u8, 20u8]};

    // epand it to fill the image (20x10)
    dataset.write_raster(1, Point(0, 0), Point(20, 10), raster);

    // read a pixel from the left side
    let left = dataset.read_raster(1, Point(5, 5), Point(1, 1), Point(1, 1));
    assert_eq!(left.data[0], 50u8);

    // read a pixel from the right side
    let right = dataset.read_raster(1, Point(15, 5), Point(1, 1), Point(1, 1));
    assert_eq!(right.data[0], 20u8);
}

#[test]
fn test_get_dataset_driver() {
    let dataset = open(&fixture_path("tinymarble.jpeg")).unwrap();
    let driver = dataset.get_driver();
    assert_eq!(driver.get_short_name(), ~"JPEG");
    assert_eq!(driver.get_long_name(), ~"JPEG JFIF");
}

#[test]
fn test_create() {
    use gdal::driver::get_driver;
    let driver = get_driver("MEM").unwrap();
    let dataset = driver.create("", 10, 20, 3).unwrap();
    assert_eq!(dataset.get_raster_size(), (10, 20));
    assert_eq!(dataset.get_raster_count(), 3);
    assert_eq!(dataset.get_driver().get_short_name(), ~"MEM");
}

#[test]
fn test_create_copy() {
    use gdal::driver::get_driver;
    let driver = get_driver("MEM").unwrap();
    let dataset = open(&fixture_path("tinymarble.jpeg")).unwrap();
    let copy = dataset.create_copy(driver, "").unwrap();
    assert_eq!(copy.get_raster_size(), (100, 50));
    assert_eq!(copy.get_raster_count(), 3);
}

#[test]
fn test_geo_transform() {
    use gdal::driver::get_driver;
    let driver = get_driver("MEM").unwrap();
    let dataset = driver.create("", 20, 10, 1).unwrap();
    let transform = (0., 1., 0., 0., 0., 1.);
    dataset.set_geo_transform(transform);
    assert_eq!(dataset.get_geo_transform(), transform);
}
<|endoftext|>


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

use std::collections::HashSet;

use petgraph::{graph::NodeIndex, visit::EdgeRef, Direction};

use crate::graph::Graph;

pub(crate) struct GraphWalker {
    pub(crate) nodes_visited: HashSet<NodeIndex>,
}

impl GraphWalker {
    pub(crate) fn new() -> Self {
        let nodes_visited: HashSet<NodeIndex> = HashSet::new();

        Self { nodes_visited }
    }

    pub(crate) fn walk_graph(
        &mut self,
        graph: &Graph,
        origin_node_idx: NodeIndex,
        max_depth: usize,
    ) -> Graph {
        self.visit_node_recursively(graph, origin_node_idx, max_depth, 0);

        graph.to_owned()
    }

    pub(crate) fn visit_node_recursively(
        &mut self,
        graph: &Graph,
        node_idx: NodeIndex,
        max_depth: usize,
        depth: usize,
    ) {
        if depth > max_depth {
            return;
        }

        if self.nodes_visited.contains(&node_idx) {
            return;
        }

        self.nodes_visited.insert(node_idx);

        let edges_directed = graph.edges_directed(node_idx, Direction::Outgoing);

        for edge_ref in edges_directed {
            if depth > max_depth {
                return;
            }

            let target_node_idx = edge_ref.target();

            self.visit_node_recursively(graph, target_node_idx, max_depth, depth + 1);
        }
    }
}
<commit_msg>Refactor graph walker to be more versatile<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

use std::collections::HashSet;

use petgraph::{graph::NodeIndex, visit::EdgeRef, Direction};

use crate::graph::Graph;

use super::{edge::Edge, node::Node};

pub(crate) struct GraphWalker {
    direction: Direction,
    pub(crate) nodes_visited: HashSet<NodeIndex>,
}

impl GraphWalker {
    pub(crate) fn new(direction: Direction) -> Self {
        let nodes_visited = HashSet::default();

        Self {
            direction,
            nodes_visited,
        }
    }

    pub(crate) fn walk_graph<F>(&mut self, graph: &Graph, origin_node_idx: NodeIndex, predicate: F)
    where
        F: Fn(&Edge, &Node, usize) -> bool,
    {
        self.visit_node_recursively(graph, origin_node_idx, 0, &predicate);
    }

    pub(crate) fn visit_node_recursively<F>(
        &mut self,
        graph: &Graph,
        node_idx: NodeIndex,
        depth: usize,
        predicate: &F,
    ) where
        F: Fn(&Edge, &Node, usize) -> bool,
    {
        if self.nodes_visited.contains(&node_idx) {
            return;
        }

        self.nodes_visited.insert(node_idx);

        let edges_directed = graph.edges_directed(node_idx, self.direction);

        for edge_ref in edges_directed {
            let node_idx = match self.direction {
                Direction::Outgoing => edge_ref.target(),
                Direction::Incoming => edge_ref.source(),
            };

            if !predicate(edge_ref.weight(), &graph[node_idx], depth + 1) {
                continue;
            }

            self.visit_node_recursively(graph, node_idx, depth + 1, predicate);
        }
    }
}
<|endoftext|>


<commit_before>extern crate schannel;

use std::io;
use std::fmt;
use std::error;
use std::sync::Arc;
use self::schannel::cert_store::{PfxImportOptions, Memory, CertStore, CertAdd};
use self::schannel::cert_context::CertContext;
use self::schannel::schannel_cred::{Direction, SchannelCred, Protocol};
use self::schannel::tls_stream;

fn convert_protocols(protocols: &[::Protocol]) -> Vec<Protocol> {
    protocols
        .iter()
        .map(|p| match *p {
            ::Protocol::Sslv3 => Protocol::Ssl3,
            ::Protocol::Tlsv10 => Protocol::Tls10,
            ::Protocol::Tlsv11 => Protocol::Tls11,
            ::Protocol::Tlsv12 => Protocol::Tls12,
            ::Protocol::__NonExhaustive => unreachable!(),
        })
        .collect()
}

pub struct Error(io::Error);

impl error::Error for Error {
    fn description(&self) -> &str {
        error::Error::description(&self.0)
    }

    fn cause(&self) -> Option<&error::Error> {
        error::Error::cause(&self.0)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, fmt)
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl From<io::Error> for Error {
    fn from(error: io::Error) -> Error {
        Error(error)
    }
}

pub struct Pkcs12 {
    cert: CertContext,
}

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let mut store = try!(PfxImportOptions::new().password(pass).import(buf));
        let mut identity = None;

        for cert in store.certs() {
            if cert.private_key()
                .silent(true)
                .compare_key(true)
                .acquire()
                .is_ok()
            {
                identity = Some(cert);
            }
        }

        let identity = match identity {
            Some(identity) => identity,
            None => {
                return Err(
                    io::Error::new(
                        io::ErrorKind::InvalidInput,
                        "No identity found in PKCS #12 archive",
                    ).into(),
                );
            }
        };

        Ok(Pkcs12 { cert: identity })
    }
}

pub struct Certificate(CertContext);

impl Certificate {
    pub fn from_der(buf: &[u8]) -> Result<Certificate, Error> {
        let cert = try!(CertContext::new(buf));
        Ok(Certificate(cert))
    }
    pub fn from_pem(buf: &[u8]) -> Result<Certificate, Error> {
        match ::std::str::from_utf8(buf) {
            Ok(s) => {
                let cert = try!(CertContext::from_pem(s));
                Ok(Certificate(cert))
            }
            Err(e) => {
                io::Error::new(
                    io::ErrorKind::InvalidData,
                    "PEM representation contains non-UTF8 bytes",
                )
            }
        }
    }
}

pub struct MidHandshakeTlsStream<S>(tls_stream::MidHandshakeTlsStream<S>);

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
where
    S: fmt::Debug,
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S> MidHandshakeTlsStream<S>
where
    S: io::Read + io::Write,
{
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub enum HandshakeError<S> {
    Failure(Error),
    Interrupted(MidHandshakeTlsStream<S>),
}

impl<S> From<tls_stream::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: tls_stream::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            tls_stream::HandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            tls_stream::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream(s))
            }
        }
    }
}

impl<S> From<io::Error> for HandshakeError<S> {
    fn from(e: io::Error) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

pub struct TlsConnectorBuilder(TlsConnector);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        self.0.cert = Some(pkcs12.cert);
        Ok(())
    }

    pub fn add_root_certificate(&mut self, cert: Certificate) -> Result<(), Error> {
        try!(self.0.roots.add_cert(&cert.0, CertAdd::ReplaceExisting));
        Ok(())
    }

    pub fn supported_protocols(&mut self, protocols: &[::Protocol]) -> Result<(), Error> {
        self.0.protocols = convert_protocols(protocols);
        Ok(())
    }

    pub fn build(self) -> Result<TlsConnector, Error> {
        Ok(self.0)
    }
}

#[derive(Clone)]
pub struct TlsConnector {
    cert: Option<CertContext>,
    roots: CertStore,
    protocols: Vec<Protocol>,
    callback: Option<Arc<Fn(tls_stream::CertValidationResult) -> io::Result<()> + Sync + Send>>,
}

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        Ok(TlsConnectorBuilder(TlsConnector {
            cert: None,
            roots: try!(Memory::new()).into_store(),
            protocols: vec![Protocol::Tls10, Protocol::Tls11, Protocol::Tls12],
            callback: None,
        }))
    }

    pub fn connect<S>(&self, domain: &str, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        self._connect(Some(domain), stream)
    }

    pub fn connect_no_domain<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        self._connect(None, stream)
    }

    fn _connect<S>(
        &self,
        domain: Option<&str>,
        stream: S,
    ) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let mut builder = SchannelCred::builder();
        builder.enabled_protocols(&self.protocols);
        if let Some(cert) = self.cert.as_ref() {
            builder.cert(cert.clone());
        }
        let cred = try!(builder.acquire(Direction::Outbound));
        let mut builder = tls_stream::Builder::new();
        if let Some(domain) = domain {
            builder.domain(domain);
        }
        if let Some(ref callback) = self.callback {
            let callback = callback.clone();
            builder.verify_callback(move |r| callback(r));
        }
        builder.cert_store(self.roots.clone());
        match builder.connect(cred, stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsAcceptorBuilder(TlsAcceptor);

impl TlsAcceptorBuilder {
    pub fn supported_protocols(&mut self, protocols: &[::Protocol]) -> Result<(), Error> {
        self.0.protocols = convert_protocols(protocols);
        Ok(())
    }

    pub fn build(self) -> Result<TlsAcceptor, Error> {
        Ok(self.0)
    }
}

#[derive(Clone)]
pub struct TlsAcceptor {
    cert: CertContext,
    protocols: Vec<Protocol>,
}

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        Ok(TlsAcceptorBuilder(TlsAcceptor {
            cert: pkcs12.cert,
            protocols: vec![Protocol::Tls10, Protocol::Tls11, Protocol::Tls12],
        }))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let mut builder = SchannelCred::builder();
        builder.enabled_protocols(&self.protocols);
        builder.cert(self.cert.clone());
        // FIXME we're probably missing the certificate chain?
        let cred = try!(builder.acquire(Direction::Inbound));
        match tls_stream::Builder::new().accept(cred, stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsStream<S>(tls_stream::TlsStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(self.0.get_buf().len())
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        try!(self.0.shutdown());
        Ok(())
    }

    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }
}

impl<S: io::Read + io::Write> io::Read for TlsStream<S> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

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

/// SChannel-specific extensions to `TlsConnectorBuilder`.
pub trait TlsConnectorBuilderExt {
    /// Sets a callback function which decides if the server's certificate chain
    /// is to be trusted.
    fn verify_callback<F>(&mut self, callback: F)
    where
        F: Fn(tls_stream::CertValidationResult) -> io::Result<()> + 'static + Send + Sync;
}

impl TlsConnectorBuilderExt for ::TlsConnectorBuilder {
    fn verify_callback<F>(&mut self, callback: F)
    where
        F: Fn(tls_stream::CertValidationResult) -> io::Result<()> + 'static + Send + Sync,
    {
        (self.0).0.callback = Some(Arc::new(callback));
    }
}

/// SChannel-specific extensions to `TlsStream`.
pub trait TlsStreamExt<S> {
    /// Returns a shared reference to the SChannel `TlsStream`.
    fn raw_stream(&self) -> &tls_stream::TlsStream<S>;

    /// Returns a mutable reference to the SChannel `TlsStream`.
    fn raw_stream_mut(&mut self) -> &mut tls_stream::TlsStream<S>;
}

impl<S> TlsStreamExt<S> for ::TlsStream<S> {
    fn raw_stream(&self) -> &tls_stream::TlsStream<S> {
        &(self.0).0
    }

    fn raw_stream_mut(&mut self) -> &mut tls_stream::TlsStream<S> {
        &mut (self.0).0
    }
}

/// SChannel-specific extensions to `Error`
pub trait ErrorExt {
    /// Extract the underlying SChannel error for inspection.
    fn schannel_error(&self) -> &io::Error;
}

impl ErrorExt for ::Error {
    fn schannel_error(&self) -> &io::Error {
        &(self.0).0
    }
}

/// SChannel-specific extensions to `Certificate`.
pub trait CertificateExt {
    /// builds a native_Tls `Certificate` from an schannel `CertContext`
    fn from_cert_context(CertContext) -> ::Certificate;
}

impl CertificateExt for ::Certificate {
    fn from_cert_context(cert: CertContext) -> ::Certificate {
        ::Certificate(Certificate(cert))
    }
}
<commit_msg>Fix another schannel build error<commit_after>extern crate schannel;

use std::io;
use std::fmt;
use std::error;
use std::sync::Arc;
use self::schannel::cert_store::{PfxImportOptions, Memory, CertStore, CertAdd};
use self::schannel::cert_context::CertContext;
use self::schannel::schannel_cred::{Direction, SchannelCred, Protocol};
use self::schannel::tls_stream;

fn convert_protocols(protocols: &[::Protocol]) -> Vec<Protocol> {
    protocols
        .iter()
        .map(|p| match *p {
            ::Protocol::Sslv3 => Protocol::Ssl3,
            ::Protocol::Tlsv10 => Protocol::Tls10,
            ::Protocol::Tlsv11 => Protocol::Tls11,
            ::Protocol::Tlsv12 => Protocol::Tls12,
            ::Protocol::__NonExhaustive => unreachable!(),
        })
        .collect()
}

pub struct Error(io::Error);

impl error::Error for Error {
    fn description(&self) -> &str {
        error::Error::description(&self.0)
    }

    fn cause(&self) -> Option<&error::Error> {
        error::Error::cause(&self.0)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, fmt)
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl From<io::Error> for Error {
    fn from(error: io::Error) -> Error {
        Error(error)
    }
}

pub struct Pkcs12 {
    cert: CertContext,
}

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let mut store = try!(PfxImportOptions::new().password(pass).import(buf));
        let mut identity = None;

        for cert in store.certs() {
            if cert.private_key()
                .silent(true)
                .compare_key(true)
                .acquire()
                .is_ok()
            {
                identity = Some(cert);
            }
        }

        let identity = match identity {
            Some(identity) => identity,
            None => {
                return Err(
                    io::Error::new(
                        io::ErrorKind::InvalidInput,
                        "No identity found in PKCS #12 archive",
                    ).into(),
                );
            }
        };

        Ok(Pkcs12 { cert: identity })
    }
}

pub struct Certificate(CertContext);

impl Certificate {
    pub fn from_der(buf: &[u8]) -> Result<Certificate, Error> {
        let cert = try!(CertContext::new(buf));
        Ok(Certificate(cert))
    }
    pub fn from_pem(buf: &[u8]) -> Result<Certificate, Error> {
        match ::std::str::from_utf8(buf) {
            Ok(s) => {
                let cert = try!(CertContext::from_pem(s));
                Ok(Certificate(cert))
            }
            Err(e) => {
                Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!(
                        "PEM representation contains non-UTF8 bytes {}",
                        e.description()
                    ),
                ))
            }
        }
    }
}

pub struct MidHandshakeTlsStream<S>(tls_stream::MidHandshakeTlsStream<S>);

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
where
    S: fmt::Debug,
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S> MidHandshakeTlsStream<S>
where
    S: io::Read + io::Write,
{
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub enum HandshakeError<S> {
    Failure(Error),
    Interrupted(MidHandshakeTlsStream<S>),
}

impl<S> From<tls_stream::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: tls_stream::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            tls_stream::HandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            tls_stream::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream(s))
            }
        }
    }
}

impl<S> From<io::Error> for HandshakeError<S> {
    fn from(e: io::Error) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

pub struct TlsConnectorBuilder(TlsConnector);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        self.0.cert = Some(pkcs12.cert);
        Ok(())
    }

    pub fn add_root_certificate(&mut self, cert: Certificate) -> Result<(), Error> {
        try!(self.0.roots.add_cert(&cert.0, CertAdd::ReplaceExisting));
        Ok(())
    }

    pub fn supported_protocols(&mut self, protocols: &[::Protocol]) -> Result<(), Error> {
        self.0.protocols = convert_protocols(protocols);
        Ok(())
    }

    pub fn build(self) -> Result<TlsConnector, Error> {
        Ok(self.0)
    }
}

#[derive(Clone)]
pub struct TlsConnector {
    cert: Option<CertContext>,
    roots: CertStore,
    protocols: Vec<Protocol>,
    callback: Option<Arc<Fn(tls_stream::CertValidationResult) -> io::Result<()> + Sync + Send>>,
}

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        Ok(TlsConnectorBuilder(TlsConnector {
            cert: None,
            roots: try!(Memory::new()).into_store(),
            protocols: vec![Protocol::Tls10, Protocol::Tls11, Protocol::Tls12],
            callback: None,
        }))
    }

    pub fn connect<S>(&self, domain: &str, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        self._connect(Some(domain), stream)
    }

    pub fn connect_no_domain<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        self._connect(None, stream)
    }

    fn _connect<S>(
        &self,
        domain: Option<&str>,
        stream: S,
    ) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let mut builder = SchannelCred::builder();
        builder.enabled_protocols(&self.protocols);
        if let Some(cert) = self.cert.as_ref() {
            builder.cert(cert.clone());
        }
        let cred = try!(builder.acquire(Direction::Outbound));
        let mut builder = tls_stream::Builder::new();
        if let Some(domain) = domain {
            builder.domain(domain);
        }
        if let Some(ref callback) = self.callback {
            let callback = callback.clone();
            builder.verify_callback(move |r| callback(r));
        }
        builder.cert_store(self.roots.clone());
        match builder.connect(cred, stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsAcceptorBuilder(TlsAcceptor);

impl TlsAcceptorBuilder {
    pub fn supported_protocols(&mut self, protocols: &[::Protocol]) -> Result<(), Error> {
        self.0.protocols = convert_protocols(protocols);
        Ok(())
    }

    pub fn build(self) -> Result<TlsAcceptor, Error> {
        Ok(self.0)
    }
}

#[derive(Clone)]
pub struct TlsAcceptor {
    cert: CertContext,
    protocols: Vec<Protocol>,
}

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        Ok(TlsAcceptorBuilder(TlsAcceptor {
            cert: pkcs12.cert,
            protocols: vec![Protocol::Tls10, Protocol::Tls11, Protocol::Tls12],
        }))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let mut builder = SchannelCred::builder();
        builder.enabled_protocols(&self.protocols);
        builder.cert(self.cert.clone());
        // FIXME we're probably missing the certificate chain?
        let cred = try!(builder.acquire(Direction::Inbound));
        match tls_stream::Builder::new().accept(cred, stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsStream<S>(tls_stream::TlsStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(self.0.get_buf().len())
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        try!(self.0.shutdown());
        Ok(())
    }

    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }
}

impl<S: io::Read + io::Write> io::Read for TlsStream<S> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

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

/// SChannel-specific extensions to `TlsConnectorBuilder`.
pub trait TlsConnectorBuilderExt {
    /// Sets a callback function which decides if the server's certificate chain
    /// is to be trusted.
    fn verify_callback<F>(&mut self, callback: F)
    where
        F: Fn(tls_stream::CertValidationResult) -> io::Result<()> + 'static + Send + Sync;
}

impl TlsConnectorBuilderExt for ::TlsConnectorBuilder {
    fn verify_callback<F>(&mut self, callback: F)
    where
        F: Fn(tls_stream::CertValidationResult) -> io::Result<()> + 'static + Send + Sync,
    {
        (self.0).0.callback = Some(Arc::new(callback));
    }
}

/// SChannel-specific extensions to `TlsStream`.
pub trait TlsStreamExt<S> {
    /// Returns a shared reference to the SChannel `TlsStream`.
    fn raw_stream(&self) -> &tls_stream::TlsStream<S>;

    /// Returns a mutable reference to the SChannel `TlsStream`.
    fn raw_stream_mut(&mut self) -> &mut tls_stream::TlsStream<S>;
}

impl<S> TlsStreamExt<S> for ::TlsStream<S> {
    fn raw_stream(&self) -> &tls_stream::TlsStream<S> {
        &(self.0).0
    }

    fn raw_stream_mut(&mut self) -> &mut tls_stream::TlsStream<S> {
        &mut (self.0).0
    }
}

/// SChannel-specific extensions to `Error`
pub trait ErrorExt {
    /// Extract the underlying SChannel error for inspection.
    fn schannel_error(&self) -> &io::Error;
}

impl ErrorExt for ::Error {
    fn schannel_error(&self) -> &io::Error {
        &(self.0).0
    }
}

/// SChannel-specific extensions to `Certificate`.
pub trait CertificateExt {
    /// builds a native_Tls `Certificate` from an schannel `CertContext`
    fn from_cert_context(CertContext) -> ::Certificate;
}

impl CertificateExt for ::Certificate {
    fn from_cert_context(cert: CertContext) -> ::Certificate {
        ::Certificate(Certificate(cert))
    }
}
<|endoftext|>


<commit_before>use super::internal::*;
use super::*;
use std::cmp;

/// `MinLen` is an iterator that imposes a minimum length on iterator splits.
/// This struct is created by the [`min_len()`] method on [`IndexedParallelIterator`]
///
/// [`min_len()`]: trait.IndexedParallelIterator.html#method.min_len
/// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html
pub struct MinLen<M: IndexedParallelIterator> {
    base: M,
    min: usize,
}

/// Create a new `MinLen` iterator.
///
/// NB: a free fn because it is NOT part of the end-user API.
pub fn new_min_len<M>(base: M, min: usize) -> MinLen<M>
    where M: IndexedParallelIterator
{
    MinLen { base: base, min: min }
}

impl<M> ParallelIterator for MinLen<M>
    where M: IndexedParallelIterator
{
    type Item = M::Item;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge(self, consumer)
    }

    fn opt_len(&mut self) -> Option<usize> {
        Some(self.len())
    }
}

impl<M> BoundedParallelIterator for MinLen<M>
    where M: IndexedParallelIterator
{
    fn upper_bound(&mut self) -> usize {
        self.len()
    }

    fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
        bridge(self, consumer)
    }
}

impl<M> ExactParallelIterator for MinLen<M>
    where M: IndexedParallelIterator
{
    fn len(&mut self) -> usize {
        self.base.len()
    }
}

impl<M> IndexedParallelIterator for MinLen<M>
    where M: IndexedParallelIterator
{
    fn with_producer<CB>(self, callback: CB) -> CB::Output
        where CB: ProducerCallback<Self::Item>
    {
        return self.base.with_producer(Callback {
            callback: callback,
            min: self.min,
        });

        struct Callback<CB> {
            callback: CB,
            min: usize,
        }

        impl<ITEM, CB> ProducerCallback<ITEM> for Callback<CB>
            where CB: ProducerCallback<ITEM>
        {
            type Output = CB::Output;
            fn callback
(self, base: P) -> CB::Output
                where P: Producer<Item = ITEM>
            {
                let producer = MinLenProducer {
                    base: base,
                    min: self.min,
                };
                self.callback.callback(producer)
            }
        }
    }
}

/// ////////////////////////////////////////////////////////////////////////
/// `MinLenProducer` implementation

struct MinLenProducer
 {
    base: P,
    min: usize,
}

impl
 Producer for MinLenProducer

    where P: Producer
{
    type Item = P::Item;
    type IntoIter = P::IntoIter;

    fn into_iter(self) -> Self::IntoIter {
        self.base.into_iter()
    }

    fn min_len(&self) -> usize {
        cmp::max(self.min, self.base.min_len())
    }

    fn max_len(&self) -> usize { self.base.max_len() }

    fn split_at(self, index: usize) -> (Self, Self) {
        let (left, right) = self.base.split_at(index);
        (MinLenProducer {
             base: left,
             min: self.min,
         },
         MinLenProducer {
             base: right,
             min: self.min,
         })
    }
}

/// `MaxLen` is an iterator that imposes a maximum length on iterator splits.
/// This struct is created by the [`max_len()`] method on [`IndexedParallelIterator`]
///
/// [`max_len()`]: trait.IndexedParallelIterator.html#method.max_len
/// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html
pub struct MaxLen<M: IndexedParallelIterator> {
    base: M,
    max: usize,
}

/// Create a new `MaxLen` iterator.
///
/// NB: a free fn because it is NOT part of the end-user API.
pub fn new_max_len<M>(base: M, max: usize) -> MaxLen<M>
    where M: IndexedParallelIterator
{
    MaxLen { base: base, max: max }
}

impl<M> ParallelIterator for MaxLen<M>
    where M: IndexedParallelIterator
{
    type Item = M::Item;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge(self, consumer)
    }

    fn opt_len(&mut self) -> Option<usize> {
        Some(self.len())
    }
}

impl<M> BoundedParallelIterator for MaxLen<M>
    where M: IndexedParallelIterator
{
    fn upper_bound(&mut self) -> usize {
        self.len()
    }

    fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
        bridge(self, consumer)
    }
}

impl<M> ExactParallelIterator for MaxLen<M>
    where M: IndexedParallelIterator
{
    fn len(&mut self) -> usize {
        self.base.len()
    }
}

impl<M> IndexedParallelIterator for MaxLen<M>
    where M: IndexedParallelIterator
{
    fn with_producer<CB>(self, callback: CB) -> CB::Output
        where CB: ProducerCallback<Self::Item>
    {
        return self.base.with_producer(Callback {
            callback: callback,
            max: self.max,
        });

        struct Callback<CB> {
            callback: CB,
            max: usize,
        }

        impl<ITEM, CB> ProducerCallback<ITEM> for Callback<CB>
            where CB: ProducerCallback<ITEM>
        {
            type Output = CB::Output;
            fn callback
(self, base: P) -> CB::Output
                where P: Producer<Item = ITEM>
            {
                let producer = MaxLenProducer {
                    base: base,
                    max: self.max,
                };
                self.callback.callback(producer)
            }
        }
    }
}

/// ////////////////////////////////////////////////////////////////////////
/// `MaxLenProducer` implementation

struct MaxLenProducer
 {
    base: P,
    max: usize,
}

impl
 Producer for MaxLenProducer

    where P: Producer
{
    type Item = P::Item;
    type IntoIter = P::IntoIter;

    fn into_iter(self) -> Self::IntoIter {
        self.base.into_iter()
    }

    fn min_len(&self) -> usize { self.base.min_len() }

    fn max_len(&self) -> usize {
        cmp::min(self.max, self.base.max_len())
    }

    fn split_at(self, index: usize) -> (Self, Self) {
        let (left, right) = self.base.split_at(index);
        (MaxLenProducer {
             base: left,
             max: self.max,
         },
         MaxLenProducer {
             base: right,
             max: self.max,
         })
    }
}
<commit_msg>Use short type params like #261<commit_after>use super::internal::*;
use super::*;
use std::cmp;

/// `MinLen` is an iterator that imposes a minimum length on iterator splits.
/// This struct is created by the [`min_len()`] method on [`IndexedParallelIterator`]
///
/// [`min_len()`]: trait.IndexedParallelIterator.html#method.min_len
/// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html
pub struct MinLen {
    base: I,
    min: usize,
}

/// Create a new `MinLen` iterator.
///
/// NB: a free fn because it is NOT part of the end-user API.
pub fn new_min_len(base: I, min: usize) -> MinLen
    where I: IndexedParallelIterator
{
    MinLen { base: base, min: min }
}

impl ParallelIterator for MinLen
    where I: IndexedParallelIterator
{
    type Item = I::Item;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge(self, consumer)
    }

    fn opt_len(&mut self) -> Option<usize> {
        Some(self.len())
    }
}

impl BoundedParallelIterator for MinLen
    where I: IndexedParallelIterator
{
    fn upper_bound(&mut self) -> usize {
        self.len()
    }

    fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
        bridge(self, consumer)
    }
}

impl ExactParallelIterator for MinLen
    where I: IndexedParallelIterator
{
    fn len(&mut self) -> usize {
        self.base.len()
    }
}

impl IndexedParallelIterator for MinLen
    where I: IndexedParallelIterator
{
    fn with_producer<CB>(self, callback: CB) -> CB::Output
        where CB: ProducerCallback<Self::Item>
    {
        return self.base.with_producer(Callback {
            callback: callback,
            min: self.min,
        });

        struct Callback<CB> {
            callback: CB,
            min: usize,
        }

        impl<T, CB> ProducerCallback<T> for Callback<CB>
            where CB: ProducerCallback<T>
        {
            type Output = CB::Output;
            fn callback
(self, base: P) -> CB::Output
                where P: Producer<Item = T>
            {
                let producer = MinLenProducer {
                    base: base,
                    min: self.min,
                };
                self.callback.callback(producer)
            }
        }
    }
}

/// ////////////////////////////////////////////////////////////////////////
/// `MinLenProducer` implementation

struct MinLenProducer
 {
    base: P,
    min: usize,
}

impl
 Producer for MinLenProducer

    where P: Producer
{
    type Item = P::Item;
    type IntoIter = P::IntoIter;

    fn into_iter(self) -> Self::IntoIter {
        self.base.into_iter()
    }

    fn min_len(&self) -> usize {
        cmp::max(self.min, self.base.min_len())
    }

    fn max_len(&self) -> usize { self.base.max_len() }

    fn split_at(self, index: usize) -> (Self, Self) {
        let (left, right) = self.base.split_at(index);
        (MinLenProducer {
             base: left,
             min: self.min,
         },
         MinLenProducer {
             base: right,
             min: self.min,
         })
    }
}

/// `MaxLen` is an iterator that imposes a maximum length on iterator splits.
/// This struct is created by the [`max_len()`] method on [`IndexedParallelIterator`]
///
/// [`max_len()`]: trait.IndexedParallelIterator.html#method.max_len
/// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html
pub struct MaxLen {
    base: I,
    max: usize,
}

/// Create a new `MaxLen` iterator.
///
/// NB: a free fn because it is NOT part of the end-user API.
pub fn new_max_len(base: I, max: usize) -> MaxLen
    where I: IndexedParallelIterator
{
    MaxLen { base: base, max: max }
}

impl ParallelIterator for MaxLen
    where I: IndexedParallelIterator
{
    type Item = I::Item;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge(self, consumer)
    }

    fn opt_len(&mut self) -> Option<usize> {
        Some(self.len())
    }
}

impl BoundedParallelIterator for MaxLen
    where I: IndexedParallelIterator
{
    fn upper_bound(&mut self) -> usize {
        self.len()
    }

    fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
        bridge(self, consumer)
    }
}

impl ExactParallelIterator for MaxLen
    where I: IndexedParallelIterator
{
    fn len(&mut self) -> usize {
        self.base.len()
    }
}

impl IndexedParallelIterator for MaxLen
    where I: IndexedParallelIterator
{
    fn with_producer<CB>(self, callback: CB) -> CB::Output
        where CB: ProducerCallback<Self::Item>
    {
        return self.base.with_producer(Callback {
            callback: callback,
            max: self.max,
        });

        struct Callback<CB> {
            callback: CB,
            max: usize,
        }

        impl<T, CB> ProducerCallback<T> for Callback<CB>
            where CB: ProducerCallback<T>
        {
            type Output = CB::Output;
            fn callback
(self, base: P) -> CB::Output
                where P: Producer<Item = T>
            {
                let producer = MaxLenProducer {
                    base: base,
                    max: self.max,
                };
                self.callback.callback(producer)
            }
        }
    }
}

/// ////////////////////////////////////////////////////////////////////////
/// `MaxLenProducer` implementation

struct MaxLenProducer
 {
    base: P,
    max: usize,
}

impl
 Producer for MaxLenProducer

    where P: Producer
{
    type Item = P::Item;
    type IntoIter = P::IntoIter;

    fn into_iter(self) -> Self::IntoIter {
        self.base.into_iter()
    }

    fn min_len(&self) -> usize { self.base.min_len() }

    fn max_len(&self) -> usize {
        cmp::min(self.max, self.base.max_len())
    }

    fn split_at(self, index: usize) -> (Self, Self) {
        let (left, right) = self.base.split_at(index);
        (MaxLenProducer {
             base: left,
             max: self.max,
         },
         MaxLenProducer {
             base: right,
             max: self.max,
         })
    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![stable]

//! Concurrency-enabled mechanisms for sharing mutable and/or immutable state
//! between tasks.

use core::atomic;
use core::clone::Clone;
use core::kinds::{Sync, Send};
use core::mem::{min_align_of, size_of, drop};
use core::mem;
use core::ops::{Drop, Deref};
use core::option::{Some, None, Option};
use core::ptr;
use core::ptr::RawPtr;
use heap::deallocate;

/// An atomically reference counted wrapper for shared state.
///
/// # Example
///
/// In this example, a large vector of floats is shared between several tasks.
/// With simple pipes, without `Arc`, a copy would have to be made for each
/// task.
///
/// ```rust
/// use std::sync::Arc;
///
/// fn main() {
///     let numbers = Vec::from_fn(100, |i| i as f32);
///     let shared_numbers = Arc::new(numbers);
///
///     for _ in range(0u, 10) {
///         let child_numbers = shared_numbers.clone();
///
///         spawn(proc() {
///             let local_numbers = child_numbers.as_slice();
///
///             // Work with the local numbers
///         });
///     }
/// }
/// ```
#[unsafe_no_drop_flag]
#[stable]
pub struct Arc<T> {
    // FIXME #12808: strange name to try to avoid interfering with
    // field accesses of the contained type via Deref
    _ptr: *mut ArcInner<T>,
}

/// A weak pointer to an `Arc`.
///
/// Weak pointers will not keep the data inside of the `Arc` alive, and can be
/// used to break cycles between `Arc` pointers.
#[unsafe_no_drop_flag]
#[experimental = "Weak pointers may not belong in this module."]
pub struct Weak<T> {
    // FIXME #12808: strange name to try to avoid interfering with
    // field accesses of the contained type via Deref
    _ptr: *mut ArcInner<T>,
}

struct ArcInner<T> {
    strong: atomic::AtomicUint,
    weak: atomic::AtomicUint,
    data: T,
}

impl<T: Sync + Send> Arc<T> {
    /// Creates an atomically reference counted wrapper.
    #[inline]
    #[stable]
    pub fn new(data: T) -> Arc<T> {
        // Start the weak pointer count as 1 which is the weak pointer that's
        // held by all the strong pointers (kinda), see std/rc.rs for more info
        let x = box ArcInner {
            strong: atomic::AtomicUint::new(1),
            weak: atomic::AtomicUint::new(1),
            data: data,
        };
        Arc { _ptr: unsafe { mem::transmute(x) } }
    }

    #[inline]
    fn inner(&self) -> &ArcInner<T> {
        // This unsafety is ok because while this arc is alive we're guaranteed
        // that the inner pointer is valid. Furthermore, we know that the
        // `ArcInner` structure itself is `Sync` because the inner data is
        // `Sync` as well, so we're ok loaning out an immutable pointer to
        // these contents.
        unsafe { &*self._ptr }
    }

    /// Downgrades a strong pointer to a weak pointer.
    ///
    /// Weak pointers will not keep the data alive. Once all strong references
    /// to the underlying data have been dropped, the data itself will be
    /// destroyed.
    #[experimental = "Weak pointers may not belong in this module."]
    pub fn downgrade(&self) -> Weak<T> {
        // See the clone() impl for why this is relaxed
        self.inner().weak.fetch_add(1, atomic::Relaxed);
        Weak { _ptr: self._ptr }
    }
}

#[unstable = "waiting on stability of Clone"]
impl<T: Sync + Send> Clone for Arc<T> {
    /// Duplicate an atomically reference counted wrapper.
    ///
    /// The resulting two `Arc` objects will point to the same underlying data
    /// object. However, one of the `Arc` objects can be sent to another task,
    /// allowing them to share the underlying data.
    #[inline]
    fn clone(&self) -> Arc<T> {
        // Using a relaxed ordering is alright here, as knowledge of the
        // original reference prevents other threads from erroneously deleting
        // the object.
        //
        // As explained in the [Boost documentation][1], Increasing the
        // reference counter can always be done with memory_order_relaxed: New
        // references to an object can only be formed from an existing
        // reference, and passing an existing reference from one thread to
        // another must already provide any required synchronization.
        //
        // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html)
        self.inner().strong.fetch_add(1, atomic::Relaxed);
        Arc { _ptr: self._ptr }
    }
}

#[experimental = "Deref is experimental."]
impl<T: Send + Sync> Deref<T> for Arc<T> {
    #[inline]
    fn deref(&self) -> &T {
        &self.inner().data
    }
}

impl<T: Send + Sync + Clone> Arc<T> {
    /// Acquires a mutable pointer to the inner contents by guaranteeing that
    /// the reference count is one (no sharing is possible).
    ///
    /// This is also referred to as a copy-on-write operation because the inner
    /// data is cloned if the reference count is greater than one.
    #[inline]
    #[experimental]
    pub fn make_unique(&mut self) -> &mut T {
        // Note that we hold a strong reference, which also counts as
        // a weak reference, so we only clone if there is an
        // additional reference of either kind.
        if self.inner().strong.load(atomic::SeqCst) != 1 ||
           self.inner().weak.load(atomic::SeqCst) != 1 {
            *self = Arc::new(self.deref().clone())
        }
        // This unsafety is ok because we're guaranteed that the pointer
        // returned is the *only* pointer that will ever be returned to T. Our
        // reference count is guaranteed to be 1 at this point, and we required
        // the Arc itself to be `mut`, so we're returning the only possible
        // reference to the inner data.
        let inner = unsafe { &mut *self._ptr };
        &mut inner.data
    }
}

#[unsafe_destructor]
#[experimental = "waiting on stability of Drop"]
impl<T: Sync + Send> Drop for Arc<T> {
    fn drop(&mut self) {
        // This structure has #[unsafe_no_drop_flag], so this drop glue may run
        // more than once (but it is guaranteed to be zeroed after the first if
        // it's run more than once)
        if self._ptr.is_null() { return }

        // Because `fetch_sub` is already atomic, we do not need to synchronize
        // with other threads unless we are going to delete the object. This
        // same logic applies to the below `fetch_sub` to the `weak` count.
        if self.inner().strong.fetch_sub(1, atomic::Release) != 1 { return }

        // This fence is needed to prevent reordering of use of the data and
        // deletion of the data. Because it is marked `Release`, the
        // decreasing of the reference count synchronizes with this `Acquire`
        // fence. This means that use of the data happens before decreasing
        // the reference count, which happens before this fence, which
        // happens before the deletion of the data.
        //
        // As explained in the [Boost documentation][1],
        //
        // It is important to enforce any possible access to the object in
        // one thread (through an existing reference) to *happen before*
        // deleting the object in a different thread. This is achieved by a
        // "release" operation after dropping a reference (any access to the
        // object through this reference must obviously happened before),
        // and an "acquire" operation before deleting the object.
        //
        // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html)
        atomic::fence(atomic::Acquire);

        // Destroy the data at this time, even though we may not free the box
        // allocation itself (there may still be weak pointers lying around).
        unsafe { drop(ptr::read(&self.inner().data)); }

        if self.inner().weak.fetch_sub(1, atomic::Release) == 1 {
            atomic::fence(atomic::Acquire);
            unsafe { deallocate(self._ptr as *mut u8, size_of::<ArcInner<T>>(),
                                min_align_of::<ArcInner<T>>()) }
        }
    }
}

#[experimental = "Weak pointers may not belong in this module."]
impl<T: Sync + Send> Weak<T> {
    /// Attempts to upgrade this weak reference to a strong reference.
    ///
    /// This method will fail to upgrade this reference if the strong reference
    /// count has already reached 0, but if there are still other active strong
    /// references this function will return a new strong reference to the data.
    pub fn upgrade(&self) -> Option<Arc<T>> {
        // We use a CAS loop to increment the strong count instead of a
        // fetch_add because once the count hits 0 is must never be above 0.
        let inner = self.inner();
        loop {
            let n = inner.strong.load(atomic::SeqCst);
            if n == 0 { return None }
            let old = inner.strong.compare_and_swap(n, n + 1, atomic::SeqCst);
            if old == n { return Some(Arc { _ptr: self._ptr }) }
        }
    }

    #[inline]
    fn inner(&self) -> &ArcInner<T> {
        // See comments above for why this is "safe"
        unsafe { &*self._ptr }
    }
}

#[experimental = "Weak pointers may not belong in this module."]
impl<T: Sync + Send> Clone for Weak<T> {
    #[inline]
    fn clone(&self) -> Weak<T> {
        // See comments in Arc::clone() for why this is relaxed
        self.inner().weak.fetch_add(1, atomic::Relaxed);
        Weak { _ptr: self._ptr }
    }
}

#[unsafe_destructor]
#[experimental = "Weak pointers may not belong in this module."]
impl<T: Sync + Send> Drop for Weak<T> {
    fn drop(&mut self) {
        // see comments above for why this check is here
        if self._ptr.is_null() { return }

        // If we find out that we were the last weak pointer, then its time to
        // deallocate the data entirely. See the discussion in Arc::drop() about
        // the memory orderings
        if self.inner().weak.fetch_sub(1, atomic::Release) == 1 {
            atomic::fence(atomic::Acquire);
            unsafe { deallocate(self._ptr as *mut u8, size_of::<ArcInner<T>>(),
                                min_align_of::<ArcInner<T>>()) }
        }
    }
}

#[cfg(test)]
#[allow(experimental)]
mod tests {
    use std::clone::Clone;
    use std::comm::channel;
    use std::mem::drop;
    use std::ops::Drop;
    use std::option::{Option, Some, None};
    use std::sync::atomic;
    use std::task;
    use std::vec::Vec;
    use super::{Arc, Weak};
    use std::sync::Mutex;

    struct Canary(*mut atomic::AtomicUint);

    impl Drop for Canary
    {
        fn drop(&mut self) {
            unsafe {
                match *self {
                    Canary(c) => {
                        (*c).fetch_add(1, atomic::SeqCst);
                    }
                }
            }
        }
    }

    #[test]
    fn manually_share_arc() {
        let v = vec!(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
        let arc_v = Arc::new(v);

        let (tx, rx) = channel();

        task::spawn(proc() {
            let arc_v: Arc<Vec<int>> = rx.recv();
            assert_eq!((*arc_v)[3], 4);
        });

        tx.send(arc_v.clone());

        assert_eq!((*arc_v)[2], 3);
        assert_eq!((*arc_v)[4], 5);
    }

    #[test]
    fn test_cowarc_clone_make_unique() {
        let mut cow0 = Arc::new(75u);
        let mut cow1 = cow0.clone();
        let mut cow2 = cow1.clone();

        assert!(75 == *cow0.make_unique());
        assert!(75 == *cow1.make_unique());
        assert!(75 == *cow2.make_unique());

        *cow0.make_unique() += 1;
        *cow1.make_unique() += 2;
        *cow2.make_unique() += 3;

        assert!(76 == *cow0);
        assert!(77 == *cow1);
        assert!(78 == *cow2);

        // none should point to the same backing memory
        assert!(*cow0 != *cow1);
        assert!(*cow0 != *cow2);
        assert!(*cow1 != *cow2);
    }

    #[test]
    fn test_cowarc_clone_unique2() {
        let mut cow0 = Arc::new(75u);
        let cow1 = cow0.clone();
        let cow2 = cow1.clone();

        assert!(75 == *cow0);
        assert!(75 == *cow1);
        assert!(75 == *cow2);

        *cow0.make_unique() += 1;

        assert!(76 == *cow0);
        assert!(75 == *cow1);
        assert!(75 == *cow2);

        // cow1 and cow2 should share the same contents
        // cow0 should have a unique reference
        assert!(*cow0 != *cow1);
        assert!(*cow0 != *cow2);
        assert!(*cow1 == *cow2);
    }

    #[test]
    fn test_cowarc_clone_weak() {
        let mut cow0 = Arc::new(75u);
        let cow1_weak = cow0.downgrade();

        assert!(75 == *cow0);
        assert!(75 == *cow1_weak.upgrade().unwrap());

        *cow0.make_unique() += 1;

        assert!(76 == *cow0);
        assert!(cow1_weak.upgrade().is_none());
    }

    #[test]
    fn test_live() {
        let x = Arc::new(5i);
        let y = x.downgrade();
        assert!(y.upgrade().is_some());
    }

    #[test]
    fn test_dead() {
        let x = Arc::new(5i);
        let y = x.downgrade();
        drop(x);
        assert!(y.upgrade().is_none());
    }

    #[test]
    fn weak_self_cyclic() {
        struct Cycle {
            x: Mutex<Option<Weak<Cycle>>>
        }

        let a = Arc::new(Cycle { x: Mutex::new(None) });
        let b = a.clone().downgrade();
        *a.x.lock() = Some(b);

        // hopefully we don't double-free (or leak)...
    }

    #[test]
    fn drop_arc() {
        let mut canary = atomic::AtomicUint::new(0);
        let x = Arc::new(Canary(&mut canary as *mut atomic::AtomicUint));
        drop(x);
        assert!(canary.load(atomic::Acquire) == 1);
    }

    #[test]
    fn drop_arc_weak() {
        let mut canary = atomic::AtomicUint::new(0);
        let arc = Arc::new(Canary(&mut canary as *mut atomic::AtomicUint));
        let arc_weak = arc.downgrade();
        assert!(canary.load(atomic::Acquire) == 0);
        drop(arc);
        assert!(canary.load(atomic::Acquire) == 1);
        drop(arc_weak);
    }
}
<commit_msg>Implement Show for `Arc<T>`<commit_after>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![stable]

//! Concurrency-enabled mechanisms for sharing mutable and/or immutable state
//! between tasks.

use core::atomic;
use core::clone::Clone;
use core::fmt::{mod, Show};
use core::kinds::{Sync, Send};
use core::mem::{min_align_of, size_of, drop};
use core::mem;
use core::ops::{Drop, Deref};
use core::option::{Some, None, Option};
use core::ptr;
use core::ptr::RawPtr;
use heap::deallocate;

/// An atomically reference counted wrapper for shared state.
///
/// # Example
///
/// In this example, a large vector of floats is shared between several tasks.
/// With simple pipes, without `Arc`, a copy would have to be made for each
/// task.
///
/// ```rust
/// use std::sync::Arc;
///
/// fn main() {
///     let numbers = Vec::from_fn(100, |i| i as f32);
///     let shared_numbers = Arc::new(numbers);
///
///     for _ in range(0u, 10) {
///         let child_numbers = shared_numbers.clone();
///
///         spawn(proc() {
///             let local_numbers = child_numbers.as_slice();
///
///             // Work with the local numbers
///         });
///     }
/// }
/// ```
#[unsafe_no_drop_flag]
#[stable]
pub struct Arc<T> {
    // FIXME #12808: strange name to try to avoid interfering with
    // field accesses of the contained type via Deref
    _ptr: *mut ArcInner<T>,
}

/// A weak pointer to an `Arc`.
///
/// Weak pointers will not keep the data inside of the `Arc` alive, and can be
/// used to break cycles between `Arc` pointers.
#[unsafe_no_drop_flag]
#[experimental = "Weak pointers may not belong in this module."]
pub struct Weak<T> {
    // FIXME #12808: strange name to try to avoid interfering with
    // field accesses of the contained type via Deref
    _ptr: *mut ArcInner<T>,
}

struct ArcInner<T> {
    strong: atomic::AtomicUint,
    weak: atomic::AtomicUint,
    data: T,
}

impl<T: Sync + Send> Arc<T> {
    /// Creates an atomically reference counted wrapper.
    #[inline]
    #[stable]
    pub fn new(data: T) -> Arc<T> {
        // Start the weak pointer count as 1 which is the weak pointer that's
        // held by all the strong pointers (kinda), see std/rc.rs for more info
        let x = box ArcInner {
            strong: atomic::AtomicUint::new(1),
            weak: atomic::AtomicUint::new(1),
            data: data,
        };
        Arc { _ptr: unsafe { mem::transmute(x) } }
    }

    #[inline]
    fn inner(&self) -> &ArcInner<T> {
        // This unsafety is ok because while this arc is alive we're guaranteed
        // that the inner pointer is valid. Furthermore, we know that the
        // `ArcInner` structure itself is `Sync` because the inner data is
        // `Sync` as well, so we're ok loaning out an immutable pointer to
        // these contents.
        unsafe { &*self._ptr }
    }

    /// Downgrades a strong pointer to a weak pointer.
    ///
    /// Weak pointers will not keep the data alive. Once all strong references
    /// to the underlying data have been dropped, the data itself will be
    /// destroyed.
    #[experimental = "Weak pointers may not belong in this module."]
    pub fn downgrade(&self) -> Weak<T> {
        // See the clone() impl for why this is relaxed
        self.inner().weak.fetch_add(1, atomic::Relaxed);
        Weak { _ptr: self._ptr }
    }
}

#[unstable = "waiting on stability of Clone"]
impl<T: Sync + Send> Clone for Arc<T> {
    /// Duplicate an atomically reference counted wrapper.
    ///
    /// The resulting two `Arc` objects will point to the same underlying data
    /// object. However, one of the `Arc` objects can be sent to another task,
    /// allowing them to share the underlying data.
    #[inline]
    fn clone(&self) -> Arc<T> {
        // Using a relaxed ordering is alright here, as knowledge of the
        // original reference prevents other threads from erroneously deleting
        // the object.
        //
        // As explained in the [Boost documentation][1], Increasing the
        // reference counter can always be done with memory_order_relaxed: New
        // references to an object can only be formed from an existing
        // reference, and passing an existing reference from one thread to
        // another must already provide any required synchronization.
        //
        // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html)
        self.inner().strong.fetch_add(1, atomic::Relaxed);
        Arc { _ptr: self._ptr }
    }
}

#[experimental = "Deref is experimental."]
impl<T: Send + Sync> Deref<T> for Arc<T> {
    #[inline]
    fn deref(&self) -> &T {
        &self.inner().data
    }
}

impl<T: Send + Sync + Show> Show for Arc<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        (**self).fmt(f)
    }
}

impl<T: Send + Sync + Clone> Arc<T> {
    /// Acquires a mutable pointer to the inner contents by guaranteeing that
    /// the reference count is one (no sharing is possible).
    ///
    /// This is also referred to as a copy-on-write operation because the inner
    /// data is cloned if the reference count is greater than one.
    #[inline]
    #[experimental]
    pub fn make_unique(&mut self) -> &mut T {
        // Note that we hold a strong reference, which also counts as
        // a weak reference, so we only clone if there is an
        // additional reference of either kind.
        if self.inner().strong.load(atomic::SeqCst) != 1 ||
           self.inner().weak.load(atomic::SeqCst) != 1 {
            *self = Arc::new(self.deref().clone())
        }
        // This unsafety is ok because we're guaranteed that the pointer
        // returned is the *only* pointer that will ever be returned to T. Our
        // reference count is guaranteed to be 1 at this point, and we required
        // the Arc itself to be `mut`, so we're returning the only possible
        // reference to the inner data.
        let inner = unsafe { &mut *self._ptr };
        &mut inner.data
    }
}

#[unsafe_destructor]
#[experimental = "waiting on stability of Drop"]
impl<T: Sync + Send> Drop for Arc<T> {
    fn drop(&mut self) {
        // This structure has #[unsafe_no_drop_flag], so this drop glue may run
        // more than once (but it is guaranteed to be zeroed after the first if
        // it's run more than once)
        if self._ptr.is_null() { return }

        // Because `fetch_sub` is already atomic, we do not need to synchronize
        // with other threads unless we are going to delete the object. This
        // same logic applies to the below `fetch_sub` to the `weak` count.
        if self.inner().strong.fetch_sub(1, atomic::Release) != 1 { return }

        // This fence is needed to prevent reordering of use of the data and
        // deletion of the data. Because it is marked `Release`, the
        // decreasing of the reference count synchronizes with this `Acquire`
        // fence. This means that use of the data happens before decreasing
        // the reference count, which happens before this fence, which
        // happens before the deletion of the data.
        //
        // As explained in the [Boost documentation][1],
        //
        // It is important to enforce any possible access to the object in
        // one thread (through an existing reference) to *happen before*
        // deleting the object in a different thread. This is achieved by a
        // "release" operation after dropping a reference (any access to the
        // object through this reference must obviously happened before),
        // and an "acquire" operation before deleting the object.
        //
        // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html)
        atomic::fence(atomic::Acquire);

        // Destroy the data at this time, even though we may not free the box
        // allocation itself (there may still be weak pointers lying around).
        unsafe { drop(ptr::read(&self.inner().data)); }

        if self.inner().weak.fetch_sub(1, atomic::Release) == 1 {
            atomic::fence(atomic::Acquire);
            unsafe { deallocate(self._ptr as *mut u8, size_of::<ArcInner<T>>(),
                                min_align_of::<ArcInner<T>>()) }
        }
    }
}

#[experimental = "Weak pointers may not belong in this module."]
impl<T: Sync + Send> Weak<T> {
    /// Attempts to upgrade this weak reference to a strong reference.
    ///
    /// This method will fail to upgrade this reference if the strong reference
    /// count has already reached 0, but if there are still other active strong
    /// references this function will return a new strong reference to the data.
    pub fn upgrade(&self) -> Option<Arc<T>> {
        // We use a CAS loop to increment the strong count instead of a
        // fetch_add because once the count hits 0 is must never be above 0.
        let inner = self.inner();
        loop {
            let n = inner.strong.load(atomic::SeqCst);
            if n == 0 { return None }
            let old = inner.strong.compare_and_swap(n, n + 1, atomic::SeqCst);
            if old == n { return Some(Arc { _ptr: self._ptr }) }
        }
    }

    #[inline]
    fn inner(&self) -> &ArcInner<T> {
        // See comments above for why this is "safe"
        unsafe { &*self._ptr }
    }
}

#[experimental = "Weak pointers may not belong in this module."]
impl<T: Sync + Send> Clone for Weak<T> {
    #[inline]
    fn clone(&self) -> Weak<T> {
        // See comments in Arc::clone() for why this is relaxed
        self.inner().weak.fetch_add(1, atomic::Relaxed);
        Weak { _ptr: self._ptr }
    }
}

#[unsafe_destructor]
#[experimental = "Weak pointers may not belong in this module."]
impl<T: Sync + Send> Drop for Weak<T> {
    fn drop(&mut self) {
        // see comments above for why this check is here
        if self._ptr.is_null() { return }

        // If we find out that we were the last weak pointer, then its time to
        // deallocate the data entirely. See the discussion in Arc::drop() about
        // the memory orderings
        if self.inner().weak.fetch_sub(1, atomic::Release) == 1 {
            atomic::fence(atomic::Acquire);
            unsafe { deallocate(self._ptr as *mut u8, size_of::<ArcInner<T>>(),
                                min_align_of::<ArcInner<T>>()) }
        }
    }
}

#[cfg(test)]
#[allow(experimental)]
mod tests {
    use std::clone::Clone;
    use std::comm::channel;
    use std::mem::drop;
    use std::ops::Drop;
    use std::option::{Option, Some, None};
    use std::str::Str;
    use std::sync::atomic;
    use std::task;
    use std::vec::Vec;
    use super::{Arc, Weak};
    use std::sync::Mutex;

    struct Canary(*mut atomic::AtomicUint);

    impl Drop for Canary
    {
        fn drop(&mut self) {
            unsafe {
                match *self {
                    Canary(c) => {
                        (*c).fetch_add(1, atomic::SeqCst);
                    }
                }
            }
        }
    }

    #[test]
    fn manually_share_arc() {
        let v = vec!(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
        let arc_v = Arc::new(v);

        let (tx, rx) = channel();

        task::spawn(proc() {
            let arc_v: Arc<Vec<int>> = rx.recv();
            assert_eq!((*arc_v)[3], 4);
        });

        tx.send(arc_v.clone());

        assert_eq!((*arc_v)[2], 3);
        assert_eq!((*arc_v)[4], 5);
    }

    #[test]
    fn test_cowarc_clone_make_unique() {
        let mut cow0 = Arc::new(75u);
        let mut cow1 = cow0.clone();
        let mut cow2 = cow1.clone();

        assert!(75 == *cow0.make_unique());
        assert!(75 == *cow1.make_unique());
        assert!(75 == *cow2.make_unique());

        *cow0.make_unique() += 1;
        *cow1.make_unique() += 2;
        *cow2.make_unique() += 3;

        assert!(76 == *cow0);
        assert!(77 == *cow1);
        assert!(78 == *cow2);

        // none should point to the same backing memory
        assert!(*cow0 != *cow1);
        assert!(*cow0 != *cow2);
        assert!(*cow1 != *cow2);
    }

    #[test]
    fn test_cowarc_clone_unique2() {
        let mut cow0 = Arc::new(75u);
        let cow1 = cow0.clone();
        let cow2 = cow1.clone();

        assert!(75 == *cow0);
        assert!(75 == *cow1);
        assert!(75 == *cow2);

        *cow0.make_unique() += 1;

        assert!(76 == *cow0);
        assert!(75 == *cow1);
        assert!(75 == *cow2);

        // cow1 and cow2 should share the same contents
        // cow0 should have a unique reference
        assert!(*cow0 != *cow1);
        assert!(*cow0 != *cow2);
        assert!(*cow1 == *cow2);
    }

    #[test]
    fn test_cowarc_clone_weak() {
        let mut cow0 = Arc::new(75u);
        let cow1_weak = cow0.downgrade();

        assert!(75 == *cow0);
        assert!(75 == *cow1_weak.upgrade().unwrap());

        *cow0.make_unique() += 1;

        assert!(76 == *cow0);
        assert!(cow1_weak.upgrade().is_none());
    }

    #[test]
    fn test_live() {
        let x = Arc::new(5i);
        let y = x.downgrade();
        assert!(y.upgrade().is_some());
    }

    #[test]
    fn test_dead() {
        let x = Arc::new(5i);
        let y = x.downgrade();
        drop(x);
        assert!(y.upgrade().is_none());
    }

    #[test]
    fn weak_self_cyclic() {
        struct Cycle {
            x: Mutex<Option<Weak<Cycle>>>
        }

        let a = Arc::new(Cycle { x: Mutex::new(None) });
        let b = a.clone().downgrade();
        *a.x.lock() = Some(b);

        // hopefully we don't double-free (or leak)...
    }

    #[test]
    fn drop_arc() {
        let mut canary = atomic::AtomicUint::new(0);
        let x = Arc::new(Canary(&mut canary as *mut atomic::AtomicUint));
        drop(x);
        assert!(canary.load(atomic::Acquire) == 1);
    }

    #[test]
    fn drop_arc_weak() {
        let mut canary = atomic::AtomicUint::new(0);
        let arc = Arc::new(Canary(&mut canary as *mut atomic::AtomicUint));
        let arc_weak = arc.downgrade();
        assert!(canary.load(atomic::Acquire) == 0);
        drop(arc);
        assert!(canary.load(atomic::Acquire) == 1);
        drop(arc_weak);
    }

    #[test]
    fn show_arc() {
        let a = Arc::new(5u32);
        assert!(format!("{}", a).as_slice() == "5")
    }
}
<|endoftext|>