{-# OPTIONS -XRecordWildCards -XNamedFieldPuns #-} 
module Scp (
           scpProgram
          ) where

#include "HsVersions.h"

-- Define to 0 for a sequential run
#define PARALLEL_SCP 0
#define OLDBETA 0
#define MATCHTYPES 0
#define FOLDVARS 1
#define PATTERNSUBST 1
-- #undef PATTERNSUBST
-- #undef FOLDVARS
-- #define OLD_HOMEMB 1
#undef OLD_HOMEMB
-- #define NOBORING 1
-- #define NO_NORMALIZE 1
#undef NO_NORMALIZE
#undef NOBORING

{-

Note: [Regression Testing]

THIS IS IMPORTANT.

Even small changes to the supercompiler can prevent optimizations from
happening, and it takes literally days to track down
afterwards. Verify that this does not happen, in the following order:

* The output from append (append xs ys) zs contains a new function
  $sappend that takes 3 arguments, and calls a different $sappend that
  takes two arguments from the first nil-branch. 

* Supercompile Charcount9, and compare the runtime on a 1M file
  between the supercompiled version and the ordinary version. The
  supercompiled version should complete in half the time.

Note: [Logging]

0: Bug reporting worthy things: missing rules, failed things that should work.

2: Whistles, folds, new funtcions

3: Print out whistle candidates and free variables

4: Functions that are going to be inlined, error prone steps in drive.

5: All steps in drive (SLOW).

6: Type checking of individual steps.

-}

{-

Note: [New Design]

* Switch Expression/Context format in rho to the one on Peter's notes.

* What about the whistle?

* Alpha-conversion in msg necessary

-}

import CoreMonad
import CoreSyn
import HscTypes
import BasicTypes       ( isActive )
import Rules		( RuleBase, mkRuleBase, lookupRule, getRules )
import CoreSubst        hiding (IdSubstEnv, substExpr)
import qualified CoreSubst        (substExpr)
import CoreUtils
import CoreFVs 		( exprFreeVars )
import Id		( realIdUnfolding, idUnfolding, mkUserLocal, idName, idDetails, idType, setIdUnfolding, isDataConWorkId_maybe, isDataConId_maybe, zapIdOccInfo, idUnique, idInlineActivation, mkLocalId )
import IdInfo           ( IdDetails(..) )
import VarSet           ( varSetElems, mkVarSet, elemVarSet )
import Type             ( Type, coreEqType, tcPartOfType, splitForAllTy_maybe, isFunTy, mkTyVarTy, splitTyConApp, splitTyConApp_maybe, zipOpenTvSubst, isUnLiftedType, isUnboxedTupleType, tyVarsOfTypes, coreView, tyVarsOfType, tcEqTypeX, repSplitAppTy_maybe, tyConAppArgs, applyTy, funResultTy, funArgTy, splitFunTy )
import qualified Type (substTy) 
import DataCon          ( dataConUnivTyVars, dataConRepArgTys, dataConExTyVars, dataConEqSpec, dataConRepArity, dataConTyCon, dataConEqTheta, DataCon() )
import VarEnv           ( emptyInScopeSet, emptyVarEnv, uniqAway, extendInScopeSet, extendInScopeSetList, mkRnEnv2, varEnvElts, lookupVarEnv, inRnEnvR, extendVarEnv, nukeRnEnvL, rnOccR, rnOccL, IdEnv, rnInScope, rnBndr2, rnBndrs2, nukeRnEnvR, lookupRnInScope, extendRnInScopeList, elemInScopeSet, getInScopeVars, extendInScopeSetSet, delInScopeSet, lookupInScope )
import Unify            ( MatchEnv(..) ) --, ruleMatchTyX )
import qualified Unify  ( ruleMatchTyX )
import Var              ( isLocalVar, Var(..), TyVar(), mkTyVar, tyVarKind, Id(..), isCoVar, isLocalId, isGlobalId, zapIdInfo )
import Util             ( equalLength, lengthAtLeast, mapAndUnzip )
import TypeRep          ( Type(..), PredType(..) )
import TyCon            ( TyCon(..), tyConArity, isFunTyCon, isAlgTyCon, isTupleTyCon, isSynTyCon, isPrimTyCon, isCoercionTyCon, isAnyTyCon, isSuperKindTyCon )
import OccurAnal        ( occurAnalysePgm )
import FiniteMap        ( FiniteMap, listToFM, addListToFM, lookupFM, lookupWithDefaultFM, emptyFM, addToFM_C, addToFM, fmToList, elemFM )
import TysPrim          ( realWorldTy, realWorldStatePrimTy )


-- For realEqExpr
import Coercion         ( coreEqCoercion, mkTransCoercion, coercionKind, mkInstCoercion, decomposeCo, mkSymCoercion, isIdentityCoercion, mkEqPred, getEqPredTys )

import PrelRules        ( primOpRules )
import MonadUtils       ( mapAndUnzipM )
import MkId             ( mkPrimOpId, mkImpossibleExpr, realWorldPrimId, voidArgId )
import PrimOp           ( allThePrimOps )
import Literal
import Coercion

import PprCore	--	( pprCoreBindings )
import Outputable
-- For finding main
import OccName hiding (varName)
import PrelNames        ( main_RDR_Unqual )
import FastString       ( mkFastString, fsLit )
import RdrName
import Name             ( Name, nameOccName, nameSrcSpan, mkSysTvName )

import DynFlags		( DynFlags(..), DynFlag(..)) 
import UniqSupply
import Maybes           ( orElse, fromJust, isJust )
import ScpMonad



import Data.List        ( intersect, nub, sortBy ) -- , unzip3 )
import Control.Monad    ( filterM )
#if PARALLEL_SCP
import Control.Parallel ( par )
#endif

import Debug.Trace      ( trace )

tr :: (Monad m) => String -> m ()
tr _ = return ()
-- tr m = trace m (return ())

substExpr :: Subst -> CoreExpr -> CoreExpr
substExpr s e = CoreSubst.substExpr empty s e

inlineFunctions :: [String]
inlineFunctions = ["append", "lvl",
--                   "fromInteger", "fromInteger2", "smallInteger", 
--                   "plusInteger", "timesInteger", "minusInteger", 
--                   "plusInt", "p2Num", "$p2Num", "f3", "$f3",
--                   "negateInteger", "absInteger", "signumInteger",
--                   "+", "f2", "$f2", "$f6", "f6", "D:Num", "Num",
--                   "Show", "D:Show", "a12", "a15",
--                   "a33", "wa11", "$wa11", "a11", "return",
--                   "a13", "getArgs", "lvl11", "a28",
--                   "stdout", 
--                   "lvl15", "lvl10", "a10",
--                   "f25", "$f25", "print", "a6", "a23", "a96",
--                   ">>", 
                   ">>="]

approvedFunctions :: [String]
approvedFunctions = inlineFunctions 

forbiddenFunctions :: [String]
forbiddenFunctions = [-- "lvl9", 
--                      "$wa2", "$wa6", "$wa1", "a12", 
--                      "a14", 
--                      "a7", "a48", "$wa", "a15", "alloca", "lvl6", "a35",
--                     "a22", "$wa22", "lvl4", 
--                      "lvl3", "lvl", "$wa7", 
--                      "stdout", "stdin", "errnoToIOError", "a13", "alloca"]
			"$fMonadIO_$c>>=", "$fMonadIO_$c>>"]

savingsThreshold = 29
-- savingsThreshold = -1000
taintThreshold = 3
-- taintThreshold = 10000

{-

Environment: goes downwards

-}

type InExpr = CoreExpr -- _Before_ transformation

type OutExpr = CoreExpr -- _After_ transformation
type OutId = Id
type OutVar = Var

data AExpr = AVar Id
           | ALit Literal
           | ALam Var ZExpr
           | AType Type

data ABind = ANonRec Var ZExpr

type AAlt = (AltCon, [Var], ZExpr)

data AFrame = AAppCtxt ZExpr
            | APrimOpCtxt [ZExpr] [ZExpr] 
            | ACaseCtxt Var Type [AAlt]
            | ACastCtxt Coercion

type AContext = [AFrame]

type ZExpr = (Int, AExpr, AContext, [ABind])

