-- | 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
  [BindModifier]
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 (Exp
e', [BindModifier]
modifs') = Exp -> [BindModifier] -> (Exp, [BindModifier])
applyModifiers Exp
e [BindModifier]
modifs
  ((([Binding], [BindModifier]), VNameSource)
 -> (([Binding], [BindModifier]), VNameSource))
-> OrderingM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((([Binding], [BindModifier]), VNameSource)
  -> (([Binding], [BindModifier]), VNameSource))
 -> OrderingM ())
-> ((([Binding], [BindModifier]), VNameSource)
    -> (([Binding], [BindModifier]), VNameSource))
-> OrderingM ()
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])
-> ([BindModifier] -> [BindModifier])
-> ([Binding], [BindModifier])
-> ([Binding], [BindModifier])
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap ([SizeBinder VName] -> Pat StructType -> Exp -> Binding
PatBind ([SizeBinder VName]
s [SizeBinder VName] -> [SizeBinder VName] -> [SizeBinder VName]
forall a. Semigroup a => a -> a -> a
<> [SizeBinder VName]
implicit) Pat StructType
p Exp
e' :) ([BindModifier] -> [BindModifier] -> [BindModifier]
forall a b. a -> b -> a
const [BindModifier]
modifs')
  where
    implicit :: [SizeBinder VName]
implicit = case Exp
e of
      (AppExp AppExpBase Info VName
_ (Info (AppRes StructType
_ [VName]
ext))) -> (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 :)

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
  VName
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 :: StructType
ty = Exp -> StructType
typeOf Exp
e
      loc :: SrcLoc
loc = Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e
      pat :: Pat StructType
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
  Binding -> OrderingM ()
addBind (Binding -> OrderingM ()) -> Binding -> OrderingM ()
forall a b. (a -> b) -> a -> b
$ [SizeBinder VName] -> Pat StructType -> Exp -> Binding
PatBind [] Pat StructType
pat Exp
e
  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 (VName -> QualName VName
forall v. v -> QualName v
qualName VName
name) (StructType -> Info StructType
forall a. a -> Info a
Info StructType
ty) SrcLoc
loc

-- An evocative name to use when naming subexpressions of the
-- expression bound to this pattern.
patRepName :: Pat t -> String
patRepName :: forall t. Pat t -> [Char]
patRepName (PatAscription PatBase Info VName t
p TypeExp 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
  Exp
ass' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
ass
  Int
l_prev <- StateT
  (([Binding], [BindModifier]), VNameSource) (Reader [Char]) Int
-> OrderingM Int
forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM (StateT
   (([Binding], [BindModifier]), VNameSource) (Reader [Char]) Int
 -> OrderingM Int)
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) Int
-> OrderingM Int
forall a b. (a -> b) -> a -> b
$ ((([Binding], [BindModifier]), VNameSource) -> Int)
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) Int
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (((([Binding], [BindModifier]), VNameSource) -> Int)
 -> StateT
      (([Binding], [BindModifier]), VNameSource) (Reader [Char]) Int)
-> ((([Binding], [BindModifier]), VNameSource) -> Int)
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) Int
forall a b. (a -> b) -> a -> b
$ [BindModifier] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([BindModifier] -> Int)
-> ((([Binding], [BindModifier]), VNameSource) -> [BindModifier])
-> (([Binding], [BindModifier]), VNameSource)
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([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
  BindModifier -> OrderingM ()
addModifier (BindModifier -> OrderingM ()) -> BindModifier -> OrderingM ()
forall a b. (a -> b) -> a -> b
$ Exp -> Info Text -> SrcLoc -> BindModifier
Ass Exp
ass' Info Text
txt SrcLoc
loc
  Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
final Exp
e
  Int
l_after <- StateT
  (([Binding], [BindModifier]), VNameSource) (Reader [Char]) Int
-> OrderingM Int
forall a.
StateT (([Binding], [BindModifier]), VNameSource) (Reader [Char]) a
-> OrderingM a
OrderingM (StateT
   (([Binding], [BindModifier]), VNameSource) (Reader [Char]) Int
 -> OrderingM Int)
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) Int
-> OrderingM Int
forall a b. (a -> b) -> a -> b
$ ((([Binding], [BindModifier]), VNameSource) -> Int)
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) Int
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (((([Binding], [BindModifier]), VNameSource) -> Int)
 -> StateT
      (([Binding], [BindModifier]), VNameSource) (Reader [Char]) Int)
-> ((([Binding], [BindModifier]), VNameSource) -> Int)
-> StateT
     (([Binding], [BindModifier]), VNameSource) (Reader [Char]) Int
