Refactor the `MatchResult` type in the desugarer

This way, it does a better job of proving whether or not the fail operator is used.

3 files changed, 110 insertions, 87 deletions

diff --git a/compiler/GHC/HsToCore/Match.hs b/compiler/GHC/HsToCore/Match.hs
index c847bca068..dbdd24cbac 100644
--- a/compiler/GHC/HsToCore/Match.hs
+++ b/compiler/GHC/HsToCore/Match.hs
@@ -248,7 +248,7 @@ match (v:vs) ty eqns    -- Eqns *can* be empty
 matchEmpty :: MatchId -> Type -> DsM (NonEmpty MatchResult)
 -- See Note [Empty case expressions]
 matchEmpty var res_ty
-  = return [MatchResult CanFail mk_seq]
+  = return [MR_Fallible mk_seq]
   where
     mk_seq fail = return $ mkWildCase (Var var) (idType var) res_ty
                                       [(DEFAULT, [], fail)]
diff --git a/compiler/GHC/HsToCore/Monad.hs b/compiler/GHC/HsToCore/Monad.hs
index d09473798a..25d1fcfeb4 100644
--- a/compiler/GHC/HsToCore/Monad.hs
+++ b/compiler/GHC/HsToCore/Monad.hs
@@ -6,10 +6,14 @@
 Monadery used in desugaring
 -}
 
