git.iko.soy/typerep-map
on main help sign in
[#30] Rewrite fromSortedList to use arrays (#51)
32ebb33Veronika Romashkina
CHANGELOG.md+4
compare with main
unifiedsplit
@@ -11,6 +11,10 @@ The change log is available [on GitHub][2].
11 Add `keys` function.
11 Add `keys` function.
12* [#48](https://github.com/kowainik/typerep-map/issues/48):
12* [#48](https://github.com/kowainik/typerep-map/issues/48):
13 Add `adjust` function for `TypeRepMap` and `TMap`.
13 Add `adjust` function for `TypeRepMap` and `TMap`.
14+* [#30](https://github.com/kowainik/typerep-map/issues/30):
15+ Rewrite `fromSortedList` to use `Array` and `MutableArray`
16+ instead of `IntMap`.
17+
14
18
15# 0.2.0
19# 0.2.0
16
20
src/Data/TypeRepMap/Internal.hs+21−16
@@ -23,15 +23,13 @@ module Data.TypeRepMap.Internal where
23
23
24import Prelude hiding (lookup)
24import Prelude hiding (lookup)
25
25
26−import Control.Monad.ST (runST)
26+import Control.Monad.ST (ST, runST)
27import Control.Monad.Zip (mzip)
27import Control.Monad.Zip (mzip)
28import Data.Function (on)
28import Data.Function (on)
29−import Data.IntMap.Strict (IntMap)
30import Data.Kind (Type)
29import Data.Kind (Type)
31import Data.List (intercalate, nubBy)
30import Data.List (intercalate, nubBy)
32−import Data.Maybe (fromJust)
33−import Data.Primitive.Array (Array, freezeArray, indexArray, mapArray', readArray, sizeofArray,
31+import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', readArray, sizeofArray,
34− thawArray, writeArray)
32+ thawArray, unsafeFreezeArray, writeArray)
35import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
33import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
36import Data.Semigroup (Semigroup (..))
34import Data.Semigroup (Semigroup (..))
37import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
35import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
@@ -43,7 +41,6 @@ import Type.Reflection (SomeTypeRep (..), TypeRep, Typeable, typeRep, withTypeab
43import Type.Reflection.Unsafe (typeRepFingerprint)
41import Type.Reflection.Unsafe (typeRepFingerprint)
44import Unsafe.Coerce (unsafeCoerce)
42import Unsafe.Coerce (unsafeCoerce)
45
43
46−import qualified Data.IntMap.Strict as IM
47import qualified Data.Map.Strict as Map
44import qualified Data.Map.Strict as Map
48import qualified GHC.Exts as GHC (fromList, toList)
45import qualified GHC.Exts as GHC (fromList, toList)
49
46
@@ -183,7 +180,7 @@ adjust fun tr = case cachedBinarySearch (typeFp @a) (fingerprintAs tr) (fingerpr
183 mutArr <- thawArray trAs 0 n
180 mutArr <- thawArray trAs 0 n
184 a <- toAny . fun . fromAny <$> readArray mutArr i
181 a <- toAny . fun . fromAny <$> readArray mutArr i
185 writeArray mutArr i a
182 writeArray mutArr i a
186− freezeArray mutArr 0 n
183+ unsafeFreezeArray mutArr
187{-# INLINE adjust #-}
184{-# INLINE adjust #-}
188
185
189{- | Map over the elements of a 'TypeRepMap'.
186{- | Map over the elements of a 'TypeRepMap'.
@@ -390,14 +387,22 @@ fromTriples kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromLi
390----------------------------------------------------------------------------
387----------------------------------------------------------------------------
391
388
392fromSortedList :: forall a . [a] -> [a]
389fromSortedList :: forall a . [a] -> [a]
393−fromSortedList l = IM.elems $ fst $ go 0 0 mempty (IM.fromList $ zip [0..] l)
390+fromSortedList l = runST $ do
391+ let n = length l
392+ let arrOrigin = fromListN n l
393+ arrResult <- thawArray arrOrigin 0 n
394+ go n arrResult arrOrigin
395+ toList <$> unsafeFreezeArray arrResult
394 where
396 where
395 -- state monad could be used here, but it's another dependency
397 -- state monad could be used here, but it's another dependency
396− go :: Int -> Int -> IntMap a -> IntMap a -> (IntMap a, Int)
398+ go :: forall s . Int -> MutableArray s a -> Array a -> ST s ()
397− go i first result vector =
399+ go len result origin = () <$ loop 0 0
398− if i >= IM.size vector
400+ where
401+ loop :: Int -> Int -> ST s Int
402+ loop i first =
403+ if i >= len
399− then (result, first)
404+ then pure first
400− else do
405+ else do
401− let (newResult, newFirst) = go (2 * i + 1) first result vector
406+ newFirst <- loop (2 * i + 1) first
402− let withCur = IM.insert i (fromJust $ IM.lookup newFirst vector) newResult
407+ writeArray result i (indexArray origin newFirst)
403− go (2 * i + 2) (newFirst + 1) withCur vector
408+ loop (2 * i + 2) (newFirst + 1)