Indexed update

examples/example0003-linear-algebra.lhs

@@ -98,6 +98,18 @@ Instead of requiring the us to have this knowledge, we can offload the work to t
 
 I've found this distinction between vectors and arrays greatly simplifies the syntax when using linear algebra.
 
+
+Instances of `IxContainers` are also updateable at a specific index with the `(!~)` operator. This is especially useful with the
+`(&)`-Operator which is just flipped function-application:
+
+>   putStrLn $ "vec & 3 !~ 42 = " + (show $ vec & 3 !~ 42)
+
+Similarly we can also modify the entry with a function using `(%~)` and even combine this with the use of function-composition:
+
+>   putStrLn $ "vec & 0 %~ (*5) = " + (show $ vec & 0 %~ (5))
+>   putStrLn $ "vec & 4 %~ (\\x -> x*x) . 3 !~ 42 = " + (show $ vec & (4 %~ (\x -> x*x)) . (3 !~ 42))
+
+
 Linear Algebra
 =======================================
 
@@ -113,6 +125,7 @@ Let's create two vectors and show all the vector operations you might want to pe
 >   putStrLn $ "component mul: " + show (u .*. v)
 
 Because `SVector` is not just a vector space but also a [hilbert space][hilbert-wiki] (i.e. instance of `Hilbert`),
+
 we get the following operations as well:
 
 >   putStrLn ""

src/SubHask/Algebra.hs

@@ -2763,6 +2763,18 @@ class (ValidLogic s, Monoid s, ValidSetElem s{-, ValidSetIndex s-}) => IxContain
 
     imap :: (ValidElem s (Elem s), ValidElem s b) => (Index s -> Elem s -> b) -> s -> SetElem s b
 
