-- | 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 syntactic operations.
--
-- Also, it performs various kinds of desugaring:
--
-- * Turns operator sections into explicit lambdas.
--
-- * Rewrites BinOp nodes to Apply nodes (&& and || are converted to conditionals).
--
-- * Turns `let x [i] = e1` into `let x = x with [i] = e1`.
--
-- * Binds all implicit sizes.
--
-- * Turns implicit record fields into explicit record fields.
--
-- 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 Futhark.MonadFreshNames
import Futhark.Util (showText)
import Futhark.Util.Loc (srcspan)
import Language.Futhark
import Language.Futhark.Traversals
import Language.Futhark.TypeChecker.Types

-- Modifier to apply on bindings, this is used to propagate attributes.
newtype BindModifier
  = 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 =
  (BindModifier -> (Exp, [BindModifier]) -> (Exp, [BindModifier]))
-> (Exp, [BindModifier]) -> [BindModifier] -> (Exp, [BindModifier])
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr BindModifier -> (Exp, [BindModifier]) -> (Exp, [BindModifier])
forall {f :: * -> *}.
BindModifier
-> (ExpBase f VName, [BindModifier])
-> (ExpBase f VName, [BindModifier])
f ((Exp, [BindModifier]) -> [BindModifier] -> (Exp, [BindModifier]))
-> (Exp -> (Exp, [BindModifier]))
-> Exp
-> [BindModifier]
-> (Exp, [BindModifier])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (,[])
  where
    f :: BindModifier
-> (ExpBase f VName, [BindModifier])
-> (ExpBase f VName, [BindModifier])
f (Att AttrInfo VName
attr) (ExpBase f VName
body, [BindModifier]
modifs) =
      (AttrInfo VName -> ExpBase f VName -> SrcLoc -> ExpBase f VName
forall (f :: * -> *) vn.
AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo VName
attr ExpBase f VName
body SrcLoc
forall a. Monoid a => a
mempty, AttrInfo VName -> BindModifier
Att AttrInfo VName
attr BindModifier -> [BindModifier] -> [BindModifier]
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 Exp 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 Name) a)
  deriving
    ((forall a b. (a -> b) -> OrderingM a -> OrderingM b)
-> (forall a b. a -> OrderingM b -> OrderingM a)
-> Functor OrderingM
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
$cfmap :: forall a b. (a -> b) -> OrderingM a -> OrderingM b
fmap :: forall a b. (a -> b) -> OrderingM a -> OrderingM b
$c<$ :: forall a b. a -> OrderingM b -> OrderingM a
<$ :: forall a b. a -> OrderingM b -> OrderingM a
Functor, Functor OrderingM
Functor OrderingM =>
(forall a. a -> OrderingM a)
-> (forall a b. OrderingM (a -> b) -> OrderingM a -> OrderingM b)
-> (forall a b c.
    (a -> b -> c) -> OrderingM a -> OrderingM b -> OrderingM c)
-> (forall a b. OrderingM a -> OrderingM b -> OrderingM b)
-> (forall a b. OrderingM a -> OrderingM b -> OrderingM a)
-> Applicative 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
$cpure :: forall a. a -> OrderingM a
pure :: forall a. a -> OrderingM a
$c<*> :: forall a b. OrderingM (a -> b) -> OrderingM a -> OrderingM b
<*> :: forall a b. OrderingM (a -> b) -> OrderingM a -> OrderingM b
$cliftA2 :: forall a b c.
(a -> b -> c) -> OrderingM a -> OrderingM b -> OrderingM c
liftA2 :: forall a b c.
(a -> b -> c) -> OrderingM a -> OrderingM b -> OrderingM c
$c*> :: forall a b. OrderingM a -> OrderingM b -> OrderingM b
*> :: forall a b. OrderingM a -> OrderingM b -> OrderingM b
$c<* :: forall a b. OrderingM a -> OrderingM b -> OrderingM a
<* :: forall a b. OrderingM a -> OrderingM b -> OrderingM a
Applicative, Applicative OrderingM
Applicative OrderingM =>
(forall a b. OrderingM a -> (a -> OrderingM b) -> OrderingM b)
-> (forall a b. OrderingM a -> OrderingM b -> OrderingM b)
-> (forall a. a -> OrderingM a)
-> Monad 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
$c>>= :: forall a b. OrderingM a -> (a -> OrderingM b) -> OrderingM b
>>= :: forall a b. OrderingM a -> (a -> OrderingM b) -> OrderingM b
$c>> :: forall a b. OrderingM a -> OrderingM b -> OrderingM b
>> :: forall a b. OrderingM a -> OrderingM b -> OrderingM b
$creturn :: forall a. a -> OrderingM a
return :: forall a. a -> OrderingM a
Monad, MonadReader Name, MonadState NormState)

instance MonadFreshNames OrderingM where
  getNameSource :: OrderingM VNameSource
getNameSource = StateT
  (([Binding], [BindModifier]), VNameSource)
  (Reader Name)
  VNameSource
-> OrderingM VNameSource
forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) a
-> OrderingM a
OrderingM (StateT
   (([Binding], [BindModifier]), VNameSource)
   (Reader Name)
   VNameSource
 -> OrderingM VNameSource)
-> StateT
     (([Binding], [BindModifier]), VNameSource)
     (Reader Name)
     VNameSource
-> OrderingM VNameSource
forall a b. (a -> b) -> a -> b
$ ((([Binding], [BindModifier]), VNameSource) -> VNameSource)
-> StateT
     (([Binding], [BindModifier]), VNameSource)
     (Reader Name)
     VNameSource
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (([Binding], [BindModifier]), VNameSource) -> VNameSource
forall a b. (a, b) -> b
snd
  putNameSource :: VNameSource -> OrderingM ()
putNameSource = StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
-> OrderingM ()
forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) a
-> OrderingM a
OrderingM (StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
 -> OrderingM ())
-> (VNameSource
    -> StateT
         (([Binding], [BindModifier]), VNameSource) (Reader Name) ())
-> VNameSource
-> OrderingM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((([Binding], [BindModifier]), VNameSource)
 -> (([Binding], [BindModifier]), VNameSource))
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((([Binding], [BindModifier]), VNameSource)
  -> (([Binding], [BindModifier]), VNameSource))
 -> StateT
      (([Binding], [BindModifier]), VNameSource) (Reader Name) ())
