<commit_before>use std::{f64, mem};
use std::ops::{Add, Mul};
use super::basis::occ;
use super::j_scheme::{DiagOpJ10, MopJ012, OpBlockJ200, OpJ100, OpJ200};
use super::linalg::{self, Conj, EigenvalueRange, Part, Transpose};
use super::mat::Mat;
use super::op::{Op, VectorMut};
use super::utils::Toler;

/// Hartree–Fock with ad hoc linear mixing.
#[derive(Clone, Copy, Debug)]
pub struct Conf {
    pub init_mix_factor: f64,
    pub mix_modifier: f64,
    pub toler: Toler,
    pub heevr_abstol: f64,
}

impl Default for Conf {
    fn default() -> Self {
        Self {
            init_mix_factor: 0.5,
            mix_modifier: 2.0,
            toler: Default::default(),
            heevr_abstol: Default::default(),
        }
    }
}

impl Conf {
    pub fn make_run<'a>(
        self,
        h1: &'a OpJ100<f64>,
        h2: &'a OpJ200<f64>,
    ) -> Run<'a> {
        let scheme = h1.scheme();
        let mut dcoeff = Op::new(scheme.clone());
        for p in scheme.states_10(&occ::ALL1) {
            dcoeff.set(p, p, 1.0);
        }
        Run {
            conf: self,
            h1,
            h2,
            energies: Op::new_vec(scheme.clone()),
            dcoeff,
            qcoeff: Op::new(scheme.clone()),
            fock: Op::new(scheme.clone()),
            fock_old: Op::new(scheme.clone()),
            energy_sum: f64::NAN,
            energy_change: 0.0,
            mix_factor: self.init_mix_factor,
            first: true,
        }
    }

    pub fn adjust_mix_factor(
        &self,
        de: f64,
        de_new: f64,
        mix_factor: &mut f64,
    ) {
        if de == 0.0 {
            return;
        }
        // adjust mixing depending on whether oscillation is present
        if de_new * de < 0.0 {
            // oscillating (don't make it too high though!)
            *mix_factor = f64::min(0.5, *mix_factor * self.mix_modifier);
        } else {
            // not oscillating
            *mix_factor *= 1.0 / self.mix_modifier;
        }
    }
}

pub struct Run<'a> {
    pub conf: Conf,
    pub h1: &'a OpJ100<f64>,
    pub h2: &'a OpJ200<f64>,
    pub energies: DiagOpJ10<f64>,
    pub dcoeff: OpJ100<f64>,
    pub qcoeff: OpJ100<f64>,
    pub fock: OpJ100<f64>,
    pub fock_old: OpJ100<f64>,
    pub energy_sum: f64,
    pub energy_change: f64,
    pub mix_factor: f64,
    pub first: bool,
}

impl<'a> Run<'a> {
    pub fn step(&mut self) -> Result<(), ()> {
        // Q[λ; v x] = ∑[i] D[λ; x i] D†[λ; i v]
        qcoeff(&self.dcoeff, &mut self.qcoeff);

        mem::swap(&mut self.fock_old, &mut self.fock);

        // compute Fock matrix
        self.fock.clone_from(&self.h1);
        fock2(&self.h2, &self.qcoeff, &mut self.fock);

        // mix in a little bit of the previous matrix to dampen oscillations
        // that can sometimes occur
        if self.first {
            self.first = false;
            block_mat_axpby(
                self.mix_factor,
                &self.fock_old.data.0,
                1.0 - self.mix_factor,
                &mut self.fock.data.0,
            );
        }

        // solve the HF eigenvalue problem.
        // ∑[w] f[u w] D[w g] = ε[g] D[u g]
        block_heevr(
            false,
            Part::Lower,
            &self.fock.data.0,
            self.conf.heevr_abstol,
            &mut self.energies.data.0,
            &mut self.dcoeff.data.0,
        ).unwrap();

        // test convergence using energy sum and adjust mixing
        let old_energy_sum = mem::replace(
            &mut self.energy_sum,
            weighted_sum(&self.energies),
        );
        let old_energy_change = mem::replace(
            &mut self.energy_change,
            self.energy_sum - old_energy_sum,
        );
        self.conf.adjust_mix_factor(
            old_energy_change,
            self.energy_change,
            &mut self.mix_factor,
        );
        if self.conf.toler.is_eq(self.energy_sum, old_energy_sum) {
            Ok(())
        } else {
            Err(())
        }
    }

    /// Iterates until the convergence criterion has been met.
    pub fn do_run(&mut self) -> Result<(), ()> {
        println!("hf:");
        for i in 0 .. 1024 {
            if self.step().is_ok() {
                println!("hf_converged: true");
                return Ok(());
            }
            println!("- {{iter: {}, energy_sum: {}, mix_factor: {}}}",
                     i, self.energy_sum, self.mix_factor);
        }
        println!("hf_converged: false");
        Err(())
    }
}

/// `y ← α × x + β × y`
pub fn block_mat_axpby<T>(
    alpha: T,
    x: &[Mat<T>],
    beta: T,
    y: &mut [Mat<T>],
) where
    T: Add<Output = T> + Mul<Output = T> + Clone,
{
    assert_eq!(x.len(), y.len());
    for l in 0 .. x.len() {
        linalg::mat_axpby(
            alpha.clone(),
            x[l].as_ref(),
            beta.clone(),
            y[l].as_mut(),
        );
    }
}

pub fn block_heevr<T: linalg::Heevr>(
    left: bool,
    part: Part,
    a: &[Mat<T>],
    abstol: T::Real,
    w: &mut [Vec<T::Real>],
    z: &mut [Mat<T>],
) -> Result<(), i32> where
    T::Real: Clone,
{
    assert_eq!(a.len(), w.len());
    assert_eq!(w.len(), z.len());
    let mut isuppz = Vec::new();
    for l in 0 .. a.len() {
        // make sure matrix `a` is not modified!
        let mut a_l = a[l].clone();
        linalg::heevr(
            left,
            EigenvalueRange::All,
            part,
            a_l.as_mut(),
            abstol.clone(),
            &mut w[l],
            z[l].as_mut(),
            Err(&mut isuppz),
        )?;
    }
    Ok(())
}

/// ```text
/// R = ∑[p] Jp E[p]
/// ```
pub fn weighted_sum(e1: &DiagOpJ10<f64>) -> f64 {
    let scheme = e1.scheme();
    let mut r = 0.0;
    for p in scheme.states_10(&occ::ALL1) {
        r += p.jweight(2) * e1.at(p, p);
    }
    r
}

/// ```text
/// Q[v x] = ∑[i] D[x i] D†[i v]
/// ```
pub fn qcoeff(
    d1: &OpJ100<f64>,
    q1: &mut OpJ100<f64>,
)
{
    let scheme = d1.scheme();
    q1.set_zero();
    for p in scheme.states_10(&occ::ALL1) {
        for q in p.costates_10(&occ::ALL1) {
            for i in p.costates_10(&[occ::I]) {
                q1.add(p, q, d1.at(q, i) * d1.at(p, i).conj());
            }
        }
    }
}

pub fn fock2(
    h2: &OpJ200<f64>,
    q1: &OpJ100<f64>,
    f1: &mut OpJ100<f64>,
)
{
    let scheme = h2.scheme();
    for pq in scheme.states_20(&occ::ALL2) {
        let (p, q) = pq.split_to_10_10();
        for r in p.costates_10(&occ::ALL1) {
            for s in q.costates_10(&occ::ALL1) {
                for rs in r.combine_with_10(s, pq.j()) {
                    f1.add(p, r, (
                        pq.jweight(2) / p.jweight(2)
                            * h2.at(pq, rs)
                            * q1.at(q, s)
                    ));
                }
            }
        }
    }
}

pub fn transform_h1(
    h1: &OpJ100<f64>,
    d1: &OpJ100<f64>,
    r1: &mut OpJ100<f64>,
)
{
    let scheme = h1.scheme();
    for p in scheme.states_10(&occ::ALL1) {
        for q in p.costates_10(&occ::ALL1) {
            for r in p.costates_10(&occ::ALL1) {
                for s in p.costates_10(&occ::ALL1) {
                    r1.add(p, s, (
                        h1.at(q, r)
                            * d1.at(q, p).conj()
                            * d1.at(r, s)
                    ));
                }
            }
        }
    }
}

pub fn transform_h2(
    h2: &OpJ200<f64>,
    d1: &OpJ100<f64>,
    r2: &mut OpJ200<f64>,
)
{
    let scheme = h2.scheme();
    for l in 0 .. scheme.basis_20.num_chans() {
        let mut og: OpBlockJ200<f64> = Op::new_block(scheme.clone(), l);
        let mut ohg: OpBlockJ200<f64> = Op::new_block(scheme.clone(), l);
        for pq in scheme.costates_20(l, &occ::ALL2) {
            let (p, q) = pq.split_to_10_10();
            let weight = pq.num_permut as f64;
            for rs in pq.costates_20(&occ::ALL2) {
                let (r, s) = rs.split_to_10_10();
                og.add(pq, rs, weight * d1.at(p, r) * d1.at(q, s));
            }
        }
        linalg::gemm(
            Transpose::None,
            Transpose::None,
            1.0,
            h2.data.0[l as usize].as_ref(),
            og.data.data.as_ref(),
            0.0,
            ohg.data.data.as_mut(),
        );
        linalg::gemm(
            Transpose::Conjugate,
            Transpose::None,
            1.0,
            og.data.data.as_ref(),
            ohg.data.data.as_ref(),
            1.0,
            r2.data.0[l as usize].as_mut(),
        );
    }
}

pub fn hf_energy(
    h1: &OpJ100<f64>,
    h2: &OpJ200<f64>,
) -> f64
{
    // ZN[p q] =
    //     ∑[i] Ji² U[i i]
    //     + 1/2 ∑[i j] Jij^2 V[i j i j]
    //     + 1/6 ∑[i j k] Jijk^2 W[i j k i j k]    (NYI)
    let scheme = h1.scheme();
    let mut r = 0.0;
    for i in scheme.states_10(&[occ::I]) {
        r += i.jweight(2) * h1.at(i, i);
    }
    for ij in scheme.states_20(&[occ::II]) {
        r += ij.jweight(2) * h2.at(ij, ij) / 2.0;
    }
    r
}

