Implemented 'clear' (requires testing), and began implementation of 'balance'.

use Collection;
use BlockDist;

/*
  A highly parallel segmented multiset. Each node gets its own bag, and in each bag
  it is segmented into 'here.maxTaskPar' segments. Segments allow for actual parallelism
  while operating in that it enables us to manage 'best-case', 'average-case',
  and 'worst-case' scenarios by making multiple passes over each segment. In the
  case where there is no oversubscription, the best-case will always be achieved
  (considering any other conditions are also met), while in the case of oversubscription
  or, for example a near empty bag, we fall into the 'average-case', etc. Examples
  of 'best-case' scenarios for a removal would be when a segment is unlocked
  and contains some elements we can drain, and the 'average-case' would be to find
  any segment, unlocked or not, that contains elements we can drain, and so on.

  This data structure also employs its own load balancing algorithm, beginning
  with a best-effort round-robin algorithm such that each task begins searching
  at another segment (with the added benefit of reducing overall contention), in
  particular is useful for locally distributing insertion operations. Next it also
  employs a work-stealing algorithm that horizontally steals a ratio of elements
  across segments. A segment is considered 'horizontal' if it shares the same
  segment index within [0, here.maxTaskPar); those other segments likely are also
  benefitted from local distribution from the round-robin algorithm discussed above.
  We also steal a ratio of elements, say 25%, because it leaves all other segments
  with 75% of their work; the more elements they have, the more we take, but the
  less they have, the less we steal; this also has the added benefit of reducing
  unnecessary work stealing between segments when the bag is nearly emptied. Furthermore,
  stealing horizontally ensures that all segments remain filled and that we still
  achieve parallelism across segments for removal operations. Lastly, we attempt
  to steal a maximum of `N / sizeof(eltType)`, where N is some size in megabytes
  (representing how much data can be sent in one network request), which keeps
  down excessive communication.

  This data structure does not come without flaws; as work stealing is dynamic
  and trigered on demand, work stealing can still be performed in excess, which
  dramatically causes a performance drop; this data structure scales in terms of
  both processors per node, and in terms of raw nodes themselves. As well, due
  to how work stealing is, we perform the best when we have large amounts of data
  to work on, but do not do well with smaller amounts due to excessive work stealing.
  As well, to achieve true parallelism, one must use a privatized instance, specific
  to their node to avoid communication overhead of accessing class fields.
*/

/*
  Below are segment statuses, which is a way to make visible to outsiders the
  current ongoing operation.
*/
private const STATUS_UNLOCKED : uint = 0;
private const STATUS_ADD : uint = 1;
private const STATUS_REMOVE : uint = 2;
private const STATUS_LOOKUP : uint = 3;
private const STATUS_BALANCE : uint = 4;

/*
  Below are statuses specific to the work stealing algorithm.
*/
private const WS_INITIALIZED : int = -1;
private const WS_FINISHED_WITH_NO_WORK : int = 0;
private const WS_FINISHED_WITH_WORK : int = 1;

/*
  Frozen states... If we are FREEZE_UNFROZEN, we are mutable. If we are FREEZE_FROZEN,
  we are immutable. If we are FREEZE_MARKED, we are in the middle of a state change.
*/
private const FREEZE_UNFROZEN = 0;
private const FREEZE_MARKED = 1;
private const FREEZE_FROZEN = 2;

/*
  The phases for operations.
*/
private const ADD_BEST_CASE = 0;
private const ADD_AVERAGE_CASE = 1;
private const REMOVE_BEST_CASE = 2;
private const REMOVE_AVERAGE_CASE = 3;
private const REMOVE_WORST_CASE = 4;

/*
  The initial amount of elements in an unroll block. Each successive unroll block
  is double the size of it's predecessor.
*/
config param INITIAL_BLOCK_SIZE = 1024;
// The amount we attempt to steal from each segment of other nodes. Helps to
// further ensure a proper load balance by not leaving segments empty by stealing
// on average 1/4 of their work, leaving them with 3/4. As we steal from *all* other
// segments, it is imperative that this stays low to prevent a ping-pong effect on
// work propagating from excessive work stealing. We steal at least one element,
// but no more than the amount of data that can be transferred within a single
// bulk transfer, measured in megabytes.
config param WORK_STEALING_RATIO = 0.25;
config param WORK_STEALING_MAX_MEMORY_MB : real = 1.0;
config param BAG_NO_FREEZE = false;
// The amount of elements a segment must have to be stolen from.
config param WORK_STEALING_MINIMUM = 1;

/*
  A segment block is an unrolled linked list node that holds a contiguous buffer
  of memory. Each segment block *should* be a power of two, as we increase the
  size of each subsequent unroll block by twice the size. This is so that stealing
  work is faster in that majority of elements are confined to one area.
*/
pragma "no doc"
class BagSegmentBlock {
  type eltType;

  // Contiguous memory containing all elements
  var elems :  c_ptr(eltType);
  var next : BagSegmentBlock(eltType);

  // The capacity of this block.
  var cap : int;
  // The number of occupied elements in this block.
  var size : int;

  inline proc isEmpty {
    return size == 0;
  }

  inline proc isFull {
    return size == cap;
  }

  inline proc push(elt : eltType) {
    if elems == nil {
      writeln("Cap: ", cap, "; elems: ", elems : string);
      halt("'elems' is nil");
    }

    if isFull {
      halt("SegmentBlock is Full");
    }

    elems[size] = elt;
    size = size + 1;
  }

  inline proc pop() : eltType {
    if elems == nil {
      halt("'elems' is nil");
    }

    if isEmpty {
      halt("SegmentBlock is Empty");
    }

    size = size - 1;
    var elt = elems[size];
    return elt;
  }

  proc BagSegmentBlock(type eltType, capacity) {
    if capacity == 0 {
      halt("Capacity is 0...");
    }

    cap = capacity;
    elems = c_malloc(eltType, capacity);
  }

  proc BagSegmentBlock(type eltType, ptr, capacity) {
    cap = capacity;
    elems = ptr;
    size = cap;
  }

  proc ~BagSegmentBlock() {
    c_free(elems);
  }
}

/*
  A segment is, in and of itself an unrolled linked list. We maintain one per core
  to ensure maximum parallelism.
*/
pragma "no doc"
record BagSegment {
  type eltType;

  // Used as a test-and-test-and-set spinlock.
  var status : atomic uint;

  var headBlock : BagSegmentBlock(eltType);
  var tailBlock : BagSegmentBlock(eltType);

  var nElems : atomic uint;

  inline proc isEmpty {
    return nElems.read() == 0;
  }

  inline proc acquireWithStatus(newStatus) {
    return status.compareExchangeStrong(STATUS_UNLOCKED, newStatus);
  }

  // Set status with a test-and-test-and-set loop...
  inline proc acquire(newStatus) {
    while true {
      if currentStatus == STATUS_UNLOCKED && acquireWithStatus(newStatus) {
        break;
      }

      chpl_task_yield();
    }
  }

  // Set status with a test-and-test-and-set loop, but only while it is not empty...
  inline proc acquireIfNonEmpty(newStatus) {
    while !isEmpty {
      if currentStatus == STATUS_UNLOCKED && acquireWithStatus(newStatus) {
        if isEmpty {
          releaseStatus();
          return false;
        } else {
          return true;
        }
      }

      chpl_task_yield();
    }

    return false;
  }

  inline proc isUnlocked {
    return status.read() == STATUS_UNLOCKED;
  }

  inline proc currentStatus {
    return status.read();
  }

  inline proc releaseStatus() {
    status.write(STATUS_UNLOCKED);
  }

  inline proc transferElements(destPtr, n, locId = here.id) {
    var destOffset = 0;
    var srcOffset = 0;
    while destOffset < n {
      if headBlock == nil || isEmpty {
        halt(here, ": Attempted transfer ", n, " elements to ", locId, " but failed... destOffset=", destOffset);
      }

      var len = headBlock.size;
      var need = n - destOffset;
      // If the amount in this block is greater than what is left to transfer, we
      // cannot begin transferring at the beginning, so we set our offset from the end.
      if len > need {
        srcOffset = len - need;
        headBlock.size = srcOffset;
        __primitive("chpl_comm_array_put", headBlock.elems[srcOffset], locId, destPtr[destOffset], need);
        destOffset = destOffset + need;
      } else {
        srcOffset = 0;
        headBlock.size = 0;
        __primitive("chpl_comm_array_put", headBlock.elems[srcOffset], locId, destPtr[destOffset], len);
        destOffset = destOffset + len;
      }

      // Fix list if we consumed last one...
      if headBlock.isEmpty {
        var tmp = headBlock;
        headBlock = headBlock.next;
        delete tmp;

        if headBlock == nil then tailBlock = nil;
      }
    }

    nElems.sub(n : uint);
  }

  inline proc addElementsPtr(ptr, n) {
    var offset = 0;
    while offset < n {
      var block = tailBlock;
      // Empty? Create a new one of initial size
      if block == nil then {
        tailBlock = new BagSegmentBlock(eltType, INITIAL_BLOCK_SIZE);
        headBlock = tailBlock;
        block = tailBlock;
      }

      // Full? Create a new one double the previous size
      if block.isFull {
        block.next = new BagSegmentBlock(eltType, block.cap * 2);
        tailBlock = block.next;
        block = block.next;
      }

      var nLeft = n - offset;
      var nSpace = block.cap - block.size;
      var nFill = min(nLeft, nSpace);
      __primitive("chpl_comm_array_put", ptr[offset], here.id, block.elems[block.size], nFill);
      block.size = block.size + nFill;
      offset = offset + nFill;
    }

    nElems.add(n : uint);
  }

