add more Shrimp API, tests

src/index.ts

@@ -1,14 +1,17 @@
-import { glob, readFileSync } from 'fs'
+import { readFileSync } from 'fs'
-import { VM, fromValue, type Bytecode } from 'reefvm'
+import { VM, fromValue, toValue, type Bytecode } from 'reefvm'
 import { type Tree } from '@lezer/common'
 import { Compiler } from '#compiler/compiler'
 import { parser } from '#parser/shrimp'
-import { setGlobals } from '#parser/tokenizer'
+import { globals as parserGlobals, setGlobals as setParserGlobals } from '#parser/tokenizer'
-import { globals as shrimpGlobals, colors } from '#prelude'
+import { globals as shrimpGlobals } from '#prelude'
 
 export { Compiler } from '#compiler/compiler'
 export { parser } from '#parser/shrimp'
-export { globals } from '#prelude'
+export { globals as prelude } from '#prelude'
+export type { Tree } from '@lezer/common'
+export { type Value, type Bytecode } from 'reefvm'
+export { toValue, fromValue, Scope, VM, bytecodeToString } from 'reefvm'
 
 export class Shrimp {
     vm: VM
@@ -20,6 +23,38 @@ export class Shrimp {
         this.globals = globals
     }
 
+    get(name: string): any {
+        const value = this.vm.scope.get(name)
+        return value ? fromValue(value) : null
+    }
+
+    set(name: string, value: any) {
+        this.vm.scope.set(name, toValue(value, this.vm))
+    }
+
+    has(name: string): boolean {
+        return this.vm.scope.has(name)
+    }
+
+    async call(name: string, ...args: any[]): Promise<any> {
+        const result = await this.vm.call(name, ...args)
+        // vm.call() returns Value for native functions, JS values for Reef functions
+        // Check if it's a Value object and unwrap if needed
+        if (result && typeof result === 'object' && 'type' in result) {
+            return fromValue(result)
+        }
+        return result
+    }
+
+    parse(code: string): Tree {
+        return parseCode(code, this.globals)
+    }
+
+    compile(code: string): Bytecode {
+        return compileCode(code, this.globals)
+    }
+
+
     async run(code: string | Bytecode, locals?: Record<string, any>): Promise<any> {
         let bytecode
 
@@ -38,10 +73,6 @@ export class Shrimp {
         return this.vm.stack.length ? fromValue(this.vm.stack.at(-1)!) : null
     }
 
-    get(name: string): any {
-        const value = this.vm.scope.get(name)
-        return value ? fromValue(value) : null
-    }
 }
 
 export async function runFile(path: string, globals?: Record<string, any>): Promise<any> {
@@ -54,14 +85,9 @@ export async function runCode(code: string, globals?: Record<string, any>): Prom
 }
 
 export async function runBytecode(bytecode: Bytecode, globals?: Record<string, any>): Promise<any> {
-    try {
     const vm = new VM(bytecode, Object.assign({}, shrimpGlobals, globals))
     await vm.run()
     return vm.stack.length ? fromValue(vm.stack[vm.stack.length - 1]!) : null
-    } catch (error: any) {
-        console.error(`${colors.red}Error:${colors.reset} ${error.message}`)
-        process.exit(1)
-    }
 }
 
 export function compileFile(path: string, globals?: Record<string, any>): Bytecode {
@@ -81,7 +107,12 @@ export function parseFile(path: string, globals?: Record<string, any>): Tree {
 }
 
 export function parseCode(code: string, globals?: Record<string, any>): Tree {
+    const oldGlobals = [...parserGlobals]
     const globalNames = [...Object.keys(shrimpGlobals), ...(globals ? Object.keys(globals) : [])]
-    setGlobals(globalNames)
+
-    return parser.parse(code)
+    setParserGlobals(globalNames)
+    const result = parser.parse(code)
+    setParserGlobals(oldGlobals)
+
+    return result
 }

src/tests/shrimp.test.ts

@@ -1,6 +1,6 @@
 import { describe } from 'bun:test'
 import { expect, test } from 'bun:test'
-import { Shrimp } from '..'
+import { Shrimp, runCode, compileCode, parseCode, bytecodeToString } from '..'
 
 describe('Shrimp', () => {
   test('allows running Shrimp code', async () => {
@@ -50,4 +50,403 @@ describe('Shrimp', () => {
     await shrimp.run('abc = nothing')
     expect(shrimp.get('abc')).toEqual('nothing')
   })
+
+  describe('set()', () => {
+    test('allows setting variables', async () => {
+      const shrimp = new Shrimp()
+
+      shrimp.set('foo', 42)
+      expect(shrimp.get('foo')).toEqual(42)
+
+      shrimp.set('bar', 'hello')
+      expect(shrimp.get('bar')).toEqual('hello')
+    })
+
+    test('set variables are accessible in code', async () => {
+      const shrimp = new Shrimp()
+
+      shrimp.set('x', 10)
+      shrimp.set('y', 20)
+
+      const result = await shrimp.run('x + y')
+      expect(result).toEqual(30)
+    })
+
+    test('allows setting functions', async () => {
+      const shrimp = new Shrimp()
+
+      shrimp.set('double', (n: number) => n * 2)
+
+      const result = await shrimp.run('double 21')
+      expect(result).toEqual(42)
+    })
+
+    test('overwrites existing variables', async () => {
+      const shrimp = new Shrimp()
+
+      await shrimp.run('x = 100')
+      expect(shrimp.get('x')).toEqual(100)
+
+      shrimp.set('x', 200)
+      expect(shrimp.get('x')).toEqual(200)
+    })
+  })
+
+  describe('has()', () => {
+    test('returns true for existing variables', async () => {
+      const shrimp = new Shrimp()
+
+      await shrimp.run('x = 5')
+      expect(shrimp.has('x')).toEqual(true)
+    })
+
+    test('returns false for non-existing variables', () => {
+      const shrimp = new Shrimp()
+
+      expect(shrimp.has('nonexistent')).toEqual(false)
+    })
+
+    test('returns true for globals', () => {
+      const shrimp = new Shrimp({ myGlobal: 42 })
+
+      expect(shrimp.has('myGlobal')).toEqual(true)
+    })
+
+    test('returns true for prelude functions', () => {
+      const shrimp = new Shrimp()
+
+      expect(shrimp.has('echo')).toEqual(true)
+      expect(shrimp.has('type')).toEqual(true)
+    })
+  })
+
+  describe('call()', () => {
+    test('calls Shrimp functions with positional args', async () => {
+      const shrimp = new Shrimp()
+
+      await shrimp.run(`add = do x y:
+        x + y
+      end`)
+
+      const result = await shrimp.call('add', 5, 10)
+      expect(result).toEqual(15)
+    })
+
+    test('calls Shrimp functions with named args', async () => {
+      const shrimp = new Shrimp()
+
+      await shrimp.run(`greet = do name:
+        concat 'Hello ' name
+      end`)
+
+      const result = await shrimp.call('greet', { name: 'World' })
+      expect(result).toEqual('Hello World')
+    })
+
+    test('calls native functions', async () => {
+      const shrimp = new Shrimp()
+
+      shrimp.set('multiply', (a: number, b: number) => a * b)
+
+      const result = await shrimp.call('multiply', 6, 7)
+      expect(result).toEqual(42)
+    })
+
+    test('calls prelude functions', async () => {
+      const shrimp = new Shrimp()
+
+      const result = await shrimp.call('type', 42)
+      expect(result).toEqual('number')
+    })
+
+    test('calls async functions', async () => {
+      const shrimp = new Shrimp()
+
+      shrimp.set('fetchData', async () => {
+        return await Promise.resolve('async data')
+      })
+
+      const result = await shrimp.call('fetchData')
+      expect(result).toEqual('async data')
+    })
+  })
+
+  describe('compile()', () => {
+    test('compiles code to bytecode', () => {
+      const shrimp = new Shrimp()
+
+      const bytecode = shrimp.compile('x = 5')
+
+      expect(bytecode).toHaveProperty('instructions')
+      expect(bytecode).toHaveProperty('constants')
+      expect(bytecode).toHaveProperty('labels')
+      expect(bytecode.instructions.length).toBeGreaterThan(0)
+    })
+
+    test('respects globals when compiling', () => {
+      const shrimp = new Shrimp({ customGlobal: 42 })
+
+      const bytecode = shrimp.compile('x = customGlobal')
+      expect(bytecode.instructions.length).toBeGreaterThan(0)
+    })
+
+    test('compiled bytecode can be run', async () => {
+      const shrimp = new Shrimp()
+
+      const bytecode = shrimp.compile('2 * 21')
+      const result = await shrimp.run(bytecode)
+
+      expect(result).toEqual(42)
+    })
+  })
+
+  describe('parse()', () => {
+    test('parses code to syntax tree', () => {
+      const shrimp = new Shrimp()
+
+      const tree = shrimp.parse('x = 5')
+
+      expect(tree).toHaveProperty('length')
+      expect(tree).toHaveProperty('cursor')
+      expect(tree.length).toBeGreaterThan(0)
+    })
+
+    test('respects globals when parsing', () => {
+      const shrimp = new Shrimp({ myVar: 42 })
+
+      const tree = shrimp.parse('x = myVar + 10')
+
+      // Should parse without errors
+      expect(tree).toHaveProperty('length')
+      expect(tree.length).toBeGreaterThan(0)
+    })
+
+    test('parses function definitions', () => {
+      const shrimp = new Shrimp()
+
+      const tree = shrimp.parse(`add = do x y:
+        x + y
+      end`)
+
+      expect(tree.length).toBeGreaterThan(0)
+    })
+  })
+
+  describe('get()', () => {
+    test('returns null for undefined variables', () => {
+      const shrimp = new Shrimp()
+
+      expect(shrimp.get('undefined')).toEqual(null)
+    })
+
+    test('returns values from code execution', async () => {
+      const shrimp = new Shrimp()
+
+      await shrimp.run('x = 42')
+      expect(shrimp.get('x')).toEqual(42)
+    })
+
+    test('returns arrays', async () => {
+      const shrimp = new Shrimp()
+
+      await shrimp.run('arr = [1 2 3]')
+      expect(shrimp.get('arr')).toEqual([1, 2, 3])
+    })
+
+    test('returns dicts', async () => {
+      const shrimp = new Shrimp()
+
+      await shrimp.run('dict = [a=1 b=2]')
+      expect(shrimp.get('dict')).toEqual({ a: 1, b: 2 })
+    })
+  })
+
+  describe('running bytecode directly', () => {
+    test('can run pre-compiled bytecode', async () => {
+      const shrimp = new Shrimp()
+
+      const bytecode = shrimp.compile('x = 100')
+      const result = await shrimp.run(bytecode)
+
+      expect(result).toEqual(100)
+      expect(shrimp.get('x')).toEqual(100)
+    })
+
+    test('maintains state across bytecode runs', async () => {
+      const shrimp = new Shrimp()
+
+      const bytecode1 = shrimp.compile('x = 10')
+      const bytecode2 = shrimp.compile('x + 5')
+
+      await shrimp.run(bytecode1)
+      const result = await shrimp.run(bytecode2)
+
+      expect(result).toEqual(15)
+    })
+  })
+})
+
+describe('Functional API', () => {
+  describe('runCode()', () => {
+    test('runs code and returns result', async () => {
+      const result = await runCode('1 + 1')
+      expect(result).toEqual(2)
+    })
+
+    test('works with globals', async () => {
+      const result = await runCode('greet', { greet: () => 'hello' })
+      expect(result).toEqual('hello')
+    })
+
+    test('has access to prelude', async () => {
+      const result = await runCode('type 42')
+      expect(result).toEqual('number')
+    })
+
+    test('returns null for empty code', async () => {
+      const result = await runCode('')
+      expect(result).toEqual(null)
+    })
+  })
+
+  describe('compileCode()', () => {
+    test('compiles code to bytecode', () => {
+      const bytecode = compileCode('x = 5')
+
+      expect(bytecode).toHaveProperty('instructions')
+      expect(bytecode).toHaveProperty('constants')
+      expect(bytecode.instructions.length).toBeGreaterThan(0)
+    })
+
+    test('respects globals', () => {
+      const bytecode = compileCode('x = myGlobal', { myGlobal: 42 })
+
+      expect(bytecode.instructions.length).toBeGreaterThan(0)
+    })
+
+    test('compiled bytecode is usable', async () => {
+      const bytecode = compileCode('21 * 2')
+      const result = await runCode('21 * 2')
+
+      expect(result).toEqual(42)
+    })
+  })
+
+  describe('parseCode()', () => {
+    test('parses code to syntax tree', () => {
+      const tree = parseCode('x = 5')
+
+      expect(tree).toHaveProperty('length')
+      expect(tree.length).toBeGreaterThan(0)
+    })
+
+    test('respects globals', () => {
+      const tree = parseCode('x = myGlobal', { myGlobal: 42 })
+
+      expect(tree.length).toBeGreaterThan(0)
+    })
+
+    test('handles complex expressions', () => {
+      const tree = parseCode(`add = do x y:
+        x + y
+      end
+      result = add 5 10`)
+
+      expect(tree.length).toBeGreaterThan(0)
+    })
+  })
+
+  describe('bytecodeToString()', () => {
+    test('converts bytecode to human-readable format', () => {
+      const bytecode = compileCode('x = 42')
+      const str = bytecodeToString(bytecode)
+
+      expect(typeof str).toEqual('string')
+      expect(str.length).toBeGreaterThan(0)
+    })
+
+    test('shows instructions', () => {
+      const bytecode = compileCode('1 + 1')
+      const str = bytecodeToString(bytecode)
+
+      // Should contain some opcodes
+      expect(str).toContain('PUSH')
+    })
+  })
+})
+
+describe('Integration tests', () => {
+  test('complex REPL-like workflow', async () => {
+    const shrimp = new Shrimp()
+
+    // Define a function
+    await shrimp.run(`double = do x:
+      x * 2
+    end`)
+    expect(shrimp.has('double')).toEqual(true)
+
+    // Use the function
+    const result1 = await shrimp.run('double 21')
+    expect(result1).toEqual(42)
+
+    // Call it from TypeScript
+    const result2 = await shrimp.call('double', 50)
+    expect(result2).toEqual(100)
+
+    // Define another function using the first
+    await shrimp.run(`quadruple = do x:
+      double (double x)
+    end`)
+
+    const result3 = await shrimp.run('quadruple 5')
+    expect(result3).toEqual(20)
+  })
+
+  test('mixing native and Shrimp functions', async () => {
+    const shrimp = new Shrimp({
+      log: (msg: string) => `Logged: ${msg}`,
+      multiply: (a: number, b: number) => a * b,
+    })
+
+    await shrimp.run(`greet = do name:
+      log name
+    end`)
+
+    const result1 = await shrimp.run('greet Alice')
+    expect(result1).toEqual('Logged: Alice')
+
+    await shrimp.run(`calc = do x:
+      multiply x 3
+    end`)
+
+    const result2 = await shrimp.run('calc 7')
+    expect(result2).toEqual(21)
+  })
+
+  test('working with arrays and dicts', async () => {
+    const shrimp = new Shrimp()
+
+    await shrimp.run('nums = [1 2 3 4 5]')
+    expect(shrimp.get('nums')).toEqual([1, 2, 3, 4, 5])
+
+    await shrimp.run("config = [host='localhost' port=3000]")
+    expect(shrimp.get('config')).toEqual({ host: 'localhost', port: 3000 })
+
+    const result = await shrimp.run('length nums')
+    expect(result).toEqual(5)
+  })
+
+  test('compile once, run multiple times', async () => {
+    const bytecode = compileCode('x * 2')
+
+    const shrimp1 = new Shrimp()
+    shrimp1.set('x', 10)
+    const result1 = await shrimp1.run(bytecode)
+    expect(result1).toEqual(20)
+
+    const shrimp2 = new Shrimp()
+    shrimp2.set('x', 100)
+    const result2 = await shrimp2.run(bytecode)
+    expect(result2).toEqual(200)
+  })
 })