lua-subtree: Blueprint Add a Lua built-in function to return a size of subtree

(https://blueprints.launchpad.net/tarantool/+spec/lua-builtin-size-of-subtree)

mod/box/box.lua

@@ -253,12 +253,19 @@ function box.on_reload_configuration()
     index_mt.max = function(index) return index.idx:max() end
     -- iteration
     index_mt.pairs = function(index)
-        return index.idx.next, index.idx, nil end
+        return index.idx.next, index.idx, nil
+    end
     --
     index_mt.next = function(index, ...)
-        return index.idx:next(...) end
+        return index.idx:next(...)
+    end
     index_mt.prev = function(index, ...)
-        return index.idx:prev(...) end
+        return index.idx:prev(...)
+    end
+    -- index subtree size
+    index_mt.count = function(index, ...)
+        return index.idx:count(...)
+    end
     --
     index_mt.select_range = function(index, limit, ...)
         local range = {}

mod/box/box_lua.m

@@ -562,6 +562,50 @@ lbox_index_prev_equal(struct lua_State *L)
 	return tuple ? 2 : 1;
 }
 
+/**
+ * Lua index subtree count function.
+ * Function iterates over an index, counting
+ * number of tuples which equals key search criteria.
+ * Key can be multi-parted.
+ */
+static int
+lbox_index_count(struct lua_State *L)
+{
+	Index *index = lua_checkindex(L, 1);
+	int argc = lua_gettop(L) - 1;
+	/* preparing single or multi-part key */
+	void *key;
+	int key_part_count;
+	if (argc == 1 && lua_type(L, 2) == LUA_TUSERDATA) {
+		/* Searching by tuple. */
+		struct tuple *tuple = lua_checktuple(L, 2);
+		key = tuple->data;
+		key_part_count = tuple->field_count;
+	} else {
+		/* Single or multi- part key. */
+		key_part_count = argc;
+		struct tbuf *data = tbuf_alloc(fiber->gc_pool);
+		for (int i = 0; i < argc; ++i)
+			append_key_part(L, i + 2, data,
+					index->key_def->parts[i].type);
+		key = data->data;
+	}
+	u32 count = 0;
+	/* preparing index iterator */
+	struct iterator *it = index->position;
+	[index initIteratorByKey: it :ITER_FORWARD :key :key_part_count];
+	/* iterating over the index and counting tuples */
+	struct tuple *tuple;
+	while ((tuple = it->next_equal(it)) != NULL) {
+		if (tuple->flags & GHOST)
+			continue;
+		count++;
+	}
+	/* returning subtree size */
+	lua_pushnumber(L, count);
+	return 1;
+}
+
 static const struct luaL_reg lbox_index_meta[] = {
 	{"__tostring", lbox_index_tostring},
 	{"__len", lbox_index_len},
@@ -572,6 +616,7 @@ static const struct luaL_reg lbox_index_meta[] = {
 	{"prev", lbox_index_prev},
 	{"next_equal", lbox_index_next_equal},
 	{"prev_equal", lbox_index_prev_equal},
+	{"count", lbox_index_count},
 	{NULL, NULL}
 };
 

test/box_big/lua.result

@@ -258,3 +258,55 @@ lua for k, v in box.space[16].index[1].idx.prev_equal, box.space[16].index[1].id
 'pid_5': {'sid_2', 'tid_995'}
 'pid_6': {'sid_2', 'tid_994'}
 ...
+lua box.insert(17, 1, 1, 1)
+---
+ - 1: {1, 1}
+...
+lua box.insert(17, 2, 2, 0)
+---
+ - 2: {2, 0}
+...
+lua box.insert(17, 3, 2, 1)
+---
+ - 3: {2, 1}
+...
+lua box.insert(17, 4, 3, 0)
+---
+ - 4: {3, 0}
+...
+lua box.insert(17, 5, 3, 1)
+---
+ - 5: {3, 1}
+...
+lua box.insert(17, 6, 3, 2)
+---
+ - 6: {3, 2}
+...
+lua box.space[17].index[1].idx:count(1)
+---
+ - 1
+...
+lua box.space[17].index[1].idx:count(2)
+---
+ - 2
+...
+lua box.space[17].index[1].idx:count(2, 1)
+---
+ - 1
+...
+lua box.space[17].index[1].idx:count(2, 2)
+---
+ - 0
+...
+lua box.space[17].index[1].idx:count(3)
+---
+ - 3
+...
+lua box.space[17].index[1].idx:count(3, 3)
+---
+ - 0
+...
+lua box.space[17].index[1].idx:count()
+---
+ - 6
+...

test/box_big/lua.test

@@ -89,3 +89,20 @@ exec admin "lua for k, v in box.space[16].index[1].idx.next_equal, box.space[16]
 exec admin "lua for k, v in box.space[16].index[1].idx.prev_equal, box.space[16].index[1].idx, 'sid_1' do print(v) end"
 exec admin "lua for k, v in box.space[16].index[1].idx.next_equal, box.space[16].index[1].idx, 'sid_2' do print(v) end"
 exec admin "lua for k, v in box.space[16].index[1].idx.prev_equal, box.space[16].index[1].idx, 'sid_2' do print(v) end"
+
+#
+# Tests for lua idx:count()
+#
+exec admin "lua box.insert(17, 1, 1, 1)"
+exec admin "lua box.insert(17, 2, 2, 0)"
+exec admin "lua box.insert(17, 3, 2, 1)"
+exec admin "lua box.insert(17, 4, 3, 0)"
+exec admin "lua box.insert(17, 5, 3, 1)"
+exec admin "lua box.insert(17, 6, 3, 2)"
+exec admin "lua box.space[17].index[1].idx:count(1)"
+exec admin "lua box.space[17].index[1].idx:count(2)"
+exec admin "lua box.space[17].index[1].idx:count(2, 1)"
+exec admin "lua box.space[17].index[1].idx:count(2, 2)"
+exec admin "lua box.space[17].index[1].idx:count(3)"
+exec admin "lua box.space[17].index[1].idx:count(3, 3)"
+exec admin "lua box.space[17].index[1].idx:count()"

test/box_big/tarantool.cfg

@@ -212,3 +212,17 @@ space[16].index[1].key_field[0].fieldno = 1
 space[16].index[1].key_field[0].type = "STR"
 space[16].index[1].key_field[1].fieldno = 2
 space[16].index[1].key_field[1].type = "STR"
+
+# lua index.idx:count() testing
+# https://blueprints.launchpad.net/tarantool/+spec/lua-builtin-size-of-subtree
+space[17].enabled = true
+space[17].index[0].type = "HASH"
+space[17].index[0].unique = 1
+space[17].index[0].key_field[0].fieldno = 0
+space[17].index[0].key_field[0].type = "NUM"
+space[17].index[1].type = "TREE"
+space[17].index[1].unique = 0
+space[17].index[1].key_field[0].fieldno = 1
+space[17].index[1].key_field[0].type = "NUM"
+space[17].index[1].key_field[1].fieldno = 2
+space[17].index[1].key_field[1].type = "NUM"