  inline proc takeElements(n) {
    var iterations = n : int;
    var arr : [{0..#n : int}] eltType;
    var arrIdx = 0;

    for 1 .. n : int{
      if isEmpty {
        halt("Attempted to take ", n, " elements when insufficient");
      }

      if headBlock.isEmpty {
        halt("Iterating over ", n, " elements with headBlock empty but nElems is ", nElems.read());
      }

      arr[arrIdx] = headBlock.pop();
      arrIdx = arrIdx + 1;
      nElems.sub(1);

      // Fix list if we consumed last one...
      if headBlock.isEmpty {
        var tmp = headBlock;
        headBlock = headBlock.next;
        delete tmp;

        if headBlock == nil then tailBlock = nil;
      }
    }

    return arr;
  }

  inline proc takeElement() {
    if isEmpty {
      return (false, _defaultOf(eltType));
    }

    if headBlock.isEmpty {
      halt("Iterating over 1 element with headBlock empty but nElems is ", nElems.read());
    }

    var elem = headBlock.pop();
    nElems.sub(1);

    // Fix list if we consumed last one...
    if headBlock.isEmpty {
      var tmp = headBlock;
      headBlock = headBlock.next;
      delete tmp;

      if headBlock == nil then tailBlock = nil;
    }

    return (true, elem);
  }

  inline proc addElements(elt : eltType) {
    var block = tailBlock;

    // Empty? Create a new one of initial size
    if block == nil then {
      tailBlock = new BagSegmentBlock(eltType, INITIAL_BLOCK_SIZE);
      headBlock = tailBlock;
      block = tailBlock;
    }

    // Full? Create a new one double the previous size
    if block.isFull {
      block.next = new BagSegmentBlock(eltType, block.cap * 2);
      tailBlock = block.next;
      block = block.next;
    }

    block.push(elt);
    nElems.add(1);
  }

  inline proc addElements(elts) {
    for elt in elts do addElements(elt);
  }
}

/*
  We maintain a multiset 'bag' per node. Each bag keeps a handle to it's parent,
  which is required for work stealing.
*/
pragma "no doc"
class Bag {
  type eltType;

  // A handle to our parent 'distributed' bag, which is needed for work stealing.
  var parentHandle : DistributedBag(eltType);

  /*
    Helps evenly distribute and balance placement of elements in a best-effort
    round-robin approach. In the case where we have parallel enqueues or dequeues,
    they are less likely overlap with each other. Furthermore, it increases our
    chance to find our 'ideal' segment.
  */
  var startIdxEnq : atomic uint;
  var startIdxDeq : atomic uint;

  /*
    If a task makes 2 complete passes (1 best-case, 1 average-case) and has not
    found enough items, then it may attempt to balance the load distribution.
  */
  var loadBalanceInProgress : atomic bool;

  var maxParallelSegmentSpace = {0 .. #here.maxTaskPar};
  var segments : [maxParallelSegmentSpace] BagSegment(eltType);

  inline proc nextStartIdxEnq {
    return (startIdxEnq.fetchAdd(1) % here.maxTaskPar : uint) : int;
  }

  inline proc nextStartIdxDeq {
    return (startIdxDeq.fetchAdd(1) % here.maxTaskPar : uint) : int;
  }

  proc Bag(type eltType, parentHandle) {
    this.parentHandle = parentHandle;
  }

  proc add(elt : eltType) : bool {
    var startIdx = nextStartIdxEnq : int;
    var phase = ADD_BEST_CASE;

    while true {
      select phase {
        /*
          Pass 1: Best Case

          Find a segment that is unlocked and attempt to acquire it. As we are adding
          elements, we don't care how many elements there are, just that we find
          some place to add ours.
        */
        when ADD_BEST_CASE {
          for offset in 0 .. #here.maxTaskPar {
            ref segment = segments[(startIdx + offset) % here.maxTaskPar];

            // Attempt to acquire...
            if segment.acquireWithStatus(STATUS_ADD) {
              segment.addElements(elt);
              segment.releaseStatus();
              return true;
            }
          }

          phase = ADD_AVERAGE_CASE;
        }
        /*
          Pass 2: Average Case

          Find any segment (locked or unlocked) and make an attempt to acquire it.
        */
        when ADD_AVERAGE_CASE {
          ref segment = segments[startIdx];

          while true {
            select segment.currentStatus {
              // Quick acquire...
              when STATUS_UNLOCKED do {
                if segment.acquireWithStatus(STATUS_ADD) {
                  segment.addElements(elt);
                  segment.releaseStatus();
                  return true;
                }
              }
            }
            chpl_task_yield();
          }
        }
      }
    }

    halt("DEADCODE");
  }

  proc remove() : (bool, eltType) {
    var startIdx = nextStartIdxDeq;
    var idx = startIdx;
    var iterations = 0;
    var phase = REMOVE_BEST_CASE;
    var backoff = 0;

    while true {
      select phase {
        /*
          Pass 1: Best Case

          Find the first bucket that is both unlocked and contains elements. This is
          extremely helpful in the case where we have a good distribution of elements
          in each segment.
        */
        when REMOVE_BEST_CASE {
          while iterations < here.maxTaskPar {
            ref segment = segments[idx];

            // Attempt to acquire...
            if !segment.isEmpty && segment.acquireWithStatus(STATUS_REMOVE) {
              var (hasElem, elem) : (bool, eltType) = segment.takeElement();
              segment.releaseStatus();

              if hasElem {
                return (hasElem, elem);
              }
            }

            iterations = iterations + 1;
            idx = (idx + 1) % here.maxTaskPar;
          }

          phase = REMOVE_AVERAGE_CASE;
        }

        /*
          Pass 2: Average Case

          Find the first bucket containing elements. We don't care if it is locked
          or unlocked this time, just that it contains elements; this handles majority
          of cases where we have elements anywhere in any segment.
        */
        when REMOVE_AVERAGE_CASE {
          while iterations < here.maxTaskPar {
            ref segment = segments[idx];

            // Attempt to acquire...
            while !segment.isEmpty {
              if segment.isUnlocked && segment.acquireWithStatus(STATUS_REMOVE) {
                var (hasElem, elem) : (bool, eltType) = segment.takeElement();
                segment.releaseStatus();

                if hasElem {
                  return (hasElem, elem);
                }
              }

              // Backoff
              chpl_task_yield();
            }

            iterations = iterations + 1;
            idx = (idx + 1) % here.maxTaskPar;
          }

          phase = REMOVE_WORST_CASE;
        }

        /*
          Pass 3: Worst Case

          After two full iterations, we're sure the queue is full at this point, so we
          can attempt to steal work from other nodes. In this pass, we find *any* segment
          and if it is empty, we attempt to become the work-stealer; if someone else is the
          current work stealer we assist them instead and lift an element for ourselves.

          Furthermore, in this phase we loop indefinitely until we are 100% certain it is
          empty or we get an item, so introduce some backoff here.
        */
        when REMOVE_WORST_CASE {
          while true {
            ref segment = segments[idx];

            select segment.currentStatus {
              // Quick acquire
              when STATUS_UNLOCKED {
                if segment.acquireWithStatus(STATUS_REMOVE) {
                  // We're lucky; another element has been added to the current segment,
                  // take it and leave like normal...
                  if !segment.isEmpty {
                    var (hasElem, elem) : (bool, eltType) = segment.takeElement();
                    segment.releaseStatus();
                    return (hasElem, elem);
                  }

                  if parentHandle.targetLocales.size == 1 {
                    segment.releaseStatus();
                    return (false, _defaultOf(eltType));
                  }

                  // Attempt to become the sole work stealer for this node. If we
                  // do not, we spin until they finish. We need to release the lock
                  // on our segment so our segment may be load balanced as well.
                  if loadBalanceInProgress.testAndSet() {
                    segment.releaseStatus();

                    loadBalanceInProgress.waitFor(false);
                    // Reset our phase and scan for more elements...
                    phase = REMOVE_BEST_CASE;
                    break;
                  }

                  // We are the sole work stealer, and so it is our responsibility
                  // to balance the load for our node. We fork-join new worker
                  // tasks that will check horizontally across each node (as in
                  // across each segment with the same index), and vertically across
                  // each segment (each segment in a node). Horizontally, we steal
                  // at most a % of work from other nodes to give to ourselves.
                  // As load balancer, we also are the only one who knows whether
                  // or not all bags are empty.
                  var isEmpty : atomic bool;
                  isEmpty.write(true);
                  segment.releaseStatus();
                  coforall segmentIdx in 0..#here.maxTaskPar {
                    var stolenWork : [{0..#numLocales}] (int, c_ptr(eltType));
                    coforall loc in parentHandle.targetLocales {
                      if loc != here then on loc {
                        // As we jumped to the target node, 'localBag' returns
                        // the target's bag that we are attempting to steal from.
                        var targetBag = parentHandle.bag;

                        // Only proceed if the target is not load balancing themselves...
                        if !targetBag.loadBalanceInProgress.read() {
                          ref targetSegment = targetBag.segments[segmentIdx];

                          // As we only care that the segment contains data,
                          // we test-and-test-and-set until we gain ownership.
                          while targetSegment.nElems.read() >= WORK_STEALING_MINIMUM {
                            var backoff = 0;
                            if targetSegment.currentStatus == STATUS_UNLOCKED && targetSegment.acquireWithStatus(STATUS_REMOVE) {
                              // Sanity check: ensure segment did not fall under minimum since last check
                              if  targetSegment.nElems.read() < WORK_STEALING_MINIMUM {
                                targetSegment.releaseStatus();
                                break;
                              }

                              extern proc sizeof(type x): size_t;
                              // We steal at most 1MB worth of data. If the user has less than that, we steal a %, at least 1.
                              const mb = WORK_STEALING_MAX_MEMORY_MB * 1024 * 1024;
                              var toSteal = max(WORK_STEALING_MINIMUM, min(mb / sizeof(eltType), targetSegment.nElems.read() * WORK_STEALING_RATIO)) : int;

                              // Allocate storage...
                              on stolenWork do stolenWork[loc.id] = (toSteal, c_malloc(eltType, toSteal));
                              var destPtr = stolenWork[here.id][2];
                              targetSegment.transferElements(destPtr, toSteal, stolenWork.locale.id);
                              targetSegment.releaseStatus();

                              // We are done...
                              break;
                            }

                            // Backoff...
                            chpl_task_yield();
                          }
                        }
                      }
                    }

                    // It is our job now to distribute all stolen data to the same
                    // horizontal segment on our node. Acquire lock...
                    ref recvSegment = segments[segmentIdx];
                    while true {
                      if recvSegment.currentStatus == STATUS_UNLOCKED && recvSegment.acquireWithStatus(STATUS_ADD) then break;
                      chpl_task_yield();
                    }

                    // Add stolen elements to segment...
                    for (nStolen, stolenPtr) in stolenWork {
                      if nStolen == 0 then continue;
                      recvSegment.addElementsPtr(stolenPtr, nStolen);
                      c_free(stolenPtr);

                      // Let parent know that the bag is not empty.
                      isEmpty.write(false);
                    }
                    recvSegment.releaseStatus();
                  }

                  loadBalanceInProgress.write(false);

                  // At this point, if no work has been found, we will return empty...
                  if isEmpty.read() {
                    return (false, _defaultOf(eltType));
                  } else {
                    // Otherwise, we try to get data like everyone else.
                    phase = REMOVE_BEST_CASE;
                    break;
                  }
                }
              }
            }

            // Backoff to maximum...
            chpl_task_yield();
          }
        }

        otherwise do halt("Invalid phase #", phase);
      }

      // Reset variables...
      idx = startIdx;
      iterations = 0;
      backoff = 0;
    }

    halt("DEADCODE");
  }
}

class DistributedBag : Collection {
  var targetLocDom : domain(1);
  var targetLocales : [targetLocDom] locale;
  private var pid : int = -1;

  // Node-local fields
  private var bag : Bag(eltType);
  private var concurrentTasks : atomic int;
  private var frozenState : atomic int;

  proc DistributedBag(type eltType, targetLocales : [?targetLocDom] locale = Locales) {
    bag = new Bag(eltType, this);
    this.targetLocDom = targetLocDom;
    this.targetLocales = targetLocales;
    pid = _newPrivatizedClass(this);
  }

  pragma "no doc"
  proc DistributedBag(other, pid, type eltType = other.eltType) {
    bag = new Bag(eltType, this);
    this.targetLocDom = other.targetLocDom;
    this.targetLocales = other.targetLocales;
    this.pid = pid;
  }

  pragma "no doc"
  proc dsiPrivatize(pid) {
    return new DistributedBag(this, pid);
  }

  pragma "no doc"
  proc dsiGetPrivatizeData() {
    return pid;
  }

  pragma "no doc"
  inline proc getPrivatizedThis {
    if this.locale == here {
      return this;
    }

    return chpl_getPrivatizedCopy(this.type, pid);
  }

  /*
    Obtains a privatized version of this instance. The privatized version is a
    cloned instance that is allocated on this node, which is useful for eliding
    communications generated from accessing instance fields. Using another node's
    instance will significantly penalize performance by bounding overall throughput
    on the communication between the two nodes.
  */
  proc getPrivatizedInstance() {
    return getPrivatizedThis;
  }

  /*
    Insert an element to this node's bag. The ordering is not guaranteed to be
    preserved. If this instance is privatized, it is an entirely local operation.
  */
  proc add(elt : eltType) : bool {
    var localThis = getPrivatizedThis;
    if BAG_NO_FREEZE then return localThis.bag.add(elt);

    local {
      localThis.concurrentTasks.add(1);

      // Check if the queue is now 'immutable'.
      if localThis.frozenState.read() > FREEZE_UNFROZEN {
        localThis.concurrentTasks.sub(1);
        return false;
      }
    }

    var result = localThis.bag.add(elt);
    local { localThis.concurrentTasks.sub(1); }
    return result;
  }

  /*
    Remove an element from this node's bag. The order in which elements are removed
    are not guaranteed to be the same order it has been inserted. If this node's
    bag is empty, it will attempt to steal elements from bags of other nodes.
  */
  proc remove() : (bool, eltType) {
    var localThis = getPrivatizedThis;
    if BAG_NO_FREEZE then return localThis.bag.remove();

    local {
      localThis.concurrentTasks.add(1);

      // Check if the queue is now 'immutable'.
      if localThis.frozenState.read() > FREEZE_UNFROZEN {
        localThis.concurrentTasks.sub(1);
        return (false, _defaultOf(eltType));
      }
    }

    var (elem, hasElem) = localThis.bag.remove();
    local { localThis.concurrentTasks.sub(1); }
    return (elem, hasElem);
  }

  /*
    Obtain the number of elements held in all bags across all nodes. This method
    is best-effort if the bags are unfrozen, and can be non-deterministic for
    concurrent updates across nodes, and may miss elements or even count
    duplicates resulting from any concurrent insertion or removal operations.
  */
  proc size() : int {
    var sz : atomic int;
    forall loc in targetLocales do on loc {
      var instance = getPrivatizedThis;
      forall segmentIdx in 0..#here.maxTaskPar {
        sz.add(instance.bag.segments[segmentIdx].nElems.read() : int);
      }
    }

    return sz.read();
  }

  /*
    Performs a lookup to determine if the requested element exists in this bag.
    This method is best-effort if the bags are unfrozen, and can be
    non-deterministic for concurrent updates across nodes, and may miss elements
    resulting from any concurrent insertion or removal operations.
  */
  proc contains(elt : eltType) : bool {
    var foundElt : atomic bool;
    forall elem in getPrivatizedThis {
      if elem == elt {
        foundElt.write(true);
      }
    }
    return foundElt.read();
  }

  /*
    Freeze our state, becoming immutable; we wait for any ongoing concurrent
    operations to allow them to finish.
  */
  proc freeze() {
    if BAG_NO_FREEZE then return false;

    // Check if already frozen
    if frozenState.read() == FREEZE_FROZEN {
      return true;
    }

    // Mark as transient state...
    coforall loc in targetLocales do on loc {
      var localThis = getPrivatizedThis;
      localThis.frozenState.write(FREEZE_MARKED);
      localThis.concurrentTasks.waitFor(0);
    }

    // Mark as frozen...
    coforall loc in targetLocales do on loc {
      var localThis = getPrivatizedThis;
      localThis.frozenState.write(FREEZE_FROZEN);
    }

    return true;
  }

  /*
    Unfreezes our state, allowing mutating operations; we wait on any ongoing
    concurrent operations to allow them to finish.
  */
  proc unfreeze() : bool {
    if BAG_NO_FREEZE then return false;
    // Check if already unfrozen
    if frozenState.read() == FREEZE_UNFROZEN {
      return true;
    }

    // Mark as transient state...
    coforall loc in targetLocales do on loc {
      var localThis = getPrivatizedThis;
      localThis.frozenState.write(FREEZE_MARKED);
      localThis.concurrentTasks.waitFor(0);
    }

    // Mark as unfrozen...
    coforall loc in targetLocales do on loc {
      var localThis = getPrivatizedThis;
      localThis.frozenState.write(FREEZE_UNFROZEN);
    }

    return true;
  }

  proc canFreeze() : bool {
    return !BAG_NO_FREEZE;
  }

  /*
    If we are currently frozen. If we are in the middle of a state transition, we
    wait until it has completed.
  */
  proc isFrozen() : bool {
    if BAG_NO_FREEZE then return false;

    var state = frozenState.read();

    // Current transitioning state, wait it out...
    while state == FREEZE_MARKED {
      chpl_task_yield();
      state = frozenState.read();
    }

    return state == FREEZE_FROZEN;
  }

  /*
    Clear all bags across all nodes in a best-effort approach. Elements added or
    moved around from concurrent additions or removals (including work stealing)
    may be missed while clearing.
  */
  proc clear() {
    var localThis = getPrivatizedThis;
    if !BAG_NO_FREEZE {
      local {
        localThis.concurrentTasks.add(1);

        // Check if the queue is now 'immutable'.
        if localThis.frozenState.read() > FREEZE_UNFROZEN {
          localThis.concurrentTasks.sub(1);
          return;
        }
      }
    }

    coforall loc in instance.targetLocales do on loc {
      var instance = getPrivatizedThis;
      forall segmentIdx in 0 .. #here.maxTaskPar {
        ref segment = instance.bag.segments[segmentIdx];
        if segment.acquireIfNonEmpty(STATUS_REMOVE) {
          var block = segment.headBlock;
          while block != nil {
            var tmp = block;
            block = block.next;
            delete tmp;
          }

          segment.headBlock = nil;
          segment.tailBlock = nil;
          segment.nElems.write(0);
          segment.releaseStatus();
        }
      }
    }

    localThis.concurrentTasks.sub(1);
  }

  /*
    Triggers a more static approach to load balancing; note that this operation
    is parallel-safe but is performed in parallel across all `targetLocales` and
    should only be called at a point of synchronization.
  */
  proc balance() {
    var localThis = getPrivatizedThis;
    if !BAG_NO_FREEZE {
      local {
        localThis.concurrentTasks.add(1);

        // Check if the queue is now 'immutable'.
        if localThis.frozenState.read() > FREEZE_UNFROZEN {
          localThis.concurrentTasks.sub(1);
          return;
        }
      }
    }

    // Phase 1: Acquire all locks from from first node and segment to last
    // node and segment (our global locking order...)
    for loc in localThis.targetLocales do on loc {
      var instance = getPrivatizedThis;
      for segmentIdx in 0 .. #here.maxTaskPar {
        ref segment = instance.bag.segments[segmentIdx];
        segment.acquire(STATUS_BALANCE);
      }
    }

    // Phase 2: Obtain the average of all elements held in a horizontal segment...
    var elemsPerSegment : [0..#here.maxTaskPar] int;
    coforall segmentIdx in 0 .. #here.maxTaskPar {
      for loc in localThis.targetLocales do on loc {
        elemsPerSegment[segmentIdx] += getPrivatizedThis.bag.segments[segmentIdx].nElems.read() : int;
      }
    }
    for elem in elemsPerSegment do elem /= here.maxTaskPar;

    // Phase 3: Take elements from segments who have more...

    // Phase 4: Give elements to segments who have less...

    // Phase 5: Release all locks from first node and segment to last node and segment.
  }

