content: {-# LANGUAGE ImportQualifiedPost #-}
  
  -- |
  -- Module      :  Disco.Util
  -- Copyright   :  disco team and contributors
  -- Maintainer  :  byorgey@gmail.com
  --
  -- SPDX-License-Identifier: BSD-3-Clause
  --
  -- Miscellaneous utilities.
  module Disco.Util where
  
  import Control.Applicative (Alternative)
  import Control.Monad (guard)
  import Data.Bifunctor (bimap)
  import Data.List.NonEmpty (NonEmpty)
  import Data.List.NonEmpty qualified as NE
  import Data.Map qualified as M
  
  infixr 1 ==>
  
  -- | A synonym for pairing which makes convenient syntax for
  --   constructing literal maps via M.fromList.
  (==>) :: a -> b -> (a, b)
  (==>) = (,)
  
  -- | Flipped variant of 'map'.
  for :: [a] -> (a -> b) -> [b]
  for = flip map
  
  -- | A variant of 'Map' indexing that throws a custom error message
  --   in case the key is not found, to help with debugging.
  (!) :: (Show k, Ord k) => M.Map k v -> k -> v
  m ! k = case M.lookup k m of
    Nothing -> error $ "key " ++ show k ++ " is not an element in the map"
    Just v -> v
  
  -- | Find the maximum of a list of positive numbers, yielding 0 in the
  --   case of an empty list.
  maximum0 :: (Num a, Ord a) => [a] -> a
  maximum0 [] = 0
  maximum0 xs = maximum xs
  
  -- | A variant of 'filter' that returns a @Maybe (NonEmpty a)@ instead
  --   of a regular list.
  filterNE :: (a -> Bool) -> NonEmpty a -> Maybe (NonEmpty a)
  filterNE p = NE.nonEmpty . NE.filter p
  
  -- | A variant of 'partition' that returns @Maybe (NonEmpty a)@s instead
  --   of regular lists.
  partitionNE :: (a -> Bool) -> NonEmpty a -> (Maybe (NonEmpty a), Maybe (NonEmpty a))
  partitionNE p as = (filterNE p as, filterNE (not . p) as)
  
  -- | A variant of 'partitionEithers' for nonempty lists.  If the
  --   result is Left, it means all the inputs were Left.  If the result
  --   is Right, we definitely have some Rights, and possibly some Lefts
  --   as well.  This properly encodes the fact that at least one result
  --   list must be nonempty.
  partitionEithersNE :: NonEmpty (Either a b) -> Either (NonEmpty a) ([a], NonEmpty b)
  partitionEithersNE = foldr1 combine . NE.map (bimap NE.singleton (([],) . NE.singleton))
   where
    combine :: Either (NonEmpty a) ([a], NonEmpty b) -> Either (NonEmpty a) ([a], NonEmpty b) -> Either (NonEmpty a) ([a], NonEmpty b)
    combine (Left as1) (Left as2) = Left (NE.append as1 as2)
    combine (Left as1) (Right (as2, bs)) = Right (NE.toList as1 ++ as2, bs)
    combine (Right (as1, bs)) (Left as2) = Right (as1 ++ NE.toList as2, bs)
    combine (Right (as1, bs1)) (Right (as2, bs2)) = Right (as1 ++ as2, NE.append bs1 bs2)
  
  -- | Iterate a function until finding the first value that satisfies
  --   the given predicate.  @iterUntil f p@ is equivalent to @head
  --   . filter p . iterate f@ but does not trigger a partiality
  --   warning.
  iterUntil :: (a -> a) -> (a -> Maybe b) -> a -> b
  iterUntil f p a = case p a of
    Just b -> b
    _ -> iterUntil f p (f a)
  
  -- | Allow a value through only if it satisfies the given predicate.
  gate :: Alternative f => (a -> Bool) -> a -> f a
  gate p a = a <$ guard (p a)
  