aboutsummaryrefslogtreecommitdiffstats
path: root/util/array.lua
blob: c6352de58476f5f5023bb9f36ac8297fc622690e (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
-- Prosody IM
-- Copyright (C) 2008-2010 Matthew Wild
-- Copyright (C) 2008-2010 Waqas Hussain
--
-- This project is MIT/X11 licensed. Please see the
-- COPYING file in the source package for more information.
--

local t_insert, t_sort, t_remove, t_concat
    = table.insert, table.sort, table.remove, table.concat;

local setmetatable = setmetatable;
local math_random = math.random;
local math_floor = math.floor;
local pairs, ipairs = pairs, ipairs;
local tostring = tostring;
local type = type;

local array = {};
local array_base = {};
local array_methods = {};
local array_mt = {
	__index = array_methods;
	__name = "array";
	__tostring = function (self) return "{"..self:concat(", ").."}"; end;
};

function array_mt:__freeze() return self; end

local function new_array(self, t, _s, _var)
	if type(t) == "function" then -- Assume iterator
		t = self.collect(t, _s, _var);
	end
	return setmetatable(t or {}, array_mt);
end

function array_mt.__add(a1, a2)
	local res = new_array();
	return res:append(a1):append(a2);
end

function array_mt.__eq(a, b)
	if #a == #b then
		for i = 1, #a do
			if a[i] ~= b[i] then
				return false;
			end
		end
	else
		return false;
	end
	return true;
end

setmetatable(array, { __call = new_array });

-- Read-only methods
function array_methods:random()
	return self[math_random(1, #self)];
end

-- These methods can be called two ways:
--   array.method(existing_array, [params [, ...]]) -- Create new array for result
--   existing_array:method([params, ...]) -- Transform existing array into result
--
function array_base.map(outa, ina, func)
	for k, v in ipairs(ina) do
		outa[k] = func(v);
	end
	return outa;
end

function array_base.filter(outa, ina, func)
	local inplace, start_length = ina == outa, #ina;
	local write = 1;
	for read = 1, start_length do
		local v = ina[read];
		if func(v) then
			outa[write] = v;
			write = write + 1;
		end
	end

	if inplace and write <= start_length then
		for i = write, start_length do
			outa[i] = nil;
		end
	end

	return outa;
end

function array_base.sort(outa, ina, ...)
	if ina ~= outa then
		outa:append(ina);
	end
	t_sort(outa, ...);
	return outa;
end

function array_base.unique(outa, ina)
	local seen = {};
	return array_base.filter(outa, ina, function (item)
		if seen[item] then
			return false;
		else
			seen[item] = true;
			return true;
		end
	end);
end

function array_base.pluck(outa, ina, key)
	for i = 1, #ina do
		outa[i] = ina[i][key];
	end
	return outa;
end

function array_base.reverse(outa, ina)
	local len = #ina;
	if ina == outa then
		local middle = math_floor(len/2);
		len = len + 1;
		local o; -- opposite
		for i = 1, middle do
			o = len - i;
			outa[i], outa[o] = outa[o], outa[i];
		end
	else
		local off = len + 1;
		for i = 1, len do
			outa[i] = ina[off - i];
		end
	end
	return outa;
end

--- These methods only mutate the array
function array_methods:shuffle()
	local len = #self;
	for i = 1, #self do
		local r = math_random(i, len);
		self[i], self[r] = self[r], self[i];
	end
	return self;
end

function array_methods:append(ina)
	local len, len2 = #self, #ina;
	for i = 1, len2 do
		self[len+i] = ina[i];
	end
	return self;
end

function array_methods:push(x)
	t_insert(self, x);
	return self;
end

array_methods.pop = t_remove;

function array_methods:concat(sep)
	return t_concat(array.map(self, tostring), sep);
end

function array_methods:length()
	return #self;
end

--- These methods always create a new array
function array.collect(f, s, var)
	local t = {};
	while true do
		var = f(s, var);
		if var == nil then break; end
		t_insert(t, var);
	end
	return setmetatable(t, array_mt);
end

---

-- Setup methods from array_base
for method, f in pairs(array_base) do
	local base_method = f;
	-- Setup global array method which makes new array
	array[method] = function (old_a, ...)
		local a = new_array();
		return base_method(a, old_a, ...);
	end
	-- Setup per-array (mutating) method
	array_methods[method] = function (self, ...)
		return base_method(self, self, ...);
	end
end

return array;