Trie tree arbitrary instance

Trie.hs

-- |A @'Trie' a b@ is a map with keys of type @[a]@ and values of type @b@.
data Trie a b = Fork (Maybe b) (Map a (Trie a b))
    deriving (Show, Eq)

instance (Ord a, Arbitrary a, Arbitrary b) => Arbitrary (Trie a b) where

  arbitrary :: Gen (Trie a b)
  arbitrary = sized $ \n -> if n == 0      -- We interpret the size n as maximum number of values
                                           -- stored in the trie.
    then return empty                      -- If the n == 0, the trie must be empty.
    else do
      mb <- arbitrary                      -- Chose a value for the empty key (or not).
      let n' = maybe n (const $ n - 1) mb  -- n' is the maximum number of additional elements.
      as <- arbitrary :: Gen [a]           -- Pick arbitrary key heads.
      m  <- resize n' $ f as               -- resize to n' before picking an arbitrary Map.
      return $ Fork mb m
   where
    f ::[a] -> Gen (Map a (Trie a b))
    f []       = return M.empty
    f (a : as) = sized $ \n -> do
      t <- arbitrary                       -- Recursively pick an arbitrary sub-trie.
      let n' = max 0 (n - length t)        -- update the maximum number of additional elements.
      m <- resize n' (f as)                -- Pick the rest of the dictionary recursively.
      return $ M.insert a t m

  shrink :: Trie a b -> [Trie a b]         -- We want to ensure that all shrinks of a valid trie are valid.
  shrink (Fork mb m) =
    let ms =    [M.insert a t'' m | (a, t') <- M.toList m, t'' <- shrink t', not (null t'')]
                -- Recursively shrink a sub-trie, but don't let it become empty.
             ++ [M.delete a m | a <- M.keys m]
                -- Delete a sub-trie from the Map.
    in  case mb of
          Nothing -> [Fork Nothing m' | m' <- ms]        -- If no value is stored for the empty key, we just take what we got.
          Just _  ->    [Fork mb m' | m' <- ms]          -- If there is a value, we can additionally delete that value.
                     ++ [Fork Nothing m' | m' <- m : ms] -- In that case, we can optionally keep the Map as is.

-- | A QuickCheck generator for /@'valid'@/ tries.
genValidTrie :: forall a b. (Ord a, Arbitrary a, Arbitrary b) => Gen (Trie a b)
genValidTrie = sized $ \n -> case n of
  0 -> return empty
  1 -> oneof
    [ return empty
    , arbitrary >>= \b -> return $ Fork (Just b) M.empty
    ]
  _ -> do
    mb <- arbitrary
    let n' = maybe n (const $ n - 1) mb
    as <- arbitrary :: Gen [a]
    m  <- resize n' $ f as
    return $ Fork mb m
 where
  f ::[a] -> Gen (Map a (Trie a b))
  f []       = return M.empty
  f (a : as) = sized $ \n -> do
    t <- genValidTrie
    let l = length t
    let n' = max 0 (n - l)
    m <- resize n' (f as)
    return $ if l == 0 then m else M.insert a t m