You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

fun.lua 29KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076
  1. ---
  2. --- Lua Fun - a high-performance functional programming library for LuaJIT
  3. ---
  4. --- Copyright (c) 2013-2017 Roman Tsisyk <roman@tsisyk.com>
  5. ---
  6. --- Distributed under the MIT/X11 License. See COPYING.md for more details.
  7. ---
  8. local exports = {}
  9. local methods = {}
  10. -- Rspamd specific
  11. local rspamd_text = require "rspamd_text"
  12. local text_cookie = rspamd_text.cookie
  13. -- compatibility with Lua 5.1/5.2
  14. local unpack = rawget(table, "unpack") or unpack
  15. --------------------------------------------------------------------------------
  16. -- Tools
  17. --------------------------------------------------------------------------------
  18. local return_if_not_empty = function(state_x, ...)
  19. if state_x == nil then
  20. return nil
  21. end
  22. return ...
  23. end
  24. local call_if_not_empty = function(fun, state_x, ...)
  25. if state_x == nil then
  26. return nil
  27. end
  28. return state_x, fun(...)
  29. end
  30. local function deepcopy(orig) -- used by cycle()
  31. local orig_type = type(orig)
  32. local copy
  33. if orig_type == 'table' then
  34. copy = {}
  35. for orig_key, orig_value in next, orig, nil do
  36. copy[deepcopy(orig_key)] = deepcopy(orig_value)
  37. end
  38. else
  39. copy = orig
  40. end
  41. return copy
  42. end
  43. local iterator_mt = {
  44. -- usually called by for-in loop
  45. __call = function(self, param, state)
  46. return self.gen(param, state)
  47. end;
  48. __tostring = function(self)
  49. return '<generator>'
  50. end;
  51. -- add all exported methods
  52. __index = methods;
  53. }
  54. local wrap = function(gen, param, state)
  55. return setmetatable({
  56. gen = gen,
  57. param = param,
  58. state = state
  59. }, iterator_mt), param, state
  60. end
  61. exports.wrap = wrap
  62. local unwrap = function(self)
  63. return self.gen, self.param, self.state
  64. end
  65. methods.unwrap = unwrap
  66. --------------------------------------------------------------------------------
  67. -- Basic Functions
  68. --------------------------------------------------------------------------------
  69. local nil_gen = function(_param, _state)
  70. return nil
  71. end
  72. local string_gen = function(param, state)
  73. local state = state + 1
  74. if state > #param then
  75. return nil
  76. end
  77. local r = string.sub(param, state, state)
  78. return state, r
  79. end
  80. local text_gen = function(param, state)
  81. local state = state + 1
  82. if state > #param then
  83. return nil
  84. end
  85. local r = string.char(param:byte(state))
  86. return state, r
  87. end
  88. local ipairs_gen = ipairs({}) -- get the generating function from ipairs
  89. local pairs_gen = pairs({ a = 0 }) -- get the generating function from pairs
  90. local map_gen = function(tab, key)
  91. local value
  92. local key, value = pairs_gen(tab, key)
  93. return key, key, value
  94. end
  95. local rawiter = function(obj, param, state)
  96. assert(obj ~= nil, "invalid iterator")
  97. if type(obj) == "table" then
  98. local mt = getmetatable(obj);
  99. if mt ~= nil then
  100. if mt == iterator_mt then
  101. return obj.gen, obj.param, obj.state
  102. elseif mt.__ipairs ~= nil then
  103. return mt.__ipairs(obj)
  104. elseif mt.__pairs ~= nil then
  105. return mt.__pairs(obj)
  106. end
  107. end
  108. if #obj > 0 then
  109. -- array
  110. return ipairs(obj)
  111. else
  112. -- hash
  113. return map_gen, obj, nil
  114. end
  115. elseif (type(obj) == "function") then
  116. return obj, param, state
  117. elseif (type(obj) == "string") then
  118. if #obj == 0 then
  119. return nil_gen, nil, nil
  120. end
  121. return string_gen, obj, 0
  122. elseif (type(obj) == "userdata" and obj.cookie == text_cookie) then
  123. if #obj == 0 then
  124. return nil_gen, nil, nil
  125. end
  126. return text_gen, obj, 0
  127. end
  128. error(string.format('object %s of type "%s" is not iterable',
  129. obj, type(obj)))
  130. end
  131. local iter = function(obj, param, state)
  132. return wrap(rawiter(obj, param, state))
  133. end
  134. exports.iter = iter
  135. local method0 = function(fun)
  136. return function(self)
  137. return fun(self.gen, self.param, self.state)
  138. end
  139. end
  140. local method1 = function(fun)
  141. return function(self, arg1)
  142. return fun(arg1, self.gen, self.param, self.state)
  143. end
  144. end
  145. local method2 = function(fun)
  146. return function(self, arg1, arg2)
  147. return fun(arg1, arg2, self.gen, self.param, self.state)
  148. end
  149. end
  150. local export0 = function(fun)
  151. return function(gen, param, state)
  152. return fun(rawiter(gen, param, state))
  153. end
  154. end
  155. local export1 = function(fun)
  156. return function(arg1, gen, param, state)
  157. return fun(arg1, rawiter(gen, param, state))
  158. end
  159. end
  160. local export2 = function(fun)
  161. return function(arg1, arg2, gen, param, state)
  162. return fun(arg1, arg2, rawiter(gen, param, state))
  163. end
  164. end
  165. local each = function(fun, gen, param, state)
  166. repeat
  167. state = call_if_not_empty(fun, gen(param, state))
  168. until state == nil
  169. end
  170. methods.each = method1(each)
  171. exports.each = export1(each)
  172. methods.for_each = methods.each
  173. exports.for_each = exports.each
  174. methods.foreach = methods.each
  175. exports.foreach = exports.each
  176. --------------------------------------------------------------------------------
  177. -- Generators
  178. --------------------------------------------------------------------------------
  179. local range_gen = function(param, state)
  180. local stop, step = param[1], param[2]
  181. local state = state + step
  182. if state > stop then
  183. return nil
  184. end
  185. return state, state
  186. end
  187. local range_rev_gen = function(param, state)
  188. local stop, step = param[1], param[2]
  189. local state = state + step
  190. if state < stop then
  191. return nil
  192. end
  193. return state, state
  194. end
  195. local range = function(start, stop, step)
  196. if step == nil then
  197. if stop == nil then
  198. if start == 0 then
  199. return nil_gen, nil, nil
  200. end
  201. stop = start
  202. start = stop > 0 and 1 or -1
  203. end
  204. step = start <= stop and 1 or -1
  205. end
  206. assert(type(start) == "number", "start must be a number")
  207. assert(type(stop) == "number", "stop must be a number")
  208. assert(type(step) == "number", "step must be a number")
  209. assert(step ~= 0, "step must not be zero")
  210. if (step > 0) then
  211. return wrap(range_gen, {stop, step}, start - step)
  212. elseif (step < 0) then
  213. return wrap(range_rev_gen, {stop, step}, start - step)
  214. end
  215. end
  216. exports.range = range
  217. local duplicate_table_gen = function(param_x, state_x)
  218. return state_x + 1, unpack(param_x)
  219. end
  220. local duplicate_fun_gen = function(param_x, state_x)
  221. return state_x + 1, param_x(state_x)
  222. end
  223. local duplicate_gen = function(param_x, state_x)
  224. return state_x + 1, param_x
  225. end
  226. local duplicate = function(...)
  227. if select('#', ...) <= 1 then
  228. return wrap(duplicate_gen, select(1, ...), 0)
  229. else
  230. return wrap(duplicate_table_gen, {...}, 0)
  231. end
  232. end
  233. exports.duplicate = duplicate
  234. exports.replicate = duplicate
  235. exports.xrepeat = duplicate
  236. local tabulate = function(fun)
  237. assert(type(fun) == "function")
  238. return wrap(duplicate_fun_gen, fun, 0)
  239. end
  240. exports.tabulate = tabulate
  241. local zeros = function()
  242. return wrap(duplicate_gen, 0, 0)
  243. end
  244. exports.zeros = zeros
  245. local ones = function()
  246. return wrap(duplicate_gen, 1, 0)
  247. end
  248. exports.ones = ones
  249. local rands_gen = function(param_x, _state_x)
  250. return 0, math.random(param_x[1], param_x[2])
  251. end
  252. local rands_nil_gen = function(_param_x, _state_x)
  253. return 0, math.random()
  254. end
  255. local rands = function(n, m)
  256. if n == nil and m == nil then
  257. return wrap(rands_nil_gen, 0, 0)
  258. end
  259. assert(type(n) == "number", "invalid first arg to rands")
  260. if m == nil then
  261. m = n
  262. n = 0
  263. else
  264. assert(type(m) == "number", "invalid second arg to rands")
  265. end
  266. assert(n < m, "empty interval")
  267. return wrap(rands_gen, {n, m - 1}, 0)
  268. end
  269. exports.rands = rands
  270. --------------------------------------------------------------------------------
  271. -- Slicing
  272. --------------------------------------------------------------------------------
  273. local nth = function(n, gen_x, param_x, state_x)
  274. assert(n > 0, "invalid first argument to nth")
  275. -- An optimization for arrays and strings
  276. if gen_x == ipairs_gen then
  277. return param_x[n]
  278. elseif gen_x == string_gen then
  279. if n <= #param_x then
  280. return string.sub(param_x, n, n)
  281. else
  282. return nil
  283. end
  284. end
  285. for i=1,n-1,1 do
  286. state_x = gen_x(param_x, state_x)
  287. if state_x == nil then
  288. return nil
  289. end
  290. end
  291. return return_if_not_empty(gen_x(param_x, state_x))
  292. end
  293. methods.nth = method1(nth)
  294. exports.nth = export1(nth)
  295. local head_call = function(state, ...)
  296. if state == nil then
  297. error("head: iterator is empty")
  298. end
  299. return ...
  300. end
  301. local head = function(gen, param, state)
  302. return head_call(gen(param, state))
  303. end
  304. methods.head = method0(head)
  305. exports.head = export0(head)
  306. exports.car = exports.head
  307. methods.car = methods.head
  308. local tail = function(gen, param, state)
  309. state = gen(param, state)
  310. if state == nil then
  311. return wrap(nil_gen, nil, nil)
  312. end
  313. return wrap(gen, param, state)
  314. end
  315. methods.tail = method0(tail)
  316. exports.tail = export0(tail)
  317. exports.cdr = exports.tail
  318. methods.cdr = methods.tail
  319. local take_n_gen_x = function(i, state_x, ...)
  320. if state_x == nil then
  321. return nil
  322. end
  323. return {i, state_x}, ...
  324. end
  325. local take_n_gen = function(param, state)
  326. local n, gen_x, param_x = param[1], param[2], param[3]
  327. local i, state_x = state[1], state[2]
  328. if i >= n then
  329. return nil
  330. end
  331. return take_n_gen_x(i + 1, gen_x(param_x, state_x))
  332. end
  333. local take_n = function(n, gen, param, state)
  334. assert(n >= 0, "invalid first argument to take_n")
  335. return wrap(take_n_gen, {n, gen, param}, {0, state})
  336. end
  337. methods.take_n = method1(take_n)
  338. exports.take_n = export1(take_n)
  339. local take_while_gen_x = function(fun, state_x, ...)
  340. if state_x == nil or not fun(...) then
  341. return nil
  342. end
  343. return state_x, ...
  344. end
  345. local take_while_gen = function(param, state_x)
  346. local fun, gen_x, param_x = param[1], param[2], param[3]
  347. return take_while_gen_x(fun, gen_x(param_x, state_x))
  348. end
  349. local take_while = function(fun, gen, param, state)
  350. assert(type(fun) == "function", "invalid first argument to take_while")
  351. return wrap(take_while_gen, {fun, gen, param}, state)
  352. end
  353. methods.take_while = method1(take_while)
  354. exports.take_while = export1(take_while)
  355. local take = function(n_or_fun, gen, param, state)
  356. if type(n_or_fun) == "number" then
  357. return take_n(n_or_fun, gen, param, state)
  358. else
  359. return take_while(n_or_fun, gen, param, state)
  360. end
  361. end
  362. methods.take = method1(take)
  363. exports.take = export1(take)
  364. local drop_n = function(n, gen, param, state)
  365. assert(n >= 0, "invalid first argument to drop_n")
  366. local i
  367. for i=1,n,1 do
  368. state = gen(param, state)
  369. if state == nil then
  370. return wrap(nil_gen, nil, nil)
  371. end
  372. end
  373. return wrap(gen, param, state)
  374. end
  375. methods.drop_n = method1(drop_n)
  376. exports.drop_n = export1(drop_n)
  377. local drop_while_x = function(fun, state_x, ...)
  378. if state_x == nil or not fun(...) then
  379. return state_x, false
  380. end
  381. return state_x, true, ...
  382. end
  383. local drop_while = function(fun, gen_x, param_x, state_x)
  384. assert(type(fun) == "function", "invalid first argument to drop_while")
  385. local cont, state_x_prev
  386. repeat
  387. state_x_prev = deepcopy(state_x)
  388. state_x, cont = drop_while_x(fun, gen_x(param_x, state_x))
  389. until not cont
  390. if state_x == nil then
  391. return wrap(nil_gen, nil, nil)
  392. end
  393. return wrap(gen_x, param_x, state_x_prev)
  394. end
  395. methods.drop_while = method1(drop_while)
  396. exports.drop_while = export1(drop_while)
  397. local drop = function(n_or_fun, gen_x, param_x, state_x)
  398. if type(n_or_fun) == "number" then
  399. return drop_n(n_or_fun, gen_x, param_x, state_x)
  400. else
  401. return drop_while(n_or_fun, gen_x, param_x, state_x)
  402. end
  403. end
  404. methods.drop = method1(drop)
  405. exports.drop = export1(drop)
  406. local split = function(n_or_fun, gen_x, param_x, state_x)
  407. return take(n_or_fun, gen_x, param_x, state_x),
  408. drop(n_or_fun, gen_x, param_x, state_x)
  409. end
  410. methods.split = method1(split)
  411. exports.split = export1(split)
  412. methods.split_at = methods.split
  413. exports.split_at = exports.split
  414. methods.span = methods.split
  415. exports.span = exports.split
  416. --------------------------------------------------------------------------------
  417. -- Indexing
  418. --------------------------------------------------------------------------------
  419. local index = function(x, gen, param, state)
  420. local i = 1
  421. for _k, r in gen, param, state do
  422. if r == x then
  423. return i
  424. end
  425. i = i + 1
  426. end
  427. return nil
  428. end
  429. methods.index = method1(index)
  430. exports.index = export1(index)
  431. methods.index_of = methods.index
  432. exports.index_of = exports.index
  433. methods.elem_index = methods.index
  434. exports.elem_index = exports.index
  435. local indexes_gen = function(param, state)
  436. local x, gen_x, param_x = param[1], param[2], param[3]
  437. local i, state_x = state[1], state[2]
  438. local r
  439. while true do
  440. state_x, r = gen_x(param_x, state_x)
  441. if state_x == nil then
  442. return nil
  443. end
  444. i = i + 1
  445. if r == x then
  446. return {i, state_x}, i
  447. end
  448. end
  449. end
  450. local indexes = function(x, gen, param, state)
  451. return wrap(indexes_gen, {x, gen, param}, {0, state})
  452. end
  453. methods.indexes = method1(indexes)
  454. exports.indexes = export1(indexes)
  455. methods.elem_indexes = methods.indexes
  456. exports.elem_indexes = exports.indexes
  457. methods.indices = methods.indexes
  458. exports.indices = exports.indexes
  459. methods.elem_indices = methods.indexes
  460. exports.elem_indices = exports.indexes
  461. --------------------------------------------------------------------------------
  462. -- Filtering
  463. --------------------------------------------------------------------------------
  464. local filter1_gen = function(fun, gen_x, param_x, state_x, a)
  465. while true do
  466. if state_x == nil or fun(a) then break; end
  467. state_x, a = gen_x(param_x, state_x)
  468. end
  469. return state_x, a
  470. end
  471. -- call each other
  472. local filterm_gen
  473. local filterm_gen_shrink = function(fun, gen_x, param_x, state_x)
  474. return filterm_gen(fun, gen_x, param_x, gen_x(param_x, state_x))
  475. end
  476. filterm_gen = function(fun, gen_x, param_x, state_x, ...)
  477. if state_x == nil then
  478. return nil
  479. end
  480. if fun(...) then
  481. return state_x, ...
  482. end
  483. return filterm_gen_shrink(fun, gen_x, param_x, state_x)
  484. end
  485. local filter_detect = function(fun, gen_x, param_x, state_x, ...)
  486. if select('#', ...) < 2 then
  487. return filter1_gen(fun, gen_x, param_x, state_x, ...)
  488. else
  489. return filterm_gen(fun, gen_x, param_x, state_x, ...)
  490. end
  491. end
  492. local filter_gen = function(param, state_x)
  493. local fun, gen_x, param_x = param[1], param[2], param[3]
  494. return filter_detect(fun, gen_x, param_x, gen_x(param_x, state_x))
  495. end
  496. local filter = function(fun, gen, param, state)
  497. return wrap(filter_gen, {fun, gen, param}, state)
  498. end
  499. methods.filter = method1(filter)
  500. exports.filter = export1(filter)
  501. methods.remove_if = methods.filter
  502. exports.remove_if = exports.filter
  503. local grep = function(fun_or_regexp, gen, param, state)
  504. local fun = fun_or_regexp
  505. if type(fun_or_regexp) == "string" then
  506. fun = function(x) return string.find(x, fun_or_regexp) ~= nil end
  507. end
  508. return filter(fun, gen, param, state)
  509. end
  510. methods.grep = method1(grep)
  511. exports.grep = export1(grep)
  512. local partition = function(fun, gen, param, state)
  513. local neg_fun = function(...)
  514. return not fun(...)
  515. end
  516. return filter(fun, gen, param, state),
  517. filter(neg_fun, gen, param, state)
  518. end
  519. methods.partition = method1(partition)
  520. exports.partition = export1(partition)
  521. --------------------------------------------------------------------------------
  522. -- Reducing
  523. --------------------------------------------------------------------------------
  524. local foldl_call = function(fun, start, state, ...)
  525. if state == nil then
  526. return nil, start
  527. end
  528. return state, fun(start, ...)
  529. end
  530. local foldl = function(fun, start, gen_x, param_x, state_x)
  531. while true do
  532. state_x, start = foldl_call(fun, start, gen_x(param_x, state_x))
  533. if state_x == nil then
  534. break;
  535. end
  536. end
  537. return start
  538. end
  539. methods.foldl = method2(foldl)
  540. exports.foldl = export2(foldl)
  541. methods.reduce = methods.foldl
  542. exports.reduce = exports.foldl
  543. local length = function(gen, param, state)
  544. if gen == ipairs_gen or gen == string_gen then
  545. return #param
  546. end
  547. local len = 0
  548. repeat
  549. state = gen(param, state)
  550. len = len + 1
  551. until state == nil
  552. return len - 1
  553. end
  554. methods.length = method0(length)
  555. exports.length = export0(length)
  556. local is_null = function(gen, param, state)
  557. return gen(param, deepcopy(state)) == nil
  558. end
  559. methods.is_null = method0(is_null)
  560. exports.is_null = export0(is_null)
  561. local is_prefix_of = function(iter_x, iter_y)
  562. local gen_x, param_x, state_x = iter(iter_x)
  563. local gen_y, param_y, state_y = iter(iter_y)
  564. local r_x, r_y
  565. for i=1,10,1 do
  566. state_x, r_x = gen_x(param_x, state_x)
  567. state_y, r_y = gen_y(param_y, state_y)
  568. if state_x == nil then
  569. return true
  570. end
  571. if state_y == nil or r_x ~= r_y then
  572. return false
  573. end
  574. end
  575. end
  576. methods.is_prefix_of = is_prefix_of
  577. exports.is_prefix_of = is_prefix_of
  578. local all = function(fun, gen_x, param_x, state_x)
  579. local r
  580. repeat
  581. state_x, r = call_if_not_empty(fun, gen_x(param_x, state_x))
  582. until state_x == nil or not r
  583. return state_x == nil
  584. end
  585. methods.all = method1(all)
  586. exports.all = export1(all)
  587. methods.every = methods.all
  588. exports.every = exports.all
  589. local any = function(fun, gen_x, param_x, state_x)
  590. local r
  591. repeat
  592. state_x, r = call_if_not_empty(fun, gen_x(param_x, state_x))
  593. until state_x == nil or r
  594. return not not r
  595. end
  596. methods.any = method1(any)
  597. exports.any = export1(any)
  598. methods.some = methods.any
  599. exports.some = exports.any
  600. local sum = function(gen, param, state)
  601. local s = 0
  602. local r = 0
  603. repeat
  604. s = s + r
  605. state, r = gen(param, state)
  606. until state == nil
  607. return s
  608. end
  609. methods.sum = method0(sum)
  610. exports.sum = export0(sum)
  611. local product = function(gen, param, state)
  612. local p = 1
  613. local r = 1
  614. repeat
  615. p = p * r
  616. state, r = gen(param, state)
  617. until state == nil
  618. return p
  619. end
  620. methods.product = method0(product)
  621. exports.product = export0(product)
  622. local min_cmp = function(m, n)
  623. if n < m then return n else return m end
  624. end
  625. local max_cmp = function(m, n)
  626. if n > m then return n else return m end
  627. end
  628. local min = function(gen, param, state)
  629. local state, m = gen(param, state)
  630. if state == nil then
  631. error("min: iterator is empty")
  632. end
  633. local cmp
  634. if type(m) == "number" then
  635. -- An optimization: use math.min for numbers
  636. cmp = math.min
  637. else
  638. cmp = min_cmp
  639. end
  640. for _, r in gen, param, state do
  641. m = cmp(m, r)
  642. end
  643. return m
  644. end
  645. methods.min = method0(min)
  646. exports.min = export0(min)
  647. methods.minimum = methods.min
  648. exports.minimum = exports.min
  649. local min_by = function(cmp, gen_x, param_x, state_x)
  650. local state_x, m = gen_x(param_x, state_x)
  651. if state_x == nil then
  652. error("min: iterator is empty")
  653. end
  654. for _, r in gen_x, param_x, state_x do
  655. m = cmp(m, r)
  656. end
  657. return m
  658. end
  659. methods.min_by = method1(min_by)
  660. exports.min_by = export1(min_by)
  661. methods.minimum_by = methods.min_by
  662. exports.minimum_by = exports.min_by
  663. local max = function(gen_x, param_x, state_x)
  664. local state_x, m = gen_x(param_x, state_x)
  665. if state_x == nil then
  666. error("max: iterator is empty")
  667. end
  668. local cmp
  669. if type(m) == "number" then
  670. -- An optimization: use math.max for numbers
  671. cmp = math.max
  672. else
  673. cmp = max_cmp
  674. end
  675. for _, r in gen_x, param_x, state_x do
  676. m = cmp(m, r)
  677. end
  678. return m
  679. end
  680. methods.max = method0(max)
  681. exports.max = export0(max)
  682. methods.maximum = methods.max
  683. exports.maximum = exports.max
  684. local max_by = function(cmp, gen_x, param_x, state_x)
  685. local state_x, m = gen_x(param_x, state_x)
  686. if state_x == nil then
  687. error("max: iterator is empty")
  688. end
  689. for _, r in gen_x, param_x, state_x do
  690. m = cmp(m, r)
  691. end
  692. return m
  693. end
  694. methods.max_by = method1(max_by)
  695. exports.max_by = export1(max_by)
  696. methods.maximum_by = methods.maximum_by
  697. exports.maximum_by = exports.maximum_by
  698. local totable = function(gen_x, param_x, state_x)
  699. local tab, key, val = {}
  700. while true do
  701. state_x, val = gen_x(param_x, state_x)
  702. if state_x == nil then
  703. break
  704. end
  705. table.insert(tab, val)
  706. end
  707. return tab
  708. end
  709. methods.totable = method0(totable)
  710. exports.totable = export0(totable)
  711. local tomap = function(gen_x, param_x, state_x)
  712. local tab, key, val = {}
  713. while true do
  714. state_x, key, val = gen_x(param_x, state_x)
  715. if state_x == nil then
  716. break
  717. end
  718. tab[key] = val
  719. end
  720. return tab
  721. end
  722. methods.tomap = method0(tomap)
  723. exports.tomap = export0(tomap)
  724. --------------------------------------------------------------------------------
  725. -- Transformations
  726. --------------------------------------------------------------------------------
  727. local map_gen = function(param, state)
  728. local gen_x, param_x, fun = param[1], param[2], param[3]
  729. return call_if_not_empty(fun, gen_x(param_x, state))
  730. end
  731. local map = function(fun, gen, param, state)
  732. return wrap(map_gen, {gen, param, fun}, state)
  733. end
  734. methods.map = method1(map)
  735. exports.map = export1(map)
  736. local enumerate_gen_call = function(state, i, state_x, ...)
  737. if state_x == nil then
  738. return nil
  739. end
  740. return {i + 1, state_x}, i, ...
  741. end
  742. local enumerate_gen = function(param, state)
  743. local gen_x, param_x = param[1], param[2]
  744. local i, state_x = state[1], state[2]
  745. return enumerate_gen_call(state, i, gen_x(param_x, state_x))
  746. end
  747. local enumerate = function(gen, param, state)
  748. return wrap(enumerate_gen, {gen, param}, {1, state})
  749. end
  750. methods.enumerate = method0(enumerate)
  751. exports.enumerate = export0(enumerate)
  752. local intersperse_call = function(i, state_x, ...)
  753. if state_x == nil then
  754. return nil
  755. end
  756. return {i + 1, state_x}, ...
  757. end
  758. local intersperse_gen = function(param, state)
  759. local x, gen_x, param_x = param[1], param[2], param[3]
  760. local i, state_x = state[1], state[2]
  761. if i % 2 == 1 then
  762. return {i + 1, state_x}, x
  763. else
  764. return intersperse_call(i, gen_x(param_x, state_x))
  765. end
  766. end
  767. -- TODO: interperse must not add x to the tail
  768. local intersperse = function(x, gen, param, state)
  769. return wrap(intersperse_gen, {x, gen, param}, {0, state})
  770. end
  771. methods.intersperse = method1(intersperse)
  772. exports.intersperse = export1(intersperse)
  773. --------------------------------------------------------------------------------
  774. -- Compositions
  775. --------------------------------------------------------------------------------
  776. local function zip_gen_r(param, state, state_new, ...)
  777. if #state_new == #param / 2 then
  778. return state_new, ...
  779. end
  780. local i = #state_new + 1
  781. local gen_x, param_x = param[2 * i - 1], param[2 * i]
  782. local state_x, r = gen_x(param_x, state[i])
  783. if state_x == nil then
  784. return nil
  785. end
  786. table.insert(state_new, state_x)
  787. return zip_gen_r(param, state, state_new, r, ...)
  788. end
  789. local zip_gen = function(param, state)
  790. return zip_gen_r(param, state, {})
  791. end
  792. -- A special hack for zip/chain to skip last two state, if a wrapped iterator
  793. -- has been passed
  794. local numargs = function(...)
  795. local n = select('#', ...)
  796. if n >= 3 then
  797. -- Fix last argument
  798. local it = select(n - 2, ...)
  799. if type(it) == 'table' and getmetatable(it) == iterator_mt and
  800. it.param == select(n - 1, ...) and it.state == select(n, ...) then
  801. return n - 2
  802. end
  803. end
  804. return n
  805. end
  806. local zip = function(...)
  807. local n = numargs(...)
  808. if n == 0 then
  809. return wrap(nil_gen, nil, nil)
  810. end
  811. local param = { [2 * n] = 0 }
  812. local state = { [n] = 0 }
  813. local i, gen_x, param_x, state_x
  814. for i=1,n,1 do
  815. local it = select(n - i + 1, ...)
  816. gen_x, param_x, state_x = rawiter(it)
  817. param[2 * i - 1] = gen_x
  818. param[2 * i] = param_x
  819. state[i] = state_x
  820. end
  821. return wrap(zip_gen, param, state)
  822. end
  823. methods.zip = zip
  824. exports.zip = zip
  825. local cycle_gen_call = function(param, state_x, ...)
  826. if state_x == nil then
  827. local gen_x, param_x, state_x0 = param[1], param[2], param[3]
  828. return gen_x(param_x, deepcopy(state_x0))
  829. end
  830. return state_x, ...
  831. end
  832. local cycle_gen = function(param, state_x)
  833. local gen_x, param_x, state_x0 = param[1], param[2], param[3]
  834. return cycle_gen_call(param, gen_x(param_x, state_x))
  835. end
  836. local cycle = function(gen, param, state)
  837. return wrap(cycle_gen, {gen, param, state}, deepcopy(state))
  838. end
  839. methods.cycle = method0(cycle)
  840. exports.cycle = export0(cycle)
  841. -- call each other
  842. local chain_gen_r1
  843. local chain_gen_r2 = function(param, state, state_x, ...)
  844. if state_x == nil then
  845. local i = state[1]
  846. i = i + 1
  847. if param[3 * i - 1] == nil then
  848. return nil
  849. end
  850. local state_x = param[3 * i]
  851. return chain_gen_r1(param, {i, state_x})
  852. end
  853. return {state[1], state_x}, ...
  854. end
  855. chain_gen_r1 = function(param, state)
  856. local i, state_x = state[1], state[2]
  857. local gen_x, param_x = param[3 * i - 2], param[3 * i - 1]
  858. return chain_gen_r2(param, state, gen_x(param_x, state[2]))
  859. end
  860. local chain = function(...)
  861. local n = numargs(...)
  862. if n == 0 then
  863. return wrap(nil_gen, nil, nil)
  864. end
  865. local param = { [3 * n] = 0 }
  866. local i, gen_x, param_x, state_x
  867. for i=1,n,1 do
  868. local elem = select(i, ...)
  869. gen_x, param_x, state_x = iter(elem)
  870. param[3 * i - 2] = gen_x
  871. param[3 * i - 1] = param_x
  872. param[3 * i] = state_x
  873. end
  874. return wrap(chain_gen_r1, param, {1, param[3]})
  875. end
  876. methods.chain = chain
  877. exports.chain = chain
  878. --------------------------------------------------------------------------------
  879. -- Operators
  880. --------------------------------------------------------------------------------
  881. local operator = {
  882. ----------------------------------------------------------------------------
  883. -- Comparison operators
  884. ----------------------------------------------------------------------------
  885. lt = function(a, b) return a < b end,
  886. le = function(a, b) return a <= b end,
  887. eq = function(a, b) return a == b end,
  888. ne = function(a, b) return a ~= b end,
  889. ge = function(a, b) return a >= b end,
  890. gt = function(a, b) return a > b end,
  891. ----------------------------------------------------------------------------
  892. -- Arithmetic operators
  893. ----------------------------------------------------------------------------
  894. add = function(a, b) return a + b end,
  895. div = function(a, b) return a / b end,
  896. floordiv = function(a, b) return math.floor(a/b) end,
  897. intdiv = function(a, b)
  898. local q = a / b
  899. if a >= 0 then return math.floor(q) else return math.ceil(q) end
  900. end,
  901. mod = function(a, b) return a % b end,
  902. mul = function(a, b) return a * b end,
  903. neq = function(a) return -a end,
  904. unm = function(a) return -a end, -- an alias
  905. pow = function(a, b) return a ^ b end,
  906. sub = function(a, b) return a - b end,
  907. truediv = function(a, b) return a / b end,
  908. ----------------------------------------------------------------------------
  909. -- String operators
  910. ----------------------------------------------------------------------------
  911. concat = function(a, b) return a..b end,
  912. len = function(a) return #a end,
  913. length = function(a) return #a end, -- an alias
  914. ----------------------------------------------------------------------------
  915. -- Logical operators
  916. ----------------------------------------------------------------------------
  917. land = function(a, b) return a and b end,
  918. lor = function(a, b) return a or b end,
  919. lnot = function(a) return not a end,
  920. truth = function(a) return not not a end,
  921. }
  922. exports.operator = operator
  923. methods.operator = operator
  924. exports.op = operator
  925. methods.op = operator
  926. --------------------------------------------------------------------------------
  927. -- module definitions
  928. --------------------------------------------------------------------------------
  929. -- a special syntax sugar to export all functions to the global table
  930. setmetatable(exports, {
  931. __call = function(t, override)
  932. for k, v in pairs(t) do
  933. if _G[k] ~= nil then
  934. local msg = 'function ' .. k .. ' already exists in global scope.'
  935. if override then
  936. _G[k] = v
  937. print('WARNING: ' .. msg .. ' Overwritten.')
  938. else
  939. print('NOTICE: ' .. msg .. ' Skipped.')
  940. end
  941. else
  942. _G[k] = v
  943. end
  944. end
  945. end,
  946. })
  947. return exports