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

    -- * Internalising values
    internalisePrimValue,

    -- * For internal testing
    inferAliases,
    internaliseConstructors,
  )
where

import Control.Monad
import Control.Monad.Free (Free (..))
import Control.Monad.State
import Data.Bitraversable (bitraverse)
import Data.Foldable (toList)
import Data.List (delete, find, foldl')
import Data.List qualified as L
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.IR.SOACS hiding (Free)
import Futhark.IR.SOACS qualified as I
import Futhark.Internalise.Monad
import Futhark.Util (chunkLike)
import Language.Futhark qualified 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 -> 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
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
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 = forall a. [VName] -> InternaliseTypeM a -> a
runInternaliseTypeM' 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) = forall s a. State s a -> s -> a
evalState State TypeState a
m forall a b. (a -> b) -> a -> b
$ Int -> TypeState
TypeState (forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
exts)

internaliseParamTypes ::
  [E.ParamType] ->
  InternaliseM [[Tree (I.TypeBase Shape Uniqueness)]]
internaliseParamTypes :: [ParamType] -> InternaliseM [[Tree (TypeBase Shape Uniqueness)]]
internaliseParamTypes [ParamType]
ts =
  forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall shape u. TypeBase shape u -> InternaliseM (TypeBase shape u)
mkAccCerts)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. InternaliseTypeM a -> a
runInternaliseTypeM forall a b. (a -> b) -> a -> b
$
    forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall {u}. TypeBase ExtShape u -> TypeBase Shape u
onType)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName Int
-> ResType
-> InternaliseTypeM [Tree (TypeBase ExtShape Uniqueness)]
internaliseTypeM forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParamType -> ResType
E.paramToRes) [ParamType]
ts
  where
    onType :: TypeBase ExtShape u -> TypeBase Shape u
onType = forall a. a -> Maybe a -> a
fromMaybe forall {a}. a
bad forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. TypeBase ExtShape u -> Maybe (TypeBase Shape u)
hasStaticShape
    bad :: a
bad = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseParamTypes: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString [ParamType]
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, b)] ->
  [(TypeBase shape2 u2, b)]
fixupKnownTypes :: forall shape1 u1 shape2 u2 b.
[TypeBase shape1 u1]
-> [(TypeBase shape2 u2, b)] -> [(TypeBase shape2 u2, b)]
fixupKnownTypes = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {shape} {u} {shape} {u} {b}.
TypeBase shape u -> (TypeBase shape u, b) -> (TypeBase shape u, b)
fixup
  where
    fixup :: TypeBase shape u -> (TypeBase shape u, b) -> (TypeBase shape u, b)
fixup (Acc VName
acc Shape
ispace [TypeBase Shape NoUniqueness]
ts u
_) (Acc VName
_ Shape
_ [TypeBase Shape NoUniqueness]
_ u
u2, b
b) = (forall shape u.
VName
-> Shape -> [TypeBase Shape NoUniqueness] -> u -> TypeBase shape u
Acc VName
acc Shape
ispace [TypeBase Shape NoUniqueness]
ts u
u2, b
b)
    fixup TypeBase shape u
_ (TypeBase shape u, b)
t = (TypeBase shape u, b)
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) =
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> shape -> u -> TypeBase shape u
Array PrimType
pt shape
shape u
u
mkAccCerts (Acc VName
c Shape
shape [TypeBase Shape NoUniqueness]
ts u
u) =
  forall shape u.
VName
-> Shape -> [TypeBase Shape NoUniqueness] -> u -> TypeBase shape u
Acc forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> InternaliseM VName
c' forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure Shape
shape forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure [TypeBase Shape NoUniqueness]
ts forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure u
u
  where
    c' :: InternaliseM VName
c'
      | VName -> Int
baseTag VName
c forall a. Eq a => a -> a -> Bool
== Int
0 = forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"acc_cert"
      | Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
c
mkAccCerts TypeBase shape u
t = forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase shape u
t

