Check argument list length

1 files changed, 31 insertions, 17 deletions

diff --git a/src/Main.hs b/src/Main.hs
index 2e02955..81e5976 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -15,25 +15,39 @@ import Data.Maybe (fromMaybe)
 import System.FilePath (replaceExtension, takeFileName)
 import Text.Megaparsec (runParser, errorBundlePretty)
 import qualified Data.Text.IO as Text
+import System.Exit (ExitCode(..), exitWith)
+import Control.Exception (IOException, catch)
+
+-- * Error codes
+--
+-- 1 - Command line parsing failed and other errors.
+-- 2 - The input could not be read.
+-- 3 - Parse error.
+-- 4 - Name analysis error.
+-- 5 - Type error.
 
 main :: IO ()
-main  = execParser commandLine >>= withCommandLine
+main = execParser commandLine >>= withCommandLine
   where
     withCommandLine CommandLine{..} =
-        let defaultOutput = flip fromMaybe output
-                $ replaceExtension (takeFileName input) "o"
-         in Text.readFile input
-            >>= withParsedInput defaultOutput
+        let defaultOutputName = replaceExtension (takeFileName input) "o"
+            outputName = fromMaybe defaultOutputName output
+         in catch (Text.readFile input) (printAndExit 2 :: IOException -> IO a)
+            >>= withParsedInput outputName
             . runParser programP input
-    withParsedInput output (Right program) =
-        either print (withSymbolTable output program)
-            $ nameAnalysis program
-    withParsedInput _ (Left errorBundle) = putStrLn
-        $ errorBundlePretty errorBundle
-    withSymbolTable output program symbolTable =
-        let _ = typeAnalysis symbolTable program
-            instructions = generateRiscV
-                $ allocate riscVConfiguration
-                $ glue symbolTable program
-         in elfObject output
-            $ riscv32Elf instructions
+    withParsedInput output (Right program)
+        = either (printAndExit 4) (withSymbolTable output program)
+        $ nameAnalysis program
+    withParsedInput _ (Left errorBundle)
+        = putStrLn (errorBundlePretty errorBundle)
+        >> exitWith (ExitFailure 3)
+    withSymbolTable output program symbolTable
+        | Just typeError <- typeAnalysis symbolTable program =
+            printAndExit 5 typeError
+        | otherwise =
+            let instructions = generateRiscV
+                    $ allocate riscVConfiguration
+                    $ glue symbolTable program
+             in elfObject output $ riscv32Elf instructions
+    printAndExit :: Show b => forall a. Int -> b -> IO a
+    printAndExit failureCode e = print e >> exitWith (ExitFailure failureCode)