From 8263789602783aedcf54d6bb01768fb8bddb50a8 Mon Sep 17 00:00:00 2001
From: Mike Beller <bellermike@gmail.com>
Date: Wed, 20 May 2020 13:30:48 -0400
Subject: [PATCH] Fix issue #408 -- make "next" work for typed arrays, and also
 fix bug where tarray/new failed to fully check the type of it's last
 argument.

---
 src/core/typedarray.c | 29 +++++++++++++++++++++++++++--
 test/suite5.janet     | 11 +++++++++++
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/core/typedarray.c b/src/core/typedarray.c
index 961d103e..0e91d891 100644
--- a/src/core/typedarray.c
+++ b/src/core/typedarray.c
@@ -275,6 +275,24 @@ static void ta_setter(void *p, Janet key, Janet value) {
     }
 }
 
+static Janet ta_view_next(void *p, Janet key) {
+    JanetTArrayView *view = p;
+    if (janet_checktype(key, JANET_NIL)) {
+        if (view->size > 0) {
+            return janet_wrap_number(0);
+        } else {
+            return janet_wrap_nil();
+        }
+    }
+    if (!janet_checksize(key)) janet_panic("expected size as key");
+    size_t index = janet_unwrap_number(key);
+    index++;
+    if (index < view->size) {
+        return janet_wrap_number((double) index);
+    }
+    return janet_wrap_nil();
+}
+
 const JanetAbstractType janet_ta_view_type = {
     "ta/view",
     NULL,
@@ -283,7 +301,11 @@ const JanetAbstractType janet_ta_view_type = {
     ta_setter,
     ta_view_marshal,
     ta_view_unmarshal,
-    JANET_ATEND_UNMARSHAL
+    NULL,
+    NULL,
+    NULL,
+    ta_view_next,
+    JANET_ATEND_NEXT
 };
 
 JanetTArrayBuffer *janet_tarray_buffer(size_t size) {
@@ -364,8 +386,11 @@ static Janet cfun_typed_array_new(int32_t argc, Janet *argv) {
             offset = (view->buffer->data - view->as.u8) + offset * ta_type_sizes[view->type];
             stride *= view->stride;
             buffer = view->buffer;
-        } else {
+        } else if (janet_abstract_type(p) == &janet_ta_buffer_type) {
             buffer = p;
+        } else {
+            janet_panicf("bad slot #%d, expected ta/view|ta/buffer, got %v",
+                         4, argv[4]);
         }
     }
     JanetTArrayView *view = janet_tarray_view(type, size, stride, offset, buffer);
diff --git a/test/suite5.janet b/test/suite5.janet
index c4d0f5a8..963f9d91 100644
--- a/test/suite5.janet
+++ b/test/suite5.janet
@@ -58,6 +58,17 @@
 
 (assert (= ((unmarshal (marshal b)) 3) (b 3)) "marshal")
 
+# Issue 408
+(assert-error :invalid-type (tarray/new :int32 10 1 0 (int/u64 7)) "tarray/new should only allow tarray or buffer for last argument")
+(def ta (tarray/new :int32 10))
+(assert (= (next a nil) 0) "tarray next 1")
+(assert (= (next a 0) 1) "tarray next 2")
+(assert (= (next a 8) 9) "tarray next 3")
+(assert (nil? (next a 9)) "tarray next 4")
+(put ta 3 7)
+(put ta 9 7)
+(assert (= 2 (count |(= $ 7) ta)) "tarray count")
+
 # Array remove
 
 (assert (deep= (array/remove @[1 2 3 4 5] 2) @[1 2 4 5]) "array/remove 1")