// Copyright 2019 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// 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.

#include <stddef.h>

#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE "tests/shift_test.cc"
#include "hwy/foreach_target.h"  // IWYU pragma: keep
#include "hwy/highway.h"
#include "hwy/tests/test_util-inl.h"

HWY_BEFORE_NAMESPACE();
namespace hwy {
namespace HWY_NAMESPACE {

template <bool kSigned>
struct TestLeftShifts {
  template <typename T, class D>
  HWY_NOINLINE void operator()(T t, D d) {
    if (kSigned) {
      // Also test positive values
      TestLeftShifts</*kSigned=*/false>()(t, d);
    }

    using TI = MakeSigned<T>;
    using TU = MakeUnsigned<T>;
    const size_t N = Lanes(d);
    auto expected = AllocateAligned<T>(N);
    HWY_ASSERT(expected);

    // Values to shift
    const auto values = Iota(d, kSigned ? -TI(N) : TI(0));
    constexpr size_t kMaxShift = (sizeof(T) * 8) - 1;

    // 0
    HWY_ASSERT_VEC_EQ(d, values, ShiftLeft<0>(values));
    HWY_ASSERT_VEC_EQ(d, values, ShiftLeftSame(values, 0));

    // 1
    for (size_t i = 0; i < N; ++i) {
      const T value =
          kSigned ? static_cast<T>(static_cast<T>(i) - static_cast<T>(N))
                  : static_cast<T>(i);
      expected[i] = static_cast<T>(static_cast<TU>(value) << 1);
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), ShiftLeft<1>(values));
    HWY_ASSERT_VEC_EQ(d, expected.get(), ShiftLeftSame(values, 1));

    // max
    for (size_t i = 0; i < N; ++i) {
      const T value =
          kSigned ? static_cast<T>(static_cast<T>(i) - static_cast<T>(N))
                  : static_cast<T>(i);
      expected[i] = static_cast<T>(static_cast<TU>(value) << kMaxShift);
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), ShiftLeft<kMaxShift>(values));
    HWY_ASSERT_VEC_EQ(d, expected.get(), ShiftLeftSame(values, kMaxShift));
  }
};

template <bool kSigned>
struct TestVariableLeftShifts {
  template <typename T, class D>
  HWY_NOINLINE void operator()(T t, D d) {
    if (kSigned) {
      // Also test positive values
      TestVariableLeftShifts</*kSigned=*/false>()(t, d);
    }

    using TI = MakeSigned<T>;
    using TU = MakeUnsigned<T>;
    const size_t N = Lanes(d);
    auto expected = AllocateAligned<T>(N);
    HWY_ASSERT(expected);

    const auto v0 = Zero(d);
    const auto v1 = Set(d, 1);
    const auto values = Iota(d, kSigned ? -TI(N) : TI(0));  // value to shift

    constexpr size_t kMaxShift = (sizeof(T) * 8) - 1;
    const auto max_shift = Set(d, kMaxShift);
    const auto small_shifts = And(Iota(d, 0), max_shift);
    const auto large_shifts = Sub(max_shift, small_shifts);

    // Same: 0
    HWY_ASSERT_VEC_EQ(d, values, Shl(values, v0));

    // Same: 1
    for (size_t i = 0; i < N; ++i) {
      const T value =
          kSigned ? static_cast<T>(static_cast<T>(i) - static_cast<T>(N))
                  : static_cast<T>(i);
      expected[i] = static_cast<T>(static_cast<TU>(value) << 1);
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), Shl(values, v1));

