diff --git a/native/jni/src/proximity_info_params.cpp b/native/jni/src/proximity_info_params.cpp
index d62e090..5a51f62 100644
--- a/native/jni/src/proximity_info_params.cpp
+++ b/native/jni/src/proximity_info_params.cpp
@@ -25,4 +25,6 @@
 const int ProximityInfoParams::NORMALIZED_SQUARED_DISTANCE_SCALING_FACTOR =
         1 << NORMALIZED_SQUARED_DISTANCE_SCALING_FACTOR_LOG_2;
 const float ProximityInfoParams::NOT_A_DISTANCE_FLOAT = -1.0f;
+// TODO: Investigate if this is required
+const float ProximityInfoParams::SEARCH_KEY_RADIUS_RATIO = 0.95f;
 } // namespace latinime
diff --git a/native/jni/src/proximity_info_params.h b/native/jni/src/proximity_info_params.h
index 171e6b4..b941fec 100644
--- a/native/jni/src/proximity_info_params.h
+++ b/native/jni/src/proximity_info_params.h
@@ -29,6 +29,7 @@
     static const int MIN_DOUBLE_LETTER_BEELINE_SPEED_PERCENTILE;
     static const int NORMALIZED_SQUARED_DISTANCE_SCALING_FACTOR;
     static const float NOT_A_DISTANCE_FLOAT;
+    static const float SEARCH_KEY_RADIUS_RATIO;
  private:
     DISALLOW_IMPLICIT_CONSTRUCTORS(ProximityInfoParams);
     static const int NORMALIZED_SQUARED_DISTANCE_SCALING_FACTOR_LOG_2;
diff --git a/native/jni/src/proximity_info_state.cpp b/native/jni/src/proximity_info_state.cpp
index 86fe7e1..e720275 100644
--- a/native/jni/src/proximity_info_state.cpp
+++ b/native/jni/src/proximity_info_state.cpp
@@ -116,28 +116,14 @@
                     mProximityInfo->getKeyCount(), lastSavedInputSize, mSampledInputSize,
                     &mSampledInputXs, &mSampledInputYs, &mSpeedRates, &mSampledLengthCache,
                     &mDistanceCache_G, &mNearKeysVector, &mCharProbabilities);
-
-            static const float READ_FORWORD_LENGTH_SCALE = 0.95f;
-            const int readForwordLength = static_cast<int>(
-                    hypotf(mProximityInfo->getKeyboardWidth(), mProximityInfo->getKeyboardHeight())
-                            * READ_FORWORD_LENGTH_SCALE);
-            for (int i = 0; i < mSampledInputSize; ++i) {
-                if (i >= lastSavedInputSize) {
-                    mSearchKeysVector[i].reset();
-                }
-                for (int j = max(i, lastSavedInputSize); j < mSampledInputSize; ++j) {
-                    if (mSampledLengthCache[j] - mSampledLengthCache[i] >= readForwordLength) {
-                        break;
-                    }
-                    mSearchKeysVector[i] |= mNearKeysVector[j];
-                }
-            }
+            ProximityInfoStateUtils::updateSearchKeysVector(mProximityInfo, mSampledInputSize,
+                    lastSavedInputSize, &mSampledLengthCache, &mNearKeysVector, &mSearchKeysVector);
         }
     }
 
     if (DEBUG_SAMPLING_POINTS) {
         ProximityInfoStateUtils::dump(isGeometric, inputSize, xCoordinates, yCoordinates,
-                mSampledInputSize, &mSampledInputXs, &mSampledInputYs, &mSpeedRates,
+                mSampledInputSize, &mSampledInputXs, &mSampledInputYs, &mSampledTimes, &mSpeedRates,
                 &mBeelineSpeedPercentiles);
     }
     // end
diff --git a/native/jni/src/proximity_info_state_utils.cpp b/native/jni/src/proximity_info_state_utils.cpp
index 9dd3f0b..be6cde1 100644
--- a/native/jni/src/proximity_info_state_utils.cpp
+++ b/native/jni/src/proximity_info_state_utils.cpp
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <cmath>
 #include <sstream> // for debug prints
 #include <vector>
 
@@ -37,7 +38,7 @@
         if (times) {
             for (int i = 0; i < inputSize; ++i) {
                 AKLOGI("(%d) x %d, y %d, time %d",
-                        i, xCoordinates[i], yCoordinates[i], times[i]);
+                        i, inputXCoordinates[i], inputYCoordinates[i], times[i]);
             }
         }
     }
