gosh/arith.go

package main

// Arithmetic expansion $((expression)) with 64-bit integers and the
// operators of C (and ** for powers), in the order of precedence:
//
//	x++ x--, ++x --x, + - ! ~ (unary), **, * / %, + -, << >>,
//	< <= > >=, == !=, &, ^, |, &&, ||, ?:, = += -= *= /= %= <<= >>= &= ^= |=, ","
//
// Variables can be used with or without $; their values are evaluated as
// expressions, an empty value is 0, an unset variable is an error (like
// everywhere in gosh). Numbers are decimal, 0x... hexadecimal or 0...
// octal.

import (
	"errors"
	"fmt"
	"os"
	"strconv"
	"strings"
)

// arithEnd returns the index of the last ) of $((...)) at s[i], or -1 if
// there is no arithmetic expansion (e.g. $( (cmd) ), a command substitution
// of a subshell).
func arithEnd(s string, i int) int {
	if !strings.HasPrefix(s[i:], "$((") {
		return -1
	}
	end := substEnd(s, i)
	if end < 0 {
		return -1
	}
	// The inner ( must end right before the outer )
	if inner := substEnd(s, i+1); inner != end-1 {
		return -1
	}
	return end
}

// evalArithText expands an arithmetic expression (variables, $(...), quotes)
// and evaluates it.
func evalArithText(text string) (int64, error) {
	expr, err := expandWordText(text)
	if err != nil {
		return 0, err
	}
	return evalArith(expr, 0)
}

// maxArithDepth limits the recursion of variables whose values are
// expressions with variables.
const maxArithDepth = 32

func evalArith(expr string, depth int) (int64, error) {
	if depth > maxArithDepth {
		return 0, errors.New("arithmetic: expression recursion level exceeded")
	}
	toks, err := lexArith(expr)
	if err != nil {
		return 0, err
	}
	a := &arith{toks: toks, depth: depth, expr: expr}
	if len(toks) == 0 {
		return 0, nil
	}
	v, err := a.comma()
	if err != nil {
		return 0, err
	}
	if a.pos < len(a.toks) {
		return 0, a.syntaxErr()
	}
	return v, nil
}

type arithToken struct {
	kind byte // 'n' number, 'v' variable, 'o' operator
	text string
	num  int64
}

var arithOps = []string{
	"<<=", ">>=", "**", "++", "--", "<<", ">>", "<=", ">=", "==", "!=", "&&", "||",
	"+=", "-=", "*=", "/=", "%=", "&=", "^=", "|=",
	"+", "-", "*", "/", "%", "<", ">", "&", "|", "^", "!", "~", "?", ":", "=", "(", ")", ",",
}

func lexArith(s string) ([]arithToken, error) {
	var toks []arithToken
	for i := 0; i < len(s); {
		ch := s[i]
		switch {
		case ch == ' ' || ch == '\t' || ch == '\n':
			i++
		case ch >= '0' && ch <= '9':
			j := i
			for j < len(s) && (isVarNameByte(s[j], false)) {
				j++
			}
			n, err := strconv.ParseInt(s[i:j], 0, 64)
			if err != nil {
				return nil, fmt.Errorf("%s: value too great for base (error token is \"%s\")", strings.TrimSpace(s), s[i:j])
			}
			toks = append(toks, arithToken{kind: 'n', text: s[i:j], num: n})
			i = j
		case isVarNameByte(ch, true):
			j := i
			for j < len(s) && isVarNameByte(s[j], false) {
				j++
			}
			toks = append(toks, arithToken{kind: 'v', text: s[i:j]})
			i = j
		default:
			found := false
			for _, op := range arithOps {
				if strings.HasPrefix(s[i:], op) {
					toks = append(toks, arithToken{kind: 'o', text: op})
					i += len(op)
					found = true
					break
				}
			}
			if !found {
				return nil, fmt.Errorf("%s: syntax error: invalid arithmetic operator (error token is \"%s\")", strings.TrimSpace(s), sanitize(s[i:]))
			}
		}
	}
	return toks, nil
}

