Fix scoped type variables for expression type signatures
[ghc-hetmet.git] / compiler / specialise / Rules.lhs
index 4c223d4..4f62115 100644 (file)
@@ -10,6 +10,8 @@ module Rules (
 
        mkSpecInfo, extendSpecInfo, addSpecInfo,
        rulesOfBinds, addIdSpecialisations, 
+       
+       matchN,
 
         lookupRule, mkLocalRule, roughTopNames
     ) where
@@ -18,11 +20,12 @@ module Rules (
 
 import CoreSyn         -- All of it
 import OccurAnal       ( occurAnalyseExpr )
-import CoreFVs         ( exprFreeVars, exprsFreeVars, rulesRhsFreeVars )
+import CoreFVs         ( exprFreeVars, exprsFreeVars, bindFreeVars, rulesRhsFreeVars )
 import CoreUnfold      ( isCheapUnfolding, unfoldingTemplate )
 import CoreUtils       ( tcEqExprX )
 import PprCore         ( pprRules )
 import Type            ( TvSubstEnv )
+import Coercion         ( coercionKind )
 import TcType          ( tcSplitTyConApp_maybe )
 import CoreTidy                ( tidyRules )
 import Id              ( Id, idUnfolding, isLocalId, isGlobalId, idName,
@@ -33,7 +36,8 @@ import VarEnv         ( IdEnv, InScopeSet, emptyTidyEnv,
                          emptyInScopeSet, mkInScopeSet, extendInScopeSetList, 
                          emptyVarEnv, lookupVarEnv, extendVarEnv, 
                          nukeRnEnvL, mkRnEnv2, rnOccR, rnOccL, inRnEnvR,
-                         rnBndrR, rnBndr2, rnBndrL, rnBndrs2 )
+                         rnBndrR, rnBndr2, rnBndrL, rnBndrs2,
+                         rnInScope, extendRnInScopeList, lookupRnInScope )
 import VarSet
 import Name            ( Name, NamedThing(..), nameOccName )
 import NameEnv
@@ -42,6 +46,7 @@ import BasicTypes     ( Activation, CompilerPhase, isActive )
 import Outputable
 import FastString
 import Maybes          ( isJust, orElse )
+import OrdList
 import Bag
 import Util            ( singleton )
 import List            ( isPrefixOf )
@@ -237,7 +242,7 @@ findBest :: (Id, [CoreExpr])
 findBest target (rule,ans)   [] = (rule,ans)
 findBest target (rule1,ans1) ((rule2,ans2):prs)
   | rule1 `isMoreSpecific` rule2 = findBest target (rule1,ans1) prs
-  | rule2 `isMoreSpecific` rule1 = findBest target (rule1,ans1) prs
+  | rule2 `isMoreSpecific` rule1 = findBest target (rule2,ans2) prs
 #ifdef DEBUG
   | otherwise = pprTrace "Rules.findBest: rule overlap (Rule 1 wins)"
                         (vcat [ptext SLIT("Expression to match:") <+> ppr fn <+> sep (map ppr args),
@@ -305,9 +310,10 @@ matchRule is_active in_scope args rough_args
   | otherwise
   = case matchN in_scope tpl_vars tpl_args args of
        Nothing                    -> Nothing
-       Just (tpl_vals, leftovers) -> Just (rule_fn
-                                           `mkApps` tpl_vals
-                                           `mkApps` leftovers)
+       Just (binds, tpl_vals, leftovers) -> Just (mkLets binds $
+                                                  rule_fn
+                                                   `mkApps` tpl_vals
+                                                   `mkApps` leftovers)
   where
     rule_fn = occurAnalyseExpr (mkLams tpl_vars rhs)
        -- We could do this when putting things into the rulebase, I guess
@@ -318,12 +324,16 @@ matchN    :: InScopeSet
        -> [Var]                -- Template tyvars
        -> [CoreExpr]           -- Template
        -> [CoreExpr]           -- Target; can have more elts than template
-       -> Maybe ([CoreExpr],   -- What is substituted for each template var
+       -> Maybe ([CoreBind],   -- Bindings to wrap around the entire result
+                 [CoreExpr],   -- What is substituted for each template var
                  [CoreExpr])   -- Leftover target exprs
 
 matchN in_scope tmpl_vars tmpl_es target_es
-  = do { (subst, leftover_es) <- go init_menv emptySubstEnv tmpl_es target_es
-       ; return (map (lookup_tmpl subst) tmpl_vars, leftover_es) }
+  = do { ((tv_subst, id_subst, binds), leftover_es)
+               <- go init_menv emptySubstEnv tmpl_es target_es
+       ; return (fromOL binds, 
+                 map (lookup_tmpl tv_subst id_subst) tmpl_vars, 
+                 leftover_es) }
   where
     init_menv = ME { me_tmpls = mkVarSet tmpl_vars, me_env = init_rn_env }
     init_rn_env = mkRnEnv2 (extendInScopeSetList in_scope tmpl_vars)
@@ -333,8 +343,8 @@ matchN in_scope tmpl_vars tmpl_es target_es
     go menv subst (t:ts) (e:es) = do { subst1 <- match menv subst t e 
                                     ; go menv subst1 ts es }
 
-    lookup_tmpl :: (TvSubstEnv, IdSubstEnv) -> Var -> CoreExpr
-    lookup_tmpl (tv_subst, id_subst) tmpl_var
+    lookup_tmpl :: TvSubstEnv -> IdSubstEnv -> Var -> CoreExpr
+    lookup_tmpl tv_subst id_subst tmpl_var
        | isTyVar tmpl_var = case lookupVarEnv tv_subst tmpl_var of
                                Just ty         -> Type ty
                                Nothing         -> unbound tmpl_var
@@ -353,13 +363,18 @@ matchN in_scope tmpl_vars tmpl_es target_es
 \begin{code}
 -- These two definitions are not the same as in Subst,
 -- but they simple and direct, and purely local to this module
--- The third, for TvSubstEnv, is the same as in VarEnv, but repeated here
--- for uniformity with IdSubstEnv
-type SubstEnv   = (TvSubstEnv, IdSubstEnv)     
-type IdSubstEnv = IdEnv    CoreExpr            
+--
+-- * The domain of the TvSubstEnv and IdSubstEnv are the template
+--   variables passed into the match.
+--
+-- * The (OrdList CoreBind) in a SubstEnv are the bindings floated out
+--   from nested matches; see the Let case of match, below
+--
+type SubstEnv   = (TvSubstEnv, IdSubstEnv, OrdList CoreBind)
+type IdSubstEnv = IdEnv CoreExpr               
 
 emptySubstEnv :: SubstEnv
-emptySubstEnv = (emptyVarEnv, emptyVarEnv)
+emptySubstEnv = (emptyVarEnv, emptyVarEnv, nilOL)
 
 
 --     At one stage I tried to match even if there are more 
@@ -393,38 +408,59 @@ match :: MatchEnv
 -- succeed in matching what looks like the template variable 'a' against 3.
 
 -- The Var case follows closely what happens in Unify.match
-match menv subst@(tv_subst, id_subst) (Var v1) e2 
-  | v1 `elemVarSet` me_tmpls menv
-  = case lookupVarEnv id_subst v1' of
-       Nothing | any (inRnEnvR rn_env) (varSetElems (exprFreeVars e2))
-               -> Nothing      -- Occurs check failure
-               -- e.g. match forall a. (\x-> a x) against (\y. y y)
-
-               | otherwise
-               -> Just (tv_subst, extendVarEnv id_subst v1 e2)
-
-       Just e2' | tcEqExprX (nukeRnEnvL rn_env) e2' e2 
-                -> Just subst
-
-       other -> Nothing
-
-  |    -- v1 is not a template variable; check for an exact match with e2
-    Var v2 <- e2, v1' == rnOccR rn_env v2
+match menv subst (Var v1) e2 
+  | Just subst <- match_var menv subst v1 e2
   = Just subst
 
-  where
-    rn_env = me_env menv
-    v1'    = rnOccL rn_env v1
-
 -- Here is another important rule: if the term being matched is a
 -- variable, we expand it so long as its unfolding is a WHNF
 -- (Its occurrence information is not necessarily up to date,
 --  so we don't use it.)
 match menv subst e1 (Var v2)
-  | isCheapUnfolding unfolding
+  | not (inRnEnvR rn_env v2),
+       -- If v2 is in the RnEnvR, then it's locally bound and can't
+       -- have an unfolding. We must make this check because if it
+       -- is locally bound we must not look it up in the in-scope set
+       -- E.g.         (\x->x) where x is already in scope
+    isCheapUnfolding unfolding
   = match menv subst e1 (unfoldingTemplate unfolding)
   where
-    unfolding = idUnfolding v2
+    rn_env    = me_env menv
+    unfolding = idUnfolding (lookupRnInScope rn_env v2)
+       -- Notice that we look up v2 in the in-scope set
+       -- See Note [Lookup in-scope]
+
+-- Matching a let-expression.  Consider
+--     RULE forall x.  f (g x) = <rhs>
+-- and target expression
+--     f (let { w=R } in g E))
+-- Then we'd like the rule to match, to generate
+--     let { w=R } in (\x. <rhs>) E
+-- In effect, we want to float the let-binding outward, to enable
+-- the match to happen.  This is the WHOLE REASON for accumulating
+-- bindings in the SubstEnv
+--
+-- We can only do this if
+--     (a) Widening the scope of w does not capture any variables
+--         We use a conservative test: w is not already in scope
+--     (b) The free variables of R are not bound by the part of the
+--         target expression outside the let binding; e.g.
+--             f (\v. let w = v+1 in g E)
+--         Here we obviously cannot float the let-binding for w.
+
+match menv subst@(tv_subst, id_subst, binds) e1 (Let bind e2)
+  | all freshly_bound bndrs,
+    not (any locally_bound bind_fvs)
+  = match (menv { me_env = rn_env' }) 
+         (tv_subst, id_subst, binds `snocOL` bind)
+         e1 e2
+  where
+    rn_env   = me_env menv
+    bndrs    = bindersOf bind
+    bind_fvs = varSetElems (bindFreeVars bind)
+    freshly_bound x = not (x `rnInScope` rn_env)
+    locally_bound x = inRnEnvR rn_env x
+    rn_env' = extendRnInScopeList rn_env bndrs
 
 match menv subst (Lit lit1) (Lit lit2)
   | lit1 == lit2
@@ -441,6 +477,9 @@ match menv subst (Lam x1 e1) (Lam x2 e2)
 
 -- This rule does eta expansion
 --             (\x.M)  ~  N    iff     M  ~  N x
+-- It's important that this is *after* the let rule,
+-- so that     (\x.M)  ~  (let y = e in \y.N)
+-- does the let thing, and then gets the lam/lam rule above
 match menv subst (Lam x1 e1) e2
   = match menv' subst e1 (App e2 (varToCoreExpr new_x))
   where
@@ -458,38 +497,86 @@ match menv subst e1 (Lam x2 e2)
 match menv subst (Case e1 x1 ty1 alts1) (Case e2 x2 ty2 alts2)
   = do { subst1 <- match_ty menv subst ty1 ty2
        ; subst2 <- match menv subst1 e1 e2
-       ; let menv' = menv { me_env = rnBndr2 (me_env menv) x2 x2 }
+       ; let menv' = menv { me_env = rnBndr2 (me_env menv) x1 x2 }
        ; match_alts menv' subst2 alts1 alts2   -- Alts are both sorted
        }
 
 match menv subst (Type ty1) (Type ty2)
   = match_ty menv subst ty1 ty2
 
-match menv subst (Note (Coerce to1 from1) e1) (Note (Coerce to2 from2) e2)
+match menv subst (Cast e1 co1) (Cast e2 co2)
+  | (from1, to1) <- coercionKind co1
+  , (from2, to2) <- coercionKind co2
   = do { subst1 <- match_ty menv subst  to1   to2
        ; subst2 <- match_ty menv subst1 from1 from2
        ; match menv subst2 e1 e2 }
 
+{-     REMOVING OLD CODE: I think that the above handling for let is 
+                          better than the stuff here, which looks 
+                          pretty suspicious to me.  SLPJ Sept 06
 -- This is an interesting rule: we simply ignore lets in the 
 -- term being matched against!  The unfolding inside it is (by assumption)
 -- already inside any occurrences of the bound variables, so we'll expand
--- them when we encounter them.
-match menv subst e1 (Let (NonRec x2 r2) e2)
-  = match menv' subst e1 e2
+-- them when we encounter them.  This gives a chance of matching
+--     forall x,y.  f (g (x,y))
+-- against
+--     f (let v = (a,b) in g v)
+
+match menv subst e1 (Let bind e2)
+  = match (menv { me_env = rn_env' }) subst e1 e2
   where
-    menv' = menv { me_env = fst (rnBndrR (me_env menv) x2) }
-       -- It's important to do this renaming. For example:
+    (rn_env', _bndrs') = mapAccumL rnBndrR (me_env menv) (bindersOf bind)
+       -- It's important to do this renaming, so that the bndrs
+       -- are brought into the local scope. For example:
        -- Matching
        --      forall f,x,xs. f (x:xs)
        --   against
        --      f (let y = e in (y:[]))
-       -- We must not get success with x->y!  Instead, we 
-       -- need an occurs check.
+       -- We must not get success with x->y!  So we record that y is
+       -- locally bound (with rnBndrR), and proceed.  The Var case
+       -- will fail when trying to bind x->y
+-}
 
 -- Everything else fails
 match menv subst e1 e2 = Nothing
 
 ------------------------------------------
+match_var :: MatchEnv
+         -> SubstEnv
+         -> Var                -- Template
+         -> CoreExpr           -- Target
+         -> Maybe SubstEnv
+match_var menv subst@(tv_subst, id_subst, binds) v1 e2
+  | v1' `elemVarSet` me_tmpls menv
+  = case lookupVarEnv id_subst v1' of
+       Nothing | any (inRnEnvR rn_env) (varSetElems (exprFreeVars e2))
+               -> Nothing      -- Occurs check failure
+               -- e.g. match forall a. (\x-> a x) against (\y. y y)
+
+               | otherwise     -- No renaming to do on e2
+               -> Just (tv_subst, extendVarEnv id_subst v1 e2, binds)
+
+       Just e2' | tcEqExprX (nukeRnEnvL rn_env) e2' e2 
+                -> Just subst
+
+                | otherwise
+                -> Nothing
+
+  | otherwise  -- v1 is not a template variable; check for an exact match with e2
+  = case e2 of
+       Var v2 | v1' == rnOccR rn_env v2 -> Just subst
+       other                           -> Nothing
+
+  where
+    rn_env = me_env menv
+    v1'    = rnOccL rn_env v1  
+       -- If the template is
+       --      forall x. f x (\x -> x) = ...
+       -- Then the x inside the lambda isn't the 
+       -- template x, so we must rename first!
+                               
+
+------------------------------------------
 match_alts :: MatchEnv
       -> SubstEnv
       -> [CoreAlt]             -- Template
@@ -517,12 +604,48 @@ We only want to replace (f T) with f', not (f Int).
 
 \begin{code}
 ------------------------------------------
-match_ty menv (tv_subst, id_subst) ty1 ty2
+match_ty menv (tv_subst, id_subst, binds) ty1 ty2
   = do { tv_subst' <- Unify.ruleMatchTyX menv tv_subst ty1 ty2
-       ; return (tv_subst', id_subst) }
+       ; return (tv_subst', id_subst, binds) }
 \end{code}
 
 
+Note [Lookup in-scope]
+~~~~~~~~~~~~~~~~~~~~~~
+Consider this example
+       foo :: Int -> Maybe Int -> Int
+       foo 0 (Just n) = n
+       foo m (Just n) = foo (m-n) (Just n)
+
+SpecConstr sees this fragment:
+
+       case w_smT of wild_Xf [Just A] {
+         Data.Maybe.Nothing -> lvl_smf;
+         Data.Maybe.Just n_acT [Just S(L)] ->
+           case n_acT of wild1_ams [Just A] { GHC.Base.I# y_amr [Just L] ->
+           $wfoo_smW (GHC.Prim.-# ds_Xmb y_amr) wild_Xf
+           }};
+
+and correctly generates the rule
+
+       RULES: "SC:$wfoo1" [0] __forall {y_amr [Just L] :: GHC.Prim.Int#
+                                         sc_snn :: GHC.Prim.Int#}
+         $wfoo_smW sc_snn (Data.Maybe.Just @ GHC.Base.Int (GHC.Base.I# y_amr))
+         = $s$wfoo_sno y_amr sc_snn ;]
+
+BUT we must ensure that this rule matches in the original function!
+Note that the call to $wfoo is
+           $wfoo_smW (GHC.Prim.-# ds_Xmb y_amr) wild_Xf
+
+During matching we expand wild_Xf to (Just n_acT).  But then we must also
+expand n_acT to (I# y_amr).  And we can only do that if we look up n_acT
+in the in-scope set, because in wild_Xf's unfolding it won't have an unfolding
+at all. 
+
+That is why the 'lookupRnInScope' call in the (Var v2) case of 'match'
+is so important.
+
+
 %************************************************************************
 %*                                                                     *
 \subsection{Checking a program for failing rule applications}
@@ -570,6 +693,7 @@ ruleCheck env (Lit l)           = emptyBag
 ruleCheck env (Type ty)     = emptyBag
 ruleCheck env (App f a)     = ruleCheckApp env (App f a) []
 ruleCheck env (Note n e)    = ruleCheck env e
+ruleCheck env (Cast e co)   = ruleCheck env e
 ruleCheck env (Let bd e)    = ruleCheckBind env bd `unionBags` ruleCheck env e
 ruleCheck env (Lam b e)     = ruleCheck env e
 ruleCheck env (Case e _ _ as) = ruleCheck env e `unionBags`