internaliseLoopParamType ::
  E.ParamType ->
  [TypeBase shape u] ->
  InternaliseM [I.TypeBase Shape Uniqueness]
internaliseLoopParamType :: forall shape u.
ParamType
-> [TypeBase shape u] -> InternaliseM [TypeBase Shape Uniqueness]
internaliseLoopParamType ParamType
et [TypeBase shape u]
ts =
  forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall shape1 u1 shape2 u2 b.
[TypeBase shape1 u1]
-> [(TypeBase shape2 u2, b)] -> [(TypeBase shape2 u2, b)]
fixupKnownTypes [TypeBase shape u]
ts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (,()) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall (t :: * -> *) a. Foldable t => t a -> [a]
toList)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [ParamType] -> InternaliseM [[Tree (TypeBase Shape Uniqueness)]]
internaliseParamTypes [ParamType
et]

-- Tag every sublist with its offset in corresponding flattened list.
withOffsets :: (Foldable a) => [a b] -> [(a b, Int)]
withOffsets :: forall (a :: * -> *) b. Foldable a => [a b] -> [(a b, Int)]
withOffsets [a b]
xs = forall a b. [a] -> [b] -> [(a, b)]
zip [a b]
xs (forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl forall a. Num a => a -> a -> a
(+) Int
0 forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall (t :: * -> *) a. Foldable t => t a -> Int
length [a b]
xs)

numberFrom :: Int -> Tree a -> Tree (a, Int)
numberFrom :: forall a. Int -> Tree a -> Tree (a, Int)
numberFrom Int
o = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> a
evalState Int
o forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {b} {m :: * -> *} {f :: * -> *} {a}.
(MonadState b m, Num b, Traversable f) =>
Free f a -> m (Free f (a, b))
f
  where
    f :: Free f a -> m (Free f (a, b))
f (Pure a
x) = forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state forall a b. (a -> b) -> a -> b
$ \b
i -> (forall (f :: * -> *) a. a -> Free f a
Pure (a
x, b
i), b
i forall a. Num a => a -> a -> a
+ b
1)
    f (Free f (Free f a)
xs) = forall (f :: * -> *) a. f (Free f a) -> Free f a
Free forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Free f a -> m (Free f (a, b))
f f (Free f a)
xs

numberTrees :: [Tree a] -> [Tree (a, Int)]
numberTrees :: forall a. [Tree a] -> [Tree (a, Int)]
numberTrees = forall a b. (a -> b) -> [a] -> [b]
map (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Int -> Tree a -> Tree (a, Int)
numberFrom) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (a :: * -> *) b. Foldable a => [a b] -> [(a b, Int)]
withOffsets

nonuniqueArray :: TypeBase shape Uniqueness -> Bool
nonuniqueArray :: forall shape. TypeBase shape Uniqueness -> Bool
nonuniqueArray t :: TypeBase shape Uniqueness
t@Array {} = Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall shape. TypeBase shape Uniqueness -> Bool
unique TypeBase shape Uniqueness
t
nonuniqueArray TypeBase shape Uniqueness
_ = Bool
False

matchTrees :: Tree a -> Tree b -> Maybe (Tree (a, b))
matchTrees :: forall a b. Tree a -> Tree b -> Maybe (Tree (a, b))
matchTrees (Pure a
a) (Pure b
b) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. a -> Free f a
Pure (a
a, b
b)
matchTrees (Free [Free [] a]
as) (Free [Free [] b]
bs)
  | forall (t :: * -> *) a. Foldable t => t a -> Int
length [Free [] a]
as forall a. Eq a => a -> a -> Bool
== forall (t :: * -> *) a. Foldable t => t a -> Int
length [Free [] b]
bs =
      forall (f :: * -> *) a. f (Free f a) -> Free f a
Free forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM forall a b. Tree a -> Tree b -> Maybe (Tree (a, b))
matchTrees [Free [] a]
as [Free [] b]
bs
matchTrees Free [] a
_ Free [] b
_ = forall a. Maybe a
Nothing

