From 996b1462a92ed8e186436800cbf9e7837fe99028 Mon Sep 17 00:00:00 2001
From: ThibaultBee <37510686+ThibaultBee@users.noreply.github.com>
Date: Sat, 14 Oct 2023 17:22:05 +0200
Subject: [PATCH] refactor(core): improve ByteBuffer extensions

---
 .../internal/encoders/MediaCodecEncoder.kt    |  2 +-
 .../internal/muxers/flv/packet/VideoTag.kt    |  6 +-
 .../internal/muxers/mp4/models/TrackChunks.kt |  6 +-
 .../internal/utils/av/video/H26X.kt           | 22 ------
 .../avc/AVCDecoderConfigurationRecord.kt      | 10 +--
 .../hevc/HEVCDecoderConfigurationRecord.kt    | 12 +--
 .../utils/extensions/ByteBufferExtensions.kt  | 78 ++++++++++++-------
 7 files changed, 70 insertions(+), 66 deletions(-)

diff --git a/core/src/main/java/io/github/thibaultbee/streampack/internal/encoders/MediaCodecEncoder.kt b/core/src/main/java/io/github/thibaultbee/streampack/internal/encoders/MediaCodecEncoder.kt
index 7e93bc10a..e36f85bf8 100644
--- a/core/src/main/java/io/github/thibaultbee/streampack/internal/encoders/MediaCodecEncoder.kt
+++ b/core/src/main/java/io/github/thibaultbee/streampack/internal/encoders/MediaCodecEncoder.kt
@@ -25,10 +25,10 @@ import io.github.thibaultbee.streampack.data.Config
 import io.github.thibaultbee.streampack.error.StreamPackError
 import io.github.thibaultbee.streampack.internal.data.Frame
 import io.github.thibaultbee.streampack.internal.events.EventHandler
-import io.github.thibaultbee.streampack.internal.utils.extensions.toByteArray
 import io.github.thibaultbee.streampack.internal.utils.extensions.isAudio
 import io.github.thibaultbee.streampack.internal.utils.extensions.slices
 import io.github.thibaultbee.streampack.internal.utils.extensions.startsWith
+import io.github.thibaultbee.streampack.internal.utils.extensions.toByteArray
 import io.github.thibaultbee.streampack.logger.Logger
 import io.github.thibaultbee.streampack.utils.TAG
 import java.nio.ByteBuffer
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/internal/muxers/flv/packet/VideoTag.kt b/core/src/main/java/io/github/thibaultbee/streampack/internal/muxers/flv/packet/VideoTag.kt
index e8e414ed5..8df9d2837 100644
--- a/core/src/main/java/io/github/thibaultbee/streampack/internal/muxers/flv/packet/VideoTag.kt
+++ b/core/src/main/java/io/github/thibaultbee/streampack/internal/muxers/flv/packet/VideoTag.kt
@@ -18,11 +18,11 @@ package io.github.thibaultbee.streampack.internal.muxers.flv.packet
 import android.media.MediaFormat
 import io.github.thibaultbee.streampack.data.VideoConfig
 import io.github.thibaultbee.streampack.internal.utils.av.video.avc.AVCDecoderConfigurationRecord
-import io.github.thibaultbee.streampack.internal.utils.av.video.getStartCodeSize
 import io.github.thibaultbee.streampack.internal.utils.av.video.hevc.HEVCDecoderConfigurationRecord
-import io.github.thibaultbee.streampack.internal.utils.av.video.removeStartCode
 import io.github.thibaultbee.streampack.internal.utils.extensions.put
 import io.github.thibaultbee.streampack.internal.utils.extensions.putInt24
+import io.github.thibaultbee.streampack.internal.utils.extensions.removeStartCode
+import io.github.thibaultbee.streampack.internal.utils.extensions.startCodeSize
 import java.io.IOException
 import java.nio.ByteBuffer
 
