-- | @futhark script@
module Futhark.CLI.Script (main) where

import Control.Monad.Except
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Binary qualified as Bin
import Data.ByteString.Lazy.Char8 qualified as BS
import Data.Char (chr)
import Data.Text qualified as T
import Data.Text.IO qualified as T
import Futhark.CLI.Literate
  ( Options (..),
    initialOptions,
    prepareServer,
    scriptCommandLineOptions,
  )
import Futhark.Script
import Futhark.Test.Values (Compound (..), getValue, valueType)
import Futhark.Util.Options
import Futhark.Util.Pretty (prettyText)
import System.Exit
import System.IO

commandLineOptions :: [FunOptDescr Options]
commandLineOptions :: [FunOptDescr Options]
commandLineOptions =
  [FunOptDescr Options]
scriptCommandLineOptions
    [FunOptDescr Options]
-> [FunOptDescr Options] -> [FunOptDescr Options]
forall a. [a] -> [a] -> [a]
++ [ String
-> [String]
-> ArgDescr (Either (IO ()) (Options -> Options))
-> String
-> FunOptDescr Options
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
           String
"D"
           [String
"debug"]
           ( Either (IO ()) (Options -> Options)
-> ArgDescr (Either (IO ()) (Options -> Options))
forall a. a -> ArgDescr a
NoArg (Either (IO ()) (Options -> Options)
 -> ArgDescr (Either (IO ()) (Options -> Options)))
-> Either (IO ()) (Options -> Options)
-> ArgDescr (Either (IO ()) (Options -> Options))
forall a b. (a -> b) -> a -> b
$ (Options -> Options) -> Either (IO ()) (Options -> Options)
forall a b. b -> Either a b
Right ((Options -> Options) -> Either (IO ()) (Options -> Options))
-> (Options -> Options) -> Either (IO ()) (Options -> Options)
forall a b. (a -> b) -> a -> b
$ \Options
config ->
               Options
config
                 { scriptExtraOptions = "-D" : scriptExtraOptions config,
                   scriptVerbose = scriptVerbose config + 1
                 }
           )
           String
"Enable debugging.",
         String
-> [String]
-> ArgDescr (Either (IO ()) (Options -> Options))
-> String
-> FunOptDescr Options
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
           String
"L"
           [String
"log"]
           ( Either (IO ()) (Options -> Options)
-> ArgDescr (Either (IO ()) (Options -> Options))
forall a. a -> ArgDescr a
NoArg (Either (IO ()) (Options -> Options)
 -> ArgDescr (Either (IO ()) (Options -> Options)))
-> Either (IO ()) (Options -> Options)
-> ArgDescr (Either (IO ()) (Options -> Options))
forall a b. (a -> b) -> a -> b
$ (Options -> Options) -> Either (IO ()) (Options -> Options)
forall a b. b -> Either a b
Right ((Options -> Options) -> Either (IO ()) (Options -> Options))
-> (Options -> Options) -> Either (IO ()) (Options -> Options)
forall a b. (a -> b) -> a -> b
$ \Options
config ->
               Options
config
                 { scriptExtraOptions = "-L" : scriptExtraOptions config,
                   scriptVerbose = scriptVerbose config + 1
                 }
           )
           String
"Enable logging.",
         String
-> [String]
-> ArgDescr (Either (IO ()) (Options -> Options))
-> String
-> FunOptDescr Options
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
           String
"b"
           [String
"binary"]
           (Either (IO ()) (Options -> Options)
-> ArgDescr (Either (IO ()) (Options -> Options))
forall a. a -> ArgDescr a
NoArg (Either (IO ()) (Options -> Options)
 -> ArgDescr (Either (IO ()) (Options -> Options)))
-> Either (IO ()) (Options -> Options)
-> ArgDescr (Either (IO ()) (Options -> Options))
forall a b. (a -> b) -> a -> b
$ (Options -> Options) -> Either (IO ()) (Options -> Options)
forall a b. b -> Either a b
Right ((Options -> Options) -> Either (IO ()) (Options -> Options))
-> (Options -> Options) -> Either (IO ()) (Options -> Options)
forall a b. (a -> b) -> a -> b
$ \Options
config -> Options
config {scriptBinary = True})
           String
"Produce binary output.",
         String
-> [String]
-> ArgDescr (Either (IO ()) (Options -> Options))
-> String
-> FunOptDescr Options
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
           String
"f"
           [String
"file"]
           ( (String -> Either (IO ()) (Options -> Options))
-> String -> ArgDescr (Either (IO ()) (Options -> Options))
forall a. (String -> a) -> String -> ArgDescr a
ReqArg
               (\String
f -> (Options -> Options) -> Either (IO ()) (Options -> Options)
forall a b. b -> Either a b
Right ((Options -> Options) -> Either (IO ()) (Options -> Options))
-> (Options -> Options) -> Either (IO ()) (Options -> Options)
forall a b. (a -> b) -> a -> b
$ \Options
config -> Options
config {scriptExps = scriptExps config ++ [Left f]})
               String
"FILE"
           )
           String
"Run FutharkScript from this file.",
         String
-> [String]
-> ArgDescr (Either (IO ()) (Options -> Options))
-> String
-> FunOptDescr Options
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
           String
"e"
           [String
"expression"]
           ( (String -> Either (IO ()) (Options -> Options))
-> String -> ArgDescr (Either (IO ()) (Options -> Options))
forall a. (String -> a) -> String -> ArgDescr a
ReqArg
               (\String
s -> (Options -> Options) -> Either (IO ()) (Options -> Options)
forall a b. b -> Either a b
Right ((Options -> Options) -> Either (IO ()) (Options -> Options))
-> (Options -> Options) -> Either (IO ()) (Options -> Options)
forall a b. (a -> b) -> a -> b
$ \Options
config -> Options
config {scriptExps = scriptExps config ++ [Right (T.pack s)]})
               String
"EXP"
           )
           String
"Run this expression."
       ]

