1 files changed, 68 insertions, 7 deletions

diff --git a/lib/CommandLine.hs b/lib/CommandLine.hs
index 88584ce..62f0315 100644
--- a/lib/CommandLine.hs
+++ b/lib/CommandLine.hs
@@ -9,7 +9,9 @@ module CommandLine
     ( Args
     , UsageError(..)
     , usageErrorMessage
-    , parseInvocation
+    , ArgsStyle(..)
+    , vanilla
+    , fancy
     , runArgs
     , arg
     , param
@@ -235,7 +237,7 @@ param :: Int -> Args String
 param n = Args (Prim (\_ us -> us!!n) (singleton $ exactly (n+1))) []
 
 arg :: String -> Args String
-arg optname = Args (Prim (\opts _ -> head $ concat $ take 1 opts)
+arg optname = Args (Prim (\opts _ -> concat $ take 1 $ concat $ take 1 opts)
                    (singleton $ exactly 0))
                    [optname]
 
@@ -247,7 +249,7 @@ label :: String -> Args a -> Args a
 label _ = id
 
 data ArgsStyle = ArgsStyle
-    { parseInvocation :: String -> [String] -> ([(String,[String])], [String])
+    { parseInvocation :: [String] -> ([(String,[String])], [String])
     }
 
 -- | Phase 1.  This function accepts a list of command line arguments in its
@@ -262,8 +264,8 @@ data ArgsStyle = ArgsStyle
 -- unnamed argument values.
 --
 -- The results of this function are intended to be used as input to 'runArgs'.
-vanilla :: ArgsStyle
-vanilla = ArgsStyle { parseInvocation = parse }
+vanilla :: String -> ArgsStyle
+vanilla flags  = ArgsStyle { parseInvocation = parse flags }
  where
   parse flags cli = (opts, concat nons ++ nondashed ++ drop 1 trailing)
    where
@@ -275,6 +277,65 @@ vanilla = ArgsStyle { parseInvocation = parse }
     splitParams (('-':[x]):xs) | x `elem` flags = (['-':[x]],xs)
     splitParams xs = splitAt 2 xs
 
+-- |
+-- Arguments:
+--
+--   * option-count pairs - List of option names paired with number of expected values to follow them.
+--
+--   * polyvariadic options - List of option names that can take any number of arguments.
+--
+--   * default polyvariadic - Implicit polyvariadic option if no other option is specified.
+--
+fancy sargspsec polyVariadicArgs defaultPoly = ArgsStyle
+    { parseInvocation = parse sargspsec polyVariadicArgs defaultPoly }
+ where
+  parse sargspec polyVariadicArgs defaultPoly args_raw = (Map.toList $ fst margs, snd margs)
+    where
+        (args,trail1) = break (=="--") args_raw
+        trail = drop 1 trail1
+        commonArgSpec = []
+        -- [ ("--homedir",1) , ("--passphrase-fd",1) , ("--help",0) ]
+        sargspec' = commonArgSpec ++ sargspec
+        (sargs,margs) = (sargs, tackTail mapped )
+                    where (sargs,vargs) = partitionStaticArguments sargspec' args
+                          argspec = map fst sargspec' ++ polyVariadicArgs
+                          args' = if null defaultPoly || map (take 1) (take 1 vargs) == ["-"]
+                                    then vargs
+                                    else defaultPoly:vargs
+                          -- grouped args
+                          gargs = (sargs ++)
+                                  . (if null defaultPoly then id else toLast (++trail))
+                                  . groupBy (\_ s-> take 1 s /= "-")
+                                  $ args'
+                          append (m,as) (k:xs)
+                            | k `elem` argspec = (Map.alter (appendArgs k xs) k m,as)
+                            | null defaultPoly && "-" == take 1 k
+                                               = (Map.alter (appendArgs k xs) k m,as)
+                            | otherwise        = (m,as ++ (k:xs))
+                          appendArgs k xs opt = Just . maybe xs (++xs) $ opt
+                          mapped = foldl' append (Map.empty,[]) gargs
+                          tackTail | null defaultPoly = second (++trail)
+                                   | otherwise        = id
+
+        -- partitionStaticArguments :: Ord a => [(a, Int)] -> [a] -> ([[a]], [a])
+        partitionStaticArguments specs args = psa args
+         where
+            smap = Map.fromList specs
+            psa [] = ([],[])
+            psa (a:as) =
+              case Map.lookup a smap of
+                Nothing | (k,'=':v) <- break (=='=') a
+                        , Just 1 <- Map.lookup k smap
+                    -> first ([k,v]:) $ psa as
+                Nothing -> second (a:) $ psa as
+                Just n  -> first ((a:take n as):) $ psa (drop n as)
+
+        -- | Returns the given list with its last element modified.
+        toLast :: (x -> x) -> [x] -> [x]
+        toLast f []  = []
+        toLast f [x] = [f x]
+        toLast f (x:xs) = x : toLast f xs
+
 -- | Information about how the user failed to satisfy a specified usage.
 data UsageError
     = TooManyParameters Int
@@ -409,9 +470,9 @@ runArgs (os,us) c
     dups = mapMaybe notSingle $ groupBy ((==) `on` fst) (os' :: [(String,[String])])
         where notSingle [x] = Nothing
               notSingle ((k,v):xs)  = Just (k,v : map snd xs)
-    getbit = let r = Map.fromList $ zip (accepts c) [0..] in trace ("getbit = "++show r) r
+    getbit = Map.fromList $ zip (accepts c) [0..]
     goods :: [(Int,[String])]
-    (bads,goods) = let r = partitionEithers $ map f os' in trace ("(bads,goods)="++show r) r
+    (bads,goods) = partitionEithers $ map f os'
         where f (k,v) = case Map.lookup k getbit of
                             Just b  -> Right (b,v)
                             Nothing -> Left  (k,v)