{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
module Simplify where

import Data.Monoid (Sum(..))

import AST
import ASTfunc
import MonMap (MonMap)
import qualified MonMap


simplify :: Expr env t -> Expr env t
simplify e = unSExpr (slr' . sli' . slr' $ toSExpr e)

slr' :: SExpr env t -> SExpr env t
slr' e = slr weakenId e

-- simplify using let-rotate, and work away Unused bindings
slr :: env :> env' -> SExpr env t -> SExpr env' t
slr w (Auto expr) = case expr of
    -- unused binding
    Let lhs Unused e -> slr (weakenSkip lhs w) e

    -- let-rotate
    Let lhs (Auto (Let lhs2 rhs2 e2)) e
      | Exists lhs' <- rebuildLHS lhs
      -> let sh = sinkWithLHS lhs lhs' (weakenWithLHS lhs2)
         in slr w (In (Let lhs2 rhs2 (In (Let lhs' e2 (applyShift sh e)))))

    -- let-split
    Let (L2 l1 l2) (Auto (Pair e1 e2)) e ->
        slr w (In (Let l1 e1 (In (Let l2 (applyShift (weakenWithLHS l1) e2) e))))

    -- redundant wildcard binding
    Let (L0 _) _ e -> slr w e

    Const t x -> In (Const t x)
    Pair e1 e2 -> In (Pair (slr w e1) (slr w e2))
    Nil -> In Nil
    Prim op e -> In (Prim op (slr w e))
    Var (V t i) -> In (Var (V t (w >:> i)))
    Let lhs rhs e
      | Exists lhs' <- rebuildLHS lhs
      -> In (Let lhs' (slr w rhs) (slr (sinkWithLHS lhs lhs' w) e))
slr _ Unused = error "slr: unused found outside Let RHS"

-- simplify using linear inlining
sli' :: SExpr env t -> SExpr env t
sli' expr =
    let (countmp, expr') = count 0 expr
        countmp' = ((<= 1) . getSum) <$> countmp
    in inline countmp' 0 (Inliner (In . Var)) expr'

count :: Int -> SExpr env t -> (MonMap Int (Sum Int), SExpr env t)
count d (Auto expr) = In <$> case expr of
    Const t x -> return (Const t x)
    Pair e1 e2 -> Pair <$> count d e1 <*> count d e2
    Nil -> return Nil
    Prim op e -> Prim op <$> count d e
    Var (V t i) -> (MonMap.singleton (d - 1 - idxToInt i) 1, Var (V t i))
    Let lhs rhs e -> Let lhs <$> count d rhs <*> count (d + lhsSize lhs) e
count _ Unused = return Unused

lhsSize :: LHS s t env env' -> Int
lhsSize (L0 _) = 0
lhsSize (L1 _) = 0
lhsSize (L2 l1 l2) = lhsSize l1 + lhsSize l2

data Inliner env = Inliner { unInliner :: forall t. V STy env t -> SExpr env t }

inline :: MonMap Int Bool -> Int -> Inliner env -> SExpr env t -> SExpr env t
inline countmp = go
  where
    go :: Int -> Inliner env -> SExpr env t -> SExpr env t
    go d il (Auto expr) = case expr of
        Const t x -> In (Const t x)
        Pair e1 e2 -> In (Pair (go d il e1) (go d il e2))
        Nil -> In Nil
        Prim op e -> In (Prim op (go d il e))
        Var v -> unInliner il v
        Let (L1 t) rhs e
          | Just True <- MonMap.lookup d countmp ->
              let il' = Inliner
                      (\case V _ Z -> applyShift weakenSucc rhs
                             V t' (S i) -> applyShift weakenSucc (unInliner il (V t' i)))
              in In (Let (L1 t) Unused (go (d + 1) il' e))
        Let lhs rhs e ->
            In (Let lhs (go d il rhs)
                    (go (d + lhsSize lhs) (weakenILwithLHS lhs il) e))
    go _ _ Unused = Unused

weakenILwithLHS :: LHS STy t env env' -> Inliner env -> Inliner env'
weakenILwithLHS (L0 _) il = il
weakenILwithLHS (L1 _) il =
    Inliner (\case v@(V _ Z) -> In (Var v)
                   V t (S i) -> applyShift weakenSucc (unInliner il (V t i)))
weakenILwithLHS (L2 l1 l2) il = weakenILwithLHS l2 (weakenILwithLHS l1 il)


data SExpr env t where
    Shift :: env :> env' -> PExpr SExpr env t -> SExpr env' t
    In :: PExpr SExpr env t -> SExpr env t
    Unused :: SExpr env t

toSExpr :: Expr env t -> SExpr env t
toSExpr (Expr e) = In (eInto toSExpr e)

unSExpr :: SExpr env t -> Expr env t
unSExpr (In e) = Expr (eInto unSExpr e)
unSExpr (Shift w e) = Expr (shiftTraverse (eInto unSExpr) w e)
unSExpr Unused = error "unSExpr: unused found"

applyShift :: env :> env' -> SExpr env t -> SExpr env' t
applyShift w (In e) = Shift w e
applyShift w (Shift w' e) = Shift (w >. w') e
applyShift _ Unused = Unused

pattern Auto :: PExpr SExpr env t -> SExpr env t
pattern Auto se <- (matchAuto -> Just se)
{-# COMPLETE Auto, Unused #-}

matchAuto :: SExpr env t -> Maybe (PExpr SExpr env t)
matchAuto (In e) = Just e
matchAuto (Shift w e) = Just (shiftTraverse id w e)
matchAuto Unused = Nothing

shiftTraverse :: (forall env' t'. PExpr SExpr env' t' -> PExpr expr env' t')
              -> env :> env1 -> PExpr SExpr env t -> PExpr expr env1 t
shiftTraverse f w e = case e of
    Const t x -> Const t x
    Pair e1 e2 -> f (Pair (applyShift w e1) (applyShift w e2))
    Nil -> Nil
    Prim op e1 -> f (Prim op (applyShift w e1))
    Var v -> Var (weakenVar w v)
    Let lhs rhs e1
      | Exists lhs' <- rebuildLHS lhs
      -> f (Let lhs' (applyShift w rhs)
                     (applyShift (sinkWithLHS lhs lhs' w) e1))