Bindings.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Safe #-}
-- | Internalising bindings.
module Futhark.Internalise.Bindings
  (
    bindingParams
  , bindingLoopParams
  , bindingLambdaParams
  , stmPattern
  )
  where

import Control.Monad.State  hiding (mapM)
import Control.Monad.Reader hiding (mapM)

import qualified Data.Map.Strict as M

import Language.Futhark as E hiding (matchDims)
import qualified Futhark.IR.SOACS as I
import Futhark.MonadFreshNames
import Futhark.Internalise.Monad
import Futhark.Internalise.TypesValues
import Futhark.Util

bindingParams :: [E.TypeParam] -> [E.Pattern]
              -> ([I.FParam] -> [[I.FParam]] -> InternaliseM a)
              -> InternaliseM a
bindingParams tparams params m = do
  flattened_params <- mapM flattenPattern params
  let params_idents = concat flattened_params
  params_ts <-
    internaliseParamTypes $
    map (flip E.setAliases () . E.unInfo . E.identType) params_idents
  let num_param_idents = map length flattened_params
      num_param_ts = map (sum . map length) $ chunks num_param_idents params_ts

  let shape_params = [ I.Param v $ I.Prim I.int32 | E.TypeParamDim v _ <- tparams ]
      shape_subst = M.fromList [ (I.paramName p, [I.Var $ I.paramName p]) | p <- shape_params ]
  bindingFlatPattern params_idents (concat params_ts) $ \valueparams ->
    I.localScope (I.scopeOfFParams $ shape_params++concat valueparams) $
    substitutingVars shape_subst $ m shape_params $
    chunks num_param_ts (concat valueparams)

bindingLoopParams :: [E.TypeParam] -> E.Pattern
                  -> ([I.FParam] -> [I.FParam] -> InternaliseM a)
                  -> InternaliseM a
bindingLoopParams tparams pat m = do
  pat_idents <- flattenPattern pat
  pat_ts <- internaliseLoopParamType (E.patternStructType pat)

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

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

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

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

processFlatPattern :: Show t => [E.Ident] -> [t]
                   -> InternaliseM ([[I.Param t]], VarSubstitutions)
processFlatPattern x y = processFlatPattern' [] x y
  where
    processFlatPattern' pat []       _  = do
      let (vs, substs) = unzip pat
          substs' = M.fromList substs
          idents = reverse vs
      return (idents, substs')

    processFlatPattern' pat (p:rest) ts = do
      (ps, subst, rest_ts) <- handleMapping ts <$> internaliseBindee p
      processFlatPattern' ((ps, (E.identName p, map (I.Var . I.paramName) subst)) : pat) rest rest_ts

    handleMapping ts [] =
      ([], [], ts)
    handleMapping ts (r:rs) =
        let (ps, reps, ts')    = handleMapping' ts r
            (pss, repss, ts'') = handleMapping ts' rs
        in (ps++pss, reps:repss, ts'')

    handleMapping' (t:ts) vname =
      let v' = I.Param vname t
      in ([v'], v', ts)
    handleMapping' [] _ =
      error $ "processFlatPattern: insufficient identifiers in pattern." ++ show (x, y)

    internaliseBindee :: E.Ident -> InternaliseM [VName]
    internaliseBindee bindee = do
      let name = E.identName bindee
      n <- internalisedTypeSize $ flip E.setAliases () $ E.unInfo $ E.identType bindee
      case n of
        1 -> return [name]
        _ -> replicateM n $ newVName $ baseString name

bindingFlatPattern :: Show t => [E.Ident] -> [t]
                   -> ([[I.Param t]] -> InternaliseM a)
                   -> InternaliseM a
bindingFlatPattern idents ts m = do
  (ps, substs) <- processFlatPattern idents ts
  local (\env -> env { envSubsts = substs `M.union` envSubsts env}) $
    m ps

-- | Flatten a pattern.  Returns a list of identifiers.  The
-- structural type of each identifier is returned separately.
flattenPattern :: MonadFreshNames m => E.Pattern -> m [E.Ident]
flattenPattern = flattenPattern'
  where flattenPattern' (E.PatternParens p _) =
          flattenPattern' p
        flattenPattern' (E.Wildcard t loc) = do
          name <- newVName "nameless"
          flattenPattern' $ E.Id name t loc
        flattenPattern' (E.Id v (Info t) loc) =
          return [E.Ident v (Info t) loc]
        -- XXX: treat empty tuples and records as bool.
        flattenPattern' (E.TuplePattern [] loc) =
          flattenPattern' (E.Wildcard (Info $ E.Scalar $ E.Prim E.Bool) loc)
        flattenPattern' (E.RecordPattern [] loc) =
          flattenPattern' (E.Wildcard (Info $ E.Scalar $ E.Prim E.Bool) loc)
        flattenPattern' (E.TuplePattern pats _) =
          concat <$> mapM flattenPattern' pats
        flattenPattern' (E.RecordPattern fs loc) =
          flattenPattern' $ E.TuplePattern (map snd $ sortFields $ M.fromList fs) loc
        flattenPattern' (E.PatternAscription p _ _) =
          flattenPattern' p
        flattenPattern' (E.PatternLit _ t loc) =
          flattenPattern' $ E.Wildcard t loc
        flattenPattern' (E.PatternConstr _ _ ps _) =
          concat <$> mapM flattenPattern' ps

stmPattern :: E.Pattern -> [I.Type]
           -> ([VName] -> InternaliseM a)
           -> InternaliseM a
stmPattern pat ts m = do
  pat' <- flattenPattern pat
  let addShapeStms l =
        m (map I.paramName $ concat l)
  bindingFlatPattern pat' ts addShapeStms