sort.h

// Copyright (C) 2021-2025 Igor Baratta and Paul T. Kühner
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include <algorithm>
#include <bit>
#include <cassert>
#include <concepts>
#include <cstdint>
#include <functional>
#include <iterator>
#include <limits>
#include <numeric>
#include <optional>
#include <span>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace dolfinx
{
struct _unsigned_projection
{
  // Transforms the projected value to an unsigned int (if signed),
  // while maintaining relative order by
  //    x ↦ x + |std::numeric_limits<I>::min()|
  template <std::signed_integral T>
  constexpr std::make_unsigned_t<T> operator()(T e) const noexcept
  {
    using uT = std::make_unsigned_t<T>;
 
    // Assert binary structure for bit shift
    static_assert(static_cast<uT>(std::numeric_limits<T>::min())
                      + static_cast<uT>(std::numeric_limits<T>::max())
                  == static_cast<uT>(T(-1)));
    static_assert(std::numeric_limits<uT>::digits
                  == std::numeric_limits<T>::digits + 1);
    static_assert(std::bit_cast<uT>(std::numeric_limits<T>::min())
                  == (uT(1) << (sizeof(T) * 8 - 1)));
 
    return std::bit_cast<uT>(std::forward<T>(e))
           ^ (uT(1) << (sizeof(T) * 8 - 1));
  }
};

inline constexpr _unsigned_projection unsigned_projection{};

template <int BITS = 8, typename P = std::identity,
          std::ranges::random_access_range R>
constexpr void radix_sort(R&& range, P proj = {})
{
  using bits_t = std::make_unsigned_t<
      std::remove_cvref_t<std::invoke_result_t<P, std::iter_value_t<R>>>>;
  constexpr bits_t _BITS = BITS;
 
  // Value type
  using T = std::iter_value_t<R>;
 
  // Index type (if no projection is provided it holds I == T)
  using I = std::remove_cvref_t<std::invoke_result_t<P, T>>;
  using uI = std::make_unsigned_t<I>;
 
  if constexpr (!std::is_same_v<uI, I>)
  {
    radix_sort<_BITS>(std::forward<R>(range), [&proj](const T& e) -> uI
                      { return unsigned_projection(proj(e)); });
    return;
  }
 
  if (range.size() <= 1)
    return;
 
  // Sort N bits at a time
  constexpr uI bucket_size = 1 << _BITS;
  uI mask = (uI(1) << _BITS) - 1;
  constexpr uI top_bit = uI(1) << (sizeof(uI) * 8 - 1);
 
  // Adjacency list arrays for computing insertion position. counter is
  // pre-filled below with the first pass's histogram (bucketing on the
  // low BITS bits, which is always the correct pass-0 bucket regardless
  // of the top-bit special case), so that pass can skip a second full
  // traversal to build it.
  std::array<I, bucket_size> counter{};
  std::array<I, bucket_size> offset;
 
  // Single pass computing the maximum projected value, whether all
  // elements share the top bit (in which case it carries no ordering
  // information and can be dropped, reducing the iteration count), and
  // the first pass's histogram
  uI max_value = 0;
  bool all_first_bit = true;
  for (const auto& e : range)
  {
    uI v = proj(e);
    max_value = std::max(max_value, v);
    all_first_bit = all_first_bit && (v & top_bit);
    counter[v & mask]++;
  }
 
  if (all_first_bit)
    max_value = max_value & ~top_bit;
 
  // Compute number of iterations, most significant digit (N bits) of
  // maxvalue
  I its = 0;
  while (max_value)
  {
    max_value >>= _BITS;
    its++;
  }
 
  uI mask_offset = 0;
  std::vector<T> buffer(range.size());
  std::span<T> current_perm = range;
  std::span<T> next_perm = buffer;
  for (I i = 0; i < its; i++)
  {
    if (i > 0)
    {
      // Zero counter array
      std::ranges::fill(counter, 0);
 
      // Count number of elements per bucket
      for (auto c : current_perm)
        counter[(proj(c) & mask) >> mask_offset]++;
    }
 
    // Exclusive prefix sum, used directly as the insertion cursor for
    // each bucket
    std::exclusive_scan(counter.begin(), counter.end(), offset.begin(), I(0));
    for (auto c : current_perm)
    {
      uI bucket = (proj(c) & mask) >> mask_offset;
      next_perm[offset[bucket]++] = c;
    }
 
    mask = mask << _BITS;
    mask_offset += _BITS;
 
    std::swap(current_perm, next_perm);
  }
 
  // Copy data back to array
  if (its % 2 != 0)
    std::ranges::copy(buffer, range.begin());
}

template <typename T, int BITS = 16>
std::vector<std::int32_t> sort_by_perm(std::span<const T> x, std::size_t shape1,
                                       std::optional<std::size_t> ncols
                                       = std::nullopt)
{
  static_assert(std::is_integral_v<T>, "Integral required.");
 
  if (x.empty())
    return std::vector<std::int32_t>{};
 
  assert(shape1 > 0);
  assert(x.size() % shape1 == 0);
  std::size_t n = ncols.value_or(shape1);
  assert(n <= shape1);
  const std::size_t shape0 = x.size() / shape1;
  std::vector<std::int32_t> perm(shape0);
  std::iota(perm.begin(), perm.end(), 0);
 
  // Sort by each of the leading `n` columns, right to left. Col 0 has
  // the most significant "digit"; any columns from `n` to `shape1 - 1`
  // are excluded from the key.
  std::vector<T> column(shape0);
  for (std::size_t i = 0; i < n; ++i)
  {
    std::size_t col = n - 1 - i;
    for (std::size_t j = 0; j < shape0; ++j)
      column[j] = x[j * shape1 + col];
    radix_sort<BITS>(perm, [column = std::cref(column)](auto index)
                     { return column.get()[index]; });
  }
 
  return perm;
}

template <typename T, int BITS = 16>
std::vector<std::int32_t> sort_by_perm(std::span<std::span<const T>> x)
{
  static_assert(std::is_integral_v<T>, "Integral required.");
  if (x.empty())
    return {};
 
  std::size_t shape1 = x.size();
  std::size_t shape0 = x.front().size();
 
  std::vector<std::int32_t> perm(shape0);
  std::iota(perm.begin(), perm.end(), 0);
 
  // Each column is already its own contiguous span -- no copy needed,
  // just index into `x` for the column itself.
  for (std::size_t i = 0; i < shape1; ++i)
  {
    std::size_t col = shape1 - 1 - i;
    assert(x[col].size() == shape0);
    std::span<const T> column = x[col];
    radix_sort<BITS>(perm, [column](auto index) { return column[index]; });
  }
  return perm;
}
 
} // namespace dolfinx