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 (Int -> HuffmanTree a -> ShowS
[HuffmanTree a] -> ShowS
HuffmanTree a -> String
(Int -> HuffmanTree a -> ShowS)
-> (HuffmanTree a -> String)
-> ([HuffmanTree a] -> ShowS)
-> Show (HuffmanTree a)
forall a. Show a => Int -> HuffmanTree a -> ShowS
forall a. Show a => [HuffmanTree a] -> ShowS
forall a. Show a => HuffmanTree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [HuffmanTree a] -> ShowS
$cshowList :: forall a. Show a => [HuffmanTree a] -> ShowS
show :: HuffmanTree a -> String
$cshow :: forall a. Show a => HuffmanTree a -> String
showsPrec :: Int -> HuffmanTree a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> HuffmanTree a -> ShowS
Show)

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

emptyHuffmanTree :: HuffmanTree a
emptyHuffmanTree :: HuffmanTree a
emptyHuffmanTree = HuffmanTree a
forall a. HuffmanTree a
HuffmanEmpty

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

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

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