type arith struct {
	toks  []arithToken
	pos   int
	skip  int // > 0 while a branch is evaluated that has no effect (&&, ||, ?:)
	depth int
	expr  string
}

func (a *arith) peekOp(ops ...string) string {
	if a.pos < len(a.toks) && a.toks[a.pos].kind == 'o' {
		for _, op := range ops {
			if a.toks[a.pos].text == op {
				return op
			}
		}
	}
	return ""
}

func (a *arith) syntaxErr() error {
	token := "end of expression"
	if a.pos < len(a.toks) {
		token = a.toks[a.pos].text
	}
	return fmt.Errorf("%s: arithmetic syntax error (error token is \"%s\")", strings.TrimSpace(a.expr), sanitize(token))
}

// variable returns the value of a variable in an expression.
func (a *arith) variable(name string) (int64, error) {
	if a.skip > 0 {
		return 0, nil
	}
	value, ok := os.LookupEnv(name)
	if !ok {
		return 0, fmt.Errorf("%s: unbound variable", name)
	}
	if strings.TrimSpace(value) == "" {
		return 0, nil
	}
	if n, err := strconv.ParseInt(strings.TrimSpace(value), 0, 64); err == nil {
		return n, nil
	}
	return evalArith(value, a.depth+1)
}

func (a *arith) setVariable(name string, v int64) {
	if a.skip == 0 {
		os.Setenv(name, strconv.FormatInt(v, 10))
	}
}

func (a *arith) comma() (int64, error) {
	v, err := a.assign()
	for err == nil && a.peekOp(",") != "" {
		a.pos++
		v, err = a.assign()
	}
	return v, err
}

func (a *arith) assign() (int64, error) {
	if a.pos+1 < len(a.toks) && a.toks[a.pos].kind == 'v' && a.toks[a.pos+1].kind == 'o' {
		name, op := a.toks[a.pos].text, a.toks[a.pos+1].text
		if op == "=" || len(op) >= 2 && strings.HasSuffix(op, "=") && op != "==" && op != "!=" && op != "<=" && op != ">=" {
			a.pos += 2
			rhs, err := a.assign()
			if err != nil {
				return 0, err
			}
			v := rhs
			if op != "=" {
				old, err := a.variable(name)
				if err != nil {
					return 0, err
				}
				if v, err = a.apply(strings.TrimSuffix(op, "="), old, rhs); err != nil {
					return 0, err
				}
			}
			a.setVariable(name, v)
			return v, nil
		}
	}
	return a.ternary()
}

func (a *arith) ternary() (int64, error) {
	cond, err := a.binary(0)
	if err != nil || a.peekOp("?") == "" {
		return cond, err
	}
	a.pos++
	if cond == 0 {
		a.skip++
	}
	yes, err := a.assign()
	if cond == 0 {
		a.skip--
	}
	if err != nil {
		return 0, err
	}
	if a.peekOp(":") == "" {
		return 0, a.syntaxErr()
	}
	a.pos++
	if cond != 0 {
		a.skip++
	}
	no, err := a.assign()
	if cond != 0 {
		a.skip--
	}
	if err != nil {
		return 0, err
	}
	if cond != 0 {
		return yes, nil
	}
	return no, nil
}

// binaryLevels are the binary operators from the lowest precedence on.
var binaryLevels = [][]string{
	{"||"}, {"&&"}, {"|"}, {"^"}, {"&"}, {"==", "!="}, {"<", "<=", ">", ">="},
	{"<<", ">>"}, {"+", "-"}, {"*", "/", "%"},
}

func (a *arith) binary(level int) (int64, error) {
	if level == len(binaryLevels) {
		return a.power()
	}
	lhs, err := a.binary(level + 1)
	if err != nil {
		return 0, err
	}
	for {
		op := a.peekOp(binaryLevels[level]...)
		if op == "" {
			return lhs, nil
		}
		a.pos++
		// && and || do not evaluate the right side if the result is known
		short := op == "&&" && lhs == 0 || op == "||" && lhs != 0
		if short {
			a.skip++
		}
		rhs, err := a.binary(level + 1)
		if short {
			a.skip--
		}
		if err != nil {
			return 0, err
		}
		if lhs, err = a.apply(op, lhs, rhs); err != nil {
			return 0, err
		}
	}
}

