{-# LANGUAGE TypeFamilies #-}

-- | This module provides various simple ways to query and manipulate
-- fundamental Futhark terms, such as types and values.  The intent is
-- to keep "Futhark.IR.Syntax" simple, and put whatever embellishments
-- we need here.  This is an internal, desugared representation.
module Futhark.IR.Prop
  ( module Futhark.IR.Prop.Reshape,
    module Futhark.IR.Prop.Rearrange,
    module Futhark.IR.Prop.Types,
    module Futhark.IR.Prop.Constants,
    module Futhark.IR.Prop.TypeOf,
    module Futhark.IR.Prop.Pat,
    module Futhark.IR.Prop.Names,
    module Futhark.IR.RetType,
    module Futhark.IR.Rephrase,

    -- * Built-in functions
    isBuiltInFunction,
    builtInFunctions,

    -- * Extra tools
    asBasicOp,
    safeExp,
    subExpVars,
    subExpVar,
    commutativeLambda,
    defAux,
    stmCerts,
    certify,
    expExtTypesFromPat,
    attrsForAssert,
    lamIsBinOp,
    ASTConstraints,
    IsOp (..),
    ASTRep (..),
  )
where

import Control.Monad
import Data.List (elemIndex, find)
import Data.Map.Strict qualified as M
import Data.Maybe (isJust, mapMaybe)
import Data.Set qualified as S
import Futhark.IR.Pretty
import Futhark.IR.Prop.Constants
import Futhark.IR.Prop.Names
import Futhark.IR.Prop.Pat
import Futhark.IR.Prop.Rearrange
import Futhark.IR.Prop.Reshape
import Futhark.IR.Prop.TypeOf
import Futhark.IR.Prop.Types
import Futhark.IR.Rephrase
import Futhark.IR.RetType
import Futhark.IR.Syntax
import Futhark.Transform.Rename (Rename, Renameable)
import Futhark.Transform.Substitute (Substitutable, Substitute)
import Futhark.Util (maybeNth)

-- | @isBuiltInFunction k@ is 'True' if @k@ is an element of 'builtInFunctions'.
isBuiltInFunction :: Name -> Bool
isBuiltInFunction :: Name -> Bool
isBuiltInFunction Name
fnm = Name
fnm Name -> Map Name (PrimType, [PrimType]) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map Name (PrimType, [PrimType])
builtInFunctions

-- | A map of all built-in functions and their types.
builtInFunctions :: M.Map Name (PrimType, [PrimType])
builtInFunctions :: Map Name (PrimType, [PrimType])
builtInFunctions = [(Name, (PrimType, [PrimType]))] -> Map Name (PrimType, [PrimType])
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, (PrimType, [PrimType]))]
 -> Map Name (PrimType, [PrimType]))
