-- | 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.Reader
import Control.Monad.State
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.  This is used because the monomorphiser also expands all
-- record patterns.
type RecordReplacements = M.Map VName RecordReplacement

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

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

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

instance MonadFreshNames RecordM where
  getNameSource :: RecordM VNameSource
getNameSource = forall a. ReaderT Env (State VNameSource) a -> RecordM a
RecordM forall s (m :: * -> *). MonadState s m => m s
get
  putNameSource :: VNameSource -> RecordM ()
putNameSource = forall a. ReaderT Env (State VNameSource) a -> RecordM a
RecordM forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (m :: * -> *). MonadState s m => s -> m ()
put

runRecordM :: VNameSource -> RecordM a -> (a, VNameSource)
runRecordM :: forall a. VNameSource -> RecordM a -> (a, VNameSource)
runRecordM VNameSource
src (RecordM ReaderT Env (State VNameSource) a
m) =
  forall s a. State s a -> s -> (a, s)
runState (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT Env (State VNameSource) a
m (RecordReplacements -> Env
Env forall a. Monoid a => a
mempty)) VNameSource
src

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

lookupRecordReplacement :: VName -> RecordM (Maybe RecordReplacement)
lookupRecordReplacement :: VName -> RecordM (Maybe RecordReplacement)
lookupRecordReplacement VName
v = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v 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 =
  forall (f :: * -> *) vn t.
[(Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat (forall a b. [a] -> [b] -> [(a, b)]
zip (forall k a. Map k a -> [k]
M.keys Map Name (TypeBase Exp u)
fs) forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map ((forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
`Wildcard` SrcLoc
loc) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Info a
Info) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [a]
M.elems Map Name (TypeBase Exp u)
fs) SrcLoc
loc
wildcard TypeBase Exp u
t SrcLoc
loc =
  forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
Wildcard (forall a. a -> Info a
Info TypeBase Exp u
t) SrcLoc
loc

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

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

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

-- | 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 =
  forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource forall a b. (a -> b) -> a -> b
$ \VNameSource
namesrc ->
    forall a. VNameSource -> RecordM a -> (a, VNameSource)
runRecordM VNameSource
namesrc forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ValBind -> RecordM ValBind
onValBind [ValBind]
vbs