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

import Control.Monad
import Control.Monad.State
import Data.List (find)
import Data.Map qualified as M
import Futhark.IR qualified as I
import Futhark.Internalise.TypesValues (internalisedTypeSize)
import Futhark.Util.Pretty (prettyText, prettyTextOneLine)
import Language.Futhark qualified as E hiding (TypeArg)
import Language.Futhark.Core (Name, Uniqueness (..), VName, nameFromText)
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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (FileModule -> [TypeBind]
modTypes forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd)
  where
    modTypes :: FileModule -> [TypeBind]
modTypes = forall {f :: * -> *} {vn}. ProgBase f vn -> [TypeBindBase f vn]
progTypes forall b c a. (b -> c) -> (a -> b) -> a -> c
. FileModule -> Prog
E.fileProg
    progTypes :: ProgBase f vn -> [TypeBindBase f vn]
progTypes = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall {f :: * -> *} {vn}. DecBase f vn -> [TypeBindBase f vn]
decTypes forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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.Info VName)
findType :: VName -> VisibleTypes -> Maybe (TypeExp Info VName)
findType VName
v (VisibleTypes [TypeBind]
ts) = forall (f :: * -> *) vn. TypeBindBase f vn -> TypeExp f vn
E.typeExp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== VName
v) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) vn. TypeBindBase f vn -> vn
E.typeAlias) [TypeBind]
ts

valueType :: I.TypeBase I.Rank Uniqueness -> I.ValueType
valueType :: TypeBase Rank Uniqueness -> 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 Uniqueness
_) = Signedness -> Rank -> PrimType -> ValueType
I.ValueType Signedness
I.Signed Rank
rank PrimType
pt
valueType I.Acc {} = forall a. HasCallStack => [Char] -> a
error [Char]
"valueType Acc"
valueType I.Mem {} = forall a. HasCallStack => [Char] -> a
error [Char]
"valueType Mem"

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

rootType :: E.TypeExp E.Info VName -> E.TypeExp E.Info VName
rootType :: TypeExp Info VName -> TypeExp Info VName
rootType (E.TEApply TypeExp Info VName
te E.TypeArgExpSize {} SrcLoc
_) = TypeExp Info VName -> TypeExp Info VName
rootType TypeExp Info VName
te
rootType (E.TEUnique TypeExp Info VName
te SrcLoc
_) = TypeExp Info VName -> TypeExp Info VName
rootType TypeExp Info VName
te
rootType TypeExp Info VName
te = TypeExp Info VName
te

typeExpOpaqueName :: E.TypeExp E.Info VName -> Name
typeExpOpaqueName :: TypeExp Info VName -> Name
typeExpOpaqueName = TypeExp Info VName -> Name
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeExp Info VName -> TypeExp Info VName
rootType
  where
    f :: TypeExp Info VName -> Name
