content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  module Util.List
    ( blend
    , chunk
    , count
    , diff
    , equalChunksOf
    , atMostChunksOf
    , splitChunks
    , chunkBy
    , mapHead
    , mapLast
    , mostCommonElem
    , replace
    , sortOnM
    , stripSuffix
    , uniq
    , uniqBy
    ) where
  
  import Data.List
  import Data.Ord (comparing)
  import qualified Data.Set as Set
  
  -- This is an O(n+m) implementation for computing difference between two lists,
  -- when the order of the result is irrelevant.
  -- Returns those elements in the first list which are not present in the second.
  -- All occurences of repeated elements are removed.
  -- eg. [1,1,2,4,3] `List.diff` [1,2] -> [3,4]
  diff :: (Eq a, Ord a) => [a] -> [a] -> [a]
  diff xs ys = Set.toList $ Set.fromList xs `Set.difference` Set.fromList ys
  
  -- | Returns the number of elements that satisfy the predicate
  count :: (a -> Bool) -> [a] -> Int
  count f = length . filter f
  
  mapHead :: (a -> a) -> [a] -> [a]
  mapHead f (x:xs) = f x : xs
  mapHead _ [] = []
  
  mapLast :: (a -> a) -> [a] -> [a]
  mapLast f = reverse . mapHead f . reverse
  
  replace :: Eq a => a -> a -> [a] -> [a]
  replace a b = map (\x -> if x == a then b else x)
  
  sortOnM :: (Applicative f, Ord b) => (a -> f b) -> [a] -> f [a]
  sortOnM f xs =
    map fst . sortOn snd <$> traverse (\x -> (x,) <$> f x) xs
  
  uniq :: Ord a => [a] -> [a]
  uniq = Set.toList . Set.fromList
  
  uniqBy :: (a -> a -> Ordering) -> [a] -> [a]
  uniqBy cmp = map head . groupBy (\a b -> cmp a b == EQ) . sortBy cmp
  
  -- | Splits list into chunks of roughly equal size 'Int'
  -- optimising for variance, which means the chunks can be bigger than the 'Int'
  -- and the last chunk can be 50% bigger or 50% smaller than the others.
  --
  -- > equalChunksOf 25 [1..99] == [ [1..33], [34..66], [67..99]]
  -- > equalChunksOf 34 [1..99] == [ [1..49], [50..99] ]
  -- > equalChunksOf 49 [1..99] == [ [1..49], [50..99] ]
  -- > equalChunksOf 50 [1..99] == [ [1..99] ]
  -- > equalChunksOf 51 [1..99] == [ [1..99] ]
  --
  -- > equalChunksOf 24 [1..99] == [ [1..24], [25..48], [49..72], [73..99] ]
  -- >   -- lengths are 24, 24, 24, 27 where last chunk is 3 larger.
  -- > map length $ equalChunksOf 6 [1..81] == [6,6,6,6,6,6,6,6,6,6,6,6,9]
  -- > last $ map length $ equalChunksOf 24 [1..348] == 36 -- 50% bigger
  -- > last $ map length $ equalChunksOf 24 [1..349] == 13 -- 50% smaller
  --
  -- >
  -- > equalChunksOf i [] = [ [] ]
  -- > equalChunksOf 0 xs -- *** Exception: divide by zero
  -- > equalChunksOf (-1) xs == [ xs ]
  equalChunksOf :: Int -> [a] -> [[a]]
  equalChunksOf n as = go as size []
    where
      size = length as
      bestLen = size `div` max 1 (size `div` n)
      go bs bsLen acc | bsLen <= bestLen * 3 `div` 2 = reverse $ bs:acc
      go bs bsLen acc =
        let (chunk, rest) = splitAt bestLen bs in
          go rest (bsLen - bestLen) $ chunk:acc
  
  -- | Splits list into non-empty pieces which are at most 'Int' size limit.
  -- The actual chunk size will often be smaller (see splitChunks), using
  -- the minimum number of chunks to not go over the 'Int' size limit.  The size
  -- of chunks will all be equal (if possible) and otherwise vary by at most 1.
  --
  -- > concat (atMostChunksOf i xs) == xs
  -- > all (not . null) (atMostChunksOf i xs)
  --
  -- > atMostChunksOf 24 [1..99] == [[1..20],[21..40],[41..60],[61..80],[81..99]]
  -- > atMostChunksOf 34 [1..99] == [ [1..33], [34..66], [67..99] ]
  -- > atMostChunksOf 49 [1..99] == [ [1..33], [34..66], [67..99] ]
  -- > atMostChunksOf 50 [1..99] == [ [1..50], [51..99] ]
  -- > atMostChunksOf 51 [1..99] == [ [1..50], [51..99] ]
  -- >
  -- > atMostChunksOf i [] = [] -- never make an empty chunk by returning [ [] ]
  -- > atMostChunksOf 1 [x,y,z] = [ [x], [y], [z] ]
  -- > atMostChunksOf 0 [x,y,z] = [ [x,y,z] ] -- nonsense, non-empty list
  -- > atMostChunksOf (-1) [x,y,z] = [ [x,y,z] ]-- nonsense, non-empty list
  atMostChunksOf :: Int -> [a] -> [[a]]
  atMostChunksOf i as
      | null as = []
      | i <= 0 = [as]
      | size <= i = [as]
      | otherwise = splitChunks pieces as
    where
      size = length as
      pieces = (size + pred i) `div` i   -- assert size <= pieces * i
  
  -- | Splits list into at most 'Int' non-empty chunks of nearly equal size.
  -- Initial larger chunks can be 1 longer than later smaller chunks.
  --
  -- > concat (splitChunks i xs) == xs
  -- > all (not . null) (splitChunks i xs)
  --
  -- > splitChunks _ [] = [] -- never make an empty chunk by returning [ [] ]
  -- > splitChunks 4 [x,y,z] = [ [x], [y], [z] ] -- only 3 non-empty chunks
  -- > splitChunks 3 [x,y,z] = [ [x], [y], [z] ]
  -- > splitChunks 2 [x,y,z] = [ [x,y], [z] ] -- larger groups first
  -- > splitChunks 1 [x,y,z] = [ [x,y,z] ]
  -- > splitChunks 0 [x,y,z] = [ [x,y,z] ] -- nonsense, non-empty list
  -- > splitChunks (-1) [x,y,z] = [ [x,y,z] ] -- nonsense, non-empty list
  splitChunks :: Int -> [a] -> [[a]]
  splitChunks b as
      | null as = []
      | b <= 0 = [as]
      | size <= b = map (:[]) as
      | otherwise = go as sizes []
    where
      size = length as
      larges = size `mod` b   -- assert 0 <= larges < b
      smalls = b - larges     -- assert 0 < smalls <= b
      -- assert larges + smalls == b
      smaller = size `div` b  -- assert 1 <= smaller ; since 0 < b < size
      larger = succ smaller   -- assert 2 <= larger
      sizes = replicate larges larger ++ replicate smalls smaller
      -- assert 0 < length sizes == b
      -- assert sum sizes == size
      go [] _ acc = reverse acc      -- base case
      go bs (x:xs) acc =
        let (chunk, rest) = splitAt x bs -- assert 1 <= length chunk
        in go rest xs (chunk:acc)
      go bs [] acc = reverse $ bs:acc -- should not happen
  
  -- | @chunkBy n f xs@: divide @xs@ into chunks of at most size @n@,
  -- where the size of an element @x@ of @xs@ is given by @f x@. If the size of
  -- any element is greater than @n@, then it will be in a chunk on its own.
  --
  -- > chunkBy 4 length ["ab", "cde", "f", "ghijk"]
  -- > [["ab"],["cde","f"],["ghijk"]]
  --
  chunkBy :: Int -> (a -> Int) -> [a] -> [[a]]
  chunkBy _ _ [] = []
  chunkBy chunkSize getSize (x:xs)
    | chunkSize <= 0 = [x:xs]
    | otherwise = go (getSize x) [x] xs
    where
      -- invariant: chunk has at least one element
      go _ chunk [] = [reverse chunk]
      go n chunk (x:xs)
        | n + xSize > chunkSize = reverse chunk : go xSize [x] xs
        | otherwise = go (n + xSize) (x:chunk) xs
        where xSize = getSize x
  
  stripSuffix :: Eq a => [a] -> [a] -> Maybe [a]
  stripSuffix xs ys
    | Just zs <- stripPrefix (reverse xs) (reverse ys) = Just (reverse zs)
    | otherwise = Nothing
  
  -- | Returns the most common element in the list.
  -- Throws if given an empty list.
  mostCommonElem :: Ord a => [a] -> a
  mostCommonElem = head . maximumBy (comparing length) . group . sort
  
  blend :: [a] -> [a] -> [a]
  blend [] _ = []
  blend _ [] = []
  blend (x:xs) (y:ys) = x:y:blend xs ys
  
  -- | Divide a range @0..(k-1)@ into @n@ chunks. Each pair in the result
  -- list is the (start,length) of a chunk.
  chunk
    :: Int                                -- ^ @n@
    -> Int                                -- ^ @k@
    -> [(Int,Int)]                        -- ^ @[(start,length)]@
  chunk !n !k = zip (scanl (+) 0 lens) lens
    where
      lens = replicate r (q+1) ++ if q == 0 then [] else replicate (n-r) q
      !q = k `div` n
      !r = k `mod` n
  