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

-- | Inspecing and modifying t'Pattern's, function parameters and
-- pattern elements.
module Futhark.IR.Prop.Patterns
  ( -- * Function parameters
    paramIdent,
    paramType,
    paramDeclType,

    -- * Pattern elements
    patElemIdent,
    patElemType,
    setPatElemLore,
    patternElements,
    patternIdents,
    patternContextIdents,
    patternValueIdents,
    patternNames,
    patternValueNames,
    patternContextNames,
    patternTypes,
    patternValueTypes,
    patternSize,

    -- * Pattern construction
    basicPattern,
  )
where

import Futhark.IR.Prop.Types (DeclTyped (..), Typed (..))
import Futhark.IR.Syntax

-- | The 'Type' of a parameter.
paramType :: Typed dec => Param dec -> Type
paramType :: forall dec. Typed dec => Param dec -> Type
paramType = Param dec -> Type
forall t. Typed t => t -> Type
typeOf

-- | The 'DeclType' of a parameter.
paramDeclType :: DeclTyped dec => Param dec -> DeclType
paramDeclType :: forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType = Param dec -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf

-- | An 'Ident' corresponding to a parameter.
paramIdent :: Typed dec => Param dec -> Ident
paramIdent :: forall dec. Typed dec => Param dec -> Ident
paramIdent Param dec
param = VName -> Type -> Ident
Ident (Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
param) (Param dec -> Type
forall t. Typed t => t -> Type
typeOf Param dec
param)

-- | An 'Ident' corresponding to a pattern element.
patElemIdent :: Typed dec => PatElemT dec -> Ident
patElemIdent :: forall dec. Typed dec => PatElemT dec -> Ident
patElemIdent PatElemT dec
pelem = VName -> Type -> Ident
Ident (PatElemT dec -> VName
forall dec. PatElemT dec -> VName
patElemName PatElemT dec
pelem) (PatElemT dec -> Type
forall t. Typed t => t -> Type
typeOf PatElemT dec
pelem)

-- | The type of a name bound by a t'PatElem'.
patElemType :: Typed dec => PatElemT dec -> Type
patElemType :: forall dec. Typed dec => PatElemT dec -> Type
patElemType = PatElemT dec -> Type
forall t. Typed t => t -> Type
typeOf

-- | Set the lore of a t'PatElem'.
setPatElemLore :: PatElemT oldattr -> newattr -> PatElemT newattr
setPatElemLore :: forall oldattr newattr.
PatElemT oldattr -> newattr -> PatElemT newattr
setPatElemLore PatElemT oldattr
pe newattr
x = (oldattr -> newattr) -> PatElemT oldattr -> PatElemT newattr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (newattr -> oldattr -> newattr
forall a b. a -> b -> a
const newattr
x) PatElemT oldattr
pe

-- | All pattern elements in the pattern - context first, then values.
patternElements :: PatternT dec -> [PatElemT dec]
patternElements :: forall dec. PatternT dec -> [PatElemT dec]
patternElements PatternT dec
pat = PatternT dec -> [PatElemT dec]
forall dec. PatternT dec -> [PatElemT dec]
patternContextElements PatternT dec
pat [PatElemT dec] -> [PatElemT dec] -> [PatElemT dec]
forall a. [a] -> [a] -> [a]
++ PatternT dec -> [PatElemT dec]
forall dec. PatternT dec -> [PatElemT dec]
patternValueElements PatternT dec
pat

-- | Return a list of the 'Ident's bound by the t'Pattern'.
patternIdents :: Typed dec => PatternT dec -> [Ident]
patternIdents :: forall dec. Typed dec => PatternT dec -> [Ident]
patternIdents PatternT dec
pat = PatternT dec -> [Ident]
forall dec. Typed dec => PatternT dec -> [Ident]
patternContextIdents PatternT dec
pat [Ident] -> [Ident] -> [Ident]
forall a. [a] -> [a] -> [a]
++ PatternT dec -> [Ident]
forall dec. Typed dec => PatternT dec -> [Ident]
patternValueIdents PatternT dec
pat

-- | Return a list of the context 'Ident's bound by the t'Pattern'.
patternContextIdents :: Typed dec => PatternT dec -> [Ident]
patternContextIdents :: forall dec. Typed dec => PatternT dec -> [Ident]
patternContextIdents = (PatElemT dec -> Ident) -> [PatElemT dec] -> [Ident]
forall a b. (a -> b) -> [a] -> [b]
map PatElemT dec -> Ident
forall dec. Typed dec => PatElemT dec -> Ident
patElemIdent ([PatElemT dec] -> [Ident])
-> (PatternT dec -> [PatElemT dec]) -> PatternT dec -> [Ident]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternT dec -> [PatElemT dec]
forall dec. PatternT dec -> [PatElemT dec]
patternContextElements

