Huffman Encoding in Haskell (or the beginnings of it, at least)

Lib.hs

{-# LANGUAGE ScopedTypeVariables #-}

module Lib
  ( toBasicTree,
    toTree,
    sortTreeList,
    weight,
    Tree (..),
  )
where

import Data.Function (on)
import Data.List (sortOn)
import Data.Map ()
import qualified Data.Map as Map
import Data.Map.Lazy (Map)
import Data.Maybe (fromMaybe)
import Data.PQueue.Min (MinQueue)
import qualified Data.PQueue.Min as PQueue

data Tree a
  = Leaf a Int
  | Node (Tree a) (Tree a)
  | Empty
  deriving (Show, Eq)

instance Eq a => Ord (Tree a) where
  compare = compare `on` weight

weight :: Tree a -> Int
weight (Leaf _ n) = n
weight (Node x y) = weight x + weight y
weight Empty = 0

sortTreeList :: [Tree a] -> [Tree a]
sortTreeList = sortOn weight

toBasicTree :: (Foldable f) => f a -> MinQueue (Tree a)
toBasicTree = weightMapToLeaves . foldr sumWeights mempty
  where
    weightMapToLeaves :: Ord a => Map a Int -> MinQueue (Tree a)
    weightMapToLeaves = PQueue.fromList . fmap (uncurry Leaf) . Map.toList

    sumWeights :: Ord a => a -> Map a Int -> Map a Int
    sumWeights = Map.alter (Just . (+ 1) . fromMaybe 0)

toTree :: (Ord a, Foldable f) => f a -> Tree a
toTree str = go $ toBasicTree str
  where
    go :: Eq a => MinQueue (Tree a) -> Tree a
    go q = case PQueue.take 2 q of
      [] -> Empty
      [x] -> x
      (x : y : _) -> go $ PQueue.insert (Node x y) (PQueue.drop 2 q)

LibSpec.hs

module LibSpec (spec) where

import qualified Data.PQueue.Min as PQueue
import Lib (Tree (Leaf, Node), toBasicTree, toTree, weight)
import Test.Hspec (Spec, describe, it, shouldBe)
import Test.QuickCheck (Testable (property))

spec :: Spec
spec =
  describe "Huffman" $ do
    it "Can create a tree" $ do
      toBasicTree "hello" `shouldBe` PQueue.fromList [Leaf 'e' 1, Leaf 'h' 1, Leaf 'l' 2, Leaf 'o' 1]
    it "The weight of the whole tree is the length of the input string." $
      property $
        \str -> weight (toTree str) == length (str :: String)
    it "Can reduce a tree" $ do
      toTree "hello"
        `shouldBe` Node
          (Leaf 'l' 2)
          ( Node
              (Leaf 'e' 1)
              ( Node
                  (Leaf 'o' 1)
                  (Leaf 'h' 1)
              )
          )

      toTree "The quick and rather charming brown fox jumped over the lazy evaluation algorithm"
        `shouldBe` Node
          ( Node
              ( Node
                  ( Node
                      ( Node
                          ( Node (Leaf 'k' 1) (Leaf 'j' 1)
                          )
                          ( Node (Leaf 'q' 1) (Leaf 'p' 1)
                          )
                      )
                      (Leaf 'i' 4)
                  )
                  ( Node
                      ( Node
                          ( Node (Leaf 'z' 1) (Leaf 'y' 1)
                          )
                          (Leaf 'v' 2)
                      )
                      (Leaf 't' 4)
                  )
              )
              ( Node
                  ( Node
                      (Leaf 'n' 4)
                      ( Node (Leaf 'd' 2) (Leaf 'c' 2)
                      )
                  )
                  ( Node
                      ( Node
                          (Leaf 'g' 2)
                          ( Node (Leaf 'x' 1) (Leaf 'w' 1)
                          )
                      )
                      (Leaf 'h' 5)
                  )
              )
          )
          ( Node
              ( Node
                  ( Node
                      (Leaf 'o' 5)
                      ( Node
                          (Leaf 'l' 3)
                          ( Node
                              (Leaf 'T' 1)
                              ( Node (Leaf 'f' 1) (Leaf 'b' 1)
                              )
                          )
                      )
                  )
                  (Leaf ' ' 12)
              )
              ( Node
                  ( Node
                      (Leaf 'r' 6)
                      ( Node (Leaf 'u' 3) (Leaf 'm' 3)
                      )
                  )
                  ( Node (Leaf 'e' 6) (Leaf 'a' 7)
                  )
              )
          )