mol-script: macro support

src/mol-script/script/mol-script/macro.ts

-// /**
-//  * Copyright (c) 2018 Mol* contributors, licensed under MIT, See LICENSE file for more info.
-//  *
-//  * @author David Sehnal <david.sehnal@gmail.com>
-//  * @author Alexander Rose <alexander.rose@weirdbyte.de>
-//  */
+/**
+ * Copyright (c) 2018-2019 Mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
 
-// import B from '../../language/builder'
+import { MolScriptBuilder as B } from '../../language/builder';
 
-// export function getPositionalArgs(args: any) {
-//     return Object.keys(args)
-//         .filter(k => !isNaN(k as any))
-//         .map(k => +k)
-//         .sort((a, b) => a - b)
-//         .map(k => args[k]);
-// }
+export function getPositionalArgs(args: any) {
+    return Object.keys(args)
+        .filter(k => !isNaN(k as any))
+        .map(k => +k)
+        .sort((a, b) => a - b)
+        .map(k => args[k]);
+}
 
-// export function tryGetArg(args: any, name: string | number, defaultValue?: any) {
-//     return (args && args[name] !== void 0) ? args[name] : defaultValue;
-// }
+export function tryGetArg(args: any, name: string | number, defaultValue?: any) {
+    return (args && args[name] !== void 0) ? args[name] : defaultValue;
+}
 
-// export function pickArgs(args: any, ...names: string[]) {
-//     const ret = Object.create(null);
-//     let count = 0;
-//     for (let k of Object.keys(args)) {
-//         if (names.indexOf(k) >= 0) {
-//             ret[k] = args[k];
-//             count++;
-//         }
-//     }
-//     return count ? ret : void 0;
-// }
+export function pickArgs(args: any, ...names: string[]) {
+    const ret = Object.create(null);
+    let count = 0;
+    for (let k of Object.keys(args)) {
+        if (names.indexOf(k) >= 0) {
+            ret[k] = args[k];
+            count++;
+        }
+    }
+    return count ? ret : void 0;
+}
 
-// export function aggregate(property: any, fn: any, initial?: any){
-//     return B.struct.atomSet.reduce({
-//         initial: initial !== void 0 ? initial : property,
-//         value: fn([ B.struct.slot.elementSetReduce(), property ])
-//     });
-// }
+export function aggregate(property: any, fn: any, initial?: any) {
+    return B.struct.atomSet.reduce({
+        initial: initial !== void 0 ? initial : property,
+        value: fn([ B.struct.slot.elementSetReduce(), property ])
+    });
+}
\ No newline at end of file

src/mol-script/script/mol-script/symbols.ts

@@ -6,20 +6,23 @@
 
 import { UniqueArray } from '../../../mol-data/generic';
 import Expression from '../../language/expression';
-import { Argument, MSymbol } from '../../language/symbol';
+import { Argument, MSymbol, Arguments } from '../../language/symbol';
 import { MolScriptSymbolTable as MolScript } from '../../language/symbol-table';
 import Type from '../../language/type';
+import { Types as StructureQueryTypes } from '../../language/symbol-table/structure-query';
+import { MolScriptBuilder as B } from '../../language/builder';
+import { getPositionalArgs, tryGetArg } from './macro';
 
 export type MolScriptSymbol =
     | { kind: 'alias', aliases: string[], symbol: MSymbol }
     | { kind: 'macro', aliases: string[], symbol: MSymbol, translate: (args: any) => Expression }
 
 function Alias(symbol: MSymbol<any>, ...aliases: string[]): MolScriptSymbol { return { kind: 'alias', aliases, symbol }; }
-// function Macro(symbol: MSymbol<any>, translate: (args: any) => Expression, ...aliases: string[]): MolScriptSymbol {
-//     symbol.info.namespace = 'molscript-macro';
-//     symbol.id = `molscript-macro.${symbol.info.name}`;
-//     return { kind: 'macro', symbol, translate, aliases: [symbol.info.name, ...aliases] };
-// }
+function Macro(symbol: MSymbol<any>, translate: (args: any) => Expression, ...aliases: string[]): MolScriptSymbol {
+    symbol.info.namespace = 'molscript-macro';
+    symbol.id = `molscript-macro.${symbol.info.name}`;
+    return { kind: 'macro', symbol, translate, aliases: [symbol.info.name, ...aliases] };
+}
 
 export function isMolScriptSymbol(x: any): x is MolScriptSymbol {
     return x.kind === 'alias' || x.kind === 'macro';
@@ -104,26 +107,26 @@ export const SymbolTable = [
             Alias(MolScript.structureQuery.generator.empty, 'sel.atom.empty'),
             Alias(MolScript.structureQuery.generator.all, 'sel.atom.all'),
 
-            // Macro(MSymbol('sel.atom.atoms', Arguments.Dictionary({
-            //     0: Argument(Type.Bool, { isOptional: true, defaultValue: true, description: 'Test applied to each atom.' })
-            // }), Struct.Types.ElementSelection, 'A selection of singleton atom sets.'),
-            // args => B.struct.generator.atomGroups({ 'atom-test':  M.tryGetArg(args, 0, true) })),
-
-            // Macro(MSymbol('sel.atom.res', Arguments.Dictionary({
-            //     0: Argument(Type.Bool, { isOptional: true, defaultValue: true, description: 'Test applied to the 1st atom of each residue.' })
-            // }), Struct.Types.ElementSelection, 'A selection of atom sets grouped by residue.'),
-            // args => B.struct.generator.atomGroups({
-            //     'residue-test':  M.tryGetArg(args, 0, true),
-            //     'group-by': B.ammp('residueKey')
-            // })),
-
-            // Macro(MSymbol('sel.atom.chains', Arguments.Dictionary({
-            //     0: Argument(Type.Bool, { isOptional: true, defaultValue: true, description: 'Test applied to the 1st atom of each chain.' })
-            // }), Struct.Types.ElementSelection, 'A selection of atom sets grouped by chain.'),
-            // args => B.struct.generator.atomGroups({
-            //     'chain-test': M.tryGetArg(args, 0, true),
-            //     'group-by': B.ammp('chainKey')
-            // })),
+            Macro(MSymbol('sel.atom.atoms', Arguments.Dictionary({
+                0: Argument(Type.Bool, { isOptional: true, defaultValue: true, description: 'Test applied to each atom.' })
+            }), StructureQueryTypes.ElementSelection, 'A selection of singleton atom sets.'),
+            args => B.struct.generator.atomGroups({ 'atom-test':  tryGetArg(args, 0, true) })),
+
+            Macro(MSymbol('sel.atom.res', Arguments.Dictionary({
+                0: Argument(Type.Bool, { isOptional: true, defaultValue: true, description: 'Test applied to the 1st atom of each residue.' })
+            }), StructureQueryTypes.ElementSelection, 'A selection of atom sets grouped by residue.'),
+            args => B.struct.generator.atomGroups({
+                'residue-test': tryGetArg(args, 0, true),
+                'group-by': B.ammp('residueKey')
+            })),
+
+            Macro(MSymbol('sel.atom.chains', Arguments.Dictionary({
+                0: Argument(Type.Bool, { isOptional: true, defaultValue: true, description: 'Test applied to the 1st atom of each chain.' })
+            }), StructureQueryTypes.ElementSelection, 'A selection of atom sets grouped by chain.'),
+            args => B.struct.generator.atomGroups({
+                'chain-test': tryGetArg(args, 0, true),
+                'group-by': B.ammp('chainKey')
+            })),
         ],
         [
             'Modifiers',
@@ -244,9 +247,9 @@ export const SymbolTable = [
         [
             'Link Properties',
             Alias(MolScript.structureQuery.linkProperty.order, 'link.order'),
-            // Macro(MSymbol('bond.is', Arguments.List(Struct.Types.LinkFlag), Type.Bool,
-            //     `Test if the current bond has at least one (or all if partial = false) of the specified flags: ${Type.oneOfValues(Struct.Types.LinkFlag).join(', ')}`),
-            // args => B.core.flags.hasAny([B.struct.bondProperty.flags(), B.struct.type.linkFlags(M.getPositionalArgs(args))])),
+            Macro(MSymbol('link.is', Arguments.List(StructureQueryTypes.LinkFlag), Type.Bool,
+                `Test if the current link has at least one (or all if partial = false) of the specified flags: ${Type.oneOfValues(StructureQueryTypes.LinkFlag).join(', ')}`),
+            args => B.core.flags.hasAny([B.struct.linkProperty.flags(), B.struct.type.linkFlags(getPositionalArgs(args))])),
         ]
     ]
 ];
@@ -305,19 +308,25 @@ function substSymbols(expr: Expression): Expression {
     }
     if (Expression.isSymbol(expr)) {
         if (!SymbolMap[expr.name]) return expr;
-        return Expression.Symbol(SymbolMap[expr.name]!.symbol.id);
+        const s = SymbolMap[expr.name]!;
+        if (s.kind === 'alias') return Expression.Symbol(SymbolMap[expr.name]!.symbol.id);
+        throw s.translate([]);
     }
 
-    const head = substSymbols(expr.head);
+    const isMacro = Expression.isSymbol(expr.head) && !!SymbolMap[expr.head.name] && SymbolMap[expr.head.name]!.kind === 'macro';
+
+    const head = isMacro ? expr.head : substSymbols(expr.head);
     const headChanged = head !== expr.head;
     if (!expr.args) {
+        if (isMacro) return substSymbols(expr.head); // TODO: is this correct?
         return headChanged ? Expression.Apply(head) : expr;
     }
 
     let argsChanged = false;
+    let newArgs: any;
 
     if (Expression.isArgumentsArray(expr.args)) {
-        let newArgs: Expression[] = [];
+        newArgs = [];
         for (let i = 0, _i = expr.args.length; i < _i; i++) {
             const oldArg = expr.args[i];
             const newArg = substSymbols(oldArg);
@@ -325,21 +334,27 @@ function substSymbols(expr: Expression): Expression {
             newArgs[newArgs.length] = newArg;
         }
         if (!argsChanged) newArgs = expr.args;
-        if (!headChanged && !argsChanged) return expr;
-        return Expression.Apply(head, newArgs);
+        if (!isMacro && !headChanged && !argsChanged) return expr;
     } else {
-        let newArgs: any = {}
+        newArgs = {};
         for (const key of Object.keys(expr.args)) {
             const oldArg = expr.args[key];
             const newArg = substSymbols(oldArg);
             if (oldArg !== newArg) argsChanged = true;
             newArgs[key] = newArg;
         }
-        if (!headChanged && !argsChanged) return expr;
+        if (!isMacro && !headChanged && !argsChanged) return expr;
         if (!argsChanged) newArgs = expr.args;
+    }
 
-        return Expression.Apply(head, newArgs);
+    if (isMacro) {
+        const macro = SymbolMap[(expr.head as Expression.Symbol).name]!;
+        if (macro.kind !== 'macro') return Expression.Apply(head, newArgs);
+        const ret = macro.translate(newArgs);
+        return ret;
     }
+
+    return Expression.Apply(head, newArgs);
 }
 
 export function transpileMolScript(expr: Expression) {