From 7a7e0c0e90bc2b03fe6f2acab896820adfb3cd99 Mon Sep 17 00:00:00 2001
From: Kirill Zhuravlev <kazhuravlev@fastmail.com>
Date: Tue, 3 Sep 2024 20:18:43 +0200
Subject: [PATCH] add SliceIter

---
 slice.go      | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 slice_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/slice.go b/slice.go
index b2907aa..404eef3 100644
--- a/slice.go
+++ b/slice.go
@@ -725,6 +725,53 @@ func Slice2Iter[T any](in []T) func(func(int, T) bool) {
 	}
 }
 
+type IterContext interface {
+	// Idx returns an element index
+	Idx() int
+	RevIdx() int
+	IsFirst() bool
+	IsLast() bool
+}
+
+type iterContext struct {
+	idx   int
+	inLen int
+}
+
+func (i iterContext) Idx() int {
+	return i.idx
+}
+
+func (i iterContext) RevIdx() int {
+	return i.inLen - i.idx - 1
+}
+
+func (i iterContext) IsFirst() bool {
+	return i.idx == 0
+}
+
+func (i iterContext) IsLast() bool {
+	return i.idx == i.inLen-1
+}
+
+var _ IterContext = (*iterContext)(nil)
+
+// SliceIter create an iterator from slice. The first argument will contain a useful context struct.
+func SliceIter[T any](in []T) func(func(IterContext, T) bool) {
+	return func(yield func(loop IterContext, elem T) bool) {
+		inLen := len(in)
+		for i := range in {
+			ctx := iterContext{
+				idx:   i,
+				inLen: inLen,
+			}
+			if !yield(ctx, in[i]) {
+				return
+			}
+		}
+	}
+}
+
 // SliceShuffle will shuffle the slice in-place.
 func SliceShuffle[T any](in []T) {
 	for i := range in {
diff --git a/slice_test.go b/slice_test.go
index eb61bef..797b9d3 100644
--- a/slice_test.go
+++ b/slice_test.go
@@ -2,6 +2,7 @@ package just_test
 
 import (
 	"errors"
+	"iter"
 	"strconv"
 	"testing"
 	"time"
@@ -1862,3 +1863,28 @@ func TestSliceLastN(t *testing.T) {
 		})
 	}
 }
+
+func TestSliceIter(t *testing.T) {
+	t.Parallel()
+
+	f := func(next func() (just.IterContext, int, bool), val, idx, revIdx int, isFirst, isLast bool) {
+		t.Helper()
+
+		iterCtx, elem, valid := next()
+		require.True(t, valid)
+		assert.Equal(t, val, elem)
+		assert.Equal(t, idx, iterCtx.Idx())
+		assert.Equal(t, revIdx, iterCtx.RevIdx())
+		assert.Equal(t, isFirst, iterCtx.IsFirst())
+		assert.Equal(t, isLast, iterCtx.IsLast())
+	}
+
+	in := []int{10, 20, 30, 40}
+	iterator := just.SliceIter(in)
+	next, _ := iter.Pull2(iterator)
+
+	f(next, 10, 0, 3, true, false)
+	f(next, 20, 1, 2, false, false)
+	f(next, 30, 2, 1, false, false)
+	f(next, 40, 3, 0, false, true)
+}