From 7c990b3bb521312c7497ff21cb4d4ffe268573ea Mon Sep 17 00:00:00 2001
From: Trek H <trek.hopton@gmail.com>
Date: Wed, 10 Apr 2019 17:18:42 +0930
Subject: [PATCH] mts: reordered, neatened and clarified code.

---
 container/mts/audio_test.go | 77 ++++++++++++++++++-------------------
 container/mts/encoder.go    |  7 ++--
 container/mts/pes/pes.go    |  2 +-
 3 files changed, 43 insertions(+), 43 deletions(-)

diff --git a/container/mts/audio_test.go b/container/mts/audio_test.go
index f4aaf902..23ba16e6 100644
--- a/container/mts/audio_test.go
+++ b/container/mts/audio_test.go
@@ -29,9 +29,10 @@ import (
 	"io/ioutil"
 	"testing"
 
-	"bitbucket.org/ausocean/av/container/mts/meta"
 	"github.com/Comcast/gots/packet"
 	"github.com/Comcast/gots/pes"
+
+	"bitbucket.org/ausocean/av/container/mts/meta"
 )
 
 // TestEncodePcm tests the mpegts encoder's ability to encode pcm audio data.
@@ -39,13 +40,12 @@ import (
 func TestEncodePcm(t *testing.T) {
 	Meta = meta.New()
 
-	var b []byte
-	buf := bytes.NewBuffer(b)
+	var buf bytes.Buffer
 	sampleRate := 48000
 	sampleSize := 2
 	blockSize := 16000
 	writeFreq := float64(sampleRate*sampleSize) / float64(blockSize)
-	e := NewEncoder(buf, writeFreq, Audio)
+	e := NewEncoder(&buf, writeFreq, Audio)
 
 	inPath := "../../../test/test-data/av/input/sweep_400Hz_20000Hz_-3dBFS_5s_48khz.pcm"
 	inPcm, err := ioutil.ReadFile(inPath)
@@ -84,11 +84,34 @@ func TestEncodePcm(t *testing.T) {
 	for i+PacketSize <= len(clip) {
 
 		// Check MTS packet
-		if pkt.PID() == audioPid {
-			if pkt.PayloadUnitStartIndicator() {
+		if !(pkt.PID() == audioPid) {
+			i += PacketSize
+			if i+PacketSize <= len(clip) {
+				copy(pkt[:], clip[i:i+PacketSize])
+			}
+			continue
+		}
+		if !pkt.PayloadUnitStartIndicator() {
+			i += PacketSize
+			if i+PacketSize <= len(clip) {
+				copy(pkt[:], clip[i:i+PacketSize])
+			}
+		} else {
+			// Copy the first MTS payload
+			payload, err := pkt.Payload()
+			if err != nil {
+				t.Errorf("unable to get MTS payload: %v", err)
+			}
+			pesPacket = append(pesPacket, payload...)
 
-				// Copy the first MTS payload
-				payload, err := pkt.Payload()
+			i += PacketSize
+			if i+PacketSize <= len(clip) {
+				copy(pkt[:], clip[i:i+PacketSize])
+			}
+
+			// Copy the rest of the MTS payloads that are part of the same PES packet
+			for (!pkt.PayloadUnitStartIndicator()) && i+PacketSize <= len(clip) {
+				payload, err = pkt.Payload()
 				if err != nil {
 					t.Errorf("unable to get MTS payload: %v", err)
 				}
@@ -98,39 +121,15 @@ func TestEncodePcm(t *testing.T) {
 				if i+PacketSize <= len(clip) {
 					copy(pkt[:], clip[i:i+PacketSize])
 				}
-
-				// Copy the rest of the MTS payloads that are part of the same PES packet
-				for (!pkt.PayloadUnitStartIndicator()) && i+PacketSize <= len(clip) {
-					payload, err = pkt.Payload()
-					if err != nil {
-						t.Errorf("unable to get MTS payload: %v", err)
-					}
-					pesPacket = append(pesPacket, payload...)
-
-					i += PacketSize
-					if i+PacketSize <= len(clip) {
-						copy(pkt[:], clip[i:i+PacketSize])
-					}
-				}
-			} else {
-				i += PacketSize
-				if i+PacketSize <= len(clip) {
-					copy(pkt[:], clip[i:i+PacketSize])
-				}
-			}
-			// Get the audio data from the current PES packet
-			pesHeader, err := pes.NewPESHeader(pesPacket)
-			if err != nil {
-				t.Errorf("unable to read PES packet: %v", err)
-			}
-			got = append(got, pesHeader.Data()...)
-			pesPacket = pesPacket[:0]
-		} else {
-			i += PacketSize
-			if i+PacketSize <= len(clip) {
-				copy(pkt[:], clip[i:i+PacketSize])
 			}
 		}
+		// Get the audio data from the current PES packet
+		pesHeader, err := pes.NewPESHeader(pesPacket)
+		if err != nil {
+			t.Errorf("unable to read PES packet: %v", err)
+		}
+		got = append(got, pesHeader.Data()...)
+		pesPacket = pesPacket[:0]
 	}
 
 	// Compare data from MTS with original data.
diff --git a/container/mts/encoder.go b/container/mts/encoder.go
index 6ccde63e..e72cfebf 100644
--- a/container/mts/encoder.go
+++ b/container/mts/encoder.go
@@ -147,8 +147,9 @@ type Encoder struct {
 	psiLastTime time.Time
 }
 
-// NewEncoder returns an Encoder with the specified frame rate.
-func NewEncoder(dst io.Writer, writeFreq float64, mediaType int) *Encoder {
+// NewEncoder returns an Encoder with the specified media type and rate eg. if a video stream
+// calls write for every frame, the rate will be the frame rate of the video.
+func NewEncoder(dst io.Writer, rate float64, mediaType int) *Encoder {
 	var mPid int
 	var sid byte
 	switch mediaType {
@@ -163,7 +164,7 @@ func NewEncoder(dst io.Writer, writeFreq float64, mediaType int) *Encoder {
 	return &Encoder{
 		dst: dst,
 
-		writePeriod: time.Duration(float64(time.Second) / writeFreq),
+		writePeriod: time.Duration(float64(time.Second) / rate),
 		ptsOffset:   ptsOffset,
 
 		timeBasedPsi: true,
diff --git a/container/mts/pes/pes.go b/container/mts/pes/pes.go
index 50544aa2..b0e40f86 100644
--- a/container/mts/pes/pes.go
+++ b/container/mts/pes/pes.go
@@ -26,7 +26,7 @@ LICENSE
 
 package pes
 
-const MaxPesSize = 65536
+const MaxPesSize = 64 * 1 << 10
 
 /*
 The below data struct encapsulates the fields of an PES packet. Below is