@@ -45,7 +46,10 @@
     if (times) {
         for (int i = 0; i < inputSize; ++i) {
             if (i > 0) {
-                ASSERT(times[i] >= times[i - 1]);
+                if (times[i] < times[i - 1]) {
+                    AKLOGI("Invalid time sequence. %d, %d", times[i], times[i - 1]);
+                    ASSERT(false);
+                }
             }
         }
     }
@@ -469,11 +473,10 @@
                 // This point is not used because it's too close to the previous point.
                 if (DEBUG_GEO_FULL) {
                     AKLOGI("p0: size = %zd, x = %d, y = %d, lx = %d, ly = %d, dist = %d, "
-                           "width = %d", size, x, y, mSampledInputXs.back(),
-                           mSampledInputYs.back(), ProximityInfoUtils::getDistanceInt(
-                                   x, y, mSampledInputXs.back(), mSampledInputYs.back()),
-                           mProximityInfo->getMostCommonKeyWidth()
-                                   / LAST_POINT_SKIP_DISTANCE_SCALE);
+                           "width = %d", size, x, y, sampledInputXs->back(),
+                           sampledInputYs->back(), getDistanceInt(
+                                   x, y, sampledInputXs->back(), sampledInputYs->back()),
+                           mostCommonKeyWidth / LAST_POINT_SKIP_DISTANCE_SCALE);
                 }
                 return popped;
             }
@@ -511,11 +514,12 @@
         const float averageSpeed, const int id, const int inputSize, const int *const xCoordinates,
         const int *const yCoordinates, const int *times, const int sampledInputSize,
         const std::vector<int> *const sampledInputXs,
-        const std::vector<int> *const sampledInputYs, const std::vector<int> *const inputIndice) {
+        const std::vector<int> *const sampledInputYs,
+        const std::vector<int> *const sampledInputIndices) {
     if (sampledInputSize <= 0 || averageSpeed < 0.001f) {
         if (DEBUG_SAMPLING_POINTS) {
             AKLOGI("--- invalid state: cancel. size = %d, ave = %f",
-                    mSampledInputSize, mAverageSpeed);
+                    sampledInputSize, averageSpeed);
         }
         return 1.0f;
     }
@@ -523,7 +527,7 @@
             * ProximityInfoParams::LOOKUP_RADIUS_PERCENTILE / MAX_PERCENTILE;
     const int x0 = (*sampledInputXs)[id];
     const int y0 = (*sampledInputYs)[id];
-    const int actualInputIndex = (*inputIndice)[id];
+    const int actualInputIndex = (*sampledInputIndices)[id];
     int tempTime = 0;
     int tempBeelineDistance = 0;
     int start = actualInputIndex;
@@ -582,10 +586,10 @@
     if (DEBUG_DOUBLE_LETTER) {
         AKLOGI("--- (%d, %d) double letter: start = %d, end = %d, dist = %d, time = %d,"
                 " speed = %f, ave = %f, val = %f, start time = %d, end time = %d",
-                id, mSampledInputIndice[id], start, end, beelineDistance, time,
-                (static_cast<float>(beelineDistance) / static_cast<float>(time)), mAverageSpeed,
+                id, (*sampledInputIndices)[id], start, end, beelineDistance, time,
+                (static_cast<float>(beelineDistance) / static_cast<float>(time)), averageSpeed,
                 ((static_cast<float>(beelineDistance) / static_cast<float>(time))
-                        / mAverageSpeed), adjustedStartTime, adjustedEndTime);
+                        / averageSpeed), adjustedStartTime, adjustedEndTime);
     }
     // Offset 1%
     // TODO: Detect double letter more smartly
@@ -850,7 +854,6 @@
         }
     }
 
-
     if (DEBUG_POINTS_PROBABILITY) {
         for (int i = 0; i < sampledInputSize; ++i) {
             std::stringstream sstream;
@@ -916,6 +919,29 @@
     }
 }
 