@@ -139,7 +139,7 @@ class VideoTag(
                 }
             }
         } else {
-            return buffers[0].remaining() - buffers[0].getStartCodeSize() + 4 // Replace start code with annex B
+            return buffers[0].remaining() - buffers[0].startCodeSize + 4 // Replace start code with annex B
         }
     }
 }
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/internal/muxers/mp4/models/TrackChunks.kt b/core/src/main/java/io/github/thibaultbee/streampack/internal/muxers/mp4/models/TrackChunks.kt
index f635f8240..6c6d9be28 100644
--- a/core/src/main/java/io/github/thibaultbee/streampack/internal/muxers/mp4/models/TrackChunks.kt
+++ b/core/src/main/java/io/github/thibaultbee/streampack/internal/muxers/mp4/models/TrackChunks.kt
@@ -58,10 +58,10 @@ import io.github.thibaultbee.streampack.internal.utils.av.descriptors.AudioSpeci
 import io.github.thibaultbee.streampack.internal.utils.av.descriptors.ESDescriptor
 import io.github.thibaultbee.streampack.internal.utils.av.descriptors.SLConfigDescriptor
 import io.github.thibaultbee.streampack.internal.utils.av.video.avc.AVCDecoderConfigurationRecord
-import io.github.thibaultbee.streampack.internal.utils.av.video.getStartCodeSize
 import io.github.thibaultbee.streampack.internal.utils.av.video.hevc.HEVCDecoderConfigurationRecord
-import io.github.thibaultbee.streampack.internal.utils.av.video.removeStartCode
 import io.github.thibaultbee.streampack.internal.utils.extensions.clone
+import io.github.thibaultbee.streampack.internal.utils.extensions.removeStartCode
+import io.github.thibaultbee.streampack.internal.utils.extensions.startCodeSize
 import java.nio.ByteBuffer
 
 /**
@@ -108,7 +108,7 @@ class TrackChunks(
             MediaFormat.MIMETYPE_VIDEO_HEVC,
             MediaFormat.MIMETYPE_VIDEO_AVC -> {
                 // Replace start code with size (from Annex B to AVCC)
-                4 + buffer.remaining() - buffer.getStartCodeSize()
+                4 + buffer.remaining() - buffer.startCodeSize
             }
 
             else -> {
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/video/H26X.kt b/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/video/H26X.kt
index d1804fe9d..28a7d45ec 100644
--- a/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/video/H26X.kt
+++ b/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/video/H26X.kt
@@ -15,28 +15,6 @@
  */
 package io.github.thibaultbee.streampack.internal.utils.av.video
 
