<commit_before>use na::{DVector, RealField, Unit};
use std::ops::Range;

use crate::joint::{unit_constraint, JointConstraint};
use crate::math::{AngularVector, Point, Vector, DIM, SPATIAL_DIM};
use crate::object::{BodyPartHandle, BodySet, Body, BodyHandle};
use crate::solver::helper;
use crate::solver::{LinearConstraints, GenericNonlinearConstraint, IntegrationParameters,
             NonlinearConstraintGenerator};

/// A constraint that remove all be one translational degrees of freedom.
pub struct PrismaticConstraint<N: RealField, Handle: BodyHandle> {
    b1: BodyPartHandle<Handle>,
    b2: BodyPartHandle<Handle>,
    anchor1: Point<N>,
    anchor2: Point<N>,
    axis1: Unit<Vector<N>>,
    lin_impulses: Vector<N>,
    ang_impulses: AngularVector<N>,
    limit_impulse: N,
    bilateral_ground_rng: Range<usize>,
    bilateral_rng: Range<usize>,

    min_offset: Option<N>,
    max_offset: Option<N>,
}

impl<N: RealField, Handle: BodyHandle> PrismaticConstraint<N, Handle> {
    /// Create a new prismatic constraint that ensures the relative motion between the two
    /// body parts are restricted to a single translation along the `axis1` axis (expressed in
    /// the local coordinates frame of `b1`).
    pub fn new(
        b1: BodyPartHandle<Handle>,
        b2: BodyPartHandle<Handle>,
        anchor1: Point<N>,
        axis1: Unit<Vector<N>>,
        anchor2: Point<N>,
    ) -> Self {
        let min_offset = None;
        let max_offset = None;

        PrismaticConstraint {
            b1,
            b2,
            anchor1,
            anchor2,
            axis1,
            lin_impulses: Vector::zeros(),
            ang_impulses: AngularVector::zeros(),
            limit_impulse: N::zero(),
            bilateral_ground_rng: 0..0,
            bilateral_rng: 0..0,
            min_offset,
            max_offset,
        }
    }

    /// The lower limit, if any, of the relative translation (along the joint axis) of the body parts attached to this joint.
    pub fn min_offset(&self) -> Option<N> {
        self.min_offset
    }

    /// The upper limit, if any, of the relative translation (along the joint axis) of the body parts attached to this joint.
    pub fn max_offset(&self) -> Option<N> {
        self.max_offset
    }

    /// Disable the lower limit of the relative translational motion along the joint axis.
    pub fn disable_min_offset(&mut self) {
        self.min_offset = None;
    }

    /// Disable the upper limit of the relative translational motion along the joint axis.
    pub fn disable_max_offset(&mut self) {
        self.max_offset = None;
    }

    /// Enables the lower limit of the relative translational motion along the joint axis.
    pub fn enable_min_offset(&mut self, limit: N) {
        self.min_offset = Some(limit);
        self.assert_limits();
    }

    /// Disable the lower limit of the relative translational motion along the joint axis.
    pub fn enable_max_offset(&mut self, limit: N) {
        self.max_offset = Some(limit);
        self.assert_limits();
    }

    fn assert_limits(&self) {
        if let (Some(min_offset), Some(max_offset)) = (self.min_offset, self.max_offset) {
            assert!(
                min_offset <= max_offset,
                "RevoluteJoint constraint limits: the min angle must be larger than (or equal to) the max angle.");
        }
    }
}

impl<N: RealField, Handle: BodyHandle, Bodies: BodySet<N, Handle = Handle>> JointConstraint<N, Bodies> for PrismaticConstraint<N, Handle> {
    fn num_velocity_constraints(&self) -> usize {
        (SPATIAL_DIM - 1) + 2
    }

    fn anchors(&self) -> (BodyPartHandle<Handle>, BodyPartHandle<Handle>) {
        (self.b1, self.b2)
    }

    fn velocity_constraints(
        &mut self,
        _: &IntegrationParameters<N>,
        bodies: &Bodies,
        ext_vels: &DVector<N>,
        ground_j_id: &mut usize,
        j_id: &mut usize,
        jacobians: &mut [N],
        constraints: &mut LinearConstraints<N, usize>,
    ) {
        let body1 = try_ret!(bodies.get(self.b1.0));
        let body2 = try_ret!(bodies.get(self.b2.0));
        let part1 = try_ret!(body1.part(self.b1.1));
        let part2 = try_ret!(body2.part(self.b2.1));

        /*
         *
         * Joint constraints.
         *
         */
        let pos1 = body1.position_at_material_point(part1, &self.anchor1);
        let pos2 = body2.position_at_material_point(part2, &self.anchor2);

        let anchor1 = Point::from(pos1.translation.vector);
        let anchor2 = Point::from(pos2.translation.vector);

        let assembly_id1 = body1.companion_id();
        let assembly_id2 = body2.companion_id();

        let first_bilateral_ground = constraints.bilateral_ground.len();
        let first_bilateral = constraints.bilateral.len();

        let axis = pos1 * self.axis1;

        helper::restrict_relative_linear_velocity_to_axis(
            body1,
            part1,
            self.b1,
            body2,
            part2,
            self.b2,
            assembly_id1,
            assembly_id2,
            &anchor1,
            &anchor2,
            &axis,
            ext_vels,
            self.lin_impulses.as_slice(),
            0,
            ground_j_id,
            j_id,
            jacobians,
            constraints,
        );

        helper::cancel_relative_angular_velocity(
            body1,
            part1,
            self.b1,
            body2,
            part2,
            self.b2,
            assembly_id1,
            assembly_id2,
            &anchor1,
            &anchor2,
            ext_vels,
            &self.ang_impulses,
            DIM - 1,
            ground_j_id,
            j_id,
            jacobians,
            constraints,
        );

        /*
         *
         * Limit constraints.
         *
         */
        unit_constraint::build_linear_limits_velocity_constraint(
            body1,
            part1,
            self.b1,
            body2,
            part2,
            self.b2,
            assembly_id1,
            assembly_id2,
            &anchor1,
            &anchor2,
            &axis,
            self.min_offset,
            self.max_offset,
            ext_vels,
            self.limit_impulse,
            SPATIAL_DIM - 1,
            ground_j_id,
            j_id,
            jacobians,
            constraints,
        );

        self.bilateral_ground_rng =
            first_bilateral_ground..constraints.bilateral_ground.len();
        self.bilateral_rng = first_bilateral..constraints.bilateral.len();
    }

    fn cache_impulses(&mut self, constraints: &LinearConstraints<N, usize>) {
        for c in &constraints.bilateral_ground[self.bilateral_ground_rng.clone()] {
            if c.impulse_id < DIM - 1 {
                self.lin_impulses[c.impulse_id] = c.impulse;
            } else if c.impulse_id < SPATIAL_DIM - 1 {
                self.ang_impulses[c.impulse_id - DIM + 1] = c.impulse;
            } else {
                self.limit_impulse = c.impulse
            }
        }

        for c in &constraints.bilateral[self.bilateral_rng.clone()] {
            if c.impulse_id < DIM - 1 {
                self.lin_impulses[c.impulse_id] = c.impulse;
            } else if c.impulse_id < SPATIAL_DIM - 1 {
                self.ang_impulses[c.impulse_id - DIM + 1] = c.impulse;
            } else {
                self.limit_impulse = c.impulse
            }
        }
    }
}

impl<N: RealField, Handle: BodyHandle, Bodies: BodySet<N, Handle = Handle>> NonlinearConstraintGenerator<N, Bodies> for PrismaticConstraint<N, Handle> {
    fn num_position_constraints(&self, bodies: &Bodies) -> usize {
        // FIXME: calling this at each iteration of the non-linear resolution is costly.
        if self.is_active(bodies) {
            if self.min_offset.is_some() || self.max_offset.is_some() {
                3
            } else {
                2
            }
        } else {
            0
        }
    }

    fn position_constraint(
        &self,
        parameters: &IntegrationParameters<N>,
        i: usize,
        bodies: &mut Bodies,
        jacobians: &mut [N],
    ) -> Option<GenericNonlinearConstraint<N, Handle>> {
        let body1 = bodies.get(self.b1.0)?;
        let body2 = bodies.get(self.b2.0)?;
        let part1 = body1.part(self.b1.1)?;
        let part2 = body2.part(self.b2.1)?;

        let pos1 = body1.position_at_material_point(part1, &self.anchor1);
        let pos2 = body2.position_at_material_point(part2, &self.anchor2);

        let anchor1 = Point::from(pos1.translation.vector);
        let anchor2 = Point::from(pos2.translation.vector);

        if i == 0 {
            return helper::cancel_relative_rotation(
                parameters,
                body1,
                part1,
                self.b1,
                body2,
                part2,
                self.b2,
                &anchor1,
                &anchor2,
                &pos1.rotation,
                &pos2.rotation,
                jacobians,
            );
        } else if i == 1 {
            let axis = pos1 * self.axis1;

            return helper::project_anchor_to_axis(
                parameters,
                body1,
                part1,
                self.b1,
                body2,
                part2,
                self.b2,
                &anchor1,
                &anchor2,
                &axis,
                jacobians,
            );
        } else if i == 2 {
            let axis = pos1 * self.axis1;

            return unit_constraint::build_linear_limits_position_constraint(
                parameters,
                body1,
                part1,
                self.b1,
                body2,
                part2,
                self.b2,
                &anchor1,
                &anchor2,
                &axis,
                self.min_offset,
                self.max_offset,
                jacobians,
            );
        }

        return None;
    }
}
<commit_msg>Fix underflow panic in debug mode.<commit_after>use na::{DVector, RealField, Unit};
use std::ops::Range;

use crate::joint::{unit_constraint, JointConstraint};
use crate::math::{AngularVector, Point, Vector, DIM, SPATIAL_DIM};
use crate::object::{BodyPartHandle, BodySet, Body, BodyHandle};
use crate::solver::helper;
use crate::solver::{LinearConstraints, GenericNonlinearConstraint, IntegrationParameters,
             NonlinearConstraintGenerator};

/// A constraint that remove all be one translational degrees of freedom.
pub struct PrismaticConstraint<N: RealField, Handle: BodyHandle> {
    b1: BodyPartHandle<Handle>,
    b2: BodyPartHandle<Handle>,
    anchor1: Point<N>,
    anchor2: Point<N>,
    axis1: Unit<Vector<N>>,
    lin_impulses: Vector<N>,
    ang_impulses: AngularVector<N>,
    limit_impulse: N,
    bilateral_ground_rng: Range<usize>,
    bilateral_rng: Range<usize>,

    min_offset: Option<N>,
    max_offset: Option<N>,
}

impl<N: RealField, Handle: BodyHandle> PrismaticConstraint<N, Handle> {
    /// Create a new prismatic constraint that ensures the relative motion between the two
    /// body parts are restricted to a single translation along the `axis1` axis (expressed in
    /// the local coordinates frame of `b1`).
    pub fn new(
        b1: BodyPartHandle<Handle>,
        b2: BodyPartHandle<Handle>,
        anchor1: Point<N>,
        axis1: Unit<Vector<N>>,
        anchor2: Point<N>,
    ) -> Self {
        let min_offset = None;
        let max_offset = None;

        PrismaticConstraint {
            b1,
            b2,
            anchor1,
            anchor2,
            axis1,
            lin_impulses: Vector::zeros(),
            ang_impulses: AngularVector::zeros(),
            limit_impulse: N::zero(),
            bilateral_ground_rng: 0..0,
            bilateral_rng: 0..0,
            min_offset,
            max_offset,
        }
    }

    /// The lower limit, if any, of the relative translation (along the joint axis) of the body parts attached to this joint.
    pub fn min_offset(&self) -> Option<N> {
        self.min_offset
    }

    /// The upper limit, if any, of the relative translation (along the joint axis) of the body parts attached to this joint.
    pub fn max_offset(&self) -> Option<N> {
        self.max_offset
    }

    /// Disable the lower limit of the relative translational motion along the joint axis.
    pub fn disable_min_offset(&mut self) {
        self.min_offset = None;
    }

    /// Disable the upper limit of the relative translational motion along the joint axis.
    pub fn disable_max_offset(&mut self) {
        self.max_offset = None;
    }

    /// Enables the lower limit of the relative translational motion along the joint axis.
    pub fn enable_min_offset(&mut self, limit: N) {
        self.min_offset = Some(limit);
        self.assert_limits();
    }

    /// Disable the lower limit of the relative translational motion along the joint axis.
    pub fn enable_max_offset(&mut self, limit: N) {
        self.max_offset = Some(limit);
        self.assert_limits();
    }

