<commit_before>//! Codegen vtables and vtable accesses.
//!
//! See `rustc_codegen_ssa/src/meth.rs` for reference.
// FIXME dedup this logic between miri, cg_llvm and cg_clif

use crate::prelude::*;

const DROP_FN_INDEX: usize = 0;
const SIZE_INDEX: usize = 1;
const ALIGN_INDEX: usize = 2;

fn vtable_memflags() -> MemFlags {
    let mut flags = MemFlags::trusted(); // A vtable access is always aligned and will never trap.
    flags.set_readonly(); // A vtable is always read-only.
    flags
}

pub(crate) fn drop_fn_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (DROP_FN_INDEX * usize_size) as i32,
    )
}

pub(crate) fn size_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (SIZE_INDEX * usize_size) as i32,
    )
}

pub(crate) fn min_align_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (ALIGN_INDEX * usize_size) as i32,
    )
}

pub(crate) fn get_ptr_and_method_ref<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    arg: CValue<'tcx>,
    idx: usize,
) -> (Value, Value) {
    let (ptr, vtable) = if let Abi::ScalarPair(_, _) = arg.layout().abi {
        arg.load_scalar_pair(fx)
    } else {
        let (ptr, vtable) = arg.try_to_ptr().unwrap();
        (ptr.get_addr(fx), vtable.unwrap())
    };

    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes();
    let func_ref = fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        ((idx + 3) * usize_size as usize) as i32,
    );
    (ptr, func_ref)
}

pub(crate) fn get_vtable<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    layout: TyAndLayout<'tcx>,
    trait_ref: Option<ty::PolyExistentialTraitRef<'tcx>>,
) -> Value {
    let data_id = if let Some(data_id) = fx.vtables.get(&(layout.ty, trait_ref)) {
        *data_id
    } else {
        let data_id = build_vtable(fx, layout, trait_ref);
        fx.vtables.insert((layout.ty, trait_ref), data_id);
        data_id
    };

    let local_data_id = fx.module.declare_data_in_func(data_id, &mut fx.bcx.func);
    fx.bcx.ins().global_value(fx.pointer_type, local_data_id)
}

fn build_vtable<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    layout: TyAndLayout<'tcx>,
    trait_ref: Option<ty::PolyExistentialTraitRef<'tcx>>,
) -> DataId {
    let tcx = fx.tcx;
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;

    let drop_in_place_fn = import_function(
        tcx,
        fx.module,
        Instance::resolve_drop_in_place(tcx, layout.ty).polymorphize(fx.tcx),
    );

    let mut components: Vec<_> = vec![Some(drop_in_place_fn), None, None];

    let methods_root;
    let methods = if let Some(trait_ref) = trait_ref {
        methods_root = tcx.vtable_methods(trait_ref.with_self_ty(tcx, layout.ty));
        methods_root.iter()
    } else {
        (&[]).iter()
    };
    let methods = methods.cloned().map(|opt_mth| {
        opt_mth.map(|(def_id, substs)| {
            import_function(
                tcx,
                fx.module,
                Instance::resolve_for_vtable(tcx, ParamEnv::reveal_all(), def_id, substs)
                    .unwrap()
                    .polymorphize(fx.tcx),
            )
        })
    });
    components.extend(methods);

    let mut data_ctx = DataContext::new();
    let mut data = ::std::iter::repeat(0u8)
        .take(components.len() * usize_size)
        .collect::<Vec<u8>>()
        .into_boxed_slice();

    write_usize(fx.tcx, &mut data, SIZE_INDEX, layout.size.bytes());
    write_usize(fx.tcx, &mut data, ALIGN_INDEX, layout.align.abi.bytes());
    data_ctx.define(data);

    for (i, component) in components.into_iter().enumerate() {
        if let Some(func_id) = component {
            let func_ref = fx.module.declare_func_in_data(func_id, &mut data_ctx);
            data_ctx.write_function_addr((i * usize_size) as u32, func_ref);
        }
    }

    data_ctx.set_align(fx.tcx.data_layout.pointer_align.pref.bytes());

    let data_id = fx.module.declare_anonymous_data(false, false).unwrap();

    fx.module.define_data(data_id, &data_ctx).unwrap();

    data_id
}

fn write_usize(tcx: TyCtxt<'_>, buf: &mut [u8], idx: usize, num: u64) {
    let pointer_size =
        tcx.layout_of(ParamEnv::reveal_all().and(tcx.types.usize)).unwrap().size.bytes() as usize;
    let target = &mut buf[idx * pointer_size..(idx + 1) * pointer_size];

    match tcx.data_layout.endian {
        rustc_target::abi::Endian::Little => match pointer_size {
            4 => target.copy_from_slice(&(num as u32).to_le_bytes()),
            8 => target.copy_from_slice(&(num as u64).to_le_bytes()),
            _ => todo!("pointer size {} is not yet supported", pointer_size),
        },
        rustc_target::abi::Endian::Big => match pointer_size {
            4 => target.copy_from_slice(&(num as u32).to_be_bytes()),
            8 => target.copy_from_slice(&(num as u64).to_be_bytes()),
            _ => todo!("pointer size {} is not yet supported", pointer_size),
        },
    }
}
<commit_msg>Refactor to make interpreter and codegen backend neutral to vtable internal representation.<commit_after>//! Codegen vtables and vtable accesses.
//!
//! See `rustc_codegen_ssa/src/meth.rs` for reference.
// FIXME dedup this logic between miri, cg_llvm and cg_clif

use crate::prelude::*;
use ty::VtblEntry;

fn vtable_memflags() -> MemFlags {
    let mut flags = MemFlags::trusted(); // A vtable access is always aligned and will never trap.
    flags.set_readonly(); // A vtable is always read-only.
    flags
}

pub(crate) fn drop_fn_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (ty::COMMON_VTABLE_ENTRIES_DROPINPLACE * usize_size) as i32,
    )
}

pub(crate) fn size_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (ty::COMMON_VTABLE_ENTRIES_SIZE * usize_size) as i32,
    )
}

pub(crate) fn min_align_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (ty::COMMON_VTABLE_ENTRIES_SIZE * usize_size) as i32,
    )
}

pub(crate) fn get_ptr_and_method_ref<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    arg: CValue<'tcx>,
    idx: usize,
) -> (Value, Value) {
    let (ptr, vtable) = if let Abi::ScalarPair(_, _) = arg.layout().abi {
        arg.load_scalar_pair(fx)
    } else {
        let (ptr, vtable) = arg.try_to_ptr().unwrap();
        (ptr.get_addr(fx), vtable.unwrap())
    };

    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes();
    let func_ref = fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (idx * usize_size as usize) as i32,
    );
    (ptr, func_ref)
}

pub(crate) fn get_vtable<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    layout: TyAndLayout<'tcx>,
    trait_ref: Option<ty::PolyExistentialTraitRef<'tcx>>,
) -> Value {
    let data_id = if let Some(data_id) = fx.vtables.get(&(layout.ty, trait_ref)) {
        *data_id
    } else {
        let data_id = build_vtable(fx, layout, trait_ref);
        fx.vtables.insert((layout.ty, trait_ref), data_id);
        data_id
    };

    let local_data_id = fx.module.declare_data_in_func(data_id, &mut fx.bcx.func);
    fx.bcx.ins().global_value(fx.pointer_type, local_data_id)
}

fn build_vtable<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    layout: TyAndLayout<'tcx>,
    trait_ref: Option<ty::PolyExistentialTraitRef<'tcx>>,
) -> DataId {
    let tcx = fx.tcx;
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;

    let drop_in_place_fn = import_function(
        tcx,
        fx.module,
        Instance::resolve_drop_in_place(tcx, layout.ty).polymorphize(fx.tcx),
    );

    let vtable_entries = if let Some(trait_ref) = trait_ref {
        tcx.vtable_entries(trait_ref.with_self_ty(tcx, layout.ty))
    } else {
        ty::COMMON_VTABLE_ENTRIES
    };

    let mut data_ctx = DataContext::new();
    let mut data = ::std::iter::repeat(0u8)
        .take(vtable_entries.len() * usize_size)
        .collect::<Vec<u8>>()
        .into_boxed_slice();

    for (idx, entry) in vtable_entries.iter().enumerate() {
        match entry {
            VtblEntry::MetadataSize => {
                write_usize(fx.tcx, &mut data, idx, layout.size.bytes());
            }
            VtblEntry::MetadataAlign => {
                write_usize(fx.tcx, &mut data, idx, layout.align.abi.bytes());
            }
            VtblEntry::MetadataDropInPlace | VtblEntry::Vacant | VtblEntry::Method(_, _) => {}
        }
    }
    data_ctx.define(data);

    for (idx, entry) in vtable_entries.iter().enumerate() {
        match entry {
            VtblEntry::MetadataDropInPlace => {
                let func_ref = fx.module.declare_func_in_data(drop_in_place_fn, &mut data_ctx);
                data_ctx.write_function_addr((idx * usize_size) as u32, func_ref);
            }
            VtblEntry::Method(def_id, substs) => {
                let func_id = import_function(
                    tcx,
                    fx.module,
                    Instance::resolve_for_vtable(tcx, ParamEnv::reveal_all(), *def_id, substs)
                        .unwrap()
                        .polymorphize(fx.tcx),
                );
                let func_ref = fx.module.declare_func_in_data(func_id, &mut data_ctx);
                data_ctx.write_function_addr((idx * usize_size) as u32, func_ref);
            }
            VtblEntry::MetadataSize | VtblEntry::MetadataAlign | VtblEntry::Vacant => {}
        }
    }

    data_ctx.set_align(fx.tcx.data_layout.pointer_align.pref.bytes());

    let data_id = fx.module.declare_anonymous_data(false, false).unwrap();

    fx.module.define_data(data_id, &data_ctx).unwrap();

    data_id
}

