expressir 1.2.4 → 1.2.5

data/ext/express-parser/antlr4-upstream/runtime/Cpp/runtime/src/atn/PredictionContext.cpp

@@ -3,35 +3,136 @@
  * can be found in the LICENSE.txt file in the project root.
  */
 
-#include "atn/EmptyPredictionContext.h"
+#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"
 #include "support/Arrays.h"
 #include "support/CPPUtils.h"
+#include "support/Casts.h"
 
 #include "atn/PredictionContext.h"
 
 using namespace antlr4;
 using namespace antlr4::misc;
 using namespace antlr4::atn;
-
 using namespace antlrcpp;
 
-size_t PredictionContext::globalNodeCount = 0;
-const Ref<PredictionContext> PredictionContext::EMPTY = std::make_shared<EmptyPredictionContext>();
+namespace {
 
-//----------------- PredictionContext ----------------------------------------------------------------------------------
+  void combineCommonParents(std::vector<Ref<const PredictionContext>> &parents) {
+    std::unordered_set<Ref<const PredictionContext>> uniqueParents;
+    uniqueParents.reserve(parents.size());
+    for (const auto &parent : parents) {
+      uniqueParents.insert(parent);
+    }
+    for (auto &parent : parents) {
+      parent = *uniqueParents.find(parent);
+    }
+  }
 
-PredictionContext::PredictionContext(size_t cachedHashCode) : id(globalNodeCount++), cachedHashCode(cachedHashCode)  {
-}
+  Ref<const PredictionContext> getCachedContextImpl(const Ref<const PredictionContext> &context,
+                                                    PredictionContextCache &contextCache,
+                                                    std::unordered_map<Ref<const PredictionContext>,
+                                                    Ref<const PredictionContext>> &visited) {
+    if (context->isEmpty()) {
+      return context;
+    }
+
+    {
+      auto iterator = visited.find(context);
+      if (iterator != visited.end()) {
+        return iterator->second; // Not necessarly the same as context.
+      }
+    }
+
+    auto cached = contextCache.get(context);
+    if (cached) {
+      visited[context] = cached;
+      return cached;
+    }
+
+    bool changed = false;
+
+    std::vector<Ref<const PredictionContext>> parents(context->size());
+    for (size_t i = 0; i < parents.size(); i++) {
+      auto parent = getCachedContextImpl(context->getParent(i), contextCache, visited);
+      if (changed || parent != context->getParent(i)) {
+        if (!changed) {
+          parents.clear();
+          for (size_t j = 0; j < context->size(); j++) {
+            parents.push_back(context->getParent(j));
+          }
+
+          changed = true;
+        }
+
+        parents[i] = std::move(parent);
+      }
+    }
+
+    if (!changed) {
+      visited[context] = context;
+      contextCache.put(context);
+      return context;
+    }
+
+    Ref<const PredictionContext> updated;
+    if (parents.empty()) {
+      updated = PredictionContext::EMPTY;
+    } else if (parents.size() == 1) {
+      updated = SingletonPredictionContext::create(std::move(parents[0]), context->getReturnState(0));
+      contextCache.put(updated);
+    } else {
+      updated = std::make_shared<ArrayPredictionContext>(std::move(parents), downCast<const ArrayPredictionContext*>(context.get())->returnStates);
+      contextCache.put(updated);
+    }
+
+    visited[updated] = updated;
+    visited[context] = updated;
+
+    return updated;
+  }
+
+  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()) {
+      return; // Already done.
+    }
+
+    visited.insert(context.get());
+    nodes.push_back(context);
+
+    for (size_t i = 0; i < context->size(); i++) {
+      getAllContextNodesImpl(context->getParent(i), nodes, visited);
+    }
+  }
+
+  size_t insertOrAssignNodeId(std::unordered_map<const PredictionContext*, size_t> &nodeIds, size_t &nodeId, const PredictionContext *node) {
+    auto existing = nodeIds.find(node);
+    if (existing != nodeIds.end()) {
+      return existing->second;
+    }
+    return nodeIds.insert({node, nodeId++}).first->second;
+  }
 
-PredictionContext::~PredictionContext() {
 }
 