pub fn normord(
    h: &MopJ012<f64>,
    r: &mut MopJ012<f64>,
)
{
    // UN[p q] =
    //     U[p q]
    //     + ∑[i] (Jpi / Jp)² V[p i q i]
    //     + 1/2 ∑[i j] (Jpij/Jp)² W[p i j q i j]    (NYI)
    //
    // VN[p q r s] =
    //     V[p q r s]
    //     + ∑[i] (Jpqi / Jpq)² W[p q i r s i]    (NYI)
    let scheme = h.1.scheme();
    r.0 += h.0 + hf_energy(&h.1, &h.2);
    r.1.clone_from(&h.1);
    for pi in scheme.states_20(&[occ::II, occ::AI]) {
        let (p, i) = pi.split_to_10_10();
        for q in p.costates_10(&occ::ALL1) {
            for qi in q.combine_with_10(i, pi.j()) {
                r.1.add(p, q, pi.jweight(2) / p.jweight(2) * h.2.at(pi, qi));
            }
        }
    }
    r.2.clone_from(&h.2);
}
<commit_msg>normord should set, not add<commit_after>use std::{f64, mem};
use std::ops::{Add, Mul};
use super::basis::occ;
use super::j_scheme::{DiagOpJ10, MopJ012, OpBlockJ200, OpJ100, OpJ200};
use super::linalg::{self, Conj, EigenvalueRange, Part, Transpose};
use super::mat::Mat;
use super::op::{Op, VectorMut};
use super::utils::Toler;

/// Hartree–Fock with ad hoc linear mixing.
#[derive(Clone, Copy, Debug)]
pub struct Conf {
    pub init_mix_factor: f64,
    pub mix_modifier: f64,
    pub toler: Toler,
    pub heevr_abstol: f64,
}

impl Default for Conf {
    fn default() -> Self {
        Self {
            init_mix_factor: 0.5,
            mix_modifier: 2.0,
            toler: Default::default(),
            heevr_abstol: Default::default(),
        }
    }
}

impl Conf {
    pub fn make_run<'a>(
        self,
        h1: &'a OpJ100<f64>,
        h2: &'a OpJ200<f64>,
    ) -> Run<'a> {
        let scheme = h1.scheme();
        let mut dcoeff = Op::new(scheme.clone());
        for p in scheme.states_10(&occ::ALL1) {
            dcoeff.set(p, p, 1.0);
        }
        Run {
            conf: self,
            h1,
            h2,
            energies: Op::new_vec(scheme.clone()),
            dcoeff,
            qcoeff: Op::new(scheme.clone()),
            fock: Op::new(scheme.clone()),
            fock_old: Op::new(scheme.clone()),
            energy_sum: f64::NAN,
            energy_change: 0.0,
            mix_factor: self.init_mix_factor,
            first: true,
        }
    }

    pub fn adjust_mix_factor(
        &self,
        de: f64,
        de_new: f64,
        mix_factor: &mut f64,
    ) {
        if de == 0.0 {
            return;
        }
        // adjust mixing depending on whether oscillation is present
        if de_new * de < 0.0 {
            // oscillating (don't make it too high though!)
            *mix_factor = f64::min(0.5, *mix_factor * self.mix_modifier);
        } else {
            // not oscillating
            *mix_factor *= 1.0 / self.mix_modifier;
        }
    }
}

pub struct Run<'a> {
    pub conf: Conf,
    pub h1: &'a OpJ100<f64>,
    pub h2: &'a OpJ200<f64>,
    pub energies: DiagOpJ10<f64>,
    pub dcoeff: OpJ100<f64>,
    pub qcoeff: OpJ100<f64>,
    pub fock: OpJ100<f64>,
    pub fock_old: OpJ100<f64>,
    pub energy_sum: f64,
    pub energy_change: f64,
    pub mix_factor: f64,
    pub first: bool,
}

impl<'a> Run<'a> {
    pub fn step(&mut self) -> Result<(), ()> {
        // Q[λ; v x] = ∑[i] D[λ; x i] D†[λ; i v]
        qcoeff(&self.dcoeff, &mut self.qcoeff);

        mem::swap(&mut self.fock_old, &mut self.fock);

        // compute Fock matrix
        self.fock.clone_from(&self.h1);
        fock2(&self.h2, &self.qcoeff, &mut self.fock);

        // mix in a little bit of the previous matrix to dampen oscillations
        // that can sometimes occur
        if self.first {
            self.first = false;
            block_mat_axpby(
                self.mix_factor,
                &self.fock_old.data.0,
                1.0 - self.mix_factor,
                &mut self.fock.data.0,
            );
        }

        // solve the HF eigenvalue problem.
        // ∑[w] f[u w] D[w g] = ε[g] D[u g]
        block_heevr(
            false,
            Part::Lower,
            &self.fock.data.0,
            self.conf.heevr_abstol,
            &mut self.energies.data.0,
            &mut self.dcoeff.data.0,
        ).unwrap();

        // test convergence using energy sum and adjust mixing
        let old_energy_sum = mem::replace(
            &mut self.energy_sum,
            weighted_sum(&self.energies),
        );
        let old_energy_change = mem::replace(
            &mut self.energy_change,
            self.energy_sum - old_energy_sum,
        );
        self.conf.adjust_mix_factor(
            old_energy_change,
            self.energy_change,
            &mut self.mix_factor,
        );
        if self.conf.toler.is_eq(self.energy_sum, old_energy_sum) {
            Ok(())
        } else {
            Err(())
        }
    }

    /// Iterates until the convergence criterion has been met.
    pub fn do_run(&mut self) -> Result<(), ()> {
        println!("hf:");
        for i in 0 .. 1024 {
            if self.step().is_ok() {
                println!("hf_converged: true");
                return Ok(());
            }
            println!("- {{iter: {}, energy_sum: {}, mix_factor: {}}}",
                     i, self.energy_sum, self.mix_factor);
        }
        println!("hf_converged: false");
        Err(())
    }
}

/// `y ← α × x + β × y`
pub fn block_mat_axpby<T>(
    alpha: T,
    x: &[Mat<T>],
    beta: T,
    y: &mut [Mat<T>],
) where
    T: Add<Output = T> + Mul<Output = T> + Clone,
{
    assert_eq!(x.len(), y.len());
    for l in 0 .. x.len() {
        linalg::mat_axpby(
            alpha.clone(),
            x[l].as_ref(),
            beta.clone(),
            y[l].as_mut(),
        );
    }
}

pub fn block_heevr<T: linalg::Heevr>(
    left: bool,
    part: Part,
    a: &[Mat<T>],
    abstol: T::Real,
    w: &mut [Vec<T::Real>],
    z: &mut [Mat<T>],
) -> Result<(), i32> where
    T::Real: Clone,
{
    assert_eq!(a.len(), w.len());
    assert_eq!(w.len(), z.len());
    let mut isuppz = Vec::new();
    for l in 0 .. a.len() {
        // make sure matrix `a` is not modified!
        let mut a_l = a[l].clone();
        linalg::heevr(
            left,
            EigenvalueRange::All,
            part,
            a_l.as_mut(),
            abstol.clone(),
            &mut w[l],
            z[l].as_mut(),
            Err(&mut isuppz),
        )?;
    }
    Ok(())
}

/// ```text
/// R = ∑[p] Jp E[p]
/// ```
pub fn weighted_sum(e1: &DiagOpJ10<f64>) -> f64 {
    let scheme = e1.scheme();
    let mut r = 0.0;
    for p in scheme.states_10(&occ::ALL1) {
        r += p.jweight(2) * e1.at(p, p);
    }
    r
}

/// ```text
/// Q[v x] = ∑[i] D[x i] D†[i v]
/// ```
pub fn qcoeff(
    d1: &OpJ100<f64>,
    q1: &mut OpJ100<f64>,
)
{
    let scheme = d1.scheme();
    q1.set_zero();
    for p in scheme.states_10(&occ::ALL1) {
        for q in p.costates_10(&occ::ALL1) {
            for i in p.costates_10(&[occ::I]) {
                q1.add(p, q, d1.at(q, i) * d1.at(p, i).conj());
            }
        }
    }
}

pub fn fock2(
    h2: &OpJ200<f64>,
    q1: &OpJ100<f64>,
    f1: &mut OpJ100<f64>,
)
{
    let scheme = h2.scheme();
    for pq in scheme.states_20(&occ::ALL2) {
        let (p, q) = pq.split_to_10_10();
        for r in p.costates_10(&occ::ALL1) {
            for s in q.costates_10(&occ::ALL1) {
                for rs in r.combine_with_10(s, pq.j()) {
                    f1.add(p, r, (
                        pq.jweight(2) / p.jweight(2)
                            * h2.at(pq, rs)
                            * q1.at(q, s)
                    ));
                }
            }
        }
    }
}

pub fn transform_h1(
    h1: &OpJ100<f64>,
    d1: &OpJ100<f64>,
    r1: &mut OpJ100<f64>,
)
{
    let scheme = h1.scheme();
    for p in scheme.states_10(&occ::ALL1) {
        for q in p.costates_10(&occ::ALL1) {
            for r in p.costates_10(&occ::ALL1) {
                for s in p.costates_10(&occ::ALL1) {
                    r1.add(p, s, (
                        h1.at(q, r)
                            * d1.at(q, p).conj()
                            * d1.at(r, s)
                    ));
                }
            }
        }
    }
}

pub fn transform_h2(
    h2: &OpJ200<f64>,
    d1: &OpJ100<f64>,
    r2: &mut OpJ200<f64>,
)
{
    let scheme = h2.scheme();
    for l in 0 .. scheme.basis_20.num_chans() {
        let mut og: OpBlockJ200<f64> = Op::new_block(scheme.clone(), l);
        let mut ohg: OpBlockJ200<f64> = Op::new_block(scheme.clone(), l);
        for pq in scheme.costates_20(l, &occ::ALL2) {
            let (p, q) = pq.split_to_10_10();
            let weight = pq.num_permut as f64;
            for rs in pq.costates_20(&occ::ALL2) {
                let (r, s) = rs.split_to_10_10();
                og.add(pq, rs, weight * d1.at(p, r) * d1.at(q, s));
            }
        }
        linalg::gemm(
            Transpose::None,
            Transpose::None,
            1.0,
            h2.data.0[l as usize].as_ref(),
            og.data.data.as_ref(),
            0.0,
            ohg.data.data.as_mut(),
        );
        linalg::gemm(
            Transpose::Conjugate,
            Transpose::None,
            1.0,
            og.data.data.as_ref(),
            ohg.data.data.as_ref(),
            1.0,
            r2.data.0[l as usize].as_mut(),
        );
    }
}

