@@ -18,6 +18,17 @@ type internal WhileKind =
1818 /// The item under test is always excluded
1919| Exclusive
2020
21+ [<Struct>]
22+ type internal TakeOrSkipKind =
23+ /// use the Seq.take semantics, raises exception if not enough elements
24+ | Take
25+ /// use the Seq.skip semantics, raises exception if not enough elements
26+ | Skip
27+ /// use the Seq.truncate semantics, safe operation, returns all if count exceeds the seq
28+ | Truncate
29+ /// no Seq equiv, but like Stream.drop in Scala: safe operation, return empty if not enough elements
30+ | Drop
31+
2132[<Struct>]
2233type internal Action < 'T , 'U , 'TaskU when 'TaskU :> Task < 'U >> =
2334 | CountableAction of countable_action : ( int -> 'T -> 'U )
@@ -52,19 +63,16 @@ module internal TaskSeqInternal =
5263 nullArg argName
5364
5465 let inline raiseEmptySeq () =
55- ArgumentException( " The asynchronous input sequence was empty." , " source" )
56- |> raise
66+ invalidArg " source" " The input task sequence was empty."
5767
58- let inline raiseCannotBeNegative ( name : string ) =
59- ArgumentException( " The value cannot be negative" , name)
60- |> raise
68+ let inline raiseCannotBeNegative name =
69+ invalidArg name " The value must be non-negative"
6170
6271 let inline raiseInsufficient () =
63- ArgumentException( " The asynchronous input sequence was has an insufficient number of elements." , " source" )
64- |> raise
72+ invalidArg " source" " The input task sequence was has an insufficient number of elements."
6573
6674 let inline raiseNotFound () =
67- KeyNotFoundException( " The predicate function or index did not satisfy any item in the async sequence." )
75+ KeyNotFoundException( " The predicate function or index did not satisfy any item in the task sequence." )
6876 |> raise
6977
7078 let isEmpty ( source : TaskSeq < _ >) =
@@ -76,6 +84,16 @@ module internal TaskSeqInternal =
7684 return not step
7785 }
7886
87+ let empty < 'T > =
88+ { new IAsyncEnumerable< 'T> with
89+ member _.GetAsyncEnumerator ( _ ) =
90+ { new IAsyncEnumerator< 'T> with
91+ member _.MoveNextAsync () = ValueTask.False
92+ member _.Current = Unchecked.defaultof< 'T>
93+ member _.DisposeAsync () = ValueTask.CompletedTask
94+ }
95+ }
96+
7997 let singleton ( value : 'T ) =
8098 { new IAsyncEnumerable< 'T> with
8199 member _.GetAsyncEnumerator ( _ ) =
@@ -613,9 +631,96 @@ module internal TaskSeqInternal =
613631 | false -> ()
614632 }
615633
616- let takeWhile whileKind predicate ( source : TaskSeq < _ >) =
634+
635+ let skipOrTake skipOrTake count ( source : TaskSeq < _ >) =
617636 checkNonNull ( nameof source) source
637+ if count < 0 then raiseCannotBeNegative ( nameof count)
638+
639+ match skipOrTake with
640+ | Skip ->
641+ // don't create a new sequence if count = 0
642+ if count = 0 then source
643+ else
644+ taskSeq {
645+ use e = source.GetAsyncEnumerator CancellationToken.None
646+
647+ for _ in 1 .. count do
648+ let! step = e.MoveNextAsync()
649+ if not step then
650+ raiseInsufficient()
651+
652+ let mutable cont = true
653+
654+ while cont do
655+ yield e.Current
656+ let! moveNext = e.MoveNextAsync()
657+ cont <- moveNext
658+
659+ }
660+ | Drop ->
661+ // don't create a new sequence if count = 0
662+ if count = 0 then source
663+ else
664+ taskSeq {
665+ use e = source.GetAsyncEnumerator CancellationToken.None
666+
667+ let! step = e.MoveNextAsync()
668+ let mutable cont = step
669+ let mutable pos = 0
670+
671+ // skip, or stop looping if we reached the end
672+ while cont do
673+ pos <- pos + 1
674+ let! moveNext = e.MoveNextAsync()
675+ cont <- moveNext && pos <= count
676+
677+ // return the rest
678+ while cont do
679+ yield e.Current
680+ let! moveNext = e.MoveNextAsync()
681+ cont <- moveNext
682+
683+ }
684+ | Take ->
685+ // don't initialize an empty task sequence
686+ if count = 0 then empty
687+ else
688+ taskSeq {
689+ use e = source.GetAsyncEnumerator CancellationToken.None
690+
691+ for _ in count .. - 1 .. 1 do
692+ let! step = e.MoveNextAsync()
693+ if not step then
694+ raiseInsufficient()
618695
696+ yield e.Current
697+ }
698+
699+ | Truncate ->
700+ // don't create a new sequence if count = 0
701+ if count = 0 then empty
702+ else
703+ taskSeq {
704+ use e = source.GetAsyncEnumerator CancellationToken.None
705+
706+ let! step = e.MoveNextAsync()
707+ let mutable cont = step
708+ let mutable pos = 0
709+
710+ // return items until we've exhausted the seq
711+ // report this line, weird error:
712+ //while! e.MoveNextAsync() && pos < 1 do
713+ while cont do
714+ yield e.Current
715+ pos <- pos + 1
716+ let! moveNext = e.MoveNextAsync()
717+ cont <- moveNext && pos <= count
718+
719+ }
720+
721+ let takeWhile whileKind predicate ( source : TaskSeq < _ >) =
722+ checkNonNull ( nameof source) source
723+
619724 taskSeq {
620725 use e = source.GetAsyncEnumerator CancellationToken.None
621726 let! step = e.MoveNextAsync()
0 commit comments