    // Same: max
    for (size_t i = 0; i < N; ++i) {
      const T value =
          kSigned ? static_cast<T>(static_cast<T>(i) - static_cast<T>(N))
                  : static_cast<T>(i);
      expected[i] = static_cast<T>(static_cast<TU>(value) << kMaxShift);
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), Shl(values, max_shift));

    // Variable: small
    for (size_t i = 0; i < N; ++i) {
      const T value =
          kSigned ? static_cast<T>(static_cast<T>(i) - static_cast<T>(N))
                  : static_cast<T>(i);
      expected[i] = static_cast<T>(static_cast<TU>(value) << (i & kMaxShift));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), Shl(values, small_shifts));

    // Variable: large
    for (size_t i = 0; i < N; ++i) {
      expected[i] =
          static_cast<T>(static_cast<TU>(1) << (kMaxShift - (i & kMaxShift)));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), Shl(v1, large_shifts));
  }
};

struct TestUnsignedRightShifts {
  template <typename T, class D>
  HWY_NOINLINE void operator()(T /*unused*/, D d) {
    const size_t N = Lanes(d);
    auto expected = AllocateAligned<T>(N);
    HWY_ASSERT(expected);

    const auto values = Iota(d, 0);

    const T kMax = LimitsMax<T>();
    constexpr size_t kMaxShift = (sizeof(T) * 8) - 1;

    // Shift by 0
    HWY_ASSERT_VEC_EQ(d, values, ShiftRight<0>(values));
    HWY_ASSERT_VEC_EQ(d, values, ShiftRightSame(values, 0));

    // Shift by 1
    for (size_t i = 0; i < N; ++i) {
      expected[i] = static_cast<T>(static_cast<T>(i & kMax) >> 1);
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), ShiftRight<1>(values));
    HWY_ASSERT_VEC_EQ(d, expected.get(), ShiftRightSame(values, 1));

    // max
    for (size_t i = 0; i < N; ++i) {
      expected[i] = static_cast<T>(static_cast<T>(i & kMax) >> kMaxShift);
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), ShiftRight<kMaxShift>(values));
    HWY_ASSERT_VEC_EQ(d, expected.get(), ShiftRightSame(values, kMaxShift));
  }
};

struct TestRotateLeft {
  template <typename T, class D>
  HWY_NOINLINE void operator()(T /*unused*/, D d) {
    using TU = MakeUnsigned<T>;

    const size_t N = Lanes(d);
    auto expected = AllocateAligned<T>(N);
    HWY_ASSERT(expected);

    constexpr size_t kBits = sizeof(T) * 8;
    const Vec<D> mask_shift = Set(d, static_cast<T>(kBits - 1));
    // Cover as many bit positions as possible to test shifting out
    const Vec<D> values =
        Shl(Set(d, static_cast<T>(1)), And(Iota(d, 0), mask_shift));
    const Vec<D> values2 = Xor(values, SignBit(d));

    // Rotate by 0
    HWY_ASSERT_VEC_EQ(d, values, RotateLeft<0>(values));
    HWY_ASSERT_VEC_EQ(d, values2, RotateLeft<0>(values2));

    // Rotate by 1
    Store(values, d, expected.get());
    for (size_t i = 0; i < N; ++i) {
      expected[i] =
          ConvertScalarTo<T>((static_cast<TU>(expected[i]) << 1) |
                             (static_cast<TU>(expected[i]) >> (kBits - 1)));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), RotateLeft<1>(values));

    for (size_t i = 0; i < N; ++i) {
      expected[i] = ConvertScalarTo<T>(expected[i] ^ static_cast<T>(1));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), RotateLeft<1>(values2));

    // Rotate by half
    Store(values, d, expected.get());
    for (size_t i = 0; i < N; ++i) {
      expected[i] =
          ConvertScalarTo<T>((static_cast<TU>(expected[i]) << (kBits / 2)) |
                             (static_cast<TU>(expected[i]) >> (kBits / 2)));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), RotateLeft<kBits / 2>(values));

    for (size_t i = 0; i < N; ++i) {
      expected[i] = ConvertScalarTo<T>(
          expected[i] ^ (static_cast<T>(1) << ((kBits / 2) - 1)));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), RotateLeft<kBits / 2>(values2));

    // Rotate by max
    Store(values, d, expected.get());
    for (size_t i = 0; i < N; ++i) {
      expected[i] =
          ConvertScalarTo<T>((static_cast<TU>(expected[i]) << (kBits - 1)) |
                             (static_cast<TU>(expected[i]) >> 1));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), RotateLeft<kBits - 1>(values));

    for (size_t i = 0; i < N; ++i) {
      expected[i] =
          ConvertScalarTo<T>(expected[i] ^ (static_cast<T>(1) << (kBits - 2)));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), RotateLeft<kBits - 1>(values2));
  }
};

