Entry.hs

-- | 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 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, Uniqueness (..), 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 = VisibleTypes . foldMap (modTypes . snd)
  where
    modTypes = progTypes . E.fileProg
    progTypes = foldMap decTypes . E.progDecs
    decTypes (E.TypeDec tb) = [tb]
    decTypes _ = []

findType :: VName -> VisibleTypes -> Maybe (E.TypeExp E.Exp VName)
findType v (VisibleTypes ts) = E.typeExp <$> find ((== v) . E.typeAlias) ts

valueType :: I.TypeBase I.Rank Uniqueness -> I.ValueType
valueType (I.Prim pt) = I.ValueType I.Signed (I.Rank 0) pt
valueType (I.Array pt rank _) = I.ValueType I.Signed rank pt
valueType I.Acc {} = error "valueType Acc"
valueType I.Mem {} = error "valueType Mem"

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

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

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

type GenOpaque = State I.OpaqueTypes

runGenOpaque :: GenOpaque a -> (a, I.OpaqueTypes)
runGenOpaque = flip runState mempty

addType :: Name -> I.OpaqueType -> GenOpaque ()
addType name t = modify $ \(I.OpaqueTypes ts) ->
  case find ((== name) . fst) ts of
    Just (_, t')
      | t /= t' ->
          error . unlines $
            [ "Duplicate definition of entry point type " <> E.prettyString name,
              show t,
              show t'
            ]
    _ -> I.OpaqueTypes ts <> I.OpaqueTypes [(name, t)]

isRecord ::
  VisibleTypes ->
  E.TypeExp E.Exp VName ->
  Maybe (M.Map Name (E.TypeExp E.Exp VName))
isRecord _ (E.TERecord fs _) = Just $ M.fromList $ map (first unLoc) fs
isRecord _ (E.TETuple fs _) = Just $ E.tupleFields fs
isRecord types (E.TEVar v _) = isRecord types =<< findType (E.qualLeaf v) types
isRecord _ _ = Nothing

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

opaqueRecord ::
  VisibleTypes ->
  [(Name, E.EntryType)] ->
  [I.TypeBase I.Rank Uniqueness] ->
  GenOpaque [(Name, I.EntryPointType)]
opaqueRecord _ [] _ = pure []
opaqueRecord types ((f, t) : fs) ts = do
  let (f_ts, ts') = splitAt (internalisedTypeSize $ E.entryType t) ts
  f' <- opaqueField t f_ts
  ((f, f') :) <$> opaqueRecord types fs ts'
  where
    opaqueField e_t i_ts = snd <$> entryPointType types e_t i_ts

opaqueRecordArray ::
  VisibleTypes ->
  Int ->
  [(Name, E.EntryType)] ->
  [I.TypeBase I.Rank Uniqueness] ->
  GenOpaque [(Name, I.EntryPointType)]
opaqueRecordArray _ _ [] _ = pure []
opaqueRecordArray types rank ((f, t) : fs) ts = do
  let (f_ts, ts') = splitAt (internalisedTypeSize $ E.entryType t) ts
  f' <- opaqueField t f_ts
  ((f, f') :) <$> opaqueRecordArray types rank fs ts'
  where
    opaqueField (E.EntryType e_t _) i_ts =
      snd <$> entryPointType types (E.EntryType e_t' Nothing) i_ts
      where
        e_t' = E.arrayOf (E.Shape (replicate rank $ E.anySize 0)) e_t

isSum ::
  VisibleTypes ->
  E.TypeExp E.Exp VName ->
  Maybe (M.Map Name [E.TypeExp E.Exp VName])
isSum _ (E.TESum cs _) = Just $ M.fromList cs
isSum types (E.TEVar v _) = isSum types =<< findType (E.qualLeaf v) types
isSum _ _ = Nothing

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

opaqueSum ::
  VisibleTypes ->
  [(Name, ([E.EntryType], [Int]))] ->
  [I.TypeBase I.Rank Uniqueness] ->
  GenOpaque [(Name, [(I.EntryPointType, [Int])])]
opaqueSum types cs ts = mapM (traverse f) cs
  where
    f (ets, is) = do
      let ns = map (internalisedTypeSize . E.entryType) ets
          is' = chunks ns is
      ets' <- map snd <$> zipWithM (entryPointType types) ets (map (map (ts !!)) is')
      pure . zip ets' $
        if length cs == 1
          then is'
          else map (map (+ 1)) is' -- Adjust for tag.

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

elemTypeExp :: E.TypeExp E.Exp VName -> Maybe (E.TypeExp E.Exp VName)
elemTypeExp (E.TEArray _ te _) = Just te
elemTypeExp (E.TEUnique te _) = elemTypeExp te
elemTypeExp (E.TEParens te _) = elemTypeExp te
elemTypeExp _ = Nothing

rowTypeExp :: Int -> E.TypeExp E.Exp VName -> Maybe (E.TypeExp E.Exp VName)
rowTypeExp 0 te = Just te
rowTypeExp r te = rowTypeExp (r - 1) =<< elemTypeExp te

entryPointType ::
  VisibleTypes ->
  E.EntryType ->
  [I.TypeBase I.Rank Uniqueness] ->
  GenOpaque (Uniqueness, I.EntryPointType)
entryPointType types t ts
  | E.Scalar (E.Prim E.Unsigned {}) <- E.entryType t,
    [I.Prim ts0] <- ts =
      pure (u, I.TypeTransparent $ I.ValueType I.Unsigned (I.Rank 0) ts0)
  | E.Array _ _ (E.Prim E.Unsigned {}) <- E.entryType t,
    [I.Array ts0 r _] <- ts =
      pure (u, I.TypeTransparent $ I.ValueType I.Unsigned r ts0)
  | E.Scalar E.Prim {} <- E.entryType t,
    [I.Prim ts0] <- ts =
      pure (u, I.TypeTransparent $ I.ValueType I.Signed (I.Rank 0) ts0)
  | E.Array _ _ E.Prim {} <- E.entryType t,
    [I.Array ts0 r _] <- ts =
      pure (u, I.TypeTransparent $ I.ValueType I.Signed r ts0)
  | otherwise = do
      case E.entryType t of
        E.Scalar (E.Record fs)
          | not $ null fs -> do
              let fs' = recordFields types fs $ E.entryAscribed t
              addType desc . I.OpaqueRecord =<< opaqueRecord types fs' ts
        E.Scalar (E.Sum cs) -> do
          let (_, places) = internaliseSumTypeRep cs
              cs' = sumConstrs types cs $ E.entryAscribed t
              cs'' = zip (map fst cs') (zip (map snd cs') (map snd places))
              ts' = if length cs == 1 then ts else drop 1 ts
          addType desc . I.OpaqueSum (map valueType ts)
            =<< opaqueSum types cs'' ts'
        E.Array _ shape (E.Record fs)
          | not $ null fs -> do
              let fs' = recordFields types fs $ E.entryAscribed t
                  rank = E.shapeRank shape
                  ts' = map (strip rank) ts
                  record_t = E.Scalar (E.Record fs)
                  record_te = rowTypeExp rank =<< E.entryAscribed t
              ept <- snd <$> entryPointType types (E.EntryType record_t record_te) ts'
              addType desc . I.OpaqueRecordArray rank (entryPointTypeName ept)
                =<< opaqueRecordArray types rank fs' ts
        E.Array _ shape et -> do
          let ts' = map (strip (E.shapeRank shape)) ts
              rank = E.shapeRank shape
              elem_te = rowTypeExp rank =<< E.entryAscribed t
          ept <- snd <$> entryPointType types (E.EntryType (E.Scalar et) elem_te) ts'
          addType desc . I.OpaqueArray (E.shapeRank shape) (entryPointTypeName ept) $
            map valueType ts
        _ -> addType desc $ I.OpaqueType $ map valueType ts

      pure (u, I.TypeOpaque desc)
  where
    u = foldl max Nonunique $ map I.uniqueness ts
    desc =
      maybe (nameFromText $ prettyTextOneLine t') typeExpOpaqueName $
        E.entryAscribed t
    t' = E.noSizes (E.entryType t) `E.setUniqueness` Nonunique
    strip k (I.Array pt (I.Rank r) t_u) =
      I.arrayOf (I.Prim pt) (I.Rank (r - k)) t_u
    strip _ ts_t = ts_t

entryPoint ::
  VisibleTypes ->
  Name ->
  [(E.EntryParam, [I.Param I.DeclType])] ->
  ( E.EntryType,
    [[I.TypeBase I.Rank I.Uniqueness]]
  ) ->
  (I.EntryPoint, I.OpaqueTypes)
entryPoint types name params (eret, crets) =
  runGenOpaque $
    (name,,)
      <$> mapM onParam params
      <*> ( uncurry I.EntryResult
              <$> entryPointType types eret (concat crets)
          )
  where
    onParam (E.EntryParam e_p e_t, ps) =
      uncurry (I.EntryParam e_p)
        <$> entryPointType types e_t (map (I.rankShaped . I.paramDeclType) ps)