{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE Trustworthy #-}
module Futhark.Internalise.TypesValues
  (
   -- * Internalising types
    BoundInTypes
  , boundInTypes
  , internaliseReturnType
  , internaliseEntryReturnType
  , internaliseParamTypes
  , internaliseType
  , internalisePrimType
  , internalisedTypeSize
  , internaliseSumType

  -- * Internalising values
  , internalisePrimValue
  )
  where

import Control.Monad.State
import Data.List (delete, find, foldl')
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import Data.Maybe

import qualified Language.Futhark as E
import Futhark.IR.SOACS as I
import Futhark.Internalise.Monad

internaliseUniqueness :: E.Uniqueness -> I.Uniqueness
internaliseUniqueness :: Uniqueness -> Uniqueness
internaliseUniqueness Uniqueness
E.Nonunique = Uniqueness
I.Nonunique
internaliseUniqueness Uniqueness
E.Unique = Uniqueness
I.Unique

-- | The names that are bound for some types, either implicitly or
-- explicitly.
newtype BoundInTypes = BoundInTypes (S.Set VName)
                       deriving (b -> BoundInTypes -> BoundInTypes
NonEmpty BoundInTypes -> BoundInTypes
BoundInTypes -> BoundInTypes -> BoundInTypes
(BoundInTypes -> BoundInTypes -> BoundInTypes)
-> (NonEmpty BoundInTypes -> BoundInTypes)
-> (forall b. Integral b => b -> BoundInTypes -> BoundInTypes)
-> Semigroup BoundInTypes
forall b. Integral b => b -> BoundInTypes -> BoundInTypes
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
stimes :: b -> BoundInTypes -> BoundInTypes
$cstimes :: forall b. Integral b => b -> BoundInTypes -> BoundInTypes
sconcat :: NonEmpty BoundInTypes -> BoundInTypes
$csconcat :: NonEmpty BoundInTypes -> BoundInTypes
<> :: BoundInTypes -> BoundInTypes -> BoundInTypes
$c<> :: BoundInTypes -> BoundInTypes -> BoundInTypes
Semigroup, Semigroup BoundInTypes
BoundInTypes
Semigroup BoundInTypes
-> BoundInTypes
-> (BoundInTypes -> BoundInTypes -> BoundInTypes)
-> ([BoundInTypes] -> BoundInTypes)
-> Monoid BoundInTypes
[BoundInTypes] -> BoundInTypes
BoundInTypes -> BoundInTypes -> BoundInTypes
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
mconcat :: [BoundInTypes] -> BoundInTypes
$cmconcat :: [BoundInTypes] -> BoundInTypes
mappend :: BoundInTypes -> BoundInTypes -> BoundInTypes
$cmappend :: BoundInTypes -> BoundInTypes -> BoundInTypes
mempty :: BoundInTypes
$cmempty :: BoundInTypes
$cp1Monoid :: Semigroup BoundInTypes
Monoid)

-- | Determine the names bound for some types.
boundInTypes :: [E.TypeParam] -> BoundInTypes
boundInTypes :: [TypeParam] -> BoundInTypes
boundInTypes = Set VName -> BoundInTypes
BoundInTypes (Set VName -> BoundInTypes)
-> ([TypeParam] -> Set VName) -> [TypeParam] -> BoundInTypes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName)
-> ([TypeParam] -> [VName]) -> [TypeParam] -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeParam -> Maybe VName) -> [TypeParam] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe TypeParam -> Maybe VName
forall a. TypeParamBase a -> Maybe a
isTypeParam
  where isTypeParam :: TypeParamBase a -> Maybe a
isTypeParam (E.TypeParamDim a
v SrcLoc
_) = a -> Maybe a
forall a. a -> Maybe a
Just a
v
        isTypeParam TypeParamBase a
_ = Maybe a
forall a. Maybe a
Nothing

internaliseParamTypes :: BoundInTypes
                      -> M.Map VName VName
                      -> [E.TypeBase (E.DimDecl VName) ()]
                      -> InternaliseM [[I.TypeBase ExtShape Uniqueness]]
