{-# LANGUAGE ScopedTypeVariables #-}

module General.Extra(
    getProcessorCount,
    findGcc,
    withResultType,
    whenLeft,
    randomElem,
    wrapQuote, showBracket,
    withs,
    maximum', maximumBy',
    fastAt,
    forkFinallyUnmasked,
    isAsyncException,
    doesFileExist_,
    removeFile_, createDirectoryRecursive,
    catchIO, tryIO, handleIO
    ) where

import Control.Exception
import Data.Char
import Data.List
import System.Environment.Extra
import System.IO.Extra
import System.IO.Unsafe
import System.Info.Extra
import System.FilePath
import System.Random
import System.Directory
import System.Exit
import Control.Concurrent
import Data.Maybe
import Data.Functor
import Data.Primitive.Array
import Control.Applicative
import Control.Monad
import Control.Monad.ST
import GHC.Conc(getNumProcessors)
import Prelude


---------------------------------------------------------------------
-- Prelude

-- See https://ghc.haskell.org/trac/ghc/ticket/10830 - they broke maximumBy
maximumBy' :: (a -> a -> Ordering) -> [a] -> a
maximumBy' :: (a -> a -> Ordering) -> [a] -> a
maximumBy' cmp :: a -> a -> Ordering
cmp = (a -> a -> a) -> [a] -> a
forall a. (a -> a -> a) -> [a] -> a
foldl1' ((a -> a -> a) -> [a] -> a) -> (a -> a -> a) -> [a] -> a
forall a b. (a -> b) -> a -> b
$ \x :: a
x y :: a
y -> if a -> a -> Ordering
cmp a
x a
y Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
GT then a
x else a
y

maximum' :: Ord a => [a] -> a
maximum' :: [a] -> a
maximum' = (a -> a -> Ordering) -> [a] -> a
forall a. (a -> a -> Ordering) -> [a] -> a
maximumBy' a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare


---------------------------------------------------------------------
-- Data.List

-- | If a string has any spaces then put quotes around and double up all internal quotes.
--   Roughly the inverse of Windows command line parsing.
wrapQuote :: String -> String
wrapQuote :: String -> String
wrapQuote xs :: String
xs | (Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Char -> Bool
isSpace String
xs = "\"" String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Char -> String) -> String -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\x :: Char
x -> if Char
x Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== '\"' then "\"\"" else [Char
x]) String
xs String -> String -> String
forall a. [a] -> [a] -> [a]
++ "\""
             | Bool
otherwise = String
xs

-- | If a string has any spaces then put brackets around it.
wrapBracket :: String -> String
wrapBracket :: String -> String
wrapBracket xs :: String
xs | (Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Char -> Bool
isSpace String
xs = "(" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
xs String -> String -> String
forall a. [a] -> [a] -> [a]
++ ")"
               | Bool
otherwise = String
xs

-- | Alias for @wrapBracket . show@.
showBracket :: Show a => a -> String
showBracket :: a -> String
showBracket = String -> String
wrapBracket (String -> String) -> (a -> String) -> a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> String
forall a. Show a => a -> String
show


-- | Version of '!!' which is fast and returns 'Nothing' if the index is not present.
fastAt :: [a] -> (Int -> Maybe a)
fastAt :: [a] -> Int -> Maybe a
fastAt xs :: [a]
xs = \i :: Int
i -> if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< 0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n then Maybe a
forall a. Maybe a
Nothing else a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ Array a -> Int -> a
forall a. Array a -> Int -> a
indexArray Array a
arr Int
i
    where
        n :: Int
n = [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
        arr :: Array a
arr = (forall s. ST s (Array a)) -> Array a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array a)) -> Array a)
-> (forall s. ST s (Array a)) -> Array a
forall a b. (a -> b) -> a -> b
$ do
            let n :: Int
