-- | Particularly simple simplification rules.
module Futhark.Optimise.Simplify.Rules.Simple
  ( TypeLookup,
    VarLookup,
    applySimpleRules,
  )
where

import Control.Monad
import Data.List (isSuffixOf)
import Data.List.NonEmpty qualified as NE
import Futhark.Analysis.PrimExp.Convert
import Futhark.IR
import Futhark.Util (focusNth)

-- | A function that, given a variable name, returns its definition.
type VarLookup rep = VName -> Maybe (Exp rep, Certs)

-- | A function that, given a subexpression, returns its type.
type TypeLookup = SubExp -> Maybe Type

-- | A simple rule is a top-down rule that can be expressed as a pure
-- function.
type SimpleRule rep = VarLookup rep -> TypeLookup -> BasicOp -> Maybe (BasicOp, Certs)

isCt1 :: SubExp -> Bool
isCt1 :: SubExp -> Bool
isCt1 (Constant PrimValue
v) = PrimValue -> Bool
oneIsh PrimValue
v
isCt1 SubExp
_ = Bool
False

isCt0 :: SubExp -> Bool
isCt0 :: SubExp -> Bool
isCt0 (Constant PrimValue
v) = PrimValue -> Bool
zeroIsh PrimValue
v
isCt0 SubExp
_ = Bool
False

simplifyCmpOp :: SimpleRule rep
simplifyCmpOp :: forall rep. SimpleRule rep
simplifyCmpOp VarLookup rep
_ TypeLookup
_ (CmpOp CmpOp
cmp SubExp
e1 SubExp
e2)
  | SubExp
e1 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 = PrimValue -> Maybe (BasicOp, Certs)
constRes (PrimValue -> Maybe (BasicOp, Certs))
-> PrimValue -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$
      Bool -> PrimValue
BoolValue (Bool -> PrimValue) -> Bool -> PrimValue
forall a b. (a -> b) -> a -> b
$
        case CmpOp
cmp of
          CmpEq {} -> Bool
True
          CmpSlt {} -> Bool
False
          CmpUlt {} -> Bool
False
          CmpSle {} -> Bool
True
          CmpUle {} -> Bool
True
          FCmpLt {} -> Bool
False
          FCmpLe {} -> Bool
True
          CmpOp
CmpLlt -> Bool
False
          CmpOp
CmpLle -> Bool
True
simplifyCmpOp VarLookup rep
_ TypeLookup
_ (CmpOp CmpOp
cmp (Constant PrimValue
v1) (Constant PrimValue
v2)) =
  PrimValue -> Maybe (BasicOp, Certs)
constRes (PrimValue -> Maybe (BasicOp, Certs))
-> (Bool -> PrimValue) -> Bool -> Maybe (BasicOp, Certs)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> PrimValue
BoolValue (Bool -> Maybe (BasicOp, Certs))
-> Maybe Bool -> Maybe (BasicOp, Certs)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< CmpOp -> PrimValue -> PrimValue -> Maybe Bool
doCmpOp CmpOp
cmp PrimValue
v1 PrimValue
v2
simplifyCmpOp VarLookup rep
look TypeLookup
_ (CmpOp CmpEq {} (Constant (IntValue IntValue
x)) (Var VName
v))
  | Just (BasicOp (ConvOp BToI {} SubExp
b), Certs
cs) <- VarLookup rep
look VName
v =
      case IntValue -> Int
forall int. Integral int => IntValue -> int
valueIntegral IntValue
x :: Int of
        Int
