@@ -23,15 +23,13 @@ module Data.TypeRepMap.Internal where
2323
2424import Prelude hiding (lookup)
2525
26−import Control.Monad.ST (runST)
26+import Control.Monad.ST (ST, runST)
2727import Control.Monad.Zip (mzip)
2828import Data.Function (on)
29−import Data.IntMap.Strict (IntMap)
3029import Data.Kind (Type)
3130import Data.List (intercalate, nubBy)
32−import Data.Maybe (fromJust)
33−import Data.Primitive.Array (Array, freezeArray, indexArray, mapArray', readArray, sizeofArray,
34− thawArray, writeArray)
31+import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', readArray, sizeofArray,
32+ thawArray, unsafeFreezeArray, writeArray)
3533import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
3634import Data.Semigroup (Semigroup (..))
3735import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
@@ -43,7 +41,6 @@ import Type.Reflection (SomeTypeRep (..), TypeRep, Typeable, typeRep, withTypeab
4341import Type.Reflection.Unsafe (typeRepFingerprint)
4442import Unsafe.Coerce (unsafeCoerce)
4543
46−import qualified Data.IntMap.Strict as IM
4744import qualified Data.Map.Strict as Map
4845import qualified GHC.Exts as GHC (fromList, toList)
4946
@@ -183,7 +180,7 @@ adjust fun tr = case cachedBinarySearch (typeFp @a) (fingerprintAs tr) (fingerpr
183180 mutArr <- thawArray trAs 0 n
184181 a <- toAny . fun . fromAny <$> readArray mutArr i
185182 writeArray mutArr i a
186− freezeArray mutArr 0 n
183+ unsafeFreezeArray mutArr
187184{-# INLINE adjust #-}
188185
189186{- | Map over the elements of a 'TypeRepMap'.
@@ -390,14 +387,22 @@ fromTriples kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromLi
390387----------------------------------------------------------------------------
391388
392389fromSortedList :: 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
394396 where
395397 -- state monad could be used here, but it's another dependency
396− go :: Int -> Int -> IntMap a -> IntMap a -> (IntMap a, Int)
397− go i first result vector =
398− if i >= IM.size vector
399− then (result, first)
400− else do
401− let (newResult, newFirst) = go (2 * i + 1) first result vector
402− let withCur = IM.insert i (fromJust $ IM.lookup newFirst vector) newResult
403− go (2 * i + 2) (newFirst + 1) withCur vector
398+ go :: forall s . Int -> MutableArray s a -> Array a -> ST s ()
399+ go len result origin = () <$ loop 0 0
400+ where
401+ loop :: Int -> Int -> ST s Int
402+ loop i first =
403+ if i >= len
404+ then pure first
405+ else do
406+ newFirst <- loop (2 * i + 1) first
407+ writeArray result i (indexArray origin newFirst)
408+ loop (2 * i + 2) (newFirst + 1)