Test that list synonym samplers respect size constraints

internal/Data/Boltzmann/System/TH.hs

@@ -100,20 +100,6 @@ getWeight constr = do
   weightResolver <- asks constructorWeight
   lift $ weightResolver `unWeightResolver` constr
 
-mkCoerce :: TypeVariant -> SamplerGen Exp
-mkCoerce tv = do
-  typSynonym <- findTypeSyn tv
-  let fromType = convert tv
-      toType = typSynonym
-
-  coerce' <- lift [|coerce|]
-  pure $ AppTypeE (AppTypeE coerce' fromType) toType
-  where
-    convert :: TypeVariant -> Type
-    convert = \case
-      Plain tn -> ConT $ coerce tn
-      List tn -> AppT ListT (ConT $ coerce tn)
-
 toTypeVariant :: Type -> SamplerGen TypeVariant
 toTypeVariant (ConT tn) = pure . Plain $ coerce tn
 toTypeVariant (AppT ListT (ConT tn)) = pure . List $ coerce tn
@@ -160,17 +146,25 @@ mkCaseConstr = \case
 
     caseMatches <- mapM (mkCaseMatch tv) constrGroup
     pure $ LamCaseE caseMatches
-  tv -> do
-    coerceExp <- mkCoerce tv
-    lift
-      [|
-        \case
-          0 -> pure ([], 0)
-          1 -> do
-            (x, w) <- sample ub
-            (xs, ws) <- sample (ub - w)
-            pure ($(pure coerceExp) (x : xs), w + ws)
-        |]
+  tv@(List tn) ->
+    do
+      typSynonym <- findTypeSyn (Plain tn)
+      listTypSynonym <- findTypeSyn tv
+
+      lift
+        [|
+          \case
+            0 -> pure ([], 0)
+            1 -> do
+              (x, w) <- $(sampleExp typSynonym) ub
+              (xs, ws) <- $(sampleExp listTypSynonym) (ub - w)
+              pure ((x : xs), w + ws)
+          |]
+
+sampleExp :: Type -> Q Exp
+sampleExp t = do
+  sample' <- [|sample|]
+  pure $ AppTypeE sample' t
 
 mkCaseMatch :: TypeVariant -> (ConstructorInfo, Integer) -> SamplerGen Match
 mkCaseMatch tv (constr, idx) = do

test/Test/Sampler.hs

@@ -10,7 +10,7 @@ import Data.Boltzmann (
 
 import Test.Samplers.BinTree (BinTree)
 import Test.Samplers.Lambda (BinLambda, Lambda, abstractions)
-import Test.Samplers.Tree (Tree)
+import Test.Samplers.Tree (Tree, Tree')
 
 import System.Random.SplitMix (SMGen)
 
@@ -35,6 +35,10 @@ tests =
         \tree ->
           let s = size @Tree tree
            in 1600 <= s && s <= 2400
+    , QC.testProperty "Tree' sampler respects size constraints" $
+        \tree ->
+          let s = size @Tree' tree
+           in 8500 <= s && s <= 11_150
     , QC.testProperty "Lambda sampler respects size constraints" $
         \term ->
           let s = size @Lambda term

test/Test/Samplers/Tree.hs

@@ -1,18 +1,23 @@
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE TemplateHaskell #-}
 
-module Test.Samplers.Tree (Tree (..)) where
+module Test.Samplers.Tree (Tree (..), Tree' (..)) where
 
 import Data.Boltzmann (
   BoltzmannSampler (..),
+  Constructable (..),
   LowerBound (MkLowerBound),
+  System (..),
   UpperBound (MkUpperBound),
   hoistRejectionSampler,
+  mkBoltzmannSampler,
   mkDefBoltzmannSampler,
+  mkDefWeights,
  )
+import Data.Default (def)
 import GHC.Generics (Generic)
 import Test.QuickCheck (Arbitrary (arbitrary, shrink))
-import Test.Utils (Size(size))
+import Test.Utils (Size (size))
 
 data Tree = T [Tree]
   deriving (Generic, Show)
@@ -21,10 +26,33 @@ mkDefBoltzmannSampler ''Tree 2000
 
 instance Size Tree where
   size = \case
-    T ts -> 1 + sum (map size ts)
+    T ts -> 1 + sum (size <$> ts)
 
 instance Arbitrary Tree where
   arbitrary =
     hoistRejectionSampler $
       const (MkLowerBound 1600, MkUpperBound 2400)
   shrink = const []
+
+newtype Tree' = MkTree' Tree
+  deriving (Generic, Show)
+
+instance Size Tree' where
+  size = \case
+    (MkTree' (T ts)) -> 2 + sum (size . MkTree' <$> ts)
+
+mkBoltzmannSampler
+  System
+    { targetType = ''Tree'
+    , meanSize = 10_000
+    , frequencies = def
+    , weights =
+        ('T, 2)
+          <:> $(mkDefWeights ''Tree)
+    }
+
+instance Arbitrary Tree' where
+  arbitrary =
+    hoistRejectionSampler $
+      const (MkLowerBound 8500, MkUpperBound 11_150)
+  shrink = const []