internaliseParamTypes :: BoundInTypes
-> Map VName VName
-> [TypeBase (DimDecl VName) ()]
-> InternaliseM [[TypeBase ExtShape Uniqueness]]
internaliseParamTypes (BoundInTypes Set VName
bound) Map VName VName
pnames [TypeBase (DimDecl VName) ()]
ts =
  InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
-> InternaliseM [[TypeBase ExtShape Uniqueness]]
forall a. InternaliseTypeM a -> InternaliseM a
runInternaliseTypeM (InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
 -> InternaliseM [[TypeBase ExtShape Uniqueness]])
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
-> InternaliseM [[TypeBase ExtShape Uniqueness]]
forall a b. (a -> b) -> a -> b
$ DimTable
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
forall a. DimTable -> InternaliseTypeM a -> InternaliseTypeM a
withDims (DimTable
bound' DimTable -> DimTable -> DimTable
forall a. Semigroup a => a -> a -> a
<> (VName -> Ext SubExp) -> Map VName VName -> DimTable
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (SubExp -> Ext SubExp
forall a. a -> Ext a
Free (SubExp -> Ext SubExp) -> (VName -> SubExp) -> VName -> Ext SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) Map VName VName
pnames) (InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
 -> InternaliseTypeM [[TypeBase ExtShape Uniqueness]])
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
forall a b. (a -> b) -> a -> b
$
  (TypeBase (DimDecl VName) ()
 -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> [TypeBase (DimDecl VName) ()]
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeBase (DimDecl VName) ()
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM [TypeBase (DimDecl VName) ()]
ts
  where bound' :: DimTable
bound' = [(VName, Ext SubExp)] -> DimTable
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([VName] -> [Ext SubExp] -> [(VName, Ext SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Set VName -> [VName]
forall a. Set a -> [a]
S.toList Set VName
bound)
                                 ((VName -> Ext SubExp) -> [VName] -> [Ext SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> Ext SubExp
forall a. a -> Ext a
Free (SubExp -> Ext SubExp) -> (VName -> SubExp) -> VName -> Ext SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var) ([VName] -> [Ext SubExp]) -> [VName] -> [Ext SubExp]
forall a b. (a -> b) -> a -> b
$ Set VName -> [VName]
forall a. Set a -> [a]
S.toList Set VName
bound))

internaliseReturnType :: E.TypeBase (E.DimDecl VName) ()
                      -> InternaliseM [I.TypeBase ExtShape Uniqueness]
internaliseReturnType :: TypeBase (DimDecl VName) ()
-> InternaliseM [TypeBase ExtShape Uniqueness]
internaliseReturnType = ([[TypeBase ExtShape Uniqueness]]
 -> [TypeBase ExtShape Uniqueness])
-> InternaliseM [[TypeBase ExtShape Uniqueness]]
-> InternaliseM [TypeBase ExtShape Uniqueness]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[TypeBase ExtShape Uniqueness]] -> [TypeBase ExtShape Uniqueness]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (InternaliseM [[TypeBase ExtShape Uniqueness]]
 -> InternaliseM [TypeBase ExtShape Uniqueness])
-> (TypeBase (DimDecl VName) ()
    -> InternaliseM [[TypeBase ExtShape Uniqueness]])
-> TypeBase (DimDecl VName) ()
-> InternaliseM [TypeBase ExtShape Uniqueness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase (DimDecl VName) ()
-> InternaliseM [[TypeBase ExtShape Uniqueness]]
internaliseEntryReturnType

-- | As 'internaliseReturnType', but returns components of a top-level
-- tuple type piecemeal.
internaliseEntryReturnType :: E.TypeBase (E.DimDecl VName) ()
                           -> InternaliseM [[I.TypeBase ExtShape Uniqueness]]
internaliseEntryReturnType :: TypeBase (DimDecl VName) ()
-> InternaliseM [[TypeBase ExtShape Uniqueness]]
internaliseEntryReturnType TypeBase (DimDecl VName) ()
t = do
  let ts :: [TypeBase (DimDecl VName) ()]
ts = case TypeBase (DimDecl VName) () -> Maybe [TypeBase (DimDecl VName) ()]
forall dim as. TypeBase dim as -> Maybe [TypeBase dim as]
E.isTupleRecord TypeBase (DimDecl VName) ()
t of Just [TypeBase (DimDecl VName) ()]
tts | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [TypeBase (DimDecl VName) ()] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeBase (DimDecl VName) ()]
tts -> [TypeBase (DimDecl VName) ()]
tts
                                     Maybe [TypeBase (DimDecl VName) ()]
_ -> [TypeBase (DimDecl VName) ()
t]
  InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
-> InternaliseM [[TypeBase ExtShape Uniqueness]]
forall a. InternaliseTypeM a -> InternaliseM a
runInternaliseTypeM (InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
 -> InternaliseM [[TypeBase ExtShape Uniqueness]])
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
-> InternaliseM [[TypeBase ExtShape Uniqueness]]
forall a b. (a -> b) -> a -> b
$ (TypeBase (DimDecl VName) ()
 -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> [TypeBase (DimDecl VName) ()]
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeBase (DimDecl VName) ()
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM [TypeBase (DimDecl VName) ()]
ts

internaliseType :: E.TypeBase (E.DimDecl VName) ()
                -> InternaliseM [I.TypeBase I.ExtShape Uniqueness]
internaliseType :: TypeBase (DimDecl VName) ()
-> InternaliseM [TypeBase ExtShape Uniqueness]
internaliseType = InternaliseTypeM [TypeBase ExtShape Uniqueness]
-> InternaliseM [TypeBase ExtShape Uniqueness]
forall a. InternaliseTypeM a -> InternaliseM a
runInternaliseTypeM (InternaliseTypeM [TypeBase ExtShape Uniqueness]
 -> InternaliseM [TypeBase ExtShape Uniqueness])
-> (TypeBase (DimDecl VName) ()
    -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> TypeBase (DimDecl VName) ()
-> InternaliseM [TypeBase ExtShape Uniqueness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase (DimDecl VName) ()
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM

newId :: InternaliseTypeM Int
newId :: InternaliseTypeM Int
newId = do Int
i <- InternaliseTypeM Int
forall s (m :: * -> *). MonadState s m => m s
get
           Int -> InternaliseTypeM ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int -> InternaliseTypeM ()) -> Int -> InternaliseTypeM ()
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
           Int -> InternaliseTypeM Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
i

internaliseDim :: E.DimDecl VName
               -> InternaliseTypeM ExtSize
internaliseDim :: DimDecl VName -> InternaliseTypeM (Ext SubExp)
internaliseDim DimDecl VName
d =
  case DimDecl VName
d of
    DimDecl VName
E.AnyDim -> Int -> Ext SubExp
forall a. Int -> Ext a
Ext (Int -> Ext SubExp)
-> InternaliseTypeM Int -> InternaliseTypeM (Ext SubExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> InternaliseTypeM Int
newId
    E.ConstDim Int
n -> Ext SubExp -> InternaliseTypeM (Ext SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (Ext SubExp -> InternaliseTypeM (Ext SubExp))
-> Ext SubExp -> InternaliseTypeM (Ext SubExp)
forall a b. (a -> b) -> a -> b
$ SubExp -> Ext SubExp
forall a. a -> Ext a
Free (SubExp -> Ext SubExp) -> SubExp -> Ext SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
I.Int32 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n
    E.NamedDim QualName VName
name -> QualName VName -> InternaliseTypeM (Ext SubExp)
namedDim QualName VName
name
  where namedDim :: QualName VName -> InternaliseTypeM (Ext SubExp)
namedDim (E.QualName [VName]
_ VName
name) = do
          Maybe [SubExp]
subst <- InternaliseM (Maybe [SubExp]) -> InternaliseTypeM (Maybe [SubExp])
forall a. InternaliseM a -> InternaliseTypeM a
liftInternaliseM (InternaliseM (Maybe [SubExp])
 -> InternaliseTypeM (Maybe [SubExp]))
-> InternaliseM (Maybe [SubExp])
-> InternaliseTypeM (Maybe [SubExp])
forall a b. (a -> b) -> a -> b
$ VName -> InternaliseM (Maybe [SubExp])
lookupSubst VName
name
          Maybe (Ext SubExp)
is_dim <- VName -> InternaliseTypeM (Maybe (Ext SubExp))
lookupDim VName
name

          case (Maybe (Ext SubExp)
is_dim, Maybe [SubExp]
subst) of
            (Just Ext SubExp
dim, Maybe [SubExp]
_) -> Ext SubExp -> InternaliseTypeM (Ext SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return Ext SubExp
dim
            (Maybe (Ext SubExp)
Nothing, Just [SubExp
v]) -> Ext SubExp -> InternaliseTypeM (Ext SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (Ext SubExp -> InternaliseTypeM (Ext SubExp))
-> Ext SubExp -> InternaliseTypeM (Ext SubExp)
forall a b. (a -> b) -> a -> b
$ SubExp -> Ext SubExp
forall a. a -> Ext a
I.Free SubExp
v
            (Maybe (Ext SubExp), Maybe [SubExp])
_ -> Ext SubExp -> InternaliseTypeM (Ext SubExp)
forall (m :: * -> *) a. Monad m => a -> m a
return (Ext SubExp -> InternaliseTypeM (Ext SubExp))
-> Ext SubExp -> InternaliseTypeM (Ext SubExp)
forall a b. (a -> b) -> a -> b
$ SubExp -> Ext SubExp
forall a. a -> Ext a
I.Free (SubExp -> Ext SubExp) -> SubExp -> Ext SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
name

internaliseTypeM :: E.StructType
                 -> InternaliseTypeM [I.TypeBase ExtShape Uniqueness]
internaliseTypeM :: TypeBase (DimDecl VName) ()
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM TypeBase (DimDecl VName) ()
orig_t =
  case TypeBase (DimDecl VName) ()
orig_t of
    E.Array ()
_ Uniqueness
u ScalarTypeBase (DimDecl VName) ()
et ShapeDecl (DimDecl VName)
shape -> do
      [Ext SubExp]
dims <- ShapeDecl (DimDecl VName) -> InternaliseTypeM [Ext SubExp]
internaliseShape ShapeDecl (DimDecl VName)
shape
      [TypeBase ExtShape Uniqueness]
ets <- TypeBase (DimDecl VName) ()
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM (TypeBase (DimDecl VName) ()
 -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> TypeBase (DimDecl VName) ()
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar ScalarTypeBase (DimDecl VName) ()
et
      [TypeBase ExtShape Uniqueness]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall (m :: * -> *) a. Monad m => a -> m a
return [TypeBase ExtShape Uniqueness
-> ExtShape -> Uniqueness -> TypeBase ExtShape Uniqueness
forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
I.arrayOf TypeBase ExtShape Uniqueness
et' ([Ext SubExp] -> ExtShape
forall d. [d] -> ShapeBase d
Shape [Ext SubExp]
dims) (Uniqueness -> TypeBase ExtShape Uniqueness)
-> Uniqueness -> TypeBase ExtShape Uniqueness
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Uniqueness
internaliseUniqueness Uniqueness
u | TypeBase ExtShape Uniqueness
et' <- [TypeBase ExtShape Uniqueness]
ets ]
    E.Scalar (E.Prim PrimType
bt) ->
      [TypeBase ExtShape Uniqueness]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimType -> TypeBase ExtShape Uniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim (PrimType -> TypeBase ExtShape Uniqueness)
-> PrimType -> TypeBase ExtShape Uniqueness
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
internalisePrimType PrimType
bt]
    E.Scalar (E.Record Map Name (TypeBase (DimDecl VName) ())
ets)
      -- XXX: we map empty records to bools, because otherwise
      -- arrays of unit will lose their sizes.
      | Map Name (TypeBase (DimDecl VName) ()) -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Map Name (TypeBase (DimDecl VName) ())
ets -> [TypeBase ExtShape Uniqueness]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall (m :: * -> *) a. Monad m => a -> m a
return [PrimType -> TypeBase ExtShape Uniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Bool]
      | Bool
otherwise ->
          [[TypeBase ExtShape Uniqueness]] -> [TypeBase ExtShape Uniqueness]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[TypeBase ExtShape Uniqueness]]
 -> [TypeBase ExtShape Uniqueness])
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Name, TypeBase (DimDecl VName) ())
 -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> [(Name, TypeBase (DimDecl VName) ())]
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (TypeBase (DimDecl VName) ()
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM (TypeBase (DimDecl VName) ()
 -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> ((Name, TypeBase (DimDecl VName) ())
    -> TypeBase (DimDecl VName) ())
-> (Name, TypeBase (DimDecl VName) ())
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, TypeBase (DimDecl VName) ()) -> TypeBase (DimDecl VName) ()
forall a b. (a, b) -> b
snd) (Map Name (TypeBase (DimDecl VName) ())
-> [(Name, TypeBase (DimDecl VName) ())]
forall a. Map Name a -> [(Name, a)]
E.sortFields Map Name (TypeBase (DimDecl VName) ())
ets)
    E.Scalar E.TypeVar{} ->
      [Char] -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseTypeM: cannot handle type variable."
    E.Scalar E.Arrow{} ->
      [Char] -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> [Char] -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseTypeM: cannot handle function type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ TypeBase (DimDecl VName) () -> [Char]
forall a. Pretty a => a -> [Char]
pretty TypeBase (DimDecl VName) ()
orig_t
    E.Scalar (E.Sum Map Name [TypeBase (DimDecl VName) ()]
cs) -> do
      ([TypeBase ExtShape Uniqueness]
ts, Map Name (Int, [Int])
_) <- Map Name [TypeBase ExtShape Uniqueness]
-> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
internaliseConstructors (Map Name [TypeBase ExtShape Uniqueness]
 -> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int])))
