{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
-- | Utility declarations for performing range analysis.  The ranges
-- computed here are /local/ (does not take range of subexpressions
-- into account), which is probably not very interesting.  See
-- "Futhark.Analysis.Range" for a more comprehensive analysis built on
-- these building blocks.
module Futhark.IR.Prop.Ranges
       ( Bound
       , KnownBound (..)
       , boundToScalExp
       , minimumBound
       , maximumBound
       , Range
       , unknownRange
       , ScalExpRange
       , Ranged
       , RangeOf (..)
       , RangesOf (..)
       , expRanges
       , RangedOp (..)
       , CanBeRanged (..)
       )
       where

import qualified Data.Kind
import qualified Data.Map.Strict as M

import Futhark.IR.Prop
import Futhark.IR.Syntax
import qualified Futhark.Analysis.ScalExp as SE
import qualified Futhark.Analysis.AlgSimplify as AS
import Futhark.Transform.Substitute
import Futhark.Transform.Rename
import qualified Futhark.Util.Pretty as PP

-- | A known bound on a value.
data KnownBound = VarBound VName
                  -- ^ Has the same bounds as this variable.  VERY
                  -- IMPORTANT: this variable may be an array, so it
                  -- cannot be immediately translated to a 'SE.ScalExp'.
                | MinimumBound KnownBound KnownBound
                  -- ^ Bounded by the minimum of these two bounds.
                | MaximumBound KnownBound KnownBound
                  -- ^ Bounded by the maximum of these two bounds.
                | ScalarBound SE.ScalExp
                  -- ^ Bounded by this scalar expression.
                deriving (KnownBound -> KnownBound -> Bool
(KnownBound -> KnownBound -> Bool)
-> (KnownBound -> KnownBound -> Bool) -> Eq KnownBound
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: KnownBound -> KnownBound -> Bool
$c/= :: KnownBound -> KnownBound -> Bool
== :: KnownBound -> KnownBound -> Bool
$c== :: KnownBound -> KnownBound -> Bool
Eq, Eq KnownBound
Eq KnownBound
-> (KnownBound -> KnownBound -> Ordering)
-> (KnownBound -> KnownBound -> Bool)
-> (KnownBound -> KnownBound -> Bool)
-> (KnownBound -> KnownBound -> Bool)
-> (KnownBound -> KnownBound -> Bool)
-> (KnownBound -> KnownBound -> KnownBound)
-> (KnownBound -> KnownBound -> KnownBound)
-> Ord KnownBound
KnownBound -> KnownBound -> Bool
KnownBound -> KnownBound -> Ordering
KnownBound -> KnownBound -> KnownBound
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 :: KnownBound -> KnownBound -> KnownBound
$cmin :: KnownBound -> KnownBound -> KnownBound
max :: KnownBound -> KnownBound -> KnownBound
$cmax :: KnownBound -> KnownBound -> KnownBound
>= :: KnownBound -> KnownBound -> Bool
$c>= :: KnownBound -> KnownBound -> Bool
> :: KnownBound -> KnownBound -> Bool
$c> :: KnownBound -> KnownBound -> Bool
<= :: KnownBound -> KnownBound -> Bool
$c<= :: KnownBound -> KnownBound -> Bool
< :: KnownBound -> KnownBound -> Bool
$c< :: KnownBound -> KnownBound -> Bool
compare :: KnownBound -> KnownBound -> Ordering
$ccompare :: KnownBound -> KnownBound -> Ordering
$cp1Ord :: Eq KnownBound
Ord, Int -> KnownBound -> ShowS
[KnownBound] -> ShowS
KnownBound -> String
(Int -> KnownBound -> ShowS)
-> (KnownBound -> String)
-> ([KnownBound] -> ShowS)
-> Show KnownBound
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [KnownBound] -> ShowS
$cshowList :: [KnownBound] -> ShowS
show :: KnownBound -> String
$cshow :: KnownBound -> String
showsPrec :: Int -> KnownBound -> ShowS
$cshowsPrec :: Int -> KnownBound -> ShowS
Show)

instance Substitute KnownBound where
  substituteNames :: Map VName VName -> KnownBound -> KnownBound
substituteNames Map VName VName
substs (VarBound VName
name) =
    VName -> KnownBound
VarBound (VName -> KnownBound) -> VName -> KnownBound
forall a b. (a -> b) -> a -> b
$ Map VName VName -> VName -> VName
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs VName
name
  substituteNames Map VName VName
substs (MinimumBound KnownBound
b1 KnownBound
b2) =
    KnownBound -> KnownBound -> KnownBound
MinimumBound (Map VName VName -> KnownBound -> KnownBound
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs KnownBound
b1) (Map VName VName -> KnownBound -> KnownBound
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs KnownBound
b2)
  substituteNames Map VName VName
substs (MaximumBound KnownBound
b1 KnownBound
b2) =
    KnownBound -> KnownBound -> KnownBound
MaximumBound (Map VName VName -> KnownBound -> KnownBound
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs KnownBound
b1) (Map VName VName -> KnownBound -> KnownBound
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs KnownBound
b2)
  substituteNames Map VName VName
substs (ScalarBound ScalExp
se) =
    ScalExp -> KnownBound
ScalarBound (ScalExp -> KnownBound) -> ScalExp -> KnownBound
forall a b. (a -> b) -> a -> b
$ Map VName VName -> ScalExp -> ScalExp
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs ScalExp
se

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

instance FreeIn KnownBound where
  freeIn' :: KnownBound -> FV
freeIn' (VarBound VName
v)         = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
v
  freeIn' (MinimumBound KnownBound
