From 70eb8193cb8eda885bed1c51056d93940a6fe602 Mon Sep 17 00:00:00 2001
From: Saxon <saxon@ausocean.org>
Date: Fri, 14 Jun 2019 01:09:55 +0930
Subject: [PATCH] container/mts: wrote BytesForPTSInterval

This function will return the media between two provided PTS. Binary search has been used to
find the corresponding Frames to 'from' and 'to'.
---
 container/mts/payload.go | 54 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/container/mts/payload.go b/container/mts/payload.go
index f1f6bb6a..a824cb64 100644
--- a/container/mts/payload.go
+++ b/container/mts/payload.go
@@ -2,6 +2,7 @@ package mts
 
 import (
 	"errors"
+	"sort"
 
 	"github.com/Comcast/gots/packet"
 	"github.com/Comcast/gots/pes"
@@ -76,6 +77,7 @@ func Extract(p []byte) (*Clip, error) {
 				// and can add this data to the frame pertaining to the finish frame.
 				if !firstPUSI {
 					clip.frames[len(clip.frames)-2].Media = clip.backing[frameStart:lenOfFrame]
+					clip.frames[len(clip.frames)-2].idx = frameStart
 					frameStart = lenOfFrame
 				}
 				firstPUSI = false
@@ -90,6 +92,7 @@ func Extract(p []byte) (*Clip, error) {
 	}
 	// We're finished up with media frames, so give the final Frame it's data.
 	clip.frames[len(clip.frames)-1].Media = clip.backing[frameStart:lenOfFrame]
+	clip.frames[len(clip.frames)-1].idx = frameStart
 	return clip, nil
 }
 
@@ -105,6 +108,7 @@ type Frame struct {
 	PTS   uint64            // PTS from PES packet (this gives time relative from start of stream).
 	ID    uint8             // StreamID from the PES packet, identifying media codec.
 	Meta  map[string]string // Contains metadata from PMT relevant to this frame.
+	idx   int               // Index in the backing slice.
 }
 
 // Bytes returns the concatentated media bytes from each frame in the Clip c.
@@ -114,3 +118,53 @@ func (c *Clip) Bytes() []byte {
 	}
 	return c.backing
 }
+
+// BytesForPTSInterval returns the media data between PTS' from and to. If from
+// sits between two PTS, the Frame posessing lower PTS will be considered the start.
+// The Frame before the Frame corresponding to to will be considered the final
+// Frame.
+func (c *Clip) BytesForPTSInterval(from, to uint64) ([]byte, error) {
+	// First check that the interval makes sense.
+	if from >= to {
+		return nil, errors.New("PTS interval is invalid")
+	}
+
+	// Use binary search to find 'from'.
+	n := len(c.frames)
+	idx := sort.Search(
+		n,
+		func(i int) bool {
+			if from >= c.frames[i].PTS && from < c.frames[i].PTS {
+				return true
+			}
+			return false
+		},
+	)
+	if idx == n {
+		return nil, errors.New("'from' cannot be found")
+	}
+
+	// Now get the start index for the backing slice from this Frame.
+	start := c.frames[idx].idx
+
+	// Now use binary search again to find 'to'.
+	off := idx + 1
+	n = n - (off)
+	idx = sort.Search(
+		n,
+		func(i int) bool {
+			if to >= c.frames[i+off].PTS && from < c.frames[i+off].PTS {
+				return true
+			}
+			return false
+		},
+	)
+	if idx == n {
+		return nil, errors.New("'to' cannot be found")
+	}
+
+	// Now get the end index for the backing slice from this Frame, and return
+	// segment from backing slice corresponding to start and end.
+	end := c.frames[idx+off].idx
+	return c.backing[start:end], nil
+}