pub fn hf_energy(
    h1: &OpJ100<f64>,
    h2: &OpJ200<f64>,
) -> f64
{
    // ZN[p q] =
    //     ∑[i] Ji² U[i i]
    //     + 1/2 ∑[i j] Jij^2 V[i j i j]
    //     + 1/6 ∑[i j k] Jijk^2 W[i j k i j k]    (NYI)
    let scheme = h1.scheme();
    let mut r = 0.0;
    for i in scheme.states_10(&[occ::I]) {
        r += i.jweight(2) * h1.at(i, i);
    }
    for ij in scheme.states_20(&[occ::II]) {
        r += ij.jweight(2) * h2.at(ij, ij) / 2.0;
    }
    r
}

pub fn normord(
    h: &MopJ012<f64>,
    r: &mut MopJ012<f64>,
)
{
    // UN[p q] =
    //     U[p q]
    //     + ∑[i] (Jpi / Jp)² V[p i q i]
    //     + 1/2 ∑[i j] (Jpij/Jp)² W[p i j q i j]    (NYI)
    //
    // VN[p q r s] =
    //     V[p q r s]
    //     + ∑[i] (Jpqi / Jpq)² W[p q i r s i]    (NYI)
    let scheme = h.1.scheme();
    r.0 = h.0 + hf_energy(&h.1, &h.2);
    r.1.clone_from(&h.1);
    for pi in scheme.states_20(&[occ::II, occ::AI]) {
        let (p, i) = pi.split_to_10_10();
        for q in p.costates_10(&occ::ALL1) {
            for qi in q.combine_with_10(i, pi.j()) {
                r.1.add(p, q, pi.jweight(2) / p.jweight(2) * h.2.at(pi, qi));
            }
        }
    }
    r.2.clone_from(&h.2);
}
<|endoftext|>


<commit_before>use grammar::*;
use asm_ops::*;
use asm_ops::AsmOperand::*;
use asm_ops::AsmOp::*;

struct AsmProgram {
    contents : Vec<AsmOp>
}

impl AsmProgram{
    fn add(&mut self, op: AsmOp){
        println!("{:?}", &op);
        self.contents.push(op);
    }

    fn new() -> AsmProgram {
        AsmProgram {contents: Vec::new()}
    }
}

impl Extend<AsmOp> for AsmProgram {
    fn extend<T: IntoIterator<Item=AsmOp>>(&mut self, iterable: T) {
        for elem in iterable {
            self.add(elem);
        }
    }
}

fn move_into_expr(expr: &Expr, var_store: &VarStore) -> Vec<AsmOp>{
    let mut ops = Vec::new();
    match *expr {
        Expr::Num(ref num) => ops.push(Add(Register::RAX, Value(*num))),
        Expr::AddSub(ref terms)  => {
            ops.push(Push(RegisterOperand(Register::RAX)));
            ops.extend(add(terms.as_slice(), var_store));
            ops.push(Pop(RegisterOperand(Register::RBX)));
            ops.push(Pop(RegisterOperand(Register::RAX)));
            ops.push(Add(Register::RAX, RegisterOperand(Register::RBX)));
        },
        Expr::MultDiv(ref terms) => {
            ops.push(Push(RegisterOperand(Register::RAX)));
            ops.extend(mult(terms.as_slice(), var_store));
            ops.push(Pop(RegisterOperand(Register::RBX)));
            ops.push(Pop(RegisterOperand(Register::RAX)));
            ops.push(Mul(Register::RAX, RegisterOperand(Register::RBX)));
        },
        Expr::Variable(ref name) => ops.push(Add(Register::RAX, Memory(var_store.get_var_address_r(name))))
    }
    ops
}

fn move_into_first_term_expr(expr: &Expr, var_store: &VarStore) -> Vec<AsmOp>{
    let mut ops = Vec::new();
    match *expr {
        Expr::Num(ref num) => ops.push(Mov(Register::RAX, Value(*num))),
        Expr::AddSub(ref terms) => {
            ops.extend(add(terms.as_slice(), var_store));
            ops.push(Pop(RegisterOperand(Register::RAX)));
        },
        Expr::MultDiv(ref terms) => {
            ops.extend(mult(terms.as_slice(), var_store));
            ops.push(Pop(RegisterOperand(Register::RAX)));
        },
        Expr::Variable(ref name) => ops.push(Mov(Register::RAX, Memory(var_store.get_var_address_r(name))))
    }
    ops
}

fn add(terms: &[AddTerm], var_store: &VarStore) -> Vec<AsmOp>{
    let mut ops = Vec::new();
    let mut counter = 0;
    for term in terms {
        let &AddTerm(_, ref expr) = term;
        if counter == 0{
            ops.extend(move_into_first_term_expr(expr, var_store))
        } else {
            ops.extend(move_into_expr(expr, var_store));
        }
        counter +=1;
    }
    ops.push(Push(RegisterOperand(Register::RAX)));
    ops
}

fn mult(terms: &[MultTerm], var_store: &VarStore) -> Vec<AsmOp>{
    let mut ops = Vec::new();
    let mut counter = 0;
    for term in terms {
        let &MultTerm(_, ref expr) = term;
        if counter == 0{
            ops.extend(move_into_first_term_expr(expr, var_store))
        } else {
            ops.extend(move_into_expr(expr, var_store));
        }
        counter +=1;
    }
    ops.push(Push(RegisterOperand(Register::RAX)));
    ops
}

use std::collections::HashMap;

struct VarStore{
    variables: HashMap<String, u16>,
    current_address: u16
}

impl VarStore{
    fn new() -> VarStore{
        VarStore{variables: HashMap::new(), current_address: 0}
    }

    fn get_var_address_r(&self, name: &String) -> u16{
        match self.variables.get(name){
            Some(address) => *address,
            None => panic!("No variable named {}.", name)
        }
    }
    fn get_var_address_l(&mut self, name: &String) -> u16{
        if !self.variables.contains_key(name){
            let result = self.variables.insert(name.to_string(), self.current_address + 100);
            match result{
                Some(_) => panic!(),
                None =>{
                    self.current_address += 8;
                    self.get_var_address_r(name)
                }
            }
        }else{
            self.get_var_address_r( name)
        }
    }
}

fn calculate_expr(expr: &Expr, var_store: &VarStore) -> Vec<AsmOp>{
    let mut ops = Vec::new();
    match *expr {
        Expr::AddSub(ref terms) => ops.extend(add(terms.as_slice(), var_store)),
        Expr::MultDiv(ref terms) => ops.extend(mult(terms.as_slice(), var_store)),
        Expr::Num(ref num) => ops.push(Push(Value(*num))),
        Expr::Variable(ref name) => ops.push(Push(Memory(var_store.get_var_address_r(name)))),
    }
    ops
}

pub fn translate(block: &Vec<Statement>) -> Box<Vec<AsmOp>>{
    let mut ops =  AsmProgram::new();
    let mut var_store = VarStore::new();
    for stmt in block {
        println!("\n{:?}\nIs translated into:", stmt);
        match *stmt {
            Statement::Assign(ref name, ref expr) => {
                ops.extend(calculate_expr(expr, &var_store));
                ops.add(Pop(Memory(var_store.get_var_address_l(name))));
            }
            Statement::Output(ref expr) => {
                ops.extend(calculate_expr(expr, &var_store));
                ops.add(Pop(RegisterOperand(Register::RAX)));
            }
            _ => {}
        }
    }
    Box::new(ops.contents)
}
<commit_msg>Increased an offset for variable storages<commit_after>use grammar::*;
use asm_ops::*;
use asm_ops::AsmOperand::*;
use asm_ops::AsmOp::*;

struct AsmProgram {
    contents : Vec<AsmOp>
}

impl AsmProgram{
    fn add(&mut self, op: AsmOp){
        println!("{:?}", &op);
        self.contents.push(op);
    }

    fn new() -> AsmProgram {
        AsmProgram {contents: Vec::new()}
    }
}

impl Extend<AsmOp> for AsmProgram {
    fn extend<T: IntoIterator<Item=AsmOp>>(&mut self, iterable: T) {
        for elem in iterable {
            self.add(elem);
        }
    }
}

fn move_into_expr(expr: &Expr, var_store: &VarStore) -> Vec<AsmOp>{
    let mut ops = Vec::new();
    match *expr {
        Expr::Num(ref num) => ops.push(Add(Register::RAX, Value(*num))),
        Expr::AddSub(ref terms)  => {
            ops.push(Push(RegisterOperand(Register::RAX)));
            ops.extend(add(terms.as_slice(), var_store));
            ops.push(Pop(RegisterOperand(Register::RBX)));
            ops.push(Pop(RegisterOperand(Register::RAX)));
            ops.push(Add(Register::RAX, RegisterOperand(Register::RBX)));
        },
        Expr::MultDiv(ref terms) => {
            ops.push(Push(RegisterOperand(Register::RAX)));
            ops.extend(mult(terms.as_slice(), var_store));
            ops.push(Pop(RegisterOperand(Register::RBX)));
            ops.push(Pop(RegisterOperand(Register::RAX)));
            ops.push(Mul(Register::RAX, RegisterOperand(Register::RBX)));
        },
        Expr::Variable(ref name) => ops.push(Add(Register::RAX, Memory(var_store.get_var_address_r(name))))
    }
    ops
}

fn move_into_first_term_expr(expr: &Expr, var_store: &VarStore) -> Vec<AsmOp>{
    let mut ops = Vec::new();
    match *expr {
        Expr::Num(ref num) => ops.push(Mov(Register::RAX, Value(*num))),
        Expr::AddSub(ref terms) => {
            ops.extend(add(terms.as_slice(), var_store));
            ops.push(Pop(RegisterOperand(Register::RAX)));
        },
        Expr::MultDiv(ref terms) => {
            ops.extend(mult(terms.as_slice(), var_store));
            ops.push(Pop(RegisterOperand(Register::RAX)));
        },
        Expr::Variable(ref name) => ops.push(Mov(Register::RAX, Memory(var_store.get_var_address_r(name))))
    }
    ops
}

fn add(terms: &[AddTerm], var_store: &VarStore) -> Vec<AsmOp>{
    let mut ops = Vec::new();
    let mut counter = 0;
    for term in terms {
        let &AddTerm(_, ref expr) = term;
        if counter == 0{
            ops.extend(move_into_first_term_expr(expr, var_store))
        } else {
            ops.extend(move_into_expr(expr, var_store));
        }
        counter +=1;
    }
    ops.push(Push(RegisterOperand(Register::RAX)));
    ops
}

