tree-sitter fun 👀

Author: Strum355

client/src/commands.ts

@@ -1,3 +1,4 @@
+import path = require('path')
 import * as vscode from 'vscode'
 import * as lsp from 'vscode-languageclient'
 import { Extension } from './extension'
@@ -30,7 +31,7 @@ export function virtualMergedDocument(e: Extension): Command {
         command: 'virtualMerge',
         arguments: [path]
       })
-    } catch(e) {}
+    } catch (e) { }
 
     return content
   }
@@ -40,17 +41,67 @@ export function virtualMergedDocument(e: Extension): Command {
     onDidChange = this.onDidChangeEmitter.event
 
     provideTextDocumentContent(uri: vscode.Uri, __: vscode.CancellationToken): vscode.ProviderResult<string> {
-      return getVirtualDocument(uri.path)
+      return getVirtualDocument(uri.path.replace('.flattened' + path.extname(uri.path), path.extname(uri.path)))
     }
   }
 
   e.context.subscriptions.push(vscode.workspace.registerTextDocumentContentProvider('mcglsl', docProvider))
 
   return async () => {
+    if (vscode.window.activeTextEditor.document.languageId != 'glsl') return
+
+    const uri = vscode.window.activeTextEditor.document.uri.path
+      .substring(0, vscode.window.activeTextEditor.document.uri.path.lastIndexOf('.'))
+      + '.flattened.'
+      + vscode.window.activeTextEditor.document.uri.path
+        .slice(vscode.window.activeTextEditor.document.uri.path.lastIndexOf('.') + 1)
+    const path = vscode.Uri.parse(`mcglsl:${uri}`)
+    
+    const doc = await vscode.workspace.openTextDocument(path)
+    docProvider.onDidChangeEmitter.fire(path)
+    await vscode.window.showTextDocument(doc, {
+      viewColumn: vscode.ViewColumn.Two,
+      preview: true
+    })
+  }
+}
+
+export function parseTree(e: Extension): Command {
+  const getVirtualDocument = async (path: string): Promise<string | null> => {
+    let content: string = ''
+    try {
+      content = await e.lspClient.sendRequest<string>(lsp.ExecuteCommandRequest.type.method, {
+        command: 'parseTree',
+        arguments: [path]
+      })
+    } catch (e) { }
+
+    return content
+  }
+
+  const docProvider = new class implements vscode.TextDocumentContentProvider {
+    onDidChangeEmitter = new vscode.EventEmitter<vscode.Uri>()
+    onDidChange = this.onDidChangeEmitter.event
+
+    provideTextDocumentContent(uri: vscode.Uri, __: vscode.CancellationToken): vscode.ProviderResult<string> {
+      if (uri.path.includes('.flattened.')) return ''
+      return getVirtualDocument(uri.path.substring(0, uri.path.lastIndexOf('.')))
+    }
+  }
+
+  e.context.subscriptions.push(vscode.workspace.registerTextDocumentContentProvider('mcglsl', docProvider))
+
+  return async () => {
+    if (vscode.window.activeTextEditor.document.languageId != 'glsl') return
+
     const uri = vscode.window.activeTextEditor.document.uri
-    const path = vscode.Uri.parse('mcglsl:' + uri.path)
+    const path = vscode.Uri.parse(`mcglsl:${uri.path}.ast`)
+    
     const doc = await vscode.workspace.openTextDocument(path)
     docProvider.onDidChangeEmitter.fire(path)
-    await vscode.window.showTextDocument(doc, {preview: true})
+    await vscode.window.showTextDocument(doc, {
+      viewColumn: vscode.ViewColumn.Two,
+      preview: true
+    })
   }
 }

client/src/extension.ts

