-- |
-- Module:      Data.Poly.Internal.Multi.Core
-- Copyright:   (c) 2019 Andrew Lelechenko
-- Licence:     BSD3
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Sparse polynomials of one variable.
--

{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE UndecidableInstances       #-}

module Data.Poly.Internal.Multi.Core
  ( normalize
  , plusPoly
  , minusPoly
  , convolution
  , scaleInternal
  , derivPoly
  ) where

import Control.Monad
import Control.Monad.ST
import Data.Bits
import Data.Ord
import qualified Data.Vector.Algorithms.Tim as Tim
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as MG
import qualified Data.Vector.Unboxed as U

normalize
  :: (G.Vector v (t, a), Ord t)
  => (a -> Bool)
  -> (a -> a -> a)
  -> v (t, a)
  -> v (t, a)
normalize :: forall (v :: * -> *) t a.
(Vector v (t, a), Ord t) =>
(a -> Bool) -> (a -> a -> a) -> v (t, a) -> v (t, a)
normalize a -> Bool
p a -> a -> a
add v (t, a)
vs
  | forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v (t, a)
vs = v (t, a)
vs
  | Bool
otherwise = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
    Mutable v s (t, a)
ws <- forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
G.thaw v (t, a)
vs
    Int
l' <- forall (v :: * -> *) t a s.
(Vector v (t, a), Ord t) =>
(a -> Bool) -> (a -> a -> a) -> Mutable v s (t, a) -> ST s Int
normalizeM a -> Bool
p a -> a -> a
add Mutable v s (t, a)
ws
    forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze forall a b. (a -> b) -> a -> b
$ forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
MG.unsafeSlice Int
0 Int
l' Mutable v s (t, a)
ws
{-# INLINABLE normalize #-}

normalizeM
  :: (G.Vector v (t, a), Ord t)
  => (a -> Bool)
  -> (a -> a -> a)
  -> G.Mutable v s (t, a)
  -> ST s Int
normalizeM :: forall (v :: * -> *) t a s.
(Vector v (t, a), Ord t) =>
(a -> Bool) -> (a -> a -> a) -> Mutable v s (t, a) -> ST s Int
normalizeM a -> Bool
p a -> a -> a
add Mutable v s (t, a)
ws = do
    let l :: Int
l = forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
MG.length Mutable v s (t, a)
ws
    let go :: Int -> Int -> (t, a) -> ST s Int
go Int
i Int
j acc :: (t, a)
acc@(t
accP, a
accC)
          | Int
j forall a. Ord a => a -> a -> Bool
>= Int
l =
            if a -> Bool
p a
accC
              then do
                forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.write Mutable v s (t, a)
ws Int
i (t, a)
acc
                forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Int
i forall a. Num a => a -> a -> a
+ Int
1
              else forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i
          | Bool
otherwise = do
            v :: (t, a)
v@(t
vp, a
vc) <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s (t, a)
ws Int
j
            if t
vp forall a. Eq a => a -> a -> Bool
== t
accP
              then Int -> Int -> (t, a) -> ST s Int
go Int
i (Int
j forall a. Num a => a -> a -> a
+ Int
1) (t
accP, a
accC a -> a -> a
`add` a
vc)
              else if a -> Bool
p a
accC
                then do
                  forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.write Mutable v s (t, a)
ws Int
i (t, a)
acc
                  Int -> Int -> (t, a) -> ST s Int
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) (Int
j forall a. Num a => a -> a -> a
+ Int
1) (t, a)
v
                else Int -> Int -> (t, a) -> ST s Int
go Int
i (Int
j forall a. Num a => a -> a -> a
+ Int
1) (t, a)
v
    forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m ()
Tim.sortBy (forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing forall a b. (a, b) -> a
fst) Mutable v s (t, a)
ws
    (t, a)
wsHead <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s (t, a)
ws Int
0
    Int -> Int -> (t, a) -> ST s Int
go Int
0 Int
1 (t, a)
wsHead
{-# INLINABLE normalizeM #-}

plusPoly
  :: (G.Vector v (t, a), Ord t)
  => (a -> Bool)
  -> (a -> a -> a)
  -> v (t, a)
  -> v (t, a)
  -> v (t, a)
plusPoly :: forall (v :: * -> *) t a.
(Vector v (t, a), Ord t) =>
(a -> Bool) -> (a -> a -> a) -> v (t, a) -> v (t, a) -> v (t, a)
plusPoly a -> Bool
p a -> a -> a
add = \v (t, a)
xs v (t, a)
ys -> forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  Mutable v s (t, a)
zs <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew (forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (t, a)
xs forall a. Num a => a -> a -> a
+ forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (t, a)
ys)
  Int
lenZs <- forall (v :: * -> *) t a s.
(Vector v (t, a), Ord t) =>
(a -> Bool)
-> (a -> a -> a)
-> v (t, a)
-> v (t, a)
-> Mutable v s (t, a)
-> ST s Int
plusPolyM a -> Bool
p a -> a -> a
add v (t, a)
xs v (t, a)
ys Mutable v s (t, a)
zs
  forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze forall a b. (a -> b) -> a -> b
$ forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
MG.unsafeSlice Int
0 Int
lenZs Mutable v s (t, a)
zs
{-# INLINABLE plusPoly #-}

plusPolyM
  :: (G.Vector v (t, a), Ord t)
  => (a -> Bool)
  -> (a -> a -> a)
  -> v (t, a)
  -> v (t, a)
  -> G.Mutable v s (t, a)
  -> ST s Int
plusPolyM :: forall (v :: * -> *) t a s.
(Vector v (t, a), Ord t) =>
(a -> Bool)
-> (a -> a -> a)
-> v (t, a)
-> v (t, a)
-> Mutable v s (t, a)
-> ST s Int
plusPolyM a -> Bool
p a -> a -> a
add v (t, a)
xs v (t, a)
ys Mutable v s (t, a)
zs = Int -> Int -> Int -> ST s Int
go Int
0 Int
0 Int
0
  where
    lenXs :: Int
lenXs = forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (t, a)
xs
    lenYs :: Int
lenYs = forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (t, a)
ys

    go :: Int -> Int -> Int -> ST s Int
go Int
ix Int
iy Int
iz
      | Int
ix forall a. Eq a => a -> a -> Bool
== Int
lenXs, Int
iy forall a. Eq a => a -> a -> Bool
== Int
lenYs = forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
iz
      | Int
ix forall a. Eq a => a -> a -> Bool
== Int
lenXs = do
        forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> v a -> m ()
G.unsafeCopy
          (forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
MG.unsafeSlice Int
iz (Int
lenYs forall a. Num a => a -> a -> a
- Int
iy) Mutable v s (t, a)
zs)
          (forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
G.unsafeSlice Int
iy (Int
lenYs forall a. Num a => a -> a -> a
- Int
iy) v (t, a)
ys)
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Int
iz forall a. Num a => a -> a -> a
+ Int
lenYs forall a. Num a => a -> a -> a
- Int
iy
      | Int
iy forall a. Eq a => a -> a -> Bool
== Int
lenYs = do
        forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> v a -> m ()
G.unsafeCopy
          (forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
MG.unsafeSlice Int
iz (Int
lenXs forall a. Num a => a -> a -> a
- Int
ix) Mutable v s (t, a)
zs)
          (forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
G.unsafeSlice Int
ix (Int
lenXs forall a. Num a => a -> a -> a
- Int
ix) v (t, a)
xs)
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Int
iz forall a. Num a => a -> a -> a
+ Int
lenXs forall a. Num a => a -> a -> a
- Int
ix
      | (t
xp, a
xc) <- forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v (t, a)
xs Int
ix
      , (t
yp, a
yc) <- forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v (t, a)
ys Int
iy
      = case t
xp forall a. Ord a => a -> a -> Ordering
`compare` t
yp of
        Ordering
LT -> do
          forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s (t, a)
zs Int
iz (t
xp, a
xc)
          Int -> Int -> Int -> ST s Int
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1) Int
iy (Int
iz forall a. Num a => a -> a -> a
+ Int
1)
        Ordering
EQ -> do
          let zc :: a
zc = a
xc a -> a -> a
`add` a
yc
          if a -> Bool
p a
zc then do
            forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s (t, a)
zs Int
iz (t
xp, a
zc)
            Int -> Int -> Int -> ST s Int
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1) (Int
iy forall a. Num a => a -> a -> a
+ Int
1) (Int
iz forall a. Num a => a -> a -> a
+ Int
1)
          else
            Int -> Int -> Int -> ST s Int
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1) (Int
iy forall a. Num a => a -> a -> a
+ Int
1) Int
iz
        Ordering
