Add null-checks for each function that takes an IAsyncEnumerable or otherwise nullable type

Author: abelbraaksma

src/FSharp.Control.TaskSeq/TaskSeq.fs

@@ -30,6 +30,7 @@ module TaskSeq =
     //
 
     let toList (source: taskSeq<'T>) = [
+        Internal.checkNonNull (nameof source) source
         let e = source.GetAsyncEnumerator(CancellationToken())
 
         try
@@ -40,6 +41,7 @@ module TaskSeq =
     ]
 
     let toArray (source: taskSeq<'T>) = [|
+        Internal.checkNonNull (nameof source) source
         let e = source.GetAsyncEnumerator(CancellationToken())
 
         try
@@ -50,6 +52,7 @@ module TaskSeq =
     |]
 
     let toSeq (source: taskSeq<'T>) = seq {
+        Internal.checkNonNull (nameof source) source
         let e = source.GetAsyncEnumerator(CancellationToken())
 
         try
@@ -74,6 +77,8 @@ module TaskSeq =
     //
 
     let ofArray (source: 'T[]) = taskSeq {
+        Internal.checkNonNull (nameof source) source
+
         for c in source do
             yield c
     }
@@ -84,16 +89,22 @@ module TaskSeq =
     }
 
     let ofSeq (source: 'T seq) = taskSeq {
+        Internal.checkNonNull (nameof source) source
+
         for c in source do
             yield c
     }
 
     let ofResizeArray (source: 'T ResizeArray) = taskSeq {
+        Internal.checkNonNull (nameof source) source
+
         for c in source do
             yield c
     }
 
     let ofTaskSeq (source: #Task<'T> seq) = taskSeq {
+        Internal.checkNonNull (nameof source) source
+
         for c in source do
             let! c = c
             yield c
@@ -106,12 +117,16 @@ module TaskSeq =
     }
 
     let ofTaskArray (source: #Task<'T> array) = taskSeq {
+        Internal.checkNonNull (nameof source) source
+
         for c in source do
             let! c = c
             yield c
     }
 
     let ofAsyncSeq (source: Async<'T> seq) = taskSeq {
+        Internal.checkNonNull (nameof source) source
+
         for c in source do
             let! c = task { return! c }
             yield c
@@ -124,6 +139,8 @@ module TaskSeq =
     }
 
     let ofAsyncArray (source: Async<'T> array) = taskSeq {
+        Internal.checkNonNull (nameof source) source
+
         for c in source do
             let! c = Async.toTask c
             yield c
@@ -148,21 +165,29 @@ module TaskSeq =
         }
 
     let concat (sources: taskSeq<#taskSeq<'T>>) = taskSeq {
+        Internal.checkNonNull (nameof sources) sources
+
         for ts in sources do
             yield! (ts :> taskSeq<'T>)
     }
 
     let append (source1: #taskSeq<'T>) (source2: #taskSeq<'T>) = taskSeq {
+        Internal.checkNonNull (nameof source1) source1
+        Internal.checkNonNull (nameof source2) source2
         yield! (source1 :> IAsyncEnumerable<'T>)
         yield! (source2 :> IAsyncEnumerable<'T>)
     }
 
     let appendSeq (source1: #taskSeq<'T>) (source2: #seq<'T>) = taskSeq {
+        Internal.checkNonNull (nameof source1) source1
+        Internal.checkNonNull (nameof source2) source2
         yield! (source1 :> IAsyncEnumerable<'T>)
         yield! (source2 :> seq<'T>)
     }
 
     let prependSeq (source1: #seq<'T>) (source2: #taskSeq<'T>) = taskSeq {
+        Internal.checkNonNull (nameof source1) source1
+        Internal.checkNonNull (nameof source2) source2
         yield! (source1 :> seq<'T>)
         yield! (source2 :> IAsyncEnumerable<'T>)
     }
@@ -242,6 +267,7 @@ module TaskSeq =
     }
 
     let indexed (source: taskSeq<'T>) = taskSeq {
+        Internal.checkNonNull (nameof source) source
         let mutable i = 0
 
         for x in source do

src/FSharp.Control.TaskSeq/TaskSeqInternal.fs

@@ -39,6 +39,11 @@ type internal InitAction<'T, 'TaskT when 'TaskT :> Task<'T>> =
     | InitActionAsync of async_init_item: (int -> 'TaskT)
 
 module internal TaskSeqInternal =
+    /// Raise an NRE for arguments that are null. Only used for 'source' parameters, never for function parameters.
+    let inline checkNonNull argName arg =
+        if isNull arg then
+            nullArg argName
+
     let inline raiseEmptySeq () =
         ArgumentException("The asynchronous input sequence was empty.", "source")
         |> raise
@@ -56,6 +61,7 @@ module internal TaskSeqInternal =
         |> raise
 
     let isEmpty (source: taskSeq<_>) = task {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
         let! step = e.MoveNextAsync()
         return not step
@@ -88,6 +94,7 @@ module internal TaskSeqInternal =
 
     /// Returns length unconditionally, or based on a predicate
     let lengthBy predicate (source: taskSeq<_>) = task {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
         let mutable go = true
         let mutable i = 0
@@ -123,6 +130,7 @@ module internal TaskSeqInternal =
 
     /// Returns length unconditionally, or based on a predicate
     let lengthBeforeMax max (source: taskSeq<_>) = task {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
         let mutable go = true
         let mutable i = 0
@@ -138,6 +146,7 @@ module internal TaskSeqInternal =
     }
 
     let tryExactlyOne (source: taskSeq<_>) = task {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
 
         match! e.MoveNextAsync() with
@@ -194,6 +203,7 @@ module internal TaskSeqInternal =
     }
 
     let iter action (source: taskSeq<_>) = task {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
         let mutable go = true
         let! step = e.MoveNextAsync()
@@ -235,6 +245,7 @@ module internal TaskSeqInternal =
     }
 
     let fold folder initial (source: taskSeq<_>) = task {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
         let mutable go = true
         let mutable result = initial
@@ -259,44 +270,49 @@ module internal TaskSeqInternal =
     }
 
     let toResizeArrayAsync source = task {
+        checkNonNull (nameof source) source
         let res = ResizeArray()
         do! source |> iter (SimpleAction(fun item -> res.Add item))
         return res
     }
 
     let toResizeArrayAndMapAsync mapper source = (toResizeArrayAsync >> Task.map mapper) source
 
-    let map mapper (taskSequence: taskSeq<_>) =
+    let map mapper (source: taskSeq<_>) =
+        checkNonNull (nameof source) source
+
         match mapper with
         | CountableAction mapper -> taskSeq {
             let mutable i = 0
 
-            for c in taskSequence do
+            for c in source do
                 yield mapper i c
                 i <- i + 1
           }
 
         | SimpleAction mapper -> taskSeq {
-            for c in taskSequence do
+            for c in source do
                 yield mapper c
           }
 
         | AsyncCountableAction mapper -> taskSeq {
             let mutable i = 0
 
-            for c in taskSequence do
+            for c in source do
                 let! result = mapper i c
                 yield result
                 i <- i + 1
           }
 
         | AsyncSimpleAction mapper -> taskSeq {
-            for c in taskSequence do
+            for c in source do
                 let! result = mapper c
                 yield result
           }
 
     let zip (source1: taskSeq<_>) (source2: taskSeq<_>) = taskSeq {
+        checkNonNull (nameof source1) source1
+        checkNonNull (nameof source2) source2
         use e1 = source1.GetAsyncEnumerator(CancellationToken())
         use e2 = source2.GetAsyncEnumerator(CancellationToken())
         let mutable go = true
@@ -312,28 +328,37 @@ module internal TaskSeqInternal =
     }
 
     let collect (binder: _ -> #IAsyncEnumerable<_>) (source: taskSeq<_>) = taskSeq {
+        checkNonNull (nameof source) source
+
         for c in source do
             yield! binder c :> IAsyncEnumerable<_>
     }
 
     let collectSeq (binder: _ -> #seq<_>) (source: taskSeq<_>) = taskSeq {
+        checkNonNull (nameof source) source
+
         for c in source do
             yield! binder c :> seq<_>
     }
 
     let collectAsync (binder: _ -> #Task<#IAsyncEnumerable<_>>) (source: taskSeq<_>) = taskSeq {
+        checkNonNull (nameof source) source
+
         for c in source do
             let! result = binder c
             yield! result :> IAsyncEnumerable<_>
     }
 
     let collectSeqAsync (binder: _ -> #Task<#seq<_>>) (source: taskSeq<_>) = taskSeq {
+        checkNonNull (nameof source) source
+
         for c in source do
             let! result = binder c
             yield! result :> seq<_>
     }
 
     let tryLast (source: taskSeq<_>) = task {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
         let mutable go = true
         let mutable last = ValueNone
@@ -351,6 +376,7 @@ module internal TaskSeqInternal =
     }
 
     let tryHead (source: taskSeq<_>) = task {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
 
         match! e.MoveNextAsync() with
@@ -359,6 +385,7 @@ module internal TaskSeqInternal =
     }
 
     let tryTail (source: taskSeq<_>) = task {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
 
         match! e.MoveNextAsync() with
@@ -379,6 +406,8 @@ module internal TaskSeqInternal =
     }
 
     let tryItem index (source: taskSeq<_>) = task {
+        checkNonNull (nameof source) source
+
         if index < 0 then
             // while the loop below wouldn't run anyway, we don't want to call MoveNext in this case
             // to prevent side effects hitting unnecessarily
@@ -404,6 +433,7 @@ module internal TaskSeqInternal =
     }
 
     let tryPick chooser (source: taskSeq<_>) = task {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
 
         let mutable go = true
@@ -436,6 +466,7 @@ module internal TaskSeqInternal =
     }
 
     let tryFind predicate (source: taskSeq<_>) = task {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
 
         let mutable go = true
@@ -472,6 +503,7 @@ module internal TaskSeqInternal =
     }
 
     let tryFindIndex predicate (source: taskSeq<_>) = task {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
 
         let mutable go = true
@@ -504,6 +536,8 @@ module internal TaskSeqInternal =
     }
 
     let choose chooser (source: taskSeq<_>) = taskSeq {
+        checkNonNull (nameof source) source
+
         match chooser with
         | TryPick picker ->
             for item in source do
@@ -519,6 +553,8 @@ module internal TaskSeqInternal =
     }
 
     let filter predicate (source: taskSeq<_>) = taskSeq {
+        checkNonNull (nameof source) source
+
         match predicate with
         | Predicate predicate ->
             for item in source do
@@ -585,6 +621,7 @@ module internal TaskSeqInternal =
                 ValueTask.CompletedTask
 
     let except itemsToExclude (source: taskSeq<_>) = taskSeq {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
         let mutable go = true
         let! step = e.MoveNextAsync()
@@ -609,6 +646,7 @@ module internal TaskSeqInternal =
     }
 
     let exceptOfSeq itemsToExclude (source: taskSeq<_>) = taskSeq {
+        checkNonNull (nameof source) source
         use e = source.GetAsyncEnumerator(CancellationToken())
         let mutable go = true
         let! step = e.MoveNextAsync()