{-# 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 a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  MVector s Word8
coefficients <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MUV.new Int
degree
  -- Start with the monomial x^0
  forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> a -> m ()
MUV.set MVector s Word8
coefficients Word8
0
  forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUV.write MVector s Word8
coefficients (Int
degreeforall 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).
  forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => Int -> a -> (a -> m a) -> m a
iterateNM Int
degree Word8
1 forall a b. (a -> b) -> a -> b
$ \Word8
root -> do
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
degreeforall a. Num a => a -> a -> a
-Int
2] forall a b. (a -> b) -> a -> b
$ \Int
j -> do
      Word8
next <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUV.read MVector s Word8
coefficients (Int
jforall a. Num a => a -> a -> a
+Int
1)
      forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MUV.modify MVector s Word8
coefficients (\Word8
c -> Word8 -> Word8 -> Word8
multiply Word8
c Word8
root forall a. Bits a => a -> a -> a
`xor` Word8
next) Int
j
    -- calc last (does not have a next)
    forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MUV.modify MVector s Word8
coefficients (Word8 -> Word8 -> Word8
multiply Word8
root) (Int
degreeforall a. Num a => a -> a -> a
-Int
1)
    forall (m :: * -> *) a. Monad m => a -> m a
return (Word8 -> Word8 -> Word8
multiply Word8
root Word8
0x02)
  Vector Word8 -> RsGeneratorPolynomial
RsGeneratorPolynomial forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
UV.unsafeFreeze MVector s Word8
coefficients

  where
    iterateNM :: Monad m => Int -> a -> (a -> m a) -> m a
    iterateNM :: forall (m :: * -> *) a. Monad m => Int -> a -> (a -> m a) -> m a
iterateNM Int
n0 a
i0 a -> m a
f = 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 forall a. Ord a => a -> a -> Bool
<= t
0 = forall (m :: * -> *) a. Monad m => a -> m a
return a
i
          | Bool
otherwise = t -> a -> m a
go (t
nforall a. Num a => a -> a -> a
-t
1) 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 a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  let
    len :: Int
len = forall a. Unbox a => Vector a -> Int
UV.length Vector Word8
coefficients
  MVector s Word8
result <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MUV.new Int
len
  forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> a -> m ()
MUV.set MVector s Word8
result Word8
0
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Word8]
input forall a b. (a -> b) -> a -> b
$ \Word8
b -> do
    Word8
r0 <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUV.read MVector s Word8
result Int
0
    let
      factor :: Word8
factor = Word8
b forall a. Bits a => a -> a -> a
`xor` Word8
r0
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
1 .. Int
lenforall a. Num a => a -> a -> a
-Int
1] forall a b. (a -> b) -> a -> b
$ \Int
i -> do
      Word8
t <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUV.read MVector s Word8
result Int
i
      forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUV.write MVector s Word8
result (Int
iforall a. Num a => a -> a -> a
-Int
1) Word8
t
    forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUV.write MVector s Word8
result (Int
lenforall a. Num a => a -> a -> a
-Int
1) Word8
0
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
lenforall a. Num a => a -> a -> a
-Int
1] forall a b. (a -> b) -> a -> b
$ \Int
i ->
      forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
MUV.modify MVector s Word8
result (\Word8
rx -> Word8
rx forall a. Bits a => a -> a -> a
`xor` Word8 -> Word8 -> Word8
multiply (Vector Word8
coefficients forall a. Unbox a => Vector a -> Int -> a
UV.! Int
i) Word8
factor) Int
i
  Vector Word8
result' <- forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
UV.unsafeFreeze MVector s Word8
result
  forall (m :: * -> *) a. Monad m => a -> m a
return (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 forall a. Bits a => a -> Int -> a
`shiftL` Int
1) forall a. Bits a => a -> a -> a
`xor` ((Word8
z forall a. Bits a => a -> Int -> a
`shiftR` Int
7) forall a. Num a => a -> a -> a
* Word8
0x1d)
      forall a. Bits a => a -> a -> a
`xor` (((Word8
y forall a. Bits a => a -> Int -> a
`shiftR` Int
i) forall a. Bits a => a -> a -> a
.&. Word8
1) forall a. Num a => a -> a -> a
* Word8
x)
  in
    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]