From d620396b9d653ddace830955fded4455a4e48550 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=A4lim=C3=A4ki=20Tuomas?= <tuomas.valimaki@insta.fi>
Date: Fri, 16 Aug 2024 14:41:28 +0300
Subject: [PATCH] Refactor opus writing logic to OpusFrameWriter

---
 .../streampack/internal/muxers/ts/TSMuxer.kt  | 34 +--------------
 .../utils/av/audio/opus/OpusFrameWriter.kt    | 42 +++++++++++++++++++
 2 files changed, 44 insertions(+), 32 deletions(-)
 create mode 100644 core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/audio/opus/OpusFrameWriter.kt

diff --git a/core/src/main/java/io/github/thibaultbee/streampack/internal/muxers/ts/TSMuxer.kt b/core/src/main/java/io/github/thibaultbee/streampack/internal/muxers/ts/TSMuxer.kt
index 235912fdf..7768270b6 100644
--- a/core/src/main/java/io/github/thibaultbee/streampack/internal/muxers/ts/TSMuxer.kt
+++ b/core/src/main/java/io/github/thibaultbee/streampack/internal/muxers/ts/TSMuxer.kt
@@ -34,6 +34,7 @@ import io.github.thibaultbee.streampack.internal.muxers.ts.utils.MuxerConst
 import io.github.thibaultbee.streampack.internal.muxers.ts.utils.TSConst
 import io.github.thibaultbee.streampack.internal.utils.av.audio.aac.ADTSFrameWriter
 import io.github.thibaultbee.streampack.internal.utils.av.audio.aac.LATMFrameWriter
+import io.github.thibaultbee.streampack.internal.utils.av.audio.opus.OpusFrameWriter
 import java.nio.ByteBuffer
 import java.util.MissingFormatArgumentException
 import kotlin.random.Random
@@ -150,13 +151,7 @@ class TSMuxer(
             }
 
             MediaFormat.MIMETYPE_AUDIO_OPUS -> {
-                val payload = frame.rawBuffer
-                val controlHeader = opusControlHeader(payload.remaining())
-                val opusAccessUnit = ByteBuffer.allocate(controlHeader.remaining() + payload.remaining())
-                opusAccessUnit.put(controlHeader)
-                opusAccessUnit.put(payload)
-                opusAccessUnit.rewind()
-                frame.copy(rawBuffer = opusAccessUnit)
+                frame.copy(rawBuffer = OpusFrameWriter.fromPayload(frame.rawBuffer).toByteBuffer())
             }
             else -> throw IllegalArgumentException("Unsupported mimeType ${frame.mimeType}")
         }
@@ -166,31 +161,6 @@ class TSMuxer(
         }
     }
 
-    private fun opusControlHeader(payloadSize: Int): ByteBuffer {
-        val payloadSizeFullBytesCount = payloadSize / 255
-        val payloadSizeRemainderByte = payloadSize % 255
-        val headerSize = 2 + payloadSizeFullBytesCount + payloadSizeRemainderByte.coerceAtMost(1)
-        val header = ByteBuffer.allocate(headerSize)
-
-        // control_header_prefix 11 bits (0x3FF or 01111111111)
-        // start_trim_flag 1 bit (0)
-        // end_trim_flag 1 bit (0)
-        // control_extension_flag 1 bit (0)
-        // reserved 2 bits (0)
-        header.put(0x7F.toByte())
-        header.put(0xE0.toByte())
-
-        repeat(payloadSizeFullBytesCount) {
-            header.put(0xFF.toByte())
-        }
-        if (payloadSizeRemainderByte > 0) {
-            header.put(payloadSizeRemainderByte.toByte())
-        }
-
-        header.rewind()
-        return header
-    }
-
     /**
      * Generate MPEG-TS table and elementary stream from the frame
      * @param pes Pes containing infos on the stream
diff --git a/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/audio/opus/OpusFrameWriter.kt b/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/audio/opus/OpusFrameWriter.kt
new file mode 100644
index 000000000..747ff6c65
--- /dev/null
+++ b/core/src/main/java/io/github/thibaultbee/streampack/internal/utils/av/audio/opus/OpusFrameWriter.kt
@@ -0,0 +1,42 @@
+package io.github.thibaultbee.streampack.internal.utils.av.audio.opus
+
+import io.github.thibaultbee.streampack.internal.utils.av.buffer.ByteBufferWriter
+import java.nio.ByteBuffer
+
+class OpusFrameWriter(private val frameBuffer: ByteBuffer) : ByteBufferWriter() {
+    private val payloadSize: Int
+        get() = frameBuffer.remaining()
+    private val payloadSizeFullBytesCount: Int
+        get() = payloadSize / 255
+    private val payloadSizeRemainderByte: Int
+        get() = payloadSize % 255
+    private val controlHeaderSize: Int
+        get() = 2 + payloadSizeFullBytesCount + payloadSizeRemainderByte.coerceAtMost(1)
+
+    override val size = controlHeaderSize + payloadSize
+
+    override fun write(output: ByteBuffer) {
+        // control_header_prefix 11 bits (0x3FF or 01111111111)
+        // start_trim_flag 1 bit (0)
+        // end_trim_flag 1 bit (0)
+        // control_extension_flag 1 bit (0)
+        // reserved 2 bits (0)
+        output.put(0x7F.toByte())
+        output.put(0xE0.toByte())
+
+        repeat(payloadSizeFullBytesCount) {
+            output.put(0xFF.toByte())
+        }
+        if (payloadSizeRemainderByte > 0) {
+            output.put(payloadSizeRemainderByte.toByte())
+        }
+
+        output.put(frameBuffer)
+    }
+
+    companion object {
+        fun fromPayload(frameBuffer: ByteBuffer): OpusFrameWriter {
+            return OpusFrameWriter(frameBuffer)
+        }
+    }
+}
\ No newline at end of file