module Futhark.Internalise.TypesValues
  ( -- * Internalising types
    internaliseReturnType,
    internaliseCoerceType,
    internaliseLambdaReturnType,
    internaliseEntryReturnType,
    internaliseType,
    internaliseParamTypes,
    internaliseLoopParamType,
    internalisePrimType,
    internalisedTypeSize,
    internaliseSumTypeRep,
    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.Bifunctor
import Data.Bitraversable (bitraverse)
import Data.Foldable (toList)
import Data.List (delete, find)
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

-- | The 'E.Diet' of a source parameter becomes the 'I.Diet' of the
-- internalised parameter; the two types coincide.
internaliseDiet :: E.Diet -> I.Diet
internaliseDiet :: Diet -> Diet
internaliseDiet Diet
E.Observe = Diet
I.Observe
internaliseDiet Diet
E.Consume = Diet
I.Consume

-- | The mode of an array type, if it is an array. Used to ask whether aliasing
-- is permitted at a position, which is 'E.Nonfresh' for a result and
-- 'I.Observe' for a parameter.
arrayMode :: TypeBase shape o -> Maybe o
arrayMode :: forall shape o. TypeBase shape o -> Maybe o
arrayMode (Array PrimType
_ shape
_ o
o) = o -> Maybe o
forall a. a -> Maybe a
Just o
o
arrayMode TypeBase shape o
_ = Maybe o
forall a. Maybe a
Nothing

aliasableArray :: (o -> Bool) -> TypeBase shape o -> Bool
aliasableArray :: forall o shape. (o -> Bool) -> TypeBase shape o -> Bool
aliasableArray o -> Bool
p = Bool -> (o -> Bool) -> Maybe o -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False o -> Bool
p (Maybe o -> Bool)
-> (TypeBase shape o -> Maybe o) -> TypeBase shape o -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase shape o -> Maybe o
forall shape o. TypeBase shape o -> Maybe o
arrayMode

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
$cfmap :: forall a b. (a -> b) -> InternaliseTypeM a -> InternaliseTypeM b
fmap :: forall a b. (a -> b) -> InternaliseTypeM a -> InternaliseTypeM b
$c<$ :: forall a b. a -> InternaliseTypeM b -> InternaliseTypeM a
<$ :: forall a b. a -> InternaliseTypeM b -> InternaliseTypeM a
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
$cpure :: forall a. a -> InternaliseTypeM a
pure :: forall a. a -> InternaliseTypeM a
$c<*> :: forall a b.
InternaliseTypeM (a -> b)
-> InternaliseTypeM a -> InternaliseTypeM b
<*> :: forall a b.
InternaliseTypeM (a -> b)
-> InternaliseTypeM a -> InternaliseTypeM b
$cliftA2 :: forall a b c.
(a -> b -> c)
-> InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM c
liftA2 :: forall a b c.
(a -> b -> c)
-> InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM c
$c*> :: forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM b
*> :: forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM b
$c<* :: forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM a
<* :: forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> 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
$c>>= :: forall a b.
InternaliseTypeM a
-> (a -> InternaliseTypeM b) -> InternaliseTypeM b
>>= :: forall a b.
InternaliseTypeM a
-> (a -> InternaliseTypeM b) -> InternaliseTypeM b
$c>> :: forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM b
>> :: forall a b.
InternaliseTypeM a -> InternaliseTypeM b -> InternaliseTypeM b
$creturn :: forall a. a -> InternaliseTypeM a
return :: forall a. a -> InternaliseTypeM a
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 a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VName]
exts)

internaliseParamTypes ::
  [E.ParamType] ->
  InternaliseM [[Tree (I.TypeBase Shape I.Diet)]]
internaliseParamTypes :: [ParamType] -> InternaliseM [[Tree (TypeBase Shape Diet)]]
internaliseParamTypes [ParamType]
ts =
  ([Tree (TypeBase Shape Diet)]
 -> InternaliseM [Tree (TypeBase Shape Diet)])
-> [[Tree (TypeBase Shape Diet)]]
-> InternaliseM [[Tree (TypeBase Shape Diet)]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((Tree (TypeBase Shape Diet)
 -> InternaliseM (Tree (TypeBase Shape Diet)))
-> [Tree (TypeBase Shape Diet)]
-> InternaliseM [Tree (TypeBase Shape Diet)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((TypeBase Shape Diet -> InternaliseM (TypeBase Shape Diet))
-> Tree (TypeBase Shape Diet)
-> InternaliseM (Tree (TypeBase Shape Diet))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Free [] a -> m (Free [] b)
mapM TypeBase Shape Diet -> InternaliseM (TypeBase Shape Diet)
forall shape o. TypeBase shape o -> InternaliseM (TypeBase shape o)
mkAccCerts)) ([[Tree (TypeBase Shape Diet)]]
 -> InternaliseM [[Tree (TypeBase Shape Diet)]])
-> (InternaliseTypeM [[Tree (TypeBase Shape Diet)]]
    -> [[Tree (TypeBase Shape Diet)]])
-> InternaliseTypeM [[Tree (TypeBase Shape Diet)]]
-> InternaliseM [[Tree (TypeBase Shape Diet)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseTypeM [[Tree (TypeBase Shape Diet)]]
-> [[Tree (TypeBase Shape Diet)]]
forall a. InternaliseTypeM a -> a
runInternaliseTypeM (InternaliseTypeM [[Tree (TypeBase Shape Diet)]]
 -> InternaliseM [[Tree (TypeBase Shape Diet)]])
-> InternaliseTypeM [[Tree (TypeBase Shape Diet)]]
-> InternaliseM [[Tree (TypeBase Shape Diet)]]
forall a b. (a -> b) -> a -> b
$
    (ParamType -> InternaliseTypeM [Tree (TypeBase Shape Diet)])
-> [ParamType] -> InternaliseTypeM [[Tree (TypeBase Shape Diet)]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (([Free [] (TypeBase ExtShape Diet)]
 -> [Tree (TypeBase Shape Diet)])
-> InternaliseTypeM [Free [] (TypeBase ExtShape Diet)]
-> InternaliseTypeM [Tree (TypeBase Shape Diet)]
forall a b. (a -> b) -> InternaliseTypeM a -> InternaliseTypeM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Free [] (TypeBase ExtShape Diet) -> Tree (TypeBase Shape Diet))
-> [Free [] (TypeBase ExtShape Diet)]
-> [Tree (TypeBase Shape Diet)]
forall a b. (a -> b) -> [a] -> [b]
map ((TypeBase ExtShape Diet -> TypeBase Shape Diet)
-> Free [] (TypeBase ExtShape Diet) -> Tree (TypeBase Shape Diet)
forall a b. (a -> b) -> Free [] a -> Free [] b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase ExtShape Diet -> TypeBase Shape Diet
onType)) (InternaliseTypeM [Free [] (TypeBase ExtShape Diet)]
 -> InternaliseTypeM [Tree (TypeBase Shape Diet)])
-> (ParamType
    -> InternaliseTypeM [Free [] (TypeBase ExtShape Diet)])
-> ParamType
-> InternaliseTypeM [Tree (TypeBase Shape Diet)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName Int
-> ParamType -> InternaliseTypeM [Free [] (TypeBase ExtShape Diet)]
forall o.
Pretty o =>
Map VName Int
-> TypeBase Size o -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
internaliseTypeM Map VName Int
forall a. Monoid a => a
mempty) [ParamType]
ts
  where
    onType :: TypeBase ExtShape Diet -> TypeBase Shape Diet
onType = (Diet -> Diet) -> TypeBase Shape Diet -> TypeBase Shape Diet
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Diet -> Diet
internaliseDiet (TypeBase Shape Diet -> TypeBase Shape Diet)
-> (TypeBase ExtShape Diet -> TypeBase Shape Diet)
-> TypeBase ExtShape Diet
-> TypeBase Shape Diet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Shape Diet
-> Maybe (TypeBase Shape Diet) -> TypeBase Shape Diet
forall a. a -> Maybe a -> a
fromMaybe TypeBase Shape Diet
forall {b}. b
bad (Maybe (TypeBase Shape Diet) -> TypeBase Shape Diet)
-> (TypeBase ExtShape Diet -> Maybe (TypeBase Shape Diet))
-> TypeBase ExtShape Diet
-> TypeBase Shape Diet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase ExtShape Diet -> Maybe (TypeBase Shape Diet)
forall o. TypeBase ExtShape o -> Maybe (TypeBase Shape o)
hasStaticShape
    bad :: b
bad = [Char] -> b
forall a. HasCallStack => [Char] -> a
error ([Char] -> b) -> [Char] -> b
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseParamTypes: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [ParamType] -> [Char]
forall a. Pretty a => a -> [Char]
prettyString [ParamType]
ts

-- Replace an accumulator's token, index space, and element types with those of
-- a known accumulator type. We must do this because these components cannot be
-- recovered from a source type: 'internaliseTypeM' produces a placeholder token
-- and a guessed index space. The known type is computed elsewhere (from
-- concrete loop values, or from an accumulator parameter).
fixupAcc :: TypeBase shape1 u1 -> (TypeBase shape2 u2, b) -> (TypeBase shape2 u2, b)
fixupAcc :: forall shape1 u1 shape2 u2 b.
TypeBase shape1 u1
-> (TypeBase shape2 u2, b) -> (TypeBase shape2 u2, b)
fixupAcc (Acc VName
acc Shape
ispace [Type]
ts) (Acc {}, b
b) = (VName -> Shape -> [Type] -> TypeBase shape2 u2
forall shape o. VName -> Shape -> [Type] -> TypeBase shape o
Acc VName
acc Shape
ispace [Type]
ts, b
b)
fixupAcc TypeBase shape1 u1
_ (TypeBase shape2 u2, b)
t = (TypeBase shape2 u2, b)
t

-- Fix up accumulators using a positionally-matching list of concrete
-- types (e.g. the actual types of loop values).
fixupKnownTypes ::
  [TypeBase shape1 o1] ->
  [(TypeBase shape2 o2, b)] ->
  [(TypeBase shape2 o2, b)]
fixupKnownTypes :: forall shape1 o1 shape2 o2 b.
[TypeBase shape1 o1]
-> [(TypeBase shape2 o2, b)] -> [(TypeBase shape2 o2, b)]
fixupKnownTypes = (TypeBase shape1 o1
 -> (TypeBase shape2 o2, b) -> (TypeBase shape2 o2, b))
-> [TypeBase shape1 o1]
-> [(TypeBase shape2 o2, b)]
-> [(TypeBase shape2 o2, b)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase shape1 o1
-> (TypeBase shape2 o2, b) -> (TypeBase shape2 o2, b)
forall shape1 u1 shape2 u2 b.
TypeBase shape1 u1
-> (TypeBase shape2 u2, b) -> (TypeBase shape2 u2, b)
fixupAcc

-- 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 o -> InternaliseM (TypeBase shape o)
mkAccCerts :: forall shape o. TypeBase shape o -> InternaliseM (TypeBase shape o)
mkAccCerts (Array PrimType
pt shape
shape o
o) =
  TypeBase shape o -> InternaliseM (TypeBase shape o)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase shape o -> InternaliseM (TypeBase shape o))
-> TypeBase shape o -> InternaliseM (TypeBase shape o)
forall a b. (a -> b) -> a -> b
$ PrimType -> shape -> o -> TypeBase shape o
forall shape o. PrimType -> shape -> o -> TypeBase shape o
Array PrimType
pt shape
shape o
o
mkAccCerts (Acc VName
c Shape
shape [Type]
ts) =
  VName -> Shape -> [Type] -> TypeBase shape o
forall shape o. VName -> Shape -> [Type] -> TypeBase shape o
Acc (VName -> Shape -> [Type] -> TypeBase shape o)
-> InternaliseM VName
-> InternaliseM (Shape -> [Type] -> TypeBase shape o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> InternaliseM VName
c' InternaliseM (Shape -> [Type] -> TypeBase shape o)
-> InternaliseM Shape -> InternaliseM ([Type] -> TypeBase shape o)
forall a b.
InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Shape -> InternaliseM Shape
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Shape
shape InternaliseM ([Type] -> TypeBase shape o)
-> InternaliseM [Type] -> InternaliseM (TypeBase shape o)
forall a b.
InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Type] -> InternaliseM [Type]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Type]
ts
  where
    c' :: InternaliseM VName
c'
      | VName -> Int
baseTag VName
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Name -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"acc_cert"
      | Bool
otherwise = VName -> InternaliseM VName
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
c
mkAccCerts TypeBase shape o
t = TypeBase shape o -> InternaliseM (TypeBase shape o)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase shape o
t

internaliseLoopParamType ::
  E.ParamType ->
  [TypeBase shape o] ->
  InternaliseM [I.TypeBase Shape Diet]
internaliseLoopParamType :: forall shape o.
ParamType
-> [TypeBase shape o] -> InternaliseM [TypeBase Shape Diet]
internaliseLoopParamType ParamType
et [TypeBase shape o]
ts =
  ((TypeBase Shape Diet, ()) -> TypeBase Shape Diet)
-> [(TypeBase Shape Diet, ())] -> [TypeBase Shape Diet]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase Shape Diet, ()) -> TypeBase Shape Diet
forall a b. (a, b) -> a
fst ([(TypeBase Shape Diet, ())] -> [TypeBase Shape Diet])
-> ([[Tree (TypeBase Shape Diet)]] -> [(TypeBase Shape Diet, ())])
-> [[Tree (TypeBase Shape Diet)]]
-> [TypeBase Shape Diet]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeBase shape o]
-> [(TypeBase Shape Diet, ())] -> [(TypeBase Shape Diet, ())]
forall shape1 o1 shape2 o2 b.
[TypeBase shape1 o1]
-> [(TypeBase shape2 o2, b)] -> [(TypeBase shape2 o2, b)]
fixupKnownTypes [TypeBase shape o]
ts ([(TypeBase Shape Diet, ())] -> [(TypeBase Shape Diet, ())])
-> ([[Tree (TypeBase Shape Diet)]] -> [(TypeBase Shape Diet, ())])
-> [[Tree (TypeBase Shape Diet)]]
-> [(TypeBase Shape Diet, ())]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase Shape Diet -> (TypeBase Shape Diet, ()))
-> [TypeBase Shape Diet] -> [(TypeBase Shape Diet, ())]
forall a b. (a -> b) -> [a] -> [b]
map (,()) ([TypeBase Shape Diet] -> [(TypeBase Shape Diet, ())])
-> ([[Tree (TypeBase Shape Diet)]] -> [TypeBase Shape Diet])
-> [[Tree (TypeBase Shape Diet)]]
-> [(TypeBase Shape Diet, ())]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Tree (TypeBase Shape Diet)] -> [TypeBase Shape Diet])
-> [[Tree (TypeBase Shape Diet)]] -> [TypeBase Shape Diet]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((Tree (TypeBase Shape Diet) -> [TypeBase Shape Diet])
-> [Tree (TypeBase Shape Diet)] -> [TypeBase Shape Diet]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree (TypeBase Shape Diet) -> [TypeBase Shape Diet]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList)
    ([[Tree (TypeBase Shape Diet)]] -> [TypeBase Shape Diet])
-> InternaliseM [[Tree (TypeBase Shape Diet)]]
-> InternaliseM [TypeBase Shape Diet]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [ParamType] -> InternaliseM [[Tree (TypeBase Shape Diet)]]
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 = [a b] -> [Int] -> [(a b, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a b]
xs ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (a b -> Int) -> [a b] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map a b -> Int
forall a. a a -> Int
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 = (State Int (Tree (a, Int)) -> Int -> Tree (a, Int))
-> Int -> State Int (Tree (a, Int)) -> Tree (a, Int)
forall a b c. (a -> b -> c) -> b -> a -> c
flip State Int (Tree (a, Int)) -> Int -> Tree (a, Int)
forall s a. State s a -> s -> a
evalState Int
o (State Int (Tree (a, Int)) -> Tree (a, Int))
-> (Tree a -> State Int (Tree (a, Int))) -> Tree a -> Tree (a, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree a -> State Int (Tree (a, Int))
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) = (b -> (Free f (a, b), b)) -> m (Free f (a, b))
forall a. (b -> (a, b)) -> m a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((b -> (Free f (a, b), b)) -> m (Free f (a, b)))
-> (b -> (Free f (a, b), b)) -> m (Free f (a, b))
forall a b. (a -> b) -> a -> b
$ \b
i -> ((a, b) -> Free f (a, b)
forall (f :: * -> *) a. a -> Free f a
Pure (a
x, b
i), b
i b -> b -> b
forall a. Num a => a -> a -> a
+ b
1)
    f (Free f (Free f a)
xs) = f (Free f (a, b)) -> Free f (a, b)
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free (f (Free f (a, b)) -> Free f (a, b))
-> m (f (Free f (a, b))) -> m (Free f (a, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Free f a -> m (Free f (a, b)))
-> f (Free f a) -> m (f (Free f (a, b)))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> f a -> f (f 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 = ((Tree a, Int) -> Tree (a, Int))
-> [(Tree a, Int)] -> [Tree (a, Int)]
forall a b. (a -> b) -> [a] -> [b]
map ((Tree a -> Int -> Tree (a, Int)) -> (Tree a, Int) -> Tree (a, Int)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Tree a -> Int -> Tree (a, Int))
 -> (Tree a, Int) -> Tree (a, Int))
-> (Tree a -> Int -> Tree (a, Int))
-> (Tree a, Int)
-> Tree (a, Int)
forall a b. (a -> b) -> a -> b
$ (Int -> Tree a -> Tree (a, Int)) -> Tree a -> Int -> Tree (a, Int)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Tree a -> Tree (a, Int)
forall a. Int -> Tree a -> Tree (a, Int)
numberFrom) ([(Tree a, Int)] -> [Tree (a, Int)])
-> ([Tree a] -> [(Tree a, Int)]) -> [Tree a] -> [Tree (a, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Tree a] -> [(Tree a, Int)]
forall (a :: * -> *) b. Foldable a => [a b] -> [(a b, Int)]
withOffsets

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) = Tree (a, b) -> Maybe (Tree (a, b))
forall a. a -> Maybe a
Just (Tree (a, b) -> Maybe (Tree (a, b)))
-> Tree (a, b) -> Maybe (Tree (a, b))
forall a b. (a -> b) -> a -> b
$ (a, b) -> Tree (a, b)
forall (f :: * -> *) a. a -> Free f a
Pure (a
a, b
b)
matchTrees (Free [Free [] a]
as) (Free [Free [] b]
bs)
  | [Free [] a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Free [] a]
as Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Free [] b] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Free [] b]
bs =
      [Tree (a, b)] -> Tree (a, b)
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free ([Tree (a, b)] -> Tree (a, b))
-> Maybe [Tree (a, b)] -> Maybe (Tree (a, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Free [] a -> Free [] b -> Maybe (Tree (a, b)))
-> [Free [] a] -> [Free [] b] -> Maybe [Tree (a, b)]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Free [] a -> Free [] b -> Maybe (Tree (a, b))
forall a b. Tree a -> Tree b -> Maybe (Tree (a, b))
matchTrees [Free [] a]
as [Free [] b]
bs
matchTrees Free [] a
_ Free [] b
_ = Maybe (Tree (a, 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 Tree a -> Tree b -> Maybe (Tree (a, b))
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' -> (Tree b -> [Tree (a, b)]) -> [Tree b] -> [Tree (a, b)]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Tree a -> Tree b -> [Tree (a, b)]
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 I.Diet)] ->
  [Tree (I.TypeBase ExtShape E.Freshness)] ->
  [[(I.TypeBase ExtShape E.Freshness, RetAls)]]
inferAliases :: [Tree (TypeBase Shape Diet)]
-> [Tree (TypeBase ExtShape Freshness)]
-> [[(TypeBase ExtShape Freshness, RetAls)]]
inferAliases [Tree (TypeBase Shape Diet)]
all_param_ts [Tree (TypeBase ExtShape Freshness)]
all_res_ts =
  (Tree (TypeBase ExtShape Freshness)
 -> [(TypeBase ExtShape Freshness, RetAls)])
-> [Tree (TypeBase ExtShape Freshness)]
-> [[(TypeBase ExtShape Freshness, RetAls)]]
forall a b. (a -> b) -> [a] -> [b]
map Tree (TypeBase ExtShape Freshness)
-> [(TypeBase ExtShape Freshness, RetAls)]
forall {shape}.
Free [] (TypeBase shape Freshness)
-> [(TypeBase shape Freshness, RetAls)]
onRes [Tree (TypeBase ExtShape Freshness)]
all_res_ts
  where
    all_res_ts' :: [Tree (TypeBase ExtShape Freshness, Int)]
all_res_ts' = [Tree (TypeBase ExtShape Freshness)]
-> [Tree (TypeBase ExtShape Freshness, Int)]
forall a. [Tree a] -> [Tree (a, Int)]
numberTrees [Tree (TypeBase ExtShape Freshness)]
all_res_ts
    all_param_ts' :: [Tree (TypeBase Shape Diet, Int)]
all_param_ts' = [Tree (TypeBase Shape Diet)] -> [Tree (TypeBase Shape Diet, Int)]
forall a. [Tree a] -> [Tree (a, Int)]
numberTrees [Tree (TypeBase Shape Diet)]
all_param_ts
    observed :: TypeBase shape Diet -> Bool
observed = (Diet -> Bool) -> TypeBase shape Diet -> Bool
forall o shape. (o -> Bool) -> TypeBase shape o -> Bool
aliasableArray (Diet -> Diet -> Bool
forall a. Eq a => a -> a -> Bool
== Diet
I.Observe)
    nonfresh :: TypeBase shape Freshness -> Bool
nonfresh = (Freshness -> Bool) -> TypeBase shape Freshness -> Bool
forall o shape. (o -> Bool) -> TypeBase shape o -> Bool
aliasableArray (Freshness -> Freshness -> Bool
forall a. Eq a => a -> a -> Bool
== Freshness
E.Nonfresh)
    aliasable_param_ts :: [Tree (TypeBase Shape Diet, Int)]
aliasable_param_ts = (Tree (TypeBase Shape Diet, Int) -> Bool)
-> [Tree (TypeBase Shape Diet, Int)]
-> [Tree (TypeBase Shape Diet, Int)]
forall a. (a -> Bool) -> [a] -> [a]
filter (((TypeBase Shape Diet, Int) -> Bool)
-> Tree (TypeBase Shape Diet, Int) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (((TypeBase Shape Diet, Int) -> Bool)
 -> Tree (TypeBase Shape Diet, Int) -> Bool)
-> ((TypeBase Shape Diet, Int) -> Bool)
-> Tree (TypeBase Shape Diet, Int)
-> Bool
forall a b. (a -> b) -> a -> b
$ TypeBase Shape Diet -> Bool
forall {shape}. TypeBase shape Diet -> Bool
observed (TypeBase Shape Diet -> Bool)
-> ((TypeBase Shape Diet, Int) -> TypeBase Shape Diet)
-> (TypeBase Shape Diet, Int)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase Shape Diet, Int) -> TypeBase Shape Diet
forall a b. (a, b) -> a
fst) [Tree (TypeBase Shape Diet, Int)]
all_param_ts'
    aliasable_res_ts :: [Tree (TypeBase ExtShape Freshness, Int)]
aliasable_res_ts = (Tree (TypeBase ExtShape Freshness, Int) -> Bool)
-> [Tree (TypeBase ExtShape Freshness, Int)]
-> [Tree (TypeBase ExtShape Freshness, Int)]
forall a. (a -> Bool) -> [a] -> [a]
filter (((TypeBase ExtShape Freshness, Int) -> Bool)
-> Tree (TypeBase ExtShape Freshness, Int) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (((TypeBase ExtShape Freshness, Int) -> Bool)
 -> Tree (TypeBase ExtShape Freshness, Int) -> Bool)
-> ((TypeBase ExtShape Freshness, Int) -> Bool)
-> Tree (TypeBase ExtShape Freshness, Int)
-> Bool
forall a b. (a -> b) -> a -> b
$ TypeBase ExtShape Freshness -> Bool
forall {shape}. TypeBase shape Freshness -> Bool
nonfresh (TypeBase ExtShape Freshness -> Bool)
-> ((TypeBase ExtShape Freshness, Int)
    -> TypeBase ExtShape Freshness)
-> (TypeBase ExtShape Freshness, Int)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase ExtShape Freshness, Int) -> TypeBase ExtShape Freshness
forall a b. (a, b) -> a
fst) [Tree (TypeBase ExtShape Freshness, Int)]
all_res_ts'
    onRes :: Free [] (TypeBase shape Freshness)
-> [(TypeBase shape Freshness, RetAls)]
onRes (Pure TypeBase shape Freshness
res_t) =
      -- Necessarily a non-array.
      [(TypeBase shape Freshness
res_t, [Int] -> [Int] -> RetAls
RetAls [Int]
forall a. Monoid a => a
mempty [Int]
forall a. Monoid a => a
mempty)]
    onRes (Free [Free [] (TypeBase shape Freshness)]
res_ts) =
      [ if TypeBase shape Freshness -> Bool
forall {shape}. TypeBase shape Freshness -> Bool
nonfresh TypeBase shape Freshness
res_t
          then (TypeBase shape Freshness
res_t, [Int] -> [Int] -> RetAls
RetAls [Int]
pals [Int]
rals)
          else (TypeBase shape Freshness
res_t, RetAls
forall a. Monoid a => a
mempty)
      | (TypeBase shape Freshness
res_t, [Int]
pals, [Int]
rals) <- [TypeBase shape Freshness]
-> [[Int]] -> [[Int]] -> [(TypeBase shape Freshness, [Int], [Int])]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (Free [] (TypeBase shape Freshness) -> [TypeBase shape Freshness]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList ([Free [] (TypeBase shape Freshness)]
-> Free [] (TypeBase shape Freshness)
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free [Free [] (TypeBase shape Freshness)]
res_ts)) [[Int]]
palss [[Int]]
ralss
      ]
      where
        reorder :: [[a]] -> [[a]]
reorder [] = Int -> [a] -> [[a]]
forall a. Int -> a -> [a]
replicate (Free [] (TypeBase shape Freshness) -> Int
forall a. Free [] a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Free [] (TypeBase shape Freshness)]
-> Free [] (TypeBase shape Freshness)
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free [Free [] (TypeBase shape Freshness)]
res_ts)) []
        reorder [[a]]
xs = [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
L.transpose [[a]]
xs
        infer :: t (Tree (a, a)) -> [[a]]
infer t (Tree (a, a))
ts =
          [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
reorder ([[a]] -> [[a]])
-> ([Free [] (TypeBase shape Freshness, (a, a))] -> [[a]])
-> [Free [] (TypeBase shape Freshness, (a, a))]
-> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Free [] (TypeBase shape Freshness, (a, a)) -> [a])
-> [Free [] (TypeBase shape Freshness, (a, a))] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (Free [] a -> [a]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Free [] a -> [a])
-> (Free [] (TypeBase shape Freshness, (a, a)) -> Free [] a)
-> Free [] (TypeBase shape Freshness, (a, a))
-> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((TypeBase shape Freshness, (a, a)) -> a)
-> Free [] (TypeBase shape Freshness, (a, a)) -> Free [] a
forall a b. (a -> b) -> Free [] a -> Free [] b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a, a) -> a
forall a b. (a, b) -> b
snd ((a, a) -> a)
-> ((TypeBase shape Freshness, (a, a)) -> (a, a))
-> (TypeBase shape Freshness, (a, a))
-> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase shape Freshness, (a, a)) -> (a, a)
forall a b. (a, b) -> b
snd)) ([Free [] (TypeBase shape Freshness, (a, a))] -> [[a]])
-> [Free [] (TypeBase shape Freshness, (a, a))] -> [[a]]
forall a b. (a -> b) -> a -> b
$
            (Tree (a, a) -> [Free [] (TypeBase shape Freshness, (a, a))])
-> t (Tree (a, a)) -> [Free [] (TypeBase shape Freshness, (a, a))]
forall m a. Monoid m => (a -> m) -> t a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Free [] (TypeBase shape Freshness)
-> Tree (a, a) -> [Free [] (TypeBase shape Freshness, (a, a))]
forall a b. Tree a -> Tree b -> [Tree (a, b)]
subtreesMatching ([Free [] (TypeBase shape Freshness)]
-> Free [] (TypeBase shape Freshness)
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free [Free [] (TypeBase shape Freshness)]
res_ts)) t (Tree (a, a))
ts
        palss :: [[Int]]
palss = [Tree (TypeBase Shape Diet, Int)] -> [[Int]]
forall {t :: * -> *} {a} {a}.
Foldable t =>
t (Tree (a, a)) -> [[a]]
infer [Tree (TypeBase Shape Diet, Int)]
aliasable_param_ts
        ralss :: [[Int]]
ralss = [Tree (TypeBase ExtShape Freshness, Int)] -> [[Int]]
forall {t :: * -> *} {a} {a}.
Foldable t =>
t (Tree (a, a)) -> [[a]]
infer [Tree (TypeBase ExtShape Freshness, Int)]
aliasable_res_ts

-- | The mode of the source-level return type is used to compute the
-- 'RetAls', and then discarded: an IR return type says nothing about
-- aliasing on its own.
internaliseReturnType ::
  [Tree (I.TypeBase Shape I.Diet)] ->
  E.ResRetType ->
  [(I.TypeBase ExtShape NoMode, RetAls)]
internaliseReturnType :: [Tree (TypeBase Shape Diet)]
-> ResRetType -> [(TypeBase ExtShape NoMode, RetAls)]
internaliseReturnType [Tree (TypeBase Shape Diet)]
paramts (E.RetType [VName]
dims TypeBase Size Freshness
et) =
  ((TypeBase ExtShape Freshness, RetAls)
 -> (TypeBase ExtShape NoMode, RetAls))
-> [(TypeBase ExtShape Freshness, RetAls)]
-> [(TypeBase ExtShape NoMode, RetAls)]
forall a b. (a -> b) -> [a] -> [b]
map ((TypeBase ExtShape Freshness -> TypeBase ExtShape NoMode)
-> (TypeBase ExtShape Freshness, RetAls)
-> (TypeBase ExtShape NoMode, RetAls)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first TypeBase ExtShape Freshness -> TypeBase ExtShape NoMode
forall shape o. TypeBase shape o -> TypeBase shape NoMode
I.fromDecl) ([(TypeBase ExtShape Freshness, RetAls)]
 -> [(TypeBase ExtShape NoMode, RetAls)])
-> ([Tree (TypeBase ExtShape Freshness)]
    -> [(TypeBase ExtShape Freshness, RetAls)])
-> [Tree (TypeBase ExtShape Freshness)]
-> [(TypeBase ExtShape NoMode, RetAls)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(TypeBase ExtShape Freshness, RetAls)]
-> [(TypeBase ExtShape Freshness, RetAls)]
forall {shape2} {u2} {b}.
[(TypeBase shape2 u2, b)] -> [(TypeBase shape2 u2, b)]
fixupAccs ([(TypeBase ExtShape Freshness, RetAls)]
 -> [(TypeBase ExtShape Freshness, RetAls)])
-> ([Tree (TypeBase ExtShape Freshness)]
    -> [(TypeBase ExtShape Freshness, RetAls)])
-> [Tree (TypeBase ExtShape Freshness)]
-> [(TypeBase ExtShape Freshness, RetAls)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[(TypeBase ExtShape Freshness, RetAls)]]
-> [(TypeBase ExtShape Freshness, RetAls)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(TypeBase ExtShape Freshness, RetAls)]]
 -> [(TypeBase ExtShape Freshness, RetAls)])
-> ([Tree (TypeBase ExtShape Freshness)]
    -> [[(TypeBase ExtShape Freshness, RetAls)]])
-> [Tree (TypeBase ExtShape Freshness)]
-> [(TypeBase ExtShape Freshness, RetAls)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Tree (TypeBase Shape Diet)]
-> [Tree (TypeBase ExtShape Freshness)]
-> [[(TypeBase ExtShape Freshness, RetAls)]]
inferAliases [Tree (TypeBase Shape Diet)]
paramts ([Tree (TypeBase ExtShape Freshness)]
 -> [(TypeBase ExtShape NoMode, RetAls)])
-> [Tree (TypeBase ExtShape Freshness)]
-> [(TypeBase ExtShape NoMode, RetAls)]
forall a b. (a -> b) -> a -> b
$
    [VName]
-> InternaliseTypeM [Tree (TypeBase ExtShape Freshness)]
-> [Tree (TypeBase ExtShape Freshness)]
forall a. [VName] -> InternaliseTypeM a -> a
runInternaliseTypeM' [VName]
dims (Map VName Int
-> TypeBase Size Freshness
-> InternaliseTypeM [Tree (TypeBase ExtShape Freshness)]
forall o.
Pretty o =>
Map VName Int
-> TypeBase Size o -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
internaliseTypeM Map VName Int
exts TypeBase Size Freshness
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 ..]
    -- Any 'Acc' in the return type must (by the type rules) be the function's
    -- single 'Acc' parameter, so we substitute its known accumulator type.
    fixupAccs :: [(TypeBase shape2 u2, b)] -> [(TypeBase shape2 u2, b)]
fixupAccs = case [TypeBase Shape Diet
t | t :: TypeBase Shape Diet
t@Acc {} <- (Tree (TypeBase Shape Diet) -> [TypeBase Shape Diet])
-> [Tree (TypeBase Shape Diet)] -> [TypeBase Shape Diet]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (TypeBase Shape Diet) -> [TypeBase Shape Diet]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [Tree (TypeBase Shape Diet)]
paramts] of
      TypeBase Shape Diet
acc : [TypeBase Shape Diet]
_ -> ((TypeBase shape2 u2, b) -> (TypeBase shape2 u2, b))
-> [(TypeBase shape2 u2, b)] -> [(TypeBase shape2 u2, b)]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase Shape Diet
-> (TypeBase shape2 u2, b) -> (TypeBase shape2 u2, b)
forall shape1 u1 shape2 u2 b.
TypeBase shape1 u1
-> (TypeBase shape2 u2, b) -> (TypeBase shape2 u2, b)
fixupAcc TypeBase Shape Diet
acc)
      [] -> [(TypeBase shape2 u2, b)] -> [(TypeBase shape2 u2, b)]
forall a. a -> a
id

-- | As 'internaliseReturnType', but returns components of a top-level
-- tuple type piecemeal.
internaliseEntryReturnType ::
  [Tree (I.TypeBase Shape I.Diet)] ->
  E.ResRetType ->
  [[(I.TypeBase ExtShape E.Freshness, RetAls)]]
internaliseEntryReturnType :: [Tree (TypeBase Shape Diet)]
-> ResRetType -> [[(TypeBase ExtShape Freshness, RetAls)]]
internaliseEntryReturnType [Tree (TypeBase Shape Diet)]
paramts (E.RetType [VName]
dims TypeBase Size Freshness
et) =
  let et' :: [[Tree (TypeBase ExtShape Freshness)]]
et' = [VName]
-> InternaliseTypeM [[Tree (TypeBase ExtShape Freshness)]]
-> [[Tree (TypeBase ExtShape Freshness)]]
forall a. [VName] -> InternaliseTypeM a -> a
runInternaliseTypeM' [VName]
dims (InternaliseTypeM [[Tree (TypeBase ExtShape Freshness)]]
 -> [[Tree (TypeBase ExtShape Freshness)]])
-> ([TypeBase Size Freshness]
    -> InternaliseTypeM [[Tree (TypeBase ExtShape Freshness)]])
-> [TypeBase Size Freshness]
-> [[Tree (TypeBase ExtShape Freshness)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase Size Freshness
 -> InternaliseTypeM [Tree (TypeBase ExtShape Freshness)])
-> [TypeBase Size Freshness]
-> InternaliseTypeM [[Tree (TypeBase ExtShape Freshness)]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Map VName Int
-> TypeBase Size Freshness
-> InternaliseTypeM [Tree (TypeBase ExtShape Freshness)]
forall o.
Pretty o =>
Map VName Int
-> TypeBase Size o -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
internaliseTypeM Map VName Int
exts) ([TypeBase Size Freshness]
 -> [[Tree (TypeBase ExtShape Freshness)]])
-> [TypeBase Size Freshness]
-> [[Tree (TypeBase ExtShape Freshness)]]
forall a b. (a -> b) -> a -> b
$
        case TypeBase Size Freshness -> Maybe [TypeBase Size Freshness]
forall dim o. TypeBase dim o -> Maybe [TypeBase dim o]
E.isTupleRecord TypeBase Size Freshness
et of
          Just [TypeBase Size Freshness]
ets | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [TypeBase Size Freshness] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeBase Size Freshness]
ets -> [TypeBase Size Freshness]
ets
          Maybe [TypeBase Size Freshness]
_ -> [TypeBase Size Freshness
et]
   in ([[(TypeBase ExtShape Freshness, RetAls)]]
 -> [(TypeBase ExtShape Freshness, RetAls)])
-> [[[(TypeBase ExtShape Freshness, RetAls)]]]
-> [[(TypeBase ExtShape Freshness, RetAls)]]
forall a b. (a -> b) -> [a] -> [b]
map [[(TypeBase ExtShape Freshness, RetAls)]]
-> [(TypeBase ExtShape Freshness, RetAls)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[[(TypeBase ExtShape Freshness, RetAls)]]]
 -> [[(TypeBase ExtShape Freshness, RetAls)]])
-> [[[(TypeBase ExtShape Freshness, RetAls)]]]
-> [[(TypeBase ExtShape Freshness, RetAls)]]
forall a b. (a -> b) -> a -> b
$ [[Tree (TypeBase ExtShape Freshness)]]
-> [[(TypeBase ExtShape Freshness, RetAls)]]
-> [[[(TypeBase ExtShape Freshness, RetAls)]]]
forall a b. [[a]] -> [b] -> [[b]]
chunkLike [[Tree (TypeBase ExtShape Freshness)]]
et' ([[(TypeBase ExtShape Freshness, RetAls)]]
 -> [[[(TypeBase ExtShape Freshness, RetAls)]]])
-> [[(TypeBase ExtShape Freshness, RetAls)]]
-> [[[(TypeBase ExtShape Freshness, RetAls)]]]
forall a b. (a -> b) -> a -> b
$ [Tree (TypeBase Shape Diet)]
-> [Tree (TypeBase ExtShape Freshness)]
-> [[(TypeBase ExtShape Freshness, RetAls)]]
inferAliases [Tree (TypeBase Shape Diet)]
paramts ([Tree (TypeBase ExtShape Freshness)]
 -> [[(TypeBase ExtShape Freshness, RetAls)]])
-> [Tree (TypeBase ExtShape Freshness)]
-> [[(TypeBase ExtShape Freshness, RetAls)]]
forall a b. (a -> b) -> a -> b
$ [[Tree (TypeBase ExtShape Freshness)]]
-> [Tree (TypeBase ExtShape Freshness)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Tree (TypeBase ExtShape Freshness)]]
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 ..]

internaliseCoerceType ::
  E.StructType ->
  [I.TypeBase ExtShape NoMode]
internaliseCoerceType :: StructType -> [TypeBase ExtShape NoMode]
internaliseCoerceType StructType
et =
  ((TypeBase ExtShape NoMode, RetAls) -> TypeBase ExtShape NoMode)
-> [(TypeBase ExtShape NoMode, RetAls)]
-> [TypeBase ExtShape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase ExtShape NoMode, RetAls) -> TypeBase ExtShape NoMode
forall a b. (a, b) -> a
fst ([(TypeBase ExtShape NoMode, RetAls)]
 -> [TypeBase ExtShape NoMode])
-> [(TypeBase ExtShape NoMode, RetAls)]
-> [TypeBase ExtShape NoMode]
forall a b. (a -> b) -> a -> b
$ [Tree (TypeBase Shape Diet)]
-> ResRetType -> [(TypeBase ExtShape NoMode, RetAls)]
internaliseReturnType [] ([VName] -> TypeBase Size Freshness -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
E.RetType [] (TypeBase Size Freshness -> ResRetType)
-> TypeBase Size Freshness -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> TypeBase Size Freshness
forall o. Freshness -> TypeBase Size o -> TypeBase Size Freshness
E.toRes Freshness
E.Nonfresh StructType
et)

internaliseLambdaReturnType ::
  E.ResType ->
  [TypeBase shape o] ->
  InternaliseM [I.TypeBase Shape NoMode]
internaliseLambdaReturnType :: forall shape o.
TypeBase Size Freshness
-> [TypeBase shape o] -> InternaliseM [Type]
internaliseLambdaReturnType TypeBase Size Freshness
et [TypeBase shape o]
ts =
  (TypeBase Shape Diet -> Type) -> [TypeBase Shape Diet] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape Diet -> Type
forall shape o. TypeBase shape o -> TypeBase shape NoMode
fromDecl ([TypeBase Shape Diet] -> [Type])
-> InternaliseM [TypeBase Shape Diet] -> InternaliseM [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ParamType
-> [TypeBase shape o] -> InternaliseM [TypeBase Shape Diet]
forall shape o.
ParamType
-> [TypeBase shape o] -> InternaliseM [TypeBase Shape Diet]
internaliseLoopParamType (TypeBase Size Freshness -> ParamType
E.resToParam TypeBase Size Freshness
et) [TypeBase shape o]
ts

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

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

internaliseDim ::
  M.Map VName Int ->
  E.Size ->
  InternaliseTypeM ExtSize
internaliseDim :: Map VName Int -> Size -> InternaliseTypeM ExtSize
internaliseDim Map VName Int
exts Size
d =
  case Size
d of
    Size
e | Just Int
_ <- Size -> Maybe Int
E.isAnySize Size
e -> Int -> ExtSize
forall a. Int -> Ext a
Ext (Int -> ExtSize)
-> InternaliseTypeM Int -> InternaliseTypeM ExtSize
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> InternaliseTypeM Int
newId
    (E.IntLit Integer
n Info StructType
_ SrcLoc
_) -> ExtSize -> InternaliseTypeM ExtSize
forall a. a -> InternaliseTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExtSize -> InternaliseTypeM ExtSize)
-> ExtSize -> InternaliseTypeM ExtSize
forall a b. (a -> b) -> a -> b
$ SubExp -> ExtSize
forall a. a -> Ext a
I.Free (SubExp -> ExtSize) -> SubExp -> ExtSize
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
I.Int64 Integer
n
    (E.Var QualName VName
name Info StructType
_ SrcLoc
_) -> ExtSize -> InternaliseTypeM ExtSize
forall a. a -> InternaliseTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExtSize -> InternaliseTypeM ExtSize)
-> ExtSize -> InternaliseTypeM ExtSize
forall a b. (a -> b) -> a -> b
$ QualName VName -> ExtSize
namedDim QualName VName
name
    Size
e -> [Char] -> InternaliseTypeM ExtSize
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseTypeM ExtSize)
-> [Char] -> InternaliseTypeM ExtSize
forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpected size expression: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Size -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Size
e
  where
    namedDim :: QualName VName -> ExtSize
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 -> ExtSize
forall a. Int -> Ext a
I.Ext Int
x
      | Bool
otherwise = SubExp -> ExtSize
forall a. a -> Ext a
I.Free (SubExp -> ExtSize) -> SubExp -> ExtSize
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 []

-- | Internalise a source type, preserving whatever mode it carries: a
-- 'E.Diet' for a parameter type, a 'E.Freshness' for a return type.
-- Which of the two it is matters to 'inferAliases', so it is not
-- collapsed here.
internaliseTypeM ::
  (Pretty o) =>
  M.Map VName Int ->
  E.TypeBase E.Size o ->
  InternaliseTypeM [Tree (I.TypeBase ExtShape o)]
internaliseTypeM :: forall o.
Pretty o =>
Map VName Int
-> TypeBase Size o -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
internaliseTypeM Map VName Int
exts TypeBase Size o
orig_t =
  case TypeBase Size o
orig_t of
    E.Array o
o Shape Size
shape ScalarTypeBase Size NoMode
et -> do
      dims <- Shape Size -> InternaliseTypeM [ExtSize]
internaliseShape Shape Size
shape
      ets <- internaliseTypeM exts $ E.Scalar et
      let f TypeBase ExtShape o_unused
et' = TypeBase ExtShape o_unused -> ExtShape -> o -> TypeBase ExtShape o
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
I.arrayOf TypeBase ExtShape o_unused
et' ([ExtSize] -> ExtShape
forall d. [d] -> ShapeBase d
Shape [ExtSize]
dims) o
o
      pure [array $ map (fmap f) ets]
    E.Scalar (E.Prim PrimType
bt) ->
      [Tree (TypeBase ExtShape o)]
-> InternaliseTypeM [Tree (TypeBase ExtShape o)]
forall a. a -> InternaliseTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [TypeBase ExtShape o -> Tree (TypeBase ExtShape o)
forall (f :: * -> *) a. a -> Free f a
Pure (TypeBase ExtShape o -> Tree (TypeBase ExtShape o))
-> TypeBase ExtShape o -> Tree (TypeBase ExtShape o)
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase ExtShape o
forall shape o. PrimType -> TypeBase shape o
I.Prim (PrimType -> TypeBase ExtShape o)
-> PrimType -> TypeBase ExtShape o
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
internalisePrimType PrimType
bt]
    E.Scalar (E.Record Map Name (TypeBase Size o)
ets)
      -- We map empty records to units, because otherwise arrays of
      -- unit will lose their sizes.
      | Map Name (TypeBase Size o) -> Bool
forall a. Map Name a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Map Name (TypeBase Size o)
ets -> [Tree (TypeBase ExtShape o)]
-> InternaliseTypeM [Tree (TypeBase ExtShape o)]
forall a. a -> InternaliseTypeM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [TypeBase ExtShape o -> Tree (TypeBase ExtShape o)
forall (f :: * -> *) a. a -> Free f a
Pure (TypeBase ExtShape o -> Tree (TypeBase ExtShape o))
-> TypeBase ExtShape o -> Tree (TypeBase ExtShape o)
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase ExtShape o
forall shape o. PrimType -> TypeBase shape o
I.Prim PrimType
I.Unit]
      | Bool
otherwise ->
          [[Tree (TypeBase ExtShape o)]] -> [Tree (TypeBase ExtShape o)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Tree (TypeBase ExtShape o)]] -> [Tree (TypeBase ExtShape o)])
-> InternaliseTypeM [[Tree (TypeBase ExtShape o)]]
-> InternaliseTypeM [Tree (TypeBase ExtShape o)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Name, TypeBase Size o)
 -> InternaliseTypeM [Tree (TypeBase ExtShape o)])
-> [(Name, TypeBase Size o)]
-> InternaliseTypeM [[Tree (TypeBase ExtShape o)]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Map VName Int
-> TypeBase Size o -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
forall o.
Pretty o =>
Map VName Int
-> TypeBase Size o -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
internaliseTypeM Map VName Int
exts (TypeBase Size o -> InternaliseTypeM [Tree (TypeBase ExtShape o)])
-> ((Name, TypeBase Size o) -> TypeBase Size o)
-> (Name, TypeBase Size o)
-> InternaliseTypeM [Tree (TypeBase ExtShape o)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, TypeBase Size o) -> TypeBase Size o
forall a b. (a, b) -> b
snd) (Map Name (TypeBase Size o) -> [(Name, TypeBase Size o)]
forall a. Map Name a -> [(Name, a)]
E.sortFields Map Name (TypeBase Size o)
ets)
    E.Scalar (E.TypeVar o
_ QualName VName
tn [E.TypeArgType StructType
arr_t])
      | VName -> Bool
E.isIntrinsic (QualName VName -> VName
forall vn. QualName vn -> vn
E.qualLeaf QualName VName
tn),
        VName -> Name
baseName (QualName VName -> VName
forall vn. QualName vn -> vn
E.qualLeaf QualName VName
tn) Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
"acc" -> do
          ts <-
            (Tree (TypeBase ExtShape NoMode) -> [Type])
-> [Tree (TypeBase ExtShape NoMode)] -> [Type]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Free [] Type -> [Type]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Free [] Type -> [Type])
-> (Tree (TypeBase ExtShape NoMode) -> Free [] Type)
-> Tree (TypeBase ExtShape NoMode)
-> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase ExtShape NoMode -> Type)
-> Tree (TypeBase ExtShape NoMode) -> Free [] Type
forall a b. (a -> b) -> Free [] a -> Free [] b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase ExtShape NoMode -> Type
forall {o}. TypeBase ExtShape o -> TypeBase Shape o
onAccType)
              ([Tree (TypeBase ExtShape NoMode)] -> [Type])
-> InternaliseTypeM [Tree (TypeBase ExtShape NoMode)]
-> InternaliseTypeM [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map VName Int
-> StructType -> InternaliseTypeM [Tree (TypeBase ExtShape NoMode)]
forall o.
Pretty o =>
Map VName Int
-> TypeBase Size o -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
internaliseTypeM Map VName Int
exts StructType
arr_t
          let acc_param = Name -> Int -> VName
VName Name
"PLACEHOLDER" Int
0 -- See mkAccCerts.
              acc_shape = [SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [Int -> [Type] -> SubExp
forall o. Int -> [TypeBase Shape o] -> SubExp
arraysSize Int
0 [Type]
ts]
              acc_t = VName -> Shape -> [Type] -> TypeBase shape o
forall shape o. VName -> Shape -> [Type] -> TypeBase shape o
Acc VName
acc_param Shape
acc_shape ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType [Type]
ts)
          pure [Pure acc_t]
    E.Scalar E.TypeVar {} ->
      [Char] -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseTypeM [Tree (TypeBase ExtShape o)])
-> [Char] -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseTypeM: cannot handle type variable: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ TypeBase Size o -> [Char]
forall a. Pretty a => a -> [Char]
prettyString TypeBase Size o
orig_t
    E.Scalar E.Arrow {} ->
      [Char] -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseTypeM [Tree (TypeBase ExtShape o)])
-> [Char] -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseTypeM: cannot handle function type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ TypeBase Size o -> [Char]
forall a. Pretty a => a -> [Char]
prettyString TypeBase Size o
orig_t
    E.Scalar (E.Sum Map Name [TypeBase Size o]
cs) -> do
      (ts, _) <-
        Map Name [Tree (TypeBase ExtShape o)]
-> ([Tree (TypeBase ExtShape o)], [(Name, [Int])])
forall o.
Map Name [Tree (TypeBase ExtShape o)]
-> ([Tree (TypeBase ExtShape o)], [(Name, [Int])])
internaliseConstructors
          (Map Name [Tree (TypeBase ExtShape o)]
 -> ([Tree (TypeBase ExtShape o)], [(Name, [Int])]))
-> InternaliseTypeM (Map Name [Tree (TypeBase ExtShape o)])
-> InternaliseTypeM ([Tree (TypeBase ExtShape o)], [(Name, [Int])])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([TypeBase Size o]
 -> InternaliseTypeM [Tree (TypeBase ExtShape o)])
-> Map Name [TypeBase Size o]
-> InternaliseTypeM (Map Name [Tree (TypeBase ExtShape o)])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse (([[Tree (TypeBase ExtShape o)]] -> [Tree (TypeBase ExtShape o)])
-> InternaliseTypeM [[Tree (TypeBase ExtShape o)]]
-> InternaliseTypeM [Tree (TypeBase ExtShape o)]
forall a b. (a -> b) -> InternaliseTypeM a -> InternaliseTypeM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Tree (TypeBase ExtShape o)]] -> [Tree (TypeBase ExtShape o)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (InternaliseTypeM [[Tree (TypeBase ExtShape o)]]
 -> InternaliseTypeM [Tree (TypeBase ExtShape o)])
-> ([TypeBase Size o]
    -> InternaliseTypeM [[Tree (TypeBase ExtShape o)]])
-> [TypeBase Size o]
-> InternaliseTypeM [Tree (TypeBase ExtShape o)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase Size o -> InternaliseTypeM [Tree (TypeBase ExtShape o)])
-> [TypeBase Size o]
-> InternaliseTypeM [[Tree (TypeBase ExtShape o)]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Map VName Int
-> TypeBase Size o -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
forall o.
Pretty o =>
Map VName Int
-> TypeBase Size o -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
internaliseTypeM Map VName Int
exts)) Map Name [TypeBase Size o]
cs
      pure $
        if length cs == 1
          then ts
          else Pure (I.Prim (I.IntType I.Int8)) : ts
  where
    internaliseShape :: Shape Size -> InternaliseTypeM [ExtSize]
internaliseShape = (Size -> InternaliseTypeM ExtSize)
-> [Size] -> InternaliseTypeM [ExtSize]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Map VName Int -> Size -> InternaliseTypeM ExtSize
internaliseDim Map VName Int
exts) ([Size] -> InternaliseTypeM [ExtSize])
-> (Shape Size -> [Size])
-> Shape Size
-> InternaliseTypeM [ExtSize]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape Size -> [Size]
forall dim. Shape dim -> [dim]
E.shapeDims
    array :: [Free [] a] -> Free [] a
array [Free [Free [] a]
ts] = [Free [] a] -> Free [] a
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free [Free [] a]
ts
    array [Free [] a]
ts = [Free [] a] -> Free [] a
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free [Free [] a]
ts

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

-- | Only exposed for testing purposes.
internaliseConstructors ::
  M.Map Name [Tree (I.TypeBase ExtShape o)] ->
  ( [Tree (I.TypeBase ExtShape o)],
    [(Name, [Int])]
  )
internaliseConstructors :: forall o.
Map Name [Tree (TypeBase ExtShape o)]
-> ([Tree (TypeBase ExtShape o)], [(Name, [Int])])
internaliseConstructors Map Name [Tree (TypeBase ExtShape o)]
cs =
  ([Tree (TypeBase ExtShape o)]
 -> (Name, [Tree (TypeBase ExtShape o)])
 -> ([Tree (TypeBase ExtShape o)], (Name, [Int])))
-> [Tree (TypeBase ExtShape o)]
-> [(Name, [Tree (TypeBase ExtShape o)])]
-> ([Tree (TypeBase ExtShape o)], [(Name, [Int])])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
L.mapAccumL [Tree (TypeBase ExtShape o)]
-> (Name, [Tree (TypeBase ExtShape o)])
-> ([Tree (TypeBase ExtShape o)], (Name, [Int]))
forall {t :: * -> *} {shape} {t :: * -> *} {o} {a}.
(Eq (t (TypeBase shape NoMode)), Foldable t, Foldable t,
 Functor t) =>
[t (TypeBase shape o)]
-> (a, t (t (TypeBase shape o)))
-> ([t (TypeBase shape o)], (a, [Int]))
onConstructor [Tree (TypeBase ExtShape o)]
forall a. Monoid a => a
mempty ([(Name, [Tree (TypeBase ExtShape o)])]
 -> ([Tree (TypeBase ExtShape o)], [(Name, [Int])]))
-> [(Name, [Tree (TypeBase ExtShape o)])]
-> ([Tree (TypeBase ExtShape o)], [(Name, [Int])])
forall a b. (a -> b) -> a -> b
$ Map Name [Tree (TypeBase ExtShape o)]
-> [(Name, [Tree (TypeBase ExtShape o)])]
forall a. Map Name a -> [(Name, a)]
E.sortConstrs Map Name [Tree (TypeBase ExtShape o)]
cs
  where
    unmoded :: TypeBase shape o -> TypeBase shape NoMode
unmoded = TypeBase shape o -> TypeBase shape NoMode
forall shape o. TypeBase shape o -> TypeBase shape NoMode
fromDecl
    onConstructor :: [t (TypeBase shape o)]
-> (a, t (t (TypeBase shape o)))
-> ([t (TypeBase shape o)], (a, [Int]))
onConstructor [t (TypeBase shape o)]
ts (a
c, t (t (TypeBase shape o))
c_ts) =
      let ([(t (TypeBase shape NoMode), Int)]
_, [Int]
js, [t (TypeBase shape o)]
new_ts) =
            (([(t (TypeBase shape NoMode), Int)], [Int],
  [t (TypeBase shape o)])
 -> t (TypeBase shape o)
 -> ([(t (TypeBase shape NoMode), Int)], [Int],
     [t (TypeBase shape o)]))
-> ([(t (TypeBase shape NoMode), Int)], [Int],
    [t (TypeBase shape o)])
-> t (t (TypeBase shape o))
-> ([(t (TypeBase shape NoMode), Int)], [Int],
    [t (TypeBase shape o)])
forall b a. (b -> a -> b) -> b -> t a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ([(t (TypeBase shape NoMode), Int)], [Int], [t (TypeBase shape o)])
-> t (TypeBase shape o)
-> ([(t (TypeBase shape NoMode), Int)], [Int],
    [t (TypeBase shape o)])
forall {shape} {o}.
Eq (t (TypeBase shape NoMode)) =>
([(t (TypeBase shape NoMode), Int)], [Int], [t (TypeBase shape o)])
-> t (TypeBase shape o)
-> ([(t (TypeBase shape NoMode), Int)], [Int],
    [t (TypeBase shape o)])
f ([t (TypeBase shape NoMode)] -> [(t (TypeBase shape NoMode), Int)]
forall (a :: * -> *) b. Foldable a => [a b] -> [(a b, Int)]
withOffsets ((t (TypeBase shape o) -> t (TypeBase shape NoMode))
-> [t (TypeBase shape o)] -> [t (TypeBase shape NoMode)]
forall a b. (a -> b) -> [a] -> [b]
map ((TypeBase shape o -> TypeBase shape NoMode)
-> t (TypeBase shape o) -> t (TypeBase shape NoMode)
forall a b. (a -> b) -> t a -> t b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase shape o -> TypeBase shape NoMode
forall shape o. TypeBase shape o -> TypeBase shape NoMode
unmoded) [t (TypeBase shape o)]
ts), [Int]
forall a. Monoid a => a
mempty, [t (TypeBase shape o)]
forall a. Monoid a => a
mempty) t (t (TypeBase shape o))
c_ts
       in ([t (TypeBase shape o)]
ts [t (TypeBase shape o)]
-> [t (TypeBase shape o)] -> [t (TypeBase shape o)]
forall a. [a] -> [a] -> [a]
++ [t (TypeBase shape o)]
new_ts, (a
c, [Int]
js))
      where
        size :: [t a] -> Int
size = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> ([t a] -> [Int]) -> [t a] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (t a -> Int) -> [t a] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map t a -> Int
forall a. t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length
        f :: ([(t (TypeBase shape NoMode), Int)], [Int], [t (TypeBase shape o)])
-> t (TypeBase shape o)
-> ([(t (TypeBase shape NoMode), Int)], [Int],
    [t (TypeBase shape o)])
f ([(t (TypeBase shape NoMode), Int)]
ts', [Int]
js, [t (TypeBase shape o)]
new_ts) t (TypeBase shape o)
t
          | (TypeBase shape o -> Bool) -> t (TypeBase shape o) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all TypeBase shape o -> Bool
forall shape o. TypeBase shape o -> Bool
primType t (TypeBase shape o)
t,
            Just (t (TypeBase shape NoMode)
_, Int
j) <- ((t (TypeBase shape NoMode), Int) -> Bool)
-> [(t (TypeBase shape NoMode), Int)]
-> Maybe (t (TypeBase shape NoMode), Int)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((t (TypeBase shape NoMode) -> t (TypeBase shape NoMode) -> Bool
forall a. Eq a => a -> a -> Bool
== (TypeBase shape o -> TypeBase shape NoMode)
-> t (TypeBase shape o) -> t (TypeBase shape NoMode)
forall a b. (a -> b) -> t a -> t b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase shape o -> TypeBase shape NoMode
forall shape o. TypeBase shape o -> TypeBase shape NoMode
unmoded t (TypeBase shape o)
t) (t (TypeBase shape NoMode) -> Bool)
-> ((t (TypeBase shape NoMode), Int) -> t (TypeBase shape NoMode))
-> (t (TypeBase shape NoMode), Int)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (t (TypeBase shape NoMode), Int) -> t (TypeBase shape NoMode)
forall a b. (a, b) -> a
fst) [(t (TypeBase shape NoMode), Int)]
ts' =
              ( (t (TypeBase shape NoMode), Int)
-> [(t (TypeBase shape NoMode), Int)]
-> [(t (TypeBase shape NoMode), Int)]
forall a. Eq a => a -> [a] -> [a]
delete ((TypeBase shape o -> TypeBase shape NoMode)
-> t (TypeBase shape o) -> t (TypeBase shape NoMode)
forall a b. (a -> b) -> t a -> t b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase shape o -> TypeBase shape NoMode
forall shape o. TypeBase shape o -> TypeBase shape NoMode
unmoded t (TypeBase shape o)
t, Int
j) [(t (TypeBase shape NoMode), Int)]
ts',
                [Int]
js [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take (t (TypeBase shape o) -> Int
forall a. t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t (TypeBase shape o)
t) [Int
j ..],
                [t (TypeBase shape o)]
new_ts
              )
          | Bool
otherwise =
              ( [(t (TypeBase shape NoMode), Int)]
ts',
                [Int]
js [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take (t (TypeBase shape o) -> Int
forall a. t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t (TypeBase shape o)
t) [[t (TypeBase shape o)] -> Int
forall {a}. [t a] -> Int
size [t (TypeBase shape o)]
ts Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [t (TypeBase shape o)] -> Int
forall {a}. [t a] -> Int
size [t (TypeBase shape o)]
new_ts ..],
                [t (TypeBase shape o)]
new_ts [t (TypeBase shape o)]
-> [t (TypeBase shape o)] -> [t (TypeBase shape o)]
forall a. [a] -> [a] -> [a]
++ [t (TypeBase shape o)
t]
              )

internaliseSumTypeRep ::
  M.Map Name [E.StructType] ->
  ( [I.TypeBase ExtShape NoMode],
    [(Name, [Int])]
  )
internaliseSumTypeRep :: Map Name [StructType]
-> ([TypeBase ExtShape NoMode], [(Name, [Int])])
internaliseSumTypeRep Map Name [StructType]
cs =
  ([Tree (TypeBase ExtShape NoMode)] -> [TypeBase ExtShape NoMode])
-> ([Tree (TypeBase ExtShape NoMode)], [(Name, [Int])])
-> ([TypeBase ExtShape NoMode], [(Name, [Int])])
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((Tree (TypeBase ExtShape NoMode) -> [TypeBase ExtShape NoMode])
-> [Tree (TypeBase ExtShape NoMode)] -> [TypeBase ExtShape NoMode]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (TypeBase ExtShape NoMode) -> [TypeBase ExtShape NoMode]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) (([Tree (TypeBase ExtShape NoMode)], [(Name, [Int])])
 -> ([TypeBase ExtShape NoMode], [(Name, [Int])]))
-> (InternaliseTypeM
      ([Tree (TypeBase ExtShape NoMode)], [(Name, [Int])])
    -> ([Tree (TypeBase ExtShape NoMode)], [(Name, [Int])]))
-> InternaliseTypeM
     ([Tree (TypeBase ExtShape NoMode)], [(Name, [Int])])
-> ([TypeBase ExtShape NoMode], [(Name, [Int])])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternaliseTypeM
  ([Tree (TypeBase ExtShape NoMode)], [(Name, [Int])])
-> ([Tree (TypeBase ExtShape NoMode)], [(Name, [Int])])
forall a. InternaliseTypeM a -> a
runInternaliseTypeM (InternaliseTypeM
   ([Tree (TypeBase ExtShape NoMode)], [(Name, [Int])])
 -> ([TypeBase ExtShape NoMode], [(Name, [Int])]))
-> InternaliseTypeM
     ([Tree (TypeBase ExtShape NoMode)], [(Name, [Int])])
-> ([TypeBase ExtShape NoMode], [(Name, [Int])])
forall a b. (a -> b) -> a -> b
$
    Map Name [Tree (TypeBase ExtShape NoMode)]
-> ([Tree (TypeBase ExtShape NoMode)], [(Name, [Int])])
forall o.
Map Name [Tree (TypeBase ExtShape o)]
-> ([Tree (TypeBase ExtShape o)], [(Name, [Int])])
internaliseConstructors
      (Map Name [Tree (TypeBase ExtShape NoMode)]
 -> ([Tree (TypeBase ExtShape NoMode)], [(Name, [Int])]))
-> InternaliseTypeM (Map Name [Tree (TypeBase ExtShape NoMode)])
-> InternaliseTypeM
     ([Tree (TypeBase ExtShape NoMode)], [(Name, [Int])])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([StructType]
 -> InternaliseTypeM [Tree (TypeBase ExtShape NoMode)])
-> Map Name [StructType]
-> InternaliseTypeM (Map Name [Tree (TypeBase ExtShape NoMode)])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse (([[Tree (TypeBase ExtShape NoMode)]]
 -> [Tree (TypeBase ExtShape NoMode)])
-> InternaliseTypeM [[Tree (TypeBase ExtShape NoMode)]]
-> InternaliseTypeM [Tree (TypeBase ExtShape NoMode)]
forall a b. (a -> b) -> InternaliseTypeM a -> InternaliseTypeM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Tree (TypeBase ExtShape NoMode)]]
-> [Tree (TypeBase ExtShape NoMode)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (InternaliseTypeM [[Tree (TypeBase ExtShape NoMode)]]
 -> InternaliseTypeM [Tree (TypeBase ExtShape NoMode)])
-> ([StructType]
    -> InternaliseTypeM [[Tree (TypeBase ExtShape NoMode)]])
-> [StructType]
-> InternaliseTypeM [Tree (TypeBase ExtShape NoMode)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StructType -> InternaliseTypeM [Tree (TypeBase ExtShape NoMode)])
-> [StructType]
-> InternaliseTypeM [[Tree (TypeBase ExtShape NoMode)]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Map VName Int
-> StructType -> InternaliseTypeM [Tree (TypeBase ExtShape NoMode)]
forall o.
Pretty o =>
Map VName Int
-> TypeBase Size o -> InternaliseTypeM [Tree (TypeBase ExtShape o)]
internaliseTypeM Map VName Int
forall a. Monoid a => a
mempty)) Map Name [StructType]
cs

internaliseSumType ::
  M.Map Name [E.StructType] ->
  InternaliseM
    ( [I.TypeBase ExtShape NoMode],
      [(Name, [Int])]
    )
internaliseSumType :: Map Name [StructType]
-> InternaliseM ([TypeBase ExtShape NoMode], [(Name, [Int])])
internaliseSumType =
  ([TypeBase ExtShape NoMode]
 -> InternaliseM [TypeBase ExtShape NoMode])
-> ([(Name, [Int])] -> InternaliseM [(Name, [Int])])
-> ([TypeBase ExtShape NoMode], [(Name, [Int])])
-> InternaliseM ([TypeBase ExtShape NoMode], [(Name, [Int])])
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> (a, b) -> f (c, d)
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 NoMode
 -> InternaliseM (TypeBase ExtShape NoMode))
-> [TypeBase ExtShape NoMode]
-> InternaliseM [TypeBase ExtShape NoMode]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TypeBase ExtShape NoMode -> InternaliseM (TypeBase ExtShape NoMode)
forall shape o. TypeBase shape o -> InternaliseM (TypeBase shape o)
mkAccCerts) [(Name, [Int])] -> InternaliseM [(Name, [Int])]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (([TypeBase ExtShape NoMode], [(Name, [Int])])
 -> InternaliseM ([TypeBase ExtShape NoMode], [(Name, [Int])]))
-> (Map Name [StructType]
    -> ([TypeBase ExtShape NoMode], [(Name, [Int])]))
-> Map Name [StructType]
-> InternaliseM ([TypeBase ExtShape NoMode], [(Name, [Int])])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [StructType]
-> ([TypeBase ExtShape NoMode], [(Name, [Int])])
internaliseSumTypeRep

-- | How many core language values are needed to represent one source
-- language value of the given type?
internalisedTypeSize :: E.TypeBase E.Size o -> Int
-- A few special cases for performance.
internalisedTypeSize :: forall o. TypeBase Size o -> Int
internalisedTypeSize (E.Scalar (E.Prim PrimType
_)) = Int
1
internalisedTypeSize (E.Array o
_ Shape Size
_ (E.Prim PrimType
_)) = Int
1
internalisedTypeSize TypeBase Size o
t = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (Tree (TypeBase ExtShape NoMode) -> Int)
-> [Tree (TypeBase ExtShape NoMode)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Tree (TypeBase ExtShape NoMode) -> Int
forall a. Free [] a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Tree (TypeBase ExtShape NoMode)] -> [Int])
-> [Tree (TypeBase ExtShape NoMode)] -> [Int]
forall a b. (a -> b) -> a -> b
$ StructType -> [Tree (TypeBase ExtShape NoMode)]
internaliseType (StructType -> [Tree (TypeBase ExtShape NoMode)])
-> StructType -> [Tree (TypeBase ExtShape NoMode)]
forall a b. (a -> b) -> a -> b
$ TypeBase Size o -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
E.toStruct TypeBase Size o
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 nonfresh
-- 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.