    fn assert_limits(&self) {
        if let (Some(min_offset), Some(max_offset)) = (self.min_offset, self.max_offset) {
            assert!(
                min_offset <= max_offset,
                "RevoluteJoint constraint limits: the min angle must be larger than (or equal to) the max angle.");
        }
    }
}

impl<N: RealField, Handle: BodyHandle, Bodies: BodySet<N, Handle = Handle>> JointConstraint<N, Bodies> for PrismaticConstraint<N, Handle> {
    fn num_velocity_constraints(&self) -> usize {
        (SPATIAL_DIM - 1) + 2
    }

    fn anchors(&self) -> (BodyPartHandle<Handle>, BodyPartHandle<Handle>) {
        (self.b1, self.b2)
    }

    fn velocity_constraints(
        &mut self,
        _: &IntegrationParameters<N>,
        bodies: &Bodies,
        ext_vels: &DVector<N>,
        ground_j_id: &mut usize,
        j_id: &mut usize,
        jacobians: &mut [N],
        constraints: &mut LinearConstraints<N, usize>,
    ) {
        let body1 = try_ret!(bodies.get(self.b1.0));
        let body2 = try_ret!(bodies.get(self.b2.0));
        let part1 = try_ret!(body1.part(self.b1.1));
        let part2 = try_ret!(body2.part(self.b2.1));

        /*
         *
         * Joint constraints.
         *
         */
        let pos1 = body1.position_at_material_point(part1, &self.anchor1);
        let pos2 = body2.position_at_material_point(part2, &self.anchor2);

        let anchor1 = Point::from(pos1.translation.vector);
        let anchor2 = Point::from(pos2.translation.vector);

        let assembly_id1 = body1.companion_id();
        let assembly_id2 = body2.companion_id();

        let first_bilateral_ground = constraints.bilateral_ground.len();
        let first_bilateral = constraints.bilateral.len();

        let axis = pos1 * self.axis1;

        helper::restrict_relative_linear_velocity_to_axis(
            body1,
            part1,
            self.b1,
            body2,
            part2,
            self.b2,
            assembly_id1,
            assembly_id2,
            &anchor1,
            &anchor2,
            &axis,
            ext_vels,
            self.lin_impulses.as_slice(),
            0,
            ground_j_id,
            j_id,
            jacobians,
            constraints,
        );

        helper::cancel_relative_angular_velocity(
            body1,
            part1,
            self.b1,
            body2,
            part2,
            self.b2,
            assembly_id1,
            assembly_id2,
            &anchor1,
            &anchor2,
            ext_vels,
            &self.ang_impulses,
            DIM - 1,
            ground_j_id,
            j_id,
            jacobians,
            constraints,
        );

        /*
         *
         * Limit constraints.
         *
         */
        unit_constraint::build_linear_limits_velocity_constraint(
            body1,
            part1,
            self.b1,
            body2,
            part2,
            self.b2,
            assembly_id1,
            assembly_id2,
            &anchor1,
            &anchor2,
            &axis,
            self.min_offset,
            self.max_offset,
            ext_vels,
            self.limit_impulse,
            SPATIAL_DIM - 1,
            ground_j_id,
            j_id,
            jacobians,
            constraints,
        );

        self.bilateral_ground_rng =
            first_bilateral_ground..constraints.bilateral_ground.len();
        self.bilateral_rng = first_bilateral..constraints.bilateral.len();
    }

    fn cache_impulses(&mut self, constraints: &LinearConstraints<N, usize>) {
        for c in &constraints.bilateral_ground[self.bilateral_ground_rng.clone()] {
            if c.impulse_id < DIM - 1 {
                self.lin_impulses[c.impulse_id] = c.impulse;
            } else if c.impulse_id < SPATIAL_DIM - 1 {
                self.ang_impulses[c.impulse_id + 1 - DIM] = c.impulse;
            } else {
                self.limit_impulse = c.impulse
            }
        }

        for c in &constraints.bilateral[self.bilateral_rng.clone()] {
            if c.impulse_id < DIM - 1 {
                self.lin_impulses[c.impulse_id] = c.impulse;
            } else if c.impulse_id < SPATIAL_DIM - 1 {
                self.ang_impulses[c.impulse_id + 1 - DIM] = c.impulse;
            } else {
                self.limit_impulse = c.impulse
            }
        }
    }
}

impl<N: RealField, Handle: BodyHandle, Bodies: BodySet<N, Handle = Handle>> NonlinearConstraintGenerator<N, Bodies> for PrismaticConstraint<N, Handle> {
    fn num_position_constraints(&self, bodies: &Bodies) -> usize {
        // FIXME: calling this at each iteration of the non-linear resolution is costly.
        if self.is_active(bodies) {
            if self.min_offset.is_some() || self.max_offset.is_some() {
                3
            } else {
                2
            }
        } else {
            0
        }
    }

    fn position_constraint(
        &self,
        parameters: &IntegrationParameters<N>,
        i: usize,
        bodies: &mut Bodies,
        jacobians: &mut [N],
    ) -> Option<GenericNonlinearConstraint<N, Handle>> {
        let body1 = bodies.get(self.b1.0)?;
        let body2 = bodies.get(self.b2.0)?;
        let part1 = body1.part(self.b1.1)?;
        let part2 = body2.part(self.b2.1)?;

        let pos1 = body1.position_at_material_point(part1, &self.anchor1);
        let pos2 = body2.position_at_material_point(part2, &self.anchor2);

        let anchor1 = Point::from(pos1.translation.vector);
        let anchor2 = Point::from(pos2.translation.vector);

        if i == 0 {
            return helper::cancel_relative_rotation(
                parameters,
                body1,
                part1,
                self.b1,
                body2,
                part2,
                self.b2,
                &anchor1,
                &anchor2,
                &pos1.rotation,
                &pos2.rotation,
                jacobians,
            );
        } else if i == 1 {
            let axis = pos1 * self.axis1;

            return helper::project_anchor_to_axis(
                parameters,
                body1,
                part1,
                self.b1,
                body2,
                part2,
                self.b2,
                &anchor1,
                &anchor2,
                &axis,
                jacobians,
            );
        } else if i == 2 {
            let axis = pos1 * self.axis1;

            return unit_constraint::build_linear_limits_position_constraint(
                parameters,
                body1,
                part1,
                self.b1,
                body2,
                part2,
                self.b2,
                &anchor1,
                &anchor2,
                &axis,
                self.min_offset,
                self.max_offset,
                jacobians,
            );
        }

        return None;
    }
}
<|endoftext|>


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

// The classification code for the x86_64 ABI is taken from the clay language
// https://github.com/jckarter/clay/blob/master/compiler/src/externals.cpp

use abi::{ArgType, CastTarget, FnType, LayoutExt, Reg, RegKind};
use context::CodegenCx;

use rustc::ty::layout::{self, TyLayout, Size};

/// Classification of "eightbyte" components.
// NB: the order of the variants is from general to specific,
// such that `unify(a, b)` is the "smaller" of `a` and `b`.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
enum Class {
    Int,
    Sse,
    SseUp
}

#[derive(Clone, Copy, Debug)]
struct Memory;

// Currently supported vector size (AVX-512).
const LARGEST_VECTOR_SIZE: usize = 512;
const MAX_EIGHTBYTES: usize = LARGEST_VECTOR_SIZE / 64;

fn classify_arg<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>, arg: &ArgType<'tcx>)
                          -> Result<[Option<Class>; MAX_EIGHTBYTES], Memory> {
    fn classify<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>,
                          layout: TyLayout<'tcx>,
                          cls: &mut [Option<Class>],
                          off: Size)
                          -> Result<(), Memory> {
        if !off.is_abi_aligned(layout.align) {
            if !layout.is_zst() {
                return Err(Memory);
            }
            return Ok(());
        }

        let mut c = match layout.abi {
            layout::Abi::Uninhabited => return Ok(()),

            layout::Abi::Scalar(ref scalar) => {
                match scalar.value {
                    layout::Int(..) |
                    layout::Pointer => Class::Int,
                    layout::F32 |
                    layout::F64 => Class::Sse
                }
            }

            layout::Abi::Vector { .. } => Class::Sse,

            layout::Abi::ScalarPair(..) |
            layout::Abi::Aggregate { .. } => {
                match layout.variants {
                    layout::Variants::Single { .. } => {
                        for i in 0..layout.fields.count() {
                            let field_off = off + layout.fields.offset(i);
                            classify(cx, layout.field(cx, i), cls, field_off)?;
                        }
                        return Ok(());
                    }
                    layout::Variants::Tagged { .. } |
                    layout::Variants::NicheFilling { .. } => return Err(Memory),
                }
            }

        };

        // Fill in `cls` for scalars (Int/Sse) and vectors (Sse).
        let first = (off.bytes() / 8) as usize;
        let last = ((off.bytes() + layout.size.bytes() - 1) / 8) as usize;
        for cls in &mut cls[first..=last] {
            *cls = Some(cls.map_or(c, |old| old.min(c)));

            // Everything after the first Sse "eightbyte"
            // component is the upper half of a register.
            if c == Class::Sse {
                c = Class::SseUp;
            }
        }

        Ok(())
    }

    let n = ((arg.layout.size.bytes() + 7) / 8) as usize;
    if n > MAX_EIGHTBYTES {
        return Err(Memory);
    }

    let mut cls = [None; MAX_EIGHTBYTES];
    classify(cx, arg.layout, &mut cls, Size::from_bytes(0))?;
    if n > 2 {
        if cls[0] != Some(Class::Sse) {
            return Err(Memory);
        }
        if cls[1..n].iter().any(|&c| c != Some(Class::SseUp)) {
            return Err(Memory);
        }
    } else {
        let mut i = 0;
        while i < n {
            if cls[i] == Some(Class::SseUp) {
                cls[i] = Some(Class::Sse);
            } else if cls[i] == Some(Class::Sse) {
                i += 1;
                while i != n && cls[i] == Some(Class::SseUp) { i += 1; }
            } else {
                i += 1;
            }
        }
    }

    Ok(cls)
}

fn reg_component(cls: &[Option<Class>], i: &mut usize, size: Size) -> Option<Reg> {
    if *i >= cls.len() {
        return None;
    }

    match cls[*i] {
        None => None,
        Some(Class::Int) => {
            *i += 1;
            Some(match size.bytes() {
                1 => Reg::i8(),
                2 => Reg::i16(),
                3 |
                4 => Reg::i32(),
                _ => Reg::i64()
            })
        }
        Some(Class::Sse) => {
            let vec_len = 1 + cls[*i+1..].iter()
                .take_while(|&&c| c == Some(Class::SseUp))
                .count();
            *i += vec_len;
            Some(if vec_len == 1 {
                match size.bytes() {
                    4 => Reg::f32(),
                    _ => Reg::f64()
                }
            } else {
                Reg {
                    kind: RegKind::Vector,
                    size: Size::from_bytes(8) * (vec_len as u64)
                }
            })
        }
        Some(c) => bug!("reg_component: unhandled class {:?}", c)
    }
}

fn cast_target(cls: &[Option<Class>], size: Size) -> CastTarget {
    let mut i = 0;
    let lo = reg_component(cls, &mut i, size).unwrap();
    let offset = Size::from_bytes(8) * (i as u64);
    let mut target = CastTarget::from(lo);
    if size > offset {
        if let Some(hi) = reg_component(cls, &mut i, size - offset) {
            target = CastTarget::Pair(lo, hi);
        }
    }
    assert_eq!(reg_component(cls, &mut i, Size::from_bytes(0)), None);
    target
}

pub fn compute_abi_info<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>, fty: &mut FnType<'tcx>) {
    let mut int_regs = 6; // RDI, RSI, RDX, RCX, R8, R9
    let mut sse_regs = 8; // XMM0-7

    let mut x86_64_ty = |arg: &mut ArgType<'tcx>, is_arg: bool| {
        let cls = classify_arg(cx, arg);

        let mut needed_int = 0;
        let mut needed_sse = 0;
        let in_mem = match cls {
            Err(Memory) => true,
            Ok(ref cls) if is_arg => {
                for &c in cls {
                    match c {
                        Some(Class::Int) => needed_int += 1,
                        Some(Class::Sse) => needed_sse += 1,
                        _ => {}
                    }
                }
                arg.layout.is_aggregate() &&
                    (int_regs < needed_int || sse_regs < needed_sse)
            }
            Ok(_) => false
        };

        if in_mem {
            if is_arg {
                arg.make_indirect_byval();
            } else {
                // `sret` parameter thus one less integer register available
                arg.make_indirect();
                int_regs -= 1;
            }
        } else {
            // split into sized chunks passed individually
            int_regs -= needed_int;
            sse_regs -= needed_sse;

            if arg.layout.is_aggregate() {
                let size = arg.layout.size;
                arg.cast_to(cast_target(cls.as_ref().unwrap(), size))
            } else {
                arg.extend_integer_width_to(32);
            }
        }
    };

    if !fty.ret.is_ignore() {
        x86_64_ty(&mut fty.ret, false);
    }

    for arg in &mut fty.args {
        if arg.is_ignore() { continue; }
        x86_64_ty(arg, true);
    }
}
<commit_msg>Rollup merge of #47626 - eddyb:one-less-unwrap, r=nagisa<commit_after>// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// The classification code for the x86_64 ABI is taken from the clay language
// https://github.com/jckarter/clay/blob/master/compiler/src/externals.cpp

