Enable highlighting for commands

kshell-repl-api/src/main/kotlin/sparklin/kshell/repl/CodeAnalyser.kt

@@ -59,7 +59,7 @@ class CodeAnalyzer(environment: KotlinCoreEnvironment) {
         class Success(diagnostics: Diagnostics) : AnalyzerResult(diagnostics)
     }
 
-    fun doAnalyze(linePsi: KtFile, codeLine: CodeLine): AnalyzerResult {
+    fun doAnalyze(linePsi: KtFile): AnalyzerResult {
         trace.clearDiagnostics()
         scriptDeclarationFactory.setDelegateFactory(FileBasedDeclarationProviderFactory(resolveSession.storageManager, listOf(linePsi)))
         /*val context = */topDownAnalyzer.analyzeDeclarations(topDownAnalysisContext.topDownAnalysisMode, listOf(linePsi))

kshell-repl-api/src/main/kotlin/sparklin/kshell/repl/ReplApi.kt

@@ -85,14 +85,19 @@ open class ReplState(val lock: ReentrantReadWriteLock) {
     val history: MutableList<Snippet> = mutableListOf()
 }
 
-interface Code {
+interface SourceCode {
     fun mkFileName(): String
-    fun source(): String
+    fun nextPart(codePart: String): SourceCode
+    fun replace(code: String): SourceCode
+    val no: Int
+    val code: String
+    val part: Int
 }
 
-data class CodeLine(val no: Int, val code: String, val part: Int = 0): Code {
+data class CodeLine(override val no: Int, override val code: String, override val part: Int = 0): SourceCode {
     override fun mkFileName(): String = "Line_$no" + if (part != 0) "_$part" else ""
-    override fun source(): String = code
+    override fun nextPart(codePart: String): CodeLine = CodeLine(no, codePart, part + 1)
+    override fun replace(code: String): CodeLine = CodeLine(no, code, part)
 }
 
 data class CompilationData(val snippets: List<Snippet>, val classes: CompiledClasses)

kshell-repl-api/src/main/kotlin/sparklin/kshell/repl/ReplChecker.kt

@@ -46,16 +46,16 @@ class ReplChecker(
 
     private fun createDiagnosticHolder() = ConsoleDiagnosticMessageHolder()
 
-    fun check(state: ReplState, code: Code, isScript: Boolean): Result<CheckedCode, EvalError.CompileError> {
+    fun check(state: ReplState, code: SourceCode, isScript: Boolean): Result<CheckedCode, EvalError.CompileError> {
         state.lock.write {
             val fileName = code.mkFileName() + (if (isScript) ".kts" else ".kt")
             val virtualFile =
-                    LightVirtualFile(fileName, KotlinLanguage.INSTANCE, StringUtil.convertLineSeparators(code.source())).apply {
+                    LightVirtualFile(fileName, KotlinLanguage.INSTANCE, StringUtil.convertLineSeparators(code.code)).apply {
                         charset = CharsetToolkit.UTF8_CHARSET
                     }
 
             val psiFile: KtFile =  psiFileFactory.trySetupPsiForFile(virtualFile, KotlinLanguage.INSTANCE, true, false) as KtFile?
-                        ?: error("File not analyzed ${code.source()}")
+                        ?: error("File not analyzed ${code.code}")
 
             val errorHolder = createDiagnosticHolder()
 

kshell-repl-api/src/main/kotlin/sparklin/kshell/repl/ReplCompiler.kt

@@ -22,7 +22,7 @@ class ReplCompiler(disposable: Disposable,
 
     private val analyzerEngine = CodeAnalyzer(checker.environment)
 
-    fun compile(state: ReplState, codeLine: CodeLine, previousStage: List<Snippet> = listOf()): Result<CompilationData, EvalError.CompileError> {
+    fun compile(state: ReplState, codeLine: SourceCode, previousStage: List<Snippet> = listOf()): Result<CompilationData, EvalError.CompileError> {
         state.lock.write {
             val lineResult = checker.check(state, codeLine, true)
             val checkedLine = when (lineResult) {
@@ -52,13 +52,13 @@ class ReplCompiler(disposable: Disposable,
 
             val code = generateKotlinCodeFor(generatedClassname, import + previousStage, actualSnippets)
 
-            val result = checker.check(state, CodeLine(codeLine.no, code, codeLine.part), false)
+            val result = checker.check(state, codeLine.replace(code), false)
             val psiForObject = when (result) {
                 is Result.Success -> result.data.psiFile
                 else -> throw IllegalStateException("Should never happen")
             }
 
-            val analysisResult = analyzerEngine.doAnalyze(psiForObject, codeLine)
+            val analysisResult = analyzerEngine.doAnalyze(psiForObject)
 
             AnalyzerWithCompilerReport.reportDiagnostics(analysisResult.diagnostics, errorHolder)
 
@@ -117,7 +117,7 @@ class ReplCompiler(disposable: Disposable,
 
             return if (deferredSnippets.isNotEmpty()) {
                 val otherCode = deferredSnippets.joinToString(separator = "\n") { it.code() }
-                val otherResult = compile(state, CodeLine(codeLine.no, otherCode, codeLine.part + 1), actualSnippets)
+                val otherResult = compile(state, codeLine.nextPart(otherCode), actualSnippets)
 
                 when (otherResult) {
                     is Result.Error -> otherResult

kshell/src/main/kotlin/sparklin/kshell/Command.kt

@@ -1,5 +1,7 @@
 package sparklin.kshell
 
+import sparklin.kshell.org.jline.reader.Highlighter
+
 interface Command {
     fun execute(line: String)
 
@@ -10,11 +12,20 @@ interface Command {
     fun help(): String
 
     fun desc(): String
+
+    fun highlighter(): Highlighter? = null
 }
 
 fun Command.match(line: String): Boolean {
     val ind = line.indexOf(' ')
     val command = if (ind < 0) line else line.substring(0, ind)
     return (short != null && command.equals(":$short", ignoreCase = true)) ||
             command.equals(":$name", ignoreCase = true)
+}
+
+fun Command.weakMatch(line: String): Boolean {
+    val ind = line.indexOf(' ')
+    val command = if (ind < 0) line else line.substring(0, ind)
+    return (short != null && command.startsWith(":$short", ignoreCase = true)) ||
+            command.startsWith(":$name", ignoreCase = true)
 }

kshell/src/main/kotlin/sparklin/kshell/ContextHighlighter.kt

@@ -4,35 +4,28 @@ import sparklin.kshell.org.jline.reader.Highlighter
 import sparklin.kshell.org.jline.reader.LineReader
 import sparklin.kshell.org.jline.utils.AttributedString
 import sparklin.kshell.org.jline.utils.AttributedStringBuilder
+import sparklin.kshell.plugins.BaseHighlighter
 
-class ContextHighlighter(private val isIncomplete: () -> Boolean): Highlighter {
-    private val highlighters = mutableMapOf<Context, Highlighter>()
+class ContextHighlighter(private val isSyntaxMode: (String) -> Boolean,
+                         private val findCommand: (String) -> Command?): Highlighter {
+    var enabled: Boolean = false
 
-    override fun highlight(reader: LineReader, buffer: String): AttributedString {
-        val context = if (isIncomplete() &&
-                buffer.startsWith(":") &&
-                !buffer.startsWith("::")) Context.COMMAND else Context.CODE
-        return highlighters.getOrDefault(context, DEFAULT).highlight(reader, buffer)
-    }
-
-    fun bind(context: Context, highlighter: Highlighter) {
-        highlighters[context] = highlighter
-    }
+    var syntaxHighlighter: BaseHighlighter = DEFAULT
 
-    fun default(context: Context) {
-        bind(context, DEFAULT)
+    override fun highlight(reader: LineReader, buffer: String): AttributedString {
+        if (!enabled) return DEFAULT.highlight(reader, buffer)
+        if (isSyntaxMode(buffer)) return syntaxHighlighter.highlight(reader, buffer)
+        val highlighter = findCommand(buffer)?.highlighter() ?: DEFAULT
+        return highlighter.highlight(reader, buffer)
     }
 
     companion object {
-        private val DEFAULT = Highlighter { _, buffer ->
-            val sb = AttributedStringBuilder()
-            sb.append(buffer)
-            sb.toAttributedString()
+        private val DEFAULT = object : BaseHighlighter {
+            override fun highlight(buffer: String, offset: Int): AttributedString {
+                val sb = AttributedStringBuilder()
+                sb.append(buffer)
+                return sb.toAttributedString()
+            }
         }
     }
-
-    enum class Context {
-        CODE,
-        COMMAND
-    }
 }

kshell/src/main/kotlin/sparklin/kshell/KShell.kt

@@ -46,7 +46,7 @@ open class KShell(val disposable: Disposable,
     val term = TerminalBuilder.builder().build()
     lateinit var readerBuilder: LineReaderBuilder
     lateinit var reader: LineReader
-    val highlighter = ContextHighlighter(incompleteLines::isNotEmpty)
+    val highlighter = ContextHighlighter({ s -> !isCommandMode(s)}, { s -> commands.firstOrNull { it.weakMatch(s) } })
 
     val commands = mutableListOf<sparklin.kshell.Command>(FakeQuit())
     val eventManager = EventManager()
@@ -72,6 +72,10 @@ open class KShell(val disposable: Disposable,
         configuration.plugins().forEach { it.init(this, configuration) }
     }
 
+    private fun isCommandMode(buffer: String): Boolean = incompleteLines.isEmpty()
+                && buffer.startsWith(":")
+                && !buffer.startsWith("::")
+
     fun doRun() {
         initEngine()
 
@@ -80,7 +84,7 @@ open class KShell(val disposable: Disposable,
 
             if (line == null || isQuitAction(line)) break
 
-            if (incompleteLines.isEmpty() && line.startsWith(":") && !line.startsWith("::")) {
+            if (isCommandMode(line)) {
                 try {
                     val action = commands.first { it.match(line) }
                     action.execute(line)
@@ -128,7 +132,9 @@ open class KShell(val disposable: Disposable,
 
     private fun nextLine(code: String) = CodeLine(state.lineIndex.getAndIncrement(), code)
 
-    fun compile(code: String) = compiler.compile(state, nextLine(code))
+    private fun compile(code: String) = compiler.compile(state, nextLine(code))
+
+    fun compile(code: SourceCode) = compiler.compile(state, code)
 
     fun eval(source: String): ResultWrapper =
         state.lock.write {

kshell/src/main/kotlin/sparklin/kshell/plugins/BaseHighlighter.kt

@@ -0,0 +1,13 @@
+package sparklin.kshell.plugins
+
+import sparklin.kshell.org.jline.reader.Highlighter
+import sparklin.kshell.org.jline.reader.LineReader
+import sparklin.kshell.org.jline.utils.AttributedString
+
+interface BaseHighlighter: Highlighter {
+    fun highlight(buffer: String, offset: Int = 0): AttributedString
+
+    override fun highlight(reader: LineReader, buffer: String): AttributedString {
+        return highlight(buffer)
+    }
+}

kshell/src/main/kotlin/sparklin/kshell/plugins/KotlinHighlighter.kt

@@ -3,32 +3,19 @@ package sparklin.kshell.plugins
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.lexer.KtKeywordToken
 import org.jetbrains.kotlin.psi.*
-import sparklin.kshell.org.jline.reader.Highlighter
-import sparklin.kshell.org.jline.reader.LineReader
 import sparklin.kshell.org.jline.utils.AttributedString
 import sparklin.kshell.org.jline.utils.AttributedStringBuilder
 import sparklin.kshell.org.jline.utils.AttributedStyle
-import sparklin.kshell.repl.Code
-import sparklin.kshell.repl.ReplChecker
-import sparklin.kshell.repl.ReplState
-import sparklin.kshell.repl.Result
+import sparklin.kshell.repl.*
 import java.util.concurrent.atomic.AtomicInteger
 
 class KotlinHighlighter(private val state: ReplState,
                         private val checker: ReplChecker,
-                        private val styles: SyntaxPlugin.HighlightStyles): Highlighter {
-    private val counter = AtomicInteger(0)
-
-    override fun highlight(reader: LineReader, buffer: String): AttributedString {
-        return highlight(buffer)
-    }
-
-    fun highlight(buffer: String): AttributedString {
-        if (buffer.startsWith(":") && !buffer.startsWith("::")) {
-            return AttributedStringBuilder().append(buffer).toAttributedString()
-        }
-
-        val fragment = CodeFragment(counter.getAndIncrement(), buffer)
+                        private val styles: SyntaxPlugin.HighlightStyles): BaseHighlighter {
+    override fun highlight(buffer: String, offset: Int): AttributedString {
+        require(offset >= 0)
+        val code = buffer.substring(offset)
+        val fragment = CodeFragment(code)
         val lineResult = checker.check(state, fragment, true)
 
         val psi = when (lineResult) {
@@ -37,7 +24,8 @@ class KotlinHighlighter(private val state: ReplState,
         }
 
         val sb = AttributedStringBuilder()
-        for (i in buffer.indices) {
+        if (offset != 0) sb.append(buffer.substring(0, offset))
+        for (i in code.indices) {
             psi.findElementAt(i)?.let { element ->
                 val st = when {
                     element.isKeyword() -> styles.keyword
@@ -49,7 +37,7 @@ class KotlinHighlighter(private val state: ReplState,
                     else -> null
                 } ?: AttributedStyle.DEFAULT
                 sb.style(st)
-                sb.append(buffer[i])
+                sb.append(code[i])
             }
         }
         return sb.toAttributedString()
@@ -69,9 +57,11 @@ class KotlinHighlighter(private val state: ReplState,
         return ss.any { it == e }
     }
 
-    class CodeFragment(val no: Int, val src: String): Code {
-        override fun mkFileName(): String = "Fragment_$no"
-
-        override fun source(): String = src
+    private class CodeFragment(override val code: String): SourceCode {
+        override val no: Int = 0
+        override val part: Int = 0
+        override fun mkFileName(): String = "Fragment"
+        override fun nextPart(codePart: String): SourceCode = throw UnsupportedOperationException("Should never happen")
+        override fun replace(code: String): SourceCode = throw UnsupportedOperationException("Should never happen")
     }
 }

kshell/src/main/kotlin/sparklin/kshell/plugins/PastePlugin.kt

@@ -6,7 +6,6 @@ import sparklin.kshell.BaseCommand
 import sparklin.kshell.KShell
 import sparklin.kshell.Plugin
 import sparklin.kshell.configuration.Configuration
-import sparklin.kshell.repl.EvalError
 import sparklin.kshell.repl.Result
 
 class PastePlugin : Plugin {
@@ -45,7 +44,7 @@ class PastePlugin : Plugin {
     override fun init(repl: KShell, config: Configuration) {
         this.repl = repl
         this.console = repl.reader
-        this.pasteConsole = repl.readerBuilder.build()
+        this.pasteConsole = repl.readerBuilder.highlighter(console.highlighter).build()
 
         repl.registerCommand(Paste(config))
     }

kshell/src/main/kotlin/sparklin/kshell/plugins/RuntimePlugin.kt

@@ -2,9 +2,13 @@ package sparklin.kshell.plugins
 
 import sparklin.kshell.*
 import sparklin.kshell.configuration.Configuration
+import sparklin.kshell.org.jline.reader.Highlighter
+import sparklin.kshell.org.jline.reader.LineReader
+import sparklin.kshell.org.jline.utils.AttributedString
 import sparklin.kshell.repl.*
 import sparklin.kshell.repl.ReplCompiler.Companion.RESULT_FIELD_NAME
 import sparklin.kshell.repl.ReplCompiler.Companion.RUN_FIELD_NAME
+import java.util.concurrent.atomic.AtomicInteger
 import kotlin.reflect.KClass
 import kotlin.reflect.KFunction
 import kotlin.reflect.KProperty1
@@ -19,22 +23,32 @@ class RuntimePlugin : Plugin {
         override val description: String = "display the type of an expression without evaluating it"
 
         override val params = "<expr>"
-
         override fun execute(line: String) {
             val p = line.indexOf(' ')
             val expr = line.substring(p + 1).trim()
 
-            val compileResult = repl.compile(expr)
+            val compileResult = repl.compile(CodeExpr(counter.getAndIncrement(), expr))
             when (compileResult) {
-                is Result.Error ->
-                    repl.handleError(compileResult.error)
-                is Result.Success -> {
-                    compileResult.data.classes.type?.let {
-                        println(it)
-                    }
-                }
+                is Result.Error -> repl.handleError(compileResult.error)
+                is Result.Success -> compileResult.data.classes.type?.let { println(it) }
             }
         }
+
+        override fun highlighter(): Highlighter = customHighlighter
+    }
+
+    private class CustomHighlighter(val baseHighlighter: () -> BaseHighlighter): Highlighter {
+        override fun highlight(reader: LineReader, buffer: String): AttributedString {
+            val p = buffer.indexOf(' ')
+            return baseHighlighter().highlight(buffer, p + 1)
+        }
+    }
+
+    private data class CodeExpr(override val no: Int, override val code: String): SourceCode {
+        override val part: Int = 0
+        override fun mkFileName(): String = "TypeInference_$no"
+        override fun nextPart(codePart: String): SourceCode = throw UnsupportedOperationException("Should never happen")
+        override fun replace(code: String): CodeExpr = CodeExpr(no, code)
     }
 
     inner class ListSymbols(conf: Configuration) : sparklin.kshell.BaseCommand() {
@@ -52,6 +66,8 @@ class RuntimePlugin : Plugin {
     private lateinit var repl: KShell
     private var lastCompiledClasses: CompiledClasses? = null
     private lateinit var table: SymbolsTable
+    private lateinit var customHighlighter: CustomHighlighter
+    private val counter = AtomicInteger(0)
 
     override fun init(repl: KShell, config: Configuration) {
         this.repl = repl
@@ -65,6 +81,7 @@ class RuntimePlugin : Plugin {
         })
 
         repl.invokeWrapper = ExtractSymbols(repl.invokeWrapper, table)
+        customHighlighter = CustomHighlighter({ repl.highlighter.syntaxHighlighter })
 
         repl.registerCommand(InferType(config))
         repl.registerCommand(ListSymbols(config))