From 6cf5697ae550ae6534031832cb8b0cf031d0b191 Mon Sep 17 00:00:00 2001
From: Ishank <ishank@hotstar.com>
Date: Tue, 7 Jan 2025 18:43:13 +0530
Subject: [PATCH] Feature: User specified initial capacity for slices

Added struct tag to enable custom initial capacity for slices
---
 Makefile               |  3 ++-
 gen/decoder.go         |  4 +++-
 gen/encoder.go         | 29 ++++++++++++++++++++++-------
 tests/basic_test.go    |  5 ++++-
 tests/init_cap.go      | 25 +++++++++++++++++++++++++
 tests/init_cap_test.go | 40 ++++++++++++++++++++++++++++++++++++++++
 6 files changed, 96 insertions(+), 10 deletions(-)
 create mode 100644 tests/init_cap.go
 create mode 100644 tests/init_cap_test.go

diff --git a/Makefile b/Makefile
index cc5ebbad..5709a990 100644
--- a/Makefile
+++ b/Makefile
@@ -32,7 +32,8 @@ generate: build
  		./tests/nothing.go \
  		./tests/errors.go \
  		./tests/html.go \
- 		./tests/type_declaration_skip.go
+ 		./tests/type_declaration_skip.go \
+ 		./tests/init_cap.go
 	bin/easyjson \
 		./tests/nested_easy.go \
 		./tests/named_type.go \
diff --git a/gen/decoder.go b/gen/decoder.go
index df1a4bd3..eb2c9998 100644
--- a/gen/decoder.go
+++ b/gen/decoder.go
@@ -148,11 +148,13 @@ func (g *Generator) genTypeDecoderNoCheck(t reflect.Type, out string, tags field
 			fmt.Fprintln(g.out, ws+"}")
 
 		} else {
-
 			capacity := 1
 			if elem.Size() > 0 {
 				capacity = minSliceBytes / int(elem.Size())
 			}
+			if tags.initialCapacity > 0 {
+				capacity = tags.initialCapacity
+			}
 
 			fmt.Fprintln(g.out, ws+"if in.IsNull() {")
 			fmt.Fprintln(g.out, ws+"  in.Skip()")
diff --git a/gen/encoder.go b/gen/encoder.go
index 22db5e98..7a5d54da 100644
--- a/gen/encoder.go
+++ b/gen/encoder.go
@@ -53,18 +53,22 @@ var primitiveStringEncoders = map[reflect.Kind]string{
 type fieldTags struct {
 	name string
 
-	omit        bool
-	omitEmpty   bool
-	noOmitEmpty bool
-	asString    bool
-	required    bool
-	intern      bool
-	noCopy      bool
+	omit            bool
+	omitEmpty       bool
+	noOmitEmpty     bool
+	asString        bool
+	required        bool
+	intern          bool
+	noCopy          bool
+	initialCapacity int
 }
 
+const defaultCapacity int = -1
+
 // parseFieldTags parses the json field tag into a structure.
 func parseFieldTags(f reflect.StructField) fieldTags {
 	var ret fieldTags
+	ret.initialCapacity = defaultCapacity
 
 	for i, s := range strings.Split(f.Tag.Get("json"), ",") {
 		switch {
@@ -87,6 +91,17 @@ func parseFieldTags(f reflect.StructField) fieldTags {
 		}
 	}
 
+	for i, s := range strings.Split(f.Tag.Get("initialCap"), ",") {
+		switch {
+		case i == 0:
+			capacity, err := strconv.Atoi(s)
+			if err != nil || capacity <= 0 {
+				capacity = defaultCapacity
+			}
+			ret.initialCapacity = capacity
+		}
+	}
+
 	return ret
 }
 
diff --git a/tests/basic_test.go b/tests/basic_test.go
index 190f6d01..98c0b431 100644
--- a/tests/basic_test.go
+++ b/tests/basic_test.go
@@ -59,6 +59,9 @@ var testCases = []struct {
 	{&myTypeDeclaredValue, myTypeDeclaredString},
 	{&myTypeNotSkippedValue, myTypeNotSkippedString},
 	{&intern, internString},
+	{&noInitCap, noInitCapString},
+	{&initCap, initCapString},
+	{&invalidInitCap, invalidInitCapString},
 }
 
 func TestMarshal(t *testing.T) {
@@ -243,7 +246,7 @@ func TestNestedMarshaler(t *testing.T) {
 		t.Errorf("Can't marshal NestedMarshaler: %s", err)
 	}
 
-	s2 := NestedMarshaler {
+	s2 := NestedMarshaler{
 		Value: &StructWithMarshaler{},
 	}
 
diff --git a/tests/init_cap.go b/tests/init_cap.go
new file mode 100644
index 00000000..0e8a6807
--- /dev/null
+++ b/tests/init_cap.go
@@ -0,0 +1,25 @@
+package tests
+
+//easyjson:json
+type NoInitCap struct {
+	Field []string `json:"field"`
+}
+
+//easyjson:json
+type InitCap struct {
+	Field []string `json:"field" initialCap:"10"`
+}
+
+//easyjson:json
+type InvalidInitCap struct {
+	Field []string `json:"field" initialCap:"-3"`
+}
+
+var noInitCap = NoInitCap{Field: []string{"elem"}}
+var noInitCapString = `{"field":["elem"]}`
+
+var initCap = InitCap{Field: []string{"elem"}}
+var initCapString = `{"field":["elem"]}`
+
+var invalidInitCap = InvalidInitCap{Field: []string{"elem"}}
+var invalidInitCapString = `{"field":["elem"]}`
diff --git a/tests/init_cap_test.go b/tests/init_cap_test.go
new file mode 100644
index 00000000..15374e0d
--- /dev/null
+++ b/tests/init_cap_test.go
@@ -0,0 +1,40 @@
+package tests
+
+import (
+	"github.com/mailru/easyjson"
+	"testing"
+)
+
+func TestTagInitialCapacity(t *testing.T) {
+	data := []byte(`{"field":["elem"]}`)
+
+	var n NoInitCap
+	err := easyjson.Unmarshal(data, &n)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if cap(n.Field) != 4 {
+		t.Fatalf("wrong slice capacity: %d", cap(n.Field))
+	}
+
+	var c InitCap
+	err = easyjson.Unmarshal(data, &c)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if cap(c.Field) != 10 {
+		t.Fatalf("wrong slice capacity: %d", cap(c.Field))
+	}
+
+	var i InvalidInitCap
+	err = easyjson.Unmarshal(data, &i)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if cap(i.Field) != 4 {
+		t.Fatalf("wrong slice capacity: %d", cap(i.Field))
+	}
+}