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

-- | Representation used by the simplification engine.  It contains
-- aliasing information and a bit of caching for various information
-- that is looked up frequently.  The name is an old relic; feel free
-- to suggest a better one.
module Futhark.Optimise.Simplify.Rep
  ( Wise,
    VarWisdom (..),
    ExpWisdom,
    removeStmWisdom,
    removeLambdaWisdom,
    removeFunDefWisdom,
    removeExpWisdom,
    removePatWisdom,
    removeBodyWisdom,
    removeScopeWisdom,
    addScopeWisdom,
    addWisdomToPat,
    mkWiseBody,
    mkWiseStm,
    mkWiseExpDec,
    CanBeWise (..),

    -- * Constructing representation
    Informing,
    informLambda,
    informFunDef,
    informStms,
    informBody,
  )
where

import Control.Category
import Control.Monad.Identity
import Control.Monad.Reader
import Data.Map.Strict qualified as M
import Futhark.Builder
import Futhark.IR
import Futhark.IR.Aliases
  ( AliasDec (..),
    ConsumedInExp,
    VarAliases,
    unAliases,
  )
import Futhark.IR.Aliases qualified as Aliases
import Futhark.IR.Prop.Aliases
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Util.Pretty
import Prelude hiding (id, (.))

-- | Representative phantom type for the simplifier representation.
data Wise rep

-- | The information associated with a let-bound variable.
newtype VarWisdom = VarWisdom {VarWisdom -> VarAliases
varWisdomAliases :: VarAliases}
  deriving (VarWisdom -> VarWisdom -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: VarWisdom -> VarWisdom -> Bool
$c/= :: VarWisdom -> VarWisdom -> Bool
== :: VarWisdom -> VarWisdom -> Bool
$c== :: VarWisdom -> VarWisdom -> Bool
Eq, Eq VarWisdom
VarWisdom -> VarWisdom -> Bool
VarWisdom -> VarWisdom -> Ordering
VarWisdom -> VarWisdom -> VarWisdom
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: VarWisdom -> VarWisdom -> VarWisdom
$cmin :: VarWisdom -> VarWisdom -> VarWisdom
max :: VarWisdom -> VarWisdom -> VarWisdom
$cmax :: VarWisdom -> VarWisdom -> VarWisdom
>= :: VarWisdom -> VarWisdom -> Bool
$c>= :: VarWisdom -> VarWisdom -> Bool
> :: VarWisdom -> VarWisdom -> Bool
$c> :: VarWisdom -> VarWisdom -> Bool
<= :: VarWisdom -> VarWisdom -> Bool
$c<= :: VarWisdom -> VarWisdom -> Bool
< :: VarWisdom -> VarWisdom -> Bool
$c< :: VarWisdom -> VarWisdom -> Bool
compare :: VarWisdom -> VarWisdom -> Ordering
$ccompare :: VarWisdom -> VarWisdom -> Ordering
Ord, Int -> VarWisdom -> ShowS
[VarWisdom] -> ShowS
VarWisdom -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [VarWisdom] -> ShowS
$cshowList :: [VarWisdom] -> ShowS
show :: VarWisdom -> String
$cshow :: VarWisdom -> String
showsPrec :: Int -> VarWisdom -> ShowS
$cshowsPrec :: Int -> VarWisdom -> ShowS
Show)

instance Rename VarWisdom where
  rename :: VarWisdom -> RenameM VarWisdom
rename = forall a. Substitute a => a -> RenameM a
substituteRename

instance Substitute VarWisdom where
  substituteNames :: Map VName VName -> VarWisdom -> VarWisdom
substituteNames Map VName VName
substs (VarWisdom VarAliases
als) =
    VarAliases -> VarWisdom
VarWisdom (forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs VarAliases
als)

instance FreeIn VarWisdom where
  freeIn' :: VarWisdom -> FV
freeIn' (VarWisdom VarAliases
als) = forall a. FreeIn a => a -> FV
freeIn' VarAliases
als

