{-# 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 :: m Type
genType = TypeName -> Type
VarTy (TypeName -> Type) -> (Int -> TypeName) -> Int -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OccName -> NameFlavour -> TypeName
TypeName OccName
forall a. Maybe a
Nothing (NameFlavour -> TypeName)
-> (Int -> NameFlavour) -> Int -> TypeName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> NameFlavour
forall a. a -> Maybe a
Just (Int -> Type) -> m Int -> m Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Int
forall (m :: * -> *). MonadAlpha m => m Int
nextCounter

genVarName :: MonadAlpha m => VarName -> m VarName
genVarName :: VarName -> m VarName
genVarName (VarName OccName
x NameFlavour
_) = OccName -> NameFlavour -> VarName
VarName OccName
x (NameFlavour -> VarName) -> (Int -> NameFlavour) -> Int -> VarName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> NameFlavour
forall a. a -> Maybe a
Just (Int -> VarName) -> m Int -> m VarName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Int
forall (m :: * -> *). MonadAlpha m => m Int
nextCounter

genVarName' :: MonadAlpha m => m VarName
genVarName' :: m VarName
genVarName' = VarName -> m VarName
forall (m :: * -> *). MonadAlpha m => VarName -> m VarName
genVarName (OccName -> NameFlavour -> VarName
VarName OccName
forall a. Maybe a
Nothing NameFlavour
forall a. Maybe a
Nothing)

genVarName'' :: MonadAlpha m => Expr -> m VarName
genVarName'' :: Expr -> m VarName
genVarName'' = \case
  Var VarName
x -> VarName -> m VarName
forall (m :: * -> *). MonadAlpha m => VarName -> m VarName
genVarName VarName
x
  Expr
_ -> m VarName
forall (m :: * -> *). MonadAlpha m => m VarName
genVarName'

mapSubTypesM :: Monad m => (Type -> m Type) -> Type -> m Type
mapSubTypesM :: (Type -> m Type) -> Type -> m Type
mapSubTypesM Type -> m Type
f = Type -> m Type
go
  where
    go :: Type -> m Type
go = \case
      VarTy TypeName
x -> Type -> m Type
f (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ TypeName -> Type
VarTy TypeName
x
      Type
IntTy -> Type -> m Type
f Type
IntTy
      Type
BoolTy -> Type -> m Type
f Type
BoolTy
      ListTy Type
t -> Type -> m Type
f (Type -> m Type) -> (Type -> Type) -> Type -> m Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type
ListTy (Type -> m Type) -> m Type -> m Type
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Type -> m Type
f Type
t
      TupleTy [Type]
ts -> Type -> m Type
f (Type -> m Type) -> ([Type] -> Type) -> [Type] -> m Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Type] -> Type
TupleTy ([Type] -> m Type) -> m [Type] -> m Type
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Type -> m Type) -> [Type] -> m [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Type -> m Type
f [Type]
ts
      FunTy Type
t1 Type
t2 -> Type -> m Type
f (Type -> m Type) -> m Type -> m Type
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Type -> Type -> Type
FunTy (Type -> Type -> Type) -> m Type -> m (Type -> Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> m Type
f Type
t1 m (Type -> Type) -> m Type -> m Type
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> m Type
f Type
t2)
      DataStructureTy DataStructure
ds -> Type -> m Type
f (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ DataStructure -> Type
DataStructureTy DataStructure
ds

mapTypeLiteralM :: Monad m => (Type -> m Type) -> Literal -> m Literal
mapTypeLiteralM :: (Type -> m Type) -> Literal -> m Literal
mapTypeLiteralM Type -> m Type
f = \case
  LitBuiltin Builtin
builtin [Type]
ts -> Builtin -> [Type] -> Literal
LitBuiltin Builtin
builtin ([Type] -> Literal) -> m [Type] -> m Literal
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Type -> m Type) -> [Type] -> m [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Type -> m Type
f [Type]
ts
  LitInt Integer
n -> Literal -> m Literal
forall (m :: * -> *) a. Monad m => a -> m a
return (Literal -> m Literal) -> Literal -> m Literal
forall a b. (a -> b) -> a -> b
$ Integer -> Literal
LitInt Integer
n
  LitBool Bool
p -> Literal -> m Literal
forall (m :: * -> *) a. Monad m => a -> m a
return (Literal -> m Literal) -> Literal -> m Literal
forall a b. (a -> b) -> a -> b
$ Bool -> Literal
LitBool Bool
p
  LitNil Type
t -> Type -> Literal
LitNil (Type -> Literal) -> m Type -> m Literal
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> m Type
f Type
t
  LitBottom Type
t String
err -> Type -> String -> Literal
LitBottom (Type -> String -> Literal) -> m Type -> m (String -> Literal)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> m Type
f Type
t m (String -> Literal) -> m String -> m Literal
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> m String
forall (f :: * -> *) a. Applicative f => a -> f a
pure String
err

mapTypeExprM :: Monad m => (Type -> m Type) -> Expr -> m Expr
mapTypeExprM :: (Type -> m Type) -> Expr -> m Expr
mapTypeExprM Type -> m Type
f = Expr -> m Expr
go
  where
    go :: Expr -> m Expr
go = \case
      Var VarName
x -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ VarName -> Expr
Var VarName
x
      Lit Literal
lit -> Literal -> Expr
Lit (Literal -> Expr) -> m Literal -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Type -> m Type) -> Literal -> m Literal
forall (m :: * -> *).
Monad m =>
(Type -> m Type) -> Literal -> m Literal
mapTypeLiteralM Type -> m Type
f Literal
lit
      App Expr
f Expr
e -> Expr -> Expr -> Expr
App (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> m Expr
go Expr
f m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
e
      Lam VarName
x Type
t Expr
body -> VarName -> Type -> Expr -> Expr
Lam VarName
x (Type -> Expr -> Expr) -> m Type -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> m Type
f Type
t m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
body
      Let VarName
x Type
t Expr
e1 Expr
e2 -> VarName -> Type -> Expr -> Expr -> Expr
Let VarName
x (Type -> Expr -> Expr -> Expr)
-> m Type -> m (Expr -> Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> m Type
f Type
t m (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
e1 m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
e2
      Assert Expr
e1 Expr
e2 -> Expr -> Expr -> Expr
Assert (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> m Expr
go Expr
e1 m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
e2

mapTypeExpr :: (Type -> Type) -> Expr -> Expr
mapTypeExpr :: (Type -> Type) -> Expr -> Expr
mapTypeExpr Type -> Type
f Expr
e = Identity Expr -> Expr
forall a. Identity a -> a
runIdentity ((Type -> Identity Type) -> Expr -> Identity Expr
forall (m :: * -> *). Monad m => (Type -> m Type) -> Expr -> m Expr
mapTypeExprM (Type -> Identity Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Identity Type) -> (Type -> Type) -> Type -> Identity Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type
f) Expr
e)

mapTypeToplevelExprM :: Monad m => (Type -> m Type) -> ToplevelExpr -> m ToplevelExpr
mapTypeToplevelExprM :: (Type -> m Type) -> ToplevelExpr -> m ToplevelExpr
mapTypeToplevelExprM Type -> m Type
f = \case
  ResultExpr Expr
e -> Expr -> ToplevelExpr
ResultExpr (Expr -> ToplevelExpr) -> m Expr -> m ToplevelExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Type -> m Type) -> Expr -> m Expr
forall (m :: * -> *). Monad m => (Type -> m Type) -> Expr -> m Expr
mapTypeExprM Type -> m Type
f Expr
e
  ToplevelLet VarName
x Type
t Expr
e ToplevelExpr
cont -> VarName -> Type -> Expr -> ToplevelExpr -> ToplevelExpr
ToplevelLet VarName
x (Type -> Expr -> ToplevelExpr -> ToplevelExpr)
-> m Type -> m (Expr -> ToplevelExpr -> ToplevelExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> m Type
f Type
t m (Expr -> ToplevelExpr -> ToplevelExpr)
-> m Expr -> m (ToplevelExpr -> ToplevelExpr)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Type -> m Type) -> Expr -> m Expr
forall (m :: * -> *). Monad m => (Type -> m Type) -> Expr -> m Expr
mapTypeExprM Type -> m Type
f Expr
e m (ToplevelExpr -> ToplevelExpr)
-> m ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Type -> m Type) -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
Monad m =>
(Type -> m Type) -> ToplevelExpr -> m ToplevelExpr
mapTypeToplevelExprM Type -> m Type
f ToplevelExpr
cont
  ToplevelLetRec VarName
g [(VarName, Type)]
args Type
ret Expr
body ToplevelExpr
cont -> VarName
-> [(VarName, Type)]
-> Type
-> Expr
-> ToplevelExpr
-> ToplevelExpr
ToplevelLetRec VarName
g ([(VarName, Type)] -> Type -> Expr -> ToplevelExpr -> ToplevelExpr)
-> m [(VarName, Type)]
-> m (Type -> Expr -> ToplevelExpr -> ToplevelExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((VarName, Type) -> m (VarName, Type))
-> [(VarName, Type)] -> m [(VarName, Type)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\(VarName
x, Type
t) -> (VarName
x,) (Type -> (VarName, Type)) -> m Type -> m (VarName, Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> m Type
f Type
t) [(VarName, Type)]
args m (Type -> Expr -> ToplevelExpr -> ToplevelExpr)
-> m Type -> m (Expr -> ToplevelExpr -> ToplevelExpr)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> m Type
f Type
ret m (Expr -> ToplevelExpr -> ToplevelExpr)
-> m Expr -> m (ToplevelExpr -> ToplevelExpr)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Type -> m Type) -> Expr -> m Expr
forall (m :: * -> *). Monad m => (Type -> m Type) -> Expr -> m Expr
mapTypeExprM Type -> m Type
f Expr
body m (ToplevelExpr -> ToplevelExpr)
-> m ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Type -> m Type) -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
Monad m =>
(Type -> m Type) -> ToplevelExpr -> m ToplevelExpr
mapTypeToplevelExprM Type -> m Type
f ToplevelExpr
cont
  ToplevelAssert Expr
e ToplevelExpr
cont -> Expr -> ToplevelExpr -> ToplevelExpr
ToplevelAssert (Expr -> ToplevelExpr -> ToplevelExpr)
-> m Expr -> m (ToplevelExpr -> ToplevelExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Type -> m Type) -> Expr -> m Expr
forall (m :: * -> *). Monad m => (Type -> m Type) -> Expr -> m Expr
mapTypeExprM Type -> m Type
f Expr
e m (ToplevelExpr -> ToplevelExpr)
-> m ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Type -> m Type) -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
Monad m =>
(Type -> m Type) -> ToplevelExpr -> m ToplevelExpr
mapTypeToplevelExprM Type -> m Type
f ToplevelExpr
cont

mapTypeProgramM :: Monad m => (Type -> m Type) -> Program -> m Program
mapTypeProgramM :: (Type -> m Type) -> ToplevelExpr -> m ToplevelExpr
mapTypeProgramM = (Type -> m Type) -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
Monad m =>
(Type -> m Type) -> ToplevelExpr -> m ToplevelExpr
mapTypeToplevelExprM

mapTypeProgram :: (Type -> Type) -> Program -> Program
mapTypeProgram :: (Type -> Type) -> ToplevelExpr -> ToplevelExpr
mapTypeProgram Type -> Type
f ToplevelExpr
prog = Identity ToplevelExpr -> ToplevelExpr
forall a. Identity a -> a
runIdentity ((Type -> Identity Type) -> ToplevelExpr -> Identity ToplevelExpr
forall (m :: * -> *).
Monad m =>
(Type -> m Type) -> ToplevelExpr -> m ToplevelExpr
mapTypeProgramM (Type -> Identity Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Identity Type) -> (Type -> Type) -> Type -> Identity Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type
f) ToplevelExpr
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' :: ([(VarName, Type)] -> Expr -> m Expr)
-> ([(VarName, Type)] -> Expr -> m Expr)
-> [(VarName, Type)]
-> Expr
-> m Expr
mapSubExprM' [(VarName, Type)] -> Expr -> m Expr
pre [(VarName, Type)] -> Expr -> m Expr
post [(VarName, Type)]
env Expr
e = do
  Expr
e <- [(VarName, Type)] -> Expr -> m Expr
pre [(VarName, Type)]
env Expr
e
  let go :: [(VarName, Type)] -> Expr -> m Expr
go = ([(VarName, Type)] -> Expr -> m Expr)
-> ([(VarName, Type)] -> Expr -> m Expr)
-> [(VarName, Type)]
-> Expr
-> m Expr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> Expr -> m Expr)
-> ([(VarName, Type)] -> Expr -> m Expr)
-> [(VarName, Type)]
-> Expr
-> m Expr
mapSubExprM' [(VarName, Type)] -> Expr -> m Expr
pre [(VarName, Type)] -> Expr -> m Expr
post
  Expr
