1 files changed, 43 insertions, 26 deletions

diff --git a/native/src/correction.h b/native/src/correction.h
index ae6c7a421..9d385a44e 100644
--- a/native/src/correction.h
+++ b/native/src/correction.h
@@ -18,6 +18,7 @@
 #define LATINIME_CORRECTION_H
 
 #include <stdint.h>
+#include "correction_state.h"
 
 #include "defines.h"
 
@@ -39,16 +40,16 @@ public:
     Correction(const int typedLetterMultiplier, const int fullWordMultiplier);
     void initCorrection(
             const ProximityInfo *pi, const int inputLength, const int maxWordLength);
+    void initCorrectionState(const int rootPos, const int childCount, const bool traverseAll);
+
+    // TODO: remove
     void setCorrectionParams(const int skipPos, const int excessivePos, const int transposedPos,
             const int spaceProximityPos, const int missingSpacePos);
     void checkState();
-    void initProcessState(const int matchCount, const int inputIndex, const int outputIndex,
-            const bool traverseAllNodes, const int diffs);
-    void getProcessState(int *matchedCount, int *inputIndex, int *outputIndex,
-            bool *traverseAllNodes, int *diffs);
+    bool initProcessState(const int index);
+
     int getOutputIndex();
     int getInputIndex();
-    bool needsToTraverseAll();
 
     virtual ~Correction();
     int getSpaceProximityPos() const {
@@ -77,52 +78,68 @@ public:
 
     CorrectionType processCharAndCalcState(const int32_t c, const bool isTerminal);
 
-    int getDiffs() const {
-        return mDiffs;
+    /////////////////////////
+    // Tree helper methods
+    int goDownTree(const int parentIndex, const int childCount, const int firstChildPos);
+
+    inline int getTreeSiblingPos(const int index) const {
+        return mCorrectionStates[index].mSiblingPos;
+    }
+
+    inline void setTreeSiblingPos(const int index, const int pos) {
+        mCorrectionStates[index].mSiblingPos = pos;
+    }
+
+    inline int getTreeParentIndex(const int index) const {
+        return mCorrectionStates[index].mParentIndex;
     }
 private:
-    void charMatched();
-    void incrementInputIndex();
-    void incrementOutputIndex();
-    void startTraverseAll();
+    inline void incrementInputIndex();
+    inline void incrementOutputIndex();
+    inline bool needsToTraverseAllNodes();
+    inline void startToTraverseAllNodes();
+    inline bool isQuote(const unsigned short c);
+    inline CorrectionType processSkipChar(const int32_t c, const bool isTerminal);
 
     // TODO: remove
-
-    void incrementDiffs() {
-        ++mDiffs;
+    inline void incrementProximityCount() {
+        ++mProximityCount;
     }
 
     const int TYPED_LETTER_MULTIPLIER;
     const int FULL_WORD_MULTIPLIER;
-
     const ProximityInfo *mProximityInfo;
 
     int mMaxEditDistance;
     int mMaxDepth;
     int mInputLength;
-    int mSkipPos;
-    int mSkippedOutputIndex;
     int mExcessivePos;
     int mTransposedPos;
     int mSpaceProximityPos;
     int mMissingSpacePos;
-
-    int mMatchedCharCount;
-    int mInputIndex;
-    int mOutputIndex;
     int mTerminalInputIndex;
     int mTerminalOutputIndex;
-    int mDiffs;
-    bool mTraverseAllNodes;
     unsigned short mWord[MAX_WORD_LENGTH_INTERNAL];
+    // Caveat: Do not create multiple tables per thread as this table eats up RAM a lot.
+    int mEditDistanceTable[MAX_WORD_LENGTH_INTERNAL * MAX_WORD_LENGTH_INTERNAL];
 
-    inline bool isQuote(const unsigned short c);
-    inline CorrectionType processSkipChar(const int32_t c, const bool isTerminal);
+    CorrectionState mCorrectionStates[MAX_WORD_LENGTH_INTERNAL];
+
+    // The following member variables are being used as cache values of the correction state.
+    int mOutputIndex;
+    int mInputIndex;
+    int mProximityCount;
+    int mSkippedCount;
+    int mSkipPos;
+    bool mNeedsToTraverseAllNodes;
+    bool mMatching;
+    bool mSkipping;
+    bool mProximityMatching;
 
     class RankingAlgorithm {
     public:
         static int calculateFinalFreq(const int inputIndex, const int depth,
-                const int matchCount, const int freq, const bool sameLength,
+                const int freq, const bool sameLength, int *editDistanceTable,
                 const Correction* correction);
         static int calcFreqForSplitTwoWords(const int firstFreq, const int secondFreq,
                 const Correction* correction);