Merge pull request #1214 from joto/prep-for-changes-multistage-proc

Prep for changes multistage proc

Author: lonvia

docs/flex.md

@@ -37,8 +37,6 @@ The following functions are defined:
   more control than the more convenient other functions.
 * `osm2pgsql.mark_way(id)`: Mark the OSM way with the specified id. This way
   will be processed (again) in stage 2.
-* `osm2pgsql.mark_relation(id)`: Mark the OSM relation with the specified id.
-  This relation will be processed (again) in stage 2.
 
 You are expected to define one or more of the following functions:
 
@@ -226,10 +224,9 @@ a default transformation. These are the defaults:
 
 ## Stages
 
-Osm2pgsql processes the data in up to two stages. You can mark ways or
-relations in stage 1 for processing in stage 2 by calling
-`osm2pgsql.mark_way(id)` or `osm2pgsql.mark_relation(id)`, respectively. If you
-don't mark any objects, nothing will be done in stage 2.
+Osm2pgsql processes the data in up to two stages. You can mark ways in stage 1
+for processing in stage 2 by calling `osm2pgsql.mark_way(id)`. If you don't
+mark any ways, nothing will be done in stage 2.
 
 You can look at `osm2pgsql.stage` to see in which stage you are.
 

src/init.lua

@@ -33,10 +33,6 @@ function osm2pgsql.mark_way(id)
     return osm2pgsql.mark('w', id)
 end
 
-function osm2pgsql.mark_relation(id)
-    return osm2pgsql.mark('r', id)
-end
-
 function osm2pgsql.clamp(value, low, high)
     return math.min(math.max(value, low), high)
 end

src/output-flex.cpp

