{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE Trustworthy #-}
-- | 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
       ) where

import Control.Monad.Writer
import Data.Text (Text)
import qualified Data.Text as T
import Data.List (tails)
import qualified Data.Map.Strict as M

import Futhark.IR

-- | AST metrics are simply a collection from identifiable node names
-- to the number of times that node appears.
newtype AstMetrics = AstMetrics (M.Map Text Int)

instance Show AstMetrics where
  show :: AstMetrics -> String
show (AstMetrics Map Text Int
m) = [String] -> String
unlines ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ ((Text, Int) -> String) -> [(Text, Int)] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Int) -> String
forall a a. (Pretty a, Pretty a) => (a, a) -> String
metric ([(Text, Int)] -> [String]) -> [(Text, Int)] -> [String]
forall a b. (a -> b) -> a -> b
$ Map Text Int -> [(Text, Int)]
forall k a. Map k a -> [(k, a)]
M.toList Map Text Int
m
    where metric :: (a, a) -> String
metric (a
k, a
v) = a -> String
forall a. Pretty a => a -> String
pretty a
k String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Pretty a => a -> String
pretty a
v

instance Read AstMetrics where
  readsPrec :: Int -> ReadS AstMetrics
readsPrec Int
_ String
s =
    [(AstMetrics, String)]
-> ([(Text, Int)] -> [(AstMetrics, String)])
-> Maybe [(Text, Int)]
-> [(AstMetrics, String)]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] [(Text, Int)] -> [(AstMetrics, String)]
forall b. IsString b => [(Text, Int)] -> [(AstMetrics, b)]
success (Maybe [(Text, Int)] -> [(AstMetrics, String)])
-> Maybe [(Text, Int)] -> [(AstMetrics, String)]
forall a b. (a -> b) -> a -> b
$ (String -> Maybe (Text, Int)) -> [String] -> Maybe [(Text, Int)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM String -> Maybe (Text, Int)
forall b. Read b => String -> Maybe (Text, b)
onLine ([String] -> Maybe [(Text, Int)])
-> [String] -> Maybe [(Text, Int)]
forall a b. (a -> b) -> a -> b
$ String -> [String]
lines String
s
    where onLine :: String -> Maybe (Text, b)
onLine String
l = case String -> [String]
words String
l of
                       [String
k, String
x] | [(b
n, String
"")] <- ReadS b
forall a. Read a => ReadS a
reads String
x -> (Text, b) -> Maybe (Text, b)
forall a. a -> Maybe a
Just (String -> Text
T.pack String
k, b
n)
                       [String]
_ -> Maybe (Text, b)
forall a. Maybe a
Nothing
          success :: [(Text, Int)] -> [(AstMetrics, b)]
success [(Text, Int)]
m = [(Map Text Int -> AstMetrics
AstMetrics (Map Text Int -> AstMetrics) -> Map Text Int -> AstMetrics
forall a b. (a -> b) -> a -> b
$ [(Text, Int)] -> Map Text Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Text, Int)]
m, b
"")]

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

instance OpMetrics () where
  opMetrics :: () -> MetricsM ()