  /*
    Iterate over each bag in each node. If the bag is frozen, we elide the need
    to acquire any locks and iterate directly over the bag. If the bag is not
    frozen, we must acquire not only the needed locks, but also take a snapshot
    approach, increasing memory consumption.
  */
  iter these() : eltType {
    // Only allowed while frozen due to issue where breaking from a serial iterator
    // does not cleanup resources, leaving the bag in an undefined state.
    var frozen = isFrozen();

    for loc in targetLocales {
      for segmentIdx in 0..#here.maxTaskPar {

        // If the data structure is frozen, we may elide the need to acquire any locks,
        // and we can iterate directly over the data.
        if frozen {
          var instance : DistributedBag(int);
          on loc do instance = getPrivatizedThis;
          ref segment = instance.bag.segments[segmentIdx];
          var block = segment.headBlock;

          while block != nil {
            for i in 0 .. #block.size {
              var elt : eltType;
              on loc do elt = block.elems[i];
              yield elt;
            }
            block = block.next;
          }
        }
        // If the data structure is not frozen, then we must acquire the lock and
        // produce a snapshot (to reduce the time we hold the lock)
        else {
          // The size of the snapshot is only known once we have the lock.
          // TODO/BUG: c_ptr does not seem to work here and causes a segmentation fault at runtime.
          var dom : domain(1) = {0..-1};
          var buffer : [dom] eltType;

          on loc {
            ref segment = getPrivatizedThis.bag.segments[segmentIdx];

            if segment.acquireIfNonEmpty(STATUS_LOOKUP) {
              var block = segment.headBlock;
              while block != nil {
                for i in 0 .. #block.size {
                    buffer.push_back(block.elems[i]);
                }
                block = block.next;
              }
              segment.releaseStatus();

            }
          }

          // Process this chunk if we have one...
          for elem in buffer {
            yield elem;
          }
        }
      }
    }
  }

  iter these(param tag : iterKind) where tag == iterKind.leader {
    var frozen = isFrozen();
    coforall loc in targetLocales do on loc {
      var instance = getPrivatizedThis;
      coforall segmentIdx in 0 .. #here.maxTaskPar {
        ref segment = instance.bag.segments[segmentIdx];

        // If the data structure is frozen, we elide the need to acquire any locks,
        // and we may iterate directly over the data as well.
        if frozen {
          var block = segment.headBlock;

          while block != nil {
            yield (block.size, block.elems);
            block = block.next;
          }
        } else if segment.acquireIfNonEmpty(STATUS_LOOKUP) {
          // Create a snapshot...
          var block = segment.headBlock;
          var bufferSz = segment.nElems.read() : int;
          var buffer = c_malloc(eltType, bufferSz);
          var bufferOffset = 0;

          while block != nil {
            if bufferOffset + block.size > bufferSz {
              halt("Snapshot attempt with bufferSz(", bufferSz, ") with offset bufferOffset(", bufferOffset + block.size, ")");
            }
            __primitive("chpl_comm_array_put", block.elems[0], here.id, buffer[bufferOffset], block.size);
            bufferOffset += block.size;
            block = block.next;
          }

          // Yield this chunk to be process...
          segment.releaseStatus();
          yield (bufferSz, buffer);
          c_free(buffer);
        }
      }
    }
  }

  iter these(param tag : iterKind, followThis) where tag == iterKind.follower {
    var (bufferSz, buffer) = followThis;
    for i in 0 .. #bufferSz {
      yield buffer[i];
    }
  }
}

use Collection;
use BlockDist;

/*
  A highly parallel segmented multiset. Each node gets its own bag, and in each bag
  it is segmented into 'here.maxTaskPar' segments. Segments allow for actual parallelism
  while operating in that it enables us to manage 'best-case', 'average-case',
  and 'worst-case' scenarios by making multiple passes over each segment. In the
  case where there is no oversubscription, the best-case will always be achieved
  (considering any other conditions are also met), while in the case of oversubscription
  or, for example a near empty bag, we fall into the 'average-case', etc. Examples
  of 'best-case' scenarios for a removal would be when a segment is unlocked
  and contains some elements we can drain, and the 'average-case' would be to find
  any segment, unlocked or not, that contains elements we can drain, and so on.

  This data structure also employs its own load balancing algorithm, beginning
  with a best-effort round-robin algorithm such that each task begins searching
  at another segment (with the added benefit of reducing overall contention), in
  particular is useful for locally distributing insertion operations. Next it also
  employs a work-stealing algorithm that horizontally steals a ratio of elements
  across segments. A segment is considered 'horizontal' if it shares the same
  segment index within [0, here.maxTaskPar); those other segments likely are also
  benefitted from local distribution from the round-robin algorithm discussed above.
  We also steal a ratio of elements, say 25%, because it leaves all other segments
  with 75% of their work; the more elements they have, the more we take, but the
  less they have, the less we steal; this also has the added benefit of reducing
  unnecessary work stealing between segments when the bag is nearly emptied. Furthermore,
  stealing horizontally ensures that all segments remain filled and that we still
  achieve parallelism across segments for removal operations. Lastly, we attempt
  to steal a maximum of `N / sizeof(eltType)`, where N is some size in megabytes
  (representing how much data can be sent in one network request), which keeps
  down excessive communication.

  This data structure does not come without flaws; as work stealing is dynamic
  and trigered on demand, work stealing can still be performed in excess, which
  dramatically causes a performance drop; this data structure scales in terms of
  both processors per node, and in terms of raw nodes themselves. As well, due
  to how work stealing is, we perform the best when we have large amounts of data
  to work on, but do not do well with smaller amounts due to excessive work stealing.
  As well, to achieve true parallelism, one must use a privatized instance, specific
  to their node to avoid communication overhead of accessing class fields.
*/

/*
  Below are segment statuses, which is a way to make visible to outsiders the
  current ongoing operation.
*/
private const STATUS_UNLOCKED : uint = 0;
private const STATUS_ADD : uint = 1;
private const STATUS_REMOVE : uint = 2;
private const STATUS_LOOKUP : uint = 3;

