{-# LANGUAGE TypeFamilies #-}

module Futhark.Analysis.SymbolTable
  ( SymbolTable (bindings, loopDepth, availableAtClosestLoop, simplifyMemory),
    empty,
    fromScope,
    toScope,

    -- * Entries
    Entry,
    deepen,
    entryAccInput,
    entryDepth,
    entryLetBoundDec,
    entryIsSize,
    entryStm,
    entryFParam,
    entryLParam,

    -- * Lookup
    elem,
    lookup,
    lookupStm,
    lookupExp,
    lookupBasicOp,
    lookupType,
    lookupSubExp,
    lookupAliases,
    lookupLoopVar,
    lookupLoopParam,
    aliases,
    available,
    subExpAvailable,
    consume,
    index,
    index',
    Indexed (..),
    indexedAddCerts,
    IndexOp (..),

    -- * Insertion
    insertStm,
    insertStms,
    insertFParams,
    insertLParam,
    insertLoopVar,
    insertLoopMerge,

    -- * Misc
    hideCertified,
    noteAccTokens,
  )
where

import Control.Arrow ((&&&))
import Control.Monad
import Data.List (elemIndex, foldl')
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Ord
import Futhark.Analysis.PrimExp.Convert
import Futhark.IR hiding (FParam, lookupType)
import Futhark.IR qualified as AST
import Futhark.IR.Prop.Aliases qualified as Aliases
import Prelude hiding (elem, lookup)

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

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

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

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

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

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

deepen :: SymbolTable rep -> SymbolTable rep
deepen :: forall rep. SymbolTable rep -> SymbolTable rep
deepen SymbolTable rep
vtable =
  SymbolTable rep
vtable
    { loopDepth :: Int
loopDepth = SymbolTable rep -> Int
forall rep. SymbolTable rep -> Int
loopDepth SymbolTable rep
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 rep) -> [VName]
forall k a. Map k a -> [k]
M.keys (Map VName (Entry rep) -> [VName])
-> Map VName (Entry rep) -> [VName]
forall a b. (a -> b) -> a -> b
$ SymbolTable rep -> Map VName (Entry rep)
forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable
    }

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

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

instance FreeIn Indexed where
  freeIn' :: Indexed -> FV
freeIn' (Indexed Certs
cs PrimExp VName
v) = Certs -> FV
forall a. FreeIn a => a -> FV
freeIn' Certs
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 Certs
cs VName
arr [TPrimExp Int64 VName]
v) = Certs -> FV
forall a. FreeIn a => a -> FV
freeIn' Certs
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
<> [TPrimExp Int64 VName] -> FV
forall a. FreeIn a => a -> FV
freeIn' [TPrimExp Int64 VName]
v

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

data Entry rep = Entry
  { -- | True if consumed.
    forall rep. Entry rep -> Bool
entryConsumed :: Bool,
    forall rep. Entry rep -> Int
entryDepth :: Int,
    -- | True if this name has been used as an array size,
    -- implying that it is non-negative.
    forall rep. Entry rep -> Bool
entryIsSize :: Bool,
    -- | For names that are tokens of an accumulator, this is the
    -- corresponding combining function and neutral element.
    forall rep. Entry rep -> Maybe (WithAccInput rep)
entryAccInput :: Maybe (WithAccInput rep),
    forall rep. Entry rep -> EntryType rep
entryType :: EntryType rep
  }

data EntryType rep
  = LoopVar (LoopVarEntry rep)
  | LetBound (LetBoundEntry rep)
  | FParam (FParamEntry rep)
  | LParam (LParamEntry rep)
  | FreeVar (FreeVarEntry rep)

data LoopVarEntry rep = LoopVarEntry
  { forall {k} (rep :: k). LoopVarEntry rep -> IntType
loopVarType :: IntType,
    forall {k} (rep :: k). LoopVarEntry rep -> SubExp
loopVarBound :: SubExp
  }

data LetBoundEntry rep = LetBoundEntry
  { forall rep. LetBoundEntry rep -> LetDec rep
letBoundDec :: LetDec rep,
    forall rep. LetBoundEntry rep -> Names
letBoundAliases :: Names,
    forall rep. LetBoundEntry rep -> Stm rep
letBoundStm :: Stm rep,
    -- | Index a delayed array, if possible.
    forall rep. LetBoundEntry rep -> Int -> IndexArray
letBoundIndex :: Int -> IndexArray
  }

data FParamEntry rep = FParamEntry
  { forall rep. FParamEntry rep -> FParamInfo rep
fparamDec :: FParamInfo rep,
    forall rep. FParamEntry rep -> Names
fparamAliases :: Names,
    -- | If a loop parameter, the initial value and the eventual
    -- result.  The result need not be in scope in the symbol table.
    forall rep. FParamEntry rep -> Maybe (SubExp, SubExp)
fparamMerge :: Maybe (SubExp, SubExp)
  }

data LParamEntry rep = LParamEntry
  { forall rep. LParamEntry rep -> LParamInfo rep
lparamDec :: LParamInfo rep,
    forall rep. LParamEntry rep -> Names
lparamAliases :: Names,
    forall rep. LParamEntry rep -> IndexArray
lparamIndex :: IndexArray
  }

data FreeVarEntry rep = FreeVarEntry
  { forall rep. FreeVarEntry rep -> NameInfo rep
freeVarDec :: NameInfo rep,
    forall rep. FreeVarEntry rep -> Names
freeVarAliases :: Names,
    -- | Index a delayed array, if possible.
    forall rep. FreeVarEntry rep -> VName -> IndexArray
freeVarIndex :: VName -> IndexArray
  }

instance (ASTRep rep) => Typed (Entry rep) where
  typeOf :: Entry rep -> Type
typeOf = NameInfo rep -> Type
forall t. Typed t => t -> Type
typeOf (NameInfo rep -> Type)
-> (Entry rep -> NameInfo rep) -> Entry rep -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry rep -> NameInfo rep
forall rep. Entry rep -> NameInfo rep
entryInfo

entryInfo :: Entry rep -> NameInfo rep
entryInfo :: forall rep. Entry rep -> NameInfo rep
entryInfo Entry rep
e = case Entry rep -> EntryType rep
forall rep. Entry rep -> EntryType rep
entryType Entry rep
e of
  LetBound LetBoundEntry rep
entry -> LetDec rep -> NameInfo rep
forall rep. LetDec rep -> NameInfo rep
LetName (LetDec rep -> NameInfo rep) -> LetDec rep -> NameInfo rep
forall a b. (a -> b) -> a -> b
$ LetBoundEntry rep -> LetDec rep
forall rep. LetBoundEntry rep -> LetDec rep
letBoundDec LetBoundEntry rep
entry
  LoopVar LoopVarEntry rep
entry -> IntType -> NameInfo rep
forall rep. IntType -> NameInfo rep
IndexName (IntType -> NameInfo rep) -> IntType -> NameInfo rep
forall a b. (a -> b) -> a -> b
$ LoopVarEntry rep -> IntType
forall {k} (rep :: k). LoopVarEntry rep -> IntType
loopVarType LoopVarEntry rep
entry
  FParam FParamEntry rep
entry -> FParamInfo rep -> NameInfo rep
forall rep. FParamInfo rep -> NameInfo rep
FParamName (FParamInfo rep -> NameInfo rep) -> FParamInfo rep -> NameInfo rep
forall a b. (a -> b) -> a -> b
$ FParamEntry rep -> FParamInfo rep
forall rep. FParamEntry rep -> FParamInfo rep
fparamDec FParamEntry rep
entry
  LParam LParamEntry rep
entry -> LParamInfo rep -> NameInfo rep
forall rep. LParamInfo rep -> NameInfo rep
LParamName (LParamInfo rep -> NameInfo rep) -> LParamInfo rep -> NameInfo rep
forall a b. (a -> b) -> a -> b
$ LParamEntry rep -> LParamInfo rep
forall rep. LParamEntry rep -> LParamInfo rep
lparamDec LParamEntry rep
entry
  FreeVar FreeVarEntry rep
