Adding swap function to strict trees

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 1.1.1.0 -- December 2024
+
+* Added `swap` function to strict PATRICIA and radix tree modules.
+
 ## 1.1.0.0 -- October 2024
 
 * Added `Zipper` modules for non-empty radix trees;

no/No/Tree.hs

@@ -35,6 +35,7 @@ module No.Tree
 
   , insert
   , insertWith
+  , swap
   , adjust
   , delete
   , update
@@ -221,6 +222,11 @@ insert k a = alter (\_ -> Just a) k
 insertWith :: Ord k => (a -> a) -> k -> a -> NoTree k a -> NoTree k a
 insertWith f k a = alter (Just . maybe a f) k
 
+swap :: Ord k => k -> a -> NoTree k a -> (Maybe a, NoTree k a)
+swap k a no =
+  let ~(NoTree as, mb, NoTree bs) = splitLookup k no
+  in (mb, NoTree $ as <> ((k, a) :<| bs))
+
 adjust :: Ord k => (a -> a) -> k -> NoTree k a -> NoTree k a
 adjust f = alter (fmap f)
 

radix-tree.cabal

@@ -1,5 +1,5 @@
 name:          radix-tree
-version:       1.1.0.0
+version:       1.1.1.0
 
 category:      Data Structures
 synopsis:      Radix trees

src/Data/Patricia/Word/Strict.hs

@@ -61,6 +61,10 @@ module Data.Patricia.Word.Strict
   , insertWith
   , insertWith'
 
+    -- | === Swap
+  , Swap (..)
+  , swap
+
     -- ** Map
   , adjust
   , adjust'

src/Data/Patricia/Word/Strict/Internal.hs

@@ -65,6 +65,9 @@ module Data.Patricia.Word.Strict.Internal
   , insertWith
   , insertWith'
 
+  , Swap (..)
+  , swap
+
   , adjust
   , adjust'
 
@@ -1459,6 +1462,38 @@ insertWith' f !w b = go
 
 
 
