Fix CaseCopyPropagation

grin/src/ParseGrin.hs

@@ -1,6 +1,6 @@
 {-# LANGUAGE TupleSections #-}
 
-module ParseGrin (parseGrin) where
+module ParseGrin (parseGrin, parseDef) where
 
 import Data.Void
 import Control.Applicative (empty)
@@ -112,3 +112,6 @@ grinModule = Program <$> some def <* sc <* eof
 
 parseGrin :: String -> String -> Either (ParseError Char Void) Exp
 parseGrin filename content = runParser grinModule filename content
+
+parseDef :: String -> Exp
+parseDef = either (error . show) id . runParser def ""

grin/src/Transformations/Optimising/CaseCopyPropagation.hs

@@ -4,8 +4,6 @@ module Transformations.Optimising.CaseCopyPropagation where
 import Control.Arrow
 import Grin
 import Data.Functor.Foldable
-import Control.Monad.Free
-import Control.Comonad.Cofree
 import Data.Monoid hiding (Alt)
 import Debug.Trace
 import Data.Maybe
@@ -85,6 +83,7 @@ collectInfo = para convert where
     ECaseF v ((mconcat . map snd) -> info) -> case info of
       Info returns cases -> case (allTheSame $ map snd returns) of
         Just (Just tag) -> Info mempty (([], tag):cases) -- The case is the same return values
+        Just Nothing    -> Info mempty cases
         Nothing         -> Info mempty cases
 
     AltF pat (ei, info) -> info
@@ -133,6 +132,21 @@ caseCopyPropagation e = apo builder (Build e info False) where
         i1 = zoom 2 i
         newVar = Var (n <> "'")
 
+    -- Exp: The last statement is a case
+    EBind lhs pat rhs@(ECase var@(Var n) alts)
+      | caseFocusOnStep (toStepE rhs) i1 ->
+          EBindF
+            (stepIn lhs)
+            pat
+            (Right $ Skip 3
+              (EBind (SBlock rhs) newVar (SReturn (ConstTagNode (fromJust (caseTagOnStep (toStepE rhs) i1)) [newVar])))
+              i1
+              True)
+      | otherwise -> EBindF (stepIn lhs) pat (stepIn rhs)
+      where
+        i1 = zoom 1 i
+        newVar = Var (n <> "'")
+
     EBind lhs pat rhs | isSimpleExp lhs -> EBindF (Left lhs)   pat (stepIn rhs)
     EBind lhs pat rhs                   -> EBindF (stepIn lhs) pat (stepIn rhs)
 

grin/test/Transformations/Optimising/CaseCopyPropagationSpec.hs

@@ -8,6 +8,7 @@ import Free
 import Grin
 import Test
 import Assertions
+import ParseGrin
 
 
 spec :: Spec
@@ -190,5 +191,33 @@ spec = do
 
     caseCopyPropagation before `sameAs` after
 
+  it "last expression is a case" $ do
+
+    let before = parseDef $ unlines
+          [ "sum l ="
+          , "  l2 <- eval l"
+          , "  case l2 of"
+          , "    (CNil) -> pure (CInt 0)"
+          , "    (CCons x xs) -> (CInt x') <- eval x"
+          , "                    (CInt s') <- sum xs"
+          , "                    ax' <- _prim_int_add x' s'"
+          , "                    pure (CInt ax')"
+          ]
+
+    let after = parseDef $ unlines
+          [ "sum l ="
+          , "  l2 <- eval l"
+          , "  l2' <- do"
+          , "    case l2 of"
+          , "      (CNil) -> pure 0"
+          , "      (CCons x xs) -> (CInt x') <- eval x"
+          , "                      (CInt s') <- sum xs"
+          , "                      ax' <- _prim_int_add x' s'"
+          , "                      pure ax'"
+          , "  pure (CInt l2')"
+          ]
+
+    caseCopyPropagation before `sameAs` after
+
 runTests :: IO ()
 runTests = hspec spec