Make size run in constant time

Author: rockbmb

Data/HashMap/Array.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE BangPatterns, CPP, MagicHash, Rank2Types, UnboxedTuples, ScopedTypeVariables #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# OPTIONS_GHC -fno-full-laziness -funbox-strict-fields #-}
 
 -- | Zero based arrays.
@@ -7,6 +8,11 @@
 module Data.HashMap.Array
     ( Array
     , MArray
+    , RunRes (..)
+    , RunResA
+    , RunResM
+    , Size (..)
+    , Sized (..)
 
       -- * Creation
     , new
@@ -25,6 +31,7 @@ module Data.HashMap.Array
     , index#
     , update
     , updateWith'
+    , updateWithInternal'
     , unsafeUpdateM
     , insert
     , insertM
@@ -36,6 +43,7 @@ module Data.HashMap.Array
     , unsafeThaw
     , unsafeSameArray
     , run
+    , runInternal
     , run2
     , copy
     , copyM
@@ -317,10 +325,24 @@ unsafeThaw ary
                    (# s', mary #) -> (# s', marray mary (length ary) #)
 {-# INLINE unsafeThaw #-}
 
+-- | Helper datatype used in 'runInternal' and 'updateWithInternal'
+data RunRes f e = RunRes {-# UNPACK #-} !Size !(f e)
+
+type RunResA e = RunRes Array e
+
+type RunResM s e = RunRes (MArray s) e
+
 run :: (forall s . ST s (MArray s e)) -> Array e
 run act = runST $ act >>= unsafeFreeze
 {-# INLINE run #-}
 
+runInternal :: (forall s . ST s (RunResM s e)) -> RunResA e
+runInternal act = runST $ do
+    RunRes s mary <- act
+    ary <- unsafeFreeze mary
+    return (RunRes s ary)
+{-# INLINE runInternal #-}
+
 run2 :: (forall s. ST s (MArray s e, a)) -> (Array e, a)
 run2 k = runST (do
                  (marr,b) <- k
@@ -392,7 +414,7 @@ updateM ary idx b =
   where !count = length ary
 {-# INLINE updateM #-}
 
--- | /O(n)/ Update the element at the given positio in this array, by
+-- | /O(n)/ Update the element at the given position in this array, by
 -- applying a function to it.  Evaluates the element to WHNF before
 -- inserting it into the array.
 updateWith' :: Array e -> Int -> (e -> e) -> Array e
@@ -401,6 +423,26 @@ updateWith' ary idx f
   = update ary idx $! f x
 {-# INLINE updateWith' #-}
 
+-- | This newtype wrapper is to avoid confusion when local functions
+-- take more than one paramenter of 'Int' type (see 'go' in
+-- 'Data.HashMap.Base.unionWithKeyInternal').
+newtype Size = Size { unSize :: Int }
+    deriving (Eq, Ord, Num, Integral, Enum, Real)
+
+-- | Helper datatype used in 'updateWithInternal''. Used when a change in
+-- a value's size must be returned along with the value itself (typically
+-- a hashmap).
+data Sized a = Sized {-# UNPACK #-} !Size !a
+
+-- | /O(n)/ Update the element at the given position in this array, by
+-- applying a function to it.  Evaluates the element to WHNF before
+-- inserting it into the array.
+updateWithInternal' :: Array e -> Int -> (e -> Sized e) -> RunResA e
+updateWithInternal' ary idx f =
+    let Sized sz e = f (index ary idx)
+    in RunRes sz (update ary idx e)
+{-# INLINE updateWithInternal' #-}
+
 -- | /O(1)/ Update the element at the given position in this array,
 -- without copying.
 unsafeUpdateM :: Array e -> Int -> e -> ST s ()