parseScriptFile :: FilePath -> IO Exp
parseScriptFile :: String -> IO Exp
parseScriptFile String
f = do
  Text
s <- String -> IO Text
T.readFile String
f
  case String -> Text -> Either Text Exp
parseExpFromText String
f Text
s of
    Left Text
e -> do
      Handle -> Text -> IO ()
T.hPutStrLn Handle
stderr Text
e
      IO Exp
forall a. IO a
exitFailure
    Right Exp
e -> Exp -> IO Exp
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e

getExp :: Either FilePath T.Text -> IO Exp
getExp :: Either String Text -> IO Exp
getExp (Left String
f) = String -> IO Exp
parseScriptFile String
f
getExp (Right Text
s) = case String -> Text -> Either Text Exp
parseExpFromText String
"command line option" Text
s of
  Left Text
e -> do
    Handle -> Text -> IO ()
T.hPutStrLn Handle
stderr Text
e
    IO Exp
forall a. IO a
exitFailure
  Right Exp
e -> Exp -> IO Exp
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e

-- A few extra procedures that are not handled by scriptBuiltin.
extScriptBuiltin :: (MonadError T.Text m, MonadIO m) => EvalBuiltin m
extScriptBuiltin :: forall (m :: * -> *).
(MonadError Text m, MonadIO m) =>
EvalBuiltin m
extScriptBuiltin Text
"store" [ValueAtom Value
fv, ValueAtom Value
vv]
  | Just [Word8]
path <- Value -> Maybe [Word8]
forall t. GetValue t => Value -> Maybe t
getValue Value
fv = do
      let path' :: String
path' = (Word8 -> Char) -> [Word8] -> String
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Char
chr (Int -> Char) -> (Word8 -> Int) -> Word8 -> Char
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral) ([Word8]
path :: [Bin.Word8])
      IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ String -> ByteString -> IO ()
BS.writeFile String
path' (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ Value -> ByteString
forall a. Binary a => a -> ByteString
Bin.encode Value
vv
      CompoundValue -> m CompoundValue
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CompoundValue -> m CompoundValue)
-> CompoundValue -> m CompoundValue
forall a b. (a -> b) -> a -> b
$ [CompoundValue] -> CompoundValue
forall v. [Compound v] -> Compound v
ValueTuple []
extScriptBuiltin Text
"store" [CompoundValue]
vs =
  Text -> m CompoundValue
forall a. Text -> m a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Text -> m CompoundValue) -> Text -> m CompoundValue
forall a b. (a -> b) -> a -> b
$
    Text
"$store does not accept arguments of types: "
      Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> [Text] -> Text
T.intercalate Text
", " ((CompoundValue -> Text) -> [CompoundValue] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Compound ValueType -> Text
forall a. Pretty a => a -> Text
prettyText (Compound ValueType -> Text)
-> (CompoundValue -> Compound ValueType) -> CompoundValue -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> ValueType) -> CompoundValue -> Compound ValueType
forall a b. (a -> b) -> Compound a -> Compound b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Value -> ValueType
valueType) [CompoundValue]
vs)
extScriptBuiltin Text
f [CompoundValue]
vs =
  String -> EvalBuiltin m
forall (m :: * -> *).
(MonadIO m, MonadError Text m) =>
String -> EvalBuiltin m
scriptBuiltin String
"." Text
f [CompoundValue]
vs

-- | Run @futhark script@.
main :: String -> [String] -> IO ()
main :: String -> [String] -> IO ()
main = Options
-> [FunOptDescr Options]
-> String
-> ([String] -> Options -> Maybe (IO ()))
-> String
-> [String]
-> IO ()
forall cfg.
cfg
-> [FunOptDescr cfg]
-> String
-> ([String] -> cfg -> Maybe (IO ()))
-> String
-> [String]
-> IO ()
mainWithOptions Options
initialOptions [FunOptDescr Options]
commandLineOptions String
"PROGRAM [EXP]" (([String] -> Options -> Maybe (IO ()))
 -> String -> [String] -> IO ())
-> ([String] -> Options -> Maybe (IO ()))
-> String
-> [String]
-> IO ()
forall a b. (a -> b) -> a -> b
$ \[String]
args Options
opts ->
  case [String]
