git.iko.soy/reflex
on js-mutable-intmap help sign in
reflex/src/Data/FastMutableIntMap.hs
237 lines 8.6 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE DeriveFunctor #-}
4{-# LANGUAGE DeriveFoldable #-}
5{-# LANGUAGE DeriveTraversable #-}
6{-# LANGUAGE GeneralizedNewtypeDeriving #-}
7{-# LANGUAGE LambdaCase #-}
8{-# LANGUAGE ScopedTypeVariables #-}
9{-# LANGUAGE TypeFamilies #-}
10#undef ghcjs_HOST_OS
11#ifdef ghcjs_HOST_OS
12{-# LANGUAGE ForeignFunctionInterface #-}
13{-# LANGUAGE JavaScriptFFI #-}
14{-# LANGUAGE MagicHash #-}
15{-# LANGUAGE UnliftedFFITypes #-}
16#endif
17module Data.FastMutableIntMap
18 ( FastMutableIntMap
19 , new
20 , newEmpty
21 , insert
22 , insertLookup
23 , deleteLookup
24 , isEmpty
25 , getFrozenAndClear
26 , size
27 , applyPatch
28 , PatchIntMap (..)
29 , traverseIntMapPatchWithKey
30 , lookup
31 , forIntersectionWithImmutable_
32 , for_
33 , patchIntMapNewElements
34 , patchIntMapNewElementsMap
35 , getDeletions
36 ) where
38--TODO: Pure JS version
39--TODO: Fast copy to FastIntMap
40--TODO: Fast patch type
42import Prelude hiding (lookup)
44import Data.IntMap.Strict (IntMap)
45import qualified Data.IntMap.Strict as IntMap
46import Data.Maybe
47import Reflex.Patch.Class
48import Control.Monad.IO.Class
50#ifdef ghcjs_HOST_OS
51import Control.Lens (imapM_, iforM_, iforM)
52import Control.Monad
53import GHCJS.Types
54import GHCJS.Marshal
55import JavaScript.Array (JSArray)
56import qualified JavaScript.Array as JSArray
57import GHCJS.Prim (JSVal (..))
58import GHC.Exts (ByteArray#)
59import GHC.Prim (unsafeCoerce#)
60#else
61import Data.IORef
62import Data.Foldable (traverse_)
63#endif
65newtype FastMutableIntMap a =
66#ifdef ghcjs_HOST_OS
67 FastMutableIntMap JSVal
68#else
69 FastMutableIntMap (IORef (IntMap a))
70#endif
72new :: IntMap a -> IO (FastMutableIntMap a)
73#ifdef ghcjs_HOST_OS
74new m = do
75 result <- newEmpty
76 imapM_ (insert result) m
77 return result
78#else
79new m = FastMutableIntMap <$> newIORef m
80#endif
82#ifdef ghcjs_HOST_OS
83foreign import javascript unsafe "var a = []; a.__ghcjsArray = true; $r = { f: { n: 'FastMutableIntMap', size: 2 }, d1: a, d2: { d1: 0 }, m: 0 }" newEmpty :: IO (FastMutableIntMap a)
84#else
85newEmpty :: IO (FastMutableIntMap a)
86newEmpty = FastMutableIntMap <$> newIORef IntMap.empty
87#endif
89insert :: forall a. FastMutableIntMap a -> Int -> a -> IO ()
90#ifdef ghcjs_HOST_OS
91insert m k !v = insert_ m k ((unsafeCoerce# :: a -> ByteArray#) v)
92foreign import javascript unsafe "if($1.d1[$2 & 0x7FFFFFFF] === undefined) { ++$1.d2.d1; } $1.d1[$2 & 0x7FFFFFFF] = $3;" insert_ :: FastMutableIntMap a -> Int -> ByteArray# -> IO ()
93#else
94insert (FastMutableIntMap r) k v = modifyIORef' r $ IntMap.insert k v
95#endif
97insertLookup :: forall a. FastMutableIntMap a -> Int -> a -> IO (Maybe a)
98#ifdef ghcjs_HOST_OS
99insertLookup m k !newVal = do
100 oldVal <- insertLookup_ m k ((unsafeCoerce# :: a -> ByteArray#) newVal)
101 return $ elemFromJSVal oldVal
102foreign import javascript unsafe "$r = $1.d1[$2 & 0x7FFFFFFF]; if($r === undefined) { ++$1.d2.d1; } $1.d1[$2 & 0x7FFFFFFF] = $3;" insertLookup_ :: FastMutableIntMap a -> Int -> ByteArray# -> IO JSVal
103#else
104insertLookup (FastMutableIntMap r) k v = do
105 oldMap <- readIORef r
106 let (oldVal, newMap) = IntMap.updateLookupWithKey (\_ _ -> Just v) k oldMap
107 writeIORef r $! newMap
108 return oldVal
109#endif
111deleteLookup :: FastMutableIntMap a -> Int -> IO (Maybe a)
112#ifdef ghcjs_HOST_OS
113deleteLookup m k = do
114 v <- deleteLookup_ m k
115 return $ elemFromJSVal v
116foreign import javascript unsafe "$r = $1.d1[$2 & 0x7FFFFFFF]; if($r !== undefined) { --$1.d2.d1; } delete $1.d1[$2 & 0x7FFFFFFF];" deleteLookup_ :: FastMutableIntMap a -> Int -> IO JSVal
117#else
118deleteLookup (FastMutableIntMap r) k = do
119 oldMap <- readIORef r
120 let (oldVal, newMap) = IntMap.updateLookupWithKey (\_ _ -> Nothing) k oldMap
121 writeIORef r $! newMap
122 return oldVal
123#endif
125lookup :: forall a. FastMutableIntMap a -> Int -> IO (Maybe a)
126#ifdef ghcjs_HOST_OS
127lookup m k = do
128 v <- lookup_ m k
129 return $ elemFromJSVal v
130foreign import javascript unsafe "$r = $1.d1[$2 & 0x7FFFFFFF];" lookup_ :: FastMutableIntMap a -> Int -> IO JS
132--TODO: Find out whether it's possible for Haskell objects to be represented as 'undefined'; if so, we need to do something slightly better to represent missing values
134elemFromJSVal :: JSVal -> Maybe a
135elemFromJSVal v = case isUndefined v of
136 True -> Nothing
137 False -> Just $ (unsafeCoerce# :: ByteArray# -> a) (case v of JSVal raw -> raw)
138Val
139#else
140lookup (FastMutableIntMap r) k = IntMap.lookup k <$> readIORef r
141#endif
143--TODO: Do something faster than trying every one
144{-# INLINE forIntersectionWithImmutable_ #-}
145forIntersectionWithImmutable_ :: MonadIO m => FastMutableIntMap a -> IntMap b -> (a -> b -> m ()) -> m ()
146#ifdef ghcjs_HOST_OS
147forIntersectionWithImmutable_ ma mb f = iforM_ mb $ \k b -> do
148 liftIO (lookup ma k) >>= \case
149 Nothing -> return ()
150 Just a -> f a b
151#else
152forIntersectionWithImmutable_ (FastMutableIntMap r) b f = do
153 a <- liftIO $ readIORef r
154 traverse_ (uncurry f) $ IntMap.intersectionWith (,) a b
155#endif
157{-# INLINE for_ #-}
158for_ :: MonadIO m => FastMutableIntMap a -> (a -> m ()) -> m ()
159#ifdef ghcjs_HOST_OS
160for_ m f = do
161 raw <- liftIO $ packFastMutableIntMap m --TODO: Don't pack the keys, since we don't need them
162 forM_ (JSArray.toList raw) $ \p -> do
163 (_ :: Int, JSVal v) <- liftIO $ fromJSValUnchecked p
164 f $ (unsafeCoerce# :: ByteArray# -> a) v
165#else
166for_ (FastMutableIntMap r) f = do
167 a <- liftIO $ readIORef r
168 traverse_ f a
169#endif
171#ifdef ghcjs_HOST_OS
172foreign import javascript unsafe "$r = $1.d2.d1 === 0;" isEmpty :: FastMutableIntMap a -> IO Bool
173#else
174isEmpty :: FastMutableIntMap a -> IO Bool
175isEmpty (FastMutableIntMap r) = IntMap.null <$> readIORef r
176#endif
178#ifdef ghcjs_HOST_OS
179foreign import javascript unsafe "$r = $1.d2.d1;" size :: FastMutableIntMap a -> IO Int
180#else
181size :: FastMutableIntMap a -> IO Int
182size (FastMutableIntMap r) = IntMap.size <$> readIORef r
183#endif
185-- | Make an immutable snapshot of the datastructure and clear it
186getFrozenAndClear :: FastMutableIntMap a -> IO (IntMap a)
187#ifdef ghcjs_HOST_OS
188getFrozenAndClear m = do
189 raw <- packFastMutableIntMap m
190 clear m
191 pairs <- forM (JSArray.toList raw) $ \p -> do
192 (k, JSVal v) <- fromJSValUnchecked p
193 return (k, (unsafeCoerce# :: ByteArray# -> a) v)
194 return $ IntMap.fromDistinctAscList pairs
195-- | Make a copy of a FastMutableIntMap's contents, as an array of key/value pairs
196foreign import javascript unsafe "$r = new Array($1.d2.d1); $r.__ghcjsArray = true; var i = 0; for(var k in $1.d1) { if(k !== '__ghcjsArray') { var p = [k, $1.d1[k]]; p.__ghcjsArray = true; $r[i++] = p; } } $r.sort(function(a, b) { return a[0] - b[0]; });" packFastMutableIntMap :: FastMutableIntMap a -> IO JSArray
197foreign import javascript unsafe "$1.d1 = []; $1.d1.__ghcjsArray = true; $1.d2.d1 = 0;" clear :: FastMutableIntMap a -> IO ()
198--NOTE: We would like to use Array.forEach here, but, at least on Chrome 59, forEach appears to be linear in the maximum key in the array. Since we sometimes use very large indexes, this is unusable. for(... in ...), on the other hand, appears to actually operate sparsely - but this requires us to sort the array in a second step.
199#else
200getFrozenAndClear (FastMutableIntMap r) = do
201 result <- readIORef r
202 writeIORef r IntMap.empty
203 return result
204#endif
206applyPatch :: FastMutableIntMap a -> PatchIntMap a -> IO (IntMap a)
207#ifdef ghcjs_HOST_OS
208applyPatch f (PatchIntMap m) = fmap (IntMap.mapMaybe id) $ iforM m $ \k mv -> case mv of
209 Nothing -> deleteLookup f k
210 Just v -> insertLookup f k v
211#else
212applyPatch (FastMutableIntMap r) p@(PatchIntMap m) = do
213 v <- readIORef r
214 writeIORef r $! applyAlways p v
215 return $ IntMap.intersection v m
216#endif
218newtype PatchIntMap a = PatchIntMap (IntMap (Maybe a)) deriving (Functor, Foldable, Traversable, Monoid)
220instance Patch (PatchIntMap a) where
221 type PatchTarget (PatchIntMap a) = IntMap a
222 apply (PatchIntMap p) v = if IntMap.null p then Nothing else Just $
223 let removes = IntMap.filter isNothing p
224 adds = IntMap.mapMaybe id p
225 in IntMap.union adds $ v `IntMap.difference` removes
227traverseIntMapPatchWithKey :: Applicative t => (Int -> a -> t b) -> PatchIntMap a -> t (PatchIntMap b)
228traverseIntMapPatchWithKey f (PatchIntMap m) = PatchIntMap <$> IntMap.traverseWithKey (\k mv -> traverse (f k) mv) m
230patchIntMapNewElements :: PatchIntMap a -> [a]
231patchIntMapNewElements (PatchIntMap m) = catMaybes $ IntMap.elems m
233patchIntMapNewElementsMap :: PatchIntMap a -> IntMap a
234patchIntMapNewElementsMap (PatchIntMap m) = IntMap.mapMaybe id m
236getDeletions :: PatchIntMap v -> IntMap v' -> IntMap v'
237getDeletions (PatchIntMap m) v = IntMap.intersection v m