{-# LANGUAGE NoImplicitPrelude #-}

-- | Computes the Reed-Solomon error correction code words for a sequence of data code words at a given degree.

module Codec.QRCode.Code.ReedSolomonEncoder
  ( RsGeneratorPolynomial
  , rsGeneratorPolynomial
  , rsEncode
  ) where

import           Codec.QRCode.Base

import qualified Data.Vector.Unboxed         as UV
import qualified Data.Vector.Unboxed.Mutable as MUV

newtype RsGeneratorPolynomial
  = RsGeneratorPolynomial (UV.Vector Word8)

-- | Creates a Reed-Solomon ECC generator for the specified degree.
rsGeneratorPolynomial :: Int -> RsGeneratorPolynomial
rsGeneratorPolynomial :: Int -> RsGeneratorPolynomial
rsGeneratorPolynomial Int
degree = (forall s. ST s RsGeneratorPolynomial) -> RsGeneratorPolynomial
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s RsGeneratorPolynomial) -> RsGeneratorPolynomial)
-> (forall s. ST s RsGeneratorPolynomial) -> RsGeneratorPolynomial
forall a b. (a -> b) -> a -> b
$ do
  MVector s Word8
coefficients <- Int -> ST s (MVector (PrimState (ST s)) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MUV.new Int
degree
  -- Start with the monomial x^0
  MVector (PrimState (ST s)) Word8 -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> a -> m ()
MUV.set MVector s Word8
MVector (PrimState (ST s)) Word8
coefficients Word8
0
  MVector (PrimState (ST s)) Word8 -> Int -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUV.write MVector s Word8
MVector (PrimState (ST s)) Word8
coefficients (Int
degreeInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Word8
1

  -- Compute the product polynomial (x - r^0) * (x - r^1) * (x - r^2) * ... * (x - r^{degree-1}),
  -- drop the highest term, and store the rest of the coefficients in order of descending powers.
  -- Note that r = 0x02, which is a generator element of this field GF(2^8/0x11D).
  ST s Word8 -> ST s ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ST s Word8 -> ST s ()) -> ST s Word8 -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int -> Word8 -> (Word8 -> ST s Word8) -> ST s Word8
forall (m :: * -> *) a. Monad m => Int -> a -> (a -> m a) -> m a
iterateNM Int
degree Word8
1 ((Word8 -> ST s Word8) -> ST s Word8)
-> (Word8 -> ST s Word8) -> ST s Word8
forall a b. (a -> b) -> a -> b
$ \Word8
root -> do
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
degreeInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
2] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> do
      Word8
next <- MVector (PrimState (ST s)) Word8 -> Int -> ST s Word8
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUV.read MVector s Word8
MVector (PrimState (ST s)) Word8
coefficients (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
      MVector (PrimState (ST s)) Word8
-> (Word8 -> Word8) -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MUV.modify MVector s Word8
MVector (PrimState (ST s)) Word8
coefficients (\Word8
c -> Word8 -> Word8 -> Word8
multiply Word8
c Word8
root Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
`xor` Word8
next) Int
j
    -- calc last (does not have a next)
    MVector (PrimState (ST s)) Word8
-> (Word8 -> Word8) -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MUV.modify MVector s Word8
MVector (PrimState (ST s)) Word8
coefficients (Word8 -> Word8 -> Word8
multiply Word8
root) (Int
degreeInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
    Word8 -> ST s Word8
forall (m :: * -> *) a. Monad m => a -> m a
return (Word8 -> Word8 -> Word8
multiply Word8
root Word8
0x02)
  Vector Word8 -> RsGeneratorPolynomial
RsGeneratorPolynomial (Vector Word8 -> RsGeneratorPolynomial)
-> ST s (Vector Word8) -> ST s RsGeneratorPolynomial
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Word8 -> ST s (Vector Word8)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
UV.unsafeFreeze MVector s Word8
MVector (PrimState (ST s)) Word8
coefficients

  where
    iterateNM :: Monad m => Int -> a -> (a -> m a) -> m a
    iterateNM :: Int -> a -> (a -> m a) -> m a
iterateNM Int
n0 a
i0 a -> m a
f = Int -> a -> m a
forall t. (Ord t, Num t) => t -> a -> m a
go Int
n0 a
i0
      where
        go :: t -> a -> m a
go t
n a
i
          | t
n t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
0 = a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
i
          | Bool
otherwise = t -> a -> m a
go (t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
1) (a -> m a) -> m a -> m a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< a -> m a
f a
i

-- | Computes and returns the Reed-Solomon error correction code words for the specified sequence of data codewords.
rsEncode :: RsGeneratorPolynomial -> [Word8] -> [Word8]
rsEncode :: RsGeneratorPolynomial -> [Word8] -> [Word8]
rsEncode (RsGeneratorPolynomial Vector Word8
coefficients) [Word8]
input = (forall s. ST s [Word8]) -> [Word8]
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s [Word8]) -> [Word8])
-> (forall s. ST s [Word8]) -> [Word8]
forall a b. (a -> b) -> a -> b
$ do
  let
    len :: Int
