From 967aa22330ce62d2ffe5f293946925f5c5eef0c7 Mon Sep 17 00:00:00 2001
From: krzys-h <krzys_h@interia.pl>
Date: Sun, 13 Mar 2016 20:17:10 +0100
Subject: [PATCH] Allow assigning to array in initialization from function
 (#624, #728)

---
 src/CBot/CBotInstr/CBotDefArray.cpp | 14 ++++++-
 test/unit/CBot/CBot_test.cpp        | 58 +++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/src/CBot/CBotInstr/CBotDefArray.cpp b/src/CBot/CBotInstr/CBotDefArray.cpp
index 6be430d4..2cb05aa0 100644
--- a/src/CBot/CBotInstr/CBotDefArray.cpp
+++ b/src/CBot/CBotInstr/CBotDefArray.cpp
@@ -23,6 +23,7 @@
 #include "CBot/CBotInstr/CBotExpression.h"
 #include "CBot/CBotInstr/CBotListArray.h"
 #include "CBot/CBotInstr/CBotEmpty.h"
+#include "CBot/CBotInstr/CBotTwoOpExpr.h"
 
 #include "CBot/CBotStack.h"
 #include "CBot/CBotCStack.h"
@@ -95,7 +96,18 @@ CBotInstr* CBotDefArray::Compile(CBotToken* &p, CBotCStack* pStack, CBotTypResul
 
         if (IsOfType(p, ID_ASS))                                        // with an assignment
         {
-            inst->m_listass = CBotListArray::Compile(p, pStk, type.GetTypElem());
+            if ((inst->m_listass = CBotTwoOpExpr::Compile(p, pStk)) != nullptr)
+            {
+                if (!pStk->GetTypResult().Compare(type))  // compatible type ?
+                {
+                    pStk->SetError(CBotErrBadType1, p->GetStart());
+                    goto error;
+                }
+            }
+            else
+            {
+                inst->m_listass = CBotListArray::Compile(p, pStk, type.GetTypElem());
+            }
         }
 
         if (pStk->IsOk()) return pStack->Return(inst, pStk);
diff --git a/test/unit/CBot/CBot_test.cpp b/test/unit/CBot/CBot_test.cpp
index b8cccd7b..10faf257 100644
--- a/test/unit/CBot/CBot_test.cpp
+++ b/test/unit/CBot/CBot_test.cpp
@@ -995,3 +995,61 @@ TEST_F(CBotUT, DISABLED_TestNANParam_Issue642)
         "}\n"
     );
 }
+
+TEST_F(CBotUT, TestArrayInitialization)
+{
+    ExecuteTest(
+        "extern void TestArrayInitialization() {\n"
+        "    int[] a = {1, 2, 3};\n"
+        "    ASSERT(sizeof(a) == 3);\n"
+        "    ASSERT(a[0] == 1);\n"
+        "    ASSERT(a[1] == 2);\n"
+        "    ASSERT(a[2] == 3);\n"
+        "}\n"
+    );
+
+    ExecuteTest(
+        "extern void TestArrayInitializationOutOfRange() {\n"
+        "    int a[2] = {1, 2, 3};\n"
+        "}\n",
+        CBotErrOutArray
+    );
+
+    ExecuteTest(
+        "extern void TestArrayInitializationSmallerThanRange() {\n"
+        "    int a[4] = {1, 2, 3};\n"
+        "    ASSERT(sizeof(a) == 3);\n"
+        "    ASSERT(a[0] == 1);\n"
+        "    ASSERT(a[1] == 2);\n"
+        "    ASSERT(a[2] == 3);\n"
+        "    a[3] = 4;\n"
+        "    ASSERT(sizeof(a) == 4);\n"
+        "    ASSERT(a[3] == 4);\n"
+        "}\n"
+    );
+
+    ExecuteTest(
+        "extern void TestArrayInitializationLimitUnchanged() {\n"
+        "    int a[4] = {1, 2, 3};\n"
+        "    a[4] = 5;\n"
+        "}\n",
+        CBotErrOutArray
+    );
+}
+
+TEST_F(CBotUT, TestArrayFunctionReturn)
+{
+    ExecuteTest(
+        "int[] test() {\n"
+        "    int[] a = {1, 2, 3};\n"
+        "    return a;"
+        "}\n"
+        "extern void TestArrayFunctionReturn() {\n"
+        "    int[] b = test();\n"
+        "    ASSERT(sizeof(b) == 3);\n"
+        "    ASSERT(b[0] == 1);\n"
+        "    ASSERT(b[1] == 2);\n"
+        "    ASSERT(b[2] == 3);\n"
+        "}\n"
+    );
+}
\ No newline at end of file