/*
  Below are statuses specific to the work stealing algorithm.
*/
private const WS_INITIALIZED : int = -1;
private const WS_FINISHED_WITH_NO_WORK : int = 0;
private const WS_FINISHED_WITH_WORK : int = 1;

/*
  Frozen states... If we are FREEZE_UNFROZEN, we are mutable. If we are FREEZE_FROZEN,
  we are immutable. If we are FREEZE_MARKED, we are in the middle of a state change.
*/
private const FREEZE_UNFROZEN = 0;
private const FREEZE_MARKED = 1;
private const FREEZE_FROZEN = 2;

/*
  The phases for operations.
*/
private const ADD_BEST_CASE = 0;
private const ADD_AVERAGE_CASE = 1;
private const REMOVE_BEST_CASE = 2;
private const REMOVE_AVERAGE_CASE = 3;
private const REMOVE_WORST_CASE = 4;

/*
  The initial amount of elements in an unroll block. Each successive unroll block
  is double the size of it's predecessor.
*/
config param INITIAL_BLOCK_SIZE = 1024;
// The amount we attempt to steal from each segment of other nodes. Helps to
// further ensure a proper load balance by not leaving segments empty by stealing
// on average 1/4 of their work, leaving them with 3/4. As we steal from *all* other
// segments, it is imperative that this stays low to prevent a ping-pong effect on
// work propagating from excessive work stealing. We steal at least one element,
// but no more than the amount of data that can be transferred within a single
// bulk transfer, measured in megabytes.
config param WORK_STEALING_RATIO = 0.25;
config param WORK_STEALING_MAX_MEMORY_MB : real = 1.0;
config param BAG_NO_FREEZE = false;
// The amount of elements a segment must have to be stolen from.
config param WORK_STEALING_MINIMUM = 1;

/*
  A segment block is an unrolled linked list node that holds a contiguous buffer
  of memory. Each segment block *should* be a power of two, as we increase the
  size of each subsequent unroll block by twice the size. This is so that stealing
  work is faster in that majority of elements are confined to one area.
*/
pragma "no doc"
class BagSegmentBlock {
  type eltType;

  // Contiguous memory containing all elements
  var elems :  c_ptr(eltType);
  var next : BagSegmentBlock(eltType);

  // The capacity of this block.
  var cap : int;
  // The number of occupied elements in this block.
  var size : int;

  inline proc isEmpty {
    return size == 0;
  }

  inline proc isFull {
    return size == cap;
  }

  inline proc push(elt : eltType) {
    if elems == nil {
      writeln("Cap: ", cap, "; elems: ", elems : string);
      halt("'elems' is nil");
    }

    if isFull {
      halt("SegmentBlock is Full");
    }

    elems[size] = elt;
    size = size + 1;
  }

  inline proc pop() : eltType {
    if elems == nil {
      halt("'elems' is nil");
    }

    if isEmpty {
      halt("SegmentBlock is Empty");
    }

    size = size - 1;
    var elt = elems[size];
    return elt;
  }

  proc BagSegmentBlock(type eltType, capacity) {
    if capacity == 0 {
      halt("Capacity is 0...");
    }

    cap = capacity;
    elems = c_malloc(eltType, capacity);
  }

  proc BagSegmentBlock(type eltType, ptr, capacity) {
    cap = capacity;
    elems = ptr;
    size = cap;
  }

  proc ~BagSegmentBlock() {
    c_free(elems);
  }
}

/*
  A segment is, in and of itself an unrolled linked list. We maintain one per core
  to ensure maximum parallelism.
*/
pragma "no doc"
record BagSegment {
  type eltType;

  // Used as a test-and-test-and-set spinlock.
  var status : atomic uint;

  var headBlock : BagSegmentBlock(eltType);
  var tailBlock : BagSegmentBlock(eltType);

  var nElems : atomic uint;

  inline proc isEmpty {
    return nElems.read() == 0;
  }

  inline proc acquireWithStatus(newStatus) {
    return status.compareExchangeStrong(STATUS_UNLOCKED, newStatus);
  }

  // Set status with a test-and-test-and-set loop...
  inline proc acquire(newStatus) {
    while true {
      if currentStatus == STATUS_UNLOCKED && acquireWithStatus(newStatus) {
        break;
      }

      chpl_task_yield();
    }
  }

  // Set status with a test-and-test-and-set loop, but only while it is not empty...
  inline proc acquireIfNonEmpty(newStatus) {
    while !isEmpty {
      if currentStatus == STATUS_UNLOCKED && acquireWithStatus(newStatus) {
        if isEmpty {
          releaseStatus();
          return false;
        } else {
          return true;
        }
      }

      chpl_task_yield();
    }

    return false;
  }

  inline proc isUnlocked {
    return status.read() == STATUS_UNLOCKED;
  }

  inline proc currentStatus {
    return status.read();
  }

  inline proc releaseStatus() {
    status.write(STATUS_UNLOCKED);
  }

  inline proc transferElements(destPtr, n, locId = here.id) {
    var destOffset = 0;
    var srcOffset = 0;
    while destOffset < n {
      if headBlock == nil || isEmpty {
        halt(here, ": Attempted transfer ", n, " elements to ", locId, " but failed... destOffset=", destOffset);
      }

      var len = headBlock.size;
      var need = n - destOffset;
      // If the amount in this block is greater than what is left to transfer, we
      // cannot begin transferring at the beginning, so we set our offset from the end.
      if len > need {
        srcOffset = len - need;
        headBlock.size = srcOffset;
        __primitive("chpl_comm_array_put", headBlock.elems[srcOffset], locId, destPtr[destOffset], need);
        destOffset = destOffset + need;
      } else {
        srcOffset = 0;
        headBlock.size = 0;
        __primitive("chpl_comm_array_put", headBlock.elems[srcOffset], locId, destPtr[destOffset], len);
        destOffset = destOffset + len;
      }

      // Fix list if we consumed last one...
      if headBlock.isEmpty {
        var tmp = headBlock;
        headBlock = headBlock.next;
        delete tmp;

        if headBlock == nil then tailBlock = nil;
      }
    }

    nElems.sub(n : uint);
  }

  inline proc addElementsPtr(ptr, n) {
    var offset = 0;
    while offset < n {
      var block = tailBlock;
      // Empty? Create a new one of initial size
      if block == nil then {
        tailBlock = new BagSegmentBlock(eltType, INITIAL_BLOCK_SIZE);
        headBlock = tailBlock;
        block = tailBlock;
      }

      // Full? Create a new one double the previous size
      if block.isFull {
        block.next = new BagSegmentBlock(eltType, block.cap * 2);
        tailBlock = block.next;
        block = block.next;
      }

      var nLeft = n - offset;
      var nSpace = block.cap - block.size;
      var nFill = min(nLeft, nSpace);
      __primitive("chpl_comm_array_put", ptr[offset], here.id, block.elems[block.size], nFill);
      block.size = block.size + nFill;
      offset = offset + nFill;
    }

    nElems.add(n : uint);
  }

