Fix negative ints and add MARK ops

Author: gleb-modal

modal-js/src/pickle.test.ts

@@ -1,5 +1,6 @@
-import { expect, test } from "vitest";
+import { describe, expect, test } from "vitest";
 import { dumps, loads, type Protocol } from "./pickle";
+import { Buffer } from "buffer";
 
 test("PickleUnpickle", () => {
   const sample = {
@@ -14,3 +15,58 @@ test("PickleUnpickle", () => {
     expect(back).toEqual(sample);
   }
 });
+
+// Python pickle compatibility tests (v4)
+// Using `python -c "import pickle, base64; print(base64.b64encode(pickle.dumps(...), protocol=4).decode())"`
+const testCases = [
+  {
+    name: "MinusOne",
+    b64: "gASVBgAAAAAAAABK/////y4=",
+    expected: -1,
+  },
+  {
+    // [b'1', b'2', b'3']
+    name: "BytesList - uses MARK and APPENDS",
+    b64: "gASVEQAAAAAAAABdlChDATGUQwEylEMBM5RlLg==",
+    expected: [
+      new Uint8Array([49]),
+      new Uint8Array([50]),
+      new Uint8Array([51]),
+    ],
+  },
+  {
+    name: "SimpleList",
+    b64: "gASVCwAAAAAAAABdlChLAUsCSwNlLg==",
+    expected: [1, 2, 3],
+  },
+  {
+    name: "SimpleDict",
+    b64: "gASVEQAAAAAAAAB9lCiMAWGUSwGMAWKUSwJ1Lg==",
+    expected: { a: 1, b: 2 },
+  },
+  // Integer edge cases
+  { name: "BININT1_0", b64: "gARLAC4=", expected: 0 },
+  { name: "BININT1_255", b64: "gARL/y4=", expected: 255 },
+  { name: "BININT2_-32768", b64: "gASVBgAAAAAAAABKAID//y4=", expected: -32768 },
+  { name: "BININT2_32767", b64: "gASVBAAAAAAAAABN/38u", expected: 32767 },
+  {
+    name: "BININT4_-2147483648",
+    b64: "gASVBgAAAAAAAABKAAAAgC4=",
+    expected: -2147483648,
+  },
+  {
+    name: "BININT4_2147483647",
+    b64: "gASVBgAAAAAAAABK////fy4=",
+    expected: 2147483647,
+  },
+];
+
+describe("Python compatibility", () => {
+  for (const { name, b64, expected } of testCases) {
+    test(name, () => {
+      const buf = Buffer.from(b64, "base64");
+      const val = loads(new Uint8Array(buf));
+      expect(val).toEqual(expected);
+    });
+  }
+});

modal-js/src/pickle.ts

@@ -48,6 +48,7 @@ const enum Op {
   APPEND = 0x61, // a
   EMPTY_DICT = 0x7d, // }
   SETITEM = 0x73, // s
+  MARK = 0x28, // (  (mark stack position)
 
   // Memo / frame machinery
   BINPUT = 0x71, // q  idx(1)
@@ -56,6 +57,8 @@ const enum Op {
   LONG_BINGET = 0x6a, // j  idx(4)
   MEMOIZE = 0x94, // \x94 (≥4)
   FRAME = 0x95, // \x95 size(8) (proto‑5)
+  APPENDS = 0x65, // e
+  SETITEMS = 0x75, // u
 }
 
 // ─── Binary helpers ─────────────────────────────────────────
@@ -118,6 +121,15 @@ class Reader {
     const hi = this.uint32LE() >>> 0;
     return hi * 2 ** 32 + lo;
   }
+  int32LE() {
+    const v = new DataView(
+      this.buf.buffer,
+      this.buf.byteOffset + this.pos,
+      4,
+    ).getInt32(0, true);
+    this.pos += 4;
+    return v;
+  }
   float64BE() {
     const v = new DataView(
       this.buf.buffer,
@@ -276,6 +288,9 @@ export function loads(buf: Uint8Array): any {
     void size; // silence tsclint
   }
 
+  // Unique marker for stack operations (cannot be confused with user data)
+  const MARK = Symbol("pickle-mark");
+
   while (!r.eof()) {
     const op = r.byte();
     switch (op) {
@@ -302,8 +317,7 @@ export function loads(buf: Uint8Array): any {
         break;
       }
       case Op.BININT4: {
-        const n = r.uint32LE();
-        push(n >>> 31 ? n - 0x1_0000_0000 : n);
+        push(r.int32LE());
         break;
       }
       case Op.BINFLOAT:
@@ -386,6 +400,51 @@ export function loads(buf: Uint8Array): any {
         /* ignore */ break;
       }
 
+      case Op.MARK:
+        push(MARK);
+        break;
+
+      case Op.APPENDS: {
+        // Pops all items after the last MARK and appends them to the list below the MARK
+        // Find the last MARK
+        const markIndex = stack.lastIndexOf(MARK);
+        if (markIndex === -1) {
+          throw new PickleError("APPENDS without MARK");
+        }
+        const lst = stack[markIndex - 1];
+        if (!Array.isArray(lst)) {
+          throw new PickleError("APPENDS expects a list below MARK");
+        }
+        const items = stack.slice(markIndex + 1);
+        lst.push(...items);
+        stack.length = markIndex - 1; // Remove everything after the list
+        push(lst);
+        break;
+      }
+
+      case Op.SETITEMS: {
+        // Sets multiple key-value pairs in a dict after the last MARK
+        // Find the last MARK
+        const markIndex = stack.lastIndexOf(MARK);
+        if (markIndex === -1) {
+          throw new PickleError("SETITEMS without MARK");
+        }
+        const d = stack[markIndex - 1];
+        if (typeof d !== "object" || d === null || Array.isArray(d)) {
+          throw new PickleError("SETITEMS expects a dict below MARK");
+        }
+        const items = stack.slice(markIndex + 1);
+        // Set key-value pairs (items come in pairs: key, value, key, value, ...)
+        for (let i = 0; i < items.length; i += 2) {
+          if (i + 1 < items.length) {
+            d[items[i]] = items[i + 1];
+          }
+        }
+        stack.length = markIndex - 1; // Remove everything after the dict
+        push(d);
+        break;
+      }
+
       default:
         throw new PickleError(`Unsupported opcode 0x${op.toString(16)}`);
     }