subtreesMatching :: Tree a -> Tree b -> [Tree (a, b)]
subtreesMatching :: forall a b. Tree a -> Tree b -> [Tree (a, b)]
subtreesMatching Tree a
as Tree b
bs =
  case forall a b. Tree a -> Tree b -> Maybe (Tree (a, b))
matchTrees Tree a
as Tree b
bs of
    Just Tree (a, b)
m -> [Tree (a, b)
m]
    Maybe (Tree (a, b))
Nothing -> case Tree b
bs of
      Pure b
_ -> []
      Free [Tree b]
bs' -> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall a b. Tree a -> Tree b -> [Tree (a, b)]
subtreesMatching Tree a
as) [Tree b]
bs'

-- See Note [Alias Inference].
inferAliases ::
  [Tree (I.TypeBase Shape Uniqueness)] ->
  [Tree (I.TypeBase ExtShape Uniqueness)] ->
  [[(I.TypeBase ExtShape Uniqueness, RetAls)]]
inferAliases :: [Tree (TypeBase Shape Uniqueness)]
-> [Tree (TypeBase ExtShape Uniqueness)]
-> [[(TypeBase ExtShape Uniqueness, RetAls)]]
inferAliases [Tree (TypeBase Shape Uniqueness)]
all_param_ts [Tree (TypeBase ExtShape Uniqueness)]
all_res_ts =
  forall a b. (a -> b) -> [a] -> [b]
map forall {shape}.
Free [] (TypeBase shape Uniqueness)
-> [(TypeBase shape Uniqueness, RetAls)]
onRes [Tree (TypeBase ExtShape Uniqueness)]
all_res_ts
  where
    all_res_ts' :: [Tree (TypeBase ExtShape Uniqueness, Int)]
all_res_ts' = forall a. [Tree a] -> [Tree (a, Int)]
numberTrees [Tree (TypeBase ExtShape Uniqueness)]
all_res_ts
    all_param_ts' :: [Tree (TypeBase Shape Uniqueness, Int)]
all_param_ts' = forall a. [Tree a] -> [Tree (a, Int)]
numberTrees [Tree (TypeBase Shape Uniqueness)]
all_param_ts
    aliasable_param_ts :: [Tree (TypeBase Shape Uniqueness, Int)]
aliasable_param_ts = forall a. (a -> Bool) -> [a] -> [a]
filter (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall a b. (a -> b) -> a -> b
$ forall shape. TypeBase shape Uniqueness -> Bool
nonuniqueArray forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [Tree (TypeBase Shape Uniqueness, Int)]
all_param_ts'
    aliasable_res_ts :: [Tree (TypeBase ExtShape Uniqueness, Int)]
aliasable_res_ts = forall a. (a -> Bool) -> [a] -> [a]
filter (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall a b. (a -> b) -> a -> b
$ forall shape. TypeBase shape Uniqueness -> Bool
nonuniqueArray forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [Tree (TypeBase ExtShape Uniqueness, Int)]
all_res_ts'
    onRes :: Free [] (TypeBase shape Uniqueness)
-> [(TypeBase shape Uniqueness, RetAls)]
onRes (Pure TypeBase shape Uniqueness
res_t) =
      -- Necessarily a non-array.
      [(TypeBase shape Uniqueness
res_t, [Int] -> [Int] -> RetAls
RetAls forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty)]
    onRes (Free [Free [] (TypeBase shape Uniqueness)]
res_ts) =
      [ if forall shape. TypeBase shape Uniqueness -> Bool
nonuniqueArray TypeBase shape Uniqueness
res_t
          then (TypeBase shape Uniqueness
res_t, [Int] -> [Int] -> RetAls
RetAls [Int]
pals [Int]
rals)
          else (TypeBase shape Uniqueness
res_t, forall a. Monoid a => a
mempty)
        | (TypeBase shape Uniqueness
res_t, [Int]
pals, [Int]
rals) <- forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (forall (f :: * -> *) a. f (Free f a) -> Free f a
Free [Free [] (TypeBase shape Uniqueness)]
res_ts)) [[Int]]
palss [[Int]]
ralss
      ]
      where
        reorder :: [[a]] -> [[a]]
reorder [] = forall a. Int -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall (f :: * -> *) a. f (Free f a) -> Free f a
Free [Free [] (TypeBase shape Uniqueness)]
res_ts)) []
        reorder [[a]]
