{-# LANGUAGE BangPatterns #-}
-- |
-- Module      : Crypto.Number.Polynomial
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : experimental
-- Portability : Good

module Crypto.Number.Polynomial
    ( Monomial(..)
    -- * polynomial operations
    , Polynomial
    , toList
    , fromList
    , addPoly
    , subPoly
    , mulPoly
    , squarePoly
    , expPoly
    , divPoly
    , negPoly
    ) where

import Data.List (intercalate, sort)
import Data.Vector ((!), Vector)
import qualified Data.Vector as V
import Control.Arrow (first)

data Monomial = Monomial {-# UNPACK #-} !Int !Integer
    deriving (Monomial -> Monomial -> Bool
(Monomial -> Monomial -> Bool)
-> (Monomial -> Monomial -> Bool) -> Eq Monomial
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Monomial -> Monomial -> Bool
$c/= :: Monomial -> Monomial -> Bool
== :: Monomial -> Monomial -> Bool
$c== :: Monomial -> Monomial -> Bool
Eq)

data Polynomial = Polynomial (Vector Monomial)
    deriving (Polynomial -> Polynomial -> Bool
(Polynomial -> Polynomial -> Bool)
-> (Polynomial -> Polynomial -> Bool) -> Eq Polynomial
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Polynomial -> Polynomial -> Bool
$c/= :: Polynomial -> Polynomial -> Bool
== :: Polynomial -> Polynomial -> Bool
$c== :: Polynomial -> Polynomial -> Bool
Eq)

instance Ord Monomial where
    compare :: Monomial -> Monomial -> Ordering
compare (Monomial w1 :: Int
w1 v1 :: Integer
v1) (Monomial w2 :: Int
w2 v2 :: Integer
v2) =
        case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
w1 Int
w2 of
            EQ -> Integer -> Integer -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Integer
v1 Integer
v2
            r :: Ordering
r  -> Ordering
r

instance Show Monomial where
    show :: Monomial -> String
show (Monomial w :: Int
w v :: Integer
v) = Integer -> String
forall a. Show a => a -> String
show Integer
v String -> ShowS
forall a. [a] -> [a] -> [a]
++ "x^" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
w

instance Show Polynomial where
    show :: Polynomial -> String
show (Polynomial p :: Vector Monomial
p) = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate "+" ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ (Monomial -> String) -> [Monomial] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Monomial -> String
forall a. Show a => a -> String
show ([Monomial] -> [String]) -> [Monomial] -> [String]
forall a b. (a -> b) -> a -> b
$ Vector Monomial -> [Monomial]
forall a. Vector a -> [a]
V.toList Vector Monomial
p

toList :: Polynomial -> [Monomial]
toList :: Polynomial -> [Monomial]
toList (Polynomial p :: Vector Monomial
p) = Vector Monomial -> [Monomial]
forall a. Vector a -> [a]
V.toList Vector Monomial
p

fromList :: [Monomial] -> Polynomial
fromList :: [Monomial] -> Polynomial
fromList = Vector Monomial -> Polynomial
Polynomial (Vector Monomial -> Polynomial)
-> ([Monomial] -> Vector Monomial) -> [Monomial] -> Polynomial
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Monomial] -> Vector Monomial
forall a. [a] -> Vector a
V.fromList ([Monomial] -> Vector Monomial)
-> ([Monomial] -> [Monomial]) -> [Monomial] -> Vector Monomial
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Monomial] -> [Monomial]
forall a. [a] -> [a]
reverse ([Monomial] -> [Monomial])
-> ([Monomial] -> [Monomial]) -> [Monomial] -> [Monomial]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Monomial] -> [Monomial]
forall a. Ord a => [a] -> [a]
sort ([Monomial] -> [Monomial])
-> ([Monomial] -> [Monomial]) -> [Monomial] -> [Monomial]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Monomial] -> [Monomial]
filterZero
    where
        filterZero :: [Monomial] -> [Monomial]
