shifts.go

     1  // Copyright 2026 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  //specgen:category Shifts
     6  
     7  package spec
     8  
     9  import "math/bits"
    10  
    11  // ScaleSaturated multiplies x[i] by 2^scale[i], with saturation.
    12  //
    13  // Positive exponents scale up (shift left with saturation). Negative exponents
    14  // scale down ({{if eq .xB "int"}}arithmetic{{else}}logical{{end}} shift right).
    15  //
    16  //	z[i] = sat(x[i] * 2^scale[i])
    17  //
    18  //specgen:require scale=Int{xN}x{xL}
    19  func ScaleSaturated[E Ints | Uints, W Width, yE Ints](x Vec[E, W], scale Vec[yE, W]) (z Vec[E, W]) {
    20  	z = makeVec[E, W]()
    21  	for i := range z {
    22  		z[i] = scaleSaturated(x[i], scale[i])
    23  	}
    24  	return z
    25  }
    26  
    27  // {{if eq .xB "int" -}}
    28  // ShiftAllRight arithmetically shifts each element of x right by shift bits.
    29  // If shift is greater than the element width, the result is 0 or -1.
    30  // {{- else -}}
    31  // ShiftAllRight logically shifts each element of x right by shift bits.
    32  // If shift is greater than the element width, the result is 0.
    33  // {{- end}}
    34  //
    35  //	z[i] = x[i] >> shift
    36  func ShiftAllRight[E Ints | Uints, W Width](x Vec[E, W], shift uint64) (z Vec[E, W]) {
    37  	return map1[E, W, E, W](x, func(v E) E {
    38  		return v >> shift
    39  	})
    40  }
    41  
    42  // ShiftAllLeft shifts each element of x left by shift bits.
    43  // If shift is greater than the element width, the result is 0.
    44  //
    45  //	z[i] = x[i] << shift
    46  func ShiftAllLeft[E Ints | Uints, W Width](x Vec[E, W], shift uint64) (z Vec[E, W]) {
    47  	return map1[E, W, E, W](x, func(v E) E {
    48  		return v << shift
    49  	})
    50  }
    51  
    52  // {{if eq .xB "int" -}}
    53  // ShiftRight arithmetically shifts x[i] right by shift[i] bits.
    54  // If shift[i] is greater than the element width, the result is 0 or -1.
    55  // {{- else -}}
    56  // ShiftRight logically shifts x[i] right by shift[i] bits.
    57  // If shift[i] is greater than the element width, the result is 0.
    58  // {{- end}}
    59  //
    60  //	z[i] = x[i] >> shift[i]
    61  //
    62  //specgen:require shift=Uint{xN}x{xL}
    63  func ShiftRight[E Ints | Uints, W Width, yE Uints](x Vec[E, W], shift Vec[yE, W]) (z Vec[E, W]) {
    64  	z = makeVec[E, W]()
    65  	for i := range z {
    66  		z[i] = x[i] >> shift[i]
    67  	}
    68  	return z
    69  }
    70  
    71  // ShiftLeft shifts x[i] left by shift[i] bits.
    72  // If shift[i] is greater than the element width, the result is 0.
    73  //
    74  //	z[i] = x[i] << shift[i]
    75  //
    76  //specgen:require shift=Uint{xN}x{xL}
    77  func ShiftLeft[E Ints | Uints, W Width, yE Uints](x Vec[E, W], shift Vec[yE, W]) (z Vec[E, W]) {
    78  	z = makeVec[E, W]()
    79  	for i := range z {
    80  		z[i] = x[i] << shift[i]
    81  	}
    82  	return z
    83  }
    84  
    85  // RotateAllLeft rotates each element of x left by shift bits.
    86  //
    87  //	z[i] = rotateLeft(x[i], shift)
    88  func RotateAllLeft[E Uints, W Width](x Vec[E, W], shift uint64) (z Vec[E, W]) {
    89  	// Compute mod before casting to int.
    90  	k := int(shift % uint64(elemBits[E]()))
    91  	return map1[E, W, E, W](x, func(x E) E {
    92  		switch any(x).(type) {
    93  		case int8, uint8:
    94  			return E(bits.RotateLeft8(uint8(x), k))
    95  		case int16, uint16:
    96  			return E(bits.RotateLeft16(uint16(x), k))
    97  		case int32, uint32:
    98  			return E(bits.RotateLeft32(uint32(x), k))
    99  		case int64, uint64:
   100  			return E(bits.RotateLeft64(uint64(x), k))
   101  		}
   102  		panic("unreachable")
   103  	})
   104  }
   105  
   106  // RotateLeft rotates x left by shift bits, elementwise.
   107  //
   108  //	z[i] = rotateLeft(x[i], shift[i])
   109  func RotateLeft[E Uints, W Width](x, shift Vec[E, W]) (z Vec[E, W]) {
   110  	return map2[E, W, E, W](x, shift, func(x, shift E) E {
   111  		// Compute mod before casting to int.
   112  		k := int(shift % E(elemBits[E]()))
   113  		switch any(x).(type) {
   114  		case int8, uint8:
   115  			return E(bits.RotateLeft8(uint8(x), k))
   116  		case int16, uint16:
   117  			return E(bits.RotateLeft16(uint16(x), k))
   118  		case int32, uint32:
   119  			return E(bits.RotateLeft32(uint32(x), k))
   120  		case int64, uint64:
   121  			return E(bits.RotateLeft64(uint64(x), k))
   122  		}
   123  		panic("unreachable")
   124  	})
   125  }
   126  
   127  // RotateAllRight rotates each element of x right by shift bits.
   128  //
   129  //	z[i] = rotateRight(x[i], shift)
   130  func RotateAllRight[E Uints, W Width](x Vec[E, W], shift uint64) (z Vec[E, W]) {
   131  	// Compute mod before casting to int.
   132  	k := int(shift % uint64(elemBits[E]()))
   133  	return map1[E, W, E, W](x, func(x E) E {
   134  		switch any(x).(type) {
   135  		case int8, uint8:
   136  			return E(bits.RotateLeft8(uint8(x), -k))
   137  		case int16, uint16:
   138  			return E(bits.RotateLeft16(uint16(x), -k))
   139  		case int32, uint32:
   140  			return E(bits.RotateLeft32(uint32(x), -k))
   141  		case int64, uint64:
   142  			return E(bits.RotateLeft64(uint64(x), -k))
   143  		}
   144  		panic("unreachable")
   145  	})
   146  }
   147  
   148  // RotateRight rotates x right by shift bits, elementwise.
   149  //
   150  //	z[i] = rotateRight(x[i], shift[i])
   151  func RotateRight[E Uints, W Width](x, shift Vec[E, W]) (z Vec[E, W]) {
   152  	return map2[E, W, E, W](x, shift, func(x, shift E) E {
   153  		// Compute mod before casting to int.
   154  		k := int(shift % E(elemBits[E]()))
   155  		switch any(x).(type) {
   156  		case int8, uint8:
   157  			return E(bits.RotateLeft8(uint8(x), -k))
   158  		case int16, uint16:
   159  			return E(bits.RotateLeft16(uint16(x), -k))
   160  		case int32, uint32:
   161  			return E(bits.RotateLeft32(uint32(x), -k))
   162  		case int64, uint64:
   163  			return E(bits.RotateLeft64(uint64(x), -k))
   164  		}
   165  		panic("unreachable")
   166  	})
   167  }
   168
View as plain text