xs = forall a. [[a]] -> [[a]]
L.transpose [[a]]
xs
        infer :: t (Tree (a, a)) -> [[a]]
infer t (Tree (a, a))
ts =
          forall a. [[a]] -> [[a]]
reorder forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd)) forall a b. (a -> b) -> a -> b
$
            forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall a b. Tree a -> Tree b -> [Tree (a, b)]
subtreesMatching (forall (f :: * -> *) a. f (Free f a) -> Free f a
Free [Free [] (TypeBase shape Uniqueness)]
res_ts)) t (Tree (a, a))
ts
        palss :: [[Int]]
palss = forall {t :: * -> *} {a} {a}.
Foldable t =>
t (Tree (a, a)) -> [[a]]
infer [Tree (TypeBase Shape Uniqueness, Int)]
aliasable_param_ts
        ralss :: [[Int]]
ralss = forall {t :: * -> *} {a} {a}.
Foldable t =>
t (Tree (a, a)) -> [[a]]
infer [Tree (TypeBase ExtShape Uniqueness, Int)]
aliasable_res_ts

internaliseReturnType ::
  [Tree (I.TypeBase Shape Uniqueness)] ->
  E.ResRetType ->
  [TypeBase shape u] ->
  [(I.TypeBase ExtShape Uniqueness, RetAls)]
internaliseReturnType :: forall shape u.
[Tree (TypeBase Shape Uniqueness)]
-> ResRetType
-> [TypeBase shape u]
-> [(TypeBase ExtShape Uniqueness, RetAls)]
internaliseReturnType [Tree (TypeBase Shape Uniqueness)]
paramts (E.RetType [VName]
dims ResType
et) [TypeBase shape u]
ts =
  forall shape1 u1 shape2 u2 b.
[TypeBase shape1 u1]
-> [(TypeBase shape2 u2, b)] -> [(TypeBase shape2 u2, b)]
fixupKnownTypes [TypeBase shape u]
ts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Tree (TypeBase Shape Uniqueness)]
-> [Tree (TypeBase ExtShape Uniqueness)]
-> [[(TypeBase ExtShape Uniqueness, RetAls)]]
inferAliases [Tree (TypeBase Shape Uniqueness)]
paramts forall a b. (a -> b) -> a -> b
$
    forall a. [VName] -> InternaliseTypeM a -> a
runInternaliseTypeM' [VName]
dims (Map VName Int
-> ResType
-> InternaliseTypeM [Tree (TypeBase ExtShape Uniqueness)]
internaliseTypeM Map VName Int
exts ResType
et)
  where
    exts :: Map VName Int
exts = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
dims [Int
0 ..]

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

internaliseCoerceType ::
  E.StructType ->
  [TypeBase shape u] ->
  [I.TypeBase ExtShape Uniqueness]
internaliseCoerceType :: forall shape u.
StructType -> [TypeBase shape u] -> [TypeBase ExtShape Uniqueness]
internaliseCoerceType StructType
et [TypeBase shape u]
ts =
  forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall shape u.
[Tree (TypeBase Shape Uniqueness)]
-> ResRetType
-> [TypeBase shape u]
-> [(TypeBase ExtShape Uniqueness, RetAls)]
internaliseReturnType [] (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
E.RetType [] forall a b. (a -> b) -> a -> b
$ forall u. Uniqueness -> TypeBase Size u -> ResType
E.toRes Uniqueness
E.Nonunique StructType
et) [TypeBase shape u]
ts

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

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