mkAExpr :: CoreExpr -> ZExpr
mkAExpr e = goE e [] []
    where goE (Var x) c bs = (i, AVar x, c, reverse bs)
              where i = 1 + zcsize c + zbsize bs
          goE (Lam b e) c bs = (i, ALam b e', c, reverse bs)
              where e' = mkAExpr e
                    i = 1 + zsize e' + zcsize c + zbsize bs
          goE (App e1 e2) c bs = goE e1 ((AAppCtxt (mkAExpr e2)):c) bs
          goE (Let (NonRec b e) body) c bs = goE body c ((ANonRec b (mkAExpr e)):bs)
          goE (Let (Rec {}) _) _ _ = error "Letrec, mkAExpr UPS!"
          goE (Case scrut b t alts) c bs = goE scrut ((ACaseCtxt b t (map goAlt alts)):c) bs
          goE (Cast e' co) c bs = goE e' ((ACastCtxt co):c) bs
          goE (Note _ e) c bs = goE e c bs
          goE (Lit l) c bs = (i, ALit l, c, reverse bs)
              where i = 1 + zcsize c + zbsize bs
          goE (Type t) c bs = (i, AType t, c, reverse bs)
              where i = 1 + zcsize c + zbsize bs

          goAlt (c, bs, e) = (c, bs, mkAExpr e)

zsize :: ZExpr -> Int
zsize (i, _, _, _) = i

zasize :: AExpr -> Int
zasize (AVar {}) = 1
zasize (ALit {}) = 1
zasize (ALam _ e) = 1 + zsize e
zasize (AType {}) = 1


zcsize :: AContext -> Int
zcsize cs = sum (map goC cs)
    where goC (AAppCtxt e) = 1 + zsize e
          goC (APrimOpCtxt ies oes) = 1 + sum (map zsize (ies ++ oes))
          goC (ACaseCtxt _ _ alts) = 1 +  sum (map goAlt alts)
          goC (ACastCtxt _) = 1
                              
          goAlt (_, _, e) = zsize e

zbsize :: [ABind] -> Int
zbsize bs = sum (map go bs)
    where go (ANonRec _ e) = zsize e

aConToCon :: AContext -> Context
aConToCon cs = go cs []
    where go [] res = reverse res
          go ((AAppCtxt arg):c) res = go c ((AppCtxt (mkE arg)):res)
          go ((APrimOpCtxt ies oes):c) res = go c ((PrimOpCtxt (map mkE ies) (map mkE oes)):res)
          go ((ACaseCtxt b t alts):c) res = go c ((CaseCtxt b t (map goAlt alts)):res)
          go ((ACastCtxt co):c) res = go c ((CastCtxt co):res)
          goAlt (c, bs, e) = (c, bs, mkE e)

mkE :: ZExpr -> CoreExpr
mkE (_, e, c, bs) = mkLets (mkBsE bs) (plugA c e)

mkBsE :: [ABind] -> [CoreBind]
mkBsE [] = []
mkBsE ((ANonRec b e):t) = (NonRec b (mkE e)):mkBsE t

plugA :: AContext -> AExpr -> CoreExpr
plugA c e = go c (mkAE e)
    where go [] e = e
          go ((AAppCtxt arg):c) e = go c (mkApps e [mkE arg])
          go ((APrimOpCtxt oes ies):c) e = go c (mkApps (mkE $ head oes) ((map mkE (tail oes))  ++ e:map mkE ies))
          go ((ACaseCtxt v t alts):c) e = go c (Case e v t (map goAlt alts))
          go ((ACastCtxt co):c) e = go c (Cast e co)

          goAlt (c, bs, e) = (c, bs, mkE e)

mkAE :: AExpr -> CoreExpr
mkAE (AVar x) = Var x
mkAE (ALit l) = Lit l
mkAE (ALam b e) = Lam b (mkE e)
mkAE (AType t) = Type t

type Context = [CFrame]

data CFrame = AppCtxt InExpr
            | PrimOpCtxt [OutExpr] [InExpr] 
            | CaseCtxt Var Type [CoreAlt]
            | CastCtxt Coercion


plug :: Context -> CoreExpr -> CoreExpr
plug [] e = e
plug ((AppCtxt arg):c) e = plug c (mkApps e [arg])
plug ((PrimOpCtxt (v:oes) ies):c) e = plug c (mkApps v (oes ++ e:ies))
plug ((CaseCtxt b t alts):c) e = plug c (Case e b t alts)
plug ((CastCtxt co):c) e = plug c (Cast e co)

depth :: Context -> Int
depth = length

emptyContext :: Context
emptyContext = []

emptyAContext :: AContext
emptyAContext = []

makePrimOpCtxt :: Var -> Context -> (CoreExpr, Context)
makePrimOpCtxt v c = go (PrimOpCtxt [Var v] []) c
 where go (PrimOpCtxt is oes) ((AppCtxt arg):c') = go (PrimOpCtxt is (oes ++ [arg])) c'
       go (PrimOpCtxt is oes) c = (head oes, (PrimOpCtxt is (tail oes)):c)


-- XXXpj: Why is PrimOpCtxt matched here?
splitCaseCtxt :: CoreExpr -> Context -> Maybe (CoreExpr, Context)
splitCaseCtxt _ [] = Nothing
splitCaseCtxt e ((AppCtxt arg):c) = splitCaseCtxt (App e arg) c
splitCaseCtxt _ ((PrimOpCtxt {}):_) = Nothing
splitCaseCtxt e c@((CaseCtxt {}):_) = Just (e, c)
splitCaseCtxt e ((CastCtxt co):c) = splitCaseCtxt (Cast e co) c

collectCtxtArgs :: Context -> Maybe ([CoreExpr], Context)
collectCtxtArgs ((AppCtxt arg):c) = go [arg] c
    where go es ((AppCtxt arg):c) = go (arg:es) c
          go es c = Just (reverse es, c)
collectCtxtArgs _ = Nothing

splitTerm :: CoreExpr -> (CoreExpr, Context, [CoreBind])
splitTerm e = splitTerm' e emptyContext

splitTerm' :: CoreExpr -> Context -> (CoreExpr, Context, [CoreBind])
splitTerm' e c = (oute, outc, reverse outbs)
    where (oute, outc, outbs) = go e c []
          go v@(Var x) c bs = (v, c, bs)
          go e@(Lam {}) c bs = (e, c, bs)
          go (App e1 e2) c bs = go e1 ((AppCtxt e2):c) bs
          go e@(Let b1 body) c bs = go body c (b1:bs)
          go e@(Case scrut b t alts) c bs = go scrut ((CaseCtxt b t alts):c) bs
          go e@(Cast e' co) c bs = go e' ((CastCtxt co):c) bs
          go (Note _ e) c bs = go e c bs
          go e@(Lit {}) c bs = (e, c, bs)
          go e@(Type {}) c bs = (e, c, bs)

data RhoElement = RhoE {
      freshName :: Var,
      inFvs :: [Var],
      stateHack :: Bool,
      zExp :: ZExpr,
      restExp :: Context,
      headExp :: CoreExpr,
      compSize :: Integer
}

data MemoHead = MemoH {
      varMap :: FiniteMap Var [RhoElement],
      varHeadMap :: [RhoElement],
      expMap :: [RhoElement]
    }                
                

type Store = [(RhoElement, CoreExpr)]

data ScpEnv = ScpE {
      ls :: MemoHead, -- This is \rho
--      scp_subst :: Subst,
      hasH :: Bool,
      inSet :: InScopeSet, 
      allBinds :: FiniteMap Var CoreExpr,    -- This is \mathcal{F}
      binds :: FiniteMap Var CoreExpr    -- This is \mathcal{F} lambda lifted
    }

env0 :: ScpEnv
env0 = ScpE { 
	   ls = MemoH {
                varMap = emptyFM,
                varHeadMap = [],
                expMap = []
                },
           hasH = False,
           inSet = emptyInScopeSet,
           allBinds = emptyFM,
           binds = emptyFM
         }

initScpEnv :: DynFlags -> [CoreBind] -> ScpEnv
initScpEnv dflags b = env0 {allBinds = listToFM (flattenBinds b)}

extendSimilarExprs :: ScpEnv -> FiniteMap (Var, Int) [[CoreExpr]] -> (CoreExpr, Int) -> CoreExpr -> (CoreExpr, Context) -> ScpM s ()
extendSimilarExprs env m (v@(Var fun), i)  e e' = do
  let key = (fun, i)
      se = lookupWithDefaultFM m ([[]]::[[CoreExpr]]) key
--  trace ("Extending" ++ (showSDoc $ ppr fun)) (return ())
  se' <- updateSimilarExprsElem env se e e'
  putSimilarExprs (addToFM m key se')
  return ()

updateSimilarExprsElem :: ScpEnv -> [[CoreExpr]] -> CoreExpr -> (CoreExpr, Context) -> ScpM s [[CoreExpr]]
updateSimilarExprsElem _ [] e _ = return [[e]]
updateSimilarExprsElem _ [[]] e _ = return [[e]]
updateSimilarExprsElem env ((h@(e1:_)):t) e p@(fun@(Var f), ctxt) 
 | shallowEqualCtxt ctxt' ctxt = do
  tmp <- msg env f (fun, ctxt, []) (fun', ctxt', [])
  
  case tmp of
--    Just (_, _, _, tp) | trace (if length ctxt == 6 && occNameString (nameOccName (idName f)) == "gcdInteger2" then "Testing" ++ (showSDoc $ ppr (plug ctxt' fun') <+> ppr (plug ctxt fun) <+> ppr tp) else "") $  similarEnough env tp -> do
    Just (_, _, _, tp) |  similarEnough env tp -> do
                                                  tmp' <- msg env f (fun', ctxt', []) (fun, ctxt, [])
                                                  case tmp' of
                                                   Just (_, _, _, tp') | similarEnough env tp' -> return ((e:h):t)
                                                   _ -> do
                                                     l' <- updateSimilarExprsElem env t e p
                                                     return (h:l')
 

-- return ((e:h):t) -- if (length h > 3) then trace ("lots of: " ++ (showSDoc (ppr (length h) <+> ppr fun))) $ return ((e:h):t) else return ((e:h):t)
--                       | otherwise -> do
--                           if length ctxt == 11 then  trace ("not similar:" ++ (showSDoc (ppr fun <+> ppr (length ctxt) <+> ppr tp))) (return ()) else return ()
--                           l' <- updateSimilarExprsElem env t e p
--                           return (h:l')
    _ -> do
      l' <- updateSimilarExprsElem env t e p
      return (h:l')
 | otherwise = do
  l' <- updateSimilarExprsElem env t e p
  return (h:l')
    where (fun', ctxt', []) = splitTerm e1

shallowEqualCtxt :: Context -> Context -> Bool
shallowEqualCtxt [] [] = True
shallowEqualCtxt ((AppCtxt {}):t1) ((AppCtxt {}):t2) = shallowEqualCtxt t1 t2
shallowEqualCtxt ((CastCtxt {}):t1) ((CastCtxt {}):t2) = shallowEqualCtxt t1 t2
shallowEqualCtxt ((CaseCtxt {}):t1) ((CaseCtxt {}):t2) = shallowEqualCtxt t1 t2
shallowEqualCtxt ((PrimOpCtxt {}):t1) ((PrimOpCtxt {}):t2) = shallowEqualCtxt t1 t2
shallowEqualCtxt _ _ = False

similarEnough :: ScpEnv -> TermParts -> Bool
similarEnough env (_, es) = all (simpleExpression env) es

simpleExpression :: ScpEnv -> (Var, (InScopeSet, CoreExpr)) -> Bool
simpleExpression env (v, (iss, e)) = simpleExpression' (env {inSet = iss}) e

simpleExpression' :: ScpEnv -> CoreExpr -> Bool
simpleExpression' env (Var v) | elemInScopeSet v (inSet env) = True
simpleExpression' env (Lam b e) = simpleExpression' env' e
    where env' = env {inSet = extendInScopeSet (inSet env) b}
simpleExpression' _ (Lit {}) = True
simpleExpression' _ (Type _) = True
simpleExpression' env e@(App e1 e2)
 | simpleExpression' env e1
 , simpleExpression' env e2 = True
 | (Var fun, es) <- ps
 , PrimOpId _ <- idDetails fun = all (simpleExpression' env) es
 | (Var fun, es) <- ps
 ,  elem (occNameString (nameOccName (idName fun))) primopFunctions
  = -- case es of
--     [Var x, n] | numLike env n && elemInScopeSet x (inSet env) -> trace ("Discarding " ++ (showSDoc $ ppr e)) True
--     [n, Var x] | numLike env n && elemInScopeSet x (inSet env) -> trace ("Discarding " ++ (showSDoc $ ppr e)) True
--    -> 
      all (simpleExpression' env) es
  where ps = collectArgs e
simpleExpression' env (Cast e _) = simpleExpression' env e
simpleExpression' env (Note _ e) = simpleExpression' env e
simpleExpression' _ _ = False

extendLs :: MemoHead -> RhoElement -> MemoHead
extendLs h@(MemoH { varMap = m, ..}) e@(RhoE {..}) = case headExp of
                                      Var v -> newVarMap v
                                      _ -> h { expMap = e:expMap}
    where newVarMap v = h {varMap = addToFM_C (flip (++)) m v [e]}

extendLs' :: MemoHead -> RhoElement -> MemoHead
extendLs' h@(MemoH { varHeadMap = m,..}) e = h { varHeadMap = e:m }

getLs :: Bool -> MemoHead -> CoreExpr -> [RhoElement]
getLs b (MemoH {..}) (Var v) | b = varHeadMap
                             | otherwise = lookupWithDefaultFM varMap [] v
getLs _ (MemoH {..}) _ = expMap
                       
-- scpSubstId :: ScpEnv -> Id -> Maybe CoreExpr
-- scpSubstId _ _ = Nothing
-- scpSubstId env v = case lookupIdSubst (scp_subst env) v of
--                      e@(Var v') -> if v == v' then Nothing else Just e
--                      e -> Just e

-- scpSubstTy :: ScpEnv -> Type -> Type
-- scpSubstTy env ty = substTy (scp_subst env) ty

-- zapScpSubst :: ScpEnv -> ScpEnv
-- zapScpSubst env = env { scp_subst = zapSubstEnv (scp_subst env)} 

-- extendScpInScope :: ScpEnv -> [Var] -> ScpEnv
-- 	-- Bring the quantified variables into scope
-- extendScpInScope env qvars = env'
--     where env' = env { scp_subst = extendInScopeList (scp_subst env) qvars,
--                      inSet = qvars ++ inSet env}

	-- Extend the substitution
-- extendScpSubst :: ScpEnv -> Var -> OutExpr -> ScpEnv
-- extendScpSubst env var expr = env'
--     where env' = env { scp_subst = extendSubst (scp_subst env) var expr,
--                      inSet = var:inSet env }

-- extendScpSubstList :: ScpEnv -> [(Var,OutExpr)] -> ScpEnv
-- extendScpSubstList env prs = env'
--     where env' = env { scp_subst = extendSubstList (scp_subst env) prs,
--                      inSet = (map fst prs) ++ inSet env }

-- extendBndr :: ScpEnv -> Var -> (ScpEnv, Var)
-- extendBndr env bndr = (env { scp_subst = subst', inSet = bndr':inSet env}, bndr')
--     where
--       (subst', bndr') = substBndr (scp_subst env) bndr

-- extendRecBndrs :: ScpEnv -> [Var] -> (ScpEnv, [Var])
-- extendRecBndrs env bndrs = (env { scp_subst = subst', 
--                                   inSet = bndrs' ++ inSet env }, bndrs')
--     where
--       (subst', bndrs') = substRecBndrs (scp_subst env) bndrs

pjBndr :: ScpEnv -> Var -> (ScpEnv, Maybe (Subst, Var))
pjBndr env bndr = case pjBndrs env [bndr] of
                    (env, Just (s', bndrs)) -> (env, Just (s', head bndrs))
                    (env, Nothing) -> (env, Nothing)

pjBndrs :: ScpEnv -> [Var] -> (ScpEnv, Maybe (Subst, [Var]))
pjBndrs env bndrs
 | collision = (env, Just (go env (mkSubst (inSet env) emptyVarEnv emptyVarEnv, []) bndrs))
 | otherwise = (env', Nothing)
  where collision = any ((flip elemInScopeSet) (inSet env)) bndrs

        go _ (s, b) [] = (extendInScopeList s b, reverse b)
        go env (s, b) (h:t) | elemInScopeSet h (inSet env)  = go (env {inSet = extendInScopeSet (inSet env) h'}) (s', h':b) t
                            | otherwise = go (env {inSet = extendInScopeSet (inSet env) hc}) (s, hc:b) t
                            where hc = h -- | isTyVar h = h 
--                                     | otherwise = zapIdOccInfo (h `setIdUnfolding` NoUnfolding)
                                  h' = uniqAway iss hc
                                  iss = inSet env
                                  s' | isTyVar h = extendTvSubst s hc (mkTyVarTy h')
                                     | otherwise = extendSubst s hc (Var h')
        env' = env {inSet = extendInScopeSetList (inSet env) bndrs}

{-

The main recursive function, D[]

-}


drive :: ScpEnv -> CoreExpr -> Context -> ScpM Store OutExpr

drive env (Lit l) ((CaseCtxt b t alts):c) = do
  let (_, _, rhs) = findAlt (LitAlt l) alts
                    `orElse` (DEFAULT, [], mkImpossibleExpr (coreAltsType alts))
      newrhs = substExpr (newExprSubst env b (Lit l)) rhs
  scpLog 6 "R7:" (ppr (checkType "CL" (Case (Lit l) b t alts) newrhs))
  e' <- drive env newrhs c
  cheapSaving
  return e'
drive env (Lit l) context@((PrimOpCtxt ((Var fun):oes) []):c) | all isValue oes = do
  scpLog 4 ("R8 (" ++ show (length oes) ++ "):") (ppr (plug context (Lit l)))
  scpLog 4 "l: " (ppr l)
  scpLog 4 "oes: " (ppr oes)
  let rule_base = localRuleBase
  let rules = getRules rule_base fun
      iss = inSet env
      -- XXXpj: This might fail. Print out a debug message when it does.
  case lookupRule (const True) (const NoUnfolding) iss fun (oes ++ [Lit l]) rules of
     Just (_, rule_rhs) -> do
                            e' <- drive env rule_rhs c
                            cheapSaving
			    return e'
     Nothing -> build env (Lit l) context
--  scpLog 5 "R7" (ppr (checkType "R7" l rule_rhs))
 
drive env l@(Var v) context
 | Just _ <- isDataConId_maybe v,
   Just (e, (CaseCtxt b _ alts):c') <- splitCaseCtxt l context,
   Just (dc, e') <- exprIsConApp_maybe e = do
  scpLog 5 "R27/R28:" (ppr e)
  let (dc', bs, rhs) = findAlt (DataAlt dc) alts
                     `orElse` (DEFAULT, [], mkImpossibleExpr (coreAltsType alts))
      n_drop_tys = case dc' of
                     DataAlt dc'' -> length (dataConUnivTyVars dc'')
                     _ -> 0
      newargs = drop n_drop_tys e'

      newe = case bs of
               [] -> mkLets [NonRec b e] rhs
               _ -> let (newbs, rhs') = case pjBndrs env bs of
                                          (_, Just (s, bs')) -> (bs', substExpr s rhs)
                                          _ -> (bs, rhs)
                    in mkLets [NonRec b e] (mkLets [NonRec b r | (b, r) <- (zip newbs newargs)] rhs')

  scpLog 4 "Case selection, new branch:" (ppr newe)
--  scpLog 5 "Binds in CASE:" (ppr (checkType "BCC" c newe))
  ret <- drive env newe c'
  expensiveSaving
  return ret
 | PrimOpId _ <- idDetails v = do
  let (e, context') = makePrimOpCtxt v context
  scpLog 4 "PrimopCtxt" (ppr v <+> ppr e)
  drive env e context'
#ifdef FOLDVARS
 | not (null context)
 , Just _ <- lookupInScope (inSet env) v = maybeFoldVar env v context
#endif
 | otherwise = do
  bs <- getBinds
  taint <- getTaintedExprs
  let env' = env {binds = bs}
  case (lookupWithDefaultFM taint 0 v >= taintThreshold, maybeInline env' v) of
    (False, Just b) | not (boringContext env context) -> do
               (body', env'') <- lambdaLift env' b v
               scpLog 4 "R3:" (ppr v)
               scpLog 3 "R3 ctxt:" (if slask context then ppr (plug context (Var v)) else empty)
               maybeFold env'' v context body'
--	   | otherwise -> trace ("Boring body: " ++ (showSDocDebug $ ppr l) ++ (showSDoc $ ppr l <+> ppr b)) $ build env l context
    _ -> do
                scpLogDebug 3 "No unfolding:" (ppr v)
                build env l context
drive env c@(Cast e co1) ((CastCtxt co2):c') = do
  let co' = mkTransCoercion co1 co2
  scpLog 6 "R6:" (ppr (checkType "R9" (Cast c co2) (Cast e co')))
  ret <- drive env e ((CastCtxt co'):c')
  cheapSaving
  return ret
drive env (Cast e co) context
 | isIdentityCoercion co = drive env e context
 | otherwise = drive env e ((CastCtxt co):context)
drive env l@(Lam b e) c 
 | Just (s, b') <- s' = do
  let (s', bs') = substBndr s b
      newexp = Lam bs' (substExpr s' e)
  drive env newexp c
 | ((AppCtxt _):_) <- c = do
  let Just (args, newc) = collectCtxtArgs c
      newexp = doBeta env l args
  scpLog 6 "LamBeta" (ppr (checkType "LamBeta" newexp (mkApps l args)))
  ret <- drive env newexp newc
  cheapSaving
  return ret
 | ((CastCtxt co):c') <- c,
   Just (args, newc) <- collectCtxtArgs c' = do
  scpLog 5 "R17/R18/R19:" (ppr l)
  let newexp = doBeta env (Cast l co) args
  scpLog 6 "LamCBeta" (ppr (checkType "LamBeta" newexp (mkApps (Cast l co) args)))
  ret <- drive env newexp newc
  cheapSaving
  return ret
 | otherwise = do
  replaceSavings SZero
  e' <- drive (env' {hasH = False}) e emptyContext
  build env (Lam b e') c
    where (env', s') = pjBndr env b
drive env l@(Let (NonRec b e) body) c
 | Just (s, b') <- s' = do
  let (s', bs') = substBndr s b
      newexp = Let (NonRec bs' (substExpr s' e)) (substExpr s' body)
  drive env newexp c
-- | isValue e || linear b body || exprIsTrivial e || exprIsCheap e = do
 | linear b body || exprIsTrivial e = do
  let newexp = substExpr (newExprSubst env b e) body
  scpLog 4 ("R22/R23/R24:" ++ (showSDocDebug $ ppr b)) (ppr l)
  scpLog 6 "R22/R23/R24:" (ppr (checkType "R22" l newexp))
  ret <- drive env newexp c
  cheapSaving
  return ret
 | otherwise = do
  scpLog 4 ("R23/R24(nl):" ++ (showSDocDebug $ ppr b)) (ppr l)
  replaceSavings SZero
  e' <- drive (env' {hasH = False}) e emptyContext
  sav <- getSavings
  replaceSavings SZero
  scpLog 6 "R23(nl):" (ppr (checkType "R23nle" e e'))
  body' <- drive env' body c
  sav' <- getSavings
  replaceSavings (SSum [sav, sav'])
  scpLog 6 "R23(nl):" (ppr (checkType "R23nlb" (plug c body) body'))
  let t = exprType e
--  n' <- newName b t
  let newexp | isUnLiftedType t = Case e' b (exprType body') [(DEFAULT, [], body')]
             | otherwise = Let (NonRec b e') body'
  return newexp
   where (env', s') = pjBndr env b 
drive env l@(Let b@(Rec p) body) c = error "This should not happen, letrec"
drive env (Note _ e) context = drive env e context
drive env (Case e b t alts) context 
 | Just (s, b') <- s' = do
  let alts' = map (\(c, bs, e) -> let (s', bs') = substBndrs s bs in (c, bs', substExpr s' e)) alts
      newexp = Case e b' t alts'
  drive env newexp context
 | otherwise = drive env' e ((CaseCtxt b t alts):context)
   where (env', s') = pjBndr env b 
drive env (App e1 e2) c = drive env e1 ((AppCtxt e2):c)
drive env e c = do
  scpLog 5 "Fallthrough:" (ppr (plug c e))
  build env e c

build :: ScpEnv -> OutExpr -> Context -> ScpM Store OutExpr
build env e ((PrimOpCtxt comp@(o:oes) ies):c)
 | hasH env = do -- Make sure our context does not contain h functions
  scpLog 2 "h function in context" (ppr e)
  sav' <- getSavings
  (ies', savb) <- mapAndUnzipM (\e -> do replaceSavings SZero; ret <- drive (env {hasH = False}) e emptyContext; sav <- getSavings; return (ret, sav)) ies
  replaceSavings (SSum (sav':savb))
  build (env {hasH = True}) (mkApps o (oes ++ e:ies')) c
 | null ies = build env (mkApps o (oes ++ [e])) c
 | all isValue (e:oes) = drive env (head ies) ((PrimOpCtxt (comp ++ [e]) (tail ies)):c)
 | otherwise = do -- Make sure we don't repeatedly fold against ourselves
  scpLog 2 "Variable in prim context" (ppr e)
  sav' <- getSavings
  (ies', savb) <- mapAndUnzipM (\e -> do replaceSavings SZero; ret <- drive env e emptyContext; sav <- getSavings; return (ret, sav)) ies
  replaceSavings (SSum (sav':savb))
  build env (mkApps o (oes ++ e:ies')) c
build env e ((AppCtxt arg):c) = do
  sav <- getSavings
  replaceSavings SZero; 
  arg' <- drive (env {hasH = False}) arg emptyContext
  sav' <- getSavings
  replaceSavings (SSum [sav, sav'])
  build (env {hasH = True}) (mkApps e [arg']) c
build env e ((CastCtxt co):c) = build env (Cast e co) c
build env l@(Var x) con@((CaseCtxt b t alts):c) 
 | Just _ <- lookupInScope (inSet env) x = do
  scpLogDebug 5 "VarSwitching:" (ppr x)
  scpLogDebug 5 "VarSwitching2:" (ppr (realFvs env (plug con (Var x))))
  let s = newExprSubst env b l
  let alts' = map (\(c, bs, e) -> let (s', bs') = substBndrs s bs in (c, bs', substExpr s' e)) alts
      env' = env {hasH = False}
  oldsav <- getSavings
  (alts'', savb) <- mapAndUnzipM (\e -> do replaceSavings SZero; ret <- driveAlt env' x c e; sav <- getSavings; return (ret, sav)) alts'
  let t' = exprType (trd3 . head $ alts'')
--      newcase = (Case (Var x) (zapIdOccInfo (b `setIdUnfolding` NoUnfolding)) t' alts'')
      newcase = (Case l b t' alts'')
  scpLog 5 "VarSwitching'" (ppr x <+> ppr newcase)
  replaceSavings (SSum [oldsav, SChoice savb])
  return newcase
build env e ((CaseCtxt b _ alts):c) = do
  scpLog 5 "Switching:" (ppr e)
  let env' = env {hasH = False}
  oldsav <- getSavings
  (alts', savb) <- mapAndUnzipM (\e -> do replaceSavings SZero; ret <- driveAlt env' b c e; sav <- getSavings; return (ret, sav)) alts
  let t' = exprType (trd3 . head $ alts')
--  return (Case e (zapIdOccInfo (b `setIdUnfolding` NoUnfolding)) t' alts')
  replaceSavings (SSum [oldsav, SChoice savb])
  return (Case e b t' alts')
build _ e [] = do
  return e


gen :: ScpEnv -> Id -> CoreExpr -> Context -> CoreExpr -> Context -> Integer -> ScpM Store CoreExpr
gen env fun e1 c1 e2 c2 sz = do
  tmp <- msg env fun (e1, c1, []) (e2, c2, [])
  (term, c1', bs1, tps) <- case tmp of
                   Just res@(ground, c1', bs1, _) | realExprSize (mkLets bs1 (plug c1' ground)) < sz -> return res
                   _ ->
                             do
--                              let (e1', c1, []) = splitTerm e1 
                              (e1'', tps') <- split' env fun e1 c1
                              return (e1'', emptyContext, [], tps')
  let env' = env {inSet = extendInScopeSetList (inSet env) (getTpBinders tps)}
  scpLog 3 "msg ground, full" (ppr term)
  scpLog 3 "msg ground:" (ppr (plug c1' term))
  scpLog 3 "msg subst:" (ppr tps)
  term' <- drive env' (mkLets bs1 (plug c1' term)) emptyContext
  sav <- getSavings
  (tps',sav') <- driveTps env' tps
  let rterm = plugTpTerm tps' term'
  scpLog 3 "gen3:" (ppr rterm)
  replaceSavings (SSplit (sav:sav'))
  return rterm


driveTps :: ScpEnv -> TermParts -> ScpM Store (TermParts, [Savings])
driveTps env tps@(ts, s) = do
 let (bs, tmp1) = unzip s
     (ins1, _) = unzip tmp1
 (tmp1', savb) <- mapAndUnzipM (\(is, e) -> do ret <- drive (env {inSet = extendInScopeSetList is bs, hasH = False}) e emptyContext; sav <- getSavings; return (ret, sav)) tmp1
 let s' = zip bs (zip ins1 tmp1')
 return ((ts, s'), savb)

split' :: ScpEnv -> Id -> CoreExpr -> Context -> ScpM s (CoreExpr, TermParts)
split' _ fun _ [] = panic ("split' called on emtpy context: " ++ (showSDoc $ ppr fun))
split' env fun e1 c = go c emptyContext
    where 
          go ((c'@(AppCtxt e)):c) oc
           | [] <- c = do
            let tp = plug (reverse oc) e1
            x <- newName fun (exprType tp)
            return (App (Var x) e, newTpExp env x tp)
           | otherwise = go c (c':oc)
          go ((c'@(PrimOpCtxt (e:oes) ies)):c) oc
           | [] <- c = do
            let tp = plug (reverse oc) e1
            xs <- mapM (newName fun) (map exprType (tp:oes ++ ies))
            return (mkApps e (map Var xs), newTpExps env xs (tp:oes ++ ies))
           | otherwise = go c (c':oc)
          go ((c'@(CaseCtxt b t alts)):c) oc
           | [] <- c = do
            let tp = plug (reverse oc) e1
            x <- newName fun (exprType tp) 
            return (Case (Var x) b t alts, newTpExp env x tp)
           | otherwise = go c (c':oc)
          go ((c'@(CastCtxt co)):c) oc
           | [] <- c = do
            let tp = plug (reverse oc) e1
            x <- newName fun (exprType tp)
            return (Cast (Var x) co, newTpExp env x tp)
           | otherwise = go c (c':oc)
            

driveAlt :: ScpEnv -> Var -> Context -> CoreAlt -> ScpM Store CoreAlt
driveAlt env v c (con, bs, e) 
 | Just (s, bs') <- s' = do
  let (s', bs') = substBndrs s bs
      newtrip = (con, bs', substExpr s' e)
  driveAlt env v c newtrip 
 | otherwise  = do
  scpLogDebug 4 "Branch fvs" (ppr bs)
  scpLog 4 ("Driving leg:") (ppr $ plug c e)
  e' <- drive env' (substExpr (mkPatternSubst env con bs v) e) c
  return (con, bs, e')
   where (env', s') = pjBndrs env bs

mkPatternSubst :: ScpEnv -> AltCon -> [CoreBndr] -> Var -> Subst
#if PATTERNSUBST
mkPatternSubst _ DEFAULT _ _ = emptySubst
mkPatternSubst env (LitAlt l) _ v = newExprSubst env v (Lit l)
mkPatternSubst env (DataAlt con) vs v = newExprSubst env v  (mkConApp con con_args)
    where con_args = map Type inst_tys' ++ varsToCoreExprs vs
          inst_tys' = tyConAppArgs (idType v)
#else
mkPatternSubst _ _ _ _ = emptySubst
#endif

maybeFold :: ScpEnv -> Id -> Context -> CoreExpr -> ScpM Store CoreExpr
maybeFold env fun myctxt e = do
  scpLog 5 ("Unfolding found, depth " ++ show (depth myctxt)) (ppr fun)
  scpLog 5 "Unfolds to:" (ppr e)
  ml <- getStore
--  simExp <- getSimilarExprs
  let -- l = substExpr (scp_subst env) (mkApps (Var fun) (args))
      -- l' = substExpr (scp_subst env) (plug (ctxt env)  l)
      -- Normalize the expression
      l' = normalize env (plug myctxt (Var fun)) emptyContext
      ae = mkAExpr l'
      (fun', myctxt', []) = splitTerm l'
      sz = realExprSize l'
      res = renamings env False ml fun' myctxt' sz ae
--  scpLog 5 "maybeFold" (ppr (checkType "maybeFold" l unfolded))
  res' <- filterM (\(n, _, _, _) -> do b <- isTaintedExpr n; return (not b)) res
  if (not $ null res)
    then do
      if null res'
        then do
         -- trace ("res' null") (return ())
         build env (Var fun) myctxt
        else do
--      scpLog 4 "Renamings:" (ppr res)
         let (n', in_fvs, state_hack, ((ts, es), newctxt, t)) = head res'
         tmpnam <- newName fun t
         let restterm = plug newctxt (Var tmpnam)
         scpLog 2 "Rest term:" (ppr restterm)
         let in_scope = inSet env
             s = mkSubst in_scope ts es
             in_fvs' = (if state_hack then [realWorldPrimId] else []) ++ in_fvs
             newexp = substExpr s (mkVarApps (Var n') in_fvs')
         scpLogDebug 2 ("Folding" ++ (showSDocDebug $ ppr n')) (ppr fun <+> ppr n')
         scpLog 6 "Renaming TC:" (ppr (checkType "renaming" l' newexp))
         let env' = env {inSet = extendInScopeSet (inSet env) tmpnam}
         restterm' <- if null newctxt then do replaceSavings SZero; return (Var tmpnam) else drive env' (Var tmpnam) newctxt
         let newexp' = substExpr (newExprSubst env tmpnam newexp) restterm'
         newsav <- getSavings
         replaceSavings (SSplit [SFold, newsav])
         return newexp'
    else do
      scpLogDebug 1 "Inlining" (ppr fun)
--      scpLog 0 "Unfolding" (ppr e)
      scpLog 1 "No renaming:" (ppr l')
      bs <- getBinds
      let fvs = realFvs env l'
          body_ty = exprType l'
          state_hack = all isTyVar fvs && isUnLiftedType body_ty
          fvs' = (if state_hack then [realWorldPrimId] else []) ++ fvs
          full_ty = mkPiTypes ((if state_hack then [voidArgId] else []) ++ fvs) body_ty
      let env' = env {binds = bs}
#if OLD_HOMEMB
          nont = homemb env' fun' sz l'
#else
          nont = homemb env' (fun', myctxt', []) sz ae
#endif
      scpLog 5 "body_ty:" (ppr body_ty)
      scpLog 5 "full_ty:" (ppr full_ty)
      scpLogDebug 4 "Fvs:" (ppr fvs)
      if (not $ null nont)
        then do
          let (e1, c1) = head nont
          scpLog 2 "Whistle:" (ppr  l')
          scpLog 3 "Whistle against:" (ppr nont)
          scpLog 3 "First term" (ppr (plug c1 e1))
          -- Check type for g
          gen env fun fun' myctxt' e1 c1 sz
        else do
--          extendSimilarExprs env simExp (fun', length myctxt') l' (fun', myctxt')
          fname <- newName fun full_ty
          scpLog 5 "Inserting to rho:" (ppr l')
          scpLogDebug 1 "Fresh name:" (ppr fname)
          scpLog 2 "New function:" (ppr fname <+> ppr fvs <+> ppr fun 
                                    <+>  ppr (depth myctxt))

          let newquad = RhoE { freshName = fname, inFvs = fvs,
                               stateHack = state_hack,
                               zExp = ae,
                               restExp = myctxt', headExp = fun', 
                               compSize = sz}
              env'' = env' {ls = extendLs (ls env) newquad}
          scpLog 4 "DriveApp:" (ppr e)
          scpLog 4 "DriveApp2:" (ppr (plug myctxt e))
          scpLog 4 "DriveApp3:" (ppr (mkE ae))
	  oldStore <- getStore
	  reallyOldBinds <- getBinds
          e' <- drive env'' e myctxt
	  newStore <- getStore
	  sav <- getSavings
--	  trace ("Savings for " ++ (showSDoc $ ppr fname <+> ppr sav)) (return ())
--	  trace ("Saving term: " ++ (showSDoc $ ppr l')) (return ())
--	  trace ("Saved body:" ++ (showSDoc $ ppr e')) (return ())
	  let totsav = calculateSavings sav sz (realExprSize e')
	  if totsav <= savingsThreshold
            then do
	          scpLog 1 "Discarding term" (ppr fun <+> ppr totsav) -- <+> ppr (calculateSavings' sav))
                  scpLog 1 "Body: " (ppr l')
		  let purgedstore | null newStore = newStore
		                  | otherwise = pruneStore (head oldStore) newStore [fname]
                  addToStore (newquad, e')
                  replaceSavings SZero
--		  replaceStore purgedstore
--		  putBinds reallyOldBinds
                  incTaintedExpr fun
                  build env (Var fun) myctxt
            else do
	          scpLog 1 "Keeping term" (ppr fun <+> ppr totsav <+> ppr (calculateSavings' sav))
--                  scpLog 0 "Body: " (ppr l')
                  addToStore (newquad, e')
	          replaceSavings (SEnables (findFold sav))
                  scpLog 6 "DriveApp:" (ppr (checkType "Driveapp" (plug myctxt (Var fun)) e'))
                  return (mkVarApps (Var fname) fvs')

#ifdef FOLDVARS
-- Like maybeFold, but for variables in the expression head.
maybeFoldVar :: ScpEnv -> Id -> Context -> ScpM Store CoreExpr
maybeFoldVar env fun myctxt = do
  ml <- getStore
  let -- Normalize the expression
      l' = normalize env (plug myctxt (Var fun)) emptyContext
      ae = mkAExpr l'
      (fun', myctxt', []) = splitTerm l'
      sz = realExprSize l'
      res = findNullCtxt (renamings env True ml fun' myctxt' sz ae)
      (n', in_fvs, state_hack, ((ts, es), newctxt, t)) = head res
--  scpLog 5 "maybeFold" (ppr (checkType "maybeFold" l unfolded))
  if (not (null res))
    then do
--      scpLog 4 "Renamings:" (ppr res)
      let in_scope = inSet env
          s = mkSubst in_scope ts es
          in_fvs' = (if state_hack then [realWorldPrimId] else []) ++ in_fvs
          newexp = substExpr s (mkVarApps (Var n') in_fvs')
      scpLogDebug 2 "VarFolding" (ppr fun <+> ppr n')
      scpLog 2 "VarFolding2" (ppr l')
      scpLog 2 ("VarFolding3" ++ (showSDocDebug $ ppr n')) (ppr newexp)
      scpLog 6 "Renaming TC:" (ppr (checkType "renaming" l' newexp))
      replaceSavings SFold
      return newexp
    else do
      let fvs = realFvs env l'
          body_ty = exprType l'
          state_hack' = all isTyVar fvs && isUnLiftedType body_ty
          fvs' = (if state_hack' then [realWorldPrimId] else []) ++ fvs
          full_ty = mkPiTypes ((if state_hack' then [voidArgId] else []) ++ fvs) body_ty
          boxedtuples = any isUnboxedTupleType (map exprType (map Var fvs))
      fname <- newName fun full_ty
      scpLog 5 "Inserting to var-rho:" (ppr l' <+> ppr body_ty)
      scpLog 5 "Inserting to var-rho (t):" (ppr (exprType (Var fun)))
      scpLogDebug 1 "Fresh var-name:" (ppr fname)
      scpLog 2 "New function var:" (ppr fname <+> ppr fvs <+> ppr fun 
                                <+>  ppr (depth myctxt))
      scpLog 2 "New function var comp:" (ppr ((mkVarApps (Var fname) fvs') :: CoreExpr))

      let newquad = RhoE { freshName = fname, inFvs = fvs,
                           stateHack = state_hack',
                           zExp = ae,
                           restExp = myctxt', headExp = fun', 
                           compSize = sz}
          env' | boxedtuples = env
               | otherwise = env {ls = extendLs' (ls env) newquad}
      scpLogDebug 4 "DriveApp':" (ppr fun)
      scpLog 4 "DriveApp2':" (ppr (plug myctxt (Var fun)))
      scpLog 4 "DriveApp3':" (ppr (mkE ae))
      oldStore <- getStore
      e' <- build env' (Var fun) myctxt
      newStore <- getStore
      sav <- getSavings
      let totsav = calculateSavings sav sz (realExprSize e')
      case boxedtuples of
        False -> do
                  addToStore (newquad, e')
                  if totsav <= savingsThreshold
                    then do
                      scpLogDebug 1 "Discarding var-term" (ppr fun <+> ppr fname)
		      let purgedstore | null newStore = newStore
		                      | otherwise = pruneStore (head oldStore) newStore [fname]
                      replaceSavings SZero
--		      replaceStore purgedstore
                      build env (Var fun) myctxt
                    else do
	              replaceSavings (SEnables (findFold sav))
                      scpLog 6 "DriveApp:" (ppr (checkType "Driveapp" (plug myctxt (Var fun)) e'))
                      return (mkVarApps (Var fname) fvs')
        _ -> do
              scpLogDebug 1 "Discarding var-term" (ppr fun <+> ppr fname)
              return e'
                 

findNullCtxt ::  [(Var, [Var], Bool, (SubstEnv, Context, Type))] -> [(Var, [Var], Bool, (SubstEnv, Context, Type))]
findNullCtxt [] = []
findNullCtxt (v@(_, _, _, (_, c, _)):t) | null c = v:findNullCtxt t
                                        | otherwise = findNullCtxt t
#endif

#if PARALLEL_SCP
renamings :: ScpEnv -> Store -> CoreExpr -> Context -> Integer -> [(Var, [Var], Bool, (SubstEnv, Context, Type))]
renamings env ml head con msize = (concat $ zipWith fixup1 b1list ml) ++ (concat $ zipWith fixup2 b2list myls)

  where
     in_scope = inSet env 
     myls = getLs (ls env) head
     -- See comment above "data MatchEnv" in Unify.lhs
     lmenv = mkRnEnv2 in_scope
     renaming t1 in_fvs t2 = match menv emptySubstEnv t1 t2 
          where menv = ME { me_env = lmenv
	                   , me_tmpls = mkVarSet in_fvs }
     b1list = parBufferWHNF 100 (map p1 ml)
     b2list = parBufferWHNF 100 (map p2 myls)
--     p1 :: (RhoElement, CoreExpr) -> Maybe SubstEnv
     p1 (RhoE {..}, _) | msize == compSize,
                         head `weakUnsoundEqExpr` headExp,
                         s@(Just (_, [], _)) <- renaming (headExp, restExp,[]) inFvs (head, con, []) = s
                       | otherwise = Nothing
--     p2 :: RhoElement -> Maybe SubstEnv
     p2 (RhoE {..}) | msize == compSize,
                      Just s <- renaming (headExp, restExp, []) inFvs (head, con, []) = Just s
                    | otherwise = Nothing
--     fixup1 :: Maybe SubstEnv -> (RhoElement, CoreExpr) -> [(Var, [Var], SubstEnv)]
     fixup1 (Just s) ((RhoE {..}), _) = [(freshName, inFvs, stateHack, s)]
     fixup1 Nothing _ = []
--     fixup2 :: Maybe SubstEnv -> RhoElement -> [(Var, [Var], SubstEnv)]
     fixup2 (Just s) (RhoE {..}) = [(freshName, inFvs, stateHack, s)]
     fixup2 Nothing _ = []
#else
renamings :: ScpEnv -> Bool -> Store -> CoreExpr -> Context -> Integer -> ZExpr -> [(Var, [Var], Bool, (SubstEnv, Context, Type))]
renamings env varHead ml head con msize ze = 
  [(freshName, inFvs, stateHack, s) | (RhoE {..}, _) <- ml,
                                        msize == compSize,
--                                        head `weakUnsoundEqExpr` headExp,
                                        Just s <- [renaming zExp inFvs ze] ] ++  rhs
  where
     rhs =   [(freshName, inFvs, stateHack, s) | RhoE {..} <- getLs varHead (ls env) head,
                                        msize == compSize,
                                        Just s <- [renaming zExp inFvs ze] ]

     in_scope = inSet env 
     -- See comment above "data MatchEnv" in Unify.lhs
     lmenv = mkRnEnv2 in_scope 
     renaming e1 in_fvs e2 = match menv emptySubstEnv e1 e2 
          where menv = ME { me_env = lmenv
	                   , me_tmpls = mkVarSet in_fvs }
#endif

scpProgram :: ModGuts -> CoreM ModGuts
scpProgram mod_guts = do
        dflags <- getDynFlags
        us <- getUniqueSupplyM
        rule_base <- getRuleBase
--        liftIO $ showPass dflags "Scp2"

        let main_is_flag = mainFunIs dflags
            main_fn' = case main_is_flag of
 		  Just fn -> mkRdrUnqual (mkVarOccFS (mkFastString fn))
 		  Nothing -> main_RDR_Unqual
            main_name = mainName mod_guts main_fn'
--            e = findMain mod_guts main_fn' my_binds

--        trace (showSDoc $ ppr my_binds) (return ())
--        tr "---------------------------- End binds ----------\n"

	let binds' = initScp dflags rule_base us (findFunAndScp (initScpEnv dflags my_binds) (bindersOfBinds my_binds) main_name my_binds)
-- 	let binds' = case e of
--                        Nothing -> my_binds
--                        Just e' -> initScp dflags rule_base us (findFunAndScp (initScpEnv dflags my_binds) (bindersOfBinds my_binds) main_name e')

        -- The simplifier blindly trusts occurence information, and any
        -- occurence information still in the program is *very* wrong.
--        let allBinds = occurAnalysePgm binds' []
        let allBinds = washBinds binds'

        dumpIfSet_dyn Opt_D_dump_scp "Scp binds"
                           (pprCoreBindings allBinds)
        tr "------------ End binds ---------\n"
        tr (showSDoc $ ppr allBinds)
        tr "------------ End binds2 ---------\n"
        tr (showSDocDebug $ ppr binds')
        tr "------------ End binds3 ---------\n"

        return (mod_guts { mg_binds = allBinds})
   where
     my_binds = mg_binds mod_guts

findFunAndScp :: ScpEnv -> [Var] -> [Var] -> [CoreBind] -> ScpM Store [CoreBind]
findFunAndScp env modulebinders [main_name] binds = do
  (mfunc, newbinds) <- findFunAndScp' env modulebinders [main_name] [] binds []
  ml <- getStore
--  taint <- getTaintedExprs
  let taint = emptyFM
  let newbinds' = makeAllFunctions taint (main_name:mfunc) ml newbinds []
  return newbinds'
findFunAndScp _ _ funs binds = return binds


findFunAndScp' :: ScpEnv -> [Var] -> [Var] -> [Var] -> [CoreBind] -> [CoreBind] -> ScpM Store ([Var], [CoreBind])
-- findFunAndScp' _ _ [] _ bs c = do
--                              trace "Empty list of names" (return ())
-- 	                     ml <- getStore
--                              return $ reverse c ++ bs ++ [Rec (createBinds ml)]
findFunAndScp' _ _ _ done [] c = return (done, reverse c) -- do
--                            ml <- getStore
--                            return $ makeAllFunctions ml c ++ [Rec (createBinds ml)]
findFunAndScp' env modulebinders ns done ((NonRec b e):rest_binds) c | b `elem` ns = do
                          (e', env') <- lambdaLift env (False, e) b
                          scpLog 1 ("Driving nonrec fun") (ppr b) -- <+> ppr e')
			  oldStore <- getStore
                          replaceSavings SZero
                          exp <- drive env' e' emptyContext
                          ml <- getStore
                          taint <- getTaintedExprs
                          let ns' = createCandidates taint modulebinders (delete [b] ns) (b:done) exp oldStore ml
--			  trace ("New names: " ++ (showSDoc $ ppr ns')) (return ())
--                          trace ("Storelength:" ++ (show $ length ml)) (return [])
			  findFunAndScp' env modulebinders ns' (b:done) (reverse c ++ (NonRec b exp):rest_binds) []

--                          simExp <- getSimilarExprs
--                          trace "Beraknar langder" (return ())
--			  trace ("Store:\n" ++ (showSDoc $ (ppr (printPrep ml)))) (return ())
--		          trace ("Lengths:\n" ++ (showSDoc $ (ppr (slask2 simExp)))) (return ())
--                          trace "Beraknar langder igen" (return ())
--			  trace ("Lengths2:\n" ++ (showSDoc $ (ppr (slask3 simExp)))) (return ())
--                          return $ rest_binds ++ [(Rec (createBinds ml)), (NonRec b exp)]
findFunAndScp' env modulebinders ns done (p@(Rec bs):rest_binds) c | any ((flip elem) ns) (bindersOf p) = do
--	      	          oldStore <- getStore
                          bs' <- findAndTransformMain bs
			  let compiled = intersect ns (bindersOf p)
			      ns' = delete compiled ns
--                          ml <- getStore
--                          return $ reverse c ++ rest_binds ++ [(Rec (createBinds ml)), (Rec bs')]
			  findFunAndScp' env modulebinders ns' (compiled ++ done) (reverse c ++ (Rec bs'):rest_binds) []
         where findAndTransformMain :: [(Var, CoreExpr)] -> ScpM Store [(Var, CoreExpr)]
               findAndTransformMain [] = return []
               findAndTransformMain ((n, e):binds) = do
                                              if elem n ns
                                                 then do
                                                   (e', env') <- lambdaLift env (False, e) n
                                                   scpLog 1 ("Driving rec fun") (ppr n) --  <+> ppr e')
                                                   replaceSavings SZero
                                                   exp <- drive env' e' emptyContext
                                                   return ((n, exp):binds)
                                                 else do
                                                   binds' <- findAndTransformMain binds
                                                   return ((n, e):binds')
findFunAndScp' env modulebinders ns done (h:t) c = findFunAndScp' env modulebinders ns done t (h:c)


-- findMain :: ModGuts -> RdrName -> [CoreBind] -> Maybe (CoreBind, [CoreBind])
-- findMain mod_guts main_fn binds = ret p
--     where 
--           p = case main_name of
--                 Nothing -> Nothing
--                 Just n -> go binds n

--           ret Nothing = Nothing
--           ret (Just b) = Just (b, dropBind b binds)

--           go :: [CoreBind] -> Name -> Maybe CoreBind
--           go [] _ = Nothing
--           go (b:bs) n | hasfun b n = Just b
--                       | otherwise = go bs n
--               where hasfun bs n = n `elem` (map varName (bindersOf bs)) 

dropBind :: CoreBind -> [CoreBind] -> [CoreBind]
dropBind _ [] = []
dropBind b@(NonRec b1 _) (l@(NonRec b2 _):t) | b1 == b2 = dropBind b t
                                             | otherwise = l:dropBind b t
dropBind b@(Rec b1) (l@(Rec b2):t) | bs1 == bs2 = dropBind b t
                                   | otherwise = l:dropBind b t
 where bs1 = map fst b1
       bs2 = map fst b2
dropBind b (l:t) = l:dropBind b t

mainName :: ModGuts -> RdrName -> [Var]
mainName mod_guts main_fn = name -- main_name
  where env = mg_rdr_env mod_guts
        rmain = rdrNameSpace main_fn
        name = [mkLocalId name realWorldStatePrimTy | gre <- lookupGRE_RdrName main_fn env,
                              let name = gre_name gre,
                              rmain == occNameSpace (nameOccName name)]
        main_name | null name = Nothing
                  | otherwise = Just (head name)

--- XXXpj: Should this go into CoreUtils.lhs?
realNoteEq :: Note -> Note -> Bool
realNoteEq (SCC c1) (SCC c2) = c1 == c2
realNoteEq (CoreNote s1) (CoreNote s2) = s1 == s2
realNoteEq _ _ = False


newName :: Id -> Type -> ScpM s Var
newName e t = do
  uniq <- getUniqueM -- Us
  let fn_name = idName e
      fn_loc = nameSrcSpan fn_name
      spec_occ = mkSpecOcc (nameOccName fn_name)
      spec_id = mkUserLocal spec_occ uniq t fn_loc
  return spec_id

-- newTName :: Id -> Kind -> ScpM s Var
--newTName e k = do
newTName :: Kind -> ScpM s Var
newTName k = do
  uniq <- getUniqueM
  let co_name = mkSysTvName uniq (fsLit "scpa")
      spec_id = mkTyVar co_name k
  return spec_id
--      typeKind fran Coercion

newName' :: ScpEnv -> Id -> Type -> (Var, ScpEnv)
newName' env fun t = (new_name, env')
    where fn_name = idName fun
          fn_loc = nameSrcSpan fn_name
          spec_occ = mkSpecOcc (nameOccName fn_name)
          iss = inSet env
          fun' = uniqAway iss fun
          new_name = mkUserLocal spec_occ (idUnique fun') t fn_loc
          env' = env {inSet = extendInScopeSet (inSet env) new_name}

maybeInline :: ScpEnv -> Id -> Maybe (Bool, CoreExpr)
maybeInline env fun
 | not (isActive 0 (idInlineActivation fun)) = Nothing
 | elem (occNameString (nameOccName (idName fun))) forbiddenFunctions = Nothing
 | Just unf <- lookupFM (binds env) fun = Just (True, unf)
-- | isGlobalId fun,
--   Just unf <- maybeUnfoldingTemplate (realIdUnfolding fun) = let (e, bs) = lambdaLift emptyInScopeSet emptyInScopeSet unf in Just (e, (env {binds = addListToFM (binds env) ((fun, e):bs)})) -- maybeUnfoldingTemplate (realIdUnfolding fun) = trace ("maybeInline: " ++ (showSDoc $ ppr fun <+> ppr unf)) $ let (e, bs) = lambdaLift emptyInScopeSet emptyInScopeSet unf in Just (e, (env {binds = addListToFM (binds env) ((fun, e):bs)}))
 | Just unf <- lookupFM (allBinds env) fun = Just (False, unf) -- let (e, bs) = lambdaLift emptyInScopeSet emptyInScopeSet unf in Just (e, (env {binds = addListToFM (binds env) ((fun, e):bs)}))
-- | elem (occNameString (nameOccName (idName fun))) forbiddenFunctions = Nothing
 | isGlobalId fun,
   Just unf <- maybeUnfoldingTemplate (realIdUnfolding fun) = Just (False, unf) --let (e, bs) = lambdaLift emptyInScopeSet emptyInScopeSet unf in Just (e, (env {binds = addListToFM (binds env) ((fun, e):bs)}))
 | otherwise = Nothing
--  | Just e <- go binds' fun
--  , elem (occNameString (nameOccName (idName fun))) approvedFunctions = Just e -- trace (showSDoc $ ppr $ idName fun) Just e
--  | Just e <- maybeUnfoldingTemplate (idUnfolding fun)
--  , not (isLocalId fun)
--  , elem (occNameString (nameOccName (idName fun))) approvedFunctions = Just e -- trace (showSDoc $ ppr $ idName fun) Just e
--  | Just e <- go binds' fun
--  , not (elem (occNameString (nameOccName (idName fun))) forbiddenFunctions) = Just e -- trace (showSDoc $ ppr $ idName fun) Just e
--  | not (isLocalId fun),
--    Just e <- maybeUnfoldingTemplate (idUnfolding fun),
--  , not (elem (occNameString (nameOccName (idName fun))) forbiddenFunctions) = Just e -- trace (showSDoc $ ppr $ idName fun) Just e
--  | Just e <- go binds' fun = Just e
--  | Just e <- maybeUnfoldingTemplate (idUnfolding fun) = Just e
--  | otherwise = Nothing


realContextSize :: Context -> Integer
realContextSize [] = 0
realContextSize ((AppCtxt e):c) = 1 + realExprSize e + realContextSize c
realContextSize ((PrimOpCtxt oes ies):c) = 1 + sum (map realExprSize (oes ++ ies)) + realContextSize c
realContextSize ((CaseCtxt _ _ alts):c) = 1 + sum (map realAltSize alts)  + realContextSize c
realContextSize ((CastCtxt _):c) = 1 + realContextSize c

realAltSize :: CoreAlt -> Integer
realAltSize (_, _, e) = realExprSize e

realExprSize :: CoreExpr -> Integer
realExprSize (Var _) = 1
realExprSize (Lit _) = 1
realExprSize (App e1 e2) = 1 + realExprSize e1 + realExprSize e2
realExprSize (Lam _ e) = 1 + realExprSize e
realExprSize (Let (NonRec _ e) body) = 1 + realExprSize e + realExprSize body
realExprSize (Let (Rec p) body) = 1 + sum (map (realExprSize . snd) p)
                                    + realExprSize body
realExprSize (Case e _ _ alts) = 1 + realExprSize e 
                                   + sum (map (realExprSize . trd3) alts)
realExprSize (Cast e _) = 1 + realExprSize e
realExprSize (Note _ e) = 1 + realExprSize e
realExprSize (Type _) = 1

realGroupFvs :: ScpEnv -> [CoreExpr] -> [Var]
realGroupFvs env es = nub (tvs ++ fvs)
  where fvs = concatMap (realFvs' env) es
        tvs = varSetElems (tyVarsOfTypes (map idType fvs))

-- We distinguish between variables and names.
realFvs :: ScpEnv -> CoreExpr -> [Var]
realFvs env e = nub (tvs ++ fvs)
  where fvs = realFvs' env e
        tvs = varSetElems (tyVarsOfTypes (map idType fvs))

realFvs' :: ScpEnv -> CoreExpr -> [Var]
realFvs' env (Var v) | elemInScopeSet v (inSet env) = [v]
                     | otherwise = []
realFvs' _ (Lit {}) = []
realFvs' env (App e1 e2) = realFvs' env e1 ++ realFvs' env e2
realFvs' env (Lam b e) = delete [b] (realFvs' env e)
realFvs' env (Let b e') = delete (bindersOf b) (concatMap (realFvs' env) (rhssOfBind b) ++ realFvs' env e')
realFvs' env (Case e b t alts) = e1 ++ t1 ++ e2
  where e1 = realFvs' env e
        t1 = realFvs'_type env t
        e2 = delete [b] (concatMap (realFvs'_alt env) alts)
realFvs' env (Cast e co) = realFvs'_type env co ++ realFvs' env e
realFvs' env (Note _ e) = realFvs' env e
realFvs' env (Type t) = realFvs'_type env t

realFvs'_type :: ScpEnv -> Type -> [Var]
realFvs'_type env (TyVarTy v) | elemInScopeSet v (inSet env) = [v]
                              | otherwise = []
realFvs'_type env (AppTy t1 t2) = realFvs'_type env t1 ++ realFvs'_type env t2
realFvs'_type env (TyConApp tc ts) = concatMap (realFvs'_type env) ts
realFvs'_type env (FunTy t1 t2) = realFvs'_type env t1 ++ realFvs'_type env t2
realFvs'_type env (ForAllTy tv t) = delete [tv] (realFvs'_type env t)
realFvs'_type env (PredTy pt) = realFvs'_pred env pt

realFvs'_pred :: ScpEnv -> PredType -> [Var]
realFvs'_pred env (IParam _ ty)    = realFvs'_type env ty
realFvs'_pred env (ClassP _ tys)   = concatMap (realFvs'_type env) tys
realFvs'_pred env (EqPred ty1 ty2) = realFvs'_type env ty1  ++ realFvs'_type env ty2

realFvs'_alt :: ScpEnv -> CoreAlt -> [Var]
realFvs'_alt env (_, bs, e) = delete bs (realFvs env e)

localRuleBase :: RuleBase
localRuleBase = mkRuleBase (concatMap func allThePrimOps)
  where func op = primOpRules op (idName (mkPrimOpId op))

#if OLDBETA
-- XXXpj: Could merge alt 1 and alt 2 into one subst; still need 3 guards.
doBeta :: ScpEnv -> CoreExpr -> [CoreExpr] -> CoreExpr
doBeta env (Lam b body) args@(a:as) 
 | Just (s, b') <- s' = let (s'', bs') = substBndr s b in doBeta env (Lam bs' (substExpr s'' body)) args
 | isTypeArg a = doBeta env (substExpr (newExprSubst env b a) body) as
 | otherwise = Let (NonRec b a) (doBeta env' body as)
    where (env', s') = pjBndr env b
doBeta env e@(Cast (Lam b body) co) args@(a:as) 
 | Just (s, b') <- s' = let (s'', bs') = substBndr s b in doBeta env (Cast (Lam bs' (substExpr s'' body)) co) args
 | otherwise = doBeta env e' (a':as)
   where (_, s') = pjBndr env b
         (App e' a') = evalPush env (App e a)
doBeta _ fn args = mkApps fn args
#else
-- XXXpj: Could merge alt 1 and alt 2 into one subst; still need 3 guards.
doBeta :: ScpEnv -> CoreExpr -> [CoreExpr] -> CoreExpr
doBeta env e args = floatRight oute
  where 
        (oute, outbs) = go env e args [] 

        floatRight (Cast e'@(Lam {}) c) = Cast (mkLams outbss' (mkLets outbs oute'')) c
            where (outbss', oute'') = collectBinders e'
        floatRight e' = mkLams outbss' (mkLets outbs oute'')
            where (outbss', oute'') = collectBinders e'

        go env (Lam b body) args@(a:as) inbs
          | Just (s, b') <- s' = let (s'', bs') = substBndr s b in go env (Lam bs' (substExpr s'' body)) args inbs
          | isTypeArg a = go env (substExpr (newExprSubst env b a) body) as inbs
          | otherwise = go env' body as ((NonRec b a):inbs)
            where (env', s') = pjBndr env b
        go env e@(Cast (Lam b body) co) args@(a:as) inbs
          | Just (s, b') <- s' = let (s'', bs') = substBndr s b in go env (Cast (Lam bs' (substExpr s'' body)) co) args inbs
          | otherwise = go env e' (a':as) inbs
            where (_, s') = pjBndr env b
                  (App e' a') = evalPush env (App e a)
        go env fn args inbs = (mkApps (deShadowExpr env fn) args, reverse inbs)
#endif

isValue :: CoreExpr -> Bool
isValue (Var {}) = False
isValue (Lit {}) = True
isValue (Lam {}) = True
isValue (Type {}) = True
isValue (App {}) = False
isValue (Let {}) = False
isValue (Case {}) = False
isValue (Note _ e) = isValue e
isValue (Cast e _) = isValue e

{-

  The whistle (homeomorphic embedding)

-}

#if 0
parBufferWHNF n xs = return xs (start n xs)
    where return (x:xs) (y:ys) = y `par` (x:return xs ys)
          return xs [] = xs
          start n [] = []
          start 0 ys = ys
          start n (y:ys) = y `par` start (n-1) ys
#else
parBufferWHNF n xs = xs
#endif

#if PARALLEL_SCP

homemb :: ScpEnv -> (CoreExpr, Context, [CoreBind]) -> Integer -> [(CoreExpr, Context)]
homemb env e1@(head1, _, _) msize = concat $ zipWith fixup blist myls
    where myls = getLs (ls env) head1
          blist = parBufferWHNF 100 (map p myls)
          p :: RhoElement -> Bool
          p (RhoE {..}) = msize > compSize && isHomemb env (headExp, restExp, []) e1
          fixup True (RhoE {..}) = [(headExp, restExp)]
          fixup False  _ = []
#else

#if OLD_HOMEMB

homemb :: ScpEnv -> CoreExpr -> Integer -> CoreExpr -> [(CoreExpr, Context)]
homemb env head1 msize ae = [ (headExp, restExp) |
                    RhoE {..} <- getLs False (ls env) head1, msize > compSize,
                    isHomemb env (plug restExp headExp) ae]


isHomemb :: ScpEnv -> CoreExpr -> CoreExpr -> Bool
isHomemb env e1 e2 = peel env e1 e2 || any (isHomemb env e1) (dive e2)


dive :: CoreExpr -> [CoreExpr]
dive (Lam _ e) = [e]
dive (App e1 e2) = [e1, e2]
dive (Let b body) = body:rhssOfBind b
dive (Case cond _ _ alts) = cond:rhssOfAlts alts
dive (Cast e _) = [e]
dive (Note _ e) = [e]
dive (Lit {}) = []
dive (Var {}) = []
dive (Type {}) = []

peel :: ScpEnv -> CoreExpr -> CoreExpr -> Bool
peel env (Var v1) (Var v2) 
  | v1 == v2 = True
  | PrimOpId f1 <- idDetails v1
  , PrimOpId f2 <- idDetails v2 = f1 == f2
  | isJust (maybeInline env v1) || isJust (maybeInline env v2) = False
  | otherwise = True
peel _ (Lit {}) (Lit {}) = True
peel env (Lam _ e1) (Lam _ e2) = isHomemb env e1 e2
peel env e1@(App {}) e2@(App {})
  | (f1, args1) <- collectArgs e1
  , (f2, args2) <- collectArgs e2
  , length args1 == length args2 =
  isHomemb env f1 f2 && and (zipWith (isHomemb env) args1 args2)
peel env (Let (NonRec _ e1) body1) (Let (NonRec _ e2) body2) =
  isHomemb env e1 e2 && isHomemb env body1 body2 
peel env (Let b1@(Rec p1) body1) (Let b2@(Rec p2) body2)
  | length p1 == length p2 && bindersOf b1 == bindersOf b2 =
      isHomemb env body1 body2 && and (zipWith (\(_, e1) (_, e2) -> isHomemb env e1 e2) p1 p2)
peel env (Case c1 _ _ alts1) (Case c2 _ _ alts2) 
  | length alts1 == length alts2 = isHomemb env c1 c2 && 
      and (zipWith (\(_, _, e1) (_, _, e2) -> isHomemb env e1 e2) alts1 alts2)
peel env (Cast e1 _) (Cast e2 _) = isHomemb env e1 e2
peel env (Note _ e1) (Note _ e2) = isHomemb env e1 e2
peel _ (Type _) (Type _) = True
peel _ _ _ = False

#else
homemb :: ScpEnv -> (CoreExpr, Context, [CoreBind]) -> Integer -> ZExpr -> [(CoreExpr, Context)]
homemb env e1@(head1, _, _) msize ae = [ (headExp, restExp) |
                    RhoE {..} <- getLs False (ls env) head1, msize > compSize,
                    isHomemb env zExp ae]


isHomemb :: ScpEnv -> ZExpr -> ZExpr -> Bool
isHomemb env t1@(i1, _, _, _) t2@(i2, _, _, _) 
 | i1 <= i2 = new_peel env t1 t2 || any (isHomemb env t1) (subterms i1 t2)
 | otherwise = False

subterms :: Int -> ZExpr -> [ZExpr]
subterms sz (i, e, c, bs) | sz > i = []
                          | otherwise = subterm e ++ subterms_ctxt sz c ++ concatMap rhssOfBind bs
 where subterm (ALam _ e@(i', _, _, _)) = if i' < sz then [] else [e]
       subterm _ = []

       rhssOfBind (ANonRec _ e@(i', _, _, _)) = if i' < sz then [] else [e]


subterms_ctxt :: Int -> AContext -> [ZExpr]
subterms_ctxt _ [] = []
subterms_ctxt sz ((AAppCtxt e@(i, _, _, _)):c) | sz > i = subterms_ctxt sz c
                                               | otherwise = e:subterms_ctxt sz c
subterms_ctxt sz ((APrimOpCtxt oes ies):c) = select oes ++ select ies ++ subterms_ctxt sz c
    where select [] = []
          select (e@(i, _, _, _):t) | sz > i = select t
                                    | otherwise = e:select t
subterms_ctxt sz ((ACaseCtxt _ _ alts):c) = rhssOfAlts alts ++ subterms_ctxt sz c
    where rhssOfAlts [] = []
          rhssOfAlts ((c, bs, e@(i, _, _, _)):t) | sz > i = rhssOfAlts t
                                                 | otherwise = e:rhssOfAlts t
subterms_ctxt sz ((ACastCtxt _):c) = subterms_ctxt sz c


new_peel :: ScpEnv -> ZExpr
                   -> ZExpr -> Bool
new_peel env t1@(i1, AVar v1, c1, b1) (i2, AVar v2, c2, b2)
 | v1 == v2 -- f(e1..en) < f(e1'..en')
 , Just _ <- e = peel_con env c1 c2 && peel_bind env b1 b2
 | v1 == v2 -- k(e1..en) < k(e1'..en')
 , Just _ <- conv1 = peel_con env c1 c2 && peel_bind env b1 b2 
 | PrimOpId f1 <- detv1
 , PrimOpId f2 <- detv2
 , f1 /= f2 = False -- Distinguish + from -
 | isJust e || isJust (maybeInline env v2) = False -- Distinguish f from g.
 | isJust conv1 && not (isJust conv2) || not (isJust conv1) && isJust conv2 = False
 | otherwise = peel_con env c1 c2 && peel_bind env b1 b2 -- x is embedded in y.
       where e = maybeInline env v1
             detv1 = idDetails v1
             detv2 = idDetails v2
             conv1 = isDataConId_maybe v1
             conv2 = isDataConId_maybe v2
new_peel env t1@(i1, ALit v1, c1, b1) (i2, ALit v2, c2, b2) 
  = peel_con env c1 c2 && peel_bind env b1 b2 -- 9 embedded in 2
new_peel env t1@(i1, ALam _ e1, c1, b1) (i2, ALam _ e2, c2, b2)
 | isHomemb env e1 e2 = peel_con env c1 c2 && peel_bind env b1 b2
new_peel env t1@(i1, AType t1', c1, b1) (i2, AType t2', c2, b2) 
 | isHomemb_ty env t1' t2' = peel_con env c1 c2 && peel_bind env b1 b2
new_peel _ t1@(i1, e1, _, _) t2 = False

peel_con :: ScpEnv -> AContext -> AContext -> Bool
peel_con _ [] _ = True
peel_con _ _ [] = False
peel_con env ((AAppCtxt e1):c1) ((AAppCtxt e2):c2) 
 | isHomemb env e1 e2 = peel_con env c1 c2
peel_con env ((APrimOpCtxt (e1:oes1) ies1):c1) ((APrimOpCtxt (e2:oes2) ies2):c2)
 | v1 == v2 && length oes1 == length oes2 && length ies1 == length ies2 &&
     and (zipWith (isHomemb env) oes1 oes2) && and (zipWith (isHomemb env) ies1 ies2) = peel_con env c1 c2
  where (_, AVar v1, _, _) = e1
        (_, AVar v2, _, _) = e2
peel_con env ((ACaseCtxt b1 t1 alts1):c1) ((ACaseCtxt b2 t2 alts2):c2)
 | length alts1 == length alts2 && 
      and (zipWith (\(_, _, e1) (_, _, e2) -> isHomemb env e1 e2) alts1 alts2) = peel_con env c1 c2
peel_con env ((ACastCtxt _):c1) ((ACastCtxt _):c2) = peel_con env c1 c2
peel_con env c1 ((AAppCtxt e):c2) = peel_con env c1 c2
peel_con env c1 ((APrimOpCtxt _ _):c2) = peel_con env c1 c2
peel_con env c1 ((ACaseCtxt _ _ _):c2) = peel_con env c1 c2
peel_con env c1 ((ACastCtxt _):c2) = peel_con env c1 c2


peel_bind :: ScpEnv -> [ABind] -> [ABind] -> Bool
peel_bind _ [] _ = True
peel_bind _ _ [] = False
peel_bind env ((ANonRec _ e1):t1) ((ANonRec _ e2):t2) 
 | isHomemb env e1 e2 = peel_bind env t1 t2
-- peel_bind env ((Rec p1):t1) ((Rec p2):t2) 
--  | length p1 == length p2 && 
--    and (zipWith (\(b1, e1) (b2, e2) -> isHomemb env (splitTerm e1) (splitTerm e2)) p1 p2) = peel_bind env t1 t2
peel_bind env e (h:t) = peel_bind env e t


isHomemb_ty :: ScpEnv -> Type -> Type -> Bool
-- isHomemb_ty env t1 t2 = trace ("HETY: " ++ (showSDoc $ ppr t1) ++ (showSDoc $ ppr t2) ++ (showSDoc $ ppr res)) $ res
#if MATCHTYPES
isHomemb_ty env t1 t2 = res
  where res = new_peel_ty env t1 t2 || any (isHomemb_ty env t1) (subterms_ty t2)
#else
isHomemb_ty _ _ _ = True
#endif

subterms_ty :: Type -> [Type]
subterms_ty (TyVarTy _) = []
subterms_ty (AppTy t1 t2) = [t1, t2]
subterms_ty (TyConApp _ ts) = ts
subterms_ty (FunTy t1 t2) = [t1, t2]
subterms_ty (ForAllTy _ t) = [t]
subterms_ty (PredTy _) = []

new_peel_ty :: ScpEnv -> Type -> Type -> Bool
new_peel_ty _ (TyVarTy {}) (TyVarTy {}) = True
new_peel_ty env (AppTy t1 t2) (AppTy t1' t2') = isHomemb_ty env t1 t1' && isHomemb_ty env t2 t2'
--new_peel_ty env t1@(TyConApp tc1 ts1) t2@(TyConApp tc2 ts2) | tc1 == tc2 = trace ("Tyconapp: " ++ (showSDoc $ ppr t1) ++ (showSDoc $ ppr t2)) $ and (zipWith (isHomemb_ty env) ts1 ts2)
new_peel_ty env t1@(TyConApp tc1 ts1) t2@(TyConApp tc2 ts2) 
  | tc1 == tc2 = and (zipWith (isHomemb_ty env) ts1 ts2)
new_peel_ty env (FunTy t1 t2) (FunTy t1' t2') = isHomemb_ty env t1 t1' && isHomemb_ty env t2 t2'
new_peel_ty env (ForAllTy _ t1) (ForAllTy _ t2) = isHomemb_ty env t1 t2
new_peel_ty env t1@(PredTy _) t2@(PredTy _) = True
new_peel_ty _ _ _ = False

sameTyCon t1 t2 = isFunTyCon t1 && isFunTyCon t2 || 
                  isAlgTyCon t1 && isAlgTyCon t2 ||
                  isTupleTyCon t1 && isTupleTyCon t2 ||
                  isSynTyCon t1 && isSynTyCon t2 ||
                  isPrimTyCon t1 && isPrimTyCon t2 ||
                  isCoercionTyCon t1 && isCoercionTyCon t2 ||
                  isAnyTyCon t1 && isAnyTyCon t2 ||
                  isSuperKindTyCon t1 && isSuperKindTyCon t2

#endif

#endif

{-
sameTyCon (FunTyCon {}) (FunTyCon {}) = True
sameTyCon (AlgTyCon {}) (AlgTyCon {}) = True
sameTyCon (TupleTyCon {}) (TupleTyCon {}) = True
sameTyCon (SynTyCon {}) (SynTyCon {}) = True
sameTyCon (PrimTyCon {}) (PrimTyCon {}) = True
sameTyCon (CoTyCon {}) (CoTyCon {}) = True
sameTyCon (AnyTyCon {}) (AnyTyCon {}) = True
sameTyCon (SuperKindTyCon {}) (SuperKindTyCon {}) = True
sameTyCon _ _ = False
-}

rwCtxt :: Context -> Bool
rwCtxt ((CastCtxt {}):c) = rwCtxt c
rwCtxt ((AppCtxt e):c) 
 | coreEqType (exprType e) realWorldStatePrimTy && notAppCtxt c = True
rwCtxt _ = False

notAppCtxt :: Context -> Bool
notAppCtxt ((AppCtxt {}):c) = False
notAppCtxt _ = True

boringContext :: ScpEnv -> Context -> Bool
#if NOBORING
boringContext _ _ = False
#else
boringContext env c = rwCtxt c || boringContext' env c

boringContext' :: ScpEnv -> Context -> Bool
boringContext' _ [] = True
boringContext' env ((AppCtxt e):c) | boringExpr env e = boringContext' env c
boringContext' env ((CastCtxt {}):c) = boringContext' env c
boringContext' env ((PrimOpCtxt (_:oes) ies):c) 
 | all (boringExpr env) (oes ++ ies) = boringContext' env c
boringContext' _ _ = False

boringExpr env e = -- trace ((if res then "BORING: " else "FUNNY: ") ++ (showSDoc $ ppr e)) 
 res
  where res = boringExpr' env e

boringExpr' :: ScpEnv -> CoreExpr -> Bool
boringExpr' env (Var v) | elemInScopeSet v (inSet env) = True
boringExpr' _ (Type _) = True
boringExpr' env e@(App e1 e2)
 | boringExpr' env e1
 , boringExpr' env e2 = True
--  | (Var fun, es) <- ps
--  , PrimOpId _ <- idDetails fun = case es of
--      [Var x, n] | numLike env n && elemInScopeSet x (inSet env) -> True
--      [n, Var x] | numLike env n && elemInScopeSet x (inSet env) -> True
--      _ -> all (boringExpr' env) es
 | (Var fun, es) <- ps
 ,  elem (occNameString (nameOccName (idName fun))) primopFunctions
  = -- case es of
--     [Var x, n] | numLike env n && elemInScopeSet x (inSet env) -> trace ("Discarding " ++ (showSDoc $ ppr e)) True
--     [n, Var x] | numLike env n && elemInScopeSet x (inSet env) -> trace ("Discarding " ++ (showSDoc $ ppr e)) True
--    -> 
      all (boringExpr' env) es
  where ps = collectArgs e
boringExpr' env (Cast e _) = boringExpr' env e
boringExpr' env (Note _ e) = boringExpr' env e
boringExpr' _ _ = False
#endif

-- XXX libraries/integer-gmp/GHC for the full list

primopFunctions = ["+", "-", "*", "/", "plusInteger", "timesInteger", "eqInteger", "absInteger", "signumInteger", "ltInteger", "gtInteger", "leInteger", "geInteger"]

-- XXX: S# and J# does not work.
intTypes :: [String]
intTypes = ["smallInteger", "S#", "J#"]

numLike :: ScpEnv -> CoreExpr -> Bool
numLike _ l@(Lit {}) = True
numLike env e@(App (Var fun) (Lit l)) | elem (occNameString (nameOccName (idName fun))) intTypes = True
numLike _ e = False


{-

  The most specific generalisation

-}
msg :: ScpEnv -> Id -> (CoreExpr, Context, [CoreBind]) -> (CoreExpr, Context, [CoreBind]) -> ScpM s (Maybe (CoreExpr, Context, [CoreBind], TermParts))
msg env fun p1@(e1, c1, b1) p2 = do
  let fvs = realFvs env (mkLets b1 (plug c1 e1))
  msg' env emptyFM True fun fvs emptyTp p1 p2

msg' :: ScpEnv -> FiniteMap Var Var -> Bool -> Id -> [Var] -> TermParts -> (CoreExpr, Context, [CoreBind]) -> (CoreExpr, Context, [CoreBind]) -> ScpM s (Maybe (CoreExpr, Context, [CoreBind], TermParts))
msg' env vars toplevel fun fvs subst (e1, c1, b1) (e2, c2, b2) = do
  let env' = env {inSet = extendInScopeSetList (inSet env) (bindersOfBinds b1)}
  tmp <- msg_headexp env' vars fun fvs subst (e1, e2)
  case tmp of
    Nothing -> if toplevel
                 then return Nothing
                 else do
                   let plugexp = mkLets b1 (plug c1 e1)
                       fvs2 = realFvs env plugexp
                       fvs1' = delete fvs fvs2
                   n <- newName fun (mkPiTypes fvs1' (exprType plugexp))
                   let (lterm, newterm) = (mkLams fvs1' plugexp, mkVarApps (Var n) fvs1')
                       ret | realExprSize plugexp < realExprSize newterm = Nothing
                           | otherwise = Just (newterm, emptyContext, [], extendTpExp env subst n lterm)
                   return ret
    Just (e1', subst1) -> rest env' e1' (exprType e1) subst1
   where rest env' e1' t1 subst1 = do
                           tmp2 <- msg_binds env vars fun fvs subst1 b1 b2
                           case tmp2 of
                             Nothing -> return Nothing
                             Just (b1', subst2) -> do
                                                    tmp3 <- if toplevel
                                                             then msg_context env' vars fun fvs e1' subst2 c1 c2 t1
                                                             else msg_context_helper env' vars fun fvs e1' subst2 c1 c2 t1
                                                    case tmp3 of
                                                      Nothing -> return Nothing
                                                      Just (e1'', c1', subst3) -> do
                                                                                   return $ Just (e1'', c1', b1', subst3)

msg_headexp :: ScpEnv -> FiniteMap Var Var -> Id -> [Var] -> TermParts -> (CoreExpr, CoreExpr) -> ScpM s (Maybe (CoreExpr, TermParts))
msg_headexp env vars fun fvs1 subst (Var n1, Var n2)
 | n1 == n2 = return $ Just (Var n1, subst)
 | Just v' <- lookupFM vars n1
 , n2 == v' = return $ Just (Var n1, subst)
 | Just v' <- lookupTpExp' subst n1 = return $ Just (Var v', subst)
msg_headexp _ _ _ _ subst (Lit l1, Lit l2) | l1 == l2 = return $ Just (Lit l1, subst)
msg_headexp env vars fun fvs1 subst (l1@(Lam b1 e1), l2@(Lam b2 e2))
  | b1 == b2 = do
    let env' = env {inSet = extendInScopeSet (inSet env) b1}
    tmp <- msg' env' vars False fun fvs1 subst (splitTerm e1) (splitTerm e2)
    case tmp of
      Just (e', c', b', p) -> return $ Just (Lam b1 (mkLets b' (plug c' e')), p)
      Nothing -> return Nothing
msg_headexp env vars fun fvs1 subst (l1@(Type t1), l2@(Type t2)) 
 | t1 `coreEqType`  t2 = do
  return $ Just (Type t1, subst)
 | otherwise = return Nothing 
msg_headexp env vars fun fvs1 subst (e1, e2) = return Nothing

msg_binds :: ScpEnv -> FiniteMap Var Var -> Id -> [Var] -> TermParts -> [CoreBind] -> [CoreBind] -> ScpM s (Maybe ([CoreBind], TermParts))
msg_binds env vars fun fvs subst b1 b2 = go subst b1 b2 []
    where go subst [] [] bs = return $ Just (reverse bs, subst)
          go subst (b1:t1) (b2:t2) bs = do
            tmp <- msg'_bind env vars fun fvs subst b1 b2
            case tmp of
              Just (b1', subst1) -> go subst1 t1 t2 (b1':bs)
              Nothing -> return Nothing
          go _ _ _ _ = return Nothing

msg'_bind :: ScpEnv -> FiniteMap Var Var -> Id -> [Var] -> TermParts -> CoreBind -> CoreBind -> ScpM s (Maybe (CoreBind, TermParts))
msg'_bind env vars fun fvs subst b1 b2 = go subst b1 b2
  where go subst ps1@(Rec p1) ps2@(Rec p2)
         | equalLength p1 p2 = do
          tmp <- msg'_exps env' vars fun fvs subst (rhssOfBind ps1) (rhssOfBind ps2)
          case tmp of
            Just (e1', subst1) -> return $ Just (Rec (zip bs1 e1'), subst1)
            Nothing -> return Nothing
         where bs1 = bindersOf ps1
               env' = env {inSet = extendInScopeSetList (inSet env) bs1 }
        go subst (NonRec b1 e1) (NonRec b2 e2)
         | b1 == b2 = do
          tmp <- msg' env' vars False fun fvs subst (splitTerm e1) (splitTerm e2)
          case tmp of
            Just (e1', c1, bs1, subst1) -> return $ Just (NonRec b1 (mkLets bs1 (plug c1 e1')), subst1)
            Nothing -> return Nothing
         where env' = env {inSet = extendInScopeSet (inSet env) b1 }
        go _ _ _ = return Nothing

msg'_exps :: ScpEnv -> FiniteMap Var Var -> Id -> [Var] -> TermParts -> [CoreExpr] -> [CoreExpr] -> ScpM s (Maybe ([CoreExpr], TermParts))
msg'_exps env vars fun fvs subst es1 es2 = go subst es1 es2 []
    where go subst [] [] es = return $ Just (reverse es, subst)
          go subst (e1:t1) (e2:t2) es = do
           tmp <- msg' env vars False fun fvs subst (splitTerm e1) (splitTerm e2)
           case tmp of
             Just (e1', c1, b1, subst1) -> go subst1 t1 t2 ((mkLets b1 (plug c1 e1')):es)
             Nothing -> return Nothing
          go _ _ _ _ = return Nothing

msg_context :: ScpEnv -> FiniteMap Var Var -> Id -> [Var] -> CoreExpr -> TermParts -> Context -> Context -> Type -> ScpM s (Maybe (CoreExpr, Context, TermParts))
msg_context env vars fun fvs e1 subst c1@((AppCtxt {}):t1) c2@((AppCtxt {}):t2) _t = msg_context_helper env vars fun fvs e1 subst c1 c2 _t
msg_context env vars fun fvs e1 subst c1@((PrimOpCtxt {}):t1) c2@((PrimOpCtxt {}):t2) _t = msg_context_helper env vars fun fvs e1 subst c1 c2 _t
msg_context env vars fun fvs e1 subst c1@((CaseCtxt {}):t1) c2@((CaseCtxt {}):t2) _t = msg_context_helper env vars fun fvs e1 subst c1 c2 _t
msg_context env vars fun fvs e1 subst c1@((CastCtxt {}):t1) c2@((CastCtxt {}):t2) _t = msg_context_helper env vars fun fvs e1 subst c1 c2 _t
msg_context _ _ _ _ _ _ _ _ _ = return Nothing

msg_context_helper :: ScpEnv -> FiniteMap Var Var -> Id -> [Var] -> CoreExpr -> TermParts -> Context -> Context -> Type -> ScpM s (Maybe (CoreExpr, Context, TermParts))
msg_context_helper env vars fun fvs e1 subst c1 c2 _int = go subst _int c1 c2 []
    where go subst _ [] [] cs = return $ Just (e1, reverse cs, subst)
          go subst _int ((AppCtxt l@(Type t)):t1) ((AppCtxt arg2):t2) cs 
           | (Type _t) <- arg2
           , t `coreEqType` _t = do
            let _int' = applyTy _int _t
            go subst _int' t1 t2 ((AppCtxt l):cs)
           | (Type _t) <- arg2 = do
--            trace ("Generalising different types:" ++ (showSDoc $ ppr l <+> ppr arg2)) (return ())
            nam <- newTName (typeKind _t)
            let newtv = mkTyVarTy nam
                _int' = applyTy _int newtv
                subst' = extendTpType env subst nam t
            go subst' _int' t1 t2 ((AppCtxt (Type newtv)):cs)
           | otherwise = do -- We have different number of type arguments - bad.
            return Nothing
          go subst _int ((AppCtxt arg1):t1) ((AppCtxt arg2):t2) cs
           | coreEqType (exprType arg1) (exprType arg2) = do
            tmp <- msg' env vars False fun fvs subst (splitTerm arg1) (splitTerm arg2)
            let _int' | isFunTy _int = funResultTy _int
                      | otherwise = _int
            case tmp of
              Nothing -> return Nothing
              Just (e1', c1', bs', subst1) -> go subst1 _int' t1 t2 ((AppCtxt (mkLets bs' (plug c1' e1'))):cs)
           | isFunTy _int = do
            let (argty, futy) = splitFunTy _int
                fvs2 = realFvs env arg1
                fvs1' = delete fvs fvs2
            nam <- newName fun (mkPiTypes fvs1' argty)
            let (lterm, newterm) = (mkLams fvs1' arg1, mkVarApps (Var nam) fvs1')
                subst1 = extendTpExp env subst nam lterm
--            trace ("Ny typgeneraliserad typ:" ++ (showSDoc $ ppr (exprType arg1) <+> ppr (exprType arg2) <+> ppr arg1 <+> ppr arg2)) (return ())
            go subst1 futy t1 t2 ((AppCtxt newterm):cs)
           | otherwise = do
            return Nothing
          go subst _int ((PrimOpCtxt ((Var v1):oes1) ies1):t1) ((PrimOpCtxt ((Var v2):oes2) ies2):t2) cs 
           | v1 == v2 = do
             tmp1 <- msg'_exps env vars fun fvs subst oes1 oes2
             case tmp1 of
               Nothing -> return Nothing
               Just (oes1', subst1) -> do
                                        tmp2 <- msg'_exps env vars fun fvs subst1 ies1 ies2
                                        case tmp2 of
                                          Nothing -> return Nothing
                                          Just (ies1', subst2) -> go subst2 _int t1 t2 ((PrimOpCtxt ((Var v1):oes1') ies1'):cs)
          go subst _int ((CaseCtxt b1 t1 alts1):tl1) ((CaseCtxt b2 t2 alts2):tl2) cs 
           | t1 `coreEqType` t2 && map fst3 alts1 == map fst3 alts2 = do
            let env' = env {inSet = extendInScopeSet (inSet env) b1}
            tmp <- msg'_alts env' vars fun fvs subst alts1 alts2
            case tmp of
              Nothing -> return Nothing
              Just (alts1', subst1) -> go subst1 _int tl1 tl2 ((CaseCtxt b1 t1 alts1'):cs)
          go subst _int ((CastCtxt co1):t1) ((CastCtxt co2):t2) cs 
           | co1 `coreEqCoercion` co2 = go subst _int t1 t2 ((CastCtxt co1):cs)
          go subst _int c1 c2 cs = do
            let plugexp = plug (reverse cs) e1
                fvs2 = realFvs env plugexp
                fvs1' = delete fvs fvs2
            n <- newName fun (mkPiTypes fvs1' (exprType plugexp))
            let (lterm, newterm) = (mkLams fvs1' plugexp, mkVarApps (Var n) fvs1')                
            let subst1 = extendTpExp env subst n lterm
            return $ Just (newterm, c1, subst1)


msg'_alts :: ScpEnv -> FiniteMap Var Var -> Id -> [Var] -> TermParts -> [CoreAlt] -> [CoreAlt] -> ScpM s (Maybe ([CoreAlt], TermParts))
msg'_alts env vars fun fvs subst alts1 alts2 = go subst alts1 alts2 []
  where go subst [] [] alts = return $ Just (reverse alts, subst)
        go subst ((c1,vs1,r1):alts1) ((c2,vs2,r2):alts2) alts
         | c1 == c2 = do
          let vars' = addListToFM vars (zip vs1 vs2)
          tmp <- msg' env' vars' False fun fvs subst (splitTerm r1) (splitTerm r2)
          case tmp of
            Just (e1', c1', bs1, subst1) -> go subst1 alts1 alts2 ((c1, vs1, (mkLets bs1 (plug c1' e1'))):alts)
            Nothing -> return Nothing
         where env' = env { inSet = extendInScopeSetList (inSet env) vs1 }
        go _ _ _ _ = return Nothing
            

fst3 :: (a, b, c) -> a
fst3 (a, _, _) = a

snd3 :: (a, b, c) -> b
snd3 (_, b, _) = b

trd3 :: (a, b, c) -> c
trd3 (_, _, c) = c

-- Does one of TPush, CPush or Push evaluation steps
evalPush :: ScpEnv -> CoreExpr -> CoreExpr
evalPush env (App (Cast (Lam b e) co) (Type ty)) 
 | (s1s2, _t1t2) <- coercionKind co     -- TPush
 , Just (tyvar, _) <- splitForAllTy_maybe s1s2
 , not (isCoVar tyvar) = 
   let newexp =  mkApps (Lam b (Cast e co')) [Type ty]
   in newexp
 | otherwise = 
   let newexp =  mkApps (Lam b (Cast e' co2)) [Type ty]
   in newexp
         where co' = mkInstCoercion co ty
               e'  = substExpr (newExprSubst env b (Cast (Var b) g1)) e
               [co1, co2] = decomposeCo 1 co
               g1   = mkSymCoercion co1
evalPush env (App (Cast (Lam b e) co) arg) 
 | (s1s2, _t1t2) <- coercionKind co    -- Push
 , not (isTypeArg arg)
 , isFunTy s1s2 = 
   let newexp = mkApps (Lam b (Cast e' co2)) [arg]
   in newexp
         where 
               [co1, co2] = decomposeCo 2 co -- co1 = right (left co)
               g1         = mkSymCoercion co1
               e'         = substExpr (newExprSubst env b (Cast (Var b) g1)) e 
evalPush _ e = error ("evalPush: " ++ (showSDoc $ ppr e))

checkType :: String -> CoreExpr -> CoreExpr  -> Bool
checkType s e1 e2 
  | t1 `coreEqType` t2 = True
  | t1 `tcPartOfType` t2 = trace ("checkType: " ++ s) True
  | otherwise = trace ("CHECKTYPE ASSERTION: " ++ s ++ "\n" ++ (showSDoc $ ppr e1) ++ "::" ++ (showSDoc $ ppr t1) ++ "\n" ++ (showSDoc $ ppr e2)++ "::" ++ (showSDoc $ ppr t2)) True
    where t1 = exprType e1
          t2 = exprType e2
          

newExprSubst :: ScpEnv -> Var -> CoreArg -> Subst
newExprSubst env id expr = extendSubst es id expr
 where es = mkEmptySubst (inSet env)

{-
newTvSubst :: ScpEnv -> TyVar -> Type -> Subst
newTvSubst env id tv = extendTvSubst es id tv
 where es = mkEmptySubst (extendInScopeSetSet (inSet env) (getInScopeVars (funSet env)))
-}

linear :: Id -> CoreExpr -> Bool
linear id e = linear' e <= 1
    where linear' :: CoreExpr -> Integer
          linear' (Var v) | id == v = 1
                          | otherwise = 0
          linear' (Lit {}) = 0
          linear' (App e1 e2) = linear' e1 + linear' e2
          linear' (Lam b e) | id == b = 0
                            | otherwise = 2 * linear' e
--           linear' (Let (NonRec b e) (Lam b' e')) | id == b = 0
--                                                  | id == b' = linear' e
--                                                  | otherwise = linear' e + linear' e'
          linear' (Let (NonRec b e) body) | id == b = 0
                                          | otherwise = linear' e + linear' body
          linear' (Let p body) | elem id (bindersOf p) = 0
                               | otherwise = linear' body + sum (map linear' (rhssOfBind p))
          linear' (Case e b _ alts) | b == id = linear' e
                                    | otherwise = linear' e + maximum (linear_alts' alts)
          linear' (Cast e _) = linear' e
          linear' (Type _) = 0
          linear' (Note _ e) = linear' e
                               
          linear_alts' [] = [0]
          linear_alts' ((_, bs, e):t) | elem id bs = linear_alts' t
                                      | otherwise = linear' e:linear_alts' t

{-

  The TermParts structure and helper functions operating on it.

-}


--                  TypeSubst,   ExpSubst,                          BinderSubst
type TermParts = ([(Var, Type)], [(Var,(InScopeSet, CoreExpr))])

emptyTp :: TermParts
emptyTp = ([], [])

extendTpExp :: ScpEnv -> TermParts -> Var -> CoreExpr -> TermParts
extendTpExp env (t, s) x e = (t, (x, (inSet env, e)):s)

extendTpExps :: ScpEnv -> TermParts -> [Var] -> [CoreExpr] -> TermParts
extendTpExps env (t, s) xs es = (t, zip xs (map ((,) (inSet env)) es) ++ s)

extendTpType :: ScpEnv -> TermParts -> Var -> Type -> TermParts
extendTpType _ (s, e) x t = ((x, t):s, e)

delTp :: TermParts -> [Var] -> TermParts
delTp (t, e) xs = (t, go xs e)
    where go _ [] = []
          go xs (p@(b, _):t) | elem b xs = go xs t
                             | otherwise = p:go xs t

newTpExp :: ScpEnv -> Var -> CoreExpr -> TermParts
newTpExp env x e = extendTpExp env emptyTp x e

newTpExps :: ScpEnv -> [Var] -> [CoreExpr] -> TermParts
newTpExps env xs es = extendTpExps env emptyTp xs es

newTpType :: ScpEnv -> Var -> Type -> TermParts
newTpType env x t = extendTpType env emptyTp x t

getTpBinders :: TermParts -> [Var]
getTpBinders (t, e) = map fst t ++ map fst e

lookupTpExp' :: TermParts -> Var -> Maybe Var
lookupTpExp' (_, env) v = go env
    where go [] = Nothing
          go ((outv, (_, e)):t) 
            | (Var v') <- e
            , v == v' = Just outv
            | otherwise = go t

lookupTpExp :: TermParts -> Var -> Maybe CoreExpr
lookupTpExp (_, env) x 
  | Just (_, e) <- lookup x env = Just e
  | otherwise = Nothing

lookupTpType :: TermParts -> Var -> Maybe Type
lookupTpType (env, _) x = lookup x env

plugTpTerm :: TermParts -> CoreExpr -> CoreExpr
plugTpTerm tp l@(Var x) 
  | Just e <- lookupTpExp tp x = plugTpTerm tp e
  | otherwise = l
plugTpTerm _ l@(Lit {}) = l
plugTpTerm tp (App e1 e2) = App (plugTpTerm tp e1) (plugTpTerm tp e2)
plugTpTerm tp (Lam b e) = Lam b (plugTpTerm (delTp tp [b]) e)
plugTpTerm tp (Let (NonRec b e) e') = Let (NonRec b (plugTpTerm tp' e)) e''
   where e'' = plugTpTerm tp' e'
         tp' = delTp tp [b]
plugTpTerm tp (Let (Rec p) e) = Let (Rec p') (plugTpTerm tp' e)
   where (bs, es) = unzip p
         tp' = delTp tp bs
         p' = zip bs (map (plugTpTerm tp') es)
plugTpTerm tp (Case e b t alts) = Case (plugTpTerm tp e) b t (map helper alts)
   where helper (c, bs, e) = (c, bs, plugTpTerm tp' e)
          where tp' = delTp tp (b:bs)
plugTpTerm tp (Cast e co) = Cast (plugTpTerm tp e) co
plugTpTerm tp (Note n e) = Note n (plugTpTerm tp e)
plugTpTerm tp (Type (TyVarTy x)) | Just e <- lookupTpType tp x = Type e
plugTpTerm _ l@(Type {}) = l

-- Stolen from Rules.lhs

type SubstEnv   = (TvSubstEnv, IdSubstEnv)
type IdSubstEnv = IdEnv CoreExpr		

emptySubstEnv :: SubstEnv
emptySubstEnv = (emptyVarEnv, emptyVarEnv)


match :: MatchEnv
      -> SubstEnv
      -> ZExpr -- Template
      -> ZExpr -- Target
      -> Maybe (SubstEnv, Context, Type)

-- See the notes with Unify.match, which matches types
-- Everything is very similar for terms

-- Interesting examples:
-- Consider matching
--	\x->f 	   against    \f->f
-- When we meet the lambdas we must remember to rename f to f' in the
-- second expresion.  The RnEnv2 does that.
--
-- Consider matching 
--	forall a. \b->b	   against   \a->3
-- We must rename the \a.  Otherwise when we meet the lambdas we 
-- might substitute [a/b] in the template, and then erroneously 
-- 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 (i1, e1, c1, b1) (i2, e2, c2, b2) = do
  subst1 <- match_headexp menv subst e1 e2
  subst2 <- match_binds menv subst1 b1 b2
  match_context menv subst2 (e2, emptyAContext) c1 c2

match_headexp :: MatchEnv -> SubstEnv -> AExpr -> AExpr -> Maybe SubstEnv
match_headexp menv subst (AVar v1) e2
  | Just subst' <- match_var menv subst v1 e2 = Just subst'

-- match_headexp menv subst e1 (AVar v2)     -- Note [Expanding variables]
--   | not (inRnEnvR rn_env v2) -- Note [Do not expand locally-bound variables]
--   , Just e2' <- expandId v2'
--   , Just s@(subst', [], _) <- match (menv { me_env = nukeRnEnvR rn_env }) subst (e1, emptyAContext, []) (splitTerm e2') = Just subst'
--   where
--     v2'    = lookupRnInScope rn_env v2
--     rn_env = me_env menv
-- 	-- Notice that we look up v2 in the in-scope set
-- 	-- See Note [Lookup in-scope]
-- 	-- No need to apply any renaming first (hence no rnOccR)
-- 	-- becuase of the not-locallyBoundR

match_headexp _ subst (ALit lit1) (ALit lit2)
  | lit1 == lit2 = Just subst

match_headexp menv subst (ALam x1 e1) (ALam x2 e2)
  | Just (subst', [], _) <- match menv' subst e1 e2 =
    Just subst'
     where
      menv' = menv { me_env = rnBndr2 (me_env menv) x1 x2 }

match_headexp menv subst (AType ty1) (AType ty2)
  = match_ty menv subst ty1 ty2

-- Everything else fails
match_headexp _ _ _e1 _e2 = -- pprTrace "Failing at" ((text "e1:" <+> ppr _e1) $$ (text "e2:" <+> ppr _e2)) $ 
			 Nothing


match_context :: MatchEnv -> SubstEnv -> (AExpr, AContext) -> AContext -> AContext -> Maybe (SubstEnv, Context, Type)
-- match menv subst (Let b1 body1) (Let b2 body2)
--   | not (any (locallyBoundR rn_env) bs2)
--   = do 	{ subst' <- match_binds menv subst b1 b2
-- 	; match menv' subst' body1 body2 }
--   where
--     rn_env   = me_env menv
--     bs2 = bindersOf b2
--     menv' = menv { me_env = rnBndrs2 rn_env (bindersOf b1) bs2 }

-- match menv subst (Note n1 e1) (Note n2 e2) 
--   | n1 `realNoteEq` n2 = match menv subst e1 e2

match_context _ subst (e, c') [] c = Just (subst, aConToCon c, exprType (plugA (reverse c') e))

match_context menv subst (e, c') ((AAppCtxt e1):c1) (c@(AAppCtxt e2):c2)
 | Just (subst', [], _) <- match menv subst e1 e2
  = match_context menv subst' (e, c:c') c1 c2

match_context menv subst (e, c') ((APrimOpCtxt (e1:oes1) ies1):c1) (c@(APrimOpCtxt (e2:oes2) ies2):c2)
 | v1 == v2
  = do 	{ subst1 <- match_exps menv subst oes1 oes2 -- Is oes always empty?
        ; subst2 <- match_exps menv subst1 ies1 ies2
	; match_context menv subst2 (e, c:c') c1 c2 }
  where (_, AVar v1, _, _) = e1
        (_, AVar v2, _, _) = e2

match_context menv subst (e, c') ((ACaseCtxt x1 ty1 alts1):c1) (c@(ACaseCtxt x2 ty2 alts2):c2)
  = do	{ subst1 <- match_ty menv subst ty1 ty2
	; let menv' = menv { me_env = rnBndr2 (me_env menv) x1 x2 }
	; subst2 <- match_alts menv' subst1 alts1 alts2	-- Alts are both sorted
        ; match_context menv subst2 (e, c:c') c1 c2
	}

match_context menv subst (e, c') ((ACastCtxt co1):c1) (c@(ACastCtxt co2):c2)
  = do	{ subst1 <- match_ty menv subst co1 co2
	; match_context menv subst1 (e, c:c') c1 c2 }


match_context _ _ _ _ _ = Nothing

------------------------------------------
match_var :: MatchEnv
      	  -> SubstEnv
      	  -> Var		-- Template
      	  -> AExpr		-- Target
      	  -> Maybe SubstEnv
match_var menv subst@(tv_subst, id_subst) v1 e2@(AVar v2')
  | v1' `elemVarSet` me_tmpls menv
  = case lookupVarEnv id_subst v1' of
	Nothing	| inRnEnvR rn_env v2' || not (rnInScope v2' rn_env)
		-> Nothing	-- Occurs check failure
		-- e.g. match forall a. (\x-> a x) against (\y. y y)

		| otherwise	-- No renaming to do on e2, because no free var
				-- of e2 is in the rnEnvR of the envt
		-- Note [Matching variable types]
		-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
		-- However, we must match the *types*; e.g.
		--   forall (c::Char->Int) (x::Char). 
		--	f (c x) = "RULE FIRED"
		-- We must only match on args that have the right type
		-- It's actually quite difficult to come up with an example that shows
		-- you need type matching, esp since matching is left-to-right, so type
		-- args get matched first.  But it's possible (e.g. simplrun008) and
		-- this is the Right Thing to do
		-> do	{ tv_subst' <- ruleMatchTyX menv tv_subst (idType v1') (exprType (Var v2'))
						-- c.f. match_ty below
                          -- XXXpj: Make sure we only match renamable vars here.
			; return (tv_subst', extendVarEnv id_subst v1' (Var v2')) }

	Just e1' | eqExprX (\_ -> NoUnfolding) (nukeRnEnvL rn_env) e1' (Var v2') 
		 -> Just subst

		 | otherwise
		 -> Nothing

  | otherwise	-- v1 is not a template variable; check for an exact match with e2
  = case e2 of
       AVar v2 | v1' == rnOccR rn_env v2 -> Just subst
       _    				-> 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_var _ _ _ _ = Nothing

match_alts :: MatchEnv
      -> SubstEnv
      -> [AAlt]		-- Template
      -> [AAlt]		-- Target
      -> Maybe SubstEnv
match_alts _ subst [] []
  = return subst
match_alts menv subst ((c1,vs1,r1):alts1) ((c2,vs2,r2):alts2)
  | c1 == c2
  , Just (subst1, [], _) <- match menv' subst r1 r2 
  = match_alts menv subst1 alts1 alts2 
  where
    menv' :: MatchEnv
    menv' = menv { me_env = rnBndrs2 (me_env menv) vs1 vs2 }

match_alts _ _ _ _
  = Nothing


match_ty :: MatchEnv
      	 -> SubstEnv
      	 -> Type		-- Template
      	 -> Type		-- Target
      	 -> Maybe SubstEnv
match_ty menv (tv_subst, id_subst) ty1 ty2
  = do	{ tv_subst' <- ruleMatchTyX menv tv_subst ty1 ty2
	; return (tv_subst', id_subst) }




match_binds :: MatchEnv -> SubstEnv -> [ABind] -> [ABind] -> Maybe SubstEnv
match_binds _ subst [] [] = return subst
match_binds menv subst ((ANonRec _ e1):t1) ((ANonRec _ e2):t2) 
 | Just (subst1, [], _) <- match menv subst e1 e2 = do
  match_binds menv subst1 t1 t2
-- match_binds menv subst (ps1@(Rec b1):t1) (ps2@(Rec b2):t2)
--  | equalLength b1 b2 = do
--   subst1 <- match_exps menv' subst (rhssOfBind ps1) (rhssOfBind ps2)
--   match_binds menv subst1 t1 t2
--    where menv' = menv { me_env = rnBndrs2 (me_env menv) (bindersOf ps1) (bindersOf ps2)}

match_binds _ _ _ _ = Nothing

match_exps :: MatchEnv -> SubstEnv -> [ZExpr] -> [ZExpr] -> Maybe SubstEnv
match_exps _ subst [] [] = Just subst
match_exps menv subst (e1:t1) (e2:t2) 
  | Just (subst', [], _) <- match menv subst e1 e2 = do
  match_exps menv subst' t1 t2
match_exps _ _ _ _ = Nothing


#if MATCHTYPES
-- ruleMatchTyX and helpers are stolen from Unify.lhs. Only allow renamings
-- instead of arbitrary substitutions in matching.

ruleMatchTyX :: MatchEnv	-- For the most part this is pushed downwards
      -> TvSubstEnv 	-- Substitution so far:
			--   Domain is subset of template tyvars
			--   Free vars of range is subset of 
			--	in-scope set of the RnEnv2
      -> Type -> Type	-- Template and target respectively
      -> Maybe TvSubstEnv
-- This matcher works on core types; that is, it ignores PredTypes
-- Watch out if newtypes become transparent agin!
-- 	this matcher must respect newtypes

ruleMatchTyX menv subst ty1 ty2 | Just ty1' <- coreView ty1 = ruleMatchTyX menv subst ty1' ty2
			 | Just ty2' <- coreView ty2 = ruleMatchTyX menv subst ty1 ty2'

ruleMatchTyX menv subst (TyVarTy tv1) ty2@(TyVarTy {})
  | Just ty1' <- lookupVarEnv subst tv1'	-- tv1' is already bound
  = if tcEqTypeX (nukeRnEnvL rn_env) ty1' ty2
	-- ty1 has no locally-bound variables, hence nukeRnEnvL
	-- Note tcEqType...we are doing source-type matching here
    then Just subst
    else Nothing	-- ty2 doesn't match

  | tv1' `elemVarSet` me_tmpls menv
  = if any (inRnEnvR rn_env) (varSetElems (tyVarsOfType ty2))
    then Nothing	-- Occurs check
    else do { subst1 <- match_kind menv subst tv1 ty2
			-- Note [Matching kinds]
	    ; return (extendVarEnv subst1 tv1' ty2) }

   | otherwise	-- tv1 is not a template tyvar
   = case ty2 of
	TyVarTy tv2 | tv1' == rnOccR rn_env tv2 -> Just subst
	_                                       -> Nothing
  where
    rn_env = me_env menv
    tv1' = rnOccL rn_env tv1

ruleMatchTyX menv subst (ForAllTy tv1 ty1) (ForAllTy tv2 ty2) 
  = ruleMatchTyX menv' subst ty1 ty2
  where		-- Use the magic of rnBndr2 to go under the binders
    menv' = menv { me_env = rnBndr2 (me_env menv) tv1 tv2 }

ruleMatchTyX menv subst (PredTy p1) (PredTy p2) 
  = match_pred menv subst p1 p2
ruleMatchTyX menv subst (TyConApp tc1 tys1) (TyConApp tc2 tys2) 
  | tc1 == tc2 = match_tys menv subst tys1 tys2
ruleMatchTyX menv subst (FunTy ty1a ty1b) (FunTy ty2a ty2b) 
  = do { subst' <- ruleMatchTyX menv subst ty1a ty2a
       ; ruleMatchTyX menv subst' ty1b ty2b }
ruleMatchTyX menv subst (AppTy ty1a ty1b) ty2@(AppTy ty2a ty2b)
--  | Just (ty2a, ty2b) <- repSplitAppTy_maybe ty2
	-- 'repSplit' used because the tcView stuff is done above
  = do { subst' <- ruleMatchTyX menv subst ty1a ty2a
       ; ruleMatchTyX menv subst' ty1b ty2b }

ruleMatchTyX _ _ _ _
  = Nothing

--------------
match_kind :: MatchEnv -> TvSubstEnv -> TyVar -> Type -> Maybe TvSubstEnv
-- Match the kind of the template tyvar with the kind of Type
-- Note [Matching kinds]
match_kind menv subst tv ty
  | isCoVar tv = do { let (ty1,ty2) = coVarKind tv
			  (ty3,ty4) = coercionKind ty
		    ; subst1 <- ruleMatchTyX menv subst ty1 ty3
		    ; ruleMatchTyX menv subst1 ty2 ty4 }
  | otherwise  = if typeKind ty `isSubKind` tyVarKind tv
		 then Just subst
		 else Nothing

-- Note [Matching kinds]
-- ~~~~~~~~~~~~~~~~~~~~~
-- For ordinary type variables, we don't want (m a) to match (n b) 
-- if say (a::*) and (b::*->*).  This is just a yes/no issue. 
--
-- For coercion kinds matters are more complicated.  If we have a 
-- coercion template variable co::a~[b], where a,b are presumably also
-- template type variables, then we must match co's kind against the 
-- kind of the actual argument, so as to give bindings to a,b.  
--
-- In fact I have no example in mind that *requires* this kind-matching
-- to instantiate template type variables, but it seems like the right
-- thing to do.  C.f. Note [Matching variable types] in Rules.lhs

--------------
match_tys :: MatchEnv -> TvSubstEnv -> [Type] -> [Type] -> Maybe TvSubstEnv
match_tys menv subst tys1 tys2 = match_list (ruleMatchTyX menv) subst tys1 tys2

--------------
match_list :: (TvSubstEnv -> a -> a -> Maybe TvSubstEnv)
	   -> TvSubstEnv -> [a] -> [a] -> Maybe TvSubstEnv
match_list _  subst []         []         = Just subst
match_list fn subst (ty1:tys1) (ty2:tys2) = do	{ subst' <- fn subst ty1 ty2
						; match_list fn subst' tys1 tys2 }
match_list _  _     _          _          = Nothing

--------------
match_pred :: MatchEnv -> TvSubstEnv -> PredType -> PredType -> Maybe TvSubstEnv
match_pred menv subst (ClassP c1 tys1) (ClassP c2 tys2)
  | c1 == c2 = match_tys menv subst tys1 tys2
match_pred menv subst (IParam n1 t1) (IParam n2 t2)
  | n1 == n2 = ruleMatchTyX menv subst t1 t2
match_pred _    _     _ _ = Nothing
#else
ruleMatchTyX = Unify.ruleMatchTyX
#endif



-- | Returns @Just (dc, [x1..xn])@ if the argument expression is 
-- a constructor application of the form @dc x1 .. xn@
exprIsConApp_maybe :: CoreExpr -> Maybe (DataCon, [CoreExpr])
exprIsConApp_maybe (Cast expr co)
  =     -- Here we do the KPush reduction rule as described in the FC paper
    case exprIsConApp_maybe expr of {
	Nothing 	   -> Nothing ;
	Just (dc, dc_args) -> 

	-- The transformation applies iff we have
	--	(C e1 ... en) `cast` co
	-- where co :: (T t1 .. tn) ~ (T s1 ..sn)
	-- That is, with a T at the top of both sides
	-- The left-hand one must be a T, because exprIsConApp returned True
	-- but the right-hand one might not be.  (Though it usually will.)

    let (from_ty, to_ty)	   = coercionKind co
	(from_tc, from_tc_arg_tys) = splitTyConApp from_ty
  		-- The inner one must be a TyConApp
    in
    case splitTyConApp_maybe to_ty of {
	Nothing -> Nothing ;
	Just (to_tc, to_tc_arg_tys) 
		| from_tc /= to_tc -> Nothing
		-- These two Nothing cases are possible; we might see 
		--	(C x y) `cast` (g :: T a ~ S [a]),
		-- where S is a type function.  In fact, exprIsConApp
		-- will probably not be called in such circumstances,
		-- but there't nothing wrong with it 

	 	| otherwise  ->
    let
	tc_arity = tyConArity from_tc

        (univ_args, rest1)        = splitAt tc_arity dc_args
        (ex_args, rest2)          = splitAt n_ex_tvs rest1
	(co_args_spec, rest3)     = splitAt n_cos_spec rest2
	(co_args_theta, val_args) = splitAt n_cos_theta rest3

        arg_tys 	    = dataConRepArgTys dc
	dc_univ_tyvars	    = dataConUnivTyVars dc
        dc_ex_tyvars        = dataConExTyVars dc
	dc_eq_spec	    = dataConEqSpec dc
        dc_eq_theta         = dataConEqTheta dc
        dc_tyvars           = dc_univ_tyvars ++ dc_ex_tyvars
        n_ex_tvs            = length dc_ex_tyvars
	n_cos_spec	    = length dc_eq_spec
	n_cos_theta	    = length dc_eq_theta

	-- Make the "theta" from Fig 3 of the paper
        gammas              = decomposeCo tc_arity co
        new_tys             = gammas ++ map (\ (Type t) -> t) ex_args
        theta               = zipOpenTvSubst dc_tyvars new_tys

          -- First we cast the existential coercion arguments
        cast_co_spec (tv, ty) co 
          = cast_co_theta (mkEqPred (mkTyVarTy tv, ty)) co
        cast_co_theta eqPred (Type co) 
          | (ty1, ty2) <- getEqPredTys eqPred
          = Type $ mkSymCoercion (Type.substTy theta ty1)
		   `mkTransCoercion` co
		   `mkTransCoercion` (Type.substTy theta ty2)
        new_co_args = zipWith cast_co_spec  dc_eq_spec  co_args_spec ++
                      zipWith cast_co_theta dc_eq_theta co_args_theta
  
          -- ...and now value arguments
	new_val_args = zipWith cast_arg arg_tys val_args
	cast_arg arg_ty arg = mkCoerce (Type.substTy theta arg_ty) arg

    in
    ASSERT( length univ_args == tc_arity )
    ASSERT( from_tc == dataConTyCon dc )
    ASSERT( and (zipWith coreEqType [t | Type t <- univ_args] from_tc_arg_tys) )
    ASSERT( all isTypeArg (univ_args ++ ex_args) )
    ASSERT2( equalLength val_args arg_tys, ppr dc $$ ppr dc_tyvars $$ ppr dc_ex_tyvars $$ ppr arg_tys $$ ppr dc_args $$ ppr univ_args $$ ppr ex_args $$ ppr val_args $$ ppr arg_tys  )

    Just (dc, map Type to_tc_arg_tys ++ ex_args ++ new_co_args ++ new_val_args)
    }}

exprIsConApp_maybe expr = analyse (collectArgs expr)
  where
    analyse (Var fun, args)
	| Just con <- isDataConWorkId_maybe fun,
	  args `lengthAtLeast` dataConRepArity con
		-- Might be > because the arity excludes type args
	= Just (con,args)

    analyse _ = Nothing

delete :: Eq a => [a] -> [a] -> [a]
delete _ [] = []
delete es (h:t) | h `elem` es = delete es t
                | otherwise = h:delete es t

sortGT (a1,b1,c1) (a2,b2,c2) = case guff a1 a2 of
                              EQ -> compare b1 b2
                              LT -> GT
                              GT -> LT

guff :: CoreExpr -> CoreExpr -> Ordering
guff (Var h1) (Var h2) = compare h1 h2
guff _ _ = EQ



printPrep :: Store -> [(CoreExpr, Int, CoreExpr)]
printPrep [] = []
printPrep (r:t) = sortBy sortGT (newe r:printPrep t)
  where newe ((RhoE {..}), _) = (headExp, length restExp, plug restExp headExp)

createBinds :: FiniteMap Var Int -> Store -> [(Var, CoreExpr)]
createBinds _ [] = []
createBinds taint (r@(RhoE {..}, _):t) -- | elemFM freshName taint = createBinds taint t
                                       | otherwise = newe r:createBinds taint t
    where newe (RhoE {freshName = h, ..}, e) = (h, mkLams ((if stateHack then [voidArgId] else []) ++ inFvs) e)

beautify :: Store -> [SDoc]
beautify [] = []
beautify (((RhoE {freshName = h, inFvs = xs, restExp = c, headExp = f}), e):t) = (text "(" <+> ppr h <+> ppr xs <+> ppr (plug c f) <+> text "body:" <+> ppr e <+> text ")"):beautify t


weakUnsoundEqExpr :: CoreExpr -> CoreExpr -> Bool
weakUnsoundEqExpr (Var v1)   (Var v2)   = v1==v2
weakUnsoundEqExpr (Lit lit1) (Lit lit2) = lit1 == lit2
weakUnsoundEqExpr (Note n1 e1) (Note n2 e2) = 
  e1 `weakUnsoundEqExpr` e2 && n1 `realNoteEq` n2
weakUnsoundEqExpr (Type t1)  (Type t2)  = t1 `coreEqType` t2
weakUnsoundEqExpr (Lam _ _) (Lam _ _) = True
weakUnsoundEqExpr _ _ = False

instance Outputable Savings where
  ppr SZero = text "Zero"
  ppr (SSmall s) = text "Small->" <+> ppr s
  ppr (SBig s) = text "Big->" <+> ppr s
  ppr (SSum s) = text "Sum(" <+> vcat (map ppr s) <+> text ")"
  ppr SFold = text "FOLD"
  ppr (SChoice s) = text "Choice(" <+> vcat (map ppr s) <+> text ")"
  ppr (SSplit s) = text "Split(" <+> vcat (map ppr s) <+> text ")"
  ppr (SEnables b) | b = text "Enables fold"
                   | otherwise = empty


instance Outputable CFrame where
  ppr (AppCtxt {}) = text "AppCtxt"
  ppr (PrimOpCtxt {}) = text "PrimOpCtxt"
  ppr (CaseCtxt {}) = text "CaseCtxt"
  ppr (CastCtxt {}) = text "CastCtxt"

instance Outputable AExpr where
  ppr (AVar id) = ppr id
  ppr (ALit l) = ppr l
  ppr (ALam v ze) = text "/" <+> ppr v <+> text "." <+> ppr ze
  ppr (AType t) = ppr t

instance Outputable ABind where
  ppr (ANonRec v e) = ppr v <+> text "=" <+> ppr e

instance Outputable AFrame where
  ppr (AAppCtxt e) = ppr e
  ppr (APrimOpCtxt oes ies) =  ppr oes <+> ppr ies
  ppr (ACaseCtxt v t alts) = ppr v <+> ppr t <+> ppr alts
  ppr (ACastCtxt c) = ppr c

instance Outputable Integer where
  ppr l = ppr ((fromIntegral l) :: Int)

deShadowExpr :: ScpEnv -> CoreExpr -> CoreExpr
deShadowExpr env e = substExpr s e
    where s =  mkEmptySubst (inSet env)

washBinds :: [CoreBind] -> [CoreBind]
washBinds [] = []
washBinds ((NonRec b e):t) = (NonRec b' (washExpression env e)):washBinds t
    where b' = zapIdInfo b
          env = extendInScopeSet emptyInScopeSet b'
washBinds ((Rec ps):t) = (Rec (zip bs' es')):washBinds t
    where (bs, es) = unzip ps
          bs' = map zapIdInfo bs
          env = extendInScopeSetList emptyInScopeSet bs'
          es' = map (washExpression env) es

washExpression :: InScopeSet -> CoreExpr -> CoreExpr
washExpression env l@(Var v) 
 | Just v' <- lookupInScope env v = Var v'
 | otherwise = l
washExpression _ l@(Lit {}) = l
washExpression env (App e1 e2) = App (washExpression env e1) (washExpression env e2)
washExpression env (Lam b e) = Lam b' (washExpression env' e)
    where b' = zapIdInfo b
          env' = extendInScopeSet env b'
washExpression env (Let (NonRec b e) body) = Let (NonRec b' e') body'
    where b' = zapIdInfo b
          env' = extendInScopeSet env b'
          e' = washExpression env e
          body' = washExpression env' body
washExpression env (Let (Rec p) body) = Let (Rec (zip bs' es')) body'
    where (bs, es) = unzip p
          bs' = map zapIdInfo bs
          env' = extendInScopeSetList env bs'
          es' = map (washExpression env') es
          body' = washExpression env' body
washExpression env (Case e b t alts) = Case (washExpression env e) b' t alts'
    where b' = zapIdInfo b
          alts' = map (washAlt env') alts
          env' = extendInScopeSet env b'
washExpression env (Cast e c) = Cast (washExpression env e) c
washExpression env (Note n e) = Note n (washExpression env e)
washExpression _ t@(Type {}) = t

washAlt :: InScopeSet -> CoreAlt -> CoreAlt
washAlt env (c, bs, e) = (c, bs', washExpression env' e)
    where bs' = map zapIdInfo bs
          env' = extendInScopeSetList env bs'


expandId :: Id -> Maybe CoreExpr
expandId id
  | isExpandableUnfolding unfolding = Just (unfoldingTemplate unfolding)
  | otherwise                       = Nothing
  where
    unfolding = idUnfolding id

findCtxtBinders :: ScpEnv -> (CoreExpr, Context, [CoreBind]) -> (ScpEnv, (CoreExpr, Context, [CoreBind]) )
findCtxtBinders env e@(_, c, _) = (env { inSet = extendInScopeSetList (inSet env) newbs}, e)
  where newbs = go c []
        go [] bs = reverse bs
        go ((CaseCtxt b t alts):cs) bs = go cs (b:bs)
        go (_:cs) bs = go cs bs

lambdaLift :: ScpEnv -> (Bool, CoreExpr) -> Id -> ScpM Store (CoreExpr, ScpEnv)
lambdaLift env (True, e) _ = return (e, env)
lambdaLift env (False, e) fun = do
  scpLog 6 "Lifting" (ppr fun <+> ppr e)
  (e', bs) <- lambdaLift' env0 e
  scpLog 6 "Lifting, 2" (ppr bs)
  let binds' = addListToFM (binds env) ((fun, e'):bs)
      env' = env {binds = binds'}
  putBinds binds'
  mapM (\(b, e) -> let newquad = RhoE { freshName = b, inFvs = [], stateHack = False, zExp = (1, (AVar b), [], []), restExp = [], headExp = (Var b), compSize = 1} in addToStore (newquad, e)) bs
  return (e', env')

lambdaLift' :: ScpEnv -> CoreExpr -> ScpM s (CoreExpr, [(Var, CoreExpr)])
lambdaLift' _ v@(Var {}) = return (v, [])
lambdaLift' _ l@(Lit {}) = return (l, [])
lambdaLift' env (App e1 e2) = do
  (e1', b1) <- lambdaLift' env e1
  (e2', b2) <- lambdaLift' env e2
  return (App e1' e2', b1 ++ b2)
lambdaLift' env (Lam b e) = do
  (e', b1) <- lambdaLift' (env {inSet = extendInScopeSet (inSet env) b}) e
  return (Lam b e', b1)
lambdaLift' env (Let (NonRec b e) body) = do
  let env' = env {inSet = extendInScopeSet (inSet env) b}
  (e', b1) <- lambdaLift' env' e
  (body', b2) <- lambdaLift' env' body
  return (Let (NonRec b e') body', b1 ++ b2)
lambdaLift' env (Let (Rec b) body) = do
  let (bs, es) = unzip b
  (s_tmp, p) <- closeBody env bs es
  let (bs', es') = unzip p
      s' = extendInScopeList s_tmp bs'
  (body', b1) <- lambdaLift' env (substExpr s' body)
  (es'', b2) <- mapAndUnzipM (lambdaLift' env) es'
  scpLog 6 "Lifting, Fore" (ppr b)
  scpLog 6 "Lifting, Efter" (ppr (zip bs' es''))
  return (body', b1 ++ zip bs' es'' ++ concat b2)
lambdaLift' env (Case e b t alts) = do
  let env' = env {inSet = extendInScopeSet (inSet env) b}
  (e', b1) <- lambdaLift' env' e
  (alts', bs2) <- mapAndUnzipM (lambdaLiftAlt env') alts
  return (Case e' b t alts', b1 ++ concat bs2)
lambdaLift' env (Cast e co) = do
  (e', b1) <- lambdaLift' env e
  return (Cast e' co, b1)
lambdaLift' env (Note n e) = do
  (e', b1) <- lambdaLift' env e
  return (Note n e', b1)
lambdaLift' _ t@(Type {}) = return (t, [])

lambdaLiftAlt :: ScpEnv -> CoreAlt -> ScpM s (CoreAlt, [(Var, CoreExpr)])
lambdaLiftAlt env (c, bs, e) = do
  let env' = env {inSet = extendInScopeSetList (inSet env) bs}
  (e', b1) <- lambdaLift' env' e
  return ((c, bs, e'), b1) 

closeBody :: ScpEnv -> [Var] -> [CoreExpr] -> ScpM s (Subst, [(Var, CoreExpr)])
closeBody env bs es = go emptySubst bs es []
    where fvs = realGroupFvs env es
          prep _ [] res = res
          prep s ((b,e):es) res = prep s es ((b,(mkLams fvs (substExpr s e))):res)
          go s [] [] res = return (s, prep s res [])
          go s (b:t1) (e:t2) res = do
           let full_ty = mkPiTypes fvs (exprType e)
           b' <- newName b full_ty
           let s' = extendSubst s b (mkVarApps (Var b') fvs)
               ret = (b', e)
           go s' t1 t2 (ret:res)

                    
slask [] = True
slask ((CaseCtxt {}):_) = False
slask ((CastCtxt _):c) = slask c
slask ((AppCtxt {}):c) = slask c
slask ((PrimOpCtxt {}):c) = slask c

-- slask2 :: FiniteMap (Var, Int) [[CoreExpr]] -> [(Var, Int, Int, [CoreExpr])]
slask2 m = sortBy hupp (concatMap go (fmToList m))
-- slask2 m = filter (\(_, _, i, _) -> i > 3) (concatMap go (fmToList m))
  where go ((f, l), elt) = map (go' f l) elt
--        go' f l elt = (f, l, length elt, elt)
        go' f l elt = (f, l, length elt)

--        hupp (_, _, l1, _) (_, _, l2, _) = compare l1 l2
        hupp (_, _, l1) (_, _, l2) = compare l1 l2

-- slask2 :: FiniteMap (Var, Int) [[CoreExpr]] -> [(Var, Int, Int, [CoreExpr])]
slask3 m = sortBy sortQ (concatMap go (fmToList m))
-- slask3 m = filter (\(_, _, i, _) -> i > 3) (concatMap go (fmToList m))
  where go ((f, l), elt) = go' f l elt
        go' _ _ [] = []
        go' f l (elt:t) | length elt > 2 = (f, l, length elt, elt):go' f l t
                        | otherwise = go' f l t
--        go' f l elt = (f, l, length elt)

        hupp (_, _, l1, _) (_, _, l2, _) = compare l1 l2
--        hupp (_, _, l1) (_, _, l2) = compare l1 l2

sortQ (a1,b1,c1,_) (a2,b2,c2,_) = case compare a1 a2 of
                              EQ -> case compare b1 b2 of
                                     EQ -> compare c1 c2
                                     LT -> GT
                                     GT -> LT
                              LT -> GT
                              GT -> LT


normalize :: ScpEnv -> CoreExpr -> Context -> CoreExpr
#if NO_NORMALIZE
normalize _ e c = plug c e
#else
normalize env  l@(Var v) context
 | Just _ <- isDataConId_maybe v,
   Just (e, (CaseCtxt b _ alts):c') <- splitCaseCtxt l context,
   Just (dc, e') <- exprIsConApp_maybe e = 
  let (dc', bs, rhs) = findAlt (DataAlt dc) alts
                     `orElse` (DEFAULT, [], mkImpossibleExpr (coreAltsType alts))
      n_drop_tys = case dc' of
                     DataAlt dc'' -> length (dataConUnivTyVars dc'')
                     _ -> 0
      newargs = drop n_drop_tys e'

      newe = case bs of
               [] -> mkLets [NonRec b e] rhs
               _ -> let (newbs, rhs') = case pjBndrs env bs of
                                          (_, Just (s, bs')) -> (bs', substExpr s rhs)
                                          _ -> (bs, rhs)
                    in mkLets [NonRec b e] (mkLets [NonRec b r | (b, r) <- (zip newbs newargs)] rhs')
  in normalize env newe c'
 | PrimOpId _ <- idDetails v = 
  let (e, context') = makePrimOpCtxt v context
  in normalize env e context'
normalize env (Lit l) ((CaseCtxt b t alts):c) = normalize env newrhs c
  where newrhs = substExpr (newExprSubst env b (Lit l)) rhs
        (_, _, rhs) = findAlt (LitAlt l) alts
                    `orElse` (DEFAULT, [], mkImpossibleExpr (coreAltsType alts))
-- normalize env (Lit l) context@((PrimOpCtxt ((Var fun):oes) []):c) 
--  | all isValue oes =
--   let rule_base = localRuleBase
--       rules = getRules rule_base fun
--       iss = inSet env
--   in case lookupRule (const True) (const NoUnfolding) iss fun (oes ++ [Lit l]) rules of
--      Just (_, rule_rhs) -> normalize env rule_rhs c
--      Nothing -> rebuild env (Lit l) context
normalize env l@(Lam b e) c
 | Just (s, b') <- s' =
  let (s', bs') = substBndr s b
      newexp = Lam bs' (substExpr s' e)
  in normalize env newexp c
 | ((AppCtxt _):_) <- c =
  let Just (args, newc) = collectCtxtArgs c
      newexp = doBeta env l args
  in normalize env newexp newc
 | ((CastCtxt co):c') <- c,
   Just (args, newc) <- collectCtxtArgs c' = 
  let newexp = doBeta env (Cast l co) args
  in normalize env newexp newc
 | otherwise = rebuild env (Lam b (normalize env' e emptyContext)) c
  where (env', s') = pjBndr env b 
normalize env (Let (NonRec b e) body) c 
 | linear b body' || exprIsTrivial e' = 
  let newexp = substExpr (newExprSubst env b e') body'
  in normalize env newexp c
 | otherwise = rebuild env (Let (NonRec b e') body') c
  where env' = env { inSet = extendInScopeSet (inSet env) b}
        e' = normalize env e emptyContext
	body' = normalize env' body emptyContext
normalize _ (Let (Rec p) e') _ = error "Letrec"
normalize env (Case e b t alts) c = normalize env' e ((CaseCtxt b t alts):c) 
  where env' = env { inSet = extendInScopeSet (inSet env) b}
normalize env e@(App e1 e2) c = normalize env e1 ((AppCtxt e2):c)
normalize env c@(Cast e co1) ((CastCtxt co2):c') =
  let co' = mkTransCoercion co1 co2
  in normalize env e ((CastCtxt co'):c') 
normalize env (Cast e co) c
 | isIdentityCoercion co = normalize env e c
 | otherwise = normalize env e ((CastCtxt co):c)
normalize env (Note _ e) c = normalize env e c
normalize env e c = rebuild env e c

rebuild :: ScpEnv -> CoreExpr -> Context -> CoreExpr
rebuild env e ((PrimOpCtxt comp@(o:oes) ies):c)
  | null ies = rebuild env (mkApps o (oes ++ [e])) c
  | all isValue (e:oes) = normalize env (head ies) ((PrimOpCtxt (comp ++ [e]) (tail ies)):c)
  | otherwise = rebuild env (mkApps o (oes ++ e:ies')) c
   where ies' = map (\e -> normalize env e emptyContext) ies
rebuild env e ((AppCtxt arg):c) = rebuild env (mkApps e [arg']) c
  where arg' = normalize env arg emptyContext
rebuild env e ((CastCtxt co):c) = rebuild env (Cast e co) c
rebuild env e ((CaseCtxt b t alts):c) = rebuild env (Case e b t alts') c
  where alts' = map (normAlt env) alts
rebuild _ e [] = e

normAlt :: ScpEnv -> CoreAlt -> CoreAlt
normAlt env (c, bs, e) = (c, bs, normalize env' e emptyContext)
  where env' = env { inSet = extendInScopeSetList (inSet env) bs}

#endif


calculateSavings :: Savings -> Integer -> Integer -> Integer
calculateSavings s n1 n2 | n2 > 40 = 0
                         | n2 > 20 = n - 100
                         | n1 >= n2 = n + (savingsThreshold - 3)
                         | otherwise = n - 50
  where n = calculateSavings' s 

calculateSavings' :: Savings -> Integer
calculateSavings' SZero = 0
calculateSavings' (SSmall s) = 1 + calculateSavings' s
calculateSavings' (SBig s) = 5 + calculateSavings' s
calculateSavings' (SSum s) = sum (map calculateSavings' s)
calculateSavings' SFold = 1000
calculateSavings' (SChoice s) =  maximum (map calculateSavings' s)
calculateSavings' (SSplit s) = sum (map calculateSavings' s)
calculateSavings' (SEnables b) = if b then 500 else 0

findFold :: Savings -> Bool
findFold SZero = False
findFold (SSmall s) = findFold s
findFold (SBig s) = findFold s
findFold (SSum s) = or (map findFold s)
findFold SFold = True
findFold (SChoice s) = or (map findFold s)
findFold (SSplit s) = or (map findFold s)
findFold (SEnables b) = b

createCandidates :: FiniteMap Var Int -> [Var] -> [Var] -> [Var] -> CoreExpr -> Store -> Store -> [Var]
createCandidates taint modulebinders ns done e oldstore ml = nub (ns ++ e_newfuns ++ store_newfuns)
  where e_newfuns = findLocalFuns modulebinders done e
        store_newfuns = scanStore modulebinders done diffstore
        diffstore = tempstore
-- filter (\(RhoE {..}, _) -> not (elemFM freshName taint)) tempstore
        tempstore = take (length ml - length oldstore) ml


scanStore  :: [Var] -> [Var] -> Store -> [Var]
scanStore _ _ [] = []
scanStore modulebinders done ((_, e):t) = findLocalFuns modulebinders done e ++ scanStore modulebinders done t

findLocalFuns :: [Var] -> [Var] -> CoreExpr -> [Var]
findLocalFuns ns done e = findLocalFuns' ns done e

findLocalFuns' :: [Var] -> [Var] -> CoreExpr -> [Var]
findLocalFuns' _ _ (Lit {}) = []
findLocalFuns' ns done (Var n) | elem n ns && not (elem n done) = [n]
                               | otherwise = []
findLocalFuns' ns done (App e1 e2) = findLocalFuns' ns done e1 ++ findLocalFuns' ns done e2
findLocalFuns' ns done (Lam _ e) = findLocalFuns' ns done e
findLocalFuns' ns done (Case e _ _ alts) = findLocalFuns' ns done e ++ concatMap (findLocalFunsAlt' ns done) alts
findLocalFuns' ns done (Cast e _) = findLocalFuns' ns done e
findLocalFuns' ns done (Let bs body) = findLocalFuns' ns done body ++ concatMap (findLocalFuns' ns done) (rhssOfBind bs)
findLocalFuns' ns done (Note _ e) = findLocalFuns' ns done e
findLocalFuns' _ _ (Type {}) = []

findLocalFunsAlt' :: [Var] -> [Var] -> CoreAlt -> [Var]
findLocalFunsAlt' ns done (_, _, e) = findLocalFuns' ns done e

makeAllFunctions :: FiniteMap Var Int -> [Var] -> Store -> [CoreBind] -> [CoreBind] -> [CoreBind]
#if 0
makeAllFunctions _ ml [] mn = [Rec ((createBinds ml) ++ (flattenBinds (reverse mn)))]
makeAllFunctions done ml (m@(NonRec b e):t) mn 
 | elem b done = makeAllFunctions done ml t (m:mn)
makeAllFunctions done ml (m@(Rec bs):t) mn
 | any ((flip elem) done) (bindersOf m) = makeAllFunctions done ml t (m:mn)
makeAllFunctions done ml (h:t) mn = h:makeAllFunctions done ml t mn
#else
makeAllFunctions taint _ ml bs [] = [Rec (flattenBinds bs ++ createBinds taint ml)]
#endif

pruneStore :: (RhoElement, CoreExpr) -> Store -> [Var] -> Store
pruneStore (RhoE {freshName = stopelem, ..}, _) ml infuns = trace ("pruneStore" ++ (showSDoc $ ppr infuns <+> ppr work)) $ prune ml work
  where work = todolist infuns []
        allfuns = newfuns ml

	findfuns [] _ _ = []
	findfuns _ _ [] = []
	findfuns ns done ((RhoE {..},e):t) | freshName == stopelem = []
		    	 		   | otherwise = if null (findLocalFuns ns done e) then findfuns ns done t else freshName:findfuns ns done t

        newfuns [] = []
        newfuns ((RhoE {..},_):t) | freshName == stopelem = []
		       		  | otherwise = freshName:newfuns t

        prune _ [] = error "Emtpy prune list"
        prune [] _ = []
  	prune l@(e@(RhoE {..},_):t) funs | freshName == stopelem = l
  	   	    	                 | elem freshName funs = trace "PRUNED" $ prune t funs
			                 | otherwise = e:prune t funs

        todolist infuns done | null rest = done
		 	     | otherwise = trace ("todolist: " ++ (showSDoc $ ppr rest)) $ todolist rest (infuns ++ done)
            where rest = nub (findfuns infuns done ml)