Experiment: Viper integration

.gitignore

@@ -2,3 +2,5 @@
 /out/
 result
 *.vsix
+
+/src/generated/**/*

package.json

@@ -2,7 +2,7 @@
     "name": "vscode-motoko",
     "displayName": "Motoko",
     "description": "Motoko language support",
-    "version": "0.5.3",
+    "version": "0.5.3-viper",
     "publisher": "dfinity-foundation",
     "repository": "https://github.com/dfinity/vscode-motoko",
     "engines": {
@@ -155,9 +155,14 @@
         "typescript": "^4.8.2",
         "vsce": "^2.11.0"
     },
+    "extensionDependencies": [
+        "viper-admin.viper"
+    ],
     "scripts": {
-        "vscode:prepublish": "npm run compile",
-        "compile": "rimraf ./out && tsc -p .",
+        "vscode:prepublish": "npm run generate && npm run compile",
+        "generate": "npm run download && cd ../motoko && nix-shell --command 'make -C src moc.js' && cp -vL src/moc.js ../vscode-motoko/src/generated",
+        "download": "curl -L -o src/generated/viperserver.jar https://github.com/viperproject/viperserver/releases/download/v-2022-10-18-0728/viperserver.jar -C -",
+        "compile": "rimraf ./out && tsc -p . && cp -vr src/generated out/generated",
         "test": "jest",
         "lint": "tslint -p .",
         "package": "vsce package",

src/server/ast.ts

@@ -1,7 +1,7 @@
 import { AST } from 'motoko/lib/ast';
 import { resolveVirtualPath, tryGetFileText } from './utils';
 import { fromAST, Program } from './syntax';
-import mo from 'motoko';
+import mo from './motoko';
 
 export interface AstStatus {
     uri: string;

src/server/motoko.ts

@@ -0,0 +1,5 @@
+import getMotoko from 'motoko/lib';
+
+const mo = getMotoko(require('../generated/moc.js').Motoko as any);
+
+export default mo;

src/server/server.ts

@@ -1,7 +1,7 @@
 import { execSync } from 'child_process';
 import * as glob from 'fast-glob';
-import { existsSync, readFileSync } from 'fs';
-import mo from 'motoko';
+import { existsSync, readFileSync, unlinkSync, writeFileSync } from 'fs';
+import mo from './motoko';
 import { Node } from 'motoko/lib/ast';
 import { keywords } from 'motoko/lib/keywords';
 import * as baseLibrary from 'motoko/packages/latest/base.json';
@@ -41,6 +41,7 @@ import {
     resolveFilePath,
     resolveVirtualPath,
 } from './utils';
+import { compileViper, invalidateViper } from './viper';
 
 interface Settings {
     motoko: MotokoSettings;
@@ -500,7 +501,38 @@ function check(uri: string | TextDocument): boolean {
         }
 
         console.log('~', virtualPath);
-        let diagnostics = mo.check(virtualPath) as any as Diagnostic[];
+
+        let diagnostics: Diagnostic[] = [];
+        let regularTypeCheck = true;
+
+        // Check for a `// @viper` comment at the beginning of the Motoko file
+        if (/^\s*\/\/ *@viper/i.test(getFileText(resolvedUri))) {
+            const viperFile = resolveFilePath(
+                resolvedUri.replace(/\.mo$/, '.vpr'),
+            );
+            try {
+                // @ts-ignore
+                const result = compileViper(resolvedUri, virtualPath);
+                console.log('Viper file:', viperFile);
+                if (result.diagnostics) {
+                    diagnostics = result.diagnostics;
+                }
+                if (result.code) {
+                    const [source, _lookup] = result.code;
+                    writeFileSync(viperFile, source, 'utf8');
+                } else if (existsSync(viperFile)) {
+                    unlinkSync(viperFile);
+                }
+                regularTypeCheck = false; // Already type checked in `compileViper()`
+            } catch (err) {
+                console.error(`Error while translating to Viper: ${err}`);
+                writeFileSync(viperFile, err?.toString() || '', 'utf8');
+            }
+        }
+
+        if (regularTypeCheck) {
+            diagnostics = mo.check(virtualPath) as any as Diagnostic[];
+        }
 
         if (settings) {
             if (settings.maxNumberOfProblems > 0) {
@@ -593,10 +625,12 @@ function writeVirtual(path: string, content: string) {
     //     content = preprocessMotoko(content);
     // }
     mo.write(path, content);
+    invalidateViper(path);
 }
 
 function deleteVirtual(path: string) {
     mo.delete(path);
+    invalidateViper(path);
 }
 
 connection.onCodeAction((event) => {

src/server/viper.ts

@@ -0,0 +1,86 @@
+// Load custom `moc.js` with Viper integration
+import { Span } from 'motoko/lib/ast';
+import mo from './motoko';
+import { spawn } from 'child_process';
+import * as rpc from 'vscode-jsonrpc/node';
+import { resolve } from 'path';
+import { connect } from 'net';
+
+const serverPort = 54816; // TODO: config
+
+let connection: rpc.MessageConnection | undefined;
+
+try {
+    spawn('java', [
+        '-Xmx2048m',
+        '-Xss16m',
+        '-jar',
+        resolve(__dirname, '../generated/viperserver.jar'),
+        '--singleClient',
+        '--serverMode',
+        'LSP',
+        '--port',
+        String(serverPort),
+    ]).on('error', console.error);
+
+    const socket = connect(serverPort);
+    connection = rpc.createMessageConnection(
+        new rpc.SocketMessageReader(socket),
+        new rpc.SocketMessageWriter(socket),
+    );
+    connection.listen();
+
+    console.log('Listening to Viper LSP');
+
+    connection.sendNotification(new rpc.NotificationType('initialize'));
+} catch (err) {
+    console.error(`Error while initializing Viper LSP: ${err}`);
+}
+
+type Range = { start: Span; end: Span };
+type CompilerRange = [number, number, number, number];
+export type Lookup = (pos: CompilerRange) => CompilerRange;
+
+// Viper -> Motoko source map cache
+const sourceLookupMap = new Map<string, Lookup>();
+
+export function getMotokoSourceRange(
+    virtualPath: string,
+    { start: [a, b], end: [c, d] }: Range,
+): Range | undefined {
+    const lookup = sourceLookupMap.get(virtualPath);
+    if (!lookup) {
+        return;
+    }
+    const result = lookup([a, b, c, d]);
+    if (!result) {
+        return;
+    }
+    return {
+        start: [result[0], result[1]],
+        end: [result[2], result[3]],
+    };
+}
+
+export function compileViper(uri: string, virtualPath: string) {
+    const result = mo.compiler.viper([virtualPath]);
+    const lookup = result?.code?.[1];
+    if (lookup) {
+        sourceLookupMap.set(virtualPath, lookup);
+    }
+    if (connection) {
+        // TODO
+        connection.sendNotification(
+            new rpc.NotificationType('textDocument/didOpen'),
+            {
+                textDocument: { languageId: 'viper', version: 0, uri },
+            },
+        );
+    }
+
+    return result;
+}
+
+export function invalidateViper(virtualPath: string) {
+    sourceLookupMap.delete(virtualPath);
+}