treesitter vibes?

3 files changed, 214 insertions, 8 deletions

diff --git a/build.zig b/build.zig
index 25520c27..2e54c7f7 100644
--- a/build.zig
+++ b/build.zig
@@ -3,6 +3,9 @@ const std = @import("std");
 pub fn build(b: *std.Build) void {
     const optimize = b.standardOptimizeOption(.{});
     const target = b.standardTargetOptions(.{});
+    var threaded: std.Io.Threaded = .init(b.allocator, .{});
+    defer threaded.deinit();
+    const io = threaded.io();
 
     const exe = b.addExecutable(.{
         .name = "pardes",
@@ -10,6 +13,7 @@ pub fn build(b: *std.Build) void {
             .target = target,
             .optimize = optimize,
             .root_source_file = b.path("main.zig"),
+            .link_libc = true,
         }),
     });
     const tests = b.addExecutable(.{
@@ -73,6 +77,67 @@ pub fn build(b: *std.Build) void {
     vaxis_mod.addImport("uucode", uucode_mod);
     exe.root_module.addImport("vaxis", vaxis_mod);
 
+    const tree_sitter = b.dependency("tree_sitter", .{
+        .target = target,
+        .optimize = optimize,
+    });
+    exe.root_module.addImport("tree-sitter", tree_sitter.module("tree_sitter"));
+
+    const ts_rust_dep = b.dependency("ts_rust", .{});
+    const ts_c_dep = b.dependency("ts_c", .{});
+    const ts_zig_dep = b.dependency("ts_zig", .{});
+    const c_query_path = ts_c_dep.path("queries/highlights.scm").getPath(b);
+    const rust_query_path = ts_rust_dep.path("queries/highlights.scm").getPath(b);
+    const zig_query_path = ts_zig_dep.path("queries/highlights.scm").getPath(b);
+    const c_query = std.Io.Dir.cwd().readFileAlloc(io, c_query_path, b.allocator, .limited(0x100000)) catch @panic("read c highlights");
+    const rust_query = std.Io.Dir.cwd().readFileAlloc(io, rust_query_path, b.allocator, .limited(0x100000)) catch @panic("read rust highlights");
+    const zig_query = std.Io.Dir.cwd().readFileAlloc(io, zig_query_path, b.allocator, .limited(0x100000)) catch @panic("read zig highlights");
+    const ts_queries = b.addOptions();
+    ts_queries.addOption([]const u8, "c_highlights", c_query);
+    ts_queries.addOption([]const u8, "rust_highlights", rust_query);
+    ts_queries.addOption([]const u8, "zig_highlights", zig_query);
+    exe.root_module.addOptions("ts_queries", ts_queries);
+
+    const rust_parser_lib = b.addLibrary(.{
+        .name = "tree-sitter-rust",
+        .root_module = b.createModule(.{
+            .target = target,
+            .optimize = optimize,
+            .link_libc = true,
+        }),
+        .linkage = .static,
+    });
+    rust_parser_lib.root_module.addCSourceFile(.{ .file = ts_rust_dep.path("src/parser.c") });
+    rust_parser_lib.root_module.addCSourceFile(.{ .file = ts_rust_dep.path("src/scanner.c") });
+    rust_parser_lib.root_module.addIncludePath(ts_rust_dep.path("src"));
+    exe.root_module.linkLibrary(rust_parser_lib);
+
+    const c_parser_lib = b.addLibrary(.{
+        .name = "tree-sitter-c",
+        .root_module = b.createModule(.{
+            .target = target,
+            .optimize = optimize,
+            .link_libc = true,
+        }),
+        .linkage = .static,
+    });
+    c_parser_lib.root_module.addCSourceFile(.{ .file = ts_c_dep.path("src/parser.c") });
+    c_parser_lib.root_module.addIncludePath(ts_c_dep.path("src"));
+    exe.root_module.linkLibrary(c_parser_lib);
+
+    const zig_parser_lib = b.addLibrary(.{
+        .name = "tree-sitter-zig",
+        .root_module = b.createModule(.{
+            .target = target,
+            .optimize = optimize,
+            .link_libc = true,
+        }),
+        .linkage = .static,
+    });
+    zig_parser_lib.root_module.addCSourceFile(.{ .file = ts_zig_dep.path("src/parser.c") });
+    zig_parser_lib.root_module.addIncludePath(ts_zig_dep.path("src"));
+    exe.root_module.linkLibrary(zig_parser_lib);
+
     const run_exe = b.addRunArtifact(exe);
 
     const run = b.step("run", "descend to the pardes");
diff --git a/build.zig.zon b/build.zig.zon
index 2484f418..a0ebb8d8 100644
--- a/build.zig.zon
+++ b/build.zig.zon
@@ -20,6 +20,22 @@
             .url = "git+https://github.com/jacobsandlund/uucode#2826a37a4562284fdacd8fa029d49509cc9bffcd",
             .hash = "uucode-0.2.0-ZZjBPlK5VADj7fdoq7G8LIHzD5o6FSkcBXXrRWr4jnrA",
         },
+        .tree_sitter = .{
+            .url = "git+https://github.com/tree-sitter/zig-tree-sitter#b55d4977d61330390fc733944da06106ff972660",
+            .hash = "tree_sitter-0.26.0-8heIf3CaAQB4SSumoGzq3gLSz2rkvXxcKFyskRWqDGBz",
+        },
+        .ts_rust = .{
+            .url = "git+https://github.com/tree-sitter/tree-sitter-rust#77a3747266f4d621d0757825e6b11edcbf991ca5",
+            .hash = "N-V-__8AAHJYbgAipcanYj2A6rnj9wtGzjGzANkrOsqiFAxm",
+        },
+        .ts_c = .{
+            .url = "git+https://github.com/tree-sitter/tree-sitter-c#b780e47fc780ddc8da13afa35a3f4ed5c157823d",
+            .hash = "tree_sitter_c-0.24.2-y5boS-ptQADHoCoVfjGT_nFtFQ5LbomIkW0fxG3_cmdB",
+        },
+        .ts_zig = .{
+            .url = "git+https://github.com/maxxnino/tree-sitter-zig#a80a6e9be81b33b182ce6305ae4ea28e29211bd5",
+            .hash = "N-V-__8AADLAWgBszUZPv5YTGyO5ZNyxNcuuR8MB3mIMJrZf",
+        },
     },
     .minimum_zig_version = "0.16.0",
     .paths = .{""},
diff --git a/main.zig b/main.zig
index 5ac51a42..13450063 100644
--- a/main.zig
+++ b/main.zig
@@ -1,6 +1,12 @@
 const std = @import("std");
 const vaxis = @import("vaxis");
 const ghostty_vt = @import("ghostty-vt");
+const ts = @import("tree-sitter");
+const ts_queries = @import("ts_queries");
+
+extern fn tree_sitter_c() callconv(.c) *const ts.Language;
+extern fn tree_sitter_rust() callconv(.c) *const ts.Language;
+extern fn tree_sitter_zig() callconv(.c) *const ts.Language;
 
 extern "c" fn forkpty(
     amaster: *c_int,
@@ -46,6 +52,13 @@ const MAX_LINE_LEN: usize = 0x200;
 const MAX_HISTORY: usize = 0x100;
 const FILE_PANE: usize = 0;
 const TTY_PANE: usize = 1;
+const TS_COMMENT: u8 = 1;
+const TS_STRING: u8 = 2;
+const TS_NUMBER: u8 = 3;
+const TS_KEYWORD: u8 = 4;
+const TS_TYPE: u8 = 5;
+const TS_FUNCTION: u8 = 6;
+const TS_BUILTIN: u8 = 7;
 
 fn ticker(io: std.Io, loop: *Loop) void {
     while (true) {
@@ -146,7 +159,10 @@ fn loadFile(
     rendered_lines: *usize,
     current_path: *[MAX_LINE_LEN]u8,
     current_path_len: *usize,
+    file_text: *[]u8,
+    highlight_styles: *[MAX_LINES][MAX_LINE_LEN]u8,
 ) !void {
+    for (highlight_styles) |*row| @memset(row, 0);
     const path_len = @min(path.len, MAX_LINE_LEN - 1);
     @memcpy(current_path[0..path_len], path[0..path_len]);
     current_path[path_len] = 0;
@@ -169,9 +185,69 @@ fn loadFile(
         rendered_lines.* = 1;
         return;
     };
-    defer gpa.free(text);
 
-    rendered_lines.* = splitLines(lines, text);
+    if (file_text.*.len > 0) gpa.free(file_text.*);
+    file_text.* = text;
+    rendered_lines.* = splitLines(lines, file_text.*);
+
+    var language: ?*const ts.Language = null;
+    var query_text: []const u8 = "";
+    const ext = std.fs.path.extension(path);
+    if (std.mem.eql(u8, ext, ".c") or std.mem.eql(u8, ext, ".h")) {
+        language = tree_sitter_c();
+        query_text = ts_queries.c_highlights;
+    }
+    if (std.mem.eql(u8, ext, ".rs")) {
+        language = tree_sitter_rust();
+        query_text = ts_queries.rust_highlights;
+    }
+    if (std.mem.eql(u8, ext, ".zig") or std.mem.eql(u8, ext, ".zon")) {
+        language = tree_sitter_zig();
+        query_text = ts_queries.zig_highlights;
+    }
+
+    if (language) |lang| {
+        defer lang.destroy();
+        const parser = ts.Parser.create();
+        defer parser.destroy();
+        parser.setLanguage(lang) catch return;
+        const tree = parser.parseString(file_text.*, null) orelse return;
+        defer tree.destroy();
+
+        var error_offset: u32 = 0;
+        const query = ts.Query.create(lang, query_text, &error_offset) catch return;
+        defer query.destroy();
+        const cursor = ts.QueryCursor.create();
+        defer cursor.destroy();
+        cursor.exec(query, tree.rootNode());
+
+        while (cursor.nextMatch()) |match| {
+            for (match.captures) |capture| {
+                const capture_name = query.captureNameForId(capture.index) orelse "";
+                var style: u8 = 0;
+                if (std.mem.indexOf(u8, capture_name, "comment") != null) style = TS_COMMENT;
+                if (std.mem.indexOf(u8, capture_name, "string") != null or std.mem.indexOf(u8, capture_name, "character") != null) style = TS_STRING;
+                if (std.mem.indexOf(u8, capture_name, "number") != null or std.mem.indexOf(u8, capture_name, "float") != null or std.mem.indexOf(u8, capture_name, "boolean") != null) style = TS_NUMBER;
+                if (std.mem.indexOf(u8, capture_name, "keyword") != null or std.mem.indexOf(u8, capture_name, "include") != null or std.mem.indexOf(u8, capture_name, "operator") != null) style = TS_KEYWORD;
+                if (std.mem.indexOf(u8, capture_name, "type") != null or std.mem.indexOf(u8, capture_name, "constructor") != null) style = TS_TYPE;
+                if (std.mem.indexOf(u8, capture_name, "function") != null or std.mem.indexOf(u8, capture_name, "method") != null or std.mem.indexOf(u8, capture_name, "macro") != null) style = TS_FUNCTION;
+                if (std.mem.indexOf(u8, capture_name, "builtin") != null) style = TS_BUILTIN;
+                if (style == 0) continue;
+
+                const start = capture.node.startPoint();
+                const end = capture.node.endPoint();
+                var row: usize = start.row;
+                while (row <= end.row and row < MAX_LINES) : (row += 1) {
+                    const start_col: usize = if (row == start.row) start.column else 0;
+                    const end_col: usize = if (row == end.row) end.column else MAX_LINE_LEN - 1;
+                    var col = @min(start_col, MAX_LINE_LEN - 1);
+                    while (col < end_col and col < MAX_LINE_LEN - 1) : (col += 1) {
+                        highlight_styles[row][col] = style;
+                    }
+                }
+            }
+        }
+    }
 }
 
 fn ghostColor(color: ghostty_vt.Style.Color) vaxis.Color {
@@ -248,6 +324,9 @@ pub fn main(init: std.process.Init) !void {
     var tty_buf: [0x10000]u8 = undefined;
     var display_lines: [2][MAX_LINES][MAX_LINE_LEN]u8 = undefined;
     var tty_render_lines: [MAX_LINES][MAX_LINE_LEN]u8 = undefined;
+    var highlight_styles: [MAX_LINES][MAX_LINE_LEN]u8 = undefined;
+    var file_text: []u8 = &.{};
+    defer if (file_text.len > 0) gpa.free(file_text);
     var rendered_lines = [_]usize{ 0, 0 };
     var line_nums: [MAX_LINES][8]u8 = undefined;
     var file_scroll: usize = 0;
@@ -266,6 +345,8 @@ pub fn main(init: std.process.Init) !void {
         &rendered_lines[FILE_PANE],
         &current_path,
         &current_path_len,
+        &file_text,
+        &highlight_styles,
     );
 
     var tty = try vaxis.Tty.init(io, &tty_buf);
@@ -394,7 +475,7 @@ pub fn main(init: std.process.Init) !void {
                                     try writer.flush();
                                     if (!mouse_on_left) term.scrollViewport(.bottom);
                                 } else {
-                                    loadFile(io, gpa, found_path, &display_lines[FILE_PANE], &rendered_lines[FILE_PANE], &current_path, &current_path_len) catch {};
+                                    loadFile(io, gpa, found_path, &display_lines[FILE_PANE], &rendered_lines[FILE_PANE], &current_path, &current_path_len, &file_text, &highlight_styles) catch {};
                                     file_scroll = 0;
                                 }
                             }
@@ -479,11 +560,55 @@ pub fn main(init: std.process.Init) !void {
                 .col_offset = 0,
                 .wrap = .none,
             });
-            _ = win_file.printSegment(.{ .text = display_lines[FILE_PANE][line_idx][0..len] }, .{
-                .row_offset = @intCast(i + 1),
-                .col_offset = 5,
-                .wrap = .none,
-            });
+            var col: usize = 0;
+            while (col < len) {
+                const style_id = highlight_styles[line_idx][col];
+                var end = col + 1;
+                while (end < len and highlight_styles[line_idx][end] == style_id) : (end += 1) {}
+
+                var style: vaxis.Style = .{};
+                var rgb = term.colors.palette.current[0];
+                switch (style_id) {
+                    TS_COMMENT => {
+                        rgb = term.colors.palette.current[8];
+                        style = .{ .fg = .{ .rgb = .{ rgb.r, rgb.g, rgb.b } }, .dim = true, .italic = true };
+                    },
+                    TS_STRING => {
+                        rgb = term.colors.palette.current[2];
+                        style = .{ .fg = .{ .rgb = .{ rgb.r, rgb.g, rgb.b } } };
+                    },
+                    TS_NUMBER => {
+                        rgb = term.colors.palette.current[3];
+                        style = .{ .fg = .{ .rgb = .{ rgb.r, rgb.g, rgb.b } } };
+                    },
+                    TS_KEYWORD => {
+                        rgb = term.colors.palette.current[5];
+                        style = .{ .fg = .{ .rgb = .{ rgb.r, rgb.g, rgb.b } }, .bold = true };
+                    },
+                    TS_TYPE => {
+                        rgb = term.colors.palette.current[6];
+                        style = .{ .fg = .{ .rgb = .{ rgb.r, rgb.g, rgb.b } } };
+                    },
+                    TS_FUNCTION => {
+                        rgb = term.colors.palette.current[4];
+                        style = .{ .fg = .{ .rgb = .{ rgb.r, rgb.g, rgb.b } }, .bold = true };
+                    },
+                    TS_BUILTIN => {
+                        rgb = term.colors.palette.current[13];
+                        style = .{ .fg = .{ .rgb = .{ rgb.r, rgb.g, rgb.b } } };
+                    },
+                    else => {},
+                }
+                _ = win_file.printSegment(.{
+                    .text = display_lines[FILE_PANE][line_idx][col..end],
+                    .style = style,
+                }, .{
+                    .row_offset = @intCast(i + 1),
+                    .col_offset = @intCast(5 + col),
+                    .wrap = .none,
+                });
+                col = end;
+            }
         }
 
         i = 0;