content: module Language.Haskell.Exference.Core.ExpressionSimplify
    ( simplifyExpression
    )
  where
  
  
  
  import Language.Haskell.Exference.Core.Types
  import Language.Haskell.Exference.Core.Expression
  
  import Data.Function ( on )
  
  
  
  simplifyExpression :: Expression -> Expression
  simplifyExpression = simplifyId . simplifyCompose . simplifyEta . simplifyLets
  
  
  
  simplifyLets :: Expression -> Expression
  simplifyLets e@ExpVar{}       = e
  simplifyLets e@ExpName{}      = e
  simplifyLets (ExpLambda i ty e)  = ExpLambda i ty $ simplifyLets e
  simplifyLets (ExpApply e1 e2) = ExpApply (simplifyLets e1) (simplifyLets e2)
  simplifyLets e@ExpHole{}      = e
  simplifyLets (ExpLetMatch name vids bindExp inExp) =
    ExpLetMatch name vids (simplifyLets bindExp) (simplifyLets inExp)
  simplifyLets (ExpLet i ty bindExp inExp) = case countUses i inExp of
    0 -> simplifyLets inExp
    1 -> simplifyLets $ replaceVar i bindExp inExp
    _ -> ExpLet i ty (simplifyLets bindExp) (simplifyLets inExp)
  simplifyLets (ExpCaseMatch bindExp alts) =
    ExpCaseMatch (simplifyLets bindExp) [ (c, vars, simplifyLets expr)
                                        | (c, vars, expr) <- alts
                                        ]
  
  simplifyEta :: Expression -> Expression
  simplifyEta e@ExpVar{}         = e
  simplifyEta e@ExpName{}        = e
  simplifyEta (ExpLambda i ty e) = simplifyEta' $ ExpLambda i ty $ simplifyEta e
  simplifyEta (ExpApply e1 e2)   = ExpApply (simplifyEta e1) (simplifyEta e2)
  simplifyEta e@ExpHole{}        = e
  simplifyEta (ExpLetMatch name vids bindExp inExp) =
    ExpLetMatch name vids (simplifyEta bindExp) (simplifyEta inExp)
  simplifyEta (ExpLet i ty bindExp inExp) =
    ExpLet i ty (simplifyEta bindExp) (simplifyEta inExp)
  simplifyEta (ExpCaseMatch bindExp alts) =
    ExpCaseMatch (simplifyEta bindExp) [ (c, vars, simplifyEta expr)
                                       | (c, vars, expr) <- alts
                                       ]
  
  simplifyEta' :: Expression -> Expression
  simplifyEta' (ExpLambda i _ (ExpApply e1 (ExpVar j _)))
    | i==j && 0==countUses i e1 = e1
  simplifyEta' e = e
  
  
  simplifyId :: Expression -> Expression
  simplifyId e@ExpVar{}                 = e
  simplifyId e@ExpName{}                = e
  simplifyId e@(ExpLambda i _ (ExpVar j _))
    | i == j                            = ExpName (QualifiedName [] "id")
    | otherwise                         = e
  simplifyId (ExpLambda i ty e)         = ExpLambda i ty $ simplifyId e
  simplifyId (ExpApply e1 e2)           = ExpApply (simplifyId e1) (simplifyId e2)
  simplifyId e@ExpHole{}                = e
  simplifyId (ExpLetMatch name vids bindExp inExp)
                                        = ExpLetMatch name
                                                      vids
                                                      (simplifyId bindExp)
                                                      (simplifyId inExp)
  simplifyId (ExpLet i ty bindExp inExp)   =
    ExpLet i ty (simplifyId bindExp) (simplifyId inExp)
  simplifyId (ExpCaseMatch bindExp alts) =
    ExpCaseMatch (simplifyId bindExp) [ (c, vars, simplifyId expr)
                                      | (c, vars, expr) <- alts
                                      ]
  
  simplifyCompose :: Expression -> Expression
  simplifyCompose e@ExpVar{}         = e
  simplifyCompose e@ExpName{}        = e
  simplifyCompose (ExpLambda i ty e) = simplifyCompose' i ty [] e
  simplifyCompose (ExpApply e1 e2)   = ExpApply (simplifyCompose e1) (simplifyCompose e2)
  simplifyCompose e@ExpHole{}        = e
  simplifyCompose (ExpLetMatch name vids bindExp inExp)
                                   = ExpLetMatch name
                                                 vids
                                                 (simplifyCompose bindExp)
                                                 (simplifyCompose inExp)
  simplifyCompose (ExpLet i ty bindExp inExp)
                                   = ExpLet i
                                            ty
                                            (simplifyCompose bindExp)
                                            (simplifyCompose inExp)
  simplifyCompose (ExpCaseMatch bindExp alts)
                                   = ExpCaseMatch
                                       (simplifyCompose bindExp)
                                       [ (c, vars, simplifyCompose expr)
                                       | (c, vars, expr) <- alts
                                       ]
  
  simplifyCompose' :: Int -> HsType -> [Expression] -> Expression -> Expression
  simplifyCompose' i ty [] e@ExpVar{} = ExpLambda i ty e
  simplifyCompose' i ty l e@(ExpVar j _)
    | i == j    = foldl1 (\e1 e2 -> ExpApply
                                      (ExpApply
                                        (ExpName (QualifiedName [] "(.)")) e2) e1)
                         l
    | otherwise = ExpLambda i ty $ foldl (flip ExpApply) e l
  simplifyCompose' i ty l (ExpApply e1 e2)
    | countUses i e1==0 = simplifyCompose' i ty (simplifyCompose e1:l) e2
  simplifyCompose' i ty l e = ExpLambda i ty
                            $ foldl (flip ExpApply) (simplifyCompose e) l
  
  replaceVar :: TVarId -> Expression -> Expression -> Expression
  replaceVar vid t orig@(ExpVar j _) | vid==j = t
                                     | otherwise = orig
  replaceVar vid t (ExpLambda i ty e) = ExpLambda i ty $ replaceVar vid t e
  replaceVar vid t (ExpApply e1 e2) = ExpApply (replaceVar vid t e1)
                                               (replaceVar vid t e2)
  replaceVar vid t (ExpLetMatch n vars bindExp inExp) =
    ExpLetMatch n vars (replaceVar vid t bindExp) (replaceVar vid t inExp)
  replaceVar vid t (ExpLet i ty bindExp inExp) =
    ExpLet i ty (replaceVar vid t bindExp) (replaceVar vid t inExp)
  replaceVar vid t (ExpCaseMatch bindExp alts) =
    ExpCaseMatch (replaceVar vid t bindExp) [ (c, vars, replaceVar vid t expr)
                                            | (c, vars, expr) <- alts
                                            ]
  replaceVar _ _ t@(ExpName _) = t
  replaceVar _ _ t@(ExpHole _) = t
  
  
  countUses :: TVarId -> Expression -> Int
  countUses i (ExpVar j _) | i==j           = 1
                           | otherwise      = 0
  countUses _ ExpName{}                     = 0
  countUses i (ExpLambda j _ e) | i==j      = 0
                                | otherwise = countUses i e
  countUses i (ExpApply e1 e2)              = ((+) `on` countUses i) e1 e2
  countUses _ ExpHole{}                     = 0
  countUses i (ExpLetMatch _ _ bindExp inExp) = ((+) `on` countUses i) bindExp inExp
  countUses i (ExpLet _ _ bindExp inExp)    = ((+) `on` countUses i) bindExp inExp
  countUses i (ExpCaseMatch bindExp alts)   = sum $ countUses i bindExp
                                                  : [ countUses i expr
                                                    | (_, _, expr) <- alts
                                                    ]
  
  