-- Copyright 2018-2021 Google LLC
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--      http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

-- | Lenses and related optics for 'Vec's.

-- Work around <https://ghc.haskell.org/trac/ghc/ticket/14511>
{-# OPTIONS_GHC -fno-float-in #-}
{-# OPTIONS_GHC -Wno-orphans #-}

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE UndecidableInstances #-}

module Data.Vec.Short.Lens
    (
    -- ** List-based lenses
      list
    -- ** Arity-based lenses
    , paired

    -- * List-like operators
    -- ** Constructor views
    -- *** The cons view
    , consed
    -- *** The snoc view
    , snoced
    -- ** Operator views
    -- *** The append view
    , chopped
    -- *** The concat view
    , subVecs
    -- *** The reverse view
    , reversed
    -- *** The transposition view
    , vtransposed

    -- * Misc lenses
    , midElem, ixElem, ix
    , sliced
    , rotated
    , vdiagonal
    ) where

import Prelude hiding ((++), concat, concatMap, iterate)

import qualified Data.Foldable as F
import GHC.ST (runST)
import GHC.TypeLits (KnownNat, type (+), type (<=), type (-))
import qualified GHC.TypeLits

import Control.Lens (Iso, Lens', iso, lens, from, swapped)

import Data.Fin.Int.Explicit (Fin, complementFin, finToInt, unsafeFin)
import Data.SInt (sintVal, unSInt)
import Data.Vec.Short.Internal

#if !MIN_VERSION_lens(5,0,0)
import qualified Control.Lens as L
import Data.Foldable.WithIndex (FoldableWithIndex(..))
import Data.Traversable.WithIndex (TraversableWithIndex(..))
#endif

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

-- | A list (of the right length) is isomorphic to a vector.
-- The list-to-vector direction is partial.
list :: (KnownNat n) => Iso (Vec n a) (Vec n b)
                            [a]       [b]
list :: Iso (Vec n a) (Vec n b) [a] [b]
list = (Vec n a -> [a])
-> ([b] -> Vec n b) -> Iso (Vec n a) (Vec n b) [a] [b]
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso Vec n a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList (SInt n -> [b] -> Vec n b
forall (n :: Nat) a. HasCallStack => SInt n -> [a] -> Vec n a
fromList SInt n
forall (n :: Nat). (HasCallStack, KnownNat n) => SInt n
sintVal)

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

-- | Lens on a single element.
ix :: Fin n -> Lens' (Vec n a) a
ix :: Fin n -> Lens' (Vec n a) a
ix Fin n
i = Fin n
i Fin n
-> ((a -> f a) -> Vec n a -> f (Vec n a))
-> (a -> f a)
-> Vec n a
-> f (Vec n a)
`seq` (Vec n a -> a) -> (Vec n a -> a -> Vec n a) -> Lens' (Vec n a) a
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens (Vec n a -> Fin n -> a
forall (n :: Nat) a. Vec n a -> Fin n -> a
! Fin n
i) (Fin n -> Vec n a -> a -> Vec n a
forall (n :: Nat) a. Fin n -> Vec n a -> a -> Vec n a
upd Fin n
i)
{-# INLINE ix #-}

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

-- | An isomorphism with an element added at the beginning of a vector.
consed :: Iso (a, Vec n a)    (b, Vec m b)
              (Vec (n + 1) a) (Vec (m + 1) b)
consed :: p (Vec (n + 1) a) (f (Vec (m + 1) b))
-> p (a, Vec n a) (f (b, Vec m b))
consed = AnIso (Vec (m + 1) b) (Vec (n + 1) a) (b, Vec m b) (a, Vec n a)
-> Iso (a, Vec n a) (b, Vec m b) (Vec (n + 1) a) (Vec (m + 1) b)
forall s t a b. AnIso s t a b -> Iso b a t s
from ((Int -> Int)
-> Iso (Vec (m + 1) b) (Vec (n + 1) a) (b, Vec m b) (a, Vec n a)
forall (n :: Nat) a (m :: Nat) b.
(Int -> Int)
-> Iso (Vec (n + 1) a) (Vec (m + 1) b) (a, Vec n a) (b, Vec m b)
unsafeIxElem (Int -> Int -> Int
forall a b. a -> b -> a
const Int
0))

-- | An isomorphism with an element added at the end of a vector.
snoced :: Iso (Vec n a, a)    (Vec m b, b)
              (Vec (n + 1) a) (Vec (m + 1) b)
snoced :: p (Vec (n + 1) a) (f (Vec (m + 1) b))
-> p (Vec n a, a) (f (Vec m b, b))
snoced = AnIso (Vec (m + 1) b) (Vec (n + 1) a) (Vec m b, b) (Vec n a, a)
-> Iso (Vec n a, a) (Vec m b, b) (Vec (n + 1) a) (Vec (m + 1) b)
forall s t a b. AnIso s t a b -> Iso b a t s
from ((Int -> Int)
-> Iso (Vec (m + 1) b) (Vec (n + 1) a) (b, Vec m b) (a, Vec n a)
forall (n :: Nat) a (m :: Nat) b.
(Int -> Int)
-> Iso (Vec (n + 1) a) (Vec (m + 1) b) (a, Vec n a) (b, Vec m b)
unsafeIxElem Int -> Int
forall a. a -> a
id (Exchange
   (Vec m b, b) (Vec n a, a) (b, Vec m b) (Identity (a, Vec n a))
 -> Exchange
      (Vec m b, b)
      (Vec n a, a)
      (Vec (m + 1) b)
      (Identity (Vec (n + 1) a)))
-> (Exchange
      (Vec m b, b) (Vec n a, a) (Vec m b, b) (Identity (Vec n a, a))
    -> Exchange
         (Vec m b, b) (Vec n a, a) (b, Vec m b) (Identity (a, Vec n a)))
-> AnIso (Vec (m + 1) b) (Vec (n + 1) a) (Vec m b, b) (Vec n a, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exchange
  (Vec m b, b) (Vec n a, a) (Vec m b, b) (Identity (Vec n a, a))
-> Exchange
     (Vec m b, b) (Vec n a, a) (b, Vec m b) (Identity (a, Vec n a))
forall (p :: * -> * -> *) a b c d.
Swap p =>
Iso (p a b) (p c d) (p b a) (p d c)
swapped)

-- | Isomorphism between a vector with and without its middle element.
midElem :: forall m n a b. Iso (Vec (n+1) a) (Vec (m+1) b)
                               (a, Vec n a)  (b, Vec m b)
midElem :: p (a, Vec n a) (f (b, Vec m b))
-> p (Vec (n + 1) a) (f (Vec (m + 1) b))
midElem = (Int -> Int)
-> Iso (Vec (n + 1) a) (Vec (m + 1) b) (a, Vec n a) (b, Vec m b)
forall (n :: Nat) a (m :: Nat) b.
(Int -> Int)
-> Iso (Vec (n + 1) a) (Vec (m + 1) b) (a, Vec n a) (b, Vec m b)
unsafeIxElem (Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2)

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

-- | An isomorphism with a 'split' vector.
chopped :: (KnownNat m) => Iso (Vec (m + n) a)    (Vec (o + p) b)
                               (Vec m a, Vec n a) (Vec o b, Vec p b)
chopped :: Iso
  (Vec (m + n) a)
  (Vec (o + p) b)
  (Vec m a, Vec n a)
  (Vec o b, Vec p b)
chopped = (Vec (m + n) a -> (Vec m a, Vec n a))
-> ((Vec o b, Vec p b) -> Vec (o + p) b)
-> Iso
     (Vec (m + n) a)
     (Vec (o + p) b)
     (Vec m a, Vec n a)
     (Vec o b, Vec p b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso (SInt m -> Vec (m + n) a -> (Vec m a, Vec n a)
forall (m :: Nat) (n :: Nat) a.
SInt m -> Vec (m + n) a -> (Vec m a, Vec n a)
split SInt m
forall (n :: Nat). (HasCallStack, KnownNat n) => SInt n
sintVal) ((Vec o b -> Vec p b -> Vec (o + p) b)
-> (Vec o b, Vec p b) -> Vec (o + p) b
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Vec o b -> Vec p b -> Vec (o + p) b
forall (n :: Nat) a (m :: Nat). Vec n a -> Vec m a -> Vec (n + m) a
(++))

-- | A vector can be split (isomorphically) into a vector of vectors.
subVecs :: (KnownNat m, (n GHC.TypeLits.* m) ~ nm, (p GHC.TypeLits.* o) ~ po)
        => Iso (Vec nm a)        (Vec po b)
               (Vec n (Vec m a)) (Vec p (Vec o b))
subVecs :: Iso (Vec nm a) (Vec po b) (Vec n (Vec m a)) (Vec p (Vec o b))
subVecs = (Vec nm a -> Vec n (Vec m a))
-> (Vec p (Vec o b) -> Vec po b)
-> Iso (Vec nm a) (Vec po b) (Vec n (Vec m a)) (Vec p (Vec o b))
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso (SInt m -> Vec (n * m) a -> Vec n (Vec m a)
forall (m :: Nat) (n :: Nat) a.
SInt m -> Vec (n * m) a -> Vec n (Vec m a)
reshape SInt m
forall (n :: Nat). (HasCallStack, KnownNat n) => SInt n
sintVal) Vec p (Vec o b) -> Vec po b
forall (m :: Nat) (n :: Nat) a. Vec n (Vec m a) -> Vec (n * m) a
concat

-- | A vector is isomorphic to its reversal.
reversed :: Iso (Vec n a) (Vec m b)
                (Vec n a) (Vec m b)
reversed :: p (Vec n a) (f (Vec m b)) -> p (Vec n a) (f (Vec m b))
reversed = (Vec n a -> Vec n a)
-> (Vec m b -> Vec m b)
-> Iso (Vec n a) (Vec m b) (Vec n a) (Vec m b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso Vec n a -> Vec n a
forall (n :: Nat) a. Vec n a -> Vec n a
rev Vec m b -> Vec m b
forall (n :: Nat) a. Vec n a -> Vec n a
rev

--------------------------------
-- Other misc lenses

-- Unsafe version of ixElem, but with a function computing the
-- index from @valueOf @n@.
-- This allows unsafeIxElem to be used easily for midElem and snoced.
unsafeIxElem :: (Int -> Int) -> Iso (Vec (n+1) a) (Vec (m+1) b)
                                    (a, Vec n a)  (b, Vec m b)
unsafeIxElem :: (Int -> Int)
-> Iso (Vec (n + 1) a) (Vec (m + 1) b) (a, Vec n a) (b, Vec m b)
unsafeIxElem Int -> Int
fi = (Vec (n + 1) a -> (a, Vec n a))
-> ((b, Vec m b) -> Vec (m + 1) b)
-> Iso (Vec (n + 1) a) (Vec (m + 1) b) (a, Vec n a) (b, Vec m b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso Vec (n + 1) a -> (a, Vec n a)
getf (b, Vec m b) -> Vec (m + 1) b
setf
  where getf :: Vec (n + 1) a -> (a, Vec n a)
getf Vec (n + 1) a
xs =
          let !i :: Int
i = Int -> Int
fi (Vec (n + 1) a -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize Vec (n + 1) a
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
          in  Vec (n + 1) a -> Int -> (a -> (a, Vec n a)) -> (a, Vec n a)
forall (n :: Nat) a r. Vec n a -> Int -> (a -> r) -> r
unsafeIndexK Vec (n + 1) a
xs Int
i ((a -> (a, Vec n a)) -> (a, Vec n a))
-> (a -> (a, Vec n a)) -> (a, Vec n a)
forall a b. (a -> b) -> a -> b
$ \ a
xi ->
                let !xs' :: Vec n a
xs' = Int -> Vec (n + 1) a -> Vec n a
forall (n :: Nat) a. Int -> Vec (n + 1) a -> Vec n a
unsafeRemove Int
i Vec (n + 1) a
xs
                in  (a
xi, Vec n a
xs')
        setf :: (b, Vec m b) -> Vec (m + 1) b
setf (b
xi, Vec m b
xs) =
          let !i :: Int
i = Int -> Int
fi (Vec m b -> Int
forall (n :: Nat) a. Vec n a -> Int
vSize Vec m b
xs)
          in  Int -> b -> Vec m b -> Vec (m + 1) b
forall a (n :: Nat). Int -> a -> Vec n a -> Vec (n + 1) a
unsafeInsert Int
i b
xi Vec m b
xs
{-# INLINE unsafeIxElem #-}

-- | Isomorphism between a vector with and without a single element at
-- the given index.
ixElem :: forall n a b. Fin (n+1) -> Iso (Vec (n+1) a) (Vec (n+1) b)
                                         (a, Vec n a)  (b, Vec n b)
ixElem :: Fin (n + 1)
-> Iso (Vec (n + 1) a) (Vec (n + 1) b) (a, Vec n a) (b, Vec n b)
ixElem Fin (n + 1)
i = (Int -> Int)
-> Iso (Vec (n + 1) a) (Vec (n + 1) b) (a, Vec n a) (b, Vec n b)
forall (n :: Nat) a (m :: Nat) b.
(Int -> Int)
-> Iso (Vec (n + 1) a) (Vec (m + 1) b) (a, Vec n a) (b, Vec m b)
unsafeIxElem (Int -> Int -> Int
forall a b. a -> b -> a
const (Fin (n + 1) -> Int
forall (n :: Nat). Fin n -> Int
finToInt Fin (n + 1)
i))


-- | A lens to a slice of the vector.
sliced
    :: forall m n a
    .  (KnownNat m, KnownNat n, m <= n)
    => Fin (n - m + 1)
    -> Lens' (Vec n a) (Vec m a)
sliced :: Fin ((n - m) + 1) -> Lens' (Vec n a) (Vec m a)
sliced (Fin ((n - m) + 1) -> Int
forall (n :: Nat). Fin n -> Int
finToInt -> !Int
start) = (Vec n a -> Vec m a)
-> (Vec n a -> Vec m a -> Vec n a) -> Lens' (Vec n a) (Vec m a)
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens Vec n a -> Vec m a
getf Vec n a -> Vec m a -> Vec n a
setf
    where
    !m :: Int
m    = SInt m -> Int
forall (n :: Nat). SInt n -> Int
unSInt @m SInt m
forall (n :: Nat). (HasCallStack, KnownNat n) => SInt n
sintVal
    !n :: Int
n    = SInt n -> Int
forall (n :: Nat). SInt n -> Int
unSInt @n SInt n
forall (n :: Nat). (HasCallStack, KnownNat n) => SInt n
sintVal
    !end :: Int
end  = Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
m
    !rest :: Int
rest = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
end -- the length of the post-slice portion of the vector

    getf :: Vec n a -> Vec m a
getf Vec n a
xs    = Vec n a -> Int -> SInt m -> Vec m a
forall (n :: Nat) a (m :: Nat). Vec n a -> Int -> SInt m -> Vec m a
sliceVec Vec n a
xs Int
start SInt m
forall (n :: Nat). (HasCallStack, KnownNat n) => SInt n
sintVal
    setf :: Vec n a -> Vec m a -> Vec n a
setf Vec n a
xs Vec m a
ys =
        SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
forall (n :: Nat) a.
SInt n -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
createVec SInt n
forall (n :: Nat). (HasCallStack, KnownNat n) => SInt n
sintVal ((forall s. MutableVec s n a -> ST s ()) -> Vec n a)
-> (forall s. MutableVec s n a -> ST s ()) -> Vec n a
forall a b. (a -> b) -> a -> b
$ \MutableVec s n a
mv -> do
            Vec n a -> Int -> MutableVec s n a -> Int -> Int -> ST s ()
forall (n :: Nat) a s (m :: Nat).
Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec Vec n a
xs Int
0   MutableVec s n a
mv Int
0     Int
start
            Vec m a -> Int -> MutableVec s n a -> Int -> Int -> ST s ()
forall (n :: Nat) a s (m :: Nat).
Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec Vec m a
ys Int
0   MutableVec s n a
mv Int
start Int
m
            Vec n a -> Int -> MutableVec s n a -> Int -> Int -> ST s ()
forall (n :: Nat) a s (m :: Nat).
Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s ()
unsafeCopyVec Vec n a
xs Int
end MutableVec s n a
mv Int
end   Int
rest


-- | A two-element vector is isomorphic to a pair.
paired :: Iso (Vec 2 a) (Vec 2 b) (a, a) (b, b)
paired :: p (a, a) (f (b, b)) -> p (Vec 2 a) (f (Vec 2 b))
paired = (Vec 2 a -> (a, a))
-> ((b, b) -> Vec 2 b) -> Iso (Vec 2 a) (Vec 2 b) (a, a) (b, b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso Vec 2 a -> (a, a)
forall (n :: Nat) b. Vec n b -> (b, b)
unvec2 ((b -> b -> Vec 2 b) -> (b, b) -> Vec 2 b
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry b -> b -> Vec 2 b
forall a. a -> a -> Vec 2 a
vec2)
    where
    unvec2 :: Vec n b -> (b, b)
unvec2 Vec n b
v =
      Vec n b -> Fin n -> (b -> (b, b)) -> (b, b)
forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n b
v (Int -> Fin n
forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin @Int Int
0) ((b -> (b, b)) -> (b, b)) -> (b -> (b, b)) -> (b, b)
forall a b. (a -> b) -> a -> b
$ \b
x ->
      Vec n b -> Fin n -> (b -> (b, b)) -> (b, b)
forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n b
v (Int -> Fin n
forall a (n :: Nat). Integral a => a -> Fin n
unsafeFin @Int Int
1) ((b -> (b, b)) -> (b, b)) -> (b -> (b, b)) -> (b, b)
forall a b. (a -> b) -> a -> b
$ \b
y -> (b
x,b
y)

-- | Isomorphism between a vector and a vector rotated @i@ steps.
-- The element at index 0 in the first vector is at index @i@ in the second.
-- E.g., @view (rotated 1) (fromList "ABCD") == fromList "DABC"@
rotated :: forall n a b. Fin n -> Iso (Vec n a) (Vec n b) (Vec n a) (Vec n b)
rotated :: Fin n -> Iso (Vec n a) (Vec n b) (Vec n a) (Vec n b)
rotated Fin n
i = (Vec n a -> Vec n a)
-> (Vec n b -> Vec n b)
-> Iso (Vec n a) (Vec n b) (Vec n a) (Vec n b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso (Fin n -> Vec n a -> Vec n a
forall (n :: Nat) a. Fin n -> Vec n a -> Vec n a
rot Fin n
i) (\Vec n b
v -> Fin n -> Vec n b -> Vec n b
forall (n :: Nat) a. Fin n -> Vec n a -> Vec n a
rot (SInt n -> Fin n -> Fin n
forall (n :: Nat). SInt n -> Fin n -> Fin n
complementFin (Vec n b -> SInt n
forall (n :: Nat) a. Vec n a -> SInt n
svSize Vec n b
v) Fin n
i) Vec n b
v)
{-# INLINE rotated #-}

-- | Isomorphism of transposed vectors.
vtransposed :: (KnownNat m, KnownNat p)
            => Iso (Vec n (Vec m a)) (Vec p (Vec o b))
                   (Vec m (Vec n a)) (Vec o (Vec p b))
vtransposed :: Iso
  (Vec n (Vec m a))
  (Vec p (Vec o b))
  (Vec m (Vec n a))
  (Vec o (Vec p b))
vtransposed = (Vec n (Vec m a) -> Vec m (Vec n a))
-> (Vec o (Vec p b) -> Vec p (Vec o b))
-> Iso
     (Vec n (Vec m a))
     (Vec p (Vec o b))
     (Vec m (Vec n a))
     (Vec o (Vec p b))
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso (SInt m -> Vec n (Vec m a) -> Vec m (Vec n a)
forall (m :: Nat) (n :: Nat) a.
SInt m -> Vec n (Vec m a) -> Vec m (Vec n a)
vtranspose SInt m
forall (n :: Nat). (HasCallStack, KnownNat n) => SInt n
sintVal) (SInt p -> Vec o (Vec p b) -> Vec p (Vec o b)
forall (m :: Nat) (n :: Nat) a.
SInt m -> Vec n (Vec m a) -> Vec m (Vec n a)
vtranspose SInt p
forall (n :: Nat). (HasCallStack, KnownNat n) => SInt n
sintVal)

-- TODO: KnownNat not needed.
-- | Lens on the main diagonal.
vdiagonal :: forall n a. KnownNat n => Lens' (Vec n (Vec n a)) (Vec n a)
vdiagonal :: Lens' (Vec n (Vec n a)) (Vec n a)
vdiagonal = (Vec n (Vec n a) -> Vec n a)
-> (Vec n (Vec n a) -> Vec n a -> Vec n (Vec n a))
-> Lens' (Vec n (Vec n a)) (Vec n a)
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens Vec n (Vec n a) -> Vec n a
getf Vec n (Vec n a) -> Vec n a -> Vec n (Vec n a)
setf
    where
    getf :: Vec n (Vec n a) -> Vec n a
    getf :: Vec n (Vec n a) -> Vec n a
getf = (Fin n -> Vec n a -> a) -> Vec n (Vec n a) -> Vec n a
forall (n :: Nat) a b. (Fin n -> a -> b) -> Vec n a -> Vec n b
imap ((Vec n a -> Fin n -> a) -> Fin n -> Vec n a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip (!))

    setf :: Vec n (Vec n a) -> Vec n a -> Vec n (Vec n a)
    setf :: Vec n (Vec n a) -> Vec n a -> Vec n (Vec n a)
setf Vec n (Vec n a)
m Vec n a
d =
        SInt n -> (Fin n -> Vec n a) -> Vec n (Vec n a)
forall (n :: Nat) a. SInt n -> (Fin n -> a) -> Vec n a
mkVec SInt n
forall (n :: Nat). (HasCallStack, KnownNat n) => SInt n
sintVal ((Fin n -> Vec n a) -> Vec n (Vec n a))
-> (Fin n -> Vec n a) -> Vec n (Vec n a)
forall a b. (a -> b) -> a -> b
$ \Fin n
i ->
            Vec n (Vec n a) -> Fin n -> (Vec n a -> Vec n a) -> Vec n a
forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n (Vec n a)
m Fin n
i ((Vec n a -> Vec n a) -> Vec n a)
-> (Vec n a -> Vec n a) -> Vec n a
forall a b. (a -> b) -> a -> b
$ \Vec n a
mi ->
            Vec n a -> Fin n -> (a -> Vec n a) -> Vec n a
forall (n :: Nat) a r. Vec n a -> Fin n -> (a -> r) -> r
indexK Vec n a
d Fin n
i ((a -> Vec n a) -> Vec n a) -> (a -> Vec n a) -> Vec n a
forall a b. (a -> b) -> a -> b
$ \a
di ->
            (forall s. ST s (Vec n a)) -> Vec n a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vec n a)) -> Vec n a)
-> (forall s. ST s (Vec n a)) -> Vec n a
forall a b. (a -> b) -> a -> b
$ do
                MutableVec s n a
mi' <- Vec n a -> ST s (MutableVec s n a)
forall (n :: Nat) a s. Vec n a -> ST s (MutableVec s n a)
safeThawMV Vec n a
mi
                MutableVec s n a -> Fin n -> a -> ST s ()
forall s (n :: Nat) a. MutableVec s n a -> Fin n -> a -> ST s ()
writeMV MutableVec s n a
mi' Fin n
i a
di
                MutableVec s n a -> ST s (Vec n a)
forall s (n :: Nat) a. MutableVec s n a -> ST s (Vec n a)
unsafeFreezeMV MutableVec s n a
mi'

#if !MIN_VERSION_lens(5,0,0)
instance L.FunctorWithIndex (Fin n) (Vec n) where imap = imap
instance KnownNat n => L.FoldableWithIndex (Fin n) (Vec n) where
  ifoldMap = ifoldMap
instance KnownNat n => L.TraversableWithIndex (Fin n) (Vec n) where
  itraverse = itraverse
#endif