From 09865015c9880271e21fca8d0161fb1bf4b6d939 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Wed, 22 Jun 2022 13:58:49 +0200
Subject: [PATCH] h264: move SPS bitstream restriction into dedicated struct

---
 pkg/h264/dtsextractor.go |  11 +++--
 pkg/h264/sps.go          | 103 ++++++++++++++++++++++-----------------
 pkg/h264/sps_test.go     |  37 +++++++-------
 3 files changed, 86 insertions(+), 65 deletions(-)

diff --git a/pkg/h264/dtsextractor.go b/pkg/h264/dtsextractor.go
index af3c841d..e3908789 100644
--- a/pkg/h264/dtsextractor.go
+++ b/pkg/h264/dtsextractor.go
@@ -131,9 +131,11 @@ func (d *DTSExtractor) extractInner(
 				d.spsp = &spsp
 
 				// in case of B-frames, we have to subtract from DTS the maximum number of reordered frames
-				if d.spsp.VUI != nil && d.spsp.VUI.TimingInfo != nil {
-					d.ptsDTSOffset = time.Duration(math.Round(float64(time.Duration(d.spsp.VUI.MaxNumReorderFrames)*time.Second*
-						time.Duration(d.spsp.VUI.TimingInfo.NumUnitsInTick)*2) / float64(d.spsp.VUI.TimingInfo.TimeScale)))
+				if d.spsp.VUI != nil && d.spsp.VUI.TimingInfo != nil &&
+					d.spsp.VUI.BitstreamRestriction != nil {
+					d.ptsDTSOffset = time.Duration(math.Round(float64(
+						time.Duration(d.spsp.VUI.BitstreamRestriction.MaxNumReorderFrames)*time.Second*
+							time.Duration(d.spsp.VUI.TimingInfo.NumUnitsInTick)*2) / float64(d.spsp.VUI.TimingInfo.TimeScale)))
 				} else {
 					d.ptsDTSOffset = 0
 				}
@@ -170,7 +172,8 @@ func (d *DTSExtractor) extractInner(
 	}
 
 	// special case to eliminate errors near 0
-	if d.spsp.VUI != nil && d.spsp.VUI.TimingInfo != nil && pocDiff == -int32(d.spsp.VUI.MaxNumReorderFrames)*2 {
+	if d.spsp.VUI != nil && d.spsp.VUI.TimingInfo != nil && d.spsp.VUI.BitstreamRestriction != nil &&
+		pocDiff == -int32(d.spsp.VUI.BitstreamRestriction.MaxNumReorderFrames)*2 {
 		return pts, pocDiff, nil
 	}
 
diff --git a/pkg/h264/sps.go b/pkg/h264/sps.go
index 5588b332..b81b522e 100644
--- a/pkg/h264/sps.go
+++ b/pkg/h264/sps.go
@@ -221,6 +221,57 @@ func (t *SPS_TimingInfo) unmarshal(br *bitio.Reader) error {
 	return nil
 }
 
+// SPS_BitstreamRestriction are bitstream restriction infos.
+type SPS_BitstreamRestriction struct { //nolint:revive
+	MotionVectorsOverPicBoundariesFlag bool
+	MaxBytesPerPicDenom                uint32
+	MaxBitsPerMbDenom                  uint32
+	Log2MaxMvLengthHorizontal          uint32
+	Log2MaxMvLengthVertical            uint32
+	MaxNumReorderFrames                uint32
+	MaxDecFrameBuffering               uint32
+}
+
+func (r *SPS_BitstreamRestriction) unmarshal(br *bitio.Reader) error {
+	var err error
+	r.MotionVectorsOverPicBoundariesFlag, err = readFlag(br)
+	if err != nil {
+		return err
+	}
+
+	r.MaxBytesPerPicDenom, err = readGolombUnsigned(br)
+	if err != nil {
+		return err
+	}
+
+	r.MaxBitsPerMbDenom, err = readGolombUnsigned(br)
+	if err != nil {
+		return err
+	}
+
+	r.Log2MaxMvLengthHorizontal, err = readGolombUnsigned(br)
+	if err != nil {
+		return err
+	}
+
+	r.Log2MaxMvLengthVertical, err = readGolombUnsigned(br)
+	if err != nil {
+		return err
+	}
+
+	r.MaxNumReorderFrames, err = readGolombUnsigned(br)
+	if err != nil {
+		return err
+	}
+
+	r.MaxDecFrameBuffering, err = readGolombUnsigned(br)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // SPS_VUI is a video usability information.
 type SPS_VUI struct { //nolint:revive
 	AspectRatioInfoPresentFlag bool
@@ -256,18 +307,11 @@ type SPS_VUI struct { //nolint:revive
 	// vclHrdParametersPresentFlag == true
 	VclHRD *SPS_HRD
 
-	LowDelayHrdFlag          bool
-	PicStructPresentFlag     bool
-	BitstreamRestrictionFlag bool
+	LowDelayHrdFlag      bool
+	PicStructPresentFlag bool
 
-	// BitstreamRestrictionFlag == true
-	MotionVectorsOverPicBoundariesFlag bool
-	MaxBytesPerPicDenom                uint32
-	MaxBitsPerMbDenom                  uint32
-	Log2MaxMvLengthHorizontal          uint32
-	Log2MaxMvLengthVertical            uint32
-	MaxNumReorderFrames                uint32
-	MaxDecFrameBuffering               uint32
+	// bitstreamRestrictionFlag == true
+	BitstreamRestriction *SPS_BitstreamRestriction
 }
 
 func (v *SPS_VUI) unmarshal(br *bitio.Reader) error {
@@ -416,43 +460,14 @@ func (v *SPS_VUI) unmarshal(br *bitio.Reader) error {
 		return err
 	}
 
-	v.BitstreamRestrictionFlag, err = readFlag(br)
+	bitstreamRestrictionFlag, err := readFlag(br)
 	if err != nil {
 		return err
 	}
 
-	if v.BitstreamRestrictionFlag {
-		v.MotionVectorsOverPicBoundariesFlag, err = readFlag(br)
-		if err != nil {
-			return err
-		}
-
-		v.MaxBytesPerPicDenom, err = readGolombUnsigned(br)
-		if err != nil {
-			return err
-		}
-
-		v.MaxBitsPerMbDenom, err = readGolombUnsigned(br)
-		if err != nil {
-			return err
-		}
-
-		v.Log2MaxMvLengthHorizontal, err = readGolombUnsigned(br)
-		if err != nil {
-			return err
-		}
-
-		v.Log2MaxMvLengthVertical, err = readGolombUnsigned(br)
-		if err != nil {
-			return err
-		}
-
-		v.MaxNumReorderFrames, err = readGolombUnsigned(br)
-		if err != nil {
-			return err
-		}
-
-		v.MaxDecFrameBuffering, err = readGolombUnsigned(br)
+	if bitstreamRestrictionFlag {
+		v.BitstreamRestriction = &SPS_BitstreamRestriction{}
+		err := v.BitstreamRestriction.unmarshal(br)
 		if err != nil {
 			return err
 		}
diff --git a/pkg/h264/sps_test.go b/pkg/h264/sps_test.go
index 60b4a338..3722f341 100644
--- a/pkg/h264/sps_test.go
+++ b/pkg/h264/sps_test.go
@@ -74,12 +74,13 @@ func TestSPSUnmarshal(t *testing.T) {
 						NumUnitsInTick: 1,
 						TimeScale:      60,
 					},
-					BitstreamRestrictionFlag:           true,
-					MotionVectorsOverPicBoundariesFlag: true,
-					Log2MaxMvLengthHorizontal:          11,
-					Log2MaxMvLengthVertical:            11,
-					MaxNumReorderFrames:                2,
-					MaxDecFrameBuffering:               4,
+					BitstreamRestriction: &SPS_BitstreamRestriction{
+						MotionVectorsOverPicBoundariesFlag: true,
+						Log2MaxMvLengthHorizontal:          11,
+						Log2MaxMvLengthVertical:            11,
+						MaxNumReorderFrames:                2,
+						MaxDecFrameBuffering:               4,
+					},
 				},
 			},
 			1280,
@@ -112,11 +113,12 @@ func TestSPSUnmarshal(t *testing.T) {
 						NumUnitsInTick: 1,
 						TimeScale:      60,
 					},
-					BitstreamRestrictionFlag:           true,
-					MotionVectorsOverPicBoundariesFlag: true,
-					Log2MaxMvLengthHorizontal:          11,
-					Log2MaxMvLengthVertical:            11,
-					MaxDecFrameBuffering:               3,
+					BitstreamRestriction: &SPS_BitstreamRestriction{
+						MotionVectorsOverPicBoundariesFlag: true,
+						Log2MaxMvLengthHorizontal:          11,
+						Log2MaxMvLengthVertical:            11,
+						MaxDecFrameBuffering:               3,
+					},
 				},
 			},
 			1920,
@@ -149,12 +151,13 @@ func TestSPSUnmarshal(t *testing.T) {
 						NumUnitsInTick: 1,
 						TimeScale:      60,
 					},
-					BitstreamRestrictionFlag:           true,
-					MotionVectorsOverPicBoundariesFlag: true,
-					Log2MaxMvLengthHorizontal:          11,
-					Log2MaxMvLengthVertical:            11,
-					MaxNumReorderFrames:                2,
-					MaxDecFrameBuffering:               4,
+					BitstreamRestriction: &SPS_BitstreamRestriction{
+						MotionVectorsOverPicBoundariesFlag: true,
+						Log2MaxMvLengthHorizontal:          11,
+						Log2MaxMvLengthVertical:            11,
+						MaxNumReorderFrames:                2,
+						MaxDecFrameBuffering:               4,
+					},
 				},
 			},
 			1920,