    -- | Abstract Syntax Tree metrics.  This is used in the @futhark test@
    -- program, for the @structure@ stanzas.
    module Futhark.Analysis.Metrics
      ( AstMetrics (..),
        progMetrics,

        -- * Extensibility
        OpMetrics (..),
        seen,
        inside,
        MetricsM,
        stmMetrics,
        lambdaMetrics,
        bodyMetrics,
      )
    where

    import Control.Monad
    import Control.Monad.Writer
    import Data.List (tails)
    import Data.Map.Strict qualified as M
    import Data.Text (Text)
    import Data.Text qualified as T
    import Futhark.Analysis.Metrics.Type
    import Futhark.IR
    import Futhark.Util (showText)

    -- | Compute the metrics for some operation.
    class OpMetrics op where
      opMetrics :: op -> MetricsM ()

    instance (OpMetrics a) => OpMetrics (Maybe a) where
      opMetrics :: Maybe a -> MetricsM ()
    opMetrics Maybe a
    Nothing = () -> MetricsM ()
    forall a. a -> MetricsM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
      opMetrics (Just a
    x) = a -> MetricsM ()
    forall op. OpMetrics op => op -> MetricsM ()
    opMetrics a
    x

    instance OpMetrics (NoOp rep) where
      opMetrics :: NoOp rep -> MetricsM ()
    opMetrics NoOp rep
    NoOp = () -> MetricsM ()
    forall a. a -> MetricsM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

    newtype CountMetrics = CountMetrics [([Text], Text)]

    instance Semigroup CountMetrics where
      CountMetrics [([Text], Text)]
    x <> :: CountMetrics -> CountMetrics -> CountMetrics
    <> CountMetrics [([Text], Text)]
    y = [([Text], Text)] -> CountMetrics
    CountMetrics ([([Text], Text)] -> CountMetrics)
    -> [([Text], Text)] -> CountMetrics
    forall a b. (a -> b) -> a -> b
    $ [([Text], Text)]
    x [([Text], Text)] -> [([Text], Text)] -> [([Text], Text)]
    forall a. Semigroup a => a -> a -> a
    <> [([Text], Text)]
    y

    instance Monoid CountMetrics where
      mempty :: CountMetrics
    mempty = [([Text], Text)] -> CountMetrics
    CountMetrics [([Text], Text)]
    forall a. Monoid a => a
    mempty

    actualMetrics :: CountMetrics -> AstMetrics
    actualMetrics :: CountMetrics -> AstMetrics
    actualMetrics (CountMetrics [([Text], Text)]
    metrics) =
      Map Text Int -> AstMetrics
    AstMetrics (Map Text Int -> AstMetrics) -> Map Text Int -> AstMetrics
    forall a b. (a -> b) -> a -> b
    $ (Int -> Int -> Int) -> [(Text, Int)] -> Map Text Int
    forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
    M.fromListWith Int -> Int -> Int
    forall a. Num a => a -> a -> a
    (+) ([(Text, Int)] -> Map Text Int) -> [(Text, Int)] -> Map Text Int
    forall a b. (a -> b) -> a -> b
    $ (([Text], Text) -> [(Text, Int)])
    -> [([Text], Text)] -> [(Text, Int)]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap ([Text], Text) -> [(Text, Int)]
    forall {b}. Num b => ([Text], Text) -> [(Text, b)]
    expand [([Text], Text)]
    metrics
      where
        expand :: ([Text], Text) -> [(Text, b)]
    expand ([Text]
    ctx, Text
    k) =
          [ (Text -> [Text] -> Text
    T.intercalate Text
    "/" ([Text]
    ctx' [Text] -> [Text] -> [Text]
    forall a. [a] -> [a] -> [a]
    ++ [Text
    k]), b
    1)
            | [Text]
    ctx' <- [Text] -> [[Text]]
    forall a. [a] -> [[a]]
    tails ([Text] -> [[Text]]) -> [Text] -> [[Text]]
    forall a b. (a -> b) -> a -> b
    $ Text
    "" Text -> [Text] -> [Text]
    forall a. a -> [a] -> [a]
    : [Text]
    ctx
          ]

    -- | This monad is used for computing metrics.  It internally keeps
    -- track of what we've seen so far.  Use 'seen' to add more stuff.
    newtype MetricsM a = MetricsM {forall a. MetricsM a -> Writer CountMetrics a
    runMetricsM :: Writer CountMetrics a}
      deriving
        ( Applicative MetricsM
    Applicative MetricsM =>
    (forall a b. MetricsM a -> (a -> MetricsM b) -> MetricsM b)
    -> (forall a b. MetricsM a -> MetricsM b -> MetricsM b)
    -> (forall a. a -> MetricsM a)
    -> Monad MetricsM
    forall a. a -> MetricsM a
    forall a b. MetricsM a -> MetricsM b -> MetricsM b
    forall a b. MetricsM a -> (a -> MetricsM b) -> MetricsM b
    forall (m :: * -> *).
    Applicative m =>
    (forall a b. m a -> (a -> m b) -> m b)
    -> (forall a b. m a -> m b -> m b)
    -> (forall a. a -> m a)
    -> Monad m
    $c>>= :: forall a b. MetricsM a -> (a -> MetricsM b) -> MetricsM b
    >>= :: forall a b. MetricsM a -> (a -> MetricsM b) -> MetricsM b
    $c>> :: forall a b. MetricsM a -> MetricsM b -> MetricsM b
    >> :: forall a b. MetricsM a -> MetricsM b -> MetricsM b
    $creturn :: forall a. a -> MetricsM a
    return :: forall a. a -> MetricsM a
    Monad,
          Functor MetricsM
    Functor MetricsM =>
    (forall a. a -> MetricsM a)
    -> (forall a b. MetricsM (a -> b) -> MetricsM a -> MetricsM b)
    -> (forall a b c.
        (a -> b -> c) -> MetricsM a -> MetricsM b -> MetricsM c)
    -> (forall a b. MetricsM a -> MetricsM b -> MetricsM b)
    -> (forall a b. MetricsM a -> MetricsM b -> MetricsM a)
    -> Applicative MetricsM
    forall a. a -> MetricsM a
    forall a b. MetricsM a -> MetricsM b -> MetricsM a
    forall a b. MetricsM a -> MetricsM b -> MetricsM b
    forall a b. MetricsM (a -> b) -> MetricsM a -> MetricsM b
    forall a b c.
    (a -> b -> c) -> MetricsM a -> MetricsM b -> MetricsM c
    forall (f :: * -> *).
    Functor f =>
    (forall a. a -> f a)
    -> (forall a b. f (a -> b) -> f a -> f b)
    -> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
    -> (forall a b. f a -> f b -> f b)
    -> (forall a b. f a -> f b -> f a)
    -> Applicative f
    $cpure :: forall a. a -> MetricsM a
    pure :: forall a. a -> MetricsM a
    $c<*> :: forall a b. MetricsM (a -> b) -> MetricsM a -> MetricsM b
    <*> :: forall a b. MetricsM (a -> b) -> MetricsM a -> MetricsM b
    $cliftA2 :: forall a b c.
    (a -> b -> c) -> MetricsM a -> MetricsM b -> MetricsM c
    liftA2 :: forall a b c.
    (a -> b -> c) -> MetricsM a -> MetricsM b -> MetricsM c
    $c*> :: forall a b. MetricsM a -> MetricsM b -> MetricsM b
    *> :: forall a b. MetricsM a -> MetricsM b -> MetricsM b
    $c<* :: forall a b. MetricsM a -> MetricsM b -> MetricsM a
    <* :: forall a b. MetricsM a -> MetricsM b -> MetricsM a
    Applicative,
          (forall a b. (a -> b) -> MetricsM a -> MetricsM b)
    -> (forall a b. a -> MetricsM b -> MetricsM a) -> Functor MetricsM
    forall a b. a -> MetricsM b -> MetricsM a
    forall a b. (a -> b) -> MetricsM a -> MetricsM b
    forall (f :: * -> *).
    (forall a b. (a -> b) -> f a -> f b)
    -> (forall a b. a -> f b -> f a) -> Functor f
    $cfmap :: forall a b. (a -> b) -> MetricsM a -> MetricsM b
    fmap :: forall a b. (a -> b) -> MetricsM a -> MetricsM b
    $c<$ :: forall a b. a -> MetricsM b -> MetricsM a
    <$ :: forall a b. a -> MetricsM b -> MetricsM a
    Functor,
          MonadWriter CountMetrics
        )

    -- | Add this node to the current tally.
    seen :: Text -> MetricsM ()
    seen :: Text -> MetricsM ()
    seen Text
    k = CountMetrics -> MetricsM ()
    forall w (m :: * -> *). MonadWriter w m => w -> m ()
    tell (CountMetrics -> MetricsM ()) -> CountMetrics -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ [([Text], Text)] -> CountMetrics
    CountMetrics [([], Text
    k)]

    -- | Enclose a metrics counting operation.  Most importantly, this
    -- prefixes the name of the context to all the metrics computed in the
    -- enclosed operation.
    inside :: Text -> MetricsM () -> MetricsM ()
    inside :: Text -> MetricsM () -> MetricsM ()
    inside Text
    what MetricsM ()
    m = Text -> MetricsM ()
    seen Text
    what MetricsM () -> MetricsM () -> MetricsM ()
    forall a b. MetricsM a -> MetricsM b -> MetricsM b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> (CountMetrics -> CountMetrics) -> MetricsM () -> MetricsM ()
    forall w (m :: * -> *) a. MonadWriter w m => (w -> w) -> m a -> m a
    censor CountMetrics -> CountMetrics
    addWhat MetricsM ()
    m
      where
        addWhat :: CountMetrics -> CountMetrics
    addWhat (CountMetrics [([Text], Text)]
    metrics) =
          [([Text], Text)] -> CountMetrics
    CountMetrics ((([Text], Text) -> ([Text], Text))
    -> [([Text], Text)] -> [([Text], Text)]
    forall a b. (a -> b) -> [a] -> [b]
    map ([Text], Text) -> ([Text], Text)
    forall {b}. ([Text], b) -> ([Text], b)
    addWhat' [([Text], Text)]
    metrics)
        addWhat' :: ([Text], b) -> ([Text], b)
    addWhat' ([Text]
    ctx, b
    k) = (Text
    what Text -> [Text] -> [Text]
    forall a. a -> [a] -> [a]
    : [Text]
    ctx, b
    k)

    -- | Compute the metrics for a program.
    progMetrics :: (OpMetrics (Op rep)) => Prog rep -> AstMetrics
    progMetrics :: forall rep. OpMetrics (Op rep) => Prog rep -> AstMetrics
    progMetrics Prog rep
    prog =
      CountMetrics -> AstMetrics
    actualMetrics (CountMetrics -> AstMetrics) -> CountMetrics -> AstMetrics
    forall a b. (a -> b) -> a -> b
    $
        WriterT CountMetrics Identity () -> CountMetrics
    forall w a. Writer w a -> w
    execWriter (WriterT CountMetrics Identity () -> CountMetrics)
    -> WriterT CountMetrics Identity () -> CountMetrics
    forall a b. (a -> b) -> a -> b
    $
          MetricsM () -> WriterT CountMetrics Identity ()
    forall a. MetricsM a -> Writer CountMetrics a
    runMetricsM (MetricsM () -> WriterT CountMetrics Identity ())
    -> MetricsM () -> WriterT CountMetrics Identity ()
    forall a b. (a -> b) -> a -> b
    $ do
            (FunDef rep -> MetricsM ()) -> [FunDef rep] -> MetricsM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ FunDef rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => FunDef rep -> MetricsM ()
    funDefMetrics ([FunDef rep] -> MetricsM ()) -> [FunDef rep] -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Prog rep -> [FunDef rep]
    forall rep. Prog rep -> [FunDef rep]
    progFuns Prog rep
    prog
            (Stm rep -> MetricsM ()) -> Seq (Stm rep) -> MetricsM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ Stm rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Stm rep -> MetricsM ()
    stmMetrics (Seq (Stm rep) -> MetricsM ()) -> Seq (Stm rep) -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Prog rep -> Seq (Stm rep)
    forall rep. Prog rep -> Stms rep
    progConsts Prog rep
    prog

    funDefMetrics :: (OpMetrics (Op rep)) => FunDef rep -> MetricsM ()
    funDefMetrics :: forall rep. OpMetrics (Op rep) => FunDef rep -> MetricsM ()
    funDefMetrics = Body rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Body rep -> MetricsM ()
    bodyMetrics (Body rep -> MetricsM ())
    -> (FunDef rep -> Body rep) -> FunDef rep -> MetricsM ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . FunDef rep -> Body rep
    forall rep. FunDef rep -> Body rep
    funDefBody

    -- | Compute metrics for this body.
    bodyMetrics :: (OpMetrics (Op rep)) => Body rep -> MetricsM ()
    bodyMetrics :: forall rep. OpMetrics (Op rep) => Body rep -> MetricsM ()
    bodyMetrics = (Stm rep -> MetricsM ()) -> Seq (Stm rep) -> MetricsM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ Stm rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Stm rep -> MetricsM ()
    stmMetrics (Seq (Stm rep) -> MetricsM ())
    -> (Body rep -> Seq (Stm rep)) -> Body rep -> MetricsM ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Body rep -> Seq (Stm rep)
    forall rep. Body rep -> Stms rep
    bodyStms

    -- | Compute metrics for this statement.
    stmMetrics :: (OpMetrics (Op rep)) => Stm rep -> MetricsM ()
    stmMetrics :: forall rep. OpMetrics (Op rep) => Stm rep -> MetricsM ()
    stmMetrics = Exp rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Exp rep -> MetricsM ()
    expMetrics (Exp rep -> MetricsM ())
    -> (Stm rep -> Exp rep) -> Stm rep -> MetricsM ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stm rep -> Exp rep
    forall rep. Stm rep -> Exp rep
    stmExp

    expMetrics :: (OpMetrics (Op rep)) => Exp rep -> MetricsM ()
    expMetrics :: forall rep. OpMetrics (Op rep) => Exp rep -> MetricsM ()
    expMetrics (BasicOp BasicOp
    op) =
      Text -> MetricsM ()
    seen Text
    "BasicOp" MetricsM () -> MetricsM () -> MetricsM ()
    forall a b. MetricsM a -> MetricsM b -> MetricsM b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> BasicOp -> MetricsM ()
    basicOpMetrics BasicOp
    op
    expMetrics (Loop [(FParam rep, SubExp)]
    _ ForLoop {} Body rep
    body) =
      Text -> MetricsM () -> MetricsM ()
    inside Text
    "Loop" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Text -> MetricsM ()
    seen Text
    "ForLoop" MetricsM () -> MetricsM () -> MetricsM ()
    forall a b. MetricsM a -> MetricsM b -> MetricsM b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Body rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Body rep -> MetricsM ()
    bodyMetrics Body rep
    body
    expMetrics (Loop [(FParam rep, SubExp)]
    _ WhileLoop {} Body rep
    body) =
      Text -> MetricsM () -> MetricsM ()
    inside Text
    "Loop" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Text -> MetricsM ()
    seen Text
    "WhileLoop" MetricsM () -> MetricsM () -> MetricsM ()
    forall a b. MetricsM a -> MetricsM b -> MetricsM b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Body rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Body rep -> MetricsM ()
    bodyMetrics Body rep
    body
    expMetrics (Match [SubExp]
    _ [Case [Just (BoolValue Bool
    True)] Body rep
    tb] Body rep
    fb MatchDec (BranchType rep)
    _) =
      Text -> MetricsM () -> MetricsM ()
    inside Text
    "If" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ do
        Text -> MetricsM () -> MetricsM ()
    inside Text
    "True" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Body rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Body rep -> MetricsM ()
    bodyMetrics Body rep
    tb
        Text -> MetricsM () -> MetricsM ()
    inside Text
    "False" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Body rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Body rep -> MetricsM ()
    bodyMetrics Body rep
    fb
    expMetrics (Match [SubExp]
    _ [Case (Body rep)]
    cases Body rep
    defbody MatchDec (BranchType rep)
    _) =
      Text -> MetricsM () -> MetricsM ()
    inside Text
    "Match" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ do
        [(Int, Case (Body rep))]
    -> ((Int, Case (Body rep)) -> MetricsM ()) -> MetricsM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([Int] -> [Case (Body rep)] -> [(Int, Case (Body rep))]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Int
    0 ..] [Case (Body rep)]
    cases) (((Int, Case (Body rep)) -> MetricsM ()) -> MetricsM ())
    -> ((Int, Case (Body rep)) -> MetricsM ()) -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ \(Int
    i, Case (Body rep)
    c) ->
          Text -> MetricsM () -> MetricsM ()
    inside (Int -> Text
    forall a. Show a => a -> Text
    showText (Int
    i :: Int)) (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Body rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Body rep -> MetricsM ()
    bodyMetrics (Body rep -> MetricsM ()) -> Body rep -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Case (Body rep) -> Body rep
    forall body. Case body -> body
    caseBody Case (Body rep)
    c
        Text -> MetricsM () -> MetricsM ()
    inside Text
    "default" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Body rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Body rep -> MetricsM ()
    bodyMetrics Body rep
    defbody
    expMetrics Apply {} =
      Text -> MetricsM ()
    seen Text
    "Apply"
    expMetrics (WithAcc [WithAccInput rep]
    _ Lambda rep
    lam) =
      Text -> MetricsM () -> MetricsM ()
    inside Text
    "WithAcc" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Lambda rep -> MetricsM ()
    lambdaMetrics Lambda rep
    lam
    expMetrics (Op Op rep
    op) =
      Op rep -> MetricsM ()
    forall op. OpMetrics op => op -> MetricsM ()
    opMetrics Op rep
    op

    basicOpMetrics :: BasicOp -> MetricsM ()
    basicOpMetrics :: BasicOp -> MetricsM ()
    basicOpMetrics (SubExp SubExp
    _) = Text -> MetricsM ()
    seen Text
    "SubExp"
    basicOpMetrics (Opaque OpaqueOp
    _ SubExp
    _) = Text -> MetricsM ()
    seen Text
    "Opaque"
    basicOpMetrics ArrayVal {} = Text -> MetricsM ()
    seen Text
    "ArrayVal"
    basicOpMetrics ArrayLit {} = Text -> MetricsM ()
    seen Text
    "ArrayLit"
    basicOpMetrics BinOp {} = Text -> MetricsM ()
    seen Text
    "BinOp"
    basicOpMetrics UnOp {} = Text -> MetricsM ()
    seen Text
    "UnOp"
    basicOpMetrics ConvOp {} = Text -> MetricsM ()
    seen Text
    "ConvOp"
    basicOpMetrics CmpOp {} = Text -> MetricsM ()
    seen Text
    "CmpOp"
    basicOpMetrics Assert {} = Text -> MetricsM ()
    seen Text
    "Assert"
    basicOpMetrics Index {} = Text -> MetricsM ()
    seen Text
    "Index"
    basicOpMetrics Update {} = Text -> MetricsM ()
    seen Text
    "Update"
    basicOpMetrics FlatIndex {} = Text -> MetricsM ()
    seen Text
    "FlatIndex"
    basicOpMetrics FlatUpdate {} = Text -> MetricsM ()
    seen Text
    "FlatUpdate"
    basicOpMetrics Concat {} = Text -> MetricsM ()
    seen Text
    "Concat"
    basicOpMetrics Manifest {} = Text -> MetricsM ()
    seen Text
    "Manifest"
    basicOpMetrics Iota {} = Text -> MetricsM ()
    seen Text
    "Iota"
    basicOpMetrics Replicate {} = Text -> MetricsM ()
    seen Text
    "Replicate"
    basicOpMetrics Scratch {} = Text -> MetricsM ()
    seen Text
    "Scratch"
    basicOpMetrics Reshape {} = Text -> MetricsM ()
    seen Text
    "Reshape"
    basicOpMetrics Rearrange {} = Text -> MetricsM ()
    seen Text
    "Rearrange"
    basicOpMetrics UpdateAcc {} = Text -> MetricsM ()
    seen Text
    "UpdateAcc"

    -- | Compute metrics for this lambda.
    lambdaMetrics :: (OpMetrics (Op rep)) => Lambda rep -> MetricsM ()
    lambdaMetrics :: forall rep. OpMetrics (Op rep) => Lambda rep -> MetricsM ()
    lambdaMetrics = Body rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Body rep -> MetricsM ()
    bodyMetrics (Body rep -> MetricsM ())
    -> (Lambda rep -> Body rep) -> Lambda rep -> MetricsM ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Lambda rep -> Body rep
    forall rep. Lambda rep -> Body rep
    lambdaBody