-- | Converts identifiers of record type into record patterns (and
-- similarly for tuples).  This is to ensure that the closures
-- produced in lambda lifting and defunctionalisation do not carry
-- around huge records of which only a tiny part is needed.
module Futhark.Internalise.ReplaceRecords (transformProg) where

import Control.Monad
import Control.Monad.Reader
import Control.Monad.State
import Data.Bifunctor
import Data.Bitraversable
import Data.Map.Strict qualified as M
import Futhark.MonadFreshNames
import Language.Futhark
import Language.Futhark.Traversals

-- Mapping from record names to the variable names that contain the
-- fields, as well as an expression for the entire record. This is
-- used because the monomorphiser also expands all, record patterns.
type RecordReplacements = M.Map VName RecordReplacement

type RecordReplacement = (M.Map Name (VName, StructType), Exp)

newtype Env = Env
  { Env -> RecordReplacements
envRecordReplacements :: RecordReplacements
  }

data S = S
  { S -> VNameSource
stateNameSource :: VNameSource,
    S -> Map StructType StructType
stateStructTypeMemo :: M.Map StructType StructType,
    S -> Map ParamType ParamType
stateParamTypeMemo :: M.Map ParamType ParamType
  }

-- The monomorphization monad.
newtype RecordM a
  = RecordM (ReaderT Env (State S) a)
  deriving
    ( (forall a b. (a -> b) -> RecordM a -> RecordM b)
-> (forall a b. a -> RecordM b -> RecordM a) -> Functor RecordM
forall a b. a -> RecordM b -> RecordM a
forall a b. (a -> b) -> RecordM a -> RecordM 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) -> RecordM a -> RecordM b
fmap :: forall a b. (a -> b) -> RecordM a -> RecordM b
$c<$ :: forall a b. a -> RecordM b -> RecordM a
<$ :: forall a b. a -> RecordM b -> RecordM a
Functor,
      Functor RecordM
Functor RecordM =>
(forall a. a -> RecordM a)
-> (forall a b. RecordM (a -> b) -> RecordM a -> RecordM b)
-> (forall a b c.
    (a -> b -> c) -> RecordM a -> RecordM b -> RecordM c)
-> (forall a b. RecordM a -> RecordM b -> RecordM b)
-> (forall a b. RecordM a -> RecordM b -> RecordM a)
-> Applicative RecordM
forall a. a -> RecordM a
forall a b. RecordM a -> RecordM b -> RecordM a
forall a b. RecordM a -> RecordM b -> RecordM b
forall a b. RecordM (a -> b) -> RecordM a -> RecordM b
forall a b c. (a -> b -> c) -> RecordM a -> RecordM b -> RecordM 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 -> RecordM a
pure :: forall a. a -> RecordM a
$c<*> :: forall a b. RecordM (a -> b) -> RecordM a -> RecordM b
<*> :: forall a b. RecordM (a -> b) -> RecordM a -> RecordM b
$cliftA2 :: forall a b c. (a -> b -> c) -> RecordM a -> RecordM b -> RecordM c
liftA2 :: forall a b c. (a -> b -> c) -> RecordM a -> RecordM b -> RecordM c
$c*> :: forall a b. RecordM a -> RecordM b -> RecordM b
*> :: forall a b. RecordM a -> RecordM b -> RecordM b
$c<* :: forall a b. RecordM a -> RecordM b -> RecordM a
<* :: forall a b. RecordM a -> RecordM b -> RecordM a
Applicative,
      Applicative RecordM
Applicative RecordM =>
(forall a b. RecordM a -> (a -> RecordM b) -> RecordM b)
-> (forall a b. RecordM a -> RecordM b -> RecordM b)
-> (forall a. a -> RecordM a)
-> Monad RecordM
forall a. a -> RecordM a
forall a b. RecordM a -> RecordM b -> RecordM b
forall a b. RecordM a -> (a -> RecordM b) -> RecordM 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. RecordM a -> (a -> RecordM b) -> RecordM b
>>= :: forall a b. RecordM a -> (a -> RecordM b) -> RecordM b
$c>> :: forall a b. RecordM a -> RecordM b -> RecordM b
>> :: forall a b. RecordM a -> RecordM b -> RecordM b
$creturn :: forall a. a -> RecordM a
return :: forall a. a -> RecordM a
Monad,
      MonadReader Env,
      MonadState S
    )

instance MonadFreshNames RecordM where
  getNameSource :: RecordM VNameSource
getNameSource = ReaderT Env (State S) VNameSource -> RecordM VNameSource
forall a. ReaderT Env (State S) a -> RecordM a
RecordM (ReaderT Env (State S) VNameSource -> RecordM VNameSource)
-> ReaderT Env (State S) VNameSource -> RecordM VNameSource
forall a b. (a -> b) -> a -> b
$ (S -> VNameSource) -> ReaderT Env (State S) VNameSource
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets S -> VNameSource
stateNameSource
  putNameSource :: VNameSource -> RecordM ()