fn mult(terms: &[MultTerm], var_store: &VarStore) -> Vec<AsmOp>{
    let mut ops = Vec::new();
    let mut counter = 0;
    for term in terms {
        let &MultTerm(_, ref expr) = term;
        if counter == 0{
            ops.extend(move_into_first_term_expr(expr, var_store))
        } else {
            ops.extend(move_into_expr(expr, var_store));
        }
        counter +=1;
    }
    ops.push(Push(RegisterOperand(Register::RAX)));
    ops
}

use std::collections::HashMap;

struct VarStore{
    variables: HashMap<String, u16>,
    current_address: u16
}

impl VarStore{
    fn new() -> VarStore{
        VarStore{variables: HashMap::new(), current_address: 0}
    }

    fn get_var_address_r(&self, name: &String) -> u16{
        match self.variables.get(name){
            Some(address) => *address,
            None => panic!("No variable named {}.", name)
        }
    }
    fn get_var_address_l(&mut self, name: &String) -> u16{
        if !self.variables.contains_key(name){
            let result = self.variables.insert(name.to_string(), self.current_address + 500);
            match result{
                Some(_) => panic!(),
                None =>{
                    self.current_address += 8;
                    self.get_var_address_r(name)
                }
            }
        }else{
            self.get_var_address_r( name)
        }
    }
}

fn calculate_expr(expr: &Expr, var_store: &VarStore) -> Vec<AsmOp>{
    let mut ops = Vec::new();
    match *expr {
        Expr::AddSub(ref terms) => ops.extend(add(terms.as_slice(), var_store)),
        Expr::MultDiv(ref terms) => ops.extend(mult(terms.as_slice(), var_store)),
        Expr::Num(ref num) => ops.push(Push(Value(*num))),
        Expr::Variable(ref name) => ops.push(Push(Memory(var_store.get_var_address_r(name)))),
    }
    ops
}

pub fn translate(block: &Vec<Statement>) -> Box<Vec<AsmOp>>{
    let mut ops =  AsmProgram::new();
    let mut var_store = VarStore::new();
    for stmt in block {
        println!("\n{:?}\nIs translated into:", stmt);
        match *stmt {
            Statement::Assign(ref name, ref expr) => {
                ops.extend(calculate_expr(expr, &var_store));
                ops.add(Pop(Memory(var_store.get_var_address_l(name))));
            }
            Statement::Output(ref expr) => {
                ops.extend(calculate_expr(expr, &var_store));
                ops.add(Pop(RegisterOperand(Register::RAX)));
            }
            _ => {}
        }
    }
    Box::new(ops.contents)
}
<|endoftext|>


<commit_before>use std::iter::repeat;
use std::io::{self, Write};
use time::{self, Timespec, Duration};
use tty::{Width, terminal_size};

macro_rules! printfl {
    ($($tt:tt)*) => {{
        print!($($tt)*);
        io::stdout().flush().ok().expect("flush() fail");
    }}
}

macro_rules! kb_fmt {
    ($n: ident) => {{
        let kb = 1024f64;
        match $n {
            $n if $n >= kb.powf(4_f64) => format!("{:.*} TB", 2, $n / kb.powf(4_f64)),
            $n if $n >= kb.powf(3_f64) => format!("{:.*} GB", 2, $n / kb.powf(3_f64)),
            $n if $n >= kb.powf(2_f64) => format!("{:.*} MB", 2, $n / kb.powf(2_f64)),
            $n if $n >= kb => format!("{:.*} KB", 2, $n / kb),
            _ => format!("{:.*} B", 0, $n)
        }
    }}
}

macro_rules! repeat {
    ($s: expr, $n: expr) => {{
        &repeat($s).take($n).collect::<String>()
    }}
}

static FORMAT: &'static str = "[=>-]";

// Output type format, indicate which format wil be used in
// the speed box.
#[derive(Debug)]
pub enum Units {
    Default,
    Bytes,
}

#[derive(Debug)]
pub struct ProgressBar {
    start_time: Timespec,
    units: Units,
    total: usize,
    current: usize,
    bar_start: String,
    bar_current: String,
    bar_current_n: String,
    bar_remain: String,
    bar_end: String,
    pub is_finish: bool,
    pub show_bar: bool,
    pub show_speed: bool,
    pub show_percent: bool,
    pub show_counter: bool,
    pub show_time_left: bool,
}

impl ProgressBar {
    /// Create a new ProgressBar with default configuration.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::thread;
    /// use pbr::{ProgressBar, Units};
    ///
    /// let count = 1000;
    /// let mut pb = ProgressBar::new(count);
    /// pb.set_units(Units::Bytes);
    ///
    /// for _ in 0..count {
    ///    pb.inc();
    ///    thread::sleep_ms(100);
    /// }
    /// ```
    pub fn new(total: usize) -> ProgressBar {
        let mut pb = ProgressBar {
            total: total,
            current: 0,
            start_time: time::get_time(),
            units: Units::Default,
            is_finish: false,
            show_bar: true,
            show_speed: true,
            show_percent: true,
            show_counter: true,
            show_time_left: true,
            bar_start: String::new(),
            bar_current: String::new(),
            bar_current_n: String::new(),
            bar_remain: String::new(),
            bar_end: String::new(),
        };
        pb.format(FORMAT);
        pb
    }

    /// set units, default is simple numbers
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pbr::{ProgressBar, Units};
    ///
    /// let mut pb = ProgressBar::new(n_bytes);
    /// pb.set_units(Units::Bytes);
    /// ```
    pub fn set_units(&mut self, u: Units) {
        self.units = u;
    }

    /// Set custom format to the drawing bar, default is `[=>-]`
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut pb = ProgressBar::new(...);
    /// pb.format("[=>_]");
    /// ```
    pub fn format(&mut self, fmt: &str) {
        if fmt.len() >= 5 {
            let v: Vec<&str> = fmt.split("").collect();
            self.bar_start = v[1].to_string();
            self.bar_current = v[2].to_string();
            self.bar_current_n = v[3].to_string();
            self.bar_remain = v[4].to_string();
            self.bar_end = v[5].to_string();
        }
    }

    /// Add to current value
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut pb = ProgressBar::new(10);
    /// pb.add(5);
    /// pb.finish();
    /// ```
    pub fn add(&mut self, i: usize) -> usize {
        self.current += i;
        if self.current <= self.total {
            self.draw()
        };
        self.current
    }

    /// Increment current value
    pub fn inc(&mut self) -> usize {
        self.add(1)
    }

    fn draw(&self) {
        let tty_size = terminal_size();
        let width = if let Some((Width(w), _)) = tty_size {
            w as usize
        } else {
            80
        };
        let mut base = String::new();
        let mut suffix = String::new();
        let mut prefix = String::new();
        let mut out;
        // precent box
        if self.show_percent {
            let percent = self.current as f64 / (self.total as f64 / 100f64);
            suffix = suffix + &format!(" {:.*} % ", 2, percent);
        }
        // speed box
        if self.show_speed {
            let from_start = (time::get_time() - self.start_time).num_nanoseconds().unwrap() as f64;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as f64;
            let speed = self.current as f64 / (from_start / sec_nano);
            suffix = match self.units {
                Units::Default => suffix + &format!("{:.*}/s ", 2, speed),
                Units::Bytes => suffix + &format!("{}/s ", kb_fmt!(speed)),
            };
        }
        // time left box
        if self.show_time_left {
            let from_start = time::get_time() - self.start_time;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as i32;
            let per_entry = from_start / self.current as i32;
            let mut left = per_entry * (self.total - self.current) as i32;
            left = (left / sec_nano) * sec_nano;
            if left.num_seconds() > 0 {
                if left.num_seconds() < Duration::minutes(1).num_seconds() {
                    suffix = suffix + &format!("{}s", left.num_seconds());
                } else {
                    suffix = suffix + &format!("{}m", left.num_minutes());
                }
            }
        }
        // counter box
        if self.show_counter {
            let (c, t) = (self.current as f64, self.total as f64);
            prefix = match self.units {
                Units::Default => format!("{} / {} ", c, t),
                Units::Bytes => format!("{} / {} ", kb_fmt!(c), kb_fmt!(t)),
            };
        }
        // bar box
        if self.show_bar {
            let size = width - (prefix.len() + suffix.len() + 3);
            if size > 0 {
                let curr_count =
                    ((self.current as f64 / self.total as f64) * size as f64).ceil() as usize;
                let rema_count = size - curr_count;
                base = self.bar_start.to_string();
                if rema_count > 0 {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count - 1) +
                           &self.bar_current_n;
                } else {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count);
                }
                base = base + repeat!(self.bar_remain.as_ref(), rema_count) + &self.bar_end;
            }
        }
        out = prefix + &base + &suffix;
        // pad
        if out.len() < width {
            let gap = width - out.len();
            out = out + repeat!(" ", gap);
        }
        // print
        printfl!("\r{}", out);
    }

    /// Calling finish manually will set current to total and draw
    /// the last time
    pub fn finish(&mut self) {
        if self.current < self.total {
            self.current = self.total;
            self.draw();
        }
        println!("");
        self.is_finish = true;
    }
}

// Implement io::Writer
impl Write for ProgressBar {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let n = buf.len();
        self.add(n);
        Ok(n)
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use pb::ProgressBar;

    #[test]
    fn add() {
        let mut pb = ProgressBar::new(10);
        pb.add(2);
        assert!(pb.current == 2,
                "should add the given `n` to current");
        assert!(pb.add(2) == pb.current,
                "add should return the current value");
    }

    #[test]
    fn inc() {
        let mut pb = ProgressBar::new(10);
        pb.inc();
        assert!(pb.current == 1, "should increment current by 1");
    }

    #[test]
    fn format() {
        let fmt = "[~> ]";
        let mut pb = ProgressBar::new(1);
        pb.format(fmt);
        assert!(pb.bar_start + &pb.bar_current + &pb.bar_current_n + &pb.bar_remain +
                &pb.bar_end == fmt);
    }

    #[test]
    fn finish() {
        let mut pb = ProgressBar::new(10);
        pb.finish();
        assert!(pb.current == pb.total,
                "should set current to total");
        assert!(pb.is_finish, "should set is_finish to true");
    }