e <- case Expr
e of
    Var VarName
y -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ VarName -> Expr
Var VarName
y
    Lit Literal
lit -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ Literal -> Expr
Lit Literal
lit
    App Expr
g Expr
e -> Expr -> Expr -> Expr
App (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> m Expr
go [(VarName, Type)]
env Expr
g m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [(VarName, Type)] -> Expr -> m Expr
go [(VarName, Type)]
env Expr
e
    Lam VarName
x Type
t Expr
body -> VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> m Expr
go ((VarName
x, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
env) Expr
body
    Let VarName
y Type
t Expr
e1 Expr
e2 -> VarName -> Type -> Expr -> Expr -> Expr
Let VarName
y Type
t (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> m Expr
go [(VarName, Type)]
env Expr
e1 m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [(VarName, Type)] -> Expr -> m Expr
go ((VarName
y, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
env) Expr
e2
    Assert Expr
e1 Expr
e2 -> Expr -> Expr -> Expr
Assert (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> m Expr
go [(VarName, Type)]
env Expr
e1 m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [(VarName, Type)] -> Expr -> m Expr
go [(VarName, Type)]
env Expr
e2
  [(VarName, Type)] -> Expr -> m Expr
post [(VarName, Type)]
env Expr
e

mapToplevelExprM' :: Monad m => ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr) -> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr) -> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
mapToplevelExprM' :: ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)]
-> ToplevelExpr
-> m ToplevelExpr
mapToplevelExprM' [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
pre [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
post [(VarName, Type)]
env ToplevelExpr
e = do
  ToplevelExpr
e <- [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
pre [(VarName, Type)]
env ToplevelExpr
e
  ToplevelExpr
e <- case ToplevelExpr
e of
    ResultExpr Expr
e -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (ToplevelExpr -> m ToplevelExpr) -> ToplevelExpr -> m ToplevelExpr
forall a b. (a -> b) -> a -> b
$ Expr -> ToplevelExpr
ResultExpr Expr
e
    ToplevelLet VarName
y Type
t Expr
e ToplevelExpr
cont ->
      VarName -> Type -> Expr -> ToplevelExpr -> ToplevelExpr
ToplevelLet VarName
y Type
t Expr
e (ToplevelExpr -> ToplevelExpr) -> m ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)]
-> ToplevelExpr
-> m ToplevelExpr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)]
-> ToplevelExpr
-> m ToplevelExpr
mapToplevelExprM' [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
pre [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
post ((VarName
y, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
env) ToplevelExpr
cont
    ToplevelLetRec VarName
g [(VarName, Type)]
args Type
ret Expr
body ToplevelExpr
cont ->
      let env' :: [(VarName, Type)]
env' = (VarName
g, ((VarName, Type) -> Type -> Type)
-> Type -> [(VarName, Type)] -> Type
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Type -> Type -> Type
FunTy (Type -> Type -> Type)
-> ((VarName, Type) -> Type) -> (VarName, Type) -> Type -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VarName, Type) -> Type
forall a b. (a, b) -> b
snd) Type
ret [(VarName, Type)]
args) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
env
       in VarName
-> [(VarName, Type)]
-> Type
-> Expr
-> ToplevelExpr
-> ToplevelExpr
ToplevelLetRec VarName
g [(VarName, Type)]
args Type
ret Expr
body (ToplevelExpr -> ToplevelExpr) -> m ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)]
-> ToplevelExpr
-> m ToplevelExpr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)]
-> ToplevelExpr
-> m ToplevelExpr
mapToplevelExprM' [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
pre [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
post [(VarName, Type)]
env' ToplevelExpr
cont
    ToplevelAssert Expr
e ToplevelExpr
cont ->
      Expr -> ToplevelExpr -> ToplevelExpr
ToplevelAssert Expr
e (ToplevelExpr -> ToplevelExpr) -> m ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)]
-> ToplevelExpr
-> m ToplevelExpr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)]
-> ToplevelExpr
-> m ToplevelExpr
mapToplevelExprM' [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
pre [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
post [(VarName, Type)]
env ToplevelExpr
cont
  [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
post [(VarName, Type)]
env ToplevelExpr
e

mapExprToplevelExprM :: Monad m => ([(VarName, Type)] -> Expr -> m Expr) -> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
mapExprToplevelExprM :: ([(VarName, Type)] -> Expr -> m Expr)
-> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
mapExprToplevelExprM [(VarName, Type)] -> Expr -> m Expr
f [(VarName, Type)]
env = ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)]
-> ToplevelExpr
-> m ToplevelExpr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)]
-> ToplevelExpr
-> m ToplevelExpr
mapToplevelExprM' [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
pre' (\[(VarName, Type)]
_ ToplevelExpr
e -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ToplevelExpr
e) [(VarName, Type)]
env
  where
    pre' :: [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
