From 072e822020f88f8888cf4ae460c0846890d346c8 Mon Sep 17 00:00:00 2001
From: chrchr-github <78114321+chrchr-github@users.noreply.github.com>
Date: Thu, 2 Mar 2023 21:01:20 +0100
Subject: [PATCH] Set ValueType for free function returning iterator (#4837)

---
 lib/symboldatabase.cpp      | 22 ++++++++++++++++++-
 test/testsymboldatabase.cpp | 42 +++++++++++++++++++++++++++++--------
 2 files changed, 54 insertions(+), 10 deletions(-)

diff --git a/lib/symboldatabase.cpp b/lib/symboldatabase.cpp
index a124f4f4f..46e2251e3 100644
--- a/lib/symboldatabase.cpp
+++ b/lib/symboldatabase.cpp
@@ -7045,7 +7045,8 @@ void SymbolDatabase::setValueTypeInTokenList(bool reportDebugWarnings, Token *to
                     }
                 }
 
-                if (typestr.empty() || typestr == "iterator") {
+                const bool isReturnIter = typestr == "iterator";
+                if (typestr.empty() || isReturnIter) {
                     if (Token::simpleMatch(tok->astOperand1(), ".") &&
                         tok->astOperand1()->astOperand1() &&
                         tok->astOperand1()->astOperand2() &&
@@ -7063,6 +7064,25 @@ void SymbolDatabase::setValueTypeInTokenList(bool reportDebugWarnings, Token *to
                                 vt.containerTypeToken =
                                     tok->astOperand1()->astOperand1()->valueType()->containerTypeToken;
                                 setValueType(tok, vt);
+                                continue;
+                            }
+                        }
+                    }
+                    if (isReturnIter) {
+                        const std::vector<const Token*> args = getArguments(tok);
+                        if (!args.empty()) {
+                            const Library::ArgumentChecks::IteratorInfo* info = mSettings->library.getArgIteratorInfo(tok->previous(), 1);
+                            if (info && info->it) {
+                                const Token* contTok = args[0];
+                                if (Token::simpleMatch(args[0]->astOperand1(), ".") && args[0]->astOperand1()->astOperand1())
+                                    contTok = args[0]->astOperand1()->astOperand1();
+                                if (contTok && contTok->variable() && contTok->variable()->valueType() && contTok->variable()->valueType()->container) {
+                                    ValueType vt;
+                                    vt.type = ValueType::Type::ITERATOR;
+                                    vt.container = contTok->variable()->valueType()->container;
+                                    vt.containerTypeToken = contTok->variable()->valueType()->containerTypeToken;
+                                    setValueType(tok, vt);
+                                }
                             }
                         }
                     }
diff --git a/test/testsymboldatabase.cpp b/test/testsymboldatabase.cpp
index 805a587bd..8fdf1f34e 100644
--- a/test/testsymboldatabase.cpp
+++ b/test/testsymboldatabase.cpp
@@ -8111,16 +8111,40 @@ private:
     }
 
     void valueType3() {
-        GET_SYMBOL_DB("void f(std::vector<std::unordered_map<int, std::unordered_set<int>>>& v, int i, int j) {\n"
-                      "    auto& s = v[i][j];\n"
-                      "    s.insert(0);\n"
-                      "}\n");
-        ASSERT_EQUALS("", errout.str());
+        {
+            GET_SYMBOL_DB("void f(std::vector<std::unordered_map<int, std::unordered_set<int>>>& v, int i, int j) {\n"
+                          "    auto& s = v[i][j];\n"
+                          "    s.insert(0);\n"
+                          "}\n");
+            ASSERT_EQUALS("", errout.str());
 
-        const Token* tok = tokenizer.tokens();
-        tok = Token::findsimplematch(tok, "s .");
-        ASSERT(tok && tok->valueType());
-        ASSERT_EQUALS("container(std :: set|unordered_set <)", tok->valueType()->str());
+            const Token* tok = tokenizer.tokens();
+            tok = Token::findsimplematch(tok, "s .");
+            ASSERT(tok && tok->valueType());
+            ASSERT_EQUALS("container(std :: set|unordered_set <)", tok->valueType()->str());
+        }
+        {
+            GET_SYMBOL_DB("void f(std::vector<int> v) {\n"
+                          "    auto it = std::find(v.begin(), v.end(), 0);\n"
+                          "}\n");
+            ASSERT_EQUALS("", errout.str());
+
+            const Token* tok = tokenizer.tokens();
+            tok = Token::findsimplematch(tok, "auto");
+            ASSERT(tok && tok->valueType());
+            ASSERT_EQUALS("iterator(std :: vector <)", tok->valueType()->str());
+        }
+        {
+            GET_SYMBOL_DB("void f(std::vector<int>::iterator beg, std::vector<int>::iterator end) {\n"
+                          "    auto it = std::find(beg, end, 0);\n"
+                          "}\n");
+            ASSERT_EQUALS("", errout.str());
+
+            const Token* tok = tokenizer.tokens();
+            tok = Token::findsimplematch(tok, "auto");
+            ASSERT(tok && tok->valueType());
+            ASSERT_EQUALS("iterator(std :: vector <)", tok->valueType()->str());
+        }
     }
 
     void valueTypeThis() {