Implement add and sub for P

1 files changed, 62 insertions, 1 deletions

diff --git a/src/poly.rs b/src/poly.rs
index 2128d38..34444b3 100644
--- a/src/poly.rs
+++ b/src/poly.rs
@@ -1,6 +1,53 @@
+mod iter;
+
+use crate::number::Number;
+use iter::Iter;
+use std::cmp;
 use std::cmp::Ordering;
+use std::ops::{Add, Sub};
+use std::iter::{Zip, Take};
+
+#[derive(PartialEq, Debug)]
+pub struct P(Vec<Number>);
+
+impl P {
+    fn degree(&self) -> usize {
+        let mut i = self.0.len() - 1;
+        while self.0[i] == 0.0 && i > 0 {
+            i -= 1;
+        }
+        i
+    }
+
+    fn iter(&self) -> Iter {
+        Iter::new(self.0.clone())
+    }
+
+    fn zip(&self, other: &Self) -> Zip<Take<Iter>, Take<Iter>> {
+        let deg = cmp::max(self.degree(), other.degree()) + 1;
+        let a = self.iter().take(deg);
+        let b = other.iter().take(deg);
+        a.zip(b)
+    }
+}
+
+impl Add for &P {
+    type Output = P;
+
+    fn add(self, other: Self) -> P {
+        P(self.zip(other).map(|(x, y)| {x + y}).collect())
+    }
+}
+
+impl Sub for &P {
+    type Output = P;
+
+    fn sub(self, other: Self) -> P {
+        P(self.zip(other).map(|(x, y)| {x - y}).collect())
+    }
+}
 
-pub type Number = f32;
+//****************************REFACTORING************************************
 
 pub type Poly = Vec<Number>;
 
@@ -88,6 +135,20 @@ mod tests {
     use super::*;
 
     #[test]
+    fn add_test() {
+        let a = P(vec![1.0]);
+        let b = P(vec![0.0, 0.0, 0.0, 1.0]);
+        assert_eq!(&a + &b, P(vec![1.0, 0.0, 0.0, 1.0]));
+    }
+
+    #[test]
+    fn sub_test() {
+        let a = P(vec![1.0, 2.0, 3.0]);
+        let b = P(vec![1.0]);
+        assert_eq!(&a - &b, P(vec![0.0, 2.0, 3.0]));
+    }
+
+    #[test]
     fn poly_is_zero() {
         let p = vec![0.0; 5];
         assert_eq!(is_zero(&p), true);