+-- | Result of a value swap.
+data Swap a = Swap !(Maybe a) !(Patricia a)
+              deriving Show
+
+-- | \(\mathcal{O}(\min(n,W))\).
+--   Insert a new value in the tree at the given key.
+--   If a value already exists at that key, it is returned alongside the tree.
+swap :: Word -> a -> Patricia a -> Swap a
+swap !w a = \t ->
+  let (# mb, t' #) = go t
+  in Swap mb t'
+  where
+    go t =
+      case t of
+        Bin p l r
+          | beyond p w -> (# Nothing, join w (Tip w a) p t #)
+
+          | w < p      -> let !(# mb, l' #) = go l
+                          in (# mb, Bin p l' r #)
+
+          | otherwise  -> let !(# mb, r' #) = go r
+                          in (# mb, Bin p l r' #)
+
+        Tip k b
+          | k == w    -> (# Just b, Tip k a #)
+          | otherwise -> (# Nothing, join w (Tip w a) k t #)
+
+        Nil -> (# Nothing, Tip w a #)
+
+
+
+
 -- | \(\mathcal{O}(\min(n,W))\).
 --   Apply a function to a value in the tree at the given key.
 adjust :: (a -> a) -> Word -> Patricia a -> Patricia a

src/Data/Radix1Tree/Word8/Strict.hs

@@ -73,6 +73,10 @@ module Data.Radix1Tree.Word8.Strict
   , insertWith
   , insertWith'
 
+    -- | === Swap
+  , Swap1 (..)
+  , swap
+
     -- ** Map
   , adjust
   , adjust'
@@ -500,6 +504,14 @@ insertWith' :: (a -> a) -> Feed1 -> a -> Radix1Tree a -> Radix1Tree a
 insertWith' = insertWith1'
 
 
+{-# INLINE swap #-}
+-- | \(\mathcal{O}(\min(x,k))\).
+--   Insert a new value in the tree at the given key.
+--   If a value already exists at that key, it is returned alongside the tree.
+swap :: Feed1 -> a -> Radix1Tree a -> Swap1 a
+swap = swap1
+
+
 {-# INLINE adjust #-}
 -- | \(\mathcal{O}(\min(x,k))\).
 --   Apply a function to a value in the tree at the given key.

src/Data/RadixNTree/Word8/Strict.hs

@@ -172,6 +172,11 @@ module Data.RadixNTree.Word8.Strict
   , insertWith1
   , insertWith1'
 
+  , Swap (..)
+  , Swap1 (..)
+  , swap0
+  , swap1
+
   , adjust0
   , adjust0'
 
@@ -4126,6 +4131,98 @@ insertWith'_ f a step = go
 
 
 
+-- | Result of a value swap.
+data Swap a = Swap !(Maybe a) !(RadixTree a)
+              deriving Show
+
+{-# INLINE swap0 #-}
+swap0 :: Feed -> a -> RadixTree a -> Swap a
+swap0 (Feed feed) a = \(RadixTree mx t) ->
+  feed $ \step s ->
+    case step s of
+      More w z -> let !(# mb, t' #) = swap_ a step w z t
+                  in Swap mb $ RadixTree mx t'
+
+      Done     -> Swap mx $ RadixTree (Just a) t
+
+
+-- | Result of a value swap.
+data Swap1 a = Swap1 !(Maybe a) !(Radix1Tree a)
+               deriving Show
+
+{-# INLINE swap1 #-}
+swap1 :: Feed1 -> a -> Radix1Tree a -> Swap1 a
+swap1 (Feed1 w feed) a = \t ->
+  feed $ \step s ->
+    let !(# mb, t' #) = swap_ a step w s t
+    in Swap1 mb t'
+
+
+{-# INLINE swap_ #-}
+swap_
+  :: a -> (x -> Step Word8 x) -> Word8 -> x -> Radix1Tree a -> (# Maybe a, Radix1Tree a #)
+swap_ a step = go
+  where
+    go !w !s t =
+      case t of
+        Bin p l r
+          | beyond p w -> let !t' = join
+                                      w (singleton_ step w s a)
+                                      p t
+
+                          in (# Nothing, t' #)
+
+          | w < p      -> let !(# mb, l' #) = go w s l
+                          in (# mb, Bin p l' r #)
+
+          | otherwise  -> let !(# mb, r' #) = go w s r
+                          in (# mb, Bin p l r' #)
+
+        Tip arr mx dx -> goarr w s 0
+          where
+            goarr v !z n
+              | v == indexByteArray arr n =
+                  if n + 1 >= sizeofByteArray arr
+                    then case step z of
+                           More u z' -> let !(# mb, dy #) = go u z' dx
+                                        in (# mb, Tip arr mx dy #)
+
+                           Done      -> let !dy = Tip arr (Just a) dx
+                                        in (# mx, dy #)
+
+                    else case step z of
+                           More u z' -> goarr u z' (n + 1)
+                           Done      ->
+                             let !(# !brr, !crr #) = splitByteArray 0 (n + 1) arr
+
+                                 !dy = Tip brr (Just a) (Tip crr mx dx)
+
+                             in (# Nothing, dy #)
+
+              | n == 0    =
+                  let !dy = join
+                              (indexByteArray arr 0) t
+                              w (singleton_ step w s a)
+
+                  in (# Nothing, dy #)
+
+              | otherwise =
+                  let !(# !brr, !crr #) = splitByteArray 0 n arr
+
+                      !dy = Tip brr Nothing $
+                              join
+                                (indexByteArray crr 0) (Tip crr mx dx)
+                                v (singleton_ step v z a)
+
+                  in (# Nothing, dy #)
+
+        Nil -> let !t' = singleton_ step w s a
+               in (# Nothing, t' #)
+
+
+
+
+
 {-# INLINE adjust0 #-}
 adjust0 :: (a -> a) -> Feed -> RadixTree a -> RadixTree a
 adjust0 f (Feed feed) = \(RadixTree mx t) ->

src/Data/RadixTree/Word8/Strict.hs

@@ -71,6 +71,10 @@ module Data.RadixTree.Word8.Strict
   , insertWith
   , insertWith'
 
+    -- | === Swap
+  , Swap (..)
+  , swap
+
     -- ** Map
   , adjust
   , adjust'
@@ -498,6 +502,14 @@ insertWith' :: (a -> a) -> Feed -> a -> RadixTree a -> RadixTree a
 insertWith' = insertWith0'
 
 
+{-# INLINE swap #-}
+-- | \(\mathcal{O}(\min(x,k))\).
+--   Insert a new value in the tree at the given key.
+--   If a value already exists at that key, it is returned alongside the tree.
+swap :: Feed -> a -> RadixTree a -> Swap a
+swap = swap0
+
+
 {-# INLINE adjust #-}
 -- | \(\mathcal{O}(\min(x,k))\).
 --   Apply a function to a value in the tree at the given key.

test/properties/Test/Patricia/Word/Strict.hs

@@ -126,6 +126,13 @@ maxViewEq (Just (Pat.ViewR pat (Pat.Lookup k a))) (Just (no, l, b)) =
 maxViewEq Nothing Nothing = True
 maxViewEq _       _       = False
 
+type SwapT s a =
+       Test s (Patricia a) (NoTree Word a)
+         (Pat.Swap a) (Maybe a, NoTree Word a)
+
+swapEq :: Eq a => Pat.Swap a -> (Maybe a, NoTree Word a) -> Bool
+swapEq (Pat.Swap mx pat) (my, no) = mx == my && treeEq pat no
+
 
 
 lookupT, dirtyLookupT :: Eq a => IdT Word a (Maybe a)
@@ -165,6 +172,9 @@ insertWithT_ g =
   let f x = (+ fromIntegral x)
   in Test treeEq (\(k, a) -> g (f a) k a) (\(k, a) -> No.insertWith (f a) k a)
 
+swapT :: Eq a => SwapT (Word, a) a
+swapT = Test swapEq (uncurry Pat.swap) (uncurry No.swap)
+
 adjustT, adjustT' :: (Eq a, Integral a) => TreeT (Word, a) a
 adjustT  = adjustT_ Pat.adjust
 adjustT' = adjustT_ Pat.adjust'
@@ -773,6 +783,7 @@ test = do
     it "insert"           $ run unary1  insertT
     it "insertWith"       $ run unary1  insertWithT
     it "insertWith'"      $ run unary1  insertWithT'
+    it "swap"             $ run unary1  swapT
     it "adjust"           $ run unary1  adjustT
     it "adjust'"          $ run unary1  adjustT'
     it "delete"           $ run unary1_ deleteT

test/properties/Test/RadixTree/Word8/Strict.hs

@@ -123,6 +123,13 @@ maxViewEq (Just (Radix.ViewR t k a)) (Just (no, l, b)) =
 maxViewEq Nothing                    Nothing           = True
 maxViewEq _                          _                 = False
 
+type SwapT s a =
+       Test s (RadixTree a) (NoTree [Word8] a)
+         (Radix.Swap a) (Maybe a, NoTree [Word8] a)
+
+swapEq :: Eq a => Radix.Swap a -> (Maybe a, NoTree [Word8] a) -> Bool
+swapEq (Radix.Swap mx pat) (my, no) = mx == my && treeEq pat no
+
 
 
 lookupT :: Eq a => IdT [Word8] a (Maybe a)
@@ -227,6 +234,9 @@ insertWithT_ g =
   in Test treeEq (\(k, a) -> g (f a) (Radix.feedBytes k) a)
                  (\(k, a) -> No.insertWith (f a) k a)
 
+swapT :: Eq a => SwapT ([Word8], a) a
+swapT = Test swapEq (\(k, i) -> Radix.swap (Radix.feedBytes k) i) (uncurry No.swap)
+
 adjustT, adjustT' :: (Eq a, Integral a) => TreeT ([Word8], a) a
 adjustT  = adjustT_ Radix.adjust
 adjustT' = adjustT_ Radix.adjust'
@@ -783,6 +793,7 @@ test = do
     it "insert"           $ run unary1  insertT
     it "insertWith"       $ run unary1  insertWithT
     it "insertWith'"      $ run unary1  insertWithT'
+    it "swap"             $ run unary1  swapT
     it "adjust"           $ run unary1  adjustT
     it "adjust'"          $ run unary1  adjustT'
     it "delete"           $ run unary1_ deleteT