From 1f3ab6e473d6f2084fd1a94eb8408e9351e9b978 Mon Sep 17 00:00:00 2001
From: Thulinma <jaron@vietors.com>
Date: Mon, 15 May 2017 17:11:24 +0200
Subject: [PATCH] Improved NAL unit and keyframe handling

---
 lib/dtsc.h       |  2 +-
 lib/dtscmeta.cpp | 19 +++++++++++--------
 lib/nal.cpp      | 37 ++++++++-----------------------------
 lib/nal.h        |  2 +-
 4 files changed, 21 insertions(+), 39 deletions(-)

diff --git a/lib/dtsc.h b/lib/dtsc.h
index 71c1edd7..11bdd91f 100644
--- a/lib/dtsc.h
+++ b/lib/dtsc.h
@@ -127,7 +127,7 @@ namespace DTSC {
       bool getFlag(const char * identifier) const;
       bool hasMember(const char * identifier) const;
       void appendNal(const char * appendData, uint32_t appendLen, uint32_t totalLen);
-      void clearKeyFrame();
+      void setKeyFrame(bool kf);
       long long unsigned int getTime() const;
       long int getTrackId() const;
       char * getData() const;
diff --git a/lib/dtscmeta.cpp b/lib/dtscmeta.cpp
index 2bb6d2fc..b0191fbd 100644
--- a/lib/dtscmeta.cpp
+++ b/lib/dtscmeta.cpp
@@ -271,21 +271,23 @@ namespace DTSC {
     memcpy(data+offset+11+packDataSize, "\000\000\356", 3);
   }
 
-  ///clear the keyframe byte.
-  void Packet::clearKeyFrame(){
+  ///sets the keyframe byte.
+  void Packet::setKeyFrame(bool kf){
     uint32_t offset = 23;
-    while (data[offset] != 'd' && data[offset] != 'k'){
+    while (data[offset] != 'd' && data[offset] != 'k' && data[offset] != 'K'){
       switch (data[offset]){
         case 'o': offset += 17; break;
         case 'b': offset += 15; break;
         default:
-          FAIL_MSG("Errrrrrr");
+          FAIL_MSG("Unknown field: %c", data[offset]);
       }
     }
 
-    if(data[offset] == 'k'){
-      data[offset] = 'K';
-      data[offset+16] = 0;
+    if(data[offset] == 'k' || data[offset] == 'K'){
+      data[offset] = (kf?'k':'K');
+      data[offset+16] = (kf?1:0);
+    }else{
+      ERROR_MSG("Could not set keyframe - field not found!");
     }
   }
 
@@ -327,8 +329,9 @@ namespace DTSC {
         case 'o': offset += 17; break;
         case 'b': offset += 15; break;
         case 'k': offset += 19; break;
+        case 'K': offset += 19; break;
         default:
-          FAIL_MSG("Errrrrrr");
+          FAIL_MSG("Unknown field: %c", data[offset]);
           return -1;
       }
     }
diff --git a/lib/nal.cpp b/lib/nal.cpp
index f4d75d0d..4a502005 100644
--- a/lib/nal.cpp
+++ b/lib/nal.cpp
@@ -69,54 +69,33 @@ namespace nalu {
     return dataSize;
   }
 
+  ///Scans data for the last non-zero byte, returning a pointer to it.
   const char* nalEndPosition(const char * data, uint32_t dataSize){
     while(dataSize > 0 && memcmp(data+dataSize-1, "\000",1) == 0 ){
       dataSize--;
     }
-    
     return data+dataSize;
   }
 
-  ///scan data stream for startcode. return pointer to location when found, NULL otherwise
-  void scanAnnexB(const char * data, uint32_t dataSize, const char *& packetPointer){
+  ///Scan data for Annex B start code. Returns pointer to it when found, null otherwise.
+  const char * scanAnnexB(const char * data, uint32_t dataSize){
     int offset = 0;
-
     while(offset+2 < dataSize){
       const char * begin = data + offset;
-//      int t = ((((int*)begin)[0]) >> 8) & 0x00FFFFFF;
       int t = (int)((begin[0] << 8)|((begin[1]) << 8)|(begin[2])); 
-      //int t = (int)((begin[0]|begin[1]) << 1)|(begin[2]); 
-      //search for startcode
-
-      //if(memcmp(begin, "\000\000\001",3) != 0){
       if(t != 1){
-
-        //if((t & 0x0000FF != 0 )) 
-        if((int)begin[2] != 0 )                                   //XX1
-        {
+        if (begin[2]){//skip three bytes if the last one isn't zero
           offset +=3;
-        }else if(((int)begin[1] == 1) && ((int)begin[2] ==0)){    //X10 
+        }else if (begin[1]){//skip two bytes if the second one isn't zero
           offset +=2; 
-        }else{
-          offset++;   //[X00]? incr with 1 because the startcode could be one at 1byte offset.
-        }
-/*
-        if(t != 0 )
-        {
-          offset += 3;
-        }else{
+        }else{//All other cases, skip one byte
           offset++;
         }
-*/
-
-//        offset++;
       }else{
-        packetPointer = begin;
-        return;
+        return begin;
       }
     }
-
-    packetPointer = NULL;
+    return 0;
   }
 
   unsigned long fromAnnexB(const char * data, unsigned long dataSize, char *& result){
diff --git a/lib/nal.h b/lib/nal.h
index 28a3e246..524b670c 100644
--- a/lib/nal.h
+++ b/lib/nal.h
@@ -16,6 +16,6 @@ namespace nalu {
 
   unsigned long toAnnexB(const char * data, unsigned long dataSize, char *& result);
   unsigned long fromAnnexB(const char * data, unsigned long dataSize, char *& result);
-  void scanAnnexB(const char * data, uint32_t dataSize, const char *& packetPointer);
+  const char* scanAnnexB(const char * data, uint32_t dataSize);
   const char* nalEndPosition(const char * data, uint32_t dataSize);
 }