1 -> (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (SubExp -> BasicOp
SubExp SubExp
b, Certs
cs)
        Int
0 -> (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (UnOp -> SubExp -> BasicOp
UnOp (PrimType -> UnOp
Neg PrimType
Bool) SubExp
b, Certs
cs)
        Int
_ -> (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (SubExp -> BasicOp
SubExp (PrimValue -> SubExp
Constant (Bool -> PrimValue
BoolValue Bool
False)), Certs
cs)
simplifyCmpOp VarLookup rep
_ TypeLookup
_ BasicOp
_ = Maybe (BasicOp, Certs)
forall a. Maybe a
Nothing

simplifyBinOp :: SimpleRule rep
simplifyBinOp :: forall rep. SimpleRule rep
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp BinOp
op (Constant PrimValue
v1) (Constant PrimValue
v2))
  | Just PrimValue
res <- BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
doBinOp BinOp
op PrimValue
v1 PrimValue
v2 =
      PrimValue -> Maybe (BasicOp, Certs)
constRes PrimValue
res
-- By normalisation, constants are always on the left.
--
-- x+(y+z) = (x+y)+z (where x and y are constants).
simplifyBinOp VarLookup rep
look TypeLookup
_ (BinOp BinOp
op1 (Constant PrimValue
x1) (Var VName
y1))
  | BinOp -> Bool
associativeBinOp BinOp
op1,
    Just (BasicOp (BinOp BinOp
op2 (Constant PrimValue
x2) SubExp
y2), Certs
cs) <- VarLookup rep
look VName
y1,
    BinOp
op1 BinOp -> BinOp -> Bool
forall a. Eq a => a -> a -> Bool
== BinOp
op2,
    Just PrimValue
res <- BinOp -> PrimValue -> PrimValue -> Maybe PrimValue
doBinOp BinOp
op1 PrimValue
x1 PrimValue
x2 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (BinOp -> SubExp -> SubExp -> BasicOp
BinOp BinOp
op1 (PrimValue -> SubExp
Constant PrimValue
res) SubExp
y2, Certs
cs)
simplifyBinOp VarLookup rep
look TypeLookup
_ (BinOp (Add IntType
it Overflow
ovf) SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e2
  | SubExp -> Bool
isCt0 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  -- x+(y-x) => y
  | Var VName
v2 <- SubExp
e2,
    Just (BasicOp (BinOp Sub {} SubExp
e2_a SubExp
e2_b), Certs
cs) <- VarLookup rep
look VName
v2,
    SubExp
e2_b SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e1 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (SubExp -> BasicOp
SubExp SubExp
e2_a, Certs
cs)
  -- x+(-1*y) => x-y
  | Var VName
v2 <- SubExp
e2,
    Just (BasicOp (BinOp Mul {} (Constant (IntValue IntValue
x)) SubExp
e3), Certs
cs) <- VarLookup rep
look VName
v2,
    IntValue -> Int
forall int. Integral int => IntValue -> int
valueIntegral IntValue
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== (-Int
1 :: Int) =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> Overflow -> BinOp
Sub IntType
it Overflow
ovf) SubExp
e1 SubExp
e3, Certs
cs)
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp FAdd {} SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e2
  | SubExp -> Bool
isCt0 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
simplifyBinOp VarLookup rep
look TypeLookup
_ (BinOp sub :: BinOp
sub@(Sub IntType
t Overflow
_) SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp
e1 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 = (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (SubExp -> BasicOp
SubExp (IntType -> Integer -> SubExp
intConst IntType
t Integer
0), Certs
forall a. Monoid a => a
mempty)
  --
  -- Below are cases for simplifying (a+b)-b and permutations.
  --
  -- (e1_a+e1_b)-e1_a == e1_b
  | Var VName
v1 <- SubExp
e1,
    Just (BasicOp (BinOp Add {} SubExp
e1_a SubExp
e1_b), Certs
cs) <- VarLookup rep
look VName
v1,
    SubExp
e1_a SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (SubExp -> BasicOp
SubExp SubExp
e1_b, Certs
cs)
  -- (e1_a+e1_b)-e1_b == e1_a
  | Var VName
v1 <- SubExp
e1,
    Just (BasicOp (BinOp Add {} SubExp
e1_a SubExp
e1_b), Certs
cs) <- VarLookup rep
look VName
v1,
    SubExp
e1_b SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (SubExp -> BasicOp
SubExp SubExp
e1_a, Certs
cs)
  -- e2_a-(e2_a+e2_b) == 0-e2_b
  | Var VName
v2 <- SubExp
e2,
    Just (BasicOp (BinOp Add {} SubExp
e2_a SubExp
e2_b), Certs
cs) <- VarLookup rep
look VName
v2,
    SubExp
e2_a SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e1 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (BinOp -> SubExp -> SubExp -> BasicOp
BinOp BinOp
sub (IntType -> Integer -> SubExp
intConst IntType
t Integer
0) SubExp
e2_b, Certs
cs)
  -- e2_b-(e2_a+e2_b) == 0-e2_a
  | Var VName
v2 <- SubExp
e2,
    Just (BasicOp (BinOp Add {} SubExp
e2_a SubExp
e2_b), Certs
cs) <- VarLookup rep
look VName
v2,
    SubExp
e2_b SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e1 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (BinOp -> SubExp -> SubExp -> BasicOp
BinOp BinOp
sub (IntType -> Integer -> SubExp
intConst IntType
t Integer
0) SubExp
e2_a, Certs
cs)
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp FSub {} SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp Mul {} SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp -> Bool
isCt0 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e2
  | SubExp -> Bool
isCt1 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e2
  | SubExp -> Bool
isCt1 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp FMul {} SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt1 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e2
  | SubExp -> Bool
isCt1 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
simplifyBinOp VarLookup rep
look TypeLookup
_ (BinOp (SMod IntType
t Safety
_) SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt1 SubExp
e2 = PrimValue -> Maybe (BasicOp, Certs)
constRes (PrimValue -> Maybe (BasicOp, Certs))
-> PrimValue -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Int -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
t (Int
0 :: Int)
  | SubExp
e1 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 = PrimValue -> Maybe (BasicOp, Certs)
constRes (PrimValue -> Maybe (BasicOp, Certs))
-> PrimValue -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Int -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
t (Int
0 :: Int)
  | Var VName
v1 <- SubExp
e1,
    Just (BasicOp (BinOp SMod {} SubExp
_ SubExp
e4), Certs
v1_cs) <- VarLookup rep
look VName
v1,
    SubExp
e4 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (SubExp -> BasicOp
SubExp SubExp
e1, Certs
v1_cs)
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp SDiv {} SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp -> Bool
isCt1 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp -> Bool
isCt0 SubExp
e2 = Maybe (BasicOp, Certs)
forall a. Maybe a
Nothing
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp SDivUp {} SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp -> Bool
isCt1 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp -> Bool
isCt0 SubExp
e2 = Maybe (BasicOp, Certs)
forall a. Maybe a
Nothing
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp FDiv {} SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp -> Bool
isCt1 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp -> Bool
isCt0 SubExp
e2 = Maybe (BasicOp, Certs)
forall a. Maybe a
Nothing
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp (SRem IntType
t Safety
_) SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt1 SubExp
e2 = PrimValue -> Maybe (BasicOp, Certs)
constRes (PrimValue -> Maybe (BasicOp, Certs))
-> PrimValue -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Int -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
t (Int
0 :: Int)
  | SubExp
e1 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 = PrimValue -> Maybe (BasicOp, Certs)
constRes (PrimValue -> Maybe (BasicOp, Certs))
-> PrimValue -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Int -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
t (Int
1 :: Int)
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp SQuot {} SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt1 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp -> Bool
isCt0 SubExp
e2 = Maybe (BasicOp, Certs)
forall a. Maybe a
Nothing
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp (Pow IntType
t) SubExp
e1 SubExp
e2)
  | SubExp
e1 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== IntType -> Integer -> SubExp
intConst IntType
t Integer
2 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> BinOp
Shl IntType
t) (IntType -> Integer -> SubExp
intConst IntType
t Integer
1) SubExp
e2, Certs
forall a. Monoid a => a
mempty)
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp (FPow FloatType
t) SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp (SubExp -> Maybe (BasicOp, Certs))
-> SubExp -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ FloatType -> Double -> SubExp
floatConst FloatType
t Double
1
  | SubExp -> Bool
isCt0 SubExp
e1 Bool -> Bool -> Bool
|| SubExp -> Bool
isCt1 SubExp
e1 Bool -> Bool -> Bool
|| SubExp -> Bool
isCt1 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp (Shl IntType
t) SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp -> Bool
isCt0 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp (SubExp -> Maybe (BasicOp, Certs))
-> SubExp -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
t Integer
0
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp AShr {} SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp (And IntType
t) SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp (SubExp -> Maybe (BasicOp, Certs))
-> SubExp -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
t Integer
0
  | SubExp -> Bool
isCt0 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp (SubExp -> Maybe (BasicOp, Certs))
-> SubExp -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
t Integer
0
  | SubExp
e1 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp Or {} SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e2
  | SubExp -> Bool
isCt0 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp
e1 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
simplifyBinOp VarLookup rep
_ TypeLookup
_ (BinOp (Xor IntType
t) SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e2
  | SubExp -> Bool
isCt0 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp
e1 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp (SubExp -> Maybe (BasicOp, Certs))
-> SubExp -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
t Integer
0
simplifyBinOp VarLookup rep
defOf TypeLookup
_ (BinOp BinOp
LogAnd SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e1 = PrimValue -> Maybe (BasicOp, Certs)
constRes (PrimValue -> Maybe (BasicOp, Certs))
-> PrimValue -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
False
  | SubExp -> Bool
isCt0 SubExp
e2 = PrimValue -> Maybe (BasicOp, Certs)
constRes (PrimValue -> Maybe (BasicOp, Certs))
-> PrimValue -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
False
  | SubExp -> Bool
isCt1 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e2
  | SubExp -> Bool
isCt1 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp
e1 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | Var VName
v <- SubExp
e1,
    Just (BasicOp (UnOp (Neg PrimType
Bool) SubExp
e1'), Certs
v_cs) <- VarLookup rep
defOf VName
v,
    SubExp
e1' SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
False, Certs
v_cs)
  | Var VName
v <- SubExp
e2,
    Just (BasicOp (UnOp (Neg PrimType
Bool) SubExp
e2'), Certs
v_cs) <- VarLookup rep
defOf VName
v,
    SubExp
e2' SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e1 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
False, Certs
v_cs)
simplifyBinOp VarLookup rep
defOf TypeLookup
_ (BinOp BinOp
LogOr SubExp
e1 SubExp
e2)
  | SubExp -> Bool
isCt0 SubExp
e1 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e2
  | SubExp -> Bool
isCt0 SubExp
e2 = SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | SubExp -> Bool
isCt1 SubExp
e1 = PrimValue -> Maybe (BasicOp, Certs)
constRes (PrimValue -> Maybe (BasicOp, Certs))
-> PrimValue -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
True
  | SubExp -> Bool
isCt1 SubExp
e2 = PrimValue -> Maybe (BasicOp, Certs)
constRes (PrimValue -> Maybe (BasicOp, Certs))
-> PrimValue -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
True
  | Var VName
v <- SubExp
e1,
    Just (BasicOp (UnOp (Neg PrimType
Bool) SubExp
e1'), Certs
v_cs) <- VarLookup rep
defOf VName
v,
    SubExp
e1' SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
True, Certs
v_cs)
  | Var VName
v <- SubExp
e2,
    Just (BasicOp (UnOp (Neg PrimType
Bool) SubExp
e2'), Certs
v_cs) <- VarLookup rep
defOf VName
v,
    SubExp
e2' SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e1 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ Bool -> PrimValue
BoolValue Bool
True, Certs
v_cs)
simplifyBinOp VarLookup rep
defOf TypeLookup
_ (BinOp (SMax IntType
it) SubExp
e1 SubExp
e2)
  | SubExp
e1 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 =
      SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
e1
  | Var VName
v1 <- SubExp
e1,
    Just (BasicOp (BinOp (SMax IntType
_) SubExp
e1_1 SubExp
e1_2), Certs
v1_cs) <- VarLookup rep
defOf VName
v1,
    SubExp
e1_1 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> BinOp
SMax IntType
it) SubExp
e1_2 SubExp
e2, Certs
v1_cs)
  | Var VName
v1 <- SubExp
e1,
    Just (BasicOp (BinOp (SMax IntType
_) SubExp
e1_1 SubExp
e1_2), Certs
v1_cs) <- VarLookup rep
defOf VName
v1,
    SubExp
e1_2 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e2 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> BinOp
SMax IntType
it) SubExp
e1_1 SubExp
e2, Certs
v1_cs)
  | Var VName
