Map local variables in IR to their original names

1 files changed, 47 insertions, 25 deletions

diff --git a/lib/Language/Elna/Glue.hs b/lib/Language/Elna/Glue.hs
index 8f28696..9101ca5 100644
--- a/lib/Language/Elna/Glue.hs
+++ b/lib/Language/Elna/Glue.hs
@@ -4,7 +4,7 @@ module Language.Elna.Glue
 
 import Control.Monad.Trans.State (State, gets, modify', runState)
 import Data.Bifunctor (Bifunctor(..))
-import Data.Foldable (Foldable(..))
+import Data.Foldable (Foldable(..), traverse_)
 import Data.HashMap.Strict (HashMap)
 import qualified Data.HashMap.Strict as HashMap
 import Data.Maybe (catMaybes)
@@ -25,10 +25,12 @@ import Language.Elna.Backend.Intermediate
 import Language.Elna.Frontend.SymbolTable (Info(..), SymbolTable)
 import qualified Language.Elna.Frontend.SymbolTable as SymbolTable
 import GHC.Records (HasField(..))
+import Language.Elna.Frontend.AST (Identifier(..))
 
 data Paste = Paste
     { temporaryCounter :: Word32
     , labelCounter :: Word32
+    , localMap :: HashMap Identifier Variable
     }
 
 newtype Glue a = Glue
@@ -47,31 +49,46 @@ instance Monad Glue
   where
     (Glue x) >>= f = Glue $ x >>= (runGlue . f)
 
-glue :: SymbolTable -> AST.Program -> HashMap AST.Identifier (Vector (Quadruple Variable))
+glue :: SymbolTable -> AST.Program -> HashMap Identifier (Vector (Quadruple Variable))
 glue globalTable
     = fst
-    . flip runState Paste{ temporaryCounter = 0, labelCounter = 0 }
+    . flip runState emptyPaste
     . runGlue
     . program globalTable
+  where
+    emptyPaste = Paste
+        { temporaryCounter = 0
+        , labelCounter = 0
+        , localMap = mempty
+        }
 
-program
-    :: SymbolTable
-    -> AST.Program
-    -> Glue (HashMap AST.Identifier (Vector (Quadruple Variable)))
-program globalTable (AST.Program declarations) = HashMap.fromList . catMaybes
+program :: SymbolTable -> AST.Program -> Glue (HashMap Identifier (Vector (Quadruple Variable)))
+program globalTable (AST.Program declarations)
+    = HashMap.fromList . catMaybes
     <$> traverse (declaration globalTable) declarations
 
 declaration
     :: SymbolTable
     -> AST.Declaration
     -> Glue (Maybe (AST.Identifier, Vector (Quadruple Variable)))
-declaration globalTable (AST.ProcedureDeclaration procedureName _ _ statements)
-    = Just
-    . (procedureName,)
-    . Vector.cons StartQuadruple
-    . flip Vector.snoc StopQuadruple
-    . fold
-    <$> traverse (statement globalTable) statements
+declaration globalTable (AST.ProcedureDeclaration procedureName _ variableDeclarations statements)
+    = traverse_ registerVariable variableDeclarations
+    >> nameQuadruplesTuple <$> traverse (statement globalTable) statements
+  where
+    registerVariable (AST.VariableDeclaration identifier _) = do
+        currentCounter <- fmap (fromIntegral . HashMap.size)
+            $ Glue $ gets $ getField @"localMap"
+        Glue $ modify' $ modifier identifier $ LocalVariable currentCounter
+    modifier identifier currentCounter generator = generator
+        { localMap = HashMap.insert identifier currentCounter
+            $ getField @"localMap" generator
+        }
+    nameQuadruplesTuple quadrupleList = Just
+        ( procedureName
+        , Vector.cons StartQuadruple
+            $ flip Vector.snoc StopQuadruple
+            $ fold quadrupleList
+        )
 declaration _ (AST.TypeDefinition _ _) = pure Nothing
 
 statement :: SymbolTable -> AST.Statement -> Glue (Vector (Quadruple Variable))
@@ -104,15 +121,16 @@ statement localTable (AST.AssignmentStatement variableAccess' assignee) = do
     (rhsOperand, rhsStatements) <- expression localTable assignee
     let variableType' = variableType variableAccess' localTable
     accessResult <- variableAccess localTable variableAccess' Nothing variableType' mempty
-    pure $ rhsStatements <> case accessResult of
+    lhsStatements <- case accessResult of
             {-(AST.Identifier identifier, Just accumulatedIndex, accumulatedStatements) ->
                 Vector.snoc accumulatedStatements
                     $ ArrayAssignQuadruple rhsOperand accumulatedIndex
                     $ LocalVariable identifier -}
-            (AST.Identifier identifier, Nothing, accumulatedStatements) ->
-                Vector.snoc accumulatedStatements
-                    $ AssignQuadruple rhsOperand
-                    $ LocalVariable 0
+            (identifier, _Nothing, accumulatedStatements)
+                -> Vector.snoc accumulatedStatements
+                . AssignQuadruple rhsOperand
+                <$> lookupLocal identifier
+    pure $ rhsStatements <> lhsStatements
 {- statement localTable (AST.WhileStatement whileCondition whileStatement) = do
     (conditionStatements, jumpConstructor) <- condition localTable whileCondition
     startLabel <- createLabel
@@ -135,6 +153,10 @@ createTemporary = do
         { temporaryCounter = getField @"temporaryCounter" generator + 1
         }
 
+lookupLocal :: Identifier -> Glue Variable
+lookupLocal identifier =
+    fmap (HashMap.! identifier) $ Glue $ gets $ getField @"localMap"
+
 createLabel :: Glue Label
 createLabel = do
     currentCounter <- Glue $ gets $ getField @"labelCounter"
@@ -242,13 +264,13 @@ expression localTable = \case
     (AST.DivisionExpression lhs rhs) ->
         binaryExpression DivisionQuadruple lhs rhs
     (AST.VariableExpression variableExpression) -> do
-        pure (VariableOperand (LocalVariable 0), mempty)
-        {- let variableType' = variableType variableExpression localTable
+        let variableType' = variableType variableExpression localTable
         variableAccess' <- variableAccess localTable variableExpression Nothing variableType' mempty
         case variableAccess' of
-            (AST.Identifier identifier, Nothing, statements) ->
-                pure (VariableOperand (Variable identifier), statements)
-            (AST.Identifier identifier, Just operand, statements) -> do
+            (identifier, _Nothing, statements)
+                -> (, statements) . VariableOperand 
+                <$> lookupLocal identifier
+            {-(AST.Identifier identifier, Just operand, statements) -> do
                 arrayAddress <- createTemporary
                 let arrayStatement = ArrayQuadruple (Variable identifier) operand arrayAddress
                 pure