  inline proc takeElements(n) {
    var iterations = n : int;
    var arr : [{0..#n : int}] eltType;
    var arrIdx = 0;

    for 1 .. n : int{
      if isEmpty {
        halt("Attempted to take ", n, " elements when insufficient");
      }

      if headBlock.isEmpty {
        halt("Iterating over ", n, " elements with headBlock empty but nElems is ", nElems.read());
      }

      arr[arrIdx] = headBlock.pop();
      arrIdx = arrIdx + 1;
      nElems.sub(1);

      // Fix list if we consumed last one...
      if headBlock.isEmpty {
        var tmp = headBlock;
        headBlock = headBlock.next;
        delete tmp;

        if headBlock == nil then tailBlock = nil;
      }
    }

    return arr;
  }

  inline proc takeElement() {
    if isEmpty {
      return (false, _defaultOf(eltType));
    }

    if headBlock.isEmpty {
      halt("Iterating over 1 element with headBlock empty but nElems is ", nElems.read());
    }

    var elem = headBlock.pop();
    nElems.sub(1);

    // Fix list if we consumed last one...
    if headBlock.isEmpty {
      var tmp = headBlock;
      headBlock = headBlock.next;
      delete tmp;

      if headBlock == nil then tailBlock = nil;
    }

    return (true, elem);
  }

  inline proc addElements(elt : eltType) {
    var block = tailBlock;

    // Empty? Create a new one of initial size
    if block == nil then {
      tailBlock = new BagSegmentBlock(eltType, INITIAL_BLOCK_SIZE);
      headBlock = tailBlock;
      block = tailBlock;
    }

    // Full? Create a new one double the previous size
    if block.isFull {
      block.next = new BagSegmentBlock(eltType, block.cap * 2);
      tailBlock = block.next;
      block = block.next;
    }

    block.push(elt);
    nElems.add(1);
  }

  inline proc addElements(elts) {
    for elt in elts do addElements(elt);
  }
}

/*
  We maintain a multiset 'bag' per node. Each bag keeps a handle to it's parent,
  which is required for work stealing.
*/
pragma "no doc"
class Bag {
  type eltType;

  // A handle to our parent 'distributed' bag, which is needed for work stealing.
  var parentHandle : DistributedBag(eltType);

  /*
    Helps evenly distribute and balance placement of elements in a best-effort
    round-robin approach. In the case where we have parallel enqueues or dequeues,
    they are less likely overlap with each other. Furthermore, it increases our
    chance to find our 'ideal' segment.
  */
  var startIdxEnq : atomic uint;
  var startIdxDeq : atomic uint;

  /*
    If a task makes 2 complete passes (1 best-case, 1 average-case) and has not
    found enough items, then it may attempt to balance the load distribution.
  */
  var loadBalanceInProgress : atomic bool;

  var maxParallelSegmentSpace = {0 .. #here.maxTaskPar};
  var segments : [maxParallelSegmentSpace] BagSegment(eltType);

  inline proc nextStartIdxEnq {
    return (startIdxEnq.fetchAdd(1) % here.maxTaskPar : uint) : int;
  }

  inline proc nextStartIdxDeq {
    return (startIdxDeq.fetchAdd(1) % here.maxTaskPar : uint) : int;
  }

  proc Bag(type eltType, parentHandle) {
    this.parentHandle = parentHandle;
  }

  proc add(elt : eltType) : bool {
    var startIdx = nextStartIdxEnq : int;
    var phase = ADD_BEST_CASE;

    while true {
      select phase {
        /*
          Pass 1: Best Case

          Find a segment that is unlocked and attempt to acquire it. As we are adding
          elements, we don't care how many elements there are, just that we find
          some place to add ours.
        */
        when ADD_BEST_CASE {
          for offset in 0 .. #here.maxTaskPar {
            ref segment = segments[(startIdx + offset) % here.maxTaskPar];

            // Attempt to acquire...
            if segment.acquireWithStatus(STATUS_ADD) {
              segment.addElements(elt);
              segment.releaseStatus();
              return true;
            }
          }

          phase = ADD_AVERAGE_CASE;
        }
        /*
          Pass 2: Average Case

          Find any segment (locked or unlocked) and make an attempt to acquire it.
        */
        when ADD_AVERAGE_CASE {
          ref segment = segments[startIdx];

          while true {
            select segment.currentStatus {
              // Quick acquire...
              when STATUS_UNLOCKED do {
                if segment.acquireWithStatus(STATUS_ADD) {
                  segment.addElements(elt);
                  segment.releaseStatus();
                  return true;
                }
              }
            }
            chpl_task_yield();
          }
        }
      }
    }

    halt("DEADCODE");
  }

  proc remove() : (bool, eltType) {
    var startIdx = nextStartIdxDeq;
    var idx = startIdx;
    var iterations = 0;
    var phase = REMOVE_BEST_CASE;
    var backoff = 0;

    while true {
      select phase {
        /*
          Pass 1: Best Case

          Find the first bucket that is both unlocked and contains elements. This is
          extremely helpful in the case where we have a good distribution of elements
          in each segment.
        */
        when REMOVE_BEST_CASE {
          while iterations < here.maxTaskPar {
            ref segment = segments[idx];

            // Attempt to acquire...
            if !segment.isEmpty && segment.acquireWithStatus(STATUS_REMOVE) {
              var (hasElem, elem) : (bool, eltType) = segment.takeElement();
              segment.releaseStatus();

              if hasElem {
                return (hasElem, elem);
              }
            }

            iterations = iterations + 1;
            idx = (idx + 1) % here.maxTaskPar;
          }

          phase = REMOVE_AVERAGE_CASE;
        }

        /*
          Pass 2: Average Case

          Find the first bucket containing elements. We don't care if it is locked
          or unlocked this time, just that it contains elements; this handles majority
          of cases where we have elements anywhere in any segment.
        */
        when REMOVE_AVERAGE_CASE {
          while iterations < here.maxTaskPar {
            ref segment = segments[idx];

            // Attempt to acquire...
            while !segment.isEmpty {
              if segment.isUnlocked && segment.acquireWithStatus(STATUS_REMOVE) {
                var (hasElem, elem) : (bool, eltType) = segment.takeElement();
                segment.releaseStatus();

                if hasElem {
                  return (hasElem, elem);
                }
              }

              // Backoff
              chpl_task_yield();
            }

            iterations = iterations + 1;
            idx = (idx + 1) % here.maxTaskPar;
          }

          phase = REMOVE_WORST_CASE;
        }

        /*
          Pass 3: Worst Case

          After two full iterations, we're sure the queue is full at this point, so we
          can attempt to steal work from other nodes. In this pass, we find *any* segment
          and if it is empty, we attempt to become the work-stealer; if someone else is the
          current work stealer we assist them instead and lift an element for ourselves.

          Furthermore, in this phase we loop indefinitely until we are 100% certain it is
          empty or we get an item, so introduce some backoff here.
        */
        when REMOVE_WORST_CASE {
          while true {
            ref segment = segments[idx];

            select segment.currentStatus {
              // Quick acquire
              when STATUS_UNLOCKED {
                if segment.acquireWithStatus(STATUS_REMOVE) {
                  // We're lucky; another element has been added to the current segment,
                  // take it and leave like normal...
                  if !segment.isEmpty {
                    var (hasElem, elem) : (bool, eltType) = segment.takeElement();
                    segment.releaseStatus();
                    return (hasElem, elem);
                  }

                  if parentHandle.targetLocales.size == 1 {
                    segment.releaseStatus();
                    return (false, _defaultOf(eltType));
                  }

                  // Attempt to become the sole work stealer for this node. If we
                  // do not, we spin until they finish. We need to release the lock
                  // on our segment so our segment may be load balanced as well.
                  if loadBalanceInProgress.testAndSet() {
                    segment.releaseStatus();

                    loadBalanceInProgress.waitFor(false);
                    // Reset our phase and scan for more elements...
                    phase = REMOVE_BEST_CASE;
                    break;
                  }

                  // We are the sole work stealer, and so it is our responsibility
                  // to balance the load for our node. We fork-join new worker
                  // tasks that will check horizontally across each node (as in
                  // across each segment with the same index), and vertically across
                  // each segment (each segment in a node). Horizontally, we steal
                  // at most a % of work from other nodes to give to ourselves.
                  // As load balancer, we also are the only one who knows whether
                  // or not all bags are empty.
                  var isEmpty : atomic bool;
                  isEmpty.write(true);
                  segment.releaseStatus();
                  coforall segmentIdx in 0..#here.maxTaskPar {
                    var stolenWork : [{0..#numLocales}] (int, c_ptr(eltType));
                    coforall loc in parentHandle.targetLocales {
                      if loc != here then on loc {
                        // As we jumped to the target node, 'localBag' returns
                        // the target's bag that we are attempting to steal from.
                        var targetBag = parentHandle.bag;

                        // Only proceed if the target is not load balancing themselves...
                        if !targetBag.loadBalanceInProgress.read() {
                          ref targetSegment = targetBag.segments[segmentIdx];

                          // As we only care that the segment contains data,
                          // we test-and-test-and-set until we gain ownership.
                          while targetSegment.nElems.read() >= WORK_STEALING_MINIMUM {
                            var backoff = 0;
                            if targetSegment.currentStatus == STATUS_UNLOCKED && targetSegment.acquireWithStatus(STATUS_REMOVE) {
                              // Sanity check: ensure segment did not fall under minimum since last check
                              if  targetSegment.nElems.read() < WORK_STEALING_MINIMUM {
                                targetSegment.releaseStatus();
                                break;
                              }

                              extern proc sizeof(type x): size_t;
                              // We steal at most 1MB worth of data. If the user has less than that, we steal a %, at least 1.
                              const mb = WORK_STEALING_MAX_MEMORY_MB * 1024 * 1024;
                              var toSteal = max(WORK_STEALING_MINIMUM, min(mb / sizeof(eltType), targetSegment.nElems.read() * WORK_STEALING_RATIO)) : int;

                              // Allocate storage...
                              on stolenWork do stolenWork[loc.id] = (toSteal, c_malloc(eltType, toSteal));
                              var destPtr = stolenWork[here.id][2];
                              targetSegment.transferElements(destPtr, toSteal, stolenWork.locale.id);
                              targetSegment.releaseStatus();

                              // We are done...
                              break;
                            }

                            // Backoff...
                            chpl_task_yield();
                          }
                        }
                      }
                    }

                    // It is our job now to distribute all stolen data to the same
                    // horizontal segment on our node. Acquire lock...
                    ref recvSegment = segments[segmentIdx];
                    while true {
                      if recvSegment.currentStatus == STATUS_UNLOCKED && recvSegment.acquireWithStatus(STATUS_ADD) then break;
                      chpl_task_yield();
                    }

                    // Add stolen elements to segment...
                    for (nStolen, stolenPtr) in stolenWork {
                      if nStolen == 0 then continue;
                      recvSegment.addElementsPtr(stolenPtr, nStolen);
                      c_free(stolenPtr);

                      // Let parent know that the bag is not empty.
                      isEmpty.write(false);
                    }
                    recvSegment.releaseStatus();
                  }

                  loadBalanceInProgress.write(false);

                  // At this point, if no work has been found, we will return empty...
                  if isEmpty.read() {
                    return (false, _defaultOf(eltType));
                  } else {
                    // Otherwise, we try to get data like everyone else.
                    phase = REMOVE_BEST_CASE;
                    break;
                  }
                }
              }
            }

            // Backoff to maximum...
            chpl_task_yield();
          }
        }

        otherwise do halt("Invalid phase #", phase);
      }

      // Reset variables...
      idx = startIdx;
      iterations = 0;
      backoff = 0;
    }

    halt("DEADCODE");
  }
}

class DistributedBag : Collection {
  var targetLocDom : domain(1);
  var targetLocales : [targetLocDom] locale;
  private var pid : int = -1;

  // Node-local fields
  private var bag : Bag(eltType);
  private var concurrentTasks : atomic int;
  private var frozenState : atomic int;

  proc DistributedBag(type eltType, targetLocales : [?targetLocDom] locale = Locales) {
    bag = new Bag(eltType, this);
    this.targetLocDom = targetLocDom;
    this.targetLocales = targetLocales;
    pid = _newPrivatizedClass(this);
  }

  pragma "no doc"
  proc DistributedBag(other, pid, type eltType = other.eltType) {
    bag = new Bag(eltType, this);
    this.targetLocDom = other.targetLocDom;
    this.targetLocales = other.targetLocales;
    this.pid = pid;
  }

  pragma "no doc"
  proc dsiPrivatize(pid) {
    return new DistributedBag(this, pid);
  }

  pragma "no doc"
  proc dsiGetPrivatizeData() {
    return pid;
  }

  pragma "no doc"
  inline proc getPrivatizedThis {
    if this.locale == here {
      return this;
    }

    return chpl_getPrivatizedCopy(this.type, pid);
  }

  /*
    Obtains a privatized version of this instance. The privatized version is a
    cloned instance that is allocated on this node, which is useful for eliding
    communications generated from accessing instance fields. Using another node's
    instance will significantly penalize performance by bounding overall throughput
    on the communication between the two nodes.
  */
  proc getPrivatizedInstance() {
    return getPrivatizedThis;
  }

  /*
    Insert an element to this node's bag. The ordering is not guaranteed to be
    preserved. If this instance is privatized, it is an entirely local operation.
  */
  proc add(elt : eltType) : bool {
    var localThis = getPrivatizedThis;
    if BAG_NO_FREEZE then return localThis.bag.add(elt);

    local {
      localThis.concurrentTasks.add(1);

      // Check if the queue is now 'immutable'.
      if localThis.frozenState.read() > FREEZE_UNFROZEN {
        localThis.concurrentTasks.sub(1);
        return false;
      }
    }

    var result = localThis.bag.add(elt);
    local { localThis.concurrentTasks.sub(1); }
    return result;
  }

  /*
    Remove an element from this node's bag. The order in which elements are removed
    are not guaranteed to be the same order it has been inserted. If this node's
    bag is empty, it will attempt to steal elements from bags of other nodes.
  */
  proc remove() : (bool, eltType) {
    var localThis = getPrivatizedThis;
    if BAG_NO_FREEZE then return localThis.bag.remove();

    local {
      localThis.concurrentTasks.add(1);

      // Check if the queue is now 'immutable'.
      if localThis.frozenState.read() > FREEZE_UNFROZEN {
        localThis.concurrentTasks.sub(1);
        return (false, _defaultOf(eltType));
      }
    }

    var (elem, hasElem) = localThis.bag.remove();
    local { localThis.concurrentTasks.sub(1); }
    return (elem, hasElem);
  }

  /*
    Obtain the number of elements held in all bags across all nodes. This method
    is best-effort if the bags are unfrozen, and can be non-deterministic for
    concurrent updates across nodes, and may miss elements or even count
    duplicates resulting from any concurrent insertion or removal operations.
  */
  proc size() : int {
    var sz : atomic int;
    forall loc in targetLocales do on loc {
      var instance = getPrivatizedThis;
      forall segmentIdx in 0..#here.maxTaskPar {
        sz.add(instance.bag.segments[segmentIdx].nElems.read() : int);
      }
    }

    return sz.read();
  }

  /*
    Performs a lookup to determine if the requested element exists in this bag.
    This method is best-effort if the bags are unfrozen, and can be
    non-deterministic for concurrent updates across nodes, and may miss elements
    resulting from any concurrent insertion or removal operations.
  */
  proc contains(elt : eltType) : bool {
    var foundElt : atomic bool;
    forall elem in getPrivatizedThis {
      if elem == elt {
        foundElt.write(true);
      }
    }
    return foundElt.read();
  }

  /*
    Freeze our state, becoming immutable; we wait for any ongoing concurrent
    operations to allow them to finish.
  */
  proc freeze() {
    if BAG_NO_FREEZE then return false;

    // Check if already frozen
    if frozenState.read() == FREEZE_FROZEN {
      return true;
    }

    // Mark as transient state...
    coforall loc in targetLocales do on loc {
      var localThis = getPrivatizedThis;
      localThis.frozenState.write(FREEZE_MARKED);
      localThis.concurrentTasks.waitFor(0);
    }

    // Mark as frozen...
    coforall loc in targetLocales do on loc {
      var localThis = getPrivatizedThis;
      localThis.frozenState.write(FREEZE_FROZEN);
    }

    return true;
  }

  /*
    Unfreezes our state, allowing mutating operations; we wait on any ongoing
    concurrent operations to allow them to finish.
  */
  proc unfreeze() : bool {
    if BAG_NO_FREEZE then return false;
    // Check if already unfrozen
    if frozenState.read() == FREEZE_UNFROZEN {
      return true;
    }

    // Mark as transient state...
    coforall loc in targetLocales do on loc {
      var localThis = getPrivatizedThis;
      localThis.frozenState.write(FREEZE_MARKED);
      localThis.concurrentTasks.waitFor(0);
    }

    // Mark as unfrozen...
    coforall loc in targetLocales do on loc {
      var localThis = getPrivatizedThis;
      localThis.frozenState.write(FREEZE_UNFROZEN);
    }

    return true;
  }

  proc canFreeze() : bool {
    return !BAG_NO_FREEZE;
  }

  /*
    If we are currently frozen. If we are in the middle of a state transition, we
    wait until it has completed.
  */
  proc isFrozen() : bool {
    if BAG_NO_FREEZE then return false;

    var state = frozenState.read();

    // Current transitioning state, wait it out...
    while state == FREEZE_MARKED {
      chpl_task_yield();
      state = frozenState.read();
    }

    return state == FREEZE_FROZEN;
  }

  /*
    Iterate over each bag in each node. If the bag is frozen, we elide the need
    to acquire any locks and iterate directly over the bag. If the bag is not
    frozen, we must acquire not only the needed locks, but also take a snapshot
    approach, increasing memory consumption.
  */
  iter these() : eltType {
    // Only allowed while frozen due to issue where breaking from a serial iterator
    // does not cleanup resources, leaving the bag in an undefined state.
    var frozen = isFrozen();

    for loc in targetLocales {
      for segmentIdx in 0..#here.maxTaskPar {

        // If the data structure is frozen, we may elide the need to acquire any locks,
        // and we can iterate directly over the data.
        if frozen {
          var instance : DistributedBag(int);
          on loc do instance = getPrivatizedThis;
          ref segment = instance.bag.segments[segmentIdx];
          var block = segment.headBlock;

          while block != nil {
            for i in 0 .. #block.size {
              var elt : eltType;
              on loc do elt = block.elems[i];
              yield elt;
            }
            block = block.next;
          }
        }
        // If the data structure is not frozen, then we must acquire the lock and
        // produce a snapshot (to reduce the time we hold the lock)
        else {
          // The size of the snapshot is only known once we have the lock.
          // TODO/BUG: c_ptr does not seem to work here and causes a segmentation fault at runtime.
          var dom : domain(1) = {0..-1};
          var buffer : [dom] eltType;

          on loc {
            ref segment = getPrivatizedThis.bag.segments[segmentIdx];

            if segment.acquireIfNonEmpty(STATUS_LOOKUP) {
              var block = segment.headBlock;
              while block != nil {
                for i in 0 .. #block.size {
                    buffer.push_back(block.elems[i]);
                }
                block = block.next;
              }
              segment.releaseStatus();

            }
          }

          // Process this chunk if we have one...
          for elem in buffer {
            yield elem;
          }
        }
      }
    }
  }

  iter these(param tag : iterKind) where tag == iterKind.leader {
    var frozen = isFrozen();
    coforall loc in targetLocales do on loc {
      var instance = getPrivatizedThis;
      coforall segmentIdx in 0 .. #here.maxTaskPar {
        ref segment = instance.bag.segments[segmentIdx];

        // If the data structure is frozen, we elide the need to acquire any locks,
        // and we may iterate directly over the data as well.
        if frozen {
          var block = segment.headBlock;

          while block != nil {
            yield (block.size, block.elems);
            block = block.next;
          }
        } else if segment.acquireIfNonEmpty(STATUS_LOOKUP) {
          // Create a snapshot...
          var block = segment.headBlock;
          var bufferSz = segment.nElems.read() : int;
          var buffer = c_malloc(eltType, bufferSz);
          var bufferOffset = 0;

          while block != nil {
            if bufferOffset + block.size > bufferSz {
              halt("Snapshot attempt with bufferSz(", bufferSz, ") with offset bufferOffset(", bufferOffset + block.size, ")");
            }
            __primitive("chpl_comm_array_put", block.elems[0], here.id, buffer[bufferOffset], block.size);
            bufferOffset += block.size;
            block = block.next;
          }

          // Yield this chunk to be process...
          segment.releaseStatus();
          yield (bufferSz, buffer);
          c_free(buffer);
        }
      }
    }
  }

  iter these(param tag : iterKind, followThis) where tag == iterKind.follower {
    var (bufferSz, buffer) = followThis;
    for i in 0 .. #bufferSz {
      yield buffer[i];
    }
  }
}


