Huffman

Huff.hs

{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TemplateHaskell            #-}

module Huff where

import           Control.Applicative
import           Control.Arrow
import           Control.Comonad
import           Control.Lens
import           Data.Foldable
import           Data.Function
import           Data.List (sortBy, sort)
import qualified Data.Map as Map
import           Data.Monoid
import           Data.Ord
import           Prelude hiding (sum, foldr)

newtype Density a
  = Density { unDensity :: [(a, Double)] }
  deriving (Functor, Show, Eq, Ord)

density :: [(a, Double)] -> Density a
density probs = Density (map (second (/2)) probs)
  where z = sum (map snd probs)

--------------------------------------------------------------------------------

data Store s a = Store (s -> a) s deriving Functor

instance Comonad (Store s) where
  extract   (Store f s) = f s
  extend g  (Store f s) = Store (g . Store f) s
  duplicate (Store f s) = Store (Store f) s

instance Eq a => Eq (Store s a) where
  (==) = (==) `on` extract

instance Ord a => Ord (Store s a) where
  compare = comparing extract

store :: (s -> a) -> s -> Store s a
store = Store

context :: Store s a -> s
context (Store f s) = s

--------------------------------------------------------------------------------

newtype Sorted a
  = Sorted { sorts :: [a] }
  deriving (Functor, Show, Eq, Ord)

instance Ord a => Cons (Sorted a) (Sorted a) a a where
  _Cons = prism' (\(a, Sorted as) -> sortFast (a:as))
                 (\(Sorted as) -> second Sorted <$> preview _Cons as)

sortFast :: Ord a => [a] -> Sorted a
sortFast = Sorted . sort

instance Ord a => Monoid (Sorted a) where
  mempty = Sorted []
  mappend (Sorted as) (Sorted bs) = Sorted (as /\/ bs) where
    (/\/) :: Ord a => [a] -> [a] -> [a]
    [] /\/ bs = bs
    as /\/ [] = as
    (a : as) /\/ (b : bs) =
      case a `compare` b of
        EQ -> a : b : (as /\/ bs)
        LT -> a : as /\/ (b : bs)
        GT -> b : (a : as) /\/ bs

sorted :: (Ord a, Foldable t) => t a -> Sorted a
sorted = foldMap (Sorted . pure)

--------------------------------------------------------------------------------

class Finite a where
  every :: [a]

--------------------------------------------------------------------------------

newtype Code = Code [Bool] deriving (Eq, Ord, Monoid)

instance Show Code where
  show (Code c) = case c of
    [] -> "·"
    _  -> map (\x -> if x then '1' else '0') c

newtype Coding a = Coding { encode :: a -> Maybe Code }

instance (Show a, Finite a) => Show (Coding a) where
  show c = show $ map (\x -> (x, encode c x)) every

--------------------------------------------------------------------------------

data Bin a
  = Node a
  | Bin (Bin a) (Bin a)
  deriving (Functor, Show, Eq, Ord)

bin :: (r -> r -> r) -> (a -> r) -> (Bin a -> r)
bin two one = go where
  go = \case
    Bin l r -> two (go l) (go r)
    Node  a -> one a

iterMaybe :: (a -> Maybe a) -> (a -> a)
iterMaybe f = go where go a0 = maybe a0 go (f a0)

huffmanTree :: forall a . Density a -> Maybe (Bin a)
huffmanTree (Density probs) =
  fmap (fst . context) . fst <$> uncons res
  where
    res :: Sorted (Bin (Store (a, Double) Double))
    res = iterMaybe go ps

    ps :: Sorted (Bin (Store (a, Double) Double))
    ps  = sortFast (Node . store snd <$> probs)
    
    go :: Sorted (Bin (Store (a, Double) Double))
          -> Maybe (Sorted (Bin (Store (a, Double) Double)))
    go s = do
      (a, s')  <- uncons s
      (b, s'') <- uncons s'
      return (cons (Bin a b) s'')

huffmanCode :: Ord a => Bin a -> Coding a
huffmanCode = Coding . flip Map.lookup . fmap (Code . reverse) . go [] where
  go c = \case
    Node a  -> Map.singleton a c
    Bin l r -> go (False : c) l <> go (True : c) r

huffmanDecode :: Bin a -> Code -> Maybe a
huffmanDecode b (Code c) = go c b where
  go []           (Node b)  = Just b
  go (False : c') (Bin l r) = go c' l
  go (True  : c') (Bin l r) = go c' r
  go _            _         = Nothing
    
huffman :: Ord a => Density a -> Maybe (Coding a)
huffman = fmap huffmanCode . huffmanTree

--------------------------------------------------------------------------------

data Symb = A | B | C | D | E deriving ( Show, Eq, Ord, Enum )

instance Finite Symb where
  every = [A .. E]