feat(hogvm): bytecode in state (#22740)

PostHog · Jun 6, 2024 · 9cbf986 · 9cbf986
1 parent 85526c0
commit 9cbf986
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 9 deletions.
diff --git a/hogvm/typescript/src/__tests__/execute.test.ts b/hogvm/typescript/src/__tests__/execute.test.ts
@@ -377,6 +377,7 @@ describe('HogQL Bytecode', () => {
             finished: false,
             result: undefined,
             state: {
+                bytecode,
                 asyncSteps: 1,
                 callStack: [],
                 declaredFunctions: {},

diff --git a/hogvm/typescript/src/execute.ts b/hogvm/typescript/src/execute.ts
@@ -52,6 +52,8 @@ function setNestedValue(obj: any, chain: any[], value: any): void {
 }
 
 interface VMState {
+    /** Bytecode running in the VM */
+    bytecode: any[]
     /** Stack of the VM */
     stack: any[]
     /** Call stack of the VM */
@@ -93,24 +95,24 @@ export function execSync(bytecode: any[], options?: ExecOptions): any {
 }
 
 export async function execAsync(bytecode: any[], options?: ExecOptions): Promise<any> {
-    let lastState: VMState | undefined = undefined
+    let vmState: VMState | undefined = undefined
     while (true) {
-        const response = exec(bytecode, options, lastState)
+        const response = exec(vmState ?? bytecode, options)
         if (response.finished) {
             return response.result
         }
         if (response.state && response.asyncFunctionName && response.asyncFunctionArgs) {
-            lastState = response.state
+            vmState = response.state
             if (options?.asyncFunctions && response.asyncFunctionName in options.asyncFunctions) {
                 const result = await options?.asyncFunctions[response.asyncFunctionName](...response.asyncFunctionArgs)
-                lastState.stack.push(result)
+                vmState.stack.push(result)
             } else if (response.asyncFunctionName in ASYNC_STL) {
                 const result = await ASYNC_STL[response.asyncFunctionName](
                     response.asyncFunctionArgs,
                     response.asyncFunctionName,
                     options?.timeout ?? DEFAULT_TIMEOUT
                 )
-                lastState.stack.push(result)
+                vmState.stack.push(result)
             } else {
                 throw new Error('Invalid async function call: ' + response.asyncFunctionName)
             }
@@ -120,8 +122,17 @@ export async function execAsync(bytecode: any[], options?: ExecOptions): Promise
     }
 }
 
-export function exec(bytecode: any[], options?: ExecOptions, vmState?: VMState): ExecResult {
-    if (bytecode.length === 0 || bytecode[0] !== '_h') {
+export function exec(code: any[] | VMState, options?: ExecOptions): ExecResult {
+    let vmState: VMState | undefined = undefined
+    let bytecode: any[] | undefined = undefined
+    if (!Array.isArray(code)) {
+        vmState = code
+        bytecode = vmState.bytecode
+    } else {
+        bytecode = code
+    }
+
+    if (!bytecode || bytecode.length === 0 || bytecode[0] !== '_h') {
         throw new Error("Invalid HogQL bytecode, must start with '_h'")
     }
 
@@ -149,10 +160,10 @@ export function exec(bytecode: any[], options?: ExecOptions, vmState?: VMState):
     }
 
     function next(): any {
-        if (ip >= bytecode.length - 1) {
+        if (ip >= bytecode!.length - 1) {
             throw new Error('Unexpected end of bytecode')
         }
-        return bytecode[++ip]
+        return bytecode![++ip]
     }
     function checkTimeout(): void {
         if (syncDuration + Date.now() - startTime > timeout * 1000) {
@@ -383,6 +394,7 @@ export function exec(bytecode: any[], options?: ExecOptions, vmState?: VMState):
                             asyncFunctionName: name,
                             asyncFunctionArgs: args,
                             state: {
+                                bytecode,
                                 stack,
                                 callStack,
                                 declaredFunctions,