@@ -69,14 +69,18 @@ static char const osm2pgsql_table_name[] = "osm2pgsql.table";
 static char const osm2pgsql_object_metatable[] = "osm2pgsql.object_metatable";
 
 prepared_lua_function_t::prepared_lua_function_t(lua_State *lua_state,
-                                                 char const *name)
+                                                 calling_context context,
+                                                 char const *name, int nresults)
 {
     int const index = lua_gettop(lua_state);
 
     lua_getfield(lua_state, 1, name);
 
     if (lua_type(lua_state, -1) == LUA_TFUNCTION) {
         m_index = index;
+        m_name = name;
+        m_nresults = nresults;
+        m_calling_context = context;
         return;
     }
 
@@ -518,9 +522,7 @@ int output_flex_t::app_mark()
     osmium::object_id_type const id = luaL_checkinteger(lua_state(), 2);
 
     if (type_name[0] == 'w') {
-        m_stage2_ways_tracker->mark(id);
-    } else if (type_name[0] == 'r') {
-        m_stage2_rels_tracker->mark(id);
+        m_stage2_way_ids->set(id);
     }
 
     return 0;
@@ -539,6 +541,12 @@ std::size_t output_flex_t::get_way_nodes()
 
 int output_flex_t::app_get_bbox()
 {
+    if (m_calling_context != calling_context::process_node &&
+        m_calling_context != calling_context::process_way) {
+        throw std::runtime_error{"The function get_bbox() can only be called"
+                                 " from process_node() or process_way()"};
+    }
+
     if (lua_gettop(lua_state()) > 1) {
         throw std::runtime_error{"No parameter(s) needed for get_box()"};
     }
@@ -725,9 +733,9 @@ void output_flex_t::setup_flex_table_columns(flex_table_t *table)
 
 int output_flex_t::app_define_table()
 {
-    if (m_context_node || m_context_way || m_context_relation) {
-        throw std::runtime_error{"Tables have to be defined before calling any "
-                                 "of the process callbacks"};
+    if (m_calling_context != calling_context::main) {
+        throw std::runtime_error{"Database tables have to be defined in the"
+                                 " main Lua code, not in any of the callbacks"};
     }
 
     luaL_checktype(lua_state(), 1, LUA_TTABLE);
@@ -790,8 +798,17 @@ int output_flex_t::table_tostring()
 
 int output_flex_t::table_add_row()
 {
+    if (m_calling_context != calling_context::process_node &&
+        m_calling_context != calling_context::process_way &&
+        m_calling_context != calling_context::process_relation) {
+        throw std::runtime_error{
+            "The function add_row() can only be called from the "
+            "process_node/way/relation() functions"};
+    }
+
     if (lua_gettop(lua_state()) != 2) {
-        throw std::runtime_error{"Need two parameters: The osm2pgsql.table and the row data"};
+        throw std::runtime_error{
+            "Need two parameters: The osm2pgsql.table and the row data"};
     }
 
     auto &table_connection =
@@ -821,14 +838,7 @@ int output_flex_t::table_add_row()
             throw std::runtime_error{
                 "Trying to add relation to table '{}'"_format(table.name())};
         }
-        if (m_in_stage2) {
-            delete_from_table(&table_connection, osmium::item_type::relation,
-                              m_context_relation->id());
-        }
         add_row(&table_connection, *m_context_relation);
-    } else {
-        throw std::runtime_error{"The add_row() function can only be called "
-                                 "from inside a process function"};
     }
 
     return 0;
@@ -1012,16 +1022,21 @@ void output_flex_t::call_lua_function(prepared_lua_function_t func,
 {
     std::lock_guard<std::mutex> guard{lua_mutex};
 
+    m_calling_context = func.context();
+
     lua_pushvalue(lua_state(), func.index()); // the function to call
     push_osm_object_to_lua_stack(
         lua_state(), object,
         get_options()->extra_attributes); // the single argument
 
     luaX_set_context(lua_state(), this);
-    if (luaX_pcall(lua_state(), 1, 0)) {
-        throw std::runtime_error{"Failed to execute lua processing function:"
-                                 " {}"_format(lua_tostring(lua_state(), -1))};
+    if (luaX_pcall(lua_state(), 1, func.nresults())) {
+        throw std::runtime_error{
+            "Failed to execute Lua function 'osm2pgsql.{}':"
+            " {}"_format(func.name(), lua_tostring(lua_state(), -1))};
     }
+
+    m_calling_context = calling_context::main;
 }
 
 void output_flex_t::pending_way(osmid_t id)
@@ -1213,8 +1228,7 @@ output_flex_t::clone(std::shared_ptr<middle_query_t> const &mid,
 {
     return std::make_shared<output_flex_t>(
         mid, *get_options(), copy_thread, true, m_lua_state, m_process_node,
-        m_process_way, m_process_relation, m_tables,
-        m_stage2_ways_tracker, m_stage2_rels_tracker);
+        m_process_way, m_process_relation, m_tables, m_stage2_way_ids);
 }
 
 output_flex_t::output_flex_t(
@@ -1224,11 +1238,9 @@ output_flex_t::output_flex_t(
     prepared_lua_function_t process_way,
     prepared_lua_function_t process_relation,
     std::shared_ptr<std::vector<flex_table_t>> tables,
-    std::shared_ptr<id_tracker> ways_tracker,
-    std::shared_ptr<id_tracker> rels_tracker)
+    std::shared_ptr<idset_t> stage2_way_ids)
 : output_t(mid, o), m_tables(std::move(tables)),
-  m_stage2_ways_tracker(std::move(ways_tracker)),
-  m_stage2_rels_tracker(std::move(rels_tracker)), m_copy_thread(copy_thread),
+  m_stage2_way_ids(std::move(stage2_way_ids)), m_copy_thread(copy_thread),
   m_lua_state(std::move(lua_state)), m_builder(o.projection),
   m_expire(o.expire_tiles_zoom, o.expire_tiles_max_bbox, o.projection),
   m_buffer(32768, osmium::memory::Buffer::auto_grow::yes),
@@ -1323,21 +1335,20 @@ void output_flex_t::init_lua(std::string const &filename)
     // Check whether the process_* functions are available and store them on
     // the Lua stack for fast access later
     lua_getglobal(lua_state(), "osm2pgsql");
-    m_process_node = prepared_lua_function_t{lua_state(), "process_node"};
-    m_process_way = prepared_lua_function_t{lua_state(), "process_way"};
-    m_process_relation =
-        prepared_lua_function_t{lua_state(), "process_relation"};
+    m_process_node = prepared_lua_function_t{
+        lua_state(), calling_context::process_node, "process_node"};
+    m_process_way = prepared_lua_function_t{
+        lua_state(), calling_context::process_way, "process_way"};
+    m_process_relation = prepared_lua_function_t{
+        lua_state(), calling_context::process_relation, "process_relation"};
 
     lua_remove(lua_state(), 1); // global "osm2pgsql"
 }
 
 void output_flex_t::stage2_proc()
 {
-    bool const has_marked_ways = !m_stage2_ways_tracker->empty();
-    bool const has_marked_rels = !m_stage2_rels_tracker->empty();
-
-    if (!has_marked_ways && !has_marked_rels) {
-        fmt::print(stderr, "Skipping stage 2 (no marked objects).\n");
+    if (m_stage2_way_ids->empty()) {
+        fmt::print(stderr, "Skipping stage 2 (no marked ways).\n");
         return;
     }
 
@@ -1349,10 +1360,7 @@ void output_flex_t::stage2_proc()
         util::timer_t timer;
 
         for (auto &table : m_table_connections) {
-            if ((has_marked_ways &&
-                 table.table().matches_type(osmium::item_type::way)) ||
-                (has_marked_rels &&
-                 table.table().matches_type(osmium::item_type::relation))) {
+            if (table.table().matches_type(osmium::item_type::way)) {
                 fmt::print(stderr, "  Creating id index on table '{}'...\n",
                            table.table().name());
                 table.create_id_index();
@@ -1372,32 +1380,19 @@ void output_flex_t::stage2_proc()
     lua_setfield(lua_state(), -2, "stage");
     lua_pop(lua_state(), 1); // osm2pgsql
 
-    osmid_t id;
-
     fmt::print(stderr, "Entering stage 2 processing of {} ways...\n"_format(
-                           m_stage2_ways_tracker->size()));
+                           m_stage2_way_ids->size()));
 
-    while (id_tracker::is_valid((id = m_stage2_ways_tracker->pop_mark()))) {
+    m_stage2_way_ids->sort_unique();
+    for (osmid_t const id : *m_stage2_way_ids) {
         m_buffer.clear();
         if (!m_mid->way_get(id, m_buffer)) {
             continue;
         }
         auto &way = m_buffer.get<osmium::Way>(0);
         way_add(&way);
     }
-
-    fmt::print(stderr,
-               "Entering stage 2 processing of {} relations...\n"_format(
-                   m_stage2_rels_tracker->size()));
-
-    while (id_tracker::is_valid((id = m_stage2_rels_tracker->pop_mark()))) {
-        m_rels_buffer.clear();
-        if (!m_mid->relation_get(id, m_rels_buffer)) {
-            continue;
-        }
-        auto const &relation = m_rels_buffer.get<osmium::Relation>(0);
-        relation_add(relation);
-    }
+    m_stage2_way_ids->clear();
 }
 
 void output_flex_t::merge_expire_trees(output_t *other)

src/output-flex.hpp

@@ -7,12 +7,12 @@
 #include "flex-table.hpp"
 #include "format.hpp"
 #include "geom-transform.hpp"
-#include "id-tracker.hpp"
 #include "osmium-builder.hpp"
 #include "output.hpp"
 #include "table.hpp"
 #include "tagtransform.hpp"
 
+#include <osmium/index/id_set.hpp>
 #include <osmium/osm/item_type.hpp>
 
 extern "C"
@@ -26,6 +26,20 @@ extern "C"
 #include <utility>
 #include <vector>
 
+using idset_t = osmium::index::IdSetSmall<osmid_t>;
+
+/**
+ * When C++ code is called from the Lua code we sometimes need to know
+ * in what context this happens. These are the possible contexts.
+ */
+enum class calling_context
+{
+    main = 0, ///< In main context, i.e. the Lua script outside any callbacks
+    process_node = 1, ///< In the process_node() callback
+    process_way = 2, ///< In the process_way() callback
+    process_relation = 3 ///< In the process_relation() callback
+};
+
 /**
  * The flex output calls several user-defined Lua functions. They are
  * "prepared" by putting the function pointers on the Lua stack. Objects
@@ -41,37 +55,49 @@ class prepared_lua_function_t
     /**
      * Get function with the name "osm2pgsql.name" from Lua and put pointer
      * to it on the Lua stack.
+     *
+     * \param lua_state Current Lua state.
+     * \param name Name of the function.
+     * \param nresults The number of results this function is supposed to have.
      */
-    prepared_lua_function_t(lua_State *lua_state, const char *name);
+    prepared_lua_function_t(lua_State *lua_state, calling_context context,
+                            const char *name, int nresults = 0);
 
     /// Return the index of the function on the Lua stack.
     int index() const noexcept { return m_index; }
 
+    /// The name of the function.
+    char const* name() const noexcept { return m_name; }
+
+    /// The number of results this function is expected to have.
+    int nresults() const noexcept { return m_nresults; }
+
+    calling_context context() const noexcept { return m_calling_context; }
+
     /// Is this function defined in the users Lua code?
     operator bool() const noexcept { return m_index != 0; }
 
 private:
+    char const *m_name = nullptr;
     int m_index = 0;
+    int m_nresults = 0;
+    calling_context m_calling_context = calling_context::main;
 };
 
 class output_flex_t : public output_t
 {
 
 public:
-    output_flex_t(std::shared_ptr<middle_query_t> const &mid,
-                  options_t const &options,
-                  std::shared_ptr<db_copy_thread_t> const &copy_thread,
-                  bool is_clone = false,
-                  std::shared_ptr<lua_State> lua_state = nullptr,
-                  prepared_lua_function_t process_node = {},
-                  prepared_lua_function_t process_way = {},
-                  prepared_lua_function_t process_relation = {},
-                  std::shared_ptr<std::vector<flex_table_t>> tables =
-                      std::make_shared<std::vector<flex_table_t>>(),
-                  std::shared_ptr<id_tracker> ways_tracker =
-                      std::make_shared<id_tracker>(),
-                  std::shared_ptr<id_tracker> rels_tracker =
-                      std::make_shared<id_tracker>());
+    output_flex_t(
+        std::shared_ptr<middle_query_t> const &mid, options_t const &options,
+        std::shared_ptr<db_copy_thread_t> const &copy_thread,
+        bool is_clone = false, std::shared_ptr<lua_State> lua_state = nullptr,
+        prepared_lua_function_t process_node = {},
+        prepared_lua_function_t process_way = {},
+        prepared_lua_function_t process_relation = {},
+        std::shared_ptr<std::vector<flex_table_t>> tables =
+            std::make_shared<std::vector<flex_table_t>>(),
+        std::shared_ptr<idset_t> stage2_way_ids = std::make_shared<idset_t>());
 
     output_flex_t(output_flex_t const &) = delete;
     output_flex_t &operator=(output_flex_t const &) = delete;
@@ -166,8 +192,7 @@ class output_flex_t : public output_t
     std::shared_ptr<std::vector<flex_table_t>> m_tables;
     std::vector<table_connection_t> m_table_connections;
 
-    std::shared_ptr<id_tracker> m_stage2_ways_tracker;
-    std::shared_ptr<id_tracker> m_stage2_rels_tracker;
+    std::shared_ptr<idset_t> m_stage2_way_ids;
 
     std::shared_ptr<db_copy_thread_t> m_copy_thread;
 
@@ -185,10 +210,13 @@ class output_flex_t : public output_t
 
     std::size_t m_num_way_nodes = std::numeric_limits<std::size_t>::max();
 
-    bool m_in_stage2 = false;
     prepared_lua_function_t m_process_node;
     prepared_lua_function_t m_process_way;
     prepared_lua_function_t m_process_relation;
+
+    calling_context m_calling_context = calling_context::main;
+
+    bool m_in_stage2 = false;
 };
 
 #endif // OSM2PGSQL_OUTPUT_FLEX_HPP