{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}

module Jikka.RestrictedPython.Language.Util
  ( -- * generating symbols
    genType,
    genVarName,
    genVarName',

    -- * free variables
    freeTyVars,
    freeVars,
    freeVars',
    freeVarsTarget,
    freeVarsTarget',

    -- * return-statements
    doesAlwaysReturn,
    doesPossiblyReturn,

    -- * traversing statements
    mapStatement,
    mapStatementM,
    mapLargeStatement,
    mapLargeStatementM,
    mapStatements,
    mapStatementsM,
    listStatements,

    -- * traversing sub exprs
    mapSubExprM,
    mapSubExpr,
    listSubExprs,

    -- * traversing exprs
    mapExprTargetM,
    mapExprStatementM,
    mapExprM,
    listExprs,

    -- * exprs
    hasFunctionCall,
    isSmallExpr,
    dropLocation,

    -- * targets
    targetVars,
    targetVars',
    hasSubscriptTrg,
    hasBareNameTrg,
    exprToTarget,
    targetToExpr,

    -- * programs
    toplevelMainDef,
  )
where

import Control.Monad.Identity
import Control.Monad.Writer.Strict
import Data.List (delete, nub)
import Jikka.Common.Alpha
import Jikka.Common.Location
import Jikka.RestrictedPython.Language.Expr

genType :: MonadAlpha m => m Type
genType :: m Type
genType = do
  Int
i <- m Int
forall (m :: * -> *). MonadAlpha m => m Int
nextCounter
  Type -> m Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ TypeName -> Type
VarTy (String -> TypeName
TypeName (Char
'$' Char -> String -> String
forall a. a -> [a] -> [a]
: Int -> String
forall a. Show a => a -> String
show Int
i))

genVarName :: MonadAlpha m => VarName' -> m VarName'
genVarName :: VarName' -> m VarName'
genVarName VarName'
x = do
  Int
i <- m Int
forall (m :: * -> *). MonadAlpha m => m Int
nextCounter
  let base :: String
base = if VarName -> String
unVarName (VarName' -> VarName
forall a. WithLoc' a -> a
value' VarName'
x) String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"_" then String
"" else (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'$') (VarName -> String
unVarName (VarName' -> VarName
forall a. WithLoc' a -> a
value' VarName'
x))
  VarName' -> m VarName'
forall (m :: * -> *) a. Monad m => a -> m a
return (VarName' -> m VarName') -> VarName' -> m VarName'
forall a b. (a -> b) -> a -> b
$ Maybe Loc -> VarName -> VarName'
forall a. Maybe Loc -> a -> WithLoc' a
WithLoc' (VarName' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' VarName'
x) (String -> VarName
VarName (String
base String -> String -> String
forall a. [a] -> [a] -> [a]
++ Char
'$' Char -> String -> String
forall a. a -> [a] -> [a]
: Int -> String
forall a. Show a => a -> String
show Int
i))

genVarName' :: MonadAlpha m => m VarName'
genVarName' :: m VarName'
genVarName' = VarName' -> m VarName'
forall (m :: * -> *). MonadAlpha m => VarName' -> m VarName'
genVarName (VarName -> VarName'
forall a. a -> WithLoc' a
withoutLoc (String -> VarName
VarName String
"_"))

freeTyVars :: Type -> [TypeName]
freeTyVars :: Type -> [TypeName]
freeTyVars = [TypeName] -> [TypeName]
forall a. Eq a => [a] -> [a]
nub ([TypeName] -> [TypeName])
-> (Type -> [TypeName]) -> Type -> [TypeName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> [TypeName]
go
  where
    go :: Type -> [TypeName]
go = \case
      VarTy TypeName
x -> [TypeName
x]
      Type
IntTy -> []
      Type
BoolTy -> []
      ListTy Type
t -> Type -> [TypeName]
go Type
t
      TupleTy [Type]
ts -> [[TypeName]] -> [TypeName]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[TypeName]] -> [TypeName]) -> [[TypeName]] -> [TypeName]
forall a b. (a -> b) -> a -> b
$ (Type -> [TypeName]) -> [Type] -> [[TypeName]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Type -> [TypeName]
go [Type]
ts
      CallableTy [Type]
ts Type
ret -> [[TypeName]] -> [TypeName]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[TypeName]] -> [TypeName]) -> [[TypeName]] -> [TypeName]
forall a b. (a -> b) -> a -> b
$ (Type -> [TypeName]) -> [Type] -> [[TypeName]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Type -> [TypeName]
go (Type
ret Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: [Type]
ts)
      Type
StringTy -> []
      Type
SideEffectTy -> []

-- | `freeVars'` reports all free variables.
freeVars :: Expr' -> [VarName]
freeVars :: Expr' -> [VarName]
freeVars = [VarName] -> [VarName]
forall a. Eq a => [a] -> [a]
nub ([VarName] -> [VarName])
-> (Expr' -> [VarName]) -> Expr' -> [VarName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VarName' -> VarName) -> [VarName'] -> [VarName]
forall a b. (a -> b) -> [a] -> [b]
map VarName' -> VarName
forall a. WithLoc' a -> a
value' ([VarName'] -> [VarName])
-> (Expr' -> [VarName']) -> Expr' -> [VarName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr' -> [VarName']
freeVars'

-- | `freeVars'` reports all free variables with their locations, i.e. occurrences.
-- For examples, @x + x@ and @x@ have the same free variables @x@ but they have different sets of occurrences of free variable.
freeVars' :: Expr' -> [VarName']
freeVars' :: Expr' -> [VarName']
freeVars' (WithLoc' Maybe Loc
_ Expr
e0) = case Expr
e0 of
  BoolOp Expr'
e1 BoolOp
_ Expr'
e2 -> Expr' -> [VarName']
freeVars' Expr'
e1 [VarName'] -> [VarName'] -> [VarName']
forall a. [a] -> [a] -> [a]
++ Expr' -> [VarName']
freeVars' Expr'
e2
  BinOp Expr'
e1 Operator
_ Expr'
e2 -> Expr' -> [VarName']
freeVars' Expr'
e1 [VarName'] -> [VarName'] -> [VarName']
forall a. [a] -> [a] -> [a]
++ Expr' -> [VarName']
freeVars' Expr'
e2
  UnaryOp UnaryOp
