Vim treesitter prototype

highlight.vim

vim9script

ts_load("c", "/usr/lib/tree_sitter/c.so")

prop_type_add('macro', {highlight: 'Macro', override: true})

def GetBufferText(node: opaque<TSNode>, bufnr: number): string
    var end: tuple<number, number> = node.end_point
    var start: tuple<number, number> = node.start_point
    var lines: list<string> = getbufline(bufnr, start[0], end[0])

    lines[0] = strpart(lines[0], start[1])
    lines[-1] = strpart(lines[-1], 0, end[1])

    return join(lines, "\n")
enddef

def Predicate_eq(matches: dict<list<opaque<TSNode>>>, bufnr: number,
                predicate: list<any>, any: bool): bool
    var capid: number = predicate[1]

    if !has_key(matches, capid)
        return true
    endif

    var nodes: list<opaque<TSNode>> = matches[capid]

    var target: string

    if type(predicate[2]) == v:t_string
        target = predicate[2]
    else
        var target_id: number = predicate[2]
        var target_nodes = matches[target_id]

        if empty(target_nodes) 
            return false 
        endif
        target = GetBufferText(target_nodes[0], bufnr)
    endif

    # Comparison logic
    for node in nodes
        var node_text: string = GetBufferText(node, bufnr)
        var is_equal: bool = node_text == target

        if any && is_equal
            return true
        elseif !any && !is_equal
            return false
        endif
    endfor

    return !any
enddef

var predicates: dict<any> = {
    "eq?": Predicate_eq
}

class LanguageTree
    var parser: opaque<TSParser>
    var tree: opaque<TSTree>
    var children: list<object<LanguageTree>>
    var buf: number
    var lang: string

    def new(this.buf, this.lang)
        this.parser = tsparser_new()

        tsparser_set_language(this.parser, this.lang)
    enddef

    def Parse(): void
        this.tree = tsparser_parse_buf(this.parser, this.buf, 99999, this.tree)
    enddef
endclass

class Highlighter
    var captures: tuple<...list<string>>
    var patterns: dict<list<list<any>>>
    var languagetree: object<LanguageTree>
    var querystr: string
    var query: opaque<TSQuery>

    def new(this.languagetree, this.querystr)
        this.query = tsquery_new(this.languagetree.lang, this.querystr)

        var info: dict<any> = tsquery_inspect(this.query)

        this.captures = info.captures
        this.patterns = info.patterns
    enddef

    def Run()
        if this.languagetree.tree == null_opaque
            return
        endif

        var cursor: opaque<TSQueryCursor> = tsquerycursor_new()
        var root: opaque<TSNode> = tstree_root_node(this.languagetree.tree)

        tsquerycursor_exec(cursor, this.query, root)

        prop_clear(1, line('$'))

        while tsquerycursor_next_match(cursor)
            var captures: dict<list<opaque<TSNode>>> = {}
            var capturesid: dict<number>
            
            for capture in cursor.match_captures
                var node: opaque<TSNode> = capture[0]
                var id: number = capture[1]

                if has_key(captures, id)
                    captures[id] += [node]
                else
                    captures[id] = [node]
                endif

                capturesid[id] = id
            endfor

            var valid: bool = true

            if has_key(this.patterns, cursor.match_pattern_index)
                var preds = this.patterns[cursor.match_pattern_index]

                for pred in preds
                    var op: string = pred[0]

                    if has_key(predicates, op)
                        if !predicates[op](captures, this.languagetree.buf, pred, false)
                            valid = false
                            break
                        endif
                    endif
                endfor
            endif

            if valid
                for id in keys(captures)
                    for node in captures[id]
                        var startp: tuple<number, number> = node.start_point
                        var endp: tuple<number, number> = node.end_point

                        prop_add(startp[0], startp[1], {
                            type: this.captures[capturesid[id]],
                            end_lnum: endp[0],
                            end_col: endp[1],
                        })
                    endfor
                endfor
            endif
        endwhile
    enddef
endclass

def StartHighlight(bufnr: number): void
    var querystr: list<string> =<< trim END
    [
    "#if"
    "#ifdef"
    "#ifndef"
    "#else"
    "#elif"
    "#endif"
    "#elifdef"
    "#elifndef"
    (preproc_directive)
    ] @macro
    END
    var lt: object<LanguageTree> = LanguageTree.new(bufnr, "c")
    var hi: object<Highlighter> = Highlighter.new(lt, join(querystr, "\n"))

    setbufvar(bufnr, "language_tree", lt)
    setbufvar(bufnr, "highlighter", hi)

    lt.Parse()
    hi.Run()
enddef

command! -nargs=1 StartHighlight {
    var nr: number = str2nr(<q-args>)
    call StartHighlight(nr)
}