int.go

// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package grumpy

import (
	"fmt"
	"math/big"
	"reflect"
	"strconv"
)

const (
	internedIntMin = -2
	internedIntMax = 300
)

var internedInts [internedIntMax - internedIntMin + 1]Int

func init() {
	for i := internedIntMin; i <= internedIntMax; i++ {
		internedInts[i-internedIntMin] = Int{Object{typ: IntType, self: &internedInts[i-internedIntMin]}, i}
	}
}

// Int represents Python 'int' objects.
type Int struct {
	Object
	value int
}

// NewInt returns a new Int holding the given integer value.
func NewInt(value int) *Int {
	if value >= internedIntMin && value <= internedIntMax {
		return &internedInts[value-internedIntMin]
	}
	i := &Int{Object{typ: IntType}, value}
	i.self = i
	return i
}

func toIntUnsafe(o *Object) *Int {
	return o.self.(*Int)
}

// ToObject upcasts i to an Object.
func (i *Int) ToObject() *Object {
	return &i.Object
}

// Value returns the underlying integer value held by i.
func (i *Int) Value() int {
	return i.value
}

// IsTrue returns false if i is zero, true otherwise.
func (i *Int) IsTrue() bool {
	return i.Value() != 0
}

// IntType is the object representing the Python 'int' type.
var IntType = newBasisType("int", reflect.TypeOf(Int{}), toIntUnsafe, ObjectType)

func intAbs(f *Frame, o *Object) (*Object, *BaseException) {
	z := toIntUnsafe(o)
	if z.Value() > 0 {
		return z.ToObject(), nil
	}
	if z.Value() == MinInt {
		nz := big.NewInt(int64(z.Value()))
		return NewLong(nz.Neg(nz)).ToObject(), nil
	}
	return NewInt(-z.Value()).ToObject(), nil
}

func intAdd(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intAddMulOp(f, "__add__", v, w, intCheckedAdd, longAdd)
}

func intAnd(f *Frame, v, w *Object) (*Object, *BaseException) {
	if !w.isInstance(IntType) {
		return NotImplemented, nil
	}
	return NewInt(toIntUnsafe(v).Value() & toIntUnsafe(w).Value()).ToObject(), nil
}

func intDiv(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intDivModOp(f, "__div__", v, w, intCheckedDiv, longDiv)
}

func intEq(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intCompare(compareOpEq, toIntUnsafe(v), w), nil
}

func intGE(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intCompare(compareOpGE, toIntUnsafe(v), w), nil
}

func intGetNewArgs(f *Frame, args Args, _ KWArgs) (*Object, *BaseException) {
	if raised := checkMethodArgs(f, "__getnewargs__", args, IntType); raised != nil {
		return nil, raised
	}
	return NewTuple(args[0]).ToObject(), nil
}

func intGT(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intCompare(compareOpGT, toIntUnsafe(v), w), nil
}

func intFloat(f *Frame, o *Object) (*Object, *BaseException) {
	i := toIntUnsafe(o).Value()
	return NewFloat(float64(i)).ToObject(), nil
}

func intHash(f *Frame, o *Object) (*Object, *BaseException) {
	return o, nil
}

func intIndex(f *Frame, o *Object) (*Object, *BaseException) {
	return o, nil
}

func intInt(f *Frame, o *Object) (*Object, *BaseException) {
	if o.typ == IntType {
		return o, nil
	}
	return NewInt(toIntUnsafe(o).Value()).ToObject(), nil
}

func intInvert(f *Frame, o *Object) (*Object, *BaseException) {
	return NewInt(^toIntUnsafe(o).Value()).ToObject(), nil
}

func intLE(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intCompare(compareOpLE, toIntUnsafe(v), w), nil
}

func intLong(f *Frame, o *Object) (*Object, *BaseException) {
	return NewLong(big.NewInt(int64(toIntUnsafe(o).Value()))).ToObject(), nil
}

func intLShift(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intShiftOp(f, v, w, func(v, w int) (int, int, bool) { return v, w, false })
}

func intLT(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intCompare(compareOpLT, toIntUnsafe(v), w), nil
}

func intMod(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intDivModOp(f, "__mod__", v, w, intCheckedMod, longMod)
}

func intMul(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intAddMulOp(f, "__mul__", v, w, intCheckedMul, longMul)
}

