feat(shader-ast): add forLoop(), ternary(), fix float/int casts, docs

packages/shader-ast/src/api.ts

@@ -17,6 +17,7 @@ export type Tag =
     | "call_i"
     | "decl"
     | "fn"
+    | "for"
     | "idx"
     | "if"
     | "lit"
@@ -25,7 +26,8 @@ export type Tag =
     | "ret"
     | "scope"
     | "swizzle"
-    | "sym";
+    | "sym"
+    | "ternary";
 
 export type Type =
     | "void"
@@ -393,6 +395,12 @@ export interface Branch extends Term<"void"> {
     f?: Scope;
 }
 
+export interface Ternary<T extends Type> extends Term<T> {
+    test: Term<"bool">;
+    t: Term<T>;
+    f: Term<T>;
+}
+
 export interface FuncReturn<T extends Type> extends Term<T> {
     val?: Term<any>;
 }
@@ -499,6 +507,13 @@ export interface FnCall<T extends Type> extends Term<T> {
     info?: string;
 }
 
+export interface ForLoop extends Term<"void"> {
+    init?: Decl<any>;
+    test: Term<"bool">;
+    iter?: Term<any>;
+    body: Scope;
+}
+
 export interface TargetImpl extends Record<Tag, Fn<any, string>> {
     arg: Fn<FuncArg<any>, string>;
     assign: Fn<Assign<any>, string>;

packages/shader-ast/src/ast.ts

@@ -33,6 +33,7 @@ import {
     FnBody7,
     FnBody8,
     FnCall,
+    ForLoop,
     Func,
     FuncArg,
     FuncReturn,
@@ -74,6 +75,7 @@ import {
     TaggedFn7,
     TaggedFn8,
     Term,
+    Ternary,
     Type,
     Vec
 } from "./api";
@@ -611,3 +613,38 @@ export const ifThen = (
     t: scope(truthy),
     f: falsey ? scope(falsey) : undefined
 });
+
+export const ternary = <A extends Type, B extends A>(
+    test: Term<"bool">,
+    t: Term<A>,
+    f: Term<B>
+): Ternary<A> => ({
+    tag: "ternary",
+    type: t.type,
+    test,
+    t,
+    f
+});
+
+// prettier-ignore
+export function forLoop<T extends Type>(test: Fn<Sym<T>, Term<"bool">>, body: FnBody1<T>): ForLoop;
+// prettier-ignore
+export function forLoop<T extends Type>(init: Sym<T> | undefined, test: Fn<Sym<T>, Term<"bool">>, body: FnBody1<T>): ForLoop;
+// prettier-ignore
+export function forLoop<T extends Type>(init: Sym<T> | undefined, test: Fn<Sym<T>, Term<"bool">>, iter: Fn<Sym<T>, Term<any>>, body: FnBody1<T>): ForLoop;
+export function forLoop(...xs: any[]): ForLoop {
+    const [init, test, iter, body] =
+        xs.length === 2
+            ? [, xs[0], , xs[1]]
+            : xs.length === 3
+            ? [xs[0], xs[1], , xs[2]]
+            : xs;
+    return {
+        tag: "for",
+        type: "void",
+        init: init ? decl(init) : undefined,
+        test: test(init!),
+        iter: iter ? iter(init!) : undefined,
+        body: scope(body(init!))
+    };
+}

packages/shader-ast/src/codegen/glsl.ts

@@ -70,6 +70,11 @@ export const targetGLSL = (_ = 300) => {
         fn: (t) =>
             `${$type(t.type)} ${t.id}(${$list(t.args)}) ${emit(t.scope)}`,
 
+        for: (t) =>
+            `for(${t.init ? emit(t.init) : ""}; ${emit(t.test)}; ${
+                t.iter ? emit(t.iter) : ""
+            }) ${emit(t.body)}`,
+
         idx: (t) => `${emit(t.val)}[${emit(t.id)}]`,
 
         if: (t) => {
@@ -87,10 +92,10 @@ export const targetGLSL = (_ = 300) => {
                         ? v === Math.trunc(v)
                             ? v + ".0"
                             : String(v)
-                        : emit(v);
+                        : `float(${emit(v)})`;
                 case "i32":
                 case "u32":
-                    return isNumber(v) ? String(v) : emit(v);
+                    return isNumber(v) ? String(v) : `int(${emit(v)})`;
                 case "vec2":
                 case "vec3":
                 case "vec4":
@@ -117,7 +122,9 @@ export const targetGLSL = (_ = 300) => {
 
         swizzle: (t) => `${emit(t.val)}.${t.id}`,
 
-        sym: (t) => t.id
+        sym: (t) => t.id,
+
+        ternary: (t) => `(${emit(t.test)} ? ${emit(t.t)} : ${emit(t.f)})`
     });
 
     Object.assign(emit, <GLSLTarget>{

packages/shader-ast/src/codegen/js.ts

@@ -21,6 +21,7 @@ import {
     Func,
     Lit,
     Operator,
+    Scope,
     Swizzle,
     Sym,
     Term
@@ -31,6 +32,37 @@ import { defTarget } from "./target";
 type Mat = m.Mat;
 type Vec = v.Vec;
 
+export interface JSTarget extends Fn<Term<any>, string> {
+    /**
+     * Compiles given AST to JavaScript, using optional `env` as backend
+     * for various operators / builtins. If `env` is not given the
+     * bundled `JS_DEFAULT_ENV` is used (based on thi.ng/vectors and
+     * thi.ng/matrices packages).
+     *
+     * Any functions defined in the given AST will be exported using
+     * their defined name via the returned object.
+     *
+     * ```
+     * const js = targetJS();
+     * const module = js.compile(
+     *   defn("f32", "foo", [["f32"]], (x)=> [ret(mul(x, float(10)))])
+     * );
+     *
+     * module.foo(42)
+     * // 420
+     *
+     * module.foo.toString()
+     * // function foo(_sym0) {
+     * // return (_sym0 * 10);
+     * // }
+     * ```
+     *
+     * @param tree
+     * @param env
+     */
+    compile(tree: Term<any>, env?: JSEnv): any;
+}
+
 export interface JSBuiltins<T> {
     abs: Fn<T, T>;
     acos: Fn<T, T>;
@@ -132,10 +164,6 @@ export interface JSEnv {
     mat4: JSBuiltinsMat;
 }
 
-export interface JSTarget extends Fn<Term<any>, string> {
-    compile(tree: Term<any>, env: JSEnv, exports: Func<any>[]): any;
-}
-
 export const JS_DEFAULT_ENV: JSEnv = {
     vec2n: (n) => v.setN2([], n),
     vec3n: (n) => v.setN3([], n),
@@ -396,6 +424,7 @@ export const targetJS = () => {
         "/": "div",
         "||": "or",
         "&&": "and",
+        // TODO below
         "|": "bitor",
         "&": "bitand",
         "^": "bitxor",
@@ -463,6 +492,11 @@ const mat4 = env.mat4;
             "\n */\n" +
             `function ${t.id}(${$list(t.args)}) ${emit(t.scope)}`,
 
+        for: (t) =>
+            `for(${t.init ? emit(t.init) : ""}; ${emit(t.test)}; ${
+                t.iter ? emit(t.iter) : ""
+            }) ${emit(t.body)}`,
+
         idx: (t) => `${emit(t.val)}[${emit(t.id)}]`,
 
         if: (t) => {
@@ -476,9 +510,11 @@ const mat4 = env.mat4;
                 case "bool":
                     return String(!!v);
                 case "f32":
+                    return isNumber(v) ? String(v) : emit(v);
                 case "i32":
+                    return isNumber(v) ? String(v) : `(${emit(v)} | 0)`;
                 case "u32":
-                    return isNumber(v) ? String(v) : emit(v);
+                    return isNumber(v) ? String(v) : `(${emit(v)} >>> 0)`;
                 case "vec2":
                 case "vec3":
                 case "vec4":
@@ -518,12 +554,25 @@ const mat4 = env.mat4;
         swizzle: (t) =>
             `env.swizzle${t.id.length}(${emit(t.val)}, ${$swizzle(t.id)})`,
 
-        sym: (t) => t.id
+        sym: (t) => t.id,
+
+        ternary: (t) => `(${emit(t.test)} ? ${emit(t.t)} : ${emit(t.f)})`
     });
 
     Object.assign(emit, <JSTarget>{
-        compile: (tree, env, fns) => {
-            const exports = fns.map((f) => f.id + ": " + f.id).join(",\n");
+        compile: (tree, env = JS_DEFAULT_ENV) => {
+            const exports =
+                tree.tag === "scope"
+                    ? (<Scope>tree).body
+                          .filter((x) => x.tag === "fn")
+                          .map(
+                              (f) =>
+                                  (<Func<any>>f).id + ": " + (<Func<any>>f).id
+                          )
+                          .join(",\n")
+                    : tree.tag === "fn"
+                    ? `${(<Func<any>>tree).id}: ${(<Func<any>>tree).id}`
+                    : "";
             return new Function(
                 "env",
                 PRELUDE + emit(tree) + "\nreturn {\n" + exports + "\n};"