{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Optimization.AD
-- Copyright   :  (c) Masahiro Sakai 2023
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  non-portable
--
-- This module is a wrapper of "Numeric.Optimization" that uses
-- [ad](https://hackage.haskell.org/package/ad)'s automatic differentiation.
--
-----------------------------------------------------------------------------
module Numeric.Optimization.AD
  (
  -- * Main function
    minimize
  , minimizeReverse
  , minimizeSparse

  -- * Problem specification
  , Constraint (..)

  -- * Algorithm selection
  , Method (..)
  , isSupportedMethod
  , Params (..)

  -- * Result
  , Result (..)
  , Statistics (..)
  , OptimizationException (..)

  -- * Utilities and Re-exports
  , Default (..)
  , AD
  , auto
  , Reverse
  , Reifies
  , Tape
  , Sparse
  ) where


import Control.Monad.Primitive
import Data.Default.Class
import Data.Foldable (foldlM, toList)
import Data.Functor.Contravariant
import Data.Reflection (Reifies)
import Data.Traversable
import qualified Data.Vector as V
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM
import Numeric.AD (AD, auto)
import Numeric.AD.Internal.Reverse (Tape)
import Numeric.AD.Mode.Reverse (Reverse)
import qualified Numeric.AD.Mode.Reverse as Reverse
import Numeric.AD.Mode.Sparse (Sparse)
import qualified Numeric.AD.Mode.Sparse as Sparse
import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.Optimization as Opt
import Numeric.Optimization hiding (minimize, IsProblem (..))

-- ------------------------------------------------------------------------

data ProblemReverse f
  = ProblemReverse
      (forall s. Reifies s Tape => f (Reverse s Double) -> Reverse s Double)
      (Maybe (V.Vector (Double, Double)))
      [Constraint]
      Int
      (f Int)

instance Traversable f => Opt.IsProblem (ProblemReverse f) where
  func :: ProblemReverse f -> Vector Double -> Double
func (ProblemReverse forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
f Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
_size f Int
template) Vector Double
x =
    forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s.
 (Reifies s Tape, Typeable s) =>
 f (Reverse s a) -> Reverse s a)
-> f a -> (a, f a)
Reverse.grad' forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
f (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template Vector Double
x)

  bounds :: ProblemReverse f -> Maybe (Vector (Double, Double))
bounds (ProblemReverse forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
_f Maybe (Vector (Double, Double))
bounds [Constraint]
_constraints Int
_size f Int
_template) = Maybe (Vector (Double, Double))
bounds

  constraints :: ProblemReverse f -> [Constraint]
constraints (ProblemReverse forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
_f Maybe (Vector (Double, Double))
_bounds [Constraint]
constraints Int
_size f Int
_template) = [Constraint]
constraints

instance Traversable f => Opt.HasGrad (ProblemReverse f) where
  grad :: ProblemReverse f -> Vector Double -> Vector Double
grad (ProblemReverse forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
func Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
size f Int
template) =
    forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s.
 (Reifies s Tape, Typeable s) =>
 f (Reverse s a) -> Reverse s a)
-> f a -> f a
Reverse.grad forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
func forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template

  grad'M :: forall (m :: * -> *).
PrimMonad m =>
ProblemReverse f
-> Vector Double -> MVector (PrimState m) Double -> m Double
grad'M (ProblemReverse forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
f Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
_size f Int
template) Vector Double
x MVector (PrimState m) Double
gvec = do
    case forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s.
 (Reifies s Tape, Typeable s) =>
 f (Reverse s a) -> Reverse s a)
-> f a -> (a, f a)
Reverse.grad' forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
f (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template Vector Double
x) of
      (Double
y, f Double
g) -> do
        forall (m :: * -> *) (mv :: * -> * -> *) a (f :: * -> *).
(PrimMonad m, MVector mv a, Traversable f) =>
f a -> mv (PrimState m) a -> m ()
writeToMVector f Double
g MVector (PrimState m) Double
gvec
        forall (m :: * -> *) a. Monad m => a -> m a
return Double
y

instance Traversable f => Opt.Optionally (Opt.HasGrad (ProblemReverse f)) where
  optionalDict :: Maybe (Dict (HasGrad (ProblemReverse f)))
optionalDict = forall (c :: Constraint). c => Maybe (Dict c)
hasOptionalDict

instance Opt.Optionally (Opt.HasHessian (ProblemReverse f)) where
  optionalDict :: Maybe (Dict (HasHessian (ProblemReverse f)))
optionalDict = forall a. Maybe a
Nothing

