Refactor completion monad.

Options/Applicative/BashCompletion.hs

@@ -23,12 +23,12 @@ bashCompletionParser parser = complParser
       <*> option (long "bash-completion-index") )
 
 bashCompletionQuery :: Parser a -> [String] -> Int -> [String]
-bashCompletionQuery parser ws i = case runCompletion compl ws i of
-  Nothing -> []
-  Just p  -> list_options p
+bashCompletionQuery parser ws i = case runCompletion compl ws i parser of
+  Left ComplExit -> []
+  _              -> []
   where
     list_options = concat . mapParser (\_ -> map show_name . optionNames . optMain)
     show_name (OptShort c) = '-':[c]
     show_name (OptLong name) = "--" ++ name
 
-    compl = runParserWith pure parser
+    compl = runParserWith (\_ _ -> exitCompletion) parser

Options/Applicative/Common.hs

@@ -46,6 +46,8 @@ module Options.Applicative.Common (
   optionNames
   ) where
 
+import Debug.Trace
+
 import Control.Applicative
 import Control.Monad
 import Data.Maybe
@@ -79,7 +81,8 @@ optMatches rdr arg = case rdr of
     , arg1 `elem` names
     -> Just $ do
          arg' <- nextArg val
-         liftMaybe $ f arg'
+         trace ("arg' nothing: " ++ show (isNothing arg')) $ return ()
+         liftMaybe $ arg' >>= f
     | otherwise -> Nothing
   FlagReader names x
     | Just (arg1, Nothing) <- parsed
@@ -128,17 +131,24 @@ stepParser (BindP p k) arg = do
   x <- liftMaybe $ evalParser p'
   return $ k x
 
-runParserWith :: MonadP m => (Parser a -> m b) -> Parser a -> m b
-runParserWith h p = tryP (h p) $ do
-    arg <- nextArg Nothing
-    p' <- stepParser p arg
-    runParserWith h p'
+runParserWith :: MonadP m => (Parser a -> Maybe String -> m b) -> Parser a -> m b
+runParserWith h p = do
+  a <- nextArg Nothing
+  case a of
+    Nothing -> h p Nothing
+    Just arg -> do
+      r <- tryP $ stepParser p arg
+      case r of
+        Left e -> h p (Just arg) <|> errorP e
+        Right p' -> do
+          setParser (Just arg) p'
+          runParserWith h p'
 
 -- | Apply a 'Parser' to a command line, and return a result and leftover
 -- arguments.  This function returns an error if any parsing error occurs, or
 -- if any options are missing and don't have a default value.
 runParser :: MonadP m => Parser a -> m a
-runParser = runParserWith (liftMaybe . evalParser)
+runParser = runParserWith $ \p _ -> liftMaybe (evalParser p)
 
 runParserFully :: Parser a -> P a
 runParserFully p = do

Options/Applicative/Internal.hs

@@ -1,4 +1,4 @@
-{-# LANGUAGE GADTs, FlexibleInstances #-}
+{-# LANGUAGE GADTs, FlexibleInstances, TypeFamilies #-}
 module Options.Applicative.Internal
   ( P
   , Context(..)
@@ -6,18 +6,20 @@ module Options.Applicative.Internal
 
   , uncons
   , liftMaybe
+
   , getArgs
   , setArgs
-
   , runP
+
   , runCompletion
+  , ComplError(..)
+  , exitCompletion
   ) where
 
 import Control.Applicative
 import Control.Monad
 import Control.Monad.Trans.Class
 import Control.Monad.Trans.Error
-import Control.Monad.Trans.Maybe
 import Control.Monad.Trans.State
 import Control.Monad.Trans.Writer
 import Data.Maybe
@@ -26,9 +28,14 @@ import Data.Monoid
 import Options.Applicative.Types
 
 class (Alternative m, MonadPlus m) => MonadP m where
-  nextArg :: Maybe String -> m String
+  type PError m
+
+  nextArg :: Maybe String -> m (Maybe String)
   setContext :: Maybe String -> ParserInfo a -> m ()
-  tryP :: m a -> m a -> m a
+  setParser :: Maybe String -> Parser a -> m ()
+
+  tryP :: m a -> m (Either (PError m) a)
+  errorP :: PError m -> m a
 
 type P = StateT [String] (ErrorT String (Writer Context))
 
@@ -42,23 +49,30 @@ instance Monoid Context where
   mappend c _ = c
 
 instance MonadP P where
+  type PError P = String
+
   nextArg val = do
     args <- getArgs
-    (arg', args') <- liftMaybe . uncons $ maybeToList val ++ args
-    setArgs args'
-    return arg'
+    case maybeToList val ++ args of
+      [] -> return Nothing
+      (arg':args') -> do
+        setArgs args'
+        return $ Just arg'
 
   setContext name = lift . lift . tell . Context name
+  setParser _ _ = return ()
 
-  tryP result p = do
+  errorP = lift . throwError
+
+  tryP p = do
     args <- getArgs
     let (r, ctx) = runP p args
     lift . lift . tell $ ctx
     case r of
-      Left e -> result `mplus` lift (throwError e)
+      Left e -> return (Left e)
       Right (x, args') -> do
         setArgs args'
-        return x
+        return (Right x)
 
 liftMaybe :: MonadPlus m => Maybe a -> m a
 liftMaybe = maybe mzero return
@@ -76,27 +90,58 @@ uncons :: [a] -> Maybe (a, [a])
 uncons [] = Nothing
 uncons (x : xs) = Just (x, xs)
 
-data ComplState = ComplState [String] !Int
-type Completion = MaybeT (State ComplState)
+data SomeParser where
+  SomeParser :: Parser a -> SomeParser
+
+data ComplState = ComplState
+  { complWords :: [String]
+  , complIndex :: !Int
+  , complParser :: SomeParser
+  , complArg :: String }
+
+data ComplError
+  = ComplParseError String
+  | ComplExit
+
+instance Error ComplError where
+  strMsg = ComplParseError
+
+type Completion = ErrorT ComplError (State ComplState)
 
 instance MonadP Completion where
+  type PError Completion = ComplError
+
   nextArg val = do
-    ComplState ws i <- lift get
-    guard $ i > 0
+    st <- lift get
+    let i = complIndex st
+    unless (i > 0) exitCompletion
     case val of
-      Just arg -> return arg
+      Just arg -> return $ Just arg
       Nothing -> do
+        let ws = complWords st
         (arg, ws') <- liftMaybe (uncons ws)
-        lift . put $ ComplState ws' (i - 1)
-        return arg
-  setContext _ _ = return ()
+        lift . modify $ \s -> s
+          { complWords = ws'
+          , complIndex = i - 1 }
+        return $ Just arg
 
-  tryP result p = do
-    r <- lift $ runMaybeT p
+  setContext val i = setParser val (infoParser i)
+  setParser val p = lift . modify $ \s -> s
+    { complParser = SomeParser p
+    , complArg = fromMaybe "" val }
+
+  errorP = throwError
+
+  tryP p = do
+    r <- lift $ runErrorT p
     case r of
-      Nothing -> result
-      Just x -> return x
+      Left e@(ComplParseError _) -> return (Left e)
+      Left e -> throwError e
+      Right x -> return (Right x)
 
-runCompletion :: Completion a -> [String] -> Int -> Maybe a
-runCompletion c ws i = evalState (runMaybeT c) s
-  where s = ComplState ws i
+runCompletion :: Completion r -> [String] -> Int -> Parser a -> Either ComplError r
+runCompletion c ws i p = evalState (runErrorT c) s
+  where s = ComplState ws i (SomeParser p) ""
+
+exitCompletion :: Completion ()
+exitCompletion = throwError ComplExit