*

2 files changed, 42 insertions, 16 deletions

diff --git a/lua/chupacabra/chupacabra.lua b/lua/chupacabra/chupacabra.lua
index 645d327..e4211e5 100644
--- a/lua/chupacabra/chupacabra.lua
+++ b/lua/chupacabra/chupacabra.lua
@@ -1,12 +1,21 @@
 local chupacabra = {}
 
--- TODO: add the ability to push and pop arrays from the stack
-
 -- Parse a string into a list of tokens
-function chupacabra.tokenize(str)
+function chupacabra.tokenize(s)
     local tokens = {}
-    for token in string.gmatch(str, "%S+") do
-        table.insert(tokens, token)
+    local array = nil
+    for token in s:gmatch("%S+") do
+        if array then
+            array = array .. " " .. token
+            if token:find("]$") then
+                table.insert(tokens, array)
+                array = nil
+            end
+        elseif token:find("^%[") then
+            array = token
+        else
+            table.insert(tokens, token)
+        end
     end
     return tokens
 end
@@ -17,6 +26,12 @@ function chupacabra.evaluate(tokens, context)
     for _, token in ipairs(tokens) do
         if tonumber(token) then
             table.insert(stack, tonumber(token))
+        elseif token:match("^%b[]$") then -- If the token is an array
+            local array = {}
+            for number in token:sub(2, -2):gmatch("%S+") do
+                table.insert(array, tonumber(number))
+            end
+            table.insert(stack, array)
         elseif token == "pop" then
             table.remove(stack) -- testing is fun
         elseif token == "+" then
@@ -35,29 +50,25 @@ function chupacabra.evaluate(tokens, context)
             local b = table.remove(stack)
             local a = table.remove(stack)
             table.insert(stack, a / b)
-        elseif token == "map+" then
-            -- If the token is "map+", pop an array from the stack, add 1 to each element, and push the result
+        elseif token == "map+" then -- FIXME: fix all map functions to either add 2 arrays or spread
             local a = table.remove(stack)
             for i, v in ipairs(a) do
                 a[i] = v + 1
             end
             table.insert(stack, a)
         elseif token == "map*" then
-            -- If the token is "map*", pop an array from the stack, multiply each element by 2, and push the result
             local a = table.remove(stack)
             for i, v in ipairs(a) do
                 a[i] = v * 2
             end
             table.insert(stack, a)
         elseif token == "map-" then
-            -- If the token is "map-", pop an array from the stack, subtract 1 from each element, and push the result
             local a = table.remove(stack)
             for i, v in ipairs(a) do
                 a[i] = v - 1
             end
             table.insert(stack, a)
         elseif token == "map/" then
-            -- If the token is "map/", pop an array from the stack, divide each element by 2, and push the result
             local a = table.remove(stack)
             for i, v in ipairs(a) do
                 a[i] = v / 2
diff --git a/lua/chupacabra/test_chupacabra.lua b/lua/chupacabra/test_chupacabra.lua
index 4d75f70..1eb5a48 100644
--- a/lua/chupacabra/test_chupacabra.lua
+++ b/lua/chupacabra/test_chupacabra.lua
@@ -1,22 +1,37 @@
 local chupacabra = require("chupacabra") 
 
--- Define a function to run a test case
+-- convert a table to a string
+local function table_to_string(t)
+    local str = "{"
+    for i, v in ipairs(t) do
+        if i > 1 then
+            str = str .. ", "
+        end
+        str = str .. tostring(v)
+    end
+    str = str .. "}"
+    return str
+end
+
 local function test_case(input, expected_output)
     local output = chupacabra.run(input, {})
-    assert(output == expected_output, "Test failed: " .. input .. " => " .. output .. ", expected: " .. expected_output)
-    print("Test passed: " .. input .. " => " .. output)
+    local expected_output_str = type(expected_output) == "table" and table_to_string(expected_output) or tostring(expected_output)
+    local output_str = type(output) == "table" and table_to_string(output) or tostring(output)
+    assert(output_str == expected_output_str, "Test failed: " .. input .. " => " .. output_str .. ", expected: " .. expected_output_str)
+    print("Test passed: " .. input .. " => " .. output_str)
 end
 
 
--- Run some test cases
 test_case("1", 1)  -- 1
 test_case("2 1 pop", 2) -- 2
+test_case("[1 1]", {1, 1})
+test_case("[1 2 3] map+", {2, 3, 4})
 test_case("3 4 +", 7)  -- 3 + 4 = 7
 test_case("5 2 -", 3)  -- 5 - 2 = 3
 test_case("2 3 *", 6)  -- 2 * 3 = 6
-test_case("8 2 /", 4)  -- 8 / 2 = 4
+test_case("8 2 /", 4.0)  -- 8 / 2 = 4
 test_case("2 3 4 + *", 14)  -- 2 * (3 + 4) = 14
 test_case("5 2 3 + 4 + +", 14)  -- 5 + 2 + 3 + 4 = 14
 test_case("7 2 + 3 4 + +", 16)  -- 7 + 2 + 3 + 4 = 16
 test_case("5 2 - 3 4 * +", 15)  -- (5 - 2) + (3 * 4) = 15
-test_case("8 2 / 3 4 * +", 16)  -- (8 / 2) + (3 * 4) = 16
\ No newline at end of file
+test_case("8 2 / 3 4 * +", 16.0)  -- (8 / 2) + (3 * 4) = 16
\ No newline at end of file