codec/h264/h264dec: added functions to get TrailingOnes and TotalCoeff as well as testing

The functions firstly derive the value of nC and then read the value of coeff_token from the BitReader. The table read prior and loaded into
a 'map' is then used to get the corresponding values of TrailingOnes and TotalCoef.

codec/h264/h264dec/cavlc.go

@@ -55,6 +55,201 @@ func init() {
 	}
 }
 
+const (
+	chromaDCLevel = iota
+	intra16x16DCLevel
+	intra16x16ACLevel
+	cbIntra16x16DCLevel
+	cbIntra16x16ACLevel
+	crIntra16x16DCLevel
+	crIntra16x16ACLevel
+	lumaLevel4x4
+	cbLevel4x4
+	crLevel4x4
+)
+
+// TODO: put this somewhere more appropriate once context is understood.
+type block struct {
+	usingInterMbPredMode bool
+	mbType               int
+	totalCoef            int
+}
+
+// parseTotalCoeffAndTrailingOnes will use logic provided in section 9.2.1 of
+// the specifications to obtain a value of nC, parse coeff_token from br and
+// then use table 9-5 to find corresponding values of TrailingOnes(coeff_token)
+// and TotalCoeff(coeff_token) which are then subsequently returned.
+func parseTotalCoeffAndTrailingOnes(
+	br *bits.BitReader,
+	ctx *SliceContext,
+	usingMbPredMode bool,
+	level,
+	maxNumCoef,
+	inBlockIdx int,
+) (totalCoeff, trailingOnes, nC, outBlockIdx int, err error) {
+	if level == chromaDCLevel {
+		if ctx.chromaArrayType == 1 {
+			nC = -1
+		} else {
+			nC = -2
+		}
+	} else {
+		// Steps 1,2 and 3.
+		if level == intra16x16DCLevel || level == cbIntra16x16DCLevel || level == crIntra16x16DCLevel {
+			outBlockIdx = 0
+		}
+
+		// Step 4 derive blkA and blkB variables (blockA and blockB here).
+		var mbAddr, blk [2]block
+		switch level {
+		case intra16x16DCLevel, intra16x16ACLevel, lumaLevel4x4:
+			// TODO: clause 6.4.11.4
+			panic("not implemented")
+		case cbIntra16x16DCLevel, cbIntra16x16ACLevel, cbLevel4x4:
+			// TODO: clause 6.4.11.6
+			panic("not implemented")
+		case crIntra16x16DCLevel, crIntra16x16ACLevel, crLevel4x4:
+			// TODO: clause 6.4.11.6
+			panic("not implemented")
+		default:
+			// TODO: clause 6.4.11.5
+			panic("not implemented")
+		}
+
+		var availableFlag [2]bool
+		var n [2]int
+		for i := range availableFlag {
+			// Step 5.
+			if !(!available(mbAddr[i]) || usingMbPredMode || ctx.ConstrainedIntraPred ||
+				mbAddr[i].usingInterMbPredMode || ctx.nalType == 2 || ctx.nalType == 3 || ctx.nalType == 4) {
+				availableFlag[i] = true
+			}
+
+			// Step 6.
+			if availableFlag[i] {
+				switch {
+				case mbAddr[i].mbType == pSkip || mbAddr[i].mbType == bSkip || (mbAddr[i].mbType != iPCM && resTransformCoeffLevelsZero(blk[i])):
+					n[i] = 0
+				case mbAddr[i].mbType == iPCM:
+					n[i] = 16
+				default:
+					// TODO: how is this set ?
+					// "Otherwise, nN is set equal to the value TotalCoeff( coeff_token ) of
+					// the neighbouring block blkN."
+					// Do we need to run this same process for this block, or assume it's
+					// already happened?
+					n[i] = blk[i].totalCoef
+				}
+			}
+		}
+
+		// Step 7 derive
+		switch {
+		case availableFlag[0] && availableFlag[1]:
+			nC = (n[0] + n[1] + 1) >> 1
+		case availableFlag[0]:
+			nC = n[0]
+		case availableFlag[1]:
+			nC = n[1]
+		default:
+			nC = 0
+		}
+	}
+
+	trailingOnes, totalCoeff, _, err = readCoeffToken(br, nC)
+	if err != nil {
+		err = fmt.Errorf("could not get trailingOnes and totalCoeff vars, failed with error: %v", err)
+		return
+	}
+	return
+}
+
+var (
+	errInvalidNC = errors.New("invalid value of nC")
+	errBadToken  = errors.New("could not find coeff_token value in map")
+)
+
+// readCoeffToken will read the coeff_token from br and find a match in the
+// coeff_token mapping table (table 9-5 in the specifications) given also nC.
+// The resultant TrailingOnes(coeff_token) and TotalCoeff(coeff_token) are
+// returned as well as the value of coeff_token.
+func readCoeffToken(br *bits.BitReader, nC int) (trailingOnes, totalCoeff, coeffToken int, err error) {
+	// Get the number of leading zeros.
+	var b uint64
+	nZeros := -1
+	for ; b == 0; nZeros++ {
+		b, err = br.ReadBits(1)
+		if err != nil {
+			err = fmt.Errorf("could not read coeff_token leading zeros, failed with error: %v", err)
+			return
+		}
+	}
+
+	// Get the column idx for the map.
+	var nCIdx int
+	switch {
+	case 0 <= nC && nC < 2:
+		nCIdx = 0
+	case 2 <= nC && nC < 4:
+		nCIdx = 1
+	case 4 <= nC && nC < 8:
+		nCIdx = 2
+	case 8 <= nC:
+		nCIdx = 3
+	case nC == -1:
+		nCIdx = 4
+	case nC == -2:
+		nCIdx = 5
+	default:
+		err = errInvalidNC
+		return
+	}
+
+	// Get the value of coeff_token.
+	val := b
+	nBits := nZeros
+	for {
+		vars, ok := coeffTokenMap[nCIdx][nZeros][int(val)]
+		if ok {
+			trailingOnes = vars[0]
+			totalCoeff = vars[1]
+			coeffToken = int(val)
+			return
+		}
+
+		const maxCoeffTokenBits = 16
+		if !ok && nBits == maxCoeffTokenBits {
+			err = errBadToken
+			return
+		}
+
+		b, err = br.ReadBits(1)
+		if err != nil {
+			err = fmt.Errorf("could not read next bit of coeff_token, failed with error: %v", err)
+			return
+		}
+
+		nBits++
+		val <<= 1
+		val |= b
+	}
+}
+
+// TODO: implement this. From step 6 section 9.2.1. Will return true if "AC
+// residual transform coefficient levels of the neighbouring block blkN are
+// equal to 0 due to the corresponding bit of CodedBlockPatternLuma or
+// CodedBlockPatternChroma being equal to 0"
+func resTransformCoeffLevelsZero(b block) bool {
+	panic("not implemented")
+	return true
+}
+
+// TODO: implement this
+func available(b block) bool {
+	panic("not implemented")
+	return true
+}
+
 // The number of columns in the coeffTokenMap defined below. This is
 // representative of the number of defined nC ranges defined in table 9-5.
 const nColumns = 6