filterZero = (Monomial -> Bool) -> [Monomial] -> [Monomial]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(Monomial _ v :: Integer
v) -> Integer
v Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= 0)

getWeight :: Polynomial -> Int -> Maybe Integer
getWeight :: Polynomial -> Int -> Maybe Integer
getWeight (Polynomial p :: Vector Monomial
p) n :: Int
n = Int -> Maybe Integer
look 0
    where
        plen :: Int
plen = Vector Monomial -> Int
forall a. Vector a -> Int
V.length Vector Monomial
p
        look :: Int -> Maybe Integer
look !Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
plen = Maybe Integer
forall a. Maybe a
Nothing
            | Bool
otherwise =
                let (Monomial w :: Int
w v :: Integer
v) = Vector Monomial
p Vector Monomial -> Int -> Monomial
forall a. Vector a -> Int -> a
! Int
i in
                case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
w Int
n of
                    LT -> Maybe Integer
forall a. Maybe a
Nothing
                    EQ -> Integer -> Maybe Integer
forall a. a -> Maybe a
Just Integer
v
                    GT -> Int -> Maybe Integer
look (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1)
        

mergePoly :: (Integer -> Integer -> Integer) -> Polynomial -> Polynomial -> Polynomial
mergePoly :: (Integer -> Integer -> Integer)
-> Polynomial -> Polynomial -> Polynomial
mergePoly f :: Integer -> Integer -> Integer
f (Polynomial p1 :: Vector Monomial
p1) (Polynomial p2 :: Vector Monomial
p2) = [Monomial] -> Polynomial
fromList ([Monomial] -> Polynomial) -> [Monomial] -> Polynomial
forall a b. (a -> b) -> a -> b
$ Int -> Int -> [Monomial]
loop 0 0
    where
        l1 :: Int
l1 = Vector Monomial -> Int
forall a. Vector a -> Int
V.length Vector Monomial
p1
        l2 :: Int
l2 = Vector Monomial -> Int
forall a. Vector a -> Int
V.length Vector Monomial
p2
        loop :: Int -> Int -> [Monomial]
loop !Int
i1 !Int
i2
            | Int
i1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
l1 Bool -> Bool -> Bool
&& Int
i2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
l2 = []
            | Int
i1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
l1             = (Vector Monomial
p2 Vector Monomial -> Int -> Monomial
forall a. Vector a -> Int -> a
! Int
i2) Monomial -> [Monomial] -> [Monomial]
forall a. a -> [a] -> [a]
: Int -> Int -> [Monomial]
loop Int
i1 (Int
i2Int -> Int -> Int
forall a. Num a => a -> a -> a
+1)
            | Int
i2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
l2             = (Vector Monomial
p1 Vector Monomial -> Int -> Monomial
forall a. Vector a -> Int -> a
! Int
i1) Monomial -> [Monomial] -> [Monomial]
forall a. a -> [a] -> [a]
: Int -> Int -> [Monomial]
loop (Int
i1Int -> Int -> Int
forall a. Num a => a -> a -> a
+1) Int
i2
            | Bool
otherwise            =
                let (coef :: Monomial
coef, i1inc :: Int
i1inc, i2inc :: Int
i2inc) = Monomial -> Monomial -> (Monomial, Int, Int)
forall b c.
(Num b, Num c) =>
Monomial -> Monomial -> (Monomial, b, c)
addCoef (Vector Monomial
p1 Vector Monomial -> Int -> Monomial
forall a. Vector a -> Int -> a
! Int
i1) (Vector Monomial
p2 Vector Monomial -> Int -> Monomial
forall a. Vector a -> Int -> a
! Int
i2) in
                Monomial
coef Monomial -> [Monomial] -> [Monomial]
forall a. a -> [a] -> [a]
: Int -> Int -> [Monomial]
loop (Int
i1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
i1inc) (Int
i2Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
i2inc)
        addCoef :: Monomial -> Monomial -> (Monomial, b, c)
addCoef m1 :: Monomial
m1@(Monomial w1 :: Int
w1 v1 :: Integer
v1) (Monomial w2 :: Int
w2 v2 :: Integer
v2) =
            case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
