Make SpecConstr work right for nullary constructors
[ghc-hetmet.git] / compiler / specialise / SpecConstr.lhs
index e5583e1..10010cc 100644 (file)
@@ -19,6 +19,7 @@ import PprCore                ( pprRules )
 import WwLib           ( mkWorkerArgs )
 import DataCon         ( dataConRepArity, dataConUnivTyVars )
 import Type            ( Type, tyConAppArgs )
+import Coercion                ( coercionKind )
 import Rules           ( matchN )
 import Id              ( Id, idName, idType, isDataConWorkId_maybe, 
                          mkUserLocal, mkSysLocal, idUnfolding, isLocalId )
@@ -298,6 +299,24 @@ may avoid allocating it altogether.  Just like for constructors.
 
 Looks cool, but probably rare...but it might be easy to implement.
 
+
+Note [SpecConstr for casts]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider 
+    data family T a :: *
+    data instance T Int = T Int
+
+    foo n = ...
+       where
+         go (T 0) = 0
+         go (T n) = go (T (n-1))
+
+The recursive call ends up looking like 
+       go (T (I# ...) `cast` g)
+So we want to spot the construtor application inside the cast.
+That's why we have the Cast case in argToPat
+
+
 -----------------------------------------------------
                Stuff not yet handled
 -----------------------------------------------------
@@ -466,14 +485,19 @@ extendCaseBndrs env case_bndr scrut con alt_bndrs
                                [(b,how_bound) | b <- case_bndr:alt_bndrs] }
 
        -- Record RecArg for the components iff the scrutinee is RecArg
+       -- I think the only reason for this is to keep the usage envt small
+       -- so is it worth it at all?
        --      [This comment looks plain wrong to me, so I'm ignoring it
        --           "Also forget if the scrutinee is a RecArg, because we're
        --           now in the branch of a case, and we don't want to
        --           record a non-scrutinee use of v if we have
        --              case v of { (a,b) -> ...(f v)... }" ]
-    how_bound = case scrut of
-                 Var v -> lookupVarEnv cur_scope v `orElse` Other
-                 other -> Other
+    how_bound = get_how scrut
+       where
+           get_how (Var v)    = lookupVarEnv cur_scope v `orElse` Other
+           get_how (Cast e _) = get_how e
+           get_how (Note _ e) = get_how e
+           get_how other      = Other
 
     extend_data_con data_con = 
       extendCons env1 scrut case_bndr (CV con vanilla_args)
@@ -547,9 +571,10 @@ data ArgOcc = NoOcc        -- Doesn't occur at all; or a type argument
 
 {-     Note  [ScrutOcc]
 
-An occurrence of ScrutOcc indicates that the thing is *only* taken apart or applied.
+An occurrence of ScrutOcc indicates that the thing, or a `cast` version of the thing,
+is *only* taken apart or applied.
 
-  Functions, litersl: ScrutOcc emptyUFM
+  Functions, literal: ScrutOcc emptyUFM
   Data constructors:  ScrutOcc subs,
 
 where (subs :: UniqFM [ArgOcc]) gives usage of the *pattern-bound* components,
@@ -563,14 +588,20 @@ A pattern binds b, x::a, y::b, z::b->a, but not 'a'!
 -}
 
 instance Outputable ArgOcc where
-  ppr (ScrutOcc xs) = ptext SLIT("scrut-occ") <> parens (ppr xs)
+  ppr (ScrutOcc xs) = ptext SLIT("scrut-occ") <> ppr xs
   ppr UnkOcc       = ptext SLIT("unk-occ")
   ppr BothOcc      = ptext SLIT("both-occ")
   ppr NoOcc                = ptext SLIT("no-occ")
 
+-- Experimentally, this vresion of combineOcc makes ScrutOcc "win", so
+-- that if the thing is scrutinised anywhere then we get to see that
+-- in the overall result, even if it's also used in a boxed way
+-- This might be too agressive; see Note [Reboxing]
 combineOcc NoOcc        occ           = occ
 combineOcc occ                  NoOcc         = occ
 combineOcc (ScrutOcc xs) (ScrutOcc ys) = ScrutOcc (plusUFM_C combineOccs xs ys)
+combineOcc occ           (ScrutOcc ys) = ScrutOcc ys
+combineOcc (ScrutOcc xs) occ          = ScrutOcc xs
 combineOcc UnkOcc        UnkOcc        = UnkOcc
 combineOcc _       _                  = BothOcc
 
@@ -663,9 +694,12 @@ scExpr env e@(App _ _)
 ----------------------
 scScrut :: ScEnv -> CoreExpr -> ArgOcc -> UniqSM (ScUsage, CoreExpr)
 -- Used for the scrutinee of a case, 
--- or the function of an application
-scScrut env e@(Var v) occ = returnUs (varUsage env v occ, e)
-scScrut env e        occ = scExpr env e
+-- or the function of an application.
+-- Remember to look through casts
+scScrut env e@(Var v)   occ = returnUs (varUsage env v occ, e)
+scScrut env (Cast e co) occ = do { (usg, e') <- scScrut env e occ
+                                ; returnUs (usg, Cast e' co) }
+scScrut env e          occ = scExpr env e
 
 
 ----------------------
@@ -725,21 +759,25 @@ specialise :: ScEnv
 
 specialise env fn bndrs body body_usg
   = do { let (_, bndr_occs) = lookupOccs body_usg bndrs
+             all_calls = lookupVarEnv (calls body_usg) fn `orElse` []
 
-       ; mb_calls <- mapM (callToPats (scope env) bndr_occs)
-                          (lookupVarEnv (calls body_usg) fn `orElse` [])
+       ; mb_pats <- mapM (callToPats (scope env) bndr_occs) all_calls
 
-       ; let good_calls :: [([Var], [CoreArg])]
-             good_calls = catMaybes mb_calls
+       ; let good_pats :: [([Var], [CoreArg])]
+             good_pats = catMaybes mb_pats
              in_scope = mkInScopeSet $ unionVarSets $
                         [ exprsFreeVars pats `delVarSetList` vs 
-                        | (vs,pats) <- good_calls ]
-             uniq_calls = nubBy (same_call in_scope) good_calls
-       ; mapAndUnzipUs (spec_one env fn (mkLams bndrs body)) 
-                       (uniq_calls `zip` [1..]) }
+                        | (vs,pats) <- good_pats ]
+             uniq_pats = nubBy (same_pat in_scope) good_pats
+       ; -- pprTrace "specialise" (vcat [ppr fn <+> ppr bndrs <+> ppr bndr_occs,
+         --                            text "calls" <+> ppr all_calls,
+         --                            text "good pats" <+> ppr good_pats,
+         --                            text "uniq pats" <+> ppr uniq_pats])  $
+         mapAndUnzipUs (spec_one env fn (mkLams bndrs body)) 
+                       (uniq_pats `zip` [1..]) }
   where
-       -- Two calls are the same if they match both ways
-    same_call in_scope (vs1,as1)(vs2,as2)
+       -- Two pats are the same if they match both ways
+    same_pat in_scope (vs1,as1)(vs2,as2)
         =  isJust (matchN in_scope vs1 as1 as2)
         && isJust (matchN in_scope vs2 as2 as1)
 
@@ -758,7 +796,8 @@ callToPats in_scope bndr_occs (con_env, args)
                -- Quantify over variables that are not in sccpe
                -- See Note [Shadowing] at the top
                
-       ; if or good_pats 
+       ; -- pprTrace "callToPats"  (ppr args $$ ppr prs $$ ppr bndr_occs) $
+         if or good_pats 
          then return (Just (qvars, pats))
          else return Nothing }
 
@@ -863,18 +902,6 @@ argToPat :: InScopeEnv                     -- What's in scope at the fn defn site
 argToPat in_scope con_env arg@(Type ty) arg_occ
   = return (False, arg)
 
-argToPat in_scope con_env (Var v) arg_occ
-  | not (isLocalId v) || v `elemVarEnv` in_scope
-  =    -- The recursive call passes a variable that 
-       -- is in scope at the function definition site
-       -- It's worth specialising on this if
-       --      (a) it's used in an interesting way in the body
-       --      (b) we know what its value is
-    if    (case arg_occ of { UnkOcc -> False; other -> True }) -- (a)
-       && isValueUnfolding (idUnfolding v)                     -- (b)
-    then return (True, Var v)
-    else wildCardPat (idType v)
-
 argToPat in_scope con_env (Let _ arg) arg_occ
   = argToPat in_scope con_env arg arg_occ
        -- Look through let expressions
@@ -882,6 +909,13 @@ argToPat in_scope con_env (Let _ arg) arg_occ
        -- Here we can specialise for f (\y -> ...)
        -- because the rule-matcher will look through the let.
 
+argToPat in_scope con_env (Cast arg co) arg_occ
+  = do { (interesting, arg') <- argToPat in_scope con_env arg arg_occ
+       ; if interesting then 
+               return (interesting, Cast arg' co)
+         else 
+               wildCardPat (snd (coercionKind co)) }
+
 argToPat in_scope con_env arg arg_occ
   | is_value_lam arg
   = return (True, arg)
@@ -891,26 +925,40 @@ argToPat in_scope con_env arg arg_occ
        | otherwise = is_value_lam e
     is_value_lam other = False
 
+  -- Check for a constructor application
+  -- NB: this *precedes* the Var case, so that we catch nullary constrs
 argToPat in_scope con_env arg arg_occ
   | Just (CV dc args) <- is_con_app_maybe con_env arg
   , case arg_occ of
        ScrutOcc _ -> True              -- Used only by case scrutinee
-       BothOcc    -> case arg of       -- Used by case scrut
-                       App {} -> True  -- ...and elsewhere...
+       BothOcc    -> case arg of       -- Used elsewhere
+                       App {} -> True  --     see Note [Reboxing]
                        other  -> False
        other      -> False     -- No point; the arg is not decomposed
   = do { args' <- argsToPats in_scope con_env (args `zip` conArgOccs arg_occ dc)
        ; return (True, mk_con_app dc (map snd args')) }
 
+  -- Check if the argument is a variable that 
+  -- is in scope at the function definition site
+  -- It's worth specialising on this if
+  --   (a) it's used in an interesting way in the body
+  --   (b) we know what its value is
 argToPat in_scope con_env (Var v) arg_occ
-  =    -- A variable bound inside the function. 
-       -- Don't make a wild-card, because we may usefully share
-       --      e.g.  f a = let x = ... in f (x,x)
-       -- NB: this case follows the lambda and con-app cases!!
-    return (False, Var v)
+  | not (isLocalId v) || v `elemVarEnv` in_scope,
+    case arg_occ of { UnkOcc -> False; other -> True },        -- (a)
+    isValueUnfolding (idUnfolding v)                   -- (b)
+  = return (True, Var v)
+
+  -- Check for a variable bound inside the function. 
+  -- Don't make a wild-card, because we may usefully share
+  --   e.g.  f a = let x = ... in f (x,x)
+  -- NB: this case follows the lambda and con-app cases!!
+argToPat in_scope con_env (Var v) arg_occ
+  = return (False, Var v)
 
--- The default case: make a wild-card
-argToPat in_scope con_env arg arg_occ = wildCardPat (exprType arg)
+  -- The default case: make a wild-card
+argToPat in_scope con_env arg arg_occ
+  = wildCardPat (exprType arg)
 
 wildCardPat :: Type -> UniqSM (Bool, CoreArg)
 wildCardPat ty = do { uniq <- getUniqueUs
@@ -929,33 +977,30 @@ argsToPats in_scope con_env args
 
 \begin{code}
 is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe ConValue
+is_con_app_maybe env (Lit lit)
+  = Just (CV (LitAlt lit) [])
+
+is_con_app_maybe env expr      -- Maybe it's a constructor application
+  | (Var fun, args) <- collectArgs expr,
+    Just con <- isDataConWorkId_maybe fun,
+    args `lengthAtLeast` dataConRepArity con
+       -- Might be > because the arity excludes type args
+  = Just (CV (DataAlt con) args)
+
 is_con_app_maybe env (Var v)
-  = case lookupVarEnv env v of
-       Just stuff -> Just stuff
-               -- You might think we could look in the idUnfolding here
+  | Just stuff <- lookupVarEnv env v
+  = Just stuff -- You might think we could look in the idUnfolding here
                -- but that doesn't take account of which branch of a 
                -- case we are in, which is the whole point
 
-       Nothing | isCheapUnfolding unf
-               -> is_con_app_maybe env (unfoldingTemplate unf)
-               where
-                 unf = idUnfolding v
-               -- However we do want to consult the unfolding 
-               -- as well, for let-bound constructors!
-
-       other  -> Nothing
-
-is_con_app_maybe env (Lit lit)
-  = Just (CV (LitAlt lit) [])
-
-is_con_app_maybe env expr
-  = case collectArgs expr of
-       (Var fun, args) | Just con <- isDataConWorkId_maybe fun,
-                         args `lengthAtLeast` dataConRepArity con
-               -- Might be > because the arity excludes type args
-                       -> Just (CV (DataAlt con) args)
+  | isCheapUnfolding unf
+  = is_con_app_maybe env (unfoldingTemplate unf)
+  where
+    unf = idUnfolding v
+       -- However we do want to consult the unfolding 
+       -- as well, for let-bound constructors!
 
-       other -> Nothing
+is_con_app_maybe env expr = Nothing
 
 mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
 mk_con_app (LitAlt lit)  []   = Lit lit