@@ -18,6 +18,17 @@ type internal WhileKind =
18
18
/// The item under test is always excluded
19
19
| Exclusive
20
20
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
+
21
32
[<Struct>]
22
33
type internal Action < 'T , 'U , 'TaskU when 'TaskU :> Task < 'U >> =
23
34
| CountableAction of countable_action : ( int -> 'T -> 'U )
@@ -51,20 +62,15 @@ module internal TaskSeqInternal =
51
62
if isNull arg then
52
63
nullArg argName
53
64
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."
57
66
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"
61
68
62
69
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."
65
71
66
72
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." )
68
74
|> raise
69
75
70
76
let isEmpty ( source : TaskSeq < _ >) =
@@ -76,6 +82,16 @@ module internal TaskSeqInternal =
76
82
return not step
77
83
}
78
84
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
+
79
95
let singleton ( value : 'T ) =
80
96
{ new IAsyncEnumerable< 'T> with
81
97
member _.GetAsyncEnumerator ( _ ) =
@@ -613,6 +629,101 @@ module internal TaskSeqInternal =
613
629
| false -> ()
614
630
}
615
631
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
+
616
727
let takeWhile whileKind predicate ( source : TaskSeq < _ >) =
617
728
checkNonNull ( nameof source) source
618
729
0 commit comments