Implement AudioFrameObu::ValidateAndReadPayload.

PiperOrigin-RevId: 632187932
AOMediaCodec · May 14, 2024 · 8c6c639 · 8c6c639
1 parent d3d998d
commit 8c6c639
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 11 deletions.
diff --git a/iamf/obu/audio_frame.cc b/iamf/obu/audio_frame.cc
@@ -51,6 +51,7 @@ absl::StatusOr<AudioFrameObu> AudioFrameObu::CreateFromBuffer(
     const ObuHeader& header, const int64_t obu_payload_serialized_size,
     ReadBitBuffer& rb) {
   AudioFrameObu audio_frame_obu(header);
+  audio_frame_obu.payload_serialized_size_ = obu_payload_serialized_size;
   RETURN_IF_NOT_OK(audio_frame_obu.ValidateAndReadPayload(rb));
   return audio_frame_obu;
 }
@@ -67,8 +68,17 @@ absl::Status AudioFrameObu::ValidateAndWritePayload(WriteBitBuffer& wb) const {
 }
 
 absl::Status AudioFrameObu::ValidateAndReadPayload(ReadBitBuffer& rb) {
-  return absl::UnimplementedError(
-      "AudioFrameOBU ValidateAndReadPayload not yet implemented.");
+  int8_t encoded_uleb128_size = 0;
+  if (header_.obu_type == kObuIaAudioFrame) {
+    // The ID is explicitly in the bitstream when `kObuIaAudioFrame`. Otherwise
+    // it is implied by `obu_type`.
+    RETURN_IF_NOT_OK(rb.ReadULeb128(audio_substream_id_, encoded_uleb128_size));
+  } else {
+    audio_substream_id_ = header_.obu_type;
+  }
+  RETURN_IF_NOT_OK(rb.ReadUint8Vector(
+      payload_serialized_size_ - encoded_uleb128_size, audio_frame_));
+  return absl::OkStatus();
 }
 
 void AudioFrameObu::PrintObu() const {

diff --git a/iamf/obu/audio_frame.h b/iamf/obu/audio_frame.h
@@ -83,11 +83,14 @@ class AudioFrameObu : public ObuBase {
 
  private:
   // This field is not serialized when in the range [0, 17].
-  const DecodedUleb128 audio_substream_id_;
+  DecodedUleb128 audio_substream_id_;
+
+  // Size of the obu in bytes. Used for the decoder only.
+  int64_t payload_serialized_size_ = 0;
 
   // Used only by the factory create function.
   explicit AudioFrameObu(const ObuHeader& header)
-      : ObuBase(header, kObuIaAudioElement),
+      : ObuBase(header, header.obu_type),
         audio_frame_({}),
         audio_substream_id_(DecodedUleb128()) {}
   /*\!brief Writes the OBU payload to the buffer.

diff --git a/iamf/obu/tests/audio_frame_test.cc b/iamf/obu/tests/audio_frame_test.cc
@@ -287,14 +287,46 @@ TEST_F(AudioFrameObuTest, ValidateAndWriteObuFailsWithIllegalRedundantCopy) {
 }
 
 // --- Begin CreateFromBuffer tests ---
-// TODO(b/329700769): Update test once ValidateAndReadPayload is implemented.
-TEST(CreateFromBuffer, IsNotSupported) {
-  std::vector<uint8_t> source;
+TEST(CreateFromBuffer, ValidAudioFrameWithExplicitId) {
+  std::vector<uint8_t> source = {// `explicit_audio_substream_id`
+                                 18,
+                                 // `audio_frame`.
+                                 8, 6, 24, 55, 11};
   ReadBitBuffer buffer(1024, &source);
-  ObuHeader header;
-  int64_t obu_payload_size = 16;
-  EXPECT_FALSE(
-      AudioFrameObu::CreateFromBuffer(header, obu_payload_size, buffer).ok());
+  ObuHeader header = {.obu_type = kObuIaAudioFrame};
+  const int64_t obu_payload_size = 6;
+  auto obu = AudioFrameObu::CreateFromBuffer(header, obu_payload_size, buffer);
+  EXPECT_TRUE(obu.ok());
+  EXPECT_EQ(obu->GetSubstreamId(), 18);
+  EXPECT_EQ(obu->audio_frame_, std::vector<uint8_t>({8, 6, 24, 55, 11}));
+}
+
+TEST(CreateFromBuffer, ValidAudioFrameWithImplicitId) {
+  std::vector<uint8_t> source = {// `audio_frame`.
+                                 8, 6, 24, 55, 11};
+  ReadBitBuffer buffer(1024, &source);
+  ObuHeader header = {.obu_type = kObuIaAudioFrameId0};
+  int64_t obu_payload_size = 5;
+  auto obu = AudioFrameObu::CreateFromBuffer(header, obu_payload_size, buffer);
+  EXPECT_TRUE(obu.ok());
+  // audio_substream_id is set implicitly to the value of `obu_type`
+  AudioFrameObu expected_obu =
+      AudioFrameObu(header, /*audio_substream_id=*/6,
+                    /*audio_frame=*/{8, 6, 24, 55, 11});
+  EXPECT_EQ(obu->GetSubstreamId(), 6);
+  EXPECT_EQ(obu->audio_frame_, std::vector<uint8_t>({8, 6, 24, 55, 11}));
+}
+
+TEST(CreateFromBuffer, FailsWithPayloadSizeTooLarge) {
+  std::vector<uint8_t> source = {// `explicit_audio_substream_id`
+                                 18,
+                                 // `audio_frame`.
+                                 8, 6, 24, 55, 11};
+  ReadBitBuffer buffer(1024, &source);
+  ObuHeader header = {.obu_type = kObuIaAudioFrame};
+  int64_t obu_payload_size = 7;
+  auto obu = AudioFrameObu::CreateFromBuffer(header, obu_payload_size, buffer);
+  EXPECT_FALSE(obu.ok());
 }
 
 }  // namespace