{-# 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)
-> (t -> t)
-> (t -> a -> a)
-> 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 #-}