Add traverseWithIndex to Data.Sequence

David Feuer david.feuer at gmail.com
Tue Jan 19 16:55:17 UTC 2016


Indexed traversals seem to be pretty popular these days, and sequences
should be able to support them efficiently. At present, the best
options are

1. Thread a counter through. This doesn't work well for weird things
like lazy State.

2. Use zipWith first to add indices, then traverse. This adds a
smaller, but still asymptotic, penalty to the same sorts of unusual
functors, and also adds constant-factor overhead to strict ones that
the counter threading wouldn't.

I propose the following, modified mechanically from mapWithIndex. It
should be about as efficient as possible in all cases.

traverseWithIndex :: Applicative f => (Int -> a -> f b) -> Seq a -> f (Seq b)
traverseWithIndex f' (Seq xs') = Seq <$> traverseWithIndexTree (\s
(Elem a) -> Elem <$> f' s a) 0 xs'
 where
  traverseWithIndexTree :: (Applicative f, Sized a) => (Int -> a -> f
b) -> Int -> FingerTree a -> f (FingerTree b)
  traverseWithIndexTree _ s Empty = s `seq` pure Empty
  traverseWithIndexTree f s (Single xs) = Single <$> f s xs
  traverseWithIndexTree f s (Deep n pr m sf) = sPspr `seq` sPsprm `seq`
          Deep n <$>
               traverseWithIndexDigit f s pr <*>
               traverseWithIndexTree (traverseWithIndexNode f) sPspr m <*>
               traverseWithIndexDigit f sPsprm sf
    where
      sPspr = s + size pr
      sPsprm = s + n - size sf

  traverseWithIndexDigit :: (Applicative f, Sized a) => (Int -> a -> f
b) -> Int -> Digit a -> f (Digit b)
  traverseWithIndexDigit f s (One a) = One <$> f s a
  traverseWithIndexDigit f s (Two a b) = sPsa `seq` Two <$> f s a <*> f sPsa b
    where
      sPsa = s + size a
  traverseWithIndexDigit f s (Three a b c) = sPsa `seq` sPsab `seq`
                                      Three <$> f s a <*> f sPsa b <*> f sPsab c
    where
      sPsa = s + size a
      sPsab = sPsa + size b
  traverseWithIndexDigit f s (Four a b c d) = sPsa `seq` sPsab `seq`
sPsabc `seq`
                          Four <$> f s a <*> f sPsa b <*> f sPsab c
<*> f sPsabc d
    where
      sPsa = s + size a
      sPsab = sPsa + size b
      sPsabc = sPsab + size c

  traverseWithIndexNode :: (Applicative f, Sized a) => (Int -> a -> f
b) -> Int -> Node a -> f (Node b)
  traverseWithIndexNode f s (Node2 ns a b) = sPsa `seq` Node2 ns <$> f
s a <*> f sPsa b
    where
      sPsa = s + size a
  traverseWithIndexNode f s (Node3 ns a b c) = sPsa `seq` sPsab `seq`
                                     Node3 ns <$> f s a <*> f sPsa b
<*> f sPsab c
    where
      sPsa = s + size a
      sPsab = sPsa + size b


More information about the Libraries mailing list