{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-redundant-constraints #-}

-- | A representation where all patterns are annotated with aliasing
-- information.  It also records consumption of variables in bodies.
--
-- Note that this module is mostly not concerned with actually
-- /computing/ the aliasing information; only with shuffling it around
-- and providing some basic building blocks.  See modules such as
-- "Futhark.Analysis.Alias" for computing the aliases in the first
-- place.
module Futhark.IR.Aliases
  ( -- * The representation definition
    Aliases,
    AliasDec (..),
    VarAliases,
    ConsumedInExp,
    BodyAliasing,
    module Futhark.IR.Prop.Aliases,

    -- * Module re-exports
    module Futhark.IR.Prop,
    module Futhark.IR.Traversals,
    module Futhark.IR.Pretty,
    module Futhark.IR.Syntax,

    -- * Adding aliases
    mkAliasedBody,
    mkAliasedPat,
    mkBodyAliasing,
    CanBeAliased (..),
    AliasableRep,

    -- * Removing aliases
    removeProgAliases,
    removeFunDefAliases,
    removeExpAliases,
    removeStmAliases,
    removeBodyAliases,
    removeLambdaAliases,
    removePatAliases,
    removeScopeAliases,

    -- * Tracking aliases
    AliasesAndConsumed,
    trackAliases,
    mkStmsAliases,
    consumedInStms,
  )
where

import Control.Monad.Identity
import Control.Monad.Reader
import Data.Kind qualified
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.Builder
import Futhark.IR.Pretty
import Futhark.IR.Prop
import Futhark.IR.Prop.Aliases
import Futhark.IR.Syntax
import Futhark.IR.Traversals
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Util.Pretty qualified as PP

-- | The rep for the basic representation.
data Aliases (rep :: Data.Kind.Type)

-- | A wrapper around 'AliasDec' to get around the fact that we need an
-- 'Ord' instance, which 'AliasDec does not have.
newtype AliasDec = AliasDec {AliasDec -> Names
unAliases :: Names}
  deriving (Int -> AliasDec -> ShowS
[AliasDec] -> ShowS
AliasDec -> String
(Int -> AliasDec -> ShowS)
-> (AliasDec -> String) -> ([AliasDec] -> ShowS) -> Show AliasDec
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> AliasDec -> ShowS
showsPrec :: Int -> AliasDec -> ShowS
$cshow :: AliasDec -> String
show :: AliasDec -> String
$cshowList :: [AliasDec] -> ShowS
showList :: [AliasDec] -> ShowS
Show)

instance Semigroup AliasDec where
  AliasDec
x <> :: AliasDec -> AliasDec -> AliasDec
<> AliasDec
y = Names -> AliasDec
AliasDec (Names -> AliasDec) -> Names -> AliasDec
forall a b. (a -> b) -> a -> b
$ AliasDec -> Names
unAliases AliasDec
x Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> AliasDec -> Names
unAliases AliasDec
y

instance Monoid AliasDec where
  mempty :: AliasDec
mempty = Names -> AliasDec
AliasDec Names
forall a. Monoid a => a
mempty

instance Eq AliasDec where
  AliasDec
_ == :: AliasDec -> AliasDec -> Bool
== AliasDec
_ = Bool
True

instance Ord AliasDec where
  AliasDec
_ compare :: AliasDec -> AliasDec -> Ordering
`compare` AliasDec
_ = Ordering
EQ

instance Rename AliasDec where
  rename :: AliasDec -> RenameM AliasDec
rename (AliasDec Names
names) = Names -> AliasDec
AliasDec (Names -> AliasDec) -> RenameM Names -> RenameM AliasDec
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Names -> RenameM Names
forall a. Rename a => a -> RenameM a
rename Names
names

instance Substitute AliasDec where
  substituteNames :: Map VName VName -> AliasDec -> AliasDec
substituteNames Map VName VName
substs (AliasDec Names
names) = Names -> AliasDec
AliasDec (Names -> AliasDec) -> Names -> AliasDec
forall a b. (a -> b) -> a -> b
$ Map VName VName -> Names -> Names
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs Names
names

instance FreeIn AliasDec where
  freeIn' :: AliasDec -> FV
freeIn' = FV -> AliasDec -> FV
forall a b. a -> b -> a
const FV
forall a. Monoid a => a
mempty

instance PP.Pretty AliasDec where
  pretty :: forall ann. AliasDec -> Doc ann
pretty = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.braces (Doc ann -> Doc ann)
-> (AliasDec -> Doc ann) -> AliasDec -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
PP.commasep ([Doc ann] -> Doc ann)
-> (AliasDec -> [Doc ann]) -> AliasDec -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Doc ann) -> [VName] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
PP.pretty ([VName] -> [Doc ann])
-> (AliasDec -> [VName]) -> AliasDec -> [Doc ann]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> [VName]
namesToList (Names -> [VName]) -> (AliasDec -> Names) -> AliasDec -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AliasDec -> Names
unAliases

-- | The aliases of the let-bound variable.
type VarAliases = AliasDec

-- | Everything consumed in the expression.
type ConsumedInExp = AliasDec

-- | The aliases of what is returned by the t'Body', and what is
-- consumed inside of it.
type BodyAliasing = ([VarAliases], ConsumedInExp)

instance (RepTypes rep, ASTConstraints (OpC rep (Aliases rep))) => RepTypes (Aliases rep) where
  type LetDec (Aliases rep) = (VarAliases, LetDec rep)
  type ExpDec (Aliases rep) = (ConsumedInExp, ExpDec rep)
  type BodyDec (Aliases rep) = (BodyAliasing, BodyDec rep)
  type FParamInfo (Aliases rep) = FParamInfo rep
  type LParamInfo (Aliases rep) = LParamInfo rep
  type RetType (Aliases rep) = RetType rep
  type BranchType (Aliases rep) = BranchType rep
  type OpC (Aliases rep) = OpC rep

instance AliasesOf (VarAliases, dec) where
  aliasesOf :: (AliasDec, dec) -> Names
aliasesOf = AliasDec -> Names
unAliases (AliasDec -> Names)
-> ((AliasDec, dec) -> AliasDec) -> (AliasDec, dec) -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AliasDec, dec) -> AliasDec
forall a b. (a, b) -> a
fst

instance FreeDec AliasDec

