branch out variables

src/Juvix/Compiler/Reg/Language.hs

@@ -206,9 +206,9 @@ data InstrBranch = InstrBranch
   { _instrBranchValue :: Value,
     _instrBranchTrue :: Code,
     _instrBranchFalse :: Code,
-    -- | Live variable storing the result (corresponds to the top of the value
+    -- | Output variable storing the result (corresponds to the top of the value
     -- stack in JuvixAsm after executing the branches)
-    _instrBranchVar :: Maybe VarRef
+    _instrBranchOutVar :: Maybe VarRef
   }
   deriving stock (Eq)
 
@@ -218,7 +218,7 @@ data InstrCase = InstrCase
     _instrCaseIndRep :: IndRep,
     _instrCaseBranches :: [CaseBranch],
     _instrCaseDefault :: Maybe Code,
-    _instrCaseVar :: Maybe VarRef
+    _instrCaseOutVar :: Maybe VarRef
   }
   deriving stock (Eq)
 

src/Juvix/Compiler/Reg/Pretty/Base.hs

@@ -108,12 +108,12 @@ ppLiveVars vars
       vars' <- mapM ppCode vars
       return $ comma <+> primitive "live:" <+> arglist vars'
 
-ppBranchVar :: (Member (Reader Options) r) => Maybe VarRef -> Sem r (Doc Ann)
-ppBranchVar = \case
+ppOutVar :: (Member (Reader Options) r) => Maybe VarRef -> Sem r (Doc Ann)
+ppOutVar = \case
   Nothing -> return mempty
   Just var -> do
     var' <- ppCode var
-    return $ comma <+> primitive "var:" <+> var'
+    return $ comma <+> primitive "out:" <+> var'
 
 instance PrettyCode InstrPrealloc where
   ppCode InstrPrealloc {..} = do
@@ -217,7 +217,7 @@ instance PrettyCode InstrBranch where
     val <- ppCode _instrBranchValue
     br1 <- ppCodeCode _instrBranchTrue
     br2 <- ppCodeCode _instrBranchFalse
-    var <- ppBranchVar _instrBranchVar
+    var <- ppOutVar _instrBranchOutVar
     return $
       primitive Str.br
         <+> val
@@ -249,7 +249,7 @@ instance PrettyCode InstrCase where
     val <- ppCode _instrCaseValue
     brs <- mapM ppCode _instrCaseBranches
     def <- maybe (return Nothing) (fmap Just . ppDefaultBranch) _instrCaseDefault
-    var <- ppBranchVar _instrCaseVar
+    var <- ppOutVar _instrCaseOutVar
     let brs' = brs ++ catMaybes [def]
     return $ primitive Str.case_ <> brackets ind <+> val <> var <+> braces' (vsep brs')
 

src/Juvix/Compiler/Reg/Transformation/SSA.hs

@@ -29,43 +29,63 @@ computeFunctionSSA =
       where
         instr' = overValueRefs (adjustVarRef mp) 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)
     combine instr mps = case instr of
       Branch InstrBranch {..} -> case mps of
         mp1 :| mp2 : []
-          | isNothing _instrBranchVar || idx1 == idx2 ->
+          | isNothing _instrBranchOutVar ->
               (mp, instr)
+          | idx1 == idx2 ->
+              ( mp,
+                Branch
+                  InstrBranch
+                    { _instrBranchOutVar = Just $ mkVarRef VarGroupLocal idx1,
+                      ..
+                    }
+              )
           | otherwise ->
               ( mp',
                 Branch
                   InstrBranch
                     { _instrBranchTrue = assignInBranch _instrBranchTrue idx' idx1,
                       _instrBranchFalse = assignInBranch _instrBranchFalse idx' idx2,
+                      _instrBranchOutVar = Just $ mkVarRef VarGroupLocal idx',
                       ..
                     }
               )
           where
-            var = fromJust _instrBranchVar
+            var = fromJust _instrBranchOutVar
             idx1 = IndexMap.lookup mp1 var
             idx2 = IndexMap.lookup mp2 var
             mp = IndexMap.combine mp1 mp2
             (idx', mp') = IndexMap.assign mp var
         _ -> impossible
       Case InstrCase {..} -> case mps of
         mp0 :| mps'
-          | isNothing _instrCaseVar || all (== head idxs) (NonEmpty.tail idxs) ->
+          | isNothing _instrCaseOutVar ->
               (mp, instr)
+          | all (== head idxs) (NonEmpty.tail idxs) ->
+              ( mp,
+                Case
+                  InstrCase
+                    { _instrCaseOutVar = Just $ mkVarRef VarGroupLocal (head idxs),
+                      ..
+                    }
+              )
           | otherwise ->
               ( mp',
                 Case
                   InstrCase
                     { _instrCaseBranches = brs',
                       _instrCaseDefault = def',
+                      _instrCaseOutVar = Just $ mkVarRef VarGroupLocal idx',
                       ..
                     }
               )
           where
-            var = fromJust _instrCaseVar
+            var = fromJust _instrCaseOutVar
             idxs = fmap (flip IndexMap.lookup var) mps
             mp = foldr IndexMap.combine mp0 mps'
             (idx', mp') = IndexMap.assign mp var