-- | Simplifier information about an expression.
data ExpWisdom = ExpWisdom
  { ExpWisdom -> VarAliases
_expWisdomConsumed :: ConsumedInExp,
    -- | The free variables in the expression.
    ExpWisdom -> VarAliases
expWisdomFree :: AliasDec
  }
  deriving (ExpWisdom -> ExpWisdom -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ExpWisdom -> ExpWisdom -> Bool
$c/= :: ExpWisdom -> ExpWisdom -> Bool
== :: ExpWisdom -> ExpWisdom -> Bool
$c== :: ExpWisdom -> ExpWisdom -> Bool
Eq, Eq ExpWisdom
ExpWisdom -> ExpWisdom -> Bool
ExpWisdom -> ExpWisdom -> Ordering
ExpWisdom -> ExpWisdom -> ExpWisdom
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ExpWisdom -> ExpWisdom -> ExpWisdom
$cmin :: ExpWisdom -> ExpWisdom -> ExpWisdom
max :: ExpWisdom -> ExpWisdom -> ExpWisdom
$cmax :: ExpWisdom -> ExpWisdom -> ExpWisdom
>= :: ExpWisdom -> ExpWisdom -> Bool
$c>= :: ExpWisdom -> ExpWisdom -> Bool
> :: ExpWisdom -> ExpWisdom -> Bool
$c> :: ExpWisdom -> ExpWisdom -> Bool
<= :: ExpWisdom -> ExpWisdom -> Bool
$c<= :: ExpWisdom -> ExpWisdom -> Bool
< :: ExpWisdom -> ExpWisdom -> Bool
$c< :: ExpWisdom -> ExpWisdom -> Bool
compare :: ExpWisdom -> ExpWisdom -> Ordering
$ccompare :: ExpWisdom -> ExpWisdom -> Ordering
Ord, Int -> ExpWisdom -> ShowS
[ExpWisdom] -> ShowS
ExpWisdom -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ExpWisdom] -> ShowS
$cshowList :: [ExpWisdom] -> ShowS
show :: ExpWisdom -> String
$cshow :: ExpWisdom -> String
showsPrec :: Int -> ExpWisdom -> ShowS
$cshowsPrec :: Int -> ExpWisdom -> ShowS
Show)

instance FreeIn ExpWisdom where
  freeIn' :: ExpWisdom -> FV
freeIn' = forall a. Monoid a => a
mempty

instance FreeDec ExpWisdom where
  precomputed :: ExpWisdom -> FV -> FV
precomputed = forall a b. a -> b -> a
const forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Names -> FV
fvNames forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. VarAliases -> Names
unAliases forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ExpWisdom -> VarAliases
expWisdomFree

instance Substitute ExpWisdom where
  substituteNames :: Map VName VName -> ExpWisdom -> ExpWisdom
substituteNames Map VName VName
substs (ExpWisdom VarAliases
cons VarAliases
free) =
    VarAliases -> VarAliases -> ExpWisdom
ExpWisdom
      (forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs VarAliases
cons)
      (forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs VarAliases
free)

instance Rename ExpWisdom where
  rename :: ExpWisdom -> RenameM ExpWisdom
rename = forall a. Substitute a => a -> RenameM a
substituteRename

