Huffman coding algorithm

huffman.hs

import Data.List (sort, insert, elem, concat)

data HuffmanTree = Leaf { symbol :: Char, weight :: Int} | Node { left :: HuffmanTree, right :: HuffmanTree, symbols :: [Char], weight :: Int} deriving (Show)

instance Eq HuffmanTree where
    tree1 == tree2 = (weight tree1) == (weight tree2)

instance Ord HuffmanTree where
    tree1 `compare` tree2 = (weight tree1) `compare` (weight tree2)

data Bit = Zero | One deriving (Show)

fromBool :: Bool -> Bit
fromBool True = One
fromBool False = Zero

mergeBranches :: HuffmanTree -> HuffmanTree -> HuffmanTree
mergeBranches left right = Node left right (getSymbols left ++ getSymbols right) (weight left + weight right)

getSymbols :: HuffmanTree -> [Char]
getSymbols (Leaf s _) = [s]
getSymbols node = symbols node

decode :: HuffmanTree -> [Bit] -> [Char]
decode _ [] = []
decode tree bits =
    let (symbol, remainingBits) = decodeSymbol bits tree
    in symbol:(decode tree remainingBits)
                   
decodeSymbol :: [Bit] -> HuffmanTree -> (Char, [Bit])
decodeSymbol bits (Leaf symbol _) = (symbol, bits)
decodeSymbol (b:bs) branch = decodeSymbol bs $ chooseBranch b branch
                   
chooseBranch :: Bit -> HuffmanTree -> HuffmanTree
chooseBranch One = right
chooseBranch Zero = left


buildTree :: [HuffmanTree] -> HuffmanTree
buildTree [leaf] = leaf
buildTree nodes = (buildTree . mergeTwoFirst . sort) nodes
                    where mergeTwoFirst (node1:node2:other) = (mergeBranches node1 node2):other
                  
encode :: HuffmanTree -> String -> [Bit]
encode _ "" = []
encode tree symbols = concat $ map (encodeSymbol tree) symbols

encodeSymbol :: HuffmanTree -> Char -> [Bit]
encodeSymbol (Leaf _ _) _ = []
encodeSymbol tree symbol =
    let bit = fromBool $ symbol `elem` (getSymbols $ right tree)
    in bit:(encodeSymbol (chooseBranch bit tree) symbol)
                      

*Main> let tree = buildTree [(Leaf 't' 2), (Leaf 'e' 1), (Leaf 's' 1)]
*Main> encode tree "test"
[One,Zero,Zero,Zero,One,One]
*Main> decode tree $ encode tree "test"
"test"