module ISimplify(iSimplify) where
import Data.List((\\))
import qualified Data.Map as M
import PPrint
import IntLit
import ErrorUtil
import Id
import ISyntax
import Prim
import ISyntaxUtil(aitBit, isTrue, isFalse, isitActionValue_, isitActionValue, iDefsMap)
import ISyntaxXRef(mapIExprPosition)
import Eval
--import Debug.Trace
--import Util(traces)
--import Position

--
-- TODO
--  Useful transformations (from looking at .bo)
--   * let x = (a,b) in ... fst x ... snd x ...  -->  ... a ... b ...
--     note that the tuple transformation is important to inline dictionaries (I think)
--   * PrimConcat 0 n n _ x  -->  x
--   * let d = Literal-Bits n in ... .fromInteger d NNN  -->   NNN

iSimplify :: (Hyper a) => IPackage a -> IPackage a
iSimplify (IPackage pi lps ps ds) =
    IPackage pi lps ps ({-iSimpDefs-} (iSimpDefs (iSimpDefs ds)))        -- XXX

iSimpDefs :: Hyper a => [IDef a] -> [IDef a]
iSimpDefs ds = fixUpDefs $ iDefsMap (iSimp True) ds

iSimp :: (Hyper a) => Bool -> IExpr a -> IExpr a
iSimp n (ILam i t e) = ILam i t (iSimp n e)
iSimp n (IAps e ts as) = iSimpAp' n (iSimp n (expDef e)) ts (map (iSimp n) as)
iSimp _ e@(IVar _) = e
iSimp n (ILAM i k e) = ILAM i k (iSimp n e)
iSimp _ e@(ICon _ _) = e
iSimp _ e = internalError ("iSimp: " ++ ppReadable e)

expDef :: IExpr a -> IExpr a
expDef (ICon _ (ICDef _ e)) | isHarmless e = e
expDef e = e

iSimpAp' :: Hyper a => Bool -> IExpr a -> [IType] -> [IExpr a] -> IExpr a
iSimpAp' b f ts es = mapIExprPosition True (f, iSimpAp b f ts es)

iSimpAp :: (Hyper a) => Bool -> IExpr a -> [IType] -> [IExpr a] -> IExpr a
iSimpAp n (ILAM i _ e) (t:ts) as = iSimpAp n (etSubst i t e) ts as
iSimpAp n (ILam i _ e) [] (a:as)
    | not (isKeepId i) && (isTriv a || countOcc i e <= 1) =
        let e' = eSubst i a e
        in iSimpAp n e' [] as
iSimpAp _ (ICon _ (ICPrim _ prim)) ts es | m /= Nothing = r
  where m = doPrim prim ts es
        Just r = m