-> (VNameSource
    -> (([Binding], [BindModifier]), VNameSource)
    -> (([Binding], [BindModifier]), VNameSource))
-> VNameSource
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VNameSource -> VNameSource)
-> (([Binding], [BindModifier]), VNameSource)
-> (([Binding], [BindModifier]), VNameSource)
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 ((VNameSource -> VNameSource)
 -> (([Binding], [BindModifier]), VNameSource)
 -> (([Binding], [BindModifier]), VNameSource))
-> (VNameSource -> VNameSource -> VNameSource)
-> VNameSource
-> (([Binding], [BindModifier]), VNameSource)
-> (([Binding], [BindModifier]), VNameSource)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VNameSource -> VNameSource -> VNameSource
forall a b. a -> b -> a
const

addModifier :: BindModifier -> OrderingM ()
addModifier :: BindModifier -> OrderingM ()
addModifier = StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
-> OrderingM ()
forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) a
-> OrderingM a
OrderingM (StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
 -> OrderingM ())
-> (BindModifier
    -> StateT
         (([Binding], [BindModifier]), VNameSource) (Reader Name) ())
-> BindModifier
-> OrderingM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((([Binding], [BindModifier]), VNameSource)
 -> (([Binding], [BindModifier]), VNameSource))
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((([Binding], [BindModifier]), VNameSource)
  -> (([Binding], [BindModifier]), VNameSource))
 -> StateT
      (([Binding], [BindModifier]), VNameSource) (Reader Name) ())
-> (BindModifier
    -> (([Binding], [BindModifier]), VNameSource)
    -> (([Binding], [BindModifier]), VNameSource))
-> BindModifier
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (([Binding], [BindModifier]) -> ([Binding], [BindModifier]))
-> (([Binding], [BindModifier]), VNameSource)
-> (([Binding], [BindModifier]), VNameSource)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((([Binding], [BindModifier]) -> ([Binding], [BindModifier]))
 -> (([Binding], [BindModifier]), VNameSource)
 -> (([Binding], [BindModifier]), VNameSource))
-> (BindModifier
    -> ([Binding], [BindModifier]) -> ([Binding], [BindModifier]))
-> BindModifier
-> (([Binding], [BindModifier]), VNameSource)
-> (([Binding], [BindModifier]), VNameSource)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([BindModifier] -> [BindModifier])
-> ([Binding], [BindModifier]) -> ([Binding], [BindModifier])
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 (([BindModifier] -> [BindModifier])
 -> ([Binding], [BindModifier]) -> ([Binding], [BindModifier]))
-> (BindModifier -> [BindModifier] -> [BindModifier])
-> BindModifier
-> ([Binding], [BindModifier])
-> ([Binding], [BindModifier])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (:)

rmModifier :: OrderingM ()
rmModifier :: OrderingM ()
rmModifier = StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
-> OrderingM ()
forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) a
-> OrderingM a
OrderingM (StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
 -> OrderingM ())
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
-> OrderingM ()
forall a b. (a -> b) -> a -> b
$ ((([Binding], [BindModifier]), VNameSource)
 -> (([Binding], [BindModifier]), VNameSource))
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((([Binding], [BindModifier]), VNameSource)
  -> (([Binding], [BindModifier]), VNameSource))
 -> StateT
      (([Binding], [BindModifier]), VNameSource) (Reader Name) ())
-> ((([Binding], [BindModifier]), VNameSource)
    -> (([Binding], [BindModifier]), VNameSource))
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
forall a b. (a -> b) -> a -> b
$ (([Binding], [BindModifier]) -> ([Binding], [BindModifier]))
-> (([Binding], [BindModifier]), VNameSource)
-> (([Binding], [BindModifier]), VNameSource)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((([Binding], [BindModifier]) -> ([Binding], [BindModifier]))
 -> (([Binding], [BindModifier]), VNameSource)
 -> (([Binding], [BindModifier]), VNameSource))
-> (([Binding], [BindModifier]) -> ([Binding], [BindModifier]))
-> (([Binding], [BindModifier]), VNameSource)
-> (([Binding], [BindModifier]), VNameSource)
forall a b. (a -> b) -> a -> b
$ ([BindModifier] -> [BindModifier])
-> ([Binding], [BindModifier]) -> ([Binding], [BindModifier])
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 [BindModifier] -> [BindModifier]
forall a. HasCallStack => [a] -> [a]
tail

addBind :: Binding -> OrderingM ()
addBind :: Binding -> OrderingM ()
addBind (PatBind [SizeBinder VName]
s Pat StructType
p Exp
e) = do
  modifs <- ((([Binding], [BindModifier]), VNameSource) -> [BindModifier])
-> OrderingM [BindModifier]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (((([Binding], [BindModifier]), VNameSource) -> [BindModifier])
 -> OrderingM [BindModifier])
-> ((([Binding], [BindModifier]), VNameSource) -> [BindModifier])
-> OrderingM [BindModifier]
forall a b. (a -> b) -> a -> b
$ ([Binding], [BindModifier]) -> [BindModifier]
forall a b. (a, b) -> b
snd (([Binding], [BindModifier]) -> [BindModifier])
-> ((([Binding], [BindModifier]), VNameSource)
    -> ([Binding], [BindModifier]))
-> (([Binding], [BindModifier]), VNameSource)
-> [BindModifier]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (([Binding], [BindModifier]), VNameSource)
-> ([Binding], [BindModifier])
forall a b. (a, b) -> a
fst
  let (e', modifs') = applyModifiers e modifs
  modify $ first $ bimap (PatBind (s <> implicit) p e' :) (const modifs')
  where
    implicit :: [SizeBinder VName]
implicit = case Exp
e of
      (AppExp AppExpBase Info VName
_ (Info (AppRes StructType
_ [VName]
ext))) -> (VName -> SizeBinder VName) -> [VName] -> [SizeBinder VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SrcLoc -> SizeBinder VName
forall vn. vn -> SrcLoc -> SizeBinder vn
`SizeBinder` SrcLoc
forall a. Monoid a => a
mempty) [VName]
ext
      Exp
_ -> []
addBind b :: Binding
b@FunBind {} =
  StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