HWY_NOINLINE void TestAllRotateLeft() {
  ForIntegerTypes(ForPartialVectors<TestRotateLeft>());
}

struct TestRotateRight {
  template <typename T, class D>
  HWY_NOINLINE void operator()(T /*unused*/, D d) {
    using TU = MakeUnsigned<T>;

    const size_t N = Lanes(d);
    auto expected = AllocateAligned<T>(N);
    HWY_ASSERT(expected);

    constexpr size_t kBits = sizeof(T) * 8;
    const Vec<D> mask_shift = Set(d, static_cast<T>(kBits - 1));
    // Cover as many bit positions as possible to test shifting out
    const Vec<D> values =
        Shl(Set(d, static_cast<T>(1)), And(Iota(d, 0), mask_shift));
    const Vec<D> values2 = Xor(values, SignBit(d));

    // Rotate by 0
    HWY_ASSERT_VEC_EQ(d, values, RotateRight<0>(values));
    HWY_ASSERT_VEC_EQ(d, values2, RotateRight<0>(values2));

    // Rotate by 1
    Store(values, d, expected.get());
    for (size_t i = 0; i < N; ++i) {
      expected[i] =
          ConvertScalarTo<T>((static_cast<TU>(expected[i]) >> 1) |
                             (static_cast<TU>(expected[i]) << (kBits - 1)));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), RotateRight<1>(values));

    for (size_t i = 0; i < N; ++i) {
      expected[i] =
          ConvertScalarTo<T>(expected[i] ^ (static_cast<T>(1) << (kBits - 2)));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), RotateRight<1>(values2));

    // Rotate by half
    Store(values, d, expected.get());
    for (size_t i = 0; i < N; ++i) {
      expected[i] =
          ConvertScalarTo<T>((static_cast<TU>(expected[i]) >> (kBits / 2)) |
                             (static_cast<TU>(expected[i]) << (kBits / 2)));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), RotateRight<kBits / 2>(values));

    for (size_t i = 0; i < N; ++i) {
      expected[i] = ConvertScalarTo<T>(
          expected[i] ^ (static_cast<T>(1) << ((kBits / 2) - 1)));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), RotateRight<kBits / 2>(values2));

    // Rotate by max
    Store(values, d, expected.get());
    for (size_t i = 0; i < N; ++i) {
      expected[i] =
          ConvertScalarTo<T>((static_cast<TU>(expected[i]) >> (kBits - 1)) |
                             (static_cast<TU>(expected[i]) << 1));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), RotateRight<kBits - 1>(values));

    for (size_t i = 0; i < N; ++i) {
      expected[i] = ConvertScalarTo<T>(expected[i] ^ static_cast<T>(1));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), RotateRight<kBits - 1>(values2));
  }
};

HWY_NOINLINE void TestAllRotateRight() {
  ForIntegerTypes(ForPartialVectors<TestRotateRight>());
}

