@@ -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 )
@@ -51,20 +62,15 @@ module internal TaskSeqInternal =
5162 if isNull arg then
5263 nullArg argName
5364
54- let inline raiseEmptySeq () =
55- ArgumentException( " The asynchronous input sequence was empty." , " source" )
56- |> raise
65+ let inline raiseEmptySeq () = invalidArg " source" " The input task sequence was empty."
5766
58- let inline raiseCannotBeNegative ( name : string ) =
59- ArgumentException( " The value cannot be negative" , name)
60- |> raise
67+ let inline raiseCannotBeNegative name = invalidArg name " The value must be non-negative"
6168
6269 let inline raiseInsufficient () =
63- ArgumentException( " The asynchronous input sequence was has an insufficient number of elements." , " source" )
64- |> raise
70+ invalidArg " source" " The input task sequence was has an insufficient number of elements."
6571
6672 let inline raiseNotFound () =
67- KeyNotFoundException( " The predicate function or index did not satisfy any item in the async sequence." )
73+ KeyNotFoundException( " The predicate function or index did not satisfy any item in the task sequence." )
6874 |> raise
6975
7076 let isEmpty ( source : TaskSeq < _ >) =
@@ -76,6 +82,16 @@ module internal TaskSeqInternal =
7682 return not step
7783 }
7884
85+ let empty < 'T > =
86+ { new IAsyncEnumerable< 'T> with
87+ member _.GetAsyncEnumerator ( _ ) =
88+ { new IAsyncEnumerator< 'T> with
89+ member _.MoveNextAsync () = ValueTask.False
90+ member _.Current = Unchecked.defaultof< 'T>
91+ member _.DisposeAsync () = ValueTask.CompletedTask
92+ }
93+ }
94+
7995 let singleton ( value : 'T ) =
8096 { new IAsyncEnumerable< 'T> with
8197 member _.GetAsyncEnumerator ( _ ) =
@@ -613,6 +629,101 @@ module internal TaskSeqInternal =
613629 | false -> ()
614630 }
615631
632+
633+ let skipOrTake skipOrTake count ( source : TaskSeq < _ >) =
634+ checkNonNull ( nameof source) source
635+
636+ if count < 0 then
637+ 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
643+ source
644+ else
645+ taskSeq {
646+ use e = source.GetAsyncEnumerator CancellationToken.None
647+
648+ for _ in 1 .. count do
649+ let! step = e.MoveNextAsync()
650+
651+ if not step then
652+ raiseInsufficient ()
653+
654+ let mutable cont = true
655+
656+ while cont do
657+ yield e.Current
658+ let! moveNext = e.MoveNextAsync()
659+ cont <- moveNext
660+
661+ }
662+ | Drop ->
663+ // don't create a new sequence if count = 0
664+ if count = 0 then
665+ source
666+ else
667+ taskSeq {
668+ use e = source.GetAsyncEnumerator CancellationToken.None
669+
670+ let! step = e.MoveNextAsync()
671+ let mutable cont = step
672+ let mutable pos = 0
673+
674+ // skip, or stop looping if we reached the end
675+ while cont do
676+ pos <- pos + 1
677+ let! moveNext = e.MoveNextAsync()
678+ cont <- moveNext && pos <= count
679+
680+ // return the rest
681+ while cont do
682+ yield e.Current
683+ let! moveNext = e.MoveNextAsync()
684+ cont <- moveNext
685+
686+ }
687+ | Take ->
688+ // don't initialize an empty task sequence
689+ if count = 0 then
690+ empty
691+ else
692+ taskSeq {
693+ use e = source.GetAsyncEnumerator CancellationToken.None
694+
695+ for _ in count .. - 1 .. 1 do
696+ let! step = e.MoveNextAsync()
697+
698+ if not step then
699+ raiseInsufficient ()
700+
701+ yield e.Current
702+ }
703+
704+ | Truncate ->
705+ // don't create a new sequence if count = 0
706+ if count = 0 then
707+ empty
708+ else
709+ taskSeq {
710+ use e = source.GetAsyncEnumerator CancellationToken.None
711+
712+ let! step = e.MoveNextAsync()
713+ let mutable cont = step
714+ let mutable pos = 0
715+
716+ // return items until we've exhausted the seq
717+ // report this line, weird error:
718+ //while! e.MoveNextAsync() && pos < 1 do
719+ while cont do
720+ yield e.Current
721+ pos <- pos + 1
722+ let! moveNext = e.MoveNextAsync()
723+ cont <- moveNext && pos <= count
724+
725+ }
726+
616727 let takeWhile whileKind predicate ( source : TaskSeq < _ >) =
617728 checkNonNull ( nameof source) source
618729
0 commit comments