<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 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));
    }
}
<commit_msg>Use Compression::fast<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::fast());
        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>//! Utility functions to write nifti images.

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

use byteordered::{ByteOrdered, Endian};
use flate2::write::GzEncoder;
use flate2::Compression;
use ndarray::{ArrayBase, ArrayView, Axis, Data, Dimension, RemoveAxis, ScalarOperand};
use num_traits::FromPrimitive;
use safe_transmute::{guarded_transmute_to_bytes_pod_many, PodTransmutable};

use {
    header::{MAGIC_CODE_NI1, MAGIC_CODE_NIP1},
    util::{is_gz_file, is_hdr_file, adapt_bytes},
    volume::element::DataElement,
    NiftiHeader, NiftiType, Result,
};

/// Write a nifti file (.nii or .nii.gz).
///
/// If a `reference` is given, it will be used to fill most of the header's fields. Otherwise, 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. This means that the `datatype`
/// defined in `reference` will be ignored. Because of this, `scl_slope` will be set to 1.0 and
/// `scl_inter` to 0.0.
pub fn write_nifti<P, A, S, D>(
    header_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 compression_level = Compression::fast();
    let is_gz = is_gz_file(&header_path);
    let (header, data_path) =
        prepare_header_and_paths(&header_path, data, reference, A::DATA_TYPE)?;

    // Need the transpose for fortran ordering used in nifti file format.
    let data = data.t();

    let header_file = File::create(header_path)?;
    if header.vox_offset > 0.0 {
        if is_gz {
            let mut writer = ByteOrdered::runtime(
                GzEncoder::new(header_file, compression_level),
                header.endianness,
            );
            write_header(writer.as_mut(), &header)?;
            write_data(writer.as_mut(), &header, data)?;
            let _ = writer.into_inner().finish()?;
        } else {
            let mut writer = ByteOrdered::runtime(BufWriter::new(header_file), header.endianness);
            write_header(writer.as_mut(), &header)?;
            write_data(writer, &header, data)?;
        }
    } else {
        let data_file = File::create(&data_path)?;
        if is_gz {
            let mut writer = ByteOrdered::runtime(
                GzEncoder::new(header_file, compression_level),
                header.endianness,
            );
            write_header(writer.as_mut(), &header)?;
            let _ = writer.into_inner().finish()?;

            let mut writer = ByteOrdered::runtime(
                GzEncoder::new(data_file, compression_level),
                header.endianness,
            );
            write_data(writer.as_mut(), &header, data)?;
            let _ = writer.into_inner().finish()?;
        } else {
            let header_writer =
                ByteOrdered::runtime(BufWriter::new(header_file), header.endianness);
            write_header(header_writer, &header)?;
            let data_writer =
                ByteOrdered::runtime(BufWriter::new(data_file), header.endianness);
            write_data(data_writer, &header, data)?;
        }
    }

    Ok(())
}

/// Write a RGB nifti file (.nii or .nii.gz).
///
/// If a `reference` is given, it will be used to fill most of the header's fields, except those
/// necessary to be recognized as a RGB image. `scl_slope` will be set to 1.0 and `scl_inter` to
/// 0.0. If `reference` is not given, a default `NiftiHeader` will be built and written.
pub fn write_rgb_nifti<P, S, D>(
    header_path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = [u8; 3]>,
    D: Dimension + RemoveAxis,
{
    let compression_level = Compression::fast();
    let is_gz = is_gz_file(&header_path);
    let (header, data_path) =
        prepare_header_and_paths(&header_path, data, reference, NiftiType::Rgb24)?;

    // Need the transpose for fortran used in nifti file format.
    let data = data.t();

    let header_file = File::create(header_path)?;
    if header.vox_offset > 0.0 {
        if is_gz {
            let mut writer = ByteOrdered::runtime(
                GzEncoder::new(header_file, compression_level),
                header.endianness,
            );
            write_header(writer.as_mut(), &header)?;
            write_slices::<_, u8, _, _, _, _>(writer.as_mut(), data)?;
            let _ = writer.into_inner().finish()?;
        } else {
            let mut writer = ByteOrdered::runtime(BufWriter::new(header_file), header.endianness);
            write_header(writer.as_mut(), &header)?;
            write_slices::<_, u8, _, _, _, _>(writer, data)?;
        }
    } else {
        let data_file = File::create(&data_path)?;
        if is_gz {
            let mut writer = ByteOrdered::runtime(
                GzEncoder::new(header_file, compression_level),
                header.endianness,
            );
            write_header(writer.as_mut(), &header)?;
            let _ = writer.into_inner().finish()?;

            let mut writer = ByteOrdered::runtime(
                GzEncoder::new(data_file, compression_level),
                header.endianness,
            );
            write_slices::<_, u8, _, _, _, _>(writer.as_mut(), data)?;
            let _ = writer.into_inner().finish()?;
        } else {
            let header_writer =
                ByteOrdered::runtime(BufWriter::new(header_file), header.endianness);
            write_header(header_writer, &header)?;

            let data_writer =
                ByteOrdered::runtime(BufWriter::new(data_file), header.endianness);
            write_slices::<_, u8, _, _, _, _>(data_writer, data)?;
        }
    }

    Ok(())
}

fn prepare_header_and_paths<P, T, D>(
    header_path: P,
    data: &ArrayBase<T, D>,
    reference: Option<&NiftiHeader>,
    datatype: NiftiType,
) -> Result<(NiftiHeader, PathBuf)>
where
    P: AsRef<Path>,
    T: Data,
    D: Dimension,
{
    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 => {
            let mut header = NiftiHeader::default();
            header.pixdim = [1.0; 8];
            header.sform_code = 2;
            header.srow_x = [1.0, 0.0, 0.0, 0.0];
            header.srow_y = [0.0, 1.0, 0.0, 0.0];
            header.srow_z = [0.0, 0.0, 1.0, 0.0];
            header
        }
    };

    let mut header = NiftiHeader {
        dim,
        sizeof_hdr: 348,
        datatype: datatype as i16,
        bitpix: (datatype.size_of() * 8) as i16,
        vox_offset: 352.0,
        scl_inter: 0.0,
        scl_slope: 1.0,
        magic: *MAGIC_CODE_NIP1,
        // All other fields are copied from reference header
        ..reference
    };

    // The only acceptable length is 80. If different, try to set it.
    header.validate_description()?;

    let mut path_buf = PathBuf::from(header_path.as_ref());
    let data_path = if is_hdr_file(&header_path) {
        header.vox_offset = 0.0;
        header.magic = *MAGIC_CODE_NI1;
        let data_path = if is_gz_file(&header_path) {
            let _ = path_buf.set_extension("");
            path_buf.with_extension("img.gz")
        } else {
            path_buf.with_extension("img")
        };
        data_path
    } else {
        path_buf
    };

    Ok((header, data_path))
}

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

    writer.write_all(&header.descrip)?;
    writer.write_all(&header.aux_file)?;
    writer.write_i16(header.qform_code)?;
    writer.write_i16(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(*f)?;
    }
    for f in header
        .srow_x
        .iter()
        .chain(&header.srow_y)
        .chain(&header.srow_z)
    {
        writer.write_f32(*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(0)?;

    Ok(())
}

/// Write the data in 'f' order.
///
/// Like NiBabel, we iterate by "slice" to improve speed and use less memory.
fn write_data<T, D, W, E>(mut writer: ByteOrdered<W, E>, header: &NiftiHeader, data: ArrayView<T, D>) -> Result<()>
where
    T: Clone + PodTransmutable,
    T: Div<Output = T>,
    T: FromPrimitive,
    T: PodTransmutable,
    T: ScalarOperand,
    T: Sub<Output = T>,
    D: Dimension + RemoveAxis,
    W: Write,
    E: Endian + Copy,
{
    // `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 {
        // TODO Use linear transformation like when reading. An scl_slope of 0.5 would turn all
        // voxel values to 0 if we pass an ndarray of integers.
        let slope = T::from_f32(slope).unwrap();
        let inter = T::from_f32(header.scl_inter).unwrap();
        for arr_data in data.axis_iter(Axis(0)) {
            write_slice::<T, T, _, _, _, _>(writer.as_mut(), arr_data.sub(inter).div(slope))?;
        }
    } else {
        write_slices::<T, T, _, _, _, _>(writer, data)?;
    }
    Ok(())
}

