Even nicer lifted code

Author: Jaro Reinders

fleet-array.cabal

@@ -7,7 +7,7 @@ author:          Jaro Reinders
 maintainer:      [email protected]
 build-type:      Simple
 category: Data, Data Structures, Array
-synopsis: Fleet arrays are pure, but support fast updates in single-threaded use
+synopsis: Fleet arrays are pure, but support fast updates if used linearly
 tested-with: GHC ==9.10.1 || ==9.8.4 || ==9.6.6
 description:
   Updating a pure array in Haskell usually requires copying the whole array,
@@ -19,6 +19,9 @@ description:
   performance overhead. We hope this package can form the basis of a pure
   and performant array library.
 
+  Fleet arrays can be more than 10x faster than 'IntMap' if used densly and
+  linearly.
+
 common warnings
     ghc-options: -Wall
 

src/Fleet/Array.hs

@@ -1,5 +1,7 @@
 {-# LANGUAGE MagicHash, UnboxedTuples, UnliftedDatatypes #-}
-{-# OPTIONS_GHC -Wno-name-shadowing -ddump-simpl -ddump-to-file -dsuppress-all -dno-suppress-type-signatures -dno-typeable-binds #-}
+{-# OPTIONS_GHC -ddump-simpl -ddump-to-file -dsuppress-all -dno-suppress-type-signatures -dno-typeable-binds #-}
+-- O2 is necessary to get the right call pattern specializations and remove all the lifted abstractions
+{-# OPTIONS_GHC -O2 #-}
 {-# LANGUAGE LambdaCase #-}
 
 {-|
@@ -25,16 +27,26 @@ import GHC.Exts hiding (fromList, toList, Lifted)
 
 import Data.Kind (Type)
 import GHC.IO.Unsafe (unsafeDupablePerformIO)
-import GHC.Base (IO(IO), unIO)
+import GHC.Base (IO(IO))
 
-data Op a = Set Int# a | Swap Int# Int#
+data Op a = Set {-# UNPACK #-} !Int a | Swap {-# UNPACK #-} !Int {-# UNPACK #-} !Int
 
 -- | Fleet arrays.
-data Array a = DA (MutVar# RealWorld (ArrayData a))
-type ArrayData :: Type -> UnliftedType
-data ArrayData a
-  = Current (MutableArray# RealWorld a)
-  | Diff {-# UNPACK #-} !(Op a) (MutVar# RealWorld (ArrayData a))
+data Array a = A {-# UNPACK #-} !(ArrayVar a)
+type ArrayData# :: Type -> UnliftedType
+data ArrayData# a
+  = Current# {-# UNPACK #-} !(MutArray a)
+  | Diff# {-# UNPACK #-} !(Op a) {-# UNPACK #-} !(ArrayVar a)
+
+data ArrayData a = Current (MutArray a) | Diff !(Op a) !(ArrayVar a)
+
+to# :: ArrayData a -> ArrayData# a
+to# (Current x) = Current# x
+to# (Diff op v) = Diff# op v
+
+from# :: ArrayData# a -> ArrayData a
+from# (Current# x) = Current x
+from# (Diff# op v) = Diff op v
 
 instance Show a => Show (Array a) where
   show xs = "fromList " ++ show (toList xs)
@@ -43,150 +55,162 @@ instance Show a => Show (Array a) where
 aseq :: a -> b -> b
 aseq x y = x `seq` lazy y
 
-type Lifted :: UnliftedType -> Type
-data Lifted a = Lifted a
+-- ArrayVar
+data ArrayVar a = AV (MutVar# RealWorld (ArrayData# a))
+newArrayVar :: ArrayData a -> IO (ArrayVar a)
+newArrayVar x = IO $ \s ->
+  case newMutVar# (to# x) s of
+    (# s', v #) -> (# s', AV v #)
 
-{-# INLINE newMutVarIO #-}
-newMutVarIO :: forall (a :: UnliftedType). a -> IO (Lifted (MutVar# RealWorld a))
-newMutVarIO x = IO $ \s ->
-  case newMutVar# x s of
-    (# s', v #) -> (# s', Lifted v #)
+readArrayVar :: ArrayVar a -> IO (ArrayData a)
+readArrayVar (AV v) = IO $ \s -> case readMutVar# v s of (# s', x #) -> (# s', from# x #)
 
-{-# INLINE readMutVarIO #-}
-readMutVarIO :: forall (a :: UnliftedType) b. MutVar# RealWorld a -> (a -> IO b) -> IO b
-readMutVarIO v f = IO (\s -> case readMutVar# v s of (# s', x #) -> unIO (f x) s')
+writeArrayVar :: ArrayVar a -> ArrayData a -> IO ()
+writeArrayVar (AV v) x = IO $ \s -> (# writeMutVar# v (to# x) s, () #)
 
-{-# INLINE writeMutVarIO #-}
-writeMutVarIO :: forall (a :: UnliftedType). MutVar# RealWorld a -> a -> IO ()
-writeMutVarIO v x = IO (\s -> (# writeMutVar# v x s, () #))
+-- MutArray
 
-readArrayIO :: MutableArray# RealWorld a -> Int# -> IO a
-readArrayIO arr i = IO (readArray# arr i)
+data MutArray a = MA (MutableArray# RealWorld a)
 
-writeArrayIO :: MutableArray# RealWorld a -> Int# -> a -> IO ()
-writeArrayIO arr i x = IO (\s -> (# writeArray# arr i x s, () #))
-
-newArrayIO :: Int# -> a -> IO (Lifted (MutableArray# RealWorld a))
-newArrayIO n x = IO $ \s ->
+newMutArray :: Int -> a -> IO (MutArray a)
+newMutArray (I# n) x = IO $ \s ->
   case newArray# n x s of
-    (# s', arr #) -> (# s', Lifted arr #)
+    (# s', arr #) -> (# s', MA arr #)
+
+readMutArray :: MutArray a -> Int -> IO a
+readMutArray (MA arr) (I# i) = IO (readArray# arr i)
+
+writeMutArray :: MutArray a -> Int -> a -> IO ()
+writeMutArray (MA arr) (I# i) x = IO (\s -> (# writeArray# arr i x s, () #))
 
 -- | Convert a list into an array. O(n)
 fromList :: [a] -> Array a
 fromList xs = unsafeDupablePerformIO $ do
-  let !(I# n) = length xs
-  Lifted arr <- newArrayIO n undefined
+  arr0 <- newMutArray (length xs) undefined
   let go _ _ [] = pure ()
-      go arr i (x:xs') = writeArrayIO arr i x >> go arr (i +# 1#) xs'
-  go arr 0# xs
-  Lifted var <- newMutVarIO (Current arr)
-  pure (DA var)
+      go arr i (x:xs') = writeMutArray arr i x *> go arr (i + 1) xs'
+  go arr0 0 xs
+  v <- newArrayVar (Current arr0)
+  pure (A v)
 
-cloneMutableArrayIO :: MutableArray# RealWorld a -> Int# -> Int# -> IO (Lifted (MutableArray# RealWorld a))
-cloneMutableArrayIO arr off len = IO $ \s ->
+cloneMutArray :: MutArray a -> Int -> Int -> IO (MutArray a)
+cloneMutArray (MA arr) (I# off) (I# len) = IO $ \s ->
   case cloneMutableArray# arr off len s of
-    (# s', arr' #) -> (# s', Lifted arr' #)
+    (# s', arr' #) -> (# s', MA arr' #)
+
+sizeofMutArray :: MutArray a -> Int
+sizeofMutArray (MA x) = I# (sizeofMutableArray# x)
 
-copyInternalIO :: MutVar# RealWorld (ArrayData a) -> IO (Lifted (MutableArray# RealWorld a))
-copyInternalIO v =
-  readMutVarIO v $ \case
-    Current arr -> cloneMutableArrayIO arr 0# (sizeofMutableArray# arr)
+copyInternal :: ArrayVar a -> IO (MutArray a)
+copyInternal v = do
+  av <- readArrayVar v
+  case av of
+    Current arr -> cloneMutArray arr 0 (sizeofMutArray arr)
     Diff op v' -> do
-      Lifted clone <- copyInternalIO v'
-      appOpIO clone op
-      pure (Lifted clone)
+      clone <- copyInternal v'
+      appOp clone op
+      pure clone
 
 -- | Converting an array into a list. O(n)
 toList :: Array a -> [a]
-toList (DA v) = unsafeDupablePerformIO $ do
-  Lifted arr <- copyInternalIO v
-  let n = sizeofMutableArray# arr
+toList (A v) = unsafeDupablePerformIO $ do
+  arr <- copyInternal v
+  let n = sizeofMutArray arr
       go i
-        | isTrue# (i >=# n) = pure []
+        | i >= n = pure []
         | otherwise = do
-            x <- readArrayIO arr i
-            xs <- go (i +# 1#)
+            x <- readMutArray arr i
+            xs <- go (i + 1)
             pure (x : xs)
-  go 0#
+  go 0
 
 -- | Indexing an array. O(1)
 {-# INLINE (!) #-}
 (!) :: Array a -> Int -> a
-DA v ! I# i = unsafeDupablePerformIO (helper v i) where
-    helper v i = readMutVarIO v $ \case
-      Current arr -> readArrayIO arr i
-      Diff (Set j x) xs
-        | isTrue# (i ==# j) -> pure x
-        | otherwise -> helper xs i
-      Diff (Swap j1 j2) xs
-        | isTrue# (i ==# j1) -> helper xs j2
-        | isTrue# (i ==# j2) -> helper xs j1
-        | otherwise -> helper xs i
+A v0 ! i0 = unsafeDupablePerformIO (go v0 i0) where
+  go v i = do
+    dat <- readArrayVar v
+    case dat of
+      Current arr -> readMutArray arr i
+      Diff (Set j x) v'
+        | i == j -> pure x
+        | otherwise -> go v' i
+      Diff (Swap j1 j2) v'
+        | i == j1 -> go v' j2
+        | i == j2 -> go v' j1
+        | otherwise -> go v' i
 
 -- | Indexing an array. O(1)
 -- Using the 'Solo' constructor, you can sequence indexing to happen before
 -- future updates without having to evaluate the element itself.
 {-# INLINE index #-}
 index :: Int -> Array a -> Solo a
-index (I# i) (DA v) = unsafeDupablePerformIO (helper v i) where
-  helper v i = readMutVarIO v $ \case
-      Current arr -> MkSolo <$> readArrayIO arr i
+index i0 (A v0) = unsafeDupablePerformIO (go v0 i0) where
+  go v i = do
+    dat <- readArrayVar v
+    case dat of
+      Current arr -> MkSolo <$> readMutArray arr i
       Diff (Set j x) xs
-        | isTrue# (i ==# j) -> pure (MkSolo x)
-        | otherwise -> helper xs i
+        | i == j -> pure (MkSolo x)
+        | otherwise -> go xs i
       Diff (Swap j1 j2) xs
-        | isTrue# (i ==# j1) -> helper xs j2
-        | isTrue# (i ==# j2) -> helper xs j1
-        | otherwise -> helper xs i
-
-{-# INLINE invertIO #-}
-invertIO :: MutableArray# RealWorld a -> Op a -> IO (Op a)
-invertIO _ (Swap i j) = pure (Swap i j)
-invertIO arr (Set i _) = do
-  y <- readArrayIO arr i
+        | i == j1 -> go xs j2
+        | i == j2 -> go xs j1
+        | otherwise -> go xs i
+
+{-# INLINE invert #-}
+invert :: MutArray a -> Op a -> IO (Op a)
+invert _ (Swap i j) = pure (Swap i j)
+invert arr (Set i _) = do
+  y <- readMutArray arr i
   pure (Set i y)
 
-{-# INLINE appOpIO #-}
-appOpIO :: MutableArray# RealWorld a -> Op a -> IO ()
-appOpIO arr (Set i x) = writeArrayIO arr i x
-appOpIO arr (Swap i j) = do
-  x <- readArrayIO arr i
-  y <- readArrayIO arr j
-  writeArrayIO arr i y
-  writeArrayIO arr j x
+{-# INLINE appOp #-}
+appOp :: MutArray a -> Op a -> IO ()
+appOp arr (Set i x) = writeMutArray arr i x
+appOp arr (Swap i j) = do
+  x <- readMutArray arr i
+  y <- readMutArray arr j
+  writeMutArray arr i y
+  writeMutArray arr j x
 
 {-# INLINE appDiffOp #-}
 appDiffOp :: Op a -> Array a -> Array a
-appDiffOp op (DA v) = unsafeDupablePerformIO $
-  readMutVarIO v $ \case
+appDiffOp op (A v) = unsafeDupablePerformIO $ do
+  dat <- readArrayVar v
+  case dat of
     xs@(Current arr) -> do
-      op' <- invertIO arr op
-      appOpIO arr op
-      Lifted v' <- newMutVarIO xs
-      writeMutVarIO v (Diff op' v')
-      pure (DA v')
+      op' <- invert arr op
+      appOp arr op
+      v' <- newArrayVar xs
+      writeArrayVar v (Diff op' v')
+      pure (A v')
     Diff op' v' -> do
-      Lifted arr <- copyInternalIO v'
-      appOpIO arr op'
-      appOpIO arr op
-      Lifted v'' <- newMutVarIO (Current arr)
-      pure (DA v'')
+      -- TODO: pointer inversion instead of copy
+      -- first invert all pointers until Current
+      -- then apply all updates until back at v
+      -- then do the same as above
+      arr <- copyInternal v'
+      appOp arr op'
+      appOp arr op
+      v'' <- newArrayVar (Current arr)
+      pure (A v'')
 
 -- | Update the array element at a given position to a new value. O(1)
 {-# INLINE set #-}
 set :: Int -> a -> Array a -> Array a
-set (I# i) x = appDiffOp (Set i x)
+set i x = appDiffOp (Set i x)
 
 -- | Swap two elements in an array. O(1)
 {-# INLINE swap #-}
 swap :: Int -> Int -> Array a -> Array a
-swap (I# i) (I# j) = appDiffOp (Swap i j)
+swap i j = appDiffOp (Swap i j)
 
 -- | Copy an array. O(n)
 -- This detaches any future updates from old versions of the array.
 -- Use this when you know you will be updating a large part of an array.
 copy :: Array a -> Array a
-copy (DA v) = unsafeDupablePerformIO $ do
-  Lifted arr <- copyInternalIO v
-  Lifted var <- newMutVarIO (Current arr)
-  pure (DA var)
+copy (A v) = unsafeDupablePerformIO $ do
+  arr <- copyInternal v
+  var <- newArrayVar (Current arr)
+  pure (A var)