module Futhark.Internalise.FullNormalise (transformProg) where
import Control.Monad.Reader
import Control.Monad.State
import Data.Bifunctor
import Data.List.NonEmpty qualified as NE
import Data.Map qualified as M
import Data.Text qualified as T
import Futhark.MonadFreshNames
import Language.Futhark
import Language.Futhark.Traversals
import Language.Futhark.TypeChecker.Types
data BindModifier
= Ass Exp (Info T.Text) SrcLoc
| Att (AttrInfo VName)
applyModifiers :: Exp -> [BindModifier] -> (Exp, [BindModifier])
applyModifiers :: Exp -> [BindModifier] -> (Exp, [BindModifier])
applyModifiers =
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr BindModifier -> (Exp, [BindModifier]) -> (Exp, [BindModifier])
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. (,[])
where
f :: BindModifier -> (Exp, [BindModifier]) -> (Exp, [BindModifier])
f (Ass Exp
ass Info Text
txt SrcLoc
loc) (Exp
body, [BindModifier]
modifs) =
(forall (f :: * -> *) vn.
ExpBase f vn -> ExpBase f vn -> f Text -> SrcLoc -> ExpBase f vn
Assert Exp
ass Exp
body Info Text
txt SrcLoc
loc, [BindModifier]
modifs)
f (Att AttrInfo VName
attr) (Exp
body, [BindModifier]
modifs) =
(forall (f :: * -> *) vn.
AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo VName
attr Exp
body forall a. Monoid a => a
mempty, AttrInfo VName -> BindModifier
Att AttrInfo VName
attr forall a. a -> [a] -> [a]
: [BindModifier]
modifs)
data Binding
= PatBind [SizeBinder VName] (Pat StructType) Exp
| FunBind VName ([TypeParam], [Pat ParamType], Maybe (TypeExp Info VName), Info ResRetType, Exp)
type NormState = (([Binding], [BindModifier]), VNameSource)
newtype OrderingM a = OrderingM (StateT NormState (Reader String) a)
deriving
(forall a b. a -> OrderingM b -> OrderingM a
forall a b. (a -> b) -> OrderingM a -> OrderingM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> OrderingM b -> OrderingM a
$c<$ :: forall a b. a -> OrderingM b -> OrderingM a
fmap :: forall a b. (a -> b) -> OrderingM a -> OrderingM b
$cfmap :: forall a b. (a -> b) -> OrderingM a -> OrderingM b
Functor, Functor OrderingM
forall a. a -> OrderingM a
forall a b. OrderingM a -> OrderingM b -> OrderingM a
forall a b. OrderingM a -> OrderingM b -> OrderingM b
forall a b. OrderingM (a -> b) -> OrderingM a -> OrderingM b
forall a b c.
(a -> b -> c) -> OrderingM a -> OrderingM b -> OrderingM c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. OrderingM a -> OrderingM b -> OrderingM a
$c<* :: forall a b. OrderingM a -> OrderingM b -> OrderingM a
*> :: forall a b. OrderingM a -> OrderingM b -> OrderingM b
$c*> :: forall a b. OrderingM a -> OrderingM b -> OrderingM b
liftA2 :: forall a b c.
(a -> b -> c) -> OrderingM a -> OrderingM b -> OrderingM c
$cliftA2 :: forall a b c.
(a -> b -> c) -> OrderingM a -> OrderingM b -> OrderingM c
<*> :: forall a b. OrderingM (a -> b) -> OrderingM a -> OrderingM b
$c<*> :: forall a b. OrderingM (a -> b) -> OrderingM a -> OrderingM b
pure :: forall a. a -> OrderingM a
$cpure :: forall a. a -> OrderingM a
Applicative, Applicative OrderingM
forall a. a -> OrderingM a
forall a b. OrderingM a -> OrderingM b -> OrderingM b
forall a b. OrderingM a -> (a -> OrderingM b) -> OrderingM b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> OrderingM a
$creturn :: forall a. a -> OrderingM a
>> :: forall a b. OrderingM a -> OrderingM b -> OrderingM b
$c>> :: forall a b. OrderingM a -> OrderingM b -> OrderingM b
>>= :: forall a b. OrderingM a -> (a -> OrderingM b) -> OrderingM b
$c>>= :: forall a b. OrderingM a -> (a -> OrderingM b) -> OrderingM b
Monad, MonadReader String, MonadState NormState)
instance MonadFreshNames OrderingM where
getNameSource :: OrderingM VNameSource
getNameSource = forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM forall a b. (a -> b) -> a -> b
$ forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. (a, b) -> b
snd
putNameSource :: VNameSource -> OrderingM ()
putNameSource = forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> b -> a
const
addModifier :: BindModifier -> OrderingM ()
addModifier :: BindModifier -> OrderingM ()
addModifier = forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second forall b c a. (b -> c) -> (a -> b) -> a -> c
. (:)
rmModifier :: OrderingM ()
rmModifier :: OrderingM ()
rmModifier = forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM forall a b. (a -> b) -> a -> b
$ forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second forall a. [a] -> [a]
tail
addBind :: Binding -> OrderingM ()
addBind :: Binding -> OrderingM ()
addBind (PatBind [SizeBinder VName]
s Pat StructType
p Exp
e) = do
[BindModifier]
modifs <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst
let (Exp
e', [BindModifier]
modifs') = Exp -> [BindModifier] -> (Exp, [BindModifier])
applyModifiers Exp
e [BindModifier]
modifs
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap ([SizeBinder VName] -> Pat StructType -> Exp -> Binding
PatBind ([SizeBinder VName]
s forall a. Semigroup a => a -> a -> a
<> [SizeBinder VName]
implicit) Pat StructType
p Exp
e' :) (forall a b. a -> b -> a
const [BindModifier]
modifs')
where
implicit :: [SizeBinder VName]
implicit = case Exp
e of
(AppExp AppExpBase Info VName
_ (Info (AppRes StructType
_ [VName]
ext))) -> forall a b. (a -> b) -> [a] -> [b]
map (forall vn. vn -> SrcLoc -> SizeBinder vn
`SizeBinder` forall a. Monoid a => a
mempty) [VName]
ext
Exp
_ -> []
addBind b :: Binding
b@FunBind {} =
forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM forall a b. (a -> b) -> a -> b
$ forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Binding
b :)
runOrdering :: MonadFreshNames m => OrderingM a -> m (a, [Binding])
runOrdering :: forall (m :: * -> *) a.
MonadFreshNames m =>
OrderingM a -> m (a, [Binding])
runOrdering (OrderingM StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
m) =
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource forall a b. (a -> b) -> a -> b
$ forall {t :: * -> *} {a} {b} {a} {b}.
Foldable t =>
(a, ((b, t a), b)) -> ((a, b), b)
mod_tup forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall r a. Reader r a -> r -> a
runReader [Char]
"tmp" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
m forall b c a. (b -> c) -> (a -> b) -> a -> c
. (([], []),)
where
mod_tup :: (a, ((b, t a), b)) -> ((a, b), b)
mod_tup (a
a, ((b
binds, t a
modifs), b
src)) =
if forall (t :: * -> *) a. Foldable t => t a -> Bool
null t a
modifs
then ((a
a, b
binds), b
src)
else forall a. HasCallStack => [Char] -> a
error [Char]
"not all bind modifiers were freed"
naming :: String -> OrderingM a -> OrderingM a
naming :: forall a. [Char] -> OrderingM a -> OrderingM a
naming [Char]
s = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (forall a b. a -> b -> a
const [Char]
s)
nameExp :: Bool -> Exp -> OrderingM Exp
nameExp :: Bool -> Exp -> OrderingM Exp
nameExp Bool
True Exp
e = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
nameExp Bool
False Exp
e = do
VName
name <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall r (m :: * -> *). MonadReader r m => m r
ask
let ty :: StructType
ty = Exp -> StructType
typeOf Exp
e
loc :: SrcLoc
loc = forall a. Located a => a -> SrcLoc
srclocOf Exp
e
pat :: Pat StructType
pat = forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
name (forall a. a -> Info a
Info StructType
ty) SrcLoc
loc
Binding -> OrderingM ()
addBind forall a b. (a -> b) -> a -> b
$ [SizeBinder VName] -> Pat StructType -> Exp -> Binding
PatBind [] Pat StructType
pat Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
name) (forall a. a -> Info a
Info StructType
ty) SrcLoc
loc
patRepName :: Pat t -> String
patRepName :: forall t. Pat t -> [Char]
patRepName (PatAscription PatBase Info VName t
p TypeExp Info VName
_ SrcLoc
_) = forall t. Pat t -> [Char]
patRepName PatBase Info VName t
p
patRepName (Id VName
v Info t
_ SrcLoc
_) = VName -> [Char]
baseString VName
v
patRepName PatBase Info VName t
_ = [Char]
"tmp"
expRepName :: Exp -> String
expRepName :: Exp -> [Char]
expRepName (Var QualName VName
v Info StructType
_ SrcLoc
_) = forall a. Pretty a => a -> [Char]
prettyString QualName VName
v
expRepName Exp
e = [Char]
"d<{" forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString (Exp -> ExpBase NoInfo VName
bareExp Exp
e) forall a. [a] -> [a] -> [a]
++ [Char]
"}>"
argRepName :: Exp -> Int -> String
argRepName :: Exp -> Int -> [Char]
argRepName Exp
e Int
i = Exp -> [Char]
expRepName Exp
e forall a. Semigroup a => a -> a -> a
<> [Char]
"_arg" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> [Char]
show Int
i
getOrdering :: Bool -> Exp -> OrderingM Exp
getOrdering :: Bool -> Exp -> OrderingM Exp
getOrdering Bool
final (Assert Exp
ass Exp
e Info Text
txt SrcLoc
loc) = do
Exp
ass' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
ass
Int
l_prev <- forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM forall a b. (a -> b) -> a -> b
$ forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Int
length forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst
BindModifier -> OrderingM ()
addModifier forall a b. (a -> b) -> a -> b
$ Exp -> Info Text -> SrcLoc -> BindModifier
Ass Exp
ass' Info Text
txt SrcLoc
loc
Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
final Exp
e
Int
l_after <- forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM forall a b. (a -> b) -> a -> b
$ forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Int
length forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst
if Int
l_after forall a. Ord a => a -> a -> Bool
<= Int
l_prev
then forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e'
else do
OrderingM ()
rmModifier
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
ExpBase f vn -> ExpBase f vn -> f Text -> SrcLoc -> ExpBase f vn
Assert Exp
ass' Exp
e' Info Text
txt SrcLoc
loc
getOrdering Bool
final (Attr AttrInfo VName
attr Exp
e SrcLoc
loc) = do
BindModifier -> OrderingM ()
addModifier forall a b. (a -> b) -> a -> b
$ AttrInfo VName -> BindModifier
Att AttrInfo VName
attr
Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
final Exp
e
OrderingM ()
rmModifier
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo VName
attr Exp
e' SrcLoc
loc
getOrdering Bool
_ e :: Exp
e@Literal {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
getOrdering Bool
_ e :: Exp
e@IntLit {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
getOrdering Bool
_ e :: Exp
e@FloatLit {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
getOrdering Bool
_ e :: Exp
e@StringLit {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
getOrdering Bool
_ e :: Exp
e@Hole {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
getOrdering Bool
_ e :: Exp
e@Var {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
getOrdering Bool
final (Parens Exp
e SrcLoc
_) = Bool -> Exp -> OrderingM Exp
getOrdering Bool
final Exp
e
getOrdering Bool
final (QualParens (QualName VName, SrcLoc)
_ Exp
e SrcLoc
_) = Bool -> Exp -> OrderingM Exp
getOrdering Bool
final Exp
e
getOrdering Bool
_ (TupLit [Exp]
es SrcLoc
loc) = do
[Exp]
es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Bool -> Exp -> OrderingM Exp
getOrdering Bool
False) [Exp]
es
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit [Exp]
es' SrcLoc
loc
getOrdering Bool
_ (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) = do
[FieldBase Info VName]
fs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName -> OrderingM (FieldBase Info VName)
f [FieldBase Info VName]
fs
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit [FieldBase Info VName]
fs' SrcLoc
loc
where
f :: FieldBase Info VName -> OrderingM (FieldBase Info VName)
f (RecordFieldExplicit Name
n Exp
e SrcLoc
floc) = do
Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
n Exp
e' SrcLoc
floc
f field :: FieldBase Info VName
field@RecordFieldImplicit {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure FieldBase Info VName
field
getOrdering Bool
_ (ArrayLit [Exp]
es Info StructType
ty SrcLoc
loc) = do
[Exp]
es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Bool -> Exp -> OrderingM Exp
getOrdering Bool
False) [Exp]
es
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
ArrayLit [Exp]
es' Info StructType
ty SrcLoc
loc
getOrdering Bool
_ (Project Name
n Exp
e Info StructType
ty SrcLoc
loc) = do
Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
n Exp
e' Info StructType
ty SrcLoc
loc
getOrdering Bool
_ (Negate Exp
e SrcLoc
loc) = do
Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Negate Exp
e' SrcLoc
loc
getOrdering Bool
_ (Not Exp
e SrcLoc
loc) = do
Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Not Exp
e' SrcLoc
loc
getOrdering Bool
final (Constr Name
n [Exp]
es Info StructType
ty SrcLoc
loc) = do
[Exp]
es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Bool -> Exp -> OrderingM Exp
getOrdering Bool
False) [Exp]
es
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
Name -> [ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
Constr Name
n [Exp]
es' Info StructType
ty SrcLoc
loc
getOrdering Bool
final (Update Exp
eb SliceBase Info VName
slice Exp
eu SrcLoc
loc) = do
Exp
eu' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
eu
SliceBase Info VName
slice' <- forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper OrderingM
mapper SliceBase Info VName
slice
Exp
eb' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
eb
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update Exp
eb' SliceBase Info VName
slice' Exp
eu' SrcLoc
loc
where
mapper :: ASTMapper OrderingM
mapper = forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp :: Exp -> OrderingM Exp
mapOnExp = Bool -> Exp -> OrderingM Exp
getOrdering Bool
False}
getOrdering Bool
final (RecordUpdate Exp
eb [Name]
ns Exp
eu Info StructType
ty SrcLoc
loc) = do
Exp
eb' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
eb
Exp
eu' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
eu
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
ExpBase f vn
-> [Name] -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
RecordUpdate Exp
eb' [Name]
ns Exp
eu' Info StructType
ty SrcLoc
loc
getOrdering Bool
final (Lambda [PatBase Info VName ParamType]
params Exp
body Maybe (TypeExp Info VName)
mte Info ResRetType
ret SrcLoc
loc) = do
Exp
body' <- forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
body
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [PatBase Info VName ParamType]
params Exp
body' Maybe (TypeExp Info VName)
mte Info ResRetType
ret SrcLoc
loc
getOrdering Bool
_ (OpSection QualName VName
qn Info StructType
ty SrcLoc
loc) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn Info StructType
ty SrcLoc
loc
getOrdering Bool
final (OpSectionLeft QualName VName
op Info StructType
ty Exp
e (Info (PName
xp, ParamType
_, Maybe VName
xext), Info (PName
yp, ParamType
yty)) (Info (RetType [VName]
dims TypeBase Exp Uniqueness
ret), Info [VName]
exts) SrcLoc
loc) = do
Exp
x <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
VName
yn <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"y"
let y :: Exp
y = forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
yn) (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ParamType
yty) forall a. Monoid a => a
mempty
ret' :: TypeBase Exp Uniqueness
ret' = forall a. Substitutable a => TypeSubs -> a -> a
applySubst (forall {t}. Exp -> Exp -> VName -> Maybe (Subst t)
pSubst Exp
x Exp
y) TypeBase Exp Uniqueness
ret
body :: Exp
body =
forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply (forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
op Info StructType
ty forall a. Monoid a => a
mempty) [(Diet
Observe, Maybe VName
xext, Exp
x), (Diet
Observe, forall a. Maybe a
Nothing, Exp
y)] forall a b. (a -> b) -> a -> b
$
StructType -> [VName] -> AppRes
AppRes (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp Uniqueness
ret') [VName]
exts
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
yn (forall a. a -> Info a
Info ParamType
yty) forall a. Monoid a => a
mempty] Exp
body forall a. Maybe a
Nothing (forall a. a -> Info a
Info (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase Exp Uniqueness
ret')) SrcLoc
loc
where
pSubst :: Exp -> Exp -> VName -> Maybe (Subst t)
pSubst Exp
x Exp
y VName
vn
| Named VName
p <- PName
xp, VName
p forall a. Eq a => a -> a -> Bool
== VName
vn = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall t. Exp -> Subst t
ExpSubst Exp
x
| Named VName
p <- PName
yp, VName
p forall a. Eq a => a -> a -> Bool
== VName
vn = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall t. Exp -> Subst t
ExpSubst Exp
y
| Bool
otherwise = forall a. Maybe a
Nothing
getOrdering Bool
final (OpSectionRight QualName VName
op Info StructType
ty Exp
e (Info (PName
xp, ParamType
xty), Info (PName
yp, ParamType
_, Maybe VName
yext)) (Info (RetType [VName]
dims TypeBase Exp Uniqueness
ret)) SrcLoc
loc) = do
VName
xn <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"x"
Exp
y <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
let x :: Exp
x = forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
xn) (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ParamType
xty) forall a. Monoid a => a
mempty
ret' :: TypeBase Exp Uniqueness
ret' = forall a. Substitutable a => TypeSubs -> a -> a
applySubst (forall {t}. Exp -> Exp -> VName -> Maybe (Subst t)
pSubst Exp
x Exp
y) TypeBase Exp Uniqueness
ret
body :: Exp
body = forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply (forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
op Info StructType
ty forall a. Monoid a => a
mempty) [(Diet
Observe, forall a. Maybe a
Nothing, Exp
x), (Diet
Observe, Maybe VName
yext, Exp
y)] forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp Uniqueness
ret') []
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
xn (forall a. a -> Info a
Info ParamType
xty) forall a. Monoid a => a
mempty] Exp
body forall a. Maybe a
Nothing (forall a. a -> Info a
Info (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase Exp Uniqueness
ret')) SrcLoc
loc
where
pSubst :: Exp -> Exp -> VName -> Maybe (Subst t)
pSubst Exp
x Exp
y VName
vn
| Named VName
p <- PName
xp, VName
p forall a. Eq a => a -> a -> Bool
== VName
vn = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall t. Exp -> Subst t
ExpSubst Exp
x
| Named VName
p <- PName
yp, VName
p forall a. Eq a => a -> a -> Bool
== VName
vn = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall t. Exp -> Subst t
ExpSubst Exp
y
| Bool
otherwise = forall a. Maybe a
Nothing
getOrdering Bool
final (ProjectSection [Name]
names (Info StructType
ty) SrcLoc
loc) = do
VName
xn <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"x"
let (ParamType
xty, RetType [VName]
dims TypeBase Exp Uniqueness
ret) = case StructType
ty of
Scalar (Arrow NoUniqueness
_ PName
_ Diet
d StructType
xty' ResRetType
ret') -> (forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
d StructType
xty', ResRetType
ret')
StructType
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"not a function type for project section: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString StructType
ty
x :: Exp
x = forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
xn) (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ParamType
xty) forall a. Monoid a => a
mempty
body :: Exp
body = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Exp -> Name -> Exp
project Exp
x [Name]
names
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
xn (forall a. a -> Info a
Info ParamType
xty) forall a. Monoid a => a
mempty] Exp
body forall a. Maybe a
Nothing (forall a. a -> Info a
Info (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase Exp Uniqueness
ret)) SrcLoc
loc
where
project :: Exp -> Name -> Exp
project Exp
e Name
field =
case Exp -> StructType
typeOf Exp
e of
Scalar (Record Map Name StructType
fs)
| Just StructType
t <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
field Map Name StructType
fs ->
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
field Exp
e (forall a. a -> Info a
Info StructType
t) forall a. Monoid a => a
mempty
StructType
t ->
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"desugar ProjectSection: type "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString StructType
t
forall a. [a] -> [a] -> [a]
++ [Char]
" does not have field "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Name
field
getOrdering Bool
final (IndexSection SliceBase Info VName
slice (Info StructType
ty) SrcLoc
loc) = do
SliceBase Info VName
slice' <- forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper OrderingM
mapper SliceBase Info VName
slice
VName
xn <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"x"
let (ParamType
xty, RetType [VName]
dims TypeBase Exp Uniqueness
ret) = case StructType
ty of
Scalar (Arrow NoUniqueness
_ PName
_ Diet
d StructType
xty' ResRetType
ret') -> (forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
d StructType
xty', ResRetType
ret')
StructType
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"not a function type for index section: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString StructType
ty
x :: Exp
x = forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
xn) (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ParamType
xty) forall a. Monoid a => a
mempty
body :: Exp
body = forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
Index Exp
x SliceBase Info VName
slice' SrcLoc
loc) (forall a. a -> Info a
Info (StructType -> [VName] -> AppRes
AppRes (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp Uniqueness
ret) []))
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
xn (forall a. a -> Info a
Info ParamType
xty) forall a. Monoid a => a
mempty] Exp
body forall a. Maybe a
Nothing (forall a. a -> Info a
Info (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase Exp Uniqueness
ret)) SrcLoc
loc
where
mapper :: ASTMapper OrderingM
mapper = forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp :: Exp -> OrderingM Exp
mapOnExp = Bool -> Exp -> OrderingM Exp
getOrdering Bool
False}
getOrdering Bool
_ (Ascript Exp
e TypeExp Info VName
_ SrcLoc
_) = Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
getOrdering Bool
final (AppExp (Apply Exp
f NonEmpty (Info (Diet, Maybe VName), Exp)
args SrcLoc
loc) Info AppRes
resT) = do
NonEmpty (Info (Diet, Maybe VName), Exp)
args' <-
forall a. NonEmpty a -> NonEmpty a
NE.reverse forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {t}. ((t, Exp), Int) -> OrderingM (t, Exp)
onArg (forall a. NonEmpty a -> NonEmpty a
NE.reverse (forall a b. NonEmpty a -> NonEmpty b -> NonEmpty (a, b)
NE.zip NonEmpty (Info (Diet, Maybe VName), Exp)
args (forall a. [a] -> NonEmpty a
NE.fromList [Int
0 ..])))
Exp
f' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
f
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (f (Diet, Maybe VName), ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
f' NonEmpty (Info (Diet, Maybe VName), Exp)
args' SrcLoc
loc) Info AppRes
resT
where
onArg :: ((t, Exp), Int) -> OrderingM (t, Exp)
onArg ((t
d, Exp
e), Int
i) =
forall a. [Char] -> OrderingM a -> OrderingM a
naming (Exp -> Int -> [Char]
argRepName Exp
f Int
i) forall a b. (a -> b) -> a -> b
$ (t
d,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
getOrdering Bool
final (Coerce Exp
e TypeExp Info VName
te Info StructType
t SrcLoc
loc) = do
Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
ExpBase f vn
-> TypeExp f vn -> f StructType -> SrcLoc -> ExpBase f vn
Coerce Exp
e' TypeExp Info VName
te Info StructType
t SrcLoc
loc
getOrdering Bool
final (AppExp (Range Exp
start Maybe Exp
stride Inclusiveness Exp
end SrcLoc
loc) Info AppRes
resT) = do
Exp
start' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
start
Maybe Exp
stride' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Bool -> Exp -> OrderingM Exp
getOrdering Bool
False) Maybe Exp
stride
Inclusiveness Exp
end' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Bool -> Exp -> OrderingM Exp
getOrdering Bool
False) Inclusiveness Exp
end
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> Maybe (ExpBase f vn)
-> Inclusiveness (ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Range Exp
start' Maybe Exp
stride' Inclusiveness Exp
end' SrcLoc
loc) Info AppRes
resT
getOrdering Bool
final (AppExp (LetPat [SizeBinder VName]
sizes Pat StructType
pat Exp
expr Exp
body SrcLoc
_) Info AppRes
_) = do
Exp
expr' <- forall a. [Char] -> OrderingM a -> OrderingM a
naming (forall t. Pat t -> [Char]
patRepName Pat StructType
pat) forall a b. (a -> b) -> a -> b
$ Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
expr
Binding -> OrderingM ()
addBind forall a b. (a -> b) -> a -> b
$ [SizeBinder VName] -> Pat StructType -> Exp -> Binding
PatBind [SizeBinder VName]
sizes Pat StructType
pat Exp
expr'
Bool -> Exp -> OrderingM Exp
getOrdering Bool
final Exp
body
getOrdering Bool
final (AppExp (LetFun VName
vn ([TypeParamBase VName]
tparams, [PatBase Info VName ParamType]
params, Maybe (TypeExp Info VName)
mrettype, Info ResRetType
rettype, Exp
body) Exp
e SrcLoc
_) Info AppRes
_) = do
Exp
body' <- forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
body
Binding -> OrderingM ()
addBind forall a b. (a -> b) -> a -> b
$ VName
-> ([TypeParamBase VName], [PatBase Info VName ParamType],
Maybe (TypeExp Info VName), Info ResRetType, Exp)
-> Binding
FunBind VName
vn ([TypeParamBase VName]
tparams, [PatBase Info VName ParamType]
params, Maybe (TypeExp Info VName)
mrettype, Info ResRetType
rettype, Exp
body')
Bool -> Exp -> OrderingM Exp
getOrdering Bool
final Exp
e
getOrdering Bool
final (AppExp (If Exp
cond Exp
et Exp
ef SrcLoc
loc) Info AppRes
resT) = do
Exp
cond' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
cond
Exp
et' <- forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
et
Exp
ef' <- forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
ef
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
If Exp
cond' Exp
et' Exp
ef' SrcLoc
loc) Info AppRes
resT
getOrdering Bool
final (AppExp (DoLoop [VName]
sizes PatBase Info VName ParamType
pat Exp
einit LoopFormBase Info VName
form Exp
body SrcLoc
loc) Info AppRes
resT) = do
Exp
einit' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
einit
LoopFormBase Info VName
form' <- case LoopFormBase Info VName
form of
For IdentBase Info VName StructType
ident Exp
e -> forall (f :: * -> *) vn.
IdentBase f vn StructType -> ExpBase f vn -> LoopFormBase f vn
For IdentBase Info VName StructType
ident forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
e
ForIn Pat StructType
fpat Exp
e -> forall (f :: * -> *) vn.
PatBase f vn StructType -> ExpBase f vn -> LoopFormBase f vn
ForIn Pat StructType
fpat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
e
While Exp
e -> forall (f :: * -> *) vn. ExpBase f vn -> LoopFormBase f vn
While forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
e
Exp
body' <- forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
body
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
[VName]
-> PatBase f vn ParamType
-> ExpBase f vn
-> LoopFormBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
DoLoop [VName]
sizes PatBase Info VName ParamType
pat Exp
einit' LoopFormBase Info VName
form' Exp
body' SrcLoc
loc) Info AppRes
resT
getOrdering Bool
final (AppExp (BinOp (QualName VName
op, SrcLoc
oloc) Info StructType
opT (Exp
el, Info Maybe VName
elp) (Exp
er, Info Maybe VName
erp) SrcLoc
loc) (Info AppRes
resT)) = do
Exp
expr' <- case (Bool
isOr, Bool
isAnd) of
(Bool
True, Bool
_) -> do
Exp
el' <- forall a. [Char] -> OrderingM a -> OrderingM a
naming [Char]
"or_lhs" forall a b. (a -> b) -> a -> b
$ Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
el
Exp
er' <- forall a. [Char] -> OrderingM a -> OrderingM a
naming [Char]
"or_rhs" forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
er
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
If Exp
el' (forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (Bool -> PrimValue
BoolValue Bool
True) forall a. Monoid a => a
mempty) Exp
er' SrcLoc
loc) (forall a. a -> Info a
Info AppRes
resT)
(Bool
_, Bool
True) -> do
Exp
el' <- forall a. [Char] -> OrderingM a -> OrderingM a
naming [Char]
"and_lhs" forall a b. (a -> b) -> a -> b
$ Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
el
Exp
er' <- forall a. [Char] -> OrderingM a -> OrderingM a
naming [Char]
"and_rhs" forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
er
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
If Exp
el' Exp
er' (forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (Bool -> PrimValue
BoolValue Bool
False) forall a. Monoid a => a
mempty) SrcLoc
loc) (forall a. a -> Info a
Info AppRes
resT)
(Bool
False, Bool
False) -> do
Exp
el' <- forall a. [Char] -> OrderingM a -> OrderingM a
naming (forall a. Pretty a => a -> [Char]
prettyString QualName VName
op forall a. Semigroup a => a -> a -> a
<> [Char]
"_lhs") forall a b. (a -> b) -> a -> b
$ Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
el
Exp
er' <- forall a. [Char] -> OrderingM a -> OrderingM a
naming (forall a. Pretty a => a -> [Char]
prettyString QualName VName
op forall a. Semigroup a => a -> a -> a
<> [Char]
"_rhs") forall a b. (a -> b) -> a -> b
$ Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
er
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply (forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
op Info StructType
opT SrcLoc
oloc) [(Diet
Observe, Maybe VName
elp, Exp
el'), (Diet
Observe, Maybe VName
erp, Exp
er')] AppRes
resT
Bool -> Exp -> OrderingM Exp
nameExp Bool
final Exp
expr'
where
isOr :: Bool
isOr = VName -> Name
baseName (forall vn. QualName vn -> vn
qualLeaf QualName VName
op) forall a. Eq a => a -> a -> Bool
== Name
"||"
isAnd :: Bool
isAnd = VName -> Name
baseName (forall vn. QualName vn -> vn
qualLeaf QualName VName
op) forall a. Eq a => a -> a -> Bool
== Name
"&&"
getOrdering Bool
final (AppExp (LetWith (Ident VName
dest Info StructType
dty SrcLoc
dloc) (Ident VName
src Info StructType
sty SrcLoc
sloc) SliceBase Info VName
slice Exp
e Exp
body SrcLoc
loc) Info AppRes
_) = do
Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
SliceBase Info VName
slice' <- forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper OrderingM
mapper SliceBase Info VName
slice
Binding -> OrderingM ()
addBind forall a b. (a -> b) -> a -> b
$ [SizeBinder VName] -> Pat StructType -> Exp -> Binding
PatBind [] (forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
dest Info StructType
dty SrcLoc
dloc) (forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update (forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
src) Info StructType
sty SrcLoc
sloc) SliceBase Info VName
slice' Exp
e' SrcLoc
loc)
Bool -> Exp -> OrderingM Exp
getOrdering Bool
final Exp
body
where
mapper :: ASTMapper OrderingM
mapper = forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp :: Exp -> OrderingM Exp
mapOnExp = Bool -> Exp -> OrderingM Exp
getOrdering Bool
False}
getOrdering Bool
final (AppExp (Index Exp
e SliceBase Info VName
slice SrcLoc
loc) Info AppRes
resT) = do
Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
SliceBase Info VName
slice' <- forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper OrderingM
mapper SliceBase Info VName
slice
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
Index Exp
e' SliceBase Info VName
slice' SrcLoc
loc) Info AppRes
resT
where
mapper :: ASTMapper OrderingM
mapper = forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp :: Exp -> OrderingM Exp
mapOnExp = Bool -> Exp -> OrderingM Exp
getOrdering Bool
False}
getOrdering Bool
final (AppExp (Match Exp
expr NonEmpty (CaseBase Info VName)
cs SrcLoc
loc) Info AppRes
resT) = do
Exp
expr' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
expr
NonEmpty (CaseBase Info VName)
cs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {m :: * -> *}.
MonadFreshNames m =>
CaseBase Info VName -> m (CaseBase Info VName)
f NonEmpty (CaseBase Info VName)
cs
Bool -> Exp -> OrderingM Exp
nameExp Bool
final forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (CaseBase f vn) -> SrcLoc -> AppExpBase f vn
Match Exp
expr' NonEmpty (CaseBase Info VName)
cs' SrcLoc
loc) Info AppRes
resT
where
f :: CaseBase Info VName -> m (CaseBase Info VName)
f (CasePat Pat StructType
pat Exp
body SrcLoc
cloc) = do
Exp
body' <- forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
body
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
PatBase f vn StructType -> ExpBase f vn -> SrcLoc -> CaseBase f vn
CasePat Pat StructType
pat Exp
body' SrcLoc
cloc)
transformBody :: MonadFreshNames m => Exp -> m Exp
transformBody :: forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
e = do
(Exp
e', [Binding]
pre_eval) <- forall (m :: * -> *) a.
MonadFreshNames m =>
OrderingM a -> m (a, [Binding])
runOrdering (Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Exp -> Binding -> Exp
f Exp
e' [Binding]
pre_eval
where
appRes :: Info AppRes
appRes = case Exp
e of
(AppExp AppExpBase Info VName
_ Info AppRes
r) -> Info AppRes
r
Exp
_ -> forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (Exp -> StructType
typeOf Exp
e) []
f :: Exp -> Binding -> Exp
f Exp
body (PatBind [SizeBinder VName]
sizes Pat StructType
p Exp
expr) =
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
( forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn StructType
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat
[SizeBinder VName]
sizes
Pat StructType
p
Exp
expr
Exp
body
forall a. Monoid a => a
mempty
)
Info AppRes
appRes
f Exp
body (FunBind VName
vn ([TypeParamBase VName], [PatBase Info VName ParamType],
Maybe (TypeExp Info VName), Info ResRetType, Exp)
infos) =
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
vn
-> ([TypeParamBase vn], [PatBase f vn ParamType],
Maybe (TypeExp f vn), f ResRetType, ExpBase f vn)
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetFun VName
vn ([TypeParamBase VName], [PatBase Info VName ParamType],
Maybe (TypeExp Info VName), Info ResRetType, Exp)
infos Exp
body forall a. Monoid a => a
mempty) Info AppRes
appRes
transformDec :: MonadFreshNames m => Dec -> m Dec
transformDec :: forall (m :: * -> *). MonadFreshNames m => Dec -> m Dec
transformDec (ValDec ValBindBase Info VName
valbind) = do
Exp
body' <- forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. ValBindBase f vn -> ExpBase f vn
valBindBody ValBindBase Info VName
valbind
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. ValBindBase f vn -> DecBase f vn
ValDec (ValBindBase Info VName
valbind {valBindBody :: Exp
valBindBody = Exp
body'})
transformDec Dec
d = forall (f :: * -> *) a. Applicative f => a -> f a
pure Dec
d
transformProg :: MonadFreshNames m => [Dec] -> m [Dec]
transformProg :: forall (m :: * -> *). MonadFreshNames m => [Dec] -> m [Dec]
transformProg = forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall (m :: * -> *). MonadFreshNames m => Dec -> m Dec
transformDec