fn write_usize(tcx: TyCtxt<'_>, buf: &mut [u8], idx: usize, num: u64) {
    let pointer_size =
        tcx.layout_of(ParamEnv::reveal_all().and(tcx.types.usize)).unwrap().size.bytes() as usize;
    let target = &mut buf[idx * pointer_size..(idx + 1) * pointer_size];

    match tcx.data_layout.endian {
        rustc_target::abi::Endian::Little => match pointer_size {
            4 => target.copy_from_slice(&(num as u32).to_le_bytes()),
            8 => target.copy_from_slice(&(num as u64).to_le_bytes()),
            _ => todo!("pointer size {} is not yet supported", pointer_size),
        },
        rustc_target::abi::Endian::Big => match pointer_size {
            4 => target.copy_from_slice(&(num as u32).to_be_bytes()),
            8 => target.copy_from_slice(&(num as u64).to_be_bytes()),
            _ => todo!("pointer size {} is not yet supported", pointer_size),
        },
    }
}
<|endoftext|>


<commit_before>//! Utility functions to write nifti images.

use std::fs::File;
use std::io::BufWriter;
use std::mem;
use std::ops::{Div, Sub};

use byteorder::{LittleEndian, WriteBytesExt};
use flate2::write::GzEncoder;
use flate2::Compression;
use ndarray::{Array, ArrayBase, Axis, Data, Dimension, RemoveAxis, ScalarOperand};
use num_traits::FromPrimitive;
use safe_transmute::{guarded_transmute_to_bytes_pod_many, PodTransmutable};

use {util::is_gz_file, volume::element::DataElement, NiftiHeader, Result};

// TODO make this configurable. The Nifti standard does not specify a specific field for endianness,
// but it is encoded in `dim[0]`. "if dim[0] is outside range 1..7, then swap".
type B = LittleEndian;