    #[test]
    fn kb_fmt() {
        let kb = 1024f64;
        let mb = kb.powf(2f64);
        let gb = kb.powf(3f64);
        let tb = kb.powf(4f64);
        assert_eq!(kb_fmt!(kb), "1.00 KB");
        assert_eq!(kb_fmt!(mb), "1.00 MB");
        assert_eq!(kb_fmt!(gb), "1.00 GB");
        assert_eq!(kb_fmt!(tb), "1.00 TB");
    }
}
<commit_msg>refac(pb): fix code linting<commit_after>use std::iter::repeat;
use std::io::{self, Write};
use time::{self, Timespec, Duration};
use tty::{Width, terminal_size};

macro_rules! printfl {
    ($($tt:tt)*) => {{
        print!($($tt)*);
        io::stdout().flush().ok().expect("flush() fail");
    }}
}

macro_rules! kb_fmt {
    ($n: ident) => {{
        let kb = 1024f64;
        match $n {
            $n if $n >= kb.powf(4_f64) => format!("{:.*} TB", 2, $n / kb.powf(4_f64)),
            $n if $n >= kb.powf(3_f64) => format!("{:.*} GB", 2, $n / kb.powf(3_f64)),
            $n if $n >= kb.powf(2_f64) => format!("{:.*} MB", 2, $n / kb.powf(2_f64)),
            $n if $n >= kb => format!("{:.*} KB", 2, $n / kb),
            _ => format!("{:.*} B", 0, $n)
        }
    }}
}

macro_rules! repeat {
    ($s: expr, $n: expr) => {{
        &repeat($s).take($n).collect::<String>()
    }}
}

static FORMAT: &'static str = "[=>-]";

// Output type format, indicate which format wil be used in
// the speed box.
#[derive(Debug)]
pub enum Units {
    Default,
    Bytes,
}

#[derive(Debug)]
pub struct ProgressBar {
    start_time: Timespec,
    units: Units,
    total: usize,
    current: usize,
    bar_start: String,
    bar_current: String,
    bar_current_n: String,
    bar_remain: String,
    bar_end: String,
    pub is_finish: bool,
    pub show_bar: bool,
    pub show_speed: bool,
    pub show_percent: bool,
    pub show_counter: bool,
    pub show_time_left: bool,
}

impl ProgressBar {
    /// Create a new ProgressBar with default configuration.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::thread;
    /// use pbr::{ProgressBar, Units};
    ///
    /// let count = 1000;
    /// let mut pb = ProgressBar::new(count);
    /// pb.set_units(Units::Bytes);
    ///
    /// for _ in 0..count {
    ///    pb.inc();
    ///    thread::sleep_ms(100);
    /// }
    /// ```
    pub fn new(total: usize) -> ProgressBar {
        let mut pb = ProgressBar {
            total: total,
            current: 0,
            start_time: time::get_time(),
            units: Units::Default,
            is_finish: false,
            show_bar: true,
            show_speed: true,
            show_percent: true,
            show_counter: true,
            show_time_left: true,
            bar_start: String::new(),
            bar_current: String::new(),
            bar_current_n: String::new(),
            bar_remain: String::new(),
            bar_end: String::new(),
        };
        pb.format(FORMAT);
        pb
    }

    /// set units, default is simple numbers
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pbr::{ProgressBar, Units};
    ///
    /// let mut pb = ProgressBar::new(n_bytes);
    /// pb.set_units(Units::Bytes);
    /// ```
    pub fn set_units(&mut self, u: Units) {
        self.units = u;
    }

    /// Set custom format to the drawing bar, default is `[=>-]`
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut pb = ProgressBar::new(...);
    /// pb.format("[=>_]");
    /// ```
    pub fn format(&mut self, fmt: &str) {
        if fmt.len() >= 5 {
            let v: Vec<&str> = fmt.split("").collect();
            self.bar_start = v[1].to_owned();
            self.bar_current = v[2].to_owned();
            self.bar_current_n = v[3].to_owned();
            self.bar_remain = v[4].to_owned();
            self.bar_end = v[5].to_owned();
        }
    }

    /// Add to current value
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut pb = ProgressBar::new(10);
    /// pb.add(5);
    /// pb.finish();
    /// ```
    pub fn add(&mut self, i: usize) -> usize {
        self.current += i;
        if self.current <= self.total {
            self.draw()
        };
        self.current
    }

    /// Increment current value
    pub fn inc(&mut self) -> usize {
        self.add(1)
    }

    fn draw(&self) {
        let tty_size = terminal_size();
        let width = if let Some((Width(w), _)) = tty_size {
            w as usize
        } else {
            80
        };
        let mut base = String::new();
        let mut suffix = String::new();
        let mut prefix = String::new();
        let mut out;
        // precent box
        if self.show_percent {
            let percent = self.current as f64 / (self.total as f64 / 100f64);
            suffix = suffix + &format!(" {:.*} % ", 2, percent);
        }
        // speed box
        if self.show_speed {
            let from_start = (time::get_time() - self.start_time).num_nanoseconds().unwrap() as f64;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as f64;
            let speed = self.current as f64 / (from_start / sec_nano);
            suffix = match self.units {
                Units::Default => suffix + &format!("{:.*}/s ", 2, speed),
                Units::Bytes => suffix + &format!("{}/s ", kb_fmt!(speed)),
            };
        }
        // time left box
        if self.show_time_left {
            let from_start = time::get_time() - self.start_time;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as i32;
            let per_entry = from_start / self.current as i32;
            let mut left = per_entry * (self.total - self.current) as i32;
            left = (left / sec_nano) * sec_nano;
            if left.num_seconds() > 0 {
                if left.num_seconds() < Duration::minutes(1).num_seconds() {
                    suffix = suffix + &format!("{}s", left.num_seconds());
                } else {
                    suffix = suffix + &format!("{}m", left.num_minutes());
                }
            }
        }
        // counter box
        if self.show_counter {
            let (c, t) = (self.current as f64, self.total as f64);
            prefix = match self.units {
                Units::Default => format!("{} / {} ", c, t),
                Units::Bytes => format!("{} / {} ", kb_fmt!(c), kb_fmt!(t)),
            };
        }
        // bar box
        if self.show_bar {
            let size = width - (prefix.len() + suffix.len() + 3);
            if size > 0 {
                let curr_count =
                    ((self.current as f64 / self.total as f64) * size as f64).ceil() as usize;
                let rema_count = size - curr_count;
                base = self.bar_start.clone();
                if rema_count > 0 {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count - 1) +
                           &self.bar_current_n;
                } else {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count);
                }
                base = base + repeat!(self.bar_remain.as_ref(), rema_count) + &self.bar_end;
            }
        }
        out = prefix + &base + &suffix;
        // pad
        if out.len() < width {
            let gap = width - out.len();
            out = out + repeat!(" ", gap);
        }
        // print
        printfl!("\r{}", out);
    }

    /// Calling finish manually will set current to total and draw
    /// the last time
    pub fn finish(&mut self) {
        if self.current < self.total {
            self.current = self.total;
            self.draw();
        }
        println!("");
        self.is_finish = true;
    }
}

// Implement io::Writer
impl Write for ProgressBar {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let n = buf.len();
        self.add(n);
        Ok(n)
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use pb::ProgressBar;

    #[test]
    fn add() {
        let mut pb = ProgressBar::new(10);
        pb.add(2);
        assert!(pb.current == 2,
                "should add the given `n` to current");
        assert!(pb.add(2) == pb.current,
                "add should return the current value");
    }

    #[test]
    fn inc() {
        let mut pb = ProgressBar::new(10);
        pb.inc();
        assert!(pb.current == 1, "should increment current by 1");
    }

    #[test]
    fn format() {
        let fmt = "[~> ]";
        let mut pb = ProgressBar::new(1);
        pb.format(fmt);
        assert!(pb.bar_start + &pb.bar_current + &pb.bar_current_n + &pb.bar_remain +
                &pb.bar_end == fmt);
    }

    #[test]
    fn finish() {
        let mut pb = ProgressBar::new(10);
        pb.finish();
        assert!(pb.current == pb.total,
                "should set current to total");
        assert!(pb.is_finish, "should set is_finish to true");
    }

    #[test]
    fn kb_fmt() {
        let kb = 1024f64;
        let mb = kb.powf(2f64);
        let gb = kb.powf(3f64);
        let tb = kb.powf(4f64);
        assert_eq!(kb_fmt!(kb), "1.00 KB");
        assert_eq!(kb_fmt!(mb), "1.00 MB");
        assert_eq!(kb_fmt!(gb), "1.00 GB");
        assert_eq!(kb_fmt!(tb), "1.00 TB");
    }
}
<|endoftext|>


<commit_before>use ffi::prelude::LLVMTypeRef;
use ffi::{core, target, LLVMTypeKind};
use libc::{c_int, c_uint};
use compile::Compile;
use context::{Context, GetContext};
use target::TargetData;
use util::{self, CastFrom};
use std::{fmt, mem};
use std::iter::Iterator;
use std::ops::Deref;

/// Defines how a value should be laid out in memory.
pub struct Type;
native_ref!(&Type = LLVMTypeRef);
impl Type {
    #[inline(always)]
    /// Get the type given as an LLVM type descriptor in the context given.
    pub fn get<'a, T>(context:&'a Context) -> &'a Type where T:Compile<'a> {
        T::get_type(context)
    }
    /// Make a new function signature with the return type and arguments given.
    pub fn new_function<'a>(ret: &'a Type, args: &[&'a Type]) -> &'a FunctionType {
        unsafe { core::LLVMFunctionType(ret.into(), args.as_ptr() as *mut LLVMTypeRef, args.len() as c_uint, 0) }.into()
    }
    /// Make a new array with the length given.
    pub fn new_array<'a>(element: &'a Type, length: usize) -> &'a Type {
        unsafe { core::LLVMArrayType(element.into(), length as c_uint) }.into()
    }
    /// Make a new vector with the length given.
    pub fn new_vector<'a>(element: &'a Type, length: usize) -> &'a Type {
        unsafe { core::LLVMVectorType(element.into(), length as c_uint) }.into()
    }
    /// Make a new pointer with the given element type.
    pub fn new_pointer<'a>(elem: &'a Type) -> &'a Type {
        unsafe { core::LLVMPointerType(elem.into(), 0 as c_uint) }.into()
    }
    /// Make a new structure type with the given types.
    pub fn new_struct<'a>(elems: &[&'a Type], packed: bool) -> &'a Type {
        unsafe { core::LLVMStructType(elems.as_ptr() as *mut LLVMTypeRef, elems.len() as c_uint, packed as c_int) }.into()
    }
    /// Returns true if the size of the type is known at compile-time.
    ///
    /// This is equivalent to the type implementing `Sized` in Rust
    pub fn is_sized(&self) -> bool {
        unsafe { core::LLVMTypeIsSized(self.into()) != 0 }
    }
    /// Returns true if this type is a function.
    pub fn is_function(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMFunctionTypeKind as c_uint
    }
    /// Returns true if this type is a struct.
    pub fn is_struct(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMStructTypeKind as c_uint
    }
    /// Returns true if this type is void.
    pub fn is_void(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMVoidTypeKind as c_uint
    }
    /// Returns the size of the type in bytes.
    pub fn get_size(&self, target: &TargetData) -> usize {
        unsafe { target::LLVMABISizeOfType(target.into(), self.into()) as usize }
    }
    /// Returns the element of this pointer type.
    pub fn get_element(&self) -> Option<&Type> {
        unsafe { mem::transmute(core::LLVMGetElementType(self.into())) }
    }
}
get_context!(Type, LLVMGetTypeContext);
to_str!(Type, LLVMPrintTypeToString);

