ts_indexer: support more index for CBS. [1/1]

PD#SWPL-128709

Problem:
CBS needs more index of TS/MPEG2/AVC/HEVC.

Solution:
add support.

Verify:
ubuntu.

Change-Id: I1a2e0c4cd5e1dfc3a07ff1b9f726b416d1c5b685
Signed-off-by: Yahui Han <yahui.han@amlogic.com>
diff --git a/include/ts_indexer.h b/include/ts_indexer.h
index 147ef18..8de8ea0 100644
--- a/include/ts_indexer.h
+++ b/include/ts_indexer.h
@@ -21,9 +21,26 @@
 
 /**Event type.*/
 typedef enum {
-  TS_INDEXER_EVENT_TYPE_VIDEO_I_FRAME, /**< Video I frame PTS.*/
-  TS_INDEXER_EVENT_TYPE_VIDEO_PTS,     /**< Video PTS.*/
-  TS_INDEXER_EVENT_TYPE_AUDIO_PTS      /**< Audio PTS.*/
+  TS_INDEXER_EVENT_TYPE_START_INDICATOR,            /**< TS start indicator.*/
+  TS_INDEXER_EVENT_TYPE_DISCONTINUITY_INDICATOR,    /**< TS discontinuity indicator.*/
+  TS_INDEXER_EVENT_TYPE_MPEG2_I_FRAME,              /**< MPEG2 I frame.*/
+  TS_INDEXER_EVENT_TYPE_MPEG2_P_FRAME,              /**< MPEG2 P frame.*/
+  TS_INDEXER_EVENT_TYPE_MPEG2_B_FRAME,              /**< MPEG2 B frame.*/
+  TS_INDEXER_EVENT_TYPE_MPEG2_SEQUENCE,             /**< MPEG2 Video Sequence header.*/
+  TS_INDEXER_EVENT_TYPE_AVC_I_SLICE,                /**< AVC I slice.*/
+  TS_INDEXER_EVENT_TYPE_AVC_B_SLICE,                /**< AVC B slice.*/
+  TS_INDEXER_EVENT_TYPE_AVC_P_SLICE,                /**< AVC P slice.*/
+  TS_INDEXER_EVENT_TYPE_AVC_SI_SLICE,               /**< AVC SI slice.*/
+  TS_INDEXER_EVENT_TYPE_AVC_SP_SLICE,               /**< AVC SP slice.*/
+  TS_INDEXER_EVENT_TYPE_HEVC_SPS,                   /**< HEVC NAL unit type SPS_NUT.*/
+  TS_INDEXER_EVENT_TYPE_HEVC_AUD,                   /**< HEVC NAL unit type AUD_NUT.*/
+  TS_INDEXER_EVENT_TYPE_HEVC_BLA_W_LP,              /**< HEVC NAL unit type BLA_W_LP.*/
+  TS_INDEXER_EVENT_TYPE_HEVC_BLA_W_RADL,            /**< HEVC NAL unit type BLA_W_RADL.*/
+  TS_INDEXER_EVENT_TYPE_HEVC_IDR_W_RADL,            /**< HEVC NAL unit type IDR_W_RADL.*/
+  TS_INDEXER_EVENT_TYPE_HEVC_IDR_N_LP,              /**< HEVC NAL unit type IDR_N_LP.*/
+  TS_INDEXER_EVENT_TYPE_HEVC_TRAIL_CRA,             /**< HEVC NAL unit type TRAIL_CRA.*/
+  TS_INDEXER_EVENT_TYPE_VIDEO_PTS,                  /**< MEPG2/AVC/HEVC PTS.*/
+  TS_INDEXER_EVENT_TYPE_AUDIO_PTS                   /**< Audio PTS.*/
 } TS_Indexer_EventType_t;
 
 /**Stream Parser state.*/
@@ -39,7 +56,7 @@
 typedef struct {
   TS_Indexer_EventType_t type;   /**< Event type.*/
   int                    pid;    /**< The PID of the stream.*/
-  uint64_t               offset; /**< The offset of the first TS packet of this frame.*/
+  uint64_t               offset; /**< the TS packet offset of this event.*/
   uint64_t               pts;    /**< The PTS of this frame.*/
 } TS_Indexer_Event_t;
 
diff --git a/src/ts_indexer.c b/src/ts_indexer.c
index c0cef04..a17fd37 100644
--- a/src/ts_indexer.c
+++ b/src/ts_indexer.c
@@ -16,6 +16,14 @@
 #define NAL_TYPE_IDR 5 // IDR NALU 类型
 #define NAL_TYPE_NON_IDR 1 // 非 IDR NALU 类型
 
+#define HEVC_NALU_BLA_W_LP      16
+#define HEVC_NALU_BLA_W_RADL    17
+#define HEVC_NALU_IDR_W_RADL    19
+#define HEVC_NALU_IDR_N_LP      20
+#define HEVC_NALU_TRAIL_CRA     21
+#define HEVC_NALU_SPS           33
+#define HEVC_NALU_AUD           35
+
 /**
  * Initialize the TS indexer.
  * \param ts_indexer The TS indexer to be initialized.
@@ -143,11 +151,14 @@
 {
   int i;
   uint32_t needle = 0;
+  uint32_t needle1 = 0;
   uint8_t *haystack = data;
   int haystack_len = len;
   int left = len;
   // start code of picture header
   uint8_t arr[4] = {0x00, 0x00, 0x01, 0x00};
+  // start code of sequence header
+  uint8_t arr1[4] = {0x00, 0x00, 0x01, 0xb3};
   TS_Indexer_Event_t event;
 
   /* mpeg header needs at least 4 bytes */
@@ -165,31 +176,59 @@
     needle += (arr[i] << (8 * i));
   }
 