GT -> do
          forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s (t, a)
zs Int
iz (t
yp, a
yc)
          Int -> Int -> Int -> ST s Int
go Int
ix (Int
iy forall a. Num a => a -> a -> a
+ Int
1) (Int
iz forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE plusPolyM #-}

minusPoly
  :: (G.Vector v (t, a), Ord t)
  => (a -> Bool)
  -> (a -> a)
  -> (a -> a -> a)
  -> v (t, a)
  -> v (t, a)
  -> v (t, a)
minusPoly :: forall (v :: * -> *) t a.
(Vector v (t, a), Ord t) =>
(a -> Bool)
-> (a -> a) -> (a -> a -> a) -> v (t, a) -> v (t, a) -> v (t, a)
minusPoly a -> Bool
p a -> a
neg a -> a -> a
sub = \v (t, a)
xs v (t, a)
ys -> forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  let lenXs :: Int
lenXs = forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (t, a)
xs
      lenYs :: Int
lenYs = forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (t, a)
ys
  Mutable v s (t, a)
zs <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew (Int
lenXs forall a. Num a => a -> a -> a
+ Int
lenYs)
  let go :: Int -> Int -> Int -> ST s Int
go Int
ix Int
iy Int
iz
        | Int
ix forall a. Eq a => a -> a -> Bool
== Int
lenXs, Int
iy forall a. Eq a => a -> a -> Bool
== Int
lenYs = forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
iz
        | Int
ix forall a. Eq a => a -> a -> Bool
== Int
lenXs = do
          forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
iy .. Int
lenYs forall a. Num a => a -> a -> a
- Int
1] forall a b. (a -> b) -> a -> b
$ \Int
i ->
            forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s (t, a)