use abi::{ArgType, CastTarget, FnType, LayoutExt, Reg, RegKind};
use context::CodegenCx;

use rustc::ty::layout::{self, TyLayout, Size};

/// Classification of "eightbyte" components.
// NB: the order of the variants is from general to specific,
// such that `unify(a, b)` is the "smaller" of `a` and `b`.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
enum Class {
    Int,
    Sse,
    SseUp
}

#[derive(Clone, Copy, Debug)]
struct Memory;

// Currently supported vector size (AVX-512).
const LARGEST_VECTOR_SIZE: usize = 512;
const MAX_EIGHTBYTES: usize = LARGEST_VECTOR_SIZE / 64;

fn classify_arg<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>, arg: &ArgType<'tcx>)
                          -> Result<[Option<Class>; MAX_EIGHTBYTES], Memory> {
    fn classify<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>,
                          layout: TyLayout<'tcx>,
                          cls: &mut [Option<Class>],
                          off: Size)
                          -> Result<(), Memory> {
        if !off.is_abi_aligned(layout.align) {
            if !layout.is_zst() {
                return Err(Memory);
            }
            return Ok(());
        }

        let mut c = match layout.abi {
            layout::Abi::Uninhabited => return Ok(()),

            layout::Abi::Scalar(ref scalar) => {
                match scalar.value {
                    layout::Int(..) |
                    layout::Pointer => Class::Int,
                    layout::F32 |
                    layout::F64 => Class::Sse
                }
            }

            layout::Abi::Vector { .. } => Class::Sse,

            layout::Abi::ScalarPair(..) |
            layout::Abi::Aggregate { .. } => {
                match layout.variants {
                    layout::Variants::Single { .. } => {
                        for i in 0..layout.fields.count() {
                            let field_off = off + layout.fields.offset(i);
                            classify(cx, layout.field(cx, i), cls, field_off)?;
                        }
                        return Ok(());
                    }
                    layout::Variants::Tagged { .. } |
                    layout::Variants::NicheFilling { .. } => return Err(Memory),
                }
            }

        };

        // Fill in `cls` for scalars (Int/Sse) and vectors (Sse).
        let first = (off.bytes() / 8) as usize;
        let last = ((off.bytes() + layout.size.bytes() - 1) / 8) as usize;
        for cls in &mut cls[first..=last] {
            *cls = Some(cls.map_or(c, |old| old.min(c)));

            // Everything after the first Sse "eightbyte"
            // component is the upper half of a register.
            if c == Class::Sse {
                c = Class::SseUp;
            }
        }

        Ok(())
    }

    let n = ((arg.layout.size.bytes() + 7) / 8) as usize;
    if n > MAX_EIGHTBYTES {
        return Err(Memory);
    }

    let mut cls = [None; MAX_EIGHTBYTES];
    classify(cx, arg.layout, &mut cls, Size::from_bytes(0))?;
    if n > 2 {
        if cls[0] != Some(Class::Sse) {
            return Err(Memory);
        }
        if cls[1..n].iter().any(|&c| c != Some(Class::SseUp)) {
            return Err(Memory);
        }
    } else {
        let mut i = 0;
        while i < n {
            if cls[i] == Some(Class::SseUp) {
                cls[i] = Some(Class::Sse);
            } else if cls[i] == Some(Class::Sse) {
                i += 1;
                while i != n && cls[i] == Some(Class::SseUp) { i += 1; }
            } else {
                i += 1;
            }
        }
    }

    Ok(cls)
}

fn reg_component(cls: &[Option<Class>], i: &mut usize, size: Size) -> Option<Reg> {
    if *i >= cls.len() {
        return None;
    }

    match cls[*i] {
        None => None,
        Some(Class::Int) => {
            *i += 1;
            Some(match size.bytes() {
                1 => Reg::i8(),
                2 => Reg::i16(),
                3 |
                4 => Reg::i32(),
                _ => Reg::i64()
            })
        }
        Some(Class::Sse) => {
            let vec_len = 1 + cls[*i+1..].iter()
                .take_while(|&&c| c == Some(Class::SseUp))
                .count();
            *i += vec_len;
            Some(if vec_len == 1 {
                match size.bytes() {
                    4 => Reg::f32(),
                    _ => Reg::f64()
                }
            } else {
                Reg {
                    kind: RegKind::Vector,
                    size: Size::from_bytes(8) * (vec_len as u64)
                }
            })
        }
        Some(c) => bug!("reg_component: unhandled class {:?}", c)
    }
}

fn cast_target(cls: &[Option<Class>], size: Size) -> CastTarget {
    let mut i = 0;
    let lo = reg_component(cls, &mut i, size).unwrap();
    let offset = Size::from_bytes(8) * (i as u64);
    let mut target = CastTarget::from(lo);
    if size > offset {
        if let Some(hi) = reg_component(cls, &mut i, size - offset) {
            target = CastTarget::Pair(lo, hi);
        }
    }
    assert_eq!(reg_component(cls, &mut i, Size::from_bytes(0)), None);
    target
}

pub fn compute_abi_info<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>, fty: &mut FnType<'tcx>) {
    let mut int_regs = 6; // RDI, RSI, RDX, RCX, R8, R9
    let mut sse_regs = 8; // XMM0-7

    let mut x86_64_ty = |arg: &mut ArgType<'tcx>, is_arg: bool| {
        let mut cls_or_mem = classify_arg(cx, arg);

        let mut needed_int = 0;
        let mut needed_sse = 0;
        if is_arg {
            if let Ok(cls) = cls_or_mem {
                for &c in &cls {
                    match c {
                        Some(Class::Int) => needed_int += 1,
                        Some(Class::Sse) => needed_sse += 1,
                        _ => {}
                    }
                }
                if arg.layout.is_aggregate() {
                    if int_regs < needed_int || sse_regs < needed_sse {
                        cls_or_mem = Err(Memory);
                    }
                }
            }
        }

        match cls_or_mem {
            Err(Memory) => {
                if is_arg {
                    arg.make_indirect_byval();
                } else {
                    // `sret` parameter thus one less integer register available
                    arg.make_indirect();
                    int_regs -= 1;
                }
            }
            Ok(ref cls) => {
                // split into sized chunks passed individually
                int_regs -= needed_int;
                sse_regs -= needed_sse;

                if arg.layout.is_aggregate() {
                    let size = arg.layout.size;
                    arg.cast_to(cast_target(cls, size))
                } else {
                    arg.extend_integer_width_to(32);
                }
            }
        }
    };

    if !fty.ret.is_ignore() {
        x86_64_ty(&mut fty.ret, false);
    }

    for arg in &mut fty.args {
        if arg.is_ignore() { continue; }
        x86_64_ty(arg, true);
    }
}
<|endoftext|>


<commit_before>#![warn(clippy::needless_pass_by_value)]
#![allow(
    dead_code,
    clippy::single_match,
    clippy::redundant_pattern_matching,
    clippy::many_single_char_names,
    clippy::option_option
)]

use std::borrow::Borrow;
use std::collections::HashSet;
use std::convert::AsRef;

// `v` should be warned
// `w`, `x` and `y` are allowed (moved or mutated)
fn foo<T: Default>(v: Vec<T>, w: Vec<T>, mut x: Vec<T>, y: Vec<T>) -> Vec<T> {
    assert_eq!(v.len(), 42);

    consume(w);

    x.push(T::default());

    y
}

fn consume<T>(_: T) {}

struct Wrapper(String);

fn bar(x: String, y: Wrapper) {
    assert_eq!(x.len(), 42);
    assert_eq!(y.0.len(), 42);
}

// V implements `Borrow<V>`, but should be warned correctly
fn test_borrow_trait<T: Borrow<str>, U: AsRef<str>, V>(t: T, u: U, v: V) {
    println!("{}", t.borrow());
    println!("{}", u.as_ref());
    consume(&v);
}

// ok
fn test_fn<F: Fn(i32) -> i32>(f: F) {
    f(1);
}

// x should be warned, but y is ok
fn test_match(x: Option<Option<String>>, y: Option<Option<String>>) {
    match x {
        Some(Some(_)) => 1, // not moved
        _ => 0,
    };

    match y {
        Some(Some(s)) => consume(s), // moved
        _ => (),
    };
}

// x and y should be warned, but z is ok
fn test_destructure(x: Wrapper, y: Wrapper, z: Wrapper) {
    let Wrapper(s) = z; // moved
    let Wrapper(ref t) = y; // not moved
    let Wrapper(_) = y; // still not moved

    assert_eq!(x.0.len(), s.len());
    println!("{}", t);
}

trait Foo {}

// `S: Serialize` is allowed to be passed by value, since a caller can pass `&S` instead
trait Serialize {}
impl<'a, T> Serialize for &'a T where T: Serialize {}
impl Serialize for i32 {}

fn test_blanket_ref<T: Foo, S: Serialize>(_foo: T, _serializable: S) {}

fn issue_2114(s: String, t: String, u: Vec<i32>, v: Vec<i32>) {
    s.capacity();
    let _ = t.clone();
    u.capacity();
    let _ = v.clone();
}

struct S<T, U>(T, U);

impl<T: Serialize, U> S<T, U> {
    fn foo(
        self,
        // taking `self` by value is always allowed
        s: String,
        t: String,
    ) -> usize {
        s.len() + t.capacity()
    }

    fn bar(_t: T // Ok, since `&T: Serialize` too
    ) {
    }

    fn baz(&self, _u: U, _s: Self) {}
}

trait FalsePositive {
    fn visit_str(s: &str);
    fn visit_string(s: String) {
        Self::visit_str(&s);
    }
}

// shouldn't warn on extern funcs
extern "C" fn ext(x: String) -> usize {
    x.len()
}

// whitelist RangeArgument
fn range<T: ::std::ops::RangeBounds<usize>>(range: T) {
    let _ = range.start_bound();
}

struct CopyWrapper(u32);

fn bar_copy(x: u32, y: CopyWrapper) {
    assert_eq!(x, 42);
    assert_eq!(y.0, 42);
}

// x and y should be warned, but z is ok
fn test_destructure_copy(x: CopyWrapper, y: CopyWrapper, z: CopyWrapper) {
    let CopyWrapper(s) = z; // moved
    let CopyWrapper(ref t) = y; // not moved
    let CopyWrapper(_) = y; // still not moved

    assert_eq!(x.0, s);
    println!("{}", t);
}

// The following 3 lines should not cause an ICE. See #2831
trait Bar<'a, A> {}
impl<'b, T> Bar<'b, T> for T {}
fn some_fun<'b, S: Bar<'b, ()>>(_item: S) {}

// Also this should not cause an ICE. See #2831
trait Club<'a, A> {}
impl<T> Club<'static, T> for T {}
fn more_fun(_item: impl Club<'static, i32>) {}

fn is_sync<T>(_: T)
where
    T: Sync,
{
}

fn main() {
    // This should not cause an ICE either
    // https://github.com/rust-lang/rust-clippy/issues/3144
    is_sync(HashSet::<usize>::new());
}
<commit_msg>Run rustfmt<commit_after>#![warn(clippy::needless_pass_by_value)]
#![allow(
    dead_code,
    clippy::single_match,
    clippy::redundant_pattern_matching,
    clippy::many_single_char_names,
    clippy::option_option
)]

use std::borrow::Borrow;
use std::collections::HashSet;
use std::convert::AsRef;

// `v` should be warned
// `w`, `x` and `y` are allowed (moved or mutated)
fn foo<T: Default>(v: Vec<T>, w: Vec<T>, mut x: Vec<T>, y: Vec<T>) -> Vec<T> {
    assert_eq!(v.len(), 42);

    consume(w);

    x.push(T::default());

    y
}

fn consume<T>(_: T) {}

struct Wrapper(String);