/// Write a nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields. The voxels
/// intensity will be subtracted by `scl_slope` and divided by `scl_inter`. If `reference` is not
/// given, a default `NiftiHeader` will be built and written.
///
/// In all cases, the `dim`, `datatype` and `bitpix` fields will depend only on `data`, not on the
/// header. In other words, the `datatype` defined in `reference` will be ignored.
pub fn write_nifti<P, A, S, D>(
    path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<::std::path::Path>,
    S: Data<Elem = A>,
    A: Copy,
    A: DataElement,
    A: Div<Output = A>,
    A: FromPrimitive,
    A: PodTransmutable,
    A: ScalarOperand,
    A: Sub<Output = A>,
    D: Dimension + RemoveAxis,
{
    let mut dim = [1; 8];
    dim[0] = data.ndim() as u16;
    for (i, s) in data.shape().iter().enumerate() {
        dim[i + 1] = *s as u16;
    }

    // If no reference header is given, use the default.
    let reference = match reference {
        Some(r) => r.clone(),
        None => NiftiHeader::default(),
    };
    let header = NiftiHeader {
        dim,
        datatype: A::DATA_TYPE as i16,
        bitpix: (mem::size_of::<A>() * 8) as i16,
        // All other fields are copied from reference header
        ..reference
    };

    let f = File::create(&path).expect("Can't create new nifti file");
    let mut writer = BufWriter::new(f);
    if is_gz_file(&path) {
        let mut e = GzEncoder::new(writer, Compression::default());
        write_header(&mut e, &header)?;
        write_data(&mut e, header, data)?;
        let _ = e.finish()?; // Must use result
    } else {
        write_header(&mut writer, &header)?;
        write_data(&mut writer, header, data)?;
    }
    Ok(())
}

fn write_header<W: WriteBytesExt>(writer: &mut W, header: &NiftiHeader) -> Result<()> {
    writer.write_i32::<B>(header.sizeof_hdr)?;
    writer.write_all(&header.data_type)?;
    writer.write_all(&header.db_name)?;
    writer.write_i32::<B>(header.extents)?;
    writer.write_i16::<B>(header.session_error)?;
    writer.write_u8(header.regular)?;
    writer.write_u8(header.dim_info)?;
    for s in &header.dim {
        writer.write_u16::<B>(*s)?;
    }
    writer.write_f32::<B>(header.intent_p1)?;
    writer.write_f32::<B>(header.intent_p2)?;
    writer.write_f32::<B>(header.intent_p3)?;
    writer.write_i16::<B>(header.intent_code)?;
    writer.write_i16::<B>(header.datatype)?;
    writer.write_i16::<B>(header.bitpix)?;
    writer.write_i16::<B>(header.slice_start)?;
    for f in &header.pixdim {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_f32::<B>(header.vox_offset)?;
    writer.write_f32::<B>(header.scl_slope)?;
    writer.write_f32::<B>(header.scl_inter)?;
    writer.write_i16::<B>(header.slice_end)?;
    writer.write_u8(header.slice_code)?;
    writer.write_u8(header.xyzt_units)?;
    writer.write_f32::<B>(header.cal_max)?;
    writer.write_f32::<B>(header.cal_min)?;
    writer.write_f32::<B>(header.slice_duration)?;
    writer.write_f32::<B>(header.toffset)?;
    writer.write_i32::<B>(header.glmax)?;
    writer.write_i32::<B>(header.glmin)?;

    writer.write_all(&header.descrip)?;
    writer.write_all(&header.aux_file)?;
    writer.write_i16::<B>(header.qform_code)?;
    writer.write_i16::<B>(header.sform_code)?;
    for f in &[
        header.quatern_b,
        header.quatern_c,
        header.quatern_d,
        header.quatern_x,
        header.quatern_y,
        header.quatern_z,
    ] {
        writer.write_f32::<B>(*f)?;
    }
    for f in header
        .srow_x
        .iter()
        .chain(&header.srow_y)
        .chain(&header.srow_z)
    {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_all(&header.intent_name)?;
    writer.write_all(&header.magic)?;

    // Empty 4 bytes after the header
    // TODO Support writing extension data.
    writer.write_u32::<B>(0)?;

    Ok(())
}

/// Write the data in 'f' order.
///
/// Like NiBabel, we iterate by "slice" to improve speed and use less memory.
fn write_data<A, S, D, W>(writer: &mut W, header: NiftiHeader, data: &ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Copy,
    A: DataElement,
    A: Div<Output = A>,
    A: FromPrimitive,
    A: PodTransmutable,
    A: ScalarOperand,
    A: Sub<Output = A>,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    // Need the transpose for fortran ordering used in nifti file format.
    let data = data.t();

    let mut write_slice = |data: Array<A, D::Smaller>| -> Result<()> {
        let len = data.len();
        let arr_data = data.into_shape(len).unwrap();
        let slice = arr_data.as_slice().unwrap();
        writer.write_all(guarded_transmute_to_bytes_pod_many(slice))?;
        Ok(())
    };

    // `1.0x + 0.0` would give the same results, but we avoid a lot of divisions
    let slope = if header.scl_slope == 0.0 {
        1.0
    } else {
        header.scl_slope
    };
    if slope != 1.0 || header.scl_inter != 0.0 {
        let slope = A::from_f32(slope).unwrap();
        let inter = A::from_f32(header.scl_inter).unwrap();
        for arr_data in data.axis_iter(Axis(0)) {
            write_slice(arr_data.sub(inter).div(slope))?;
        }
    } else {
        for arr_data in data.axis_iter(Axis(0)) {
            // We need to own the data because of the into_shape() for `c` ordering.
            write_slice(arr_data.to_owned())?;
        }
    }
    Ok(())
}

#[cfg(test)]
pub mod tests {
    extern crate tempfile;

    use super::*;
    use std::ops::{Add, Mul};
    use std::path::PathBuf;

    use self::tempfile::tempdir;
    use ndarray::{Array, Array2, Ix2, IxDyn, ShapeBuilder};

    use {header::MAGIC_CODE_NIP1, object::NiftiObject, InMemNiftiObject, IntoNdArray};

    fn get_random_path(ext: &str) -> PathBuf {
        let dir = tempdir().unwrap();
        let path = if ext == "" {
            dir.into_path()
        } else {
            dir.into_path().join(ext)
        };
        PathBuf::from(path.to_str().unwrap())
    }

    pub fn generate_nifti_header(
        dim: [u16; 8],
        scl_slope: f32,
        scl_inter: f32,
        datatype: i16,
    ) -> NiftiHeader {
        NiftiHeader {
            dim,
            datatype,
            bitpix: (mem::size_of::<f32>() * 8) as i16,
            magic: *MAGIC_CODE_NIP1,
            scl_slope,
            scl_inter,
            ..NiftiHeader::default()
        }
    }

    fn read_2d_image<P: AsRef<::std::path::Path>>(path: P) -> Array2<f32> {
        let nifti_object = InMemNiftiObject::from_file(path).expect("Nifti file is unreadable.");
        let volume = nifti_object.into_volume();
        let dyn_data = volume.into_ndarray().unwrap();
        dyn_data.into_dimensionality::<Ix2>().unwrap()
    }

    fn test_write_read(arr: &Array<f32, IxDyn>, path: &str) {
        let path = get_random_path(path);
        let mut dim = [1; 8];
        dim[0] = arr.ndim() as u16;
        for (i, s) in arr.shape().iter().enumerate() {
            dim[i + 1] = *s as u16;
        }
        let header = generate_nifti_header(dim, 1.0, 0.0, 16);
        write_nifti(&path, &arr, Some(&header)).unwrap();

        let read_nifti = read_2d_image(path);
        assert!(read_nifti.all_close(&arr, 1e-10));
    }

    fn f_order_array() -> Array<f32, IxDyn> {
        let dim = vec![4, 4];
        let vec = (0..16).map(|x| x as f32).collect();
        Array::from_shape_vec(IxDyn(&dim).f(), vec).unwrap()
    }

    fn c_order_array() -> Array<f32, IxDyn> {
        let dim = vec![4, 4];
        let vec = (0..16).map(|x| x as f32).collect();
        Array::from_shape_vec(IxDyn(&dim), vec).unwrap()
    }

    #[test]
    fn test_fortran_writing() {
        // Test .nii
        let arr = f_order_array();
        test_write_read(&arr, "test.nii");
        let mut arr = f_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii");

        // Test .nii.gz
        let arr = f_order_array();
        test_write_read(&arr, "test.nii.gz");
        let mut arr = f_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii.gz");
    }

    #[test]
    fn test_c_writing() {
        // Test .nii
        let arr = c_order_array();
        test_write_read(&arr, "test.nii");

        let mut arr = c_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii");

        // Test .nii.gz
        let arr = c_order_array();
        test_write_read(&arr, "test.nii.gz");

        let mut arr = c_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii.gz");
    }

    #[test]
    fn test_header_slope_inter() {
        let arr = f_order_array();
        let slope = 2.2;
        let inter = 101.1;

        let path = get_random_path("test_slope_inter.nii");
        let mut dim = [1; 8];
        dim[0] = arr.ndim() as u16;
        for (i, s) in arr.shape().iter().enumerate() {
            dim[i + 1] = *s as u16;
        }
        let header = generate_nifti_header(dim, slope, inter, 16);
        let transformed_data = arr.mul(slope).add(inter);
        write_nifti(&path, &transformed_data, Some(&header)).unwrap();

        let read_nifti = read_2d_image(path);
        assert!(read_nifti.all_close(&transformed_data, 1e-10));
    }
}
<commit_msg>Various cleaning in writer.rs<commit_after>//! Utility functions to write nifti images.

use std::fs::File;
use std::io::BufWriter;
use std::mem;
use std::ops::{Div, Sub};
use std::path::Path;

use byteorder::{LittleEndian, WriteBytesExt};
use flate2::write::GzEncoder;
use flate2::Compression;
use ndarray::{Array, ArrayBase, Axis, Data, Dimension, RemoveAxis, ScalarOperand};
use num_traits::FromPrimitive;
use safe_transmute::{guarded_transmute_to_bytes_pod_many, PodTransmutable};

use {util::is_gz_file, volume::element::DataElement, NiftiHeader, Result};

// TODO make this configurable. The Nifti standard does not specify a specific field for endianness,
// but it is encoded in `dim[0]`. "if dim[0] is outside range 1..7, then swap".
type B = LittleEndian;

/// Write a nifti file (.nii or .nii.gz) in Little Endian.
///
/// If a `reference` is given, it will be used to fill most of the header's fields. The voxels
/// intensity will be subtracted by `scl_slope` and divided by `scl_inter`. If `reference` is not
/// given, a default `NiftiHeader` will be built and written.
///
/// In all cases, the `dim`, `datatype` and `bitpix` fields will depend only on `data`, not on the
/// header. In other words, the `datatype` defined in `reference` will be ignored.
pub fn write_nifti<P, A, S, D>(
    path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = A>,
    A: Copy,
    A: DataElement,
    A: Div<Output = A>,
    A: FromPrimitive,
    A: PodTransmutable,
    A: ScalarOperand,
    A: Sub<Output = A>,
    D: Dimension + RemoveAxis,
{
    let mut dim = [1; 8];
    dim[0] = data.ndim() as u16;
    for (i, s) in data.shape().iter().enumerate() {
        dim[i + 1] = *s as u16;
    }

    // If no reference header is given, use the default.
    let reference = match reference {
        Some(r) => r.clone(),
        None => NiftiHeader::default(),
    };
    let header = NiftiHeader {
        dim,
        datatype: A::DATA_TYPE as i16,
        bitpix: (mem::size_of::<A>() * 8) as i16,
        // All other fields are copied from reference header
        ..reference
    };

    let f = File::create(&path).expect("Can't create new nifti file");
    let mut writer = BufWriter::new(f);
    if is_gz_file(&path) {
        let mut e = GzEncoder::new(writer, Compression::default());
        write_header(&mut e, &header)?;
        write_data(&mut e, header, data)?;
        let _ = e.finish()?; // Must use result
    } else {
        write_header(&mut writer, &header)?;
        write_data(&mut writer, header, data)?;
    }
    Ok(())
}

fn write_header<W>(writer: &mut W, header: &NiftiHeader) -> Result<()>
    where W: WriteBytesExt
{
    writer.write_i32::<B>(header.sizeof_hdr)?;
    writer.write_all(&header.data_type)?;
    writer.write_all(&header.db_name)?;
    writer.write_i32::<B>(header.extents)?;
    writer.write_i16::<B>(header.session_error)?;
    writer.write_u8(header.regular)?;
    writer.write_u8(header.dim_info)?;
    for s in &header.dim {
        writer.write_u16::<B>(*s)?;
    }
    writer.write_f32::<B>(header.intent_p1)?;
    writer.write_f32::<B>(header.intent_p2)?;
    writer.write_f32::<B>(header.intent_p3)?;
    writer.write_i16::<B>(header.intent_code)?;
    writer.write_i16::<B>(header.datatype)?;
    writer.write_i16::<B>(header.bitpix)?;
    writer.write_i16::<B>(header.slice_start)?;
    for f in &header.pixdim {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_f32::<B>(header.vox_offset)?;
    writer.write_f32::<B>(header.scl_slope)?;
    writer.write_f32::<B>(header.scl_inter)?;
    writer.write_i16::<B>(header.slice_end)?;
    writer.write_u8(header.slice_code)?;
    writer.write_u8(header.xyzt_units)?;
    writer.write_f32::<B>(header.cal_max)?;
    writer.write_f32::<B>(header.cal_min)?;
    writer.write_f32::<B>(header.slice_duration)?;
    writer.write_f32::<B>(header.toffset)?;
    writer.write_i32::<B>(header.glmax)?;
    writer.write_i32::<B>(header.glmin)?;

    writer.write_all(&header.descrip)?;
    writer.write_all(&header.aux_file)?;
    writer.write_i16::<B>(header.qform_code)?;
    writer.write_i16::<B>(header.sform_code)?;
    for f in &[
        header.quatern_b,
        header.quatern_c,
        header.quatern_d,
        header.quatern_x,
        header.quatern_y,
        header.quatern_z,
    ] {
        writer.write_f32::<B>(*f)?;
    }
    for f in header
        .srow_x
        .iter()
        .chain(&header.srow_y)
        .chain(&header.srow_z)
    {
        writer.write_f32::<B>(*f)?;
    }
    writer.write_all(&header.intent_name)?;
    writer.write_all(&header.magic)?;

    // Empty 4 bytes after the header
    // TODO Support writing extension data.
    writer.write_u32::<B>(0)?;

    Ok(())
}

/// Write the data in 'f' order.
///
/// Like NiBabel, we iterate by "slice" to improve speed and use less memory.
fn write_data<A, S, D, W>(writer: &mut W, header: NiftiHeader, data: &ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Copy,
    A: DataElement,
    A: Div<Output = A>,
    A: FromPrimitive,
    A: PodTransmutable,
    A: ScalarOperand,
    A: Sub<Output = A>,
    D: Dimension + RemoveAxis,
    W: WriteBytesExt,
{
    // Need the transpose for fortran ordering used in nifti file format.
    let data = data.t();

    let mut write_slice = |data: Array<A, D::Smaller>| -> Result<()> {
        let len = data.len();
        let arr_data = data.into_shape(len).unwrap();
        let slice = arr_data.as_slice().unwrap();
        writer.write_all(guarded_transmute_to_bytes_pod_many(slice))?;
        Ok(())
    };

    // `1.0x + 0.0` would give the same results, but we avoid a lot of divisions
    let slope = if header.scl_slope == 0.0 {
        1.0
    } else {
        header.scl_slope
    };
    if slope != 1.0 || header.scl_inter != 0.0 {
        let slope = A::from_f32(slope).unwrap();
        let inter = A::from_f32(header.scl_inter).unwrap();
        for arr_data in data.axis_iter(Axis(0)) {
            write_slice(arr_data.sub(inter).div(slope))?;
        }
    } else {
        for arr_data in data.axis_iter(Axis(0)) {
            // We need to own the data because of the into_shape() for `c` ordering.
            write_slice(arr_data.to_owned())?;
        }
    }
    Ok(())
}

#[cfg(test)]
pub mod tests {
    extern crate tempfile;

    use super::*;
    use std::ops::{Add, Mul};
    use std::path::PathBuf;

    use self::tempfile::tempdir;
    use ndarray::{Array, Array2, Ix2, IxDyn, ShapeBuilder};

    use {header::MAGIC_CODE_NIP1, object::NiftiObject, InMemNiftiObject, IntoNdArray};

    fn get_random_path(ext: &str) -> PathBuf {
        let dir = tempdir().unwrap();
        let mut path = dir.into_path();
        if !ext.is_empty() {
            path.push(ext);
        }
        path
    }

    pub fn generate_nifti_header(
        dim: [u16; 8],
        scl_slope: f32,
        scl_inter: f32,
        datatype: i16,
    ) -> NiftiHeader {
        NiftiHeader {
            dim,
            datatype,
            bitpix: (mem::size_of::<f32>() * 8) as i16,
            magic: *MAGIC_CODE_NIP1,
            scl_slope,
            scl_inter,
            ..NiftiHeader::default()
        }
    }

    fn read_2d_image<P>(path: P) -> Array2<f32>
        where P: AsRef<Path>
    {
        let nifti_object = InMemNiftiObject::from_file(path).expect("Nifti file is unreadable.");
        let volume = nifti_object.into_volume();
        let dyn_data = volume.into_ndarray().unwrap();
        dyn_data.into_dimensionality::<Ix2>().unwrap()
    }

    fn test_write_read(arr: &Array<f32, IxDyn>, path: &str) {
        let path = get_random_path(path);
        let mut dim = [1; 8];
        dim[0] = arr.ndim() as u16;
        for (i, s) in arr.shape().iter().enumerate() {
            dim[i + 1] = *s as u16;
        }
        let header = generate_nifti_header(dim, 1.0, 0.0, 16);
        write_nifti(&path, &arr, Some(&header)).unwrap();

        let read_nifti = read_2d_image(path);
        assert!(read_nifti.all_close(&arr, 1e-10));
    }

    fn f_order_array() -> Array<f32, IxDyn> {
        let dim = vec![4, 4];
        let vec = (0..16).map(|x| x as f32).collect();
        Array::from_shape_vec(IxDyn(&dim).f(), vec).unwrap()
    }

    fn c_order_array() -> Array<f32, IxDyn> {
        let dim = vec![4, 4];
        let vec = (0..16).map(|x| x as f32).collect();
        Array::from_shape_vec(IxDyn(&dim), vec).unwrap()
    }

    #[test]
    fn test_fortran_writing() {
        // Test .nii
        let arr = f_order_array();
        test_write_read(&arr, "test.nii");
        let mut arr = f_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii");

        // Test .nii.gz
        let arr = f_order_array();
        test_write_read(&arr, "test.nii.gz");
        let mut arr = f_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii.gz");
    }

    #[test]
    fn test_c_writing() {
        // Test .nii
        let arr = c_order_array();
        test_write_read(&arr, "test.nii");

        let mut arr = c_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii");

        // Test .nii.gz
        let arr = c_order_array();
        test_write_read(&arr, "test.nii.gz");

        let mut arr = c_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii.gz");
    }

    #[test]
    fn test_header_slope_inter() {
        let arr = f_order_array();
        let slope = 2.2;
        let inter = 101.1;

        let path = get_random_path("test_slope_inter.nii");
        let mut dim = [1; 8];
        dim[0] = arr.ndim() as u16;
        for (i, s) in arr.shape().iter().enumerate() {
            dim[i + 1] = *s as u16;
        }
        let header = generate_nifti_header(dim, slope, inter, 16);
        let transformed_data = arr.mul(slope).add(inter);
        write_nifti(&path, &transformed_data, Some(&header)).unwrap();

        let read_nifti = read_2d_image(path);
        assert!(read_nifti.all_close(&transformed_data, 1e-10));
    }
}
<|endoftext|>


<commit_before>// spell-checker:ignore regfile
use crate::common::util::*;

static GIBBERISH: &str = "supercalifragilisticexpialidocious";

#[cfg(not(windows))]
static NOT_A_DIRECTORY: &str = "Not a directory";
#[cfg(windows)]
static NOT_A_DIRECTORY: &str = "The directory name is invalid.";

#[test]
fn test_resolve() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;

    at.touch("foo");
    at.symlink_file("foo", "bar");

    scene.ucmd().arg("bar").succeeds().stdout_contains("foo\n");
}

#[test]
fn test_canonicalize() {
    let (at, mut ucmd) = at_and_ucmd!();
    let actual = ucmd.arg("-f").arg(".").run().stdout_move_str();
    let expect = at.root_dir_resolved() + "\n";
    println!("actual: {:?}", actual);
    println!("expect: {:?}", expect);
    assert_eq!(actual, expect);
}

#[test]
fn test_canonicalize_existing() {
    let (at, mut ucmd) = at_and_ucmd!();
    let actual = ucmd.arg("-e").arg(".").run().stdout_move_str();
    let expect = at.root_dir_resolved() + "\n";
    println!("actual: {:?}", actual);
    println!("expect: {:?}", expect);
    assert_eq!(actual, expect);
}

#[test]
fn test_canonicalize_missing() {
    let (at, mut ucmd) = at_and_ucmd!();
    let actual = ucmd.arg("-m").arg(GIBBERISH).run().stdout_move_str();
    let expect = path_concat!(at.root_dir_resolved(), GIBBERISH) + "\n";
    println!("actual: {:?}", actual);
    println!("expect: {:?}", expect);
    assert_eq!(actual, expect);
}

#[test]
fn test_long_redirection_to_current_dir() {
    let (at, mut ucmd) = at_and_ucmd!();
    // Create a 256-character path to current directory
    let dir = path_concat!(".", ..128);
    let actual = ucmd.arg("-n").arg("-m").arg(dir).run().stdout_move_str();
    let expect = at.root_dir_resolved();
    println!("actual: {:?}", actual);
    println!("expect: {:?}", expect);
    assert_eq!(actual, expect);
}

#[test]
fn test_long_redirection_to_root() {
    // Create a 255-character path to root
    let dir = path_concat!("..", ..85);
    let actual = new_ucmd!()
        .arg("-n")
        .arg("-m")
        .arg(dir)
        .run()
        .stdout_move_str();
    let expect = get_root_path();
    println!("actual: {:?}", actual);
    println!("expect: {:?}", expect);
    assert_eq!(actual, expect);
}

#[test]
fn test_symlink_to_itself_verbose() {
    let (at, mut ucmd) = at_and_ucmd!();
    at.relative_symlink_file("a", "a");
    ucmd.args(&["-ev", "a"])
        .fails()
        .code_is(1)
        .stderr_contains("Too many levels of symbolic links");
}

#[test]
fn test_trailing_slash_regular_file() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.touch("regfile");
    scene
        .ucmd()
        .args(&["-ev", "./regfile/"])
        .fails()
        .code_is(1)
        .stderr_contains(NOT_A_DIRECTORY)
        .no_stdout();
    scene
        .ucmd()
        .args(&["-e", "./regfile"])
        .succeeds()
        .stdout_contains("regfile");
}

#[test]
fn test_trailing_slash_symlink_to_regular_file() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.touch("regfile");
    at.relative_symlink_file("regfile", "link");
    scene
        .ucmd()
        .args(&["-ev", "./link/"])
        .fails()
        .code_is(1)
        .stderr_contains(NOT_A_DIRECTORY)
        .no_stdout();
    scene
        .ucmd()
        .args(&["-e", "./link"])
        .succeeds()
        .stdout_contains("regfile");
    scene
        .ucmd()
        .args(&["-ev", "./link/more"])
        .fails()
        .code_is(1)
        .stderr_contains(NOT_A_DIRECTORY)
        .no_stdout();
}

#[test]
fn test_trailing_slash_directory() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.mkdir("directory");
    for query in ["./directory", "./directory/"] {
        scene
            .ucmd()
            .args(&["-e", query])
            .succeeds()
            .stdout_contains("directory");
    }
}

#[test]
fn test_trailing_slash_symlink_to_directory() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.mkdir("directory");
    at.relative_symlink_dir("directory", "link");
    for query in ["./link", "./link/"] {
        scene
            .ucmd()
            .args(&["-e", query])
            .succeeds()
            .stdout_contains("directory");
    }
    scene
        .ucmd()
        .args(&["-ev", "./link/more"])
        .fails()
        .code_is(1)
        .stderr_contains("No such file or directory");
}

#[test]
fn test_trailing_slash_symlink_to_missing() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.mkdir("subdir");
    at.relative_symlink_file("missing", "link");
    at.relative_symlink_file("subdir/missing", "link2");
    for query in [
        "missing",
        "./missing/",
        "link",
        "./link/",
        "link/more",
        "link2",
        "./link2/",
        "link2/more",
    ] {
        scene
            .ucmd()
            .args(&["-ev", query])
            .fails()
            .code_is(1)
            .stderr_contains("No such file or directory");
    }
}

