1 files changed, 71 insertions, 7 deletions

diff --git a/src/Dom/Dom.c b/src/Dom/Dom.c
index f0d4f36..4eb3ad4 100644
--- a/src/Dom/Dom.c
+++ b/src/Dom/Dom.c
@@ -10,8 +10,75 @@ static int absolute_index(lua_State *lua, int index) {
   return index < 0 ? lua_gettop(lua) + index + 1 : index;
 }
 
+static int table_is_node(lua_State *lua, int table_index) {
+  int is_node;
+
+  table_index = absolute_index(lua, table_index);
+  lua_getfield(lua, table_index, "tag");
+  is_node = !lua_isnil(lua, -1);
+  lua_pop(lua, 1);
+  if (!is_node) {
+    lua_getfield(lua, table_index, "kind");
+    is_node = !lua_isnil(lua, -1);
+    lua_pop(lua, 1);
+  }
+  return is_node;
+}
+
+static void append_element_child(lua_State *lua, int children_index,
+                                 int child_index, size_t *child_count,
+                                 unsigned int depth) {
+  size_t array_length;
+  size_t index;
+  int child_type;
+
+  if (!lua_checkstack(lua, 2)) {
+    luaL_error(lua, "html.element children are nested too deeply");
+  }
+  children_index = absolute_index(lua, children_index);
+  child_index = absolute_index(lua, child_index);
+  child_type = lua_type(lua, child_index);
+
+  if (child_type == LUA_TNIL) {
+    return;
+  }
+  if (child_type == LUA_TSTRING ||
+      (child_type == LUA_TTABLE && table_is_node(lua, child_index))) {
+    if (*child_count == INT_MAX) {
+      luaL_error(lua, "html.element has too many children");
+    }
+    ++*child_count;
+    lua_pushvalue(lua, child_index);
+    lua_rawseti(lua, children_index, (int)*child_count);
+    return;
+  }
+  if (child_type != LUA_TTABLE) {
+    luaL_error(lua,
+               "html.element child must be a string, node, child array, or "
+               "nil (got %s)",
+               lua_typename(lua, child_type));
+  }
+  if (depth == 100) {
+    luaL_error(lua, "html.element children are nested too deeply");
+  }
+
+  array_length = lua_objlen(lua, child_index);
+  if (array_length > INT_MAX) {
+    luaL_error(lua, "html.element child array is too large");
+  }
+  for (index = 1; index <= array_length; ++index) {
+    lua_rawgeti(lua, child_index, (int)index);
+    append_element_child(lua, children_index, -1, child_count, depth + 1);
+    lua_pop(lua, 1);
+  }
+}
+
 static int html_element(lua_State *lua) {
   const char *tag;
+  size_t child_count = 0;
+  int argument_count = lua_gettop(lua);
+  int children_index;
+  int index;
 
   luaL_checktype(lua, 1, LUA_TSTRING);
   tag = lua_tostring(lua, 1);
@@ -21,9 +88,6 @@ static int html_element(lua_State *lua) {
   if (!lua_isnoneornil(lua, 2)) {
     luaL_checktype(lua, 2, LUA_TTABLE);
   }
-  if (!lua_isnoneornil(lua, 3)) {
-    luaL_checktype(lua, 3, LUA_TTABLE);
-  }
 
   lua_newtable(lua);
 
@@ -37,10 +101,10 @@ static int html_element(lua_State *lua) {
   }
   lua_setfield(lua, -2, "attributes");
 
-  if (lua_istable(lua, 3)) {
-    lua_pushvalue(lua, 3);
-  } else {
-    lua_newtable(lua);
+  lua_newtable(lua);
+  children_index = lua_gettop(lua);
+  for (index = 3; index <= argument_count; ++index) {
+    append_element_child(lua, children_index, index, &child_count, 0);
   }
   lua_setfield(lua, -2, "children");