fn bar(x: String, y: Wrapper) {
    assert_eq!(x.len(), 42);
    assert_eq!(y.0.len(), 42);
}

// V implements `Borrow<V>`, but should be warned correctly
fn test_borrow_trait<T: Borrow<str>, U: AsRef<str>, V>(t: T, u: U, v: V) {
    println!("{}", t.borrow());
    println!("{}", u.as_ref());
    consume(&v);
}

// ok
fn test_fn<F: Fn(i32) -> i32>(f: F) {
    f(1);
}

// x should be warned, but y is ok
fn test_match(x: Option<Option<String>>, y: Option<Option<String>>) {
    match x {
        Some(Some(_)) => 1, // not moved
        _ => 0,
    };

    match y {
        Some(Some(s)) => consume(s), // moved
        _ => (),
    };
}

// x and y should be warned, but z is ok
fn test_destructure(x: Wrapper, y: Wrapper, z: Wrapper) {
    let Wrapper(s) = z; // moved
    let Wrapper(ref t) = y; // not moved
    let Wrapper(_) = y; // still not moved

    assert_eq!(x.0.len(), s.len());
    println!("{}", t);
}

trait Foo {}

// `S: Serialize` is allowed to be passed by value, since a caller can pass `&S` instead
trait Serialize {}
impl<'a, T> Serialize for &'a T where T: Serialize {}
impl Serialize for i32 {}

fn test_blanket_ref<T: Foo, S: Serialize>(_foo: T, _serializable: S) {}

fn issue_2114(s: String, t: String, u: Vec<i32>, v: Vec<i32>) {
    s.capacity();
    let _ = t.clone();
    u.capacity();
    let _ = v.clone();
}

struct S<T, U>(T, U);

impl<T: Serialize, U> S<T, U> {
    fn foo(
        self,
        // taking `self` by value is always allowed
        s: String,
        t: String,
    ) -> usize {
        s.len() + t.capacity()
    }

    fn bar(_t: T, // Ok, since `&T: Serialize` too
    ) {
    }

    fn baz(&self, _u: U, _s: Self) {}
}

trait FalsePositive {
    fn visit_str(s: &str);
    fn visit_string(s: String) {
        Self::visit_str(&s);
    }
}

// shouldn't warn on extern funcs
extern "C" fn ext(x: String) -> usize {
    x.len()
}

// whitelist RangeArgument
fn range<T: ::std::ops::RangeBounds<usize>>(range: T) {
    let _ = range.start_bound();
}

struct CopyWrapper(u32);

fn bar_copy(x: u32, y: CopyWrapper) {
    assert_eq!(x, 42);
    assert_eq!(y.0, 42);
}

// x and y should be warned, but z is ok
fn test_destructure_copy(x: CopyWrapper, y: CopyWrapper, z: CopyWrapper) {
    let CopyWrapper(s) = z; // moved
    let CopyWrapper(ref t) = y; // not moved
    let CopyWrapper(_) = y; // still not moved

    assert_eq!(x.0, s);
    println!("{}", t);
}

// The following 3 lines should not cause an ICE. See #2831
trait Bar<'a, A> {}
impl<'b, T> Bar<'b, T> for T {}
fn some_fun<'b, S: Bar<'b, ()>>(_item: S) {}

// Also this should not cause an ICE. See #2831
trait Club<'a, A> {}
impl<T> Club<'static, T> for T {}
fn more_fun(_item: impl Club<'static, i32>) {}

fn is_sync<T>(_: T)
where
    T: Sync,
{
}

fn main() {
    // This should not cause an ICE either
    // https://github.com/rust-lang/rust-clippy/issues/3144
    is_sync(HashSet::<usize>::new());
}
<|endoftext|>


<commit_before>use std::fmt;
use std::hash;
use std::mem;
use std::ptr;

use runtime::{Class, Messageable, Object};
use foundation::INSObject;

#[unsafe_no_drop_flag]
pub struct Id<T> {
	ptr: *T,
}

impl<T> Id<T> {
	pub fn nil() -> Id<T> {
		Id { ptr: ptr::null() }
	}

	unsafe fn retain(&self) {
		msg_send![self.as_ptr() retain];
	}

	unsafe fn release(&self) {
		msg_send![self.as_ptr() release];
	}

	pub unsafe fn from_ptr(ptr: *T) -> Id<T> {
		let id = Id { ptr: ptr };
		id.retain();
		id
	}

	pub unsafe fn from_retained_ptr(ptr: *T) -> Id<T> {
		Id { ptr: ptr }
	}
}

impl<T> Messageable for Id<T> {
	unsafe fn as_ptr(&self) -> *Object {
		self.ptr as *Object
	}
}

impl<T> Clone for Id<T> {
	fn clone(&self) -> Id<T> {
		unsafe {
			Id::from_ptr(self.ptr)
		}
	}
}

#[unsafe_destructor]
impl<T> Drop for Id<T> {
	fn drop(&mut self) {
		if !self.ptr.is_null() {
			unsafe {
				self.release();
			}
			self.ptr = ptr::null();
		}
	}
}

impl<T> Deref<T> for Id<T> {
	fn deref<'a>(&'a self) -> &'a T {
		unsafe { &*self.ptr }
	}
}

impl<T: PartialEq> PartialEq for Id<T> {
	fn eq(&self, other: &Id<T>) -> bool {
		self.deref() == other.deref()
	}

	fn ne(&self, other: &Id<T>) -> bool {
		self.deref() != other.deref()
	}
}

impl<T: Eq> Eq for Id<T> { }

impl<S: hash::Writer, T: hash::Hash<S>> hash::Hash<S> for Id<T> {
	fn hash(&self, state: &mut S) {
		self.deref().hash(state)
	}
}

impl<T: fmt::Show> fmt::Show for Id<T> {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		self.deref().fmt(f)
	}
}

pub trait IdVector<T> {
	fn as_refs_slice<'a>(&'a self) -> &'a [&'a T];

	fn as_ptrs_slice<'a>(&'a self) -> &'a [*T] {
		unsafe {
			mem::transmute(self.as_refs_slice())
		}
	}
}

impl<T> IdVector<T> for Vec<Id<T>> {
	fn as_refs_slice<'a>(&'a self) -> &'a [&'a T] {
		unsafe {
			mem::transmute(self.as_slice())
		}
	}
}

pub struct ClassName<T> {
	name: &'static str,
}

impl<T> ClassName<T> {
	pub fn from_str(name: &'static str) -> ClassName<T> {
		ClassName { name: name }
	}

	pub fn as_str(&self) -> &'static str {
		self.name
	}
}

pub fn class<T: INSObject>() -> Class {
	let name: ClassName<T> = INSObject::class_name();
	Class::get(name.as_str())
}
<commit_msg>Removed retain and release methods from Id.<commit_after>use std::fmt;
use std::hash;
use std::mem;
use std::ptr;

use runtime::{Class, Messageable, Object};
use foundation::INSObject;

#[unsafe_no_drop_flag]
pub struct Id<T> {
	ptr: *T,
}

impl<T> Id<T> {
	pub unsafe fn from_ptr(ptr: *T) -> Id<T> {
		msg_send![ptr as *Object retain];
		Id::from_retained_ptr(ptr)
	}

	pub unsafe fn from_retained_ptr(ptr: *T) -> Id<T> {
		Id { ptr: ptr }
	}
}

impl<T> Messageable for Id<T> {
	unsafe fn as_ptr(&self) -> *Object {
		self.ptr as *Object
	}
}

impl<T> Clone for Id<T> {
	fn clone(&self) -> Id<T> {
		unsafe {
			Id::from_ptr(self.ptr)
		}
	}
}

#[unsafe_destructor]
impl<T> Drop for Id<T> {
	fn drop(&mut self) {
		if !self.ptr.is_null() {
			unsafe {
				msg_send![self.as_ptr() release];
			}
			self.ptr = ptr::null();
		}
	}
}

impl<T> Deref<T> for Id<T> {
	fn deref<'a>(&'a self) -> &'a T {
		unsafe { &*self.ptr }
	}
}

impl<T: PartialEq> PartialEq for Id<T> {
	fn eq(&self, other: &Id<T>) -> bool {
		self.deref() == other.deref()
	}

	fn ne(&self, other: &Id<T>) -> bool {
		self.deref() != other.deref()
	}
}

impl<T: Eq> Eq for Id<T> { }

impl<S: hash::Writer, T: hash::Hash<S>> hash::Hash<S> for Id<T> {
	fn hash(&self, state: &mut S) {
		self.deref().hash(state)
	}
}

impl<T: fmt::Show> fmt::Show for Id<T> {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		self.deref().fmt(f)
	}
}

pub trait IdVector<T> {
	fn as_refs_slice<'a>(&'a self) -> &'a [&'a T];

	fn as_ptrs_slice<'a>(&'a self) -> &'a [*T] {
		unsafe {
			mem::transmute(self.as_refs_slice())
		}
	}
}

impl<T> IdVector<T> for Vec<Id<T>> {
	fn as_refs_slice<'a>(&'a self) -> &'a [&'a T] {
		unsafe {
			mem::transmute(self.as_slice())
		}
	}
}

pub struct ClassName<T> {
	name: &'static str,
}

impl<T> ClassName<T> {
	pub fn from_str(name: &'static str) -> ClassName<T> {
		ClassName { name: name }
	}

	pub fn as_str(&self) -> &'static str {
		self.name
	}
}

pub fn class<T: INSObject>() -> Class {
	let name: ClassName<T> = INSObject::class_name();
	Class::get(name.as_str())
}
<|endoftext|>


<commit_before>// check-pass

// This test exhibits undefined behavior, but it is impossible to prevent generally during
// const eval, even if it could be prevented in the cases here if we added expensive and
// complex checks in rustc.
// The reason it's impossible in general
// is that we run into query cycles even *without* UB, just because we're checking for UB.
// We do not detect it if you create references to statics
// in ways that are UB.

enum Foo {
    A = 5,
    B = 42,
}
enum Bar {
    C = 42,
    D = 99,
}
#[repr(C)]
union Union {
    foo: &'static Foo,
    bar: &'static Bar,
    u8: &'static u8,
}
static BAR: u8 = 5;
static FOO: (&Foo, &Bar) = unsafe {
    (
        // undefined behavior
        Union { u8: &BAR }.foo,
        Union { u8: &BAR }.bar,
    )
};
static FOO2: (&Foo, &Bar) = unsafe { (std::mem::transmute(&BAR), std::mem::transmute(&BAR)) };
//^ undefined behavior

fn main() {}
<commit_msg>We can make const eval sound, it's just super expensive<commit_after>// check-pass

// This test exhibits undefined behavior, but it is very expensive and complex to check for such
// UB in constants.
// Thus, we do not detect it if you create references to statics in ways that are UB.

enum Foo {
    A = 5,
    B = 42,
}
enum Bar {
    C = 42,
    D = 99,
}
#[repr(C)]
union Union {
    foo: &'static Foo,
    bar: &'static Bar,
    u8: &'static u8,
}
static BAR: u8 = 5;
static FOO: (&Foo, &Bar) = unsafe {
    (
        // undefined behavior
        Union { u8: &BAR }.foo,
        Union { u8: &BAR }.bar,
    )
};
static FOO2: (&Foo, &Bar) = unsafe { (std::mem::transmute(&BAR), std::mem::transmute(&BAR)) };
//^ undefined behavior

fn main() {}
<|endoftext|>


<commit_before>mod extractor;

extern crate num_cpus;

use flate2::write::GzEncoder;
use rayon::prelude::*;
use std::fs;
use std::io::{BufRead, BufWriter};
use std::path::{Path, PathBuf};
use tree_sitter::{Language, Parser, Range};

enum TrapCompression {
    None,
    Gzip,
}

impl TrapCompression {
    fn from_env() -> TrapCompression {
        match std::env::var("CODEQL_RUBY_TRAP_COMPRESSION") {
            Ok(method) => match TrapCompression::from_string(&method) {
                Some(c) => c,
                None => {
                    tracing::error!("Unknown compression method '{}'; using gzip.", &method);
                    TrapCompression::Gzip
                }
            },
            // Default compression method if the env var isn't set:
            Err(_) => TrapCompression::Gzip,
        }
    }

    fn from_string(s: &str) -> Option<TrapCompression> {
        match s.to_lowercase().as_ref() {
            "none" => Some(TrapCompression::None),
            "gzip" => Some(TrapCompression::Gzip),
            _ => None,
        }
    }

    fn extension(&self) -> &str {
        match self {
            TrapCompression::None => "trap",
            TrapCompression::Gzip => "trap.gz",
        }
    }
}