-> OrderingM ()
forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) a
-> OrderingM a
OrderingM (StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
 -> OrderingM ())
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
-> OrderingM ()
forall a b. (a -> b) -> a -> b
$ ((([Binding], [BindModifier]), VNameSource)
 -> (([Binding], [BindModifier]), VNameSource))
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((([Binding], [BindModifier]), VNameSource)
  -> (([Binding], [BindModifier]), VNameSource))
 -> StateT
      (([Binding], [BindModifier]), VNameSource) (Reader Name) ())
-> ((([Binding], [BindModifier]), VNameSource)
    -> (([Binding], [BindModifier]), VNameSource))
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader Name) ()
forall a b. (a -> b) -> a -> b
$ (([Binding], [BindModifier]) -> ([Binding], [BindModifier]))
-> (([Binding], [BindModifier]), VNameSource)
-> (([Binding], [BindModifier]), VNameSource)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((([Binding], [BindModifier]) -> ([Binding], [BindModifier]))
 -> (([Binding], [BindModifier]), VNameSource)
 -> (([Binding], [BindModifier]), VNameSource))
-> (([Binding], [BindModifier]) -> ([Binding], [BindModifier]))
-> (([Binding], [BindModifier]), VNameSource)
-> (([Binding], [BindModifier]), VNameSource)
forall a b. (a -> b) -> a -> b
$ ([Binding] -> [Binding])
-> ([Binding], [BindModifier]) -> ([Binding], [BindModifier])
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Binding
b Binding -> [Binding] -> [Binding]
forall a. a -> [a] -> [a]
:)

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 Name) a
m) =
  (VNameSource -> ((a, [Binding]), VNameSource)) -> m (a, [Binding])
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> ((a, [Binding]), VNameSource))
 -> m (a, [Binding]))
-> (VNameSource -> ((a, [Binding]), VNameSource))
-> m (a, [Binding])
forall a b. (a -> b) -> a -> b
$ (a, (([Binding], [BindModifier]), VNameSource))
-> ((a, [Binding]), VNameSource)
forall {t :: * -> *} {a} {b} {a} {b}.
Foldable t =>
(a, ((b, t a), b)) -> ((a, b), b)
mod_tup ((a, (([Binding], [BindModifier]), VNameSource))
 -> ((a, [Binding]), VNameSource))
-> (VNameSource -> (a, (([Binding], [BindModifier]), VNameSource)))
-> VNameSource
-> ((a, [Binding]), VNameSource)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Reader Name (a, (([Binding], [BindModifier]), VNameSource))
 -> Name -> (a, (([Binding], [BindModifier]), VNameSource)))
-> Name
-> Reader Name (a, (([Binding], [BindModifier]), VNameSource))
-> (a, (([Binding], [BindModifier]), VNameSource))
forall a b c. (a -> b -> c) -> b -> a -> c
flip Reader Name (a, (([Binding], [BindModifier]), VNameSource))
-> Name -> (a, (([Binding], [BindModifier]), VNameSource))
forall r a. Reader r a -> r -> a
runReader Name
"tmp" (Reader Name (a, (([Binding], [BindModifier]), VNameSource))
 -> (a, (([Binding], [BindModifier]), VNameSource)))
-> (VNameSource
    -> Reader Name (a, (([Binding], [BindModifier]), VNameSource)))
-> VNameSource
-> (a, (([Binding], [BindModifier]), VNameSource))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) a
-> (([Binding], [BindModifier]), VNameSource)
-> Reader Name (a, (([Binding], [BindModifier]), VNameSource))
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT (([Binding], [BindModifier]), VNameSource) (Reader Name) a
m ((([Binding], [BindModifier]), VNameSource)
 -> Reader Name (a, (([Binding], [BindModifier]), VNameSource)))
-> (VNameSource -> (([Binding], [BindModifier]), VNameSource))
-> VNameSource
-> Reader Name (a, (([Binding], [BindModifier]), VNameSource))
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 t a -> Bool
forall a. t a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null t a
modifs
        then ((a
a, b
binds), b
src)
        else String -> ((a, b), b)
forall a. HasCallStack => String -> a
error String
"not all bind modifiers were freed"

naming :: Name -> OrderingM a -> OrderingM a
naming :: forall a. Name -> OrderingM a -> OrderingM a
naming Name
s = (Name -> Name) -> OrderingM a -> OrderingM a
forall a. (Name -> Name) -> OrderingM a -> OrderingM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Name -> Name -> Name
forall a b. a -> b -> a
const Name
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 = Exp -> OrderingM Exp
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
nameExp Bool
False Exp
e = do
  name <- Name -> OrderingM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName (Name -> OrderingM VName) -> OrderingM Name -> OrderingM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< OrderingM Name
forall r (m :: * -> *). MonadReader r m => m r
ask -- "e<{" ++ prettyString e ++ "}>"
  let ty = Exp -> StructType
typeOf Exp
e
      loc = Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e
      pat = VName -> Info StructType -> SrcLoc -> Pat StructType
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
name (StructType -> Info StructType
forall a. a -> Info a
Info StructType
ty) SrcLoc
loc
  addBind $ PatBind [] pat e
  pure $ Var (qualName name) (Info ty) loc

-- An evocative name to use when naming subexpressions of the
-- expression bound to this pattern.
patRepName :: Pat t -> Name
patRepName :: forall t. Pat t -> Name
patRepName (PatAscription PatBase Info VName t
p TypeExp Exp VName
_ SrcLoc
_) = PatBase Info VName t -> Name
forall t. Pat t -> Name
patRepName PatBase Info VName t
p
patRepName (Id VName
v Info t
_ SrcLoc
_) = VName -> Name
baseName VName
v
patRepName PatBase Info VName t
_ = Name
"tmp"

expRepName :: Exp -> Name
expRepName :: Exp -> Name
expRepName (Var QualName VName
v Info StructType
_ SrcLoc
_) = Text -> Name
nameFromText (Text -> Name) -> Text -> Name
forall a b. (a -> b) -> a -> b
$ QualName VName -> Text
forall a. Pretty a => a -> Text
prettyText QualName VName
v
expRepName Exp
e = Name
"d<{" Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Text -> Name
nameFromText (ExpBase NoInfo VName -> Text
forall a. Pretty a => a -> Text
prettyText (Exp -> ExpBase NoInfo VName
bareExp Exp
e)) Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"}>"