len = Vector Word8 -> Int
forall a. Unbox a => Vector a -> Int
UV.length Vector Word8
coefficients
  MVector s Word8
result <- Int -> ST s (MVector (PrimState (ST s)) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MUV.new Int
len
  MVector (PrimState (ST s)) Word8 -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> a -> m ()
MUV.set MVector s Word8
MVector (PrimState (ST s)) Word8
result Word8
0
  [Word8] -> (Word8 -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Word8]
input ((Word8 -> ST s ()) -> ST s ()) -> (Word8 -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Word8
b -> do
    Word8
r0 <- MVector (PrimState (ST s)) Word8 -> Int -> ST s Word8
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUV.read MVector s Word8
MVector (PrimState (ST s)) Word8
result Int
0
    let
      factor :: Word8
factor = Word8
b Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
`xor` Word8
r0
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
1 .. Int
lenInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
      Word8
t <- MVector (PrimState (ST s)) Word8 -> Int -> ST s Word8
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUV.read MVector s Word8
MVector (PrimState (ST s)) Word8
result Int
i
      MVector (PrimState (ST s)) Word8 -> Int -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUV.write MVector s Word8
MVector (PrimState (ST s)) Word8
result (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Word8
t
    MVector (PrimState (ST s)) Word8 -> Int -> Word8 -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUV.write MVector s Word8
MVector (PrimState (ST s)) Word8
result (Int
lenInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Word8
0
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
lenInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
      MVector (PrimState (ST s)) Word8
-> (Word8 -> Word8) -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MUV.modify MVector s Word8
MVector (PrimState (ST s)) Word8
result (\Word8
rx -> Word8
rx Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
`xor` Word8 -> Word8 -> Word8
multiply (Vector Word8
coefficients Vector Word8 -> Int -> Word8
forall a. Unbox a => Vector a -> Int -> a
UV.! Int
i) Word8
factor) Int
i
  Vector Word8
result' <- MVector (PrimState (ST s)) Word8 -> ST s (Vector Word8)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
UV.unsafeFreeze MVector s Word8
MVector (PrimState (ST s)) Word8
result
  [Word8] -> ST s [Word8]
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector Word8 -> [Word8]
forall a. Unbox a => Vector a -> [a]
UV.toList Vector Word8
result')

-- | Returns the product of the two given field elements modulo GF(2^8/0x11D).
multiply :: Word8 -> Word8 -> Word8
{-# INLINABLE multiply #-}
multiply :: Word8 -> Word8 -> Word8
multiply Word8
x Word8
y =
  let
    step :: Word8 -> Int -> Word8
step Word8
z Int
i =
      (Word8
z Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
1) Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
`xor` ((Word8
z Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftR` Int
7) Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
* Word8
0x1d)
      Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
`xor` (((Word8
y Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftR` Int
i) Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
1) Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
* Word8
x)
  in
    (Word8 -> Int -> Word8) -> Word8 -> [Int] -> Word8
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Word8 -> Int -> Word8
step Word8
0 [Int
7, Int
6 .. Int
0]