Add alter

Author: bagl

docs/Data/IntMap.md

@@ -143,6 +143,19 @@ updateWithKey :: forall a. (Int -> a -> Maybe a) -> Int -> IntMap a -> IntMap a
 at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is
 deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
 
+#### `alter`
+
+``` purescript
+alter :: forall a. (Maybe a -> Maybe a) -> Int -> IntMap a -> IntMap a
+```
+
+/O(min(n,W))/. The expresion (@'alter' f k m@) alters the value @x@
+at key @k@, or absence thereof.
+'alter' can be used to insert, delete, or update the value under given
+key in the 'IntMap'.
+The following property holds:
+@'lookup' k ('alter' f k m) = f ('lookup' k m)@.
+
 #### `unionWith`
 
 ``` purescript

docs/Data/IntMap/Internal.md

@@ -79,6 +79,18 @@ branchingBit' :: Int -> Mask -> Int -> Mask -> Mask
 branchingBit :: Int -> Int -> Mask
 ```
 
+#### `branchMask`
+
+``` purescript
+branchMask :: Int -> Int -> Mask
+```
+
+#### `highestBitMask`
+
+``` purescript
+highestBitMask :: Int -> Int
+```
+
 #### `dec2bin`
 
 ``` purescript

src/Data/IntMap.purs

@@ -35,6 +35,7 @@ module Data.IntMap (
   , adjustWithKey
   , update
   , updateWithKey
+  , alter
 
   , unionWith
   , unionLeft
@@ -58,7 +59,7 @@ module Data.IntMap (
   ) where
 
 import Data.Foldable (class Foldable, foldMap, foldl)
-import Data.IntMap.Internal (Prefix, Mask(Mask), mask, branchLeft, branchingBit', prefixAsKey, matchPrefix, maskLonger)
+import Data.IntMap.Internal (Prefix(Prefix), Mask(Mask), mask, branchLeft, branchingBit', prefixAsKey, matchPrefix, maskLonger, branchMask)
 import Data.Maybe (Maybe(Nothing, Just))
 import Data.Monoid (class Monoid, mempty)
 import Data.Traversable (class Traversable)
@@ -226,8 +227,38 @@ updateWithKey f k t = go t where
         | branchLeft m k -> br p m (go l) r
         | otherwise -> br p m l (go r)
 
--- | Unions two `IntMap`s together using a splatting function. If 
--- | a key is present in both constituent lists then the resulting 
+-- | /O(min(n,W))/. The expresion (@'alter' f k m@) alters the value @x@
+-- | at key @k@, or absence thereof.
+-- | 'alter' can be used to insert, delete, or update the value under given
+-- | key in the 'IntMap'.
+-- | The following property holds:
+-- | @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
+alter :: forall a. (Maybe a -> Maybe a) -> Int -> IntMap a -> IntMap a
+alter f k t =
+  case t of
+    Br p@(Prefix p') m l r
+      | not (matchPrefix p m k) ->
+        case f Nothing of
+          Nothing -> t
+          Just a  -> link k (Lf k a) p' t
+      | branchLeft m k -> br p m (alter f k l) r
+      | otherwise      -> br p m l (alter f k r)
+    Lf ky y
+      | k == ky ->
+        case f (Just y) of
+          Just x  -> Lf ky x
+          Nothing -> Empty
+      | otherwise ->
+        case f Nothing of
+          Just x  -> link k (Lf k x) ky t
+          Nothing -> Lf ky y
+    Empty ->
+      case f Nothing of
+        Just a  -> Lf k a
+        Nothing -> Empty
+
+-- | Unions two `IntMap`s together using a splatting function. If
+-- | a key is present in both constituent lists then the resulting
 -- | list will be the splat of the values from each constituent. If the key
 -- | was available in only one constituent then it is available unmodified 
 -- | in the result.
@@ -378,3 +409,12 @@ join k1 m1 t1 k2 m2 t2 =
    in if branchLeft m k1
          then Br (mask m k1) m t1 t2
          else Br (mask m k1) m t2 t1
+
+link :: forall a. Int -> IntMap a -> Int -> IntMap a -> IntMap a
+link k1 t1 k2 t2 =
+  if branchLeft m k1
+  then Br p m t1 t2
+  else Br p m t2 t1
+  where
+    m = branchMask k1 k2
+    p = mask m k1

src/Data/IntMap/Internal.purs

@@ -2,7 +2,7 @@
 module Data.IntMap.Internal where
 
 import Data.Function (Fn2 (), runFn2)
-import Data.Int.Bits ((.^.), complement, (.&.), (.|.))
+import Data.Int.Bits ((.^.), complement, (.&.), (.|.), zshr)
 import Prelude (class Ord, class Eq, (*), (-), otherwise, (==), (<), compare, eq)
 
 -- Newtypes
@@ -57,6 +57,19 @@ branchingBit' k1 (Mask m1) k2 (Mask m2) =
 branchingBit :: Int -> Int -> Mask
 branchingBit k1 k2 = branchingBit' k1 (Mask 0) k2 (Mask 0)
 
+branchMask :: Int -> Int -> Mask
+branchMask x1 x2 =
+  Mask (highestBitMask (x1 .^. x2))
+
+highestBitMask :: Int -> Int
+highestBitMask x1 =
+  let x2 = x1 .|. x1 `zshr` 1
+      x3 = x2 .|. x2 `zshr` 2
+      x4 = x3 .|. x3 `zshr` 4
+      x5 = x4 .|. x4 `zshr` 8
+      x6 = x5 .|. x5 `zshr` 16
+  in x6 .^. (x6 `zshr` 1)
+
 foreign import dec2bin :: Int -> String
 foreign import bin2dec :: String -> Int
 foreign import pow :: Fn2 Int Int Int

test/Data/IntMap.purs

@@ -9,6 +9,7 @@ import           Prelude
 import qualified Test.Data.IntMap.Internal as Internal
 import           Test.Unit                 (Test (), test)
 import           Test.Unit.Assert          as Assert
+import Data.Tuple (Tuple(Tuple))
 
 (>|) :: forall a b . a -> (a -> b) -> b
 (>|) a f = f a
@@ -49,4 +50,28 @@ tests = do
       Assert.equal empty (filter (const false) ex2)
     test "filter by key" $
       Assert.equal (delete 20 ex2) (filterWithKey (\i _ -> i /= 20) ex2)
+    testAlter
     Internal.tests
+
+testAlter :: Test ()
+testAlter = do
+  test "alter" do
+    test "adding" do
+       Assert.equal (singleton 1 10) (alterIns 10 1 empty)
+       Assert.equal (singleton 1 10) (alterIns 10 1 (singleton 1 10))
+       Assert.equal (fromAssocArray [Tuple 1 10, Tuple 2 20, Tuple 3 30])
+         (alterIns 20 2 (fromAssocArray [Tuple 1 10, Tuple 3 30]))
+    test "deleting" do
+      Assert.equal empty (alterDel 1 empty)
+      Assert.equal empty (alterDel 1 (singleton 1 10))
+      Assert.equal (singleton 2 20) (alterDel 1 (singleton 2 20))
+      Assert.equal (singleton 2 20) (alterDel 1 (fromAssocArray [Tuple 1 10, Tuple 2 20]))
+    test "updating" do
+      Assert.equal empty (alterUpd 1 empty)
+      Assert.equal (singleton 2 20) (alterUpd 1 (singleton 2 20))
+      Assert.equal (fromAssocArray [Tuple 1 10, Tuple 2 21, Tuple 3 30])
+        (alterUpd 2 (fromAssocArray [Tuple 1 10, Tuple 2 20, Tuple 3 30]))
+  where
+    alterIns a = alter (maybe (Just a) Just)
+    alterDel   = alter (const Nothing)
+    alterUpd   = alter (map (_ + 1))

test/Data/IntMap/Internal.purs

@@ -26,6 +26,7 @@ tests = do
         Assert.equal "10000" (dec2bin (highestBit (bin2dec "10100") 1))
         Assert.equal "10000" (binBranchingBit "01010101" "01000001")
       Assert.equal "1000000" (dec2bin $ highestBit (bin2dec "1010101") (bin2dec "00000000001"))
+    testHighestBitMask
 
 testInversionTrick :: Test ()
 testInversionTrick =
@@ -41,3 +42,13 @@ binBranchingBit :: String -> String -> String
 binBranchingBit s1 s2 =
   case branchingBit (bin2dec s1) (bin2dec s2) of
     Mask b -> dec2bin b
+
+testHighestBitMask :: Test ()
+testHighestBitMask =
+  test "highest bit mask" do
+    eq "000000" "000000"
+    eq "000001" "000001"
+    eq "000010" "000011"
+    eq "010000" "010110"
+    eq "100000" "100101"
+  where eq m k = Assert.equal (bin2dec m) (highestBitMask $ bin2dec k)