{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE TypeFamilies #-}
module Futhark.Analysis.SymbolTable
  ( SymbolTable (bindings, loopDepth, availableAtClosestLoop, simplifyMemory)
  , empty
  , fromScope
  , toScope
    -- * Entries
  , Entry
  , deepen
  , bindingDepth
  , valueRange
  , entryStm
  , entryLetBoundDec
  , entryType
  , asScalExp
    -- * Lookup
  , elem
  , lookup
  , lookupStm
  , lookupExp
  , lookupBasicOp
  , lookupType
  , lookupSubExp
  , lookupScalExp
  , lookupAliases
  , available
  , consume
  , index
  , index'
  , Indexed(..)
  , indexedAddCerts
  , IndexOp(..)
    -- * Insertion
  , insertStm
  , insertStms
  , insertFParams
  , insertLParam
  , insertArrayLParam
  , insertLoopVar
    -- * Bounds
  , updateBounds
  , setUpperBound
  , setLowerBound
  , isAtLeast
    -- * Misc
  , rangesRep
  , hideIf
  , hideCertified
  )
  where

import Control.Arrow ((&&&))
import Control.Monad
import Control.Monad.Reader
import Data.Ord
import Data.Maybe
import Data.List (foldl', elemIndex)
import qualified Data.List as L
import qualified Data.Set        as S
import qualified Data.Map.Strict as M

import Prelude hiding (elem, lookup)

import Futhark.Analysis.PrimExp.Convert
import Futhark.IR hiding (FParam, lookupType)
import qualified Futhark.IR as AST
import Futhark.Analysis.ScalExp

import qualified Futhark.Analysis.AlgSimplify as AS
import Futhark.IR.Prop.Ranges
  (Range, ScalExpRange, Ranged)
import qualified Futhark.IR.Prop.Ranges as Ranges
import qualified Futhark.IR.Prop.Aliases as Aliases

data SymbolTable lore = SymbolTable {
    SymbolTable lore -> Int
loopDepth :: Int
  , SymbolTable lore -> Map VName (Entry lore)
bindings :: M.Map VName (Entry lore)
  , SymbolTable lore -> Names
availableAtClosestLoop :: Names
    -- ^ Which names are available just before the most enclosing
    -- loop?
  , SymbolTable lore -> Bool
simplifyMemory :: Bool
    -- ^ We are in a situation where we should
    -- simplify/hoist/un-existentialise memory as much as possible -
    -- typically, inside a kernel.
  }

instance Semigroup (SymbolTable lore) where
  SymbolTable lore
table1 <> :: SymbolTable lore -> SymbolTable lore -> SymbolTable lore
<> SymbolTable lore
table2 =
    SymbolTable :: forall lore.
Int -> Map VName (Entry lore) -> Names -> Bool -> SymbolTable lore
SymbolTable { loopDepth :: Int
loopDepth = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (SymbolTable lore -> Int
forall lore. SymbolTable lore -> Int
loopDepth SymbolTable lore
table1) (SymbolTable lore -> Int
forall lore. SymbolTable lore -> Int
loopDepth SymbolTable lore
table2)
                , bindings :: Map VName (Entry lore)
bindings = SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings SymbolTable lore
table1 Map VName (Entry lore)
-> Map VName (Entry lore) -> Map VName (Entry lore)
forall a. Semigroup a => a -> a -> a
<> SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings SymbolTable lore
table2
                , availableAtClosestLoop :: Names
availableAtClosestLoop = SymbolTable lore -> Names
forall lore. SymbolTable lore -> Names
availableAtClosestLoop SymbolTable lore
table1 Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<>
                                           SymbolTable lore -> Names
forall lore. SymbolTable lore -> Names
availableAtClosestLoop SymbolTable lore
table2
                , simplifyMemory :: Bool
simplifyMemory = SymbolTable lore -> Bool
forall lore. SymbolTable lore -> Bool
simplifyMemory SymbolTable lore
table1 Bool -> Bool -> Bool
|| SymbolTable lore -> Bool
forall lore. SymbolTable lore -> Bool
simplifyMemory SymbolTable lore
table2
                }

instance Monoid (SymbolTable lore) where
  mempty :: SymbolTable lore
mempty = SymbolTable lore
forall lore. SymbolTable lore
empty

empty :: SymbolTable lore
empty :: SymbolTable lore
empty = Int -> Map VName (Entry lore) -> Names -> Bool -> SymbolTable lore
forall lore.
Int -> Map VName (Entry lore) -> Names -> Bool -> SymbolTable lore
SymbolTable Int
0 Map VName (Entry lore)
forall k a. Map k a
M.empty Names
forall a. Monoid a => a
mempty Bool
False

fromScope :: ASTLore lore => Scope lore -> SymbolTable lore
fromScope :: Scope lore -> SymbolTable lore
fromScope = (SymbolTable lore -> VName -> NameInfo lore -> SymbolTable lore)
-> SymbolTable lore -> Scope lore -> SymbolTable lore
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey' SymbolTable lore -> VName -> NameInfo lore -> SymbolTable lore
forall lore.
ASTLore lore =>
SymbolTable lore -> VName -> NameInfo lore -> SymbolTable lore
insertFreeVar' SymbolTable lore
forall lore. SymbolTable lore
empty
  where insertFreeVar' :: SymbolTable lore -> VName -> NameInfo lore -> SymbolTable lore
insertFreeVar' SymbolTable lore
m VName
k NameInfo lore
dec = VName -> NameInfo lore -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
VName -> NameInfo lore -> SymbolTable lore -> SymbolTable lore
insertFreeVar VName
k NameInfo lore
dec SymbolTable lore
m

toScope :: SymbolTable lore -> Scope lore
toScope :: SymbolTable lore -> Scope lore
toScope = (Entry lore -> NameInfo lore)
-> Map VName (Entry lore) -> Scope lore
forall a b k. (a -> b) -> Map k a -> Map k b
M.map Entry lore -> NameInfo lore
forall lore. Entry lore -> NameInfo lore
entryInfo (Map VName (Entry lore) -> Scope lore)
-> (SymbolTable lore -> Map VName (Entry lore))
-> SymbolTable lore
-> Scope lore
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings

deepen :: SymbolTable lore -> SymbolTable lore
deepen :: SymbolTable lore -> SymbolTable lore
deepen SymbolTable lore
vtable = SymbolTable lore
vtable { loopDepth :: Int
loopDepth = SymbolTable lore -> Int
forall lore. SymbolTable lore -> Int
loopDepth SymbolTable lore
vtable Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1,
                         availableAtClosestLoop :: Names
availableAtClosestLoop = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ Map VName (Entry lore) -> [VName]
forall k a. Map k a -> [k]
M.keys (Map VName (Entry lore) -> [VName])
-> Map VName (Entry lore) -> [VName]
forall a b. (a -> b) -> a -> b
$ SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings SymbolTable lore
vtable
                       }

-- | The result of indexing a delayed array.
data Indexed = Indexed Certificates (PrimExp VName)
               -- ^ A PrimExp based on the indexes (that is, without
               -- accessing any actual array).
             | IndexedArray Certificates VName [PrimExp VName]
               -- ^ The indexing corresponds to another (perhaps more
               -- advantageous) array.

indexedAddCerts :: Certificates -> Indexed -> Indexed
indexedAddCerts :: Certificates -> Indexed -> Indexed
indexedAddCerts Certificates
cs1 (Indexed Certificates
cs2 PrimExp VName
v) = Certificates -> PrimExp VName -> Indexed
Indexed (Certificates
cs1Certificates -> Certificates -> Certificates
forall a. Semigroup a => a -> a -> a
<>Certificates
cs2) PrimExp VName
v
indexedAddCerts Certificates
cs1 (IndexedArray Certificates
cs2 VName
arr [PrimExp VName]
v) = Certificates -> VName -> [PrimExp VName] -> Indexed
IndexedArray (Certificates
cs1Certificates -> Certificates -> Certificates
forall a. Semigroup a => a -> a -> a
<>Certificates
cs2) VName
arr [PrimExp VName]
v

instance FreeIn Indexed where
  freeIn' :: Indexed -> FV
freeIn' (Indexed Certificates
cs PrimExp VName
v) = Certificates -> FV
forall a. FreeIn a => a -> FV
freeIn' Certificates
cs FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> PrimExp VName -> FV
forall a. FreeIn a => a -> FV
freeIn' PrimExp VName
v
  freeIn' (IndexedArray Certificates
cs VName
arr [PrimExp VName]
v) = Certificates -> FV
forall a. FreeIn a => a -> FV
freeIn' Certificates
cs FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
arr FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [PrimExp VName] -> FV
forall a. FreeIn a => a -> FV
freeIn' [PrimExp VName]
v

-- | Indexing a delayed array if possible.
type IndexArray = [PrimExp VName] -> Maybe Indexed

data Entry lore = LoopVar (LoopVarEntry lore)
                | LetBound (LetBoundEntry lore)
                | FParam (FParamEntry lore)
                | LParam (LParamEntry lore)
                | FreeVar (FreeVarEntry lore)

data LoopVarEntry lore =
  LoopVarEntry { LoopVarEntry lore -> ScalExpRange
loopVarRange    :: ScalExpRange
               , LoopVarEntry lore -> Int
loopVarStmDepth :: Int
               , LoopVarEntry lore -> IntType
loopVarType     :: IntType
               }

data LetBoundEntry lore =
  LetBoundEntry { LetBoundEntry lore -> ScalExpRange
letBoundRange    :: ScalExpRange
                , LetBoundEntry lore -> LetDec lore
letBoundDec      :: LetDec lore
                , LetBoundEntry lore -> Names
letBoundAliases  :: Names
                , LetBoundEntry lore -> Stm lore
letBoundStm      :: Stm lore
                , LetBoundEntry lore -> Int
letBoundStmDepth :: Int
                , LetBoundEntry lore -> Maybe ScalExp
letBoundScalExp  :: Maybe ScalExp
                , LetBoundEntry lore -> Int -> IndexArray
letBoundIndex    :: Int -> IndexArray
                -- ^ Index a delayed array, if possible.
                , LetBoundEntry lore -> Bool
letBoundConsumed :: Bool
                  -- ^ True if consumed.
                }

data FParamEntry lore =
  FParamEntry { FParamEntry lore -> ScalExpRange
fparamRange    :: ScalExpRange
              , FParamEntry lore -> FParamInfo lore
fparamDec      :: FParamInfo lore
              , FParamEntry lore -> Names
fparamAliases  :: Names
              , FParamEntry lore -> Int
fparamStmDepth :: Int
              , FParamEntry lore -> Bool
fparamConsumed :: Bool
              }

data LParamEntry lore =
  LParamEntry { LParamEntry lore -> ScalExpRange
lparamRange    :: ScalExpRange
              , LParamEntry lore -> LParamInfo lore
lparamDec      :: LParamInfo lore
              , LParamEntry lore -> Int
lparamStmDepth :: Int
              , LParamEntry lore -> IndexArray
lparamIndex    :: IndexArray
              , LParamEntry lore -> Bool
lparamConsumed :: Bool
              }

data FreeVarEntry lore =
  FreeVarEntry { FreeVarEntry lore -> NameInfo lore
freeVarDec      :: NameInfo lore
               , FreeVarEntry lore -> Int
freeVarStmDepth :: Int
               , FreeVarEntry lore -> ScalExpRange
freeVarRange    :: ScalExpRange
               , FreeVarEntry lore -> VName -> IndexArray
freeVarIndex    :: VName -> IndexArray
                -- ^ Index a delayed array, if possible.
               , FreeVarEntry lore -> Bool
freeVarConsumed :: Bool
                -- ^ True if consumed.
               }

entryInfo :: Entry lore -> NameInfo lore
entryInfo :: Entry lore -> NameInfo lore
entryInfo (LetBound LetBoundEntry lore
entry) = LetDec lore -> NameInfo lore
forall lore. LetDec lore -> NameInfo lore
LetName (LetDec lore -> NameInfo lore) -> LetDec lore -> NameInfo lore
forall a b. (a -> b) -> a -> b
$ LetBoundEntry lore -> LetDec lore
forall lore. LetBoundEntry lore -> LetDec lore
letBoundDec LetBoundEntry lore
entry
entryInfo (LoopVar LoopVarEntry lore
entry) = IntType -> NameInfo lore
forall lore. IntType -> NameInfo lore
IndexName (IntType -> NameInfo lore) -> IntType -> NameInfo lore
forall a b. (a -> b) -> a -> b
$ LoopVarEntry lore -> IntType
forall lore. LoopVarEntry lore -> IntType
loopVarType LoopVarEntry lore
entry
entryInfo (FParam FParamEntry lore
entry) = FParamInfo lore -> NameInfo lore
forall lore. FParamInfo lore -> NameInfo lore
FParamName (FParamInfo lore -> NameInfo lore)
-> FParamInfo lore -> NameInfo lore
forall a b. (a -> b) -> a -> b
$ FParamEntry lore -> FParamInfo lore
forall lore. FParamEntry lore -> FParamInfo lore
fparamDec FParamEntry lore
entry
entryInfo (LParam LParamEntry lore
entry) = LParamInfo lore -> NameInfo lore
forall lore. LParamInfo lore -> NameInfo lore
LParamName (LParamInfo lore -> NameInfo lore)
-> LParamInfo lore -> NameInfo lore
forall a b. (a -> b) -> a -> b
$ LParamEntry lore -> LParamInfo lore
forall lore. LParamEntry lore -> LParamInfo lore
lparamDec LParamEntry lore
entry
entryInfo (FreeVar FreeVarEntry lore
entry) = FreeVarEntry lore -> NameInfo lore
forall lore. FreeVarEntry lore -> NameInfo lore
freeVarDec FreeVarEntry lore
entry