pre' [(VarName, Type)]
env = \case
      ResultExpr Expr
e -> Expr -> ToplevelExpr
ResultExpr (Expr -> ToplevelExpr) -> m Expr -> m ToplevelExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> m Expr
f [(VarName, Type)]
env Expr
e
      ToplevelLet VarName
y Type
t Expr
e ToplevelExpr
cont -> VarName -> Type -> Expr -> ToplevelExpr -> ToplevelExpr
ToplevelLet VarName
y Type
t (Expr -> ToplevelExpr -> ToplevelExpr)
-> m Expr -> m (ToplevelExpr -> ToplevelExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> m Expr
f [(VarName, Type)]
env Expr
e m (ToplevelExpr -> ToplevelExpr)
-> m ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *) a. Applicative f => a -> f a
pure ToplevelExpr
cont
      ToplevelLetRec VarName
g [(VarName, Type)]
args Type
ret Expr
body ToplevelExpr
cont ->
        let env' :: [(VarName, Type)]
env' = (VarName
g, ((VarName, Type) -> Type -> Type)
-> Type -> [(VarName, Type)] -> Type
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Type -> Type -> Type
FunTy (Type -> Type -> Type)
-> ((VarName, Type) -> Type) -> (VarName, Type) -> Type -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VarName, Type) -> Type
forall a b. (a, b) -> b
snd) Type
ret [(VarName, Type)]
args) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
env
         in VarName