_ Expr'
e -> Expr' -> [VarName']
freeVars' Expr'
e
  Lambda [(VarName', Type)]
args Expr'
e -> ([VarName'] -> (VarName', Type) -> [VarName'])
-> [VarName'] -> [(VarName', Type)] -> [VarName']
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\[VarName']
vars (VarName'
x, Type
_) -> VarName' -> [VarName'] -> [VarName']
forall a. Eq a => a -> [a] -> [a]
delete VarName'
x [VarName']
vars) (Expr' -> [VarName']
freeVars' Expr'
e) [(VarName', Type)]
args
  IfExp Expr'
e1 Expr'
e2 Expr'
e3 -> Expr' -> [VarName']
freeVars' Expr'
e1 [VarName'] -> [VarName'] -> [VarName']
forall a. [a] -> [a] -> [a]
++ Expr' -> [VarName']
freeVars' Expr'
e2 [VarName'] -> [VarName'] -> [VarName']
forall a. [a] -> [a] -> [a]
++ Expr' -> [VarName']
freeVars' Expr'
e3
  ListComp Expr'
e (Comprehension Target'
x Expr'
iter Maybe Expr'
pred) -> Expr' -> [VarName']
freeVars' Expr'
iter [VarName'] -> [VarName'] -> [VarName']
forall a. [a] -> [a] -> [a]
++ ([VarName'] -> VarName' -> [VarName'])
-> [VarName'] -> [VarName'] -> [VarName']
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\[VarName']
vars VarName'
x -> VarName' -> [VarName'] -> [VarName']
forall a. Eq a => a -> [a] -> [a]
delete VarName'
x [VarName']
vars) (Expr' -> [VarName']
freeVars' Expr'
e [VarName'] -> [VarName'] -> [VarName']
forall a. [a] -> [a] -> [a]
++ (Expr' -> [VarName']) -> Maybe Expr' -> [VarName']
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr' -> [VarName']
freeVars' Maybe Expr'
pred) (Target' -> [VarName']
targetVars' Target'
x)
  Compare Expr'
e1 CmpOp'
_ Expr'
e2 -> Expr' -> [VarName']
freeVars' Expr'
e1 [VarName'] -> [VarName'] -> [VarName']
forall a. [a] -> [a] -> [a]
++ Expr' -> [VarName']
freeVars' Expr'
e2
  Call Expr'
f [Expr']
args -> (Expr' -> [VarName']) -> [Expr'] -> [VarName']
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr' -> [VarName']
freeVars' (Expr'
f Expr' -> [Expr'] -> [Expr']
forall a. a -> [a] -> [a]
: [Expr']
args)
  Constant Constant
_ -> []
  Attribute Expr'
e Attribute'
_ -> Expr' -> [VarName']
freeVars' Expr'
e
  Subscript Expr'
e1 Expr'
e2 -> Expr' -> [VarName']
freeVars' Expr'
e1 [VarName'] -> [VarName'] -> [VarName']
forall a. [a] -> [a] -> [a]
++ Expr' -> [VarName']
freeVars' Expr'
e2
  Starred Expr'
e -> Expr' -> [VarName']
freeVars' Expr'
e
  Name VarName'
x -> [VarName'
x]
  List Type
_ [Expr']
es -> (Expr' -> [VarName']) -> [Expr'] -> [VarName']
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr' -> [VarName']
freeVars' [Expr']
es
  Tuple [Expr']
es -> (Expr' -> [VarName']) -> [Expr'] -> [VarName']
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr' -> [VarName']
freeVars' [Expr']
es
  SubscriptSlice Expr'
e Maybe Expr'
from Maybe Expr'
to Maybe Expr'
step -> Expr' -> [VarName']
freeVars' Expr'
e [VarName'] -> [VarName'] -> [VarName']
forall a. [a] -> [a] -> [a]
++ (Expr' -> [VarName']) -> Maybe Expr' -> [VarName']
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr' -> [VarName']
freeVars' Maybe Expr'
from [VarName'] -> [VarName'] -> [VarName']
forall a. [a] -> [a] -> [a]
++ (Expr' -> [VarName']) -> Maybe Expr' -> [VarName']
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr' -> [VarName']
freeVars' Maybe Expr'
to [VarName'] -> [VarName'] -> [VarName']
forall a. [a] -> [a] -> [a]
++ (Expr' -> [VarName']) -> Maybe Expr' -> [VarName']
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr' -> [VarName']
freeVars' Maybe Expr'
step

freeVarsTarget :: Target' -> [VarName]
freeVarsTarget :: Target' -> [VarName]
freeVarsTarget = [VarName] -> [VarName]
forall a. Eq a => [a] -> [a]
nub ([VarName] -> [VarName])
-> (Target' -> [VarName]) -> Target' -> [VarName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VarName' -> VarName) -> [VarName'] -> [VarName]
forall a b. (a -> b) -> [a] -> [b]
map VarName' -> VarName
forall a. WithLoc' a -> a
value' ([VarName'] -> [VarName])
-> (Target' -> [VarName']) -> Target' -> [VarName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Target' -> [VarName']
freeVarsTarget'

freeVarsTarget' :: Target' -> [VarName']
freeVarsTarget' :: Target' -> [VarName']
freeVarsTarget' (WithLoc' Maybe Loc
_ Target
x) = case Target
x of
  SubscriptTrg Target'
_ Expr'
e -> Expr' -> [VarName']
freeVars' Expr'
e
  NameTrg VarName'
_ -> []
  TupleTrg [Target']
xs -> (Target' -> [VarName']) -> [Target'] -> [VarName']
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Target' -> [VarName']
freeVarsTarget' [Target']
xs

doesAlwaysReturn :: Statement -> Bool
doesAlwaysReturn :: Statement -> Bool
doesAlwaysReturn = \case
  Return Expr'
_ -> Bool
True
  AugAssign Target'
_ Operator
_ Expr'
_ -> Bool
False
  AnnAssign Target'
_ Type
_ Expr'
_ -> Bool
False
  For Target'
_ Expr'
_ [Statement]
_ -> Bool
False
  If Expr'
_ [Statement]
body1 [Statement]
body2 -> (Statement -> Bool) -> [Statement] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Statement -> Bool
doesAlwaysReturn [Statement]
body1 Bool -> Bool -> Bool
&& (Statement -> Bool) -> [Statement] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Statement -> Bool
doesAlwaysReturn [Statement]
body2
  Assert Expr'
_ -> Bool
False
  Expr' Expr'
_ -> Bool
False

doesPossiblyReturn :: Statement -> Bool
doesPossiblyReturn :: Statement -> Bool
doesPossiblyReturn = \case
  Return Expr'
_ -> Bool
True
  AugAssign Target'
_ Operator
_ Expr'
_ -> Bool
False
  AnnAssign Target'
_ Type
_ Expr'
_ -> Bool
False
  For Target'
_ Expr'
_ [Statement]
body -> (Statement -> Bool) -> [Statement] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Statement -> Bool
doesPossiblyReturn [Statement]
body
  If Expr'
_ [Statement]
body1 [Statement]
body2 -> (Statement -> Bool) -> [Statement] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Statement -> Bool
doesPossiblyReturn [Statement]
body1 Bool -> Bool -> Bool
|| (Statement -> Bool) -> [Statement] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Statement -> Bool
doesPossiblyReturn [Statement]
body2
  Assert Expr'
_ -> Bool
False
  Expr' Expr'
_ -> Bool
False