forall a b. (a -> b) -> a -> b
$ [BindModifier] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([BindModifier] -> Int)
-> ((([Binding], [BindModifier]), VNameSource) -> [BindModifier])
-> (([Binding], [BindModifier]), VNameSource)
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([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
  -- if the list of modifier has reduced in size, that means that
  -- all assertions as been inserted,
  -- else, we have to introduce the assertion ourself
  if Int
l_after Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l_prev
    then Exp -> OrderingM Exp
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e'
    else do
      OrderingM ()
rmModifier
      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
$ 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
e' Info Text
txt SrcLoc
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
  Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
final Exp
e
  OrderingM ()
rmModifier
  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
$ AttrInfo VName -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo VName
attr Exp
e' SrcLoc
loc
getOrdering Bool
_ e :: Exp
e@Literal {} = 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
  [Exp]
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
  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
$ [Exp] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit [Exp]
es' SrcLoc
loc
getOrdering Bool
_ (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) = do
  [FieldBase Info VName]
fs' <- (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
  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
$ [FieldBase Info VName] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit [FieldBase Info VName]
fs' SrcLoc
loc
  where
    f :: FieldBase Info VName -> OrderingM (FieldBase Info VName)
f (RecordFieldExplicit L Name
n Exp
e SrcLoc
floc) = do
      Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
      FieldBase Info VName -> OrderingM (FieldBase Info VName)
forall a. a -> OrderingM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (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 L Name
n Exp
e' SrcLoc
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
      [Exp]
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
      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
$ [Exp] -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
ArrayLit [Exp]
es' Info StructType
ty SrcLoc
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
  Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  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
$ Name -> Exp -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
n Exp
e' Info StructType
ty SrcLoc
loc
getOrdering Bool
_ (Negate Exp
e SrcLoc
loc) = do
  Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  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
$ Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Negate Exp
e' SrcLoc
loc
getOrdering Bool
_ (Not Exp
e SrcLoc
loc) = do
  Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  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
$ Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Not Exp
e' SrcLoc
loc
getOrdering Bool
final (Constr Name
n [Exp]
es Info StructType
ty SrcLoc
loc) = do
  [Exp]
es' <- (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
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ Name -> [Exp] -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> [ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
Constr Name
n [Exp]
es' Info StructType
ty SrcLoc
loc
getOrdering Bool
final (Update Exp
eb SliceBase Info VName
slice Exp
eu SrcLoc
loc) = do
  Exp
eu' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
eu
  SliceBase Info VName
slice' <- 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
  Exp
eb' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
eb
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> SliceBase Info VName -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update Exp
eb' SliceBase Info VName
slice' Exp
eu' SrcLoc
loc
  where
    mapper :: ASTMapper OrderingM
mapper = 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
  Exp
eb' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
eb
  Exp
eu' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
eu
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> [Name] -> Exp -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [Name] -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
RecordUpdate Exp
eb' [Name]
ns Exp
eu' Info StructType
ty SrcLoc
loc
getOrdering Bool
final (Lambda [PatBase Info VName ParamType]
params Exp
body Maybe (TypeExp Exp VName)
mte Info ResRetType
ret SrcLoc
loc) = do
  Exp
body' <- Exp -> OrderingM Exp
forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
body
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ [PatBase Info VName ParamType]
-> Exp
-> Maybe (TypeExp Exp VName)
-> Info ResRetType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp (ExpBase f vn) vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [PatBase Info VName ParamType]
params Exp
body' Maybe (TypeExp Exp VName)
mte Info ResRetType
ret SrcLoc
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
  Exp
x <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  VName
yn <- [Char] -> OrderingM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"y"
  let y :: Exp
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' :: TypeBase Exp Uniqueness
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
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
forall a. Monoid a => a
mempty) [(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
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ [PatBase Info VName ParamType]
-> Exp
-> Maybe (TypeExp Exp VName)
-> Info ResRetType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp (ExpBase f vn) vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [VName -> Info ParamType -> SrcLoc -> PatBase Info VName ParamType
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
yn (ParamType -> Info ParamType
forall a. a -> Info a
Info ParamType
yty) SrcLoc
forall a. Monoid a => a
mempty] Exp
body Maybe (TypeExp Exp VName)
forall a. Maybe a
Nothing (ResRetType -> Info ResRetType
forall a. a -> Info a
Info ([VName] -> TypeBase Exp Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase Exp Uniqueness
ret')) SrcLoc
loc
  where
    pSubst :: Exp -> Exp -> VName -> Maybe (Subst t)
pSubst Exp
x Exp
y VName
vn
      | Named VName
p <- PName
xp, VName
p 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
  VName
xn <- [Char] -> OrderingM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"x"
  Exp
y <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  let x :: Exp
x = 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' :: TypeBase Exp Uniqueness
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
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
forall a. Monoid a => a
mempty) [(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') []
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ [PatBase Info VName ParamType]
-> Exp
-> Maybe (TypeExp Exp VName)
-> Info ResRetType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp (ExpBase f vn) vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [VName -> Info ParamType -> SrcLoc -> PatBase Info VName ParamType
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
xn (ParamType -> Info ParamType
forall a. a -> Info a
Info ParamType
xty) SrcLoc
forall a. Monoid a => a
mempty] Exp
body Maybe (TypeExp Exp VName)
forall a. Maybe a
Nothing (ResRetType -> Info ResRetType
forall a. a -> Info a
Info ([VName] -> TypeBase Exp Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase Exp Uniqueness
ret')) SrcLoc
loc
  where
    pSubst :: Exp -> Exp -> VName -> Maybe (Subst t)
pSubst Exp
x Exp
y VName
vn
      | Named VName
p <- PName
xp, VName
p 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
  VName
xn <- [Char] -> OrderingM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"x"
  let (ParamType
xty, RetType [VName]
dims TypeBase Exp Uniqueness
ret) = case StructType
ty of
        Scalar (Arrow NoUniqueness
_ PName
_ Diet
d StructType
xty' ResRetType
ret') -> (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 :: Exp
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
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
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ [PatBase Info VName ParamType]
-> Exp
-> Maybe (TypeExp Exp VName)
-> Info ResRetType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp (ExpBase f vn) vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [VName -> Info ParamType -> SrcLoc -> PatBase Info VName ParamType
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
xn (ParamType -> Info ParamType
forall a. a -> Info a
Info ParamType
xty) SrcLoc
forall a. Monoid a => a
mempty] Exp
body Maybe (TypeExp Exp VName)
forall a. Maybe a
Nothing (ResRetType -> Info ResRetType
forall a. a -> Info a
Info ([VName] -> TypeBase Exp Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase Exp Uniqueness
ret)) SrcLoc
loc
  where
    project :: Exp -> Name -> Exp
project Exp
e Name
field =
      case Exp -> StructType
typeOf Exp
e of
        Scalar (Record Map Name StructType
fs)
          | Just StructType
t <- 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
  SliceBase Info VName
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
  VName
xn <- [Char] -> OrderingM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"x"
  let (ParamType
xty, RetType [VName]
dims TypeBase Exp Uniqueness
ret) = case StructType
ty of
        Scalar (Arrow NoUniqueness
_ PName
_ Diet
d StructType
xty' ResRetType
ret') -> (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 :: Exp
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
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) []))
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ [PatBase Info VName ParamType]
-> Exp
-> Maybe (TypeExp Exp VName)
-> Info ResRetType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp (ExpBase f vn) vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [VName -> Info ParamType -> SrcLoc -> PatBase Info VName ParamType
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
xn (ParamType -> Info ParamType
forall a. a -> Info a
Info ParamType
xty) SrcLoc
forall a. Monoid a => a
mempty] Exp
body Maybe (TypeExp Exp VName)
forall a. Maybe a
Nothing (ResRetType -> Info ResRetType
forall a. a -> Info a
Info ([VName] -> TypeBase Exp Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase Exp Uniqueness
ret)) SrcLoc
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
  NonEmpty (Info (Maybe VName), Exp)
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 ..])))
  Exp
f' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
f
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp
-> NonEmpty (Info (Maybe VName), Exp)
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (f (Maybe VName), ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
f' NonEmpty (Info (Maybe VName), Exp)
args' SrcLoc
loc) Info AppRes
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
  Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> TypeExp Exp VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> TypeExp (ExpBase f vn) vn
-> f StructType
-> SrcLoc
-> ExpBase f vn
Coerce Exp
e' TypeExp Exp VName
te Info StructType
t SrcLoc
loc
getOrdering Bool
final (AppExp (Range Exp
start Maybe Exp
stride Inclusiveness Exp
end SrcLoc
loc) Info AppRes
resT) = do
  Exp
start' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
start
  Maybe Exp
stride' <- (Exp -> OrderingM Exp) -> Maybe Exp -> OrderingM (Maybe 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) -> Maybe a -> m (Maybe b)
mapM (Bool -> Exp -> OrderingM Exp
getOrdering Bool
False) Maybe Exp
stride
  Inclusiveness Exp
end' <- (Exp -> OrderingM Exp)
-> Inclusiveness Exp -> OrderingM (Inclusiveness 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) -> Inclusiveness a -> m (Inclusiveness b)
mapM (Bool -> Exp -> OrderingM Exp
getOrdering Bool
False) Inclusiveness Exp
end
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp
-> Maybe Exp
-> Inclusiveness Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> Maybe (ExpBase f vn)
-> Inclusiveness (ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Range Exp
start' Maybe Exp
stride' Inclusiveness Exp
end' SrcLoc
loc) Info AppRes
resT
getOrdering Bool
final (AppExp (LetPat [SizeBinder VName]
sizes Pat StructType
pat Exp
expr Exp
body SrcLoc
_) Info AppRes
_) = do
  Exp
expr' <- [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
  Binding -> OrderingM ()
addBind (Binding -> OrderingM ()) -> Binding -> OrderingM ()
forall a b. (a -> b) -> a -> b
$ [SizeBinder VName] -> Pat StructType -> Exp -> Binding
PatBind [SizeBinder VName]
sizes Pat StructType
pat Exp
expr'
  Bool -> Exp -> OrderingM Exp
getOrdering Bool
final Exp
body
getOrdering Bool
final (AppExp (LetFun VName
vn ([TypeParamBase VName]
tparams, [PatBase Info VName ParamType]
params, Maybe (TypeExp Exp VName)
mrettype, Info ResRetType
rettype, Exp
body) Exp
e SrcLoc
_) Info AppRes
_) = do
  Exp
body' <- Exp -> OrderingM Exp
forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
body
  Binding -> OrderingM ()
addBind (Binding -> OrderingM ()) -> Binding -> OrderingM ()
forall a b. (a -> b) -> a -> b
$ VName
-> ([TypeParamBase VName], [PatBase Info VName ParamType],
    Maybe (TypeExp Exp VName), Info ResRetType, Exp)
-> Binding
FunBind VName
vn ([TypeParamBase VName]
tparams, [PatBase Info VName ParamType]
params, Maybe (TypeExp Exp VName)
mrettype, Info ResRetType
rettype, Exp
body')
  Bool -> Exp -> OrderingM Exp
getOrdering Bool
final Exp
e
getOrdering Bool
final (AppExp (If Exp
cond Exp
et Exp
ef SrcLoc
loc) Info AppRes
resT) = do
  Exp
cond' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
True Exp
cond
  Exp
et' <- Exp -> OrderingM Exp
forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
et
  Exp
ef' <- Exp -> OrderingM Exp
forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
ef
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp -> Exp -> Exp -> SrcLoc -> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
If Exp
cond' Exp
et' Exp
ef' SrcLoc
loc) Info AppRes
resT
getOrdering Bool
final (AppExp (Loop [VName]
sizes PatBase Info VName ParamType
pat LoopInitBase Info VName
einit LoopFormBase Info VName
form Exp
body SrcLoc
loc) Info AppRes
resT) = do
  Exp
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
  LoopFormBase Info VName
form' <- case LoopFormBase Info VName
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
  Exp
body' <- Exp -> OrderingM Exp
forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
body
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp ([VName]
-> PatBase Info VName ParamType
-> LoopInitBase Info VName
-> LoopFormBase Info VName
-> Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
[VName]
-> PatBase f vn ParamType
-> LoopInitBase f vn
-> LoopFormBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
Loop [VName]
sizes PatBase Info VName ParamType
pat (Exp -> LoopInitBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> LoopInitBase f vn
LoopInitExplicit Exp
einit') LoopFormBase Info VName
form' Exp
body' SrcLoc
loc) Info AppRes
resT
getOrdering Bool
final (AppExp (BinOp (QualName VName
op, SrcLoc
oloc) Info StructType
opT (Exp
el, Info Maybe VName
elp) (Exp
er, Info Maybe VName
erp) SrcLoc
loc) (Info AppRes
resT)) = do
  Exp
expr' <- case (Bool
isOr, Bool
isAnd) of
    (Bool
True, Bool
_) -> do
      Exp
el' <- [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
      Exp
er' <- [Char] -> OrderingM Exp -> OrderingM Exp
forall a. [Char] -> OrderingM a -> OrderingM a
naming [Char]
"or_rhs" (OrderingM Exp -> OrderingM Exp) -> OrderingM Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> OrderingM Exp
forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
er
      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
$ AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp -> Exp -> Exp -> SrcLoc -> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
If Exp
el' (PrimValue -> SrcLoc -> Exp
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (Bool -> PrimValue
BoolValue Bool
True) SrcLoc
forall a. Monoid a => a
mempty) Exp
er' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
resT)
    (Bool
_, Bool
True) -> do
      Exp
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
      Exp
er' <- [Char] -> OrderingM Exp -> OrderingM Exp
forall a. [Char] -> OrderingM a -> OrderingM a
naming [Char]
"and_rhs" (OrderingM Exp -> OrderingM Exp) -> OrderingM Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> OrderingM Exp
forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
er
      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
$ AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp -> Exp -> Exp -> SrcLoc -> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
If Exp
el' Exp
er' (PrimValue -> SrcLoc -> Exp
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (Bool -> PrimValue
BoolValue Bool
False) SrcLoc
forall a. Monoid a => a
mempty) SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
resT)
    (Bool
False, Bool
False) -> do
      Exp
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
      Exp
er' <- [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]
"_rhs") (OrderingM Exp -> OrderingM Exp) -> OrderingM Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
er
      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
