feat(shader-ast): add/update vec coercions

- add $info() helper to allow new coercion details: - b = bvec - i = ivec - u = uvec - v = vec - add tests
thi-ng · Aug 10, 2020 · 764f4e5 · 764f4e5
1 parent 20a9d20
commit 764f4e5
Show file tree

Hide file tree

Showing 2 changed files with 115 additions and 8 deletions.
diff --git a/packages/shader-ast/src/ast/lit.ts b/packages/shader-ast/src/ast/lit.ts
@@ -4,13 +4,16 @@ import type { Lit, Term } from "../api/nodes";
 import type {
     BVec2Term,
     BVec3Term,
+    BVec4Term,
     FloatTerm,
     IntTerm,
     IVec2Term,
     IVec3Term,
+    IVec4Term,
     UintTerm,
     UVec2Term,
     UVec3Term,
+    UVec4Term,
     Vec2Term,
     Vec3Term,
     Vec4Term,
@@ -112,17 +115,20 @@ const $uvec = $gvec(wrapUint, UINT0);
 
 const $bvec = $gvec(wrapBool, FALSE);
 
+const $info = (xs: any[], info: (string | undefined)[]) =>
+    isVec(xs[0]) ? xs[0].type[0] : info[xs.length];
+
 const $gvec2 = <T extends Type>(
     type: T,
     ctor: Fn<any[], (Term<any> | undefined)[]>,
     xs: any[]
-) => lit(type, (xs = ctor(xs)), ["n", "n"][xs.length]);
+) => lit(type, (xs = ctor(xs)), $info(xs, ["n", "n"]));
 
 const $gvec3 = <T extends Type>(
     type: T,
     ctor: Fn<any[], (Term<any> | undefined)[]>,
     xs: any[]
-) => lit(type, (xs = ctor(xs)), ["n", "n", "vn"][xs.length]);
+) => lit(type, (xs = ctor(xs)), $info(xs, ["n", "n", "vn"]));
 
 const $gvec4 = <T extends Type>(
     type: T,
@@ -136,7 +142,7 @@ const $gvec4 = <T extends Type>(
             ? isVec(xs[1])
                 ? "vv"
                 : "vn"
-            : ["n", "n", , "vnn"][xs.length]
+            : $info(xs, ["n", "n", , "vnn"])
     );
 
 const $gmat = <T extends Type>(
@@ -146,8 +152,8 @@ const $gmat = <T extends Type>(
 ) => lit(type, (xs = $vec(xs)), info[xs.length]);
 
 export function vec2(): Lit<"vec2">;
-export function vec2(x: NumericB): Lit<"vec2">;
-// export function vec2(x: Term<Vec | IVec | BVec>): Lit<"vec2">;
+// prettier-ignore
+export function vec2(x: NumericB | IVec2Term | UVec2Term | BVec2Term): Lit<"vec2">;
 // prettier-ignore
 export function vec2(x: NumericB, y: NumericB): Lit<"vec2">;
 // prettier-ignore
@@ -156,7 +162,8 @@ export function vec2(...xs: any[]): Lit<"vec2"> {
 }
 
 export function vec3(): Lit<"vec3">;
-export function vec3(x: NumericB): Lit<"vec3">;
+// prettier-ignore
+export function vec3(x: NumericB | IVec3Term | UVec3Term | BVec3Term): Lit<"vec3">;
 // prettier-ignore
 export function vec3(xy: Vec2Term | IVec2Term | UVec2Term | BVec2Term, z: NumericB): Lit<"vec3">;
 // prettier-ignore
@@ -166,7 +173,8 @@ export function vec3(...xs: any[]): Lit<"vec3"> {
 }
 
 export function vec4(): Lit<"vec4">;
-export function vec4(x: NumericB): Lit<"vec4">;
+// prettier-ignore
+export function vec4(x: NumericB | IVec4Term | UVec4Term | BVec4Term): Lit<"vec4">;
 // prettier-ignore
 export function vec4(xyz: Vec3Term | IVec3Term | UVec3Term | BVec3Term, w: NumericB): Lit<"vec4">;
 // prettier-ignore

diff --git a/packages/shader-ast/test/index.ts b/packages/shader-ast/test/index.ts
@@ -1,5 +1,16 @@
 import * as assert from "assert";
-import { defn, float, isTerm, mul, ret, TRUE, vec2 } from "../src";
+import {
+    bvec2,
+    defn,
+    float,
+    isTerm,
+    ivec2,
+    Lit,
+    mul,
+    ret,
+    TRUE,
+    vec2,
+} from "../src";
 
 describe("shader-ast", () => {
     it("op2 type infer mulvv", () => {
@@ -70,4 +81,92 @@ describe("shader-ast", () => {
         assert.equal(bar.deps.length, 1);
         assert.equal(bar.deps[0].id, "foo");
     });
+
+    it("vec2 ctor", () => {
+        assert.deepEqual(vec2(), <Lit<"vec2">>{
+            tag: "lit",
+            type: "vec2",
+            info: "n",
+            val: [
+                {
+                    info: undefined,
+                    tag: "lit",
+                    type: "float",
+                    val: 0,
+                },
+            ],
+        });
+        assert.deepEqual(vec2(1), <Lit<"vec2">>{
+            tag: "lit",
+            type: "vec2",
+            info: "n",
+            val: [
+                {
+                    info: undefined,
+                    tag: "lit",
+                    type: "float",
+                    val: 1,
+                },
+            ],
+        });
+        assert.deepEqual(vec2(1, 2), <Lit<"vec2">>{
+            tag: "lit",
+            type: "vec2",
+            info: undefined,
+            val: [
+                {
+                    info: undefined,
+                    tag: "lit",
+                    type: "float",
+                    val: 1,
+                },
+                {
+                    info: undefined,
+                    tag: "lit",
+                    type: "float",
+                    val: 2,
+                },
+            ],
+        });
+        assert.deepEqual(vec2(bvec2(true)), <Lit<"vec2">>{
+            tag: "lit",
+            type: "vec2",
+            info: "b",
+            val: [
+                {
+                    info: "n",
+                    tag: "lit",
+                    type: "bvec2",
+                    val: [
+                        {
+                            info: undefined,
+                            tag: "lit",
+                            type: "bool",
+                            val: true,
+                        },
+                    ],
+                },
+            ],
+        });
+        assert.deepEqual(vec2(ivec2(1)), <Lit<"vec2">>{
+            tag: "lit",
+            type: "vec2",
+            info: "i",
+            val: [
+                {
+                    info: "n",
+                    tag: "lit",
+                    type: "ivec2",
+                    val: [
+                        {
+                            info: undefined,
+                            tag: "lit",
+                            type: "int",
+                            val: 1,
+                        },
+                    ],
+                },
+            ],
+        });
+    });
 });