-- | Simplifier information about a body.
data BodyWisdom = BodyWisdom
  { BodyWisdom -> [VarAliases]
bodyWisdomAliases :: [VarAliases],
    BodyWisdom -> VarAliases
bodyWisdomConsumed :: ConsumedInExp,
    BodyWisdom -> VarAliases
bodyWisdomFree :: AliasDec
  }
  deriving (BodyWisdom -> BodyWisdom -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: BodyWisdom -> BodyWisdom -> Bool
$c/= :: BodyWisdom -> BodyWisdom -> Bool
== :: BodyWisdom -> BodyWisdom -> Bool
$c== :: BodyWisdom -> BodyWisdom -> Bool
Eq, Eq BodyWisdom
BodyWisdom -> BodyWisdom -> Bool
BodyWisdom -> BodyWisdom -> Ordering
BodyWisdom -> BodyWisdom -> BodyWisdom
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: BodyWisdom -> BodyWisdom -> BodyWisdom
$cmin :: BodyWisdom -> BodyWisdom -> BodyWisdom
max :: BodyWisdom -> BodyWisdom -> BodyWisdom
$cmax :: BodyWisdom -> BodyWisdom -> BodyWisdom
>= :: BodyWisdom -> BodyWisdom -> Bool
$c>= :: BodyWisdom -> BodyWisdom -> Bool
> :: BodyWisdom -> BodyWisdom -> Bool
$c> :: BodyWisdom -> BodyWisdom -> Bool
<= :: BodyWisdom -> BodyWisdom -> Bool
$c<= :: BodyWisdom -> BodyWisdom -> Bool
< :: BodyWisdom -> BodyWisdom -> Bool
$c< :: BodyWisdom -> BodyWisdom -> Bool
compare :: BodyWisdom -> BodyWisdom -> Ordering
$ccompare :: BodyWisdom -> BodyWisdom -> Ordering
Ord, Int -> BodyWisdom -> ShowS
[BodyWisdom] -> ShowS
BodyWisdom -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [BodyWisdom] -> ShowS
$cshowList :: [BodyWisdom] -> ShowS
show :: BodyWisdom -> String
$cshow :: BodyWisdom -> String
showsPrec :: Int -> BodyWisdom -> ShowS
$cshowsPrec :: Int -> BodyWisdom -> ShowS
Show)

instance Rename BodyWisdom where
  rename :: BodyWisdom -> RenameM BodyWisdom
rename = forall a. Substitute a => a -> RenameM a
substituteRename

instance Substitute BodyWisdom where
  substituteNames :: Map VName VName -> BodyWisdom -> BodyWisdom
substituteNames Map VName VName
substs (BodyWisdom [VarAliases]
als VarAliases
cons VarAliases
free) =
    [VarAliases] -> VarAliases -> VarAliases -> BodyWisdom
BodyWisdom
      (forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs [VarAliases]
als)
      (forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs VarAliases
cons)
      (forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs VarAliases
free)

instance FreeIn BodyWisdom where
  freeIn' :: BodyWisdom -> FV
freeIn' (BodyWisdom [VarAliases]
als VarAliases
cons VarAliases
free) =
    forall a. FreeIn a => a -> FV
freeIn' [VarAliases]
als forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' VarAliases
cons forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> FV
freeIn' VarAliases
free

instance FreeDec BodyWisdom where
  precomputed :: BodyWisdom -> FV -> FV
precomputed = forall a b. a -> b -> a
const forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Names -> FV
fvNames forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. VarAliases -> Names
unAliases forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. BodyWisdom -> VarAliases
bodyWisdomFree

instance
  ( Informing rep,
    Ord (OpC rep (Wise rep)),
    Eq (OpC rep (Wise rep)),
    Show (OpC rep (Wise rep)),
    IsOp (OpC rep (Wise rep)),
    Pretty (OpC rep (Wise rep))
  ) =>
  RepTypes (Wise rep)
  where
  type LetDec (Wise rep) = (VarWisdom, LetDec rep)
  type ExpDec (Wise rep) = (ExpWisdom, ExpDec rep)
  type BodyDec (Wise rep) = (BodyWisdom, BodyDec rep)
  type FParamInfo (Wise rep) = FParamInfo rep
  type LParamInfo (Wise rep) = LParamInfo rep
  type RetType (Wise rep) = RetType rep
  type BranchType (Wise rep) = BranchType rep
  type OpC (Wise rep) = OpC rep

withoutWisdom ::
  (HasScope (Wise rep) m, Monad m) =>
  ReaderT (Scope rep) m a ->
  m a
withoutWisdom :: forall rep (m :: * -> *) a.
(HasScope (Wise rep) m, Monad m) =>
ReaderT (Scope rep) m a -> m a
withoutWisdom ReaderT (Scope rep) m a
m = do
  Scope rep
