From 39d2cc084e25684095bc38e43740345ac6afb17c Mon Sep 17 00:00:00 2001
From: Supun Setunga <supun.setunga@gmail.com>
Date: Tue, 27 Feb 2024 11:19:37 +0530
Subject: [PATCH] Reject reference to optionals in the checker

---
 runtime/sema/check_reference_expression.go  | 14 +++++++-
 runtime/sema/errors.go                      | 36 +++++++++++++++++++++
 runtime/tests/checker/reference_test.go     | 19 ++++++++++-
 runtime/tests/interpreter/reference_test.go | 19 ++---------
 4 files changed, 69 insertions(+), 19 deletions(-)

diff --git a/runtime/sema/check_reference_expression.go b/runtime/sema/check_reference_expression.go
index e4f6dc3634..ecbef01961 100644
--- a/runtime/sema/check_reference_expression.go
+++ b/runtime/sema/check_reference_expression.go
@@ -112,12 +112,24 @@ func (checker *Checker) expectedTypeForReferencedExpr(
 
 	switch expectedType := expectedType.(type) {
 	case *OptionalType:
-		expectedLeftType, returnType, referenceType = checker.expectedTypeForReferencedExpr(expectedType.Type, hasPosition)
+		expectedLeftType, returnType, referenceType =
+			checker.expectedTypeForReferencedExpr(expectedType.Type, hasPosition)
+
 		// Re-wrap with an optional
 		expectedLeftType = &OptionalType{Type: expectedLeftType}
 		returnType = &OptionalType{Type: returnType}
 
 	case *ReferenceType:
+		referencedType := expectedType.Type
+		if referencedOptionalType, referenceToOptional := referencedType.(*OptionalType); referenceToOptional {
+			checker.report(
+				&ReferenceToAnOptionalError{
+					ReferencedOptionalType: referencedOptionalType,
+					Range:                  ast.NewRangeFromPositioned(checker.memoryGauge, hasPosition),
+				},
+			)
+		}
+
 		return expectedType.Type, expectedType, expectedType
 
 	default:
diff --git a/runtime/sema/errors.go b/runtime/sema/errors.go
index 17299b4870..6a4153be1b 100644
--- a/runtime/sema/errors.go
+++ b/runtime/sema/errors.go
@@ -2795,6 +2795,42 @@ func (e *NonReferenceTypeReferenceError) SecondaryError() string {
 	)
 }
 
+// ReferenceToAnOptionalError
+
+type ReferenceToAnOptionalError struct {
+	ReferencedOptionalType *OptionalType
+	ast.Range
+}
+
+var _ SemanticError = &ReferenceToAnOptionalError{}
+var _ errors.UserError = &ReferenceToAnOptionalError{}
+var _ errors.SecondaryError = &ReferenceToAnOptionalError{}
+
+func (*ReferenceToAnOptionalError) isSemanticError() {}
+
+func (*ReferenceToAnOptionalError) IsUserError() {}
+
+func (e *ReferenceToAnOptionalError) Error() string {
+	return "cannot create reference"
+}
+
+func (e *ReferenceToAnOptionalError) SecondaryError() string {
+	return fmt.Sprintf(
+		"expected non-optional type, got `%s`. Consider taking a reference with type `%s`",
+		e.ReferencedOptionalType.QualifiedString(),
+
+		// Suggest taking the optional out of the reference type.
+		NewOptionalType(
+			nil,
+			NewReferenceType(
+				nil,
+				UnauthorizedAccess,
+				e.ReferencedOptionalType.Type,
+			),
+		),
+	)
+}
+
 // InvalidResourceCreationError
 
 type InvalidResourceCreationError struct {
diff --git a/runtime/tests/checker/reference_test.go b/runtime/tests/checker/reference_test.go
index 112d0d7c8f..b06d1caee2 100644
--- a/runtime/tests/checker/reference_test.go
+++ b/runtime/tests/checker/reference_test.go
@@ -3823,6 +3823,22 @@ func TestCheckOptionalReference(t *testing.T) {
 		require.NoError(t, err)
 	})
 
+	t.Run("reference to optional", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := ParseAndCheck(t, `
+            fun main() {
+                var dict: {String: Foo} = {}
+                var ref: &(Foo?) = &dict["foo"] as &(Foo?)
+            }
+
+            struct Foo {}
+        `)
+
+		errs := RequireCheckerErrors(t, err, 1)
+		assert.IsType(t, &sema.ReferenceToAnOptionalError{}, errs[0])
+	})
+
 	t.Run("reference to nested optional", func(t *testing.T) {
 		t.Parallel()
 
@@ -3835,6 +3851,7 @@ func TestCheckOptionalReference(t *testing.T) {
             struct Foo {}
         `)
 
-		require.NoError(t, err)
+		errs := RequireCheckerErrors(t, err, 1)
+		assert.IsType(t, &sema.ReferenceToAnOptionalError{}, errs[0])
 	})
 }
diff --git a/runtime/tests/interpreter/reference_test.go b/runtime/tests/interpreter/reference_test.go
index 4b4d23128b..c3b423b4b2 100644
--- a/runtime/tests/interpreter/reference_test.go
+++ b/runtime/tests/interpreter/reference_test.go
@@ -3197,23 +3197,8 @@ func TestInterpretOptionalReference(t *testing.T) {
 		inter := parseCheckAndInterpret(t, `
             fun main() {
                 var dict: {String: Foo?} = {}
-                var ref: (&Foo)?? = &dict["foo"] as &Foo??
-            }
-
-            struct Foo {}
-        `)
-
-		_, err := inter.Invoke("main")
-		require.NoError(t, err)
-	})
-
-	t.Run("reference to nested optional", func(t *testing.T) {
-		t.Parallel()
-
-		inter := parseCheckAndInterpret(t, `
-            fun main() {
-                var dict: {String: Foo?} = {}
-                var ref: &(Foo??) = &dict["foo"] as &(Foo??)
+                var element: Foo?? = dict["foo"]
+                var ref: &(Foo)?? = &element as &(Foo)??
             }
 
             struct Foo {}