entryType :: ASTLore lore => Entry lore -> Type
entryType :: Entry lore -> Type
entryType = NameInfo lore -> Type
forall t. Typed t => t -> Type
typeOf (NameInfo lore -> Type)
-> (Entry lore -> NameInfo lore) -> Entry lore -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry lore -> NameInfo lore
forall lore. Entry lore -> NameInfo lore
entryInfo

isVarBound :: Entry lore -> Maybe (LetBoundEntry lore)
isVarBound :: Entry lore -> Maybe (LetBoundEntry lore)
isVarBound (LetBound LetBoundEntry lore
entry) = LetBoundEntry lore -> Maybe (LetBoundEntry lore)
forall a. a -> Maybe a
Just LetBoundEntry lore
entry
isVarBound Entry lore
_ = Maybe (LetBoundEntry lore)
forall a. Maybe a
Nothing

asScalExp :: Entry lore -> Maybe ScalExp
asScalExp :: Entry lore -> Maybe ScalExp
asScalExp = LetBoundEntry lore -> Maybe ScalExp
forall lore. LetBoundEntry lore -> Maybe ScalExp
letBoundScalExp (LetBoundEntry lore -> Maybe ScalExp)
-> (Entry lore -> Maybe (LetBoundEntry lore))
-> Entry lore
-> Maybe ScalExp
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Entry lore -> Maybe (LetBoundEntry lore)
forall lore. Entry lore -> Maybe (LetBoundEntry lore)
isVarBound

bindingDepth :: Entry lore -> Int
bindingDepth :: Entry lore -> Int
bindingDepth (LetBound LetBoundEntry lore
entry) = LetBoundEntry lore -> Int
forall lore. LetBoundEntry lore -> Int
letBoundStmDepth LetBoundEntry lore
entry
bindingDepth (FParam FParamEntry lore
entry) = FParamEntry lore -> Int
forall lore. FParamEntry lore -> Int
fparamStmDepth FParamEntry lore
entry
bindingDepth (LParam LParamEntry lore
entry) = LParamEntry lore -> Int
forall lore. LParamEntry lore -> Int
lparamStmDepth LParamEntry lore
entry
bindingDepth (LoopVar LoopVarEntry lore
entry) = LoopVarEntry lore -> Int
forall lore. LoopVarEntry lore -> Int
loopVarStmDepth LoopVarEntry lore
entry
bindingDepth (FreeVar FreeVarEntry lore
_) = Int
0

setStmDepth :: Int -> Entry lore -> Entry lore
setStmDepth :: Int -> Entry lore -> Entry lore
setStmDepth Int
d (LetBound LetBoundEntry lore
entry) =
  LetBoundEntry lore -> Entry lore
forall lore. LetBoundEntry lore -> Entry lore
LetBound (LetBoundEntry lore -> Entry lore)
-> LetBoundEntry lore -> Entry lore
forall a b. (a -> b) -> a -> b
$ LetBoundEntry lore
entry { letBoundStmDepth :: Int
letBoundStmDepth = Int
d }
setStmDepth Int
d (FParam FParamEntry lore
entry) =
  FParamEntry lore -> Entry lore
forall lore. FParamEntry lore -> Entry lore
FParam (FParamEntry lore -> Entry lore) -> FParamEntry lore -> Entry lore
forall a b. (a -> b) -> a -> b
$ FParamEntry lore
entry { fparamStmDepth :: Int
fparamStmDepth = Int
d }
setStmDepth Int
d (LParam LParamEntry lore
entry) =
  LParamEntry lore -> Entry lore
forall lore. LParamEntry lore -> Entry lore
LParam (LParamEntry lore -> Entry lore) -> LParamEntry lore -> Entry lore
forall a b. (a -> b) -> a -> b
$ LParamEntry lore
entry { lparamStmDepth :: Int
lparamStmDepth = Int
d }
setStmDepth Int
d (LoopVar LoopVarEntry lore
entry) =
  LoopVarEntry lore -> Entry lore
forall lore. LoopVarEntry lore -> Entry lore
LoopVar (LoopVarEntry lore -> Entry lore)
-> LoopVarEntry lore -> Entry lore
forall a b. (a -> b) -> a -> b
$ LoopVarEntry lore
entry { loopVarStmDepth :: Int
loopVarStmDepth = Int
d }
setStmDepth Int
d (FreeVar FreeVarEntry lore
entry) =
  FreeVarEntry lore -> Entry lore
forall lore. FreeVarEntry lore -> Entry lore
FreeVar (FreeVarEntry lore -> Entry lore)
-> FreeVarEntry lore -> Entry lore
forall a b. (a -> b) -> a -> b
$ FreeVarEntry lore
entry { freeVarStmDepth :: Int
freeVarStmDepth = Int
d }

valueRange :: Entry lore -> ScalExpRange
valueRange :: Entry lore -> ScalExpRange
valueRange (LetBound LetBoundEntry lore
entry) = LetBoundEntry lore -> ScalExpRange
forall lore. LetBoundEntry lore -> ScalExpRange
letBoundRange LetBoundEntry lore
entry
valueRange (FParam FParamEntry lore
entry)   = FParamEntry lore -> ScalExpRange
forall lore. FParamEntry lore -> ScalExpRange
fparamRange FParamEntry lore
entry
valueRange (LParam LParamEntry lore
entry)   = LParamEntry lore -> ScalExpRange
forall lore. LParamEntry lore -> ScalExpRange
lparamRange LParamEntry lore
entry
valueRange (LoopVar LoopVarEntry lore
entry)  = LoopVarEntry lore -> ScalExpRange
forall lore. LoopVarEntry lore -> ScalExpRange
loopVarRange LoopVarEntry lore
entry
valueRange (FreeVar FreeVarEntry lore
entry)  = FreeVarEntry lore -> ScalExpRange
forall lore. FreeVarEntry lore -> ScalExpRange
freeVarRange FreeVarEntry lore
entry

setValueRange :: ScalExpRange -> Entry lore -> Entry lore
setValueRange :: ScalExpRange -> Entry lore -> Entry lore
setValueRange ScalExpRange
range (LetBound LetBoundEntry lore
entry) =
  LetBoundEntry lore -> Entry lore
forall lore. LetBoundEntry lore -> Entry lore
LetBound (LetBoundEntry lore -> Entry lore)
-> LetBoundEntry lore -> Entry lore
forall a b. (a -> b) -> a -> b
$ LetBoundEntry lore
entry { letBoundRange :: ScalExpRange
letBoundRange = ScalExpRange
range }
setValueRange ScalExpRange
range (FParam FParamEntry lore
entry) =
  FParamEntry lore -> Entry lore
forall lore. FParamEntry lore -> Entry lore
FParam (FParamEntry lore -> Entry lore) -> FParamEntry lore -> Entry lore
forall a b. (a -> b) -> a -> b
$ FParamEntry lore
entry { fparamRange :: ScalExpRange
fparamRange = ScalExpRange
range }
setValueRange ScalExpRange
range (LParam LParamEntry lore
entry) =
  LParamEntry lore -> Entry lore
forall lore. LParamEntry lore -> Entry lore
LParam (LParamEntry lore -> Entry lore) -> LParamEntry lore -> Entry lore
forall a b. (a -> b) -> a -> b
$ LParamEntry lore
entry { lparamRange :: ScalExpRange
lparamRange = ScalExpRange
range }
setValueRange ScalExpRange
range (LoopVar LoopVarEntry lore
entry) =
  LoopVarEntry lore -> Entry lore
forall lore. LoopVarEntry lore -> Entry lore
LoopVar (LoopVarEntry lore -> Entry lore)
-> LoopVarEntry lore -> Entry lore
forall a b. (a -> b) -> a -> b
$ LoopVarEntry lore
entry { loopVarRange :: ScalExpRange
loopVarRange = ScalExpRange
range }
setValueRange ScalExpRange
range (FreeVar FreeVarEntry lore
entry) =
  FreeVarEntry lore -> Entry lore
forall lore. FreeVarEntry lore -> Entry lore
FreeVar (FreeVarEntry lore -> Entry lore)
-> FreeVarEntry lore -> Entry lore
forall a b. (a -> b) -> a -> b
$ FreeVarEntry lore
entry { freeVarRange :: ScalExpRange
freeVarRange = ScalExpRange
range }

consumed :: Entry lore -> Bool
consumed :: Entry lore -> Bool
consumed (LetBound LetBoundEntry lore
entry) = LetBoundEntry lore -> Bool
forall lore. LetBoundEntry lore -> Bool
letBoundConsumed LetBoundEntry lore
entry
consumed (FParam FParamEntry lore
entry)   = FParamEntry lore -> Bool
forall lore. FParamEntry lore -> Bool
fparamConsumed FParamEntry lore
entry
consumed (LParam LParamEntry lore
entry)   = LParamEntry lore -> Bool
forall lore. LParamEntry lore -> Bool
lparamConsumed LParamEntry lore
entry
consumed LoopVar{}        = Bool
False
consumed (FreeVar FreeVarEntry lore
entry)  = FreeVarEntry lore -> Bool
forall lore. FreeVarEntry lore -> Bool
freeVarConsumed FreeVarEntry lore
entry

entryStm :: Entry lore -> Maybe (Stm lore)
entryStm :: Entry lore -> Maybe (Stm lore)
entryStm (LetBound LetBoundEntry lore
entry) = Stm lore -> Maybe (Stm lore)
forall a. a -> Maybe a
Just (Stm lore -> Maybe (Stm lore)) -> Stm lore -> Maybe (Stm lore)
forall a b. (a -> b) -> a -> b
$ LetBoundEntry lore -> Stm lore
forall lore. LetBoundEntry lore -> Stm lore
letBoundStm LetBoundEntry lore
entry
entryStm Entry lore
_                = Maybe (Stm lore)
forall a. Maybe a
Nothing

entryLetBoundDec :: Entry lore -> Maybe (LetDec lore)
entryLetBoundDec :: Entry lore -> Maybe (LetDec lore)
entryLetBoundDec (LetBound LetBoundEntry lore
entry) = LetDec lore -> Maybe (LetDec lore)
forall a. a -> Maybe a
Just (LetDec lore -> Maybe (LetDec lore))
-> LetDec lore -> Maybe (LetDec lore)
forall a b. (a -> b) -> a -> b
$ LetBoundEntry lore -> LetDec lore
forall lore. LetBoundEntry lore -> LetDec lore
letBoundDec LetBoundEntry lore
entry
entryLetBoundDec Entry lore
_                = Maybe (LetDec lore)
forall a. Maybe a
Nothing

asStm :: Entry lore -> Maybe (Stm lore)
asStm :: Entry lore -> Maybe (Stm lore)
asStm = (LetBoundEntry lore -> Stm lore)
-> Maybe (LetBoundEntry lore) -> Maybe (Stm lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap LetBoundEntry lore -> Stm lore
forall lore. LetBoundEntry lore -> Stm lore
letBoundStm (Maybe (LetBoundEntry lore) -> Maybe (Stm lore))
-> (Entry lore -> Maybe (LetBoundEntry lore))
-> Entry lore
-> Maybe (Stm lore)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry lore -> Maybe (LetBoundEntry lore)
forall lore. Entry lore -> Maybe (LetBoundEntry lore)
isVarBound

elem :: VName -> SymbolTable lore -> Bool
elem :: VName -> SymbolTable lore -> Bool
elem VName
name = Maybe (Entry lore) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Entry lore) -> Bool)
-> (SymbolTable lore -> Maybe (Entry lore))
-> SymbolTable lore
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SymbolTable lore -> Maybe (Entry lore)
forall lore. VName -> SymbolTable lore -> Maybe (Entry lore)
lookup VName
name

lookup :: VName -> SymbolTable lore -> Maybe (Entry lore)
lookup :: VName -> SymbolTable lore -> Maybe (Entry lore)
lookup VName
name = VName -> Map VName (Entry lore) -> Maybe (Entry lore)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (Map VName (Entry lore) -> Maybe (Entry lore))
-> (SymbolTable lore -> Map VName (Entry lore))
-> SymbolTable lore
-> Maybe (Entry lore)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings

lookupStm :: VName -> SymbolTable lore -> Maybe (Stm lore)
lookupStm :: VName -> SymbolTable lore -> Maybe (Stm lore)
lookupStm VName
name SymbolTable lore
vtable = Entry lore -> Maybe (Stm lore)
forall lore. Entry lore -> Maybe (Stm lore)
asStm (Entry lore -> Maybe (Stm lore))
-> Maybe (Entry lore) -> Maybe (Stm lore)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> SymbolTable lore -> Maybe (Entry lore)
forall lore. VName -> SymbolTable lore -> Maybe (Entry lore)
lookup VName
name SymbolTable lore
vtable

lookupExp :: VName -> SymbolTable lore -> Maybe (Exp lore, Certificates)
lookupExp :: VName -> SymbolTable lore -> Maybe (Exp lore, Certificates)
lookupExp VName
name SymbolTable lore
vtable = (Stm lore -> Exp lore
forall lore. Stm lore -> Exp lore
stmExp (Stm lore -> Exp lore)
-> (Stm lore -> Certificates)
-> Stm lore
-> (Exp lore, Certificates)
forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& Stm lore -> Certificates
forall lore. Stm lore -> Certificates
stmCerts) (Stm lore -> (Exp lore, Certificates))
-> Maybe (Stm lore) -> Maybe (Exp lore, Certificates)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> SymbolTable lore -> Maybe (Stm lore)
forall lore. VName -> SymbolTable lore -> Maybe (Stm lore)
lookupStm VName
name SymbolTable lore
vtable

