{-# LANGUAGE Trustworthy #-}

-- | Non-Futhark-specific utilities.  If you find yourself writing
-- general functions on generic data structures, consider putting them
-- here.
--
-- Sometimes it is also preferable to copy a small function rather
-- than introducing a large dependency.  In this case, make sure to
-- note where you got it from (and make sure that the license is
-- compatible).
module Futhark.Util
  ( mapAccumLM,
    maxinum,
    chunk,
    chunks,
    dropAt,
    takeLast,
    dropLast,
    mapEither,
    maybeNth,
    maybeHead,
    splitFromEnd,
    splitAt3,
    focusNth,
    unixEnvironment,
    isEnvVarSet,
    fancyTerminal,
    runProgramWithExitCode,
    directoryContents,
    roundFloat,
    ceilFloat,
    floorFloat,
    roundDouble,
    ceilDouble,
    floorDouble,
    lgamma,
    lgammaf,
    tgamma,
    tgammaf,
    fromPOSIX,
    toPOSIX,
    trim,
    pmapIO,
    UserString,
    EncodedString,
    zEncodeString,
  )
where

import Control.Concurrent
import Control.Exception
import Control.Monad
import qualified Data.ByteString as BS
import Data.Char
import Data.Either
import Data.List (foldl', genericDrop, genericSplitAt)
import Data.Maybe
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.Text.Encoding.Error as T
import Numeric
import qualified System.Directory.Tree as Dir
import System.Environment
import System.Exit
import qualified System.FilePath as Native
import qualified System.FilePath.Posix as Posix
import System.IO (hIsTerminalDevice, stdout)
import System.IO.Unsafe
import System.Process.ByteString

-- | Like 'Data.Traversable.mapAccumL', but monadic.
mapAccumLM ::
  Monad m =>
  (acc -> x -> m (acc, y)) ->
  acc ->
  [x] ->
  m (acc, [y])
mapAccumLM :: (acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
mapAccumLM acc -> x -> m (acc, y)
_ acc
acc [] = (acc, [y]) -> m (acc, [y])
forall (m :: * -> *) a. Monad m => a -> m a
return (acc
acc, [])
mapAccumLM acc -> x -> m (acc, y)
f acc
acc (x
x : [x]
xs) = do
  (acc
acc', y
x') <- acc -> x -> m (acc, y)
f acc
acc x
x
  (acc
acc'', [y]
xs') <- (acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
forall (m :: * -> *) acc x y.
Monad m =>
(acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
mapAccumLM acc -> x -> m (acc, y)
f acc
acc' [x]
xs
  (acc, [y]) -> m (acc, [y])
forall (m :: * -> *) a. Monad m => a -> m a
return (acc
acc'', y
x' y -> [y] -> [y]
forall a. a -> [a] -> [a]
: [y]
xs')

-- | @chunk n a@ splits @a@ into @n@-size-chunks.  If the length of
-- @a@ is not divisible by @n@, the last chunk will have fewer than
-- @n@ elements (but it will never be empty).
chunk :: Int -> [a] -> [[a]]
chunk :: Int -> [a] -> [[a]]
chunk Int
_ [] = []
chunk Int
n [a]
xs =
  let ([a]
bef, [a]
aft) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [a]
xs
   in [a]
bef [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
chunk Int
n [a]
aft

-- | @chunks ns a@ splits @a@ into chunks determined by the elements
-- of @ns@.  It must hold that @sum ns == length a@, or the resulting
-- list may contain too few chunks, or not all elements of @a@.
chunks :: [Int] -> [a] -> [[a]]
chunks :: [Int] -> [a] -> [[a]]
chunks [] [a]
_ = []
chunks (Int
n : [Int]
ns) [a]
xs =
  let ([a]
bef, [a]
aft) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [a]
xs
   in [a]
bef [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [Int] -> [a] -> [[a]]
forall a. [Int] -> [a] -> [[a]]
chunks [Int]
ns [a]
aft

-- | Like 'maximum', but returns zero for an empty list.
maxinum :: (Num a, Ord a, Foldable f) => f a -> a
maxinum :: f a -> a
maxinum = (a -> a -> a) -> a -> f a -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' a -> a -> a
forall a. Ord a => a -> a -> a
max a
0

-- | @dropAt i n@ drops @n@ elements starting at element @i@.
dropAt :: Int -> Int -> [a] -> [a]
dropAt :: Int -> Int -> [a] -> [a]
dropAt Int
i Int
n [a]
xs = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
i [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) [a]
xs

-- | @takeLast n l@ takes the last @n@ elements of @l@.
takeLast :: Int -> [a] -> [a]
takeLast :: Int -> [a] -> [a]
takeLast Int
n = [a] -> [a]
forall a. [a] -> [a]
reverse ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
n ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. [a] -> [a]
reverse

-- | @dropLast n l@ drops the last @n@ elements of @l@.
dropLast :: Int -> [a] -> [a]
dropLast :: Int -> [a] -> [a]
dropLast Int
n = [a] -> [a]
forall a. [a] -> [a]
reverse ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
n ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. [a] -> [a]
reverse

-- | A combination of 'map' and 'partitionEithers'.
mapEither :: (a -> Either b c) -> [a] -> ([b], [c])
mapEither :: (a -> Either b c) -> [a] -> ([b], [c])
mapEither a -> Either b c
f [a]
l = [Either b c] -> ([b], [c])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ([Either b c] -> ([b], [c])) -> [Either b c] -> ([b], [c])
forall a b. (a -> b) -> a -> b
$ (a -> Either b c) -> [a] -> [Either b c]
forall a b. (a -> b) -> [a] -> [b]
map a -> Either b c
f [a]
l

-- | Return the list element at the given index, if the index is valid.
maybeNth :: Integral int => int -> [a] -> Maybe a
maybeNth :: int -> [a] -> Maybe a
maybeNth int
i [a]
l
  | int
i int -> int -> Bool
forall a. Ord a => a -> a -> Bool
>= int
0, a
v : [a]
_ <- int -> [a] -> [a]
forall i a. Integral i => i -> [a] -> [a]
genericDrop int
i [a]
l = a -> Maybe a
forall a. a -> Maybe a
Just a
v
  | Bool
otherwise = Maybe a
forall a. Maybe a
Nothing

-- | Return the first element of the list, if it exists.
maybeHead :: [a] -> Maybe a
maybeHead :: [a] -> Maybe a
maybeHead [] = Maybe a
forall a. Maybe a
Nothing
maybeHead (a
x : [a]
_) = a -> Maybe a
forall a. a -> Maybe a
Just a
x

-- | Like 'splitAt', but from the end.
splitFromEnd :: Int -> [a] -> ([a], [a])
splitFromEnd :: Int -> [a] -> ([a], [a])
splitFromEnd Int
i [a]
l = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) [a]
l

-- | Like 'splitAt', but produces three lists.
splitAt3 :: Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 :: Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 Int
n Int
m [a]
l =
  let ([a]
xs, [a]
l') = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [a]
l
      ([a]
ys, [a]
zs) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
m [a]
l'
   in ([a]
xs, [a]
ys, [a]
zs)

-- | Return the list element at the given index, if the index is
-- valid, along with the elements before and after.
focusNth :: Integral int => int -> [a] -> Maybe ([a], a, [a])
focusNth :: int -> [a] -> Maybe ([a], a, [a])
focusNth int
i [a]
xs
  | ([a]
bef, a
x : [a]
aft) <- int -> [a] -> ([a], [a])
forall i a. Integral i => i -> [a] -> ([a], [a])
genericSplitAt int
i [a]
xs = ([a], a, [a]) -> Maybe ([a], a, [a])
forall a. a -> Maybe a
Just ([a]
bef, a
x, [a]
aft)
  | Bool
otherwise = Maybe ([a], a, [a])
forall a. Maybe a
Nothing

{-# NOINLINE unixEnvironment #-}

-- | The Unix environment when the Futhark compiler started.
unixEnvironment :: [(String, String)]
unixEnvironment :: [(String, String)]
unixEnvironment = IO [(String, String)] -> [(String, String)]
forall a. IO a -> a
unsafePerformIO IO [(String, String)]
getEnvironment

-- | Is an environment variable set to 0 or 1?  If 0, return False; if
-- 1, True; otherwise the default value.
isEnvVarSet :: String -> Bool -> Bool
isEnvVarSet :: String -> Bool -> Bool
isEnvVarSet String
name Bool
default_val = Bool -> Maybe Bool -> Bool
forall a. a -> Maybe a -> a
fromMaybe Bool
default_val (Maybe Bool -> Bool) -> Maybe Bool -> Bool
forall a b. (a -> b) -> a -> b
$ do
  String
val <- String -> [(String, String)] -> Maybe String
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup String
name [(String, String)]
unixEnvironment
  case String
val of
    String
"0" -> Bool -> Maybe Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
    String
"1" -> Bool -> Maybe Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
    String
_ -> Maybe Bool
forall a. Maybe a
Nothing

{-# NOINLINE fancyTerminal #-}

-- | Are we running in a terminal capable of fancy commands and
-- visualisation?
fancyTerminal :: Bool
fancyTerminal :: Bool
fancyTerminal = IO Bool -> Bool
forall a. IO a -> a
unsafePerformIO (IO Bool -> Bool) -> IO Bool -> Bool
forall a b. (a -> b) -> a -> b
$ do
  Bool
isTTY <- Handle -> IO Bool
hIsTerminalDevice Handle
stdout
  Bool
isDumb <- (String -> Maybe String
forall a. a -> Maybe a
Just String
"dumb" Maybe String -> Maybe String -> Bool
forall a. Eq a => a -> a -> Bool
==) (Maybe String -> Bool) -> IO (Maybe String) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO (Maybe String)
lookupEnv String
"TERM"
  Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ Bool
isTTY Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
isDumb

-- | Like 'readProcessWithExitCode', but also wraps exceptions when
-- the indicated binary cannot be launched, or some other exception is
-- thrown.  Also does shenanigans to handle improperly encoded outputs.
runProgramWithExitCode ::
  FilePath ->
  [String] ->
  BS.ByteString ->
  IO (Either IOException (ExitCode, String, String))
runProgramWithExitCode :: String
-> [String]
-> ByteString
-> IO (Either IOException (ExitCode, String, String))
runProgramWithExitCode String
exe [String]
args ByteString
inp =
  ((ExitCode, String, String)
-> Either IOException (ExitCode, String, String)
forall a b. b -> Either a b
Right ((ExitCode, String, String)
 -> Either IOException (ExitCode, String, String))
-> ((ExitCode, ByteString, ByteString)
    -> (ExitCode, String, String))
-> (ExitCode, ByteString, ByteString)
-> Either IOException (ExitCode, String, String)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ExitCode, ByteString, ByteString) -> (ExitCode, String, String)
forall a. (a, ByteString, ByteString) -> (a, String, String)
postprocess ((ExitCode, ByteString, ByteString)
 -> Either IOException (ExitCode, String, String))
-> IO (ExitCode, ByteString, ByteString)
-> IO (Either IOException (ExitCode, String, String))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String
-> [String] -> ByteString -> IO (ExitCode, ByteString, ByteString)
readProcessWithExitCode String
exe [String]
args ByteString
inp)
    IO (Either IOException (ExitCode, String, String))
-> (IOException
    -> IO (Either IOException (ExitCode, String, String)))
-> IO (Either IOException (ExitCode, String, String))
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \IOException
e -> Either IOException (ExitCode, String, String)
-> IO (Either IOException (ExitCode, String, String))
forall (m :: * -> *) a. Monad m => a -> m a
return (IOException -> Either IOException (ExitCode, String, String)
forall a b. a -> Either a b
Left IOException
e)
  where
    decode :: ByteString -> String
decode = Text -> String
T.unpack (Text -> String) -> (ByteString -> Text) -> ByteString -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OnDecodeError -> ByteString -> Text
T.decodeUtf8With OnDecodeError
T.lenientDecode
    postprocess :: (a, ByteString, ByteString) -> (a, String, String)
postprocess (a
code, ByteString
out, ByteString
err) =
      (a
code, ByteString -> String
decode ByteString
out, ByteString -> String
decode ByteString
err)

-- | Every non-directory file contained in a directory tree.
directoryContents :: FilePath -> IO [FilePath]
directoryContents :: String -> IO [String]
directoryContents String
dir = do
  String
_ Dir.:/ DirTree String
tree <- (String -> IO String) -> String -> IO (AnchoredDirTree String)
forall a. (String -> IO a) -> String -> IO (AnchoredDirTree a)
Dir.readDirectoryWith String -> IO String
forall (m :: * -> *) a. Monad m => a -> m a
return String
dir
  case DirTree String -> [DirTree String]
forall a. DirTree a -> [DirTree a]
Dir.failures DirTree String
tree of
    Dir.Failed String
_ IOException
err : [DirTree String]
_ -> IOException -> IO [String]
forall a e. Exception e => e -> a
throw IOException
err
    [DirTree String]
_ -> [String] -> IO [String]
forall (m :: * -> *) a. Monad m => a -> m a
return ([String] -> IO [String]) -> [String] -> IO [String]
forall a b. (a -> b) -> a -> b
$ (DirTree String -> Maybe String) -> [DirTree String] -> [String]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe DirTree String -> Maybe String
forall a. DirTree a -> Maybe a
isFile ([DirTree String] -> [String]) -> [DirTree String] -> [String]
forall a b. (a -> b) -> a -> b
$ DirTree String -> [DirTree String]
forall a. DirTree a -> [DirTree a]
Dir.flattenDir DirTree String
tree
  where
    isFile :: DirTree a -> Maybe a
isFile (Dir.File String
_ a
path) = a -> Maybe a
forall a. a -> Maybe a
Just a
path
    isFile DirTree a
_ = Maybe a
forall a. Maybe a
Nothing

foreign import ccall "nearbyint" c_nearbyint :: Double -> Double

foreign import ccall "nearbyintf" c_nearbyintf :: Float -> Float

foreign import ccall "ceil" c_ceil :: Double -> Double

foreign import ccall "ceilf" c_ceilf :: Float -> Float

foreign import ccall "floor" c_floor :: Double -> Double

foreign import ccall "floorf" c_floorf :: Float -> Float

-- | Round a single-precision floating point number correctly.
roundFloat :: Float -> Float
roundFloat :: Float -> Float
roundFloat = Float -> Float
c_nearbyintf

-- | Round a single-precision floating point number upwards correctly.
ceilFloat :: Float -> Float
ceilFloat :: Float -> Float
ceilFloat = Float -> Float
c_ceilf

-- | Round a single-precision floating point number downwards correctly.
floorFloat :: Float -> Float
floorFloat :: Float -> Float
floorFloat = Float -> Float
c_floorf

-- | Round a double-precision floating point number correctly.
roundDouble :: Double -> Double
roundDouble :: Double -> Double
roundDouble = Double -> Double
c_nearbyint

-- | Round a double-precision floating point number upwards correctly.
ceilDouble :: Double -> Double
ceilDouble :: Double -> Double
ceilDouble = Double -> Double
c_ceil

-- | Round a double-precision floating point number downwards correctly.
floorDouble :: Double -> Double
floorDouble :: Double -> Double
floorDouble = Double -> Double
c_floor

foreign import ccall "lgamma" c_lgamma :: Double -> Double

foreign import ccall "lgammaf" c_lgammaf :: Float -> Float

foreign import ccall "tgamma" c_tgamma :: Double -> Double

foreign import ccall "tgammaf" c_tgammaf :: Float -> Float

-- | The system-level @lgamma()@ function.
lgamma :: Double -> Double
lgamma :: Double -> Double
lgamma = Double -> Double
c_lgamma

-- | The system-level @lgammaf()@ function.
lgammaf :: Float -> Float
lgammaf :: Float -> Float
lgammaf = Float -> Float
c_lgammaf

-- | The system-level @tgamma()@ function.
tgamma :: Double -> Double
tgamma :: Double -> Double
tgamma = Double -> Double
c_tgamma

-- | The system-level @tgammaf()@ function.
tgammaf :: Float -> Float
tgammaf :: Float -> Float
tgammaf = Float -> Float
c_tgammaf

-- | Turn a POSIX filepath into a filepath for the native system.
toPOSIX :: Native.FilePath -> Posix.FilePath
toPOSIX :: String -> String
toPOSIX = [String] -> String
Posix.joinPath ([String] -> String) -> (String -> [String]) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [String]
Native.splitDirectories

-- | Some bad operating systems do not use forward slash as
-- directory separator - this is where we convert Futhark includes
-- (which always use forward slash) to native paths.
fromPOSIX :: Posix.FilePath -> Native.FilePath
fromPOSIX :: String -> String
fromPOSIX = [String] -> String
Native.joinPath ([String] -> String) -> (String -> [String]) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [String]
Posix.splitDirectories

-- | Remove leading and trailing whitespace from a string.  Not an
-- efficient implementation!
trim :: String -> String
trim :: String -> String
trim = String -> String
forall a. [a] -> [a]
reverse (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
forall a. [a] -> [a]
reverse (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace

-- | Run various 'IO' actions concurrently, possibly with a bound on
-- the number of threads.  The list must be finite.  The ordering of
-- the result list is not deterministic - add your own sorting if
-- needed.  If any of the actions throw an exception, then that
-- exception is propagated to this function.
pmapIO :: Maybe Int -> (a -> IO b) -> [a] -> IO [b]
pmapIO :: Maybe Int -> (a -> IO b) -> [a] -> IO [b]
pmapIO Maybe Int
concurrency a -> IO b
f [a]
elems = do
  MVar [a]
tasks <- [a] -> IO (MVar [a])
forall a. a -> IO (MVar a)
newMVar [a]
elems
  MVar (Either SomeException b)
results <- IO (MVar (Either SomeException b))
forall a. IO (MVar a)
newEmptyMVar
  Int
num_threads <- IO Int -> (Int -> IO Int) -> Maybe Int -> IO Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO Int
getNumCapabilities Int -> IO Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
concurrency
  Int -> IO ThreadId -> IO ()
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ Int
num_threads (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ MVar [a] -> MVar (Either SomeException b) -> IO ()
forall a. Exception a => MVar [a] -> MVar (Either a b) -> IO ()
worker MVar [a]
tasks MVar (Either SomeException b)
results
  Int -> IO b -> IO [b]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
elems) (IO b -> IO [b]) -> IO b -> IO [b]
forall a b. (a -> b) -> a -> b
$ MVar (Either SomeException b) -> IO b
forall b. MVar (Either SomeException b) -> IO b
getResult MVar (Either SomeException b)
results
  where
    worker :: MVar [a] -> MVar (Either a b) -> IO ()
worker MVar [a]
tasks MVar (Either a b)
results = do
      Maybe a
task <- MVar [a] -> ([a] -> IO ([a], Maybe a)) -> IO (Maybe a)
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar [a]
tasks [a] -> IO ([a], Maybe a)
forall (f :: * -> *) a. Applicative f => [a] -> f ([a], Maybe a)
getTask
      case Maybe a
task of
        Maybe a
Nothing -> () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        Just a
x -> do
          Either a b
y <- (b -> Either a b
forall a b. b -> Either a b
Right (b -> Either a b) -> IO b -> IO (Either a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> IO b
f a
x) IO (Either a b) -> (a -> IO (Either a b)) -> IO (Either a b)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` (Either a b -> IO (Either a b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either a b -> IO (Either a b))
-> (a -> Either a b) -> a -> IO (Either a b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Either a b
forall a b. a -> Either a b
Left)
          MVar (Either a b) -> Either a b -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (Either a b)
results Either a b
y
          MVar [a] -> MVar (Either a b) -> IO ()
worker MVar [a]
tasks MVar (Either a b)
results

    getTask :: [a] -> f ([a], Maybe a)
getTask [] = ([a], Maybe a) -> f ([a], Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Maybe a
forall a. Maybe a
Nothing)
    getTask (a
task : [a]
tasks) = ([a], Maybe a) -> f ([a], Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([a]
tasks, a -> Maybe a
forall a. a -> Maybe a
Just a
task)

    getResult :: MVar (Either SomeException b) -> IO b
getResult MVar (Either SomeException b)
results = do
      Either SomeException b
res <- MVar (Either SomeException b) -> IO (Either SomeException b)
forall a. MVar a -> IO a
takeMVar MVar (Either SomeException b)
results
      case Either SomeException b
res of
        Left SomeException
err -> SomeException -> IO b
forall a e. Exception e => e -> a
throw (SomeException
err :: SomeException)
        Right b
v -> b -> IO b
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
v

-- Z-encoding from https://ghc.haskell.org/trac/ghc/wiki/Commentary/Compiler/SymbolNames
--
-- Slightly simplified as we do not need it to deal with tuples and
-- the like.
--
-- (c) The University of Glasgow, 1997-2006

-- | As the user typed it.
type UserString = String

-- | Encoded form.
type EncodedString = String

-- | Z-encode a string using a slightly simplified variant of GHC
-- Z-encoding.  The encoded string is a valid identifier in most
-- programming languages.
zEncodeString :: UserString -> EncodedString
zEncodeString :: String -> String
zEncodeString String
"" = String
""
zEncodeString (Char
c : String
cs) = Char -> String
encodeDigitChar Char
c String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Char -> String) -> String -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Char -> String
encodeChar String
cs

unencodedChar :: Char -> Bool -- True for chars that don't need encoding
unencodedChar :: Char -> Bool
unencodedChar Char
'Z' = Bool
False
unencodedChar Char
'z' = Bool
False
unencodedChar Char
'_' = Bool
True
unencodedChar Char
c =
  Char -> Bool
isAsciiLower Char
c
    Bool -> Bool -> Bool
|| Char -> Bool
isAsciiUpper Char
c
    Bool -> Bool -> Bool
|| Char -> Bool
isDigit Char
c

-- If a digit is at the start of a symbol then we need to encode it.
-- Otherwise names like 9pH-0.1 give linker errors.
encodeDigitChar :: Char -> EncodedString
encodeDigitChar :: Char -> String
encodeDigitChar Char
c
  | Char -> Bool
isDigit Char
c = Char -> String
encodeAsUnicodeCharar Char
c
  | Bool
otherwise = Char -> String
encodeChar Char
c

encodeChar :: Char -> EncodedString
encodeChar :: Char -> String
encodeChar Char
c | Char -> Bool
unencodedChar Char
c = [Char
c] -- Common case first

-- Constructors
encodeChar Char
'(' = String
"ZL" -- Needed for things like (,), and (->)
encodeChar Char
')' = String
"ZR" -- For symmetry with (
encodeChar Char
'[' = String
"ZM"
encodeChar Char
']' = String
"ZN"
encodeChar Char
':' = String
"ZC"
encodeChar Char
'Z' = String
"ZZ"
-- Variables
encodeChar Char
'z' = String
"zz"
encodeChar Char
'&' = String
"za"
encodeChar Char
'|' = String
"zb"
encodeChar Char
'^' = String
"zc"
encodeChar Char
'$' = String
"zd"
encodeChar Char
'=' = String
"ze"
encodeChar Char
'>' = String
"zg"
encodeChar Char
'#' = String
"zh"
encodeChar Char
'.' = String
"zi"
encodeChar Char
'<' = String
"zl"
encodeChar Char
'-' = String
"zm"
encodeChar Char
'!' = String
"zn"
encodeChar Char
'+' = String
"zp"
encodeChar Char
'\'' = String
"zq"
encodeChar Char
'\\' = String
"zr"
encodeChar Char
'/' = String
"zs"
encodeChar Char
'*' = String
"zt"
encodeChar Char
'_' = String
"zu"
encodeChar Char
'%' = String
"zv"
encodeChar Char
c = Char -> String
encodeAsUnicodeCharar Char
c

encodeAsUnicodeCharar :: Char -> EncodedString
encodeAsUnicodeCharar :: Char -> String
encodeAsUnicodeCharar Char
c =
  Char
'z' Char -> String -> String
forall a. a -> [a] -> [a]
:
  if Char -> Bool
isDigit (String -> Char
forall a. [a] -> a
head String
hex_str)
    then String
hex_str
    else Char
'0' Char -> String -> String
forall a. a -> [a] -> [a]
: String
hex_str
  where
    hex_str :: String
hex_str = Int -> String -> String
forall a. (Integral a, Show a) => a -> String -> String
showHex (Char -> Int
ord Char
c) String
"U"