#[test]
fn test_canonicalize_trailing_slash_regfile() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.touch("regfile");
    at.relative_symlink_file("regfile", "link1");
    for name in ["regfile", "link1"] {
        scene
            .ucmd()
            .args(&["-f", name])
            .succeeds()
            .stdout_contains("regfile");
        scene
            .ucmd()
            .args(&["-fv", &format!("./{}/", name)])
            .fails()
            .code_is(1)
            .stderr_contains(NOT_A_DIRECTORY);
        scene
            .ucmd()
            .args(&["-fv", &format!("{}/more", name)])
            .fails()
            .code_is(1)
            .stderr_contains(NOT_A_DIRECTORY);
        scene
            .ucmd()
            .args(&["-fv", &format!("./{}/more/", name)])
            .fails()
            .code_is(1)
            .stderr_contains(NOT_A_DIRECTORY);
    }
}

#[test]
fn test_canonicalize_trailing_slash_subdir() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.mkdir("subdir");
    at.relative_symlink_dir("subdir", "link2");
    for name in ["subdir", "link2"] {
        scene
            .ucmd()
            .args(&["-f", name])
            .succeeds()
            .stdout_contains("subdir");
        scene
            .ucmd()
            .args(&["-f", &format!("./{}/", name)])
            .succeeds()
            .stdout_contains("subdir");
        scene
            .ucmd()
            .args(&["-f", &format!("{}/more", name)])
            .succeeds()
            .stdout_contains(path_concat!("subdir", "more"));
        scene
            .ucmd()
            .args(&["-f", &format!("./{}/more/", name)])
            .succeeds()
            .stdout_contains(path_concat!("subdir", "more"));
        scene
            .ucmd()
            .args(&["-f", &format!("{}/more/more2", name)])
            .fails()
            .code_is(1);
        scene
            .ucmd()
            .args(&["-f", &format!("./{}/more/more2/", name)])
            .fails()
            .code_is(1);
    }
}