/**
 * Gets the number of threads the extractor should use, by reading the
 * CODEQL_THREADS environment variable and using it as described in the
 * extractor spec:
 *
 * "If the number is positive, it indicates the number of threads that should
 * be used. If the number is negative or zero, it should be added to the number
 * of cores available on the machine to determine how many threads to use
 * (minimum of 1). If unspecified, should be considered as set to -1."
 */
fn num_codeql_threads() -> usize {
    let threads_str = std::env::var("CODEQL_THREADS").unwrap_or_else(|_| "-1".to_owned());
    match threads_str.parse::<i32>() {
        Ok(num) if num <= 0 => {
            let reduction = -num as usize;
            std::cmp::max(1, num_cpus::get() - reduction)
        }
        Ok(num) => num as usize,

        Err(_) => {
            tracing::error!(
                "Unable to parse CODEQL_THREADS value '{}'; defaulting to 1 thread.",
                &threads_str
            );
            1
        }
    }
}

fn main() -> std::io::Result<()> {
    tracing_subscriber::fmt()
        .with_target(false)
        .without_time()
        .with_level(true)
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or(tracing_subscriber::EnvFilter::new("ruby_extractor=warn")),
        )
        .init();
    tracing::warn!("Support for Ruby is currently in Beta: https://git.io/codeql-language-support");
    let num_threads = num_codeql_threads();
    tracing::info!(
        "Using {} {}",
        num_threads,
        if num_threads == 1 {
            "thread"
        } else {
            "threads"
        }
    );
    rayon::ThreadPoolBuilder::new()
        .num_threads(num_threads)
        .build_global()
        .unwrap();

    let matches = clap::App::new("Ruby extractor")
        .version("1.0")
        .author("GitHub")
        .about("CodeQL Ruby extractor")
        .args_from_usage(
            "--source-archive-dir=<DIR> 'Sets a custom source archive folder'
                    --output-dir=<DIR>         'Sets a custom trap folder'
                    --file-list=<FILE_LIST>    'A text files containing the paths of the files to extract'",
        )
        .get_matches();
    let src_archive_dir = matches
        .value_of("source-archive-dir")
        .expect("missing --source-archive-dir");
    let src_archive_dir = PathBuf::from(src_archive_dir);

    let trap_dir = matches
        .value_of("output-dir")
        .expect("missing --output-dir");
    let trap_dir = PathBuf::from(trap_dir);
    let trap_compression = TrapCompression::from_env();

    let file_list = matches.value_of("file-list").expect("missing --file-list");
    let file_list = fs::File::open(file_list)?;

    let language = tree_sitter_ruby::language();
    let erb = tree_sitter_embedded_template::language();
    // Look up tree-sitter kind ids now, to avoid string comparisons when scanning ERB files.
    let erb_directive_id = erb.id_for_node_kind("directive", true);
    let erb_output_directive_id = erb.id_for_node_kind("output_directive", true);
    let erb_code_id = erb.id_for_node_kind("code", true);
    let schema = node_types::read_node_types_str("ruby", tree_sitter_ruby::NODE_TYPES)?;
    let erb_schema =
        node_types::read_node_types_str("erb", tree_sitter_embedded_template::NODE_TYPES)?;
    let lines: std::io::Result<Vec<String>> = std::io::BufReader::new(file_list).lines().collect();
    let lines = lines?;
    lines
        .par_iter()
        .try_for_each(|line| {
            let path = PathBuf::from(line).canonicalize()?;
            let src_archive_file = path_for(&src_archive_dir, &path, "");
            let mut source = std::fs::read(&path)?;
            let code_ranges;
            let mut trap_writer = extractor::new_trap_writer();
            if path.extension().map_or(false, |x| x == "erb") {
                tracing::info!("scanning: {}", path.display());
                extractor::extract(
                    erb,
                    "erb",
                    &erb_schema,
                    &mut trap_writer,
                    &path,
                    &source,
                    &[],
                )?;

                let (ranges, line_breaks) = scan_erb(
                    erb,
                    &source,
                    erb_directive_id,
                    erb_output_directive_id,
                    erb_code_id,
                );
                for i in line_breaks {
                    if i < source.len() {
                        source[i] = b'\n';
                    }
                }
                code_ranges = ranges;
            } else {
                code_ranges = vec![];
            }
            extractor::extract(
                language,
                "ruby",
                &schema,
                &mut trap_writer,
                &path,
                &source,
                &code_ranges,
            )?;
            std::fs::create_dir_all(&src_archive_file.parent().unwrap())?;
            std::fs::copy(&path, &src_archive_file)?;
            write_trap(&trap_dir, path, trap_writer, &trap_compression)
        })
        .expect("failed to extract files");

    let path = PathBuf::from("extras");
    let mut trap_writer = extractor::new_trap_writer();
    trap_writer.populate_empty_location();
    write_trap(&trap_dir, path, trap_writer, &trap_compression)
}

fn write_trap(
    trap_dir: &Path,
    path: PathBuf,
    trap_writer: extractor::TrapWriter,
    trap_compression: &TrapCompression,
) -> std::io::Result<()> {
    let trap_file = path_for(trap_dir, &path, trap_compression.extension());
    std::fs::create_dir_all(&trap_file.parent().unwrap())?;
    let trap_file = std::fs::File::create(&trap_file)?;
    let mut trap_file = BufWriter::new(trap_file);
    match trap_compression {
        TrapCompression::None => trap_writer.output(&mut trap_file),
        TrapCompression::Gzip => {
            let mut compressed_writer = GzEncoder::new(trap_file, flate2::Compression::fast());
            trap_writer.output(&mut compressed_writer)
        }
    }
}

fn scan_erb(
    erb: Language,
    source: &[u8],
    directive_id: u16,
    output_directive_id: u16,
    code_id: u16,
) -> (Vec<Range>, Vec<usize>) {
    let mut parser = Parser::new();
    parser.set_language(erb).unwrap();
    let tree = parser.parse(&source, None).expect("Failed to parse file");
    let mut result = Vec::new();
    let mut line_breaks = vec![];

    for n in tree.root_node().children(&mut tree.walk()) {
        let kind_id = n.kind_id();
        if kind_id == directive_id || kind_id == output_directive_id {
            for c in n.children(&mut tree.walk()) {
                if c.kind_id() == code_id {
                    let mut range = c.range();
                    if range.end_byte < source.len() {
                        line_breaks.push(range.end_byte);
                        range.end_byte += 1;
                        range.end_point.column += 1;
                    }
                    result.push(range);
                }
            }
        }
    }
    if result.is_empty() {
        let root = tree.root_node();
        // Add an empty range at the end of the file
        result.push(Range {
            start_byte: root.end_byte(),
            end_byte: root.end_byte(),
            start_point: root.end_position(),
            end_point: root.end_position(),
        });
    }
    (result, line_breaks)
}

fn path_for(dir: &Path, path: &Path, ext: &str) -> PathBuf {
    let mut result = PathBuf::from(dir);
    for component in path.components() {
        match component {
            std::path::Component::Prefix(prefix) => match prefix.kind() {
                std::path::Prefix::Disk(letter) | std::path::Prefix::VerbatimDisk(letter) => {
                    result.push(format!("{}_", letter as char))
                }
                std::path::Prefix::Verbatim(x) | std::path::Prefix::DeviceNS(x) => {
                    result.push(x);
                }
                std::path::Prefix::UNC(server, share)
                | std::path::Prefix::VerbatimUNC(server, share) => {
                    result.push("unc");
                    result.push(server);
                    result.push(share);
                }
            },
            std::path::Component::RootDir => {
                // skip
            }
            std::path::Component::Normal(_) => {
                result.push(component);
            }
            std::path::Component::CurDir => {
                // skip
            }
            std::path::Component::ParentDir => {
                result.pop();
            }
        }
    }
    if !ext.is_empty() {
        match result.extension() {
            Some(x) => {
                let mut new_ext = x.to_os_string();
                new_ext.push(".");
                new_ext.push(ext);
                result.set_extension(new_ext);
            }
            None => {
                result.set_extension(ext);
            }
        }
    }
    result
}
<commit_msg>Ruby: use unwrap_or_else to construct object only when needed<commit_after>mod extractor;

extern crate num_cpus;

use flate2::write::GzEncoder;
use rayon::prelude::*;
use std::fs;
use std::io::{BufRead, BufWriter};
use std::path::{Path, PathBuf};
use tree_sitter::{Language, Parser, Range};

enum TrapCompression {
    None,
    Gzip,
}

impl TrapCompression {
    fn from_env() -> TrapCompression {
        match std::env::var("CODEQL_RUBY_TRAP_COMPRESSION") {
            Ok(method) => match TrapCompression::from_string(&method) {
                Some(c) => c,
                None => {
                    tracing::error!("Unknown compression method '{}'; using gzip.", &method);
                    TrapCompression::Gzip
                }
            },
            // Default compression method if the env var isn't set:
            Err(_) => TrapCompression::Gzip,
        }
    }

    fn from_string(s: &str) -> Option<TrapCompression> {
        match s.to_lowercase().as_ref() {
            "none" => Some(TrapCompression::None),
            "gzip" => Some(TrapCompression::Gzip),
            _ => None,
        }
    }

    fn extension(&self) -> &str {
        match self {
            TrapCompression::None => "trap",
            TrapCompression::Gzip => "trap.gz",
        }
    }
}

/**
 * Gets the number of threads the extractor should use, by reading the
 * CODEQL_THREADS environment variable and using it as described in the
 * extractor spec:
 *
 * "If the number is positive, it indicates the number of threads that should
 * be used. If the number is negative or zero, it should be added to the number
 * of cores available on the machine to determine how many threads to use
 * (minimum of 1). If unspecified, should be considered as set to -1."
 */
fn num_codeql_threads() -> usize {
    let threads_str = std::env::var("CODEQL_THREADS").unwrap_or_else(|_| "-1".to_owned());
    match threads_str.parse::<i32>() {
        Ok(num) if num <= 0 => {
            let reduction = -num as usize;
            std::cmp::max(1, num_cpus::get() - reduction)
        }
        Ok(num) => num as usize,

        Err(_) => {
            tracing::error!(
                "Unable to parse CODEQL_THREADS value '{}'; defaulting to 1 thread.",
                &threads_str
            );
            1
        }
    }
}

fn main() -> std::io::Result<()> {
    tracing_subscriber::fmt()
        .with_target(false)
        .without_time()
        .with_level(true)
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("ruby_extractor=warn")),
        )
        .init();
    tracing::warn!("Support for Ruby is currently in Beta: https://git.io/codeql-language-support");
    let num_threads = num_codeql_threads();
    tracing::info!(
        "Using {} {}",
        num_threads,
        if num_threads == 1 {
            "thread"
        } else {
            "threads"
        }
    );
    rayon::ThreadPoolBuilder::new()
        .num_threads(num_threads)
        .build_global()
        .unwrap();

    let matches = clap::App::new("Ruby extractor")
        .version("1.0")
        .author("GitHub")
        .about("CodeQL Ruby extractor")
        .args_from_usage(
            "--source-archive-dir=<DIR> 'Sets a custom source archive folder'
                    --output-dir=<DIR>         'Sets a custom trap folder'
                    --file-list=<FILE_LIST>    'A text files containing the paths of the files to extract'",
        )
        .get_matches();
    let src_archive_dir = matches
        .value_of("source-archive-dir")
        .expect("missing --source-archive-dir");
    let src_archive_dir = PathBuf::from(src_archive_dir);

    let trap_dir = matches
        .value_of("output-dir")
        .expect("missing --output-dir");
    let trap_dir = PathBuf::from(trap_dir);
    let trap_compression = TrapCompression::from_env();

    let file_list = matches.value_of("file-list").expect("missing --file-list");
    let file_list = fs::File::open(file_list)?;

    let language = tree_sitter_ruby::language();
    let erb = tree_sitter_embedded_template::language();
    // Look up tree-sitter kind ids now, to avoid string comparisons when scanning ERB files.
    let erb_directive_id = erb.id_for_node_kind("directive", true);
    let erb_output_directive_id = erb.id_for_node_kind("output_directive", true);
    let erb_code_id = erb.id_for_node_kind("code", true);
    let schema = node_types::read_node_types_str("ruby", tree_sitter_ruby::NODE_TYPES)?;
    let erb_schema =
        node_types::read_node_types_str("erb", tree_sitter_embedded_template::NODE_TYPES)?;
    let lines: std::io::Result<Vec<String>> = std::io::BufReader::new(file_list).lines().collect();
    let lines = lines?;
    lines
        .par_iter()
        .try_for_each(|line| {
            let path = PathBuf::from(line).canonicalize()?;
            let src_archive_file = path_for(&src_archive_dir, &path, "");
            let mut source = std::fs::read(&path)?;
            let code_ranges;
            let mut trap_writer = extractor::new_trap_writer();
            if path.extension().map_or(false, |x| x == "erb") {
                tracing::info!("scanning: {}", path.display());
                extractor::extract(
                    erb,
                    "erb",
                    &erb_schema,
                    &mut trap_writer,
                    &path,
                    &source,
                    &[],
                )?;

                let (ranges, line_breaks) = scan_erb(
                    erb,
                    &source,
                    erb_directive_id,
                    erb_output_directive_id,
                    erb_code_id,
                );
                for i in line_breaks {
                    if i < source.len() {
                        source[i] = b'\n';
                    }
                }
                code_ranges = ranges;
            } else {
                code_ranges = vec![];
            }
            extractor::extract(
                language,
                "ruby",
                &schema,
                &mut trap_writer,
                &path,
                &source,
                &code_ranges,
            )?;
            std::fs::create_dir_all(&src_archive_file.parent().unwrap())?;
            std::fs::copy(&path, &src_archive_file)?;
            write_trap(&trap_dir, path, trap_writer, &trap_compression)
        })
        .expect("failed to extract files");

    let path = PathBuf::from("extras");
    let mut trap_writer = extractor::new_trap_writer();
    trap_writer.populate_empty_location();
    write_trap(&trap_dir, path, trap_writer, &trap_compression)
}