struct TestVariableUnsignedRightShifts {
  template <typename T, class D>
  HWY_NOINLINE void operator()(T /*unused*/, D d) {
    const size_t N = Lanes(d);
    auto expected = AllocateAligned<T>(N);
    HWY_ASSERT(expected);

    const auto v0 = Zero(d);
    const auto v1 = Set(d, 1);
    const auto values = Iota(d, 0);

    const T kMax = LimitsMax<T>();
    const auto max = Set(d, kMax);

    constexpr size_t kMaxShift = (sizeof(T) * 8) - 1;
    const auto max_shift = Set(d, kMaxShift);
    const auto small_shifts = And(Iota(d, 0), max_shift);
    const auto large_shifts = Sub(max_shift, small_shifts);

    // Same: 0
    HWY_ASSERT_VEC_EQ(d, values, Shr(values, v0));

    // Same: 1
    for (size_t i = 0; i < N; ++i) {
      expected[i] = static_cast<T>(static_cast<T>(i & kMax) >> 1);
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), Shr(values, v1));

    // Same: max
    for (size_t i = 0; i < N; ++i) {
      expected[i] = static_cast<T>(static_cast<T>(i & kMax) >> kMaxShift);
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), Shr(values, max_shift));

    // Variable: small
    for (size_t i = 0; i < N; ++i) {
      expected[i] = ConvertScalarTo<T>(T(i) >> (i & kMaxShift));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), Shr(values, small_shifts));

    // Variable: Large
    for (size_t i = 0; i < N; ++i) {
      expected[i] = ConvertScalarTo<T>(kMax >> (kMaxShift - (i & kMaxShift)));
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), Shr(max, large_shifts));
  }
};

template <int kAmount, typename T>
T RightShiftNegative(T val) {
  // C++ shifts are implementation-defined for negative numbers, and we have
  // seen divisions replaced with shifts, so resort to bit operations.
  using TU = hwy::MakeUnsigned<T>;
  TU bits;
  CopySameSize(&val, &bits);

  const TU shifted = TU(bits >> kAmount);

  const TU all = TU(~TU(0));
  const size_t num_zero = sizeof(TU) * 8 - 1 - kAmount;
  const TU sign_extended = static_cast<TU>((all << num_zero) & LimitsMax<TU>());

  bits = shifted | sign_extended;
  CopySameSize(&bits, &val);
  return val;
}

class TestSignedRightShifts {
 public:
  template <typename T, class D>
  HWY_NOINLINE void operator()(T /*unused*/, D d) {
    const size_t N = Lanes(d);
    auto expected = AllocateAligned<T>(N);
    HWY_ASSERT(expected);
    constexpr T kMin = LimitsMin<T>();
    constexpr T kMax = LimitsMax<T>();
    constexpr size_t kMaxShift = (sizeof(T) * 8) - 1;

    // First test positive values, negative are checked below.
    const auto v0 = Zero(d);
    const auto values = And(Iota(d, 0), Set(d, kMax));

    // Shift by 0
    HWY_ASSERT_VEC_EQ(d, values, ShiftRight<0>(values));
    HWY_ASSERT_VEC_EQ(d, values, ShiftRightSame(values, 0));

    // Shift by 1
    for (size_t i = 0; i < N; ++i) {
      expected[i] = static_cast<T>(static_cast<T>(i & kMax) >> 1);
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), ShiftRight<1>(values));
    HWY_ASSERT_VEC_EQ(d, expected.get(), ShiftRightSame(values, 1));

    // max
    HWY_ASSERT_VEC_EQ(d, v0, ShiftRight<kMaxShift>(values));
    HWY_ASSERT_VEC_EQ(d, v0, ShiftRightSame(values, kMaxShift));

    // Even negative value
    Test<0>(kMin, d, __LINE__);
    Test<1>(kMin, d, __LINE__);
    Test<2>(kMin, d, __LINE__);
    Test<kMaxShift>(kMin, d, __LINE__);