-- | Minimization of scalar function of one or more variables.
--
-- This is a wrapper of 'Opt.minimize' and use "Numeric.AD.Mode.Reverse" to compute gradient.
--
-- It cannot be used with methods that requires hessian (e.g. 'Newton').
--
-- Example:
--
-- > {-# LANGUAGE FlexibleContexts #-}
-- > import Numeric.Optimization.AD
-- > 
-- > main :: IO ()
-- > main = do
-- >   (x, result, stat) <- minimizeReverse LBFGS def rosenbrock Nothing [] [-3,-4]
-- >   print (resultSuccess result)  -- True
-- >   print (resultSolution result)  -- [0.999999999009131,0.9999999981094296]
-- >   print (resultValue result)  -- 1.8129771632403013e-18
-- > 
-- > -- https://en.wikipedia.org/wiki/Rosenbrock_function
-- > rosenbrock :: Floating a => [a] -> a
-- > -- rosenbrock :: Reifies s Tape => [Reverse s Double] -> Reverse s Double
-- > rosenbrock [x,y] = sq (1 - x) + 100 * sq (y - sq x)
-- > 
-- > sq :: Floating a => a -> a
-- > sq x = x ** 2
minimizeReverse
  :: forall f. Traversable f
  => Method  -- ^ Numerical optimization algorithm to use
  -> Params (f Double)  -- ^ Parameters for optimization algorithms. Use 'def' as a default.
  -> (forall s. Reifies s Tape => f (Reverse s Double) -> Reverse s Double)  -- ^ Function to be minimized.
  -> Maybe (f (Double, Double))  -- ^ Bounds
  -> [Constraint]  -- ^ Constraints
  -> f Double -- ^ Initial value
  -> IO (Result (f Double))
minimizeReverse :: forall (f :: * -> *).
Traversable f =>
Method
-> Params (f Double)
-> (forall s.
    Reifies s Tape =>
    f (Reverse s Double) -> Reverse s Double)
-> Maybe (f (Double, Double))
-> [Constraint]
-> f Double
-> IO (Result (f Double))
minimizeReverse Method
method Params (f Double)
params forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
f Maybe (f (Double, Double))
bounds [Constraint]
constraints f Double
x0 = do
  let size :: Int
      template :: f Int
      (Int
size, f Int
template) = forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (\Int
i Double
_ -> Int
i seq :: forall a b. a -> b -> b
`seq` (Int
iforall a. Num a => a -> a -> a
+Int
1, Int
i)) Int
0 f Double
x0

      bounds' :: Maybe (V.Vector (Double, Double))
      bounds' :: Maybe (Vector (Double, Double))
bounds' = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size) Maybe (f (Double, Double))
bounds

      prob :: ProblemReverse f
prob = forall (f :: * -> *).
(forall s.
 Reifies s Tape =>
 f (Reverse s Double) -> Reverse s Double)
-> Maybe (Vector (Double, Double))
-> [Constraint]
-> Int
-> f Int
-> ProblemReverse f
ProblemReverse forall s.
Reifies s Tape =>
f (Reverse s Double) -> Reverse s Double
f Maybe (Vector (Double, Double))
bounds' [Constraint]
constraints Int
size f Int
template
      params' :: Params (Vector Double)
params' = forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template) Params (f Double)
params

  Result (Vector Double)
result <- forall prob.
(IsProblem prob, Optionally (HasGrad prob),
 Optionally (HasHessian prob)) =>
Method
-> Params (Vector Double)
-> prob
-> Vector Double
-> IO (Result (Vector Double))
Opt.minimize Method
method Params (Vector Double)
params' ProblemReverse f
prob (forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size f Double
x0)
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template) Result (Vector Double)
result

-- ------------------------------------------------------------------------

data ProblemSparse f
  = ProblemSparse
      (forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double))
      (Maybe (V.Vector (Double, Double)))
      [Constraint]
      Int
      (f Int)

instance Traversable f => Opt.IsProblem (ProblemSparse f) where
  func :: ProblemSparse f -> Vector Double -> Double
func (ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
f Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
_size f Int
template) Vector Double
x =
    forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s. f (AD s (Sparse a)) -> AD s (Sparse a))
-> f a -> (a, f a)
Sparse.grad' forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
f (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template Vector Double
x)

  bounds :: ProblemSparse f -> Maybe (Vector (Double, Double))
bounds (ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
_f Maybe (Vector (Double, Double))
bounds [Constraint]
_constraints Int
_size f Int
_template) = Maybe (Vector (Double, Double))
bounds

  constraints :: ProblemSparse f -> [Constraint]
constraints (ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
_f Maybe (Vector (Double, Double))
_bounds [Constraint]
constraints Int
_size f Int
_template) = [Constraint]
constraints