fn write_trap(
    trap_dir: &Path,
    path: PathBuf,
    trap_writer: extractor::TrapWriter,
    trap_compression: &TrapCompression,
) -> std::io::Result<()> {
    let trap_file = path_for(trap_dir, &path, trap_compression.extension());
    std::fs::create_dir_all(&trap_file.parent().unwrap())?;
    let trap_file = std::fs::File::create(&trap_file)?;
    let mut trap_file = BufWriter::new(trap_file);
    match trap_compression {
        TrapCompression::None => trap_writer.output(&mut trap_file),
        TrapCompression::Gzip => {
            let mut compressed_writer = GzEncoder::new(trap_file, flate2::Compression::fast());
            trap_writer.output(&mut compressed_writer)
        }
    }
}

fn scan_erb(
    erb: Language,
    source: &[u8],
    directive_id: u16,
    output_directive_id: u16,
    code_id: u16,
) -> (Vec<Range>, Vec<usize>) {
    let mut parser = Parser::new();
    parser.set_language(erb).unwrap();
    let tree = parser.parse(&source, None).expect("Failed to parse file");
    let mut result = Vec::new();
    let mut line_breaks = vec![];

    for n in tree.root_node().children(&mut tree.walk()) {
        let kind_id = n.kind_id();
        if kind_id == directive_id || kind_id == output_directive_id {
            for c in n.children(&mut tree.walk()) {
                if c.kind_id() == code_id {
                    let mut range = c.range();
                    if range.end_byte < source.len() {
                        line_breaks.push(range.end_byte);
                        range.end_byte += 1;
                        range.end_point.column += 1;
                    }
                    result.push(range);
                }
            }
        }
    }
    if result.is_empty() {
        let root = tree.root_node();
        // Add an empty range at the end of the file
        result.push(Range {
            start_byte: root.end_byte(),
            end_byte: root.end_byte(),
            start_point: root.end_position(),
            end_point: root.end_position(),
        });
    }
    (result, line_breaks)
}

fn path_for(dir: &Path, path: &Path, ext: &str) -> PathBuf {
    let mut result = PathBuf::from(dir);
    for component in path.components() {
        match component {
            std::path::Component::Prefix(prefix) => match prefix.kind() {
                std::path::Prefix::Disk(letter) | std::path::Prefix::VerbatimDisk(letter) => {
                    result.push(format!("{}_", letter as char))
                }
                std::path::Prefix::Verbatim(x) | std::path::Prefix::DeviceNS(x) => {
                    result.push(x);
                }
                std::path::Prefix::UNC(server, share)
                | std::path::Prefix::VerbatimUNC(server, share) => {
                    result.push("unc");
                    result.push(server);
                    result.push(share);
                }
            },
            std::path::Component::RootDir => {
                // skip
            }
            std::path::Component::Normal(_) => {
                result.push(component);
            }
            std::path::Component::CurDir => {
                // skip
            }
            std::path::Component::ParentDir => {
                result.pop();
            }
        }
    }
    if !ext.is_empty() {
        match result.extension() {
            Some(x) => {
                let mut new_ext = x.to_os_string();
                new_ext.push(".");
                new_ext.push(ext);
                result.set_extension(new_ext);
            }
            None => {
                result.set_extension(ext);
            }
        }
    }
    result
}
<|endoftext|>


<commit_before>// before struct
pub struct MyStruct {
	member : i32 // member line
} // last struct line

// before impl MyStruct
impl MyStruct {} // newline between { and }
// some line

// before impl MyStruct
impl MyStruct {} // newline after }
// some line
<commit_msg>improved autoindent test sample<commit_after>namespace {
  // before struct
  pub struct MyStruct {
    member : i32 // member line
  } // last struct line

  // before impl MyStruct
  impl MyStruct {} // newline between { and }
  // some line

  // before impl MyStruct
  impl MyStruct {} // newline after }
  // some line
}<|endoftext|>


<commit_before>use std::fmt;
use std::fmt::{Display, Formatter};
use joker::track::{Span, TrackingRef};
use expr::Expr;
use patt::{Patt, AssignTarget, CompoundPatt, PropPatt};
use obj::{Prop, PropVal};

#[derive(Debug, PartialEq, Clone)]
pub enum Error {
    InvalidAssignTarget(Option<Span>),
    InvalidPropPatt(Option<Span>)
}

impl Display for Error {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        match self {
            &Error::InvalidAssignTarget(_) => {
                fmt.write_str("invalid assignment pattern")
            }
            &Error::InvalidPropPatt(_) => {
                fmt.write_str("invalid object property in assignment pattern")
            }
        }
    }
}

pub trait IntoAssignTarget {
    fn into_assign_target(self) -> Result<AssignTarget, Error>;
}

impl IntoAssignTarget for Expr {
    fn into_assign_target(self) -> Result<AssignTarget, Error> {
        Ok(match self {
            Expr::Id(id)                     => AssignTarget::Id(id),
            Expr::Dot(location, obj, key)    => AssignTarget::Dot(location, obj, key),
            Expr::Brack(location, obj, prop) => AssignTarget::Brack(location, obj, prop),
            _ => { return Err(Error::InvalidAssignTarget(*self.tracking_ref())); }
        })
    }
}

pub trait IntoAssignPatt : IntoAssignTarget {
    fn into_assign_patt(self) -> Result<Patt<AssignTarget>, Error>;
}

impl IntoAssignPatt for Expr {
    fn into_assign_patt(self) -> Result<Patt<AssignTarget>, Error> {
        Ok(match self {
            Expr::Obj(location, props) => {
                let mut prop_patts = Vec::with_capacity(props.len());
                for prop in props {
                    prop_patts.push(try!(prop.into_assign_prop()));
                }
                Patt::Compound(CompoundPatt::Obj(location, prop_patts))
            }
            Expr::Arr(location, exprs) => {
                let mut patts = Vec::with_capacity(exprs.len());
                for expr in exprs {
                    patts.push(match expr {
                        Some(expr) => Some(try!(expr.into_assign_patt())),
                        None => None
                    });
                }
                Patt::Compound(CompoundPatt::Arr(location, patts))
            }
            _ => { return self.into_assign_target().map(Patt::Simple); }
        })
    }
}

pub trait IntoAssignProp {
    fn into_assign_prop(self) -> Result<PropPatt<AssignTarget>, Error>;
}

impl IntoAssignProp for Prop {
    fn into_assign_prop(self) -> Result<PropPatt<AssignTarget>, Error> {
        let location = *self.tracking_ref();
        let key = self.key;
        let patt = match self.val {
            PropVal::Init(expr) => try!(expr.into_assign_patt()),
            _ => { return Err(Error::InvalidPropPatt(*self.val.tracking_ref())); }
        };
        Ok(PropPatt { location: location, key: key, patt: patt })
    }
}
<commit_msg>replace `try!` with `?` in easter crate<commit_after>use std::fmt;
use std::fmt::{Display, Formatter};
use joker::track::{Span, TrackingRef};
use expr::Expr;
use patt::{Patt, AssignTarget, CompoundPatt, PropPatt};
use obj::{Prop, PropVal};

#[derive(Debug, PartialEq, Clone)]
pub enum Error {
    InvalidAssignTarget(Option<Span>),
    InvalidPropPatt(Option<Span>)
}

impl Display for Error {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        match self {
            &Error::InvalidAssignTarget(_) => {
                fmt.write_str("invalid assignment pattern")
            }
            &Error::InvalidPropPatt(_) => {
                fmt.write_str("invalid object property in assignment pattern")
            }
        }
    }
}

pub trait IntoAssignTarget {
    fn into_assign_target(self) -> Result<AssignTarget, Error>;
}

impl IntoAssignTarget for Expr {
    fn into_assign_target(self) -> Result<AssignTarget, Error> {
        Ok(match self {
            Expr::Id(id)                     => AssignTarget::Id(id),
            Expr::Dot(location, obj, key)    => AssignTarget::Dot(location, obj, key),
            Expr::Brack(location, obj, prop) => AssignTarget::Brack(location, obj, prop),
            _ => { return Err(Error::InvalidAssignTarget(*self.tracking_ref())); }
        })
    }
}

pub trait IntoAssignPatt : IntoAssignTarget {
    fn into_assign_patt(self) -> Result<Patt<AssignTarget>, Error>;
}

impl IntoAssignPatt for Expr {
    fn into_assign_patt(self) -> Result<Patt<AssignTarget>, Error> {
        Ok(match self {
            Expr::Obj(location, props) => {
                let mut prop_patts = Vec::with_capacity(props.len());
                for prop in props {
                    prop_patts.push(prop.into_assign_prop()?);
                }
                Patt::Compound(CompoundPatt::Obj(location, prop_patts))
            }
            Expr::Arr(location, exprs) => {
                let mut patts = Vec::with_capacity(exprs.len());
                for expr in exprs {
                    patts.push(match expr {
                        Some(expr) => Some(expr.into_assign_patt()?),
                        None => None
                    });
                }
                Patt::Compound(CompoundPatt::Arr(location, patts))
            }
            _ => { return self.into_assign_target().map(Patt::Simple); }
        })
    }
}

pub trait IntoAssignProp {
    fn into_assign_prop(self) -> Result<PropPatt<AssignTarget>, Error>;
}

impl IntoAssignProp for Prop {
    fn into_assign_prop(self) -> Result<PropPatt<AssignTarget>, Error> {
        let location = *self.tracking_ref();
        let key = self.key;
        let patt = match self.val {
            PropVal::Init(expr) => expr.into_assign_patt()?,
            _ => { return Err(Error::InvalidPropPatt(*self.val.tracking_ref())); }
        };
        Ok(PropPatt { location: location, key: key, patt: patt })
    }
}
<|endoftext|>


<commit_before>// This is actually a test that runs on the HOST (!). I had to hack things
// around quite a bit to make this work; that's way this ended up in this
// directory :-).

extern crate f3;

use f3::peripheral;

macro_rules! offset {
    ($register:expr, $offset:expr) => {
        if $offset != &$register as *const _ as usize {
            panic!("{}: expected 0x{:x}, got 0x{:x}. Off by {}",
                   stringify!($register), $offset,
                   &$register as *const _ as usize,
                   (&$register as *const _ as isize - $offset).abs())
        }
    }
}