#[test]
fn test_canonicalize_trailing_slash_missing() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.relative_symlink_file("missing", "link3");
    for name in ["missing", "link3"] {
        scene
            .ucmd()
            .args(&["-f", name])
            .succeeds()
            .stdout_contains("missing");
        scene
            .ucmd()
            .args(&["-f", &format!("./{}/", name)])
            .succeeds()
            .stdout_contains("missing");
        scene
            .ucmd()
            .args(&["-f", &format!("{}/more", name)])
            .fails()
            .code_is(1);
        scene
            .ucmd()
            .args(&["-f", &format!("./{}/more/", name)])
            .fails()
            .code_is(1);
    }
}

#[test]
fn test_canonicalize_trailing_slash_subdir_missing() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.mkdir("subdir");
    at.relative_symlink_file("subdir/missing", "link4");
    let name = "link4";
    scene
        .ucmd()
        .args(&["-f", name])
        .succeeds()
        .stdout_contains(path_concat!("subdir", "missing"));
    scene
        .ucmd()
        .args(&["-f", &format!("./{}/", name)])
        .succeeds()
        .stdout_contains(path_concat!("subdir", "missing"));
    scene
        .ucmd()
        .args(&["-f", &format!("{}/more", name)])
        .fails()
        .code_is(1);
    scene
        .ucmd()
        .args(&["-f", &format!("./{}/more/", name)])
        .fails()
        .code_is(1);
}

#[test]
fn test_canonicalize_trailing_slash_symlink_loop() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    let name = "link5";
    at.relative_symlink_file(name, name);
    scene.ucmd().args(&["-f", name]).fails().code_is(1);
    scene
        .ucmd()
        .args(&["-f", &format!("./{}/", name)])
        .fails()
        .code_is(1);
    scene
        .ucmd()
        .args(&["-f", &format!("{}/more", name)])
        .fails()
        .code_is(1);
    scene
        .ucmd()
        .args(&["-f", &format!("./{}/more/", name)])
        .fails()
        .code_is(1);
}
<commit_msg>tests/readlink: fix for Windows<commit_after>// spell-checker:ignore regfile
use crate::common::util::*;

static GIBBERISH: &str = "supercalifragilisticexpialidocious";

#[cfg(not(windows))]
static NOT_A_DIRECTORY: &str = "Not a directory";
#[cfg(windows)]
static NOT_A_DIRECTORY: &str = "The directory name is invalid.";

#[test]
fn test_resolve() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;

    at.touch("foo");
    at.symlink_file("foo", "bar");

    scene.ucmd().arg("bar").succeeds().stdout_contains("foo\n");
}

#[test]
fn test_canonicalize() {
    let (at, mut ucmd) = at_and_ucmd!();
    let actual = ucmd.arg("-f").arg(".").run().stdout_move_str();
    let expect = at.root_dir_resolved() + "\n";
    println!("actual: {:?}", actual);
    println!("expect: {:?}", expect);
    assert_eq!(actual, expect);
}

#[test]
fn test_canonicalize_existing() {
    let (at, mut ucmd) = at_and_ucmd!();
    let actual = ucmd.arg("-e").arg(".").run().stdout_move_str();
    let expect = at.root_dir_resolved() + "\n";
    println!("actual: {:?}", actual);
    println!("expect: {:?}", expect);
    assert_eq!(actual, expect);
}

#[test]
fn test_canonicalize_missing() {
    let (at, mut ucmd) = at_and_ucmd!();
    let actual = ucmd.arg("-m").arg(GIBBERISH).run().stdout_move_str();
    let expect = path_concat!(at.root_dir_resolved(), GIBBERISH) + "\n";
    println!("actual: {:?}", actual);
    println!("expect: {:?}", expect);
    assert_eq!(actual, expect);
}

#[test]
fn test_long_redirection_to_current_dir() {
    let (at, mut ucmd) = at_and_ucmd!();
    // Create a 256-character path to current directory
    let dir = path_concat!(".", ..128);
    let actual = ucmd.arg("-n").arg("-m").arg(dir).run().stdout_move_str();
    let expect = at.root_dir_resolved();
    println!("actual: {:?}", actual);
    println!("expect: {:?}", expect);
    assert_eq!(actual, expect);
}

#[test]
fn test_long_redirection_to_root() {
    // Create a 255-character path to root
    let dir = path_concat!("..", ..85);
    let actual = new_ucmd!()
        .arg("-n")
        .arg("-m")
        .arg(dir)
        .run()
        .stdout_move_str();
    let expect = get_root_path();
    println!("actual: {:?}", actual);
    println!("expect: {:?}", expect);
    assert_eq!(actual, expect);
}

#[test]
fn test_symlink_to_itself_verbose() {
    let (at, mut ucmd) = at_and_ucmd!();
    at.relative_symlink_file("a", "a");
    ucmd.args(&["-ev", "a"])
        .fails()
        .code_is(1)
        .stderr_contains("Too many levels of symbolic links");
}

#[test]
fn test_trailing_slash_regular_file() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.touch("regfile");
    scene
        .ucmd()
        .args(&["-ev", "./regfile/"])
        .fails()
        .code_is(1)
        .stderr_contains(NOT_A_DIRECTORY)
        .no_stdout();
    scene
        .ucmd()
        .args(&["-e", "./regfile"])
        .succeeds()
        .stdout_contains("regfile");
}

#[test]
fn test_trailing_slash_symlink_to_regular_file() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.touch("regfile");
    at.relative_symlink_file("regfile", "link");
    scene
        .ucmd()
        .args(&["-ev", "./link/"])
        .fails()
        .code_is(1)
        .stderr_contains(NOT_A_DIRECTORY)
        .no_stdout();
    scene
        .ucmd()
        .args(&["-e", "./link"])
        .succeeds()
        .stdout_contains("regfile");
    scene
        .ucmd()
        .args(&["-e", "./link/more"])
        .fails()
        .code_is(1)
        .no_stdout();
}

#[test]
fn test_trailing_slash_directory() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.mkdir("directory");
    for query in ["./directory", "./directory/"] {
        scene
            .ucmd()
            .args(&["-e", query])
            .succeeds()
            .stdout_contains("directory");
    }
}

#[test]
fn test_trailing_slash_symlink_to_directory() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.mkdir("directory");
    at.relative_symlink_dir("directory", "link");
    for query in ["./link", "./link/"] {
        scene
            .ucmd()
            .args(&["-e", query])
            .succeeds()
            .stdout_contains("directory");
    }
    scene
        .ucmd()
        .args(&["-ev", "./link/more"])
        .fails()
        .code_is(1)
        .stderr_contains("No such file or directory");
}

#[test]
fn test_trailing_slash_symlink_to_missing() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.mkdir("subdir");
    at.relative_symlink_file("missing", "link");
    at.relative_symlink_file("subdir/missing", "link2");
    for query in [
        "missing",
        "./missing/",
        "link",
        "./link/",
        "link/more",
        "link2",
        "./link2/",
        "link2/more",
    ] {
        scene
            .ucmd()
            .args(&["-ev", query])
            .fails()
            .code_is(1)
            .stderr_contains("No such file or directory");
    }
}

#[test]
fn test_canonicalize_trailing_slash_regfile() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.touch("regfile");
    at.relative_symlink_file("regfile", "link1");
    for name in ["regfile", "link1"] {
        scene
            .ucmd()
            .args(&["-f", name])
            .succeeds()
            .stdout_contains("regfile");
        scene
            .ucmd()
            .args(&["-fv", &format!("./{}/", name)])
            .fails()
            .code_is(1)
            .stderr_contains(NOT_A_DIRECTORY);
        scene
            .ucmd()
            .args(&["-fv", &format!("{}/more", name)])
            .fails()
            .code_is(1)
            .stderr_contains(NOT_A_DIRECTORY);
        scene
            .ucmd()
            .args(&["-fv", &format!("./{}/more/", name)])
            .fails()
            .code_is(1)
            .stderr_contains(NOT_A_DIRECTORY);
    }
}

#[test]
fn test_canonicalize_trailing_slash_subdir() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.mkdir("subdir");
    at.relative_symlink_dir("subdir", "link2");
    for name in ["subdir", "link2"] {
        scene
            .ucmd()
            .args(&["-f", name])
            .succeeds()
            .stdout_contains("subdir");
        scene
            .ucmd()
            .args(&["-f", &format!("./{}/", name)])
            .succeeds()
            .stdout_contains("subdir");
        scene
            .ucmd()
            .args(&["-f", &format!("{}/more", name)])
            .succeeds()
            .stdout_contains(path_concat!("subdir", "more"));
        scene
            .ucmd()
            .args(&["-f", &format!("./{}/more/", name)])
            .succeeds()
            .stdout_contains(path_concat!("subdir", "more"));
        scene
            .ucmd()
            .args(&["-f", &format!("{}/more/more2", name)])
            .fails()
            .code_is(1);
        scene
            .ucmd()
            .args(&["-f", &format!("./{}/more/more2/", name)])
            .fails()
            .code_is(1);
    }
}

#[test]
fn test_canonicalize_trailing_slash_missing() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.relative_symlink_file("missing", "link3");
    for name in ["missing", "link3"] {
        scene
            .ucmd()
            .args(&["-f", name])
            .succeeds()
            .stdout_contains("missing");
        scene
            .ucmd()
            .args(&["-f", &format!("./{}/", name)])
            .succeeds()
            .stdout_contains("missing");
        scene
            .ucmd()
            .args(&["-f", &format!("{}/more", name)])
            .fails()
            .code_is(1);
        scene
            .ucmd()
            .args(&["-f", &format!("./{}/more/", name)])
            .fails()
            .code_is(1);
    }
}