-> [(Name, (PrimType, [PrimType]))]
-> Map Name (PrimType, [PrimType])
forall a b. (a -> b) -> a -> b
$ ((String, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
 -> (Name, (PrimType, [PrimType])))
-> [(String,
     ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
-> [(Name, (PrimType, [PrimType]))]
forall a b. (a -> b) -> [a] -> [b]
map (String, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))
-> (Name, (PrimType, [PrimType]))
forall {b} {a} {c}. (String, (b, a, c)) -> (Name, (a, b))
namify ([(String, ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
 -> [(Name, (PrimType, [PrimType]))])
-> [(String,
     ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
-> [(Name, (PrimType, [PrimType]))]
forall a b. (a -> b) -> a -> b
$ Map String ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
-> [(String,
     ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue))]
forall k a. Map k a -> [(k, a)]
M.toList Map String ([PrimType], PrimType, [PrimValue] -> Maybe PrimValue)
primFuns
  where
    namify :: (String, (b, a, c)) -> (Name, (a, b))
namify (String
k, (b
paramts, a
ret, c
_)) = (String -> Name
nameFromString String
k, (a
ret, b
paramts))

-- | If the expression is a t'BasicOp', return it, otherwise 'Nothing'.
asBasicOp :: Exp rep -> Maybe BasicOp
asBasicOp :: forall rep. Exp rep -> Maybe BasicOp
asBasicOp (BasicOp BasicOp
op) = BasicOp -> Maybe BasicOp
forall a. a -> Maybe a
Just BasicOp
op
asBasicOp Exp rep
_ = Maybe BasicOp
forall a. Maybe a
Nothing

-- | An expression is safe if it is always well-defined (assuming that
-- any required certificates have been checked) in any context.  For
-- example, array indexing is not safe, as the index may be out of
-- bounds.  On the other hand, adding two numbers cannot fail.
safeExp :: (ASTRep rep) => Exp rep -> Bool
safeExp :: forall rep. ASTRep rep => Exp rep -> Bool
safeExp (BasicOp BasicOp
op) = BasicOp -> Bool
safeBasicOp BasicOp
op
  where
    safeBasicOp :: BasicOp -> Bool
safeBasicOp (BinOp (SDiv IntType
_ Safety
Safe) SubExp
_ SubExp
_) = Bool
True
    safeBasicOp (BinOp (SDivUp IntType
_ Safety
Safe) SubExp
_ SubExp
_) = Bool
True
    safeBasicOp (BinOp (SQuot IntType
_ Safety
Safe) SubExp
_ SubExp
_) = Bool
True
    safeBasicOp (BinOp (UDiv IntType
_ Safety
Safe) SubExp
_ SubExp
_) = Bool
True
    safeBasicOp (BinOp (UDivUp IntType
_ Safety
Safe) SubExp
_ SubExp
_) = Bool
True
    safeBasicOp (BinOp (SMod IntType
_ Safety
Safe) SubExp
_ SubExp
_) = Bool
True
    safeBasicOp (BinOp (SRem IntType
_ Safety
Safe) SubExp
_ SubExp
_) = Bool
True
    safeBasicOp (BinOp (UMod IntType
_ Safety
Safe) SubExp
_ SubExp
_) = Bool
True
    safeBasicOp (BinOp SDiv {} SubExp
_ (Constant PrimValue
y)) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ PrimValue -> Bool
zeroIsh PrimValue
y
    safeBasicOp (BinOp SDiv {} SubExp
_ SubExp
_) = Bool
False
    safeBasicOp (BinOp SDivUp {} SubExp
_ (Constant PrimValue
y)) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ PrimValue -> Bool
zeroIsh PrimValue
y
    safeBasicOp (BinOp SDivUp {} SubExp
_ SubExp
_) = Bool
False
    safeBasicOp (BinOp UDiv {} SubExp
_ (Constant PrimValue
y)) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ PrimValue -> Bool
zeroIsh PrimValue
y
    safeBasicOp (BinOp UDiv {} SubExp
_ SubExp
_) = Bool
False
    safeBasicOp (BinOp UDivUp {} SubExp
_ (Constant PrimValue
y)) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ PrimValue -> Bool
zeroIsh PrimValue
y
    safeBasicOp (BinOp UDivUp {} SubExp
_ SubExp
_) = Bool
False
    safeBasicOp (BinOp SMod {} SubExp
_ (Constant PrimValue
y)) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ PrimValue -> Bool
zeroIsh PrimValue
y
    safeBasicOp (BinOp SMod {} SubExp
_ SubExp
_) = Bool
False
    safeBasicOp (BinOp UMod {} SubExp
_ (Constant PrimValue
y)) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ PrimValue -> Bool
zeroIsh PrimValue
y
    safeBasicOp (BinOp UMod {} SubExp
_ SubExp
_) = Bool
False
    safeBasicOp (BinOp SQuot {} SubExp
_ (Constant PrimValue
y)) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ PrimValue -> Bool
zeroIsh PrimValue
y
    safeBasicOp (BinOp SQuot {} SubExp
_ SubExp
_) = Bool
False
    safeBasicOp (BinOp SRem {} SubExp
_ (Constant PrimValue
y)) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ PrimValue -> Bool
zeroIsh PrimValue
y
    safeBasicOp (BinOp SRem {} SubExp
_ SubExp
_) = Bool
False
    safeBasicOp (BinOp Pow {} SubExp
_ (Constant PrimValue
y)) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ PrimValue -> Bool
negativeIsh PrimValue
y
    safeBasicOp (BinOp Pow {} SubExp
_ SubExp
_) = Bool
False
    safeBasicOp ArrayLit {} = Bool
True
    safeBasicOp BinOp {} = Bool
True
    safeBasicOp SubExp {} = Bool
True
    safeBasicOp UnOp {} = Bool
True
    safeBasicOp CmpOp {} = Bool
True
    safeBasicOp ConvOp {} = Bool
True
    safeBasicOp Scratch {} = Bool
True
    safeBasicOp Concat {} = Bool
True
    safeBasicOp Reshape {} = Bool
True
    safeBasicOp Rearrange {} = Bool
True
    safeBasicOp Manifest {} = Bool
True
    safeBasicOp Iota {} = Bool
True
    safeBasicOp Replicate {} = Bool
True
    safeBasicOp BasicOp
_ = Bool
False
safeExp (Loop [(FParam rep, SubExp)]
_ LoopForm
_ Body rep
body) = Body rep -> Bool
forall rep. ASTRep rep => Body rep -> Bool
safeBody Body rep
body
safeExp (Apply Name
fname [(SubExp, Diet)]
_ [(RetType rep, RetAls)]
_ (Safety, SrcLoc, [SrcLoc])
_) =
  Name -> Bool
isBuiltInFunction Name
fname
safeExp (Match [SubExp]
_ [Case (Body rep)]
cases Body rep
def_case MatchDec (BranchType rep)
_) =
  (Case (Body rep) -> Bool) -> [Case (Body rep)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Stm rep -> Bool) -> Seq (Stm rep) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Exp rep -> Bool
forall rep. ASTRep rep => Exp rep -> Bool
safeExp (Exp rep -> Bool) -> (Stm rep -> Exp rep) -> Stm rep -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp) (Seq (Stm rep) -> Bool)
-> (Case (Body rep) -> Seq (Stm rep)) -> Case (Body rep) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body rep -> Seq (Stm rep)
forall rep. Body rep -> Stms rep
bodyStms (Body rep -> Seq (Stm rep))
-> (Case (Body rep) -> Body rep)
-> Case (Body rep)
-> Seq (Stm rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Case (Body rep) -> Body rep
forall body. Case body -> body
caseBody) [Case (Body rep)]
cases
    Bool -> Bool -> Bool
&& (Stm rep -> Bool) -> Seq (Stm rep) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Exp rep -> Bool
forall rep. ASTRep rep => Exp rep -> Bool
safeExp (Exp rep -> Bool) -> (Stm rep -> Exp rep) -> Stm rep -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp) (Body rep -> Seq (Stm rep)
forall rep. Body rep -> Stms rep
bodyStms Body rep
def_case)
safeExp WithAcc {} = Bool
True -- Although unlikely to matter.
safeExp (Op Op rep
op) = Op rep -> Bool
forall rep. ASTRep rep => OpC rep rep -> Bool
forall (op :: * -> *) rep. (IsOp op, ASTRep rep) => op rep -> Bool
safeOp Op rep
op

safeBody :: (ASTRep rep) => Body rep -> Bool
safeBody :: forall rep. ASTRep rep => Body rep -> Bool
safeBody = (Stm rep -> Bool) -> Seq (Stm rep) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Exp rep -> Bool
forall rep. ASTRep rep => Exp rep -> Bool
safeExp (Exp rep -> Bool) -> (Stm rep -> Exp rep) -> Stm rep -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp) (Seq (Stm rep) -> Bool)
-> (Body rep -> Seq (Stm rep)) -> Body rep -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body rep -> Seq (Stm rep)
forall rep. Body rep -> Stms rep
bodyStms

-- | Return the variable names used in 'Var' subexpressions.  May contain
-- duplicates.
subExpVars :: [SubExp] -> [VName]
subExpVars :: [SubExp] -> [VName]
subExpVars = (SubExp -> Maybe VName) -> [SubExp] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe SubExp -> Maybe VName
subExpVar

-- | If the t'SubExp' is a 'Var' return the variable name.
subExpVar :: SubExp -> Maybe VName
subExpVar :: SubExp -> Maybe VName
subExpVar (Var VName
v) = VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v
subExpVar Constant {} = Maybe VName
forall a. Maybe a
Nothing

-- | Does the given lambda represent a known commutative function?
-- Based on pattern matching and checking whether the lambda
-- represents a known arithmetic operator; don't expect anything
-- clever here.
commutativeLambda :: Lambda rep -> Bool
commutativeLambda :: forall rep. Lambda rep -> Bool
commutativeLambda Lambda rep
lam =
  let body :: Body rep
body = Lambda rep -> Body rep
forall rep. Lambda rep -> Body rep
lambdaBody Lambda rep
lam
      n2 :: Int
n2 = [Param (LParamInfo rep)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Lambda rep -> [Param (LParamInfo rep)]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
      ([Param (LParamInfo rep)]
xps, [Param (LParamInfo rep)]
yps) = Int
-> [Param (LParamInfo rep)]
-> ([Param (LParamInfo rep)], [Param (LParamInfo rep)])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n2 (Lambda rep -> [Param (LParamInfo rep)]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam)

      okComponent :: (Param (LParamInfo rep), Param (LParamInfo rep), SubExpRes) -> Bool
okComponent (Param (LParamInfo rep), Param (LParamInfo rep), SubExpRes)
c = Maybe (Stm rep) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Stm rep) -> Bool) -> Maybe (Stm rep) -> Bool
forall a b. (a -> b) -> a -> b
$ (Stm rep -> Bool) -> Seq (Stm rep) -> Maybe (Stm rep)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Param (LParamInfo rep), Param (LParamInfo rep), SubExpRes)
-> Stm rep -> Bool
forall {dec} {dec} {rep}.
(Param dec, Param dec, SubExpRes) -> Stm rep -> Bool
okBinOp (Param (LParamInfo rep), Param (LParamInfo rep), SubExpRes)
c) (Seq (Stm rep) -> Maybe (Stm rep))
-> Seq (Stm rep) -> Maybe (Stm rep)
forall a b. (a -> b) -> a -> b
$ Body rep -> Seq (Stm rep)
forall rep. Body rep -> Stms rep
bodyStms Body rep
body
      okBinOp :: (Param dec, Param dec, SubExpRes) -> Stm rep -> Bool
