{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
module Simplify where

import Data.Monoid

import AST
import AST.Count
import Data


simplifyN :: KnownEnv env => Int -> Ex env t -> Ex env t
simplifyN 0 = id
simplifyN n = simplifyN (n - 1) . simplify

simplify :: forall env t. KnownEnv env => Ex env t -> Ex env t
simplify = let ?accumInScope = checkAccumInScope @env knownEnv in simplify'

simplify' :: (?accumInScope :: Bool) => Ex env t -> Ex env t
simplify' = \case
  -- inlining
  ELet _ rhs body
    | not ?accumInScope || not (hasAdds rhs)  -- cannot discard effectful computations
    , Occ lexOcc runOcc <- occCount IZ body
    , lexOcc <= One  -- prevent code size blowup
    , runOcc <= One  -- prevent runtime increase
    -> simplify' (subst1 rhs body)
    | cheapExpr rhs
    -> simplify' (subst1 rhs body)

  -- let splitting
  ELet _ (EPair _ a b) body ->
    simplify' $
      ELet ext a $
      ELet ext (weakenExpr WSink b) $
        subst (\_ t -> \case IZ -> EPair ext (EVar ext (typeOf a) (IS IZ)) (EVar ext (typeOf b) IZ)
                             IS i -> EVar ext t (IS (IS i)))
              body

  -- let rotation
  ELet _ (ELet _ rhs a) b ->
    ELet ext (simplify' rhs) $
    ELet ext (simplify' a) $
      weakenExpr (WCopy WSink) (simplify' b)

  -- beta rules for products
  EFst _ (EPair _ e _) -> simplify' e
  ESnd _ (EPair _ _ e) -> simplify' e

  -- beta rules for coproducts
  ECase _ (EInl _ _ e) rhs _ -> simplify' (ELet ext e rhs)
  ECase _ (EInr _ _ e) _ rhs -> simplify' (ELet ext e rhs)

  -- TODO: array indexing (index of build, index of fold)

  -- TODO: constant folding for operations

  EVar _ t i -> EVar ext t i
  ELet _ a b -> ELet ext (simplify' a) (simplify' b)
  EPair _ a b -> EPair ext (simplify' a) (simplify' b)
  EFst _ e -> EFst ext (simplify' e)
  ESnd _ e -> ESnd ext (simplify' e)
  ENil _ -> ENil ext
  EInl _ t e -> EInl ext t (simplify' e)
  EInr _ t e -> EInr ext t (simplify' e)
  ECase _ e a b -> ECase ext (simplify' e) (simplify' a) (simplify' b)
  EBuild1 _ a b -> EBuild1 ext (simplify' a) (simplify' b)
  EBuild _ es e -> EBuild ext (fmap simplify' es) (simplify' e)
  EFold1 _ a b -> EFold1 ext (simplify' a) (simplify' b)
  EUnit _ e -> EUnit ext (simplify' e)
  EConst _ t v -> EConst ext t v
  EIdx0 _ e -> EIdx0 ext (simplify' e)
  EIdx1 _ a b -> EIdx1 ext (simplify' a) (simplify' b)
  EIdx _ e es -> EIdx ext (simplify' e) (fmap simplify' es)
  EOp _ op e -> EOp ext op (simplify' e)
  EWith e1 e2 -> EWith (simplify' e1) (let ?accumInScope = True in simplify' e2)
  EAccum1 e1 e2 e3 -> EAccum1 (simplify' e1) (simplify' e2) (simplify' e3)
  EError t s -> EError t s

cheapExpr :: Expr x env t -> Bool
cheapExpr = \case
  EVar{} -> True
  ENil{} -> True
  EConst{} -> True
  _ -> False

-- | This can be made more precise by tracking (and not counting) adds on
-- locally eliminated accumulators.
hasAdds :: Expr x env t -> Bool
hasAdds = \case
  EVar _ _ _ -> False
  ELet _ rhs body -> hasAdds rhs || hasAdds body
  EPair _ a b -> hasAdds a || hasAdds b
  EFst _ e -> hasAdds e
  ESnd _ e -> hasAdds e
  ENil _ -> False
  EInl _ _ e -> hasAdds e
  EInr _ _ e -> hasAdds e
  ECase _ e a b -> hasAdds e || hasAdds a || hasAdds b
  EBuild1 _ a b -> hasAdds a || hasAdds b
  EBuild _ es e -> getAny (foldMap (Any . hasAdds) es) || hasAdds e
  EFold1 _ a b -> hasAdds a || hasAdds b
  EUnit _ e -> hasAdds e
  EConst _ _ _ -> False
  EIdx0 _ e -> hasAdds e
  EIdx1 _ a b -> hasAdds a || hasAdds b
  EIdx _ e es -> hasAdds e || getAny (foldMap (Any . hasAdds) es)
  EOp _ _ e -> hasAdds e
  EWith a b -> hasAdds a || hasAdds b
  EAccum1 _ _ _ -> True
  EError _ _ -> False

checkAccumInScope :: SList STy env -> Bool
checkAccumInScope = \case SNil -> False
                          SCons t env -> check t || checkAccumInScope env
  where
    check :: STy t -> Bool
    check STNil = False
    check (STPair s t) = check s || check t
    check (STEither s t) = check s || check t
    check (STArr _ t) = check t
    check (STScal _) = False
    check STAccum{} = True