content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE TupleSections #-}
  
  module Jikka.Core.Language.Util where
  
  import Control.Arrow
  import Control.Monad.Identity
  import Control.Monad.Trans.Maybe
  import Control.Monad.Writer (execWriter, tell)
  import Data.Maybe
  import Data.Monoid (Dual (..))
  import qualified Data.Vector as V
  import Jikka.Common.Alpha
  import Jikka.Common.Error
  import Jikka.Core.Language.BuiltinPatterns
  import Jikka.Core.Language.Expr
  
  genType :: MonadAlpha m => m Type
  genType = VarTy . TypeName Nothing . Just <$> nextCounter
  
  genVarName :: MonadAlpha m => VarName -> m VarName
  genVarName (VarName x _) = VarName x . Just <$> nextCounter
  
  genVarName' :: MonadAlpha m => m VarName
  genVarName' = genVarName (VarName Nothing Nothing)
  
  genVarName'' :: MonadAlpha m => Expr -> m VarName
  genVarName'' = \case
    Var x -> genVarName x
    _ -> genVarName'
  
  mapSubTypesM :: Monad m => (Type -> m Type) -> Type -> m Type
  mapSubTypesM f = go
    where
      go = \case
        VarTy x -> f $ VarTy x
        IntTy -> f IntTy
        BoolTy -> f BoolTy
        ListTy t -> f . ListTy =<< f t
        TupleTy ts -> f . TupleTy =<< mapM f ts
        FunTy t1 t2 -> f =<< (FunTy <$> f t1 <*> f t2)
        DataStructureTy ds -> f $ DataStructureTy ds
  
  mapTypeLiteralM :: Monad m => (Type -> m Type) -> Literal -> m Literal
  mapTypeLiteralM f = \case
    LitBuiltin builtin ts -> LitBuiltin builtin <$> mapM f ts
    LitInt n -> return $ LitInt n
    LitBool p -> return $ LitBool p
    LitNil t -> LitNil <$> f t
    LitBottom t err -> LitBottom <$> f t <*> pure err
  
  mapTypeExprM :: Monad m => (Type -> m Type) -> Expr -> m Expr
  mapTypeExprM f = go
    where
      go = \case
        Var x -> return $ Var x
        Lit lit -> Lit <$> mapTypeLiteralM f lit
        App f e -> App <$> go f <*> go e
        Lam x t body -> Lam x <$> f t <*> go body
        Let x t e1 e2 -> Let x <$> f t <*> go e1 <*> go e2
        Assert e1 e2 -> Assert <$> go e1 <*> go e2
  
  mapTypeExpr :: (Type -> Type) -> Expr -> Expr
  mapTypeExpr f e = runIdentity (mapTypeExprM (return . f) e)
  
  mapTypeToplevelExprM :: Monad m => (Type -> m Type) -> ToplevelExpr -> m ToplevelExpr
  mapTypeToplevelExprM f = \case
    ResultExpr e -> ResultExpr <$> mapTypeExprM f e
    ToplevelLet x t e cont -> ToplevelLet x <$> f t <*> mapTypeExprM f e <*> mapTypeToplevelExprM f cont
    ToplevelLetRec g args ret body cont -> ToplevelLetRec g <$> mapM (\(x, t) -> (x,) <$> f t) args <*> f ret <*> mapTypeExprM f body <*> mapTypeToplevelExprM f cont
    ToplevelAssert e cont -> ToplevelAssert <$> mapTypeExprM f e <*> mapTypeToplevelExprM f cont
  
  mapTypeProgramM :: Monad m => (Type -> m Type) -> Program -> m Program
  mapTypeProgramM = mapTypeToplevelExprM
  
  mapTypeProgram :: (Type -> Type) -> Program -> Program
  mapTypeProgram f prog = runIdentity (mapTypeProgramM (return . f) prog)
  
  -- | `mapSubExprM'` substitutes exprs using given two functions, which are called in pre-order and post-order.
  mapSubExprM' :: Monad m => ([(VarName, Type)] -> Expr -> m Expr) -> ([(VarName, Type)] -> Expr -> m Expr) -> [(VarName, Type)] -> Expr -> m Expr
  mapSubExprM' pre post env e = do
    e <- pre env e
    let go = mapSubExprM' pre post
    e <- case e of
      Var y -> return $ Var y
      Lit lit -> return $ Lit lit
      App g e -> App <$> go env g <*> go env e
      Lam x t body -> Lam x t <$> go ((x, t) : env) body
      Let y t e1 e2 -> Let y t <$> go env e1 <*> go ((y, t) : env) e2
      Assert e1 e2 -> Assert <$> go env e1 <*> go env e2
    post env e
  
  mapToplevelExprM' :: Monad m => ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr) -> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr) -> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
  mapToplevelExprM' pre post env e = do
    e <- pre env e
    e <- case e of
      ResultExpr e -> return $ ResultExpr e
      ToplevelLet y t e cont ->
        ToplevelLet y t e <$> mapToplevelExprM' pre post ((y, t) : env) cont
      ToplevelLetRec g args ret body cont ->
        let env' = (g, foldr (FunTy . snd) ret args) : env
         in ToplevelLetRec g args ret body <$> mapToplevelExprM' pre post env' cont
      ToplevelAssert e cont ->
        ToplevelAssert e <$> mapToplevelExprM' pre post env cont
    post env e
  
  mapExprToplevelExprM :: Monad m => ([(VarName, Type)] -> Expr -> m Expr) -> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
  mapExprToplevelExprM f env = mapToplevelExprM' pre' (\_ e -> return e) env
    where
      pre' env = \case
        ResultExpr e -> ResultExpr <$> f env e
        ToplevelLet y t e cont -> ToplevelLet y t <$> f env e <*> pure cont
        ToplevelLetRec g args ret body cont ->
          let env' = (g, foldr (FunTy . snd) ret args) : env
           in ToplevelLetRec g args ret <$> f (reverse args ++ env') body <*> pure cont
        ToplevelAssert e cont -> ToplevelAssert <$> f env e <*> pure cont
  
  mapExprProgramM :: Monad m => ([(VarName, Type)] -> Expr -> m Expr) -> Program -> m Program
  mapExprProgramM f = mapExprToplevelExprM f []
  
  -- | `mapSubExprM` is a wrapper of `mapSubExprM'`. This function works in post-order.
  mapSubExprM :: Monad m => ([(VarName, Type)] -> Expr -> m Expr) -> [(VarName, Type)] -> Expr -> m Expr
  mapSubExprM f = mapSubExprM' (\_ e -> return e) f
  
  mapSubExpr :: ([(VarName, Type)] -> Expr -> Expr) -> [(VarName, Type)] -> Expr -> Expr
  mapSubExpr f env e = runIdentity $ mapSubExprM (\env e -> return $ f env e) env e
  
  mapExprToplevelExpr :: ([(VarName, Type)] -> Expr -> Expr) -> [(VarName, Type)] -> ToplevelExpr -> ToplevelExpr
  mapExprToplevelExpr f env e = runIdentity $ mapExprToplevelExprM (\env e -> return $ f env e) env e
  
  -- | @mapExprProgram f prog@ applies @f@ to each root exprs in @prog@.
  -- This doesn't run into sub-exprs. For example, @toplevel-let x = (e1 + e2) in ...@ becomes @toplevel-let x = (f (e1 + e2)) in ...@, instead of @toplevel-let x = (f (f e1 + f e2)) in ...@
  mapExprProgram :: ([(VarName, Type)] -> Expr -> Expr) -> Program -> Program
  mapExprProgram f prog = runIdentity $ mapExprProgramM (\env e -> return $ f env e) prog
  
  -- | `mapToplevelExprM` is a wrapper of `mapToplevelExprM'`. This function works in post-order.
  mapToplevelExprM :: Monad m => ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr) -> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
  mapToplevelExprM f env e = mapToplevelExprM' (\_ e -> return e) f env e
  
  mapToplevelExprProgramM :: Monad m => ([(VarName, Type)] -> Program -> m Program) -> Program -> m Program
  mapToplevelExprProgramM f prog = mapToplevelExprM f [] prog
  
  mapToplevelExprProgram :: ([(VarName, Type)] -> Program -> Program) -> Program -> Program
  mapToplevelExprProgram f prog = runIdentity $ mapToplevelExprProgramM (\env e -> return $ f env e) prog
  
  listSubExprs :: Expr -> [Expr]
  listSubExprs e = getDual . execWriter $ mapSubExprM go [] e
    where
      go _ e = do
        tell $ Dual [e]
        return e
  
  uncurryFunTy :: Type -> ([Type], Type)
  uncurryFunTy = \case
    (FunTy t t') -> let (ts, ret) = uncurryFunTy t' in (t : ts, ret)
    ret -> ([], ret)
  
  uncurryLam :: Expr -> ([(VarName, Type)], Expr)
  uncurryLam = \case
    Lam x t body -> let (args, body') = uncurryLam body in ((x, t) : args, body')
    body -> ([], body)
  
  curryApp :: Expr -> (Expr, [Expr])
  curryApp = \case
    App f e -> let (f', e') = curryApp f in (f', e' ++ [e])
    f -> (f, [])
  
  curryFunTy :: [Type] -> Type -> Type
  curryFunTy ts ret = foldr FunTy ret ts
  
  curryLam :: [(VarName, Type)] -> Expr -> Expr
  curryLam args body = foldr (uncurry Lam) body args
  
  uncurryApp :: Expr -> [Expr] -> Expr
  uncurryApp = foldl App
  
  isVectorTy :: Type -> Bool
  isVectorTy = isJust . sizeOfVectorTy
  
  isVectorTy' :: [Type] -> Bool
  isVectorTy' = isVectorTy . TupleTy
  
  sizeOfVectorTy :: Type -> Maybe Int
  sizeOfVectorTy = \case
    TupleTy ts | all (== IntTy) ts -> Just (length ts)
    _ -> Nothing
  
  isMatrixTy :: Type -> Bool
  isMatrixTy = isJust . sizeOfMatrixTy
  
  isMatrixTy' :: [Type] -> Bool
  isMatrixTy' = isMatrixTy . TupleTy
  
  sizeOfMatrixTy :: Type -> Maybe (Int, Int)
  sizeOfMatrixTy = \case
    TupleTy ts@(TupleTy ts' : _) | all (== IntTy) ts' && all (== TupleTy ts') ts -> Just (length ts, length ts')
    _ -> Nothing
  
  isConstantTimeBuiltin :: Builtin -> Bool
  isConstantTimeBuiltin = \case
    -- arithmetical functions
    Negate -> True
    Plus -> True
    Minus -> True
    Mult -> True
    FloorDiv -> True
    FloorMod -> True
    CeilDiv -> True
    CeilMod -> True
    JustDiv -> True
    Pow -> True
    -- advanced arithmetical functions
    Abs -> True
    Gcd -> True
    Lcm -> True
    Min2 -> True
    Max2 -> True
    Iterate -> False
    -- logical functions
    Not -> True
    And -> True
    Or -> True
    Implies -> True
    If -> True
    -- bitwise functions
    BitNot -> True
    BitAnd -> True
    BitOr -> True
    BitXor -> True
    BitLeftShift -> True
    BitRightShift -> True
    -- matrix functions
    MatAp _ _ -> True
    MatZero _ _ -> True
    MatOne _ -> True
    MatAdd _ _ -> True
    MatMul _ _ _ -> True
    MatPow _ -> True
    VecFloorMod _ -> True
    MatFloorMod _ _ -> True
    -- modular functions
    ModNegate -> True
    ModPlus -> True
    ModMinus -> True
    ModMult -> True
    ModInv -> True
    ModPow -> True
    ModMatAp _ _ -> True
    ModMatAdd _ _ -> True
    ModMatMul _ _ _ -> True
    ModMatPow _ -> True
    -- list functions
    Cons -> False
    Snoc -> False
    Foldl -> False
    Scanl -> False
    Build -> False
    Len -> True
    Map -> False
    Filter -> False
    At -> True
    SetAt -> False
    Elem -> False
    Sum -> False
    Product -> False
    ModSum -> False
    ModProduct -> False
    Min1 -> False
    Max1 -> False
    ArgMin -> False
    ArgMax -> False
    Gcd1 -> False
    Lcm1 -> False
    All -> False
    Any -> False
    Sorted -> False
    Reversed -> False
    Range1 -> False
    Range2 -> False
    Range3 -> False
    -- tuple functions
    Tuple -> True
    Proj _ -> True
    -- comparison
    LessThan -> True
    LessEqual -> True
    GreaterThan -> True
    GreaterEqual -> True
    Equal -> True
    NotEqual -> True
    -- combinational functions
    Fact -> True
    Choose -> True
    Permute -> True
    MultiChoose -> True
    -- data structures
    ConvexHullTrickInit -> False
    ConvexHullTrickInsert -> False
    ConvexHullTrickGetMin -> False
    SegmentTreeInitList _ -> False
    SegmentTreeGetRange _ -> False
    SegmentTreeSetPoint _ -> False
  
  isLiteral :: Expr -> Bool
  isLiteral = \case
    Lit _ -> True
    _ -> False
  
  -- | `isConstantTimeExpr` checks whether given exprs are suitable to propagate.
  isConstantTimeExpr :: Expr -> Bool
  isConstantTimeExpr = \case
    Var _ -> True
    Lit _ -> True
    e@(App _ _) -> case curryApp e of
      (Lit (LitBuiltin f _), args) -> isConstantTimeBuiltin f && all isConstantTimeExpr args
      _ -> False
    Lam _ _ _ -> True
    Let _ _ e1 e2 -> isConstantTimeExpr e1 && isConstantTimeExpr e2
    Assert e1 e2 -> isConstantTimeExpr e1 && isConstantTimeExpr e2
  
  -- | `replaceLenF` replaces @len(f)@ in an expr with @i + k@.
  -- * This assumes that there are no name conflicts.
  replaceLenF :: MonadError Error m => VarName -> VarName -> Integer -> Expr -> m Expr
  replaceLenF f i k = go
    where
      go = \case
        Len' _ (Var f') | f' == f -> return $ Plus' (Var i) (LitInt' k)
        Var y -> return $ Var y
        Lit lit -> return $ Lit lit
        App g e -> App <$> go g <*> go e
        Lam x _ _ | x == i -> throwInternalError "Jikka.Core.Language.Util.replaceLenF: name conflict"
        Lam x t body -> Lam x t <$> (if x == f then return body else go body)
        Let y _ _ _ | y == i -> throwInternalError "Jikka.Core.Language.Util.replaceLenF: name conflict"
        Let y t e1 e2 -> Let y t <$> go e1 <*> (if y == f then return e2 else go e2)
        Assert e1 e2 -> Assert <$> go e1 <*> go e2
  
  -- | `getRecurrenceFormulaBase` makes a pair @((a_0, ..., a_{k - 1}), a)@ from @setat (... (setat a 0 a_0) ...) (k - 1) a_{k - 1})@.
  getRecurrenceFormulaBase :: Expr -> ([Expr], Expr)
  getRecurrenceFormulaBase = go (V.replicate recurrenceLimit Nothing)
    where
      recurrenceLimit :: Num a => a
      recurrenceLimit = 20
      go :: V.Vector (Maybe (Expr, Type)) -> Expr -> ([Expr], Expr)
      go base = \case
        SetAt' t e (LitInt' i) e'
          | 0 <= i && i < recurrenceLimit -> go (base V.// [(fromInteger i, Just (e', t))]) e
          | otherwise -> second (\e -> SetAt' t e (LitInt' i) e') $ go base e
        e ->
          let (base', base'') = span isJust (V.toList base)
              base''' = map (fst . fromJust) base'
              e'' = foldr (\(i, e') e -> maybe id (\(e', t) e -> SetAt' t e (LitInt' i) e') e' e) e (zip [toInteger (length base') ..] base'')
           in (base''', e'')
  
  hoistMaybe :: Applicative m => Maybe a -> MaybeT m a
  hoistMaybe = MaybeT . pure
  