lookupBasicOp :: VName -> SymbolTable lore -> Maybe (BasicOp, Certificates)
lookupBasicOp :: VName -> SymbolTable lore -> Maybe (BasicOp, Certificates)
lookupBasicOp VName
name SymbolTable lore
vtable = case VName -> SymbolTable lore -> Maybe (Exp lore, Certificates)
forall lore.
VName -> SymbolTable lore -> Maybe (Exp lore, Certificates)
lookupExp VName
name SymbolTable lore
vtable of
  Just (BasicOp BasicOp
e, Certificates
cs) -> (BasicOp, Certificates) -> Maybe (BasicOp, Certificates)
forall a. a -> Maybe a
Just (BasicOp
e, Certificates
cs)
  Maybe (Exp lore, Certificates)
_                    -> Maybe (BasicOp, Certificates)
forall a. Maybe a
Nothing

lookupType :: ASTLore lore => VName -> SymbolTable lore -> Maybe Type
lookupType :: VName -> SymbolTable lore -> Maybe Type
lookupType VName
name SymbolTable lore
vtable = Entry lore -> Type
forall lore. ASTLore lore => Entry lore -> Type
entryType (Entry lore -> Type) -> Maybe (Entry lore) -> Maybe Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> SymbolTable lore -> Maybe (Entry lore)
forall lore. VName -> SymbolTable lore -> Maybe (Entry lore)
lookup VName
name SymbolTable lore
vtable

lookupSubExpType :: ASTLore lore => SubExp -> SymbolTable lore -> Maybe Type
lookupSubExpType :: SubExp -> SymbolTable lore -> Maybe Type
lookupSubExpType (Var VName
v) = VName -> SymbolTable lore -> Maybe Type
forall lore.
ASTLore lore =>
VName -> SymbolTable lore -> Maybe Type
lookupType VName
v
lookupSubExpType (Constant PrimValue
v) = Maybe Type -> SymbolTable lore -> Maybe Type
forall a b. a -> b -> a
const (Maybe Type -> SymbolTable lore -> Maybe Type)
-> Maybe Type -> SymbolTable lore -> Maybe Type
forall a b. (a -> b) -> a -> b
$ Type -> Maybe Type
forall a. a -> Maybe a
Just (Type -> Maybe Type) -> Type -> Maybe Type
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimType
primValueType PrimValue
v

lookupSubExp :: VName -> SymbolTable lore -> Maybe (SubExp, Certificates)
lookupSubExp :: VName -> SymbolTable lore -> Maybe (SubExp, Certificates)
lookupSubExp VName
name SymbolTable lore
vtable = do
  (Exp lore
e,Certificates
cs) <- VName -> SymbolTable lore -> Maybe (Exp lore, Certificates)
forall lore.
VName -> SymbolTable lore -> Maybe (Exp lore, Certificates)
lookupExp VName
name SymbolTable lore
vtable
  case Exp lore
e of
    BasicOp (SubExp SubExp
se) -> (SubExp, Certificates) -> Maybe (SubExp, Certificates)
forall a. a -> Maybe a
Just (SubExp
se,Certificates
cs)
    Exp lore
_                   -> Maybe (SubExp, Certificates)
forall a. Maybe a
Nothing

lookupScalExp :: ASTLore lore => VName -> SymbolTable lore -> Maybe ScalExp
lookupScalExp :: VName -> SymbolTable lore -> Maybe ScalExp
lookupScalExp VName
name SymbolTable lore
vtable =
  case (VName -> SymbolTable lore -> Maybe (Entry lore)
forall lore. VName -> SymbolTable lore -> Maybe (Entry lore)
lookup VName
name SymbolTable lore
vtable, VName -> SymbolTable lore -> ScalExpRange
forall lore. VName -> SymbolTable lore -> ScalExpRange
lookupRange VName
name SymbolTable lore
vtable) of
    -- If we know the lower and upper bound, and these are the same,
    -- then we morally know the ScalExp, but only if the variable has
    -- the right type.
    (Just Entry lore
entry, (Just ScalExp
lower, Just ScalExp
upper))
      | Entry lore -> Type
forall lore. ASTLore lore => Entry lore -> Type
entryType Entry lore
entry Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int32,
        ScalExp
lower ScalExp -> ScalExp -> Bool
forall a. Eq a => a -> a -> Bool
== ScalExp
upper, ScalExp -> PrimType
scalExpType ScalExp
lower PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
int32 ->
          ScalExp -> Maybe ScalExp
forall a. a -> Maybe a
Just (ScalExp -> Maybe ScalExp) -> ScalExp -> Maybe ScalExp
forall a b. (a -> b) -> a -> b
$ LookupVar -> ScalExp -> ScalExp
expandScalExp (VName -> SymbolTable lore -> Maybe ScalExp
forall lore.
ASTLore lore =>
VName -> SymbolTable lore -> Maybe ScalExp
`lookupScalExp` SymbolTable lore
vtable) ScalExp
lower
    (Just Entry lore
entry, ScalExpRange
_) -> Entry lore -> Maybe ScalExp
forall lore. Entry lore -> Maybe ScalExp
asScalExp Entry lore
entry
    (Maybe (Entry lore), ScalExpRange)
_ -> Maybe ScalExp
forall a. Maybe a
Nothing

lookupAliases :: VName -> SymbolTable lore -> Names
lookupAliases :: VName -> SymbolTable lore -> Names
lookupAliases VName
name SymbolTable lore
vtable = case VName -> Map VName (Entry lore) -> Maybe (Entry lore)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (Map VName (Entry lore) -> Maybe (Entry lore))
-> Map VName (Entry lore) -> Maybe (Entry lore)
forall a b. (a -> b) -> a -> b
$ SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings SymbolTable lore
vtable of
                              Just (LetBound LetBoundEntry lore
e) -> LetBoundEntry lore -> Names
forall lore. LetBoundEntry lore -> Names
letBoundAliases LetBoundEntry lore
e
                              Just (FParam FParamEntry lore
e)   -> FParamEntry lore -> Names
forall lore. FParamEntry lore -> Names
fparamAliases FParamEntry lore
e
                              Maybe (Entry lore)
_                 -> Names
forall a. Monoid a => a
mempty

-- | In symbol table and not consumed.
available :: VName -> SymbolTable lore -> Bool
available :: VName -> SymbolTable lore -> Bool
available VName
name = Bool -> (Entry lore -> Bool) -> Maybe (Entry lore) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (Bool -> Bool
not (Bool -> Bool) -> (Entry lore -> Bool) -> Entry lore -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry lore -> Bool
forall lore. Entry lore -> Bool
consumed) (Maybe (Entry lore) -> Bool)
-> (SymbolTable lore -> Maybe (Entry lore))
-> SymbolTable lore
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Map VName (Entry lore) -> Maybe (Entry lore)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (Map VName (Entry lore) -> Maybe (Entry lore))
-> (SymbolTable lore -> Map VName (Entry lore))
-> SymbolTable lore
-> Maybe (Entry lore)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings

index :: ASTLore lore => VName -> [SubExp] -> SymbolTable lore
      -> Maybe Indexed
index :: VName -> [SubExp] -> SymbolTable lore -> Maybe Indexed
index VName
name [SubExp]
is SymbolTable lore
table = do
  [PrimExp VName]
is' <- (SubExp -> Maybe (PrimExp VName))
-> [SubExp] -> Maybe [PrimExp VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> Maybe (PrimExp VName)
asPrimExp [SubExp]
is
  VName -> [PrimExp VName] -> SymbolTable lore -> Maybe Indexed
forall lore.
VName -> [PrimExp VName] -> SymbolTable lore -> Maybe Indexed
index' VName
name [PrimExp VName]
is' SymbolTable lore
table
  where asPrimExp :: SubExp -> Maybe (PrimExp VName)
asPrimExp SubExp
i = do
          Prim PrimType
t <- SubExp -> SymbolTable lore -> Maybe Type
forall lore.
ASTLore lore =>
SubExp -> SymbolTable lore -> Maybe Type
lookupSubExpType SubExp
i SymbolTable lore
table
          PrimExp VName -> Maybe (PrimExp VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (PrimExp VName -> Maybe (PrimExp VName))
-> PrimExp VName -> Maybe (PrimExp VName)
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
t SubExp
i

index' :: VName -> [PrimExp VName] -> SymbolTable lore
       -> Maybe Indexed
index' :: VName -> [PrimExp VName] -> SymbolTable lore -> Maybe Indexed
index' VName
name [PrimExp VName]
is SymbolTable lore
vtable = do
  Entry lore
entry <- VName -> SymbolTable lore -> Maybe (Entry lore)
forall lore. VName -> SymbolTable lore -> Maybe (Entry lore)
lookup VName
name SymbolTable lore
vtable
  case Entry lore
entry of
    LetBound LetBoundEntry lore
entry' |
      Just Int
k <- VName -> [VName] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex VName
name ([VName] -> Maybe Int) -> [VName] -> Maybe Int
forall a b. (a -> b) -> a -> b
$ PatternT (LetDec lore) -> [VName]
forall dec. PatternT dec -> [VName]
patternValueNames (PatternT (LetDec lore) -> [VName])
-> PatternT (LetDec lore) -> [VName]
forall a b. (a -> b) -> a -> b
$
                Stm lore -> PatternT (LetDec lore)
forall lore. Stm lore -> Pattern lore
stmPattern (Stm lore -> PatternT (LetDec lore))
-> Stm lore -> PatternT (LetDec lore)
forall a b. (a -> b) -> a -> b
$ LetBoundEntry lore -> Stm lore
forall lore. LetBoundEntry lore -> Stm lore
letBoundStm LetBoundEntry lore
entry' ->
        LetBoundEntry lore -> Int -> IndexArray
forall lore. LetBoundEntry lore -> Int -> IndexArray
letBoundIndex LetBoundEntry lore
entry' Int
k [PrimExp VName]
is
    FreeVar FreeVarEntry lore
entry' ->
      FreeVarEntry lore -> VName -> IndexArray
forall lore. FreeVarEntry lore -> VName -> IndexArray
freeVarIndex FreeVarEntry lore
entry' VName
name [PrimExp VName]
is
    LParam LParamEntry lore
entry' -> LParamEntry lore -> IndexArray
forall lore. LParamEntry lore -> IndexArray
lparamIndex LParamEntry lore
entry' [PrimExp VName]
is
    Entry lore
_ -> Maybe Indexed
forall a. Maybe a
Nothing

lookupRange :: VName -> SymbolTable lore -> ScalExpRange
lookupRange :: VName -> SymbolTable lore -> ScalExpRange
lookupRange VName
name SymbolTable lore
vtable =
  ScalExpRange
-> (Entry lore -> ScalExpRange)
-> Maybe (Entry lore)
-> ScalExpRange
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Maybe ScalExp
forall a. Maybe a
Nothing, Maybe ScalExp
forall a. Maybe a
Nothing) Entry lore -> ScalExpRange
forall lore. Entry lore -> ScalExpRange
valueRange (Maybe (Entry lore) -> ScalExpRange)
-> Maybe (Entry lore) -> ScalExpRange
forall a b. (a -> b) -> a -> b
$ VName -> SymbolTable lore -> Maybe (Entry lore)
forall lore. VName -> SymbolTable lore -> Maybe (Entry lore)
lookup VName
name SymbolTable lore
vtable

rangesRep :: SymbolTable lore -> AS.RangesRep
rangesRep :: SymbolTable lore -> RangesRep
rangesRep = ((Int, Maybe ScalExp, Maybe ScalExp) -> Bool)
-> RangesRep -> RangesRep
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (Int, Maybe ScalExp, Maybe ScalExp) -> Bool
forall a a a. (a, Maybe a, Maybe a) -> Bool
knownRange (RangesRep -> RangesRep)
-> (SymbolTable lore -> RangesRep) -> SymbolTable lore -> RangesRep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Entry lore -> (Int, Maybe ScalExp, Maybe ScalExp))
-> Map VName (Entry lore) -> RangesRep
forall a b k. (a -> b) -> Map k a -> Map k b
M.map Entry lore -> (Int, Maybe ScalExp, Maybe ScalExp)
forall lore. Entry lore -> (Int, Maybe ScalExp, Maybe ScalExp)
toRep (Map VName (Entry lore) -> RangesRep)
-> (SymbolTable lore -> Map VName (Entry lore))
-> SymbolTable lore
-> RangesRep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings
  where toRep :: Entry lore -> (Int, Maybe ScalExp, Maybe ScalExp)
toRep Entry lore
entry = (Entry lore -> Int
forall lore. Entry lore -> Int
bindingDepth Entry lore
entry, Maybe ScalExp
lower, Maybe ScalExp
upper)
          where (Maybe ScalExp
lower, Maybe ScalExp
upper) = Entry lore -> ScalExpRange
forall lore. Entry lore -> ScalExpRange
valueRange Entry lore
entry
        knownRange :: (a, Maybe a, Maybe a) -> Bool
knownRange (a
_, Maybe a
lower, Maybe a
upper) = Maybe a -> Bool
forall a. Maybe a -> Bool
isJust Maybe a
lower Bool -> Bool -> Bool
|| Maybe a -> Bool
forall a. Maybe a -> Bool
isJust Maybe a
upper

class IndexOp op where
  indexOp :: (ASTLore lore, IndexOp (Op lore)) =>
             SymbolTable lore -> Int -> op
          -> [PrimExp VName] -> Maybe Indexed
  indexOp SymbolTable lore
