wesl 0.7.27 → 0.7.28

package/src/LinkerUtil.ts

@@ -1,22 +1,40 @@
 import type {
   AbstractElem,
-  ContainerElem,
   DeclIdentElem,
+  ExpressionElem,
+  HasAttributes,
+  ModuleElem,
   RefIdentElem,
 } from "./AbstractElems.ts";
+import { assertUnreachableSilent } from "./Assertions.ts";
 import { srcLog } from "./Logging.ts";
 
+/** A module's top-level declarations of a given kind, narrowed to that elem type. */
+export function declsOfKind<K extends AbstractElem["kind"]>(
+  moduleElem: ModuleElem,
+  kind: K,
+): Extract<AbstractElem, { kind: K }>[] {
+  return moduleElem.decls.filter(
+    (e): e is Extract<AbstractElem, { kind: K }> => e.kind === kind,
+  );
+}
+
+/** Visit an elem and all its descendants, parent before children (pre-order). */
 export function visitAst(
   elem: AbstractElem,
   visitor: (elem: AbstractElem) => void,
 ) {
   visitor(elem);
-  if ((elem as ContainerElem).contents) {
-    const container = elem as ContainerElem;
-    container.contents.forEach(child => {
-      visitAst(child, visitor);
-    });
-  }
+  for (const child of childElems(elem)) visitAst(child, visitor);
+}
+
+/**
+ * Child elems of any AST node: the typed structural fields (attributes plus
+ * body / condition / decls / ...) in source order. Returns [] for leaf elems.
+ */
+export function childElems(elem: AbstractElem): readonly AbstractElem[] {
+  const attributes = (elem as HasAttributes).attributes ?? [];
+  return [...attributes, ...structuralFields(elem)];
 }
 
 export function identElemLog(
@@ -29,3 +47,119 @@ export function identElemLog(
     ...messages,
   );
 }
+
+/** The child elems held in a node's typed fields, in source order (attributes
+ *  are added separately by {@link childElems}). Empty for leaf kinds. The
+ *  switch is exhaustive: a new elem kind without a case fails to compile in the
+ *  default, so children can't be silently dropped. */
+function structuralFields(elem: AbstractElem): AbstractElem[] {
+  switch (elem.kind) {
+    case "module":
+      return elem.decls;
+    case "var":
+    case "gvar":
+      return [
+        ...(elem.template ?? []),
+        elem.name,
+        ...(elem.init ? [elem.init] : []),
+      ];
+    case "let":
+    case "const":
+    case "override":
+      return [elem.name, ...(elem.init ? [elem.init] : [])];
+    case "alias":
+      return [elem.name, elem.typeRef];
+    case "assert":
+      return [elem.expression];
+    case "struct":
+      return [elem.name, ...elem.members];
+    case "member":
+      return [elem.name, elem.typeRef];
+    case "type":
+      return [elem.name.refIdentElem, ...(elem.templateParams ?? [])];
+    case "expression":
+      return [elem.expression];
+    case "binary-expression":
+      return [elem.left, elem.right];
+    case "unary-expression":
+    case "parenthesized-expression":
+      return [elem.expression];
+    case "component-expression":
+      return [elem.base, elem.access];
+    case "component-member-expression":
+      return [elem.base, elem.access];
+    case "call-expression":
+      return [elem.function, ...(elem.templateArgs ?? []), ...elem.arguments];
+    case "param":
+      return [elem.name];
+    case "typeDecl":
+      return elem.typeRef ? [elem.decl, elem.typeRef] : [elem.decl];
+    case "fn":
+      return [
+        elem.name,
+        ...elem.params,
+        ...(elem.returnType ? [elem.returnType] : []),
+        elem.body,
+      ];
+    case "block":
+      return elem.body;
+    case "if":
+      return [elem.condition, elem.body, ...(elem.else ? [elem.else] : [])];
+    case "for":
+      return [elem.init, elem.condition, elem.update, elem.body].filter(
+        isDefined,
+      );
+    case "while":
+      return [elem.condition, elem.body];
+    case "loop":
+    case "continuing":
+      return [elem.body];
+    case "switch":
+      return [elem.selector, ...(elem.bodyAttributes ?? []), ...elem.clauses];
+    case "switch-clause":
+      return [...exprSelectors(elem.selectors), elem.body];
+    case "return":
+      return elem.value ? [elem.value] : [];
+    case "break":
+      return elem.condition ? [elem.condition] : [];
+    case "assign":
+      return elem.lhs.kind === "phony" ? [elem.rhs] : [elem.lhs, elem.rhs];
+    case "increment":
+    case "decrement":
+      return [elem.target];
+    case "call":
+      return [elem.call];
+    case "continue":
+    case "discard":
+    case "empty":
+      return [];
+    case "do":
+      return [elem.name, ...elem.params, elem.body];
+    // import/directive hold no child elems in typed fields (their payload isn't
+    // an AbstractElem), but they carry attributes that childElems prepends.
+    case "import":
+    case "directive":
+    // leaf kinds: no child elems and no attributes.
+    case "attribute":
+    case "name":
+    case "literal":
+    case "ref":
+    case "decl":
+    case "synthetic":
+      return [];
+    default:
+      assertUnreachableSilent(elem);
+      return [];
+  }
+}
+
+function isDefined<T>(value: T | undefined): value is T {
+  return value !== undefined;
+}
+
+/** Drop the `"default"` sentinel, keeping only real case-selector expressions. */
+function exprSelectors(
+  selectors: (ExpressionElem | "default")[],
+): ExpressionElem[] {
+  return selectors.filter((s): s is ExpressionElem => s !== "default");
+}