/// A structure type, such as a tuple or struct
pub struct StructType;
native_ref!(&StructType = LLVMTypeRef);
impl StructType {
    /// Make a new struct with the given fields and packed representation.
    pub fn new<'a>(context: &'a Context, fields: &[&'a Type], packed: bool) -> &'a StructType {
        unsafe { core::LLVMStructTypeInContext(context.into(), fields.as_ptr() as *mut LLVMTypeRef, fields.len() as c_uint, packed as c_int) }.into()
    }
    /// Make a new named struct with the given fields and packed representation.
    pub fn new_named<'a>(context: &'a Context, name: &str, fields: &[&'a Type], packed: bool) -> &'a StructType {
        util::with_cstr(name, |name| unsafe {
            let ty = core::LLVMStructCreateNamed(context.into(), name);
            core::LLVMStructSetBody(ty, fields.as_ptr() as *mut LLVMTypeRef, fields.len() as c_uint, packed as c_int);
            ty.into()
        })
    }
    /// Returns the elements (AKA fields) that make up this struct.
    pub fn get_elements(&self) -> Vec<Type> {
        unsafe {
            let size = core::LLVMCountStructElementTypes(self.into());
            let mut els:Vec<_> = (0..size).map(|_| mem::uninitialized()).collect();
            core::LLVMGetStructElementTypes(self.into(), els.as_mut_ptr() as *mut LLVMTypeRef);
            els
        }
    }
}
impl CastFrom for StructType {
    type From = Type;
    fn cast(ty: &Type) -> Option<&StructType> {
        unsafe {
            let kind = core::LLVMGetTypeKind(ty.into());
            if kind as c_uint == LLVMTypeKind::LLVMStructTypeKind as c_uint {
                mem::transmute(ty)
            } else {
                None
            }
        }
    }
}
deref!(StructType, Type);
get_context!(StructType, LLVMGetTypeContext);
to_str!(StructType, LLVMPrintTypeToString);

/// A function signature type
pub struct FunctionType;
native_ref!(&FunctionType = LLVMTypeRef);
deref!(FunctionType, Type);
impl CastFrom for FunctionType {
    type From = Type;
    fn cast(mut ty: &Type) -> Option<&FunctionType> {
        unsafe {
            use libc::c_uint;
            while let Some(elem) = ty.get_element() {
                ty = elem;
            }
            let kind = core::LLVMGetTypeKind(ty.into());
            if kind as c_uint == LLVMTypeKind::LLVMFunctionTypeKind as c_uint {
                mem::transmute(ty)
            } else {
                None
            }
        }
    }
}
impl FunctionType {
    /// Returns how many parameters this signature takes.
    pub fn num_params(&self) -> usize {
        unsafe { core::LLVMCountParamTypes(self.into()) as usize }
    }
    /// Returns a vector of this signature's parameters' types.
    pub fn get_params(&self) -> Vec<&Type> {
        unsafe {
            let count = core::LLVMCountParamTypes(self.into());
            let mut types:Vec<_> = (0..count).map(|_| mem::uninitialized()).collect();
            core::LLVMGetParamTypes(self.into(), types.as_mut_ptr() as *mut LLVMTypeRef);
            types
        }
    }
    /// Returns the type that this function returns.
    pub fn get_return(&self) -> &Type {
        unsafe { core::LLVMGetReturnType(self.into()).into() }
    }
}
get_context!(FunctionType, LLVMGetTypeContext);
to_str!(FunctionType, LLVMPrintTypeToString);
<commit_msg>Add is_pointer function to Type<commit_after>use ffi::prelude::LLVMTypeRef;
use ffi::{core, target, LLVMTypeKind};
use libc::{c_int, c_uint};
use compile::Compile;
use context::{Context, GetContext};
use target::TargetData;
use util::{self, CastFrom};
use std::{fmt, mem};
use std::iter::Iterator;
use std::ops::Deref;

/// Defines how a value should be laid out in memory.
pub struct Type;
native_ref!(&Type = LLVMTypeRef);
impl Type {
    #[inline(always)]
    /// Get the type given as an LLVM type descriptor in the context given.
    pub fn get<'a, T>(context:&'a Context) -> &'a Type where T:Compile<'a> {
        T::get_type(context)
    }
    /// Make a new function signature with the return type and arguments given.
    pub fn new_function<'a>(ret: &'a Type, args: &[&'a Type]) -> &'a FunctionType {
        unsafe { core::LLVMFunctionType(ret.into(), args.as_ptr() as *mut LLVMTypeRef, args.len() as c_uint, 0) }.into()
    }
    /// Make a new array with the length given.
    pub fn new_array<'a>(element: &'a Type, length: usize) -> &'a Type {
        unsafe { core::LLVMArrayType(element.into(), length as c_uint) }.into()
    }
    /// Make a new vector with the length given.
    pub fn new_vector<'a>(element: &'a Type, length: usize) -> &'a Type {
        unsafe { core::LLVMVectorType(element.into(), length as c_uint) }.into()
    }
    /// Make a new pointer with the given element type.
    pub fn new_pointer<'a>(elem: &'a Type) -> &'a Type {
        unsafe { core::LLVMPointerType(elem.into(), 0 as c_uint) }.into()
    }
    /// Make a new structure type with the given types.
    pub fn new_struct<'a>(elems: &[&'a Type], packed: bool) -> &'a Type {
        unsafe { core::LLVMStructType(elems.as_ptr() as *mut LLVMTypeRef, elems.len() as c_uint, packed as c_int) }.into()
    }
    /// Returns true if the size of the type is known at compile-time.
    ///
    /// This is equivalent to the type implementing `Sized` in Rust
    pub fn is_sized(&self) -> bool {
        unsafe { core::LLVMTypeIsSized(self.into()) != 0 }
    }
    /// Returns true if this type is a function.
    pub fn is_function(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMFunctionTypeKind as c_uint
    }
    /// Returns true if this type is a struct.
    pub fn is_struct(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMStructTypeKind as c_uint
    }
    /// Returns true if this type is void.
    pub fn is_void(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMVoidTypeKind as c_uint
    }
    /// Returns true if this type is a pointer.
    pub fn is_pointer(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMPointerTypeKind as c_uint
    }
    /// Returns the size of the type in bytes.
    pub fn get_size(&self, target: &TargetData) -> usize {
        unsafe { target::LLVMABISizeOfType(target.into(), self.into()) as usize }
    }
    /// Returns the element of this pointer type.
    pub fn get_element(&self) -> Option<&Type> {
        unsafe { mem::transmute(core::LLVMGetElementType(self.into())) }
    }
}
get_context!(Type, LLVMGetTypeContext);
to_str!(Type, LLVMPrintTypeToString);

