Move all the computation in receive_packet

Author: lu-zero

src/api/context.rs

@@ -122,12 +122,8 @@ impl<T: Pixel> Context<T> {
     }
 
     let inner = &mut self.inner;
-    let run = move || inner.send_frame(frame, params);
 
-    match &self.pool {
-      Some(pool) => pool.install(run),
-      None => run(),
-    }
+    inner.send_frame(frame, params)
   }
 
   /// Returns the first-pass data of a two-pass encode for the frame that was

src/api/internal.rs

@@ -258,6 +258,7 @@ pub(crate) struct ContextInner<T: Pixel> {
   next_lookahead_output_frameno: u64,
   /// Optional opaque to be sent back to the user
   opaque_q: BTreeMap<u64, Opaque>,
+  is_flushing: bool,
 }
 
 impl<T: Pixel> ContextInner<T> {
@@ -325,6 +326,7 @@ impl<T: Pixel> ContextInner<T> {
       next_lookahead_frame: 1,
       next_lookahead_output_frameno: 0,
       opaque_q: BTreeMap::new(),
+      is_flushing: false,
     }
   }
 
@@ -333,8 +335,9 @@ impl<T: Pixel> ContextInner<T> {
     &mut self, frame: Option<Arc<Frame<T>>>, params: Option<FrameParameters>,
   ) -> Result<(), EncoderStatus> {
     let input_frameno = self.frame_count;
-    let is_flushing = frame.is_none();
-    if !is_flushing {
+
+    self.is_flushing = frame.is_none();
+    if !self.is_flushing {
       self.frame_count += 1;
     }
     self.frame_q.insert(input_frameno, frame);
@@ -348,33 +351,6 @@ impl<T: Pixel> ContextInner<T> {
       }
     }
 
-    if !self.needs_more_frame_q_lookahead(self.next_lookahead_frame) {
-      let lookahead_frames = self
-        .frame_q
-        .range(self.next_lookahead_frame - 1..)
-        .filter_map(|(&_input_frameno, frame)| frame.clone())
-        .collect::<Vec<_>>();
-
-      if is_flushing {
-        // This is the last time send_frame is called, process all the
-        // remaining frames.
-        for cur_lookahead_frames in
-          std::iter::successors(Some(&lookahead_frames[..]), |s| s.get(1..))
-        {
-          if cur_lookahead_frames.len() < 2 {
-            // All frames have been processed
-            break;
-          }
-
-          self.compute_keyframe_placement(cur_lookahead_frames);
-        }
-      } else {
-        self.compute_keyframe_placement(&lookahead_frames);
-      }
-    }
-
-    self.compute_frame_invariants();
-
     Ok(())
   }
 
@@ -1288,12 +1264,44 @@ impl<T: Pixel> ContextInner<T> {
     }
   }
 
+  // lookahead computations
+  pub(crate) fn compute_fi(&mut self) {
+    if !self.needs_more_frame_q_lookahead(self.next_lookahead_frame) {
+      let lookahead_frames = self
+        .frame_q
+        .range(self.next_lookahead_frame - 1..)
+        .filter_map(|(&_input_frameno, frame)| frame.clone())
+        .collect::<Vec<_>>();
+
+      if self.is_flushing {
+        // This is the last time send_frame is called, process all the
+        // remaining frames.
+        for cur_lookahead_frames in
+          std::iter::successors(Some(&lookahead_frames[..]), |s| s.get(1..))
+        {
+          if cur_lookahead_frames.len() < 2 {
+            // All frames have been processed
+            break;
+          }
+
+          self.compute_keyframe_placement(cur_lookahead_frames);
+        }
+      } else {
+        self.compute_keyframe_placement(&lookahead_frames);
+      }
+    }
+
+    self.compute_frame_invariants();
+  }
+
   #[hawktracer(receive_packet)]
   pub fn receive_packet(&mut self) -> Result<Packet<T>, EncoderStatus> {
     if self.done_processing() {
       return Err(EncoderStatus::LimitReached);
     }
 
+    self.compute_fi();
+
     if self.needs_more_fi_lookahead() {
       return Err(EncoderStatus::NeedMoreData);
     }

src/api/test.rs

@@ -274,8 +274,9 @@ fn send_test_frame<T: Pixel>(ctx: &mut Context<T>, content_value: T) {
 }
 
 fn get_frame_invariants<T: Pixel>(
-  ctx: Context<T>,
+  mut ctx: Context<T>,
 ) -> impl Iterator<Item = FrameInvariants<T>> {
+  ctx.inner.compute_fi();
   ctx.inner.frame_data.into_iter().map(|(_, v)| v.fi)
 }
 
@@ -1777,6 +1778,7 @@ fn lookahead_size_properly_bounded(
   for i in 0..LIMIT {
     let input = ctx.new_frame();
     let _ = ctx.send_frame(input);
+    ctx.inner.compute_fi();
     pre_receive_frame_q_lens[i] = ctx.inner.frame_q.len();
     pre_receive_fi_lens[i] = ctx.inner.frame_data.len();
     while ctx.receive_packet().is_ok() {
@@ -2047,6 +2049,7 @@ fn min_quantizer_bounds_correctly() {
   ctx.flush();
 
   for i in 0..limit {
+    ctx.inner.compute_fi();
     ctx.inner.encode_packet(i).unwrap();
     let frame_data = ctx.inner.frame_data.get(&i).unwrap();
     if i == 0 {
@@ -2078,6 +2081,7 @@ fn min_quantizer_bounds_correctly() {
   ctx.flush();
 
   for i in 0..limit {
+    ctx.inner.compute_fi();
     ctx.inner.encode_packet(i).unwrap();
     let frame_data = ctx.inner.frame_data.get(&i).unwrap();
     if i == 0 {
@@ -2112,6 +2116,7 @@ fn max_quantizer_bounds_correctly() {
   ctx.flush();
 
   for i in 0..limit {
+    ctx.inner.compute_fi();
     ctx.inner.encode_packet(i).unwrap();
     let frame_data = ctx.inner.frame_data.get(&i).unwrap();
     if i == 0 {
@@ -2143,6 +2148,7 @@ fn max_quantizer_bounds_correctly() {
   ctx.flush();
 
   for i in 0..limit {
+    ctx.inner.compute_fi();
     ctx.inner.encode_packet(i).unwrap();
     let frame_data = ctx.inner.frame_data.get(&i).unwrap();
     if i == 0 {