okBinOp
        (Param dec
xp, Param dec
yp, SubExpRes Certs
_ (Var VName
r))
        (Let (Pat [PatElem (LetDec rep)
pe]) StmAux (ExpDec rep)
_ (BasicOp (BinOp BinOp
op (Var VName
x) (Var VName
y)))) =
          PatElem (LetDec rep) -> VName
forall dec. PatElem dec -> VName
patElemName PatElem (LetDec rep)
pe VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
r
            Bool -> Bool -> Bool
&& BinOp -> Bool
commutativeBinOp BinOp
op
            Bool -> Bool -> Bool
&& ( (VName
x VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
xp Bool -> Bool -> Bool
&& VName
y VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
yp)
                   Bool -> Bool -> Bool
|| (VName
y VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
xp Bool -> Bool -> Bool
&& VName
x VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
yp)
               )
      okBinOp (Param dec, Param dec, SubExpRes)
_ Stm rep
_ = Bool
False
   in Int
n2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Param (LParamInfo rep)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Lambda rep -> [Param (LParamInfo rep)]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam)
        Bool -> Bool -> Bool
&& Int
n2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [SubExpRes] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Body rep -> [SubExpRes]
forall rep. Body rep -> [SubExpRes]
bodyResult Body rep
body)
        Bool -> Bool -> Bool
