From 8e08d982aba3edf87c74f6462a99cadc793fa487 Mon Sep 17 00:00:00 2001
From: olly <olly@google.com>
Date: Wed, 2 Aug 2017 03:07:56 -0700
Subject: [PATCH] Further fix H262 segmentation

Issue: #2891

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=163951910
---
 .../androidTest/assets/ts/sample.ts.0.dump    |  2 +-
 .../exoplayer2/extractor/ts/H262Reader.java   | 46 ++++++++-----------
 2 files changed, 21 insertions(+), 27 deletions(-)

diff --git a/library/core/src/androidTest/assets/ts/sample.ts.0.dump b/library/core/src/androidTest/assets/ts/sample.ts.0.dump
index 83f1337816a..91e48b17220 100644
--- a/library/core/src/androidTest/assets/ts/sample.ts.0.dump
+++ b/library/core/src/androidTest/assets/ts/sample.ts.0.dump
@@ -28,7 +28,7 @@ track 256:
       data = length 22, hash CE183139
   sample count = 2
   sample 0:
-    time = 0
+    time = 33366
     flags = 1
     data = length 20711, hash 34341E8
   sample 1:
diff --git a/library/core/src/main/java/com/google/android/exoplayer2/extractor/ts/H262Reader.java b/library/core/src/main/java/com/google/android/exoplayer2/extractor/ts/H262Reader.java
index 160a9c5a717..a3502a3242d 100644
--- a/library/core/src/main/java/com/google/android/exoplayer2/extractor/ts/H262Reader.java
+++ b/library/core/src/main/java/com/google/android/exoplayer2/extractor/ts/H262Reader.java
@@ -51,17 +51,17 @@ public final class H262Reader implements ElementaryStreamReader {
   // State that should be reset on seek.
   private final boolean[] prefixFlags;
   private final CsdBuffer csdBuffer;
-  private boolean foundPicture;
   private long totalBytesWritten;
+  private boolean startedFirstSample;
 
   // Per packet state that gets reset at the start of each packet.
   private long pesTimeUs;
-  private boolean pesPtsUsAvailable;
 
-  // Per sample state that gets reset at the start of each frame.
-  private boolean isKeyframe;
+  // Per sample state that gets reset at the start of each sample.
   private long samplePosition;
   private long sampleTimeUs;
+  private boolean sampleIsKeyframe;
+  private boolean sampleHasPicture;
 
   public H262Reader() {
     prefixFlags = new boolean[4];
@@ -72,10 +72,8 @@ public H262Reader() {
   public void seek() {
     NalUnitUtil.clearPrefixFlags(prefixFlags);
     csdBuffer.reset();
-    pesPtsUsAvailable = false;
-    foundPicture = false;
-    samplePosition = C.POSITION_UNSET;
     totalBytesWritten = 0;
+    startedFirstSample = false;
   }
 
   @Override
@@ -87,10 +85,7 @@ public void createTracks(ExtractorOutput extractorOutput, TrackIdGenerator idGen
 
   @Override
   public void packetStarted(long pesTimeUs, boolean dataAlignmentIndicator) {
-    pesPtsUsAvailable = pesTimeUs != C.TIME_UNSET;
-    if (pesPtsUsAvailable) {
-      this.pesTimeUs = pesTimeUs;
-    }
+    this.pesTimeUs = pesTimeUs;
   }
 
   @Override
@@ -136,27 +131,26 @@ public void consume(ParsableByteArray data) {
         }
       }
 
-      if (hasOutputFormat
-          && (startCodeValue == START_PICTURE || startCodeValue == START_SEQUENCE_HEADER)) {
+      if (startCodeValue == START_PICTURE || startCodeValue == START_SEQUENCE_HEADER) {
         int bytesWrittenPastStartCode = limit - startCodeOffset;
-        boolean resetSample = samplePosition == C.POSITION_UNSET;
-        if (foundPicture) {
-          @C.BufferFlags int flags = isKeyframe ? C.BUFFER_FLAG_KEY_FRAME : 0;
+        if (startedFirstSample && sampleHasPicture && hasOutputFormat) {
+          // Output the sample.
+          @C.BufferFlags int flags = sampleIsKeyframe ? C.BUFFER_FLAG_KEY_FRAME : 0;
           int size = (int) (totalBytesWritten - samplePosition) - bytesWrittenPastStartCode;
           output.sampleMetadata(sampleTimeUs, flags, size, bytesWrittenPastStartCode, null);
-          isKeyframe = false;
-          resetSample = true;
         }
-        foundPicture = startCodeValue == START_PICTURE;
-        if (resetSample) {
+        if (!startedFirstSample || sampleHasPicture) {
+          // Start the next sample.
           samplePosition = totalBytesWritten - bytesWrittenPastStartCode;
-          sampleTimeUs = (pesPtsUsAvailable ? pesTimeUs : sampleTimeUs + frameDurationUs);
-          pesPtsUsAvailable = false;
+          sampleTimeUs = pesTimeUs != C.TIME_UNSET ? pesTimeUs
+              : (startedFirstSample ? (sampleTimeUs + frameDurationUs) : 0);
+          sampleIsKeyframe = false;
+          pesTimeUs = C.TIME_UNSET;
+          startedFirstSample = true;
         }
-      }
-
-      if (hasOutputFormat && startCodeValue == START_GROUP) {
-        isKeyframe = true;
+        sampleHasPicture = startCodeValue == START_PICTURE;
+      } else if (startCodeValue == START_GROUP) {
+        sampleIsKeyframe = true;
       }
 
       offset = startCodeOffset + 3;