CABI: simplify future logic by splitting out from streams

Author: lukewagner

design/mvp/canonical-abi/definitions.py

@@ -751,7 +751,7 @@ def drop(self):
 
 RevokeBuffer = Callable[[], None]
 OnPartialCopy = Callable[[RevokeBuffer], None]
-OnCopyDone = Callable[[Literal['completed','cancelled']], None]
+OnCopyDone = Callable[[Literal['completed','cancelled','closed']], None]
 
 class ReadableStream:
   t: ValType
@@ -793,7 +793,7 @@ def close(self):
     if not self.closed_:
       self.closed_ = True
       if self.pending_buffer:
-        self.reset_and_notify_pending('completed')
+        self.reset_and_notify_pending('closed')
 
   def closed(self):
     return self.closed_
@@ -853,26 +853,96 @@ 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
+LiftValue = Callable[[], any]
+LowerValue = Callable[[any], None]
+
+class ReadableFuture:
+  t: ValType
+  read: Callable[[ComponentInstance, LowerValue, OnCopyDone], Literal['done','blocked']]
+  cancel: Callable[[], None]
+  drop: Callable[[]]
+
+class SharedFutureImpl(ReadableFuture):
+  dropped: bool
+  pending_inst: Optional[ComponentInstance]
+  pending_copy_value: Optional[LiftValue|LowerValue]
+  pending_on_copy_done: Optional[OnCopyDone]
+
+  def __init__(self, t):
+    self.t = t
+    self.dropped = False
+    self.reset_pending()
+
+  def reset_pending(self):
+    self.set_pending(None, None, None)
+
+  def set_pending(self, inst, copy_value, on_copy_done):
+    self.pending_inst = inst
+    self.pending_copy_value = copy_value
+    self.pending_on_copy_done = on_copy_done
+
+  def reset_and_notify_pending(self, why):
+    pending_on_copy_done = self.pending_on_copy_done
+    self.reset_pending()
+    pending_on_copy_done(why)
+
+  def cancel(self):
+    self.reset_and_notify_pending('cancelled')
+
+  def drop(self):
+    assert(not self.dropped)
+    self.dropped = True
+    if self.pending_on_copy_done:
+      self.reset_and_notify_pending('closed')
+
+  def read(self, inst, lower_value, on_copy_done):
+    assert(not self.dropped)
+    return self.copy(inst, lower_value, on_copy_done, self.pending_copy_value, lower_value)
+
+  def write(self, inst, lift_value, on_copy_done):
+    if self.dropped:
+      return 'closed'
+    return self.copy(inst, lift_value, on_copy_done, lift_value, self.pending_copy_value)
+
+  def copy(self, inst, copy_value, on_copy_done, lift_value, lower_value):
+    if not self.pending_copy_value:
+      self.set_pending(inst, copy_value, on_copy_done)
+      return 'blocked'
+    else:
+      trap_if(inst is self.pending_inst and self.t is not None) # temporary
+      lower_value(lift_value())
+      self.reset_and_notify_pending('done')
+      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, lower_value, on_copy_done):
+    return self.shared.read(inst, lower_value, on_copy_done)
+
+  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, lift_value, on_copy_done):
+    return self.shared.write(inst, lift_value, on_copy_done)
+
   def drop(self):
+    trap_if(not self.done)
     FutureEnd.drop(self)
 
 ### Despecialization
@@ -1201,19 +1271,20 @@ def lift_borrow(cx, i, t):
   return h.rep
 
 def lift_stream(cx, i, t):
-  return lift_async_value(ReadableStreamEnd, cx, i, t)
+  assert(not contains_borrow(t))
+  e = cx.inst.table.remove(i)
+  trap_if(not isinstance(e, ReadableStreamEnd))
+  trap_if(e.shared.t != t)
+  trap_if(e.copying)
+  return e.shared
 
 def lift_future(cx, i, t):
-  v = lift_async_value(ReadableFutureEnd, cx, i, t)
-  trap_if(v.closed())
-  return v
-
-def lift_async_value(ReadableEndT, cx, i, t):
   assert(not contains_borrow(t))
   e = cx.inst.table.remove(i)
