Add OutputStream.getWritableBytes(n)

README.md

@@ -473,6 +473,10 @@ bytes before continuing. The rationale here is that introducing an interruption
 point at every `write` produces less optimal code, but if this is desired you
 can use the `stream.writeAndWait` API.
 
+If you have existing algorithms that output data to an `openArray`, you can use
+the `stream.getWritableBytes` API to continue using them without introducing any
+intermediate buffers.
+
 #### Delayed Writes
 
 Many protocols and formats employ fixed-size and variable-size length prefixes

faststreams/outputs.nim

@@ -136,7 +136,7 @@ template flushImpl(s: OutputStream, awaiter, writeOp, flushOp: untyped) =
       trackWrittenTo(s.buffers, s.span.startAddr)
       awaiter s.vtable.writeOp(s, nil, 0)
       s.span = getWritableSpan(s.buffers)
-      s.spanEndPos += s.span.len - runway
+      s.spanEndPos = s.spanEndPos - runway + s.span.len
 
     if s.vtable.flushOp != nil:
       awaiter s.vtable.flushOp(s)
@@ -223,7 +223,7 @@ proc ensureRunway*(s: OutputStream, neededRunway: Natural) =
     # feeding to it.
     fsAssert s.buffers != nil, "Unsafe memory output of insufficient size"
     s.buffers.ensureRunway(s.span, neededRunway)
-    s.spanEndPos += (s.span.len - runway)
+    s.spanEndPos = s.spanEndPos - runway + s.span.len
 
 when fsAsyncSupport:
   template ensureRunway*(s: AsyncOutputStream, neededRunway: Natural) =
@@ -472,7 +472,55 @@ proc delayVarSizeWrite*(s: OutputStream, maxSize: Natural): VarSizeWriteCursor =
 
     s.span = PageSpan(startAddr: cursorEndAddr,
                       endAddr: nextPageSpan.endAddr)
-    s.spanEndPos += nextPageSize - runway
+    s.spanEndPos = s.spanEndPos - runway + nextPageSize
+
+proc getWritableBytesOnANewPage(s: OutputStream, spanSize: Natural): ptr byte =
+  fsAssert s.buffers != nil
+
+  trackWrittenTo(s.buffers, s.span.startAddr)
+
+  let
+    nextPageSize = nextAlignedSize(spanSize, s.buffers.pageSize)
+    nextPage = allocWritablePage(nextPageSize, spanSize)
+
+  s.buffers.queue.addLast nextPage
+
+  let retiredSpanRunway = s.span.len
+  s.span = nextPage.fullSpan
+  s.spanEndPos = s.spanEndPos - retiredSpanRunway + spanSize
+  s.span.startAddr
+
+template getWritableBytes*(sp: OutputStream, spanSizeParam: Natural): openArray[byte] =
+  ## Returns a contiguous range of memory that the caller is free to populate fully
+  ## or partially. The caller indicates how many bytes were written to the span by
+  ## calling `advance(numberOfBytes)` once or multiple times. Advancing the stream
+  ## past the allocated span size is considered a defect. The typical usage pattern
+  ## of this API looks as follows:
+  ##
+  ## stream.advance(myComponent.writeBlock(stream.getWritableBytes(maxBlockSize)))
+  ##
+  ## In the example, `writeBlock` would be a function returning the number of bytes
+  ## written to the span.
+  ##
+  ## While it's not illegal to issue other writing operations to the stream during
+  ## the `getWritetableSpan` -> `advance` sequence, doing this is not recommended
+  ## because it will result in overwriting the same range of bytes.
+  let
+    s = sp
+    spanSize = spanSizeParam
+    runway = s.span.len
+    startAddr = if spanSize <= runway:
+      s.span.startAddr
+    else:
+      getWritableBytesOnANewPage(s, spanSize)
+
+  makeOpenArray(startAddr, spanSize)
+
+proc advance*(s: OutputStream, bytesWrittenToWritableSpan: Natural) =
+  ## Advance the stream write cursor.
+  ## Typically used after a previous call to `getWritableBytes`.
+  fsAssert bytesWrittenToWritableSpan <= s.span.len
+  s.span.startAddr = offset(s.span.startAddr, bytesWrittenToWritableSpan)
 
 proc finalize*(cursor: var WriteCursor) =
   fsAssert cursor.stream.extCursorsCount > 0
@@ -776,7 +824,7 @@ proc consumeOutputsImpl(s: OutputStream, consumer: OutputConsumingProc) =
     consumer(makeOpenArray(pageReadableStart, pageLen))
 
   s.span = getWritableSpan(s.buffers)
-  s.spanEndPos += s.span.len - runway
+  s.spanEndPos = s.spanEndPos - runway + s.span.len
 
 template consumeOutputs*(s: OutputStream, bytesVar, body: untyped) =
   ## Please note that calling `consumeOutputs` on an unbuffered stream
@@ -821,7 +869,7 @@ proc consumeContiguousOutputImpl(s: OutputStream, consumer: OutputConsumingProc)
   consumer(makeOpenArray(bytesPtr, bytesLen))
 
   s.span = s.buffers.getWritableSpan()
-  s.spanEndPos += s.span.len - runway
+  s.spanEndPos = s.spanEndPos - runway + s.span.len
 
 template consumeContiguousOutput*(s: OutputStream, bytesVar, body: untyped) =
   ## Please note that calling `consumeContiguousOutput` on an unbuffered stream

tests/test_outputs.nim

@@ -213,6 +213,11 @@ suite "output stream":
 
     checkOutputsMatch()
 
+proc writeBlock(data: openArray[byte], output: openArray[byte]): int =
+  doAssert data.len <= output.len
+  copyMem(unsafeAddr output[0], unsafeAddr data[0], data.len)
+  data.len
+
 suite "randomized tests":
   type
     WriteTypes = enum
@@ -357,6 +362,13 @@ suite "randomized tests":
       for i in 0 .. 3000:
         let bytes = randomBytes(rand(9999) + 1)
         referenceBytes.add bytes
+
+        var openArraySize = rand(12000)
+        if openArraySize >= bytes.len:
+          # Make sure that sometimes `writeBlock` populates the entire span
+          if i < 100: openArraySize = bytes.len
+          output.advance writeBlock(bytes, output.getWritableBytes(openArraySize))
+        else:
           output.write bytes
 
       close output