adjust live variables

anoma · Feb 19, 2024 · 902767a · 902767a
1 parent 86524e4
commit 902767a
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 deletions.
diff --git a/src/Juvix/Compiler/Reg/Data/IndexMap.hs b/src/Juvix/Compiler/Reg/Data/IndexMap.hs
@@ -33,8 +33,11 @@ assign IndexMap {..} k =
       }
   )
 
+lookup' :: (Hashable k) => IndexMap k -> k -> Maybe Index
+lookup' IndexMap {..} k = HashMap.lookup k _indexMapTable
+
 lookup :: (Hashable k) => IndexMap k -> k -> Index
-lookup IndexMap {..} k = fromJust $ HashMap.lookup k _indexMapTable
+lookup mp = fromJust . lookup' mp
 
 combine :: forall k. (Hashable k) => IndexMap k -> IndexMap k -> IndexMap k
 combine mp1 mp2 =

diff --git a/src/Juvix/Compiler/Reg/Transformation/SSA.hs b/src/Juvix/Compiler/Reg/Transformation/SSA.hs
@@ -22,13 +22,20 @@ computeFunctionSSA =
   where
     go :: Instruction -> IndexMap VarRef -> (IndexMap VarRef, Instruction)
     go instr mp = case getResultVar instr' of
-      Just vref -> (mp', setResultVar instr' (mkVarRef VarGroupLocal idx))
+      Just vref -> (mp', updateLiveVars mp' (setResultVar instr' (mkVarRef VarGroupLocal idx)))
         where
           (idx, mp') = IndexMap.assign mp vref
-      Nothing -> (mp, instr')
+      Nothing -> (mp, updateLiveVars mp instr')
       where
         instr' = overValueRefs (adjustVarRef mp) instr
 
+    updateLiveVars :: IndexMap VarRef -> Instruction -> Instruction
+    updateLiveVars mp = \case
+      Prealloc x -> Prealloc $ over instrPreallocLiveVars (mapMaybe (adjustVarRef' mp)) x
+      Call x -> Call $ over instrCallLiveVars (mapMaybe (adjustVarRef' mp)) x
+      CallClosures x -> CallClosures $ over instrCallClosuresLiveVars (mapMaybe (adjustVarRef' mp)) x
+      instr -> instr
+
     -- For branches, when necessary we insert assignments unifying the renamed
     -- output variables into a single output variable for both branches.
     combine :: Instruction -> NonEmpty (IndexMap VarRef) -> (IndexMap VarRef, Instruction)
@@ -103,6 +110,13 @@ computeFunctionSSA =
       VarGroupArgs -> vref
       VarGroupLocal -> mkVarRef VarGroupLocal (IndexMap.lookup mpv vref)
 
+    adjustVarRef' :: IndexMap VarRef -> VarRef -> Maybe VarRef
+    adjustVarRef' mpv vref = case IndexMap.lookup' mpv vref of
+      Just idx -> Just $ mkVarRef VarGroupLocal idx
+      Nothing -> case vref ^. varRefGroup of
+        VarGroupArgs -> Just vref
+        VarGroupLocal -> Nothing
+
     assignInBranch :: Code -> Index -> Index -> Code
     assignInBranch is idx idx' =
       is