b1 KnownBound
b2) = KnownBound -> FV
forall a. FreeIn a => a -> FV
freeIn' KnownBound
b1 FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> KnownBound -> FV
forall a. FreeIn a => a -> FV
freeIn' KnownBound
b2
  freeIn' (MaximumBound KnownBound
b1 KnownBound
b2) = KnownBound -> FV
forall a. FreeIn a => a -> FV
freeIn' KnownBound
b1 FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> KnownBound -> FV
forall a. FreeIn a => a -> FV
freeIn' KnownBound
b2
  freeIn' (ScalarBound ScalExp
e)      = ScalExp -> FV
forall a. FreeIn a => a -> FV
freeIn' ScalExp
e

instance FreeDec KnownBound where
  precomputed :: KnownBound -> FV -> FV
precomputed KnownBound
_ = FV -> FV
forall a. a -> a
id

instance PP.Pretty KnownBound where
  ppr :: KnownBound -> Doc
ppr (VarBound VName
v) =
    String -> Doc
PP.text String
"variable " Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> VName -> Doc
forall a. Pretty a => a -> Doc
PP.ppr VName
v
  ppr (MinimumBound KnownBound
b1 KnownBound
b2) =
    String -> Doc
PP.text String
"min" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
PP.parens (KnownBound -> Doc
forall a. Pretty a => a -> Doc
PP.ppr KnownBound
b1 Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
PP.comma Doc -> Doc -> Doc
PP.<+> KnownBound -> Doc
forall a. Pretty a => a -> Doc
PP.ppr KnownBound
b2)
  ppr (MaximumBound KnownBound
b1 KnownBound
b2) =
    String -> Doc
PP.text String
"max" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
PP.parens (KnownBound -> Doc
forall a. Pretty a => a -> Doc
PP.ppr KnownBound
b1 Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
PP.comma Doc -> Doc -> Doc
PP.<+> KnownBound -> Doc
forall a. Pretty a => a -> Doc
PP.ppr KnownBound
b2)
  ppr (ScalarBound ScalExp
e) =
    ScalExp -> Doc
forall a. Pretty a => a -> Doc
PP.ppr ScalExp
e

