-- | 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
  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 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
      t' <- RecordM ParamType
m
      modify $ \S
s -> S
s {stateParamTypeMemo = M.insert t t' $ stateParamTypeMemo s}
      pure t'

memoStructType :: StructType -> RecordM StructType -> RecordM StructType
memoStructType :: StructType -> RecordM StructType -> RecordM StructType
memoStructType StructType
t RecordM StructType
m = do
  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 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
      t' <- RecordM StructType
m
      modify $ \S
s -> S
s {stateStructTypeMemo = M.insert t t' $ stateStructTypeMemo s}
      pure 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
  (fs_ks, 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
  pure
    ( RecordPat
        (zip (map (L noLoc . fst) fs') (zipWith3 Id fs_ks (map Info fs_ts) $ repeat loc))
        loc,
      M.singleton
        v
        ( M.fromList $ zip (map fst fs') $ zip fs_ks $ map toStruct fs_ts,
          RecordLit (zipWith3 toField (map fst fs') fs_ks fs_ts) 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
  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
  pure (Id v t' loc, mempty)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType (TuplePat [PatBase Info VName (TypeBase Exp u)]
pats SrcLoc
loc) = do
  (pats', 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
  pure (TuplePat pats' loc, mconcat 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
  (field_pats', 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
  pure (RecordPat (zip field_names field_pats') loc, mconcat rrs)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType (PatParens PatBase Info VName (TypeBase Exp u)
pat SrcLoc
loc) = do
  (pat', 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
  pure (PatParens pat' loc, rr)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType (PatAttr AttrInfo VName
attr PatBase Info VName (TypeBase Exp u)
pat SrcLoc
loc) = do
  (pat', 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
  pure (PatAttr attr pat' loc, rr)
transformPat TypeBase Exp u -> RecordM (TypeBase Exp u)
onType (Wildcard (Info TypeBase Exp u
t) SrcLoc
loc) = do
  t' <- TypeBase Exp u -> RecordM (TypeBase Exp u)
onType TypeBase Exp u
t
  pure (wildcard t' loc, 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
  (all_ps', 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
  pure (PatConstr name t all_ps' loc, mconcat 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_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_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
      e' <- Exp -> RecordM Exp
transformExp Exp
e
      pure $ Project n e' t loc
transformExp (Var QualName VName
fname Info StructType
t SrcLoc
loc) = do
  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_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
  e' <- Exp -> RecordM Exp
transformExp Exp
e
  (pat', rr) <- transformPat transformStructType pat
  body' <- withRecordReplacements rr $ transformExp body
  pure $ AppExp (LetPat sizes pat' e' body' loc) 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
  (params', 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
  funbody' <- withRecordReplacements (mconcat rrs) $ transformExp funbody
  body' <- transformExp body
  pure $ AppExp (LetFun fname (tparams, params', retdecl, ret, funbody') body' loc) res
transformExp (Lambda [PatBase Info VName ParamType]
params Exp
body Maybe (TypeExp Exp VName)
retdecl Info ResRetType
ret SrcLoc
loc) = do
  (params', 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
  body' <- withRecordReplacements (mconcat rrs) $ transformExp body
  pure $ Lambda params' body' retdecl ret 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,
          mapOnStructType = transformStructType,
          mapOnParamType = transformParamType
        }

onValBind :: ValBind -> RecordM ValBind
onValBind :: ValBind -> RecordM ValBind
onValBind ValBind
vb = do
  (params', 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
  e' <- withRecordReplacements (mconcat rrs) $ transformExp $ valBindBody vb
  ret <- traverse (bitraverse transformExp pure) $ valBindRetType vb
  memoClear
  pure $
    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