-{-# LANGUAGE FlexibleInstances, FlexibleContexts #-}
-{-# OPTIONS_GHC -fno-warn-orphans #-}  -- instance MonadThings is necessarily an orphan
+{-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE ViewPatterns #-}
 
+{-# OPTIONS_GHC -fno-warn-orphans #-}  -- instance MonadThings is necessarily an orphan
+
 module GHC.HsToCore.Monad (
         DsM, mapM, mapAndUnzipM,
         initDs, initDsTc, initTcDsForSolver, initDsWithModGuts, fixDs,
@@ -42,8 +46,7 @@ module GHC.HsToCore.Monad (
 
         -- Data types
         DsMatchContext(..),
-        EquationInfo(..), MatchResult(..), DsWrapper, idDsWrapper,
-        CanItFail(..), orFail,
+        EquationInfo(..), MatchResult'(..), MatchResult, runMatchResult, DsWrapper, idDsWrapper,
 
         -- Levity polymorphism
         dsNoLevPoly, dsNoLevPolyExpr, dsWhenNoErrs,
@@ -134,21 +137,37 @@ idDsWrapper e = e
 --      \fail. wrap (case vs of { pats -> rhs fail })
 -- where vs are not bound by wrap
 
-
--- A MatchResult is an expression with a hole in it
-data MatchResult
-  = MatchResult
-        CanItFail       -- Tells whether the failure expression is used
-        (CoreExpr -> DsM CoreExpr)
-                        -- Takes a expression to plug in at the
-                        -- failure point(s). The expression should
-                        -- be duplicatable!
-
-data CanItFail = CanFail | CantFail
-
-orFail :: CanItFail -> CanItFail -> CanItFail
-orFail CantFail CantFail = CantFail
-orFail _        _        = CanFail
+-- | This is a value of type a with potentially a CoreExpr-shaped hole in it.
+-- This is used to deal with cases where we are potentially handling pattern
+-- match failure, and want to later specify how failure is handled.
+data MatchResult' a
+  -- | We represent the case where there is no hole without a function from
+  -- 'CoreExpr', like this, because sometimes we have nothing to put in the
+  -- hole and so want to be sure there is in fact no hole.
+  = MR_Infallible (DsM a)
+  | MR_Fallible (CoreExpr -> DsM a)
+  deriving (Functor)
+
+-- | Product is an "or" on falliblity---the combined match result is infallible
+-- only if the left and right argument match results both were.
+--
+-- This is useful for combining a bunch of alternatives together and then
+-- getting the overall falliblity of the entire group. See 'mkDataConCase' for
+-- an example.
+instance Applicative MatchResult' where
+  pure v = MR_Infallible (pure v)
+  MR_Infallible f <*> MR_Infallible x = MR_Infallible (f <*> x)
+  f <*> x = MR_Fallible $ \fail -> runMatchResult fail f <*> runMatchResult fail x
+
+-- This is a CoreExpr with potentially a CoreExpr hole in it, which is the most common case.
+type MatchResult = MatchResult' CoreExpr
+
+-- Given a fail expression to use, and a MatchResult, compute the filled CoreExpr whether
+-- the MatchResult was failable or not.
+runMatchResult :: CoreExpr -> MatchResult' a -> DsM a
+runMatchResult fail = \case
+  MR_Infallible body -> body
+  MR_Fallible body_fn -> body_fn fail
 
 {-
 ************************************************************************
diff --git a/compiler/GHC/HsToCore/Utils.hs b/compiler/GHC/HsToCore/Utils.hs
index 3f0637f350..74c31f0a0f 100644
--- a/compiler/GHC/HsToCore/Utils.hs
+++ b/compiler/GHC/HsToCore/Utils.hs
@@ -10,6 +10,7 @@ This module exports some utility functions of no great interest.
 
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE ViewPatterns #-}
 
@@ -18,7 +19,7 @@ module GHC.HsToCore.Utils (
         EquationInfo(..),
         firstPat, shiftEqns,
 
-        MatchResult(..), CanItFail(..), CaseAlt(..),
+        MatchResult'(..), MatchResult, CaseAlt(..),
         cantFailMatchResult, alwaysFailMatchResult,
         extractMatchResult, combineMatchResults,
         adjustMatchResult,  adjustMatchResultDs,
@@ -85,6 +86,7 @@ import GHC.Tc.Types.Evidence
 
 import Control.Monad    ( zipWithM )
 import Data.List.NonEmpty (NonEmpty(..))
+import Data.Maybe (maybeToList)
 import qualified Data.List.NonEmpty as NEL
 
 {-
@@ -195,45 +197,50 @@ shiftEqns = fmap $ \eqn -> eqn { eqn_pats = tail (eqn_pats eqn) }
 -- Functions on MatchResults
 
 matchCanFail :: MatchResult -> Bool
-matchCanFail (MatchResult CanFail _)  = True
-matchCanFail (MatchResult CantFail _) = False
+matchCanFail (MR_Fallible {})  = True
+matchCanFail (MR_Infallible {}) = False
 
 alwaysFailMatchResult :: MatchResult
-alwaysFailMatchResult = MatchResult CanFail (\fail -> return fail)
+alwaysFailMatchResult = MR_Fallible $ \fail -> return fail
 
 cantFailMatchResult :: CoreExpr -> MatchResult
-cantFailMatchResult expr = MatchResult CantFail (\_ -> return expr)
+cantFailMatchResult expr = MR_Infallible $ return expr
 
 extractMatchResult :: MatchResult -> CoreExpr -> DsM CoreExpr
-extractMatchResult (MatchResult CantFail match_fn) _
-  = match_fn (error "It can't fail!")
-
-extractMatchResult (MatchResult CanFail match_fn) fail_expr = do
+extractMatchResult (MR_Infallible match_fn) _
+  = match_fn
+extractMatchResult (MR_Fallible match_fn) fail_expr = do
     (fail_bind, if_it_fails) <- mkFailurePair fail_expr
     body <- match_fn if_it_fails
     return (mkCoreLet fail_bind body)
 
-
 combineMatchResults :: MatchResult -> MatchResult -> MatchResult
-combineMatchResults (MatchResult CanFail      body_fn1)
-                    (MatchResult can_it_fail2 body_fn2)
-  = MatchResult can_it_fail2 body_fn
-  where
-    body_fn fail = do body2 <- body_fn2 fail
-                      (fail_bind, duplicatable_expr) <- mkFailurePair body2
-                      body1 <- body_fn1 duplicatable_expr
-                      return (Let fail_bind body1)
-
-combineMatchResults match_result1@(MatchResult CantFail _) _
+combineMatchResults (MR_Fallible body_fn1)
+                    (MR_Fallible body_fn2)
+  = MR_Fallible $ \fail -> do
+      body2 <- body_fn2 fail
+      (fail_bind, duplicatable_expr) <- mkFailurePair body2
+      body1 <- body_fn1 duplicatable_expr
+      return (Let fail_bind body1)
+combineMatchResults (MR_Fallible body_fn1)
+                    (MR_Infallible body_fn2)
+  = MR_Infallible $ do
+      body2 <- body_fn2
+      (fail_bind, duplicatable_expr) <- mkFailurePair body2
+      body1 <- body_fn1 duplicatable_expr
+      return (Let fail_bind body1)
+combineMatchResults match_result1@(MR_Infallible _) _
   = match_result1
 
-adjustMatchResult :: DsWrapper -> MatchResult -> MatchResult
-adjustMatchResult encl_fn (MatchResult can_it_fail body_fn)
-  = MatchResult can_it_fail (\fail -> encl_fn <$> body_fn fail)
+adjustMatchResult :: (a -> b) -> MatchResult' a -> MatchResult' b
+adjustMatchResult = fmap
 
-adjustMatchResultDs :: (CoreExpr -> DsM CoreExpr) -> MatchResult -> MatchResult
-adjustMatchResultDs encl_fn (MatchResult can_it_fail body_fn)
-  = MatchResult can_it_fail (\fail -> encl_fn =<< body_fn fail)
+adjustMatchResultDs :: (a -> DsM b) -> MatchResult' a -> MatchResult' b
+adjustMatchResultDs encl_fn = \case
+  MR_Infallible body_fn -> MR_Infallible $
+    encl_fn =<< body_fn
+  MR_Fallible body_fn -> MR_Fallible $ \fail ->
+    encl_fn =<< body_fn fail
 
 wrapBinds :: [(Var,Var)] -> CoreExpr -> CoreExpr
 wrapBinds [] e = e
@@ -261,25 +268,25 @@ mkEvalMatchResult var ty
   = adjustMatchResult (\e -> Case (Var var) var ty [(DEFAULT, [], e)])
 
 mkGuardedMatchResult :: CoreExpr -> MatchResult -> MatchResult
-mkGuardedMatchResult pred_expr (MatchResult _ body_fn)
-  = MatchResult CanFail (\fail -> do body <- body_fn fail
-                                     return (mkIfThenElse pred_expr body fail))
+mkGuardedMatchResult pred_expr mr = MR_Fallible $ \fail -> do
+  body <- runMatchResult fail mr
+  return (mkIfThenElse pred_expr body fail)
 
 mkCoPrimCaseMatchResult :: Id                  -- Scrutinee
                         -> Type                      -- Type of the case
                         -> [(Literal, MatchResult)]  -- Alternatives
                         -> MatchResult               -- Literals are all unlifted
 mkCoPrimCaseMatchResult var ty match_alts
-  = MatchResult CanFail mk_case
+  = MR_Fallible mk_case
   where
     mk_case fail = do
         alts <- mapM (mk_alt fail) sorted_alts
         return (Case (Var var) var ty ((DEFAULT, [], fail) : alts))
 
     sorted_alts = sortWith fst match_alts       -- Right order for a Case
-    mk_alt fail (lit, MatchResult _ body_fn)
+    mk_alt fail (lit, mr)
        = ASSERT( not (litIsLifted lit) )
-         do body <- body_fn fail
+         do body <- runMatchResult fail mr
             return (LitAlt lit, [], body)
 
 data CaseAlt a = MkCaseAlt{ alt_pat :: a,
@@ -315,14 +322,13 @@ mkCoAlgCaseMatchResult var ty match_alts
     newtype_rhs = unwrapNewTypeBody tc ty_args (Var var)
 
 mkCoSynCaseMatchResult :: Id -> Type -> CaseAlt PatSyn -> MatchResult
-mkCoSynCaseMatchResult var ty alt = MatchResult CanFail $ mkPatSynCase var ty alt
+mkCoSynCaseMatchResult var ty alt = MR_Fallible $ mkPatSynCase var ty alt
 
 mkPatSynCase :: Id -> Type -> CaseAlt PatSyn -> CoreExpr -> DsM CoreExpr
 mkPatSynCase var ty alt fail = do
     matcher <- dsLExpr $ mkLHsWrap wrapper $
                          nlHsTyApp matcher [getRuntimeRep ty, ty]
-    let MatchResult _ mkCont = match_result
-    cont <- mkCoreLams bndrs <$> mkCont fail
+    cont <- mkCoreLams bndrs <$> runMatchResult fail match_result
     return $ mkCoreAppsDs (text "patsyn" <+> ppr var) matcher [Var var, ensure_unstrict cont, Lam voidArgId fail]
   where
     MkCaseAlt{ alt_pat = psyn,
@@ -337,48 +343,46 @@ mkPatSynCase var ty alt fail = do
                          | otherwise      = cont
 
 mkDataConCase :: Id -> Type -> NonEmpty (CaseAlt DataCon) -> MatchResult
-mkDataConCase var ty alts@(alt1 :| _) = MatchResult fail_flag mk_case
+mkDataConCase var ty alts@(alt1 :| _)
+    = liftA2 mk_case mk_default mk_alts
+    -- The liftA2 combines the failability of all the alternatives and the default
   where
     con1          = alt_pat alt1
     tycon         = dataConTyCon con1
     data_cons     = tyConDataCons tycon
-    match_results = fmap alt_result alts
 
-    sorted_alts :: NonEmpty (CaseAlt DataCon)
-    sorted_alts  = NEL.sortWith (dataConTag . alt_pat) alts
+    sorted_alts :: [ CaseAlt DataCon ]
+    sorted_alts  = sortWith (dataConTag . alt_pat) $ NEL.toList alts
 
     var_ty       = idType var
     (_, ty_args) = tcSplitTyConApp var_ty -- Don't look through newtypes
                                           -- (not that splitTyConApp does, these days)
 
-    mk_case :: CoreExpr -> DsM CoreExpr
-    mk_case fail = do
-        alts <- mapM (mk_alt fail) sorted_alts
-        return $ mkWildCase (Var var) (idType var) ty (mk_default fail ++ NEL.toList alts)
-
-    mk_alt :: CoreExpr -> CaseAlt DataCon -> DsM CoreAlt
-    mk_alt fail MkCaseAlt{ alt_pat = con,
-                           alt_bndrs = args,
-                           alt_result = MatchResult _ body_fn }
-      = do { body <- body_fn fail
-           ; case dataConBoxer con of {
-                Nothing -> return (DataAlt con, args, body) ;
-                Just (DCB boxer) ->
-        do { us <- newUniqueSupply
-           ; let (rep_ids, binds) = initUs_ us (boxer ty_args args)
-           ; return (DataAlt con, rep_ids, mkLets binds body) } } }
-
-    mk_default :: CoreExpr -> [CoreAlt]
-    mk_default fail | exhaustive_case = []
-                    | otherwise       = [(DEFAULT, [], fail)]
-
-    fail_flag :: CanItFail
-    fail_flag | exhaustive_case
-              = foldr orFail CantFail [can_it_fail | MatchResult can_it_fail _ <- NEL.toList match_results]
-              | otherwise
-              = CanFail
-
-    mentioned_constructors = mkUniqSet $ map alt_pat $ NEL.toList alts
+    mk_case :: Maybe CoreAlt -> [CoreAlt] -> CoreExpr
+    mk_case def alts = mkWildCase (Var var) (idType var) ty $
+      maybeToList def ++ alts
+
+    mk_alts :: MatchResult' [CoreAlt]
+    mk_alts = traverse mk_alt sorted_alts
+
+    mk_alt :: CaseAlt DataCon -> MatchResult' CoreAlt
+    mk_alt MkCaseAlt { alt_pat = con
+                     , alt_bndrs = args
+                     , alt_result = match_result } =
+      flip adjustMatchResultDs match_result $ \body -> do
+        case dataConBoxer con of
+          Nothing -> return (DataAlt con, args, body)
+          Just (DCB boxer) -> do
+            us <- newUniqueSupply
+            let (rep_ids, binds) = initUs_ us (boxer ty_args args)
+            return (DataAlt con, rep_ids, mkLets binds body)
+
+    mk_default :: MatchResult' (Maybe CoreAlt)
+    mk_default
+      | exhaustive_case = MR_Infallible $ return Nothing
+      | otherwise       = MR_Fallible $ \fail -> return $ Just (DEFAULT, [], fail)
+
+    mentioned_constructors = mkUniqSet $ map alt_pat sorted_alts
     un_mentioned_constructors
         = mkUniqSet data_cons `minusUniqSet` mentioned_constructors
     exhaustive_case = isEmptyUniqSet un_mentioned_constructors