zs (Int
iz forall a. Num a => a -> a -> a
+ Int
i forall a. Num a => a -> a -> a
- Int
iy)
              (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
neg (forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v (t, a)
ys Int
i))
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Int
iz forall a. Num a => a -> a -> a
+ Int
lenYs forall a. Num a => a -> a -> a
- Int
iy
        | Int
iy forall a. Eq a => a -> a -> Bool
== Int
lenYs = do
          forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> v a -> m ()
G.unsafeCopy
            (forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
MG.unsafeSlice Int
iz (Int
lenXs forall a. Num a => a -> a -> a
- Int
ix) Mutable v s (t, a)
zs)
            (forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
G.unsafeSlice Int
ix (Int
lenXs forall a. Num a => a -> a -> a
- Int
ix) v (t, a)
xs)
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Int
iz forall a. Num a => a -> a -> a
+ Int
lenXs forall a. Num a => a -> a -> a
- Int
ix
        | (t
xp, a
xc) <- forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v (t, a)
xs Int
ix
        , (t
yp, a
yc) <- forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v (t, a)
ys Int
iy
        = case t
xp forall a. Ord a => a -> a -> Ordering
`compare` t
yp of
          Ordering
LT -> do
            forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s (t, a)
zs Int
iz (t
xp, a
xc)
            Int -> Int -> Int -> ST s Int
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1) Int
iy (Int
iz forall a. Num a => a -> a -> a
+ Int
1)
          Ordering
EQ -> do
            let zc :: a
zc = a
xc a -> a -> a
`sub` a
yc
            if a -> Bool
p a
zc then do
              forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s (t, a)
zs Int
iz (t
xp, a
zc)
              Int -> Int -> Int -> ST s Int
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1) (Int
iy forall a. Num a => a -> a -> a
+ Int
1) (Int
iz forall a. Num a => a -> a -> a
+ Int
1)
            else
              Int -> Int -> Int -> ST s Int
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1) (Int
iy forall a. Num a => a -> a -> a
+ Int
1) Int
iz
          Ordering
GT -> do
            forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s (t, a)
zs Int
iz (t
yp, a -> a
neg a
yc)
            Int -> Int -> Int -> ST s Int