withoutAliases ::
  (HasScope (Aliases rep) m, Monad m) =>
  ReaderT (Scope rep) m a ->
  m a
withoutAliases :: forall rep (m :: * -> *) a.
(HasScope (Aliases rep) m, Monad m) =>
ReaderT (Scope rep) m a -> m a
withoutAliases ReaderT (Scope rep) m a
m = do
  Scope rep
scope <- (Scope (Aliases rep) -> Scope rep) -> m (Scope rep)
forall a. (Scope (Aliases rep) -> a) -> m a
forall rep (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope Scope (Aliases rep) -> Scope rep
forall rep. Scope (Aliases rep) -> Scope rep
removeScopeAliases
  ReaderT (Scope rep) m a -> Scope rep -> m a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (Scope rep) m a
m Scope rep
scope

instance
  ( ASTRep rep,
    AliasedOp (OpC rep (Aliases rep)),
    IsOp (OpC rep (Aliases rep))
  ) =>
  ASTRep (Aliases rep)
  where
  expTypesFromPat :: forall (m :: * -> *).
(HasScope (Aliases rep) m, Monad m) =>
Pat (LetDec (Aliases rep)) -> m [BranchType (Aliases rep)]
expTypesFromPat =
    ReaderT (Scope rep) m [BranchType rep] -> m [BranchType rep]
forall rep (m :: * -> *) a.
(HasScope (Aliases rep) m, Monad m) =>
ReaderT (Scope rep) m a -> m a
withoutAliases (ReaderT (Scope rep) m [BranchType rep] -> m [BranchType rep])
-> (Pat (AliasDec, LetDec rep)
    -> ReaderT (Scope rep) m [BranchType rep])
-> Pat (AliasDec, LetDec rep)
-> m [BranchType rep]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (LetDec rep) -> ReaderT (Scope rep) m [BranchType rep]
forall rep (m :: * -> *).
(ASTRep rep, HasScope rep m, Monad m) =>
Pat (LetDec rep) -> m [BranchType rep]
forall (m :: * -> *).
(HasScope rep m, Monad m) =>
Pat (LetDec rep) -> m [BranchType rep]
expTypesFromPat (Pat (LetDec rep) -> ReaderT (Scope rep) m [BranchType rep])
-> (Pat (AliasDec, LetDec rep) -> Pat (LetDec rep))
-> Pat (AliasDec, LetDec rep)
-> ReaderT (Scope rep) m [BranchType rep]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (AliasDec, LetDec rep) -> Pat (LetDec rep)
forall a. Pat (AliasDec, a) -> Pat a
removePatAliases

instance (ASTRep rep, AliasedOp (OpC rep (Aliases rep))) => Aliased (Aliases rep) where
  bodyAliases :: Body (Aliases rep) -> [Names]
bodyAliases = (AliasDec -> Names) -> [AliasDec] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map AliasDec -> Names
unAliases ([AliasDec] -> [Names])
-> (Body (Aliases rep) -> [AliasDec])
-> Body (Aliases rep)
-> [Names]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BodyAliasing -> [AliasDec]
forall a b. (a, b) -> a
fst (BodyAliasing -> [AliasDec])
-> (Body (Aliases rep) -> BodyAliasing)
-> Body (Aliases rep)
-> [AliasDec]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (BodyAliasing, BodyDec rep) -> BodyAliasing
forall a b. (a, b) -> a
fst ((BodyAliasing, BodyDec rep) -> BodyAliasing)
-> (Body (Aliases rep) -> (BodyAliasing, BodyDec rep))
-> Body (Aliases rep)
-> BodyAliasing
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body (Aliases rep) -> (BodyAliasing, BodyDec rep)
Body (Aliases rep) -> BodyDec (Aliases rep)
forall rep. Body rep -> BodyDec rep
bodyDec
  consumedInBody :: Body (Aliases rep) -> Names
consumedInBody = AliasDec -> Names
unAliases (AliasDec -> Names)
-> (Body (Aliases rep) -> AliasDec) -> Body (Aliases rep) -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BodyAliasing -> AliasDec
forall a b. (a, b) -> b
snd (BodyAliasing -> AliasDec)
-> (Body (Aliases rep) -> BodyAliasing)
-> Body (Aliases rep)
-> AliasDec
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (BodyAliasing, BodyDec rep) -> BodyAliasing
forall a b. (a, b) -> a
fst ((BodyAliasing, BodyDec rep) -> BodyAliasing)
-> (Body (Aliases rep) -> (BodyAliasing, BodyDec rep))
-> Body (Aliases rep)
-> BodyAliasing
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body (Aliases rep) -> (BodyAliasing, BodyDec rep)
Body (Aliases rep) -> BodyDec (Aliases rep)
forall rep. Body rep -> BodyDec rep
bodyDec

instance
  ( ASTRep rep,
    AliasedOp (OpC rep (Aliases rep)),
    Pretty (OpC rep (Aliases rep))
  ) =>
  PrettyRep (Aliases rep)
  where
  ppExpDec :: forall a.
ExpDec (Aliases rep) -> Exp (Aliases rep) -> Maybe (Doc a)
ppExpDec (AliasDec
consumed, ExpDec rep
inner) Exp (Aliases rep)
e =
    [Doc a] -> Maybe (Doc a)
forall a. [Doc a] -> Maybe (Doc a)
maybeComment ([Doc a] -> Maybe (Doc a))
-> ([Maybe (Doc a)] -> [Doc a]) -> [Maybe (Doc a)] -> Maybe (Doc a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe (Doc a)] -> [Doc a]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (Doc a)] -> Maybe (Doc a))
-> [Maybe (Doc a)] -> Maybe (Doc a)
forall a b. (a -> b) -> a -> b
$
      [Maybe (Doc a)
exp_dec, Maybe (Doc a)
merge_dec, ExpDec rep -> Exp rep -> Maybe (Doc a)
forall a. ExpDec rep -> Exp rep -> Maybe (Doc a)
forall rep a.
PrettyRep rep =>
ExpDec rep -> Exp rep -> Maybe (Doc a)
ppExpDec ExpDec rep
inner (Exp rep -> Maybe (Doc a)) -> Exp rep -> Maybe (Doc a)
forall a b. (a -> b) -> a -> b
$ Exp (Aliases rep) -> Exp rep
forall rep. RephraseOp (OpC rep) => Exp (Aliases rep) -> Exp rep
removeExpAliases Exp (Aliases rep)
e]
    where
      merge_dec :: Maybe (Doc a)
