-- |
-- Module      : Conjure.Utils
-- Copyright   : (c) 2021 Rudy Matela
-- License     : 3-Clause BSD  (see the file LICENSE)
-- Maintainer  : Rudy Matela <rudy@matela.com.br>
--
-- An internal module of "Conjure".
-- This exports 'Data.List', 'Data.Maybe', 'Data.Function'
-- and a few other simple utitilites.
{-# LANGUAGE CPP #-}
module Conjure.Utils
  ( module Data.List
  , module Data.Function
  , module Data.Maybe
  , module Data.Monoid
  , module Data.Tuple
  , module Data.Typeable

  , count
  , nubOn
  , nubSort
  , mzip
  , groupOn
#if __GLASGOW_HASKELL__ < 710
  , sortOn
#endif
  , idIO
  , mapHead
  , sets
  , headOr
  , allEqual
  , allEqualOn
  , allEqual2
  , choices
  , choicesThat
  , foldr0
  , indent
  , indentBy
  , classify   -- from LeanCheck.Stats
  , classifyBy -- from LeanCheck.Stats
  , classifyOn -- from LeanCheck.Stats
  , none
  , updateAt
  , first
  , second
  , both
  , (+++)
  )
where

import Data.List
import Data.Function
import Data.Maybe
import Data.Monoid
import Data.Tuple
import Data.Typeable

import System.IO.Unsafe

import Test.LeanCheck.Stats (classify, classifyBy, classifyOn)


-- | Checks if all elements of a list are equal.
allEqual :: Eq a => [a] -> Bool
allEqual :: forall a. Eq a => [a] -> Bool
allEqual (a
x:a
y:[a]
etc)  =  a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y Bool -> Bool -> Bool
&& [a] -> Bool
forall a. Eq a => [a] -> Bool
allEqual (a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
etc)
allEqual [a]
_          =  Bool
True


allEqualOn :: Eq b => (a -> b) -> [a] -> Bool
allEqualOn :: forall b a. Eq b => (a -> b) -> [a] -> Bool
allEqualOn a -> b
f  =  [b] -> Bool
forall a. Eq a => [a] -> Bool
allEqual ([b] -> Bool) -> ([a] -> [b]) -> [a] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> [a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map a -> b
f


-- | Checks if all elements of a list are equal.
--
-- Exceptionally this function returns false for an empty or unit list.
--
-- This returns true when all elements are equal and the list
-- has a length greater than or equal to two.
allEqual2 :: Eq a => [a] -> Bool
allEqual2 :: forall a. Eq a => [a] -> Bool
allEqual2 []  =  Bool
False
allEqual2 [a
_]  =  Bool
False
allEqual2 [a]
xs  =  [a] -> Bool
forall a. Eq a => [a] -> Bool
allEqual [a]
xs

-- | Counts the number of occurrences on a list.
count :: (a -> Bool) -> [a] -> Int
count :: forall a. (a -> Bool) -> [a] -> Int
count a -> Bool
p  =  [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([a] -> Int) -> ([a] -> [a]) -> [a] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
filter a -> Bool
p

-- | Nubs using a given field.
nubOn :: Eq b => (a -> b) -> [a] -> [a]
nubOn :: forall b a. Eq b => (a -> b) -> [a] -> [a]
nubOn a -> b
f  =  (a -> a -> Bool) -> [a] -> [a]
forall a. (a -> a -> Bool) -> [a] -> [a]
nubBy (b -> b -> Bool
forall a. Eq a => a -> a -> Bool
(==) (b -> b -> Bool) -> (a -> b) -> a -> a -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` a -> b
f)

-- | Equivalent to @nub . sort@ but running in /O(n log n)/.
nubSort :: Ord a => [a] -> [a]
nubSort :: forall a. Ord a => [a] -> [a]
nubSort  =  [a] -> [a]
forall {a}. Eq a => [a] -> [a]
nnub ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. Ord a => [a] -> [a]
sort
  where
  -- linear nub of adjacent values
  nnub :: [a] -> [a]
nnub [] = []
  nnub [a
x] = [a
x]
  nnub (a
x:[a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a]
nnub ((a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
==a
x) [a]
xs)

-- | Zips 'Monoid' values leaving trailing values.
--
-- > > mzip ["ab","cd"] ["ef"]
-- > ["abef","cd"]
mzip :: Monoid a => [a] -> [a] -> [a]
mzip :: forall a. Monoid a => [a] -> [a] -> [a]
mzip [] []  =  []
mzip [] [a]
ys  =  [a]
ys
mzip [a]
xs []  =  [a]
xs
mzip (a
x:[a]
xs) (a
y:[a]
ys)  =  a
x a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
y a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
forall a. Monoid a => [a] -> [a] -> [a]
mzip [a]
xs [a]
ys

-- | Group values using a given field selector.
groupOn :: Eq b => (a -> b) -> [a] -> [[a]]
groupOn :: forall b a. Eq b => (a -> b) -> [a] -> [[a]]
groupOn a -> b
f = (a -> a -> Bool) -> [a] -> [[a]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (b -> b -> Bool
forall a. Eq a => a -> a -> Bool
(==) (b -> b -> Bool) -> (a -> b) -> a -> a -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` a -> b
f)

#if __GLASGOW_HASKELL__ < 710
-- | 'sort' on a given field.
sortOn :: Ord b => (a -> b) -> [a] -> [a]
sortOn f = sortBy (compare `on` f)
#endif

-- | __WARNING:__
--   uses 'unsafePerformIO' and should only be used for debugging!
--
-- > > idIO print 10
-- > 10
-- > 10
idIO :: (a -> IO ()) -> a -> a
idIO :: forall a. (a -> IO ()) -> a -> a
idIO a -> IO ()
action a
x  =  IO a -> a
forall a. IO a -> a
unsafePerformIO (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ do
  a -> IO ()
action a
x
  a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x

-- | Applies a function to the head of a list.
mapHead :: (a -> a) -> [a] -> [a]
mapHead :: forall a. (a -> a) -> [a] -> [a]
mapHead a -> a
f (a
x:[a]
xs)  =  a -> a
f a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs

-- | Return sets of values based on the list.
--
-- The values in the list must me unique.
sets :: [a] -> [[a]]
sets :: forall a. [a] -> [[a]]
sets []  =  [[]]
sets (a
x:[a]
xs)  =  ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) [[a]]
ss [[a]] -> [[a]] -> [[a]]
forall a. [a] -> [a] -> [a]
++ [[a]]
ss
  where
  ss :: [[a]]
ss  =  [a] -> [[a]]
forall a. [a] -> [[a]]
sets [a]
xs

-- | Like 'head' but allows providing a default value.
headOr :: a -> [a] -> a
headOr :: forall a. a -> [a] -> a
headOr a
x []  =  a
x
headOr a
_ (a
x:[a]
xs)  =  a
x

-- | Lists choices of values.
choices :: [a] -> [(a,[a])]
choices :: forall a. [a] -> [(a, [a])]
choices []  =  []
choices (a
x:[a]
xs)  =  (a
x,[a]
xs) (a, [a]) -> [(a, [a])] -> [(a, [a])]
forall a. a -> [a] -> [a]
: ((a, [a]) -> (a, [a])) -> [(a, [a])] -> [(a, [a])]
forall a b. (a -> b) -> [a] -> [b]
map (([a] -> [a]) -> (a, [a]) -> (a, [a])
forall {t} {b} {a}. (t -> b) -> (a, t) -> (a, b)
second (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:)) ([a] -> [(a, [a])]
forall a. [a] -> [(a, [a])]
choices [a]
xs)
  where
  second :: (t -> b) -> (a, t) -> (a, b)
second t -> b
f (a
x,t
y)  =  (a
x,t -> b
f t
y)

-- | Lists choices of values that follow a property.
choicesThat :: (a -> [a] -> Bool) -> [a] -> [(a,[a])]
choicesThat :: forall a. (a -> [a] -> Bool) -> [a] -> [(a, [a])]
choicesThat a -> [a] -> Bool
(?)  =  ((a, [a]) -> Bool) -> [(a, [a])] -> [(a, [a])]
forall a. (a -> Bool) -> [a] -> [a]
filter ((a -> [a] -> Bool) -> (a, [a]) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> [a] -> Bool
(?)) ([(a, [a])] -> [(a, [a])])
-> ([a] -> [(a, [a])]) -> [a] -> [(a, [a])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [(a, [a])]
forall a. [a] -> [(a, [a])]
choices

-- | A variation of foldr that only uses "zero" when the list is empty
foldr0 :: (a -> a -> a) -> a -> [a] -> a
foldr0 :: forall a. (a -> a -> a) -> a -> [a] -> a
foldr0 a -> a -> a
f a
z [a]
xs | [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
xs   = a
z
              | Bool
otherwise = (a -> a -> a) -> [a] -> a
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 a -> a -> a
f [a]
xs

-- | Indents a block of text by 4 spaces
indent :: String -> String
indent :: String -> String
indent  =  Int -> String -> String
indentBy Int
4

-- | Indents a block of text with the provided amount of spaces
indentBy :: Int -> String -> String
indentBy :: Int -> String -> String
indentBy Int
n  =  [String] -> String
unlines ([String] -> String) -> (String -> [String]) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String -> String) -> [String] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Char -> String
forall a. Int -> a -> [a]
replicate Int
n Char
' ' String -> String -> String
forall a. [a] -> [a] -> [a]
++) ([String] -> [String])
-> (String -> [String]) -> String -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [String]
lines

none :: (a -> Bool) -> [a] -> Bool
none :: forall a. (a -> Bool) -> [a] -> Bool
none a -> Bool
p  =  Bool -> Bool
not (Bool -> Bool) -> ([a] -> Bool) -> [a] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any a -> Bool
p

-- | Updates the value in a list at a given position.
--
-- > > updateAt 2 (*10) [1,2,3,4]
-- > [1,2,30,4]
updateAt :: Int -> (a -> a) -> [a] -> [a]
updateAt :: forall a. Int -> (a -> a) -> [a] -> [a]
updateAt Int
_ a -> a
_ []  =  []
updateAt Int
0 a -> a
f (a
x:[a]
xs)  =  a -> a
f a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs
updateAt Int
n a -> a
f (a
x:[a]
xs)  =  a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> (a -> a) -> [a] -> [a]
forall a. Int -> (a -> a) -> [a] -> [a]
updateAt (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a -> a
f [a]
xs

-- | Applies a function to the first element of a pair.
--   Often known on the wild as @mapFst@.
--
-- > > first (*10) (1,2)
-- > (10,2)
first :: (a -> a') -> (a,b) -> (a',b)
first :: forall a a' b. (a -> a') -> (a, b) -> (a', b)
first a -> a'
f (a
x,b
y)  =  (a -> a'
f a
x, b
y)

-- | Applies a function to the second element of a pair.
--   Often known on the wild as @mapSnd@.
--
-- > > second (*100) (1,2)
-- > (1,200)
second :: (b -> b') -> (a,b) -> (a,b')
second :: forall {t} {b} {a}. (t -> b) -> (a, t) -> (a, b)
second b -> b'
f (a
x,b
y)  =  (a
x, b -> b'
f b
y)

both :: (a -> b) -> (a,a) -> (b,b)
both :: forall a b. (a -> b) -> (a, a) -> (b, b)
both a -> b
f (a
x,a
y)  =  (a -> b
f a
x, a -> b
f a
y)

nubMergeBy :: (a -> a -> Ordering) -> [a] -> [a] -> [a]
nubMergeBy :: forall a. (a -> a -> Ordering) -> [a] -> [a] -> [a]
nubMergeBy a -> a -> Ordering
cmp (a
x:[a]
xs) (a
y:[a]
ys) = case a
x a -> a -> Ordering
`cmp` a
y of
                                 Ordering
LT -> a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:(a -> a -> Ordering) -> [a] -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a] -> [a]
nubMergeBy a -> a -> Ordering
cmp [a]
xs (a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys)
                                 Ordering
GT -> a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:(a -> a -> Ordering) -> [a] -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a] -> [a]
nubMergeBy a -> a -> Ordering
cmp (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs) [a]
ys
                                 Ordering
EQ -> a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:(a -> a -> Ordering) -> [a] -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a] -> [a]
nubMergeBy a -> a -> Ordering
cmp [a]
xs [a]
ys
nubMergeBy a -> a -> Ordering
_ [a]
xs [a]
ys = [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
ys

nubMergeOn :: Ord b => (a -> b) -> [a] -> [a] -> [a]
nubMergeOn :: forall b a. Ord b => (a -> b) -> [a] -> [a] -> [a]
nubMergeOn a -> b
f = (a -> a -> Ordering) -> [a] -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a] -> [a]
nubMergeBy (b -> b -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (b -> b -> Ordering) -> (a -> b) -> a -> a -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` a -> b
f)

nubMerge :: Ord a => [a] -> [a] -> [a]
nubMerge :: forall a. Ord a => [a] -> [a] -> [a]
nubMerge = (a -> a -> Ordering) -> [a] -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a] -> [a]
nubMergeBy a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare

(+++) :: Ord a => [a] -> [a] -> [a]
+++ :: forall a. Ord a => [a] -> [a] -> [a]
(+++) = [a] -> [a] -> [a]
forall a. Ord a => [a] -> [a] -> [a]
nubMerge
infixr 5 +++