-- | `mapSubExprM` replaces all exprs in a given expr using a given function.
-- This may breaks various constraints.
mapSubExprM :: Monad m => (Expr' -> m Expr') -> Expr' -> m Expr'
mapSubExprM :: (Expr' -> m Expr') -> Expr' -> m Expr'
mapSubExprM Expr' -> m Expr'
f = Expr' -> m Expr'
go
  where
    go :: Expr' -> m Expr'
go Expr'
e0 =
      Expr' -> m Expr'
f (Expr' -> m Expr') -> (Expr -> Expr') -> Expr -> m Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe Loc -> Expr -> Expr'
forall a. Maybe Loc -> a -> WithLoc' a
WithLoc' (Expr' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Expr'
e0) (Expr -> m Expr') -> m Expr -> m Expr'
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< case Expr' -> Expr
forall a. WithLoc' a -> a
value' Expr'
e0 of
        BoolOp Expr'
e1 BoolOp
op Expr'
e2 -> Expr' -> BoolOp -> Expr' -> Expr
BoolOp (Expr' -> BoolOp -> Expr' -> Expr)
-> m Expr' -> m (BoolOp -> Expr' -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
go Expr'
e1 m (BoolOp -> Expr' -> Expr) -> m BoolOp -> m (Expr' -> Expr)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> BoolOp -> m BoolOp
forall (m :: * -> *) a. Monad m => a -> m a
return BoolOp
op 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
        BinOp Expr'
e1 Operator
op Expr'
e2 -> Expr' -> Operator -> Expr' -> Expr
BinOp (Expr' -> Operator -> Expr' -> Expr)
-> m Expr' -> m (Operator -> Expr' -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
go Expr'
e1 m (Operator -> Expr' -> Expr) -> m Operator -> m (Expr' -> Expr)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Operator -> m Operator
forall (m :: * -> *) a. Monad m => a -> m a
return Operator
op 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
        UnaryOp UnaryOp
op Expr'
e -> UnaryOp -> Expr' -> Expr
UnaryOp UnaryOp
op (Expr' -> Expr) -> m Expr' -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
go Expr'
e
        Lambda [(VarName', Type)]
args Expr'
body -> [(VarName', Type)] -> Expr' -> Expr
Lambda [(VarName', Type)]
args (Expr' -> Expr) -> m Expr' -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
go Expr'
body
        IfExp Expr'
e1 Expr'
e2 Expr'
e3 -> Expr' -> Expr' -> Expr' -> Expr
IfExp (Expr' -> Expr' -> Expr' -> Expr)
-> m Expr' -> m (Expr' -> Expr' -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
go Expr'
e1 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'
e2 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'
e3
        ListComp Expr'
e (Comprehension Target'
x Expr'
iter Maybe Expr'
pred) -> do
          Expr'
e <- Expr' -> m Expr'
go Expr'
e
          Target'
x <- (Expr' -> m Expr') -> Target' -> m Target'
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> Target' -> m Target'
mapExprTargetM Expr' -> m Expr'
f Target'
x
          Expr'
iter <- Expr' -> m Expr'
go Expr'
iter
          Maybe Expr'
pred <- (Expr' -> m Expr') -> Maybe Expr' -> m (Maybe Expr')
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Expr' -> m Expr'
go Maybe Expr'
pred
          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' -> Comprehension -> Expr
ListComp Expr'
e (Target' -> Expr' -> Maybe Expr' -> Comprehension
Comprehension Target'
x Expr'
iter Maybe Expr'
pred)
        Compare Expr'
e1 CmpOp'
op Expr'
e2 -> Expr' -> CmpOp' -> Expr' -> Expr
Compare (Expr' -> CmpOp' -> Expr' -> Expr)
-> m Expr' -> m (CmpOp' -> Expr' -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
go Expr'
e1 m (CmpOp' -> Expr' -> Expr) -> m CmpOp' -> m (Expr' -> Expr)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> CmpOp' -> m CmpOp'
forall (m :: * -> *) a. Monad m => a -> m a
return CmpOp'
op 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
        Call Expr'
g [Expr']
args -> Expr' -> [Expr'] -> Expr
Call (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') -> [Expr'] -> m [Expr']
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Expr' -> m Expr'
go [Expr']
args
        Constant Constant
const -> 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
$ Constant -> Expr
Constant Constant
const
        Attribute Expr'
e Attribute'
x -> Expr' -> Attribute' -> Expr
Attribute (Expr' -> Attribute' -> Expr) -> m Expr' -> m (Attribute' -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
go Expr'
e m (Attribute' -> Expr) -> m Attribute' -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Attribute' -> m Attribute'
forall (f :: * -> *) a. Applicative f => a -> f a
pure Attribute'
x
        Subscript Expr'
e1 Expr'
e2 -> Expr' -> Expr' -> Expr
Subscript (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
        Starred Expr'
e -> Expr' -> Expr
Starred (Expr' -> Expr) -> m Expr' -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
go Expr'
e
        Name 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
Name VarName'
x
        List Type
t [Expr']
es -> Type -> [Expr'] -> Expr
List Type
t ([Expr'] -> Expr) -> m [Expr'] -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Expr' -> m Expr') -> [Expr'] -> m [Expr']
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Expr' -> m Expr'
go [Expr']
es
        Tuple [Expr']
es -> [Expr'] -> Expr
Tuple ([Expr'] -> Expr) -> m [Expr'] -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Expr' -> m Expr') -> [Expr'] -> m [Expr']
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Expr' -> m Expr'
go [Expr']
es
        SubscriptSlice Expr'
e Maybe Expr'
from Maybe Expr'
to Maybe Expr'
step -> Expr' -> Maybe Expr' -> Maybe Expr' -> Maybe Expr' -> Expr
SubscriptSlice (Expr' -> Maybe Expr' -> Maybe Expr' -> Maybe Expr' -> Expr)
-> m Expr' -> m (Maybe Expr' -> Maybe Expr' -> Maybe Expr' -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
go Expr'
e m (Maybe Expr' -> Maybe Expr' -> Maybe Expr' -> Expr)
-> m (Maybe Expr') -> m (Maybe Expr' -> Maybe Expr' -> Expr)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Expr' -> m Expr') -> Maybe Expr' -> m (Maybe Expr')
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Expr' -> m Expr'
go Maybe Expr'
from m (Maybe Expr' -> Maybe Expr' -> Expr)
-> m (Maybe Expr') -> m (Maybe Expr' -> Expr)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Expr' -> m Expr') -> Maybe Expr' -> m (Maybe Expr')
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Expr' -> m Expr'
go Maybe Expr'
to m (Maybe Expr' -> Expr) -> m (Maybe Expr') -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Expr' -> m Expr') -> Maybe Expr' -> m (Maybe Expr')
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Expr' -> m Expr'
go Maybe Expr'
step

mapSubExpr :: (Expr' -> Expr') -> Expr' -> Expr'
mapSubExpr :: (Expr' -> Expr') -> Expr' -> Expr'
mapSubExpr Expr' -> Expr'
f = Identity Expr' -> Expr'
forall a. Identity a -> a
runIdentity (Identity Expr' -> Expr')
-> (Expr' -> Identity Expr') -> Expr' -> Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Expr' -> Identity Expr') -> Expr' -> Identity Expr'
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> Expr' -> m Expr'
mapSubExprM (Expr' -> Identity Expr'
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr' -> Identity Expr')
-> (Expr' -> Expr') -> Expr' -> Identity Expr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr' -> Expr'
f)

listSubExprs :: Expr' -> [Expr']
listSubExprs :: Expr' -> [Expr']
listSubExprs = [Expr'] -> [Expr']
forall a. [a] -> [a]
reverse ([Expr'] -> [Expr']) -> (Expr' -> [Expr']) -> Expr' -> [Expr']
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dual [Expr'] -> [Expr']
forall a. Dual a -> a
getDual (Dual [Expr'] -> [Expr'])
-> (Expr' -> 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' -> Dual [Expr'])
-> (Expr' -> Writer (Dual [Expr']) Expr') -> Expr' -> Dual [Expr']
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Expr' -> Writer (Dual [Expr']) Expr')
-> Expr' -> Writer (Dual [Expr']) Expr'
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> Expr' -> m Expr'
mapSubExprM Expr' -> Writer (Dual [Expr']) Expr'
forall (m :: * -> *) b. MonadWriter (Dual [b]) m => b -> m b
go
  where
    go :: b -> m b
go 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

mapExprTargetM :: Monad m => (Expr' -> m Expr') -> Target' -> m Target'
mapExprTargetM :: (Expr' -> m Expr') -> Target' -> m Target'
mapExprTargetM Expr' -> m Expr'
f Target'
x =
  Maybe Loc -> Target -> Target'
forall a. Maybe Loc -> a -> WithLoc' a
WithLoc' (Target' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Target'
x) (Target -> Target') -> m Target -> m Target'
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case Target' -> Target
forall a. WithLoc' a -> a
value' Target'
x of
    SubscriptTrg Target'
x Expr'
e -> Target' -> Expr' -> Target
SubscriptTrg (Target' -> Expr' -> Target) -> m Target' -> m (Expr' -> Target)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Expr' -> m Expr') -> Target' -> m Target'
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> Target' -> m Target'
mapExprTargetM Expr' -> m Expr'
f Target'
x m (Expr' -> Target) -> m Expr' -> m Target
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr' -> m Expr'
f Expr'
e
    NameTrg VarName'
x -> Target -> m Target
forall (m :: * -> *) a. Monad m => a -> m a
return (Target -> m Target) -> Target -> m Target
forall a b. (a -> b) -> a -> b
$ VarName' -> Target
NameTrg VarName'
x
    TupleTrg [Target']
xs -> [Target'] -> Target
TupleTrg ([Target'] -> Target) -> m [Target'] -> m Target
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Target' -> m Target') -> [Target'] -> m [Target']
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Expr' -> m Expr') -> Target' -> m Target'
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> Target' -> m Target'
mapExprTargetM Expr' -> m Expr'
f) [Target']
xs

mapExprStatementM :: Monad m => (Expr' -> m Expr') -> Statement -> m Statement
mapExprStatementM :: (Expr' -> m Expr') -> Statement -> m Statement
mapExprStatementM Expr' -> m Expr'
f = \case
  Return Expr'
e -> Expr' -> Statement
Return (Expr' -> Statement) -> m Expr' -> m Statement
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
f Expr'
e
  AugAssign Target'
x Operator
op Expr'
e -> Target' -> Operator -> Expr' -> Statement
AugAssign (Target' -> Operator -> Expr' -> Statement)
-> m Target' -> m (Operator -> Expr' -> Statement)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Expr' -> m Expr') -> Target' -> m Target'
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> Target' -> m Target'
mapExprTargetM Expr' -> m Expr'
f Target'
x m (Operator -> Expr' -> Statement)
-> m Operator -> m (Expr' -> Statement)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Operator -> m Operator
forall (f :: * -> *) a. Applicative f => a -> f a
pure Operator
op m (Expr' -> Statement) -> m Expr' -> m Statement
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr' -> m Expr'
f Expr'
e
  AnnAssign Target'
x Type
t Expr'
e -> Target' -> Type -> Expr' -> Statement
AnnAssign (Target' -> Type -> Expr' -> Statement)
-> m Target' -> m (Type -> Expr' -> Statement)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Expr' -> m Expr') -> Target' -> m Target'
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> Target' -> m Target'
mapExprTargetM Expr' -> m Expr'
f Target'
x m (Type -> Expr' -> Statement) -> m Type -> m (Expr' -> Statement)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
t m (Expr' -> Statement) -> m Expr' -> m Statement
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr' -> m Expr'
f Expr'
e
  For Target'
x Expr'
iter [Statement]
body -> Target' -> Expr' -> [Statement] -> Statement
For (Target' -> Expr' -> [Statement] -> Statement)
-> m Target' -> m (Expr' -> [Statement] -> Statement)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Expr' -> m Expr') -> Target' -> m Target'
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> Target' -> m Target'
mapExprTargetM Expr' -> m Expr'
f Target'
x m (Expr' -> [Statement] -> Statement)
-> m Expr' -> m ([Statement] -> Statement)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr' -> m Expr'
f Expr'
iter m ([Statement] -> Statement) -> m [Statement] -> m Statement
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Statement -> m Statement) -> [Statement] -> m [Statement]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Expr' -> m Expr') -> Statement -> m Statement
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> Statement -> m Statement
mapExprStatementM Expr' -> m Expr'
f) [Statement]
body
  If Expr'
e [Statement]
body1 [Statement]
body2 -> Expr' -> [Statement] -> [Statement] -> Statement
If (Expr' -> [Statement] -> [Statement] -> Statement)
-> m Expr' -> m ([Statement] -> [Statement] -> Statement)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
f Expr'
e m ([Statement] -> [Statement] -> Statement)
-> m [Statement] -> m ([Statement] -> Statement)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Statement -> m Statement) -> [Statement] -> m [Statement]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Expr' -> m Expr') -> Statement -> m Statement
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> Statement -> m Statement
mapExprStatementM Expr' -> m Expr'
f) [Statement]
body1 m ([Statement] -> Statement) -> m [Statement] -> m Statement
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Statement -> m Statement) -> [Statement] -> m [Statement]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Expr' -> m Expr') -> Statement -> m Statement
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> Statement -> m Statement
mapExprStatementM Expr' -> m Expr'
f) [Statement]
body2
  Assert Expr'
e -> Expr' -> Statement
Assert (Expr' -> Statement) -> m Expr' -> m Statement
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
f Expr'
e
  Expr' Expr'
e -> Expr' -> Statement
Expr' (Expr' -> Statement) -> m Expr' -> m Statement
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
f Expr'
e

mapExprToplevelStatementM :: Monad m => (Expr' -> m Expr') -> ToplevelStatement -> m ToplevelStatement
mapExprToplevelStatementM :: (Expr' -> m Expr') -> ToplevelStatement -> m ToplevelStatement
mapExprToplevelStatementM Expr' -> m Expr'
f = \case
  ToplevelAnnAssign VarName'
x Type
t Expr'
e -> VarName' -> Type -> Expr' -> ToplevelStatement
ToplevelAnnAssign VarName'
x Type
t (Expr' -> ToplevelStatement) -> m Expr' -> m ToplevelStatement
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
f Expr'
e
  ToplevelFunctionDef VarName'
g [(VarName', Type)]
args Type
ret [Statement]
body -> VarName'
-> [(VarName', Type)] -> Type -> [Statement] -> ToplevelStatement
ToplevelFunctionDef VarName'
g [(VarName', Type)]
args Type
ret ([Statement] -> ToplevelStatement)
-> m [Statement] -> m ToplevelStatement
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Statement -> m Statement) -> [Statement] -> m [Statement]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Expr' -> m Expr') -> Statement -> m Statement
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> Statement -> m Statement
mapExprStatementM Expr' -> m Expr'
f) [Statement]
body
  ToplevelAssert Expr'
e -> Expr' -> ToplevelStatement
ToplevelAssert (Expr' -> ToplevelStatement) -> m Expr' -> m ToplevelStatement
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> m Expr'
f Expr'
e

mapExprM :: Monad m => (Expr' -> m Expr') -> Program -> m Program
mapExprM :: (Expr' -> m Expr') -> Program -> m Program
mapExprM Expr' -> m Expr'
f = (ToplevelStatement -> m ToplevelStatement) -> Program -> m Program
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Expr' -> m Expr') -> ToplevelStatement -> m ToplevelStatement
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> ToplevelStatement -> m ToplevelStatement
mapExprToplevelStatementM Expr' -> m Expr'
f)

listExprs :: Program -> [Expr']
listExprs :: Program -> [Expr']
listExprs = [Expr'] -> [Expr']
forall a. [a] -> [a]
reverse ([Expr'] -> [Expr']) -> (Program -> [Expr']) -> Program -> [Expr']
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dual [Expr'] -> [Expr']
forall a. Dual a -> a
getDual (Dual [Expr'] -> [Expr'])
-> (Program -> Dual [Expr']) -> Program -> [Expr']
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Writer (Dual [Expr']) Program -> Dual [Expr']
forall w a. Writer w a -> w
execWriter (Writer (Dual [Expr']) Program -> Dual [Expr'])
-> (Program -> Writer (Dual [Expr']) Program)
-> Program
-> Dual [Expr']
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Expr' -> Writer (Dual [Expr']) Expr')
-> Program -> Writer (Dual [Expr']) Program
forall (m :: * -> *).
Monad m =>
(Expr' -> m Expr') -> Program -> m Program
mapExprM Expr' -> Writer (Dual [Expr']) Expr'
forall (m :: * -> *) b. MonadWriter (Dual [b]) m => b -> m b
go
  where
    go :: b -> m b
go 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

mapStatementStatementM :: Monad m => (Statement -> m [Statement]) -> Statement -> m [Statement]
mapStatementStatementM :: (Statement -> m [Statement]) -> Statement -> m [Statement]
mapStatementStatementM Statement -> m [Statement]
f = \case
  Return Expr'
e -> Statement -> m [Statement]
f (Statement -> m [Statement]) -> Statement -> m [Statement]
forall a b. (a -> b) -> a -> b
$ Expr' -> Statement
Return Expr'
e
  AugAssign Target'
x Operator
op Expr'
e -> Statement -> m [Statement]
f (Statement -> m [Statement]) -> Statement -> m [Statement]
forall a b. (a -> b) -> a -> b
$ Target' -> Operator -> Expr' -> Statement
AugAssign Target'
x Operator
op Expr'
e
  AnnAssign Target'
x Type
t Expr'
e -> Statement -> m [Statement]
f (Statement -> m [Statement]) -> Statement -> m [Statement]
forall a b. (a -> b) -> a -> b
$ Target' -> Type -> Expr' -> Statement
AnnAssign Target'
x Type
t Expr'
e
  For Target'
x Expr'
iter [Statement]
body -> do
    [Statement]
body <- [[Statement]] -> [Statement]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Statement]] -> [Statement]) -> m [[Statement]] -> m [Statement]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Statement -> m [Statement]) -> [Statement] -> m [[Statement]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Statement -> m [Statement]) -> Statement -> m [Statement]
forall (m :: * -> *).
Monad m =>
(Statement -> m [Statement]) -> Statement -> m [Statement]
mapStatementStatementM Statement -> m [Statement]
f) [Statement]
body
    Statement -> m [Statement]
f (Statement -> m [Statement]) -> Statement -> m [Statement]
forall a b. (a -> b) -> a -> b
$ Target' -> Expr' -> [Statement] -> Statement
For Target'
x Expr'
iter [Statement]
body
  If Expr'
e [Statement]
body1 [Statement]
body2 -> do
    [Statement]
body1 <- [[Statement]] -> [Statement]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Statement]] -> [Statement]) -> m [[Statement]] -> m [Statement]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Statement -> m [Statement]) -> [Statement] -> m [[Statement]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Statement -> m [Statement]) -> Statement -> m [Statement]
forall (m :: * -> *).
Monad m =>
(Statement -> m [Statement]) -> Statement -> m [Statement]
mapStatementStatementM Statement -> m [Statement]
f) [Statement]
body1
    [Statement]
body2 <- [[Statement]] -> [Statement]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Statement]] -> [Statement]) -> m [[Statement]] -> m [Statement]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Statement -> m [Statement]) -> [Statement] -> m [[Statement]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Statement -> m [Statement]) -> Statement -> m [Statement]
forall (m :: * -> *).
Monad m =>
(Statement -> m [Statement]) -> Statement -> m [Statement]
mapStatementStatementM Statement -> m [Statement]
f) [Statement]
body2
    Statement -> m [Statement]
