{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE Trustworthy #-}

-- | 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 qualified Data.DList as DL
import qualified Data.Map as M
import Data.Maybe
import qualified Data.Set as S
import Futhark.MonadFreshNames
import Language.Futhark
import Language.Futhark.Semantic (FileModule (..), Imports)
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 -> Map VName VName -> VName
lookupSubst VName
v Map VName VName
substs = case VName -> Map VName VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName VName
substs of
  Just VName
v' | VName
v' VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
/= VName
v -> VName -> Map VName VName -> VName
lookupSubst VName
v' Map VName VName
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 -> String
(Int -> Mod -> ShowS)
-> (Mod -> String) -> ([Mod] -> ShowS) -> Show Mod
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Mod] -> ShowS
$cshowList :: [Mod] -> ShowS
show :: Mod -> String
$cshow :: Mod -> String
showsPrec :: Int -> Mod -> ShowS
$cshowsPrec :: Int -> 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 -> Map VName VName
scopeSubsts :: Substitutions,
    Scope -> Map VName Mod
scopeMods :: M.Map VName Mod
  }
  deriving (Int -> Scope -> ShowS
[Scope] -> ShowS
Scope -> String
(Int -> Scope -> ShowS)
-> (Scope -> String) -> ([Scope] -> ShowS) -> Show Scope
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Scope] -> ShowS
$cshowList :: [Scope] -> ShowS
show :: Scope -> String
$cshow :: Scope -> String
showsPrec :: Int -> Scope -> ShowS
$cshowsPrec :: Int -> 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 Map VName VName
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 -> Map VName VName -> VName
lookupSubst VName
name Map VName VName
substs, Scope
scope)
    VName
q : [VName]
qs ->
      let q' :: VName
q' = VName -> Map VName VName -> VName
lookupSubst VName
q Map VName VName
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 Map VName VName
ss1 Map VName Mod
mt1 <> :: Scope -> Scope -> Scope
<> Scope Map VName VName
ss2 Map VName Mod
mt2 = Map VName VName -> Map VName Mod -> Scope
Scope (Map VName VName
ss1 Map VName VName -> Map VName VName -> Map VName VName
forall a. Semigroup a => a -> a -> a
<> Map VName VName
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 = Map VName VName -> Map VName Mod -> Scope
Scope Map VName VName
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 String Scope
envImports :: M.Map String 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
<* :: forall a b. TransformM a -> TransformM b -> TransformM a
$c<* :: forall a b. TransformM a -> TransformM b -> TransformM a
*> :: forall a b. TransformM a -> TransformM b -> TransformM b
$c*> :: forall a b. TransformM a -> TransformM b -> TransformM b
liftA2 :: forall a b c.
(a -> b -> c) -> TransformM a -> TransformM b -> TransformM c
$cliftA2 :: forall a b c.
(a -> b -> c) -> TransformM a -> TransformM b -> TransformM c
<*> :: forall a b. TransformM (a -> b) -> TransformM a -> TransformM b
$c<*> :: forall a b. TransformM (a -> b) -> TransformM a -> TransformM b
pure :: forall a. a -> TransformM a
$cpure :: forall a. a -> 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
<$ :: forall a b. a -> TransformM b -> TransformM a
$c<$ :: forall a b. a -> TransformM b -> TransformM a
fmap :: forall a b. (a -> b) -> TransformM a -> TransformM b
$cfmap :: forall a b. (a -> b) -> TransformM a -> TransformM b
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
return :: forall a. a -> TransformM a
$creturn :: forall a. a -> TransformM a
>> :: forall a b. TransformM a -> TransformM b -> TransformM b
$c>> :: forall a b. TransformM a -> TransformM b -> TransformM b
>>= :: forall a b. TransformM a -> (a -> TransformM b) -> TransformM b
$c>>= :: forall a b. TransformM a -> (a -> TransformM b) -> TransformM b
Monad,
      Monad TransformM
Applicative TransformM
TransformM VNameSource
Applicative TransformM
-> Monad TransformM
-> TransformM VNameSource
-> (VNameSource -> TransformM ())
-> MonadFreshNames TransformM
VNameSource -> TransformM ()
forall (m :: * -> *).
Applicative m
-> Monad m
-> m VNameSource
-> (VNameSource -> m ())
-> MonadFreshNames m
putNameSource :: VNameSource -> TransformM ()
$cputNameSource :: VNameSource -> TransformM ()
getNameSource :: TransformM VNameSource
$cgetNameSource :: TransformM VNameSource
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 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 Map VName VName
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 :: Map VName VName
scopeSubsts = (VName -> VName) -> Map VName VName -> Map VName VName
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Map VName VName -> VName -> VName
forall {k}. Ord k => Map k k -> k -> k
forward (Scope -> Map VName VName
scopeSubsts Scope
scope)) Map VName VName
substs Map VName VName -> Map VName VName -> Map VName VName
forall a. Semigroup a => a -> a -> a
<> Scope -> Map VName VName
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. Map VName VName -> TransformM a -> TransformM a
substituting Map VName VName
substs = Scope -> TransformM a -> TransformM a
forall a. Scope -> TransformM a -> TransformM a
extendScope Scope
forall a. Monoid a => a
mempty {scopeSubsts :: Map VName VName
scopeSubsts = Map VName VName
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 (m :: * -> *) a. Monad m => a -> m a
return 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)
mapM VName -> TransformM VName
boundName [VName]
names
  let substs :: Map VName VName
