github.com/teddydd/sh@v2.6.4+incompatible/expand/arith.go (about)

     1  // Copyright (c) 2017, Daniel Martí <mvdan@mvdan.cc>
     2  // See LICENSE for licensing information
     3  
     4  package expand
     5  
     6  import (
     7  	"fmt"
     8  	"strconv"
     9  
    10  	"mvdan.cc/sh/syntax"
    11  )
    12  
    13  func Arithm(cfg *Config, expr syntax.ArithmExpr) (int, error) {
    14  	switch x := expr.(type) {
    15  	case *syntax.Word:
    16  		str, err := Literal(cfg, x)
    17  		if err != nil {
    18  			return 0, err
    19  		}
    20  		// recursively fetch vars
    21  		i := 0
    22  		for str != "" && syntax.ValidName(str) {
    23  			val := cfg.envGet(str)
    24  			if val == "" {
    25  				break
    26  			}
    27  			if i++; i >= maxNameRefDepth {
    28  				break
    29  			}
    30  			str = val
    31  		}
    32  		// default to 0
    33  		return atoi(str), nil
    34  	case *syntax.ParenArithm:
    35  		return Arithm(cfg, x.X)
    36  	case *syntax.UnaryArithm:
    37  		switch x.Op {
    38  		case syntax.Inc, syntax.Dec:
    39  			name := x.X.(*syntax.Word).Lit()
    40  			old := atoi(cfg.envGet(name))
    41  			val := old
    42  			if x.Op == syntax.Inc {
    43  				val++
    44  			} else {
    45  				val--
    46  			}
    47  			cfg.envSet(name, strconv.Itoa(val))
    48  			if x.Post {
    49  				return old, nil
    50  			}
    51  			return val, nil
    52  		}
    53  		val, err := Arithm(cfg, x.X)
    54  		if err != nil {
    55  			return 0, err
    56  		}
    57  		switch x.Op {
    58  		case syntax.Not:
    59  			return oneIf(val == 0), nil
    60  		case syntax.Plus:
    61  			return val, nil
    62  		default: // syntax.Minus
    63  			return -val, nil
    64  		}
    65  	case *syntax.BinaryArithm:
    66  		switch x.Op {
    67  		case syntax.Assgn, syntax.AddAssgn, syntax.SubAssgn,
    68  			syntax.MulAssgn, syntax.QuoAssgn, syntax.RemAssgn,
    69  			syntax.AndAssgn, syntax.OrAssgn, syntax.XorAssgn,
    70  			syntax.ShlAssgn, syntax.ShrAssgn:
    71  			return cfg.assgnArit(x)
    72  		case syntax.Quest: // Colon can't happen here
    73  			cond, err := Arithm(cfg, x.X)
    74  			if err != nil {
    75  				return 0, err
    76  			}
    77  			b2 := x.Y.(*syntax.BinaryArithm) // must have Op==Colon
    78  			if cond == 1 {
    79  				return Arithm(cfg, b2.X)
    80  			}
    81  			return Arithm(cfg, b2.Y)
    82  		}
    83  		left, err := Arithm(cfg, x.X)
    84  		if err != nil {
    85  			return 0, err
    86  		}
    87  		right, err := Arithm(cfg, x.Y)
    88  		if err != nil {
    89  			return 0, err
    90  		}
    91  		return binArit(x.Op, left, right), nil
    92  	default:
    93  		panic(fmt.Sprintf("unexpected arithm expr: %T", x))
    94  	}
    95  }
    96  
    97  func oneIf(b bool) int {
    98  	if b {
    99  		return 1
   100  	}
   101  	return 0
   102  }
   103  
   104  // atoi is just a shorthand for strconv.Atoi that ignores the error,
   105  // just like shells do.
   106  func atoi(s string) int {
   107  	n, _ := strconv.Atoi(s)
   108  	return n
   109  }
   110  
   111  func (cfg *Config) assgnArit(b *syntax.BinaryArithm) (int, error) {
   112  	name := b.X.(*syntax.Word).Lit()
   113  	val := atoi(cfg.envGet(name))
   114  	arg, err := Arithm(cfg, b.Y)
   115  	if err != nil {
   116  		return 0, err
   117  	}
   118  	switch b.Op {
   119  	case syntax.Assgn:
   120  		val = arg
   121  	case syntax.AddAssgn:
   122  		val += arg
   123  	case syntax.SubAssgn:
   124  		val -= arg
   125  	case syntax.MulAssgn:
   126  		val *= arg
   127  	case syntax.QuoAssgn:
   128  		val /= arg
   129  	case syntax.RemAssgn:
   130  		val %= arg
   131  	case syntax.AndAssgn:
   132  		val &= arg
   133  	case syntax.OrAssgn:
   134  		val |= arg
   135  	case syntax.XorAssgn:
   136  		val ^= arg
   137  	case syntax.ShlAssgn:
   138  		val <<= uint(arg)
   139  	case syntax.ShrAssgn:
   140  		val >>= uint(arg)
   141  	}
   142  	cfg.envSet(name, strconv.Itoa(val))
   143  	return val, nil
   144  }
   145  
   146  func intPow(a, b int) int {
   147  	p := 1
   148  	for b > 0 {
   149  		if b&1 != 0 {
   150  			p *= a
   151  		}
   152  		b >>= 1
   153  		a *= a
   154  	}
   155  	return p
   156  }
   157  
   158  func binArit(op syntax.BinAritOperator, x, y int) int {
   159  	switch op {
   160  	case syntax.Add:
   161  		return x + y
   162  	case syntax.Sub:
   163  		return x - y
   164  	case syntax.Mul:
   165  		return x * y
   166  	case syntax.Quo:
   167  		return x / y
   168  	case syntax.Rem:
   169  		return x % y
   170  	case syntax.Pow:
   171  		return intPow(x, y)
   172  	case syntax.Eql:
   173  		return oneIf(x == y)
   174  	case syntax.Gtr:
   175  		return oneIf(x > y)
   176  	case syntax.Lss:
   177  		return oneIf(x < y)
   178  	case syntax.Neq:
   179  		return oneIf(x != y)
   180  	case syntax.Leq:
   181  		return oneIf(x <= y)
   182  	case syntax.Geq:
   183  		return oneIf(x >= y)
   184  	case syntax.And:
   185  		return x & y
   186  	case syntax.Or:
   187  		return x | y
   188  	case syntax.Xor:
   189  		return x ^ y
   190  	case syntax.Shr:
   191  		return x >> uint(y)
   192  	case syntax.Shl:
   193  		return x << uint(y)
   194  	case syntax.AndArit:
   195  		return oneIf(x != 0 && y != 0)
   196  	case syntax.OrArit:
   197  		return oneIf(x != 0 || y != 0)
   198  	default: // syntax.Comma
   199  		// x is executed but its result discarded
   200  		return y
   201  	}
   202  }