Skip to content

Commit c341dd1

Browse files
noughtmarenoughtmare
authored
Merge pull request #11 from LukaHorvat/io
Wrap RealWorld state passing
2 parents c7eb86c + d7f4482 commit c341dd1

File tree

1 file changed

+109
-105
lines changed

1 file changed

+109
-105
lines changed

‎src/Fleet/Array.hs‎

Lines changed: 109 additions & 105 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
{-# LANGUAGE MagicHash, UnboxedTuples, UnliftedDatatypes #-}
22
{-# OPTIONS_GHC -Wno-name-shadowing -ddump-simpl -ddump-to-file -dsuppress-all -dno-suppress-type-signatures -dno-typeable-binds #-}
3+
{-# LANGUAGE LambdaCase #-}
34

45
{-|
56
Module : Fleet.Array
@@ -19,10 +20,12 @@ latest version.
1920
-}
2021
module Fleet.Array (Array, fromList, toList, (!), index, set, copy, swap, aseq) where
2122

22-
import GHC.Exts hiding (fromList, toList)
2323
import Data.Tuple (Solo (MkSolo))
24+
import GHC.Exts hiding (fromList, toList, Lifted)
2425

2526
import Data.Kind (Type)
27+
import GHC.IO.Unsafe (unsafeDupablePerformIO)
28+
import GHC.Base (IO(IO), unIO)
2629

2730
data Op a = Set Int# a | Swap Int# Int#
2831

@@ -40,42 +43,82 @@ instance Show a => Show (Array a) where
4043
aseq :: a -> b -> b
4144
aseq x y = x `seq` lazy y
4245

46+
type Lifted :: UnliftedType -> Type
47+
data Lifted a = Lifted a
48+
49+
{-# INLINE newMutVarIO #-}
50+
newMutVarIO :: forall (a :: UnliftedType). a -> IO (Lifted (MutVar# RealWorld a))
51+
newMutVarIO x = IO $ \s ->
52+
case newMutVar# x s of
53+
(# s', v #) -> (# s', Lifted v #)
54+
55+
{-# INLINE readMutVarIO #-}
56+
readMutVarIO :: forall (a :: UnliftedType) b. MutVar# RealWorld a -> (a -> IO b) -> IO b
57+
readMutVarIO v f = IO (\s -> case readMutVar# v s of (# s', x #) -> unIO (f x) s')
58+
59+
{-# INLINE writeMutVarIO #-}
60+
writeMutVarIO :: forall (a :: UnliftedType). MutVar# RealWorld a -> a -> IO ()
61+
writeMutVarIO v x = IO (\s -> (# writeMutVar# v x s, () #))
62+
63+
readArrayIO :: MutableArray# RealWorld a -> Int# -> IO a
64+
readArrayIO arr i = IO (readArray# arr i)
65+
66+
writeArrayIO :: MutableArray# RealWorld a -> Int# -> a -> IO ()
67+
writeArrayIO arr i x = IO (\s -> (# writeArray# arr i x s, () #))
68+
69+
newArrayIO :: Int# -> a -> IO (Lifted (MutableArray# RealWorld a))
70+
newArrayIO n x = IO $ \s ->
71+
case newArray# n x s of
72+
(# s', arr #) -> (# s', Lifted arr #)
73+
4374
-- | Convert a list into an array. O(n)
4475
fromList :: [a] -> Array a
45-
fromList xs = DA (runRW# $ \s ->
46-
case newArray# (case length xs of (I# n) -> n) undefined s of { (# s , arr #) ->
47-
case newMutVar# (Current arr) (go arr 0# xs s) of { (# _ , x #) -> x
48-
}}) where
49-
go _ _ [] s = s
50-
go arr i (x:xs) s = go arr (i +# 1#) xs (writeArray# arr i x s)
76+
fromList xs = unsafeDupablePerformIO $ do
77+
let !(I# n) = length xs
78+
Lifted arr <- newArrayIO n undefined
79+
let go _ _ [] = pure ()
80+
go arr i (x:xs') = writeArrayIO arr i x >> go arr (i +# 1#) xs'
81+
go arr 0# xs
82+
Lifted var <- newMutVarIO (Current arr)
83+
pure (DA var)
84+
85+
cloneMutableArrayIO :: MutableArray# RealWorld a -> Int# -> Int# -> IO (Lifted (MutableArray# RealWorld a))
86+
cloneMutableArrayIO arr off len = IO $ \s ->
87+
case cloneMutableArray# arr off len s of
88+
(# s', arr' #) -> (# s', Lifted arr' #)
89+
90+
copyInternalIO :: MutVar# RealWorld (ArrayData a) -> IO (Lifted (MutableArray# RealWorld a))
91+
copyInternalIO v =
92+
readMutVarIO v $ \case
93+
Current arr -> cloneMutableArrayIO arr 0# (sizeofMutableArray# arr)
94+
Diff op v' -> do
95+
Lifted clone <- copyInternalIO v'
96+
appOpIO clone op
97+
pure (Lifted clone)
5198

5299
-- | Converting an array into a list. O(n)
53100
toList :: Array a -> [a]
54-
toList (DA v) = runRW# $ \s ->
55-
case copyInternal v s of { (# s, arr #) ->
56-
let
57-
n = sizeofMutableArray# arr
58-
go i s
59-
| isTrue# (i >=# n) = []
60-
| otherwise =
61-
case readArray# arr i s of
62-
{ (# s, x #) -> x : go (i +# 1#) s
63-
}
64-
in go 0# s
65-
}
101+
toList (DA v) = unsafeDupablePerformIO $ do
102+
Lifted arr <- copyInternalIO v
103+
let n = sizeofMutableArray# arr
104+
go i
105+
| isTrue# (i >=# n) = pure []
106+
| otherwise = do
107+
x <- readArrayIO arr i
108+
xs <- go (i +# 1#)
109+
pure (x : xs)
110+
go 0#
66111

67112
-- | Indexing an array. O(1)
68113
{-# INLINE (!) #-}
69114
(!) :: Array a -> Int -> a
70-
DA v ! I# i = helper v i where
71-
helper v i = runRW# $ \s ->
72-
case readMutVar# v s of
73-
(# s , Current arr #) ->
74-
case readArray# arr i s of (# _ , x #) -> x
75-
(# _ , Diff (Set j x) xs #)
76-
| isTrue# (i ==# j) -> x
115+
DA v ! I# i = unsafeDupablePerformIO (helper v i) where
116+
helper v i = readMutVarIO v $ \case
117+
Current arr -> readArrayIO arr i
118+
Diff (Set j x) xs
119+
| isTrue# (i ==# j) -> pure x
77120
| otherwise -> helper xs i
78-
(# _ , Diff (Swap j1 j2) xs #)
121+
Diff (Swap j1 j2) xs
79122
| isTrue# (i ==# j1) -> helper xs j2
80123
| isTrue# (i ==# j2) -> helper xs j1
81124
| otherwise -> helper xs i
@@ -85,78 +128,49 @@ DA v ! I# i = helper v i where
85128
-- future updates without having to evaluate the element itself.
86129
{-# INLINE index #-}
87130
index :: Int -> Array a -> Solo a
88-
index (I# i) (DA v) = helper v i where
89-
helper v i = runRW# $ \s ->
90-
case readMutVar# v s of
91-
(# s , Current arr #) ->
92-
case readArray# arr i s of (# _ , x #) -> MkSolo x
93-
(# _ , Diff (Set j x) xs #)
94-
| isTrue# (i ==# j) -> MkSolo x
131+
index (I# i) (DA v) = unsafeDupablePerformIO (helper v i) where
132+
helper v i = readMutVarIO v $ \case
133+
Current arr -> MkSolo <$> readArrayIO arr i
134+
Diff (Set j x) xs
135+
| isTrue# (i ==# j) -> pure (MkSolo x)
95136
| otherwise -> helper xs i
96-
(# _ , Diff (Swap j1 j2) xs #)
137+
Diff (Swap j1 j2) xs
97138
| isTrue# (i ==# j1) -> helper xs j2
98139
| isTrue# (i ==# j2) -> helper xs j1
99140
| otherwise -> helper xs i
100141

101-
{-# INLINE appOp #-}
102-
appOp :: MutableArray# RealWorld a -> Op a -> State# RealWorld -> State# RealWorld
103-
appOp arr (Set i x) s = writeArray# arr i x s
104-
appOp arr (Swap i j) s =
105-
case readArray# arr i s of { (# s, x #) ->
106-
case readArray# arr j s of { (# s, y #) ->
107-
case writeArray# arr i y s of { s ->
108-
writeArray# arr j x s
109-
}}}
110-
111-
{-# INLINE invert #-}
112-
invert :: MutableArray# RealWorld a -> Op a -> State# RealWorld -> (# State# RealWorld, Op a #)
113-
invert _ (Swap i j) s = (# s , Swap i j #)
114-
invert arr (Set i _) s =
115-
case readArray# arr i s of { (# s, y #) ->
116-
(# s, Set i y #) }
142+
{-# INLINE invertIO #-}
143+
invertIO :: MutableArray# RealWorld a -> Op a -> IO (Op a)
144+
invertIO _ (Swap i j) = pure (Swap i j)
145+
invertIO arr (Set i _) = do
146+
y <- readArrayIO arr i
147+
pure (Set i y)
148+
149+
{-# INLINE appOpIO #-}
150+
appOpIO :: MutableArray# RealWorld a -> Op a -> IO ()
151+
appOpIO arr (Set i x) = writeArrayIO arr i x
152+
appOpIO arr (Swap i j) = do
153+
x <- readArrayIO arr i
154+
y <- readArrayIO arr j
155+
writeArrayIO arr i y
156+
writeArrayIO arr j x
117157

118158
{-# INLINE appDiffOp #-}
119159
appDiffOp :: Op a -> Array a -> Array a
120-
appDiffOp op (DA v) = runRW# $ \s ->
121-
case noDuplicate# s of { s ->
122-
case readMutVar# v s of
123-
(# s , xs@(Current arr) #) ->
124-
case invert arr op s of { (# s, op' #) ->
125-
case appOp arr op s of { s ->
126-
case newMutVar# xs s of { (# s , v' #) ->
127-
case writeMutVar# v (Diff op' v') s of { !_ ->
128-
DA v'
129-
}}}}
130-
-- making a change to an old version of the array
131-
-- we copy to anticipate more usage
132-
(# s, Diff op' v' #) ->
133-
case copyInternal v' s of { (# s , arr #) ->
134-
case appOp arr op' s of { s ->
135-
case appOp arr op s of { s ->
136-
case newMutVar# (Current arr) s of { (# _ , v'' #) ->
137-
DA v''
138-
}}}}}
139-
140-
-- {-# INLINE appUndoOp #-}
141-
-- appUndoOp :: Op a -> Array a -> Array a
142-
-- appUndoOp op (DA v) = runRW# $ \s ->
143-
-- case readMutVar# v s of
144-
-- (# s , xs@(Current arr) #) ->
145-
-- case invert arr op s of { (# s, op' #) ->
146-
-- case appOp arr op s of { s ->
147-
-- case newMutVar# xs s of { (# s , v' #) ->
148-
-- case writeMutVar# v (Diff op' v') s of { !_ ->
149-
-- DA v'
150-
-- }}}}
151-
-- -- making a change to an old version of the array
152-
-- -- we roll back all changes
153-
-- (# s, Diff op' v' #) ->
154-
-- case appUndoOp ov' s of { (# s , arr #) ->
155-
-- case appOp arr op s of { s ->
156-
-- case appOp arr op s of { s ->
157-
-- case newMutVar# (Current arr) s of { (# _ , v'' #) ->
158-
-- DA v''
159-
-- }}}}
160+
appDiffOp op (DA v) = unsafeDupablePerformIO $
161+
readMutVarIO v $ \case
162+
xs@(Current arr) -> do
163+
op' <- invertIO arr op
164+
appOpIO arr op
165+
Lifted v' <- newMutVarIO xs
166+
writeMutVarIO v (Diff op' v')
167+
pure (DA v')
168+
Diff op' v' -> do
169+
Lifted arr <- copyInternalIO v'
170+
appOpIO arr op'
171+
appOpIO arr op
172+
Lifted v'' <- newMutVarIO (Current arr)
173+
pure (DA v'')
160174

161175
-- | Update the array element at a given position to a new value. O(1)
162176
{-# INLINE set #-}
@@ -168,21 +182,11 @@ set (I# i) x = appDiffOp (Set i x)
168182
swap :: Int -> Int -> Array a -> Array a
169183
swap (I# i) (I# j) = appDiffOp (Swap i j)
170184

171-
copyInternal :: MutVar# RealWorld (ArrayData a) -> State# RealWorld -> (# State# RealWorld, MutableArray# RealWorld a #)
172-
copyInternal v s =
173-
case readMutVar# v s of
174-
(# s , Current arr #) ->
175-
cloneMutableArray# arr 0# (sizeofMutableArray# arr) s
176-
(# s , Diff op v #) ->
177-
case copyInternal v s of { (# s , arr #) ->
178-
case appOp arr op s of { s -> (# s , arr #)
179-
}}
180-
181185
-- | Copy an array. O(n)
182186
-- This detaches any future updates from old versions of the array.
183187
-- Use this when you know you will be updating a large part of an array.
184188
copy :: Array a -> Array a
185-
copy (DA v) = runRW# $ \s ->
186-
case copyInternal v s of { (# s , arr #) ->
187-
case newMutVar# (Current arr) s of { (# _, v #) -> DA v
188-
}}
189+
copy (DA v) = unsafeDupablePerformIO $ do
190+
Lifted arr <- copyInternalIO v
191+
Lifted var <- newMutVarIO (Current arr)
192+
pure (DA var)

0 commit comments

Comments
 (0)