1// Copyright 2014 The Chromium Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#include <algorithm>
6#include <vector>
7
8#include "base/bind.h"
9#include "base/command_line.h"
10#include "base/files/memory_mapped_file.h"
11#include "base/logging.h"
12#include "base/path_service.h"
13#include "base/time/time.h"
14#include "media/base/stream_parser_buffer.h"
15#include "media/base/test_data_util.h"
16#include "media/filters/h264_parser.h"
17#include "media/formats/mp2t/es_parser_h264.h"
18#include "testing/gtest/include/gtest/gtest.h"
19
20namespace media {
21class VideoDecoderConfig;
22
23namespace mp2t {
24
25namespace {
26
27struct Packet {
28  // Offset in the stream.
29  size_t offset;
30
31  // Size of the packet.
32  size_t size;
33
34  // Timestamp of the packet.
35  base::TimeDelta pts;
36};
37
38// Compute the size of each packet assuming packets are given in stream order
39// and the last packet covers the end of the stream.
40void ComputePacketSize(std::vector<Packet>& packets, size_t stream_size) {
41  for (size_t k = 0; k < packets.size() - 1; k++) {
42    DCHECK_GE(packets[k + 1].offset, packets[k].offset);
43    packets[k].size = packets[k + 1].offset - packets[k].offset;
44  }
45  packets[packets.size() - 1].size =
46      stream_size - packets[packets.size() - 1].offset;
47}
48
49// Get the offset of the start of each access unit.
50// This function assumes there is only one slice per access unit.
51// This is a very simplified access unit segmenter that is good
52// enough for unit tests.
53std::vector<Packet> GetAccessUnits(const uint8* stream, size_t stream_size) {
54  std::vector<Packet> access_units;
55  bool start_access_unit = true;
56
57  // In a first pass, retrieve the offsets of all access units.
58  size_t offset = 0;
59  while (true) {
60    // Find the next start code.
61    off_t relative_offset = 0;
62    off_t start_code_size = 0;
63    bool success = H264Parser::FindStartCode(
64        &stream[offset], stream_size - offset,
65        &relative_offset, &start_code_size);
66    if (!success)
67      break;
68    offset += relative_offset;
69
70    if (start_access_unit) {
71      Packet cur_access_unit;
72      cur_access_unit.offset = offset;
73      access_units.push_back(cur_access_unit);
74      start_access_unit = false;
75    }
76
77    // Get the NALU type.
78    offset += start_code_size;
79    if (offset >= stream_size)
80      break;
81    int nal_unit_type = stream[offset] & 0x1f;
82
83    // We assume there is only one slice per access unit.
84    if (nal_unit_type == H264NALU::kIDRSlice ||
85        nal_unit_type == H264NALU::kNonIDRSlice) {
86      start_access_unit = true;
87    }
88  }
89
90  ComputePacketSize(access_units, stream_size);
91  return access_units;
92}
93
94// Append an AUD NALU at the beginning of each access unit
95// needed for streams which do not already have AUD NALUs.
96void AppendAUD(
97    const uint8* stream, size_t stream_size,
98    const std::vector<Packet>& access_units,
99    std::vector<uint8>& stream_with_aud,
100    std::vector<Packet>& access_units_with_aud) {
101  uint8 aud[] = { 0x00, 0x00, 0x01, 0x09 };
102  stream_with_aud.resize(stream_size + access_units.size() * sizeof(aud));
103  access_units_with_aud.resize(access_units.size());
104
105  size_t offset = 0;
106  for (size_t k = 0; k < access_units.size(); k++) {
107    access_units_with_aud[k].offset = offset;
108    access_units_with_aud[k].size = access_units[k].size + sizeof(aud);
109
110    memcpy(&stream_with_aud[offset], aud, sizeof(aud));
111    offset += sizeof(aud);
112
113    memcpy(&stream_with_aud[offset],
114           &stream[access_units[k].offset], access_units[k].size);
115    offset += access_units[k].size;
116  }
117}
118
119}  // namespace
120
121class EsParserH264Test : public testing::Test {
122 public:
123  EsParserH264Test() : buffer_count_(0) {
124  }
125  virtual ~EsParserH264Test() {}
126
127 protected:
128  void LoadStream(const char* filename);
129  void GetPesTimestamps(std::vector<Packet>& pes_packets);
130  void ProcessPesPackets(const std::vector<Packet>& pes_packets,
131                         bool force_timing);
132
133  // Stream with AUD NALUs.
134  std::vector<uint8> stream_;
135
136  // Access units of the stream with AUD NALUs.
137  std::vector<Packet> access_units_;
138
139  // Number of buffers generated while parsing the H264 stream.
140  size_t buffer_count_;
141
142 private:
143  void EmitBuffer(scoped_refptr<StreamParserBuffer> buffer);
144
145  void NewVideoConfig(const VideoDecoderConfig& config) {
146  }
147
148  DISALLOW_COPY_AND_ASSIGN(EsParserH264Test);
149};
150
151void EsParserH264Test::LoadStream(const char* filename) {
152  base::FilePath file_path = GetTestDataFilePath(filename);
153
154  base::MemoryMappedFile stream_without_aud;
155  ASSERT_TRUE(stream_without_aud.Initialize(file_path))
156      << "Couldn't open stream file: " << file_path.MaybeAsASCII();
157
158  // The input file does not have AUDs.
159  std::vector<Packet> access_units_without_aud = GetAccessUnits(
160      stream_without_aud.data(), stream_without_aud.length());
161  ASSERT_GT(access_units_without_aud.size(), 0u);
162  AppendAUD(stream_without_aud.data(), stream_without_aud.length(),
163            access_units_without_aud,
164            stream_, access_units_);
165
166  // Generate some timestamps based on a 25fps stream.
167  for (size_t k = 0; k < access_units_.size(); k++)
168    access_units_[k].pts = base::TimeDelta::FromMilliseconds(k * 40u);
169}
170
171void EsParserH264Test::GetPesTimestamps(std::vector<Packet>& pes_packets) {
172  // Default: set to a negative timestamp to be able to differentiate from
173  // real timestamps.
174  // Note: we don't use kNoTimestamp() here since this one has already
175  // a special meaning in EsParserH264. The negative timestamps should be
176  // ultimately discarded by the H264 parser since not relevant.
177  for (size_t k = 0; k < pes_packets.size(); k++) {
178    pes_packets[k].pts = base::TimeDelta::FromMilliseconds(-1);
179  }
180
181  // Set a valid timestamp for PES packets which include the start
182  // of an H264 access unit.
183  size_t pes_idx = 0;
184  for (size_t k = 0; k < access_units_.size(); k++) {
185    for (; pes_idx < pes_packets.size(); pes_idx++) {
186      size_t pes_start = pes_packets[pes_idx].offset;
187      size_t pes_end = pes_packets[pes_idx].offset + pes_packets[pes_idx].size;
188      if (pes_start <= access_units_[k].offset &&
189          pes_end > access_units_[k].offset) {
190        pes_packets[pes_idx].pts = access_units_[k].pts;
191        break;
192      }
193    }
194  }
195}
196
197void EsParserH264Test::ProcessPesPackets(
198    const std::vector<Packet>& pes_packets,
199    bool force_timing) {
200  EsParserH264 es_parser(
201      base::Bind(&EsParserH264Test::NewVideoConfig, base::Unretained(this)),
202      base::Bind(&EsParserH264Test::EmitBuffer, base::Unretained(this)));
203
204  for (size_t k = 0; k < pes_packets.size(); k++) {
205    size_t cur_pes_offset = pes_packets[k].offset;
206    size_t cur_pes_size = pes_packets[k].size;
207
208    base::TimeDelta pts = kNoTimestamp();
209    base::TimeDelta dts = kNoTimestamp();
210    if (pes_packets[k].pts >= base::TimeDelta() || force_timing)
211      pts = pes_packets[k].pts;
212
213    ASSERT_TRUE(
214        es_parser.Parse(&stream_[cur_pes_offset], cur_pes_size, pts, dts));
215  }
216  es_parser.Flush();
217}
218
219void EsParserH264Test::EmitBuffer(scoped_refptr<StreamParserBuffer> buffer) {
220  ASSERT_LT(buffer_count_, access_units_.size());
221  EXPECT_EQ(buffer->timestamp(), access_units_[buffer_count_].pts);
222  buffer_count_++;
223}
224
225TEST_F(EsParserH264Test, OneAccessUnitPerPes) {
226  LoadStream("bear.h264");
227
228  // One to one equivalence between PES packets and access units.
229  std::vector<Packet> pes_packets(access_units_);
230  GetPesTimestamps(pes_packets);
231
232  // Process each PES packet.
233  ProcessPesPackets(pes_packets, false);
234  EXPECT_EQ(buffer_count_, access_units_.size());
235}
236
237TEST_F(EsParserH264Test, NonAlignedPesPacket) {
238  LoadStream("bear.h264");
239
240  // Generate the PES packets.
241  std::vector<Packet> pes_packets;
242  Packet cur_pes_packet;
243  cur_pes_packet.offset = 0;
244  for (size_t k = 0; k < access_units_.size(); k++) {
245    pes_packets.push_back(cur_pes_packet);
246
247    // The current PES packet includes the remaining bytes of the previous
248    // access unit and some bytes of the current access unit
249    // (487 bytes in this unit test but no more than the current access unit
250    // size).
251    cur_pes_packet.offset = access_units_[k].offset +
252        std::min<size_t>(487u, access_units_[k].size);
253  }
254  ComputePacketSize(pes_packets, stream_.size());
255  GetPesTimestamps(pes_packets);
256
257  // Process each PES packet.
258  ProcessPesPackets(pes_packets, false);
259  EXPECT_EQ(buffer_count_, access_units_.size());
260}
261
262TEST_F(EsParserH264Test, SeveralPesPerAccessUnit) {
263  LoadStream("bear.h264");
264
265  // Get the minimum size of an access unit.
266  size_t min_access_unit_size = stream_.size();
267  for (size_t k = 0; k < access_units_.size(); k++) {
268    if (min_access_unit_size >= access_units_[k].size)
269      min_access_unit_size = access_units_[k].size;
270  }
271
272  // Use a small PES packet size or the minimum access unit size
273  // if it is even smaller.
274  size_t pes_size = 512;
275  if (min_access_unit_size < pes_size)
276    pes_size = min_access_unit_size;
277
278  std::vector<Packet> pes_packets;
279  Packet cur_pes_packet;
280  cur_pes_packet.offset = 0;
281  while (cur_pes_packet.offset < stream_.size()) {
282    pes_packets.push_back(cur_pes_packet);
283    cur_pes_packet.offset += pes_size;
284  }
285  ComputePacketSize(pes_packets, stream_.size());
286  GetPesTimestamps(pes_packets);
287
288  // Process each PES packet.
289  ProcessPesPackets(pes_packets, false);
290  EXPECT_EQ(buffer_count_, access_units_.size());
291
292  // Process PES packets forcing timings for each PES packet.
293  buffer_count_ = 0;
294  ProcessPesPackets(pes_packets, true);
295  EXPECT_EQ(buffer_count_, access_units_.size());
296}
297
298}  // namespace mp2t
299}  // namespace media
300
301