[DS-5493] add field NotebookCellLines.Interval.data: KeyFMap

Different notebook implementation could store their own values in intervals.

GitOrigin-RevId: e1f666d746eb80179f5f45d590d613cceebe225d

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/NotebookCellLines.kt

@@ -4,7 +4,9 @@ import com.intellij.lang.Language
 import com.intellij.openapi.actionSystem.DataKey
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.util.Key
 import com.intellij.util.EventDispatcher
+import com.intellij.util.keyFMap.KeyFMap
 import java.util.*
 
 val NOTEBOOK_CELL_LINES_INTERVAL_DATA_KEY = DataKey.create<NotebookCellLines.Interval>("NOTEBOOK_CELL_LINES_INTERVAL")
@@ -36,8 +38,12 @@ interface NotebookCellLines {
     val type: CellType,
     val lines: IntRange,
     val markers: MarkersAtLines,
-    val language: Language,
+    val data: KeyFMap, // different notebook implementations could store their own values in this map
   ) : Comparable<Interval> {
+    val language: Language = data.get(INTERVAL_LANGUAGE_KEY)!!
+
+    operator fun <V> get(key: Key<V>): V? = data.get(key)
+
     override fun compareTo(other: Interval): Int = lines.first - other.lines.first
   }
 
@@ -78,5 +84,7 @@ interface NotebookCellLines {
 
     fun hasSupport(editor: Editor): Boolean =
       hasSupport(editor.document)
+
+    val INTERVAL_LANGUAGE_KEY = Key.create<Language>("org.jetbrains.plugins.notebooks.visualization.NotebookCellLines.Interval.language")
   }
 }

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/NotebookCellLinesLexer.kt

@@ -5,6 +5,7 @@ import com.intellij.lexer.Lexer
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.fileTypes.PlainTextLanguage
 import com.intellij.psi.tree.IElementType
+import com.intellij.util.keyFMap.KeyFMap
 import org.jetbrains.plugins.notebooks.visualization.NotebookCellLines.CellType
 import org.jetbrains.plugins.notebooks.visualization.NotebookCellLines.MarkersAtLines
 import kotlin.math.max