-> InternaliseTypeM (Map Name [TypeBase ExtShape Uniqueness])
-> InternaliseTypeM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
                 ([TypeBase (DimDecl VName) ()]
 -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> Map Name [TypeBase (DimDecl VName) ()]
-> InternaliseTypeM (Map Name [TypeBase ExtShape Uniqueness])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (([[TypeBase ExtShape Uniqueness]]
 -> [TypeBase ExtShape Uniqueness])
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[TypeBase ExtShape Uniqueness]] -> [TypeBase ExtShape Uniqueness]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
 -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> ([TypeBase (DimDecl VName) ()]
    -> InternaliseTypeM [[TypeBase ExtShape Uniqueness]])
-> [TypeBase (DimDecl VName) ()]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase (DimDecl VName) ()
 -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> [TypeBase (DimDecl VName) ()]
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeBase (DimDecl VName) ()
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM) Map Name [TypeBase (DimDecl VName) ()]
cs
      [TypeBase ExtShape Uniqueness]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall (m :: * -> *) a. Monad m => a -> m a
return ([TypeBase ExtShape Uniqueness]
 -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> [TypeBase ExtShape Uniqueness]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase ExtShape Uniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim (IntType -> PrimType
I.IntType IntType
I.Int8) TypeBase ExtShape Uniqueness
-> [TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness]
forall a. a -> [a] -> [a]
: [TypeBase ExtShape Uniqueness]
ts

  where internaliseShape :: ShapeDecl (DimDecl VName) -> InternaliseTypeM [Ext SubExp]
internaliseShape = (DimDecl VName -> InternaliseTypeM (Ext SubExp))
-> [DimDecl VName] -> InternaliseTypeM [Ext SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimDecl VName -> InternaliseTypeM (Ext SubExp)
internaliseDim ([DimDecl VName] -> InternaliseTypeM [Ext SubExp])
-> (ShapeDecl (DimDecl VName) -> [DimDecl VName])
-> ShapeDecl (DimDecl VName)
-> InternaliseTypeM [Ext SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeDecl (DimDecl VName) -> [DimDecl VName]
forall dim. ShapeDecl dim -> [dim]
E.shapeDims

internaliseConstructors :: M.Map Name [I.TypeBase ExtShape Uniqueness]
                        -> ([I.TypeBase ExtShape Uniqueness],
                            M.Map Name (Int, [Int]))
internaliseConstructors :: Map Name [TypeBase ExtShape Uniqueness]
-> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
internaliseConstructors Map Name [TypeBase ExtShape Uniqueness]
cs =
  (([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
 -> ((Name, [TypeBase ExtShape Uniqueness]), Int)
 -> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int])))
-> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
-> [((Name, [TypeBase ExtShape Uniqueness]), Int)]
-> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
-> ((Name, [TypeBase ExtShape Uniqueness]), Int)
-> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
forall k (t :: * -> *) a a.
(Ord k, Foldable t, Eq a) =>
([a], Map k (a, [Int])) -> ((k, t a), a) -> ([a], Map k (a, [Int]))
onConstructor ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
forall a. Monoid a => a
mempty ([((Name, [TypeBase ExtShape Uniqueness]), Int)]
 -> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int])))
