-- | 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.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 -> 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
    )

instance MonadFreshNames RecordM where
  getNameSource :: RecordM VNameSource
getNameSource = ReaderT Env (State VNameSource) VNameSource -> RecordM VNameSource
forall a. ReaderT Env (State VNameSource) a -> RecordM a
RecordM ReaderT Env (State VNameSource) VNameSource
forall s (m :: * -> *). MonadState s m => m s
get
  putNameSource :: VNameSource -> RecordM ()
putNameSource = ReaderT Env (State VNameSource) () -> RecordM ()
forall a. ReaderT Env (State VNameSource) a -> RecordM a
RecordM (ReaderT Env (State VNameSource) () -> RecordM ())
-> (VNameSource -> ReaderT Env (State VNameSource) ())
-> VNameSource
-> RecordM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VNameSource -> ReaderT Env (State VNameSource) ()
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) =
  State VNameSource a -> VNameSource -> (a, VNameSource)
forall s a. State s a -> s -> (a, s)
runState (ReaderT Env (State VNameSource) a -> Env -> State VNameSource a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT Env (State VNameSource) a
m (RecordReplacements -> Env
Env RecordReplacements
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 = (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 :: RecordReplacements
envRecordReplacements = RecordReplacements
rr RecordReplacements -> RecordReplacements -> RecordReplacements
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 = (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 =
  [(Name, PatBase Info VName (TypeBase Exp u))]
-> SrcLoc -> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t.
[(Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat ([Name]
-> [PatBase Info VName (TypeBase Exp u)]
-> [(Name, PatBase Info VName (TypeBase Exp u))]
forall a b. [a] -> [b] -> [(a, b)]
zip (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)]
 -> [(Name, PatBase Info VName (TypeBase Exp u))])
-> [PatBase Info VName (TypeBase Exp u)]
-> [(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

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' = 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
    ( [(Name, PatBase Info VName (TypeBase Exp u))]
-> SrcLoc -> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t.
[(Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat
        ([Name]
-> [PatBase Info VName (TypeBase Exp u)]
-> [(Name, PatBase Info VName (TypeBase Exp u))]
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
 -> 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 (RecordReplacement -> RecordReplacements)
-> RecordReplacement -> RecordReplacements
forall a b. (a -> b) -> a -> b
$ [(Name, (VName, StructType))] -> RecordReplacement
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, (VName, StructType))] -> RecordReplacement)
-> [(Name, (VName, StructType))] -> RecordReplacement
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
    )
transformPat (Id VName
v 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 (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 (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 PatBase Info VName (TypeBase Exp u)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall u.
Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat [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 (RecordPat [(Name, PatBase Info VName (TypeBase Exp u))]
fields SrcLoc
loc) = do
  let ([Name]
field_names, [PatBase Info VName (TypeBase Exp u)]
field_pats) = [(Name, PatBase Info VName (TypeBase Exp u))]
-> ([Name], [PatBase Info VName (TypeBase Exp u)])
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) <- (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 PatBase Info VName (TypeBase Exp u)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall u.
Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat [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 ([(Name, PatBase Info VName (TypeBase Exp u))]
-> SrcLoc -> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t.
[(Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat ([Name]
-> [PatBase Info VName (TypeBase Exp u)]
-> [(Name, PatBase Info VName (TypeBase Exp u))]
forall a b. [a] -> [b] -> [(a, b)]
zip [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 (PatParens PatBase Info VName (TypeBase Exp u)
pat SrcLoc
loc) = do
  (PatBase Info VName (TypeBase Exp u)
pat', RecordReplacements
rr) <- PatBase Info VName (TypeBase Exp u)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall u.
Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat 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 (PatAttr AttrInfo VName
attr PatBase Info VName (TypeBase Exp u)
pat SrcLoc
loc) = do
  (PatBase Info VName (TypeBase Exp u)
pat', RecordReplacements
rr) <- PatBase Info VName (TypeBase Exp u)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall u.
Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat 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 (Wildcard (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 (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 (PatAscription PatBase Info VName (TypeBase Exp u)
pat TypeExp Exp VName
_ SrcLoc
_) =
  PatBase Info VName (TypeBase Exp u)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
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) = (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 (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 PatBase Info VName (TypeBase Exp u)
-> RecordM
     (PatBase Info VName (TypeBase Exp u), RecordReplacements)
forall u.
Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat [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)

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 RecordReplacement
m
      | Just (VName
v, StructType
_) <- Name -> RecordReplacement -> Maybe (VName, StructType)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
n RecordReplacement
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 e :: Exp
e@(Var QualName VName
fname Info StructType
_ 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 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' = 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
f_t) SrcLoc
loc
            FieldBase Info vn -> f (FieldBase Info vn)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FieldBase Info vn -> f (FieldBase Info vn))
-> FieldBase Info vn -> f (FieldBase Info vn)
forall a b. (a -> b) -> a -> b
$ Name -> ExpBase Info vn -> SrcLoc -> FieldBase Info vn
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
f ExpBase Info vn
f_v' SrcLoc
loc
      [FieldBase Info VName] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ([FieldBase Info VName] -> SrcLoc -> Exp)
-> RecordM [FieldBase Info VName] -> RecordM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Name, (VName, StructType)) -> RecordM (FieldBase Info VName))
-> [(Name, (VName, StructType))] -> RecordM [FieldBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Name, (VName, StructType)) -> RecordM (FieldBase Info VName)
forall {f :: * -> *} {vn}.
Applicative f =>
(Name, (vn, StructType)) -> f (FieldBase Info vn)
toField (RecordReplacement -> [(Name, (VName, StructType))]
forall k a. Map k a -> [(k, a)]
M.toList RecordReplacement
fs) 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
    Maybe RecordReplacement
Nothing ->
      Exp -> RecordM Exp
forall a. a -> RecordM a
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) <- PatBase Info VName StructType
-> RecordM (PatBase Info VName StructType, RecordReplacements)
forall u.
Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat 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 {}) Info AppRes
_) = do
  String -> RecordM Exp
forall a. HasCallStack => String -> a
error String
"transformExp: LetFun is not supposed to occur"
transformExp (Lambda {}) =
  String -> RecordM Exp
forall a. HasCallStack => String -> a
error String
"transformExp: Lambda is not supposed to occur"
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 :: Exp -> RecordM Exp
mapOnExp = Exp -> RecordM Exp
transformExp}

onValBind :: ValBind -> RecordM ValBind
onValBind :: ValBind -> RecordM ValBind
onValBind ValBind
vb = do
  ([Pat (TypeBase Exp Diet)]
params', [RecordReplacements]
rrs) <- (Pat (TypeBase Exp Diet)
 -> RecordM (Pat (TypeBase Exp Diet), RecordReplacements))
-> [Pat (TypeBase Exp Diet)]
-> RecordM ([Pat (TypeBase Exp Diet)], [RecordReplacements])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM Pat (TypeBase Exp Diet)
-> RecordM (Pat (TypeBase Exp Diet), RecordReplacements)
forall u.
Pat (TypeBase Exp u)
-> RecordM (Pat (TypeBase Exp u), RecordReplacements)
transformPat ([Pat (TypeBase Exp Diet)]
 -> RecordM ([Pat (TypeBase Exp Diet)], [RecordReplacements]))
-> [Pat (TypeBase Exp Diet)]
-> RecordM ([Pat (TypeBase Exp Diet)], [RecordReplacements])
forall a b. (a -> b) -> a -> b
$ ValBind -> [Pat (TypeBase Exp Diet)]
forall (f :: * -> *) vn.
ValBindBase f vn -> [PatBase f vn (TypeBase Exp Diet)]
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
  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 :: Exp
valBindBody = Exp
e', valBindParams :: [Pat (TypeBase Exp Diet)]
valBindParams = [Pat (TypeBase Exp Diet)]
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 =
  (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