1 files changed, 78 insertions, 7 deletions

diff --git a/src/client/util/Scripting.ts b/src/client/util/Scripting.ts
index 62c2cfe85..1d0916ac0 100644
--- a/src/client/util/Scripting.ts
+++ b/src/client/util/Scripting.ts
@@ -1,5 +1,7 @@
-// import * as ts from "typescript"
-let ts = (window as any).ts;
+import * as ts from "typescript";
+export { ts };
+// export const ts = (window as any).ts;
+
 // // @ts-ignore
 // import * as typescriptlib from '!!raw-loader!../../../node_modules/typescript/lib/lib.d.ts'
 // // @ts-ignore
@@ -50,10 +52,31 @@ export namespace Scripting {
         } else {
             throw new Error("Must either register an object with a name, or give a name and an object");
         }
-        if (scriptingGlobals.hasOwnProperty(n)) {
+        if (_scriptingGlobals.hasOwnProperty(n)) {
             throw new Error(`Global with name ${n} is already registered, choose another name`);
         }
-        scriptingGlobals[n] = obj;
+        _scriptingGlobals[n] = obj;
+    }
+
+    export function makeMutableGlobalsCopy(globals?: { [name: string]: any }) {
+        return { ..._scriptingGlobals, ...(globals || {}) };
+    }
+
+    export function setScriptingGlobals(globals: { [key: string]: any }) {
+        scriptingGlobals = globals;
+    }
+
+    export function resetScriptingGlobals() {
+        scriptingGlobals = _scriptingGlobals;
+    }
+
+    // const types = Object.keys(ts.SyntaxKind).map(kind => ts.SyntaxKind[kind]);
+    export function printNodeType(node: any, indentation = "") {
+        console.log(indentation + ts.SyntaxKind[node.kind]);
+    }
+
+    export function getGlobals() {
+        return Object.keys(scriptingGlobals);
     }
 }
 
@@ -61,7 +84,8 @@ export function scriptingGlobal(constructor: { new(...args: any[]): any }) {
     Scripting.addGlobal(constructor);
 }
 
-const scriptingGlobals: { [name: string]: any } = {};
+const _scriptingGlobals: { [name: string]: any } = {};
+let scriptingGlobals: { [name: string]: any } = _scriptingGlobals;
 
 function Run(script: string | undefined, customParams: string[], diagnostics: any[], originalScript: string, options: ScriptOptions): CompileResult {
     const errors = diagnostics.some(diag => diag.category === ts.DiagnosticCategory.Error);
@@ -162,6 +186,12 @@ class ScriptingCompilerHost {
     }
 }
 
+export type Traverser = (node: ts.Node, indentation: string) => boolean | void;
+export type TraverserParam = Traverser | { onEnter: Traverser, onLeave: Traverser };
+export type Transformer = {
+    transformer: ts.TransformerFactory<ts.SourceFile>,
+    getVars?: () => { capturedVariables: { [name: string]: Field } }
+};
 export interface ScriptOptions {
     requiredType?: string;
     addReturn?: boolean;
@@ -169,11 +199,45 @@ export interface ScriptOptions {
     capturedVariables?: { [name: string]: Field };
     typecheck?: boolean;
     editable?: boolean;
+    traverser?: TraverserParam;
+    transformer?: Transformer;
+    globals?: { [name: string]: any };
+}
+
+// function forEachNode(node:ts.Node, fn:(node:any) => void);
+function forEachNode(node: ts.Node, onEnter: Traverser, onExit?: Traverser, indentation = "") {
+    return onEnter(node, indentation) || ts.forEachChild(node, (n: any) => {
+        forEachNode(n, onEnter, onExit, indentation + "    ");
+    }) || (onExit && onExit(node, indentation));
 }
 
 export function CompileScript(script: string, options: ScriptOptions = {}): CompileResult {
     const { requiredType = "", addReturn = false, params = {}, capturedVariables = {}, typecheck = true } = options;
+    if (options.globals) {
+        Scripting.setScriptingGlobals(options.globals);
+    }
     let host = new ScriptingCompilerHost;
+    if (options.traverser) {
+        const sourceFile = ts.createSourceFile('script.ts', script, ts.ScriptTarget.ES2015, true);
+        const onEnter = typeof options.traverser === "object" ? options.traverser.onEnter : options.traverser;
+        const onLeave = typeof options.traverser === "object" ? options.traverser.onLeave : undefined;
+        forEachNode(sourceFile, onEnter, onLeave);
+    }
+    if (options.transformer) {
+        const sourceFile = ts.createSourceFile('script.ts', script, ts.ScriptTarget.ES2015, true);
+        const result = ts.transform(sourceFile, [options.transformer.transformer]);
+        if (options.transformer.getVars) {
+            const newCaptures = options.transformer.getVars();
+            // tslint:disable-next-line: prefer-object-spread
+            options.capturedVariables = Object.assign(capturedVariables, newCaptures.capturedVariables) as any;
+        }
+        const transformed = result.transformed;
+        const printer = ts.createPrinter({
+            newLine: ts.NewLineKind.LineFeed
+        });
+        script = printer.printFile(transformed[0]);
+        result.dispose();
+    }
     let paramNames: string[] = [];
     if ("this" in params || "this" in capturedVariables) {
         paramNames.push("this");
@@ -188,14 +252,16 @@ export function CompileScript(script: string, options: ScriptOptions = {}): Comp
     });
     for (const key in capturedVariables) {
         if (key === "this") continue;
+        const val = capturedVariables[key];
         paramNames.push(key);
-        paramList.push(`${key}: ${capturedVariables[key].constructor.name}`);
+        paramList.push(`${key}: ${typeof val === "object" ? Object.getPrototypeOf(val).constructor.name : typeof val}`);
     }
     let paramString = paramList.join(", ");
     let funcScript = `(function(${paramString})${requiredType ? `: ${requiredType}` : ''} {
         ${addReturn ? `return ${script};` : script}
     })`;
     host.writeFile("file.ts", funcScript);
+
     if (typecheck) host.writeFile('node_modules/typescript/lib/lib.d.ts', typescriptlib);
     let program = ts.createProgram(["file.ts"], {}, host);
     let testResult = program.emit();
@@ -203,7 +269,12 @@ export function CompileScript(script: string, options: ScriptOptions = {}): Comp
 
     let diagnostics = ts.getPreEmitDiagnostics(program).concat(testResult.diagnostics);
 
-    return Run(outputText, paramNames, diagnostics, script, options);
+    const result = Run(outputText, paramNames, diagnostics, script, options);
+
+    if (options.globals) {
+        Scripting.resetScriptingGlobals();
+    }
+    return result;
 }
 
 Scripting.addGlobal(CompileScript);
\ No newline at end of file