feat: add html.find_by_class

2 files changed, 121 insertions, 21 deletions

diff --git a/examples/site/build.lua b/examples/site/build.lua
index 4d0939a..f3b07e1 100644
--- a/examples/site/build.lua
+++ b/examples/site/build.lua
@@ -1,17 +1,5 @@
 -- External scripts are resolved relative to the HTML file that includes them :)
 
-local function find_by_class(node, wanted)
-  if type(node) ~= "table" then return nil end
-  local classes = node.attributes and node.attributes.class or ""
-  if (" " .. classes .. " "):find(" " .. wanted .. " ", 1, true) then
-    return node
-  end
-  for _, child in ipairs(node.children or {}) do
-    local match = find_by_class(child, wanted)
-    if match then return match end
-  end
-end
-
 local features = {
   "inline Lua",
   "external Lua",
@@ -21,7 +9,7 @@ local features = {
   "mixed content"
 }
 
-local feature_list = assert(find_by_class(page, "feature-list"))
+local feature_list = assert(html.find_by_class(page, "feature-list"))
 for _, feature in ipairs(features) do
   table.insert(feature_list.children, html.element("li", nil, { feature }))
 end
@@ -32,10 +20,10 @@ for entry in fs.walk(".") do
 end
 table.sort(entries, function(left, right) return left.path < right.path end)
 
-local summary = assert(find_by_class(page, "file-summary"))
+local summary = assert(html.find_by_class(page, "file-summary"))
 summary.text = tostring(#entries) .. " entries"
 
-local file_list = assert(find_by_class(page, "file-list"))
+local file_list = assert(html.find_by_class(page, "file-list"))
 for _, entry in ipairs(entries) do
   table.insert(file_list.children, html.element("li", nil, {
     html.element("code", nil, { entry.path }),
@@ -43,10 +31,10 @@ for _, entry in ipairs(entries) do
   }))
 end
 
-local build_summary = assert(find_by_class(page, "build-summary"))
+local build_summary = assert(html.find_by_class(page, "build-summary"))
 build_summary.text = "Built with " .. site.name .. " " .. site.version
 
-local mixed = assert(find_by_class(page, "mixed-content"))
+local mixed = assert(html.find_by_class(page, "mixed-content"))
 mixed.children = {
   { kind = "comment", text = "generated by Lua" },
   "Text with ",
diff --git a/src/Dom/Dom.c b/src/Dom/Dom.c
index 6dc6760..f0d4f36 100644
--- a/src/Dom/Dom.c
+++ b/src/Dom/Dom.c
@@ -6,6 +6,10 @@
 #include <string.h>
 #include <strings.h>
 
+static int absolute_index(lua_State *lua, int index) {
+  return index < 0 ? lua_gettop(lua) + index + 1 : index;
+}
+
 static int html_element(lua_State *lua) {
   const char *tag;
 
@@ -43,17 +47,125 @@ static int html_element(lua_State *lua) {
   return 1;
 }
 
+static int is_html_whitespace(char character) {
+  return character == ' ' || character == '\t' || character == '\n' ||
+         character == '\f' || character == '\r';
+}
+
+static int class_list_contains(const char *classes, size_t classes_length,
+                               const char *wanted, size_t wanted_length) {
+  size_t index = 0;
+
+  while (index < classes_length) {
+    size_t start;
+
+    while (index < classes_length && is_html_whitespace(classes[index])) {
+      ++index;
+    }
+    start = index;
+    while (index < classes_length && !is_html_whitespace(classes[index])) {
+      ++index;
+    }
+    if (index - start == wanted_length &&
+        memcmp(classes + start, wanted, wanted_length) == 0) {
+      return 1;
+    }
+  }
+  return 0;
+}
+
+static int node_has_class(lua_State *lua, int node_index, const char *wanted,
+                          size_t wanted_length) {
+  const char *classes;
+  size_t classes_length;
+  int matches = 0;
+
+  node_index = absolute_index(lua, node_index);
+  lua_getfield(lua, node_index, "attributes");
+  if (lua_istable(lua, -1)) {
+    lua_getfield(lua, -1, "class");
+    if (lua_type(lua, -1) == LUA_TSTRING) {
+      classes = lua_tolstring(lua, -1, &classes_length);
+      matches = class_list_contains(classes, classes_length, wanted,
+                                    wanted_length);
+    }
+    lua_pop(lua, 1);
+  }
+  lua_pop(lua, 1);
+  return matches;
+}
+
+static int push_first_by_class(lua_State *lua, int node_index,
+                               const char *wanted, size_t wanted_length) {
+  size_t child_count;
+  size_t index;
+
+  if (!lua_checkstack(lua, 3)) {
+    return luaL_error(lua, "HTML tree is too deep to search");
+  }
+  node_index = absolute_index(lua, node_index);
+  if (node_has_class(lua, node_index, wanted, wanted_length)) {
+    lua_pushvalue(lua, node_index);
+    return 1;
+  }
+
+  lua_getfield(lua, node_index, "children");
+  if (!lua_istable(lua, -1)) {
+    lua_pop(lua, 1);
+    return 0;
+  }
+
+  child_count = lua_objlen(lua, -1);
+  if (child_count > INT_MAX) {
+    return luaL_error(lua, "HTML node has too many children to search");
+  }
+  for (index = 1; index <= child_count; ++index) {
+    lua_rawgeti(lua, -1, (int)index);
+    if (lua_istable(lua, -1) &&
+        push_first_by_class(lua, -1, wanted, wanted_length)) {
+      lua_remove(lua, -2);
+      lua_remove(lua, -2);
+      return 1;
+    }
+    lua_pop(lua, 1);
+  }
+
+  lua_pop(lua, 1);
+  return 0;
+}
+
+static int html_find_by_class(lua_State *lua) {
+  const char *wanted;
+  size_t wanted_length;
+  size_t index;
+
+  luaL_checktype(lua, 1, LUA_TTABLE);
+  luaL_checktype(lua, 2, LUA_TSTRING);
+  wanted = lua_tolstring(lua, 2, &wanted_length);
+  if (wanted_length == 0) {
+    return luaL_argerror(lua, 2, "class cannot be empty");
+  }
+  for (index = 0; index < wanted_length; ++index) {
+    if (is_html_whitespace(wanted[index])) {
+      return luaL_argerror(lua, 2, "class must be a single token");
+    }
+  }
+
+  if (!push_first_by_class(lua, 1, wanted, wanted_length)) {
+    lua_pushnil(lua);
+  }
+  return 1;
+}
+
 void dom_register_html(lua_State *lua) {
   lua_newtable(lua);
   lua_pushcfunction(lua, html_element);
   lua_setfield(lua, -2, "element");
+  lua_pushcfunction(lua, html_find_by_class);
+  lua_setfield(lua, -2, "find_by_class");
   lua_setglobal(lua, "html");
 }
 
-static int absolute_index(lua_State *lua, int index) {
-  return index < 0 ? lua_gettop(lua) + index + 1 : index;
-}
-
 static int is_lua_element(const xmlNode *node) {
   return node->type == XML_ELEMENT_NODE &&
          strcasecmp((const char *)node->name, "lua") == 0;