newId :: InternaliseTypeM Int
newId :: InternaliseTypeM Int
newId = do
  Int
i <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets TypeState -> Int
typeCounter
  forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \TypeState
s -> TypeState
s {typeCounter :: Int
typeCounter = Int
i forall a. Num a => a -> a -> a
+ Int
1}
  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
    Size
e | Size
e forall a. Eq a => a -> a -> Bool
== Size
E.anySize -> forall a. Int -> Ext a
Ext forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> InternaliseTypeM Int
newId
    (E.IntLit Integer
n Info StructType
_ SrcLoc
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Ext a
I.Free forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
I.Int64 Integer
n
    (E.Var QualName VName
name Info StructType
_ SrcLoc
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> Ext SubExp
namedDim QualName VName
name
    Size
e -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpected size expression: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Size
e
  where
    namedDim :: QualName VName -> Ext SubExp
namedDim (E.QualName [VName]
_ VName
name)
      | Just Int
x <- VName
name forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName Int
exts = forall a. Int -> Ext a
I.Ext Int
x
      | Bool
otherwise = forall a. a -> Ext a
I.Free forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
name

-- | A tree is just an instantiation of the free monad with a list
-- monad.
--
-- The important thing is that we use it to represent the original
-- structure of arrayss, as this matters for aliasing.  Each 'Free'
-- constructor corresponds to an array dimension.  Only non-arrays
-- have a 'Pure' at the top level.  See Note [Alias Inference].
type Tree = Free []

internaliseTypeM ::
  M.Map VName Int ->
  E.ResType ->
  InternaliseTypeM [Tree (I.TypeBase ExtShape Uniqueness)]
internaliseTypeM :: Map VName Int
-> ResType
-> InternaliseTypeM [Tree (TypeBase ExtShape Uniqueness)]
internaliseTypeM Map VName Int
exts ResType
orig_t =
  case ResType
orig_t of
    E.Array Uniqueness
u Shape Size
shape ScalarTypeBase Size NoUniqueness
et -> do
      [Ext SubExp]
dims <- Shape Size -> InternaliseTypeM [Ext SubExp]
internaliseShape Shape Size
shape
      [Tree (TypeBase ExtShape Uniqueness)]
ets <- Map VName Int
-> ResType
-> InternaliseTypeM [Tree (TypeBase ExtShape Uniqueness)]
internaliseTypeM Map VName Int
exts forall a b. (a -> b) -> a -> b
$ forall u. Uniqueness -> TypeBase Size u -> ResType
E.toRes Uniqueness
E.Nonunique forall a b. (a -> b) -> a -> b
$ forall dim u. ScalarTypeBase dim u -> TypeBase dim u
E.Scalar ScalarTypeBase Size NoUniqueness
et
      let f :: TypeBase ExtShape u_unused -> TypeBase ExtShape Uniqueness
f TypeBase ExtShape u_unused
et' = forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
I.arrayOf TypeBase ExtShape u_unused
et' (forall d. [d] -> ShapeBase d
Shape [Ext SubExp]
dims) forall a b. (a -> b) -> a -> b
$ Uniqueness -> Uniqueness
internaliseUniqueness Uniqueness
u
      forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall {a}. [Free [] a] -> Free [] a
array forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall {u_unused}.
TypeBase ExtShape u_unused -> TypeBase ExtShape Uniqueness
f) [Tree (TypeBase ExtShape Uniqueness)]
ets]
    E.Scalar (E.Prim PrimType
bt) ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall (f :: * -> *) a. a -> Free f a
Pure forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> TypeBase shape u
I.Prim forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
internalisePrimType PrimType
bt]
    E.Scalar (E.Record Map Name ResType
ets)
      -- We map empty records to units, because otherwise arrays of
      -- unit will lose their sizes.
      | forall (t :: * -> *) a. Foldable t => t a -> Bool
