huffman.hs

{-# LANGUAGE DeriveFoldable    #-}
{-# LANGUAGE DeriveFunctor     #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE KindSignatures    #-}
{-# LANGUAGE TemplateHaskell   #-}
{-# LANGUAGE TupleSections     #-}
{-# LANGUAGE TypeApplications  #-}
{-# LANGUAGE TypeFamilies      #-}

module Lib where

import           Control.Arrow (second)
import           Data.Foldable (toList)
import           Data.Functor.Foldable
import           Data.Functor.Foldable.TH
import           Data.Map (Map)
import qualified Data.Map as M
import           Data.Maybe (isNothing, fromJust, fromMaybe)
import           Data.PriorityQueue.FingerTree (PQueue)
import qualified Data.PriorityQueue.FingerTree as PQ
import           Data.Tuple (swap)


data Tree a = Leaf a | Branch (Tree a) (Tree a)
  deriving (Functor, Traversable, Foldable, Eq, Ord, Show)

makeBaseFunctor ''Tree

prefixTree :: Ord a => Tree a -> Map a [Bool]
prefixTree = M.fromList . cata alg
  where
    alg (LeafF a) = pure (a, [])
    alg (BranchF l r) = fmap (second (False :)) l
                     ++ fmap (second (True  :)) r

mergeMin :: (Ord k, Num k) => PQueue k (Tree a) -> Maybe (PQueue k (Tree a))
mergeMin pq = do
  ((k1, v1), pq')  <- PQ.minViewWithKey pq
  ((k2, v2), pq'') <- PQ.minViewWithKey pq'
  let k' = k1 + k2
  pure $ PQ.insert k' (Branch v1 v2) pq''

reduce :: (Ord k, Num k) => PQueue k (Tree a) -> Maybe (Tree a)
reduce pq =
  let z = iterate (mergeMin =<<) $ pure pq
   in fmap fst
    . PQ.minView
    . fromJust
    . snd
    . head
    . filter (isNothing . fst)
    $ zip (tail z) z

weigh :: (Ord k, Num k, Ord a, Foldable t) => t a -> PQueue k a
weigh = PQ.fromList
      . fmap swap
      . M.toList
      . M.fromListWith (+)
      . fmap (,1)
      . toList

encode :: (Monoid z, Foldable t, Ord a) => Map a z -> t a -> z
encode m = foldMap (fromMaybe mempty . flip M.lookup m)

huffman :: (Ord a, Foldable t) => t a -> ([Bool], Map a [Bool])
huffman z =
  let tree = prefixTree
           . fromJust
           . reduce @Float
           . fmap Leaf
           $ weigh z
   in (encode tree z, tree)