-> [(VarName, Type)]
-> Type
-> Expr
-> ToplevelExpr
-> ToplevelExpr
ToplevelLetRec VarName
g [(VarName, Type)]
args Type
ret (Expr -> ToplevelExpr -> ToplevelExpr)
-> m Expr -> m (ToplevelExpr -> ToplevelExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> m Expr
f ([(VarName, Type)] -> [(VarName, Type)]
forall a. [a] -> [a]
reverse [(VarName, Type)]
args [(VarName, Type)] -> [(VarName, Type)] -> [(VarName, Type)]
forall a. [a] -> [a] -> [a]
++ [(VarName, Type)]
env') Expr
body m (ToplevelExpr -> ToplevelExpr)
-> m ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *) a. Applicative f => a -> f a
pure ToplevelExpr
cont
      ToplevelAssert Expr
e ToplevelExpr
cont -> Expr -> ToplevelExpr -> ToplevelExpr
ToplevelAssert (Expr -> ToplevelExpr -> ToplevelExpr)
-> m Expr -> m (ToplevelExpr -> ToplevelExpr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(VarName, Type)] -> Expr -> m Expr
f [(VarName, Type)]
env Expr
e m (ToplevelExpr -> ToplevelExpr)
-> m ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ToplevelExpr -> m ToplevelExpr
forall (f :: * -> *) a. Applicative f => a -> f a
pure ToplevelExpr
cont

mapExprProgramM :: Monad m => ([(VarName, Type)] -> Expr -> m Expr) -> Program -> m Program
mapExprProgramM :: ([(VarName, Type)] -> Expr -> m Expr)
-> ToplevelExpr -> m ToplevelExpr
mapExprProgramM [(VarName, Type)] -> Expr -> m Expr
f = ([(VarName, Type)] -> Expr -> m Expr)
-> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> Expr -> m Expr)
-> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
mapExprToplevelExprM [(VarName, Type)] -> Expr -> m Expr
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 :: ([(VarName, Type)] -> Expr -> m Expr)
-> [(VarName, Type)] -> Expr -> m Expr
mapSubExprM [(VarName, Type)] -> Expr -> m Expr
f = ([(VarName, Type)] -> Expr -> m Expr)
-> ([(VarName, Type)] -> Expr -> m Expr)
-> [(VarName, Type)]
-> Expr
-> m Expr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> Expr -> m Expr)
-> ([(VarName, Type)] -> Expr -> m Expr)
-> [(VarName, Type)]
-> Expr
-> m Expr
mapSubExprM' (\[(VarName, Type)]
_ Expr
e -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
e) [(VarName, Type)] -> Expr -> m Expr
f

mapSubExpr :: ([(VarName, Type)] -> Expr -> Expr) -> [(VarName, Type)] -> Expr -> Expr
mapSubExpr :: ([(VarName, Type)] -> Expr -> Expr)
-> [(VarName, Type)] -> Expr -> Expr
mapSubExpr [(VarName, Type)] -> Expr -> Expr
f [(VarName, Type)]
env Expr
e = Identity Expr -> Expr
forall a. Identity a -> a
runIdentity (Identity Expr -> Expr) -> Identity Expr -> Expr
forall a b. (a -> b) -> a -> b
$ ([(VarName, Type)] -> Expr -> Identity Expr)
-> [(VarName, Type)] -> Expr -> Identity Expr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> Expr -> m Expr)
-> [(VarName, Type)] -> Expr -> m Expr
mapSubExprM (\[(VarName, Type)]
env Expr
e -> Expr -> Identity Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> Identity Expr) -> Expr -> Identity Expr
forall a b. (a -> b) -> a -> b
$ [(VarName, Type)] -> Expr -> Expr
f [(VarName, Type)]
env Expr
e) [(VarName, Type)]
env Expr
e

