From 97ebcc72c36f8ea81ed8bbfac8e165f7ef74495f Mon Sep 17 00:00:00 2001
From: ruby0x1 <ruby0x1@pm.me>
Date: Thu, 3 Dec 2020 10:30:47 -0800
Subject: [PATCH] Add wrenSetListElement, correctly allow negative indices on
 wrenGetListElement

---
 src/include/wren.h  |  4 ++++
 src/vm/wren_utils.c | 11 +++++++++++
 src/vm/wren_utils.h |  5 +++++
 src/vm/wren_vm.c    | 25 ++++++++++++++++++++++---
 test/api/lists.c    | 38 ++++++++++++++++++++++++++++++++++++++
 test/api/lists.wren |  6 ++++++
 6 files changed, 86 insertions(+), 3 deletions(-)

diff --git a/src/include/wren.h b/src/include/wren.h
index d6e58360e..a7db435c9 100644
--- a/src/include/wren.h
+++ b/src/include/wren.h
@@ -484,6 +484,10 @@ int wrenGetListCount(WrenVM* vm, int slot);
 // [elementSlot].
 void wrenGetListElement(WrenVM* vm, int listSlot, int index, int elementSlot);
 
+// Sets the value stored at [index] in the list at [listSlot], 
+// to the value from [elementSlot]. 
+void wrenSetListElement(WrenVM* vm, int listSlot, int index, int elementSlot);
+
 // Takes the value stored at [elementSlot] and inserts it into the list stored
 // at [listSlot] at [index].
 //
diff --git a/src/vm/wren_utils.c b/src/vm/wren_utils.c
index 0989fefeb..92d56f3f9 100644
--- a/src/vm/wren_utils.c
+++ b/src/vm/wren_utils.c
@@ -194,3 +194,14 @@ int wrenPowerOf2Ceil(int n)
   
   return n;
 }
+
+uint32_t wrenValidateIndex(uint32_t count, int64_t value)
+{
+  // Negative indices count from the end.
+  if (value < 0) value = count + value;
+
+  // Check bounds.
+  if (value >= 0 && value < count) return value;
+
+  return UINT32_MAX;
+}
diff --git a/src/vm/wren_utils.h b/src/vm/wren_utils.h
index b6c5a6d6c..c2a6e4a36 100644
--- a/src/vm/wren_utils.h
+++ b/src/vm/wren_utils.h
@@ -118,4 +118,9 @@ int wrenUtf8DecodeNumBytes(uint8_t byte);
 // Returns the smallest power of two that is equal to or greater than [n].
 int wrenPowerOf2Ceil(int n);
 
+// Validates that [value] is within `[0, count)`. Also allows
+// negative indices which map backwards from the end. Returns the valid positive
+// index value. If invalid, returns `UINT32_MAX`.
+uint32_t wrenValidateIndex(uint32_t count, int64_t value);
+
 #endif
diff --git a/src/vm/wren_vm.c b/src/vm/wren_vm.c
index a1cd3f32f..3f19d45da 100644
--- a/src/vm/wren_vm.c
+++ b/src/vm/wren_vm.c
@@ -1762,9 +1762,27 @@ void wrenGetListElement(WrenVM* vm, int listSlot, int index, int elementSlot)
   validateApiSlot(vm, listSlot);
   validateApiSlot(vm, elementSlot);
   ASSERT(IS_LIST(vm->apiStack[listSlot]), "Slot must hold a list.");
-  
+
   ValueBuffer elements = AS_LIST(vm->apiStack[listSlot])->elements;