fn main() {
    let btim = unsafe { &*(0x0 as *const peripheral::tim::BTim) };

    offset!(btim.cr1, 0x0);
    offset!(btim.cr2, 0x4);
    offset!(btim.dier, 0xc);
    offset!(btim.sr, 0x10);
    offset!(btim.egr, 0x14);
    offset!(btim.cnt, 0x24);
    offset!(btim.psc, 0x28);
    offset!(btim.arr, 0x2c);

    let dbgmcu = unsafe { &*(0x0 as *const peripheral::dbgmcu::Dbgmcu) };

    offset!(dbgmcu.idcode, 0x0);
    offset!(dbgmcu.cr, 0x4);
    offset!(dbgmcu.apb1fz, 0x8);
    offset!(dbgmcu.apb2fz, 0xc);

    let gpio = unsafe { &*(0x0 as *const peripheral::gpio::Gpio) };

    offset!(gpio.moder, 0x0);
    offset!(gpio.otyper, 0x4);
    offset!(gpio.ospeedr, 0x8);
    offset!(gpio.pupdr, 0xc);
    offset!(gpio.idr, 0x10);
    offset!(gpio.odr, 0x14);
    offset!(gpio.bsrr, 0x18);
    offset!(gpio.lckr, 0x1c);
    offset!(gpio.afrl, 0x20);
    offset!(gpio.afrh, 0x24);
    offset!(gpio.brr, 0x28);

    let gptim = unsafe { &*(0x0 as *const peripheral::tim::GpTim) };

    offset!(gptim.cr1, 0x0);
    offset!(gptim.cr2, 0x4);
    offset!(gptim.smcr, 0x8);
    offset!(gptim.dier, 0xc);
    offset!(gptim.sr, 0x10);
    offset!(gptim.egr, 0x14);
    offset!(gptim.ccmr1, 0x18);
    offset!(gptim.ccmr2, 0x1c);
    offset!(gptim.ccer, 0x20);
    offset!(gptim.cnt, 0x24);
    offset!(gptim.psc, 0x28);
    offset!(gptim.arr, 0x2c);
    offset!(gptim.ccr1, 0x34);
    offset!(gptim.ccr2, 0x38);
    offset!(gptim.ccr3, 0x3c);
    offset!(gptim.ccr4, 0x40);
    offset!(gptim.dcr, 0x48);
    offset!(gptim.dmar, 0x4c);

    let i2c = unsafe { &*(0x0 as *const peripheral::i2c::I2c) };

    offset!(i2c.cr1, 0x0);
    offset!(i2c.cr2, 0x4);
    offset!(i2c.oar1, 0x8);
    offset!(i2c.oar2, 0xc);
    offset!(i2c.timingr, 0x10);
    offset!(i2c.timeoutr, 0x14);
    offset!(i2c.isr, 0x18);
    offset!(i2c.icr, 0x1c);
    offset!(i2c.pecr, 0x20);
    offset!(i2c.rxdr, 0x24);
    offset!(i2c.txdr, 0x28);

    let rcc = unsafe { &*(0x0 as *const peripheral::rcc::Rcc) };

    offset!(rcc.cr, 0x0);
    offset!(rcc.cfgr, 0x4);
    offset!(rcc.cir, 0x8);
    offset!(rcc.apb2rstr, 0xc);
    offset!(rcc.apb1rstr, 0x10);
    offset!(rcc.ahbenr, 0x14);
    offset!(rcc.apb2enr, 0x18);
    offset!(rcc.apb1enr, 0x1c);
    offset!(rcc.bdcr, 0x20);
    offset!(rcc.csr, 0x24);
    offset!(rcc.ahbrstr, 0x28);
    offset!(rcc.cfgr2, 0x2c);
    offset!(rcc.cfgr3, 0x30);

    let spi = unsafe { &*(0x0 as *const peripheral::spi::Spi) };

    offset!(spi.cr1, 0x00);
    offset!(spi.cr2, 0x04);
    offset!(spi.sr, 0x08);
    offset!(spi.dr, 0x0C);
    offset!(spi.crcpr, 0x10);
    offset!(spi.rxcrcr, 0x14);
    offset!(spi.txcrcr, 0x18);
    offset!(spi.i2scfgr, 0x1C);
    offset!(spi.i2spr, 0x20);

    let usart = unsafe { &*(0x0 as *const peripheral::usart::Usart) };

    offset!(usart.cr1, 0x0);
    offset!(usart.cr2, 0x4);
    offset!(usart.cr3, 0x8);
    offset!(usart.brr, 0xc);
    offset!(usart.gtpr, 0x10);
    offset!(usart.rtor, 0x14);
    offset!(usart.rqr, 0x18);
    offset!(usart.isr, 0x1c);
    offset!(usart.icr, 0x20);
    offset!(usart.rdr, 0x24);
    offset!(usart.tdr, 0x28);
}
<commit_msg>fix the test for real<commit_after>// This is actually a test that runs on the HOST (!). I had to hack things
// around quite a bit to make this work; that's way this ended up in this
// directory :-).

extern crate f3;

use f3::peripheral;

macro_rules! offset {
    ($register:expr, $offset:expr) => {
        if $offset != &$register as *const _ as usize {
            panic!("{}: expected 0x{:x}, got 0x{:x}. Off by {}",
                   stringify!($register), $offset,
                   &$register as *const _ as usize,
                   (&$register as *const _ as isize - $offset).abs())
        }
    }
}