fn write_slices<A, B, S, D, W, E>(mut writer: ByteOrdered<W, E>, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension + RemoveAxis,
    W: Write,
    E: Endian + Copy,
{
    let mut iter = data.axis_iter(Axis(0));
    if let Some(arr_data) = iter.next() {
        // Keep slice voxels in a separate array to ensure `C` ordering even after `into_shape`.
        let mut slice = arr_data.to_owned();
        write_slice::<_, B, _, _, _, _>(writer.as_mut(), slice.view())?;
        for arr_data in iter {
            slice.assign(&arr_data);
            write_slice::<_, B, _, _, _, _>(writer.as_mut(), slice.view())?;
        }
    }
    Ok(())
}

fn write_slice<A, B, S, D, W, E>(writer: ByteOrdered<&mut W, E>, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension,
    W: Write,
    E: Endian,
{
    let len = data.len();
    let arr_data = data.into_shape(len).unwrap();
    let slice = arr_data.as_slice().unwrap();
    let bytes = guarded_transmute_to_bytes_pod_many(slice);
    let (writer, endianness) = writer.into_parts();
    let bytes = adapt_bytes::<B, _>(&bytes, endianness);
    writer.write_all(&*bytes)?;
    Ok(())
}
<commit_msg>Remove slope&inter management in writer<commit_after>//! Utility functions to write nifti images.

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

use byteordered::{ByteOrdered, Endian};
use flate2::write::GzEncoder;
use flate2::Compression;
use ndarray::{ArrayBase, Axis, Data, Dimension, RemoveAxis};
use safe_transmute::{guarded_transmute_to_bytes_pod_many, PodTransmutable};

use {
    header::{MAGIC_CODE_NI1, MAGIC_CODE_NIP1},
    util::{adapt_bytes, is_gz_file, is_hdr_file},
    volume::element::DataElement,
    NiftiHeader, NiftiType, Result,
};

/// Write a nifti file (.nii or .nii.gz).
///
/// If a `reference` is given, it will be used to fill most of the header's fields. Otherwise, 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. This means that the `datatype`
/// defined in `reference` will be ignored. Because of this, `scl_slope` will be set to 1.0 and
/// `scl_inter` to 0.0.
pub fn write_nifti<P, A, S, D>(
    header_path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = A>,
    A: DataElement,
    A: PodTransmutable,
    D: Dimension + RemoveAxis,
{
    let compression_level = Compression::fast();
    let is_gz = is_gz_file(&header_path);
    let (header, data_path) =
        prepare_header_and_paths(&header_path, data, reference, A::DATA_TYPE)?;

    // Need the transpose for fortran ordering used in nifti file format.
    let data = data.t();

    let header_file = File::create(header_path)?;
    if header.vox_offset > 0.0 {
        if is_gz {
            let mut writer = ByteOrdered::runtime(
                GzEncoder::new(header_file, compression_level),
                header.endianness,
            );
            write_header(writer.as_mut(), &header)?;
            write_data::<_, A, _, _, _, _>(writer.as_mut(), data)?;
            let _ = writer.into_inner().finish()?;
        } else {
            let mut writer = ByteOrdered::runtime(BufWriter::new(header_file), header.endianness);
            write_header(writer.as_mut(), &header)?;
            write_data::<_, A, _, _, _, _>(writer, data)?;
        }
    } else {
        let data_file = File::create(&data_path)?;
        if is_gz {
            let mut writer = ByteOrdered::runtime(
                GzEncoder::new(header_file, compression_level),
                header.endianness,
            );
            write_header(writer.as_mut(), &header)?;
            let _ = writer.into_inner().finish()?;

            let mut writer = ByteOrdered::runtime(
                GzEncoder::new(data_file, compression_level),
                header.endianness,
            );
            write_data::<_, A, _, _, _, _>(writer.as_mut(), data)?;
            let _ = writer.into_inner().finish()?;
        } else {
            let header_writer =
                ByteOrdered::runtime(BufWriter::new(header_file), header.endianness);
            write_header(header_writer, &header)?;
            let data_writer = ByteOrdered::runtime(BufWriter::new(data_file), header.endianness);
            write_data::<_, A, _, _, _, _>(data_writer, data)?;
        }
    }

    Ok(())
}

/// Write a RGB nifti file (.nii or .nii.gz).
///
/// If a `reference` is given, it will be used to fill most of the header's fields, except those
/// necessary to be recognized as a RGB image. `scl_slope` will be set to 1.0 and `scl_inter` to
/// 0.0. If `reference` is not given, a default `NiftiHeader` will be built and written.
pub fn write_rgb_nifti<P, S, D>(
    header_path: P,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>,
) -> Result<()>
where
    P: AsRef<Path>,
    S: Data<Elem = [u8; 3]>,
    D: Dimension + RemoveAxis,
{
    let compression_level = Compression::fast();
    let is_gz = is_gz_file(&header_path);
    let (header, data_path) =
        prepare_header_and_paths(&header_path, data, reference, NiftiType::Rgb24)?;

    // Need the transpose for fortran used in nifti file format.
    let data = data.t();

    let header_file = File::create(header_path)?;
    if header.vox_offset > 0.0 {
        if is_gz {
            let mut writer = ByteOrdered::runtime(
                GzEncoder::new(header_file, compression_level),
                header.endianness,
            );
            write_header(writer.as_mut(), &header)?;
            write_data::<_, u8, _, _, _, _>(writer.as_mut(), data)?;
            let _ = writer.into_inner().finish()?;
        } else {
            let mut writer = ByteOrdered::runtime(BufWriter::new(header_file), header.endianness);
            write_header(writer.as_mut(), &header)?;
            write_data::<_, u8, _, _, _, _>(writer, data)?;
        }
    } else {
        let data_file = File::create(&data_path)?;
        if is_gz {
            let mut writer = ByteOrdered::runtime(
                GzEncoder::new(header_file, compression_level),
                header.endianness,
            );
            write_header(writer.as_mut(), &header)?;
            let _ = writer.into_inner().finish()?;

            let mut writer = ByteOrdered::runtime(
                GzEncoder::new(data_file, compression_level),
                header.endianness,
            );
            write_data::<_, u8, _, _, _, _>(writer.as_mut(), data)?;
            let _ = writer.into_inner().finish()?;
        } else {
            let header_writer =
                ByteOrdered::runtime(BufWriter::new(header_file), header.endianness);
            write_header(header_writer, &header)?;

            let data_writer = ByteOrdered::runtime(BufWriter::new(data_file), header.endianness);
            write_data::<_, u8, _, _, _, _>(data_writer, data)?;
        }
    }

    Ok(())
}

fn prepare_header_and_paths<P, T, D>(
    header_path: P,
    data: &ArrayBase<T, D>,
    reference: Option<&NiftiHeader>,
    datatype: NiftiType,
) -> Result<(NiftiHeader, PathBuf)>
where
    P: AsRef<Path>,
    T: Data,
    D: Dimension,
{
    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 => {
            let mut header = NiftiHeader::default();
            header.pixdim = [1.0; 8];
            header.sform_code = 2;
            header.srow_x = [1.0, 0.0, 0.0, 0.0];
            header.srow_y = [0.0, 1.0, 0.0, 0.0];
            header.srow_z = [0.0, 0.0, 1.0, 0.0];
            header
        }
    };

    let mut header = NiftiHeader {
        dim,
        sizeof_hdr: 348,
        datatype: datatype as i16,
        bitpix: (datatype.size_of() * 8) as i16,
        vox_offset: 352.0,
        scl_inter: 0.0,
        scl_slope: 1.0,
        magic: *MAGIC_CODE_NIP1,
        // All other fields are copied from reference header
        ..reference
    };

    // The only acceptable length is 80. If different, try to set it.
    header.validate_description()?;

    let mut path_buf = PathBuf::from(header_path.as_ref());
    let data_path = if is_hdr_file(&header_path) {
        header.vox_offset = 0.0;
        header.magic = *MAGIC_CODE_NI1;
        let data_path = if is_gz_file(&header_path) {
            let _ = path_buf.set_extension("");
            path_buf.with_extension("img.gz")
        } else {
            path_buf.with_extension("img")
        };
        data_path
    } else {
        path_buf
    };

    Ok((header, data_path))
}

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

    writer.write_all(&header.descrip)?;
    writer.write_all(&header.aux_file)?;
    writer.write_i16(header.qform_code)?;
    writer.write_i16(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(*f)?;
    }
    for f in header
        .srow_x
        .iter()
        .chain(&header.srow_y)
        .chain(&header.srow_z)
    {
        writer.write_f32(*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(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, B, S, D, W, E>(mut writer: ByteOrdered<W, E>, data: ArrayBase<S, D>) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension + RemoveAxis,
    W: Write,
    E: Endian + Copy,
{
    let mut iter = data.axis_iter(Axis(0));
    if let Some(arr_data) = iter.next() {
        // Keep slice voxels in a separate array to ensure `C` ordering even after `into_shape`.
        let mut slice = arr_data.to_owned();
        write_slice::<_, B, _, _, _, _>(writer.as_mut(), slice.view())?;
        for arr_data in iter {
            slice.assign(&arr_data);
            write_slice::<_, B, _, _, _, _>(writer.as_mut(), slice.view())?;
        }
    }
    Ok(())
}

fn write_slice<A, B, S, D, W, E>(
    writer: ByteOrdered<&mut W, E>,
    data: ArrayBase<S, D>,
) -> Result<()>
where
    S: Data<Elem = A>,
    A: Clone + PodTransmutable,
    D: Dimension,
    W: Write,
    E: Endian,
{
    let len = data.len();
    let arr_data = data.into_shape(len).unwrap();
    let slice = arr_data.as_slice().unwrap();
    let bytes = guarded_transmute_to_bytes_pod_many(slice);
    let (writer, endianness) = writer.into_parts();
    let bytes = adapt_bytes::<B, _>(&bytes, endianness);
    writer.write_all(&*bytes)?;
    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2012-2013 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.

/*!

The Rust compiler.

# Note

This API is completely unstable and subject to change.

*/

#[crate_id = "rustc#0.10-pre"];
#[comment = "The Rust compiler"];
#[license = "MIT/ASL2"];
#[crate_type = "dylib"];
#[crate_type = "rlib"];
#[doc(html_logo_url = "http://www.rust-lang.org/logos/rust-logo-128x128-blk.png",
      html_favicon_url = "http://www.rust-lang.org/favicon.ico",
      html_root_url = "http://static.rust-lang.org/doc/master")];

#[feature(macro_rules, globs, struct_variant, managed_boxes)];
#[allow(unknown_features)]; // Note: remove it after a snapshot.
#[feature(quote)];

extern mod extra;
extern mod flate;
extern mod arena;
extern mod syntax;
extern mod serialize;
extern mod sync;
extern mod getopts;
extern mod collections;

use back::link;
use driver::session;
use middle::lint;

use d = driver::driver;

use std::io;
use std::num;
use std::os;
use std::str;
use std::task;
use std::vec;
use syntax::ast;
use syntax::attr;
use syntax::diagnostic::Emitter;
use syntax::diagnostic;
use syntax::parse;

pub mod middle {
    pub mod trans;
    pub mod ty;
    pub mod ty_fold;
    pub mod subst;
    pub mod resolve;
    pub mod resolve_lifetime;
    pub mod typeck;
    pub mod check_loop;
    pub mod check_match;
    pub mod check_const;
    pub mod lint;
    pub mod borrowck;
    pub mod dataflow;
    pub mod mem_categorization;
    pub mod liveness;
    pub mod kind;
    pub mod freevars;
    pub mod pat_util;
    pub mod region;
    pub mod const_eval;
    pub mod astencode;
    pub mod lang_items;
    pub mod privacy;
    pub mod moves;
    pub mod entry;
    pub mod effect;
    pub mod reachable;
    pub mod graph;
    pub mod cfg;
    pub mod dead;
}

pub mod front {
    pub mod config;
    pub mod test;
    pub mod std_inject;
    pub mod assign_node_ids_and_map;
    pub mod feature_gate;
    pub mod show_span;
}

pub mod back {
    pub mod archive;
    pub mod link;
    pub mod abi;
    pub mod arm;
    pub mod mips;
    pub mod x86;
    pub mod x86_64;
    pub mod rpath;
    pub mod target_strs;
    pub mod lto;
}

pub mod metadata;

pub mod driver;

pub mod util {
    pub mod common;
    pub mod ppaux;
    pub mod sha2;
}

pub mod lib {
    pub mod llvm;
    pub mod llvmdeps;
}

pub fn version(argv0: &str) {
    let vers = match option_env!("CFG_VERSION") {
        Some(vers) => vers,
        None => "unknown version"
    };
    println!("{} {}", argv0, vers);
    println!("host: {}", d::host_triple());
}

pub fn usage(argv0: &str) {
    let message = format!("Usage: {} [OPTIONS] INPUT", argv0);
    println!("{}\n\
Additional help:
    -W help             Print 'lint' options and default settings
    -Z help             Print internal options for debugging rustc\n",
              getopts::usage(message, d::optgroups()));
}

pub fn describe_warnings() {
    println!("
Available lint options:
    -W <foo>           Warn about <foo>
    -A <foo>           Allow <foo>
    -D <foo>           Deny <foo>
    -F <foo>           Forbid <foo> (deny, and deny all overrides)
");

    let lint_dict = lint::get_lint_dict();
    let mut lint_dict = lint_dict.move_iter()
                                 .map(|(k, v)| (v, k))
                                 .collect::<~[(lint::LintSpec, &'static str)]>();
    lint_dict.sort();

    let mut max_key = 0;
    for &(_, name) in lint_dict.iter() {
        max_key = num::max(name.len(), max_key);
    }
    fn padded(max: uint, s: &str) -> ~str {
        " ".repeat(max - s.len()) + s
    }
    println!("\nAvailable lint checks:\n");
    println!("    {}  {:7.7s}  {}",
             padded(max_key, "name"), "default", "meaning");
    println!("    {}  {:7.7s}  {}\n",
             padded(max_key, "----"), "-------", "-------");
    for (spec, name) in lint_dict.move_iter() {
        let name = name.replace("_", "-");
        println!("    {}  {:7.7s}  {}",
                 padded(max_key, name),
                 lint::level_to_str(spec.default),
                 spec.desc);
    }
    println!("");
}

pub fn describe_debug_flags() {
    println!("\nAvailable debug options:\n");
    let r = session::debugging_opts_map();
    for tuple in r.iter() {
        match *tuple {
            (ref name, ref desc, _) => {
                println!("    -Z {:>20s} -- {}", *name, *desc);
            }
        }
    }
}

pub fn describe_codegen_flags() {
    println!("\nAvailable codegen options:\n");
    let mut cg = session::basic_codegen_options();
    for &(name, parser, desc) in session::CG_OPTIONS.iter() {
        // we invoke the parser function on `None` to see if this option needs
        // an argument or not.
        let (width, extra) = if parser(&mut cg, None) {
            (25, "")
        } else {
            (21, "=val")
        };
        println!("    -C {:>width$s}{} -- {}", name.replace("_", "-"),
                 extra, desc, width=width);
    }
}

pub fn run_compiler(args: &[~str]) {
    let mut args = args.to_owned();
    let binary = args.shift().unwrap();

    if args.is_empty() { usage(binary); return; }

    let matches =
        &match getopts::getopts(args, d::optgroups()) {
          Ok(m) => m,
          Err(f) => {
            d::early_error(f.to_err_msg());
          }
        };

    if matches.opt_present("h") || matches.opt_present("help") {
        usage(binary);
        return;
    }

    let lint_flags = vec::append(matches.opt_strs("W"),
                                 matches.opt_strs("warn"));
    if lint_flags.iter().any(|x| x == &~"help") {
        describe_warnings();
        return;
    }

    let r = matches.opt_strs("Z");
    if r.iter().any(|x| x == &~"help") {
        describe_debug_flags();
        return;
    }

    let cg_flags = matches.opt_strs("C");
    if cg_flags.iter().any(|x| x == &~"help") {
        describe_codegen_flags();
        return;
    }

    if cg_flags.contains(&~"passes=list") {
        unsafe { lib::llvm::llvm::LLVMRustPrintPasses(); }
        return;
    }

    if matches.opt_present("v") || matches.opt_present("version") {
        version(binary);
        return;
    }
    let (input, input_file_path) = match matches.free.len() {
      0u => d::early_error("no input filename given"),
      1u => {
        let ifile = matches.free[0].as_slice();
        if ifile == "-" {
            let contents = io::stdin().read_to_end().unwrap();
            let src = str::from_utf8_owned(contents).unwrap();
            (d::StrInput(src), None)
        } else {
            (d::FileInput(Path::new(ifile)), Some(Path::new(ifile)))
        }
      }
      _ => d::early_error("multiple input filenames provided")
    };

    let sopts = d::build_session_options(binary, matches);
    let sess = d::build_session(sopts, input_file_path);
    let odir = matches.opt_str("out-dir").map(|o| Path::new(o));
    let ofile = matches.opt_str("o").map(|o| Path::new(o));
    let cfg = d::build_configuration(sess);
    let pretty = matches.opt_default("pretty", "normal").map(|a| {
        d::parse_pretty(sess, a)
    });
    match pretty {
      Some::<d::PpMode>(ppm) => {
        d::pretty_print_input(sess, cfg, &input, ppm);
        return;
      }
      None::<d::PpMode> => {/* continue */ }
    }
    let ls = matches.opt_present("ls");
    if ls {
        match input {
          d::FileInput(ref ifile) => {
            let mut stdout = io::stdout();
            d::list_metadata(sess, &(*ifile),
                             &mut stdout as &mut io::Writer).unwrap();
          }
          d::StrInput(_) => {
            d::early_error("can not list metadata for stdin");
          }
        }
        return;
    }
    let (crate_id, crate_name, crate_file_name) = sopts.print_metas;
    // these nasty nested conditions are to avoid doing extra work
    if crate_id || crate_name || crate_file_name {
        let attrs = parse_crate_attrs(sess, &input);
        let t_outputs = d::build_output_filenames(&input, &odir, &ofile,
                                                  attrs, sess);
        if crate_id || crate_name {
            let crateid = match attr::find_crateid(attrs) {
                Some(crateid) => crateid,
                None => {
                    sess.fatal("No crate_id and --crate-id or \
                                --crate-name requested")
                }
            };
            if crate_id {
                println!("{}", crateid.to_str());
            }
            if crate_name {
                println!("{}", crateid.name);
            }
        }

        if crate_file_name {
            let lm = link::build_link_meta(attrs, &t_outputs,
                                           &mut ::util::sha2::Sha256::new());
            let crate_types = session::collect_crate_types(&sess, attrs);
            for &style in crate_types.iter() {
                let fname = link::filename_for_input(&sess, style, &lm,
                                                     &t_outputs.with_extension(""));
                println!("{}", fname.filename_display());
            }
        }

        return;
    }

    d::compile_input(sess, cfg, &input, &odir, &ofile);
}

fn parse_crate_attrs(sess: session::Session,
                     input: &d::Input) -> ~[ast::Attribute] {
    match *input {
        d::FileInput(ref ifile) => {
            parse::parse_crate_attrs_from_file(ifile, ~[], sess.parse_sess)
        }
        d::StrInput(ref src) => {
            parse::parse_crate_attrs_from_source_str(d::anon_src(),
                                                     (*src).clone(),
                                                     ~[],
                                                     sess.parse_sess)
        }
    }
}

/// Run a procedure which will detect failures in the compiler and print nicer
/// error messages rather than just failing the test.
///
/// The diagnostic emitter yielded to the procedure should be used for reporting
/// errors of the compiler.
pub fn monitor(f: proc()) {
    // FIXME: This is a hack for newsched since it doesn't support split stacks.
    // rustc needs a lot of stack! When optimizations are disabled, it needs
    // even *more* stack than usual as well.
    #[cfg(rtopt)]
    static STACK_SIZE: uint = 6000000;  // 6MB
    #[cfg(not(rtopt))]
    static STACK_SIZE: uint = 20000000; // 20MB

    let mut task_builder = task::task();
    task_builder.name("rustc");

    // FIXME: Hacks on hacks. If the env is trying to override the stack size
    // then *don't* set it explicitly.
    if os::getenv("RUST_MIN_STACK").is_none() {
        task_builder.opts.stack_size = Some(STACK_SIZE);
    }

    let (p, c) = Chan::new();
    let w = io::ChanWriter::new(c);
    let mut r = io::PortReader::new(p);

    match task_builder.try(proc() {
        io::stdio::set_stderr(~w as ~io::Writer);
        f()
    }) {
        Ok(()) => { /* fallthrough */ }
        Err(value) => {
            // Task failed without emitting a fatal diagnostic
            if !value.is::<diagnostic::FatalError>() {
                diagnostic::DefaultEmitter.emit(
                    None,
                    diagnostic::ice_msg("unexpected failure"),
                    diagnostic::Error);

                let xs = [
                    ~"the compiler hit an unexpected failure path. \
                     this is a bug",
                ];
                for note in xs.iter() {
                    diagnostic::DefaultEmitter.emit(None,
                                                    *note,
                                                    diagnostic::Note)
                }

                println!("{}", r.read_to_str());
            }

            // Fail so the process returns a failure code, but don't pollute the
            // output with some unnecessary failure messages, we've already
            // printed everything that we needed to.
            io::stdio::set_stderr(~io::util::NullWriter as ~io::Writer);
            fail!();
        }
    }
}

pub fn main() {
    std::os::set_exit_status(main_args(std::os::args()));
}

pub fn main_args(args: &[~str]) -> int {
    let owned_args = args.to_owned();
    monitor(proc() run_compiler(owned_args));
    0
}
<commit_msg>auto merge of #12153 : pnkfelix/rust/fsk-add-dash-C-help-pointer, r=sanxiyn<commit_after>// Copyright 2012-2013 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.

/*!

The Rust compiler.

# Note

This API is completely unstable and subject to change.

*/

#[crate_id = "rustc#0.10-pre"];
#[comment = "The Rust compiler"];
#[license = "MIT/ASL2"];
#[crate_type = "dylib"];
#[crate_type = "rlib"];
#[doc(html_logo_url = "http://www.rust-lang.org/logos/rust-logo-128x128-blk.png",
      html_favicon_url = "http://www.rust-lang.org/favicon.ico",
      html_root_url = "http://static.rust-lang.org/doc/master")];

#[feature(macro_rules, globs, struct_variant, managed_boxes)];
#[allow(unknown_features)]; // Note: remove it after a snapshot.
#[feature(quote)];

extern mod extra;
extern mod flate;
extern mod arena;
extern mod syntax;
extern mod serialize;
extern mod sync;
extern mod getopts;
extern mod collections;

use back::link;
use driver::session;
use middle::lint;

use d = driver::driver;

use std::io;
use std::num;
use std::os;
use std::str;
use std::task;
use std::vec;
use syntax::ast;
use syntax::attr;
use syntax::diagnostic::Emitter;
use syntax::diagnostic;
use syntax::parse;

pub mod middle {
    pub mod trans;
    pub mod ty;
    pub mod ty_fold;
    pub mod subst;
    pub mod resolve;
    pub mod resolve_lifetime;
    pub mod typeck;
    pub mod check_loop;
    pub mod check_match;
    pub mod check_const;
    pub mod lint;
    pub mod borrowck;
    pub mod dataflow;
    pub mod mem_categorization;
    pub mod liveness;
    pub mod kind;
    pub mod freevars;
    pub mod pat_util;
    pub mod region;
    pub mod const_eval;
    pub mod astencode;
    pub mod lang_items;
    pub mod privacy;
    pub mod moves;
    pub mod entry;
    pub mod effect;
    pub mod reachable;
    pub mod graph;
    pub mod cfg;
    pub mod dead;
}

pub mod front {
    pub mod config;
    pub mod test;
    pub mod std_inject;
    pub mod assign_node_ids_and_map;
    pub mod feature_gate;
    pub mod show_span;
}

pub mod back {
    pub mod archive;
    pub mod link;
    pub mod abi;
    pub mod arm;
    pub mod mips;
    pub mod x86;
    pub mod x86_64;
    pub mod rpath;
    pub mod target_strs;
    pub mod lto;
}

pub mod metadata;

pub mod driver;

pub mod util {
    pub mod common;
    pub mod ppaux;
    pub mod sha2;
}

pub mod lib {
    pub mod llvm;
    pub mod llvmdeps;
}

pub fn version(argv0: &str) {
    let vers = match option_env!("CFG_VERSION") {
        Some(vers) => vers,
        None => "unknown version"
    };
    println!("{} {}", argv0, vers);
    println!("host: {}", d::host_triple());
}

pub fn usage(argv0: &str) {
    let message = format!("Usage: {} [OPTIONS] INPUT", argv0);
    println!("{}\n\
Additional help:
    -C help             Print codegen options
    -W help             Print 'lint' options and default settings
    -Z help             Print internal options for debugging rustc\n",
              getopts::usage(message, d::optgroups()));
}

pub fn describe_warnings() {
    println!("
Available lint options:
    -W <foo>           Warn about <foo>
    -A <foo>           Allow <foo>
    -D <foo>           Deny <foo>
    -F <foo>           Forbid <foo> (deny, and deny all overrides)
");

    let lint_dict = lint::get_lint_dict();
    let mut lint_dict = lint_dict.move_iter()
                                 .map(|(k, v)| (v, k))
                                 .collect::<~[(lint::LintSpec, &'static str)]>();
    lint_dict.sort();

    let mut max_key = 0;
    for &(_, name) in lint_dict.iter() {
        max_key = num::max(name.len(), max_key);
    }
    fn padded(max: uint, s: &str) -> ~str {
        " ".repeat(max - s.len()) + s
    }
    println!("\nAvailable lint checks:\n");
    println!("    {}  {:7.7s}  {}",
             padded(max_key, "name"), "default", "meaning");
    println!("    {}  {:7.7s}  {}\n",
             padded(max_key, "----"), "-------", "-------");
    for (spec, name) in lint_dict.move_iter() {
        let name = name.replace("_", "-");
        println!("    {}  {:7.7s}  {}",
                 padded(max_key, name),
                 lint::level_to_str(spec.default),
                 spec.desc);
    }
    println!("");
}

pub fn describe_debug_flags() {
    println!("\nAvailable debug options:\n");
    let r = session::debugging_opts_map();
    for tuple in r.iter() {
        match *tuple {
            (ref name, ref desc, _) => {
                println!("    -Z {:>20s} -- {}", *name, *desc);
            }
        }
    }
}

pub fn describe_codegen_flags() {
    println!("\nAvailable codegen options:\n");
    let mut cg = session::basic_codegen_options();
    for &(name, parser, desc) in session::CG_OPTIONS.iter() {
        // we invoke the parser function on `None` to see if this option needs
        // an argument or not.
        let (width, extra) = if parser(&mut cg, None) {
            (25, "")
        } else {
            (21, "=val")
        };
        println!("    -C {:>width$s}{} -- {}", name.replace("_", "-"),
                 extra, desc, width=width);
    }
}

pub fn run_compiler(args: &[~str]) {
    let mut args = args.to_owned();
    let binary = args.shift().unwrap();

    if args.is_empty() { usage(binary); return; }

    let matches =
        &match getopts::getopts(args, d::optgroups()) {
          Ok(m) => m,
          Err(f) => {
            d::early_error(f.to_err_msg());
          }
        };

    if matches.opt_present("h") || matches.opt_present("help") {
        usage(binary);
        return;
    }

    let lint_flags = vec::append(matches.opt_strs("W"),
                                 matches.opt_strs("warn"));
    if lint_flags.iter().any(|x| x == &~"help") {
        describe_warnings();
        return;
    }

    let r = matches.opt_strs("Z");
    if r.iter().any(|x| x == &~"help") {
        describe_debug_flags();
        return;
    }

    let cg_flags = matches.opt_strs("C");
    if cg_flags.iter().any(|x| x == &~"help") {
        describe_codegen_flags();
        return;
    }

    if cg_flags.contains(&~"passes=list") {
        unsafe { lib::llvm::llvm::LLVMRustPrintPasses(); }
        return;
    }

    if matches.opt_present("v") || matches.opt_present("version") {
        version(binary);
        return;
    }
    let (input, input_file_path) = match matches.free.len() {
      0u => d::early_error("no input filename given"),
      1u => {
        let ifile = matches.free[0].as_slice();
        if ifile == "-" {
            let contents = io::stdin().read_to_end().unwrap();
            let src = str::from_utf8_owned(contents).unwrap();
            (d::StrInput(src), None)
        } else {
            (d::FileInput(Path::new(ifile)), Some(Path::new(ifile)))
        }
      }
      _ => d::early_error("multiple input filenames provided")
    };

    let sopts = d::build_session_options(binary, matches);
    let sess = d::build_session(sopts, input_file_path);
    let odir = matches.opt_str("out-dir").map(|o| Path::new(o));
    let ofile = matches.opt_str("o").map(|o| Path::new(o));
    let cfg = d::build_configuration(sess);
    let pretty = matches.opt_default("pretty", "normal").map(|a| {
        d::parse_pretty(sess, a)
    });
    match pretty {
      Some::<d::PpMode>(ppm) => {
        d::pretty_print_input(sess, cfg, &input, ppm);
        return;
      }
      None::<d::PpMode> => {/* continue */ }
    }
    let ls = matches.opt_present("ls");
    if ls {
        match input {
          d::FileInput(ref ifile) => {
            let mut stdout = io::stdout();
            d::list_metadata(sess, &(*ifile),
                             &mut stdout as &mut io::Writer).unwrap();
          }
          d::StrInput(_) => {
            d::early_error("can not list metadata for stdin");
          }
        }
        return;
    }
    let (crate_id, crate_name, crate_file_name) = sopts.print_metas;
    // these nasty nested conditions are to avoid doing extra work
    if crate_id || crate_name || crate_file_name {
        let attrs = parse_crate_attrs(sess, &input);
        let t_outputs = d::build_output_filenames(&input, &odir, &ofile,
                                                  attrs, sess);
        if crate_id || crate_name {
            let crateid = match attr::find_crateid(attrs) {
                Some(crateid) => crateid,
                None => {
                    sess.fatal("No crate_id and --crate-id or \
                                --crate-name requested")
                }
            };
            if crate_id {
                println!("{}", crateid.to_str());
            }
            if crate_name {
                println!("{}", crateid.name);
            }
        }

        if crate_file_name {
            let lm = link::build_link_meta(attrs, &t_outputs,
                                           &mut ::util::sha2::Sha256::new());
            let crate_types = session::collect_crate_types(&sess, attrs);
            for &style in crate_types.iter() {
                let fname = link::filename_for_input(&sess, style, &lm,
                                                     &t_outputs.with_extension(""));
                println!("{}", fname.filename_display());
            }
        }

        return;
    }

    d::compile_input(sess, cfg, &input, &odir, &ofile);
}

fn parse_crate_attrs(sess: session::Session,
                     input: &d::Input) -> ~[ast::Attribute] {
    match *input {
        d::FileInput(ref ifile) => {
            parse::parse_crate_attrs_from_file(ifile, ~[], sess.parse_sess)
        }
        d::StrInput(ref src) => {
            parse::parse_crate_attrs_from_source_str(d::anon_src(),
                                                     (*src).clone(),
                                                     ~[],
                                                     sess.parse_sess)
        }
    }
}

/// Run a procedure which will detect failures in the compiler and print nicer
/// error messages rather than just failing the test.
///
/// The diagnostic emitter yielded to the procedure should be used for reporting
/// errors of the compiler.
pub fn monitor(f: proc()) {
    // FIXME: This is a hack for newsched since it doesn't support split stacks.
    // rustc needs a lot of stack! When optimizations are disabled, it needs
    // even *more* stack than usual as well.
    #[cfg(rtopt)]
    static STACK_SIZE: uint = 6000000;  // 6MB
    #[cfg(not(rtopt))]
    static STACK_SIZE: uint = 20000000; // 20MB

    let mut task_builder = task::task();
    task_builder.name("rustc");

    // FIXME: Hacks on hacks. If the env is trying to override the stack size
    // then *don't* set it explicitly.
    if os::getenv("RUST_MIN_STACK").is_none() {
        task_builder.opts.stack_size = Some(STACK_SIZE);
    }

    let (p, c) = Chan::new();
    let w = io::ChanWriter::new(c);
    let mut r = io::PortReader::new(p);

    match task_builder.try(proc() {
        io::stdio::set_stderr(~w as ~io::Writer);
        f()
    }) {
        Ok(()) => { /* fallthrough */ }
        Err(value) => {
            // Task failed without emitting a fatal diagnostic
            if !value.is::<diagnostic::FatalError>() {
                diagnostic::DefaultEmitter.emit(
                    None,
                    diagnostic::ice_msg("unexpected failure"),
                    diagnostic::Error);

                let xs = [
                    ~"the compiler hit an unexpected failure path. \
                     this is a bug",
                ];
                for note in xs.iter() {
                    diagnostic::DefaultEmitter.emit(None,
                                                    *note,
                                                    diagnostic::Note)
                }

                println!("{}", r.read_to_str());
            }

            // Fail so the process returns a failure code, but don't pollute the
            // output with some unnecessary failure messages, we've already
            // printed everything that we needed to.
            io::stdio::set_stderr(~io::util::NullWriter as ~io::Writer);
            fail!();
        }
    }
}

pub fn main() {
    std::os::set_exit_status(main_args(std::os::args()));
}

pub fn main_args(args: &[~str]) -> int {
    let owned_args = args.to_owned();
    monitor(proc() run_compiler(owned_args));
    0
}
<|endoftext|>


<commit_before>use serde::{Deserialize, Serialize};
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};
use shakmaty::{uci::Uci, ByColor, Color};

use crate::{model::GameId, util::ByColorDef};

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct MasterGameWithId {
    #[serde_as(as = "DisplayFromStr")]
    pub id: GameId,
    #[serde(flatten)]
    pub game: MasterGame,
}

#[serde_as]
#[derive(Serialize, Deserialize, Debug)]
pub struct MasterGame {
    pub event: String,
    pub site: String,
    pub date: String,
    pub round: String,
    #[serde(flatten, with = "ByColorDef")]
    pub players: ByColor<MasterGamePlayer>,
    #[serde_as(as = "Option<DisplayFromStr>")]
    pub winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, Uci>")]
    pub moves: Vec<Uci>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct MasterGamePlayer {
    pub name: String,
    pub rating: u16,
}
<commit_msg>finish master entry reading<commit_after>use std::{
    io,
    io::{Read, Write},
    ops::AddAssign,
};

use byteorder::{LittleEndian, ReadBytesExt};
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};
use shakmaty::{uci::Uci, ByColor, Color, Outcome};
use smallvec::{smallvec, SmallVec};

use crate::{
    model::{read_uci, GameId, Stats},
    util::ByColorDef,
};

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct MasterGameWithId {
    #[serde_as(as = "DisplayFromStr")]
    pub id: GameId,
    #[serde(flatten)]
    pub game: MasterGame,
}

#[serde_as]
#[derive(Serialize, Deserialize, Debug)]
pub struct MasterGame {
    pub event: String,
    pub site: String,
    pub date: String,
    pub round: String,
    #[serde(flatten, with = "ByColorDef")]
    pub players: ByColor<MasterGamePlayer>,
    #[serde_as(as = "Option<DisplayFromStr>")]
    pub winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, Uci>")]
    pub moves: Vec<Uci>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct MasterGamePlayer {
    pub name: String,
    pub rating: u16,
}

#[derive(Debug, Default)]
struct Group {
    stats: Stats,
    games: SmallVec<[(u16, GameId); 1]>,
}

impl AddAssign for Group {
    fn add_assign(&mut self, rhs: Group) {
        self.stats += rhs.stats;
        self.games.extend(rhs.games);
    }
}

#[derive(Default, Debug)]
pub struct MasterEntry {
    groups: FxHashMap<Uci, Group>,
}

impl MasterEntry {
    pub const SIZE_HINT: usize = 14;

    pub fn new_single(
        uci: Uci,
        id: GameId,
        outcome: Outcome,
        mover_rating: u16,
        opponent_rating: u16,
    ) -> MasterEntry {
        let mut groups = FxHashMap::with_capacity_and_hasher(1, Default::default());
        groups.insert(
            uci,
            Group {
                stats: Stats::new_single(outcome, mover_rating),
                games: smallvec![(mover_rating.saturating_add(opponent_rating), id)],
            },
        );
        MasterEntry { groups }
    }

    pub fn extend_from_reader<R: Read>(&mut self, reader: &mut R) -> io::Result<()> {
        loop {
            let uci = match read_uci(reader) {
                Ok(uci) => uci,
                Err(err) if err.kind() == io::ErrorKind::UnexpectedEof => return Ok(()),
                Err(err) => return Err(err),
            };

            let group = self.groups.entry(uci).or_default();

            group.stats += Stats::read(reader)?;
            let num_games = reader.read_u8()?;
            group.games.reserve_exact(usize::from(num_games));
            for _ in 0..num_games {
                group
                    .games
                    .push((reader.read_u16::<LittleEndian>()?, GameId::read(reader)?));
            }
        }
    }

    pub fn write<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        todo!()
    }
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use error::{
    Result as MyResult,
};
use error::Error::DriverError;
use error::DriverError::MixedParams;

enum ParserState {
    TopLevel,
    // (string_delimiter, last_char)
    InStringLiteral(char, char),
    MaybeInNamedParam,
    InNamedParam,
}

use self::ParserState::*;

pub fn parse_named_params<'a>(query: &'a str) -> MyResult<(Option<Vec<String>>, Cow<'a, str>)> {
    let mut state = TopLevel;
    let mut have_positional = false;
    let mut cur_param = 0;
    // Vec<(start_offset, end_offset, name)>
    let mut params = Vec::new();
    for (i, c) in query.chars().enumerate() {
        let mut rematch = false;
        match state {
            TopLevel => {
                match c {
                    ':' => state = MaybeInNamedParam,
                    '\'' => state = InStringLiteral('\'', '\''),
                    '"' => state = InStringLiteral('"', '"'),
                    '?' => have_positional = true,
                    _ => (),
                }
            },
            InStringLiteral(separator, prev_char) => {
                match c {
                    x if x == separator && prev_char != '\\' => state = TopLevel,
                    x => state = InStringLiteral(separator, x),
                }
            },
            MaybeInNamedParam => {
                match c {
                    'a' ... 'z' | '_' => {
                        params.push((i - 1, 0, String::with_capacity(16)));
                        params[cur_param].2.push(c);
                        state = InNamedParam;
                    },
                    _ => rematch = true,
                }
            },
            InNamedParam => {
                match c {
                    'a' ... 'z' | '0' ... '9' | '_' => params[cur_param].2.push(c),
                    _ => {
                        params[cur_param].1 = i;
                        cur_param += 1;
                        rematch = true;
                    }
                }
            },
        }
        if rematch {
            match c {
                ':' => state = MaybeInNamedParam,
                '\'' => state = InStringLiteral('\'', '\''),
                '"' => state = InStringLiteral('"', '"'),
                _ => state = TopLevel,
            }
        }
    }
    match state {
        InNamedParam => params[cur_param].1 = query.len(),
        _ => (),
    }
    if params.len() > 0 {
        if have_positional {
            return Err(DriverError(MixedParams));
        }
        let mut real_query = String::with_capacity(query.len());
        let mut last = 0;
        let mut out_params = Vec::with_capacity(params.len());
        for (start, end, name) in params.into_iter() {
            real_query.push_str(&query[last..start]);
            real_query.push('?');
            last = end;
            out_params.push(name);
        }
        real_query.push_str(&query[last..]);
        Ok((Some(out_params), real_query.into()))
    } else {
        Ok((None, query.into()))
    }
}

#[cfg(test)]
mod test {
    use error::Error::DriverError;
    use error::DriverError::MixedParams;
    use named_params::parse_named_params;

    #[test]
    fn should_parse_named_params() {
        let result = parse_named_params(":a :b").unwrap();
        assert_eq!((Some(vec!["a".to_string(), "b".into()]), "? ?".into()), result);

        let result = parse_named_params("SELECT (:a-10)").unwrap();
        assert_eq!((Some(vec!["a".to_string()]), "SELECT (?-10)".into()), result);

        let result = parse_named_params(r#"SELECT '"\':a' "'\"':c" :b"#).unwrap();
        assert_eq!((Some(vec!["b".to_string()]), r#"SELECT '"\':a' "'\"':c" ?"#.into()), result);

        let result = parse_named_params(r":a_Aa:b").unwrap();
        assert_eq!((Some(vec!["a_".to_string(), "b".into()]), r"?Aa?".into()), result);

        let result = parse_named_params(r"::b").unwrap();
        assert_eq!((Some(vec!["b".to_string()]), r":?".into()), result);

        let err = parse_named_params(r":a ?").unwrap_err();
        match err {
            DriverError(MixedParams) => (),
            _ => unreachable!(),
        }
    }

    #[test]
    fn should_allow_numbers_in_param_name() {
        let result = parse_named_params(":a1 :a2").unwrap();
        assert_eq!((Some(vec!["a1".to_string(), "a2".to_string()]), "? ?".into()), result);

        let result = parse_named_params(":1a :2a").unwrap();
        assert_eq!((None, ":1a :2a".into()), result);
    }

    #[cfg(feature = "nightly")]
    mod bench {
        use test;
        use named_params::parse_named_params;

        #[bench]
        fn parse_ten_named_params(bencher: &mut test::Bencher) {
            bencher.iter(|| {
                let result = parse_named_params(r#"
                SELECT :one, :two, :three, :four, :five, :six, :seven, :eight, :nine, :ten
                "#).unwrap();
                test::black_box(result);
            });
        }

        #[bench]
        fn parse_zero_named_params(bencher: &mut test::Bencher) {
            bencher.iter(|| {
                let result = parse_named_params(r"
                SELECT one, two, three, four, five, six, seven, eight, nine, ten
                ").unwrap();
                test::black_box(result);
            });
        }
    }
}
<commit_msg>Fix utf8 handling in parse_named_params<commit_after>use std::borrow::Cow;
use error::{
    Result as MyResult,
};
use error::Error::DriverError;
use error::DriverError::MixedParams;

enum ParserState {
    TopLevel,
    // (string_delimiter, last_char)
    InStringLiteral(char, char),
    MaybeInNamedParam,
    InNamedParam,
}

use self::ParserState::*;

pub fn parse_named_params<'a>(query: &'a str) -> MyResult<(Option<Vec<String>>, Cow<'a, str>)> {
    let mut state = TopLevel;
    let mut have_positional = false;
    let mut cur_param = 0;
    // Vec<(start_offset, end_offset, name)>
    let mut params = Vec::new();
    for (i, c) in query.char_indices() {
        let mut rematch = false;
        match state {
            TopLevel => {
                match c {
                    ':' => state = MaybeInNamedParam,
                    '\'' => state = InStringLiteral('\'', '\''),
                    '"' => state = InStringLiteral('"', '"'),
                    '?' => have_positional = true,
                    _ => (),
                }
            },
            InStringLiteral(separator, prev_char) => {
                match c {
                    x if x == separator && prev_char != '\\' => state = TopLevel,
                    x => state = InStringLiteral(separator, x),
                }
            },
            MaybeInNamedParam => {
                match c {
                    'a' ... 'z' | '_' => {
                        params.push((i - 1, 0, String::with_capacity(16)));
                        params[cur_param].2.push(c);
                        state = InNamedParam;
                    },
                    _ => rematch = true,
                }
            },
            InNamedParam => {
                match c {
                    'a' ... 'z' | '0' ... '9' | '_' => params[cur_param].2.push(c),
                    _ => {
                        params[cur_param].1 = i;
                        cur_param += 1;
                        rematch = true;
                    }
                }
            },
        }
        if rematch {
            match c {
                ':' => state = MaybeInNamedParam,
                '\'' => state = InStringLiteral('\'', '\''),
                '"' => state = InStringLiteral('"', '"'),
                _ => state = TopLevel,
            }
        }
    }
    match state {
        InNamedParam => params[cur_param].1 = query.len(),
        _ => (),
    }
    if params.len() > 0 {
        if have_positional {
            return Err(DriverError(MixedParams));
        }
        let mut real_query = String::with_capacity(query.len());
        let mut last = 0;
        let mut out_params = Vec::with_capacity(params.len());
        for (start, end, name) in params.into_iter() {
            real_query.push_str(&query[last..start]);
            real_query.push('?');
            last = end;
            out_params.push(name);
        }
        real_query.push_str(&query[last..]);
        Ok((Some(out_params), real_query.into()))
    } else {
        Ok((None, query.into()))
    }
}

#[cfg(test)]
mod test {
    use error::Error::DriverError;
    use error::DriverError::MixedParams;
    use named_params::parse_named_params;

    #[test]
    fn should_parse_named_params() {
        let result = parse_named_params(":a :b").unwrap();
        assert_eq!((Some(vec!["a".to_string(), "b".into()]), "? ?".into()), result);

        let result = parse_named_params("SELECT (:a-10)").unwrap();
        assert_eq!((Some(vec!["a".to_string()]), "SELECT (?-10)".into()), result);

        let result = parse_named_params(r#"SELECT '"\':a' "'\"':c" :b"#).unwrap();
        assert_eq!((Some(vec!["b".to_string()]), r#"SELECT '"\':a' "'\"':c" ?"#.into()), result);

        let result = parse_named_params(r":a_Aa:b").unwrap();
        assert_eq!((Some(vec!["a_".to_string(), "b".into()]), r"?Aa?".into()), result);

        let result = parse_named_params(r"::b").unwrap();
        assert_eq!((Some(vec!["b".to_string()]), r":?".into()), result);

        let err = parse_named_params(r":a ?").unwrap_err();
        match err {
            DriverError(MixedParams) => (),
            _ => unreachable!(),
        }
    }

    #[test]
    fn should_allow_numbers_in_param_name() {
        let result = parse_named_params(":a1 :a2").unwrap();
        assert_eq!((Some(vec!["a1".to_string(), "a2".to_string()]), "? ?".into()), result);

        let result = parse_named_params(":1a :2a").unwrap();
        assert_eq!((None, ":1a :2a".into()), result);
    }

    #[test]
    fn special_characters_in_query() {
        let result = parse_named_params(r"SELECT 1 FROM été WHERE thing = :param;").unwrap();
        assert_eq!((Some(vec!["param".to_string()]), "SELECT 1 FROM été WHERE thing = ?;".into()), result);
    }

    #[cfg(feature = "nightly")]
    mod bench {
        use test;
        use named_params::parse_named_params;

        #[bench]
        fn parse_ten_named_params(bencher: &mut test::Bencher) {
            bencher.iter(|| {
                let result = parse_named_params(r#"
                SELECT :one, :two, :three, :four, :five, :six, :seven, :eight, :nine, :ten
                "#).unwrap();
                test::black_box(result);
            });
        }

        #[bench]
        fn parse_zero_named_params(bencher: &mut test::Bencher) {
            bencher.iter(|| {
                let result = parse_named_params(r"
                SELECT one, two, three, four, five, six, seven, eight, nine, ten
                ").unwrap();
                test::black_box(result);
            });
        }
    }
}
<|endoftext|>


<commit_before>use gobject::GObject;
use libc::c_char;
use glib::GError;

pub type NotifyNotification = _NotifyNotification;

#[repr(C)]
struct NotifyNotificationPrivate;

#[repr(C)]
struct _NotifyNotification {
    _parent_object: GObject,
    _priv_: NotifyNotificationPrivate
}

extern "C" {
    pub fn notify_notification_new(summary: *const c_char,
                                   body: *const c_char,
                                   icon: *const c_char)
                                   -> *mut NotifyNotification;
    pub fn notify_notification_show(notification: *mut NotifyNotification,
                                    error: *mut *mut GError);
}
<commit_msg>Define NotifyNotification directly instead of using a type alias<commit_after>use gobject::GObject;
use libc::c_char;
use glib::GError;

#[repr(C)]
struct NotifyNotificationPrivate;

#[repr(C)]
pub struct NotifyNotification {
    _parent_object: GObject,
    _priv_: NotifyNotificationPrivate
}

extern "C" {
    pub fn notify_notification_new(summary: *const c_char,
                                   body: *const c_char,
                                   icon: *const c_char)
                                   -> *mut NotifyNotification;
    pub fn notify_notification_show(notification: *mut NotifyNotification,
                                    error: *mut *mut GError);
}
<|endoftext|>