Fix handling of known tag in case with default in simplifier (#564)

This was already mostly fixed on new-pipeline, except for the scenario where the tag matches none of the alternatives and no default is given (which should not occur?). This commit backports that fix, and handles the case where no default is present.

It also renames 'x' to 'tag' to not suggest that it has the same type and role as 'xs'.

Author: ivogabe

src/Data/Array/Accelerate/AST.hs

@@ -82,7 +82,7 @@ module Data.Array.Accelerate.AST (
   ALeftHandSide, ArrayVar, ArrayVars,
 
   -- ** Scalar expressions
-  ELeftHandSide, ExpVar, ExpVars, expVars,
+  ELeftHandSide, ExpVar, ExpVars, expVars, undefs,
   Fun, OpenFun(..),
   Exp, OpenExp(..),
   Boundary(..),
@@ -513,6 +513,10 @@ expVars TupRunit         = Nil
 expVars (TupRsingle var) = Evar var
 expVars (TupRpair v1 v2) = expVars v1 `Pair` expVars v2
 
+undefs :: TypeR t -> OpenExp env aenv t
+undefs (TupRsingle tp) = Undef tp
+undefs (TupRpair t1 t2) = undefs t1 `Pair` undefs t2
+undefs TupRunit = Nil
 
 -- | Vanilla open expressions using de Bruijn indices for variables ranging
 -- over tuples of scalars and arrays of tuples. All code, except Cond, is

src/Data/Array/Accelerate/Trafo/Simplify.hs

@@ -337,26 +337,34 @@ simplifyOpenExp env = first getAny . cvtE
            -> (Any, [(TAG, OpenExp env aenv b)])
            -> (Any, Maybe (OpenExp env aenv b))
            -> (Any, OpenExp env aenv b)
-    caseof x@(_,x') xs@(_,xs') md@(_,md')
-      | Const _ t   <- x'
-      = Stats.caseElim "known" (yes (fromJust $ lookup t xs'))
+    caseof tag@(_,tag') xs@(_,xs') md@(_,md')
+      | Const _ t   <- tag'
+      = Stats.caseElim "known" $ yes $ fromMaybe
+        -- If the tag matches no alternative, then return the default (if present)
+        -- or undef.
+        (fromMaybe (undefs tp) md')
+        $ lookup t xs'
       | Just d      <- md'
       , []          <- xs'
       = Stats.caseElim "redundant" (yes d)
       | Just d      <- md'
       , [(_,(_,u))] <- us
       , Just Refl   <- matchOpenExp d u
-      = Stats.caseDefault "merge" $ yes (Case x' (map snd vs) (Just u))
+      = Stats.caseDefault "merge" $ yes (Case tag' (map snd vs) (Just u))
       | Nothing     <- md'
       , []          <- vs
       , [(_,(_,u))] <- us
       = Stats.caseElim "overlap" (yes u)
       | Nothing     <- md'
       , [(_,(_,u))] <- us
-      = Stats.caseDefault "introduction" $ yes (Case x' (map snd vs) (Just u))
+      = Stats.caseDefault "introduction" $ yes (Case tag' (map snd vs) (Just u))
       | otherwise
-      = Case <$> x <*> xs <*> md
+      = Case <$> tag <*> xs <*> md
       where
+        tp
+          | Just d <- md' = expType d
+          | (_, x) : _ <- xs' = expType x
+          | otherwise = internalError "Empty case statement should not occur"
         (us,vs) = partition (\(n,_) -> n > 1)
                 $ Map.elems
                 . Map.fromListWith merge