// moderngpu copyright (c) 2016, Sean Baxter http://www.moderngpu.com
#pragma once

#include "cta_scan.hxx"

BEGIN_MGPU_NAMESPACE

template<typename type_t>
struct segscan_result_t {
  type_t scan;
  type_t reduction;
  bool has_carry_in;
  int left_lane;
};

template<int nt, typename type_t>
struct cta_segscan_t {
  enum { num_warps = nt / warp_size };

  union storage_t {
    int delta[num_warps + nt]; 
    struct { type_t values[2 * nt]; int packed[nt]; };
  };

  MGPU_DEVICE int find_left_lane(int tid, bool has_head_flag, 
    storage_t& storage) const {

    int warp = tid / warp_size;
    int lane = (warp_size - 1) & tid;
    int warp_mask = 0xffffffff>> (31 - lane);   // inclusive search.
    int cta_mask = 0x7fffffff>> (31 - lane);    // exclusive search.

    // Build a head flag bitfield and store it into shared memory.
    int warp_bits = ballot(has_head_flag);
    storage.delta[warp] = warp_bits;
    __syncthreads();

    if(tid < num_warps) {
      unsigned mask = __activemask();
      int cta_bits = ballot(0 != storage.delta[tid], mask);
      int warp_segment = 31 - clz(cta_mask & cta_bits);
      int start = (-1 != warp_segment) ?
        (31 - clz(storage.delta[warp_segment]) + 32 * warp_segment) : 0;
      storage.delta[num_warps + tid] = start;
    }
    __syncthreads();

    // Find the closest flag to the left of this thread within the warp.
    // Include the flag for this thread.
    int start = 31 - clz(warp_mask & warp_bits);
    if(-1 != start) start += ~31 & tid;
    else start = storage.delta[num_warps + warp];
    __syncthreads();

    return start;
  }

  template<typename op_t = plus_t<type_t> >
  MGPU_DEVICE segscan_result_t<type_t> segscan(int tid, bool has_head_flag,
    bool has_carry_out, type_t x, storage_t& storage, type_t init = type_t(),
    op_t op = op_t()) const {

    if(!has_carry_out) x = init;

    int left_lane = find_left_lane(tid, has_head_flag, storage);
    int tid_delta = tid - left_lane;

    // Store the has_carry_out flag.
    storage.packed[tid] = (int)has_carry_out | (left_lane<< 1);

    // Run an inclusive scan.
    int first = 0;
    storage.values[first + tid] = x;
    __syncthreads();

    int packed = storage.packed[left_lane];
    left_lane = packed>> 1;
    tid_delta = tid - left_lane;
    if(0 == (1 & packed)) --tid_delta;

    iterate<s_log2(nt)>([&](int pass) {
      int offset = 1<< pass;
      if(tid_delta >= offset)
        x = op(x, storage.values[first + tid - offset]);
      first = nt - first;
      storage.values[first + tid] = x;
      __syncthreads();
    });

    // Get the exclusive scan by fetching the preceding element. Also return
    // the carry-out value as the total.
    bool has_carry_in = tid ? (0 != (1 & storage.packed[tid - 1])) : false;

    segscan_result_t<type_t> result { 
      (has_carry_in && tid) ? storage.values[first + tid - 1] : init,
      storage.values[first + nt - 1],
      has_carry_in,
      left_lane
    };
    __syncthreads();

    return result;
  }
};

END_MGPU_NAMESPACE