v2 <- SubExp
e2,
    Just (BasicOp (BinOp (SMax IntType
_) SubExp
e2_1 SubExp
e2_2), Certs
v2_cs) <- VarLookup rep
defOf VName
v2,
    SubExp
e2_1 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e1 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> BinOp
SMax IntType
it) SubExp
e2_2 SubExp
e1, Certs
v2_cs)
  | Var VName
v2 <- SubExp
e2,
    Just (BasicOp (BinOp (SMax IntType
_) SubExp
e2_1 SubExp
e2_2), Certs
v2_cs) <- VarLookup rep
defOf VName
v2,
    SubExp
e2_2 SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
e1 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (BinOp -> SubExp -> SubExp -> BasicOp
BinOp (IntType -> BinOp
SMax IntType
it) SubExp
e2_1 SubExp
e1, Certs
v2_cs)
simplifyBinOp VarLookup rep
_ TypeLookup
_ BasicOp
_ = Maybe (BasicOp, Certs)
forall a. Maybe a
Nothing

constRes :: PrimValue -> Maybe (BasicOp, Certs)
constRes :: PrimValue -> Maybe (BasicOp, Certs)
constRes = (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just ((BasicOp, Certs) -> Maybe (BasicOp, Certs))
-> (PrimValue -> (BasicOp, Certs))
-> PrimValue
-> Maybe (BasicOp, Certs)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (,Certs
forall a. Monoid a => a
mempty) (BasicOp -> (BasicOp, Certs))
-> (PrimValue -> BasicOp) -> PrimValue -> (BasicOp, Certs)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> BasicOp)
-> (PrimValue -> SubExp) -> PrimValue -> BasicOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> SubExp
Constant

