@thi.ng/shader-ast/ast/function.js

DSL to define shader code in TypeScript and cross-compile to GLSL, JS and other targets

import { isString } from "@thi.ng/checks/is-string";
import { assert } from "@thi.ng/errors/assert";
import { gensym } from "./idgen.js";
import { allChildren, scope, scopedChildren, walk } from "./scope.js";
import { sym } from "./sym.js";
const __defArg = (a) => {
  const [type, id, opts] = isString(a) ? [a] : a;
  return {
    tag: "arg",
    type,
    id: id || gensym(),
    opts: { q: "in", ...opts }
  };
};
function defn(type, id, _args, _body) {
  id = id || gensym();
  const args = _args.map(__defArg);
  const body = _body(...args.map((x) => sym(x.type, x.id, x.opts))).filter(
    (x) => x != null
  );
  const returns = walk(
    (n, t) => {
      if (t.tag === "ret") {
        assert(
          t.type === type,
          `wrong return type for function '${id}', expected ${type}, got ${t.type}`
        );
        n++;
      }
      return n;
    },
    scopedChildren,
    0,
    body
  );
  if (type !== "void" && !returns) {
    throw new Error(`function '${id}' must return a value of type ${type}`);
  }
  const deps = walk(
    (acc, t) => {
      if (t.tag === "call" && t.fn) {
        acc.push(t.fn);
      }
      return acc;
    },
    allChildren,
    [],
    body
  );
  const $ = (...xs) => funcall($, ...xs);
  return Object.assign($, {
    tag: "fn",
    type,
    id,
    args,
    deps,
    scope: scope(body)
  });
}
const defMain = (body) => defn("void", "main", [], body);
function ret(val) {
  return {
    tag: "ret",
    type: val ? val.type : "void",
    val
  };
}
function funcall(fn, ...args) {
  return isString(fn) ? {
    tag: "call",
    type: args[0],
    id: fn,
    args: args.slice(1)
  } : {
    tag: "call",
    type: fn.type,
    id: fn.id,
    args,
    fn
  };
}
const builtinCall = (id, type, ...args) => ({
  tag: "call_i",
  type,
  id,
  args
});
export {
  builtinCall,
  defMain,
  defn,
  funcall,
  ret
};