-- | This module defines a generator for @getopt@ based command
-- line argument parsing.  Each option is associated with arbitrary
-- Python code that will perform side effects, usually by setting some
-- global variables.
module Futhark.CodeGen.Backends.GenericPython.Options
  ( Option (..),
    OptionArgument (..),
    generateOptionParser,
  )
where

import Data.Text qualified as T
import Futhark.CodeGen.Backends.GenericPython.AST

-- | Specification if a single command line option.  The option must
-- have a long name, and may also have a short name.
--
-- When the statement is being executed, the argument (if any) will be
-- stored in the variable @optarg@.
data Option = Option
  { Option -> Text
optionLongName :: T.Text,
    Option -> Maybe Char
optionShortName :: Maybe Char,
    Option -> OptionArgument
optionArgument :: OptionArgument,
    Option -> [PyStmt]
optionAction :: [PyStmt]
  }

-- | Whether an option accepts an argument.
data OptionArgument
  = NoArgument
  | RequiredArgument String
  | OptionalArgument

-- | Generate option parsing code that accepts the given command line options.  Will read from @sys.argv@.
--
-- If option parsing fails for any reason, the entire process will
-- terminate with error code 1.
generateOptionParser :: [Option] -> [PyStmt]
generateOptionParser :: [Option] -> [PyStmt]
generateOptionParser [Option]
options =
  [ PyExp -> PyExp -> PyStmt
Assign
      (String -> PyExp
Var String
"parser")
      ( PyExp -> [PyArg] -> PyExp
Call
          (String -> PyExp
Var String
"argparse.ArgumentParser")
          [ String -> PyExp -> PyArg
ArgKeyword String
"description" forall a b. (a -> b) -> a -> b
$
              Text -> PyExp
String Text
"A compiled Futhark program."
          ]
      )
  ]
    forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map Option -> PyStmt
parseOption [Option]
options
    forall a. [a] -> [a] -> [a]
++ [ PyExp -> PyExp -> PyStmt
Assign (String -> PyExp
Var String
"parser_result") forall a b. (a -> b) -> a -> b
$
           PyExp -> [PyArg] -> PyExp
Call (String -> PyExp
Var String
"vars") [PyExp -> PyArg
Arg forall a b. (a -> b) -> a -> b
$ PyExp -> [PyArg] -> PyExp
Call (String -> PyExp
Var String
"parser.parse_args") [PyExp -> PyArg
Arg forall a b. (a -> b) -> a -> b
$ String -> PyExp
Var String
"sys.argv[1:]"]]
       ]
    forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map Option -> PyStmt
executeOption [Option]
options
  where
    parseOption :: Option -> PyStmt
parseOption Option
option =
      PyExp -> PyStmt
Exp forall a b. (a -> b) -> a -> b
$
        PyExp -> [PyArg] -> PyExp
Call (String -> PyExp
Var String
"parser.add_argument") forall a b. (a -> b) -> a -> b
$
          forall a b. (a -> b) -> [a] -> [b]
map (PyExp -> PyArg
Arg forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> PyExp
String) [Text]
name_args forall a. [a] -> [a] -> [a]
++ [PyArg]
argument_args
      where
        name_args :: [Text]
name_args =
          forall b a. b -> (a -> b) -> Maybe a -> b
maybe
            forall a. a -> a
id
            (\Char
x [Text]
l -> (Text
"-" forall a. Semigroup a => a -> a -> a
<> Char -> Text
T.singleton Char
x) forall a. a -> [a] -> [a]
: [Text]
l)
            (Option -> Maybe Char
optionShortName Option
option)
            [Text
"--" forall a. Semigroup a => a -> a -> a
<> Option -> Text
optionLongName Option
option]
        argument_args :: [PyArg]
argument_args = case Option -> OptionArgument
optionArgument Option
option of
          RequiredArgument String
t ->
            [ String -> PyExp -> PyArg
ArgKeyword String
"action" (Text -> PyExp
String Text
"append"),
              String -> PyExp -> PyArg
ArgKeyword String
"default" forall a b. (a -> b) -> a -> b
$ [PyExp] -> PyExp
List [],
              String -> PyExp -> PyArg
ArgKeyword String
"type" forall a b. (a -> b) -> a -> b
$ String -> PyExp
Var String
t
            ]
          OptionArgument
NoArgument ->
            [ String -> PyExp -> PyArg
ArgKeyword String
"action" (Text -> PyExp
String Text
"append_const"),
              String -> PyExp -> PyArg
ArgKeyword String
"default" forall a b. (a -> b) -> a -> b
$ [PyExp] -> PyExp
List [],
              String -> PyExp -> PyArg
ArgKeyword String
"const" PyExp
None
            ]
          OptionArgument
OptionalArgument ->
            [ String -> PyExp -> PyArg
ArgKeyword String
"action" (Text -> PyExp
String Text
"append"),
              String -> PyExp -> PyArg
ArgKeyword String
"default" forall a b. (a -> b) -> a -> b
$ [PyExp] -> PyExp
List [],
              String -> PyExp -> PyArg
ArgKeyword String
"nargs" forall a b. (a -> b) -> a -> b
$ Text -> PyExp
String Text
"?"
            ]

    executeOption :: Option -> PyStmt
executeOption Option
option =
      String -> PyExp -> [PyStmt] -> PyStmt
For
        String
"optarg"
        ( PyExp -> PyIdx -> PyExp
Index (String -> PyExp
Var String
"parser_result") forall a b. (a -> b) -> a -> b
$ PyExp -> PyIdx
IdxExp forall a b. (a -> b) -> a -> b
$ Text -> PyExp
String forall a b. (a -> b) -> a -> b
$ Option -> Text
fieldName Option
option
        )
        forall a b. (a -> b) -> a -> b
$ Option -> [PyStmt]
optionAction Option
option

    fieldName :: Option -> Text
fieldName = (Char -> Char) -> Text -> Text
T.map Char -> Char
escape forall b c a. (b -> c) -> (a -> b) -> a -> c
. Option -> Text
optionLongName
      where
        escape :: Char -> Char
escape Char
'-' = Char
'_'
        escape Char
c = Char
c