Huffman.hs

module Huffman where

import Data.Array
import Data.Char
import Data.List

a ~> b = b a

data Huffman_Tree =
  HLeaf Int Char
  | HNode Int Huffman_Tree Huffman_Tree
  deriving (Eq, Show)

data Bit  = Zero | One deriving (Eq, Show)

instance Ord Huffman_Tree where
  compare (HNode i _ _) (HNode j _ _) = compare i j
  compare (HLeaf i _)   (HNode j _ _) = if (i == j) then LT          else compare i j
  compare (HNode i _ _) (HLeaf j _)   = if (i == j) then GT          else compare i j
  compare (HLeaf i a)   (HLeaf j m)   = if (i == j) then compare a m else compare i j

countchars s = accumArray (+) 0 (foldr1 min s, foldr1 max s) [(x, 1) | x <- s]

tree s = foldtree $ sort [HLeaf n c | (c, n) <- assocs $ countchars s, n > 0]
  where
    foldtree (a:[]) = a
    foldtree (a:b:ls) = foldtree $ (HNode ((weight a) + (weight b)) a b) : ls

    weight (HLeaf i _)   = i
    weight (HNode i _ _) = i

codebook tr = walk [] tr
  where
    walk pref (HLeaf _ ch) = [(ch, pref)]
    walk pref (HNode _ t1 t2) = (walk (pref ++ [Zero]) t1) ++ (walk (pref ++ [One]) t2)

encode str = (tree str, concatMap tryLookup str)
  where
    cb = codebook $ tree str
    tryLookup x = maybe undefined id $ lookup x cb

decode (tr, code) = decodeInner tr code
  where
    decodeInner _  [] = []
    decodeInner tr as =
      getFirstChar tr as
      ~> (\it -> fst it : (decodeInner tr (snd it)))

    getFirstChar (HNode _ t0 t1) (Zero:as) = getFirstChar t0 as
    getFirstChar (HNode _ t0 t1)  (One:as) = getFirstChar t1 as
    getFirstChar (HLeaf _ ch)          as  = (ch, as)