func intNative(f *Frame, o *Object) (reflect.Value, *BaseException) {
	return reflect.ValueOf(toIntUnsafe(o).Value()), nil
}

func intNE(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intCompare(compareOpNE, toIntUnsafe(v), w), nil
}

func intNew(f *Frame, t *Type, args Args, _ KWArgs) (*Object, *BaseException) {
	if len(args) == 0 {
		return newObject(t), nil
	}
	o := args[0]
	if len(args) == 1 && o.typ.slots.Int != nil {
		result, raised := intFromObject(f, o)
		if result != nil || raised != nil {
			if t == IntType || raised != nil {
				return result, raised
			}
			o := newObject(t)
			toIntUnsafe(o).value = toIntUnsafe(result).value
			return o, nil
		}
	}
	if len(args) > 2 {
		return nil, f.RaiseType(TypeErrorType, fmt.Sprintf("int() takes at most 2 arguments (%d given)", len(args)))
	}
	if !o.isInstance(StrType) {
		if len(args) == 2 {
			return nil, f.RaiseType(TypeErrorType, "int() can't convert non-string with explicit base")
		}
		return nil, f.RaiseType(TypeErrorType, fmt.Sprintf("int() argument must be a string or a number, not '%s'", o.typ.Name()))
	}
	s := toStrUnsafe(o).Value()
	base := 10
	if len(args) == 2 {
		if args[1].typ.slots.Int == nil {
			return nil, f.RaiseType(TypeErrorType, "an integer is required")
		}
		baseArg, raised := intFromObject(f, args[1])
		if raised != nil {
			return nil, raised
		}
		base = toIntUnsafe(baseArg).Value()
		if base < 0 || base == 1 || base > 36 {
			return nil, f.RaiseType(ValueErrorType, "int() base must be >= 2 and <= 36")
		}
	}
	i, ok := numParseInteger(new(big.Int), s, base)
	if !ok {
		format := "invalid literal for int() with base %d: %s"
		return nil, f.RaiseType(ValueErrorType, fmt.Sprintf(format, base, s))
	}
	if !numInIntRange(i) {
		if t == IntType {
			return NewLong(i).ToObject(), nil
		}
		return nil, f.RaiseType(OverflowErrorType, "Python int too large to convert to a Go int")
	}
	if t != IntType {
		o := newObject(t)
		toIntUnsafe(o).value = int(i.Int64())
		return o, nil
	}
	return NewInt(int(i.Int64())).ToObject(), nil
}

func intNonZero(f *Frame, o *Object) (*Object, *BaseException) {
	return GetBool(toIntUnsafe(o).Value() != 0).ToObject(), nil
}

func intOr(f *Frame, v, w *Object) (*Object, *BaseException) {
	if !w.isInstance(IntType) {
		return NotImplemented, nil
	}
	return NewInt(toIntUnsafe(v).Value() | toIntUnsafe(w).Value()).ToObject(), nil
}

func intRAdd(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intAddMulOp(f, "__radd__", v, w, intCheckedAdd, longAdd)
}

func intRDiv(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intDivModOp(f, "__rdiv__", v, w, func(v, w int) (int, divModResult) {
		return intCheckedDiv(w, v)
	}, func(z, x, y *big.Int) {
		longDiv(z, y, x)
	})
}

func intRepr(f *Frame, o *Object) (*Object, *BaseException) {
	return NewStr(strconv.FormatInt(int64(toIntUnsafe(o).Value()), 10)).ToObject(), nil
}

func intRMod(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intDivModOp(f, "__rmod__", v, w, func(v, w int) (int, divModResult) {
		return intCheckedMod(w, v)
	}, func(z, x, y *big.Int) {
		longMod(z, y, x)
	})
}

func intRMul(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intAddMulOp(f, "__rmul__", v, w, intCheckedMul, longMul)
}

func intRLShift(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intShiftOp(f, v, w, func(v, w int) (int, int, bool) { return w, v, false })
}

func intRRShift(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intShiftOp(f, v, w, func(v, w int) (int, int, bool) { return w, v, true })
}

func intRShift(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intShiftOp(f, v, w, func(v, w int) (int, int, bool) { return v, w, true })
}

func intRSub(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intAddMulOp(f, "__rsub__", v, w, func(v, w int) (int, bool) {
		return intCheckedSub(w, v)
	}, func(z, x, y *big.Int) {
		longSub(z, y, x)
	})
}

