Remove length from SingleProcessEncodedImageDataInjector::ExtractionInfo, use SpatialLayerFrameSize instead

Bug: webrtc:11632
Change-Id: I8fea71e130df9894f26287ce94cd8bb05da3a69a
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/176331
Reviewed-by: Ilya Nikolaevskiy <ilnik@webrtc.org>
Reviewed-by: Artem Titov <titovartem@webrtc.org>
Commit-Queue: Andrey Logvin <landrey@webrtc.org>
Cr-Commit-Position: refs/heads/master@{#31398}

3 files changed, 112 insertions, 26 deletions

diff --git a/test/pc/e2e/analyzer/video/single_process_encoded_image_data_injector.cc b/test/pc/e2e/analyzer/video/single_process_encoded_image_data_injector.cc
index ec0d26b780..f85da3786c 100644
--- a/test/pc/e2e/analyzer/video/single_process_encoded_image_data_injector.cc
+++ b/test/pc/e2e/analyzer/video/single_process_encoded_image_data_injector.cc
@@ -40,7 +40,6 @@ EncodedImage SingleProcessEncodedImageDataInjector::InjectData(
   RTC_CHECK(source.size() >= kUsedBufferSize);
 
   ExtractionInfo info;
-  info.length = source.size();
   info.discard = discard;
   size_t insertion_pos = source.size() - kUsedBufferSize;
   memcpy(info.origin_data, &source.data()[insertion_pos], kUsedBufferSize);
@@ -69,15 +68,26 @@ EncodedImageExtractionResult SingleProcessEncodedImageDataInjector::ExtractData(
   uint8_t* buffer = out.data();
   size_t size = out.size();
 
-  // |pos| is pointing to end of current encoded image.
-  size_t pos = size - 1;
+  std::vector<size_t> frame_sizes;
+  std::vector<size_t> frame_sl_index;
+  size_t max_spatial_index = out.SpatialIndex().value_or(0);
+  for (size_t i = 0; i <= max_spatial_index; ++i) {
+    auto frame_size = source.SpatialLayerFrameSize(i);
+    if (frame_size.value_or(0)) {
+      frame_sl_index.push_back(i);
+      frame_sizes.push_back(frame_size.value());
+    }
+  }
+  if (frame_sizes.empty()) {
+    frame_sizes.push_back(size);
+  }
+
+  size_t prev_frames_size = 0;
   absl::optional<uint16_t> id = absl::nullopt;
   bool discard = true;
   std::vector<ExtractionInfo> extraction_infos;
-  // Go through whole buffer and find all related extraction infos in
-  // order from 1st encoded image to the last.
-  while (true) {
-    size_t insertion_pos = pos - kUsedBufferSize + 1;
+  for (size_t frame_size : frame_sizes) {
+    size_t insertion_pos = prev_frames_size + frame_size - kUsedBufferSize;
     // Extract frame id from first 2 bytes starting from insertion pos.
     uint16_t next_id = buffer[insertion_pos] + (buffer[insertion_pos + 1] << 8);
     // Extract frame sub id from second 3 byte starting from insertion pos.
@@ -99,41 +109,45 @@ EncodedImageExtractionResult SingleProcessEncodedImageDataInjector::ExtractData(
       info = info_it->second;
       ext_vector_it->second.infos.erase(info_it);
     }
-    extraction_infos.push_back(info);
     // We need to discard encoded image only if all concatenated encoded images
     // have to be discarded.
     discard = discard && info.discard;
-    if (pos < info.length) {
-      break;
-    }
-    pos -= info.length;
+
+    extraction_infos.push_back(info);
+    prev_frames_size += frame_size;
   }
   RTC_CHECK(id);
-  std::reverse(extraction_infos.begin(), extraction_infos.end());
+
   if (discard) {
     out.set_size(0);
+    for (size_t i = 0; i <= max_spatial_index; ++i) {
+      out.SetSpatialLayerFrameSize(i, 0);
+    }
     return EncodedImageExtractionResult{*id, out, true};
   }
 
   // Make a pass from begin to end to restore origin payload and erase discarded
   // encoded images.
-  pos = 0;
-  auto extraction_infos_it = extraction_infos.begin();
-  while (pos < size) {
-    RTC_DCHECK(extraction_infos_it != extraction_infos.end());
-    const ExtractionInfo& info = *extraction_infos_it;
+  size_t pos = 0;
+  for (size_t frame_index = 0; frame_index < frame_sizes.size();
+       ++frame_index) {
+    RTC_CHECK(pos < size);
+    const size_t frame_size = frame_sizes[frame_index];
+    const ExtractionInfo& info = extraction_infos[frame_index];
     if (info.discard) {
       // If this encoded image is marked to be discarded - erase it's payload
       // from the buffer.
-      memmove(&buffer[pos], &buffer[pos + info.length],
-              size - pos - info.length);
-      size -= info.length;
+      memmove(&buffer[pos], &buffer[pos + frame_size], size - pos - frame_size);
+      RTC_CHECK_LT(frame_index, frame_sl_index.size())
+          << "codec doesn't support discard option or the image, that was "
+             "supposed to be discarded, is lost";
+      out.SetSpatialLayerFrameSize(frame_sl_index[frame_index], 0);
+      size -= frame_size;
     } else {
-      memcpy(&buffer[pos + info.length - kUsedBufferSize], info.origin_data,
+      memcpy(&buffer[pos + frame_size - kUsedBufferSize], info.origin_data,
              kUsedBufferSize);
-      pos += info.length;
+      pos += frame_size;
     }
-    ++extraction_infos_it;
   }
   out.set_size(pos);
 
diff --git a/test/pc/e2e/analyzer/video/single_process_encoded_image_data_injector.h b/test/pc/e2e/analyzer/video/single_process_encoded_image_data_injector.h
index 3787cc51aa..c69cc9a2c0 100644
--- a/test/pc/e2e/analyzer/video/single_process_encoded_image_data_injector.h
+++ b/test/pc/e2e/analyzer/video/single_process_encoded_image_data_injector.h
@@ -59,8 +59,6 @@ class SingleProcessEncodedImageDataInjector : public EncodedImageDataInjector,
   struct ExtractionInfo {
     // Frame sub id to distinguish encoded images for different spatial layers.
     uint8_t sub_id;
-    // Length of the origin buffer encoded image.
-    size_t length;
     // Flag to show is this encoded images should be discarded by analyzing
     // decoder because of not required spatial layer/simulcast stream.
     bool discard;
diff --git a/test/pc/e2e/analyzer/video/single_process_encoded_image_data_injector_unittest.cc b/test/pc/e2e/analyzer/video/single_process_encoded_image_data_injector_unittest.cc
index 67cafa75a6..e25361e337 100644
--- a/test/pc/e2e/analyzer/video/single_process_encoded_image_data_injector_unittest.cc
+++ b/test/pc/e2e/analyzer/video/single_process_encoded_image_data_injector_unittest.cc
@@ -44,6 +44,7 @@ TEST(SingleProcessEncodedImageDataInjector, InjectExtractDiscardFalse) {
   EXPECT_FALSE(out.discard);
   EXPECT_EQ(out.image.size(), 10ul);
   EXPECT_EQ(out.image.capacity(), 10ul);
+  EXPECT_EQ(out.image.SpatialLayerFrameSize(0).value_or(0), 0ul);
   for (int i = 0; i < 10; ++i) {
     EXPECT_EQ(out.image.data()[i], i + 1);
   }
@@ -63,6 +64,60 @@ TEST(SingleProcessEncodedImageDataInjector, InjectExtractDiscardTrue) {
   EXPECT_TRUE(out.discard);
   EXPECT_EQ(out.image.size(), 0ul);
   EXPECT_EQ(out.image.capacity(), 10ul);
+  EXPECT_EQ(out.image.SpatialLayerFrameSize(0).value_or(0), 0ul);
+}
+
+TEST(SingleProcessEncodedImageDataInjector, InjectWithUnsetSpatialLayerSizes) {
+  SingleProcessEncodedImageDataInjector injector;
+
+  rtc::Buffer buffer = CreateBufferOfSizeNFilledWithValuesFromX(10, 1);
+
+  EncodedImage source(buffer.data(), 10, 10);
+  source.SetTimestamp(123456789);
+
+  EncodedImage intermediate = injector.InjectData(512, false, source, 1);
+  intermediate.SetSpatialIndex(2);
+
+  EncodedImageExtractionResult out = injector.ExtractData(intermediate, 2);
+  EXPECT_EQ(out.id, 512);
+  EXPECT_FALSE(out.discard);
+  EXPECT_EQ(out.image.size(), 10ul);
+  EXPECT_EQ(out.image.capacity(), 10ul);
+  for (int i = 0; i < 10; ++i) {
+    EXPECT_EQ(out.image.data()[i], i + 1);
+  }
+  EXPECT_EQ(out.image.SpatialIndex().value_or(0), 2);
+  for (int i = 0; i < 3; ++i) {
+    EXPECT_EQ(out.image.SpatialLayerFrameSize(i).value_or(0), 0ul);
+  }
+}
+
+TEST(SingleProcessEncodedImageDataInjector, InjectWithZeroSpatialLayerSizes) {
+  SingleProcessEncodedImageDataInjector injector;
+
+  rtc::Buffer buffer = CreateBufferOfSizeNFilledWithValuesFromX(10, 1);
+
+  EncodedImage source(buffer.data(), 10, 10);
+  source.SetTimestamp(123456789);
+
+  EncodedImage intermediate = injector.InjectData(512, false, source, 1);
+  intermediate.SetSpatialIndex(2);
+  intermediate.SetSpatialLayerFrameSize(0, 0);
+  intermediate.SetSpatialLayerFrameSize(1, 0);
+  intermediate.SetSpatialLayerFrameSize(2, 0);
+
+  EncodedImageExtractionResult out = injector.ExtractData(intermediate, 2);
+  EXPECT_EQ(out.id, 512);
+  EXPECT_FALSE(out.discard);
+  EXPECT_EQ(out.image.size(), 10ul);
+  EXPECT_EQ(out.image.capacity(), 10ul);
+  for (int i = 0; i < 10; ++i) {
+    EXPECT_EQ(out.image.data()[i], i + 1);
+  }
+  EXPECT_EQ(out.image.SpatialIndex().value_or(0), 2);
+  for (int i = 0; i < 3; ++i) {
+    EXPECT_EQ(out.image.SpatialLayerFrameSize(i).value_or(0), 0ul);
+  }
 }
 
 TEST(SingleProcessEncodedImageDataInjector, Inject3Extract3) {
@@ -95,6 +150,7 @@ TEST(SingleProcessEncodedImageDataInjector, Inject3Extract3) {
   EXPECT_FALSE(out1.discard);
   EXPECT_EQ(out1.image.size(), 10ul);
   EXPECT_EQ(out1.image.capacity(), 10ul);
+  EXPECT_EQ(out1.image.SpatialLayerFrameSize(0).value_or(0), 0ul);
   for (int i = 0; i < 10; ++i) {
     EXPECT_EQ(out1.image.data()[i], i + 1);
   }
@@ -102,10 +158,12 @@ TEST(SingleProcessEncodedImageDataInjector, Inject3Extract3) {
   EXPECT_TRUE(out2.discard);
   EXPECT_EQ(out2.image.size(), 0ul);
   EXPECT_EQ(out2.image.capacity(), 10ul);
+  EXPECT_EQ(out2.image.SpatialLayerFrameSize(0).value_or(0), 0ul);
   EXPECT_EQ(out3.id, 520);
   EXPECT_FALSE(out3.discard);
   EXPECT_EQ(out3.image.size(), 10ul);
   EXPECT_EQ(out3.image.capacity(), 10ul);
+  EXPECT_EQ(out3.image.SpatialLayerFrameSize(0).value_or(0), 0ul);
   for (int i = 0; i < 10; ++i) {
     EXPECT_EQ(out3.image.data()[i], i + 21);
   }
@@ -140,6 +198,10 @@ TEST(SingleProcessEncodedImageDataInjector, InjectExtractFromConcatenated) {
   concatenated_buffer.AppendData(intermediate3.data(), intermediate3.size());
   EncodedImage concatenated(concatenated_buffer.data(), concatenated_length,
                             concatenated_length);
+  concatenated.SetSpatialIndex(2);
+  concatenated.SetSpatialLayerFrameSize(0, intermediate1.size());
+  concatenated.SetSpatialLayerFrameSize(1, intermediate2.size());
+  concatenated.SetSpatialLayerFrameSize(2, intermediate3.size());
 
   // Extract frame id from concatenated image
   EncodedImageExtractionResult out = injector.ExtractData(concatenated, 2);
@@ -152,6 +214,10 @@ TEST(SingleProcessEncodedImageDataInjector, InjectExtractFromConcatenated) {
     EXPECT_EQ(out.image.data()[i], i + 1);
     EXPECT_EQ(out.image.data()[i + 10], i + 21);
   }
+  EXPECT_EQ(out.image.SpatialIndex().value_or(0), 2);
+  EXPECT_EQ(out.image.SpatialLayerFrameSize(0).value_or(0), 10ul);
+  EXPECT_EQ(out.image.SpatialLayerFrameSize(1).value_or(0), 0ul);
+  EXPECT_EQ(out.image.SpatialLayerFrameSize(2).value_or(0), 10ul);
 }
 
 TEST(SingleProcessEncodedImageDataInjector,
@@ -184,6 +250,10 @@ TEST(SingleProcessEncodedImageDataInjector,
   concatenated_buffer.AppendData(intermediate3.data(), intermediate3.size());
   EncodedImage concatenated(concatenated_buffer.data(), concatenated_length,
                             concatenated_length);
+  concatenated.SetSpatialIndex(2);
+  concatenated.SetSpatialLayerFrameSize(0, intermediate1.size());
+  concatenated.SetSpatialLayerFrameSize(1, intermediate2.size());
+  concatenated.SetSpatialLayerFrameSize(2, intermediate3.size());
 
   // Extract frame id from concatenated image
   EncodedImageExtractionResult out = injector.ExtractData(concatenated, 2);
@@ -192,6 +262,10 @@ TEST(SingleProcessEncodedImageDataInjector,
   EXPECT_TRUE(out.discard);
   EXPECT_EQ(out.image.size(), 0ul);
   EXPECT_EQ(out.image.capacity(), 3 * 10ul);
+  EXPECT_EQ(out.image.SpatialIndex().value_or(0), 2);
+  for (int i = 0; i < 3; ++i) {
+    EXPECT_EQ(out.image.SpatialLayerFrameSize(i).value_or(0), 0ul);
+  }
 }
 
 }  // namespace webrtc_pc_e2e