From fc58c984807bb40be83ebf1bf4a93bad72a02ba8 Mon Sep 17 00:00:00 2001
From: wenlingang <lingang.wen@longbridge.sg>
Date: Thu, 3 Aug 2023 13:20:24 +0800
Subject: [PATCH] Add slice Insert & InsertSlice, Inserts an element or
 multiple elements at index i.

---
 README.md     | 24 ++++++++++++++++++++++++
 slice.go      | 22 ++++++++++++++++++++++
 slice_test.go | 26 ++++++++++++++++++++++++--
 3 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index eddbdecb..5019b890 100644
--- a/README.md
+++ b/README.md
@@ -116,6 +116,8 @@ Supported helpers for slices:
 - [Compact](#compact)
 - [IsSorted](#issorted)
 - [IsSortedByKey](#issortedbykey)
+- [Insert](#insert)
+- [InsertSlice](#insertslice)
 
 Supported helpers for maps:
 
@@ -923,6 +925,28 @@ slice := lo.IsSortedByKey([]string{"a", "bb", "ccc"}, func(s string) int {
 
 [[play](https://go.dev/play/p/wiG6XyBBu49)]
 
+### Insert
+
+Inserts an element at index i.
+
+```go
+result := lo.Insert(["a", "c", "d"], 1, "b")
+// []string{"a", "b", "c", "d"}
+```
+
+[[play](https://go.dev/play/p/Eu3pk6_ODVQ)]
+
+### InsertSlice
+
+Inserts multiple elements at index i.
+
+```go
+result := lo.InsertSlice([]string{"a", "d"}, 1, []string{"b", "c"})
+// []string{"a", "b", "c", "d"}
+```
+
+[[play](https://go.dev/play/p/8gXdkTovsIC)]
+
 ### Keys
 
 Creates an array of the map keys.
diff --git a/slice.go b/slice.go
index 49c991f8..95c193ba 100644
--- a/slice.go
+++ b/slice.go
@@ -592,3 +592,25 @@ func IsSortedByKey[T any, K constraints.Ordered](collection []T, iteratee func(i
 
 	return true
 }
+
+// Insert inserts an element at index i.
+// Play: https://go.dev/play/p/Eu3pk6_ODVQ
+func Insert[T any](collection []T, i int, element T) []T {
+	result := make([]T, len(collection)+1)
+	copy(result, collection[:i])
+	result[i] = element
+	copy(result[i+1:], collection[i:])
+
+	return result
+}
+
+// InsertSlice inserts multiple elements at index i.
+// Play: https://go.dev/play/p/8gXdkTovsIC
+func InsertSlice[T any](collection []T, i int, elements []T) []T {
+	result := make([]T, len(collection)+len(elements))
+	copy(result, collection[:i])
+	copy(result[i:], elements)
+	copy(result[i+len(elements):], collection[i:])
+
+	return result
+}
diff --git a/slice_test.go b/slice_test.go
index 03326a01..f9869cd8 100644
--- a/slice_test.go
+++ b/slice_test.go
@@ -388,7 +388,7 @@ func TestAssociate(t *testing.T) {
 
 func TestSliceToMap(t *testing.T) {
 	t.Parallel()
-	
+
 	type foo struct {
 		baz string
 		bar int
@@ -626,7 +626,7 @@ func TestSlice(t *testing.T) {
 	out16 := Slice(in, -10, 1)
 	out17 := Slice(in, -1, 3)
 	out18 := Slice(in, -10, 7)
-	
+
 	is.Equal([]int{}, out1)
 	is.Equal([]int{0}, out2)
 	is.Equal([]int{0, 1, 2, 3, 4}, out3)
@@ -759,3 +759,25 @@ func TestIsSortedByKey(t *testing.T) {
 		return ret
 	}))
 }
+
+func TestInsert(t *testing.T) {
+	t.Parallel()
+	is := assert.New(t)
+
+	inserts := Insert([]string{"a", "c", "d"}, 1, "b")
+	is.True(len(inserts) == 4 && inserts[1] == "b")
+
+	inserti := Insert([]int{1, 3, 4}, 1, 2)
+	is.True(len(inserti) == 4 && inserti[1] == 2)
+}
+
+func TestInsertSlice(t *testing.T) {
+	t.Parallel()
+	is := assert.New(t)
+
+	inserts := InsertSlice([]string{"a", "d"}, 1, []string{"b", "c"})
+	is.True(len(inserts) == 4 && inserts[1] == "b" && inserts[2] == "c")
+
+	inserti := InsertSlice([]int{1, 4}, 1, []int{2, 3})
+	is.True(len(inserti) == 4 && inserti[1] == 2 && inserti[2] == 3)
+}