putNameSource VNameSource
src = ReaderT Env (State S) () -> RecordM ()
forall a. ReaderT Env (State S) a -> RecordM a
RecordM (ReaderT Env (State S) () -> RecordM ())
-> ReaderT Env (State S) () -> RecordM ()
forall a b. (a -> b) -> a -> b
$ (S -> S) -> ReaderT Env (State S) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((S -> S) -> ReaderT Env (State S) ())
-> (S -> S) -> ReaderT Env (State S) ()
forall a b. (a -> b) -> a -> b
$ \S
s -> S
s {stateNameSource = src}

runRecordM :: VNameSource -> RecordM a -> (a, VNameSource)
runRecordM :: forall a. VNameSource -> RecordM a -> (a, VNameSource)
runRecordM VNameSource
src (RecordM ReaderT Env (State S) a
m) =
  (S -> VNameSource) -> (a, S) -> (a, 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 S -> VNameSource
stateNameSource ((a, S) -> (a, VNameSource)) -> (a, S) -> (a, VNameSource)
forall a b. (a -> b) -> a -> b
$
    State S a -> S -> (a, S)
forall s a. State s a -> s -> (a, s)
runState (ReaderT Env (State S) a -> Env -> State S a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT Env (State S) a
m (RecordReplacements -> Env
Env RecordReplacements
forall a. Monoid a => a
mempty)) (VNameSource
-> Map StructType StructType -> Map ParamType ParamType -> S
S VNameSource
src Map StructType StructType
forall a. Monoid a => a
mempty Map ParamType ParamType
forall a. Monoid a => a
mempty)

withRecordReplacements :: RecordReplacements -> RecordM a -> RecordM a
withRecordReplacements :: forall a. RecordReplacements -> RecordM a -> RecordM a
withRecordReplacements RecordReplacements
rr = (Env -> Env) -> RecordM a -> RecordM a
forall a. (Env -> Env) -> RecordM a -> RecordM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env -> Env) -> RecordM a -> RecordM a)
-> (Env -> Env) -> RecordM a -> RecordM a
forall a b. (a -> b) -> a -> b
$ \Env
env ->
  Env
env {envRecordReplacements = rr <> envRecordReplacements env}

lookupRecordReplacement :: VName -> RecordM (Maybe RecordReplacement)
lookupRecordReplacement :: VName -> RecordM (Maybe RecordReplacement)
lookupRecordReplacement VName
v = (Env -> Maybe RecordReplacement)
-> RecordM (Maybe RecordReplacement)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Maybe RecordReplacement)
 -> RecordM (Maybe RecordReplacement))
-> (Env -> Maybe RecordReplacement)
-> RecordM (Maybe RecordReplacement)
forall a b. (a -> b) -> a -> b
$ VName -> RecordReplacements -> Maybe RecordReplacement
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (RecordReplacements -> Maybe RecordReplacement)
-> (Env -> RecordReplacements) -> Env -> Maybe RecordReplacement
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> RecordReplacements
envRecordReplacements

wildcard :: TypeBase Size u -> SrcLoc -> Pat (TypeBase Size u)
wildcard :: forall u. TypeBase Exp u -> SrcLoc -> Pat (TypeBase Exp u)
wildcard (Scalar (Record Map Name (TypeBase Exp u)
fs)) SrcLoc
loc =
  [(L Name, PatBase Info VName (TypeBase Exp u))]
-> SrcLoc -> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t.
[(L Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat ([L Name]
-> [PatBase Info VName (TypeBase Exp u)]
-> [(L Name, PatBase Info VName (TypeBase Exp u))]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Name -> L Name) -> [Name] -> [L Name]
forall a b. (a -> b) -> [a] -> [b]
map (Loc -> Name -> L Name
forall a. Loc -> a -> L a
L Loc
forall a. IsLocation a => a
noLoc) (Map Name (TypeBase Exp u) -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name (TypeBase Exp u)
fs)) ([PatBase Info VName (TypeBase Exp u)]
 -> [(L Name, PatBase Info VName (TypeBase Exp u))])
-> [PatBase Info VName (TypeBase Exp u)]
-> [(L Name, PatBase Info VName (TypeBase Exp u))]
forall a b. (a -> b) -> a -> b
$ (TypeBase Exp u -> PatBase Info VName (TypeBase Exp u))
-> [TypeBase Exp u] -> [PatBase Info VName (TypeBase Exp u)]
forall a b. (a -> b) -> [a] -> [b]
map ((Info (TypeBase Exp u)
-> SrcLoc -> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
`Wildcard` SrcLoc
loc) (Info (TypeBase Exp u) -> PatBase Info VName (TypeBase Exp u))
-> (TypeBase Exp u -> Info (TypeBase Exp u))
-> TypeBase Exp u
-> PatBase Info VName (TypeBase Exp u)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Exp u -> Info (TypeBase Exp u)
forall a. a -> Info a
Info) ([TypeBase Exp u] -> [PatBase Info VName (TypeBase Exp u)])
-> [TypeBase Exp u] -> [PatBase Info VName (TypeBase Exp u)]
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase Exp u) -> [TypeBase Exp u]
forall k a. Map k a -> [a]
M.elems Map Name (TypeBase Exp u)
fs) SrcLoc
loc
wildcard TypeBase Exp u
t SrcLoc
loc =
  Info (TypeBase Exp u)
