content: {-# LANGUAGE DeriveDataTypeable #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE PatternSynonyms #-}
  {-# LANGUAGE ViewPatterns #-}
  
  -- |
  -- Module      : Jikka.Core.Language.Expr
  -- Description : has data types of our core language. / core 言語のためのデータ型を持ちます。
  -- Copyright   : (c) Kimiyuki Onaka, 2020
  -- License     : Apache License 2.0
  -- Maintainer  : kimiyuki95@gmail.com
  -- Stability   : experimental
  -- Portability : portable
  --
  -- `Jikka.Core.Language.Expr` module has the basic data types for our core language.
  -- They are similar to the GHC Core language.
  module Jikka.Core.Language.Expr where
  
  import Data.Data
  import Data.String
  import Jikka.Common.Name
  
  -- NOTE: @VarName occ flavour == VarName occ' flavour'@ is defined as @occ == occ' && flavour == flavour'@ but we assume that this is equivalent to @flavour == flavour'@ unless they are @Nothing@.
  data VarName = VarName OccName NameFlavour deriving (Eq, Ord, Show, Read, Data, Typeable)
  
  instance IsString VarName where
    fromString = uncurry VarName . toFlavouredName
  
  formatVarName :: VarName -> String
  formatVarName (VarName occ flavour) = formatFlavouredName occ flavour
  
  data TypeName = TypeName OccName NameFlavour deriving (Eq, Ord, Show, Read, Data, Typeable)
  
  instance IsString TypeName where
    fromString = uncurry TypeName . toFlavouredName
  
  formatTypeName :: TypeName -> String
  formatTypeName (TypeName occ flavour) = formatFlavouredName occ flavour
  
  -- | `Type` represents the types of our core language. This is similar to the `Type` of GHC Core.
  -- See also [commentary/compiler/type-type](https://gitlab.haskell.org/ghc/ghc/-/wikis/commentary/compiler/type-type).
  --
  -- \[
  --     \newcommand\int{\mathbf{int}}
  --     \newcommand\bool{\mathbf{bool}}
  --     \newcommand\list{\mathbf{list}}
  --     \begin{array}{rl}
  --         \tau ::= & \alpha \\
  --         \vert & \int \\
  --         \vert & \bool \\
  --         \vert & \list(\tau) \\
  --         \vert & \tau \times \tau \times \dots \times \tau \\
  --         \vert & \tau \to \tau \\
  --         \vert & \mathrm{data\_structure}
  --     \end{array}
  -- \]
  data Type
    = VarTy TypeName
    | IntTy
    | BoolTy
    | ListTy Type
    | TupleTy [Type]
    | FunTy Type Type
    | DataStructureTy DataStructure
    deriving (Eq, Ord, Show, Read, Data, Typeable)
  
  data DataStructure
    = ConvexHullTrick
    | SegmentTree Semigroup'
    deriving (Eq, Ord, Show, Read, Data, Typeable)
  
  data Semigroup'
    = SemigroupIntPlus
    | SemigroupIntMin
    | SemigroupIntMax
    | SemigroupIntGcd
    | SemigroupIntLcm
    deriving (Eq, Ord, Show, Read, Data, Typeable)
  
  -- | TODO: What is the difference between `Literal` and `Builtin`?
  data Builtin
    = -- arithmetical functions
  
      -- | \(: \int \to \int\)
      Negate
    | -- | \(: \int \to \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
    | -- | divison which has no remainder \(: \int \to \int \to \int\)
      JustDiv
    | -- | \(: \int \to \int \to \int\)
      Pow
    | -- advanced arithmetical functions
  
      -- | \(: \int \to \int\)
      Abs
    | -- | \(: \int \to \int \to \int\)
      Gcd
    | -- | \(: \int \to \int \to \int\)
      Lcm
    | -- | \(: \forall \alpha. \alpha \to \alpha \to \alpha\)
      Min2
    | -- | \(: \forall \alpha. \alpha \to \alpha \to \alpha\)
      Max2
    | -- | iterated application \((\lambda k f x. f^k(x)): \forall \alpha. \int \to (\alpha \to \alpha) \to \alpha \to \alpha\)
      Iterate
    | -- logical functions
  
      -- | \(: \bool \to \bool\)
      Not
    | -- | \(: \bool \to \bool \to \bool\)
      And
    | -- | \(: \bool \to \bool \to \bool\)
      Or
    | -- | \(: \bool \to \bool \to \bool\)
      Implies
    | -- | \(: \forall \alpha. \bool \to \alpha \to \alpha \to \alpha\)
      If
    | -- bitwise functions
  
      -- | \(: \int \to \int\)
      BitNot
    | -- | \(: \int \to \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
    | -- matrix functions
  
      -- | matrix application \(: \int^{H \times W} \to \int^W \to \int^H\)
      MatAp Integer Integer
    | -- | zero matrix \(: \to \int^{h \times w}\)
      MatZero Integer Integer
    | -- | unit matrix \(: \to \int^{n \times n}\)
      MatOne Integer
    | -- | matrix addition \(: \int^{H \times W} \to \int^{H \times W} \to \int^{H \times W}\)
      MatAdd Integer Integer
    | -- | matrix multiplication \(: \int^{H \times n} \to \int^{n \times W} \to \int^{H \times W}\)
      MatMul Integer Integer Integer
    | -- | matrix power \(: \int^{n \times n} \to \int \to \int^{n \times n}\)
      MatPow Integer
    | -- | vector point-wise floor-mod \(: \int^{n} \to \int \to \int^{n}\)
      VecFloorMod Integer
    | -- | matrix point-wise floor-mod \(: \int^{H \times W} \to \int \to \int^{H \times W}\)
      MatFloorMod Integer Integer
    | -- modular functions
  
      -- | \(: \int \to \int \to \int\)
      ModNegate
    | -- | \(: \int \to \int \to \int \to \int\)
      ModPlus
    | -- | \(: \int \to \int \to \int \to \int\)
      ModMinus
    | -- | \(: \int \to \int \to \int \to \int\)
      ModMult
    | -- | \(: \int \to \int \to \int\)
      ModInv
    | -- | \(: \int \to \int \to \int \to \int\)
      ModPow
    | -- | matrix application \(: \int^{H \times W} \to \int^W \to \int \to \int^H\)
      ModMatAp Integer Integer
    | -- | matrix addition \(: \int^{H \times W} \to \int^{H \times W} \to \int \to \int^{H \times W}\)
      ModMatAdd Integer Integer
    | -- | matrix multiplication \(: \int^{H \times n} \to \int^{n \times W} \to \int \to \int^{H \times W}\)
      ModMatMul Integer Integer Integer
    | -- | matrix power \(: \int^{n \times n} \to \int \to \int^{n \times n}\)
      ModMatPow Integer
    | -- list functions
  
      -- | \(: \forall \alpha. \alpha \to \list(\alpha) \to \list(\alpha)\)
      Cons
    | -- | \(: \forall \alpha. \list(alpha) \to \alpha \to \list(\alpha)\)
      Snoc
    | -- | \(: \forall \alpha \beta. (\beta \to \alpha \to \beta) \to \beta \to \list(\alpha) \to \beta\)
      Foldl
    | -- | \(: \forall \alpha \beta. (\beta \to \alpha \to \beta) \to \beta \to \list(\alpha) \to \list(\beta)\)
      Scanl
    | -- | \(\lambda f a n.\) repeat @a <- snoc a (f a)@ @n@ times \(: \forall \alpha. (\list(\alpha) \to \alpha) \to \list(\alpha) \to \int \to \list(\alpha)\)
      Build
    | -- | \(: \forall \alpha. \list(\alpha) \to \int\)
      Len
    | -- | \(: \forall \alpha \beta. (\alpha \to \beta) \to \list(\alpha) \to \list(\beta)\)
      Map
    | -- | \(: \forall \alpha \beta. (\alpha \to \bool) \to \list(\alpha) \to \list(\beta)\)
      Filter
    | -- | \(: \forall \alpha. \list(\alpha) \to \int \to \alpha\)
      At
    | -- | \(: \forall \alpha. \list(\alpha) \to \int \to \alpha \to \list(\alpha)\)
      SetAt
    | -- | \(: \forall \alpha. \alpha \to \list(\alpha) \to \bool\)
      Elem
    | -- | \(: \list(\int) \to \int\)
      Sum
    | -- | \(: \list(\int) \to \int\)
      Product
    | -- | \(: \list(\int) \to \int \to \int\)
      ModSum
    | -- | \(: \list(\int) \to \int \to \int\)
      ModProduct
    | -- | \(: \forall \alpha. \list(\alpha) \to \alpha\)
      Min1
    | -- | \(: \forall \alpha. \list(\alpha) \to \alpha\)
      Max1
    | -- | \(: \forall \alpha. \list(\alpha) \to \int\)
      ArgMin
    | -- | \(: \forall \alpha. \list(\alpha) \to \int\)
      ArgMax
    | -- | \(: \list(\int) \to \int\)
      Gcd1
    | -- | \(: \list(\int) \to \int\)
      Lcm1
    | -- | \(: \list(\bool) \to \bool\)
      All
    | -- | \(: \list(\bool) \to \bool\)
      Any
    | -- | \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\)
      Sorted
    | -- | \(: \forall \alpha. \list(\alpha) \to \list(\alpha)\)
      Reversed
    | -- | \(: \int \to \list(\int)\)
      Range1
    | -- | \(: \int \to \int \to \list(\int)\)
      Range2
    | -- | \(: \int \to \int \to \int \to \list(\int)\)
      Range3
    | -- tuple functions
  
      -- | \(: \forall \alpha_0 \alpha_1 \dots \alpha _ {n - 1}. \alpha_0 \to \dots \to \alpha _ {n - 1} \to \alpha_0 \times \dots \times \alpha _ {n - 1}\)
      Tuple
    | -- | \(: \forall \alpha_0 \alpha_1 \dots \alpha _ {n - 1}. \alpha_0 \times \dots \times \alpha _ {n - 1} \to \alpha_i\)
      --
      -- `Jikka.Core.Parse` may make broken `Proj` with its list of type arguments is empty. This is fixed by `Jikka.Core.Convert.TypeInfer` module.
      Proj Integer
    | -- comparison
  
      -- | \(: \forall \alpha. \alpha \to \alpha \to \bool\)
      LessThan
    | -- | \(: \forall \alpha. \alpha \to \alpha \to \bool\)
      LessEqual
    | -- | \(: \forall \alpha. \alpha \to \alpha \to \bool\)
      GreaterThan
    | -- | \(: \forall \alpha. \alpha \to \alpha \to \bool\)
      GreaterEqual
    | -- | \(: \forall \alpha. \alpha \to \alpha \to \bool\)
      Equal
    | -- | \(: \forall \alpha. \alpha \to \alpha \to \bool\)
      NotEqual
    | -- combinational functions
  
      -- | \(: \int \to \int\)
      Fact
    | -- | \(: \int \to \int \to \int\)
      Choose
    | -- | \(: \int \to \int \to \int\)
      Permute
    | -- | \(: \int \to \int \to \int\)
      MultiChoose
    | -- data structures
  
      -- | \(: \mathrm{convex\_hull\_trick}\)
      ConvexHullTrickInit
    | -- | \(: \mathrm{convex\_hull\_trick} \to \int \to \int\)
      ConvexHullTrickGetMin
    | -- | \(: \mathrm{convex\_hull\_trick} \to \int \to \int \to \mathrm{convex\_hull\_trick}\)
      ConvexHullTrickInsert
    | -- | \(: \list(S) \to \mathrm{segment\_tree}(S)\) for a semigroup \(S)
      SegmentTreeInitList Semigroup'
    | -- | \(: \mathrm{segment\_tree}(S) \to \int \to \int \to S\) for a semigroup \(S)
      SegmentTreeGetRange Semigroup'
    | -- | \(: \mathrm{segment\_tree}(S) \to \int \to S \to \mathrm{segment\_tree}(S)\) for a semigroup \(S)
      SegmentTreeSetPoint Semigroup'
    deriving (Eq, Ord, Show, Read, Data, Typeable)
  
  data Literal
    = LitBuiltin Builtin [Type]
    | -- | \(: \forall \alpha. \int\)
      LitInt Integer
    | -- | \(: \forall \alpha. \bool\)
      LitBool Bool
    | -- | \(: \forall \alpha. \list(\alpha)\)
      LitNil Type
    | -- | \(: \bot : \forall \alpha. \alpha\). The second argument is its error message.
      LitBottom Type String
    deriving (Eq, Ord, Show, Read, Data, Typeable)
  
  -- | `Expr` represents the exprs of our core language. This is similar to the `Expr` of GHC Core.
  -- See also [commentary/compiler/core-syn-type](https://gitlab.haskell.org/ghc/ghc/-/wikis/commentary/compiler/core-syn-type).
  --
  -- \[
  --     \begin{array}{rl}
  --         e ::= & x \\
  --         \vert & \mathrm{literal}\ldots \\
  --         \vert & e_0(e_1, e_2, \dots, e_n) \\
  --         \vert & \lambda ~ x_0\colon \tau_0, x_1\colon \tau_1, \dots, x_{n-1}\colon \tau_{n-1}. ~ e \\
  --         \vert & \mathbf{let} ~ x\colon \tau = e_1 ~ \mathbf{in} ~ e_2 \\
  --         \vert & \mathbf{assert} ~ e_1 ~ \mathbf{in} ~ e_2
  --     \end{array}
  -- \]
  data Expr
    = Var VarName
    | Lit Literal
    | App Expr Expr
    | Lam VarName Type Expr
    | -- | This "let" is not recursive.
      Let VarName Type Expr Expr
    | Assert Expr Expr
    deriving (Eq, Ord, Show, Read, Data, Typeable)
  
  pattern Fun2Ty t1 t2 ret = FunTy t1 (FunTy t2 ret)
  
  pattern Fun3Ty t1 t2 t3 ret = FunTy t1 (FunTy t2 (FunTy t3 ret))
  
  pattern Fun1STy t <-
    (\case FunTy t1 ret | t1 == ret -> Just ret; _ -> Nothing -> Just t)
    where
      Fun1STy t = FunTy t t
  
  pattern Fun2STy t <-
    (\case Fun2Ty t1 t2 ret | t1 == ret && t2 == ret -> Just ret; _ -> Nothing -> Just t)
    where
      Fun2STy t = Fun2Ty t t t
  
  pattern Fun3STy t <-
    (\case Fun3Ty t1 t2 t3 ret | t1 == ret && t2 == ret && t3 == ret -> Just ret; _ -> Nothing -> Just t)
    where
      Fun3STy t = Fun3Ty t t t t
  
  pattern FunLTy t <-
    (\case FunTy (ListTy t1) ret | t1 == ret -> Just ret; _ -> Nothing -> Just t)
    where
      FunLTy t = FunTy (ListTy t) t
  
  vectorTy :: Integer -> Type
  vectorTy n
    | 0 <= n && n < 10000 = TupleTy (replicate (fromInteger n) IntTy)
    | otherwise = error $ "Jikka.Core.Language.Expr.vectorTy: invalid size: " ++ show n
  
  matrixTy :: Integer -> Integer -> Type
  matrixTy h w
    | 0 <= h && h < 10000 && 0 <= w && w < 10000 = TupleTy (replicate (fromInteger h) (TupleTy (replicate (fromInteger w) IntTy)))
    | otherwise = error $ "Jikka.Core.Language.Expr.matrixTy: invalid size: " ++ show (h, w)
  
  pattern UnitTy = TupleTy []
  
  pattern ConvexHullTrickTy = DataStructureTy ConvexHullTrick
  
  pattern SegmentTreeTy semigrp = DataStructureTy (SegmentTree semigrp)
  
  pattern LitInt' n = Lit (LitInt n)
  
  pattern Lit0 = Lit (LitInt 0)
  
  pattern Lit1 = Lit (LitInt 1)
  
  pattern Lit2 = Lit (LitInt 2)
  
  pattern LitMinus1 = Lit (LitInt (-1))
  
  pattern LitBool' p = Lit (LitBool p)
  
  pattern LitTrue = Lit (LitBool True)
  
  pattern LitFalse = Lit (LitBool False)
  
  pattern Builtin builtin = Lit (LitBuiltin builtin [])
  
  pattern Builtin1 builtin t1 = Lit (LitBuiltin builtin [t1])
  
  pattern Builtin2 builtin t1 t2 = Lit (LitBuiltin builtin [t1, t2])
  
  pattern App2 f e1 e2 = App (App f e1) e2
  
  pattern App3 f e1 e2 e3 = App (App (App f e1) e2) e3
  
  pattern App4 f e1 e2 e3 e4 = App (App (App (App f e1) e2) e3) e4
  
  pattern AppBuiltin1 builtin e1 = App (Lit (LitBuiltin builtin [])) e1
  
  pattern AppBuiltin11 builtin t1 e1 = App (Lit (LitBuiltin builtin [t1])) e1
  
  pattern AppBuiltin2 builtin e1 e2 = App2 (Lit (LitBuiltin builtin [])) e1 e2
  
  pattern AppBuiltin12 builtin t1 e1 e2 = App2 (Lit (LitBuiltin builtin [t1])) e1 e2
  
  pattern AppBuiltin22 builtin t1 t2 e1 e2 = App2 (Lit (LitBuiltin builtin [t1, t2])) e1 e2
  
  pattern AppBuiltin3 builtin e1 e2 e3 = App3 (Lit (LitBuiltin builtin [])) e1 e2 e3
  
  pattern AppBuiltin13 builtin t1 e1 e2 e3 = App3 (Lit (LitBuiltin builtin [t1])) e1 e2 e3
  
  pattern AppBuiltin23 builtin t1 t2 e1 e2 e3 = App3 (Lit (LitBuiltin builtin [t1, t2])) e1 e2 e3
  
  pattern AppBuiltin14 builtin t1 t2 e1 e2 e3 = App3 (Lit (LitBuiltin builtin [t1, t2])) e1 e2 e3
  
  pattern Lam2 x1 t1 x2 t2 e = Lam x1 t1 (Lam x2 t2 e)
  
  pattern Lam3 x1 t1 x2 t2 x3 t3 e = Lam x1 t1 (Lam x2 t2 (Lam x3 t3 e))
  
  -- | `ToplevelExpr` is the toplevel exprs. In our core, "let rec" is allowed only on the toplevel.
  --
  -- \[
  --     \begin{array}{rl}
  --         \mathrm{tle} ::= & e \\
  --         \vert & \mathbf{let}~ x: \tau = e ~\mathbf{in}~ \mathrm{tle} \\
  --         \vert & \mathbf{let~rec}~ x(x: \tau, x: \tau, \dots, x: \tau): \tau = e ~\mathbf{in}~ \mathrm{tle} \\
  --         \vert & \mathbf{assert}~ e ~\mathbf{in}~ \mathrm{tle}
  --     \end{array}
  -- \]
  data ToplevelExpr
    = ResultExpr Expr
    | ToplevelLet VarName Type Expr ToplevelExpr
    | ToplevelLetRec VarName [(VarName, Type)] Type Expr ToplevelExpr
    | ToplevelAssert Expr ToplevelExpr
    deriving (Eq, Ord, Show, Read, Data, Typeable)
  
  type Program = ToplevelExpr
  