merge_dec =
        case Exp (Aliases rep)
e of
          Loop [(FParam (Aliases rep), SubExp)]
merge LoopForm
_ Body (Aliases rep)
body ->
            let mergeParamAliases :: Param dec -> Names -> Maybe (Doc ann)
mergeParamAliases Param dec
fparam Names
als
                  | TypeBase Shape NoUniqueness -> Bool
forall shape u. TypeBase shape u -> Bool
primType (Param dec -> TypeBase Shape NoUniqueness
forall dec. Typed dec => Param dec -> TypeBase Shape NoUniqueness
paramType Param dec
fparam) =
                      Maybe (Doc ann)
forall a. Maybe a
Nothing
                  | Bool
otherwise =
                      VName -> Names -> Maybe (Doc ann)
forall a ann. Pretty a => a -> Names -> Maybe (Doc ann)
resultAliasComment (Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
fparam) Names
als
             in [Doc a] -> Maybe (Doc a)
forall a. [Doc a] -> Maybe (Doc a)
maybeComment ([Doc a] -> Maybe (Doc a))
-> ([Maybe (Doc a)] -> [Doc a]) -> [Maybe (Doc a)] -> Maybe (Doc a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe (Doc a)] -> [Doc a]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (Doc a)] -> Maybe (Doc a))
-> [Maybe (Doc a)] -> Maybe (Doc a)
forall a b. (a -> b) -> a -> b
$
                  (Param (FParamInfo rep) -> Names -> Maybe (Doc a))
