From a535dc9605f86483200da586ad77407ae9cc9a90 Mon Sep 17 00:00:00 2001
From: Chris Wanstrath <chris@ozmm.org>
Date: Wed, 5 Nov 2025 13:41:08 -0800
Subject: [PATCH] Shrimp accepts custom globals

---
 src/index.ts             |  9 ++++++---
 src/tests/shrimp.test.ts | 14 +++++++++++++-
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index dad8935..963ba25 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -9,16 +9,19 @@ export { globals } from '#prelude'
 
 export class Shrimp {
     vm: VM
+    private globals?: Record<string, any>
 
-    constructor() {
-        this.vm = new VM({ instructions: [], constants: [], labels: new Map() }, shrimpGlobals)
+    constructor(globals?: Record<string, any>) {
+        const emptyBytecode = { instructions: [], constants: [], labels: new Map() }
+        this.vm = new VM(emptyBytecode, Object.assign({}, shrimpGlobals, globals ?? {}))
+        this.globals = globals
     }
 
     async run(code: string | Bytecode): Promise<any> {
         let bytecode
 
         if (typeof code === 'string') {
-            const compiler = new Compiler(code, Object.keys(shrimpGlobals))
+            const compiler = new Compiler(code, Object.keys(Object.assign({}, shrimpGlobals, this.globals ?? {})))
             bytecode = compiler.bytecode
         } else {
             bytecode = code
diff --git a/src/tests/shrimp.test.ts b/src/tests/shrimp.test.ts
index b71c458..2713f3e 100644
--- a/src/tests/shrimp.test.ts
+++ b/src/tests/shrimp.test.ts
@@ -5,8 +5,12 @@ import { Shrimp } from '..'
 describe('Shrimp', () => {
   test('allows running Shrimp code', async () => {
     const shrimp = new Shrimp()
-
     expect(await shrimp.run(`1 + 5`)).toEqual(6)
+    expect(await shrimp.run(`type 5`)).toEqual('number')
+  })
+
+  test('maintains state across runs', async () => {
+    const shrimp = new Shrimp()
 
     await shrimp.run(`abc = true`)
     expect(shrimp.get('abc')).toEqual(true)
@@ -18,4 +22,12 @@ describe('Shrimp', () => {
     await shrimp.run(`abc = false`)
     expect(shrimp.get('abc')).toEqual(false)
   })
+
+  test('allows setting your own globals', async () => {
+    const shrimp = new Shrimp({ hiya: () => 'hey there' })
+
+    await shrimp.run('abc = hiya')
+    expect(shrimp.get('abc')).toEqual('hey there')
+    expect(await shrimp.run('type abc')).toEqual('string')
+  })
 })
\ No newline at end of file