{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}

module Futhark.Internalise.TypesValues
  ( -- * Internalising types
    internaliseReturnType,
    internaliseLambdaReturnType,
    internaliseEntryReturnType,
    internaliseType,
    internaliseParamTypes,
    internaliseLoopParamType,
    internalisePrimType,
    internalisedTypeSize,
    internaliseSumType,

    -- * Internalising values
    internalisePrimValue,
  )
where

import Control.Monad.State
import Data.Bitraversable (bitraverse)
import Data.List (delete, find, foldl')
import qualified Data.Map.Strict as M
import Data.Maybe
import Futhark.IR.SOACS as I
import Futhark.Internalise.Monad
import qualified Language.Futhark as E

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

newtype TypeState = TypeState {TypeState -> Int
typeCounter :: Int}

newtype InternaliseTypeM a
  = InternaliseTypeM (State TypeState a)
  deriving ((forall a b. (a -> b) -> InternaliseTypeM a -> InternaliseTypeM b)
-> (forall a b. a -> InternaliseTypeM b -> InternaliseTypeM a)
-> Functor InternaliseTypeM
forall a b. a -> InternaliseTypeM b -> InternaliseTypeM a
forall a b. (a -> b) -> InternaliseTypeM a -> InternaliseTypeM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> InternaliseTypeM b -> InternaliseTypeM a
$c<$ :: forall a b. a -> InternaliseTypeM b -> InternaliseTypeM a
fmap :: forall a b. (a -> b) -> InternaliseTypeM a -> InternaliseTypeM b
$cfmap :: forall a b. (a -> b) -> InternaliseTypeM a -> InternaliseTypeM b
Functor, Functor InternaliseTypeM
Functor InternaliseTypeM
-> (forall a. a -> InternaliseTypeM a)
-> (forall a b.
    InternaliseTypeM (a -> b)
    -> InternaliseTypeM a -> InternaliseTypeM b)
-> (forall a b c.
    (a -> b -> c)
    -> InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM c)
-> (forall a b.
    InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM b)
-> (forall a b.
    InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM a)
-> Applicative InternaliseTypeM
forall a. a -> InternaliseTypeM a
forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM a
forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM b
forall a b.
InternaliseTypeM (a -> b)
-> InternaliseTypeM a -> InternaliseTypeM b
forall a b c.
(a -> b -> c)
-> InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM a
$c<* :: forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM a
*> :: forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM b
$c*> :: forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM b
liftA2 :: forall a b c.
(a -> b -> c)
-> InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM c
$cliftA2 :: forall a b c.
(a -> b -> c)
-> InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM c
<*> :: forall a b.
InternaliseTypeM (a -> b)
-> InternaliseTypeM a -> InternaliseTypeM b
$c<*> :: forall a b.
InternaliseTypeM (a -> b)
-> InternaliseTypeM a -> InternaliseTypeM b
pure :: forall a. a -> InternaliseTypeM a
$cpure :: forall a. a -> InternaliseTypeM a
Applicative, Applicative InternaliseTypeM
Applicative InternaliseTypeM
-> (forall a b.
    InternaliseTypeM a
    -> (a -> InternaliseTypeM b) -> InternaliseTypeM b)
-> (forall a b.
    InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM b)
-> (forall a. a -> InternaliseTypeM a)
-> Monad InternaliseTypeM
forall a. a -> InternaliseTypeM a
forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM b
forall a b.
InternaliseTypeM a
-> (a -> InternaliseTypeM b) -> InternaliseTypeM b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> InternaliseTypeM a
$creturn :: forall a. a -> InternaliseTypeM a
>> :: forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM b
$c>> :: forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM b
>>= :: forall a b.
InternaliseTypeM a
-> (a -> InternaliseTypeM b) -> InternaliseTypeM b
$c>>= :: forall a b.
InternaliseTypeM a
-> (a -> InternaliseTypeM b) -> InternaliseTypeM b
Monad, MonadState TypeState)

runInternaliseTypeM :: InternaliseTypeM a -> a
runInternaliseTypeM :: forall a. InternaliseTypeM a -> a
runInternaliseTypeM = [VName] -> InternaliseTypeM a -> a
forall a. [VName] -> InternaliseTypeM a -> a
runInternaliseTypeM' [VName]
forall a. Monoid a => a
mempty

runInternaliseTypeM' :: [VName] -> InternaliseTypeM a -> a
runInternaliseTypeM' :: forall a. [VName] -> InternaliseTypeM a -> a
runInternaliseTypeM' [VName]
exts (InternaliseTypeM State TypeState a
m) = State TypeState a -> TypeState -> a
forall s a. State s a -> s -> a
evalState State TypeState a
m (TypeState -> a) -> TypeState -> a
forall a b. (a -> b) -> a -> b
$ Int -> TypeState
TypeState ([VName] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
exts)

internaliseParamTypes ::
  [E.TypeBase E.Size ()] ->
  InternaliseM [[I.TypeBase Shape Uniqueness]]
internaliseParamTypes :: [StructType] -> InternaliseM [[TypeBase Shape Uniqueness]]
internaliseParamTypes [StructType]
ts =
  ([TypeBase Shape Uniqueness]
 -> InternaliseM [TypeBase Shape Uniqueness])
-> [[TypeBase Shape Uniqueness]]
-> InternaliseM [[TypeBase Shape Uniqueness]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((TypeBase Shape Uniqueness
 -> InternaliseM (TypeBase Shape Uniqueness))
-> [TypeBase Shape Uniqueness]
-> InternaliseM [TypeBase Shape Uniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeBase Shape Uniqueness
-> InternaliseM (TypeBase Shape Uniqueness)
forall shape u. TypeBase shape u -> InternaliseM (TypeBase shape u)
mkAccCerts) ([[TypeBase Shape Uniqueness]]
 -> InternaliseM [[TypeBase Shape Uniqueness]])
-> (InternaliseTypeM [[TypeBase Shape Uniqueness]]
    -> [[TypeBase Shape Uniqueness]])
-> InternaliseTypeM [[TypeBase Shape Uniqueness]]
-> InternaliseM [[TypeBase Shape Uniqueness]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseTypeM [[TypeBase Shape Uniqueness]]
-> [[TypeBase Shape Uniqueness]]
forall a. InternaliseTypeM a -> a
runInternaliseTypeM (InternaliseTypeM [[TypeBase Shape Uniqueness]]
 -> InternaliseM [[TypeBase Shape Uniqueness]])
-> InternaliseTypeM [[TypeBase Shape Uniqueness]]
-> InternaliseM [[TypeBase Shape Uniqueness]]
forall a b. (a -> b) -> a -> b
$
    (StructType -> InternaliseTypeM [TypeBase Shape Uniqueness])
-> [StructType] -> InternaliseTypeM [[TypeBase Shape Uniqueness]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (([TypeBase ExtShape Uniqueness] -> [TypeBase Shape Uniqueness])
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
-> InternaliseTypeM [TypeBase Shape Uniqueness]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeBase ExtShape Uniqueness -> TypeBase Shape Uniqueness)
-> [TypeBase ExtShape Uniqueness] -> [TypeBase Shape Uniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase ExtShape Uniqueness -> TypeBase Shape Uniqueness
forall {u}. TypeBase ExtShape u -> TypeBase Shape u
onType) (InternaliseTypeM [TypeBase ExtShape Uniqueness]
 -> InternaliseTypeM [TypeBase Shape Uniqueness])
-> (StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> StructType
-> InternaliseTypeM [TypeBase Shape Uniqueness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName Int
-> StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM Map VName Int
forall a. Monoid a => a
mempty) [StructType]
ts
  where
    onType :: TypeBase ExtShape u -> TypeBase Shape u
onType = TypeBase Shape u -> Maybe (TypeBase Shape u) -> TypeBase Shape u
forall a. a -> Maybe a -> a
fromMaybe TypeBase Shape u
forall {a}. a
bad (Maybe (TypeBase Shape u) -> TypeBase Shape u)
-> (TypeBase ExtShape u -> Maybe (TypeBase Shape u))
-> TypeBase ExtShape u
-> TypeBase Shape u
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
forall u. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
hasStaticShape
    bad :: a
bad = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ String
"internaliseParamTypes: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [StructType] -> String
forall a. Pretty a => a -> String
pretty [StructType]
ts

-- We need to fix up the arrays for any Acc return values or loop
-- parameters.  We look at the concrete types for this, since the Acc
-- parameter name in the second list will just be something we made up.
fixupKnownTypes :: [TypeBase shape1 u1] -> [TypeBase shape2 u2] -> [TypeBase shape2 u2]
fixupKnownTypes :: forall shape1 u1 shape2 u2.
[TypeBase shape1 u1]
-> [TypeBase shape2 u2] -> [TypeBase shape2 u2]
fixupKnownTypes = (TypeBase shape1 u1 -> TypeBase shape2 u2 -> TypeBase shape2 u2)
-> [TypeBase shape1 u1]
-> [TypeBase shape2 u2]
-> [TypeBase shape2 u2]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase shape1 u1 -> TypeBase shape2 u2 -> TypeBase shape2 u2
forall {shape} {u} {shape} {u}.
TypeBase shape u -> TypeBase shape u -> TypeBase shape u
fixup
  where
    fixup :: TypeBase shape u -> TypeBase shape u -> TypeBase shape u
fixup (Acc VName
acc Shape
ispace [Type]
ts u
_) (Acc VName
_ Shape
_ [Type]
_ u
u2) = VName -> Shape -> [Type] -> u -> TypeBase shape u
forall shape u. VName -> Shape -> [Type] -> u -> TypeBase shape u
Acc VName
acc Shape
ispace [Type]
ts u
u2
    fixup TypeBase shape u
_ TypeBase shape u
t = TypeBase shape u
t

-- Generate proper certificates for the placeholder accumulator
-- certificates produced by internaliseType (identified with tag 0).
-- Only needed when we cannot use 'fixupKnownTypes'.
mkAccCerts :: TypeBase shape u -> InternaliseM (TypeBase shape u)
mkAccCerts :: forall shape u. TypeBase shape u -> InternaliseM (TypeBase shape u)
mkAccCerts (Array PrimType
pt shape
shape u
u) =
  TypeBase shape u -> InternaliseM (TypeBase shape u)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase shape u -> InternaliseM (TypeBase shape u))
-> TypeBase shape u -> InternaliseM (TypeBase shape u)
forall a b. (a -> b) -> a -> b
$ PrimType -> shape -> u -> TypeBase shape u
forall shape u. PrimType -> shape -> u -> TypeBase shape u
Array PrimType
pt shape
shape u
u
mkAccCerts (Acc VName
c Shape
shape [Type]
ts u
u) =
  VName -> Shape -> [Type] -> u -> TypeBase shape u
forall shape u. VName -> Shape -> [Type] -> u -> TypeBase shape u
Acc (VName -> Shape -> [Type] -> u -> TypeBase shape u)
-> InternaliseM VName
-> InternaliseM (Shape -> [Type] -> u -> TypeBase shape u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> InternaliseM VName
c' InternaliseM (Shape -> [Type] -> u -> TypeBase shape u)
-> InternaliseM Shape
-> InternaliseM ([Type] -> u -> TypeBase shape u)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Shape -> InternaliseM Shape
forall (f :: * -> *) a. Applicative f => a -> f a
pure Shape
shape InternaliseM ([Type] -> u -> TypeBase shape u)
-> InternaliseM [Type] -> InternaliseM (u -> TypeBase shape u)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Type] -> InternaliseM [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Type]
ts InternaliseM (u -> TypeBase shape u)
-> InternaliseM u -> InternaliseM (TypeBase shape u)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> u -> InternaliseM u
forall (f :: * -> *) a. Applicative f => a -> f a
pure u
u
  where
    c' :: InternaliseM VName
c'
      | VName -> Int
baseTag VName
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"acc_cert"
      | Bool
otherwise = VName -> InternaliseM VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
c
mkAccCerts TypeBase shape u
t = TypeBase shape u -> InternaliseM (TypeBase shape u)
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase shape u
t

internaliseLoopParamType ::
  E.TypeBase E.Size () ->
  [TypeBase shape u] ->
  InternaliseM [I.TypeBase Shape Uniqueness]
internaliseLoopParamType :: forall shape u.
StructType
-> [TypeBase shape u] -> InternaliseM [TypeBase Shape Uniqueness]
internaliseLoopParamType StructType
et [TypeBase shape u]
ts =
  [TypeBase shape u]
-> [TypeBase Shape Uniqueness] -> [TypeBase Shape Uniqueness]
forall shape1 u1 shape2 u2.
[TypeBase shape1 u1]
-> [TypeBase shape2 u2] -> [TypeBase shape2 u2]
fixupKnownTypes [TypeBase shape u]
ts ([TypeBase Shape Uniqueness] -> [TypeBase Shape Uniqueness])
-> ([[TypeBase Shape Uniqueness]] -> [TypeBase Shape Uniqueness])
-> [[TypeBase Shape Uniqueness]]
-> [TypeBase Shape Uniqueness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[TypeBase Shape Uniqueness]] -> [TypeBase Shape Uniqueness]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[TypeBase Shape Uniqueness]] -> [TypeBase Shape Uniqueness])
-> InternaliseM [[TypeBase Shape Uniqueness]]
-> InternaliseM [TypeBase Shape Uniqueness]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [StructType] -> InternaliseM [[TypeBase Shape Uniqueness]]
internaliseParamTypes [StructType
et]

internaliseReturnType ::
  E.StructRetType ->
  [TypeBase shape u] ->
  [I.TypeBase ExtShape Uniqueness]
internaliseReturnType :: forall shape u.
StructRetType
-> [TypeBase shape u] -> [TypeBase ExtShape Uniqueness]
internaliseReturnType (E.RetType [VName]
dims StructType
et) [TypeBase shape u]
ts =
  [TypeBase shape u]
-> [TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness]
forall shape1 u1 shape2 u2.
[TypeBase shape1 u1]
-> [TypeBase shape2 u2] -> [TypeBase shape2 u2]
fixupKnownTypes [TypeBase shape u]
ts ([TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness])
-> [TypeBase ExtShape Uniqueness] -> [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> a -> b
$ [VName]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
-> [TypeBase ExtShape Uniqueness]
forall a. [VName] -> InternaliseTypeM a -> a
runInternaliseTypeM' [VName]
dims (Map VName Int
-> StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM Map VName Int
exts StructType
et)
  where
    exts :: Map VName Int
exts = [(VName, Int)] -> Map VName Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Int)] -> Map VName Int)
-> [(VName, Int)] -> Map VName Int
forall a b. (a -> b) -> a -> b
$ [VName] -> [Int] -> [(VName, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
dims [Int
0 ..]

internaliseLambdaReturnType ::
  E.TypeBase E.Size () ->
  [TypeBase shape u] ->
  InternaliseM [I.TypeBase Shape NoUniqueness]
internaliseLambdaReturnType :: forall shape u.
StructType -> [TypeBase shape u] -> InternaliseM [Type]
internaliseLambdaReturnType StructType
et [TypeBase shape u]
ts =
  (TypeBase Shape Uniqueness -> Type)
-> [TypeBase Shape Uniqueness] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape Uniqueness -> Type
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl ([TypeBase Shape Uniqueness] -> [Type])
-> InternaliseM [TypeBase Shape Uniqueness] -> InternaliseM [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructType
-> [TypeBase shape u] -> InternaliseM [TypeBase Shape Uniqueness]
forall shape u.
StructType
-> [TypeBase shape u] -> InternaliseM [TypeBase Shape Uniqueness]
internaliseLoopParamType StructType
et [TypeBase shape u]
ts

-- | As 'internaliseReturnType', but returns components of a top-level
-- tuple type piecemeal.
internaliseEntryReturnType ::
  E.StructRetType ->
  [[I.TypeBase ExtShape Uniqueness]]
internaliseEntryReturnType :: StructRetType -> [[TypeBase ExtShape Uniqueness]]
internaliseEntryReturnType (E.RetType [VName]
dims StructType
et) =
  [VName]
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
-> [[TypeBase ExtShape Uniqueness]]
forall a. [VName] -> InternaliseTypeM a -> a
runInternaliseTypeM' [VName]
dims (InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
 -> [[TypeBase ExtShape Uniqueness]])
-> ([StructType]
    -> InternaliseTypeM [[TypeBase ExtShape Uniqueness]])
-> [StructType]
-> [[TypeBase ExtShape Uniqueness]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> [StructType]
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Map VName Int
-> StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM Map VName Int
exts) ([StructType] -> [[TypeBase ExtShape Uniqueness]])
-> [StructType] -> [[TypeBase ExtShape Uniqueness]]
forall a b. (a -> b) -> a -> b
$
    case StructType -> Maybe [StructType]
forall dim as. TypeBase dim as -> Maybe [TypeBase dim as]
E.isTupleRecord StructType
et of
      Just [StructType]
ets | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [StructType] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [StructType]
ets -> [StructType]
ets
      Maybe [StructType]
_ -> [StructType
et]
  where
    exts :: Map VName Int
exts = [(VName, Int)] -> Map VName Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Int)] -> Map VName Int)
-> [(VName, Int)] -> Map VName Int
forall a b. (a -> b) -> a -> b
$ [VName] -> [Int] -> [(VName, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
dims [Int
0 ..]

internaliseType ::
  E.TypeBase E.Size () ->
  [I.TypeBase I.ExtShape Uniqueness]
internaliseType :: StructType -> [TypeBase ExtShape Uniqueness]
internaliseType = InternaliseTypeM [TypeBase ExtShape Uniqueness]
-> [TypeBase ExtShape Uniqueness]
forall a. InternaliseTypeM a -> a
runInternaliseTypeM (InternaliseTypeM [TypeBase ExtShape Uniqueness]
 -> [TypeBase ExtShape Uniqueness])
-> (StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> StructType
-> [TypeBase ExtShape Uniqueness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName Int
-> StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM Map VName Int
forall a. Monoid a => a
mempty

newId :: InternaliseTypeM Int
newId :: InternaliseTypeM Int
newId = do
  Int
i <- (TypeState -> Int) -> InternaliseTypeM Int
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets TypeState -> Int
typeCounter
  (TypeState -> TypeState) -> InternaliseTypeM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((TypeState -> TypeState) -> InternaliseTypeM ())
-> (TypeState -> TypeState) -> InternaliseTypeM ()
forall a b. (a -> b) -> a -> b
$ \TypeState
s -> TypeState
s {typeCounter :: Int
typeCounter = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1}
  Int -> InternaliseTypeM Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i

internaliseDim ::
  M.Map VName Int ->
  E.Size ->
  InternaliseTypeM ExtSize
internaliseDim :: Map VName Int -> Size -> InternaliseTypeM (Ext SubExp)
internaliseDim Map VName Int
exts Size
d =
  case Size
d of
    E.AnySize Maybe VName
_ -> 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.ConstSize Int
n -> Ext SubExp -> InternaliseTypeM (Ext SubExp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (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.Int64 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n
    E.NamedSize QualName VName
name -> Ext SubExp -> InternaliseTypeM (Ext SubExp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ext SubExp -> InternaliseTypeM (Ext SubExp))
-> Ext SubExp -> InternaliseTypeM (Ext SubExp)
forall a b. (a -> b) -> a -> b
$ QualName VName -> Ext SubExp
namedDim QualName VName
name
  where
    namedDim :: QualName VName -> Ext SubExp
namedDim (E.QualName [VName]
_ VName
name)
      | Just Int
x <- VName
name VName -> Map VName Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName Int
exts = Int -> Ext SubExp
forall a. Int -> Ext a
I.Ext Int
x
      | Bool
otherwise = 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 ::
  M.Map VName Int ->
  E.StructType ->
  InternaliseTypeM [I.TypeBase ExtShape Uniqueness]
internaliseTypeM :: Map VName Int
-> StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM Map VName Int
exts StructType
orig_t =
  case StructType
orig_t of
    E.Array ()
_ Uniqueness
u Shape Size
shape ScalarTypeBase Size ()
et -> do
      [Ext SubExp]
dims <- Shape Size -> InternaliseTypeM [Ext SubExp]
internaliseShape Shape Size
shape
      [TypeBase ExtShape Uniqueness]
ets <- Map VName Int
-> StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM Map VName Int
exts (StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Size () -> StructType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
E.Scalar ScalarTypeBase Size ()
et
      [TypeBase ExtShape Uniqueness]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [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 (f :: * -> *) a. Applicative f => a -> f a
pure [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 StructType
ets)
      -- XXX: we map empty records to units, because otherwise
      -- arrays of unit will lose their sizes.
      | Map Name StructType -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Map Name StructType
ets -> [TypeBase ExtShape Uniqueness]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimType -> TypeBase ExtShape Uniqueness
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Unit]
      | 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, StructType)
 -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> [(Name, StructType)]
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Map VName Int
-> StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM Map VName Int
exts (StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> ((Name, StructType) -> StructType)
-> (Name, StructType)
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> StructType
forall a b. (a, b) -> b
snd) (Map Name StructType -> [(Name, StructType)]
forall a. Map Name a -> [(Name, a)]
E.sortFields Map Name StructType
ets)
    E.Scalar (E.TypeVar ()
_ Uniqueness
u QualName VName
tn [E.TypeArgType StructType
arr_t SrcLoc
_])
      | VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
E.qualLeaf QualName VName
tn) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
E.maxIntrinsicTag,
        VName -> String
baseString (QualName VName -> VName
forall vn. QualName vn -> vn
E.qualLeaf QualName VName
tn) String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"acc" -> do
          [Type]
ts <- (TypeBase ExtShape Uniqueness -> Type)
-> [TypeBase ExtShape Uniqueness] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase Shape Uniqueness -> Type
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl (TypeBase Shape Uniqueness -> Type)
-> (TypeBase ExtShape Uniqueness -> TypeBase Shape Uniqueness)
-> TypeBase ExtShape Uniqueness
-> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase ExtShape Uniqueness -> TypeBase Shape Uniqueness
forall {u}. TypeBase ExtShape u -> TypeBase Shape u
onAccType) ([TypeBase ExtShape Uniqueness] -> [Type])
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
-> InternaliseTypeM [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map VName Int
-> StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM Map VName Int
exts StructType
arr_t
          let acc_param :: VName
acc_param = Name -> Int -> VName
VName Name
"PLACEHOLDER" Int
0 -- See mkAccCerts.
              acc_t :: TypeBase shape Uniqueness
acc_t = VName -> Shape -> [Type] -> Uniqueness -> TypeBase shape Uniqueness
forall shape u. VName -> Shape -> [Type] -> u -> TypeBase shape u
Acc VName
acc_param ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [Int -> [Type] -> SubExp
forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [Type]
ts]) ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall u. TypeBase Shape u -> TypeBase Shape u
rowType [Type]
ts) (Uniqueness -> TypeBase shape Uniqueness)
-> Uniqueness -> TypeBase shape Uniqueness
forall a b. (a -> b) -> a -> b
$ Uniqueness -> Uniqueness
internaliseUniqueness Uniqueness
u
          [TypeBase ExtShape Uniqueness]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [TypeBase ExtShape Uniqueness
forall {shape}. TypeBase shape Uniqueness
acc_t]
    E.Scalar E.TypeVar {} ->
      String -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall a. HasCallStack => String -> a
error (String -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> String -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> a -> b
$ String
"internaliseTypeM: cannot handle type variable: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ StructType -> String
forall a. Pretty a => a -> String
pretty StructType
orig_t
    E.Scalar E.Arrow {} ->
      String -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall a. HasCallStack => String -> a
error (String -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> String -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall a b. (a -> b) -> a -> b
$ String
"internaliseTypeM: cannot handle function type: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ StructType -> String
forall a. Pretty a => a -> String
pretty StructType
orig_t
    E.Scalar (E.Sum Map Name [StructType]
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
<$> ([StructType] -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> Map Name [StructType]
-> 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])
-> ([StructType]
    -> InternaliseTypeM [[TypeBase ExtShape Uniqueness]])
-> [StructType]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> [StructType]
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Map VName Int
-> StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM Map VName Int
exts)) Map Name [StructType]
cs
      [TypeBase ExtShape Uniqueness]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([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 :: Shape Size -> InternaliseTypeM [Ext SubExp]
internaliseShape = (Size -> InternaliseTypeM (Ext SubExp))
-> [Size] -> InternaliseTypeM [Ext SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Map VName Int -> Size -> InternaliseTypeM (Ext SubExp)
internaliseDim Map VName Int
exts) ([Size] -> InternaliseTypeM [Ext SubExp])
-> (Shape Size -> [Size])
-> Shape Size
-> InternaliseTypeM [Ext SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape Size -> [Size]
forall dim. Shape dim -> [dim]
E.shapeDims

    onAccType :: TypeBase ExtShape u -> TypeBase Shape u
onAccType = TypeBase Shape u -> Maybe (TypeBase Shape u) -> TypeBase Shape u
forall a. a -> Maybe a -> a
fromMaybe TypeBase Shape u
forall {a}. a
bad (Maybe (TypeBase Shape u) -> TypeBase Shape u)
-> (TypeBase ExtShape u -> Maybe (TypeBase Shape u))
-> TypeBase ExtShape u
-> TypeBase Shape u
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
forall u. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
hasStaticShape
    bad :: a
bad = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ String
"internaliseTypeM Acc: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ StructType -> String
forall a. Pretty a => a -> String
pretty StructType
orig_t

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 :: * -> *} {shape} {a}.
(Ord k, Foldable t, Eq shape) =>
([TypeBase shape Uniqueness], Map k (a, [Int]))
-> ((k, t (TypeBase shape Uniqueness)), a)
-> ([TypeBase shape Uniqueness], 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 :: ([TypeBase shape Uniqueness], Map k (a, [Int]))
-> ((k, t (TypeBase shape Uniqueness)), a)
-> ([TypeBase shape Uniqueness], Map k (a, [Int]))
onConstructor ([TypeBase shape Uniqueness]
ts, Map k (a, [Int])
mapping) ((k
c, t (TypeBase shape Uniqueness)
c_ts), a
i) =
      let ([(TypeBase shape NoUniqueness, Int)]
_, [Int]
js, [TypeBase shape Uniqueness]
new_ts) =
            (([(TypeBase shape NoUniqueness, Int)], [Int],
  [TypeBase shape Uniqueness])
 -> TypeBase shape Uniqueness
 -> ([(TypeBase shape NoUniqueness, Int)], [Int],
     [TypeBase shape Uniqueness]))
-> ([(TypeBase shape NoUniqueness, Int)], [Int],
    [TypeBase shape Uniqueness])
-> t (TypeBase shape Uniqueness)
-> ([(TypeBase shape NoUniqueness, Int)], [Int],
    [TypeBase shape Uniqueness])
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ([(TypeBase shape NoUniqueness, Int)], [Int],
 [TypeBase shape Uniqueness])
-> TypeBase shape Uniqueness
-> ([(TypeBase shape NoUniqueness, Int)], [Int],
    [TypeBase shape Uniqueness])
forall {shape}.
Eq shape =>
([(TypeBase shape NoUniqueness, Int)], [Int],
 [TypeBase shape Uniqueness])
-> TypeBase shape Uniqueness
-> ([(TypeBase shape NoUniqueness, Int)], [Int],
    [TypeBase shape Uniqueness])
f ([TypeBase shape NoUniqueness]
-> [Int] -> [(TypeBase shape NoUniqueness, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((TypeBase shape Uniqueness -> TypeBase shape NoUniqueness)
-> [TypeBase shape Uniqueness] -> [TypeBase shape NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl [TypeBase shape Uniqueness]
ts) [Int
0 ..], [Int]
forall a. Monoid a => a
mempty, [TypeBase shape Uniqueness]
forall a. Monoid a => a
mempty) t (TypeBase shape Uniqueness)
c_ts
       in ([TypeBase shape Uniqueness]
ts [TypeBase shape Uniqueness]
-> [TypeBase shape Uniqueness] -> [TypeBase shape Uniqueness]
forall a. [a] -> [a] -> [a]
++ [TypeBase shape Uniqueness]
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 :: ([(TypeBase shape NoUniqueness, Int)], [Int],
 [TypeBase shape Uniqueness])
-> TypeBase shape Uniqueness
-> ([(TypeBase shape NoUniqueness, Int)], [Int],
    [TypeBase shape Uniqueness])
f ([(TypeBase shape NoUniqueness, Int)]
ts', [Int]
js, [TypeBase shape Uniqueness]
new_ts) TypeBase shape Uniqueness
t
          | Just (TypeBase shape NoUniqueness
_, Int
j) <- ((TypeBase shape NoUniqueness, Int) -> Bool)
-> [(TypeBase shape NoUniqueness, Int)]
-> Maybe (TypeBase shape NoUniqueness, Int)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((TypeBase shape NoUniqueness -> TypeBase shape NoUniqueness -> Bool
forall a. Eq a => a -> a -> Bool
== TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl TypeBase shape Uniqueness
t) (TypeBase shape NoUniqueness -> Bool)
-> ((TypeBase shape NoUniqueness, Int)
    -> TypeBase shape NoUniqueness)
-> (TypeBase shape NoUniqueness, Int)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase shape NoUniqueness, Int) -> TypeBase shape NoUniqueness
forall a b. (a, b) -> a
fst) [(TypeBase shape NoUniqueness, Int)]
ts' =
              ( (TypeBase shape NoUniqueness, Int)
-> [(TypeBase shape NoUniqueness, Int)]
-> [(TypeBase shape NoUniqueness, Int)]
forall a. Eq a => a -> [a] -> [a]
delete (TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl TypeBase shape Uniqueness
t, Int
j) [(TypeBase shape NoUniqueness, Int)]
ts',
                [Int]
js [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
j],
                [TypeBase shape Uniqueness]
new_ts
              )
          | Bool
otherwise =
              ( [(TypeBase shape NoUniqueness, Int)]
ts',
                [Int]
js [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [[TypeBase shape Uniqueness] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase shape Uniqueness]
ts Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [TypeBase shape Uniqueness] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase shape Uniqueness]
new_ts],
                [TypeBase shape Uniqueness]
new_ts [TypeBase shape Uniqueness]
-> [TypeBase shape Uniqueness] -> [TypeBase shape Uniqueness]
forall a. [a] -> [a] -> [a]
++ [TypeBase shape Uniqueness
t]
              )

internaliseSumType ::
  M.Map Name [E.StructType] ->
  InternaliseM
    ( [I.TypeBase ExtShape Uniqueness],
      M.Map Name (Int, [Int])
    )
internaliseSumType :: Map Name [StructType]
-> InternaliseM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
internaliseSumType Map Name [StructType]
cs =
  ([TypeBase ExtShape Uniqueness]
 -> InternaliseM [TypeBase ExtShape Uniqueness])
-> (Map Name (Int, [Int]) -> InternaliseM (Map Name (Int, [Int])))
-> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
-> InternaliseM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse ((TypeBase ExtShape Uniqueness
 -> InternaliseM (TypeBase ExtShape Uniqueness))
-> [TypeBase ExtShape Uniqueness]
-> InternaliseM [TypeBase ExtShape Uniqueness]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeBase ExtShape Uniqueness
-> InternaliseM (TypeBase ExtShape Uniqueness)
forall shape u. TypeBase shape u -> InternaliseM (TypeBase shape u)
mkAccCerts) Map Name (Int, [Int]) -> InternaliseM (Map Name (Int, [Int]))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
 -> InternaliseM
      ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int])))
-> (InternaliseTypeM
      ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
    -> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int])))
-> InternaliseTypeM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
-> InternaliseM
     ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseTypeM
  ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
-> ([TypeBase ExtShape Uniqueness], Map Name (Int, [Int]))
forall a. InternaliseTypeM a -> 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
<$> ([StructType] -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> Map Name [StructType]
-> 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])
-> ([StructType]
    -> InternaliseTypeM [[TypeBase ExtShape Uniqueness]])
-> [StructType]
-> InternaliseTypeM [TypeBase ExtShape Uniqueness]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness])
-> [StructType]
-> InternaliseTypeM [[TypeBase ExtShape Uniqueness]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Map VName Int
-> StructType -> InternaliseTypeM [TypeBase ExtShape Uniqueness]
internaliseTypeM Map VName Int
forall a. Monoid a => a
mempty)) Map Name [StructType]
cs

-- | How many core language values are needed to represent one source
-- language value of the given type?
internalisedTypeSize :: E.TypeBase E.Size als -> Int
-- A few special cases for performance.
internalisedTypeSize :: forall als. TypeBase Size als -> Int
internalisedTypeSize (E.Scalar (E.Prim PrimType
_)) = Int
1
internalisedTypeSize (E.Array als
_ Uniqueness
_ Shape Size
_ (E.Prim PrimType
_)) = Int
1
internalisedTypeSize TypeBase Size als
t = [TypeBase ExtShape Uniqueness] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([TypeBase ExtShape Uniqueness] -> Int)
-> [TypeBase ExtShape Uniqueness] -> Int
forall a b. (a -> b) -> a -> b
$ StructType -> [TypeBase ExtShape Uniqueness]
internaliseType (TypeBase Size als
t TypeBase Size als -> () -> StructType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`E.setAliases` ())

-- | 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