<commit_before>use std::cmp::max;
use std::f32::consts::PI;

use collect_slice::CollectSlice;

use consts::SAMPLES;
use descramble::VoiceDecisions;
use enhance::EnhancedSpectrals;
use noise::Noise;
use params::BaseParams;
use prev::PrevFrame;
use window;

pub struct PhaseBase([f32; 56]);

impl PhaseBase {
    pub fn new(params: &BaseParams, prev: &PrevFrame) -> PhaseBase {
        let mut phase_base = [0.0; 56];
        let common = (prev.params.fundamental + params.fundamental) *
            SAMPLES as f32 / 2.0;

        (1..57).map(|l| {
            prev.phase_base.get(l) + common * l as f32
        }).collect_slice_checked(&mut phase_base[..]);

        PhaseBase(phase_base)
    }

    pub fn get(&self, l: usize) -> f32 { self.0[l - 1] }
}

impl Default for PhaseBase {
    fn default() -> PhaseBase {
        PhaseBase([0.0; 56])
    }
}

pub struct Phase([f32; 56]);

impl Phase {
    pub fn new(params: &BaseParams, voice: &VoiceDecisions, base: &PhaseBase) -> Phase {
        let mut noise = Noise::new();
        let mut phase = [0.0; 56];

        let trans = params.harmonics as usize / 4;

        (1..trans+1).map(|l| {
            base.get(l)
        }).chain((trans+1..57).map(|l| {
            base.get(l) + voice.unvoiced_count as f32 * noise.next_phase() /
                params.harmonics as f32
        })).collect_slice_checked(&mut phase[..]);

        Phase(phase)
    }

    pub fn get(&self, l: usize) -> f32 { self.0[l - 1] }
}

impl Default for Phase {
    fn default() -> Phase {
        Phase([0.0; 56])
    }
}

pub struct Voiced<'a, 'b, 'c, 'd> {
    prev: &'a PrevFrame,
    phase: &'b Phase,
    enhanced: &'c EnhancedSpectrals,
    voice: &'d VoiceDecisions,
    window: window::Window,
    fundamental: f32,
    end: usize,
    freq_changed: bool,
}

impl<'a, 'b, 'c, 'd> Voiced<'a, 'b, 'c, 'd> {
    pub fn new(params: &BaseParams, prev: &'a PrevFrame, phase: &'b Phase,
               enhanced: &'c EnhancedSpectrals, voice: &'d VoiceDecisions)
        -> Voiced<'a, 'b, 'c, 'd>
    {
        let freq_diff = (params.fundamental - prev.params.fundamental).abs();

        Voiced {
            prev: prev,
            phase: phase,
            enhanced: enhanced,
            voice: voice,
            window: window::synthesis_full(),
            fundamental: params.fundamental,
            end: max(params.harmonics, prev.params.harmonics) as usize + 1,
            freq_changed: freq_diff >= 0.1 * params.fundamental,
        }
    }

    fn get_pair(&self, l: usize, n: isize) -> f32 {
        match (self.voice.is_voiced(l), self.prev.voice.is_voiced(l)) {
            (false, false) => 0.0,
            (false, true) => self.sig_prev(l, n),
            (true, false) => self.sig_cur(l, n),
            (true, true) => if l >= 8 || self.freq_changed {
                self.sig_prev(l, n) + self.sig_cur(l, n)
            } else {
                self.amplitude(l, n) * self.theta(l, n).cos()
            },
        }
    }

    fn sig_cur(&self, l: usize, n: isize) -> f32 {
        self.window.get(n - SAMPLES as isize) * self.enhanced.get(l) * (
            self.fundamental * (n - SAMPLES as isize) as f32 * l as f32 +
                self.phase.get(l)
        ).cos()
    }

    fn sig_prev(&self, l: usize, n: isize) -> f32 {
        self.window.get(n) * self.prev.enhanced.get(l) * (
            self.prev.params.fundamental * n as f32 * l as f32 +
                self.prev.phase.get(l)
        ).cos()
    }

    fn amplitude(&self, l: usize, n: isize) -> f32 {
        self.prev.enhanced.get(l) + n as f32 / SAMPLES as f32 *
            (self.enhanced.get(l) - self.prev.enhanced.get(l))
    }

    fn theta(&self, l: usize, n: isize) -> f32 {
        self.prev.phase.get(l) +
            (self.prev.params.fundamental * l as f32 + self.freq_change(l)) * n as f32 +
            (self.fundamental - self.prev.params.fundamental) *
                l as f32 * (n as f32).powi(2) / (2.0 * SAMPLES as f32)
    }

    fn phase_change(&self, l: usize) -> f32 {
        self.phase.get(l) - self.prev.phase.get(l) - (
            self.prev.params.fundamental + self.fundamental
        ) * l as f32 * SAMPLES as f32 / 2.0
    }

    fn freq_change(&self, l: usize) -> f32 {
        (SAMPLES as f32).recip() * (self.phase_change(l) - 2.0 * PI * (
            (self.phase_change(l) + PI) / (2.0 * PI)
        ).floor())
    }

    pub fn get(&self, n: usize) -> f32 {
        2.0 * (1..self.end).map(|l| {
            self.get_pair(l, n as isize)
        }).fold(0.0, |s, x| s + x)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use params::BaseParams;
    use prev::PrevFrame;
    use descramble::{descramble, Bootstrap};

    #[test]
    fn test_phase_base() {
        let chunks = [
            0b001000010010,
            0b110011001100,
            0b111000111000,
            0b111111111111,
            0b10101110101,
            0b00101111010,
            0b01110111011,
            0b00001000,
        ];

        let b = Bootstrap::new(&chunks);
        let p = BaseParams::new(b.unwrap_period());
        let prev = PrevFrame::default();

        assert!((p.fundamental - 0.17575344).abs() < 0.000001);
        assert!((prev.params.fundamental - 0.0937765407).abs() < 0.000001);

        let pb = PhaseBase::new(&p, &prev);

        assert!((pb.get(1) - 21.56239846).abs() < 0.0001);
        assert!((pb.get(2) - 43.12479691).abs() < 0.0001);
        assert!((pb.get(3) - 64.68719537).abs() < 0.0001);
        assert!((pb.get(4) - 86.24959382).abs() < 0.0001);
        assert!((pb.get(5) - 107.8119923).abs() < 0.0001);
        assert!((pb.get(6) - 129.3743907).abs() < 0.0001);
        assert!((pb.get(20) - 431.2479691).abs() < 0.0001);
        assert!((pb.get(56) - 1207.494314).abs() < 0.001);
    }
}
<commit_msg>Simplify voiced computation if both frames are voiced<commit_after>use std::cmp::max;

use collect_slice::CollectSlice;

use consts::SAMPLES;
use descramble::VoiceDecisions;
use enhance::EnhancedSpectrals;
use noise::Noise;
use params::BaseParams;
use prev::PrevFrame;
use window;

pub struct PhaseBase([f32; 56]);

impl PhaseBase {
    pub fn new(params: &BaseParams, prev: &PrevFrame) -> PhaseBase {
        let mut phase_base = [0.0; 56];
        let common = (prev.params.fundamental + params.fundamental) *
            SAMPLES as f32 / 2.0;

        (1..57).map(|l| {
            prev.phase_base.get(l) + common * l as f32
        }).collect_slice_checked(&mut phase_base[..]);

        PhaseBase(phase_base)
    }

    pub fn get(&self, l: usize) -> f32 { self.0[l - 1] }
}

impl Default for PhaseBase {
    fn default() -> PhaseBase {
        PhaseBase([0.0; 56])
    }
}

pub struct Phase([f32; 56]);

impl Phase {
    pub fn new(params: &BaseParams, voice: &VoiceDecisions, base: &PhaseBase) -> Phase {
        let mut noise = Noise::new();
        let mut phase = [0.0; 56];

        let trans = params.harmonics as usize / 4;

        (1..trans+1).map(|l| {
            base.get(l)
        }).chain((trans+1..57).map(|l| {
            base.get(l) + voice.unvoiced_count as f32 * noise.next_phase() /
                params.harmonics as f32
        })).collect_slice_checked(&mut phase[..]);

        Phase(phase)
    }

    pub fn get(&self, l: usize) -> f32 { self.0[l - 1] }
}

impl Default for Phase {
    fn default() -> Phase {
        Phase([0.0; 56])
    }
}

pub struct Voiced<'a, 'b, 'c, 'd> {
    prev: &'a PrevFrame,
    phase: &'b Phase,
    enhanced: &'c EnhancedSpectrals,
    voice: &'d VoiceDecisions,
    window: window::Window,
    fundamental: f32,
    end: usize,
}

impl<'a, 'b, 'c, 'd> Voiced<'a, 'b, 'c, 'd> {
    pub fn new(params: &BaseParams, prev: &'a PrevFrame, phase: &'b Phase,
               enhanced: &'c EnhancedSpectrals, voice: &'d VoiceDecisions)
        -> Voiced<'a, 'b, 'c, 'd>
    {
        Voiced {
            prev: prev,
            phase: phase,
            enhanced: enhanced,
            voice: voice,
            window: window::synthesis_full(),
            fundamental: params.fundamental,
            end: max(params.harmonics, prev.params.harmonics) as usize + 1,
        }
    }

