CABI: simplify stream code a bit more (no change)

Author: lukewagner

design/mvp/CanonicalABI.md

@@ -2827,16 +2827,13 @@ of component-level values with types `ts`.
 def lift_flat_values(cx, max_flat, vi, ts):
   flat_types = flatten_types(ts)
   if len(flat_types) > max_flat:
-    return lift_heap_values(cx, vi, ts)
+    ptr = vi.next('i32')
+    tuple_type = TupleType(ts)
+    trap_if(ptr != align_to(ptr, alignment(tuple_type)))
+    trap_if(ptr + elem_size(tuple_type) > len(cx.opts.memory))
+    return list(load(cx, ptr, tuple_type).values())
   else:
     return [ lift_flat(cx, vi, t) for t in ts ]
-
-def lift_heap_values(cx, vi, ts):
-  ptr = vi.next('i32')
-  tuple_type = TupleType(ts)
-  trap_if(ptr != align_to(ptr, alignment(tuple_type)))
-  trap_if(ptr + elem_size(tuple_type) > len(cx.opts.memory))
-  return list(load(cx, ptr, tuple_type).values())
 ```
 
 Symmetrically, the `lower_flat_values` function defines how to lower a
@@ -2850,27 +2847,23 @@ def lower_flat_values(cx, max_flat, vs, ts, out_param = None):
   cx.inst.may_leave = False
   flat_types = flatten_types(ts)
   if len(flat_types) > max_flat:
-    flat_vals = lower_heap_values(cx, vs, ts, out_param)
+    tuple_type = TupleType(ts)
+    tuple_value = {str(i): v for i,v in enumerate(vs)}
+    if out_param is None:
+      ptr = cx.opts.realloc(0, 0, alignment(tuple_type), elem_size(tuple_type))
+      flat_vals = [ptr]
+    else:
+      ptr = out_param.next('i32')
+      flat_vals = []
+    trap_if(ptr != align_to(ptr, alignment(tuple_type)))
+    trap_if(ptr + elem_size(tuple_type) > len(cx.opts.memory))
+    store(cx, tuple_value, tuple_type, ptr)
   else:
     flat_vals = []
     for i in range(len(vs)):
       flat_vals += lower_flat(cx, vs[i], ts[i])
   cx.inst.may_leave = True
   return flat_vals
-
-def lower_heap_values(cx, vs, ts, out_param):
-  tuple_type = TupleType(ts)
-  tuple_value = {str(i): v for i,v in enumerate(vs)}
-  if out_param is None:
-    ptr = cx.opts.realloc(0, 0, alignment(tuple_type), elem_size(tuple_type))
-    flat_vals = [ptr]
-  else:
-    ptr = out_param.next('i32')
-    flat_vals = []
-  trap_if(ptr != align_to(ptr, alignment(tuple_type)))
-  trap_if(ptr + elem_size(tuple_type) > len(cx.opts.memory))
-  store(cx, tuple_value, tuple_type, ptr)
-  return flat_vals
 ```
 The `may_leave` flag is guarded by `canon_lower` below to prevent a component
 from calling out of the component while in the middle of lowering, ensuring

design/mvp/canonical-abi/definitions.py

@@ -307,6 +307,7 @@ class Buffer:
   MAX_LENGTH = 2**28 - 1
   t: ValType
   remain: Callable[[], int]
+  is_zero_length: Callable[[], bool]
 
 class ReadableBuffer(Buffer):
   read: Callable[[int], list[any]]
@@ -335,6 +336,9 @@ def __init__(self, t, cx, ptr, length):
   def remain(self):
     return self.length - self.progress
 
+  def is_zero_length(self):
+    return self.length == 0
+
 class ReadableBufferGuestImpl(BufferGuestImpl):
   def read(self, n):
     assert(n <= self.remain())
@@ -749,13 +753,18 @@ def drop(self):
 
 #### Stream State
 
+class CopyResult(IntEnum):
+  COMPLETED = 0
+  CLOSED = 1
+  CANCELLED = 2
+
 RevokeBuffer = Callable[[], None]
-OnPartialCopy = Callable[[RevokeBuffer], None]
-OnCopyDone = Callable[[Literal['completed','cancelled']], None]
+OnCopy = Callable[[RevokeBuffer], None]
+OnCopyResult = Callable[[CopyResult], None]
 
 class ReadableStream:
   t: ValType
-  read: Callable[[ComponentInstance, WritableBuffer, OnPartialCopy, OnCopyDone], Literal['done','blocked']]
+  read: Callable[[ComponentInstance, WritableBuffer, OnCopy, OnCopyResult], Optional[CopyResult]]
   cancel: Callable[[], None]
   close: Callable[[]]
   closed: Callable[[], bool]
@@ -764,8 +773,8 @@ class SharedStreamImpl(ReadableStream):
   closed_: bool
   pending_inst: Optional[ComponentInstance]
   pending_buffer: Optional[Buffer]
-  pending_on_partial_copy: Optional[OnPartialCopy]
-  pending_on_copy_done: Optional[OnCopyDone]
+  pending_on_copy: Optional[OnCopy]
+  pending_on_copy_result: Optional[OnCopyResult]
 
   def __init__(self, t):
     self.t = t
@@ -775,59 +784,55 @@ def __init__(self, t):
   def reset_pending(self):
     self.set_pending(None, None, None, None)
 
-  def set_pending(self, inst, buffer, on_partial_copy, on_copy_done):
+  def set_pending(self, inst, buffer, on_copy, on_copy_result):
     self.pending_inst = inst
     self.pending_buffer = buffer
-    self.pending_on_partial_copy = on_partial_copy
-    self.pending_on_copy_done = on_copy_done
+    self.pending_on_copy = on_copy
+    self.pending_on_copy_result = on_copy_result
 
-  def reset_and_notify_pending(self, why):
-    pending_on_copy_done = self.pending_on_copy_done
+  def reset_and_notify_pending(self, result):
+    pending_on_copy_result = self.pending_on_copy_result
     self.reset_pending()
-    pending_on_copy_done(why)
+    pending_on_copy_result(result)
 
   def cancel(self):
-    self.reset_and_notify_pending('cancelled')
+    self.reset_and_notify_pending(CopyResult.CANCELLED)
 
   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(CopyResult.CLOSED)
 
   def closed(self):
     return self.closed_
 
-  def read(self, inst, dst, on_partial_copy, on_copy_done):
-    return self.copy(inst, dst, on_partial_copy, on_copy_done, self.pending_buffer, dst)
+  def read(self, inst, dst, on_copy, on_copy_result):
+    return self.copy(inst, dst, on_copy, on_copy_result, self.pending_buffer, dst)
 
-  def write(self, inst, src, on_partial_copy, on_copy_done):
-    return self.copy(inst, src, on_partial_copy, on_copy_done, src, self.pending_buffer)
+  def write(self, inst, src, on_copy, on_copy_result):
+    return self.copy(inst, src, on_copy, on_copy_result, src, self.pending_buffer)
 
-  def copy(self, inst, buffer, on_partial_copy, on_copy_done, src, dst):
+  def copy(self, inst, buffer, on_copy, on_copy_result, src, dst):
     if self.closed_:
-      return 'done'
+      return CopyResult.CLOSED
     elif not self.pending_buffer:
-      self.set_pending(inst, buffer, on_partial_copy, on_copy_done)
-      return 'blocked'
+      self.set_pending(inst, buffer, on_copy, on_copy_result)
+      return None
     else:
       assert(self.t == src.t == dst.t)
       trap_if(inst is self.pending_inst and self.t is not None) # temporary
       if self.pending_buffer.remain() > 0:
         if buffer.remain() > 0:
           dst.write(src.read(min(src.remain(), dst.remain())))