+    -- | Updates value at specific index. O(n), but can be everwritten with more efficient versions
+    infixr 4 !~
+    {-# INLINEABLE (!~) #-}
+    (!~) :: (Eq_ (Index s), Logic (Index s)~Bool, ValidElem s (Elem s)) => Index s -> Elem s -> s -> s
+    (!~) i e c = imap (\ix el -> if ix == i then e else el) c
+
+    -- | Similar to "!~" the operation "%~" applies a given function to the value
+    infixr 4 %~
+    {-# INLINEABLE (%~) #-}
+    (%~) :: (Eq_ (Index s), Logic (Index s)~Bool, ValidElem s (Elem s)) => Index s -> (Elem s -> Elem s) -> s -> s
+    (%~) i f c = imap (\ix el -> if ix == i then f el else el) c
+
     toIxList :: s -> [(Index s, Elem s)]
 
     indices :: s -> [Index s]

src/SubHask/Algebra/Vector.hs

@@ -285,6 +285,27 @@ instance (Monoid r, ValidLogic r, Prim r, IsScalar r) => IxContainer (UVector (n
     {-# INLINE (!) #-}
     (!) (UVector_Dynamic arr off _) i = indexByteArray arr (off+i)
 
+    {-# INLINE (!~) #-}
+    (!~) i e (UVector_Dynamic arr off n) =
+                unsafeInlineIO $ do
+                        let b = n*Prim.sizeOf(undefined::r)
+                        marr <- newByteArray b
+                        copyByteArray marr 0 arr off b
+                        writeByteArray marr i e
+                        arr' <- unsafeFreezeByteArray marr
+                        return $ UVector_Dynamic arr' 0 n
+
+    {-# INLINE (%~) #-}
+    (%~) i f (UVector_Dynamic arr off n) =
+                unsafeInlineIO $ do
+                        let b = n*Prim.sizeOf(undefined::r)
+                        marr <- newByteArray b
+                        copyByteArray marr 0 arr off b
+                        e <- readByteArray marr i
+                        writeByteArray marr i (f e)
+                        arr' <- unsafeFreezeByteArray marr
+                        return $ UVector_Dynamic arr' 0 n
+
     {-# INLINABLE toIxList #-}
     toIxList (UVector_Dynamic arr off n) = P.zip [0..] $ go (n-1) []
         where
@@ -772,6 +793,29 @@ instance
     {-# INLINE (!) #-}
     (!) (SVector_Dynamic fp off _) i = unsafeInlineIO $ withForeignPtr fp $ \p -> peekElemOff p (off+i)
 
+    {-# INLINE (!~) #-}
+    (!~) i e (SVector_Dynamic fp1 off n) =
+            unsafeInlineIO $ do
+                let b = n*sizeOf(undefined::r)
+                fp2 <- mallocForeignPtrBytes b
+                withForeignPtr fp1 $ \ptr1 ->
+                    withForeignPtr fp2 $ \ptr2 -> do
+                        copyBytes ptr2 (plusPtr ptr1 off) b
+                        pokeElemOff ptr2 i e
+                return $ (SVector_Dynamic fp2 0 n)
+
+    {-# INLINE (%~) #-}
+    (%~) i f (SVector_Dynamic fp1 off n) =
+            unsafeInlineIO $ do
+                let b = n*sizeOf(undefined::r)
+                fp2 <- mallocForeignPtrBytes b
+                withForeignPtr fp1 $ \ptr1 ->
+                    withForeignPtr fp2 $ \ptr2 -> do
+                        copyBytes ptr2 (plusPtr ptr1 off) b
+                        e <- peekElemOff ptr2 i
+                        pokeElemOff ptr2 i (f e)
+                return $ (SVector_Dynamic fp2 0 n)
+
     {-# INLINABLE toIxList #-}
     toIxList v = P.zip [0..] $ go (dim v-1) []
         where
@@ -1155,6 +1199,32 @@ instance
     {-# INLINE (!) #-}
     (!) (SVector_Nat fp) i = unsafeInlineIO $ withForeignPtr fp $ \p -> peekElemOff p i
 
+    {-# INLINE (!~) #-}
+    (!~) i e (SVector_Nat fp1) =
+            unsafeInlineIO $ do
+                let b = n*sizeOf(undefined::r)
+                    n = nat2int (Proxy::Proxy n)
+                fp2 <- mallocForeignPtrBytes b
+                withForeignPtr fp1 $ \ptr1 ->
+                    withForeignPtr fp2 $ \ptr2 -> do
+                        copyBytes ptr2 ptr1 b
+                        pokeElemOff ptr2 i e
+                return $ (SVector_Nat fp2)
+
+    {-# INLINE (%~) #-}
+    (%~) i f (SVector_Nat fp1) =
+            unsafeInlineIO $ do
+                let b = n*sizeOf(undefined::r)
+                    n = nat2int (Proxy::Proxy n)
+                fp2 <- mallocForeignPtrBytes b
+                withForeignPtr fp1 $ \ptr1 ->
+                    withForeignPtr fp2 $ \ptr2 -> do
+                        copyBytes ptr2 ptr1 b
+                        e <- peekElemOff ptr2 i
+                        pokeElemOff ptr2 i (f e)
+                return $ (SVector_Nat fp2)
+
+
     {-# INLINABLE toIxList #-}
     toIxList v = P.zip [0..] $ go (dim v-1) []
         where

src/SubHask/Category.hs

@@ -56,6 +56,7 @@ module SubHask.Category
     -- * Hask
     , Hask
     , ($)
+    , (&)
     , ($!)
     , embedHask
     , embedHask2
@@ -221,6 +222,22 @@ infixr 0 $
 ($) :: Concrete subcat => subcat a b -> a -> b
 ($) = embedType2
 
+-- | Like in lens "&" is just "flip ($)" for reverse application.
+--
+-- This allows us to take advantage of function-composition when working on a single object, i.e. given
+--
+-- > vector :: Vector 5 Int
+--
+-- we can update the 3rd and 4th entry by
+--
+-- > vector & 3 !~ 23 . 4 !~ 42
+--
+-- without traversing the whole structure as (!~) may have a more performant implementation then "updating by traversing"
+
+infixr 1 &
+(&) :: Concrete subcat => a -> subcat a b -> b
+(&) = flip ($)
+
 -- | A strict version of '$'
 infixr 0 $!
 ($!) :: Concrete subcat => subcat a b -> a -> b