How 99 Uses TreeSitter to Determine Code Boundaries for Visual Selection
99 uses TreeSitter to locate the smallest enclosing function node that contains the user's visual selection, enabling context-aware AI prompts that include surrounding function signatures and body content.
ThePrimeagen's 99 is a Neovim plugin that bridges visual selections with AI-powered code generation. While raw visual marks define the initial text boundaries, the plugin leverages TreeSitter to determine the syntactic context—specifically the containing function—that surrounds the selected code. This article examines how 99 converts Vim visual selections into range objects and uses TreeSitter queries to identify enclosing code boundaries.
The Visual Selection Workflow in 99
When you invoke the visual command in 99, the plugin executes a precise sequence to transform Vim's native selection into a structured range object that can be enriched with TreeSitter metadata.
Converting Vim Marks to Range Objects
The process begins in lua/99/geo.lua, where the Range.from_visual_selection() function captures the current visual marks ('< and '>) and converts them into a structured Range object containing start and end Points:
function Range.from_visual_selection()
local buffer = vim.api.nvim_get_current_buf()
local start_pos = vim.fn.getpos("'<") -- [buf, line, col, off]
local end_pos = vim.fn.getpos("'>")
local start = Point:from_1_based(start_pos[2], start_pos[3])
local end_ = Point:from_1_based(end_pos[2], end_pos[3])
-- Adjust for line-wise visual mode where the end column can be past the line length
local end_row, _ = end_:to_vim()
local end_line = vim.api.nvim_buf_get_lines(buffer, end_row, end_row + 1, false)
local end_col
if #end_line == 0 then
end_col = 1
else
end_col = #end_line[1]
end
local actual_end = Point.from_0_based(end_row, end_col)
return Range:new(buffer, start, actual_end)
end
This function handles edge cases like line-wise visual mode and empty lines, ensuring the range accurately reflects the user's selection regardless of how the text was highlighted.
Processing the Over-Range Operation
Once the range is established, the workflow delegates to lua/99/ops/over-range.lua. The ops.over_range() function receives the context and range, sends the selected text to the AI backend, and replaces the buffer content upon receiving a response:
function _99.visual(opts)
opts = process_opts(opts)
local context = get_context("visual")
local function perform_range()
set_selection_marks()
local range = Range.from_visual_selection()
ops.over_range(context, range, opts)
end
if opts.additional_prompt then
perform_range()
else
capture_prompt(perform_range, "Visual", context, opts)
end
end
At this stage, the plugin has the raw text boundaries but lacks syntactic context. This is where TreeSitter integration becomes critical for determining code boundaries.
Using TreeSitter to Find Code Boundaries
While Vim marks provide character-level coordinates, TreeSitter provides semantic understanding of where those coordinates fall within the code's structure. 99 uses this to identify the enclosing function that contains the visual selection.
The Containing Function Algorithm
The core logic resides in lua/99/editor/treesitter.lua within the M.containing_function() method. This function loads the language-specific TreeSitter parser, executes a predefined query that captures function nodes, and selects the smallest function that contains the given cursor point:
function M.containing_function(context, cursor)
local root = tree_root(context.buffer, context.file_type)
local query = vim.treesitter.query.get(context.file_type, function_query)
local best_range, best_node = nil, nil
for id, node, _ in query:iter_captures(root, context.buffer, 0, -1, {all=true}) do
local range = Range:from_ts_node(node, context.buffer)
if query.captures[id] == "context.function" and range:contains(cursor) then
if not best_range or best_range:area() > range:area() then
best_range, best_node = range, node
end
end
end
return best_range and Function.from_ts_node(best_node, cursor, context) or nil
end
This algorithm prioritizes the smallest enclosing function by comparing node areas, ensuring that nested functions are handled correctly. The function returns a Function object that exposes both the function_range (including signature) and body_range (code block only).
TreeSitter Queries and Node Capture
The TreeSitter integration relies on language-specific queries that tag function definitions with the capture group @context.function. When containing_function executes, it iterates through all captured nodes and checks if the range contains the cursor point using the Range:contains() method.
This approach allows 99 to determine code boundaries dynamically without regex parsing, handling language-specific syntax variations across Python, JavaScript, Rust, and other supported languages.
Integration with AI Prompts
The TreeSitter-determined boundaries serve a specific purpose: enriching AI prompts with contextual information. When the visual selection is sent to the AI backend, 99 can optionally include the containing function's signature and body, providing the model with structural context that raw text coordinates cannot convey.
This workflow demonstrates how TreeSitter bridges the gap between raw buffer coordinates and semantic code structure, enabling 99 to generate more accurate, context-aware code suggestions based on visual selections.
Summary
- 99 converts Vim visual marks into structured Range objects using
Range.from_visual_selection()inlua/99/geo.lua, handling edge cases like line-wise selection and empty lines. - TreeSitter determines semantic code boundaries through
containing_function()inlua/99/editor/treesitter.lua, which finds the smallest enclosing function node that contains the selection. - The plugin uses TreeSitter queries with the
@context.functioncapture group to identify function nodes across different programming languages. - Visual selections flow through
ops.over_range()inlua/99/ops/over-range.lua, which sends the range to AI backends and replaces text upon completion.
Frequently Asked Questions
How does 99 handle visual selections in line-wise mode?
In lua/99/geo.lua, the Range.from_visual_selection() function detects line-wise visual mode by checking if the end column extends past the actual line length. It adjusts the end point to the last character of the line using vim.api.nvim_buf_get_lines(), ensuring the range accurately captures the selected lines regardless of how Vim represents the marks.
What TreeSitter query does 99 use to find containing functions?
The plugin uses a language-specific query identified by the constant function_query (typically "function" or similar) that captures nodes with the @context.function tag. In lua/99/editor/treesitter.lua, the containing_function() method loads this query via vim.treesitter.query.get() and iterates through captures to find the smallest function node that contains the cursor point.
Can 99 determine boundaries for nested functions?
Yes. The containing_function() algorithm in lua/99/editor/treesitter.lua compares the area of each function node that contains the cursor using range:area(). It selects the node with the smallest area, which corresponds to the innermost nested function, ensuring accurate boundary detection even in complex nested scopes.
How does the visual selection flow from marks to AI prompts?
The workflow begins in lua/99/init.lua with _99.visual(), which calls Range.from_visual_selection() to convert Vim's '< and '> marks into a structured range. This range passes to ops.over_range() in lua/99/ops/over-range.lua, which sends the selected text to the AI backend. Optionally, containing_function() from lua/99/editor/treesitter.lua enriches the prompt with the surrounding function context before the AI request is made.
Have a question about this repo?
These articles cover the highlights, but your codebase questions are specific. Give your agent direct access to the source. Share this with your agent to get started:
curl -s "https://instagit.com/install.md" Maintain an open-source project? Get it listed too →