codec/h264/h264dec: added subMbTypeBinarization and TestSubMbTypeBinarization

codec/h264/h264dec/cabac.go

@@ -280,7 +280,7 @@ var (
 // Errors used by mbTypeBinarization.
 var (
 	errBadMbType      = errors.New("macroblock type outside of valid range")
-	errInvalidSliceType = errors.New("slice type does not exist")
+	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
 	}
 }
 

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

codec/h264/h264dec/mbtype.go

@@ -34,6 +34,10 @@ const (
 	maxBMbType        = 48
 	minSIMbType       = 1
 	maxSIMbType       = 26
+	minPOrSPSubMbType = 0
+	maxPOrSPSubMbType = 3
+	minBSubMbType     = 0
+	maxBSubMbType     = 12
 )
 
 const MB_TYPE_INFERRED = 1000