-> [Param (FParamInfo rep)] -> [Names] -> [Maybe (Doc a)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Param (FParamInfo rep) -> Names -> Maybe (Doc a)
forall {dec} {ann}.
Typed dec =>
Param dec -> Names -> Maybe (Doc ann)
mergeParamAliases (((Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep))
-> [(Param (FParamInfo rep), SubExp)] -> [Param (FParamInfo rep)]
forall a b. (a -> b) -> [a] -> [b]
map (Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep)
forall a b. (a, b) -> a
fst [(Param (FParamInfo rep), SubExp)]
[(FParam (Aliases rep), SubExp)]
merge) ([Names] -> [Maybe (Doc a)]) -> [Names] -> [Maybe (Doc a)]
forall a b. (a -> b) -> a -> b
$
                    Body (Aliases rep) -> [Names]
forall rep. Aliased rep => Body rep -> [Names]
bodyAliases Body (Aliases rep)
body
          Exp (Aliases rep)
_ -> Maybe (Doc a)
forall a. Maybe a
Nothing

      exp_dec :: Maybe (Doc a)
exp_dec = case Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ AliasDec -> Names
unAliases AliasDec
consumed of
        [] -> Maybe (Doc a)
forall a. Maybe a
Nothing
        [VName]
als ->
          Doc a -> Maybe (Doc a)
forall a. a -> Maybe a
Just (Doc a -> Maybe (Doc a)) -> Doc a -> Maybe (Doc a)
forall a b. (a -> b) -> a -> b
$
            Doc a -> Doc a
forall ann. Doc ann -> Doc ann
PP.oneLine (Doc a -> Doc a) -> Doc a -> Doc a
forall a b. (a -> b) -> a -> b
$
              Doc a
"-- Consumes " Doc a -> Doc a -> Doc a
forall a. Semigroup a => a -> a -> a
<> [Doc a] -> Doc a
forall a. [Doc a] -> Doc a
PP.commasep ((VName -> Doc a) -> [VName] -> [Doc a]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Doc a
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
PP.pretty [VName]
als)

maybeComment :: [PP.Doc a] -> Maybe (PP.Doc a)
maybeComment :: forall a. [Doc a] -> Maybe (Doc a)
maybeComment [] = Maybe (Doc a)
forall a. Maybe a
Nothing
maybeComment [Doc a]
cs = Doc a -> Maybe (Doc a)
forall a. a -> Maybe a
Just (Doc a -> Maybe (Doc a)) -> Doc a -> Maybe (Doc a)
forall a b. (a -> b) -> a -> b
$ [Doc a] -> Doc a
forall a. [Doc a] -> Doc a
PP.stack [Doc a]
cs

resultAliasComment :: (PP.Pretty a) => a -> Names -> Maybe (PP.Doc ann)
resultAliasComment :: forall a ann. Pretty a => a -> Names -> Maybe (Doc ann)
resultAliasComment a
name Names
als =
  case Names -> [VName]
namesToList Names
als of
    [] -> Maybe (Doc ann)
forall a. Maybe a
Nothing
    [VName]
als' ->
      Doc ann -> Maybe (Doc ann)
forall a. a -> Maybe a
Just (Doc ann -> Maybe (Doc ann)) -> Doc ann -> Maybe (Doc ann)
forall a b. (a -> b) -> a -> b
$
        Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
PP.oneLine (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$
          Doc ann
"-- Result for "
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
PP.pretty a
name
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
" aliases "
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
PP.commasep ((VName -> Doc ann) -> [VName] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. VName -> Doc ann
PP.pretty [VName]
als')

removeAliases :: (RephraseOp (OpC rep)) => Rephraser Identity (Aliases rep) rep
removeAliases :: forall rep.
RephraseOp (OpC rep) =>
Rephraser Identity (Aliases rep) rep
removeAliases =
  Rephraser
    { rephraseExpDec :: ExpDec (Aliases rep) -> Identity (ExpDec rep)
rephraseExpDec = ExpDec rep -> Identity (ExpDec rep)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpDec rep -> Identity (ExpDec rep))
-> ((AliasDec, ExpDec rep) -> ExpDec rep)
-> (AliasDec, ExpDec rep)
-> Identity (ExpDec rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AliasDec, ExpDec rep) -> ExpDec rep
forall a b. (a, b) -> b
snd,
      rephraseLetBoundDec :: LetDec (Aliases rep) -> Identity (LetDec rep)
rephraseLetBoundDec = LetDec rep -> Identity (LetDec rep)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LetDec rep -> Identity (LetDec rep))
-> ((AliasDec, LetDec rep) -> LetDec rep)
-> (AliasDec, LetDec rep)
-> Identity (LetDec rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AliasDec, LetDec rep) -> LetDec rep
forall a b. (a, b) -> b
snd,
      rephraseBodyDec :: BodyDec (Aliases rep) -> Identity (BodyDec rep)
rephraseBodyDec = BodyDec rep -> Identity (BodyDec rep)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BodyDec rep -> Identity (BodyDec rep))
-> ((BodyAliasing, BodyDec rep) -> BodyDec rep)
-> (BodyAliasing, BodyDec rep)
-> Identity (BodyDec rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (BodyAliasing, BodyDec rep) -> BodyDec rep
forall a b. (a, b) -> b
snd,
      rephraseFParamDec :: FParamInfo (Aliases rep) -> Identity (FParamInfo rep)
rephraseFParamDec = FParamInfo rep -> Identity (FParamInfo rep)
FParamInfo (Aliases rep) -> Identity (FParamInfo rep)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      rephraseLParamDec :: LParamInfo (Aliases rep) -> Identity (LParamInfo rep)
rephraseLParamDec = LParamInfo rep -> Identity (LParamInfo rep)
LParamInfo (Aliases rep) -> Identity (LParamInfo rep)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      rephraseRetType :: RetType (Aliases rep) -> Identity (RetType rep)
rephraseRetType = RetType rep -> Identity (RetType rep)
RetType (Aliases rep) -> Identity (RetType rep)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      rephraseBranchType :: BranchType (Aliases rep) -> Identity (BranchType rep)
rephraseBranchType = BranchType rep -> Identity (BranchType rep)
BranchType (Aliases rep) -> Identity (BranchType rep)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
      rephraseOp :: Op (Aliases rep) -> Identity (Op rep)
rephraseOp = Rephraser Identity (Aliases rep) rep
-> OpC rep (Aliases rep) -> Identity (Op rep)
forall (op :: * -> *) (m :: * -> *) from to.
(RephraseOp op, Monad m) =>
Rephraser m from to -> op from -> m (op to)
forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> OpC rep from -> m (OpC rep to)
rephraseInOp Rephraser Identity (Aliases rep) rep
forall rep.
RephraseOp (OpC rep) =>
Rephraser Identity (Aliases rep) rep
removeAliases
    }

-- | Remove alias information from an aliased scope.
removeScopeAliases :: Scope (Aliases rep) -> Scope rep
removeScopeAliases :: forall rep. Scope (Aliases rep) -> Scope rep
removeScopeAliases = (NameInfo (Aliases rep) -> NameInfo rep)
-> Map VName (NameInfo (Aliases rep)) -> Map VName (NameInfo rep)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map NameInfo (Aliases rep) -> NameInfo rep
forall {rep} {a} {rep}.
(LetDec rep ~ (a, LetDec rep), FParamInfo rep ~ FParamInfo rep,
 LParamInfo rep ~ LParamInfo rep) =>
NameInfo rep -> NameInfo rep
unAlias
  where
    unAlias :: NameInfo rep -> NameInfo rep
unAlias (LetName (a
_, LetDec rep
dec)) = LetDec rep -> NameInfo rep
forall rep. LetDec rep -> NameInfo rep
LetName LetDec rep
dec
    unAlias (FParamName FParamInfo rep
dec) = FParamInfo rep -> NameInfo rep
forall rep. FParamInfo rep -> NameInfo rep
FParamName FParamInfo rep
FParamInfo rep
dec
    unAlias (LParamName LParamInfo rep
dec) = LParamInfo rep -> NameInfo rep
forall rep. LParamInfo rep -> NameInfo rep
LParamName LParamInfo rep
LParamInfo rep
dec
    unAlias (IndexName IntType
it) = IntType -> NameInfo rep
forall rep. IntType -> NameInfo rep
IndexName IntType
it

-- | Remove alias information from a program.
removeProgAliases ::
  (RephraseOp (OpC rep)) =>
  Prog (Aliases rep) ->
  Prog rep
removeProgAliases :: forall rep. RephraseOp (OpC rep) => Prog (Aliases rep) -> Prog rep
removeProgAliases = Identity (Prog rep) -> Prog rep
forall a. Identity a -> a
runIdentity (Identity (Prog rep) -> Prog rep)
-> (Prog (Aliases rep) -> Identity (Prog rep))
-> Prog (Aliases rep)
-> Prog rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rephraser Identity (Aliases rep) rep
-> Prog (Aliases rep) -> Identity (Prog rep)
forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Prog from -> m (Prog to)
rephraseProg Rephraser Identity (Aliases rep) rep
forall rep.
RephraseOp (OpC rep) =>
Rephraser Identity (Aliases rep) rep
removeAliases

-- | Remove alias information from a function.
removeFunDefAliases ::
  (RephraseOp (OpC rep)) =>
  FunDef (Aliases rep) ->
  FunDef rep
removeFunDefAliases :: forall rep.
RephraseOp (OpC rep) =>
FunDef (Aliases rep) -> FunDef rep
removeFunDefAliases = Identity (FunDef rep) -> FunDef rep
forall a. Identity a -> a
runIdentity (Identity (FunDef rep) -> FunDef rep)
-> (FunDef (Aliases rep) -> Identity (FunDef rep))
-> FunDef (Aliases rep)
-> FunDef rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rephraser Identity (Aliases rep) rep
-> FunDef (Aliases rep) -> Identity (FunDef rep)
forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> FunDef from -> m (FunDef to)
rephraseFunDef Rephraser Identity (Aliases rep) rep
forall rep.
RephraseOp (OpC rep) =>
Rephraser Identity (Aliases rep) rep
removeAliases

-- | Remove alias information from an expression.
removeExpAliases ::
  (RephraseOp (OpC rep)) =>
  Exp (Aliases rep) ->
  Exp rep
removeExpAliases :: forall rep. RephraseOp (OpC rep) => Exp (Aliases rep) -> Exp rep
removeExpAliases = Identity (Exp rep) -> Exp rep
forall a. Identity a -> a
runIdentity (Identity (Exp rep) -> Exp rep)
-> (Exp (Aliases rep) -> Identity (Exp rep))
-> Exp (Aliases rep)
-> Exp rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rephraser Identity (Aliases rep) rep
-> Exp (Aliases rep) -> Identity (Exp rep)
forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Exp from -> m (Exp to)
rephraseExp Rephraser Identity (Aliases rep) rep
forall rep.
RephraseOp (OpC rep) =>
Rephraser Identity (Aliases rep) rep
removeAliases

-- | Remove alias information from statements.
removeStmAliases ::
  (RephraseOp (OpC rep)) =>
  Stm (Aliases rep) ->
  Stm rep
removeStmAliases :: forall rep. RephraseOp (OpC rep) => Stm (Aliases rep) -> Stm rep
removeStmAliases = Identity (Stm rep) -> Stm rep
forall a. Identity a -> a
runIdentity (Identity (Stm rep) -> Stm rep)
-> (Stm (Aliases rep) -> Identity (Stm rep))
-> Stm (Aliases rep)
-> Stm rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rephraser Identity (Aliases rep) rep
-> Stm (Aliases rep) -> Identity (Stm rep)
forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Stm from -> m (Stm to)
rephraseStm Rephraser Identity (Aliases rep) rep
forall rep.
RephraseOp (OpC rep) =>
Rephraser Identity (Aliases rep) rep
removeAliases

-- | Remove alias information from body.
removeBodyAliases ::
  (RephraseOp (OpC rep)) =>
  Body (Aliases rep) ->
  Body rep
removeBodyAliases :: forall rep. RephraseOp (OpC rep) => Body (Aliases rep) -> Body rep
removeBodyAliases = Identity (Body rep) -> Body rep
forall a. Identity a -> a
runIdentity (Identity (Body rep) -> Body rep)
-> (Body (Aliases rep) -> Identity (Body rep))
-> Body (Aliases rep)
-> Body rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rephraser Identity (Aliases rep) rep
-> Body (Aliases rep) -> Identity (Body rep)
forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Body from -> m (Body to)
rephraseBody Rephraser Identity (Aliases rep) rep
forall rep.
RephraseOp (OpC rep) =>
Rephraser Identity (Aliases rep) rep
removeAliases

-- | Remove alias information from lambda.
removeLambdaAliases ::
  (RephraseOp (OpC rep)) =>
  Lambda (Aliases rep) ->
  Lambda rep
removeLambdaAliases :: forall rep.
RephraseOp (OpC rep) =>
Lambda (Aliases rep) -> Lambda rep
removeLambdaAliases = Identity (Lambda rep) -> Lambda rep
forall a. Identity a -> a
runIdentity (Identity (Lambda rep) -> Lambda rep)
-> (Lambda (Aliases rep) -> Identity (Lambda rep))
-> Lambda (Aliases rep)
-> Lambda rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rephraser Identity (Aliases rep) rep
-> Lambda (Aliases rep) -> Identity (Lambda rep)
forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Lambda from -> m (Lambda to)
rephraseLambda Rephraser Identity (Aliases rep) rep
forall rep.
RephraseOp (OpC rep) =>
Rephraser Identity (Aliases rep) rep
removeAliases

-- | Remove alias information from pattern.
removePatAliases ::
  Pat (AliasDec, a) ->
  Pat a
removePatAliases :: forall a. Pat (AliasDec, a) -> Pat a
removePatAliases = Identity (Pat a) -> Pat a
forall a. Identity a -> a
runIdentity (Identity (Pat a) -> Pat a)
-> (Pat (AliasDec, a) -> Identity (Pat a))
-> Pat (AliasDec, a)
-> Pat a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((AliasDec, a) -> Identity a)
-> Pat (AliasDec, a) -> Identity (Pat a)
forall (m :: * -> *) from to.
Monad m =>
(from -> m to) -> Pat from -> m (Pat to)
rephrasePat (a -> Identity a
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Identity a)
-> ((AliasDec, a) -> a) -> (AliasDec, a) -> Identity a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AliasDec, a) -> a
forall a b. (a, b) -> b
snd)