#[test]
fn test_canonicalize_trailing_slash_subdir_missing() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    at.mkdir("subdir");
    at.relative_symlink_file("subdir/missing", "link4");
    let name = "link4";
    scene
        .ucmd()
        .args(&["-f", name])
        .succeeds()
        .stdout_contains(path_concat!("subdir", "missing"));
    scene
        .ucmd()
        .args(&["-f", &format!("./{}/", name)])
        .succeeds()
        .stdout_contains(path_concat!("subdir", "missing"));
    scene
        .ucmd()
        .args(&["-f", &format!("{}/more", name)])
        .fails()
        .code_is(1);
    scene
        .ucmd()
        .args(&["-f", &format!("./{}/more/", name)])
        .fails()
        .code_is(1);
}

#[test]
fn test_canonicalize_trailing_slash_symlink_loop() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    let name = "link5";
    at.relative_symlink_file(name, name);
    scene.ucmd().args(&["-f", name]).fails().code_is(1);
    scene
        .ucmd()
        .args(&["-f", &format!("./{}/", name)])
        .fails()
        .code_is(1);
    scene
        .ucmd()
        .args(&["-f", &format!("{}/more", name)])
        .fails()
        .code_is(1);
    scene
        .ucmd()
        .args(&["-f", &format!("./{}/more/", name)])
        .fails()
        .code_is(1);
}
<|endoftext|>