scope <- forall rep (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope forall rep. Scope (Wise rep) -> Scope rep
removeScopeWisdom
  forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (Scope rep) m a
m Scope rep
scope

instance (Informing rep, IsOp (OpC rep (Wise rep))) => ASTRep (Wise rep) where
  expTypesFromPat :: forall (m :: * -> *).
(HasScope (Wise rep) m, Monad m) =>
Pat (LetDec (Wise rep)) -> m [BranchType (Wise rep)]
expTypesFromPat =
    forall rep (m :: * -> *) a.
(HasScope (Wise rep) m, Monad m) =>
ReaderT (Scope rep) m a -> m a
withoutWisdom forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall rep (m :: * -> *).
(ASTRep rep, HasScope rep m, Monad m) =>
Pat (LetDec rep) -> m [BranchType rep]
expTypesFromPat forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. Pat (VarWisdom, a) -> Pat a
removePatWisdom

instance Pretty VarWisdom where
  pretty :: forall ann. VarWisdom -> Doc ann
pretty VarWisdom
_ = forall a ann. Pretty a => a -> Doc ann
pretty ()

instance (Informing rep, Pretty (OpC rep (Wise rep))) => PrettyRep (Wise rep) where
  ppExpDec :: forall a. ExpDec (Wise rep) -> Exp (Wise rep) -> Maybe (Doc a)
ppExpDec (ExpWisdom
_, ExpDec rep
dec) = forall rep a.
PrettyRep rep =>
ExpDec rep -> Exp rep -> Maybe (Doc a)
ppExpDec ExpDec rep
dec forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall rep. RephraseOp (OpC rep) => Exp (Wise rep) -> Exp rep
removeExpWisdom

instance AliasesOf (VarWisdom, dec) where
  aliasesOf :: (VarWisdom, dec) -> Names
aliasesOf = VarAliases -> Names
unAliases forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. VarWisdom -> VarAliases
varWisdomAliases forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a b. (a, b) -> a
fst

instance Informing rep => Aliased (Wise rep) where
  bodyAliases :: Body (Wise rep) -> [Names]
bodyAliases = forall a b. (a -> b) -> [a] -> [b]
map VarAliases -> Names
unAliases forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. BodyWisdom -> [VarAliases]
bodyWisdomAliases forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a b. (a, b) -> a
fst forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall rep. Body rep -> BodyDec rep
bodyDec
  consumedInBody :: Body (Wise rep) -> Names
consumedInBody = VarAliases -> Names
unAliases forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. BodyWisdom -> VarAliases
bodyWisdomConsumed forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a b. (a, b) -> a
fst forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall rep. Body rep -> BodyDec rep
bodyDec

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

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

-- | Add simplifier information to scope.  All the aliasing
-- information will be vacuous, however.
addScopeWisdom :: Scope rep -> Scope (Wise rep)
addScopeWisdom :: forall rep. Scope rep -> Scope (Wise rep)
addScopeWisdom = forall a b k. (a -> b) -> Map k a -> Map k b
M.map forall {rep} {rep}.
(FParamInfo rep ~ FParamInfo rep,
 LetDec rep ~ (VarWisdom, LetDec rep),
 LParamInfo rep ~ LParamInfo rep) =>
NameInfo rep -> NameInfo rep
alias
  where
    alias :: NameInfo rep -> NameInfo rep
alias (LetName LetDec rep
dec) = forall rep. LetDec rep -> NameInfo rep
LetName (VarAliases -> VarWisdom
VarWisdom forall a. Monoid a => a
mempty, LetDec rep
dec)
    alias (FParamName FParamInfo rep
dec) = forall rep. FParamInfo rep -> NameInfo rep
FParamName FParamInfo rep
dec
    alias (LParamName LParamInfo rep
dec) = forall rep. LParamInfo rep -> NameInfo rep
LParamName LParamInfo rep
dec
    alias (IndexName IntType
it) = forall rep. IntType -> NameInfo rep
IndexName IntType
it

-- | Remove simplifier information from function.
removeFunDefWisdom :: RephraseOp (OpC rep) => FunDef (Wise rep) -> FunDef rep
removeFunDefWisdom :: forall rep. RephraseOp (OpC rep) => FunDef (Wise rep) -> FunDef rep
removeFunDefWisdom = forall a. Identity a -> a
runIdentity forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> FunDef from -> m (FunDef to)
rephraseFunDef forall rep.
RephraseOp (OpC rep) =>
Rephraser Identity (Wise rep) rep
removeWisdom

-- | Remove simplifier information from statement.
removeStmWisdom :: RephraseOp (OpC rep) => Stm (Wise rep) -> Stm rep
removeStmWisdom :: forall rep. RephraseOp (OpC rep) => Stm (Wise rep) -> Stm rep
removeStmWisdom = forall a. Identity a -> a
runIdentity forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Stm from -> m (Stm to)
rephraseStm forall rep.
RephraseOp (OpC rep) =>
Rephraser Identity (Wise rep) rep
removeWisdom

-- | Remove simplifier information from lambda.
removeLambdaWisdom :: RephraseOp (OpC rep) => Lambda (Wise rep) -> Lambda rep
removeLambdaWisdom :: forall rep. RephraseOp (OpC rep) => Lambda (Wise rep) -> Lambda rep
removeLambdaWisdom = forall a. Identity a -> a
runIdentity forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Lambda from -> m (Lambda to)
rephraseLambda forall rep.
RephraseOp (OpC rep) =>
Rephraser Identity (Wise rep) rep
removeWisdom

-- | Remove simplifier information from body.
removeBodyWisdom :: RephraseOp (OpC rep) => Body (Wise rep) -> Body rep
removeBodyWisdom :: forall rep. RephraseOp (OpC rep) => Body (Wise rep) -> Body rep
removeBodyWisdom = forall a. Identity a -> a
runIdentity forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Body from -> m (Body to)
rephraseBody forall rep.
RephraseOp (OpC rep) =>
Rephraser Identity (Wise rep) rep
removeWisdom

-- | Remove simplifier information from expression.
removeExpWisdom :: RephraseOp (OpC rep) => Exp (Wise rep) -> Exp rep
removeExpWisdom :: forall rep. RephraseOp (OpC rep) => Exp (Wise rep) -> Exp rep
removeExpWisdom = forall a. Identity a -> a
runIdentity forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Exp from -> m (Exp to)
rephraseExp forall rep.
RephraseOp (OpC rep) =>
Rephraser Identity (Wise rep) rep
removeWisdom

-- | Remove simplifier information from pattern.
removePatWisdom :: Pat (VarWisdom, a) -> Pat a
removePatWisdom :: forall a. Pat (VarWisdom, a) -> Pat a
removePatWisdom = forall a. Identity a -> a
runIdentity forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall (m :: * -> *) from to.
Monad m =>
(from -> m to) -> Pat from -> m (Pat to)
rephrasePat (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a b. (a, b) -> b
snd)

-- | Add simplifier information to pattern.
addWisdomToPat ::
  Informing rep =>
  Pat (LetDec rep) ->
  Exp (Wise rep) ->
  Pat (LetDec (Wise rep))
addWisdomToPat :: forall rep.
Informing rep =>
Pat (LetDec rep) -> Exp (Wise rep) -> Pat (LetDec (Wise rep))
addWisdomToPat Pat (LetDec rep)
pat Exp (Wise rep)
e =
  forall {b}. (VarAliases, b) -> (VarWisdom, b)
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep dec.
(Aliased rep, Typed dec) =>
Pat dec -> Exp rep -> Pat (VarAliases, dec)
Aliases.mkAliasedPat Pat (LetDec rep)
pat Exp (Wise rep)
e
  where
    f :: (VarAliases, b) -> (VarWisdom, b)
f (VarAliases
als, b
dec) = (VarAliases -> VarWisdom
VarWisdom VarAliases
als, b
dec)

-- | Produce a body with simplifier information.
mkWiseBody ::
  Informing rep =>
  BodyDec rep ->
  Stms (Wise rep) ->
  Result ->
  Body (Wise rep)
mkWiseBody :: forall rep.
Informing rep =>
BodyDec rep -> Stms (Wise rep) -> Result -> Body (Wise rep)
mkWiseBody BodyDec rep
dec Stms (Wise rep)
stms Result
res =
  forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body
    ( [VarAliases] -> VarAliases -> VarAliases -> BodyWisdom
BodyWisdom [VarAliases]
aliases VarAliases
consumed (Names -> VarAliases
AliasDec forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn forall a b. (a -> b) -> a -> b
$ forall rep.
(FreeIn (Op rep), FreeIn (LetDec rep), FreeIn (LParamInfo rep),
 FreeIn (FParamInfo rep), FreeDec (BodyDec rep),
 FreeIn (RetType rep), FreeIn (BranchType rep),
 FreeDec (ExpDec rep)) =>
Stms rep -> Result -> FV
freeInStmsAndRes Stms (Wise rep)
stms Result
res),
      BodyDec rep
dec
    )
    Stms (Wise rep)
stms
    Result
res
  where
    ([VarAliases]
aliases, VarAliases
consumed) = forall rep.
Aliased rep =>
Stms rep -> Result -> ([VarAliases], VarAliases)
Aliases.mkBodyAliasing Stms (Wise rep)
stms Result
res

-- | Produce a statement with simplifier information.
mkWiseStm ::
  Informing rep =>
  Pat (LetDec rep) ->
  StmAux (ExpDec rep) ->
  Exp (Wise rep) ->
  Stm (Wise rep)
mkWiseStm :: forall rep.
Informing rep =>
Pat (LetDec rep)
-> StmAux (ExpDec rep) -> Exp (Wise rep) -> Stm (Wise rep)
mkWiseStm Pat (LetDec rep)
pat (StmAux Certs
cs Attrs
attrs ExpDec rep
dec) Exp (Wise rep)
e =
  let pat' :: Pat (LetDec (Wise rep))
pat' = forall rep.
Informing rep =>
Pat (LetDec rep) -> Exp (Wise rep) -> Pat (LetDec (Wise rep))
addWisdomToPat Pat (LetDec rep)
pat Exp (Wise rep)
e
   in forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec (Wise rep))
