{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Definitions for multicore operations.
--
-- Most of the interesting stuff is in "Futhark.IR.SegOp", which is
-- also re-exported from here.
module Futhark.IR.MC.Op
  ( MCOp (..),
    traverseMCOpStms,
    typeCheckMCOp,
    simplifyMCOp,
    module Futhark.IR.SegOp,
  )
where

import Data.Bifunctor (first)
import Futhark.Analysis.Metrics
import Futhark.Analysis.SymbolTable qualified as ST
import Futhark.IR
import Futhark.IR.Aliases (Aliases, CanBeAliased (..))
import Futhark.IR.Prop.Aliases
import Futhark.IR.SegOp
import Futhark.IR.TypeCheck qualified as TC
import Futhark.Optimise.Simplify qualified as Simplify
import Futhark.Optimise.Simplify.Engine qualified as Engine
import Futhark.Optimise.Simplify.Rep
import Futhark.Transform.Rename
import Futhark.Transform.Substitute
import Futhark.Util.Pretty
  ( nestedBlock,
    pretty,
    (<+>),
    (</>),
  )
import Prelude hiding (id, (.))

-- | An operation for the multicore representation.  Feel free to
-- extend this on an ad hoc basis as needed.  Parameterised with some
-- other operation.
data MCOp op rep
  = -- | The first 'SegOp' (if it exists) contains nested parallelism,
    -- while the second one has a fully sequential body.  They are
    -- semantically fully equivalent.
    ParOp
      (Maybe (SegOp () rep))
      (SegOp () rep)
  | -- | Something else (in practice often a SOAC).
    OtherOp (op rep)
  deriving (Eq, Ord, Show)

traverseMCOpStms ::
  Monad m =>
  OpStmsTraverser m (op rep) rep ->
  OpStmsTraverser m (MCOp op rep) rep
traverseMCOpStms _ f (ParOp par_op op) =
  ParOp <$> traverse (traverseSegOpStms f) par_op <*> traverseSegOpStms f op
traverseMCOpStms onInner f (OtherOp op) = OtherOp <$> onInner f op

instance (ASTRep rep, Substitute (op rep)) => Substitute (MCOp op rep) where
  substituteNames substs (ParOp par_op op) =
    ParOp (substituteNames substs <$> par_op) (substituteNames substs op)
  substituteNames substs (OtherOp op) =
    OtherOp $ substituteNames substs op

instance (ASTRep rep, Rename (op rep)) => Rename (MCOp op rep) where
  rename (ParOp par_op op) = ParOp <$> rename par_op <*> rename op
  rename (OtherOp op) = OtherOp <$> rename op

instance (ASTRep rep, FreeIn (op rep)) => FreeIn (MCOp op rep) where
  freeIn' (ParOp par_op op) = freeIn' par_op <> freeIn' op
  freeIn' (OtherOp op) = freeIn' op

instance (ASTRep rep, IsOp (op rep)) => IsOp (MCOp op rep) where
  safeOp (ParOp _ op) = safeOp op
  safeOp (OtherOp op) = safeOp op

  cheapOp (ParOp _ op) = cheapOp op
  cheapOp (OtherOp op) = cheapOp op

instance TypedOp (op rep) => TypedOp (MCOp op rep) where
  opType (ParOp _ op) = opType op
  opType (OtherOp op) = opType op

instance (Aliased rep, AliasedOp (op rep)) => AliasedOp (MCOp op rep) where
  opAliases (ParOp _ op) = opAliases op
  opAliases (OtherOp op) = opAliases op

  consumedInOp (ParOp _ op) = consumedInOp op
  consumedInOp (OtherOp op) = consumedInOp op

instance CanBeAliased op => CanBeAliased (MCOp op) where
  addOpAliases aliases (ParOp par_op op) =
    ParOp (addOpAliases aliases <$> par_op) (addOpAliases aliases op)
  addOpAliases aliases (OtherOp op) =
    OtherOp $ addOpAliases aliases op

instance CanBeWise op => CanBeWise (MCOp op) where
  addOpWisdom (ParOp par_op op) =
    ParOp (addOpWisdom <$> par_op) (addOpWisdom op)
  addOpWisdom (OtherOp op) =
    OtherOp $ addOpWisdom op

instance (ASTRep rep, ST.IndexOp (op rep)) => ST.IndexOp (MCOp op rep) where
  indexOp vtable k (ParOp _ op) is = ST.indexOp vtable k op is
  indexOp vtable k (OtherOp op) is = ST.indexOp vtable k op is

instance (PrettyRep rep, Pretty (op rep)) => Pretty (MCOp op rep) where
  pretty (ParOp Nothing op) = pretty op
  pretty (ParOp (Just par_op) op) =
    "par"
      <+> nestedBlock "{" "}" (pretty par_op)
      </> "seq"
      <+> nestedBlock "{" "}" (pretty op)
  pretty (OtherOp op) = pretty op

instance (OpMetrics (Op rep), OpMetrics (op rep)) => OpMetrics (MCOp op rep) where
  opMetrics (ParOp par_op op) = opMetrics par_op >> opMetrics op
  opMetrics (OtherOp op) = opMetrics op

instance RephraseOp op => RephraseOp (MCOp op) where
  rephraseInOp r (ParOp par_op op) =
    ParOp <$> traverse (rephraseInOp r) par_op <*> rephraseInOp r op
  rephraseInOp r (OtherOp op) = OtherOp <$> rephraseInOp r op

typeCheckMCOp ::
  TC.Checkable rep =>
  (op (Aliases rep) -> TC.TypeM rep ()) ->
  MCOp op (Aliases rep) ->
  TC.TypeM rep ()
typeCheckMCOp _ (ParOp (Just par_op) op) = do
  -- It is valid for the same array to be consumed in both par_op and op.
  _ <- typeCheckSegOp pure par_op `TC.alternative` typeCheckSegOp pure op
  pure ()
typeCheckMCOp _ (ParOp Nothing op) =
  typeCheckSegOp pure op
typeCheckMCOp f (OtherOp op) = f op

simplifyMCOp ::
  ( Engine.SimplifiableRep rep,
    BodyDec rep ~ ()
  ) =>
  Simplify.SimplifyOp rep (op (Wise rep)) ->
  MCOp op (Wise rep) ->
  Engine.SimpleM rep (MCOp op (Wise rep), Stms (Wise rep))
simplifyMCOp f (OtherOp op) = do
  (op', stms) <- f op
  pure (OtherOp op', stms)
simplifyMCOp _ (ParOp par_op op) = do
  (par_op', par_op_hoisted) <-
    case par_op of
      Nothing -> pure (Nothing, mempty)
      Just x -> first Just <$> simplifySegOp x

  (op', op_hoisted) <- simplifySegOp op

  pure (ParOp par_op' op', par_op_hoisted <> op_hoisted)