Clean up and add validation to range funcs

pkg/expressions/stdlib/funcs_test.go

@@ -3,6 +3,7 @@ package stdlib
 import (
 	"fmt"
 	. "rare/pkg/expressions"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -33,6 +34,18 @@ func testExpression(t *testing.T, context KeyBuilderContext, expression string,
 	}
 }
 
+// make a delim-separated array
+func mockArray(args ...string) string {
+	var sb strings.Builder
+	for i := 0; i < len(args); i++ {
+		if i > 0 {
+			sb.WriteRune(ArraySeparator)
+		}
+		sb.WriteString(args[i])
+	}
+	return sb.String()
+}
+
 func TestCoalesce(t *testing.T) {
 	testExpression(t,
 		mockContext("", "a", "b"),

pkg/expressions/stdlib/range.go

@@ -6,35 +6,39 @@ import (
 	"strings"
 )
 
+// helper context to allow evaluating a sub-expression within a new context
 type subContext struct {
-	parent     KeyBuilderContext
-	val0, val1 string
+	parent KeyBuilderContext
+	vals   [2]string
 }
 
 var _ KeyBuilderContext = &subContext{}
 
 func (s *subContext) GetMatch(idx int) string {
-	if idx == 0 {
-		return s.val0
+	if idx < len(s.vals) {
+		return s.vals[idx]
 	}
-	if idx == 1 {
-		return s.val1
-	}
-	return s.parent.GetMatch(idx)
+	return ""
 }
 
 func (s *subContext) GetKey(k string) string {
 	return s.parent.GetKey(k)
 }
 
 func (s *subContext) Eval(stage KeyBuilderStage, v0, v1 string) string {
-	s.val0, s.val1 = v0, v1
+	s.vals[0] = v0
+	s.vals[1] = v1
+
 	return stage(s)
 }
 
-// {split <string> "delim"}
+// {$split <string> "delim"}
 func kfArraySplit(args []KeyBuilderStage) KeyBuilderStage {
-	byVal := EvalStageOrDefault(args[1], " ")
+	if !isArgCountBetween(args, 1, 2) {
+		return stageLiteral(ErrorArgCount)
+	}
+
+	byVal := EvalStageIndexOrDefault(args, 1, " ")
 	return func(context KeyBuilderContext) string {
 		return arrayOperator(
 			args[0](context),
@@ -45,8 +49,12 @@ func kfArraySplit(args []KeyBuilderStage) KeyBuilderStage {
 	}
 }
 
-// {join <array> "by"}
+// {$join <array> "by"}
 func kfArrayJoin(args []KeyBuilderStage) KeyBuilderStage {
+	if !isArgCountBetween(args, 1, 2) {
+		return stageLiteral(ErrorArgCount)
+	}
+
 	delim := EvalStageIndexOrDefault(args, 1, " ")
 	return func(context KeyBuilderContext) string {
 		return arrayOperator(
@@ -58,8 +66,12 @@ func kfArrayJoin(args []KeyBuilderStage) KeyBuilderStage {
 	}
 }
 
-// {map <arr> <mapFunc>}
+// {$map <arr> <mapFunc>}
 func kfArrayMap(args []KeyBuilderStage) KeyBuilderStage {
+	if len(args) != 2 {
+		return stageLiteral(ErrorArgCount)
+	}
+
 	return func(context KeyBuilderContext) string {
 		mapperContext := subContext{
 			parent: context,
@@ -75,8 +87,12 @@ func kfArrayMap(args []KeyBuilderStage) KeyBuilderStage {
 	}
 }
 
-// {reduce <arr> <reducer>}
+// {$reduce <arr> <reducer>}
 func kfArrayReduce(args []KeyBuilderStage) KeyBuilderStage {
+	if len(args) != 2 {
+		return stageLiteral(ErrorArgCount)
+	}
+
 	return func(context KeyBuilderContext) string {
 		mapperContext := subContext{
 			parent: context,
@@ -87,18 +103,17 @@ func kfArrayReduce(args []KeyBuilderStage) KeyBuilderStage {
 			Delim: ArraySeparatorString,
 		}
 
-		mapperContext.val0 = splitter.Next()
+		memo := splitter.Next()
 		for !splitter.Done() {
-			mapperContext.val1 = splitter.Next()
-			mapperContext.val0 = args[1](&mapperContext)
+			memo = mapperContext.Eval(args[1], memo, splitter.Next())
 		}
-		return mapperContext.val0
+		return memo
 	}
 }
 
 // {slice <arr> start len}
 func kfArraySlice(args []KeyBuilderStage) KeyBuilderStage {
-	if len(args) < 2 {
+	if !isArgCountBetween(args, 2, 3) {
 		return stageLiteral(ErrorArgCount)
 	}
 
@@ -132,6 +147,10 @@ func kfArraySlice(args []KeyBuilderStage) KeyBuilderStage {
 
 // {filter <arr> <truthy-statement>}
 func kfArrayFilter(args []KeyBuilderStage) KeyBuilderStage {
+	if len(args) != 2 {
+		return stageLiteral(ErrorArgCount)
+	}
+
 	return func(context KeyBuilderContext) string {
 		var sb strings.Builder
 

pkg/expressions/stdlib/range_test.go

@@ -5,7 +5,7 @@ import "testing"
 func TestArrayMap(t *testing.T) {
 	testExpression(
 		t,
-		mockContext(makeArray("joe", "is", "cool")),
+		mockContext(mockArray("joe", "is", "cool")),
 		`{$join {$map {0} "{0}bob"} ", "}`,
 		"joebob, isbob, coolbob",
 	)
@@ -32,7 +32,7 @@ func TestArraySlice(t *testing.T) {
 func TestArrayFilter(t *testing.T) {
 	testExpression(
 		t,
-		mockContext(makeArray("a", "123", "b", "455")),
+		mockContext(mockArray("a", "123", "b", "455")),
 		`{$join {$filter {0} "{isnum {0}}"}}`,
 		"123 455",
 	)

pkg/expressions/stdlib/util.go

@@ -2,7 +2,6 @@ package stdlib
 
 import (
 	"rare/pkg/expressions"
-	"strings"
 )
 
 // Checks if word starts with s
@@ -20,14 +19,7 @@ func isPartialString(s, word string) bool {
 	return true
 }
 
-// make a delim-separated array
-func makeArray(args ...string) string {
-	var sb strings.Builder
-	for i := 0; i < len(args); i++ {
-		if i > 0 {
-			sb.WriteRune(expressions.ArraySeparator)
-		}
-		sb.WriteString(args[i])
-	}
-	return sb.String()
+// Helper to check if number of arguments is between a min and max
+func isArgCountBetween(args []expressions.KeyBuilderStage, min, max int) bool {
+	return len(args) >= min && len(args) <= max
 }