{-# 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)
==> :: forall a b. a -> b -> (a, b)
(==>) = (,)

-- | Flipped variant of 'map'.
for :: [a] -> (a -> b) -> [b]
for :: forall a b. [a] -> (a -> b) -> [b]
for = ((a -> b) -> [a] -> [b]) -> [a] -> (a -> b) -> [b]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> b) -> [a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
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
Map k v
m ! :: forall k v. (Show k, Ord k) => Map k v -> k -> v
! k
k = case k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
k Map k v
m of
  Maybe v
Nothing -> [Char] -> v
forall a. HasCallStack => [Char] -> a
error ([Char] -> v) -> [Char] -> v
forall a b. (a -> b) -> a -> b
$ [Char]
"key " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ k -> [Char]
forall a. Show a => a -> [Char]
show k
k [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" is not an element in the map"
  Just v
v -> 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 :: forall a. (Num a, Ord a) => [a] -> a
maximum0 [] = a
0
maximum0 [a]
xs = [a] -> a
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [a]
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 :: forall a. (a -> Bool) -> NonEmpty a -> Maybe (NonEmpty a)
filterNE a -> Bool
p = [a] -> Maybe (NonEmpty a)
forall a. [a] -> Maybe (NonEmpty a)
NE.nonEmpty ([a] -> Maybe (NonEmpty a))
-> (NonEmpty a -> [a]) -> NonEmpty a -> Maybe (NonEmpty a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Bool) -> NonEmpty a -> [a]
forall a. (a -> Bool) -> NonEmpty a -> [a]
NE.filter a -> Bool
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 :: forall a.
(a -> Bool)
-> NonEmpty a -> (Maybe (NonEmpty a), Maybe (NonEmpty a))
partitionNE a -> Bool
p NonEmpty a
as = ((a -> Bool) -> NonEmpty a -> Maybe (NonEmpty a)
forall a. (a -> Bool) -> NonEmpty a -> Maybe (NonEmpty a)
filterNE a -> Bool
p NonEmpty a
as, (a -> Bool) -> NonEmpty a -> Maybe (NonEmpty a)
forall a. (a -> Bool) -> NonEmpty a -> Maybe (NonEmpty a)
filterNE (Bool -> Bool
not (Bool -> Bool) -> (a -> Bool) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
p) NonEmpty a
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 :: forall a b.
NonEmpty (Either a b) -> Either (NonEmpty a) ([a], NonEmpty b)
partitionEithersNE = (Either (NonEmpty a) ([a], NonEmpty b)
 -> Either (NonEmpty a) ([a], NonEmpty b)
 -> Either (NonEmpty a) ([a], NonEmpty b))
-> NonEmpty (Either (NonEmpty a) ([a], NonEmpty b))
-> Either (NonEmpty a) ([a], NonEmpty b)
forall a. (a -> a -> a) -> NonEmpty a -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 Either (NonEmpty a) ([a], NonEmpty b)
-> Either (NonEmpty a) ([a], NonEmpty b)
-> Either (NonEmpty a) ([a], NonEmpty b)
forall a b.
Either (NonEmpty a) ([a], NonEmpty b)
-> Either (NonEmpty a) ([a], NonEmpty b)
-> Either (NonEmpty a) ([a], NonEmpty b)
combine (NonEmpty (Either (NonEmpty a) ([a], NonEmpty b))
 -> Either (NonEmpty a) ([a], NonEmpty b))
-> (NonEmpty (Either a b)
    -> NonEmpty (Either (NonEmpty a) ([a], NonEmpty b)))
-> NonEmpty (Either a b)
-> Either (NonEmpty a) ([a], NonEmpty b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Either a b -> Either (NonEmpty a) ([a], NonEmpty b))
-> NonEmpty (Either a b)
-> NonEmpty (Either (NonEmpty a) ([a], NonEmpty b))
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NE.map ((a -> NonEmpty a)
-> (b -> ([a], NonEmpty b))
-> Either a b
-> Either (NonEmpty a) ([a], NonEmpty b)
forall a b c d. (a -> b) -> (c -> d) -> Either a c -> Either b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap a -> NonEmpty a
forall a. a -> NonEmpty a
NE.singleton (([],) (NonEmpty b -> ([a], NonEmpty b))
-> (b -> NonEmpty b) -> b -> ([a], NonEmpty b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> NonEmpty b
forall a. a -> NonEmpty a
NE.singleton))
 where
  combine :: Either (NonEmpty a) ([a], NonEmpty b) -> Either (NonEmpty a) ([a], NonEmpty b) -> Either (NonEmpty a) ([a], NonEmpty b)
  combine :: forall a b.
Either (NonEmpty a) ([a], NonEmpty b)
-> Either (NonEmpty a) ([a], NonEmpty b)
-> Either (NonEmpty a) ([a], NonEmpty b)
combine (Left NonEmpty a
as1) (Left NonEmpty a
as2) = NonEmpty a -> Either (NonEmpty a) ([a], NonEmpty b)
forall a b. a -> Either a b
Left (NonEmpty a -> NonEmpty a -> NonEmpty a
forall a. NonEmpty a -> NonEmpty a -> NonEmpty a
NE.append NonEmpty a
as1 NonEmpty a
as2)
  combine (Left NonEmpty a
as1) (Right ([a]
as2, NonEmpty b
bs)) = ([a], NonEmpty b) -> Either (NonEmpty a) ([a], NonEmpty b)
forall a b. b -> Either a b
Right (NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty a
as1 [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
as2, NonEmpty b
bs)
  combine (Right ([a]
as1, NonEmpty b
bs)) (Left NonEmpty a
as2) = ([a], NonEmpty b) -> Either (NonEmpty a) ([a], NonEmpty b)
forall a b. b -> Either a b
Right ([a]
as1 [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty a
as2, NonEmpty b
bs)
  combine (Right ([a]
as1, NonEmpty b
bs1)) (Right ([a]
as2, NonEmpty b
bs2)) = ([a], NonEmpty b) -> Either (NonEmpty a) ([a], NonEmpty b)
forall a b. b -> Either a b
Right ([a]
as1 [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
as2, NonEmpty b -> NonEmpty b -> NonEmpty b
forall a. NonEmpty a -> NonEmpty a -> NonEmpty a
NE.append NonEmpty b
bs1 NonEmpty b
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 :: forall a b. (a -> a) -> (a -> Maybe b) -> a -> b
iterUntil a -> a
f a -> Maybe b
p a
a = case a -> Maybe b
p a
a of
  Just b
b -> b
b
  Maybe b
_ -> (a -> a) -> (a -> Maybe b) -> a -> b
forall a b. (a -> a) -> (a -> Maybe b) -> a -> b
iterUntil a -> a
f a -> Maybe b
p (a -> a
f a
a)

-- | Allow a value through only if it satisfies the given predicate.
gate :: Alternative f => (a -> Bool) -> a -> f a
gate :: forall (f :: * -> *) a. Alternative f => (a -> Bool) -> a -> f a
gate a -> Bool
p a
a = a
a a -> f () -> f a
forall a b. a -> f b -> f a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Bool -> f ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (a -> Bool
p a
a)