Add deleteAt to Data.Sequence
authorDavid Feuer <David.Feuer@gmail.com>
Mon, 30 May 2016 22:44:00 +0000 (18:44 -0400)
committerDavid Feuer <David.Feuer@gmail.com>
Mon, 30 May 2016 22:55:49 +0000 (18:55 -0400)
This is messy, and may remain so, but it would be nice to find
some way to clean it up a bit. Also, we can and should be
more eager about rebuilding and should look for opportunities to
use arithmetic to figure out sizes.

Data/Sequence.hs
benchmarks/Sequence.hs
tests/seq-properties.hs

index 46a6e0f..ef4e3cc 100644 (file)
@@ -153,6 +153,7 @@ module Data.Sequence (
     take,           -- :: Int -> Seq a -> Seq a
     drop,           -- :: Int -> Seq a -> Seq a
     insertAt,       -- :: Int -> a -> Seq a -> Seq a
+    deleteAt,       -- :: Int -> Seq a -> Seq a
     splitAt,        -- :: Int -> Seq a -> (Seq a, Seq a)
     -- ** Indexing with predicates
     -- | These functions perform sequential searches from the left
@@ -1824,6 +1825,243 @@ insNode f i (Node3 s a b c)
   where sa = size a
         sab = sa + size b
 
+-- | /O(log(min(i,n-i)))/. Delete the element of a sequence at a given
+-- index. Return the original sequence if the index is out of range.
+--
+-- @since 0.5.8
+deleteAt :: Int -> Seq a -> Seq a
+deleteAt i (Seq xs)
+  | fromIntegral i < (fromIntegral (size xs) :: Word) = Seq $ delTreeE i xs
+  | otherwise = Seq xs
+
+delTreeE :: Int -> FingerTree (Elem a) -> FingerTree (Elem a)
+delTreeE !_i EmptyT = EmptyT -- Unreachable
+delTreeE _i Single{} = EmptyT
+delTreeE i (Deep s pr m sf)
+  | i < spr = delLeftDigitE i pr m sf
+  | i < spm = case delTree delNodeE (i - spr) m of
+     FullTree m' -> Deep (s - 1) pr m' sf
+     DefectTree e -> delRebuildMiddle  (s - 1) pr e sf
+  | otherwise = delRightDigitE (i - spm) pr m sf
+  where spr = size pr
+        spm = spr + size m
+
+delNodeE :: Int -> Node (Elem a) -> Del (Elem a)
+delNodeE i (Node3 _ a b c) = case i of
+  0 -> Full $ Node2 2 b c
+  1 -> Full $ Node2 2 a c
+  _ -> Full $ Node2 2 a b
+delNodeE i (Node2 _ a b) = case i of
+  0 -> Defect b
+  _ -> Defect a
+
+
+delLeftDigitE :: Int -> Digit (Elem a) -> FingerTree (Node (Elem a)) -> Digit (Elem a) -> FingerTree (Elem a)
+delLeftDigitE !_i One{} m sf = pullL (size m + size sf) m sf
+delLeftDigitE i (Two a b) m sf
+  | i == 0 = deep (One b) m sf
+  | otherwise = deep (One a) m sf
+delLeftDigitE i (Three a b c) m sf
+  | i == 0 = deep (Two b c) m sf
+  | i == 1 = deep (Two a c) m sf
+  | otherwise = deep (Two a b) m sf
+delLeftDigitE i (Four a b c d) m sf
+  | i == 0 = deep (Three b c d) m sf
+  | i == 1 = deep (Three a c d) m sf
+  | i == 2 = deep (Three a b d) m sf
+  | otherwise = deep (Three a b c) m sf
+
+delRightDigitE :: Int -> Digit (Elem a) -> FingerTree (Node (Elem a)) -> Digit (Elem a) -> FingerTree (Elem a)
+delRightDigitE !_i pr m One{} = pullR (size pr + size m) pr m
+delRightDigitE i pr m (Two a b)
+  | i == 0 = deep pr m (One b)
+  | otherwise = deep pr m (One a)
+delRightDigitE i pr m (Three a b c)
+  | i == 0 = deep pr m (Two b c)
+  | i == 1 = deep pr m (Two a c)
+  | otherwise = deep pr m (Two a b)
+delRightDigitE i pr m (Four a b c d)
+  | i == 0 = deep pr m (Three b c d)
+  | i == 1 = deep pr m (Three a c d)
+  | i == 2 = deep pr m (Three a b d)
+  | otherwise = deep pr m (Three a b c)
+
+data DelTree a = FullTree !(FingerTree (Node a)) | DefectTree a
+
+delTree :: Sized a => (Int -> Node a -> Del a) -> Int -> FingerTree (Node a) -> DelTree a
+delTree _f !_i EmptyT = FullTree EmptyT -- Unreachable
+delTree f i (Single a) = case f i a of
+  Full a' -> FullTree (Single a')
+  Defect e -> DefectTree e
+delTree f i (Deep s pr m sf)
+  | i < spr = case delDigit f i pr of
+     FullDig pr' -> FullTree $ Deep (s - 1) pr' m sf
+     DefectDig e -> case viewLTree m of
+                      EmptyLTree -> FullTree $ delRebuildRightDigit e sf
+                      ConsLTree n m' -> FullTree $ delRebuildLeftSide e n m' sf
+  | i < spm = case delTree (delNode f) (i - spr) m of
+     FullTree m' -> FullTree (Deep (s - 1) pr m' sf)
+     DefectTree e -> FullTree $ delRebuildMiddle (s - 1) pr e sf
+  | otherwise = case delDigit f (i - spm) sf of
+     FullDig sf' -> FullTree $ Deep (s - 1) pr m sf'
+     DefectDig e -> case viewRTree m of
+                      EmptyRTree -> FullTree $ delRebuildLeftDigit pr e
+                      SnocRTree m' n -> FullTree $ delRebuildRightSide pr m' n e
+  where spr = size pr
+        spm = spr + size m
+
+data Del a = Full !(Node a) | Defect a
+-- TODO: strictify node construction and use arithmetic for sizes
+delNode :: Sized a => (Int -> Node a -> Del a) -> Int -> Node (Node a) -> Del (Node a)
+delNode f i (Node3 s a b c)
+  | i < sa = case f i a of
+     Full a' -> Full $ Node3 (s - 1) a' b c
+     Defect e -> case b of
+       Node3 _ x y z -> Full $ Node3 (s - 1) (node2 e x) (node2 y z) c
+       Node2 _ x y -> Full $ Node2 (s - 1) (node3 e x y) c
+  | i < sab = case f (i - sa) b of
+     Full b' -> Full $ Node3 (s - 1) a b' c
+     Defect e -> case a of
+       Node3 _ x y z -> Full $ Node3 (s - 1) (node2 x y) (node2 z e) c
+       Node2 _ x y -> Full $ Node2 (s - 1) (node3 x y e) c
+  | otherwise = case f (i - sab) c of
+     Full c' -> Full $ Node3 (s - 1) a b c'
+     Defect e -> case b of
+       Node3 _ x y z -> Full $ Node3 (s - 1) a (node2 x y) (node2 z e)
+       Node2 _ x y -> Full $ Node2 (s - 1) a (node3 x y e)
+  where sa = size a
+        sab = sa + size b
+delNode f i (Node2 s a b)
+  | i < sa = case f i a of
+     Full a' -> Full $ Node2 (s - 1) a' b
+     Defect e -> case b of
+       Node3 _ x y z -> Full $ Node2 (s - 1) (node2 e x) (node2 y z)
+       Node2 _ x y -> Defect $ Node3 (s - 1) e x y
+  | otherwise = case f (i - sa) b of
+     Full b' -> Full $ Node2 (s - 1) a b'
+     Defect e -> case a of
+       Node3 _ x y z -> Full $ Node2 (s - 1) (node2 x y) (node2 z e)
+       Node2 _ x y -> Defect $ Node3 (s - 1) x y e
+  where sa = size a
+
+delRebuildRightDigit :: Sized a => a -> Digit (Node a) -> FingerTree (Node a)
+delRebuildRightDigit p (One a) = case a of
+  Node3 _ x y z -> deep (One (node2 p x)) EmptyT (One (node2 y z))
+  Node2 _ x y -> Single (node3 p x y)
+delRebuildRightDigit p (Two a b) = case a of
+  Node3 _ x y z -> deep (Two (node2 p x) (node2 y z)) EmptyT (One b)
+  Node2 _ x y -> deep (One (node3 p x y)) EmptyT (One b)
+delRebuildRightDigit p (Three a b c) = case a of
+  Node3 _ x y z -> deep (Two (node2 p x) (node2 y z)) EmptyT (Two b c)
+  Node2 _ x y -> deep (Two (node3 p x y) b) EmptyT (One c)
+delRebuildRightDigit p (Four a b c d) = case a of
+  Node3 _ x y z -> deep (Three (node2 p x) (node2 y z) b) EmptyT (Two c d)
+  Node2 _ x y -> deep (Two (node3 p x y) b) EmptyT (Two c d)
+
+delRebuildLeftDigit :: Sized a => Digit (Node a) -> a -> FingerTree (Node a)
+delRebuildLeftDigit (One a) p = case a of
+  Node3 _ x y z -> deep (One (node2 x y)) EmptyT (One (node2 z p))
+  Node2 _ x y -> Single (node3 x y p)
+delRebuildLeftDigit (Two a b) p = case b of
+  Node3 _ x y z -> deep (Two a (node2 x y)) EmptyT (One (node2 z p))
+  Node2 _ x y -> deep (One a) EmptyT (One (node3 x y p))
+delRebuildLeftDigit (Three a b c) p = case c of
+  Node3 _ x y z -> deep (Two a b) EmptyT (Two (node2 x y) (node2 z p))
+  Node2 _ x y -> deep (Two a b) EmptyT (One (node3 x y p))
+delRebuildLeftDigit (Four a b c d) p = case d of
+  Node3 _ x y z -> deep (Three a b c) EmptyT (Two (node2 x y) (node2 z p))
+  Node2 _ x y -> deep (Two a b) EmptyT (Two c (node3 x y p))
+
+delRebuildLeftSide :: Sized a
+                   => a -> Node (Node a) -> FingerTree (Node (Node a)) -> Digit (Node a)
+                   -> FingerTree (Node a)
+delRebuildLeftSide p (Node2 _ a b) m sf = case a of
+  Node2 _ x y -> deep (Two (node3 p x y) b) m sf
+  Node3 _ x y z -> deep (Three (node2 p x) (node2 y z) b) m sf
+delRebuildLeftSide p (Node3 _ a b c) m sf = case a of
+  Node2 _ x y -> deep (Three (node3 p x y) b c) m sf
+  Node3 _ x y z -> deep (Four (node2 p x) (node2 y z) b c) m sf
+
+delRebuildRightSide :: Sized a
+                    => Digit (Node a) -> FingerTree (Node (Node a)) -> Node (Node a) -> a
+                    -> FingerTree (Node a)
+delRebuildRightSide pr m (Node2 _ a b) p = case b of
+  Node2 _ x y -> deep pr m (Two a (node3 x y p))
+  Node3 _ x y z -> deep pr m (Three a (node2 x y) (node2 z p))
+delRebuildRightSide pr m (Node3 _ a b c) p = case c of
+  Node2 _ x y -> deep pr m (Three a b (node3 x y p))
+  Node3 _ x y z -> deep pr m (Four a b (node2 x y) (node2 z p))
+
+delRebuildMiddle :: Sized a
+                 => Int -> Digit a -> a -> Digit a
+                 -> FingerTree a
+delRebuildMiddle s (One a) e sf = Deep s (Two a e) EmptyT sf
+delRebuildMiddle s (Two a b) e sf = Deep s (Three a b e) EmptyT sf
+delRebuildMiddle s (Three a b c) e sf = Deep s (Four a b c e) EmptyT sf
+delRebuildMiddle s (Four a b c d) e sf = Deep s (Two a b) (Single (node3 c d e)) sf
+
+
+data DelDig a = FullDig !(Digit (Node a)) | DefectDig a
+delDigit :: Sized a => (Int -> Node a -> Del a) -> Int -> Digit (Node a) -> DelDig a
+delDigit f !i (One a) = case f i a of
+  Full a' -> FullDig $ One a'
+  Defect e -> DefectDig e
+delDigit f i (Two a b)
+  | i < sa = case f i a of
+     Full a' -> FullDig $ Two a' b
+     Defect e -> case b of
+       Node3 _ x y z -> FullDig $ Two (node2 e x) (node2 y z)
+       Node2 _ x y -> FullDig $ One (node3 e x y)
+  | otherwise = case f (i - sa) b of
+     Full b' -> FullDig $ Two a b'
+     Defect e -> case a of
+       Node3 _ x y z -> FullDig $ Two (node2 x y) (node2 z e)
+       Node2 _ x y -> FullDig $ One (node3 x y e)
+  where sa = size a
+delDigit f i (Three a b c)
+  | i < sa = case f i a of
+     Full a' -> FullDig $ Three a' b c
+     Defect e -> case b of
+       Node3 _ x y z -> FullDig $ Three (node2 e x) (node2 y z) c
+       Node2 _ x y -> FullDig $ Two (node3 e x y) c
+  | i < sab = case f (i - sa) b of
+     Full b' -> FullDig $ Three a b' c
+     Defect e -> case a of
+       Node3 _ x y z -> FullDig $ Three (node2 x y) (node2 z e) c
+       Node2 _ x y -> FullDig $ Two (node3 x y e) c
+  | otherwise = case f (i - sab) c of
+     Full c' -> FullDig $ Three a b c'
+     Defect e -> case b of
+       Node3 _ x y z -> FullDig $ Three a (node2 x y) (node2 z e)
+       Node2 _ x y -> FullDig $ Two a (node3 x y e)
+  where sa = size a
+        sab = sa + size b
+delDigit f i (Four a b c d)
+  | i < sa = case f i a of
+     Full a' -> FullDig $ Four a' b c d
+     Defect e -> case b of
+       Node3 _ x y z -> FullDig $ Four (node2 e x) (node2 y z) c d
+       Node2 _ x y -> FullDig $ Three (node3 e x y) c d
+  | i < sab = case f (i - sa) b of
+     Full b' -> FullDig $ Four a b' c d
+     Defect e -> case a of
+       Node3 _ x y z -> FullDig $ Four (node2 x y) (node2 z e) c d
+       Node2 _ x y -> FullDig $ Three (node3 x y e) c d
+  | i < sabc = case f (i - sab) c of
+     Full c' -> FullDig $ Four a b c' d
+     Defect e -> case b of
+       Node3 _ x y z -> FullDig $ Four a (node2 x y) (node2 z e) d
+       Node2 _ x y -> FullDig $ Three a (node3 x y e) d
+  | otherwise = case f (i - sabc) d of
+     Full d' -> FullDig $ Four a b c d'
+     Defect e -> case c of
+       Node3 _ x y z -> FullDig $ Four a b (node2 x y) (node2 z e)
+       Node2 _ x y -> FullDig $ Three a b (node3 x y e)
+  where sa = size a
+        sab = sa + size b
+        sabc = sab + size c
+
+
 data InsDigNode a = InsLeftDig !(Digit a) | InsDigNode !(Digit a) !(Node a)
 {-# SPECIALIZE insLeftDigit :: (Int -> Elem a -> Ins (Elem a)) -> Int -> Digit (Elem a) -> InsDigNode (Elem a) #-}
 {-# SPECIALIZE insLeftDigit :: (Int -> Node a -> Ins (Node a)) -> Int -> Digit (Node a) -> InsDigNode (Node a) #-}
index 717daf9..280a947 100644 (file)
@@ -35,6 +35,11 @@ main = do
          , bench "100" $ nf (shuffle r100) s100
          , bench "1000" $ nf (shuffle r1000) s1000
          ]
+      , bgroup "deleteAt"
+         [ bench "10" $ nf (deleteAtPoints r10) s10
+         , bench "100" $ nf (deleteAtPoints r100) s100
+         , bench "1000" $ nf (deleteAtPoints r1000) s1000
+         ]
       , bgroup "insertAt"
          [ bench "10" $ nf (insertAtPoints r10 10) s10
          , bench "100" $ nf (insertAtPoints r100 10) s100
@@ -101,6 +106,20 @@ insertAtPoints :: [Int] -> a -> S.Seq a -> S.Seq a
 insertAtPoints points x xs =
   foldl' (\acc k -> S.insertAt k x acc) xs points
 
+deleteAtPoints :: [Int] -> S.Seq a -> S.Seq a
+deleteAtPoints points xs =
+  foldl' (\acc k -> S.deleteAt k acc) xs points
+
+{-
+-- For comparison with deleteAt. deleteAt is several
+-- times faster for long sequences.
+fakeDeleteAt :: Int -> S.Seq a -> S.Seq a
+fakeDeleteAt i xs
+  | 0 < i && i < S.length xs = case S.splitAt i xs of
+                               (before, after) -> before S.>< S.drop 1 after
+  | otherwise = xs
+-}
+
 -- splitAt+append: repeatedly cut the sequence at a random point
 -- and rejoin the pieces in the opposite order.
 -- Finally getting the middle element forces the whole spine.
index c24c7e8..e76cbc6 100644 (file)
@@ -75,6 +75,7 @@ main = defaultMain
        , testProperty "(!?)" prop_safeIndex
        , testProperty "adjust" prop_adjust
        , testProperty "insertAt" prop_insertAt
+       , testProperty "deleteAt" prop_deleteAt
        , testProperty "update" prop_update
        , testProperty "take" prop_take
        , testProperty "drop" prop_drop
@@ -496,6 +497,14 @@ prop_insertAt x xs =
       let res = insertAt i x xs
       in valid res .&&. res === case splitAt i xs of (front, back) -> front >< x <| back
 
+prop_deleteAt :: Seq A -> Property
+prop_deleteAt xs =
+  forAll (choose (-3, length xs + 3)) $ \i ->
+      let res = deleteAt i xs
+      in valid res .&&.
+          (((0 <= i && i < length xs) .&&. res === case splitAt i xs of (front, back) -> front >< drop 1 back)
+            .||. ((i < 0 || i >= length xs) .&&. res === xs))
+
 prop_adjust :: Int -> Int -> Seq Int -> Bool
 prop_adjust n i xs =
     toList' (adjust f i xs) ~= adjustList f i (toList xs)