-          if self.pending_buffer.remain() > 0:
-            self.pending_on_partial_copy(self.reset_pending)
-          else:
-            self.reset_and_notify_pending('completed')
-        return 'done'
+          self.pending_on_copy(self.reset_pending)
+        return CopyResult.COMPLETED
+      elif buffer is src and buffer.remain() == 0 and self.pending_buffer.is_zero_length():
+        return CopyResult.COMPLETED
       else:
-        if buffer.remain() > 0 or buffer is dst:
-          self.reset_and_notify_pending('completed')
-          self.set_pending(inst, buffer, on_partial_copy, on_copy_done)
-          return 'blocked'
-        else:
-          return 'done'
+        self.reset_and_notify_pending(CopyResult.COMPLETED)
+        self.set_pending(inst, buffer, on_copy, on_copy_result)
+        return None
 
 class StreamEnd(Waitable):
   shared: ReadableStream
@@ -844,34 +849,35 @@ def drop(self):
     Waitable.drop(self)
 
 class ReadableStreamEnd(StreamEnd):
-  def copy(self, inst, dst, on_partial_copy, on_copy_done):
-    return self.shared.read(inst, dst, on_partial_copy, on_copy_done)
+  def copy(self, inst, dst, on_copy, on_copy_result):
+    return self.shared.read(inst, dst, on_copy, on_copy_result)
 
 class WritableStreamEnd(StreamEnd):
-  def copy(self, inst, src, on_partial_copy, on_copy_done):
-    return self.shared.write(inst, src, on_partial_copy, on_copy_done)
+  def copy(self, inst, src, on_copy, on_copy_result):
+    return self.shared.write(inst, src, on_copy, on_copy_result)
 
 #### Future State
 
 class FutureEnd(StreamEnd):
-  def close_after_copy(self, copy_op, inst, buffer, on_copy_done):
+  def close_after_copy(self, copy_op, inst, buffer, on_copy_result):
     assert(buffer.remain() == 1)
-    def on_copy_done_wrapper(why):
+    def on_copy_result_wrapper(result):
       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:
+      on_copy_result(result)
+    ret = copy_op(inst, buffer, on_copy = lambda _:(), on_copy_result = on_copy_result_wrapper)
+    if ret is not None and buffer.remain() == 0:
       self.shared.close()
+      return CopyResult.CLOSED
     return ret
 
 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, dst, on_copy, on_copy_result):
+    return self.close_after_copy(self.shared.read, inst, dst, on_copy_result)
 
 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, src, on_copy, on_copy_result):
+    return self.close_after_copy(self.shared.write, inst, src, on_copy_result)
   def drop(self):
     trap_if(not self.shared.closed())
     FutureEnd.drop(self)
@@ -1802,43 +1808,36 @@ def lower_flat_flags(v, labels):
 def lift_flat_values(cx, max_flat, vi, ts):
   flat_types = flatten_types(ts)
   if len(flat_types) > max_flat:
-    return lift_heap_values(cx, vi, ts)
+    ptr = vi.next('i32')
+    tuple_type = TupleType(ts)
+    trap_if(ptr != align_to(ptr, alignment(tuple_type)))
+    trap_if(ptr + elem_size(tuple_type) > len(cx.opts.memory))
+    return list(load(cx, ptr, tuple_type).values())
   else:
     return [ lift_flat(cx, vi, t) for t in ts ]
 
-def lift_heap_values(cx, vi, ts):
-  ptr = vi.next('i32')
-  tuple_type = TupleType(ts)
-  trap_if(ptr != align_to(ptr, alignment(tuple_type)))
-  trap_if(ptr + elem_size(tuple_type) > len(cx.opts.memory))
-  return list(load(cx, ptr, tuple_type).values())
-
 def lower_flat_values(cx, max_flat, vs, ts, out_param = None):
   cx.inst.may_leave = False
   flat_types = flatten_types(ts)
   if len(flat_types) > max_flat:
-    flat_vals = lower_heap_values(cx, vs, ts, out_param)
+    tuple_type = TupleType(ts)
+    tuple_value = {str(i): v for i,v in enumerate(vs)}
+    if out_param is None:
+      ptr = cx.opts.realloc(0, 0, alignment(tuple_type), elem_size(tuple_type))
+      flat_vals = [ptr]
+    else:
+      ptr = out_param.next('i32')
+      flat_vals = []
+    trap_if(ptr != align_to(ptr, alignment(tuple_type)))
+    trap_if(ptr + elem_size(tuple_type) > len(cx.opts.memory))
+    store(cx, tuple_value, tuple_type, ptr)
   else:
     flat_vals = []
     for i in range(len(vs)):
       flat_vals += lower_flat(cx, vs[i], ts[i])
   cx.inst.may_leave = True
   return flat_vals
 
-def lower_heap_values(cx, vs, ts, out_param):
-  tuple_type = TupleType(ts)
-  tuple_value = {str(i): v for i,v in enumerate(vs)}
-  if out_param is None:
-    ptr = cx.opts.realloc(0, 0, alignment(tuple_type), elem_size(tuple_type))
-    flat_vals = [ptr]
-  else:
-    ptr = out_param.next('i32')
-    flat_vals = []
-  trap_if(ptr != align_to(ptr, alignment(tuple_type)))
-  trap_if(ptr + elem_size(tuple_type) > len(cx.opts.memory))
-  store(cx, tuple_value, tuple_type, ptr)
-  return flat_vals
-
 ### `canon lift`
 
 async def canon_lift(opts, inst, ft, callee, caller, on_start, on_resolve, on_block):
@@ -2169,6 +2168,8 @@ async def canon_future_new(future_t, task):
 
 ### 🔀 `canon {stream,future}.{read,write}`
 
+BLOCKED = 0xffff_ffff
+
 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)
@@ -2196,20 +2197,21 @@ async def copy(EndT, BufferT, event_code, stream_or_future_t, opts, task, i, ptr
   cx = LiftLowerContext(opts, task.inst, borrow_scope = None)
   buffer = BufferT(stream_or_future_t.t, cx, ptr, n)
 
-  def copy_event(why, revoke_buffer):
+  def copy_event(result, 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_copy_result(result, buffer))
 
-  def on_partial_copy(revoke_buffer):
-    e.set_event(partial(copy_event, 'completed', revoke_buffer))
+  def on_copy(revoke_buffer):
+    e.set_event(partial(copy_event, CopyResult.COMPLETED, revoke_buffer))
 
-  def on_copy_done(why):
-    e.set_event(partial(copy_event, why, revoke_buffer = lambda:()))
+  def on_copy_result(result):
+    e.set_event(partial(copy_event, result, revoke_buffer = lambda:()))
 
-  if e.copy(task.inst, buffer, on_partial_copy, on_copy_done) == 'done':
-    return [pack_copy_result(task, e, buffer, 'completed')]
+  result = e.copy(task.inst, buffer, on_copy, on_copy_result)
+  if result is not None:
+    return [pack_copy_result(result, buffer)]
   else:
     e.copying = True
     if opts.sync:
@@ -2220,20 +2222,8 @@ def on_copy_done(why):
     else:
       return [BLOCKED]
 
-BLOCKED   = 0xffff_ffff
-COMPLETED = 0x0
-CLOSED    = 0x1
-CANCELLED = 0x2
-
-def pack_copy_result(task, e, buffer, why):
-  if e.shared.closed():
-    result = CLOSED
-  elif why == 'cancelled':
-    result = CANCELLED
-  else:
-    assert(why == 'completed')
-    assert(not isinstance(e, FutureEnd))
-    result = COMPLETED
+def pack_copy_result(result, buffer):
+  assert(0 <= result < 2**4)
   assert(buffer.progress <= Buffer.MAX_LENGTH < 2**28)
   packed = result | (buffer.progress << 4)
   assert(packed != BLOCKED)