_ Int
_ op
_ [PrimExp VName]
_ = Maybe Indexed
forall a. Maybe a
Nothing

instance IndexOp () where

indexExp :: (IndexOp (Op lore), ASTLore lore) =>
            SymbolTable lore -> Exp lore -> Int -> IndexArray

indexExp :: SymbolTable lore -> Exp lore -> Int -> IndexArray
indexExp SymbolTable lore
vtable (Op Op lore
op) Int
k [PrimExp VName]
is =
  SymbolTable lore -> Int -> Op lore -> IndexArray
forall op lore.
(IndexOp op, ASTLore lore, IndexOp (Op lore)) =>
SymbolTable lore -> Int -> op -> IndexArray
indexOp SymbolTable lore
vtable Int
k Op lore
op [PrimExp VName]
is

indexExp SymbolTable lore
_ (BasicOp (Iota SubExp
_ SubExp
x SubExp
s IntType
to_it)) Int
_ [PrimExp VName
i]
  | IntType IntType
from_it <- PrimExp VName -> PrimType
forall v. PrimExp v -> PrimType
primExpType PrimExp VName
i =
      Indexed -> Maybe Indexed
forall a. a -> Maybe a
Just (Indexed -> Maybe Indexed) -> Indexed -> Maybe Indexed
forall a b. (a -> b) -> a -> b
$ Certificates -> PrimExp VName -> Indexed
Indexed Certificates
forall a. Monoid a => a
mempty (PrimExp VName -> Indexed) -> PrimExp VName -> Indexed
forall a b. (a -> b) -> a -> b
$
       ConvOp -> PrimExp VName -> PrimExp VName
forall v. ConvOp -> PrimExp v -> PrimExp v
ConvOpExp (IntType -> IntType -> ConvOp
SExt IntType
from_it IntType
to_it) PrimExp VName
i
       PrimExp VName -> PrimExp VName -> PrimExp VName
forall a. Num a => a -> a -> a
* PrimType -> SubExp -> PrimExp VName
primExpFromSubExp (IntType -> PrimType
IntType IntType
to_it) SubExp
s
       PrimExp VName -> PrimExp VName -> PrimExp VName
forall a. Num a => a -> a -> a
+ PrimType -> SubExp -> PrimExp VName
primExpFromSubExp (IntType -> PrimType
IntType IntType
to_it) SubExp
x

indexExp SymbolTable lore
table (BasicOp (Replicate (Shape [SubExp]
ds) SubExp
v)) Int
_ [PrimExp VName]
is
  | [SubExp] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
ds Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [PrimExp VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PrimExp VName]
is,
    Just (Prim PrimType
t) <- SubExp -> SymbolTable lore -> Maybe Type
forall lore.
ASTLore lore =>
SubExp -> SymbolTable lore -> Maybe Type
lookupSubExpType SubExp
v SymbolTable lore
table =
      Indexed -> Maybe Indexed
forall a. a -> Maybe a
Just (Indexed -> Maybe Indexed) -> Indexed -> Maybe Indexed
forall a b. (a -> b) -> a -> b
$ Certificates -> PrimExp VName -> Indexed
Indexed Certificates
forall a. Monoid a => a
mempty (PrimExp VName -> Indexed) -> PrimExp VName -> Indexed
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
t SubExp
v

indexExp SymbolTable lore
table (BasicOp (Replicate (Shape [SubExp
_]) (Var VName
v))) Int
_ (PrimExp VName
_:[PrimExp VName]
is) =
  VName -> [PrimExp VName] -> SymbolTable lore -> Maybe Indexed
forall lore.
VName -> [PrimExp VName] -> SymbolTable lore -> Maybe Indexed
index' VName
v [PrimExp VName]
is SymbolTable lore
table

indexExp SymbolTable lore
table (BasicOp (Reshape ShapeChange SubExp
newshape VName
v)) Int
_ [PrimExp VName]
is
  | Just [SubExp]
oldshape <- 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
<$> VName -> SymbolTable lore -> Maybe Type
forall lore.
ASTLore lore =>
VName -> SymbolTable lore -> Maybe Type
lookupType VName
v SymbolTable lore
table =
      let is' :: [PrimExp VName]
is' =
            [PrimExp VName]
-> [PrimExp VName] -> [PrimExp VName] -> [PrimExp VName]
forall num. IntegralExp num => [num] -> [num] -> [num] -> [num]
reshapeIndex ((SubExp -> PrimExp VName) -> [SubExp] -> [PrimExp VName]
forall a b. (a -> b) -> [a] -> [b]
map (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int32) [SubExp]
oldshape)
                         ((SubExp -> PrimExp VName) -> [SubExp] -> [PrimExp VName]
forall a b. (a -> b) -> [a] -> [b]
map (PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
int32) ([SubExp] -> [PrimExp VName]) -> [SubExp] -> [PrimExp VName]
forall a b. (a -> b) -> a -> b
$ ShapeChange SubExp -> [SubExp]
forall d. ShapeChange d -> [d]
newDims ShapeChange SubExp
newshape)
                         [PrimExp VName]
is
      in VName -> [PrimExp VName] -> SymbolTable lore -> Maybe Indexed
forall lore.
VName -> [PrimExp VName] -> SymbolTable lore -> Maybe Indexed
index' VName
v [PrimExp VName]
is' SymbolTable lore
table

indexExp SymbolTable lore
table (BasicOp (Index VName
v Slice SubExp
slice)) Int
_ [PrimExp VName]
is =
  VName -> [PrimExp VName] -> SymbolTable lore -> Maybe Indexed
forall lore.
VName -> [PrimExp VName] -> SymbolTable lore -> Maybe Indexed
index' VName
v (Slice SubExp -> [PrimExp VName] -> [PrimExp VName]
adjust Slice SubExp
slice [PrimExp VName]
is) SymbolTable lore
table
  where adjust :: Slice SubExp -> [PrimExp VName] -> [PrimExp VName]
