Fixed crash in slang-ir-autodiff-loop-analysis.cpp (#6831)

* Added Dictionary::erase(iterator) and fixed crashing when filtering a dictionary in slang-ir-autodiff-loop-analysis.cpp

* Added Dictionary::removeIf(Predicate)

* Removed Dictionary::erase(it)

---------

Co-authored-by: Julius Ikkala <julius.ikkala@gmail.com>

2 files changed, 21 insertions, 5 deletions

diff --git a/source/core/slang-dictionary.h b/source/core/slang-dictionary.h
index 639978a08..d35255a8c 100644
--- a/source/core/slang-dictionary.h
+++ b/source/core/slang-dictionary.h
@@ -148,6 +148,25 @@ public:
     // Erases the value at the specified key if it exists
     void remove(const TKey& key) { map.erase(key); }
 
+    // Removes all values satifying the predicate:
+    // bool predicate(pair<Key, Value>)
+    template<typename Predicate>
+    void removeIf(Predicate&& predicate)
+    {
+        auto it = begin();
+        while (it != end())
+        {
+            if (predicate(*it))
+            {
+                it = map.erase(it);
+            }
+            else
+            {
+                ++it;
+            }
+        }
+    }
+
     // Reserves enough space for the specified number of values
     void reserve(Index size) { map.reserve(std::size_t(size)); };
 
diff --git a/source/slang/slang-ir-autodiff-loop-analysis.cpp b/source/slang/slang-ir-autodiff-loop-analysis.cpp
index d4ff631a6..137ab7775 100644
--- a/source/slang/slang-ir-autodiff-loop-analysis.cpp
+++ b/source/slang/slang-ir-autodiff-loop-analysis.cpp
@@ -252,11 +252,8 @@ void StatementSet::disjunct(StatementSet other)
     // Remove any insts that don't have a corresponding statement in the other set,
     // since this effectively means "any".
     //
-    for (auto& statement : statements)
-    {
-        if (!other.statements.containsKey(statement.first))
-            statements.remove(statement.first);
-    }
+    statements.removeIf([&](auto const& statement)
+                        { return !other.statements.containsKey(statement.first); });
 }
 
 void StatementSet::conjunct(StatementSet other)