f (E.TEArray SizeExp Info VName
_ TypeExp Info VName
te SrcLoc
_) =
      let (Int
d, TypeExp Info VName
te') = TypeExp Info VName -> (Int, TypeExp Info VName)
withoutDims TypeExp Info VName
te
       in Name
"arr_" forall a. Semigroup a => a -> a -> a
<> TypeExp Info VName -> Name
typeExpOpaqueName TypeExp Info VName
te' forall a. Semigroup a => a -> a -> a
<> Name
"_" forall a. Semigroup a => a -> a -> a
<> Text -> Name
nameFromText (forall a. Pretty a => a -> Text
prettyText (Int
1 forall a. Num a => a -> a -> a
+ Int
d)) forall a. Semigroup a => a -> a -> a
<> Name
"d"
    f TypeExp Info VName
te = Text -> Name
nameFromText forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Text
prettyTextOneLine TypeExp Info VName
te

type GenOpaque = State I.OpaqueTypes

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

addType :: Name -> I.OpaqueType -> GenOpaque ()
addType :: Name -> OpaqueType -> GenOpaque ()
addType Name
s OpaqueType
t = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (forall a. Semigroup a => a -> a -> a
<> [(Name, OpaqueType)] -> OpaqueTypes
I.OpaqueTypes [(Name
s, OpaqueType
t)])

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

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

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

entryPointType ::
  VisibleTypes ->
  E.EntryType ->
  [I.TypeBase I.Rank Uniqueness] ->
  GenOpaque (Uniqueness, I.EntryPointType)
entryPointType :: VisibleTypes
-> EntryType
-> [TypeBase Rank Uniqueness]
-> GenOpaque (Uniqueness, EntryPointType)
entryPointType VisibleTypes
types EntryType
t [TypeBase Rank Uniqueness]
ts
  | E.Scalar (E.Prim E.Unsigned {}) <- EntryType -> StructType
E.entryType EntryType
t,
    [I.Prim PrimType
ts0] <- [TypeBase Rank Uniqueness]
ts =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (Uniqueness
u, ValueType -> EntryPointType
I.TypeTransparent 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 ()
_ Uniqueness
_ Shape Size
_ (E.Prim E.Unsigned {}) <- EntryType -> StructType
E.entryType EntryType
t,
    [I.Array PrimType
ts0 Rank
r Uniqueness
_] <- [TypeBase Rank Uniqueness]
ts =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (Uniqueness
u, ValueType -> EntryPointType
I.TypeTransparent 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 Uniqueness]
ts =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (Uniqueness
u, ValueType -> EntryPointType
I.TypeTransparent 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 ()
_ Uniqueness
_ Shape Size
_ E.Prim {} <- EntryType -> StructType
E.entryType EntryType
t,
    [I.Array PrimType
ts0 Rank
r Uniqueness
_] <- [TypeBase Rank Uniqueness]
ts =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (Uniqueness
u, ValueType -> EntryPointType
I.TypeTransparent 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)
          | Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Bool
null Map Name StructType
fs ->
              let fs' :: [(Name, EntryType)]
fs' = VisibleTypes
-> Map Name StructType
-> Maybe (TypeExp Info VName)
-> [(Name, EntryType)]
recordFields VisibleTypes
types Map Name StructType
fs forall a b. (a -> b) -> a -> b
$ EntryType -> Maybe (TypeExp Info VName)
E.entryAscribed EntryType
t
               in Name -> OpaqueType -> GenOpaque ()
addType Name
desc forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Name, EntryPointType)] -> OpaqueType
I.OpaqueRecord forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VisibleTypes
-> [(Name, EntryType)]
-> [TypeBase Rank Uniqueness]
-> GenOpaque [(Name, EntryPointType)]
opaqueRecord VisibleTypes
types [(Name, EntryType)]
fs' [TypeBase Rank Uniqueness]
ts
        StructType
_ -> Name -> OpaqueType -> GenOpaque ()
addType Name
desc forall a b. (a -> b) -> a -> b
$ [ValueType] -> OpaqueType
I.OpaqueType forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map TypeBase Rank Uniqueness -> ValueType
valueType [TypeBase Rank Uniqueness]
ts
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (Uniqueness
u, Name -> EntryPointType
I.TypeOpaque Name
desc)
  where
    u :: Uniqueness
u = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall a. Ord a => a -> a -> a
max Uniqueness
Nonunique forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall shape. TypeBase shape Uniqueness -> Uniqueness
I.uniqueness [TypeBase Rank Uniqueness]
ts
    desc :: Name
desc =
      forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Text -> Name
nameFromText forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> Text
prettyTextOneLine TypeBase () ()
t') TypeExp Info VName -> Name
typeExpOpaqueName forall a b. (a -> b) -> a -> b
$
        EntryType -> Maybe (TypeExp Info VName)
E.entryAscribed EntryType
t
    t' :: TypeBase () ()
t' = forall as. TypeBase Size as -> TypeBase () as
E.noSizes (EntryType -> StructType
E.entryType EntryType
t) forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`E.setUniqueness` Uniqueness
Nonunique

entryPoint ::
  VisibleTypes ->
  Name ->
  [(E.EntryParam, [I.Param I.DeclType])] ->
  ( E.EntryType,
    [[I.TypeBase I.Rank I.Uniqueness]]
  ) ->
  (I.EntryPoint, I.OpaqueTypes)
entryPoint :: VisibleTypes
-> Name
-> [(EntryParam, [Param DeclType])]
-> (EntryType, [[TypeBase Rank Uniqueness]])
-> (EntryPoint, OpaqueTypes)
entryPoint VisibleTypes
types Name
name [(EntryParam, [Param DeclType])]
params (EntryType
eret, [[TypeBase Rank Uniqueness]]
crets) =
  forall a. GenOpaque a -> (a, OpaqueTypes)
runGenOpaque forall a b. (a -> b) -> a -> b
$
    (Name
name,,)
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {dec}.
DeclTyped dec =>
(EntryParam, [Param dec]) -> StateT OpaqueTypes Identity EntryParam
onParam [(EntryParam, [Param DeclType])]
params
      forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ( forall a b. (a -> b) -> [a] -> [b]
map (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Uniqueness -> EntryPointType -> EntryResult
I.EntryResult)
              forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case ( forall dim as. TypeBase dim as -> Maybe [TypeBase dim as]
E.isTupleRecord forall a b. (a -> b) -> a -> b
$ EntryType -> StructType
E.entryType EntryType
eret,
                         EntryType -> Maybe (TypeExp Info VName)
E.entryAscribed EntryType
eret
                       ) of
                (Just [StructType]
ts, Just (E.TETuple [TypeExp Info VName]
e_ts SrcLoc
_)) ->
                  forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (VisibleTypes
-> EntryType
-> [TypeBase Rank Uniqueness]
-> GenOpaque (Uniqueness, EntryPointType)
entryPointType VisibleTypes
types) (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith StructType -> Maybe (TypeExp Info VName) -> EntryType
E.EntryType [StructType]
ts (forall a b. (a -> b) -> [a] -> [b]
map forall a. a -> Maybe a
Just [TypeExp Info VName]
e_ts)) [[TypeBase Rank Uniqueness]]
crets
                (Just [StructType]
ts, Maybe (TypeExp Info VName)
Nothing) ->
                  forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (VisibleTypes
-> EntryType
-> [TypeBase Rank Uniqueness]
-> GenOpaque (Uniqueness, EntryPointType)
entryPointType VisibleTypes
types) (forall a b. (a -> b) -> [a] -> [b]
map (StructType -> Maybe (TypeExp Info VName) -> EntryType
`E.EntryType` forall a. Maybe a
Nothing) [StructType]
ts) [[TypeBase Rank Uniqueness]]
crets
                (Maybe [StructType], Maybe (TypeExp Info VName))
_ ->
                  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VisibleTypes
-> EntryType
-> [TypeBase Rank Uniqueness]
-> GenOpaque (Uniqueness, EntryPointType)
entryPointType VisibleTypes
types EntryType
eret (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[TypeBase Rank Uniqueness]]
crets)
          )
  where
    onParam :: (EntryParam, [Param dec]) -> StateT OpaqueTypes Identity EntryParam
onParam (E.EntryParam Name
e_p EntryType
e_t, [Param dec]
ps) =
      forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (Name -> Uniqueness -> EntryPointType -> EntryParam
I.EntryParam Name
e_p)
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VisibleTypes
-> EntryType
-> [TypeBase Rank Uniqueness]
-> GenOpaque (Uniqueness, EntryPointType)
entryPointType VisibleTypes
types EntryType
e_t (forall a b. (a -> b) -> [a] -> [b]
map (forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase Rank u
I.rankShaped forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. DeclTyped dec => Param dec -> DeclType
I.paramDeclType) [Param dec]
ps)