From 92205d608b5c1fa5f5c39b9c53909000c7544cbe Mon Sep 17 00:00:00 2001
From: Zahary Karadjov <zahary@gmail.com>
Date: Wed, 6 Jan 2021 21:20:27 +0200
Subject: [PATCH] Fix #14 (don't use alloca)

---
 faststreams/inputs.nim | 50 ++++++++++++++++++++++++------------------
 tests/test_outputs.nim | 43 ++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 21 deletions(-)

diff --git a/faststreams/inputs.nim b/faststreams/inputs.nim
index b461486..2ceaa7e 100644
--- a/faststreams/inputs.nim
+++ b/faststreams/inputs.nim
@@ -812,21 +812,22 @@ template readInto*(sp: AsyncInputStream, dst: var openarray[byte]): bool =
   let (dstAddr, dstLen) = openArrayToPair(dst)
   readIntoExImpl(s, dstAddr, dstLen, fsAwait, readAsync) == dstLen
 
-when defined(windows):
-  proc alloca(n: int): ptr byte {.importc, header: "<malloc.h>".}
-else:
-  proc alloca(n: int): ptr byte {.importc, header: "<alloca.h>".}
+template useHeapMem(_: Natural) =
+  var buffer: seq[byte]
 
-template allocHeapMem(tmpSeq: var seq[byte], n, _: Natural): ptr byte =
-  tmpSeq.setLen(n)
-  addr tmpSeq[0]
+  template allocMem(n: Natural): ptr byte =
+    buffer.setLen(n)
+    addr buffer[0]
 
-template allocStackMem(tmpSeq: var seq[byte], _, n: Natural): ptr byte =
-  alloca(n)
+template useStackMem(n: static Natural) =
+  var buffer: array[n, byte]
+
+  template allocMem(_: Natural): ptr byte =
+    addr buffer[0]
 
 template readNImpl(sp: InputStream,
                    np: Natural,
-                   allocMem: untyped): openarray[byte] =
+                   createAllocMemOp: untyped): openarray[byte] =
   let
     s = sp
     n = np
@@ -837,28 +838,35 @@ template readNImpl(sp: InputStream,
   # be written in this branch-free linear fashion. The `dataCopy` seq
   # may remain empty in the case where we use stack memory or return
   # an `openarray` from the existing span.
-  var tmpSeq: seq[byte]
   var startAddr: ptr byte
 
-  if n > runway:
-    startAddr = allocMem(tmpSeq, n, np)
-    let drained {.used.} = drainBuffersInto(s, startAddr, n)
-    fsAssert drained == n
-  else:
-    startAddr = s.span.startAddr
-    bumpPointer s.span, n
+  block:
+    # This defines the `allocMem` operation used below.
+    # See `useHeapMem` and `useStackMem` for the possible definitions.
+    # We are creating a block scope in order to allow multiple usages
+    # of `read` within a single scope (avoiding duplicate definitions
+    # of `allocMem`)
+    createAllocMemOp(np)
+
+    if n > runway:
+      startAddr = allocMem(n)
+      let drained {.used.} = drainBuffersInto(s, startAddr, n)
+      fsAssert drained == n
+    else:
+      startAddr = s.span.startAddr
+      bumpPointer s.span, n
 
   makeOpenArray(startAddr, n)
 
 template read*(sp: InputStream, np: static Natural): openarray[byte] =
   const n = np
   when n < maxStackUsage:
-    readNImpl(sp, n, allocStackMem)
+    readNImpl(sp, n, useStackMem)
   else:
-    readNImpl(sp, n, allocHeapMem)
+    readNImpl(sp, n, useHeapMem)
 
 template read*(s: InputStream, n: Natural): openarray[byte] =
-  readNImpl(s, n, allocHeapMem)
+  readNImpl(s, n, useHeapMem)
 
 template read*(s: AsyncInputStream, n: Natural): openarray[byte] =
   read InputStream(s), n
diff --git a/tests/test_outputs.nim b/tests/test_outputs.nim
index ed1d079..82ce17a 100644
--- a/tests/test_outputs.nim
+++ b/tests/test_outputs.nim
@@ -325,3 +325,46 @@ suite "randomized tests":
   randomizedCursorsTest(memoryOutput(), VarSize, 100)
   randomizedCursorsTest(memoryOutput(pageSize = 10), Mixed, 10)
 
+  test "randomized file roundtrip":
+    const randomBytesFileName = "random_bytes_file"
+
+    var
+      referenceBytes, restoredBytes: seq[byte]
+
+    try:
+      let output = fileOutput randomBytesFileName
+
+      for i in 0 .. 3000:
+        let bytes = randomBytes(rand(9999) + 1)
+        referenceBytes.add bytes
+        output.write bytes
+
+      close output
+
+      let input = fileInput randomBytesFileName
+
+      while input.readable(10000):
+        let r = 1 + rand(9999)
+        if r < 5000:
+          restoredBytes.add input.read(8000)
+          restoredBytes.add input.read(500)
+        elif r < 7000:
+          restoredBytes.add input.read(1)
+          restoredBytes.add input.read(2)
+          restoredBytes.add input.read(5)
+          restoredBytes.add input.read(17)
+          restoredBytes.add input.read(128)
+        else:
+          restoredBytes.add input.read(r)
+
+      while input.readable:
+        restoredBytes.add input.read
+
+      close input
+
+      doAssert referenceBytes == restoredBytes
+
+    finally:
+      if fileExists(randomBytesFileName):
+        removeFile randomBytesFileName
+