{-# LANGUAGE TypeFamilies #-}

-- | Provides last-use analysis for Futhark programs.
module Futhark.Analysis.LastUse
  ( LastUseMap,
    LastUse,
    Used,
    analyseGPUMem,
    analyseSeqMem,
  )
where

import Control.Monad.Reader
import Data.Bifunctor (bimap, first)
import Data.Function ((&))
import Data.Map (Map)
import Data.Map qualified as M
import Data.Tuple
import Futhark.Analysis.Alias (aliasAnalysis)
import Futhark.IR.Aliases
import Futhark.IR.GPUMem
import Futhark.IR.SeqMem

-- | 'LastUseMap' tells which names were last used in a given statement.
-- Statements are uniquely identified by the 'VName' of the first value
-- parameter in the statement pattern. 'Names' is the set of names last used.
type LastUseMap = Map VName Names

-- | 'LastUse' is a mapping from a 'VName' to the statement identifying it's
-- last use. 'LastUseMap' is the inverse of 'LastUse'.
type LastUse = Map VName VName

-- | 'Used' is the set of 'VName' that were used somewhere in a
-- statement, body or otherwise.
type Used = Names

type LastUseOp rep =
  VName -> (LastUse, Used) -> Op (Aliases rep) -> LastUseM rep

newtype Env rep = Env {envLastUseOp :: LastUseOp rep}

type LastUseM rep = Reader (Env rep) (LastUse, Used)

analyseGPUMem :: Prog GPUMem -> (LastUseMap, Used)
analyseGPUMem = analyseProg analyseGPUOp

analyseSeqMem :: Prog SeqMem -> (LastUseMap, Used)
analyseSeqMem = analyseProg analyseSeqOp

analyseGPUOp :: LastUseOp GPUMem
analyseGPUOp pat_name (lumap, used) (Alloc se sp) = do
  let nms = (freeIn se <> freeIn sp) `namesSubtract` used
  pure (insertNames pat_name nms lumap, used <> nms)
analyseGPUOp pat_name (lumap, used) (Inner (SizeOp sop)) = do
  let nms = freeIn sop `namesSubtract` used
  pure (insertNames pat_name nms lumap, used <> nms)
analyseGPUOp _ (lumap, used) (Inner (OtherOp ())) =
  pure (lumap, used)
analyseGPUOp pat_name (lumap, used) (Inner (SegOp (SegMap lvl _ tps body))) = do
  (lumap', used') <- analyseKernelBody (lumap, used) body
  let nms = (freeIn lvl <> freeIn tps) `namesSubtract` used'
  pure (insertNames pat_name nms lumap', used' <> nms)
analyseGPUOp pat_name (lumap, used) (Inner (SegOp (SegRed lvl _ binops tps body))) =
  segOpHelper pat_name lumap used lvl binops tps body
analyseGPUOp pat_name (lumap, used) (Inner (SegOp (SegScan lvl _ binops tps body))) =
  segOpHelper pat_name lumap used lvl binops tps body
analyseGPUOp pat_name (lumap, used) (Inner (SegOp (SegHist lvl _ binops tps body))) = do
  (lumap', used') <- foldM analyseHistOp (lumap, used) $ reverse binops
  (lumap'', used'') <- analyseKernelBody (lumap', used') body
  let nms = (freeIn lvl <> freeIn tps) `namesSubtract` used''
  pure (insertNames pat_name nms lumap'', used'' <> nms)
analyseGPUOp pat_name (lumap, used) (Inner (GPUBody ts body)) = do
  (lumap', used') <- analyseBody lumap used body
  let nms = freeIn ts
  pure (insertNames pat_name nms lumap', used' <> nms)

segOpHelper ::
  (FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
  VName ->
  LastUse ->
  Used ->
  SegLevel ->
  [SegBinOp (Aliases rep)] ->
  [Type] ->
  KernelBody (Aliases rep) ->
  LastUseM rep
segOpHelper pat_name lumap used lvl binops tps body = do
  (lumap', used') <- foldM analyseSegBinOp (lumap, used) $ reverse binops
  (lumap'', used'') <- analyseKernelBody (lumap', used') body
  let nms = (freeIn lvl <> freeIn tps) `namesSubtract` used''
  pure (insertNames pat_name nms lumap'', used'' <> nms)

analyseSeqOp :: LastUseOp SeqMem
analyseSeqOp pat_name (lumap, used) (Alloc se sp) = do
  let nms = (freeIn se <> freeIn sp) `namesSubtract` used
  pure (insertNames pat_name nms lumap, used <> nms)
analyseSeqOp _ (lumap, used) (Inner ()) =
  pure (lumap, used)

-- | Analyses a program to return a last-use map, mapping each simple statement
-- in the program to the values that were last used within that statement, and
-- the set of all `VName` that were used inside.
analyseProg :: (CanBeAliased (Op rep), Mem rep inner) => LastUseOp rep -> Prog rep -> (LastUseMap, Used)
analyseProg onOp prog =
  runReader helper (Env onOp)
  where
    helper = do
      let bound_in_consts =
            progConsts prog
              & concatMap (patNames . stmPat)
              & namesFromList
          prog_alias = aliasAnalysis prog
          consts = progConsts prog_alias
          funs = progFuns prog_alias
      (consts_lu, consts_used) <- analyseStms mempty mempty consts
      (lus, used) <- mconcat <$> mapM (analyseFun mempty bound_in_consts) funs
      pure (flipMap $ consts_lu <> lus, consts_used <> used)

analyseFun :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => LastUse -> Used -> FunDef (Aliases rep) -> LastUseM rep
analyseFun lumap used fun = do
  (lumap', used') <- analyseBody lumap used $ funDefBody fun
  pure (lumap', used' <> freeIn (funDefParams fun))

analyseStms :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => LastUse -> Used -> Stms (Aliases rep) -> LastUseM rep
analyseStms lumap used stms = foldM analyseStm (lumap, used) $ reverse $ stmsToList stms

analyseStm :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => (LastUse, Used) -> Stm (Aliases rep) -> LastUseM rep
analyseStm (lumap0, used0) (Let pat _ e) = do
  let (lumap', used') = patElems pat & foldl helper (lumap0, used0)
  analyseExp (lumap', used') e
  where
    helper (lumap_acc, used_acc) (PatElem name (aliases, _)) =
      -- Any aliases of `name` should have the same last-use as `name`
      ( case M.lookup name lumap_acc of
          Just name' ->
            insertNames name' (unAliases aliases) lumap_acc
          Nothing -> lumap_acc,
        used_acc <> unAliases aliases
      )

    addAliasesFromBodyRes (lumap_acc, used_acc) (PatElem {}, Constant _) = (lumap_acc, used_acc)
    addAliasesFromBodyRes (lumap_acc, used_acc) (PatElem name _, Var body_res) =
      -- Any aliases of `name` should have the same last-use as `name`
      ( case M.lookup name lumap_acc of
          Just name' ->
            insertNames name' (oneName body_res) lumap_acc
          Nothing -> lumap_acc,
        used_acc <> oneName body_res
      )

    pat_name = patElemName $ head $ patElems pat

    analyseExp (lumap, used) (BasicOp _) = do
      let nms = freeIn e `namesSubtract` used
      pure (insertNames pat_name nms lumap, used <> nms)
    analyseExp (lumap, used) (Apply _ args _ _) = do
      let nms = freeIn $ map fst args
      pure (insertNames pat_name nms lumap, used <> nms)
    analyseExp (lumap, used) (Match ses cases defbody dec) = do
      (lumap_cases, used_cases) <-
        bimap mconcat mconcat . unzip
          <$> mapM (analyseBody lumap used . caseBody) cases
      (lumap_defbody, used_defbody) <- analyseBody lumap used defbody
      let (lumap', used') =
            (lumap_defbody <> lumap_cases, used_cases <> used_defbody)
              & flip (foldl addAliasesFromBodyRes) (zip (patElems pat) (map resSubExp $ bodyResult defbody))
          nms = (freeIn ses <> freeIn dec) `namesSubtract` used'
      pure
        ( insertNames pat_name nms lumap',
          used' <> nms
        )
    analyseExp (lumap, used) (DoLoop merge form body) = do
      let (lumap', used') =
            zip (patElems pat) (map resSubExp $ bodyResult body)
              & foldl addAliasesFromBodyRes (lumap, used)
      (lumap'', used'') <- analyseBody lumap' used' body
      let nms = (freeIn merge <> freeIn form) `namesSubtract` used''
      pure (insertNames pat_name nms lumap'', used'' <> nms)
    analyseExp (lumap, used) (Op op) = do
      onOp <- asks envLastUseOp
      onOp pat_name (lumap, used) op
    analyseExp (lumap, used) (WithAcc _ l) =
      analyseLambda (lumap, used) l

analyseBody :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => LastUse -> Used -> Body (Aliases rep) -> LastUseM rep
analyseBody lumap used (Body _ stms result) = do
  let used' = used <> freeIn result
  analyseStms lumap used' stms

analyseKernelBody ::
  (FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
  (LastUse, Used) ->
  KernelBody (Aliases rep) ->
  LastUseM rep
analyseKernelBody (lumap, used) (KernelBody _ stms result) =
  let used' = used <> freeIn result
   in analyseStms lumap used' stms

analyseSegBinOp ::
  (FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
  (LastUse, Used) ->
  SegBinOp (Aliases rep) ->
  LastUseM rep
analyseSegBinOp (lumap, used) (SegBinOp _ lambda neutral shp) = do
  (lumap', used') <- analyseLambda (lumap, used) lambda
  let nms = (freeIn neutral <> freeIn shp) `namesSubtract` used'
  pure (lumap', used' <> nms)

analyseHistOp ::
  (FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>
  (LastUse, Used) ->
  HistOp (Aliases rep) ->
  LastUseM rep
analyseHistOp (lumap, used) (HistOp width race dest neutral shp lambda) = do
  (lumap', used') <- analyseLambda (lumap, used) lambda
  let nms =
        ( freeIn width
            <> freeIn race
            <> freeIn dest
            <> freeIn neutral
            <> freeIn shp
        )
          `namesSubtract` used'
  pure (lumap', used' <> nms)

analyseLambda :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => (LastUse, Used) -> Lambda (Aliases rep) -> LastUseM rep
analyseLambda (lumap, used) (Lambda params body ret) = do
  (lumap', used') <- analyseBody lumap used body
  let used'' = used' <> freeIn params <> freeIn ret
  pure (lumap', used'')

flipMap :: Map VName VName -> Map VName Names
flipMap m =
  M.toList m
    & fmap (swap . first oneName)
    & foldr (uncurry $ M.insertWith (<>)) mempty

insertNames :: VName -> Names -> LastUse -> LastUse
insertNames name names lumap =
  foldr (flip (M.insertWith $ \_ x -> x) name) lumap $ namesToList names