From 3d4bcfd4cc334314bdec4361fd882b1adbec551e Mon Sep 17 00:00:00 2001
From: Robin Eklind <rnd0x00@gmail.com>
Date: Mon, 30 Oct 2023 23:33:44 +0100
Subject: [PATCH] flac: add escaped parition encoding

---
 enc_test.go        |  9 +++------
 encode_subframe.go | 23 +++++++++++++++++++++--
 frame/subframe.go  |  3 +++
 3 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/enc_test.go b/enc_test.go
index 640d24a..11b9620 100644
--- a/enc_test.go
+++ b/enc_test.go
@@ -53,8 +53,7 @@ func TestEncode(t *testing.T) {
 		"testdata/flac-test-files/subset/13 - qlp precision 2 bit.flac",
 		"testdata/flac-test-files/subset/14 - wasted bits.flac",
 		"testdata/flac-test-files/subset/15 - only verbatim subframes.flac",
-		// TODO: implement encoding of escaped partitions
-		//"testdata/flac-test-files/subset/16 - partition order 8 containing escaped partitions.flac",
+		"testdata/flac-test-files/subset/16 - partition order 8 containing escaped partitions.flac",
 		"testdata/flac-test-files/subset/17 - all fixed orders.flac",
 		"testdata/flac-test-files/subset/18 - precision search.flac",
 		"testdata/flac-test-files/subset/19 - samplerate 35467Hz.flac",
@@ -73,8 +72,7 @@ func TestEncode(t *testing.T) {
 		"testdata/flac-test-files/subset/29 - high resolution audio, blocksize 16384.flac",
 		"testdata/flac-test-files/subset/30 - high resolution audio, blocksize 13456.flac",
 		"testdata/flac-test-files/subset/31 - high resolution audio, using only 32nd order predictors.flac",
-		// TODO: implement encoding of escaped partitions.
-		//"testdata/flac-test-files/subset/32 - high resolution audio, partition order 8 containing escaped partitions.flac",
+		"testdata/flac-test-files/subset/32 - high resolution audio, partition order 8 containing escaped partitions.flac",
 		"testdata/flac-test-files/subset/33 - samplerate 192kHz.flac",
 		// TODO: fix diff.
 		//"testdata/flac-test-files/subset/34 - samplerate 192kHz, using only 32nd order predictors.flac",
@@ -108,8 +106,7 @@ func TestEncode(t *testing.T) {
 		"testdata/flac-test-files/subset/61 - predictor overflow check, 16-bit.flac",
 		"testdata/flac-test-files/subset/62 - predictor overflow check, 20-bit.flac",
 		"testdata/flac-test-files/subset/63 - predictor overflow check, 24-bit.flac",
-		// TODO: implement encoding of escaped partitions.
-		//"testdata/flac-test-files/subset/64 - rice partitions with escape code zero.flac",
+		"testdata/flac-test-files/subset/64 - rice partitions with escape code zero.flac",
 	}
 	for _, path := range paths {
 		t.Run(path, func(t *testing.T) {
diff --git a/encode_subframe.go b/encode_subframe.go
index 3bc1cf0..30f0c20 100644
--- a/encode_subframe.go
+++ b/encode_subframe.go
@@ -278,8 +278,27 @@ func encodeRicePart(bw *bitio.Writer, subframe *frame.Subframe, paramSize uint,
 		}
 
 		if paramSize == 4 && param == 0xF || paramSize == 5 && param == 0x1F {
-			// TODO: implement encoding of escaped partitions.
-			panic("TODO: implement encoding of escaped partitions.")
+			// 1111 or 11111: Escape code, meaning the partition is in unencoded
+			// binary form using n bits per sample; n follows as a 5-bit number.
+			if err := bw.WriteBits(uint64(partition.EscapedBitsPerSample), 5); err != nil {
+				return errutil.Err(err)
+			}
+			for j := 0; j < nsamples; j++ {
+				// ref: https://datatracker.ietf.org/doc/draft-ietf-cellar-flac/
+				//
+				// From section 9.2.7.1.  Escaped partition:
+				//
+				// The residual samples themselves are stored signed two's
+				// complement.  For example, when a partition is escaped and each
+				// residual sample is stored with 3 bits, the number -1 is
+				// represented as 0b111.
+				residual := residuals[curResidualIndex]
+				curResidualIndex++
+				if err := bw.WriteBits(uint64(residual), uint8(partition.EscapedBitsPerSample)); err != nil {
+					return errutil.Err(err)
+				}
+			}
+			continue
 		}
 
 		// Encode the Rice residuals of the partition.
diff --git a/frame/subframe.go b/frame/subframe.go
index 2ef5cb1..b906609 100644
--- a/frame/subframe.go
+++ b/frame/subframe.go
@@ -97,6 +97,8 @@ type RiceSubframe struct {
 type RicePartition struct {
 	// Rice parameter.
 	Param uint
+	// Residual sample size in bits-per-sample used by escaped partitions.
+	EscapedBitsPerSample uint
 }
 
 // parseHeader reads and parses the header of a subframe.
@@ -455,6 +457,7 @@ func (subframe *Subframe) decodeRicePart(br *bits.Reader, paramSize uint) error
 				return unexpected(err)
 			}
 			n := uint(x)
+			partition.EscapedBitsPerSample = n
 			for j := 0; j < nsamples; j++ {
 				sample, err := br.Read(n)
 				if err != nil {