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 #-}