func intSub(f *Frame, v, w *Object) (*Object, *BaseException) {
	return intAddMulOp(f, "__sub__", v, w, intCheckedSub, longSub)
}

func intXor(f *Frame, v, w *Object) (*Object, *BaseException) {
	if !w.isInstance(IntType) {
		return NotImplemented, nil
	}
	return NewInt(toIntUnsafe(v).Value() ^ toIntUnsafe(w).Value()).ToObject(), nil
}

func initIntType(dict map[string]*Object) {
	dict["__getnewargs__"] = newBuiltinFunction("__getnewargs__", intGetNewArgs).ToObject()
	IntType.slots.Abs = &unaryOpSlot{intAbs}
	IntType.slots.Add = &binaryOpSlot{intAdd}
	IntType.slots.And = &binaryOpSlot{intAnd}
	IntType.slots.Div = &binaryOpSlot{intDiv}
	IntType.slots.Eq = &binaryOpSlot{intEq}
	IntType.slots.GE = &binaryOpSlot{intGE}
	IntType.slots.GT = &binaryOpSlot{intGT}
	IntType.slots.Float = &unaryOpSlot{intFloat}
	IntType.slots.Hash = &unaryOpSlot{intHash}
	IntType.slots.Index = &unaryOpSlot{intIndex}
	IntType.slots.Int = &unaryOpSlot{intInt}
	IntType.slots.Invert = &unaryOpSlot{intInvert}
	IntType.slots.LE = &binaryOpSlot{intLE}
	IntType.slots.LShift = &binaryOpSlot{intLShift}
	IntType.slots.LT = &binaryOpSlot{intLT}
	IntType.slots.Long = &unaryOpSlot{intLong}
	IntType.slots.Mod = &binaryOpSlot{intMod}
	IntType.slots.Mul = &binaryOpSlot{intMul}
	IntType.slots.Native = &nativeSlot{intNative}
	IntType.slots.NE = &binaryOpSlot{intNE}
	IntType.slots.New = &newSlot{intNew}
	IntType.slots.NonZero = &unaryOpSlot{intNonZero}
	IntType.slots.Or = &binaryOpSlot{intOr}
	IntType.slots.RAdd = &binaryOpSlot{intRAdd}
	IntType.slots.RAnd = &binaryOpSlot{intAnd}
	IntType.slots.RDiv = &binaryOpSlot{intRDiv}
	IntType.slots.Repr = &unaryOpSlot{intRepr}
	IntType.slots.RMod = &binaryOpSlot{intRMod}
	IntType.slots.RMul = &binaryOpSlot{intRMul}
	IntType.slots.ROr = &binaryOpSlot{intOr}
	IntType.slots.RLShift = &binaryOpSlot{intRLShift}
	IntType.slots.RRShift = &binaryOpSlot{intRRShift}
	IntType.slots.RShift = &binaryOpSlot{intRShift}
	IntType.slots.RSub = &binaryOpSlot{intRSub}
	IntType.slots.RXor = &binaryOpSlot{intXor}
	IntType.slots.Sub = &binaryOpSlot{intSub}
	IntType.slots.Xor = &binaryOpSlot{intXor}
}

type divModResult int

const (
	divModOK           divModResult = iota
	divModOverflow                  = iota
	divModZeroDivision              = iota
)

func intCompare(op compareOp, v *Int, w *Object) *Object {
	if !w.isInstance(IntType) {
		return NotImplemented
	}
	lhs, rhs := v.Value(), toIntUnsafe(w).Value()
	result := false
	switch op {
	case compareOpLT:
		result = lhs < rhs
	case compareOpLE:
		result = lhs <= rhs
	case compareOpEq:
		result = lhs == rhs
	case compareOpNE:
		result = lhs != rhs
	case compareOpGE:
		result = lhs >= rhs
	case compareOpGT:
		result = lhs > rhs
	}
	return GetBool(result).ToObject()
}

func intAddMulOp(f *Frame, method string, v, w *Object, fun func(v, w int) (int, bool), bigFun func(z, x, y *big.Int)) (*Object, *BaseException) {
	if !w.isInstance(IntType) {
		return NotImplemented, nil
	}
	r, ok := fun(toIntUnsafe(v).Value(), toIntUnsafe(w).Value())
	if !ok {
		return longCallBinary(bigFun, intToLong(toIntUnsafe(v)), intToLong(toIntUnsafe(w))), nil
	}
	return NewInt(r).ToObject(), nil
}

