{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
module Data.Poly.Internal.Convert
( denseToSparse
, denseToSparse'
, sparseToDense
, sparseToDense'
) where
import Control.Monad.ST
import Data.Semiring (Semiring(..))
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as MG
import qualified Data.Vector.Unboxed.Sized as SU
import qualified Data.Poly.Internal.Dense as Dense
import qualified Data.Poly.Internal.Multi as Sparse
denseToSparse
:: (Eq a, Num a, G.Vector v a, G.Vector v (SU.Vector 1 Word, a))
=> Dense.Poly v a
-> Sparse.Poly v a
denseToSparse :: forall a (v :: * -> *).
(Eq a, Num a, Vector v a, Vector v (Vector 1 Word, a)) =>
Poly v a -> Poly v a
denseToSparse = forall a (v :: * -> *).
(Eq a, Vector v a, Vector v (Vector 1 Word, a)) =>
a -> Poly v a -> Poly v a
denseToSparseInternal a
0
denseToSparse'
:: (Eq a, Semiring a, G.Vector v a, G.Vector v (SU.Vector 1 Word, a))
=> Dense.Poly v a
-> Sparse.Poly v a
denseToSparse' :: forall a (v :: * -> *).
(Eq a, Semiring a, Vector v a, Vector v (Vector 1 Word, a)) =>
Poly v a -> Poly v a
denseToSparse' = forall a (v :: * -> *).
(Eq a, Vector v a, Vector v (Vector 1 Word, a)) =>
a -> Poly v a -> Poly v a
denseToSparseInternal forall a. Semiring a => a
zero
denseToSparseInternal
:: (Eq a, G.Vector v a, G.Vector v (SU.Vector 1 Word, a))
=> a
-> Dense.Poly v a
-> Sparse.Poly v a
denseToSparseInternal :: forall a (v :: * -> *).
(Eq a, Vector v a, Vector v (Vector 1 Word, a)) =>
a -> Poly v a -> Poly v a
denseToSparseInternal a
z = forall (v :: * -> *) (n :: Nat) a.
v (Vector n Word, a) -> MultiPoly v n a
Sparse.MultiPoly forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(Int -> a -> Maybe b) -> v a -> v b
G.imapMaybe (\Int
i a
c -> if a
c forall a. Eq a => a -> a -> Bool
== a
z then forall a. Maybe a
Nothing else forall a. a -> Maybe a
Just (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i, a
c)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a. Poly v a -> v a
Dense.unPoly
sparseToDense
:: (Num a, G.Vector v a, G.Vector v (SU.Vector 1 Word, a))
=> Sparse.Poly v a
-> Dense.Poly v a
sparseToDense :: forall a (v :: * -> *).
(Num a, Vector v a, Vector v (Vector 1 Word, a)) =>
Poly v a -> Poly v a
sparseToDense = forall (v :: * -> *) a.
(Vector v a, Vector v (Vector 1 Word, a)) =>
a -> Poly v a -> Poly v a
sparseToDenseInternal a
0
sparseToDense'
:: (Semiring a, G.Vector v a, G.Vector v (SU.Vector 1 Word, a))
=> Sparse.Poly v a
-> Dense.Poly v a
sparseToDense' :: forall a (v :: * -> *).
(Semiring a, Vector v a, Vector v (Vector 1 Word, a)) =>
Poly v a -> Poly v a
sparseToDense' = forall (v :: * -> *) a.
(Vector v a, Vector v (Vector 1 Word, a)) =>
a -> Poly v a -> Poly v a
sparseToDenseInternal forall a. Semiring a => a
zero
sparseToDenseInternal
:: (G.Vector v a, G.Vector v (SU.Vector 1 Word, a))
=> a
-> Sparse.Poly v a
-> Dense.Poly v a
sparseToDenseInternal :: forall (v :: * -> *) a.
(Vector v a, Vector v (Vector 1 Word, a)) =>
a -> Poly v a -> Poly v a
sparseToDenseInternal a
z (Sparse.MultiPoly v (Vector 1 Word, a)
xs)
| forall (v :: * -> *) a. Vector v a => v a -> Bool
G.null v (Vector 1 Word, a)
xs = forall (v :: * -> *) a. v a -> Poly v a
Dense.Poly 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 len :: Int
len = forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall (n :: Nat) a. Unbox a => Vector (1 + n) a -> a
SU.head (forall a b. (a, b) -> a
fst (forall (v :: * -> *) a. Vector v a => v a -> a
G.unsafeLast v (Vector 1 Word, a)
xs)) forall a. Num a => a -> a -> a
+ Word
1)
Mutable v (PrimState (ST s)) a
ys <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew Int
len
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> a -> m ()
MG.set Mutable v (PrimState (ST s)) a
ys a
z
let go :: Int -> Int -> ST s ()
go Int
xi Int
yi
| Int
xi forall a. Ord a => a -> a -> Bool
>= forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v (Vector 1 Word, a)
xs = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
| (Vector 1 Word
yi', a
c) <- forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v (Vector 1 Word, a)
xs Int
xi
, Int
yi forall a. Eq a => a -> a -> Bool
== forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall (n :: Nat) a. Unbox a => Vector (1 + n) a -> a
SU.head Vector 1 Word
yi')
= forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v (PrimState (ST s)) a
ys Int
yi a
c forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> ST s ()
go (Int
xi forall a. Num a => a -> a -> a
+ Int
1) (Int
yi forall a. Num a => a -> a -> a
+ Int
1)
| Bool
otherwise = Int -> Int -> ST s ()
go Int
xi (Int
yi forall a. Num a => a -> a -> a
+ Int
1)
Int -> Int -> ST s ()
go Int
0 Int
0
forall (v :: * -> *) a. v a -> Poly v a
Dense.Poly forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze Mutable v (PrimState (ST s)) a
ys