Prep for step-reversing protocols

src/fibers/io/fd_backend/core.lua

@@ -77,8 +77,8 @@ end
 ---
 --- Required ops:
 ---   set_nonblock(fd) -> ok:boolean, err|nil
----   read(fd, max)    -> s|nil, err|nil
----   write(fd, s, len)-> n|nil, err|nil
+---   read(fd, max)    -> s|nil, err|nil, want?
+---   write(fd, s, len)-> n|nil, err|nil, want?
 ---   seek(fd, whence, off) -> pos|nil, err|nil
 ---   close(fd)        -> ok:boolean, err|nil
 ---

src/fibers/io/fd_backend/ffi.lua

@@ -167,7 +167,7 @@ local function read_fd(fd, max)
 	if n < 0 then
 		local e = get_errno()
 		if e == EAGAIN or e == EWOULDBLOCK then
-			return nil, nil -- would block
+			return nil, nil, 'rd' -- would block
 		end
 		return nil, strerror(e)
 	end
@@ -191,7 +191,7 @@ local function write_fd(fd, str, len)
 	if n < 0 then
 		local e = get_errno()
 		if e == EAGAIN or e == EWOULDBLOCK then
-			return nil, nil -- would block
+			return nil, nil, 'wr' -- would block
 		end
 		return nil, strerror(e)
 	end

src/fibers/io/fd_backend/nixio.lua

@@ -77,7 +77,7 @@ local function read_fd(fd, max)
 
 	if eno == EAGAIN or eno == EWOULDBLOCK then
 		-- Would block, signal “not ready yet”.
-		return nil, nil
+		return nil, nil, 'rd'
 	end
 
 	if not eno or eno == 0 then
@@ -110,7 +110,7 @@ local function write_fd(fd, str, len)
 
 	if eno == EAGAIN or eno == EWOULDBLOCK then
 		-- Would block.
-		return nil, nil
+		return nil, nil, 'wr'
 	end
 
 	return nil, errno_msg(msg or 'write failed', eno)

src/fibers/io/fd_backend/posix.lua

@@ -48,7 +48,7 @@ local function read_fd(fd, max)
 	local s, err, eno = unistd.read(fd, max)
 	if s == nil then
 		if eno == errno.EAGAIN or eno == errno.EWOULDBLOCK then
-			return nil, nil
+			return nil, nil, 'rd'
 		end
 		return nil, errno_msg('read failed', err, eno)
 	end
@@ -59,7 +59,7 @@ local function write_fd(fd, str, _)
 	local n, err, eno = unistd.write(fd, str)
 	if n == nil then
 		if eno == errno.EAGAIN or eno == errno.EWOULDBLOCK then
-			return nil, nil
+			return nil, nil, 'wr'
 		end
 		return nil, errno_msg('write failed', err, eno)
 	end

src/fibers/io/stream.lua

@@ -129,15 +129,15 @@ local function make_read_step(stream, buf, min, max, terminator)
 				return true, buf, tally, 'buffer capacity exhausted'
 			end
 
-			local data, err = stream.io:read_string(room)
+			local data, err, want = stream.io:read_string(room)
 			if err then
 				return true, buf, tally, err
 			end
 			if not data then
 				if tally >= min then
 					return true, buf, tally
 				end
-				return false
+				return false, want
 			end
 			if #data == 0 then
 				return true, buf, tally
@@ -165,12 +165,12 @@ local function make_write_step(stream, src_str)
 		end
 
 		local chunk = src_str:sub(offset + 1)
-		local n, err = stream.io:write_string(chunk)
+		local n, err, want = stream.io:write_string(chunk)
 		if err then
 			return true, offset, err
 		end
 		if n == nil then
-			return false
+			return false, want
 		end
 		if n == 0 then
 			return true, offset