f (Statement -> m [Statement]) -> Statement -> m [Statement]
forall a b. (a -> b) -> a -> b
$ Expr' -> [Statement] -> [Statement] -> Statement
If Expr'
e [Statement]
body1 [Statement]
body2
  Assert Expr'
e -> Statement -> m [Statement]
f (Statement -> m [Statement]) -> Statement -> m [Statement]
forall a b. (a -> b) -> a -> b
$ Expr' -> Statement
Assert Expr'
e
  Expr' Expr'
e -> Statement -> m [Statement]
f (Statement -> m [Statement]) -> Statement -> m [Statement]
forall a b. (a -> b) -> a -> b
$ Expr' -> Statement
Expr' Expr'
e

mapStatementToplevelStatementM :: Monad m => (Statement -> m [Statement]) -> ToplevelStatement -> m ToplevelStatement
mapStatementToplevelStatementM :: (Statement -> m [Statement])
-> ToplevelStatement -> m ToplevelStatement
mapStatementToplevelStatementM Statement -> m [Statement]
go = \case
  ToplevelAnnAssign VarName'
x Type
t Expr'
e -> ToplevelStatement -> m ToplevelStatement
forall (m :: * -> *) a. Monad m => a -> m a
return (ToplevelStatement -> m ToplevelStatement)
-> ToplevelStatement -> m ToplevelStatement
forall a b. (a -> b) -> a -> b
$ VarName' -> Type -> Expr' -> ToplevelStatement
ToplevelAnnAssign VarName'
x Type
t Expr'
e
  ToplevelFunctionDef VarName'
