Remove (most of) the FiniteMap wrapper
[ghc-hetmet.git] / compiler / specialise / SpecConstr.lhs
index e2eda2b..f214f0c 100644 (file)
@@ -31,11 +31,12 @@ import Coercion
 import Rules
 import Type            hiding( substTy )
 import Id
-import MkId            ( mkImpossibleExpr )
+import MkCore          ( mkImpossibleExpr )
 import Var
 import VarEnv
 import VarSet
 import Name
+import BasicTypes
 import DynFlags                ( DynFlags(..) )
 import StaticFlags     ( opt_PprStyle_Debug )
 import Maybes          ( orElse, catMaybes, isJust, isNothing )
@@ -47,15 +48,10 @@ import UniqSupply
 import Outputable
 import FastString
 import UniqFM
-import qualified LazyUniqFM as L
 import MonadUtils
 import Control.Monad   ( zipWithM )
 import Data.List
-#if __GLASGOW_HASKELL__ > 609
 import Data.Data        ( Data, Typeable )
-#else
-import Data.Generics    ( Data, Typeable )
-#endif
 \end{code}
 
 -----------------------------------------------------
@@ -389,6 +385,38 @@ But fspec doesn't have decent strictnes info.  As it happened,
 and hence f.  But now f's strictness is less than its arity, which
 breaks an invariant.
 
+Note [Forcing specialisation]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+With stream fusion and in other similar cases, we want to fully specialise
+some (but not necessarily all!) loops regardless of their size and the
+number of specialisations. We allow a library to specify this by annotating
+a type with ForceSpecConstr and then adding a parameter of that type to the
+loop. Here is a (simplified) example from the vector library:
+
+  data SPEC = SPEC | SPEC2
+  {-# ANN type SPEC ForceSpecConstr #-}
+
+  foldl :: (a -> b -> a) -> a -> Stream b -> a
+  {-# INLINE foldl #-}
+  foldl f z (Stream step s _) = foldl_loop SPEC z s
+    where
+      foldl_loop SPEC z s = case step s of
+                              Yield x s' -> foldl_loop SPEC (f z x) s'
+                              Skip       -> foldl_loop SPEC z s'
+                              Done       -> z
+
+SpecConstr will spot the SPEC parameter and always fully specialise
+foldl_loop. Note that we can't just annotate foldl_loop since it isn't a
+top-level function but even if we could, inlining etc. could easily drop the
+annotation. We also have to prevent the SPEC argument from being removed by
+w/w which is why SPEC is a sum type. This is all quite ugly; we ought to come
+up with a better design.
+
+ForceSpecConstr arguments are spotted in scExpr' and scTopBinds which then set
+force_spec to True when calling specLoop. This flag makes specLoop and
+specialise ignore specConstrCount and specConstrThreshold when deciding
+whether to specialise a function.
+
 -----------------------------------------------------
                Stuff not yet handled
 -----------------------------------------------------
@@ -523,7 +551,7 @@ data ScEnv = SCE { sc_size  :: Maybe Int,   -- Size threshold
                        -- Domain is OutIds (*after* applying the substitution)
                        -- Used even for top-level bindings (but not imported ones)
 
-                   sc_annotations :: L.UniqFM SpecConstrAnnotation
+                   sc_annotations :: UniqFM SpecConstrAnnotation
             }
 
 ---------------------
@@ -548,7 +576,7 @@ instance Outputable Value where
    ppr LambdaVal        = ptext (sLit "<Lambda>")
 
 ---------------------
-initScEnv :: DynFlags -> L.UniqFM SpecConstrAnnotation -> ScEnv
+initScEnv :: DynFlags -> UniqFM SpecConstrAnnotation -> ScEnv
 initScEnv dflags anns
   = SCE { sc_size = specConstrThreshold dflags,
          sc_count = specConstrCount dflags,
@@ -644,7 +672,7 @@ extendCaseBndrs env case_bndr con alt_bndrs
        --      Var v  -> extendValEnv env1 v cval
        --      _other -> env1
  where
-   zap v | isTyVar v = v               -- See NB2 above
+   zap v | isTyCoVar v = v             -- See NB2 above
          | otherwise = zapIdOccInfo v
    env1 = extendValEnv env case_bndr cval
    cval = case con of
@@ -657,7 +685,7 @@ extendCaseBndrs env case_bndr con alt_bndrs
 
 ignoreTyCon :: ScEnv -> TyCon -> Bool
 ignoreTyCon env tycon
-  = L.lookupUFM (sc_annotations env) tycon == Just NoSpecConstr
+  = lookupUFM (sc_annotations env) tycon == Just NoSpecConstr
 
 ignoreType :: ScEnv -> Type -> Bool
 ignoreType env ty
@@ -671,7 +699,7 @@ ignoreAltCon env (LitAlt lit) = ignoreType env (literalType lit)
 ignoreAltCon _   DEFAULT      = True
 
 forceSpecBndr :: ScEnv -> Var -> Bool
-forceSpecBndr env var = forceSpecFunTy env . varType $ var
+forceSpecBndr env var = forceSpecFunTy env . snd . splitForAllTys . varType $ var
 
 forceSpecFunTy :: ScEnv -> Type -> Bool
 forceSpecFunTy env = any (forceSpecArgTy env) . fst . splitFunTys
@@ -683,7 +711,7 @@ forceSpecArgTy env ty
 forceSpecArgTy env ty
   | Just (tycon, tys) <- splitTyConApp_maybe ty
   , tycon /= funTyCon
-      = L.lookupUFM (sc_annotations env) tycon == Just ForceSpecConstr
+      = lookupUFM (sc_annotations env) tycon == Just ForceSpecConstr
         || any (forceSpecArgTy env) tys
 
 forceSpecArgTy _ _ = False
@@ -908,7 +936,7 @@ scExpr' env (Case scrut b ty alts)
           ; return (usg', scrut_occ, (con, bs2, rhs')) }
 
 scExpr' env (Let (NonRec bndr rhs) body)
-  | isTyVar bndr       -- Type-lets may be created by doBeta
+  | isTyCoVar bndr     -- Type-lets may be created by doBeta
   = scExpr' (extendScSubst env bndr rhs) body
 
   | otherwise             -- Note [Local let bindings]
@@ -918,6 +946,9 @@ scExpr' env (Let (NonRec bndr rhs) body)
 
        ; (rhs_usg, rhs_info) <- scRecRhs env (bndr',rhs)
 
+          -- NB: We don't use the ForceSpecConstr mechanism (see
+          -- Note [Forcing specialisation]) for non-recursive bindings
+          -- at the moment. I'm not sure if this is the right thing to do.
        ; let force_spec = False
        ; (spec_usg, specs) <- specialise env force_spec 
                                           (scu_calls body_usg) 
@@ -936,6 +967,7 @@ scExpr' env (Let (Rec prs) body)
              (rhs_env1,bndrs') = extendRecBndrs env bndrs
              rhs_env2 = extendHowBound rhs_env1 bndrs' RecFun
               force_spec = any (forceSpecBndr env) bndrs'
+                -- Note [Forcing specialisation]
 
        ; (rhs_usgs, rhs_infos) <- mapAndUnzipM (scRecRhs rhs_env2) (bndrs' `zip` rhss)
        ; (body_usg, body')     <- scExpr rhs_env2 body
@@ -1039,6 +1071,7 @@ scTopBind env (Rec prs)
   where
     (bndrs,rhss) = unzip prs
     force_spec = any (forceSpecBndr env) bndrs
+      -- Note [Forcing specialisation]
 
 scTopBind env (NonRec bndr rhs)
   = do { (_, rhs') <- scExpr env rhs
@@ -1108,6 +1141,7 @@ data OneSpec  = OS CallPat                -- Call pattern that generated this specialisation
 
 specLoop :: ScEnv
          -> Bool                                -- force specialisation?
+                                                -- Note [Forcing specialisation]
         -> CallEnv
         -> [RhsInfo]
         -> ScUsage -> [SpecInfo]               -- One per binder; acccumulating parameter
@@ -1126,6 +1160,7 @@ specLoop env force_spec all_calls rhs_infos usg_so_far specs_so_far
 specialise 
    :: ScEnv
    -> Bool                              -- force specialisation?
+                                        --   Note [Forcing specialisation]
    -> CallEnv                          -- Info on calls
    -> RhsInfo
    -> SpecInfo                         -- Original RHS plus patterns dealt with
@@ -1138,6 +1173,7 @@ specialise
 specialise env force_spec bind_calls (RI fn _ arg_bndrs body arg_occs) 
                          spec_info@(SI specs spec_count mb_unspec)
   | not (isBottomingId fn)      -- Note [Do not specialise diverging functions]
+  , not (isNeverActive (idInlineActivation fn))        -- See Note [Transfer activation]
   , notNull arg_bndrs          -- Only specialise functions
   , Just all_calls <- lookupVarEnv bind_calls fn
   = do { (boring_call, pats) <- callsToPats env specs arg_occs all_calls
@@ -1244,6 +1280,7 @@ spec_one env fn arg_bndrs body (call_pat@(qvars, pats), rule_number)
              rule_rhs   = mkVarApps (Var spec_id) spec_call_args
               inline_act = idInlineActivation fn
              rule       = mkLocalRule rule_name inline_act fn_name qvars pats rule_rhs
+                          -- See Note [Transfer activation]
        ; return (spec_usg, OS call_pat rule spec_id spec_rhs) }
 
 calcSpecStrictness :: Id                    -- The original function
@@ -1272,13 +1309,17 @@ calcSpecStrictness fn qvars pats
 
 Note [Specialise original body]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-The RhsInfo for a binding keeps the *oringal* body of the binding.  We
+The RhsInfo for a binding keeps the *original* body of the binding.  We
 must specialise that, *not* the result of applying specExpr to the RHS
 (which is also kept in RhsInfo). Otherwise we end up specialising a
 specialised RHS, and that can lead directly to exponential behaviour.
 
 Note [Transfer activation]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~
+  This note is for SpecConstr, but exactly the same thing
+  happens in the overloading specialiser; see
+  Note [Auto-specialisation and RULES] in Specialise.
+
 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.  Then I made them active only
@@ -1289,8 +1330,9 @@ simplCore/should_compile/spec-inline.
 
 So now I just use the inline-activation of the parent Id, as the
 activation for the specialiation RULE, just like the main specialiser;
-see Note [Auto-specialisation and RULES] in Specialise.
 
+This in turn means there is no point in specialising NOINLINE things,
+so we test for that.
 
 Note [Transfer strictness]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -1356,7 +1398,7 @@ callToPats env bndr_occs (con_env, args)
                -- at the call site
                -- See Note [Shadowing] at the top
                
-             (tvs, ids) = partition isTyVar qvars
+             (tvs, ids) = partition isTyCoVar qvars
              qvars'     = tvs ++ ids
                -- Put the type variables first; the type of a term
                -- variable may mention a type variable
@@ -1402,11 +1444,18 @@ argToPat env in_scope val_env (Note _ arg) arg_occ
 
 argToPat env in_scope val_env (Let _ arg) arg_occ
   = argToPat env in_scope val_env arg arg_occ
+       -- See Note [Matching lets] in Rule.lhs
        -- Look through let expressions
-       -- e.g.         f (let v = rhs in \y -> ...v...)
-       -- Here we can specialise for f (\y -> ...)
+       -- e.g.         f (let v = rhs in (v,w))
+       -- Here we can specialise for f (v,w)
        -- because the rule-matcher will look through the let.
 
+{- Disabled; see Note [Matching cases] in Rule.lhs
+argToPat env in_scope val_env (Case scrut _ _ [(_, _, rhs)]) arg_occ
+  | exprOkForSpeculation scrut -- See Note [Matching cases] in Rule.hhs
+  = argToPat env in_scope val_env rhs arg_occ
+-}
+
 argToPat env in_scope val_env (Cast arg co) arg_occ
   | not (ignoreType env ty2)
   = do { (interesting, arg') <- argToPat env in_scope val_env arg arg_occ
@@ -1523,7 +1572,7 @@ isValue env (Var v)
        -- as well, for let-bound constructors!
 
 isValue env (Lam b e)
-  | isTyVar b = case isValue env e of
+  | isTyCoVar b = case isValue env e of
                  Just _  -> Just LambdaVal
                  Nothing -> Nothing
   | otherwise = Just LambdaVal