1 files changed, 24 insertions, 9 deletions
diff --git a/lib/Language/Elna/Backend/Allocator.hs b/lib/Language/Elna/Backend/Allocator.hs
index 15e8b6f..a56a73b 100644
--- a/lib/Language/Elna/Backend/Allocator.hs
+++ b/lib/Language/Elna/Backend/Allocator.hs
@@ -1,3 +1,7 @@
+{- This Source Code Form is subject to the terms of the Mozilla Public License,
+   v. 2.0. If a copy of the MPL was not distributed with this file, You can
+   obtain one at https://mozilla.org/MPL/2.0/. -}
+
 module Language.Elna.Backend.Allocator
     ( MachineConfiguration(..)
     , Store(..)
@@ -28,12 +32,19 @@ data Store r
     = RegisterStore r
     | StackStore Int32 r
 
-data AllocationError = AllocationError
+data AllocationError
+    = OutOfRegistersError
+    | UnexpectedProcedureInfoError Info
+    | UndefinedSymbolError Identifier
     deriving Eq
 
 instance Show AllocationError
   where
-    show = const "Ran out of registers during register allocation"
+    show OutOfRegistersError = "Ran out of registers during register allocation"
+    show (UnexpectedProcedureInfoError info) =
+        "Expected to encounter a procedure, got: " <> show info
+    show (UndefinedSymbolError identifier) =
+        concat ["Symbol \"", show identifier, "\" is not defined"]
 
 newtype MachineConfiguration r = MachineConfiguration
     { temporaryRegisters :: [r]
@@ -74,10 +85,14 @@ allocate machineConfiguration globalTable = HashMap.traverseWithKey function
         . runAllocator
         . mapM quadruple
     function :: Identifier -> Vector (Quadruple Variable) -> Either AllocationError (ProcedureQuadruples (Store r))
-    function identifier quadruples' =
-        let Just (ProcedureInfo localTable _) = SymbolTable.lookup identifier globalTable
-            (result, lastState) = run localTable quadruples'
-         in makeResult lastState <$> result
+    function procedureName quadruples' =
+        case SymbolTable.lookup procedureName globalTable of
+            Just (ProcedureInfo localTable _) ->
+                let (result, lastState) = run localTable quadruples'
+                 in makeResult lastState <$> result
+            Just anotherInfo -> Left $ UnexpectedProcedureInfoError anotherInfo
+            Nothing -> Left $ UndefinedSymbolError procedureName
+
     makeResult MachineState{ symbolTable } result = ProcedureQuadruples
         { quadruples = result
         , stackSize = fromIntegral $ SymbolTable.size symbolTable * 4
@@ -155,13 +170,13 @@ operand (VariableOperand variableOperand) =
 storeVariable :: Variable -> Allocator r (Store r)
 storeVariable (TempVariable index) = do
     temporaryRegisters' <- Allocator $ lift $ asks $ getField @"temporaryRegisters"
-    maybe (Allocator $ throwE AllocationError) (pure . RegisterStore)
+    maybe (Allocator $ throwE OutOfRegistersError) (pure . RegisterStore)
         $ temporaryRegisters' !? fromIntegral index
 storeVariable (LocalVariable index) = do
     temporaryRegisters' <- Allocator $ lift $ asks $ getField @"temporaryRegisters"
-    maybe (Allocator $ throwE AllocationError) (pure . StackStore (fromIntegral (succ index) * (-4)))
+    maybe (Allocator $ throwE OutOfRegistersError) (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))
+    maybe (Allocator $ throwE OutOfRegistersError) (pure . StackStore (fromIntegral index * 4))
         $ temporaryRegisters' !? fromIntegral index