content: -- | Compute a mapping from variables to their corresponding (fully
  -- expanded) PrimExps.
  module Futhark.Analysis.PrimExp.Table
    ( primExpTable,
      PrimExpTable,
  
      -- * Extensibility
      PrimExpAnalysis (..),
  
      -- * Testing
      stmToPrimExps,
    )
  where
  
  import Control.Monad.State.Strict
  import Data.Foldable
  import Data.Map.Strict qualified as M
  import Futhark.Analysis.PrimExp
  import Futhark.Analysis.PrimExp.Convert
  import Futhark.IR.Aliases
  import Futhark.IR.GPU
  import Futhark.IR.GPUMem
  import Futhark.IR.MC
  import Futhark.IR.MCMem
  
  -- | Maps variables to maybe PrimExps. Will map to nothing if it
  -- cannot be resolved to a PrimExp. For all uses of this analysis atm.
  -- a variable can be considered inscrutable if it cannot be resolved
  -- to a primexp.
  type PrimExpTable = M.Map VName (Maybe (PrimExp VName))
  
  -- | A class for extracting PrimExps from what is inside an op.
  class PrimExpAnalysis rep where
    opPrimExp :: Scope rep -> Op rep -> State PrimExpTable ()
  
  primExpTable :: (PrimExpAnalysis rep, RepTypes rep) => Prog rep -> PrimExpTable
  primExpTable prog = initialState <> foldMap' (uncurry funToPrimExp) scopesAndFuns
    where
      scopesAndFuns = do
        let fun_defs = progFuns prog
        let scopes = map getScope fun_defs
        zip scopes fun_defs
  
      getScope funDef = scopeOf (progConsts prog) <> scopeOfFParams (funDefParams funDef)
  
      -- We need to have the dummy "slice" in the analysis for our "slice hack".
      initialState =
        M.singleton (VName "slice" 0) $ Just $ LeafExp (VName "slice" 0) $ IntType Int64
  
  funToPrimExp ::
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep ->
    FunDef rep ->
    PrimExpTable
  funToPrimExp scope fundef = execState (bodyToPrimExps scope (funDefBody fundef)) mempty
  
  -- | Adds the statements of a body to the PrimExpTable
  bodyToPrimExps ::
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep ->
    Body rep ->
    State PrimExpTable ()
  bodyToPrimExps scope body = mapM_ (stmToPrimExps scope') (bodyStms body)
    where
      scope' = scope <> scopeOf (bodyStms body)
  
  -- | Adds the statements of a kernel body to the PrimExpTable
  kernelToBodyPrimExps ::
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep ->
    KernelBody rep ->
    State PrimExpTable ()
  kernelToBodyPrimExps scope kbody = mapM_ (stmToPrimExps scope') (bodyStms kbody)
    where
      scope' = scope <> scopeOf (bodyStms kbody)
  
  -- | Adds a statement to the PrimExpTable. If it can't be resolved as a `PrimExp`,
  -- it will be added as `Nothing`.
  stmToPrimExps ::
    forall rep.
    (PrimExpAnalysis rep, RepTypes rep) =>
    Scope rep ->
    Stm rep ->
    State PrimExpTable ()
  stmToPrimExps scope stm = do
    table <- get
    case stm of
      (Let (Pat pat_elems) _ e)
        | Just primExp <- primExpFromExp (toPrimExp scope table) e ->
            -- The statement can be resolved as a `PrimExp`.
            -- For each pattern element, insert the PrimExp in the table
            forM_ pat_elems $ \pe ->
              modify $ M.insert (patElemName pe) (Just primExp)
        | otherwise -> do
            -- The statement can't be resolved as a `PrimExp`.
            walk $ stmExp stm -- Traverse the rest of the AST Get the
            -- updated PrimExpTable after traversing the AST
            table' <- get
  
            -- Add pattern elements that can't be resolved as `PrimExp`
            -- to the `PrimExpTable` as `Nothing`
            forM_ pat_elems $ \pe ->
              case M.lookup (patElemName pe) table' of
                Nothing -> modify $ M.insert (patElemName pe) Nothing
                Just _ -> pure ()
    where
      walk e = do
        -- Handle most cases using the walker
        walkExpM walker e
        -- Additionally, handle loop parameters
        case e of
          Loop _ (ForLoop i t _) _ ->
            modify $ M.insert i $ Just $ LeafExp i $ IntType t
          _ -> pure ()
  
      walker =
        (identityWalker @rep)
          { walkOnBody = \body_scope -> bodyToPrimExps (scope <> body_scope),
            walkOnOp = opPrimExp scope,
            walkOnFParam = paramToPrimExp -- Loop parameters
          }
  
      -- Adds a loop parameter to the PrimExpTable
      paramToPrimExp :: FParam rep -> State PrimExpTable ()
      paramToPrimExp param = do
        let name = paramName param
        -- Construct a `PrimExp` from the type of the parameter
        -- and add it to the `PrimExpTable`
        case typeOf $ paramDec param of
          -- TODO: Handle other types?
          Prim pt ->
            modify $ M.insert name (Just $ LeafExp name pt)
          _ -> pure ()
  
  -- | Checks if a name is in the PrimExpTable and construct a `PrimExp`
  -- if it is not
  toPrimExp :: (RepTypes rep) => Scope rep -> PrimExpTable -> VName -> Maybe (PrimExp VName)
  toPrimExp scope table name = case M.lookup name table of
    Just maybePrimExp
      | Just primExp <- maybePrimExp -> Just primExp -- Already in the table
    _ -> case fmap typeOf . M.lookup name $ scope of
      (Just (Prim pt)) -> Just $ LeafExp name pt
      _ -> Nothing
  
  -- | Adds the parameters of a SegOp as well as the statements in its
  -- body to the PrimExpTable
  segOpToPrimExps :: (PrimExpAnalysis rep, RepTypes rep) => Scope rep -> SegOp lvl rep -> State PrimExpTable ()
  segOpToPrimExps scope op = do
    forM_ (map fst $ unSegSpace $ segSpace op) $ \name ->
      modify $ M.insert name $ Just $ LeafExp name int64
    kernelToBodyPrimExps scope (segBody op)
  
  instance PrimExpAnalysis GPU where
    opPrimExp scope gpu_op
      | (SegOp op) <- gpu_op = segOpToPrimExps scope op
      | (SizeOp _) <- gpu_op = pure ()
      | (GPUBody _ body) <- gpu_op = bodyToPrimExps scope body
      | (Futhark.IR.GPUMem.OtherOp _) <- gpu_op = pure ()
  
  instance PrimExpAnalysis MC where
    opPrimExp scope mc_op
      | (ParOp maybe_par_segop seq_segop) <- mc_op = do
          -- Add the statements in the parallel part of the ParOp to the PrimExpTable
          case maybe_par_segop of
            Nothing -> pure ()
            Just _ -> forM_ maybe_par_segop $ segOpToPrimExps scope
          -- Add the statements in the sequential part of the ParOp to the PrimExpTable
          segOpToPrimExps scope seq_segop
      | (Futhark.IR.MCMem.OtherOp _) <- mc_op = pure ()
  