-> SrcLoc -> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
Wildcard (TypeBase Exp u -> Info (TypeBase Exp u)
forall a. a -> Info a
Info TypeBase Exp u
t) SrcLoc
loc

memoParamType :: ParamType -> RecordM ParamType -> RecordM ParamType
memoParamType :: ParamType -> RecordM ParamType -> RecordM ParamType
memoParamType ParamType
t RecordM ParamType
m = do
  Maybe ParamType
prev <- (S -> Maybe ParamType) -> RecordM (Maybe ParamType)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((S -> Maybe ParamType) -> RecordM (Maybe ParamType))
-> (S -> Maybe ParamType) -> RecordM (Maybe ParamType)
forall a b. (a -> b) -> a -> b
$ ParamType -> Map ParamType ParamType -> Maybe ParamType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup ParamType
t (Map ParamType ParamType -> Maybe ParamType)
-> (S -> Map ParamType ParamType) -> S -> Maybe ParamType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. S -> Map ParamType ParamType
stateParamTypeMemo
  case Maybe ParamType
prev of
    Just ParamType
t' -> ParamType -> RecordM ParamType
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ParamType
t'
    Maybe ParamType
Nothing -> do
      ParamType
t' <- RecordM ParamType
m
      (S -> S) -> RecordM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((S -> S) -> RecordM ()) -> (S -> S) -> RecordM ()
forall a b. (a -> b) -> a -> b
$ \S
s -> S
s {stateParamTypeMemo = M.insert t t' $ stateParamTypeMemo s}
      ParamType -> RecordM ParamType
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ParamType
t'

memoStructType :: StructType -> RecordM StructType -> RecordM StructType
memoStructType :: StructType -> RecordM StructType -> RecordM StructType
memoStructType StructType
t RecordM StructType
m = do
  Maybe StructType
prev <- (S -> Maybe StructType) -> RecordM (Maybe StructType)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((S -> Maybe StructType) -> RecordM (Maybe StructType))
-> (S -> Maybe StructType) -> RecordM (Maybe StructType)
forall a b. (a -> b) -> a -> b
$ StructType -> Map StructType StructType -> Maybe StructType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup StructType
t (Map StructType StructType -> Maybe StructType)
-> (S -> Map StructType StructType) -> S -> Maybe StructType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. S -> Map StructType StructType
stateStructTypeMemo
  case Maybe StructType
prev of
    Just StructType
t' -> StructType -> RecordM StructType
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure StructType
t'
    Maybe StructType
Nothing -> do
      StructType
t' <- RecordM StructType
m
      (S -> S) -> RecordM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((S -> S) -> RecordM ()) -> (S -> S) -> RecordM ()
forall a b. (a -> b) -> a -> b
$ \S
s -> S
s {stateStructTypeMemo = M.insert t t' $ stateStructTypeMemo s}
      StructType -> RecordM StructType
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure StructType
t'

-- No need to keep memoisation cache between top level functions.
memoClear :: RecordM ()
memoClear :: RecordM ()
memoClear = (S -> S) -> RecordM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((S -> S) -> RecordM ()) -> (S -> S) -> RecordM ()
forall a b. (a -> b) -> a -> b
$ \S
s ->
  S
s
    { stateStructTypeMemo = mempty,
      stateParamTypeMemo = mempty
    }

transformPat ::
  (TypeBase Size u -> RecordM (TypeBase Size u)) ->
  Pat (TypeBase Size u) ->
  RecordM (Pat (TypeBase Size u), RecordReplacements)
transformPat :: forall u.
(TypeBase Exp u -> RecordM (TypeBase Exp u))
-> Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
_ (Id VName
v (Info (Scalar (Record Map Name (TypeBase Exp u)
fs))) SrcLoc
loc) = do
  let fs' :: [(Name, TypeBase Exp u)]
fs' = Map Name (TypeBase Exp u) -> [(Name, TypeBase Exp u)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name (TypeBase Exp u)
fs
  ([VName]
fs_ks, [TypeBase Exp u]
fs_ts) <- ([(VName, TypeBase Exp u)] -> ([VName], [TypeBase Exp u]))
-> RecordM [(VName, TypeBase Exp u)]
-> RecordM ([VName], [TypeBase Exp u])
forall a b. (a -> b) -> RecordM a -> RecordM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(VName, TypeBase Exp u)] -> ([VName], [TypeBase Exp u])
forall a b. [(a, b)] -> ([a], [b])
unzip (RecordM [(VName, TypeBase Exp u)]
 -> RecordM ([VName], [TypeBase Exp u]))
-> RecordM [(VName, TypeBase Exp u)]
-> RecordM ([VName], [TypeBase Exp u])
forall a b. (a -> b) -> a -> b
$
    [(Name, TypeBase Exp u)]
-> ((Name, TypeBase Exp u) -> RecordM (VName, TypeBase Exp u))
-> RecordM [(VName, TypeBase Exp u)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Name, TypeBase Exp u)]
fs' (((Name, TypeBase Exp u) -> RecordM (VName, TypeBase Exp u))
 -> RecordM [(VName, TypeBase Exp u)])
