Reed-Solomon encoding / decoding module

symbolic-base/src/ZkFold/Base/Algorithm/ReedSolomon.hs

@@ -0,0 +1,97 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+
+module ZkFold.Base.Algorithm.ReedSolomon where
+
+
+import           Data.Bool                                  (bool)
+import           Data.Vector                                as V hiding (sum)
+import           GHC.Natural                                (Natural)
+import           Prelude                                    (Eq, Int, Integer, Maybe (..), error, fromIntegral, iterate,
+                                                             min, ($), (.), (<=), (==))
+import qualified Prelude                                    as P
+
+import           ZkFold.Base.Algebra.Basic.Class
+import           ZkFold.Base.Algebra.Basic.Number           (KnownNat, value)
+import           ZkFold.Base.Algebra.Polynomials.Univariate
+
+
+numberOfError :: forall n k. (KnownNat n, KnownNat k) => Natural
+numberOfError = (value @n -! value @k) `div` 2
+
+generator :: (Field a, Eq a) => Int -> a -> Poly a
+generator r a = V.foldr (\ai pi -> toLinPoly ai * pi) one roots
+    where
+        roots = V.iterateN r (* a) a
+        toLinPoly p = toPoly $ fromList [negate p, one]
+
+encode :: (Field c, Eq c) => [c] -> c -> Int -> Poly c
+encode msg prim_elem r = msg_padded - reminder
+    where
+        g_x = generator r prim_elem
+        poly_msg = toPoly $ fromList msg
+        msg_padded = scaleP one (fromIntegral r) poly_msg
+        (_, reminder) = qr msg_padded g_x
+
+-- beta = one
+decode :: (Field c, Eq c) => Poly c -> c -> Int -> Int -> Poly c
+decode encoded primeElement r n = bool decoded encoded' isCorrect
+    where
+        rElems = iterateN r (* primeElement) primeElement
+
+        encoded' = toPoly . V.drop r $ fromPoly encoded
+        syndromes = toPoly . V.map (evalPoly encoded) $ rElems
+        isCorrect = zero == syndromes
+
+        (_, lx) = berlekamp syndromes r
+
+        invPE = finv primeElement
+        es1 = V.indexed . V.fromList $ one : P.take (n P.- 1) (iterate (* invPE) invPE)
+        iroots = mapMaybe (\(i,x) -> bool Nothing (Just (i , x)) (evalPoly lx x == zero)) es1
+
+        omega = toPoly $ take r $ fromPoly (lx * syndromes)
+        lx'= diff lx
+
+        err = V.foldl (+) zero $ map (\(i,x) ->
+            let xi = bool (monomial (fromIntegral i) one) (constant one) (i == 0)
+                ei = evalPoly omega x * finv (evalPoly lx' x)
+            in constant ei * xi) iroots
+
+        fx = encoded + err
+        checkSum = V.map (evalPoly fx) rElems
+
+        decoded = bool (error "Can't decode") (toPoly $ V.drop r $ fromPoly fx) (all (== zero) checkSum)
+
+
+berlekamp :: forall c . (Field c, Eq c) => Poly c -> Int -> (Integer, Poly c)
+berlekamp s r
+    | deg s == -1 = (0, one)
+    | P.otherwise = go сx0 bx0 0 0 1 one
+    where
+        sv = fromPoly s
+        сx0 = one :: Poly c
+        bx0 = one :: Poly c
+        lenS = fromIntegral r
+
+        go :: Poly c -> Poly c -> Integer -> Integer -> Natural -> c -> (Integer, Poly c)
+        go cx bx n l m b
+            | n == lenS = bool (error "locators didn't find") (l, cx) (deg cx == l)
+            | P.otherwise = bool innerChoice (go cx bx n' l (m+1) b) (d == zero)
+            where
+                d = scalarN (fromIntegral n P.+ 1) (fromIntegral l P.+1) lxv sv
+
+                lxv = fromPoly cx
+                cx' = cx - fromConstant d * constant (finv b) * bx * monomial m one
+
+                n' = n + 1
+                innerChoice = bool (go cx' bx n' l (m+1) b) (go cx' cx n' (n+1-l) 1 d) (2*l <= n)
+
+
+scalarN :: (Semiring c) => Int -> Int -> Vector c -> Vector c -> c
+scalarN q l lv rv = bool (sum $ V.zipWith (*) lPadded rPadded) zero (min l (length lv) <= q P.- length rv)
+    where
+        lPadded = V.drop (q P.- V.length rv) lv
+        rPadded = V.reverse $ V.take q rv
+
+
+diff :: ( Field c, Eq c) =>Poly c -> Poly c
+diff p = let cs = fromPoly p in toPoly $ V.tail $ V.imap (\i c -> scale (fromIntegral i :: Integer) c) cs

symbolic-base/symbolic-base.cabal

@@ -109,6 +109,7 @@ library
       ZkFold.Base.Algebra.Polynomials.Multivariate.Polynomial
       ZkFold.Base.Algebra.Polynomials.Multivariate.Substitution
       ZkFold.Base.Algebra.Polynomials.Univariate
+      ZkFold.Base.Algorithm.ReedSolomon
       ZkFold.Base.Control.HApplicative
       ZkFold.Base.Data.ByteString
       ZkFold.Base.Data.Functor.Rep
@@ -271,6 +272,7 @@ test-suite symbolic-base-test
       Tests.Pairing
       Tests.Permutations
       Tests.Plonkup
+      Tests.ReedSolomon
       Tests.RSA
       Tests.SHA2
       Tests.UInt

symbolic-base/test/Main.hs

@@ -19,6 +19,7 @@ import           Tests.NonInteractiveProof (specNonInteractiveProof)
 import           Tests.Pairing             (specPairing)
 import           Tests.Permutations        (specPermutations)
 import           Tests.Plonkup             (specPlonkup)
+import           Tests.ReedSolomon         (specReedSolomon)
 import           Tests.RSA                 (specRSA)
 import           Tests.SHA2                (specSHA2, specSHA2Natural)
 import           Tests.UInt                (specUInt)
@@ -37,6 +38,7 @@ main = do
     specUnivariate
     specGroebner
     specEllipticCurve
+    specReedSolomon
 
     -- Compiler spec
     specCompiler

symbolic-base/test/Tests/ReedSolomon.hs

@@ -0,0 +1,125 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+
+
+module Tests.ReedSolomon where
+
+import           Data.Bool                                   (bool)
+import           Data.List                                   (sort)
+import           Data.Typeable                               (Typeable, typeOf)
+import qualified Data.Vector                                 as V
+import           GHC.Natural                                 (Natural)
+import           Prelude
+import qualified Prelude                                     as P
+import           Test.Hspec
+import           Test.QuickCheck
+
+import qualified ZkFold.Base.Algebra.Basic.Class             as C
+import           ZkFold.Base.Algebra.Basic.Class             hiding ((*), (+))
+import qualified ZkFold.Base.Algebra.EllipticCurve.BLS12_381 as BLS12_381
+import qualified ZkFold.Base.Algebra.EllipticCurve.BN254     as BN254
+import qualified ZkFold.Base.Algebra.EllipticCurve.Pasta     as Pasta
+import           ZkFold.Base.Algebra.Polynomials.Univariate
+import           ZkFold.Base.Algorithm.ReedSolomon
+
+
+data ReedSolomonExample f = ReedSolomonExample
+    {
+        pe  :: f,       -- primitive element
+        k   :: Int,     -- message length
+        r   :: Int,     -- number of errors * 2
+        msg :: [f],     -- message
+        err :: [f]      -- error polynomial
+    } deriving (Show)
+
+instance (Arbitrary f, FiniteField f, Eq f) => Arbitrary (ReedSolomonExample f) where
+  arbitrary = do
+    let pe = fromConstant (3 :: Natural)
+        c = fromIntegral $ min 100 (order @f)
+    t <- chooseInt (1,P.div c 3)
+    let r = 2*t
+        k = c P.- r P.- 1
+    msg <- vector @f k
+    er <- polyErr (k+t) t
+    return $ ReedSolomonExample pe k r msg er
+
+polyErr :: forall c. (Arbitrary c, Field c) => Int -> Int -> Gen [c]
+polyErr kt t = do
+    pos <- shuffle (replicate t (1 :: Integer) P.++ replicate kt 0)
+    mapM (\p -> bool arbitrary (pure zero) (p==0) ) pos
+
+----------------------------------------------------------------------------------------
+
+propGenerator :: forall c. (FiniteField c, Eq c) => ReedSolomonExample c -> Bool
+propGenerator ReedSolomonExample {..} =
+    let vals = take r $ iterate (C.* pe) (pe :: c)
+        polyGen = generator r pe
+    in all (\x -> evalPoly polyGen x == zero) vals
+
+propEncoder :: forall c. (Field c, Eq c) => ReedSolomonExample c -> Bool
+propEncoder ReedSolomonExample {..} =
+    let encodedMsg = encode msg pe r
+        reminder = snd $ qr encodedMsg (generator r pe)
+    in deg reminder == -1
+
+propBerlekampNoError :: forall c. (FiniteField c, Eq c) => ReedSolomonExample c -> Bool
+propBerlekampNoError ReedSolomonExample {..} =
+    let
+        vals = V.iterateN r (C.* pe) pe
+        encodedMsg = encode msg pe r
+        syndromes = toPoly $ V.map (evalPoly encodedMsg) vals
+    in syndromes == zero && berlekamp syndromes r == (0, one)
+
+propBerlekampWithErrors :: forall c. (FiniteField c, Ord c) => ReedSolomonExample c -> Bool
+propBerlekampWithErrors ReedSolomonExample {..} =
+    let
+        vals = V.iterateN r (C.* pe) pe
+        encoded' = encode msg pe r
+        errorMsg = toPoly $ V.fromList err
+        encoded = encoded' C.+ errorMsg
+        syndromes = toPoly $ V.map (evalPoly encoded) vals
+        bl = snd $ berlekamp syndromes r
+        roots = filter (/= zero) $ map (\x -> bool zero x (evalPoly bl x == zero)) $ takeWhile (/= pe) $ iterate (C.* pe) pe
+
+        es1 = takeWhile (/= one) $ iterate (C.* pe) one
+        rightLocators = filter (/= zero) $ zipWith (\e p -> bool zero (finv p) (e /= zero)) err es1
+    in sort roots == sort rightLocators
+
+propDecodeWithoutError :: forall c. (FiniteField c, Eq c) => ReedSolomonExample c -> Bool
+propDecodeWithoutError ReedSolomonExample {..} =
+    let encoded = encode msg pe r
+        decoded = decode encoded pe r (k+r)
+    in toPoly (V.fromList msg) == decoded
+
+propDecodeWithError :: forall c. (FiniteField c, Eq c) => ReedSolomonExample c -> Bool
+propDecodeWithError ReedSolomonExample {..} =
+    let encoded' = encode msg pe r
+        errorMsg = toPoly $ V.fromList err
+        encoded = encoded' C.+ errorMsg
+        decoded = decode encoded pe r (k+r)
+    in toPoly (V.fromList msg) == decoded
+
+specReedSolomon':: forall a . (FiniteField a, Ord a, Arbitrary a, Show a, Typeable a) => IO ()
+specReedSolomon' = hspec $ do
+    describe "Reed-Solomon" $ do
+        describe ("Type: " ++ show (typeOf @a zero)) $ do
+            it "generator function is correct" $ do
+                property $ propGenerator @a
+            it "encoder function is correct" $ do
+                property $ propEncoder @a
+            it "berlekamp function is correct without errors" $ do
+                property $ propBerlekampNoError @a
+            it "berlekamp function is correct with errors" $ do
+                property $ propBerlekampWithErrors @a
+            it "decode function is correct without errors" $ do
+                property $ propDecodeWithoutError @a
+            it "decode function is correct with errors" $ do
+                property $ propDecodeWithError @a
+
+specReedSolomon :: IO ()
+specReedSolomon = do
+    specReedSolomon' @BN254.Fr
+    specReedSolomon' @BN254.Fp
+    specReedSolomon' @BLS12_381.Fr
+    specReedSolomon' @BLS12_381.Fq
+    specReedSolomon' @Pasta.Fp
+    specReedSolomon' @Pasta.Fq