/// A structure type, such as a tuple or struct
pub struct StructType;
native_ref!(&StructType = LLVMTypeRef);
impl StructType {
    /// Make a new struct with the given fields and packed representation.
    pub fn new<'a>(context: &'a Context, fields: &[&'a Type], packed: bool) -> &'a StructType {
        unsafe { core::LLVMStructTypeInContext(context.into(), fields.as_ptr() as *mut LLVMTypeRef, fields.len() as c_uint, packed as c_int) }.into()
    }
    /// Make a new named struct with the given fields and packed representation.
    pub fn new_named<'a>(context: &'a Context, name: &str, fields: &[&'a Type], packed: bool) -> &'a StructType {
        util::with_cstr(name, |name| unsafe {
            let ty = core::LLVMStructCreateNamed(context.into(), name);
            core::LLVMStructSetBody(ty, fields.as_ptr() as *mut LLVMTypeRef, fields.len() as c_uint, packed as c_int);
            ty.into()
        })
    }
    /// Returns the elements (AKA fields) that make up this struct.
    pub fn get_elements(&self) -> Vec<Type> {
        unsafe {
            let size = core::LLVMCountStructElementTypes(self.into());
            let mut els:Vec<_> = (0..size).map(|_| mem::uninitialized()).collect();
            core::LLVMGetStructElementTypes(self.into(), els.as_mut_ptr() as *mut LLVMTypeRef);
            els
        }
    }
}
impl CastFrom for StructType {
    type From = Type;
    fn cast(ty: &Type) -> Option<&StructType> {
        unsafe {
            let kind = core::LLVMGetTypeKind(ty.into());
            if kind as c_uint == LLVMTypeKind::LLVMStructTypeKind as c_uint {
                mem::transmute(ty)
            } else {
                None
            }
        }
    }
}
deref!(StructType, Type);
get_context!(StructType, LLVMGetTypeContext);
to_str!(StructType, LLVMPrintTypeToString);

/// A function signature type
pub struct FunctionType;
native_ref!(&FunctionType = LLVMTypeRef);
deref!(FunctionType, Type);
impl CastFrom for FunctionType {
    type From = Type;
    fn cast(mut ty: &Type) -> Option<&FunctionType> {
        unsafe {
            use libc::c_uint;
            while let Some(elem) = ty.get_element() {
                ty = elem;
            }
            let kind = core::LLVMGetTypeKind(ty.into());
            if kind as c_uint == LLVMTypeKind::LLVMFunctionTypeKind as c_uint {
                mem::transmute(ty)
            } else {
                None
            }
        }
    }
}
impl FunctionType {
    /// Returns how many parameters this signature takes.
    pub fn num_params(&self) -> usize {
        unsafe { core::LLVMCountParamTypes(self.into()) as usize }
    }
    /// Returns a vector of this signature's parameters' types.
    pub fn get_params(&self) -> Vec<&Type> {
        unsafe {
            let count = core::LLVMCountParamTypes(self.into());
            let mut types:Vec<_> = (0..count).map(|_| mem::uninitialized()).collect();
            core::LLVMGetParamTypes(self.into(), types.as_mut_ptr() as *mut LLVMTypeRef);
            types
        }
    }
    /// Returns the type that this function returns.
    pub fn get_return(&self) -> &Type {
        unsafe { core::LLVMGetReturnType(self.into()).into() }
    }
}
get_context!(FunctionType, LLVMGetTypeContext);
to_str!(FunctionType, LLVMPrintTypeToString);
<|endoftext|>


<commit_before>use af;
use af::{Array, MatProp};

use activations;
use params::{Input, Params};
use layer::Layer;

pub struct Dense {
  pub input_size: usize,
  pub output_size: usize,
}

impl Layer for Dense {

  fn forward(&self, params: &mut Params
             , inputs: &Input
             , recurrence: &Option<Input>) -> (Input, Option<Input>)
  {
    // keep previous_activation
    params.inputs = vec![inputs.clone()];

    // apply the activation to the previous layer [Optimization: Memory saving]
    let activated_input = activations::get_activation(&inputs.activation
                                                      , &inputs.data).unwrap();

    // sigma(Wx + b)
    let mul = af::matmul(&params.weights[0]
                         , &activated_input
                         , MatProp::NONE
                         , MatProp::NONE).unwrap();
    (Input { data: af::add(&mul, &params.biases[0], true).unwrap()
             , activation: params.activations[0].clone() }
     , None)
  }

  fn backward(&self, params: &mut Params, delta: &Array) -> Array {
    // d_l = (transpose(W) * d_{l}) .* dActivation(z-1) where z = activation w/out non-linearity
    params.deltas = vec![delta.clone()];
    let activation_prev = activations::get_activation(&params.inputs[0].activation, &params.inputs[0].data).unwrap();
    let d_activation_prev = activations::get_activation_derivative(&params.inputs[0].activation, &activation_prev).unwrap();
    let delta_prev = af::mul(&af::matmul(&params.weights[0], delta, af::MatProp::TRANS, af::MatProp::NONE).unwrap()
                             , &d_activation_prev, false).unwrap();
    delta_prev
  }

}
<commit_msg>update comment for delta<commit_after>use af;
use af::{Array, MatProp};

use activations;
use params::{Input, Params};
use layer::Layer;

pub struct Dense {
  pub input_size: usize,
  pub output_size: usize,
}

impl Layer for Dense {

  fn forward(&self, params: &mut Params
             , inputs: &Input
             , recurrence: &Option<Input>) -> (Input, Option<Input>)
  {
    // keep previous_activation
    params.inputs = vec![inputs.clone()];

    // apply the activation to the previous layer [Optimization: Memory saving]
    let activated_input = activations::get_activation(&inputs.activation
                                                      , &inputs.data).unwrap();

    // sigma(Wx + b)
    let mul = af::matmul(&params.weights[0]
                         , &activated_input
                         , MatProp::NONE
                         , MatProp::NONE).unwrap();
    (Input { data: af::add(&mul, &params.biases[0], true).unwrap()
             , activation: params.activations[0].clone() }
     , None)
  }

  fn backward(&self, params: &mut Params, delta: &Array) -> Array {
    // d_lm1 = (transpose(W) * d_{l}) .* dActivation(z-1) where z = activation w/out non-linearity
    params.deltas = vec![delta.clone()];
    let activation_prev = activations::get_activation(&params.inputs[0].activation, &params.inputs[0].data).unwrap();
    let d_activation_prev = activations::get_activation_derivative(&params.inputs[0].activation, &activation_prev).unwrap();
    let delta_prev = af::mul(&af::matmul(&params.weights[0], delta, af::MatProp::TRANS, af::MatProp::NONE).unwrap()
                             , &d_activation_prev, false).unwrap();
    delta_prev
  }

}
<|endoftext|>


<commit_before>// Copyright (C) 2019 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

pub(crate) use self::ffi::is_nand;
use super::{Error, Result};
use std::{
    fs,
    io::{BufRead, BufReader},
    path::PathBuf,
};

pub(crate) fn target_device_from_ubi_volume_name(volume: &str) -> Result<PathBuf> {
    let re = regex::Regex::new(r"^Volume ID:   (?P<volume>\d+) \(on ubi(\d+)\)$").unwrap();
    walkdir::WalkDir::new("/dev")
        .min_depth(1)
        .into_iter()
        .filter_entry(|p| {
            p.file_name()
                .to_str()
                .map(|n| n.starts_with("ubi") && !n.contains('_'))
                .unwrap_or(false)
        })
        .filter_map(std::result::Result::ok)
        .find_map(|entry| {
            let path = entry.path();
            let output =
                easy_process::run(&format!("ubinfo {} -N {}", path.display(), volume)).ok()?;

            let line = output.stdout.lines().next()?;
            let re_match = re.captures(line)?;

            Some(PathBuf::from(format!(
                "{}_{}",
                path.display(),
                &re_match.name("volume").unwrap().as_str()
            )))
        })
        .ok_or_else(|| Error::NoUbiVolume(volume.to_owned()))
}

pub(crate) fn target_device_from_mtd_name(name: &str) -> Result<PathBuf> {
    let re =
        regex::Regex::new(r#"^(?P<dev>mtd\d): ([[:xdigit:]]+) ([[:xdigit:]]+) "(?P<name>.*)"$"#)
            .unwrap();
    let proc = fs::File::open("/proc/mtd")?;

    BufReader::new(proc)
        .lines()
        .filter_map(std::result::Result::ok)
        .find_map(|line| {
            re.captures(&line).and_then(|re_match| {
                let re_dev = re_match.name("dev").unwrap().as_str();
                let re_name = re_match.name("name").unwrap().as_str();
                if re_name == name { Some(PathBuf::from(format!("/dev/{}", re_dev))) } else { None }
            })
        })
        .ok_or_else(|| Error::NoMtdDevice(name.to_owned()))
}

mod ffi {
    use crate::utils::Result;
    use nix::ioctl_read;
    use std::{mem::MaybeUninit, os::unix::io::AsRawFd, path::Path};

    // From https://github.com/torvalds/linux/blob/master/include/uapi/mtd/mtd-abi.h
    const MTD_NANDFLASH: u8 = 4;
    const MTD_MLCNANDFLASH: u8 = 8;
    const MEMGETINFO: u8 = b'M';
    const MEMGETINFO_MODE: u8 = 1;

    #[repr(C)]
    pub struct mtd_info_user {
        kind: u8,
        flags: u32,
        size: u32,
        erasesize: u32,
        writesize: u32,
        oobsize: u32,
        padding: u64,
    }

    ioctl_read!(mtd_get_info, MEMGETINFO, MEMGETINFO_MODE, mtd_info_user);

    pub fn is_nand(device: &Path) -> Result<bool> {
        let device = std::fs::File::open(device)?;
        let info = unsafe {
            let mut info = MaybeUninit::<mtd_info_user>::uninit();
            mtd_get_info(device.as_raw_fd(), info.as_mut_ptr())?;
            info.assume_init()
        };

        Ok(info.kind == MTD_NANDFLASH || info.kind == MTD_MLCNANDFLASH)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use lazy_static::lazy_static;
    use pretty_assertions::assert_eq;
    use std::sync::{Arc, Mutex};

    pub(crate) struct FakeUbi {
        #[allow(dead_code)]
        pub(crate) mtd_guard: FakeMtd,
    }

    impl FakeUbi {
        pub(crate) fn new(names: &[&str], kind: MtdKind) -> Result<FakeUbi> {
            let mtd_guard = FakeMtd::new(&["system"], kind)?;
            easy_process::run("modprobe ubi mtd=0")?;

            // Ubi created here so if anything fails the Drop will still be executed
            let ubi = FakeUbi { mtd_guard };

            for name in names {
                easy_process::run(&format!("ubimkvol /dev/ubi0 -N {} -s 1MiB", name))?;
            }

            Ok(ubi)
        }
    }

    impl Drop for FakeUbi {
        fn drop(&mut self) {
            if let Err(e) = easy_process::run("rmmod ubi") {
                eprintln!("Failed to cleanup FakeUbi, Error: {}", e);
            }
        }
    }

    pub(crate) struct FakeMtd {
        pub(crate) devices: Vec<PathBuf>,
        pub(crate) kind: MtdKind,
    }

    pub(crate) enum MtdKind {
        Nand,
        Nor,
    }

    impl FakeMtd {
        pub(crate) fn new(names: &[&str], kind: MtdKind) -> Result<FakeMtd> {
            match kind {
                MtdKind::Nand => easy_process::run("modprobe nandsim second_id_byte=0x36"),
                MtdKind::Nor => easy_process::run("modprobe mtdram total_size=20000"),
            }?;

            // FakeMtd created here so if any subsequent command fails the drop will still
            // be called to cleanup mtd devices
            let mut mtd = FakeMtd { devices: vec![], kind };
            for (i, name) in names.iter().enumerate() {
                easy_process::run(&format!("mtdpart add /dev/mtd0 {} {} {}", name, i * 100, 100))?;
                mtd.devices.push(PathBuf::from(format!("/dev/mtd{}", i + 1)));
            }

            Ok(mtd)
        }
    }

    impl Drop for FakeMtd {
        fn drop(&mut self) {
            let module = match self.kind {
                MtdKind::Nand => "nandsim",
                MtdKind::Nor => "mtdram",
            };

            // Sleep time for nandsim to sync and avoid errors
            std::thread::sleep(std::time::Duration::from_millis(500));

            if let Err(e) = easy_process::run(&format!("rmmod {}", module)) {
                eprintln!("Failed to cleanup FakeMtd, Error: {}", e);
            }
        }
    }

    // Used to serialize access to MTD devices
    lazy_static! {
        pub static ref SERIALIZE: Arc<Mutex<()>> = Arc::new(Mutex::default());
    }

    #[test]
    #[ignore]
    fn device_from_mtd_name() {
        let _lock = SERIALIZE.lock();
        let dev_names = vec!["system0", "system1"];

        let mtd = FakeMtd::new(&dev_names, MtdKind::Nor).unwrap();

        assert_eq!(
            dev_names
                .into_iter()
                .map(target_device_from_mtd_name)
                .map(Result::unwrap)
                .collect::<Vec<_>>(),
            mtd.devices,
        );
        assert!(target_device_from_mtd_name("some_inexistent_device").is_err());
    }

    #[test]
    #[ignore]
    fn test_is_nand() {
        let _lock = SERIALIZE.lock();

        {
            let _mtd = FakeMtd::new(&[], MtdKind::Nand).unwrap();
            assert_eq!(is_nand(&PathBuf::from("/dev/mtd0")).unwrap(), true);
        }
        {
            let _mtd = FakeMtd::new(&[], MtdKind::Nor).unwrap();
            assert_eq!(is_nand(&PathBuf::from("/dev/mtd0")).unwrap(), false);
        }
    }

    #[test]
    #[ignore]
    fn device_from_ubi_volume_name() {
        let _lock = SERIALIZE.lock();
        let volume_names = vec!["some_ui_volume", "another_ubi_volume"];

        let _ubi = FakeUbi::new(&volume_names, MtdKind::Nor).unwrap();
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[1]).unwrap(),
            PathBuf::from("/dev/ubi0_1")
        );
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[0]).unwrap(),
            PathBuf::from("/dev/ubi0_0")
        );
    }