    fn get_pair(&self, l: usize, n: isize) -> f32 {
        match (self.voice.is_voiced(l), self.prev.voice.is_voiced(l)) {
            (false, false) => 0.0,
            (false, true) => self.sig_prev(l, n),
            (true, false) => self.sig_cur(l, n),
            (true, true) => self.sig_prev(l, n) + self.sig_cur(l, n)
        }
    }

    fn sig_cur(&self, l: usize, n: isize) -> f32 {
        self.window.get(n - SAMPLES as isize) * self.enhanced.get(l) * (
            self.fundamental * (n - SAMPLES as isize) as f32 * l as f32 +
                self.phase.get(l)
        ).cos()
    }

    fn sig_prev(&self, l: usize, n: isize) -> f32 {
        self.window.get(n) * self.prev.enhanced.get(l) * (
            self.prev.params.fundamental * n as f32 * l as f32 +
                self.prev.phase.get(l)
        ).cos()
    }

    pub fn get(&self, n: usize) -> f32 {
        2.0 * (1..self.end).map(|l| {
            self.get_pair(l, n as isize)
        }).fold(0.0, |s, x| s + x)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use params::BaseParams;
    use prev::PrevFrame;
    use descramble::{descramble, Bootstrap};

    #[test]
    fn test_phase_base() {
        let chunks = [
            0b001000010010,
            0b110011001100,
            0b111000111000,
            0b111111111111,
            0b10101110101,
            0b00101111010,
            0b01110111011,
            0b00001000,
        ];

        let b = Bootstrap::new(&chunks);
        let p = BaseParams::new(b.unwrap_period());
        let prev = PrevFrame::default();

        assert!((p.fundamental - 0.17575344).abs() < 0.000001);
        assert!((prev.params.fundamental - 0.0937765407).abs() < 0.000001);

        let pb = PhaseBase::new(&p, &prev);

        assert!((pb.get(1) - 21.56239846).abs() < 0.0001);
        assert!((pb.get(2) - 43.12479691).abs() < 0.0001);
        assert!((pb.get(3) - 64.68719537).abs() < 0.0001);
        assert!((pb.get(4) - 86.24959382).abs() < 0.0001);
        assert!((pb.get(5) - 107.8119923).abs() < 0.0001);
        assert!((pb.get(6) - 129.3743907).abs() < 0.0001);
        assert!((pb.get(20) - 431.2479691).abs() < 0.0001);
        assert!((pb.get(56) - 1207.494314).abs() < 0.001);
    }
}
<|endoftext|>


<commit_before>//! Write las files.

use std::path::Path;

use Result;
use file::File;
use header::Header;
use point::Point;

/// A las writer.
///
/// This wrapper conforms to the more standard structure of requiring a filename on create, not on
/// close.
///
/// I recognize that it's pretty messy to have both this and `File`, and TODO I need to clean
/// things up.
#[derive(Debug)]
pub struct Writer<P: AsRef<Path>> {
    auto_offsets: bool,
    file: File,
    header: Header,
    path: P,
}

impl<P: AsRef<Path>> Writer<P> {
    /// Creates a new writer that will write las data to the given path.
    ///
    /// This won't actually write anything until the writer is closed.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::writer::Writer;
    /// let writer = Writer::from_path("temp.las");
    /// ```
    pub fn from_path(path: P) -> Writer<P> {
        Writer {
            auto_offsets: false,
            file: File::new(),
            header: Header::new(),
            path: path,
        }
    }

    /// Sets the scale factors on a writer.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::writer::Writer;
    /// let writer = Writer::from_path("temp.las").scale_factors(0.01, 0.01, 0.01);
    /// ```
    pub fn scale_factors(mut self,
                         x_scale_factor: f64,
                         y_scale_factor: f64,
                         z_scale_factor: f64)
                         -> Writer<P> {
        self.header.x_scale_factor = x_scale_factor;
        self.header.y_scale_factor = y_scale_factor;
        self.header.z_scale_factor = z_scale_factor;
        self
    }

    /// Sets the offset values for a file.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::writer::Writer;
    /// let writer = Writer::from_path("temp.las").offsets(1000.0, 2000.0, 100.0);
    /// ```
    pub fn offsets(mut self, x_offset: f64, y_offset: f64, z_offset: f64) -> Writer<P> {
        self.header.x_offset = x_offset;
        self.header.y_offset = y_offset;
        self.header.z_offset = z_offset;
        self
    }

    /// Enables auto-offsetting.
    ///
    /// If auto-offsetting is enabled, this file will set the header offset values to sensible
    /// values before writing anything. This is usually easier than calculating the offsets
    /// yourself.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::writer::Writer;
    /// let writer = Writer::from_path("temp.las").auto_offsets(true);
    /// ```
    pub fn auto_offsets(mut self, enable: bool) -> Writer<P> {
        self.auto_offsets = enable;
        self
    }

    /// Writes a point to this writer.
    ///
    /// Note that this point won't actually be written until the writer is closed.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::writer::Writer;
    /// use las::point::Point;
    /// let mut writer = Writer::from_path("temp.las");
    /// writer.write_point(Point::new());
    /// ```
    pub fn write_point(&mut self, point: Point) {
        self.file.add_point(point)
    }

    /// Closes this writer and actually writes data out to disc.
    ///
    /// Since we need to calculate some stats on the points for the header, we delay writing until
    /// the very last minute. If you don't want to hold all those points in memory, we'll need to
    /// come up with some other way to do that.
    ///
    /// This function consumes the writer.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::fs::remove_file;
    /// use las::writer::Writer;
    /// use las::point::Point;
    /// let mut writer = Writer::from_path("temp.las");
    /// writer.write_point(Point::new());
    /// writer.close().unwrap();
    /// remove_file("temp.las").unwrap();
    /// ```
    pub fn close(mut self) -> Result<()> {
        self.file.set_header(self.header);
        self.file.to_path(self.path, self.auto_offsets)
    }
}
<commit_msg>Writer close takes mut ref, not consuming<commit_after>//! Write las files.

use std::path::Path;

use Result;
use file::File;
use header::Header;
use point::Point;

/// A las writer.
///
/// This wrapper conforms to the more standard structure of requiring a filename on create, not on
/// close.
///
/// I recognize that it's pretty messy to have both this and `File`, and TODO I need to clean
/// things up.
#[derive(Debug)]
pub struct Writer<P: AsRef<Path>> {
    auto_offsets: bool,
    file: File,
    header: Header,
    path: P,
}

impl<P: AsRef<Path>> Writer<P> {
    /// Creates a new writer that will write las data to the given path.
    ///
    /// This won't actually write anything until the writer is closed.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::writer::Writer;
    /// let writer = Writer::from_path("temp.las");
    /// ```
    pub fn from_path(path: P) -> Writer<P> {
        Writer {
            auto_offsets: false,
            file: File::new(),
            header: Header::new(),
            path: path,
        }
    }

    /// Sets the scale factors on a writer.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::writer::Writer;
    /// let writer = Writer::from_path("temp.las").scale_factors(0.01, 0.01, 0.01);
    /// ```
    pub fn scale_factors(mut self,
                         x_scale_factor: f64,
                         y_scale_factor: f64,
                         z_scale_factor: f64)
                         -> Writer<P> {
        self.header.x_scale_factor = x_scale_factor;
        self.header.y_scale_factor = y_scale_factor;
        self.header.z_scale_factor = z_scale_factor;
        self
    }

    /// Sets the offset values for a file.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::writer::Writer;
    /// let writer = Writer::from_path("temp.las").offsets(1000.0, 2000.0, 100.0);
    /// ```
    pub fn offsets(mut self, x_offset: f64, y_offset: f64, z_offset: f64) -> Writer<P> {
        self.header.x_offset = x_offset;
        self.header.y_offset = y_offset;
        self.header.z_offset = z_offset;
        self
    }

    /// Enables auto-offsetting.
    ///
    /// If auto-offsetting is enabled, this file will set the header offset values to sensible
    /// values before writing anything. This is usually easier than calculating the offsets
    /// yourself.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::writer::Writer;
    /// let writer = Writer::from_path("temp.las").auto_offsets(true);
    /// ```
    pub fn auto_offsets(mut self, enable: bool) -> Writer<P> {
        self.auto_offsets = enable;
        self
    }

    /// Writes a point to this writer.
    ///
    /// Note that this point won't actually be written until the writer is closed.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::writer::Writer;
    /// use las::point::Point;
    /// let mut writer = Writer::from_path("temp.las");
    /// writer.write_point(Point::new());
    /// ```
    pub fn write_point(&mut self, point: Point) {
        self.file.add_point(point)
    }

    /// Closes this writer and actually writes data out to disc.
    ///
    /// Since we need to calculate some stats on the points for the header, we delay writing until
    /// the very last minute. If you don't want to hold all those points in memory, we'll need to
    /// come up with some other way to do that.
    ///
    /// This function consumes the writer.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::fs::remove_file;
    /// use las::writer::Writer;
    /// use las::point::Point;
    /// let mut writer = Writer::from_path("temp.las");
    /// writer.write_point(Point::new());
    /// writer.close().unwrap();
    /// remove_file("temp.las").unwrap();
    /// ```
    pub fn close(&mut self) -> Result<()> {
        self.file.set_header(self.header);
        self.file.to_path(&self.path, self.auto_offsets)
    }
}
<|endoftext|>


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

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

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

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

type B = LittleEndian;

/// 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. 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<A, S, D>(
    path: &str,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>)
    where S: Data<Elem=A>,
          A: Copy + DataElement + Div<Output = A> + FromPrimitive + ScalarOperand + 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 buffer = Vec::new();
        write_header(&mut buffer, &header);
        write_data(&mut buffer, header, data);

