huffman encoding by fingertrees

huffman.hs

{-# language NoMonomorphismRestriction, ViewPatterns, MultiParamTypeClasses, FlexibleInstances #-}
import Control.Arrow ((&&&))
import Data.List (group,unfoldr,sort)
import Data.Set (Set)
import qualified Data.Set as S
import Data.FingerTree (FingerTree,split,fromList,ViewL(..),Measured(..),(<|), viewl)
import Data.Monoid ((<>))
import Data.Semigroup (Min(..))
import Control.Applicative ((<|>))

-----   FT lib
pick
  :: (Measured v t, Eq v) =>
     FingerTree v t -> Maybe (t, FingerTree v t)
pick ft = let
  (as,bs) = split (== measure ft) ft
  in case viewl bs of
      EmptyL -> Nothing
      (b :< bs) -> Just (b,as <> bs)
---

data Bin a = Bin (Bin a) (Set a) Int  (Bin a) | Leaf a Int deriving Show

instance Measured (Min Int) (Bin a) where
  measure (Bin _  _ n _) = Min n
  measure (Leaf _ n) = Min n

contrib (Bin _ s n _ ) = (s,n)
contrib (Leaf x n) = (S.singleton x, n)

merge x@(contrib -> (sx,ix)) y@(contrib -> (sy,iy)) =
  Bin x (sx <> sy) (ix + iy) y

type FT a = FingerTree (Min Int) (Bin a)

mkFT :: Ord a => [a] -> FT a
mkFT = fromList . map (uncurry Leaf  . (head &&& length)) . group . sort

mkHuffman :: Ord a => FT a -> Maybe (Bin a)
mkHuffman ft = case pick ft of
    Nothing -> Nothing
    Just (t1,ft') -> case pick ft' of
                       Nothing -> Just (t1)
                       Just (t2,ft'') -> mkHuffman $ t1 `merge` t2 <| ft''

encoding :: Ord a => b -> b -> a -> Bin a -> Maybe [b]
encoding z o x (Bin l (S.member x -> True) _ r) = (z :) <$> encoding z o x l <|> (o:) <$> encoding z o x r
encoding _ _ x (Leaf ((==) x -> True) _) = Just []
encoding _ _ x _ = Nothing

decoding _  _ es (Leaf x _) = (x,es)
decoding z o (e:es) (Bin l _ _ r) = decoding z o es $ if e == z then l else r

encode z o ph = do
  h <- mkHuffman (mkFT ph)
  e <- mapM (flip (encoding z o) h) ph
  return (h,concat e)

decode z o (h,zs) = unfoldr f zs where
  f [] = Nothing
  f zs = Just $ decoding z o zs h

main = print $ decode 0 1 <$>  encode 0 1 "fingertrees rock"