-- | Partially evaluate all modules away from a source Futhark
-- program.  This is implemented as a source-to-source transformation.
module Futhark.Internalise.Defunctorise (transformProg) where

import Control.Monad.Identity
import Control.Monad.RWS.Strict
import Data.DList qualified as DL
import Data.Map qualified as M
import Data.Maybe
import Data.Set qualified as S
import Futhark.MonadFreshNames
import Language.Futhark
import Language.Futhark.Semantic (FileModule (..), Imports, includeToString)
import Language.Futhark.Traversals
import Prelude hiding (abs, mod)

-- | A substitution from names in the original program to names in the
-- generated/residual program.
type Substitutions = M.Map VName VName

lookupSubst :: VName -> Substitutions -> VName
lookupSubst :: VName -> Substitutions -> VName
lookupSubst VName
v Substitutions
substs = case VName -> Substitutions -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Substitutions
substs of
  Just VName
v' | VName
v' VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
/= VName
v -> VName -> Substitutions -> VName
lookupSubst VName
v' Substitutions
substs
  Maybe VName
_ -> VName
v

data Mod
  = -- | A pairing of a lexical closure and a module function.
    ModFun TySet Scope ModParam ModExp
  | -- | A non-parametric module.
    ModMod Scope
  deriving (Int -> Mod -> ShowS
[Mod] -> ShowS
Mod -> [Char]
(Int -> Mod -> ShowS)
-> (Mod -> [Char]) -> ([Mod] -> ShowS) -> Show Mod
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Mod -> ShowS
showsPrec :: Int -> Mod -> ShowS
$cshow :: Mod -> [Char]
show :: Mod -> [Char]
$cshowList :: [Mod] -> ShowS
showList :: [Mod] -> ShowS
Show)

modScope :: Mod -> Scope
modScope :: Mod -> Scope
modScope (ModMod Scope
scope) = Scope
scope
modScope ModFun {} = Scope
forall a. Monoid a => a
mempty

data Scope = Scope
  { Scope -> Substitutions
scopeSubsts :: Substitutions,
    Scope -> Map VName Mod
scopeMods :: M.Map VName Mod
  }
  deriving (Int -> Scope -> ShowS
[Scope] -> ShowS
Scope -> [Char]
(Int -> Scope -> ShowS)
-> (Scope -> [Char]) -> ([Scope] -> ShowS) -> Show Scope
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Scope -> ShowS
showsPrec :: Int -> Scope -> ShowS
$cshow :: Scope -> [Char]
show :: Scope -> [Char]
$cshowList :: [Scope] -> ShowS
showList :: [Scope] -> ShowS
Show)

lookupSubstInScope :: QualName VName -> Scope -> (QualName VName, Scope)
lookupSubstInScope :: QualName VName -> Scope -> (QualName VName, Scope)
lookupSubstInScope qn :: QualName VName
qn@(QualName [VName]
quals VName
name) scope :: Scope
scope@(Scope Substitutions
substs Map VName Mod
mods) =
  case [VName]
quals of
    [] -> (VName -> QualName VName
forall v. v -> QualName v
qualName (VName -> QualName VName) -> VName -> QualName VName
forall a b. (a -> b) -> a -> b
$ VName -> Substitutions -> VName
lookupSubst VName
name Substitutions
substs, Scope
scope)
    VName
q : [VName]
qs ->
      let q' :: VName
q' = VName -> Substitutions -> VName
lookupSubst VName
q Substitutions
substs
       in case VName -> Map VName Mod -> Maybe Mod
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
q' Map VName Mod
mods of
            Just (ModMod Scope
mod_scope) -> QualName VName -> Scope -> (QualName VName, Scope)
lookupSubstInScope ([VName] -> VName -> QualName VName
forall vn. [vn] -> vn -> QualName vn
QualName [VName]
qs VName
name) Scope
mod_scope
            Maybe Mod
_ -> (QualName VName
qn, Scope
scope)

instance Semigroup Scope where
  Scope Substitutions
ss1 Map VName Mod
mt1 <> :: Scope -> Scope -> Scope
<> Scope Substitutions
ss2 Map VName Mod
mt2 = Substitutions -> Map VName Mod -> Scope
Scope (Substitutions
ss1 Substitutions -> Substitutions -> Substitutions
forall a. Semigroup a => a -> a -> a
<> Substitutions
ss2) (Map VName Mod
mt1 Map VName Mod -> Map VName Mod -> Map VName Mod
forall a. Semigroup a => a -> a -> a
<> Map VName Mod
mt2)

instance Monoid Scope where
  mempty :: Scope
mempty = Substitutions -> Map VName Mod -> Scope
Scope Substitutions
forall a. Monoid a => a
mempty Map VName Mod
forall a. Monoid a => a
mempty

type TySet = S.Set VName

data Env = Env
  { Env -> Scope
envScope :: Scope,
    Env -> Bool
envGenerating :: Bool,
    Env -> Map ImportName Scope
envImports :: M.Map ImportName Scope,
    Env -> TySet
envAbs :: TySet
  }

newtype TransformM a = TransformM (RWS Env (DL.DList Dec) VNameSource a)
  deriving
    ( Functor TransformM
Functor TransformM
-> (forall a. a -> TransformM a)
-> (forall a b.
    TransformM (a -> b) -> TransformM a -> TransformM b)
-> (forall a b c.
    (a -> b -> c) -> TransformM a -> TransformM b -> TransformM c)
-> (forall a b. TransformM a -> TransformM b -> TransformM b)
-> (forall a b. TransformM a -> TransformM b -> TransformM a)
-> Applicative TransformM
forall a. a -> TransformM a
forall a b. TransformM a -> TransformM b -> TransformM a
forall a b. TransformM a -> TransformM b -> TransformM b
forall a b. TransformM (a -> b) -> TransformM a -> TransformM b
forall a b c.
(a -> b -> c) -> TransformM a -> TransformM b -> TransformM 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 -> TransformM a
pure :: forall a. a -> TransformM a
$c<*> :: forall a b. TransformM (a -> b) -> TransformM a -> TransformM b
<*> :: forall a b. TransformM (a -> b) -> TransformM a -> TransformM b
$cliftA2 :: forall a b c.
(a -> b -> c) -> TransformM a -> TransformM b -> TransformM c
liftA2 :: forall a b c.
(a -> b -> c) -> TransformM a -> TransformM b -> TransformM c
$c*> :: forall a b. TransformM a -> TransformM b -> TransformM b
*> :: forall a b. TransformM a -> TransformM b -> TransformM b
$c<* :: forall a b. TransformM a -> TransformM b -> TransformM a
<* :: forall a b. TransformM a -> TransformM b -> TransformM a
Applicative,
      (forall a b. (a -> b) -> TransformM a -> TransformM b)
-> (forall a b. a -> TransformM b -> TransformM a)
-> Functor TransformM
forall a b. a -> TransformM b -> TransformM a
forall a b. (a -> b) -> TransformM a -> TransformM 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) -> TransformM a -> TransformM b
fmap :: forall a b. (a -> b) -> TransformM a -> TransformM b
$c<$ :: forall a b. a -> TransformM b -> TransformM a
<$ :: forall a b. a -> TransformM b -> TransformM a
Functor,
      Applicative TransformM
Applicative TransformM
-> (forall a b.
    TransformM a -> (a -> TransformM b) -> TransformM b)
-> (forall a b. TransformM a -> TransformM b -> TransformM b)
-> (forall a. a -> TransformM a)
-> Monad TransformM
forall a. a -> TransformM a
forall a b. TransformM a -> TransformM b -> TransformM b
forall a b. TransformM a -> (a -> TransformM b) -> TransformM 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. TransformM a -> (a -> TransformM b) -> TransformM b
>>= :: forall a b. TransformM a -> (a -> TransformM b) -> TransformM b
$c>> :: forall a b. TransformM a -> TransformM b -> TransformM b
>> :: forall a b. TransformM a -> TransformM b -> TransformM b
$creturn :: forall a. a -> TransformM a
return :: forall a. a -> TransformM a
Monad,
      Monad TransformM