pat' (forall dec. Certs -> Attrs -> dec -> StmAux dec
StmAux Certs
cs Attrs
attrs forall a b. (a -> b) -> a -> b
$ forall rep.
Informing rep =>
Pat (LetDec (Wise rep))
-> ExpDec rep -> Exp (Wise rep) -> ExpDec (Wise rep)
mkWiseExpDec Pat (LetDec (Wise rep))
pat' ExpDec rep
dec Exp (Wise rep)
e) Exp (Wise rep)
e

-- | Produce simplifier information for an expression.
mkWiseExpDec ::
  Informing rep =>
  Pat (LetDec (Wise rep)) ->
  ExpDec rep ->
  Exp (Wise rep) ->
  ExpDec (Wise rep)
mkWiseExpDec :: forall rep.
Informing rep =>
Pat (LetDec (Wise rep))
-> ExpDec rep -> Exp (Wise rep) -> ExpDec (Wise rep)
mkWiseExpDec Pat (LetDec (Wise rep))
pat ExpDec rep
expdec Exp (Wise rep)
e =
  ( VarAliases -> VarAliases -> ExpWisdom
ExpWisdom
      (Names -> VarAliases
AliasDec forall a b. (a -> b) -> a -> b
$ forall rep. Aliased rep => Exp rep -> Names
consumedInExp Exp (Wise rep)
e)
      (Names -> VarAliases
AliasDec forall a b. (a -> b) -> a -> b
$ forall a. FreeIn a => a -> Names
freeIn Pat (LetDec (Wise rep))
pat forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn ExpDec rep
expdec forall a. Semigroup a => a -> a -> a
<> forall a. FreeIn a => a -> Names
freeIn Exp (Wise rep)
e),
    ExpDec rep
expdec
  )