-> [((Name, [TypeBase ExtShape Uniqueness]), Int)]
-> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
forall a b. (a -> b) -> a -> b
$ [(Name, [TypeBase ExtShape Uniqueness])]
-> [Int] -> [((Name, [TypeBase ExtShape Uniqueness]), Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Map Name [TypeBase ExtShape Uniqueness]
-> [(Name, [TypeBase ExtShape Uniqueness])]
forall a. Map Name a -> [(Name, a)]
E.sortConstrs Map Name [TypeBase ExtShape Uniqueness]
cs) [Int
0..]
  where onConstructor :: ([a], Map k (a, [Int])) -> ((k, t a), a) -> ([a], Map k (a, [Int]))
onConstructor ([a]
ts, Map k (a, [Int])
mapping) ((k
c, t a
c_ts), a
i) =
          let ([(a, Int)]
_, [Int]
js, [a]
new_ts) =
                (([(a, Int)], [Int], [a]) -> a -> ([(a, Int)], [Int], [a]))
-> ([(a, Int)], [Int], [a]) -> t a -> ([(a, Int)], [Int], [a])
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ([(a, Int)], [Int], [a]) -> a -> ([(a, Int)], [Int], [a])
forall a.
Eq a =>
([(a, Int)], [Int], [a]) -> a -> ([(a, Int)], [Int], [a])
f ([a] -> [Int] -> [(a, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
ts [Int
0..], [Int]
forall a. Monoid a => a
mempty, [a]
forall a. Monoid a => a
mempty) t a
c_ts
          in ([a]
ts [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
new_ts, k -> (a, [Int]) -> Map k (a, [Int]) -> Map k (a, [Int])
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
c (a
i, [Int]
js) Map k (a, [Int])
mapping)
          where f :: ([(a, Int)], [Int], [a]) -> a -> ([(a, Int)], [Int], [a])
f ([(a, Int)]
ts', [Int]
js, [a]
new_ts) a
t
                  | Just (a
_, Int
j) <- ((a, Int) -> Bool) -> [(a, Int)] -> Maybe (a, Int)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((a -> a -> Bool
forall a. Eq a => a -> a -> Bool
==a
t) (a -> Bool) -> ((a, Int) -> a) -> (a, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, Int) -> a
forall a b. (a, b) -> a
fst) [(a, Int)]
ts' =
                      ((a, Int) -> [(a, Int)] -> [(a, Int)]
forall a. Eq a => a -> [a] -> [a]
delete (a
t, Int
j) [(a, Int)]
ts',
                       [Int]
js [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
j],
                       [a]
new_ts)
                  | Bool
otherwise =
                      ([(a, Int)]
ts',
                       [Int]
js [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [[a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ts Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
new_ts],
                       [a]
new_ts [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
t])

internaliseSumType :: M.Map Name [E.StructType]
                   -> InternaliseM ([I.TypeBase ExtShape Uniqueness],
                                    M.Map Name (Int, [Int]))
internaliseSumType :: Map Name [TypeBase (DimDecl VName) ()]
-> InternaliseM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
internaliseSumType Map Name [TypeBase (DimDecl VName) ()]
cs =
  InternaliseTypeM
  ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
-> InternaliseM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
forall a. InternaliseTypeM a -> InternaliseM a
runInternaliseTypeM (InternaliseTypeM
   ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
 -> InternaliseM
      ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int])))
-> InternaliseTypeM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
-> InternaliseM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase ExtShape Uniqueness]
-> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
internaliseConstructors (Map Name [TypeBase ExtShape Uniqueness]
 -> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int])))
