content: -----------------------------------------------------------------------------
  -- |
  -- Module    : Data.SBV.Tools.Polynomial
  -- Copyright : (c) Levent Erkok
  -- License   : BSD3
  -- Maintainer: erkokl@gmail.com
  -- Stability : experimental
  --
  -- Implementation of polynomial arithmetic
  -----------------------------------------------------------------------------
  
  {-# LANGUAGE DataKinds            #-}
  {-# LANGUAGE FlexibleContexts     #-}
  {-# LANGUAGE FlexibleInstances    #-}
  {-# LANGUAGE PatternGuards        #-}
  {-# LANGUAGE TypeFamilies         #-}
  {-# LANGUAGE UndecidableInstances #-}
  
  {-# OPTIONS_GHC -Wall -Werror #-}
  
  module Data.SBV.Tools.Polynomial (
          -- * Polynomial arithmetic and CRCs
          Polynomial(..), crc, crcBV, ites, mdp, addPoly
          ) where
  
  import Data.Bits  (Bits(..))
  import Data.List  (genericTake)
  import Data.Maybe (fromJust, fromMaybe)
  import Data.Word  (Word8, Word16, Word32, Word64)
  
  import Data.SBV.Core.Data
  import Data.SBV.Core.Kind
  import Data.SBV.Core.Model
  
  import GHC.TypeLits
  
  -- | Implements polynomial addition, multiplication, division, and modulus operations
  -- over GF(2^n).  NB. Similar to 'sQuotRem', division by @0@ is interpreted as follows:
  --
  --     @x `pDivMod` 0 = (0, x)@
  --
  -- for all @x@ (including @0@)
  --
  -- Minimal complete definition: 'pMult', 'pDivMod', 'showPolynomial'
  class (Num a, Bits a) => Polynomial a where
   -- | Given bit-positions to be set, create a polynomial
   -- For instance
   --
   --     @polynomial [0, 1, 3] :: SWord8@
   --
   -- will evaluate to @11@, since it sets the bits @0@, @1@, and @3@. Mathematicians would write this polynomial
   -- as @x^3 + x + 1@. And in fact, 'showPoly' will show it like that.
   polynomial :: [Int] -> a
   -- | Add two polynomials in GF(2^n).
   pAdd  :: a -> a -> a
   -- | Multiply two polynomials in GF(2^n), and reduce it by the irreducible specified by
   -- the polynomial as specified by coefficients of the third argument. Note that the third
   -- argument is specifically left in this form as it is usually in GF(2^(n+1)), which is not available in our
   -- formalism. (That is, we would need SWord9 for SWord8 multiplication, etc.) Also note that we do not
   -- support symbolic irreducibles, which is a minor shortcoming. (Most GF's will come with fixed irreducibles,
   -- so this should not be a problem in practice.)
   --
   -- Passing [] for the third argument will multiply the polynomials and then ignore the higher bits that won't
   -- fit into the resulting size.
   pMult :: (a, a, [Int]) -> a
   -- | Divide two polynomials in GF(2^n), see above note for division by 0.
   pDiv  :: a -> a -> a
   -- | Compute modulus of two polynomials in GF(2^n), see above note for modulus by 0.
   pMod  :: a -> a -> a
   -- | Division and modulus packed together.
   pDivMod :: a -> a -> (a, a)
   -- | Display a polynomial like a mathematician would (over the monomial @x@), with a type.
   showPoly :: a -> String
   -- | Display a polynomial like a mathematician would (over the monomial @x@), the first argument
   -- controls if the final type is shown as well.
   showPolynomial :: Bool -> a -> String
  
   {-# MINIMAL pMult, pDivMod, showPolynomial #-}
   polynomial = foldr (flip setBit) 0
   pAdd       = xor
   pDiv x y   = fst (pDivMod x y)
   pMod x y   = snd (pDivMod x y)
   showPoly   = showPolynomial False
  
  instance Polynomial Word8   where {showPolynomial   = sp;           pMult = lift polyMult; pDivMod = liftC polyDivMod}
  instance Polynomial Word16  where {showPolynomial   = sp;           pMult = lift polyMult; pDivMod = liftC polyDivMod}
  instance Polynomial Word32  where {showPolynomial   = sp;           pMult = lift polyMult; pDivMod = liftC polyDivMod}
  instance Polynomial Word64  where {showPolynomial   = sp;           pMult = lift polyMult; pDivMod = liftC polyDivMod}
  instance Polynomial SWord8  where {showPolynomial b = liftS (sp b); pMult = polyMult;      pDivMod = polyDivMod}
  instance Polynomial SWord16 where {showPolynomial b = liftS (sp b); pMult = polyMult;      pDivMod = polyDivMod}
  instance Polynomial SWord32 where {showPolynomial b = liftS (sp b); pMult = polyMult;      pDivMod = polyDivMod}
  instance Polynomial SWord64 where {showPolynomial b = liftS (sp b); pMult = polyMult;      pDivMod = polyDivMod}
  
  instance (KnownNat n, BVIsNonZero n) => Polynomial (SWord n) where {showPolynomial b = liftS (sp b); pMult = polyMult;      pDivMod = polyDivMod}
  
  lift :: SymVal a => ((SBV a, SBV a, [Int]) -> SBV a) -> (a, a, [Int]) -> a
  lift f (x, y, z) = fromJust $ unliteral $ f (literal x, literal y, z)
  liftC :: SymVal a => (SBV a -> SBV a -> (SBV a, SBV a)) -> a -> a -> (a, a)
  liftC f x y = let (a, b) = f (literal x) (literal y) in (fromJust (unliteral a), fromJust (unliteral b))
  liftS :: SymVal a => (a -> String) -> SBV a -> String
  liftS f s
    | Just x <- unliteral s = f x
    | True                  = show s
  
  -- | Pretty print as a polynomial
  sp :: Bits a => Bool -> a -> String
  sp st a
   | null cs = '0' : t
   | True    = foldr (\x y -> sh x ++ " + " ++ y) (sh (last cs)) (init cs) ++ t
   where t | st   = " :: GF(2^" ++ show n ++ ")"
           | True = ""
         n  = fromMaybe (error "SBV.Polynomial.sp: Unexpected non-finite usage!") (bitSizeMaybe a)
         is = [n-1, n-2 .. 0]
         cs = map fst $ filter snd $ zip is (map (testBit a) is)
         sh 0 = "1"
         sh 1 = "x"
         sh i = "x^" ++ show i
  
  -- | Add two polynomials
  addPoly :: [SBool] -> [SBool] -> [SBool]
  addPoly xs    []      = xs
  addPoly []    ys      = ys
  addPoly (x:xs) (y:ys) = x .<+> y : addPoly xs ys
  
  -- | Run down a boolean condition over two lists. Note that this is
  -- different than zipWith as shorter list is assumed to be filled with
  -- sFalse at the end (i.e., zero-bits); which nicely pads it when
  -- considered as an unsigned number in little-endian form.
  ites :: SBool -> [SBool] -> [SBool] -> [SBool]
  ites s xs ys
   | Just t <- unliteral s
   = if t then xs else ys
   | True
   = go xs ys
   where go []     []     = []
         go []     (b:bs) = ite s sFalse b : go [] bs
         go (a:as) []     = ite s a sFalse : go as []
         go (a:as) (b:bs) = ite s a b : go as bs
  
  -- | Multiply two polynomials and reduce by the third (concrete) irreducible, given by its coefficients.
  -- See the remarks for the 'pMult' function for this design choice
  polyMult :: SFiniteBits a => (SBV a, SBV a, [Int]) -> SBV a
  polyMult (x, y, red)
    | isReal x
    = error $ "SBV.polyMult: Received a real value: " ++ show x
    | not (isBounded x)
    = error $ "SBV.polyMult: Received infinite precision value: " ++ show x
    | True
    = fromBitsLE $ genericTake sz $ r ++ repeat sFalse
    where (_, r) = mdp ms rs
          ms = genericTake (2*sz) $ mul (blastLE x) (blastLE y) [] ++ repeat sFalse
          rs = genericTake (2*sz) $ [fromBool (i `elem` red) |  i <- [0 .. foldr max 0 red] ] ++ repeat sFalse
          sz = intSizeOf x
          mul _  []     ps = ps
          mul as (b:bs) ps = mul (sFalse:as) bs (ites b (as `addPoly` ps) ps)
  
  polyDivMod :: SFiniteBits a => SBV a -> SBV a -> (SBV a, SBV a)
  polyDivMod x y
     | isReal x
     = error $ "SBV.polyDivMod: Received a real value: " ++ show x
     | not (isBounded x)
     = error $ "SBV.polyDivMod: Received infinite precision value: " ++ show x
     | True
     = ite (y .== 0) (0, x) (adjust d, adjust r)
     where adjust xs = fromBitsLE $ genericTake sz $ xs ++ repeat sFalse
           sz        = intSizeOf x
           (d, r)    = mdp (blastLE x) (blastLE y)
  
  -- conservative over-approximation of the degree
  degree :: [SBool] -> Int
  degree xs = walk (length xs - 1) $ reverse xs
    where walk n []     = n
          walk n (b:bs)
           | Just t <- unliteral b
           = if t then n else walk (n-1) bs
           | True
           = n -- over-estimate
  
  -- | Compute modulus/remainder of polynomials on bit-vectors.
  mdp :: [SBool] -> [SBool] -> ([SBool], [SBool])
  mdp xs ys = go (length ys - 1) (reverse ys)
    where degTop  = degree xs
          go _ []     = error "SBV.Polynomial.mdp: Impossible happened; exhausted ys before hitting 0"
          go n (b:bs)
           | n == 0   = (reverse qs, rs)
           | True     = let (rqs, rrs) = go (n-1) bs
                        in (ites b (reverse qs) rqs, ites b rs rrs)
           where degQuot = degTop - n
                 ys' = replicate degQuot sFalse ++ ys
                 (qs, rs) = divx (degQuot+1) degTop xs ys'
  
  -- return the element at index i; if not enough elements, return sFalse
  -- N.B. equivalent to '(xs ++ repeat sFalse) !! i', but more efficient
  idx :: [SBool] -> Int -> SBool
  idx []     _ = sFalse
  idx (x:_)  0 = x
  idx (_:xs) i = idx xs (i-1)
  
  divx :: Int -> Int -> [SBool] -> [SBool] -> ([SBool], [SBool])
  divx n _ xs _ | n <= 0 = ([], xs)
  divx n i xs ys'        = (q:qs, rs)
    where q        = xs `idx` i
          xs'      = ites q (xs `addPoly` ys') xs
          (qs, rs) = divx (n-1) (i-1) xs' (drop 1 ys')
  
  -- | Compute CRCs over bit-vectors. The call @crcBV n m p@ computes
  -- the CRC of the message @m@ with respect to polynomial @p@. The
  -- inputs are assumed to be blasted big-endian. The number
  -- @n@ specifies how many bits of CRC is needed. Note that @n@
  -- is actually the degree of the polynomial @p@, and thus it seems
  -- redundant to pass it in. However, in a typical proof context,
  -- the polynomial can be symbolic, so we cannot compute the degree
  -- easily. While this can be worked-around by generating code that
  -- accounts for all possible degrees, the resulting code would
  -- be unnecessarily big and complicated, and much harder to reason
  -- with. (Also note that a CRC is just the remainder from the
  -- polynomial division, but this routine is much faster in practice.)
  --
  -- NB. The @n@th bit of the polynomial @p@ /must/ be set for the CRC
  -- to be computed correctly. Note that the polynomial argument @p@ will
  -- not even have this bit present most of the time, as it will typically
  -- contain bits @0@ through @n-1@ as usual in the CRC literature. The higher
  -- order @n@th bit is simply assumed to be set, as it does not make
  -- sense to use a polynomial of a lesser degree. This is usually not a problem
  -- since CRC polynomials are designed and expressed this way.
  --
  -- NB. The literature on CRC's has many variants on how CRC's are computed.
  -- We follow the following simple procedure:
  --
  --     * Extend the message @m@ by adding @n@ 0 bits on the right
  --
  --     * Divide the polynomial thus obtained by the @p@
  --
  --     * The remainder is the CRC value.
  --
  -- There are many variants on final XOR's, reversed polynomials etc., so
  -- it is essential to double check you use the correct /algorithm/.
  crcBV :: Int -> [SBool] -> [SBool] -> [SBool]
  crcBV n m p = take n $ go (replicate n sFalse) (m ++ replicate n sFalse)
    where mask = drop (length p - n) p
          go c []     = c
          go c (b:bs) = go next bs
            where c' = drop 1 c ++ [b]
                  next = ite (hd c) (zipWith (.<+>) c' mask) c'
  
                  hd (f:_) = f
                  hd []    = error "crcBV: Impossible, prefix is empty"
  
  -- | Compute CRC's over polynomials, i.e., symbolic words. The first
  -- 'Int' argument plays the same role as the one in the 'crcBV' function.
  crc :: (SFiniteBits a, SFiniteBits b) => Int -> SBV a -> SBV b -> SBV b
  crc n m p
    | isReal m || isReal p
    = error $ "SBV.crc: Received a real value: " ++ show (m, p)
    | not (isBounded m) || not (isBounded p)
    = error $ "SBV.crc: Received an infinite precision value: " ++ show (m, p)
    | True
    = fromBitsBE $ replicate (sz - n) sFalse ++ crcBV n (blastBE m) (blastBE p)
    where sz = intSizeOf p
  