Record.hs

{-# LANGUAGE ScopedTypeVariables, TypeFamilies, MultiParamTypeClasses, DeriveDataTypeable,
             Rank2Types, FlexibleContexts, UndecidableInstances #-}
module System.Console.CmdLib.Record where
import System.Console.CmdLib.Attribute
import System.Console.CmdLib.Command
import System.Console.GetOpt
import Data.Data
import Data.IORef
import Data.List( nub, elemIndex )
import Control.Exception ( throw, evaluate )
import Control.Monad ( when )
import Control.Monad.State( evalState, get, put, State, execState )
import Data.Maybe( fromMaybe, fromJust )
import System.Exit

import System.Console.CmdLib.Flag

append :: forall a. (Data a) => a -> a -> a
append a = any <+< list
  where any = const a
        list :: (Typeable b) => b -> [String]
        list = (++(fromJust $ cast a)) . fromJust . cast

-- | This wrapper type allows use of record types (single or multi-constructor)
-- for handling flags. Each field of the record is made into a single flag of
-- the corresponding type. The record needs to be made an instance of the
-- "Attributes" class. That way, attributes can be attached to the field
-- selectors, although when used with RecordCommand, its "rec_options" method
-- can be used as well and the Attributes instance left empty.
--
-- > data Flags = Flags { wibblify :: Int, simplify :: Bool }
-- > instance Attributes Flags where
-- >     attributes _ =
-- >        wibblify %> Help "Add a wibblification pass." %+ ArgHelp "intensity" %%
-- >        simplify %> Help "Enable a two-pass simplifier."
--
-- A single value of the Flags type will then be passed to the "Command"
-- instances (those that use @Record Flags@ as their second type parameter),
-- containing the value of the rightmost occurence for each of the flags.
--
-- TODO: List-based option types should be accumulated instead of overriden.
data (Attributes rec) => Record rec = Record rec String deriving Eq

instance (Eq rec, Eq (Record rec), Data rec, Attributes rec) => FlagType (Record rec) where
  type Folded (Record rec) = rec

  flag_attrkey (Record _ f) = KeyF (typeOf (undefined :: rec)) f

  flag_list = [ Record undefined field |
                field <- nub $ concatMap constrFields
                         (dataTypeConstrs $ dataTypeOf (undefined :: rec)) ]

  flag_type f@(Record rec field) attr =
      fixup $ optionType [gmapQi idx dataTypeOf x] [gmapQi idx typeOf x]
    where constr = head [ ctor | ctor <- dataTypeConstrs $ dataTypeOf rec
                               , field `elem` constrFields ctor ]
          idx = fromMaybe (error $ "BUG: Getting type of nonexistent field " ++ field) $
                    elemIndex field $ constrFields constr
          x :: rec = fromConstr constr
          fixup BooleanOption | False <- invertible attr = SimpleOption
          fixup x = x

  flag_defaults attr = map setfield flag_list
    where setdef :: forall a. (Data a) => Record rec -> a
          setdef flag = defvalue $ attrs flag attr
          setfield flag@(Record _ field) rec =
            setField field rec (if required $ attr flag
                                  then throw RequiredArgException
                                  else setdef flag)

  flag_attrs (Record _ name) = (attrFun $ attributes (undefined :: rec))
                               (KeyF (typeOf (undefined :: rec)) name) %+
                               long (hyphenate name) %+ defaults

  flag_value (Record _ field) = getField field
  flag_eval (Record _ field) = evalField field

  flag_set (Record _ field) v = \x -> setField field x (errcast v)
     where errcast :: (Typeable a, Typeable b) => a -> b
           errcast x = fromMaybe
             (error "BUG: flag_set in Record used with wrong value type")
             (cast x)

  flag_parse f@(Record _ field) str v = setField field v (readFlag (undefined :: rec) str)

  flag_args f@(Record flag field) attr = case flag_type (Record flag field) attr of
    BooleanOption -> OptArg setoptional ""
    OptionalArgument -> OptArg setoptional ""
    RequiredArgument -> ReqArg setlist ""
    SimpleOption -> NoArg $ flag_set f True
   where set str v = setField field v (readFlag (undefined :: rec) str)
         setlist str v = setField field v (append (readFlag (undefined :: rec) str) (flag_value f v))
         setoptional str = set (fromMaybe "" str)

  flag_empty _ = fromConstr $ head $ dataTypeConstrs $ dataTypeOf (undefined :: rec)

-- | A bridge that allows multi-constructor record types to be used as a
-- description of a command set. In such a type, each constructor corresponds
-- to a single command and its fields to its options. To describe a program
-- with two commands, @foo@ and @bar@, each taking a @--wibble@ boolean option
-- and @bar@ also taking a @--text=<string>@ option, you can write:
--
-- > data Commands = Foo { wibble :: Bool }
-- >               | Bar { wibble :: Bool, text :: String }
-- >
-- > instance RecordCommand Commands where (...)
--
-- You should at least implement @run'@, @rec_options@ and @mode_summary@ are optional.

class (Data cmd) => RecordCommand cmd where
  -- | @run'@ is your entrypoint into the whole set of commands. You can
  -- dispatch on the command by looking at the constructor in @cmd@:
  --
  -- > run' cmd@(Foo {}) _ = putStrLn $ "Foo running. Wibble = " ++ show (wibble cmd)
  -- > run' cmd@(Bar {}) _ = putStrLn "This is bar."
  run' :: cmd -> [String] -> IO ()

  -- | You can also provide extra per-command flag attributes (match on the
  -- constructor like with @run'@). The attributes shared by various commands
  -- can be set in "rec_attrs" in "Attributes" instead.
  rec_options :: cmd -> AttributeMap Key
  rec_options _ = EmptyMap

  -- | Set the per-command option style, useful for supercommands to pass their
  -- options through to another dispatch, by using NoOptions.
  rec_optionStyle :: cmd -> OptionStyle
  rec_optionStyle _ = Permuted

  -- | Pattern match like in @run'@ to identify any supercommands, which will
  -- allow --help flags to be passed through to the sub-commands.
  rec_superCommand :: cmd -> Bool
  rec_superCommand _ = False

  -- | Provide a summary help string for each mode. Used in help output. Again,
  -- pattern match like in @run'@.
  mode_summary :: cmd -> String
  mode_summary _ = ""

  -- | Provide a help blurb for each mode. Use patterns like in @run'@.
  mode_help :: cmd -> String
  mode_help _ = ""

  -- | Optionally override the default usage string for each mode. Use patterns
  -- like in @run'@.
  mode_synopsis :: cmd -> Maybe String
  mode_synopsis _ = Nothing

data RecordMode cmd = RecordMode { rec_cmdname :: String
                                 , rec_initial :: Constr }
                    deriving (Typeable)

instance (Eq cmd, Eq (Record cmd), RecordCommand cmd, Data cmd, Attributes cmd)
         => Command (RecordMode cmd) (Record cmd) where
  cmdname = rec_cmdname
  run _ = run'

  summary cmd = mode_summary (fromConstr $ rec_initial cmd :: cmd)
  help cmd = mode_help (fromConstr $ rec_initial cmd :: cmd)
  synopsis cmd = fromMaybe
    (unwords $ cmdname cmd : usageDescr (attrFun $ cmdattrs cmd %% flag_attrs))
    (mode_synopsis (fromConstr $ rec_initial cmd :: cmd))

  options cmd = rec_options ctor %% available %% everywhere disable
    where available = [ (KeyF (typeOf (undefined :: cmd)) opt, [enable])
                      | opt <- constrFields $ rec_initial cmd]
          ctor = fromConstr $ rec_initial cmd :: cmd

  optionStyle cmd = rec_optionStyle (fromConstr $ rec_initial cmd :: cmd)

  supercommand cmd = rec_superCommand (fromConstr $ rec_initial cmd :: cmd)

  cmd_flag_empty cmd = fromConstr $ rec_initial cmd

-- | Construct a command list (for "dispatch"/"helpCommands") from a
-- multi-constructor record data type. See also "RecordCommand". Alternatively,
-- you can use "dispatchR" directly.
recordCommands :: forall cmd. (Eq cmd, Eq (Record cmd), Data cmd, RecordCommand cmd, Attributes cmd)
               => cmd -> [CommandWrap]
recordCommands _ = go $ dataTypeConstrs $ dataTypeOf (undefined :: cmd)
    where go = map $ CommandWrap . rec_cmd
          rec_cmd :: Constr -> RecordMode cmd
          rec_cmd x = RecordMode { rec_cmdname = nameFromConstr x
                                 , rec_initial = x }

-- | Record field update using a string field name. Sets a field value in a
-- record, using a (string) name of the field,
setField :: forall rec a. (Data rec) => String -> rec -> (forall b. (Data b) => b) -> rec
setField field current value = evalState (gmapM subst current) (constrFields $ toConstr current)
    where subst :: Data x => x -> State [String] x
          subst f = do x:xs <- get
                       put xs
                       if x == field then return value else return f

data Imp = Imp Int (forall a. (Typeable a) => a)

getField :: forall rec a. (Data rec, Typeable a) => String -> rec -> a
getField field value = case execState (gmapM find value) (Imp 0 (error "")) of Imp _ v -> v
    where idx = elemIndex field $ constrFields (toConstr value)
          find :: Data x => x -> State Imp x
          find f = do Imp i val <- get
                      if Just i == idx then put $ Imp (i + 1) (errcast f)
                                         else put $ Imp (i + 1) val
                      return $ error "find"
          errcast :: forall a b. (Typeable a, Typeable b) => a -> b
          errcast = fromMaybe (error errorMsg) . cast where
            errorMsg = "BUG: getField used with wrong type on " ++ field
              ++ " expected: " ++ show (typeOf (undefined :: a))
              ++ " got: " ++ show (typeOf (undefined :: b))

evalField :: forall rec. (Data rec) => String -> rec -> IO ()
evalField field value = do ref <- newIORef 0
                           gmapM (find ref) value
                           return ()
    where idx = elemIndex field $ constrFields (toConstr value)
          find :: Data x => IORef Int -> x -> IO x
          find ref f = do i <- readIORef ref
                          writeIORef ref $ i + 1
                          when (Just i == idx) (evaluate f >> return ())
                          return f

-- | A command parsing & dispatch entry point for record-based
-- commands. Ex. (see "RecordCommand"):
--
-- > main = getArgs >>= dispatchR [] >>= \x -> case x of
-- >   Foo {} -> putStrLn $ "You asked for foo. Wibble = " ++ show (wibble x)
-- >   Bar {} -> putStrLn $ "You asked for bar. ..."
dispatchR :: forall cmd f. (Eq cmd, Eq (Record cmd), Attributes cmd, RecordCommand cmd,
                            Command (RecordMode cmd) f,
                            Folded f ~ cmd) => [DispatchOpt] -> [String] -> IO cmd
dispatchR dopt opts = dispatch' die dopt (recordCommands (undefined :: cmd)) opts >>= \c -> case c of
  Nothing -> exitWith ExitSuccess >> return undefined
  Just (CommandWrap x, opts') -> execute' x opts' >>= \c -> case c of
    Just (command, opts') -> case cast command of
      Just comm -> return comm
      Nothing -> execute x opts' >> exitWith ExitSuccess >> return undefined
    Nothing -> exitWith ExitSuccess >> return undefined

-- | Like "execute", but you get the flags as a return value. This is useful to
-- implement non-modal applications with record-based flags, eg.:
--
-- > data Main = Main { greeting :: String, again :: Bool }
-- >     deriving (Typeable, Data, Eq)
-- > instance Attributes Main where -- (...)
-- > instance RecordCommand Main
-- > main = getArgs >>= executeR Main {} >>= \opts -> do
-- >    putStrLn (greeting opts) -- (...)
executeR :: forall cmd. (Eq cmd, Eq (Record cmd), Attributes cmd, RecordCommand cmd)
         => cmd -> [String] -> IO cmd
executeR cmd opts = execute' cmd' opts >>= \c -> case c of
  Just (command, _) -> return command
  Nothing -> exitWith ExitSuccess >> return undefined
  where cmd' = RecordMode (nameFromConstr $ toConstr cmd) (toConstr cmd)