Add and AddByte

R=rsc
DELTA=83  (83 added, 0 deleted, 0 changed)
OCL=34584
CL=34584

2 files changed, 83 insertions, 0 deletions

diff --git a/src/pkg/bytes/bytes.go b/src/pkg/bytes/bytes.go
index 5375fecaa..52aa8cdf4 100644
--- a/src/pkg/bytes/bytes.go
+++ b/src/pkg/bytes/bytes.go
@@ -260,3 +260,45 @@ func TrimSpace(s []byte) []byte {
 	}
 	return s[start:end];
 }
+
+// How big to make a byte array when growing.
+// Heuristic: Scale by 50% to give n log n time.
+func resize(n int) int {
+	if n < 16 {
+		n = 16
+	}
+	return n + n/2;
+}
+
+// Add appends the contents of t to the end of s and returns the result.
+// If s has enough capacity, it is extended in place; otherwise a
+// new array is allocated and returned.
+func Add(s, t []byte) []byte {
+	lens := len(s);
+	lent := len(t);
+	if lens + lent <= cap(s) {
+		s = s[0:lens+lent];
+	} else {
+		news := make([]byte, lens+lent, resize(lens+lent));
+		Copy(news, s);
+		s = news;
+	}
+	Copy(s[lens:lens+lent], t);
+	return s;
+}
+
+// AddByte appends byte b to the end of s and returns the result.
+// If s has enough capacity, it is extended in place; otherwise a
+// new array is allocated and returned.
+func AddByte(s []byte, t byte) []byte {
+	lens := len(s);
+	if lens + 1 <= cap(s) {
+		s = s[0:lens+1];
+	} else {
+		news := make([]byte, lens+1, resize(lens+1));
+		Copy(news, s);
+		s = news;
+	}
+	s[lens] = t;
+	return s;
+}
diff --git a/src/pkg/bytes/bytes_test.go b/src/pkg/bytes/bytes_test.go
index a7667ec21..8443480e5 100644
--- a/src/pkg/bytes/bytes_test.go
+++ b/src/pkg/bytes/bytes_test.go
@@ -121,6 +121,7 @@ var splittests = []SplitTest {
 	SplitTest{ "123",	"",	2, []string{"1", "23"} },
 	SplitTest{ "123",	"",	17, []string{"1", "2", "3"} },
 }
+
 func TestSplit(t *testing.T) {
 	for _, tt := range splittests {
 		a := Split(strings.Bytes(tt.s), strings.Bytes(tt.sep), tt.n);
@@ -261,3 +262,43 @@ func TestToLower(t *testing.T) {
 func TestTrimSpace(t *testing.T) {
 	runStringTests(t, TrimSpace, "TrimSpace", trimSpaceTests);
 }
+
+type AddTest struct {
+	s, t string;
+	cap	int;
+}
+var addtests = []AddTest {
+	AddTest{ "", "", 0 },
+	AddTest{ "a", "", 1 },
+	AddTest{ "a", "b", 1 },
+	AddTest{ "abc", "def", 100 },
+}
+
+func TestAdd(t *testing.T) {
+	for i, test := range addtests {
+		b := make([]byte, len(test.s), test.cap);
+		for i := 0; i < len(test.s); i++ {
+			b[i] = test.s[i]
+		}
+		b = Add(b, strings.Bytes(test.t));
+		if string(b) != test.s+test.t {
+			t.Errorf("Add(%q,%q) = %q", test.s, test.t, string(b));
+		}
+	}
+}
+
+func TestAddByte(t *testing.T) {
+	const N = 2e5;
+	b := make([]byte, 0);
+	for i := 0; i < N; i++ {
+		b = AddByte(b, byte(i))
+	}
+	if len(b) != N {
+		t.Errorf("AddByte: too small; expected %d got %d", N, len(b));
+	}
+	for i, c := range b {
+		if c != byte(i) {
+			t.Fatalf("AddByte: b[%d] should be %d is %d", i, c, byte(i));
+		}
+	}
+}