Repository: abs-lang/abs
Pull request #353 (2020-03-28 13:31:18)

array.partition(fn)
Partitions the array by applying `fn` to all of its elements:

```py
f odd(n) {
  return !!(n % 2)
}

[0, 1, 2, 3, 4, 5].partition(odd) # [[0, 2, 4], [1, 3, 5]]
[1, "1", {}].partition(str) # [[1, "1"], [{}]]
```


Patch:
diff --git a/docs/types/array.md b/docs/types/array.md
index 9399c6f..b25044c 100644
--- a/docs/types/array.md
+++ b/docs/types/array.md
@@ -459,6 +459,21 @@ Reduces the array to a value by iterating through its elements and applying `fn`
 [1, 2, 3, 4].reduce(f(value, element) { return value + element }, 10) # 20
 ```
 
+
+### partition(fn)
+
+Partitions the array by applying `fn` to all of its elements
+and using the result of the function invocation as the key to partition by:
+
+```py
+f odd(n) {
+  return !!(n % 2)
+}
+
+[0, 1, 2, 3, 4, 5].partition(odd) # [[0, 2, 4], [1, 3, 5]]
+[1, "1", {}].partition(str) # [[1, "1"], [{}]]
+```
+
 ## Next
 
 That's about it for this section!
diff --git a/evaluator/functions.go b/evaluator/functions.go
index b17ee5d..f2ca697 100644
--- a/evaluator/functions.go
+++ b/evaluator/functions.go
@@ -233,16 +233,21 @@ func getFns() map[string]*object.Builtin {
 			Types: []string{object.ARRAY_OBJ},
 			Fn:    diffSymmetricFn,
 		},
-		// flattem(array:[1, 2, 3])
+		// flatten(array:[1, 2, 3])
 		"flatten": &object.Builtin{
 			Types: []string{object.ARRAY_OBJ},
 			Fn:    flattenFn,
 		},
-		// flattem(array:[1, 2, 3])
+		// flatten(array:[1, 2, 3])
 		"flatten_deep": &object.Builtin{
 			Types: []string{object.ARRAY_OBJ},
 			Fn:    flattenDeepFn,
 		},
+		// partition(array:[1, 2, 3])
+		"partition": &object.Builtin{
+			Types: []string{object.ARRAY_OBJ},
+			Fn:    partitionFn,
+		},
 		// map(array:[1, 2, 3], function:f(x) { x + 1 })
 		"map": &object.Builtin{
 			Types: []string{object.ARRAY_OBJ},
@@ -1407,7 +1412,7 @@ func flattenDeepFn(tok token.Token, env *object.Environment, args ...object.Obje
 }
 
 func flatten(fnName string, deep bool, tok token.Token, env *object.Environment, args ...object.Object) object.Object {
-	err := validateArgs(tok, "flatten", args, 1, [][]string{{object.ARRAY_OBJ}})
+	err := validateArgs(tok, fnName, args, 1, [][]string{{object.ARRAY_OBJ}})
 	if err != nil {
 		return err
 	}
@@ -1433,6 +1438,49 @@ func flatten(fnName string, deep bool, tok token.Token, env *object.Environment,
 	return &object.Array{Elements: elements}
 }
 
+func partitionFn(tok token.Token, env *object.Environment, args ...object.Object) object.Object {
+	err := validateArgs(tok, "partition", args, 2, [][]string{{object.ARRAY_OBJ}, {object.FUNCTION_OBJ, object.BUILTIN_OBJ}})
+	if err != nil {
+		return err
+	}
+
+	partitions := map[string][]object.Object{}
+	elements := args[0].(*object.Array).Elements
+	// This will allows us to preserve the order
+	// of partitions based on the order of elements.
+	//
+	// When we run the partitioning function, we store
+	// it's results in a map of result{list_of_values...}.
+	// When we loop over that map, Go doesn't guarantee
+	// order of results (https://nathanleclaire.com/blog/2014/04/27/a-surprising-feature-of-golang-that-colored-me-impressed/)
+	// but we want to, so
+	// we use the partitionOrder list to extract values
+	// from the map based on the order they were
+	// inserted in.
+	partitionOrder := []string{}
+	scanned := map[string]bool{}
+
+	for _, v := range elements {
+		res := applyFunction(tok, args[1], env, []object.Object{v})
+		eqs := object.GenerateEqualityString(res)
+
+		partitions[eqs] = append(partitions[eqs], v)
+
+		if _, ok := scanned[eqs]; !ok {
+			partitionOrder = append(partitionOrder, eqs)
+			scanned[eqs] = true
+		}
+	}
+
+	result := &object.Array{Elements: []object.Object{}}
+	for _, eqs := range partitionOrder {
+		partition := partitions[eqs]
+		result.Elements = append(result.Elements, &object.Array{Elements: partition})
+	}
+
+	return result
+}
+
 // map(array:[1, 2, 3], function:f(x) { x + 1 })
 func mapFn(tok token.Token, env *object.Environment, args ...object.Object) object.Object {
 	err := validateArgs(tok, "map", args, 2, [][]string{{object.ARRAY_OBJ}, {object.FUNCTION_OBJ, object.BUILTIN_OBJ}})


Test patch:
diff --git a/evaluator/builtin_functions_test.go b/evaluator/builtin_functions_test.go
index dbcd2f7..46b960c 100644
--- a/evaluator/builtin_functions_test.go
+++ b/evaluator/builtin_functions_test.go
@@ -725,6 +725,17 @@ func TestReduce(t *testing.T) {
 	testBuiltinFunction(tests, t)
 }
 
+func TestPartition(t *testing.T) {
+	tests := []Tests{
+		{`[1, 1, 2, 2, 0].partition(f(x) { return x == 0 })[0]`, []int{1, 1, 2, 2}},
+		{`[1, 1, 2, 2, 0].partition(f(x) { return x == 0 })[1]`, []int{0}},
+		{`[1, "1"].partition(str)[0][0]`, 1},
+		{`[1, "1"].partition(str)[0][1]`, "1"},
+	}
+
+	testBuiltinFunction(tests, t)
+}
+
 func testBuiltinFunction(tests []Tests, t *testing.T) {
 	for _, tt := range tests {
 		evaluated := testEval(tt.input)