From 6df4157bffbfbf57c28c795c4a2d353ed2313cc7 Mon Sep 17 00:00:00 2001
From: swh <sokwhan@google.com>
Date: Thu, 22 Aug 2024 14:15:54 -0700
Subject: [PATCH] Add sets and lists extensions to REPL (#1005)

---
 repl/evaluator.go      | 49 +++++++++++++++++++++++++-----------------
 repl/evaluator_test.go |  2 +-
 2 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/repl/evaluator.go b/repl/evaluator.go
index 4a9cc416..a5f8814e 100644
--- a/repl/evaluator.go
+++ b/repl/evaluator.go
@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"sort"
 	"strings"
 
 	"github.com/google/cel-go/cel"
@@ -41,6 +42,19 @@ import (
 	descpb "google.golang.org/protobuf/types/descriptorpb"
 )
 
+var (
+	extensionMap = map[string]cel.EnvOption{
+		"optional": cel.OptionalTypes(),
+		"bindings": ext.Bindings(),
+		"strings":  ext.Strings(),
+		"protos":   ext.Protos(),
+		"math":     ext.Math(),
+		"encoders": ext.Encoders(),
+		"sets":     ext.Sets(),
+		"lists":    ext.Lists(),
+	}
+)
+
 // letVariable let variable representation
 type letVariable struct {
 	identifier string
@@ -723,26 +737,22 @@ func (o extensionOption) Option() cel.EnvOption {
 }
 
 func newExtensionOption(extType string) (*extensionOption, error) {
-	var extOption cel.EnvOption
 	extType = strings.ToLower(extType)
-	switch op := extType; op {
-	case "bindings":
-		extOption = ext.Bindings()
-	case "optional":
-		extOption = cel.OptionalTypes()
-	case "strings":
-		extOption = ext.Strings()
-	case "protos":
-		extOption = ext.Protos()
-	case "math":
-		extOption = ext.Math()
-	case "encoders":
-		extOption = ext.Encoders()
-	default:
-		return nil, fmt.Errorf("Unknown option: %s. Available options are: ['strings', 'protos', 'math', 'encoders', 'bindings', 'optional', 'all']", op)
+	if extOption, found := extensionMap[extType]; found {
+		return &extensionOption{extensionType: extType, option: extOption}, nil
+	} else {
+		keys := make([]string, 0)
+		for k, _ := range extensionMap {
+			keys = append(keys, k)
+		}
+		sort.Strings(keys)
+		extKeyName := make([]string, 0, len(keys))
+		for _, k := range keys {
+			extKeyName = append(extKeyName, "'"+k+"'")
+		}
+		joinedOptions := "['all', " + strings.Join(extKeyName, ", ") + "]"
+		return nil, fmt.Errorf("Unknown option: %s. Available options are: %s", extType, joinedOptions)
 	}
-
-	return &extensionOption{extensionType: extType, option: extOption}, nil
 }
 
 // setOption sets a number of options on the environment. returns an error if
@@ -811,8 +821,7 @@ func (e *Evaluator) loadExtensionOption(idx int, args []string) error {
 	argExtType := args[idx]
 	if argExtType == "all" {
 		// Load all extension types as a convenience
-		var extensionTypes = []string{"optional", "strings", "protos", "math", "encoders", "bindings"}
-		for _, val := range extensionTypes {
+		for val, _ := range extensionMap {
 			err := e.loadExtensionOptionType(val)
 			if err != nil {
 				return err
diff --git a/repl/evaluator_test.go b/repl/evaluator_test.go
index 4e49603d..9f3e2a31 100644
--- a/repl/evaluator_test.go
+++ b/repl/evaluator_test.go
@@ -1040,7 +1040,7 @@ func TestProcessOptionError(t *testing.T) {
 					"'bogus'",
 				},
 			},
-			errorMsg: "extension: Unknown option: 'bogus'. Available options are: ['strings', 'protos', 'math', 'encoders', 'bindings', 'optional', 'all']",
+			errorMsg: "extension: Unknown option: 'bogus'. Available options are: ['all', 'bindings', 'encoders', 'lists', 'math', 'optional', 'protos', 'sets', 'strings']",
 		},
 	}