instance Traversable f => Opt.HasGrad (ProblemSparse f) where
  grad :: ProblemSparse f -> Vector Double -> Vector Double
grad (ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
func Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
size f Int
template) =
    forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s. f (AD s (Sparse a)) -> AD s (Sparse a)) -> f a -> f a
Sparse.grad forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
func forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template

  grad'M :: forall (m :: * -> *).
PrimMonad m =>
ProblemSparse f
-> Vector Double -> MVector (PrimState m) Double -> m Double
grad'M (ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
f Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
_size f Int
template) Vector Double
x MVector (PrimState m) Double
gvec = do
    case forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s. f (AD s (Sparse a)) -> AD s (Sparse a))
-> f a -> (a, f a)
Sparse.grad' forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
f (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template Vector Double
x) of
      (Double
y, f Double
g) -> do
        forall (m :: * -> *) (mv :: * -> * -> *) a (f :: * -> *).
(PrimMonad m, MVector mv a, Traversable f) =>
f a -> mv (PrimState m) a -> m ()
writeToMVector f Double
g MVector (PrimState m) Double
gvec
        forall (m :: * -> *) a. Monad m => a -> m a
return Double
y

instance Traversable f => Opt.HasHessian (ProblemSparse f) where
  hessian :: ProblemSparse f -> Vector Double -> Matrix Double
hessian (ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
func Maybe (Vector (Double, Double))
_bounds [Constraint]
_constraints Int
size f Int
template) =
    forall {a} {t :: * -> *} {t :: * -> *}.
(Storable a, Foldable t, Foldable t) =>
Int -> t (t a) -> Matrix a
toMatrix Int
size forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s. f (AD s (Sparse a)) -> AD s (Sparse a))
-> f a -> f (f a)
Sparse.hessian forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
func forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template
    where
      toMatrix :: Int -> t (t a) -> Matrix a
toMatrix Int
n t (t a)
xss = (Int
n forall a. Storable a => Int -> Int -> [a] -> Matrix a
LA.>< Int
n) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall (t :: * -> *) a. Foldable t => t a -> [a]
toList forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t (t a)
xss

instance Traversable f => Opt.Optionally (Opt.HasGrad (ProblemSparse f)) where
  optionalDict :: Maybe (Dict (HasGrad (ProblemSparse f)))
optionalDict = forall (c :: Constraint). c => Maybe (Dict c)
hasOptionalDict

instance Traversable f => Opt.Optionally (Opt.HasHessian (ProblemSparse f)) where
  optionalDict :: Maybe (Dict (HasHessian (ProblemSparse f)))
optionalDict = forall (c :: Constraint). c => Maybe (Dict c)
hasOptionalDict

-- | Minimization of scalar function of one or more variables.
--
-- This is a wrapper of 'Opt.minimize' and use "Numeric.AD.Mode.Sparse" to compute gradient
-- and hessian.
--
-- Unlike 'minimizeReverse', it can be used with methods that requires hessian (e.g. 'Newton').
--
-- Example:
--
-- > {-# LANGUAGE FlexibleContexts #-}
-- > import Numeric.Optimization.AD
-- >
-- > main :: IO ()
-- > main = do
-- >   (x, result, stat) <- minimizeSparse Newton def rosenbrock Nothing [] [-3,-4]
-- >   print (resultSuccess result)  -- True
-- >   print (resultSolution result)  -- [0.9999999999999999,0.9999999999999998]
-- >   print (resultValue result)  -- 1.232595164407831e-32
-- >
-- > -- https://en.wikipedia.org/wiki/Rosenbrock_function
-- > rosenbrock :: Floating a => [a] -> a
-- > -- rosenbrock :: [AD s (Sparse Double)] -> AD s (Sparse Double)
-- > rosenbrock [x,y] = sq (1 - x) + 100 * sq (y - sq x)
-- >
-- > sq :: Floating a => a -> a
-- > sq x = x ** 2
minimizeSparse
  :: forall f. Traversable f
  => Method  -- ^ Numerical optimization algorithm to use
  -> Params (f Double)  -- ^ Parameters for optimization algorithms. Use 'def' as a default.
  -> (forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double))  -- ^ Function to be minimized.
  -> Maybe (f (Double, Double))  -- ^ Bounds
  -> [Constraint]  -- ^ Constraints
  -> f Double -- ^ Initial value
  -> IO (Result (f Double))
minimizeSparse :: forall (f :: * -> *).
Traversable f =>
Method
-> Params (f Double)
-> (forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double))
-> Maybe (f (Double, Double))
-> [Constraint]
-> f Double
-> IO (Result (f Double))
minimizeSparse Method
method Params (f Double)
params forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
f Maybe (f (Double, Double))
bounds [Constraint]
constraints f Double
x0 = do
  let size :: Int
      template :: f Int
      (Int
size, f Int
template) = forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (\Int
i Double
_ -> Int
i seq :: forall a b. a -> b -> b
`seq` (Int
iforall a. Num a => a -> a -> a
+Int
1, Int
i)) Int
0 f Double
x0

      bounds' :: Maybe (V.Vector (Double, Double))
      bounds' :: Maybe (Vector (Double, Double))
bounds' = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size) Maybe (f (Double, Double))
bounds

      prob :: ProblemSparse f
prob = forall (f :: * -> *).
(forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double))
-> Maybe (Vector (Double, Double))
-> [Constraint]
-> Int
-> f Int
-> ProblemSparse f
ProblemSparse forall s. f (AD s (Sparse Double)) -> AD s (Sparse Double)
f Maybe (Vector (Double, Double))
bounds' [Constraint]
constraints Int
size f Int
template
      params' :: Params (Vector Double)
params' = forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template) Params (f Double)
params

  Result (Vector Double)
result <- forall prob.
(IsProblem prob, Optionally (HasGrad prob),
 Optionally (HasHessian prob)) =>
Method
-> Params (Vector Double)
-> prob
-> Vector Double
-> IO (Result (Vector Double))
Opt.minimize Method
method Params (Vector Double)
params' ProblemSparse f
prob (forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size f Double
x0)
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template) Result (Vector Double)
result

-- ------------------------------------------------------------------------

-- | Synonym of 'minimizeReverse'
minimize
  :: forall f. Traversable f
  => Method  -- ^ Numerical optimization algorithm to use
  -> Params (f Double)  -- ^ Parameters for optimization algorithms. Use 'def' as a default.
  -> (forall s. Reifies s Tape => f (Reverse s Double) -> Reverse s Double)  -- ^ Function to be minimized.
  -> Maybe (f (Double, Double))  -- ^ Bounds
  -> [Constraint]  -- ^ Constraints
  -> f Double -- ^ Initial value
  -> IO (Result (f Double))
minimize :: forall (f :: * -> *).
Traversable f =>
Method
-> Params (f Double)
-> (forall s.
    Reifies s Tape =>
    f (Reverse s Double) -> Reverse s Double)
-> Maybe (f (Double, Double))
-> [Constraint]
-> f Double
-> IO (Result (f Double))
minimize = forall (f :: * -> *).
Traversable f =>
Method
-> Params (f Double)
-> (forall s.
    Reifies s Tape =>
    f (Reverse s Double) -> Reverse s Double)
-> Maybe (f (Double, Double))
-> [Constraint]
-> f Double
-> IO (Result (f Double))
minimizeReverse

-- ------------------------------------------------------------------------

fromVector :: (Functor f, VG.Vector v a) => f Int -> v a -> f a
fromVector :: forall (f :: * -> *) (v :: * -> *) a.
(Functor f, Vector v a) =>
f Int -> v a -> f a
fromVector f Int
template v a
x = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (v a
x forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.!) f Int
template

toVector :: (Traversable f, VG.Vector v a) => Int -> f a -> v a
toVector :: forall (f :: * -> *) (v :: * -> *) a.
(Traversable f, Vector v a) =>
Int -> f a -> v a
toVector Int
size f a
x = forall (v :: * -> *) a.
Vector v a =>
(forall s. ST s (Mutable v s a)) -> v a
VG.create forall a b. (a -> b) -> a -> b
$ do
  Mutable v s a
vec <- forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
VGM.new Int
size
  forall (m :: * -> *) (mv :: * -> * -> *) a (f :: * -> *).
(PrimMonad m, MVector mv a, Traversable f) =>
f a -> mv (PrimState m) a -> m ()
writeToMVector f a
x Mutable v s a
vec
  forall (m :: * -> *) a. Monad m => a -> m a
return Mutable v s a
vec

writeToMVector :: (PrimMonad m, VGM.MVector mv a, Traversable f) => f a -> mv (PrimState m) a -> m ()
writeToMVector :: forall (m :: * -> *) (mv :: * -> * -> *) a (f :: * -> *).
(PrimMonad m, MVector mv a, Traversable f) =>
f a -> mv (PrimState m) a -> m ()
writeToMVector f a
x mv (PrimState m) a
vec = do
  Int
_ <- forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM (\Int
i a
v -> forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write mv (PrimState m) a
vec Int
i a
v forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return (Int
iforall a. Num a => a -> a -> a
+Int
1)) Int
0 f a
x
  forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- ------------------------------------------------------------------------