-> ((Name, TypeBase Exp u) -> RecordM (VName, TypeBase Exp u))
-> RecordM [(VName, TypeBase Exp u)]
forall a b. (a -> b) -> a -> b
$ \(Name
f, TypeBase Exp u
ft) ->
      (,) (VName -> TypeBase Exp u -> (VName, TypeBase Exp u))
-> RecordM VName
-> RecordM (TypeBase Exp u -> (VName, TypeBase Exp u))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> RecordM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (Name -> String
nameToString Name
f) RecordM (TypeBase Exp u -> (VName, TypeBase Exp u))
-> RecordM (TypeBase Exp u) -> RecordM (VName, TypeBase Exp u)
forall a b. RecordM (a -> b) -> RecordM a -> RecordM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeBase Exp u -> RecordM (TypeBase Exp u)
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase Exp u
ft
  (PatBase Info VName (TypeBase Exp u), RecordReplacements)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( [(L Name, PatBase Info VName (TypeBase Exp u))]
-> SrcLoc -> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t.
[(L Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat
        ([L Name]
-> [PatBase Info VName (TypeBase Exp u)]
-> [(L Name, PatBase Info VName (TypeBase Exp u))]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Name, TypeBase Exp u) -> L Name)
-> [(Name, TypeBase Exp u)] -> [L Name]
forall a b. (a -> b) -> [a] -> [b]
map (Loc -> Name -> L Name
forall a. Loc -> a -> L a
L Loc
forall a. IsLocation a => a
noLoc (Name -> L Name)
-> ((Name, TypeBase Exp u) -> Name)
-> (Name, TypeBase Exp u)
-> L Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, TypeBase Exp u) -> Name
forall a b. (a, b) -> a
fst) [(Name, TypeBase Exp u)]
fs') ((VName
 -> Info (TypeBase Exp u)
 -> SrcLoc
 -> PatBase Info VName (TypeBase Exp u))
-> [VName]
-> [Info (TypeBase Exp u)]
-> [SrcLoc]
-> [PatBase Info VName (TypeBase Exp u)]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 VName
-> Info (TypeBase Exp u)
-> SrcLoc
-> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id [VName]
fs_ks ((TypeBase Exp u -> Info (TypeBase Exp u))
-> [TypeBase Exp u] -> [Info (TypeBase Exp u)]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Exp u -> Info (TypeBase Exp u)
forall a. a -> Info a
Info [TypeBase Exp u]
fs_ts) ([SrcLoc] -> [PatBase Info VName (TypeBase Exp u)])
-> [SrcLoc] -> [PatBase Info VName (TypeBase Exp u)]
forall a b. (a -> b) -> a -> b
$ SrcLoc -> [SrcLoc]
forall a. a -> [a]
repeat SrcLoc
loc))
        SrcLoc
loc,
      VName -> RecordReplacement -> RecordReplacements
forall k a. k -> a -> Map k a
M.singleton
        VName
v
        ( [(Name, (VName, StructType))] -> Map Name (VName, StructType)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, (VName, StructType))] -> Map Name (VName, StructType))
-> [(Name, (VName, StructType))] -> Map Name (VName, StructType)
forall a b. (a -> b) -> a -> b
$ [Name] -> [(VName, StructType)] -> [(Name, (VName, StructType))]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Name, TypeBase Exp u) -> Name)
-> [(Name, TypeBase Exp u)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, TypeBase Exp u) -> Name
forall a b. (a, b) -> a
fst [(Name, TypeBase Exp u)]
fs') ([(VName, StructType)] -> [(Name, (VName, StructType))])
-> [(VName, StructType)] -> [(Name, (VName, StructType))]
forall a b. (a -> b) -> a -> b
$ [VName] -> [StructType] -> [(VName, StructType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
fs_ks ([StructType] -> [(VName, StructType)])
-> [StructType] -> [(VName, StructType)]
forall a b. (a -> b) -> a -> b
$ (TypeBase Exp u -> StructType) -> [TypeBase Exp u] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Exp u -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct [TypeBase Exp u]
fs_ts,
          [FieldBase Info VName] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ((Name -> VName -> TypeBase Exp u -> FieldBase Info VName)
-> [Name] -> [VName] -> [TypeBase Exp u] -> [FieldBase Info VName]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 Name -> VName -> TypeBase Exp u -> FieldBase Info VName
forall {vn} {u}. Name -> vn -> TypeBase Exp u -> FieldBase Info vn
toField (((Name, TypeBase Exp u) -> Name)
-> [(Name, TypeBase Exp u)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, TypeBase Exp u) -> Name
forall a b. (a, b) -> a
fst [(Name, TypeBase Exp u)]
fs') [VName]
fs_ks [TypeBase Exp u]
fs_ts) SrcLoc
loc
        )
    )
  where
    toField :: Name -> vn -> TypeBase Exp u -> FieldBase Info vn
toField Name
f vn
f_v TypeBase Exp u
f_t =
      let f_v' :: ExpBase Info vn
f_v' = QualName vn -> Info StructType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (vn -> QualName vn
forall v. v -> QualName v
qualName vn
f_v) (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ TypeBase Exp u -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp u
f_t) SrcLoc
loc
       in L Name -> ExpBase Info vn -> SrcLoc -> FieldBase Info vn
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
forall a. IsLocation a => a
noLoc Name
f) ExpBase Info vn
f_v' SrcLoc
loc
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType (Id VName
v Info (TypeBase Exp u)
t SrcLoc
loc) = do
  Info (TypeBase Exp u)
