feat(shader-ast): add defMain, allow null values in scope bodies

thi-ng · Jun 29, 2019 · de0a3da · de0a3da
1 parent 482489e
commit de0a3da
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 15 deletions.
diff --git a/packages/shader-ast/src/api.ts b/packages/shader-ast/src/api.ts
@@ -266,40 +266,40 @@ export type Arg8<
     H extends Type
 > = [Arg<A>, Arg<B>, Arg<C>, Arg<D>, Arg<E>, Arg<F>, Arg<G>, Arg<H>];
 
-export type FnBody0 = Fn0<Term<any>[]>;
-export type FnBody1<A extends Type> = Fn<Sym<A>, Term<any>[]>;
+export type FnBody0 = Fn0<ScopeBody>;
+export type FnBody1<A extends Type> = Fn<Sym<A>, ScopeBody>;
 export type FnBody2<A extends Type, B extends Type> = Fn2<
     Sym<A>,
     Sym<B>,
-    Term<any>[]
+    ScopeBody
 >;
 export type FnBody3<A extends Type, B extends Type, C extends Type> = Fn3<
     Sym<A>,
     Sym<B>,
     Sym<C>,
-    Term<any>[]
+    ScopeBody
 >;
 export type FnBody4<
     A extends Type,
     B extends Type,
     C extends Type,
     D extends Type
-> = Fn4<Sym<A>, Sym<B>, Sym<C>, Sym<D>, Term<any>[]>;
+> = Fn4<Sym<A>, Sym<B>, Sym<C>, Sym<D>, ScopeBody>;
 export type FnBody5<
     A extends Type,
     B extends Type,
     C extends Type,
     D extends Type,
     E extends Type
-> = Fn5<Sym<A>, Sym<B>, Sym<C>, Sym<D>, Sym<E>, Term<any>[]>;
+> = Fn5<Sym<A>, Sym<B>, Sym<C>, Sym<D>, Sym<E>, ScopeBody>;
 export type FnBody6<
     A extends Type,
     B extends Type,
     C extends Type,
     D extends Type,
     E extends Type,
     F extends Type
-> = Fn6<Sym<A>, Sym<B>, Sym<C>, Sym<D>, Sym<E>, Sym<F>, Term<any>[]>;
+> = Fn6<Sym<A>, Sym<B>, Sym<C>, Sym<D>, Sym<E>, Sym<F>, ScopeBody>;
 export type FnBody7<
     A extends Type,
     B extends Type,
@@ -308,7 +308,7 @@ export type FnBody7<
     E extends Type,
     F extends Type,
     G extends Type
-> = Fn7<Sym<A>, Sym<B>, Sym<C>, Sym<D>, Sym<E>, Sym<F>, Sym<G>, Term<any>[]>;
+> = Fn7<Sym<A>, Sym<B>, Sym<C>, Sym<D>, Sym<E>, Sym<F>, Sym<G>, ScopeBody>;
 export type FnBody8<
     A extends Type,
     B extends Type,
@@ -327,7 +327,7 @@ export type FnBody8<
     Sym<F>,
     Sym<G>,
     Sym<H>,
-    Term<any>[]
+    ScopeBody
 >;
 
 export type Func0<T extends Type> = Fn0<FnCall<T>>;
@@ -414,6 +414,8 @@ export type SymType = "in" | "out" | "uni";
 
 export type Precision = "lowp" | "mediump" | "highp";
 
+export type ScopeBody = (Term<any> | null | undefined)[];
+
 export interface Term<T extends Type> {
     tag: Tag;
     type: T;

diff --git a/packages/shader-ast/src/ast.ts b/packages/shader-ast/src/ast.ts
@@ -68,6 +68,7 @@ import {
     Operator,
     Prim,
     Scope,
+    ScopeBody,
     Swizzle,
     Swizzle2,
     Swizzle2_1,
@@ -817,10 +818,14 @@ export function bitxor(a: Term<any>, b: Term<any>): Op2<any> {
  * @param body
  * @param global
  */
-export const scope = (body: Term<any>[], global = false): Scope => ({
+export const scope = (body: (Term<any> | null)[], global = false): Scope => ({
     tag: "scope",
     type: "void",
-    body: body.map((x) => (x.tag === "sym" ? decl(<Sym<any>>x) : x)),
+    body: <Term<any>[]>(
+        body
+            .filter((x) => x != null)
+            .map((x) => (x!.tag === "sym" ? decl(<Sym<any>>x) : x))
+    ),
     global
 });
 
@@ -883,13 +888,17 @@ export function defn<T extends Type, A extends Type, B extends Type, C extends T
 // prettier-ignore
 export function defn<T extends Type, A extends Type, B extends Type, C extends Type, D extends Type, E extends Type, F extends Type, G extends Type, H extends Type>(type: T, name: string, args: Arg8<A,B,C,D,E,F,G,H>, body: FnBody8<A,B,C,D,E,F,G,H>): TaggedFn8<A,B,C,D,E,F,G,H,T>;
 // prettier-ignore
-export function defn(type: Type, id: string, _args: Arg<any>[], _body: (...xs: Sym<any>[]) => Term<any>[]): Func<any> {
+export function defn(type: Type, id: string, _args: Arg<any>[], _body: (...xs: Sym<any>[]) => ScopeBody): Func<any> {
     const args = _args.map(defArg);
-    const body = _body(...args.map((x) => sym(x.type, x.id, x.opts)));
+    const body = <Term<any>[]>(
+        _body(...args.map((x) => sym(x.type, x.id, x.opts))).filter(
+            (x) => x != null
+        )
+    );
     // count & check returns
     const returns = walk(
         (n, t) => {
-            if(t.tag === "ret") {
+            if (t.tag === "ret") {
                 assert(
                     t.type === type,
                     `wrong return type for function '${id}', expected ${type}, got ${
@@ -927,10 +936,17 @@ export function defn(type: Type, id: string, _args: Arg<any>[], _body: (...xs: S
         id,
         args,
         deps,
-        scope: scope(body),
+        scope: scope(body)
     });
 }
 
+/**
+ * Syntax sugar for defining `void main()` functions.
+ *
+ * @param body
+ */
+export const defMain = (body: FnBody0) => defn("void", "main", [], body);
+
 export function ret(): FuncReturn<"void">;
 export function ret<T extends Type>(val: Term<T>): FuncReturn<T>;
 export function ret(val?: Term<any>): FuncReturn<any> {