From 8e148a79a987a93067d16681d4eb7ec41e09f605 Mon Sep 17 00:00:00 2001
From: Glen Keane <glenkeane.94@gmail.com>
Date: Thu, 8 Oct 2026 13:30:37 +0100
Subject: [PATCH] fix: invalidate checks after predicate cancellation or reset

---
 test/test-reentrancy.js | 92 ++++++++++++++++++++++++++++++++++++++++
 when.js                 | 94 ++++++++++++++++++-----------------------
 2 files changed, 134 insertions(+), 52 deletions(-)
 create mode 100644 test/test-reentrancy.js

diff --git a/test/test-reentrancy.js b/test/test-reentrancy.js
new file mode 100644
index 0000000..3682df9
--- /dev/null
+++ b/test/test-reentrancy.js
@@ -0,0 +1,92 @@
+const test = require("node:test");
+const assert = require("node:assert/strict");
+const { setImmediate } = require("node:timers/promises");
+const when = require("../");
+
+test("clear inside a false predicate stops future checks", async (t) => {
+  let checks = 0;
+  let calls = 0;
+  const handle = when(
+    () => {
+      checks++;
+      if (checks === 2) handle.clear();
+      return checks >= 4;
+    },
+    () => calls++,
+  );
+  t.after(() => handle.clear());
+
+  await setImmediate();
+  await setImmediate();
+  await setImmediate();
+
+  assert.equal(checks, 2);
+  assert.equal(calls, 0);
+});
+
+test("clear inside a truthy predicate suppresses its callback", async (t) => {
+  let checks = 0;
+  let calls = 0;
+  const handle = when(
+    () => {
+      checks++;
+      if (checks === 2) {
+        handle.clear();
+        return true;
+      }
+      return false;
+    },
+    () => calls++,
+  );
+  t.after(() => handle.clear());
+
+  await setImmediate();
+
+  assert.equal(calls, 0);
+});
+
+test("reset ignores a truthy result from the previous check", async (t) => {
+  let checks = 0;
+  let calls = 0;
+  const handle = when(
+    () => {
+      checks++;
+      if (checks === 2) {
+        handle.reset();
+        return true;
+      }
+      return false;
+    },
+    () => calls++,
+  );
+  t.after(() => handle.clear());
+
+  await setImmediate();
+
+  assert.equal(checks, 3);
+  assert.equal(calls, 0);
+  handle.setCondition(() => true);
+  await setImmediate();
+  assert.equal(calls, 1);
+});
+
+test("reset inside a false predicate leaves only one cancellable poll", async (t) => {
+  let checks = 0;
+  let calls = 0;
+  const handle = when(
+    () => {
+      checks++;
+      if (checks === 2) handle.reset();
+      return checks >= 4;
+    },
+    () => calls++,
+  );
+  t.after(() => handle.clear());
+
+  await setImmediate();
+  handle.clear();
+  await setImmediate();
+
+  assert.equal(checks, 3);
+  assert.equal(calls, 0);
+});
diff --git a/when.js b/when.js
index 6a474b4..84b612a 100644
--- a/when.js
+++ b/when.js
@@ -1,52 +1,42 @@
-// the require('setImmediate') adds a setImmediate polyfill to the global
-// scope. (this is for backwards and browser compatibility)
-require('setimmediate');
-
-function when(condition, code){
-  // set immediate to a blank function, I don't want to push the first
-  // check condition to a back of a queue
-  var immediate = setImmediate(function(){});
-
-  checkCondition(); //do a check straight away.
-
-
-  function checkCondition(){
-    if(condition()){
-      clear(); //there might be nothing to clear, but theres no harm calling it.
-      code(); //if the condition is true, run the code striaght away.
-      
-      // in previous releases, I pushed the code() to the back of the io queue again with a setImmediate,
-      // but there is no point in doing that, making the program wait longer to execute something
-      // which might change the state to fire another when function.
-    } else {
-      immediate = setImmediate(checkCondition);
-    }
-  }
-
-  function clear(){
-    clearImmediate(immediate);
-  }
-
-  function reset(){
-    clear();
-    immediate = setImmediate(function(){});
-    checkCondition(); 
-  }
-
-  function setCondition(newCondition){
-    condition = newCondition;   
-  }
-
-  function setCode(newCode){
-    code = newCode;
-  }
-
-  return {
-    'clear': clear, 
-    'reset': reset,
-    'setCondition': setCondition,
-    'setCode': setCode
-  };
-};
-
-module.exports = exports = when;
+require("setimmediate");
+
+function when(condition, code) {
+  let immediate;
+  let generation = 0;
+
+  function checkCondition() {
+    const currentGeneration = generation;
+    const ready = condition();
+    if (currentGeneration !== generation) return;
+
+    if (ready) {
+      clear();
+      code();
+    } else {
+      immediate = setImmediate(checkCondition);
+    }
+  }
+
+  function clear() {
+    generation++;
+    clearImmediate(immediate);
+  }
+
+  function reset() {
+    clear();
+    checkCondition();
+  }
+
+  function setCondition(newCondition) {
+    condition = newCondition;
+  }
+
+  function setCode(newCode) {
+    code = newCode;
+  }
+
+  checkCondition();
+  return { clear, reset, setCondition, setCode };
+}
+
+module.exports = when;