entry -> FreeVarEntry rep -> NameInfo rep
forall rep. FreeVarEntry rep -> NameInfo rep
freeVarDec FreeVarEntry rep
entry

isLetBound :: Entry rep -> Maybe (LetBoundEntry rep)
isLetBound :: forall rep. Entry rep -> Maybe (LetBoundEntry rep)
isLetBound Entry rep
e = case Entry rep -> EntryType rep
forall rep. Entry rep -> EntryType rep
entryType Entry rep
e of
  LetBound LetBoundEntry rep
entry -> LetBoundEntry rep -> Maybe (LetBoundEntry rep)
forall a. a -> Maybe a
Just LetBoundEntry rep
entry
  EntryType rep
_ -> Maybe (LetBoundEntry rep)
forall a. Maybe a
Nothing

entryStm :: Entry rep -> Maybe (Stm rep)
entryStm :: forall rep. Entry rep -> Maybe (Stm rep)
entryStm = (LetBoundEntry rep -> Stm rep)
-> Maybe (LetBoundEntry rep) -> Maybe (Stm rep)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap LetBoundEntry rep -> Stm rep
forall rep. LetBoundEntry rep -> Stm rep
letBoundStm (Maybe (LetBoundEntry rep) -> Maybe (Stm rep))
-> (Entry rep -> Maybe (LetBoundEntry rep))
-> Entry rep
-> Maybe (Stm rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry rep -> Maybe (LetBoundEntry rep)
forall rep. Entry rep -> Maybe (LetBoundEntry rep)
isLetBound

entryFParam :: Entry rep -> Maybe (FParamInfo rep)
entryFParam :: forall rep. Entry rep -> Maybe (FParamInfo rep)
entryFParam Entry rep
e = case Entry rep -> EntryType rep
forall rep. Entry rep -> EntryType rep
entryType Entry rep
e of
  FParam FParamEntry rep
e' -> FParamInfo rep -> Maybe (FParamInfo rep)
forall a. a -> Maybe a
Just (FParamInfo rep -> Maybe (FParamInfo rep))
-> FParamInfo rep -> Maybe (FParamInfo rep)
forall a b. (a -> b) -> a -> b
$ FParamEntry rep -> FParamInfo rep
forall rep. FParamEntry rep -> FParamInfo rep
fparamDec FParamEntry rep
e'
  EntryType rep
_ -> Maybe (FParamInfo rep)
forall a. Maybe a
Nothing

entryLParam :: Entry rep -> Maybe (LParamInfo rep)
entryLParam :: forall rep. Entry rep -> Maybe (LParamInfo rep)
entryLParam Entry rep
e = case Entry rep -> EntryType rep
forall rep. Entry rep -> EntryType rep
entryType Entry rep
e of
  LParam LParamEntry rep
e' -> LParamInfo rep -> Maybe (LParamInfo rep)
forall a. a -> Maybe a
Just (LParamInfo rep -> Maybe (LParamInfo rep))
-> LParamInfo rep -> Maybe (LParamInfo rep)
forall a b. (a -> b) -> a -> b
$ LParamEntry rep -> LParamInfo rep
forall rep. LParamEntry rep -> LParamInfo rep
lparamDec LParamEntry rep
e'
  EntryType rep
_ -> Maybe (LParamInfo rep)
forall a. Maybe a
Nothing

entryLetBoundDec :: Entry rep -> Maybe (LetDec rep)
entryLetBoundDec :: forall rep. Entry rep -> Maybe (LetDec rep)
entryLetBoundDec = (LetBoundEntry rep -> LetDec rep)
-> Maybe (LetBoundEntry rep) -> Maybe (LetDec rep)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap LetBoundEntry rep -> LetDec rep
forall rep. LetBoundEntry rep -> LetDec rep
letBoundDec (Maybe (LetBoundEntry rep) -> Maybe (LetDec rep))
-> (Entry rep -> Maybe (LetBoundEntry rep))
-> Entry rep
-> Maybe (LetDec rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry rep -> Maybe (LetBoundEntry rep)
forall rep. Entry rep -> Maybe (LetBoundEntry rep)
isLetBound

entryAliases :: EntryType rep -> Names
entryAliases :: forall rep. EntryType rep -> Names
entryAliases (LetBound LetBoundEntry rep
e) = LetBoundEntry rep -> Names
forall rep. LetBoundEntry rep -> Names
letBoundAliases LetBoundEntry rep
e
entryAliases (FParam FParamEntry rep
e) = FParamEntry rep -> Names
forall rep. FParamEntry rep -> Names
fparamAliases FParamEntry rep
e
entryAliases (LParam LParamEntry rep
e) = LParamEntry rep -> Names
forall rep. LParamEntry rep -> Names
lparamAliases LParamEntry rep
e
entryAliases (FreeVar FreeVarEntry rep
e) = FreeVarEntry rep -> Names
forall rep. FreeVarEntry rep -> Names
freeVarAliases FreeVarEntry rep
e
entryAliases (LoopVar LoopVarEntry rep
_) = Names
forall a. Monoid a => a
mempty -- Integers have no aliases.

-- | You almost always want 'available' instead of this one.
elem :: VName -> SymbolTable rep -> Bool
elem :: forall rep. VName -> SymbolTable rep -> Bool
elem VName
name = Maybe (Entry rep) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Entry rep) -> Bool)
-> (SymbolTable rep -> Maybe (Entry rep))
-> SymbolTable rep
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SymbolTable rep -> Maybe (Entry rep)
forall rep. VName -> SymbolTable rep -> Maybe (Entry rep)
lookup VName
name

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

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

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

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

lookupType :: (ASTRep rep) => VName -> SymbolTable rep -> Maybe Type
lookupType :: forall rep. ASTRep rep => VName -> SymbolTable rep -> Maybe Type
lookupType VName
name SymbolTable rep
vtable = Entry rep -> Type
forall t. Typed t => t -> Type
typeOf (Entry rep -> Type) -> Maybe (Entry rep) -> Maybe Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> SymbolTable rep -> Maybe (Entry rep)
forall rep. VName -> SymbolTable rep -> Maybe (Entry rep)
lookup VName
name SymbolTable rep
vtable

lookupSubExpType :: (ASTRep rep) => SubExp -> SymbolTable rep -> Maybe Type
lookupSubExpType :: forall rep. ASTRep rep => SubExp -> SymbolTable rep -> Maybe Type
lookupSubExpType (Var VName
v) = VName -> SymbolTable rep -> Maybe Type
forall rep. ASTRep rep => VName -> SymbolTable rep -> Maybe Type
lookupType VName
v
lookupSubExpType (Constant PrimValue
v) = Maybe Type -> SymbolTable rep -> Maybe Type
forall a b. a -> b -> a
const (Maybe Type -> SymbolTable rep -> Maybe Type)
-> Maybe Type -> SymbolTable rep -> 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 rep -> Maybe (SubExp, Certs)
lookupSubExp :: forall rep. VName -> SymbolTable rep -> Maybe (SubExp, Certs)
lookupSubExp VName
name SymbolTable rep
vtable = do
  (Exp rep
e, Certs
cs) <- VName -> SymbolTable rep -> Maybe (Exp rep, Certs)
forall rep. VName -> SymbolTable rep -> Maybe (Exp rep, Certs)
lookupExp VName
name SymbolTable rep
vtable
  case Exp rep
e of
    BasicOp (SubExp SubExp
se) -> (SubExp, Certs) -> Maybe (SubExp, Certs)
forall a. a -> Maybe a
Just (SubExp
se, Certs
cs)
    Exp rep
_ -> Maybe (SubExp, Certs)
forall a. Maybe a
Nothing

