{-# LANGUAGE FlexibleContexts #-}

-- | Alias analysis of a full Futhark program.  Takes as input a
-- program with an arbitrary rep and produces one with aliases.  This
-- module does not implement the aliasing logic itself, and derives
-- its information from definitions in
-- "Futhark.IR.Prop.Aliases" and
-- "Futhark.IR.Aliases".  The alias information computed
-- here will include transitive aliases (note that this is not what
-- the building blocks do).
module Futhark.Analysis.Alias
  ( aliasAnalysis,

    -- * Ad-hoc utilities
    analyseFun,
    analyseStms,
    analyseExp,
    analyseBody,
    analyseLambda,
  )
where

import Data.List (foldl')
import qualified Data.Map as M
import Futhark.IR.Aliases

-- | Perform alias analysis on a Futhark program.
aliasAnalysis ::
  (ASTRep rep, CanBeAliased (Op rep)) =>
  Prog rep ->
  Prog (Aliases rep)
aliasAnalysis :: Prog rep -> Prog (Aliases rep)
aliasAnalysis (Prog Stms rep
consts [FunDef rep]
funs) =
  Stms (Aliases rep) -> [FunDef (Aliases rep)] -> Prog (Aliases rep)
forall rep. Stms rep -> [FunDef rep] -> Prog rep
Prog ((Stms (Aliases rep), AliasesAndConsumed) -> Stms (Aliases rep)
forall a b. (a, b) -> a
fst (AliasTable -> Stms rep -> (Stms (Aliases rep), AliasesAndConsumed)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
AliasTable -> Stms rep -> (Stms (Aliases rep), AliasesAndConsumed)
analyseStms AliasTable
forall a. Monoid a => a
mempty Stms rep
consts)) ((FunDef rep -> FunDef (Aliases rep))
-> [FunDef rep] -> [FunDef (Aliases rep)]
forall a b. (a -> b) -> [a] -> [b]
map FunDef rep -> FunDef (Aliases rep)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
FunDef rep -> FunDef (Aliases rep)
analyseFun [FunDef rep]
funs)

analyseFun ::
  (ASTRep rep, CanBeAliased (Op rep)) =>
  FunDef rep ->
  FunDef (Aliases rep)
analyseFun :: FunDef rep -> FunDef (Aliases rep)
analyseFun (FunDef Maybe EntryPoint
entry Attrs
attrs Name
fname [RetType rep]
restype [FParam rep]
params Body rep
body) =
  Maybe EntryPoint
-> Attrs
-> Name
-> [RetType (Aliases rep)]
-> [FParam (Aliases rep)]
-> Body (Aliases rep)
-> FunDef (Aliases rep)
forall rep.
Maybe EntryPoint
-> Attrs
-> Name
-> [RetType rep]
-> [FParam rep]
-> Body rep
-> FunDef rep
FunDef Maybe EntryPoint
entry Attrs
attrs Name
fname [RetType rep]
[RetType (Aliases rep)]
restype [FParam rep]
[FParam (Aliases rep)]
params Body (Aliases rep)
body'
  where
    body' :: Body (Aliases rep)
body' = AliasTable -> Body rep -> Body (Aliases rep)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
AliasTable -> Body rep -> Body (Aliases rep)
analyseBody AliasTable
forall a. Monoid a => a
mempty Body rep
body

analyseBody ::
  ( ASTRep rep,
    CanBeAliased (Op rep)
  ) =>
  AliasTable ->
  Body rep ->
  Body (Aliases rep)
analyseBody :: AliasTable -> Body rep -> Body (Aliases rep)
analyseBody AliasTable
atable (Body BodyDec rep
rep Stms rep
stms Result
result) =
  let (Stms (Aliases rep)
stms', AliasesAndConsumed
_atable') = AliasTable -> Stms rep -> (Stms (Aliases rep), AliasesAndConsumed)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
AliasTable -> Stms rep -> (Stms (Aliases rep), AliasesAndConsumed)
analyseStms AliasTable
atable Stms rep
stms
   in BodyDec rep -> Stms (Aliases rep) -> Result -> Body (Aliases rep)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
BodyDec rep -> Stms (Aliases rep) -> Result -> Body (Aliases rep)
mkAliasedBody BodyDec rep
rep Stms (Aliases rep)
stms' Result
result

analyseStms ::
  (ASTRep rep, CanBeAliased (Op rep)) =>
  AliasTable ->
  Stms rep ->
  (Stms (Aliases rep), AliasesAndConsumed)
analyseStms :: AliasTable -> Stms rep -> (Stms (Aliases rep), AliasesAndConsumed)
analyseStms AliasTable
orig_aliases =
  ((Stms (Aliases rep), AliasesAndConsumed)
 -> Stm rep -> (Stms (Aliases rep), AliasesAndConsumed))
-> (Stms (Aliases rep), AliasesAndConsumed)
-> [Stm rep]
-> (Stms (Aliases rep), AliasesAndConsumed)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (Stms (Aliases rep), AliasesAndConsumed)
-> Stm rep -> (Stms (Aliases rep), AliasesAndConsumed)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
(Stms (Aliases rep), AliasesAndConsumed)
-> Stm rep -> (Stms (Aliases rep), AliasesAndConsumed)
f (Stms (Aliases rep)
forall a. Monoid a => a
mempty, (AliasTable
orig_aliases, Names
forall a. Monoid a => a
mempty)) ([Stm rep] -> (Stms (Aliases rep), AliasesAndConsumed))
-> (Stms rep -> [Stm rep])
-> Stms rep
-> (Stms (Aliases rep), AliasesAndConsumed)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList
  where
    f :: (Stms (Aliases rep), AliasesAndConsumed)
-> Stm rep -> (Stms (Aliases rep), AliasesAndConsumed)
f (Stms (Aliases rep)
stms, AliasesAndConsumed
aliases) Stm rep
stm =
      let stm' :: Stm (Aliases rep)
stm' = AliasTable -> Stm rep -> Stm (Aliases rep)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
AliasTable -> Stm rep -> Stm (Aliases rep)
analyseStm (AliasesAndConsumed -> AliasTable
forall a b. (a, b) -> a
fst AliasesAndConsumed
aliases) Stm rep
stm
          atable' :: AliasesAndConsumed
atable' = AliasesAndConsumed -> Stm (Aliases rep) -> AliasesAndConsumed
forall rep.
Aliased rep =>
AliasesAndConsumed -> Stm rep -> AliasesAndConsumed
trackAliases AliasesAndConsumed
aliases Stm (Aliases rep)
stm'
       in (Stms (Aliases rep)
stms Stms (Aliases rep) -> Stms (Aliases rep) -> Stms (Aliases rep)
forall a. Semigroup a => a -> a -> a
<> Stm (Aliases rep) -> Stms (Aliases rep)
forall rep. Stm rep -> Stms rep
oneStm Stm (Aliases rep)
stm', AliasesAndConsumed
atable')

analyseStm ::
  (ASTRep rep, CanBeAliased (Op rep)) =>
  AliasTable ->
  Stm rep ->
  Stm (Aliases rep)
analyseStm :: AliasTable -> Stm rep -> Stm (Aliases rep)
analyseStm AliasTable
aliases (Let Pat (LetDec rep)
pat (StmAux Certs
cs Attrs
attrs ExpDec rep
dec) Exp rep
e) =
  let e' :: Exp (Aliases rep)
e' = AliasTable -> Exp rep -> Exp (Aliases rep)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
AliasTable -> Exp rep -> Exp (Aliases rep)
analyseExp AliasTable
aliases Exp rep
e
      pat' :: Pat (VarAliases, LetDec rep)
pat' = Pat (LetDec rep)
-> Exp (Aliases rep) -> Pat (VarAliases, LetDec rep)
forall rep dec.
(ASTRep rep, CanBeAliased (Op rep), Typed dec) =>
Pat dec -> Exp (Aliases rep) -> Pat (VarAliases, dec)
addAliasesToPat Pat (LetDec rep)
pat Exp (Aliases rep)
e'
      rep' :: (VarAliases, ExpDec rep)
rep' = (Names -> VarAliases
AliasDec (Names -> VarAliases) -> Names -> VarAliases
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)
   in 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 (VarAliases, LetDec rep)
Pat (LetDec (Aliases rep))
pat' (Certs
-> Attrs
-> (VarAliases, ExpDec rep)
-> StmAux (VarAliases, ExpDec rep)
forall dec. Certs -> Attrs -> dec -> StmAux dec
StmAux Certs
cs Attrs
attrs (VarAliases, ExpDec rep)
rep') Exp (Aliases rep)
e'

analyseExp ::
  (ASTRep rep, CanBeAliased (Op rep)) =>
  AliasTable ->
  Exp rep ->
  Exp (Aliases rep)
-- Would be better to put this in a BranchType annotation, but that
-- requires a lot of other work.
analyseExp :: AliasTable -> Exp rep -> Exp (Aliases rep)
analyseExp AliasTable
aliases (If SubExp
cond Body rep
tb Body rep
fb IfDec (BranchType rep)
dec) =
  let Body ((tb_als, tb_cons), tb_dec) Stms (Aliases rep)
tb_stms Result
tb_res = AliasTable -> Body rep -> Body (Aliases rep)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
AliasTable -> Body rep -> Body (Aliases rep)
analyseBody AliasTable
aliases Body rep
tb
      Body ((fb_als, fb_cons), fb_dec) Stms (Aliases rep)
fb_stms Result
fb_res = AliasTable -> Body rep -> Body (Aliases rep)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
AliasTable -> Body rep -> Body (Aliases rep)
analyseBody AliasTable
aliases Body rep
fb
      cons :: VarAliases
cons = VarAliases
tb_cons VarAliases -> VarAliases -> VarAliases
forall a. Semigroup a => a -> a -> a
<> VarAliases
fb_cons
      isConsumed :: VName -> Bool
isConsumed VName
v =
        (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Names -> Bool
`nameIn` VarAliases -> Names
unAliases VarAliases
cons) ([VName] -> Bool) -> [VName] -> Bool
forall a b. (a -> b) -> a -> b
$
          VName
v VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: Names -> [VName]
namesToList (Names -> VName -> AliasTable -> Names
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Names
forall a. Monoid a => a
mempty VName
v AliasTable
aliases)
      notConsumed :: VarAliases -> VarAliases
notConsumed =
        Names -> VarAliases
AliasDec (Names -> VarAliases)
-> (VarAliases -> Names) -> VarAliases -> VarAliases
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> Names
namesFromList
          ([VName] -> Names)
-> (VarAliases -> [VName]) -> VarAliases -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (VName -> Bool) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Bool
isConsumed)
          ([VName] -> [VName])
-> (VarAliases -> [VName]) -> VarAliases -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> [VName]
namesToList
          (Names -> [VName])
-> (VarAliases -> Names) -> VarAliases -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VarAliases -> Names
unAliases
      tb_als' :: [VarAliases]
tb_als' = (VarAliases -> VarAliases) -> [VarAliases] -> [VarAliases]
forall a b. (a -> b) -> [a] -> [b]
map VarAliases -> VarAliases
notConsumed [VarAliases]
tb_als
      fb_als' :: [VarAliases]
fb_als' = (VarAliases -> VarAliases) -> [VarAliases] -> [VarAliases]
forall a b. (a -> b) -> [a] -> [b]
map VarAliases -> VarAliases
notConsumed [VarAliases]
fb_als
      tb' :: Body (Aliases rep)
tb' = BodyDec (Aliases rep)
-> Stms (Aliases rep) -> Result -> Body (Aliases rep)
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body (([VarAliases]
tb_als', VarAliases
tb_cons), BodyDec rep
tb_dec) Stms (Aliases rep)
tb_stms Result
tb_res
      fb' :: Body (Aliases rep)
fb' = BodyDec (Aliases rep)
-> Stms (Aliases rep) -> Result -> Body (Aliases rep)
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body (([VarAliases]
fb_als', VarAliases
fb_cons), BodyDec rep
fb_dec) Stms (Aliases rep)
fb_stms Result
fb_res
   in SubExp
-> Body (Aliases rep)
-> Body (Aliases rep)
-> IfDec (BranchType (Aliases rep))
-> Exp (Aliases rep)
forall rep.
SubExp -> Body rep -> Body rep -> IfDec (BranchType rep) -> Exp rep
If SubExp
cond Body (Aliases rep)
tb' Body (Aliases rep)
fb' IfDec (BranchType rep)
IfDec (BranchType (Aliases rep))
dec
analyseExp AliasTable
aliases Exp rep
e = Mapper rep (Aliases rep) Identity -> Exp rep -> Exp (Aliases rep)
forall frep trep. Mapper frep trep Identity -> Exp frep -> Exp trep
mapExp Mapper rep (Aliases rep) Identity
analyse Exp rep
e
  where
    analyse :: Mapper rep (Aliases rep) Identity
analyse =
      Mapper :: forall frep trep (m :: * -> *).
(SubExp -> m SubExp)
-> (Scope trep -> Body frep -> m (Body trep))
-> (VName -> m VName)
-> (RetType frep -> m (RetType trep))
-> (BranchType frep -> m (BranchType trep))
-> (FParam frep -> m (FParam trep))
-> (LParam frep -> m (LParam trep))
-> (Op frep -> m (Op trep))
-> Mapper frep trep m
Mapper
        { mapOnSubExp :: SubExp -> Identity SubExp
mapOnSubExp = SubExp -> Identity SubExp
forall (m :: * -> *) a. Monad m => a -> m a
return,
          mapOnVName :: VName -> Identity VName
mapOnVName = VName -> Identity VName
forall (m :: * -> *) a. Monad m => a -> m a
return,
          mapOnBody :: Scope (Aliases rep) -> Body rep -> Identity (Body (Aliases rep))
mapOnBody = (Body rep -> Identity (Body (Aliases rep)))
-> Scope (Aliases rep) -> Body rep -> Identity (Body (Aliases rep))
forall a b. a -> b -> a
const ((Body rep -> Identity (Body (Aliases rep)))
 -> Scope (Aliases rep)
 -> Body rep
 -> Identity (Body (Aliases rep)))
-> (Body rep -> Identity (Body (Aliases rep)))
-> Scope (Aliases rep)
-> Body rep
-> Identity (Body (Aliases rep))
forall a b. (a -> b) -> a -> b
$ Body (Aliases rep) -> Identity (Body (Aliases rep))
forall (m :: * -> *) a. Monad m => a -> m a
return (Body (Aliases rep) -> Identity (Body (Aliases rep)))
-> (Body rep -> Body (Aliases rep))
-> Body rep
-> Identity (Body (Aliases rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AliasTable -> Body rep -> Body (Aliases rep)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
AliasTable -> Body rep -> Body (Aliases rep)
analyseBody AliasTable
aliases,
          mapOnRetType :: RetType rep -> Identity (RetType (Aliases rep))
mapOnRetType = RetType rep -> Identity (RetType (Aliases rep))
forall (m :: * -> *) a. Monad m => a -> m a
return,
          mapOnBranchType :: BranchType rep -> Identity (BranchType (Aliases rep))
mapOnBranchType = BranchType rep -> Identity (BranchType (Aliases rep))
forall (m :: * -> *) a. Monad m => a -> m a
return,
          mapOnFParam :: FParam rep -> Identity (FParam (Aliases rep))
mapOnFParam = FParam rep -> Identity (FParam (Aliases rep))
forall (m :: * -> *) a. Monad m => a -> m a
return,
          mapOnLParam :: LParam rep -> Identity (LParam (Aliases rep))
mapOnLParam = LParam rep -> Identity (LParam (Aliases rep))
forall (m :: * -> *) a. Monad m => a -> m a
return,
          mapOnOp :: Op rep -> Identity (Op (Aliases rep))
mapOnOp = OpWithAliases (Op rep) -> Identity (OpWithAliases (Op rep))
forall (m :: * -> *) a. Monad m => a -> m a
return (OpWithAliases (Op rep) -> Identity (OpWithAliases (Op rep)))
-> (Op rep -> OpWithAliases (Op rep))
-> Op rep
-> Identity (OpWithAliases (Op rep))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AliasTable -> Op rep -> OpWithAliases (Op rep)
forall op. CanBeAliased op => AliasTable -> op -> OpWithAliases op
addOpAliases AliasTable
aliases
        }

analyseLambda ::
  (ASTRep rep, CanBeAliased (Op rep)) =>
  AliasTable ->
  Lambda rep ->
  Lambda (Aliases rep)
analyseLambda :: AliasTable -> Lambda rep -> Lambda (Aliases rep)
analyseLambda AliasTable
aliases Lambda rep
lam =
  let body :: Body (Aliases rep)
body = AliasTable -> Body rep -> Body (Aliases rep)
forall rep.
(ASTRep rep, CanBeAliased (Op rep)) =>
AliasTable -> Body rep -> Body (Aliases rep)
analyseBody AliasTable
aliases (Body rep -> Body (Aliases rep)) -> Body rep -> Body (Aliases rep)
forall a b. (a -> b) -> a -> b
$ Lambda rep -> Body rep
forall rep. Lambda rep -> Body rep
lambdaBody Lambda rep
lam
   in Lambda rep
lam
        { lambdaBody :: Body (Aliases rep)
lambdaBody = Body (Aliases rep)
body,
          lambdaParams :: [LParam (Aliases rep)]
lambdaParams = Lambda rep -> [LParam rep]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam
        }