IPE/ChapelBase.chpl: Drop some function types that are not currently supported

/*
 * Copyright 2004-2015 Cray Inc.
 * Other additional copyright holders may be indicated within.
 *
 * The entirety of this work is 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.
 */

// ChapelBase.chpl
//

module ChapelBase
{
  //
  // Assignment operators
  //
  inline proc =  (ref a : bool, b : bool) : void { __primitive("=", a, b); }
  inline proc =  (ref a : int,  b : int ) : void { __primitive("=", a, b); }
  inline proc =  (ref a : real, b : real) : void { __primitive("=", a, b); }

  //
  // Basic arithmetic operators
  //
  inline proc +  (    a : int,  b : int ) : int  return __primitive("+",  a, b);
  inline proc -  (    a : int,  b : int ) : int  return __primitive("-",  a, b);
  inline proc *  (    a : int,  b : int ) : int  return __primitive("*",  a, b);
  inline proc /  (    a : int,  b : int ) : int  return __primitive("/",  a, b);

  inline proc +  (    a : real, b : real) : real return __primitive("+",  a, b);
  inline proc -  (    a : real, b : real) : real return __primitive("+",  a, b);
  inline proc *  (    a : real, b : real) : real return __primitive("+",  a, b);
  inline proc /  (    a : real, b : real) : real return __primitive("+",  a, b);

  //
  // Equality comparison on primitive types
  //
  inline proc == (    a : bool, b : bool) : bool return __primitive("==", a, b);
  inline proc == (    a : int,  b : int ) : bool return __primitive("==", a, b);
  inline proc == (    a : real, b : real) : bool return __primitive("==", a, b);

  inline proc != (    a : bool, b : bool) : bool return __primitive("!=", a, b);
  inline proc != (    a : int,  b : int ) : bool return __primitive("!=", a, b);
  inline proc != (    a : real, b : real) : bool return __primitive("!=", a, b);

  //
  // Ordered comparison on primitive types
  //
  inline proc <  (    a : int,  b : int ) : bool return __primitive("<",  a, b);
  inline proc <  (    a : real, b : real) : bool return __primitive("<",  a, b);

  inline proc >  (    a : int,  b : int ) : bool return __primitive(">",  a, b);
  inline proc >  (    a : real, b : real) : bool return __primitive(">",  a, b);

  inline proc <= (    a : int,  b : int ) : bool return __primitive("<=", a, b);
  inline proc <= (    a : real, b : real) : bool return __primitive("<=", a, b);

  inline proc >= (    a : int,  b : int ) : bool return __primitive(">=", a, b);
  inline proc >= (    a : real, b : real) : bool return __primitive(">=", a, b);

  //
  // Support for implicit boolean conversion for conditional expressions
  //
  inline proc _cond_test(x : bool) : bool return x;
  inline proc _cond_test(x : int ) : bool return x != 0;
}

/*
 * Copyright 2004-2015 Cray Inc.
 * Other additional copyright holders may be indicated within.
 *
 * The entirety of this work is 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.
 */

// ChapelBase.chpl
//

module ChapelBase
{
  param bitsPerByte = 8;

  proc numBits(type t)  param where t == bool(64) return 64;
  proc numBits(type t)  param where t == int(64)  return 64;
  proc numBits(type t)  param where t == real(64) return 64;

  proc numBytes(type t) param return numBits(t) / 8;

  //
  // Cond test
  //
  inline proc _cond_test (x : bool)     return x;
  inline proc _cond_test (x : int )     return x != 0;

  //
  // Assignment operators
  //
  inline proc =  (ref a : bool, b : bool) : void { __primitive("=", a, b); }
  inline proc =  (ref a : int,  b : int ) : void { __primitive("=", a, b); }
  inline proc =  (ref a : real, b : real) : void { __primitive("=", a, b); }

  //
  // Basic arithmetic operators
  //
  inline proc +  (a: int,  b: int)        : int  return __primitive("+", a, b);
  inline proc -  (a: int,  b: int)        : int  return __primitive("-", a, b);
  inline proc *  (a: int,  b: int)        : int  return __primitive("*", a, b);
  inline proc /  (a: int,  b: int)        : int  return __primitive("/", a, b);

  inline proc +  (a: real, b: real)       : real return __primitive("+", a, b);
  inline proc -  (a: real, b: real)       : real return __primitive("+", a, b);
  inline proc *  (a: real, b: real)       : real return __primitive("+", a, b);
  inline proc /  (a: real, b: real)       : real return __primitive("+", a, b);

  //
  // equality comparison on primitive types
  //
  inline proc == (a: bool, b: bool)       : bool return __primitive("==", a, b);
  inline proc == (a: int,  b: int)        : bool return __primitive("==", a, b);
  inline proc == (a: real, b: real)       : bool return __primitive("==", a, b);

  inline proc != (a: bool, b: bool)       : bool return __primitive("!=", a, b);
  inline proc != (a: int,  b: int)        : bool return __primitive("!=", a, b);
  inline proc != (a: real, b: real)       : bool return __primitive("!=", a, b);

  //
  // ordered comparison on primitive types
  //
  inline proc <  (a: int,  b: int)        : bool return __primitive("<",  a, b);
  inline proc <  (a: real, b: real)       : bool return __primitive("<",  a, b);

  inline proc >  (a: int,  b: int)        : bool return __primitive(">",  a, b);
  inline proc >  (a: real, b: real)       : bool return __primitive(">",  a, b);

  inline proc <= (a: int,  b: int)        : bool return __primitive("<=", a, b);
  inline proc <= (a: real, b: real)       : bool return __primitive("<=", a, b);

  inline proc >= (a: int,  b: int)        : bool return __primitive(">=", a, b);
  inline proc >= (a: real, b: real)       : bool return __primitive(">=", a, b);
}


Reduce number of spawned processes

Reduce number of spawned processes

use Regexp;
use Histogram;
use Time;

// Use the test/twopt code
config const isTest=false;
config const doBrute=false;

config const fn1 = "test.dat";
config const fn2 = "test.dat";
config const pairfn = "test-DD.dat";