@@ -48,6 +48,7 @@ export class Extension {
     this.registerCommand('graphDot', commands.generateGraphDot)
     this.registerCommand('restart', commands.restartExtension)
     this.registerCommand('virtualMerge', commands.virtualMergedDocument)
+    this.registerCommand('parseTree', commands.parseTree)
 
     log.info('starting language server...')
 

package.json

@@ -41,6 +41,11 @@
         "command": "mcglsl.virtualMerge",
         "title": "Show flattened file",
         "category": "Minecraft Shader"
+      },
+      {
+        "command": "mcglsl.parseTree",
+        "title": "Show parse tree for file",
+        "category": "Minecraft Shader"
       }
     ],
     "languages": [

server/Cargo.lock

@@ -30,6 +30,8 @@ slog-atomic = "3.1"
 once_cell = "1.7"
 rand = "0.8"
 arc-swap = "1.5.0"
+tree-sitter = "0.20.6"
+tree-sitter-glsl = "0.1.2"
 
 [dev-dependencies]
 tempdir = "0.3"

server/Cargo.toml

@@ -3,9 +3,11 @@ use std::{collections::HashMap, path::Path};
 use serde_json::Value;
 
 use anyhow::{format_err, Result};
+use slog_scope::info;
 
 pub mod graph_dot;
 pub mod merged_includes;
+pub mod parse_tree;
 
 pub struct CustomCommandProvider {
     commands: HashMap<String, Box<dyn Invokeable>>,
@@ -20,6 +22,9 @@ impl CustomCommandProvider {
 
     pub fn execute(&self, command: &str, args: &[Value], root_path: &Path) -> Result<Value> {
         if self.commands.contains_key(command) {
+            info!("running command";
+                "command" => command,
+                "args" => format!("[{}]", args.iter().map(|v| serde_json::to_string(v).unwrap()).collect::<Vec<String>>().join(", ")));
             return self.commands.get(command).unwrap().run_command(root_path, args);
         }
         Err(format_err!("command doesn't exist"))

server/src/commands/mod.rs

@@ -0,0 +1,94 @@
+use std::{
+    cell::RefCell,
+    fs,
+    path::{Path, PathBuf},
+    rc::Rc,
+};
+
+use anyhow::{format_err, Result};
+use serde_json::Value;
+use slog_scope::warn;
+use tree_sitter::{Parser, TreeCursor};
+
+use crate::url_norm::FromJson;
+
+use super::Invokeable;
+
+pub struct TreeSitterSExpr {
+    pub tree_sitter: Rc<RefCell<Parser>>,
+}
+
+impl Invokeable for TreeSitterSExpr {
+    fn run_command(&self, _: &Path, arguments: &[Value]) -> Result<Value> {
+        let path = PathBuf::from_json(arguments.get(0).unwrap())?;
+
+        warn!("parsing"; "path" => path.to_str().unwrap().to_string());
+
+        let source = fs::read_to_string(path)?;
+
+        let tree = match self.tree_sitter.borrow_mut().parse(source, None) {
+            Some(tree) => tree,
+            None => return Err(format_err!("tree-sitter parsing resulted in no parse tree")),
+        };
+
+        let mut cursor = tree.walk();
+
+        let rendered = render_parse_tree(&mut cursor);
+
+        Ok(serde_json::value::Value::String(rendered))
+    }
+}
+
+fn render_parse_tree(cursor: &mut TreeCursor) -> String {
+    let mut string = String::new();
+
+    let mut indent = 0;
+    let mut visited_children = false;
+
+    loop {
+        let node = cursor.node();
+
+        let display_name = if node.is_missing() {
+            format!("MISSING {}", node.kind())
+        } else if node.is_named() {
+            node.kind().to_string()
+        } else {
+            "".to_string()
+        };
+
+        if visited_children {
+            if cursor.goto_next_sibling() {
+                visited_children = false;
+            } else if cursor.goto_parent() {
+                visited_children = true;
+                indent -= 1;
+            } else {
+                break;
+            }
+        } else {
+            if !display_name.is_empty() {
+                let start = node.start_position();
+                let end = node.end_position();
+
+                let field_name = match cursor.field_name() {
+                    Some(name) => name.to_string() + ": ",
+                    None => "".to_string(),
+                };
+
+                string += ("  ".repeat(indent)
+                    + format!("{}{} [{}, {}] - [{}, {}]\n", field_name, display_name, start.row, start.column, end.row, end.column)
+                        .trim_start())
+                .as_str();
+            }
+
+            if cursor.goto_first_child() {
+                visited_children = false;
+                indent += 1;
+            } else {
+                visited_children = true;
+            }
+        }
+    }
+
+    string
+}

server/src/commands/parse_tree.rs

@@ -7,6 +7,7 @@ use petgraph::stable_graph::NodeIndex;
 use serde::Deserialize;
 use serde_json::{from_value, Value};
 
+use tree_sitter::Parser;
 use url_norm::FromUrl;
 
 use walkdir::WalkDir;
@@ -65,12 +66,16 @@ fn main() {
 
     let cache_graph = graph::CachedStableGraph::new();
 
+    let mut parser = Parser::new();
+    parser.set_language(tree_sitter_glsl::language()).unwrap();
+
     let mut langserver = MinecraftShaderLanguageServer {
         endpoint: endpoint_output.clone(),
         graph: Rc::new(RefCell::new(cache_graph)),
         root: "".into(),
         command_provider: None,
         opengl_context: Rc::new(opengl::OpenGlContext::new()),
+        tree_sitter: Rc::new(RefCell::new(parser)),
         log_guard: Some(guard),
     };
 
@@ -87,6 +92,12 @@ fn main() {
                 graph: langserver.graph.clone(),
             }),
         ),
+        (
+            "parseTree",
+            Box::new(commands::parse_tree::TreeSitterSExpr {
+                tree_sitter: langserver.tree_sitter.clone(),
+            }),
+        ),
     ]));
 
     LSPEndpoint::run_server_from_input(&mut stdin().lock(), endpoint_output, langserver);
@@ -98,6 +109,7 @@ struct MinecraftShaderLanguageServer {
     root: PathBuf,
     command_provider: Option<commands::CustomCommandProvider>,
     opengl_context: Rc<dyn opengl::ShaderValidator>,
+    tree_sitter: Rc<RefCell<Parser>>,
     log_guard: Option<slog_scope::GlobalLoggerGuard>,
 }
 

server/src/main.rs

@@ -50,6 +50,7 @@ fn new_temp_server(opengl_context: Option<Box<dyn opengl::ShaderValidator>>) ->
         command_provider: None,
         opengl_context: context.into(),
         log_guard: Some(guard),
+        tree_sitter: Rc::new(RefCell::new(Parser::new())),
     }
 }