-- # Radix sort by key
--
-- We often wish to sort an array by some property of the elements,
-- rather than the elements themselves.  The [radix
-- sort](radix-sort.html) example can be modified to allow this by
-- parameterising over a function `f` of type `t -> u32` to sort an
-- array of type-`t` elements by the integer produced by `f`.

def radix_sort_step [n] 't (f: t -> u32) (xs: [n]t) (b: i32): [n]t =
  let bits = map (\x -> (i32.u32 (f x >> u32.i32 b)) & 1) xs
  let bits_neg = map (1-) bits
  let offs = reduce (+) 0 bits_neg
  let idxs0 = map2 (*) bits_neg (scan (+) 0 bits_neg)
  let idxs1 = map2 (*) bits (map (+offs) (scan (+) 0 bits))
  let idxs2 = map2 (+) idxs0 idxs1
  let idxs  = map (\x->x-1) idxs2
  let xs' = scatter (copy xs) (map i64.i32 idxs) xs
  in xs'

def radix_sort [n] 't (f: t -> u32) (xs: [n]t): [n]t =
  loop xs for i < 32 do radix_sort_step f xs i

-- The only change compared to the original radix sort is that we use
-- `f` to extract a representative integer.

-- # See also
--
-- [Merge sort](merge-sort.html), [matching parentheses](parens.html).
--
-- The [sorts](https://github.com/diku-dk/sorts) library.