-- | 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 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 =
  (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])
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 -> (Exp, [BindModifier]) -> (Exp, [BindModifier])
f (Ass Exp
ass Info Text
txt SrcLoc
loc) (Exp
body, [BindModifier]
modifs) =
      (Exp -> Exp -> Info Text -> SrcLoc -> Exp
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) =
      (AttrInfo VName -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo VName
attr Exp
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 String) 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 String, MonadState NormState)

instance MonadFreshNames OrderingM where
  getNameSource :: OrderingM VNameSource
getNameSource = StateT
  (([Binding], [BindModifier]), VNameSource)
  (Reader [Char])
  VNameSource
-> OrderingM VNameSource
forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM (StateT
   (([Binding], [BindModifier]), VNameSource)
   (Reader [Char])
   VNameSource
 -> OrderingM VNameSource)
-> StateT
     (([Binding], [BindModifier]), VNameSource)
     (Reader [Char])
     VNameSource
-> OrderingM VNameSource
forall a b. (a -> b) -> a -> b
$ ((([Binding], [BindModifier]), VNameSource) -> VNameSource)
-> StateT
     (([Binding], [BindModifier]), VNameSource)
     (Reader [Char])
     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 [Char]) ()
-> OrderingM ()
forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM (StateT
   (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
 -> OrderingM ())
-> (VNameSource
    -> StateT
         (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ())
-> VNameSource
-> OrderingM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((([Binding], [BindModifier]), VNameSource)
 -> (([Binding], [BindModifier]), VNameSource))
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((([Binding], [BindModifier]), VNameSource)
  -> (([Binding], [BindModifier]), VNameSource))
 -> StateT
      (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ())
-> (VNameSource
    -> (([Binding], [BindModifier]), VNameSource)
    -> (([Binding], [BindModifier]), VNameSource))
-> VNameSource
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
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 [Char]) ()
-> OrderingM ()
forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM (StateT
   (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
 -> OrderingM ())
-> (BindModifier
    -> StateT
         (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ())
-> BindModifier
-> OrderingM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((([Binding], [BindModifier]), VNameSource)
 -> (([Binding], [BindModifier]), VNameSource))
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((([Binding], [BindModifier]), VNameSource)
  -> (([Binding], [BindModifier]), VNameSource))
 -> StateT
      (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ())
-> (BindModifier
    -> (([Binding], [BindModifier]), VNameSource)
    -> (([Binding], [BindModifier]), VNameSource))
-> BindModifier
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
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 [Char]) ()
-> OrderingM ()
forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM (StateT
   (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
 -> OrderingM ())
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
-> OrderingM ()
forall a b. (a -> b) -> a -> b
$ ((([Binding], [BindModifier]), VNameSource)
 -> (([Binding], [BindModifier]), VNameSource))
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((([Binding], [BindModifier]), VNameSource)
  -> (([Binding], [BindModifier]), VNameSource))
 -> StateT
      (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ())
-> ((([Binding], [BindModifier]), VNameSource)
    -> (([Binding], [BindModifier]), VNameSource))
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
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 [Char]) ()
-> OrderingM ()
forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM (StateT
   (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
 -> OrderingM ())
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
-> OrderingM ()
forall a b. (a -> b) -> a -> b
$ ((([Binding], [BindModifier]), VNameSource)
 -> (([Binding], [BindModifier]), VNameSource))
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((([Binding], [BindModifier]), VNameSource)
  -> (([Binding], [BindModifier]), VNameSource))
 -> StateT
      (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ())
-> ((([Binding], [BindModifier]), VNameSource)
    -> (([Binding], [BindModifier]), VNameSource))
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) ()
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 [Char]) 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 [Char] (a, (([Binding], [BindModifier]), VNameSource))
 -> [Char] -> (a, (([Binding], [BindModifier]), VNameSource)))