-  vm->apiStack[elementSlot] = elements.data[index];
+
+  uint32_t usedIndex = wrenValidateIndex(elements.count, index);
+  ASSERT(usedIndex != UINT32_MAX, "Index out of bounds.");
+
+  vm->apiStack[elementSlot] = elements.data[usedIndex];
+}
+
+void wrenSetListElement(WrenVM* vm, int listSlot, int index, int elementSlot)
+{
+  validateApiSlot(vm, listSlot);
+  validateApiSlot(vm, elementSlot);
+  ASSERT(IS_LIST(vm->apiStack[listSlot]), "Slot must hold a list.");
+
+  ObjList* list = AS_LIST(vm->apiStack[listSlot]);
+
+  uint32_t usedIndex = wrenValidateIndex(list->elements.count, index);
+  ASSERT(usedIndex != UINT32_MAX, "Index out of bounds.");
+  
+  list->elements.data[usedIndex] = vm->apiStack[elementSlot];
 }
 
 void wrenInsertInList(WrenVM* vm, int listSlot, int index, int elementSlot)
@@ -1775,7 +1793,8 @@ void wrenInsertInList(WrenVM* vm, int listSlot, int index, int elementSlot)
   
   ObjList* list = AS_LIST(vm->apiStack[listSlot]);
   
-  // Negative indices count from the end.
+  // Negative indices count from the end. 
+  // We don't use wrenValidateIndex here because insert allows 1 past the end.
   if (index < 0) index = list->elements.count + 1 + index;
   
   ASSERT(index <= list->elements.count, "Index out of bounds.");
diff --git a/test/api/lists.c b/test/api/lists.c
index d763c0cad..d7eee2f0d 100644
--- a/test/api/lists.c
+++ b/test/api/lists.c
@@ -15,6 +15,14 @@ static void insertNumber(WrenVM* vm, int index, double value)
   wrenInsertInList(vm, 0, index, 1);
 }
 
+// Helper function to append a double in a slot then insert it into the list at
+// slot zero.
+static void appendNumber(WrenVM* vm, double value)
+{
+  wrenSetSlotDouble(vm, 1, value);
+  wrenInsertInList(vm, 0, -1, 1);
+}
+
 static void insert(WrenVM* vm)
 {
   wrenSetSlotNewList(vm, 0);
@@ -37,10 +45,40 @@ static void insert(WrenVM* vm)
   insertNumber(vm, -3, 9.0);
 }
 
+static void get(WrenVM* vm)
+{
+  int listSlot = 1;
+  int index = (int)wrenGetSlotDouble(vm, 2);
+
+  wrenGetListElement(vm, listSlot, index, 0);
+}
+
+static void set(WrenVM* vm)
+{
+  wrenSetSlotNewList(vm, 0);
+
+  wrenEnsureSlots(vm, 2);
+
+  appendNumber(vm, 1.0);
+  appendNumber(vm, 2.0);
+  appendNumber(vm, 3.0);
+  appendNumber(vm, 4.0);
+  
+  //list[2] = 33
+  wrenSetSlotDouble(vm, 1, 33);
+  wrenSetListElement(vm, 0, 2, 1);
+
+  //list[-1] = 44
+  wrenSetSlotDouble(vm, 1, 44);
+  wrenSetListElement(vm, 0, -1, 1);
+}
+
 WrenForeignMethodFn listsBindMethod(const char* signature)
 {
   if (strcmp(signature, "static Lists.newList()") == 0) return newList;
   if (strcmp(signature, "static Lists.insert()") == 0) return insert;
+  if (strcmp(signature, "static Lists.set()") == 0) return set;
+  if (strcmp(signature, "static Lists.get(_,_)") == 0) return get;
 
   return NULL;
 }
diff --git a/test/api/lists.wren b/test/api/lists.wren
index 64135a5b7..4ecd16d8f 100644
--- a/test/api/lists.wren
+++ b/test/api/lists.wren
@@ -1,6 +1,8 @@
 class Lists {
   foreign static newList()
   foreign static insert()
+  foreign static set()
+  foreign static get(list, index)
 }
 
 var list = Lists.newList()
@@ -8,3 +10,7 @@ System.print(list is List) // expect: true
 System.print(list.count) // expect: 0
 
 System.print(Lists.insert()) // expect: [4, 5, 6, 1, 2, 3, 9, 8, 7]
+
+System.print(Lists.set()) // expect: [1, 2, 33, 44]
+System.print(Lists.get([1,2,3,4], -2)) // expect: 3
+System.print(Lists.get([1,2,3,4], 1)) // expect: 2