adjust (DimFix SubExp
j:Slice SubExp
js') [PrimExp VName]
is' =
          SubExp -> PrimExp VName
pe SubExp
j PrimExp VName -> [PrimExp VName] -> [PrimExp VName]
forall a. a -> [a] -> [a]
: Slice SubExp -> [PrimExp VName] -> [PrimExp VName]
adjust Slice SubExp
js' [PrimExp VName]
is'
        adjust (DimSlice SubExp
j SubExp
_ SubExp
s:Slice SubExp
js') (PrimExp VName
i:[PrimExp VName]
is') =
          let i_t_s :: PrimExp VName
i_t_s = PrimExp VName
i PrimExp VName -> PrimExp VName -> PrimExp VName
forall a. Num a => a -> a -> a
* SubExp -> PrimExp VName
pe SubExp
s
              j_p_i_t_s :: PrimExp VName
j_p_i_t_s = SubExp -> PrimExp VName
pe SubExp
j PrimExp VName -> PrimExp VName -> PrimExp VName
forall a. Num a => a -> a -> a
+ PrimExp VName
i_t_s
          in PrimExp VName
j_p_i_t_s PrimExp VName -> [PrimExp VName] -> [PrimExp VName]
forall a. a -> [a] -> [a]
: Slice SubExp -> [PrimExp VName] -> [PrimExp VName]
adjust Slice SubExp
js' [PrimExp VName]
is'
        adjust Slice SubExp
_ [PrimExp VName]
_ = []

        pe :: SubExp -> PrimExp VName
pe = PrimType -> SubExp -> PrimExp VName
primExpFromSubExp (IntType -> PrimType
IntType IntType
Int32)

indexExp SymbolTable lore
_ Exp lore
_ Int
_ [PrimExp VName]
_ = Maybe Indexed
forall a. Maybe a
Nothing

defBndEntry :: (ASTLore lore, IndexOp (Op lore)) =>
               SymbolTable lore
            -> PatElem lore
            -> Range
            -> Names
            -> Stm lore
            -> LetBoundEntry lore
defBndEntry :: SymbolTable lore
-> PatElem lore -> Range -> Names -> Stm lore -> LetBoundEntry lore
defBndEntry SymbolTable lore
vtable PatElem lore
patElem Range
range Names
als Stm lore
bnd =
  LetBoundEntry :: forall lore.
ScalExpRange
-> LetDec lore
-> Names
-> Stm lore
-> Int
-> Maybe ScalExp
-> (Int -> IndexArray)
-> Bool
-> LetBoundEntry lore
LetBoundEntry {
      letBoundRange :: ScalExpRange
letBoundRange = ScalExpRange -> ScalExpRange
simplifyRange (ScalExpRange -> ScalExpRange) -> ScalExpRange -> ScalExpRange
forall a b. (a -> b) -> a -> b
$ Range -> ScalExpRange
scalExpRange Range
range
    , letBoundDec :: LetDec lore
letBoundDec = PatElem lore -> LetDec lore
forall dec. PatElemT dec -> dec
patElemDec PatElem lore
patElem
    , letBoundAliases :: Names
letBoundAliases = Names
als
    , letBoundStm :: Stm lore
letBoundStm = Stm lore
bnd
    , letBoundScalExp :: Maybe ScalExp
letBoundScalExp =
      Reader (Scope lore) (Maybe ScalExp) -> Scope lore -> Maybe ScalExp
forall r a. Reader r a -> r -> a
runReader (LookupVar -> Exp lore -> Reader (Scope lore) (Maybe ScalExp)
forall t (f :: * -> *) lore.
(HasScope t f, Monad f) =>
LookupVar -> Exp lore -> f (Maybe ScalExp)
toScalExp (VName -> SymbolTable lore -> Maybe ScalExp
forall lore.
ASTLore lore =>
VName -> SymbolTable lore -> Maybe ScalExp
`lookupScalExp` SymbolTable lore
vtable) (Stm lore -> Exp lore
forall lore. Stm lore -> Exp lore
stmExp Stm lore
bnd)) Scope lore
types
    , letBoundStmDepth :: Int
letBoundStmDepth = Int
0
    , letBoundIndex :: Int -> IndexArray
letBoundIndex = \Int
k -> (Indexed -> Indexed) -> Maybe Indexed -> Maybe Indexed
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Certificates -> Indexed -> Indexed
indexedAddCerts (StmAux (ExpDec lore) -> Certificates
forall dec. StmAux dec -> Certificates
stmAuxCerts (StmAux (ExpDec lore) -> Certificates)
-> StmAux (ExpDec lore) -> Certificates
forall a b. (a -> b) -> a -> b
$ Stm lore -> StmAux (ExpDec lore)
forall lore. Stm lore -> StmAux (ExpDec lore)
stmAux Stm lore
bnd)) (Maybe Indexed -> Maybe Indexed) -> IndexArray -> IndexArray
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                            SymbolTable lore -> Exp lore -> Int -> IndexArray
forall lore.
(IndexOp (Op lore), ASTLore lore) =>
SymbolTable lore -> Exp lore -> Int -> IndexArray
indexExp SymbolTable lore
vtable (Stm lore -> Exp lore
forall lore. Stm lore -> Exp lore
stmExp Stm lore
bnd) Int
k
    , letBoundConsumed :: Bool
letBoundConsumed = Bool
False
    }
  where ranges :: AS.RangesRep
        ranges :: RangesRep
ranges = SymbolTable lore -> RangesRep
forall lore. SymbolTable lore -> RangesRep
rangesRep SymbolTable lore
vtable

        types :: Scope lore
types = SymbolTable lore -> Scope lore
forall lore. SymbolTable lore -> Scope lore
toScope SymbolTable lore
vtable

        scalExpRange :: Range -> ScalExpRange
        scalExpRange :: Range -> ScalExpRange
scalExpRange (Bound
lower, Bound
upper) =
          ((ScalExpRange -> Maybe ScalExp) -> KnownBound -> Maybe ScalExp
scalExpBound ScalExpRange -> Maybe ScalExp
forall a b. (a, b) -> a
fst (KnownBound -> Maybe ScalExp) -> Bound -> Maybe ScalExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Bound
lower,
           (ScalExpRange -> Maybe ScalExp) -> KnownBound -> Maybe ScalExp
scalExpBound ScalExpRange -> Maybe ScalExp
forall a b. (a, b) -> b
snd (KnownBound -> Maybe ScalExp) -> Bound -> Maybe ScalExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Bound
upper)

        scalExpBound :: (ScalExpRange -> Maybe ScalExp)
                     -> Ranges.KnownBound
                     -> Maybe ScalExp
        scalExpBound :: (ScalExpRange -> Maybe ScalExp) -> KnownBound -> Maybe ScalExp
scalExpBound ScalExpRange -> Maybe ScalExp
pick (Ranges.VarBound VName
v) =
          ScalExpRange -> Maybe ScalExp
pick (ScalExpRange -> Maybe ScalExp) -> ScalExpRange -> Maybe ScalExp
forall a b. (a -> b) -> a -> b
$ VName -> SymbolTable lore -> ScalExpRange
forall lore. VName -> SymbolTable lore -> ScalExpRange
lookupRange VName
v SymbolTable lore
vtable
        scalExpBound ScalExpRange -> Maybe ScalExp
_ (Ranges.ScalarBound ScalExp
se) =
          ScalExp -> Maybe ScalExp
forall a. a -> Maybe a
Just ScalExp
se
        scalExpBound ScalExpRange -> Maybe ScalExp
_ (Ranges.MinimumBound KnownBound
b1 KnownBound
b2) = do
          ScalExp
b1' <- (ScalExpRange -> Maybe ScalExp) -> KnownBound -> Maybe ScalExp
scalExpBound ScalExpRange -> Maybe ScalExp
forall a b. (a, b) -> a
fst KnownBound
b1
          ScalExp
b2' <- (ScalExpRange -> Maybe ScalExp) -> KnownBound -> Maybe ScalExp
scalExpBound ScalExpRange -> Maybe ScalExp
forall a b. (a, b) -> a
fst 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
MaxMin Bool
True [ScalExp
b1', ScalExp
b2']
        scalExpBound ScalExpRange -> Maybe ScalExp
_ (Ranges.MaximumBound KnownBound
b1 KnownBound
b2) = do
          ScalExp
b1' <- (ScalExpRange -> Maybe ScalExp) -> KnownBound -> Maybe ScalExp
scalExpBound ScalExpRange -> Maybe ScalExp
forall a b. (a, b) -> b
snd KnownBound
b1
          ScalExp
b2' <- (ScalExpRange -> Maybe ScalExp) -> KnownBound -> Maybe ScalExp
scalExpBound ScalExpRange -> Maybe ScalExp
forall a b. (a, b) -> b
snd 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
MaxMin Bool
False [ScalExp
b1', ScalExp
b2']

        simplifyRange :: ScalExpRange -> ScalExpRange
        simplifyRange :: ScalExpRange -> ScalExpRange
simplifyRange (Maybe ScalExp
lower, Maybe ScalExp
upper) =
          (Maybe ScalExp -> Maybe ScalExp
simplifyBound Maybe ScalExp
lower,
           Maybe ScalExp -> Maybe ScalExp
simplifyBound Maybe ScalExp
upper)

        simplifyBound :: Maybe ScalExp -> Maybe ScalExp
simplifyBound (Just ScalExp
se) | ScalExp -> PrimType
scalExpType ScalExp
se PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
== PrimType
int32 =
          ScalExp -> Maybe ScalExp
forall a. a -> Maybe a
Just (ScalExp -> Maybe ScalExp) -> ScalExp -> Maybe ScalExp
forall a b. (a -> b) -> a -> b
$ ScalExp -> RangesRep -> ScalExp
AS.simplify ScalExp
se RangesRep
ranges
        simplifyBound Maybe ScalExp
_ =
          Maybe ScalExp
forall a. Maybe a
Nothing

bindingEntries :: (Ranged lore, Aliases.Aliased lore, IndexOp (Op lore)) =>
                  Stm lore -> SymbolTable lore
               -> [LetBoundEntry lore]
bindingEntries :: Stm lore -> SymbolTable lore -> [LetBoundEntry lore]
bindingEntries bnd :: Stm lore
bnd@(Let Pattern lore
pat StmAux (ExpDec lore)
_ Exp lore
_) SymbolTable lore
vtable = do
  PatElemT (LetDec lore)
pat_elem <- Pattern lore -> [PatElemT (LetDec lore)]
forall dec. PatternT dec -> [PatElemT dec]
patternElements Pattern lore
pat
  LetBoundEntry lore -> [LetBoundEntry lore]
forall (m :: * -> *) a. Monad m => a -> m a
return (LetBoundEntry lore -> [LetBoundEntry lore])
-> LetBoundEntry lore -> [LetBoundEntry lore]
forall a b. (a -> b) -> a -> b
$ SymbolTable lore
-> PatElemT (LetDec lore)
-> Range
-> Names
-> Stm lore
-> LetBoundEntry lore
forall lore.
(ASTLore lore, IndexOp (Op lore)) =>
SymbolTable lore
-> PatElem lore -> Range -> Names -> Stm lore -> LetBoundEntry lore
defBndEntry SymbolTable lore
vtable PatElemT (LetDec lore)
pat_elem
    (PatElemT (LetDec lore) -> Range
forall a. RangeOf a => a -> Range
Ranges.rangeOf PatElemT (LetDec lore)
pat_elem) (PatElemT (LetDec lore) -> Names
forall a. AliasesOf a => a -> Names
Aliases.aliasesOf PatElemT (LetDec lore)
pat_elem) Stm lore
bnd

insertEntry :: ASTLore lore =>
               VName -> Entry lore -> SymbolTable lore
            -> SymbolTable lore
insertEntry :: VName -> Entry lore -> SymbolTable lore -> SymbolTable lore
insertEntry VName
name Entry lore
entry =
  [(VName, Entry lore)] -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
[(VName, Entry lore)] -> SymbolTable lore -> SymbolTable lore
insertEntries [(VName
name,Entry lore
entry)]

insertEntries :: ASTLore lore =>
                 [(VName, Entry lore)] -> SymbolTable lore
              -> SymbolTable lore
insertEntries :: [(VName, Entry lore)] -> SymbolTable lore -> SymbolTable lore
insertEntries [(VName, Entry lore)]
entries SymbolTable lore
vtable =
  let vtable' :: SymbolTable lore
vtable' = SymbolTable lore
vtable { bindings :: Map VName (Entry lore)
bindings = (Map VName (Entry lore)
 -> (VName, Entry lore) -> Map VName (Entry lore))
-> Map VName (Entry lore)
-> [(VName, Entry lore)]
-> Map VName (Entry lore)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Map VName (Entry lore)
-> (VName, Entry lore) -> Map VName (Entry lore)
insertWithDepth (SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings SymbolTable lore
vtable) [(VName, Entry lore)]
entries }
  in (VName -> SymbolTable lore -> SymbolTable lore)
-> SymbolTable lore -> [VName] -> SymbolTable lore
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (VName -> Int -> SymbolTable lore -> SymbolTable lore
forall lore. VName -> Int -> SymbolTable lore -> SymbolTable lore
`isAtLeast` Int
0) SymbolTable lore
vtable' [VName]
dim_vars
  where insertWithDepth :: Map VName (Entry lore)
-> (VName, Entry lore) -> Map VName (Entry lore)
insertWithDepth Map VName (Entry lore)
bnds (VName
name, Entry lore
entry) =
          let entry' :: Entry lore
entry' = Int -> Entry lore -> Entry lore
forall lore. Int -> Entry lore -> Entry lore
setStmDepth (SymbolTable lore -> Int
forall lore. SymbolTable lore -> Int
loopDepth SymbolTable lore
vtable) Entry lore
entry
          in VName
-> Entry lore -> Map VName (Entry lore) -> Map VName (Entry lore)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
name Entry lore
entry' Map VName (Entry lore)
bnds
        dim_vars :: [VName]
dim_vars = [SubExp] -> [VName]
subExpVars ([SubExp] -> [VName]) -> [SubExp] -> [VName]
forall a b. (a -> b) -> a -> b
$ ((VName, Entry lore) -> [SubExp])
-> [(VName, Entry lore)] -> [SubExp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Type -> [SubExp]
forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
arrayDims (Type -> [SubExp])
-> ((VName, Entry lore) -> Type) -> (VName, Entry lore) -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry lore -> Type
forall lore. ASTLore lore => Entry lore -> Type
entryType (Entry lore -> Type)
-> ((VName, Entry lore) -> Entry lore)
-> (VName, Entry lore)
-> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, Entry lore) -> Entry lore
forall a b. (a, b) -> b
snd) [(VName, Entry lore)]
entries

insertStm :: (IndexOp (Op lore), Ranged lore, Aliases.Aliased lore) =>
             Stm lore
          -> SymbolTable lore
          -> SymbolTable lore
insertStm :: Stm lore -> SymbolTable lore -> SymbolTable lore
insertStm Stm lore
stm SymbolTable lore
vtable =
  (SymbolTable lore -> [VName] -> SymbolTable lore)
-> [VName] -> SymbolTable lore -> SymbolTable lore
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((SymbolTable lore -> VName -> SymbolTable lore)
-> SymbolTable lore -> [VName] -> SymbolTable lore
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((SymbolTable lore -> VName -> SymbolTable lore)
 -> SymbolTable lore -> [VName] -> SymbolTable lore)
-> (SymbolTable lore -> VName -> SymbolTable lore)
-> SymbolTable lore
-> [VName]
-> SymbolTable lore
forall a b. (a -> b) -> a -> b
$ (VName -> SymbolTable lore -> SymbolTable lore)
-> SymbolTable lore -> VName -> SymbolTable lore
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
VName -> SymbolTable lore -> SymbolTable lore
consume) (Names -> [VName]
namesToList Names
stm_consumed) (SymbolTable lore -> SymbolTable lore)
-> SymbolTable lore -> SymbolTable lore
forall a b. (a -> b) -> a -> b
$
  (SymbolTable lore -> [PatElemT (LetDec lore)] -> SymbolTable lore)
-> [PatElemT (LetDec lore)] -> SymbolTable lore -> SymbolTable lore
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((SymbolTable lore -> PatElemT (LetDec lore) -> SymbolTable lore)
-> SymbolTable lore -> [PatElemT (LetDec lore)] -> SymbolTable lore
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' SymbolTable lore -> PatElemT (LetDec lore) -> SymbolTable lore
forall dec lore.
AliasesOf dec =>
SymbolTable lore -> PatElemT dec -> SymbolTable lore
addRevAliases) (PatternT (LetDec lore) -> [PatElemT (LetDec lore)]
forall dec. PatternT dec -> [PatElemT dec]
patternElements (PatternT (LetDec lore) -> [PatElemT (LetDec lore)])
-> PatternT (LetDec lore) -> [PatElemT (LetDec lore)]
forall a b. (a -> b) -> a -> b
$ Stm lore -> PatternT (LetDec lore)
forall lore. Stm lore -> Pattern lore
stmPattern Stm lore
stm) (SymbolTable lore -> SymbolTable lore)
-> SymbolTable lore -> SymbolTable lore
forall a b. (a -> b) -> a -> b
$
  [(VName, Entry lore)] -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
[(VName, Entry lore)] -> SymbolTable lore -> SymbolTable lore
insertEntries ([VName] -> [Entry lore] -> [(VName, Entry lore)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
names ([Entry lore] -> [(VName, Entry lore)])
-> [Entry lore] -> [(VName, Entry lore)]
forall a b. (a -> b) -> a -> b
$ (LetBoundEntry lore -> Entry lore)
-> [LetBoundEntry lore] -> [Entry lore]
forall a b. (a -> b) -> [a] -> [b]
map LetBoundEntry lore -> Entry lore
forall lore. LetBoundEntry lore -> Entry lore
LetBound ([LetBoundEntry lore] -> [Entry lore])
-> [LetBoundEntry lore] -> [Entry lore]
forall a b. (a -> b) -> a -> b
$ Stm lore -> SymbolTable lore -> [LetBoundEntry lore]
forall lore.
(Ranged lore, Aliased lore, IndexOp (Op lore)) =>
Stm lore -> SymbolTable lore -> [LetBoundEntry lore]
bindingEntries Stm lore
stm SymbolTable lore
vtable) SymbolTable lore
vtable
  where names :: [VName]
names = PatternT (LetDec lore) -> [VName]
forall dec. PatternT dec -> [VName]
patternNames (PatternT (LetDec lore) -> [VName])
-> PatternT (LetDec lore) -> [VName]
forall a b. (a -> b) -> a -> b
$ Stm lore -> PatternT (LetDec lore)
forall lore. Stm lore -> Pattern lore
stmPattern Stm lore
stm
        adjustSeveral :: (a -> a) -> t k -> Map k a -> Map k a
adjustSeveral a -> a
f = (Map k a -> t k -> Map k a) -> t k -> Map k a -> Map k a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Map k a -> t k -> Map k a) -> t k -> Map k a -> Map k a)
-> (Map k a -> t k -> Map k a) -> t k -> Map k a -> Map k a
forall a b. (a -> b) -> a -> b
$ (Map k a -> k -> Map k a) -> Map k a -> t k -> Map k a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((Map k a -> k -> Map k a) -> Map k a -> t k -> Map k a)
-> (Map k a -> k -> Map k a) -> Map k a -> t k -> Map k a
forall a b. (a -> b) -> a -> b
$ (k -> Map k a -> Map k a) -> Map k a -> k -> Map k a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((k -> Map k a -> Map k a) -> Map k a -> k -> Map k a)
-> (k -> Map k a -> Map k a) -> Map k a -> k -> Map k a
forall a b. (a -> b) -> a -> b
$ (a -> a) -> k -> Map k a -> Map k a
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust a -> a
f
        stm_consumed :: Names
stm_consumed = Names -> SymbolTable lore -> Names
forall lore. Names -> SymbolTable lore -> Names
expandAliases (Stm lore -> Names
forall lore. Aliased lore => Stm lore -> Names
Aliases.consumedInStm Stm lore
stm) SymbolTable lore
vtable
        addRevAliases :: SymbolTable lore -> PatElemT dec -> SymbolTable lore
addRevAliases SymbolTable lore
vtable' PatElemT dec
pe =
          SymbolTable lore
vtable' { bindings :: Map VName (Entry lore)
bindings = (Entry lore -> Entry lore)
-> [VName] -> Map VName (Entry lore) -> Map VName (Entry lore)
forall (t :: * -> *) k a.
(Foldable t, Ord k) =>
(a -> a) -> t k -> Map k a -> Map k a
adjustSeveral Entry lore -> Entry lore
update [VName]
inedges (Map VName (Entry lore) -> Map VName (Entry lore))
-> Map VName (Entry lore) -> Map VName (Entry lore)
forall a b. (a -> b) -> a -> b
$ SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings SymbolTable lore
vtable' }
          where inedges :: [VName]
inedges = Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Names -> SymbolTable lore -> Names
forall lore. Names -> SymbolTable lore -> Names
expandAliases (PatElemT dec -> Names
forall a. AliasesOf a => a -> Names
Aliases.aliasesOf PatElemT dec
pe) SymbolTable lore
vtable'
                update :: Entry lore -> Entry lore
update (LetBound LetBoundEntry lore
entry) =
                  LetBoundEntry lore -> Entry lore
forall lore. LetBoundEntry lore -> Entry lore
LetBound LetBoundEntry lore
entry
                  { letBoundAliases :: Names
letBoundAliases = VName -> Names
oneName (PatElemT dec -> VName
forall dec. PatElemT dec -> VName
patElemName PatElemT dec
pe) Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> LetBoundEntry lore -> Names
forall lore. LetBoundEntry lore -> Names
letBoundAliases LetBoundEntry lore
entry }
                update (FParam FParamEntry lore
entry) =
                  FParamEntry lore -> Entry lore
forall lore. FParamEntry lore -> Entry lore
FParam FParamEntry lore
entry
                  { fparamAliases :: Names
fparamAliases = VName -> Names
oneName (PatElemT dec -> VName
forall dec. PatElemT dec -> VName
patElemName PatElemT dec
pe) Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> FParamEntry lore -> Names
forall lore. FParamEntry lore -> Names
fparamAliases FParamEntry lore
entry }
                update Entry lore
e = Entry lore
e

insertStms :: (IndexOp (Op lore), Ranged lore, Aliases.Aliased lore) =>
              Stms lore
           -> SymbolTable lore
           -> SymbolTable lore
insertStms :: Stms lore -> SymbolTable lore -> SymbolTable lore
insertStms Stms lore
stms SymbolTable lore
vtable = (SymbolTable lore -> Stm lore -> SymbolTable lore)
-> SymbolTable lore -> [Stm lore] -> SymbolTable lore
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((Stm lore -> SymbolTable lore -> SymbolTable lore)
-> SymbolTable lore -> Stm lore -> SymbolTable lore
forall a b c. (a -> b -> c) -> b -> a -> c
flip Stm lore -> SymbolTable lore -> SymbolTable lore
forall lore.
(IndexOp (Op lore), Ranged lore, Aliased lore) =>
Stm lore -> SymbolTable lore -> SymbolTable lore
insertStm) SymbolTable lore
vtable ([Stm lore] -> SymbolTable lore) -> [Stm lore] -> SymbolTable lore
forall a b. (a -> b) -> a -> b
$ Stms lore -> [Stm lore]
forall lore. Stms lore -> [Stm lore]
stmsToList Stms lore
stms

expandAliases :: Names -> SymbolTable lore -> Names
expandAliases :: Names -> SymbolTable lore -> Names
expandAliases Names
names SymbolTable lore
vtable = Names
names Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
aliasesOfAliases
  where aliasesOfAliases :: Names
aliasesOfAliases =
          [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ([Names] -> Names) -> (Names -> [Names]) -> Names -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Names) -> [VName] -> [Names]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SymbolTable lore -> Names
forall lore. VName -> SymbolTable lore -> Names
`lookupAliases` SymbolTable lore
vtable) ([VName] -> [Names]) -> (Names -> [VName]) -> Names -> [Names]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> [VName]
namesToList (Names -> Names) -> Names -> Names
forall a b. (a -> b) -> a -> b
$ Names
names

