{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}

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

    -- * Entries
    Entry,
    deepen,
    entryDepth,
    entryLetBoundDec,
    entryIsSize,
    entryStm,

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

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

    -- * Misc
    hideCertified,
  )
where

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

data SymbolTable lore = SymbolTable
  { forall lore. SymbolTable lore -> Int
loopDepth :: Int,
    forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings :: M.Map VName (Entry lore),
    -- | Which names are available just before the most enclosing
    -- loop?
    forall lore. SymbolTable lore -> 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 lore. SymbolTable lore -> Bool
simplifyMemory :: Bool
  }

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 :: forall lore. 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 :: forall lore. ASTLore lore => Scope lore -> SymbolTable lore
fromScope = (SymbolTable lore -> VName -> NameInfo lore -> SymbolTable lore)
-> SymbolTable lore
-> Map VName (NameInfo 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 :: forall lore. SymbolTable lore -> Scope lore
toScope = (Entry lore -> NameInfo lore)
-> Map VName (Entry lore) -> Map VName (NameInfo 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) -> Map VName (NameInfo lore))
-> (SymbolTable lore -> Map VName (Entry lore))
-> SymbolTable lore
-> Map VName (NameInfo 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 :: forall lore. 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
  = -- | A PrimExp based on the indexes (that is, without
    -- accessing any actual array).
    Indexed Certificates (PrimExp VName)
  | -- | The indexing corresponds to another (perhaps more
    -- advantageous) array.
    IndexedArray Certificates VName [TPrimExp Int64 VName]

indexedAddCerts :: Certificates -> Indexed -> Indexed
indexedAddCerts :: Certificates -> Indexed -> Indexed
indexedAddCerts Certificates
cs1 (Indexed Certificates
cs2 PrimExp VName
v) = Certificates -> PrimExp VName -> Indexed
Indexed (Certificates
cs1 Certificates -> Certificates -> Certificates
forall a. Semigroup a => a -> a -> a
<> Certificates
cs2) PrimExp VName
v
indexedAddCerts Certificates
cs1 (IndexedArray Certificates
cs2 VName
arr [TPrimExp Int64 VName]
v) = Certificates -> VName -> [TPrimExp Int64 VName] -> Indexed
IndexedArray (Certificates
cs1 Certificates -> Certificates -> Certificates
forall a. Semigroup a => a -> a -> a
<> Certificates
cs2) VName
arr [TPrimExp Int64 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 [TPrimExp Int64 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
<> [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 lore = Entry
  { -- | True if consumed.
    forall lore. Entry lore -> Bool
entryConsumed :: Bool,
    forall lore. Entry lore -> Int
entryDepth :: Int,
    -- | True if this name has been used as an array size,
    -- implying that it is non-negative.
    forall lore. Entry lore -> Bool
entryIsSize :: Bool,
    forall lore. Entry lore -> EntryType lore
entryType :: EntryType lore
  }

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

data LoopVarEntry lore = LoopVarEntry
  { forall lore. LoopVarEntry lore -> IntType
loopVarType :: IntType,
    forall lore. LoopVarEntry lore -> SubExp
loopVarBound :: SubExp
  }

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

data FParamEntry lore = FParamEntry
  { forall lore. FParamEntry lore -> FParamInfo lore
fparamDec :: FParamInfo lore,
    forall lore. FParamEntry lore -> 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 lore. FParamEntry lore -> Maybe (SubExp, SubExp)
fparamMerge :: Maybe (SubExp, SubExp)
  }

data LParamEntry lore = LParamEntry
  { forall lore. LParamEntry lore -> LParamInfo lore
lparamDec :: LParamInfo lore,
    forall lore. LParamEntry lore -> IndexArray
lparamIndex :: IndexArray
  }

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

instance ASTLore lore => Typed (Entry lore) where
  typeOf :: Entry lore -> Type
typeOf = 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

entryInfo :: Entry lore -> NameInfo lore
entryInfo :: forall lore. Entry lore -> NameInfo lore
entryInfo Entry lore
e = case Entry lore -> EntryType lore
forall lore. Entry lore -> EntryType lore
entryType Entry lore
e of
  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
  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
  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
  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
  FreeVar FreeVarEntry lore
entry -> FreeVarEntry lore -> NameInfo lore
forall lore. FreeVarEntry lore -> NameInfo lore
freeVarDec FreeVarEntry lore
entry

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

entryStm :: Entry lore -> Maybe (Stm lore)
entryStm :: forall lore. Entry lore -> Maybe (Stm lore)
entryStm = (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)
isLetBound

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

elem :: VName -> SymbolTable lore -> Bool
elem :: forall lore. 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 :: forall lore. 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 :: forall lore. VName -> SymbolTable lore -> Maybe (Stm lore)
lookupStm VName
name SymbolTable lore
vtable = Entry lore -> Maybe (Stm lore)
forall lore. Entry lore -> Maybe (Stm lore)
entryStm (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 :: forall lore.
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 :: forall lore.
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 :: forall lore.
ASTLore lore =>
VName -> SymbolTable lore -> Maybe Type
lookupType VName
name SymbolTable lore
vtable = Entry lore -> Type
forall t. Typed t => t -> Type
typeOf (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 :: forall lore.
ASTLore lore =>
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 :: forall lore.
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

lookupAliases :: VName -> SymbolTable lore -> Names
lookupAliases :: forall lore. VName -> SymbolTable lore -> Names
lookupAliases VName
name SymbolTable lore
vtable =
  case Entry lore -> EntryType lore
forall lore. Entry lore -> EntryType lore
entryType (Entry lore -> EntryType lore)
-> Maybe (Entry lore) -> Maybe (EntryType lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Map VName (Entry lore) -> Maybe (Entry lore)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (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 (EntryType lore)
_ -> Names
forall a. Monoid a => a
mempty

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

lookupLoopParam :: VName -> SymbolTable lore -> Maybe (SubExp, SubExp)
lookupLoopParam :: forall lore. VName -> SymbolTable lore -> Maybe (SubExp, SubExp)
lookupLoopParam VName
name SymbolTable lore
vtable = do
  FParam FParamEntry lore
e <- Entry lore -> EntryType lore
forall lore. Entry lore -> EntryType lore
entryType (Entry lore -> EntryType lore)
-> Maybe (Entry lore) -> Maybe (EntryType lore)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Map VName (Entry lore) -> Maybe (Entry lore)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (SymbolTable lore -> Map VName (Entry lore)
forall lore. SymbolTable lore -> Map VName (Entry lore)
bindings SymbolTable lore
vtable)
  FParamEntry lore -> Maybe (SubExp, SubExp)
forall lore. FParamEntry lore -> Maybe (SubExp, SubExp)
fparamMerge FParamEntry lore
e

-- | In symbol table and not consumed.
available :: VName -> SymbolTable lore -> Bool
available :: forall lore. 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
entryConsumed) (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 :: forall lore.
ASTLore lore =>
VName -> [SubExp] -> SymbolTable lore -> Maybe Indexed
index VName
name [SubExp]
is SymbolTable lore
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)
mapM SubExp -> Maybe (TPrimExp Int64 VName)
asPrimExp [SubExp]
is
  VName
-> [TPrimExp Int64 VName] -> SymbolTable lore -> Maybe Indexed
forall lore.
VName
-> [TPrimExp Int64 VName] -> SymbolTable lore -> Maybe Indexed
index' VName
name [TPrimExp Int64 VName]
is' SymbolTable lore
table
  where
    asPrimExp :: SubExp -> Maybe (TPrimExp Int64 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
      TPrimExp Int64 VName -> Maybe (TPrimExp Int64 VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (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 t 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 lore ->
  Maybe Indexed
index' :: forall lore.
VName
-> [TPrimExp Int64 VName] -> SymbolTable lore -> Maybe Indexed
index' VName
name [TPrimExp Int64 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 -> EntryType lore
forall lore. Entry lore -> EntryType lore
entryType 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 [TPrimExp Int64 VName]
is
    FreeVar FreeVarEntry lore
entry' ->
      FreeVarEntry lore -> VName -> IndexArray
forall lore. FreeVarEntry lore -> VName -> IndexArray
freeVarIndex FreeVarEntry lore
entry' VName
name [TPrimExp Int64 VName]
is
    LParam LParamEntry lore
entry' -> LParamEntry lore -> IndexArray
forall lore. LParamEntry lore -> IndexArray
lparamIndex LParamEntry lore
entry' [TPrimExp Int64 VName]
is
    EntryType lore
_ -> Maybe Indexed
forall a. Maybe a
Nothing

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

instance IndexOp ()

indexExp ::
  (IndexOp (Op lore), ASTLore lore) =>
  SymbolTable lore ->
  Exp lore ->
  Int ->
  IndexArray
indexExp :: forall lore.
(IndexOp (Op lore), ASTLore lore) =>
SymbolTable lore -> Exp lore -> Int -> IndexArray
indexExp SymbolTable lore
vtable (Op Op lore
op) Int
k [TPrimExp Int64 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 [TPrimExp Int64 VName]
is
indexExp SymbolTable lore
_ (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
$
    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
$
      ( IntType -> PrimExp VName -> PrimExp VName
forall v. IntType -> PrimExp v -> PrimExp v
sExt IntType
to_it (TPrimExp Int64 VName -> PrimExp VName
forall t 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 lore
table (BasicOp (Replicate (Shape [SubExp]
ds) SubExp
v)) Int
_ [TPrimExp Int64 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
== [TPrimExp Int64 VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TPrimExp Int64 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
_ (TPrimExp Int64 VName
_ : [TPrimExp Int64 VName]
is) =
  VName
-> [TPrimExp Int64 VName] -> SymbolTable lore -> Maybe Indexed
forall lore.
VName
-> [TPrimExp Int64 VName] -> SymbolTable lore -> Maybe Indexed
index' VName
v [TPrimExp Int64 VName]
is SymbolTable lore
table
indexExp SymbolTable lore
table (BasicOp (Reshape ShapeChange SubExp
newshape VName
v)) Int
_ [TPrimExp Int64 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' :: [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
$ ShapeChange SubExp -> [SubExp]
forall d. ShapeChange d -> [d]
newDims ShapeChange SubExp
newshape)
            [TPrimExp Int64 VName]
is
     in VName
-> [TPrimExp Int64 VName] -> SymbolTable lore -> Maybe Indexed
forall lore.
VName
-> [TPrimExp Int64 VName] -> SymbolTable lore -> Maybe Indexed
index' VName
v [TPrimExp Int64 VName]
is' SymbolTable lore
table
indexExp SymbolTable lore
table (BasicOp (Index VName
v Slice SubExp
slice)) Int
_ [TPrimExp Int64 VName]
is =
  VName
-> [TPrimExp Int64 VName] -> SymbolTable lore -> Maybe Indexed
forall lore.
VName
-> [TPrimExp Int64 VName] -> SymbolTable lore -> Maybe Indexed
index' VName
v (Slice SubExp -> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
adjust Slice SubExp
slice [TPrimExp Int64 VName]
is) SymbolTable lore
table
  where
    adjust :: Slice SubExp -> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
adjust (DimFix SubExp
j : Slice 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]
: Slice SubExp -> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
adjust Slice SubExp
js' [TPrimExp Int64 VName]
is'
    adjust (DimSlice SubExp
j SubExp
_ SubExp
s : Slice 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]
: Slice SubExp -> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
adjust Slice SubExp
js' [TPrimExp Int64 VName]
is'
    adjust Slice SubExp
_ [TPrimExp Int64 VName]
_ = []
indexExp SymbolTable lore
_ ExpT lore
_ Int
_ [TPrimExp Int64 VName]
_ = Maybe Indexed
forall a. Maybe a
Nothing

defBndEntry ::
  (ASTLore lore, IndexOp (Op lore)) =>
  SymbolTable lore ->
  PatElem lore ->
  Names ->
  Stm lore ->
  LetBoundEntry lore
defBndEntry :: forall lore.
(ASTLore lore, IndexOp (Op lore)) =>
SymbolTable lore
-> PatElem lore -> Names -> Stm lore -> LetBoundEntry lore
defBndEntry SymbolTable lore
vtable PatElem lore
patElem Names
als Stm lore
bnd =
  LetBoundEntry :: forall lore.
LetDec lore
-> Names -> Stm lore -> (Int -> IndexArray) -> LetBoundEntry lore
LetBoundEntry
    { 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,
      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
    }

bindingEntries ::
  (ASTLore lore, Aliases.Aliased lore, IndexOp (Op lore)) =>
  Stm lore ->
  SymbolTable lore ->
  [LetBoundEntry lore]
bindingEntries :: forall lore.
(ASTLore lore, Aliased lore, IndexOp (Op lore)) =>
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)
-> Names
-> Stm lore
-> LetBoundEntry lore
forall lore.
(ASTLore lore, IndexOp (Op lore)) =>
SymbolTable lore
-> PatElem lore -> Names -> Stm lore -> LetBoundEntry lore
defBndEntry SymbolTable lore
vtable 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

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 (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 ::
  ASTLore lore =>
  VName ->
  EntryType lore ->
  SymbolTable lore ->
  SymbolTable lore
insertEntry :: forall lore.
ASTLore lore =>
VName -> EntryType lore -> SymbolTable lore -> SymbolTable lore
insertEntry VName
name EntryType lore
entry SymbolTable lore
vtable =
  let entry' :: Entry lore
entry' =
        Entry :: forall lore. Bool -> Int -> Bool -> EntryType lore -> Entry lore
Entry
          { entryConsumed :: Bool
entryConsumed = Bool
False,
            entryDepth :: Int
entryDepth = SymbolTable lore -> Int
forall lore. SymbolTable lore -> Int
loopDepth SymbolTable lore
vtable,
            entryIsSize :: Bool
entryIsSize = Bool
False,
            entryType :: EntryType lore
entryType = EntryType lore
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 (ShapeBase SubExp) u -> [SubExp]
arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Entry lore -> Type
forall t. Typed t => t -> Type
typeOf Entry lore
entry'
      isSize :: Entry lore -> Entry lore
isSize Entry lore
e = Entry lore
e {entryIsSize :: Bool
entryIsSize = Bool
True}
   in SymbolTable lore
vtable
        { bindings :: Map VName (Entry lore)
bindings =
            (Entry lore -> Entry lore)
-> [VName] -> Map VName (Entry lore) -> Map VName (Entry lore)
forall k v. Ord k => (v -> v) -> [k] -> Map k v -> Map k v
adjustSeveral Entry lore -> Entry lore
forall {lore}. Entry lore -> Entry lore
isSize [VName]
dims (Map VName (Entry lore) -> Map VName (Entry lore))
-> Map VName (Entry lore) -> Map VName (Entry lore)
forall a b. (a -> b) -> a -> b
$
              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) -> 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
        }

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

insertStm ::
  (ASTLore lore, IndexOp (Op lore), Aliases.Aliased lore) =>
  Stm lore ->
  SymbolTable lore ->
  SymbolTable lore
insertStm :: forall lore.
(ASTLore lore, IndexOp (Op lore), Aliased lore) =>
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. 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, EntryType lore)] -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
[(VName, EntryType lore)] -> SymbolTable lore -> SymbolTable lore
insertEntries ([VName] -> [EntryType lore] -> [(VName, EntryType lore)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
names ([EntryType lore] -> [(VName, EntryType lore)])
-> [EntryType lore] -> [(VName, EntryType lore)]
forall a b. (a -> b) -> a -> b
$ (LetBoundEntry lore -> EntryType lore)
-> [LetBoundEntry lore] -> [EntryType lore]
forall a b. (a -> b) -> [a] -> [b]
map LetBoundEntry lore -> EntryType lore
forall lore. LetBoundEntry lore -> EntryType lore
LetBound ([LetBoundEntry lore] -> [EntryType lore])
-> [LetBoundEntry lore] -> [EntryType lore]
forall a b. (a -> b) -> a -> b
$ Stm lore -> SymbolTable lore -> [LetBoundEntry lore]
forall lore.
(ASTLore 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
    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 k v. Ord k => (v -> v) -> [k] -> Map k v -> Map k v
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 Entry lore
e = Entry lore
e {entryType :: EntryType lore
entryType = EntryType lore -> EntryType lore
update' (EntryType lore -> EntryType lore)
-> EntryType lore -> EntryType lore
forall a b. (a -> b) -> a -> b
$ Entry lore -> EntryType lore
forall lore. Entry lore -> EntryType lore
entryType Entry lore
e}
        update' :: EntryType lore -> EntryType lore
update' (LetBound LetBoundEntry lore
entry) =
          LetBoundEntry lore -> EntryType lore
forall lore. LetBoundEntry lore -> EntryType 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 -> EntryType lore
forall lore. FParamEntry lore -> EntryType 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' EntryType lore
e = EntryType lore
e

insertStms ::
  (ASTLore lore, IndexOp (Op lore), Aliases.Aliased lore) =>
  Stms lore ->
  SymbolTable lore ->
  SymbolTable lore
insertStms :: forall lore.
(ASTLore lore, IndexOp (Op lore), Aliased lore) =>
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.
(ASTLore lore, IndexOp (Op 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 :: forall lore. 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 :: forall lore.
ASTLore lore =>
FParam lore -> SymbolTable lore -> SymbolTable lore
insertFParam FParam lore
fparam = VName -> EntryType lore -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
VName -> EntryType lore -> SymbolTable lore -> SymbolTable lore
insertEntry VName
name EntryType lore
entry
  where
    name :: VName
name = FParam lore -> VName
forall dec. Param dec -> VName
AST.paramName FParam lore
fparam
    entry :: EntryType lore
entry =
      FParamEntry lore -> EntryType lore
forall lore. FParamEntry lore -> EntryType lore
FParam
        FParamEntry :: forall lore.
FParamInfo lore
-> Names -> Maybe (SubExp, SubExp) -> FParamEntry lore
FParamEntry
          { 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,
            fparamMerge :: Maybe (SubExp, SubExp)
fparamMerge = Maybe (SubExp, SubExp)
forall a. Maybe a
Nothing
          }

insertFParams ::
  ASTLore lore =>
  [AST.FParam lore] ->
  SymbolTable lore ->
  SymbolTable lore
insertFParams :: forall lore.
ASTLore lore =>
[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

insertLParam :: ASTLore lore => LParam lore -> SymbolTable lore -> SymbolTable lore
insertLParam :: forall lore.
ASTLore lore =>
LParam lore -> SymbolTable lore -> SymbolTable lore
insertLParam LParam lore
param = VName -> EntryType lore -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
VName -> EntryType lore -> SymbolTable lore -> SymbolTable lore
insertEntry VName
name EntryType lore
bind
  where
    bind :: EntryType lore
bind =
      LParamEntry lore -> EntryType lore
forall lore. LParamEntry lore -> EntryType lore
LParam
        LParamEntry :: forall lore. LParamInfo lore -> IndexArray -> LParamEntry lore
LParamEntry
          { lparamDec :: LParamInfo lore
lparamDec = LParam lore -> LParamInfo lore
forall dec. Param dec -> dec
AST.paramDec LParam lore
param,
            lparamIndex :: IndexArray
lparamIndex = Maybe Indexed -> IndexArray
forall a b. a -> b -> a
const Maybe Indexed
forall a. Maybe a
Nothing
          }
    name :: VName
name = LParam lore -> VName
forall dec. Param dec -> VName
AST.paramName LParam lore
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 ::
  ASTLore lore =>
  [(AST.FParam lore, SubExp, SubExp)] ->
  SymbolTable lore ->
  SymbolTable lore
insertLoopMerge :: forall lore.
ASTLore lore =>
[(FParam lore, SubExp, SubExp)]
-> SymbolTable lore -> SymbolTable lore
insertLoopMerge = (SymbolTable lore
 -> [(Param (FParamInfo lore), SubExp, SubExp)] -> SymbolTable lore)
-> [(Param (FParamInfo lore), SubExp, SubExp)]
-> SymbolTable lore
-> SymbolTable lore
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((SymbolTable lore
  -> [(Param (FParamInfo lore), SubExp, SubExp)] -> SymbolTable lore)
 -> [(Param (FParamInfo lore), SubExp, SubExp)]
 -> SymbolTable lore
 -> SymbolTable lore)
-> (SymbolTable lore
    -> [(Param (FParamInfo lore), SubExp, SubExp)] -> SymbolTable lore)
-> [(Param (FParamInfo lore), SubExp, SubExp)]
-> SymbolTable lore
-> SymbolTable lore
forall a b. (a -> b) -> a -> b
$ (SymbolTable lore
 -> (Param (FParamInfo lore), SubExp, SubExp) -> SymbolTable lore)
-> SymbolTable lore
-> [(Param (FParamInfo lore), SubExp, SubExp)]
-> SymbolTable lore
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((SymbolTable lore
  -> (Param (FParamInfo lore), SubExp, SubExp) -> SymbolTable lore)
 -> SymbolTable lore
 -> [(Param (FParamInfo lore), SubExp, SubExp)]
 -> SymbolTable lore)
-> (SymbolTable lore
    -> (Param (FParamInfo lore), SubExp, SubExp) -> SymbolTable lore)
-> SymbolTable lore
-> [(Param (FParamInfo lore), SubExp, SubExp)]
-> SymbolTable lore
forall a b. (a -> b) -> a -> b
$ ((Param (FParamInfo lore), SubExp, SubExp)
 -> SymbolTable lore -> SymbolTable lore)
-> SymbolTable lore
-> (Param (FParamInfo lore), SubExp, SubExp)
-> SymbolTable lore
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Param (FParamInfo lore), SubExp, SubExp)
-> SymbolTable lore -> SymbolTable lore
forall {lore}.
ASTLore lore =>
(Param (FParamInfo lore), SubExp, SubExp)
-> SymbolTable lore -> SymbolTable lore
bind
  where
    bind :: (Param (FParamInfo lore), SubExp, SubExp)
-> SymbolTable lore -> SymbolTable lore
bind (Param (FParamInfo lore)
p, SubExp
initial, SubExp
res) =
      VName -> EntryType lore -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
VName -> EntryType lore -> SymbolTable lore -> SymbolTable lore
insertEntry (Param (FParamInfo lore) -> VName
forall dec. Param dec -> VName
paramName Param (FParamInfo lore)
p) (EntryType lore -> SymbolTable lore -> SymbolTable lore)
-> EntryType lore -> SymbolTable lore -> SymbolTable lore
forall a b. (a -> b) -> a -> b
$
        FParamEntry lore -> EntryType lore
forall lore. FParamEntry lore -> EntryType lore
FParam
          FParamEntry :: forall lore.
FParamInfo lore
-> Names -> Maybe (SubExp, SubExp) -> FParamEntry lore
FParamEntry
            { fparamDec :: FParamInfo lore
fparamDec = Param (FParamInfo lore) -> FParamInfo lore
forall dec. Param dec -> dec
AST.paramDec Param (FParamInfo lore)
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 :: ASTLore lore => VName -> IntType -> SubExp -> SymbolTable lore -> SymbolTable lore
insertLoopVar :: forall lore.
ASTLore lore =>
VName -> IntType -> SubExp -> SymbolTable lore -> SymbolTable lore
insertLoopVar VName
name IntType
it SubExp
bound = VName -> EntryType lore -> SymbolTable lore -> SymbolTable lore
forall lore.
ASTLore lore =>
VName -> EntryType lore -> SymbolTable lore -> SymbolTable lore
insertEntry VName
name EntryType lore
bind
  where
    bind :: EntryType lore
bind =
      LoopVarEntry lore -> EntryType lore
forall lore. LoopVarEntry lore -> EntryType lore
LoopVar
        LoopVarEntry :: forall lore. IntType -> SubExp -> LoopVarEntry lore
LoopVarEntry
          { loopVarType :: IntType
loopVarType = IntType
it,
            loopVarBound :: SubExp
loopVarBound = SubExp
bound
          }

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

consume :: VName -> SymbolTable lore -> SymbolTable lore
consume :: forall lore. 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
forall {lore}. 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 =
      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
forall {lore}. Entry lore -> Entry lore
consume'' VName
v (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'}
    consume'' :: Entry lore -> Entry lore
consume'' Entry lore
e = Entry lore
e {entryConsumed :: Bool
entryConsumed = Bool
True}

-- | Hide definitions of those entries that satisfy some predicate.
hideIf :: (Entry lore -> Bool) -> SymbolTable lore -> SymbolTable lore
hideIf :: forall lore.
(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 =
        Entry lore
entry
          { entryType :: EntryType lore
entryType =
              FreeVarEntry lore -> EntryType lore
forall lore. FreeVarEntry lore -> EntryType lore
FreeVar
                FreeVarEntry :: forall lore.
NameInfo lore -> (VName -> IndexArray) -> FreeVarEntry lore
FreeVarEntry
                  { freeVarDec :: NameInfo lore
freeVarDec = Entry lore -> NameInfo lore
forall lore. Entry lore -> NameInfo lore
entryInfo Entry lore
entry,
                    freeVarIndex :: VName -> IndexArray
freeVarIndex = \VName
_ [TPrimExp Int64 VName]
_ -> Maybe Indexed
forall a. Maybe a
Nothing
                  }
          }
      | 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 :: forall lore. 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