{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
module Codec.Picture.Jpg.Internal.FastIdct( MutableMacroBlock
, fastIdct
, mutableLevelShift
, createEmptyMutableMacroBlock
) where
import qualified Data.Vector.Storable as V
import Control.Monad.ST( ST )
import Data.Bits( unsafeShiftL, unsafeShiftR )
import Data.Int( Int16 )
import qualified Data.Vector.Storable.Mutable as M
import Codec.Picture.Jpg.Internal.Types
iclip :: V.Vector Int16
iclip :: Vector Int16
iclip = Int -> [Int16] -> Vector Int16
forall a. Storable a => Int -> [a] -> Vector a
V.fromListN Int
1024 [ Int16 -> Int16
forall {a}. (Ord a, Num a) => a -> a
val Int16
i| Int16
i <- [(-Int16
512) .. Int16
511] ]
where val :: a -> a
val a
i | a
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< (-a
256) = -a
256
| a
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
255 = a
255
| Bool
otherwise = a
i
data IDctStage = IDctStage {
IDctStage -> Int
x0 :: {-# UNPACK #-} !Int,
IDctStage -> Int
x1 :: {-# UNPACK #-} !Int,
IDctStage -> Int
x2 :: {-# UNPACK #-} !Int,
IDctStage -> Int
x3 :: {-# UNPACK #-} !Int,
IDctStage -> Int
x4 :: {-# UNPACK #-} !Int,
IDctStage -> Int
x5 :: {-# UNPACK #-} !Int,
IDctStage -> Int
x6 :: {-# UNPACK #-} !Int,
IDctStage -> Int
x7 :: {-# UNPACK #-} !Int,
IDctStage -> Int
x8 :: {-# UNPACK #-} !Int
}
w1, w2, w3, w5, w6, w7 :: Int
w1 :: Int
w1 = Int
2841
w2 :: Int
w2 = Int
2676
w3 :: Int
w3 = Int
2408
w5 :: Int
w5 = Int
1609
w6 :: Int
w6 = Int
1108
w7 :: Int
w7 = Int
565
idctRow :: MutableMacroBlock s Int16 -> Int -> ST s ()
idctRow :: forall s. MutableMacroBlock s Int16 -> Int -> ST s ()
idctRow MutableMacroBlock s Int16
blk Int
idx = do
xx0 <- MutableMacroBlock s Int16
MVector (PrimState (ST s)) Int16
blk MVector (PrimState (ST s)) Int16 -> Int -> ST s Int16
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
`M.unsafeRead` (Int
0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
idx)
xx1 <- blk `M.unsafeRead` (4 + idx)
xx2 <- blk `M.unsafeRead` (6 + idx)
xx3 <- blk `M.unsafeRead` (2 + idx)
xx4 <- blk `M.unsafeRead` (1 + idx)
xx5 <- blk `M.unsafeRead` (7 + idx)
xx6 <- blk `M.unsafeRead` (5 + idx)
xx7 <- blk `M.unsafeRead` (3 + idx)
let initialState = IDctStage { x0 :: Int
x0 = (Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx0 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
11) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
128
, x1 :: Int
x1 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
11
, x2 :: Int
x2 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx2
, x3 :: Int
x3 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx3
, x4 :: Int
x4 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx4
, x5 :: Int
x5 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx5
, x6 :: Int
x6 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx6
, x7 :: Int
x7 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx7
, x8 :: Int
x8 = Int
0
}
firstStage IDctStage
c = IDctStage
c { x4 = x8' + (w1 - w7) * x4 c
, x5 = x8' - (w1 + w7) * x5 c
, x6 = x8'' - (w3 - w5) * x6 c
, x7 = x8'' - (w3 + w5) * x7 c
, x8 = x8''
}
where x8' :: Int
x8' = Int
w7 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (IDctStage -> Int
x4 IDctStage
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ IDctStage -> Int
x5 IDctStage
c)
x8'' :: Int
x8'' = Int
w3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (IDctStage -> Int
x6 IDctStage
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ IDctStage -> Int
x7 IDctStage
c)
secondStage IDctStage
c = IDctStage
c { x0 = x0 c - x1 c
, x8 = x0 c + x1 c
, x1 = x1''
, x2 = x1' - (w2 + w6) * x2 c
, x3 = x1' + (w2 - w6) * x3 c
, x4 = x4 c - x6 c
, x6 = x5 c + x7 c
, x5 = x5 c - x7 c
}
where x1' :: Int
x1' = Int
w6 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (IDctStage -> Int
x3 IDctStage
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ IDctStage -> Int
x2 IDctStage
c)
x1'' :: Int
x1'' = IDctStage -> Int
x4 IDctStage
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ IDctStage -> Int
x6 IDctStage
c
thirdStage IDctStage
c = IDctStage
c { x7 = x8 c + x3 c
, x8 = x8 c - x3 c
, x3 = x0 c + x2 c
, x0 = x0 c - x2 c
, x2 = (181 * (x4 c + x5 c) + 128) `unsafeShiftR` 8
, x4 = (181 * (x4 c - x5 c) + 128) `unsafeShiftR` 8
}
scaled IDctStage
c = IDctStage
c { x0 = (x7 c + x1 c) `unsafeShiftR` 8
, x1 = (x3 c + x2 c) `unsafeShiftR` 8
, x2 = (x0 c + x4 c) `unsafeShiftR` 8
, x3 = (x8 c + x6 c) `unsafeShiftR` 8
, x4 = (x8 c - x6 c) `unsafeShiftR` 8
, x5 = (x0 c - x4 c) `unsafeShiftR` 8
, x6 = (x3 c - x2 c) `unsafeShiftR` 8
, x7 = (x7 c - x1 c) `unsafeShiftR` 8
}
transformed = IDctStage -> IDctStage
scaled (IDctStage -> IDctStage)
-> (IDctStage -> IDctStage) -> IDctStage -> IDctStage
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IDctStage -> IDctStage
thirdStage (IDctStage -> IDctStage)
-> (IDctStage -> IDctStage) -> IDctStage -> IDctStage
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IDctStage -> IDctStage
secondStage (IDctStage -> IDctStage) -> IDctStage -> IDctStage
forall a b. (a -> b) -> a -> b
$ IDctStage -> IDctStage
firstStage IDctStage
initialState
(blk `M.unsafeWrite` (0 + idx)) . fromIntegral $ x0 transformed
(blk `M.unsafeWrite` (1 + idx)) . fromIntegral $ x1 transformed
(blk `M.unsafeWrite` (2 + idx)) . fromIntegral $ x2 transformed
(blk `M.unsafeWrite` (3 + idx)) . fromIntegral $ x3 transformed
(blk `M.unsafeWrite` (4 + idx)) . fromIntegral $ x4 transformed
(blk `M.unsafeWrite` (5 + idx)) . fromIntegral $ x5 transformed
(blk `M.unsafeWrite` (6 + idx)) . fromIntegral $ x6 transformed
(blk `M.unsafeWrite` (7 + idx)) . fromIntegral $ x7 transformed
idctCol :: MutableMacroBlock s Int16 -> Int -> ST s ()
idctCol :: forall s. MutableMacroBlock s Int16 -> Int -> ST s ()
idctCol MutableMacroBlock s Int16
blk Int
idx = do
xx0 <- MutableMacroBlock s Int16
MVector (PrimState (ST s)) Int16
blk MVector (PrimState (ST s)) Int16 -> Int -> ST s Int16
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
`M.unsafeRead` ( Int
0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
idx)
xx1 <- blk `M.unsafeRead` (8 * 4 + idx)
xx2 <- blk `M.unsafeRead` (8 * 6 + idx)
xx3 <- blk `M.unsafeRead` (8 * 2 + idx)
xx4 <- blk `M.unsafeRead` (8 + idx)
xx5 <- blk `M.unsafeRead` (8 * 7 + idx)
xx6 <- blk `M.unsafeRead` (8 * 5 + idx)
xx7 <- blk `M.unsafeRead` (8 * 3 + idx)
let initialState = IDctStage { x0 :: Int
x0 = (Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx0 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
8) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
8192
, x1 :: Int
x1 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
8
, x2 :: Int
x2 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx2
, x3 :: Int
x3 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx3
, x4 :: Int
x4 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx4
, x5 :: Int
x5 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx5
, x6 :: Int
x6 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx6
, x7 :: Int
x7 = Int16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
xx7
, x8 :: Int
x8 = Int
0
}
firstStage IDctStage
c = IDctStage
c { x4 = (x8' + (w1 - w7) * x4 c) `unsafeShiftR` 3
, x5 = (x8' - (w1 + w7) * x5 c) `unsafeShiftR` 3
, x6 = (x8'' - (w3 - w5) * x6 c) `unsafeShiftR` 3
, x7 = (x8'' - (w3 + w5) * x7 c) `unsafeShiftR` 3
, x8 = x8''
}
where x8' :: Int
x8' = Int
w7 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (IDctStage -> Int
x4 IDctStage
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ IDctStage -> Int
x5 IDctStage
c) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4
x8'' :: Int
x8'' = Int
w3 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (IDctStage -> Int
x6 IDctStage
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ IDctStage -> Int
x7 IDctStage
c) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4
secondStage IDctStage
c = IDctStage
c { x8 = x0 c + x1 c
, x0 = x0 c - x1 c
, x2 = (x1' - (w2 + w6) * x2 c) `unsafeShiftR` 3
, x3 = (x1' + (w2 - w6) * x3 c) `unsafeShiftR` 3
, x4 = x4 c - x6 c
, x1 = x1''
, x6 = x5 c + x7 c
, x5 = x5 c - x7 c
}
where x1' :: Int
x1' = Int
w6 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (IDctStage -> Int
x3 IDctStage
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ IDctStage -> Int
x2 IDctStage
c) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4
x1'' :: Int
x1'' = IDctStage -> Int
x4 IDctStage
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ IDctStage -> Int
x6 IDctStage
c
thirdStage IDctStage
c = IDctStage
c { x7 = x8 c + x3 c
, x8 = x8 c - x3 c
, x3 = x0 c + x2 c
, x0 = x0 c - x2 c
, x2 = (181 * (x4 c + x5 c) + 128) `unsafeShiftR` 8
, x4 = (181 * (x4 c - x5 c) + 128) `unsafeShiftR` 8
}
clip Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
511 = if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> -Int
512 then Vector Int16
iclip Vector Int16 -> Int -> Int16
forall a. Storable a => Vector a -> Int -> a
`V.unsafeIndex` (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
512)
else Vector Int16
iclip Vector Int16 -> Int -> Int16
forall a. Storable a => Vector a -> Int -> a
`V.unsafeIndex` Int
0
| Bool
otherwise = Vector Int16
iclip Vector Int16 -> Int -> Int16
forall a. Storable a => Vector a -> Int -> a
`V.unsafeIndex` Int
1023
f = IDctStage -> IDctStage
thirdStage (IDctStage -> IDctStage)
-> (IDctStage -> IDctStage) -> IDctStage -> IDctStage
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IDctStage -> IDctStage
secondStage (IDctStage -> IDctStage) -> IDctStage -> IDctStage
forall a b. (a -> b) -> a -> b
$ IDctStage -> IDctStage
firstStage IDctStage
initialState
(blk `M.unsafeWrite` (idx + 8*0)) . clip $ (x7 f + x1 f) `unsafeShiftR` 14
(blk `M.unsafeWrite` (idx + 8 )) . clip $ (x3 f + x2 f) `unsafeShiftR` 14
(blk `M.unsafeWrite` (idx + 8*2)) . clip $ (x0 f + x4 f) `unsafeShiftR` 14
(blk `M.unsafeWrite` (idx + 8*3)) . clip $ (x8 f + x6 f) `unsafeShiftR` 14
(blk `M.unsafeWrite` (idx + 8*4)) . clip $ (x8 f - x6 f) `unsafeShiftR` 14
(blk `M.unsafeWrite` (idx + 8*5)) . clip $ (x0 f - x4 f) `unsafeShiftR` 14
(blk `M.unsafeWrite` (idx + 8*6)) . clip $ (x3 f - x2 f) `unsafeShiftR` 14
(blk `M.unsafeWrite` (idx + 8*7)) . clip $ (x7 f - x1 f) `unsafeShiftR` 14
{-# INLINE fastIdct #-}
fastIdct :: MutableMacroBlock s Int16
-> ST s (MutableMacroBlock s Int16)
fastIdct :: forall s.
MutableMacroBlock s Int16 -> ST s (MutableMacroBlock s Int16)
fastIdct MutableMacroBlock s Int16
block = Int -> ST s (MutableMacroBlock s Int16)
rows Int
0
where rows :: Int -> ST s (MutableMacroBlock s Int16)
rows Int
8 = Int -> ST s (MutableMacroBlock s Int16)
cols Int
0
rows Int
i = MutableMacroBlock s Int16 -> Int -> ST s ()
forall s. MutableMacroBlock s Int16 -> Int -> ST s ()
idctRow MutableMacroBlock s Int16
block (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i) ST s ()
-> ST s (MutableMacroBlock s Int16)
-> ST s (MutableMacroBlock s Int16)
forall a b. ST s a -> ST s b -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> ST s (MutableMacroBlock s Int16)
rows (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
cols :: Int -> ST s (MutableMacroBlock s Int16)
cols Int
8 = MutableMacroBlock s Int16 -> ST s (MutableMacroBlock s Int16)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return MutableMacroBlock s Int16
block
cols Int
i = MutableMacroBlock s Int16 -> Int -> ST s ()
forall s. MutableMacroBlock s Int16 -> Int -> ST s ()
idctCol MutableMacroBlock s Int16
block Int
i ST s ()
-> ST s (MutableMacroBlock s Int16)
-> ST s (MutableMacroBlock s Int16)
forall a b. ST s a -> ST s b -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> ST s (MutableMacroBlock s Int16)
cols (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE mutableLevelShift #-}
mutableLevelShift :: MutableMacroBlock s Int16
-> ST s (MutableMacroBlock s Int16)
mutableLevelShift :: forall s.
MutableMacroBlock s Int16 -> ST s (MutableMacroBlock s Int16)
mutableLevelShift MutableMacroBlock s Int16
block = Int -> ST s (MutableMacroBlock s Int16)
update Int
0
where update :: Int -> ST s (MutableMacroBlock s Int16)
update Int
64 = MutableMacroBlock s Int16 -> ST s (MutableMacroBlock s Int16)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return MutableMacroBlock s Int16
block
update Int
idx = do
val <- MutableMacroBlock s Int16
MVector (PrimState (ST s)) Int16
block MVector (PrimState (ST s)) Int16 -> Int -> ST s Int16
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> m a
`M.unsafeRead` Int
idx
(block `M.unsafeWrite` idx) $ val + 128
update $ idx + 1