go Int
ix (Int
iy forall a. Num a => a -> a -> a
+ Int
1) (Int
iz forall a. Num a => a -> a -> a
+ Int
1)
  Int
lenZs <- Int -> Int -> Int -> ST s Int
go Int
0 Int
0 Int
0
  forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze forall a b. (a -> b) -> a -> b
$ forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
MG.unsafeSlice Int
0 Int
lenZs Mutable v s (t, a)
zs
{-# INLINABLE minusPoly #-}

scaleM
  :: (G.Vector v (t, a), Num t)
  => (a -> Bool)
  -> (a -> a -> a)
  -> v (t, a)
  -> (t, a)
  -> G.Mutable v s (t, a)
  -> ST s Int
scaleM :: forall (v :: * -> *) t a s.
(Vector v (t, a), Num t) =>
(a -> Bool)
-> (a -> a -> a)
-> v (t, a)
-> (t, a)
-> Mutable v s (t, a)
-> ST s Int
scaleM a -> Bool
p a -> a -> a
mul v (t, a)
xs (t
yp, a
yc) Mutable v s (t, a)
zs = Int -> Int -> ST s Int
go Int
0 Int
0
  where
    lenXs :: Int
lenXs = forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (t, a)
xs

    go :: Int -> Int -> ST s Int
go Int
ix Int
iz
      | Int
ix forall a. Eq a => a -> a -> Bool
== Int
lenXs = forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
iz
      | (t
xp, a
xc) <- forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v (t, a)
xs Int
ix
      = do
        let zc :: a
zc = a
xc a -> a -> a
`mul` a
yc
        if a -> Bool
p a
zc then do
          forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s (t, a)
zs Int
iz (t
xp forall a. Num a => a -> a -> a
+ t
yp, a
zc)
          Int -> Int -> ST s Int
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1) (Int
iz forall a. Num a => a -> a -> a
+ Int
1)
        else
          Int -> Int -> ST s Int
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1) Int
iz
{-# INLINABLE scaleM #-}

scaleInternal
  :: (G.Vector v (t, a), Num t)
  => (a -> Bool)
  -> (a -> a -> a)
  -> t
  -> a
  -> v (t, a)
  -> v (t, a)
scaleInternal :: forall (v :: * -> *) t a.
(Vector v (t, a), Num t) =>
(a -> Bool) -> (a -> a -> a) -> t -> a -> v (t, a) -> v (t, a)
scaleInternal a -> Bool
p a -> a -> a
mul t
yp a
yc v (t, a)
xs = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  Mutable v s (t, a)
zs <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew (forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (t, a)
xs)
  Int
len <- forall (v :: * -> *) t a s.
(Vector v (t, a), Num t) =>
(a -> Bool)
-> (a -> a -> a)
-> v (t, a)
-> (t, a)
-> Mutable v s (t, a)
-> ST s Int
scaleM a -> Bool
p (forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> a -> a
mul) v (t, a)
xs (t
yp, a
yc) Mutable v s (t, a)
zs
  forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze forall a b. (a -> b) -> a -> b
$ forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
MG.unsafeSlice Int
0 Int
len Mutable v s (t, a)
zs
{-# INLINABLE scaleInternal #-}

convolution
  :: forall v t a.
     (G.Vector v (t, a), Ord t, Num t)
  => (a -> Bool)
  -> (a -> a -> a)
  -> (a -> a -> a)
  -> v (t, a)
  -> v (t, a)
  -> v (t, a)
convolution :: forall (v :: * -> *) t a.
(Vector v (t, a), Ord t, Num t) =>
(a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> v (t, a)
-> v (t, a)
-> v (t, a)
convolution a -> Bool
p a -> a -> a
add a -> a -> a
mult = \v (t, a)
xs v (t, a)
ys ->
  if forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (t, a)
xs forall a. Ord a => a -> a -> Bool
>= forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (t, a)
ys
  then (a -> a -> a) -> v (t, a) -> v (t, a) -> v (t, a)
go a -> a -> a
mult v (t, a)
xs v (t, a)
ys
  else (a -> a -> a) -> v (t, a) -> v (t, a) -> v (t, a)
go (forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> a -> a
mult) v (t, a)
ys v (t, a)
xs
  where
    go :: (a -> a -> a) -> v (t, a) -> v (t, a) -> v (t, a)
    go :: (a -> a -> a) -> v (t, a) -> v (t, a) -> v (t, a)
go a -> a -> a
mul v (t, a)
long v (t, a)
short = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
      let lenLong :: Int
lenLong   = forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (t, a)
long
          lenShort :: Int
lenShort  = forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (t, a)
short
          lenBuffer :: Int
lenBuffer = Int
lenLong forall a. Num a => a -> a -> a
* Int
lenShort
      MVector s (Int, Int)
slices <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew Int
lenShort
      Mutable v s (t, a)
buffer <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew Int
lenBuffer

      forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. Int
lenShort forall a. Num a => a -> a -> a
- Int
1] forall a b. (a -> b) -> a -> b
$ \Int
iShort -> do
        let (t
pShort, a
cShort) = forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v (t, a)
short Int
iShort
            from :: Int
from = Int
iShort forall a. Num a => a -> a -> a
* Int
lenLong
            bufferSlice :: Mutable v s (t, a)
bufferSlice = forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
MG.unsafeSlice Int
from Int
lenLong Mutable v s (t, a)
buffer
        Int
len <- forall (v :: * -> *) t a s.
(Vector v (t, a), Num t) =>
(a -> Bool)
-> (a -> a -> a)
-> v (t, a)
-> (t, a)
-> Mutable v s (t, a)
-> ST s Int
scaleM a -> Bool
p a -> a -> a
mul v (t, a)
long (t
pShort, a
cShort) Mutable v s (t, a)
bufferSlice
        forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite MVector s (Int, Int)
slices Int
iShort (Int
from, Int
len)

      Vector (Int, Int)
slices' <- forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze MVector s (Int, Int)
slices
      v (t, a)
buffer' <- forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze Mutable v s (t, a)
buffer
      Mutable v s (t, a)
bufferNew <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew Int
lenBuffer
      forall s.
Vector (Int, Int)
-> v (t, a) -> Mutable v s (t, a) -> ST s (v (t, a))
gogo Vector (Int, Int)
slices' v (t, a)
buffer' Mutable v s (t, a)
bufferNew

    gogo
      :: U.Vector (Int, Int)
      -> v (t, a)
      -> G.Mutable v s (t, a)
      -> ST s (v (t, a))
    gogo :: forall s.
Vector (Int, Int)
-> v (t, a) -> Mutable v s (t, a) -> ST s (v (t, a))
gogo Vector (Int, Int)
slices v (t, a)
buffer Mutable v s (t, a)
bufferNew
      | forall (v :: * -> *) a. Vector v a => v a -> Int
G.length Vector (Int, Int)
slices forall a. Eq a => a -> a -> Bool
== Int
0
      = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (v :: * -> *) a. Vector v a => v a
G.empty
      | forall (v :: * -> *) a. Vector v a => v a -> Int
G.length Vector (Int, Int)
slices forall a. Eq a => a -> a -> Bool
== Int
1
      , (Int
from, Int
len) <- forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex Vector (Int, Int)
slices Int
0
      = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
G.unsafeSlice Int
from Int
len v (t, a)
buffer
      | Bool
otherwise = do
        let nSlices :: Int
nSlices = forall (v :: * -> *) a. Vector v a => v a -> Int
G.length Vector (Int, Int)
slices
        MVector s (Int, Int)
slicesNew <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew ((Int
nSlices forall a. Num a => a -> a -> a
+ Int
1) forall a. Bits a => a -> Int -> a
`shiftR` Int
1)
        forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0 .. (Int
nSlices forall a. Num a => a -> a -> a
- Int
2) forall a. Bits a => a -> Int -> a
`shiftR` Int
1] forall a b. (a -> b) -> a -> b
$ \Int
i -> do
          let (Int
from1, Int
len1) = forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex Vector (Int, Int)
slices (Int
2 forall a. Num a => a -> a -> a
* Int
i)
              (Int
from2, Int
len2) = forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex Vector (Int, Int)
slices (Int
2 forall a. Num a => a -> a -> a
* Int
i forall a. Num a => a -> a -> a
+ Int
1)
              slice1 :: v (t, a)
slice1 = forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
G.unsafeSlice Int
from1 Int
len1 v (t, a)
buffer
              slice2 :: v (t, a)
slice2 = forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
G.unsafeSlice Int
from2 Int
len2 v (t, a)
buffer
              slice3 :: Mutable v s (t, a)
slice3 = forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
MG.unsafeSlice Int
from1 (Int
len1 forall a. Num a => a -> a -> a
+ Int
len2) Mutable v s (t, a)
bufferNew
          Int
len3 <- forall (v :: * -> *) t a s.
(Vector v (t, a), Ord t) =>
(a -> Bool)
-> (a -> a -> a)
-> v (t, a)
-> v (t, a)
-> Mutable v s (t, a)
-> ST s Int
plusPolyM a -> Bool
p a -> a -> a
add v (t, a)
slice1 v (t, a)
slice2 Mutable v s (t, a)
slice3
          forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite MVector s (Int, Int)
slicesNew Int
i (Int
from1, Int
len3)

        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. Integral a => a -> Bool
odd Int
nSlices) forall a b. (a -> b) -> a -> b
$ do
          let (Int
from, Int
len) = forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex Vector (Int, Int)
slices (Int
nSlices forall a. Num a => a -> a -> a
- Int
1)
              slice1 :: v (t, a)
