{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
-------------------------------------------------------------------------------
-- |
-- Module      :  ParamTree
-- Copyright   :  (C) 2017-2023 Merijn Verstraaten
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Merijn Verstraaten <merijn@inconsistent.nl>
-- Stability   :  experimental
-- Portability :  haha
--
-- Easily generate a labelled tree of tests/benchmarks from a generation
-- function and sets of parameters to use for each of that functions arguments.
-- Example usecases include criterion benchmark trees or tasty test trees.
-------------------------------------------------------------------------------
module ParamTree
    ( Params
    , ParamFun
    , growTree
    , simpleParam
    , derivedParam
    , displayParam
    , customParam
    , paramSets
    ) where

import Data.Kind (Type)
import Data.Map (Map)
import qualified Data.Map as M
import Data.Monoid (Endo(..))

-- | Type family that converts a type level list into a function type:
--
-- @'ParamFun' ['Char', 'Int', 'Bool'] r@ =
-- @'Char' -> 'Int' -> 'Bool' -> 'String' -> r@
type family ParamFun (l :: [Type]) r where
    ParamFun '[] r = String -> r
    ParamFun (h ': t) r = h -> ParamFun t r

-- | Sets of parameters to generate the tree from.
data Params :: [Type] -> Type where
    Nil :: Params '[]
    Sets :: [Params l] -> Params l
    Param :: Eq r
          => (a -> String) -- Display parameter
          -> (a -> r)      -- Derive value from parameter
          -> String        -- Parameter name
          -> [a]           -- Parameter values
          -> Params l
          -> Params (r ': l)

data Tree :: [Type] -> Type where
    None :: Tree '[]
    Empty :: Tree l
    Grouped :: Eq r
            => Map (String, String) [(r, Tree l)]
            -> Tree (r ': l)

-- | A simple parameter set. The tree label is a combination of 'show'ing the
-- value and the parameter name.
simpleParam
    :: (Eq a, Show a)
    => String -- ^ Name of the parameter
    -> [a] -- ^ Set of values to use
    -> Params l
    -> Params (a ': l)
simpleParam :: forall a (l :: [*]).
(Eq a, Show a) =>
String -> [a] -> Params l -> Params (a : l)
simpleParam = forall r l (l :: [*]).
Eq r =>
(l -> String)
-> (l -> r) -> String -> [l] -> Params l -> Params (r : l)
Param forall a. Show a => a -> String
show forall a. a -> a
id

-- | A derived parameter set. Useful when the input expected by your function
-- can't be conveniently rendered as a string label. For example:
--
-- @'derivedParam' ('enumFromTo' 0) \"My Parameter\" [1,2,5]@
--
-- The above passed @'enumFromTo' 0 1@, @'enumFromTo' 0 2@, etc. to your
-- function, while labelling them as \"1 My Parameter\" and \"2 My Parameter\"
-- respectively.
derivedParam
    :: (Eq r, Show a)
    => (a -> r) -- ^ Parameter derivation function
    -> String -- ^ Name of the parameter
    -> [a] -- ^ Set of values to derive from
    -> Params l
    -> Params (r ': l)
derivedParam :: forall r a (l :: [*]).
(Eq r, Show a) =>
(a -> r) -> String -> [a] -> Params l -> Params (r : l)
derivedParam a -> r
f = forall r l (l :: [*]).
Eq r =>
(l -> String)
-> (l -> r) -> String -> [l] -> Params l -> Params (r : l)
Param forall a. Show a => a -> String
show a -> r
f

-- | A simple parameter set with a more flexible way of showing values,
-- 'simpleParam' is equivalent to @displayParam show@.
displayParam
    :: Eq a
    => (a -> String)
    -> String
    -> [a]
    -> Params l
    -> Params (a ': l)
displayParam :: forall a (l :: [*]).
Eq a =>
(a -> String) -> String -> [a] -> Params l -> Params (a : l)
displayParam a -> String
display = forall r l (l :: [*]).
Eq r =>
(l -> String)
-> (l -> r) -> String -> [l] -> Params l -> Params (r : l)
Param a -> String
display forall a. a -> a
id

-- | A completely customisable parameter set, allows specification of how to
-- display values and how to derive values. Equivalencies:
--
-- @'simpleParam' = 'customParam' 'show' 'id'@
--
-- @'derivedParam' = 'customParam' 'show'@
--
-- @'displayParam' = \\f -> 'customParam' f 'id'@
customParam
    :: Eq r
    => (a -> String)
    -> (a -> r)
    -> String
    -> [a]
    -> Params l
    -> Params (r ': l)
customParam :: forall r l (l :: [*]).
Eq r =>
(l -> String)
-> (l -> r) -> String -> [l] -> Params l -> Params (r : l)
customParam = forall r l (l :: [*]).
Eq r =>
(l -> String)
-> (l -> r) -> String -> [l] -> Params l -> Params (r : l)
Param

-- | Combine multiple sets of parameters into one. Allows a limited amount of
-- control over which combinations appear. For example:
--
-- @
-- 'paramSets'
--     [ 'simpleParam' \"Bool\" [True] . 'simpleParam' \"Char\" \"xy\"
--     , 'simpleParam' \"Bool\" [True,False] . 'simpleParam' \"Char\" \"a\"
--     ]
-- @
--
-- The result is \"axy\" being used in groups where the \"Bool\" parameter is
-- @True@, if the \"Bool\" parameter is @False@ only \"a\" is used.
paramSets :: [Params r -> Params l] -> Params r -> Params l
paramSets :: forall (r :: [*]) (l :: [*]).
[Params r -> Params l] -> Params r -> Params l
paramSets [Params r -> Params l]
prefixes Params r
rest = forall (l :: [*]). [Params l] -> Params l
Sets forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a -> b) -> a -> b
$ Params r
rest) [Params r -> Params l]
prefixes

trim :: [Tree l] -> Tree l
trim :: forall (l :: [*]). [Tree l] -> Tree l
trim [] = forall (l :: [*]). Tree l
Empty
trim (Tree l
None:[Tree l]
_) = Tree '[]
None
trim (Tree l
Empty:[Tree l]
l) = forall (l :: [*]). [Tree l] -> Tree l
trim [Tree l]
l
trim l :: [Tree l]
l@(Grouped{}:[Tree l]
_) = forall r (l :: [*]).
Eq r =>
Map (String, String) [(r, Tree l)] -> Tree (r : l)
Grouped forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
M.unionsWith forall x (l :: [*]).
Eq x =>
[(x, Tree l)] -> [(x, Tree l)] -> [(x, Tree l)]
fuse forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall h (t :: [*]).
Tree (h : t) -> Map (String, String) [(h, Tree t)]
explode [Tree l]
l
  where
    explode :: Tree (h ': t) -> Map (String, String) [(h, Tree t)]
    explode :: forall h (t :: [*]).
Tree (h : t) -> Map (String, String) [(h, Tree t)]
explode Tree (h : t)
Empty = forall k a. Map k a
M.empty
    explode (Grouped Map (String, String) [(r, Tree l)]
m) = Map (String, String) [(r, Tree l)]
m

sprout :: Params l -> Tree l
sprout :: forall (l :: [*]). Params l -> Tree l
sprout Params l
Nil = Tree '[]
None
sprout (Sets [Params l]
l) = forall (l :: [*]). [Tree l] -> Tree l
trim forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall (l :: [*]). Params l -> Tree l
sprout [Params l]
l
sprout (Param a -> String
display a -> r
derive String
name [a]
values Params l
remainder) =
    forall r (l :: [*]).
Eq r =>
Map (String, String) [(r, Tree l)] -> Tree (r : l)
Grouped forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith forall x (l :: [*]).
Eq x =>
[(x, Tree l)] -> [(x, Tree l)] -> [(x, Tree l)]
fuse forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map a -> ((String, String), [(r, Tree l)])
convert forall a b. (a -> b) -> a -> b
$ [a]
values
  where
    convert :: a -> ((String, String), [(r, Tree l)])
convert a
x = ((a -> String
display a
x, String
name), [(a -> r
derive a
x, forall (l :: [*]). Params l -> Tree l
sprout Params l
remainder)])

fuse :: Eq x => [(x, Tree l)] -> [(x, Tree l)] -> [(x, Tree l)]
fuse :: forall x (l :: [*]).
Eq x =>
[(x, Tree l)] -> [(x, Tree l)] -> [(x, Tree l)]
fuse = forall a. Endo a -> a -> a
appEndo forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Monoid a => [a] -> a
mconcat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall a. (a -> a) -> Endo a
Endo forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a} {l :: [*]}.
Eq a =>
(a, Tree l) -> [(a, Tree l)] -> [(a, Tree l)]
insert)
  where
    insert :: (a, Tree l) -> [(a, Tree l)] -> [(a, Tree l)]
insert (a
x, Tree l
params) [] = [(a
x, Tree l
params)]
    insert new :: (a, Tree l)
new@(a
x1, Tree l
params1) ((a
x2, Tree l
params2):[(a, Tree l)]
l)
        | a
x1 forall a. Eq a => a -> a -> Bool
== a
x2 = (a
x1, forall (l :: [*]). [Tree l] -> Tree l
trim [Tree l
params1, Tree l
params2])forall a. a -> [a] -> [a]
:[(a, Tree l)]
l
        | Bool
otherwise = (a
x2, Tree l
params2) forall a. a -> [a] -> [a]
: (a, Tree l) -> [(a, Tree l)] -> [(a, Tree l)]
insert (a, Tree l)
new [(a, Tree l)]
l

-- | Generate a tree from a function that produces a leaf and sets of
-- parameters. Useful for generating tasty TestTrees or criterion benchmark
-- trees from a function and a set of parameter. For example:
--
-- @
-- import Test.Tasty
-- import Test.Tasty.HUnit
--
-- genTestCase :: Int -> Bool -> Char -> String -> TestTree
--
-- params = 'simpleParam' \"Int\" [1,2]
--        . 'simpleParam' \"Bool\" [True]
--        . 'simpleParam' \"Char\" "xyz"
--
-- main :: IO ()
-- main = defaultMain $ testTree genTestCase params
--   where
--     testTree = growTree (Just "/") testGroup "my tests"
-- @
--
-- This generates a tasty TestTree with all combinations of values passed to
-- @genTestCase@. If the 'Maybe' 'String' argument is provided like in the
-- above example, groups with a single entry, such as \"Bool\" get collapsed
-- into their parent groups. So instead of a \"1 Int\" group containing a
-- \"True Bool\" group they get collapsed into a single \"1 Int/True Bool\"
-- group, where the \"/\" separator is the one specified by @'Just' \"\/\"@
growTree
    :: forall a l
     . Maybe String -- ^ Groups containing a single entry are skipped and their
                    -- label is appended to their child, separated by this
                    -- 'String' if not 'Nothing'.
    -> (String -> [a] -> a) -- ^ Tree labelling function, e.g. tasty's
                            -- @testGroup@ or criterion's @bgroup@
    -> String -- ^ Label for the root of the tree
    -> ParamFun l a -- ^ Function that produces leafs, such as tasty tests or
                    -- criterion benchmarks
    -> (Params '[] -> Params l) -- ^ Parameter sets to grow tree from
    -> a
growTree :: forall a (l :: [*]).
Maybe String
-> (String -> [a] -> a)
-> String
-> ParamFun l a
-> (Params '[] -> Params l)
-> a
growTree Maybe String
collapse String -> [a] -> a
labelFun String
label ParamFun l a
fun Params '[] -> Params l
params =
    forall (k :: [*]). Tree k -> ParamFun k a -> String -> a
go (forall (l :: [*]). Params l -> Tree l
sprout forall a b. (a -> b) -> a -> b
$ Params '[] -> Params l
params Params '[]
Nil) ParamFun l a
fun String
label
  where
    go :: Tree k -> ParamFun k a -> String -> a
    go :: forall (k :: [*]). Tree k -> ParamFun k a -> String -> a
go Tree k
None ParamFun k a
result String
lbl = ParamFun k a
result String
lbl
    go Tree k
Empty ParamFun k a
_ String
lbl = String -> [a] -> a
labelFun String
lbl []
    go (Grouped Map (String, String) [(r, Tree l)]
l) ParamFun k a
f String
lbl = case forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall h (t :: [*]).
((String, String), [(h, Tree t)]) -> [(String, h, Tree t)]
flatten (forall k a. Map k a -> [(k, a)]
M.toList Map (String, String) [(r, Tree l)]
l) of
        [(String, r, Tree l)
x] | Just String
sep <- Maybe String
collapse -> forall h (t :: [*]).
ParamFun (h : t) a
-> (String -> String) -> (String, h, Tree t) -> a
buildBranch ParamFun k a
f (\String
n -> String
lbl forall a. [a] -> [a] -> [a]
++ String
sep forall a. [a] -> [a] -> [a]
++ String
n)  (String, r, Tree l)
x
        [(String, r, Tree l)]
branches -> String -> [a] -> a
labelFun String
lbl forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall h (t :: [*]).
ParamFun (h : t) a
-> (String -> String) -> (String, h, Tree t) -> a
buildBranch ParamFun k a
f forall a. a -> a
id) [(String, r, Tree l)]
branches

    flatten :: ((String, String), [(h, Tree t)]) -> [(String, h, Tree t)]
    flatten :: forall h (t :: [*]).
((String, String), [(h, Tree t)]) -> [(String, h, Tree t)]
flatten ((String
param, String
name), [(h, Tree t)]
rest) = forall a b. (a -> b) -> [a] -> [b]
map (\(h
v, Tree t
r) -> (String
nextLabel, h
v, Tree t
r)) [(h, Tree t)]
rest
      where
        nextLabel :: String
nextLabel | forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
name = String
param
                  | Bool
otherwise = String
param forall a. [a] -> [a] -> [a]
++ String
" " forall a. [a] -> [a] -> [a]
++ String
name

    buildBranch
        :: ParamFun (h ': t) a
        -> (String -> String)
        -> (String, h, Tree t)
        -> a
    buildBranch :: forall h (t :: [*]).
ParamFun (h : t) a
-> (String -> String) -> (String, h, Tree t) -> a
buildBranch ParamFun (h : t) a
f String -> String
namer (String
name, h
val, Tree t
rest) = forall (k :: [*]). Tree k -> ParamFun k a -> String -> a
go Tree t
rest (ParamFun (h : t) a
f h
val) (String -> String
namer String
name)