%
% (c) The University of Glasgow 2006
%

\begin{code}
{-# OPTIONS_GHC -w #-}
module OptCoercion (
	optCoercion
   ) where 

#include "HsVersions.h"

import Unify	( tcMatchTy )
import Coercion
import Type
import TypeRep
import TyCon
import Var
import VarSet
import PrelNames
import Util
import Outputable
\end{code}

%************************************************************************
%*                                                                      *
                 Optimising coercions									
%*                                                                      *
%************************************************************************

\begin{code}
optCoercion :: TvSubst -> Coercion -> NormalCo
-- ^ optCoercion applies a substitution to a coercion, 
--   *and* optimises it to reduce its size
optCoercion env co = opt_co env False co

type NormalCo = Coercion
  -- Invariants: 
  --  * The substitution has been fully applied
  --  * For trans coercions (co1 `trans` co2)
  --       co1 is not a trans, and neither co1 nor co2 is identity
  --  * If the coercion is the identity, it has no CoVars of CoTyCons in it (just types)

type NormalNonIdCo = NormalCo  -- Extra invariant: not the identity

opt_co, opt_co' :: TvSubst
       		-> Bool	       -- True <=> return (sym co)
       		-> Coercion
       		-> NormalCo	
opt_co = opt_co'
-- opt_co sym co = pprTrace "opt_co {" (ppr sym <+> ppr co) $
--       	        co1 `seq` 
--                pprTrace "opt_co done }" (ppr co1) 
--               WARN( not same_co_kind, ppr co  <+> dcolon <+> pprEqPred (s1,t1) 
--                                     $$ ppr co1 <+> dcolon <+> pprEqPred (s2,t2) )
--                co1
--  where
--    co1 = opt_co' sym co
--    same_co_kind = s1 `coreEqType` s2 && t1 `coreEqType` t2
--    (s,t) = coercionKind co
--    (s1,t1) | sym = (t,s)
--            | otherwise = (s,t)
--    (s2,t2) = coercionKind co1

opt_co' env sym (AppTy ty1 ty2) 	  = mkAppTy (opt_co env sym ty1) (opt_co env sym ty2)
opt_co' env sym (FunTy ty1 ty2) 	  = FunTy (opt_co env sym ty1) (opt_co env sym ty2)
opt_co' env sym (PredTy (ClassP cls tys)) = PredTy (ClassP cls (map (opt_co env sym) tys))
opt_co' env sym (PredTy (IParam n ty))    = PredTy (IParam n (opt_co env sym ty))
opt_co' _   _   co@(PredTy (EqPred {}))   = pprPanic "optCoercion" (ppr co)

opt_co' env sym co@(TyVarTy tv)
  | Just ty <- lookupTyVar env tv = opt_co' (zapTvSubstEnv env) sym ty
  | not (isCoVar tv)     = co   -- Identity; does not mention a CoVar
  | ty1 `coreEqType` ty2 = ty1	-- Identity; ..ditto..
  | not sym              = co
  | otherwise            = mkSymCoercion co
  where
    (ty1,ty2) = coVarKind tv

opt_co' env sym (ForAllTy tv cor) 
  | isCoVar tv = mkCoPredTy (opt_co env sym co1) (opt_co env sym co2) (opt_co env sym cor)
  | otherwise  = case substTyVarBndr env tv of
                   (env', tv') -> ForAllTy tv' (opt_co env' sym cor)
  where
    (co1,co2) = coVarKind tv

opt_co' env sym (TyConApp tc cos)
  | Just (arity, desc) <- isCoercionTyCon_maybe tc
  = mkAppTys (opt_co_tc_app env sym tc desc (take arity cos))
             (map (opt_co env sym) (drop arity cos))
  | otherwise
  = TyConApp tc (map (opt_co env sym) cos)