+/* static */ void ProximityInfoStateUtils::updateSearchKeysVector(
+        const ProximityInfo *const proximityInfo, const int sampledInputSize,
+        const int lastSavedInputSize,
+        const std::vector<int> *const sampledLengthCache,
+        const std::vector<NearKeycodesSet> *const nearKeysVector,
+        std::vector<NearKeycodesSet> *searchKeysVector) {
+    const int readForwordLength = static_cast<int>(
+            hypotf(proximityInfo->getKeyboardWidth(), proximityInfo->getKeyboardHeight())
+                    * ProximityInfoParams::SEARCH_KEY_RADIUS_RATIO);
+    for (int i = 0; i < sampledInputSize; ++i) {
+        if (i >= lastSavedInputSize) {
+            (*searchKeysVector)[i].reset();
+        }
+        for (int j = max(i, lastSavedInputSize); j < sampledInputSize; ++j) {
+            // TODO: Investigate if this is required. This may not fail.
+            if ((*sampledLengthCache)[j] - (*sampledLengthCache)[i] >= readForwordLength) {
+                break;
+            }
+            (*searchKeysVector)[i] |= (*nearKeysVector)[j];
+        }
+    }
+}
+
 // Decreases char probabilities of index0 by checking probabilities of a near point (index1) and
 // increases char probabilities of index1 by checking probabilities of index0.
 /* static */ bool ProximityInfoStateUtils::suppressCharProbabilities(const int mostCommonKeyWidth,
@@ -963,12 +989,13 @@
         const int *const inputXCoordinates, const int *const inputYCoordinates,
         const int sampledInputSize, const std::vector<int> *const sampledInputXs,
         const std::vector<int> *const sampledInputYs,
+        const std::vector<int> *const sampledTimes,
         const std::vector<float> *const sampledSpeedRates,
         const std::vector<int> *const sampledBeelineSpeedPercentiles) {
     if (DEBUG_GEO_FULL) {
         for (int i = 0; i < sampledInputSize; ++i) {
-            AKLOGI("Sampled(%d): x = %d, y = %d, time = %d", i, mSampledInputXs[i],
-                    mSampledInputYs[i], mSampledTimes ? mSampledTimes[i], -1);
+            AKLOGI("Sampled(%d): x = %d, y = %d, time = %d", i, (*sampledInputXs)[i],
+                    (*sampledInputYs)[i], sampledTimes ? (*sampledTimes)[i] : -1);
         }
     }
 
@@ -985,8 +1012,8 @@
     for (int i = 0; i < sampledInputSize; ++i) {
         if (isGeometric) {
             AKLOGI("%d: x = %d, y = %d, time = %d, relative speed = %.4f, beeline speed = %d",
-                    i, mSampledInputXs[i], mSampledInputYs[i], mSampledTimes[i], mSpeedRates[i],
-                    getBeelineSpeedPercentile(i));
+                    i, (*sampledInputXs)[i], (*sampledInputYs)[i], (*sampledTimes)[i],
+                    (*sampledSpeedRates)[i], (*sampledBeelineSpeedPercentiles)[i]);
         }
         sampledX << (*sampledInputXs)[i];
         sampledY << (*sampledInputYs)[i];
diff --git a/native/jni/src/proximity_info_state_utils.h b/native/jni/src/proximity_info_state_utils.h
index b67f82c..b70121a 100644
--- a/native/jni/src/proximity_info_state_utils.h
+++ b/native/jni/src/proximity_info_state_utils.h
@@ -72,6 +72,12 @@
             const std::vector<float> *const distanceCache_G,
             std::vector<NearKeycodesSet> *nearKeysVector,
             std::vector<hash_map_compat<int, float> > *charProbabilities);
+    static void updateSearchKeysVector(
+            const ProximityInfo *const proximityInfo, const int sampledInputSize,
+            const int lastSavedInputSize,
+            const std::vector<int> *const sampledLengthCache,
+            const std::vector<NearKeycodesSet> *const nearKeysVector,
+            std::vector<NearKeycodesSet> *searchKeysVector);
     static float getPointToKeyByIdLength(const float maxPointToKeyLength,
             const std::vector<float> *const distanceCache_G, const int keyCount,
             const int inputIndex, const int keyId, const float scale);
@@ -99,6 +105,7 @@
             const int *const inputXCoordinates, const int *const inputYCoordinates,
             const int sampledInputSize, const std::vector<int> *const sampledInputXs,
             const std::vector<int> *const sampledInputYs,
+            const std::vector<int> *const sampledTimes,
             const std::vector<float> *const sampledSpeedRates,
             const std::vector<int> *const sampledBeelineSpeedPercentiles);
  private: