{-# 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 -> IntMap TrieNode
trieSub   :: !Trie
    }

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

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

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

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

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

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

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

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

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

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

    writeBitsGif bitWriter (fromIntegral clearCode) startCodeSize
    go 0 (startCodeSize, firstFreeIndex, initialTrie)

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

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

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