TransformM VNameSource
Monad TransformM
-> TransformM VNameSource
-> (VNameSource -> TransformM ())
-> MonadFreshNames TransformM
VNameSource -> TransformM ()
forall (m :: * -> *).
Monad m
-> m VNameSource -> (VNameSource -> m ()) -> MonadFreshNames m
$cgetNameSource :: TransformM VNameSource
getNameSource :: TransformM VNameSource
$cputNameSource :: VNameSource -> TransformM ()
putNameSource :: VNameSource -> TransformM ()
MonadFreshNames,
      MonadReader Env,
      MonadWriter (DL.DList Dec)
    )

emit :: Dec -> TransformM ()
emit :: Dec -> TransformM ()
emit = DList Dec -> TransformM ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (DList Dec -> TransformM ())
-> (Dec -> DList Dec) -> Dec -> TransformM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dec -> DList Dec
forall a. a -> DList a
DL.singleton

askScope :: TransformM Scope
askScope :: TransformM Scope
askScope = (Env -> Scope) -> TransformM Scope
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> Scope
envScope

localScope :: (Scope -> Scope) -> TransformM a -> TransformM a
localScope :: forall a. (Scope -> Scope) -> TransformM a -> TransformM a
localScope Scope -> Scope
f = (Env -> Env) -> TransformM a -> TransformM a
forall a. (Env -> Env) -> TransformM a -> TransformM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env -> Env) -> TransformM a -> TransformM a)
-> (Env -> Env) -> TransformM a -> TransformM a
forall a b. (a -> b) -> a -> b
$ \Env
env -> Env
env {envScope :: Scope
envScope = Scope -> Scope
f (Scope -> Scope) -> Scope -> Scope
forall a b. (a -> b) -> a -> b
$ Env -> Scope
envScope Env
env}

extendScope :: Scope -> TransformM a -> TransformM a
extendScope :: forall a. Scope -> TransformM a -> TransformM a
extendScope (Scope Substitutions
substs Map VName Mod
mods) = (Scope -> Scope) -> TransformM a -> TransformM a
forall a. (Scope -> Scope) -> TransformM a -> TransformM a
localScope ((Scope -> Scope) -> TransformM a -> TransformM a)
-> (Scope -> Scope) -> TransformM a -> TransformM a
forall a b. (a -> b) -> a -> b
$ \Scope
scope ->
  Scope
scope
    { scopeSubsts :: Substitutions
scopeSubsts = (VName -> VName) -> Substitutions -> Substitutions
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Substitutions -> VName -> VName
forall {k}. Ord k => Map k k -> k -> k
forward (Scope -> Substitutions
scopeSubsts Scope
scope)) Substitutions
substs Substitutions -> Substitutions -> Substitutions
forall a. Semigroup a => a -> a -> a
<> Scope -> Substitutions
scopeSubsts Scope
scope,
      scopeMods :: Map VName Mod
scopeMods = Map VName Mod
mods Map VName Mod -> Map VName Mod -> Map VName Mod
forall a. Semigroup a => a -> a -> a
<> Scope -> Map VName Mod
scopeMods Scope
scope
    }
  where
    forward :: Map k k -> k -> k
forward Map k k
old_substs k
v = k -> Maybe k -> k
forall a. a -> Maybe a -> a
fromMaybe k
v (Maybe k -> k) -> Maybe k -> k
forall a b. (a -> b) -> a -> b
$ k -> Map k k -> Maybe k
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
v Map k k
old_substs

substituting :: Substitutions -> TransformM a -> TransformM a
substituting :: forall a. Substitutions -> TransformM a -> TransformM a
substituting Substitutions
substs = Scope -> TransformM a -> TransformM a
forall a. Scope -> TransformM a -> TransformM a
extendScope Scope
forall a. Monoid a => a
mempty {scopeSubsts :: Substitutions
scopeSubsts = Substitutions
substs}

boundName :: VName -> TransformM VName
boundName :: VName -> TransformM VName
boundName VName
v = do
  Bool
g <- (Env -> Bool) -> TransformM Bool
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> Bool
envGenerating
  if Bool
g then VName -> TransformM VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName VName
v else VName -> TransformM VName
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v

bindingNames :: [VName] -> TransformM Scope -> TransformM Scope
bindingNames :: [VName] -> TransformM Scope -> TransformM Scope
bindingNames [VName]
names TransformM Scope
m = do
  [VName]
names' <- (VName -> TransformM VName) -> [VName] -> TransformM [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 VName -> TransformM VName
boundName [VName]
names
  let substs :: Substitutions
substs = [(VName, VName)] -> Substitutions
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([VName] -> [VName] -> [(VName, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
names [VName]
names')
  Substitutions -> TransformM Scope -> TransformM Scope
forall a. Substitutions -> TransformM a -> TransformM a
substituting Substitutions
substs (TransformM Scope -> TransformM Scope)
-> TransformM Scope -> TransformM Scope
forall a b. (a -> b) -> a -> b
$ Scope -> Scope -> Scope
forall a. Monoid a => a -> a -> a
mappend (Scope -> Scope -> Scope)
-> TransformM Scope -> TransformM (Scope -> Scope)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TransformM Scope
m TransformM (Scope -> Scope) -> TransformM Scope -> TransformM Scope
forall a b. TransformM (a -> b) -> TransformM a -> TransformM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Scope -> TransformM Scope
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Substitutions -> Map VName Mod -> Scope
Scope Substitutions
substs Map VName Mod
forall a. Monoid a => a
mempty)

generating :: TransformM a -> TransformM a
generating :: forall a. TransformM a -> TransformM a
generating = (Env -> Env) -> TransformM a -> TransformM a
forall a. (Env -> Env) -> TransformM a -> TransformM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env -> Env) -> TransformM a -> TransformM a)
-> (Env -> Env) -> TransformM a -> TransformM a
forall a b. (a -> b) -> a -> b
$ \Env
env -> Env
env {envGenerating :: Bool
envGenerating = Bool
True}

bindingImport :: ImportName -> Scope -> TransformM a -> TransformM a
bindingImport :: forall a. ImportName -> Scope -> TransformM a -> TransformM a
bindingImport ImportName
name Scope
scope = (Env -> Env) -> TransformM a -> TransformM a
forall a. (Env -> Env) -> TransformM a -> TransformM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env -> Env) -> TransformM a -> TransformM a)
-> (Env -> Env) -> TransformM a -> TransformM a
forall a b. (a -> b) -> a -> b
$ \Env
env ->
  Env
env {envImports :: Map ImportName Scope
envImports = ImportName -> Scope -> Map ImportName Scope -> Map ImportName Scope
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert ImportName
name Scope
scope (Map ImportName Scope -> Map ImportName Scope)
-> Map ImportName Scope -> Map ImportName Scope
forall a b. (a -> b) -> a -> b
$ Env -> Map ImportName Scope
envImports Env
env}

bindingAbs :: TySet -> TransformM a -> TransformM a
bindingAbs :: forall a. TySet -> TransformM a -> TransformM a
bindingAbs TySet
abs = (Env -> Env) -> TransformM a -> TransformM a
forall a. (Env -> Env) -> TransformM a -> TransformM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env -> Env) -> TransformM a -> TransformM a)
-> (Env -> Env) -> TransformM a -> TransformM a
forall a b. (a -> b) -> a -> b
$ \Env
env ->
  Env
env {envAbs :: TySet
envAbs = TySet
abs TySet -> TySet -> TySet
forall a. Semigroup a => a -> a -> a
<> Env -> TySet
envAbs Env
env}

