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