Introduce a Size type class.

generic-boltzmann-brain.cabal

@@ -92,6 +92,7 @@ test-suite generic-boltzmann-brain-test
       Test.Samplers.Tree
       Test.Unit.BuffonMachine
       Test.Unit.Sampler
+      Test.Utils
       Paths_generic_boltzmann_brain
   hs-source-dirs:
       test

test/Test/Samplers/BinTree.hs

@@ -1,7 +1,7 @@
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE TemplateHaskell #-}
 
-module Test.Samplers.BinTree (BinTree (..), size) where
+module Test.Samplers.BinTree (BinTree (..)) where
 
 import Data.Boltzmann (
   BoltzmannSampler (..),
@@ -16,16 +16,17 @@ import Data.Boltzmann (
 import Data.Default (def)
 import GHC.Generics (Generic)
 import Test.QuickCheck (Arbitrary (arbitrary, shrink))
+import Test.Utils (Size(size))
 
 data BinTree
   = Leaf
   | Node BinTree BinTree
   deriving (Generic, Show)
 
-size :: BinTree -> Int
-size = \case
-  Leaf -> 0
-  Node lt rt -> 1 + size lt + size rt
+instance Size BinTree where
+  size = \case
+    Leaf -> 0
+    Node lt rt -> 1 + size lt + size rt
 
 mkBoltzmannSampler
   System

test/Test/Samplers/Lambda.hs

@@ -5,8 +5,6 @@ module Test.Samplers.Lambda (
   DeBruijn (..),
   Lambda (..),
   BinLambda (..),
-  size,
-  sizeBin,
 ) where
 
 import Data.Boltzmann (
@@ -22,6 +20,7 @@ import Data.Boltzmann (
 import Data.Default (def)
 import GHC.Generics (Generic)
 import Test.QuickCheck (Arbitrary (arbitrary, shrink))
+import Test.Utils (Size(size))
 
 data DeBruijn
   = Z
@@ -48,23 +47,16 @@ mkBoltzmannSampler
           <:> $(mkDefWeights ''Lambda)
     }
 
-size' :: DeBruijn -> Int
-size' = \case
-  Z -> 1
-  S n -> 1 + size' n
+instance Size DeBruijn where
+  size = \case
+    Z -> 1
+    S n -> 1 + size n
 
-size :: Lambda -> Int
-size = \case
-  Index n -> size' n
-  App lt rt -> 1 + size lt + size rt
-  Abs t -> 1 + size t
-
-sizeBin :: BinLambda -> Int
-sizeBin = \case
-  MkBinLambda (Index n) -> size' n
-  MkBinLambda (App lt rt) ->
-    2 + sizeBin (MkBinLambda lt) + sizeBin (MkBinLambda rt)
-  MkBinLambda (Abs t) -> 2 + sizeBin (MkBinLambda t)
+instance Size Lambda where
+  size = \case
+    Index n -> size n
+    App lt rt -> 1 + size lt + size rt
+    Abs t -> 1 + size t
 
 instance Arbitrary Lambda where
   arbitrary =
@@ -75,6 +67,13 @@ instance Arbitrary Lambda where
 newtype BinLambda = MkBinLambda Lambda
   deriving (Generic, Show)
 
+instance Size BinLambda where
+  size = \case
+    MkBinLambda (Index n) -> size n
+    MkBinLambda (App lt rt) ->
+      2 + size (MkBinLambda lt) + size (MkBinLambda rt)
+    MkBinLambda (Abs t) -> 2 + size (MkBinLambda t)
+
 mkBoltzmannSampler
   System
     { targetType = ''BinLambda

test/Test/Samplers/Tree.hs

@@ -1,7 +1,7 @@
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE TemplateHaskell #-}
 
-module Test.Samplers.Tree (Tree (..), size) where
+module Test.Samplers.Tree (Tree (..)) where
 
 import Data.Boltzmann (
   BoltzmannSampler (..),
@@ -12,15 +12,16 @@ import Data.Boltzmann (
  )
 import GHC.Generics (Generic)
 import Test.QuickCheck (Arbitrary (arbitrary, shrink))
+import Test.Utils (Size(size))
 
 data Tree = T [Tree]
   deriving (Generic, Show)
 
 mkDefBoltzmannSampler ''Tree 2000
 
-size :: Tree -> Int
-size = \case
-  T ts -> 1 + sum (map size ts)
+instance Size Tree where
+  size = \case
+    T ts -> 1 + sum (map size ts)
 
 instance Arbitrary Tree where
   arbitrary =

test/Test/Unit/Sampler.hs

@@ -1,8 +1,9 @@
 module Test.Unit.Sampler (unitTests) where
 
-import qualified Test.Samplers.BinTree as BinTree
-import qualified Test.Samplers.Lambda as Lambda
-import qualified Test.Samplers.Tree as Tree
+import Test.Samplers.BinTree (BinTree)
+import Test.Samplers.Lambda (BinLambda, Lambda)
+import Test.Samplers.Tree (Tree)
+import Test.Utils (Size (size))
 
 import Test.Tasty (TestTree, testGroup)
 import Test.Tasty.QuickCheck as QC (testProperty)
@@ -13,18 +14,18 @@ unitTests =
     "Sampler tests"
     [ QC.testProperty "BinTree sampler respects size constraints" $
         \binTree ->
-          let s = BinTree.size binTree
+          let s = size @BinTree binTree
            in 800 <= s && s <= 1200
     , QC.testProperty "Tree sampler respects size constraints" $
         \tree ->
-          let s = Tree.size tree
+          let s = size @Tree tree
            in 1600 <= s && s <= 2400
     , QC.testProperty "Lambda sampler respects size constraints" $
         \term ->
-          let s = Lambda.size term
+          let s = size @Lambda term
            in 8_000 <= s && s <= 12_000
     , QC.testProperty "BinLambda sampler respects size constraints" $
         \term ->
-          let s = Lambda.sizeBin term
+          let s = size @BinLambda term
            in 5_000 <= s && s <= 6_400
     ]

test/Test/Utils.hs

@@ -0,0 +1,5 @@
+module Test.Utils (Size(..)) where
+
+-- | Objects with size.
+class Size a where
+  size :: a -> Int