content: {-# LANGUAGE Strict #-}
  
  -- | Internalising bindings.
  module Futhark.Internalise.Bindings
    ( internaliseAttrs,
      internaliseAttr,
      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)
  
  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
    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
  
    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 params_idents (concatMap (concatMap toList) params_ts) $ \valueparams -> do
      let (certparams, valueparams') =
            first concat $ unzip $ map fixAccParams valueparams
          all_params = certparams ++ shape_params ++ concat valueparams'
      I.localScope (I.scopeOfFParams all_params) $
        substitutingVars shape_subst $ do
          let values_grouped_by_params = chunks num_param_idents valueparams'
              types_grouped_by_params = chunks num_param_idents params_ts
  
          m (certparams ++ shape_params) $
            zipWith chunkValues types_grouped_by_params values_grouped_by_params
    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)
  
  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
  