From a991c689c24b83a123e5c96d6a19d26680ebb088 Mon Sep 17 00:00:00 2001
From: Luca Barbato <lu_zero@gentoo.org>
Date: Tue, 25 Sep 2018 00:48:51 +0200
Subject: [PATCH] Provide iterators for PlaneSlices

And use it to compute SAD.
---
 benches/me.rs | 23 ++++++++++++++++++++++-
 src/me.rs     | 20 ++++++++++++++++++++
 src/plane.rs  | 30 ++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/benches/me.rs b/benches/me.rs
index 0beac7b1..84f9fcc2 100644
--- a/benches/me.rs
+++ b/benches/me.rs
@@ -51,6 +51,26 @@ fn bench_get_sad(b: &mut Bencher, bs: &BlockSize) {
   })
 }
 
+fn bench_get_sad_iter(b: &mut Bencher, bs: &BlockSize) {
+  let mut ra = ChaChaRng::from_seed([0; 32]);
+  let bsw = bs.width();
+  let bsh = bs.height();
+  let w = 640;
+  let h = 480;
+  let input_plane = new_plane(&mut ra, w, h);
+  let rec_plane = new_plane(&mut ra, w, h);
+  let po = PlaneOffset { x: 0, y: 0 };
+
+  let mut plane_org = input_plane.slice(&po);
+  let mut plane_ref = rec_plane.slice(&po);
+
+  b.iter(|| {
+      let _ = me::get_sad_iter(&mut plane_org, &mut plane_ref, bsw, bsh);
+      plane_org.y = 0;
+      plane_ref.y = 0;
+  })
+}
+
 pub fn get_sad(c: &mut Criterion) {
   use partition::BlockSize::*;
   let blocks = vec![
@@ -78,5 +98,6 @@ pub fn get_sad(c: &mut Criterion) {
     BLOCK_64X16,
   ];
 
-  c.bench_function_over_inputs("get_sad", bench_get_sad, blocks);
+  c.bench_function_over_inputs("get_sad", bench_get_sad, blocks.clone());
+  c.bench_function_over_inputs("get_sad_iter", bench_get_sad_iter, blocks);
 }
diff --git a/src/me.rs b/src/me.rs
index bd81b984..60e0a45f 100644
--- a/src/me.rs
+++ b/src/me.rs
@@ -39,6 +39,26 @@ pub fn get_sad(
   sum
 }
 
+#[inline(always)]
+pub fn get_sad_iter(
+  plane_org: &mut PlaneSlice, plane_ref: &mut PlaneSlice, blk_h: usize,
+  blk_w: usize
+) -> u32 {
+  let mut sum = 0 as u32;
+
+  let org_iter = plane_org.iter_width(blk_w);
+  let ref_iter = plane_ref.iter_width(blk_w);
+
+  for (slice_org, slice_ref) in org_iter.take(blk_h).zip(ref_iter) {
+      sum += slice_org
+        .iter()
+        .zip(slice_ref)
+        .map(|(&a, &b)| (a as i32 - b as i32).abs() as u32)
+        .sum::<u32>();
+  }
+
+  sum
+}
 pub fn motion_estimation(
   fi: &FrameInvariants, fs: &FrameState, bsize: BlockSize,
   bo: &BlockOffset, ref_frame: usize, pmv: &MotionVector
diff --git a/src/plane.rs b/src/plane.rs
index 42023b82..9c38ac7e 100644
--- a/src/plane.rs
+++ b/src/plane.rs
@@ -220,12 +220,38 @@ impl Plane {
   }
 }
 
+#[derive(Clone, Copy)]
 pub struct PlaneSlice<'a> {
   pub plane: &'a Plane,
   pub x: isize,
   pub y: isize
 }
 
+pub struct IterWidth<'a> {
+    ps: PlaneSlice<'a>,
+    width: usize,
+}
+
+// TODO: Implement more methods
+impl<'a> Iterator for IterWidth<'a> {
+    type Item = &'a [u16];
+
+    #[inline]
+    fn next(&mut self) -> Option<&'a [u16]> {
+        let x = self.ps.plane.cfg.xorigin + self.ps.x;
+        let y = self.ps.plane.cfg.yorigin + self.ps.y;
+        let stride = self.ps.plane.cfg.stride;
+        let base = y as usize * stride + x as usize;
+
+        if self.ps.plane.data.len() < base + self.width {
+            None
+        } else {
+            self.ps.y += 1;
+            Some(&self.ps.plane.data[base..base + self.width])
+        }
+    }
+}
+
 impl<'a> PlaneSlice<'a> {
   pub fn as_slice(&'a self) -> &'a [u16] {
     let stride = self.plane.cfg.stride;
@@ -252,6 +278,10 @@ impl<'a> PlaneSlice<'a> {
     &self.plane.data[base..base + width]
   }
 
+  pub fn iter_width(&self, width: usize) -> IterWidth<'a> {
+    IterWidth { ps: *self, width }
+  }
+
   pub fn subslice(&'a self, xo: usize, yo: usize) -> PlaneSlice<'a> {
     PlaneSlice {
       plane: self.plane,
-- 
GitLab