diff --git a/CHANGELOG.md b/CHANGELOG.md index f725902..9aae74c 100644 --- a/CHANGELOG.md +++ b/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; diff --git a/no/No/Tree.hs b/no/No/Tree.hs index c4e7acf..2a4c80b 100644 --- a/no/No/Tree.hs +++ b/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) diff --git a/radix-tree.cabal b/radix-tree.cabal index b7b6cbe..9fa36d0 100644 --- a/radix-tree.cabal +++ b/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 diff --git a/src/Data/Patricia/Word/Strict.hs b/src/Data/Patricia/Word/Strict.hs index ae4afd5..ae64b35 100644 --- a/src/Data/Patricia/Word/Strict.hs +++ b/src/Data/Patricia/Word/Strict.hs @@ -61,6 +61,10 @@ module Data.Patricia.Word.Strict , insertWith , insertWith' + -- | === Swap + , Swap (..) + , swap + -- ** Map , adjust , adjust' diff --git a/src/Data/Patricia/Word/Strict/Internal.hs b/src/Data/Patricia/Word/Strict/Internal.hs index d504144..c292b17 100644 --- a/src/Data/Patricia/Word/Strict/Internal.hs +++ b/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,44 @@ 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 -> let !t' = join w (Tip w a) p t + in (# Nothing, t'#) + + | w < p -> let !(# mb, l' #) = go l + !t' = Bin p l' r + in (# mb, t' #) + + | otherwise -> let !(# mb, r' #) = go r + !t' = Bin p l r' + in (# mb, t' #) + + Tip k b + | k == w -> let !t' = Tip k a + in (# Just b, t' #) + | otherwise -> let !t' = join w (Tip w a) k t + in (# Nothing, t' #) + + Nil -> let !t' = Tip w a + in (# Nothing, t' #) + + + + -- | \(\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 diff --git a/src/Data/Radix1Tree/Word8/Strict.hs b/src/Data/Radix1Tree/Word8/Strict.hs index f3dc241..4e3a7fa 100644 --- a/src/Data/Radix1Tree/Word8/Strict.hs +++ b/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. diff --git a/src/Data/RadixNTree/Word8/Strict.hs b/src/Data/RadixNTree/Word8/Strict.hs index 1b92d33..8d2c2c0 100644 --- a/src/Data/RadixNTree/Word8/Strict.hs +++ b/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,101 @@ 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 + !t' = Bin p l' r + in (# mb, t' #) + + | otherwise -> let !(# mb, r' #) = go w s r + !t' = Bin p l r' + in (# mb, t' #) + + 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 + !t' = Tip arr mx dy + in (# mb, t' #) + + Done -> let !t' = Tip arr (Just a) dx + in (# mx, t' #) + + else case step z of + More u z' -> goarr u z' (n + 1) + Done -> + let !(# !brr, !crr #) = splitByteArray 0 (n + 1) arr + + !t' = Tip brr (Just a) (Tip crr mx dx) + + in (# Nothing, t' #) + + | n == 0 = + let !t' = join + (indexByteArray arr 0) t + w (singleton_ step w s a) + + in (# Nothing, t' #) + + | otherwise = + let !(# !brr, !crr #) = splitByteArray 0 n arr + + !t' = Tip brr Nothing $ + join + (indexByteArray crr 0) (Tip crr mx dx) + v (singleton_ step v z a) + + in (# Nothing, t' #) + + 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) -> diff --git a/src/Data/RadixTree/Word8/Strict.hs b/src/Data/RadixTree/Word8/Strict.hs index 5ba827e..3cef7f0 100644 --- a/src/Data/RadixTree/Word8/Strict.hs +++ b/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. diff --git a/test/properties/Test/Patricia/Word/Strict.hs b/test/properties/Test/Patricia/Word/Strict.hs index 2873085..0a1af9d 100644 --- a/test/properties/Test/Patricia/Word/Strict.hs +++ b/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 diff --git a/test/properties/Test/RadixTree/Word8/Strict.hs b/test/properties/Test/RadixTree/Word8/Strict.hs index f945e50..d83e87b 100644 --- a/test/properties/Test/RadixTree/Word8/Strict.hs +++ b/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