func boolInt(b bool) int64 {
	if b {
		return 1
	}
	return 0
}

// apply applies a binary operator.
func (a *arith) apply(op string, x, y int64) (int64, error) {
	switch op {
	case "||":
		return boolInt(x != 0 || y != 0), nil
	case "&&":
		return boolInt(x != 0 && y != 0), nil
	case "|":
		return x | y, nil
	case "^":
		return x ^ y, nil
	case "&":
		return x & y, nil
	case "==":
		return boolInt(x == y), nil
	case "!=":
		return boolInt(x != y), nil
	case "<":
		return boolInt(x < y), nil
	case "<=":
		return boolInt(x <= y), nil
	case ">":
		return boolInt(x > y), nil
	case ">=":
		return boolInt(x >= y), nil
	case "<<":
		return x << uint64(y&63), nil
	case ">>":
		return x >> uint64(y&63), nil
	case "+":
		return x + y, nil
	case "-":
		return x - y, nil
	case "*":
		return x * y, nil
	case "/", "%":
		if y == 0 {
			if a.skip > 0 {
				return 0, nil
			}
			return 0, fmt.Errorf("%s: division by 0", strings.TrimSpace(a.expr))
		}
		if op == "/" {
			return x / y, nil
		}
		return x % y, nil
	case "**":
		if y < 0 {
			if a.skip > 0 {
				return 0, nil
			}
			return 0, fmt.Errorf("%s: exponent less than 0", strings.TrimSpace(a.expr))
		}
		result := int64(1)
		for ; y > 0; y-- {
			result *= x
		}
		return result, nil
	}
	return 0, a.syntaxErr()
}

// power parses x ** y (right associative, binds tighter than * but less
// than the unary operators).
func (a *arith) power() (int64, error) {
	base, err := a.unary()
	if err != nil || a.peekOp("**") == "" {
		return base, err
	}
	a.pos++
	exp, err := a.power()
	if err != nil {
		return 0, err
	}
	return a.apply("**", base, exp)
}

func (a *arith) unary() (int64, error) {
	switch op := a.peekOp("+", "-", "!", "~", "++", "--"); op {
	case "":
		return a.postfix()
	case "++", "--":
		a.pos++
		if a.pos >= len(a.toks) || a.toks[a.pos].kind != 'v' {
			return 0, a.syntaxErr()
		}
		name := a.toks[a.pos].text
		a.pos++
		v, err := a.variable(name)
		if err != nil {
			return 0, err
		}
		if op == "++" {
			v++
		} else {
			v--
		}
		a.setVariable(name, v)
		return v, nil
	default:
		a.pos++
		v, err := a.unary()
		if err != nil {
			return 0, err
		}
		switch op {
		case "-":
			return -v, nil
		case "!":
			return boolInt(v == 0), nil
		case "~":
			return ^v, nil
		}
		return v, nil
	}
}

func (a *arith) postfix() (int64, error) {
	if a.pos < len(a.toks) && a.toks[a.pos].kind == 'v' {
		name := a.toks[a.pos].text
		a.pos++
		v, err := a.variable(name)
		if err != nil {
			return 0, err
		}
		if op := a.peekOp("++", "--"); op != "" {
			a.pos++
			n := v + 1
			if op == "--" {
				n = v - 1
			}
			a.setVariable(name, n)
		}
		return v, nil
	}
	return a.primary()
}

func (a *arith) primary() (int64, error) {
	if a.pos >= len(a.toks) {
		return 0, a.syntaxErr()
	}
	t := a.toks[a.pos]
	switch {
	case t.kind == 'n':
		a.pos++
		return t.num, nil
	case t.kind == 'o' && t.text == "(":
		a.pos++
		v, err := a.comma()
		if err != nil {
			return 0, err
		}
		if a.peekOp(")") == "" {
			return 0, a.syntaxErr()
		}
		a.pos++
		return v, nil
	}
	return 0, a.syntaxErr()
}