h264: Print NAL units, parse SPS,PPS

3 files changed, 254 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
index a54c865..60adada 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+*.264
 *.DS_Store
 *.cache
 *.cso
diff --git a/h264/CMakeLists.txt b/h264/CMakeLists.txt
new file mode 100644
index 0000000..9ba7a72
--- /dev/null
+++ b/h264/CMakeLists.txt
@@ -0,0 +1,7 @@
+cmake_minimum_required(VERSION 3.15 FATAL_ERROR)
+project(h264 C)
+
+include("${CMAKE_CURRENT_LIST_DIR}/../common/c_cpp/CMakeLists.txt")
+
+add_executable(h264 h264.c)
+target_link_libraries(h264 PRIVATE common)
diff --git a/h264/h264.c b/h264/h264.c
new file mode 100644
index 0000000..52aee9e
--- /dev/null
+++ b/h264/h264.c
@@ -0,0 +1,246 @@
+// Barebones H.264 (baseline) bitstream parser. P-slices only.
+//
+// SPDX-License-Identifier: 0BSD
+//
+// ref: https://www.itu.int/rec/T-REC-H.264
+
+// ffmpeg -i ~/badapple.mp4 -c:v libx264 -profile:v baseline -level 3.0 -preset fast -g 1 -f h264 test.264
+
+#include <stdbool.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#define ASSERT(cond) { if (!(cond)) { printf("%s:%d: %s\n", __FILE__, __LINE__, #cond); abort(); } }
+
+static struct
+{
+    const uint8_t* buf;
+    size_t len;
+    size_t cur;
+    uint8_t rbsp_buf[1 << 20];
+    size_t rbsp_len;
+    size_t rbsp_cbyte;
+    size_t rbsp_cbit;
+} G = { 0 };
+
+// Search for the next Annex-B start code (ref: §B.1)
+// Returns the size of the start code
+static size_t FindStartCode(size_t start, size_t* off)
+{
+    for (size_t i = start; i < G.len; ++i)
+    {
+        const uint8_t* p = &G.buf[i];
+        const bool start3 = G.len - i >= 3 && p[0] == 0x00 && p[1] == 0x00 && p[2] == 0x01;
+        const bool start4 = G.len - i >= 4 && p[0] == 0x00 && p[1] == 0x00 && p[2] == 0x00 && p[3] == 0x01;
+        if (start3) { *off = i; return 3; }
+        if (start4) { *off = i; return 4; }
+    }
+    return 0;
+}
+
+// Search for the next NAL unit
+static bool FindNalUnit(void)
+{
+    // Find the next start code
+    size_t start_len = 0;
+    if (!(start_len = FindStartCode(G.cur, &G.cur)))
+    {
+        return false;
+    }
+    G.cur += start_len;
+    ASSERT(G.cur < G.len);
+
+    // Find end of RBSP
+    size_t end = 0;
+    if (!FindStartCode(G.cur, &end))
+    {
+        end = G.len;
+    }
+    ASSERT(end - G.cur < sizeof(G.rbsp_buf));
+    
+    // Remove emulation bytes
+    size_t z = 0;
+    G.rbsp_len = 0;
+    G.rbsp_buf[G.rbsp_len++] = G.buf[G.cur];
+    for (size_t i = G.cur + 1; i < end; ++i)
+    {
+        const uint8_t b = G.buf[i];
+
+        // skip third byte of 0x00 0x00 0x03
+        if (b == 0x00)
+        {
+            ++z;
+        }
+        else
+        {
+            bool skip = b == 0x03 && z == 2;
+            z = 0;
+            if (skip)
+            {
+                continue;
+            }
+        }
+
+        G.rbsp_buf[G.rbsp_len++] = b;
+    }
+    ASSERT(G.rbsp_len > 0);
+
+    // Trim trailing bits (ref: §7.3.2.11)
+#if 0
+    const uint8_t tail = G.rbsp_buf[G.rbsp_len - 1];
+    int stop_bit_pos = -1;
+    for (int i = 7; i >= 0; --i)
+    {
+        if (tail & (1 << i))
+        {
+            stop_bit_pos = i;
+            break;
+        }
+    }
+    ASSERT(stop_bit_pos >= 0 && stop_bit_pos < 8);
+    G.rbsp_buf[G.rbsp_len - 1] &= ~((1 << stop_bit_pos) - 1);
+#endif
+
+    G.cur = end;
+
+    return true;
+}
+
+// Read bits from RBSP stream
+static uint32_t ReadRbspBits(size_t n)
+{
+    // @@ check overrun
+    uint32_t result = 0;
+    for (size_t i = 0; i < n; ++i)
+    {
+        const uint8_t bit = (G.rbsp_buf[G.rbsp_cbyte] >> (7 - G.rbsp_cbit)) & 0x01;
+        if (++G.rbsp_cbit >= 8)
+        {
+            ++G.rbsp_cbyte;
+            G.rbsp_cbit = 0;
+        }
+        result <<= 1;
+        result |= bit;
+    }
+    return result;
+}
+
+// Read bits from RBSP stream without advancing the cursor
+static uint32_t PeekRbspBits(size_t n)
+{
+    const size_t cbyte = G.rbsp_cbyte;
+    const size_t cbit = G.rbsp_cbit;
+    const uint32_t r = ReadRbspBits(n);
+    G.rbsp_cbyte = cbyte;
+    G.rbsp_cbit = cbit;
+    return r;
+}
+
+// Seek forward in RBSP stream
+static void SkipRbspBits(size_t n)
+{
+    // @@ check overrun
+    G.rbsp_cbit += n;
+    G.rbsp_cbyte += G.rbsp_cbit / 8;
+    G.rbsp_cbit %= 8;
+}
+
+// Read bits (Exponential-Golomb) from RBSP stream (ref: §9.1)
+static uint32_t ReadRbspBitsEG(void)
+{
+    // @@ check overrun
+    uint32_t z = 0;
+    while (PeekRbspBits(1) == 0)
+    {
+        ++z;
+        SkipRbspBits(1);
+    }
+
+    SkipRbspBits(1);
+    if (z == 0)
+    {
+        return 0;
+    }
+
+    const uint32_t suffix = ReadRbspBits(z);
+    return (1u << z) - 1u + suffix;
+}
+
+static void SpewNalParam(const char* name, uint32_t val)
+{
+    printf("  %-20s %d\n", name, val);
+}
+#define SPEW_NAL_PARAM(x) SpewNalParam(#x, x)
+
+int main(int argc, const char* argv[])
+{
+    (void)argc; (void)argv;
+
+    // To keep things simple, just load the entire file into memory.
+    // Bad Apple is only about 16MB
+    {
+        FILE* f = fopen("test.264", "rb");
+        ASSERT(f && "Test file not found");
+        fseek(f, 0, SEEK_END);
+        G.len = ftell(f);
+        fseek(f, 0, SEEK_SET);
+        G.buf = malloc(G.len); ASSERT(G.buf);
+        fread((void*)G.buf, G.len, 1, f);
+        fclose(f);
+    }
+
+    for (int i = 0; i < 10; ++i)
+    {
+        if (!FindNalUnit())
+        {
+            ASSERT(0);
+        }
+        // Reset RBSP bit stream
+        G.rbsp_cbyte = 0;
+        G.rbsp_cbit = 0;
+
+        // Read NAL unit header (ref: §7.3.1)
+        if (ReadRbspBits(1) != 0)
+        {
+            ASSERT(0 && "Got non-zero first bit of NAL header. Something is wrong.");
+        }
+        SkipRbspBits(2); // nal_ref_idc
+        const uint32_t nal_type = ReadRbspBits(5);
+
+        // ref: Table 7-1
+        switch (nal_type)
+        {
+        case 7:
+        {
+            // ref: §7.3.2.1
+            printf("NAL UNIT(Sequence parameter set)\n");
+            const uint32_t profile_idc = ReadRbspBits(8);
+            SkipRbspBits(8); // constraint_setX_flag
+            const uint32_t level_idc = ReadRbspBits(8);
+            const uint32_t seq_parameter_set_id = ReadRbspBitsEG();
+            SPEW_NAL_PARAM(profile_idc);
+            SPEW_NAL_PARAM(level_idc);
+            SPEW_NAL_PARAM(seq_parameter_set_id);
+            break;
+        }
+        case 8:
+        {
+            // ref: §7.3.2.2
+            printf("NAL UNIT(Picture parameter set)\n");
+            const uint32_t pic_parameter_set_id = ReadRbspBitsEG();
+            const uint32_t seq_parameter_set_id = ReadRbspBitsEG();
+            SPEW_NAL_PARAM(pic_parameter_set_id);
+            SPEW_NAL_PARAM(seq_parameter_set_id);
+            break;
+        }
+        default:
+        {
+            printf("NAL UNIT(%d)\n", nal_type);
+            break;
+        }  
+        }
+    }
+
+    return EXIT_SUCCESS;
+}