shrimp/src/testSetup.ts

import { expect } from 'bun:test'
import { parser } from '#parser/shrimp'
import { setGlobals } from '#parser/tokenizer'
import { globals as prelude } from '#prelude'
import { $ } from 'bun'
import { assert, errorMessage } from '#utils/utils'
import { Compiler } from '#compiler/compiler'
import { run, VM } from 'reefvm'
import { treeToString, VMResultToValue } from '#utils/tree'

const regenerateParser = async () => {
  let generate = true
  try {
    const grammarStat = await Bun.file('./src/parser/shrimp.grammar').stat()
    const tokenizerStat = await Bun.file('./src/parser/tokenizer.ts').stat()
    const parserStat = await Bun.file('./src/parser/shrimp.ts').stat()

    if (grammarStat.mtime <= parserStat.mtime && tokenizerStat.mtime <= parserStat.mtime) {
      generate = false
    }
  } catch (e) {
    console.error('Error checking or regenerating parser:', e)
  } finally {
    if (generate) {
      await $`bun generate-parser`
    }
  }
}

await regenerateParser()

// Type declaration for TypeScript
declare module 'bun:test' {
  interface Matchers<T> {
    toMatchTree(expected: string, globals?: Record<string, any>): T
    toMatchExpression(expected: string): T
    toFailParse(): T
    toEvaluateTo(expected: unknown, globals?: Record<string, any>): Promise<T>
    toFailEvaluation(): Promise<T>
  }
}

expect.extend({
  toMatchTree(received: unknown, expected: string, globals?: Record<string, any>) {
    assert(typeof received === 'string', 'toMatchTree can only be used with string values')

    const allGlobals = { ...prelude, ...(globals || {}) }
    setGlobals(Object.keys(allGlobals))
    const tree = parser.parse(received)
    const actual = treeToString(tree, received)
    const normalizedExpected = trimWhitespace(expected)

    try {
      // A hacky way to show the colorized diff in the test output
      expect(actual).toEqual(normalizedExpected)
      return { pass: true, message: () => '' }
    } catch (error) {
      return {
        message: () => (error as Error).message,
        pass: false,
      }
    }
  },

  toFailParse(received) {
    assert(typeof received === 'string', 'toFailParse can only be used with string values')

    try {
      const tree = parser.parse(received)
      let hasErrors = false
      tree.iterate({
        enter(n) {
          if (n.type.isError) {
            hasErrors = true
            return false
          }
        },
      })

      if (hasErrors) {
        return {
          message: () => `Expected input to fail parsing, and it did.`,
          pass: true,
        }
      } else {
        const actual = treeToString(tree, received)
        return {
          message: () => `Expected input to fail parsing, but it parsed successfully:\n${actual}`,
          pass: false,
        }
      }
    } catch (error) {
      return {
        message: () => `Parsing threw an error: ${(error as Error).message}`,
        pass: false,
      }
    }
  },

  async toEvaluateTo(received: unknown, expected: unknown, globals: Record<string, any> = {}) {
    assert(typeof received === 'string', 'toEvaluateTo can only be used with string values')

    try {
      const allGlobals = { ...prelude, ...(globals || {}) }
      setGlobals(Object.keys(allGlobals))
      const compiler = new Compiler(received)
      const result = await run(compiler.bytecode, allGlobals)
      let value = VMResultToValue(result)

      // Just treat regex as strings for comparison purposes
      if (expected instanceof RegExp) expected = String(expected)
      if (value instanceof RegExp) value = String(value)

      expect(value).toEqual(expected)
      return {
        message: () => `Expected evaluation to be ${expected}, but got ${value}`,
        pass: true,
      }
    } catch (error) {
      return {
        message: () => `Evaluation threw an error:\n${(error as Error).message}`,
        pass: false,
      }
    }
  },

  async toFailEvaluation(received) {
    assert(typeof received === 'string', 'toFailEvaluation can only be used with string values')

    try {
      const compiler = new Compiler(received)
      const vm = new VM(compiler.bytecode)
      const value = await vm.run()

      return {
        message: () =>
          `Expected evaluation to fail, but it succeeded with ${JSON.stringify(value)}`,
        pass: false,
      }
    } catch (error) {
      return {
        message: () => `Evaluation failed as expected: ${errorMessage(error)}`,
        pass: true,
      }
    }
  },
})

const trimWhitespace = (str: string): string => {
  const lines = str.split('\n').filter((line) => line.trim().length > 0)
  const firstLine = lines[0]
  if (!firstLine) return ''

  const leadingWhitespace = firstLine.match(/^(\s*)/)?.[1] || ''
  return lines
    .map((line) => {
      if (!line.startsWith(leadingWhitespace)) {
        let foundWhitespace = line.match(/^(\s*)/)?.[1] || ''
        throw new Error(
          `Line has inconsistent leading whitespace: "${line}" (found "${foundWhitespace}", expected "${leadingWhitespace}")`
        )
      }
      return line.slice(leadingWhitespace.length)
    })
    .join('\n')
}