1 files changed, 14 insertions, 15 deletions
diff --git a/Main.hs b/Main.hs
index e3ed454..20f479b 100644
--- a/Main.hs
+++ b/Main.hs
@@ -22,25 +22,24 @@ usage = do
     putStrLn $ "  -irpre  Print IR after optimisation, before stackification"
     putStrLn $ "  -ir     Print IR after optimisation and stackification"
 
-data Options = Options { optAST :: Bool, optIRPre :: Bool, optIR :: Bool }
+data Options = Options { optAST :: Bool, optIRPre :: Bool, optIR :: Bool, optFiles :: [FilePath] }
 
--- TODO: clean this function up
-parseOptions' :: (Options, Maybe FilePath) -> [String] -> IO (Options, Maybe FilePath)
-parseOptions' pair [] = return pair
-parseOptions' _ ("-h":_) = usage >> exitSuccess
-parseOptions' (opt, fp) ("-ast":as) = parseOptions' (opt { optAST = True }, fp) as
-parseOptions' (opt, fp) ("-irpre":as) = parseOptions' (opt { optIRPre = True }, fp) as
-parseOptions' (opt, fp) ("-ir":as) = parseOptions' (opt { optIR = True }, fp) as
-parseOptions' _ (('-':a):_) = putStrLn ("Unknown option '" ++ a ++ "'") >> usage >> exitFailure
-parseOptions' (opt, Nothing) (f:as) = parseOptions' (opt, Just f) as
-parseOptions' (_, Just _) (_:_) = putStrLn "At most one filename argument expected" >> usage >> exitFailure
-
-parseOptions :: [String] -> IO (Options, Maybe FilePath)
-parseOptions = parseOptions' (Options False False False, Nothing)
+parseOptions :: Options -> [String] -> IO Options
+parseOptions o []            = return o
+parseOptions _ ("-h":_)      = usage >> exitSuccess
+parseOptions o ("-ast":as)   = parseOptions (o { optAST = True }) as
+parseOptions o ("-irpre":as) = parseOptions (o { optIRPre = True }) as
+parseOptions o ("-ir":as)    = parseOptions (o { optIR = True }) as
+parseOptions _ (('-':a):_)   = putStrLn ("Unknown option '" ++ a ++ "'") >> usage >> exitFailure
+parseOptions o (f:as)        = parseOptions (o { optFiles = optFiles o ++ [f] }) as
 
 main :: IO ()
 main = do
-    (opts, mfname) <- getArgs >>= parseOptions
+    opts <- getArgs >>= parseOptions (Options False False False [])
+    mfname <- case optFiles opts of
+                  [] -> return Nothing
+                  [fname] -> return (Just fname)
+                  _ -> putStrLn "At most one filename argument expected" >> usage >> exitFailure
     source <- maybe getContents readFile mfname
 
     prog <- parseProgram mfname source >>= either (die . show) return