Adds syntax highlighting to the CLI (#639)

* Adds syntax highlighting to the CLI
* Implements a shell highlighter by leveraging PartiQL Lexer and Parser
* Adds support for the AST token

cli/src/org/partiql/cli/main.kt

@@ -34,6 +34,7 @@ import org.partiql.lang.eval.ExprValueFactory
 import org.partiql.lang.eval.TypingMode
 import org.partiql.lang.syntax.SqlParser
 import org.partiql.shell.Shell
+import org.partiql.shell.Shell.ShellConfiguration
 import java.io.File
 import java.io.FileInputStream
 import java.io.FileOutputStream
@@ -109,6 +110,8 @@ private val inputFormatOpt = optParser.acceptsAll(listOf("input-format", "if"),
 private val wrapIonOpt = optParser.acceptsAll(listOf("wrap-ion", "w"), "wraps Ion input file values in a bag, requires the input format to be ION, requires the query option")
     .availableIf(queryOpt)
 
+private val monochromeOpt = optParser.acceptsAll(listOf("monochrome", "m"), "removes syntax highlighting for the REPL")
+
 private val outputFileOpt = optParser.acceptsAll(listOf("output", "o"), "output file, requires the query option (default: stdout)")
     .availableIf(queryOpt)
     .withRequiredArg()
@@ -132,6 +135,8 @@ private val outputFormatOpt = optParser.acceptsAll(listOf("output-format", "of")
  * * -e --environment: takes an environment file to load as the initial global environment
  * * -p --permissive: run the query in permissive typing mode (returns MISSING rather than error for data type
  * mismatches)
+ * * Interactive only:
+ *      * -m --monochrome: removes syntax highlighting for the REPL
  * * Non interactive only:
  *      * -q --query: PartiQL query
  *      * -i --input: input file
@@ -182,7 +187,7 @@ fun main(args: Array<String>) = try {
     if (optionSet.has(queryOpt)) {
         runCli(environment, optionSet, compilerPipeline)
     } else {
-        runShell(environment, compilerPipeline)
+        runShell(environment, optionSet, compilerPipeline)
     }
 } catch (e: OptionException) {
     System.err.println("${e.message}\n")
@@ -193,8 +198,9 @@ fun main(args: Array<String>) = try {
     exitProcess(1)
 }
 
-private fun runShell(environment: Bindings<ExprValue>, compilerPipeline: CompilerPipeline) {
-    Shell(valueFactory, System.out, parser, compilerPipeline, environment).start()
+private fun runShell(environment: Bindings<ExprValue>, optionSet: OptionSet, compilerPipeline: CompilerPipeline) {
+    val config = ShellConfiguration(isMonochrome = optionSet.has(monochromeOpt))
+    Shell(valueFactory, System.out, parser, compilerPipeline, environment, config).start()
 }
 
 private fun runCli(environment: Bindings<ExprValue>, optionSet: OptionSet, compilerPipeline: CompilerPipeline) {

cli/src/org/partiql/shell/Shell.kt

@@ -79,6 +79,7 @@ class Shell(
     private val parser: Parser,
     private val compiler: CompilerPipeline,
     private val initialGlobal: Bindings<ExprValue>,
+    private val config: ShellConfiguration = ShellConfiguration()
 ) {
 
     private val homeDir: Path = Paths.get(System.getProperty("user.home"))
@@ -122,11 +123,15 @@ class Shell(
     }
 
     private fun run(exiting: AtomicBoolean) = TerminalBuilder.builder().build().use { terminal ->
-
+        val highlighter = when {
+            this.config.isMonochrome -> null
+            else -> ShellHighlighter()
+        }
         val reader = LineReaderBuilder.builder()
             .terminal(terminal)
             .parser(ShellParser)
             .completer(NullCompleter())
+            .highlighter(highlighter)
             .expander(ShellExpander)
             .variable(LineReader.HISTORY_FILE, homeDir.resolve(".partiql/.history"))
             .variable(LineReader.SECONDARY_PROMPT_PATTERN, PROMPT_2)
@@ -261,6 +266,12 @@ class Shell(
             }
         }
     }
+
+    /**
+     * A configuration class representing any configurations specified by the user
+     * @param isMonochrome specifies the removal of syntax highlighting
+     */
+    class ShellConfiguration(val isMonochrome: Boolean = false)
 }
 
 /**

cli/src/org/partiql/shell/ShellHighlighter.kt

@@ -11,32 +11,181 @@
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific
  * language governing permissions and limitations under the License.
  */
-
 package org.partiql.shell
 
-import org.jline.builtins.Nano.SyntaxHighlighter
+import com.amazon.ion.IonString
+import com.amazon.ion.system.IonSystemBuilder
 import org.jline.reader.Highlighter
 import org.jline.reader.LineReader
 import org.jline.utils.AttributedString
+import org.jline.utils.AttributedStringBuilder
 import org.jline.utils.AttributedStyle
+import org.partiql.lang.errors.Property
+import org.partiql.lang.syntax.ParserException
+import org.partiql.lang.syntax.SqlLexer
+import org.partiql.lang.syntax.SqlParser
+import org.partiql.lang.syntax.Token
+import org.partiql.lang.syntax.TokenType
 import java.io.PrintStream
-import java.nio.file.Path
 import java.util.regex.Pattern
 
 private val SUCCESS: AttributedStyle = AttributedStyle.DEFAULT.foreground(AttributedStyle.GREEN)
 private val ERROR: AttributedStyle = AttributedStyle.DEFAULT.foreground(AttributedStyle.RED)
 private val INFO: AttributedStyle = AttributedStyle.DEFAULT
 private val WARN: AttributedStyle = AttributedStyle.DEFAULT.foreground(AttributedStyle.YELLOW)
 
-class ShellHighlighter(private val syntaxFile: Path) : Highlighter {
+private val ALLOWED_SUFFIXES = setOf("!!")
+
+internal class ShellHighlighter() : Highlighter {
+
+    private val ion = IonSystemBuilder.standard().build()
+    private val lexer = SqlLexer(ion)
+    private val parser = SqlParser(ion)
+
+    /**
+     * Returns a highlighted string by passing the [input] string through the [lexer] and [parser] to identify and
+     * highlight tokens
+     */
+    override fun highlight(reader: LineReader, input: String): AttributedString {
+        if (input.isEmpty()) return AttributedString(input)
+
+        // Map Between Line Number and Index
+        val lineIndexesMap = mutableMapOf<Int, Int>()
+        var currentLine = 0
+        lineIndexesMap[currentLine] = -1
+        input.forEachIndexed { index, char ->
+            if (char == '\n') lineIndexesMap[++currentLine] = index
+        }
+
+        // Check for Non-PartiQL Suffixes (REPL Only)
+        val lastNewlineIndex = lineIndexesMap[currentLine]!!
+        val suffixString = input.substring(lastNewlineIndex + 1, input.lastIndex + 1)
+        val lastValidQueryIndex = when (currentLine > 0 && ALLOWED_SUFFIXES.contains(suffixString)) {
+            true -> lastNewlineIndex
+            false -> input.length
+        }
 
-    private val syntax = SyntaxHighlighter.build(syntaxFile, "PartiQL")
+        // Get Tokens
+        val tokens: List<Token>
+        try {
+            tokens = lexer.tokenize(input.substring(0, lastValidQueryIndex))
+        } catch (e: Exception) {
+            return AttributedString(input, AttributedStyle().foreground(AttributedStyle.RED))
+        }
 
-    override fun highlight(reader: LineReader, line: String): AttributedString = syntax.highlight(line)
+        // Build Token Colors (Last Token is EOF)
+        var builder = AttributedStringBuilder()
+        for (tokenIndex in 0..(tokens.size - 2)) {
+            if (tokenIndex == tokens.lastIndex) break
+            val currentToken = tokens[tokenIndex]
+            val preIndex = when (tokenIndex) {
+                0 -> 0
+                else -> {
+                    val prevToken = tokens[tokenIndex - 1]
+                    (getTokenIndex(prevToken, lineIndexesMap) ?: 0) + prevToken.span.length.toInt()
+                }
+            }
+            val postIndex = when (tokenIndex) {
+                tokens.lastIndex - 1 -> input.lastIndex + 1
+                else -> (getTokenIndex(currentToken, lineIndexesMap) ?: input.lastIndex) + currentToken.span.length.toInt()
+            }
+            addToAttributeStringBuilder(currentToken, lineIndexesMap, builder, input, preIndex, postIndex)
+        }
+
+        // Parse and Replace Token Style if Failures
+        try {
+            parser.parseAstStatement(input.substring(0, lastValidQueryIndex))
+        } catch (e: ParserException) {
+            val column =
+                e.errorContext[Property.COLUMN_NUMBER]?.longValue()?.toInt() ?: return builder.toAttributedString()
+            val lineNumber =
+                e.errorContext[Property.LINE_NUMBER]?.longValue()?.toInt() ?: return builder.toAttributedString()
+            val token = tokens.find {
+                it.span.column.toInt() == column && it.span.line.toInt() == lineNumber
+            } ?: return builder.toAttributedString()
+            builder = createAttributeStringBuilder(token, lineIndexesMap, builder)
+        }
+        return builder.toAttributedString()
+    }
 
     override fun setErrorPattern(errorPattern: Pattern?) {}
 
     override fun setErrorIndex(errorIndex: Int) {}
+
+    /**
+     * Based on the [token] type and location, this function will return a replica of the [stringBuilder] with the
+     * token's location having an updated color-scheme (of color RED). This is used if the [parser] throws a
+     * [ParserException].
+     */
+    private fun createAttributeStringBuilder(
+        token: Token,
+        lineIndexes: Map<Int, Int>,
+        stringBuilder: AttributedStringBuilder
+    ): AttributedStringBuilder {
+        val a = AttributedStringBuilder()
+        val style = AttributedStyle().foreground(AttributedStyle.RED)
+        val length = token.span.length.toInt()
+        val index = getTokenIndex(token, lineIndexes) ?: return a
+        a.append(stringBuilder.subSequence(0, index))
+        a.append(stringBuilder.subSequence(index, index + length), style)
+        a.append(stringBuilder.subSequence(index + length, stringBuilder.length))
+        return a
+    }
+
+    /**
+     * Based on the [token] type and location, this function will modify the [stringBuilder] in place by adding a
+     * [token] and its associated [AttributedStyle]. This is used by the [lexer] to create the highlighted string.
+     */
+    private fun addToAttributeStringBuilder(
+        token: Token,
+        lineIndexes: Map<Int, Int>,
+        stringBuilder: AttributedStringBuilder,
+        input: String,
+        preIndex: Int,
+        postIndex: Int
+    ) {
+        val style = getStyle(token)
+        val length = token.span.length.toInt()
+        val index = getTokenIndex(token, lineIndexes) ?: return
+        stringBuilder.append(input.subSequence(preIndex, index))
+        stringBuilder.append(input.subSequence(index, index + length), style)
+        stringBuilder.append(input.subSequence(index + length, postIndex))
+    }
+
+    /**
+     * Gets the index of a specific token
+     */
+    private fun getTokenIndex(token: Token, lineIndexes: Map<Int, Int>): Int? {
+        val column = token.span.column.toInt()
+        val lineNumber = token.span.line.toInt() - 1
+        return lineIndexes[lineNumber]?.plus(column)
+    }
+
+    /**
+     * Sets the color and thickness of the string based on the [token] type
+     */
+    private fun getStyle(token: Token): AttributedStyle {
+        var style = AttributedStyle()
+        val attrCode = when (token.type) {
+            TokenType.KEYWORD -> if (token.isDataType) AttributedStyle.GREEN else AttributedStyle.CYAN
+            TokenType.AS, TokenType.FOR, TokenType.ASC, TokenType.LAST, TokenType.DESC,
+            TokenType.BY, TokenType.FIRST -> AttributedStyle.CYAN
+            TokenType.LITERAL -> if (token.value is IonString) AttributedStyle.YELLOW else AttributedStyle.BLUE
+            TokenType.ION_LITERAL -> AttributedStyle.YELLOW
+            TokenType.OPERATOR -> AttributedStyle.WHITE
+            TokenType.QUOTED_IDENTIFIER -> AttributedStyle.BRIGHT
+            TokenType.IDENTIFIER -> AttributedStyle.BRIGHT
+            TokenType.MISSING -> AttributedStyle.BLUE
+            TokenType.NULL -> AttributedStyle.BLUE
+            else -> AttributedStyle.WHITE
+        }
+        style = style.foreground(attrCode)
+
+        return when (token.type) {
+            TokenType.IDENTIFIER, TokenType.OPERATOR -> style.bold()
+            else -> style
+        }
+    }
 }
 
 private fun ansi(string: String, style: AttributedStyle) = AttributedString(string, style).toAnsi()

docs/user/CLI.md

@@ -25,6 +25,7 @@ Option                                Description
 -i, --input <File>                    input file, requires the query option (optional)
 -if, --input-format <InputFormat>     input format, requires the query option (default: ION) [ION, PARTIQL]
 -w, --wrap-ion                        wraps Ion input file values in a bag, requires the input format to be ION, requires the query option
+-m, --monochrome                      removes syntax highlighting for the REPL
 -o, --output <File>                   output file, requires the query option (default: stdout)
 -of, --output-format <OutputFormat>   output format, requires the query option (default: PARTIQL) [PARTIQL, PARTIQL_PRETTY, ION_TEXT, ION_BINARY]
 -p, --permissive                      run the PartiQL query in PERMISSIVE typing mode
@@ -61,7 +62,7 @@ To start an interactive read, eval, print loop (REPL) execute:
 > Note that running directly with Gradle will eat arrow keys and control sequences due to the Gradle daemon.
 
 ```shell
-./cli/shell
+./cli/shell.sh
 ```
 
 You will see a prompt that looks as follows:

lang/src/org/partiql/lang/syntax/Token.kt

@@ -66,4 +66,11 @@ data class Token(
             isBinaryOperator || isSpecialOperator -> OPERATOR_PRECEDENCE[text] ?: 0
             else -> 0
         }
+
+    val isDataType: Boolean
+        get() = when {
+            type == TokenType.KEYWORD && CORE_TYPE_NAME_ARITY_MAP.keys.union(TYPE_ALIASES.keys)
+                .contains(text?.toLowerCase()) -> true
+            else -> false
+        }
 }