src/Juvix/Compiler/Reg/Translation/FromAsm.hs

@@ -268,7 +268,7 @@ fromAsmBranch fi isTail si Asm.CmdBranch {} codeTrue codeFalse =
         { _instrBranchValue = VRef $ mkVarRef VarGroupLocal topIdx,
           _instrBranchTrue = codeTrue,
           _instrBranchFalse = codeFalse,
-          _instrBranchVar = if isTail then Nothing else Just $ mkVarRef VarGroupLocal topIdx
+          _instrBranchOutVar = if isTail then Nothing else Just $ mkVarRef VarGroupLocal topIdx
         }
   where
     topIdx :: Int
@@ -307,7 +307,7 @@ fromAsmCase fi tab isTail si Asm.CmdCase {..} brs def =
               _cmdCaseBranches
               brs,
           _instrCaseDefault = def,
-          _instrCaseVar = if isTail then Nothing else Just $ mkVarRef VarGroupLocal topIdx
+          _instrCaseOutVar = if isTail then Nothing else Just $ mkVarRef VarGroupLocal topIdx
         }
   where
     topIdx = fromJust (fi ^. Asm.functionExtra) ^. Asm.functionMaxTempStackHeight + si ^. Asm.stackInfoValueStackHeight - 1

src/Juvix/Compiler/Reg/Translation/FromSource.hs

@@ -261,11 +261,11 @@ liveVars = do
   P.try (comma >> symbol "live:")
   parens (P.sepBy varRef comma)
 
-branchVar ::
+outVar ::
   (Members '[Reader ParserSig, InfoTableBuilder, State LocalParams] r) =>
   ParsecS r VarRef
-branchVar = do
-  P.try (comma >> symbol "var:")
+outVar = do
+  P.try (comma >> symbol "out:")
   varRef
 
 parseArgs ::
@@ -358,7 +358,7 @@ instrBranch ::
 instrBranch = do
   kw kwBr
   val <- value
-  var <- optional branchVar
+  var <- optional outVar
   (br1, br2) <- braces $ do
     symbol "true:"
     br1 <- braces parseCode
@@ -372,7 +372,7 @@ instrBranch = do
       { _instrBranchValue = val,
         _instrBranchTrue = br1,
         _instrBranchFalse = br2,
-        _instrBranchVar = var
+        _instrBranchOutVar = var
       }
 
 instrCase ::
@@ -382,7 +382,7 @@ instrCase = do
   kw kwCase
   sym <- brackets (indSymbol @Code @() @VarRef)
   val <- value
-  var <- optional branchVar
+  var <- optional outVar
   lbrace
   brs <- many caseBranch
   def <- optional defaultBranch
@@ -394,7 +394,7 @@ instrCase = do
         _instrCaseIndRep = IndRepStandard,
         _instrCaseBranches = brs,
         _instrCaseDefault = def,
-        _instrCaseVar = var
+        _instrCaseOutVar = var
       }
 
 caseBranch ::

tests/Reg/positive/test008.jvr

@@ -10,7 +10,7 @@ function main() : * {
   tmp[0] = 3;
   tmp[1] = 0;
   tmp[0] = lt tmp[1] tmp[0];
-  br tmp[0], var: tmp[0] {
+  br tmp[0], out: tmp[0] {
     true: {
       tmp[0] = 1;
     };
@@ -21,15 +21,15 @@ function main() : * {
   tmp[1] = 1;
   tmp[2] = 2;
   tmp[1] = le tmp[2] tmp[1];
-  br tmp[1], var: tmp[1] {
+  br tmp[1], out: tmp[1] {
     true: {
       tmp[1] = call loop (), live: (tmp[0]);
     };
     false: {
       tmp[1] = 7;
       tmp[2] = 8;
       tmp[1] = le tmp[2] tmp[1];
-      br tmp[1], var: tmp[1] {
+      br tmp[1], out: tmp[1] {
         true: {
           tmp[1] = call loop (), live: (tmp[0]);
         };

tests/Reg/positive/test010.jvr

@@ -6,7 +6,7 @@ function sum(integer) : integer {
   tmp[0] = arg[0];
   tmp[1] = 0;
   tmp[0] = eq tmp[1] tmp[0];
-  br tmp[0], var: tmp[0] {
+  br tmp[0], out: tmp[0] {
     true: {
       tmp[0] = 0;
     };

tests/Reg/positive/test031.jvr

@@ -56,7 +56,7 @@ function f(tree) : integer {
           {
             tmp[2] = tmp[5];
             tmp[5] = tmp[1];
-            case[tree] tmp[5], var: tmp[5] {
+            case[tree] tmp[5], out: tmp[5] {
               leaf: {
                 nop;
                 tmp[5] = 3;
@@ -79,7 +79,7 @@ function f(tree) : integer {
             {
               tmp[3] = tmp[5];
               tmp[5] = tmp[2];
-              case[tree] tmp[5], var: tmp[5] {
+              case[tree] tmp[5], out: tmp[5] {
                 node: {
                   {
                     tmp[4] = tmp[5];