Switch syntax highlighting to tree-sitter via Highlights.jl

Project.toml

@@ -9,6 +9,7 @@ projects = ["test"]
 [deps]
 CompilerCaching = "9db33cc3-5358-4881-8759-fa4194144afd"
 ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
+Highlights = "eafb193a-b7ab-5a9e-9068-77385905fa72"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
@@ -19,17 +20,21 @@ REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 Tracy = "e689c965-62c8-4b79-b2c5-8359227902fd"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
+tree_sitter_gcn_jll = "8b5cbfcf-8811-596e-8e87-0e51e05ee4b2"
+tree_sitter_llvm_jll = "44208993-ee63-5069-9443-8e43b04a9b30"
+tree_sitter_ptx_jll = "71e3f6e6-c059-5e7c-a2f2-d560fdad7ce5"
+tree_sitter_spirv_jll = "f0e86581-c468-54df-a4de-3266e11a3c86"
 
 [weakdeps]
 AMDGPU_LLVM_Backend_jll = "cc5c0156-bd05-5a77-8a68-bb0aafb29019"
 LLVMDowngrader_jll = "f52de702-fb25-5922-94ba-81dd59b07444"
 NVPTX_LLVM_Backend_jll = "ef6e0fe3-e6ef-59c0-bde6-4989574699e0"
 
-
 [compat]
 AMDGPU_LLVM_Backend_jll = "22"
 CompilerCaching = "0.3"
 ExprTools = "0.1"
+Highlights = "0.6"
 InteractiveUtils = "1"
 LLVM = "9.9"
 LLVMDowngrader_jll = "0.8"
@@ -43,3 +48,7 @@ TOML = "1"
 Tracy = "0.1.4"
 UUIDs = "1"
 julia = "1.10"
+tree_sitter_gcn_jll = "0.1"
+tree_sitter_llvm_jll = "1.1"
+tree_sitter_ptx_jll = "0.1"
+tree_sitter_spirv_jll = "0.1"

src/reflection.jl