-- An evocative name to use when naming arguments to an application.
argRepName :: Exp -> Int -> Name
argRepName :: Exp -> Int -> Name
argRepName Exp
e Int
i = Exp -> Name
expRepName Exp
e Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_arg" Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Text -> Name
nameFromText (Int -> Text
forall a. Show a => a -> Text
showText 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
  ass' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
ass
  e' <- getOrdering final e
  pure $ Assert ass' e' txt loc
getOrdering Bool
final (Attr AttrInfo VName
attr Exp
e SrcLoc
loc) = do
  -- propagate attribute
  BindModifier -> OrderingM ()
addModifier (BindModifier -> OrderingM ()) -> BindModifier -> OrderingM ()
forall a b. (a -> b) -> a -> b
$ AttrInfo VName -> BindModifier
Att AttrInfo VName
attr
  e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
final Exp
e
  rmModifier
  pure $ Attr attr e' loc
getOrdering Bool
_ e :: Exp
e@Literal {} = Exp -> OrderingM Exp
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
getOrdering Bool
_ e :: Exp
e@IntLit {} = Exp -> OrderingM Exp
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
getOrdering Bool
_ e :: Exp
e@FloatLit {} = Exp -> OrderingM Exp
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
getOrdering Bool
_ e :: Exp
e@StringLit {} = Exp -> OrderingM Exp
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
getOrdering Bool
_ e :: Exp
e@Hole {} = Exp -> OrderingM Exp
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e -- can we still have some ?
getOrdering Bool
_ e :: Exp
e@Var {} = Exp -> OrderingM Exp
forall a. a -> OrderingM a
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
  es' <- (Exp -> OrderingM Exp) -> [Exp] -> OrderingM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Bool -> Exp -> OrderingM Exp
getOrdering Bool
False) [Exp]
es
  pure $ TupLit es' loc
getOrdering Bool
_ (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) = do
  fs' <- (FieldBase Info VName -> OrderingM (FieldBase Info VName))
-> [FieldBase Info VName] -> OrderingM [FieldBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM FieldBase Info VName -> OrderingM (FieldBase Info VName)
f [FieldBase Info VName]
fs
  pure $ RecordLit fs' loc
  where
    f :: FieldBase Info VName -> OrderingM (FieldBase Info VName)
f (RecordFieldExplicit L Name
n Exp
e SrcLoc
floc) = do
      e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
      pure $ RecordFieldExplicit n e' floc
    f (RecordFieldImplicit (L Loc
vloc VName
v) Info StructType
t SrcLoc
_) =
      FieldBase Info VName -> OrderingM (FieldBase Info VName)
f (FieldBase Info VName -> OrderingM (FieldBase Info VName))
-> FieldBase Info VName -> OrderingM (FieldBase Info VName)
forall a b. (a -> b) -> a -> b
$ L Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
L Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit (Loc -> Name -> L Name
forall a. Loc -> a -> L a
L Loc
vloc (VName -> Name
baseName VName
v)) (QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) Info StructType
t SrcLoc
loc) SrcLoc
loc
getOrdering Bool
_ (ArrayVal [PrimValue]
vs PrimType
t SrcLoc
loc) =
  Exp -> OrderingM Exp
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ [PrimValue] -> PrimType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[PrimValue] -> PrimType -> SrcLoc -> ExpBase f vn
ArrayVal [PrimValue]
vs PrimType
t SrcLoc
loc
getOrdering Bool
_ (ArrayLit [Exp]
es Info StructType
ty SrcLoc
loc)
  | Just [PrimValue]
vs <- (Exp -> Maybe PrimValue) -> [Exp] -> Maybe [PrimValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> Maybe PrimValue
forall {f :: * -> *} {vn}. ExpBase f vn -> Maybe PrimValue
isLiteral [Exp]
es,
    Info (Array NoMode
_ (Shape [Exp
_]) (Prim PrimType
t)) <- Info StructType
ty =
      Exp -> OrderingM Exp
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ [PrimValue] -> PrimType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[PrimValue] -> PrimType -> SrcLoc -> ExpBase f vn
ArrayVal [PrimValue]
vs PrimType
t SrcLoc
loc
  | Bool
otherwise = do
      es' <- (Exp -> OrderingM Exp) -> [Exp] -> OrderingM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Bool -> Exp -> OrderingM Exp
getOrdering Bool
False) [Exp]
es
      pure $ ArrayLit es' ty loc
  where
    isLiteral :: ExpBase f vn -> Maybe PrimValue
isLiteral (Literal PrimValue
v SrcLoc
_) = PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just PrimValue
v
    isLiteral ExpBase f vn
_ = Maybe PrimValue
forall a. Maybe a
Nothing
getOrdering Bool
_ (Project Name
n Exp
e Info StructType
ty SrcLoc
loc) = do
  e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  pure $ Project n e' ty loc
getOrdering Bool
_ (Negate Exp
e SrcLoc
loc) = do
  e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  pure $ Negate e' loc
getOrdering Bool
_ (Not Exp
e SrcLoc
loc) = do
  e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  pure $ Not e' loc
getOrdering Bool
final (Constr Name
n [Exp]
es Info StructType
ty SrcLoc
loc) = do
  es' <- (Exp -> OrderingM Exp) -> [Exp] -> OrderingM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Bool -> Exp -> OrderingM Exp
getOrdering Bool
False) [Exp]
es
  nameExp final $ Constr n es' ty loc
getOrdering Bool
final (Update Exp
eb [UpdateStep Info VName]
steps Exp
eu Info StructType
ty SrcLoc
loc) = do
  steps' <- (UpdateStep Info VName -> OrderingM (UpdateStep Info VName))
-> [UpdateStep Info VName] -> OrderingM [UpdateStep Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM UpdateStep Info VName -> OrderingM (UpdateStep Info VName)
forall {f :: * -> *} {vn}.
ASTMappable (DimIndexBase f vn) =>
UpdateStep f vn -> OrderingM (UpdateStep f vn)
onStep [UpdateStep Info VName]
steps
  eu' <- getOrdering False eu
  eb' <- getOrdering False eb
  nameExp final $ Update eb' steps' eu' ty loc
  where
    mapper :: ASTMapper OrderingM
mapper = ASTMapper OrderingM
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = getOrdering False}
    onStep :: UpdateStep f vn -> OrderingM (UpdateStep f vn)
