Add SubtitleExtractor which wraps a SubtitleDecoder.

SubtitleExtractor is a component that extracts subtitle data taken from
ExtractorInput into samples. Samples are pushed into an ExtractorOutput
(usually SampleQueue). As a temporary solution SubtitleExtractor uses
SubtitleDecoder to extract Cues from input data.

PiperOrigin-RevId: 390319875

library/extractor/src/main/java/com/google/android/exoplayer2/extractor/subtitle/SubtitleExtractor.java

@@ -0,0 +1,210 @@
+/*
+ * Copyright 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.android.exoplayer2.extractor.subtitle;
+
+import static com.google.android.exoplayer2.util.Assertions.checkState;
+import static com.google.android.exoplayer2.util.Assertions.checkStateNotNull;
+
+import androidx.annotation.IntDef;
+import androidx.annotation.Nullable;
+import com.google.android.exoplayer2.C;
+import com.google.android.exoplayer2.Format;
+import com.google.android.exoplayer2.ParserException;
+import com.google.android.exoplayer2.extractor.Extractor;
+import com.google.android.exoplayer2.extractor.ExtractorInput;
+import com.google.android.exoplayer2.extractor.ExtractorOutput;
+import com.google.android.exoplayer2.extractor.PositionHolder;
+import com.google.android.exoplayer2.extractor.TrackOutput;
+import com.google.android.exoplayer2.text.Cue;
+import com.google.android.exoplayer2.text.CueEncoder;
+import com.google.android.exoplayer2.text.SubtitleDecoder;
+import com.google.android.exoplayer2.text.SubtitleDecoderException;
+import com.google.android.exoplayer2.text.SubtitleInputBuffer;
+import com.google.android.exoplayer2.text.SubtitleOutputBuffer;
+import com.google.android.exoplayer2.util.ParsableByteArray;
+import com.google.common.primitives.Ints;
+import java.io.IOException;
+import java.io.InterruptedIOException;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.util.List;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
+
+/** Generic extractor for extracting subtitles from various subtitle formats. */
+public class SubtitleExtractor implements Extractor {
+  @IntDef(
+      value = {
+        STATE_CREATED,
+        STATE_INITIALIZED,
+        STATE_READING,
+        STATE_DECODING,
+        STATE_WRITING,
+        STATE_FINISHED
+      })
+  @Retention(RetentionPolicy.SOURCE)
+  private @interface State {}
+
+  /** The extractor has been created. */
+  private static final int STATE_CREATED = 0;
+  /** The extractor has been initialized. */
+  private static final int STATE_INITIALIZED = 1;
+  /** The extractor is reading data from the input. */
+  private static final int STATE_READING = 2;
+  /** The extractor is queueing data for decoding. */
+  private static final int STATE_DECODING = 3;
+  /** The extractor is writing data to the output. */
+  private static final int STATE_WRITING = 4;
+  /** The extractor has finished writing. */
+  private static final int STATE_FINISHED = 5;
+
+  private static final int DEFAULT_BUFFER_SIZE = 1024;
+
+  private final SubtitleDecoder subtitleDecoder;
+  private final CueEncoder cueEncoder;
+  private final ParsableByteArray subtitleData;
+  private final Format format;
+
+  private @MonotonicNonNull ExtractorOutput extractorOutput;
+  private @MonotonicNonNull TrackOutput trackOutput;
+  private int bytesRead;
+  @State private int state;
+
+  public SubtitleExtractor(SubtitleDecoder subtitleDecoder, Format format) {
+    this.subtitleDecoder = subtitleDecoder;
+    cueEncoder = new CueEncoder();
+    subtitleData = new ParsableByteArray();
+    this.format = format;
+    state = STATE_CREATED;
+  }
+
+  @Override
+  public boolean sniff(ExtractorInput input) throws IOException {
+    // TODO: Implement sniff() according to the Extractor interface documentation. For now sniff()
+    // can safely return true because we plan to use this class in an ExtractorFactory that returns
+    // exactly one Extractor implementation.
+    return true;
+  }
+
+  @Override
+  public void init(ExtractorOutput output) {
+    checkState(state == STATE_CREATED);
+    extractorOutput = output;
+    trackOutput = extractorOutput.track(/* id= */ 0, C.TRACK_TYPE_TEXT);
+    trackOutput.format(format);
+    state = STATE_INITIALIZED;
+  }
+
+  @Override
+  public int read(ExtractorInput input, PositionHolder seekPosition) throws IOException {
+    switch (state) {
+      case STATE_INITIALIZED:
+        prepareMemory(input);
+        state = readFromInput(input) ? STATE_DECODING : STATE_READING;
+        return Extractor.RESULT_CONTINUE;
+      case STATE_READING:
+        state = readFromInput(input) ? STATE_DECODING : STATE_READING;
+        return Extractor.RESULT_CONTINUE;
+      case STATE_DECODING:
+        queueDataToDecoder();
+        state = STATE_WRITING;
+        return RESULT_CONTINUE;
+      case STATE_WRITING:
+        writeToOutput();
+        state = STATE_FINISHED;
+        return Extractor.RESULT_END_OF_INPUT;
+      case STATE_FINISHED:
+        return Extractor.RESULT_END_OF_INPUT;
+      case STATE_CREATED:
+      default:
+        throw new IllegalStateException();
+    }
+  }
+
+  @Override
+  public void seek(long position, long timeUs) {}
+
+  @Override
+  public void release() {
+    // TODO: Proper implementation of this method is missing. Implement release() according to the
+    // Extractor interface documentation.
+  }
+
+  private void prepareMemory(ExtractorInput input) {
+    subtitleData.reset(
+        input.getLength() != C.LENGTH_UNSET
+            ? Ints.checkedCast(input.getLength())
+            : DEFAULT_BUFFER_SIZE);
+  }
+
+  /** Returns whether reading has been finished. */
+  private boolean readFromInput(ExtractorInput input) throws IOException {
+    if (subtitleData.capacity() == bytesRead) {
+      subtitleData.ensureCapacity(bytesRead + DEFAULT_BUFFER_SIZE);
+    }
+    int readResult =
+        input.read(subtitleData.getData(), bytesRead, subtitleData.capacity() - bytesRead);
+    if (readResult != C.RESULT_END_OF_INPUT) {
+      bytesRead += readResult;
+    }
+    return readResult == C.RESULT_END_OF_INPUT;
+  }
+
+  private void queueDataToDecoder() throws IOException {
+    try {
+      @Nullable SubtitleInputBuffer inputBuffer = subtitleDecoder.dequeueInputBuffer();
+      while (inputBuffer == null) {
+        inputBuffer = subtitleDecoder.dequeueInputBuffer();
+        Thread.sleep(5);
+      }
+      inputBuffer.ensureSpaceForWrite(bytesRead);
+      inputBuffer.data.put(subtitleData.getData(), /* offset= */ 0, bytesRead);
+      subtitleDecoder.queueInputBuffer(inputBuffer);
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new InterruptedIOException();
+    } catch (SubtitleDecoderException e) {
+      throw ParserException.createForMalformedContainer("SubtitleDecoder failed.", e);
+    }
+  }
+
+  private void writeToOutput() throws IOException {
+    checkStateNotNull(this.trackOutput);
+    try {
+      @Nullable SubtitleOutputBuffer outputBuffer = subtitleDecoder.dequeueOutputBuffer();
+      while (outputBuffer == null) {
+        outputBuffer = subtitleDecoder.dequeueOutputBuffer();
+        Thread.sleep(5);
+      }
+
+      for (int i = 0; i < outputBuffer.getEventTimeCount(); i++) {
+        List<Cue> cues = outputBuffer.getCues(outputBuffer.getEventTime(i));
+        byte[] cuesSample = cueEncoder.encode(cues);
+        trackOutput.sampleData(new ParsableByteArray(cuesSample), cuesSample.length);
+        trackOutput.sampleMetadata(
+            /* timeUs= */ outputBuffer.getEventTime(i),
+            /* flags= */ C.BUFFER_FLAG_KEY_FRAME,
+            /* size= */ cuesSample.length,
+            /* offset= */ 0,
+            /* cryptoData= */ null);
+      }
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new InterruptedIOException();
+    } catch (SubtitleDecoderException e) {
+      throw ParserException.createForMalformedContainer("SubtitleDecoder failed.", e);
+    }
+  }
+}

library/extractor/src/main/java/com/google/android/exoplayer2/extractor/subtitle/package-info.java

@@ -0,0 +1,19 @@
+/*
+ * Copyright 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+@NonNullApi
+package com.google.android.exoplayer2.extractor.subtitle;
+
+import com.google.android.exoplayer2.util.NonNullApi;

library/extractor/src/test/java/com/google/android/exoplayer2/extractor/subtitle/SubtitleExtractorTest.java

@@ -0,0 +1,101 @@
+/*
+ * Copyright 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.android.exoplayer2.extractor.subtitle;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertThrows;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.android.exoplayer2.Format;
+import com.google.android.exoplayer2.extractor.Extractor;
+import com.google.android.exoplayer2.testutil.FakeExtractorInput;
+import com.google.android.exoplayer2.testutil.FakeExtractorOutput;
+import com.google.android.exoplayer2.testutil.FakeTrackOutput;
+import com.google.android.exoplayer2.text.Cue;
+import com.google.android.exoplayer2.text.CueDecoder;
+import com.google.android.exoplayer2.text.webvtt.WebvttDecoder;
+import com.google.android.exoplayer2.util.Util;
+import java.util.List;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/** Tests for {@link SubtitleExtractor}. */
+@RunWith(AndroidJUnit4.class)
+public class SubtitleExtractorTest {
+  @Test
+  public void extractor_outputsCues() throws Exception {
+    String testData =
+        "WEBVTT\n"
+            + "\n"
+            + "00:00.000 --> 00:01.234\n"
+            + "This is the first subtitle.\n"
+            + "\n"
+            + "00:02.345 --> 00:03.456\n"
+            + "This is the second subtitle.\n"
+            + "\n"
+            + "00:02.600 --> 00:04.567\n"
+            + "This is the third subtitle.";
+    CueDecoder decoder = new CueDecoder();
+    FakeExtractorOutput output = new FakeExtractorOutput();
+    FakeExtractorInput input =
+        new FakeExtractorInput.Builder()
+            .setData(Util.getUtf8Bytes(testData))
+            .setSimulatePartialReads(true)
+            .build();
+    SubtitleExtractor extractor =
+        new SubtitleExtractor(new WebvttDecoder(), new Format.Builder().build());
+    extractor.init(output);
+
+    while (extractor.read(input, null) != Extractor.RESULT_END_OF_INPUT) {}
+
+    FakeTrackOutput trackOutput = output.trackOutputs.get(0);
+    assertThat(trackOutput.getSampleCount()).isEqualTo(6);
+    // Check sample timestamps.
+    assertThat(trackOutput.getSampleTimeUs(0)).isEqualTo(0L);
+    assertThat(trackOutput.getSampleTimeUs(1)).isEqualTo(1_234_000L);
+    assertThat(trackOutput.getSampleTimeUs(2)).isEqualTo(2_345_000L);
+    assertThat(trackOutput.getSampleTimeUs(3)).isEqualTo(2_600_000L);
+    assertThat(trackOutput.getSampleTimeUs(4)).isEqualTo(3_456_000L);
+    assertThat(trackOutput.getSampleTimeUs(5)).isEqualTo(4_567_000L);
+    // Check sample content.
+    List<Cue> cues0 = decoder.decode(trackOutput.getSampleData(0));
+    assertThat(cues0).hasSize(1);
+    assertThat(cues0.get(0).text.toString()).isEqualTo("This is the first subtitle.");
+    List<Cue> cues1 = decoder.decode(trackOutput.getSampleData(1));
+    assertThat(cues1).isEmpty();
+    List<Cue> cues2 = decoder.decode(trackOutput.getSampleData(2));
+    assertThat(cues2).hasSize(1);
+    assertThat(cues2.get(0).text.toString()).isEqualTo("This is the second subtitle.");
+    List<Cue> cues3 = decoder.decode(trackOutput.getSampleData(3));
+    assertThat(cues3).hasSize(2);
+    assertThat(cues3.get(0).text.toString()).isEqualTo("This is the second subtitle.");
+    assertThat(cues3.get(1).text.toString()).isEqualTo("This is the third subtitle.");
+    List<Cue> cues4 = decoder.decode(trackOutput.getSampleData(4));
+    assertThat(cues4).hasSize(1);
+    assertThat(cues4.get(0).text.toString()).isEqualTo("This is the third subtitle.");
+    List<Cue> cues5 = decoder.decode(trackOutput.getSampleData(5));
+    assertThat(cues5).isEmpty();
+  }
+
+  @Test
+  public void read_notInitialized_fails() {
+    FakeExtractorInput input = new FakeExtractorInput.Builder().setData(new byte[0]).build();
+    SubtitleExtractor extractor =
+        new SubtitleExtractor(new WebvttDecoder(), new Format.Builder().build());
+
+    assertThrows(IllegalStateException.class, () -> extractor.read(input, null));
+  }
+}