f [(VarName', Type)]
args Type
ret [Statement]
body -> do
    [Statement]
body <- [[Statement]] -> [Statement]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Statement]] -> [Statement]) -> m [[Statement]] -> m [Statement]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Statement -> m [Statement]) -> [Statement] -> m [[Statement]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Statement -> m [Statement]) -> Statement -> m [Statement]
forall (m :: * -> *).
Monad m =>
(Statement -> m [Statement]) -> Statement -> m [Statement]
mapStatementStatementM Statement -> m [Statement]
go) [Statement]
body
    ToplevelStatement -> m ToplevelStatement
forall (m :: * -> *) a. Monad m => a -> m a
return (ToplevelStatement -> m ToplevelStatement)
-> ToplevelStatement -> m ToplevelStatement
forall a b. (a -> b) -> a -> b
$ VarName'
-> [(VarName', Type)] -> Type -> [Statement] -> ToplevelStatement
ToplevelFunctionDef VarName'
f [(VarName', Type)]
args Type
ret [Statement]
body
  ToplevelAssert Expr'
e -> ToplevelStatement -> m ToplevelStatement
forall (m :: * -> *) a. Monad m => a -> m a
return (ToplevelStatement -> m ToplevelStatement)
-> ToplevelStatement -> m ToplevelStatement
forall a b. (a -> b) -> a -> b
$ Expr' -> ToplevelStatement
ToplevelAssert Expr'
e

-- | `mapStatementM` replaces all statements in a given program using a given function.
-- This may breaks various constraints.
mapStatementM :: Monad m => (Statement -> m [Statement]) -> Program -> m Program
mapStatementM :: (Statement -> m [Statement]) -> Program -> m Program
mapStatementM Statement -> m [Statement]
f = (ToplevelStatement -> m ToplevelStatement) -> Program -> m Program
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Statement -> m [Statement])
-> ToplevelStatement -> m ToplevelStatement
forall (m :: * -> *).
Monad m =>
(Statement -> m [Statement])
-> ToplevelStatement -> m ToplevelStatement
mapStatementToplevelStatementM Statement -> m [Statement]
f)

mapStatement :: (Statement -> [Statement]) -> Program -> Program
mapStatement :: (Statement -> [Statement]) -> Program -> Program
mapStatement Statement -> [Statement]
f = Identity Program -> Program
forall a. Identity a -> a
runIdentity (Identity Program -> Program)
-> (Program -> Identity Program) -> Program -> Program
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Statement -> Identity [Statement]) -> Program -> Identity Program
forall (m :: * -> *).
Monad m =>
(Statement -> m [Statement]) -> Program -> m Program
mapStatementM ([Statement] -> Identity [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Statement] -> Identity [Statement])
-> (Statement -> [Statement]) -> Statement -> Identity [Statement]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Statement -> [Statement]
f)

