diff --git a/app/src/main/java/com/itsaky/androidide/language/incremental/BaseIncrementalAnalyzeManager.java b/app/src/main/java/com/itsaky/androidide/language/incremental/BaseIncrementalAnalyzeManager.java
index 3e5b43136..b23d05563 100644
--- a/app/src/main/java/com/itsaky/androidide/language/incremental/BaseIncrementalAnalyzeManager.java
+++ b/app/src/main/java/com/itsaky/androidide/language/incremental/BaseIncrementalAnalyzeManager.java
@@ -209,17 +209,19 @@ public abstract class BaseIncrementalAnalyzeManager
     var newState = 0;
     var stateObj = new LineState();
     if (state.state == NORMAL) {
-      newState = tokenizeNormal(line, 0, tokens, stateObj);
+      newState = tokenizeNormal(line, 0, tokens, stateObj, state.lexerMode);
     } else if (state.state == INCOMPLETE) {
-      final var result = fillIncomplete(line, tokens);
+      final var result = fillIncomplete(line, tokens, state.lexerMode);
       newState = IntPair.getFirst(result);
       if (newState == NORMAL) {
-        newState = tokenizeNormal(line, IntPair.getSecond(result), tokens, stateObj);
+        newState =
+            tokenizeNormal(line, IntPair.getSecond(result), tokens, stateObj, state.lexerMode);
       } else {
         newState = INCOMPLETE;
       }
     }
     stateObj.state = newState;
+    stateObj.lexerMode = lexer._mode;
     return new LineTokenizeResult<>(stateObj, tokens);
   }
 
@@ -239,8 +241,10 @@ public abstract class BaseIncrementalAnalyzeManager
       final CharSequence line,
       final int column,
       final List<IncrementalToken> tokens,
-      final LineState st) {
+      final LineState st,
+      final int lexerMode) {
     lexer.setInputStream(createStream(line));
+    lexer.pushMode(lexerMode);
     final var queues = createEvictingQueueForTokens();
     final var start = queues.getFirst();
     final var end = queues.getSecond();
@@ -310,8 +314,10 @@ public abstract class BaseIncrementalAnalyzeManager
    * @return The state and offset.
    */
   @SuppressWarnings("UnstableApiUsage")
-  protected long fillIncomplete(CharSequence line, final List<IncrementalToken> tokens) {
+  protected long fillIncomplete(
+      CharSequence line, final List<IncrementalToken> tokens, final int lexerMode) {
     lexer.setInputStream(createStream(line));
+    lexer.pushMode(lexerMode);
     final var queue = createEvictingQueueForTokens();
     final var end = queue.getSecond();
     final var allTokens =
diff --git a/app/src/main/java/com/itsaky/androidide/language/incremental/LineState.kt b/app/src/main/java/com/itsaky/androidide/language/incremental/LineState.kt
index 496432095..f0854e4b1 100644
--- a/app/src/main/java/com/itsaky/androidide/language/incremental/LineState.kt
+++ b/app/src/main/java/com/itsaky/androidide/language/incremental/LineState.kt
@@ -17,50 +17,24 @@
 
 package com.itsaky.androidide.language.incremental
 
-import com.itsaky.lsp.models.DiagnosticItem
+import org.antlr.v4.runtime.Lexer
 
 /**
  * Tokenization state of a line.
  *
+ * @param state The state of the line.
+ * @param hasBraces `true` if the line has braces. `false` otherwise.
+ * @param lexerMode The mode of the lexer. This MUST be preserved in the lexer.
+ *
  * @author Akash Yadav
  */
-class LineState {
-
+data class LineState(
+  @JvmField var state: Int = NORMAL,
+  @JvmField var hasBraces: Boolean = false,
+  @JvmField var lexerMode: Int = Lexer.DEFAULT_MODE
+) {
   companion object {
     const val NORMAL = 0
     const val INCOMPLETE = 1
   }
-
-  @JvmField var state = NORMAL
-  @JvmField var hasBraces = false
-  @JvmField var diagnostics: DiagnosticsState? = null
-
-  override fun equals(other: Any?): Boolean {
-    if (this === other) return true
-    if (javaClass != other?.javaClass) return false
-
-    other as LineState
-
-    if (state != other.state) return false
-    if (hasBraces != other.hasBraces) return false
-    if (diagnostics != other.diagnostics) return false
-
-    return true
-  }
-
-  override fun hashCode(): Int {
-    var result = state
-    result = 31 * result + hasBraces.hashCode()
-    result = 31 * result + (diagnostics?.hashCode() ?: 0)
-    return result
-  }
-
-  override fun toString(): String {
-    return "LineState(state=$state, hasBraces=$hasBraces, diagnostics=$diagnostics)"
-  }
 }
-
-data class DiagnosticsState(
-  val continueDiagnostic: Boolean = false,
-  val diagnostic: DiagnosticItem? = null
-)
diff --git a/build-tools/tooling-api-impl/src/test/test-home/.androidide/init/model.jar b/build-tools/tooling-api-impl/src/test/test-home/.androidide/init/model.jar
index 2e54e4583..b9ea432ae 100644
Binary files a/build-tools/tooling-api-impl/src/test/test-home/.androidide/init/model.jar and b/build-tools/tooling-api-impl/src/test/test-home/.androidide/init/model.jar differ