@@ -17,14 +18,14 @@ interface NotebookCellLinesLexer {
     val type: CellType,
     val offset: Int,
     val length: Int,
-    val language: Language,
+    val data: KeyFMap,
   ) : Comparable<Marker> {
     override fun compareTo(other: Marker): Int = offset - other.offset
   }
 
   companion object {
     fun defaultMarkerSequence(underlyingLexerFactory: () -> Lexer,
-                              getCellLanguageAndType: (IElementType, lexer: Lexer) -> Pair<Language, CellType>?,
+                              getCellTypeAndData: (IElementType, lexer: Lexer) -> Pair<CellType, KeyFMap>?,
                               chars: CharSequence,
                               ordinalIncrement: Int,
                               offsetIncrement: Int): Sequence<Marker> = sequence {
@@ -33,16 +34,13 @@ interface NotebookCellLinesLexer {
       var ordinal = 0
       while (true) {
         val tokenType = lexer.tokenType ?: break
-        val cellLanguageAndType = getCellLanguageAndType(tokenType, lexer)
-
-        if (cellLanguageAndType != null) {
-          val (language, type) = cellLanguageAndType
+        getCellTypeAndData(tokenType, lexer)?.let { (type, data) ->
           yield(Marker(
             ordinal = ordinal++ + ordinalIncrement,
             type = type,
             offset = lexer.currentPosition.offset + offsetIncrement,
             length = lexer.tokenText.length,
-            language = language,
+            data = data,
           ))
         }
         lexer.advance()
@@ -50,14 +48,15 @@ interface NotebookCellLinesLexer {
     }
 
     fun defaultIntervals(document: Document, markers: List<Marker>): List<NotebookCellLines.Interval> {
-      val intervals = toIntervalsInfo(document, markers)
+      val data = KeyFMap.EMPTY_MAP.plus(NotebookCellLines.INTERVAL_LANGUAGE_KEY, PlainTextLanguage.INSTANCE)
+      val intervals = toIntervalsInfo(document, markers, firstMarkerData = data, lastMarkerData = data)
 
       val result = mutableListOf<NotebookCellLines.Interval>()
       for (i in 0 until (intervals.size - 1)) {
         result += NotebookCellLines.Interval(ordinal = i, type = intervals[i].cellType,
                                              lines = intervals[i].lineNumber until intervals[i + 1].lineNumber,
                                              markers = intervals[i].markersAtLInes,
-                                             intervals[i].language)
+                                             intervals[i].data)
       }
       return result
     }
@@ -65,22 +64,24 @@ interface NotebookCellLinesLexer {
   }
 }
 
-private data class IntervalInfo(val lineNumber: Int, val cellType: CellType, val markersAtLInes: MarkersAtLines, val language: Language)
+private data class IntervalInfo(val lineNumber: Int, val cellType: CellType, val markersAtLInes: MarkersAtLines, val data: KeyFMap)
 
 private fun toIntervalsInfo(document: Document,
-                            markers: List<NotebookCellLinesLexer.Marker>): List<IntervalInfo> {
+                            markers: List<NotebookCellLinesLexer.Marker>,
+                            firstMarkerData: KeyFMap,
+                            lastMarkerData: KeyFMap): List<IntervalInfo> {
   val m = mutableListOf<IntervalInfo>()
 
   // add first if necessary
   if (markers.isEmpty() || document.getLineNumber(markers.first().offset) != 0) {
-    m += IntervalInfo(0, CellType.RAW, MarkersAtLines.NO, PlainTextLanguage.INSTANCE)
+    m += IntervalInfo(0, CellType.RAW, MarkersAtLines.NO, firstMarkerData)
   }
 
   for (marker in markers) {
-    m += IntervalInfo(document.getLineNumber(marker.offset), marker.type, MarkersAtLines.TOP, marker.language) // marker.language is provided in makeIntervals
+    m += IntervalInfo(document.getLineNumber(marker.offset), marker.type, MarkersAtLines.TOP, marker.data)
   }
 
   // marker for the end
-  m += IntervalInfo(max(document.lineCount, 1), CellType.RAW, MarkersAtLines.NO, PlainTextLanguage.INSTANCE)
+  m += IntervalInfo(max(document.lineCount, 1), CellType.RAW, MarkersAtLines.NO, lastMarkerData)
   return m
 }

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/NotebookVisualizationUiUtil.kt

@@ -1,6 +1,5 @@
 package org.jetbrains.plugins.notebooks.visualization
 
-import com.intellij.lang.Language
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.LogicalPosition
@@ -9,6 +8,7 @@ import com.intellij.openapi.editor.impl.EditorImpl
 import com.intellij.openapi.util.TextRange
 import com.intellij.util.SmartList
 import com.intellij.util.containers.ContainerUtil
+import com.intellij.util.keyFMap.KeyFMap
 import java.awt.Graphics
 import javax.swing.JComponent
 import kotlin.math.max
@@ -101,21 +101,21 @@ val NotebookCellLines.Interval.contentLines: IntRange
 fun makeMarkersFromIntervals(document: Document, intervals: Iterable<NotebookCellLines.Interval>): List<NotebookCellLinesLexer.Marker> {
   val markers = ArrayList<NotebookCellLinesLexer.Marker>()
 
-  fun addMarker(line: Int, type: NotebookCellLines.CellType, language: Language) {
+  fun addMarker(line: Int, type: NotebookCellLines.CellType, data: KeyFMap) {
     val startOffset = document.getLineStartOffset(line)
     val endOffset =
       if (line + 1 < document.lineCount) document.getLineStartOffset(line + 1)
       else document.getLineEndOffset(line)
     val length = endOffset - startOffset
-    markers.add(NotebookCellLinesLexer.Marker(markers.size, type, startOffset, length, language))
+    markers.add(NotebookCellLinesLexer.Marker(markers.size, type, startOffset, length, data))
   }
 
   for (interval in intervals) {
     if (interval.markers.hasTopLine) {
-      addMarker(interval.lines.first, interval.type, interval.language)
+      addMarker(interval.lines.first, interval.type, interval.data)
     }
     if (interval.markers.hasBottomLine) {
-      addMarker(interval.lines.last, interval.type, interval.language)
+      addMarker(interval.lines.last, interval.type, interval.data)
     }
   }
 

notebooks/visualization/test/org/jetbrains/plugins/notebooks/visualization/CodeCellLinesChecker.kt

@@ -2,6 +2,7 @@ package org.jetbrains.plugins.notebooks.visualization
 
 import com.intellij.lang.Language
 import com.intellij.openapi.editor.impl.EditorImpl
+import com.intellij.util.keyFMap.KeyFMap
 import org.assertj.core.api.Assertions.assertThat
 
 class CodeCellLinesChecker(private val description: String,
@@ -21,7 +22,7 @@ class CodeCellLinesChecker(private val description: String,
     fun marker(cellType: NotebookCellLines.CellType, offset: Int, length: Int, language: Language) {
       markers!!.add(
         NotebookCellLinesLexer.Marker(ordinal = markers!!.size + markersStartOrdinal, type = cellType, offset = offset, length = length,
-                                      language = language))
+                                      data = makeLanguageData(language)))
     }
   }
 
@@ -37,7 +38,7 @@ class CodeCellLinesChecker(private val description: String,
                  lines: IntRange,
                  markers: NotebookCellLines.MarkersAtLines,
                  language: Language) {
-      list += NotebookCellLines.Interval(list.size + startOrdinal, cellType, lines, markers, language)
+      list += NotebookCellLines.Interval(list.size + startOrdinal, cellType, lines, markers, makeLanguageData(language))
     }
 
     fun interval(cellType: NotebookCellLines.CellType, lines: IntRange, language: Language) {
@@ -149,4 +150,7 @@ class CodeCellLinesChecker(private val description: String,
         """.trimMargin("|||"))
       .isEqualTo(expectedIntervalListenerCalls.prettyListeners())
   }
-}
+}
+
+private fun makeLanguageData(language: Language) =
+  KeyFMap.EMPTY_MAP.plus(NotebookCellLines.INTERVAL_LANGUAGE_KEY, language)