&& ((Param (LParamInfo rep), Param (LParamInfo rep), SubExpRes)
 -> Bool)
-> [(Param (LParamInfo rep), Param (LParamInfo rep), SubExpRes)]
-> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Param (LParamInfo rep), Param (LParamInfo rep), SubExpRes) -> Bool
okComponent ([Param (LParamInfo rep)]
-> [Param (LParamInfo rep)]
-> [SubExpRes]
-> [(Param (LParamInfo rep), Param (LParamInfo rep), SubExpRes)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Param (LParamInfo rep)]
xps [Param (LParamInfo rep)]
yps ([SubExpRes]
 -> [(Param (LParamInfo rep), Param (LParamInfo rep), SubExpRes)])
-> [SubExpRes]
-> [(Param (LParamInfo rep), Param (LParamInfo rep), SubExpRes)]
forall a b. (a -> b) -> a -> b
$ Body rep -> [SubExpRes]
forall rep. Body rep -> [SubExpRes]
bodyResult Body rep
body)

-- | A 'StmAux' with empty 'Certs'.
defAux :: dec -> StmAux dec
defAux :: forall dec. dec -> StmAux dec
defAux = Certs -> Attrs -> dec -> StmAux dec
forall dec. Certs -> Attrs -> dec -> StmAux dec
StmAux Certs
forall a. Monoid a => a
mempty Attrs
forall a. Monoid a => a
mempty