opMetrics () = () -> MetricsM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

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 { MetricsM a -> Writer CountMetrics a
runMetricsM :: Writer CountMetrics a }
                   deriving (Applicative MetricsM
a -> MetricsM a
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
MetricsM a -> (a -> MetricsM b) -> MetricsM b
MetricsM a -> MetricsM b -> MetricsM b
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
return :: a -> MetricsM a
$creturn :: forall a. a -> MetricsM a
>> :: MetricsM a -> MetricsM b -> MetricsM b
$c>> :: forall a b. MetricsM a -> MetricsM b -> MetricsM b
>>= :: MetricsM a -> (a -> MetricsM b) -> MetricsM b
$c>>= :: forall a b. MetricsM a -> (a -> MetricsM b) -> MetricsM b
$cp1Monad :: Applicative MetricsM
Monad, Functor MetricsM
a -> MetricsM a
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
MetricsM a -> MetricsM b -> MetricsM b
MetricsM a -> MetricsM b -> MetricsM a
MetricsM (a -> b) -> MetricsM a -> MetricsM b
(a -> b -> c) -> MetricsM a -> MetricsM b -> MetricsM c
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
<* :: MetricsM a -> MetricsM b -> MetricsM a
$c<* :: forall a b. MetricsM a -> MetricsM b -> MetricsM a
*> :: MetricsM a -> MetricsM b -> MetricsM b
$c*> :: forall a b. MetricsM a -> MetricsM b -> MetricsM b
liftA2 :: (a -> b -> c) -> MetricsM a -> MetricsM b -> MetricsM c
$cliftA2 :: forall a b c.
(a -> b -> c) -> MetricsM a -> MetricsM b -> MetricsM c
<*> :: MetricsM (a -> b) -> MetricsM a -> MetricsM b
$c<*> :: forall a b. MetricsM (a -> b) -> MetricsM a -> MetricsM b
pure :: a -> MetricsM a
$cpure :: forall a. a -> MetricsM a
$cp1Applicative :: Functor MetricsM
Applicative, a -> MetricsM b -> MetricsM a
(a -> b) -> MetricsM a -> MetricsM b
(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
<$ :: a -> MetricsM b -> MetricsM a
$c<$ :: forall a b. a -> MetricsM b -> MetricsM a
fmap :: (a -> b) -> MetricsM a -> MetricsM b
$cfmap :: forall a b. (a -> b) -> MetricsM a -> MetricsM b
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 (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 lore) => Prog lore -> AstMetrics
progMetrics :: Prog lore -> AstMetrics
progMetrics Prog lore
prog =
  CountMetrics -> AstMetrics
actualMetrics (CountMetrics -> AstMetrics) -> CountMetrics -> AstMetrics
forall a b. (a -> b) -> a -> b
$ Writer CountMetrics () -> CountMetrics
forall w a. Writer w a -> w
execWriter (Writer CountMetrics () -> CountMetrics)
-> Writer CountMetrics () -> CountMetrics
forall a b. (a -> b) -> a -> b
$ MetricsM () -> Writer CountMetrics ()
forall a. MetricsM a -> Writer CountMetrics a
runMetricsM (MetricsM () -> Writer CountMetrics ())
-> MetricsM () -> Writer CountMetrics ()
forall a b. (a -> b) -> a -> b
$ do
  (FunDef lore -> MetricsM ()) -> [FunDef lore] -> MetricsM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ FunDef lore -> MetricsM ()
forall lore. OpMetrics (Op lore) => FunDef lore -> MetricsM ()
funDefMetrics ([FunDef lore] -> MetricsM ()) -> [FunDef lore] -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ Prog lore -> [FunDef lore]
forall lore. Prog lore -> [FunDef lore]
progFuns Prog lore
prog
  (Stm lore -> MetricsM ()) -> Seq (Stm lore) -> MetricsM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm lore -> MetricsM ()
forall lore. OpMetrics (Op lore) => Stm lore -> MetricsM ()
stmMetrics (Seq (Stm lore) -> MetricsM ()) -> Seq (Stm lore) -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ Prog lore -> Seq (Stm lore)
forall lore. Prog lore -> Stms lore
progConsts Prog lore
prog

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

bodyMetrics :: OpMetrics (Op lore) => Body lore -> MetricsM ()
bodyMetrics :: Body lore -> MetricsM ()
bodyMetrics = (Stm lore -> MetricsM ()) -> Seq (Stm lore) -> MetricsM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Stm lore -> MetricsM ()
forall lore. OpMetrics (Op lore) => Stm lore -> MetricsM ()
stmMetrics (Seq (Stm lore) -> MetricsM ())
-> (Body lore -> Seq (Stm lore)) -> Body lore -> MetricsM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Body lore -> Seq (Stm lore)
forall lore. BodyT lore -> Stms lore
bodyStms

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

expMetrics :: OpMetrics (Op lore) => Exp lore -> MetricsM ()
expMetrics :: Exp lore -> MetricsM ()
expMetrics (BasicOp BasicOp
op) =
  Text -> MetricsM ()
seen Text
"BasicOp" MetricsM () -> MetricsM () -> MetricsM ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> BasicOp -> MetricsM ()
primOpMetrics BasicOp
op
expMetrics (DoLoop [(FParam lore, SubExp)]
_ [(FParam lore, SubExp)]
_ ForLoop{} BodyT lore
body) =
  Text -> MetricsM () -> MetricsM ()
inside Text
"DoLoop" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ Text -> MetricsM ()
seen Text
"ForLoop" MetricsM () -> MetricsM () -> MetricsM ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> BodyT lore -> MetricsM ()
forall lore. OpMetrics (Op lore) => Body lore -> MetricsM ()
bodyMetrics BodyT lore
body
expMetrics (DoLoop [(FParam lore, SubExp)]
_ [(FParam lore, SubExp)]
_ WhileLoop{} BodyT lore
body) =
  Text -> MetricsM () -> MetricsM ()
inside Text
"DoLoop" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ Text -> MetricsM ()
seen Text
"WhileLoop" MetricsM () -> MetricsM () -> MetricsM ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> BodyT lore -> MetricsM ()
forall lore. OpMetrics (Op lore) => Body lore -> MetricsM ()
bodyMetrics BodyT lore
body
expMetrics (If SubExp
_ BodyT lore
tb BodyT lore
fb IfDec (BranchType lore)
_) =
  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
$ BodyT lore -> MetricsM ()
forall lore. OpMetrics (Op lore) => Body lore -> MetricsM ()
bodyMetrics BodyT lore
tb
    Text -> MetricsM () -> MetricsM ()
inside Text
"False" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
forall a b. (a -> b) -> a -> b
$ BodyT lore -> MetricsM ()
forall lore. OpMetrics (Op lore) => Body lore -> MetricsM ()
bodyMetrics BodyT lore
fb
expMetrics Apply{} =
  Text -> MetricsM ()
seen Text
"Apply"
expMetrics (Op Op lore
op) =
  Op lore -> MetricsM ()
forall op. OpMetrics op => op -> MetricsM ()
opMetrics Op lore
op

primOpMetrics :: BasicOp -> MetricsM ()
primOpMetrics :: BasicOp -> MetricsM ()
primOpMetrics (SubExp SubExp
_) = Text -> MetricsM ()
seen Text
"SubExp"
primOpMetrics (Opaque SubExp
_) = Text -> MetricsM ()
seen Text
"Opaque"
primOpMetrics ArrayLit{} = Text -> MetricsM ()
seen Text
"ArrayLit"
primOpMetrics BinOp{} = Text -> MetricsM ()
seen Text
"BinOp"
primOpMetrics UnOp{} = Text -> MetricsM ()
seen Text
"UnOp"
primOpMetrics ConvOp{} = Text -> MetricsM ()
seen Text
"ConvOp"
primOpMetrics CmpOp{} = Text -> MetricsM ()
seen Text
"ConvOp"
primOpMetrics Assert{} = Text -> MetricsM ()
seen Text
"Assert"
primOpMetrics Index{} = Text -> MetricsM ()
seen Text
"Index"
primOpMetrics Update{} = Text -> MetricsM ()
seen Text
"Update"
primOpMetrics Concat{} = Text -> MetricsM ()
seen Text
"Concat"
primOpMetrics Copy{} = Text -> MetricsM ()
seen Text
"Copy"
primOpMetrics Manifest{} = Text -> MetricsM ()
seen Text
"Manifest"
primOpMetrics Iota{} = Text -> MetricsM ()
seen Text
"Iota"
primOpMetrics Replicate{} = Text -> MetricsM ()
seen Text
"Replicate"
primOpMetrics Scratch{} = Text -> MetricsM ()
seen Text
"Scratch"
primOpMetrics Reshape{} = Text -> MetricsM ()
seen Text
"Reshape"
primOpMetrics Rearrange{} = Text -> MetricsM ()
seen Text
"Rearrange"
primOpMetrics Rotate{} = Text -> MetricsM ()
seen Text
"Rotate"

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