LI: inject Rust language to doctests

src/main/kotlin/org/rust/cargo/project/configurable/RsProjectConfigurable.kt

@@ -36,6 +36,9 @@ class RsProjectConfigurable(
     private val showTestToolWindowCheckbox: JBCheckBox = JBCheckBox()
     private var showTestToolWindow: Boolean by CheckboxDelegate(showTestToolWindowCheckbox)
 
+    private val doctestInjectionCheckbox: JBCheckBox = JBCheckBox()
+    private var doctestInjectionEnabled: Boolean by CheckboxDelegate(doctestInjectionCheckbox)
+
     private val hintProvider = InlayParameterHintsExtension.forLanguage(RsLanguage)
     private val hintCheckboxes: Map<String, JBCheckBox> =
         hintProvider.supportedOptions.associate { it.id to JBCheckBox() }
@@ -52,6 +55,7 @@ class RsProjectConfigurable(
             Show test results in run tool window when testing session begins
             instead of raw console.
         """)
+        row("Inject Rust language to documentation comments:", doctestInjectionCheckbox)
         val supportedHintOptions = hintProvider.supportedOptions
         if (supportedHintOptions.isNotEmpty()) {
             block("Hints") {
@@ -73,6 +77,7 @@ class RsProjectConfigurable(
         )
         expandMacros = settings.expandMacros
         showTestToolWindow = settings.showTestToolWindow
+        doctestInjectionEnabled = settings.doctestInjectionEnabled
 
         for (option in hintProvider.supportedOptions) {
             checkboxForOption(option).isSelected = option.get()
@@ -92,7 +97,8 @@ class RsProjectConfigurable(
             toolchain = rustProjectSettings.data.toolchain,
             explicitPathToStdlib = rustProjectSettings.data.explicitPathToStdlib,
             expandMacros = expandMacros,
-            showTestToolWindow = showTestToolWindow
+            showTestToolWindow = showTestToolWindow,
+            doctestInjectionEnabled = doctestInjectionEnabled
         )
     }
 
@@ -103,6 +109,7 @@ class RsProjectConfigurable(
             || data.explicitPathToStdlib != settings.explicitPathToStdlib
             || expandMacros != settings.expandMacros
             || showTestToolWindow != settings.showTestToolWindow
+            || doctestInjectionEnabled != settings.doctestInjectionEnabled
     }
 
     private fun checkboxForOption(opt: Option): JBCheckBox = hintCheckboxes[opt.id]!!

src/main/kotlin/org/rust/cargo/project/settings/RustProjectSettingsService.kt

@@ -25,6 +25,7 @@ interface RustProjectSettingsService {
         val useOffline: Boolean,
         val expandMacros: Boolean,
         val showTestToolWindow: Boolean,
+        val doctestInjectionEnabled: Boolean,
         val useSkipChildren: Boolean
     )
 
@@ -44,6 +45,7 @@ interface RustProjectSettingsService {
     val useOffline: Boolean get() = data.useOffline
     val expandMacros: Boolean get() = data.expandMacros
     val showTestToolWindow: Boolean get() = data.showTestToolWindow
+    val doctestInjectionEnabled: Boolean get() = data.doctestInjectionEnabled
     val useSkipChildren: Boolean get() = data.useSkipChildren
 
     /*

src/main/kotlin/org/rust/cargo/project/settings/impl/RustProjectSettingsServiceImpl.kt

@@ -33,6 +33,7 @@ class RustProjectSettingsServiceImpl(
         var useOffline: Boolean = false,
         var expandMacros: Boolean = true,
         var showTestToolWindow: Boolean = true,
+        var doctestInjectionEnabled: Boolean = true,
         var useSkipChildren: Boolean = false
     )
 
@@ -60,6 +61,7 @@ class RustProjectSettingsServiceImpl(
                 useOffline = state.useOffline,
                 expandMacros = state.expandMacros,
                 showTestToolWindow = state.showTestToolWindow,
+                doctestInjectionEnabled = state.doctestInjectionEnabled,
                 useSkipChildren = state.useSkipChildren
             )
         }
@@ -75,6 +77,7 @@ class RustProjectSettingsServiceImpl(
                 useOffline = value.useOffline,
                 expandMacros = value.expandMacros,
                 showTestToolWindow = value.showTestToolWindow,
+                doctestInjectionEnabled = value.doctestInjectionEnabled,
                 useSkipChildren = value.useSkipChildren
             )
             if (state != newState) {

src/main/kotlin/org/rust/ide/actions/RsEnterHandler.kt

@@ -0,0 +1,47 @@
+/*
+ * Use of this source code is governed by the MIT license that can be
+ * found in the LICENSE file.
+ */
+
+package org.rust.ide.actions
+
+import com.intellij.codeInsight.editorActions.EnterHandler
+import com.intellij.injected.editor.EditorWindow
+import com.intellij.openapi.actionSystem.DataContext
+import com.intellij.openapi.editor.Caret
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.actionSystem.EditorActionHandler
+import com.intellij.openapi.editor.actions.EnterAction
+import com.intellij.psi.impl.source.tree.injected.InjectedCaret
+import org.rust.ide.injected.RsDoctestLanguageInjector
+import org.rust.ide.injected.isDoctestInjection
+import org.rust.lang.core.psi.RsFile
+
+/**
+ * This class is used to handle typing enter inside doctest language injection (see [RsDoctestLanguageInjector]).
+ * Enter handlers are piped:
+ * [EnterHandler] -> [RsEnterHandler] -> --------------------> [EnterAction.Handler]
+ *  |                 |             \ -> [EnterHandler] -> /    ^ just insert new line [originalHandler]
+ *  |                 this class          ^ (the case of injected psi) [injectionEnterHandler]
+ *  front platform handler (handles indents and other complex stuff)
+ */
+class RsEnterHandler(private val originalHandler: EditorActionHandler) : EditorActionHandler() {
+    private val injectionEnterHandler = EnterHandler(object : EditorActionHandler() {
+        override fun doExecute(editor: Editor, caret: Caret?, dataContext: DataContext) {
+            originalHandler.execute(editor, caret, dataContext)
+        }
+    })
+
+    public override fun isEnabledForCaret(editor: Editor, caret: Caret, dataContext: DataContext): Boolean {
+        return originalHandler.isEnabled(editor, caret, dataContext)
+    }
+
+    override fun doExecute(editor: Editor, caret: Caret?, dataContext: DataContext) {
+        if (editor is EditorWindow && caret is InjectedCaret &&
+            (editor.injectedFile as? RsFile)?.isDoctestInjection == true) {
+            injectionEnterHandler.execute(editor.delegate, caret.delegate, dataContext)
+        } else {
+            originalHandler.execute(editor, caret, dataContext)
+        }
+    }
+}

src/main/kotlin/org/rust/ide/annotator/RsDoctestAnnotator.kt

@@ -0,0 +1,41 @@
+/*
+ * Use of this source code is governed by the MIT license that can be
+ * found in the LICENSE file.
+ */
+
+package org.rust.ide.annotator
+
+import com.intellij.lang.annotation.AnnotationHolder
+import com.intellij.lang.annotation.HighlightSeverity
+import com.intellij.openapi.editor.colors.EditorColors
+import com.intellij.psi.PsiElement
+import com.intellij.psi.impl.source.tree.injected.InjectionBackgroundSuppressor
+import org.rust.cargo.project.settings.rustSettings
+import org.rust.ide.injected.RsDoctestLanguageInjector
+import org.rust.ide.injected.findDoctestInjectableRanges
+import org.rust.lang.core.psi.RsDocCommentImpl
+import org.rust.lang.core.psi.ext.RsElement
+import org.rust.lang.core.psi.ext.ancestorStrict
+import org.rust.lang.core.psi.ext.containingCargoTarget
+
+/**
+ * Adds missing background for injections from [RsDoctestLanguageInjector].
+ * Background is disabled by [InjectionBackgroundSuppressor] marker implemented for [RsDocCommentImpl].
+ *
+ * We have to do it this way because we want to highlight fully range inside ```backticks```
+ * but a real injections is shifted by 1 character and empty lines are skipped.
+ */
+class RsDoctestAnnotator : RsAnnotatorBase() {
+    override fun annotateInternal(element: PsiElement, holder: AnnotationHolder) {
+        if (element !is RsDocCommentImpl) return
+        if (!element.project.rustSettings.doctestInjectionEnabled) return
+        // only library targets can have doctests
+        if (element.ancestorStrict<RsElement>()?.containingCargoTarget?.isLib != true) return
+
+        val startOffset = element.startOffset
+        findDoctestInjectableRanges(element).flatten().forEach {
+            holder.createAnnotation(HighlightSeverity.INFORMATION, it.shiftRight(startOffset), null)
+                .textAttributes = EditorColors.INJECTED_LANGUAGE_FRAGMENT
+        }
+    }
+}

src/main/kotlin/org/rust/ide/injected/RsDoctestLanguageInjector.kt

@@ -0,0 +1,182 @@
+/*
+ * Use of this source code is governed by the MIT license that can be
+ * found in the LICENSE file.
+ */
+
+package org.rust.ide.injected
+
+import com.intellij.injected.editor.VirtualFileWindow
+import com.intellij.lang.injection.MultiHostInjector
+import com.intellij.lang.injection.MultiHostRegistrar
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.TextRange
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiElement
+import com.intellij.psi.tree.IElementType
+import com.intellij.util.text.CharArrayUtil
+import org.rust.cargo.project.settings.rustSettings
+import org.rust.cargo.project.workspace.PackageOrigin
+import org.rust.cargo.util.AutoInjectedCrates
+import org.rust.lang.RsLanguage
+import org.rust.lang.core.psi.RS_DOC_COMMENTS
+import org.rust.lang.core.psi.RsDocCommentImpl
+import org.rust.lang.core.psi.RsFile
+import org.rust.lang.core.psi.ext.*
+import org.rust.lang.doc.psi.RsDocKind
+import org.rust.openapiext.toPsiFile
+import org.rust.stdext.nextOrNull
+import java.util.regex.Pattern
+
+// See https://github.com/rust-lang/rust/blob/5182cc1c/src/librustdoc/html/markdown.rs#L646
+private val LANG_SPLIT_REGEX = Pattern.compile("[^\\w-]+", Pattern.UNICODE_CHARACTER_CLASS)
+private val RUST_LANG_ALIASES = listOf(
+    "rust",
+    "allow_fail",
+    "should_panic",
+    "no_run",
+    "test_harness",
+//    "compile_fail", // don't highlight code that is expected to contain errors
+    "edition2018",
+    "edition2015"
+)
+
+class RsDoctestLanguageInjector : MultiHostInjector {
+    private data class CodeRange(val start: Int, val end: Int, val codeStart: Int) {
+        fun isCodeNotEmpty(): Boolean = codeStart + 1 < end
+
+        val indent: Int = codeStart - start
+
+        fun offsetIndent(indent: Int): CodeRange? =
+            if (start + indent < end) CodeRange(start + indent, end, codeStart) else null
+    }
+
+    override fun elementsToInjectIn(): List<Class<out PsiElement>> =
+        listOf(RsDocCommentImpl::class.java)
+
+    override fun getLanguagesToInject(registrar: MultiHostRegistrar, context: PsiElement) {
+        if (context !is RsDocCommentImpl) return
+        if (!context.isValidHost || context.elementType !in RS_DOC_COMMENTS) return
+        if (!context.project.rustSettings.doctestInjectionEnabled) return
+
+        val rsElement = context.ancestorStrict<RsElement>() ?: return
+        val cargoTarget = rsElement.containingCargoTarget ?: return
+        if (!cargoTarget.isLib) return // only library targets can have doctests
+        val crateName = cargoTarget.normName
+        val text = context.text
+
+        findDoctestInjectableRanges(text, context.elementType).map { ranges ->
+            ranges.map {
+                CodeRange(
+                    it.startOffset,
+                    it.endOffset,
+                    CharArrayUtil.shiftForward(text, it.startOffset, it.endOffset, " \t")
+                )
+            }
+        }.map { ranges ->
+            val commonIndent = ranges.filter { it.isCodeNotEmpty() }.map { it.indent }.min()
+            val indentedRanges = if (commonIndent != null) ranges.mapNotNull { it.offsetIndent(commonIndent) } else ranges
+
+            indentedRanges.map { (start, end, codeStart) ->
+                // `cargo doc` has special rules for code lines which start with `#`:
+                //   * `# ` prefix is used to mark lines that should be skipped in rendered documentation.
+                //   * `##` prefix is converted to `#`
+                // See https://github.com/rust-lang/rust/blob/5182cc1c/src/librustdoc/html/markdown.rs#L114
+                when {
+                    text.startsWith("##", codeStart) -> TextRange(codeStart + 1, end)
+                    text.startsWith("# ", codeStart) -> TextRange(codeStart + 2, end)
+                    else -> TextRange(start, end)
+                }
+            }
+        }.forEach { ranges ->
+            val inj = registrar.startInjecting(RsLanguage)
+
+            ranges.forEachIndexed { index, range ->
+                val isFirstIteration = index == 0
+                val isLastIteration = index == ranges.size - 1
+
+                val prefix = if (isFirstIteration) {
+                    buildString {
+                        // Yes, we want to skip the only "std" crate. Not core/alloc/etc, the "std" only
+                        val isStdCrate = crateName == AutoInjectedCrates.STD &&
+                            cargoTarget.pkg.origin == PackageOrigin.STDLIB
+                        if (!isStdCrate) {
+                            append("extern crate ")
+                            append(crateName)
+                            append("; ")
+                        }
+                        append("fn main() {")
+                    }
+                } else {
+                    null
+                }
+                val suffix = if (isLastIteration) "}" else null
+
+                inj.addPlace(prefix, suffix, context, range)
+            }
+
+            inj.doneInjecting()
+        }
+    }
+}
+
+fun findDoctestInjectableRanges(comment: RsDocCommentImpl): Sequence<List<TextRange>> =
+    findDoctestInjectableRanges(comment.text, comment.elementType)
+
+private fun findDoctestInjectableRanges(text: String, elementType: IElementType): Sequence<List<TextRange>> {
+    // TODO use markdown parser
+    val tripleBacktickIndices = text.indicesOf("```").toList()
+    if (tripleBacktickIndices.size < 2) return emptySequence() // no code blocks in the comment
+
+    val infix = RsDocKind.of(elementType).infix
+
+    return tripleBacktickIndices.asSequence().chunked(2).mapNotNull { idx ->
+        // Contains code lines inside backticks including `///` at the start and `\n` at the end.
+        // It doesn't contain the last line with /// ```
+        val lines = run {
+            val codeBlockStart = idx[0] + 3 // skip ```
+            val codeBlockEnd = idx.getOrNull(1) ?: return@mapNotNull null
+            generateSequence(codeBlockStart) { text.indexOf("\n", it) + 1 }
+                .takeWhile { it != 0 && it < codeBlockEnd }
+                .zipWithNext()
+                .iterator()
+        }
+
+        // ```rust, should_panic, edition2018
+        //     ^ this text
+        val lang = lines.nextOrNull()?.let { text.substring(it.first, it.second - 1) } ?: return@mapNotNull null
+        if (lang.isNotEmpty()) {
+            val parts = lang.split(LANG_SPLIT_REGEX).filter { it.isNotBlank() }
+            if (parts.any { it !in RUST_LANG_ALIASES }) return@mapNotNull null
+        }
+
+        // skip doc comment infix (`///`, `//!` or ` * `)
+        val ranges = lines.asSequence().mapNotNull { (lineStart, lineEnd) ->
+            val index = text.indexOf(infix, lineStart)
+            if (index != -1 && index < lineEnd) {
+                val start = index + infix.length
+                TextRange(start, lineEnd)
+            } else {
+                null
+            }
+        }.toList()
+
+        if (ranges.isEmpty()) return@mapNotNull null
+        ranges
+    }
+}
+
+private fun String.indicesOf(s: String): Sequence<Int> =
+    generateSequence(indexOf(s)) { indexOf(s, it + s.length) }.takeWhile { it != -1 }
+
+fun VirtualFile.isDoctestInjection(project: Project): Boolean {
+    val virtualFileWindow = this as? VirtualFileWindow ?: return false
+    val hostFile = virtualFileWindow.delegate.toPsiFile(project) as? RsFile ?: return false
+    val hostElement = hostFile.findElementAt(virtualFileWindow.documentWindow.injectedToHost(0)) ?: return false
+    return hostElement.elementType in RS_DOC_COMMENTS
+}
+
+val RsFile.isDoctestInjection: Boolean
+    get() = virtualFile?.isDoctestInjection(project) == true
+
+val RsElement.isDoctestInjection: Boolean
+    get() = (contextualFile as? RsFile)?.isDoctestInjection == true

src/main/kotlin/org/rust/ide/injected/RsSimpleMultiLineEscaper.kt

@@ -0,0 +1,24 @@
+/*
+ * Use of this source code is governed by the MIT license that can be
+ * found in the LICENSE file.
+ */
+
+package org.rust.ide.injected
+
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.LiteralTextEscaper
+import com.intellij.psi.PsiLanguageInjectionHost
+
+/** Same as [com.intellij.psi.LiteralTextEscaper.createSimple], but multi line */
+class RsSimpleMultiLineEscaper<T: PsiLanguageInjectionHost>(host: T) : LiteralTextEscaper<T>(host) {
+    override fun decode(rangeInsideHost: TextRange, outChars: java.lang.StringBuilder): Boolean {
+        outChars.append(rangeInsideHost.substring(myHost.text))
+        return true
+    }
+
+    override fun getOffsetInHost(offsetInDecoded: Int, rangeInsideHost: TextRange): Int {
+        return rangeInsideHost.startOffset + offsetInDecoded
+    }
+
+    override fun isOneLine(): Boolean = false
+}