shaderkit 0.2.0 → 0.3.0

package/src/parser.ts

@@ -1,568 +1,660 @@
 import {
-  type AST,
-  Type,
-  Identifier,
-  Literal,
-  CallExpression,
-  UnaryExpression,
-  MemberExpression,
-  TernaryExpression,
-  BinaryExpression,
+  ArraySpecifier,
+  AssignmentOperator,
+  BinaryOperator,
   BlockStatement,
-  FunctionDeclaration,
-  VariableDeclaration,
-  VariableDeclarator,
-  ContinueStatement,
   BreakStatement,
+  ConstantQualifier,
+  ContinueStatement,
   DiscardStatement,
-  ReturnStatement,
-  IfStatement,
-  WhileStatement,
-  ForStatement,
   DoWhileStatement,
-  SwitchStatement,
-  SwitchCase,
-  StructDeclaration,
+  Expression,
+  ForStatement,
+  FunctionDeclaration,
+  FunctionParameter,
+  Identifier,
+  IfStatement,
+  InterpolationQualifier,
+  LayoutQualifier,
+  Literal,
+  LogicalOperator,
+  ParameterQualifier,
+  PrecisionQualifier,
   PrecisionStatement,
-  ArrayExpression,
   PreprocessorStatement,
-} from './ast'
-import { type Token, tokenize } from './tokenizer'
-
-const UNARY_OPERATORS = ['+', '-', '~', '!', '++', '--']
-
-const BINARY_OPERATORS = [
-  '>>=',
-  '<<=',
-  '|=',
-  '&=',
-  '^=',
-  '%=',
-  '/=',
-  '*=',
-  '-=',
-  '+=',
-  '=',
-  '?',
-  '||',
-  '^^',
-  '&&',
-  '|',
-  '^',
-  '&',
-  '!=',
-  '==',
-  '>=',
-  '<=',
-  '>',
-  '<',
-  '>>',
-  '<<',
-  '+',
-  '-',
-  '%',
-  '/',
-  '*',
-]
+  Program,
+  ReturnStatement,
+  Statement,
+  StorageQualifier,
+  StructDeclaration,
+  SwitchCase,
+  SwitchStatement,
+  UnaryOperator,
+  UniformDeclarationBlock,
+  UpdateOperator,
+  VariableDeclaration,
+  VariableDeclarator,
+  WhileStatement,
+} from './ast.js'
+import { type Token, tokenize } from './tokenizer.js'
+
+// https://engineering.desmos.com/articles/pratt-parser
+// https://matklad.github.io/2020/04/13/simple-but-powerful-pratt-parsing.html
+
+// 5.1 Operators
+enum Precedence {
+  LOWEST,
+  COMMA,
+  ASSIGN,
+  LOGICAL_OR,
+  LOGICAL_XOR,
+  LOGICAL_AND,
+  BITWISE_OR,
+  BITWISE_XOR,
+  BITWISE_AND,
+  TERNARY,
+  COMPARE,
+  SHIFT,
+  ADD,
+  MULTIPLY,
+  UNARY_PREFIX,
+  UNARY_POSTFIX,
+  MEMBER,
+}
+
+const PREFIX_OPERATOR_PRECEDENCE: Record<string, Precedence> = {
+  '~': Precedence.UNARY_PREFIX,
+  '!': Precedence.UNARY_PREFIX,
+  '--': Precedence.UNARY_PREFIX,
+  '++': Precedence.UNARY_PREFIX,
+  '-': Precedence.UNARY_PREFIX,
+  '+': Precedence.UNARY_PREFIX,
+}
+
+const POSTFIX_OPERATOR_PRECEDENCE: Record<string, Precedence> = {
+  '--': Precedence.UNARY_POSTFIX,
+  '++': Precedence.UNARY_POSTFIX,
+  '(': Precedence.LOWEST,
+  '[': Precedence.MEMBER,
+  '.': Precedence.MEMBER,
+}
+
+const INFIX_OPERATOR_PRECEDENCE_LEFT: Record<string, Precedence> = {
+  '||': Precedence.LOGICAL_OR,
+  '^^': Precedence.LOGICAL_XOR,
+  '&&': Precedence.LOGICAL_AND,
+  '|': Precedence.BITWISE_OR,
+  '^': Precedence.BITWISE_XOR,
+  '&': Precedence.BITWISE_AND,
+  '==': Precedence.COMPARE,
+  '>': Precedence.COMPARE,
+  '>=': Precedence.COMPARE,
+  '<': Precedence.COMPARE,
+  '<=': Precedence.COMPARE,
+  '!=': Precedence.COMPARE,
+  '<<': Precedence.SHIFT,
+  '>>': Precedence.SHIFT,
+  '+': Precedence.ADD,
+  '-': Precedence.ADD,
+  '*': Precedence.MULTIPLY,
+  '/': Precedence.MULTIPLY,
+  '%': Precedence.MULTIPLY,
+}
+
+const INFIX_OPERATOR_PRECEDENCE_RIGHT: Record<string, Precedence> = {
+  '=': Precedence.ASSIGN,
+  '+=': Precedence.ASSIGN,
+  '&=': Precedence.ASSIGN,
+  '|=': Precedence.ASSIGN,
+  '^=': Precedence.ASSIGN,
+  '/=': Precedence.ASSIGN,
+  '*=': Precedence.ASSIGN,
+  '%=': Precedence.ASSIGN,
+  '<<=': Precedence.ASSIGN,
+  '>>=': Precedence.ASSIGN,
+  '-=': Precedence.ASSIGN,
+  '?': Precedence.TERNARY,
+}
 
 // TODO: this is GLSL-only, separate language constants
 const TYPE_REGEX = /^(void|bool|float|u?int|[uib]?vec\d|mat\d(x\d)?)$/
