MotzkinTree sampler tests

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

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

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 []