    #[test]
    #[ignore]
    fn device_from_ubi_volume_name_multiple_volumes() {
        let _lock = SERIALIZE.lock();
        let volume_names = vec![
            "volume0", "volume1", "volume2", "volume3", "volume4", "volume5", "volume6", "volume7",
            "volume8", "volume9", "volume10", "volume11", "volume12", "volume13",
        ];

        let _ubi = FakeUbi::new(&volume_names, MtdKind::Nor).unwrap();
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[8]).unwrap(),
            PathBuf::from("/dev/ubi0_8")
        );
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[12]).unwrap(),
            PathBuf::from("/dev/ubi0_12")
        );
    }
}
<commit_msg>uh: apply clippy notes<commit_after>// Copyright (C) 2019 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

pub(crate) use self::ffi::is_nand;
use super::{Error, Result};
use std::{
    fs,
    io::{BufRead, BufReader},
    path::PathBuf,
};

pub(crate) fn target_device_from_ubi_volume_name(volume: &str) -> Result<PathBuf> {
    let re = regex::Regex::new(r"^Volume ID:   (?P<volume>\d+) \(on ubi(\d+)\)$").unwrap();
    walkdir::WalkDir::new("/dev")
        .min_depth(1)
        .into_iter()
        .filter_entry(|p| {
            p.file_name()
                .to_str()
                .map(|n| n.starts_with("ubi") && !n.contains('_'))
                .unwrap_or(false)
        })
        .filter_map(std::result::Result::ok)
        .find_map(|entry| {
            let path = entry.path();
            let output =
                easy_process::run(&format!("ubinfo {} -N {}", path.display(), volume)).ok()?;

            let line = output.stdout.lines().next()?;
            let re_match = re.captures(line)?;

            Some(PathBuf::from(format!(
                "{}_{}",
                path.display(),
                &re_match.name("volume").unwrap().as_str()
            )))
        })
        .ok_or_else(|| Error::NoUbiVolume(volume.to_owned()))
}

pub(crate) fn target_device_from_mtd_name(name: &str) -> Result<PathBuf> {
    let re =
        regex::Regex::new(r#"^(?P<dev>mtd\d): ([[:xdigit:]]+) ([[:xdigit:]]+) "(?P<name>.*)"$"#)
            .unwrap();
    let proc = fs::File::open("/proc/mtd")?;

    BufReader::new(proc)
        .lines()
        .filter_map(std::result::Result::ok)
        .find_map(|line| {
            re.captures(&line).and_then(|re_match| {
                let re_dev = re_match.name("dev").unwrap().as_str();
                let re_name = re_match.name("name").unwrap().as_str();
                if re_name == name { Some(PathBuf::from(format!("/dev/{}", re_dev))) } else { None }
            })
        })
        .ok_or_else(|| Error::NoMtdDevice(name.to_owned()))
}

mod ffi {
    use crate::utils::Result;
    use nix::ioctl_read;
    use std::{mem::MaybeUninit, os::unix::io::AsRawFd, path::Path};

    // From https://github.com/torvalds/linux/blob/master/include/uapi/mtd/mtd-abi.h
    const MTD_NANDFLASH: u8 = 4;
    const MTD_MLCNANDFLASH: u8 = 8;
    const MEMGETINFO: u8 = b'M';
    const MEMGETINFO_MODE: u8 = 1;

    #[repr(C)]
    pub struct mtd_info_user {
        kind: u8,
        flags: u32,
        size: u32,
        erasesize: u32,
        writesize: u32,
        oobsize: u32,
        padding: u64,
    }

    ioctl_read!(mtd_get_info, MEMGETINFO, MEMGETINFO_MODE, mtd_info_user);

    pub fn is_nand(device: &Path) -> Result<bool> {
        let device = std::fs::File::open(device)?;
        let info = unsafe {
            let mut info = MaybeUninit::<mtd_info_user>::uninit();
            mtd_get_info(device.as_raw_fd(), info.as_mut_ptr())?;
            info.assume_init()
        };

        Ok(info.kind == MTD_NANDFLASH || info.kind == MTD_MLCNANDFLASH)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use lazy_static::lazy_static;
    use pretty_assertions::assert_eq;
    use std::sync::{Arc, Mutex};

    pub(crate) struct FakeUbi {
        #[allow(dead_code)]
        pub(crate) mtd_guard: FakeMtd,
    }

    impl FakeUbi {
        pub(crate) fn new(names: &[&str], kind: MtdKind) -> Result<FakeUbi> {
            let mtd_guard = FakeMtd::new(&["system"], kind)?;
            easy_process::run("modprobe ubi mtd=0")?;

            // Ubi created here so if anything fails the Drop will still be executed
            let ubi = FakeUbi { mtd_guard };

            for name in names {
                easy_process::run(&format!("ubimkvol /dev/ubi0 -N {} -s 1MiB", name))?;
            }

            Ok(ubi)
        }
    }

    impl Drop for FakeUbi {
        fn drop(&mut self) {
            if let Err(e) = easy_process::run("rmmod ubi") {
                eprintln!("Failed to cleanup FakeUbi, Error: {}", e);
            }
        }
    }

    pub(crate) struct FakeMtd {
        pub(crate) devices: Vec<PathBuf>,
        pub(crate) kind: MtdKind,
    }

    pub(crate) enum MtdKind {
        Nand,
        Nor,
    }

    impl FakeMtd {
        pub(crate) fn new(names: &[&str], kind: MtdKind) -> Result<FakeMtd> {
            match kind {
                MtdKind::Nand => easy_process::run("modprobe nandsim second_id_byte=0x36"),
                MtdKind::Nor => easy_process::run("modprobe mtdram total_size=20000"),
            }?;

            // FakeMtd created here so if any subsequent command fails the drop will still
            // be called to cleanup mtd devices
            let mut mtd = FakeMtd { devices: vec![], kind };
            for (i, name) in names.iter().enumerate() {
                easy_process::run(&format!("mtdpart add /dev/mtd0 {} {} {}", name, i * 100, 100))?;
                mtd.devices.push(PathBuf::from(format!("/dev/mtd{}", i + 1)));
            }

            Ok(mtd)
        }
    }

    impl Drop for FakeMtd {
        fn drop(&mut self) {
            let module = match self.kind {
                MtdKind::Nand => "nandsim",
                MtdKind::Nor => "mtdram",
            };

            // Sleep time for nandsim to sync and avoid errors
            std::thread::sleep(std::time::Duration::from_millis(500));

            if let Err(e) = easy_process::run(&format!("rmmod {}", module)) {
                eprintln!("Failed to cleanup FakeMtd, Error: {}", e);
            }
        }
    }

    // Used to serialize access to MTD devices
    lazy_static! {
        pub static ref SERIALIZE: Arc<Mutex<()>> = Arc::new(Mutex::default());
    }

    #[test]
    #[ignore]
    fn device_from_mtd_name() {
        let _lock = SERIALIZE.lock();
        let dev_names = vec!["system0", "system1"];

        let mtd = FakeMtd::new(&dev_names, MtdKind::Nor).unwrap();

        assert_eq!(
            dev_names
                .into_iter()
                .map(target_device_from_mtd_name)
                .map(Result::unwrap)
                .collect::<Vec<_>>(),
            mtd.devices,
        );
        assert!(target_device_from_mtd_name("some_inexistent_device").is_err());
    }

    #[test]
    #[ignore]
    fn test_is_nand() {
        let _lock = SERIALIZE.lock();

        {
            let _mtd = FakeMtd::new(&[], MtdKind::Nand).unwrap();
            assert!(is_nand(&PathBuf::from("/dev/mtd0")).unwrap());
        }
        {
            let _mtd = FakeMtd::new(&[], MtdKind::Nor).unwrap();
            assert!(!is_nand(&PathBuf::from("/dev/mtd0")).unwrap());
        }
    }

    #[test]
    #[ignore]
    fn device_from_ubi_volume_name() {
        let _lock = SERIALIZE.lock();
        let volume_names = vec!["some_ui_volume", "another_ubi_volume"];

        let _ubi = FakeUbi::new(&volume_names, MtdKind::Nor).unwrap();
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[1]).unwrap(),
            PathBuf::from("/dev/ubi0_1")
        );
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[0]).unwrap(),
            PathBuf::from("/dev/ubi0_0")
        );
    }

    #[test]
    #[ignore]
    fn device_from_ubi_volume_name_multiple_volumes() {
        let _lock = SERIALIZE.lock();
        let volume_names = vec![
            "volume0", "volume1", "volume2", "volume3", "volume4", "volume5", "volume6", "volume7",
            "volume8", "volume9", "volume10", "volume11", "volume12", "volume13",
        ];

        let _ubi = FakeUbi::new(&volume_names, MtdKind::Nor).unwrap();
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[8]).unwrap(),
            PathBuf::from("/dev/ubi0_8")
        );
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[12]).unwrap(),
            PathBuf::from("/dev/ubi0_12")
        );
    }
}
<|endoftext|>