instance (Buildable rep, Informing rep) => Buildable (Wise rep) where
  mkExpPat :: [Ident] -> Exp (Wise rep) -> Pat (LetDec (Wise rep))
mkExpPat [Ident]
ids Exp (Wise rep)
e =
    forall rep.
Informing rep =>
Pat (LetDec rep) -> Exp (Wise rep) -> Pat (LetDec (Wise rep))
addWisdomToPat (forall rep. Buildable rep => [Ident] -> Exp rep -> Pat (LetDec rep)
mkExpPat [Ident]
ids forall a b. (a -> b) -> a -> b
$ forall rep. RephraseOp (OpC rep) => Exp (Wise rep) -> Exp rep
removeExpWisdom Exp (Wise rep)
e) Exp (Wise rep)
e

  mkExpDec :: Pat (LetDec (Wise rep)) -> Exp (Wise rep) -> ExpDec (Wise rep)
mkExpDec Pat (LetDec (Wise rep))
pat Exp (Wise rep)
e =
    forall rep.
Informing rep =>
Pat (LetDec (Wise rep))
-> ExpDec rep -> Exp (Wise rep) -> ExpDec (Wise rep)
mkWiseExpDec Pat (LetDec (Wise rep))
pat (forall rep.
Buildable rep =>
Pat (LetDec rep) -> Exp rep -> ExpDec rep
mkExpDec (forall a. Pat (VarWisdom, a) -> Pat a
removePatWisdom Pat (LetDec (Wise rep))
pat) forall a b. (a -> b) -> a -> b
$ forall rep. RephraseOp (OpC rep) => Exp (Wise rep) -> Exp rep
removeExpWisdom Exp (Wise rep)
e) Exp (Wise rep)
e

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

  mkBody :: Stms (Wise rep) -> Result -> Body (Wise rep)