lookupAliases :: VName -> SymbolTable rep -> Names
lookupAliases :: forall rep. VName -> SymbolTable rep -> Names
lookupAliases VName
name SymbolTable rep
vtable =
  Names -> (Entry rep -> Names) -> Maybe (Entry rep) -> Names
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Names
forall a. Monoid a => a
mempty (EntryType rep -> Names
forall rep. EntryType rep -> Names
entryAliases (EntryType rep -> Names)
-> (Entry rep -> EntryType rep) -> Entry rep -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry rep -> EntryType rep
forall rep. Entry rep -> EntryType rep
entryType) (Maybe (Entry rep) -> Names) -> Maybe (Entry rep) -> Names
forall a b. (a -> b) -> a -> b
$ VName -> Map VName (Entry rep) -> Maybe (Entry rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (SymbolTable rep -> Map VName (Entry rep)
forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable)

-- | If the given variable name is the name of a 'ForLoop' parameter,
-- then return the bound of that loop.
lookupLoopVar :: VName -> SymbolTable rep -> Maybe SubExp
lookupLoopVar :: forall rep. VName -> SymbolTable rep -> Maybe SubExp
lookupLoopVar VName
name SymbolTable rep
vtable = do
  LoopVar LoopVarEntry rep
e <- Entry rep -> EntryType rep
forall rep. Entry rep -> EntryType rep
entryType (Entry rep -> EntryType rep)
-> Maybe (Entry rep) -> Maybe (EntryType rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Map VName (Entry rep) -> Maybe (Entry rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (SymbolTable rep -> Map VName (Entry rep)
forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable)
  SubExp -> Maybe SubExp
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> Maybe SubExp) -> SubExp -> Maybe SubExp
forall a b. (a -> b) -> a -> b
$ LoopVarEntry rep -> SubExp
forall {k} (rep :: k). LoopVarEntry rep -> SubExp
loopVarBound LoopVarEntry rep
e

-- | Look up the initial value and eventual result of a loop
-- parameter.  Note that the result almost certainly refers to
-- something that is not part of the symbol table.
lookupLoopParam :: VName -> SymbolTable rep -> Maybe (SubExp, SubExp)
lookupLoopParam :: forall rep. VName -> SymbolTable rep -> Maybe (SubExp, SubExp)
lookupLoopParam VName
name SymbolTable rep
vtable = do
  FParam FParamEntry rep
e <- Entry rep -> EntryType rep
forall rep. Entry rep -> EntryType rep
entryType (Entry rep -> EntryType rep)
-> Maybe (Entry rep) -> Maybe (EntryType rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Map VName (Entry rep) -> Maybe (Entry rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (SymbolTable rep -> Map VName (Entry rep)
forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable)
  FParamEntry rep -> Maybe (SubExp, SubExp)
forall rep. FParamEntry rep -> Maybe (SubExp, SubExp)
fparamMerge FParamEntry rep
e

-- | Do these two names alias each other?  This is expected to be a
-- commutative relationship, so the order of arguments does not
-- matter.
aliases :: VName -> VName -> SymbolTable rep -> Bool
aliases :: forall rep. VName -> VName -> SymbolTable rep -> Bool
aliases VName
x VName
y SymbolTable rep
vtable = VName
x VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
y Bool -> Bool -> Bool
|| (VName
x VName -> Names -> Bool
`nameIn` VName -> SymbolTable rep -> Names
forall rep. VName -> SymbolTable rep -> Names
lookupAliases VName
y SymbolTable rep
vtable)

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

-- | Constant or 'available'
subExpAvailable :: SubExp -> SymbolTable rep -> Bool
subExpAvailable :: forall rep. SubExp -> SymbolTable rep -> Bool
subExpAvailable (Var VName
name) = VName -> SymbolTable rep -> Bool
forall rep. VName -> SymbolTable rep -> Bool
available VName
name
subExpAvailable Constant {} = Bool -> SymbolTable rep -> Bool
forall a b. a -> b -> a
const Bool
True

index ::
  (ASTRep rep) =>
  VName ->
  [SubExp] ->
  SymbolTable rep ->
  Maybe Indexed
index :: forall rep.
ASTRep rep =>
VName -> [SubExp] -> SymbolTable rep -> Maybe Indexed
index VName
name [SubExp]
is SymbolTable rep
table = do
  [TPrimExp Int64 VName]
is' <- (SubExp -> Maybe (TPrimExp Int64 VName))
-> [SubExp] -> Maybe [TPrimExp Int64 VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SubExp -> Maybe (TPrimExp Int64 VName)
asPrimExp [SubExp]
is
  VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
forall rep.
VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
index' VName
name [TPrimExp Int64 VName]
is' SymbolTable rep
table
  where
    asPrimExp :: SubExp -> Maybe (TPrimExp Int64 VName)
asPrimExp SubExp
i = do
      Prim PrimType
t <- SubExp -> SymbolTable rep -> Maybe Type
forall rep. ASTRep rep => SubExp -> SymbolTable rep -> Maybe Type
lookupSubExpType SubExp
i SymbolTable rep
table
      TPrimExp Int64 VName -> Maybe (TPrimExp Int64 VName)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TPrimExp Int64 VName -> Maybe (TPrimExp Int64 VName))
-> TPrimExp Int64 VName -> Maybe (TPrimExp Int64 VName)
forall a b. (a -> b) -> a -> b
$ PrimExp VName -> TPrimExp Int64 VName
forall {k} (t :: k) v. PrimExp v -> TPrimExp t v
TPrimExp (PrimExp VName -> TPrimExp Int64 VName)
-> PrimExp VName -> TPrimExp Int64 VName
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> PrimExp VName
primExpFromSubExp PrimType
t SubExp
i

index' ::
  VName ->
  [TPrimExp Int64 VName] ->
  SymbolTable rep ->
  Maybe Indexed
index' :: forall rep.
VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
index' VName
name [TPrimExp Int64 VName]
is SymbolTable rep
vtable = do
  Entry rep
entry <- VName -> SymbolTable rep -> Maybe (Entry rep)
forall rep. VName -> SymbolTable rep -> Maybe (Entry rep)
lookup VName
name SymbolTable rep
vtable
  case Entry rep -> EntryType rep
forall rep. Entry rep -> EntryType rep
entryType Entry rep
entry of
    LetBound LetBoundEntry rep
entry'
      | Just Int
k <-
          VName -> [VName] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex VName
name ([VName] -> Maybe Int)
-> (Stm rep -> [VName]) -> Stm rep -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames (Pat (LetDec rep) -> [VName])
-> (Stm rep -> Pat (LetDec rep)) -> Stm rep -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm rep -> Pat (LetDec rep)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat (Stm rep -> Maybe Int) -> Stm rep -> Maybe Int
forall a b. (a -> b) -> a -> b
$
            LetBoundEntry rep -> Stm rep
forall rep. LetBoundEntry rep -> Stm rep
letBoundStm LetBoundEntry rep
entry' ->
          LetBoundEntry rep -> Int -> IndexArray
forall rep. LetBoundEntry rep -> Int -> IndexArray
letBoundIndex LetBoundEntry rep
entry' Int
k [TPrimExp Int64 VName]
is
    FreeVar FreeVarEntry rep
entry' ->
      FreeVarEntry rep -> VName -> IndexArray
forall rep. FreeVarEntry rep -> VName -> IndexArray
freeVarIndex FreeVarEntry rep
entry' VName
name [TPrimExp Int64 VName]
is
    LParam LParamEntry rep
entry' -> LParamEntry rep -> IndexArray
forall rep. LParamEntry rep -> IndexArray
lparamIndex LParamEntry rep
entry' [TPrimExp Int64 VName]
is
    EntryType rep
_ -> Maybe Indexed
forall a. Maybe a
Nothing

class IndexOp op where
  indexOp ::
    (ASTRep rep, IndexOp (Op rep)) =>
    SymbolTable rep ->
    Int ->
    op ->
    [TPrimExp Int64 VName] ->
    Maybe Indexed
  indexOp SymbolTable rep
_ Int
_ op
_ [TPrimExp Int64 VName]
_ = Maybe Indexed
forall a. Maybe a
Nothing

instance IndexOp (NoOp rep)

indexExp ::
  (IndexOp (Op rep), ASTRep rep) =>
  SymbolTable rep ->
  Exp rep ->
  Int ->
  IndexArray
indexExp :: forall rep.
(IndexOp (Op rep), ASTRep rep) =>
SymbolTable rep -> Exp rep -> Int -> IndexArray
indexExp SymbolTable rep
vtable (Op Op rep
op) Int
k [TPrimExp Int64 VName]
is =
  SymbolTable rep -> Int -> Op rep -> IndexArray
forall rep.
(ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep -> Int -> Op rep -> IndexArray
forall op rep.
(IndexOp op, ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep -> Int -> op -> IndexArray
indexOp SymbolTable rep
vtable Int
k Op rep
op [TPrimExp Int64 VName]
is
indexExp SymbolTable rep
_ (BasicOp (Iota SubExp
_ SubExp
x SubExp
s IntType
to_it)) Int
_ [TPrimExp Int64 VName
i] =
  Indexed -> Maybe Indexed
forall a. a -> Maybe a
Just (Indexed -> Maybe Indexed) -> Indexed -> Maybe Indexed
forall a b. (a -> b) -> a -> b
$
    Certs -> PrimExp VName -> Indexed
Indexed Certs
forall a. Monoid a => a
mempty (PrimExp VName -> Indexed) -> PrimExp VName -> Indexed
forall a b. (a -> b) -> a -> b
$
      ( IntType -> PrimExp VName -> PrimExp VName
forall v. IntType -> PrimExp v -> PrimExp v
sExt IntType
to_it (TPrimExp Int64 VName -> PrimExp VName
forall {k} (t :: k) v. TPrimExp t v -> PrimExp v
untyped TPrimExp Int64 VName
i)
          PrimExp VName -> PrimExp VName -> PrimExp VName
`mul` PrimType -> SubExp -> PrimExp VName
primExpFromSubExp (IntType -> PrimType
IntType IntType
to_it) SubExp
s
      )
        PrimExp VName -> PrimExp VName -> PrimExp VName
`add` PrimType -> SubExp -> PrimExp VName
primExpFromSubExp (IntType -> PrimType
IntType IntType
to_it) SubExp
x
  where
    mul :: PrimExp VName -> PrimExp VName -> PrimExp VName
mul = BinOp -> PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
BinOpExp (IntType -> Overflow -> BinOp
Mul IntType
to_it Overflow
OverflowWrap)
    add :: PrimExp VName -> PrimExp VName -> PrimExp VName
add = BinOp -> PrimExp VName -> PrimExp VName -> PrimExp VName
forall v. BinOp -> PrimExp v -> PrimExp v -> PrimExp v
BinOpExp (IntType -> Overflow -> BinOp
Add IntType
to_it Overflow
OverflowWrap)
indexExp SymbolTable rep
table (BasicOp (Replicate (Shape [SubExp]
ds) SubExp
v)) Int
_ [TPrimExp Int64 VName]
is
  | [SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
ds Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [TPrimExp Int64 VName] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TPrimExp Int64 VName]
is,
    Just (Prim PrimType
t) <- SubExp -> SymbolTable rep -> Maybe Type
forall rep. ASTRep rep => SubExp -> SymbolTable rep -> Maybe Type
lookupSubExpType SubExp
v SymbolTable rep
table =
      Indexed -> Maybe Indexed
forall a. a -> Maybe a
Just (Indexed -> Maybe Indexed) -> Indexed -> Maybe Indexed
forall a b. (a -> b) -> a -> b
$ Certs -> PrimExp VName -> Indexed
Indexed Certs
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 rep
table (BasicOp (Replicate Shape
s (Var VName
v))) Int
_ [TPrimExp Int64 VName]
is = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ VName
v VName -> SymbolTable rep -> Bool
forall rep. VName -> SymbolTable rep -> Bool
`available` SymbolTable rep
table
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Shape
s Shape -> Shape -> Bool
forall a. Eq a => a -> a -> Bool
/= Shape
forall a. Monoid a => a
mempty
  VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
forall rep.
VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
index' VName
v (Int -> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. Int -> [a] -> [a]
drop (Shape -> Int
forall a. ArrayShape a => a -> Int
shapeRank Shape
s) [TPrimExp Int64 VName]
is) SymbolTable rep
table
indexExp SymbolTable rep
table (BasicOp (Reshape ReshapeKind
_ Shape
newshape VName
v)) Int
_ [TPrimExp Int64 VName]
is
  | Just [SubExp]
oldshape <- Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims (Type -> [SubExp]) -> Maybe Type -> Maybe [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> SymbolTable rep -> Maybe Type
forall rep. ASTRep rep => VName -> SymbolTable rep -> Maybe Type
lookupType VName
v SymbolTable rep
table =
      -- TODO: handle coercions more efficiently.
      let is' :: [TPrimExp Int64 VName]
is' =
            [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
-> [TPrimExp Int64 VName]
forall num. IntegralExp num => [num] -> [num] -> [num] -> [num]
reshapeIndex
              ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 [SubExp]
oldshape)
              ((SubExp -> TPrimExp Int64 VName)
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> TPrimExp Int64 VName
pe64 ([SubExp] -> [TPrimExp Int64 VName])
-> [SubExp] -> [TPrimExp Int64 VName]
forall a b. (a -> b) -> a -> b
$ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
newshape)
              [TPrimExp Int64 VName]
is
       in VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
forall rep.
VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
index' VName
v [TPrimExp Int64 VName]
is' SymbolTable rep
table
indexExp SymbolTable rep
table (BasicOp (Index VName
v Slice SubExp
slice)) Int
_ [TPrimExp Int64 VName]
is = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ VName
v VName -> SymbolTable rep -> Bool
forall rep. VName -> SymbolTable rep -> Bool
`available` SymbolTable rep
table
  VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
forall rep.
VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
index' VName
v ([DimIndex SubExp]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
adjust (Slice SubExp -> [DimIndex SubExp]
forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slice) [TPrimExp Int64 VName]
is) SymbolTable rep
table
  where
    adjust :: [DimIndex SubExp]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
adjust (DimFix SubExp
j : [DimIndex SubExp]
js') [TPrimExp Int64 VName]
is' =
      SubExp -> TPrimExp Int64 VName
pe64 SubExp
j TPrimExp Int64 VName
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. a -> [a] -> [a]
: [DimIndex SubExp]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
adjust [DimIndex SubExp]
js' [TPrimExp Int64 VName]
is'
    adjust (DimSlice SubExp
j SubExp
_ SubExp
s : [DimIndex SubExp]
js') (TPrimExp Int64 VName
i : [TPrimExp Int64 VName]
is') =
      let i_t_s :: TPrimExp Int64 VName
i_t_s = TPrimExp Int64 VName
i TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
* SubExp -> TPrimExp Int64 VName
pe64 SubExp
s
          j_p_i_t_s :: TPrimExp Int64 VName
j_p_i_t_s = SubExp -> TPrimExp Int64 VName
pe64 SubExp
j TPrimExp Int64 VName
-> TPrimExp Int64 VName -> TPrimExp Int64 VName
forall a. Num a => a -> a -> a
+ TPrimExp Int64 VName
i_t_s
       in TPrimExp Int64 VName
j_p_i_t_s TPrimExp Int64 VName
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
forall a. a -> [a] -> [a]
: [DimIndex SubExp]
-> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
adjust [DimIndex SubExp]
js' [TPrimExp Int64 VName]
is'
    adjust [DimIndex SubExp]
_ [TPrimExp Int64 VName]
_ = []
indexExp SymbolTable rep
_ Exp rep
_ Int
_ [TPrimExp Int64 VName]
_ = Maybe Indexed
forall a. Maybe a
Nothing

defBndEntry ::
  (ASTRep rep, IndexOp (Op rep)) =>
  SymbolTable rep ->
  PatElem (LetDec rep) ->
  Names ->
  Stm rep ->
  LetBoundEntry rep
defBndEntry :: forall rep.
(ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep
-> PatElem (LetDec rep) -> Names -> Stm rep -> LetBoundEntry rep
defBndEntry SymbolTable rep
vtable PatElem (LetDec rep)
patElem Names
als Stm rep
stm =
  LetBoundEntry
    { letBoundDec :: LetDec rep
letBoundDec = PatElem (LetDec rep) -> LetDec rep
forall dec. PatElem dec -> dec
patElemDec PatElem (LetDec rep)
patElem,
      letBoundAliases :: Names
letBoundAliases = Names
als,
      letBoundStm :: Stm rep
letBoundStm = Stm rep
stm,
      letBoundIndex :: Int -> IndexArray
letBoundIndex = \Int
k ->
        (Indexed -> Indexed) -> Maybe Indexed -> Maybe Indexed
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Certs -> Indexed -> Indexed
indexedAddCerts (StmAux (ExpDec rep) -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts (StmAux (ExpDec rep) -> Certs) -> StmAux (ExpDec rep) -> Certs
forall a b. (a -> b) -> a -> b
$ Stm rep -> StmAux (ExpDec rep)
forall rep. Stm rep -> StmAux (ExpDec rep)
stmAux Stm rep
stm))
          (Maybe Indexed -> Maybe Indexed) -> IndexArray -> IndexArray
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SymbolTable rep -> Exp rep -> Int -> IndexArray
forall rep.
(IndexOp (Op rep), ASTRep rep) =>
SymbolTable rep -> Exp rep -> Int -> IndexArray
indexExp SymbolTable rep
vtable (Stm rep -> Exp rep
forall rep. Stm rep -> Exp rep
stmExp Stm rep
stm) Int
k
    }

bindingEntries ::
  (Aliases.Aliased rep, IndexOp (Op rep)) =>
  Stm rep ->
  SymbolTable rep ->
  [LetBoundEntry rep]
bindingEntries :: forall rep.
(Aliased rep, IndexOp (Op rep)) =>
Stm rep -> SymbolTable rep -> [LetBoundEntry rep]
bindingEntries stm :: Stm rep
stm@(Let Pat (LetDec rep)
pat StmAux (ExpDec rep)
_ Exp rep
_) SymbolTable rep
vtable = do
  PatElem (LetDec rep)
pat_elem <- Pat (LetDec rep) -> [PatElem (LetDec rep)]
forall dec. Pat dec -> [PatElem dec]
patElems Pat (LetDec rep)
pat
  LetBoundEntry rep -> [LetBoundEntry rep]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LetBoundEntry rep -> [LetBoundEntry rep])
-> LetBoundEntry rep -> [LetBoundEntry rep]
forall a b. (a -> b) -> a -> b
$ SymbolTable rep
-> PatElem (LetDec rep) -> Names -> Stm rep -> LetBoundEntry rep
forall rep.
(ASTRep rep, IndexOp (Op rep)) =>
SymbolTable rep
-> PatElem (LetDec rep) -> Names -> Stm rep -> LetBoundEntry rep
defBndEntry SymbolTable rep
vtable PatElem (LetDec rep)
pat_elem (Names -> SymbolTable rep -> Names
forall rep. Names -> SymbolTable rep -> Names
expandAliases (PatElem (LetDec rep) -> Names
forall a. AliasesOf a => a -> Names
Aliases.aliasesOf PatElem (LetDec rep)
pat_elem) SymbolTable rep
vtable) Stm rep
stm

adjustSeveral :: (Ord k) => (v -> v) -> [k] -> M.Map k v -> M.Map k v
adjustSeveral :: forall k v. Ord k => (v -> v) -> [k] -> Map k v -> Map k v
adjustSeveral v -> v
f = (Map k v -> [k] -> Map k v) -> [k] -> Map k v -> Map k v
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Map k v -> [k] -> Map k v) -> [k] -> Map k v -> Map k v)
-> (Map k v -> [k] -> Map k v) -> [k] -> Map k v -> Map k v
forall a b. (a -> b) -> a -> b
$ (Map k v -> k -> Map k v) -> Map k v -> [k] -> Map k v
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((Map k v -> k -> Map k v) -> Map k v -> [k] -> Map k v)
-> (Map k v -> k -> Map k v) -> Map k v -> [k] -> Map k v
forall a b. (a -> b) -> a -> b
$ (k -> Map k v -> Map k v) -> Map k v -> k -> Map k v
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((k -> Map k v -> Map k v) -> Map k v -> k -> Map k v)
-> (k -> Map k v -> Map k v) -> Map k v -> k -> Map k v
forall a b. (a -> b) -> a -> b
$ (v -> v) -> k -> Map k v -> Map k v
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust v -> v
f

insertEntry ::
  (ASTRep rep) =>
  VName ->
  EntryType rep ->
  SymbolTable rep ->
  SymbolTable rep
insertEntry :: forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry VName
name EntryType rep
entry SymbolTable rep
vtable =
  let entry' :: Entry rep
entry' =
        Entry
          { entryConsumed :: Bool
entryConsumed = Bool
False,
            entryDepth :: Int
entryDepth = SymbolTable rep -> Int
forall rep. SymbolTable rep -> Int
loopDepth SymbolTable rep
vtable,
            entryIsSize :: Bool
entryIsSize = Bool
False,
            entryAccInput :: Maybe (WithAccInput rep)
entryAccInput = Maybe (WithAccInput rep)
forall a. Maybe a
Nothing,
            entryType :: EntryType rep
entryType = EntryType rep
entry
          }
      dims :: [VName]
dims = (SubExp -> Maybe VName) -> [SubExp] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe SubExp -> Maybe VName
subExpVar ([SubExp] -> [VName]) -> [SubExp] -> [VName]
forall a b. (a -> b) -> a -> b
$ Type -> [SubExp]
forall u. TypeBase Shape u -> [SubExp]
arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Entry rep -> Type
forall t. Typed t => t -> Type
typeOf Entry rep
entry'
      isSize :: Entry rep -> Entry rep
isSize Entry rep
e = Entry rep
e {entryIsSize :: Bool
entryIsSize = Bool
True}
   in SymbolTable rep
vtable
        { bindings :: Map VName (Entry rep)
bindings =
            (Entry rep -> Entry rep)
-> [VName] -> Map VName (Entry rep) -> Map VName (Entry rep)
forall k v. Ord k => (v -> v) -> [k] -> Map k v -> Map k v
adjustSeveral Entry rep -> Entry rep
forall {rep}. Entry rep -> Entry rep
isSize [VName]
dims (Map VName (Entry rep) -> Map VName (Entry rep))
-> Map VName (Entry rep) -> Map VName (Entry rep)
forall a b. (a -> b) -> a -> b
$
              VName
-> Entry rep -> Map VName (Entry rep) -> Map VName (Entry rep)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
name Entry rep
entry' (Map VName (Entry rep) -> Map VName (Entry rep))
-> Map VName (Entry rep) -> Map VName (Entry rep)
forall a b. (a -> b) -> a -> b
$
                SymbolTable rep -> Map VName (Entry rep)
forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable
        }

insertEntries ::
  (ASTRep rep) =>
  [(VName, EntryType rep)] ->
  SymbolTable rep ->
  SymbolTable rep
insertEntries :: forall rep.
ASTRep rep =>
[(VName, EntryType rep)] -> SymbolTable rep -> SymbolTable rep
insertEntries [(VName, EntryType rep)]
entries SymbolTable rep
vtable =
  (SymbolTable rep -> (VName, EntryType rep) -> SymbolTable rep)
-> SymbolTable rep -> [(VName, EntryType rep)] -> SymbolTable rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' SymbolTable rep -> (VName, EntryType rep) -> SymbolTable rep
forall {rep}.
ASTRep rep =>
SymbolTable rep -> (VName, EntryType rep) -> SymbolTable rep
add SymbolTable rep
vtable [(VName, EntryType rep)]
entries
  where
    add :: SymbolTable rep -> (VName, EntryType rep) -> SymbolTable rep
add SymbolTable rep
vtable' (VName
name, EntryType rep
entry) = VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry VName
name EntryType rep
entry SymbolTable rep
vtable'

insertStm ::
  (IndexOp (Op rep), Aliases.Aliased rep) =>
  Stm rep ->
  SymbolTable rep ->
  SymbolTable rep
insertStm :: forall rep.
(IndexOp (Op rep), Aliased rep) =>
Stm rep -> SymbolTable rep -> SymbolTable rep
insertStm Stm rep
stm SymbolTable rep
vtable =
  (SymbolTable rep -> [VName] -> SymbolTable rep)
-> [VName] -> SymbolTable rep -> SymbolTable rep
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((SymbolTable rep -> VName -> SymbolTable rep)
-> SymbolTable rep -> [VName] -> SymbolTable rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((SymbolTable rep -> VName -> SymbolTable rep)
 -> SymbolTable rep -> [VName] -> SymbolTable rep)
-> (SymbolTable rep -> VName -> SymbolTable rep)
-> SymbolTable rep
-> [VName]
-> SymbolTable rep
forall a b. (a -> b) -> a -> b
$ (VName -> SymbolTable rep -> SymbolTable rep)
-> SymbolTable rep -> VName -> SymbolTable rep
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName -> SymbolTable rep -> SymbolTable rep
forall rep. VName -> SymbolTable rep -> SymbolTable rep
consume) (Names -> [VName]
namesToList Names
stm_consumed) (SymbolTable rep -> SymbolTable rep)
-> SymbolTable rep -> SymbolTable rep
forall a b. (a -> b) -> a -> b
$
    (SymbolTable rep
 -> [(VName, LetBoundEntry rep)] -> SymbolTable rep)
-> [(VName, LetBoundEntry rep)]
-> SymbolTable rep
-> SymbolTable rep
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((SymbolTable rep -> (VName, LetBoundEntry rep) -> SymbolTable rep)
-> SymbolTable rep
-> [(VName, LetBoundEntry rep)]
-> SymbolTable rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' SymbolTable rep -> (VName, LetBoundEntry rep) -> SymbolTable rep
forall {rep} {rep}.
SymbolTable rep -> (VName, LetBoundEntry rep) -> SymbolTable rep
addRevAliases) ([VName] -> [LetBoundEntry rep] -> [(VName, LetBoundEntry rep)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
names [LetBoundEntry rep]
entries) (SymbolTable rep -> SymbolTable rep)
-> SymbolTable rep -> SymbolTable rep
forall a b. (a -> b) -> a -> b
$
      [(VName, EntryType rep)] -> SymbolTable rep -> SymbolTable rep
forall rep.
ASTRep rep =>
[(VName, EntryType rep)] -> SymbolTable rep -> SymbolTable rep
insertEntries ([VName] -> [EntryType rep] -> [(VName, EntryType rep)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
names ([EntryType rep] -> [(VName, EntryType rep)])
-> [EntryType rep] -> [(VName, EntryType rep)]
forall a b. (a -> b) -> a -> b
$ (LetBoundEntry rep -> EntryType rep)
-> [LetBoundEntry rep] -> [EntryType rep]
forall a b. (a -> b) -> [a] -> [b]
map LetBoundEntry rep -> EntryType rep
forall rep. LetBoundEntry rep -> EntryType rep
LetBound [LetBoundEntry rep]
entries) SymbolTable rep
vtable
  where
    entries :: [LetBoundEntry rep]
entries = Stm rep -> SymbolTable rep -> [LetBoundEntry rep]
forall rep.
(Aliased rep, IndexOp (Op rep)) =>
Stm rep -> SymbolTable rep -> [LetBoundEntry rep]
bindingEntries Stm rep
stm SymbolTable rep
vtable
    names :: [VName]
names = Pat (LetDec rep) -> [VName]
forall dec. Pat dec -> [VName]
patNames (Pat (LetDec rep) -> [VName]) -> Pat (LetDec rep) -> [VName]
forall a b. (a -> b) -> a -> b
$ Stm rep -> Pat (LetDec rep)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm rep
stm
    stm_consumed :: Names
stm_consumed = Names -> SymbolTable rep -> Names
forall rep. Names -> SymbolTable rep -> Names
expandAliases (Stm rep -> Names
forall rep. Aliased rep => Stm rep -> Names
Aliases.consumedInStm Stm rep
stm) SymbolTable rep
vtable
    addRevAliases :: SymbolTable rep -> (VName, LetBoundEntry rep) -> SymbolTable rep
addRevAliases SymbolTable rep
vtable' (VName
name, LetBoundEntry {letBoundAliases :: forall rep. LetBoundEntry rep -> Names
letBoundAliases = Names
als}) =
      SymbolTable rep
vtable' {bindings :: Map VName (Entry rep)
bindings = (Entry rep -> Entry rep)
-> [VName] -> Map VName (Entry rep) -> Map VName (Entry rep)
forall k v. Ord k => (v -> v) -> [k] -> Map k v -> Map k v
adjustSeveral Entry rep -> Entry rep
update [VName]
inedges (Map VName (Entry rep) -> Map VName (Entry rep))
-> Map VName (Entry rep) -> Map VName (Entry rep)
forall a b. (a -> b) -> a -> b
$ SymbolTable rep -> Map VName (Entry rep)
forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable'}
      where
        inedges :: [VName]
inedges = Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ Names -> SymbolTable rep -> Names
forall rep. Names -> SymbolTable rep -> Names
expandAliases Names
als SymbolTable rep
vtable'
        update :: Entry rep -> Entry rep
update Entry rep
e = Entry rep
e {entryType :: EntryType rep
entryType = EntryType rep -> EntryType rep
update' (EntryType rep -> EntryType rep) -> EntryType rep -> EntryType rep
forall a b. (a -> b) -> a -> b
$ Entry rep -> EntryType rep
forall rep. Entry rep -> EntryType rep
entryType Entry rep
e}
        update' :: EntryType rep -> EntryType rep
update' (LetBound LetBoundEntry rep
entry) =
          LetBoundEntry rep -> EntryType rep
forall rep. LetBoundEntry rep -> EntryType rep
LetBound
            LetBoundEntry rep
entry
              { letBoundAliases :: Names
letBoundAliases = VName -> Names
oneName VName
name Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> LetBoundEntry rep -> Names
forall rep. LetBoundEntry rep -> Names
letBoundAliases LetBoundEntry rep
entry
              }
        update' (FParam FParamEntry rep
entry) =
          FParamEntry rep -> EntryType rep
forall rep. FParamEntry rep -> EntryType rep
FParam
            FParamEntry rep
entry
              { fparamAliases :: Names
fparamAliases = VName -> Names
oneName VName
name Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> FParamEntry rep -> Names
forall rep. FParamEntry rep -> Names
fparamAliases FParamEntry rep
entry
              }
        update' (LParam LParamEntry rep
entry) =
          LParamEntry rep -> EntryType rep
forall rep. LParamEntry rep -> EntryType rep
LParam
            LParamEntry rep
entry
              { lparamAliases :: Names
lparamAliases = VName -> Names
oneName VName
name Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> LParamEntry rep -> Names
forall rep. LParamEntry rep -> Names
lparamAliases LParamEntry rep
entry
              }
        update' (FreeVar FreeVarEntry rep
entry) =
          FreeVarEntry rep -> EntryType rep
forall rep. FreeVarEntry rep -> EntryType rep
FreeVar
            FreeVarEntry rep
entry
              { freeVarAliases :: Names
freeVarAliases = VName -> Names
oneName VName
name Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> FreeVarEntry rep -> Names
forall rep. FreeVarEntry rep -> Names
freeVarAliases FreeVarEntry rep
entry
              }
        update' EntryType rep
e = EntryType rep
e

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

expandAliases :: Names -> SymbolTable rep -> Names
expandAliases :: forall rep. Names -> SymbolTable rep -> Names
expandAliases Names
names SymbolTable rep
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 rep -> Names
forall rep. VName -> SymbolTable rep -> Names
`lookupAliases` SymbolTable rep
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 ::
  (ASTRep rep) =>
  AST.FParam rep ->
  SymbolTable rep ->
  SymbolTable rep
insertFParam :: forall rep.
ASTRep rep =>
FParam rep -> SymbolTable rep -> SymbolTable rep
insertFParam FParam rep
fparam = VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry VName
name EntryType rep
entry
  where
    name :: VName
name = FParam rep -> VName
forall dec. Param dec -> VName
AST.paramName FParam rep
fparam
    entry :: EntryType rep
entry =
      FParamEntry rep -> EntryType rep
forall rep. FParamEntry rep -> EntryType rep
FParam
        FParamEntry
          { fparamDec :: FParamInfo rep
fparamDec = FParam rep -> FParamInfo rep
forall dec. Param dec -> dec
AST.paramDec FParam rep
fparam,
            fparamAliases :: Names
fparamAliases = Names
forall a. Monoid a => a
mempty,
            fparamMerge :: Maybe (SubExp, SubExp)
fparamMerge = Maybe (SubExp, SubExp)
forall a. Maybe a
Nothing
          }

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

insertLParam :: (ASTRep rep) => LParam rep -> SymbolTable rep -> SymbolTable rep
insertLParam :: forall rep.
ASTRep rep =>
LParam rep -> SymbolTable rep -> SymbolTable rep
insertLParam LParam rep
param = VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry VName
name EntryType rep
bind
  where
    bind :: EntryType rep
bind =
      LParamEntry rep -> EntryType rep
forall rep. LParamEntry rep -> EntryType rep
LParam
        LParamEntry
          { lparamDec :: LParamInfo rep
lparamDec = LParam rep -> LParamInfo rep
forall dec. Param dec -> dec
AST.paramDec LParam rep
param,
            lparamAliases :: Names
lparamAliases = Names
forall a. Monoid a => a
mempty,
            lparamIndex :: IndexArray
lparamIndex = Maybe Indexed -> IndexArray
forall a b. a -> b -> a
const Maybe Indexed
forall a. Maybe a
Nothing
          }
    name :: VName
name = LParam rep -> VName
forall dec. Param dec -> VName
AST.paramName LParam rep
param

-- | Insert entries corresponding to the parameters of a loop (not
-- distinguishing contect and value part).  Apart from the parameter
-- itself, we also insert the initial value and the subexpression
-- providing the final value.  Note that the latter is likely not in
-- scope in the symbol at this point.  This is OK, and can still be
-- used to help some loop optimisations detect invariant loop
-- parameters.
insertLoopMerge ::
  (ASTRep rep) =>
  [(AST.FParam rep, SubExp, SubExpRes)] ->
  SymbolTable rep ->
  SymbolTable rep
insertLoopMerge :: forall rep.
ASTRep rep =>
[(FParam rep, SubExp, SubExpRes)]
-> SymbolTable rep -> SymbolTable rep
insertLoopMerge = (SymbolTable rep
 -> [(FParam rep, SubExp, SubExpRes)] -> SymbolTable rep)
-> [(FParam rep, SubExp, SubExpRes)]
-> SymbolTable rep
-> SymbolTable rep
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((SymbolTable rep
  -> [(FParam rep, SubExp, SubExpRes)] -> SymbolTable rep)
 -> [(FParam rep, SubExp, SubExpRes)]
 -> SymbolTable rep
 -> SymbolTable rep)
-> (SymbolTable rep
    -> [(FParam rep, SubExp, SubExpRes)] -> SymbolTable rep)
-> [(FParam rep, SubExp, SubExpRes)]
-> SymbolTable rep
-> SymbolTable rep
forall a b. (a -> b) -> a -> b
$ (SymbolTable rep
 -> (FParam rep, SubExp, SubExpRes) -> SymbolTable rep)
-> SymbolTable rep
-> [(FParam rep, SubExp, SubExpRes)]
-> SymbolTable rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((SymbolTable rep
  -> (FParam rep, SubExp, SubExpRes) -> SymbolTable rep)
 -> SymbolTable rep
 -> [(FParam rep, SubExp, SubExpRes)]
 -> SymbolTable rep)
-> (SymbolTable rep
    -> (FParam rep, SubExp, SubExpRes) -> SymbolTable rep)
-> SymbolTable rep
-> [(FParam rep, SubExp, SubExpRes)]
-> SymbolTable rep
forall a b. (a -> b) -> a -> b
$ ((FParam rep, SubExp, SubExpRes)
 -> SymbolTable rep -> SymbolTable rep)
-> SymbolTable rep
-> (FParam rep, SubExp, SubExpRes)
-> SymbolTable rep
forall a b c. (a -> b -> c) -> b -> a -> c
flip (FParam rep, SubExp, SubExpRes)
-> SymbolTable rep -> SymbolTable rep
forall {rep}.
ASTRep rep =>
(Param (FParamInfo rep), SubExp, SubExpRes)
-> SymbolTable rep -> SymbolTable rep
bind
  where
    bind :: (Param (FParamInfo rep), SubExp, SubExpRes)
-> SymbolTable rep -> SymbolTable rep
bind (Param (FParamInfo rep)
p, SubExp
initial, SubExpRes Certs
_ SubExp
res) =
      VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry (Param (FParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName Param (FParamInfo rep)
p) (EntryType rep -> SymbolTable rep -> SymbolTable rep)
-> EntryType rep -> SymbolTable rep -> SymbolTable rep
forall a b. (a -> b) -> a -> b
$
        FParamEntry rep -> EntryType rep
forall rep. FParamEntry rep -> EntryType rep
FParam
          FParamEntry
            { fparamDec :: FParamInfo rep
fparamDec = Param (FParamInfo rep) -> FParamInfo rep
forall dec. Param dec -> dec
AST.paramDec Param (FParamInfo rep)
p,
              fparamAliases :: Names
fparamAliases = Names
forall a. Monoid a => a
mempty,
              fparamMerge :: Maybe (SubExp, SubExp)
fparamMerge = (SubExp, SubExp) -> Maybe (SubExp, SubExp)
forall a. a -> Maybe a
Just (SubExp
initial, SubExp
res)
            }

insertLoopVar :: (ASTRep rep) => VName -> IntType -> SubExp -> SymbolTable rep -> SymbolTable rep
insertLoopVar :: forall rep.
ASTRep rep =>
VName -> IntType -> SubExp -> SymbolTable rep -> SymbolTable rep
insertLoopVar VName
name IntType
it SubExp
bound = VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry VName
name EntryType rep
bind
  where
    bind :: EntryType rep
bind =
      LoopVarEntry rep -> EntryType rep
forall rep. LoopVarEntry rep -> EntryType rep
LoopVar
        LoopVarEntry
          { loopVarType :: IntType
loopVarType = IntType
it,
            loopVarBound :: SubExp
loopVarBound = SubExp
bound
          }

insertFreeVar :: (ASTRep rep) => VName -> NameInfo rep -> SymbolTable rep -> SymbolTable rep
insertFreeVar :: forall rep.
ASTRep rep =>
VName -> NameInfo rep -> SymbolTable rep -> SymbolTable rep
insertFreeVar VName
name NameInfo rep
dec = VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
forall rep.
ASTRep rep =>
VName -> EntryType rep -> SymbolTable rep -> SymbolTable rep
insertEntry VName
name EntryType rep
entry
  where
    entry :: EntryType rep
entry =
      FreeVarEntry rep -> EntryType rep
forall rep. FreeVarEntry rep -> EntryType rep
FreeVar
        FreeVarEntry
          { freeVarDec :: NameInfo rep
freeVarDec = NameInfo rep
dec,
            freeVarIndex :: VName -> IndexArray
freeVarIndex = \VName
_ [TPrimExp Int64 VName]
_ -> Maybe Indexed
forall a. Maybe a
Nothing,
            freeVarAliases :: Names
freeVarAliases = Names
forall a. Monoid a => a
mempty
          }

consume :: VName -> SymbolTable rep -> SymbolTable rep
consume :: forall rep. VName -> SymbolTable rep -> SymbolTable rep
consume VName
consumee SymbolTable rep
vtable =
  (SymbolTable rep -> VName -> SymbolTable rep)
-> SymbolTable rep -> [VName] -> SymbolTable rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' SymbolTable rep -> VName -> SymbolTable rep
forall {rep}. SymbolTable rep -> VName -> SymbolTable rep
consume' SymbolTable rep
vtable ([VName] -> SymbolTable rep) -> [VName] -> SymbolTable rep
forall a b. (a -> b) -> a -> b
$
    Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$
      Names -> SymbolTable rep -> Names
forall rep. Names -> SymbolTable rep -> Names
expandAliases (VName -> Names
oneName VName
consumee) SymbolTable rep
vtable
  where
    consume' :: SymbolTable rep -> VName -> SymbolTable rep
consume' SymbolTable rep
vtable' VName
v =
      SymbolTable rep
vtable' {bindings :: Map VName (Entry rep)
bindings = (Entry rep -> Entry rep)
-> VName -> Map VName (Entry rep) -> Map VName (Entry rep)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust Entry rep -> Entry rep
forall {rep}. Entry rep -> Entry rep
consume'' VName
v (Map VName (Entry rep) -> Map VName (Entry rep))
-> Map VName (Entry rep) -> Map VName (Entry rep)
forall a b. (a -> b) -> a -> b
$ SymbolTable rep -> Map VName (Entry rep)
forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable'}
    consume'' :: Entry rep -> Entry rep
consume'' Entry rep
e = Entry rep
e {entryConsumed :: Bool
entryConsumed = Bool
True}

-- | Hide definitions of those entries that satisfy some predicate.
hideIf :: (Entry rep -> Bool) -> SymbolTable rep -> SymbolTable rep
hideIf :: forall rep.
(Entry rep -> Bool) -> SymbolTable rep -> SymbolTable rep
hideIf Entry rep -> Bool
hide SymbolTable rep
vtable = SymbolTable rep
vtable {bindings :: Map VName (Entry rep)
bindings = (Entry rep -> Entry rep)
-> Map VName (Entry rep) -> Map VName (Entry rep)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map Entry rep -> Entry rep
maybeHide (Map VName (Entry rep) -> Map VName (Entry rep))
-> Map VName (Entry rep) -> Map VName (Entry rep)
forall a b. (a -> b) -> a -> b
$ SymbolTable rep -> Map VName (Entry rep)
forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable}
  where
    maybeHide :: Entry rep -> Entry rep
maybeHide Entry rep
entry
      | Entry rep -> Bool
hide Entry rep
entry =
          Entry rep
entry
            { entryType :: EntryType rep
entryType =
                FreeVarEntry rep -> EntryType rep
forall rep. FreeVarEntry rep -> EntryType rep
FreeVar
                  FreeVarEntry
                    { freeVarDec :: NameInfo rep
freeVarDec = Entry rep -> NameInfo rep
forall rep. Entry rep -> NameInfo rep
entryInfo Entry rep
entry,
                      freeVarIndex :: VName -> IndexArray
freeVarIndex = \VName
_ [TPrimExp Int64 VName]
_ -> Maybe Indexed
forall a. Maybe a
Nothing,
                      freeVarAliases :: Names
freeVarAliases = EntryType rep -> Names
forall rep. EntryType rep -> Names
entryAliases (EntryType rep -> Names) -> EntryType rep -> Names
forall a b. (a -> b) -> a -> b
$ Entry rep -> EntryType rep
forall rep. Entry rep -> EntryType rep
entryType Entry rep
entry
                    }
            }
      | Bool
otherwise = Entry rep
entry

-- | Hide these definitions, if they are protected by certificates in
-- the set of names.
hideCertified :: Names -> SymbolTable rep -> SymbolTable rep
hideCertified :: forall rep. Names -> SymbolTable rep -> SymbolTable rep
hideCertified Names
to_hide = (Entry rep -> Bool) -> SymbolTable rep -> SymbolTable rep
forall rep.
(Entry rep -> Bool) -> SymbolTable rep -> SymbolTable rep
hideIf ((Entry rep -> Bool) -> SymbolTable rep -> SymbolTable rep)
-> (Entry rep -> Bool) -> SymbolTable rep -> SymbolTable rep
forall a b. (a -> b) -> a -> b
$ Bool -> (Stm rep -> Bool) -> Maybe (Stm rep) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False Stm rep -> Bool
hide (Maybe (Stm rep) -> Bool)
-> (Entry rep -> Maybe (Stm rep)) -> Entry rep -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry rep -> Maybe (Stm rep)
forall rep. Entry rep -> Maybe (Stm rep)
entryStm
  where
    hide :: Stm rep -> 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 rep -> [VName]) -> Stm rep -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Certs -> [VName]
unCerts (Certs -> [VName]) -> (Stm rep -> Certs) -> Stm rep -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm rep -> Certs
forall rep. Stm rep -> Certs
stmCerts

-- | Note that these names are tokens for the corresponding
-- accumulators.  The names must already be present in the symbol
-- table.
noteAccTokens ::
  [(VName, WithAccInput rep)] ->
  SymbolTable rep ->
  SymbolTable rep
noteAccTokens :: forall rep.
[(VName, WithAccInput rep)] -> SymbolTable rep -> SymbolTable rep
noteAccTokens = (SymbolTable rep -> [(VName, WithAccInput rep)] -> SymbolTable rep)
-> [(VName, WithAccInput rep)]
-> SymbolTable rep
-> SymbolTable rep
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((SymbolTable rep -> (VName, WithAccInput rep) -> SymbolTable rep)
-> SymbolTable rep
-> [(VName, WithAccInput rep)]
-> SymbolTable rep
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' SymbolTable rep -> (VName, WithAccInput rep) -> SymbolTable rep
forall {rep}.
SymbolTable rep -> (VName, WithAccInput rep) -> SymbolTable rep
f)
  where
    f :: SymbolTable rep -> (VName, WithAccInput rep) -> SymbolTable rep
f SymbolTable rep
vtable (VName
v, WithAccInput rep
accum) =
      case VName -> Map VName (Entry rep) -> Maybe (Entry rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Map VName (Entry rep) -> Maybe (Entry rep))
-> Map VName (Entry rep) -> Maybe (Entry rep)
forall a b. (a -> b) -> a -> b
$ SymbolTable rep -> Map VName (Entry rep)
forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable of
        Maybe (Entry rep)
Nothing -> SymbolTable rep
vtable
        Just Entry rep
e ->
          SymbolTable rep
vtable
            { bindings :: Map VName (Entry rep)
bindings =
                VName
-> Entry rep -> Map VName (Entry rep) -> Map VName (Entry rep)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v (Entry rep
e {entryAccInput :: Maybe (WithAccInput rep)
entryAccInput = WithAccInput rep -> Maybe (WithAccInput rep)
forall a. a -> Maybe a
Just WithAccInput rep
accum}) (Map VName (Entry rep) -> Map VName (Entry rep))
-> Map VName (Entry rep) -> Map VName (Entry rep)
forall a b. (a -> b) -> a -> b
$ SymbolTable rep -> Map VName (Entry rep)
forall rep. SymbolTable rep -> Map VName (Entry rep)
bindings SymbolTable rep
vtable
            }