CABI: simplify future logic by splitting out from streams

Author: lukewagner

design/mvp/canonical-abi/definitions.py

@@ -853,26 +853,98 @@ def copy(self, inst, src, on_partial_copy, on_copy_done):
 
 #### Future State
 
-class FutureEnd(StreamEnd):
-  def close_after_copy(self, copy_op, inst, buffer, on_copy_done):
-    assert(buffer.remain() == 1)
-    def on_copy_done_wrapper(why):
-      if buffer.remain() == 0:
-        self.shared.close()
-      on_copy_done(why)
-    ret = copy_op(inst, buffer, on_partial_copy = None, on_copy_done = on_copy_done_wrapper)
-    if ret == 'done' and buffer.remain() == 0:
-      self.shared.close()
-    return ret
+OnRead = Callable[[any|Literal['cancelled']]], None]
+OnWrite = Callable[[Optional[Literal['cancelled','dropped']]], Optional[any]]
+
+class ReadableFuture:
+  t: ValType
+  read: Callable[[ComponentInstance, OnRead], Literal['done','blocked']]
+  cancel: Callable[[], None]
+  drop: Callable[[]]
+
+class SharedFutureImpl(ReadableFuture):
+  reader_dropped: bool
+  pending_inst: Optional[ComponentInstance]
+  pending_on_read: Optional[OnRead]
+  pending_on_write: Optional[OnWrite]
+
+  def __init__(self, t):
+    self.t = t
+    self.reader_dropped = False
+    self.reset_pending()
+
+  def reset_pending(self):
+    self.pending_inst = None
+    self.pending_on_read = None
+    self.pending_on_write = None
+
+  def cancel(self):
+    assert(self.pending_on_read ^ self.pending_on_write)
+    if self.pending_on_read:
+      self.pending_on_read('cancelled')
+    else:
+      self.pending_on_write('cancelled')
+    self.reset_pending()
+
+  def drop(self):
+    assert(not self.reader_dropped and not self.pending_on_read)
+    self.reader_dropped = True
+    if self.pending_on_write:
+      self.pending_on_write('dropped')
+      self.reset_pending()
+
+  def read(self, inst, on_read):
+    assert(not self.reader_dropped and not self.pending_on_read)
+    if not self.pending_on_write:
+      self.pending_inst = inst
+      self.pending_on_read = on_read
+      return 'blocked'
+    else:
+      trap_if(inst is self.pending_inst and self.t is not None) # temporary
+      on_read(self.pending_write())
+      return 'done'
+
+  def write(self, inst, on_write):
+    assert(not self.pending_on_write)
+    if self.reader_dropped:
+      return 'done'
+    elif not self.pending_on_read:
+      self.pending_inst = inst
+      self.pending_on_write = on_write
+    else:
+      trap_if(inst is self.pending_inst and self.t is not None) # temporary
+      self.pending_on_read(on_write())
+      return 'done'
+
+class FutureEnd(Waitable):
+  shared: ReadableFuture
+  copying: bool
+  done: bool
+
+  def __init__(self, shared):
+    Waitable.__init__(self)
+    self.shared = shared
+    self.copying = False
+    self.done = False
+
+  def drop(self):
+    trap_if(self.copying)
+    Waitable.drop(self)
 
 class ReadableFutureEnd(FutureEnd):
-  def copy(self, inst, dst, on_partial_copy, on_copy_done):
-    return self.close_after_copy(self.shared.read, inst, dst, on_copy_done)
+  def copy(self, inst, on_read):
+    return self.shared.read(inst, on_read)
+
+  def drop(self):
+    self.shared.drop()
+    FutureEnd.drop(self)
 
 class WritableFutureEnd(FutureEnd):
-  def copy(self, inst, src, on_partial_copy, on_copy_done):
-    return self.close_after_copy(self.shared.write, inst, src, on_copy_done)
+  def copy(self, inst, on_write):
+    return self.shared.write(inst, on_write)
+
   def drop(self):
+    trap_if(not self.done)
     FutureEnd.drop(self)
 
 ### Despecialization
@@ -2158,50 +2230,42 @@ async def canon_stream_new(stream_t, task):
 
 async def canon_future_new(future_t, task):
   trap_if(not task.inst.may_leave)
-  shared = SharedStreamImpl(future_t.t)
+  shared = SharedFutureImpl(future_t.t)
   ri = task.inst.table.add(ReadableFutureEnd(shared))
   wi = task.inst.table.add(WritableFutureEnd(shared))
   return [ ri | (wi << 32) ]
 
-### 🔀 `canon {stream,future}.{read,write}`
+### 🔀 `canon stream.{read,write}`
 
 async def canon_stream_read(stream_t, opts, task, i, ptr, n):
-  return await copy(ReadableStreamEnd, WritableBufferGuestImpl, EventCode.STREAM_READ,
-                    stream_t, opts, task, i, ptr, n)
+  return await stream_copy(ReadableStreamEnd, WritableBufferGuestImpl, EventCode.STREAM_READ,
+                           stream_t, opts, task, i, ptr, n)
 
 async def canon_stream_write(stream_t, opts, task, i, ptr, n):