-- | The certificates associated with a statement.
stmCerts :: Stm rep -> Certs
stmCerts :: forall rep. Stm rep -> Certs
stmCerts = StmAux (ExpDec rep) -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts (StmAux (ExpDec rep) -> Certs)
-> (Stm rep -> StmAux (ExpDec rep)) -> Stm rep -> Certs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm rep -> StmAux (ExpDec rep)
forall rep. Stm rep -> StmAux (ExpDec rep)
stmAux

-- | Add certificates to a statement.
certify :: Certs -> Stm rep -> Stm rep
certify :: forall rep. Certs -> Stm rep -> Stm rep
certify Certs
cs1 (Let Pat (LetDec rep)
pat (StmAux Certs
cs2 Attrs
attrs ExpDec rep
dec) Exp rep
e) =
  Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec rep)
pat (Certs -> Attrs -> ExpDec rep -> StmAux (ExpDec rep)
forall dec. Certs -> Attrs -> dec -> StmAux dec
StmAux (Certs
cs2 Certs -> Certs -> Certs
forall a. Semigroup a => a -> a -> a
<> Certs
cs1) Attrs
attrs ExpDec rep
dec) Exp rep
e

-- | A handy shorthand for properties that we usually want for things
-- we stuff into ASTs.
type ASTConstraints a =
  (Eq a, Ord a, Show a, Rename a, Substitute a, FreeIn a, Pretty a)

-- | A type class for operations.
class (TypedOp op) => IsOp op where
  -- | Like 'safeExp', but for arbitrary ops.
  safeOp :: (ASTRep rep) => op rep -> Bool

  -- | Should we try to hoist this out of branches?
  cheapOp :: (ASTRep rep) => op rep -> Bool

  -- | Compute the data dependencies of an operation.
  opDependencies :: (ASTRep rep) => op rep -> [Names]

instance IsOp NoOp where
  safeOp :: forall rep. ASTRep rep => NoOp rep -> Bool
safeOp NoOp rep
NoOp = Bool
True
  cheapOp :: forall rep. ASTRep rep => NoOp rep -> Bool
cheapOp NoOp rep
NoOp = Bool
True
  opDependencies :: forall rep. ASTRep rep => NoOp rep -> [Names]
opDependencies NoOp rep
NoOp = []

