starlark: add ** operator, **= augmented assignment, and pow() (#631)

doc/spec.md

@@ -100,7 +100,7 @@ characters are tokens:
 +=   -=   *=   /=   //=  %=   ==   !=
 ^    <    >    <<   >>   &    |
 ^=   <=   >=   <<=  >>=  &=   |=
-.    ,    ;    :    ~    **
+.    ,    ;    :    ~    **   **=
 (    )    [    ]    {    }
 ```
 
@@ -1703,11 +1703,14 @@ There are three unary operators, all appearing before their operand:
 `+`, `-`, `~`, and `not`.
 
 ```grammar {.good}
-UnaryExpr = '+' PrimaryExpr
-          | '-' PrimaryExpr
-          | '~' PrimaryExpr
+UnaryExpr = '+' UnaryExpr
+          | '-' UnaryExpr
+          | '~' UnaryExpr
+          | PowerExpr
           | 'not' Test
           .
+
+PowerExpr = PrimaryExpr ['**' UnaryExpr] .
 ```
 
 ```text
@@ -1717,6 +1720,10 @@ UnaryExpr = '+' PrimaryExpr
 not x           logical negation        (any type)
 ```
 
+The `+`, `-`, and `~` operators bind less tightly than `**`, so
+`-x ** y` is parsed as `-(x ** y)`. However, they may appear on the
+right side of `**`, so `x ** -y` is parsed as `x ** (-y)`.
+
 The `+` and `-` operators may be applied to any number
 (`int` or `float`) and return the number unchanged.
 Unary `+` is never necessary in a correct program,
@@ -1767,11 +1774,13 @@ and
 <<   >>
 -    +
 *    /    //   %
+**
 ```
 
 Comparison operators, `in`, and `not in` are non-associative,
 so the parser will not accept `0 <= i < n`.
-All other binary operators of equal precedence associate to the left.
+The `**` operator is right-associative; all other binary operators of
+equal precedence associate to the left.
 
 ```grammar {.good}
 BinaryExpr = Test {Binop Test} .
@@ -1788,6 +1797,10 @@ Binop = 'or'
       .
 ```
 
+Note: the `**` operator is not listed in `Binop` because it has
+special parsing rules (right-associativity and interaction with unary
+operators); see `PowerExpr` under [Unary operators](#unary-operators).
+
 #### `or` and `and`
 
 The `or` and `and` operators yield, respectively, the logical disjunction and
@@ -1888,6 +1901,7 @@ Arithmetic (int or float; result has type float unless both operands have type i
    number / number              # real division  (result is always a float)
    number // number             # floored division
    number % number              # remainder of floored division
+   number ** number             # exponentiation
    number ^ number              # bitwise XOR
    number << number             # bitwise left shift
    number >> number             # bitwise right shift
@@ -1917,11 +1931,27 @@ Dict
       dict | dict               # ordered union
 ```
 
-The operands of the arithmetic operators `+`, `-`, `*`, `//`, and
-`%` must both be numbers (`int` or `float`) but need not have the same type.
+The operands of the arithmetic operators `+`, `-`, `*`, `//`, `%`,
+and `**` must both be numbers (`int` or `float`) but need not have the same type.
 The type of the result has type `int` only if both operands have that type.
 The result of real division `/` always has type `float`.
 
+The `**` operator computes exponentiation.
+It is right-associative: `2 ** 3 ** 2` is equal to `2 ** 9`, not `8 ** 2`.
+When both operands are integers and the exponent is non-negative, the
+result is an exact integer.
+When the exponent is negative or either operand is a float, the result
+is a float.
+It is an error to raise a negative number to a fractional power (the
+result would be complex) or to raise zero to a negative power.
+
+```python
+2 ** 10                 # 1024
+2 ** -1                 # 0.5
+4.0 ** 0.5              # 2.0
+2 ** 3 ** 2             # 512
+```
+
 The `+` operator may be applied to non-numeric operands of the same
 type, such as two lists, two tuples, or two strings, in which case it
 computes the concatenation of the two operands and yields a new value of
@@ -2510,11 +2540,11 @@ in `for` loops and in comprehensions.
 
 An augmented assignment, which has the form `lhs op= rhs` updates the
 variable `lhs` by applying a binary arithmetic operator `op` (one of
-`+`, `-`, `*`, `/`, `//`, `%`, `&`, `|`, `^`, `<<`, `>>`) to the previous
+`+`, `-`, `*`, `/`, `//`, `%`, `&`, `|`, `^`, `<<`, `>>`, `**`) to the previous
 value of `lhs` and the value of `rhs`.
 
 ```grammar {.good}
-AssignStmt = Expression ('+=' | '-=' | '*=' | '/=' | '//=' | '%=' | '&=' | '|=' | '^=' | '<<=' | '>>=') Expression .
+AssignStmt = Expression ('+=' | '-=' | '*=' | '/=' | '//=' | '%=' | '&=' | '|=' | '^=' | '<<=' | '>>=' | '**=') Expression .
 ```
 
 The left-hand side must be a simple target:
@@ -3190,6 +3220,28 @@ See also: `chr`.
 
 <b>Implementation note:</b> `ord` is not provided by the Java implementation.
 
+### pow
+
+`pow(base, exp)` returns `base` raised to the power `exp`.
+It is equivalent to the expression `base ** exp`.
+The arguments must be numbers (`int` or `float`).
+
+With three arguments, `pow(base, exp, mod)` computes modular
+exponentiation: it returns `base**exp mod mod`.
+The three-argument form requires all arguments to be integers and the
+modulus to be nonzero. When `exp` is negative, it computes the modular
+inverse of `base` with respect to `mod`, which requires that `base`
+and `mod` be coprime.
+
+```python
+pow(2, 10)              # 1024
+pow(2, -1)              # 0.5
+pow(2, 10, 1000)        # 24
+pow(3, -1, 7)           # 5 (modular inverse: 3*5 ≡ 1 mod 7)
+```
+
+<b>Implementation note:</b> `pow` is not provided by the Java implementation.
+
 ### print
 
 `print(*args, sep=" ")` prints its arguments, followed by a newline.

internal/compile/compile.go

@@ -86,6 +86,7 @@ const (
 	CIRCUMFLEX
 	LTLT
 	GTGT
+	STARSTAR
 
 	IN
 
@@ -215,6 +216,7 @@ var opcodeNames = [...]string{
 	SLASHSLASH:   "slashslash",
 	SLICE:        "slice",
 	STAR:         "star",
+	STARSTAR:     "starstar",
 	TILDE:        "tilde",
 	TRUE:         "true",
 	UMINUS:       "uminus",
@@ -289,6 +291,7 @@ var stackEffect = [...]int8{
 	SLASHSLASH:   -1,
 	SLICE:        -3,
 	STAR:         -1,
+	STARSTAR:     -1,
 	TRUE:         +1,
 	UMINUS:       0,
 	UNIVERSAL:    +1,
@@ -1107,7 +1110,8 @@ func (fcomp *fcomp) stmt(stmt syntax.Stmt) {
 			syntax.PIPE_EQ,
 			syntax.CIRCUMFLEX_EQ,
 			syntax.LTLT_EQ,
-			syntax.GTGT_EQ:
+			syntax.GTGT_EQ,
+			syntax.STARSTAR_EQ:
 			// augmented assignment: x += y
 
 			var set func()
@@ -1627,6 +1631,8 @@ func (fcomp *fcomp) binop(pos syntax.Position, op syntax.Token) {
 		fcomp.emit(LTLT)
 	case syntax.GTGT:
 		fcomp.emit(GTGT)
+	case syntax.STARSTAR:
+		fcomp.emit(STARSTAR)
 	case syntax.IN:
 		fcomp.emit(IN)
 	case syntax.NOT_IN:

starlark/eval.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"math"
 	"math/big"
 	"math/bits"
 	"sort"
@@ -1024,6 +1025,53 @@ func Binary(op syntax.Token, x, y Value) (Value, error) {
 			return interpolate(string(x), y)
 		}
 
+	case syntax.STARSTAR:
+		switch x := x.(type) {
+		case Int:
+			switch y := y.(type) {
+			case Int:
+				if y.Sign() < 0 {
+					if x.Sign() == 0 {
+						return nil, fmt.Errorf("zero raised to negative power")
+					}
+					xf, err := x.finiteFloat()
+					if err != nil {
+						return nil, err
+					}
+					yf, err := y.finiteFloat()
+					if err != nil {
+						return nil, err
+					}
+					return floatPow(xf, yf)
+				}
+				// y >= 0: integer exponentiation.
+				if x.bigInt().BitLen() > 1 {
+					yInt, err := AsInt32(y)
+					if err != nil || int64(x.bigInt().BitLen())*int64(yInt) > 4096 {
+						return nil, fmt.Errorf("exponent too large")
+					}
+				}
+				return MakeBigInt(new(big.Int).Exp(x.bigInt(), y.bigInt(), nil)), nil
+			case Float:
+				xf, err := x.finiteFloat()
+				if err != nil {
+					return nil, err
+				}
+				return floatPow(xf, y)
+			}
+		case Float:
+			switch y := y.(type) {
+			case Float:
+				return floatPow(x, y)
+			case Int:
+				yf, err := y.finiteFloat()
+				if err != nil {
+					return nil, err
+				}
+				return floatPow(x, yf)
+			}
+		}
+
 	case syntax.NOT_IN:
 		z, err := Binary(syntax.IN, x, y)
 		if err != nil {
@@ -1135,6 +1183,21 @@ unknown:
 	return nil, fmt.Errorf("unknown binary op: %s %s %s", x.Type(), op, y.Type())
 }
 
+// floatPow computes x ** y for float operands.
+func floatPow(x, y Float) (Value, error) {
+	if x == 0 && y < 0 {
+		return nil, fmt.Errorf("zero raised to negative power")
+	}
+	if x < 0 && y != Float(math.Trunc(float64(y))) {
+		return nil, fmt.Errorf("negative number raised to non-integer power")
+	}
+	rf := math.Pow(float64(x), float64(y))
+	if math.IsInf(rf, 0) {
+		return nil, fmt.Errorf("floating-point result too large")
+	}
+	return Float(rf), nil
+}
+
 // It's always possible to overeat in small bites but we'll
 // try to stop someone swallowing the world in one gulp.
 const maxAlloc = 1 << 30

starlark/interp.go

@@ -185,6 +185,7 @@ loop:
 			compile.CIRCUMFLEX,
 			compile.LTLT,
 			compile.GTGT,
+			compile.STARSTAR,
 			compile.IN:
 			binop := syntax.Token(op-compile.PLUS) + syntax.PLUS
 			if op == compile.IN {

starlark/library.go

@@ -59,6 +59,7 @@ func init() {
 		"max":       NewBuiltin("max", minmax),
 		"min":       NewBuiltin("min", minmax),
 		"ord":       NewBuiltin("ord", ord),
+		"pow":       NewBuiltin("pow", pow),
 		"print":     NewBuiltin("print", print),
 		"range":     NewBuiltin("range", range_),
 		"repr":      NewBuiltin("repr", repr),
@@ -795,6 +796,40 @@ func ord(thread *Thread, _ *Builtin, args Tuple, kwargs []Tuple) (Value, error)
 	}
 }
 
+// pow(base, exp[, mod]) — exponentiation.
+//
+// With two arguments, pow(x, y) is equivalent to x ** y.
+// With three arguments, pow(x, y, z) computes x**y mod z.
+// The three-argument form requires all arguments to be integers.
+func pow(thread *Thread, b *Builtin, args Tuple, kwargs []Tuple) (Value, error) {
+	var x, y, z Value
+	if err := UnpackPositionalArgs("pow", args, kwargs, 2, &x, &y, &z); err != nil {
+		return nil, err
+	}
+
+	// Three-argument form: all arguments must be integers.
+	if z != nil {
+		xi, xok := x.(Int)
+		yi, yok := y.(Int)
+		zi, zok := z.(Int)
+		if !xok || !yok || !zok {
+			return nil, fmt.Errorf("pow: 3rd argument not allowed unless all arguments are integers")
+		}
+		if zi.Sign() == 0 {
+			return nil, fmt.Errorf("pow: integer modulo by zero")
+		}
+		// big.Int.Exp returns nil if y < 0 and x and m are not coprime.
+		result := new(big.Int).Exp(xi.bigInt(), yi.bigInt(), zi.bigInt())
+		if result == nil {
+			return nil, fmt.Errorf("pow: base is not invertible for the given modulus")
+		}
+		return MakeBigInt(result), nil
+	}
+
+	// Two-argument form: equivalent to x ** y.
+	return Binary(syntax.STARSTAR, x, y)
+}
+
 // https://github.com/google/starlark-go/blob/master/doc/spec.md#print
 func print(thread *Thread, b *Builtin, args Tuple, kwargs []Tuple) (Value, error) {
 	sep := " "

starlark/testdata/builtins.star

@@ -88,6 +88,39 @@ assert.eq(dict([(1, 2), (3, 4)], foo="bar"), {1: 2, 3: 4, "foo": "bar"})
 assert.eq(dict({1:2, 3:4}), {1: 2, 3: 4})
 assert.eq(dict({1:2, 3:4}.items()), {1: 2, 3: 4})
 
+# pow
+# two-argument form (equivalent to **)
+assert.eq(pow(2, 0), 1)
+assert.eq(pow(2, 10), 1024)
+assert.eq(pow(0, 0), 1)
+assert.eq(pow(-2, 3), -8)
+assert.eq(pow(-2, 2), 4)
+assert.eq(pow(2, -1), 0.5)
+assert.eq(pow(2.0, 3), 8.0)
+assert.eq(pow(4.0, 0.5), 2.0)
+assert.eq(pow(4, 0.5), 2.0)
+assert.eq(type(pow(2, 3)), "int")
+assert.eq(type(pow(2, -1)), "float")
+assert.eq(type(pow(2.0, 3)), "float")
+# three-argument form (modular exponentiation)
+assert.eq(pow(2, 10, 1000), 24)
+assert.eq(pow(2, 10, 3), 1)
+assert.eq(pow(2, 100, 13), 3)
+assert.eq(pow(3, 0, 5), 1)
+assert.eq(pow(100, 2, 7), 4)
+assert.eq(pow(-2, 3, 5), 2)    # (-8) mod 5 = 2
+# three-argument form with negative exponent (modular inverse)
+assert.eq(pow(3, -1, 7), 5)    # 3*5 = 15 ≡ 1 (mod 7)
+assert.eq(pow(2, -1, 7), 4)    # 2*4 = 8 ≡ 1 (mod 7)
+# errors
+assert.fails(lambda: pow(0, -1), "zero raised to negative power")
+assert.fails(lambda: pow(2, 10000), "exponent too large")
+assert.fails(lambda: pow(1.0, 2, 3), "3rd argument not allowed unless all arguments are integers")
+assert.fails(lambda: pow(1, 2.0, 3), "3rd argument not allowed unless all arguments are integers")
+assert.fails(lambda: pow(1, 2, 0), "integer modulo by zero")
+assert.fails(lambda: pow(2, -1, 6), "base is not invertible for the given modulus")
+assert.fails(lambda: pow("a", 2), "unknown binary op: string \\*\\* int")
+
 # range
 assert.eq("range", type(range(10)))
 assert.eq("range(10)", str(range(0, 10, 1)))