fn main() {
    let btim = unsafe { &*(0x0 as *const peripheral::btim::BTim) };

    offset!(btim.cr1, 0x0);
    offset!(btim.cr2, 0x4);
    offset!(btim.dier, 0xc);
    offset!(btim.sr, 0x10);
    offset!(btim.egr, 0x14);
    offset!(btim.cnt, 0x24);
    offset!(btim.psc, 0x28);
    offset!(btim.arr, 0x2c);

    let dbgmcu = unsafe { &*(0x0 as *const peripheral::dbgmcu::Dbgmcu) };

    offset!(dbgmcu.idcode, 0x0);
    offset!(dbgmcu.cr, 0x4);
    offset!(dbgmcu.apb1fz, 0x8);
    offset!(dbgmcu.apb2fz, 0xc);

    let gpio = unsafe { &*(0x0 as *const peripheral::gpio::Gpio) };

    offset!(gpio.moder, 0x0);
    offset!(gpio.otyper, 0x4);
    offset!(gpio.ospeedr, 0x8);
    offset!(gpio.pupdr, 0xc);
    offset!(gpio.idr, 0x10);
    offset!(gpio.odr, 0x14);
    offset!(gpio.bsrr, 0x18);
    offset!(gpio.lckr, 0x1c);
    offset!(gpio.afrl, 0x20);
    offset!(gpio.afrh, 0x24);
    offset!(gpio.brr, 0x28);

    let gptim = unsafe { &*(0x0 as *const peripheral::gptim::GpTim) };

    offset!(gptim.cr1, 0x0);
    offset!(gptim.cr2, 0x4);
    offset!(gptim.smcr, 0x8);
    offset!(gptim.dier, 0xc);
    offset!(gptim.sr, 0x10);
    offset!(gptim.egr, 0x14);
    offset!(gptim.ccmr1_output, 0x18);
    offset!(gptim.ccmr2_output, 0x1c);
    offset!(gptim.ccer, 0x20);
    offset!(gptim.cnt, 0x24);
    offset!(gptim.psc, 0x28);
    offset!(gptim.arr, 0x2c);
    offset!(gptim.ccr1, 0x34);
    offset!(gptim.ccr2, 0x38);
    offset!(gptim.ccr3, 0x3c);
    offset!(gptim.ccr4, 0x40);
    offset!(gptim.dcr, 0x48);
    offset!(gptim.dmar, 0x4c);

    let i2c = unsafe { &*(0x0 as *const peripheral::i2c::I2c) };

    offset!(i2c.cr1, 0x0);
    offset!(i2c.cr2, 0x4);
    offset!(i2c.oar1, 0x8);
    offset!(i2c.oar2, 0xc);
    offset!(i2c.timingr, 0x10);
    offset!(i2c.timeoutr, 0x14);
    offset!(i2c.isr, 0x18);
    offset!(i2c.icr, 0x1c);
    offset!(i2c.pecr, 0x20);
    offset!(i2c.rxdr, 0x24);
    offset!(i2c.txdr, 0x28);

    let rcc = unsafe { &*(0x0 as *const peripheral::rcc::Rcc) };

    offset!(rcc.cr, 0x0);
    offset!(rcc.cfgr, 0x4);
    offset!(rcc.cir, 0x8);
    offset!(rcc.apb2rstr, 0xc);
    offset!(rcc.apb1rstr, 0x10);
    offset!(rcc.ahbenr, 0x14);
    offset!(rcc.apb2enr, 0x18);
    offset!(rcc.apb1enr, 0x1c);
    offset!(rcc.bdcr, 0x20);
    offset!(rcc.csr, 0x24);
    offset!(rcc.ahbrstr, 0x28);
    offset!(rcc.cfgr2, 0x2c);
    offset!(rcc.cfgr3, 0x30);

    let spi = unsafe { &*(0x0 as *const peripheral::spi::Spi) };

    offset!(spi.cr1, 0x00);
    offset!(spi.cr2, 0x04);
    offset!(spi.sr, 0x08);
    offset!(spi.dr, 0x0C);
    offset!(spi.crcpr, 0x10);
    offset!(spi.rxcrcr, 0x14);
    offset!(spi.txcrcr, 0x18);
    offset!(spi.i2scfgr, 0x1C);
    offset!(spi.i2spr, 0x20);

    let usart = unsafe { &*(0x0 as *const peripheral::usart::Usart) };

    offset!(usart.cr1, 0x0);
    offset!(usart.cr2, 0x4);
    offset!(usart.cr3, 0x8);
    offset!(usart.brr, 0xc);
    offset!(usart.gtpr, 0x10);
    offset!(usart.rtor, 0x14);
    offset!(usart.rqr, 0x18);
    offset!(usart.isr, 0x1c);
    offset!(usart.icr, 0x20);
    offset!(usart.rdr, 0x24);
    offset!(usart.tdr, 0x28);
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate modparse_derive;
#[macro_use]
extern crate derive_builder;
#[macro_use]
extern crate torchrs;
#[macro_use]
extern crate clap;

use torchrs::autograd::{Variable, VariableArgs, VarAccess};
use torchrs::tensor::Tensor;
use torchrs::optim;
use torchrs::nn;
use torchrs::nn::{ModuleStruct, ModIntf, ModDelegate, Module};
use torchrs::nn::functional as F;
use torchrs::utils::data as D;
use torchrs::utils::torchvision::{datasets, transforms};

use clap::{Arg, App, SubCommand};

use std::env;

#[derive(Clone, Builder)]
struct NetArgs {
    #[builder(default="64")]
    batch_size: usize,
    #[builder(default="1000")]
    test_batch_size: usize,
    #[builder(default="10")]
    epochs: u32,
    #[builder(default="0.01")]
    lr: f32,
    #[builder(default="0.5")]
    momentum: f32,
    #[builder(default="true")]
    cuda: bool,
    #[builder(default="1")]
    seed: usize,
    #[builder(default="10")]
    log_interval: usize,
}

impl Default for NetArgs {
    fn default() -> Self {
        NetArgsBuilder::default().build().unwrap()
    }
}

fn parse_args() -> NetArgs {
    unimplemented!();
    let cmd_args: Vec<String> = env::args().collect();
    let mut args = NetArgs::default();
    let matches = App::new("Torch.rs MNIST Example")
        .version("0.1")
        .author("Some Guy")
        .about("recognize handwritten digits")
        .arg(Arg::with_name("batch-size").takes_value(true))
        .arg(Arg::with_name("test-batch-size").takes_value(true))
        .arg(Arg::with_name("epochs").takes_value(true))
        .arg(Arg::with_name("lr").takes_value(true))
        .arg(Arg::with_name("momentum").takes_value(true))
        .arg(Arg::with_name("no-cuda").takes_value(false))
        .arg(Arg::with_name("seed").takes_value(true))
        .arg(Arg::with_name("log-interval").takes_value(true))
        .get_matches();
    let (b_size, tb_size, epochs, lr, seed, momentum, log_int) =
        (value_t!(matches.value_of("batch-size"), usize),
         value_t!(matches.value_of("test-batch-size"), usize),
         value_t!(matches.value_of("epochs"), u32),
         value_t!(matches.value_of("lr"), f32),
         value_t!(matches.value_of("momentum"), f32),
         value_t!(matches.value_of("seed"), u32),
         value_t!(matches.value_of("log-interval"), u32));
    if let Some(b_size) = b_size {
        args.batch_size = b_size
    }
    if let Some(tb_size) = tb_size {
        args.test_batch_size = tb_size
    }
    if let Some(epochs) = epochs {
        args.epochs = epochs
    }
    if let Some(lr) = lr {
        args.lr = lr
    }
    if let Some(momentum) = momentum {
        args.momentum = momentum
    }
    if let Some(seed) = seed {
        args.seed = seed
    }
    if let Some(log_int) = log_int {
        args.log_interval = log_int
    }
    //XXX CUDA?
    args
}

#[derive(ModParse)]
struct Net<T: Default + Copy> {
    delegate: nn::Module<T>,
    conv1: nn::Conv2d<T>,
    conv2: nn::Conv2d<T>,
    fc1: nn::Linear<T>,
    fc2: nn::Linear<T>,
}

impl<T: Default + Copy> Net<T> {
    pub fn new() -> Net<T> {
        Net {
                delegate: nn::Module::new(),
                conv1: nn::Conv2d::build(1, 10, (5, 5)).done(),
                conv2: nn::Conv2d::build(10, 20, (5, 5)).done(),
                fc1: nn::Linear::build(320, 50).done(),
                fc2: nn::Linear::build(50, 10).done(),
            }
            .init_module()
    }
}
impl_mod_delegate!(Net);

// The forward operations could take on one of two implementations.
// The first supporting a near verbatim version of the python
// implementation, and the second supporting a slightly more
// idiomatic to Rust method chaining.

// a) as a near verbatim implementation of the python version
impl ModIntf<f32> for Net<f32> {
    fn forward(&mut self, args: &mut Variable<f32>) -> Variable<f32> {
        let pool_val = F::MaxPool2dArgs::default();
        let mut dropout_val = F::DropoutArgs::default();
        dropout_val.training = self.delegate.training;
        let mut x = F::relu(&F::max_pool2d(&self.conv1.f(args), (2, 2), &pool_val));
        x = F::relu(&F::max_pool2d(&F::dropout2d(&self.conv2.f(&mut x), &dropout_val),
                                   (2, 2),
                                   &pool_val));
        x = x.view(&[-1, 320]);
        x = F::relu(&self.fc1.f(&mut x));
        x = F::dropout(&x, &dropout_val);
        x = self.fc2.f(&mut x);
        F::log_softmax(&x)
    }
}

fn train(model: &mut Net<f32>,
         args: &NetArgs,
         train_loader: &D::BatchLoader<f32, i64>,
         epoch: u32,
         optimizer: &mut optim::OptIntf) {
    model.train();
    for (batch_idx, (ref data, ref target)) in train_loader.iter().enumerate() {
        let (mut data, target) = if args.cuda {
            (Variable::new(data.cuda(None).clone()), Variable::new(target.cuda(None).clone()))
        } else {
            (Variable::new(data.clone()), Variable::new(target.clone()))
        };
        optimizer.zero_grad();
        let output = model.f(&mut data);
        let mut loss = F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        loss.backward();
        optimizer.step();
        if batch_idx % args.log_interval == 0 {
            println!("Train Epoch: {} [{}/{} ({:.0}%)]\tLoss: {:.6}",
                     epoch,
                     batch_idx * data.data().len(),
                     train_loader.dataset.len(),
                     100. * (batch_idx as f32) / train_loader.len() as f32,
                     loss.data()[0]);
        }
    }
}

fn test(model: &mut Net<f32>, args: &NetArgs, test_loader: &D::BatchLoader<f32, i64>) {
    model.eval();
    let mut test_loss = 0.;
    let mut correct = 0;
    for (ref data, ref target) in test_loader.iter() {
        let varargs = VariableArgs {
            volatile: true,
            ..Default::default()
        };
        let (mut data, mut target) = if args.cuda {
            (Variable::new_args(data.cuda(None).clone(), &varargs),
             Variable::new(target.cuda(None).clone()))
        } else {
            (Variable::new_args(data.clone(), &varargs), Variable::new(target.clone()))
        };
        let mut output = model.f(&mut data);
        test_loss += F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        let pred = output.data().max_reduce(1).1;
        correct += pred.eq_tensor(&*target.data()).cpu().sum() as u32;
    }
    test_loss /= test_loader.len() as f32;
    println!("\nTest set: Average loss: {:.4}, Accuracy: {}/{} ({:.0}%)\n",
             test_loss,
             correct,
             test_loader.dataset.len(),
             100. * (correct as f32) / (test_loader.dataset.len() as f32))

}

fn main() {
    // no data loaders yet so demonstrate with Vec placeholder

    let train_loader: D::BatchLoader<f32, i64> =
        D::DataLoader::new(datasets::MNIST::build("../data").done(),
                           D::DataLoaderArgs::default());
    let test_loader: D::BatchLoader<f32, i64> =
        D::DataLoader::new(datasets::MNIST::build("../data").done(),
                           D::DataLoaderArgs::default());
    let args = parse_args();
    let mut optimizer = optim::SGD::new();

    let mut model = Net::new();
    for epoch in 1..args.epochs + 1 {
        train(&mut model, &args, &train_loader, epoch, &mut optimizer);
        test(&mut model, &args, &test_loader);
    }
}
<commit_msg>rework so that we compile again with the updated data loader interfaces<commit_after>#[macro_use]
extern crate modparse_derive;
#[macro_use]
extern crate derive_builder;
#[macro_use]
extern crate torchrs;
#[macro_use]
extern crate clap;

use torchrs::autograd::{Variable, VariableArgs, VarAccess};
use torchrs::tensor::Tensor;
use torchrs::optim;
use torchrs::optim::OptimVal;
use torchrs::nn;
use torchrs::nn::{ModuleStruct, ModIntf, ModDelegate, Module, ModRefMut};
use torchrs::nn::functional as F;
use torchrs::utils::data as D;
use torchrs::utils::torchvision::{datasets, transforms};

use clap::{Arg, App, SubCommand};

use std::env;

#[derive(Clone, Builder)]
struct NetArgs {
    #[builder(default="64")]
    batch_size: usize,
    #[builder(default="1000")]
    test_batch_size: usize,
    #[builder(default="10")]
    epochs: u32,
    #[builder(default="0.01")]
    lr: f32,
    #[builder(default="0.5")]
    momentum: f32,
    #[builder(default="true")]
    cuda: bool,
    #[builder(default="1")]
    seed: usize,
    #[builder(default="10")]
    log_interval: usize,
}

impl Default for NetArgs {
    fn default() -> Self {
        NetArgsBuilder::default().build().unwrap()
    }
}

fn parse_args() -> NetArgs {
    unimplemented!();
    let cmd_args: Vec<String> = env::args().collect();
    let mut args = NetArgs::default();
    let matches = App::new("Torch.rs MNIST Example")
        .version("0.1")
        .author("Some Guy")
        .about("recognize handwritten digits")
        .arg(Arg::with_name("batch-size").takes_value(true))
        .arg(Arg::with_name("test-batch-size").takes_value(true))
        .arg(Arg::with_name("epochs").takes_value(true))
        .arg(Arg::with_name("lr").takes_value(true))
        .arg(Arg::with_name("momentum").takes_value(true))
        .arg(Arg::with_name("no-cuda").takes_value(false))
        .arg(Arg::with_name("seed").takes_value(true))
        .arg(Arg::with_name("log-interval").takes_value(true))
        .get_matches();
    let (b_size, tb_size, epochs, lr, momentum, seed, log_int) =
        (value_t!(matches.value_of("batch-size"), usize).ok(),
         value_t!(matches.value_of("test-batch-size"), usize).ok(),
         value_t!(matches.value_of("epochs"), u32).ok(),
         value_t!(matches.value_of("lr"), f32).ok(),
         value_t!(matches.value_of("momentum"), f32).ok(),
         value_t!(matches.value_of("seed"), usize).ok(),
         value_t!(matches.value_of("log-interval"), usize).ok());
    if let Some(b_size) = b_size {
        args.batch_size = b_size
    }
    if let Some(tb_size) = tb_size {
        args.test_batch_size = tb_size
    }
    if let Some(epochs) = epochs {
        args.epochs = epochs
    }
    if let Some(lr) = lr {
        args.lr = lr
    }
    if let Some(momentum) = momentum {
        args.momentum = momentum
    }
    if let Some(seed) = seed {
        args.seed = seed
    }
    if let Some(log_int) = log_int {
        args.log_interval = log_int
    }
    //XXX CUDA?
    args
}

#[derive(ModParse)]
struct Net<T: Default + Copy> {
    delegate: nn::Module<T>,
    conv1: nn::Conv2d<T>,
    conv2: nn::Conv2d<T>,
    fc1: nn::Linear<T>,
    fc2: nn::Linear<T>,
}

impl<T: Default + Copy> Net<T> {
    pub fn new() -> Net<T> {
        Net {
                delegate: nn::Module::new(),
                conv1: nn::Conv2d::build(1, 10, (5, 5)).done(),
                conv2: nn::Conv2d::build(10, 20, (5, 5)).done(),
                fc1: nn::Linear::build(320, 50).done(),
                fc2: nn::Linear::build(50, 10).done(),
            }
            .init_module()
    }
}
impl_mod_delegate!(Net);

// The forward operations could take on one of two implementations.
// The first supporting a near verbatim version of the python
// implementation, and the second supporting a slightly more
// idiomatic to Rust method chaining.

// a) as a near verbatim implementation of the python version
impl<T: Copy + Default + 'static> ModIntf<T> for Net<T> {
    fn forward(&mut self, args: &mut Variable<T>) -> Variable<T> {
        let pool_val = F::MaxPool2dArgs::default();
        let mut dropout_val = F::DropoutArgs::default();
        dropout_val.training = self.delegate.training;
        let mut x = F::relu(&F::max_pool2d(&self.conv1.f(args), (2, 2), &pool_val));
        x = F::relu(&F::max_pool2d(&F::dropout2d(&self.conv2.f(&mut x), &dropout_val),
                                   (2, 2),
                                   &pool_val));
        x = x.view(&[-1, 320]);
        x = F::relu(&self.fc1.f(&mut x));
        x = F::dropout(&x, &dropout_val);
        x = self.fc2.f(&mut x);
        F::log_softmax(&x)
    }
}

fn train<'a>(model: &mut Net<f32>,
             args: &NetArgs,
             train_loader: &D::BatchLoader<f32, i64>,
             epoch: u32,
             optimizer: &mut optim::OptIntf<Net<f32>, f32>) {
    model.train();
    for (batch_idx, (ref data, ref target)) in train_loader.iter().enumerate() {
        let (mut data, target) = if args.cuda {
            (Variable::new(data.cuda(None).clone()), Variable::new(target.cuda(None).clone()))
        } else {
            (Variable::new(data.clone()), Variable::new(target.clone()))
        };
        optimizer.zero_grad();
        let output = model.f(&mut data);
        let mut loss = F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        loss.backward();
        optimizer.step();
        if batch_idx % args.log_interval == 0 {
            println!("Train Epoch: {} [{}/{} ({:.0}%)]\tLoss: {:.6}",
                     epoch,
                     batch_idx * data.data().len(),
                     train_loader.dataset.len(),
                     100. * (batch_idx as f32) / train_loader.len() as f32,
                     loss.data()[0]);
        }
    }
}

fn test(model: &mut Net<f32>, args: &NetArgs, test_loader: &D::BatchLoader<f32, i64>) {
    model.eval();
    let mut test_loss = 0.;
    let mut correct = 0;
    for (ref data, ref target) in test_loader.iter() {
        let varargs = VariableArgs {
            volatile: true,
            ..Default::default()
        };
        let (mut data, mut target) = if args.cuda {
            (Variable::new_args(data.cuda(None).clone(), &varargs),
             Variable::new(target.cuda(None).clone()))
        } else {
            (Variable::new_args(data.clone(), &varargs), Variable::new(target.clone()))
        };
        let mut output = model.f(&mut data);
        test_loss += F::nll_loss(&output, &target, &F::NLLLossArgs::default());
        let pred = output.data().max_reduce(1).1;
        correct += pred.eq_tensor(&*target.data()).cpu().sum() as u32;
    }
    test_loss /= test_loader.len() as f32;
    println!("\nTest set: Average loss: {:.4}, Accuracy: {}/{} ({:.0}%)\n",
             test_loss,
             correct,
             test_loader.dataset.len(),
             100. * (correct as f32) / (test_loader.dataset.len() as f32))

}

fn main() {
    // no data loaders yet so demonstrate with Vec placeholder

    let train_loader: D::BatchLoader<f32, i64> =
        D::DataLoader::new(datasets::MNIST::<f32>::build("../data").done(None),
                           D::DataLoaderArgs::default());

    let test_loader: D::BatchLoader<f32, i64> =
        D::DataLoader::new(datasets::MNIST::<f32>::build("../data").done(None),
                           D::DataLoaderArgs::default());
    let args = parse_args();
    let mut model = Net::new();
    let refmut: ModRefMut<Net<f32>>;
    {
        let p = &mut model as *mut Net<f32>;
        refmut = p.into();
    }
    let mut optimizer = optim::SGD::new(refmut,
                                        map_opt!{"lr" => args.lr, "momentum" => args.momentum});
    for epoch in 1..args.epochs + 1 {
        train(&mut model, &args, &train_loader, epoch, &mut optimizer);
        test(&mut model, &args, &test_loader);
    }
}
<|endoftext|>