1 files changed, 116 insertions, 4 deletions

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;