func intCheckedDiv(v, w int) (int, divModResult) {
	q, _, r := intCheckedDivMod(v, w)
	return q, r
}

func intCheckedDivMod(v, w int) (int, int, divModResult) {
	if w == 0 {
		return 0, 0, divModZeroDivision
	}
	if v == MinInt && w == -1 {
		return 0, 0, divModOverflow
	}
	q := v / w
	m := v % w
	if m != 0 && (w^m) < 0 {
		// In Python the result of the modulo operator is always the
		// same sign as the divisor, whereas in Go, the result is
		// always the same sign as the dividend. Therefore we need to
		// do an adjustment when the sign of the modulo result differs
		// from that of the divisor.
		m += w
		// Relatedly, in Python the result of division truncates toward
		// negative infinity whereas it truncates toward zero in Go.
		// The fact that the signs of the divisor and the modulo result
		// differ implies that the quotient is also negative so we also
		// adjust the quotient here.
		q--
	}
	return q, m, divModOK
}

func intCheckedAdd(v, w int) (int, bool) {
	if (v > 0 && w > MaxInt-v) || (v < 0 && w < MinInt-v) {
		return 0, false
	}
	return v + w, true
}

func intCheckedMod(v, w int) (int, divModResult) {
	_, m, r := intCheckedDivMod(v, w)
	return m, r
}

func intCheckedMul(v, w int) (int, bool) {
	if v == 0 || w == 0 || v == 1 || w == 1 {
		return v * w, true
	}
	// Since MinInt can only be multiplied by zero and one safely and we've
	// already handled that case above, we know this multiplication will
	// overflow. Unfortunately the division check below will fail to catch
	// this by coincidence: MinInt * -1 overflows to MinInt, causing the
	// expression x/w to overflow, coincidentally producing MinInt which
	// makes it seem as though the multiplication was correct.
	if v == MinInt || w == MinInt {
		return 0, false
	}
	x := v * w
	if x/w != v {
		return 0, false
	}
	return x, true
}

func intCheckedSub(v, w int) (int, bool) {
	if (w > 0 && v < MinInt+w) || (w < 0 && v > MaxInt+w) {
		return 0, false
	}
	return v - w, true
}

func intDivModOp(f *Frame, method string, v, w *Object, fun func(v, w int) (int, divModResult), bigFun func(z, x, y *big.Int)) (*Object, *BaseException) {
	if !w.isInstance(IntType) {
		return NotImplemented, nil
	}
	x, r := fun(toIntUnsafe(v).Value(), toIntUnsafe(w).Value())
	switch r {
	case divModOverflow:
		return longCallBinary(bigFun, intToLong(toIntUnsafe(v)), intToLong(toIntUnsafe(w))), nil
	case divModZeroDivision:
		return nil, f.RaiseType(ZeroDivisionErrorType, "integer division or modulo by zero")
	}
	return NewInt(x).ToObject(), nil
}

func intShiftOp(f *Frame, v, w *Object, fun func(int, int) (int, int, bool)) (*Object, *BaseException) {
	if !w.isInstance(IntType) {
		return NotImplemented, nil
	}
	lhs, rhs, rshift := fun(toIntUnsafe(v).Value(), toIntUnsafe(w).Value())
	if rhs < 0 {
		return nil, f.RaiseType(ValueErrorType, "negative shift count")
	}
	var result int
	n := uint(rhs)
	if rshift {
		result = lhs >> n
	} else {
		result = lhs << n
		if result>>n != lhs {
			return NewLong(new(big.Int).Lsh(big.NewInt(int64(lhs)), n)).ToObject(), nil
		}
	}
	return NewInt(result).ToObject(), nil
}

func intFromObject(f *Frame, o *Object) (*Object, *BaseException) {
	if o.typ.slots.Int == nil {
		return nil, nil
	}
	r, raised := o.typ.slots.Int.Fn(f, o)
	switch {
	case raised != nil:
		return nil, raised
	case r.isInstance(IntType) || r.isInstance(LongType):
		return r, nil
	default:
		format := "__int__ returned non-int (type %s)"
		return nil, f.RaiseType(TypeErrorType, fmt.Sprintf(format, r.typ.Name()))
	}
}

func intToLong(o *Int) *Long {
	return NewLong(big.NewInt(int64(o.Value())))
}