From 96c2a982cb3a4936d1e8cdec63ed9a63b7f94c48 Mon Sep 17 00:00:00 2001
From: Vyacheslav Matyukhin <me@berekuk.ru>
Date: Sun, 26 Jun 2022 11:15:25 +0300
Subject: [PATCH] semantic highlight based on squiggle-lang AST

---
 packages/vscode-ext/client/src/extension.ts |  3 +
 packages/vscode-ext/client/src/highlight.ts | 90 +++++++++++++++++++++
 packages/vscode-ext/server/src/server.ts    | 14 +---
 3 files changed, 97 insertions(+), 10 deletions(-)
 create mode 100644 packages/vscode-ext/client/src/highlight.ts

diff --git a/packages/vscode-ext/client/src/extension.ts b/packages/vscode-ext/client/src/extension.ts
index 6ba040a4..ee8af595 100644
--- a/packages/vscode-ext/client/src/extension.ts
+++ b/packages/vscode-ext/client/src/extension.ts
@@ -4,6 +4,7 @@ import * as vscode from "vscode";
 import { startClient, stopClient } from "./client";
 
 import { SquiggleEditorProvider } from "./editor";
+import { registerSemanticHighlight } from "./highlight";
 import { registerPreviewCommand } from "./preview";
 
 // this method is called when your extension is activated
@@ -13,6 +14,8 @@ export function activate(context: vscode.ExtensionContext) {
 
   registerPreviewCommand(context);
 
+  registerSemanticHighlight();
+
   startClient(context);
 }
 
diff --git a/packages/vscode-ext/client/src/highlight.ts b/packages/vscode-ext/client/src/highlight.ts
new file mode 100644
index 00000000..05a90afe
--- /dev/null
+++ b/packages/vscode-ext/client/src/highlight.ts
@@ -0,0 +1,90 @@
+import * as vscode from "vscode";
+
+import { parse } from "@quri/squiggle-lang";
+import { AnyPeggyNode } from "@quri/squiggle-lang/dist/src/rescript/Reducer/Reducer_Peggy/helpers";
+
+const tokenTypes = ["class", "interface", "enum", "function", "variable"];
+const tokenModifiers = ["declaration", "documentation"];
+const legend = new vscode.SemanticTokensLegend(tokenTypes, tokenModifiers);
+
+const convertRange = (
+  range: Extract<AnyPeggyNode, { type: "Identifier" }>["location"]
+) =>
+  new vscode.Range(
+    new vscode.Position(range.start.line - 1, range.start.column - 1),
+    new vscode.Position(range.end.line - 1, range.end.column - 1)
+  );
+
+const populateTokensBuilder = (
+  tokensBuilder: vscode.SemanticTokensBuilder,
+  node: AnyPeggyNode
+  // bindings: { [key: string]: boolean }
+) => {
+  switch (node.type) {
+    case "Expression":
+      for (const child of node.nodes) {
+        populateTokensBuilder(tokensBuilder, child);
+      }
+      break;
+    case "Block":
+      for (const child of node.statements) {
+        populateTokensBuilder(tokensBuilder, child);
+      }
+      break;
+    case "LetStatement":
+      tokensBuilder.push(
+        convertRange(node.variable.location),
+        node.value.type === "Lambda" ? "function" : "variable",
+        ["declaration"]
+      );
+      populateTokensBuilder(tokensBuilder, node.value);
+      break;
+    case "Lambda":
+      for (const arg of node.args) {
+        populateTokensBuilder(tokensBuilder, arg);
+      }
+      populateTokensBuilder(tokensBuilder, node.body);
+      break;
+    case "Ternary":
+      populateTokensBuilder(tokensBuilder, node.condition);
+      populateTokensBuilder(tokensBuilder, node.trueExpression);
+      populateTokensBuilder(tokensBuilder, node.falseExpression);
+      break;
+    case "KeyValue":
+      populateTokensBuilder(tokensBuilder, node.key);
+      populateTokensBuilder(tokensBuilder, node.value);
+      break;
+    case "Identifier":
+      tokensBuilder.push(convertRange(node.location), "variable", []);
+      break;
+  }
+};
+
+export const registerSemanticHighlight = () => {
+  const provider: vscode.DocumentSemanticTokensProvider = {
+    provideDocumentSemanticTokens(
+      document: vscode.TextDocument
+    ): vscode.ProviderResult<vscode.SemanticTokens> {
+      const parseResult = parse(document.getText());
+
+      const tokensBuilder = new vscode.SemanticTokensBuilder(legend);
+      if (parseResult.tag === "Ok") {
+        populateTokensBuilder(
+          tokensBuilder,
+          parseResult.value
+          // {}
+        );
+      }
+
+      return tokensBuilder.build();
+    },
+  };
+
+  const selector = { language: "squiggle", scheme: "file" };
+
+  vscode.languages.registerDocumentSemanticTokensProvider(
+    selector,
+    provider,
+    legend
+  );
+};
diff --git a/packages/vscode-ext/server/src/server.ts b/packages/vscode-ext/server/src/server.ts
index dd22a0ff..84637a7d 100644
--- a/packages/vscode-ext/server/src/server.ts
+++ b/packages/vscode-ext/server/src/server.ts
@@ -13,6 +13,10 @@ import { parse } from "@quri/squiggle-lang";
 
 import { TextDocument } from "vscode-languageserver-textdocument";
 
+// Documentation:
+// - https://code.visualstudio.com/api/language-extensions/language-server-extension-guide
+// - https://microsoft.github.io/language-server-protocol/specifications/specification-current
+
 // Create a connection for the server, using Node's IPC as a transport.
 // Also include all preview / proposed LSP features.
 let connection = createConnection(ProposedFeatures.all);
@@ -23,17 +27,7 @@ documents.onDidChangeContent((change) => {
   validateSquiggleDocument(change.document);
 });
 
-let hasDiagnosticRelatedInformationCapability = false;
-
 connection.onInitialize((params: InitializeParams) => {
-  const capabilities = params.capabilities;
-
-  hasDiagnosticRelatedInformationCapability = !!(
-    capabilities.textDocument &&
-    capabilities.textDocument.publishDiagnostics &&
-    capabilities.textDocument.publishDiagnostics.relatedInformation
-  );
-
   const result: InitializeResult = {
     capabilities: {
       textDocumentSync: TextDocumentSyncKind.Incremental,