resIsSubExp :: SubExp -> Maybe (BasicOp, Certs)
resIsSubExp :: SubExp -> Maybe (BasicOp, Certs)
resIsSubExp = (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just ((BasicOp, Certs) -> Maybe (BasicOp, Certs))
-> (SubExp -> (BasicOp, Certs)) -> SubExp -> Maybe (BasicOp, Certs)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (,Certs
forall a. Monoid a => a
mempty) (BasicOp -> (BasicOp, Certs))
-> (SubExp -> BasicOp) -> SubExp -> (BasicOp, Certs)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp

simplifyUnOp :: SimpleRule rep
simplifyUnOp :: forall rep. SimpleRule rep
simplifyUnOp VarLookup rep
_ TypeLookup
_ (UnOp UnOp
op (Constant PrimValue
v)) =
  PrimValue -> Maybe (BasicOp, Certs)
constRes (PrimValue -> Maybe (BasicOp, Certs))
-> Maybe PrimValue -> Maybe (BasicOp, Certs)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< UnOp -> PrimValue -> Maybe PrimValue
doUnOp UnOp
op PrimValue
v
simplifyUnOp VarLookup rep
defOf TypeLookup
_ (UnOp (Neg PrimType
Bool) (Var VName
v))
  | Just (BasicOp (UnOp (Neg PrimType
Bool) SubExp
v2), Certs
v_cs) <- VarLookup rep
defOf VName
v =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (SubExp -> BasicOp
SubExp SubExp
v2, Certs
v_cs)
simplifyUnOp VarLookup rep
_ TypeLookup
_ BasicOp
_ =
  Maybe (BasicOp, Certs)
forall a. Maybe a
Nothing

simplifyConvOp :: SimpleRule rep
simplifyConvOp :: forall rep. SimpleRule rep
simplifyConvOp VarLookup rep
_ TypeLookup
_ (ConvOp ConvOp
op (Constant PrimValue
v)) =
  PrimValue -> Maybe (BasicOp, Certs)
constRes (PrimValue -> Maybe (BasicOp, Certs))
-> Maybe PrimValue -> Maybe (BasicOp, Certs)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ConvOp -> PrimValue -> Maybe PrimValue
doConvOp ConvOp
op PrimValue
v
simplifyConvOp VarLookup rep
_ TypeLookup
_ (ConvOp ConvOp
op SubExp
se)
  | (PrimType
from, PrimType
to) <- ConvOp -> (PrimType, PrimType)
convOpType ConvOp
op,
    PrimType
from PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
to =
      SubExp -> Maybe (BasicOp, Certs)
resIsSubExp SubExp
se
simplifyConvOp VarLookup rep
lookupVar TypeLookup
_ (ConvOp (SExt IntType
t2 IntType
t1) (Var VName
v))
  | Just (BasicOp (ConvOp (SExt IntType
t3 IntType
_) SubExp
se), Certs
v_cs) <- VarLookup rep
lookupVar VName
v,
    IntType
t2 IntType -> IntType -> Bool
forall a. Ord a => a -> a -> Bool
>= IntType
t3 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (ConvOp -> SubExp -> BasicOp
ConvOp (IntType -> IntType -> ConvOp
SExt IntType
t3 IntType
t1) SubExp
se, Certs
v_cs)
simplifyConvOp VarLookup rep
lookupVar TypeLookup
_ (ConvOp (ZExt IntType
t2 IntType
t1) (Var VName
v))
  | Just (BasicOp (ConvOp (ZExt IntType
t3 IntType
_) SubExp
se), Certs
v_cs) <- VarLookup rep
lookupVar VName
v,
    IntType
t2 IntType -> IntType -> Bool
forall a. Ord a => a -> a -> Bool
>= IntType
t3 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (ConvOp -> SubExp -> BasicOp
ConvOp (IntType -> IntType -> ConvOp
ZExt IntType
t3 IntType
t1) SubExp
se, Certs
v_cs)
simplifyConvOp VarLookup rep
lookupVar TypeLookup
_ (ConvOp (SIToFP IntType
t2 FloatType
t1) (Var VName
v))
  | Just (BasicOp (ConvOp (SExt IntType
t3 IntType
_) SubExp
se), Certs
v_cs) <- VarLookup rep
lookupVar VName
v,
    IntType
t2 IntType -> IntType -> Bool
forall a. Ord a => a -> a -> Bool
>= IntType
t3 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (ConvOp -> SubExp -> BasicOp
ConvOp (IntType -> FloatType -> ConvOp
SIToFP IntType
t3 FloatType
t1) SubExp
se, Certs
v_cs)
simplifyConvOp VarLookup rep
lookupVar TypeLookup
_ (ConvOp (UIToFP IntType
t2 FloatType
t1) (Var VName
v))
  | Just (BasicOp (ConvOp (ZExt IntType
t3 IntType
_) SubExp
se), Certs
v_cs) <- VarLookup rep
lookupVar VName
v,
    IntType
t2 IntType -> IntType -> Bool
forall a. Ord a => a -> a -> Bool
>= IntType
t3 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (ConvOp -> SubExp -> BasicOp
ConvOp (IntType -> FloatType -> ConvOp
UIToFP IntType
t3 FloatType
t1) SubExp
se, Certs
v_cs)
simplifyConvOp VarLookup rep
lookupVar TypeLookup
_ (ConvOp (FPConv FloatType
t2 FloatType
t1) (Var VName
v))
  | Just (BasicOp (ConvOp (FPConv FloatType
t3 FloatType
_) SubExp
se), Certs
v_cs) <- VarLookup rep
lookupVar VName
v,
    FloatType
t2 FloatType -> FloatType -> Bool
forall a. Ord a => a -> a -> Bool
>= FloatType
t3 =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (ConvOp -> SubExp -> BasicOp
ConvOp (FloatType -> FloatType -> ConvOp
FPConv FloatType
t3 FloatType
t1) SubExp
se, Certs
v_cs)
simplifyConvOp VarLookup rep
_ TypeLookup
_ BasicOp
_ =
  Maybe (BasicOp, Certs)
forall a. Maybe a
Nothing

-- If expression is true then just replace assertion.
simplifyAssert :: SimpleRule rep
simplifyAssert :: forall rep. SimpleRule rep
simplifyAssert VarLookup rep
_ TypeLookup
_ (Assert (Constant (BoolValue Bool
True)) ErrorMsg SubExp
_ (SrcLoc, [SrcLoc])
_) =
  PrimValue -> Maybe (BasicOp, Certs)
constRes PrimValue
UnitValue
simplifyAssert VarLookup rep
_ TypeLookup
_ BasicOp
_ =
  Maybe (BasicOp, Certs)
forall a. Maybe a
Nothing

simplifyReshape :: SimpleRule rep
simplifyReshape :: forall rep. SimpleRule rep
simplifyReshape VarLookup rep
defOf TypeLookup
seType (Reshape VName
v NewShape SubExp
newshape)
  -- No-op reshape
  | Just Type
t <- TypeLookup
seType TypeLookup -> TypeLookup
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v,
    NewShape SubExp -> ShapeBase SubExp
forall d. NewShape d -> ShapeBase d
newShape NewShape SubExp
newshape ShapeBase SubExp -> ShapeBase SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== Type -> ShapeBase SubExp
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape Type
t =
      SubExp -> Maybe (BasicOp, Certs)
resIsSubExp (SubExp -> Maybe (BasicOp, Certs))
-> SubExp -> Maybe (BasicOp, Certs)
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v
  -- Simplifying a reshape
  | Just ShapeBase SubExp
shape <- Type -> ShapeBase SubExp
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape (Type -> ShapeBase SubExp)
-> Maybe Type -> Maybe (ShapeBase SubExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeLookup
seType (VName -> SubExp
Var VName
v),
    Just NewShape SubExp
newshape' <- ShapeBase SubExp -> NewShape SubExp -> Maybe (NewShape SubExp)
forall d. Eq d => ShapeBase d -> NewShape d -> Maybe (NewShape d)
simplifyNewShape ShapeBase SubExp
shape NewShape SubExp
newshape =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (VName -> NewShape SubExp -> BasicOp
Reshape VName
v NewShape SubExp
newshape', Certs
forall a. Monoid a => a
mempty)
  -- Reshape-of-reshape
  | Just (BasicOp (Reshape VName
v2 NewShape SubExp
oldnewshape), Certs
v_cs) <- VarLookup rep
defOf VName
v =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (VName -> NewShape SubExp -> BasicOp
Reshape VName
v2 (NewShape SubExp
oldnewshape NewShape SubExp -> NewShape SubExp -> NewShape SubExp
forall a. Semigroup a => a -> a -> a
<> NewShape SubExp
newshape), Certs
v_cs)
  -- Reshape-of-scratch
  | Just (BasicOp (Scratch PrimType
bt [SubExp]
_), Certs
v_cs) <- VarLookup rep
defOf VName
v =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (PrimType -> [SubExp] -> BasicOp
Scratch PrimType
bt ([SubExp] -> BasicOp) -> [SubExp] -> BasicOp
forall a b. (a -> b) -> a -> b
$ ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims (ShapeBase SubExp -> [SubExp]) -> ShapeBase SubExp -> [SubExp]
forall a b. (a -> b) -> a -> b
$ NewShape SubExp -> ShapeBase SubExp
forall d. NewShape d -> ShapeBase d
newShape NewShape SubExp
newshape, Certs
v_cs)
  -- Reshape-of-replicate
  | Just (BasicOp (Replicate ShapeBase SubExp
_ SubExp
se), Certs
v_cs) <- VarLookup rep
defOf VName
v,
    Just ShapeBase SubExp
oldshape <- Type -> ShapeBase SubExp
forall shape u. ArrayShape shape => TypeBase shape u -> shape
arrayShape (Type -> ShapeBase SubExp)
-> Maybe Type -> Maybe (ShapeBase SubExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeLookup
seType SubExp
se,
    ShapeBase SubExp
newshape' <- NewShape SubExp -> ShapeBase SubExp
forall d. NewShape d -> ShapeBase d
newShape NewShape SubExp
newshape,
    ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
oldshape [SubExp] -> [SubExp] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isSuffixOf` ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
newshape' =
      let new :: [SubExp]
new =
            Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take (ShapeBase SubExp -> Int
forall a. ArrayShape a => a -> Int
shapeRank ShapeBase SubExp
newshape' Int -> Int -> Int
forall a. Num a => a -> a -> a
- ShapeBase SubExp -> Int
forall a. ArrayShape a => a -> Int
shapeRank ShapeBase SubExp
oldshape) ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$
              ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
newshape'
       in (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (ShapeBase SubExp -> SubExp -> BasicOp
Replicate ([SubExp] -> ShapeBase SubExp
forall d. [d] -> ShapeBase d
Shape [SubExp]
new) SubExp
se, Certs
v_cs)
  -- Reshape-of-iota
  | Just (BasicOp (Iota SubExp
_ SubExp
offset SubExp
stride IntType
it), Certs
v_cs) <- VarLookup rep
defOf VName
v,
    [SubExp
n] <- ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims (NewShape SubExp -> ShapeBase SubExp
forall d. NewShape d -> ShapeBase d
newShape NewShape SubExp
newshape) =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (SubExp -> SubExp -> SubExp -> IntType -> BasicOp
Iota SubExp
n SubExp
offset SubExp
stride IntType
it, Certs
v_cs)
  -- Reshape-of-concat
  | ReshapeKind
ReshapeCoerce <- NewShape SubExp -> ReshapeKind
reshapeKind NewShape SubExp
newshape = do
      (BasicOp (Concat Int
d NonEmpty VName
arrs SubExp
_), Certs
v_cs) <- VarLookup rep
defOf VName
v
      ([SubExp]
bef, SubExp
w', [SubExp]
aft) <- Int -> [SubExp] -> Maybe ([SubExp], SubExp, [SubExp])
forall int a. Integral int => int -> [a] -> Maybe ([a], a, [a])
focusNth Int
d ([SubExp] -> Maybe ([SubExp], SubExp, [SubExp]))
-> [SubExp] -> Maybe ([SubExp], SubExp, [SubExp])
forall a b. (a -> b) -> a -> b
$ ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims (ShapeBase SubExp -> [SubExp]) -> ShapeBase SubExp -> [SubExp]
forall a b. (a -> b) -> a -> b
$ NewShape SubExp -> ShapeBase SubExp
forall d. NewShape d -> ShapeBase d
newShape NewShape SubExp
newshape
      ([SubExp]
arr_bef, SubExp
_, [SubExp]
arr_aft) <-
        Int -> [SubExp] -> Maybe ([SubExp], SubExp, [SubExp])
forall int a. Integral int => int -> [a] -> Maybe ([a], a, [a])
focusNth Int
d ([SubExp] -> Maybe ([SubExp], SubExp, [SubExp]))
-> (Maybe Type -> Maybe [SubExp])
-> Maybe Type
-> Maybe ([SubExp], SubExp, [SubExp])
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< (Type -> [SubExp]) -> Maybe Type -> Maybe [SubExp]
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> [SubExp]
forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
arrayDims (Maybe Type -> Maybe ([SubExp], SubExp, [SubExp]))
-> Maybe Type -> Maybe ([SubExp], SubExp, [SubExp])
forall a b. (a -> b) -> a -> b
$ TypeLookup
seType TypeLookup -> TypeLookup
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ NonEmpty VName -> VName
forall a. NonEmpty a -> a
NE.head NonEmpty VName
arrs
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ [SubExp]
arr_bef [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== [SubExp]
bef
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ [SubExp]
arr_aft [SubExp] -> [SubExp] -> Bool
forall a. Eq a => a -> a -> Bool
== [SubExp]
aft
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (Int -> NonEmpty VName -> SubExp -> BasicOp
Concat Int
d NonEmpty VName
arrs SubExp
w', Certs
v_cs)
simplifyReshape VarLookup rep
_ TypeLookup
_ BasicOp
_ = Maybe (BasicOp, Certs)
forall a. Maybe a
Nothing

reshapeSlice :: [DimIndex d] -> [d] -> [DimIndex d]
reshapeSlice :: forall d. [DimIndex d] -> [d] -> [DimIndex d]
reshapeSlice (DimFix d
i : [DimIndex d]
slice') [d]
scs =
  d -> DimIndex d
forall d. d -> DimIndex d
DimFix d
i DimIndex d -> [DimIndex d] -> [DimIndex d]
forall a. a -> [a] -> [a]
: [DimIndex d] -> [d] -> [DimIndex d]
forall d. [DimIndex d] -> [d] -> [DimIndex d]
reshapeSlice [DimIndex d]
slice' [d]
scs
reshapeSlice (DimSlice d
x d
_ d
s : [DimIndex d]
slice') (d
d : [d]
ds') =
  d -> d -> d -> DimIndex d
forall d. d -> d -> d -> DimIndex d
DimSlice d
x d
d d
s DimIndex d -> [DimIndex d] -> [DimIndex d]
forall a. a -> [a] -> [a]
: [DimIndex d] -> [d] -> [DimIndex d]
forall d. [DimIndex d] -> [d] -> [DimIndex d]
reshapeSlice [DimIndex d]
slice' [d]
ds'
reshapeSlice [DimIndex d]
_ [d]
_ = []

-- If we are size-coercing a slice, then we might as well just use a
-- different slice instead.
simplifyReshapeIndex :: SimpleRule rep
simplifyReshapeIndex :: forall rep. SimpleRule rep
simplifyReshapeIndex VarLookup rep
defOf TypeLookup
_ (Reshape VName
v NewShape SubExp
newshape)
  | ReshapeKind
ReshapeCoerce <- NewShape SubExp -> ReshapeKind
reshapeKind NewShape SubExp
newshape,
    Just (BasicOp (Index VName
v' Slice SubExp
slice), Certs
v_cs) <- VarLookup rep
defOf VName
v,
    Slice SubExp
slice' <- [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice ([DimIndex SubExp] -> Slice SubExp)
-> [DimIndex SubExp] -> Slice SubExp
forall a b. (a -> b) -> a -> b
$ [DimIndex SubExp] -> [SubExp] -> [DimIndex SubExp]
forall d. [DimIndex d] -> [d] -> [DimIndex d]
reshapeSlice (Slice SubExp -> [DimIndex SubExp]
forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slice) ([SubExp] -> [DimIndex SubExp]) -> [SubExp] -> [DimIndex SubExp]
forall a b. (a -> b) -> a -> b
$ ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims (ShapeBase SubExp -> [SubExp]) -> ShapeBase SubExp -> [SubExp]
forall a b. (a -> b) -> a -> b
$ NewShape SubExp -> ShapeBase SubExp
forall d. NewShape d -> ShapeBase d
newShape NewShape SubExp
newshape,
    Slice SubExp
slice' Slice SubExp -> Slice SubExp -> Bool
forall a. Eq a => a -> a -> Bool
/= Slice SubExp
slice =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (VName -> Slice SubExp -> BasicOp
Index VName
v' Slice SubExp
slice', Certs
v_cs)
simplifyReshapeIndex VarLookup rep
_ TypeLookup
_ BasicOp
_ = Maybe (BasicOp, Certs)
forall a. Maybe a
Nothing

-- If we are updating a slice with the result of a size coercion, we
-- instead use the original array and update the slice dimensions.
simplifyUpdateReshape :: SimpleRule rep
simplifyUpdateReshape :: forall rep. SimpleRule rep
simplifyUpdateReshape VarLookup rep
defOf TypeLookup
seType (Update Safety
safety VName
dest Slice SubExp
slice (Var VName
v))
  | Just (BasicOp (Reshape VName
v' NewShape SubExp
newshape), Certs
v_cs) <- VarLookup rep
defOf VName
v,
    ReshapeKind
ReshapeCoerce <- NewShape SubExp -> ReshapeKind
reshapeKind NewShape SubExp
newshape,
    Just [SubExp]
ds <- Type -> [SubExp]
forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
arrayDims (Type -> [SubExp]) -> Maybe Type -> Maybe [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeLookup
seType (VName -> SubExp
Var VName
v'),
    Slice SubExp
slice' <- [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice ([DimIndex SubExp] -> Slice SubExp)
-> [DimIndex SubExp] -> Slice SubExp
forall a b. (a -> b) -> a -> b
$ [DimIndex SubExp] -> [SubExp] -> [DimIndex SubExp]
forall d. [DimIndex d] -> [d] -> [DimIndex d]
reshapeSlice (Slice SubExp -> [DimIndex SubExp]
forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slice) [SubExp]
ds,
    Slice SubExp
slice' Slice SubExp -> Slice SubExp -> Bool
forall a. Eq a => a -> a -> Bool
/= Slice SubExp
slice =
      (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
Just (Safety -> VName -> Slice SubExp -> SubExp -> BasicOp
Update Safety
safety VName
dest Slice SubExp
slice' (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v', Certs
v_cs)
simplifyUpdateReshape VarLookup rep
_ TypeLookup
_ BasicOp
_ = Maybe (BasicOp, Certs)
forall a. Maybe a
Nothing

-- | If we are replicating a scratch array (possibly indirectly), just
-- turn it into a scratch by itself.
repScratchToScratch :: SimpleRule rep
repScratchToScratch :: forall rep. SimpleRule rep
repScratchToScratch VarLookup rep
defOf TypeLookup
seType (Replicate ShapeBase SubExp
shape (Var VName
src)) = do
  Type
t <- TypeLookup
seType TypeLookup -> TypeLookup
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
src
  Certs
cs <- VName -> Maybe Certs
isActuallyScratch VName
src
  (BasicOp, Certs) -> Maybe (BasicOp, Certs)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimType -> [SubExp] -> BasicOp
Scratch (Type -> PrimType
forall shape u. TypeBase shape u -> PrimType
elemType Type
t) (ShapeBase SubExp -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims ShapeBase SubExp
shape [SubExp] -> [SubExp] -> [SubExp]
forall a. Semigroup a => a -> a -> a
<> Type -> [SubExp]
forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
arrayDims Type
t), Certs
cs)
  where
    isActuallyScratch :: VName -> Maybe Certs
isActuallyScratch VName
v =
      case VarLookup rep
defOf VName
v of
        Just (BasicOp Scratch {}, Certs
cs) ->
          Certs -> Maybe Certs
forall a. a -> Maybe a
Just Certs
cs
        Just (BasicOp (Rearrange VName
v' [Int]
_), Certs
cs) ->
          (Certs
cs <>) (Certs -> Certs) -> Maybe Certs -> Maybe Certs
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Maybe Certs
isActuallyScratch VName
v'
        Just (BasicOp (Reshape VName
v' NewShape SubExp
_), Certs
cs) ->
          (Certs
cs <>) (Certs -> Certs) -> Maybe Certs -> Maybe Certs
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Maybe Certs
isActuallyScratch VName
v'
        Maybe (Exp rep, Certs)
_ -> Maybe Certs
forall a. Maybe a
Nothing
repScratchToScratch VarLookup rep
_ TypeLookup
_ BasicOp
_ =
  Maybe (BasicOp, Certs)
forall a. Maybe a
Nothing

simpleRules :: [SimpleRule rep]
simpleRules :: forall rep. [SimpleRule rep]
simpleRules =
  [ SimpleRule rep
forall rep. SimpleRule rep
simplifyBinOp,
    SimpleRule rep
forall rep. SimpleRule rep
simplifyCmpOp,
    SimpleRule rep
forall rep. SimpleRule rep
simplifyUnOp,
    SimpleRule rep
forall rep. SimpleRule rep
simplifyConvOp,
    SimpleRule rep
forall rep. SimpleRule rep
simplifyAssert,
    SimpleRule rep
forall rep. SimpleRule rep
repScratchToScratch,
    SimpleRule rep
forall rep. SimpleRule rep
simplifyReshape,
    SimpleRule rep
forall rep. SimpleRule rep
simplifyReshapeIndex,
    SimpleRule rep
forall rep. SimpleRule rep
simplifyUpdateReshape
  ]

-- | Try to simplify the given t'BasicOp', returning a new t'BasicOp'
-- and certificates that it must depend on.
{-# NOINLINE applySimpleRules #-}
applySimpleRules ::
  VarLookup rep ->
  TypeLookup ->
  BasicOp ->
  Maybe (BasicOp, Certs)
applySimpleRules :: forall rep. SimpleRule rep
applySimpleRules VarLookup rep
defOf TypeLookup
seType BasicOp
op =
  [Maybe (BasicOp, Certs)] -> Maybe (BasicOp, Certs)
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum [SimpleRule rep
rule VarLookup rep
defOf TypeLookup
seType BasicOp
op | SimpleRule rep
rule <- [SimpleRule rep]
forall rep. [SimpleRule rep]
simpleRules]