$ 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
opT SrcLoc
oloc) [(Maybe VName
elp, Exp
el'), (Maybe VName
erp, Exp
er')] AppRes
resT
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final Exp
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
  Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  SliceBase Info VName
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
  Binding -> OrderingM ()
addBind (Binding -> OrderingM ()) -> Binding -> OrderingM ()
forall a b. (a -> b) -> a -> b
$ [SizeBinder VName] -> Pat StructType -> Exp -> Binding
PatBind [] (VName -> Info StructType -> SrcLoc -> Pat StructType
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
dest Info StructType
dty SrcLoc
dloc) (Exp -> SliceBase Info VName -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update (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
src) Info StructType
sty SrcLoc
sloc) SliceBase Info VName
slice' Exp
e' SrcLoc
loc)
  Bool -> Exp -> OrderingM Exp
getOrdering Bool
final Exp
body
  where
    mapper :: ASTMapper OrderingM
mapper = 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
  Exp
e' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
e
  SliceBase Info VName
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
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ 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
e' SliceBase Info VName
slice' SrcLoc
loc) Info AppRes
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
  Exp
expr' <- Bool -> Exp -> OrderingM Exp
getOrdering Bool
False Exp
expr
  NonEmpty (CaseBase Info VName)
cs' <- (CaseBase Info VName -> OrderingM (CaseBase Info VName))
-> NonEmpty (CaseBase Info VName)
-> OrderingM (NonEmpty (CaseBase 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) -> NonEmpty a -> m (NonEmpty b)
mapM CaseBase Info VName -> OrderingM (CaseBase Info VName)
forall {m :: * -> *}.
MonadFreshNames m =>
CaseBase Info VName -> m (CaseBase Info VName)
f NonEmpty (CaseBase Info VName)
cs
  Bool -> Exp -> OrderingM Exp
nameExp Bool
final (Exp -> OrderingM Exp) -> Exp -> OrderingM Exp
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp
-> NonEmpty (CaseBase Info VName)
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (CaseBase f vn) -> SrcLoc -> AppExpBase f vn
Match Exp
expr' NonEmpty (CaseBase Info VName)
cs' SrcLoc
loc) Info AppRes
resT
  where
    f :: CaseBase Info VName -> m (CaseBase Info VName)