t' <- (TypeBase Exp u -> RecordM (TypeBase Exp u))
-> Info (TypeBase Exp u) -> RecordM (Info (TypeBase Exp u))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Info a -> f (Info b)
traverse TypeBase Exp u -> RecordM (TypeBase Exp u)
onType Info (TypeBase Exp u)
t
  (PatBase Info VName (TypeBase Exp u), RecordReplacements)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
-> Info (TypeBase Exp u)
-> SrcLoc
-> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
v Info (TypeBase Exp u)
t' SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType (TuplePat [PatBase Info VName (TypeBase Exp u)]
pats SrcLoc
loc) = do
  ([PatBase Info VName (TypeBase Exp u)]
pats', [RecordReplacements]
rrs) <- (PatBase Info VName (TypeBase Exp u)
 -> RecordM
      (PatBase Info VName (TypeBase Exp u), RecordReplacements))
-> [PatBase Info VName (TypeBase Exp u)]
-> RecordM
     ([PatBase Info VName (TypeBase Exp u)], [RecordReplacements])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM ((TypeBase Exp u -> RecordM (TypeBase Exp u))
-> PatBase Info VName (TypeBase Exp u)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall u.
(TypeBase Exp u -> RecordM (TypeBase Exp u))
-> Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType) [PatBase Info VName (TypeBase Exp u)]
pats
  (PatBase Info VName (TypeBase Exp u), RecordReplacements)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([PatBase Info VName (TypeBase Exp u)]
-> SrcLoc -> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t.
[PatBase f vn t] -> SrcLoc -> PatBase f vn t
TuplePat [PatBase Info VName (TypeBase Exp u)]
pats' SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType (RecordPat [(L Name, PatBase Info VName (TypeBase Exp u))]
fields SrcLoc
loc) = do
  let ([L Name]
field_names, [PatBase Info VName (TypeBase Exp u)]
field_pats) = [(L Name, PatBase Info VName (TypeBase Exp u))]
-> ([L Name], [PatBase Info VName (TypeBase Exp u)])
forall a b. [(a, b)] -> ([a], [b])
unzip [(L Name, PatBase Info VName (TypeBase Exp u))]
fields
  ([PatBase Info VName (TypeBase Exp u)]
field_pats', [RecordReplacements]
rrs) <- (PatBase Info VName (TypeBase Exp u)
 -> RecordM
      (PatBase Info VName (TypeBase Exp u), RecordReplacements))
-> [PatBase Info VName (TypeBase Exp u)]
-> RecordM
     ([PatBase Info VName (TypeBase Exp u)], [RecordReplacements])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM ((TypeBase Exp u -> RecordM (TypeBase Exp u))
-> PatBase Info VName (TypeBase Exp u)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall u.
(TypeBase Exp u -> RecordM (TypeBase Exp u))
-> Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType) [PatBase Info VName (TypeBase Exp u)]
field_pats
  (PatBase Info VName (TypeBase Exp u), RecordReplacements)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([(L Name, PatBase Info VName (TypeBase Exp u))]
-> SrcLoc -> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t.
[(L Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat ([L Name]
-> [PatBase Info VName (TypeBase Exp u)]
-> [(L Name, PatBase Info VName (TypeBase Exp u))]
forall a b. [a] -> [b] -> [(a, b)]
zip [L Name]
field_names [PatBase Info VName (TypeBase Exp u)]
field_pats') SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType (PatParens PatBase Info VName (TypeBase Exp u)
pat SrcLoc
loc) = do
  (PatBase Info VName (TypeBase Exp u)
pat', RecordReplacements
rr) <- (TypeBase Exp u -> RecordM (TypeBase Exp u))
-> PatBase Info VName (TypeBase Exp u)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall u.
(TypeBase Exp u -> RecordM (TypeBase Exp u))
-> Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType PatBase Info VName (TypeBase Exp u)
pat
  (PatBase Info VName (TypeBase Exp u), RecordReplacements)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatBase Info VName (TypeBase Exp u)
-> SrcLoc -> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t.
PatBase f vn t -> SrcLoc -> PatBase f vn t
PatParens PatBase Info VName (TypeBase Exp u)
pat' SrcLoc
loc, RecordReplacements
rr)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType (PatAttr AttrInfo VName
attr PatBase Info VName (TypeBase Exp u)
pat SrcLoc
loc) = do
  (PatBase Info VName (TypeBase Exp u)
pat', RecordReplacements
rr) <- (TypeBase Exp u -> RecordM (TypeBase Exp u))
-> PatBase Info VName (TypeBase Exp u)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall u.
(TypeBase Exp u -> RecordM (TypeBase Exp u))
-> Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType PatBase Info VName (TypeBase Exp u)
pat
  (PatBase Info VName (TypeBase Exp u), RecordReplacements)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AttrInfo VName
-> PatBase Info VName (TypeBase Exp u)
-> SrcLoc
-> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t.
AttrInfo vn -> PatBase f vn t -> SrcLoc -> PatBase f vn t
PatAttr AttrInfo VName
attr PatBase Info VName (TypeBase Exp u)
pat' SrcLoc
loc, RecordReplacements
rr)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType (Wildcard (Info TypeBase Exp u
t) SrcLoc
loc) = do
  TypeBase Exp u
t' <- TypeBase Exp u -> RecordM (TypeBase Exp u)
onType TypeBase Exp u
t
  (PatBase Info VName (TypeBase Exp u), RecordReplacements)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase Exp u -> SrcLoc -> PatBase Info VName (TypeBase Exp u)
forall u. TypeBase Exp u -> SrcLoc -> Pat (TypeBase Exp u)
wildcard TypeBase Exp u
t' SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType (PatAscription PatBase Info VName (TypeBase Exp u)
pat TypeExp Exp VName
_ SrcLoc
_) =
  (TypeBase Exp u -> RecordM (TypeBase Exp u))
-> PatBase Info VName (TypeBase Exp u)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall u.
(TypeBase Exp u -> RecordM (TypeBase Exp u))
-> Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType PatBase Info VName (TypeBase Exp u)
pat
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
_ (PatLit PatLit
e Info (TypeBase Exp u)
t SrcLoc
loc) =
  (PatBase Info VName (TypeBase Exp u), RecordReplacements)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatLit
-> Info (TypeBase Exp u)
-> SrcLoc
-> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t.
PatLit -> f t -> SrcLoc -> PatBase f vn t
PatLit PatLit
e Info (TypeBase Exp u)
t SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType (PatConstr Name
name Info (TypeBase Exp u)
t [PatBase Info VName (TypeBase Exp u)]
all_ps SrcLoc
loc) = do
  ([PatBase Info VName (TypeBase Exp u)]
all_ps', [RecordReplacements]
rrs) <- (PatBase Info VName (TypeBase Exp u)
 -> RecordM
      (PatBase Info VName (TypeBase Exp u), RecordReplacements))
-> [PatBase Info VName (TypeBase Exp u)]
-> RecordM
     ([PatBase Info VName (TypeBase Exp u)], [RecordReplacements])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM ((TypeBase Exp u -> RecordM (TypeBase Exp u))
-> PatBase Info VName (TypeBase Exp u)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall u.
(TypeBase Exp u -> RecordM (TypeBase Exp u))
-> Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType) [PatBase Info VName (TypeBase Exp u)]
all_ps
  (PatBase Info VName (TypeBase Exp u), RecordReplacements)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name
-> Info (TypeBase Exp u)
-> [PatBase Info VName (TypeBase Exp u)]
-> SrcLoc
-> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t.
Name -> f t -> [PatBase f vn t] -> SrcLoc -> PatBase f vn t
PatConstr Name
name Info (TypeBase Exp u)
t [PatBase Info VName (TypeBase Exp u)]
all_ps' SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)

transformParamType :: ParamType -> RecordM ParamType
transformParamType :: ParamType -> RecordM ParamType
transformParamType ParamType
t = ParamType -> RecordM ParamType -> RecordM ParamType
memoParamType ParamType
t (RecordM ParamType -> RecordM ParamType)
-> RecordM ParamType -> RecordM ParamType
forall a b. (a -> b) -> a -> b
$ (Exp -> RecordM Exp)
-> (Diet -> RecordM Diet) -> ParamType -> RecordM ParamType
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> TypeBase a b -> f (TypeBase c d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse Exp -> RecordM Exp
transformExp Diet -> RecordM Diet
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ParamType
t

transformStructType :: StructType -> RecordM StructType
transformStructType :: StructType -> RecordM StructType
transformStructType StructType
t = StructType -> RecordM StructType -> RecordM StructType
memoStructType StructType
t (RecordM StructType -> RecordM StructType)
-> RecordM StructType -> RecordM StructType
forall a b. (a -> b) -> a -> b
$ (Exp -> RecordM Exp)
-> (NoUniqueness -> RecordM NoUniqueness)
-> StructType
-> RecordM StructType
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> TypeBase a b -> f (TypeBase c d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse Exp -> RecordM Exp
transformExp NoUniqueness -> RecordM NoUniqueness
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure StructType
t

transformExp :: Exp -> RecordM Exp
transformExp :: Exp -> RecordM Exp
transformExp (Project Name
n Exp
e Info StructType
t SrcLoc
loc) = do
  Maybe RecordReplacement
maybe_fs <- case Exp
e of
    Var QualName VName
qn Info StructType
_ SrcLoc
_ -> VName -> RecordM (Maybe RecordReplacement)
lookupRecordReplacement (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
    Exp
_ -> Maybe RecordReplacement -> RecordM (Maybe RecordReplacement)
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe RecordReplacement
forall a. Maybe a
Nothing
  case Maybe RecordReplacement
maybe_fs of
    Just (Map Name (VName, StructType)
m, Exp
_)
      | Just (VName
v, StructType
_) <- Name -> Map Name (VName, StructType) -> Maybe (VName, StructType)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
n Map Name (VName, StructType)
m ->
          Exp -> RecordM Exp
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> RecordM Exp) -> Exp -> RecordM 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
v) Info StructType
t SrcLoc
loc
    Maybe RecordReplacement
_ -> do
      Exp
e' <- Exp -> RecordM Exp
transformExp Exp
e
      Exp -> RecordM Exp
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> RecordM Exp) -> Exp -> RecordM 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
t SrcLoc
loc
transformExp (Var QualName VName
fname Info StructType
t SrcLoc
loc) = do
  Maybe RecordReplacement