null Map Name ResType
ets -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall (f :: * -> *) a. a -> Free f a
Pure forall a b. (a -> b) -> a -> b
$ forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
I.Unit]
      | Bool
otherwise ->
          forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Map VName Int
-> ResType
-> InternaliseTypeM [Tree (TypeBase ExtShape Uniqueness)]
internaliseTypeM Map VName Int
exts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) (forall a. Map Name a -> [(Name, a)]
E.sortFields Map Name ResType
ets)
    E.Scalar (E.TypeVar Uniqueness
u QualName VName
tn [E.TypeArgType StructType
arr_t])
      | VName -> Int
baseTag (forall vn. QualName vn -> vn
E.qualLeaf QualName VName
tn) forall a. Ord a => a -> a -> Bool
<= Int
E.maxIntrinsicTag,
        VName -> [Char]
baseString (forall vn. QualName vn -> vn
E.qualLeaf QualName VName
tn) forall a. Eq a => a -> a -> Bool
== [Char]
"acc" -> do
          [TypeBase Shape NoUniqueness]
ts <-
            forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {u}. TypeBase ExtShape u -> TypeBase Shape u
onAccType))
              forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map VName Int
-> ResType
-> InternaliseTypeM [Tree (TypeBase ExtShape Uniqueness)]
internaliseTypeM Map VName Int
exts (forall u. Uniqueness -> TypeBase Size u -> ResType
E.toRes Uniqueness
Nonunique StructType
arr_t)
          let acc_param :: VName
acc_param = Name -> Int -> VName
VName Name
"PLACEHOLDER" Int
0 -- See mkAccCerts.
              acc_shape :: Shape
acc_shape = forall d. [d] -> ShapeBase d
Shape [forall u. Int -> [TypeBase Shape u] -> SubExp
arraysSize Int
0 [TypeBase Shape NoUniqueness]
ts]
              u' :: Uniqueness
u' = Uniqueness -> Uniqueness
internaliseUniqueness Uniqueness
u
              acc_t :: TypeBase shape Uniqueness
acc_t = forall shape u.
VName
-> Shape -> [TypeBase Shape NoUniqueness] -> u -> TypeBase shape u
Acc VName
acc_param Shape
acc_shape (forall a b. (a -> b) -> [a] -> [b]
map forall u. TypeBase Shape u -> TypeBase Shape u
rowType [TypeBase Shape NoUniqueness]
ts) Uniqueness
u'
          forall (f :: * -> *) a. Applicative f => a -> f a
pure [forall (f :: * -> *) a. a -> Free f a
Pure forall {shape}. TypeBase shape Uniqueness
acc_t]
    E.Scalar E.TypeVar {} ->
      forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseTypeM: cannot handle type variable: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString ResType
orig_t
    E.Scalar E.Arrow {} ->
      forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseTypeM: cannot handle function type: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString ResType
orig_t
    E.Scalar (E.Sum Map Name [ResType]
cs) -> do
      ([Tree (TypeBase ExtShape Uniqueness)]
ts, Map Name (Int, [Int])
_) <-
        Map Name [Tree (TypeBase ExtShape Uniqueness)]
-> ([Tree (TypeBase ExtShape Uniqueness)], Map Name (Int, [Int]))
internaliseConstructors
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Map VName Int
-> ResType
-> InternaliseTypeM [Tree (TypeBase ExtShape Uniqueness)]
internaliseTypeM Map VName Int
exts)) Map Name [ResType]
cs
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. a -> Free f a
Pure (forall shape u. PrimType -> TypeBase shape u
I.Prim (IntType -> PrimType
I.IntType IntType
I.Int8)) forall a. a -> [a] -> [a]
: [Tree (TypeBase ExtShape Uniqueness)]
ts
  where
    internaliseShape :: Shape Size -> InternaliseTypeM [Ext SubExp]
internaliseShape = 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) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim. Shape dim -> [dim]
E.shapeDims
    array :: [Free [] a] -> Free [] a
