From f9bfae76cabbba453cc343717bde0f75de9ab504 Mon Sep 17 00:00:00 2001
From: Soma Simon <soma.simon@huawei.com>
Date: Thu, 9 May 2024 15:52:16 +0200
Subject: [PATCH] Fix segmentation fault on optional trailing lambda

Union type was not handled at lambda type inference.

Fixes internal issue #16858

Change-Id: If135ac929673baea2f62e084de655421e6348fe4
Signed-off-by: Soma Simon <soma.simon@huawei.com>
---
 ets2panda/checker/ETSchecker.h                |  1 +
 ets2panda/checker/ets/typeCheckingHelpers.cpp | 21 +++++++++-
 .../runtime/ets/optional_trailing_lambda.ets  | 39 +++++++++++++++++++
 3 files changed, 59 insertions(+), 2 deletions(-)
 create mode 100644 ets2panda/test/runtime/ets/optional_trailing_lambda.ets

diff --git a/ets2panda/checker/ETSchecker.h b/ets2panda/checker/ETSchecker.h
index c512dfb658..092499702c 100644
--- a/ets2panda/checker/ETSchecker.h
+++ b/ets2panda/checker/ETSchecker.h
@@ -297,6 +297,7 @@ public:
     bool NeedTypeInference(const ir::ScriptFunction *lambda);
     std::vector<bool> FindTypeInferenceArguments(const ArenaVector<ir::Expression *> &arguments);
     void InferTypesForLambda(ir::ScriptFunction *lambda, ir::ETSFunctionType *calleeType);
+    void HandleLambdaTypeInfer(ir::AstNode *typeAnn, ir::ScriptFunction *const lambda);
     bool TypeInference(Signature *signature, const ArenaVector<ir::Expression *> &arguments,
                        TypeRelationFlag flags = TypeRelationFlag::NONE);
     bool CheckLambdaAssignable(ir::Expression *param, ir::ScriptFunction *lambda);
diff --git a/ets2panda/checker/ets/typeCheckingHelpers.cpp b/ets2panda/checker/ets/typeCheckingHelpers.cpp
index ca5b3f41ca..52d8b95f6a 100644
--- a/ets2panda/checker/ets/typeCheckingHelpers.cpp
+++ b/ets2panda/checker/ets/typeCheckingHelpers.cpp
@@ -931,8 +931,8 @@ bool ETSChecker::TypeInference(Signature *signature, const ArenaVector<ir::Expre
             typeAnn = DerefETSTypeReference(typeAnn);
         }
 
-        ASSERT(typeAnn->IsETSFunctionType());
-        InferTypesForLambda(lambda, typeAnn->AsETSFunctionType());
+        HandleLambdaTypeInfer(typeAnn, lambda);
+
         Type *const argumentType = GetApparentType(arrowFuncExpr->Check(this));
         Type *const parameterType = GetApparentType(signature->Params()[index]->TsType());
         const Type *targetType = TryGettingFunctionTypeFromInvokeFunction(parameterType);
@@ -947,6 +947,23 @@ bool ETSChecker::TypeInference(Signature *signature, const ArenaVector<ir::Expre
     return invocable;
 }
 
+void ETSChecker::HandleLambdaTypeInfer(ir::AstNode *typeAnn, ir::ScriptFunction *const lambda)
+{
+    ASSERT(typeAnn->IsETSFunctionType() || typeAnn->IsETSUnionType());
+
+    if (typeAnn->IsETSFunctionType()) {
+        InferTypesForLambda(lambda, typeAnn->AsETSFunctionType());
+        return;
+    }
+
+    for (auto type : typeAnn->AsETSUnionType()->Types()) {
+        if (type->IsETSFunctionType()) {
+            InferTypesForLambda(lambda, type->AsETSFunctionType());
+            return;
+        }
+    }
+}
+
 bool ETSChecker::ExtensionETSFunctionType(checker::Type *type)
 {
     if (!type->IsETSFunctionType()) {
diff --git a/ets2panda/test/runtime/ets/optional_trailing_lambda.ets b/ets2panda/test/runtime/ets/optional_trailing_lambda.ets
new file mode 100644
index 0000000000..361f40cd87
--- /dev/null
+++ b/ets2panda/test/runtime/ets/optional_trailing_lambda.ets
@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) 2024 Huawei Device Co., Ltd.
+ * 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.
+ */
+
+let a = 0;
+
+function foo(exec?: () => void): void {
+    exec?.();
+    a += 1;
+}
+
+function bar() {
+    foo() {
+        a += 1;
+        assert a == 1;
+        foo() {
+            a += 1;
+            assert a == 2;
+            foo();
+        }
+    }
+}
+
+export function main() {
+    a = 0
+    bar();
+    assert a == 5;
+}
-- 
Gitee