From d4bd3016dac94c051c4998d23fafc8a840960e5a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Marjam=C3=A4ki?= <daniel.marjamaki@gmail.com>
Date: Sun, 28 Jun 2020 17:28:40 +0200
Subject: [PATCH] ExprEngine; Improved handling of for loop, loop variable

---
 lib/exprengine.cpp      | 12 +++++++++++-
 lib/exprengine.h        |  4 +++-
 test/testexprengine.cpp |  5 ++---
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/lib/exprengine.cpp b/lib/exprengine.cpp
index dc31b8ad8..fced6286c 100644
--- a/lib/exprengine.cpp
+++ b/lib/exprengine.cpp
@@ -1629,7 +1629,16 @@ static void assignExprValue(const Token *expr, ExprEngine::ValuePtr value, Data
                     arrayValue->assign(ExprEngine::ValuePtr(), value);
             } else {
                 auto indexValue = calculateArrayIndex(expr, data, *arrayValue);
-                arrayValue->assign(indexValue, value);
+                bool loopAssign = false;
+                if (auto loopValue = std::dynamic_pointer_cast<ExprEngine::IntRange>(indexValue)) {
+                    if (loopValue->loopScope == expr->scope()) {
+                        loopAssign = true;
+                        for (auto i = loopValue->minValue; i <= loopValue->maxValue; ++i)
+                            arrayValue->assign(std::make_shared<ExprEngine::IntRange>(ExprEngine::str(i), i, i), value);
+                    }
+                }
+                if (!loopAssign)
+                    arrayValue->assign(indexValue, value);
             }
         }
     } else if (expr->isUnaryOp("*")) {
@@ -2355,6 +2364,7 @@ static std::string execute(const Token *start, const Token *end, Data &data)
                 auto loopValues = std::make_shared<ExprEngine::IntRange>(data.getNewSymbolName(), initValue, lastValue);
                 data.assignValue(tok, varid, loopValues);
                 tok = tok->linkAt(1);
+                loopValues->loopScope = tok->next()->scope();
                 continue;
             }
         }
diff --git a/lib/exprengine.h b/lib/exprengine.h
index 52574a8ee..cff8d536f 100644
--- a/lib/exprengine.h
+++ b/lib/exprengine.h
@@ -143,7 +143,8 @@ namespace ExprEngine {
         IntRange(const std::string &name, int128_t minValue, int128_t maxValue)
             : Value(name, ValueType::IntRange)
             , minValue(minValue)
-            , maxValue(maxValue) {
+            , maxValue(maxValue)
+            , loopScope(nullptr) {
         }
         std::string getRange() const OVERRIDE {
             if (minValue == maxValue)
@@ -156,6 +157,7 @@ namespace ExprEngine {
 
         int128_t minValue;
         int128_t maxValue;
+        const Scope *loopScope;
     };
 
     class FloatRange : public Value {
diff --git a/test/testexprengine.cpp b/test/testexprengine.cpp
index b34f38e0a..333cf9dfe 100644
--- a/test/testexprengine.cpp
+++ b/test/testexprengine.cpp
@@ -490,9 +490,8 @@ private:
                             "  for (int i = 0; i < 10; i++) x[i] = 0;\n"
                             "  x[4] == 67;\n"
                             "}";
-        ASSERT_EQUALS("(and (>= $3 (- 2147483648)) (<= $3 2147483647))\n"
-                      "(= $3 67)\n"
-                      "z3::sat\n",
+        ASSERT_EQUALS("(= 0 67)\n"
+                      "z3::unsat\n",
                       expr(code, "=="));
     }