f (CasePat Pat StructType
pat Exp
body SrcLoc
cloc) = do
      Exp
body' <- Exp -> m Exp
forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
body
      CaseBase Info VName -> m (CaseBase Info VName)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Pat StructType -> Exp -> SrcLoc -> CaseBase Info VName
forall (f :: * -> *) vn.
PatBase f vn StructType -> ExpBase f vn -> SrcLoc -> CaseBase f vn
CasePat Pat StructType
pat Exp
body' SrcLoc
cloc)

-- Transform a body, e.g. the expression of a valbind,
-- branches of an if/match...
-- Note that this is not producing an OrderingM, produce
-- a complete separtion of states.
transformBody :: (MonadFreshNames m) => Exp -> m Exp
transformBody :: forall (m :: * -> *). MonadFreshNames m => Exp -> m Exp
transformBody Exp
e = do
  (Exp
e', [Binding]
pre_eval) <- 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)
  Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> m Exp) -> Exp -> m Exp
forall a b. (a -> b) -> a -> b
$ (Exp -> Binding -> Exp) -> Exp -> [Binding] -> 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 -> Binding -> Exp
f Exp
e' [Binding]
pre_eval
  where
    appRes :: Info AppRes
appRes = case Exp
e of
      (AppExp AppExpBase Info VName
_ Info AppRes
r) -> Info AppRes
r
      Exp
_ -> 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
  Exp
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
  ValBind -> m ValBind
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ValBind -> m ValBind) -> ValBind -> m ValBind
forall a b. (a -> b) -> a -> b
$ ValBind
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