content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE LambdaCase #-}
  
  -- |
  -- Module      : Jikka.Core.Convert.CumulativeSum
  -- Description : processes queries like range sum query using cumulative sums. / 累積和を用いて range sum query のようなクエリを処理します。
  -- Copyright   : (c) Kimiyuki Onaka, 2021
  -- License     : Apache License 2.0
  -- Maintainer  : kimiyuki95@gmail.com
  -- Stability   : experimental
  -- Portability : portable
  module Jikka.Core.Convert.CumulativeSum
    ( run,
  
      -- * internal rules
      rule,
    )
  where
  
  import Data.Maybe
  import Jikka.Common.Alpha
  import Jikka.Common.Error
  import Jikka.Core.Language.ArithmeticExpr
  import Jikka.Core.Language.BuiltinPatterns
  import Jikka.Core.Language.Expr
  import Jikka.Core.Language.FreeVars
  import Jikka.Core.Language.Lint
  import Jikka.Core.Language.RewriteRules
  import Jikka.Core.Language.Util
  
  cumulativeMax :: MonadAlpha m => (Expr -> Expr -> Expr) -> Type -> Maybe Expr -> Expr -> Expr -> m Expr
  cumulativeMax max2 t a0 a n = do
    b <- genVarName'
    let e = At' t (Var b) n
    x1 <- genVarName'
    x2 <- genVarName'
    let a0' = fromMaybe (At' t a (LitInt' 0)) a0
    return $ Let b (ListTy t) (Scanl' t t (Lam2 x1 t x2 t (max2 (Var x1) (Var x2))) a0' a) e
  
  rule :: MonadAlpha m => RewriteRule m
  rule = makeRewriteRule "Jikka.Core.Convert.CumulativeSum" $ \_ -> \case
    Sum' (Map' _ _ (Lam x _ (At' _ a index)) (Range1' n)) | x `isUnusedVar` a -> do
      case makeAffineFunctionFromArithmeticExpr x (parseArithmeticExpr index) of
        Just (coeff, shift) | isOneArithmeticExpr coeff -> do
          b <- genVarName'
          let e =
                if isZeroArithmeticExpr shift
                  then At' IntTy (Var b) n
                  else Minus' (At' IntTy (Var b) (Plus' n (formatArithmeticExpr shift))) (At' IntTy (Var b) (formatArithmeticExpr shift))
          return . Just $
            Let b (ListTy IntTy) (Scanl' IntTy IntTy (Builtin Plus) Lit0 a) e
        _ -> return Nothing
    Max1' t (Cons' _ a0 (Map' _ _ (Lam x _ (At' _ a (Var x'))) (Range1' n))) | x' == x && x `isUnusedVar` a -> do
      Just <$> cumulativeMax (Max2' t) t (Just a0) a n
    Max1' t (Map' _ _ (Lam x _ (At' _ a (Var x'))) (Range1' n)) | x' == x && x `isUnusedVar` a -> do
      Just <$> cumulativeMax (Max2' t) t Nothing a n
    Min1' t (Cons' _ a0 (Map' _ _ (Lam x _ (At' _ a (Var x'))) (Range1' n))) | x' == x && x `isUnusedVar` a -> do
      Just <$> cumulativeMax (Min2' t) t (Just a0) a n
    Min1' t (Map' _ _ (Lam x _ (At' _ a (Var x'))) (Range1' n)) | x' == x && x `isUnusedVar` a -> do
      Just <$> cumulativeMax (Min2' t) t Nothing a n
    Lcm1' t (Cons' _ a0 (Map' _ _ (Lam x _ (At' _ a (Var x'))) (Range1' n))) | x' == x && x `isUnusedVar` a -> do
      Just <$> cumulativeMax Lcm' t (Just a0) a n
    Lcm1' t (Map' _ _ (Lam x _ (At' _ a (Var x'))) (Range1' n)) | x' == x && x `isUnusedVar` a -> do
      Just <$> cumulativeMax Lcm' t Nothing a n
    Gcd1' t (Cons' _ a0 (Map' _ _ (Lam x _ (At' _ a (Var x'))) (Range1' n))) | x' == x && x `isUnusedVar` a -> do
      Just <$> cumulativeMax Gcd' t (Just a0) a n
    Gcd1' t (Map' _ _ (Lam x _ (At' _ a (Var x'))) (Range1' n)) | x' == x && x `isUnusedVar` a -> do
      Just <$> cumulativeMax Gcd' t Nothing a n
    _ -> return Nothing
  
  runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program
  runProgram = applyRewriteRuleProgram' rule
  
  -- | `run` introduces cumulative sums.
  --
  -- == Examples
  --
  -- Before:
  --
  -- > sum (map (fun i -> a[i]) (range n))
  --
  -- After:
  --
  -- > let b = scanl (+) 0 a in b[n]
  run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
  run prog = wrapError' "Jikka.Core.Convert.CumulativeSum" $ do
    precondition $ do
      lint prog
    prog <- runProgram prog
    postcondition $ do
      lint prog
    return prog
  