onStep (UpdateStepSlice SliceBase f vn
slice) =
      SliceBase f vn -> UpdateStep f vn
forall (f :: * -> *) vn. SliceBase f vn -> UpdateStep f vn
UpdateStepSlice (SliceBase f vn -> UpdateStep f vn)
-> OrderingM (SliceBase f vn) -> OrderingM (UpdateStep f vn)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ASTMapper OrderingM -> SliceBase f vn -> OrderingM (SliceBase f vn)
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *).
Monad m =>
ASTMapper m -> SliceBase f vn -> m (SliceBase f vn)
astMap ASTMapper OrderingM
mapper SliceBase f vn
slice
    onStep (UpdateStepField Name
f) =
      UpdateStep f vn -> OrderingM (UpdateStep f vn)
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (UpdateStep f vn -> OrderingM (UpdateStep f vn))
-> UpdateStep f vn -> OrderingM (UpdateStep f vn)
forall a b. (a -> b) -> a -> b
$ Name -> UpdateStep f vn
forall (f :: * -> *) vn. Name -> UpdateStep f vn
UpdateStepField Name
f
getOrdering Bool
final (Lambda [PatBase Info VName ParamType]
params Exp
body Maybe (TypeExp Exp VName)
mte Info ResRetType
ret SrcLoc
loc) = do
  body' <- Exp -> OrderingM Exp
forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
body
  nameExp final $ Lambda params body' mte ret loc
getOrdering Bool
_ (OpSection QualName VName
qn Info StructType
ty SrcLoc
loc) =
  Exp -> OrderingM Exp
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info StructType -> SrcLoc -> Exp
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 Freshness
ret), Info [VName]
exts) SrcLoc
loc) = do
  x <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  yn <- newVName "y"
  let y = QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
yn) (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ ParamType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ParamType
yty) SrcLoc
forall a. Monoid a => a
mempty
      ret' = TypeSubs -> TypeBase Exp Freshness -> TypeBase Exp Freshness
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (Exp -> Exp -> TypeSubs
forall {t}. Exp -> Exp -> VName -> Maybe (Subst t)
pSubst Exp
x Exp
y) TypeBase Exp Freshness
ret
      body =
        Exp -> [(Maybe VName, Exp)] -> AppRes -> Exp
forall vn.
ExpBase Info vn
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
mkApply (QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
op Info StructType
ty SrcLoc
loc) [(Maybe VName
xext, Exp
x), (Maybe VName
forall a. Maybe a
Nothing, Exp
y)] (AppRes -> Exp) -> AppRes -> Exp
forall a b. (a -> b) -> a -> b
$
          StructType -> [VName] -> AppRes
AppRes (TypeBase Exp Freshness -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase Exp Freshness
ret') [VName]
exts
  nameExp final $ Lambda [Id yn (Info yty) mempty] body Nothing (Info (RetType dims ret')) loc
  where
    pSubst :: Exp -> Exp -> VName -> Maybe (Subst t)
pSubst Exp
x Exp
y VName
vn
      | Named VName
p <- PName
xp, VName
p VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
vn = Subst t -> Maybe (Subst t)
forall a. a -> Maybe a
Just (Subst t -> Maybe (Subst t)) -> Subst t -> Maybe (Subst t)
forall a b. (a -> b) -> a -> b
$ Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst Exp
x
      | Named VName
p <- PName
yp, VName
p VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
vn = Subst t -> Maybe (Subst t)
forall a. a -> Maybe a
Just (Subst t -> Maybe (Subst t)) -> Subst t -> Maybe (Subst t)
forall a b. (a -> b) -> a -> b
$ Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst Exp
y
      | Bool
otherwise = Maybe (Subst t)
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 Freshness
ret)) SrcLoc
loc) = do
  xn <- Name -> OrderingM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"x"
  y <- getOrdering False e
  let x = QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
xn) (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ ParamType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ParamType
xty) SrcLoc
forall a. Monoid a => a
mempty
      ret' = TypeSubs -> TypeBase Exp Freshness -> TypeBase Exp Freshness
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (Exp -> Exp -> TypeSubs
forall {t}. Exp -> Exp -> VName -> Maybe (Subst t)
pSubst Exp
x Exp
y) TypeBase Exp Freshness
ret
      body = Exp -> [(Maybe VName, Exp)] -> AppRes -> Exp
