/*
 *  Copyright (c) 2012 The WebM project authors. All Rights Reserved.
 *
 *  Use of this source code is governed by a BSD-style license
 *  that can be found in the LICENSE file in the root of the source
 *  tree. An additional intellectual property rights grant can be found
 *  in the file PATENTS.  All contributing project authors may
 *  be found in the AUTHORS file in the root of the source tree.
 */

#include <smmintrin.h> /* SSE4.1 */

#include "./vp8_rtcd.h"
#include "vp8/encoder/block.h"
#include "vpx_ports/bitops.h" /* get_lsb */
#include "vpx_ports/compiler_attributes.h"

// Unsigned shift overflow is disabled for the use of ~1U << eob with ymask.
VPX_NO_UNSIGNED_SHIFT_CHECK void vp8_regular_quantize_b_sse4_1(BLOCK *b,
                                                               BLOCKD *d) {
  int eob = -1;
  short *zbin_boost_ptr = b->zrun_zbin_boost;
  __m128i zbin_boost0 = _mm_load_si128((__m128i *)(zbin_boost_ptr));
  __m128i zbin_boost1 = _mm_load_si128((__m128i *)(zbin_boost_ptr + 8));
  __m128i x0, x1, y0, y1, x_minus_zbin0, x_minus_zbin1, dqcoeff0, dqcoeff1;
  __m128i quant_shift0 = _mm_load_si128((__m128i *)(b->quant_shift));
  __m128i quant_shift1 = _mm_load_si128((__m128i *)(b->quant_shift + 8));
  __m128i z0 = _mm_load_si128((__m128i *)(b->coeff));
  __m128i z1 = _mm_load_si128((__m128i *)(b->coeff + 8));
  __m128i zbin_extra = _mm_cvtsi32_si128(b->zbin_extra);
  __m128i zbin0 = _mm_load_si128((__m128i *)(b->zbin));
  __m128i zbin1 = _mm_load_si128((__m128i *)(b->zbin + 8));
  __m128i round0 = _mm_load_si128((__m128i *)(b->round));
  __m128i round1 = _mm_load_si128((__m128i *)(b->round + 8));
  __m128i quant0 = _mm_load_si128((__m128i *)(b->quant));
  __m128i quant1 = _mm_load_si128((__m128i *)(b->quant + 8));
  __m128i dequant0 = _mm_load_si128((__m128i *)(d->dequant));
  __m128i dequant1 = _mm_load_si128((__m128i *)(d->dequant + 8));
  __m128i qcoeff0, qcoeff1, t0, t1, x_shuf0, x_shuf1;
  uint32_t mask, ymask;
  DECLARE_ALIGNED(16, static const uint8_t,
                  zig_zag_mask[16]) = { 0, 1,  4,  8,  5, 2,  3,  6,
                                        9, 12, 13, 10, 7, 11, 14, 15 };
  DECLARE_ALIGNED(16, uint16_t, qcoeff[16]) = { 0 };

  /* Duplicate to all lanes. */
  zbin_extra = _mm_shufflelo_epi16(zbin_extra, 0);
  zbin_extra = _mm_unpacklo_epi16(zbin_extra, zbin_extra);

  /* x = abs(z) */
  x0 = _mm_abs_epi16(z0);
  x1 = _mm_abs_epi16(z1);

  /* zbin[] + zbin_extra */
  zbin0 = _mm_add_epi16(zbin0, zbin_extra);
  zbin1 = _mm_add_epi16(zbin1, zbin_extra);

  /* In C x is compared to zbin where zbin = zbin[] + boost + extra. Rebalance
   * the equation because boost is the only value which can change:
   * x - (zbin[] + extra) >= boost */
  x_minus_zbin0 = _mm_sub_epi16(x0, zbin0);
  x_minus_zbin1 = _mm_sub_epi16(x1, zbin1);

  /* All the remaining calculations are valid whether they are done now with
   * simd or later inside the loop one at a time. */
  x0 = _mm_add_epi16(x0, round0);
  x1 = _mm_add_epi16(x1, round1);

  y0 = _mm_mulhi_epi16(x0, quant0);
  y1 = _mm_mulhi_epi16(x1, quant1);

  y0 = _mm_add_epi16(y0, x0);
  y1 = _mm_add_epi16(y1, x1);

  /* Instead of shifting each value independently we convert the scaling
   * factor with 1 << (16 - shift) so we can use multiply/return high half. */
  y0 = _mm_mulhi_epi16(y0, quant_shift0);
  y1 = _mm_mulhi_epi16(y1, quant_shift1);

  /* Restore the sign. */
  y0 = _mm_sign_epi16(y0, z0);
  y1 = _mm_sign_epi16(y1, z1);

  {
    const __m128i zig_zag_i16_0 =
        _mm_setr_epi8(0, 1, 2, 3, 8, 9, 14, 15, 10, 11, 4, 5, 6, 7, 12, 13);
    const __m128i zig_zag_i16_1 =
        _mm_setr_epi8(0, 1, 6, 7, 8, 9, 2, 3, 14, 15, 4, 5, 10, 11, 12, 13);

    /* The first part of the zig zag needs a value
     * from x_minus_zbin1 and vice versa. */
    t1 = _mm_alignr_epi8(x_minus_zbin1, x_minus_zbin1, 2);
    t0 = _mm_blend_epi16(x_minus_zbin0, t1, 0x80);
    t1 = _mm_blend_epi16(t1, x_minus_zbin0, 0x80);
    x_shuf0 = _mm_shuffle_epi8(t0, zig_zag_i16_0);
    x_shuf1 = _mm_shuffle_epi8(t1, zig_zag_i16_1);
  }

  /* Check if y is nonzero and put it in zig zag order. */
  t0 = _mm_packs_epi16(y0, y1);
  t0 = _mm_cmpeq_epi8(t0, _mm_setzero_si128());
  t0 = _mm_shuffle_epi8(t0, _mm_load_si128((const __m128i *)zig_zag_mask));
  ymask = _mm_movemask_epi8(t0) ^ 0xffff;

  for (;;) {
    t0 = _mm_cmpgt_epi16(zbin_boost0, x_shuf0);
    t1 = _mm_cmpgt_epi16(zbin_boost1, x_shuf1);
    t0 = _mm_packs_epi16(t0, t1);
    mask = _mm_movemask_epi8(t0);
    mask = ~mask & ymask;
    if (!mask) break;
    /* |eob| will contain the index of the next found element where:
     * boost[i - old_eob - 1] <= x[zigzag[i]] && y[zigzag[i]] != 0 */
    eob = get_lsb(mask);
    /* Need to clear the mask from processed elements so that
     * they are no longer counted in the next iteration. */
    ymask &= ~1U << eob;
    /* It's safe to read ahead of this buffer if struct VP8_COMP has at
     * least 32 bytes before the zrun_zbin_boost_* fields (it has 384).
     * Any data read outside of the buffer is masked by the updated |ymask|. */
    zbin_boost0 = _mm_loadu_si128((__m128i *)(zbin_boost_ptr - eob - 1));
    zbin_boost1 = _mm_loadu_si128((__m128i *)(zbin_boost_ptr - eob + 7));
    qcoeff[zig_zag_mask[eob]] = 0xffff;
  }

  qcoeff0 = _mm_load_si128((__m128i *)(qcoeff));
  qcoeff1 = _mm_load_si128((__m128i *)(qcoeff + 8));
  qcoeff0 = _mm_and_si128(qcoeff0, y0);
  qcoeff1 = _mm_and_si128(qcoeff1, y1);

  _mm_store_si128((__m128i *)(d->qcoeff), qcoeff0);
  _mm_store_si128((__m128i *)(d->qcoeff + 8), qcoeff1);

  dqcoeff0 = _mm_mullo_epi16(qcoeff0, dequant0);
  dqcoeff1 = _mm_mullo_epi16(qcoeff1, dequant1);

  _mm_store_si128((__m128i *)(d->dqcoeff), dqcoeff0);
  _mm_store_si128((__m128i *)(d->dqcoeff + 8), dqcoeff1);

  *d->eob = eob + 1;
}