Merge branch 'master' of https://github.com/keplerproject/luarocks

2 files changed, 99 insertions, 36 deletions

diff --git a/spec/util_spec.lua b/spec/util_spec.lua
index e6776e4b..2779b1ce 100644
--- a/spec/util_spec.lua
+++ b/spec/util_spec.lua
@@ -116,3 +116,59 @@ describe("Basic tests #blackbox #b_util", function()
       end)
    end)
 end)
+
+test_env.unload_luarocks()
+local util = require("luarocks.util")
+
+describe("Luarocks util test #whitebox #w_util", function()
+   describe("util.sortedpairs", function()
+      local function collect(iter, state, var)
+         local collected = {}
+
+         while true do
+            local returns = {iter(state, var)}
+
+            if returns[1] == nil then
+               return collected
+            else
+               table.insert(collected, returns)
+               var = returns[1]
+            end
+         end
+      end
+
+      it("default sort", function()
+         assert.are.same({}, collect(util.sortedpairs({})))
+         assert.are.same({
+            {1, "v1"},
+            {2, "v2"},
+            {3, "v3"},
+            {"bar", "v5"},
+            {"foo", "v4"}
+         }, collect(util.sortedpairs({"v1", "v2", "v3", foo = "v4", bar = "v5"})))
+      end)
+
+      it("sort by function", function()
+         local function compare(a, b) return a > b end
+         assert.are.same({}, collect(util.sortedpairs({}, compare)))
+         assert.are.same({
+            {3, "v3"},
+            {2, "v2"},
+            {1, "v1"}
+         }, collect(util.sortedpairs({"v1", "v2", "v3"}, compare)))
+      end)
+
+      it("sort by priority table", function()
+         assert.are.same({}, collect(util.sortedpairs({}, {"k1", "k2"})))
+         assert.are.same({
+            {"k3", "v3"},
+            {"k2", "v2", {"sub order"}},
+            {"k1", "v1"},
+            {"k4", "v4"},
+            {"k5", "v5"},
+         }, collect(util.sortedpairs({
+            k1 = "v1", k2 = "v2", k3 = "v3", k4 = "v4", k5 = "v5"
+         }, {"k3", {"k2", {"sub order"}}, "k1"})))
+      end)
+   end)
+end)
diff --git a/src/luarocks/util.lua b/src/luarocks/util.lua
index 532bea8b..c9fb7d63 100644
--- a/src/luarocks/util.lua
+++ b/src/luarocks/util.lua
@@ -357,52 +357,59 @@ local function default_sort(a, b)
    end
 end
 
--- The iterator function used internally by util.sortedpairs.
+--- A table iterator generator that returns elements sorted by key,
+-- to be used in "for" loops.
 -- @param tbl table: The table to be iterated.
--- @param sort_function function or nil: An optional comparison function
--- to be used by table.sort when sorting keys.
--- @see sortedpairs
-local function sortedpairs_iterator(tbl, sort_function)
-   local ks = util.keys(tbl)
-   if not sort_function or type(sort_function) == "function" then
-      table.sort(ks, sort_function or default_sort)
-      for _, k in ipairs(ks) do
-         coroutine.yield(k, tbl[k])
-      end
+-- @param sort_function function or table or nil: An optional comparison function
+-- to be used by table.sort when sorting keys, or an array listing an explicit order
+-- for keys. If a value itself is an array, it is taken so that the first element
+-- is a string representing the field name, and the second element is a priority table
+-- for that key, which is returned by the iterator as the third value after the key
+-- and the value.
+-- @return function: the iterator function.
+function util.sortedpairs(tbl, sort_function)
+   sort_function = sort_function or default_sort
+   local keys = util.keys(tbl)
+   local sub_orders = {}
+
+   if type(sort_function) == "function" then
+      table.sort(keys, sort_function)
    else
       local order = sort_function
-      local done = {}
-      for _, k in ipairs(order) do
-         local sub_order
-         if type(k) == "table" then
-            sub_order = k[2]
-            k = k[1]
+      local ordered_keys = {}
+      local all_keys = keys
+      keys = {}
+
+      for _, order_entry in ipairs(order) do
+         local key, sub_order
+         if type(order_entry) == "table" then
+            key = order_entry[1]
+            sub_order = order_entry[2]
+         else
+            key = order_entry
          end
-         if tbl[k] then
-            done[k] = true
-            coroutine.yield(k, tbl[k], sub_order)
+
+         if tbl[key] then
+            ordered_keys[key] = true
+            sub_orders[key] = sub_order
+            table.insert(keys, key)
          end
       end
-      table.sort(ks, default_sort)
-      for _, k in ipairs(ks) do
-         if not done[k] then
-            coroutine.yield(k, tbl[k])
+
+      table.sort(all_keys, default_sort)
+      for _, key in ipairs(all_keys) do
+         if not ordered_keys[key] then
+            table.insert(keys, key)
          end
       end
    end
-end
 
---- A table iterator generator that returns elements sorted by key,
--- to be used in "for" loops.
--- @param tbl table: The table to be iterated.
--- @param sort_function function or table or nil: An optional comparison function
--- to be used by table.sort when sorting keys, or an array listing an explicit order
--- for keys. If a value itself is an array, it is taken so that the first element
--- is a string representing the field name, and the second element is a priority table
--- for that key.
--- @return function: the iterator function.
-function util.sortedpairs(tbl, sort_function)
-   return coroutine.wrap(function() sortedpairs_iterator(tbl, sort_function) end)
+   local i = 1
+   return function()
+      local key = keys[i]
+      i = i + 1
+      return key, tbl[key], sub_orders[key]
+   end
 end
 
 function util.lua_versions()