From ba14292918d814eeaea4de62da2ad0daae92f8b0 Mon Sep 17 00:00:00 2001
From: Yuxin Wang <yuxinw@uber.com>
Date: Thu, 10 Oct 2024 16:24:15 -0400
Subject: [PATCH] Add initial handling for iterators (#288)

NilAway is panicking on code that ranges over iterators. This PR adds
initial handling for such that we assume the results from the iterators
are always nonnil, without panicking.

See #287
---
 assertion/function/assertiontree/backprop.go  | 26 ++++++++++
 nilaway_go123_plus_test.go                    | 52 +++++++++++++++++++
 .../looprangego123/looprangego123.go          | 33 ++++++++++++
 3 files changed, 111 insertions(+)
 create mode 100644 nilaway_go123_plus_test.go
 create mode 100644 testdata/src/go.uber.org/looprangego123/looprangego123.go

diff --git a/assertion/function/assertiontree/backprop.go b/assertion/function/assertiontree/backprop.go
index deea995d..ae0eebd0 100644
--- a/assertion/function/assertiontree/backprop.go
+++ b/assertion/function/assertiontree/backprop.go
@@ -410,8 +410,34 @@ func backpropAcrossRange(rootNode *RootAssertionNode, lhs []ast.Expr, rhs ast.Ex
 		}
 	}
 
+	// produceNonNil marks the ith lhs expression as nonnil due to limitations of NilAway.
+	produceNonNil := func(i int) {
+		if !util.IsEmptyExpr(lhs[i]) {
+			rootNode.AddProduction(&annotation.ProduceTrigger{
+				Annotation: &annotation.ProduceTriggerNever{},
+				Expr:       lhs[i],
+			})
+		}
+	}
+
 	rhsType := types.Unalias(rootNode.Pass().TypesInfo.Types[rhs].Type)
 
+	// Go 1.23 introduced the `iter` package, which provides a way to iterate over sequences
+	// in a generic way. The `iter.Seq` and `iter.Seq2` types are used to represent sequences
+	// and are used in the `range` statement. We currently do not handle these types yet, so
+	// here we assume that they are deeply non-nil (by adding nonnil producers).
+	// TODO: handle that (#287).
+	if named, ok := rhsType.(*types.Named); ok && named.Obj() != nil && named.Obj().Pkg() != nil && named.Obj().Pkg().Path() == "iter" {
+		if named.Obj().Name() == "Seq" {
+			produceNonNil(0)
+			return nil
+		} else if named.Obj().Name() == "Seq2" {
+			produceNonNil(0)
+			produceNonNil(1)
+			return nil
+		}
+	}
+
 	// This block breaks down the cases for the `range` statement being analyzed,
 	// starting by switching on how many left-hand operands there are
 	switch len(lhs) {
diff --git a/nilaway_go123_plus_test.go b/nilaway_go123_plus_test.go
new file mode 100644
index 00000000..405835d1
--- /dev/null
+++ b/nilaway_go123_plus_test.go
@@ -0,0 +1,52 @@
+//  Copyright (c) 2023 Uber Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build go1.23
+
+// This file is meant for testing features in Go 1.23 and beyond.
+// TODO: Migrate these test cases in the mainstream test files once NilAway starts to support Go 1.23 is a base version.
+
+package nilaway
+
+import (
+	"testing"
+
+	"golang.org/x/tools/go/analysis/analysistest"
+)
+
+func TestNilAway_Go123(t *testing.T) {
+	t.Parallel()
+
+	testdata := analysistest.TestData()
+
+	// For descriptions of the purpose of each of the following tests, consult their source files
+	// located in testdata/src/<package>.
+
+	tests := []struct {
+		name     string
+		patterns []string
+	}{
+		{name: "LoopRangeGo123", patterns: []string{"go.uber.org/looprangego123"}},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			t.Logf("Running test for packages %s", tt.patterns)
+
+			analysistest.Run(t, testdata, Analyzer, tt.patterns...)
+		})
+	}
+}
diff --git a/testdata/src/go.uber.org/looprangego123/looprangego123.go b/testdata/src/go.uber.org/looprangego123/looprangego123.go
new file mode 100644
index 00000000..740cb11c
--- /dev/null
+++ b/testdata/src/go.uber.org/looprangego123/looprangego123.go
@@ -0,0 +1,33 @@
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// This package aims to test nilability behavior for `range` in loops.
+
+// <nilaway no inference>
+package looprange
+
+import (
+	"maps"
+	"slices"
+)
+
+func testIter() {
+	i := 42
+	for element := range slices.Values([]*int{&i, &i, nil}) {
+		print(*element) // FN: we do not really handle iterators for now, the elements from iterators are assumed to be nonnil.
+	}
+	for k, v := range maps.All(map[string]*int{"abc": &i, "def": nil}) {
+		print(k)
+		print(*v) // FN: we do not really handle iterators for now, the elements from iterators are assumed to be nonnil.
+	}
+}