-- | Representation-specific attributes; also means the rep supports
-- some basic facilities.
class
  ( RepTypes rep,
    PrettyRep rep,
    Renameable rep,
    Substitutable rep,
    FreeDec (ExpDec rep),
    FreeIn (LetDec rep),
    FreeDec (BodyDec rep),
    FreeIn (FParamInfo rep),
    FreeIn (LParamInfo rep),
    FreeIn (RetType rep),
    FreeIn (BranchType rep),
    ASTConstraints (OpC rep rep),
    IsOp (OpC rep),
    RephraseOp (OpC rep)
  ) =>
  ASTRep rep
  where
  -- | Given a pattern, construct the type of a body that would match
  -- it.  An implementation for many representations would be
  -- 'expExtTypesFromPat'.
  expTypesFromPat ::
    (HasScope rep m, Monad m) =>
    Pat (LetDec rep) ->
    m [BranchType rep]

-- | Construct the type of an expression that would match the pattern.
expExtTypesFromPat :: (Typed dec) => Pat dec -> [ExtType]
expExtTypesFromPat :: forall dec. Typed dec => Pat dec -> [ExtType]
expExtTypesFromPat Pat dec
pat =
  [VName] -> [ExtType] -> [ExtType]
existentialiseExtTypes (Pat dec -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat dec
pat) ([ExtType] -> [ExtType]) -> [ExtType] -> [ExtType]
forall a b. (a -> b) -> a -> b
$
    [TypeBase Shape NoUniqueness] -> [ExtType]
forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
staticShapes ([TypeBase Shape NoUniqueness] -> [ExtType])
-> [TypeBase Shape NoUniqueness] -> [ExtType]
forall a b. (a -> b) -> a -> b
$
      (PatElem dec -> TypeBase Shape NoUniqueness)
-> [PatElem dec] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map PatElem dec -> TypeBase Shape NoUniqueness
forall dec. Typed dec => PatElem dec -> TypeBase Shape NoUniqueness
patElemType ([PatElem dec] -> [TypeBase Shape NoUniqueness])
-> [PatElem dec] -> [TypeBase Shape NoUniqueness]
forall a b. (a -> b) -> a -> b
$
        Pat dec -> [PatElem dec]
forall dec. Pat dec -> [PatElem dec]
patElems Pat dec
pat

-- | Keep only those attributes that are relevant for 'Assert'
-- expressions.
attrsForAssert :: Attrs -> Attrs
attrsForAssert :: Attrs -> Attrs
attrsForAssert (Attrs Set Attr
attrs) =
  Set Attr -> Attrs
Attrs (Set Attr -> Attrs) -> Set Attr -> Attrs
forall a b. (a -> b) -> a -> b
$ (Attr -> Bool) -> Set Attr -> Set Attr
forall a. (a -> Bool) -> Set a -> Set a
S.filter Attr -> Bool
attrForAssert Set Attr
attrs
  where
    attrForAssert :: Attr -> Bool
attrForAssert = (Attr -> Attr -> Bool
forall a. Eq a => a -> a -> Bool
== Name -> [Attr] -> Attr
AttrComp Name
"warn" [Attr
"safety_checks"])

