{-# LANGUAGE BangPatterns, MultiParamTypeClasses #-}
{-|
== Miscellaneous functions/recipes

-}

module Misc
    ( pairs
    , fArray
    , chunksOf
    , replicateL
    , unique
    , foldExclusive
    , modifyArray
    , modifyArray'
    , ifoldr
    , ifoldl'
    , foldMComp
    , farthest
    , foldTree'
    , Commutative
    , Group(..)
    , Idempotent
    , Action(..)
    , bitLength
    , unsafeBit
    , odds
    , evens
    , orM
    , andM
    , anyM
    , allM
    , minimumByMaybe
    , maximumByMaybe
    ) where

import Control.Monad
import Data.Array.IArray
import Data.Array.MArray
import Data.Bits
import Data.List
import Data.Semigroup
import Data.Tree

-- | Generates distinct pairs of elements from a list. O(n^2).
pairs :: [a] -> [(a, a)]
pairs :: [a] -> [(a, a)]
pairs [a]
xs = [(a
x, a
y) | a
x:[a]
ys <- [a] -> [[a]]
forall a. [a] -> [[a]]
tails [a]
xs, a
y <- [a]
ys]

-- | Generates an Array from bounds and a function. O(n) assuming f takes O(1).
fArray :: (IArray a e, Ix i) => (i, i) -> (i -> e) -> a i e
fArray :: (i, i) -> (i -> e) -> a i e
fArray (i, i)
b i -> e
f = (i, i) -> [e] -> a i e
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
listArray (i, i)
b (i -> e
f (i -> e) -> [i] -> [e]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (i, i) -> [i]
forall a. Ix a => (a, a) -> [a]
range (i, i)
b)
{-# INLINE fArray #-}

-- | Splits a list into chunks of fixed size. O(n).
chunksOf :: Int -> [a] -> [[a]]
chunksOf :: Int -> [a] -> [[a]]
chunksOf Int
n = ([a] -> Maybe ([a], [a])) -> [a] -> [[a]]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr [a] -> Maybe ([a], [a])
forall a. [a] -> Maybe ([a], [a])
f where
    f :: [a] -> Maybe ([a], [a])
f [] = Maybe ([a], [a])
forall a. Maybe a
Nothing
    f [a]
xs = ([a], [a]) -> Maybe ([a], [a])
forall a. a -> Maybe a
Just (Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [a]
xs)

-- | Replicates a list n times. O(nm) where m is the length of the list.
replicateL :: Int -> [a] -> [a]
replicateL :: Int -> [a] -> [a]
replicateL Int
n = [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[a]] -> [a]) -> ([a] -> [[a]]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [a] -> [[a]]
forall a. Int -> a -> [a]
replicate Int
n

-- | Eliminates consecutive duplicate elements. O(n).
unique :: Eq a => [a] -> [a]
unique :: [a] -> [a]
unique = ([a] -> a) -> [[a]] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map [a] -> a
forall a. [a] -> a
head ([[a]] -> [a]) -> ([a] -> [[a]]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [[a]]
forall a. Eq a => [a] -> [[a]]
group

-- | Folds strictly such that the ith element of the output list contains the fold of all elements in the
-- input list except for the ith element. The fold function f must be commutative, in the sense that
-- (b `f` a1) `f` a2 = (b `f` a2) `f` a1
-- f is called O(n log n) times.
foldExclusive :: (b -> a -> b) -> b -> [a] -> [b]
foldExclusive :: (b -> a -> b) -> b -> [a] -> [b]
foldExclusive b -> a -> b
_ b
_ [] = []
foldExclusive b -> a -> b
f b
y0 [a]
xs0 = b -> Int -> [a] -> [b] -> [b]
go b
y0 ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs0) [a]
xs0 [] where
    go :: b -> Int -> [a] -> [b] -> [b]
go !b
y Int
1 [a]
_ = (b
yb -> [b] -> [b]
forall a. a -> [a] -> [a]
:)
    go b
y Int
n [a]
xs = b -> Int -> [a] -> [b] -> [b]
go b
yr Int
n' [a]
xsl ([b] -> [b]) -> ([b] -> [b]) -> [b] -> [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Int -> [a] -> [b] -> [b]
go b
yl (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n') [a]
xsr where
        n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
        ([a]
xsl, [a]
xsr) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n' [a]
xs
        yl :: b
yl = (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' b -> a -> b
f b
y [a]
xsl
        yr :: b
yr = (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' b -> a -> b
f b
y [a]
xsr

-- | Modifies an element in a mutable array.
modifyArray :: (MArray a e m, Ix i) => a i e -> i -> (e -> e) -> m ()
modifyArray :: a i e -> i -> (e -> e) -> m ()
modifyArray a i e
a i
i e -> e
f = a i e -> i -> m e
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray a i e
a i
i m e -> (e -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a i e -> i -> e -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray a i e
a i
i (e -> m ()) -> (e -> e) -> e -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> e
f
{-# INLINE modifyArray #-}

-- | Modifies an element in a mutable array. Strict version.
modifyArray' :: (MArray a e m, Ix i) => a i e -> i -> (e -> e) -> m ()
modifyArray' :: a i e -> i -> (e -> e) -> m ()
modifyArray' a i e
a i
i e -> e
f = a i e -> i -> m e
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray a i e
a i
i m e -> (e -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (a i e -> i -> e -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray a i e
a i
i (e -> m ()) -> e -> m ()
forall a b. (a -> b) -> a -> b
$!) (e -> m ()) -> (e -> e) -> e -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> e
f
{-# INLINE modifyArray' #-}

-- | foldr with an index, starting at 0.
ifoldr :: Foldable f => (Int -> a -> b -> b) -> b -> f a -> b
ifoldr :: (Int -> a -> b -> b) -> b -> f a -> b
ifoldr Int -> a -> b -> b
f b
z0 = \f a
xs -> (a -> (Int -> b) -> Int -> b) -> (Int -> b) -> f a -> Int -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
x Int -> b
k !Int
i -> Int -> a -> b -> b
f Int
i a
x (Int -> b
k (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))) (b -> Int -> b
forall a b. a -> b -> a
const b
z0) f a
xs Int
0
{-# INLINE ifoldr #-}

-- | foldl' with an index, starting at 0.
ifoldl' :: Foldable f => (b -> Int -> a -> b) -> b -> f a -> b
ifoldl' :: (b -> Int -> a -> b) -> b -> f a -> b
ifoldl' b -> Int -> a -> b
f b
z0 = \f a
xs -> (a -> (Int -> b -> b) -> Int -> b -> b)
-> (Int -> b -> b) -> f a -> Int -> b -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
x Int -> b -> b
k !Int
i !b
z -> Int -> b -> b
k (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (b -> Int -> a -> b
f b
z Int
i a
x)) ((b -> b) -> Int -> b -> b
forall a b. a -> b -> a
const b -> b
forall a. a -> a
id) f a
xs Int
0 b
z0
{-# INLINE ifoldl' #-}

-- | Compose a strict left fold function with a mapM function. Useful for foldMs.
-- foldM (f `foldMComp` g) z = fmap (foldl' f z) . mapM g
foldMComp :: Monad m => (b -> a -> b) -> (c -> m a) -> b -> c -> m b
foldMComp :: (b -> a -> b) -> (c -> m a) -> b -> c -> m b
foldMComp b -> a -> b
f c -> m a
g = \b
z c
x -> b -> a -> b
f b
z (a -> b) -> m a -> m b
forall (m :: * -> *) a b. Monad m => (a -> b) -> m a -> m b
<$!> c -> m a
g c
x
{-# INLINE foldMComp #-}

-- | Repeatedly applies a function to a value and returns the value which gives back Nothing.
farthest :: (a -> Maybe a) -> a -> a
farthest :: (a -> Maybe a) -> a -> a
farthest a -> Maybe a
f = a -> a
go where go :: a -> a
go a
x = a -> (a -> a) -> Maybe a -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe a
x a -> a
go (a -> Maybe a
f a
x)

-- | Folds a tree. This does the same job as Data.Tree.foldTree but with different fold functions, which
-- may be preferable in cases such as CPS folds.
foldTree' :: (a -> b -> c) -> (c -> b -> b) -> b -> Tree a -> c
foldTree' :: (a -> b -> c) -> (c -> b -> b) -> b -> Tree a -> c
foldTree' a -> b -> c
f c -> b -> b
g b
z = Tree a -> c
go where go :: Tree a -> c
go (Node a
x Forest a
ts) = a -> b -> c
f a
x ((Tree a -> b -> b) -> b -> Forest a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (c -> b -> b
g (c -> b -> b) -> (Tree a -> c) -> Tree a -> b -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree a -> c
go) b
z Forest a
ts)
{-# INLINE foldTree' #-}

-- | A semigroup where the operation (<>) is commutative.
-- a <> b = b <> a
class Semigroup a => Commutative a

-- | A monoid where elements have inverses.
-- a <> invert a = mempty
-- invert a <> a = mempty
class Monoid a => Group a where
    invert :: a -> a

instance Num a => Commutative (Sum a)

instance Num a => Group (Sum a) where
    invert :: Sum a -> Sum a
invert = Sum a -> Sum a
forall a. Num a => a -> a
negate
    {-# INLINE invert #-}

-- | A semigroup where the elements are idempotent.
-- a <> a = a
class Semigroup a => Idempotent a

instance Ord a => Idempotent (Max a)

instance Ord a => Idempotent (Min a)

instance Idempotent (First a)

instance Idempotent (Last a)

-- | A right monoid action of u on a.
-- (x `act` u1) `act` u2 = x `act` (u1 <> u2)
-- x `act` mempty = x
class (Monoid u, Monoid a) => Action u a where
    act :: a -> u -> a

-- | The number of bits required to represent the value.
bitLength :: FiniteBits b => b -> Int
bitLength :: b -> Int
bitLength b
x = b -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize b
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- b -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros b
x
{-# INLINE bitLength #-}

-- | Just like bit, but skips the check that the index is in [0 .. size-1].
unsafeBit :: (Bits a, Num a) => Int -> a
unsafeBit :: Int -> a
unsafeBit = a -> Int -> a
forall a. Bits a => a -> Int -> a
unsafeShiftL a
1
{-# INLINE unsafeBit #-}

-- | The elements at odd positions of a list.
odds :: [a] -> [a]
odds :: [a] -> [a]
odds (a
_:a
x:[a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a]
forall a. [a] -> [a]
odds [a]
xs
odds [a]
_        = []

-- | The elements at even positions of a list.
evens :: [a] -> [a]
evens :: [a] -> [a]
evens (a
x:a
_:[a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a]
forall a. [a] -> [a]
evens [a]
xs
evens [a
x]      = [a
x]
evens []       = []

-- | Short-circuiting monadic ||.
orM :: Monad m => m Bool -> m Bool -> m Bool
orM :: m Bool -> m Bool -> m Bool
orM m Bool
m1 m Bool
m2 = m Bool
m1 m Bool -> (Bool -> m Bool) -> m Bool
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
x -> if Bool
x then Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True else m Bool
m2

-- | Short-circuiting monadic &&.
andM :: Monad m => m Bool -> m Bool -> m Bool
andM :: m Bool -> m Bool -> m Bool
andM m Bool
m1 m Bool
m2 = m Bool
m1 m Bool -> (Bool -> m Bool) -> m Bool
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
x -> if Bool
x then m Bool
m2 else Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

-- | Monadic version of any.
anyM :: (Monad m, Foldable f) => (a -> m Bool) -> f a -> m Bool
anyM :: (a -> m Bool) -> f a -> m Bool
anyM a -> m Bool
f = (a -> m Bool -> m Bool) -> m Bool -> f a -> m Bool
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (m Bool -> m Bool -> m Bool
forall (m :: * -> *). Monad m => m Bool -> m Bool -> m Bool
orM (m Bool -> m Bool -> m Bool)
-> (a -> m Bool) -> a -> m Bool -> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> m Bool
f) (Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False)

-- | Monadic version of all.
allM :: (Monad m, Foldable f) => (a -> m Bool) -> f a -> m Bool
allM :: (a -> m Bool) -> f a -> m Bool
allM a -> m Bool
f = (a -> m Bool -> m Bool) -> m Bool -> f a -> m Bool
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (m Bool -> m Bool -> m Bool
forall (m :: * -> *). Monad m => m Bool -> m Bool -> m Bool
andM (m Bool -> m Bool -> m Bool)
-> (a -> m Bool) -> a -> m Bool -> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> m Bool
f) (Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True)

maximumByMaybe :: Foldable f => (a -> a -> Ordering) -> f a -> Maybe a
maximumByMaybe :: (a -> a -> Ordering) -> f a -> Maybe a
maximumByMaybe a -> a -> Ordering
cmp f a
xs = if f a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null f a
xs then Maybe a
forall a. Maybe a
Nothing else a -> Maybe a
forall a. a -> Maybe a
Just ((a -> a -> Ordering) -> f a -> a
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
maximumBy a -> a -> Ordering
cmp f a
xs)
{-# INLINE maximumByMaybe #-}

minimumByMaybe :: Foldable f => (a -> a -> Ordering) -> f a -> Maybe a
minimumByMaybe :: (a -> a -> Ordering) -> f a -> Maybe a
minimumByMaybe a -> a -> Ordering
cmp f a
xs = if f a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null f a
xs then Maybe a
forall a. Maybe a
Nothing else a -> Maybe a
forall a. a -> Maybe a
Just ((a -> a -> Ordering) -> f a -> a
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
minimumBy a -> a -> Ordering
cmp f a
xs)
{-# INLINE minimumByMaybe #-}