mkBody Stms (Wise rep)
stms Result
res =
    let Body BodyDec rep
bodyrep Stms rep
_ Result
_ = forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall rep. RephraseOp (OpC rep) => Stm (Wise rep) -> Stm rep
removeStmWisdom Stms (Wise rep)
stms) Result
res
     in forall rep.
Informing rep =>
BodyDec rep -> Stms (Wise rep) -> Result -> Body (Wise rep)
mkWiseBody BodyDec rep
bodyrep Stms (Wise rep)
stms Result
res

-- | Constraints that let us transform a representation into a 'Wise'
-- representation.
type Informing rep =
  ( ASTRep rep,
    AliasedOp (OpC rep (Wise rep)),
    RephraseOp (OpC rep),
    CanBeWise (OpC rep),
    FreeIn (OpC rep (Wise rep))
  )

-- | A type class for indicating that this operation can be lifted into the simplifier representation.
class CanBeWise op where
  addOpWisdom :: Informing rep => op rep -> op (Wise rep)

instance CanBeWise NoOp where
  addOpWisdom :: forall rep. Informing rep => NoOp rep -> NoOp (Wise rep)
addOpWisdom NoOp rep
_ = forall a. HasCallStack => a
undefined

-- | Construct a 'Wise' statement.
informStm :: Informing rep => Stm rep -> Stm (Wise rep)
informStm :: forall rep. Informing rep => Stm rep -> Stm (Wise rep)
informStm (Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
aux Exp rep
e) = forall rep.
Informing rep =>
Pat (LetDec rep)
-> StmAux (ExpDec rep) -> Exp (Wise rep) -> Stm (Wise rep)
mkWiseStm Pat (LetDec rep)
pat StmAux (ExpDec rep)
aux forall a b. (a -> b) -> a -> b
$ forall rep. Informing rep => Exp rep -> Exp (Wise rep)
informExp Exp rep
e

-- | Construct 'Wise' statements.
informStms :: Informing rep => Stms rep -> Stms (Wise rep)
informStms :: forall rep. Informing rep => Stms rep -> Stms (Wise rep)
informStms = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall rep. Informing rep => Stm rep -> Stm (Wise rep)
informStm

-- | Construct a 'Wise' body.
informBody :: Informing rep => Body rep -> Body (Wise rep)
informBody :: forall rep. Informing rep => Body rep -> Body (Wise rep)
informBody (Body BodyDec rep
dec Stms rep
stms Result
res) = forall rep.
Informing rep =>
BodyDec rep -> Stms (Wise rep) -> Result -> Body (Wise rep)
mkWiseBody BodyDec rep
dec (forall rep. Informing rep => Stms rep -> Stms (Wise rep)
informStms Stms rep
stms) Result
res