-- | Augment a body decoration with aliasing information provided by
-- the statements and result of that body.
mkAliasedBody ::
  (ASTRep rep, AliasedOp (OpC rep (Aliases rep))) =>
  BodyDec rep ->
  Stms (Aliases rep) ->
  Result ->
  Body (Aliases rep)
mkAliasedBody :: forall rep.
(ASTRep rep, AliasedOp (OpC rep (Aliases rep))) =>
BodyDec rep -> Stms (Aliases rep) -> Result -> Body (Aliases rep)
mkAliasedBody BodyDec rep
dec Stms (Aliases rep)
stms Result
res =
  BodyDec (Aliases rep)
-> Stms (Aliases rep) -> Result -> Body (Aliases rep)
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body (Stms (Aliases rep) -> Result -> BodyAliasing
forall rep. Aliased rep => Stms rep -> Result -> BodyAliasing
mkBodyAliasing Stms (Aliases rep)
stms Result
res, BodyDec rep
dec) Stms (Aliases rep)
stms Result
res

-- | Augment a pattern with aliasing information provided by the
-- expression the pattern is bound to.
mkAliasedPat ::
  (Aliased rep, Typed dec) =>
  Pat dec ->
  Exp rep ->
  Pat (VarAliases, dec)
mkAliasedPat :: forall rep dec.
(Aliased rep, Typed dec) =>
Pat dec -> Exp rep -> Pat (AliasDec, dec)
mkAliasedPat (Pat [PatElem dec]
pes) Exp rep
e =
  [PatElem (AliasDec, dec)] -> Pat (AliasDec, dec)
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem (AliasDec, dec)] -> Pat (AliasDec, dec))
-> [PatElem (AliasDec, dec)] -> Pat (AliasDec, dec)
forall a b. (a -> b) -> a -> b
$ (PatElem dec -> Names -> PatElem (AliasDec, dec))
-> [PatElem dec] -> [Names] -> [PatElem (AliasDec, dec)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatElem dec -> Names -> PatElem (AliasDec, dec)
forall {b}. Typed b => PatElem b -> Names -> PatElem (AliasDec, b)
annotate [PatElem dec]
pes ([Names] -> [PatElem (AliasDec, dec)])
-> [Names] -> [PatElem (AliasDec, dec)]
forall a b. (a -> b) -> a -> b
$ [PatElem dec] -> Exp rep -> [Names]
forall rep dec. Aliased rep => [PatElem dec] -> Exp rep -> [Names]
expAliases [PatElem dec]
pes Exp rep
e
  where
    annotate :: PatElem b -> Names -> PatElem (AliasDec, b)
annotate (PatElem VName
v b
dec) Names
names = VName -> (AliasDec, b) -> PatElem (AliasDec, b)
forall dec. VName -> dec -> PatElem dec
PatElem VName
v (Names -> AliasDec
AliasDec Names
names', b
dec)
      where
        names' :: Names
names' =
          case b -> TypeBase Shape NoUniqueness
forall t. Typed t => t -> TypeBase Shape NoUniqueness
typeOf b
dec of
            Array {} -> Names
names
            Mem Space
_ -> Names
names
            TypeBase Shape NoUniqueness
_ -> Names
forall a. Monoid a => a
mempty

-- | Given statements (with aliasing information) and a body result,
-- produce aliasing information for the corresponding body as a whole.
-- The aliasing includes names bound in the body, i.e. which are not
-- in scope outside of it.  Note that this does *not* include aliases
-- of results that are not bound in the statements!
mkBodyAliasing ::
  (Aliased rep) =>
  Stms rep ->
  Result ->
  BodyAliasing
mkBodyAliasing :: forall rep. Aliased rep => Stms rep -> Result -> BodyAliasing
mkBodyAliasing Stms rep
stms Result
res =
  -- We need to remove the names that are bound in stms from the alias
  -- and consumption sets.  We do this by computing the transitive
  -- closure of the alias map (within stms), then removing anything
  -- bound in stms.
  let ([Names]
aliases, Names
consumed) = Stms rep -> Result -> ([Names], Names)
forall rep. Aliased rep => Stms rep -> Result -> ([Names], Names)
mkStmsAliases Stms rep
stms Result
res
      boundNames :: Names
boundNames = (Stm rep -> Names) -> Stms rep -> Names
forall m a. Monoid m => (a -> m) -> Seq a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ([VName] -> Names
namesFromList ([VName] -> Names) -> (Stm rep -> [VName]) -> Stm rep -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames (Pat (LetDec rep) -> [VName])
-> (Stm rep -> Pat (LetDec rep)) -> Stm rep -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm rep -> Pat (LetDec rep)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat) Stms rep
stms
      consumed' :: Names
consumed' = Names
consumed Names -> Names -> Names
`namesSubtract` Names
boundNames
   in ((Names -> AliasDec) -> [Names] -> [AliasDec]
forall a b. (a -> b) -> [a] -> [b]
map Names -> AliasDec
AliasDec [Names]
aliases, Names -> AliasDec
AliasDec Names
consumed')

-- | The aliases of the result and everything consumed in the given
-- statements.
mkStmsAliases ::
  (Aliased rep) =>
  Stms rep ->
  Result ->
  ([Names], Names)
mkStmsAliases :: forall rep. Aliased rep => Stms rep -> Result -> ([Names], Names)
mkStmsAliases Stms rep
stms Result
res = AliasesAndConsumed -> [Stm rep] -> ([Names], Names)
delve AliasesAndConsumed
forall a. Monoid a => a
mempty ([Stm rep] -> ([Names], Names)) -> [Stm rep] -> ([Names], Names)
forall a b. (a -> b) -> a -> b
$ Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms rep
stms
  where
    delve :: AliasesAndConsumed -> [Stm rep] -> ([Names], Names)
delve (Map VName Names
aliasmap, Names
consumed) [] =
      ( (SubExpRes -> Names) -> Result -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map (Map VName Names -> Names -> Names
aliasClosure Map VName Names
aliasmap (Names -> Names) -> (SubExpRes -> Names) -> SubExpRes -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> Names
subExpAliases (SubExp -> Names) -> (SubExpRes -> SubExp) -> SubExpRes -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExpRes -> SubExp
resSubExp) Result
res,
        Names
consumed
      )
    delve (Map VName Names
aliasmap, Names
consumed) (Stm rep
stm : [Stm rep]
stms') =
      AliasesAndConsumed -> [Stm rep] -> ([Names], Names)
delve (AliasesAndConsumed -> Stm rep -> AliasesAndConsumed
forall rep.
Aliased rep =>
AliasesAndConsumed -> Stm rep -> AliasesAndConsumed
trackAliases (Map VName Names
aliasmap, Names
consumed) Stm rep
stm) [Stm rep]
stms'
    aliasClosure :: Map VName Names -> Names -> Names
aliasClosure Map VName Names
aliasmap Names
names =
      Names
names Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ((VName -> Names) -> [VName] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Names
look ([VName] -> [Names]) -> [VName] -> [Names]
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList Names
names)
      where
        look :: VName -> Names
look VName
k = Names -> VName -> Map VName Names -> Names
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Names
forall a. Monoid a => a
mempty VName
k Map VName Names
aliasmap

-- | A tuple of a mapping from variable names to their aliases, and
-- the names of consumed variables.
type AliasesAndConsumed =
  ( M.Map VName Names,
    Names
  )

-- | The variables consumed in these statements.
consumedInStms :: (Aliased rep) => Stms rep -> Names
consumedInStms :: forall rep. Aliased rep => Stms rep -> Names
consumedInStms = ([Names], Names) -> Names
forall a b. (a, b) -> b
snd (([Names], Names) -> Names)
-> (Stms rep -> ([Names], Names)) -> Stms rep -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stms rep -> Result -> ([Names], Names))
-> Result -> Stms rep -> ([Names], Names)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Stms rep -> Result -> ([Names], Names)
forall rep. Aliased rep => Stms rep -> Result -> ([Names], Names)
mkStmsAliases []

-- | A helper function for computing the aliases of a sequence of
-- statements.  You'd use this while recursing down the statements
-- from first to last.  The 'AliasesAndConsumed' parameter is the
-- current "state" of aliasing, and the function then returns a new
-- state.  The main thing this function provides is proper handling of
-- transitivity and "reverse" aliases.
trackAliases ::
  (Aliased rep) =>
  AliasesAndConsumed ->
  Stm rep ->
  AliasesAndConsumed
trackAliases :: forall rep.
Aliased rep =>
AliasesAndConsumed -> Stm rep -> AliasesAndConsumed
trackAliases (Map VName Names
aliasmap, Names
consumed) Stm rep
stm =
  let pat :: Pat (LetDec rep)
pat = Stm rep -> Pat (LetDec rep)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm rep
stm
      pe_als :: [(VName, Names)]
pe_als =
        [VName] -> [Names] -> [(VName, Names)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Pat (LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat (LetDec rep)
pat) ([Names] -> [(VName, Names)]) -> [Names] -> [(VName, Names)]
forall a b. (a -> b) -> a -> b
$ (Names -> Names) -> [Names] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map Names -> Names
addAliasesOfAliases ([Names] -> [Names]) -> [Names] -> [Names]
forall a b. (a -> b) -> a -> b
$ Pat (LetDec rep) -> [Names]
forall dec. AliasesOf dec => Pat dec -> [Names]
patAliases Pat (LetDec rep)
pat
      als :: Map VName Names
als = [(VName, Names)] -> Map VName Names
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName, Names)]
pe_als
      rev_als :: Map VName Names
rev_als = ((VName, Names) -> Map VName Names)
-> [(VName, Names)] -> Map VName Names
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (VName, Names) -> Map VName Names
revAls [(VName, Names)]
pe_als
      revAls :: (VName, Names) -> Map VName Names
revAls (VName
v, Names
v_als) =
        [(VName, Names)] -> Map VName Names
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Names)] -> Map VName Names)
-> [(VName, Names)] -> Map VName Names
forall a b. (a -> b) -> a -> b
$ (VName -> (VName, Names)) -> [VName] -> [(VName, Names)]
forall a b. (a -> b) -> [a] -> [b]
map (,VName -> Names
oneName VName
v) ([VName] -> [(VName, Names)]) -> [VName] -> [(VName, Names)]
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList Names
v_als
      comb :: Map VName Names -> Map VName Names -> Map VName Names