-> InternaliseTypeM (Map Name [TypeBase ExtShape Uniqueness])
-> InternaliseTypeM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
  ([TypeBase (DimDecl VName) ()]
 -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> Map Name [TypeBase (DimDecl VName) ()]
-> InternaliseTypeM (Map Name [TypeBase ExtShape Uniqueness])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (([[TypeBase ExtShape Uniqueness]]
 -> [TypeBase ExtShape Uniqueness])
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[TypeBase ExtShape Uniqueness]] -> [TypeBase ExtShape Uniqueness]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
 -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> ([TypeBase (DimDecl VName) ()]
    -> InternaliseTypeM [[TypeBase ExtShape Uniqueness]])
-> [TypeBase (DimDecl VName) ()]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase (DimDecl VName) ()
 -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> [TypeBase (DimDecl VName) ()]
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeBase (DimDecl VName) ()
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM) Map Name [TypeBase (DimDecl VName) ()]
cs

-- | How many core language values are needed to represent one source
-- language value of the given type?
internalisedTypeSize :: E.TypeBase (E.DimDecl VName) () -> InternaliseM Int
internalisedTypeSize :: TypeBase (DimDecl VName) () -> InternaliseM Int
internalisedTypeSize = ([TypeBase ExtShape Uniqueness] -> Int)
-> InternaliseM [TypeBase ExtShape Uniqueness] -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [TypeBase ExtShape Uniqueness] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (InternaliseM [TypeBase ExtShape Uniqueness] -> InternaliseM Int)
-> (TypeBase (DimDecl VName) ()
    -> InternaliseM [TypeBase ExtShape Uniqueness])
