diff --git a/src/lib.rs b/src/lib.rs index f929151b..66fe2c39 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -32,7 +32,7 @@ extern crate serde; use std::fmt::{Debug, Display, Formatter, Result}; -use std::ops::{Add, Mul}; +use std::ops::{Add, AddAssign, Mul, MulAssign}; /// byte size for 1 byte pub const B: u64 = 1; @@ -217,14 +217,6 @@ impl Debug for ByteSize { macro_rules! commutative_op { ($t:ty) => { - impl Add<$t> for ByteSize { - type Output = ByteSize; - #[inline(always)] - fn add(self, rhs: $t) -> ByteSize { - ByteSize(self.0 + (rhs as u64)) - } - } - impl Add for $t { type Output = ByteSize; #[inline(always)] @@ -233,14 +225,6 @@ macro_rules! commutative_op { } } - impl Mul<$t> for ByteSize { - type Output = ByteSize; - #[inline(always)] - fn mul(self, rhs: $t) -> ByteSize { - ByteSize(self.0 * (rhs as u64)) - } - } - impl Mul for $t { type Output = ByteSize; #[inline(always)] @@ -265,13 +249,54 @@ impl Add for ByteSize { } } +impl AddAssign for ByteSize { + #[inline(always)] + fn add_assign(&mut self, rhs: ByteSize) { + self.0 += rhs.0 + } +} + +impl Add for ByteSize + where T: Into { + type Output = ByteSize; + #[inline(always)] + fn add(self, rhs: T) -> ByteSize { + ByteSize(self.0 + (rhs.into() as u64)) + } +} + +impl AddAssign for ByteSize + where T: Into { + #[inline(always)] + fn add_assign(&mut self, rhs: T) { + self.0 += rhs.into() as u64; + } +} + +impl Mul for ByteSize + where T: Into { + type Output = ByteSize; + #[inline(always)] + fn mul(self, rhs: T) -> ByteSize { + ByteSize(self.0 * (rhs.into() as u64)) + } +} + +impl MulAssign for ByteSize + where T: Into { + #[inline(always)] + fn mul_assign(&mut self, rhs: T) { + self.0 *= rhs.into() as u64; + } +} + #[cfg(test)] mod tests { use super::*; #[test] fn test_arithmetic_op() { - let x = ByteSize::mb(1); + let mut x = ByteSize::mb(1); let y = ByteSize::kb(100); assert_eq!((x + y).as_u64(), 1_100_000u64); @@ -279,11 +304,16 @@ mod tests { assert_eq!((x + (100 * 1000) as u64).as_u64(), 1_100_000); assert_eq!((x * 2u64).as_u64(), 2_000_000); + + x += y; + assert_eq!(x.as_u64(), 1_100_000); + x *= 2u64; + assert_eq!(x.as_u64(), 2_200_000); } #[test] fn test_arithmetic_primitives() { - let x = ByteSize::mb(1); + let mut x = ByteSize::mb(1); assert_eq!((x + MB as u64).as_u64(), 2_000_000); @@ -292,6 +322,12 @@ mod tests { assert_eq!((x + KB as u16).as_u64(), 1_001_000); assert_eq!((x + B as u8).as_u64(), 1_000_001); + + x += MB as u64; + x += MB as u32; + x += 10 as u16; + x += 1 as u8; + assert_eq!(x.as_u64(), 3_000_011); } #[test]