Fix arr-ai#128, add the slice feature to all types of Set

rel/expr_range.go

@@ -0,0 +1,73 @@
+package rel
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/go-errors/errors"
+)
+
+type SliceExpr struct {
+	setToSlice, start, end, step Expr
+	include                      bool
+}
+
+func NewSliceExpr(setToSlice, start, end, step Expr, include bool) SliceExpr {
+	return SliceExpr{setToSlice, start, end, step, include}
+}
+
+func (s SliceExpr) Eval(local Scope) (Value, error) {
+	var start, end, step Value
+	var err error
+
+	if s.start != nil {
+		start, err = s.start.Eval(local)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if s.end != nil {
+		end, err = s.end.Eval(local)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if s.step != nil {
+		step, err = s.step.Eval(local)
+		if err != nil {
+			return nil, err
+		}
+		if _, isNumber := step.(Number); !isNumber {
+			return nil, errors.Errorf("step %s must be a number", step)
+		}
+	} else {
+		step = Number(1)
+	}
+
+	set, err := s.setToSlice.Eval(local)
+	if err != nil {
+		return nil, err
+	}
+
+	return set.(Set).CallSlice(start, end, int(step.(Number)), s.include), nil
+}
+
+func (s SliceExpr) String() string {
+	str := strings.Builder{}
+	switch {
+	case s.start == nil && s.end == nil:
+		str.WriteString(";")
+	case s.start != nil && s.end == nil:
+		str.WriteString(fmt.Sprintf("%s;", s.start))
+	case s.start == nil && s.end != nil:
+		str.WriteString(fmt.Sprintf(";%s", s.end))
+	default:
+		str.WriteString(fmt.Sprintf("%s;%s", s.start, s.end))
+	}
+	if s.step != nil {
+		str.WriteString(fmt.Sprintf(";%s", s.step))
+	}
+	return str.String()
+}

rel/slice_util.go

@@ -0,0 +1,86 @@
+package rel
+
+import "math"
+
+// resolveArrayIndexes
+func resolveArrayIndexes(start, end Value, step, offset, maxLen int, inclusive bool) []int {
+	if maxLen == 0 {
+		return []int{}
+	}
+	startIndex, endIndex := initDefaultArrayIndex(start, end, offset, maxLen+offset, step)
+
+	if startIndex == endIndex {
+		if inclusive {
+			return []int{startIndex}
+		}
+		return []int{}
+	}
+
+	return getIndexes(startIndex, endIndex, step, inclusive)
+}
+
+func initDefaultArrayIndex(start, end Value, minLen, maxLen, step int) (startIndex int, endIndex int) {
+	if start != nil {
+		startIndex = resolveIndex(int(start.(Number)), minLen, maxLen)
+		if startIndex == maxLen {
+			startIndex--
+		}
+	} else {
+		if step > 0 {
+			startIndex = minLen
+		} else {
+			startIndex = maxLen - 1
+		}
+	}
+
+	if end != nil {
+		endIndex = resolveIndex(int(end.(Number)), minLen, maxLen)
+	} else {
+		if step > 0 {
+			endIndex = maxLen
+		} else {
+			endIndex = minLen - 1
+		}
+	}
+	return
+}
+
+func resolveIndex(i, minVal, maxVal int) int {
+	if i > maxVal {
+		return maxVal
+	} else if i < 0 {
+		if -i > maxVal {
+			return minVal
+		}
+		return maxVal + i
+	}
+	return i
+}
+
+func getIndexes(start, end, step int, inclusive bool) []int {
+	if !isValidRange(start, end, step) {
+		return []int{}
+	}
+	if inclusive {
+		if step > 0 {
+			end++
+		} else {
+			end--
+		}
+	}
+
+	length := int(math.Abs(float64(start - end)))
+	if step != 1 && step != -1 {
+		length = int(math.Ceil(float64(length) / math.Abs(float64(step))))
+	}
+	indexes := make([]int, 0, length)
+	for i := 0; i < length; i++ {
+		indexes = append(indexes, start+step*i)
+	}
+
+	return indexes
+}
+
+func isValidRange(start, end, step int) bool {
+	return step != 0 && ((start > end && step < 0) || (start < end && step > 0))
+}

rel/slice_util_test.go

@@ -0,0 +1,84 @@
+package rel
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetIndexes(t *testing.T) {
+	t.Parallel()
+
+	assert.Equal(t, []int{1, 2, 3, 4, 5}, getIndexes(1, 6, 1, false))
+	assert.Equal(t, []int{1, 2, 3, 4, 5}, getIndexes(1, 5, 1, true))
+	assert.Equal(t, []int{2, 4}, getIndexes(2, 5, 2, false))
+	assert.Equal(t, []int{1, 3, 5}, getIndexes(1, 5, 2, true))
+	assert.Equal(t, []int{5, 3}, getIndexes(5, 1, -2, false))
+	assert.Equal(t, []int{5, 3, 1}, getIndexes(5, 1, -2, true))
+	assert.Equal(t, []int{}, getIndexes(5, 1, 1, true))
+	assert.Equal(t, []int{}, getIndexes(5, 1, 0, false))
+}
+
+func TestInitDefaultArrayIndex(t *testing.T) {
+	t.Parallel()
+
+	start, end := initDefaultArrayIndex(Number(10), Number(20), 5, 25, 1)
+	assert.Equal(t, 10, start)
+	assert.Equal(t, 20, end)
+
+	start, end = initDefaultArrayIndex(nil, nil, 5, 25, 1)
+	assert.Equal(t, 5, start)
+	assert.Equal(t, 25, end)
+
+	start, end = initDefaultArrayIndex(nil, nil, 5, 25, -2)
+	assert.Equal(t, 24, start)
+	assert.Equal(t, 4, end)
+
+	start, end = initDefaultArrayIndex(nil, Number(12), 5, 25, -2)
+	assert.Equal(t, 24, start)
+	assert.Equal(t, 12, end)
+
+	start, end = initDefaultArrayIndex(nil, Number(12), 5, 25, 2)
+	assert.Equal(t, 5, start)
+	assert.Equal(t, 12, end)
+
+	start, end = initDefaultArrayIndex(Number(7), nil, 5, 25, -2)
+	assert.Equal(t, 7, start)
+	assert.Equal(t, 4, end)
+
+	start, end = initDefaultArrayIndex(Number(7), nil, 5, 25, 2)
+	assert.Equal(t, 7, start)
+	assert.Equal(t, 25, end)
+
+	start, end = initDefaultArrayIndex(nil, Number(42), 5, 25, -2)
+	assert.Equal(t, 24, start)
+	assert.Equal(t, 25, end)
+
+	start, end = initDefaultArrayIndex(Number(42), nil, 5, 25, 2)
+	assert.Equal(t, 24, start)
+	assert.Equal(t, 25, end)
+
+	start, end = initDefaultArrayIndex(Number(-5), nil, 5, 25, 2)
+	assert.Equal(t, 20, start)
+	assert.Equal(t, 25, end)
+
+	start, end = initDefaultArrayIndex(nil, Number(-5), 5, 25, 2)
+	assert.Equal(t, 5, start)
+	assert.Equal(t, 20, end)
+
+	start, end = initDefaultArrayIndex(Number(-30), nil, 5, 25, 2)
+	assert.Equal(t, 5, start)
+	assert.Equal(t, 25, end)
+
+	start, end = initDefaultArrayIndex(nil, Number(-30), 5, 25, 2)
+	assert.Equal(t, 5, start)
+	assert.Equal(t, 5, end)
+
+	start, end = initDefaultArrayIndex(Number(42), Number(-30), 5, 25, 2)
+	assert.Equal(t, 24, start)
+	assert.Equal(t, 5, end)
+
+	start, end = initDefaultArrayIndex(Number(-30), Number(42), 5, 25, 2)
+	assert.Equal(t, 5, start)
+	assert.Equal(t, 25, end)
+}

rel/value.go

@@ -106,6 +106,7 @@ type Set interface {
 	Map(func(Value) Value) Set
 	Where(func(Value) bool) Set
 	Call(arg Value) Value
+	CallSlice(start, end Value, step int, inclusive bool) Set
 
 	ArrayEnumerator() (OffsetValueEnumerator, bool)
 }

rel/value_set_array.go

@@ -263,6 +263,15 @@ func (a Array) Call(arg Value) Value {
 	return a.values[i-a.offset]
 }
 
+func (a Array) CallSlice(start, end Value, step int, inclusive bool) Set {
+	indexes := resolveArrayIndexes(start, end, step, a.offset, len(a.values), inclusive)
+	slice := make([]Value, 0, len(indexes))
+	for _, i := range indexes {
+		slice = append(slice, a.values[i-a.offset])
+	}
+	return NewOffsetArray(a.offset, slice...)
+}
+
 func (a Array) index(pos int) int {
 	pos -= a.offset
 	if 0 <= pos && pos <= len(a.values) {

rel/value_set_bytes.go

@@ -223,6 +223,15 @@ func (b Bytes) Call(arg Value) Value {
 	return NewNumber(float64(string(b.b)[i-b.offset]))
 }
 
+func (b Bytes) CallSlice(start, end Value, step int, inclusive bool) Set {
+	indexes := resolveArrayIndexes(start, end, step, b.offset, len(b.b), inclusive)
+	slice := make([]byte, 0, len(indexes))
+	for _, i := range indexes {
+		slice = append(slice, b.b[i-b.offset])
+	}
+	return NewOffsetBytes(slice, b.offset)
+}
+
 func (b Bytes) index(pos int) int {
 	pos -= b.offset
 	if 0 <= pos && pos <= len(b.b) {

rel/value_set_dict.go

@@ -269,6 +269,97 @@ func (d Dict) Call(arg Value) Value {
 	}
 }
 
+func (d Dict) CallSlice(start, end Value, step int, inclusive bool) Set {
+	allKeys := d.m.Keys().OrderedElements(
+		func(a, b interface{}) bool {
+			return a.(Value).Less(b.(Value))
+		},
+	)
+
+	keys := []int{}
+	for _, k := range allKeys {
+		if f, isFloat := k.(Number); isFloat {
+			keys = append(keys, int(f))
+		}
+	}
+
+	if len(keys) == 0 {
+		return None
+	}
+
+	// Add 1 to endIndex to account for undefined end so that it includes last value
+	startIndex, endIndex := keys[0], keys[len(keys)-1]+1
+
+	// TODO: apply inclusivity to the initial index
+	// if !inclusive {
+	// 	endIndex--
+	// }
+
+	if step < 0 {
+		startIndex, endIndex = endIndex, startIndex
+	}
+	if start != nil {
+		startIndex = int(start.(Number))
+	}
+	if end != nil {
+		endIndex = int(end.(Number))
+	}
+
+	if startIndex == endIndex {
+		if inclusive {
+			val, exists := d.m.Get(Number(startIndex))
+			if !exists {
+				return None
+			}
+			return NewArray(NewArrayItemTuple(0, val.(Value)))
+		}
+		return None
+	}
+
+	indexes := getIndexes(startIndex, endIndex, step, inclusive)
+	if len(indexes) == 0 {
+		return None
+	}
+
+	array := Array{[]Value{}, 0}
+	for _, i := range indexes {
+		val, exists := d.m.Get(Number(i))
+		if !exists {
+			continue
+		}
+
+		array = arrayFromDictEntry(val, array)
+	}
+	return array
+}
+
+func arrayFromDictEntry(val interface{}, array Array) Array {
+	switch v := val.(type) {
+	case Value:
+		if len(array.values) == 0 {
+			return NewArray(v).(Array)
+		}
+		return array.With(
+			NewArrayItemTuple(array.offset+len(array.values), v),
+		).(Array)
+	case multipleValues:
+		values := make([]Value, 0, frozen.Set(v).Count())
+		for s := frozen.Set(v).Range(); s.Next(); {
+			values = append(values, s.Value().(Value))
+		}
+		if len(array.values) == 0 {
+			return NewArray(values...).(Array)
+		}
+		return array.With(
+			NewArrayItemTuple(
+				array.offset+len(array.values),
+				NewArray(values...),
+			),
+		).(Array)
+	}
+	return array
+}
+
 func (d Dict) ArrayEnumerator() (OffsetValueEnumerator, bool) {
 	return nil, false
 }