array [Free [Free [] a]
ts] = forall (f :: * -> *) a. f (Free f a) -> Free f a
Free [Free [] a]
ts
    array [Free [] a]
ts = forall (f :: * -> *) a. f (Free f a) -> Free f a
Free [Free [] a]
ts

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

-- | Only exposed for testing purposes.
internaliseConstructors ::
  M.Map Name [Tree (I.TypeBase ExtShape Uniqueness)] ->
  ( [Tree (I.TypeBase ExtShape Uniqueness)],
    M.Map Name (Int, [Int])
  )
internaliseConstructors :: Map Name [Tree (TypeBase ExtShape Uniqueness)]
-> ([Tree (TypeBase ExtShape Uniqueness)], Map Name (Int, [Int]))
internaliseConstructors Map Name [Tree (TypeBase ExtShape Uniqueness)]
cs =
  forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {k} {t :: * -> *} {shape} {t :: * -> *} {a}.
(Ord k, Eq (t (TypeBase shape NoUniqueness)), Foldable t,
 Foldable t, Functor t) =>
([t (TypeBase shape Uniqueness)], Map k (a, [Int]))
-> ((k, t (t (TypeBase shape Uniqueness))), a)
-> ([t (TypeBase shape Uniqueness)], Map k (a, [Int]))
onConstructor forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip (forall a. Map Name a -> [(Name, a)]
E.sortConstrs Map Name [Tree (TypeBase ExtShape Uniqueness)]
cs) [Int
0 ..]
  where
    onConstructor :: ([t (TypeBase shape Uniqueness)], Map k (a, [Int]))
-> ((k, t (t (TypeBase shape Uniqueness))), a)
-> ([t (TypeBase shape Uniqueness)], Map k (a, [Int]))
onConstructor ([t (TypeBase shape Uniqueness)]
ts, Map k (a, [Int])
mapping) ((k
c, t (t (TypeBase shape Uniqueness))
c_ts), a
i) =
      let ([(t (TypeBase shape NoUniqueness), Int)]
_, [Int]
js, [t (TypeBase shape Uniqueness)]
new_ts) =
            forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {shape}.
Eq (t (TypeBase shape NoUniqueness)) =>
([(t (TypeBase shape NoUniqueness), Int)], [Int],
 [t (TypeBase shape Uniqueness)])
-> t (TypeBase shape Uniqueness)
-> ([(t (TypeBase shape NoUniqueness), Int)], [Int],
    [t (TypeBase shape Uniqueness)])
f (forall (a :: * -> *) b. Foldable a => [a b] -> [(a b, Int)]
withOffsets (forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl) [t (TypeBase shape Uniqueness)]
ts), forall a. Monoid a => a
mempty, forall a. Monoid a => a
mempty) t (t (TypeBase shape Uniqueness))
c_ts
       in ([t (TypeBase shape Uniqueness)]
ts forall a. [a] -> [a] -> [a]
++ [t (TypeBase shape Uniqueness)]
new_ts, 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
        size :: [t a] -> Int
size = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall (t :: * -> *) a. Foldable t => t a -> Int
length
        f :: ([(t (TypeBase shape NoUniqueness), Int)], [Int],
 [t (TypeBase shape Uniqueness)])
-> t (TypeBase shape Uniqueness)
-> ([(t (TypeBase shape NoUniqueness), Int)], [Int],
    [t (TypeBase shape Uniqueness)])
f ([(t (TypeBase shape NoUniqueness), Int)]
ts', [Int]
js, [t (TypeBase shape Uniqueness)]
new_ts) t (TypeBase shape Uniqueness)
t
          | Just (t (TypeBase shape NoUniqueness)
_, Int
j) <- forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl t (TypeBase shape Uniqueness)
t) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(t (TypeBase shape NoUniqueness), Int)]
ts' =
              ( forall a. Eq a => a -> [a] -> [a]
delete (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl t (TypeBase shape Uniqueness)
t, Int
j) [(t (TypeBase shape NoUniqueness), Int)]
ts',
                [Int]
js forall a. [a] -> [a] -> [a]
++ forall a. Int -> [a] -> [a]
take (forall (t :: * -> *) a. Foldable t => t a -> Int
length t (TypeBase shape Uniqueness)
t) [Int
j ..],
                [t (TypeBase shape Uniqueness)]
new_ts
              )
          | Bool