-  trap_if(not isinstance(e, ReadableEndT))
+  trap_if(not isinstance(e, ReadableFutureEnd))
   trap_if(e.shared.t != t)
   trap_if(e.copying)
+  trap_if(e.done)
   return e.shared
 
 ### Storing
@@ -1507,16 +1578,14 @@ def lower_borrow(cx, rep, t):
   return cx.inst.table.add(h)
 
 def lower_stream(cx, v, t):
-  return lower_async_value(ReadableStreamEnd, cx, v, t)
+  assert(isinstance(v, ReadableStream))
+  assert(not contains_borrow(t))
+  return cx.inst.table.add(ReadableStreamEnd(v))
 
 def lower_future(cx, v, t):
-  assert(not v.closed())
-  return lower_async_value(ReadableFutureEnd, cx, v, t)
-
-def lower_async_value(ReadableEndT, cx, v, t):
-  assert(isinstance(v, ReadableStream))
+  assert(isinstance(v, ReadableFuture))
   assert(not contains_borrow(t))
-  return cx.inst.table.add(ReadableEndT(v))
+  return cx.inst.table.add(ReadableFutureEnd(v))
 
 ### Flattening
 
@@ -2161,45 +2230,37 @@ 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):
     revoke_buffer()
     assert(e.copying)
     e.copying = False
-    return (event_code, i, pack_copy_result(task, e, buffer, why))
+    return (event_code, i, pack_stream_result(task, e, buffer, why))
 
   def on_partial_copy(revoke_buffer):
     e.set_event(partial(copy_event, 'completed', revoke_buffer))
@@ -2208,7 +2269,7 @@ def on_copy_done(why):
     e.set_event(partial(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')]
+    return [pack_stream_result(task, e, buffer, 'completed')]
   else:
     e.copying = True
     if opts.sync:
@@ -2224,7 +2285,7 @@ def on_copy_done(why):
 CLOSED    = 0x1
 CANCELLED = 0x2
 
-def pack_copy_result(task, e, buffer, why):
+def pack_stream_result(task, e, buffer, why):
   if e.shared.closed():
     result = CLOSED
   elif why == 'cancelled':
@@ -2238,6 +2299,66 @@ 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):
+  def lower_value(v):
+    if future_t.t:
+      assert(not contains_borrow(future_t))
+      cx = LiftLowerContext(opts, task.inst, borrow_scope = None)
+      store(cx, v, future_t.t, ptr)
+
+  return await future_copy(ReadableFutureEnd, EventCode.FUTURE_READ, lower_value,
+                           future_t, opts, task, i, ptr)
+
+async def canon_future_write(future_t, opts, task, i, ptr):
+  def lift_value():
+    if future_t.t:
+      assert(not contains_borrow(future_t))
+      cx = LiftLowerContext(opts, task.inst, borrow_scope = None)
+      return load(cx, ptr, future_t.t)
+
+  return await future_copy(WritableFutureEnd, EventCode.FUTURE_WRITE, lift_value,
+                           future_t, opts, task, i, ptr)
+
+async def future_copy(EndT, event_code, copy_value, 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 on_copy_done(why):
+    def copy_event():
+      assert(e.copying)
+      e.copying = False
+      if why != 'cancelled':
+        e.done = True
+      return (event_code, i, pack_future_result(task, e, why))
+    assert(not e.has_pending_event())
+    e.set_event(copy_event)
+
+  result = e.copy(task.inst, copy_value, on_copy_done)
+  if result != 'blocked':
+    e.done = True
+    return [pack_future_result(task, e, result)]
+  else:
+    e.copying = True
+    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:
+      return [BLOCKED]
+
+def pack_future_result(task, e, why):
+  match why:
+    case 'cancelled': return CANCELLED
+    case 'closed':    return CLOSED
+    case 'done':      return (CLOSED | (1 << 4))
+  assert(False)
+
 ### 🔀 `canon {stream,future}.cancel-{read,write}`
 
 async def canon_stream_cancel_read(stream_t, sync, task, i):