-- | Horizontally fission a lambda that models a binary operator.
lamIsBinOp :: (ASTRep rep) => Lambda rep -> Maybe [(BinOp, PrimType, VName, VName)]
lamIsBinOp :: forall rep.
ASTRep rep =>
Lambda rep -> Maybe [(BinOp, PrimType, VName, VName)]
lamIsBinOp Lambda rep
lam = (SubExpRes -> Maybe (BinOp, PrimType, VName, VName))
-> [SubExpRes] -> Maybe [(BinOp, PrimType, VName, VName)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExpRes -> Maybe (BinOp, PrimType, VName, VName)
splitStm ([SubExpRes] -> Maybe [(BinOp, PrimType, VName, VName)])
-> [SubExpRes] -> Maybe [(BinOp, PrimType, VName, VName)]
forall a b. (a -> b) -> a -> b
$ Body rep -> [SubExpRes]
forall rep. Body rep -> [SubExpRes]
bodyResult (Body rep -> [SubExpRes]) -> Body rep -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda rep -> Body rep
forall rep. Lambda rep -> Body rep
lambdaBody Lambda rep
lam
  where
    n :: Int
n = [TypeBase Shape NoUniqueness] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([TypeBase Shape NoUniqueness] -> Int)
-> [TypeBase Shape NoUniqueness] -> Int
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [TypeBase Shape NoUniqueness]
forall rep. Lambda rep -> [TypeBase Shape NoUniqueness]
lambdaReturnType Lambda rep
lam
    splitStm :: SubExpRes -> Maybe (BinOp, PrimType, VName, VName)
splitStm (SubExpRes Certs
cs (Var VName
res)) = do
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Certs
cs Certs -> Certs -> Bool
forall a. Eq a => a -> a -> Bool
== Certs
forall a. Monoid a => a
mempty
      Let (Pat [PatElem (LetDec rep)
pe]) StmAux (ExpDec rep)
_ (BasicOp (BinOp BinOp
op (Var VName
x) (Var VName
y))) <-
        (Stm rep -> Bool) -> [Stm rep] -> Maybe (Stm rep)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (([VName
res] ==) ([VName] -> Bool) -> (Stm rep -> [VName]) -> Stm rep -> Bool
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) ([Stm rep] -> Maybe (Stm rep)) -> [Stm rep] -> Maybe (Stm rep)
forall a b. (a -> b) -> a -> b
$
          Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList (Stms rep -> [Stm rep]) -> Stms rep -> [Stm rep]
forall a b. (a -> b) -> a -> b
$
            Body rep -> Stms rep
forall rep. Body rep -> Stms rep
bodyStms (Body rep -> Stms rep) -> Body rep -> Stms rep
forall a b. (a -> b) -> a -> b
$
              Lambda rep -> Body rep
forall rep. Lambda rep -> Body rep
lambdaBody Lambda rep
lam
      Int
i <- VName -> SubExp
Var VName
res SubExp -> [SubExp] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
`elemIndex` (SubExpRes -> SubExp) -> [SubExpRes] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp (Body rep -> [SubExpRes]
forall rep. Body rep -> [SubExpRes]
bodyResult (Lambda rep -> Body rep
forall rep. Lambda rep -> Body rep
lambdaBody Lambda rep
lam))
      Param (LParamInfo rep)
xp <- Int -> [Param (LParamInfo rep)] -> Maybe (Param (LParamInfo rep))
forall int a. Integral int => int -> [a] -> Maybe a
maybeNth Int
i ([Param (LParamInfo rep)] -> Maybe (Param (LParamInfo rep)))
-> [Param (LParamInfo rep)] -> Maybe (Param (LParamInfo rep))
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Param (LParamInfo rep)]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam
      Param (LParamInfo rep)
yp <- Int -> [Param (LParamInfo rep)] -> Maybe (Param (LParamInfo rep))
forall int a. Integral int => int -> [a] -> Maybe a
maybeNth (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) ([Param (LParamInfo rep)] -> Maybe (Param (LParamInfo rep)))
-> [Param (LParamInfo rep)] -> Maybe (Param (LParamInfo rep))
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Param (LParamInfo rep)]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName Param (LParamInfo rep)
xp VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
x
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName Param (LParamInfo rep)
yp VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
y
      Prim PrimType
t <- TypeBase Shape NoUniqueness -> Maybe (TypeBase Shape NoUniqueness)
forall a. a -> Maybe a
Just (TypeBase Shape NoUniqueness
 -> Maybe (TypeBase Shape NoUniqueness))
-> TypeBase Shape NoUniqueness
-> Maybe (TypeBase Shape NoUniqueness)
forall a b. (a -> b) -> a -> b
$ PatElem (LetDec rep) -> TypeBase Shape NoUniqueness
forall dec. Typed dec => PatElem dec -> TypeBase Shape NoUniqueness
patElemType PatElem (LetDec rep)
pe
      (BinOp, PrimType, VName, VName)
-> Maybe (BinOp, PrimType, VName, VName)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BinOp
op, PrimType
t, Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName Param (LParamInfo rep)
xp, Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName Param (LParamInfo rep)
yp)
    splitStm SubExpRes
_ = Maybe (BinOp, PrimType, VName, VName)
forall a. Maybe a
Nothing