insertFParam :: ASTLore lore =>
                AST.FParam lore
             -> SymbolTable lore
             -> SymbolTable lore
insertFParam :: FParam lore -> SymbolTable lore -> SymbolTable lore
insertFParam FParam lore
fparam = (SymbolTable lore -> [VName] -> SymbolTable lore)
-> [VName] -> SymbolTable lore -> SymbolTable lore
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((VName -> SymbolTable lore -> SymbolTable lore)
-> SymbolTable lore -> [VName] -> SymbolTable lore
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (VName -> Int -> SymbolTable lore -> SymbolTable lore
forall lore. VName -> Int -> SymbolTable lore -> SymbolTable lore
`isAtLeast` Int
0)) [VName]
sizes (SymbolTable lore -> SymbolTable lore)
-> (SymbolTable lore -> SymbolTable lore)
-> SymbolTable lore
-> SymbolTable lore
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Entry lore -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
VName -> Entry lore -> SymbolTable lore -> SymbolTable lore
insertEntry VName
name Entry lore
entry
  where name :: VName
name = FParam lore -> VName
forall dec. Param dec -> VName
AST.paramName FParam lore
fparam
        entry :: Entry lore
entry = FParamEntry lore -> Entry lore
forall lore. FParamEntry lore -> Entry lore
FParam FParamEntry :: forall lore.
ScalExpRange
-> FParamInfo lore -> Names -> Int -> Bool -> FParamEntry lore
FParamEntry { fparamRange :: ScalExpRange
fparamRange = (Maybe ScalExp
forall a. Maybe a
Nothing, Maybe ScalExp
forall a. Maybe a
Nothing)
                                   , fparamDec :: FParamInfo lore
fparamDec = FParam lore -> FParamInfo lore
forall dec. Param dec -> dec
AST.paramDec FParam lore
fparam
                                   , fparamAliases :: Names
fparamAliases = Names
forall a. Monoid a => a
mempty
                                   , fparamStmDepth :: Int
fparamStmDepth = Int
0
                                   , fparamConsumed :: Bool
fparamConsumed = Bool
False
                                   }
        sizes :: [VName]
sizes = [SubExp] -> [VName]
subExpVars ([SubExp] -> [VName]) -> [SubExp] -> [VName]
forall a b. (a -> b) -> a -> b
$ Type -> [SubExp]
forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ FParam lore -> Type
forall dec. Typed dec => Param dec -> Type
AST.paramType FParam lore
fparam

insertFParams :: ASTLore lore =>
                 [AST.FParam lore]
              -> SymbolTable lore
              -> SymbolTable lore
insertFParams :: [FParam lore] -> SymbolTable lore -> SymbolTable lore
insertFParams [FParam lore]
fparams SymbolTable lore
symtable = (SymbolTable lore -> FParam lore -> SymbolTable lore)
-> SymbolTable lore -> [FParam lore] -> SymbolTable lore
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((FParam lore -> SymbolTable lore -> SymbolTable lore)
-> SymbolTable lore -> FParam lore -> SymbolTable lore
forall a b c. (a -> b -> c) -> b -> a -> c
flip FParam lore -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
FParam lore -> SymbolTable lore -> SymbolTable lore
insertFParam) SymbolTable lore
symtable [FParam lore]
fparams

insertLParamWithRange :: ASTLore lore =>
                         LParam lore -> ScalExpRange -> IndexArray -> SymbolTable lore
                      -> SymbolTable lore
insertLParamWithRange :: LParam lore
-> ScalExpRange
-> IndexArray
-> SymbolTable lore
-> SymbolTable lore
insertLParamWithRange LParam lore
param ScalExpRange
range IndexArray
indexf SymbolTable lore
vtable =
  -- We know that the sizes in the type of param are at least zero,
  -- since they are array sizes.
  let vtable' :: SymbolTable lore
vtable' = VName -> Entry lore -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
VName -> Entry lore -> SymbolTable lore -> SymbolTable lore
insertEntry VName
name Entry lore
bind SymbolTable lore
vtable
  in (VName -> SymbolTable lore -> SymbolTable lore)
-> SymbolTable lore -> [VName] -> SymbolTable lore
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (VName -> Int -> SymbolTable lore -> SymbolTable lore
forall lore. VName -> Int -> SymbolTable lore -> SymbolTable lore
`isAtLeast` Int
0) SymbolTable lore
vtable' [VName]
sizevars
  where bind :: Entry lore
bind = LParamEntry lore -> Entry lore
forall lore. LParamEntry lore -> Entry lore
LParam LParamEntry :: forall lore.
ScalExpRange
-> LParamInfo lore -> Int -> IndexArray -> Bool -> LParamEntry lore
LParamEntry { lparamRange :: ScalExpRange
lparamRange = ScalExpRange
range
                                  , lparamDec :: LParamInfo lore
lparamDec = LParam lore -> LParamInfo lore
forall dec. Param dec -> dec
AST.paramDec LParam lore
param
                                  , lparamStmDepth :: Int
lparamStmDepth = Int
0
                                  , lparamIndex :: IndexArray
lparamIndex = IndexArray
indexf
                                  , lparamConsumed :: Bool
lparamConsumed = Bool
False
                                  }
        name :: VName
name = LParam lore -> VName
forall dec. Param dec -> VName
AST.paramName LParam lore
param
        sizevars :: [VName]
sizevars = [SubExp] -> [VName]
subExpVars ([SubExp] -> [VName]) -> [SubExp] -> [VName]
forall a b. (a -> b) -> a -> b
$ Type -> [SubExp]
forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ LParam lore -> Type
forall dec. Typed dec => Param dec -> Type
AST.paramType LParam lore
param

insertLParam :: ASTLore lore =>
                LParam lore -> SymbolTable lore -> SymbolTable lore
insertLParam :: LParam lore -> SymbolTable lore -> SymbolTable lore
insertLParam LParam lore
param =
  LParam lore
-> ScalExpRange
-> IndexArray
-> SymbolTable lore
-> SymbolTable lore
forall lore.
ASTLore lore =>
LParam lore
-> ScalExpRange
-> IndexArray
-> SymbolTable lore
-> SymbolTable lore
insertLParamWithRange LParam lore
param (Maybe ScalExp
forall a. Maybe a
Nothing, Maybe ScalExp
forall a. Maybe a
Nothing) (Maybe Indexed -> IndexArray
forall a b. a -> b -> a
const Maybe Indexed
forall a. Maybe a
Nothing)

insertArrayLParam :: ASTLore lore =>
                     LParam lore -> Maybe VName -> SymbolTable lore
                  -> SymbolTable lore
insertArrayLParam :: LParam lore -> Maybe VName -> SymbolTable lore -> SymbolTable lore
insertArrayLParam LParam lore
param (Just VName
array) SymbolTable lore
vtable =
  -- We now know that the outer size of 'array' is at least one, and
  -- that the inner sizes are at least zero, since they are array
  -- sizes.
  let vtable' :: SymbolTable lore
vtable' = LParam lore
-> ScalExpRange
-> IndexArray
-> SymbolTable lore
-> SymbolTable lore
forall lore.
ASTLore lore =>
LParam lore
-> ScalExpRange
-> IndexArray
-> SymbolTable lore
-> SymbolTable lore
insertLParamWithRange LParam lore
param (VName -> SymbolTable lore -> ScalExpRange
forall lore. VName -> SymbolTable lore -> ScalExpRange
lookupRange VName
array SymbolTable lore
vtable) (Maybe Indexed -> IndexArray
forall a b. a -> b -> a
const Maybe Indexed
forall a. Maybe a
Nothing) SymbolTable lore
vtable
  in case 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
<$> VName -> SymbolTable lore -> Maybe Type
forall lore.
ASTLore lore =>
VName -> SymbolTable lore -> Maybe Type
lookupType VName
array SymbolTable lore
vtable of
    Just (Var VName
v:[SubExp]
_) -> (VName
v VName -> Int -> SymbolTable lore -> SymbolTable lore
forall lore. VName -> Int -> SymbolTable lore -> SymbolTable lore
`isAtLeast` Int
1) SymbolTable lore
vtable'
    Maybe [SubExp]
_              -> SymbolTable lore
vtable'
insertArrayLParam LParam lore
param Maybe VName
Nothing SymbolTable lore
vtable =
  -- Well, we still know that it's a param...
  LParam lore -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
LParam lore -> SymbolTable lore -> SymbolTable lore
insertLParam LParam lore
param SymbolTable lore
vtable

insertLoopVar :: ASTLore lore => VName -> IntType -> SubExp -> SymbolTable lore -> SymbolTable lore
insertLoopVar :: VName -> IntType -> SubExp -> SymbolTable lore -> SymbolTable lore
insertLoopVar VName
name IntType
it SubExp
bound = VName -> Entry lore -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
VName -> Entry lore -> SymbolTable lore -> SymbolTable lore
insertEntry VName
name Entry lore
bind
  where bind :: Entry lore
bind = LoopVarEntry lore -> Entry lore
forall lore. LoopVarEntry lore -> Entry lore
LoopVar LoopVarEntry :: forall lore. ScalExpRange -> Int -> IntType -> LoopVarEntry lore
LoopVarEntry {
            loopVarRange :: ScalExpRange
loopVarRange = (ScalExp -> Maybe ScalExp
forall a. a -> Maybe a
Just ScalExp
0,
                            ScalExp -> Maybe ScalExp
forall a. a -> Maybe a
Just (ScalExp -> Maybe ScalExp) -> ScalExp -> Maybe ScalExp
forall a b. (a -> b) -> a -> b
$ SubExp -> PrimType -> ScalExp
subExpToScalExp SubExp
bound (IntType -> PrimType
IntType IntType
it) ScalExp -> ScalExp -> ScalExp
forall a. Num a => a -> a -> a
- ScalExp
1)
          , loopVarStmDepth :: Int
loopVarStmDepth = Int
0
          , loopVarType :: IntType
loopVarType = IntType
it
          }

insertFreeVar :: ASTLore lore => VName -> NameInfo lore -> SymbolTable lore -> SymbolTable lore
insertFreeVar :: VName -> NameInfo lore -> SymbolTable lore -> SymbolTable lore
insertFreeVar VName
name NameInfo lore
dec = VName -> Entry lore -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
VName -> Entry lore -> SymbolTable lore -> SymbolTable lore
insertEntry VName
name Entry lore
entry
  where entry :: Entry lore
