Generic.takeWhile is not copy-free

[gksato]

This issue is closely related to #182, but it is a contract failure, so I guess we need to do something about this.

The description of this issue is simple: the document of the function Data.Vector.Generic.takeWhile says:

O(n) Yield the longest prefix of elements satisfying the predicate without copying.

However, the function's implementation is:

takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
{-# INLINE takeWhile #-}
takeWhile f = unstream . Bundle.takeWhile f . stream

(See it on Hackage, or on GitHub)

The document says the function is copy-free, but it is obvious from the code that it requires copy when:

1. it is used against a vector actually living in the heap, and
2. the produced vector can't fuse away (for example, it is used more than once).

Note that Data.Vector.Generic.dropWhile also had this problem, and that we resolved it on PR #327 by making dropWhile actually copy-free.
On PR #327, we made it possible by letting dropWhile be fusible only in the case the argument vector to the function is already an unstreamed stream.

An obvious solution to this problem is to remove the phrase "without copying" from the documentation. It is completely sensible to choose that way.
The problem is more complex than that of dropWhile, and we cannot utilize the method same as the one used in #327.
I guess I've come up with a solution that makes takeWhile copy-free while preserving all required stream fusion, but it is rather global and complicated, and might let bugs sneak in.

I'm being lazy and I couldn't write up everything at once. I'll explain the difficulty of this problem and the proposed solution making takeWhile copy-free in subsequent comments.

See Also: #182 #327(+#141)

[Shimuuar]

Great catch!

The problem is more complex than that of dropWhile

But what the problem? Both functions are basically same. In case of no fusion: find index of first element satisfying predicate and then slice underlying vector. Ony difference is instead of unsafeDrop one'll have to use unsafeSlice

[gksato]

Think of:

import qualified Data.Vector.Unboxed as VU

xs :: VU.Vector Int
xs = [some stray vector actually living on heap]
f :: Int -> Bool
f = ...
g :: Int -> Word
g = ...
h :: Word -> Bool
h = ...
n :: Int
n = ...

val :: VU.Vector Word
val = VU.take n $ VU.filter h $ VU.map g $ VU.takeWhile f xs

Let:

t0 = fromMaybe (VU.length xs) $ VU.findIndex (not . f) xs
t1 = maybe (VU.length xs) (+1) $  VU.findIndices (h . g) xs VU.!? (n-1)

Then, the time complexity of the computation of val is O(t0) if VU.takeWhile f xs is computed with findIndex and unsafeTake, but O(min{t0,t1}) if it is processed via stream.

[gksato]

Therefore, takeWhile should be processed with stream fusion if the argument vector is generated by stream fusion, and if the resulting vector is consumed by stream fusion. That's the point.

[gksato]

My proposal is to utilize the field sVector in Bundle, currently unused, and (somehow, by redefining or by a new rule) have a equation

Data.Vector.Generic.new (Data.Vector.Generic.New.unstream Bundle{ sVector = Just v } = v

Of course, we redefine Data.Vector.Fusion.Bundle.takeWhile (+ Data.Vector.Fusion.Bundle.dropWhile, Data.Vector.Fusion.Bundle.take, etc) to set sVector.

[gksato]

Specifically, all the slice-taking functions must set sVector to be Just an actual slice whenever original sVector is Just something, and the other functions must set sVector to be Nothing.

[lehins]

Data.Vector.Generic.new (Data.Vector.Generic.New.unstream Bundle{ sVector = Just v } = v

Wouldn't that be nice. Unfortunately a rule like that is simply not possible. The problem is that Just is at a value level, while rewrite rules work at type level.

However, what I think might be possible is defining different unstreamSlice function that would get their own set of rewrite rules that would prioritize the sVector vs sElems. I am quite worried about relying on the currently unused sVector though. Someone needs to try it out and see what comes out of it.

[gksato]

@lehins Thanks for analysis.

Unfortunately a rule like that is simply not possible.

I see. I had a hunch that it's the case; it looked somewhat weird since it had pattern match in LHS 😅

However, what I think might be possible is defining different unstreamSlicefunction that would get their own set of rewrite rules that would prioritize the sVector vs sElems.

I don't really get it. Brutally roughly speaking, I was thinking of:

-- in Data.Vector.Generic.New

{-# INLINE_INNER clone' #-} -- Very creepy we have to have INLINE_INNER here.
clone' v = Data.Vector.Generic.clone v

{-# INLINE_FUSED unstream #-}
unstream (B.Bundle {sVector = Just v}) = clone' v
unstream B.Bundle{..} = [as always]

{-# RULES "new/clone'" Data.Vector.Generic.new (clone' v) = v #-} -- Of course it goes to Data.Vector.Generic

How different is your solution from mine? Where does the function unstreamSlice plays its role?

I am quite worried about relying on the currently unused sVector though. Someone needs to try it out and see what comes out of it.

I can't agree more. With my proposal, the fix would be a big deal for a single function anyhow we implement it. Let me point out that we should make sure we remove "without copying" from the documentation if we are not able to make takeWhile copy-free before the next release (when will it be?).

[gksato]

Uh, I've forgotten to say. I'd be happy to implement it. That is, I'm a laaaaazy person and it might take a month or two, so I'd be happy if you'd wait for me. This would be a big change, so I needed to discuss possible implementations of my vague idea.

[gksato]

{-# RULES "new/clone'" Data.Vector.Generic.new (clone' v) = v #-}

Oops. That doesn't work.... Data.Vector.Generic.new is INLINE_FUSED.

[lehins]

Here is what I had in mind is at first:

-- Data.Vector.Generic

unstreamSlice :: Vector v a => Bundle v a -> v a
{-# INLINE unstream #-}
unstreamSlice (B.Bundle {sVector = Just v}) = v
unstreamSlice s = unstream s

followed by addition of more rewrite rules, but ... now I think all these complications might be unnecessary.

We might get away with a simpler solution. We can keep unstream as before, except we will use the sVector if one is available in all the functions that can benefit from slicing instead of streaming, like takeWhile for example:

--- Data.Vector.Generic 
unstream :: Vector v a => Bundle v a -> v a
{-# INLINE unstream #-}
unstream (MBundle.Bundle {MBundle.sVector = Just v}) = v
unstream s = new (New.unstream s)

--- Data.Vector.Generic.Base (or some other module with all functions that don't rely on streaming)
sliceTakeWhile :: Vector v a => (a -> Bool) -> v a -> v a
{-# INLINE takeWhile #-}
sliceTakeWhile f v = take (go 0) v
  where
    k = length v
    go i
      | i < k && f (unsafeIndex v i) = go (i + 1)
      | otherwise = i

-- Data.Vector.Fusion.Bundle.Monadic
takeWhile :: Monad m => (a -> Bool) -> Bundle m v a -> Bundle m v a
{-# INLINE takeWhile #-}
takeWhile f b@(B.Bundle {sVector = mv}) = (takeWhileM (return . f) b) { sVector = fmap (sliceTakeWhile f) mv) }

I've tried this idea out here lehins@1cab907 and it seems to work, but I did not investigate it any deeper. If you are up for a challenge and would like to work on it that would be fantastic. Here are some things I think would be important to have answers to if you do decide to pursue this:

• Benchmarks before and after for at least few functions from different groups (some folds, mapping, indexing, etc.)
• Some sort of analysis that this approach does not mess up existing streaming functionality. A good solution for this could be partially ready in A test suite for fusion #229
• and of course some more tests for affected functions that validate correctness in cases when functions like takeWhile are fused with a streamed or with a sliced vector.

That is, I'm a laaaaazy person and it might take a month or two, so I'd be happy if you'd wait for me.

There is no rush. I don't have much time on my hands to work on this and I don't see anyone else eager on getting this done, so we can definitely wait for you.

Let me point out that we should make sure we remove "without copying" from the documentation if we are not able to make takeWhile copy-free before the next release (when will it be?).

We haven't discussed the timeline for the next release, but I imagine we could have one before the end of the year. So I think if you won't get it done prior to the next release we can certainly remove "without copying" before making the actual release.

[gksato]

After looking at your solution, I was thinking of utilizing sVector for a while, and now I'm at a loss what to do.
Since Data.Vector.Generic.stream, Data.Vector.Generic.new, and Data.Vector.Generic.New.unstream has INLINE_FUSED or INLINE [1], all the stream-new-unstream reductions must be performed before Phase 1 (within Phase 2 or before). Therefore every compile-time Data.Vector.Generic.unstream pattern-match must also be performed before Phase 1, so all the functions that emit pure Bundle must have INLINE (Note that sVector = Nothing must be visible to the optimizer). I guess a lot of Bundle-manipulating functions have INLINE_FUSED for a reason, and I don't think it is a good idea to replace all INLINE_FUSED with INLINE. What do you think?

[gksato]

I think I now clarified the situation a lot. I seemingly found something very inconvenient -- we can't make takeWhile copy-free safely.

For this comment, I consider any method, including the ones utilizing sVector, satisfying the following constraints:

1. takeWhile takes actual slice if and only if all functions in the stream fusion pipeline it belongs to is slice-taking functions and otherwise takeWhile yields to actual stream fusion,
2. all functions and types that are not in explicitly Internal modules are external and included in package contract (i.e. they only change according to PVP), and
3. the inline phase of unstream, new, New.unstream, stream stays as it is.

I claim that we have an only one solution satisfying the conditions above, and in the solution we rewrite New.unstream's definition with

-- in Data.Vector.Generic.New
unstream :: Vector v a => Bundle v a -> New v a
{-# INLINE_FUSED unstream #-}
unstream Bundle {sVector = Just v} = New (Data.Vector.Generic.unsafeThaw v)
unstream s = s `seq` New (MVector.vunstream s)

and redefine all slice-taking Bundle functions so that they use sVector field.

"Sketch of Proof"(I haven't proven this mathematically; I just explain):
Let's determine when the optimization rewriting of "unstream (... slice-taking functions ... (stream v)...) -> (actual slice-taking function) v" (which we call bundles-to-actual-slice reduction) happens. Before we perform a bundles-to-actual-slice reduction, we need to make sure all the Bundle-tweaking functions in the pipeline only takes slice. In order to confirm that, the pipeline should be fully concatenated, i.e. we should have run all the stream-new-unstream reductions. Now we have shown bundles-to-actual-slice reduction must be performed in Optimization Phase 1 or 0, and Generic.unstream must be inlined within Phase 2.
In Phase 1 and after, the definitions of new, New.unstream, and unstream have been or get inlined, no rewrite rules concerning these functions are possible. We should do the things in the definitions of these functions. Since by Constraint 2 we do not tweak the type of new, the definition of new looks impossible to tweak. Therefore New.unstream should emit New $ unsafeThaw (the resulting actual slice), and the resulting actual slice must be retrievable from the argument Bundle. Since the only information of type Vector in the datatype Bundle is sVector, it looks like the only solution. "QED".

Actually, this "solution" replaces new (New.unstream Bundle{sVector = Just v}) with copy-free unsafeThaw v >>= unsafeFreeze.

However, this is CREEPY. This is a highly unsafe solution. Since we regard New.unstream and Generic.stream as externally utilizable function,

Generic.new (New.modify f (New.unstream (Generic.stream v)))

is completely legal and modifies the immutable vector v. That is, there seems to be no safe solution satisfying Conditions 1,2 and 3.

[lehins]

I am not terribly convinced that this approach is the right one, stuff like unsafeThaw v >>= unsafeFreeze makes me feel a bit uncomfortable. On the other hand it might be totally fine. I think we'd have to see an implementation first before trying to argue about its correctness and safety.

I see now how my suggested change to Generic.unstream can break fusion in general, so here is the approach I see feasible and safe which is what I had mind to begin with when I was talking about unstreamSlice:

takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
{-# INLINE takeWhile #-}
takeWhile f = unstreamSlice . Bundle.takeWhile f . stream

-- | /O(n)/ Construct a vector from a 'Bundle', but unlike `unstream` prioritize getting a
-- potentially sliced vector, while falling back on construction from the stream.
unstreamSlice :: Vector v a => Bundle v a -> v a
{-# INLINE_FUSED unstreamSlice #-}
unstreamSlice (MBundle.Bundle {MBundle.sVector = Just v}) = v
unstreamSlice s = new (New.unstream s)

{-# RULES
"stream/unstreamSlice [Vector]" forall s.
  stream (unstreamSlice s) = s
  #-}

Do you think this approach will not work and if so what sort of problems do you think we can run into going this route?

The benefits I see is that:

• it is simple we don't need to rely on any unsafe functions for this to work
• only functions like take, takeWhile, drop, etc. that can rely on O(1) slicing will use unstreamSlice
• they will still fuse with other functions, albeit potentially not as good, but GHC is really good at optimizing Maybe away, which means that in practice it might work just as well as before.
• All three points are satisfied from your previous message

This approach would satisfy the no-copy part of documentation, because when the input is a stream it would either fuse if or materialize it, but if sVector is available a slice is taken. Some examples:

• takeWhile even $ runST $ new 4 >>= \mv -> write mv 0 2 >> write mv 1 3 >> unsafeFreeze mv - would materialize [2,3], but takeWhile would do a constant slice and output [2]
• takeWhile even $ fromList [2,2,2,3,3] - would materialize only [2,2,2], but there was no copy since the input was a stream
• takeWhile (< 6) $ takeWhile even $ fromList [2,4,6,8,10,11] - would only materialize [2,4]
• takeWhile (< 6) $ force $ takeWhile even $ fromList [2,4,6,8,10,11] - would materialize [2,4,6,8,10], but then take a slice [2,4]
• head $ takeWhile even $ fromList [2,2,2,3,3] - could maybe fuse completely?

All that being said, I do not have 100% confidence in either solution since we don't have any of these: actual implementation, corresponding benchmarks or tests which check that things fuse properly.

So whichever approach you choose to implement it's up to you. If it works in the end and doesn't cause any problems I am sure it will get merged. What I'd recommend is collecting a whole bunch of examples (like I listed above), define their expected behavior and try them out on some rough implementation see if it even works on individual cases. Math and proofs are nice, but without hard code and a running executable they are just theory. ;)

[gksato]

Huh... I didn't think of a verson of Generic.unstream with INLINE [1]. It looks nice. I now feel relieved, since, as I said in my previous comment, I was worried about having to give up the hope of a copy-less fusible implementation of takeWhile. Alrighty. I'm gonna try implementing it.

Also, thanks for providing a list of possible tests! In the previous comment, I just needed math because I wanted to reject some buggy implementations: I wanted to detect some bugs that may or may not be invoked depending on GHC's preference of optimization order (and I thought I ruled everything out and I was sad 😅). It's tests that'll finally matter, actually.