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

-- | 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
iInt -> 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

fork :: (a -> IO b) -> a -> IO (MVar b)
fork :: (a -> IO b) -> a -> IO (MVar b)
fork a -> IO b
f a
x = do MVar b
cell <- IO (MVar b)
forall a. IO (MVar a)
newEmptyMVar
              IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (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
$ do b
result <- a -> IO b
f a
x
                                 MVar b -> b -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar b
cell b
result
              MVar b -> IO (MVar b)
forall (m :: * -> *) a. Monad m => a -> m a
return MVar b
cell

-- | Run various 'IO' actions concurrently, possibly with a bound on
-- the number of threads.
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 = [a] -> [b] -> IO [b]
go [a]
elems []
  where
    go :: [a] -> [b] -> IO [b]
go [] [b]
res = [b] -> IO [b]
forall (m :: * -> *) a. Monad m => a -> m a
return [b]
res
    go [a]
xs [b]
res = do
      Int
numThreads <- 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
      let ([a]
e,[a]
es) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
numThreads [a]
xs
      [MVar (Either SomeException b)]
mvars  <- (a -> IO (MVar (Either SomeException b)))
-> [a] -> IO [MVar (Either SomeException b)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((a -> IO (Either SomeException b))
-> a -> IO (MVar (Either SomeException b))
forall a b. (a -> IO b) -> a -> IO (MVar b)
fork a -> IO (Either SomeException b)
forall a. Exception a => a -> IO (Either a b)
f') [a]
e
      [Either SomeException b]
result <- (MVar (Either SomeException b) -> IO (Either SomeException b))
-> [MVar (Either SomeException b)] -> IO [Either SomeException b]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM MVar (Either SomeException b) -> IO (Either SomeException b)
forall a. MVar a -> IO a
takeMVar [MVar (Either SomeException b)]
mvars
      case [Either SomeException b] -> Either SomeException [b]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [Either SomeException b]
result of
        Left SomeException
err -> SomeException -> IO [b]
forall a e. Exception e => e -> a
throw (SomeException
err :: SomeException)
        Right [b]
result' -> [a] -> [b] -> IO [b]
go [a]
es ([b]
result' [b] -> [b] -> [b]
forall a. [a] -> [a] -> [a]
++ [b]
res)

    f' :: a -> IO (Either a b)
f' a
x = (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)

-- 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"