w1 Int
w2 of
                LT -> (Int -> Integer -> Monomial
Monomial Int
w2 (Integer -> Integer -> Integer
f 0 Integer
v2), 0, 1)
                EQ -> (Int -> Integer -> Monomial
Monomial Int
w1 (Integer -> Integer -> Integer
f Integer
v1 Integer
v2), 1, 1)
                GT -> (Monomial
m1, 1, 0)

addPoly :: Polynomial -> Polynomial -> Polynomial
addPoly :: Polynomial -> Polynomial -> Polynomial
addPoly = (Integer -> Integer -> Integer)
-> Polynomial -> Polynomial -> Polynomial
mergePoly Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+)

subPoly :: Polynomial -> Polynomial -> Polynomial
subPoly :: Polynomial -> Polynomial -> Polynomial
subPoly = (Integer -> Integer -> Integer)
-> Polynomial -> Polynomial -> Polynomial
mergePoly (-)

negPoly :: Polynomial -> Polynomial
negPoly :: Polynomial -> Polynomial
negPoly (Polynomial p :: Vector Monomial
p) = Vector Monomial -> Polynomial
Polynomial (Vector Monomial -> Polynomial) -> Vector Monomial -> Polynomial
forall a b. (a -> b) -> a -> b
$ (Monomial -> Monomial) -> Vector Monomial -> Vector Monomial
forall a b. (a -> b) -> Vector a -> Vector b
V.map Monomial -> Monomial
negateMonomial Vector Monomial
p
    where negateMonomial :: Monomial -> Monomial
negateMonomial (Monomial w :: Int
w v :: Integer
v) = Int -> Integer -> Monomial
Monomial Int
w (-Integer
v)

mulPoly :: Polynomial -> Polynomial -> Polynomial
mulPoly :: Polynomial -> Polynomial -> Polynomial
mulPoly p1 :: Polynomial
p1@(Polynomial v1 :: Vector Monomial
v1) p2 :: Polynomial
p2@(Polynomial v2 :: Vector Monomial
v2) =
    [Monomial] -> Polynomial
fromList ([Monomial] -> Polynomial) -> [Monomial] -> Polynomial
forall a b. (a -> b) -> a -> b
$ (Monomial -> Bool) -> [Monomial] -> [Monomial]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(Monomial _ v :: Integer
v) -> Integer
v Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= 0) ([Monomial] -> [Monomial]) -> [Monomial] -> [Monomial]
forall a b. (a -> b) -> a -> b
$ (Int -> Monomial) -> [Int] -> [Monomial]
forall a b. (a -> b) -> [a] -> [b]
map (\i :: Int
i -> Int -> Integer -> Monomial
Monomial Int
i (Int -> Integer
c Int
i)) ([Int] -> [Monomial]) -> [Int] -> [Monomial]
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int]
forall a. [a] -> [a]
reverse [0..(Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
n)]
    where
        (Monomial m :: Int
m _) = Vector Monomial
v1 Vector Monomial -> Int -> Monomial
forall a. Vector a -> Int -> a
! 0
        (Monomial n :: Int
n _) = Vector Monomial
v2 Vector Monomial -> Int -> Monomial
forall a. Vector a -> Int -> a
! 0
        c :: Int -> Integer
