From 27c55f5d355d761bfbf8c3e5f29c26732472ccec Mon Sep 17 00:00:00 2001
From: Daniel <dhaavi@users.noreply.github.com>
Date: Tue, 24 Sep 2019 14:46:04 +0200
Subject: [PATCH] Add new functions to add and get data as a container

---
 container/container.go      | 54 ++++++++++++++++++++++++++++++++++++-
 container/container_test.go | 12 ++++++++-
 2 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/container/container.go b/container/container.go
index 2532a40..70f62de 100644
--- a/container/container.go
+++ b/container/container.go
@@ -48,12 +48,28 @@ func (c *Container) AppendNumber(n uint64) {
 	c.compartments = append(c.compartments, varint.Pack64(n))
 }
 
+// AppendInt appends an int (varint encoded).
+func (c *Container) AppendInt(n int) {
+	c.compartments = append(c.compartments, varint.Pack64(uint64(n)))
+}
+
 // AppendAsBlock appends the length of the data and the data itself. Data will NOT be copied.
 func (c *Container) AppendAsBlock(data []byte) {
 	c.AppendNumber(uint64(len(data)))
 	c.Append(data)
 }
 
+// AppendContainer appends another Container. Data will NOT be copied.
+func (c *Container) AppendContainer(data *Container) {
+	c.compartments = append(c.compartments, data.compartments...)
+}
+
+// AppendContainerAsBlock appends another Container (length and data). Data will NOT be copied.
+func (c *Container) AppendContainerAsBlock(data *Container) {
+	c.AppendNumber(uint64(data.Length()))
+	c.compartments = append(c.compartments, data.compartments...)
+}
+
 // Length returns the full length of all bytes held by the container.
 func (c *Container) Length() (length int) {
 	for i := c.offset; i < len(c.compartments); i++ {
@@ -92,6 +108,16 @@ func (c *Container) Get(n int) ([]byte, error) {
 	return buf, nil
 }
 
+// GetAsContainer returns the given amount of bytes in a new container. Data will NOT be copied and IS consumed.
+func (c *Container) GetAsContainer(n int) (*Container, error) {
+	new := c.gatherAsContainer(n)
+	if new == nil {
+		return nil, errors.New("container: not enough data to return")
+	}
+	c.skip(n)
+	return new, nil
+}
+
 // GetMax returns as much as possible, but the given amount of bytes at maximum. Data MAY be copied and IS consumed.
 func (c *Container) GetMax(n int) []byte {
 	buf := c.gather(n)
@@ -214,6 +240,23 @@ func (c *Container) gather(n int) []byte {
 	return slice[:n]
 }
 
+func (c *Container) gatherAsContainer(n int) (new *Container) {
+	new = &Container{}
+	for i := c.offset; i < len(c.compartments); i++ {
+		if n >= len(c.compartments[i]) {
+			new.compartments = append(new.compartments, c.compartments[i])
+			n -= len(c.compartments[i])
+		} else {
+			new.compartments = append(new.compartments, c.compartments[i][:n])
+			n = 0
+		}
+	}
+	if n > 0 {
+		return nil
+	}
+	return new
+}
+
 func (c *Container) skip(n int) {
 	for i := c.offset; i < len(c.compartments); i++ {
 		if len(c.compartments[i]) <= n {
@@ -233,7 +276,7 @@ func (c *Container) skip(n int) {
 	c.checkOffset()
 }
 
-// GetNextBlock returns the next block of data defined by a varint (note: data will MAY be copied and IS consumed).
+// GetNextBlock returns the next block of data defined by a varint. Data MAY be copied and IS consumed.
 func (c *Container) GetNextBlock() ([]byte, error) {
 	blockSize, err := c.GetNextN64()
 	if err != nil {
@@ -242,6 +285,15 @@ func (c *Container) GetNextBlock() ([]byte, error) {
 	return c.Get(int(blockSize))
 }
 
+// GetNextBlockAsContainer returns the next block of data as a Container defined by a varint. Data will NOT be copied and IS consumed.
+func (c *Container) GetNextBlockAsContainer() (*Container, error) {
+	blockSize, err := c.GetNextN64()
+	if err != nil {
+		return nil, err
+	}
+	return c.GetAsContainer(int(blockSize))
+}
+
 // GetNextN8 parses and returns a varint of type uint8.
 func (c *Container) GetNextN8() (uint8, error) {
 	buf := c.gather(2)
diff --git a/container/container_test.go b/container/container_test.go
index bcaa23b..96886f1 100644
--- a/container/container_test.go
+++ b/container/container_test.go
@@ -66,7 +66,12 @@ func TestContainerDataHandling(t *testing.T) {
 	}
 	c8.clean()
 
-	compareMany(t, testData, c1.CompileData(), c2.CompileData(), c3.CompileData(), d4, d5, c6.CompileData(), c7.CompileData(), c8.CompileData())
+	c9 := c8.gatherAsContainer(len(testData))
+
+	c10 := c9.gatherAsContainer(len(testData) - 1)
+	c10.Append(testData[len(testData)-1:])
+
+	compareMany(t, testData, c1.CompileData(), c2.CompileData(), c3.CompileData(), d4, d5, c6.CompileData(), c7.CompileData(), c8.CompileData(), c9.CompileData(), c10.CompileData())
 }
 
 func compareMany(t *testing.T, reference []byte, other ...[]byte) {
@@ -120,6 +125,11 @@ func TestDataFetching(t *testing.T) {
 	if err == nil {
 		t.Error("should fail")
 	}
+
+	_, err = c1.GetAsContainer(1000)
+	if err == nil {
+		t.Error("should fail")
+	}
 }
 
 func TestBlocks(t *testing.T) {