@@ -210,7 +210,10 @@ function Stream:read_into_op(buf, opts)
 	end
 
 	return wait.waitable(
-		function (task)
+		function (task, _, _, want)
+			if want == 'wr' then
+				return self.io:on_writable(task)
+			end
 			return self.io:on_readable(task)
 		end,
 		step,
@@ -249,7 +252,10 @@ function Stream:write_string_op(str)
 	end
 
 	return wait.waitable(
-		function (task)
+		function (task, _, _, want)
+			if want == 'rd' then
+				return self.io:on_readable(task)
+			end
 			return self.io:on_writable(task)
 		end,
 		step,

src/fibers/wait.lua

@@ -215,6 +215,12 @@ end
 -- waitable: (register, step, wrap_fn?) -> Op
 ----------------------------------------------------------------------
 
+local function normalise_want(want)
+	if want == 'rd' or want == 'wr' or want == 'any' then
+		return want
+	end
+	return nil
+end
 --- Build a waitable Op from a register function and step function.
 --
 --   step() -> done:boolean, ...
@@ -240,7 +246,7 @@ end
 --
 -- The op participates fully in choice/with_nack/on_abort; if it loses
 -- a choice, any outstanding registration is cancelled via token:unlink().
----@param register fun(task: Task, suspension: Suspension, leaf_wrap: WrapFn): WaitToken
+---@param register fun(task: Task, suspension: Suspension, leaf_wrap: WrapFn, want: any): WaitToken
 ---@param step fun(): boolean, ...
 ---@param wrap_fn? WrapFn
 ---@return Op
@@ -251,64 +257,68 @@ local function waitable(register, step, wrap_fn)
 	wrap_fn = wrap_fn or id_wrap
 
 	return op.guard(function ()
-		-- Token for this synchronisation (one per compiled leaf).
 		local token
 
-		-- Fast path: single non-blocking attempt.
-		-- step() must not yield; if it raises, this fiber fails.
+		local function unlink_token()
+			if token and token.unlink then
+				token:unlink()
+			end
+			token = nil
+		end
 		local function try()
 			return step()
 		end
 
-		--- Blocking path: register a task that will re-run step
-		--- after the external condition changes.
-		---
-		--- The same task is re-used across wake-ups; token is updated
-		--- to track the latest registration.
-		---@param suspension Suspension
-		---@param leaf_wrap WrapFn
 		local function block(suspension, leaf_wrap)
 			---@class WaitTask : Task
 			local task
 
+			local function register_with_want(want)
+				unlink_token()
+
+				if want == 'any' then
+					local t1 = register(task, suspension, leaf_wrap, 'rd')
+					local t2 = register(task, suspension, leaf_wrap, 'wr')
+					token = {
+						unlink = function ()
+							if t1 and t1.unlink then t1:unlink() end
+							if t2 and t2.unlink then t2:unlink() end
+						end,
+					}
+				else
+					token = register(task, suspension, leaf_wrap, want)
+				end
+			end
 			task = {
 				run = function ()
 					if not suspension:waiting() then
 						return
 					end
 
-					-- Re-check readiness.
 					local res  = pack(step())
 					local done = res[1]
 					if done then
-						-- Complete with the leaf's final wrap.
-						return suspension:complete(
-							leaf_wrap,
-							unpack(res, 2, res.n)
-						)
+						unlink_token()
+						return suspension:complete(leaf_wrap, unpack(res, 2, res.n))
 					end
 
-					-- Not done yet; re-register for another wake-up.
-					if token and token.unlink then
-						token:unlink()
-					end
-
-					token = register(task, suspension, leaf_wrap)
+					register_with_want(normalise_want(res[2]))
 				end,
 			}
 
-			-- Initial registration for this synchronisation.
-			token = register(task, suspension, leaf_wrap)
+			-- Initial drive: decide readiness *before* registering.
+			local first = pack(step())
+			if first[1] then
+				return suspension:complete(leaf_wrap, unpack(first, 2, first.n))
+			end
+
+			register_with_want(normalise_want(first[2]))
 		end
 
 		local prim = op.new_primitive(wrap_fn, try, block)
 
-		-- If this op participates in a choice and loses, ensure any
-		-- extant registration is cancelled once for this synchronisation.
 		return prim:on_abort(function ()
-			if token and token.unlink then
-				token:unlink()
-			end
+			unlink_token()
 		end)
 	end)
 end