From f8672b62e5c3c878aff532883b6c11fc37102fe6 Mon Sep 17 00:00:00 2001
From: Maciej Bendkowski <maciej.bendkowski@gmail.com>
Date: Sat, 16 Jul 2022 11:44:57 +0200
Subject: [PATCH] `MotzkinTree` sampler tests

---
 generic-boltzmann-brain.cabal     |  1 +
 test/Test/Sampler.hs              |  5 ++++
 test/Test/Samplers/MotzkinTree.hs | 44 +++++++++++++++++++++++++++++++
 3 files changed, 50 insertions(+)
 create mode 100644 test/Test/Samplers/MotzkinTree.hs

diff --git a/generic-boltzmann-brain.cabal b/generic-boltzmann-brain.cabal
index 7ebce86..fc40a80 100644
--- a/generic-boltzmann-brain.cabal
+++ b/generic-boltzmann-brain.cabal
@@ -95,6 +95,7 @@ test-suite generic-boltzmann-brain-test
       Test.Sampler
       Test.Samplers.BinTree
       Test.Samplers.Lambda
+      Test.Samplers.MotzkinTree
       Test.Samplers.Tree
       Test.Utils
       Paths_generic_boltzmann_brain
diff --git a/test/Test/Sampler.hs b/test/Test/Sampler.hs
index 6d61ce1..d0ba1c0 100644
--- a/test/Test/Sampler.hs
+++ b/test/Test/Sampler.hs
@@ -10,6 +10,7 @@ import Data.Boltzmann (
 
 import Test.Samplers.BinTree (BinTree)
 import Test.Samplers.Lambda (BinLambda, Lambda, abstractions)
+import Test.Samplers.MotzkinTree (MotzkinTree)
 import Test.Samplers.Tree (Tree, Tree')
 
 import System.Random.SplitMix (SMGen)
@@ -31,6 +32,10 @@ tests =
         \binTree ->
           let s = size @BinTree binTree
            in 800 <= s && s <= 1200
+    , QC.testProperty "MotzkinTree sampler respects size constraints" $
+        \tree ->
+          let s = size @MotzkinTree tree
+           in 800 <= s && s <= 1200
     , QC.testProperty "Tree sampler respects size constraints" $
         \tree ->
           let s = size @Tree tree
diff --git a/test/Test/Samplers/MotzkinTree.hs b/test/Test/Samplers/MotzkinTree.hs
new file mode 100644
index 0000000..1785cfa
--- /dev/null
+++ b/test/Test/Samplers/MotzkinTree.hs
@@ -0,0 +1,44 @@
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE TemplateHaskell #-}
+
+module Test.Samplers.MotzkinTree (MotzkinTree (..)) where
+
+import Data.Boltzmann (
+  BoltzmannSampler (..),
+  LowerBound (MkLowerBound),
+  System (..),
+  UpperBound (MkUpperBound),
+  mkBoltzmannSampler,
+  mkDefWeights,
+  quickCheckRejectionSampler,
+ )
+import Data.Default (def)
+import GHC.Generics (Generic)
+import Test.QuickCheck (Arbitrary (arbitrary, shrink))
+import Test.Utils (Size (size))
+
+data MotzkinTree
+  = Leaf
+  | UnaryNode MotzkinTree
+  | BinNode MotzkinTree MotzkinTree
+  deriving (Generic, Show)
+
+instance Size MotzkinTree where
+  size = \case
+    Leaf -> 1
+    UnaryNode t -> 1 + size t
+    BinNode lt rt -> 1 + size lt + size rt
+
+mkBoltzmannSampler
+  System
+    { targetType = ''MotzkinTree
+    , meanSize = 1000
+    , frequencies = def
+    , weights = $(mkDefWeights ''MotzkinTree)
+    }
+
+instance Arbitrary MotzkinTree where
+  arbitrary =
+    quickCheckRejectionSampler $
+      const (MkLowerBound 800, MkUpperBound 1200)
+  shrink = const []