args of
    [String
prog, String
script] -> IO () -> Maybe (IO ())
forall a. a -> Maybe a
Just (IO () -> Maybe (IO ())) -> IO () -> Maybe (IO ())
forall a b. (a -> b) -> a -> b
$ String -> Options -> [Either String Text] -> IO ()
main' String
prog Options
opts ([Either String Text] -> IO ()) -> [Either String Text] -> IO ()
forall a b. (a -> b) -> a -> b
$ Options -> [Either String Text]
scriptExps Options
opts [Either String Text]
-> [Either String Text] -> [Either String Text]
forall a. [a] -> [a] -> [a]
++ [Text -> Either String Text
forall a b. b -> Either a b
Right (Text -> Either String Text) -> Text -> Either String Text
forall a b. (a -> b) -> a -> b
$ String -> Text
T.pack String
script]
    [String
prog] -> IO () -> Maybe (IO ())
forall a. a -> Maybe a
Just (IO () -> Maybe (IO ())) -> IO () -> Maybe (IO ())
forall a b. (a -> b) -> a -> b
$ String -> Options -> [Either String Text] -> IO ()
main' String
prog Options
opts ([Either String Text] -> IO ()) -> [Either String Text] -> IO ()
forall a b. (a -> b) -> a -> b
$ Options -> [Either String Text]
scriptExps Options
opts
    [String]
_ -> Maybe (IO ())
forall a. Maybe a
Nothing
  where
    main' :: String -> Options -> [Either String Text] -> IO ()
main' String
prog Options
opts [Either String Text]
scripts = do
      [Exp]
scripts' <- (Either String Text -> IO Exp) -> [Either String Text] -> IO [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Either String Text -> IO Exp
getExp [Either String Text]
scripts
      String -> Options -> (ScriptServer -> IO ()) -> IO ()
forall a. String -> Options -> (ScriptServer -> IO a) -> IO a
prepareServer String
prog Options
opts ((ScriptServer -> IO ()) -> IO ())
-> (ScriptServer -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ScriptServer
s -> do
        Either Text (Maybe CompoundValue)
r <-
          ExceptT Text IO (Maybe CompoundValue)
-> IO (Either Text (Maybe CompoundValue))
forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT (ExceptT Text IO (Maybe CompoundValue)
 -> IO (Either Text (Maybe CompoundValue)))
-> ExceptT Text IO (Maybe CompoundValue)
-> IO (Either Text (Maybe CompoundValue))
forall a b. (a -> b) -> a -> b
$ do
            [ExpValue]
vs <- (Exp -> ExceptT Text IO ExpValue)
-> [Exp] -> ExceptT Text IO [ExpValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (EvalBuiltin (ExceptT Text IO)
-> ScriptServer -> Exp -> ExceptT Text IO ExpValue
forall (m :: * -> *).
(MonadError Text m, MonadIO m) =>
EvalBuiltin m -> ScriptServer -> Exp -> m ExpValue
evalExp EvalBuiltin (ExceptT Text IO)
forall (m :: * -> *).
(MonadError Text m, MonadIO m) =>
EvalBuiltin m
extScriptBuiltin ScriptServer
s) [Exp]
scripts'
            case [ExpValue] -> [ExpValue]
forall a. [a] -> [a]
reverse [ExpValue]
vs of
              [] -> Maybe CompoundValue -> ExceptT Text IO (Maybe CompoundValue)
forall a. a -> ExceptT Text IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe CompoundValue
forall a. Maybe a
Nothing
              ExpValue
v : [ExpValue]
_ -> CompoundValue -> Maybe CompoundValue
forall a. a -> Maybe a
Just (CompoundValue -> Maybe CompoundValue)
-> ExceptT Text IO CompoundValue
-> ExceptT Text IO (Maybe CompoundValue)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ScriptServer -> ExpValue -> ExceptT Text IO CompoundValue
forall (m :: * -> *).
(MonadError Text m, MonadIO m) =>
ScriptServer -> ExpValue -> m CompoundValue
getExpValue ScriptServer
s ExpValue
v
        case Either Text (Maybe CompoundValue)
r of
          Left Text
e -> do
            Handle -> Text -> IO ()
T.hPutStrLn Handle
stderr Text
e
            IO ()
forall a. IO a
exitFailure
          Right Maybe CompoundValue
Nothing ->
            () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          Right (Just CompoundValue
v) ->
            if Options -> Bool
scriptBinary Options
opts
              then case CompoundValue
v of
                ValueAtom Value
v' -> ByteString -> IO ()
BS.putStr (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ Value -> ByteString
forall a. Binary a => a -> ByteString
Bin.encode Value
v'
                CompoundValue
_ ->
                  Handle -> Text -> IO ()
T.hPutStrLn
                    Handle
stderr
                    Text
"Result value cannot be represented in binary format."
              else Text -> IO ()
T.putStrLn (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ CompoundValue -> Text
forall a. Pretty a => a -> Text
prettyText CompoundValue
v