{-# LANGUAGE BangPatterns, CPP #-}
module Codec.Picture.Gif.Internal.LZWEncoding( lzwEncode ) where

#if !MIN_VERSION_base(4,8,0)
import Control.Applicative( (<$>) )
import Data.Monoid( mempty )
#endif

import Control.Monad.ST( runST )
import qualified Data.ByteString.Lazy as L
import Data.Maybe( fromMaybe )
import Data.Word( Word8 )

#if MIN_VERSION_containers(0,5,0)
import qualified Data.IntMap.Strict as I
#else
import qualified Data.IntMap as I
#endif
import qualified Data.Vector.Storable as V

import Codec.Picture.BitWriter

type Trie = I.IntMap TrieNode

data TrieNode = TrieNode
    { TrieNode -> Int
trieIndex :: {-# UNPACK #-} !Int
    , TrieNode -> Trie
trieSub   :: !Trie
    }

emptyNode :: TrieNode
emptyNode :: TrieNode
emptyNode = TrieNode
    { trieIndex :: Int
trieIndex = -Int
1
    , trieSub :: Trie
trieSub = forall a. Monoid a => a
mempty
    }

initialTrie :: Trie
initialTrie :: Trie
initialTrie = forall a. [(Int, a)] -> IntMap a
I.fromList
    [(Int
i, TrieNode
emptyNode { trieIndex :: Int
trieIndex = Int
i }) | Int
i <- [Int
0 .. Int
255]]

lookupUpdate :: V.Vector Word8 -> Int -> Int -> Trie -> (Int, Int, Trie)
lookupUpdate :: Vector Word8 -> Int -> Int -> Trie -> (Int, Int, Trie)
lookupUpdate Vector Word8
vector Int
freeIndex Int
firstIndex Trie
trie =
    forall {a} {b}. (a, b, Maybe Trie) -> (a, b, Trie)
matchUpdate forall a b. (a -> b) -> a -> b
$ Trie -> Int -> Int -> (Int, Int, Maybe Trie)
go Trie
trie Int
0 Int
firstIndex 
  where
    matchUpdate :: (a, b, Maybe Trie) -> (a, b, Trie)
matchUpdate (a
lzwOutputIndex, b
nextReadIndex, Maybe Trie
sub) =
        (a
lzwOutputIndex, b
nextReadIndex, forall a. a -> Maybe a -> a
fromMaybe Trie
trie Maybe Trie
sub)

    maxi :: Int
maxi = forall a. Storable a => Vector a -> Int
V.length Vector Word8
vector
    go :: Trie -> Int -> Int -> (Int, Int, Maybe Trie)
go !Trie
currentTrie !Int
prevIndex !Int
index
      | Int
index forall a. Ord a => a -> a -> Bool
>= Int
maxi = (Int
prevIndex, Int
index, forall a. Maybe a
Nothing)
      | Bool
otherwise = case forall a. Int -> IntMap a -> Maybe a
I.lookup Int
val Trie
currentTrie of
          Just (TrieNode Int
ix Trie
subTable) ->
              let (Int
lzwOutputIndex, Int
nextReadIndex, Maybe Trie
newTable) =
                        Trie -> Int -> Int -> (Int, Int, Maybe Trie)
go Trie
subTable Int
ix forall a b. (a -> b) -> a -> b
$ Int
index forall a. Num a => a -> a -> a
+ Int
1
                  tableUpdater :: Trie -> Trie
tableUpdater Trie
t =
                      forall a. Int -> a -> IntMap a -> IntMap a
I.insert Int
val (Int -> Trie -> TrieNode
TrieNode Int
ix Trie
t) Trie
currentTrie
              in
              (Int
lzwOutputIndex, Int
nextReadIndex, Trie -> Trie
tableUpdater forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Trie
newTable)

          Maybe TrieNode
Nothing | Int
index forall a. Eq a => a -> a -> Bool
== Int
maxi -> (Int
prevIndex, Int
index, forall a. Maybe a
Nothing)
                  | Bool
otherwise -> (Int
prevIndex, Int
index, forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> IntMap a -> IntMap a
I.insert Int
val TrieNode
newNode Trie
currentTrie)

      where val :: Int
val = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ Vector Word8
vector forall a. Storable a => Vector a -> Int -> a
`V.unsafeIndex` Int
index
            newNode :: TrieNode
newNode = TrieNode
emptyNode { trieIndex :: Int
trieIndex = Int
freeIndex }

lzwEncode :: Int -> V.Vector Word8 -> L.ByteString
lzwEncode :: Int -> Vector Word8 -> ByteString
lzwEncode Int
initialKeySize Vector Word8
vec = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
    BoolWriteStateRef s
bitWriter <- forall s. ST s (BoolWriteStateRef s)
newWriteStateRef 

    let updateCodeSize :: Int -> Int -> Trie -> ST s (Int, Int, Trie)
updateCodeSize Int
12 Int
writeIdx Trie
_
            | Int
writeIdx forall a. Eq a => a -> a -> Bool
== Int
2 forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
12 :: Int) forall a. Num a => a -> a -> a
- Int
1 = do
               forall s. BoolWriteStateRef s -> Word32 -> Int -> ST s ()
writeBitsGif BoolWriteStateRef s
bitWriter (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
clearCode) Int
12
               forall (m :: * -> *) a. Monad m => a -> m a
return (Int
startCodeSize, Int
firstFreeIndex, Trie
initialTrie)

        updateCodeSize Int
codeSize Int
writeIdx Trie
trie
            | Int
writeIdx forall a. Eq a => a -> a -> Bool
== Int
2 forall a b. (Num a, Integral b) => a -> b -> a
^ Int
codeSize =
                forall (m :: * -> *) a. Monad m => a -> m a
return (Int
codeSize forall a. Num a => a -> a -> a
+ Int
1, Int
writeIdx forall a. Num a => a -> a -> a
+ Int
1, Trie
trie)
            | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return (Int
codeSize, Int
writeIdx forall a. Num a => a -> a -> a
+ Int
1, Trie
trie)

        go :: Int -> (Int, Int, Trie) -> ST s ()
go Int
readIndex (Int
codeSize, Int
_, Trie
_) | Int
readIndex forall a. Ord a => a -> a -> Bool
>= Int
maxi =
            forall s. BoolWriteStateRef s -> Word32 -> Int -> ST s ()
writeBitsGif BoolWriteStateRef s
bitWriter (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
endOfInfo) Int
codeSize
        go !Int
readIndex (!Int
codeSize, !Int
writeIndex, !Trie
trie) = do
            let (Int
indexToWrite, Int
endIndex, Trie
trie') =
                    Int -> Int -> Trie -> (Int, Int, Trie)
lookuper Int
writeIndex Int
readIndex Trie
trie
            forall s. BoolWriteStateRef s -> Word32 -> Int -> ST s ()
writeBitsGif BoolWriteStateRef s
bitWriter (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
indexToWrite) Int
codeSize
            Int -> Int -> Trie -> ST s (Int, Int, Trie)
updateCodeSize Int
codeSize Int
writeIndex Trie
trie'
                forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> (Int, Int, Trie) -> ST s ()
go Int
endIndex 

    forall s. BoolWriteStateRef s -> Word32 -> Int -> ST s ()
writeBitsGif BoolWriteStateRef s
bitWriter (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
clearCode) Int
startCodeSize
    Int -> (Int, Int, Trie) -> ST s ()
go Int
0 (Int
startCodeSize, Int
firstFreeIndex, Trie
initialTrie)

    forall s. BoolWriteStateRef s -> ST s ByteString
finalizeBoolWriterGif BoolWriteStateRef s
bitWriter
  where
    maxi :: Int
maxi = forall a. Storable a => Vector a -> Int
V.length Vector Word8
vec

    startCodeSize :: Int
startCodeSize = Int
initialKeySize forall a. Num a => a -> a -> a
+ Int
1

    clearCode :: Int
clearCode = Int
2 forall a b. (Num a, Integral b) => a -> b -> a
^ Int
initialKeySize :: Int
    endOfInfo :: Int
endOfInfo = Int
clearCode forall a. Num a => a -> a -> a
+ Int
1
    firstFreeIndex :: Int
firstFreeIndex = Int
endOfInfo forall a. Num a => a -> a -> a
+ Int
1
    
    lookuper :: Int -> Int -> Trie -> (Int, Int, Trie)
lookuper = Vector Word8 -> Int -> Int -> Trie -> (Int, Int, Trie)
lookupUpdate Vector Word8
vec