mapLargeStatementM :: Monad m => (Expr' -> [Statement] -> [Statement] -> m [Statement]) -> (Target' -> Expr' -> [Statement] -> m [Statement]) -> Program -> m Program
mapLargeStatementM :: (Expr' -> [Statement] -> [Statement] -> m [Statement])
-> (Target' -> Expr' -> [Statement] -> m [Statement])
-> Program
-> m Program
mapLargeStatementM Expr' -> [Statement] -> [Statement] -> m [Statement]
fIf Target' -> Expr' -> [Statement] -> m [Statement]
fFor = (Statement -> m [Statement]) -> Program -> m Program
forall (m :: * -> *).
Monad m =>
(Statement -> m [Statement]) -> Program -> m Program
mapStatementM Statement -> m [Statement]
go
  where
    go :: Statement -> m [Statement]
go = \case
      Return Expr'
e -> [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return [Expr' -> Statement
Return Expr'
e]
      AugAssign Target'
x Operator
op Expr'
e -> [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return [Target' -> Operator -> Expr' -> Statement
AugAssign Target'
x Operator
op Expr'
e]
      AnnAssign Target'
x Type
t Expr'
e -> [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return [Target' -> Type -> Expr' -> Statement
AnnAssign Target'
x Type
t Expr'
e]
      For Target'
x Expr'
iter [Statement]
body -> Target' -> Expr' -> [Statement] -> m [Statement]
fFor Target'
x Expr'
iter [Statement]
body
      If Expr'
e [Statement]
body1 [Statement]
body2 -> Expr' -> [Statement] -> [Statement] -> m [Statement]
fIf Expr'
e [Statement]
body1 [Statement]
body2
      Assert Expr'
e -> [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return [Expr' -> Statement
Assert Expr'
e]
      Expr' Expr'
e -> [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return [Expr' -> Statement
Expr' Expr'
e]

mapLargeStatement :: (Expr' -> [Statement] -> [Statement] -> [Statement]) -> (Target' -> Expr' -> [Statement] -> [Statement]) -> Program -> Program
mapLargeStatement :: (Expr' -> [Statement] -> [Statement] -> [Statement])
-> (Target' -> Expr' -> [Statement] -> [Statement])
-> Program
-> Program
mapLargeStatement Expr' -> [Statement] -> [Statement] -> [Statement]
fIf Target' -> Expr' -> [Statement] -> [Statement]
fFor = Identity Program -> Program
forall a. Identity a -> a
runIdentity (Identity Program -> Program)
-> (Program -> Identity Program) -> Program -> Program
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Expr' -> [Statement] -> [Statement] -> Identity [Statement])
-> (Target' -> Expr' -> [Statement] -> Identity [Statement])
-> Program
-> Identity Program
forall (m :: * -> *).
Monad m =>
(Expr' -> [Statement] -> [Statement] -> m [Statement])
-> (Target' -> Expr' -> [Statement] -> m [Statement])
-> Program
-> m Program
mapLargeStatementM Expr' -> [Statement] -> [Statement] -> Identity [Statement]
forall (m :: * -> *).
Monad m =>
Expr' -> [Statement] -> [Statement] -> m [Statement]
fIf' Target' -> Expr' -> [Statement] -> Identity [Statement]
forall (m :: * -> *).
Monad m =>
Target' -> Expr' -> [Statement] -> m [Statement]
fFor'
  where
    fIf' :: Expr' -> [Statement] -> [Statement] -> m [Statement]
fIf' Expr'
e [Statement]
body1 [Statement]
body2 = [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Statement] -> m [Statement]) -> [Statement] -> m [Statement]
forall a b. (a -> b) -> a -> b
$ Expr' -> [Statement] -> [Statement] -> [Statement]
fIf Expr'
e [Statement]
body1 [Statement]
body2
    fFor' :: Target' -> Expr' -> [Statement] -> m [Statement]
fFor' Target'
x Expr'
iter [Statement]
body = [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Statement] -> m [Statement]) -> [Statement] -> m [Statement]
forall a b. (a -> b) -> a -> b
$ Target' -> Expr' -> [Statement] -> [Statement]
fFor Target'
x Expr'
iter [Statement]
body

listStatements :: Program -> [Statement]
listStatements :: Program -> [Statement]
listStatements = [Statement] -> [Statement]
forall a. [a] -> [a]
reverse ([Statement] -> [Statement])
-> (Program -> [Statement]) -> Program -> [Statement]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dual [Statement] -> [Statement]
forall a. Dual a -> a
getDual (Dual [Statement] -> [Statement])
-> (Program -> Dual [Statement]) -> Program -> [Statement]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Writer (Dual [Statement]) Program -> Dual [Statement]
forall w a. Writer w a -> w
execWriter (Writer (Dual [Statement]) Program -> Dual [Statement])
-> (Program -> Writer (Dual [Statement]) Program)
-> Program
-> Dual [Statement]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Statement -> WriterT (Dual [Statement]) Identity [Statement])
-> Program -> Writer (Dual [Statement]) Program
forall (m :: * -> *).
Monad m =>
(Statement -> m [Statement]) -> Program -> m Program
mapStatementM Statement -> WriterT (Dual [Statement]) Identity [Statement]
forall (m :: * -> *) a. MonadWriter (Dual [a]) m => a -> m [a]
go
  where
    go :: a -> m [a]
go a
stmt = do
      Dual [a] -> m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Dual [a] -> m ()) -> Dual [a] -> m ()
forall a b. (a -> b) -> a -> b
$ [a] -> Dual [a]
forall a. a -> Dual a
Dual [a
stmt]
      [a] -> m [a]
forall (m :: * -> *) a. Monad m => a -> m a
return [a
stmt]

mapStatementsToplevelStatementM :: Monad m => ([Statement] -> m [Statement]) -> ToplevelStatement -> m ToplevelStatement
mapStatementsToplevelStatementM :: ([Statement] -> m [Statement])
-> ToplevelStatement -> m ToplevelStatement
mapStatementsToplevelStatementM [Statement] -> m [Statement]
go = \case
  ToplevelAnnAssign VarName'
x Type
t Expr'
e -> ToplevelStatement -> m ToplevelStatement
forall (m :: * -> *) a. Monad m => a -> m a
return (ToplevelStatement -> m ToplevelStatement)
-> ToplevelStatement -> m ToplevelStatement
forall a b. (a -> b) -> a -> b
$ VarName' -> Type -> Expr' -> ToplevelStatement
ToplevelAnnAssign VarName'
x Type
t Expr'
e
  ToplevelFunctionDef VarName'
f [(VarName', Type)]
args Type
ret [Statement]
body -> do
    let go' :: Statement -> m [Statement]
go' = \case
          Return Expr'
e -> [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return [Expr' -> Statement
Return Expr'
e]
          AugAssign Target'
x Operator
op Expr'
e -> [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return [Target' -> Operator -> Expr' -> Statement
AugAssign Target'
x Operator
op Expr'
e]
          AnnAssign Target'
x Type
t Expr'
e -> [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return [Target' -> Type -> Expr' -> Statement
AnnAssign Target'
x Type
t Expr'
e]
          For Target'
x Expr'
iter [Statement]
body -> do
            [Statement]
body <- [Statement] -> m [Statement]
go [Statement]
body
            [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return [Target' -> Expr' -> [Statement] -> Statement
For Target'
x Expr'
iter [Statement]
body]
          If Expr'
e [Statement]
body1 [Statement]
body2 -> do
            [Statement]
body1 <- [Statement] -> m [Statement]
go [Statement]
body1
            [Statement]
body2 <- [Statement] -> m [Statement]
go [Statement]
body2
            [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return [Expr' -> [Statement] -> [Statement] -> Statement
If Expr'
e [Statement]
body1 [Statement]
body2]
          Assert Expr'
e -> [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return [Expr' -> Statement
Assert Expr'
e]
          Expr' Expr'
e -> [Statement] -> m [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return [Expr' -> Statement
Expr' Expr'
e]
    [Statement]
body <- [[Statement]] -> [Statement]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Statement]] -> [Statement]) -> m [[Statement]] -> m [Statement]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Statement -> m [Statement]) -> [Statement] -> m [[Statement]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Statement -> m [Statement]) -> Statement -> m [Statement]
forall (m :: * -> *).
Monad m =>
(Statement -> m [Statement]) -> Statement -> m [Statement]
mapStatementStatementM Statement -> m [Statement]
go') [Statement]
body
    [Statement]
body <- [Statement] -> m [Statement]
go [Statement]
body
    ToplevelStatement -> m ToplevelStatement
forall (m :: * -> *) a. Monad m => a -> m a
return (ToplevelStatement -> m ToplevelStatement)
-> ToplevelStatement -> m ToplevelStatement
forall a b. (a -> b) -> a -> b
$ VarName'
-> [(VarName', Type)] -> Type -> [Statement] -> ToplevelStatement
ToplevelFunctionDef VarName'
f [(VarName', Type)]
args Type
ret [Statement]
body
  ToplevelAssert Expr'
e -> ToplevelStatement -> m ToplevelStatement
forall (m :: * -> *) a. Monad m => a -> m a
return (ToplevelStatement -> m ToplevelStatement)
-> ToplevelStatement -> m ToplevelStatement
forall a b. (a -> b) -> a -> b
$ Expr' -> ToplevelStatement
ToplevelAssert Expr'
e

mapStatementsM :: Monad m => ([Statement] -> m [Statement]) -> Program -> m Program
mapStatementsM :: ([Statement] -> m [Statement]) -> Program -> m Program
mapStatementsM [Statement] -> m [Statement]
f = (ToplevelStatement -> m ToplevelStatement) -> Program -> m Program
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (([Statement] -> m [Statement])
-> ToplevelStatement -> m ToplevelStatement
forall (m :: * -> *).
Monad m =>
([Statement] -> m [Statement])
-> ToplevelStatement -> m ToplevelStatement
mapStatementsToplevelStatementM [Statement] -> m [Statement]
f)

mapStatements :: ([Statement] -> [Statement]) -> Program -> Program
mapStatements :: ([Statement] -> [Statement]) -> Program -> Program
mapStatements [Statement] -> [Statement]
f = Identity Program -> Program
forall a. Identity a -> a
runIdentity (Identity Program -> Program)
-> (Program -> Identity Program) -> Program -> Program
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Statement] -> Identity [Statement])
-> Program -> Identity Program
forall (m :: * -> *).
Monad m =>
([Statement] -> m [Statement]) -> Program -> m Program
mapStatementsM ([Statement] -> Identity [Statement]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Statement] -> Identity [Statement])
-> ([Statement] -> [Statement])
-> [Statement]
-> Identity [Statement]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Statement] -> [Statement]
f)

hasFunctionCall :: Expr' -> Bool
hasFunctionCall :: Expr' -> Bool
hasFunctionCall = (Expr' -> Bool) -> [Expr'] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Expr -> Bool
check (Expr -> Bool) -> (Expr' -> Expr) -> Expr' -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr' -> Expr
forall a. WithLoc' a -> a
value') ([Expr'] -> Bool) -> (Expr' -> [Expr']) -> Expr' -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr' -> [Expr']
listSubExprs
  where
    check :: Expr -> Bool
check = \case
      Call Expr'
_ [Expr']
_ -> Bool
True
      Expr
_ -> Bool
False

-- | `isSmallExpr` is true if the evaluation of a given expr trivially terminates.
isSmallExpr :: Expr' -> Bool
isSmallExpr :: Expr' -> Bool
isSmallExpr = Bool -> Bool
not (Bool -> Bool) -> (Expr' -> Bool) -> Expr' -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr' -> Bool
hasFunctionCall

dropLocation :: Expr' -> Expr'
dropLocation :: Expr' -> Expr'
dropLocation = (Expr' -> Expr') -> Expr' -> Expr'
mapSubExpr Expr' -> Expr'
forall a. WithLoc' a -> WithLoc' a
go
  where
    go :: WithLoc' a -> WithLoc' a
go (WithLoc' Maybe Loc
_ a
e) = a -> WithLoc' a
forall a. a -> WithLoc' a
withoutLoc a
e

targetVars :: Target' -> [VarName]
targetVars :: Target' -> [VarName]
targetVars = [VarName] -> [VarName]
forall a. Eq a => [a] -> [a]
nub ([VarName] -> [VarName])
-> (Target' -> [VarName]) -> Target' -> [VarName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VarName' -> VarName) -> [VarName'] -> [VarName]
forall a b. (a -> b) -> [a] -> [b]
map VarName' -> VarName
forall a. WithLoc' a -> a
value' ([VarName'] -> [VarName])
-> (Target' -> [VarName']) -> Target' -> [VarName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Target' -> [VarName']
targetVars'

targetVars' :: Target' -> [VarName']
targetVars' :: Target' -> [VarName']
targetVars' (WithLoc' Maybe Loc
_ Target
x) = case Target
x of
  SubscriptTrg Target'
x Expr'
_ -> Target' -> [VarName']
targetVars' Target'
x
  NameTrg VarName'
x -> [VarName'
x]
  TupleTrg [Target']
xs -> (Target' -> [VarName']) -> [Target'] -> [VarName']
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Target' -> [VarName']
targetVars' [Target']
xs

hasSubscriptTrg :: Target' -> Bool
hasSubscriptTrg :: Target' -> Bool
hasSubscriptTrg (WithLoc' Maybe Loc
_ Target
x) = case Target
x of
  SubscriptTrg Target'
_ Expr'
_ -> Bool
True
  NameTrg VarName'
_ -> Bool
False
  TupleTrg [Target']
xs -> (Target' -> Bool) -> [Target'] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Target' -> Bool
hasSubscriptTrg [Target']
xs

hasBareNameTrg :: Target' -> Bool
hasBareNameTrg :: Target' -> Bool
hasBareNameTrg (WithLoc' Maybe Loc
_ Target
x) = case Target
x of
  SubscriptTrg Target'
_ Expr'
_ -> Bool
False
  NameTrg VarName'
_ -> Bool
True
  TupleTrg [Target']
xs -> (Target' -> Bool) -> [Target'] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Target' -> Bool
hasSubscriptTrg [Target']
xs

exprToTarget :: Expr' -> Maybe Target'
exprToTarget :: Expr' -> Maybe Target'
exprToTarget Expr'
e =
  Maybe Loc -> Target -> Target'
forall a. Maybe Loc -> a -> WithLoc' a
WithLoc' (Expr' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Expr'
e) (Target -> Target') -> Maybe Target -> Maybe Target'
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case Expr' -> Expr
forall a. WithLoc' a -> a
value' Expr'
e of
    Name VarName'
x -> Target -> Maybe Target
forall a. a -> Maybe a
Just (Target -> Maybe Target) -> Target -> Maybe Target
forall a b. (a -> b) -> a -> b
$ VarName' -> Target
NameTrg VarName'
x
    Tuple [Expr']
es -> [Target'] -> Target
TupleTrg ([Target'] -> Target) -> Maybe [Target'] -> Maybe Target
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Expr' -> Maybe Target') -> [Expr'] -> Maybe [Target']
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Expr' -> Maybe Target'
exprToTarget [Expr']
es
    Subscript Expr'
e1 Expr'
e2 -> Target' -> Expr' -> Target
SubscriptTrg (Target' -> Expr' -> Target)
-> Maybe Target' -> Maybe (Expr' -> Target)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr' -> Maybe Target'
exprToTarget Expr'
e1 Maybe (Expr' -> Target) -> Maybe Expr' -> Maybe Target
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr' -> Maybe Expr'
forall (f :: * -> *) a. Applicative f => a -> f a
pure Expr'
e2
    Expr
_ -> Maybe Target
forall a. Maybe a
Nothing

targetToExpr :: Target' -> Expr'
targetToExpr :: Target' -> Expr'
targetToExpr Target'
e =
  Maybe Loc -> Expr -> Expr'
forall a. Maybe Loc -> a -> WithLoc' a
WithLoc' (Target' -> Maybe Loc
forall a. WithLoc' a -> Maybe Loc
loc' Target'
e) (Expr -> Expr') -> Expr -> Expr'
forall a b. (a -> b) -> a -> b
$ case Target' -> Target
forall a. WithLoc' a -> a
value' Target'
e of
    NameTrg VarName'
x -> VarName' -> Expr
Name VarName'
x
    TupleTrg [Target']
es -> [Expr'] -> Expr
Tuple ((Target' -> Expr') -> [Target'] -> [Expr']
forall a b. (a -> b) -> [a] -> [b]
map Target' -> Expr'
targetToExpr [Target']
es)
    SubscriptTrg Target'
e1 Expr'
e2 -> Expr' -> Expr' -> Expr
Subscript (Target' -> Expr'
targetToExpr Target'
e1) Expr'
e2

toplevelMainDef :: [Statement] -> Program
toplevelMainDef :: [Statement] -> Program
toplevelMainDef [Statement]
body = [VarName'
-> [(VarName', Type)] -> Type -> [Statement] -> ToplevelStatement
ToplevelFunctionDef (Maybe Loc -> VarName -> VarName'
forall a. Maybe Loc -> a -> WithLoc' a
WithLoc' Maybe Loc
forall a. Maybe a
Nothing (String -> VarName
VarName String
"main")) [] Type
IntTy [Statement]
body]