From 6d06fab39f6d27366bc3eb7f60322574e01e63d4 Mon Sep 17 00:00:00 2001
From: Saxon <saxon@ausocean.org>
Date: Thu, 15 Aug 2019 00:30:17 +0930
Subject: [PATCH] codec/h264/h264dec: added subMbTypeBinarization and
 TestSubMbTypeBinarization

---
 codec/h264/h264dec/cabac.go      | 31 ++++++++++++++++++++++++++++---
 codec/h264/h264dec/cabac_test.go | 27 +++++++++++++++++++++++++++
 codec/h264/h264dec/mbtype.go     | 20 ++++++++++++--------
 3 files changed, 67 insertions(+), 11 deletions(-)

diff --git a/codec/h264/h264dec/cabac.go b/codec/h264/h264dec/cabac.go
index 12a07ddc..21ae8316 100644
--- a/codec/h264/h264dec/cabac.go
+++ b/codec/h264/h264dec/cabac.go
@@ -279,8 +279,8 @@ var (
 
 // Errors used by mbTypeBinarization.
 var (
-	errBadMbType        = errors.New("macroblock type outside of valid range")
-	errInvalidSliceType = errors.New("slice type does not exist")
+	errBadMbType      = errors.New("macroblock type outside of valid range")
+	errBadMbSliceType = errors.New("bad slice type for macroblock")
 )
 
 // mbTypeBinarization returns the macroblock type binarization for the given
@@ -322,7 +322,32 @@ func mbTypeBinarization(v, slice int) ([]int, error) {
 		return append([]int{1, 1, 1, 1, 0, 1}, binOfIMBTypes[v-23]...), nil
 
 	default:
-		return nil, errInvalidSliceType
+		return nil, errBadMbSliceType
+	}
+}
+
+// Error used by subMbTypeBinarization.
+var errBadSubMbSliceType = errors.New("bad slice type for sub-macroblock")
+
+// subMbTypeBinarization returns the binarization of a sub-macroblock type
+// given the slice in which it is in using the process defined in section
+// 9.3.2.5 of the specifications.
+func subMbTypeBinarization(v, slice int) ([]int, error) {
+	switch slice {
+	case sliceTypeP, sliceTypeSP:
+		if v < minPOrSPSubMbType || v > maxPOrSPSubMbType {
+			return nil, errBadMbType
+		}
+		return binOfPOrSPSubMBTypes[v], nil
+
+	case sliceTypeB:
+		if v < minBSubMbType || v > maxBSubMbType {
+			return nil, errBadMbType
+		}
+		return binOfBSubMBTypes[v], nil
+
+	default:
+		return nil, errBadSubMbSliceType
 	}
 }
 
diff --git a/codec/h264/h264dec/cabac_test.go b/codec/h264/h264dec/cabac_test.go
index 26350f18..0d8adaa4 100644
--- a/codec/h264/h264dec/cabac_test.go
+++ b/codec/h264/h264dec/cabac_test.go
@@ -28,6 +28,7 @@ func TestMbTypeBinarization(t *testing.T) {
 		{v: 30, slice: sliceTypeB, want: []int{1, 1, 1, 1, 0, 1, 1, 0, 0, 1, 0, 1, 0}},
 		{v: -1, slice: sliceTypeB, err: errBadMbType},
 		{v: 49, slice: sliceTypeB, err: errBadMbType},
+		{v: 6, slice: 20, err: errBadMbSliceType},
 	}
 
 	for i, test := range tests {
@@ -42,6 +43,32 @@ func TestMbTypeBinarization(t *testing.T) {
 	}
 }
 
+func TestSubMbTypeBinarization(t *testing.T) {
+	tests := []struct {
+		v, slice int
+		want     []int
+		err      error
+	}{
+		{v: 2, slice: sliceTypeP, want: []int{0, 1, 1}},
+		{v: 2, slice: sliceTypeSP, want: []int{0, 1, 1}},
+		{v: -1, slice: sliceTypeSP, err: errBadMbType},
+		{v: 4, slice: sliceTypeSP, err: errBadMbType},
+		{v: 9, slice: sliceTypeB, want: []int{1, 1, 1, 0, 1, 0}},
+		{v: 9, slice: 40, err: errBadSubMbSliceType},
+	}
+
+	for i, test := range tests {
+		got, err := subMbTypeBinarization(test.v, test.slice)
+		if err != test.err {
+			t.Errorf("did not get expected error for test %d\nGot: %v\nWant: %v", i, err, test.err)
+		}
+
+		if !reflect.DeepEqual(got, test.want) {
+			t.Errorf("did not get expected result for test %d\nGot: %v\nWant: %v", i, got, test.want)
+		}
+	}
+}
+
 var ctxIdxTests = []struct {
 	binIdx       int
 	maxBinIdxCtx int
diff --git a/codec/h264/h264dec/mbtype.go b/codec/h264/h264dec/mbtype.go
index 02466d45..eb716c17 100644
--- a/codec/h264/h264dec/mbtype.go
+++ b/codec/h264/h264dec/mbtype.go
@@ -26,14 +26,18 @@ const (
 
 // Min and max macroblock types for slice types.
 const (
-	minIMbType     = 0
-	maxIMbType     = 25
-	minPOrSPMbType = 0
-	maxPOrSPMbType = 30
-	minBMbType     = 0
-	maxBMbType     = 48
-	minSIMbType    = 1
-	maxSIMbType    = 26
+	minIMbType        = 0
+	maxIMbType        = 25
+	minPOrSPMbType    = 0
+	maxPOrSPMbType    = 30
+	minBMbType        = 0
+	maxBMbType        = 48
+	minSIMbType       = 1
+	maxSIMbType       = 26
+	minPOrSPSubMbType = 0
+	maxPOrSPSubMbType = 3
+	minBSubMbType     = 0
+	maxBSubMbType     = 12
 )
 
 const MB_TYPE_INFERRED = 1000