codec/h264/h264dec/cavlc_test.go

@@ -136,3 +136,55 @@ func TestParseLevelPrefix(t *testing.T) {
 		}
 	}
 }
+
+func TestReadCoeffToken(t *testing.T) {
+	tests := []struct {
+		// Input.
+		nC        int
+		tokenBits string
+
+		// Expected.
+		trailingOnes int
+		totalCoeff   int
+		err          error
+	}{
+		{
+			nC:           5,
+			tokenBits:    "0001001",
+			trailingOnes: 0,
+			totalCoeff:   6,
+		},
+		{
+			nC:        -1,
+			tokenBits: "0000000000111111111",
+			err:       errBadToken,
+		},
+		{
+			nC:        -3,
+			tokenBits: "0001001",
+			err:       errInvalidNC,
+		},
+	}
+
+	for i, test := range tests {
+		b, err := binToSlice(test.tokenBits)
+		if err != nil {
+			t.Errorf("converting bin string to slice failed with error: %v for test", err)
+			continue
+		}
+
+		gotTrailingOnes, gotTotalCoeff, _, gotErr := readCoeffToken(bits.NewBitReader(bytes.NewReader(b)), test.nC)
+		if gotErr != test.err {
+			t.Errorf("did not get expected error for test: %d\nGot: %v\nWant: %v\n", i, gotErr, test.err)
+			continue
+		}
+
+		if gotTrailingOnes != test.trailingOnes {
+			t.Errorf("did not get expected TrailingOnes(coeff_token) for test %d\nGot: %v\nWant: %v\n", i, gotTrailingOnes, test.trailingOnes)
+		}
+
+		if gotTotalCoeff != test.totalCoeff {
+			t.Errorf("did not get expected TotalCoeff(coeff_token) for test %d\nGot: %v\nWant: %v\n", i, gotTotalCoeff, test.totalCoeff)
+		}
+	}
+}