1 files changed, 38 insertions, 22 deletions
diff --git a/lib/Language/Elna/Backend/Allocator.hs b/lib/Language/Elna/Backend/Allocator.hs
index 9a85605..f0f285b 100644
--- a/lib/Language/Elna/Backend/Allocator.hs
+++ b/lib/Language/Elna/Backend/Allocator.hs
@@ -5,6 +5,7 @@ module Language.Elna.Backend.Allocator
     ) where
 
 import Data.HashMap.Strict (HashMap)
+import Data.Int (Int32)
 import Data.Word (Word32)
 import Data.Vector (Vector)
 import Language.Elna.Backend.Intermediate
@@ -17,11 +18,20 @@ import Language.Elna.Location (Identifier(..))
 import Control.Monad.Trans.Reader (ReaderT, asks, runReaderT)
 import Control.Monad.Trans.State (State, runState, modify')
 import GHC.Records (HasField(..))
-import Control.Monad.Trans.Class (MonadTrans(lift))
+import Control.Monad.Trans.Class (MonadTrans(..))
+import Control.Monad.Trans.Except (ExceptT(..), runExceptT, throwE)
+import Data.List ((!?))
 
 data Store r
     = RegisterStore r
-    | StackStore Word32 r
+    | StackStore Int32 r
+
+data AllocationError = AllocationError
+    deriving Eq
+
+instance Show AllocationError
+  where
+    show = const "Ran out of registers during register allocation"
 
 newtype MachineConfiguration r = MachineConfiguration
     { temporaryRegisters :: [r]
@@ -32,7 +42,7 @@ newtype MachineState = MachineState
     } deriving (Eq, Show)
 
 newtype Allocator r a = Allocator
-    { runAllocator :: ReaderT (MachineConfiguration r) (State MachineState) a
+    { runAllocator :: ExceptT AllocationError (ReaderT (MachineConfiguration r) (State MachineState)) a
     }
 
 instance forall r. Functor (Allocator r)
@@ -52,20 +62,22 @@ allocate
     :: forall r
     . MachineConfiguration r
     -> HashMap Identifier (Vector (Quadruple Variable))
-    -> HashMap Identifier (ProcedureQuadruples (Store r))
-allocate machineConfiguration = fmap function
+    -> Either AllocationError (HashMap Identifier (ProcedureQuadruples (Store r)))
+allocate machineConfiguration = traverse function
   where
-    function :: Vector (Quadruple Variable) -> ProcedureQuadruples (Store r)
+    run = flip runState (MachineState{ stackSize = 0 })
+        . flip runReaderT machineConfiguration
+        . runExceptT
+        . runAllocator
+        . mapM quadruple
+    function :: Vector (Quadruple Variable) -> Either AllocationError (ProcedureQuadruples (Store r))
     function quadruples' =
-        let (result, lastState)
-                = flip runState (MachineState{ stackSize = 0 })
-                $ flip runReaderT machineConfiguration
-                $ runAllocator
-                $ mapM quadruple quadruples'
-         in ProcedureQuadruples
-            { quadruples = result
-            , stackSize = getField @"stackSize" lastState
-            }
+        let (result, lastState) = run quadruples'
+         in makeResult lastState <$> result
+    makeResult MachineState{ stackSize } result = ProcedureQuadruples
+        { quadruples = result
+        , stackSize = stackSize
+        }
 
 quadruple :: Quadruple Variable -> Allocator r (Quadruple (Store r))
 quadruple = \case
@@ -131,11 +143,15 @@ operand (VariableOperand variableOperand) =
 
 storeVariable :: Variable -> Allocator r (Store r)
 storeVariable (TempVariable index) = do
-    temporaryRegisters' <- Allocator $ asks $ getField @"temporaryRegisters"
-    pure $ RegisterStore
-        $ temporaryRegisters' !! fromIntegral index
+    temporaryRegisters' <- Allocator $ lift $ asks $ getField @"temporaryRegisters"
+    maybe (Allocator $ throwE AllocationError) (pure . RegisterStore)
+        $ temporaryRegisters' !? fromIntegral index
 storeVariable (LocalVariable index) = do
-    temporaryRegisters' <- Allocator $ asks $ getField @"temporaryRegisters"
-    Allocator $ lift $ modify' $ MachineState . (+ 4) . getField @"stackSize"
-    pure $ StackStore (index * 4)
-        $ temporaryRegisters' !! pred (length temporaryRegisters' - fromIntegral index)
+    temporaryRegisters' <- Allocator $ lift $ asks $ getField @"temporaryRegisters"
+    Allocator $ lift $ lift $ modify' $ MachineState . (+ 4) . getField @"stackSize"
+    maybe (Allocator $ throwE AllocationError) (pure . StackStore (fromIntegral (succ index) * (-4)))
+        $ temporaryRegisters' !? pred (length temporaryRegisters' - fromIntegral index)
+storeVariable (ParameterVariable index) = do
+    temporaryRegisters' <- Allocator $ lift $ asks $ getField @"temporaryRegisters"
+    maybe (Allocator $ throwE AllocationError) (pure . StackStore (fromIntegral index * 4))
+        $ temporaryRegisters' !? fromIntegral index