Bindings.hs

{-# LANGUAGE Strict #-}

-- | Internalising bindings.
module Futhark.Internalise.Bindings
  ( internaliseAttrs,
    internaliseAttr,
    FunParams (..),
    internaliseFParams,
    bindFParams,
    bindingFParams,
    bindingLoopParams,
    bindingLambdaParams,
    stmPat,
  )
where

import Control.Monad
import Control.Monad.Free (Free (..))
import Control.Monad.Reader
import Data.Bifunctor
import Data.Foldable (toList)
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.IR.SOACS qualified as I
import Futhark.Internalise.Monad
import Futhark.Internalise.TypesValues
import Futhark.Util
import Language.Futhark as E hiding (matchDims)

internaliseAttr :: E.AttrInfo VName -> InternaliseM I.Attr
internaliseAttr (E.AttrAtom (E.AtomName v) _) =
  pure $ I.AttrName v
internaliseAttr (E.AttrAtom (E.AtomInt x) _) =
  pure $ I.AttrInt x
internaliseAttr (E.AttrComp f attrs _) =
  I.AttrComp f <$> mapM internaliseAttr attrs

internaliseAttrs :: [E.AttrInfo VName] -> InternaliseM I.Attrs
internaliseAttrs = fmap (mconcat . map I.oneAttr) . mapM internaliseAttr

treeLike :: Tree a -> [b] -> Tree b
treeLike (Pure _) [b] = Pure b
treeLike (Pure _) _ = error "treeLike: invalid input"
treeLike (Free ls) bs = Free $ zipWith treeLike ls (chunks (map length ls) bs)

-- | The internalised parameters of a function. Computing these assigns fresh
-- names, so we split it from binding them in scope ('bindFParams'): that way a
-- caller can internalise the parameters once, and both register the function's
-- calling information and later bind the parameters for its body without
-- recomputing them.
data FunParams = FunParams
  { -- | Shape (and certificate) parameters.
    funShapeParams :: [I.FParam I.SOACS],
    -- | Value parameters, grouped as the source parameters.
    funValueParams :: [[Tree (I.FParam I.SOACS)]],
    funSubsts :: VarSubsts,
    funShapeSubst :: VarSubsts
  }

-- | All parameters, flattened; used for scoping and the 'I.FunDef'.
funAllParams :: FunParams -> [I.FParam I.SOACS]
funAllParams fps =
  funShapeParams fps ++ foldMap (foldMap toList) (funValueParams fps)

-- | Internalise a function's parameters (see 'FunParams').
internaliseFParams ::
  [E.TypeParam] ->
  [E.Pat E.ParamType] ->
  InternaliseM FunParams
internaliseFParams tparams params = do
  flattened_params <- mapM flattenPat params
  let params_idents = concat flattened_params
  params_ts <-
    internaliseParamTypes $
      map (E.unInfo . E.identType . fst) params_idents
  let num_param_idents = map length flattened_params
      shape_params = [I.Param mempty v $ I.Prim I.int64 | E.TypeParamDim v _ <- tparams]
      shape_subst = M.fromList [(I.paramName p, [I.Var $ I.paramName p]) | p <- shape_params]
  (valueparams, substs) <-
    processFlatPat params_idents (concatMap (concatMap toList) params_ts)
  let (certparams, valueparams') =
        first concat $ unzip $ map fixAccParams valueparams
      values_grouped_by_params = chunks num_param_idents valueparams'
      types_grouped_by_params = chunks num_param_idents params_ts
  pure
    FunParams
      { funShapeParams = certparams ++ shape_params,
        funValueParams =
          zipWith chunkValues types_grouped_by_params values_grouped_by_params,
        funSubsts = substs,
        funShapeSubst = shape_subst
      }
  where
    fixAccParams ps =
      first catMaybes $ unzip $ map fixAccParam ps
    fixAccParam (I.Param attrs pv (I.Acc acc ispace ts u)) =
      ( Just (I.Param attrs acc $ I.Prim I.Unit),
        I.Param attrs pv (I.Acc acc ispace ts u)
      )
    fixAccParam p = (Nothing, p)

    chunkValues ::
      [[Tree (I.TypeBase I.Shape Uniqueness)]] ->
      [[I.FParam I.SOACS]] ->
      [Tree (I.FParam I.SOACS)]
    chunkValues tss vss =
      concat $ zipWith f tss vss
      where
        f ts vs = zipWith treeLike ts (chunks (map length ts) vs)

-- | Bind already-internalised parameters (see 'internaliseFParams') in scope.
bindFParams :: FunParams -> InternaliseM a -> InternaliseM a
bindFParams fps m =
  local (\env -> env {envSubsts = funSubsts fps `M.union` envSubsts env}) $
    I.localScope (I.scopeOfFParams (funAllParams fps)) $
      substitutingVars (funShapeSubst fps) m

bindingFParams ::
  [E.TypeParam] ->
  [E.Pat E.ParamType] ->
  ([I.FParam I.SOACS] -> [[Tree (I.FParam I.SOACS)]] -> InternaliseM a) ->
  InternaliseM a
bindingFParams tparams params m = do
  fps <- internaliseFParams tparams params
  bindFParams fps $ m (funShapeParams fps) (funValueParams fps)

bindingLoopParams ::
  [E.TypeParam] ->
  E.Pat E.ParamType ->
  [I.Type] ->
  ([I.FParam I.SOACS] -> [I.FParam I.SOACS] -> InternaliseM a) ->
  InternaliseM a
bindingLoopParams tparams pat ts m = do
  pat_idents <- flattenPat pat
  pat_ts <- internaliseLoopParamType (E.patternType pat) ts

  let shape_params = [I.Param mempty v $ I.Prim I.int64 | E.TypeParamDim v _ <- tparams]
      shape_subst = M.fromList [(I.paramName p, [I.Var $ I.paramName p]) | p <- shape_params]

  bindingFlatPat pat_idents pat_ts $ \valueparams ->
    I.localScope (I.scopeOfFParams $ shape_params ++ concat valueparams) $
      substitutingVars shape_subst $
        m shape_params (concat valueparams)

bindingLambdaParams ::
  [E.Pat E.ParamType] ->
  [I.Type] ->
  ([I.LParam I.SOACS] -> InternaliseM a) ->
  InternaliseM a
bindingLambdaParams params ts m = do
  params_idents <- concat <$> mapM flattenPat params

  bindingFlatPat params_idents ts $ \params' ->
    I.localScope (I.scopeOfLParams $ concat params') $
      m (concat params')

type Params t = [I.Param t]

processFlatPat ::
  (Show t) =>
  [(E.Ident ParamType, [E.AttrInfo VName])] ->
  [t] ->
  InternaliseM ([Params t], VarSubsts)
processFlatPat x y = processFlatPat' [] x y
  where
    processFlatPat' pat [] _ = do
      let (vs, substs) = unzip pat
      pure (reverse vs, M.fromList substs)
    processFlatPat' pat ((p, attrs) : rest) ts = do
      attrs' <- internaliseAttrs attrs
      (ps, rest_ts) <- handleMapping attrs' ts <$> internaliseBindee p
      processFlatPat'
        ((ps, (E.identName p, map (I.Var . I.paramName) ps)) : pat)
        rest
        rest_ts

    handleMapping _ ts [] =
      ([], ts)
    handleMapping attrs (t : ts) (r : rs) =
      let (ps, ts') = handleMapping attrs ts rs
       in (I.Param attrs r t : ps, ts')
    handleMapping _ [] _ =
      error $ "handleMapping: insufficient identifiers in pattern.\n" ++ show (x, y)

    internaliseBindee :: E.Ident E.ParamType -> InternaliseM [VName]
    internaliseBindee bindee = do
      let name = E.identName bindee
      case internalisedTypeSize $ E.unInfo $ E.identType bindee of
        1 -> pure [name]
        n -> replicateM n $ newName name

bindingFlatPat ::
  (Show t) =>
  [(E.Ident E.ParamType, [E.AttrInfo VName])] ->
  [t] ->
  ([Params t] -> InternaliseM a) ->
  InternaliseM a
bindingFlatPat idents ts m = do
  (ps, substs) <- processFlatPat idents ts
  local (\env -> env {envSubsts = substs `M.union` envSubsts env}) $
    m ps

-- | Flatten a pattern.  Returns a list of identifiers.
flattenPat :: (MonadFreshNames m) => E.Pat (TypeBase Size u) -> m [(E.Ident (TypeBase Size u), [E.AttrInfo VName])]
flattenPat = flattenPat'
  where
    flattenPat' (E.PatParens p _) =
      flattenPat' p
    flattenPat' (E.PatAttr attr p _) =
      map (second (attr :)) <$> flattenPat' p
    flattenPat' (E.Wildcard t loc) = do
      name <- newVName "nameless"
      flattenPat' $ E.Id name t loc
    flattenPat' (E.Id v (Info t) loc) =
      pure [(E.Ident v (Info t) loc, mempty)]
    flattenPat' (E.TuplePat [] loc) =
      flattenPat' (E.Wildcard (Info $ E.Scalar $ E.Record mempty) loc)
    flattenPat' (E.RecordPat [] loc) =
      flattenPat' (E.Wildcard (Info $ E.Scalar $ E.Record mempty) loc)
    flattenPat' (E.TuplePat pats _) =
      concat <$> mapM flattenPat' pats
    flattenPat' (E.RecordPat fs loc) =
      flattenPat' $ E.TuplePat (map snd $ sortFields $ M.fromList $ map (first unLoc) fs) loc
    flattenPat' (E.PatAscription p _ _) =
      flattenPat' p
    flattenPat' (E.PatLit _ t loc) =
      flattenPat' $ E.Wildcard t loc
    flattenPat' (E.PatConstr _ _ ps _) =
      concat <$> mapM flattenPat' ps

stmPat ::
  E.Pat E.ParamType ->
  [I.Type] ->
  ([VName] -> InternaliseM a) ->
  InternaliseM a
stmPat pat ts m = do
  pat' <- flattenPat pat
  bindingFlatPat pat' ts $ m . map I.paramName . concat