From 49b964a3ac69c8b03d554094d1baba9f2bf4c2e2 Mon Sep 17 00:00:00 2001
From: Alexey Kalina <alexey.kalina@jetbrains.com>
Date: Wed, 20 Mar 2024 09:57:22 +0100
Subject: [PATCH] [evaluation-plugin] check correctness of collected tokens
 after traversing step (ML-2950)

GitOrigin-RevId: aa1a2d2a685f9fc9c1441303406d6cd74bd0619c
---
 .../cce/visitor/LineCompletionVisitorHelper.kt    | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/plugins/evaluation-plugin/languages/src/com/intellij/cce/visitor/LineCompletionVisitorHelper.kt b/plugins/evaluation-plugin/languages/src/com/intellij/cce/visitor/LineCompletionVisitorHelper.kt
index 3c39dfb1a383..b31efe6ab00a 100644
--- a/plugins/evaluation-plugin/languages/src/com/intellij/cce/visitor/LineCompletionVisitorHelper.kt
+++ b/plugins/evaluation-plugin/languages/src/com/intellij/cce/visitor/LineCompletionVisitorHelper.kt
@@ -1,4 +1,4 @@
-// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.cce.visitor
 
 import com.intellij.cce.core.CodeFragment
@@ -17,6 +17,7 @@ class LineCompletionVisitorHelper {
   fun getFile(): CodeFragment {
     codeFragment?.let { file ->
       lines.filter { it.getChildren().isNotEmpty() }.forEach { file.addChild(it) }
+      file.validateCorrectness()
       return file
     }
     throw PsiConverterException("Invoke 'accept' with visitor on PSI first")
@@ -40,6 +41,18 @@ class LineCompletionVisitorHelper {
         ?.addChild(CodeToken(text, element.startOffset))
     }
   }
+
+  private fun CodeFragment.validateCorrectness() {
+    var lastEndOffset = 0
+    for (line in getChildren()) {
+      assert(line is CodeLine) { "Code fragment should only contain code lines" }
+      val tokens = (line as CodeLine).getChildren()
+      for (token in tokens) {
+        assert(lastEndOffset <= token.offset) { "Code tokens shouldn't overlap" }
+        lastEndOffset = token.offset + token.text.length
+      }
+    }
+  }
 }
 
 private const val MAX_PREFIX_LENGTH: Int = 4