        let mut e = GzEncoder::new(Vec::new(), Compression::default());
        e.write_all(&buffer).unwrap();

        let buffer = e.finish().unwrap();
        writer.write_all(&buffer).unwrap();
    } else {
        write_header(&mut writer, &header);
        write_data(&mut writer, header, data);
    }
}

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

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

    // Empty 4 bytes after the header
    writer.write_u32::<B>(0).unwrap();
}

fn write_data<A, S, D, W>(
    writer: &mut W,
    header: NiftiHeader,
    data: &ArrayBase<S, D>)
    where S: Data<Elem=A>,
          A: Copy + DataElement + Div<Output = A> + FromPrimitive + ScalarOperand + Sub<Output = A>,
          D: Dimension + RemoveAxis,
          W: WriteBytesExt
{
    // We finally write the data.
    // Like NiBabel, we iterate by "slice" to improve speed and use less memory

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

    let arr_len = data.subview(Axis(0), 0).len();
    let resolution: usize = header.dim[1..(header.dim[0] + 1) as usize]
        .iter().map(|d| *d as usize).product();
    let nb_bytes = resolution * header.bitpix as usize / 8;
    let dim0_size = data.shape()[0];

    // `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)) {
            let arr_data = arr_data.sub(inter).div(slope);
            let arr_data = arr_data.into_shape(arr_len).unwrap();
            let arr_data = arr_data.as_slice().unwrap();

            let buffer = to_bytes(&arr_data, nb_bytes / dim0_size);
            writer.write_all(&buffer).unwrap();
        }
    } else {
        for arr_data in data.axis_iter(Axis(0)) {
            // We need to own the data because of the into_shape() for `C` ordering.
            let arr_data = arr_data.to_owned();
            let arr_data = arr_data.into_shape(arr_len).unwrap();
            let arr_data = arr_data.as_slice().unwrap();

            let buffer = to_bytes(&arr_data, nb_bytes / dim0_size);
            writer.write_all(&buffer).unwrap();
        }
    }
}

fn to_bytes<T>(data: &[T], size: usize) -> &[u8] {
    unsafe { from_raw_parts::<u8>(data.as_ptr() as *const u8, size) }
}

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

    use super::*;

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

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

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

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

    fn read_2d_image(path: &str) -> 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));

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

    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() {
        use std::ops::{Add, Mul};

        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));

        let read_nifti = read_2d_image(&path);
        assert!(read_nifti.all_close(&transformed_data, 1e-10) == true);
    }
}
<commit_msg>Recoverable errors when writing images<commit_after>//! Utility functions to write nifti images.

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

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

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

type B = LittleEndian;

/// 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. 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<A, S, D>(
    path: &str,
    data: &ArrayBase<S, D>,
    reference: Option<&NiftiHeader>
) -> Result<()>
    where S: Data<Elem=A>,
          A: Copy + DataElement + Div<Output = A> + FromPrimitive + ScalarOperand + 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 buffer = Vec::new();
        write_header(&mut buffer, &header)?;
        write_data(&mut buffer, header, data)?;

        let mut e = GzEncoder::new(Vec::new(), Compression::default());
        e.write_all(&buffer)?;

        let buffer = e.finish()?;
        writer.write_all(&buffer)?;
    } 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
    writer.write_u32::<B>(0)?;

    Ok(())
}

fn write_data<A, S, D, W>(
    writer: &mut W,
    header: NiftiHeader,
    data: &ArrayBase<S, D>
) -> Result<()>
    where S: Data<Elem=A>,
          A: Copy + DataElement + Div<Output = A> + FromPrimitive + ScalarOperand + Sub<Output = A>,
          D: Dimension + RemoveAxis,
          W: WriteBytesExt
{
    // We finally write the data.
    // Like NiBabel, we iterate by "slice" to improve speed and use less memory

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

    let arr_len = data.subview(Axis(0), 0).len();
    let resolution: usize = header.dim[1..(header.dim[0] + 1) as usize]
        .iter().map(|d| *d as usize).product();
    let nb_bytes = resolution * header.bitpix as usize / 8;
    let dim0_size = data.shape()[0];

    // `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)) {
            let arr_data = arr_data.sub(inter).div(slope);
            let arr_data = arr_data.into_shape(arr_len).unwrap();
            let arr_data = arr_data.as_slice().unwrap();

            let buffer = to_bytes(&arr_data, nb_bytes / dim0_size);
            writer.write_all(&buffer)?;
        }
    } else {
        for arr_data in data.axis_iter(Axis(0)) {
            // We need to own the data because of the into_shape() for `C` ordering.
            let arr_data = arr_data.to_owned();
            let arr_data = arr_data.into_shape(arr_len).unwrap();
            let arr_data = arr_data.as_slice().unwrap();

            let buffer = to_bytes(&arr_data, nb_bytes / dim0_size);
            writer.write_all(&buffer)?;
        }
    }
    Ok(())
}

fn to_bytes<T>(data: &[T], size: usize) -> &[u8] {
    unsafe { from_raw_parts::<u8>(data.as_ptr() as *const u8, size) }
}

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

    use super::*;

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

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

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

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

    fn read_2d_image(path: &str) -> 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) == true);
    }

    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() {
        use std::ops::{Add, Mul};

        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) == true);
    }
}
<|endoftext|>


<commit_before>use anyhow::Result;
use log::debug;

use crate::{arch::Arch, module::Module};

pub fn get_disassembler(module: &Module) -> Result<zydis::Decoder> {
    let mut decoder = match module.arch {
        Arch::X64 => zydis::Decoder::new(zydis::MachineMode::LONG_64, zydis::AddressWidth::_64)?,
        Arch::X32 => zydis::Decoder::new(zydis::MachineMode::LEGACY_32, zydis::AddressWidth::_32)?,
    };

    // modes described here: https://github.com/zyantific/zydis/blob/5af06d64432aaa3f6af3cd3e120eefa061b790ab/include/Zydis/Decoder.h#L55
    //
    // performance, captured empirically:
    //  - minimal mode - 8.7M instructions/second
    //  - full mode    - 4.5M instructions/second
    decoder.enable_mode(zydis::DecoderMode::MINIMAL, false)?;

    decoder.enable_mode(zydis::DecoderMode::KNC, false)?;
    decoder.enable_mode(zydis::DecoderMode::MPX, false)?;
    decoder.enable_mode(zydis::DecoderMode::CET, false)?;
    decoder.enable_mode(zydis::DecoderMode::LZCNT, false)?;
    decoder.enable_mode(zydis::DecoderMode::TZCNT, false)?;
    decoder.enable_mode(zydis::DecoderMode::WBNOINVD, false)?;
    decoder.enable_mode(zydis::DecoderMode::CLDEMOTE, false)?;

    Ok(decoder)
}

pub fn linear_disassemble<'a>(
    decoder: &'a zydis::Decoder,
    buf: &'a [u8],
) -> Box<dyn Iterator<Item = (usize, zydis::Result<Option<zydis::DecodedInstruction>>)> + 'a> {
    let mut offset = 0usize;
    let mut insn_count = 0usize;
    let iter = std::iter::from_fn(move || {
        if offset >= buf.len() {
            debug!("decoded {} instructions", insn_count);
            return None;
        }

        let insn_offset = offset;
        let insn_buf = &buf[insn_offset..];
        let insn = decoder.decode(insn_buf);

        if let Ok(Some(insn)) = &insn {
            // see discussion of linear vs thorough disassemble in this module doc for
            // call_targets. thorough is 4x more expensive, with limited
            // results.

            // linear disassembly:
            offset += insn.length as usize;

            // thorough disassembly:
            // offset += 1;

            insn_count += 1;
        } else {
            offset += 1;
        }

        Some((insn_offset, insn))
    });

    Box::new(iter)
}
<commit_msg>dis: re-published zydis since its a part of our interface<commit_after>use anyhow::Result;
use log::debug;

// because we use zydis data structures throughout our API
// make this dependency public.
// this way, our users can do `use lancelot::analysis::dis::zydis`
// and not have any version conflicts.
pub use zydis;

use crate::{arch::Arch, module::Module};

pub fn get_disassembler(module: &Module) -> Result<zydis::Decoder> {
    let mut decoder = match module.arch {
        Arch::X64 => zydis::Decoder::new(zydis::MachineMode::LONG_64, zydis::AddressWidth::_64)?,
        Arch::X32 => zydis::Decoder::new(zydis::MachineMode::LEGACY_32, zydis::AddressWidth::_32)?,
    };

    // modes described here: https://github.com/zyantific/zydis/blob/5af06d64432aaa3f6af3cd3e120eefa061b790ab/include/Zydis/Decoder.h#L55
    //
    // performance, captured empirically:
    //  - minimal mode - 8.7M instructions/second
    //  - full mode    - 4.5M instructions/second
    decoder.enable_mode(zydis::DecoderMode::MINIMAL, false)?;

    decoder.enable_mode(zydis::DecoderMode::KNC, false)?;
    decoder.enable_mode(zydis::DecoderMode::MPX, false)?;
    decoder.enable_mode(zydis::DecoderMode::CET, false)?;
    decoder.enable_mode(zydis::DecoderMode::LZCNT, false)?;
    decoder.enable_mode(zydis::DecoderMode::TZCNT, false)?;
    decoder.enable_mode(zydis::DecoderMode::WBNOINVD, false)?;
    decoder.enable_mode(zydis::DecoderMode::CLDEMOTE, false)?;

    Ok(decoder)
}

pub fn linear_disassemble<'a>(
    decoder: &'a zydis::Decoder,
    buf: &'a [u8],
) -> Box<dyn Iterator<Item = (usize, zydis::Result<Option<zydis::DecodedInstruction>>)> + 'a> {
    let mut offset = 0usize;
    let mut insn_count = 0usize;
    let iter = std::iter::from_fn(move || {
        if offset >= buf.len() {
            debug!("decoded {} instructions", insn_count);
            return None;
        }

        let insn_offset = offset;
        let insn_buf = &buf[insn_offset..];
        let insn = decoder.decode(insn_buf);

        if let Ok(Some(insn)) = &insn {
            // see discussion of linear vs thorough disassemble in this module doc for
            // call_targets. thorough is 4x more expensive, with limited
            // results.

            // linear disassembly:
            offset += insn.length as usize;

            // thorough disassembly:
            // offset += 1;

            insn_count += 1;
        } else {
            offset += 1;
        }

        Some((insn_offset, insn))
    });

    Box::new(iter)
}
<|endoftext|>


<commit_before>//! ARM compiler specific intrinsics
//!
//! # References
//!
//! - [ARM Compiler v 6.10 - armclang Reference Guide][arm_comp_ref]
//!
//! [arm_comp_ref]: https://developer.arm.com/docs/100067/0610

#[cfg(test)]
use stdsimd_test::assert_instr;

/// Inserts a breakpoint instruction.
///
/// `val` is a compile-time constant integer in range `[0, 255]`.
///
/// The breakpoint instruction inserted is:
///
/// * `BKPT` when compiling as T32,
/// * `BRK` when compiling as A32 or A64.
///
/// # Safety
///
/// If `val` is out-of-range the behavior is **undefined**.
///
/// # Note
///
/// [ARM's documentation][arm_docs] defines that `__breakpoint` accepts the
/// following values for `val`:
///
/// - `0...65535` when compiling as A32 or A64,
/// - `0...255` when compiling as T32.
///
/// The current implementation only accepts values in range `[0, 255]` - if the
/// value is out-of-range the behavior is **undefined**.
///
/// [arm_docs]: https://developer.arm.com/docs/100067/latest/compiler-specific-intrinsics/__breakpoint-intrinsic
#[cfg_attr(all(test, target_arch = "arm"), assert_instr(bkpt, val = 0))]
#[cfg_attr(all(test, target_arch = "aarch64"), assert_instr(brk, val = 0))]
#[inline(always)]
#[rustc_args_required_const(0)]
pub unsafe fn __breakpoint(val: i32) {
    // Ensure that this compiles correctly on non-arm architectures, so libstd
    // doc builds work. The proper macro will shadow this definition below.
    #[allow(unused_macros)]
    macro_rules! call {
        ($e:expr) => {()}
    }

    #[cfg(target_arch = "arm")]
    macro_rules! call {
        ($imm8:expr) => {
            asm!(concat!("BKPT ", stringify!($imm8)) : : : : "volatile")
        }
    }

    #[cfg(target_arch = "aarch64")]
    macro_rules! call {
        ($imm8:expr) => {
            asm!(concat!("BRK ", stringify!($imm8)) : : : : "volatile")
        }
    }

    // We can't `panic!` inside this intrinsic, so we can't really validate the
    // arguments here. If `val` is out-of-range this macro uses `val == 255`:
    constify_imm8!(val, call);
}
<commit_msg>reformat<commit_after>//! ARM compiler specific intrinsics
//!
//! # References
//!
//! - [ARM Compiler v 6.10 - armclang Reference Guide][arm_comp_ref]
//!
//! [arm_comp_ref]: https://developer.arm.com/docs/100067/0610

#[cfg(test)]
use stdsimd_test::assert_instr;

/// Inserts a breakpoint instruction.
///
/// `val` is a compile-time constant integer in range `[0, 255]`.
///
/// The breakpoint instruction inserted is:
///
/// * `BKPT` when compiling as T32,
/// * `BRK` when compiling as A32 or A64.
///
/// # Safety
///
/// If `val` is out-of-range the behavior is **undefined**.
///
/// # Note
///
/// [ARM's documentation][arm_docs] defines that `__breakpoint` accepts the
/// following values for `val`:
///
/// - `0...65535` when compiling as A32 or A64,
/// - `0...255` when compiling as T32.
///
/// The current implementation only accepts values in range `[0, 255]` - if the
/// value is out-of-range the behavior is **undefined**.
///
/// [arm_docs]: https://developer.arm.com/docs/100067/latest/compiler-specific-intrinsics/__breakpoint-intrinsic
#[cfg_attr(all(test, target_arch = "arm"), assert_instr(bkpt, val = 0))]
#[cfg_attr(all(test, target_arch = "aarch64"), assert_instr(brk, val = 0))]
#[inline(always)]
#[rustc_args_required_const(0)]
pub unsafe fn __breakpoint(val: i32) {
    // Ensure that this compiles correctly on non-arm architectures, so libstd
    // doc builds work. The proper macro will shadow this definition below.
    #[allow(unused_macros)]
    macro_rules! call {
        ($e:expr) => {
            ()
        };
    }

    #[cfg(target_arch = "arm")]
    macro_rules! call {
        ($imm8:expr) => {
            asm!(concat!("BKPT ", stringify!($imm8)) : : : : "volatile")
        }
    }

    #[cfg(target_arch = "aarch64")]
    macro_rules! call {
        ($imm8:expr) => {
            asm!(concat!("BRK ", stringify!($imm8)) : : : : "volatile")
        }
    }

    // We can't `panic!` inside this intrinsic, so we can't really validate the
    // arguments here. If `val` is out-of-range this macro uses `val == 255`:
    constify_imm8!(val, call);
}
<|endoftext|>


<commit_before>#![allow(non_snake_case, non_upper_case_globals)]
// #![cfg_attr(test, feature(test))]
extern crate time;

use time::precise_time_s;
use std::mem;

///////////////////////////////////////////////////////////////////////////////
// 1. Calibration
///////////////////////////////////////////////////////////////////////////////
const aalpha:f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const bbeta:f64 = 0.95;           // Discount factor

// Productivity values
const vProductivity:[f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
const mTransition:[[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                   [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                   [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                   [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                   [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const nGridCapital: usize = 17820;
const nGridProductivity: usize = 5;

fn solve(print: bool) -> f64 {
  /////////////////////////////////////////////////////////////////////////////
  // 2. Steady State
  /////////////////////////////////////////////////////////////////////////////

    let capitalSteadyState: f64 = (aalpha * bbeta).powf(1_f64 / (1_f64 - aalpha));
    let outputSteadyState: f64 = capitalSteadyState.powf(aalpha);
    let consumptionSteadyState: f64 = outputSteadyState - capitalSteadyState;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 outputSteadyState,
                 capitalSteadyState,
                 consumptionSteadyState);
    }

    let vGridCapital = (0..nGridCapital)
                           .map(|nCapital| 0.5 * capitalSteadyState + 0.00001 * (nCapital as f64))
                           .collect::<Vec<f64>>();

    // 3. Required matrices and vectors
    let mut mValueFunction = vec![[0f64; nGridProductivity]; nGridCapital];
    let mut mValueFunctionNew = mValueFunction.clone();
    let mut mPolicyFunction = mValueFunction.clone();
    let mut expectedValueFunction = mValueFunction.clone();

    // 4. We pre-build output for each point in the grid
    let mOutput = (0..nGridCapital)
                      .map(|nCapital| {
                          (0..nGridProductivity).map(|nProductivity|
        vProductivity[nProductivity]*vGridCapital[nCapital].powf(aalpha)
      ).collect::<Vec<f64>>()
                      })
                      .collect::<Vec<Vec<f64>>>();

  // 5. Main iteration
  // TODO: one could implement a macro for the multiple declarations
    let mut maxDifference = 10_f64;
    const tolerance: f64 = 0.0000001; // compiler warn: variable does not need to be mutable
    let mut capitalChoice: f64;

    let mut iteration = 0;

    while maxDifference > tolerance {
        for nProductivity in 0..nGridProductivity {
            let transitions = &mTransition[nProductivity];

            for nCapital in 0..nGridCapital {
                let value_fns = &mValueFunction[nCapital];

                let expected_value = transitions.iter()
                                                .zip(value_fns.iter())
                                                .fold(0.0f64, |acc, (transition, value_fn)| {
                    acc + (transition * value_fn)
                });

                expectedValueFunction[nCapital][nProductivity] = expected_value;
            }
        }

        for nProductivity in 0..nGridProductivity {

        // We start from previous choice (monotonicity of policy function)
            let mut gridCapitalNextPeriod = 0;

            for nCapital in 0..nGridCapital {

                let mut valueHighSoFar = -100000.0;

                for nCapitalNextPeriod in gridCapitalNextPeriod..nGridCapital {

                    let consumption = mOutput[nCapital][nProductivity] -
                                  vGridCapital[nCapitalNextPeriod];
                    let valueProvisional = (1_f64 - bbeta) * (consumption.ln()) +
                                       bbeta *
                                       expectedValueFunction[nCapitalNextPeriod][nProductivity];

                    if valueProvisional > valueHighSoFar {
                        valueHighSoFar = valueProvisional;
                        capitalChoice = vGridCapital[nCapitalNextPeriod];
                        gridCapitalNextPeriod = nCapitalNextPeriod;
                    } else {
                        break; // We break when we have achieved the max
                    }

                    mValueFunctionNew[nCapital][nProductivity] = valueHighSoFar;
                    mPolicyFunction[nCapital][nProductivity] = capitalChoice;
                }

            }

        }

        let mut diffHighSoFar = -100000.0;
        for nProductivity in 0..nGridProductivity {
            for nCapital in 0..nGridCapital {
                let diff = (mValueFunction[nCapital][nProductivity] -
                        mValueFunctionNew[nCapital][nProductivity])
                           .abs();
                if diff > diffHighSoFar {
                    diffHighSoFar = diff;
                }
            }
        }

        // swap buffers after the loop
        mem::swap(&mut mValueFunction, &mut mValueFunctionNew);

        maxDifference = diffHighSoFar;

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1) {
            println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
        }
    }

    if print {
        println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
    }

    mPolicyFunction[999][2]
}

fn main() {
    let sample = false;

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}

// FIXME: Can't use bench as the loop is too slow (takes 8mins to get samples)
// #[cfg(test)]
// mod test {
//     extern crate test;
//     use self::test::Bencher;
//     use super::solve;

//     #[bench]
//     fn bench(b: &mut Bencher) {
//         b.iter(|| {
//             // use `test::black_box` to prevent compiler optimizations from disregarding
//             // unused values
//             test::black_box(solve(false));
//         });
//     }
// }
<commit_msg>Use statically sized array for vGridCapital<commit_after>#![allow(non_snake_case, non_upper_case_globals)]
// #![cfg_attr(test, feature(test))]
extern crate time;

use time::precise_time_s;
use std::mem;

///////////////////////////////////////////////////////////////////////////////
// 1. Calibration
///////////////////////////////////////////////////////////////////////////////
const aalpha:f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const bbeta:f64 = 0.95;           // Discount factor

// Productivity values
const vProductivity:[f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
const mTransition:[[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                   [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                   [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                   [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                   [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const nGridCapital: usize = 17820;
const nGridProductivity: usize = 5;

fn solve(print: bool) -> f64 {
  /////////////////////////////////////////////////////////////////////////////
  // 2. Steady State
  /////////////////////////////////////////////////////////////////////////////

    let capitalSteadyState: f64 = (aalpha * bbeta).powf(1_f64 / (1_f64 - aalpha));
    let outputSteadyState: f64 = capitalSteadyState.powf(aalpha);
    let consumptionSteadyState: f64 = outputSteadyState - capitalSteadyState;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 outputSteadyState,
                 capitalSteadyState,
                 consumptionSteadyState);
    }

    let mut vGridCapital = [0f64; nGridCapital];

    for (i, val) in vGridCapital.iter_mut().enumerate() {
        *val = 0.5 * capitalSteadyState + 0.00001 * (i as f64)
    }

    // 3. Required matrices and vectors
    let mut mValueFunction = vec![[0f64; nGridProductivity]; nGridCapital];
    let mut mValueFunctionNew = mValueFunction.clone();
    let mut mPolicyFunction = mValueFunction.clone();
    let mut expectedValueFunction = mValueFunction.clone();

    // 4. We pre-build output for each point in the grid
    let mOutput = (0..nGridCapital)
                      .map(|nCapital| {
                          (0..nGridProductivity).map(|nProductivity|
        vProductivity[nProductivity]*vGridCapital[nCapital].powf(aalpha)
      ).collect::<Vec<f64>>()
                      })
                      .collect::<Vec<Vec<f64>>>();

  // 5. Main iteration
  // TODO: one could implement a macro for the multiple declarations
    let mut maxDifference = 10_f64;
    const tolerance: f64 = 0.0000001; // compiler warn: variable does not need to be mutable
    let mut capitalChoice: f64;

    let mut iteration = 0;

    while maxDifference > tolerance {
        for nProductivity in 0..nGridProductivity {
            let transitions = &mTransition[nProductivity];

            for nCapital in 0..nGridCapital {
                let value_fns = &mValueFunction[nCapital];

                let expected_value = transitions.iter()
                                                .zip(value_fns.iter())
                                                .fold(0.0f64, |acc, (transition, value_fn)| {
                    acc + (transition * value_fn)
                });

                expectedValueFunction[nCapital][nProductivity] = expected_value;
            }
        }

        for nProductivity in 0..nGridProductivity {

        // We start from previous choice (monotonicity of policy function)
            let mut gridCapitalNextPeriod = 0;

            for nCapital in 0..nGridCapital {

                let mut valueHighSoFar = -100000.0;

                for nCapitalNextPeriod in gridCapitalNextPeriod..nGridCapital {

                    let consumption = mOutput[nCapital][nProductivity] -
                                  vGridCapital[nCapitalNextPeriod];
                    let valueProvisional = (1_f64 - bbeta) * (consumption.ln()) +
                                       bbeta *
                                       expectedValueFunction[nCapitalNextPeriod][nProductivity];

                    if valueProvisional > valueHighSoFar {
                        valueHighSoFar = valueProvisional;
                        capitalChoice = vGridCapital[nCapitalNextPeriod];
                        gridCapitalNextPeriod = nCapitalNextPeriod;
                    } else {
                        break; // We break when we have achieved the max
                    }

                    mValueFunctionNew[nCapital][nProductivity] = valueHighSoFar;
                    mPolicyFunction[nCapital][nProductivity] = capitalChoice;
                }

            }

        }

        let mut diffHighSoFar = -100000.0;
        for nProductivity in 0..nGridProductivity {
            for nCapital in 0..nGridCapital {
                let diff = (mValueFunction[nCapital][nProductivity] -
                        mValueFunctionNew[nCapital][nProductivity])
                           .abs();
                if diff > diffHighSoFar {
                    diffHighSoFar = diff;
                }
            }
        }

        // swap buffers after the loop
        mem::swap(&mut mValueFunction, &mut mValueFunctionNew);

        maxDifference = diffHighSoFar;

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1) {
            println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
        }
    }

    if print {
        println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
    }

    mPolicyFunction[999][2]
}

fn main() {
    let sample = false;

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}

// FIXME: Can't use bench as the loop is too slow (takes 8mins to get samples)
// #[cfg(test)]
// mod test {
//     extern crate test;
//     use self::test::Bencher;
//     use super::solve;

//     #[bench]
//     fn bench(b: &mut Bencher) {
//         b.iter(|| {
//             // use `test::black_box` to prevent compiler optimizations from disregarding
//             // unused values
//             test::black_box(solve(false));
//         });
//     }
// }
<|endoftext|>


<commit_before>// Copyright 2013-2014 The gl-rs developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! # gl_generator
//!
//! `gl_generator` is an OpenGL bindings generator plugin. It defines a macro
//!  named `generate_gl_bindings!` which can be used to generate all constants
//!  and functions of a given OpenGL version.
//!
//! ## Example
//!
//! ~~~rust
//! #[phase(plugin)]
//! extern crate gl_generator;
//!
//! use self::types::*;
//!
//! generate_gl_bindings! {
//!     api: "gl",
//!     profile: "core",
//!     version: "4.5",
//!     generator: "global",
//!     extensions: [
//!         "GL_EXT_texture_filter_anisotropic",
//!     ],
//! }
//! ~~~
//!
//! ## Arguments
//!
//! Each field can be specified at most once, or not at all. If the field is not
//! specified, then a default value will be used.
//!
//! - `api`: The API to generate. Can be either `"gl"`, `"gles1"`, `"gles2"`,
//!   `"wgl"`, `"glx"`, `"egl"`. Defaults to `"gl"`.
//! - `profile`: Can be either `"core"` or `"compatibility"`. Defaults to
//!   `"core"`. `"core"` will only include all functions supported by the
//!   requested version it self, while `"compatibility"` will include all the
//!   functions from previous versions as well.
//! - `version`: The requested API version. This is usually in the form
//!   `"major.minor"`. Defaults to `"1.0"`
//! - `generator`: The type of loader to generate. Can be either `"static"`,
//!   `"global"`, or `"struct"`. Defaults to `"static"`.
//! - `extensions`: Extra extensions to include in the bindings. These are
//!   specified as a list of strings. Defaults to `[]`.
//!
//! ## About EGL
//!
//! When you generate bindings for EGL, the following platform-specific types must be declared
//!  *at the same level where you call `generate_gl_bindings`*:
//!
//! - `khronos_utime_nanoseconds_t`
//! - `khronos_uint64_t`
//! - `khronos_ssize_t`
//! - `EGLNativeDisplayType`
//! - `EGLNativePixmapType`
//! - `EGLNativeWindowType`
//! - `EGLint`
//! - `NativeDisplayType`
//! - `NativePixmapType`
//! - `NativeWindowType`
//!

#![crate_name = "gl_generator"]
#![comment = "OpenGL function loader generator."]
#![license = "ASL2"]
#![crate_type = "dylib"]

#![feature(advanced_slice_patterns)]
#![feature(macro_rules)]
#![feature(phase)]
#![feature(plugin_registrar)]
#![feature(quote)]

#[phase(plugin, link)]
extern crate log;

#[phase(plugin)]
extern crate regex_macros;

extern crate khronos_api;
extern crate regex;
extern crate rustc;
extern crate syntax;

use generators::Generator;
use registry::{Registry, Filter, Ns};
use syntax::ast::{TokenTree, TtDelimited, TtToken};
use syntax::codemap::Span;
use syntax::ext::base::{DummyResult, ExtCtxt, MacResult, MacItems};
use syntax::parse::token;

mod generators;

#[allow(dead_code)]
mod registry;

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut ::rustc::plugin::Registry) {
    reg.register_macro("generate_gl_bindings", macro_handler);
}

/// A predicate that is useful for splitting a comma separated list of tokens
fn is_comma(tt: &TokenTree) -> bool {
    match *tt {
        TtToken(_, token::Comma) => true,
        _ => false,
    }
}

/// Drops a trailing comma if it exists
fn drop_trailing_comma(tts: &[TokenTree]) -> &[TokenTree] {
    match tts {
        [tts.., TtToken(_, token::Comma)] => tts,
        tts => tts,
    }
}

/// handler for generate_gl_bindings!
fn macro_handler(ecx: &mut ExtCtxt, span: Span, tts: &[TokenTree]) -> Box<MacResult+'static> {
    // Generator options
    let mut api = None::<(Ns, &'static [u8])>;
    let mut profile = None::<String>;
    let mut version = None::<String>;
    let mut generator = None::<Box<Generator>>;
    let mut extensions = None::<Vec<String>>;

    let tts = drop_trailing_comma(tts);

    // Iterate through the comma separated
    for tts in tts.split(is_comma) {
        let mut it = tts.iter();
        let field = match it.next() {
            Some(&TtToken(_, token::Ident(ref field, _))) => field.as_str(),
            tt => {
                let span = tt.map(|tt| tt.get_span()).unwrap_or(span);
                ecx.span_err(span, "Expected a generator argument name, \
                                    either: `api`, `profile`, `version`, \
                                    `generator`, or `extensions`.");
                return DummyResult::any(span);
            },
        };
        match it.next() {
            Some(&TtToken(_, token::Colon)) => {},
            tt => {
                let span = tt.map(|tt| tt.get_span()).unwrap_or(span);
                ecx.span_err(span, "Expected `:`");
                return DummyResult::any(span);
            },
        }
        match (field, it.collect::<Vec<_>>().as_slice()) {
            ("api", tts) => {
                if api.is_some() {
                    let span = tts.head().map_or(span, |tt| tt.get_span());
                    ecx.span_err(span, "An API was already specified.");
                    return DummyResult::any(span);
                }
                api = Some(match tts {
                    [&TtToken(_, token::LitStr(api))] if api.as_str() == "gl"
                        => (registry::Gl, khronos_api::GL_XML),
                    [&TtToken(_, token::LitStr(api))] if api.as_str() == "glx"
                        => (registry::Glx, khronos_api::GLX_XML),
                    [&TtToken(_, token::LitStr(api))] if api.as_str() == "wgl"
                        => (registry::Wgl, khronos_api::WGL_XML),
                    [&TtToken(_, token::LitStr(api))] if api.as_str() == "egl"
                        => (registry::Egl, khronos_api::EGL_XML),
                    [&TtToken(_, token::LitStr(api))] if api.as_str() == "gles1"
                        => (registry::Gles1, khronos_api::GL_XML),
                    [&TtToken(_, token::LitStr(api))] if api.as_str() == "gles2"
                        => (registry::Gles2, khronos_api::GL_XML),
                    [&TtToken(span, token::LitStr(api))] => {
                        ecx.span_err(span, format!("Unknown API \"{}\"", api.as_str()).as_slice());
                        return DummyResult::any(span);
                    },
                    _ => {
                        let span = tts.head().map_or(span, |tt| tt.get_span());
                        ecx.span_err(span, "Invalid API format, expected \
                                            string.");
                        return DummyResult::any(span);
                    }
                })
            }
            ("profile", tts) => {
                if profile.is_some() {
                    ecx.span_err(tts[0].get_span(), "A profile was already specified.");
                    return DummyResult::any(span);
                }
                profile = Some(match tts {
                    [&TtToken(_, token::LitStr(profile))] if profile.as_str() == "core"
                        => "core".to_string(),
                    [&TtToken(_, token::LitStr(profile))] if profile.as_str() == "compatibility"
                        => "compatibility".to_string(),
                    [&TtToken(_, token::LitStr(profile))] => {
                        let span = tts.head().map_or(span, |tt| tt.get_span());
                        ecx.span_err(span, format!("Unknown profile \"{}\"",
                                                   profile.as_str()).as_slice());
                        return DummyResult::any(span);
                    },
                    _ => {
                        let span = tts.head().map_or(span, |tt| tt.get_span());
                        ecx.span_err(span, "Invalid profile format, expected \
                                            string.");
                        return DummyResult::any(span);
                    },
                })
            }
            ("version", tts) => {
                if version.is_some() {
                    let span = tts.head().map_or(span, |tt| tt.get_span());
                    ecx.span_err(span, "A version was already specified.");
                    return DummyResult::any(span);
                }
                version = Some(match tts {
                    [&TtToken(_, token::LitStr(version))] => {
                        version.as_str().to_string()
                    },
                    _ => {
                        let span = tts.head().map_or(span, |tt| tt.get_span());
                        ecx.span_err(span, "Invalid version format, expected \
                                            string.");
                        return DummyResult::any(span);
                    },
                });
            }
            ("generator", tts) => {
                if generator.is_some() {
                    let span = tts.head().map_or(span, |tt| tt.get_span());
                    ecx.span_err(span, "A generator was already specified.");
                    return DummyResult::any(span);
                }
                generator = Some(match tts {
                    [&TtToken(_, token::LitStr(gen))] if gen.as_str() == "global"
                        => box generators::global_gen::GlobalGenerator as Box<Generator>,
                    [&TtToken(_, token::LitStr(gen))] if gen.as_str() == "struct"
                        => box generators::struct_gen::StructGenerator as Box<Generator>,
                    [&TtToken(_, token::LitStr(gen))] if gen.as_str() == "static"
                        => box generators::static_gen::StaticGenerator as Box<Generator>,
                    [&TtToken(span, token::LitStr(gen))] => {
                        ecx.span_err(span, format!("Unknown generator \"{}\"",
                                                   gen.as_str()).as_slice());
                        return DummyResult::any(span);
                    },
                    _ => {
                        let span = tts.head().map_or(span, |tt| tt.get_span());
                        ecx.span_err(span, "Invalid generator format, expected \
                                            string.");
                        return DummyResult::any(span);
                    },
                });
            }
            ("extensions", tts) => {
                if extensions.is_some() {
                    let span = tts.head().map_or(span, |tt| tt.get_span());
                    ecx.span_err(span, "The list of extensions were already \
                                        specified.");
                    return DummyResult::any(span);
                }
                extensions = Some(match tts {
                    [&TtDelimited(span, ref delimited)] => {
                        let (ref begin, ref tts, ref close) = **delimited;
                        match (&begin.token, &close.token) {
                            (&token::LBracket, &token::RBracket) => {
                                // Drop the trailing comma if it exists
                                let tts = drop_trailing_comma(tts.as_slice());

                                // Collect the extensions, breaking early if a parse
                                // error occurs.
                                let mut failed = false;
                                let exts = tts.split(is_comma).scan((), |_, tts| {
                                    match tts {
                                        [TtToken(_, token::LitStr(ext))] => {
                                            Some(ext.as_str().to_string())
                                        },
                                        _ => {
                                            failed = true;
                                            None
                                        },
                                    }
                                }).collect();

                                // Cause an error if there is still some leftover
                                // tokens.
                                if failed {
                                    ecx.span_err(span, "Invalid extension format, \
                                                        expected string.");
                                    return DummyResult::any(span);
                                } else {
                                    exts
                                }
                            },
                            (_, _) => {
                                ecx.span_err(span, "Expected a comma separated \
                                                    list of extension strings \
                                                    delimited by square brackets: \
                                                    `[]`");
                                return DummyResult::any(span);
                            },
                        }
                    },
                    _ => {
                        ecx.span_err(span, "Expected a comma separated list of \
                                            extension strings.");
                        return DummyResult::any(span);
                    },
                });
            }
            (field, _) => {
                ecx.span_err(span, format!("Unknown field `{}`", field).as_slice());
                return DummyResult::any(span);
            }
        }
    }

    // Use default values if the parameters have not been set
    let (ns, source) = api.unwrap_or((registry::Gl, khronos_api::GL_XML));
    let extensions = extensions.unwrap_or(vec![]);
    let version = version.unwrap_or("1.0".to_string());
    let generator = generator.unwrap_or(box generators::static_gen::StaticGenerator);
    let profile = profile.unwrap_or("core".to_string());

    // Get generator field values, using default values if they have not been
    // specified
    let filter = Some(Filter {
        api: ns.to_string(),
        extensions: extensions,
        version: version,
        profile: profile,
    });

    // Generate the registry of all bindings
    let registry = {
        use std::io::BufReader;
        use std::task;

        let result = task::try(proc() {
            let reader = BufReader::new(source);
            Registry::from_xml(reader, ns, filter)
        });

        match result {
            Ok(reg) => reg,
            Err(err) => {
                use std::any::{Any, AnyRefExt};
                let err: &Any = &*err;

                match err {
                    err if err.is::<String>() => {
                        ecx.span_err(span, "error while parsing the registry");
                        ecx.span_err(span, err.downcast_ref::<String>().unwrap().as_slice());
                    },
                    err if err.is::<&'static str>() => {
                        ecx.span_err(span, "error while parsing the registry");
                        ecx.span_err(span, err.downcast_ref::<&'static str>().unwrap().as_slice());
                    },
                    _ => {
                        ecx.span_err(span, "unknown error while parsing the registry");
                    }
                }

                return DummyResult::any(span);
            }
        }
    };

    // generating the Rust bindings as a source code into "buffer"
    let items = generator.write(ecx, &registry, ns);

    MacItems::new(items.into_iter())
}
<commit_msg>Update token tree<commit_after>// Copyright 2013-2014 The gl-rs developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! # gl_generator
//!
//! `gl_generator` is an OpenGL bindings generator plugin. It defines a macro
//!  named `generate_gl_bindings!` which can be used to generate all constants
//!  and functions of a given OpenGL version.
//!
//! ## Example
//!
//! ~~~rust
//! #[phase(plugin)]
//! extern crate gl_generator;
//!
//! use self::types::*;
//!
//! generate_gl_bindings! {
//!     api: "gl",
//!     profile: "core",
//!     version: "4.5",
//!     generator: "global",
//!     extensions: [
//!         "GL_EXT_texture_filter_anisotropic",
//!     ],
//! }
//! ~~~
//!
//! ## Arguments
//!
//! Each field can be specified at most once, or not at all. If the field is not
//! specified, then a default value will be used.
//!
//! - `api`: The API to generate. Can be either `"gl"`, `"gles1"`, `"gles2"`,
//!   `"wgl"`, `"glx"`, `"egl"`. Defaults to `"gl"`.
//! - `profile`: Can be either `"core"` or `"compatibility"`. Defaults to
//!   `"core"`. `"core"` will only include all functions supported by the
//!   requested version it self, while `"compatibility"` will include all the
//!   functions from previous versions as well.
//! - `version`: The requested API version. This is usually in the form
//!   `"major.minor"`. Defaults to `"1.0"`
//! - `generator`: The type of loader to generate. Can be either `"static"`,
//!   `"global"`, or `"struct"`. Defaults to `"static"`.
//! - `extensions`: Extra extensions to include in the bindings. These are
//!   specified as a list of strings. Defaults to `[]`.
//!
//! ## About EGL
//!
//! When you generate bindings for EGL, the following platform-specific types must be declared
//!  *at the same level where you call `generate_gl_bindings`*:
//!
//! - `khronos_utime_nanoseconds_t`
//! - `khronos_uint64_t`
//! - `khronos_ssize_t`
//! - `EGLNativeDisplayType`
//! - `EGLNativePixmapType`
//! - `EGLNativeWindowType`
//! - `EGLint`
//! - `NativeDisplayType`
//! - `NativePixmapType`
//! - `NativeWindowType`
//!

#![crate_name = "gl_generator"]
#![comment = "OpenGL function loader generator."]
#![license = "ASL2"]
#![crate_type = "dylib"]

#![feature(advanced_slice_patterns)]
#![feature(macro_rules)]
#![feature(phase)]
#![feature(plugin_registrar)]
#![feature(quote)]

#[phase(plugin, link)]
extern crate log;

#[phase(plugin)]
extern crate regex_macros;

extern crate khronos_api;
extern crate regex;
extern crate rustc;
extern crate syntax;

use generators::Generator;
use registry::{Registry, Filter, Ns};
use syntax::ast::{TokenTree, TtDelimited, TtToken};
use syntax::codemap::Span;
use syntax::ext::base::{DummyResult, ExtCtxt, MacResult, MacItems};
use syntax::parse::token;

mod generators;

#[allow(dead_code)]
mod registry;

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut ::rustc::plugin::Registry) {
    reg.register_macro("generate_gl_bindings", macro_handler);
}

/// A predicate that is useful for splitting a comma separated list of tokens
fn is_comma(tt: &TokenTree) -> bool {
    match *tt {
        TtToken(_, token::Comma) => true,
        _ => false,
    }
}

/// Drops a trailing comma if it exists
fn drop_trailing_comma(tts: &[TokenTree]) -> &[TokenTree] {
    match tts {
        [tts.., TtToken(_, token::Comma)] => tts,
        tts => tts,
    }
}

/// handler for generate_gl_bindings!
fn macro_handler(ecx: &mut ExtCtxt, span: Span, tts: &[TokenTree]) -> Box<MacResult+'static> {
    // Generator options
    let mut api = None::<(Ns, &'static [u8])>;
    let mut profile = None::<String>;
    let mut version = None::<String>;
    let mut generator = None::<Box<Generator>>;
    let mut extensions = None::<Vec<String>>;

    let tts = drop_trailing_comma(tts);

    // Iterate through the comma separated
    for tts in tts.split(is_comma) {
        let mut it = tts.iter();
        let field = match it.next() {
            Some(&TtToken(_, token::Ident(ref field, _))) => field.as_str(),
            tt => {
                let span = tt.map(|tt| tt.get_span()).unwrap_or(span);
                ecx.span_err(span, "Expected a generator argument name, \
                                    either: `api`, `profile`, `version`, \
                                    `generator`, or `extensions`.");
                return DummyResult::any(span);
            },
        };
        match it.next() {
            Some(&TtToken(_, token::Colon)) => {},
            tt => {
                let span = tt.map(|tt| tt.get_span()).unwrap_or(span);
                ecx.span_err(span, "Expected `:`");
                return DummyResult::any(span);
            },
        }
        match (field, it.collect::<Vec<_>>().as_slice()) {
            ("api", tts) => {
                if api.is_some() {
                    let span = tts.head().map_or(span, |tt| tt.get_span());
                    ecx.span_err(span, "An API was already specified.");
                    return DummyResult::any(span);
                }
                api = Some(match tts {
                    [&TtToken(_, token::LitStr(api))] if api.as_str() == "gl"
                        => (registry::Gl, khronos_api::GL_XML),
                    [&TtToken(_, token::LitStr(api))] if api.as_str() == "glx"
                        => (registry::Glx, khronos_api::GLX_XML),
                    [&TtToken(_, token::LitStr(api))] if api.as_str() == "wgl"
                        => (registry::Wgl, khronos_api::WGL_XML),
                    [&TtToken(_, token::LitStr(api))] if api.as_str() == "egl"
                        => (registry::Egl, khronos_api::EGL_XML),
                    [&TtToken(_, token::LitStr(api))] if api.as_str() == "gles1"
                        => (registry::Gles1, khronos_api::GL_XML),
                    [&TtToken(_, token::LitStr(api))] if api.as_str() == "gles2"
                        => (registry::Gles2, khronos_api::GL_XML),
                    [&TtToken(span, token::LitStr(api))] => {
                        ecx.span_err(span, format!("Unknown API \"{}\"", api.as_str()).as_slice());
                        return DummyResult::any(span);
                    },
                    _ => {
                        let span = tts.head().map_or(span, |tt| tt.get_span());
                        ecx.span_err(span, "Invalid API format, expected \
                                            string.");
                        return DummyResult::any(span);
                    }
                })
            }
            ("profile", tts) => {
                if profile.is_some() {
                    ecx.span_err(tts[0].get_span(), "A profile was already specified.");
                    return DummyResult::any(span);
                }
                profile = Some(match tts {
                    [&TtToken(_, token::LitStr(profile))] if profile.as_str() == "core"
                        => "core".to_string(),
                    [&TtToken(_, token::LitStr(profile))] if profile.as_str() == "compatibility"
                        => "compatibility".to_string(),
                    [&TtToken(_, token::LitStr(profile))] => {
                        let span = tts.head().map_or(span, |tt| tt.get_span());
                        ecx.span_err(span, format!("Unknown profile \"{}\"",
                                                   profile.as_str()).as_slice());
                        return DummyResult::any(span);
                    },
                    _ => {
                        let span = tts.head().map_or(span, |tt| tt.get_span());
                        ecx.span_err(span, "Invalid profile format, expected \
                                            string.");
                        return DummyResult::any(span);
                    },
                })
            }
            ("version", tts) => {
                if version.is_some() {
                    let span = tts.head().map_or(span, |tt| tt.get_span());
                    ecx.span_err(span, "A version was already specified.");
                    return DummyResult::any(span);
                }
                version = Some(match tts {
                    [&TtToken(_, token::LitStr(version))] => {
                        version.as_str().to_string()
                    },
                    _ => {
                        let span = tts.head().map_or(span, |tt| tt.get_span());
                        ecx.span_err(span, "Invalid version format, expected \
                                            string.");
                        return DummyResult::any(span);
                    },
                });
            }
            ("generator", tts) => {
                if generator.is_some() {
                    let span = tts.head().map_or(span, |tt| tt.get_span());
                    ecx.span_err(span, "A generator was already specified.");
                    return DummyResult::any(span);
                }
                generator = Some(match tts {
                    [&TtToken(_, token::LitStr(gen))] if gen.as_str() == "global"
                        => box generators::global_gen::GlobalGenerator as Box<Generator>,
                    [&TtToken(_, token::LitStr(gen))] if gen.as_str() == "struct"
                        => box generators::struct_gen::StructGenerator as Box<Generator>,
                    [&TtToken(_, token::LitStr(gen))] if gen.as_str() == "static"
                        => box generators::static_gen::StaticGenerator as Box<Generator>,
                    [&TtToken(span, token::LitStr(gen))] => {
                        ecx.span_err(span, format!("Unknown generator \"{}\"",
                                                   gen.as_str()).as_slice());
                        return DummyResult::any(span);
                    },
                    _ => {
                        let span = tts.head().map_or(span, |tt| tt.get_span());
                        ecx.span_err(span, "Invalid generator format, expected \
                                            string.");
                        return DummyResult::any(span);
                    },
                });
            }
            ("extensions", tts) => {
                if extensions.is_some() {
                    let span = tts.head().map_or(span, |tt| tt.get_span());
                    ecx.span_err(span, "The list of extensions were already \
                                        specified.");
                    return DummyResult::any(span);
                }
                extensions = Some(match tts {
                    [&TtDelimited(span, ref delimited)]
                    if delimited.delim == token::Bracket => {
                        // Drop the trailing comma if it exists
                        let tts = drop_trailing_comma(delimited.tts.as_slice());

                        // Collect the extensions, breaking early if a parse
                        // error occurs.
                        let mut failed = false;
                        let exts = tts.split(is_comma).scan((), |_, tts| {
                            match tts {
                                [TtToken(_, token::LitStr(ext))] => {
                                    Some(ext.as_str().to_string())
                                },
                                _ => {
                                    failed = true;
                                    None
                                },
                            }
                        }).collect();

                        // Cause an error if there is still some leftover
                        // tokens.
                        if failed {
                            ecx.span_err(span, "Invalid extension format, \
                                                expected string.");
                            return DummyResult::any(span);
                        } else {
                            exts
                        }
                    },
                    _ => {
                        ecx.span_err(span, "Expected a list of extension strings.");
                        return DummyResult::any(span);
                    }
                });
            }
            (field, _) => {
                ecx.span_err(span, format!("Unknown field `{}`", field).as_slice());
                return DummyResult::any(span);
            }
        }
    }

    // Use default values if the parameters have not been set
    let (ns, source) = api.unwrap_or((registry::Gl, khronos_api::GL_XML));
    let extensions = extensions.unwrap_or(vec![]);
    let version = version.unwrap_or("1.0".to_string());
    let generator = generator.unwrap_or(box generators::static_gen::StaticGenerator);
    let profile = profile.unwrap_or("core".to_string());

    // Get generator field values, using default values if they have not been
    // specified
    let filter = Some(Filter {
        api: ns.to_string(),
        extensions: extensions,
        version: version,
        profile: profile,
    });

    // Generate the registry of all bindings
    let registry = {
        use std::io::BufReader;
        use std::task;

        let result = task::try(proc() {
            let reader = BufReader::new(source);
            Registry::from_xml(reader, ns, filter)
        });

        match result {
            Ok(reg) => reg,
            Err(err) => {
                use std::any::{Any, AnyRefExt};
                let err: &Any = &*err;

                match err {
                    err if err.is::<String>() => {
                        ecx.span_err(span, "error while parsing the registry");
                        ecx.span_err(span, err.downcast_ref::<String>().unwrap().as_slice());
                    },
                    err if err.is::<&'static str>() => {
                        ecx.span_err(span, "error while parsing the registry");
                        ecx.span_err(span, err.downcast_ref::<&'static str>().unwrap().as_slice());
                    },
                    _ => {
                        ecx.span_err(span, "unknown error while parsing the registry");
                    }
                }

                return DummyResult::any(span);
            }
        }
    };

    // generating the Rust bindings as a source code into "buffer"
    let items = generator.write(ecx, &registry, ns);

    MacItems::new(items.into_iter())
}
<|endoftext|>


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

use crate::utils;
use serde::{Deserialize, Serialize};
use std::net::SocketAddr;

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct JsonRpcConfig {
    pub address: SocketAddr,
    pub batch_size_limit: u16,
    pub page_size_limit: u16,
    pub content_length_limit: usize,
    pub tls_cert_path: Option<String>,
    pub tls_key_path: Option<String>,
    #[serde(default)]
    pub stream_rpc: StreamConfig,
}

pub const DEFAULT_JSON_RPC_ADDRESS: &str = "127.0.0.1";
pub const DEFAULT_JSON_RPC_PORT: u16 = 8080;
pub const DEFAULT_BATCH_SIZE_LIMIT: u16 = 20;
pub const DEFAULT_PAGE_SIZE_LIMIT: u16 = 1000;
pub const DEFAULT_CONTENT_LENGTH_LIMIT: usize = 32 * 1024; // 32kb

impl Default for JsonRpcConfig {
    fn default() -> JsonRpcConfig {
        JsonRpcConfig {
            address: format!("{}:{}", DEFAULT_JSON_RPC_ADDRESS, DEFAULT_JSON_RPC_PORT)
                .parse()
                .unwrap(),
            batch_size_limit: DEFAULT_BATCH_SIZE_LIMIT,
            page_size_limit: DEFAULT_PAGE_SIZE_LIMIT,
            content_length_limit: DEFAULT_CONTENT_LENGTH_LIMIT,
            tls_cert_path: None,
            tls_key_path: None,
            stream_rpc: StreamConfig::default(),
        }
    }
}

impl JsonRpcConfig {
    pub fn randomize_ports(&mut self) {
        self.address.set_port(utils::get_available_port());
    }
}

/// This API is experimental and subject to change
/// Documentation is in /json-rpc/src/stream_rpc/README.md
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct StreamConfig {
    pub enabled: bool,
    pub subscription_fetch_size: u64,
    pub send_queue_size: usize,
    pub poll_interval_ms: u64,
    pub max_poll_interval_ms: u64,
}

pub const DEFAULT_STREAM_RPC_SUBSCRIPTION_FETCH_SIZE: u64 = 100;
pub const DEFAULT_STREAM_RPC_SEND_QUEUE_SIZE: usize = 100;
pub const DEFAULT_STREAM_RPC_POLL_INTERVAL_MS: u64 = 1000;
pub const DEFAULT_STREAM_RPC_MAX_POLL_INTERVAL_MS: u64 = 5000;

impl Default for StreamConfig {
    fn default() -> StreamConfig {
        StreamConfig {
            enabled: false,
            subscription_fetch_size: DEFAULT_STREAM_RPC_SUBSCRIPTION_FETCH_SIZE,
            send_queue_size: DEFAULT_STREAM_RPC_SEND_QUEUE_SIZE,
            poll_interval_ms: DEFAULT_STREAM_RPC_POLL_INTERVAL_MS,
            max_poll_interval_ms: DEFAULT_STREAM_RPC_MAX_POLL_INTERVAL_MS,
        }
    }
}
<commit_msg>[json-rpc] increase the default content limit size<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::utils;
use serde::{Deserialize, Serialize};
use std::net::SocketAddr;

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct JsonRpcConfig {
    pub address: SocketAddr,
    pub batch_size_limit: u16,
    pub page_size_limit: u16,
    pub content_length_limit: usize,
    pub tls_cert_path: Option<String>,
    pub tls_key_path: Option<String>,
    #[serde(default)]
    pub stream_rpc: StreamConfig,
}

pub const DEFAULT_JSON_RPC_ADDRESS: &str = "127.0.0.1";
pub const DEFAULT_JSON_RPC_PORT: u16 = 8080;
pub const DEFAULT_BATCH_SIZE_LIMIT: u16 = 20;
pub const DEFAULT_PAGE_SIZE_LIMIT: u16 = 1000;
pub const DEFAULT_CONTENT_LENGTH_LIMIT: usize = 128 * 1024; // 128kb

impl Default for JsonRpcConfig {
    fn default() -> JsonRpcConfig {
        JsonRpcConfig {
            address: format!("{}:{}", DEFAULT_JSON_RPC_ADDRESS, DEFAULT_JSON_RPC_PORT)
                .parse()
                .unwrap(),
            batch_size_limit: DEFAULT_BATCH_SIZE_LIMIT,
            page_size_limit: DEFAULT_PAGE_SIZE_LIMIT,
            content_length_limit: DEFAULT_CONTENT_LENGTH_LIMIT,
            tls_cert_path: None,
            tls_key_path: None,
            stream_rpc: StreamConfig::default(),
        }
    }
}

impl JsonRpcConfig {
    pub fn randomize_ports(&mut self) {
        self.address.set_port(utils::get_available_port());
    }
}

/// This API is experimental and subject to change
/// Documentation is in /json-rpc/src/stream_rpc/README.md
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct StreamConfig {
    pub enabled: bool,
    pub subscription_fetch_size: u64,
    pub send_queue_size: usize,
    pub poll_interval_ms: u64,
    pub max_poll_interval_ms: u64,
}

pub const DEFAULT_STREAM_RPC_SUBSCRIPTION_FETCH_SIZE: u64 = 100;
pub const DEFAULT_STREAM_RPC_SEND_QUEUE_SIZE: usize = 100;
pub const DEFAULT_STREAM_RPC_POLL_INTERVAL_MS: u64 = 1000;
pub const DEFAULT_STREAM_RPC_MAX_POLL_INTERVAL_MS: u64 = 5000;

impl Default for StreamConfig {
    fn default() -> StreamConfig {
        StreamConfig {
            enabled: false,
            subscription_fetch_size: DEFAULT_STREAM_RPC_SUBSCRIPTION_FETCH_SIZE,
            send_queue_size: DEFAULT_STREAM_RPC_SEND_QUEUE_SIZE,
            poll_interval_ms: DEFAULT_STREAM_RPC_POLL_INTERVAL_MS,
            max_poll_interval_ms: DEFAULT_STREAM_RPC_MAX_POLL_INTERVAL_MS,
        }
    }
}
<|endoftext|>