+  for (i = 0; i < 4; ++i) {
+    needle1 += (arr1[i] << (8 * i));
+  }
+
   for (i = 0; i < haystack_len - sizeof(needle) + 1;) {
+    if (left < 5) {
+      INF("MPEG2 picture header across TS Packet\n");
+
+      /* MEPG2 picture header across TS packet, should cache the left data */
+      memcpy(&stream->PES.data[0], haystack + i, left);
+      stream->PES.len = left;
+      return;
+    }
+
     if (*(uint32_t *)(haystack + i) == needle) {
       // picture header found
-      if (left < 5) {
-        INF("MPEG2 picture header across TS Packet\n");
+      int frame_type = (haystack[i + 5] >> 3) & 0x7;
+      switch (frame_type) {
+        case 1:
+            INF("I frame found, offset: %lx\n", event.offset);
+            event.type = TS_INDEXER_EVENT_TYPE_MPEG2_I_FRAME;
+            break;
 
-        /* MEPG2 picture header across TS packet, should cache the left data */
-        memcpy(&stream->PES.data[0], haystack + i, left);
-        stream->PES.len = left;
-        return;
+        case 2:
+            INF("P frame found, offset: %lx\n", event.offset);
+            event.type = TS_INDEXER_EVENT_TYPE_MPEG2_P_FRAME;
+            break;
+
+        case 3:
+            INF("B frame found, offset: %lx\n", event.offset);
+            event.type = TS_INDEXER_EVENT_TYPE_MPEG2_B_FRAME;
+            break;
+
+        default:
+            i += 5;
+            left -= 5;
+            continue;
+
       }
 
-      int frame_type = (haystack[i + 5] >> 3) & 0x7;
-      INF("frame_type: %d\n", frame_type);
-      if (frame_type == 1) {
-        INF("I-frame found, offset: %ld\n", event.offset);
-        if (stream->format == TS_INDEXER_VIDEO_FORMAT_MPEG2) {
-          event.pts = stream->PES.pts;
-          event.type = TS_INDEXER_EVENT_TYPE_VIDEO_I_FRAME;
-          if (indexer->callback) {
-            indexer->callback(indexer, &event);
-          }
-        } else {
-          ERR("%s not MPEG2 video I-frame??\n", __func__);
-        }
+      event.pts = stream->PES.pts;
+      if (indexer->callback) {
+        indexer->callback(indexer, &event);
+      }
+
+      i += 5;
+      left -= 5;
+    } else if (*(uint32_t *)(haystack + i) == needle1) {
+      // sequence header found
+      event.type = TS_INDEXER_EVENT_TYPE_MPEG2_SEQUENCE;
+      event.pts = stream->PES.pts;
+      if (indexer->callback) {
+        indexer->callback(indexer, &event);
       }
 
       i += 5;
@@ -257,7 +296,7 @@
 
   while (1) {
     int left = pes_data_len - (nalu - data);
-    if (left <= 4) {
+    if (left <= 5) {
       memcpy(&stream->PES.data[0], nalu, left);
       stream->PES.len = left;
       break;
@@ -269,16 +308,33 @@
 
     if (nalu[0] == 0x00 && nalu[1] == 0x00 && nalu[2] == 0x01) {
       int nal_type = nalu[3] & 0x1f;
-      if (nal_type == NAL_TYPE_IDR || nal_type == NAL_TYPE_NON_IDR) {
-        int nal_ref_idr = nalu[3] & 0x60;
-        if (nal_ref_idr == 0x60) { //I-frame
-          INF("H264 I-frame found\n");
-          event.pts = stream->PES.pts;
-          event.type = TS_INDEXER_EVENT_TYPE_VIDEO_I_FRAME;
-          if (indexer->callback) {
-            indexer->callback(indexer, &event);
-          }
+
+      event.pts = stream->PES.pts;
+      if (nal_type == NAL_TYPE_IDR) {
+        INF("AVC IDR found\n");
+        event.type = TS_INDEXER_EVENT_TYPE_AVC_I_SLICE;
+      } else if (nal_type == NAL_TYPE_NON_IDR) {
+        int slice_type = (nalu[4] >> 4) & 0x3;
+        if (slice_type == 0 || slice_type == 5) {
+            event.type = TS_INDEXER_EVENT_TYPE_AVC_P_SLICE;
+        } else if (slice_type == 1 || slice_type == 6) {
+            event.type = TS_INDEXER_EVENT_TYPE_AVC_B_SLICE;
+        } else if (slice_type == 2 || slice_type == 7) {
+            event.type = TS_INDEXER_EVENT_TYPE_AVC_I_SLICE;
+        } else if (slice_type == 3 || slice_type == 8) {
+            event.type = TS_INDEXER_EVENT_TYPE_AVC_SP_SLICE;
+        } else if (slice_type == 4 || slice_type == 9) {
+            event.type = TS_INDEXER_EVENT_TYPE_AVC_SI_SLICE;
+        } else {
+            ERR("%s invalid slice_type: %d\n", __func__, slice_type);
         }
+      } else {
+        nalu += nalu_len;
+        continue;
+      }
+
+      if (indexer->callback) {
+        indexer->callback(indexer, &event);
       }
     }
 
@@ -314,20 +370,50 @@
     if (nalu[0] == 0x00 && nalu[1] == 0x00 && nalu[2] == 0x01) {
       int nalu_type = (nalu[3] & 0x7E) >> 1;
       //INF("nalu[3]: %#x, nalu_type: %#x\n", nalu[3], nalu_type);
-      if (nalu_type == 19) {
-        event.pts = stream->PES.pts;
-        event.type = TS_INDEXER_EVENT_TYPE_VIDEO_I_FRAME;
-        INF("HEVC I-frame found\n");
-        if (indexer->callback) {
-          indexer->callback(indexer, &event);
-        }
-      } else {
-        //INF("%s, not HEVC video I-frame. nalu_type: %#x\n", __func__, nalu_type);
+      switch (nalu_type) {
+        case HEVC_NALU_BLA_W_LP:
+            event.type = TS_INDEXER_EVENT_TYPE_HEVC_BLA_W_LP;
+            break;
+
+        case HEVC_NALU_BLA_W_RADL:
+            event.type = TS_INDEXER_EVENT_TYPE_HEVC_BLA_W_RADL;
+            break;
+
+        case HEVC_NALU_IDR_W_RADL:
+            event.type = TS_INDEXER_EVENT_TYPE_HEVC_IDR_W_RADL;
+            //INF("HEVC I-frame found\n");
+            break;
+
+        case HEVC_NALU_IDR_N_LP:
+            event.type = TS_INDEXER_EVENT_TYPE_HEVC_IDR_N_LP;
+            break;
+
+        case HEVC_NALU_TRAIL_CRA:
+            event.type = TS_INDEXER_EVENT_TYPE_HEVC_TRAIL_CRA;
+            break;
+
+        case HEVC_NALU_SPS:
+            event.type = TS_INDEXER_EVENT_TYPE_HEVC_SPS;
+            break;
+
+        case HEVC_NALU_AUD:
+            event.type = TS_INDEXER_EVENT_TYPE_HEVC_AUD;
+            break;
+
+        default:
+            nalu += nalu_len;
+            continue;
+      }
+
+      event.pts = stream->PES.pts;
+      if (indexer->callback) {
+        indexer->callback(indexer, &event);
       }
     }
 
     nalu += nalu_len;
   }
+
   stream->PES.len = 0;
 }
 
@@ -415,10 +501,15 @@
         pi->callback(pi, &event);
       }
     }
-    stream->PES.state = TS_INDEXER_STATE_PES_PTS;
+    if (stream->format != -1) {
+      stream->PES.state = TS_INDEXER_STATE_PES_PTS;
 
-    p += header_length;
-    left -= header_length;
+      p += header_length;
+      left -= header_length;
+    } else {
+      stream->PES.state = TS_INDEXER_STATE_INIT;
+      left = 0;
+    }
   }
 
   stream->PES.len = left;
@@ -442,6 +533,8 @@
       break;
 
     default:
+      stream->PES.state = TS_INDEXER_STATE_INIT;
+      stream->PES.len = 0;
       break;
   }
 }
@@ -456,6 +549,7 @@
   TS_Indexer_t *pi = ts_indexer;
   int len;
   int is_start;
+  TS_Indexer_Event_t event;
 
   is_start = p[1] & 0x40;
   pid = ((p[1] & 0x1f) << 8) | p[2];
@@ -468,6 +562,14 @@
   }
 
   if (is_start) {
+    memset(&event, 0, sizeof(event));
+    event.pid = pid;
+    event.offset = pi->offset;
+    event.type = TS_INDEXER_EVENT_TYPE_START_INDICATOR;
+    if (pi->callback) {
+      pi->callback(pi, &event);
+    }
+
     if (pid == pi->video_parser.pid) {
       pi->video_parser.offset = pi->offset;
       pi->video_parser.PES.state = TS_INDEXER_STATE_TS_START;
@@ -485,6 +587,15 @@
 
   if (afc & 2) {
     int adp_field_len = p[0];
+    if (p[1] & 0x80) {
+      memset(&event, 0, sizeof(event));
+      event.pid = pid;
+      event.offset = pi->offset;
+      event.type = TS_INDEXER_EVENT_TYPE_DISCONTINUITY_INDICATOR;
+      if (pi->callback) {
+        pi->callback(pi, &event);
+      }
+    }
     p++;
     len--;
 
diff --git a/test/ts_indexer_test/ts_indexer_test.c b/test/ts_indexer_test/ts_indexer_test.c
index 68602e6..c1868b4 100644
--- a/test/ts_indexer_test/ts_indexer_test.c
+++ b/test/ts_indexer_test/ts_indexer_test.c
@@ -33,23 +33,45 @@
 #define INF(fmt, ...) fprintf(stdout, fmt, ##__VA_ARGS__)
 #define ERR(fmt, ...) fprintf(stderr, fmt, ##__VA_ARGS__)
 
+typedef struct
+{
+  uint8_t *ptr;
+  uint64_t last_pts;
+  uint64_t last_offset;
+  FILE *dump_file;
+} ts_indexer_ctl;
+
+static ts_indexer_ctl gControl;
+
 static void ts_indexer_event_cb(TS_Indexer_t *ts_indexer, TS_Indexer_Event_t *event)
 {
+  int write_len;
+
   if (ts_indexer == NULL || event == NULL)
     return;
 
   switch (event->type) {
-    case TS_INDEXER_EVENT_TYPE_VIDEO_I_FRAME:
+    case TS_INDEXER_EVENT_TYPE_MPEG2_I_FRAME:
+      #if 1
       INF("pid: %#x ", event->pid);
-      INF("TS_INDEXER_EVENT_TYPE_VIDEO_I_FRAME, offset: %lx, pts: %lx\n",
+      INF("TS_INDEXER_EVENT_TYPE_I_FRAME, offset: %lx, pts: %lx\n",
             event->offset, event->pts);
+      #endif
       break;
 
     case TS_INDEXER_EVENT_TYPE_VIDEO_PTS:
       #if 0
       INF("PID: %#x ", event->pid);
-      INF("TS_INDEXER_EVENT_TYPE_VIDEO_PTS, Offset: %lx, Pts: %lx\n",
-            event->offset, event->pts);
+      INF("TS_INDEXER_EVENT_TYPE_VIDEO_PTS, Offset: %lx, Lastoffset: %lx, Pts: %lx\n",
+            event->offset, gControl.last_offset, event->pts);
+      write_len = ts_indexer->offset - gControl.last_offset;
+      if (gControl.dump_file) {
+        INF("ptr: %p, write_len: %#x\n", gControl.ptr, write_len);
+        fwrite(gControl.ptr, 1, write_len, gControl.dump_file);
+        gControl.ptr += write_len;
+        INF("ptr: %p\n", gControl.ptr);
+      }
+      gControl.last_offset = ts_indexer->offset;
       #endif
       break;
 
@@ -61,6 +83,26 @@
       #endif
       break;
 
+    case TS_INDEXER_EVENT_TYPE_START_INDICATOR:
+    case TS_INDEXER_EVENT_TYPE_DISCONTINUITY_INDICATOR:
+    case TS_INDEXER_EVENT_TYPE_MPEG2_P_FRAME:
+    case TS_INDEXER_EVENT_TYPE_MPEG2_B_FRAME:
+    case TS_INDEXER_EVENT_TYPE_MPEG2_SEQUENCE:
+    case TS_INDEXER_EVENT_TYPE_AVC_I_SLICE:
+    case TS_INDEXER_EVENT_TYPE_AVC_P_SLICE:
+    case TS_INDEXER_EVENT_TYPE_AVC_B_SLICE:
+    case TS_INDEXER_EVENT_TYPE_AVC_SI_SLICE:
+    case TS_INDEXER_EVENT_TYPE_AVC_SP_SLICE:
+    case TS_INDEXER_EVENT_TYPE_HEVC_SPS:
+    case TS_INDEXER_EVENT_TYPE_HEVC_AUD:
+    case TS_INDEXER_EVENT_TYPE_HEVC_BLA_W_LP:
+    case TS_INDEXER_EVENT_TYPE_HEVC_BLA_W_RADL:
+    case TS_INDEXER_EVENT_TYPE_HEVC_IDR_W_RADL:
+    case TS_INDEXER_EVENT_TYPE_HEVC_IDR_N_LP:
+    case TS_INDEXER_EVENT_TYPE_HEVC_TRAIL_CRA:
+      INF("type: %d, offset: %lx\n", event->type, event->offset);
+      break;
+
     default:
       break;
   }
@@ -106,8 +148,7 @@
   }
 
   if (argc == 1 ||
-    (vpid == 0x1fff && apid == 0x1fff) ||
-    (vpid != 0x1fff && vfmt == -1))
+    (vpid == 0x1fff && apid == 0x1fff))
   {
     usage(argc, argv);
     exit(0);
@@ -134,13 +175,31 @@
 
   size_t len = 0;
   INF("vpid: %#x, vfmt: %d, apid:%#x\n", vpid, vfmt, apid);
+  memset(&gControl, 0, sizeof(ts_indexer_ctl));
+  gControl.last_pts = -1;
+  gControl.last_offset = 0;
+  gControl.dump_file = fopen("./dump.ts", "wb+");
+  if (gControl.dump_file == NULL) {
+    ERR("open dump file failed\n");
+    return -1;
+  }
+
+  INF("ptr: %p ~ %p\n", &data[0], &data[0] + block_size);
   while (1) {
     len = fread(data, 1, block_size, f);
     if (len <= 0)
       break;
 
+    gControl.ptr = &data[0];
     ts_indexer_parse(&ts_indexer, data, len);
+    if (gControl.ptr - &data[0] < len) {
+      int left = len - (gControl.ptr - &data[0]);
+      fwrite(gControl.ptr, 1, left, gControl.dump_file);
+      gControl.last_offset += left;
+    }
   }
+  fflush(gControl.dump_file);
+  fclose(gControl.dump_file);
 
   return 0;
 }