content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE LambdaCase #-}
  
  -- |
  -- Module      : Jikka.Core.Convert.KubaruToMorau
  -- Description : converts Kubaru DP to Morau DP. / 配る DP を貰う DP に変換します。
  -- Copyright   : (c) Kimiyuki Onaka, 2021
  -- License     : Apache License 2.0
  -- Maintainer  : kimiyuki95@gmail.com
  -- Stability   : experimental
  -- Portability : portable
  module Jikka.Core.Convert.KubaruToMorau
    ( run,
  
      -- * internal rules
      rule,
      runFunctionBody,
    )
  where
  
  import Control.Monad.Trans.Maybe
  import Jikka.Common.Alpha
  import Jikka.Common.Error
  import Jikka.Core.Language.ArithmeticExpr
  import Jikka.Core.Language.Beta
  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
  
  -- | @runFunctionBody c i j step y x k@ returns @step'(y, x, i, k)@ s.t. @step(c, i, j) = step'(c[i + j + 1], c[i], i, i + j + 1)@
  runFunctionBody :: (MonadAlpha m, MonadError Error m) => VarName -> VarName -> VarName -> Expr -> VarName -> VarName -> VarName -> MaybeT m Expr
  runFunctionBody c i j step y x k = do
    step <- lift $ substitute j (Minus' (Minus' (Var k) (Var i)) (LitInt' 1)) step
    let go = \case
          Var x
            | x == c -> hoistMaybe Nothing
            | otherwise -> return $ Var x
          Lit lit -> return $ Lit lit
          At' _ (Var c') index | c' == c -> case () of
            () | parseArithmeticExpr index == parseArithmeticExpr (Var i) -> return $ Var x
            () | parseArithmeticExpr index == parseArithmeticExpr (Var k) -> return $ Var y
            () | otherwise -> hoistMaybe Nothing
          App e1 e2 -> App <$> go e1 <*> go e2
          Let x t e1 e2
            | x == c || x == i || x == j -> throwRuntimeError "name confliction found"
            | otherwise -> Let x t <$> go e1 <*> go e2
          Lam x t e
            | x == c || x == i || x == j -> throwRuntimeError "name confliction found"
            | otherwise -> Lam x t <$> go e
          Assert e1 e2 -> Assert <$> go e1 <*> go e2
    go step
  
  -- | TODO: remove the assumption that the length of @a@ is equals to @n@
  rule :: (MonadAlpha m, MonadError Error m) => RewriteRule m
  rule = makeRewriteRule "Jikka.Core.Convert.KubaruToMorau" $ \_ -> \case
    -- foldl (fun b i -> foldl (fun c j -> setAt c index(i, j) step(c, i, j)) b (range m(i))) a (range n)
    Foldl' IntTy (ListTy t2) (Lam2 b _ i _ (Foldl' IntTy (ListTy t2') (Lam2 c _ j _ (SetAt' _ (Var c') index step)) (Var b') (Range1' m))) a (Range1' n)
      | t2' == t2 && b' == b && c == c' && b `isUnusedVar` m && b `isUnusedVar` index && b `isUnusedVar` step && c `isUnusedVar` index -> runMaybeT $ do
        -- m(i) = n - i - 1
        guard $ parseArithmeticExpr m == parseArithmeticExpr (Minus' (Minus' n (Var i)) (LitInt' 1))
        -- index(i, j) = i + j + 1
        guard $ parseArithmeticExpr index == parseArithmeticExpr (Plus' (Var i) (Plus' (Var j) (LitInt' 1)))
        x <- lift genVarName'
        y <- lift genVarName'
        k <- lift genVarName'
        -- get step'(y, x, i, k) s.t. step(c, i, j) = step'(c[i + j + 1], c[i], i, i + j + 1)
        step <- runFunctionBody c i j step y x k
        step <- lift $ substitute x (At' t2 (Var c) (Var i)) step
        step <- lift $ substitute k (Len' t2 (Var c)) step
        let base = At' t2 a (Len' t2 (Var c))
        return $ Build' t2 (Lam c (ListTy t2) (Foldl' IntTy t2 (Lam2 y t2 i IntTy step) base (Range1' (Len' t2 (Var c))))) (Nil' t2) n
    _ -> return Nothing
  
  runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program
  runProgram = applyRewriteRuleProgram' rule
  
  -- | `run` converts Kubaru DP
  -- (for each \(i\), updates \(
  --     \mathrm{dp}(j) \gets f(\mathrm{dp}(j), \mathrm{dp}(i))
  -- \) for each \(j \gt i\))
  -- to Morau DP
  -- (for each \(i\), computes \(
  --     \mathrm{dp}(i) = F(\lbrace \mathrm{dp}(j) \mid j \lt i \rbrace)
  -- \)).
  --
  -- == Examples
  --
  -- Before:
  --
  -- > foldl (fun dp i ->
  -- >     foldl (fun dp j ->
  -- >         setAt dp j (
  -- >             f dp[j] dp[i])
  -- >         ) dp (range (i + 1) n)
  -- >     ) dp (range n)
  --
  -- After:
  --
  -- > build (fun dp' ->
  -- >     foldl (fun dp_i j ->
  -- >         f dp_i dp'[j]
  -- >         ) dp[i] (range i)
  -- >     ) [] n
  run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
  run prog = wrapError' "Jikka.Core.Convert.KubaruToMorau" $ do
    precondition $ do
      lint prog
    prog <- runProgram prog
    postcondition $ do
      lint prog
    return prog
  