[C++] Implement configurable PredictionContextMergeCache

runtime/Cpp/runtime/src/LexerInterpreter.h

@@ -7,6 +7,7 @@
 
 #include "Lexer.h"
 #include "atn/PredictionContext.h"
+#include "atn/PredictionContextCache.h"
 #include "Vocabulary.h"
 
 namespace antlr4 {

runtime/Cpp/runtime/src/ParserInterpreter.h

@@ -9,6 +9,7 @@
 #include "atn/ATN.h"
 #include "support/BitSet.h"
 #include "atn/PredictionContext.h"
+#include "atn/PredictionContextCache.h"
 #include "Vocabulary.h"
 
 namespace antlr4 {

runtime/Cpp/runtime/src/antlr4-runtime.h

@@ -95,12 +95,16 @@
 #include "atn/OrderedATNConfigSet.h"
 #include "atn/ParseInfo.h"
 #include "atn/ParserATNSimulator.h"
+#include "atn/ParserATNSimulatorOptions.h"
 #include "atn/PlusBlockStartState.h"
 #include "atn/PlusLoopbackState.h"
 #include "atn/PrecedencePredicateTransition.h"
 #include "atn/PredicateEvalInfo.h"
 #include "atn/PredicateTransition.h"
 #include "atn/PredictionContext.h"
+#include "atn/PredictionContextCache.h"
+#include "atn/PredictionContextMergeCache.h"
+#include "atn/PredictionContextMergeCacheOptions.h"
 #include "atn/PredictionMode.h"
 #include "atn/ProfilingATNSimulator.h"
 #include "atn/RangeTransition.h"

runtime/Cpp/runtime/src/atn/ATNSimulator.h

@@ -7,6 +7,7 @@
 
 #include "atn/ATN.h"
 #include "atn/PredictionContext.h"
+#include "atn/PredictionContextCache.h"
 #include "misc/IntervalSet.h"
 #include "support/CPPUtils.h"
 

runtime/Cpp/runtime/src/atn/ParserATNSimulator.cpp

@@ -53,7 +53,13 @@ ParserATNSimulator::ParserATNSimulator(const ATN &atn, std::vector<dfa::DFA> &de
 
 ParserATNSimulator::ParserATNSimulator(Parser *parser, const ATN &atn, std::vector<dfa::DFA> &decisionToDFA,
                                        PredictionContextCache &sharedContextCache)
-: ATNSimulator(atn, sharedContextCache), decisionToDFA(decisionToDFA), parser(parser) {
+: ParserATNSimulator(parser, atn, decisionToDFA, sharedContextCache, ParserATNSimulatorOptions()) {}
+
+ParserATNSimulator::ParserATNSimulator(Parser *parser, const ATN &atn, std::vector<dfa::DFA> &decisionToDFA,
+                                       PredictionContextCache &sharedContextCache,
+                                       const ParserATNSimulatorOptions &options)
+: ATNSimulator(atn, sharedContextCache), decisionToDFA(decisionToDFA), parser(parser),
+  mergeCache(options.getPredictionContextMergeCacheOptions()) {
   InitializeInstanceFields();
 }
 
@@ -87,7 +93,12 @@ size_t ParserATNSimulator::adaptivePredict(TokenStream *input, size_t decision,
   // Now we are certain to have a specific decision's DFA
   // But, do we still need an initial state?
   auto onExit = finally([this, input, index, m] {
-    mergeCache.clear(); // wack cache after each prediction
+    if (mergeCache.getOptions().getClearEveryN() != 0) {
+      if (++_mergeCacheCounter == mergeCache.getOptions().getClearEveryN()) {
+        mergeCache.clear();
+        _mergeCacheCounter = 0;
+      }
+    }
     _dfa = nullptr;
     input->seek(index);
     input->release(m);
@@ -108,10 +119,9 @@ size_t ParserATNSimulator::adaptivePredict(TokenStream *input, size_t decision,
   }
 
   if (s0 == nullptr) {
-    bool fullCtx = false;
-    std::unique_ptr<ATNConfigSet> s0_closure = computeStartState(dfa.atnStartState,
-                                                                 &ParserRuleContext::EMPTY, fullCtx);
-
+    auto s0_closure = computeStartState(dfa.atnStartState, &ParserRuleContext::EMPTY, false);
+    std::unique_ptr<dfa::DFAState> newState;
+    std::unique_ptr<dfa::DFAState> oldState;
     std::unique_lock<std::shared_mutex> stateLock(atn._stateMutex);
     dfa::DFAState* ds0 = dfa.s0;
     if (dfa.isPrecedenceDfa()) {
@@ -122,25 +132,20 @@ size_t ParserATNSimulator::adaptivePredict(TokenStream *input, size_t decision,
        * than simply setting DFA.s0.
        */
       ds0->configs = std::move(s0_closure); // not used for prediction but useful to know start configs anyway
-      dfa::DFAState *newState = new dfa::DFAState(applyPrecedenceFilter(ds0->configs.get())); /* mem-check: managed by the DFA or deleted below */
-      s0 = addDFAState(dfa, newState);
+      newState = std::make_unique<dfa::DFAState>(applyPrecedenceFilter(ds0->configs.get()));
+      s0 = addDFAState(dfa, newState.get());
       std::unique_lock<std::shared_mutex> edgeLock(atn._edgeMutex);
       dfa.setPrecedenceStartState(parser->getPrecedence(), s0);
-      if (s0 != newState) {
-        delete newState; // If there was already a state with this config set we don't need the new one.
-      }
     } else {
-      dfa::DFAState *newState = new dfa::DFAState(std::move(s0_closure)); /* mem-check: managed by the DFA or deleted below */
-      s0 = addDFAState(dfa, newState);
-
+      newState = std::make_unique<dfa::DFAState>(std::move(s0_closure));
+      s0 = addDFAState(dfa, newState.get());
       if (ds0 != s0) {
-        delete ds0; // Delete existing s0 DFA state, if there's any.
-        ds0 = nullptr;
+        oldState.reset(ds0);
         dfa.s0 = s0;
       }
-      if (s0 != newState) {
-        delete newState; // If there was already a state with this config set we don't need the new one.
-      }
+    }
+    if (s0 == newState.get()) {
+      newState.release();
     }
   }
 
@@ -458,7 +463,7 @@ std::unique_ptr<ATNConfigSet> ParserATNSimulator::computeReachSet(ATNConfigSet *
 
   // First figure out where we can reach on input t
   for (const auto &c : closure_->configs) {
-    if (c->state != nullptr && c->state->getStateType() == ATNStateType::RULE_STOP) {
+    if (RuleStopState::is(c->state)) {
       assert(c->context->isEmpty());
 
       if (fullCtx || t == Token::EOF) {

runtime/Cpp/runtime/src/atn/ParserATNSimulator.h

@@ -9,6 +9,8 @@
 #include "dfa/DFAState.h"
 #include "atn/ATNSimulator.h"
 #include "atn/PredictionContext.h"
+#include "atn/PredictionContextMergeCache.h"
+#include "atn/ParserATNSimulatorOptions.h"
 #include "SemanticContext.h"
 #include "atn/ATNConfig.h"
 
@@ -251,6 +253,10 @@ namespace atn {
     ParserATNSimulator(Parser *parser, const ATN &atn, std::vector<dfa::DFA> &decisionToDFA,
                        PredictionContextCache &sharedContextCache);
 
+    ParserATNSimulator(Parser *parser, const ATN &atn, std::vector<dfa::DFA> &decisionToDFA,
+                       PredictionContextCache &sharedContextCache,
+                       const ParserATNSimulatorOptions &options);
+
     virtual void reset() override;
     virtual void clearDFA() override;
     virtual size_t adaptivePredict(TokenStream *input, size_t decision, ParserRuleContext *outerContext);
@@ -380,6 +386,7 @@ namespace atn {
     /// also be examined during cache lookup.
     /// </summary>
     PredictionContextMergeCache mergeCache;
+    size_t _mergeCacheCounter = 0;
 
     // LAME globals to avoid parameters!!!!! I need these down deep in predTransition
     TokenStream *_input;

runtime/Cpp/runtime/src/atn/ParserATNSimulatorOptions.h

@@ -0,0 +1,50 @@
+// Copyright 2012-2022 The ANTLR Project
+//
+// Redistribution and use in source and binary forms, with or without modification, are permitted
+// provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this list of conditions
+//    and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice, this list of
+//    conditions and the following disclaimer in the documentation and/or other materials provided
+//    with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its contributors may be used to
+//    endorse or promote products derived from this software without specific prior written
+//    permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+// IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+// FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
+// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
+// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY
+// WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#pragma once
+
+#include "atn/PredictionContextMergeCacheOptions.h"
+
+namespace antlr4 {
+namespace atn {
+
+  class ANTLR4CPP_PUBLIC ParserATNSimulatorOptions final {
+  public:
+    ParserATNSimulatorOptions& setPredictionContextMergeCacheOptions(
+        PredictionContextMergeCacheOptions predictionContextMergeCacheOptions) {
+      _predictionContextMergeCacheOptions = std::move(predictionContextMergeCacheOptions);
+      return *this;
+    }
+
+    const PredictionContextMergeCacheOptions& getPredictionContextMergeCacheOptions() const {
+      return _predictionContextMergeCacheOptions;
+    }
+
+  private:
+    PredictionContextMergeCacheOptions _predictionContextMergeCacheOptions;
+  };
+
+}  // namespace atn
+}  // namespace antlr4

runtime/Cpp/runtime/src/atn/PredictionContext.cpp

@@ -6,6 +6,8 @@
 #include "atn/SingletonPredictionContext.h"
 #include "misc/MurmurHash.h"
 #include "atn/ArrayPredictionContext.h"
+#include "atn/PredictionContextCache.h"
+#include "atn/PredictionContextMergeCache.h"
 #include "RuleContext.h"
 #include "ParserRuleContext.h"
 #include "atn/RuleTransition.h"
@@ -35,7 +37,8 @@ namespace {
 
   Ref<const PredictionContext> getCachedContextImpl(const Ref<const PredictionContext> &context,
                                                     PredictionContextCache &contextCache,
-                                                    std::unordered_map<Ref<const PredictionContext>, Ref<const PredictionContext>> &visited) {
+                                                    std::unordered_map<Ref<const PredictionContext>,
+                                                    Ref<const PredictionContext>> &visited) {
     if (context->isEmpty()) {
       return context;
     }
@@ -95,7 +98,8 @@ namespace {
     return updated;
   }
 
-  void getAllContextNodesImpl(const Ref<const PredictionContext> &context, std::vector<Ref<const PredictionContext>> &nodes,
+  void getAllContextNodesImpl(const Ref<const PredictionContext> &context,
+                              std::vector<Ref<const PredictionContext>> &nodes,
                               std::unordered_set<const PredictionContext*> &visited) {
 
     if (visited.find(context.get()) != visited.end()) {
@@ -573,54 +577,3 @@ std::vector<std::string> PredictionContext::toStrings(Recognizer *recognizer, co
 
   return result;
 }
-
-//----------------- PredictionContextCache ------------------------------------------------------------------------
-
-void PredictionContextCache::put(const Ref<const PredictionContext> &value) {
-  assert(value);
-
-  _data.insert(value);
-}
-
-Ref<const PredictionContext> PredictionContextCache::get(const Ref<const PredictionContext> &value) const {
-  assert(value);
-
-  auto iterator = _data.find(value);
-  if (iterator == _data.end()) {
-    return nullptr;
-  }
-  return *iterator;
-}
-
-//----------------- PredictionContextMergeCache ------------------------------------------------------------------------
-
-Ref<const PredictionContext> PredictionContextMergeCache::put(const Ref<const PredictionContext> &key1,
-                                                              const Ref<const PredictionContext> &key2,
-                                                              Ref<const PredictionContext> value) {
-  assert(key1);
-  assert(key2);
-
-  return _data.try_emplace(key1).first->second.insert_or_assign(key2, std::move(value)).first->second;
-}
-
-Ref<const PredictionContext> PredictionContextMergeCache::get(const Ref<const PredictionContext> &key1,
-                                                              const Ref<const PredictionContext> &key2) const {
-  assert(key1);
-  assert(key2);
-
-  auto iterator1 = _data.find(key1);
-  if (iterator1 == _data.end()) {
-    return nullptr;
-  }
-
-  auto iterator2 = iterator1->second.find(key2);
-  if (iterator2 == iterator1->second.end()) {
-    return nullptr;
-  }
-
-  return iterator2->second;
-}
-
-void PredictionContextMergeCache::clear() {
-  _data.clear();
-}