-- | Construct a 'Wise' lambda.
informLambda :: Informing rep => Lambda rep -> Lambda (Wise rep)
informLambda :: forall rep. Informing rep => Lambda rep -> Lambda (Wise rep)
informLambda (Lambda [LParam rep]
ps Body rep
body [Type]
ret) = forall rep. [LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda [LParam rep]
ps (forall rep. Informing rep => Body rep -> Body (Wise rep)
informBody Body rep
body) [Type]
ret

-- | Construct a 'Wise' expression.
informExp :: Informing rep => Exp rep -> Exp (Wise rep)
informExp :: forall rep. Informing rep => Exp rep -> Exp (Wise rep)
informExp (Match [SubExp]
cond [Case (Body rep)]
cases Body rep
defbody (MatchDec [BranchType rep]
ts MatchSort
ifsort)) =
  forall rep.
[SubExp]
-> [Case (Body rep)]
-> Body rep
-> MatchDec (BranchType rep)
-> Exp rep
Match [SubExp]
cond (forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall rep. Informing rep => Body rep -> Body (Wise rep)
informBody) [Case (Body rep)]
cases) (forall rep. Informing rep => Body rep -> Body (Wise rep)
informBody Body rep
defbody) (forall rt. [rt] -> MatchSort -> MatchDec rt
MatchDec [BranchType rep]
ts MatchSort
ifsort)
informExp (DoLoop [(FParam rep, SubExp)]
merge LoopForm rep
form Body rep
loopbody) =
  let form' :: LoopForm (Wise rep)
form' = case LoopForm rep
form of
        ForLoop VName
i IntType
it SubExp
bound [(LParam rep, VName)]
params -> forall rep.
VName -> IntType -> SubExp -> [(LParam rep, VName)] -> LoopForm rep
ForLoop VName
i IntType
it SubExp
bound [(LParam rep, VName)]
params
        WhileLoop VName
cond -> forall rep. VName -> LoopForm rep
WhileLoop VName
cond
   in forall rep.
[(FParam rep, SubExp)] -> LoopForm rep -> Body rep -> Exp rep
DoLoop [(FParam rep, SubExp)]
merge LoopForm (Wise rep)
form' forall a b. (a -> b) -> a -> b
$ forall rep. Informing rep => Body rep -> Body (Wise rep)
informBody Body rep
loopbody
informExp Exp rep
e = forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) frep trep.
Monad m =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper rep (Wise rep) Identity
mapper Exp rep
e
  where
    mapper :: Mapper rep (Wise rep) Identity
mapper =
      Mapper
        { mapOnBody :: Scope (Wise rep) -> Body rep -> Identity (Body (Wise rep))
mapOnBody = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall rep. Informing rep => Body rep -> Body (Wise rep)
informBody,
          mapOnSubExp :: SubExp -> Identity SubExp
mapOnSubExp = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnVName :: VName -> Identity VName
mapOnVName = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnRetType :: RetType rep -> Identity (RetType (Wise rep))
mapOnRetType = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnBranchType :: BranchType rep -> Identity (BranchType (Wise rep))
mapOnBranchType = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnFParam :: FParam rep -> Identity (FParam (Wise rep))
mapOnFParam = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnLParam :: LParam rep -> Identity (LParam (Wise rep))
mapOnLParam = forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnOp :: Op rep -> Identity (Op (Wise rep))
mapOnOp = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall (op :: * -> *) rep.
(CanBeWise op, Informing rep) =>
op rep -> op (Wise rep)
addOpWisdom
        }

-- | Construct a 'Wise' function definition.
informFunDef :: Informing rep => FunDef rep -> FunDef (Wise rep)
informFunDef :: forall rep. Informing rep => FunDef rep -> FunDef (Wise rep)
informFunDef (FunDef Maybe EntryPoint
entry Attrs
attrs Name
fname [RetType rep]
rettype [FParam rep]
params Body rep
body) =
  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 [FParam rep]
params forall a b. (a -> b) -> a -> b
$ forall rep. Informing rep => Body rep -> Body (Wise rep)
informBody Body rep
body