codec/h264/h264dec: wrote TestDecRefPicMarking and fixed bugs

codec/h264/h264dec/slice.go

@@ -258,6 +258,10 @@ type DecRefPicMarking struct {
 	NoOutputOfPriorPicsFlag          bool
 	LongTermReferenceFlag            bool
 	AdaptiveRefPicMarkingModeFlag    bool
+	elements []drpmElement
+}
+
+type drpmElement struct {
 	MemoryManagementControlOperation int
 	DifferenceOfPicNumsMinus1        int
 	LongTermPicNum                   int
@@ -268,7 +272,7 @@ type DecRefPicMarking struct {
 // NewDecRefPicMarking parses elements of a dec_ref_pic_marking following the
 // syntax structure defined in section 7.3.3.3, and returns as a new
 // DecRefPicMarking.
-func NewDecRefPicMarking(br *bits.BitReader, idrPic bool, h *SliceHeader) (*DecRefPicMarking, error) {
+func NewDecRefPicMarking(br *bits.BitReader, idrPic bool) (*DecRefPicMarking, error) {
 	d := &DecRefPicMarking{}
 	if idrPic {
 		b, err := br.ReadBits(1)
@@ -290,35 +294,42 @@ func NewDecRefPicMarking(br *bits.BitReader, idrPic bool, h *SliceHeader) (*DecR
 		d.AdaptiveRefPicMarkingModeFlag = b == 1
 
 		if d.AdaptiveRefPicMarkingModeFlag {
-			d.MemoryManagementControlOperation, err = readUe(br)
+			for i := 0; ; i++ {
+				d.elements = append(d.elements,drpmElement{})
+
+				d.elements[i].MemoryManagementControlOperation, err = readUe(br)
 				if err != nil {
 					return nil, errors.Wrap(err, "could not parse MemoryManagementControlOperation")
 				}
-			for d.MemoryManagementControlOperation != 0 {
-				if d.MemoryManagementControlOperation == 1 || d.MemoryManagementControlOperation == 3 {
-					d.DifferenceOfPicNumsMinus1, err = readUe(br)
+
+				if d.elements[i].MemoryManagementControlOperation == 1 || d.elements[i].MemoryManagementControlOperation == 3 {
+					d.elements[i].DifferenceOfPicNumsMinus1, err = readUe(br)
 					if err != nil {
 						return nil, errors.Wrap(err, "could not parse MemoryManagementControlOperation")
 					}
 				}
-				if d.MemoryManagementControlOperation == 2 {
-					d.LongTermPicNum, err = readUe(br)
+				if d.elements[i].MemoryManagementControlOperation == 2 {
+					d.elements[i].LongTermPicNum, err = readUe(br)
 					if err != nil {
 						return nil, errors.Wrap(err, "could not parse LongTermPicNum")
 					}
 				}
-				if d.MemoryManagementControlOperation == 3 || d.MemoryManagementControlOperation == 6 {
-					d.LongTermFrameIdx, err = readUe(br)
+				if d.elements[i].MemoryManagementControlOperation == 3 || d.elements[i].MemoryManagementControlOperation == 6 {
+					d.elements[i].LongTermFrameIdx, err = readUe(br)
 					if err != nil {
 						return nil, errors.Wrap(err, "could not parse LongTermFrameIdx")
 					}
 				}
-				if d.MemoryManagementControlOperation == 4 {
-					d.MaxLongTermFrameIdxPlus1, err = readUe(br)
+				if d.elements[i].MemoryManagementControlOperation == 4 {
+					d.elements[i].MaxLongTermFrameIdxPlus1, err = readUe(br)
 					if err != nil {
 						return nil, errors.Wrap(err, "could not parse MaxLongTermFrameIdxPlus1")
 					}
 				}
+
+				if d.elements[i].MemoryManagementControlOperation == 0 {
+					break
+				}
 			}
 		}
 	}
@@ -1358,7 +1369,7 @@ func NewSliceContext(vid *VideoStream, nalUnit *NALUnit, rbsp []byte, showPacket
 	}
 	if nalUnit.RefIdc != 0 {
 		// devRefPicMarking()
-		header.DecRefPicMarking, err = NewDecRefPicMarking(br, idrPic, &header)
+		header.DecRefPicMarking, err = NewDecRefPicMarking(br, idrPic)
 		if err != nil {
 			return nil, errors.Wrap(err, "could not parse DecRefPicMarking")
 		}

codec/h264/h264dec/slice_test.go

@@ -236,3 +236,83 @@ func TestNewPredWeightTable(t *testing.T) {
 		}
 	}
 }
+
+func TestDecRefPicMarking(t *testing.T) {
+	tests := []struct {
+		in     string
+		idrPic bool
+		want   DecRefPicMarking
+	}{
+		{
+			in: "0" + // u(1) no_output_of_prior_pics_flag = false
+				"1", // u(1) long_term_reference_flag = true
+			idrPic: true,
+			want: DecRefPicMarking{
+				NoOutputOfPriorPicsFlag: false,
+				LongTermReferenceFlag:   true,
+			},
+		},
+		{
+			in: "1" + // u(1) adaptive_ref_pic_marking_mode_flag = true
+
+				"010" + // ue(v) memory_management_control_operation = 1
+				"011" + // ue(v) difference_of_pic_nums_minus1 = 2
+
+				"00100" + // ue(v) memory_management_control_operation = 3
+				"1" + // ue(v) difference_of_pic_nums_minus1 = 0
+				"011" + // ue(v) long_term_frame_idx = 2
+
+				"011" + // ue(v) memory_management_control_operation = 2
+				"00100" + // ue(v) long_term_pic_num = 3
+
+				"00101" + // ue(v) memory_management_control_operation = 4
+				"010" + // ue(v) max_long_term_frame_idx_plus1 = 1
+
+				"1", // ue(v) memory_management_control_operation = 0
+
+			idrPic: false,
+
+			want: DecRefPicMarking{
+				AdaptiveRefPicMarkingModeFlag: true,
+				elements: []drpmElement{
+					{
+						MemoryManagementControlOperation: 1,
+						DifferenceOfPicNumsMinus1:        2,
+					},
+					{
+						MemoryManagementControlOperation: 3,
+						DifferenceOfPicNumsMinus1:        0,
+						LongTermFrameIdx:                 2,
+					},
+					{
+						MemoryManagementControlOperation: 2,
+						LongTermPicNum:                   3,
+					},
+					{
+						MemoryManagementControlOperation: 4,
+						MaxLongTermFrameIdxPlus1:         1,
+					},
+					{
+						MemoryManagementControlOperation: 0,
+					},
+				},
+			},
+		},
+	}
+
+	for i, test := range tests {
+		inBytes, err := binToSlice(test.in)
+		if err != nil {
+			t.Fatalf("unexpected error %v for binToSlice in test %d", err, i)
+		}
+
+		got, err := NewDecRefPicMarking(bits.NewBitReader(bytes.NewReader(inBytes)), test.idrPic)
+		if err != nil {
+			t.Fatalf("unexpected error %v for NewPredWeightTable in test %d", err, i)
+		}
+
+		if !reflect.DeepEqual(*got, test.want) {
+			t.Errorf("did not get expected result for test %d\nGot: %v\nWant: %v\n", i, got, test.want)
+		}
+	}
+}