-  return await copy(WritableStreamEnd, ReadableBufferGuestImpl, EventCode.STREAM_WRITE,
-                    stream_t, opts, task, i, ptr, n)
-
-async def canon_future_read(future_t, opts, task, i, ptr):
-  return await copy(ReadableFutureEnd, WritableBufferGuestImpl, EventCode.FUTURE_READ,
-                    future_t, opts, task, i, ptr, 1)
-
-async def canon_future_write(future_t, opts, task, i, ptr):
-  return await copy(WritableFutureEnd, ReadableBufferGuestImpl, EventCode.FUTURE_WRITE,
-                    future_t, opts, task, i, ptr, 1)
+  return await stream_copy(WritableStreamEnd, ReadableBufferGuestImpl, EventCode.STREAM_WRITE,
+                           stream_t, opts, task, i, ptr, n)
 
-async def copy(EndT, BufferT, event_code, stream_or_future_t, opts, task, i, ptr, n):
+async def stream_copy(EndT, BufferT, event_code, stream_t, opts, task, i, ptr, n):
   trap_if(not task.inst.may_leave)
   e = task.inst.table.get(i)
   trap_if(not isinstance(e, EndT))
-  trap_if(e.shared.t != stream_or_future_t.t)
+  trap_if(e.shared.t != stream_t.t)
   trap_if(e.copying)
 
-  assert(not contains_borrow(stream_or_future_t))
+  assert(not contains_borrow(stream_t))
   cx = LiftLowerContext(opts, task.inst, borrow_scope = None)
-  buffer = BufferT(stream_or_future_t.t, cx, ptr, n)
+  buffer = BufferT(stream_t.t, cx, ptr, n)
 
-  def copy_event(why, revoke_buffer):
+  def stream_copy_event(why, revoke_buffer):
     revoke_buffer()
     e.copying = False
     return (event_code, i, pack_copy_result(task, e, buffer, why))
 
   def on_partial_copy(revoke_buffer):
-    e.set_event(partial(copy_event, 'completed', revoke_buffer))
+    e.set_event(partial(stream_copy_event, 'completed', revoke_buffer))
 
   def on_copy_done(why):
-    e.set_event(partial(copy_event, why, revoke_buffer = lambda:()))
+    e.set_event(partial(stream_copy_event, why, revoke_buffer = lambda:()))
 
   if e.copy(task.inst, buffer, on_partial_copy, on_copy_done) == 'done':
     return [pack_copy_result(task, e, buffer, 'completed')]
@@ -2234,6 +2298,55 @@ def pack_copy_result(task, e, buffer, why):
   assert(packed != BLOCKED)
   return packed
 
+### 🔀 `canon future.{read,write}`
+
+async def canon_future_read(future_t, opts, task, i, ptr):
+  assert(not contains_borrow(stream_t))
+  cx = LiftLowerContext(opts, task.inst, borrow_scope = None)
+  def lower(v):
+    store(cx, v, future_t.t, ptr)
+  return await future_copy(ReadableFutureEnd, EventCode.FUTURE_READ, lower,
+                           future_t, opts, task, i, ptr)
+
+async def canon_future_write(future_t, opts, task, i, ptr):
+  assert(not contains_borrow(stream_t))
+  cx = LiftLowerContext(opts, task.inst, borrow_scope = None)
+  def lift(_):
+    return load(cx, ptr, future_t.t)
+  return await future_copy(WritableFutureEnd, EventCode.FUTURE_WRITE, lift,
+                           future_t, opts, task, i, ptr)
+
+async def future_copy(EndT, event_code, copy, future_t, opts, task, i, ptr):
+  trap_if(not task.inst.may_leave)
+  e = task.inst.table.get(i)
+  trap_if(not isinstance(e, EndT))
+  trap_if(e.shared.t != future_t.t)
+  trap_if(e.copying or e.done)
+
+  def future_copy_event(why):
+    e.copying = False
+    if why != 'cancelled':
+      e.done = True
+    return (event_code, i, pack_copy_result(task, e, buffer, why))
+
+  def on_copy(why, v):
+    assert(not e.has_pending_event())
+    e.set_event(partial(future_copy_event, why))
+    if why != 'dropped' and why != 'cancelled':
+      return copy(v)
+
+  if e.copy(task.inst, on_copy) == 'done':
+    return [pack_copy_result(task, e, buffer, 'completed')]
+  else:
+    if opts.sync:
+      await task.wait_on(e.wait_for_pending_event(), sync = True)
+      code,index,payload = e.get_event()
+      assert(code == event_code and index == i)
+      return [payload]
+    else:
+      e.copying = True
+      return [BLOCKED]
+
 ### 🔀 `canon {stream,future}.cancel-{read,write}`
 
 async def canon_stream_cancel_read(stream_t, sync, task, i):

design/mvp/canonical-abi/run_tests.py

@@ -1815,7 +1815,7 @@ def __init__(self, t):
     self.pending_on_copy_done = None
   def closed(self):
     return self.is_closed
-  def read(self, inst, dst, on_partial_copy, on_copy_done):
+  def read(self, inst, dst, on_copy_done):
     if self.is_closed:
       return 'done'
     elif self.v:
@@ -1855,7 +1855,7 @@ async def host_func(task, on_start, on_resolve, on_block):
     outgoing = HostFutureSource(U8Type())
     on_resolve([outgoing])
     incoming = HostFutureSink(U8Type())
-    future.read(None, incoming, lambda:(), lambda why:())
+    future.read(None, incoming, lambda why:())
     wait = asyncio.create_task(incoming.has_v.wait())
     await on_block(wait)
     assert(incoming.v == 42)