otherwise =
              ( [(t (TypeBase shape NoUniqueness), Int)]
ts',
                [Int]
js forall a. [a] -> [a] -> [a]
++ forall a. Int -> [a] -> [a]
take (forall (t :: * -> *) a. Foldable t => t a -> Int
length t (TypeBase shape Uniqueness)
t) [forall {a}. [t a] -> Int
size [t (TypeBase shape Uniqueness)]
ts forall a. Num a => a -> a -> a
+ forall {a}. [t a] -> Int
size [t (TypeBase shape Uniqueness)]
new_ts ..],
                [t (TypeBase shape Uniqueness)]
new_ts forall a. [a] -> [a] -> [a]
++ [t (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 =
  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 (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall shape u. TypeBase shape u -> InternaliseM (TypeBase shape u)
mkAccCerts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. InternaliseTypeM a -> a
runInternaliseTypeM forall a b. (a -> b) -> a -> b
$
    Map Name [Tree (TypeBase ExtShape Uniqueness)]
-> ([Tree (TypeBase ExtShape Uniqueness)], Map Name (Int, [Int]))
internaliseConstructors
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Map VName Int
-> ResType
-> InternaliseTypeM [Tree (TypeBase ExtShape Uniqueness)]
internaliseTypeM forall a. Monoid a => a
mempty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall u. Uniqueness -> TypeBase Size u -> ResType
E.toRes Uniqueness
E.Nonunique)) 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
_ Shape Size
_ (E.Prim PrimType
_)) = Int
1
internalisedTypeSize TypeBase Size als
t = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall (t :: * -> *) a. Foldable t => t a -> Int
length forall a b. (a -> b) -> a -> b
$ StructType -> [Tree (TypeBase ExtShape Uniqueness)]
internaliseType forall a b. (a -> b) -> a -> b
$ forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
E.toStruct TypeBase Size als
t

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

-- Note [Alias Inference]
--
-- The core language requires us to precisely indicate the aliasing of
-- function results (the RetAls type).  This is a problem when coming
-- from the source language, where it is implicit: a non-unique
-- function return value aliases every function argument.  The problem
-- now occurs because the core language uses a different value
-- representation than the source language - in particular, we do not
-- have arrays of tuples. E.g. @([]i32,[]i32)@ and @[](i32,i32)@ both
-- have the same core representation, but their implications for
-- aliasing are different.
--
--
-- To understand why this is a problem, consider a source program
--
--     def id (x: [](i32,i32)) = x
--
--     def f n =
--       let x = replicate n (0,0)
--       let x' = id x
--       let x'' = x' with [0] = (1,1)
--       in x''
--
-- With the core language value representation, it will be this:
--
--   def id (x1: []i32) (x2: []i32) = (x1,x2)
--
--   def f n =
--     let x1 = replicate n 0
--     let x2 = replicate n 0
--     let (x1', x2') = id x1 x2
--     let x1'' = x1' with [0] = 1
--     let x2'' = x2' with [0] = 1
--     in (x1'', x2'')
--
-- The results of 'id' alias *both* of the arguments, so x1' aliases
-- x1 and x2, and x2' also aliases x1 and x2.  This means that the
-- first with-expression will consume all of x1/x2/x1'/x2', and then
-- the second with-expression is a type error, as it references a
-- consumed variable.
--
-- Our solution is to deduce the possible aliasing such that
-- components that originally constituted the same array-of-tuples are
-- not aliased.  The main complexity is that we have to keep
-- information on the original (source) type structure around for a
-- while.  This is done with the Tree type.