From 6f8d6b19d20ce9cc42eaab7f21aa0fc6c937003a Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@rspamd.com>
Date: Sat, 10 Dec 2022 22:40:31 +0000
Subject: [PATCH] [Feature] Some rework of the selectors framework

* Add `apply_methods` transform
* Rework userdata application of the transform functor
* Add more join methods
---
 lualib/lua_selectors/init.lua       | 11 ++++---
 lualib/lua_selectors/transforms.lua | 48 +++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/lualib/lua_selectors/init.lua b/lualib/lua_selectors/init.lua
index e4d76471e..f85b9a487 100644
--- a/lualib/lua_selectors/init.lua
+++ b/lualib/lua_selectors/init.lua
@@ -97,7 +97,8 @@ local function process_selector(task, sel)
   local pipe = sel.processor_pipe or E
   local first_elt = pipe[1]
 
-  if first_elt and first_elt.method then
+  if first_elt and (first_elt.method or
+      fun.any(function(t) return t == 'userdata' or t == 'table' end, first_elt.types)) then
     -- Explicit conversion
     local meth = first_elt
 
@@ -114,7 +115,7 @@ local function process_selector(task, sel)
         -- Map method to a list of inputs, excluding empty elements
         input = fun.filter(function(map_elt) return map_elt end,
             fun.map(function(list_elt)
-              local ret, _ = meth.process(list_elt, pt)
+              local ret, _ = meth.process(list_elt, pt, meth.args)
               return ret
             end, input))
         etype = 'string_list'
@@ -124,7 +125,6 @@ local function process_selector(task, sel)
     pipe = fun.drop_n(1, pipe)
   elseif etype:match('^userdata') or etype:match('^table') then
     -- Implicit conversion
-
     local pt = pure_type(etype)
 
     if not pt then
@@ -140,6 +140,8 @@ local function process_selector(task, sel)
           end, input))
       etype = 'string_list'
     end
+  else
+    lua_util.debugm(M, task, 'avoid implicit conversion as the transformer accepts complex input')
   end
 
   -- Now we fold elements using left fold
@@ -189,11 +191,11 @@ local function process_selector(task, sel)
   if not res or not res[1] then return nil end -- Pipeline failed
 
   if not allowed_type(res[2]) then
-
     -- Search for implicit conversion
     local pt = pure_type(res[2])
 
     if pt then
+
       lua_util.debugm(M, task, 'apply implicit map %s->string_list', pt)
       res[1] = fun.map(function(e) return implicit_tostring(pt, e) end, res[1])
       res[2] = 'string_list'
@@ -382,6 +384,7 @@ exports.parse_selector = function(cfg, str)
         local processor = lua_util.shallowcopy(transform_function[proc_name])
         processor.name = proc_name
         processor.args = proc_tbl[2] or E
+        logger.errx('hui: %s -> %s', proc_name, processor.args)
 
         if not check_args(processor.name, processor.args_schema, processor.args) then
           pipeline_error = 'args schema for ' .. proc_name
diff --git a/lualib/lua_selectors/transforms.lua b/lualib/lua_selectors/transforms.lua
index fc929175f..f6a0b0a37 100644
--- a/lualib/lua_selectors/transforms.lua
+++ b/lualib/lua_selectors/transforms.lua
@@ -113,6 +113,37 @@ local transform_function = {
     ['description'] = 'Joins strings into a single string using separator in the argument',
     ['args_schema'] = {ts.string:is_optional()}
   },
+  -- Joins strings into a set of strings using N elements and a separator in the argument
+  ['join_nth'] = {
+    ['types'] = {
+      ['string_list'] = true
+    },
+    ['process'] = function(inp, _, args)
+      local step = args[1]
+      local sep = args[2] or ''
+      local inp_t = fun.totable(inp)
+      local res = {}
+
+      for i=1, #inp_t, step do
+        table.insert(res, table.concat(inp_t, sep, i, i + step))
+      end
+      return res,'string_list'
+    end,
+    ['description'] = 'Joins strings into a set of strings using N elements and a separator in the argument',
+    ['args_schema'] = {ts.number + ts.string / tonumber, ts.string:is_optional()}
+  },
+  -- Joins tables into a table of strings
+  ['join_tables'] = {
+    ['types'] = {
+      ['string_list'] = true
+    },
+    ['process'] = function(inp, _, args)
+      local sep = args[1] or ''
+      return fun.map(function(t) return table.concat(t, sep) end, inp), 'string_list'
+    end,
+    ['description'] = 'Joins tables into a table of strings',
+    ['args_schema'] = {ts.string:is_optional()}
+  },
   -- Sort strings
   ['sort'] = {
     ['types'] = {
@@ -480,6 +511,23 @@ Empty string comes the first argument or 'true', non-empty string comes nil]],
     ['description'] = 'Removes all nils from a list of strings (when converted implicitly)',
     ['args_schema'] = {}
   },
+  -- Call a set of methods on a userdata object
+  ['apply_methods'] = {
+    ['types'] = {
+      ['userdata'] = true,
+    },
+    ['process'] = function(inp, _, args)
+      local res = {}
+      for _,arg in ipairs(args) do
+        local meth = inp[arg]
+        local ret = meth(inp)
+        if not ret then return nil end
+        table.insert(res, tostring(ret))
+      end
+      return res,'string_list'
+    end,
+    ['description'] = 'Apply a list of method calls to the userdata object',
+  }
 }
 
 transform_function.match = transform_function.regexp
-- 
2.39.5