entry = FreeVarEntry lore -> Entry lore
forall lore. FreeVarEntry lore -> Entry lore
FreeVar FreeVarEntry :: forall lore.
NameInfo lore
-> Int
-> ScalExpRange
-> (VName -> IndexArray)
-> Bool
-> FreeVarEntry lore
FreeVarEntry {
            freeVarDec :: NameInfo lore
freeVarDec = NameInfo lore
dec
          , freeVarRange :: ScalExpRange
freeVarRange = (Maybe ScalExp
forall a. Maybe a
Nothing, Maybe ScalExp
forall a. Maybe a
Nothing)
          , freeVarStmDepth :: Int
freeVarStmDepth = Int
0
          , freeVarIndex :: VName -> IndexArray
freeVarIndex  = \VName
_ [PrimExp VName]
_ -> Maybe Indexed
forall a. Maybe a
Nothing
          , freeVarConsumed :: Bool
freeVarConsumed = Bool
False
          }

updateBounds :: ASTLore lore => Bool -> SubExp -> SymbolTable lore -> SymbolTable lore
updateBounds :: Bool -> SubExp -> SymbolTable lore -> SymbolTable lore
updateBounds Bool
isTrue SubExp
cond SymbolTable lore
vtable =
  case Reader (Scope lore) (Maybe ScalExp) -> Scope lore -> Maybe ScalExp
forall r a. Reader r a -> r -> a
runReader (LookupVar -> Exp Any -> Reader (Scope lore) (Maybe ScalExp)
forall t (f :: * -> *) lore.
(HasScope t f, Monad f) =>
LookupVar -> Exp lore -> f (Maybe ScalExp)
toScalExp (VName -> SymbolTable lore -> Maybe ScalExp
forall lore.
ASTLore lore =>
VName -> SymbolTable lore -> Maybe ScalExp
`lookupScalExp` SymbolTable lore
vtable) (Exp Any -> Reader (Scope lore) (Maybe ScalExp))
-> Exp Any -> Reader (Scope lore) (Maybe ScalExp)
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp Any
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp Any) -> BasicOp -> Exp Any
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
cond) Scope lore
types of
    Maybe ScalExp
Nothing    -> SymbolTable lore
vtable
    Just ScalExp
cond' ->
      let cond'' :: ScalExp
cond'' | Bool
isTrue    = ScalExp
cond'
                 | Bool
otherwise = ScalExp -> ScalExp
SNot ScalExp
cond'
      in ScalExp -> SymbolTable lore -> SymbolTable lore
forall lore. ScalExp -> SymbolTable lore -> SymbolTable lore
updateBounds' ScalExp
cond'' SymbolTable lore
vtable
  where types :: Scope lore
types = SymbolTable lore -> Scope lore
forall lore. SymbolTable lore -> Scope lore
toScope SymbolTable lore
vtable

-- | Updating the ranges of all symbols whenever we enter a branch is
-- presently too expensive, and disabled here.
noUpdateBounds :: Bool
noUpdateBounds :: Bool
noUpdateBounds = Bool
True

-- | Refines the ranges in the symbol table with
--     ranges extracted from branch conditions.
--   @cond@ is the condition of the if-branch.
updateBounds' :: ScalExp -> SymbolTable lore -> SymbolTable lore
updateBounds' :: ScalExp -> SymbolTable lore -> SymbolTable lore
updateBounds' ScalExp
_ SymbolTable lore
sym_tab | Bool
noUpdateBounds = SymbolTable lore
sym_tab
updateBounds' ScalExp
cond SymbolTable lore
sym_tab =
  ((VName, Bool, ScalExp) -> SymbolTable lore -> SymbolTable lore)
-> SymbolTable lore -> [(VName, Bool, ScalExp)] -> SymbolTable lore
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (VName, Bool, ScalExp) -> SymbolTable lore -> SymbolTable lore
forall lore.
(VName, Bool, ScalExp) -> SymbolTable lore -> SymbolTable lore
updateBound SymbolTable lore
sym_tab ([(VName, Bool, ScalExp)] -> SymbolTable lore)
-> [(VName, Bool, ScalExp)] -> SymbolTable lore
forall a b. (a -> b) -> a -> b
$ (ScalExp -> Maybe (VName, Bool, ScalExp))
-> [ScalExp] -> [(VName, Bool, ScalExp)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ScalExp -> Maybe (VName, Bool, ScalExp)
solve_leq0 ([ScalExp] -> [(VName, Bool, ScalExp)])
-> [ScalExp] -> [(VName, Bool, ScalExp)]
forall a b. (a -> b) -> a -> b
$
  ScalExp -> [ScalExp]
getNotFactorsLEQ0 (ScalExp -> [ScalExp]) -> ScalExp -> [ScalExp]
forall a b. (a -> b) -> a -> b
$ ScalExp -> RangesRep -> ScalExp
AS.simplify (ScalExp -> ScalExp
SNot ScalExp
cond) RangesRep
ranges
    where
      updateBound :: (VName, Bool, ScalExp) -> SymbolTable lore -> SymbolTable lore
updateBound (VName
sym,Bool
True ,ScalExp
bound) = VName -> ScalExp -> SymbolTable lore -> SymbolTable lore
forall lore.
VName -> ScalExp -> SymbolTable lore -> SymbolTable lore
setUpperBound VName
sym ScalExp
bound
      updateBound (VName
sym,Bool
False,ScalExp
bound) = VName -> ScalExp -> SymbolTable lore -> SymbolTable lore
forall lore.
VName -> ScalExp -> SymbolTable lore -> SymbolTable lore
setLowerBound VName
sym ScalExp
bound

      ranges :: RangesRep
ranges = ((Int, Maybe ScalExp, Maybe ScalExp) -> Bool)
-> RangesRep -> RangesRep
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter (Int, Maybe ScalExp, Maybe ScalExp) -> Bool
forall a a a. (a, Maybe a, Maybe a) -> Bool
nonEmptyRange (RangesRep -> RangesRep) -> RangesRep -> RangesRep
forall a b. (a -> b) -> a -> b
$ (Entry lore -> (Int, Maybe ScalExp, Maybe ScalExp))
-> Map VName (Entry lore) -> RangesRep
forall a b k. (a -> b) -> Map k a -> Map k b
M.map Entry lore -> (Int, Maybe ScalExp, Maybe ScalExp)
forall lore. Entry lore -> (Int, Maybe ScalExp, Maybe ScalExp)
toRep (Map VName (Entry lore) -> RangesRep)
-> Map VName (Entry lore) -> RangesRep
forall a b. (a -> b) -> a -> b
$ SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings SymbolTable lore
sym_tab
      toRep :: Entry lore -> (Int, Maybe ScalExp, Maybe ScalExp)
toRep Entry lore
entry = (Entry lore -> Int
forall lore. Entry lore -> Int
bindingDepth Entry lore
entry, Maybe ScalExp
lower, Maybe ScalExp
upper)
        where (Maybe ScalExp
lower, Maybe ScalExp
upper) = Entry lore -> ScalExpRange
forall lore. Entry lore -> ScalExpRange
valueRange Entry lore
entry
      nonEmptyRange :: (a, Maybe a, Maybe a) -> Bool
nonEmptyRange (a
_, Maybe a
lower, Maybe a
upper) = Maybe a -> Bool
forall a. Maybe a -> Bool
isJust Maybe a
lower Bool -> Bool -> Bool
|| Maybe a -> Bool
forall a. Maybe a -> Bool
isJust Maybe a
upper

      -- | Input: a bool exp in DNF form, named @cond@
      --   It gets the terms of the argument,
      --         i.e., cond = c1 || ... || cn
      --   and negates them.
      --   Returns [not c1, ..., not cn], i.e., the factors
      --   of @not cond@ in CNF form: not cond = (not c1) && ... && (not cn)
      getNotFactorsLEQ0 :: ScalExp -> [ScalExp]
      getNotFactorsLEQ0 :: ScalExp -> [ScalExp]
getNotFactorsLEQ0 (RelExp RelOp0
rel ScalExp
e_scal) =
          if ScalExp -> PrimType
scalExpType ScalExp
e_scal PrimType -> PrimType -> Bool
forall a. Eq a => a -> a -> Bool
/= PrimType
int32 then []
          else let leq0_escal :: ScalExp
leq0_escal = if RelOp0
rel RelOp0 -> RelOp0 -> Bool
forall a. Eq a => a -> a -> Bool
== RelOp0
LTH0
                                then ScalExp -> ScalExp -> ScalExp
SMinus ScalExp
0 ScalExp
e_scal
                                else ScalExp -> ScalExp -> ScalExp
SMinus ScalExp
1 ScalExp
e_scal

               in  [ScalExp -> RangesRep -> ScalExp
AS.simplify ScalExp
leq0_escal RangesRep
ranges]
      getNotFactorsLEQ0 (SLogOr  ScalExp
e1 ScalExp
e2) = ScalExp -> [ScalExp]
getNotFactorsLEQ0 ScalExp
e1 [ScalExp] -> [ScalExp] -> [ScalExp]
forall a. [a] -> [a] -> [a]
++ ScalExp -> [ScalExp]
getNotFactorsLEQ0 ScalExp
e2
      getNotFactorsLEQ0 ScalExp
_ = []

      -- | Argument is scalar expression @e@.
      --    Implementation finds the symbol defined at
      --    the highest depth in expression @e@, call it @i@,
      --    and decomposes e = a*i + b.  If @a@ and @b@ are
      --    free of @i@, AND @a == 1 or -1@ THEN the upper/lower
      --    bound can be improved. Otherwise Nothing.
      --
      --  Returns: Nothing or
      --  Just (i, a == 1, -a*b), i.e., (symbol, isUpperBound, bound)
      solve_leq0 :: ScalExp -> Maybe (VName, Bool, ScalExp)
      solve_leq0 :: ScalExp -> Maybe (VName, Bool, ScalExp)
solve_leq0 ScalExp
e_scal = do
        VName
sym <- Set VName -> ScalExp -> Maybe VName
pickRefinedSym Set VName
forall a. Set a
S.empty ScalExp
e_scal
        (ScalExp
a,ScalExp
b) <- (Error -> Maybe (ScalExp, ScalExp))
-> (Maybe (ScalExp, ScalExp) -> Maybe (ScalExp, ScalExp))
-> Either Error (Maybe (ScalExp, ScalExp))
-> Maybe (ScalExp, ScalExp)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe (ScalExp, ScalExp) -> Error -> Maybe (ScalExp, ScalExp)
forall a b. a -> b -> a
const Maybe (ScalExp, ScalExp)
forall a. Maybe a
Nothing) Maybe (ScalExp, ScalExp) -> Maybe (ScalExp, ScalExp)
forall a. a -> a
id (Either Error (Maybe (ScalExp, ScalExp))
 -> Maybe (ScalExp, ScalExp))
-> Either Error (Maybe (ScalExp, ScalExp))
-> Maybe (ScalExp, ScalExp)
forall a b. (a -> b) -> a -> b
$ VName
-> ScalExp -> RangesRep -> Either Error (Maybe (ScalExp, ScalExp))
AS.linFormScalE VName
sym ScalExp
e_scal RangesRep
ranges
        case ScalExp
a of
          -1 ->
            (VName, Bool, ScalExp) -> Maybe (VName, Bool, ScalExp)
forall a. a -> Maybe a
Just (VName
sym, Bool
False, ScalExp
b)
          ScalExp
1  ->
            let mb :: ScalExp
mb = ScalExp -> RangesRep -> ScalExp
AS.simplify (ScalExp -> ScalExp
forall a. Num a => a -> a
negate ScalExp
b) RangesRep
ranges
            in (VName, Bool, ScalExp) -> Maybe (VName, Bool, ScalExp)
forall a. a -> Maybe a
Just (VName
sym, Bool
True, ScalExp
mb)
          ScalExp
_ -> Maybe (VName, Bool, ScalExp)
forall a. Maybe a
Nothing

      -- When picking a symbols, @sym@ whose bound it is to be refined:
      -- make sure that @sym@ does not belong to the transitive closure
      -- of the symbols apearing in the ranges of all the other symbols
      -- in the sclar expression (themselves included).
      -- If this does not hold, pick another symbol, rinse and repeat.
      pickRefinedSym :: S.Set VName -> ScalExp -> Maybe VName
      pickRefinedSym :: Set VName -> ScalExp -> Maybe VName
pickRefinedSym Set VName
elsyms0 ScalExp
e_scal = do
        let candidates :: Names
candidates = ScalExp -> Names
forall a. FreeIn a => a -> Names
freeIn ScalExp
e_scal
            sym0 :: Maybe VName
sym0 = RangesRep -> Set VName -> ScalExp -> Maybe VName
AS.pickSymToElim RangesRep
ranges Set VName
elsyms0 ScalExp
e_scal
        case Maybe VName
sym0 of
            Just VName
sy -> let trclsyms :: Names
trclsyms = (Names -> VName -> Names) -> Names -> [VName] -> Names
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Names -> VName -> Names
trClSymsInRange Names
forall a. Monoid a => a
mempty ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$
                                      Names
candidates Names -> Names -> Names
`namesSubtract` VName -> Names
oneName VName
sy
                       in  if   VName
sy VName -> Names -> Bool
`nameIn` Names
trclsyms
                           then Set VName -> ScalExp -> Maybe VName
pickRefinedSym (VName -> Set VName -> Set VName
forall a. Ord a => a -> Set a -> Set a
S.insert VName
sy Set VName
elsyms0) ScalExp
e_scal
                           else Maybe VName
sym0
            Maybe VName
Nothing -> Maybe VName
sym0

      -- computes the transitive closure of the symbols appearing
      -- in the ranges of a symbol
      trClSymsInRange :: Names -> VName -> Names
      trClSymsInRange :: Names -> VName -> Names
