1 files changed, 23 insertions, 19 deletions
diff --git a/src/Main.hs b/src/Main.hs
index aff5360..9e70cd6 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -6,7 +6,11 @@ module Main
     ( main
     ) where
 
-import Language.Elna.CommandLine (CommandLine(..), commandLine, execParser)
+import Language.Elna.Driver
+    ( Driver(..)
+    , IntermediateStage(..)
+    , drive
+    )
 import Language.Elna.Object.ElfCoder (elfObject)
 import Language.Elna.Backend.Allocator (allocate)
 import Language.Elna.Glue (glue)
@@ -15,8 +19,6 @@ import Language.Elna.Frontend.Parser (programP)
 import Language.Elna.Frontend.TypeAnalysis (typeAnalysis)
 import Language.Elna.RiscV.CodeGenerator (generateRiscV, riscVConfiguration)
 import Language.Elna.RiscV.ElfWriter (riscv32Elf)
-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)
@@ -32,27 +34,29 @@ import Control.Exception (IOException, catch)
 -- 6 - Register allocation error.
 
 main :: IO ()
-main = execParser commandLine >>= withCommandLine
+main = drive >>= withCommandLine
   where
-    withCommandLine CommandLine{..} =
-        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 (printAndExit 4) (withSymbolTable output program)
-        $ nameAnalysis program
+    withCommandLine driver@Driver{ input }
+        = catch (Text.readFile input) (printAndExit 2 :: IOException -> IO a)
+        >>= withParsedInput driver
+        . runParser programP input
+    withParsedInput driver@Driver{ intermediateStage } (Right program)
+        | Just ParseStage <- intermediateStage = pure ()
+        | otherwise
+            = either (printAndExit 4) (withSymbolTable driver program)
+            $ nameAnalysis program 
     withParsedInput _ (Left errorBundle)
         = putStrLn (errorBundlePretty errorBundle)
         >> exitWith (ExitFailure 3)
-    withSymbolTable output program symbolTable
+    withSymbolTable driver@Driver{ intermediateStage } program symbolTable
         | Just typeError <- typeAnalysis symbolTable program =
             printAndExit 5 typeError
-        | otherwise =
-            let makeObject = elfObject output . riscv32Elf . generateRiscV
-             in either (printAndExit 6) makeObject
-                $ allocate riscVConfiguration symbolTable
-                $ glue symbolTable program
+        | Just ValidateStage <- intermediateStage = pure ()
+        | otherwise = either (printAndExit 6) (withTac driver)
+            $ allocate riscVConfiguration symbolTable
+            $ glue symbolTable program
+    withTac Driver{ intermediateStage, output } tac
+        | Just CodeGenStage <- intermediateStage = pure ()
+        | otherwise = elfObject output $ riscv32Elf $ generateRiscV tac
     printAndExit :: Show b => forall a. Int -> b -> IO a
     printAndExit failureCode e = print e >> exitWith (ExitFailure failureCode)