--------
opt_co_tc_app :: TvSubst -> Bool -> TyCon -> CoTyConDesc -> [Coercion] -> NormalCo
-- Used for CoercionTyCons only
-- Arguments are *not* already simplified/substituted
opt_co_tc_app env sym tc desc cos
  = case desc of
      CoAxiom {} -- Do *not* push sym inside top-level axioms
    		 -- e.g. if g is a top-level axiom
    		 --   g a : F a ~ a
		 -- Then (sym (g ty)) /= g (sym ty) !!
        | sym       -> mkSymCoercion the_co  
        | otherwise -> the_co
        where
           the_co = TyConApp tc (map (opt_co env False) cos)
           -- Note that the_co does *not* have sym pushed into it
    
      CoTrans 
        | sym       -> opt_trans opt_co2 opt_co1   -- sym (g `o` h) = sym h `o` sym g
        | otherwise -> opt_trans opt_co1 opt_co2

      CoUnsafe
        | sym       -> mkUnsafeCoercion ty2' ty1'
        | otherwise -> mkUnsafeCoercion ty1' ty2'

      CoSym   -> opt_co env (not sym) co1
      CoLeft  -> opt_lr fst
      CoRight -> opt_lr snd
      CoCsel1 -> opt_csel fstOf3
      CoCsel2 -> opt_csel sndOf3
      CoCselR -> opt_csel thirdOf3

      CoInst        -- See if the first arg is already a forall
		    -- ...then we can just extend the current substitution
        | Just (tv, co1_body) <- splitForAllTy_maybe co1
        -> opt_co (extendTvSubst env tv ty2') sym co1_body

                    -- See if is *now* a forall
        | Just (tv, opt_co1_body) <- splitForAllTy_maybe opt_co1
        -> substTyWith [tv] [ty2'] opt_co1_body	-- An inefficient one-variable substitution

        | otherwise
        -> TyConApp tc [opt_co1, ty2']

  where
    (co1 : cos1) = cos
    (co2 : _)    = cos1

    ty1' = substTy env co1
    ty2' = substTy env co2

	-- These opt_cos have the sym pushed into them
    opt_co1 = opt_co env sym co1
    opt_co2 = opt_co env sym co2

    the_unary_opt_co = TyConApp tc [opt_co1]

    opt_lr   sel = case splitAppTy_maybe opt_co1 of
                     Nothing -> the_unary_opt_co 
                     Just lr -> sel lr
    opt_csel sel = case splitCoPredTy_maybe opt_co1 of
                     Nothing -> the_unary_opt_co 
                     Just lr -> sel lr

-------------
opt_transL :: [NormalCo] -> [NormalCo] -> [NormalCo]
opt_transL = zipWith opt_trans

opt_trans :: NormalCo -> NormalCo -> NormalCo
opt_trans co1 co2
  | isIdNormCo co1 = co2
  | otherwise      = opt_trans1 co1 co2

opt_trans1 :: NormalNonIdCo -> NormalCo -> NormalCo
-- First arg is not the identity
opt_trans1 co1 co2
  | isIdNormCo co2 = co1
  | otherwise      = opt_trans2 co1 co2

opt_trans2 :: NormalNonIdCo -> NormalNonIdCo -> NormalCo
-- Neither arg is the identity
opt_trans2 (TyConApp tc [co1a,co1b]) co2
  | tc `hasKey` transCoercionTyConKey
  = opt_trans1 co1a (opt_trans2 co1b co2)

opt_trans2 co1 co2 
  | Just co <- opt_trans_rule co1 co2
  = co

opt_trans2 co1 (TyConApp tc [co2a,co2b])
  | tc `hasKey` transCoercionTyConKey
  , Just co1_2a <- opt_trans_rule co1 co2a
  = if isIdNormCo co1_2a
    then co2b
    else opt_trans2 co1_2a co2b

opt_trans2 co1 co2
  = mkTransCoercion co1 co2

------
opt_trans_rule :: NormalNonIdCo -> NormalNonIdCo -> Maybe NormalCo
opt_trans_rule (TyConApp tc1 args1) (TyConApp tc2 args2)
  | tc1 == tc2
  = case isCoercionTyCon_maybe tc1 of
      Nothing 
        -> Just (TyConApp tc1 (opt_transL args1 args2))
      Just (arity, desc) 
        | arity == length args1
        -> opt_trans_rule_equal_tc desc args1 args2
        | otherwise
        -> case opt_trans_rule_equal_tc desc 
                         (take arity args1) 
                         (take arity args2) of
              Just co -> Just $ mkAppTys co $ 
                         opt_transL (drop arity args1) (drop arity args2)
	      Nothing -> Nothing 
 
-- Push transitivity inside apply
opt_trans_rule co1 co2
  | Just (co1a, co1b) <- splitAppTy_maybe co1
  , Just (co2a, co2b) <- etaApp_maybe co2
  = Just (mkAppTy (opt_trans co1a co2a) (opt_trans co1b co2b))

  | Just (co2a, co2b) <- splitAppTy_maybe co2
  , Just (co1a, co1b) <- etaApp_maybe co1
  = Just (mkAppTy (opt_trans co1a co2a) (opt_trans co1b co2b))

-- Push transitivity inside (s~t)=>r
opt_trans_rule co1 co2
  | Just (s1,t1,r1) <- splitCoPredTy_maybe co1
  , Just (s2,t2,r2) <- etaCoPred_maybe co2
  = Just (mkCoPredTy (opt_trans s1 s2) (opt_trans t1 t2) (opt_trans r1 r2))

  | Just (s2,t2,r2) <- splitCoPredTy_maybe co2
  , Just (s1,t1,r1) <- etaCoPred_maybe co1
  = Just (mkCoPredTy (opt_trans s1 s2) (opt_trans t1 t2) (opt_trans r1 r2))

-- Push transitivity inside forall
opt_trans_rule co1 co2
  | Just (tv1,r1) <- splitTypeForAll_maybe co1
  , Just (tv2,r2) <- etaForAll_maybe co2
  , let r2' = substTyWith [tv2] [TyVarTy tv1] r2
  = Just (ForAllTy tv1 (opt_trans2 r1 r2'))

  | Just (tv2,r2) <- splitTypeForAll_maybe co2
  , Just (tv1,r1) <- etaForAll_maybe co1
  , let r1' = substTyWith [tv1] [TyVarTy tv2] r1
  = Just (ForAllTy tv1 (opt_trans2 r1' r2))

opt_trans_rule co1 co2
{- 	Omitting for now, because unsound
  | Just (sym1, (ax_tc1, ax1_args, ax_tvs, ax_lhs, ax_rhs)) <- co1_is_axiom_maybe
  , Just (sym2, (ax_tc2, ax2_args, _, _, _)) <- co2_is_axiom_maybe
  , ax_tc1 == ax_tc2
  , sym1 /= sym2
  = Just $
    if sym1 
    then substTyWith ax_tvs (opt_transL (map mkSymCoercion ax1_args) ax2_args) ax_rhs
    else substTyWith ax_tvs (opt_transL ax1_args (map mkSymCoercion ax2_args)) ax_lhs
-}

  | Just (sym, (ax_tc, ax_args, ax_tvs, ax_lhs, _)) <- co1_is_axiom_maybe
  , Just cos <- matchesAxiomLhs ax_tvs ax_lhs co2
  = Just $ 
    if sym 
    then mkSymCoercion $ TyConApp ax_tc (opt_transL (map mkSymCoercion cos) ax_args)
    else                 TyConApp ax_tc (opt_transL ax_args cos)

  | Just (sym, (ax_tc, ax_args, ax_tvs, ax_lhs, _)) <- isAxiom_maybe co2
  , Just cos <- matchesAxiomLhs ax_tvs ax_lhs co1
  = Just $ 
    if sym 
    then mkSymCoercion $ TyConApp ax_tc (opt_transL ax_args (map mkSymCoercion cos))
    else                 TyConApp ax_tc (opt_transL cos ax_args)
  where
    co1_is_axiom_maybe = isAxiom_maybe co1
    co2_is_axiom_maybe = isAxiom_maybe co2

opt_trans_rule co1 co2	-- Identity rule
  | (ty1,_) <- coercionKind co1
  , (_,ty2) <- coercionKind co2
  , ty1 `coreEqType` ty2
  = Just ty2

opt_trans_rule _ _ = Nothing

-----------  
isAxiom_maybe :: Coercion -> Maybe (Bool, (TyCon, [Coercion], [TyVar], Type, Type))
isAxiom_maybe co
  | Just (tc, args) <- splitTyConApp_maybe co
  , Just (_, desc)  <- isCoercionTyCon_maybe tc
  = case desc of
      CoAxiom { co_ax_tvs = tvs, co_ax_lhs = lhs, co_ax_rhs = rhs } 
            -> Just (False, (tc, args, tvs, lhs, rhs))
      CoSym | (arg1:_) <- args  
            -> case isAxiom_maybe arg1 of
                 Nothing           -> Nothing
                 Just (sym, stuff) -> Just (not sym, stuff)
      _ -> Nothing
  | otherwise
  = Nothing

matchesAxiomLhs :: [TyVar] -> Type -> Type -> Maybe [Type]
matchesAxiomLhs tvs ty_tmpl ty 
  = case tcMatchTy (mkVarSet tvs) ty_tmpl ty of
      Nothing    -> Nothing
      Just subst -> Just (map (substTyVar subst) tvs)

-----------  
opt_trans_rule_equal_tc :: CoTyConDesc -> [Coercion] -> [Coercion] -> Maybe Coercion
-- Rules for Coercion TyCons only

-- Push transitivity inside instantiation
opt_trans_rule_equal_tc desc [co1,ty1] [co2,ty2]
  | CoInst <- desc
  , ty1 `coreEqType` ty2
  , co1 `compatible_co` co2
  = Just (mkInstCoercion (opt_trans2 co1 co2) ty1) 

opt_trans_rule_equal_tc desc [co1] [co2]
  | CoLeft  <- desc, is_compat = Just (mkLeftCoercion res_co)
  | CoRight <- desc, is_compat = Just (mkRightCoercion res_co)
  | CoCsel1 <- desc, is_compat = Just (mkCsel1Coercion res_co)
  | CoCsel2 <- desc, is_compat = Just (mkCsel2Coercion res_co)
  | CoCselR <- desc, is_compat = Just (mkCselRCoercion res_co)
  where
    is_compat = co1 `compatible_co` co2
    res_co    = opt_trans2 co1 co2

opt_trans_rule_equal_tc _ _ _ = Nothing

-------------
compatible_co :: Coercion -> Coercion -> Bool
-- Check whether (co1 . co2) will be well-kinded
compatible_co co1 co2
  = x1 `coreEqType` x2		
  where
    (_,x1) = coercionKind co1
    (x2,_) = coercionKind co2

-------------
etaForAll_maybe :: Coercion -> Maybe (TyVar, Coercion)
-- Try to make the coercion be of form (forall tv. co)
etaForAll_maybe co
  | Just (tv, r) <- splitForAllTy_maybe co
  , not (isCoVar tv)	-- Check it is a *type* forall, not a (t1~t2)=>co
  = Just (tv, r)

  | (ty1,ty2) <- coercionKind co
  , Just (tv1, _) <- splitTypeForAll_maybe ty1
  , Just (tv2, _) <- splitTypeForAll_maybe ty2
  , tyVarKind tv1 `eqKind` tyVarKind tv2
  = Just (tv1, mkInstCoercion co (mkTyVarTy tv1))

  | otherwise
  = Nothing

etaCoPred_maybe :: Coercion -> Maybe (Coercion, Coercion, Coercion)
etaCoPred_maybe co 
  | Just (s,t,r) <- splitCoPredTy_maybe co
  = Just (s,t,r)
  
  --  co :: (s1~t1)=>r1 ~ (s2~t2)=>r2
  | (ty1,ty2) <- coercionKind co	-- We know ty1,ty2 have same kind
  , Just (s1,_,_) <- splitCoPredTy_maybe ty1
  , Just (s2,_,_) <- splitCoPredTy_maybe ty2
  , typeKind s1 `eqKind` typeKind s2	-- t1,t2 have same kinds
  = Just (mkCsel1Coercion co, mkCsel2Coercion co, mkCselRCoercion co)
  
  | otherwise
  = Nothing

etaApp_maybe :: Coercion -> Maybe (Coercion, Coercion)
-- Split a coercion g :: t1a t1b ~ t2a t2b
-- into (left g, right g) if possible
etaApp_maybe co
  | Just (co1, co2) <- splitAppTy_maybe co
  = Just (co1, co2)

  | (ty1,ty2) <- coercionKind co
  , Just (ty1a, _) <- splitAppTy_maybe ty1
  , Just (ty2a, _) <- splitAppTy_maybe ty2
  , typeKind ty1a `eqKind` typeKind ty2a
  = Just (mkLeftCoercion co, mkRightCoercion co)

  | otherwise
  = Nothing

-------------
splitTypeForAll_maybe :: Type -> Maybe (TyVar, Type)
-- Returns Just only for a *type* forall, not a (t1~t2)=>co
splitTypeForAll_maybe ty
  | Just (tv, rty) <- splitForAllTy_maybe ty
  , not (isCoVar tv)
  = Just (tv, rty)

  | otherwise
  = Nothing

-------------
isIdNormCo :: NormalCo -> Bool
-- Cheap identity test: look for coercions with no coercion variables at all
-- So it'll return False for (sym g `trans` g)
isIdNormCo ty = go ty
  where
    go (TyVarTy tv)  	       = not (isCoVar tv)
    go (AppTy t1 t2) 	       = go t1 && go t2
    go (FunTy t1 t2) 	       = go t1 && go t2
    go (ForAllTy tv ty)        = go (tyVarKind tv) && go ty
    go (TyConApp tc tys)       = not (isCoercionTyCon tc) && all go tys
    go (PredTy (IParam _ ty))  = go ty
    go (PredTy (ClassP _ tys)) = all go tys
    go (PredTy (EqPred t1 t2)) = go t1 && go t2
\end{code}