<commit_before>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::BytesMut;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use std::time::Duration;
use futures::{Async,Poll,Sink,Stream,StartSend,Future,future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use tokio_timer::{Interval,Timer};
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec;

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("decoded frame: {:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let length = buf.len();
      if length < 8192 {
        //reserve more capacity and intialize it
        buf.extend(repeat(0).take(8192 - length));
      }
      trace!("will encode and write frame: {:?}", frame);

      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, 0), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, 0), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, 0), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            buf.truncate(sz);
            trace!("serialized frame: {} bytes", sz);
            return Ok(());
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            match e {
              GenError::BufferTooSmall(sz) => {
                buf.extend(repeat(0).take(sz - length));
                //return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
              },
              GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
                return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
              }
            }
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream: Framed<T,AMQPCodec>,
  heartbeat: Interval,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(upstream: Framed<T,AMQPCodec>, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error>> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_heartbeat(options.heartbeat);
    if let Err(e) = conn.connect() {
      let err = format!("Failed to connect: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }

    let mut t = AMQPTransport {
      upstream:  upstream,
      heartbeat: Timer::default().interval(Duration::from_secs(conn.configuration.heartbeat as u64)),
      conn:      conn,
    };

    t.send_and_handle_frames();

    let mut connector = AMQPTransportConnector {
      transport: Some(t)
    };

    trace!("pre-poll");
    connector.poll();
    trace!("post-poll");

    Box::new(connector)
  }

  fn poll_heartbeat(&mut self) {
    if let Ok(Async::Ready(_)) = self.heartbeat.poll() {
      trace!("Sending heartbeat");
      if let Err(e) = self.send_frame(Frame::Heartbeat(0)) {
        error!("Failed to send heartbeat: {:?}", e);
      }
    }
  }

  fn poll_upstream(&mut self) -> Poll<Option<()>, io::Error> {
    let value = match self.upstream.poll() {
      Ok(Async::Ready(t)) => t,
      Ok(Async::NotReady) => {
        trace!("upstream poll gave NotReady");
        return Ok(Async::NotReady);
      },
      Err(e) => {
        error!("upstream poll gave error: {:?}", e);
        return Err(From::from(e));
      },
    };

    if let Some(frame) = value {
      trace!("upstream poll gave frame: {:?}", frame);
      self.conn.handle_frame(frame);
      self.send_frames();
      Ok(Async::Ready(Some(())))
    } else {
      error!("upstream poll gave Ready(None)");
      Ok(Async::Ready(None))
    }
  }

  pub fn send_and_handle_frames(&mut self) {
    self.send_frames();
    self.handle_frames();
  }

  pub fn send_frames(&mut self) {
    //FIXME: find a way to use a future here
    while let Some(f) = self.conn.next_frame() {
      if let Err(e) = self.send_frame(f) {
        error!("Failed to send frame: {:?}", e);
        break;
      }
    }
  }

  fn send_frame(&mut self, frame: Frame) -> Poll<(), io::Error> {
      self.start_send(frame).and_then(|_| self.poll_complete())
  }

  pub fn handle_frames(&mut self) {
    while let Ok(Async::Ready(Some(_))) = self.poll() {}
  }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    debug!("AMQPTransportConnector poll transport is none? {}", self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    //we might have received a frame before here
    transport.send_and_handle_frames();

    debug!("conn state: {:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      debug!("already connected");
      return Ok(Async::Ready(transport))
    }

    trace!("waiting before poll");
    transport.poll().map(|value| {
      match value {
        Async::Ready(Some(_)) => {
          if transport.conn.state == ConnectionState::Connected {
            // Upstream had frames available and we're connected, the transport is ready
            Async::Ready(transport)
          } else {
            // Upstream had frames but we're not yet connected, continue polling
            transport.poll();
            self.transport = Some(transport);
            Async::NotReady
          }
        },
        _ => {
          // Upstream had no frames
          self.transport = Some(transport);
          Async::NotReady
        },
      }
    })
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<()>, io::Error> {
      trace!("stream poll");
      self.poll_heartbeat();
      self.poll_upstream()
    }
}

impl<T> Sink for AMQPTransport<T>
    where T: AsyncWrite {
    type SinkItem = Frame;
    type SinkError = io::Error;

    fn start_send(&mut self, item: Frame) -> StartSend<Frame, io::Error> {
        trace!("sink start send");
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        trace!("sink poll_complete");
        self.upstream.poll_complete()
    }
}

<commit_msg>futures: limit the number of frames we handle in handle_frames<commit_after>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use lapin_async::format::frame::*;

use nom::{IResult,Offset};
use cookie_factory::GenError;
use bytes::BytesMut;
use std::iter::repeat;
use std::io::{self,Error,ErrorKind};
use std::time::Duration;
use futures::{Async,Poll,Sink,Stream,StartSend,Future,future};
use tokio_io::{AsyncRead,AsyncWrite};
use tokio_io::codec::{Decoder,Encoder,Framed};
use tokio_timer::{Interval,Timer};
use client::ConnectionOptions;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec;

impl Decoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Frame>, io::Error> {
        let (consumed, f) = match frame(buf) {
          IResult::Incomplete(_) => {
            return Ok(None)
          },
          IResult::Error(e) => {
            return Err(io::Error::new(io::ErrorKind::Other, format!("parse error: {:?}", e)))
          },
          IResult::Done(i, frame) => {
            (buf.offset(i), frame)
          }
        };

        trace!("decoded frame: {:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = Frame;
    type Error = io::Error;

    fn encode(&mut self, frame: Frame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let length = buf.len();
      if length < 8192 {
        //reserve more capacity and intialize it
        buf.extend(repeat(0).take(8192 - length));
      }
      trace!("will encode and write frame: {:?}", frame);

      loop {
        let gen_res = match &frame {
          &Frame::ProtocolHeader => {
            gen_protocol_header((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Heartbeat(_) => {
            gen_heartbeat_frame((buf, 0)).map(|tup| tup.1)
          },
          &Frame::Method(channel, ref method) => {
            gen_method_frame((buf, 0), channel, method).map(|tup| tup.1)
          },
          &Frame::Header(channel_id, class_id, ref header) => {
            gen_content_header_frame((buf, 0), channel_id, class_id, header.body_size, &header.properties).map(|tup| tup.1)
          },
          &Frame::Body(channel_id, ref data) => {
            gen_content_body_frame((buf, 0), channel_id, data).map(|tup| tup.1)
          }
        };

        match gen_res {
          Ok(sz) => {
            buf.truncate(sz);
            trace!("serialized frame: {} bytes", sz);
            return Ok(());
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            match e {
              GenError::BufferTooSmall(sz) => {
                buf.extend(repeat(0).take(sz - length));
                //return Err(Error::new(ErrorKind::InvalidData, "send buffer too small"));
              },
              GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
                return Err(Error::new(ErrorKind::InvalidData, "could not generate"));
              }
            }
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream: Framed<T,AMQPCodec>,
  heartbeat: Interval,
  pub conn: Connection,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(upstream: Framed<T,AMQPCodec>, options: &ConnectionOptions) -> Box<Future<Item = AMQPTransport<T>, Error = io::Error>> {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_heartbeat(options.heartbeat);
    if let Err(e) = conn.connect() {
      let err = format!("Failed to connect: {:?}", e);
      return Box::new(future::err(Error::new(ErrorKind::ConnectionAborted, err)));
    }

    let mut t = AMQPTransport {
      upstream:  upstream,
      heartbeat: Timer::default().interval(Duration::from_secs(conn.configuration.heartbeat as u64)),
      conn:      conn,
    };

    t.send_and_handle_frames();

    let mut connector = AMQPTransportConnector {
      transport: Some(t)
    };

    trace!("pre-poll");
    connector.poll();
    trace!("post-poll");

    Box::new(connector)
  }

  fn poll_heartbeat(&mut self) {
    if let Ok(Async::Ready(_)) = self.heartbeat.poll() {
      trace!("Sending heartbeat");
      if let Err(e) = self.send_frame(Frame::Heartbeat(0)) {
        error!("Failed to send heartbeat: {:?}", e);
      }
    }
  }

  fn poll_upstream(&mut self) -> Poll<Option<()>, io::Error> {
    let value = match self.upstream.poll() {
      Ok(Async::Ready(t)) => t,
      Ok(Async::NotReady) => {
        trace!("upstream poll gave NotReady");
        return Ok(Async::NotReady);
      },
      Err(e) => {
        error!("upstream poll gave error: {:?}", e);
        return Err(From::from(e));
      },
    };

    if let Some(frame) = value {
      trace!("upstream poll gave frame: {:?}", frame);
      self.conn.handle_frame(frame);
      self.send_frames();
      Ok(Async::Ready(Some(())))
    } else {
      error!("upstream poll gave Ready(None)");
      Ok(Async::Ready(None))
    }
  }

  pub fn send_and_handle_frames(&mut self) {
    self.send_frames();
    self.handle_frames();
  }

  pub fn send_frames(&mut self) {
    //FIXME: find a way to use a future here
    while let Some(f) = self.conn.next_frame() {
      if let Err(e) = self.send_frame(f) {
        error!("Failed to send frame: {:?}", e);
        break;
      }
    }
  }

  fn send_frame(&mut self, frame: Frame) -> Poll<(), io::Error> {
      self.start_send(frame).and_then(|_| self.poll_complete())
  }

  pub fn handle_frames(&mut self) {
    for _ in 0..30 {
      if let Ok(Async::Ready(Some(_))) = self.poll() {
        // Do nothing
      } else {
        break;
      }
    }
  }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = io::Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    debug!("AMQPTransportConnector poll transport is none? {}", self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    //we might have received a frame before here
    transport.send_and_handle_frames();

    debug!("conn state: {:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      debug!("already connected");
      return Ok(Async::Ready(transport))
    }

    trace!("waiting before poll");
    transport.poll().map(|value| {
      match value {
        Async::Ready(Some(_)) => {
          if transport.conn.state == ConnectionState::Connected {
            // Upstream had frames available and we're connected, the transport is ready
            Async::Ready(transport)
          } else {
            // Upstream had frames but we're not yet connected, continue polling
            transport.poll();
            self.transport = Some(transport);
            Async::NotReady
          }
        },
        _ => {
          // Upstream had no frames
          self.transport = Some(transport);
          Async::NotReady
        },
      }
    })
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: 'static {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<()>, io::Error> {
      trace!("stream poll");
      self.poll_heartbeat();
      self.poll_upstream()
    }
}

impl<T> Sink for AMQPTransport<T>
    where T: AsyncWrite {
    type SinkItem = Frame;
    type SinkError = io::Error;

    fn start_send(&mut self, item: Frame) -> StartSend<Frame, io::Error> {
        trace!("sink start send");
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), io::Error> {
        trace!("sink poll_complete");
        self.upstream.poll_complete()
    }
}

<|endoftext|>


<commit_before>extern crate pencil;

use pencil::Pencil;
use pencil::{Request, Params, PencilResult, PenValue};

fn user(_: Request, params: Params) -> PencilResult {
    PenValue(format!("user {}", params[0]))
}

fn main() {
    let mut app: Pencil = Pencil::new("/web/example");
    app.route(r"/user/(\d+)", &["GET"], "user", user);
    app.run();
}
<commit_msg>Drop type<commit_after>extern crate pencil;

use pencil::Pencil;
use pencil::{Request, Params, PencilResult, PenValue};

fn user(_: Request, params: Params) -> PencilResult {
    PenValue(format!("user {}", params[0]))
}

fn main() {
    let mut app = Pencil::new("/web/example");
    app.route(r"/user/(\d+)", &["GET"], "user", user);
    app.run();
}
<|endoftext|>


<commit_before>extern crate rand;
extern crate timely;
extern crate differential_dataflow;

use rand::{Rng, SeedableRng, StdRng};

use differential_dataflow::input::Input;
use differential_dataflow::operators::count::CountTotal;

fn main() {

    let nodes: u32 = std::env::args().nth(1).unwrap().parse().unwrap();
    let edges: usize = std::env::args().nth(2).unwrap().parse().unwrap();
    let batch: usize = std::env::args().nth(3).unwrap().parse().unwrap();

    let inspect: bool = std::env::args().find(|x| x == "inspect").is_some();

    // define a new computational scope, in which to run BFS
    timely::execute_from_args(std::env::args().skip(4), move |worker| {

        let timer = ::std::time::Instant::now();

        let index = worker.index();
        let peers = worker.peers();

        // create a a degree counting differential dataflow
        let (mut input, probe) = worker.dataflow(|scope| {

            // create edge input, count a few ways.
            let (input, edges) = scope.new_collection();

            let degrs = edges.map(|(src, _dst)| src)
                             .count_total();

            // pull of count, and count.
            let distr = degrs.map(|(_src, cnt)| cnt as usize)
                             .count_total();

            // show us something about the collection, notice when done.
            let probe = if inspect {
                distr.inspect(|x| println!("observed: {:?}", x))
                     .probe()
            }
            else { distr.probe() };

            (input, probe)
        });

        let seed: &[_] = &[1, 2, 3, index];
        let mut rng1: StdRng = SeedableRng::from_seed(seed);    // rng for edge additions
        let mut rng2: StdRng = SeedableRng::from_seed(seed);    // rng for edge additions

        // load up graph dataz
        for _ in 0 .. (edges / peers) + if index < (edges % peers) { 1 } else { 0 } {
            input.insert((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)))
        }

        input.advance_to(1u64);
        input.flush();
        worker.step_while(|| probe.less_than(input.time()));

        println!("round 0 finished after {:?} (loading)", timer.elapsed());

        if batch > 0 {

            if !::std::env::args().any(|x| x == "open-loop") {

                // closed-loop latency-throughput test, parameterized by batch size.
                let timer = ::std::time::Instant::now();
                let mut wave = 1;
                while timer.elapsed().as_secs() < 10 {
                    for round in 0 .. batch {
                        input.advance_to((((wave * batch) + round) * peers + index) as u64);
                        input.insert((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)));
                        input.remove((rng2.gen_range(0, nodes), rng2.gen_range(0, nodes)));
                    }

                    wave += 1;
                    input.advance_to((wave * batch * peers) as u64);
                    input.flush();
                    worker.step_while(|| probe.less_than(input.time()));
                }

                let elapsed = timer.elapsed();
                let seconds = elapsed.as_secs() as f64 + (elapsed.subsec_nanos() as f64) / 1000000000.0;
                println!("{:?}, {:?}", seconds / (wave - 1) as f64, ((wave - 1) * batch * peers) as f64 / seconds);

            }
            else {

                let requests_per_sec = batch;
                let ns_per_request = 1_000_000_000 / requests_per_sec;
                let mut request_counter = peers + index;    // skip first request for each.
                let mut measurements = Vec::with_capacity(20 * requests_per_sec / peers);

                let timer = ::std::time::Instant::now();

                let measurements_per_sec = requests_per_sec / peers;

                while measurements.len() < measurements_per_sec * 21 {

                    // Open-loop latency-throughput test, parameterized by offered rate `ns_per_request`.
                    let elapsed = timer.elapsed();
                    let elapsed_ns = elapsed.as_secs() * 1_000_000_000 + (elapsed.subsec_nanos() as u64);

                    // Introduce any requests that have "arrived" since last we were here.
                    // Request i "arrives" at `index + ns_per_request * (i + 1)`.
                    while (((index + request_counter) * ns_per_request) as u64) < elapsed_ns {
                        input.advance_to(((index + request_counter) * ns_per_request) as u64);
                        input.insert((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)));
                        input.remove((rng2.gen_range(0, nodes), rng2.gen_range(0, nodes)));
                        request_counter += peers;
                    }
                    input.advance_to(elapsed_ns);
                    input.flush();

                    while probe.less_than(input.time()) {
                        worker.step();
                    }

                    // Determine completed ns.
                    let acknowledged_ns: u64 = probe.with_frontier(|frontier| frontier[0].inner);

                    let elapsed = timer.elapsed();
                    let elapsed_ns = elapsed.as_secs() * 1_000_000_000 + (elapsed.subsec_nanos() as u64);

                    // any un-recorded measurements that are complete should be recorded.
                    while (((index + peers * (measurements.len() + 1)) * ns_per_request) as u64) < acknowledged_ns {
                        let requested_at = ((index + peers * (measurements.len() + 1)) * ns_per_request) as u64;
                        measurements.push(elapsed_ns - requested_at);
                    }
                }

                measurements.truncate(20 * measurements_per_sec);
                measurements.drain(0 .. (10 * measurements_per_sec));
                measurements.sort();

                let min = measurements[0];
                let med = measurements[measurements.len() / 2];
                let p99 = measurements[99 * measurements.len() / 100];
                let max = measurements[measurements.len() - 1];

                println!("worker {}:\t{}\t{}\t{}\t{}\t(of {} measurements)", index, min, med, p99, max, measurements.len());
            }
        }
    }).unwrap();
}<commit_msg>change degrees reporting<commit_after>extern crate rand;
extern crate timely;
extern crate differential_dataflow;

use rand::{Rng, SeedableRng, StdRng};

use differential_dataflow::input::Input;
use differential_dataflow::operators::count::CountTotal;

fn main() {

    let nodes: u32 = std::env::args().nth(1).unwrap().parse().unwrap();
    let edges: usize = std::env::args().nth(2).unwrap().parse().unwrap();
    let batch: usize = std::env::args().nth(3).unwrap().parse().unwrap();

    let inspect: bool = std::env::args().find(|x| x == "inspect").is_some();

    // define a new computational scope, in which to run BFS
    timely::execute_from_args(std::env::args().skip(4), move |worker| {

        let timer = ::std::time::Instant::now();

        let index = worker.index();
        let peers = worker.peers();

        // create a a degree counting differential dataflow
        let (mut input, probe) = worker.dataflow(|scope| {

            // create edge input, count a few ways.
            let (input, edges) = scope.new_collection();

            let degrs = edges.map(|(src, _dst)| src)
                             .count_total();

            // pull of count, and count.
            let distr = degrs.map(|(_src, cnt)| cnt as usize)
                             .count_total();

            // show us something about the collection, notice when done.
            let probe = if inspect {
                distr.inspect(|x| println!("observed: {:?}", x))
                     .probe()
            }
            else { distr.probe() };

            (input, probe)
        });

        let seed: &[_] = &[1, 2, 3, index];
        let mut rng1: StdRng = SeedableRng::from_seed(seed);    // rng for edge additions
        let mut rng2: StdRng = SeedableRng::from_seed(seed);    // rng for edge additions

        // load up graph dataz
        for _ in 0 .. (edges / peers) + if index < (edges % peers) { 1 } else { 0 } {
            input.insert((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)))
        }

        input.advance_to(1u64);
        input.flush();
        worker.step_while(|| probe.less_than(input.time()));

        println!("round 0 finished after {:?} (loading)", timer.elapsed());

        if batch > 0 {

            if !::std::env::args().any(|x| x == "open-loop") {

                // closed-loop latency-throughput test, parameterized by batch size.
                let timer = ::std::time::Instant::now();
                let mut wave = 1;
                while timer.elapsed().as_secs() < 10 {
                    for round in 0 .. batch {
                        input.advance_to((((wave * batch) + round) * peers + index) as u64);
                        input.insert((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)));
                        input.remove((rng2.gen_range(0, nodes), rng2.gen_range(0, nodes)));
                    }

                    wave += 1;
                    input.advance_to((wave * batch * peers) as u64);
                    input.flush();
                    worker.step_while(|| probe.less_than(input.time()));
                }

                let elapsed = timer.elapsed();
                let seconds = elapsed.as_secs() as f64 + (elapsed.subsec_nanos() as f64) / 1000000000.0;
                println!("{:?}, {:?}", seconds / (wave - 1) as f64, ((wave - 1) * batch * peers) as f64 / seconds);

            }
            else {

                let requests_per_sec = batch;
                let ns_per_request = 1_000_000_000 / requests_per_sec;
                let mut request_counter = peers + index;    // skip first request for each.
                let mut ack_counter = peers + index;
                let mut counts = vec![0u64; 64];
                let timer = ::std::time::Instant::now();

                loop {

                    // Open-loop latency-throughput test, parameterized by offered rate `ns_per_request`.
                    let elapsed = timer.elapsed();
                    let elapsed_ns = elapsed.as_secs() * 1_000_000_000 + (elapsed.subsec_nanos() as u64);

                    // Introduce any requests that have "arrived" since last we were here.
                    // Request i "arrives" at `index + ns_per_request * (i + 1)`.
                    while (((index + request_counter) * ns_per_request) as u64) < elapsed_ns {
                        input.advance_to(((index + request_counter) * ns_per_request) as u64);
                        input.insert((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)));
                        input.remove((rng2.gen_range(0, nodes), rng2.gen_range(0, nodes)));
                        request_counter += peers;
                    }
                    input.advance_to(elapsed_ns);
                    input.flush();

                    while probe.less_than(input.time()) {
                        worker.step();
                    }

                    // Determine completed ns.
                    let acknowledged_ns: u64 = probe.with_frontier(|frontier| frontier[0].inner);

                    let elapsed = timer.elapsed();
                    let elapsed_ns = elapsed.as_secs() * 1_000_000_000 + (elapsed.subsec_nanos() as u64);

                    // any un-recorded measurements that are complete should be recorded.
                    while ((ack_counter * ns_per_request) as u64) < acknowledged_ns {
                        let requested_at = ((index + peers * ack_counter) * ns_per_request) as u64;
                        let count_index = (elapsed_ns - requested_at).next_power_of_two().trailing_zeros() as usize;
                        counts[count_index] += 1;
                        ack_counter += peers;

                        if (ack_counter & ((1 << 20) - 1)) == 0 {
                            println!("latencies:");
                            for index in 0 .. counts.len() {
                                if counts[index] > 0 {
                                    println!("\tcount[{}]:\t{}", index, counts[index]);
                                }
                            }
                            counts = vec![0u64; 64];
                        }
                    }
                }
            }
        }
    }).unwrap();
}<|endoftext|>


<commit_before>extern crate orbclient;

use std::cmp::max;
use std::env;

use orbclient::{BmpFile, Color, EventOption, Window};

fn main() {
    let path = match env::args().nth(1) {
        Some(arg) => arg,
        None => "res/redox.bmp".to_string(),
    };

    let bmp = BmpFile::from_path(&path);
    let mut window = Window::new(-1,
                                 -1,
                                 max(32, bmp.width() as u32),
                                 max(32, bmp.height() as u32),
                                 &path)
                         .unwrap();
    window.set(Color::rgb(0, 0, 0));
    window.image(0, 0, bmp.width() as u32, bmp.height() as u32, &bmp);
    window.sync();

    loop {
        for event in window.events() {
            println!("{:?}", event.to_option());
            if let EventOption::Quit(_) = event.to_option() {
                return;
            }
        }
    }
}
<commit_msg>This file has been renamed<commit_after><|endoftext|>


<commit_before>extern crate linjs;
#[macro_use] extern crate linjs_derive;

use linjs::CanAlloc;
use linjs::CanInitialize;
use linjs::JSContext;
use linjs::JSManaged;
use linjs::JSManageable;
use linjs::JSRunnable;

use std::marker::PhantomData;

// -------------------------------------------------------------------

type Window<'a, C> = JSManaged<'a, C, NativeWindow<'a, C>>;

#[derive(JSManageable)]
struct NativeWindow<'a, C> {
    console: Console<'a, C>,
    body: Element<'a, C>,
}

// -------------------------------------------------------------------

type Console<'a, C> = JSManaged<'a, C, NativeConsole<'a, C>>;

#[derive(JSManageable)]
struct NativeConsole<'a, C>(PhantomData<(&'a(), C)>);

fn new_console<'a, C, S>(cx: &'a mut JSContext<S>) -> Console<'a, C> where
    C: 'a,
    S: CanAlloc<C>,
{
    cx.manage(NativeConsole(PhantomData))
}

// -------------------------------------------------------------------

type Element<'a, C> = JSManaged<'a, C, NativeElement<'a, C>>;

#[derive(JSManageable)]
struct NativeElement<'a, C> {
    parent: Option<Element<'a, C>>,
    children: Vec<Element<'a, C>>,
}

fn new_element<'a, C, S>(cx: &'a mut JSContext<S>) -> Element<'a, C> where
    S: CanAlloc<C>,
    C: 'a,
{
    cx.manage(NativeElement {
        parent: None,
        children: Vec::new(),
    })
}

// -------------------------------------------------------------------

struct Main;

impl JSRunnable for Main {
    fn run<C, S>(self, cx: JSContext<S>) where S: CanInitialize<C> {
        println!("Hello, world.");
    }
}

fn main() {
    Main.start();
}
<commit_msg>Updated mini DOM example.<commit_after>#![allow(dead_code)]

extern crate linjs;
#[macro_use] extern crate linjs_derive;

use linjs::CanAlloc;
use linjs::CanInitialize;
use linjs::Initialized;
use linjs::JSContext;
use linjs::JSManaged;
use linjs::JSManageable;
use linjs::JSRunnable;

use std::marker::PhantomData;

// -------------------------------------------------------------------

type Window<'a, C> = JSManaged<'a, C, NativeWindow<'a, C>>;

#[derive(JSManageable)]
struct NativeWindow<'a, C> {
    console: Console<'a, C>,
    body: Element<'a, C>,
}

type DOMContext<'a, C> = JSContext<Initialized<Window<'a, C>>>;

fn init_window<'a, C, S>(cx: JSContext<S>) -> DOMContext<'a, C> where
    C: 'a,
    S: CanInitialize<C>,
{
    let mut cx = cx.pre_init();
    let ref mut roots = cx.roots();
    let console = new_console(&mut cx).root(roots);
    let body = new_element(&mut cx).root(roots);
    cx.post_init(NativeWindow {
        console: console,
        body: body,
    })
}

// -------------------------------------------------------------------

type Console<'a, C> = JSManaged<'a, C, NativeConsole<'a, C>>;

#[derive(JSManageable)]
struct NativeConsole<'a, C>(PhantomData<(&'a(), C)>);

fn new_console<'a, C, S>(cx: &'a mut JSContext<S>) -> Console<'a, C> where
    C: 'a,
    S: CanAlloc<C>,
{
    cx.manage(NativeConsole(PhantomData))
}

// -------------------------------------------------------------------

type Element<'a, C> = JSManaged<'a, C, NativeElement<'a, C>>;

#[derive(JSManageable)]
struct NativeElement<'a, C> {
    parent: Option<Element<'a, C>>,
    children: Vec<Element<'a, C>>,
}

fn new_element<'a, C, S>(cx: &'a mut JSContext<S>) -> Element<'a, C> where
    S: CanAlloc<C>,
    C: 'a,
{
    cx.manage(NativeElement {
        parent: None,
        children: Vec::new(),
    })
}

// -------------------------------------------------------------------

struct Main;

impl JSRunnable for Main {
    fn run<C, S>(self, cx: JSContext<S>) where S: CanInitialize<C> {
        let ref mut _cx = init_window(cx);
    }
}

fn main() {
    Main.start();
}
<|endoftext|>