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