c r :: Int
r = (Integer -> Int -> Integer) -> Integer -> [Int] -> Integer
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\acc :: Integer
acc i :: Int
i -> (Int -> Integer
b (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Int -> Integer
a (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ Int
i) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
acc) 0 [0..Int
r]
            where
                a :: Int -> Integer
a = Integer -> (Integer -> Integer) -> Maybe Integer -> Integer
forall b a. b -> (a -> b) -> Maybe a -> b
maybe 0 Integer -> Integer
forall a. a -> a
id (Maybe Integer -> Integer)
-> (Int -> Maybe Integer) -> Int -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Polynomial -> Int -> Maybe Integer
getWeight Polynomial
p1
                b :: Int -> Integer
b = Integer -> (Integer -> Integer) -> Maybe Integer -> Integer
forall b a. b -> (a -> b) -> Maybe a -> b
maybe 0 Integer -> Integer
forall a. a -> a
id (Maybe Integer -> Integer)
-> (Int -> Maybe Integer) -> Int -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Polynomial -> Int -> Maybe Integer
getWeight Polynomial
p2

squarePoly :: Polynomial -> Polynomial
squarePoly :: Polynomial -> Polynomial
squarePoly p :: Polynomial
p = Polynomial
p Polynomial -> Polynomial -> Polynomial
`mulPoly` Polynomial
p

expPoly :: Polynomial -> Integer -> Polynomial
expPoly :: Polynomial -> Integer -> Polynomial
expPoly p :: Polynomial
p e :: Integer
e = Polynomial -> Integer -> Polynomial
forall t. (Eq t, Num t) => Polynomial -> t -> Polynomial
loop Polynomial
p Integer
e
    where
        loop :: Polynomial -> t -> Polynomial
loop t :: Polynomial
t 0 = Polynomial
t
        loop t :: Polynomial
t n :: t
n = Polynomial -> t -> Polynomial
loop (Polynomial -> Polynomial
squarePoly Polynomial
t) (t
nt -> t -> t
forall a. Num a => a -> a -> a
-1)

divPoly :: Polynomial -> Polynomial -> (Polynomial, Polynomial)
divPoly :: Polynomial -> Polynomial -> (Polynomial, Polynomial)
divPoly p1 :: Polynomial
p1 p2 :: Polynomial
p2@(Polynomial pp2 :: Vector Monomial
pp2) = ([Monomial] -> Polynomial)
-> ([Monomial], Polynomial) -> (Polynomial, Polynomial)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first [Monomial] -> Polynomial
fromList (([Monomial], Polynomial) -> (Polynomial, Polynomial))
-> ([Monomial], Polynomial) -> (Polynomial, Polynomial)
forall a b. (a -> b) -> a -> b
$ Polynomial -> ([Monomial], Polynomial)
divLoop Polynomial
p1
    where divLoop :: Polynomial -> ([Monomial], Polynomial)
divLoop d1 :: Polynomial
d1@(Polynomial pp1 :: Vector Monomial
pp1)
            | Vector Monomial -> Bool
forall a. Vector a -> Bool
V.null Vector Monomial
pp1 = ([], Polynomial
d1)
            | Bool
otherwise  =
                let (Monomial w1 :: Int
w1 v1 :: Integer
v1) = Vector Monomial
pp1 Vector Monomial -> Int -> Monomial
forall a. Vector a -> Int -> a
! 0 in
                let (Monomial w2 :: Int
w2 v2 :: Integer
v2) = Vector Monomial
pp2 Vector Monomial -> Int -> Monomial
forall a. Vector a -> Int -> a
! 0 in
                let w :: Int
w = Int
w1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
w2 in
                let (v :: Integer
v,r :: Integer
r) = Integer
v1 Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
`divMod` Integer
v2 in
                if Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= 0 Bool -> Bool -> Bool
&& Integer
r Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== 0
                    then
                        let mono :: Monomial
mono = (Int -> Integer -> Monomial
Monomial Int
w Integer
v) in
                        let remain :: Polynomial
remain = Polynomial
d1 Polynomial -> Polynomial -> Polynomial
`subPoly` (Polynomial
p2 Polynomial -> Polynomial -> Polynomial
`mulPoly` ([Monomial] -> Polynomial
fromList [Monomial
mono])) in
                        let (l :: [Monomial]
l, finalRem :: Polynomial
finalRem) = Polynomial -> ([Monomial], Polynomial)
divLoop Polynomial
remain in
                        (Monomial
mono Monomial -> [Monomial] -> [Monomial]
forall a. a -> [a] -> [a]
: [Monomial]
l, Polynomial
finalRem)
                    else
                        ([], Polynomial
d1)