HuffmanTree.hs

module Codec.Compression.Zlib.HuffmanTree (
  HuffmanTree,
  AdvanceResult (..),
  createHuffmanTree,
  advanceTree,
) where

import Data.Bits (testBit)
import Data.Word (Word8)

data HuffmanTree a
  = HuffmanNode (HuffmanTree a) (HuffmanTree a)
  | HuffmanValue a
  | HuffmanEmpty
  deriving (Show)

data AdvanceResult a
  = AdvanceError String
  | NewTree (HuffmanTree a)
  | Result a

emptyHuffmanTree :: HuffmanTree a
emptyHuffmanTree = HuffmanEmpty

createHuffmanTree ::
  Show a =>
  [(a, Int, Int)] ->
  Either String (HuffmanTree a)
createHuffmanTree = foldr addHuffmanNode' (Right emptyHuffmanTree)
 where
  addHuffmanNode' (a, b, c) acc =
    case acc of
      Left err -> Left err
      Right tree -> addHuffmanNode a b c tree

addHuffmanNode ::
  Show a =>
  a ->
  Int ->
  Int ->
  HuffmanTree a ->
  Either String (HuffmanTree a)
addHuffmanNode val len code node =
  case node of
    HuffmanEmpty
      | len == 0 ->
        Right (HuffmanValue val)
    HuffmanEmpty ->
      case addHuffmanNode val (len - 1) code HuffmanEmpty of
        Left err -> Left err
        Right newNode
          | testBit code (len - 1) -> Right (HuffmanNode HuffmanEmpty newNode)
          | otherwise -> Right (HuffmanNode newNode HuffmanEmpty)
    --
    HuffmanValue _
      | len == 0 ->
        Left "Two values point to the same place!"
    HuffmanValue _ ->
      Left "HuffmanValue hit while inserting a value!"
    --
    HuffmanNode _ _
      | len == 0 ->
        Left ("Tried to add where the leaf is a node: " ++ show val)
    HuffmanNode l r | testBit code (len - 1) ->
      case addHuffmanNode val (len - 1) code r of
        Left err -> Left err
        Right r' -> Right (HuffmanNode l r')
    HuffmanNode l r ->
      case addHuffmanNode val (len - 1) code l of
        Left err -> Left err
        Right l' -> Right (HuffmanNode l' r)

advanceTree :: Word8 -> HuffmanTree a -> AdvanceResult a
advanceTree x node =
  case node of
    HuffmanEmpty -> AdvanceError "Tried to advance empty tree!"
    HuffmanValue _ -> AdvanceError "Tried to advance value!"
    HuffmanNode l r ->
      case if (x == 1) then r else l of
        HuffmanEmpty -> AdvanceError "Advanced to empty tree!"
        HuffmanValue y -> Result y
        t -> NewTree t
{-# INLINE advanceTree #-}