slice1 = forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
G.unsafeSlice Int
from Int
len v (t, a)
buffer
              slice3 :: Mutable v s (t, a)
slice3 = forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
MG.unsafeSlice Int
from Int
len Mutable v s (t, a)
bufferNew
          forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> v a -> m ()
G.unsafeCopy Mutable v s (t, a)
slice3 v (t, a)
slice1
          forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite MVector s (Int, Int)
slicesNew (Int
nSlices forall a. Bits a => a -> Int -> a
`shiftR` Int
1) (Int
from, Int
len)

        Vector (Int, Int)
slicesNew' <- forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze MVector s (Int, Int)
slicesNew
        Mutable v s (t, a)
buffer'    <- forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
G.unsafeThaw   v (t, a)
buffer
        v (t, a)
bufferNew' <- forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze Mutable v s (t, a)
bufferNew
        forall s.
Vector (Int, Int)
-> v (t, a) -> Mutable v s (t, a) -> ST s (v (t, a))
gogo Vector (Int, Int)
slicesNew' v (t, a)
bufferNew' Mutable v s (t, a)
buffer'
{-# INLINABLE convolution #-}

derivPoly
  :: (G.Vector v (t, a))
  => (a -> Bool)   -- ^ is coefficient non-zero?
  -> (t -> t)      -- ^ how to modify powers?
  -> (t -> a -> a) -- ^ how to modify coefficient?
  -> v (t, a)
  -> v (t, a)
derivPoly :: forall (v :: * -> *) t a.
Vector v (t, a) =>
(a -> Bool) -> (t -> t) -> (t -> a -> a) -> v (t, a) -> v (t, a)
derivPoly a -> Bool
p t -> t
dec t -> a -> a
mul v (t, a)
xs
  | forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v (t, a)
xs = forall (v :: * -> *) a. Vector v a => v a
G.empty
  | Bool
otherwise = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
    let lenXs :: Int
lenXs = forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (t, a)
xs
    Mutable v s (t, a)
zs <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew Int
lenXs
    let go :: Int -> Int -> ST s Int
go Int
ix Int
iz
          | Int
ix forall a. Eq a => a -> a -> Bool
== Int
lenXs = forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
iz
          | (t
xp, a
xc) <- forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v (t, a)
xs Int
ix
          = do
            let zc :: a
zc = t
xp t -> a -> a
`mul` a
xc
            if a -> Bool
p a
zc then do
              forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s (t, a)
zs Int
iz (t -> t
dec t
xp, a
zc)
              Int -> Int -> ST s Int
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1) (Int
iz forall a. Num a => a -> a -> a
+ Int
1)
            else
              Int -> Int -> ST s Int
go (Int
ix forall a. Num a => a -> a -> a
+ Int
1) Int
iz
    Int
lenZs <- Int -> Int -> ST s Int
go Int
0 Int
0
    forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze forall a b. (a -> b) -> a -> b
$ forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
MG.unsafeSlice Int
0 Int
lenZs Mutable v s (t, a)
zs
{-# INLINABLE derivPoly #-}