lookupImport :: ImportName -> TransformM Scope
lookupImport :: ImportName -> TransformM Scope
lookupImport ImportName
name = TransformM Scope
-> (Scope -> TransformM Scope) -> Maybe Scope -> TransformM Scope
forall b a. b -> (a -> b) -> Maybe a -> b
maybe TransformM Scope
forall {a}. a
bad Scope -> TransformM Scope
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Scope -> TransformM Scope)
-> TransformM (Maybe Scope) -> TransformM Scope
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Env -> Maybe Scope) -> TransformM (Maybe Scope)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (ImportName -> Map ImportName Scope -> Maybe Scope
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup ImportName
name (Map ImportName Scope -> Maybe Scope)
-> (Env -> Map ImportName Scope) -> Env -> Maybe Scope
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map ImportName Scope
envImports)
  where
    bad :: a
bad = [Char] -> a
forall a. HasCallStack => [Char] -> a
error ([Char] -> a) -> [Char] -> a
forall a b. (a -> b) -> a -> b
$ [Char]
"Defunctorise: unknown import: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ ImportName -> [Char]
includeToString ImportName
name

lookupMod' :: QualName VName -> Scope -> Either String Mod
lookupMod' :: QualName VName -> Scope -> Either [Char] Mod
lookupMod' QualName VName
mname Scope
scope =
  let (QualName VName
mname', Scope
scope') = QualName VName -> Scope -> (QualName VName, Scope)
lookupSubstInScope QualName VName
mname Scope
scope
   in Either [Char] Mod
-> (Mod -> Either [Char] Mod) -> Maybe Mod -> Either [Char] Mod
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([Char] -> Either [Char] Mod
forall a b. a -> Either a b
Left ([Char] -> Either [Char] Mod) -> [Char] -> Either [Char] Mod
forall a b. (a -> b) -> a -> b
$ QualName VName -> [Char]
forall {a}. Pretty a => a -> [Char]
bad QualName VName
mname') (Mod -> Either [Char] Mod
forall a b. b -> Either a b
Right (Mod -> Either [Char] Mod)
-> (Mod -> Mod) -> Mod -> Either [Char] Mod
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mod -> Mod
extend) (Maybe Mod -> Either [Char] Mod) -> Maybe Mod -> Either [Char] Mod
forall a b. (a -> b) -> a -> b
$ VName -> Map VName Mod -> Maybe Mod
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
mname') (Map VName Mod -> Maybe Mod) -> Map VName Mod -> Maybe Mod
forall a b. (a -> b) -> a -> b
$ Scope -> Map VName Mod
scopeMods Scope
scope'
  where
    bad :: a -> [Char]
bad a
mname' = [Char]
"Unknown module: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ QualName VName -> [Char]
forall {a}. Pretty a => a -> [Char]
prettyString QualName VName
mname [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
" (" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> [Char]
forall {a}. Pretty a => a -> [Char]
prettyString a
mname' [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
")"
    extend :: Mod -> Mod
extend (ModMod (Scope Substitutions
inner_scope Map VName Mod
inner_mods)) =
      -- XXX: perhaps hacky fix for #1653.  We need to impose the
      -- substitutions of abstract types from outside, because the
      -- inner module may have some incorrect substitutions in some
      -- cases.  Our treatment of abstract types is completely whack
      -- and should be fixed.
      Scope -> Mod
ModMod (Scope -> Mod) -> Scope -> Mod
forall a b. (a -> b) -> a -> b
$ Substitutions -> Map VName Mod -> Scope
Scope (Scope -> Substitutions
scopeSubsts Scope
scope Substitutions -> Substitutions -> Substitutions
forall a. Semigroup a => a -> a -> a
<> Substitutions
inner_scope) Map VName Mod
inner_mods
    extend Mod
m = Mod
m

lookupMod :: QualName VName -> TransformM Mod
lookupMod :: QualName VName -> TransformM Mod
lookupMod QualName VName
mname = ([Char] -> TransformM Mod)
-> (Mod -> TransformM Mod) -> Either [Char] Mod -> TransformM Mod
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either [Char] -> TransformM Mod
forall a. HasCallStack => [Char] -> a
error Mod -> TransformM Mod
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either [Char] Mod -> TransformM Mod)
-> (Scope -> Either [Char] Mod) -> Scope -> TransformM Mod
forall b c a. (b -> c) -> (a -> b) -> a -> c
. QualName VName -> Scope -> Either [Char] Mod
lookupMod' QualName VName
mname (Scope -> TransformM Mod) -> TransformM Scope -> TransformM Mod
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TransformM Scope
askScope

runTransformM :: VNameSource -> TransformM a -> (a, VNameSource, DL.DList Dec)
runTransformM :: forall a.
VNameSource -> TransformM a -> (a, VNameSource, DList Dec)
runTransformM VNameSource
src (TransformM RWS Env (DList Dec) VNameSource a
m) = RWS Env (DList Dec) VNameSource a
-> Env -> VNameSource -> (a, VNameSource, DList Dec)
forall r w s a. RWS r w s a -> r -> s -> (a, s, w)
runRWS RWS Env (DList Dec) VNameSource a
m Env
env VNameSource
src
  where
    env :: Env
env = Scope -> Bool -> Map ImportName Scope -> TySet -> Env
Env Scope
forall a. Monoid a => a
mempty Bool
False Map ImportName Scope
forall a. Monoid a => a
mempty TySet
forall a. Monoid a => a
mempty

maybeAscript ::
  SrcLoc ->
  Maybe (SigExp, Info (M.Map VName VName)) ->
  ModExp ->
  ModExp
maybeAscript :: SrcLoc -> Maybe (SigExp, Info Substitutions) -> ModExp -> ModExp
maybeAscript SrcLoc
loc (Just (SigExp
mtye, Info Substitutions
substs)) ModExp
me = ModExp -> SigExp -> Info Substitutions -> SrcLoc -> ModExp
forall (f :: * -> *) vn.
ModExpBase f vn
-> SigExpBase f vn -> f Substitutions -> SrcLoc -> ModExpBase f vn
ModAscript ModExp
me SigExp
mtye Info Substitutions
substs SrcLoc
loc
maybeAscript SrcLoc
_ Maybe (SigExp, Info Substitutions)
Nothing ModExp
me = ModExp
me

substituteInMod :: Substitutions -> Mod -> Mod
substituteInMod :: Substitutions -> Mod -> Mod
substituteInMod Substitutions
substs (ModMod (Scope Substitutions
mod_substs Map VName Mod
mod_mods)) =
  -- Forward all substitutions.
  Scope -> Mod
ModMod (Scope -> Mod) -> Scope -> Mod
forall a b. (a -> b) -> a -> b
$ Substitutions -> Map VName Mod -> Scope
Scope Substitutions
substs' (Map VName Mod -> Scope) -> Map VName Mod -> Scope
forall a b. (a -> b) -> a -> b
$ (Mod -> Mod) -> Map VName Mod -> Map VName Mod
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Substitutions -> Mod -> Mod
substituteInMod Substitutions
substs) Map VName Mod
mod_mods
  where
    forward :: VName -> VName
forward VName
v = VName -> Substitutions -> VName
lookupSubst VName
v (Substitutions -> VName) -> Substitutions -> VName
forall a b. (a -> b) -> a -> b
$ Substitutions
mod_substs Substitutions -> Substitutions -> Substitutions
forall a. Semigroup a => a -> a -> a
<> Substitutions
substs
    substs' :: Substitutions
substs' = (VName -> VName) -> Substitutions -> Substitutions
forall a b k. (a -> b) -> Map k a -> Map k b
M.map VName -> VName
forward Substitutions
substs
substituteInMod Substitutions
substs (ModFun TySet
abs (Scope Substitutions
mod_substs Map VName Mod
mod_mods) ModParam
mparam ModExp
mbody) =
  TySet -> Scope -> ModParam -> ModExp -> Mod
ModFun TySet
abs (Substitutions -> Map VName Mod -> Scope
Scope (Substitutions
substs' Substitutions -> Substitutions -> Substitutions
forall a. Semigroup a => a -> a -> a
<> Substitutions
mod_substs) Map VName Mod
mod_mods) ModParam
mparam ModExp
mbody
  where
    forward :: VName -> VName
forward VName
v = VName -> Substitutions -> VName
lookupSubst VName
v Substitutions
mod_substs
    substs' :: Substitutions
substs' = (VName -> VName) -> Substitutions -> Substitutions
forall a b k. (a -> b) -> Map k a -> Map k b
M.map VName -> VName
forward Substitutions
substs

extendAbsTypes :: Substitutions -> TransformM a -> TransformM a
extendAbsTypes :: forall a. Substitutions -> TransformM a -> TransformM a
extendAbsTypes Substitutions
ascript_substs TransformM a
m = do
  TySet
abs <- (Env -> TySet) -> TransformM TySet
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> TySet
envAbs
  -- Some abstract types may have a different name on the inside, and
  -- we need to make them visible, because substitutions involving
  -- abstract types must be lifted out in transformModBind.
  let subst_abs :: TySet
subst_abs =
        [VName] -> TySet
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> TySet)
-> ([(VName, VName)] -> [VName]) -> [(VName, VName)] -> TySet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, VName) -> VName) -> [(VName, VName)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, VName) -> VName
forall a b. (a, b) -> b
snd ([(VName, VName)] -> [VName])
-> ([(VName, VName)] -> [(VName, VName)])
-> [(VName, VName)]
-> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, VName) -> Bool) -> [(VName, VName)] -> [(VName, VName)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> TySet -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` TySet
abs) (VName -> Bool)
-> ((VName, VName) -> VName) -> (VName, VName) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, VName) -> VName
forall a b. (a, b) -> a
fst) ([(VName, VName)] -> TySet) -> [(VName, VName)] -> TySet
forall a b. (a -> b) -> a -> b
$
          Substitutions -> [(VName, VName)]
forall k a. Map k a -> [(k, a)]
M.toList Substitutions
ascript_substs
  TySet -> TransformM a -> TransformM a
forall a. TySet -> TransformM a -> TransformM a
bindingAbs TySet
subst_abs TransformM a
m

evalModExp :: ModExp -> TransformM Mod
evalModExp :: ModExp -> TransformM Mod
evalModExp (ModVar QualName VName
qn SrcLoc
_) = QualName VName -> TransformM Mod
lookupMod QualName VName
qn
evalModExp (ModParens ModExp
e SrcLoc
_) = ModExp -> TransformM Mod
evalModExp ModExp
e
evalModExp (ModDecs [Dec]
decs SrcLoc
_) = Scope -> Mod
ModMod (Scope -> Mod) -> TransformM Scope -> TransformM Mod
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Dec] -> TransformM Scope
transformDecs [Dec]
decs
evalModExp (ModImport [Char]
_ (Info ImportName
fpath) SrcLoc
_) = Scope -> Mod
ModMod (Scope -> Mod) -> TransformM Scope -> TransformM Mod
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ImportName -> TransformM Scope
lookupImport ImportName
fpath
evalModExp (ModAscript ModExp
me SigExp
_ (Info Substitutions
ascript_substs) SrcLoc
_) =
  Substitutions -> TransformM Mod -> TransformM Mod
forall a. Substitutions -> TransformM a -> TransformM a
extendAbsTypes Substitutions
ascript_substs (TransformM Mod -> TransformM Mod)
-> TransformM Mod -> TransformM Mod
forall a b. (a -> b) -> a -> b
$
    Substitutions -> Mod -> Mod
substituteInMod Substitutions
ascript_substs (Mod -> Mod) -> TransformM Mod -> TransformM Mod
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ModExp -> TransformM Mod
evalModExp ModExp
me
evalModExp (ModApply ModExp
f ModExp
arg (Info Substitutions
p_substs) (Info Substitutions
b_substs) SrcLoc
loc) = do
  Mod
f_mod <- ModExp -> TransformM Mod
evalModExp ModExp
f
  Mod
arg_mod <- ModExp -> TransformM Mod
evalModExp ModExp
arg
  case Mod
f_mod of
    ModMod Scope
_ ->
      [Char] -> TransformM Mod
forall a. HasCallStack => [Char] -> a
error ([Char] -> TransformM Mod) -> [Char] -> TransformM Mod
forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot apply non-parametric module at " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr SrcLoc
loc
    ModFun TySet
f_abs Scope
f_closure ModParam
f_p ModExp
f_body ->
      TySet -> TransformM Mod -> TransformM Mod
forall a. TySet -> TransformM a -> TransformM a
bindingAbs (TySet
f_abs TySet -> TySet -> TySet
forall a. Semigroup a => a -> a -> a
<> [VName] -> TySet
forall a. Ord a => [a] -> Set a
S.fromList (Info [VName] -> [VName]
forall a. Info a -> a
unInfo (ModParam -> Info [VName]
forall (f :: * -> *) vn. ModParamBase f vn -> f [VName]
modParamAbs ModParam
f_p)))
        (TransformM Mod -> TransformM Mod)
-> (TransformM Mod -> TransformM Mod)
-> TransformM Mod
-> TransformM Mod
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Substitutions -> TransformM Mod -> TransformM Mod
forall a. Substitutions -> TransformM a -> TransformM a
extendAbsTypes Substitutions
b_substs
        (TransformM Mod -> TransformM Mod)
-> (TransformM Mod -> TransformM Mod)
-> TransformM Mod
-> TransformM Mod
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Scope -> Scope) -> TransformM Mod -> TransformM Mod
forall a. (Scope -> Scope) -> TransformM a -> TransformM a
localScope (Scope -> Scope -> Scope
forall a b. a -> b -> a
const Scope
f_closure) -- Start afresh.
        (TransformM Mod -> TransformM Mod)
-> (TransformM Mod -> TransformM Mod)
-> TransformM Mod
-> TransformM Mod
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TransformM Mod -> TransformM Mod
forall a. TransformM a -> TransformM a
generating
        (TransformM Mod -> TransformM Mod)
-> TransformM Mod -> TransformM Mod
forall a b. (a -> b) -> a -> b
$ do
          TySet
abs <- (Env -> TySet) -> TransformM TySet
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> TySet
envAbs
          let keep :: VName -> p -> Bool
keep VName
k p
_ = VName
k VName -> Substitutions -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Substitutions
p_substs Bool -> Bool -> Bool
|| VName
k VName -> TySet -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` TySet
abs
              abs_substs :: Substitutions
abs_substs =
                (VName -> VName -> Bool) -> Substitutions -> Substitutions
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey VName -> VName -> Bool
forall {p}. VName -> p -> Bool
keep (Substitutions -> Substitutions) -> Substitutions -> Substitutions
forall a b. (a -> b) -> a -> b
$
                  (VName -> VName) -> Substitutions -> Substitutions
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (VName -> Substitutions -> VName
`lookupSubst` Scope -> Substitutions
scopeSubsts (Mod -> Scope
modScope Mod
arg_mod)) Substitutions
p_substs
                    Substitutions -> Substitutions -> Substitutions
forall a. Semigroup a => a -> a -> a
<> Scope -> Substitutions
scopeSubsts Scope
f_closure
                    Substitutions -> Substitutions -> Substitutions
forall a. Semigroup a => a -> a -> a
<> Scope -> Substitutions
scopeSubsts (Mod -> Scope
modScope Mod
arg_mod)
          Scope -> TransformM Mod -> TransformM Mod
forall a. Scope -> TransformM a -> TransformM a
extendScope
            ( Substitutions -> Map VName Mod -> Scope
Scope
                Substitutions
abs_substs
                ( VName -> Mod -> Map VName Mod
forall k a. k -> a -> Map k a
M.singleton (ModParam -> VName
forall (f :: * -> *) vn. ModParamBase f vn -> vn
modParamName ModParam
f_p) (Mod -> Map VName Mod) -> Mod -> Map VName Mod
forall a b. (a -> b) -> a -> b
$
                    Substitutions -> Mod -> Mod
substituteInMod Substitutions
p_substs Mod
arg_mod
                )
            )
            (TransformM Mod -> TransformM Mod)
-> TransformM Mod -> TransformM Mod
forall a b. (a -> b) -> a -> b
$ do
              Substitutions
substs <- Scope -> Substitutions
scopeSubsts (Scope -> Substitutions)
-> TransformM Scope -> TransformM Substitutions
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TransformM Scope
askScope
              Mod
x <- ModExp -> TransformM Mod
evalModExp ModExp
f_body
              Mod -> TransformM Mod
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Mod -> TransformM Mod) -> Mod -> TransformM Mod
forall a b. (a -> b) -> a -> b
$
                TySet -> Substitutions -> Mod -> Mod
addSubsts TySet
abs Substitutions
abs_substs (Mod -> Mod) -> Mod -> Mod
forall a b. (a -> b) -> a -> b
$
                  -- The next one is dubious, but is necessary to
                  -- propagate substitutions from the argument (see
                  -- modules/functor24.fut).
                  Substitutions -> Mod -> Mod
addSubstsModMod (Scope -> Substitutions
scopeSubsts (Scope -> Substitutions) -> Scope -> Substitutions
forall a b. (a -> b) -> a -> b
$ Mod -> Scope
modScope Mod
arg_mod) (Mod -> Mod) -> Mod -> Mod
forall a b. (a -> b) -> a -> b
$
                    Substitutions -> Mod -> Mod
substituteInMod (Substitutions
b_substs Substitutions -> Substitutions -> Substitutions
forall a. Semigroup a => a -> a -> a
<> Substitutions
substs) Mod
x
  where
    addSubsts :: TySet -> Substitutions -> Mod -> Mod
addSubsts TySet
abs Substitutions
substs (ModFun TySet
mabs (Scope Substitutions
msubsts Map VName Mod
mods) ModParam
mp ModExp
me) =
      TySet -> Scope -> ModParam -> ModExp -> Mod
ModFun (TySet
abs TySet -> TySet -> TySet
forall a. Semigroup a => a -> a -> a
<> TySet
mabs) (Substitutions -> Map VName Mod -> Scope
Scope (Substitutions
substs Substitutions -> Substitutions -> Substitutions
forall a. Semigroup a => a -> a -> a
<> Substitutions
msubsts) Map VName Mod
mods) ModParam
mp ModExp
me
    addSubsts TySet
_ Substitutions
substs (ModMod (Scope Substitutions
msubsts Map VName Mod
mods)) =
      Scope -> Mod
ModMod (Scope -> Mod) -> Scope -> Mod
forall a b. (a -> b) -> a -> b
$ Substitutions -> Map VName Mod -> Scope
Scope (Substitutions
substs Substitutions -> Substitutions -> Substitutions
forall a. Semigroup a => a -> a -> a
<> Substitutions
msubsts) Map VName Mod
mods
    addSubstsModMod :: Substitutions -> Mod -> Mod
addSubstsModMod Substitutions
substs (ModMod (Scope Substitutions
msubsts Map VName Mod
mods)) =
      Scope -> Mod
ModMod (Scope -> Mod) -> Scope -> Mod
forall a b. (a -> b) -> a -> b
$ Substitutions -> Map VName Mod -> Scope
Scope (Substitutions
substs Substitutions -> Substitutions -> Substitutions
forall a. Semigroup a => a -> a -> a
<> Substitutions
msubsts) Map VName Mod
mods
    addSubstsModMod Substitutions
_ Mod
m = Mod
m
evalModExp (ModLambda ModParam
p Maybe (SigExp, Info Substitutions)
ascript ModExp
e SrcLoc
loc) = do
  Scope
scope <- TransformM Scope
askScope
  TySet
abs <- (Env -> TySet) -> TransformM TySet
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> TySet
envAbs
  Mod -> TransformM Mod
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Mod -> TransformM Mod) -> Mod -> TransformM Mod
forall a b. (a -> b) -> a -> b
$ TySet -> Scope -> ModParam -> ModExp -> Mod
ModFun TySet
abs Scope
scope ModParam
p (ModExp -> Mod) -> ModExp -> Mod
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Maybe (SigExp, Info Substitutions) -> ModExp -> ModExp
maybeAscript SrcLoc
loc Maybe (SigExp, Info Substitutions)
ascript ModExp
e

transformName :: VName -> TransformM VName
transformName :: VName -> TransformM VName
transformName VName
v = VName -> Substitutions -> VName
lookupSubst VName
v (Substitutions -> VName)
-> (Scope -> Substitutions) -> Scope -> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope -> Substitutions
scopeSubsts (Scope -> VName) -> TransformM Scope -> TransformM VName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TransformM Scope
askScope

-- | A general-purpose substitution of names.
transformNames :: (ASTMappable x) => x -> TransformM x
transformNames :: forall x. ASTMappable x => x -> TransformM x
transformNames x
x = do
  Scope
scope <- TransformM Scope
askScope
  x -> TransformM x
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (x -> TransformM x) -> x -> TransformM x
forall a b. (a -> b) -> a -> b
$ Identity x -> x
forall a. Identity a -> a
runIdentity (Identity x -> x) -> Identity x -> x
forall a b. (a -> b) -> a -> b
$ ASTMapper Identity -> x -> Identity x
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> x -> m x
astMap (Scope -> ASTMapper Identity
forall {m :: * -> *}. Monad m => Scope -> ASTMapper m
substituter Scope
scope) x
x
  where
    substituter :: Scope -> ASTMapper m
substituter Scope
scope =
      ASTMapper
        { mapOnExp :: Exp -> m Exp
mapOnExp = Scope -> Exp -> m Exp
onExp Scope
scope,
          mapOnName :: QualName VName -> m (QualName VName)
mapOnName = \QualName VName
v -> QualName VName -> m (QualName VName)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName VName -> m (QualName VName))
-> QualName VName -> m (QualName VName)
forall a b. (a -> b) -> a -> b
$ (QualName VName, Scope) -> QualName VName
forall a b. (a, b) -> a
fst ((QualName VName, Scope) -> QualName VName)
-> (QualName VName, Scope) -> QualName VName
forall a b. (a -> b) -> a -> b
$ QualName VName -> Scope -> (QualName VName, Scope)
lookupSubstInScope QualName VName
v {qualQuals :: [VName]
qualQuals = []} Scope
scope,
          mapOnStructType :: StructType -> m StructType
mapOnStructType = ASTMapper m -> StructType -> m StructType
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *).
Monad m =>
ASTMapper m -> StructType -> m StructType
astMap (Scope -> ASTMapper m
substituter Scope
scope),
          mapOnParamType :: ParamType -> m ParamType
mapOnParamType = ASTMapper m -> ParamType -> m ParamType
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *).
Monad m =>
ASTMapper m -> ParamType -> m ParamType
astMap (Scope -> ASTMapper m
substituter Scope
scope),
          mapOnResRetType :: ResRetType -> m ResRetType
mapOnResRetType = ASTMapper m -> ResRetType -> m ResRetType
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *).
Monad m =>
ASTMapper m -> ResRetType -> m ResRetType
astMap (Scope -> ASTMapper m
substituter Scope
scope)
        }
    onExp :: Scope -> Exp -> m Exp
onExp Scope
scope Exp
e =
      -- One expression is tricky, because it interacts with scoping rules.
      case Exp
e of
        QualParens (QualName VName
mn, SrcLoc
_) Exp
e' SrcLoc
_ ->
          case QualName VName -> Scope -> Either [Char] Mod
lookupMod' QualName VName
mn Scope
scope of
            Left [Char]
err -> [Char] -> m Exp
forall a. HasCallStack => [Char] -> a
error [Char]
err
            Right Mod
mod ->
              ASTMapper m -> Exp -> m Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap (Scope -> ASTMapper m
substituter (Scope -> ASTMapper m) -> Scope -> ASTMapper m
forall a b. (a -> b) -> a -> b
$ Mod -> Scope
modScope Mod
mod Scope -> Scope -> Scope
forall a. Semigroup a => a -> a -> a
<> Scope
scope) Exp
e'
        Exp
_ -> ASTMapper m -> Exp -> m Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap (Scope -> ASTMapper m
substituter Scope
scope) Exp
e

transformTypeExp :: TypeExp Info VName -> TransformM (TypeExp Info VName)
transformTypeExp :: TypeExp Info VName -> TransformM (TypeExp Info VName)
transformTypeExp = TypeExp Info VName -> TransformM (TypeExp Info VName)
forall x. ASTMappable x => x -> TransformM x
transformNames

transformStructType :: StructType -> TransformM StructType
transformStructType :: StructType -> TransformM StructType
transformStructType = StructType -> TransformM StructType
forall x. ASTMappable x => x -> TransformM x
transformNames

transformResType :: ResType -> TransformM ResType
transformResType :: ResType -> TransformM ResType
transformResType = ResType -> TransformM ResType
forall x. ASTMappable x => x -> TransformM x
transformNames

transformExp :: Exp -> TransformM Exp
transformExp :: Exp -> TransformM Exp
transformExp = Exp -> TransformM Exp
forall x. ASTMappable x => x -> TransformM x
transformNames

transformEntry :: EntryPoint -> TransformM EntryPoint
transformEntry :: EntryPoint -> TransformM EntryPoint
transformEntry (EntryPoint [EntryParam]
params EntryType
ret) =
  [EntryParam] -> EntryType -> EntryPoint
EntryPoint ([EntryParam] -> EntryType -> EntryPoint)
-> TransformM [EntryParam] -> TransformM (EntryType -> EntryPoint)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (EntryParam -> TransformM EntryParam)
-> [EntryParam] -> TransformM [EntryParam]
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 EntryParam -> TransformM EntryParam
onEntryParam [EntryParam]
params TransformM (EntryType -> EntryPoint)
-> TransformM EntryType -> TransformM EntryPoint
forall a b. TransformM (a -> b) -> TransformM a -> TransformM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> EntryType -> TransformM EntryType
onEntryType EntryType
ret
  where
    onEntryParam :: EntryParam -> TransformM EntryParam
onEntryParam (EntryParam Name
v EntryType
t) =
      Name -> EntryType -> EntryParam
EntryParam Name
v (EntryType -> EntryParam)
-> TransformM EntryType -> TransformM EntryParam
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EntryType -> TransformM EntryType
onEntryType EntryType
t
    onEntryType :: EntryType -> TransformM EntryType
onEntryType (EntryType StructType
t Maybe (TypeExp Info VName)
te) =
      StructType -> Maybe (TypeExp Info VName) -> EntryType
EntryType (StructType -> Maybe (TypeExp Info VName) -> EntryType)
-> TransformM StructType
-> TransformM (Maybe (TypeExp Info VName) -> EntryType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructType -> TransformM StructType
transformStructType StructType
t TransformM (Maybe (TypeExp Info VName) -> EntryType)
-> TransformM (Maybe (TypeExp Info VName)) -> TransformM EntryType
forall a b. TransformM (a -> b) -> TransformM a -> TransformM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe (TypeExp Info VName)
-> TransformM (Maybe (TypeExp Info VName))
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (TypeExp Info VName)
te

transformValBind :: ValBind -> TransformM ()
transformValBind :: ValBind -> TransformM ()
transformValBind (ValBind Maybe (Info EntryPoint)
entry VName
name Maybe (TypeExp Info VName)
tdecl (Info (RetType [VName]
dims ResType
t)) [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params Exp
e Maybe DocComment
doc [AttrInfo VName]
attrs SrcLoc
loc) = do
  Maybe (Info EntryPoint)
entry' <- (Info EntryPoint -> TransformM (Info EntryPoint))
-> Maybe (Info EntryPoint) -> TransformM (Maybe (Info EntryPoint))
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) -> Maybe a -> f (Maybe b)
traverse ((EntryPoint -> TransformM EntryPoint)
-> Info EntryPoint -> TransformM (Info EntryPoint)
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 EntryPoint -> TransformM EntryPoint
transformEntry) Maybe (Info EntryPoint)
entry
  VName
name' <- VName -> TransformM VName
transformName VName
name
  Maybe (TypeExp Info VName)
tdecl' <- (TypeExp Info VName -> TransformM (TypeExp Info VName))
-> Maybe (TypeExp Info VName)
-> TransformM (Maybe (TypeExp Info VName))
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) -> Maybe a -> f (Maybe b)
traverse TypeExp Info VName -> TransformM (TypeExp Info VName)
transformTypeExp Maybe (TypeExp Info VName)
tdecl
  ResType
t' <- ResType -> TransformM ResType
transformResType ResType
t
  Exp
e' <- Exp -> TransformM Exp
transformExp Exp
e
  [PatBase Info VName ParamType]
params' <- (PatBase Info VName ParamType
 -> TransformM (PatBase Info VName ParamType))
-> [PatBase Info VName ParamType]
-> TransformM [PatBase Info VName ParamType]
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) -> [a] -> f [b]
traverse PatBase Info VName ParamType
-> TransformM (PatBase Info VName ParamType)
forall x. ASTMappable x => x -> TransformM x
transformNames [PatBase Info VName ParamType]
params
  Dec -> TransformM ()
emit (Dec -> TransformM ()) -> Dec -> TransformM ()
forall a b. (a -> b) -> a -> b
$ ValBind -> Dec
forall (f :: * -> *) vn. ValBindBase f vn -> DecBase f vn
ValDec (ValBind -> Dec) -> ValBind -> Dec
forall a b. (a -> b) -> a -> b
$ Maybe (Info EntryPoint)
-> VName
-> Maybe (TypeExp Info VName)
-> Info ResRetType
-> [TypeParamBase VName]
-> [PatBase Info VName ParamType]
-> Exp
-> Maybe DocComment
-> [AttrInfo VName]
-> SrcLoc
-> ValBind
forall (f :: * -> *) vn.
Maybe (f EntryPoint)
-> vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> [TypeParamBase vn]
-> [PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe DocComment
-> [AttrInfo vn]
-> SrcLoc
-> ValBindBase f vn
ValBind Maybe (Info EntryPoint)
entry' VName
name' Maybe (TypeExp Info VName)
tdecl' (ResRetType -> Info ResRetType
forall a. a -> Info a
Info ([VName] -> ResType -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims ResType
t')) [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params' Exp
e' Maybe DocComment
doc [AttrInfo VName]
attrs SrcLoc
loc

transformTypeBind :: TypeBind -> TransformM ()
transformTypeBind :: TypeBind -> TransformM ()
transformTypeBind (TypeBind VName
name Liftedness
l [TypeParamBase VName]
tparams TypeExp Info VName
te (Info (RetType [VName]
dims StructType
t)) Maybe DocComment
doc SrcLoc
loc) = do
  VName
name' <- VName -> TransformM VName
transformName VName
name
  Dec -> TransformM ()
emit (Dec -> TransformM ())
-> (TypeBind -> Dec) -> TypeBind -> TransformM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBind -> Dec
forall (f :: * -> *) vn. TypeBindBase f vn -> DecBase f vn
TypeDec
    (TypeBind -> TransformM ()) -> TransformM TypeBind -> TransformM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ( VName
-> Liftedness
-> [TypeParamBase VName]
-> TypeExp Info VName
-> Info StructRetType
-> Maybe DocComment
-> SrcLoc
-> TypeBind
forall (f :: * -> *) vn.
vn
-> Liftedness
-> [TypeParamBase vn]
-> TypeExp f vn
-> f StructRetType
-> Maybe DocComment
-> SrcLoc
-> TypeBindBase f vn
TypeBind VName
name' Liftedness
l [TypeParamBase VName]
tparams
            (TypeExp Info VName
 -> Info StructRetType -> Maybe DocComment -> SrcLoc -> TypeBind)
-> TransformM (TypeExp Info VName)
-> TransformM
     (Info StructRetType -> Maybe DocComment -> SrcLoc -> TypeBind)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeExp Info VName -> TransformM (TypeExp Info VName)
transformTypeExp TypeExp Info VName
te
            TransformM
  (Info StructRetType -> Maybe DocComment -> SrcLoc -> TypeBind)
-> TransformM (Info StructRetType)
-> TransformM (Maybe DocComment -> SrcLoc -> TypeBind)
forall a b. TransformM (a -> b) -> TransformM a -> TransformM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (StructRetType -> Info StructRetType
forall a. a -> Info a
Info (StructRetType -> Info StructRetType)
-> (StructType -> StructRetType)
-> StructType
-> Info StructRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> StructType -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (StructType -> Info StructRetType)
-> TransformM StructType -> TransformM (Info StructRetType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructType -> TransformM StructType
transformStructType StructType
t)
            TransformM (Maybe DocComment -> SrcLoc -> TypeBind)
-> TransformM (Maybe DocComment) -> TransformM (SrcLoc -> TypeBind)
forall a b. TransformM (a -> b) -> TransformM a -> TransformM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe DocComment -> TransformM (Maybe DocComment)
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe DocComment
doc
            TransformM (SrcLoc -> TypeBind)
-> TransformM SrcLoc -> TransformM TypeBind
forall a b. TransformM (a -> b) -> TransformM a -> TransformM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TransformM SrcLoc
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
        )

transformModBind :: ModBind -> TransformM Scope
transformModBind :: ModBind -> TransformM Scope
transformModBind ModBind
mb = do
  let addParam :: ModParamBase f vn -> ModExpBase f vn -> ModExpBase f vn
addParam ModParamBase f vn
p ModExpBase f vn
me = ModParamBase f vn
-> Maybe (SigExpBase f vn, f Substitutions)
-> ModExpBase f vn
-> SrcLoc
-> ModExpBase f vn
forall (f :: * -> *) vn.
ModParamBase f vn
-> Maybe (SigExpBase f vn, f Substitutions)
-> ModExpBase f vn
-> SrcLoc
-> ModExpBase f vn
ModLambda ModParamBase f vn
p Maybe (SigExpBase f vn, f Substitutions)
forall a. Maybe a
Nothing ModExpBase f vn
me (SrcLoc -> ModExpBase f vn) -> SrcLoc -> ModExpBase f vn
forall a b. (a -> b) -> a -> b
$ ModExpBase f vn -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ModExpBase f vn
me
  Mod
mod <-
    ModExp -> TransformM Mod
evalModExp
      (ModExp -> TransformM Mod) -> ModExp -> TransformM Mod
forall a b. (a -> b) -> a -> b
$ (ModParam -> ModExp -> ModExp) -> ModExp -> [ModParam] -> ModExp
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
        ModParam -> ModExp -> ModExp
forall {f :: * -> *} {vn}.
ModParamBase f vn -> ModExpBase f vn -> ModExpBase f vn
addParam
        (SrcLoc -> Maybe (SigExp, Info Substitutions) -> ModExp -> ModExp
maybeAscript (ModBind -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ModBind
mb) (ModBind -> Maybe (SigExp, Info Substitutions)
forall (f :: * -> *) vn.
ModBindBase f vn -> Maybe (SigExpBase f vn, f Substitutions)
modSignature ModBind
mb) (ModExp -> ModExp) -> ModExp -> ModExp
forall a b. (a -> b) -> a -> b
$ ModBind -> ModExp
forall (f :: * -> *) vn. ModBindBase f vn -> ModExpBase f vn
modExp ModBind
mb)
      ([ModParam] -> ModExp) -> [ModParam] -> ModExp
forall a b. (a -> b) -> a -> b
$ ModBind -> [ModParam]
forall (f :: * -> *) vn. ModBindBase f vn -> [ModParamBase f vn]
modParams ModBind
mb
  VName
mname <- VName -> TransformM VName
transformName (VName -> TransformM VName) -> VName -> TransformM VName
forall a b. (a -> b) -> a -> b
$ ModBind -> VName
forall (f :: * -> *) vn. ModBindBase f vn -> vn
modName ModBind
mb
  Scope -> TransformM Scope
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Scope -> TransformM Scope) -> Scope -> TransformM Scope
forall a b. (a -> b) -> a -> b
$ Substitutions -> Map VName Mod -> Scope
Scope (Scope -> Substitutions
scopeSubsts (Scope -> Substitutions) -> Scope -> Substitutions
forall a b. (a -> b) -> a -> b
$ Mod -> Scope
modScope Mod
mod) (Map VName Mod -> Scope) -> Map VName Mod -> Scope
forall a b. (a -> b) -> a -> b
$ VName -> Mod -> Map VName Mod
forall k a. k -> a -> Map k a
M.singleton VName
mname Mod
mod

transformDecs :: [Dec] -> TransformM Scope
transformDecs :: [Dec] -> TransformM Scope
transformDecs [Dec]
ds =
  case [Dec]
ds of
    [] ->
      Scope -> TransformM Scope
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Scope
forall a. Monoid a => a
mempty
    LocalDec Dec
d SrcLoc
_ : [Dec]
ds' ->
      [Dec] -> TransformM Scope
transformDecs ([Dec] -> TransformM Scope) -> [Dec] -> TransformM Scope
forall a b. (a -> b) -> a -> b
$ Dec
d Dec -> [Dec] -> [Dec]
forall a. a -> [a] -> [a]
: [Dec]
ds'
    ValDec ValBind
fdec : [Dec]
ds' ->
      [VName] -> TransformM Scope -> TransformM Scope
bindingNames [ValBind -> VName
forall (f :: * -> *) vn. ValBindBase f vn -> vn
valBindName ValBind
fdec] (TransformM Scope -> TransformM Scope)
-> TransformM Scope -> TransformM Scope
forall a b. (a -> b) -> a -> b
$ do
        ValBind -> TransformM ()
transformValBind ValBind
fdec
        [Dec] -> TransformM Scope
transformDecs [Dec]
ds'
    TypeDec TypeBind
tb : [Dec]
ds' ->
      [VName] -> TransformM Scope -> TransformM Scope
bindingNames [TypeBind -> VName
forall (f :: * -> *) vn. TypeBindBase f vn -> vn
typeAlias TypeBind
tb] (TransformM Scope -> TransformM Scope)
-> TransformM Scope -> TransformM Scope
forall a b. (a -> b) -> a -> b
$ do
        TypeBind -> TransformM ()
transformTypeBind TypeBind
tb
        [Dec] -> TransformM Scope
transformDecs [Dec]
ds'
    SigDec {} : [Dec]
ds' ->
      [Dec] -> TransformM Scope
transformDecs [Dec]
ds'
    ModDec ModBind
mb : [Dec]
ds' ->
      [VName] -> TransformM Scope -> TransformM Scope
bindingNames [ModBind -> VName
forall (f :: * -> *) vn. ModBindBase f vn -> vn
modName ModBind
mb] (TransformM Scope -> TransformM Scope)
-> TransformM Scope -> TransformM Scope
forall a b. (a -> b) -> a -> b
$ do
        Scope
mod_scope <- ModBind -> TransformM Scope
transformModBind ModBind
mb
        Scope -> TransformM Scope -> TransformM Scope
forall a. Scope -> TransformM a -> TransformM a
extendScope Scope
mod_scope (TransformM Scope -> TransformM Scope)
-> TransformM Scope -> TransformM Scope
forall a b. (a -> b) -> a -> b
$ Scope -> Scope -> Scope
forall a. Monoid a => a -> a -> a
mappend (Scope -> Scope -> Scope)
-> TransformM Scope -> TransformM (Scope -> Scope)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Dec] -> TransformM Scope
transformDecs [Dec]
ds' TransformM (Scope -> Scope) -> TransformM Scope -> TransformM Scope
forall a b. TransformM (a -> b) -> TransformM a -> TransformM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Scope -> TransformM Scope
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Scope
mod_scope
    OpenDec ModExp
e SrcLoc
_ : [Dec]
ds' -> do
      Scope
scope <- Mod -> Scope
modScope (Mod -> Scope) -> TransformM Mod -> TransformM Scope
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ModExp -> TransformM Mod
evalModExp ModExp
e
      Scope -> TransformM Scope -> TransformM Scope
forall a. Scope -> TransformM a -> TransformM a
extendScope Scope
scope (TransformM Scope -> TransformM Scope)
-> TransformM Scope -> TransformM Scope
forall a b. (a -> b) -> a -> b
$ Scope -> Scope -> Scope
forall a. Monoid a => a -> a -> a
mappend (Scope -> Scope -> Scope)
-> TransformM Scope -> TransformM (Scope -> Scope)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Dec] -> TransformM Scope
transformDecs [Dec]
ds' TransformM (Scope -> Scope) -> TransformM Scope -> TransformM Scope
forall a b. TransformM (a -> b) -> TransformM a -> TransformM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Scope -> TransformM Scope
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Scope
scope
    ImportDec [Char]
name Info ImportName
name' SrcLoc
loc : [Dec]
ds' ->
      let d :: DecBase Info vn
d = DecBase Info vn -> SrcLoc -> DecBase Info vn
forall (f :: * -> *) vn. DecBase f vn -> SrcLoc -> DecBase f vn
LocalDec (ModExpBase Info vn -> SrcLoc -> DecBase Info vn
forall (f :: * -> *) vn. ModExpBase f vn -> SrcLoc -> DecBase f vn
OpenDec ([Char] -> Info ImportName -> SrcLoc -> ModExpBase Info vn
forall (f :: * -> *) vn.
[Char] -> f ImportName -> SrcLoc -> ModExpBase f vn
ModImport [Char]
name Info ImportName
name' SrcLoc
loc) SrcLoc
loc) SrcLoc
loc
       in [Dec] -> TransformM Scope
transformDecs ([Dec] -> TransformM Scope) -> [Dec] -> TransformM Scope
forall a b. (a -> b) -> a -> b
$ Dec
forall {vn}. DecBase Info vn
d Dec -> [Dec] -> [Dec]
forall a. a -> [a] -> [a]
: [Dec]
ds'

transformImports :: Imports -> TransformM ()
transformImports :: Imports -> TransformM ()
transformImports [] = () -> TransformM ()
forall a. a -> TransformM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
transformImports ((ImportName
name, FileModule
imp) : Imports
imps) = do
  let abs :: TySet
abs = [VName] -> TySet
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> TySet) -> [VName] -> TySet
forall a b. (a -> b) -> a -> b
$ (QualName VName -> VName) -> [QualName VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf ([QualName VName] -> [VName]) -> [QualName VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ Map (QualName VName) Liftedness -> [QualName VName]
forall k a. Map k a -> [k]
M.keys (Map (QualName VName) Liftedness -> [QualName VName])
-> Map (QualName VName) Liftedness -> [QualName VName]
forall a b. (a -> b) -> a -> b
$ FileModule -> Map (QualName VName) Liftedness
fileAbs FileModule
imp
  Scope
scope <-
    (DList Dec -> DList Dec) -> TransformM Scope -> TransformM Scope
forall w (m :: * -> *) a. MonadWriter w m => (w -> w) -> m a -> m a
censor ((Dec -> Dec) -> DList Dec -> DList Dec
forall a b. (a -> b) -> DList a -> DList b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Dec -> Dec
forall {f :: * -> *} {vn}. DecBase f vn -> DecBase f vn
maybeHideEntryPoint) (TransformM Scope -> TransformM Scope)
-> TransformM Scope -> TransformM Scope
forall a b. (a -> b) -> a -> b
$
      TySet -> TransformM Scope -> TransformM Scope
forall a. TySet -> TransformM a -> TransformM a
bindingAbs TySet
abs (TransformM Scope -> TransformM Scope)
-> TransformM Scope -> TransformM Scope
forall a b. (a -> b) -> a -> b
$
        [Dec] -> TransformM Scope
transformDecs ([Dec] -> TransformM Scope) -> [Dec] -> TransformM Scope
forall a b. (a -> b) -> a -> b
$
          ProgBase Info VName -> [Dec]
forall (f :: * -> *) vn. ProgBase f vn -> [DecBase f vn]
progDecs (ProgBase Info VName -> [Dec]) -> ProgBase Info VName -> [Dec]
forall a b. (a -> b) -> a -> b
$
            FileModule -> ProgBase Info VName
fileProg FileModule
imp
  TySet -> TransformM () -> TransformM ()
forall a. TySet -> TransformM a -> TransformM a
bindingAbs TySet
abs (TransformM () -> TransformM ()) -> TransformM () -> TransformM ()
forall a b. (a -> b) -> a -> b
$ ImportName -> Scope -> TransformM () -> TransformM ()
forall a. ImportName -> Scope -> TransformM a -> TransformM a
bindingImport ImportName
name Scope
scope (TransformM () -> TransformM ()) -> TransformM () -> TransformM ()
forall a b. (a -> b) -> a -> b
$ Imports -> TransformM ()
transformImports Imports
imps
  where
    -- Only the "main" file (last import) is allowed to have entry points.
    permit_entry_points :: Bool
permit_entry_points = Imports -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Imports
imps

    maybeHideEntryPoint :: DecBase f vn -> DecBase f vn
maybeHideEntryPoint (ValDec ValBindBase f vn
vdec) =
      ValBindBase f vn -> DecBase f vn
forall (f :: * -> *) vn. ValBindBase f vn -> DecBase f vn
ValDec
        ValBindBase f vn
vdec
          { valBindEntryPoint :: Maybe (f EntryPoint)
valBindEntryPoint =
              if Bool
permit_entry_points
                then ValBindBase f vn -> Maybe (f EntryPoint)
forall (f :: * -> *) vn. ValBindBase f vn -> Maybe (f EntryPoint)
valBindEntryPoint ValBindBase f vn
vdec
                else Maybe (f EntryPoint)
forall a. Maybe a
Nothing
          }
    maybeHideEntryPoint DecBase f vn
d = DecBase f vn
d

-- | Perform defunctorisation.
transformProg :: (MonadFreshNames m) => Imports -> m [Dec]
transformProg :: forall (m :: * -> *). MonadFreshNames m => Imports -> m [Dec]
transformProg Imports
prog = (VNameSource -> ([Dec], VNameSource)) -> m [Dec]
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> ([Dec], VNameSource)) -> m [Dec])
-> (VNameSource -> ([Dec], VNameSource)) -> m [Dec]
forall a b. (a -> b) -> a -> b
$ \VNameSource
namesrc ->
  let ((), VNameSource
namesrc', DList Dec
prog') = VNameSource -> TransformM () -> ((), VNameSource, DList Dec)
forall a.
VNameSource -> TransformM a -> (a, VNameSource, DList Dec)
runTransformM VNameSource
namesrc (TransformM () -> ((), VNameSource, DList Dec))
-> TransformM () -> ((), VNameSource, DList Dec)
forall a b. (a -> b) -> a -> b
$ Imports -> TransformM ()
transformImports Imports
prog
   in (DList Dec -> [Dec]
forall a. DList a -> [a]
DL.toList DList Dec
prog', VNameSource
namesrc')