comb = (Names -> Names -> Names)
-> Map VName Names -> Map VName Names -> Map VName Names
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
(<>)
      aliasmap' :: Map VName Names
aliasmap' = Map VName Names
rev_als Map VName Names -> Map VName Names -> Map VName Names
`comb` Map VName Names
als Map VName Names -> Map VName Names -> Map VName Names
`comb` Map VName Names
aliasmap
      consumed' :: Names
consumed' = Names
consumed Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names -> Names
addAliasesOfAliases (Stm rep -> Names
forall rep. Aliased rep => Stm rep -> Names
consumedInStm Stm rep
stm)
   in (Map VName Names
aliasmap', Names
consumed')
  where
    addAliasesOfAliases :: Names -> Names
addAliasesOfAliases Names
names = Names
names Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names -> Names
aliasesOfAliases Names
names
    aliasesOfAliases :: Names -> Names
aliasesOfAliases = [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ([Names] -> Names) -> (Names -> [Names]) -> Names -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Names) -> [VName] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Names
look ([VName] -> [Names]) -> (Names -> [VName]) -> Names -> [Names]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> [VName]
namesToList
    look :: VName -> Names
look VName
k = Names -> VName -> Map VName Names -> Names
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Names
forall a. Monoid a => a
mempty VName
k Map VName Names
aliasmap

mkAliasedStm ::
  (ASTRep rep, AliasedOp (OpC rep (Aliases rep))) =>
  Pat (LetDec rep) ->
  StmAux (ExpDec rep) ->
  Exp (Aliases rep) ->
  Stm (Aliases rep)
mkAliasedStm :: forall rep.
(ASTRep rep, AliasedOp (OpC rep (Aliases rep))) =>
Pat (LetDec rep)
-> StmAux (ExpDec rep) -> Exp (Aliases rep) -> Stm (Aliases rep)
mkAliasedStm Pat (LetDec rep)
pat (StmAux Certs
cs Attrs
attrs ExpDec rep
dec) Exp (Aliases rep)
e =
  Pat (LetDec (Aliases rep))
-> StmAux (ExpDec (Aliases rep))
-> Exp (Aliases rep)
-> Stm (Aliases rep)
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let
    (Pat (LetDec rep) -> Exp (Aliases rep) -> Pat (AliasDec, LetDec rep)
forall rep dec.
(Aliased rep, Typed dec) =>
Pat dec -> Exp rep -> Pat (AliasDec, dec)
mkAliasedPat Pat (LetDec rep)
pat Exp (Aliases rep)
e)
    (Certs
-> Attrs -> (AliasDec, ExpDec rep) -> StmAux (AliasDec, ExpDec rep)
forall dec. Certs -> Attrs -> dec -> StmAux dec
StmAux Certs
cs Attrs
attrs (Names -> AliasDec
AliasDec (Names -> AliasDec) -> Names -> AliasDec
forall a b. (a -> b) -> a -> b
$ Exp (Aliases rep) -> Names
forall rep. Aliased rep => Exp rep -> Names
consumedInExp Exp (Aliases rep)
e, ExpDec rep
dec))
    Exp (Aliases rep)
e

instance (Buildable rep, AliasedOp (OpC rep (Aliases rep))) => Buildable (Aliases rep) where
  mkExpDec :: Pat (LetDec (Aliases rep))
-> Exp (Aliases rep) -> ExpDec (Aliases rep)
mkExpDec Pat (LetDec (Aliases rep))
pat Exp (Aliases rep)
e =
    let dec :: ExpDec rep
dec = Pat (LetDec rep) -> Exp rep -> ExpDec rep
forall rep.
Buildable rep =>
Pat (LetDec rep) -> Exp rep -> ExpDec rep
mkExpDec (Pat (AliasDec, LetDec rep) -> Pat (LetDec rep)
forall a. Pat (AliasDec, a) -> Pat a
removePatAliases Pat (AliasDec, LetDec rep)
Pat (LetDec (Aliases rep))
pat) (Exp rep -> ExpDec rep) -> Exp rep -> ExpDec rep
forall a b. (a -> b) -> a -> b
$ Exp (Aliases rep) -> Exp rep
forall rep. RephraseOp (OpC rep) => Exp (Aliases rep) -> Exp rep
removeExpAliases Exp (Aliases rep)
e
     in (Names -> AliasDec
AliasDec (Names -> AliasDec) -> Names -> AliasDec
forall a b. (a -> b) -> a -> b
$ Exp (Aliases rep) -> Names
forall rep. Aliased rep => Exp rep -> Names
consumedInExp Exp (Aliases rep)
e, ExpDec rep
dec)

  mkExpPat :: [Ident] -> Exp (Aliases rep) -> Pat (LetDec (Aliases rep))
mkExpPat [Ident]
ids Exp (Aliases rep)
e =
    Pat (LetDec rep) -> Exp (Aliases rep) -> Pat (AliasDec, LetDec rep)
forall rep dec.
(Aliased rep, Typed dec) =>
Pat dec -> Exp rep -> Pat (AliasDec, dec)
mkAliasedPat ([Ident] -> Exp rep -> Pat (LetDec rep)
forall rep. Buildable rep => [Ident] -> Exp rep -> Pat (LetDec rep)
mkExpPat [Ident]
ids (Exp rep -> Pat (LetDec rep)) -> Exp rep -> Pat (LetDec rep)
forall a b. (a -> b) -> a -> b
$ Exp (Aliases rep) -> Exp rep
forall rep. RephraseOp (OpC rep) => Exp (Aliases rep) -> Exp rep
removeExpAliases Exp (Aliases rep)
e) Exp (Aliases rep)
e

  mkLetNames :: forall (m :: * -> *).
(MonadFreshNames m, HasScope (Aliases rep) m) =>
[VName] -> Exp (Aliases rep) -> m (Stm (Aliases rep))
mkLetNames [VName]
names Exp (Aliases rep)
e = do
    Scope rep
env <- (Scope (Aliases rep) -> Scope rep) -> m (Scope rep)
forall a. (Scope (Aliases rep) -> a) -> m a
forall rep (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope Scope (Aliases rep) -> Scope rep
forall rep. Scope (Aliases rep) -> Scope rep
removeScopeAliases
    (ReaderT (Scope rep) m (Stm (Aliases rep))
 -> Scope rep -> m (Stm (Aliases rep)))
-> Scope rep
-> ReaderT (Scope rep) m (Stm (Aliases rep))
-> m (Stm (Aliases rep))
forall a b c. (a -> b -> c) -> b -> a -> c
flip ReaderT (Scope rep) m (Stm (Aliases rep))
-> Scope rep -> m (Stm (Aliases rep))
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT Scope rep
env (ReaderT (Scope rep) m (Stm (Aliases rep))
 -> m (Stm (Aliases rep)))
-> ReaderT (Scope rep) m (Stm (Aliases rep))
-> m (Stm (Aliases rep))
forall a b. (a -> b) -> a -> b
$ do
      Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
dec Exp rep
_ <- [VName] -> Exp rep -> ReaderT (Scope rep) m (Stm rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m, HasScope rep m) =>
[VName] -> Exp rep -> m (Stm rep)
forall (m :: * -> *).
(MonadFreshNames m, HasScope rep m) =>
[VName] -> Exp rep -> m (Stm rep)
mkLetNames [VName]
names (Exp rep -> ReaderT (Scope rep) m (Stm rep))
-> Exp rep -> ReaderT (Scope rep) m (Stm rep)
forall a b. (a -> b) -> a -> b
$ Exp (Aliases rep) -> Exp rep
forall rep. RephraseOp (OpC rep) => Exp (Aliases rep) -> Exp rep
removeExpAliases Exp (Aliases rep)
e
      Stm (Aliases rep) -> ReaderT (Scope rep) m (Stm (Aliases rep))
forall a. a -> ReaderT (Scope rep) m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stm (Aliases rep) -> ReaderT (Scope rep) m (Stm (Aliases rep)))
-> Stm (Aliases rep) -> ReaderT (Scope rep) m (Stm (Aliases rep))
forall a b. (a -> b) -> a -> b
$ Pat (LetDec rep)
-> StmAux (ExpDec rep) -> Exp (Aliases rep) -> Stm (Aliases rep)
forall rep.
(ASTRep rep, AliasedOp (OpC rep (Aliases rep))) =>
Pat (LetDec rep)
-> StmAux (ExpDec rep) -> Exp (Aliases rep) -> Stm (Aliases rep)
mkAliasedStm Pat (LetDec rep)
pat StmAux (ExpDec rep)
dec Exp (Aliases rep)
e

  mkBody :: Stms (Aliases rep) -> Result -> Body (Aliases rep)
mkBody Stms (Aliases rep)
stms Result
res =
    let Body BodyDec rep
bodyrep Stms rep
_ Result
_ = Stms rep -> Result -> Body rep
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody ((Stm (Aliases rep) -> Stm rep) -> Stms (Aliases rep) -> Stms rep
forall a b. (a -> b) -> Seq a -> Seq b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Stm (Aliases rep) -> Stm rep
forall rep. RephraseOp (OpC rep) => Stm (Aliases rep) -> Stm rep
removeStmAliases Stms (Aliases rep)
stms) Result
res
     in BodyDec rep -> Stms (Aliases rep) -> Result -> Body (Aliases rep)
forall rep.
(ASTRep rep, AliasedOp (OpC rep (Aliases rep))) =>
BodyDec rep -> Stms (Aliases rep) -> Result -> Body (Aliases rep)
mkAliasedBody BodyDec rep
bodyrep Stms (Aliases rep)
stms Result
res

instance
  ( ASTRep rep,
    AliasedOp (OpC rep (Aliases rep)),
    Buildable (Aliases rep)
  ) =>
  BuilderOps (Aliases rep)

-- | What we require of an aliasable representation.
type AliasableRep rep =
  ( ASTRep rep,
    RephraseOp (OpC rep),
    CanBeAliased (OpC rep),
    AliasedOp (OpC rep (Aliases rep))
  )

-- | The class of operations that can be given aliasing information.
-- This is a somewhat subtle concept that is only used in the
-- simplifier and when using "rep adapters".
class CanBeAliased op where
  -- | Add aliases to this op.
  addOpAliases ::
    (AliasableRep rep) => AliasTable -> op rep -> op (Aliases rep)

instance CanBeAliased NoOp where
  addOpAliases :: forall rep.
AliasableRep rep =>
Map VName Names -> NoOp rep -> NoOp (Aliases rep)
addOpAliases Map VName Names
_ NoOp rep
NoOp = NoOp (Aliases rep)
forall {k} (rep :: k). NoOp rep
NoOp