add math module

cmd/starlark/starlark.go

@@ -16,6 +16,7 @@ import (
 	"strings"
 
 	"go.starlark.net/internal/compile"
+	"go.starlark.net/lib/math"
 	"go.starlark.net/lib/time"
 	"go.starlark.net/repl"
 	"go.starlark.net/resolve"
@@ -93,6 +94,7 @@ func doMain() int {
 	// TODO(adonovan): plumb predeclared env through to the REPL.
 	starlark.Universe["json"] = starlarkjson.Module
 	starlark.Universe["time"] = time.Module
+	starlark.Universe["math"] = math.Module
 
 	switch {
 	case flag.NArg() == 1 || *execprog != "":

lib/math/math.go

@@ -0,0 +1,100 @@
+package math
+
+import (
+	"math"
+
+	"go.starlark.net/starlark"
+	"go.starlark.net/starlarkstruct"
+)
+
+const tau = math.Pi * 2
+const oneRad = tau / 360
+
+// Module math is a Starlark module of math-related functions.
+var Module = &starlarkstruct.Module{
+	Name: "math",
+	Members: starlark.StringDict{
+		"ceil":  starlark.NewBuiltin("ceil", ceil),
+		"fabs":  starlark.NewBuiltin("fabs", fabs),
+		"floor": starlark.NewBuiltin("floor", floor),
+		"round": starlark.NewBuiltin("round", round),
+
+		"exp":  starlark.NewBuiltin("exp", exp),
+		"sqrt": starlark.NewBuiltin("sqrt", sqrt),
+
+		"acos":  starlark.NewBuiltin("acos", acos),
+		"asin":  starlark.NewBuiltin("asin", asin),
+		"atan":  starlark.NewBuiltin("atan", atan),
+		"atan2": starlark.NewBuiltin("atan2", atan2),
+		"cos":   starlark.NewBuiltin("cos", cos),
+		"hypot": starlark.NewBuiltin("hypot", hypot),
+		"sin":   starlark.NewBuiltin("sin", sin),
+		"tan":   starlark.NewBuiltin("tan", tan),
+
+		"degrees": starlark.NewBuiltin("degrees", degrees),
+		"radians": starlark.NewBuiltin("radians", radians),
+
+		"acosh": starlark.NewBuiltin("acosh", acosh),
+		"asinh": starlark.NewBuiltin("asinh", asinh),
+		"atanh": starlark.NewBuiltin("atanh", atanh),
+		"cosh":  starlark.NewBuiltin("cosh", cosh),
+		"sinh":  starlark.NewBuiltin("sinh", sinh),
+		"tanh":  starlark.NewBuiltin("tanh", tanh),
+
+		"e":   starlark.Float(math.E),
+		"pi":  starlark.Float(math.Pi),
+		"tau": starlark.Float(tau),
+		"phi": starlark.Float(math.Phi),
+		"inf": starlark.Float(math.Inf(1)),
+		"nan": starlark.Float(math.NaN()),
+	},
+}
+
+// floatFunc unpacks a starlark function call, calls a passed in float64 function
+// and returns the result as a starlark value
+func floatFunc(name string, args starlark.Tuple, kwargs []starlark.Tuple, fn func(float64) float64) (starlark.Value, error) {
+	var x starlark.Float
+	if err := starlark.UnpackPositionalArgs(name, args, kwargs, 1, &x); err != nil {
+		return nil, err
+	}
+	return starlark.Float(fn(float64(x))), nil
+}
+
+// floatFunc2 is a 2-argument float func
+func floatFunc2(name string, args starlark.Tuple, kwargs []starlark.Tuple, fn func(float64, float64) float64) (starlark.Value, error) {
+	var x, y starlark.Float
+	if err := starlark.UnpackPositionalArgs(name, args, kwargs, 2, &x, &y); err != nil {
+		return nil, err
+	}
+	return starlark.Float(fn(float64(x), float64(y))), nil
+}
+
+// Return the floor of x, the largest integer less than or equal to x.
+func floor(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("floor", args, kwargs, math.Floor)
+}
+
+// Return the ceiling of x, the smallest integer greater than or equal to x
+func ceil(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("ceil", args, kwargs, math.Ceil)
+}
+
+// Return the absolute value of x
+func fabs(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("fabs", args, kwargs, math.Abs)
+}
+
+// Round returns the nearest integer, rounding half away from zero.
+func round(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("fabs", args, kwargs, math.Round)
+}
+
+// Return e raised to the power x, where e = 2.718281… is the base of natural logarithms.
+func exp(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("exp", args, kwargs, math.Exp)
+}
+
+// Return the square root of x
+func sqrt(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("sqrt", args, kwargs, math.Sqrt)
+}

lib/math/trig.go

@@ -0,0 +1,92 @@
+package math
+
+import (
+	"math"
+
+	"go.starlark.net/starlark"
+)
+
+// Return the arc cosine of x, in radians.
+func acos(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("acos", args, kwargs, math.Acos)
+}
+
+// asin(x) - Return the arc sine of x, in radians.
+func asin(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("asin", args, kwargs, math.Asin)
+}
+
+// atan(x) - Return the arc tangent of x, in radians.
+func atan(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("atan", args, kwargs, math.Atan)
+}
+
+// atan2(y, x) - Return atan(y / x), in radians. The result is between -pi and pi.
+// The vector in the plane from the origin to point (x, y) makes this angle with the positive X axis.
+// The point of atan2() is that the signs of both inputs are known to it, so it can compute the correct quadrant for the angle.
+// For example, atan(1) and atan2(1, 1) are both pi/4, but atan2(-1, -1) is -3*pi/4.
+func atan2(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc2("atan2", args, kwargs, math.Atan2)
+}
+
+// cos(x) - Return the cosine of x radians.
+func cos(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("cos", args, kwargs, math.Cos)
+}
+
+// hypot(x, y) - Return the Euclidean norm, sqrt(x*x + y*y). This is the length of the vector from the origin to point (x, y).
+func hypot(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc2("hypot", args, kwargs, math.Hypot)
+}
+
+// sin(x) - Return the sine of x radians.
+func sin(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("sin", args, kwargs, math.Sin)
+}
+
+// tan(x) - Return the tangent of x radians.
+func tan(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("tan", args, kwargs, math.Tan)
+}
+
+// degrees(x) - Convert angle x from radians to degrees.
+func degrees(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	toDeg := func(x float64) float64 { return x / oneRad }
+	return floatFunc("degrees", args, kwargs, toDeg)
+}
+
+// radians(x) - Convert angle x from degrees to radians.
+func radians(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	toRad := func(x float64) float64 { return x * oneRad }
+	return floatFunc("radians", args, kwargs, toRad)
+}
+
+// acosh(x) - Return the inverse hyperbolic cosine of x.
+func acosh(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("acosh", args, kwargs, math.Acosh)
+}
+
+// asinh(x) - Return the inverse hyperbolic sine of x.
+func asinh(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("asinh", args, kwargs, math.Asinh)
+}
+
+// atanh(x) - Return the inverse hyperbolic tangent of x.
+func atanh(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("atanh", args, kwargs, math.Atanh)
+}
+
+// cosh(x) - Return the hyperbolic cosine of x.
+func cosh(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("cosh", args, kwargs, math.Cosh)
+}
+
+// sinh(x) - Return the hyperbolic sine of x.
+func sinh(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("sinh", args, kwargs, math.Sinh)
+}
+
+// tanh(x) - Return the hyperbolic tangent of x.
+func tanh(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	return floatFunc("tanh", args, kwargs, math.Tanh)
+}

starlark/eval_test.go

@@ -16,6 +16,7 @@ import (
 	"testing"
 
 	"go.starlark.net/internal/chunkedfile"
+	starlarkmath "go.starlark.net/lib/math"
 	"go.starlark.net/lib/time"
 	"go.starlark.net/resolve"
 	"go.starlark.net/starlark"
@@ -124,6 +125,7 @@ func TestExecFile(t *testing.T) {
 		"testdata/int.star",
 		"testdata/json.star",
 		"testdata/list.star",
+		"testdata/math.star",
 		"testdata/misc.star",
 		"testdata/set.star",
 		"testdata/string.star",
@@ -198,6 +200,9 @@ func load(thread *starlark.Thread, module string) (starlark.StringDict, error) {
 	if module == "time.star" {
 		return starlark.StringDict{"time": time.Module}, nil
 	}
+	if module == "math.star" {
+		return starlark.StringDict{"math": starlarkmath.Module}, nil
+	}
 
 	// TODO(adonovan): test load() using this execution path.
 	filename := filepath.Join(filepath.Dir(thread.CallFrame(0).Pos.Filename()), module)

starlark/testdata/math.star

@@ -0,0 +1,43 @@
+# Tests of math module.
+
+load('math.star', 'math')
+load('assert.star', 'assert')
+
+def assert_near(actual, desired, threshold):
+  val = desired - actual
+  if val < 0:
+    val = val * -1
+  return val < threshold
+
+assert.eq(math.ceil(0.5), 1.0)
+assert.eq(math.fabs(-2.0), 2.0)
+assert.eq(math.floor(0.5), 0)
+
+assert.eq(math.exp(2.0), 7.38905609893065)
+assert.eq(math.sqrt(4.0), 2.0)
+
+assert.eq(math.acos(1.0), 0)
+assert.eq(math.asin(1.0), 1.5707963267948966)
+assert.eq(math.atan(1.0), 0.7853981633974483)
+assert.eq(math.atan2(1.0,1.0), 0.7853981633974483)
+assert.eq(math.cos(1.0), 0.5403023058681398)
+assert.eq(math.hypot(1.0,1.0), 1.4142135623730951)
+assert.eq(math.sin(1.0), 0.8414709848078965)
+assert.eq(math.tan(1.0), 1.557407724654902)
+
+assert.eq(math.degrees(1.0), 57.29577951308232)
+assert.eq(math.radians(1.0), 0.017453292519943295)
+
+assert.eq(math.acosh(1.0), 0)
+assert.eq(assert_near(math.asinh(1.0), 0.8813735870195432, 0.00000001), True)
+assert.eq(math.atanh(0.5), 0.5493061443340548)
+assert.eq(math.cosh(1.0), 1.5430806348152437)
+assert.eq(math.sinh(1.0), 1.1752011936438014)
+assert.eq(math.tanh(1.0), 0.7615941559557649)
+
+assert.eq(math.e, 2.7182818284590452)
+assert.eq(math.pi, 3.1415926535897932)
+assert.eq(math.tau, 6.2831853071795864)
+assert.eq(math.phi, 1.6180339887498948)
+assert.eq(math.inf, math.inf + 1)
+assert.eq(math.nan == math.nan, False)