<commit_before>{-# 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 v substs = case M.lookup v substs of
  Just v' | v' /= v -> lookupSubst v' substs
  _ -> 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 (Show)

modScope :: Mod -> Scope
modScope (ModMod scope) = scope
modScope ModFun {} = mempty

data Scope = Scope
  { scopeSubsts :: Substitutions,
    scopeMods :: M.Map VName Mod
  }
  deriving (Show)

lookupSubstInScope :: QualName VName -> Scope -> (QualName VName, Scope)
lookupSubstInScope qn@(QualName quals name) scope@(Scope substs mods) =
  case quals of
    [] -> (qualName $ lookupSubst name substs, scope)
    q : qs ->
      let q' = lookupSubst q substs
       in case M.lookup q' mods of
            Just (ModMod mod_scope) -> lookupSubstInScope (QualName qs name) mod_scope
            _ -> (qn, scope)

instance Semigroup Scope where
  Scope ss1 mt1 <> Scope ss2 mt2 = Scope (ss1 <> ss2) (mt1 <> mt2)

instance Monoid Scope where
  mempty = Scope mempty mempty

type TySet = S.Set VName

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

newtype TransformM a = TransformM (RWS Env (DL.DList Dec) VNameSource a)
  deriving
    ( Applicative,
      Functor,
      Monad,
      MonadFreshNames,
      MonadReader Env,
      MonadWriter (DL.DList Dec)
    )

emit :: Dec -> TransformM ()
emit = tell . DL.singleton

askScope :: TransformM Scope
askScope = asks envScope

localScope :: (Scope -> Scope) -> TransformM a -> TransformM a
localScope f = local $ \env -> env {envScope = f $ envScope env}

extendScope :: Scope -> TransformM a -> TransformM a
extendScope (Scope substs mods) = localScope $ \scope ->
  scope
    { scopeSubsts = M.map (forward (scopeSubsts scope)) substs <> scopeSubsts scope,
      scopeMods = mods <> scopeMods scope
    }
  where
    forward old_substs v = fromMaybe v $ M.lookup v old_substs

substituting :: Substitutions -> TransformM a -> TransformM a
substituting substs = extendScope mempty {scopeSubsts = substs}

boundName :: VName -> TransformM VName
boundName v = do
  g <- asks envGenerating
  if g then newName v else return v

bindingNames :: [VName] -> TransformM Scope -> TransformM Scope
bindingNames names m = do
  names' <- mapM boundName names
  let substs = M.fromList (zip names names')
  substituting substs $ mappend <$> m <*> pure (Scope substs mempty)

generating :: TransformM a -> TransformM a
generating = local $ \env -> env {envGenerating = True}

bindingImport :: String -> Scope -> TransformM a -> TransformM a
bindingImport name scope = local $ \env ->
  env {envImports = M.insert name scope $ envImports env}

bindingAbs :: TySet -> TransformM a -> TransformM a
bindingAbs abs = local $ \env ->
  env {envAbs = abs <> envAbs env}

lookupImport :: String -> TransformM Scope
lookupImport name = maybe bad return =<< asks (M.lookup name . envImports)
  where
    bad = error $ "Unknown import: " ++ name

lookupMod' :: QualName VName -> Scope -> Either String Mod
lookupMod' mname scope =
  let (mname', scope') = lookupSubstInScope mname scope
   in maybe (Left $ bad mname') Right $ M.lookup (qualLeaf mname') $ scopeMods scope'
  where
    bad mname' = "Unknown module: " ++ pretty mname ++ " (" ++ pretty mname' ++ ")"

lookupMod :: QualName VName -> TransformM Mod
lookupMod mname = either error return . lookupMod' mname =<< askScope

runTransformM :: VNameSource -> TransformM a -> (a, VNameSource, DL.DList Dec)
runTransformM src (TransformM m) = runRWS m env src
  where
    env = Env mempty False mempty mempty

maybeAscript ::
  SrcLoc ->
  Maybe (SigExp, Info (M.Map VName VName)) ->
  ModExp ->
  ModExp
maybeAscript loc (Just (mtye, substs)) me = ModAscript me mtye substs loc
maybeAscript _ Nothing me = me

substituteInMod :: Substitutions -> Mod -> Mod
substituteInMod substs (ModMod (Scope mod_substs mod_mods)) =
  -- Forward all substitutions.
  ModMod $ Scope substs' $ M.map (substituteInMod substs) mod_mods
  where
    forward v = lookupSubst v $ mod_substs <> substs
    substs' = M.map forward substs
substituteInMod substs (ModFun abs (Scope mod_substs mod_mods) mparam mbody) =
  ModFun abs (Scope (substs' <> mod_substs) mod_mods) mparam mbody
  where
    forward v = lookupSubst v mod_substs
    substs' = M.map forward substs

extendAbsTypes :: Substitutions -> TransformM a -> TransformM a
extendAbsTypes ascript_substs m = do
  abs <- asks 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 =
        S.fromList $
          map snd $
            filter ((`S.member` abs) . fst) $
              M.toList ascript_substs
  bindingAbs subst_abs m

evalModExp :: ModExp -> TransformM Mod
evalModExp (ModVar qn _) = lookupMod qn
evalModExp (ModParens e _) = evalModExp e
evalModExp (ModDecs decs _) = ModMod <$> transformDecs decs
evalModExp (ModImport _ (Info fpath) _) = ModMod <$> lookupImport fpath
evalModExp (ModAscript me _ (Info ascript_substs) _) =
  extendAbsTypes ascript_substs $
    substituteInMod ascript_substs <$> evalModExp me
evalModExp (ModApply f arg (Info p_substs) (Info b_substs) loc) = do
  f_mod <- evalModExp f<commit_msg>Fewer linebreaks.<commit_after>{-# 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 v substs = case M.lookup v substs of
  Just v' | v' /= v -> lookupSubst v' substs
  _ -> 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 (Show)

modScope :: Mod -> Scope
modScope (ModMod scope) = scope
modScope ModFun {} = mempty

data Scope = Scope
  { scopeSubsts :: Substitutions,
    scopeMods :: M.Map VName Mod
  }
  deriving (Show)

lookupSubstInScope :: QualName VName -> Scope -> (QualName VName, Scope)
lookupSubstInScope qn@(QualName quals name) scope@(Scope substs mods) =
  case quals of
    [] -> (qualName $ lookupSubst name substs, scope)
    q : qs ->
      let q' = lookupSubst q substs
       in case M.lookup q' mods of
            Just (ModMod mod_scope) -> lookupSubstInScope (QualName qs name) mod_scope
            _ -> (qn, scope)

instance Semigroup Scope where
  Scope ss1 mt1 <> Scope ss2 mt2 = Scope (ss1 <> ss2) (mt1 <> mt2)

instance Monoid Scope where
  mempty = Scope mempty mempty

type TySet = S.Set VName

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

newtype TransformM a = TransformM (RWS Env (DL.DList Dec) VNameSource a)
  deriving
    ( Applicative,
      Functor,
      Monad,
      MonadFreshNames,
      MonadReader Env,
      MonadWriter (DL.DList Dec)
    )

emit :: Dec -> TransformM ()
emit = tell . DL.singleton

askScope :: TransformM Scope
askScope = asks envScope

localScope :: (Scope -> Scope) -> TransformM a -> TransformM a
localScope f = local $ \env -> env {envScope = f $ envScope env}

extendScope :: Scope -> TransformM a -> TransformM a
extendScope (Scope substs mods) = localScope $ \scope ->
  scope
    { scopeSubsts = M.map (forward (scopeSubsts scope)) substs <> scopeSubsts scope,
      scopeMods = mods <> scopeMods scope
    }
  where
    forward old_substs v = fromMaybe v $ M.lookup v old_substs

substituting :: Substitutions -> TransformM a -> TransformM a
substituting substs = extendScope mempty {scopeSubsts = substs}

boundName :: VName -> TransformM VName
boundName v = do
  g <- asks envGenerating
  if g then newName v else return v

bindingNames :: [VName] -> TransformM Scope -> TransformM Scope
bindingNames names m = do
  names' <- mapM boundName names
  let substs = M.fromList (zip names names')
  substituting substs $ mappend <$> m <*> pure (Scope substs mempty)

generating :: TransformM a -> TransformM a
generating = local $ \env -> env {envGenerating = True}

bindingImport :: String -> Scope -> TransformM a -> TransformM a
bindingImport name scope = local $ \env ->
  env {envImports = M.insert name scope $ envImports env}

bindingAbs :: TySet -> TransformM a -> TransformM a
bindingAbs abs = local $ \env ->
  env {envAbs = abs <> envAbs env}

lookupImport :: String -> TransformM Scope
lookupImport name = maybe bad return =<< asks (M.lookup name . envImports)
  where
    bad = error $ "Unknown import: " ++ name

lookupMod' :: QualName VName -> Scope -> Either String Mod
lookupMod' mname scope =
  let (mname', scope') = lookupSubstInScope mname scope
   in maybe (Left $ bad mname') Right $ M.lookup (qualLeaf mname') $ scopeMods scope'
  where
    bad mname' = "Unknown module: " ++ pretty mname ++ " (" ++ pretty mname' ++ ")"

lookupMod :: QualName VName -> TransformM Mod
lookupMod mname = either error return . lookupMod' mname =<< askScope

runTransformM :: VNameSource -> TransformM a -> (a, VNameSource, DL.DList Dec)
runTransformM src (TransformM m) = runRWS m env src
  where
    env = Env mempty False mempty mempty

maybeAscript ::
  SrcLoc ->
  Maybe (SigExp, Info (M.Map VName VName)) ->
  ModExp ->
  ModExp
maybeAscript loc (Just (mtye, substs)) me = ModAscript me mtye substs loc
maybeAscript _ Nothing me = me

substituteInMod :: Substitutions -> Mod -> Mod
substituteInMod substs (ModMod (Scope mod_substs mod_mods)) =
  -- Forward all substitutions.
  ModMod $ Scope substs' $ M.map (substituteInMod substs) mod_mods
  where
    forward v = lookupSubst v $ mod_substs <> substs
    substs' = M.map forward substs
substituteInMod substs (ModFun abs (Scope mod_substs mod_mods) mparam mbody) =
  ModFun abs (Scope (substs' <> mod_substs) mod_mods) mparam mbody
  where
    forward v = lookupSubst v mod_substs
    substs' = M.map forward substs

extendAbsTypes :: Substitutions -> TransformM a -> TransformM a
extendAbsTypes ascript_substs m = do
  abs <- asks 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 =
        S.fromList . map snd . filter ((`S.member` abs) . fst) $
          M.toList ascript_substs
  bindingAbs subst_abs m

evalModExp :: ModExp -> TransformM Mod
evalModExp (ModVar qn _) = lookupMod qn
evalModExp (ModParens e _) = evalModExp e
evalModExp (ModDecs decs _) = ModMod <$> transformDecs decs
evalModExp (ModImport _ (Info fpath) _) = ModMod <$> lookupImport fpath
evalModExp (ModAscript me _ (Info ascript_substs) _) =
  extendAbsTypes ascript_substs $
    substituteInMod ascript_substs <$> evalModExp me
evalModExp (ModApply f arg (Info p_substs) (Info b_substs) loc) = do
  f_mod <- evalModExp f
{'fraction_non_alphanumeric': 0.09348281980964672, 'fraction_numerical': 0.0012905307307630263, 'mean_word_length': 3.5103674063295744, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 123, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5923036', 'n_tokens_mistral': 4118, 'n_tokens_neox': 3998, 'n_words': 1394}