module LiberateCase ( liberateCase ) where

#include "HsVersions.h"

import DynFlags
import CoreSyn
import CoreUnfold	( couldBeSmallEnoughToInline )
import Id
import VarEnv
import Util             ( notNull )

liberateCase :: DynFlags -> [CoreBind] -> [CoreBind]
liberateCase dflags binds = do_prog (initEnv dflags) binds
  where
    do_prog _   [] = []
    do_prog env (bind:binds) = bind' : do_prog env' binds
			     where
			       (env', bind') = libCaseBind env bind

libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)

libCaseBind env (NonRec binder rhs)
  = (addBinders env [binder], NonRec binder (libCase env rhs))

libCaseBind env (Rec pairs)
  = (env_body, Rec pairs')
  where
    binders = map fst pairs

    env_body = addBinders env binders

    pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]

	-- We extend the rec-env by binding each Id to its rhs, first
	-- processing the rhs with an *un-extended* environment, so
	-- that the same process doesn't occur for ever!
    env_rhs = addRecBinds env [ (localiseId binder, libCase env_body rhs)
			      | (binder, rhs) <- pairs
 			      , rhs_small_enough binder rhs ]
	-- localiseID : see Note [Need to localiseId in libCaseBind]
		 

    rhs_small_enough id rhs	-- Note [Small enough]
	=  idArity id > 0	-- Note [Only functions!]
	&& maybe True (\size -> couldBeSmallEnoughToInline size rhs)
                      (bombOutSize env)

libCase :: LibCaseEnv
	-> CoreExpr
	-> CoreExpr

libCase env (Var v)             = libCaseId env v
libCase _   (Lit lit)           = Lit lit
libCase _   (Type ty)           = Type ty
libCase _   (Coercion co)       = Coercion co
libCase env (App fun arg)       = App (libCase env fun) (libCase env arg)
libCase env (Note note body)    = Note note (libCase env body)
libCase env (Cast e co)         = Cast (libCase env e) co

libCase env (Lam binder body)
  = Lam binder (libCase (addBinders env [binder]) body)

libCase env (Let bind body)
  = Let bind' (libCase env_body body)
  where
    (env_body, bind') = libCaseBind env bind

libCase env (Case scrut bndr ty alts)
  = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
  where
    env_alts = addBinders (mk_alt_env scrut) [bndr]
    mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
    mk_alt_env (Cast scrut _)  = mk_alt_env scrut	-- Note [Scrutinee with cast]
    mk_alt_env _    	       = env

libCaseAlt :: LibCaseEnv -> (AltCon, [CoreBndr], CoreExpr)
                         -> (AltCon, [CoreBndr], CoreExpr)
libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)

libCaseId :: LibCaseEnv -> Id -> CoreExpr
libCaseId env v
  | Just the_bind <- lookupRecId env v	-- It's a use of a recursive thing
  , notNull free_scruts 		-- with free vars scrutinised in RHS
  = Let the_bind (Var v)

  | otherwise
  = Var v

  where
    rec_id_level = lookupLevel env v
    free_scruts  = freeScruts env rec_id_level

freeScruts :: LibCaseEnv
	   -> LibCaseLevel 	-- Level of the recursive Id
	   -> [Id]		-- Ids that are scrutinised between the binding
				-- of the recursive Id and here
freeScruts env rec_bind_lvl
  = [v | (v, scrut_bind_lvl, scrut_at_lvl) <- lc_scruts env
       , scrut_bind_lvl <= rec_bind_lvl
       , scrut_at_lvl > rec_bind_lvl]
	-- Note [When to specialise]
	-- Note [Avoiding fruitless liberate-case]

addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
  = env { lc_lvl_env = lvl_env' }
  where
    lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)

addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env, 
			     lc_rec_env = rec_env}) pairs
  = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
  where
    lvl'     = lvl + 1
    lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
    rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]

addScrutedVar :: LibCaseEnv
	      -> Id		-- This Id is being scrutinised by a case expression
	      -> LibCaseEnv

addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env, 
				lc_scruts = scruts }) scrut_var
  | bind_lvl < lvl
  = env { lc_scruts = scruts' }
	-- Add to scruts iff the scrut_var is being scrutinised at
	-- a deeper level than its defn

  | otherwise = env
  where
    scruts'  = (scrut_var, bind_lvl, lvl) : scruts
    bind_lvl = case lookupVarEnv lvl_env scrut_var of
		 Just lvl -> lvl
		 Nothing  -> topLevel

lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
lookupRecId env id = lookupVarEnv (lc_rec_env env) id

lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
lookupLevel env id
  = case lookupVarEnv (lc_lvl_env env) id of
      Just lvl -> lvl
      Nothing  -> topLevel

type LibCaseLevel = Int

topLevel :: LibCaseLevel
topLevel = 0

data LibCaseEnv
  = LibCaseEnv {
	lc_size :: Maybe Int,	-- Bomb-out size for deciding if
				-- potential liberatees are too big.
				-- (passed in from cmd-line args)

	lc_lvl :: LibCaseLevel,	-- Current level
		-- The level is incremented when (and only when) going
		-- inside the RHS of a (sufficiently small) recursive
		-- function.

	lc_lvl_env :: IdEnv LibCaseLevel,  
		-- Binds all non-top-level in-scope Ids (top-level and
		-- imported things have a level of zero)

	lc_rec_env :: IdEnv CoreBind, 
		-- Binds *only* recursively defined ids, to their own
		-- binding group, and *only* in their own RHSs

	lc_scruts :: [(Id, LibCaseLevel, LibCaseLevel)]
     		-- Each of these Ids was scrutinised by an enclosing
		-- case expression, at a level deeper than its binding
		-- level.
		-- 
 		-- The first LibCaseLevel is the *binding level* of
 		--   the scrutinised Id, 
		-- The second is the level *at which it was scrutinised*.
		--   (see Note [Avoiding fruitless liberate-case])
		-- The former is a bit redundant, since you could always
		-- look it up in lc_lvl_env, but it's just cached here
		-- 
		-- The order is insignificant; it's a bag really
		-- 
		-- There's one element per scrutinisation;
		--    in principle the same Id may appear multiple times,
		--    although that'd be unusual:
		--       case x of { (a,b) -> ....(case x of ...) .. }
	}

initEnv :: DynFlags -> LibCaseEnv
initEnv dflags 
  = LibCaseEnv { lc_size = liberateCaseThreshold dflags,
		 lc_lvl = 0,
		 lc_lvl_env = emptyVarEnv, 
		 lc_rec_env = emptyVarEnv,
		 lc_scruts = [] }

bombOutSize :: LibCaseEnv -> Maybe Int
bombOutSize = lc_size