-import java.nio.ByteBuffer
-
-fun ByteBuffer.getStartCodeSize(): Int {
-    return if (this.get(0) == 0x00.toByte() && this.get(1) == 0x00.toByte()
-        && this.get(2) == 0x00.toByte() && this.get(3) == 0x01.toByte()
-    ) {
-        4
-    } else if (this.get(0) == 0x00.toByte() && this.get(1) == 0x00.toByte()
-        && this.get(2) == 0x01.toByte()
-    ) {
-        3
-    } else {
-        0
-    }
-}
-
-fun ByteBuffer.removeStartCode(): ByteBuffer {
-    val startCodeSize = this.getStartCodeSize()
-    this.position(startCodeSize)
-    return this.slice()
-}
-
 enum class ChromaFormat(val value: Byte) {
     YUV400(0),
     YUV420(1),
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/video/avc/AVCDecoderConfigurationRecord.kt b/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/video/avc/AVCDecoderConfigurationRecord.kt
index 7c8ca6d04..26c31b537 100644
--- a/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/video/avc/AVCDecoderConfigurationRecord.kt
+++ b/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/video/avc/AVCDecoderConfigurationRecord.kt
@@ -16,11 +16,11 @@
 package io.github.thibaultbee.streampack.internal.utils.av.video.avc
 
 import io.github.thibaultbee.streampack.internal.utils.av.video.ChromaFormat
-import io.github.thibaultbee.streampack.internal.utils.av.video.getStartCodeSize
-import io.github.thibaultbee.streampack.internal.utils.av.video.removeStartCode
 import io.github.thibaultbee.streampack.internal.utils.extensions.put
 import io.github.thibaultbee.streampack.internal.utils.extensions.putShort
+import io.github.thibaultbee.streampack.internal.utils.extensions.removeStartCode
 import io.github.thibaultbee.streampack.internal.utils.extensions.shl
+import io.github.thibaultbee.streampack.internal.utils.extensions.startCodeSize
 import java.nio.ByteBuffer
 
 data class AVCDecoderConfigurationRecord(
@@ -105,12 +105,12 @@ data class AVCDecoderConfigurationRecord(
             var size =
                 AVC_DECODER_CONFIGURATION_RECORD_SIZE
             sps.forEach {
-                size += it.remaining() - it.getStartCodeSize()
+                size += it.remaining() - it.startCodeSize
             }
             pps.forEach {
-                size += it.remaining() - it.getStartCodeSize()
+                size += it.remaining() - it.startCodeSize
             }
-            val spsStartCodeSize = sps[0].getStartCodeSize()
+            val spsStartCodeSize = sps[0].startCodeSize
             val profileIdc = sps[0].get(spsStartCodeSize + 1).toInt()
             if ((profileIdc == 100) || (profileIdc == 110) || (profileIdc == 122) || (profileIdc == 144)) {
                 size += 4
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/video/hevc/HEVCDecoderConfigurationRecord.kt b/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/video/hevc/HEVCDecoderConfigurationRecord.kt
index 7c1f01beb..5847d0150 100644
--- a/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/video/hevc/HEVCDecoderConfigurationRecord.kt
+++ b/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/video/hevc/HEVCDecoderConfigurationRecord.kt
@@ -16,13 +16,13 @@
 package io.github.thibaultbee.streampack.internal.utils.av.video.hevc
 
 import io.github.thibaultbee.streampack.internal.utils.av.video.ChromaFormat
-import io.github.thibaultbee.streampack.internal.utils.av.video.getStartCodeSize
-import io.github.thibaultbee.streampack.internal.utils.av.video.removeStartCode
 import io.github.thibaultbee.streampack.internal.utils.extensions.put
 import io.github.thibaultbee.streampack.internal.utils.extensions.putLong48
 import io.github.thibaultbee.streampack.internal.utils.extensions.putShort
+import io.github.thibaultbee.streampack.internal.utils.extensions.removeStartCode
 import io.github.thibaultbee.streampack.internal.utils.extensions.shl
 import io.github.thibaultbee.streampack.internal.utils.extensions.shr
+import io.github.thibaultbee.streampack.internal.utils.extensions.startCodeSize
 import java.nio.ByteBuffer
 
 data class HEVCDecoderConfigurationRecord(
@@ -125,7 +125,7 @@ data class HEVCDecoderConfigurationRecord(
             parameterSets: List<ByteBuffer>
         ): HEVCDecoderConfigurationRecord {
             val nalUnitParameterSets = parameterSets.map {
-                val nalType = (it[it.getStartCodeSize()] shr 1 and 0x3F).toByte()
+                val nalType = (it[it.startCodeSize] shr 1 and 0x3F).toByte()
                 val type = NalUnit.Type.fromValue(nalType)
                 NalUnit(type, it)
             }
@@ -168,13 +168,13 @@ data class HEVCDecoderConfigurationRecord(
             var size =
                 HEVC_DECODER_CONFIGURATION_RECORD_SIZE
             sps.forEach {
-                size += it.remaining() - it.getStartCodeSize() + HEVC_PARAMETER_SET_HEADER_SIZE
+                size += it.remaining() - it.startCodeSize + HEVC_PARAMETER_SET_HEADER_SIZE
             }
             pps.forEach {
-                size += it.remaining() - it.getStartCodeSize() + HEVC_PARAMETER_SET_HEADER_SIZE
+                size += it.remaining() - it.startCodeSize + HEVC_PARAMETER_SET_HEADER_SIZE
             }
             vps.forEach {
-                size += it.remaining() - it.getStartCodeSize() + HEVC_PARAMETER_SET_HEADER_SIZE
+                size += it.remaining() - it.startCodeSize + HEVC_PARAMETER_SET_HEADER_SIZE
             }
 
             return size
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/extensions/ByteBufferExtensions.kt b/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/extensions/ByteBufferExtensions.kt
index 21e15b461..8c53c157d 100644
--- a/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/extensions/ByteBufferExtensions.kt
+++ b/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/extensions/ByteBufferExtensions.kt
@@ -15,7 +15,6 @@
  */
 package io.github.thibaultbee.streampack.internal.utils.extensions
 
-import io.github.thibaultbee.streampack.internal.utils.av.video.getStartCodeSize
 import java.nio.ByteBuffer
 import java.nio.ByteOrder
 import java.nio.charset.StandardCharsets
@@ -95,6 +94,23 @@ fun ByteBuffer.put(buffer: ByteBuffer, offset: Int, length: Int) {
     buffer.limit(limit)
 }
 
+fun ByteBuffer.getString(size: Int = this.remaining()): String {
+    val bytes = ByteArray(size)
+    this.get(bytes)
+    return String(bytes, StandardCharsets.UTF_8)
+}
+
+fun ByteBuffer.getLong(isLittleEndian: Boolean): Long {
+    if (isLittleEndian) {
+        order(ByteOrder.LITTLE_ENDIAN)
+    }
+    val value = long
+    if (isLittleEndian) {
+        order(ByteOrder.BIG_ENDIAN)
+    }
+    return value
+}
+
 fun ByteBuffer.indicesOf(prefix: ByteArray): List<Int> {
     if (prefix.isEmpty()) {
         return emptyList()
@@ -175,12 +191,46 @@ fun ByteBuffer.toByteArray(): ByteArray {
     }
 }
 
+fun ByteBuffer.clone(): ByteBuffer {
+    val clone = ByteBuffer.allocate(this.remaining())
+    return clone.put(this).apply { rewind() }
+}
+
+/**
+ * For AVC and HEVC
+ */
+
+
+/**
+ * Get start code size of [ByteBuffer].
+ */
+val ByteBuffer.startCodeSize: Int
+    get() {
+        return if (this.get(0) == 0x00.toByte() && this.get(1) == 0x00.toByte()
+            && this.get(2) == 0x00.toByte() && this.get(3) == 0x01.toByte()
+        ) {
+            4
+        } else if (this.get(0) == 0x00.toByte() && this.get(1) == 0x00.toByte()
+            && this.get(2) == 0x01.toByte()
+        ) {
+            3
+        } else {
+            0
+        }
+    }
+
+fun ByteBuffer.removeStartCode(): ByteBuffer {
+    val startCodeSize = this.startCodeSize
+    this.position(startCodeSize)
+    return this.slice()
+}
+
 fun ByteBuffer.extractRbsp(headerLength: Int): ByteBuffer {
     val rbsp = ByteBuffer.allocateDirect(this.remaining())
 
     val indices = this.indicesOf(byteArrayOf(0x00, 0x00, 0x03))
 
-    rbsp.put(this, this.getStartCodeSize(), headerLength)
+    rbsp.put(this, this.startCodeSize, headerLength)
 
     var previous = this.position()
     indices.forEach {
@@ -193,27 +243,3 @@ fun ByteBuffer.extractRbsp(headerLength: Int): ByteBuffer {
     rbsp.rewind()
     return rbsp
 }
-
-fun ByteBuffer.clone(): ByteBuffer {
-    val clone = ByteBuffer.allocate(this.remaining())
-    return clone.put(this).apply { rewind() }
-}
-
-fun ByteBuffer.getString(size: Int = this.remaining()): String {
-    val bytes = ByteArray(size)
-    this.get(bytes)
-    return String(bytes, StandardCharsets.UTF_8)
-}
-
-
-fun ByteBuffer.getLong(isLittleEndian: Boolean): Long {
-    if (isLittleEndian) {
-        order(ByteOrder.LITTLE_ENDIAN)
-    }
-    val value = long
-    if (isLittleEndian) {
-        order(ByteOrder.BIG_ENDIAN)
-    }
-    return value
-}
-