    const T odd = ConvertScalarTo<T>(kMin + 1);
    Test<0>(odd, d, __LINE__);
    Test<1>(odd, d, __LINE__);
    Test<2>(odd, d, __LINE__);
    Test<kMaxShift>(odd, d, __LINE__);
  }

 private:
  template <int kAmount, typename T, class D>
  void Test(T val, D d, int line) {
    const auto expected = Set(d, RightShiftNegative<kAmount>(val));
    const auto in = Set(d, val);
    const char* file = __FILE__;
    AssertVecEqual(d, expected, ShiftRight<kAmount>(in), file, line);
    AssertVecEqual(d, expected, ShiftRightSame(in, kAmount), file, line);
  }
};

struct TestVariableSignedRightShifts {
  template <typename T, class D>
  HWY_NOINLINE void operator()(T /*unused*/, D d) {
    using TU = MakeUnsigned<T>;
    const size_t N = Lanes(d);
    auto expected = AllocateAligned<T>(N);
    HWY_ASSERT(expected);

    constexpr T kMin = LimitsMin<T>();
    constexpr T kMax = LimitsMax<T>();

    constexpr size_t kMaxShift = (sizeof(T) * 8) - 1;

    // First test positive values, negative are checked below.
    const auto v0 = Zero(d);
    const auto positive = And(Iota(d, 0), Set(d, kMax));

    // Shift by 0
    HWY_ASSERT_VEC_EQ(d, positive, ShiftRight<0>(positive));
    HWY_ASSERT_VEC_EQ(d, positive, ShiftRightSame(positive, 0));

    // Shift by 1
    for (size_t i = 0; i < N; ++i) {
      expected[i] = static_cast<T>(static_cast<T>(i & kMax) >> 1);
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), ShiftRight<1>(positive));
    HWY_ASSERT_VEC_EQ(d, expected.get(), ShiftRightSame(positive, 1));

    // max
    HWY_ASSERT_VEC_EQ(d, v0, ShiftRight<kMaxShift>(positive));
    HWY_ASSERT_VEC_EQ(d, v0, ShiftRightSame(positive, kMaxShift));

    const auto max_shift = Set(d, kMaxShift);
    const auto small_shifts = And(Iota(d, 0), max_shift);
    const auto large_shifts = Sub(max_shift, small_shifts);

    const auto negative = Iota(d, kMin);

    // Test varying negative to shift
    for (size_t i = 0; i < N; ++i) {
      const T val = ConvertScalarTo<T>(static_cast<TU>(kMin) + i);
      expected[i] =
          (val < 0) ? RightShiftNegative<1>(val) : ConvertScalarTo<T>(val >> 1);
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), Shr(negative, Set(d, 1)));

    // Shift MSB right by small amounts
    for (size_t i = 0; i < N; ++i) {
      const size_t amount = i & kMaxShift;
      const TU shifted = static_cast<TU>(~((1ull << (kMaxShift - amount)) - 1));
      CopySameSize(&shifted, &expected[i]);
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), Shr(Set(d, kMin), small_shifts));

    // Shift MSB right by large amounts
    for (size_t i = 0; i < N; ++i) {
      const size_t amount = kMaxShift - (i & kMaxShift);
      const TU shifted = static_cast<TU>(~((1ull << (kMaxShift - amount)) - 1));
      CopySameSize(&shifted, &expected[i]);
    }
    HWY_ASSERT_VEC_EQ(d, expected.get(), Shr(Set(d, kMin), large_shifts));
  }
};

HWY_NOINLINE void TestAllShifts() {
  ForUnsignedTypes(ForPartialVectors<TestLeftShifts</*kSigned=*/false>>());
  ForSignedTypes(ForPartialVectors<TestLeftShifts</*kSigned=*/true>>());
  ForUnsignedTypes(ForPartialVectors<TestUnsignedRightShifts>());
  ForSignedTypes(ForPartialVectors<TestSignedRightShifts>());
}

HWY_NOINLINE void TestAllVariableShifts() {
  ForUnsignedTypes(
      ForPartialVectors<TestVariableLeftShifts</*kSigned=*/false>>());
  ForSignedTypes(ForPartialVectors<TestVariableLeftShifts</*kSigned=*/true>>());
  ForUnsignedTypes(ForPartialVectors<TestVariableUnsignedRightShifts>());
  ForSignedTypes(ForPartialVectors<TestVariableSignedRightShifts>());
}

