-- | This full normalisation module converts a well-typed, polymorphic,
-- module-free Futhark program into an equivalent with only simple expresssions.
-- Notably, all non-trivial expression are converted into a list of
-- let-bindings to make them simpler, with no nested apply, nested lets...
-- This module only performs synthatic operations.
--
-- Also, it performs desugaring that is:
-- * Turn operator section into lambda
-- * turn BinOp into application (&& and || are converted to if structure)
-- * turn `let x [i] = e1` into `let x = x with [i] = e1`
-- * binds all implicit sizes
--
-- This is currently not done for expressions inside sizes, this processing
-- still needed in monomorphisation for now.
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

-- Modifier to apply on binding, this is used to propagate attributes and move assertions
data BindModifier
  = Ass Exp (Info T.Text) SrcLoc
  | Att (AttrInfo VName)

-- Apply a list of modifiers, removing the assertions as it is not needed to check them multiple times
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)

-- A binding that occurs in the calculation flow
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)

-- | Main monad of this module, the state as 3 parts:
-- * the VNameSource to produce new names
-- * the [Binding] is the accumulator for the result
--   It behave a bit like a writer
-- * the [BindModifier] is the current list of modifiers to apply to the introduced bindings
--   It behave like a reader for attributes modifier, and as a state for assertion,
--   they have to be in the same list to conserve their order
-- Direct interaction with the inside state should be done with caution, that's why their
-- no instance of `MonadState`.
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)

-- | From now, we say an expression is "final" if it's going to be stored in a let-bind
-- or is at the end of the body e.g. after all lets

-- Replace a non-final expression by a let-binded variable
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 -- "e<{" ++ prettyString e ++ "}>"
  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

-- An evocative name to use when naming subexpressions of the
-- expression bound to this pattern.
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]
"}>"

-- An evocative name to use when naming arguments to an application.
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

-- Modify an expression as describe in module introduction,
-- introducing the let-bindings in the state.
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 the list of modifier has reduced in size, that means that
  -- all assertions as been inserted,
  -- else, we have to introduce the assertion ourself
  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
  -- propagate attribute
  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 -- can we still have some ?
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)

-- Transform a body, e.g. the expression of a valbind,
-- branches of an if/match...
-- Note that this is not producing an OrderingM, produce
-- a complete separtion of states.
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