maybe_fs <- VName -> RecordM (Maybe RecordReplacement)
lookupRecordReplacement (VName -> RecordM (Maybe RecordReplacement))
-> VName -> RecordM (Maybe RecordReplacement)
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname
  case Maybe RecordReplacement
maybe_fs of
    Just (Map Name (VName, StructType)
_, Exp
e) -> Exp -> RecordM Exp
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
    Maybe RecordReplacement
Nothing ->
      QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
fname (Info StructType -> SrcLoc -> Exp)
-> RecordM (Info StructType) -> RecordM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (StructType -> RecordM StructType)
-> Info StructType -> RecordM (Info StructType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Info a -> f (Info b)
traverse StructType -> RecordM StructType
transformStructType Info StructType
t RecordM (SrcLoc -> Exp) -> RecordM SrcLoc -> RecordM Exp
forall a b. RecordM (a -> b) -> RecordM a -> RecordM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> RecordM SrcLoc
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (AppExp (LetPat [SizeBinder VName]
sizes PatBase Info VName StructType
pat Exp
e Exp
body SrcLoc
loc) Info AppRes
res) = do
  Exp
e' <- Exp -> RecordM Exp
transformExp Exp
e
  (PatBase Info VName StructType
pat', RecordReplacements
rr) <- (StructType -> RecordM StructType)
-> PatBase Info VName StructType
-> RecordM (PatBase Info VName StructType, RecordReplacements)
forall u.
(TypeBase Exp u -> RecordM (TypeBase Exp u))
-> Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat StructType -> RecordM StructType
transformStructType PatBase Info VName StructType
pat
  Exp
body' <- RecordReplacements -> RecordM Exp -> RecordM Exp
forall a. RecordReplacements -> RecordM a -> RecordM a
withRecordReplacements RecordReplacements
rr (RecordM Exp -> RecordM Exp) -> RecordM Exp -> RecordM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> RecordM Exp
transformExp Exp
body
  Exp -> RecordM Exp
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> RecordM Exp) -> Exp -> RecordM 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 ([SizeBinder VName]
-> PatBase Info VName 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 PatBase Info VName StructType
pat' Exp
e' Exp
body' SrcLoc
loc) Info AppRes
res
transformExp (AppExp (LetFun VName
fname ([TypeParamBase VName]
tparams, [PatBase Info VName ParamType]
params, Maybe (TypeExp Exp VName)
retdecl, Info ResRetType
ret, Exp
funbody) Exp
body SrcLoc
loc) Info AppRes
res) = do
  ([PatBase Info VName ParamType]
params', [RecordReplacements]
rrs) <- (PatBase Info VName ParamType
 -> RecordM (PatBase Info VName ParamType, RecordReplacements))
-> [PatBase Info VName ParamType]
-> RecordM ([PatBase Info VName ParamType], [RecordReplacements])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM ((ParamType -> RecordM ParamType)
-> PatBase Info VName ParamType
-> RecordM (PatBase Info VName ParamType, RecordReplacements)
forall u.
(TypeBase Exp u -> RecordM (TypeBase Exp u))
-> Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat ParamType -> RecordM ParamType
transformParamType) [PatBase Info VName ParamType]
params
  Exp
funbody' <- RecordReplacements -> RecordM Exp -> RecordM Exp
forall a. RecordReplacements -> RecordM a -> RecordM a
withRecordReplacements ([RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs) (RecordM Exp -> RecordM Exp) -> RecordM Exp -> RecordM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> RecordM Exp
transformExp Exp
funbody
  Exp
body' <- Exp -> RecordM Exp
transformExp Exp
body
  Exp -> RecordM Exp
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> RecordM Exp) -> Exp -> RecordM 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
-> ([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
fname ([TypeParamBase VName]
tparams, [PatBase Info VName ParamType]
params', Maybe (TypeExp Exp VName)
retdecl, Info ResRetType
ret, Exp
funbody') Exp
body' SrcLoc
loc) Info AppRes
res
transformExp (Lambda [PatBase Info VName ParamType]
params Exp
body Maybe (TypeExp Exp VName)
retdecl Info ResRetType
ret SrcLoc
loc) = do
  ([PatBase Info VName ParamType]
params', [RecordReplacements]
rrs) <- (PatBase Info VName ParamType
 -> RecordM (PatBase Info VName ParamType, RecordReplacements))
-> [PatBase Info VName ParamType]
-> RecordM ([PatBase Info VName ParamType], [RecordReplacements])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM ((ParamType -> RecordM ParamType)
-> PatBase Info VName ParamType
-> RecordM (PatBase Info VName ParamType, RecordReplacements)
forall u.
(TypeBase Exp u -> RecordM (TypeBase Exp u))
-> Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat ParamType -> RecordM ParamType
transformParamType) [PatBase Info VName ParamType]
params
  Exp
body' <- RecordReplacements -> RecordM Exp -> RecordM Exp
forall a. RecordReplacements -> RecordM a -> RecordM a
withRecordReplacements ([RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs) (RecordM Exp -> RecordM Exp) -> RecordM Exp -> RecordM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> RecordM Exp
transformExp Exp
body
  Exp -> RecordM Exp
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> RecordM Exp) -> Exp -> RecordM 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)
retdecl Info ResRetType
ret SrcLoc
loc
transformExp Exp
e = ASTMapper RecordM -> Exp -> RecordM Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap ASTMapper RecordM
m Exp
e
  where
    m :: ASTMapper RecordM
m = ASTMapper RecordM
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = transformExp}