struct TestVariableRotations {
  template <typename T, class D>
  HWY_NOINLINE void operator()(T /*unused*/, D d) {
    using TU = MakeUnsigned<T>;

    constexpr TU kBits1 = static_cast<TU>(0x7C29085C41482973ULL);
    constexpr TU kBits2 = static_cast<TU>(0xD3C8835FBD1A89BAULL);

    const auto viota0 = Iota(d, 0);
    const auto va = Xor(Set(d, static_cast<T>(kBits1)), viota0);
    const auto vb = Xor(Set(d, static_cast<T>(kBits2)), viota0);

    const size_t N = Lanes(d);
    auto expected1 = AllocateAligned<T>(N);
    auto expected2 = AllocateAligned<T>(N);
    auto expected3 = AllocateAligned<T>(N);
    auto expected4 = AllocateAligned<T>(N);
    HWY_ASSERT(expected1 && expected2 && expected3 && expected4);

    constexpr size_t kBits = sizeof(T) * 8;

    auto vrotate_amt1 = viota0;
    auto vrotate_amt2 = Sub(Set(d, static_cast<T>(kBits)), viota0);
    auto vrotate_amt_incr = Set(d, static_cast<T>(N));

    const RebindToSigned<decltype(d)> di;

    for (size_t i = 0; i < kBits; i += N) {
      for (size_t j = 0; j < N; j++) {
        const size_t shift_amt_1 = (i + j) & (kBits - 1);
        const size_t shift_amt_2 = (size_t{0} - shift_amt_1) & (kBits - 1);

        const TU val_a = static_cast<TU>(kBits1 ^ j);
        const TU val_b = static_cast<TU>(kBits2 ^ j);

        expected1[j] =
            static_cast<T>((val_a << shift_amt_1) | (val_a >> shift_amt_2));
        expected2[j] =
            static_cast<T>((val_a >> shift_amt_1) | (val_a << shift_amt_2));
        expected3[j] =
            static_cast<T>((val_b << shift_amt_1) | (val_b >> shift_amt_2));
        expected4[j] =
            static_cast<T>((val_b >> shift_amt_1) | (val_b << shift_amt_2));
      }

      const auto vrotate_amt3 = BitCast(d, Neg(BitCast(di, vrotate_amt1)));
      const auto vrotate_amt4 = BitCast(d, Neg(BitCast(di, vrotate_amt2)));

      HWY_ASSERT_VEC_EQ(d, expected1.get(), Rol(va, vrotate_amt1));
      HWY_ASSERT_VEC_EQ(d, expected2.get(), Ror(va, vrotate_amt1));
      HWY_ASSERT_VEC_EQ(d, expected3.get(), Rol(vb, vrotate_amt1));
      HWY_ASSERT_VEC_EQ(d, expected4.get(), Ror(vb, vrotate_amt1));

      HWY_ASSERT_VEC_EQ(d, expected1.get(), Ror(va, vrotate_amt2));
      HWY_ASSERT_VEC_EQ(d, expected2.get(), Rol(va, vrotate_amt2));
      HWY_ASSERT_VEC_EQ(d, expected3.get(), Ror(vb, vrotate_amt2));
      HWY_ASSERT_VEC_EQ(d, expected4.get(), Rol(vb, vrotate_amt2));

      HWY_ASSERT_VEC_EQ(d, expected1.get(), Ror(va, vrotate_amt3));
      HWY_ASSERT_VEC_EQ(d, expected2.get(), Rol(va, vrotate_amt3));
      HWY_ASSERT_VEC_EQ(d, expected3.get(), Ror(vb, vrotate_amt3));
      HWY_ASSERT_VEC_EQ(d, expected4.get(), Rol(vb, vrotate_amt3));

      HWY_ASSERT_VEC_EQ(d, expected1.get(), Rol(va, vrotate_amt4));
      HWY_ASSERT_VEC_EQ(d, expected2.get(), Ror(va, vrotate_amt4));
      HWY_ASSERT_VEC_EQ(d, expected3.get(), Rol(vb, vrotate_amt4));
      HWY_ASSERT_VEC_EQ(d, expected4.get(), Ror(vb, vrotate_amt4));

      vrotate_amt1 = Add(vrotate_amt1, vrotate_amt_incr);
      vrotate_amt2 = Sub(vrotate_amt2, vrotate_amt_incr);
    }

    for (int i = 0; i < static_cast<int>(kBits); ++i) {
      for (size_t j = 0; j < N; j++) {
        const int shift_amt_2 =
            static_cast<int>(static_cast<size_t>(-i) & (kBits - 1));

        const TU val_a = static_cast<TU>(kBits1 ^ j);
        const TU val_b = static_cast<TU>(kBits2 ^ j);

        expected1[j] = static_cast<T>((val_a << i) | (val_a >> shift_amt_2));
        expected2[j] = static_cast<T>((val_a >> i) | (val_a << shift_amt_2));
        expected3[j] = static_cast<T>((val_b << i) | (val_b >> shift_amt_2));
        expected4[j] = static_cast<T>((val_b >> i) | (val_b << shift_amt_2));
      }

      HWY_ASSERT_VEC_EQ(d, expected1.get(), RotateLeftSame(va, i));
      HWY_ASSERT_VEC_EQ(d, expected2.get(), RotateRightSame(va, i));
      HWY_ASSERT_VEC_EQ(d, expected3.get(), RotateLeftSame(vb, i));
      HWY_ASSERT_VEC_EQ(d, expected4.get(), RotateRightSame(vb, i));

      HWY_ASSERT_VEC_EQ(d, expected1.get(), RotateRightSame(va, -i));
      HWY_ASSERT_VEC_EQ(d, expected2.get(), RotateLeftSame(va, -i));
      HWY_ASSERT_VEC_EQ(d, expected3.get(), RotateRightSame(vb, -i));
      HWY_ASSERT_VEC_EQ(d, expected4.get(), RotateLeftSame(vb, -i));

      HWY_ASSERT_VEC_EQ(d, expected1.get(),
                        RotateRightSame(va, static_cast<int>(kBits) - i));
      HWY_ASSERT_VEC_EQ(d, expected2.get(),
                        RotateLeftSame(va, static_cast<int>(kBits) - i));
      HWY_ASSERT_VEC_EQ(d, expected3.get(),
                        RotateRightSame(vb, static_cast<int>(kBits) - i));
      HWY_ASSERT_VEC_EQ(d, expected4.get(),
                        RotateLeftSame(vb, static_cast<int>(kBits) - i));
    }
  }
};

HWY_NOINLINE void TestAllVariableRotations() {
  ForIntegerTypes(ForPartialVectors<TestVariableRotations>());
}

// NOLINTNEXTLINE(google-readability-namespace-comments)
}  // namespace HWY_NAMESPACE
}  // namespace hwy
HWY_AFTER_NAMESPACE();

#if HWY_ONCE

namespace hwy {
HWY_BEFORE_TEST(HwyShiftTest);
HWY_EXPORT_AND_TEST_P(HwyShiftTest, TestAllShifts);
HWY_EXPORT_AND_TEST_P(HwyShiftTest, TestAllVariableShifts);
HWY_EXPORT_AND_TEST_P(HwyShiftTest, TestAllRotateLeft);
HWY_EXPORT_AND_TEST_P(HwyShiftTest, TestAllRotateRight);
HWY_EXPORT_AND_TEST_P(HwyShiftTest, TestAllVariableRotations);
HWY_AFTER_TEST();
}  // namespace hwy

#endif