Make -fliberate-case work for GADTs
[ghc-hetmet.git] / ghc / compiler / specialise / SpecConstr.lhs
index 59fef91..74944da 100644 (file)
@@ -12,29 +12,32 @@ module SpecConstr(
 
 import CoreSyn
 import CoreLint                ( showPass, endPass )
-import CoreUtils       ( exprType, eqExpr )
+import CoreUtils       ( exprType, tcEqExpr, mkPiTypes )
 import CoreFVs                 ( exprsFreeVars )
-import DataCon         ( dataConRepArity )
-import Type            ( tyConAppArgs )
-import PprCore         ( pprCoreRules )
-import Id              ( Id, idName, idType, idSpecialisation,
-                         isDataConId_maybe,
+import CoreSubst       ( Subst, mkSubst, substExpr )
+import CoreTidy                ( tidyRules )
+import PprCore         ( pprRules )
+import WwLib           ( mkWorkerArgs )
+import DataCon         ( dataConRepArity, isVanillaDataCon )
+import Type            ( tyConAppArgs, tyVarsOfTypes )
+import Unify           ( coreRefineTys )
+import Id              ( Id, idName, idType, isDataConWorkId_maybe, 
                          mkUserLocal, mkSysLocal )
 import Var             ( Var )
 import VarEnv
 import VarSet
 import Name            ( nameOccName, nameSrcLoc )
-import Rules           ( addIdSpecialisations )
+import Rules           ( addIdSpecialisations, mkLocalRule, rulesOfBinds )
 import OccName         ( mkSpecOcc )
 import ErrUtils                ( dumpIfSet_dyn )
-import CmdLineOpts     ( DynFlags, DynFlag(..) )
-import Outputable
-
+import DynFlags                ( DynFlags, DynFlag(..) )
+import BasicTypes      ( Activation(..) )
 import Maybes          ( orElse )
-import Util            ( mapAccumL )
+import Util            ( mapAccumL, lengthAtLeast, notNull )
 import List            ( nubBy, partition )
 import UniqSupply
 import Outputable
+import FastString
 \end{code}
 
 -----------------------------------------------------
@@ -179,7 +182,7 @@ specConstrProgram dflags us binds
        endPass dflags "SpecConstr" Opt_D_dump_spec binds'
 
        dumpIfSet_dyn dflags Opt_D_dump_rules "Top-level specialisations"
-                 (vcat (map dump_specs (concat (map bindersOf binds'))))
+                 (pprRules (tidyRules emptyTidyEnv (rulesOfBinds binds')))
 
        return binds'
   where
@@ -187,8 +190,6 @@ specConstrProgram dflags us binds
     go env (bind:binds) = scBind env bind      `thenUs` \ (env', _, bind') ->
                          go env' binds         `thenUs` \ binds' ->
                          returnUs (bind' : binds')
-
-dump_specs var = pprCoreRules var (idSpecialisation var)
 \end{code}
 
 
@@ -205,10 +206,17 @@ data ScEnv = SCE { scope :: VarEnv HowBound,
                   cons  :: ConstrEnv
             }
 
-type ConstrEnv = IdEnv (AltCon, [CoreArg])
+type ConstrEnv = IdEnv ConValue
+data ConValue  = CV AltCon [CoreArg]
        -- Variables known to be bound to a constructor
        -- in a particular case alternative
 
+refineConstrEnv :: Subst -> ConstrEnv -> ConstrEnv
+-- The substitution is a type substitution only
+refineConstrEnv subst env = mapVarEnv refine_con_value env
+  where
+    refine_con_value (CV con args) = CV con (map (substExpr subst) args)
+
 emptyScEnv = SCE { scope = emptyVarEnv, cons = emptyVarEnv }
 
 data HowBound = RecFun         -- These are the recursive functions for which 
@@ -222,6 +230,11 @@ data HowBound = RecFun             -- These are the recursive functions for which
                                -- passed as a parameter and what is in scope at the 
                                -- function definition site
 
+instance Outputable HowBound where
+  ppr RecFun = text "RecFun"
+  ppr RecArg = text "RecArg"
+  ppr Other = text "Other"
+
 lookupScopeEnv env v = lookupVarEnv (scope env) v
 
 extendBndrs env bndrs = env { scope = extendVarEnvList (scope env) [(b,Other) | b <- bndrs] }
@@ -235,24 +248,47 @@ extendCaseBndrs :: ScEnv -> Id -> CoreExpr -> AltCon -> [Var] -> ScEnv
 extendCaseBndrs env case_bndr scrut DEFAULT alt_bndrs
   = extendBndrs env (case_bndr : alt_bndrs)
 
-extendCaseBndrs env case_bndr scrut con alt_bndrs
-  = case scrut of
+extendCaseBndrs env case_bndr scrut con@(LitAlt lit) alt_bndrs
+  = ASSERT( null alt_bndrs ) extendAlt env case_bndr scrut (CV con []) []
+
+extendCaseBndrs env case_bndr scrut con@(DataAlt data_con) alt_bndrs
+  | isVanillaDataCon data_con
+  = extendAlt env case_bndr scrut (CV con vanilla_args) alt_bndrs
+    
+  | otherwise  -- GADT
+  = extendAlt env1 case_bndr scrut (CV con gadt_args) alt_bndrs
+  where
+    vanilla_args = map Type (tyConAppArgs (idType case_bndr)) ++
+                  map varToCoreExpr alt_bndrs
+
+    gadt_args = map (substExpr subst . varToCoreExpr) alt_bndrs
+
+    (alt_tvs, _) = span isTyVar alt_bndrs
+    Just (tv_subst, is_local) = coreRefineTys data_con alt_tvs (idType case_bndr)
+    subst = mkSubst in_scope tv_subst emptyVarEnv      -- No Id substitition
+    in_scope = mkInScopeSet (tyVarsOfTypes (varEnvElts tv_subst))
+
+    env1 | is_local  = env
+        | otherwise = env { cons = refineConstrEnv subst (cons env) }
+
+
+
+extendAlt :: ScEnv -> Id -> CoreExpr -> ConValue -> [Var] -> ScEnv
+extendAlt env case_bndr scrut val alt_bndrs
+  = let 
+       env1 = SCE { scope = extendVarEnvList (scope env) [(b,Other) | b <- case_bndr : alt_bndrs],
+                   cons  = extendVarEnv     (cons  env) case_bndr val }
+    in
+    case scrut of
        Var v ->   -- Bind the scrutinee in the ConstrEnv if it's a variable
                   -- 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)... }
                 SCE { scope = extendVarEnv (scope env1) v Other,
-                      cons  = extendVarEnv (cons env1)  v (con,args) }
+                      cons  = extendVarEnv (cons env1)  v val }
        other -> env1
 
-  where
-    env1 = SCE { scope = extendVarEnvList (scope env) [(b,Other) | b <- case_bndr : alt_bndrs],
-                cons  = extendVarEnv     (cons  env) case_bndr (con,args) }
-
-    args = map Type (tyConAppArgs (idType case_bndr)) ++
-          map varToCoreExpr alt_bndrs
-
     -- When we encounter a recursive function binding
     -- f = \x y -> ...
     -- we want to extend the scope env with bindings 
@@ -329,11 +365,11 @@ scExpr env (Note n e) = scExpr env e      `thenUs` \ (usg,e') ->
 scExpr env (Lam b e)  = scExpr (extendBndr env b) e    `thenUs` \ (usg,e') ->
                        returnUs (usg, Lam b e')
 
-scExpr env (Case scrut b alts) 
+scExpr env (Case scrut b ty alts) 
   = sc_scrut scrut             `thenUs` \ (scrut_usg, scrut') ->
     mapAndUnzipUs sc_alt alts  `thenUs` \ (alts_usgs, alts') ->
     returnUs (combineUsages alts_usgs `combineUsage` scrut_usg,
-             Case scrut' b alts')
+             Case scrut' b ty alts')
   where
     sc_scrut e@(Var v) = returnUs (varUsage env v CaseScrut, e)
     sc_scrut e        = scExpr env e
@@ -369,19 +405,21 @@ scExpr env e@(App _ _)
 ----------------------
 scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
 scBind env (Rec [(fn,rhs)])
-  | not (null val_bndrs)
-  = scExpr env' body                   `thenUs` \ (usg, body') ->
+  | notNull val_bndrs
+  = scExpr env_fn_body body            `thenUs` \ (usg, body') ->
     let
        SCU { calls = calls, occs = occs } = usg
     in
     specialise env fn bndrs body usg   `thenUs` \ (rules, spec_prs) ->
-    returnUs (extendBndrs env bndrs,
+    returnUs (extendBndr env fn,       -- For the body of the letrec, just
+                                       -- extend the env with Other to record 
+                                       -- that it's in scope; no funny RecFun business
              SCU { calls = calls `delVarEnv` fn, occs = occs `delVarEnvList` val_bndrs},
              Rec ((fn `addIdSpecialisations` rules, mkLams bndrs body') : spec_prs))
   where
     (bndrs,body) = collectBinders rhs
     val_bndrs    = filter isId bndrs
-    env'        = extendRecBndr env fn bndrs
+    env_fn_body         = extendRecBndr env fn bndrs
 
 scBind env (Rec prs)
   = mapAndUnzipUs do_one prs   `thenUs` \ (usgs, prs') ->
@@ -424,18 +462,17 @@ specialise env fn bndrs body (SCU {calls=calls, occs=occs})
        good_calls :: [[CoreArg]]
        good_calls = [ pats
                     | (con_env, call_args) <- all_calls,
-                      length call_args >= n_bndrs,         -- App is saturated
+                      call_args `lengthAtLeast` n_bndrs,           -- App is saturated
                       let call = (bndrs `zip` call_args),
                       any (good_arg con_env occs) call,    -- At least one arg is a constr app
                       let (_, pats) = argsToPats con_env us call_args
                     ]
     in
-    pprTrace "specialise" (ppr all_calls $$ ppr good_calls) $
     mapAndUnzipUs (spec_one env fn (mkLams bndrs body)) 
                  (nubBy same_call good_calls `zip` [1..])
   where
     n_bndrs  = length bndrs
-    same_call as1 as2 = and (zipWith eqExpr as1 as2)
+    same_call as1 as2 = and (zipWith tcEqExpr as1 as2)
 
 ---------------------
 good_arg :: ConstrEnv -> IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
@@ -446,8 +483,7 @@ good_arg con_env arg_occs (bndr, arg)
 
 bndr_usg_ok :: IdEnv ArgOcc -> Var -> CoreArg -> Bool
 bndr_usg_ok arg_occs bndr arg
-  = pprTrace "bndr_ok" (ppr bndr <+> ppr (lookupVarEnv arg_occs bndr)) $
-    case lookupVarEnv arg_occs bndr of
+  = case lookupVarEnv arg_occs bndr of
        Just CaseScrut -> True                  -- Used only by case scrutiny
        Just Both      -> case arg of           -- Used by case and elsewhere
                            App _ _ -> True     -- so the arg should be an explicit con app
@@ -470,22 +506,21 @@ spec_one :: ScEnv
   Example
   
      In-scope: a, x::a   
-     f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) v (h v))...
-         [c is presumably bound by the (...) part]
+     f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) (x,v) (h w))...
+         [c::*, v::(b,c) are presumably bound by the (...) part]
   ==>
-     f_spec = /\ b c \ v::(a,(b,c)) -> 
-                 (...entire RHS of f...) (b,c) ((:) (a,(b,c)) v (h v))
+     f_spec = /\ b c \ v::(b,c) hw::[(a,(b,c))] ->
+                 (...entire RHS of f...) (b,c) ((:) (a,(b,c)) (x,v) hw)
   
-     RULE:  forall b c,
-                  y::[(a,(b,c))], 
-                  v::(a,(b,c)), 
-                  h::(a,(b,c))->[(a,(b,c))] .
+     RULE:  forall b::* c::*,          -- Note, *not* forall a, x
+                  v::(b,c),
+                  hw::[(a,(b,c))] .
   
-           f (b,c) ((:) (a,(b,c)) v (h v)) = f_spec b c v
+           f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
 -}
 
-spec_one env fn rhs (pats, n)
-  = getUniqueUs                        `thenUs` \ spec_uniq ->
+spec_one env fn rhs (pats, rule_number)
+  = getUniqueUs                `thenUs` \ spec_uniq ->
     let 
        fn_name      = idName fn
        fn_loc       = nameSrcLoc fn_name
@@ -493,16 +528,35 @@ spec_one env fn rhs (pats, n)
        pat_fvs      = varSetElems (exprsFreeVars pats)
        vars_to_bind = filter not_avail pat_fvs
        not_avail v  = not (v `elemVarEnv` scope env)
-               -- Put the type variables first just for tidiness
+               -- Put the type variables first; the type of a term
+               -- variable may mention a type variable
        (tvs, ids)   = partition isTyVar vars_to_bind
        bndrs        = tvs ++ ids
+       spec_body    = mkApps rhs pats
+       body_ty      = exprType spec_body
        
-       rule_name = _PK_ ("SC:" ++ showSDoc (ppr fn <> int n))
-       spec_rhs  = mkLams bndrs (mkApps rhs pats)
-       spec_id   = mkUserLocal spec_occ spec_uniq (exprType spec_rhs) fn_loc
-       rule      = Rule rule_name pat_fvs pats (mkVarApps (Var spec_id) bndrs)
+       (spec_lam_args, spec_call_args) = mkWorkerArgs bndrs body_ty
+               -- Usual w/w hack to avoid generating 
+               -- a spec_rhs of unlifted type and no args
+       
+       rule_name = mkFastString ("SC:" ++ showSDoc (ppr fn <> int rule_number))
+       spec_rhs  = mkLams spec_lam_args spec_body
+       spec_id   = mkUserLocal spec_occ spec_uniq (mkPiTypes spec_lam_args body_ty) fn_loc
+       rule_rhs  = mkVarApps (Var spec_id) spec_call_args
+       rule      = mkLocalRule rule_name specConstrActivation fn_name bndrs pats rule_rhs
     in
     returnUs (rule, (spec_id, spec_rhs))
+
+-- In which phase should the specialise-constructor rules be active?
+-- Originally I made them always-active, but Manuel found that
+-- this defeated some clever user-written rules.  So Plan B
+-- is to make them active only in Phase 0; after all, currently,
+-- the specConstr transformation is only run after the simplifier
+-- has reached Phase 0.  In general one would want it to be 
+-- flag-controllable, but for now I'm leaving it baked in
+--                                     [SLPJ Oct 01]
+specConstrActivation :: Activation
+specConstrActivation = ActiveAfter 0   -- Baked in; see comments above
 \end{code}
 
 %************************************************************************
@@ -521,12 +575,12 @@ they are constructor applications.
     -- placeholder variables.  For example:
     --    C a (D (f x) (g y))  ==>  C p1 (D p2 p3)
 
-argToPat   :: ConstrEnv -> UniqSupply -> CoreArg   -> (UniqSupply, CoreExpr)
+argToPat   :: ConstrEnv -> UniqSupply -> CoreArg -> (UniqSupply, CoreExpr)
 argToPat env us (Type ty) 
   = (us, Type ty)
 
 argToPat env us arg
-  | Just (dc,args) <- is_con_app_maybe env arg
+  | Just (CV dc args) <- is_con_app_maybe env arg
   = let
        (us',args') = argsToPats env us args
     in
@@ -536,7 +590,7 @@ argToPat env us (Var v)     -- Don't uniqify existing vars,
   = (us, Var v)                -- so that we can spot when we pass them twice
 
 argToPat env us arg
-  = (us1, Var (mkSysLocal SLIT("sc") (uniqFromSupply us2) (exprType arg)))
+  = (us1, Var (mkSysLocal FSLIT("sc") (uniqFromSupply us2) (exprType arg)))
   where
     (us1,us2) = splitUniqSupply us
 
@@ -546,7 +600,7 @@ argsToPats env us args = mapAccumL (argToPat env) us args
 
 
 \begin{code}
-is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe (AltCon, [CoreExpr])
+is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe ConValue
 is_con_app_maybe env (Var v)
   = lookupVarEnv env v
        -- You might think we could look in the idUnfolding here
@@ -554,14 +608,14 @@ is_con_app_maybe env (Var v)
        -- case we are in, which is the whole point
 
 is_con_app_maybe env (Lit lit)
-  = Just (LitAlt lit, [])
+  = Just (CV (LitAlt lit) [])
 
 is_con_app_maybe env expr
   = case collectArgs expr of
-       (Var fun, args) | Just con <- isDataConId_maybe fun,
-                         length args >= dataConRepArity con
+       (Var fun, args) | Just con <- isDataConWorkId_maybe fun,
+                         args `lengthAtLeast` dataConRepArity con
                -- Might be > because the arity excludes type args
-                       -> Just (DataAlt con,args)
+                       -> Just (CV (DataAlt con) args)
 
        other -> Nothing