Enumerating BinTrees by depth

BinTrees.hs

import qualified Control.Monad.WeightedSearch as W
import Data.List (sortBy)
import Data.Ord (comparing)
import Control.Applicative

data BTree = Leaf | Branch BTree BTree deriving (Show, Eq)

-- I can lazily list all the btrees
btrees :: W.T Integer BTree
btrees = pure Leaf <|> W.weight 1 (Branch <$> btrees <*> btrees)

-- Each tree t is to be weighted by 2^depth t - 1
btreesByDepth :: W.T Integer BTree
btreesByDepth = pure Leaf <|> W.weight 1 (do -- w += 1
  l <- btreesByDepth -- w += 2^dl - 1
  r <- btreesByDepth -- w += 2^dr - 1
  let dl = depth l
      dr = depth r
  -- assert (w == (2^dl + 2^dr - 1))
  -- At the end Branch l r should have weight
  -- 2 ^ (1 + max dl dr) - 1 = (2^max + 2^min - 1) + (2^max - 2^min)
  -- we just need to add the second term
  weighted (2 ^ max dl dr - 2 ^ min dl dr) (Branch l r))
  where weighted x = W.weight x . pure

-- Now, let's consider this function
depth :: BTree -> Integer
depth Leaf = 0
depth (Branch b1 b2) = 1 + max (depth b1) (depth b2)

-- I'd like to generate the btrees in an order such that `test btrees == True`, where
test :: W.T Integer BTree -> Bool
test ts = ls == sortBy (comparing depth) ls
  where ls = take 1000 $ W.toList ts