+// TODO: must be ordered: invariant interpolation storage precision storage parameter precision
+// const cannot be used with storage or parameter qualifiers
 const QUALIFIER_REGEX = /^(const|uniform|in|out|inout|centroid|flat|smooth|invariant|lowp|mediump|highp)$/
 const VARIABLE_REGEX = new RegExp(`${TYPE_REGEX.source}|${QUALIFIER_REGEX.source}|layout`)
 
-const isDeclaration = RegExp.prototype.test.bind(VARIABLE_REGEX)
-
-const isOpen = RegExp.prototype.test.bind(/^[\(\[\{]$/)
-const isClose = RegExp.prototype.test.bind(/^[\)\]\}]$/)
-
+const SCOPE_DELTAS: Record<string, number> = {
+  // Open
+  '(': 1,
+  '[': 1,
+  '{': 1,
+  // Close
+  ')': -1,
+  ']': -1,
+  '}': -1,
+}
 function getScopeDelta(token: Token): number {
-  if (isOpen(token.value)) return 1
-  if (isClose(token.value)) return -1
-  return 0
+  return SCOPE_DELTAS[token.value] ?? 0
 }
 
-let tokens: Token[] = []
-let i: number = 0
+function consume(tokens: Token[], expected?: string): Token {
+  const token = tokens.shift()
 
-function readUntil(value: string, body: Token[], offset: number = 0): Token[] {
-  const output: Token[] = []
-  let scopeIndex = 0
-
-  while (offset < body.length) {
-    const token = body[offset++]
-    output.push(token)
-
-    scopeIndex += getScopeDelta(token)
-    if (scopeIndex === 0 && token.value === value) break
+  if (token === undefined && expected !== undefined) {
+    throw new SyntaxError(`Expected "${expected}"`)
+  } else if (token === undefined) {
+    throw new SyntaxError('Unexpected end of input')
+  } else if (expected !== undefined && token.value !== expected) {
+    throw new SyntaxError(`Expected "${expected}" got "${token.value}"`)
   }
 
-  return output
+  return token
 }
 
-function consumeUntil(value: string): Token[] {
-  const output = readUntil(value, tokens, i)
-  i += output.length
-  return output
-}
-
-function parseExpression(body: Token[]): AST | null {
-  if (body.length === 0) return null
-
-  const first = body[0]
-  const last = body[body.length - 1]
-  if (UNARY_OPERATORS.includes(first.value)) {
-    const right = parseExpression(body.slice(1))!
-    return new UnaryExpression(first.value, null, right)
-  } else if (UNARY_OPERATORS.includes(last.value)) {
-    const left = parseExpression(body.slice(0, body.length - 1))!
-    return new UnaryExpression(last.value, left, null)
-  }
-
-  if (first.value === '(') {
-    const leftBody = readUntil(')', body)
-    const left = parseExpression(leftBody.slice(1, leftBody.length - 1))!
-
-    const operator = body[leftBody.length]
-    if (operator) {
-      const rightBody = body.slice(leftBody.length + 1)
-      const right = parseExpression(rightBody)!
-
-      return new BinaryExpression(operator.value, left, right)
+function parseExpression(tokens: Token[], minBindingPower: number = 0): Expression {
+  let token = consume(tokens)
+
+  let lhs: Expression
+  if (token.type === 'identifier' || token.type === 'keyword') {
+    lhs = { type: 'Identifier', name: token.value }
+  } else if (token.type === 'bool' || token.type === 'float' || token.type === 'int') {
+    lhs = { type: 'Literal', value: token.value }
+  } else if (token.type === 'symbol' && token.value === '(') {
+    lhs = parseExpression(tokens, 0)
+    consume(tokens, ')')
+  } else if (token.type === 'symbol' && token.value in PREFIX_OPERATOR_PRECEDENCE) {
+    const rightBindingPower = PREFIX_OPERATOR_PRECEDENCE[token.value]
+    const rhs = parseExpression(tokens, rightBindingPower)
+    if (token.value === '--' || token.value === '++') {
+      lhs = { type: 'UpdateExpression', operator: token.value, prefix: true, argument: rhs }
+    } else {
+      lhs = { type: 'UnaryExpression', operator: token.value as UnaryOperator, prefix: true, argument: rhs }
     }
-
-    return left
+  } else {
+    throw new SyntaxError(`Unexpected token: "${token.value}"`)
   }
 
-  let scopeIndex = 0
-
-  for (const operator of BINARY_OPERATORS) {
-    for (let i = 0; i < body.length; i++) {
-      const token = body[i]
-      if (token.type !== 'symbol') continue
+  while (tokens.length) {
+    token = tokens[0]
 
-      scopeIndex += getScopeDelta(token)
+    if (token.value in POSTFIX_OPERATOR_PRECEDENCE) {
+      const leftBindingPower = POSTFIX_OPERATOR_PRECEDENCE[token.value]
+      if (leftBindingPower < minBindingPower) break
 
-      if (scopeIndex === 0 && token.value === operator) {
-        if (operator === '?') {
-          const testBody = body.slice(0, i)
-          const consequentBody = readUntil(':', body, i + 1).slice(0, -1)
-          const alternateBody = body.slice(i + consequentBody.length + 2)
+      consume(tokens)
 
-          const test = parseExpression(testBody)!
-          const consequent = parseExpression(consequentBody)!
-          const alternate = parseExpression(alternateBody)!
+      if (token.value === '(') {
+        const args: Expression[] = []
 
-          return new TernaryExpression(test, consequent, alternate)
+        while (tokens[0]?.value !== ')') {
+          args.push(parseExpression(tokens, 0))
+          if (tokens[0]?.value !== ')') consume(tokens, ',')
+        }
+        consume(tokens, ')')
+
+        lhs = { type: 'CallExpression', callee: lhs, arguments: args }
+      } else if (token.value === '[') {
+        const rhs = tokens[0]?.value !== ']' ? parseExpression(tokens, 0) : null
+        consume(tokens, ']')
+
+        if (tokens[0]?.value === '(') {
+          consume(tokens, '(')
+          const elements: Expression[] = []
+
+          while ((tokens[0] as Token | undefined)?.value !== ')') {
+            elements.push(parseExpression(tokens, 0))
+            if ((tokens[0] as Token | undefined)?.value !== ')') consume(tokens, ',')
+          }
+          consume(tokens, ')')
+
+          const typeSpecifier: ArraySpecifier = {
+            type: 'ArraySpecifier',
+            typeSpecifier: lhs as Identifier,
+            dimensions: [rhs as Literal | null],
+          }
+          lhs = { type: 'ArrayExpression', typeSpecifier, elements }
         } else {
-          const left = parseExpression(body.slice(0, i))!
-          const right = parseExpression(body.slice(i + 1, body.length))!
-
-          return new BinaryExpression(operator, left, right)
+          lhs = { type: 'MemberExpression', object: lhs, property: rhs!, computed: true }
         }
+      } else if (token.value === '.') {
+        const rhs = parseExpression(tokens, 0)
+        lhs = { type: 'MemberExpression', object: lhs, property: rhs, computed: false }
+      } else if (token.value === '--' || token.value === '++') {
+        lhs = { type: 'UpdateExpression', operator: token.value as UpdateOperator, prefix: false, argument: lhs }
+      } else {
+        lhs = { type: 'UnaryExpression', operator: token.value as UnaryOperator, prefix: false, argument: lhs }
       }
+    } else if (token.value in INFIX_OPERATOR_PRECEDENCE_LEFT) {
+      const precedence = INFIX_OPERATOR_PRECEDENCE_LEFT[token.value]
+      const leftBindingPower = precedence - 1
+      const rightBindingPower = precedence
 
-      if (scopeIndex < 0) {
-        return parseExpression(body.slice(0, i))
-      }
-    }
-  }
-
-  if (first.type === 'bool' || first.type === 'int' || first.type === 'float') {
-    return new Literal(first.value)
-  } else if (first.type === 'identifier' || first.type === 'keyword') {
-    const second = body[1]
-
-    if (!second) {
-      return new Identifier(first.value)
-    } else if (second.value === '(') {
-      const callee = new Identifier(first.value)
-      const args: AST[] = []
+      if (leftBindingPower < minBindingPower) break
 
-      const scope = readUntil(')', body, 1).slice(1, -1)
+      consume(tokens)
 
-      let j = 0
-      while (j < scope.length) {
-        const line = readUntil(',', scope, j)
-        j += line.length
-        if (line[line.length - 1]?.value === ',') line.pop() // skip ,
-
-        const arg = parseExpression(line)
-        if (arg) args.push(arg)
+      if (token.value === '||' || token.value === '&&' || token.value === '^^') {
+        const rhs = parseExpression(tokens, rightBindingPower)
+        lhs = { type: 'LogicalExpression', operator: token.value, left: lhs, right: rhs }
+      } else {
+        const rhs = parseExpression(tokens, rightBindingPower)
+        lhs = { type: 'BinaryExpression', operator: token.value as BinaryOperator, left: lhs, right: rhs }
       }
+    } else if (token.value in INFIX_OPERATOR_PRECEDENCE_RIGHT) {
+      const precedence = INFIX_OPERATOR_PRECEDENCE_RIGHT[token.value]
+      const leftBindingPower = precedence
+      const rightBindingPower = precedence - 1
 
-      const expression = new CallExpression(callee, args)
+      if (leftBindingPower < minBindingPower) break
 
-      // e.g. texture().rgb
-      let i = 3 + j
-      if (body[i]?.value === '.') {
-        const right = parseExpression([first, ...body.slice(i)])! as MemberExpression
-        right.object = expression
-        return right
-      }
+      consume(tokens)
 
-      return expression
-    } else if (second.value === '.') {
-      const object = new Identifier(first.value)
-      const property = parseExpression([body[2]])!
-      const left = new MemberExpression(object, property)
-
-      // e.g. array.length()
-      if (body[3]?.value === '(' && last.value === ')') {
-        const right = parseExpression(body.slice(2))! as CallExpression
-        right.callee = left
-        return right
+      if (token.value === '?') {
+        const mhs = parseExpression(tokens, 0)
+        consume(tokens, ':')
+        const rhs = parseExpression(tokens, rightBindingPower)
+        lhs = { type: 'ConditionalExpression', test: lhs, alternate: mhs, consequent: rhs }
+      } else {
+        const rhs = parseExpression(tokens, rightBindingPower)
+        lhs = { type: 'AssignmentExpression', operator: token.value as AssignmentOperator, left: lhs, right: rhs }
       }
+    } else {
+      break
+    }
+  }
 
-      return left
-    } else if (second.value === '[') {
-      let i = 2
-
-      const type = new Type(first.value, [])
-
-      if (body[i].value !== ']') type.parameters!.push(parseExpression([body[i++]]) as any)
-      i++ // skip ]
-
-      const scope = readUntil(')', body, i).slice(1, -1)
-
-      const members: AST[] = []
-
-      let j = 0
-      while (j < scope.length) {
-        const next = readUntil(',', scope, j)
-        j += next.length
+  return lhs
+}
 
-        if (next[next.length - 1].value === ',') next.pop()
+function parseVariableDeclarator(
+  tokens: Token[],
+  typeSpecifier: Identifier | ArraySpecifier,
+  qualifiers: (ConstantQualifier | InterpolationQualifier | StorageQualifier | PrecisionQualifier)[],
+  layout: Record<string, string | boolean> | null,
+): VariableDeclarator {
+  let id: Identifier = { type: 'Identifier', name: consume(tokens).value }
+
+  if (tokens[0]?.value === '[') {
+    consume(tokens, '[')
+    id = {
+      type: 'ArraySpecifier',
+      typeSpecifier: id,
+      dimensions: [parseExpression(tokens) as Literal | Identifier],
+    } satisfies ArraySpecifier as unknown as Identifier
+    consume(tokens, ']')
+  }
 
-        members.push(parseExpression(next)!)
-      }
+  let init: Expression | null = null
 
-      return new ArrayExpression(type, members)
-    }
+  if (tokens[0]?.value === '=') {
+    consume(tokens, '=')
+    init = parseExpression(tokens)
   }
 
-  return null
+  return { type: 'VariableDeclarator', id, qualifiers, typeSpecifier, layout, init }
 }
 
 function parseVariable(
-  qualifiers: string[] = [],
+  tokens: Token[],
+  typeSpecifier: Identifier | ArraySpecifier,
+  qualifiers: (ConstantQualifier | InterpolationQualifier | StorageQualifier | PrecisionQualifier)[] = [],
   layout: Record<string, string | boolean> | null = null,
 ): VariableDeclaration {
-  i-- // TODO: remove backtrack hack
-
-  const kind = null // TODO: WGSL
-  const type = new Type(tokens[i++].value, null)
-
   const declarations: VariableDeclarator[] = []
 
-  const body = consumeUntil(';')
-  let j = 0
-
-  while (j < body.length) {
-    const name = body[j++].value
-
-    let prefix: AST | null = null
-    if (body[j].value === '[') {
-      j++ // skip [
-      prefix = new ArrayExpression(new Type(type.name, [parseExpression([body[j++]]) as any]), [])
-      j++ // skip ]
-    }
-
-    let value: AST | null = null
-
-    const delimiter = body[j++]
-    if (delimiter?.value === '=') {
-      const right = readUntil(',', body, j)
-      j += right.length
+  if (tokens[0]?.value !== ';') {
+    while (tokens.length) {
+      declarations.push(parseVariableDeclarator(tokens, typeSpecifier, qualifiers, layout))
 
-      value = parseExpression(right.slice(0, -1))
+      if (tokens[0]?.value === ',') {
+        consume(tokens, ',')
+      } else {
+        break
+      }
     }
-
-    declarations.push(new VariableDeclarator(name, value ?? prefix))
   }
 
-  return new VariableDeclaration(layout, qualifiers, kind, type, declarations)
+  consume(tokens, ';')
+
+  return { type: 'VariableDeclaration', declarations }
 }
 
-function parseFunction(qualifiers: string[]): FunctionDeclaration {
-  i-- // TODO: remove backtrack hack
+function parseUniformBlock(
+  tokens: Token[],
+  typeSpecifier: Identifier | ArraySpecifier,
+  qualifiers: LayoutQualifier[] = [],
+  layout: Record<string, string | boolean> | null = null,
+): UniformDeclarationBlock {
+  const members = parseBlock(tokens).body as VariableDeclaration[]
 
-  const type = new Type(tokens[i++].value, null)
-  const name = tokens[i++].value
-  const args: VariableDeclaration[] = []
+  let id: Identifier | null = null
+  if (tokens[0]?.value !== ';') id = parseExpression(tokens) as Identifier
+  consume(tokens, ';')
 
-  // TODO: merge with parseVariable
-  const header = consumeUntil(')').slice(1, -1)
-  let j = 0
-  while (j < header.length) {
-    const qualifiers: string[] = []
-    while (header[j] && header[j].type !== 'identifier') {
-      qualifiers.push(header[j++].value)
-    }
-    const type = new Type(qualifiers.pop()!, null)
+  return { type: 'UniformDeclarationBlock', id, qualifiers, typeSpecifier, layout, members }
+}
 
-    const line = readUntil(',', header, j)
-    j += line.length
+function parseFunction(
+  tokens: Token[],
+  typeSpecifier: ArraySpecifier | Identifier,
+  qualifiers: PrecisionQualifier[] = [],
+): FunctionDeclaration {
+  const id: Identifier = { type: 'Identifier', name: consume(tokens).value }
 
-    const name = line.shift()!.value
+  consume(tokens, '(')
 
-    let prefix: AST | null = null
-    if (line[0]?.value === '[') {
-      line.shift() // skip [
-      prefix = new ArrayExpression(new Type(type.name, [parseExpression([line.shift()!]) as any]), [])
-      line.shift() // skip ]
+  const params: FunctionParameter[] = []
+  while (tokens[0] && tokens[0].value !== ')') {
+    const qualifiers: (ConstantQualifier | ParameterQualifier | PrecisionQualifier)[] = []
+    while (tokens[0] && QUALIFIER_REGEX.test(tokens[0].value)) {
+      qualifiers.push(consume(tokens).value as ConstantQualifier | ParameterQualifier | PrecisionQualifier)
     }
+    const typeSpecifier = parseExpression(tokens) as ArraySpecifier | Identifier
+    const id: Identifier = { type: 'Identifier', name: consume(tokens).value }
 
-    if (line[line.length - 1]?.value === ',') line.pop() // skip ,
-
-    const value = parseExpression(line) ?? prefix
+    params.push({ type: 'FunctionParameter', id, qualifiers, typeSpecifier })
 
-    const declarations: VariableDeclarator[] = [new VariableDeclarator(name, value)]
-
-    args.push(new VariableDeclaration(null, qualifiers, null, type, declarations))
+    if (tokens[0]?.value === ',') consume(tokens, ',')
   }
 
+  consume(tokens, ')')
+
   let body = null
-  if (tokens[i].value === ';') i++ // skip ;
-  else body = parseBlock()
+  if (tokens[0].value === ';') consume(tokens, ';')
+  else body = parseBlock(tokens)
 
-  return new FunctionDeclaration(name, type, qualifiers, args, body)
+  return { type: 'FunctionDeclaration', id, qualifiers, typeSpecifier, params, body }
 }
 
-function parseIndeterminate(): VariableDeclaration | FunctionDeclaration {
-  i-- // TODO: remove backtrack hack
-
+function parseIndeterminate(tokens: Token[]): VariableDeclaration | FunctionDeclaration | UniformDeclarationBlock {
   let layout: Record<string, string | boolean> | null = null
-  if (tokens[i].value === 'layout') {
-    i++ // skip layout
+  if (tokens[0].value === 'layout') {
+    consume(tokens, 'layout')
+    consume(tokens, '(')
 
     layout = {}
 
-    let key: string | null = null
-    while (tokens[i] && tokens[i].value !== ')') {
-      const token = tokens[i++]
-
-      if (token.value === ',') key = null
-      if (token.type === 'symbol') continue
-
-      if (!key) {
-        key = token.value
-        layout[key] = true
+    while (tokens[0] && (tokens[0] as Token).value !== ')') {
+      const expression = parseExpression(tokens)
+
+      if (
+        expression.type === 'AssignmentExpression' &&
+        expression.left.type === 'Identifier' &&
+        expression.right.type === 'Literal'
+      ) {
+        layout[expression.left.name] = expression.right.value
+      } else if (expression.type === 'Identifier') {
+        layout[expression.name] = true
       } else {
-        layout[key] = token.value
+        throw new TypeError('Unexpected expression')
       }
+
+      if (tokens[0] && (tokens[0] as Token).value !== ')') consume(tokens, ',')
     }
 
-    i++ // skip )
+    consume(tokens, ')')
   }
 
+  // TODO: only precision qualifier valid for function return type
   const qualifiers: string[] = []
-  while (tokens[i] && QUALIFIER_REGEX.test(tokens[i].value)) {
-    qualifiers.push(tokens[i++].value)
+  while (tokens[0] && QUALIFIER_REGEX.test(tokens[0].value)) {
+    qualifiers.push(consume(tokens).value)
   }
-  i++
 
-  return tokens[i + 1]?.value === '(' ? parseFunction(qualifiers) : parseVariable(qualifiers, layout)
+  // TODO: cleanup handling of typeSpecifier
+  if (tokens[2]?.value === '(')
+    return parseFunction(
+      tokens,
+      parseExpression(tokens) as Identifier | ArraySpecifier,
+      qualifiers as PrecisionQualifier[],
+    )
+
+  const typeSpecifier = parseExpression(tokens) as Identifier | ArraySpecifier
+
+  if (tokens[0]?.value === '{') return parseUniformBlock(tokens, typeSpecifier, qualifiers as LayoutQualifier[], layout)
+  else
+    return parseVariable(
+      tokens,
+      typeSpecifier,
+      qualifiers as (ConstantQualifier | InterpolationQualifier | StorageQualifier | PrecisionQualifier)[],
+      layout,
+    )
 }
 
-function parseStruct(): StructDeclaration {
-  const name = tokens[i++].value
-  i++ // skip {
+function parseStruct(tokens: Token[]): StructDeclaration {
+  consume(tokens, 'struct')
+  const id: Identifier = { type: 'Identifier', name: consume(tokens).value }
+  consume(tokens, '{')
   const members: VariableDeclaration[] = []
-  while (tokens[i] && tokens[i].value !== '}') {
-    i++ // TODO: remove backtrack hack
-    members.push(parseIndeterminate() as VariableDeclaration)
+  while (tokens[0] && tokens[0].value !== '}') {
+    members.push(parseIndeterminate(tokens) as VariableDeclaration)
   }
-  i++ // skip }
-  i++ // skip ;
+  consume(tokens, '}')
+  consume(tokens, ';')
 
-  return new StructDeclaration(name, members)
+  return { type: 'StructDeclaration', id, members }
 }
 
-function parseReturn(): ReturnStatement {
-  const body = consumeUntil(';')
-  body.pop() // skip ;
+function parseContinue(tokens: Token[]): ContinueStatement {
+  consume(tokens, 'continue')
+  consume(tokens, ';')
 
-  const argument = parseExpression(body)
-
-  return new ReturnStatement(argument as any)
+  return { type: 'ContinueStatement' }
 }
 
-function parseIf(): IfStatement {
-  const test = parseExpression(consumeUntil(')'))!
-  const consequent = parseBlock()
+function parseBreak(tokens: Token[]): BreakStatement {
+  consume(tokens, 'break')
+  consume(tokens, ';')
 
-  let alternate = null
-  if (tokens[i].value === 'else') {
-    i++ // TODO: remove backtrack hack
+  return { type: 'BreakStatement' }
+}
 
-    if (tokens[i].value === 'if') {
-      i++
-      alternate = parseIf()
-    } else {
-      alternate = parseBlock()
-    }
-  }
+function parseDiscard(tokens: Token[]): DiscardStatement {
+  consume(tokens, 'discard')
+  consume(tokens, ';')
 
-  return new IfStatement(test, consequent, alternate)
+  return { type: 'DiscardStatement' }
 }
 
-function parseWhile(): WhileStatement {
-  const test = parseExpression(consumeUntil(')'))!
-  const body = parseBlock()
+function parseReturn(tokens: Token[]): ReturnStatement {
+  consume(tokens, 'return')
+
+  let argument: Expression | null = null
+  if (tokens[0]?.value !== ';') argument = parseExpression(tokens)
+  consume(tokens, ';')
 
-  return new WhileStatement(test, body)
+  return { type: 'ReturnStatement', argument }
 }
 
-function parseFor(): ForStatement {
-  const delimiterIndex = i + (readUntil(')', tokens, i).length - 1)
+function parseIf(tokens: Token[]): IfStatement {
+  consume(tokens, 'if')
+  consume(tokens, '(')
+  const test = parseExpression(tokens)
+  consume(tokens, ')')
 
-  i++ // skip (
-  i++ // TODO: remove backtrack hack
+  const consequent = parseBlock(tokens)
 
-  const init = parseVariable()
-  const test = parseExpression(consumeUntil(';').slice(0, -1))
-  const update = parseExpression(tokens.slice(i, delimiterIndex))
+  let alternate = null
+  if (tokens[0] && tokens[0].value === 'else') {
+    consume(tokens, 'else')
+
+    if (tokens[0] && (tokens[0] as Token).value === 'if') {
+      alternate = parseIf(tokens)
+    } else {
+      alternate = parseBlock(tokens)
+    }
+  }
+
+  return { type: 'IfStatement', test, consequent, alternate }
+}
 
-  i = delimiterIndex
-  i++ // skip )
+function parseWhile(tokens: Token[]): WhileStatement {
+  consume(tokens, 'while')
+  consume(tokens, '(')
+  const test = parseExpression(tokens)
+  consume(tokens, ')')
+  const body = parseBlock(tokens)
 
-  const body = parseBlock()
+  return { type: 'WhileStatement', test, body }
+}
 
-  return new ForStatement(init, test, update, body)
+function parseFor(tokens: Token[]): ForStatement {
+  consume(tokens, 'for')
+  consume(tokens, '(')
+  const typeSpecifier = parseExpression(tokens) as Identifier | ArraySpecifier
+  const init = parseVariable(tokens, typeSpecifier)
+  // consume(tokens, ';')
+  const test = parseExpression(tokens)
+  consume(tokens, ';')
+  const update = parseExpression(tokens)
+  consume(tokens, ')')
+  const body = parseBlock(tokens)
+
+  return { type: 'ForStatement', init, test, update, body }
 }
 
-function parseDoWhile(): DoWhileStatement {
-  const body = parseBlock()
-  i++ // skip while
-  const test = parseExpression(consumeUntil(')'))!
-  i++ // skip ;
+function parseDoWhile(tokens: Token[]): DoWhileStatement {
+  consume(tokens, 'do')
+  const body = parseBlock(tokens)
+  consume(tokens, 'while')
+  consume(tokens, '(')
+  const test = parseExpression(tokens)
+  consume(tokens, ')')
+  consume(tokens, ';')
 
-  return new DoWhileStatement(test, body)
+  return { type: 'DoWhileStatement', test, body }
 }
 
-function parseSwitch(): SwitchStatement {
-  const discriminant = parseExpression(consumeUntil(')'))
-  const delimiterIndex = i + readUntil('}', tokens, i).length - 1
+function parseSwitch(tokens: Token[]): SwitchStatement {
+  consume(tokens, 'switch')
+  const discriminant = parseExpression(tokens)
 
   const cases: SwitchCase[] = []
-  while (i < delimiterIndex) {
-    const token = tokens[i++]
+  while (tokens.length) {
+    const token = consume(tokens)
+    if (token.value === '}') break
 
     if (token.value === 'case') {
-      const test = parseExpression(consumeUntil(':').slice(0, -1))
-      const consequent = parseStatements()
-      cases.push(new SwitchCase(test, consequent))
+      const test = parseExpression(tokens)
+      consume(tokens, ':')
+      const consequent = parseStatements(tokens)
+      cases.push({ type: 'SwitchCase', test, consequent })
     } else if (token.value === 'default') {
-      i++ // skip :
-      const consequent = parseStatements()
-      cases.push(new SwitchCase(null, consequent))
+      const test = null
+      consume(tokens, ':')
+      const consequent = parseStatements(tokens)
+      cases.push({ type: 'SwitchCase', test, consequent })
     }
   }
 
-  return new SwitchStatement(discriminant!, cases)
+  return { type: 'SwitchStatement', discriminant, cases }
 }
 
-function parsePrecision(): PrecisionStatement {
-  const precision = tokens[i++].value
-  const type = new Type(tokens[i++].value, null)
-  i++ // skip ;
-  return new PrecisionStatement(precision as any, type)
+function parsePrecision(tokens: Token[]): PrecisionStatement {
+  consume(tokens, 'precision')
+  const precision = consume(tokens).value as PrecisionQualifier
+  const typeSpecifier: Identifier = { type: 'Identifier', name: consume(tokens).value }
+  consume(tokens, ';')
+  return { type: 'PrecisionStatement', precision, typeSpecifier }
 }
 
-function parsePreprocessor(): PreprocessorStatement {
-  const name = tokens[i++].value
+function parsePreprocessor(tokens: Token[]): PreprocessorStatement {
+  consume(tokens, '#')
 
-  const body = consumeUntil('\\').slice(0, -1)
-  let value: AST[] | null = null
+  let name = '' // name can be unset for the # directive which is ignored
+  let value: Expression[] | null = null
 
-  if (name !== 'else' && name !== 'endif') {
-    value = []
+  if (tokens[0]?.value !== '\\') {
+    name = consume(tokens).value
 
     if (name === 'define') {
-      const left = parseExpression([body.shift()!])!
-      const right = parseExpression(body)!
-      value.push(left, right)
+      const left = parseExpression(tokens)
+      if (tokens[0]?.value === '\\') value = [left]
+      else value = [left, parseExpression(tokens)]
     } else if (name === 'extension') {
-      const left = parseExpression([body.shift()!])!
-      body.shift() // skip :
-      const right = parseExpression(body)!
-      value.push(left, right)
+      // TODO: extension directives must be before declarations
+      const left = parseExpression(tokens)
+      consume(tokens, ':')
+      const right = parseExpression(tokens)
+      value = [left, right]
     } else if (name === 'include') {
-      value.push(parseExpression(body.slice(1, -1))!)
-    } else {
-      value.push(parseExpression(body)!)
+      consume(tokens, '<')
+      value = [{ type: 'Identifier', name: consume(tokens).value }]
+      consume(tokens, '>')
+    } else if (name !== 'else' && name !== 'endif') {
+      value = []
+      while (tokens.length && tokens[0].value !== '\\') {
+        value.push(parseExpression(tokens))
+      }
     }
   }
 
-  return new PreprocessorStatement(name, value)
+  consume(tokens, '\\')
+
+  return { type: 'PreprocessorStatement', name, value }
 }
 
-function parseStatements(): AST[] {
-  const body: AST[] = []
+function parseStatements(tokens: Token[]): Statement[] {
+  const body: Statement[] = []
   let scopeIndex = 0
 
-  while (i < tokens.length) {
-    const token = tokens[i++]
+  while (tokens.length) {
+    const token = tokens[0]
 
     scopeIndex += getScopeDelta(token)
     if (scopeIndex < 0) break
 
-    let statement: AST | null = null
-
-    if (token.value === '#') {
-      statement = parsePreprocessor()
-    } else if (token.type === 'keyword') {
-      if (token.value === 'case' || token.value === 'default') {
-        i--
-        break
-      } else if (token.value === 'struct') statement = parseStruct()
-      else if (token.value === 'continue') (statement = new ContinueStatement()), i++
-      else if (token.value === 'break') (statement = new BreakStatement()), i++
-      else if (token.value === 'discard') (statement = new DiscardStatement()), i++
-      else if (token.value === 'return') statement = parseReturn()
-      else if (token.value === 'if') statement = parseIf()
-      else if (token.value === 'while') statement = parseWhile()
-      else if (token.value === 'for') statement = parseFor()
-      else if (token.value === 'do') statement = parseDoWhile()
-      else if (token.value === 'switch') statement = parseSwitch()
-      else if (token.value === 'precision') statement = parsePrecision()
-      else if (isDeclaration(token.value) && tokens[i].value !== '[') statement = parseIndeterminate()
+    let statement: Statement | null = null
+
+    if (token.value === 'case' || token.value === 'default') break
+    else if (token.value === '#') statement = parsePreprocessor(tokens)
+    else if (token.value === 'struct') statement = parseStruct(tokens)
+    else if (token.value === 'continue') statement = parseContinue(tokens)
+    else if (token.value === 'break') statement = parseBreak(tokens)
+    else if (token.value === 'discard') statement = parseDiscard(tokens)
+    else if (token.value === 'return') statement = parseReturn(tokens)
+    else if (token.value === 'if') statement = parseIf(tokens)
+    else if (token.value === 'while') statement = parseWhile(tokens)
+    else if (token.value === 'for') statement = parseFor(tokens)
+    else if (token.value === 'do') statement = parseDoWhile(tokens)
+    else if (token.value === 'switch') statement = parseSwitch(tokens)
+    else if (token.value === 'precision') statement = parsePrecision(tokens)
+    else if (VARIABLE_REGEX.test(token.value) && tokens[1].value !== '[') statement = parseIndeterminate(tokens)
+    else {
+      const expression = parseExpression(tokens)
+      consume(tokens, ';')
+      statement = { type: 'ExpressionStatement', expression }
     }
 
-    if (statement) {
-      body.push(statement)
-    } else {
-      const line = [token, ...consumeUntil(';')]
-      if (line[line.length - 1].value === ';') line.pop()
-      const expression = parseExpression(line)
-      if (expression) body.push(expression)
-    }
+    body.push(statement)
   }
 
   return body
 }
 
-function parseBlock(): BlockStatement {
-  i++ // skip {
-  const body = parseStatements()
-  return new BlockStatement(body)
+// TODO: allow block versus sub-statements for GLSL/WGSL
+function parseBlock(tokens: Token[]): BlockStatement {
+  consume(tokens, '{')
+  const body = parseStatements(tokens)
+  consume(tokens, '}')
+  return { type: 'BlockStatement', body }
 }
 
+const NEWLINE_REGEX = /\\\n/gm
 const DIRECTIVE_REGEX = /(^\s*#[^\\]*?)(\n|\/[\/\*])/gm
 
 /**
  * Parses a string of GLSL (WGSL WIP) code into an [AST](https://en.wikipedia.org/wiki/Abstract_syntax_tree).
  */
-export function parse(code: string): AST[] {
-  // Remove (implicit) version header
-  code = code.replace('#version 300 es', '')
+export function parse(code: string): Program {
+  // Fold newlines
+  code = code.replace(NEWLINE_REGEX, '')
 
   // Escape newlines after directives, skip comments
   code = code.replace(DIRECTIVE_REGEX, '$1\\$2')
 
   // TODO: preserve
-  tokens = tokenize(code).filter((token) => token.type !== 'whitespace' && token.type !== 'comment')
-  i = 0
+  const tokens = tokenize(code).filter((token) => token.type !== 'whitespace' && token.type !== 'comment')
 
-  return parseStatements()
+  return { type: 'Program', body: parseStatements(tokens) }
 }