Trim the per-chunk work of compiled object-mode stream callbacks (#2523)

packages/compiler/src/backend/llvm/expr-stream-callbacks.ts

@@ -279,8 +279,6 @@ export function streamDoneFnFor(
   host.declare(
     `declare void @${entry}(ptr, ptr${twoSlot ? ", ptr, ptr" : valueData ? ", ptr" : ""})`,
   );
-  host.declare(`declare ptr @scr_box_get_ref(ptr)`);
-  host.declare(`declare void @scr_stream_release_v(ptr)`);
   host.declare(`declare ptr @scr_error_retain_v(ptr)`);
   host.declare(`declare void @scr_union_release(ptr)`);
   const params = [
@@ -295,7 +293,10 @@ export function streamDoneFnFor(
     ...(twoSlot ? [`  %bslot = alloca ptr`, `  %sslot = alloca ptr`] : []),
     `  %capsp = getelementptr inbounds %ScrClosure, ptr %self, i64 1`,
     `  %bx = load ptr, ptr %capsp`,
-    `  %s = call ptr @scr_box_get_ref(ptr %bx) ; the stream, +1`,
+    // the stream, borrowed: the box holds it, and the closure (alive for
+    // the call) holds the box
+    `  %sp = getelementptr inbounds %ScrBox, ptr %bx, i64 0, i32 5`,
+    `  %s = load ptr, ptr %sp`,
     `  store ptr null, ptr %eslot`,
     ...(twoSlot ? [`  store ptr null, ptr %bslot`, `  store ptr null, ptr %sslot`] : []),
   ];
@@ -347,7 +348,6 @@ export function streamDoneFnFor(
   } else {
     d.push(`  call void @${entry}(ptr %s, ptr %ev) ; moves err; borrows s`);
   }
-  d.push(`  call void @scr_stream_release_v(ptr %s)`);
   if (errT !== undefined) d.push(`  call void ${releaseSym(host.shapeHost, errT)}(ptr %e)`);
   if (dataT !== undefined && isRefCounted(dataT))
     d.push(`  call void ${releaseSym(host.shapeHost, dataT)}(ptr %d)`);
@@ -868,18 +868,46 @@ export function streamCbThunkFor(
       host.declare(`declare void @scr_dyn_release(ptr)`);
       host.declare(`declare zeroext i1 @scr_exc_pending()`);
       if (hasThis) host.declare(`declare void @scr_stream_release_v(ptr)`);
+      // An object-mode write without an encoding argument passes
+      // undefined: a pointer union whose unit arm is undefined holds it as
+      // the null pointer, so that case skips the conversion.
+      const pointer = p.kind === "union" ? host.nullableUnions.get(p.unionId) : null;
+      const undefinedIsNull =
+        pointer !== null &&
+        p.kind === "union" &&
+        host.unionsById.get(p.unionId)?.arms[pointer.unitTag]?.kind === "undefinedT";
+      if (undefinedIsNull) {
+        host.declare(`declare zeroext i1 @scr_stream_chunk_encoding_undefined(ptr)`);
+        d.push(
+          `  br label %etchk${i}`,
+          `etchk${i}:`,
+          `  %eu${i} = call zeroext i1 @scr_stream_chunk_encoding_undefined(ptr %s)`,
+          `  br i1 %eu${i}, label %etok${i}, label %etconv${i}`,
+          `etconv${i}:`,
+        );
+      }
       d.push(
         `  %ev${i} = call ptr @scr_stream_chunk_encoding_value(ptr %s)`,
         `  %et${i} = call ${ty === "i1" ? "zeroext i1" : ty} @${helper}(ptr %ev${i}, ptr null)`,
         `  call void @scr_dyn_release(ptr %ev${i})`,
         `  %etbad${i} = call zeroext i1 @scr_exc_pending()`,
-        `  br i1 %etbad${i}, label %etfail${i}, label %etok${i}`,
+        `  br i1 %etbad${i}, label %etfail${i}, label %${undefinedIsNull ? "etconvok" : "etok"}${i}`,
         `etfail${i}:`,
         ...(hasThis ? [`  call void @scr_stream_release_v(ptr %sr)`] : []),
         `  ret void`,
-        `etok${i}:`,
       );
-      passed.push(`${ty} %et${i}`);
+      if (undefinedIsNull) {
+        d.push(
+          `etconvok${i}:`,
+          `  br label %etok${i}`,
+          `etok${i}:`,
+          `  %etv${i} = phi ptr [ null, %etchk${i} ], [ %et${i}, %etconvok${i} ]`,
+        );
+        passed.push(`${ty} %etv${i}`);
+      } else {
+        d.push(`etok${i}:`);
+        passed.push(`${ty} %et${i}`);
+      }
     } else if (isErrPos) {
       // destroy's error argument: `Error | null` — wrap type-directedly.
       // The finished/pipeline callback ("e") may declare `Error | null |

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8520-stream-object-encoding-arg.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/8520-stream-object-encoding-arg.ts"
+ ],
+ "diags": []
+}

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8531-stream-saved-callback-receiver.mjs.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/8531-stream-saved-callback-receiver.mjs"
+ ],
+ "diags": []
+}

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8532-stream-saved-callback-receiver/main.ts.json

@@ -0,0 +1,7 @@
+{
+ "order": [
+  "<repo>/tests/corpus/8532-stream-saved-callback-receiver/label.js",
+  "<repo>/tests/corpus/8532-stream-saved-callback-receiver/main.ts"
+ ],
+ "diags": []
+}

packages/runtime/src/scr_stream.c

@@ -752,7 +752,12 @@ static void scr_stream_do_destroy(ScrStream *s, ScrError *err /*borrowed*/);
 static void scr_stream_notify_finished(ScrStream *s);
 static void *scr_stream_read_n(ScrStream *s, double size);
 static void *scr_stream_read_n_kind(ScrStream *s, double size, ScrEntryKind *kind);
-static void scr_stream_settle_next(ScrStream *s);
+static void scr_stream_settle_next_waiter(ScrStream *s);
+/* A for-await iterator's parked next(): only a stream with a waiter that
+ * is not mid-read has anything to settle. */
+static inline void scr_stream_settle_next(ScrStream *s) {
+  if (s->st->r.next_waiter && !s->st->r.next_reading) scr_stream_settle_next_waiter(s);
+}
 static ScrStream *scr_stream_alloc(const ScrVt *vt, const char *cls, bool has_r, bool has_w,
                                     double rhwm, double whwm, bool auto_destroy,
                                     bool emit_close, bool allow_half_open);
@@ -989,7 +994,12 @@ static void scr_stream_call_read(ScrStream *s) {
  * multi-byte tail held pending) adds nothing but still clears the
  * reading flag. */
 static int scr_stream_add_entry(ScrStream *s, void *entry /*moves*/, ScrEntryKind kind);