n = [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
            MutableArray s a
arr <- Int -> a -> ST s (MutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
n a
forall a. HasCallStack => a
undefined
            [(Int, a)] -> ((Int, a) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [0..] [a]
xs) (((Int, a) -> ST s ()) -> ST s ())
-> ((Int, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(i :: Int
i,x :: a
x) ->
                MutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s a
MutableArray (PrimState (ST s)) a
arr Int
i a
x
            MutableArray (PrimState (ST s)) a -> ST s (Array a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
unsafeFreezeArray MutableArray s a
MutableArray (PrimState (ST s)) a
arr


---------------------------------------------------------------------
-- System.Info

{-# NOINLINE getProcessorCount #-}
getProcessorCount :: IO Int
-- unsafePefromIO so we cache the result and only compute it once
getProcessorCount :: IO Int
getProcessorCount = let res :: Int
res = IO Int -> Int
forall a. IO a -> a
unsafePerformIO IO Int
act in Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
res
    where
        act :: IO Int
act =
            if Bool
rtsSupportsBoundThreads then
                Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int) -> IO Int -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Int
getNumProcessors
            else do
                Maybe String
env <- String -> IO (Maybe String)
lookupEnv "NUMBER_OF_PROCESSORS"
                case Maybe String
env of
                    Just s :: String
s | [(i :: Int
i,"")] <- ReadS Int
forall a. Read a => ReadS a
reads String
s -> Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
i
                    _ -> do
                        String
src <- String -> IO String
readFile' "/proc/cpuinfo" IO String -> (IOException -> IO String) -> IO String
forall a. IO a -> (IOException -> IO a) -> IO a
`catchIO` \_ -> String -> IO String
forall (m :: * -> *) a. Monad m => a -> m a
return ""
                        Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$! Int -> Int -> Int
forall a. Ord a => a -> a -> a
max 1 (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ [()] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [() | String
x <- String -> [String]
lines String
src, "processor" String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` String
x]


-- Can you find a GCC executable? return a Bool, and optionally something to add to $PATH to run it
findGcc :: IO (Bool, Maybe FilePath)
findGcc :: IO (Bool, Maybe String)
findGcc = do
    Maybe String
v <- String -> IO (Maybe String)
findExecutable "gcc"
    case Maybe String
v of
        Nothing | Bool
isWindows -> do
            Maybe String
ghc <- String -> IO (Maybe String)
findExecutable "ghc"
            case Maybe String
ghc of
                Just ghc :: String
ghc -> do
                    let gcc :: String
gcc = String -> String
takeDirectory (String -> String
takeDirectory String
ghc) String -> String -> String
</> "mingw/bin/gcc.exe"
                    Bool
b <- String -> IO Bool
doesFileExist_ String
gcc
                    (Bool, Maybe String) -> IO (Bool, Maybe String)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Bool, Maybe String) -> IO (Bool, Maybe String))
-> (Bool, Maybe String) -> IO (Bool, Maybe String)
forall a b. (a -> b) -> a -> b
$ if Bool
b then (Bool
True, String -> Maybe String
forall a. a -> Maybe a
Just (String -> Maybe String) -> String -> Maybe String
forall a b. (a -> b) -> a -> b
$ String -> String
takeDirectory String
gcc) else (Bool
False, Maybe String
forall a. Maybe a
Nothing)
                _ -> (Bool, Maybe String) -> IO (Bool, Maybe String)
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool
False, Maybe String
forall a. Maybe a
Nothing)
        _ -> (Bool, Maybe String) -> IO (Bool, Maybe String)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe String -> Bool
forall a. Maybe a -> Bool
isJust Maybe String
v, Maybe String
forall a. Maybe a
Nothing)



---------------------------------------------------------------------
-- System.Random

randomElem :: [a] -> IO a
randomElem :: [a] -> IO a
randomElem xs :: [a]
xs = do
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
xs) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> IO ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "General.Extra.randomElem called with empty list, can't pick a random element"
    Int
i <- (Int, Int) -> IO Int
forall a. Random a => (a, a) -> IO a
randomRIO (0, [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> IO a) -> a -> IO a
forall a b. (a -> b) -> a -> b
$ [a]
xs [a] -> Int -> a
forall a. [a] -> Int -> a
!! Int
i


---------------------------------------------------------------------
-- Control.Monad

withs :: [(a -> r) -> r] -> ([a] -> r) -> r
withs :: [(a -> r) -> r] -> ([a] -> r) -> r
withs [] act :: [a] -> r
act = [a] -> r
act []
withs (f :: (a -> r) -> r
f:fs :: [(a -> r) -> r]
fs) act :: [a] -> r
act = (a -> r) -> r
f ((a -> r) -> r) -> (a -> r) -> r
forall a b. (a -> b) -> a -> b
$ \a :: a
a -> [(a -> r) -> r] -> ([a] -> r) -> r
forall a r. [(a -> r) -> r] -> ([a] -> r) -> r
withs [(a -> r) -> r]
fs (([a] -> r) -> r) -> ([a] -> r) -> r
forall a b. (a -> b) -> a -> b
$ \as :: [a]
as -> [a] -> r
act ([a] -> r) -> [a] -> r
forall a b. (a -> b) -> a -> b
$ a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
as


---------------------------------------------------------------------
-- Control.Concurrent

-- | Like 'forkFinally', but the inner thread is unmasked even if you started masked.
forkFinallyUnmasked :: IO a -> (Either SomeException a -> IO ()) -> IO ThreadId
forkFinallyUnmasked :: IO a -> (Either SomeException a -> IO ()) -> IO ThreadId
forkFinallyUnmasked act :: IO a
act cleanup :: Either SomeException a -> IO ()
cleanup =
    IO ThreadId -> IO ThreadId
forall a. IO a -> IO a
mask_ (IO ThreadId -> IO ThreadId) -> IO ThreadId -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forkIOWithUnmask (((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId)
-> ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ \unmask :: forall a. IO a -> IO a
unmask ->
        IO a -> IO (Either SomeException a)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO a -> IO a
forall a. IO a -> IO a
unmask IO a
act) IO (Either SomeException a)
-> (Either SomeException a -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Either SomeException a -> IO ()
cleanup


---------------------------------------------------------------------
-- Control.Exception

-- | Is the exception asynchronous, not a "coding error" that should be ignored
isAsyncException :: SomeException -> Bool
isAsyncException :: SomeException -> Bool
isAsyncException e :: SomeException
e
    | Just (AsyncException
_ :: AsyncException) <- SomeException -> Maybe AsyncException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e = Bool
True
    | Just (ExitCode
_ :: ExitCode) <- SomeException -> Maybe ExitCode
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e = Bool
True
    | Bool
otherwise = Bool
False

catchIO :: IO a -> (IOException -> IO a) -> IO a
catchIO :: IO a -> (IOException -> IO a) -> IO a
catchIO = IO a -> (IOException -> IO a) -> IO a
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
Control.Exception.catch -- GHC 7.4 has catch in the Prelude as well

tryIO :: IO a -> IO (Either IOException a)
tryIO :: IO a -> IO (Either IOException a)
tryIO = IO a -> IO (Either IOException a)
forall e a. Exception e => IO a -> IO (Either e a)
try

handleIO :: (IOException -> IO a) -> IO a -> IO a
handleIO :: (IOException -> IO a) -> IO a -> IO a
handleIO = (IO a -> (IOException -> IO a) -> IO a)
-> (IOException -> IO a) -> IO a -> IO a
forall a b c. (a -> b -> c) -> b -> a -> c
flip IO a -> (IOException -> IO a) -> IO a
forall a. IO a -> (IOException -> IO a) -> IO a
catchIO


---------------------------------------------------------------------
-- System.Directory

doesFileExist_ :: FilePath -> IO Bool
doesFileExist_ :: String -> IO Bool
doesFileExist_ x :: String
x = String -> IO Bool
doesFileExist String
x IO Bool -> (IOException -> IO Bool) -> IO Bool
forall a. IO a -> (IOException -> IO a) -> IO a
`catchIO` \_ -> Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

-- | Remove a file, but don't worry if it fails
removeFile_ :: FilePath -> IO ()
removeFile_ :: String -> IO ()
removeFile_ x :: String
x = String -> IO ()
removeFile String
x IO () -> (IOException -> IO ()) -> IO ()
forall a. IO a -> (IOException -> IO a) -> IO a
`catchIO` \_ -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Like @createDirectoryIfMissing True@ but faster, as it avoids
--   any work in the common case the directory already exists.
createDirectoryRecursive :: FilePath -> IO ()
createDirectoryRecursive :: String -> IO ()
createDirectoryRecursive dir :: String
dir = do
    Either IOException Bool
x <- IO Bool -> IO (Either IOException Bool)
forall a. IO a -> IO (Either IOException a)
tryIO (IO Bool -> IO (Either IOException Bool))
-> IO Bool -> IO (Either IOException Bool)
forall a b. (a -> b) -> a -> b
$ String -> IO Bool
doesDirectoryExist String
dir
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Either IOException Bool
x Either IOException Bool -> Either IOException Bool -> Bool
forall a. Eq a => a -> a -> Bool
/= Bool -> Either IOException Bool
forall a b. b -> Either a b
Right Bool
True) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Bool -> String -> IO ()
createDirectoryIfMissing Bool
True String
dir


---------------------------------------------------------------------
-- Data.Either

whenLeft :: Applicative m => Either a b -> (a -> m ()) -> m ()
whenLeft :: Either a b -> (a -> m ()) -> m ()
whenLeft x :: Either a b
x f :: a -> m ()
f = (a -> m ()) -> (b -> m ()) -> Either a b -> m ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either a -> m ()
f (m () -> b -> m ()
forall a b. a -> b -> a
const (m () -> b -> m ()) -> m () -> b -> m ()
forall a b. (a -> b) -> a -> b
$ () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) Either a b
x


---------------------------------------------------------------------
-- Data.Proxy

-- Should be Proxy, but that's not available in older GHC 7.6 and before
withResultType :: (Maybe a -> a) -> a
withResultType :: (Maybe a -> a) -> a
withResultType f :: Maybe a -> a
f = Maybe a -> a
f Maybe a
forall a. Maybe a
Nothing