refactor: preprocessor statements

src/ast.ts

@@ -18,6 +18,12 @@ export class BlockStatement extends Node {
   }
 }
 
+export class PreprocessorStatement extends Node {
+  constructor(public name: string, public value: Node[] | null) {
+    super()
+  }
+}
+
 export class Type extends Node {
   constructor(public name: string, public parameters: (Type | Literal | Identifier)[] | null) {
     super()
@@ -173,3 +179,4 @@ export type AST =
   | ContinueStatement
   | BreakStatement
   | DiscardStatement
+  | PreprocessorStatement

src/generator.ts

@@ -24,6 +24,7 @@ import {
   ContinueStatement,
   BreakStatement,
   DiscardStatement,
+  PreprocessorStatement,
 } from './ast'
 
 const EOL_REGEX = /\n$/
@@ -133,6 +134,19 @@ function format(node: AST | null): string {
     line = 'break;\n'
   } else if (node instanceof DiscardStatement) {
     line = 'discard;\n'
+  } else if (node instanceof PreprocessorStatement) {
+    let value = ''
+    if (node.value) {
+      if (node.name === 'include') {
+        value = ` <${format(node.value[0])}>`
+      } else if (node.name === 'extension') {
+        value = ` ${node.value.map(format).join(': ')}`
+      } else {
+        value = ` ${node.value.map(format).join(' ')}`
+      }
+    }
+
+    line = `#${node.name}${value}\n`
   }
 
   return line

src/parser.ts

@@ -25,6 +25,7 @@ import {
   StructDeclaration,
   PrecisionStatement,
   ArrayExpression,
+  PreprocessorStatement,
 } from './ast'
 import { generate } from './generator'
 import { type Token, tokenize } from './tokenizer'
@@ -455,6 +456,34 @@ function parsePrecision(): PrecisionStatement {
   return new PrecisionStatement(precision as any, type)
 }
 
+function parsePreprocessor(): PreprocessorStatement {
+  const name = tokens[i++].value
+
+  const body = consumeUntil('\\').slice(0, -1)
+  let value: AST[] | null = null
+
+  if (name !== 'else' && name !== 'endif') {
+    value = []
+
+    if (name === 'define') {
+      const left = parseExpression([body.shift()!])!
+      const right = parseExpression(body)!
+      value.push(left, right)
+    } else if (name === 'extension') {
+      const left = parseExpression([body.shift()!])!
+      body.shift() // skip :
+      const right = parseExpression(body)!
+      value.push(left, right)
+    } else if (name === 'include') {
+      value.push(parseExpression(body.slice(1, -1))!)
+    } else {
+      value.push(parseExpression(body)!)
+    }
+  }
+
+  return new PreprocessorStatement(name, value)
+}
+
 function parseStatements(): AST[] {
   const body: AST[] = []
   let scopeIndex = 0
@@ -467,7 +496,9 @@ function parseStatements(): AST[] {
 
     let statement: AST | null = null
 
-    if (token.type === 'keyword') {
+    if (token.value === '#') {
+      statement = parsePreprocessor()
+    } else if (token.type === 'keyword') {
       if (token.value === 'case' || token.value === 'default') {
         i--
         break
@@ -504,12 +535,19 @@ function parseBlock(): BlockStatement {
   return new BlockStatement(body)
 }
 
+const DIRECTIVE_REGEX = /(^\s*#[^\\]*?)(\n|\/[\/\*])/gm
+
 /**
  * Parses a string of GLSL or WGSL code into an [AST](https://en.wikipedia.org/wiki/Abstract_syntax_tree).
  */
 export function parse(code: string): AST[] {
-  // TODO: preserve
+  // Remove (implicit) version header
   code = code.replace('#version 300 es', '')
+
+  // 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
 
@@ -553,6 +591,17 @@ void main() {
   gl_FragColor = vec4(1, 0, 0, 1);
 }
 
+#line 0
+#define TEST 1
+#extension all: disable
+#error test
+#pragma optimize(on)
+#include <chunk>
+
+#ifdef TEST
+test;
+#endif
+
 method(true);
 
 foo.bar();

tests/parser.test.ts

@@ -25,6 +25,7 @@ import {
   SwitchCase,
   PrecisionStatement,
   ArrayExpression,
+  PreprocessorStatement,
 } from 'shaderkit'
 
 describe('parser', () => {
@@ -371,4 +372,63 @@ describe('parser', () => {
     expect(statement.members[2]).toBeInstanceOf(Literal)
     expect((statement.members[2] as Literal).value).toBe('1.5')
   })
+
+  it('parses preprocessor statements', () => {
+    {
+      const statement = parse('#define TEST 1\n')[0] as PreprocessorStatement
+      expect(statement).toBeInstanceOf(PreprocessorStatement)
+      expect(statement.name).toBe('define')
+      console.log(statement.value)
+      expect(statement.value!.length).toBe(2)
+      expect(statement.value![0]).toBeInstanceOf(Identifier)
+      expect((statement.value![0] as Identifier).value).toBe('TEST')
+      expect(statement.value![1]).toBeInstanceOf(Literal)
+      expect((statement.value![1] as Literal).value).toBe('1')
+    }
+
+    {
+      const statement = parse('#extension all: disable\n')[0] as PreprocessorStatement
+      expect(statement).toBeInstanceOf(PreprocessorStatement)
+      expect(statement.name).toBe('extension')
+      expect(statement.value!.length).toBe(2)
+    }
+
+    {
+      const statement = parse('#pragma optimize(on)\n')[0] as PreprocessorStatement
+      expect(statement).toBeInstanceOf(PreprocessorStatement)
+      expect(statement.name).toBe('pragma')
+      expect(statement.value!.length).toBe(1)
+      expect(statement.value![0]).toBeInstanceOf(CallExpression)
+      expect((statement.value![0] as CallExpression).callee).toBeInstanceOf(Identifier)
+      expect(((statement.value![0] as CallExpression).callee as Identifier).value).toBe('optimize')
+      expect((statement.value![0] as CallExpression).args.length).toBe(1)
+      expect((statement.value![0] as CallExpression).args[0]).toBeInstanceOf(Identifier)
+      expect(((statement.value![0] as CallExpression).args[0] as Identifier).value).toBe('on')
+    }
+
+    {
+      const statement = parse('#include <chunk>\n')[0] as PreprocessorStatement
+      expect(statement).toBeInstanceOf(PreprocessorStatement)
+      expect(statement.name).toBe('include')
+      expect(statement.value!.length).toBe(1)
+      expect(statement.value![0]).toBeInstanceOf(Identifier)
+      expect((statement.value![0] as Identifier).value).toBe('chunk')
+    }
+
+    {
+      const statement = parse('#ifdef TEST\n')[0] as PreprocessorStatement
+      expect(statement).toBeInstanceOf(PreprocessorStatement)
+      expect(statement.name).toBe('ifdef')
+      expect(statement.value!.length).toBe(1)
+      expect(statement.value![0]).toBeInstanceOf(Identifier)
+      expect((statement.value![0] as Identifier).value).toBe('TEST')
+    }
+
+    {
+      const statement = parse('#endif\n')[0] as PreprocessorStatement
+      expect(statement).toBeInstanceOf(PreprocessorStatement)
+      expect(statement.name).toBe('endif')
+      expect(statement.value).toBe(null)
+    }
+  })
 })