-> [Char]
-> Reader [Char] (a, (([Binding], [BindModifier]), VNameSource))
-> (a, (([Binding], [BindModifier]), VNameSource))
forall a b c. (a -> b -> c) -> b -> a -> c
flip Reader [Char] (a, (([Binding], [BindModifier]), VNameSource))
-> [Char] -> (a, (([Binding], [BindModifier]), VNameSource))
forall r a. Reader r a -> r -> a
runReader [Char]
"tmp" (Reader [Char] (a, (([Binding], [BindModifier]), VNameSource))
 -> (a, (([Binding], [BindModifier]), VNameSource)))
-> (VNameSource
    -> Reader [Char] (a, (([Binding], [BindModifier]), VNameSource)))
-> VNameSource
-> (a, (([Binding], [BindModifier]), VNameSource))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> (([Binding], [BindModifier]), VNameSource)
-> Reader [Char] (a, (([Binding], [BindModifier]), VNameSource))
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
m ((([Binding], [BindModifier]), VNameSource)
 -> Reader [Char] (a, (([Binding], [BindModifier]), VNameSource)))
-> (VNameSource -> (([Binding], [BindModifier]), VNameSource))
-> VNameSource
-> Reader [Char] (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 [Char] -> ((a, b), b)
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 = ([Char] -> [Char]) -> OrderingM a -> OrderingM a
forall a. ([Char] -> [Char]) -> OrderingM a -> OrderingM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ([Char] -> [Char] -> [Char]
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 = 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 <- [Char] -> OrderingM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString ([Char] -> OrderingM VName) -> OrderingM [Char] -> OrderingM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< OrderingM [Char]
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 -> String
patRepName :: forall t. Pat t -> [Char]
patRepName (PatAscription PatBase Info VName t
p TypeExp Exp VName
_ SrcLoc
_) = PatBase Info VName t -> [Char]
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
_) = QualName VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString QualName VName
v
expRepName Exp
e = [Char]
"d<{" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ ExpBase NoInfo VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString (Exp -> ExpBase NoInfo VName
bareExp Exp
e) [Char] -> [Char] -> [Char]
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 [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_arg" [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
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
  ass' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
ass
  l_prev <- OrderingM $ gets $ length . snd . fst
  addModifier $ Ass ass' txt loc
  e' <- getOrdering final e
  l_after <- OrderingM $ gets $ length . snd . 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 l_after <= l_prev
    then pure e'
    else do
      rmModifier
      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 NoUniqueness
_ (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 SliceBase Info VName
slice Exp
eu SrcLoc
loc) = do
  eu' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
eu
  slice' <- astMap mapper slice
  eb' <- getOrdering False eb
  nameExp final $ Update eb' slice' eu' loc
  where
    mapper :: ASTMapper OrderingM
mapper = ASTMapper OrderingM
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = getOrdering False}
getOrdering Bool
final (RecordUpdate Exp
eb [Name]
ns Exp
eu Info StructType
ty SrcLoc
loc) = do
  eb' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
eb
  eu' <- getOrdering False eu
  nameExp final $ RecordUpdate eb' ns eu' ty loc
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 Uniqueness
ret), Info [VName]
exts) SrcLoc
loc) = do
  x <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  yn <- newNameFromString "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 u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ParamType
yty) SrcLoc
forall a. Monoid a => a
mempty
      ret' = TypeSubs -> TypeBase Exp Uniqueness -> TypeBase Exp Uniqueness
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 Uniqueness
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 Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp Uniqueness
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 Uniqueness
ret)) SrcLoc
loc) = do
  xn <- [Char] -> OrderingM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"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 u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ParamType
xty) SrcLoc
forall a. Monoid a => a
mempty
      ret' = TypeSubs -> TypeBase Exp Uniqueness -> TypeBase Exp Uniqueness
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 Uniqueness
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 Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp Uniqueness
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 (ProjectSection [Name]
names (Info StructType
ty) SrcLoc
loc) = do
  xn <- [Char] -> OrderingM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"x"
  let (xty, RetType dims ret) = case ty of
        Scalar (Arrow NoUniqueness
_ PName
_ Diet
d StructType
xty' ResRetType
ret') -> (Diet -> StructType -> ParamType
forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
d StructType
xty', ResRetType
ret')
        StructType
_ -> [Char] -> (ParamType, ResRetType)
forall a. HasCallStack => [Char] -> a
error ([Char] -> (ParamType, ResRetType))
-> [Char] -> (ParamType, ResRetType)
forall a b. (a -> b) -> a -> b
$ [Char]
"not a function type for project section: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
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 u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ParamType
xty) SrcLoc
forall a. Monoid a => a
mempty
      body = (Exp -> Name -> Exp) -> Exp -> [Name] -> Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Exp -> Name -> Exp
