wesl 0.6.0-pre10

package/src/Scope.ts

@@ -0,0 +1,162 @@
+import { DeclarationElem, IfAttribute, RefIdentElem } from "./AbstractElems.ts";
+import { assertThatDebug } from "./Assertions.ts";
+import { scopeValid } from "./Conditions.ts";
+import { WeslAST } from "./ParseWESL.ts";
+
+export interface SrcModule {
+  /** module path "rand_pkg::sub::foo", or "package::main" */
+  modulePath: string;
+
+  /** file path to the module for user error reporting e.g "rand_pkg:sub/foo.wesl", or "./sub/foo.wesl" */
+  debugFilePath: string;
+
+  /** original src for module */
+  src: string;
+}
+
+/** a src declaration or reference to an ident */
+export type Ident = DeclIdent | RefIdent;
+
+/** LATER change this to a Map, so that `toString` isn't accidentally a condition */
+export type Conditions = Record<string, boolean>;
+
+interface IdentBase {
+  originalName: string; // name in the source code for ident matching (may be mangled in the output)
+  id?: number; // for debugging
+}
+
+export interface RefIdent extends IdentBase {
+  kind: "ref";
+
+  // LATER these fields are set during binding, not parsing. Make a naming scheme _refersTo or a separate interface (BindingRefIdent) to make that clear
+  refersTo?: Ident; // import or decl ident in scope to which this ident refers. undefined before binding
+  std?: true; // true if this is a standard wgsl identifier (like sin, or u32)
+
+  // TODO consider tracking the current ast in BindIdents so that this field is unnecessary
+  ast: WeslAST; // AST from module that contains this ident (to find imports during decl binding)
+
+  refIdentElem: RefIdentElem; // for error reporting and mangling
+}
+
+export interface DeclIdent extends IdentBase {
+  kind: "decl";
+  mangledName?: string; // name in the output code
+  declElem?: DeclarationElem; // link to AST so that we can traverse scopes and know what elems to emit // LATER make separate GlobalDecl kind with this required
+  scope: Scope; // scope for the references within this declaration
+  isGlobal: boolean; // true if this is a global declaration (e.g. not a local variable)
+  srcModule: SrcModule; // To figure out which module this declaration is from.
+}
+
+/** tree of ident references, organized by lexical scope and partialScope . */
+export type Scope = LexicalScope | PartialScope;
+
+/** A wgsl scope */
+export interface LexicalScope extends ScopeBase {
+  kind: "scope";
+
+  /** @if condition for conditionally translating this scope */
+  ifAttribute?: IfAttribute;
+
+  /**
+   * Efficient access to declarations in this scope.
+   * constructed on demand, for module root scopes only */ // LATER consider make a special kind for root scopes
+  scopeDecls?: Map<string, DeclIdent>;
+}
+
+/** A synthetic partial scope to contain @if conditioned idents.
+ * PartialScope idents are considered to be in the wgsl lexical scope of their parent.  */
+export interface PartialScope extends ScopeBase {
+  kind: "partial";
+
+  /** @if condition for conditionally translating this scope */
+  ifAttribute?: IfAttribute; // LATER this is required, consider changing type to reflect that
+}
+
+/** common scope elements  */
+interface ScopeBase {
+  /** id for debugging */
+  id: number;
+
+  /** null for root scope in a module */
+  parent: Scope | null;
+
+  /* Child scopes and idents in lexical order  */
+  contents: (Ident | Scope)[];
+
+  /** @if conditions for conditionally translating this scope */
+  ifAttribute?: IfAttribute;
+}
+
+/** Combine two scope siblings.
+ * The first scope is mutated to append the contents of the second.  */
+export function mergeScope(a: Scope, b: Scope): void {
+  assertThatDebug(a.kind === b.kind);
+  assertThatDebug(a.parent === b.parent);
+  assertThatDebug(!b.ifAttribute);
+  a.contents = a.contents.concat(b.contents);
+}
+
+/** reset scope and ident debugging ids */
+export function resetScopeIds() {
+  scopeId = 0;
+  identId = 0;
+}
+
+let scopeId = 0;
+let identId = 0;
+
+export function nextIdentId(): number {
+  return identId++;
+}
+
+/** make a new Scope object */
+export function emptyScope(
+  parent: Scope | null,
+  kind: Scope["kind"] = "scope",
+): Omit<Scope, "ifAttribute"> {
+  const id = scopeId++;
+  return { id, kind, parent, contents: [] };
+}
+
+/** For debugging,
+ * @return true if a scope is in the rootScope tree somewhere */
+export function containsScope(rootScope: Scope, scope: Scope): boolean {
+  if (scope === rootScope) {
+    return true;
+  }
+  for (const child of rootScope.contents) {
+    if (childScope(child) && containsScope(child, scope)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+/** @returns true if the provided element of a Scope
+ *    is itself a Scope (and not an Ident) */
+export function childScope(child: Scope | Ident): child is Scope {
+  const { kind } = child;
+  return kind === "partial" || kind === "scope";
+}
+
+/** @returns true if the provided element of a Scope
+ *    is an Ident (and not a child Scope) */
+export function childIdent(child: Scope | Ident): child is Ident {
+  return !childScope(child);
+}
+
+/** find a public declaration with the given original name */
+export function publicDecl(
+  scope: Scope,
+  name: string,
+  conditions: Conditions,
+): DeclIdent | undefined {
+  for (const elem of scope.contents) {
+    if (elem.kind === "decl" && elem.originalName === name) {
+      return elem;
+    } else if (elem.kind === "partial" && scopeValid(elem, conditions)) {
+      const found = publicDecl(elem, name, conditions);
+      if (found) return found;
+    }
+  }
+}

package/src/StandardTypes.ts

@@ -0,0 +1,97 @@
+// From https://www.w3.org/TR/WGSL/#predeclared
+// Use https://github.com/wgsl-tooling-wg/wgsl-spec to regenerate these list in the future
+
+export const stdFns = `bitcast all any select arrayLength 
+  abs acos acosh asin asinh atan atanh atan2 ceil clamp cos cosh 
+  countLeadingZeros countOneBits countTrailingZeros cross 
+  degrees determinant distance dot dot4U8Packed dot4I8Packed 
+  exp exp2 extractBits faceForward firstLeadingBit firstTrailingBit 
+  floor fma fract frexp insertBits inverseSqrt ldexp length log log2
+  max min mix modf normalize pow quantizeToF16 radians reflect refract
+  reverseBits round saturate sign sin sinh smoothstep sqrt step tan tanh
+  transpose trunc
+  dpdx dpdxCoarse dpdxFine dpdy dpdyCoarse dpdyFine fwidth 
+  fwidthCoarse fwidthFine
+  textureDimensions textureGather textureGatherCompare textureLoad
+  textureNumLayers textureNumLevels textureNumSamples
+  textureSample textureSampleBias textureSampleCompare textureSampleCompareLevel
+  textureSampleGrad textureSampleLevel textureSampleBaseClampToEdge
+  textureStore
+  atomicLoad atomicStore atomicAdd atomicSub atomicMax atomicMin
+  atomicAnd atomicOr atomicXor atomicExchange atomicCompareExchangeWeak
+  pack4x8snorm pack4x8unorm pack4xI8 pack4xU8 pack4xI8Clamp pack4xU8Clamp
+  pack2x16snorm pack2x16unorm pack2x16float
+  unpack4x8snorm unpack4x8unorm unpack4xI8 unpack4xU8 
+  unpack2x16snorm unpack2x16unorm unpack2x16float
+  storageBarrier textureBarrier workgroupBarrier workgroupUniformLoad
+  subgroupAdd subgroupAll subgroupAnd subgroupAny subgroupBallot 
+  subgroupBroadcast subgroupBroadcastFirst subgroupElect 
+  subgroupExclusiveAdd subgroupExclusiveMul subgroupInclusiveAdd 
+  subgroupInclusiveMul subgroupMax subgroupMin subgroupMul subgroupOr 
+  subgroupShuffle subgroupShuffleUp subgroupShuffleXor subgroupXor
+  quadBroadcast quadSwapDiagonal quadSwapX quadSwapY`.split(/\s+/);
+
+export const sampledTextureTypes = `
+  texture_1d texture_2d texture_2d_array texture_3d 
+  texture_cube texture_cube_array
+`;
+
+export const multisampledTextureTypes = `
+  texture_multisampled_2d texture_depth_multisampled_2d
+`;
+
+export const textureStorageTypes = `
+  texture_storage_1d texture_storage_2d texture_storage_2d_array 
+  texture_storage_3d
+`;
+
+export const stdTypes = `array atomic bool f16 f32 i32 
+  mat2x2 mat2x3 mat2x4 mat3x2 mat3x3 mat3x4 mat4x2 mat4x3 mat4x4
+  mat2x2f mat2x3f mat2x4f mat3x2f mat3x3f mat3x4f
+  mat4x2f mat4x3f mat4x4f
+  mat2x2h mat2x3h mat2x4h mat3x2h mat3x3h mat3x4h
+  mat4x2h mat4x3h mat4x4h
+  u32 vec2 vec3 vec4 ptr
+  vec2i vec3i vec4i vec2u vec3u vec4u
+  vec2f vec3f vec4f vec2h vec3h vec4h
+  ${sampledTextureTypes}
+  ${multisampledTextureTypes}
+  texture_external
+  ${textureStorageTypes}
+  texture_depth_2d texture_depth_2d_array texture_depth_cube
+  texture_depth_cube_array
+  sampler sampler_comparison
+  rgba8unorm rgba8snorm rgba8uint rgba8sint
+  rgba16uint rgba16sint rgba16float 
+  r32uint r32sint r32float rg32uint rg32sint rg32float
+  rgba32uint rgba32sint rgba32float
+  bgra8unorm`.split(/\s+/);
+
+/** https://www.w3.org/TR/WGSL/#predeclared-enumerants  */
+export const stdEnumerants = `read write read_write 
+  function private workgroup uniform storage
+  rgba8unorm rgba8snorm rgba8uint rgba8sint 
+  rgba16uint rgba16sint rgba16float 
+  r32uint r32sint r32float rg32uint rg32sint rg32float
+  rgba32uint rgba32sint rgba32float bgra8unorm`.split(/\s+/);
+
+/* Note the texel formats like rgba8unorm are here because they appear in type position
+ in <templates> for texture_storage_* types. 
+ (We could parse texture_storage types specially, but user code is unlikely to alias 
+  the texture format names with e.g. a 'struct rbga8unorm .)
+*/
+
+/** return true if the name is for a built in type (not a user struct) */
+export function stdType(name: string): boolean {
+  return stdTypes.includes(name);
+}
+
+/** return true if the name is for a built in fn (not a user function) */
+export function stdFn(name: string): boolean {
+  return stdFns.includes(name) || stdType(name);
+}
+
+/** return true if the name is for a built in enumerant */
+export function stdEnumerant(name: string): boolean {
+  return stdEnumerants.includes(name);
+}

package/src/TransformBindingStructs.ts

@@ -0,0 +1,319 @@
+import { tracing } from "mini-parse";
+import {
+  AbstractElem,
+  AttributeElem,
+  BindingStructElem,
+  DeclarationElem,
+  FnElem,
+  ModuleElem,
+  SimpleMemberRef,
+  StructElem,
+  StructMemberElem,
+  SyntheticElem,
+  TypeTemplateParameter,
+} from "./AbstractElems.ts";
+import { TransformedAST, WeslJsPlugin } from "./Linker.ts";
+import { visitAst } from "./LinkerUtil.ts";
+import { findDecl } from "./LowerAndEmit.ts";
+import { minimallyMangledName } from "./Mangler.ts";
+import {
+  attributeToString,
+  contentsToString,
+  typeListToString,
+  typeParamToString,
+} from "./RawEmit.ts";
+import { textureStorage } from "./Reflection.ts";
+import { DeclIdent, RefIdent } from "./Scope.ts";
+import { filterMap } from "./Util.ts";
+
+export function bindingStructsPlugin(): WeslJsPlugin {
+  return {
+    transform: lowerBindingStructs,
+  };
+}
+
+/**
+ * Transform binding structures into binding variables by mutating the AST.
+ *
+ * First we find all the binding structs:
+ *   . find all the structs in the module by filtering the moduleElem.contents
+ *     . for each struct:
+ *       . mark any structs with that contain @group or @binding annotations as 'binding structs' and save them in a list
+ *       . (later) create reverse links from structs to struct members
+ *       . (later) visit all the binding structs and traverse to referencing structs, marking the referencing structs as binding structs too
+ * Generate synethic AST nodes for binding variables
+ *
+ * Find all references to binding struct members
+ *   . find the componound idents by traversing moduleElem.contents
+ *   . filter to find the compound idents that refer to 'binding structs'
+ *     . go from each ident to its declaration,
+ *     . declaration to typeRef reference
+ *     . typeRef to type declaration
+ *     . check type declaration to see if it's a binding struct
+ *     . record the intermediate declaration (e.g. a fn param b:Bindings from 'fn(b:Bindings)' )
+ * rewrite references to binding struct members as synthetic elements
+ *
+ * Remove the binding structs from the AST
+ * Remove the intermediate fn param declarations from the AST
+ * Add the new binding variables to the AST
+ *
+ * @return the binding structs and the mutated AST
+ */
+export function lowerBindingStructs(ast: TransformedAST): TransformedAST {
+  const clonedAst = structuredClone(ast);
+  const { moduleElem, globalNames, notableElems } = clonedAst;
+  const bindingStructs = markBindingStructs(moduleElem); // CONSIDER should we only mark bining structs referenced from the entry point?
+  markEntryTypes(moduleElem, bindingStructs);
+  const newVars = bindingStructs.flatMap(s =>
+    transformBindingStruct(s, globalNames),
+  );
+  const bindingRefs = findRefsToBindingStructs(moduleElem);
+
+  // convert references 'b.particles' to references to the synthetic var 'particles'
+  bindingRefs.forEach(({ memberRef, struct }) =>
+    transformBindingReference(memberRef, struct),
+  );
+  // remove intermediate fn param declaration b:Bindings from 'fn(b:Bindings)'
+  bindingRefs.forEach(({ intermediates }) =>
+    intermediates.forEach(e => (e.contents = [])),
+  );
+  const contents = removeBindingStructs(moduleElem);
+  moduleElem.contents = [...newVars, ...contents];
+  notableElems.bindingStructs = bindingStructs;
+  return { ...clonedAst, moduleElem };
+}
+
+export function markEntryTypes(
+  moduleElem: ModuleElem,
+  bindingStructs: BindingStructElem[],
+): void {
+  const fns = moduleElem.contents.filter(e => e.kind === "fn");
+  const fnFound = fnReferencesBindingStruct(fns, bindingStructs);
+  if (fnFound) {
+    const { fn, struct } = fnFound;
+    struct.entryFn = fn;
+  }
+}
+
+function fnReferencesBindingStruct(
+  fns: FnElem[],
+  bindingStructs: BindingStructElem[],
+): { fn: FnElem; struct: BindingStructElem } | undefined {
+  for (const fn of fns) {
+    const { params } = fn;
+    for (const p of params) {
+      const ref = p.name?.typeRef?.name as RefIdent | undefined;
+      const referencedElem = (ref?.refersTo as DeclIdent)
+        ?.declElem as StructElem;
+      const struct = bindingStructs.find(s => s === referencedElem);
+      if (struct) {
+        return { fn, struct };
+      }
+    }
+  }
+}
+
+function removeBindingStructs(moduleElem: ModuleElem): AbstractElem[] {
+  return moduleElem.contents.filter(
+    elem => elem.kind !== "struct" || !elem.bindingStruct,
+  );
+}
+
+/** mutate the AST, marking StructElems as bindingStructs
+ *  (if they contain ptrs with @group @binding annotations)
+ * @return the binding structs
+ */
+export function markBindingStructs(
+  moduleElem: ModuleElem,
+): BindingStructElem[] {
+  const structs = moduleElem.contents.filter(elem => elem.kind === "struct");
+  const bindingStructs = structs.filter(containsBinding);
+  bindingStructs.forEach(struct => (struct.bindingStruct = true));
+  // LATER also mark structs that reference a binding struct..
+  return bindingStructs as BindingStructElem[];
+}
+
+/** @return true if this struct contains a member with marked with @binding or @group */
+function containsBinding(struct: StructElem): boolean {
+  return struct.members.some(({ attributes }) => bindingAttribute(attributes));
+}
+
+function bindingAttribute(attributes?: AttributeElem[]): boolean {
+  if (!attributes) return false;
+  return attributes.some(
+    ({ attribute }) =>
+      attribute.kind === "@attribute" &&
+      (attribute.name === "binding" || attribute.name === "group"),
+  );
+}
+
+/** convert each member of the binding struct into a synthetic global variable */
+export function transformBindingStruct(
+  s: StructElem,
+  globalNames: Set<string>,
+): SyntheticElem[] {
+  return s.members.map(member => {
+    const { typeRef, name: memberName } = member;
+    const { name: typeName } = typeRef!; // members should always have a typeRef.. TODO fix typing to show this
+    const typeParameters = typeRef?.templateParams;
+
+    const varName = minimallyMangledName(memberName.name, globalNames);
+    member.mangledVarName = varName; // save new name so we can rewrite references to this member later
+    globalNames.add(varName);
+
+    const attributes =
+      member.attributes?.map(attributeToString).join(" ") ?? "";
+    const varTypes =
+      lowerPtrMember(member, typeName, typeParameters, varName) ??
+      lowerStdTypeMember(typeName, typeParameters) ??
+      lowerStorageTextureMember(typeName, typeParameters);
+    if (!varTypes) {
+      console.log("unhandled case transforming member", typeName);
+      return syntheticVar(attributes, varName, "", "??");
+    }
+
+    const { storage: storageType, varType } = varTypes;
+    return syntheticVar(attributes, varName, storageType, varType);
+  });
+}
+
+interface LoweredVarTypes {
+  storage: string;
+  varType: string;
+}
+
+function lowerPtrMember(
+  member: StructMemberElem,
+  typeName: RefIdent,
+  typeParameters: TypeTemplateParameter[] | undefined,
+  varName: string,
+): LoweredVarTypes | undefined {
+  if (typeName.originalName === "ptr") {
+    const origParams = typeParameters ?? [];
+    const newParams = [origParams[0]];
+    if (origParams[2]) newParams.push(origParams[2]);
+    const storage = typeListToString(newParams);
+
+    const varType = typeParamToString(origParams?.[1]);
+    return { storage, varType };
+  }
+}
+
+function lowerStdTypeMember(
+  typeName: RefIdent,
+  typeParameters: TypeTemplateParameter[] | undefined,
+): LoweredVarTypes | undefined {
+  if (typeof typeName !== "string") {
+    const varBaseType = typeName.std ? typeName.originalName : "??";
+    const params = typeParameters ? typeListToString(typeParameters) : "";
+    const varType = varBaseType + params;
+
+    return { varType, storage: "" };
+  }
+}
+
+function lowerStorageTextureMember(
+  typeName: RefIdent,
+  typeParameters: TypeTemplateParameter[] | undefined,
+): LoweredVarTypes | undefined {
+  if (textureStorage.test(typeName.originalName)) {
+    const params = typeParameters ? typeListToString(typeParameters) : "";
+    const varType = typeName + params;
+    return { varType, storage: "" };
+  }
+}
+
+function syntheticVar(
+  attributes: string,
+  varName: string,
+  storageTemplate: string,
+  varType: string,
+): SyntheticElem {
+  const varText = `${attributes} var${storageTemplate} ${varName} : ${varType};\n`;
+
+  const elem: SyntheticElem = {
+    kind: "synthetic",
+    text: varText,
+  };
+  return elem;
+}
+
+interface MemberRefToStruct extends StructTrace {
+  memberRef: SimpleMemberRef; // e.g. the memberRef 'b.particles'
+}
+
+interface StructTrace {
+  struct: StructElem; // e.g. the struct Bindings
+  intermediates: DeclarationElem[]; // e.g. the fn param b:Bindings from 'fn(b:Bindings)'
+}
+
+/** find all simple member references in the module that refer to binding structs */
+export function findRefsToBindingStructs(
+  moduleElem: ModuleElem,
+): MemberRefToStruct[] {
+  const members: SimpleMemberRef[] = [];
+  visitAst(moduleElem, elem => {
+    if (elem.kind === "memberRef") members.push(elem);
+  });
+  return filterMap(members, refersToBindingStruct);
+}
+
+/** @return true if this memberRef refers to a binding struct */
+function refersToBindingStruct(
+  memberRef: SimpleMemberRef,
+): MemberRefToStruct | undefined {
+  const found = traceToStruct(memberRef.name.ident);
+
+  if (found && found.struct.bindingStruct) {
+    return { memberRef, ...found };
+  }
+}
+
+/** If this identifier ultimately refers to a struct type, return the struct declaration */
+function traceToStruct(ident: RefIdent): StructTrace | undefined {
+  const decl = findDecl(ident);
+  const declElem = decl.declElem;
+  // for now only handle the case where the reference points at a fn parameter
+  if (declElem && declElem.kind === "param") {
+    const name = declElem.name.typeRef!.name;
+    if (typeof name !== "string") {
+      if (name.std) {
+        return undefined;
+      }
+
+      const paramDecl = findDecl(name);
+      const structElem = paramDecl.declElem;
+      if (structElem?.kind === "struct") {
+        return { struct: structElem, intermediates: [declElem] };
+      }
+      return undefined;
+    }
+  } else {
+    // LATER presumably handle other cases? Should this be more general, e.g. traceToType()?
+    // elemLog(
+    //   ident.refIdentElem!,
+    //   `unhandled case in traceToStruct: decl ${declElem.kind} not yet implemented`,
+    // );
+  }
+}
+
+/** Mutate the member reference elem to instead contain synthetic elem text.
+ * The new text is the mangled var name of the struct member that the memberRef refers to. */
+export function transformBindingReference(
+  memberRef: SimpleMemberRef,
+  struct: StructElem,
+): SyntheticElem {
+  const refName = memberRef.member.name;
+  const structMember = struct.members.find(m => m.name.name === refName)!;
+  if (!structMember || !structMember.mangledVarName) {
+    if (tracing) console.log(`missing mangledVarName for ${refName}`);
+    return { kind: "synthetic", text: refName };
+  }
+  const { extraComponents } = memberRef;
+  const extraText = extraComponents ? contentsToString(extraComponents) : "";
+
+  const text = structMember.mangledVarName + extraText;
+  const synthElem: SyntheticElem = { kind: "synthetic", text };
+  memberRef.contents = [synthElem];
+  return synthElem;
+}