module DsBinds ( dsTopLHsBinds, dsLHsBinds, decomposeRuleLhs, dsSpec,
                 dsHsWrapper, dsTcEvBinds, dsEvBinds, wrapDsEvBinds,
		 DsEvBind(..), AutoScc(..)
  ) where

#include "HsVersions.h"

import {-# SOURCE #-}	DsExpr( dsLExpr )
import {-# SOURCE #-}	Match( matchWrapper )

import DsMonad
import DsGRHSs
import DsUtils

import HsSyn		-- lots of things
import CoreSyn		-- lots of things
import CoreSubst
import MkCore
import CoreUtils
import CoreArity ( etaExpand )
import CoreUnfold
import CoreFVs
import Digraph

import TcType
import Type
import TysPrim  ( anyTypeOfKind )
import CostCentre
import Module
import Id
import TyCon	( tyConDataCons )
import Class
import DataCon	( dataConRepType )
import Name	( localiseName )
import MkId	( seqId )
import Var
import VarSet
import Rules
import VarEnv
import Outputable
import SrcLoc
import Maybes
import OrdList
import Bag
import BasicTypes hiding ( TopLevel )
import FastString
-- import StaticFlags	( opt_DsMultiTyVar )
import Util

import MonadUtils

dsTopLHsBinds :: AutoScc -> LHsBinds Id -> DsM (OrdList (Id,CoreExpr))
dsTopLHsBinds auto_scc binds = ds_lhs_binds auto_scc binds

dsLHsBinds :: LHsBinds Id -> DsM [(Id,CoreExpr)]
dsLHsBinds binds = do { binds' <- ds_lhs_binds NoSccs binds
                      ; return (fromOL binds') }

------------------------
ds_lhs_binds :: AutoScc -> LHsBinds Id -> DsM (OrdList (Id,CoreExpr))

	 -- scc annotation policy (see below)
ds_lhs_binds auto_scc binds = do { ds_bs <- mapBagM (dsLHsBind auto_scc) binds
                                 ; return (foldBag appOL id nilOL ds_bs) }

dsLHsBind :: AutoScc -> LHsBind Id -> DsM (OrdList (Id,CoreExpr))
dsLHsBind auto_scc (L loc bind)
  = putSrcSpanDs loc $ dsHsBind auto_scc bind

dsHsBind :: AutoScc -> HsBind Id -> DsM (OrdList (Id,CoreExpr))

dsHsBind _ (VarBind { var_id = var, var_rhs = expr, var_inline = inline_regardless })
  = do  { core_expr <- dsLExpr expr

	        -- Dictionary bindings are always VarBinds,
	        -- so we only need do this here
	; core_expr' <- addDictScc var core_expr
	; let var' | inline_regardless = var `setIdUnfolding` mkCompulsoryUnfolding core_expr'
	      	   | otherwise         = var

	; return (unitOL (var', core_expr')) }

dsHsBind auto_scc (FunBind { fun_id = L _ fun, fun_matches = matches
		    , fun_co_fn = co_fn, fun_tick = tick 
                    , fun_infix = inf }) 
 = do	{ (args, body) <- matchWrapper (FunRhs (idName fun) inf) matches
	; body'    <- mkOptTickBox tick body
	; wrap_fn' <- dsHsWrapper co_fn 
	; let rhs = addAutoScc auto_scc fun $ wrap_fn' (mkLams args body')
	; return (unitOL (makeCorePair fun False 0 rhs)) }

dsHsBind auto_scc (PatBind { pat_lhs = pat, pat_rhs = grhss, pat_rhs_ty = ty })
  = do	{ body_expr <- dsGuarded grhss ty
	; sel_binds <- mkSelectorBinds pat body_expr
	  -- We silently ignore inline pragmas; no makeCorePair
	  -- Not so cool, but really doesn't matter
    ; let sel_binds' = [ (v, addAutoScc auto_scc v expr)
                       | (v, expr) <- sel_binds ]
	; return (toOL sel_binds') }

	-- A common case: one exported variable
	-- Non-recursive bindings come through this way
	-- So do self-recursive bindings, and recursive bindings
	-- that have been chopped up with type signatures
dsHsBind auto_scc (AbsBinds { abs_tvs = all_tyvars, abs_ev_vars = dicts
               		    , abs_exports = [(tyvars, global, local, prags)]
               		    , abs_ev_binds = ev_binds, abs_binds = binds })
  = ASSERT( all (`elem` tyvars) all_tyvars )
    do	{ bind_prs    <- ds_lhs_binds NoSccs binds
        ; ds_ev_binds <- dsTcEvBinds ev_binds

	; let	core_bind = Rec (fromOL bind_prs)
	        rhs       = addAutoScc auto_scc global $
			    mkLams tyvars $ mkLams dicts $ 
	                    wrapDsEvBinds ds_ev_binds $
                            Let core_bind $
                            Var local
    
	; (spec_binds, rules) <- dsSpecs rhs prags

	; let   global'   = addIdSpecialisations global rules
		main_bind = makeCorePair global' (isDefaultMethod prags)
                                         (dictArity dicts) rhs 
    
	; return (main_bind `consOL` spec_binds) }

dsHsBind auto_scc (AbsBinds { abs_tvs = all_tyvars, abs_ev_vars = dicts
                            , abs_exports = exports, abs_ev_binds = ev_binds
                       	    , abs_binds = binds })
  = do	{ bind_prs    <- ds_lhs_binds NoSccs binds
        ; ds_ev_binds <- dsTcEvBinds ev_binds
	; let env = mkABEnv exports
	      do_one (lcl_id,rhs) | Just (_, gbl_id, _, _prags) <- lookupVarEnv env lcl_id
			          = (lcl_id, addAutoScc auto_scc gbl_id rhs)
				  | otherwise = (lcl_id,rhs)
	       
	      core_bind = Rec (map do_one (fromOL bind_prs))
	      	-- Monomorphic recursion possible, hence Rec

	      tup_expr     = mkBigCoreVarTup locals
	      tup_ty	   = exprType tup_expr
	      poly_tup_rhs = mkLams all_tyvars $ mkLams dicts $
	      		     wrapDsEvBinds ds_ev_binds $
			     Let core_bind $
	 	     	     tup_expr
	      locals       = [local | (_, _, local, _) <- exports]
	      local_tys    = map idType locals

	; poly_tup_id <- newSysLocalDs (exprType poly_tup_rhs)

	; let mk_bind ((tyvars, global, _, spec_prags), n)  -- locals!!n == local
	        = 	-- Need to make fresh locals to bind in the selector,
		      	-- because some of the tyvars will be bound to 'Any'
		  do { let ty_args = map mk_ty_arg all_tyvars
		           substitute = substTyWith all_tyvars ty_args
		     ; locals' <- newSysLocalsDs (map substitute local_tys)
		     ; tup_id  <- newSysLocalDs  (substitute tup_ty)
	             ; let rhs = mkLams tyvars $ mkLams dicts $
	      	     		 mkTupleSelector locals' (locals' !! n) tup_id $
			         mkVarApps (mkTyApps (Var poly_tup_id) ty_args)
			 		   dicts
                           full_rhs = Let (NonRec poly_tup_id poly_tup_rhs) rhs
		     ; (spec_binds, rules) <- dsSpecs full_rhs spec_prags
				                      
		     ; let global' = addIdSpecialisations global rules
		     ; return ((global', rhs) `consOL` spec_binds) }
	        where
	          mk_ty_arg all_tyvar
			| all_tyvar `elem` tyvars = mkTyVarTy all_tyvar
	      		| otherwise		  = dsMkArbitraryType all_tyvar

	; export_binds_s <- mapM mk_bind (exports `zip` [0..])
	     -- Don't scc (auto-)annotate the tuple itself.

	; return ((poly_tup_id, poly_tup_rhs) `consOL` 
		    concatOL export_binds_s) }

--------------------------------------
data DsEvBind 
  = LetEvBind		-- Dictionary or coercion
      CoreBind		-- recursive or non-recursive

  | CaseEvBind		-- Coercion binding by superclass selection
    			-- Desugars to case d of d { K _ g _ _ _ -> ... } 			
      DictId 		   -- b   The dictionary
      AltCon 		   -- K   Its constructor
      [CoreBndr] 	   -- _ g _ _ _   The binders in the alternative

wrapDsEvBinds :: [DsEvBind] -> CoreExpr -> CoreExpr
wrapDsEvBinds ds_ev_binds body = foldr wrap_one body ds_ev_binds
  where
    body_ty = exprType body
    wrap_one (LetEvBind b)       body = Let b body
    wrap_one (CaseEvBind x k xs) body = Case (Var x) x body_ty [(k,xs,body)]

dsTcEvBinds :: TcEvBinds -> DsM [DsEvBind]
dsTcEvBinds (TcEvBinds {}) = panic "dsEvBinds"	-- Zonker has got rid of this
dsTcEvBinds (EvBinds bs)   = dsEvBinds bs

dsEvBinds :: Bag EvBind -> DsM [DsEvBind]
dsEvBinds bs = return (map dsEvGroup sccs)
  where
    sccs :: [SCC EvBind]
    sccs = stronglyConnCompFromEdgedVertices edges

    edges :: [(EvBind, EvVar, [EvVar])]
    edges = foldrBag ((:) . mk_node) [] bs 

    mk_node :: EvBind -> (EvBind, EvVar, [EvVar])
    mk_node b@(EvBind var term) = (b, var, free_vars_of term)

    free_vars_of :: EvTerm -> [EvVar]
    free_vars_of (EvId v)           = [v]
    free_vars_of (EvCast v co)      = v : varSetElems (tyVarsOfType co)
    free_vars_of (EvCoercion co)    = varSetElems (tyVarsOfType co)
    free_vars_of (EvDFunApp _ _ vs) = vs
    free_vars_of (EvSuperClass d _) = [d]

dsEvGroup :: SCC EvBind -> DsEvBind
dsEvGroup (AcyclicSCC (EvBind co_var (EvSuperClass dict n)))
  | isCoVar co_var	 -- An equality superclass
  = ASSERT( null other_data_cons )
    CaseEvBind dict (DataAlt data_con) bndrs
  where
    (cls, tys) = getClassPredTys (evVarPred dict)
    (data_con:other_data_cons) = tyConDataCons (classTyCon cls)
    (ex_tvs, theta, rho) = tcSplitSigmaTy (applyTys (dataConRepType data_con) tys)
    (arg_tys, _) = splitFunTys rho
    bndrs = ex_tvs ++ map mk_wild_pred (theta `zip` [0..])
                   ++ map mkWildValBinder arg_tys
    mk_wild_pred (p, i) | i==n      = ASSERT( p `tcEqPred` (coVarPred co_var)) 
                                      co_var
                        | otherwise = mkWildEvBinder p
    
dsEvGroup (AcyclicSCC (EvBind v r))
  = LetEvBind (NonRec v (dsEvTerm r))

dsEvGroup (CyclicSCC bs)
  = LetEvBind (Rec (map ds_pair bs))
  where
    ds_pair (EvBind v r) = (v, dsEvTerm r)

dsEvTerm :: EvTerm -> CoreExpr
dsEvTerm (EvId v)                = Var v
dsEvTerm (EvCast v co)           = Cast (Var v) co
dsEvTerm (EvDFunApp df tys vars) = Var df `mkTyApps` tys `mkVarApps` vars
dsEvTerm (EvCoercion co)         = Type co
dsEvTerm (EvSuperClass d n)
  = ASSERT( isClassPred (classSCTheta cls !! n) )
    	    -- We can only select *dictionary* superclasses
	    -- in terms.  Equality superclasses are dealt with
	    -- in dsEvGroup, where they can generate a case expression
    Var sc_sel_id `mkTyApps` tys `App` Var d
  where
    sc_sel_id  = classSCSelId cls n	-- Zero-indexed
    (cls, tys) = getClassPredTys (evVarPred d)    
    
------------------------
makeCorePair :: Id -> Bool -> Arity -> CoreExpr -> (Id, CoreExpr)
makeCorePair gbl_id is_default_method dict_arity rhs
  | is_default_method		      -- Default methods are *always* inlined
  = (gbl_id `setIdUnfolding` mkCompulsoryUnfolding rhs, rhs)

  | otherwise
  = case inlinePragmaSpec inline_prag of
      	  EmptyInlineSpec -> (gbl_id, rhs)
      	  NoInline        -> (gbl_id, rhs)
      	  Inlinable       -> (gbl_id `setIdUnfolding` inlinable_unf, rhs)
          Inline          -> inline_pair

  where
    inline_prag   = idInlinePragma gbl_id
    inlinable_unf = mkInlinableUnfolding rhs
    inline_pair
       | Just arity <- inlinePragmaSat inline_prag
      	-- Add an Unfolding for an INLINE (but not for NOINLINE)
	-- And eta-expand the RHS; see Note [Eta-expanding INLINE things]
       , let real_arity = dict_arity + arity
        -- NB: The arity in the InlineRule takes account of the dictionaries
       = ( gbl_id `setIdUnfolding` mkInlineUnfolding (Just real_arity) rhs
         , etaExpand real_arity rhs)

       | otherwise
       = pprTrace "makeCorePair: arity missing" (ppr gbl_id) $
         (gbl_id `setIdUnfolding` mkInlineUnfolding Nothing rhs, rhs)


dictArity :: [Var] -> Arity
-- Don't count coercion variables in arity
dictArity dicts = count isId dicts


------------------------
type AbsBindEnv = VarEnv ([TyVar], Id, Id, TcSpecPrags)
	-- Maps the "lcl_id" for an AbsBind to
	-- its "gbl_id" and associated pragmas, if any

mkABEnv :: [([TyVar], Id, Id, TcSpecPrags)] -> AbsBindEnv
-- Takes the exports of a AbsBinds, and returns a mapping
--	lcl_id -> (tyvars, gbl_id, lcl_id, prags)
mkABEnv exports = mkVarEnv [ (lcl_id, export) | export@(_, _, lcl_id, _) <- exports]

------------------------
dsSpecs :: CoreExpr     -- Its rhs
        -> TcSpecPrags
        -> DsM ( OrdList (Id,CoreExpr) 	-- Binding for specialised Ids
	       , [CoreRule] )		-- Rules for the Global Ids
-- See Note [Implementing SPECIALISE pragmas]
dsSpecs _ IsDefaultMethod = return (nilOL, [])
dsSpecs poly_rhs (SpecPrags sps)
  = do { pairs <- mapMaybeM (dsSpec (Just poly_rhs)) sps
       ; let (spec_binds_s, rules) = unzip pairs
       ; return (concatOL spec_binds_s, rules) }

dsSpec :: Maybe CoreExpr  	-- Just rhs => RULE is for a local binding
       	  			-- Nothing => RULE is for an imported Id
				-- 	      rhs is in the Id's unfolding
       -> Located TcSpecPrag
       -> DsM (Maybe (OrdList (Id,CoreExpr), CoreRule))
dsSpec mb_poly_rhs (L loc (SpecPrag poly_id spec_co spec_inl))
  = putSrcSpanDs loc $ 
    do { let poly_name = idName poly_id
       ; spec_name <- newLocalName poly_name
       ; wrap_fn   <- dsHsWrapper spec_co
       ; let (bndrs, ds_lhs) = collectBinders (wrap_fn (Var poly_id))
             spec_ty = mkPiTypes bndrs (exprType ds_lhs)
       ; case decomposeRuleLhs bndrs ds_lhs of {
           Left msg -> do { warnDs msg; return Nothing } ;
           Right (final_bndrs, _fn, args) -> do

       { (spec_unf, unf_pairs) <- specUnfolding wrap_fn spec_ty (realIdUnfolding poly_id)

       ; let spec_id  = mkLocalId spec_name spec_ty 
         	            `setInlinePragma` inl_prag
         	 	    `setIdUnfolding`  spec_unf
             inl_prag | isDefaultInlinePragma spec_inl = idInlinePragma poly_id
         	      | otherwise                      = spec_inl
       	      -- Get the INLINE pragma from SPECIALISE declaration, or,
              -- failing that, from the original Id

             rule =  mkRule False {- Not auto -} is_local_id
                        (mkFastString ("SPEC " ++ showSDoc (ppr poly_name)))
       			AlwaysActive poly_name
       		        final_bndrs args
       			(mkVarApps (Var spec_id) bndrs)

             spec_rhs  = wrap_fn poly_rhs
             spec_pair = makeCorePair spec_id False (dictArity bndrs) spec_rhs

       ; return (Just (spec_pair `consOL` unf_pairs, rule))
       } } }
  where
    is_local_id = isJust mb_poly_rhs
    poly_rhs | Just rhs <-  mb_poly_rhs
             = rhs  	    -- Local Id; this is its rhs
             | Just unfolding <- maybeUnfoldingTemplate (realIdUnfolding poly_id)
             = unfolding    -- Imported Id; this is its unfolding
	       		    -- Use realIdUnfolding so we get the unfolding 
			    -- even when it is a loop breaker. 
			    -- We want to specialise recursive functions!
             | otherwise = pprPanic "dsImpSpecs" (ppr poly_id)
	                    -- The type checker has checked that it *has* an unfolding

specUnfolding :: (CoreExpr -> CoreExpr) -> Type 
              -> Unfolding -> DsM (Unfolding, OrdList (Id,CoreExpr))
{-   [Dec 10: TEMPORARILY commented out, until we can straighten out how to
              generate unfoldings for specialised DFuns

specUnfolding wrap_fn spec_ty (DFunUnfolding _ _ ops)
  = do { let spec_rhss = map wrap_fn ops
       ; spec_ids <- mapM (mkSysLocalM (fsLit "spec") . exprType) spec_rhss
       ; return (mkDFunUnfolding spec_ty (map Var spec_ids), toOL (spec_ids `zip` spec_rhss)) }
-}
specUnfolding _ _ _
  = return (noUnfolding, nilOL)

dsMkArbitraryType :: TcTyVar -> Type
dsMkArbitraryType tv = anyTypeOfKind (tyVarKind tv)

decomposeRuleLhs :: [Var] -> CoreExpr -> Either SDoc ([Var], Id, [CoreExpr])
-- Take apart the LHS of a RULE.  It's suuposed to look like
--     /\a. f a Int dOrdInt
-- or  /\a.\d:Ord a. let { dl::Ord [a] = dOrdList a d } in f [a] dl
-- That is, the RULE binders are lambda-bound
-- Returns Nothing if the LHS isn't of the expected shape
decomposeRuleLhs bndrs lhs 
  =  -- Note [Simplifying the left-hand side of a RULE]
    case collectArgs opt_lhs of
        (Var fn, args) -> check_bndrs fn args

        (Case scrut bndr ty [(DEFAULT, _, body)], args)
	        | isDeadBinder bndr	-- Note [Matching seqId]
		-> check_bndrs seqId (args' ++ args)
		where
		   args' = [Type (idType bndr), Type ty, scrut, body]
	   
	_other -> Left bad_shape_msg
 where
   opt_lhs = simpleOptExpr lhs

   check_bndrs fn args
     | null (dead_bndrs) = Right (extra_dict_bndrs ++ bndrs, fn, args)
     | otherwise         = Left (vcat (map dead_msg dead_bndrs))
     where
       arg_fvs = exprsFreeVars args

            -- Check for dead binders: Note [Unused spec binders]
       dead_bndrs = filterOut (`elemVarSet` arg_fvs) bndrs

            -- Add extra dict binders: Note [Constant rule dicts]
       extra_dict_bndrs = [ mkLocalId (localiseName (idName d)) (idType d)
                          | d <- varSetElems (arg_fvs `delVarSetList` bndrs)
         	          , isDictId d]


   bad_shape_msg = hang (ptext (sLit "RULE left-hand side too complicated to desugar"))
                      2 (ppr opt_lhs)
   dead_msg bndr = hang (ptext (sLit "Forall'd") <+> pp_bndr bndr
				 <+> ptext (sLit "is not bound in RULE lhs"))
                      2 (ppr opt_lhs)
   pp_bndr bndr
    | isTyVar bndr = ptext (sLit "type variable") <+> ppr bndr
    | isCoVar bndr = ptext (sLit "coercion variable") <+> ppr bndr
    | isDictId bndr = ptext (sLit "constraint") <+> ppr (get_pred bndr)
    | otherwise     = ptext (sLit "variable") <+> ppr bndr

   get_pred b = ASSERT( isId b ) expectJust "decomposeRuleLhs" 
                                 (tcSplitPredTy_maybe (idType b))

data AutoScc = NoSccs 
	     | AddSccs Module (Id -> Bool)
-- The (Id->Bool) says which Ids to add SCCs to 
-- But we never add a SCC to function marked INLINE

addAutoScc :: AutoScc	
	   -> Id	-- Binder
	   -> CoreExpr 	-- Rhs
	   -> CoreExpr	-- Scc'd Rhs

addAutoScc NoSccs _ rhs
  = rhs
addAutoScc _ id rhs | isInlinePragma (idInlinePragma id)
  = rhs
addAutoScc (AddSccs mod add_scc) id rhs
  | add_scc id = mkSCC (mkAutoCC id mod NotCafCC) rhs
  | otherwise  = rhs

addDictScc :: Id -> CoreExpr -> DsM CoreExpr
addDictScc _ rhs = return rhs

{- DISABLED for now (need to somehow make up a name for the scc) -- SDM
  | not ( opt_SccProfilingOn && opt_AutoSccsOnDicts)
    || not (isDictId var)
  = return rhs				-- That's easy: do nothing

  | otherwise
  = do (mod, grp) <- getModuleAndGroupDs
	-- ToDo: do -dicts-all flag (mark dict things with individual CCs)
       return (Note (SCC (mkAllDictsCC mod grp False)) rhs)
-}

dsHsWrapper :: HsWrapper -> DsM (CoreExpr -> CoreExpr)
dsHsWrapper WpHole 	      = return (\e -> e)
dsHsWrapper (WpTyApp ty)      = return (\e -> App e (Type ty))
dsHsWrapper (WpLet ev_binds)  = do { ds_ev_binds <- dsTcEvBinds ev_binds
                                   ; return (wrapDsEvBinds ds_ev_binds) }
dsHsWrapper (WpCompose c1 c2) = do { k1 <- dsHsWrapper c1 
                                   ; k2 <- dsHsWrapper c2
                                   ; return (k1 . k2) }
dsHsWrapper (WpCast co)       = return (\e -> Cast e co) 
dsHsWrapper (WpEvLam ev)      = return (\e -> Lam ev e) 
dsHsWrapper (WpTyLam tv)      = return (\e -> Lam tv e) 
dsHsWrapper (WpEvApp evtrm)   = return (\e -> App e (dsEvTerm evtrm))