google · TristonianJones · Oct 22, 2024 · Oct 21, 2024
diff --git a/ext/math.go b/ext/math.go
@@ -325,8 +325,12 @@ import (
 //
 //	math.isFinite(0.0/0.0)  // returns false
 //	math.isFinite(1.2)      // returns true
-func Math() cel.EnvOption {
-	return cel.Lib(&mathLib{version: math.MaxUint32})
+func Math(options ...MathOption) cel.EnvOption {
+	m := &mathLib{version: math.MaxUint32}
+	for _, o := range options {
+		m = o(m)
+	}
+	return cel.Lib(m)
 }
 
 const (
@@ -366,8 +370,10 @@ var (
 	errIntOverflow = types.NewErr("integer overflow")
 )
 
+// MathOption declares a functional operator for configuring math extensions.
 type MathOption func(*mathLib) *mathLib
 
+// MathVersion sets the library version for math extensions.
 func MathVersion(version uint32) MathOption {
 	return func(lib *mathLib) *mathLib {
 		lib.version = version

diff --git a/ext/math_test.go b/ext/math_test.go
@@ -20,6 +20,7 @@ import (
 	"testing"
 
 	"github.com/google/cel-go/cel"
+	"github.com/google/cel-go/common/types"
 )
 
 func TestMath(t *testing.T) {
@@ -566,6 +567,78 @@ func TestMathWithExtension(t *testing.T) {
 	}
 }
 
+func TestMathVersions(t *testing.T) {
+	versionCases := []struct {
+		version            uint32
+		supportedFunctions map[string]string
+	}{
+		{
+			version: 0,
+			supportedFunctions: map[string]string{
+				"greatest": `math.greatest(1, 2) == 2`,
+				"least":    `math.least(2.1, -1.0) == -1.0`,
+			},
+		},
+		{
+			version: 1,
+			supportedFunctions: map[string]string{
+				"ceil":          `math.ceil(1.5) == 2.0`,
+				"floor":         `math.floor(1.2) == 1.0`,
+				"round":         `math.round(1.5) == 2.0`,
+				"trunc":         `math.trunc(1.222) == 1.0`,
+				"isInf":         `!math.isInf(0.0)`,
+				"isNaN":         `math.isNaN(0.0/0.0)`,
+				"isFinite":      `math.isFinite(0.0)`,
+				"abs":           `math.abs(1.2) == 1.2`,
+				"sign":          `math.sign(-1) == -1`,
+				"bitAnd":        `math.bitAnd(1, 2) == 0`,
+				"bitOr":         `math.bitOr(1, 2) == 3`,
+				"bitXor":        `math.bitXor(1, 3) == 2`,
+				"bitNot":        `math.bitNot(-1) == 0`,
+				"bitShiftLeft":  `math.bitShiftLeft(4, 2) == 16`,
+				"bitShiftRight": `math.bitShiftRight(4, 2) == 1`,
+			},
+		},
+	}
+	for _, lib := range versionCases {
+		env, err := cel.NewEnv(Math(MathVersion(lib.version)))
+		if err != nil {
+			t.Fatalf("cel.NewEnv(Math(MathVersion(%d))) failed: %v", lib.version, err)
+		}
+		t.Run(fmt.Sprintf("version=%d", lib.version), func(t *testing.T) {
+			for _, tc := range versionCases {
+				for name, expr := range tc.supportedFunctions {
+					supported := lib.version >= tc.version
+					t.Run(fmt.Sprintf("%s-supported=%t", name, supported), func(t *testing.T) {
+						ast, iss := env.Compile(expr)
+						if supported {
+							if iss.Err() != nil {
+								t.Errorf("unexpected error: %v", iss.Err())
+							}
+						} else {
+							if iss.Err() == nil || !strings.Contains(iss.Err().Error(), "undeclared reference") {
+								t.Errorf("got error %v, wanted error %s for expr: %s, version: %d", iss.Err(), "undeclared reference", expr, tc.version)
+							}
+							return
+						}
+						prg, err := env.Program(ast)
+						if err != nil {
+							t.Fatalf("env.Program() failed: %v", err)
+						}
+						out, _, err := prg.Eval(cel.NoVars())
+						if err != nil {
+							t.Fatalf("prg.Eval() failed: %v", err)
+						}
+						if out != types.True {
+							t.Errorf("prg.Eval() got %v, wanted true", out)
+						}
+					})
+				}
+			}
+		})
+	}
+}
+
 func testMathEnv(t *testing.T, opts ...cel.EnvOption) *cel.Env {
 	t.Helper()
 	baseOpts := []cel.EnvOption{Math(), cel.EnableMacroCallTracking()}