content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE QuasiQuotes #-}
  {-# LANGUAGE ViewPatterns #-}
  
  -- |
  -- Module      : Jikka.Core.Convert.ConstantFolding
  -- Description : folds constants. / 定数畳み込みをします。
  -- Copyright   : (c) Kimiyuki Onaka, 2021
  -- License     : Apache License 2.0
  -- Maintainer  : kimiyuki95@gmail.com
  -- Stability   : experimental
  -- Portability : portable
  --
  -- \[
  --     \newcommand\int{\mathbf{int}}
  --     \newcommand\bool{\mathbf{bool}}
  --     \newcommand\list{\mathbf{list}}
  -- \]
  module Jikka.Core.Convert.ConstantFolding
    ( run,
  
      -- * internal rules
      rule,
      reduceConstArithmeticExpr,
      reduceConstMaxExpr,
      reduceIdempotentBooleanExpr,
      reduceUnitBooleanExpr,
      reduceConstBooleanExpr,
      reduceUnitBitExpr,
      reduceConstBitExpr,
      reduceConstIntComparison,
      reduceUnitBooleanComparison,
    )
  where
  
  import Data.Bits
  import Data.Either
  import Jikka.Common.Alpha
  import Jikka.Common.Error
  import Jikka.Core.Language.BuiltinPatterns
  import Jikka.Core.Language.Expr
  import Jikka.Core.Language.Lint
  import Jikka.Core.Language.QuasiRules
  import Jikka.Core.Language.RewriteRules
  import Jikka.Core.Language.Runtime
  
  -- |
  -- == List of functions which are reduced
  --
  -- === Basic arithmetical functions
  --
  -- * `Negate` \(: \int \to \int\)
  -- * `Plus` \(: \int \to \int \to \int\)
  -- * `Minus` \(: \int \to \int \to \int\)
  -- * `Mult` \(: \int \to \int \to \int\)
  -- * `FloorDiv` \(: \int \to \int \to \int\)
  -- * `FloorMod` \(: \int \to \int \to \int\)
  -- * `CeilDiv` \(: \int \to \int \to \int\)
  -- * `CeilMod` \(: \int \to \int \to \int\)
  -- * `Pow` \(: \int \to \int \to \int\)
  --
  -- === Advanced arithmetical functions
  --
  -- * `Abs` \(: \int \to \int\)
  -- * `Gcd` \(: \int \to \int \to \int\)
  -- * `Lcm` \(: \int \to \int \to \int\)
  reduceConstArithmeticExpr :: (MonadAlpha m, MonadError Error m) => RewriteRule m
  reduceConstArithmeticExpr =
    let return' = Just . LitInt'
     in simpleRewriteRule "reduceConstArithmeticExpr" $ \case
          Negate' (LitInt' a) -> return' $ - a
          Plus' a (LitInt' 0) -> Just a
          Plus' (LitInt' 0) b -> Just b
          Plus' (LitInt' a) (LitInt' b) -> return' $ a + b
          Minus' a (LitInt' 0) -> Just a
          Minus' (LitInt' 0) b -> Just (Negate' b)
          Minus' (LitInt' a) (LitInt' b) -> return' $ a - b
          Mult' _ (LitInt' 0) -> return' 0
          Mult' a (LitInt' 1) -> Just a
          Mult' (LitInt' 0) _ -> return' 0
          Mult' (LitInt' 1) b -> Just b
          Mult' (LitInt' a) (LitInt' b) -> return' $ a * b
          FloorDiv' a (LitInt' 1) -> Just a
          FloorDiv' (LitInt' a) (LitInt' b) -> Just . fromRight (Bottom' IntTy "division by zero") . (LitInt' <$>) $ floorDiv a b
          FloorMod' _ (LitInt' 1) -> return' 0
          FloorMod' (LitInt' a) (LitInt' b) -> Just . fromRight (Bottom' IntTy "modulo by zero") . (LitInt' <$>) $ floorMod a b
          CeilDiv' a (LitInt' 1) -> Just a
          CeilDiv' (LitInt' a) (LitInt' b) -> Just . fromRight (Bottom' IntTy "division by zero") . (LitInt' <$>) $ ceilDiv a b
          CeilMod' _ (LitInt' 1) -> return' 0
          CeilMod' (LitInt' a) (LitInt' b) -> Just . fromRight (Bottom' IntTy "modulo by zero") . (LitInt' <$>) $ ceilMod a b
          Pow' _ (LitInt' 0) -> return' 1
          Pow' a (LitInt' 1) -> Just a
          Pow' (LitInt' a) (LitInt' b) | b >= 0 && fromInteger b * log (abs (fromInteger a)) < 100 -> return' $ a ^ b
          Abs' (LitInt' a) -> return' $ abs a
          Gcd' a (LitInt' 0) -> Just a
          Gcd' _ (LitInt' 1) -> return' 1
          Gcd' (LitInt' 0) b -> Just b
          Gcd' (LitInt' 1) _ -> return' 1
          Gcd' (LitInt' a) (LitInt' b) -> return' $ gcd a b
          Lcm' _ (LitInt' 0) -> return' 0
          Lcm' a (LitInt' 1) -> Just a
          Lcm' (LitInt' 0) _ -> return' 0
          Lcm' (LitInt' 1) b -> Just b
          Lcm' (LitInt' a) (LitInt' b) -> return' $ lcm a b
          _ -> Nothing
  
  -- |
  -- == List of functions which are reduced
  --
  -- === Max functions
  --
  -- * `Min2` \(: \forall \alpha. \alpha \to \alpha \to \alpha\) (specialized to \(\alpha = \lbrace \bool, \int \rbrace\))
  -- * `Max2` \(: \forall \alpha. \alpha \to \alpha \to \alpha\) (specialized to \(\alpha = \lbrace \bool, \int \rbrace\))
  reduceConstMaxExpr :: (MonadAlpha m, MonadError Error m) => RewriteRule m
  reduceConstMaxExpr = simpleRewriteRule "reduceConstMaxExpr" $ \case
    Min2' _ (LitInt' a) (LitInt' b) -> Just . LitInt' $ min a b
    Min2' _ (LitBool' a) (LitBool' b) -> Just . LitBool' $ min a b
    Max2' _ (LitInt' a) (LitInt' b) -> Just . LitInt' $ max a b
    Max2' _ (LitBool' a) (LitBool' b) -> Just . LitBool' $ max a b
    _ -> Nothing
  
  -- |
  -- == List of functions which are reduced
  --
  -- === Boolean functions
  --
  -- * `And` \(: \bool \to \bool \to \bool\)
  -- * `Or` \(: \bool \to \bool \to \bool\)
  -- * `Implies` \(: \bool \to \bool \to \bool\)
  reduceIdempotentBooleanExpr :: (MonadAlpha m, MonadError Error m) => RewriteRule m
  reduceIdempotentBooleanExpr =
    mconcat
      [ [r| "join/and" forall x. x && x = x|],
        [r| "join/or" forall x. x || x = x|],
        [r| "join/implies" forall x. implies (not x) x = x|],
        [r| "join/implies'" forall x. implies x (not x) = not x|]
      ]
  
  -- |
  -- == List of functions which are reduced
  --
  -- === Boolean functions
  --
  -- * `Not` \(: \bool \to \bool\)
  -- * `And` \(: \bool \to \bool \to \bool\)
  -- * `Or` \(: \bool \to \bool \to \bool\)
  -- * `Implies` \(: \bool \to \bool \to \bool\)
  reduceUnitBooleanExpr :: (MonadAlpha m, MonadError Error m) => RewriteRule m
  reduceUnitBooleanExpr =
    mconcat
      [ [r| "not/true" not true = false|],
        [r| "not/false" not false = false|],
        [r| "and/false" forall x. false && x = false|],
        [r| "and/false'" forall x. x && false = false|],
        [r| "and/true" forall x. true && x = x|],
        [r| "and/true'" forall x. x && true = x|],
        [r| "or/false" forall x. false || x = x|],
        [r| "or/false'" forall x. x || false = x|],
        [r| "or/true" forall x. true || x = true|],
        [r| "or/true'" forall x. x || true = true|],
        [r| "implies/false" forall x. implies false x = true|],
        [r| "implies/false'" forall x. implies x false = not x|],
        [r| "implies/true" forall x. implies true x = x|],
        [r| "implies/true'" forall x. implies x true = true|]
      ]
  
  -- |
  -- == List of functions which are reduced
  --
  -- === Boolean functions
  --
  -- * `If` \(: \forall \alpha. \bool \to \alpha \to \alpha \to \alpha\)
  reduceConstBooleanExpr :: (MonadAlpha m, MonadError Error m) => RewriteRule m
  reduceConstBooleanExpr =
    mconcat
      [ [r| "if/true" forall e1 e2. if true then e1 else e2 = e1|],
        [r| "if/false" forall e1 e2. if false then e1 else e2 = e2|]
      ]
  
  -- |
  -- == List of functions which are reduced
  --
  -- === Bitwise boolean functions
  --
  -- * `BitNot` \(: \int \to \int\)
  -- * `BitAnd` \(: \int \to \int \to \int\)
  -- * `BitOr` \(: \int \to \int \to \int\)
  -- * `BitXor` \(: \int \to \int \to \int\)
  -- * `BitLeftShift` \(: \int \to \int \to \int\)
  -- * `BitRightShift` \(: \int \to \int \to \int\)
  reduceUnitBitExpr :: (MonadAlpha m, MonadError Error m) => RewriteRule m
  reduceUnitBitExpr =
    mconcat
      [ [r| "bitand/0" forall x. 0 & x = 0 |],
        [r| "bitand/0'" forall x. x & 0 = 0 |],
        [r| "bitand/-1" forall x. (-1) & x = x |],
        [r| "bitand/-1'" forall x. x & (-1) = x |],
        [r| "bitor/0" forall x. 0 | x = x |],
        [r| "bitor/0'" forall x. x | 0 = x |],
        [r| "bitor/-1" forall x. (-1) | x = -1 |],
        [r| "bitor/-1'" forall x. x | (-1) = -1 |],
        [r| "bitxor/0" forall x. 0 ^ x = x |],
        [r| "bitxor/0'" forall x. x ^ 0 = x |],
        [r| "bitxor/-1" forall x. (-1) ^ x = ~ x |],
        [r| "bitxor/-1'" forall x. x ^ (-1) = ~ x |],
        [r| "bitleftshift/0" forall x. 0 << x = 0 |],
        [r| "bitleftshift/0'" forall x. x << 0 = x |],
        [r| "bitrightshift/0" forall x. 0 >> x = 0 |],
        [r| "bitrightshift/0'" forall x. x >> 0 = x |]
      ]
  
  -- |
  -- == List of functions which are reduced
  --
  -- === Bitwise boolean functions
  --
  -- * `BitNot` \(: \int \to \int\)
  -- * `BitAnd` \(: \int \to \int \to \int\)
  -- * `BitOr` \(: \int \to \int \to \int\)
  -- * `BitXor` \(: \int \to \int \to \int\)
  -- * `BitLeftShift` \(: \int \to \int \to \int\)
  -- * `BitRightShift` \(: \int \to \int \to \int\)
  reduceConstBitExpr :: (MonadAlpha m, MonadError Error m) => RewriteRule m
  reduceConstBitExpr =
    let return' = Just . LitInt'
     in simpleRewriteRule "reduceConstBitExpr" $ \case
          BitNot' (LitInt' a) -> return' $ complement a
          BitAnd' (LitInt' a) (LitInt' b) -> return' $ a .&. b
          BitOr' (LitInt' a) (LitInt' b) -> return' $ a .|. b
          BitXor' (LitInt' a) (LitInt' b) -> return' $ a `xor` b
          BitLeftShift' (LitInt' a) (LitInt' b) | - 100 < b && b < 100 -> return' $ a `shift` fromInteger b
          BitRightShift' (LitInt' a) (LitInt' b) | - 100 < b && b < 100 -> return' $ a `shift` fromInteger (- b)
          _ -> Nothing
  
  -- |
  -- == List of functions which are reduced
  --
  -- === Comparison functions
  --
  -- * `LessThan` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\))
  -- * `LessEqual` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\))
  -- * `GreaterThan` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\))
  -- * `GreaterEqual` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\))
  -- * `Equal` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\))
  -- * `NotEqual` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\))
  reduceConstIntComparison :: (MonadAlpha m, MonadError Error m) => RewriteRule m
  reduceConstIntComparison =
    simpleRewriteRule "comparison/const/int" $
      (LitBool' <$>) . \case
        LessThan' _ (LitInt' a) (LitInt' b) -> Just $ a < b
        LessEqual' _ (LitInt' a) (LitInt' b) -> Just $ a <= b
        GreaterThan' _ (LitInt' a) (LitInt' b) -> Just $ a > b
        GreaterEqual' _ (LitInt' a) (LitInt' b) -> Just $ a >= b
        Equal' _ (LitInt' a) (LitInt' b) -> Just $ a == b
        NotEqual' _ (LitInt' a) (LitInt' b) -> Just $ a /= b
        _ -> Nothing
  
  -- |
  -- == List of functions which are reduced
  --
  -- === Comparison functions
  --
  -- * `LessThan` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\))
  -- * `LessEqual` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\))
  -- * `GreaterThan` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\))
  -- * `GreaterEqual` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\))
  -- * `Equal` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\))
  -- * `NotEqual` \(: \forall \alpha. \alpha \to \alpha \to \bool\) (specialized to \(\alpha \in \lbrace \bool, \int \rbrace\))
  reduceUnitBooleanComparison :: (MonadAlpha m, MonadError Error m) => RewriteRule m
  reduceUnitBooleanComparison =
    mconcat
      [ -- TODO: implement lessthan and lessequal
        -- NOTE: We can ignore greaterthan and greaterequal because EqualitySolving swaps inequalities.
        [r| "equal/true" forall x. true == x = x |],
        [r| "equal/true'" forall x. x == true = x |],
        [r| "equal/false" forall x. false == x = not x |],
        [r| "equal/false'" forall x. x == false = not x |],
        [r| "notequal/true" forall x. true /= x = not x |],
        [r| "notequal/true'" forall x. x /= true = not x |],
        [r| "notequal/false" forall x. false /= x = x |],
        [r| "notequal/false'" forall x. x /= false = x |]
      ]
  
  rule :: (MonadAlpha m, MonadError Error m) => RewriteRule m
  rule =
    mconcat
      [ reduceConstArithmeticExpr,
        reduceConstMaxExpr,
        reduceIdempotentBooleanExpr,
        reduceUnitBooleanExpr,
        reduceConstBooleanExpr,
        reduceUnitBitExpr,
        reduceConstBitExpr,
        reduceConstIntComparison,
        reduceUnitBooleanComparison
      ]
  
  runProgram :: (MonadAlpha m, MonadError Error m) => Program -> m Program
  runProgram = applyRewriteRuleProgram' rule
  
  -- | `run` folds constants in given programs.
  -- For example, this converts the following:
  --
  -- > 3 x + 2 + 1
  --
  -- to the follwoing:
  --
  -- > 3 x + 3
  run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
  run prog = wrapError' "Jikka.Core.Convert.ConstantFolding" $ do
    precondition $ do
      lint prog
    prog <- runProgram prog
    postcondition $ do
      lint prog
    return prog
  