haskell

Combination of group and take with Data.List.NonEmpty


I'm refactoring a program by changing [a] to (NonEmpty a) so as to avoid having to checks cases of empty lists, which simplifies a lot of my code. In the old code, at one point I transformed a list as follows:

myTake :: Int -> [a] -> [a]
myTake someNum = concat . take someNum . groupBy someFunc

I want to write

myTake' :: Int -> NonEmpty a -> NonEmpty a

I'm stumped. I notice there's no concat in Data.List.NonEmpty and the standard definition doesn't work with Data.List.NonEmpty.


Solution

  • Your biggest problem is what to do about myTake 0 xs. That must produce an empty list, and so you can't return a NonEmpty. If your specification requires that you handle 0, then the only thing you can do is leave the return type as [a] (or Maybe (NonEmpty a), which amounts to the same thing). If, on the other hand, 0 is an input you will never receive, you have some choices. You can simply throw an error when you receive 0, or you can shift the burden of proof to the caller by defining a new PositiveInt type that can't contain 0, and requiring that as your first argument. For details on that approach see Positive integer type.

    Whatever you do, the compiler will think that take n xs may return an empty list, and you have to handle it somehow. If you've already convinced yourself that n can't be zero, then your best bet is really just to call error in the impossible case. Something like

    let included = case take n groups of
          [] -> error "Impossible: n was known to be positive"
          (g:gs) -> g :| gs
    in ...
    

    Your other problems can be fixed by using the right functions.

    groupBy :: Foldable f => (a -> a -> Bool) -> f a -> [NonEmpty a]
    

    is provided in Data.List.NonEmpty because it's a perfectly lovely improvement on Data.List.groupBy, in that it proves at the type level something that is indeed true: each partition is non-empty. But it doesn't prove that there is at least one partition, because the input list might be empty. You do know the input list is non-empty, so you should use

    groupBy1 :: (a -> a -> Bool) -> NonEmpty a -> NonEmpty (NonEmpty a)
    

    This gets you part of the way there: you now have a NonEmpty (NonEmpty a), which even the compiler will agree definitely has at least one element. You can call take on this, and convert the result to NonEmpty using the approaches described in the first section of this answer. You still can't call concat on this, of course, because concat expects traditional lists. But concat is just a special case of the more general mconcat, which works for any Monoid (lists being one example). Unfortunately, NonEmpty is not a Monoid, because it can't support mempty. But there's a name for "things which are almost a Monoid except they don't support mempty": Semigroup. And Semigroup contains

    sconcat :: Semigroup a => NonEmpty a -> a
    

    which is exactly what you want: it converts a NonEmpty (NonEmpty a) into a simple NonEmpty a.

    To sum up, here is one possible solution you might use, where I've chosen "Throw an error" as the approach to being asked for 0 items.

    import Data.List.NonEmpty (NonEmpty((:|)), groupBy1)
    import qualified Data.List.NonEmpty as NE
    import Data.Semigroup (sconcat)
    
    myTake' :: Int -> NonEmpty a -> NonEmpty a
    myTake' n xs = case NE.take n $ groupBy1 someFunc xs of
      [] -> error "You aren't supposed to ask for 0 things"
      (g:gs) -> sconcat (g:|gs)