trClSymsInRange Names
cur_syms VName
sym =
        if VName
sym VName -> Names -> Bool
`nameIn` Names
cur_syms then Names
cur_syms
        else case VName -> RangesRep -> Maybe (Int, Maybe ScalExp, Maybe ScalExp)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
sym RangesRep
ranges of
               Just (Int
_,Maybe ScalExp
lb,Maybe ScalExp
ub) -> let sym_bds :: [VName]
sym_bds = (ScalExp -> [VName]) -> [ScalExp] -> [VName]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Names -> [VName]
namesToList (Names -> [VName]) -> (ScalExp -> Names) -> ScalExp -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ScalExp -> Names
forall a. FreeIn a => a -> Names
freeIn) ([Maybe ScalExp] -> [ScalExp]
forall a. [Maybe a] -> [a]
catMaybes [Maybe ScalExp
lb, Maybe ScalExp
ub])
                                 in  (Names -> VName -> Names) -> Names -> [VName] -> Names
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Names -> VName -> Names
trClSymsInRange
                                           (VName -> Names
oneName VName
sym Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
cur_syms)
                                           [VName]
sym_bds
               Maybe (Int, Maybe ScalExp, Maybe ScalExp)
Nothing        -> VName -> Names
oneName VName
sym Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Names
cur_syms

consume :: ASTLore lore => VName -> SymbolTable lore -> SymbolTable lore
consume :: VName -> SymbolTable lore -> SymbolTable lore
consume VName
consumee SymbolTable lore
vtable = (SymbolTable lore -> VName -> SymbolTable lore)
-> SymbolTable lore -> [VName] -> SymbolTable lore
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' SymbolTable lore -> VName -> SymbolTable lore
consume' SymbolTable lore
vtable ([VName] -> SymbolTable lore) -> [VName] -> SymbolTable lore
forall a b. (a -> b) -> a -> b
$ Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$
                          Names -> SymbolTable lore -> Names
forall lore. Names -> SymbolTable lore -> Names
expandAliases (VName -> Names
oneName VName
consumee) SymbolTable lore
vtable
  where consume' :: SymbolTable lore -> VName -> SymbolTable lore
consume' SymbolTable lore
vtable' VName
v | Just Entry lore
e <- VName -> SymbolTable lore -> Maybe (Entry lore)
forall lore. VName -> SymbolTable lore -> Maybe (Entry lore)
lookup VName
v SymbolTable lore
vtable = VName -> Entry lore -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
VName -> Entry lore -> SymbolTable lore -> SymbolTable lore
insertEntry VName
v (Entry lore -> Entry lore
forall lore. Entry lore -> Entry lore
consume'' Entry lore
e) SymbolTable lore
vtable'
                           | Bool
otherwise                 = SymbolTable lore
vtable'
        consume'' :: Entry lore -> Entry lore
consume'' (FreeVar FreeVarEntry lore
e)  = FreeVarEntry lore -> Entry lore
forall lore. FreeVarEntry lore -> Entry lore
FreeVar FreeVarEntry lore
e { freeVarConsumed :: Bool
freeVarConsumed = Bool
True }
        consume'' (LetBound LetBoundEntry lore
e) = LetBoundEntry lore -> Entry lore
forall lore. LetBoundEntry lore -> Entry lore
LetBound LetBoundEntry lore
e { letBoundConsumed :: Bool
letBoundConsumed = Bool
True }
        consume'' (FParam FParamEntry lore
e)   = FParamEntry lore -> Entry lore
forall lore. FParamEntry lore -> Entry lore
FParam FParamEntry lore
e { fparamConsumed :: Bool
fparamConsumed = Bool
True }
        consume'' (LParam LParamEntry lore
e)   = LParamEntry lore -> Entry lore
forall lore. LParamEntry lore -> Entry lore
LParam LParamEntry lore
e { lparamConsumed :: Bool
lparamConsumed = Bool
True }
        consume'' (LoopVar LoopVarEntry lore
e)  = LoopVarEntry lore -> Entry lore
forall lore. LoopVarEntry lore -> Entry lore
LoopVar LoopVarEntry lore
e

setUpperBound :: VName -> ScalExp -> SymbolTable lore
              -> SymbolTable lore
setUpperBound :: VName -> ScalExp -> SymbolTable lore -> SymbolTable lore
setUpperBound VName
name ScalExp
bound SymbolTable lore
vtable =
  SymbolTable lore
vtable { bindings :: Map VName (Entry lore)
bindings = (Entry lore -> Entry lore)
-> VName -> Map VName (Entry lore) -> Map VName (Entry lore)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust Entry lore -> Entry lore
setUpperBound' VName
name (Map VName (Entry lore) -> Map VName (Entry lore))
-> Map VName (Entry lore) -> Map VName (Entry lore)
forall a b. (a -> b) -> a -> b
$ SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings SymbolTable lore
vtable }
  where setUpperBound' :: Entry lore -> Entry lore
setUpperBound' Entry lore
entry =
          let (Maybe ScalExp
oldLowerBound, Maybe ScalExp
oldUpperBound) = Entry lore -> ScalExpRange
forall lore. Entry lore -> ScalExpRange
valueRange Entry lore
entry
          in if ScalExp -> Bool -> Maybe ScalExp -> Bool
alreadyTheBound ScalExp
bound Bool
True Maybe ScalExp
oldUpperBound
             then Entry lore
entry
             else ScalExpRange -> Entry lore -> Entry lore
forall lore. ScalExpRange -> Entry lore -> Entry lore
setValueRange
                  (Maybe ScalExp
oldLowerBound,
                   ScalExp -> Maybe ScalExp
forall a. a -> Maybe a
Just (ScalExp -> Maybe ScalExp) -> ScalExp -> Maybe ScalExp
forall a b. (a -> b) -> a -> b
$ ScalExp -> (ScalExp -> ScalExp) -> Maybe ScalExp -> ScalExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ScalExp
bound (Bool -> [ScalExp] -> ScalExp
MaxMin Bool
True ([ScalExp] -> ScalExp)
-> (ScalExp -> [ScalExp]) -> ScalExp -> ScalExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ScalExp -> [ScalExp] -> [ScalExp]
forall a. a -> [a] -> [a]
:[ScalExp
bound])) Maybe ScalExp
oldUpperBound)
                  Entry lore
entry

setLowerBound :: VName -> ScalExp -> SymbolTable lore -> SymbolTable lore
setLowerBound :: VName -> ScalExp -> SymbolTable lore -> SymbolTable lore
setLowerBound VName
name ScalExp
bound SymbolTable lore
vtable =
  SymbolTable lore
vtable { bindings :: Map VName (Entry lore)
bindings = (Entry lore -> Entry lore)
-> VName -> Map VName (Entry lore) -> Map VName (Entry lore)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust Entry lore -> Entry lore
setLowerBound' VName
name (Map VName (Entry lore) -> Map VName (Entry lore))
-> Map VName (Entry lore) -> Map VName (Entry lore)
forall a b. (a -> b) -> a -> b
$ SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings SymbolTable lore
vtable }
  where setLowerBound' :: Entry lore -> Entry lore
setLowerBound' Entry lore
entry =
          let (Maybe ScalExp
oldLowerBound, Maybe ScalExp
oldUpperBound) = Entry lore -> ScalExpRange
forall lore. Entry lore -> ScalExpRange
valueRange Entry lore
entry
          in if ScalExp -> Bool -> Maybe ScalExp -> Bool
alreadyTheBound ScalExp
bound Bool
False Maybe ScalExp
oldLowerBound
             then Entry lore
entry
             else ScalExpRange -> Entry lore -> Entry lore
forall lore. ScalExpRange -> Entry lore -> Entry lore
setValueRange
                  (ScalExp -> Maybe ScalExp
forall a. a -> Maybe a
Just (ScalExp -> Maybe ScalExp) -> ScalExp -> Maybe ScalExp
forall a b. (a -> b) -> a -> b
$ ScalExp -> (ScalExp -> ScalExp) -> Maybe ScalExp -> ScalExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ScalExp
bound (Bool -> [ScalExp] -> ScalExp
MaxMin Bool
False ([ScalExp] -> ScalExp)
-> (ScalExp -> [ScalExp]) -> ScalExp -> ScalExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ScalExp -> [ScalExp] -> [ScalExp]
forall a. a -> [a] -> [a]
:[ScalExp
bound])) Maybe ScalExp
oldLowerBound,
                   Maybe ScalExp
oldUpperBound)
                  Entry lore
entry

alreadyTheBound :: ScalExp -> Bool -> Maybe ScalExp -> Bool
alreadyTheBound :: ScalExp -> Bool -> Maybe ScalExp -> Bool
alreadyTheBound ScalExp
_ Bool
_ Maybe ScalExp
Nothing = Bool
False
alreadyTheBound ScalExp
new_bound Bool
b1 (Just ScalExp
cur_bound)
  | ScalExp
cur_bound ScalExp -> ScalExp -> Bool
forall a. Eq a => a -> a -> Bool
== ScalExp
new_bound = Bool
True
  | MaxMin Bool
b2 [ScalExp]
ses <- ScalExp
cur_bound = Bool
b1 Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
b2 Bool -> Bool -> Bool
&& (ScalExp
new_bound ScalExp -> [ScalExp] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`L.elem` [ScalExp]
ses)
  | Bool
otherwise = Bool
False

isAtLeast :: VName -> Int -> SymbolTable lore -> SymbolTable lore
isAtLeast :: VName -> Int -> SymbolTable lore -> SymbolTable lore
isAtLeast VName
name Int
x =
  VName -> ScalExp -> SymbolTable lore -> SymbolTable lore
forall lore.
VName -> ScalExp -> SymbolTable lore -> SymbolTable lore
setLowerBound VName
name (ScalExp -> SymbolTable lore -> SymbolTable lore)
-> ScalExp -> SymbolTable lore -> SymbolTable lore
forall a b. (a -> b) -> a -> b
$ Int -> ScalExp
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x

-- | Hide definitions of those entries that satisfy some predicate.
hideIf :: (Entry lore -> Bool) -> SymbolTable lore -> SymbolTable lore
hideIf :: (Entry lore -> Bool) -> SymbolTable lore -> SymbolTable lore
hideIf Entry lore -> Bool
hide SymbolTable lore
vtable = SymbolTable lore
vtable { bindings :: Map VName (Entry lore)
bindings = (Entry lore -> Entry lore)
-> Map VName (Entry lore) -> Map VName (Entry lore)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map Entry lore -> Entry lore
maybeHide (Map VName (Entry lore) -> Map VName (Entry lore))
-> Map VName (Entry lore) -> Map VName (Entry lore)
forall a b. (a -> b) -> a -> b
$ SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings SymbolTable lore
vtable }
  where maybeHide :: Entry lore -> Entry lore
maybeHide Entry lore
entry
          | Entry lore -> Bool
hide Entry lore
entry = FreeVarEntry lore -> Entry lore
forall lore. FreeVarEntry lore -> Entry lore
FreeVar FreeVarEntry :: forall lore.
NameInfo lore
-> Int
-> ScalExpRange
-> (VName -> IndexArray)
-> Bool
-> FreeVarEntry lore
FreeVarEntry { freeVarDec :: NameInfo lore
freeVarDec = Entry lore -> NameInfo lore
forall lore. Entry lore -> NameInfo lore
entryInfo Entry lore
entry
                                              , freeVarStmDepth :: Int
freeVarStmDepth = Entry lore -> Int
forall lore. Entry lore -> Int
bindingDepth Entry lore
entry
                                              , freeVarRange :: ScalExpRange
freeVarRange = Entry lore -> ScalExpRange
forall lore. Entry lore -> ScalExpRange
valueRange Entry lore
entry
                                              , freeVarIndex :: VName -> IndexArray
freeVarIndex = \VName
_ [PrimExp VName]
_ -> Maybe Indexed
forall a. Maybe a
Nothing
                                              , freeVarConsumed :: Bool
freeVarConsumed = Entry lore -> Bool
forall lore. Entry lore -> Bool
consumed Entry lore
entry
                                              }
          | Bool
otherwise = Entry lore
entry

-- | Hide these definitions, if they are protected by certificates in
-- the set of names.
hideCertified :: Names -> SymbolTable lore -> SymbolTable lore
hideCertified :: Names -> SymbolTable lore -> SymbolTable lore
hideCertified Names
to_hide = (Entry lore -> Bool) -> SymbolTable lore -> SymbolTable lore
forall lore.
(Entry lore -> Bool) -> SymbolTable lore -> SymbolTable lore
hideIf ((Entry lore -> Bool) -> SymbolTable lore -> SymbolTable lore)
-> (Entry lore -> Bool) -> SymbolTable lore -> SymbolTable lore
forall a b. (a -> b) -> a -> b
$ Bool -> (Stm lore -> Bool) -> Maybe (Stm lore) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False Stm lore -> Bool
hide (Maybe (Stm lore) -> Bool)
-> (Entry lore -> Maybe (Stm lore)) -> Entry lore -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry lore -> Maybe (Stm lore)
forall lore. Entry lore -> Maybe (Stm lore)
entryStm
  where hide :: Stm lore -> Bool
hide = (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Names -> Bool
`nameIn` Names
to_hide) ([VName] -> Bool) -> (Stm lore -> [VName]) -> Stm lore -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Certificates -> [VName]
unCertificates (Certificates -> [VName])
-> (Stm lore -> Certificates) -> Stm lore -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm lore -> Certificates
forall lore. Stm lore -> Certificates
stmCerts