-- | Return a list of the value 'Ident's bound by the t'Pattern'.
patternValueIdents :: Typed dec => PatternT dec -> [Ident]
patternValueIdents :: forall dec. Typed dec => PatternT dec -> [Ident]
patternValueIdents = (PatElemT dec -> Ident) -> [PatElemT dec] -> [Ident]
forall a b. (a -> b) -> [a] -> [b]
map PatElemT dec -> Ident
forall dec. Typed dec => PatElemT dec -> Ident
patElemIdent ([PatElemT dec] -> [Ident])
-> (PatternT dec -> [PatElemT dec]) -> PatternT dec -> [Ident]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternT dec -> [PatElemT dec]
forall dec. PatternT dec -> [PatElemT dec]
patternValueElements

-- | Return a list of the 'Name's bound by the t'Pattern'.
patternNames :: PatternT dec -> [VName]
patternNames :: forall dec. PatternT dec -> [VName]
patternNames = (PatElemT dec -> VName) -> [PatElemT dec] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map PatElemT dec -> VName
forall dec. PatElemT dec -> VName
patElemName ([PatElemT dec] -> [VName])
-> (PatternT dec -> [PatElemT dec]) -> PatternT dec -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternT dec -> [PatElemT dec]
forall dec. PatternT dec -> [PatElemT dec]
patternElements

-- | Return a list of the 'Name's bound by the context part of the t'Pattern'.
patternContextNames :: PatternT dec -> [VName]
patternContextNames :: forall dec. PatternT dec -> [VName]
patternContextNames = (PatElemT dec -> VName) -> [PatElemT dec] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map PatElemT dec -> VName
forall dec. PatElemT dec -> VName
patElemName ([PatElemT dec] -> [VName])
-> (PatternT dec -> [PatElemT dec]) -> PatternT dec -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternT dec -> [PatElemT dec]
forall dec. PatternT dec -> [PatElemT dec]
patternContextElements

-- | Return a list of the 'Name's bound by the value part of the t'Pattern'.
patternValueNames :: PatternT dec -> [VName]
patternValueNames :: forall dec. PatternT dec -> [VName]
patternValueNames = (PatElemT dec -> VName) -> [PatElemT dec] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map PatElemT dec -> VName
forall dec. PatElemT dec -> VName
patElemName ([PatElemT dec] -> [VName])
-> (PatternT dec -> [PatElemT dec]) -> PatternT dec -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternT dec -> [PatElemT dec]
forall dec. PatternT dec -> [PatElemT dec]
patternValueElements

-- | Return a list of the typess bound by the pattern.
patternTypes :: Typed dec => PatternT dec -> [Type]
patternTypes :: forall dec. Typed dec => PatternT dec -> [Type]
patternTypes = (Ident -> Type) -> [Ident] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> Type
identType ([Ident] -> [Type])
-> (PatternT dec -> [Ident]) -> PatternT dec -> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternT dec -> [Ident]
forall dec. Typed dec => PatternT dec -> [Ident]
patternIdents

-- | Return a list of the typess bound by the value part of the pattern.
patternValueTypes :: Typed dec => PatternT dec -> [Type]
patternValueTypes :: forall dec. Typed dec => PatternT dec -> [Type]
patternValueTypes = (Ident -> Type) -> [Ident] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> Type
identType ([Ident] -> [Type])
-> (PatternT dec -> [Ident]) -> PatternT dec -> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternT dec -> [Ident]
forall dec. Typed dec => PatternT dec -> [Ident]
patternValueIdents

-- | Return the number of names bound by the pattern.
patternSize :: PatternT dec -> Int
patternSize :: forall dec. PatternT dec -> Int
patternSize (Pattern [PatElemT dec]
context [PatElemT dec]
values) = [PatElemT dec] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PatElemT dec]
context Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [PatElemT dec] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PatElemT dec]
values

-- | Create a pattern using 'Type' as the attribute.
basicPattern :: [Ident] -> [Ident] -> PatternT Type
basicPattern :: [Ident] -> [Ident] -> PatternT Type
basicPattern [Ident]
context [Ident]
values =
  [PatElemT Type] -> [PatElemT Type] -> PatternT Type
forall dec. [PatElemT dec] -> [PatElemT dec] -> PatternT dec
Pattern ((Ident -> PatElemT Type) -> [Ident] -> [PatElemT Type]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> PatElemT Type
patElem [Ident]
context) ((Ident -> PatElemT Type) -> [Ident] -> [PatElemT Type]
forall a b. (a -> b) -> [a] -> [b]
map Ident -> PatElemT Type
patElem [Ident]
values)
  where
    patElem :: Ident -> PatElemT Type
patElem (Ident VName
name Type
t) = VName -> Type -> PatElemT Type
forall dec. VName -> dec -> PatElemT dec
PatElem VName
name Type
t