[SSAI-500] Allow Segment Insertion into MediaPlaylist (#11) (#12)

* Allow Segment Insertion into MediaPlaylist

go.mod

@@ -1,3 +1,5 @@
 module github.com/jwplayer/m3u8
 
 go 1.12
+
+require github.com/stretchr/testify v1.9.0

go.sum

@@ -0,0 +1,19 @@
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
+github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
+github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
+github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
+github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

structure.go

@@ -117,7 +117,6 @@ type MediaPlaylist struct {
 	StartTime           float64
 	StartTimePrecise    bool
 	durationAsInt       bool // output durations as integers of floats?
-	keyformat           int
 	winsize             uint // max number of segments displayed in an encoded playlist; need set to zero for VOD playlists
 	capacity            uint // total capacity of slice used for the playlist
 	head                uint // head of FIFO, we add segments to head

writer.go

@@ -24,6 +24,9 @@ import (
 // ErrPlaylistFull declares the playlist error.
 var ErrPlaylistFull = errors.New("playlist is full")
 
+// ErrPlaylistFull declares the provided playlist is empty
+var ErrPlaylistEmpty = errors.New("playlist is empty")
+
 // Set version of the playlist accordingly with section 7
 func version(ver *uint8, newver uint8) {
 	if *ver < newver {
@@ -335,6 +338,62 @@ func NewMediaPlaylist(winsize uint, capacity uint) (*MediaPlaylist, error) {
 	return p, nil
 }
 
+// InsertSegments allows the insertion of one or multiple MediaSegments into the MediaPlaylist.
+// seqID is the sequence ID at which the new segments should be inserted to
+// This operation does reset playlist cache.
+func (p *MediaPlaylist) InsertSegments(segments []*MediaSegment, seqID uint64) error {
+	if len(segments) == 0 {
+		return ErrPlaylistEmpty
+	}
+
+	// Determine the index where the new segments should be inserted
+	var insertIndex = 0
+	switch {
+	case seqID >= uint64(len(p.Segments)):
+		insertIndex = len(p.Segments)
+	case seqID != 0 && seqID < uint64(len(p.Segments)):
+		insertIndex = int(seqID) - 1
+	}
+
+	adjustment := uint(len(segments))
+
+	// Shift MediaPlaylist segments in preparation for insertion
+	newLength := len(p.Segments) + len(segments)
+	if cap(p.Segments) < newLength {
+		newSegments := make([]*MediaSegment, newLength)
+		copy(newSegments, p.Segments[:insertIndex])
+		copy(newSegments[insertIndex+len(segments):], p.Segments[insertIndex:])
+		p.Segments = newSegments
+	} else {
+		p.Segments = p.Segments[:newLength]
+		copy(p.Segments[insertIndex+len(segments):], p.Segments[insertIndex:])
+	}
+
+	// Insert the segments
+	copy(p.Segments[insertIndex:], segments)
+
+	iterator := 1
+	// Adjust the sequence IDs of the inserted segments
+	for i := insertIndex; i < len(p.Segments[:insertIndex])+len(segments); i++ {
+		p.Segments[i].SeqId = uint64(insertIndex + iterator)
+		iterator++
+	}
+
+	// Adjust the sequence IDs of the following segments
+	for i := insertIndex + len(segments); i < len(p.Segments); i++ {
+		if p.Segments[i] != nil {
+			p.Segments[i].SeqId += uint64(adjustment)
+		}
+	}
+
+	p.count += adjustment
+	p.capacity += adjustment
+	p.tail = p.count
+
+	p.buf.Reset()
+	return nil
+}
+
 // last returns the previously written segment's index
 func (p *MediaPlaylist) last() uint {
 	if p.tail == 0 {

writer_test.go

@@ -20,6 +20,8 @@ import (
 	"sync"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/require"
 )
 
 // Check how master and media playlists implement common Playlist interface
@@ -1053,6 +1055,110 @@ func TestEncodeMediaPlaylistDateRangeTagsForInterstitials(t *testing.T) {
 	}
 }
 
+func TestInsertSegments(t *testing.T) {
+	tests := []struct {
+		name            string
+		seqID           uint64
+		initialSegments []*MediaSegment
+		expectedResult  []*MediaSegment
+	}{
+		{
+			name:  "Insert at the beginning",
+			seqID: 0,
+			initialSegments: []*MediaSegment{
+				{SeqId: 1, URI: "original"},
+				{SeqId: 2, URI: "original"},
+				{SeqId: 3, URI: "original"},
+			},
+			expectedResult: []*MediaSegment{
+				{SeqId: 1, URI: "new"},
+				{SeqId: 2, URI: "new"},
+				{SeqId: 3, URI: "original"},
+				{SeqId: 4, URI: "original"},
+				{SeqId: 5, URI: "original"},
+			},
+		},
+		{
+			name:  "Insert in the middle",
+			seqID: 2,
+			initialSegments: []*MediaSegment{
+				{SeqId: 1, URI: "original"},
+				{SeqId: 2, URI: "original"},
+				{SeqId: 3, URI: "original"},
+			},
+			expectedResult: []*MediaSegment{
+				{SeqId: 1, URI: "original"},
+				{SeqId: 2, URI: "new"},
+				{SeqId: 3, URI: "new"},
+				{SeqId: 4, URI: "original"},
+				{SeqId: 5, URI: "original"},
+			},
+		},
+		{
+			name:  "Insert at the end",
+			seqID: 3,
+			initialSegments: []*MediaSegment{
+				{SeqId: 1, URI: "original"},
+				{SeqId: 2, URI: "original"},
+				{SeqId: 3, URI: "original"},
+			},
+			expectedResult: []*MediaSegment{
+				{SeqId: 1, URI: "original"},
+				{SeqId: 2, URI: "original"},
+				{SeqId: 3, URI: "original"},
+				{SeqId: 4, URI: "new"},
+				{SeqId: 5, URI: "new"},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			newSegments := []*MediaSegment{
+				{SeqId: 10, URI: "new"},
+				{SeqId: 11, URI: "new"},
+			}
+			playlist := &MediaPlaylist{
+				Segments: tt.initialSegments,
+				count:    uint(len(tt.initialSegments)),
+				tail:     uint(len(tt.initialSegments)),
+			}
+			err := playlist.InsertSegments(newSegments, tt.seqID)
+			require.NoError(t, err)
+			require.Equal(t, tt.expectedResult, playlist.Segments)
+		})
+	}
+	t.Run("Can fit into capacity", func(t *testing.T) {
+		playlist := &MediaPlaylist{}
+		initialSegments := make([]*MediaSegment, 3, 10)
+		initialSegments[0] = &MediaSegment{SeqId: 1, URI: "original"}
+		initialSegments[1] = &MediaSegment{SeqId: 2, URI: "original"}
+		initialSegments[2] = &MediaSegment{SeqId: 3, URI: "original"}
+		playlist.Segments = initialSegments
+
+		playlist.count = 3
+
+		newSegments := []*MediaSegment{
+			{SeqId: 10, URI: "new"},
+			{SeqId: 11, URI: "new"},
+		}
+
+		err := playlist.InsertSegments(newSegments, 4)
+		require.Equal(t, 5, len(playlist.Segments))
+		require.Equal(t, 10, cap(playlist.Segments))
+		require.NoError(t, err)
+		require.Equal(t, tests[2].expectedResult, playlist.Segments)
+
+	})
+
+	t.Run("Empty segments error", func(t *testing.T) {
+		playlist := &MediaPlaylist{}
+		err := playlist.InsertSegments([]*MediaSegment{}, 0)
+		require.Error(t, err)
+		require.Equal(t, ErrPlaylistEmpty, err)
+	})
+}
+
 // Create new media playlist
 // Add three segments to media playlist
 // Set gap tag for the 2nd segment.