param NDIM  = 3;
const Ddim = {0.. #NDIM};

config const spaces=compile("\\s+");
config const bufsize=10000;
config const minpart=500;

config const smax=200.0;
config const smax2=smax**2;
config const nmubins=5;
config const nsbins=5;

// Global variables
var gtime1 : Timer;

proc main() {
  if isTest then testPairs(); else doPairs();
}

record WeightedParticle3D {
  var x : [Ddim] real;
  var w : real;
}

proc countLines(fn : string) : int {
  var ff = open(fn, iomode.r);
  var ipart = 0;
  for iff in ff.lines() do ipart +=1;
  ff.close();
  return ipart;
}

proc readFile(fn : string, pp : []WeightedParticle3D)  {
  const maxcols=25;

  var ff = open(fn, iomode.r);
  var cols : [0.. #maxcols] real;
  var icol=0;
  var ipart = 0;
  for iff in ff.lines() {
   icol = 0;
   for col1 in iff.split(spaces) {
     if (col1.length==0) then continue;
     cols[icol] = col1 : real;
     icol += 1;
   }
   if (icol < 4) then assert(false,"malformed line...");
   pp[ipart].x = cols[Ddim];
   pp[ipart].w = cols[3];
   ipart += 1;
  }
}

class SOA_WeightedParticle3D {
  var npart : int;
  var Dpart : domain(1);
  var x,y,z,w,r2 : [Dpart] real;

  proc SOA_WeightedParticle3D(pp : []WeightedParticle3D) {
    npart = pp.size;
    Dpart = 0.. #npart;

    forall ii in Dpart {
      x[ii] = pp[ii].x[0];
      y[ii] = pp[ii].x[1];
      z[ii] = pp[ii].x[2];
      w[ii] = pp[ii].w;
      r2[ii] = + reduce (pp[ii].x**2);
    }
  }
}

proc smuAccumulate(hh : UniformBins, p1,p2 : SOA_WeightedParticle3D, d1,d2 : domain(1), scale : real) {
  forall ii in d1 { // Loop over first set of particles

    var x1,y1,z1,w1,r2 : real;
    var sl, s2, l1, s1, l2, mu, wprod : real;
    x1 = p1.x[ii]; y1 = p1.y[ii]; z1 = p1.z[ii]; w1 = p1.w[ii]; r2 = p1.r2[ii];

    for jj in d2 { // Second set of particles
      mu=2*(p2.x[jj]*x1 + p2.y[jj]*y1 + p2.z[jj]*z1);
      sl = r2 - p2.r2[jj];
      l1 = r2 + p2.r2[jj];
      s2 = l1 - mu;
      l2 = l1 + mu;
      if ((s2 >= smax2) || (s2 < 1.0e-20)) then continue;
      wprod = scale * w1 * p2.w[jj];
      s1 = sqrt(s2);
      mu = sl/(s1*sqrt(l2));
      if (mu < 0) then mu = -mu;

      hh.add((s1,mu),wprod);
    }
  }
}

proc writeHist(fn : string, hh : UniformBins) {
  var ff = openwriter(fn);
  // Dump out values
  for xx in hh.bins(1) do ff.writef("%12.4dr",xx);
  ff.writeln();
  for xx in hh.bins(2) do ff.writef("%12.4dr",xx);
  ff.writeln("\n##");
  for ii in hh.Dhist.dim(1) {
    for jj in hh.Dhist.dim(2) {
      ff.writef("%20.14er ",hh[(ii,jj)]);
    }
    ff.writeln();
  }
  ff.close();
}

proc splitOn(pp : []WeightedParticle3D, scr : []WeightedParticle3D, splitDim : int, rr : []real) : int {
  var npart, lnpart : int;
  var lo = pp.domain.low;
  var hi = pp.domain.high;
  npart = pp.domain.size;
  lnpart = 0;
  for ipp in pp {
    if (ipp.x[splitDim] < rr[splitDim]) then lnpart+=1;
  }
//  var part : [pp.domain] WeightedParticle3D;
  var li, ri : int;
  li = lo; ri = (lo+lnpart);
  for ipp in pp {
    if (ipp.x[splitDim] < rr[splitDim]) {
      scr[li] = ipp;
      li+=1;
    } else {
      scr[ri] = ipp;
      ri+=1;
    }
  }
  assert(li==(lo+lnpart), "splitOn : li != lo+npart");
  assert(ri==(hi+1), "splitOn : ri != hi+1");
  pp = scr;

  return lnpart;
}

class KDNode {
  var lo, hi,npart,id : int;
  var dom : domain(1);
  var xcen : [Ddim]real;
  var rcell : real;
  var left, right : KDNode;

  proc isLeaf() : bool {
    return (left==nil) && (right==nil);
  }

  proc ~KDNode() {
    if left then delete left;
    if right then delete right;
  }

}

proc BuildTree(pp : []WeightedParticle3D, scr : []WeightedParticle3D, id : int) : KDNode{
  gtime1.start();
  var me = new KDNode();
  var dom = pp.domain;
  me.lo = dom.low;
  me.hi = dom.high;
  me.dom = {me.lo..me.hi};
  me.id = id;
  me.npart = (me.hi-me.lo)+1;

  //  work out xcen and vantage point radius
  var pmin = pp[me.lo].x;
  var pmax = pp[me.lo].x;
  for ipp in pp {
    for idim in Ddim {
      if (ipp.x[idim] < pmin[idim]) then pmin[idim] = ipp.x[idim];
      if (ipp.x[idim] > pmax[idim]) then pmax[idim] = ipp.x[idim];
    }
  }
  me.xcen = (pmax+pmin)/2.0;
  var dx : [Ddim]real;
  me.rcell = 0.0;
  var r1 : real;
  for ipp in pp {
    dx = (ipp.x - me.xcen);
    r1 = sqrt(+ reduce dx**2);
    if (r1 > me.rcell) then me.rcell = r1;
  }

  // Continue to split
  gtime1.stop();
  if (me.npart <= minpart) then return me; // Don't split further.
  gtime1.start();

  // Find dimension to split on
  dx = pmax - pmin;
  var splitDim = 0;
  for idim in Ddim {
    if (dx[idim] > dx[splitDim]) then splitDim=idim;
  }

  // Split
  var lnpart = splitOn(pp, scr,splitDim, me.xcen);
  var ldom = {me.lo..(me.lo+lnpart-1)};
  var rdom = {(me.lo+lnpart)..me.hi};
  gtime1.stop();
  me.left = BuildTree(pp[ldom],scr[ldom],2*id+1);
  me.right = BuildTree(pp[rdom],scr[rdom],2*id+2);
  return me;
}

proc TreeAccumulate(hh : UniformBins, p1, p2 : SOA_WeightedParticle3D, node1, node2 : KDNode) {
  // Compute the distance between node1 and node2
  var rr = sqrt (+ reduce(node1.xcen - node2.xcen)**2);
  var rmin = rr - (node1.rcell+node2.rcell);

  // If distance is greater than all cases
  if (rmin > smax) then return;

  // If both nodes are leaves
  if (node1.isLeaf() & node2.isLeaf()) {
    begin smuAccumulate(hh, p1, p2, node1.dom, node2.dom,1);
    return;
  }

  // If one node is a leaf
  if (node1.isLeaf()) {
    TreeAccumulate(hh, p1, p2, node1, node2.left);
    TreeAccumulate(hh, p1, p2, node1, node2.right);
    return;
  }
  if (node2.isLeaf()) {
    TreeAccumulate(hh, p1, p2, node1.left, node2);
    TreeAccumulate(hh, p1, p2, node1.right, node2);
    return;
  }

  // Split the larger case;
  if (node1.npart > node2.npart) {
    TreeAccumulate(hh, p1, p2, node1.left, node2);
    TreeAccumulate(hh, p1, p2, node1.right, node2);
    return;
  } else {
    TreeAccumulate(hh, p1, p2, node1, node2.left);
    TreeAccumulate(hh, p1, p2, node1, node2.right);
    return;
  }

}

// A simple testHarness
proc testPairs() {

  // Read in the file
  var nlines = countLines(fn1);
  var pp1 : [0.. #nlines] WeightedParticle3D;
  readFile("test.dat",pp1);

  // Build the tree
  var scr : [0.. #nlines] WeightedParticle3D;
  var root1 = BuildTree(pp1,scr,0);

  // Set up the histogram
  var hh = new UniformBins(2,(nsbins,nmubins), ((0.0,smax),(0.0,1.0+1.e-10)));

  // AOS -> SOA
  var soa1 = new SOA_WeightedParticle3D(pp1);

  // Do the paircounts with a tree
  sync TreeAccumulate(hh, soa1, soa1, root1, root1);
  for xx in hh.bins(1) do writef("%12.4dr",xx);
  writeln();
  for xx in hh.bins(2) do writef("%12.4dr",xx);
  writeln("\n##");
  for ii in hh.Dhist.dim(1) {
    for jj in hh.Dhist.dim(2) {
      writef("%20.5er ",hh[(ii,jj)]);
    }
    writeln();
  }
}

// The main code
proc doPairs() {
  var tt : Timer;

  // Read in the file
  tt.clear(); tt.start();
  var nlines1 = countLines(fn1);
  var pp1 : [0.. #nlines1] WeightedParticle3D;
  readFile(fn1,pp1);
  var nlines2 = countLines(fn2);
  var pp2 : [0.. #nlines2] WeightedParticle3D;
  readFile(fn1,pp2);
  tt.stop();
  writef("Read in %i lines from file %s \n", pp1.size, fn1);
  writef("Read in %i lines from file %s \n", pp2.size, fn2);
  writef("Time to read : %r \n", tt.elapsed());

  // Build the tree
  tt.clear(); tt.start(); gtime1.clear();
  var scr1 : [0.. #nlines1] WeightedParticle3D;
  var root1 = BuildTree(pp1,scr1,0);
  var scr2 : [0.. #nlines2] WeightedParticle3D;
  var root2 = BuildTree(pp2,scr2,0);
  tt.stop();
  writef("Time to build trees : %r \n", tt.elapsed());
  writef("Time in splitOn : %r \n", gtime1.elapsed());

  // Set up the histogram
  var hh = new UniformBins(2,(nsbins,nmubins), ((0.0,smax),(0.0,1.0+1.e-10)));

  // AOS -> SOA
  tt.clear(); tt.start();
  var soa1 = new SOA_WeightedParticle3D(pp1);
  var soa2 = new SOA_WeightedParticle3D(pp2);
  tt.stop();
  writef("Time to SOA : %r \n", tt.elapsed());

  // Brute force paircounts
  if doBrute {
    tt.clear(); tt.start();
    smuAccumulate(hh, soa1,soa2, soa1.Dpart, soa2.Dpart, 1.0);
    tt.stop();
    writef("Time to brute force paircount : %r \n", tt.elapsed());
    writeHist("%s.brute".format(pairfn),hh);
  }

   // Do the paircounts with a tree
  hh.reset();
  tt.clear(); tt.start();
  sync TreeAccumulate(hh, soa1, soa2, root1, root2);
  tt.stop();
  writef("Time to tree paircount : %r \n", tt.elapsed());
  writeHist("%s.tree".format(pairfn),hh);
}

module Histogram {

  class UniformBins {
    param dim : int;
    var nbins : dim*int;
    var Dhist : domain(dim);
    var lo, hi,dx, invdx : [1..dim] real;
    var arr : [Dhist] atomic real;

    proc UniformBins(param dim : int, nbins : dim*int, limits : dim*(real,real)) {
      var dd : dim*range;
      this.nbins = nbins;
      for param ii in 1..dim {
        lo[ii] = limits(ii)(1);
        hi[ii] = limits(ii)(2);
        dx[ii] = (hi[ii]-lo[ii])/nbins(ii);
        dd(ii) = 0.. #nbins(ii);
      }
      invdx = 1.0/dx;
      Dhist = {(...dd)};
      [x in arr] x.write(0.0);
    }

    proc reset() {
      [x in arr] x.write(0.0);
    }

    proc bins(idim : int) {
      var bounds : [0.. #nbins(idim)+1] real;
      for i in {0.. #(nbins(idim)+1)} do bounds[i] = lo[idim] + dx[idim]*i;
      return bounds;
    }

    proc add(x : dim*real, w : real) {
      for param ii in 1..dim do
        if ((x(ii) < lo[ii]) | (x(ii) >= hi[ii])) then return;

      var pos : dim*int;
      for param ii in 1..dim do pos(ii) = ((x(ii)-lo(ii))*invdx[ii]) : int;
      arr[pos].add(w);
    }

    proc this(ndx) : real {
      return arr[ndx].read();
    }

  } // UniformBins

} // End module Histogram

use Regexp;
use Histogram;
use Time;

// Use the test/twopt code
config const isTest=false;
config const doBrute=false;

config const fn1 = "test.dat";
config const fn2 = "test.dat";
config const pairfn = "test-DD.dat";

param NDIM  = 3;
const Ddim = {0.. #NDIM};

config const spaces=compile("\\s+");
config const bufsize=10000;
config const minpart=500;

config const smax=200.0;
config const smax2=smax**2;
config const nmubins=5;
config const nsbins=5;

// Global variables
var gtime1 : Timer;

proc main() {
  if isTest then testPairs(); else doPairs();
}

record WeightedParticle3D {
  var x : [Ddim] real;
  var w : real;
}

proc countLines(fn : string) : int {
  var ff = open(fn, iomode.r);
  var ipart = 0;
  for iff in ff.lines() do ipart +=1;
  ff.close();
  return ipart;
}

proc readFile(fn : string, pp : []WeightedParticle3D)  {
  const maxcols=25;

  var ff = open(fn, iomode.r);
  var cols : [0.. #maxcols] real;
  var icol=0;
  var ipart = 0;
  for iff in ff.lines() {
   icol = 0;
   for col1 in iff.split(spaces) {
     if (col1.length==0) then continue;
     cols[icol] = col1 : real;
     icol += 1;
   }
   if (icol < 4) then assert(false,"malformed line...");
   pp[ipart].x = cols[Ddim];
   pp[ipart].w = cols[3];
   ipart += 1;
  }
}

class SOA_WeightedParticle3D {
  var npart : int;
  var Dpart : domain(1);
  var x,y,z,w,r2 : [Dpart] real;

  proc SOA_WeightedParticle3D(pp : []WeightedParticle3D) {
    npart = pp.size;
    Dpart = 0.. #npart;

    forall ii in Dpart {
      x[ii] = pp[ii].x[0];
      y[ii] = pp[ii].x[1];
      z[ii] = pp[ii].x[2];
      w[ii] = pp[ii].w;
      r2[ii] = + reduce (pp[ii].x**2);
    }
  }
}

proc smuAccumulate(hh : UniformBins, p1,p2 : SOA_WeightedParticle3D, d1,d2 : domain(1), scale : real) {
  forall ii in d1 { // Loop over first set of particles

    var x1,y1,z1,w1,r2 : real;
    var sl, s2, l1, s1, l2, mu, wprod : real;
    x1 = p1.x[ii]; y1 = p1.y[ii]; z1 = p1.z[ii]; w1 = p1.w[ii]; r2 = p1.r2[ii];

    for jj in d2 { // Second set of particles
      mu=2*(p2.x[jj]*x1 + p2.y[jj]*y1 + p2.z[jj]*z1);
      sl = r2 - p2.r2[jj];
      l1 = r2 + p2.r2[jj];
      s2 = l1 - mu;
      l2 = l1 + mu;
      if ((s2 >= smax2) || (s2 < 1.0e-20)) then continue;
      wprod = scale * w1 * p2.w[jj];
      s1 = sqrt(s2);
      mu = sl/(s1*sqrt(l2));
      if (mu < 0) then mu = -mu;

      hh.add((s1,mu),wprod);
    }
  }
}

proc writeHist(fn : string, hh : UniformBins) {
  var ff = openwriter(fn);
  // Dump out values
  for xx in hh.bins(1) do ff.writef("%12.4dr",xx);
  ff.writeln();
  for xx in hh.bins(2) do ff.writef("%12.4dr",xx);
  ff.writeln("\n##");
  for ii in hh.Dhist.dim(1) {
    for jj in hh.Dhist.dim(2) {
      ff.writef("%20.14er ",hh[(ii,jj)]);
    }
    ff.writeln();
  }
  ff.close();
}

proc splitOn(pp : []WeightedParticle3D, scr : []WeightedParticle3D, splitDim : int, rr : []real) : int {
  var npart, lnpart : int;
  var lo = pp.domain.low;
  var hi = pp.domain.high;
  npart = pp.domain.size;
  lnpart = 0;
  for ipp in pp {
    if (ipp.x[splitDim] < rr[splitDim]) then lnpart+=1;
  }
//  var part : [pp.domain] WeightedParticle3D;
  var li, ri : int;
  li = lo; ri = (lo+lnpart);
  for ipp in pp {
    if (ipp.x[splitDim] < rr[splitDim]) {
      scr[li] = ipp;
      li+=1;
    } else {
      scr[ri] = ipp;
      ri+=1;
    }
  }
  assert(li==(lo+lnpart), "splitOn : li != lo+npart");
  assert(ri==(hi+1), "splitOn : ri != hi+1");
  pp = scr;

  return lnpart;
}

class KDNode {
  var lo, hi,npart,id : int;
  var dom : domain(1);
  var xcen : [Ddim]real;
  var rcell : real;
  var left, right : KDNode;

  proc isLeaf() : bool {
    return (left==nil) && (right==nil);
  }

  proc ~KDNode() {
    if left then delete left;
    if right then delete right;
  }

}

proc BuildTree(pp : []WeightedParticle3D, scr : []WeightedParticle3D, id : int) : KDNode{
  gtime1.start();
  var me = new KDNode();
  var dom = pp.domain;
  me.lo = dom.low;
  me.hi = dom.high;
  me.dom = {me.lo..me.hi};
  me.id = id;
  me.npart = (me.hi-me.lo)+1;

  //  work out xcen and vantage point radius
  var pmin = pp[me.lo].x;
  var pmax = pp[me.lo].x;
  for ipp in pp {
    for idim in Ddim {
      if (ipp.x[idim] < pmin[idim]) then pmin[idim] = ipp.x[idim];
      if (ipp.x[idim] > pmax[idim]) then pmax[idim] = ipp.x[idim];
    }
  }
  me.xcen = (pmax+pmin)/2.0;
  var dx : [Ddim]real;
  me.rcell = 0.0;
  var r1 : real;
  for ipp in pp {
    dx = (ipp.x - me.xcen);
    r1 = sqrt(+ reduce dx**2);
    if (r1 > me.rcell) then me.rcell = r1;
  }

  // Continue to split
  gtime1.stop();
  if (me.npart <= minpart) then return me; // Don't split further.
  gtime1.start();

  // Find dimension to split on
  dx = pmax - pmin;
  var splitDim = 0;
  for idim in Ddim {
    if (dx[idim] > dx[splitDim]) then splitDim=idim;
  }

  // Split
  var lnpart = splitOn(pp, scr,splitDim, me.xcen);
  var ldom = {me.lo..(me.lo+lnpart-1)};
  var rdom = {(me.lo+lnpart)..me.hi};
  gtime1.stop();
  me.left = BuildTree(pp[ldom],scr[ldom],2*id+1);
  me.right = BuildTree(pp[rdom],scr[rdom],2*id+2);
  return me;
}

proc TreeAccumulate(hh : UniformBins, p1, p2 : SOA_WeightedParticle3D, node1, node2 : KDNode) {
  // Compute the distance between node1 and node2
  var rr = sqrt (+ reduce(node1.xcen - node2.xcen)**2);
  var rmin = rr - (node1.rcell+node2.rcell);

  // If distance is greater than all cases
  if (rmin > smax) then return;

  // If both nodes are leaves
  if (node1.isLeaf() & node2.isLeaf()) {
    begin smuAccumulate(hh, p1, p2, node1.dom, node2.dom,1);
    return;
  }

  // If one node is a leaf
  if (node1.isLeaf()) {
    begin TreeAccumulate(hh, p1, p2, node1, node2.left);
    begin TreeAccumulate(hh, p1, p2, node1, node2.right);
    return;
  }
  if (node2.isLeaf()) {
    begin TreeAccumulate(hh, p1, p2, node1.left, node2);
    begin TreeAccumulate(hh, p1, p2, node1.right, node2);
    return;
  }

  // Split the larger case;
  if (node1.npart > node2.npart) {
    begin TreeAccumulate(hh, p1, p2, node1.left, node2);
    begin TreeAccumulate(hh, p1, p2, node1.right, node2);
    return;
  } else {
    begin TreeAccumulate(hh, p1, p2, node1, node2.left);
    begin TreeAccumulate(hh, p1, p2, node1, node2.right);
    return;
  }

}

// A simple testHarness
proc testPairs() {

  // Read in the file
  var nlines = countLines(fn1);
  var pp1 : [0.. #nlines] WeightedParticle3D;
  readFile("test.dat",pp1);

  // Build the tree
  var scr : [0.. #nlines] WeightedParticle3D;
  var root1 = BuildTree(pp1,scr,0);

  // Set up the histogram
  var hh = new UniformBins(2,(nsbins,nmubins), ((0.0,smax),(0.0,1.0+1.e-10)));

  // AOS -> SOA
  var soa1 = new SOA_WeightedParticle3D(pp1);

  // Do the paircounts with a tree
  sync TreeAccumulate(hh, soa1, soa1, root1, root1);
  for xx in hh.bins(1) do writef("%12.4dr",xx);
  writeln();
  for xx in hh.bins(2) do writef("%12.4dr",xx);
  writeln("\n##");
  for ii in hh.Dhist.dim(1) {
    for jj in hh.Dhist.dim(2) {
      writef("%20.5er ",hh[(ii,jj)]);
    }
    writeln();
  }
}

// The main code
proc doPairs() {
  var tt : Timer;

  // Read in the file
  tt.clear(); tt.start();
  var nlines1 = countLines(fn1);
  var pp1 : [0.. #nlines1] WeightedParticle3D;
  readFile(fn1,pp1);
  var nlines2 = countLines(fn2);
  var pp2 : [0.. #nlines2] WeightedParticle3D;
  readFile(fn1,pp2);
  tt.stop();
  writef("Read in %i lines from file %s \n", pp1.size, fn1);
  writef("Read in %i lines from file %s \n", pp2.size, fn2);
  writef("Time to read : %r \n", tt.elapsed());

  // Build the tree
  tt.clear(); tt.start(); gtime1.clear();
  var scr1 : [0.. #nlines1] WeightedParticle3D;
  var root1 = BuildTree(pp1,scr1,0);
  var scr2 : [0.. #nlines2] WeightedParticle3D;
  var root2 = BuildTree(pp2,scr2,0);
  tt.stop();
  writef("Time to build trees : %r \n", tt.elapsed());
  writef("Time in splitOn : %r \n", gtime1.elapsed());

  // Set up the histogram
  var hh = new UniformBins(2,(nsbins,nmubins), ((0.0,smax),(0.0,1.0+1.e-10)));

  // AOS -> SOA
  tt.clear(); tt.start();
  var soa1 = new SOA_WeightedParticle3D(pp1);
  var soa2 = new SOA_WeightedParticle3D(pp2);
  tt.stop();
  writef("Time to SOA : %r \n", tt.elapsed());

  // Brute force paircounts
  if doBrute {
    tt.clear(); tt.start();
    smuAccumulate(hh, soa1,soa2, soa1.Dpart, soa2.Dpart, 1.0);
    tt.stop();
    writef("Time to brute force paircount : %r \n", tt.elapsed());
    writeHist("%s.brute".format(pairfn),hh);
  }

   // Do the paircounts with a tree
  hh.reset();
  tt.clear(); tt.start();
  sync TreeAccumulate(hh, soa1, soa2, root1, root2);
  tt.stop();
  writef("Time to tree paircount : %r \n", tt.elapsed());
  writeHist("%s.tree".format(pairfn),hh);
}

module Histogram {

  class UniformBins {
    param dim : int;
    var nbins : dim*int;
    var Dhist : domain(dim);
    var lo, hi,dx, invdx : [1..dim] real;
    var arr : [Dhist] atomic real;

    proc UniformBins(param dim : int, nbins : dim*int, limits : dim*(real,real)) {
      var dd : dim*range;
      this.nbins = nbins;
      for param ii in 1..dim {
        lo[ii] = limits(ii)(1);
        hi[ii] = limits(ii)(2);
        dx[ii] = (hi[ii]-lo[ii])/nbins(ii);
        dd(ii) = 0.. #nbins(ii);
      }
      invdx = 1.0/dx;
      Dhist = {(...dd)};
      [x in arr] x.write(0.0);
    }

    proc reset() {
      [x in arr] x.write(0.0);
    }

    proc bins(idim : int) {
      var bounds : [0.. #nbins(idim)+1] real;
      for i in {0.. #(nbins(idim)+1)} do bounds[i] = lo[idim] + dx[idim]*i;
      return bounds;
    }

    proc add(x : dim*real, w : real) {
      for param ii in 1..dim do
        if ((x(ii) < lo[ii]) | (x(ii) >= hi[ii])) then return;

      var pos : dim*int;
      for param ii in 1..dim do pos(ii) = ((x(ii)-lo(ii))*invdx[ii]) : int;
      arr[pos].add(w);
    }

    proc this(ndx) : real {
      return arr[ndx].read();
    }

  } // UniformBins

} // End module Histogram