@@ -8,101 +8,130 @@ const Cthulhu = Base.PkgId(UUID("f68482b8-f384-11e8-15f7-abe071a5a75f"), "Cthulh
 # syntax highlighting
 #
 
-const _pygmentize = Ref{Union{String,Nothing}}()
-function pygmentize()
-    if !isassigned(_pygmentize)
-        _pygmentize[] = Sys.which("pygmentize")
+import Highlights
+using tree_sitter_llvm_jll, tree_sitter_ptx_jll, tree_sitter_spirv_jll,
+      tree_sitter_gcn_jll
+
+# tree-sitter grammars for the assembly formats we emit, keyed by the
+# `source_code` of the compiler target
+const highlight_languages = Dict(
+    "llvm"  => tree_sitter_llvm_jll,
+    "ptx"   => tree_sitter_ptx_jll,
+    "spirv" => tree_sitter_spirv_jll,
+    "gcn"   => tree_sitter_gcn_jll,
+)
+
+# Set to a Highlights.jl theme, or `nothing` to detect the terminal background.
+const highlight_theme = Ref{Any}(nothing)
+const detected_highlight_theme = Ref{Union{Nothing,String}}()
+const highlight_theme_lock = ReentrantLock()
+
+function environment_highlight_theme(colorfgbg = get(ENV, "COLORFGBG", ""))
+    background = tryparse(Int, last(rsplit(colorfgbg, ';'; limit=2)))
+    isnothing(background) && return
+    rgb = if 0 <= background < 16
+        palette = (0x000000, 0x800000, 0x008000, 0x808000,
+                   0x000080, 0x800080, 0x008080, 0xc0c0c0,
+                   0x808080, 0xff0000, 0x00ff00, 0xffff00,
+                   0x0000ff, 0xff00ff, 0x00ffff, 0xffffff)
+        color = palette[background + 1]
+        ((color >> 16) & 0xff, (color >> 8) & 0xff, color & 0xff) ./ 255
+    elseif 16 <= background < 232
+        levels = (0, 95, 135, 175, 215, 255)
+        color = background - 16
+        (levels[color ÷ 36 + 1], levels[(color ÷ 6) % 6 + 1], levels[color % 6 + 1]) ./ 255
+    elseif 232 <= background < 256
+        fill((8 + 10 * (background - 232)) / 255, 3)
+    else
+        return
     end
-    return _pygmentize[]
+    light_background(rgb) ? "Monokai Pro Light" : "Monokai Pro"
 end
 
-const _pygmentize_version = Ref{Union{VersionNumber, Nothing}}()
-function pygmentize_version()
-    isassigned(_pygmentize_version) && return _pygmentize_version[]
+light_background(rgb) = 0.299 * rgb[1] + 0.587 * rgb[2] + 0.114 * rgb[3] > 0.5
 
-    pygmentize_cmd = pygmentize()
-    if isnothing(pygmentize_cmd)
-        return _pygmentize_version[] = nothing
+function terminal_background(reply)
+    match = Base.match(r"\e\]11;rgb:([0-9a-f]{1,4})/([0-9a-f]{1,4})/([0-9a-f]{1,4})(?:\a|\e\\)"i,
+                       reply)
+    isnothing(match) && return
+    map(match.captures) do component
+        parse(Int, component; base=16) / (16^length(component) - 1)
     end
+end
 
-    cmd = `$pygmentize_cmd -V`
-    @static if Sys.iswindows()
-        # Avoid encoding issues with pipes on Windows by using cmd.exe to capture stdout for us
-        cmd = `cmd.exe /C $cmd`
-        cmd = addenv(cmd, "PYTHONUTF8" => 1)
-    end
-    version_str = readchomp(cmd)
+function terminal_highlight_theme(reply)
+    rgb = terminal_background(reply)
+    isnothing(rgb) && return
+    light_background(rgb) ? "Monokai Pro Light" : "Monokai Pro"
+end
 
-    pos = findfirst("Pygments version ", version_str)
-    if !isnothing(pos)
-        version_start = last(pos) + 1
-        version_end = findnext(',', version_str, version_start) - 1
-        version = tryparse(VersionNumber, version_str[version_start:version_end])
-    else
-        version = nothing
-    end
+function repl_terminal(io)
+    repl = isdefined(Base, :active_repl) ? Base.active_repl : nothing
+    hasproperty(repl, :t) || return
+    terminal = repl.t
+    terminal isa Base.Terminals.TTYTerminal || return
+    only(Base.unwrapcontext(io)) === terminal.out_stream || return
+    terminal.in_stream isa Base.TTY || return
+    terminal
+end
 
-    if isnothing(version)
-        @warn "Could not parse Pygments version:\n$version_str"
+function query_terminal_theme(terminal; timeout=0.1)
+    input = terminal.in_stream
+    Base.Terminals.raw!(terminal, true) || return
+    try
+        Base.start_reading(input)
+        print(terminal.out_stream, "\e]11;?\a")
+        flush(terminal.out_stream)
+
+        reply = UInt8[]
+        deadline = time() + timeout
+        while time() < deadline
+            timedwait(() -> bytesavailable(input) > 0, max(0, deadline - time()); pollint=0.005) ===
+                :timed_out && break
+            append!(reply, readavailable(input))
+            theme = terminal_highlight_theme(String(reply))
+            isnothing(theme) || return theme
+        end
+    finally
+        Base.stop_reading(input)
+        Base.Terminals.raw!(terminal, false)
     end
-
-    return _pygmentize_version[] = version
 end
 
-function pygmentize_support(lexer)
-    highlighter_ver = pygmentize_version()
-    if isnothing(highlighter_ver)
-        @warn "Syntax highlighting of $lexer code relies on Pygments.\n\
-               Use `pip install pygments` to install the lastest version" maxlog = 1
-        return false
-    elseif lexer == "ptx"
-        if highlighter_ver < v"2.16"
-            @warn "Pygments supports PTX highlighting starting from version 2.16\n\
-                   Detected version $highlighter_ver\n\
-                   Please update with `pip install pygments -U`" maxlog = 1
-            return false
-        end
-        return true
-    elseif lexer == "gcn"
-        if highlighter_ver < v"2.8"
-            @warn "Pygments supports GCN highlighting starting from version 2.8\n\
-                   Detected version $highlighter_ver\n\
-                   Please update with `pip install pygments -U`" maxlog = 1
-            return false
+function selected_highlight_theme(io)
+    theme = highlight_theme[]
+    isnothing(theme) || return theme
+
+    theme = environment_highlight_theme()
+    isnothing(theme) || return theme
+
+    terminal = repl_terminal(io)
+    isnothing(terminal) && return "Monokai Pro"
+    lock(highlight_theme_lock) do
+        if !isassigned(detected_highlight_theme)
+            detected_highlight_theme[] = query_terminal_theme(terminal)
         end
-        return true
-    else
-        return false
+        something(detected_highlight_theme[], "Monokai Pro")
     end
 end
 
 function highlight(io::IO, code, lexer)
     have_color = get(io, :color, false)
-    if !have_color
+    language = get(highlight_languages, lexer, nothing)
+    if !have_color || language === nothing
         print(io, code)
-    elseif lexer == "llvm"
-        InteractiveUtils.print_llvm(io, code)
-    elseif pygmentize_support(lexer)
-        lexer = lexer == "gcn" ? "amdgpu" : lexer
-        pygments_args = String[pygmentize(), "-f", "terminal", "-P", "bg=dark", "-l", lexer]
-        @static if Sys.iswindows()
-            # Avoid encoding issues with pipes on Windows by using a temporary file
-            mktemp() do tmp_path, tmp_io
-                println(tmp_io, code)
-                close(tmp_io)
-                push!(pygments_args, "-o", tmp_path, tmp_path)
-                cmd = Cmd(pygments_args)
-                wait(open(cmd))  # stdout and stderr go to devnull
-                print(io, read(tmp_path, String))
-            end
-        else
-            cmd = Cmd(pygments_args)
-            pipe = open(cmd, "r+")
-            print(pipe, code)
-            close(pipe.in)
-            print(io, read(pipe, String))
-        end
-    else
+        return
+    end
+
+    # render to a buffer first so that highlighting failures don't result in
+    # partial output
+    buf = IOBuffer()
+    try
+        Highlights.highlight(buf, MIME("text/ansi"), code, language,
+                             selected_highlight_theme(io))
+        write(io, take!(buf))
+    catch err
+        @warn "Failed to highlight $lexer code" exception=(err, catch_backtrace()) maxlog=1
         print(io, code)
     end
     return

src/spirv.jl

@@ -219,6 +219,8 @@ end
     return output
 end
 
+source_code(target::SPIRVCompilerTarget) = "spirv"
+
 # reimplementation that uses `spirv-dis`, giving much more pleasant output
 function code_native(io::IO, job::CompilerJob{SPIRVCompilerTarget}; raw::Bool=false, dump_module::Bool=false)
     config = CompilerConfig(job.config; strip=!raw, only_entry=!dump_module, validate=false)
@@ -230,14 +232,10 @@ function code_native(io::IO, job::CompilerJob{SPIRVCompilerTarget}; raw::Bool=fa
         flush(input_io)
 
         disassembler = SPIRV_Tools_jll.spirv_dis()
-        if io == stdout
-            run(`$disassembler $input_path`)
-        else
-            mktemp() do output_path, output_io
-                run(`$disassembler $input_path -o $output_path`)
-                asm = read(output_io, String)
-                print(io, asm)
-            end
+        mktemp() do output_path, output_io
+            run(`$disassembler $input_path -o $output_path`)
+            asm = read(output_io, String)
+            highlight(io, asm, source_code(job.config.target))
         end
     end
 end

test/utils.jl

@@ -114,27 +114,80 @@ end
 @testset "highlighting" begin
     ansi_color = "\x1B[3"  # beginning of any foreground color change
 
+    @test GPUCompiler.environment_highlight_theme("15;0") == "Monokai Pro"
+    @test GPUCompiler.environment_highlight_theme("0;15") == "Monokai Pro Light"
+    @test GPUCompiler.environment_highlight_theme("0;231") == "Monokai Pro Light"
+    @test GPUCompiler.environment_highlight_theme("15;232") == "Monokai Pro"
+    @test GPUCompiler.environment_highlight_theme("0;255") == "Monokai Pro Light"
+    @test GPUCompiler.environment_highlight_theme("") === nothing
+
+    @test GPUCompiler.terminal_background("\e]11;rgb:f/f/f\a") == [1, 1, 1]
+    @test GPUCompiler.terminal_background("\e]11;rgb:00/80/ff\e\\") == [0, 128/255, 1]
+    @test GPUCompiler.terminal_background("\e]11;rgb:0000/0000/0000\a") == [0, 0, 0]
+    @test GPUCompiler.terminal_background("\e]10;rgb:ffff/ffff/ffff\a") === nothing
+    @test GPUCompiler.terminal_highlight_theme("\e]11;rgb:ffff/ffff/ffff\a") ==
+          "Monokai Pro Light"
+    @test GPUCompiler.terminal_highlight_theme("\e]11;rgb:0000/0000/0000\a") ==
+          "Monokai Pro"
+
+    old_theme = GPUCompiler.highlight_theme[]
+    try
+        GPUCompiler.highlight_theme[] = "GitHub"
+        @test GPUCompiler.selected_highlight_theme(IOBuffer()) == "GitHub"
+    finally
+        GPUCompiler.highlight_theme[] = old_theme
+    end
+
     @testset "PTX" begin
         sample = """
             max.s64         %rd24, %rd18, 0;
             add.s64         %rd7, %rd23, -1;
             setp.lt.u64     %p2, %rd7, %rd24;
             @%p2 bra        \$L__BB0_3;
         """
-        can_highlight = GPUCompiler.pygmentize_support("ptx")
         highlighted = sprint(GPUCompiler.highlight, sample, "ptx"; context = (:color => true))
-        @test occursin(ansi_color, highlighted) skip = !can_highlight
+        @test occursin(ansi_color, highlighted)
+
+        # no color without a color-enabled IO
+        plain = sprint(GPUCompiler.highlight, sample, "ptx")
+        @test plain == sample
+    end
+
+    @testset "SPIR-V" begin
+        sample = """
+                       OpCapability Kernel
+                  %1 = OpExtInstImport "OpenCL.std"
+                 %10 = OpTypeInt 32 0
+        """
+        highlighted = sprint(GPUCompiler.highlight, sample, "spirv"; context = (:color => true))
+        @test occursin(ansi_color, highlighted)
+    end
+
+    @testset "LLVM" begin
+        sample = """
+            define i32 @add(i32 %x, i32 %y) {
+              %sum = add i32 %x, %y
+              ret i32 %sum
+            }
+        """
+        highlighted = sprint(GPUCompiler.highlight, sample, "llvm"; context = (:color => true))
+        @test occursin(ansi_color, highlighted)
     end
 
     @testset "GCN" begin
         sample = """
             v_add_u32     v3, vcc, s0, v0
             v_mov_b32     v4, s1
-            v_addc_u32    v4, vcc, v4, 0, vcc
+            s_endpgm
         """
-        can_highlight = GPUCompiler.pygmentize_support("gcn")
         highlighted = sprint(GPUCompiler.highlight, sample, "gcn"; context = (:color => true))
-        @test occursin(ansi_color, highlighted) skip = !can_highlight
+        @test occursin(ansi_color, highlighted)
+    end
+
+    @testset "unknown lexers pass through" begin
+        sample = "some random text\n"
+        highlighted = sprint(GPUCompiler.highlight, sample, "metal"; context = (:color => true))
+        @test highlighted == sample
     end
 end