-Ref<PredictionContext> PredictionContext::fromRuleContext(const ATN &atn, RuleContext *outerContext) {
+const Ref<const PredictionContext> PredictionContext::EMPTY = std::make_shared<SingletonPredictionContext>(nullptr, PredictionContext::EMPTY_RETURN_STATE);
+
+//----------------- PredictionContext ----------------------------------------------------------------------------------
+
+PredictionContext::PredictionContext(PredictionContextType contextType) : _contextType(contextType), _hashCode(0) {}
+
+PredictionContext::PredictionContext(PredictionContext&& other) : _contextType(other._contextType), _hashCode(other._hashCode.exchange(0, std::memory_order_relaxed)) {}
+
+Ref<const PredictionContext> PredictionContext::fromRuleContext(const ATN &atn, RuleContext *outerContext) {
   if (outerContext == nullptr) {
     return PredictionContext::EMPTY;
   }
@@ -43,15 +144,9 @@ Ref<PredictionContext> PredictionContext::fromRuleContext(const ATN &atn, RuleCo
   }
 
   // If we have a parent, convert it to a PredictionContext graph
-  Ref<PredictionContext> parent = PredictionContext::fromRuleContext(atn, dynamic_cast<RuleContext *>(outerContext->parent));
-
-  ATNState *state = atn.states.at(outerContext->invokingState);
-  RuleTransition *transition = (RuleTransition *)state->transitions[0];
-  return SingletonPredictionContext::create(parent, transition->followState->stateNumber);
-}
-
-bool PredictionContext::isEmpty() const {
-  return this == EMPTY.get();
+  auto parent = PredictionContext::fromRuleContext(atn, RuleContext::is(outerContext->parent) ? downCast<RuleContext*>(outerContext->parent) : nullptr);
+  const auto *transition = downCast<const RuleTransition*>(atn.states[outerContext->invokingState]->transitions[0].get());
+  return SingletonPredictionContext::create(std::move(parent), transition->followState->stateNumber);
 }
 
 bool PredictionContext::hasEmptyPath() const {
@@ -60,40 +155,19 @@ bool PredictionContext::hasEmptyPath() const {
 }
 
 size_t PredictionContext::hashCode() const {
-  return cachedHashCode;
-}
-
-size_t PredictionContext::calculateEmptyHashCode() {
-  size_t hash = MurmurHash::initialize(INITIAL_HASH);
-  hash = MurmurHash::finish(hash, 0);
-  return hash;
-}
-
-size_t PredictionContext::calculateHashCode(Ref<PredictionContext> parent, size_t returnState) {
-  size_t hash = MurmurHash::initialize(INITIAL_HASH);
-  hash = MurmurHash::update(hash, parent);
-  hash = MurmurHash::update(hash, returnState);
-  hash = MurmurHash::finish(hash, 2);
-  return hash;
-}
-
-size_t PredictionContext::calculateHashCode(const std::vector<Ref<PredictionContext>> &parents,
-                                            const std::vector<size_t> &returnStates) {
-  size_t hash = MurmurHash::initialize(INITIAL_HASH);
-
-  for (auto parent : parents) {
-    hash = MurmurHash::update(hash, parent);
-  }
-
-  for (auto returnState : returnStates) {
-    hash = MurmurHash::update(hash, returnState);
+  auto hash = cachedHashCode();
+  if (hash == 0) {
+    hash = hashCodeImpl();
+    if (hash == 0) {
+      hash = std::numeric_limits<size_t>::max();
+    }
+    _hashCode.store(hash, std::memory_order_relaxed);
   }
-
-  return MurmurHash::finish(hash, parents.size() + returnStates.size());
+  return hash;
 }
 
-Ref<PredictionContext> PredictionContext::merge(const Ref<PredictionContext> &a,
-  const Ref<PredictionContext> &b, bool rootIsWildcard, PredictionContextMergeCache *mergeCache) {
+Ref<const PredictionContext> PredictionContext::merge(Ref<const PredictionContext> a, Ref<const PredictionContext> b,
+                                                      bool rootIsWildcard, PredictionContextMergeCache *mergeCache) {
   assert(a && b);
 
   // share same graph if both same
@@ -101,42 +175,45 @@ Ref<PredictionContext> PredictionContext::merge(const Ref<PredictionContext> &a,
     return a;
   }
 
-  if (is<SingletonPredictionContext>(a) && is<SingletonPredictionContext>(b)) {
-    return mergeSingletons(std::dynamic_pointer_cast<SingletonPredictionContext>(a),
-                           std::dynamic_pointer_cast<SingletonPredictionContext>(b), rootIsWildcard, mergeCache);
+  const auto aType = a->getContextType();
+  const auto bType = b->getContextType();
+
+  if (aType == PredictionContextType::SINGLETON && bType == PredictionContextType::SINGLETON) {
+    return mergeSingletons(std::static_pointer_cast<const SingletonPredictionContext>(std::move(a)),
+                           std::static_pointer_cast<const SingletonPredictionContext>(std::move(b)), rootIsWildcard, mergeCache);
   }
 
   // At least one of a or b is array.
   // If one is $ and rootIsWildcard, return $ as * wildcard.
   if (rootIsWildcard) {
-    if (is<EmptyPredictionContext>(a)) {
+    if (a == PredictionContext::EMPTY) {
       return a;
     }
-    if (is<EmptyPredictionContext>(b)) {
+    if (b == PredictionContext::EMPTY) {
       return b;
     }
   }
 
   // convert singleton so both are arrays to normalize
-  Ref<ArrayPredictionContext> left;
-  if (is<SingletonPredictionContext>(a)) {
-    left = std::make_shared<ArrayPredictionContext>(std::dynamic_pointer_cast<SingletonPredictionContext>(a));
+  Ref<const ArrayPredictionContext> left;
+  if (aType == PredictionContextType::SINGLETON) {
+    left = std::make_shared<ArrayPredictionContext>(downCast<const SingletonPredictionContext&>(*a));
   } else {
-    left = std::dynamic_pointer_cast<ArrayPredictionContext>(a);
+    left = std::static_pointer_cast<const ArrayPredictionContext>(std::move(a));
   }
-  Ref<ArrayPredictionContext> right;
-  if (is<SingletonPredictionContext>(b)) {
-    right = std::make_shared<ArrayPredictionContext>(std::dynamic_pointer_cast<SingletonPredictionContext>(b));
+  Ref<const ArrayPredictionContext> right;
+  if (bType == PredictionContextType::SINGLETON) {
+    right = std::make_shared<ArrayPredictionContext>(downCast<const SingletonPredictionContext&>(*b));
   } else {
-    right = std::dynamic_pointer_cast<ArrayPredictionContext>(b);
+    right = std::static_pointer_cast<const ArrayPredictionContext>(std::move(b));
   }
-  return mergeArrays(left, right, rootIsWildcard, mergeCache);
+  return mergeArrays(std::move(left), std::move(right), rootIsWildcard, mergeCache);
 }
 
-Ref<PredictionContext> PredictionContext::mergeSingletons(const Ref<SingletonPredictionContext> &a,
-  const Ref<SingletonPredictionContext> &b, bool rootIsWildcard, PredictionContextMergeCache *mergeCache) {
+Ref<const PredictionContext> PredictionContext::mergeSingletons(Ref<const SingletonPredictionContext> a, Ref<const SingletonPredictionContext> b,
+                                                                bool rootIsWildcard, PredictionContextMergeCache *mergeCache) {
 
-  if (mergeCache != nullptr) { // Can be null if not given to the ATNState from which this call originates.
+  if (mergeCache) {
     auto existing = mergeCache->get(a, b);
     if (existing) {
       return existing;
@@ -147,18 +224,18 @@ Ref<PredictionContext> PredictionContext::mergeSingletons(const Ref<SingletonPre
     }
   }
 
-  Ref<PredictionContext> rootMerge = mergeRoot(a, b, rootIsWildcard);
+  auto rootMerge = mergeRoot(a, b, rootIsWildcard);
   if (rootMerge) {
-    if (mergeCache != nullptr) {
-      mergeCache->put(a, b, rootMerge);
+    if (mergeCache) {
+      return mergeCache->put(a, b, std::move(rootMerge));
     }
     return rootMerge;
   }
 
-  Ref<PredictionContext> parentA = a->parent;
-  Ref<PredictionContext> parentB = b->parent;
+  const auto& parentA = a->parent;
+  const auto& parentB = b->parent;
   if (a->returnState == b->returnState) { // a == b
-    Ref<PredictionContext> parent = merge(parentA, parentB, rootIsWildcard, mergeCache);
+    auto parent = merge(parentA, parentB, rootIsWildcard, mergeCache);
 
     // If parent is same as existing a or b parent or reduced to a parent, return it.
     if (parent == parentA) { // ax + bx = ax, if a=b
@@ -172,55 +249,55 @@ Ref<PredictionContext> PredictionContext::mergeSingletons(const Ref<SingletonPre
     // merge parents x and y, giving array node with x,y then remainders
     // of those graphs.  dup a, a' points at merged array
     // new joined parent so create new singleton pointing to it, a'
-    Ref<PredictionContext> a_ = SingletonPredictionContext::create(parent, a->returnState);
-    if (mergeCache != nullptr) {
-      mergeCache->put(a, b, a_);
-    }
-    return a_;
-  } else {
-    // a != b payloads differ
-    // see if we can collapse parents due to $+x parents if local ctx
-    Ref<PredictionContext> singleParent;
-    if (a == b || (*parentA == *parentB)) { // ax + bx = [a,b]x
-      singleParent = parentA;
-    }
-    if (singleParent) { // parents are same, sort payloads and use same parent
-      std::vector<size_t> payloads = { a->returnState, b->returnState };
-      if (a->returnState > b->returnState) {
-        payloads[0] = b->returnState;
-        payloads[1] = a->returnState;
-      }
-      std::vector<Ref<PredictionContext>> parents = { singleParent, singleParent };
-      Ref<PredictionContext> a_ = std::make_shared<ArrayPredictionContext>(parents, payloads);
-      if (mergeCache != nullptr) {
-        mergeCache->put(a, b, a_);
-      }
-      return a_;
-    }
-
-    // parents differ and can't merge them. Just pack together
-    // into array; can't merge.
-    // ax + by = [ax,by]
-    Ref<PredictionContext> a_;
-    if (a->returnState > b->returnState) { // sort by payload
-      std::vector<size_t> payloads = { b->returnState, a->returnState };
-      std::vector<Ref<PredictionContext>> parents = { b->parent, a->parent };
-      a_ = std::make_shared<ArrayPredictionContext>(parents, payloads);
-    } else {
-      std::vector<size_t> payloads = {a->returnState, b->returnState};
-      std::vector<Ref<PredictionContext>> parents = { a->parent, b->parent };
-      a_ = std::make_shared<ArrayPredictionContext>(parents, payloads);
-    }
-
-    if (mergeCache != nullptr) {
-      mergeCache->put(a, b, a_);
-    }
-    return a_;
-  }
+    auto c = SingletonPredictionContext::create(std::move(parent), a->returnState);
+    if (mergeCache) {
+      return mergeCache->put(a, b, std::move(c));
+    }
+    return c;
+  }
+  // a != b payloads differ
+  // see if we can collapse parents due to $+x parents if local ctx
+  Ref<const PredictionContext> singleParent;
+  if (a == b || (*parentA == *parentB)) { // ax + bx = [a,b]x
+    singleParent = parentA;
+  }
+  if (singleParent) { // parents are same, sort payloads and use same parent
+    std::vector<size_t> payloads = { a->returnState, b->returnState };
+    if (a->returnState > b->returnState) {
+      payloads[0] = b->returnState;
+      payloads[1] = a->returnState;
+    }
+    std::vector<Ref<const PredictionContext>> parents = { singleParent, singleParent };
+    auto c = std::make_shared<ArrayPredictionContext>(std::move(parents), std::move(payloads));
+    if (mergeCache) {
+      return mergeCache->put(a, b, std::move(c));
+    }
+    return c;
+  }
+
+  // parents differ and can't merge them. Just pack together
+  // into array; can't merge.
+  // ax + by = [ax,by]
+  if (a->returnState > b->returnState) { // sort by payload
+    std::vector<size_t> payloads = { b->returnState, a->returnState };
+    std::vector<Ref<const PredictionContext>> parents = { b->parent, a->parent };
+    auto c = std::make_shared<ArrayPredictionContext>(std::move(parents), std::move(payloads));
+    if (mergeCache) {
+      return mergeCache->put(a, b, std::move(c));
+    }
+    return c;
+  }
+  std::vector<size_t> payloads = {a->returnState, b->returnState};
+  std::vector<Ref<const PredictionContext>> parents = { a->parent, b->parent };
+  auto c = std::make_shared<ArrayPredictionContext>(std::move(parents), std::move(payloads));
+  if (mergeCache) {
+    return mergeCache->put(a, b, std::move(c));
+  }
+  return c;
 }
 
-Ref<PredictionContext> PredictionContext::mergeRoot(const Ref<SingletonPredictionContext> &a,
-  const Ref<SingletonPredictionContext> &b, bool rootIsWildcard) {
+Ref<const PredictionContext> PredictionContext::mergeRoot(Ref<const SingletonPredictionContext> a, Ref<const SingletonPredictionContext> b,
+                                                          bool rootIsWildcard) {
   if (rootIsWildcard) {
     if (a == EMPTY) { // * + b = *
       return EMPTY;
@@ -234,24 +311,22 @@ Ref<PredictionContext> PredictionContext::mergeRoot(const Ref<SingletonPredictio
     }
     if (a == EMPTY) { // $ + x = [$,x]
       std::vector<size_t> payloads = { b->returnState, EMPTY_RETURN_STATE };
-      std::vector<Ref<PredictionContext>> parents = { b->parent, nullptr };
-      Ref<PredictionContext> joined = std::make_shared<ArrayPredictionContext>(parents, payloads);
-      return joined;
+      std::vector<Ref<const PredictionContext>> parents = { b->parent, nullptr };
+      return std::make_shared<ArrayPredictionContext>(std::move(parents), std::move(payloads));
     }
     if (b == EMPTY) { // x + $ = [$,x] ($ is always first if present)
       std::vector<size_t> payloads = { a->returnState, EMPTY_RETURN_STATE };
-      std::vector<Ref<PredictionContext>> parents = { a->parent, nullptr };
-      Ref<PredictionContext> joined = std::make_shared<ArrayPredictionContext>(parents, payloads);
-      return joined;
+      std::vector<Ref<const PredictionContext>> parents = { a->parent, nullptr };
+      return std::make_shared<ArrayPredictionContext>(std::move(parents), std::move(payloads));
     }
   }
   return nullptr;
 }
 
-Ref<PredictionContext> PredictionContext::mergeArrays(const Ref<ArrayPredictionContext> &a,
-  const Ref<ArrayPredictionContext> &b, bool rootIsWildcard, PredictionContextMergeCache *mergeCache) {
+Ref<const PredictionContext> PredictionContext::mergeArrays(Ref<const ArrayPredictionContext> a, Ref<const ArrayPredictionContext> b,
+                                                            bool rootIsWildcard, PredictionContextMergeCache *mergeCache) {
 
-  if (mergeCache != nullptr) {
+  if (mergeCache) {
     auto existing = mergeCache->get(a, b);
     if (existing) {
       return existing;
@@ -268,36 +343,33 @@ Ref<PredictionContext> PredictionContext::mergeArrays(const Ref<ArrayPredictionC
   size_t k = 0; // walks target M array
 
   std::vector<size_t> mergedReturnStates(a->returnStates.size() + b->returnStates.size());
-  std::vector<Ref<PredictionContext>> mergedParents(a->returnStates.size() + b->returnStates.size());
+  std::vector<Ref<const PredictionContext>> mergedParents(a->returnStates.size() + b->returnStates.size());
 
   // walk and merge to yield mergedParents, mergedReturnStates
   while (i < a->returnStates.size() && j < b->returnStates.size()) {
-    Ref<PredictionContext> a_parent = a->parents[i];
-    Ref<PredictionContext> b_parent = b->parents[j];
+    const auto& parentA = a->parents[i];
+    const auto& parentB = b->parents[j];
     if (a->returnStates[i] == b->returnStates[j]) {
       // same payload (stack tops are equal), must yield merged singleton
       size_t payload = a->returnStates[i];
       // $+$ = $
-      bool both$ = payload == EMPTY_RETURN_STATE && !a_parent && !b_parent;
-      bool ax_ax = (a_parent && b_parent) && *a_parent == *b_parent; // ax+ax -> ax
+      bool both$ = payload == EMPTY_RETURN_STATE && !parentA && !parentB;
+      bool ax_ax = (parentA && parentB) && *parentA == *parentB; // ax+ax -> ax
       if (both$ || ax_ax) {
-        mergedParents[k] = a_parent; // choose left
+        mergedParents[k] = parentA; // choose left
         mergedReturnStates[k] = payload;
-      }
-      else { // ax+ay -> a'[x,y]
-        Ref<PredictionContext> mergedParent = merge(a_parent, b_parent, rootIsWildcard, mergeCache);
-        mergedParents[k] = mergedParent;
+      } else { // ax+ay -> a'[x,y]
+        mergedParents[k] = merge(parentA, parentB, rootIsWildcard, mergeCache);
         mergedReturnStates[k] = payload;
       }
       i++; // hop over left one as usual
       j++; // but also skip one in right side since we merge
     } else if (a->returnStates[i] < b->returnStates[j]) { // copy a[i] to M
-      mergedParents[k] = a_parent;
+      mergedParents[k] = parentA;
       mergedReturnStates[k] = a->returnStates[i];
       i++;
-    }
-    else { // b > a, copy b[j] to M
-      mergedParents[k] = b_parent;
+    } else { // b > a, copy b[j] to M
+      mergedParents[k] = parentB;
       mergedReturnStates[k] = b->returnStates[j];
       j++;
     }
@@ -306,13 +378,13 @@ Ref<PredictionContext> PredictionContext::mergeArrays(const Ref<ArrayPredictionC
 
   // copy over any payloads remaining in either array
   if (i < a->returnStates.size()) {
-    for (std::vector<int>::size_type p = i; p < a->returnStates.size(); p++) {
+    for (auto p = i; p < a->returnStates.size(); p++) {
       mergedParents[k] = a->parents[p];
       mergedReturnStates[k] = a->returnStates[p];
       k++;
     }
   } else {
-    for (std::vector<int>::size_type p = j; p < b->returnStates.size(); p++) {
+    for (auto p = j; p < b->returnStates.size(); p++) {
       mergedParents[k] = b->parents[p];
       mergedReturnStates[k] = b->returnStates[p];
       k++;
@@ -322,63 +394,42 @@ Ref<PredictionContext> PredictionContext::mergeArrays(const Ref<ArrayPredictionC
   // trim merged if we combined a few that had same stack tops
   if (k < mergedParents.size()) { // write index < last position; trim
     if (k == 1) { // for just one merged element, return singleton top
-      Ref<PredictionContext> a_ = SingletonPredictionContext::create(mergedParents[0], mergedReturnStates[0]);
-      if (mergeCache != nullptr) {
-        mergeCache->put(a, b, a_);
+      auto c = SingletonPredictionContext::create(std::move(mergedParents[0]), mergedReturnStates[0]);
+      if (mergeCache) {
+        return mergeCache->put(a, b, std::move(c));
       }
-      return a_;
+      return c;
     }
     mergedParents.resize(k);
     mergedReturnStates.resize(k);
   }
 
-  Ref<ArrayPredictionContext> M = std::make_shared<ArrayPredictionContext>(mergedParents, mergedReturnStates);
+  ArrayPredictionContext m(std::move(mergedParents), std::move(mergedReturnStates));
 
   // if we created same array as a or b, return that instead
   // TODO: track whether this is possible above during merge sort for speed
-  if (*M == *a) {
-    if (mergeCache != nullptr) {
-      mergeCache->put(a, b, a);
+  if (m == *a) {
+    if (mergeCache) {
+      return mergeCache->put(a, b, a);
     }
     return a;
   }
-  if (*M == *b) {
-    if (mergeCache != nullptr) {
-      mergeCache->put(a, b, b);
+  if (m == *b) {
+    if (mergeCache) {
+      return mergeCache->put(a, b, b);
     }
     return b;
   }
 
-  // ml: this part differs from Java code. We have to recreate the context as the parents array is copied on creation.
-  if (combineCommonParents(mergedParents)) {
-    mergedReturnStates.resize(mergedParents.size());
-    M = std::make_shared<ArrayPredictionContext>(mergedParents, mergedReturnStates);
-  }
-
-  if (mergeCache != nullptr) {
-    mergeCache->put(a, b, M);
-  }
-  return M;
-}
-
-bool PredictionContext::combineCommonParents(std::vector<Ref<PredictionContext>> &parents) {
-
-  std::set<Ref<PredictionContext>> uniqueParents;
-  for (size_t p = 0; p < parents.size(); ++p) {
-    Ref<PredictionContext> parent = parents[p];
-    if (uniqueParents.find(parent) == uniqueParents.end()) { // don't replace
-      uniqueParents.insert(parent);
-    }
-  }
-
-  for (size_t p = 0; p < parents.size(); ++p) {
-    parents[p] = *uniqueParents.find(parents[p]);
+  combineCommonParents(m.parents);
+  auto c = std::make_shared<ArrayPredictionContext>(std::move(m));
+  if (mergeCache) {
+    return mergeCache->put(a, b, std::move(c));
   }
-
-  return true;
+  return c;
 }
 
-std::string PredictionContext::toDOTString(const Ref<PredictionContext> &context) {
+std::string PredictionContext::toDOTString(const Ref<const PredictionContext> &context) {
   if (context == nullptr) {
     return "";
   }
@@ -386,24 +437,23 @@ std::string PredictionContext::toDOTString(const Ref<PredictionContext> &context
   std::stringstream ss;
   ss << "digraph G {\n" << "rankdir=LR;\n";
 
-  std::vector<Ref<PredictionContext>> nodes = getAllContextNodes(context);
-  std::sort(nodes.begin(), nodes.end(), [](const Ref<PredictionContext> &o1, const Ref<PredictionContext> &o2) {
-    return o1->id - o2->id;
-  });
+  std::vector<Ref<const PredictionContext>> nodes = getAllContextNodes(context);
+  std::unordered_map<const PredictionContext*, size_t> nodeIds;
+  size_t nodeId = 0;
 
-  for (auto current : nodes) {
-    if (is<SingletonPredictionContext>(current)) {
-      std::string s = std::to_string(current->id);
+  for (const auto &current : nodes) {
+    if (current->getContextType() == PredictionContextType::SINGLETON) {
+      std::string s = std::to_string(insertOrAssignNodeId(nodeIds, nodeId, current.get()));
       ss << "  s" << s;
       std::string returnState = std::to_string(current->getReturnState(0));
-      if (is<EmptyPredictionContext>(current)) {
+      if (current == PredictionContext::EMPTY) {
         returnState = "$";
       }
       ss << " [label=\"" << returnState << "\"];\n";
       continue;
     }
-    Ref<ArrayPredictionContext> arr = std::static_pointer_cast<ArrayPredictionContext>(current);
-    ss << "  s" << arr->id << " [shape=box, label=\"" << "[";
+    Ref<const ArrayPredictionContext> arr = std::static_pointer_cast<const ArrayPredictionContext>(current);
+    ss << "  s" << insertOrAssignNodeId(nodeIds, nodeId, arr.get()) << " [shape=box, label=\"" << "[";
     bool first = true;
     for (auto inv : arr->returnStates) {
       if (!first) {
@@ -420,7 +470,7 @@ std::string PredictionContext::toDOTString(const Ref<PredictionContext> &context
     ss << "\"];\n";
   }
 
-  for (auto current : nodes) {
+  for (const auto &current : nodes) {
     if (current == EMPTY) {
       continue;
     }
@@ -428,7 +478,7 @@ std::string PredictionContext::toDOTString(const Ref<PredictionContext> &context
       if (!current->getParent(i)) {
         continue;
       }
-      ss << "  s" << current->id << "->" << "s" << current->getParent(i)->id;
+      ss << "  s" << insertOrAssignNodeId(nodeIds, nodeId, current.get()) << "->" << "s" << insertOrAssignNodeId(nodeIds, nodeId, current->getParent(i).get());
       if (current->size() > 1) {
         ss << " [label=\"parent[" << i << "]\"];\n";
       } else {
@@ -442,112 +492,31 @@ std::string PredictionContext::toDOTString(const Ref<PredictionContext> &context
 }
 
 // The "visited" map is just a temporary structure to control the retrieval process (which is recursive).
-Ref<PredictionContext> PredictionContext::getCachedContext(const Ref<PredictionContext> &context,
-  PredictionContextCache &contextCache, std::map<Ref<PredictionContext>, Ref<PredictionContext>> &visited) {
-  if (context->isEmpty()) {
-    return context;
-  }
-
-  {
-    auto iterator = visited.find(context);
-    if (iterator != visited.end())
-      return iterator->second; // Not necessarly the same as context.
-  }
-
-  auto iterator = contextCache.find(context);
-  if (iterator != contextCache.end()) {
-    visited[context] = *iterator;
-
-    return *iterator;
-  }
-
-  bool changed = false;
-
-  std::vector<Ref<PredictionContext>> parents(context->size());
-  for (size_t i = 0; i < parents.size(); i++) {
-    Ref<PredictionContext> parent = getCachedContext(context->getParent(i), contextCache, visited);
-    if (changed || parent != context->getParent(i)) {
-      if (!changed) {
-        parents.clear();
-        for (size_t j = 0; j < context->size(); j++) {
-          parents.push_back(context->getParent(j));
-        }
-
-        changed = true;
-      }
-
-      parents[i] = parent;
-    }
-  }
-
-  if (!changed) {
-    contextCache.insert(context);
-    visited[context] = context;
-
-    return context;
-  }
-
-  Ref<PredictionContext> updated;
-  if (parents.empty()) {
-    updated = EMPTY;
-  } else if (parents.size() == 1) {
-    updated = SingletonPredictionContext::create(parents[0], context->getReturnState(0));
-    contextCache.insert(updated);
-  } else {
-    updated = std::make_shared<ArrayPredictionContext>(parents, std::dynamic_pointer_cast<ArrayPredictionContext>(context)->returnStates);
-    contextCache.insert(updated);
-  }
-
-  visited[updated] = updated;
-  visited[context] = updated;
-
-  return updated;
+Ref<const PredictionContext> PredictionContext::getCachedContext(const Ref<const PredictionContext> &context,
+                                                                 PredictionContextCache &contextCache) {
+  std::unordered_map<Ref<const PredictionContext>, Ref<const PredictionContext>> visited;
+  return getCachedContextImpl(context, contextCache, visited);
 }
 
-std::vector<Ref<PredictionContext>> PredictionContext::getAllContextNodes(const Ref<PredictionContext> &context) {
-  std::vector<Ref<PredictionContext>> nodes;
-  std::set<PredictionContext *> visited;
-  getAllContextNodes_(context, nodes, visited);
+std::vector<Ref<const PredictionContext>> PredictionContext::getAllContextNodes(const Ref<const PredictionContext> &context) {
+  std::vector<Ref<const PredictionContext>> nodes;
+  std::unordered_set<const PredictionContext*> visited;
+  getAllContextNodesImpl(context, nodes, visited);
   return nodes;
 }
 
-
-void PredictionContext::getAllContextNodes_(const Ref<PredictionContext> &context, std::vector<Ref<PredictionContext>> &nodes,
-  std::set<PredictionContext *> &visited) {
-
-  if (visited.find(context.get()) != visited.end()) {
-    return; // Already done.
-  }
-
-  visited.insert(context.get());
-  nodes.push_back(context);
-
-  for (size_t i = 0; i < context->size(); i++) {
-    getAllContextNodes_(context->getParent(i), nodes, visited);
-  }
-}
-
-std::string PredictionContext::toString() const {
-
-  return antlrcpp::toString(this);
-}
-
-std::string PredictionContext::toString(Recognizer * /*recog*/) const {
-  return toString();
-}
-
-std::vector<std::string> PredictionContext::toStrings(Recognizer *recognizer, int currentState) {
+std::vector<std::string> PredictionContext::toStrings(Recognizer *recognizer, int currentState) const {
   return toStrings(recognizer, EMPTY, currentState);
 }
 
-std::vector<std::string> PredictionContext::toStrings(Recognizer *recognizer, const Ref<PredictionContext> &stop, int currentState) {
+std::vector<std::string> PredictionContext::toStrings(Recognizer *recognizer, const Ref<const PredictionContext> &stop, int currentState) const {
 
   std::vector<std::string> result;
 
   for (size_t perm = 0; ; perm++) {
     size_t offset = 0;
     bool last = true;
-    PredictionContext *p = this;
+    const PredictionContext *p = this;
     size_t stateNumber = currentState;
 
     std::stringstream ss;
@@ -608,55 +577,3 @@ std::vector<std::string> PredictionContext::toStrings(Recognizer *recognizer, co
 
   return result;
 }
-
-//----------------- PredictionContextMergeCache ------------------------------------------------------------------------
-
-Ref<PredictionContext> PredictionContextMergeCache::put(Ref<PredictionContext> const& key1, Ref<PredictionContext> const& key2,
-                                                        Ref<PredictionContext> const& value) {
-  Ref<PredictionContext> previous;
-
-  auto iterator = _data.find(key1);
-  if (iterator == _data.end())
-    _data[key1][key2] = value;
-  else {
-    auto iterator2 = iterator->second.find(key2);
-    if (iterator2 != iterator->second.end())
-      previous = iterator2->second;
-    iterator->second[key2] = value;
-  }
-
-  return previous;
-}
-
-Ref<PredictionContext> PredictionContextMergeCache::get(Ref<PredictionContext> const& key1, Ref<PredictionContext> const& key2) {
-  auto iterator = _data.find(key1);
-  if (iterator == _data.end())
-    return nullptr;
-
-  auto iterator2 = iterator->second.find(key2);
-  if (iterator2 == iterator->second.end())
-    return nullptr;
-
-  return iterator2->second;
-}
-
-void PredictionContextMergeCache::clear() {
-  _data.clear();
-}
-
-std::string PredictionContextMergeCache::toString() const {
-  std::string result;
-  for (auto pair : _data)
-    for (auto pair2 : pair.second)
-      result += pair2.second->toString() + "\n";
-
-  return result;
-}
-
-size_t PredictionContextMergeCache::count() const {
-  size_t result = 0;
-  for (auto entry : _data)
-    result += entry.second.size();
-  return result;
-}
-