From fc6c5f85831a5b4f1d32527c0e89b88b691694a8 Mon Sep 17 00:00:00 2001
From: Heshan Padmasiri <hpheshan@gmail.com>
Date: Wed, 20 Mar 2024 08:51:33 +0530
Subject: [PATCH] Refactor optional field check

---
 .../wso2/ballerinalang/compiler/desugar/Desugar.java   |  5 ++---
 .../compiler/semantics/analyzer/SemanticAnalyzer.java  |  4 ++--
 .../ballerinalang/compiler/util/CompilerUtils.java     | 10 ++++------
 3 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/Desugar.java b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/Desugar.java
index f422200808d2..ae1c3a3da03d 100644
--- a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/Desugar.java
+++ b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/desugar/Desugar.java
@@ -322,7 +322,6 @@
 import java.util.Stack;
 import java.util.TreeMap;
 import java.util.stream.Collectors;
-
 import javax.xml.XMLConstants;
 
 import static io.ballerina.runtime.api.constants.RuntimeConstants.UNDERSCORE;
@@ -337,7 +336,7 @@
 import static org.wso2.ballerinalang.compiler.desugar.ASTBuilderUtil.createVariable;
 import static org.wso2.ballerinalang.compiler.desugar.ASTBuilderUtil.createVariableRef;
 import static org.wso2.ballerinalang.compiler.util.CompilerUtils.getMajorVersion;
-import static org.wso2.ballerinalang.compiler.util.CompilerUtils.isAssignmentToOptionalField;
+import static org.wso2.ballerinalang.compiler.util.CompilerUtils.isOptionalFieldAccess;
 import static org.wso2.ballerinalang.compiler.util.Names.GENERATED_INIT_SUFFIX;
 import static org.wso2.ballerinalang.compiler.util.Names.GEN_VAR_PREFIX;
 import static org.wso2.ballerinalang.compiler.util.Names.IGNORE;
@@ -2481,7 +2480,7 @@ public void visit(BLangAssignment assignNode) {
     }
 
     private static boolean shouldWidenExpressionTypeWithNil(BLangAssignment assignNode) {
-        if (!assignNode.expr.getBType().isNullable() || !isAssignmentToOptionalField(assignNode)) {
+        if (!assignNode.expr.getBType().isNullable() || !isOptionalFieldAccess(assignNode.varRef)) {
             return false;
         }
         // If we are assigning to an optional field we have a field based access on a record
diff --git a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/SemanticAnalyzer.java b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/SemanticAnalyzer.java
index 044b030646cf..adf7c0355433 100644
--- a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/SemanticAnalyzer.java
+++ b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/semantics/analyzer/SemanticAnalyzer.java
@@ -245,7 +245,7 @@
 import static org.ballerinalang.model.tree.NodeKind.RECORD_LITERAL_EXPR;
 import static org.ballerinalang.model.tree.NodeKind.REG_EXP_CAPTURING_GROUP;
 import static org.ballerinalang.model.tree.NodeKind.REG_EXP_CHARACTER_CLASS;
-import static org.wso2.ballerinalang.compiler.util.CompilerUtils.isAssignmentToOptionalField;
+import static org.wso2.ballerinalang.compiler.util.CompilerUtils.isOptionalFieldAccess;
 
 /**
  * @since 0.94
@@ -2322,7 +2322,7 @@ public void visit(BLangAssignment assignNode, AnalyzerData data) {
 
         BLangExpression expr = assignNode.expr;
         expr.usedForOptionalFieldValue =
-                data.expType != symTable.semanticError && isAssignmentToOptionalField(assignNode);
+                data.expType != symTable.semanticError && isOptionalFieldAccess(assignNode.varRef);
         data.typeChecker.checkExpr(expr, data.env, data.expType, data.prevEnvs, data.commonAnalyzerData);
 
         validateWorkerAnnAttachments(assignNode.expr, data);
diff --git a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/util/CompilerUtils.java b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/util/CompilerUtils.java
index 74fce1a321b6..cd0601298fc6 100644
--- a/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/util/CompilerUtils.java
+++ b/compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/util/CompilerUtils.java
@@ -26,9 +26,8 @@
 import org.wso2.ballerinalang.compiler.semantics.model.types.BRecordType;
 import org.wso2.ballerinalang.compiler.semantics.model.types.BType;
 import org.wso2.ballerinalang.compiler.tree.BLangFunction;
-import org.wso2.ballerinalang.compiler.tree.BLangNode;
+import org.wso2.ballerinalang.compiler.tree.expressions.BLangExpression;
 import org.wso2.ballerinalang.compiler.tree.expressions.BLangFieldBasedAccess;
-import org.wso2.ballerinalang.compiler.tree.statements.BLangAssignment;
 
 import static org.wso2.ballerinalang.compiler.util.Constants.MAIN_FUNCTION_NAME;
 
@@ -77,12 +76,11 @@ public static String getPackageIDStringWithMajorVersion(PackageID packageID) {
         return org + packageID.name + Names.VERSION_SEPARATOR.value + getMajorVersion(packageID.version.value);
     }
 
-    public static boolean isAssignmentToOptionalField(BLangAssignment assignNode) {
-        BLangNode varRef = assignNode.varRef;
-        if (varRef.getKind() != NodeKind.FIELD_BASED_ACCESS_EXPR) {
+    public static boolean isOptionalFieldAccess(BLangExpression exprNode) {
+        if (exprNode.getKind() != NodeKind.FIELD_BASED_ACCESS_EXPR) {
             return false;
         }
-        BLangFieldBasedAccess fieldAccessNode = (BLangFieldBasedAccess) varRef;
+        BLangFieldBasedAccess fieldAccessNode = (BLangFieldBasedAccess) exprNode;
         BType targetType = Types.getImpliedType(fieldAccessNode.expr.getBType());
         if (targetType.tag != TypeTags.RECORD) {
             return false;