-static void *scr_stream_entry_convert(ScrStreamState *st, void *entry, ScrEntryKind *kind);
+static void *scr_stream_entry_convert_kind(ScrStreamState *st, void *entry, ScrEntryKind *kind);
+/* An entry already of the buffer's kind is kept as-is. */
+static inline void *scr_stream_entry_convert(ScrStreamState *st, void *entry, ScrEntryKind *kind) {
+  if (*kind == scr_stream_rkind(st)) return entry;
+  return scr_stream_entry_convert_kind(st, entry, kind);
+}
 
 static bool scr_stream_add_chunk(ScrStream *s, ScrBytes *chunk, bool front) {
   ScrStreamState *st = s->st;
@@ -1395,8 +1405,6 @@ static void scr_stream_next_fulfill(ScrStreamState *st, ScrPromise *w /*moves*/,
   scr_queue_microtask(job);
 }
 
-static void scr_stream_settle_next(ScrStream *s);
-
 /* One pass of createAsyncIterator's loop for the parked waiter:
  * `stream.read()` (a _read below the hwm, then the buffered chunk; pushes
  * the _read makes join it), else the settled states (error, EOF), else
@@ -1442,7 +1450,7 @@ static void scr_stream_next_wake_job(ScrClosure *job) {
  * error rejects, buffered content fulfills (one whole entry in from
  * mode, the whole buffer otherwise — Node's iterator read()), EOF
  * fulfills the sentinel. Called from push/eof/destroy transitions. */
-static void scr_stream_settle_next(ScrStream *s) {
+static void scr_stream_settle_next_waiter(ScrStream *s) {
   ScrStreamState *st = s->st;
   if (!st->r.next_waiter || st->r.next_reading) return;
   if (st->errored && !st->error_emitted) {
@@ -1964,7 +1972,7 @@ static void scr_stream_rbuf_to_dyn(ScrStreamState *st) {
  * A Buffer on a stream with a decoder decodes through it (Node's
  * decoder.write in readableAddChunkPushObjectMode); an entry of another
  * kind turns the buffer into dyn entries. */
-static void *scr_stream_entry_convert(ScrStreamState *st, void *entry, ScrEntryKind *kind) {
+static void *scr_stream_entry_convert_kind(ScrStreamState *st, void *entry, ScrEntryKind *kind) {
   ScrEntryKind want = scr_stream_rkind(st);
   if (*kind == want) return entry;
   if (*kind == SCR_ENTRY_BYTES && want == SCR_ENTRY_STR && st->r.decoding) {
@@ -2390,7 +2398,7 @@ static void scr_stream_writev_batch(ScrStream *s) {
 }
 
 /* The buffered writes, one at a time (Node's clearBuffer without _writev). */
-static void scr_stream_clear_buffer_each(ScrStream *s) {
+static void scr_stream_clear_buffer_entries(ScrStream *s) {
   ScrStreamState *st = s->st;
   if (st->w.buffer_processing) return;
   st->w.buffer_processing = true;
@@ -2412,6 +2420,12 @@ static void scr_stream_clear_buffer_each(ScrStream *s) {
   st->w.buffer_processing = false;
 }
 
+/* An empty write queue has nothing to clear. */
+static inline void scr_stream_clear_buffer_each(ScrStream *s) {
+  ScrStreamState *st = s->st;
+  if (st->w.head < st->w.n) scr_stream_clear_buffer_entries(s);
+}
+
 /* The fs streams' clearBuffer: nothing while constructing, and _writev
  * for more than one buffered chunk. */
 static SCR_COLD void scr_stream_clear_buffer_runtime(ScrStream *s) {
@@ -6698,6 +6712,10 @@ ScrDyn *scr_stream_chunk_encoding_value(ScrStream *s) {
   return enc ? scr_dyn_retain(enc) : scr_dyn_undefined();
 }
 
+bool scr_stream_chunk_encoding_undefined(ScrStream *s) {
+  return scr_stream_writes_values(s) && s->st->w.cur_enc == NULL;
+}
+
 /* ── the readable surface ─────────────────────────────────────────────── */
 
 bool scr_stream_push(ScrStream *s, ScrBytes *chunk) {

packages/runtime/src/scr_stream.h

@@ -556,6 +556,9 @@ ScrBytes *scr_stream_chunk_bytes(ScrStream *s, ScrBytes *chunk);
 struct ScrDyn *scr_stream_chunk_value(ScrStream *s, ScrBytes *chunk);
 ScrStr *scr_stream_chunk_encoding(ScrStream *s);
 ScrDyn *scr_stream_chunk_encoding_value(ScrStream *s);
+/* True when the write in progress is an object-mode write that passed no
+ * encoding argument (its encoding value is undefined). */
+bool scr_stream_chunk_encoding_undefined(ScrStream *s);
 /* emit('data', value) with a checked-dynamic chunk (borrowed). */
 bool scr_stream_emit_value(ScrStream *s, ScrStr *name, struct ScrDyn *v);
 /* Whether a dyn for-await step (readable.nextChunkDyn) answered the end. */

tests/corpus/8520-stream-object-encoding-arg.ts

@@ -0,0 +1,42 @@
+// The encoding argument a typed write or transform callback receives:
+// undefined for an object-mode write without one, the given string with
+// one, and 'buffer' for byte-mode chunks.
+import { Transform, Writable } from "node:stream";
+
+const seen: string[] = [];
+const objects = new Writable({
+  objectMode: true,
+  write(chunk: { n: number }, encoding, cb) {
+    seen.push(`object ${chunk.n} ${typeof encoding} ${String(encoding)}`);
+    cb();
+  },
+});
+objects.write({ n: 1 });
+objects.write({ n: 2 }, "utf8");
+objects.write({ n: 3 }, undefined);
+objects.end(() => {
+  const bytes = new Writable({
+    write(chunk: Buffer, encoding, cb) {
+      seen.push(`bytes ${chunk.length} ${typeof encoding} ${String(encoding)}`);
+      cb();
+    },
+  });
+  bytes.write("abc");
+  bytes.write(Buffer.from("de"));
+  bytes.end(() => {
+    const upper = new Transform({
+      objectMode: true,
+      transform(row: { n: number }, encoding, cb) {
+        seen.push(`transform ${row.n} ${typeof encoding} ${String(encoding)}`);
+        cb(null, { n: row.n * 10 });
+      },
+    });
+    upper.on("data", (row: { n: number }) => seen.push(`data ${row.n}`));
+    upper.on("end", () => {
+      for (const line of seen) console.log(line);
+    });
+    upper.write({ n: 4 });
+    upper.write({ n: 5 }, "latin1");
+    upper.end();
+  });
+});

tests/corpus/8531-stream-saved-callback-receiver.mjs

@@ -0,0 +1,34 @@
+// A write's saved completion called later inside another object's method:
+// the write(chunk, callback) callback it runs is called unbound, so its
+// `this` is undefined (module code is strict), and the next queued write's
+// _write sees the stream as `this`.
+import { Writable } from "node:stream";
+
+let saved = null;
+const seen = [];
+const label = (self) => (self === undefined ? "undefined" : self === w ? "stream" : self === holder ? "holder" : typeof self);
+const w = new Writable({
+  objectMode: true,
+  write(chunk, _enc, cb) {
+    seen.push(`write ${chunk.n} this=${label(this)}`);
+    if (chunk.n === 1) saved = cb;
+    else cb();
+  },
+});
+const holder = {
+  run() {
+    saved();
+  },
+};
+w.write({ n: 1 }, function () {
+  seen.push(`callback 1 this=${label(this)}`);
+});
+w.write({ n: 2 }, function () {
+  seen.push(`callback 2 this=${label(this)}`);
+});
+setTimeout(() => {
+  holder.run();
+  w.end(() => {
+    for (const line of seen) console.log(line);
+  });
+}, 0);

tests/corpus/8532-stream-saved-callback-receiver/label.js

@@ -0,0 +1,13 @@
+// write(chunk, callback) callbacks that report their `this`; module code is
+// strict, so an unbound call sees undefined.
+export function makeCallback(seen, name, holder) {
+  return function () {
+    seen.push(`${name} this=${this === undefined ? "undefined" : this === holder ? "holder" : typeof this}`);
+  };
+}
+
+// Calls holder.run() as a method from JavaScript: run's own call of the
+// saved completion happens while this call's receiver is the holder.
+export function callRun(holder) {
+  holder.run();
+}

tests/corpus/8532-stream-saved-callback-receiver/main.ts

@@ -0,0 +1,32 @@
+// A saved write completion called inside a class method, which JavaScript
+// calls as holder.run(), runs the write(chunk, callback) callbacks unbound:
+// a JavaScript callback sees `this` as undefined, not the method's
+// receiver.
+import { Writable } from "node:stream";
+import { callRun, makeCallback } from "./label.js";
+
+let saved: ((error?: Error | null) => void) | null = null;
+const seen: string[] = [];
+const w = new Writable({
+  objectMode: true,
+  write(chunk: { n: number }, _enc, cb) {
+    seen.push(`write ${chunk.n}`);
+    if (chunk.n === 1) saved = cb;
+    else cb();
+  },
+});
+class Holder {
+  label = "holder";
+  run(): void {
+    if (saved !== null) saved();
+  }
+}
+const holder = new Holder();
+w.write({ n: 1 }, makeCallback(seen, "callback 1", holder));
+w.write({ n: 2 }, makeCallback(seen, "callback 2", holder));
+setTimeout(() => {
+  callRun(holder);
+  w.end(() => {
+    for (const line of seen) console.log(line);
+  });
+}, 0);