project Exp
x [Name]
names
  nameExp final $ Lambda [Id xn (Info xty) mempty] body Nothing (Info (RetType dims ret)) 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 <- 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 -> Exp -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
field Exp
e (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t) SrcLoc
forall a. Monoid a => a
mempty
        StructType
t ->
          [Char] -> Exp
forall a. HasCallStack => [Char] -> a
error ([Char] -> Exp) -> [Char] -> Exp
forall a b. (a -> b) -> a -> b
$
            [Char]
"desugar ProjectSection: type "
              [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
t
              [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" does not have field "
              [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Name
field
getOrdering Bool
final (IndexSection SliceBase Info VName
slice (Info StructType
ty) SrcLoc
loc) = do
  slice' <- ASTMapper OrderingM
-> SliceBase Info VName -> OrderingM (SliceBase Info VName)
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *).
Monad m =>
ASTMapper m -> SliceBase Info VName -> m (SliceBase Info VName)
astMap ASTMapper OrderingM
mapper SliceBase Info VName
slice
  xn <- newNameFromString "x"
  let (xty, RetType dims ret) = case ty of
        Scalar (Arrow NoUniqueness
_ PName
_ Diet
d StructType
xty' ResRetType
ret') -> (Diet -> StructType -> ParamType
forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
d StructType
xty', ResRetType
ret')
        StructType
_ -> [Char] -> (ParamType, ResRetType)
forall a. HasCallStack => [Char] -> a
error ([Char] -> (ParamType, ResRetType))
-> [Char] -> (ParamType, ResRetType)
forall a b. (a -> b) -> a -> b
$ [Char]
"not a function type for index section: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
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 u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ParamType
xty) SrcLoc
forall a. Monoid a => a
mempty
      body = AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp -> SliceBase Info VName -> SrcLoc -> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
Index Exp
x SliceBase Info VName
slice' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info (StructType -> [VName] -> AppRes
AppRes (TypeBase Exp Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp Uniqueness
ret) []))
  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}
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) =
      [Char] -> OrderingM (t, Exp) -> OrderingM (t, Exp)
forall a. [Char] -> OrderingM a -> OrderingM a
naming (Exp -> Int -> [Char]
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' <- [Char] -> OrderingM Exp -> OrderingM Exp
forall a. [Char] -> OrderingM a -> OrderingM a
naming (Pat StructType -> [Char]
forall t. Pat t -> [Char]
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 ([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' <- [Char] -> OrderingM Exp -> OrderingM Exp
forall a. [Char] -> OrderingM a -> OrderingM a
naming [Char]
"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' <- [Char] -> OrderingM Exp -> OrderingM Exp
forall a. [Char] -> OrderingM a -> OrderingM a
naming [Char]
"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' <- [Char] -> OrderingM Exp -> OrderingM Exp
forall a. [Char] -> OrderingM a -> OrderingM a
naming (QualName VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString QualName VName
op [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_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 (prettyString 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) SliceBase Info VName
slice Exp
e Exp
body SrcLoc
loc) Info AppRes
_) = do
  e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  slice' <- astMap mapper slice
  addBind $ PatBind [] (Id dest dty dloc) (Update (Var (qualName src) sty sloc) slice' e' loc)
  getOrdering final body
  where
    mapper :: ASTMapper OrderingM
mapper = ASTMapper OrderingM
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = getOrdering False}
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
-> ([TypeParamBase VName], [PatBase Info VName ParamType],
    Maybe (TypeExp Exp VName), Info ResRetType, Exp)
-> Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
vn
-> ([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 ([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