-- | Convert the bound to a scalar expression if possible.  This is
-- possible for all bounds that do not contain 'VarBound's.
boundToScalExp :: KnownBound -> Maybe SE.ScalExp
boundToScalExp :: KnownBound -> Maybe ScalExp
boundToScalExp (VarBound VName
_) = Maybe ScalExp
forall a. Maybe a
Nothing
boundToScalExp (ScalarBound ScalExp
se) = ScalExp -> Maybe ScalExp
forall a. a -> Maybe a
Just ScalExp
se
boundToScalExp (MinimumBound KnownBound
b1 KnownBound
b2) = do
  ScalExp
b1' <- KnownBound -> Maybe ScalExp
boundToScalExp KnownBound
b1
  ScalExp
b2' <- KnownBound -> Maybe ScalExp
boundToScalExp KnownBound
b2
  ScalExp -> Maybe ScalExp
forall (m :: * -> *) a. Monad m => a -> m a
return (ScalExp -> Maybe ScalExp) -> ScalExp -> Maybe ScalExp
forall a b. (a -> b) -> a -> b
$ Bool -> [ScalExp] -> ScalExp
SE.MaxMin Bool
True [ScalExp
b1', ScalExp
b2']
boundToScalExp (MaximumBound KnownBound
b1 KnownBound
b2) = do
  ScalExp
b1' <- KnownBound -> Maybe ScalExp
boundToScalExp KnownBound
b1
  ScalExp
b2' <- KnownBound -> Maybe ScalExp
boundToScalExp KnownBound
b2
  ScalExp -> Maybe ScalExp
forall (m :: * -> *) a. Monad m => a -> m a
return (ScalExp -> Maybe ScalExp) -> ScalExp -> Maybe ScalExp
forall a b. (a -> b) -> a -> b
$ Bool -> [ScalExp] -> ScalExp
SE.MaxMin Bool
False [ScalExp
b1', ScalExp
b2']

-- | A possibly undefined bound on a value.
type Bound = Maybe KnownBound

-- | Construct a 'MinimumBound' from two possibly known bounds.  The
-- resulting bound will be unknown unless both of the given 'Bound's
-- are known.  This may seem counterintuitive, but it actually makes
-- sense when you consider the task of combining the lower bounds for
-- two different flows of execution (like an @if@ expression).  If we
-- only have knowledge about one of the branches, this means that we
-- have no useful information about the combined lower bound, as the
-- other branch may take any value.
minimumBound :: Bound -> Bound -> Bound
minimumBound :: Bound -> Bound -> Bound
minimumBound (Just KnownBound
x)  (Just KnownBound
y) = KnownBound -> Bound
forall a. a -> Maybe a
Just (KnownBound -> Bound) -> KnownBound -> Bound
forall a b. (a -> b) -> a -> b
$ KnownBound -> KnownBound -> KnownBound
MinimumBound KnownBound
x KnownBound
y
minimumBound Bound
_         Bound
_        = Bound
forall a. Maybe a
Nothing

-- | Like 'minimumBound', but constructs a 'MaximumBound'.
maximumBound :: Bound -> Bound -> Bound
maximumBound :: Bound -> Bound -> Bound
maximumBound (Just KnownBound
x)  (Just KnownBound
y) = KnownBound -> Bound
forall a. a -> Maybe a
Just (KnownBound -> Bound) -> KnownBound -> Bound
forall a b. (a -> b) -> a -> b
$ KnownBound -> KnownBound -> KnownBound
MaximumBound KnownBound
x KnownBound
y
maximumBound Bound
_         Bound
_        = Bound
forall a. Maybe a
Nothing

-- | Upper and lower bound, both inclusive.
type Range = (Bound, Bound)

-- | A range in which both upper and lower bounds are 'Nothing.
unknownRange :: Range
unknownRange :: Range
unknownRange = (Bound
forall a. Maybe a
Nothing, Bound
forall a. Maybe a
Nothing)

-- | The range as a pair of scalar expressions.
type ScalExpRange = (Maybe SE.ScalExp, Maybe SE.ScalExp)

-- | The lore has embedded range information.  Note that it may not be
-- up to date, unless whatever maintains the syntax tree is careful.
type Ranged lore = (ASTLore lore,
                    RangedOp (Op lore),
                    RangeOf (LetDec lore),
                    RangesOf (BodyDec lore))

-- | Something that contains range information.
class RangeOf a where
  -- | The range of the argument element.
  rangeOf :: a -> Range

instance RangeOf Range where
  rangeOf :: Range -> Range
rangeOf = Range -> Range
forall a. a -> a
id

instance RangeOf dec => RangeOf (PatElemT dec) where
  rangeOf :: PatElemT dec -> Range
rangeOf = dec -> Range
forall a. RangeOf a => a -> Range
rangeOf (dec -> Range) -> (PatElemT dec -> dec) -> PatElemT dec -> Range
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatElemT dec -> dec
forall dec. PatElemT dec -> dec
patElemDec

instance RangeOf SubExp where
  rangeOf :: SubExp -> Range
rangeOf SubExp
se = (KnownBound -> Bound
forall a. a -> Maybe a
Just KnownBound
lower, KnownBound -> Bound
forall a. a -> Maybe a
Just KnownBound
upper)
    where (KnownBound
lower, KnownBound
upper) = SubExp -> (KnownBound, KnownBound)
subExpKnownRange SubExp
se

-- | Something that contains range information for several things,
-- most notably t'Body' and t'Pattern'.
class RangesOf a where
  -- | The ranges of the argument.
  rangesOf :: a -> [Range]

instance RangeOf a => RangesOf [a] where
  rangesOf :: [a] -> [Range]
rangesOf = (a -> Range) -> [a] -> [Range]
forall a b. (a -> b) -> [a] -> [b]
map a -> Range
forall a. RangeOf a => a -> Range
rangeOf

instance RangeOf dec => RangesOf (PatternT dec) where
  rangesOf :: PatternT dec -> [Range]
rangesOf = (PatElemT dec -> Range) -> [PatElemT dec] -> [Range]
forall a b. (a -> b) -> [a] -> [b]
map PatElemT dec -> Range
forall a. RangeOf a => a -> Range
rangeOf ([PatElemT dec] -> [Range])
-> (PatternT dec -> [PatElemT dec]) -> PatternT dec -> [Range]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternT dec -> [PatElemT dec]
forall dec. PatternT dec -> [PatElemT dec]
patternElements

instance Ranged lore => RangesOf (Body lore) where
  rangesOf :: Body lore -> [Range]
rangesOf = BodyDec lore -> [Range]
forall a. RangesOf a => a -> [Range]
rangesOf (BodyDec lore -> [Range])
-> (Body lore -> BodyDec lore) -> Body lore -> [Range]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body lore -> BodyDec lore
forall lore. BodyT lore -> BodyDec lore
bodyDec

subExpKnownRange :: SubExp -> (KnownBound, KnownBound)
subExpKnownRange :: SubExp -> (KnownBound, KnownBound)
subExpKnownRange (Var VName
v) =
  (VName -> KnownBound
VarBound VName
v,
   VName -> KnownBound
VarBound VName
v)
subExpKnownRange (Constant PrimValue
val) =
  (ScalExp -> KnownBound
ScalarBound (ScalExp -> KnownBound) -> ScalExp -> KnownBound
forall a b. (a -> b) -> a -> b
$ PrimValue -> ScalExp
SE.Val PrimValue
val,
   ScalExp -> KnownBound
ScalarBound (ScalExp -> KnownBound) -> ScalExp -> KnownBound
forall a b. (a -> b) -> a -> b
$ PrimValue -> ScalExp
SE.Val PrimValue
val)

-- | The range of a scalar expression.
scalExpRange :: SE.ScalExp -> Range
scalExpRange :: ScalExp -> Range
scalExpRange ScalExp
se =
  (KnownBound -> Bound
forall a. a -> Maybe a
Just (KnownBound -> Bound) -> KnownBound -> Bound
forall a b. (a -> b) -> a -> b
$ ScalExp -> KnownBound
ScalarBound ScalExp
se, KnownBound -> Bound
forall a. a -> Maybe a
Just (KnownBound -> Bound) -> KnownBound -> Bound
forall a b. (a -> b) -> a -> b
$ ScalExp -> KnownBound
ScalarBound ScalExp
se)

primOpRanges :: BasicOp -> [Range]
primOpRanges :: BasicOp -> [Range]
primOpRanges (SubExp SubExp
se) =
  [SubExp -> Range
forall a. RangeOf a => a -> Range
rangeOf SubExp
se]

primOpRanges (BinOp (Add IntType
t Overflow
_) SubExp
x SubExp
y) =
  [ScalExp -> Range
scalExpRange (ScalExp -> Range) -> ScalExp -> Range
forall a b. (a -> b) -> a -> b
$ ScalExp -> ScalExp -> ScalExp
SE.SPlus (SubExp -> PrimType -> ScalExp
SE.subExpToScalExp SubExp
x (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
t) (SubExp -> PrimType -> ScalExp
SE.subExpToScalExp SubExp
y (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
t)]
primOpRanges (BinOp (Sub IntType
t Overflow
_) SubExp
x SubExp
y) =
  [ScalExp -> Range
scalExpRange (ScalExp -> Range) -> ScalExp -> Range
forall a b. (a -> b) -> a -> b
$ ScalExp -> ScalExp -> ScalExp
SE.SMinus (SubExp -> PrimType -> ScalExp
SE.subExpToScalExp SubExp
x (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
t) (SubExp -> PrimType -> ScalExp
SE.subExpToScalExp SubExp
y (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
t)]
primOpRanges (BinOp (Mul IntType
t Overflow
_) SubExp
x SubExp
y) =
  [ScalExp -> Range
scalExpRange (ScalExp -> Range) -> ScalExp -> Range
forall a b. (a -> b) -> a -> b
$ ScalExp -> ScalExp -> ScalExp
SE.STimes (SubExp -> PrimType -> ScalExp
SE.subExpToScalExp SubExp
x (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
t) (SubExp -> PrimType -> ScalExp
SE.subExpToScalExp SubExp
y (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
t)]
primOpRanges (BinOp (SDiv IntType
t) SubExp
x SubExp
y) =
  [ScalExp -> Range
scalExpRange (ScalExp -> Range) -> ScalExp -> Range
forall a b. (a -> b) -> a -> b
$ ScalExp -> ScalExp -> ScalExp
SE.SDiv (SubExp -> PrimType -> ScalExp
SE.subExpToScalExp SubExp
x (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
t) (SubExp -> PrimType -> ScalExp
SE.subExpToScalExp SubExp
y (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
t)]

primOpRanges (ConvOp (SExt IntType
from IntType
to) SubExp
x)
  | IntType
from IntType -> IntType -> Bool
forall a. Ord a => a -> a -> Bool
< IntType
to = [SubExp -> Range
forall a. RangeOf a => a -> Range
rangeOf SubExp
x]

primOpRanges (ConvOp (BToI IntType
it) SubExp
_) =
  [(KnownBound -> Bound
forall a. a -> Maybe a
Just (KnownBound -> Bound) -> KnownBound -> Bound
forall a b. (a -> b) -> a -> b
$ ScalExp -> KnownBound
ScalarBound (ScalExp -> KnownBound) -> ScalExp -> KnownBound
forall a b. (a -> b) -> a -> b
$ PrimValue -> ScalExp
SE.Val (PrimValue -> ScalExp) -> PrimValue -> ScalExp
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
it (Int
0::Int),
    KnownBound -> Bound
forall a. a -> Maybe a
Just (KnownBound -> Bound) -> KnownBound -> Bound
forall a b. (a -> b) -> a -> b
$ ScalExp -> KnownBound
ScalarBound (ScalExp -> KnownBound) -> ScalExp -> KnownBound
forall a b. (a -> b) -> a -> b
$ PrimValue -> ScalExp
SE.Val (PrimValue -> ScalExp) -> PrimValue -> ScalExp
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
it (Int
1::Int))]

primOpRanges (Iota SubExp
n SubExp
x SubExp
s IntType
Int32) =
  [(KnownBound -> Bound
forall a. a -> Maybe a
Just (KnownBound -> Bound) -> KnownBound -> Bound
forall a b. (a -> b) -> a -> b
$ ScalExp -> KnownBound
ScalarBound ScalExp
x',
    KnownBound -> Bound
forall a. a -> Maybe a
Just (KnownBound -> Bound) -> KnownBound -> Bound
forall a b. (a -> b) -> a -> b
$ ScalExp -> KnownBound
ScalarBound (ScalExp -> KnownBound) -> ScalExp -> KnownBound
forall a b. (a -> b) -> a -> b
$ ScalExp
x' ScalExp -> ScalExp -> ScalExp
forall a. Num a => a -> a -> a
+ (ScalExp
n' ScalExp -> ScalExp -> ScalExp
forall a. Num a => a -> a -> a
- ScalExp
1) ScalExp -> ScalExp -> ScalExp
forall a. Num a => a -> a -> a
* ScalExp
s')]
  where n' :: ScalExp
n' = case SubExp
n of
          Var VName
v        -> VName -> PrimType -> ScalExp
SE.Id VName
v (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
Int32
          Constant PrimValue
val -> PrimValue -> ScalExp
SE.Val PrimValue
val
        x' :: ScalExp
x' = case SubExp
x of
          Var VName
v        -> VName -> PrimType -> ScalExp
SE.Id VName
v (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
Int32
          Constant PrimValue
val -> PrimValue -> ScalExp
SE.Val PrimValue
val
        s' :: ScalExp
s' = case SubExp
s of
          Var VName
v        -> VName -> PrimType -> ScalExp
SE.Id VName
v (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
Int32
          Constant PrimValue
val -> PrimValue -> ScalExp
SE.Val PrimValue
val
primOpRanges (Replicate Shape
_ SubExp
v) =
  [SubExp -> Range
forall a. RangeOf a => a -> Range
rangeOf SubExp
v]
primOpRanges (Rearrange [Int]
_ VName
v) =
  [SubExp -> Range
forall a. RangeOf a => a -> Range
rangeOf (SubExp -> Range) -> SubExp -> Range
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v]
primOpRanges (Copy VName
se) =
  [SubExp -> Range
forall a. RangeOf a => a -> Range
rangeOf (SubExp -> Range) -> SubExp -> Range
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
se]
primOpRanges (Index VName
v Slice SubExp
_) =
  [SubExp -> Range
forall a. RangeOf a => a -> Range
rangeOf (SubExp -> Range) -> SubExp -> Range
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
v]
primOpRanges (ArrayLit (SubExp
e:[SubExp]
es) Type
_) =
  [(KnownBound -> Bound
forall a. a -> Maybe a
Just KnownBound
lower, KnownBound -> Bound
forall a. a -> Maybe a
Just KnownBound
upper)]
  where (KnownBound
e_lower, KnownBound
e_upper) = SubExp -> (KnownBound, KnownBound)
subExpKnownRange SubExp
e
        ([KnownBound]
es_lower, [KnownBound]
es_upper) = [(KnownBound, KnownBound)] -> ([KnownBound], [KnownBound])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(KnownBound, KnownBound)] -> ([KnownBound], [KnownBound]))
-> [(KnownBound, KnownBound)] -> ([KnownBound], [KnownBound])
forall a b. (a -> b) -> a -> b
$ (SubExp -> (KnownBound, KnownBound))
-> [SubExp] -> [(KnownBound, KnownBound)]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> (KnownBound, KnownBound)
subExpKnownRange [SubExp]
es
        lower :: KnownBound
lower = (KnownBound -> KnownBound -> KnownBound)
-> KnownBound -> [KnownBound] -> KnownBound
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl KnownBound -> KnownBound -> KnownBound
MinimumBound KnownBound
e_lower [KnownBound]
es_lower
        upper :: KnownBound
upper = (KnownBound -> KnownBound -> KnownBound)
-> KnownBound -> [KnownBound] -> KnownBound
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl KnownBound -> KnownBound -> KnownBound
MaximumBound KnownBound
e_upper [KnownBound]
es_upper
primOpRanges BasicOp
_ =
  [Range
unknownRange]

-- | Ranges of the value parts of the expression.
expRanges :: Ranged lore =>
             Exp lore -> [Range]
expRanges :: Exp lore -> [Range]
expRanges (BasicOp BasicOp
op) =
  BasicOp -> [Range]
primOpRanges BasicOp
op
expRanges (If SubExp
_ BodyT lore
tbranch BodyT lore
fbranch IfDec (BranchType lore)
_) =
  [Bound] -> [Bound] -> [Range]
forall a b. [a] -> [b] -> [(a, b)]
zip
  ((Bound -> Bound -> Bound) -> [Bound] -> [Bound] -> [Bound]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Bound -> Bound -> Bound
minimumBound [Bound]
t_lower [Bound]
f_lower)
  ((Bound -> Bound -> Bound) -> [Bound] -> [Bound] -> [Bound]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Bound -> Bound -> Bound
maximumBound [Bound]
t_upper [Bound]
f_upper)
  where ([Bound]
t_lower, [Bound]
t_upper) = [Range] -> ([Bound], [Bound])
forall a b. [(a, b)] -> ([a], [b])
unzip ([Range] -> ([Bound], [Bound])) -> [Range] -> ([Bound], [Bound])
forall a b. (a -> b) -> a -> b
$ BodyT lore -> [Range]
forall a. RangesOf a => a -> [Range]
rangesOf BodyT lore
tbranch
        ([Bound]
f_lower, [Bound]
f_upper) = [Range] -> ([Bound], [Bound])
forall a b. [(a, b)] -> ([a], [b])
unzip ([Range] -> ([Bound], [Bound])) -> [Range] -> ([Bound], [Bound])
forall a b. (a -> b) -> a -> b
$ BodyT lore -> [Range]
forall a. RangesOf a => a -> [Range]
rangesOf BodyT lore
fbranch
expRanges (DoLoop [(FParam lore, SubExp)]
ctxmerge [(FParam lore, SubExp)]
valmerge (ForLoop VName
i IntType
Int32 SubExp
iterations [(LParam lore, VName)]
_) BodyT lore
body) =
  ((FParam lore, SubExp) -> Range -> Range)
-> [(FParam lore, SubExp)] -> [Range] -> [Range]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (FParam lore, SubExp) -> Range -> Range
returnedRange [(FParam lore, SubExp)]
valmerge ([Range] -> [Range]) -> [Range] -> [Range]
forall a b. (a -> b) -> a -> b
$ BodyT lore -> [Range]
forall a. RangesOf a => a -> [Range]
rangesOf BodyT lore
body
  where bound_in_loop :: Names
bound_in_loop =
          [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ VName
i VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: ((FParam lore, SubExp) -> VName)
-> [(FParam lore, SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (FParam lore -> VName
forall dec. Param dec -> VName
paramName (FParam lore -> VName)
-> ((FParam lore, SubExp) -> FParam lore)
-> (FParam lore, SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FParam lore, SubExp) -> FParam lore
forall a b. (a, b) -> a
fst) ([(FParam lore, SubExp)]
ctxmerge[(FParam lore, SubExp)]
-> [(FParam lore, SubExp)] -> [(FParam lore, SubExp)]
forall a. [a] -> [a] -> [a]
++[(FParam lore, SubExp)]
valmerge) [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++
          (Stm lore -> [VName]) -> Seq (Stm lore) -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (PatternT (LetDec lore) -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (PatternT (LetDec lore) -> [VName])
-> (Stm lore -> PatternT (LetDec lore)) -> Stm lore -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm lore -> PatternT (LetDec lore)
forall lore. Stm lore -> Pattern lore
stmPattern) (BodyT lore -> Seq (Stm lore)
forall lore. BodyT lore -> Stms lore
bodyStms BodyT lore
body)

        returnedRange :: (FParam lore, SubExp) -> Range -> Range
returnedRange (FParam lore, SubExp)
mergeparam (Bound
lower, Bound
upper) =
          ((FParam lore, SubExp) -> Bound -> Bound
returnedBound (FParam lore, SubExp)
mergeparam Bound
lower,
           (FParam lore, SubExp) -> Bound -> Bound
returnedBound (FParam lore, SubExp)
mergeparam Bound
upper)

        returnedBound :: (FParam lore, SubExp) -> Bound -> Bound
returnedBound (FParam lore
param, SubExp
mergeinit) (Just KnownBound
bound)
          | FParam lore -> Type
forall dec. Typed dec => Param dec -> Type
paramType FParam lore
param Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim (IntType -> PrimType
IntType IntType
Int32),
            Just ScalExp
bound' <- KnownBound -> Maybe ScalExp
boundToScalExp KnownBound
bound,
            let se_diff :: ScalExp
se_diff =
                  ScalExp -> RangesRep -> ScalExp
AS.simplify (ScalExp -> ScalExp -> ScalExp
SE.SMinus (VName -> PrimType -> ScalExp
SE.Id (FParam lore -> VName
forall dec. Param dec -> VName
paramName FParam lore
param) (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
Int32) ScalExp
bound') RangesRep
forall k a. Map k a
M.empty,
            Names -> Names -> Bool
namesIntersect Names
bound_in_loop (Names -> Bool) -> Names -> Bool
forall a b. (a -> b) -> a -> b
$ ScalExp -> Names
forall a. FreeIn a => a -> Names
freeIn ScalExp
se_diff =
              KnownBound -> Bound
forall a. a -> Maybe a
Just (KnownBound -> Bound) -> KnownBound -> Bound
forall a b. (a -> b) -> a -> b
$ ScalExp -> KnownBound
ScalarBound (ScalExp -> KnownBound) -> ScalExp -> KnownBound
forall a b. (a -> b) -> a -> b
$ ScalExp -> ScalExp -> ScalExp
SE.SPlus (SubExp -> PrimType -> ScalExp
SE.subExpToScalExp SubExp
mergeinit (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
Int32) (ScalExp -> ScalExp) -> ScalExp -> ScalExp
forall a b. (a -> b) -> a -> b
$
              ScalExp -> ScalExp -> ScalExp
SE.STimes ScalExp
se_diff (ScalExp -> ScalExp) -> ScalExp -> ScalExp
forall a b. (a -> b) -> a -> b
$ Bool -> [ScalExp] -> ScalExp
SE.MaxMin Bool
False
              [SubExp -> PrimType -> ScalExp
SE.subExpToScalExp SubExp
iterations (PrimType -> ScalExp) -> PrimType -> ScalExp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
Int32, ScalExp
0]
        returnedBound (FParam lore, SubExp)
_ Bound
_ = Bound
forall a. Maybe a
Nothing
expRanges (Op Op lore
ranges) = Op lore -> [Range]
forall op. RangedOp op => op -> [Range]
opRanges Op lore
ranges
expRanges Exp lore
e =
  Int -> Range -> [Range]
forall a. Int -> a -> [a]
replicate (Exp lore -> Int
forall lore.
(Decorations lore, TypedOp (Op lore)) =>
Exp lore -> Int
expExtTypeSize Exp lore
e) Range
unknownRange

-- | The class of operations that can produce range information.
class IsOp op => RangedOp op where
  opRanges :: op -> [Range]

instance RangedOp () where
  opRanges :: () -> [Range]
opRanges () = []

-- | The class of operations that can be given ranging information.
-- This is a somewhat subtle concept that is only used in the
-- simplifier and when using "lore adapters".
class RangedOp (OpWithRanges op) =>
      CanBeRanged op where
  type OpWithRanges op :: Data.Kind.Type
  removeOpRanges :: OpWithRanges op -> op
  addOpRanges :: op -> OpWithRanges op

instance CanBeRanged () where
  type OpWithRanges () = ()
  removeOpRanges :: OpWithRanges () -> ()
removeOpRanges = OpWithRanges () -> ()
forall a. a -> a
id
  addOpRanges :: () -> OpWithRanges ()
addOpRanges = () -> OpWithRanges ()
forall a. a -> a
id