forall vn.
ExpBase Info vn
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
mkApply (QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
op Info StructType
ty SrcLoc
loc) [(Maybe VName
forall a. Maybe a
Nothing, Exp
x), (Maybe VName
yext, Exp
y)] (AppRes -> Exp) -> AppRes -> Exp
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (TypeBase Exp Freshness -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase Exp Freshness
ret') []
  nameExp final $ Lambda [Id xn (Info xty) mempty] body Nothing (Info (RetType dims ret')) loc
  where
    pSubst :: Exp -> Exp -> VName -> Maybe (Subst t)
pSubst Exp
x Exp
y VName
vn
      | Named VName
p <- PName
xp, VName
p VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
vn = Subst t -> Maybe (Subst t)
forall a. a -> Maybe a
Just (Subst t -> Maybe (Subst t)) -> Subst t -> Maybe (Subst t)
forall a b. (a -> b) -> a -> b
$ Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst Exp
x
      | Named VName
p <- PName
yp, VName
p VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
vn = Subst t -> Maybe (Subst t)
forall a. a -> Maybe a
Just (Subst t -> Maybe (Subst t)) -> Subst t -> Maybe (Subst t)
forall a b. (a -> b) -> a -> b
$ Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst Exp
y
      | Bool
otherwise = Maybe (Subst t)
forall a. Maybe a
Nothing
getOrdering Bool
final (UpdateSection [UpdateStep Info VName]
steps (Info StructType
ty) SrcLoc
loc) = do
  steps' <- (UpdateStep Info VName -> OrderingM (UpdateStep Info VName))
-> [UpdateStep Info VName] -> OrderingM [UpdateStep Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM UpdateStep Info VName -> OrderingM (UpdateStep Info VName)
forall {f :: * -> *} {vn}.
ASTMappable (DimIndexBase f vn) =>
UpdateStep f vn -> OrderingM (UpdateStep f vn)
transformStep [UpdateStep Info VName]
steps
  xn <- newVName "x"
  let (xty, RetType dims ret) = case ty of
        Scalar (Arrow NoMode
_ PName
_ Diet
d StructType
xty' ResRetType
ret') -> (Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
d StructType
xty', ResRetType
ret')
        StructType
_ -> String -> (ParamType, ResRetType)
forall a. HasCallStack => String -> a
error (String -> (ParamType, ResRetType))
-> String -> (ParamType, ResRetType)
forall a b. (a -> b) -> a -> b
$ String
"not a function type for section update: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ StructType -> String
forall a. Pretty a => a -> String
prettyString StructType
ty
      x = QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
xn) (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ ParamType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ParamType
xty) SrcLoc
forall a. Monoid a => a
mempty
      body = (Exp, StructType) -> Exp
forall a b. (a, b) -> a
fst ((Exp, StructType) -> Exp) -> (Exp, StructType) -> Exp
forall a b. (a -> b) -> a -> b
$ ((Exp, StructType) -> UpdateStep Info VName -> (Exp, StructType))
-> (Exp, StructType)
-> [UpdateStep Info VName]
-> (Exp, StructType)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Exp, StructType) -> UpdateStep Info VName -> (Exp, StructType)
forall {vn}.
(ExpBase Info vn, StructType)
-> UpdateStep Info vn -> (ExpBase Info vn, StructType)
applyStep (Exp
x, ParamType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ParamType
xty) [UpdateStep Info VName]
steps'
  nameExp final $ Lambda [Id xn (Info xty) mempty] body Nothing (Info (RetType dims ret)) loc
  where
    mapper :: ASTMapper OrderingM
mapper = ASTMapper OrderingM
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = getOrdering False}

    transformStep :: UpdateStep f vn -> OrderingM (UpdateStep f vn)
transformStep (UpdateStepField Name
f) =
      UpdateStep f vn -> OrderingM (UpdateStep f vn)
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (UpdateStep f vn -> OrderingM (UpdateStep f vn))
-> UpdateStep f vn -> OrderingM (UpdateStep f vn)
forall a b. (a -> b) -> a -> b
$ Name -> UpdateStep f vn
forall (f :: * -> *) vn. Name -> UpdateStep f vn
UpdateStepField Name
f
    transformStep (UpdateStepSlice SliceBase f vn
slice) =
      SliceBase f vn -> UpdateStep f vn
forall (f :: * -> *) vn. SliceBase f vn -> UpdateStep f vn
UpdateStepSlice (SliceBase f vn -> UpdateStep f vn)
-> OrderingM (SliceBase f vn) -> OrderingM (UpdateStep f vn)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ASTMapper OrderingM -> SliceBase f vn -> OrderingM (SliceBase f vn)
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *).
Monad m =>
ASTMapper m -> SliceBase f vn -> m (SliceBase f vn)
astMap ASTMapper OrderingM
mapper SliceBase f vn
slice

    applyStep :: (ExpBase Info vn, StructType)
-> UpdateStep Info vn -> (ExpBase Info vn, StructType)
applyStep (ExpBase Info vn
e, StructType
t) (UpdateStepField Name
field) =
      case StructType
t of
        Scalar (Record Map Name StructType
fs)
          | Just StructType
t' <- Name -> Map Name StructType -> Maybe StructType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
field Map Name StructType
fs ->
              (Name
-> ExpBase Info vn -> Info StructType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
field ExpBase Info vn
e (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t') SrcLoc
forall a. Monoid a => a
mempty, StructType
t')
        StructType
_ ->
          String -> (ExpBase Info vn, StructType)
forall a. HasCallStack => String -> a
error (String -> (ExpBase Info vn, StructType))
-> String -> (ExpBase Info vn, StructType)
forall a b. (a -> b) -> a -> b
$
            String
"desugar UpdateSection: type "
              String -> String -> String
forall a. [a] -> [a] -> [a]
++ StructType -> String
forall a. Pretty a => a -> String
prettyString StructType
t
              String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" does not have field "
              String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Pretty a => a -> String
prettyString Name
field
    applyStep (ExpBase Info vn
e, StructType
t) (UpdateStepSlice SliceBase Info vn
slice) =
      let t' :: StructType
t' = Int -> StructType -> StructType
forall dim o. Int -> TypeBase dim o -> TypeBase dim o
stripArray (SliceBase Info vn -> Int
forall {f :: * -> *} {vn}. [DimIndexBase f vn] -> Int
fixedDims SliceBase Info vn
slice) StructType
t
          e' :: ExpBase Info vn
e' = AppExpBase Info vn -> Info AppRes -> ExpBase Info vn
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (ExpBase Info vn
-> SliceBase Info vn -> SrcLoc -> AppExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
Index ExpBase Info vn
e SliceBase Info vn
slice SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info (StructType -> [VName] -> AppRes
AppRes StructType
t' []))
       in (ExpBase Info vn
e', StructType
t')

    fixedDims :: [DimIndexBase f vn] -> Int
fixedDims = [DimIndexBase f vn] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([DimIndexBase f vn] -> Int)
-> ([DimIndexBase f vn] -> [DimIndexBase f vn])
-> [DimIndexBase f vn]
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (DimIndexBase f vn -> Bool)
-> [DimIndexBase f vn] -> [DimIndexBase f vn]
forall a. (a -> Bool) -> [a] -> [a]
filter DimIndexBase f vn -> Bool
forall {f :: * -> *} {vn}. DimIndexBase f vn -> Bool
isFix
    isFix :: DimIndexBase f vn -> Bool
isFix DimFix {} = Bool
True
    isFix DimIndexBase f vn
_ = Bool
False
getOrdering Bool
_ (Ascript Exp
e TypeExp Exp VName
_ SrcLoc
_) = Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
getOrdering Bool
final (AppExp (Apply Exp
f NonEmpty (Info (Maybe VName), Exp)
args SrcLoc
loc) Info AppRes
resT) = do
  args' <-
    NonEmpty (Info (Maybe VName), Exp)
-> NonEmpty (Info (Maybe VName), Exp)
forall a. NonEmpty a -> NonEmpty a
NE.reverse (NonEmpty (Info (Maybe VName), Exp)
 -> NonEmpty (Info (Maybe VName), Exp))
-> OrderingM (NonEmpty (Info (Maybe VName), Exp))
-> OrderingM (NonEmpty (Info (Maybe VName), Exp))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (((Info (Maybe VName), Exp), Int)
 -> OrderingM (Info (Maybe VName), Exp))
-> NonEmpty ((Info (Maybe VName), Exp), Int)
-> OrderingM (NonEmpty (Info (Maybe VName), Exp))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> NonEmpty a -> m (NonEmpty b)
mapM ((Info (Maybe VName), Exp), Int)
-> OrderingM (Info (Maybe VName), Exp)
forall {t}. ((t, Exp), Int) -> OrderingM (t, Exp)
onArg (NonEmpty ((Info (Maybe VName), Exp), Int)
-> NonEmpty ((Info (Maybe VName), Exp), Int)
forall a. NonEmpty a -> NonEmpty a
NE.reverse (NonEmpty (Info (Maybe VName), Exp)
-> NonEmpty Int -> NonEmpty ((Info (Maybe VName), Exp), Int)
forall a b. NonEmpty a -> NonEmpty b -> NonEmpty (a, b)
NE.zip NonEmpty (Info (Maybe VName), Exp)
args ([Int] -> NonEmpty Int
forall a. HasCallStack => [a] -> NonEmpty a
NE.fromList [Int
0 ..])))
  f' <- getOrdering False f
  nameExp final $ AppExp (Apply f' args' loc) resT
  where
    onArg :: ((t, Exp), Int) -> OrderingM (t, Exp)
onArg ((t
d, Exp
e), Int
i) =
      Name -> OrderingM (t, Exp) -> OrderingM (t, Exp)
forall a. Name -> OrderingM a -> OrderingM a
naming (Exp -> Int -> Name
argRepName Exp
f Int
i) (OrderingM (t, Exp) -> OrderingM (t, Exp))
-> OrderingM (t, Exp) -> OrderingM (t, Exp)
forall a b. (a -> b) -> a -> b
$ (t
d,) (Exp -> (t, Exp)) -> OrderingM Exp -> OrderingM (t, Exp)
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 Exp VName
te Info StructType
t SrcLoc
loc) = do
  e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  nameExp final $ Coerce e' te t loc
getOrdering Bool
final (AppExp (Range Exp
start Maybe Exp
stride Inclusiveness Exp
end SrcLoc
loc) Info AppRes
resT) = do
  start' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
start
  stride' <- mapM (getOrdering False) stride
  end' <- mapM (getOrdering False) end
  nameExp final $ AppExp (Range start' stride' end' loc) resT
getOrdering Bool
final (AppExp (LetPat [SizeBinder VName]
sizes Pat StructType
pat Exp
expr Exp
body SrcLoc
_) Info AppRes
_) = do
  expr' <- Name -> OrderingM Exp -> OrderingM Exp
forall a. Name -> OrderingM a -> OrderingM a
naming (Pat StructType -> Name
forall t. Pat t -> Name
patRepName Pat StructType
pat) (OrderingM Exp -> OrderingM Exp) -> OrderingM Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
expr
  addBind $ PatBind sizes pat expr'
  getOrdering final body
getOrdering Bool
final (AppExp (LetFun (VName
vn, SrcLoc
_) ([TypeParamBase VName]
tparams, [PatBase Info VName ParamType]
params, Maybe (TypeExp Exp VName)
mrettype, Info ResRetType
rettype, Exp
body) Exp
e SrcLoc
_) Info AppRes
_) = do
  body' <- Exp -> OrderingM Exp
forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
body
  addBind $ FunBind vn (tparams, params, mrettype, rettype, body')
  getOrdering final e
getOrdering Bool
final (AppExp (If Exp
cond Exp
et Exp
ef SrcLoc
loc) Info AppRes
resT) = do
  cond' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
cond
  et' <- transformBody et
  ef' <- transformBody ef
  nameExp final $ AppExp (If cond' et' ef' loc) resT
getOrdering Bool
final (AppExp (Loop [VName]
sizes PatBase Info VName ParamType
pat LoopInitBase Info VName
einit LoopFormBase Info VName
form Exp
body SrcLoc
loc) Info AppRes
resT) = do
  einit' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ LoopInitBase Info VName -> Exp
loopInitExp LoopInitBase Info VName
einit
  form' <- case form of
    For IdentBase Info VName StructType
ident Exp
e -> IdentBase Info VName StructType -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn.
IdentBase f vn StructType -> ExpBase f vn -> LoopFormBase f vn
For IdentBase Info VName StructType
ident (Exp -> LoopFormBase Info VName)
-> OrderingM Exp -> OrderingM (LoopFormBase Info VName)
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 -> Pat StructType -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn.
PatBase f vn StructType -> ExpBase f vn -> LoopFormBase f vn
ForIn Pat StructType
fpat (Exp -> LoopFormBase Info VName)
-> OrderingM Exp -> OrderingM (LoopFormBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
e
    While Exp
e -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> LoopFormBase f vn
While (Exp -> LoopFormBase Info VName)
-> OrderingM Exp -> OrderingM (LoopFormBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> OrderingM Exp
forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
e
  body' <- transformBody body
  nameExp final $ AppExp (Loop sizes pat (LoopInitExplicit einit') form' body' loc) 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
  expr' <- case (Bool
isOr, Bool
isAnd) of
    (Bool
True, Bool
_) -> do
      el' <- Name -> OrderingM Exp -> OrderingM Exp
forall a. Name -> OrderingM a -> OrderingM a
naming Name
"or_lhs" (OrderingM Exp -> OrderingM Exp) -> OrderingM Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
el
      er' <- naming "or_rhs" $ transformBody er
      pure $ AppExp (If el' (Literal (BoolValue True) mempty) er' loc) (Info resT)
    (Bool
_, Bool
True) -> do
      el' <- Name -> OrderingM Exp -> OrderingM Exp
forall a. Name -> OrderingM a -> OrderingM a
naming Name
"and_lhs" (OrderingM Exp -> OrderingM Exp) -> OrderingM Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
el
      er' <- naming "and_rhs" $ transformBody er
      pure $ AppExp (If el' er' (Literal (BoolValue False) mempty) loc) (Info resT)
    (Bool
False, Bool
False) -> do
      el' <- Name -> OrderingM Exp -> OrderingM Exp
forall a. Name -> OrderingM a -> OrderingM a
naming (Text -> Name
nameFromText (QualName VName -> Text
forall a. Pretty a => a -> Text
prettyText QualName VName
op) Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_lhs") (OrderingM Exp -> OrderingM Exp) -> OrderingM Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
el
      er' <- naming (nameFromText (prettyText op) <> "_rhs") $ getOrdering False er
      pure $ mkApply (Var op opT oloc) [(elp, el'), (erp, er')] resT
  nameExp final expr'
  where
    isOr :: Bool
isOr = VName -> Name
baseName (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
op) Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
"||"
    isAnd :: Bool
isAnd = VName -> Name
baseName (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
op) Name -> Name -> Bool
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) [UpdateStep Info VName]
steps Exp
e Exp
body SrcLoc
_) Info AppRes
_) = do
  steps' <- (UpdateStep Info VName -> OrderingM (UpdateStep Info VName))
-> [UpdateStep Info VName] -> OrderingM [UpdateStep Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM UpdateStep Info VName -> OrderingM (UpdateStep Info VName)
forall {f :: * -> *} {vn}.
ASTMappable (DimIndexBase f vn) =>
UpdateStep f vn -> OrderingM (UpdateStep f vn)
onStep [UpdateStep Info VName]
steps
  e' <- getOrdering False e
  let loc' = SrcLoc -> Exp -> SrcLoc
forall a b. (Located a, Located b) => a -> b -> SrcLoc
srcspan SrcLoc
dloc Exp
e
  addBind $
    PatBind
      []
      (Id dest dty dloc)
      (Update (Var (qualName src) sty sloc) steps' e' (Info (unInfo sty)) loc')
  getOrdering final body
  where
    mapper :: ASTMapper OrderingM
mapper = ASTMapper OrderingM
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = getOrdering False}
    onStep :: UpdateStep f vn -> OrderingM (UpdateStep f vn)
onStep (UpdateStepSlice SliceBase f vn
slice) = SliceBase f vn -> UpdateStep f vn
forall (f :: * -> *) vn. SliceBase f vn -> UpdateStep f vn
UpdateStepSlice (SliceBase f vn -> UpdateStep f vn)
-> OrderingM (SliceBase f vn) -> OrderingM (UpdateStep f vn)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ASTMapper OrderingM -> SliceBase f vn -> OrderingM (SliceBase f vn)
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *).
Monad m =>
ASTMapper m -> SliceBase f vn -> m (SliceBase f vn)
astMap ASTMapper OrderingM
mapper SliceBase f vn
slice
    onStep (UpdateStepField Name
f) = UpdateStep f vn -> OrderingM (UpdateStep f vn)
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (UpdateStep f vn -> OrderingM (UpdateStep f vn))
-> UpdateStep f vn -> OrderingM (UpdateStep f vn)
forall a b. (a -> b) -> a -> b
$ Name -> UpdateStep f vn
forall (f :: * -> *) vn. Name -> UpdateStep f vn
UpdateStepField Name
f
getOrdering Bool
final (AppExp (Index Exp
e SliceBase Info VName
slice SrcLoc
loc) Info AppRes
resT) = do
  e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  slice' <- astMap mapper slice
  nameExp final $ AppExp (Index e' slice' loc) resT
  where
    mapper :: ASTMapper OrderingM
mapper = ASTMapper OrderingM
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = getOrdering False}
getOrdering Bool
final (AppExp (Match Exp
expr NonEmpty (CaseBase Info VName)
cs SrcLoc
loc) Info AppRes
resT) = do
  expr' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
expr
  cs' <- mapM f cs
  nameExp final $ AppExp (Match expr' cs' loc) resT
  where
    f :: CaseBase Info VName -> m (CaseBase Info VName)
f (CasePat Pat StructType
pat Exp
body SrcLoc
cloc) = do
      body' <- Exp -> m Exp
forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
body
      pure (CasePat pat body' 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
  (e', pre_eval) <- OrderingM Exp -> m (Exp, [Binding])
forall (m :: * -> *) a.
MonadFreshNames m =>
OrderingM a -> m (a, [Binding])
runOrdering (Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
e)
  pure $ foldl f e' pre_eval
  where
    appRes :: Info AppRes
appRes = case Exp
e of
      (AppExp AppExpBase Info VName
_ Info AppRes
r) -> Info AppRes
r
      Exp
_ -> AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
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) =
      AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp ([SizeBinder VName]
-> Pat StructType -> Exp -> Exp -> SrcLoc -> AppExpBase Info VName
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 SrcLoc
forall a. Monoid a => a
mempty) Info AppRes
appRes
    f Exp
body (FunBind VName
vn ([TypeParamBase VName], [PatBase Info VName ParamType],
 Maybe (TypeExp Exp VName), Info ResRetType, Exp)
infos) =
      AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp ((VName, SrcLoc)
-> ([TypeParamBase VName], [PatBase Info VName ParamType],
    Maybe (TypeExp Exp VName), Info ResRetType, Exp)
-> Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
(vn, SrcLoc)
-> ([TypeParamBase vn], [PatBase f vn ParamType],
    Maybe (TypeExp (ExpBase f vn) vn), f ResRetType, ExpBase f vn)
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetFun (VName
vn, SrcLoc
forall a. Monoid a => a
mempty) ([TypeParamBase VName], [PatBase Info VName ParamType],
 Maybe (TypeExp Exp VName), Info ResRetType, Exp)
infos Exp
body SrcLoc
forall a. Monoid a => a
mempty) Info AppRes
appRes

transformValBind :: (MonadFreshNames m) => ValBind -> m ValBind
transformValBind :: forall (m :: * -> *). MonadFreshNames m => ValBind -> m ValBind
transformValBind ValBind
valbind = do
  body' <- Exp -> m Exp
forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody (Exp -> m Exp) -> Exp -> m Exp
forall a b. (a -> b) -> a -> b
$ ValBind -> Exp
forall (f :: * -> *) vn. ValBindBase f vn -> ExpBase f vn
valBindBody ValBind
valbind
  pure $ valbind {valBindBody = body'}

-- | Fully normalise top level bindings.
transformProg :: (MonadFreshNames m) => [ValBind] -> m [ValBind]
transformProg :: forall (m :: * -> *). MonadFreshNames m => [ValBind] -> m [ValBind]
transformProg = (ValBind -> m ValBind) -> [ValBind] -> m [ValBind]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ValBind -> m ValBind
forall (m :: * -> *). MonadFreshNames m => ValBind -> m ValBind
transformValBind