iSimpAp n f@(ICon _ (ICSel { selNo = k })) ts
        es@(def : as) | n && m /= Nothing = {-trace (ppReadable (IAps f ts es, e'))-} e'
  where m = getTuple def
        Just ms = m
        e' = iSimpAp n (iSimp n (ms !! fromInteger k)) [] as
iSimpAp n e [] [] = e -- iSimp has already been called
iSimpAp n f ts es = IAps f ts es

getTuple :: (Hyper a) => IExpr a -> Maybe [IExpr a]
getTuple (ICon di (ICDef { iConDef = def@(IAps (ICon _ (ICTuple { })) _ ms) })) | di `notElem` dVars def =
        -- trace ("unfold " ++ ppReadable di) $
        Just ms
getTuple (IAps (ICon iii (ICDef { iConDef = body })) ts []) =
        -- trace ("getTuple " ++ ppReadable (iii,body)) $
        case iSimpAp False body ts [] of
        IAps (ICon _ (ICTuple { })) _ ms -> Just ms
        _ -> Nothing
getTuple _ = Nothing

-- XXX should we do more PrimOps here?
doPrim :: PrimOp -> [IType] -> [IExpr a] -> Maybe (IExpr a)
doPrim PrimIntegerToBit [t@(ITNum s)] [ICon i l@(ICInt { iVal = v })] | ilValue v >= 0 &&
                                                                        s >=0 &&
                                                                        ilValue v < 2^s = Just $ ICon i (l { iConType = aitBit t })
doPrim PrimOrd          [t,s] [IAps (ICon _ (ICPrim _ PrimChr)) [s',t'] [e]] | s == s' && t == t' = Just e
doPrim PrimIf _ [c, t, e] | isTrue  c = Just t
                          | isFalse c = Just e
doPrim _ _ _ = Nothing

{-
data Occ = None | One | Many
  deriving(Eq)
-}


countOcc :: Id -> IExpr a -> Int
countOcc i e = countOcc' i e 0

countOcc' :: Id -> IExpr a -> Int -> Int
countOcc' _ _ occ | occ > 1   = occ
countOcc' i (ILam i' _ e) occ = if i == i' then occ else countOcc' i e occ
countOcc' i (IAps f _ es) occ = foldr (countOcc' i) (countOcc' i f occ) es
countOcc' i (IVar i') occ     = if i == i' then occ + 1 else occ
countOcc' i (ILAM _ _ e) occ  = countOcc' i e occ
countOcc' _ _ occ             = occ

{-
size :: IExpr -> Int
size (ILam _ _ e) = size e
size (IAps f _ es) = size f + sum (map size es) + length es
size (ILAM _ _ e) = size e
size _ = 0
-}

isTriv :: IExpr a -> Bool
isTriv (IVar _) = True
-- do not inline ActionValue constants
-- may break correlations for foreign functions
isTriv (ICon _ ci) | isitActionValue_ t || isitActionValue t = False
  where t = iConType ci
isTriv (ICon _ (ICInt { })) = True
isTriv (ICon _ (ICReal { })) = True
isTriv (ICon _ (ICUndet { })) = True
isTriv (ICon _ (ICDef { })) = True
isTriv _ = False

isHarmless :: IExpr a -> Bool
isHarmless e =
        --trace (ppReadable (e, onlySimple e, isPerm [] e)) $
        onlySimple e && isPerm [] e

-- expression is a proper combinator: somehow permutes arguments and constants
-- has no free variables and no embedded lambda expressions
--   is: accumulated list of lambda-bound parameters
isPerm :: [Id] -> IExpr a -> Bool
isPerm is (ILAM _ _ e) = isPerm is e
isPerm is (ILam i _ e) = isPerm (i:is) e
isPerm is e = null (gVars e \\ is)

gVars :: IExpr a -> [Id]
gVars (ILam i t e) = gVars e
gVars (IVar i) = [i]
gVars (ILAM i _ e) = gVars e
gVars (IAps f ts es) = gVars f ++ concatMap gVars es
gVars (ICon _ _) = []
gVars (IRefT _ _ _) = []

-- computes the top-level definitions the expression depends onb
-- dVars :: IExpr a -> [Id]
-- dVars (ILam _ _ e) = dVars e
-- dVars (IVar _) = []
-- dVars (ILAM _ _ e) = dVars e
-- dVars (IAps f _ es) = dVars f ++ concatMap dVars es
-- dVars (ICon i (ICDef { })) = [i]
-- dVars (ICon _ _) = []
-- dVars (IRefT _ _ _) = []

dVars :: IExpr a -> [Id]
dVars e = dVars' [] e

-- auxiliary function to guard against circular traversals
dVars' :: [Id] -> IExpr a -> [Id]
dVars' ids (ILam _ _ e) = dVars' ids e
dVars' ids (IVar _) = []
dVars' ids (ILAM _ _ e) = dVars' ids e
dVars' ids (IAps f _ es) = dVars' ids f ++ concatMap (dVars' ids) es
-- guarding against circular traversal
dVars' ids (ICon i (ICDef { })) | i `elem` ids = ids
dVars' ids (ICon i (ICDef {iConDef = e})) = dVars' (i:ids) e
dVars' ids (ICon _ _) = []
dVars' ids (IRefT _ _ _) = []

onlySimple :: IExpr a -> Bool
onlySimple (ILam _ _ e) = onlySimple e
onlySimple (ILAM _ _ e) = onlySimple e
onlySimple (IAps f _ es) = onlySimple f && all onlySimple es
onlySimple (IVar _) = True
onlySimple (ICon _ (ICPrim { })) = True
onlySimple (ICon _ (ICInt { })) = True
onlySimple (ICon _ (ICReal { })) = True
onlySimple (ICon _ (ICUndet { })) = True
-- note that foreign function calls are not simple
onlySimple e = False

-----

-- should combine with FixupDefs

fixUpDefs :: [IDef a] -> [IDef a]
fixUpDefs ds =
    let m = M.fromList [ (i, e) | IDef i _ e _ <- ds ]
    in  iDefsMap (fixUp m) ds

fixUp :: M.Map Id (IExpr a) -> IExpr a -> IExpr a
fixUp m (ILam i t e) = ILam i t (fixUp m e)
fixUp m (ILAM i k e) = ILAM i k (fixUp m e)
fixUp m (IAps f ts es) = IAps (fixUp m f) ts (map (fixUp m) es)
fixUp m (ICon i (ICDef t d)) = ICon i (ICDef t (M.findWithDefault d i m))
fixUp m e = e