onValBind :: ValBind -> RecordM ValBind
onValBind :: ValBind -> RecordM ValBind
onValBind ValBind
vb = do
  ([PatBase Info VName ParamType]
params', [RecordReplacements]
rrs) <- (PatBase Info VName ParamType
 -> RecordM (PatBase Info VName ParamType, RecordReplacements))
-> [PatBase Info VName ParamType]
-> RecordM ([PatBase Info VName ParamType], [RecordReplacements])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM ((ParamType -> RecordM ParamType)
-> PatBase Info VName ParamType
-> RecordM (PatBase Info VName ParamType, RecordReplacements)
forall u.
(TypeBase Exp u -> RecordM (TypeBase Exp u))
-> Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat ParamType -> RecordM ParamType
transformParamType) ([PatBase Info VName ParamType]
 -> RecordM ([PatBase Info VName ParamType], [RecordReplacements]))
-> [PatBase Info VName ParamType]
-> RecordM ([PatBase Info VName ParamType], [RecordReplacements])
forall a b. (a -> b) -> a -> b
$ ValBind -> [PatBase Info VName ParamType]
forall (f :: * -> *) vn.
ValBindBase f vn -> [PatBase f vn ParamType]
valBindParams ValBind
vb
  Exp
e' <- RecordReplacements -> RecordM Exp -> RecordM Exp
forall a. RecordReplacements -> RecordM a -> RecordM a
withRecordReplacements ([RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs) (RecordM Exp -> RecordM Exp) -> RecordM Exp -> RecordM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> RecordM Exp
transformExp (Exp -> RecordM Exp) -> Exp -> RecordM Exp
forall a b. (a -> b) -> a -> b
$ ValBind -> Exp
forall (f :: * -> *) vn. ValBindBase f vn -> ExpBase f vn
valBindBody ValBind
vb
  Info ResRetType
ret <- (ResRetType -> RecordM ResRetType)
-> Info ResRetType -> RecordM (Info ResRetType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Info a -> f (Info b)
traverse ((Exp -> RecordM Exp)
-> (Uniqueness -> RecordM Uniqueness)
-> ResRetType
-> RecordM ResRetType
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> RetTypeBase a b -> f (RetTypeBase c d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse Exp -> RecordM Exp
transformExp Uniqueness -> RecordM Uniqueness
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure) (Info ResRetType -> RecordM (Info ResRetType))
-> Info ResRetType -> RecordM (Info ResRetType)
forall a b. (a -> b) -> a -> b
$ ValBind -> Info ResRetType
forall (f :: * -> *) vn. ValBindBase f vn -> f ResRetType
valBindRetType ValBind
vb
  RecordM ()
memoClear
  ValBind -> RecordM ValBind
forall a. a -> RecordM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ValBind -> RecordM ValBind) -> ValBind -> RecordM ValBind
forall a b. (a -> b) -> a -> b
$
    ValBind
vb
      { valBindBody = e',
        valBindParams = params',
        valBindRetType = ret
      }

-- | Monomorphise a list of top-level declarations. A module-free input program
-- is expected, so only value declarations and type declaration are accepted.
transformProg :: (MonadFreshNames m) => [ValBind] -> m [ValBind]
transformProg :: forall (m :: * -> *). MonadFreshNames m => [ValBind] -> m [ValBind]
transformProg [ValBind]
vbs =
  (VNameSource -> ([ValBind], VNameSource)) -> m [ValBind]
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> ([ValBind], VNameSource)) -> m [ValBind])
-> (VNameSource -> ([ValBind], VNameSource)) -> m [ValBind]
forall a b. (a -> b) -> a -> b
$ \VNameSource
namesrc ->
    VNameSource -> RecordM [ValBind] -> ([ValBind], VNameSource)
forall a. VNameSource -> RecordM a -> (a, VNameSource)
runRecordM VNameSource
namesrc (RecordM [ValBind] -> ([ValBind], VNameSource))
-> RecordM [ValBind] -> ([ValBind], VNameSource)
forall a b. (a -> b) -> a -> b
$ (ValBind -> RecordM ValBind) -> [ValBind] -> RecordM [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 -> RecordM ValBind
onValBind [ValBind]
vbs