mapExprToplevelExpr :: ([(VarName, Type)] -> Expr -> Expr) -> [(VarName, Type)] -> ToplevelExpr -> ToplevelExpr
mapExprToplevelExpr :: ([(VarName, Type)] -> Expr -> Expr)
-> [(VarName, Type)] -> ToplevelExpr -> ToplevelExpr
mapExprToplevelExpr [(VarName, Type)] -> Expr -> Expr
f [(VarName, Type)]
env ToplevelExpr
e = Identity ToplevelExpr -> ToplevelExpr
forall a. Identity a -> a
runIdentity (Identity ToplevelExpr -> ToplevelExpr)
-> Identity ToplevelExpr -> ToplevelExpr
forall a b. (a -> b) -> a -> b
$ ([(VarName, Type)] -> Expr -> Identity Expr)
-> [(VarName, Type)] -> ToplevelExpr -> Identity ToplevelExpr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> Expr -> m Expr)
-> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
mapExprToplevelExprM (\[(VarName, Type)]
env Expr
e -> Expr -> Identity Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> Identity Expr) -> Expr -> Identity Expr
forall a b. (a -> b) -> a -> b
$ [(VarName, Type)] -> Expr -> Expr
f [(VarName, Type)]
env Expr
e) [(VarName, Type)]
env ToplevelExpr
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 :: ([(VarName, Type)] -> Expr -> Expr) -> ToplevelExpr -> ToplevelExpr
mapExprProgram [(VarName, Type)] -> Expr -> Expr
f ToplevelExpr
prog = Identity ToplevelExpr -> ToplevelExpr
forall a. Identity a -> a
runIdentity (Identity ToplevelExpr -> ToplevelExpr)
-> Identity ToplevelExpr -> ToplevelExpr
forall a b. (a -> b) -> a -> b
$ ([(VarName, Type)] -> Expr -> Identity Expr)
-> ToplevelExpr -> Identity ToplevelExpr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> Expr -> m Expr)
-> ToplevelExpr -> m ToplevelExpr
mapExprProgramM (\[(VarName, Type)]
env Expr
e -> Expr -> Identity Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> Identity Expr) -> Expr -> Identity Expr
forall a b. (a -> b) -> a -> b
$ [(VarName, Type)] -> Expr -> Expr
f [(VarName, Type)]
env Expr
e) ToplevelExpr
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 :: ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
mapToplevelExprM [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
f [(VarName, Type)]
env ToplevelExpr
e = ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)]
-> ToplevelExpr
-> m ToplevelExpr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)]
-> ToplevelExpr
-> m ToplevelExpr
mapToplevelExprM' (\[(VarName, Type)]
_ ToplevelExpr
e -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ToplevelExpr
e) [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
f [(VarName, Type)]
env ToplevelExpr
e

mapToplevelExprProgramM :: Monad m => ([(VarName, Type)] -> Program -> m Program) -> Program -> m Program
mapToplevelExprProgramM :: ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ToplevelExpr -> m ToplevelExpr
mapToplevelExprProgramM [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
f ToplevelExpr
prog = ([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
mapToplevelExprM [(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr
f [] ToplevelExpr
prog

mapToplevelExprProgram :: ([(VarName, Type)] -> Program -> Program) -> Program -> Program
mapToplevelExprProgram :: ([(VarName, Type)] -> ToplevelExpr -> ToplevelExpr)
-> ToplevelExpr -> ToplevelExpr
mapToplevelExprProgram [(VarName, Type)] -> ToplevelExpr -> ToplevelExpr
f ToplevelExpr
prog = Identity ToplevelExpr -> ToplevelExpr
forall a. Identity a -> a
runIdentity (Identity ToplevelExpr -> ToplevelExpr)
-> Identity ToplevelExpr -> ToplevelExpr
forall a b. (a -> b) -> a -> b
$ ([(VarName, Type)] -> ToplevelExpr -> Identity ToplevelExpr)
-> ToplevelExpr -> Identity ToplevelExpr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> ToplevelExpr -> m ToplevelExpr)
-> ToplevelExpr -> m ToplevelExpr
mapToplevelExprProgramM (\[(VarName, Type)]
env ToplevelExpr
e -> ToplevelExpr -> Identity ToplevelExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (ToplevelExpr -> Identity ToplevelExpr)
-> ToplevelExpr -> Identity ToplevelExpr
forall a b. (a -> b) -> a -> b
$ [(VarName, Type)] -> ToplevelExpr -> ToplevelExpr
f [(VarName, Type)]
env ToplevelExpr
e) ToplevelExpr
prog

listSubExprs :: Expr -> [Expr]
listSubExprs :: Expr -> [Expr]
listSubExprs Expr
e = Dual [Expr] -> [Expr]
forall a. Dual a -> a
getDual (Dual [Expr] -> [Expr])
-> (Writer (Dual [Expr]) Expr -> Dual [Expr])
-> Writer (Dual [Expr]) Expr
-> [Expr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Writer (Dual [Expr]) Expr -> Dual [Expr]
forall w a. Writer w a -> w
execWriter (Writer (Dual [Expr]) Expr -> [Expr])
-> Writer (Dual [Expr]) Expr -> [Expr]
forall a b. (a -> b) -> a -> b
$ ([(VarName, Type)] -> Expr -> Writer (Dual [Expr]) Expr)
-> [(VarName, Type)] -> Expr -> Writer (Dual [Expr]) Expr
forall (m :: * -> *).
Monad m =>
([(VarName, Type)] -> Expr -> m Expr)
-> [(VarName, Type)] -> Expr -> m Expr
mapSubExprM [(VarName, Type)] -> Expr -> Writer (Dual [Expr]) Expr
forall (m :: * -> *) b p. MonadWriter (Dual [b]) m => p -> b -> m b
go [] Expr
e
  where
    go :: p -> b -> m b
go p
_ b
e = do
      Dual [b] -> m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Dual [b] -> m ()) -> Dual [b] -> m ()
forall a b. (a -> b) -> a -> b
$ [b] -> Dual [b]
forall a. a -> Dual a
Dual [b
e]
      b -> m b
forall (m :: * -> *) a. Monad m => a -> m a
return b
e

uncurryFunTy :: Type -> ([Type], Type)
uncurryFunTy :: Type -> ([Type], Type)
uncurryFunTy = \case
  (FunTy Type
t Type
t') -> let ([Type]
ts, Type
ret) = Type -> ([Type], Type)
uncurryFunTy Type
t' in (Type
t Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: [Type]
ts, Type
ret)
  Type
ret -> ([], Type
ret)

uncurryLam :: Expr -> ([(VarName, Type)], Expr)
uncurryLam :: Expr -> ([(VarName, Type)], Expr)
uncurryLam = \case
  Lam VarName
x Type
t Expr
body -> let ([(VarName, Type)]
args, Expr
body') = Expr -> ([(VarName, Type)], Expr)
uncurryLam Expr
body in ((VarName
x, Type
t) (VarName, Type) -> [(VarName, Type)] -> [(VarName, Type)]
forall a. a -> [a] -> [a]
: [(VarName, Type)]
args, Expr
body')
  Expr
body -> ([], Expr
body)

curryApp :: Expr -> (Expr, [Expr])
curryApp :: Expr -> (Expr, [Expr])
curryApp = \case
  App Expr
f Expr
e -> let (Expr
f', [Expr]
e') = Expr -> (Expr, [Expr])
curryApp Expr
f in (Expr
f', [Expr]
e' [Expr] -> [Expr] -> [Expr]
forall a. [a] -> [a] -> [a]
++ [Expr
e])
  Expr
f -> (Expr
f, [])

curryFunTy :: [Type] -> Type -> Type
curryFunTy :: [Type] -> Type -> Type
curryFunTy [Type]
ts Type
ret = (Type -> Type -> Type) -> Type -> [Type] -> Type
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Type -> Type -> Type
FunTy Type
ret [Type]
ts

curryLam :: [(VarName, Type)] -> Expr -> Expr
curryLam :: [(VarName, Type)] -> Expr -> Expr
curryLam [(VarName, Type)]
args Expr
body = ((VarName, Type) -> Expr -> Expr)
-> Expr -> [(VarName, Type)] -> Expr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((VarName -> Type -> Expr -> Expr)
-> (VarName, Type) -> Expr -> Expr
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry VarName -> Type -> Expr -> Expr
Lam) Expr
body [(VarName, Type)]
args

uncurryApp :: Expr -> [Expr] -> Expr
uncurryApp :: Expr -> [Expr] -> Expr
uncurryApp = (Expr -> Expr -> Expr) -> Expr -> [Expr] -> Expr
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Expr -> Expr -> Expr
App

isVectorTy :: Type -> Bool
isVectorTy :: Type -> Bool
isVectorTy = NameFlavour -> Bool
forall a. Maybe a -> Bool
isJust (NameFlavour -> Bool) -> (Type -> NameFlavour) -> Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> NameFlavour
sizeOfVectorTy

isVectorTy' :: [Type] -> Bool
isVectorTy' :: [Type] -> Bool
isVectorTy' = Type -> Bool
isVectorTy (Type -> Bool) -> ([Type] -> Type) -> [Type] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Type] -> Type
TupleTy

sizeOfVectorTy :: Type -> Maybe Int
sizeOfVectorTy :: Type -> NameFlavour
sizeOfVectorTy = \case
  TupleTy [Type]
ts | (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== Type
IntTy) [Type]
ts -> Int -> NameFlavour
forall a. a -> Maybe a
Just ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
ts)
  Type
_ -> NameFlavour
forall a. Maybe a
Nothing

isMatrixTy :: Type -> Bool
isMatrixTy :: Type -> Bool
isMatrixTy = Maybe (Int, Int) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Int, Int) -> Bool)
-> (Type -> Maybe (Int, Int)) -> Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Maybe (Int, Int)
sizeOfMatrixTy

isMatrixTy' :: [Type] -> Bool
isMatrixTy' :: [Type] -> Bool
isMatrixTy' = Type -> Bool
isMatrixTy (Type -> Bool) -> ([Type] -> Type) -> [Type] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Type] -> Type
TupleTy

sizeOfMatrixTy :: Type -> Maybe (Int, Int)
sizeOfMatrixTy :: Type -> Maybe (Int, Int)
sizeOfMatrixTy = \case
  TupleTy ts :: [Type]
ts@(TupleTy [Type]
ts' : [Type]
_) | (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== Type
IntTy) [Type]
ts' Bool -> Bool -> Bool
&& (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== [Type] -> Type
TupleTy [Type]
ts') [Type]
ts -> (Int, Int) -> Maybe (Int, Int)
forall a. a -> Maybe a
Just ([Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
ts, [Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
ts')
  Type
_ -> Maybe (Int, Int)
forall a. Maybe a
Nothing

isConstantTimeBuiltin :: Builtin -> Bool
isConstantTimeBuiltin :: Builtin -> Bool
isConstantTimeBuiltin = \case
  -- arithmetical functions
  Builtin
Negate -> Bool
True
  Builtin
Plus -> Bool
True
  Builtin
Minus -> Bool
True
  Builtin
Mult -> Bool
True
  Builtin
FloorDiv -> Bool
True
  Builtin
FloorMod -> Bool
True
  Builtin
CeilDiv -> Bool
True
  Builtin
CeilMod -> Bool
True
  Builtin
JustDiv -> Bool
True
  Builtin
Pow -> Bool
True
  -- advanced arithmetical functions
  Builtin
Abs -> Bool
True
  Builtin
Gcd -> Bool
True
  Builtin
Lcm -> Bool
True
  Builtin
Min2 -> Bool
True
  Builtin
Max2 -> Bool
True
  Builtin
Iterate -> Bool
False
  -- logical functions
  Builtin
Not -> Bool
True
  Builtin
And -> Bool
True
  Builtin
Or -> Bool
True
  Builtin
Implies -> Bool
True
  Builtin
If -> Bool
True
  -- bitwise functions
  Builtin
BitNot -> Bool
True
  Builtin
BitAnd -> Bool
True
  Builtin
BitOr -> Bool
True
  Builtin
BitXor -> Bool
True
  Builtin
BitLeftShift -> Bool
True
  Builtin
BitRightShift -> Bool
True
  -- matrix functions
  MatAp Integer
_ Integer
_ -> Bool
True
  MatZero Integer
_ Integer
_ -> Bool
True
  MatOne Integer
_ -> Bool
True
  MatAdd Integer
_ Integer
_ -> Bool
True
  MatMul Integer
_ Integer
_ Integer
_ -> Bool
True
  MatPow Integer
_ -> Bool
True
  VecFloorMod Integer
_ -> Bool
True
  MatFloorMod Integer
_ Integer
_ -> Bool
True
  -- modular functions
  Builtin
ModNegate -> Bool
True
  Builtin
ModPlus -> Bool
True
  Builtin
ModMinus -> Bool
True
  Builtin
ModMult -> Bool
True
  Builtin
ModInv -> Bool
True
  Builtin
ModPow -> Bool
True
  ModMatAp Integer
_ Integer
_ -> Bool
True
  ModMatAdd Integer
_ Integer
_ -> Bool
True
  ModMatMul Integer
_ Integer
_ Integer
_ -> Bool
True
  ModMatPow Integer
_ -> Bool
True
  -- list functions
  Builtin
Cons -> Bool
False
  Builtin
Snoc -> Bool
False
  Builtin
Foldl -> Bool
False
  Builtin
Scanl -> Bool
False
  Builtin
Build -> Bool
False
  Builtin
Len -> Bool
True
  Builtin
Map -> Bool
False
  Builtin
Filter -> Bool
False
  Builtin
At -> Bool
True
  Builtin
SetAt -> Bool
False
  Builtin
Elem -> Bool
False
  Builtin
Sum -> Bool
False
  Builtin
Product -> Bool
False
  Builtin
ModSum -> Bool
False
  Builtin
ModProduct -> Bool
False
  Builtin
Min1 -> Bool
False
  Builtin
Max1 -> Bool
False
  Builtin
ArgMin -> Bool
False
  Builtin
ArgMax -> Bool
False
  Builtin
Gcd1 -> Bool
False
  Builtin
Lcm1 -> Bool
False
  Builtin
All -> Bool
False
  Builtin
Any -> Bool
False
  Builtin
Sorted -> Bool
False
  Builtin
Reversed -> Bool
False
  Builtin
Range1 -> Bool
False
  Builtin
Range2 -> Bool
False
  Builtin
Range3 -> Bool
False
  -- tuple functions
  Builtin
Tuple -> Bool
True
  Proj Integer
_ -> Bool
True
  -- comparison
  Builtin
LessThan -> Bool
True
  Builtin
LessEqual -> Bool
True
  Builtin
GreaterThan -> Bool
True
  Builtin
GreaterEqual -> Bool
True
  Builtin
Equal -> Bool
True
  Builtin
NotEqual -> Bool
True
  -- combinational functions
  Builtin
Fact -> Bool
True
  Builtin
Choose -> Bool
True
  Builtin
Permute -> Bool
True
  Builtin
MultiChoose -> Bool
True
  -- data structures
  Builtin
ConvexHullTrickInit -> Bool
False
  Builtin
ConvexHullTrickInsert -> Bool
False
  Builtin
ConvexHullTrickGetMin -> Bool
False
  SegmentTreeInitList Semigroup'
_ -> Bool
False
  SegmentTreeGetRange Semigroup'
_ -> Bool
False
  SegmentTreeSetPoint Semigroup'
_ -> Bool
False

isLiteral :: Expr -> Bool
isLiteral :: Expr -> Bool
isLiteral = \case
  Lit Literal
_ -> Bool
True
  Expr
_ -> Bool
False

-- | `isConstantTimeExpr` checks whether given exprs are suitable to propagate.
isConstantTimeExpr :: Expr -> Bool
isConstantTimeExpr :: Expr -> Bool
isConstantTimeExpr = \case
  Var VarName
_ -> Bool
True
  Lit Literal
_ -> Bool
True
  e :: Expr
e@(App Expr
_ Expr
_) -> case Expr -> (Expr, [Expr])
curryApp Expr
e of
    (Lit (LitBuiltin Builtin
f [Type]
_), [Expr]
args) -> Builtin -> Bool
isConstantTimeBuiltin Builtin
f Bool -> Bool -> Bool
&& (Expr -> Bool) -> [Expr] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Expr -> Bool
isConstantTimeExpr [Expr]
args
    (Expr, [Expr])
_ -> Bool
False
  Lam VarName
_ Type
_ Expr
_ -> Bool
True
  Let VarName
_ Type
_ Expr
e1 Expr
e2 -> Expr -> Bool
isConstantTimeExpr Expr
e1 Bool -> Bool -> Bool
&& Expr -> Bool
isConstantTimeExpr Expr
e2
  Assert Expr
e1 Expr
e2 -> Expr -> Bool
isConstantTimeExpr Expr
e1 Bool -> Bool -> Bool
&& Expr -> Bool
isConstantTimeExpr Expr
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 :: VarName -> VarName -> Integer -> Expr -> m Expr
replaceLenF VarName
f VarName
i Integer
k = Expr -> m Expr
go
  where
    go :: Expr -> m Expr
go = \case
      Len' Type
_ (Var VarName
f') | VarName
f' VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
f -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ Expr -> Expr -> Expr
Plus' (VarName -> Expr
Var VarName
i) (Integer -> Expr
LitInt' Integer
k)
      Var VarName
y -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ VarName -> Expr
Var VarName
y
      Lit Literal
lit -> Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ Literal -> Expr
Lit Literal
lit
      App Expr
g Expr
e -> Expr -> Expr -> Expr
App (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> m Expr
go Expr
g m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
e
      Lam VarName
x Type
_ Expr
_ | VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
i -> String -> m Expr
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwInternalError String
"Jikka.Core.Language.Util.replaceLenF: name conflict"
      Lam VarName
x Type
t Expr
body -> VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (if VarName
x VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
f then Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
body else Expr -> m Expr
go Expr
body)
      Let VarName
y Type
_ Expr
_ Expr
_ | VarName
y VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
i -> String -> m Expr
forall (m :: * -> *) a. MonadError Error m => String -> m a
throwInternalError String
"Jikka.Core.Language.Util.replaceLenF: name conflict"
      Let VarName
y Type
t Expr
e1 Expr
e2 -> VarName -> Type -> Expr -> Expr -> Expr
Let VarName
y Type
t (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> m Expr
go Expr
e1 m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (if VarName
y VarName -> VarName -> Bool
forall a. Eq a => a -> a -> Bool
== VarName
f then Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
e2 else Expr -> m Expr
go Expr
e2)
      Assert Expr
e1 Expr
e2 -> Expr -> Expr -> Expr
Assert (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> m Expr
go Expr
e1 m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
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 :: Expr -> ([Expr], Expr)
getRecurrenceFormulaBase = Vector (Maybe (Expr, Type)) -> Expr -> ([Expr], Expr)
go (Int -> Maybe (Expr, Type) -> Vector (Maybe (Expr, Type))
forall a. Int -> a -> Vector a
V.replicate Int
forall a. Num a => a
recurrenceLimit Maybe (Expr, Type)
forall a. Maybe a
Nothing)
  where
    recurrenceLimit :: Num a => a
    recurrenceLimit :: a
recurrenceLimit = a
20
    go :: V.Vector (Maybe (Expr, Type)) -> Expr -> ([Expr], Expr)
    go :: Vector (Maybe (Expr, Type)) -> Expr -> ([Expr], Expr)
go Vector (Maybe (Expr, Type))
base = \case
      SetAt' Type
t Expr
e (LitInt' Integer
i) Expr
e'
        | Integer
0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
i Bool -> Bool -> Bool
&& Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
forall a. Num a => a
recurrenceLimit -> Vector (Maybe (Expr, Type)) -> Expr -> ([Expr], Expr)
go (Vector (Maybe (Expr, Type))
base Vector (Maybe (Expr, Type))
-> [(Int, Maybe (Expr, Type))] -> Vector (Maybe (Expr, Type))
forall a. Vector a -> [(Int, a)] -> Vector a
V.// [(Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
i, (Expr, Type) -> Maybe (Expr, Type)
forall a. a -> Maybe a
Just (Expr
e', Type
t))]) Expr
e
        | Bool
otherwise -> (Expr -> Expr) -> ([Expr], Expr) -> ([Expr], Expr)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (\Expr
e -> Type -> Expr -> Expr -> Expr -> Expr
SetAt' Type
t Expr
e (Integer -> Expr
LitInt' Integer
i) Expr
e') (([Expr], Expr) -> ([Expr], Expr))
-> ([Expr], Expr) -> ([Expr], Expr)
forall a b. (a -> b) -> a -> b
$ Vector (Maybe (Expr, Type)) -> Expr -> ([Expr], Expr)
go Vector (Maybe (Expr, Type))
base Expr
e
      Expr
e ->
        let ([Maybe (Expr, Type)]
base', [Maybe (Expr, Type)]
base'') = (Maybe (Expr, Type) -> Bool)
-> [Maybe (Expr, Type)]
-> ([Maybe (Expr, Type)], [Maybe (Expr, Type)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span Maybe (Expr, Type) -> Bool
forall a. Maybe a -> Bool
isJust (Vector (Maybe (Expr, Type)) -> [Maybe (Expr, Type)]
forall a. Vector a -> [a]
V.toList Vector (Maybe (Expr, Type))
base)
            base''' :: [Expr]
base''' = (Maybe (Expr, Type) -> Expr) -> [Maybe (Expr, Type)] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map ((Expr, Type) -> Expr
forall a b. (a, b) -> a
fst ((Expr, Type) -> Expr)
-> (Maybe (Expr, Type) -> (Expr, Type))
-> Maybe (Expr, Type)
-> Expr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe (Expr, Type) -> (Expr, Type)
forall a. HasCallStack => Maybe a -> a
fromJust) [Maybe (Expr, Type)]
base'
            e'' :: Expr
e'' = ((Integer, Maybe (Expr, Type)) -> Expr -> Expr)
-> Expr -> [(Integer, Maybe (Expr, Type))] -> Expr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(Integer
i, Maybe (Expr, Type)
e') Expr
e -> (Expr -> Expr)
-> ((Expr, Type) -> Expr -> Expr)
-> Maybe (Expr, Type)
-> Expr
-> Expr
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Expr -> Expr
forall a. a -> a
id (\(Expr
e', Type
t) Expr
e -> Type -> Expr -> Expr -> Expr -> Expr
SetAt' Type
t Expr
e (Integer -> Expr
LitInt' Integer
i) Expr
e') Maybe (Expr, Type)
e' Expr
e) Expr
e ([Integer]
-> [Maybe (Expr, Type)] -> [(Integer, Maybe (Expr, Type))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Maybe (Expr, Type)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Maybe (Expr, Type)]
base') ..] [Maybe (Expr, Type)]
base'')
         in ([Expr]
base''', Expr
e'')

hoistMaybe :: Applicative m => Maybe a -> MaybeT m a
hoistMaybe :: Maybe a -> MaybeT m a
hoistMaybe = m (Maybe a) -> MaybeT m a
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (m (Maybe a) -> MaybeT m a)
-> (Maybe a -> m (Maybe a)) -> Maybe a -> MaybeT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe a -> m (Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure