Refactor CoreArity a bit
authorsimonpj@microsoft.com <unknown>
Thu, 24 Dec 2009 15:34:48 +0000 (15:34 +0000)
committersimonpj@microsoft.com <unknown>
Thu, 24 Dec 2009 15:34:48 +0000 (15:34 +0000)
I was experimenting with making coercions opaque to
arity.  I think this is ultimately the right thing to do
but I've left the functionality unchanged for now.

compiler/coreSyn/CoreArity.lhs
compiler/types/Type.lhs

index be34b07..49106df 100644 (file)
@@ -99,29 +99,35 @@ exprArity :: CoreExpr -> Arity
 -- ^ An approximate, fast, version of 'exprEtaExpandArity'
 exprArity e = go e
   where
-    go (Var v)                          = idArity v
-    go (Lam x e) | isId x       = go e + 1
-                | otherwise     = go e
-    go (Note _ e)                = go e
-    go (Cast e co)               = trim_arity (go e) 0 (snd (coercionKind co))
-    go (App e (Type _))          = go e
-    go (App f a) | exprIsCheap a = (go f - 1) `max` 0
-       -- NB: exprIsCheap a!  
-       --      f (fac x) does not have arity 2, 
-       --      even if f has arity 3!
-       -- NB: `max 0`!  (\x y -> f x) has arity 2, even if f is
-       --               unknown, hence arity 0
+    go (Var v)                            = idArity v
+    go (Lam x e) | isId x         = go e + 1
+                | otherwise       = go e
+    go (Note _ e)                  = go e
+    go (Cast e co)                 = go e `min` typeArity (snd (coercionKind co))
+                                               -- Note [exprArity invariant]
+    go (App e (Type _))            = go e
+    go (App f a) | exprIsTrivial a = (go f - 1) `max` 0
+        -- See Note [exprArity for applications]
     go _                          = 0
-
-       -- Note [exprArity invariant]
-    trim_arity n a ty
-       | n==a                                        = a
-       | Just (_, ty') <- splitForAllTy_maybe ty     = trim_arity n a     ty'
-       | Just (_, ty') <- splitFunTy_maybe ty        = trim_arity n (a+1) ty'
-       | Just (ty',_)  <- splitNewTypeRepCo_maybe ty = trim_arity n a     ty'
-       | otherwise                                   = a
 \end{code}
 
+Note [exprArity for applications]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+When we come to an application we check that the arg is trivial.
+   eg  f (fac x) does not have arity 2, 
+                 even if f has arity 3!
+
+* We require that is trivial rather merely cheap.  Suppose f has arity 2.
+  Then    f (Just y)
+  has arity 0, because if we gave it arity 1 and then inlined f we'd get
+          let v = Just y in \w. <f-body>
+  which has arity 0.  And we try to maintain the invariant that we don't
+  have arity decreases.
+
+*  The `max 0` is important!  (\x y -> f x) has arity 2, even if f is
+   unknown, hence arity 0
+
+
 %************************************************************************
 %*                                                                     *
           Eta expansion
@@ -169,7 +175,6 @@ Or, to put it another way, in any context C
          is as efficient as
    C[ e ]
 
-
 It's all a bit more subtle than it looks:
 
 Note [Arity of case expressions]
@@ -191,7 +196,6 @@ This should diverge!  But if we eta-expand, it won't.   Again, we ignore this
 "problem", because being scrupulous would lose an important transformation for
 many programs.
 
-
 1.  Note [One-shot lambdas]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Consider one-shot lambdas
@@ -212,7 +216,6 @@ should diverge, but it'll converge if we eta-expand f.  Nevertheless, we
 do so; it improves some programs significantly, and increasing convergence
 isn't a bad thing.  Hence the ABot/ATop in ArityType.
 
-
 4. Note [Newtype arity]
 ~~~~~~~~~~~~~~~~~~~~~~~~
 Non-recursive newtypes are transparent, and should not get in the way.
@@ -233,26 +236,6 @@ we want to get:             coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
   And since negate has arity 2, you might try to eta expand.  But you can't
   decopose Int to a function type.   Hence the final case in eta_expand.
   
-Note [The state-transformer hack]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-Suppose we have 
-       f = e
-where e has arity n.  Then, if we know from the context that f has
-a usage type like
-       t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
-then we can expand the arity to m.  This usage type says that
-any application (x e1 .. en) will be applied to uniquely to (m-n) more args
-Consider f = \x. let y = <expensive> 
-                in case x of
-                     True  -> foo
-                     False -> \(s:RealWorld) -> e
-where foo has arity 1.  Then we want the state hack to
-apply to foo too, so we can eta expand the case.
-
-Then we expect that if f is applied to one arg, it'll be applied to two
-(that's the hack -- we don't really know, and sometimes it's false)
-See also Id.isOneShotBndr.
-
 \begin{code}
 applyStateHack :: CoreExpr -> ArityType -> Arity
 applyStateHack e (AT orig_arity is_bot)
@@ -264,16 +247,18 @@ applyStateHack e (AT orig_arity is_bot)
     go :: Type -> Arity -> Arity
     go ty arity                -- This case analysis should match that in eta_expand
        | Just (_, ty') <- splitForAllTy_maybe ty   = go ty' arity
+       | Just (arg,res) <- splitFunTy_maybe ty
+       , arity > 0 || isStateHackType arg = 1 + go res (arity-1)
 
+-- See Note [trimCast]
        | Just (tc,tys) <- splitTyConApp_maybe ty 
        , Just (ty', _) <- instNewTyCon_maybe tc tys
        , not (isRecursiveTyCon tc)                 = go ty' arity
                -- Important to look through non-recursive newtypes, so that, eg 
                --      (f x)   where f has arity 2, f :: Int -> IO ()
                -- Here we want to get arity 1 for the result!
+-------
 
-       | Just (arg,res) <- splitFunTy_maybe ty
-       , arity > 0 || isStateHackType arg = 1 + go res (arity-1)
 {-
        = if arity > 0 then 1 + go res (arity-1)
          else if isStateHackType arg then
@@ -285,6 +270,26 @@ applyStateHack e (AT orig_arity is_bot)
        | otherwise = WARN( arity > 0, ppr arity <+> ppr ty) 0
 \end{code}
 
+Note [The state-transformer hack]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Suppose we have 
+       f = e
+where e has arity n.  Then, if we know from the context that f has
+a usage type like
+       t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
+then we can expand the arity to m.  This usage type says that
+any application (x e1 .. en) will be applied to uniquely to (m-n) more args
+Consider f = \x. let y = <expensive> 
+                in case x of
+                     True  -> foo
+                     False -> \(s:RealWorld) -> e
+where foo has arity 1.  Then we want the state hack to
+apply to foo too, so we can eta expand the case.
+
+Then we expect that if f is applied to one arg, it'll be applied to two
+(that's the hack -- we don't really know, and sometimes it's false)
+See also Id.isOneShotBndr.
+
 Note [State hack and bottoming functions]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 It's a terrible idea to use the state hack on a bottoming function.
@@ -348,6 +353,29 @@ andArityType (AT _  ABot) (AT a2 ATop) = AT a2               ATop
 andArityType (AT a1 ATop) (AT _  ABot) = AT a1           ATop
 andArityType (AT a1 ABot) (AT a2 ABot) = AT (a1 `max` a2) ABot
 
+---------------------------
+trimCast :: Coercion -> ArityType -> ArityType
+-- Trim the arity to be no more than allowed by the
+-- arrows in ty2, where co :: ty1~ty2
+trimCast _ at = at
+
+{-        Omitting for now Note [trimCast]
+trimCast co at@(AT ar _)
+  | ar > co_arity = AT co_arity ATop
+  | otherwise     = at
+  where
+    (_, ty2) = coercionKind co
+    co_arity = typeArity ty2
+-}
+\end{code}
+
+Note [trimCast]
+~~~~~~~~~~~~~~~
+When you try putting trimCast back in, comment out the snippets
+flagged by the other references to Note [trimCast]
+
+\begin{code}
+---------------------------
 trimArity :: Bool -> ArityType -> ArityType
 -- We have something like (let x = E in b), where b has the given
 -- arity type.  Then
@@ -417,9 +445,9 @@ arityType dicts_cheap (Let b e)
        -- See Note [Dictionary-like types] in TcType.lhs for why we use
        -- isDictLikeTy here rather than isDictTy
 
-arityType dicts_cheap (Note _ e) = arityType dicts_cheap e
-arityType dicts_cheap (Cast e _) = arityType dicts_cheap e
-arityType _           _          = vanillaArityType
+arityType dicts_cheap (Note _ e)  = arityType dicts_cheap e
+arityType dicts_cheap (Cast e co) = trimCast co (arityType dicts_cheap e)
+arityType _           _           = vanillaArityType
 \end{code}
   
   
@@ -470,11 +498,9 @@ etaExpand :: Arity         -- ^ Result should have this number of value args
 -- so perhaps the extra code isn't worth it
 
 etaExpand n orig_expr
-  | manifestArity orig_expr >= n = orig_expr   -- The no-op case
-  | otherwise              
   = go n orig_expr
   where
-      -- Strip off existing lambdas
+      -- Strip off existing lambdas and casts
       -- Note [Eta expansion and SCCs]
     go 0 expr = expr
     go n (Lam v body) | isTyVar v = Lam v (go n     body)
@@ -560,8 +586,8 @@ mkEtaWW :: Arity -> InScopeSet -> Type
        -- Outgoing InScopeSet includes the EtaInfo vars
        --   and the original free vars
 
-mkEtaWW n in_scope ty
-  = go n empty_subst ty []
+mkEtaWW orig_n in_scope orig_ty
+  = go orig_n empty_subst orig_ty []
   where
     empty_subst = mkTvSubst in_scope emptyTvSubstEnv
 
@@ -579,6 +605,7 @@ mkEtaWW n in_scope ty
            -- Avoid free vars of the original expression
        = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
                                           
+-- See Note [trimCast]
        | Just(ty',co) <- splitNewTypeRepCo_maybe ty
        =       -- Given this:
                        --      newtype T = MkT ([T] -> Int)
index 8a9cf0e..8177e5a 100644 (file)
@@ -30,7 +30,7 @@ module Type (
 
        mkFunTy, mkFunTys, splitFunTy, splitFunTy_maybe, 
        splitFunTys, splitFunTysN,
-       funResultTy, funArgTy, zipFunTys,
+       funResultTy, funArgTy, zipFunTys, typeArity,
 
        mkTyConApp, mkTyConTy, 
        tyConAppTyCon, tyConAppArgs, 
@@ -141,6 +141,7 @@ import VarSet
 import Name
 import Class
 import TyCon
+import BasicTypes      ( Arity )
 
 -- others
 import StaticFlags
@@ -495,6 +496,14 @@ funArgTy :: Type -> Type
 funArgTy ty | Just ty' <- coreView ty = funArgTy ty'
 funArgTy (FunTy arg _res)  = arg
 funArgTy ty                = pprPanic "funArgTy" (ppr ty)
+
+typeArity :: Type -> Arity
+-- How many value arrows are visible in the type?
+-- We look through foralls, but not through newtypes, dictionaries etc
+typeArity ty | Just ty' <- coreView ty = typeArity ty'
+typeArity (FunTy _ ty)    = 1 + typeArity ty
+typeArity (ForAllTy _ ty) = typeArity ty
+typeArity _               = 0
 \end{code}
 
 ---------------------------------------------------------------------
@@ -1334,7 +1343,7 @@ then (substTy subst ty) does nothing.
 For example, consider:
        (/\a. /\b:(a~Int). ...b..) Int
 We substitute Int for 'a'.  The Unique of 'b' does not change, but
-nevertheless we add 'b' to the TvSubstEnv, because b's type does change
+nevertheless we add 'b' to the TvSubstEnv, because b's kind does change
 
 This invariant has several crucial consequences: