-- | Generating metadata so that programs can run at all.
module Futhark.Internalise.Entry
  ( entryPoint,
    VisibleTypes,
    visibleTypes,
  )
where

import Control.Monad
import Control.Monad.State.Strict
import Data.Bifunctor (first)
import Data.List (find, intersperse)
import Data.Map qualified as M
import Data.Text qualified as T
import Futhark.IR qualified as I
import Futhark.Internalise.TypesValues (internaliseSumTypeRep, internalisedTypeSize)
import Futhark.Util (chunks)
import Futhark.Util.Pretty (prettyTextOneLine)
import Language.Futhark qualified as E hiding (TypeArg)
import Language.Futhark.Core (L (..), Name, VName, nameFromText, unLoc)
import Language.Futhark.Semantic qualified as E

-- | The types that are visible to the outside world.
newtype VisibleTypes = VisibleTypes [E.TypeBind]

-- | Retrieve those type bindings that should be visible to the
-- outside world.  Currently that is everything at top level that does
-- not have type parameters.
visibleTypes :: E.Imports -> VisibleTypes
visibleTypes :: Imports -> VisibleTypes
visibleTypes = [TypeBind] -> VisibleTypes
VisibleTypes ([TypeBind] -> VisibleTypes)
-> (Imports -> [TypeBind]) -> Imports -> VisibleTypes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((ImportName, FileModule) -> [TypeBind]) -> Imports -> [TypeBind]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (FileModule -> [TypeBind]
modTypes (FileModule -> [TypeBind])
-> ((ImportName, FileModule) -> FileModule)
-> (ImportName, FileModule)
-> [TypeBind]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ImportName, FileModule) -> FileModule
forall a b. (a, b) -> b
snd)
  where
    modTypes :: FileModule -> [TypeBind]
modTypes = ProgBase Info VName -> [TypeBind]
forall {f :: * -> *} {vn}. ProgBase f vn -> [TypeBindBase f vn]
progTypes (ProgBase Info VName -> [TypeBind])
-> (FileModule -> ProgBase Info VName) -> FileModule -> [TypeBind]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FileModule -> ProgBase Info VName
E.fileProg
    progTypes :: ProgBase f vn -> [TypeBindBase f vn]
progTypes = (DecBase f vn -> [TypeBindBase f vn])
-> [DecBase f vn] -> [TypeBindBase f vn]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DecBase f vn -> [TypeBindBase f vn]
forall {f :: * -> *} {vn}. DecBase f vn -> [TypeBindBase f vn]
decTypes ([DecBase f vn] -> [TypeBindBase f vn])
-> (ProgBase f vn -> [DecBase f vn])
-> ProgBase f vn
-> [TypeBindBase f vn]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ProgBase f vn -> [DecBase f vn]
forall (f :: * -> *) vn. ProgBase f vn -> [DecBase f vn]
E.progDecs
    decTypes :: DecBase f vn -> [TypeBindBase f vn]
decTypes (E.TypeDec TypeBindBase f vn
tb) = [TypeBindBase f vn
tb]
    decTypes DecBase f vn
_ = []

findType :: VName -> VisibleTypes -> Maybe (E.TypeExp E.Exp VName)
findType :: VName -> VisibleTypes -> Maybe (TypeExp Exp VName)
findType VName
v (VisibleTypes [TypeBind]
ts) = TypeBind -> TypeExp Exp VName
forall (f :: * -> *) vn.
TypeBindBase f vn -> TypeExp (ExpBase f vn) vn
E.typeExp (TypeBind -> TypeExp Exp VName)
-> Maybe TypeBind -> Maybe (TypeExp Exp VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBind -> Bool) -> [TypeBind] -> Maybe TypeBind
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
v) (VName -> Bool) -> (TypeBind -> VName) -> TypeBind -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBind -> VName
forall (f :: * -> *) vn. TypeBindBase f vn -> vn
E.typeAlias) [TypeBind]
ts

-- | The mode of an array type.  Anything else is at the bottom of the
-- mode lattice: 'I.Observe' for a parameter, 'E.Nonfresh' for a result.
modeOf :: (Bounded o) => I.TypeBase shape o -> o
modeOf :: forall o shape. Bounded o => TypeBase shape o -> o
modeOf (I.Array PrimType
_ shape
_ o
o) = o
o
modeOf TypeBase shape o
_ = o
forall a. Bounded a => a
minBound

valueType :: I.TypeBase I.Rank o -> I.ValueType
valueType :: forall o. TypeBase Rank o -> ValueType
valueType (I.Prim PrimType
pt) = Signedness -> Rank -> PrimType -> ValueType
I.ValueType Signedness
I.Signed (Int -> Rank
I.Rank Int
0) PrimType
pt
valueType (I.Array PrimType
pt Rank
rank o
_) = Signedness -> Rank -> PrimType -> ValueType
I.ValueType Signedness
I.Signed Rank
rank PrimType
pt
valueType I.Acc {} = String -> ValueType
forall a. HasCallStack => String -> a
error String
"valueType Acc"
valueType I.Mem {} = String -> ValueType
forall a. HasCallStack => String -> a
error String
"valueType Mem"

withoutDims :: E.TypeExp E.Exp VName -> (Int, E.TypeExp E.Exp VName)
withoutDims :: TypeExp Exp VName -> (Int, TypeExp Exp VName)
withoutDims (E.TEArray SizeExp Exp
_ TypeExp Exp VName
te SrcLoc
_) =
  let (Int
d, TypeExp Exp VName
te') = TypeExp Exp VName -> (Int, TypeExp Exp VName)
withoutDims TypeExp Exp VName
te
   in (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, TypeExp Exp VName
te')
withoutDims TypeExp Exp VName
te = (Int
0 :: Int, TypeExp Exp VName
te)

rootType :: E.TypeExp E.Exp VName -> E.TypeExp E.Exp VName
rootType :: TypeExp Exp VName -> TypeExp Exp VName
rootType (E.TEApply TypeExp Exp VName
te E.TypeArgExpSize {} SrcLoc
_) = TypeExp Exp VName -> TypeExp Exp VName
rootType TypeExp Exp VName
te
rootType (E.TEStar TypeExp Exp VName
te SrcLoc
_) = TypeExp Exp VName -> TypeExp Exp VName
rootType TypeExp Exp VName
te
rootType (E.TEDim [VName]
_ TypeExp Exp VName
te SrcLoc
_) = TypeExp Exp VName -> TypeExp Exp VName
rootType TypeExp Exp VName
te
rootType (E.TEParens TypeExp Exp VName
te SrcLoc
_) = TypeExp Exp VName -> TypeExp Exp VName
rootType TypeExp Exp VName
te
rootType TypeExp Exp VName
te = TypeExp Exp VName
te

typeExpOpaqueName :: E.TypeExp E.Exp VName -> Name
typeExpOpaqueName :: TypeExp Exp VName -> Name
typeExpOpaqueName = Text -> Name
nameFromText (Text -> Name)
-> (TypeExp Exp VName -> Text) -> TypeExp Exp VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeExp Exp VName -> Text
f
  where
    f :: TypeExp Exp VName -> Text
f = TypeExp Exp VName -> Text
g (TypeExp Exp VName -> Text)
-> (TypeExp Exp VName -> TypeExp Exp VName)
-> TypeExp Exp VName
-> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeExp Exp VName -> TypeExp Exp VName
rootType
    g :: TypeExp Exp VName -> Text
g (E.TEArray SizeExp Exp
_ TypeExp Exp VName
te SrcLoc
_) =
      let (Int
d, TypeExp Exp VName
te') = TypeExp Exp VName -> (Int, TypeExp Exp VName)
withoutDims TypeExp Exp VName
te
       in [Text] -> Text
forall a. Monoid a => [a] -> a
mconcat (Int -> Text -> [Text]
forall a. Int -> a -> [a]
replicate (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
d) Text
"[]") Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> TypeExp Exp VName -> Text
f TypeExp Exp VName
te'
    g (E.TETuple [TypeExp Exp VName]
tes SrcLoc
_) =
      Text
"(" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Text] -> Text
forall a. Monoid a => [a] -> a
mconcat (Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
intersperse Text
", " ((TypeExp Exp VName -> Text) -> [TypeExp Exp VName] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map TypeExp Exp VName -> Text
f [TypeExp Exp VName]
tes)) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")"
    g (E.TERecord [(L Name, TypeExp Exp VName)]
tes SrcLoc
_) =
      Text
"{" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Text] -> Text
forall a. Monoid a => [a] -> a
mconcat (Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
intersperse Text
", " (((L Name, TypeExp Exp VName) -> Text)
-> [(L Name, TypeExp Exp VName)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (L Name, TypeExp Exp VName) -> Text
onField [(L Name, TypeExp Exp VName)]
tes)) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"}"
      where
        onField :: (L Name, TypeExp Exp VName) -> Text
onField (L Loc
_ Name
k, TypeExp Exp VName
te) = Name -> Text
E.nameToText Name
k Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> TypeExp Exp VName -> Text
f TypeExp Exp VName
te
    g (E.TESum [(Name, [TypeExp Exp VName])]
cs SrcLoc
_) =
      [Text] -> Text
forall a. Monoid a => [a] -> a
mconcat (Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
intersperse Text
" | " (((Name, [TypeExp Exp VName]) -> Text)
-> [(Name, [TypeExp Exp VName])] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Name, [TypeExp Exp VName]) -> Text
onConstr [(Name, [TypeExp Exp VName])]
cs))
      where
        onConstr :: (Name, [TypeExp Exp VName]) -> Text
onConstr (Name
k, [TypeExp Exp VName]
tes) =
          Name -> Text
E.nameToText Name
k Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
":" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Text] -> Text
forall a. Monoid a => [a] -> a
mconcat (Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
intersperse Text
" " ((TypeExp Exp VName -> Text) -> [TypeExp Exp VName] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map TypeExp Exp VName -> Text
f [TypeExp Exp VName]
tes))
    g (E.TEParens TypeExp Exp VName
te SrcLoc
_) = Text
"(" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> TypeExp Exp VName -> Text
f TypeExp Exp VName
te Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")"
    g TypeExp Exp VName
te = TypeExp Exp VName -> Text
forall a. Pretty a => a -> Text
prettyTextOneLine TypeExp Exp VName
te

type GenOpaque = State I.OpaqueTypes

runGenOpaque :: GenOpaque a -> (a, I.OpaqueTypes)
runGenOpaque :: forall a. GenOpaque a -> (a, OpaqueTypes)
runGenOpaque = (GenOpaque a -> OpaqueTypes -> (a, OpaqueTypes))
-> OpaqueTypes -> GenOpaque a -> (a, OpaqueTypes)
forall a b c. (a -> b -> c) -> b -> a -> c
flip GenOpaque a -> OpaqueTypes -> (a, OpaqueTypes)
forall s a. State s a -> s -> (a, s)
runState OpaqueTypes
forall a. Monoid a => a
mempty

addType :: Name -> Maybe T.Text -> I.OpaqueType -> GenOpaque ()
addType :: Name -> Maybe Text -> OpaqueType -> GenOpaque ()
addType Name
name Maybe Text
doc OpaqueType
t = (OpaqueTypes -> OpaqueTypes) -> GenOpaque ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((OpaqueTypes -> OpaqueTypes) -> GenOpaque ())
-> (OpaqueTypes -> OpaqueTypes) -> GenOpaque ()
forall a b. (a -> b) -> a -> b
$ \(I.OpaqueTypes [(Name, (OpaqueType, Maybe Text))]
ts) ->
  case Name
-> [(Name, (OpaqueType, Maybe Text))]
-> Maybe (OpaqueType, Maybe Text)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
name [(Name, (OpaqueType, Maybe Text))]
ts of
    Just (OpaqueType
t', Maybe Text
_)
      | OpaqueType
t OpaqueType -> OpaqueType -> Bool
forall a. Eq a => a -> a -> Bool
/= OpaqueType
t' ->
          String -> OpaqueTypes
forall a. HasCallStack => String -> a
error (String -> OpaqueTypes)
-> ([String] -> String) -> [String] -> OpaqueTypes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String
unlines ([String] -> OpaqueTypes) -> [String] -> OpaqueTypes
forall a b. (a -> b) -> a -> b
$
            [ String
"Duplicate definition of entry point type " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Name -> String
forall a. Pretty a => a -> String
E.prettyString Name
name,
              OpaqueType -> String
forall a. Show a => a -> String
show OpaqueType
t,
              OpaqueType -> String
forall a. Show a => a -> String
show OpaqueType
t'
            ]
    Maybe (OpaqueType, Maybe Text)
_ -> [(Name, (OpaqueType, Maybe Text))] -> OpaqueTypes
I.OpaqueTypes [(Name, (OpaqueType, Maybe Text))]
ts OpaqueTypes -> OpaqueTypes -> OpaqueTypes
forall a. Semigroup a => a -> a -> a
<> [(Name, (OpaqueType, Maybe Text))] -> OpaqueTypes
I.OpaqueTypes [(Name
name, (OpaqueType
t, Maybe Text
doc))]

isRecord ::
  VisibleTypes ->
  E.TypeExp E.Exp VName ->
  Maybe (M.Map Name (E.TypeExp E.Exp VName))
isRecord :: VisibleTypes
-> TypeExp Exp VName -> Maybe (Map Name (TypeExp Exp VName))
isRecord VisibleTypes
_ (E.TERecord [(L Name, TypeExp Exp VName)]
fs SrcLoc
_) = Map Name (TypeExp Exp VName)
-> Maybe (Map Name (TypeExp Exp VName))
forall a. a -> Maybe a
Just (Map Name (TypeExp Exp VName)
 -> Maybe (Map Name (TypeExp Exp VName)))
-> Map Name (TypeExp Exp VName)
-> Maybe (Map Name (TypeExp Exp VName))
forall a b. (a -> b) -> a -> b
$ [(Name, TypeExp Exp VName)] -> Map Name (TypeExp Exp VName)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, TypeExp Exp VName)] -> Map Name (TypeExp Exp VName))
-> [(Name, TypeExp Exp VName)] -> Map Name (TypeExp Exp VName)
forall a b. (a -> b) -> a -> b
$ ((L Name, TypeExp Exp VName) -> (Name, TypeExp Exp VName))
-> [(L Name, TypeExp Exp VName)] -> [(Name, TypeExp Exp VName)]
forall a b. (a -> b) -> [a] -> [b]
map ((L Name -> Name)
-> (L Name, TypeExp Exp VName) -> (Name, TypeExp Exp VName)
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 L Name -> Name
forall a. L a -> a
unLoc) [(L Name, TypeExp Exp VName)]
fs
isRecord VisibleTypes
_ (E.TETuple [TypeExp Exp VName]
fs SrcLoc
_) = Map Name (TypeExp Exp VName)
-> Maybe (Map Name (TypeExp Exp VName))
forall a. a -> Maybe a
Just (Map Name (TypeExp Exp VName)
 -> Maybe (Map Name (TypeExp Exp VName)))
-> Map Name (TypeExp Exp VName)
-> Maybe (Map Name (TypeExp Exp VName))
forall a b. (a -> b) -> a -> b
$ [TypeExp Exp VName] -> Map Name (TypeExp Exp VName)
forall a. [a] -> Map Name a
E.tupleFields [TypeExp Exp VName]
fs
isRecord VisibleTypes
types (E.TEVar QualName VName
v SrcLoc
_) = VisibleTypes
-> TypeExp Exp VName -> Maybe (Map Name (TypeExp Exp VName))
isRecord VisibleTypes
types (TypeExp Exp VName -> Maybe (Map Name (TypeExp Exp VName)))
-> Maybe (TypeExp Exp VName)
-> Maybe (Map Name (TypeExp Exp VName))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> VisibleTypes -> Maybe (TypeExp Exp VName)
findType (QualName VName -> VName
forall vn. QualName vn -> vn
E.qualLeaf QualName VName
v) VisibleTypes
types
isRecord VisibleTypes
_ TypeExp Exp VName
_ = Maybe (Map Name (TypeExp Exp VName))
forall a. Maybe a
Nothing

recordFields ::
  VisibleTypes ->
  M.Map Name E.StructType ->
  Maybe (E.TypeExp E.Exp VName) ->
  [(Name, E.EntryType)]
recordFields :: VisibleTypes
-> Map Name StructType
-> Maybe (TypeExp Exp VName)
-> [(Name, EntryType)]
recordFields VisibleTypes
types Map Name StructType
fs Maybe (TypeExp Exp VName)
t =
  case VisibleTypes
-> TypeExp Exp VName -> Maybe (Map Name (TypeExp Exp VName))
isRecord VisibleTypes
types (TypeExp Exp VName -> Maybe (Map Name (TypeExp Exp VName)))
-> (TypeExp Exp VName -> TypeExp Exp VName)
-> TypeExp Exp VName
-> Maybe (Map Name (TypeExp Exp VName))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeExp Exp VName -> TypeExp Exp VName
rootType (TypeExp Exp VName -> Maybe (Map Name (TypeExp Exp VName)))
-> Maybe (TypeExp Exp VName)
-> Maybe (Map Name (TypeExp Exp VName))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Maybe (TypeExp Exp VName)
t of
    Just Map Name (TypeExp Exp VName)
e_fs ->
      ((Name, StructType)
 -> (Name, TypeExp Exp VName) -> (Name, EntryType))
-> [(Name, StructType)]
-> [(Name, TypeExp Exp VName)]
-> [(Name, EntryType)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (Name, StructType)
-> (Name, TypeExp Exp VName) -> (Name, EntryType)
forall {a} {a}.
(a, StructType) -> (a, TypeExp Exp VName) -> (a, EntryType)
f (Map Name StructType -> [(Name, StructType)]
forall a. Map Name a -> [(Name, a)]
E.sortFields Map Name StructType
fs) (Map Name (TypeExp Exp VName) -> [(Name, TypeExp Exp VName)]
forall a. Map Name a -> [(Name, a)]
E.sortFields Map Name (TypeExp Exp VName)
e_fs)
      where
        f :: (a, StructType) -> (a, TypeExp Exp VName) -> (a, EntryType)
f (a
k, StructType
f_t) (a
_, TypeExp Exp VName
e_f_t) = (a
k, StructType -> Maybe (TypeExp Exp VName) -> EntryType
E.EntryType StructType
f_t (Maybe (TypeExp Exp VName) -> EntryType)
-> Maybe (TypeExp Exp VName) -> EntryType
forall a b. (a -> b) -> a -> b
$ TypeExp Exp VName -> Maybe (TypeExp Exp VName)
forall a. a -> Maybe a
Just TypeExp Exp VName
e_f_t)
    Maybe (Map Name (TypeExp Exp VName))
Nothing ->
      ((Name, StructType) -> (Name, EntryType))
-> [(Name, StructType)] -> [(Name, EntryType)]
forall a b. (a -> b) -> [a] -> [b]
map ((StructType -> EntryType)
-> (Name, StructType) -> (Name, EntryType)
forall a b. (a -> b) -> (Name, a) -> (Name, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (StructType -> Maybe (TypeExp Exp VName) -> EntryType
`E.EntryType` Maybe (TypeExp Exp VName)
forall a. Maybe a
Nothing)) ([(Name, StructType)] -> [(Name, EntryType)])
-> [(Name, StructType)] -> [(Name, EntryType)]
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> [(Name, StructType)]
forall a. Map Name a -> [(Name, a)]
E.sortFields Map Name StructType
fs

opaqueRecord ::
  (Ord o, Bounded o) =>
  VisibleTypes ->
  [(Name, E.EntryType)] ->
  [I.TypeBase I.Rank o] ->
  GenOpaque [(Name, I.EntryPointType)]
opaqueRecord :: forall o.
(Ord o, Bounded o) =>
VisibleTypes
-> [(Name, EntryType)]
-> [TypeBase Rank o]
-> GenOpaque [(Name, EntryPointType)]
opaqueRecord VisibleTypes
_ [] [TypeBase Rank o]
_ = [(Name, EntryPointType)] -> GenOpaque [(Name, EntryPointType)]
forall a. a -> StateT OpaqueTypes Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
opaqueRecord VisibleTypes
types ((Name
f, EntryType
t) : [(Name, EntryType)]
fs) [TypeBase Rank o]
ts = do
  let ([TypeBase Rank o]
f_ts, [TypeBase Rank o]
ts') = Int -> [TypeBase Rank o] -> ([TypeBase Rank o], [TypeBase Rank o])
forall a. Int -> [a] -> ([a], [a])
splitAt (StructType -> Int
forall o. TypeBase Exp o -> Int
internalisedTypeSize (StructType -> Int) -> StructType -> Int
forall a b. (a -> b) -> a -> b
$ EntryType -> StructType
E.entryType EntryType
t) [TypeBase Rank o]
ts
  f' <- EntryType
-> [TypeBase Rank o] -> StateT OpaqueTypes Identity EntryPointType
forall {a}.
(Ord a, Bounded a) =>
EntryType
-> [TypeBase Rank a] -> StateT OpaqueTypes Identity EntryPointType
opaqueField EntryType
t [TypeBase Rank o]
f_ts
  ((f, f') :) <$> opaqueRecord types fs ts'
  where
    opaqueField :: EntryType
-> [TypeBase Rank a] -> StateT OpaqueTypes Identity EntryPointType
opaqueField EntryType
e_t [TypeBase Rank a]
i_ts = (a, EntryPointType) -> EntryPointType
forall a b. (a, b) -> b
snd ((a, EntryPointType) -> EntryPointType)
-> StateT OpaqueTypes Identity (a, EntryPointType)
-> StateT OpaqueTypes Identity EntryPointType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VisibleTypes
-> EntryType
-> [TypeBase Rank a]
-> StateT OpaqueTypes Identity (a, EntryPointType)
forall o.
(Ord o, Bounded o) =>
VisibleTypes
-> EntryType -> [TypeBase Rank o] -> GenOpaque (o, EntryPointType)
entryPointType VisibleTypes
types EntryType
e_t [TypeBase Rank a]
i_ts

teArrayOf :: Int -> E.TypeExp d vn -> E.TypeExp d vn
teArrayOf :: forall d vn. Int -> TypeExp d vn -> TypeExp d vn
teArrayOf Int
rank TypeExp d vn
t =
  (TypeExp d vn -> Int -> TypeExp d vn)
-> TypeExp d vn -> [Int] -> TypeExp d vn
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
    (\TypeExp d vn
x Int
_ -> SizeExp d -> TypeExp d vn -> SrcLoc -> TypeExp d vn
forall d vn. SizeExp d -> TypeExp d vn -> SrcLoc -> TypeExp d vn
E.TEArray (SrcLoc -> SizeExp d
forall d. SrcLoc -> SizeExp d
E.SizeExpAny SrcLoc
forall a. Monoid a => a
mempty) TypeExp d vn
x SrcLoc
forall a. Monoid a => a
mempty)
    TypeExp d vn
t
    [Int
0 .. Int
rank Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

opaqueRecordArray ::
  (Ord o, Bounded o) =>
  VisibleTypes ->
  Int ->
  [(Name, E.EntryType)] ->
  [I.TypeBase I.Rank o] ->
  GenOpaque [(Name, I.EntryPointType)]
opaqueRecordArray :: forall o.
(Ord o, Bounded o) =>
VisibleTypes
-> Int
-> [(Name, EntryType)]
-> [TypeBase Rank o]
-> GenOpaque [(Name, EntryPointType)]
opaqueRecordArray VisibleTypes
_ Int
_ [] [TypeBase Rank o]
_ = [(Name, EntryPointType)] -> GenOpaque [(Name, EntryPointType)]
forall a. a -> StateT OpaqueTypes Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
opaqueRecordArray VisibleTypes
types Int
rank ((Name
f, EntryType
t) : [(Name, EntryType)]
fs) [TypeBase Rank o]
ts = do
  let ([TypeBase Rank o]
f_ts, [TypeBase Rank o]
ts') = Int -> [TypeBase Rank o] -> ([TypeBase Rank o], [TypeBase Rank o])
forall a. Int -> [a] -> ([a], [a])
splitAt (StructType -> Int
forall o. TypeBase Exp o -> Int
internalisedTypeSize (StructType -> Int) -> StructType -> Int
forall a b. (a -> b) -> a -> b
$ EntryType -> StructType
E.entryType EntryType
t) [TypeBase Rank o]
ts
  f' <- EntryType
-> [TypeBase Rank o] -> StateT OpaqueTypes Identity EntryPointType
forall {a}.
(Ord a, Bounded a) =>
EntryType
-> [TypeBase Rank a] -> StateT OpaqueTypes Identity EntryPointType
opaqueField EntryType
t [TypeBase Rank o]
f_ts
  ((f, f') :) <$> opaqueRecordArray types rank fs ts'
  where
    opaqueField :: EntryType
-> [TypeBase Rank a] -> StateT OpaqueTypes Identity EntryPointType
opaqueField (E.EntryType StructType
e_t Maybe (TypeExp Exp VName)
ascribed) [TypeBase Rank a]
i_ts =
      (a, EntryPointType) -> EntryPointType
forall a b. (a, b) -> b
snd ((a, EntryPointType) -> EntryPointType)
-> StateT OpaqueTypes Identity (a, EntryPointType)
-> StateT OpaqueTypes Identity EntryPointType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VisibleTypes
-> EntryType
-> [TypeBase Rank a]
-> StateT OpaqueTypes Identity (a, EntryPointType)
forall o.
(Ord o, Bounded o) =>
VisibleTypes
-> EntryType -> [TypeBase Rank o] -> GenOpaque (o, EntryPointType)
entryPointType VisibleTypes
types (StructType -> Maybe (TypeExp Exp VName) -> EntryType
E.EntryType StructType
e_t' Maybe (TypeExp Exp VName)
ascribed') [TypeBase Rank a]
i_ts
      where
        e_t' :: StructType
e_t' = Shape Exp -> StructType -> StructType
forall dim. Shape dim -> TypeBase dim NoMode -> TypeBase dim NoMode
E.arrayOf ([Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
E.Shape (Int -> Exp -> [Exp]
forall a. Int -> a -> [a]
replicate Int
rank (Exp -> [Exp]) -> Exp -> [Exp]
forall a b. (a -> b) -> a -> b
$ Int -> Exp
E.anySize Int
0)) StructType
e_t
        ascribed' :: Maybe (TypeExp Exp VName)
ascribed' = Int -> TypeExp Exp VName -> TypeExp Exp VName
forall d vn. Int -> TypeExp d vn -> TypeExp d vn
teArrayOf Int
rank (TypeExp Exp VName -> TypeExp Exp VName)
-> Maybe (TypeExp Exp VName) -> Maybe (TypeExp Exp VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (TypeExp Exp VName)
ascribed

isSum ::
  VisibleTypes ->
  E.TypeExp E.Exp VName ->
  Maybe (M.Map Name [E.TypeExp E.Exp VName])
isSum :: VisibleTypes
-> TypeExp Exp VName -> Maybe (Map Name [TypeExp Exp VName])
isSum VisibleTypes
_ (E.TESum [(Name, [TypeExp Exp VName])]
cs SrcLoc
_) = Map Name [TypeExp Exp VName]
-> Maybe (Map Name [TypeExp Exp VName])
forall a. a -> Maybe a
Just (Map Name [TypeExp Exp VName]
 -> Maybe (Map Name [TypeExp Exp VName]))
-> Map Name [TypeExp Exp VName]
-> Maybe (Map Name [TypeExp Exp VName])
forall a b. (a -> b) -> a -> b
$ [(Name, [TypeExp Exp VName])] -> Map Name [TypeExp Exp VName]
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, [TypeExp Exp VName])]
cs
isSum VisibleTypes
types (E.TEVar QualName VName
v SrcLoc
_) = VisibleTypes
-> TypeExp Exp VName -> Maybe (Map Name [TypeExp Exp VName])
isSum VisibleTypes
types (TypeExp Exp VName -> Maybe (Map Name [TypeExp Exp VName]))
-> Maybe (TypeExp Exp VName)
-> Maybe (Map Name [TypeExp Exp VName])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName -> VisibleTypes -> Maybe (TypeExp Exp VName)
findType (QualName VName -> VName
forall vn. QualName vn -> vn
E.qualLeaf QualName VName
v) VisibleTypes
types
isSum VisibleTypes
_ TypeExp Exp VName
_ = Maybe (Map Name [TypeExp Exp VName])
forall a. Maybe a
Nothing

sumConstrs ::
  VisibleTypes ->
  M.Map Name [E.StructType] ->
  Maybe (E.TypeExp E.Exp VName) ->
  [(Name, [E.EntryType])]
sumConstrs :: VisibleTypes
-> Map Name [StructType]
-> Maybe (TypeExp Exp VName)
-> [(Name, [EntryType])]
sumConstrs VisibleTypes
types Map Name [StructType]
cs Maybe (TypeExp Exp VName)
t =
  case VisibleTypes
-> TypeExp Exp VName -> Maybe (Map Name [TypeExp Exp VName])
isSum VisibleTypes
types (TypeExp Exp VName -> Maybe (Map Name [TypeExp Exp VName]))
-> (TypeExp Exp VName -> TypeExp Exp VName)
-> TypeExp Exp VName
-> Maybe (Map Name [TypeExp Exp VName])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeExp Exp VName -> TypeExp Exp VName
rootType (TypeExp Exp VName -> Maybe (Map Name [TypeExp Exp VName]))
-> Maybe (TypeExp Exp VName)
-> Maybe (Map Name [TypeExp Exp VName])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Maybe (TypeExp Exp VName)
t of
    Just Map Name [TypeExp Exp VName]
e_cs ->
      ((Name, [StructType])
 -> (Name, [TypeExp Exp VName]) -> (Name, [EntryType]))
-> [(Name, [StructType])]
-> [(Name, [TypeExp Exp VName])]
-> [(Name, [EntryType])]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (Name, [StructType])
-> (Name, [TypeExp Exp VName]) -> (Name, [EntryType])
forall {a} {a}.
(a, [StructType]) -> (a, [TypeExp Exp VName]) -> (a, [EntryType])
f (Map Name [StructType] -> [(Name, [StructType])]
forall a. Map Name a -> [(Name, a)]
E.sortConstrs Map Name [StructType]
cs) (Map Name [TypeExp Exp VName] -> [(Name, [TypeExp Exp VName])]
forall a. Map Name a -> [(Name, a)]
E.sortConstrs Map Name [TypeExp Exp VName]
e_cs)
      where
        f :: (a, [StructType]) -> (a, [TypeExp Exp VName]) -> (a, [EntryType])
f (a
k, [StructType]
c_ts) (a
_, [TypeExp Exp VName]
e_c_ts) = (a
k, (StructType -> Maybe (TypeExp Exp VName) -> EntryType)
-> [StructType] -> [Maybe (TypeExp Exp VName)] -> [EntryType]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith StructType -> Maybe (TypeExp Exp VName) -> EntryType
E.EntryType [StructType]
c_ts ([Maybe (TypeExp Exp VName)] -> [EntryType])
-> [Maybe (TypeExp Exp VName)] -> [EntryType]
forall a b. (a -> b) -> a -> b
$ (TypeExp Exp VName -> Maybe (TypeExp Exp VName))
-> [TypeExp Exp VName] -> [Maybe (TypeExp Exp VName)]
forall a b. (a -> b) -> [a] -> [b]
map TypeExp Exp VName -> Maybe (TypeExp Exp VName)
forall a. a -> Maybe a
Just [TypeExp Exp VName]
e_c_ts)
    Maybe (Map Name [TypeExp Exp VName])
Nothing ->
      ((Name, [StructType]) -> (Name, [EntryType]))
-> [(Name, [StructType])] -> [(Name, [EntryType])]
forall a b. (a -> b) -> [a] -> [b]
map (([StructType] -> [EntryType])
-> (Name, [StructType]) -> (Name, [EntryType])
forall a b. (a -> b) -> (Name, a) -> (Name, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((StructType -> EntryType) -> [StructType] -> [EntryType]
forall a b. (a -> b) -> [a] -> [b]
map (StructType -> Maybe (TypeExp Exp VName) -> EntryType
`E.EntryType` Maybe (TypeExp Exp VName)
forall a. Maybe a
Nothing))) ([(Name, [StructType])] -> [(Name, [EntryType])])
-> [(Name, [StructType])] -> [(Name, [EntryType])]
forall a b. (a -> b) -> a -> b
$ Map Name [StructType] -> [(Name, [StructType])]
forall a. Map Name a -> [(Name, a)]
E.sortConstrs Map Name [StructType]
cs

opaqueSum ::
  (Ord o, Bounded o) =>
  VisibleTypes ->
  [(Name, ([E.EntryType], [Int]))] ->
  [I.TypeBase I.Rank o] ->
  GenOpaque [(Name, [(I.EntryPointType, [Int])])]
opaqueSum :: forall o.
(Ord o, Bounded o) =>
VisibleTypes
-> [(Name, ([EntryType], [Int]))]
-> [TypeBase Rank o]
-> GenOpaque [(Name, [(EntryPointType, [Int])])]
opaqueSum VisibleTypes
types [(Name, ([EntryType], [Int]))]
cs [TypeBase Rank o]
ts = ((Name, ([EntryType], [Int]))
 -> StateT OpaqueTypes Identity (Name, [(EntryPointType, [Int])]))
-> [(Name, ([EntryType], [Int]))]
-> GenOpaque [(Name, [(EntryPointType, [Int])])]
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 ((([EntryType], [Int])
 -> StateT OpaqueTypes Identity [(EntryPointType, [Int])])
-> (Name, ([EntryType], [Int]))
-> StateT OpaqueTypes Identity (Name, [(EntryPointType, [Int])])
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) -> (Name, a) -> f (Name, b)
traverse ([EntryType], [Int])
-> StateT OpaqueTypes Identity [(EntryPointType, [Int])]
f) [(Name, ([EntryType], [Int]))]
cs
  where
    f :: ([EntryType], [Int])
-> StateT OpaqueTypes Identity [(EntryPointType, [Int])]
f ([EntryType]
ets, [Int]
is) = do
      let ns :: [Int]
ns = (EntryType -> Int) -> [EntryType] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (StructType -> Int
forall o. TypeBase Exp o -> Int
internalisedTypeSize (StructType -> Int)
-> (EntryType -> StructType) -> EntryType -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EntryType -> StructType
E.entryType) [EntryType]
ets
          is' :: [[Int]]
is' = [Int] -> [Int] -> [[Int]]
forall a. [Int] -> [a] -> [[a]]
chunks [Int]
ns [Int]
is
      ets' <- ((o, EntryPointType) -> EntryPointType)
-> [(o, EntryPointType)] -> [EntryPointType]
forall a b. (a -> b) -> [a] -> [b]
map (o, EntryPointType) -> EntryPointType
forall a b. (a, b) -> b
snd ([(o, EntryPointType)] -> [EntryPointType])
-> StateT OpaqueTypes Identity [(o, EntryPointType)]
-> StateT OpaqueTypes Identity [EntryPointType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (EntryType
 -> [TypeBase Rank o]
 -> StateT OpaqueTypes Identity (o, EntryPointType))
-> [EntryType]
-> [[TypeBase Rank o]]
-> StateT OpaqueTypes Identity [(o, EntryPointType)]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (VisibleTypes
-> EntryType
-> [TypeBase Rank o]
-> StateT OpaqueTypes Identity (o, EntryPointType)
forall o.
(Ord o, Bounded o) =>
VisibleTypes
-> EntryType -> [TypeBase Rank o] -> GenOpaque (o, EntryPointType)
entryPointType VisibleTypes
types) [EntryType]
ets (([Int] -> [TypeBase Rank o]) -> [[Int]] -> [[TypeBase Rank o]]
forall a b. (a -> b) -> [a] -> [b]
map ((Int -> TypeBase Rank o) -> [Int] -> [TypeBase Rank o]
forall a b. (a -> b) -> [a] -> [b]
map ([TypeBase Rank o]
ts [TypeBase Rank o] -> Int -> TypeBase Rank o
forall a. HasCallStack => [a] -> Int -> a
!!)) [[Int]]
is')
      pure . zip ets' $
        if length cs == 1
          then is'
          else map (map (+ 1)) is' -- Adjust for tag.

entryPointTypeName :: I.EntryPointType -> Name
entryPointTypeName :: EntryPointType -> Name
entryPointTypeName (I.TypeOpaque Name
v) = Name
v
entryPointTypeName (I.TypeTransparent {}) = String -> Name
forall a. HasCallStack => String -> a
error String
"entryPointTypeName: TypeTransparent"

elemTypeExp :: E.TypeExp E.Exp VName -> Maybe (E.TypeExp E.Exp VName)
elemTypeExp :: TypeExp Exp VName -> Maybe (TypeExp Exp VName)
elemTypeExp (E.TEArray SizeExp Exp
_ TypeExp Exp VName
te SrcLoc
_) = TypeExp Exp VName -> Maybe (TypeExp Exp VName)
forall a. a -> Maybe a
Just TypeExp Exp VName
te
elemTypeExp (E.TEStar TypeExp Exp VName
te SrcLoc
_) = TypeExp Exp VName -> Maybe (TypeExp Exp VName)
elemTypeExp TypeExp Exp VName
te
elemTypeExp (E.TEParens TypeExp Exp VName
te SrcLoc
_) = TypeExp Exp VName -> Maybe (TypeExp Exp VName)
elemTypeExp TypeExp Exp VName
te
elemTypeExp TypeExp Exp VName
_ = Maybe (TypeExp Exp VName)
forall a. Maybe a
Nothing

rowTypeExp :: Int -> E.TypeExp E.Exp VName -> Maybe (E.TypeExp E.Exp VName)
rowTypeExp :: Int -> TypeExp Exp VName -> Maybe (TypeExp Exp VName)
rowTypeExp Int
0 TypeExp Exp VName
te = TypeExp Exp VName -> Maybe (TypeExp Exp VName)
forall a. a -> Maybe a
Just TypeExp Exp VName
te
rowTypeExp Int
r TypeExp Exp VName
te = Int -> TypeExp Exp VName -> Maybe (TypeExp Exp VName)
rowTypeExp (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (TypeExp Exp VName -> Maybe (TypeExp Exp VName))
-> Maybe (TypeExp Exp VName) -> Maybe (TypeExp Exp VName)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TypeExp Exp VName -> Maybe (TypeExp Exp VName)
elemTypeExp TypeExp Exp VName
te

entryPointType ::
  (Ord o, Bounded o) =>
  VisibleTypes ->
  E.EntryType ->
  [I.TypeBase I.Rank o] ->
  GenOpaque (o, I.EntryPointType)
entryPointType :: forall o.
(Ord o, Bounded o) =>
VisibleTypes
-> EntryType -> [TypeBase Rank o] -> GenOpaque (o, EntryPointType)
entryPointType VisibleTypes
types EntryType
t [TypeBase Rank o]
ts
  | E.Scalar (E.Prim E.Unsigned {}) <- EntryType -> StructType
E.entryType EntryType
t,
    [I.Prim PrimType
ts0] <- [TypeBase Rank o]
ts =
      (o, EntryPointType)
-> StateT OpaqueTypes Identity (o, EntryPointType)
forall a. a -> StateT OpaqueTypes Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (o
o, ValueType -> EntryPointType
I.TypeTransparent (ValueType -> EntryPointType) -> ValueType -> EntryPointType
forall a b. (a -> b) -> a -> b
$ Signedness -> Rank -> PrimType -> ValueType
I.ValueType Signedness
I.Unsigned (Int -> Rank
I.Rank Int
0) PrimType
ts0)
  | E.Array NoMode
_ Shape Exp
_ (E.Prim E.Unsigned {}) <- EntryType -> StructType
E.entryType EntryType
t,
    [I.Array PrimType
ts0 Rank
r o
_] <- [TypeBase Rank o]
ts =
      (o, EntryPointType)
-> StateT OpaqueTypes Identity (o, EntryPointType)
forall a. a -> StateT OpaqueTypes Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (o
o, ValueType -> EntryPointType
I.TypeTransparent (ValueType -> EntryPointType) -> ValueType -> EntryPointType
forall a b. (a -> b) -> a -> b
$ Signedness -> Rank -> PrimType -> ValueType
I.ValueType Signedness
I.Unsigned Rank
r PrimType
ts0)
  | E.Scalar E.Prim {} <- EntryType -> StructType
E.entryType EntryType
t,
    [I.Prim PrimType
ts0] <- [TypeBase Rank o]
ts =
      (o, EntryPointType)
-> StateT OpaqueTypes Identity (o, EntryPointType)
forall a. a -> StateT OpaqueTypes Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (o
o, ValueType -> EntryPointType
I.TypeTransparent (ValueType -> EntryPointType) -> ValueType -> EntryPointType
forall a b. (a -> b) -> a -> b
$ Signedness -> Rank -> PrimType -> ValueType
I.ValueType Signedness
I.Signed (Int -> Rank
I.Rank Int
0) PrimType
ts0)
  | E.Array NoMode
_ Shape Exp
_ E.Prim {} <- EntryType -> StructType
E.entryType EntryType
t,
    [I.Array PrimType
ts0 Rank
r o
_] <- [TypeBase Rank o]
ts =
      (o, EntryPointType)
-> StateT OpaqueTypes Identity (o, EntryPointType)
forall a. a -> StateT OpaqueTypes Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (o
o, ValueType -> EntryPointType
I.TypeTransparent (ValueType -> EntryPointType) -> ValueType -> EntryPointType
forall a b. (a -> b) -> a -> b
$ Signedness -> Rank -> PrimType -> ValueType
I.ValueType Signedness
I.Signed Rank
r PrimType
ts0)
  | Bool
otherwise = do
      case EntryType -> StructType
E.entryType EntryType
t of
        E.Scalar (E.Record Map Name StructType
fs) -> do
          let fs' :: [(Name, EntryType)]
fs' = VisibleTypes
-> Map Name StructType
-> Maybe (TypeExp Exp VName)
-> [(Name, EntryType)]
recordFields VisibleTypes
types Map Name StructType
fs (Maybe (TypeExp Exp VName) -> [(Name, EntryType)])
-> Maybe (TypeExp Exp VName) -> [(Name, EntryType)]
forall a b. (a -> b) -> a -> b
$ EntryType -> Maybe (TypeExp Exp VName)
E.entryAscribed EntryType
t
          Name -> Maybe Text -> OpaqueType -> GenOpaque ()
addType Name
desc Maybe Text
forall a. Maybe a
doc (OpaqueType -> GenOpaque ())
-> ([(Name, EntryPointType)] -> OpaqueType)
-> [(Name, EntryPointType)]
-> GenOpaque ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Name, EntryPointType)] -> OpaqueType
I.OpaqueRecord ([(Name, EntryPointType)] -> GenOpaque ())
-> GenOpaque [(Name, EntryPointType)] -> GenOpaque ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VisibleTypes
-> [(Name, EntryType)]
-> [TypeBase Rank o]
-> GenOpaque [(Name, EntryPointType)]
forall o.
(Ord o, Bounded o) =>
VisibleTypes
-> [(Name, EntryType)]
-> [TypeBase Rank o]
-> GenOpaque [(Name, EntryPointType)]
opaqueRecord VisibleTypes
types [(Name, EntryType)]
fs' [TypeBase Rank o]
ts
        E.Scalar (E.Sum Map Name [StructType]
cs) -> do
          let ([TypeBase ExtShape NoMode]
_, [(Name, [Int])]
places) = Map Name [StructType]
-> ([TypeBase ExtShape NoMode], [(Name, [Int])])
internaliseSumTypeRep Map Name [StructType]
cs
              cs' :: [(Name, [EntryType])]
cs' = VisibleTypes
-> Map Name [StructType]
-> Maybe (TypeExp Exp VName)
-> [(Name, [EntryType])]
sumConstrs VisibleTypes
types Map Name [StructType]
cs (Maybe (TypeExp Exp VName) -> [(Name, [EntryType])])
-> Maybe (TypeExp Exp VName) -> [(Name, [EntryType])]
forall a b. (a -> b) -> a -> b
$ EntryType -> Maybe (TypeExp Exp VName)
E.entryAscribed EntryType
t
              cs'' :: [(Name, ([EntryType], [Int]))]
cs'' = [Name] -> [([EntryType], [Int])] -> [(Name, ([EntryType], [Int]))]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Name, [EntryType]) -> Name) -> [(Name, [EntryType])] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, [EntryType]) -> Name
forall a b. (a, b) -> a
fst [(Name, [EntryType])]
cs') ([[EntryType]] -> [[Int]] -> [([EntryType], [Int])]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Name, [EntryType]) -> [EntryType])
-> [(Name, [EntryType])] -> [[EntryType]]
forall a b. (a -> b) -> [a] -> [b]
map (Name, [EntryType]) -> [EntryType]
forall a b. (a, b) -> b
snd [(Name, [EntryType])]
cs') (((Name, [Int]) -> [Int]) -> [(Name, [Int])] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map (Name, [Int]) -> [Int]
forall a b. (a, b) -> b
snd [(Name, [Int])]
places))
              ts' :: [TypeBase Rank o]
ts' = if Map Name [StructType] -> Int
forall a. Map Name a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map Name [StructType]
cs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 then [TypeBase Rank o]
ts else Int -> [TypeBase Rank o] -> [TypeBase Rank o]
forall a. Int -> [a] -> [a]
drop Int
1 [TypeBase Rank o]
ts
          Name -> Maybe Text -> OpaqueType -> GenOpaque ()
addType Name
desc Maybe Text
forall a. Maybe a
doc (OpaqueType -> GenOpaque ())
-> ([(Name, [(EntryPointType, [Int])])] -> OpaqueType)
-> [(Name, [(EntryPointType, [Int])])]
-> GenOpaque ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ValueType] -> [(Name, [(EntryPointType, [Int])])] -> OpaqueType
I.OpaqueSum ((TypeBase Rank o -> ValueType) -> [TypeBase Rank o] -> [ValueType]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Rank o -> ValueType
forall o. TypeBase Rank o -> ValueType
valueType [TypeBase Rank o]
ts)
            ([(Name, [(EntryPointType, [Int])])] -> GenOpaque ())
-> GenOpaque [(Name, [(EntryPointType, [Int])])] -> GenOpaque ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VisibleTypes
-> [(Name, ([EntryType], [Int]))]
-> [TypeBase Rank o]
-> GenOpaque [(Name, [(EntryPointType, [Int])])]
forall o.
(Ord o, Bounded o) =>
VisibleTypes
-> [(Name, ([EntryType], [Int]))]
-> [TypeBase Rank o]
-> GenOpaque [(Name, [(EntryPointType, [Int])])]
opaqueSum VisibleTypes
types [(Name, ([EntryType], [Int]))]
cs'' [TypeBase Rank o]
ts'
        E.Array NoMode
_ Shape Exp
shape (E.Record Map Name StructType
fs) | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> Bool
forall a. Map Name a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Map Name StructType
fs -> do
          let rank :: Int
rank = Shape Exp -> Int
forall dim. Shape dim -> Int
E.shapeRank Shape Exp
shape
              fs' :: [(Name, EntryType)]
fs' = VisibleTypes
-> Map Name StructType
-> Maybe (TypeExp Exp VName)
-> [(Name, EntryType)]
recordFields VisibleTypes
types Map Name StructType
fs (Maybe (TypeExp Exp VName) -> [(Name, EntryType)])
-> Maybe (TypeExp Exp VName) -> [(Name, EntryType)]
forall a b. (a -> b) -> a -> b
$ Int -> TypeExp Exp VName -> Maybe (TypeExp Exp VName)
rowTypeExp Int
rank (TypeExp Exp VName -> Maybe (TypeExp Exp VName))
-> Maybe (TypeExp Exp VName) -> Maybe (TypeExp Exp VName)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< EntryType -> Maybe (TypeExp Exp VName)
E.entryAscribed EntryType
t
              ts' :: [TypeBase Rank o]
ts' = (TypeBase Rank o -> TypeBase Rank o)
-> [TypeBase Rank o] -> [TypeBase Rank o]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> TypeBase Rank o -> TypeBase Rank o
forall {o}. Int -> TypeBase Rank o -> TypeBase Rank o
strip Int
rank) [TypeBase Rank o]
ts
              record_t :: StructType
record_t = ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
E.Scalar (Map Name StructType -> ScalarTypeBase Exp NoMode
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
E.Record Map Name StructType
fs)
              record_te :: Maybe (TypeExp Exp VName)
record_te = Int -> TypeExp Exp VName -> Maybe (TypeExp Exp VName)
rowTypeExp Int
rank (TypeExp Exp VName -> Maybe (TypeExp Exp VName))
-> Maybe (TypeExp Exp VName) -> Maybe (TypeExp Exp VName)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< EntryType -> Maybe (TypeExp Exp VName)
E.entryAscribed EntryType
t
          ept <- (o, EntryPointType) -> EntryPointType
forall a b. (a, b) -> b
snd ((o, EntryPointType) -> EntryPointType)
-> StateT OpaqueTypes Identity (o, EntryPointType)
-> StateT OpaqueTypes Identity EntryPointType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VisibleTypes
-> EntryType
-> [TypeBase Rank o]
-> StateT OpaqueTypes Identity (o, EntryPointType)
forall o.
(Ord o, Bounded o) =>
VisibleTypes
-> EntryType -> [TypeBase Rank o] -> GenOpaque (o, EntryPointType)
entryPointType VisibleTypes
types (StructType -> Maybe (TypeExp Exp VName) -> EntryType
E.EntryType StructType
record_t Maybe (TypeExp Exp VName)
record_te) [TypeBase Rank o]
ts'
          addType desc doc . I.OpaqueRecordArray rank (entryPointTypeName ept)
            =<< opaqueRecordArray types rank fs' ts
        E.Array NoMode
_ Shape Exp
shape ScalarTypeBase Exp NoMode
et -> do
          let ts' :: [TypeBase Rank o]
ts' = (TypeBase Rank o -> TypeBase Rank o)
-> [TypeBase Rank o] -> [TypeBase Rank o]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> TypeBase Rank o -> TypeBase Rank o
forall {o}. Int -> TypeBase Rank o -> TypeBase Rank o
strip (Shape Exp -> Int
forall dim. Shape dim -> Int
E.shapeRank Shape Exp
shape)) [TypeBase Rank o]
ts
              rank :: Int
rank = Shape Exp -> Int
forall dim. Shape dim -> Int
E.shapeRank Shape Exp
shape
              elem_te :: Maybe (TypeExp Exp VName)
elem_te = Int -> TypeExp Exp VName -> Maybe (TypeExp Exp VName)
rowTypeExp Int
rank (TypeExp Exp VName -> Maybe (TypeExp Exp VName))
-> Maybe (TypeExp Exp VName) -> Maybe (TypeExp Exp VName)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< EntryType -> Maybe (TypeExp Exp VName)
E.entryAscribed EntryType
t
          ept <- (o, EntryPointType) -> EntryPointType
forall a b. (a, b) -> b
snd ((o, EntryPointType) -> EntryPointType)
-> StateT OpaqueTypes Identity (o, EntryPointType)
-> StateT OpaqueTypes Identity EntryPointType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VisibleTypes
-> EntryType
-> [TypeBase Rank o]
-> StateT OpaqueTypes Identity (o, EntryPointType)
forall o.
(Ord o, Bounded o) =>
VisibleTypes
-> EntryType -> [TypeBase Rank o] -> GenOpaque (o, EntryPointType)
entryPointType VisibleTypes
types (StructType -> Maybe (TypeExp Exp VName) -> EntryType
E.EntryType (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
E.Scalar ScalarTypeBase Exp NoMode
et) Maybe (TypeExp Exp VName)
elem_te) [TypeBase Rank o]
ts'
          addType desc doc . I.OpaqueArray (E.shapeRank shape) (entryPointTypeName ept) $
            map valueType ts
        StructType
_ -> String -> GenOpaque ()
forall a. HasCallStack => String -> a
error (String -> GenOpaque ()) -> String -> GenOpaque ()
forall a b. (a -> b) -> a -> b
$ String
"entryPointType: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> StructType -> String
forall a. Pretty a => a -> String
E.prettyString (EntryType -> StructType
E.entryType EntryType
t)

      (o, EntryPointType)
-> StateT OpaqueTypes Identity (o, EntryPointType)
forall a. a -> StateT OpaqueTypes Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (o
o, Name -> EntryPointType
I.TypeOpaque Name
desc)
  where
    doc :: Maybe a
doc = Maybe a
forall a. Maybe a
Nothing
    -- The mode of a composite is that of its most-restrictive
    -- component: consumed for a parameter, fresh for a result.
    o :: o
o = (o -> o -> o) -> o -> [o] -> o
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl o -> o -> o
forall a. Ord a => a -> a -> a
max o
forall a. Bounded a => a
minBound ([o] -> o) -> [o] -> o
forall a b. (a -> b) -> a -> b
$ (TypeBase Rank o -> o) -> [TypeBase Rank o] -> [o]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Rank o -> o
forall o shape. Bounded o => TypeBase shape o -> o
modeOf [TypeBase Rank o]
ts
    desc :: Name
desc =
      Name
-> (TypeExp Exp VName -> Name) -> Maybe (TypeExp Exp VName) -> Name
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Text -> Name
nameFromText (Text -> Name) -> Text -> Name
forall a b. (a -> b) -> a -> b
$ TypeBase () Freshness -> Text
forall a. Pretty a => a -> Text
prettyTextOneLine TypeBase () Freshness
t') TypeExp Exp VName -> Name
typeExpOpaqueName (Maybe (TypeExp Exp VName) -> Name)
-> Maybe (TypeExp Exp VName) -> Name
forall a b. (a -> b) -> a -> b
$
        EntryType -> Maybe (TypeExp Exp VName)
E.entryAscribed EntryType
t
    t' :: TypeBase () Freshness
t' = StructType -> TypeBase () NoMode
forall o. TypeBase Exp o -> TypeBase () o
E.noSizes (EntryType -> StructType
E.entryType EntryType
t) TypeBase () NoMode -> Freshness -> TypeBase () Freshness
forall dim o1 o2. TypeBase dim o1 -> o2 -> TypeBase dim o2
`E.setMode` Freshness
E.Nonfresh
    strip :: Int -> TypeBase Rank o -> TypeBase Rank o
strip Int
k (I.Array PrimType
pt (I.Rank Int
r) o
t_u) =
      TypeBase Rank (ZonkAny 0) -> Rank -> o -> TypeBase Rank o
forall shape o_unused o.
ArrayShape shape =>
TypeBase shape o_unused -> shape -> o -> TypeBase shape o
I.arrayOf (PrimType -> TypeBase Rank (ZonkAny 0)
forall shape o. PrimType -> TypeBase shape o
I.Prim PrimType
pt) (Int -> Rank
I.Rank (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k)) o
t_u
    strip Int
_ TypeBase Rank o
ts_t = TypeBase Rank o
ts_t

-- | An entry point result is either fresh or not; the IR has no freshness type,
-- so this is recorded as a 'I.Diet' where 'I.Consume' means fresh. See
-- 'I.EntryResult'. This is a bit wonky, but it is just an intermediate thing
-- during internalisation.
freshnessAsDiet :: E.Freshness -> I.Diet
freshnessAsDiet :: Freshness -> Diet
freshnessAsDiet Freshness
E.Nonfresh = Diet
I.Observe
freshnessAsDiet Freshness
E.Fresh = Diet
I.Consume

entryPoint ::
  VisibleTypes ->
  Name ->
  Maybe T.Text ->
  [(E.EntryParam, [I.Param I.DeclType])] ->
  ( E.EntryType,
    [[I.TypeBase I.Rank E.Freshness]]
  ) ->
  (I.EntryPoint, I.OpaqueTypes)
entryPoint :: VisibleTypes
-> Name
-> Maybe Text
-> [(EntryParam, [Param DeclType])]
-> (EntryType, [[TypeBase Rank Freshness]])
-> (EntryPoint, OpaqueTypes)
entryPoint VisibleTypes
types Name
name Maybe Text
doc [(EntryParam, [Param DeclType])]
params (EntryType
eret, [[TypeBase Rank Freshness]]
crets) =
  GenOpaque EntryPoint -> (EntryPoint, OpaqueTypes)
forall a. GenOpaque a -> (a, OpaqueTypes)
runGenOpaque (GenOpaque EntryPoint -> (EntryPoint, OpaqueTypes))
-> GenOpaque EntryPoint -> (EntryPoint, OpaqueTypes)
forall a b. (a -> b) -> a -> b
$
    (Name
name,,,Maybe Text
doc)
      ([EntryParam] -> EntryResult -> EntryPoint)
-> StateT OpaqueTypes Identity [EntryParam]
-> StateT OpaqueTypes Identity (EntryResult -> EntryPoint)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((EntryParam, [Param DeclType])
 -> StateT OpaqueTypes Identity EntryParam)
-> [(EntryParam, [Param DeclType])]
-> StateT OpaqueTypes Identity [EntryParam]
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 (EntryParam, [Param DeclType])
-> StateT OpaqueTypes Identity EntryParam
forall {dec}.
DeclTyped dec =>
(EntryParam, [Param dec]) -> StateT OpaqueTypes Identity EntryParam
onParam [(EntryParam, [Param DeclType])]
params
      StateT OpaqueTypes Identity (EntryResult -> EntryPoint)
-> StateT OpaqueTypes Identity EntryResult -> GenOpaque EntryPoint
forall a b.
StateT OpaqueTypes Identity (a -> b)
-> StateT OpaqueTypes Identity a -> StateT OpaqueTypes Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ( (Freshness -> EntryPointType -> EntryResult)
-> (Freshness, EntryPointType) -> EntryResult
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (Diet -> EntryPointType -> EntryResult
I.EntryResult (Diet -> EntryPointType -> EntryResult)
-> (Freshness -> Diet)
-> Freshness
-> EntryPointType
-> EntryResult
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Freshness -> Diet
freshnessAsDiet)
              ((Freshness, EntryPointType) -> EntryResult)
-> StateT OpaqueTypes Identity (Freshness, EntryPointType)
-> StateT OpaqueTypes Identity EntryResult
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VisibleTypes
-> EntryType
-> [TypeBase Rank Freshness]
-> StateT OpaqueTypes Identity (Freshness, EntryPointType)
forall o.
(Ord o, Bounded o) =>
VisibleTypes
-> EntryType -> [TypeBase Rank o] -> GenOpaque (o, EntryPointType)
entryPointType VisibleTypes
types EntryType
eret ([[TypeBase Rank Freshness]] -> [TypeBase Rank Freshness]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[TypeBase Rank Freshness]]
crets)
          )
  where
    onParam :: (EntryParam, [Param dec]) -> StateT OpaqueTypes Identity EntryParam
onParam (E.EntryParam Name
e_p EntryType
e_t, [Param dec]
ps) =
      (Diet -> EntryPointType -> EntryParam)
-> (Diet, EntryPointType) -> EntryParam
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (Name -> Diet -> EntryPointType -> EntryParam
I.EntryParam Name
e_p)
        ((Diet, EntryPointType) -> EntryParam)
-> StateT OpaqueTypes Identity (Diet, EntryPointType)
-> StateT OpaqueTypes Identity EntryParam
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VisibleTypes
-> EntryType
-> [TypeBase Rank Diet]
-> StateT OpaqueTypes Identity (Diet, EntryPointType)
forall o.
(Ord o, Bounded o) =>
VisibleTypes
-> EntryType -> [TypeBase Rank o] -> GenOpaque (o, EntryPointType)
entryPointType VisibleTypes
types EntryType
e_t ((Param dec -> TypeBase Rank Diet)
-> [Param dec] -> [TypeBase Rank Diet]
forall a b. (a -> b) -> [a] -> [b]
map (DeclType -> TypeBase Rank Diet
forall shape o.
ArrayShape shape =>
TypeBase shape o -> TypeBase Rank o
I.rankShaped (DeclType -> TypeBase Rank Diet)
-> (Param dec -> DeclType) -> Param dec -> TypeBase Rank Diet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param dec -> DeclType
forall dec. DeclTyped dec => Param dec -> DeclType
I.paramDeclType) [Param dec]
ps)