Another try 2

rust-av · Nov 28, 2019 · 664baf3 · 664baf3
1 parent a0e3e46
commit 664baf3
Showing 1 changed file with 18 additions and 29 deletions.
diff --git a/av_metrics/src/video/mod.rs b/av_metrics/src/video/mod.rs
@@ -200,35 +200,15 @@ pub enum ParallelMethod {
 
 macro_rules! process_video {
     ($self:ident, $decoder1:ident, $decoder2:ident, $frame_limit:ident, $type:ident, $ret:ident) => {
-        $ret = (0..$frame_limit.unwrap_or(std::usize::MAX))
-            .map(|_| {
-                (
-                    $decoder1.read_video_frame::<$type>().ok(),
-                    $decoder2.read_video_frame::<$type>().ok(),
-                )
-            })
-            .take_while(|(frame1, frame2)| frame1.is_some() && frame2.is_some())
-            .par_bridge()
-            .try_fold(
-                || vec![],
-                |mut acc, (f1, f2)| {
-                    let frame1 = f1.unwrap();
-                    let frame2 = f2.unwrap();
-                    $self
-                        .process_frame(&frame1, &frame2)
-                        .map(|(metrics, cweight)| {
-                            acc.push((metrics, cweight));
-                            acc
-                        })
-                },
-            )
-            .try_reduce(
-                || vec![],
-                |mut a, b| {
-                    a.extend_from_slice(&b);
-                    Ok(a)
-                },
-            )?;
+        let frames_pool: Vec<(FrameInfo<$type>, FrameInfo<$type>)>;
+        frames_pool = iter::from_fn(|| $decoder1.read_video_frame::<$type>().ok())
+            .zip(iter::from_fn(|| $decoder2.read_video_frame::<$type>().ok()))
+            .take($frame_limit.unwrap_or(std::usize::MAX))
+            .collect();
+        $ret = frames_pool
+            .par_iter()
+            .map(|(frame1, frame2)| $self.process_frame(&frame1, &frame2).unwrap())
+            .collect();
     };
 }
 
@@ -257,6 +237,15 @@ trait VideoMetric {
             }));
         }
 
+        // This HACK is horrible...
+        match rayon::ThreadPoolBuilder::new()
+            .num_threads(8)
+            .build_global()
+        {
+            Ok(_) => (),
+            Err(_) => (),
+        }
+
         let metrics_values: Vec<(Self::FrameResult, Option<f64>)>;
         if decoder1.get_bit_depth() > 8 {
             process_video!(self, decoder1, decoder2, frame_limit, u16, metrics_values);