feat: add query, query_all, find_by_id

1 files changed, 235 insertions, 21 deletions

diff --git a/src/Dom/Dom.c b/src/Dom/Dom.c
index 67817aa..3a70f61 100644
--- a/src/Dom/Dom.c
+++ b/src/Dom/Dom.c
@@ -243,66 +243,270 @@ static int node_has_class(lua_State *lua, int node_index, const char *wanted,
   return matches;
 }
 
-static int push_first_by_class(lua_State *lua, int node_index,
-                               const char *wanted, size_t wanted_length) {
+typedef struct {
+  const char *tag;
+  size_t tag_length;
+  const char *id;
+  size_t id_length;
+  const char *suffixes;
+  size_t suffixes_length;
+} CssSelector;
+
+static int is_selector_delimiter(char character) {
+  return character == '.' || character == '#';
+}
+
+static int is_unsupported_selector_character(char character) {
+  return is_html_whitespace(character) || character == '>' ||
+         character == '+' || character == '~' || character == ',' ||
+         character == '[' || character == ']';
+}
+
+static int parse_css_selector(lua_State *lua, const char *source,
+                              size_t length, CssSelector *selector) {
+  size_t index = 0;
+
+  memset(selector, 0, sizeof(*selector));
+  if (length == 0) {
+    return luaL_argerror(lua, 2, "selector cannot be empty");
+  }
+
+  if (source[0] == '*') {
+    index = 1;
+  } else if (!is_selector_delimiter(source[0])) {
+    selector->tag = source;
+    while (index < length && !is_selector_delimiter(source[index])) {
+      if (is_unsupported_selector_character(source[index])) {
+        return luaL_argerror(lua, 2, "unsupported CSS selector");
+      }
+      ++index;
+    }
+    selector->tag_length = index;
+  }
+
+  selector->suffixes = source + index;
+  selector->suffixes_length = length - index;
+  while (index < length) {
+    char delimiter = source[index++];
+    size_t start = index;
+
+    if (!is_selector_delimiter(delimiter)) {
+      return luaL_argerror(lua, 2, "unsupported CSS selector");
+    }
+    while (index < length && !is_selector_delimiter(source[index])) {
+      if (is_unsupported_selector_character(source[index])) {
+        return luaL_argerror(lua, 2, "unsupported CSS selector");
+      }
+      ++index;
+    }
+    if (index == start) {
+      return luaL_argerror(lua, 2, "invalid CSS selector");
+    }
+    if (delimiter == '#') {
+      if (selector->id != NULL) {
+        return luaL_argerror(lua, 2, "selector can contain only one id");
+      }
+      selector->id = source + start;
+      selector->id_length = index - start;
+    }
+  }
+  return 0;
+}
+
+static int string_field_equals(lua_State *lua, int table_index,
+                               const char *field, const char *wanted,
+                               size_t wanted_length, int ignore_case) {
+  const char *value;
+  size_t value_length;
+  int matches = 0;
+
+  table_index = absolute_index(lua, table_index);
+  lua_getfield(lua, table_index, field);
+  if (lua_type(lua, -1) == LUA_TSTRING) {
+    value = lua_tolstring(lua, -1, &value_length);
+    if (value_length == wanted_length) {
+      matches = ignore_case ? strncasecmp(value, wanted, wanted_length) == 0
+                            : memcmp(value, wanted, wanted_length) == 0;
+    }
+  }
+  lua_pop(lua, 1);
+  return matches;
+}
+
+static int node_matches_selector(lua_State *lua, int node_index,
+                                 const CssSelector *selector) {
+  size_t index = 0;
+  int attributes_index;
+
+  node_index = absolute_index(lua, node_index);
+  lua_getfield(lua, node_index, "tag");
+  if (lua_type(lua, -1) != LUA_TSTRING) {
+    lua_pop(lua, 1);
+    return 0;
+  }
+  lua_pop(lua, 1);
+  if (selector->tag != NULL &&
+      !string_field_equals(lua, node_index, "tag", selector->tag,
+                           selector->tag_length, 1)) {
+    return 0;
+  }
+
+  lua_getfield(lua, node_index, "attributes");
+  if (!lua_istable(lua, -1)) {
+    lua_pop(lua, 1);
+    return selector->suffixes_length == 0;
+  }
+  attributes_index = lua_gettop(lua);
+
+  while (index < selector->suffixes_length) {
+    char delimiter = selector->suffixes[index++];
+    size_t start = index;
+    int matches;
+
+    while (index < selector->suffixes_length &&
+           !is_selector_delimiter(selector->suffixes[index])) {
+      ++index;
+    }
+    if (delimiter == '.') {
+      matches = node_has_class(lua, node_index, selector->suffixes + start,
+                               index - start);
+    } else {
+      matches = string_field_equals(lua, attributes_index, "id",
+                                    selector->suffixes + start, index - start,
+                                    0);
+    }
+    if (!matches) {
+      lua_pop(lua, 1);
+      return 0;
+    }
+  }
+
+  lua_pop(lua, 1);
+  return 1;
+}
+
+static void append_matches_by_selector(lua_State *lua, int node_index,
+                                       const CssSelector *selector,
+                                       int matches_index,
+                                       size_t *match_count) {
   size_t child_count;
   size_t index;
 
   if (!lua_checkstack(lua, 3)) {
-    return luaL_error(lua, "HTML tree is too deep to search");
+    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)) {
+  matches_index = absolute_index(lua, matches_index);
+  if (node_matches_selector(lua, node_index, selector)) {
+    if (*match_count == INT_MAX) {
+      luaL_error(lua, "HTML query has too many matches");
+    }
+    ++*match_count;
     lua_pushvalue(lua, node_index);
-    return 1;
+    lua_rawseti(lua, matches_index, (int)*match_count);
   }
 
   lua_getfield(lua, node_index, "children");
   if (!lua_istable(lua, -1)) {
     lua_pop(lua, 1);
-    return 0;
+    return;
   }
 
   child_count = lua_objlen(lua, -1);
   if (child_count > INT_MAX) {
-    return luaL_error(lua, "HTML node has too many children to search");
+    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;
+    if (lua_istable(lua, -1)) {
+      append_matches_by_selector(lua, -1, selector, matches_index,
+                                 match_count);
     }
     lua_pop(lua, 1);
   }
 
   lua_pop(lua, 1);
-  return 0;
 }
 
-static int html_find_by_class(lua_State *lua) {
+static int query_all(lua_State *lua, const char *source, size_t length) {
+  CssSelector selector;
+  size_t match_count = 0;
+  int matches_index;
+
+  parse_css_selector(lua, source, length, &selector);
+  lua_newtable(lua);
+  matches_index = lua_gettop(lua);
+  append_matches_by_selector(lua, 1, &selector, matches_index, &match_count);
+  return 1;
+}
+
+static int query(lua_State *lua, const char *source, size_t length) {
+  query_all(lua, source, length);
+  lua_rawgeti(lua, -1, 1);
+  lua_remove(lua, -2);
+  return 1;
+}
+
+static const char *check_query_arguments(lua_State *lua, size_t *length) {
+  luaL_checktype(lua, 1, LUA_TTABLE);
+  luaL_checktype(lua, 2, LUA_TSTRING);
+  return lua_tolstring(lua, 2, length);
+}
+
+static int html_query(lua_State *lua) {
+  const char *source;
+  size_t length;
+
+  source = check_query_arguments(lua, &length);
+  return query(lua, source, length);
+}
+
+static int html_query_all(lua_State *lua) {
+  const char *source;
+  size_t length;
+
+  source = check_query_arguments(lua, &length);
+  return query_all(lua, source, length);
+}
+
+static int find_by_attribute_token(lua_State *lua, char prefix,
+                                   const char *empty_error,
+                                   const char *token_error) {
   const char *wanted;
   size_t wanted_length;
   size_t index;
+  luaL_Buffer selector;
 
   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");
+    return luaL_argerror(lua, 2, empty_error);
   }
   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 (is_html_whitespace(wanted[index]) ||
+        is_selector_delimiter(wanted[index]) ||
+        is_unsupported_selector_character(wanted[index])) {
+      return luaL_argerror(lua, 2, token_error);
     }
   }
 
-  if (!push_first_by_class(lua, 1, wanted, wanted_length)) {
-    lua_pushnil(lua);
-  }
-  return 1;
+  luaL_buffinit(lua, &selector);
+  luaL_addchar(&selector, prefix);
+  luaL_addlstring(&selector, wanted, wanted_length);
+  luaL_pushresult(&selector);
+  wanted = lua_tolstring(lua, -1, &wanted_length);
+  return query(lua, wanted, wanted_length);
+}
+
+static int html_find_by_class(lua_State *lua) {
+  return find_by_attribute_token(lua, '.', "class cannot be empty",
+                                 "class must be a single token");
+}
+
+static int html_find_by_id(lua_State *lua) {
+  return find_by_attribute_token(lua, '#', "id cannot be empty",
+                                 "id must be a single token");
 }
 
 void dom_register_html(lua_State *lua) {
@@ -310,6 +514,10 @@ void dom_register_html(lua_State *lua) {
     lua_newtable(lua);
     lua_pushcfunction(lua, html_append);
     lua_setfield(lua, -2, "append");
+    lua_pushcfunction(lua, html_query);
+    lua_setfield(lua, -2, "query");
+    lua_pushcfunction(lua, html_query_all);
+    lua_setfield(lua, -2, "query_all");
     lua_setfield(lua, -2, "__index");
     lua_pushboolean(lua, 0);
     lua_setfield(lua, -2, "__metatable");
@@ -323,6 +531,12 @@ void dom_register_html(lua_State *lua) {
   lua_setfield(lua, -2, "append");
   lua_pushcfunction(lua, html_find_by_class);
   lua_setfield(lua, -2, "find_by_class");
+  lua_pushcfunction(lua, html_find_by_id);
+  lua_setfield(lua, -2, "find_by_id");
+  lua_pushcfunction(lua, html_query);
+  lua_setfield(lua, -2, "query");
+  lua_pushcfunction(lua, html_query_all);
+  lua_setfield(lua, -2, "query_all");
   lua_setglobal(lua, "html");
 }