substs = [(VName, VName)] -> Map VName VName
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')
  Map VName VName -> TransformM Scope -> TransformM Scope
forall a. Map VName VName -> TransformM a -> TransformM a
substituting Map VName VName
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 (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Scope -> TransformM Scope
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Map VName VName -> Map VName Mod -> Scope
Scope Map VName VName
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 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 :: String -> Scope -> TransformM a -> TransformM a
bindingImport :: forall a. String -> Scope -> TransformM a -> TransformM a
bindingImport String
name Scope
scope = (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 String Scope
envImports = String -> Scope -> Map String Scope -> Map String Scope
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert String
name Scope
scope (Map String Scope -> Map String Scope)
-> Map String Scope -> Map String Scope
forall a b. (a -> b) -> a -> b
$ Env -> Map String 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 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 :: String -> TransformM Scope
lookupImport :: String -> TransformM Scope
lookupImport String
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 (m :: * -> *) a. Monad m => a -> m a
return (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 (String -> Map String Scope -> Maybe Scope
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup String
name (Map String Scope -> Maybe Scope)
-> (Env -> Map String Scope) -> Env -> Maybe Scope
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map String Scope
envImports)
  where
    bad :: a
bad = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ String
"Unknown import: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
name

lookupMod' :: QualName VName -> Scope -> Either String Mod
lookupMod' :: QualName VName -> Scope -> Either String 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 String Mod
-> (Mod -> Either String Mod) -> Maybe Mod -> Either String Mod
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (String -> Either String Mod
forall a b. a -> Either a b
Left (String -> Either String Mod) -> String -> Either String Mod
forall a b. (a -> b) -> a -> b
$ QualName VName -> String
forall {a}. Pretty a => a -> String
bad QualName VName
mname') Mod -> Either String Mod
forall a b. b -> Either a b
Right (Maybe Mod -> Either String Mod) -> Maybe Mod -> Either String 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 -> String
bad a
mname' = String
"Unknown module: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ QualName VName -> String
forall {a}. Pretty a => a -> String
pretty QualName VName
mname String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall {a}. Pretty a => a -> String
pretty a
mname' String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

lookupMod :: QualName VName -> TransformM Mod
lookupMod :: QualName VName -> TransformM Mod
lookupMod QualName VName
mname = (String -> TransformM Mod)
-> (Mod -> TransformM Mod) -> Either String Mod -> TransformM Mod
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either String -> TransformM Mod
forall a. HasCallStack => String -> a
error Mod -> TransformM Mod
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String Mod -> TransformM Mod)
-> (Scope -> Either String Mod) -> Scope -> TransformM Mod
forall b c a. (b -> c) -> (a -> b) -> a -> c
. QualName VName -> Scope -> Either String 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 String Scope -> TySet -> Env
Env Scope
forall a. Monoid a => a
mempty Bool
False Map String 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 (Map VName VName)) -> ModExp -> ModExp
maybeAscript SrcLoc
loc (Just (SigExp
mtye, Info (Map VName VName)
substs)) ModExp
me = ModExp -> SigExp -> Info (Map VName VName) -> SrcLoc -> ModExp
forall (f :: * -> *) vn.
ModExpBase f vn
-> SigExpBase f vn
-> f (Map VName VName)
-> SrcLoc
-> ModExpBase f vn
ModAscript ModExp
me SigExp
mtye Info (Map VName VName)
substs SrcLoc
loc
maybeAscript SrcLoc
_ Maybe (SigExp, Info (Map VName VName))
Nothing ModExp
me = ModExp
me

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

extendAbsTypes :: Substitutions -> TransformM a -> TransformM a
extendAbsTypes :: forall a. Map VName VName -> TransformM a -> TransformM a
extendAbsTypes Map VName VName
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] -> TySet
forall a b. (a -> b) -> a -> b
$
          ((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]
forall a b. (a -> b) -> a -> b
$
            ((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)] -> [(VName, VName)])
-> [(VName, VName)] -> [(VName, VName)]
forall a b. (a -> b) -> a -> b
$
              Map VName VName -> [(VName, VName)]
forall k a. Map k a -> [(k, a)]
M.toList Map VName VName
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 String
_ (Info String
fpath) SrcLoc
_) = Scope -> Mod
ModMod (Scope -> Mod) -> TransformM Scope -> TransformM Mod
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> TransformM Scope
lookupImport String
fpath
evalModExp (ModAscript ModExp
me SigExp
_ (Info Map VName VName
ascript_substs) SrcLoc
_) =
  Map VName VName -> TransformM Mod -> TransformM Mod
forall a. Map VName VName -> TransformM a -> TransformM a
extendAbsTypes Map VName VName
ascript_substs (TransformM Mod -> TransformM Mod)
-> TransformM Mod -> TransformM Mod
forall a b. (a -> b) -> a -> b
$
    Map VName VName -> Mod -> Mod
substituteInMod Map VName VName
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 Map VName VName
p_substs) (Info Map VName VName
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
_ ->
      String -> TransformM Mod
forall a. HasCallStack => String -> a
error (String -> TransformM Mod) -> String -> TransformM Mod
forall a b. (a -> b) -> a -> b
$ String
"Cannot apply non-parametric module at " String -> ShowS
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
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
. Map VName VName -> TransformM Mod -> TransformM Mod
forall a. Map VName VName -> TransformM a -> TransformM a
extendAbsTypes Map VName VName
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 -> TransformM Mod -> TransformM Mod
forall a. Scope -> TransformM a -> TransformM a
extendScope Scope
f_closure
        (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
          Map VName VName
outer_substs <- Scope -> Map VName VName
scopeSubsts (Scope -> Map VName VName)
-> TransformM Scope -> TransformM (Map VName VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TransformM Scope
askScope
          TySet
abs <- (Env -> TySet) -> TransformM TySet
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> TySet
envAbs
          let forward :: (VName, b) -> (VName, b)
forward (VName
k, b
v) = (VName -> Map VName VName -> VName
lookupSubst VName
k Map VName VName
outer_substs, b
v)
              p_substs' :: Map VName VName
p_substs' = [(VName, VName)] -> Map VName VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, VName)] -> Map VName VName)
-> [(VName, VName)] -> Map VName VName
forall a b. (a -> b) -> a -> b
$ ((VName, VName) -> (VName, VName))
-> [(VName, VName)] -> [(VName, VName)]
forall a b. (a -> b) -> [a] -> [b]
map (VName, VName) -> (VName, VName)
forall {b}. (VName, b) -> (VName, b)
forward ([(VName, VName)] -> [(VName, VName)])
-> [(VName, VName)] -> [(VName, VName)]
forall a b. (a -> b) -> a -> b
$ Map VName VName -> [(VName, VName)]
forall k a. Map k a -> [(k, a)]
M.toList Map VName VName
p_substs
              keep :: VName -> p -> Bool
keep VName
k p
_ =
                VName
k VName -> Map VName VName -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map VName VName
p_substs'
                  Bool -> Bool -> Bool
|| VName
k VName -> TySet -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` TySet
abs
              abs_substs :: Map VName VName
abs_substs =
                (VName -> VName -> Bool) -> Map VName VName -> Map VName VName
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey VName -> VName -> Bool
forall {p}. VName -> p -> Bool
keep (Map VName VName -> Map VName VName)
-> Map VName VName -> Map VName VName
forall a b. (a -> b) -> a -> b
$
                  (VName -> VName) -> Map VName VName -> Map VName VName
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (VName -> Map VName VName -> VName
`lookupSubst` Scope -> Map VName VName
scopeSubsts (Mod -> Scope
modScope Mod
arg_mod)) Map VName VName
p_substs'
                    Map VName VName -> Map VName VName -> Map VName VName
forall a. Semigroup a => a -> a -> a
<> Scope -> Map VName VName
scopeSubsts Scope
f_closure
                    Map VName VName -> Map VName VName -> Map VName VName
forall a. Semigroup a => a -> a -> a
<> Scope -> Map VName VName
scopeSubsts (Mod -> Scope
modScope Mod
arg_mod)
          Scope -> TransformM Mod -> TransformM Mod
forall a. Scope -> TransformM a -> TransformM a
extendScope
            ( Map VName VName -> Map VName Mod -> Scope
Scope
                Map VName VName
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
$
                    Map VName VName -> Mod -> Mod
substituteInMod Map VName VName
p_substs' Mod
arg_mod
                )
            )
            (TransformM Mod -> TransformM Mod)
-> TransformM Mod -> TransformM Mod
forall a b. (a -> b) -> a -> b
$ do
              Map VName VName
substs <- Scope -> Map VName VName
scopeSubsts (Scope -> Map VName VName)
-> TransformM Scope -> TransformM (Map VName VName)
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 (m :: * -> *) a. Monad m => a -> m a
return (Mod -> TransformM Mod) -> Mod -> TransformM Mod
forall a b. (a -> b) -> a -> b
$
                TySet -> Map VName VName -> Mod -> Mod
addSubsts TySet
abs Map VName VName
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).
                  Map VName VName -> Mod -> Mod
addSubstsModMod (Scope -> Map VName VName
scopeSubsts (Scope -> Map VName VName) -> Scope -> Map VName VName
forall a b. (a -> b) -> a -> b
$ Mod -> Scope
modScope Mod
arg_mod) (Mod -> Mod) -> Mod -> Mod
forall a b. (a -> b) -> a -> b
$
                    Map VName VName -> Mod -> Mod
substituteInMod (Map VName VName
b_substs Map VName VName -> Map VName VName -> Map VName VName
forall a. Semigroup a => a -> a -> a
<> Map VName VName
substs) Mod
x
  where
    addSubsts :: TySet -> Map VName VName -> Mod -> Mod
addSubsts TySet
abs Map VName VName
substs (ModFun TySet
mabs (Scope Map VName VName
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) (Map VName VName -> Map VName Mod -> Scope
Scope (Map VName VName
substs Map VName VName -> Map VName VName -> Map VName VName
forall a. Semigroup a => a -> a -> a
<> Map VName VName
msubsts) Map VName Mod
mods) ModParam
mp ModExp
me
    addSubsts TySet
_ Map VName VName
substs (ModMod (Scope Map VName VName
msubsts Map VName Mod
mods)) =
      Scope -> Mod
ModMod (Scope -> Mod) -> Scope -> Mod
forall a b. (a -> b) -> a -> b
$ Map VName VName -> Map VName Mod -> Scope
Scope (Map VName VName
substs Map VName VName -> Map VName VName -> Map VName VName
forall a. Semigroup a => a -> a -> a
<> Map VName VName
msubsts) Map VName Mod
mods
    addSubstsModMod :: Map VName VName -> Mod -> Mod
addSubstsModMod Map VName VName
substs (ModMod (Scope Map VName VName
msubsts Map VName Mod
mods)) =
      Scope -> Mod
ModMod (Scope -> Mod) -> Scope -> Mod
forall a b. (a -> b) -> a -> b
$ Map VName VName -> Map VName Mod -> Scope
Scope (Map VName VName
substs Map VName VName -> Map VName VName -> Map VName VName
forall a. Semigroup a => a -> a -> a
<> Map VName VName
msubsts) Map VName Mod
mods
    addSubstsModMod Map VName VName
_ Mod
m = Mod
m
evalModExp (ModLambda ModParam
p Maybe (SigExp, Info (Map VName VName))
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 (m :: * -> *) a. Monad m => a -> m a
return (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 (Map VName VName)) -> ModExp -> ModExp
maybeAscript SrcLoc
loc Maybe (SigExp, Info (Map VName VName))
ascript ModExp
e

transformName :: VName -> TransformM VName
transformName :: VName -> TransformM VName
transformName VName
v = VName -> Map VName VName -> VName
lookupSubst VName
v (Map VName VName -> VName)
-> (Scope -> Map VName VName) -> Scope -> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope -> Map VName VName
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 (m :: * -> *) a. Monad m => a -> m a
return (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
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 :: forall (m :: * -> *).
(Exp -> m Exp)
-> (VName -> m VName)
-> (QualName VName -> m (QualName VName))
-> (StructType -> m StructType)
-> (PatternType -> m PatternType)
-> ASTMapper m
ASTMapper
        { mapOnExp :: Exp -> m Exp
mapOnExp = Scope -> Exp -> m Exp
onExp Scope
scope,
          mapOnName :: VName -> m VName
mapOnName = \VName
v ->
            VName -> m VName
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> m VName) -> VName -> m VName
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf (QualName VName -> VName) -> QualName VName -> 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 (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) Scope
scope,
          mapOnQualName :: QualName VName -> m (QualName VName)
mapOnQualName = \QualName VName
v ->
            QualName VName -> m (QualName VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (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 Scope
scope,
          mapOnStructType :: StructType -> m StructType
mapOnStructType = ASTMapper m -> StructType -> m StructType
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (Scope -> ASTMapper m
substituter Scope
scope),
          mapOnPatternType :: PatternType -> m PatternType
mapOnPatternType = ASTMapper m -> PatternType -> m PatternType
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
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 String Mod
lookupMod' QualName VName
mn Scope
scope of
            Left String
err -> String -> m Exp
forall a. HasCallStack => String -> a
error String
err
            Right Mod
mod ->
              ASTMapper m -> Exp -> m Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
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
astMap (Scope -> ASTMapper m
substituter Scope
scope) Exp
e

transformTypeExp :: TypeExp VName -> TransformM (TypeExp VName)
transformTypeExp :: TypeExp VName -> TransformM (TypeExp VName)
transformTypeExp = TypeExp VName -> TransformM (TypeExp 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

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

transformValBind :: ValBind -> TransformM ()
transformValBind :: ValBind -> TransformM ()
transformValBind (ValBind Maybe (Info EntryPoint)
entry VName
name Maybe (TypeExp VName)
tdecl (Info (StructType
t, [VName]
retext)) [TypeParamBase VName]
tparams [PatternBase Info VName]
params Exp
e Maybe DocComment
doc [AttrInfo]
attrs SrcLoc
loc) = do
  VName
name' <- VName -> TransformM VName
transformName VName
name
  Maybe (TypeExp VName)
tdecl' <- (TypeExp VName -> TransformM (TypeExp VName))
-> Maybe (TypeExp VName) -> TransformM (Maybe (TypeExp VName))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse TypeExp VName -> TransformM (TypeExp VName)
transformTypeExp Maybe (TypeExp VName)
tdecl
  StructType
t' <- StructType -> TransformM StructType
transformStructType StructType
t
  Exp
e' <- Exp -> TransformM Exp
transformExp Exp
e
  [TypeParamBase VName]
tparams' <- (TypeParamBase VName -> TransformM (TypeParamBase VName))
-> [TypeParamBase VName] -> TransformM [TypeParamBase VName]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse TypeParamBase VName -> TransformM (TypeParamBase VName)
forall x. ASTMappable x => x -> TransformM x
transformNames [TypeParamBase VName]
tparams
  [PatternBase Info VName]
params' <- (PatternBase Info VName -> TransformM (PatternBase Info VName))
-> [PatternBase Info VName] -> TransformM [PatternBase Info VName]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatternBase Info VName -> TransformM (PatternBase Info VName)
forall x. ASTMappable x => x -> TransformM x
transformNames [PatternBase Info VName]
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 VName)
-> Info (StructType, [VName])
-> [TypeParamBase VName]
-> [PatternBase Info VName]
-> Exp
-> Maybe DocComment
-> [AttrInfo]
-> SrcLoc
-> ValBind
forall (f :: * -> *) vn.
Maybe (f EntryPoint)
-> vn
-> Maybe (TypeExp vn)
-> f (StructType, [VName])
-> [TypeParamBase vn]
-> [PatternBase f vn]
-> ExpBase f vn
-> Maybe DocComment
-> [AttrInfo]
-> SrcLoc
-> ValBindBase f vn
ValBind Maybe (Info EntryPoint)
entry VName
name' Maybe (TypeExp VName)
tdecl' ((StructType, [VName]) -> Info (StructType, [VName])
forall a. a -> Info a
Info (StructType
t', [VName]
retext)) [TypeParamBase VName]
tparams' [PatternBase Info VName]
params' Exp
e' Maybe DocComment
doc [AttrInfo]
attrs SrcLoc
loc

transformTypeDecl :: TypeDecl -> TransformM TypeDecl
transformTypeDecl :: TypeDecl -> TransformM TypeDecl
transformTypeDecl (TypeDecl TypeExp VName
dt (Info StructType
et)) =
  TypeExp VName -> Info StructType -> TypeDecl
forall (f :: * -> *) vn.
TypeExp vn -> f StructType -> TypeDeclBase f vn
TypeDecl (TypeExp VName -> Info StructType -> TypeDecl)
-> TransformM (TypeExp VName)
-> TransformM (Info StructType -> TypeDecl)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeExp VName -> TransformM (TypeExp VName)
transformTypeExp TypeExp VName
dt TransformM (Info StructType -> TypeDecl)
-> TransformM (Info StructType) -> TransformM TypeDecl
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType)
-> TransformM StructType -> TransformM (Info StructType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructType -> TransformM StructType
transformStructType StructType
et)

transformTypeBind :: TypeBind -> TransformM ()
transformTypeBind :: TypeBind -> TransformM ()
transformTypeBind (TypeBind VName
name Liftedness
l [TypeParamBase VName]
tparams TypeDecl
te 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]
-> TypeDecl
-> Maybe DocComment
-> SrcLoc
-> TypeBind
forall (f :: * -> *) vn.
vn
-> Liftedness
-> [TypeParamBase vn]
-> TypeDeclBase f vn
-> Maybe DocComment
-> SrcLoc
-> TypeBindBase f vn
TypeBind VName
name' Liftedness
l ([TypeParamBase VName]
 -> TypeDecl -> Maybe DocComment -> SrcLoc -> TypeBind)
-> TransformM [TypeParamBase VName]
-> TransformM (TypeDecl -> Maybe DocComment -> SrcLoc -> TypeBind)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeParamBase VName -> TransformM (TypeParamBase VName))
-> [TypeParamBase VName] -> TransformM [TypeParamBase VName]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse TypeParamBase VName -> TransformM (TypeParamBase VName)
forall x. ASTMappable x => x -> TransformM x
transformNames [TypeParamBase VName]
tparams
            TransformM (TypeDecl -> Maybe DocComment -> SrcLoc -> TypeBind)
-> TransformM TypeDecl
-> TransformM (Maybe DocComment -> SrcLoc -> TypeBind)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeDecl -> TransformM TypeDecl
transformTypeDecl TypeDecl
te
            TransformM (Maybe DocComment -> SrcLoc -> TypeBind)
-> TransformM (Maybe DocComment) -> TransformM (SrcLoc -> TypeBind)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe DocComment -> TransformM (Maybe DocComment)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe DocComment
doc
            TransformM (SrcLoc -> TypeBind)
-> TransformM SrcLoc -> TransformM TypeBind
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> TransformM SrcLoc
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 (Map VName VName))
-> ModExpBase f vn
-> SrcLoc
-> ModExpBase f vn
forall (f :: * -> *) vn.
ModParamBase f vn
-> Maybe (SigExpBase f vn, f (Map VName VName))
-> ModExpBase f vn
-> SrcLoc
-> ModExpBase f vn
ModLambda ModParamBase f vn
p Maybe (SigExpBase f vn, f (Map VName VName))
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 (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 (Map VName VName)) -> ModExp -> ModExp
maybeAscript (ModBind -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf ModBind
mb) (ModBind -> Maybe (SigExp, Info (Map VName VName))
forall (f :: * -> *) vn.
ModBindBase f vn -> Maybe (SigExpBase f vn, f (Map VName VName))
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
  TySet
abs <- (Env -> TySet) -> TransformM TySet
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> TySet
envAbs
  -- Copy substitutions involving abstract types out, because they are
  -- always resolved at the outermost level.
  let abs_substs :: Map VName VName
abs_substs =
        (VName -> VName -> Bool) -> Map VName VName -> Map VName VName
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (Bool -> VName -> Bool
forall a b. a -> b -> a
const (Bool -> VName -> Bool)
-> (VName -> Bool) -> VName -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> TySet -> Bool) -> TySet -> VName -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName -> TySet -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member TySet
abs) (Map VName VName -> Map VName VName)
-> Map VName VName -> Map VName VName
forall a b. (a -> b) -> a -> b
$
          Scope -> Map VName VName
scopeSubsts (Scope -> Map VName VName) -> Scope -> Map VName VName
forall a b. (a -> b) -> a -> b
$ Mod -> Scope
modScope Mod
mod
  Scope -> TransformM Scope
forall (m :: * -> *) a. Monad m => a -> m a
return (Scope -> TransformM Scope) -> Scope -> TransformM Scope
forall a b. (a -> b) -> a -> b
$ Map VName VName -> Map VName Mod -> Scope
Scope Map VName VName
abs_substs (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 (m :: * -> *) a. Monad m => a -> m a
return 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 (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Scope -> TransformM Scope
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 (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Scope -> TransformM Scope
forall (f :: * -> *) a. Applicative f => a -> f a
pure Scope
scope
    ImportDec String
name Info String
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 (String -> Info String -> SrcLoc -> ModExpBase Info vn
forall (f :: * -> *) vn.
String -> f String -> SrcLoc -> ModExpBase f vn
ModImport String
name Info String
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 (m :: * -> *) a. Monad m => a -> m a
return ()
transformImports ((String
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 (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
$ String -> Scope -> TransformM () -> TransformM ()
forall a. String -> Scope -> TransformM a -> TransformM a
bindingImport String
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 (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')