-> TypeBase (DimDecl VName) ()
-> InternaliseM Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase (DimDecl VName) ()
-> InternaliseM [TypeBase ExtShape Uniqueness]
internaliseType

-- | Convert an external primitive to an internal primitive.
internalisePrimType :: E.PrimType -> I.PrimType
internalisePrimType :: PrimType -> PrimType
internalisePrimType (E.Signed IntType
t) = IntType -> PrimType
I.IntType IntType
t
internalisePrimType (E.Unsigned IntType
t) = IntType -> PrimType
I.IntType IntType
t
internalisePrimType (E.FloatType FloatType
t) = FloatType -> PrimType
I.FloatType FloatType
t
internalisePrimType PrimType
E.Bool = PrimType
I.Bool

-- | Convert an external primitive value to an internal primitive value.
internalisePrimValue :: E.PrimValue -> I.PrimValue
internalisePrimValue :: PrimValue -> PrimValue
internalisePrimValue (E.SignedValue IntValue
v) = IntValue -> PrimValue
I.IntValue IntValue
v
internalisePrimValue (E.UnsignedValue IntValue
v) = IntValue -> PrimValue
I.IntValue IntValue
v
internalisePrimValue (E.FloatValue FloatValue
v) = FloatValue -> PrimValue
I.FloatValue FloatValue
v
internalisePrimValue (E.BoolValue Bool
b) = Bool -> PrimValue
I.BoolValue Bool
b