<commit_before>/*
 * Copyright (C) 2015 ScyllaDB
 */

/*
 * This file is part of Scylla.
 *
 * Scylla is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Scylla is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Scylla.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <seastar/core/thread.hh>
#include <seastar/util/defer.hh>
#include "core/do_with.hh"
#include "cql_test_env.hh"
#include "cql3/query_processor.hh"
#include "cql3/query_options.hh"
#include "core/distributed.hh"
#include "core/shared_ptr.hh"
#include "utils/UUID_gen.hh"
#include "service/migration_manager.hh"
#include "message/messaging_service.hh"
#include "service/storage_service.hh"
#include "db/config.hh"
#include "db/batchlog_manager.hh"
#include "schema_builder.hh"
#include "tmpdir.hh"
#include "db/query_context.hh"

// TODO: remove (#293)
#include "message/messaging_service.hh"
#include "gms/failure_detector.hh"
#include "gms/gossiper.hh"
#include "service/storage_service.hh"
#include "auth/auth.hh"

static future<> tst_init_ms_fd_gossiper(db::seed_provider_type seed_provider, sstring cluster_name = "Test Cluster") {
    return gms::get_failure_detector().start().then([seed_provider, cluster_name] {
        // Init gossiper
        std::set<gms::inet_address> seeds;
        if (seed_provider.parameters.count("seeds") > 0) {
            size_t begin = 0;
            size_t next = 0;
            sstring seeds_str = seed_provider.parameters.find("seeds")->second;
            while (begin < seeds_str.length() && begin != (next=seeds_str.find(",",begin))) {
                seeds.emplace(gms::inet_address(seeds_str.substr(begin,next-begin)));
                begin = next+1;
            }
        }
        if (seeds.empty()) {
            seeds.emplace(gms::inet_address("127.0.0.1"));
        }
        return gms::get_gossiper().start().then([seeds, cluster_name] {
            auto& gossiper = gms::get_local_gossiper();
            gossiper.set_seeds(seeds);
            gossiper.set_cluster_name(cluster_name);
        });
    });
}
// END TODO

class single_node_cql_env : public cql_test_env {
public:
    static auto constexpr ks_name = "ks";
    static std::atomic<bool> active;
private:
    ::shared_ptr<distributed<database>> _db;
    lw_shared_ptr<tmpdir> _data_dir;
private:
    struct core_local_state {
        service::client_state client_state;

        core_local_state()
            : client_state(service::client_state::for_external_calls()) {
        }

        future<> stop() {
            return make_ready_future<>();
        }
    };
    distributed<core_local_state> _core_local;
private:
    auto make_query_state() {
        if (_db->local().has_keyspace(ks_name)) {
            _core_local.local().client_state.set_keyspace(*_db, ks_name);
        }
        return ::make_shared<service::query_state>(_core_local.local().client_state);
    }
public:
    single_node_cql_env(::shared_ptr<distributed<database>> db) : _db(db)
    { }

    virtual future<::shared_ptr<transport::messages::result_message>> execute_cql(const sstring& text) override {
        auto qs = make_query_state();
        return local_qp().process(text, *qs, cql3::query_options::DEFAULT).finally([qs, this] {
            _core_local.local().client_state.merge(qs->get_client_state());
        });
    }

    virtual future<::shared_ptr<transport::messages::result_message>> execute_cql(
        const sstring& text,
        std::unique_ptr<cql3::query_options> qo) override
    {
        auto qs = make_query_state();
        auto& lqo = *qo;
        return local_qp().process(text, *qs, lqo).finally([qs, qo = std::move(qo), this] {
            _core_local.local().client_state.merge(qs->get_client_state());
        });
    }

    virtual future<bytes> prepare(sstring query) override {
        return qp().invoke_on_all([query, this] (auto& local_qp) {
            auto qs = this->make_query_state();
            return local_qp.prepare(query, *qs).finally([qs] {}).discard_result();
        }).then([query, this] {
            return local_qp().compute_id(query, ks_name);
        });
    }

    virtual future<::shared_ptr<transport::messages::result_message>> execute_prepared(
        bytes id,
        std::vector<bytes_opt> values) override
    {
        auto prepared = local_qp().get_prepared(id);
        if (!prepared) {
            throw not_prepared_exception(id);
        }
        auto stmt = prepared->statement;
        assert(stmt->get_bound_terms() == values.size());

        auto options = ::make_shared<cql3::query_options>(std::move(values));
        options->prepare(prepared->bound_names);

        auto qs = make_query_state();
        return local_qp().process_statement(stmt, *qs, *options)
            .finally([options, qs, this] {
                _core_local.local().client_state.merge(qs->get_client_state());
            });
    }

    virtual future<> create_table(std::function<schema(const sstring&)> schema_maker) override {
        auto id = utils::UUID_gen::get_time_UUID();
        schema_builder builder(make_lw_shared(schema_maker(ks_name)));
        builder.set_uuid(id);
        auto s = builder.build(schema_builder::compact_storage::no);
        return service::get_local_migration_manager().announce_new_column_family(s, true);
    }

    virtual future<> require_keyspace_exists(const sstring& ks_name) override {
        auto& db = _db->local();
        assert(db.has_keyspace(ks_name));
        return make_ready_future<>();
    }

    virtual future<> require_table_exists(const sstring& ks_name, const sstring& table_name) override {
        auto& db = _db->local();
        assert(db.has_schema(ks_name, table_name));
        return make_ready_future<>();
    }

    virtual future<> require_column_has_value(const sstring& table_name,
                                      std::vector<data_value> pk,
                                      std::vector<data_value> ck,
                                      const sstring& column_name,
                                      data_value expected) override {
        auto& db = _db->local();
        auto& cf = db.find_column_family(ks_name, table_name);
        auto schema = cf.schema();
        auto pkey = partition_key::from_deeply_exploded(*schema, pk);
        auto ckey = clustering_key::from_deeply_exploded(*schema, ck);
        auto exp = expected.type()->decompose(expected);
        auto dk = dht::global_partitioner().decorate_key(*schema, pkey);
        auto shard = db.shard_of(dk._token);
        return _db->invoke_on(shard, [pkey = std::move(pkey),
                                      ckey = std::move(ckey),
                                      ks_name = std::move(ks_name),
                                      column_name = std::move(column_name),
                                      exp = std::move(exp),
                                      table_name = std::move(table_name)] (database& db) mutable {
          auto& cf = db.find_column_family(ks_name, table_name);
          auto schema = cf.schema();
          return cf.find_partition_slow(schema, pkey)
                  .then([schema, ckey, column_name, exp] (column_family::const_mutation_partition_ptr p) {
            assert(p != nullptr);
            auto row = p->find_row(ckey);
            assert(row != nullptr);
            auto col_def = schema->get_column_definition(utf8_type->decompose(column_name));
            assert(col_def != nullptr);
            const atomic_cell_or_collection* cell = row->find_cell(col_def->id);
            if (!cell) {
                assert(((void)"column not set", 0));
            }
            bytes actual;
            if (!col_def->type->is_multi_cell()) {
                auto c = cell->as_atomic_cell();
                assert(c.is_live());
                actual = { c.value().begin(), c.value().end() };
            } else {
                auto c = cell->as_collection_mutation();
                auto type = dynamic_pointer_cast<const collection_type_impl>(col_def->type);
                actual = type->to_value(type->deserialize_mutation_form(c),
                                        cql_serialization_format::internal());
            }
            assert(col_def->type->equal(actual, exp));
          });
        });
    }

    virtual database& local_db() override {
        return _db->local();
    }

    cql3::query_processor& local_qp() override {
        return cql3::get_local_query_processor();
    }

    distributed<database>& db() override {
        return *_db;
    }

    distributed<cql3::query_processor>& qp() override {
        return cql3::get_query_processor();
    }

    future<> start() {
        return _core_local.start();
    }

    future<> stop() {
        return _core_local.stop();
    }

    static future<> do_with(std::function<future<>(cql_test_env&)> func, const db::config& cfg_in) {
        return seastar::async([cfg_in, func] {
            bool old_active = false;
            if (!active.compare_exchange_strong(old_active, true)) {
                throw std::runtime_error("Starting more than one cql_test_env at a time not supported due to singletons.");
            }
            auto deactivate = defer([] {
                bool old_active = true;
                auto success = active.compare_exchange_strong(old_active, false);
                assert(success);
            });

            utils::fb_utilities::set_broadcast_address(gms::inet_address("localhost"));
            utils::fb_utilities::set_broadcast_rpc_address(gms::inet_address("localhost"));
            locator::i_endpoint_snitch::create_snitch("SimpleSnitch").get();
            auto stop_snitch = defer([] { locator::i_endpoint_snitch::stop_snitch().get(); });

            auto wait_for_background_jobs = defer([] { sstables::await_background_jobs_on_all_shards().get(); });

            auto db = ::make_shared<distributed<database>>();
            auto cfg = make_lw_shared<db::config>(std::move(cfg_in));
            tmpdir data_dir;
            cfg->data_file_directories() = { data_dir.path };
            cfg->commitlog_directory() = data_dir.path + "/commitlog.dir";
            cfg->num_tokens() = 256;
            cfg->ring_delay_ms() = 500;
            cfg->experimental() = true;
            cfg->shutdown_announce_in_ms() = 0;
            boost::filesystem::create_directories((data_dir.path + "/system").c_str());
            boost::filesystem::create_directories(cfg->commitlog_directory().c_str());

            const gms::inet_address listen("127.0.0.1");
            auto& ms = net::get_messaging_service();
            ms.start(listen, std::move(7000)).get();
            auto stop_ms = defer([&ms] { ms.stop().get(); });

            auto& ss = service::get_storage_service();
            ss.start(std::ref(*db));
            auto stop_storage_service = defer([&ss] { ss.stop().get(); });

            db->start(std::move(*cfg)).get();
            auto stop_db = defer([db] { db->stop().get(); });

            // FIXME: split
            tst_init_ms_fd_gossiper(db::config::seed_provider_type()).get();
            auto stop_ms_fd_gossiper = defer([] {
                gms::get_gossiper().stop().get();
                gms::get_failure_detector().stop().get();
            });

            distributed<service::storage_proxy>& proxy = service::get_storage_proxy();
            distributed<service::migration_manager>& mm = service::get_migration_manager();
            distributed<db::batchlog_manager>& bm = db::get_batchlog_manager();

            proxy.start(std::ref(*db)).get();
            auto stop_proxy = defer([&proxy] { proxy.stop().get(); });

            mm.start().get();
            auto stop_mm = defer([&mm] { mm.stop().get(); });

            auto& qp = cql3::get_query_processor();
            qp.start(std::ref(proxy), std::ref(*db)).get();
            auto stop_qp = defer([&qp] { qp.stop().get(); });

            bm.start(std::ref(qp)).get();
            auto stop_bm = defer([&bm] { bm.stop().get(); });

            db->invoke_on_all([] (database& db) {
                return db.init_system_keyspace();
            }).get();

            auto& ks = db->local().find_keyspace(db::system_keyspace::NAME);
            parallel_for_each(ks.metadata()->cf_meta_data(), [&ks] (auto& pair) {
                auto cfm = pair.second;
                return ks.make_directory_for_column_family(cfm->cf_name(), cfm->id());
            }).get();

            // In main.cc we call db::system_keyspace::setup which calls
            // minimal_setup and init_local_cache
            db::system_keyspace::minimal_setup(*db, qp);
            auto stop_system_keyspace = defer([] { db::qctx = {}; });

            db::system_keyspace::init_local_cache().get();
            auto stop_local_cache = defer([] { db::system_keyspace::deinit_local_cache().get(); });

            service::get_local_storage_service().init_server().get();
            auto deinit_storage_service_server = defer([] {
                gms::get_local_gossiper().stop_gossiping().get();
                auth::auth::shutdown().get();
            });

            single_node_cql_env env(db);
            env.start().get();
            auto stop_env = defer([&env] { env.stop().get(); });

            auto query = sprint("create keyspace %s with replication = { 'class' : 'org.apache.cassandra.locator.SimpleStrategy', 'replication_factor' : 1 };", sstring{ks_name});
            env.execute_cql(query).get();

            func(env).get();
        });
    }
};

std::atomic<bool> single_node_cql_env::active = { false };

future<> do_with_cql_env(std::function<future<>(cql_test_env&)> func, const db::config& cfg_in) {
    return single_node_cql_env::do_with(func, cfg_in);
}

future<> do_with_cql_env(std::function<future<>(cql_test_env&)> func) {
    return do_with_cql_env(std::move(func), db::config{});
}
<commit_msg>tests::cql_test_env: Fake logged in user in case test requires is.<commit_after>/*
 * Copyright (C) 2015 ScyllaDB
 */

/*
 * This file is part of Scylla.
 *
 * Scylla is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Scylla is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Scylla.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <seastar/core/thread.hh>
#include <seastar/util/defer.hh>
#include "core/do_with.hh"
#include "cql_test_env.hh"
#include "cql3/query_processor.hh"
#include "cql3/query_options.hh"
#include "core/distributed.hh"
#include "core/shared_ptr.hh"
#include "utils/UUID_gen.hh"
#include "service/migration_manager.hh"
#include "message/messaging_service.hh"
#include "service/storage_service.hh"
#include "db/config.hh"
#include "db/batchlog_manager.hh"
#include "schema_builder.hh"
#include "tmpdir.hh"
#include "db/query_context.hh"

// TODO: remove (#293)
#include "message/messaging_service.hh"
#include "gms/failure_detector.hh"
#include "gms/gossiper.hh"
#include "service/storage_service.hh"
#include "auth/auth.hh"

static future<> tst_init_ms_fd_gossiper(db::seed_provider_type seed_provider, sstring cluster_name = "Test Cluster") {
    return gms::get_failure_detector().start().then([seed_provider, cluster_name] {
        // Init gossiper
        std::set<gms::inet_address> seeds;
        if (seed_provider.parameters.count("seeds") > 0) {
            size_t begin = 0;
            size_t next = 0;
            sstring seeds_str = seed_provider.parameters.find("seeds")->second;
            while (begin < seeds_str.length() && begin != (next=seeds_str.find(",",begin))) {
                seeds.emplace(gms::inet_address(seeds_str.substr(begin,next-begin)));
                begin = next+1;
            }
        }
        if (seeds.empty()) {
            seeds.emplace(gms::inet_address("127.0.0.1"));
        }
        return gms::get_gossiper().start().then([seeds, cluster_name] {
            auto& gossiper = gms::get_local_gossiper();
            gossiper.set_seeds(seeds);
            gossiper.set_cluster_name(cluster_name);
        });
    });
}
// END TODO

class single_node_cql_env : public cql_test_env {
public:
    static auto constexpr ks_name = "ks";
    static std::atomic<bool> active;
private:
    ::shared_ptr<distributed<database>> _db;
    lw_shared_ptr<tmpdir> _data_dir;
private:
    struct core_local_state {
        service::client_state client_state;

        core_local_state()
            : client_state(service::client_state::for_external_calls())
        {
            client_state.set_login(::make_shared<auth::authenticated_user>("cassandra"));
        }

        future<> stop() {
            return make_ready_future<>();
        }
    };
    distributed<core_local_state> _core_local;
private:
    auto make_query_state() {
        if (_db->local().has_keyspace(ks_name)) {
            _core_local.local().client_state.set_keyspace(*_db, ks_name);
        }
        return ::make_shared<service::query_state>(_core_local.local().client_state);
    }
public:
    single_node_cql_env(::shared_ptr<distributed<database>> db) : _db(db)
    { }

    virtual future<::shared_ptr<transport::messages::result_message>> execute_cql(const sstring& text) override {
        auto qs = make_query_state();
        return local_qp().process(text, *qs, cql3::query_options::DEFAULT).finally([qs, this] {
            _core_local.local().client_state.merge(qs->get_client_state());
        });
    }

    virtual future<::shared_ptr<transport::messages::result_message>> execute_cql(
        const sstring& text,
        std::unique_ptr<cql3::query_options> qo) override
    {
        auto qs = make_query_state();
        auto& lqo = *qo;
        return local_qp().process(text, *qs, lqo).finally([qs, qo = std::move(qo), this] {
            _core_local.local().client_state.merge(qs->get_client_state());
        });
    }

    virtual future<bytes> prepare(sstring query) override {
        return qp().invoke_on_all([query, this] (auto& local_qp) {
            auto qs = this->make_query_state();
            return local_qp.prepare(query, *qs).finally([qs] {}).discard_result();
        }).then([query, this] {
            return local_qp().compute_id(query, ks_name);
        });
    }

    virtual future<::shared_ptr<transport::messages::result_message>> execute_prepared(
        bytes id,
        std::vector<bytes_opt> values) override
    {
        auto prepared = local_qp().get_prepared(id);
        if (!prepared) {
            throw not_prepared_exception(id);
        }
        auto stmt = prepared->statement;
        assert(stmt->get_bound_terms() == values.size());

        auto options = ::make_shared<cql3::query_options>(std::move(values));
        options->prepare(prepared->bound_names);

        auto qs = make_query_state();
        return local_qp().process_statement(stmt, *qs, *options)
            .finally([options, qs, this] {
                _core_local.local().client_state.merge(qs->get_client_state());
            });
    }

    virtual future<> create_table(std::function<schema(const sstring&)> schema_maker) override {
        auto id = utils::UUID_gen::get_time_UUID();
        schema_builder builder(make_lw_shared(schema_maker(ks_name)));
        builder.set_uuid(id);
        auto s = builder.build(schema_builder::compact_storage::no);
        return service::get_local_migration_manager().announce_new_column_family(s, true);
    }

    virtual future<> require_keyspace_exists(const sstring& ks_name) override {
        auto& db = _db->local();
        assert(db.has_keyspace(ks_name));
        return make_ready_future<>();
    }

    virtual future<> require_table_exists(const sstring& ks_name, const sstring& table_name) override {
        auto& db = _db->local();
        assert(db.has_schema(ks_name, table_name));
        return make_ready_future<>();
    }

    virtual future<> require_column_has_value(const sstring& table_name,
                                      std::vector<data_value> pk,
                                      std::vector<data_value> ck,
                                      const sstring& column_name,
                                      data_value expected) override {
        auto& db = _db->local();
        auto& cf = db.find_column_family(ks_name, table_name);
        auto schema = cf.schema();
        auto pkey = partition_key::from_deeply_exploded(*schema, pk);
        auto ckey = clustering_key::from_deeply_exploded(*schema, ck);
        auto exp = expected.type()->decompose(expected);
        auto dk = dht::global_partitioner().decorate_key(*schema, pkey);
        auto shard = db.shard_of(dk._token);
        return _db->invoke_on(shard, [pkey = std::move(pkey),
                                      ckey = std::move(ckey),
                                      ks_name = std::move(ks_name),
                                      column_name = std::move(column_name),
                                      exp = std::move(exp),
                                      table_name = std::move(table_name)] (database& db) mutable {
          auto& cf = db.find_column_family(ks_name, table_name);
          auto schema = cf.schema();
          return cf.find_partition_slow(schema, pkey)
                  .then([schema, ckey, column_name, exp] (column_family::const_mutation_partition_ptr p) {
            assert(p != nullptr);
            auto row = p->find_row(ckey);
            assert(row != nullptr);
            auto col_def = schema->get_column_definition(utf8_type->decompose(column_name));
            assert(col_def != nullptr);
            const atomic_cell_or_collection* cell = row->find_cell(col_def->id);
            if (!cell) {
                assert(((void)"column not set", 0));
            }
            bytes actual;
            if (!col_def->type->is_multi_cell()) {
                auto c = cell->as_atomic_cell();
                assert(c.is_live());
                actual = { c.value().begin(), c.value().end() };
            } else {
                auto c = cell->as_collection_mutation();
                auto type = dynamic_pointer_cast<const collection_type_impl>(col_def->type);
                actual = type->to_value(type->deserialize_mutation_form(c),
                                        cql_serialization_format::internal());
            }
            assert(col_def->type->equal(actual, exp));
          });
        });
    }

    virtual database& local_db() override {
        return _db->local();
    }

    cql3::query_processor& local_qp() override {
        return cql3::get_local_query_processor();
    }

    distributed<database>& db() override {
        return *_db;
    }

    distributed<cql3::query_processor>& qp() override {
        return cql3::get_query_processor();
    }

    future<> start() {
        return _core_local.start();
    }

    future<> stop() {
        return _core_local.stop();
    }

    static future<> do_with(std::function<future<>(cql_test_env&)> func, const db::config& cfg_in) {
        return seastar::async([cfg_in, func] {
            bool old_active = false;
            if (!active.compare_exchange_strong(old_active, true)) {
                throw std::runtime_error("Starting more than one cql_test_env at a time not supported due to singletons.");
            }
            auto deactivate = defer([] {
                bool old_active = true;
                auto success = active.compare_exchange_strong(old_active, false);
                assert(success);
            });

            utils::fb_utilities::set_broadcast_address(gms::inet_address("localhost"));
            utils::fb_utilities::set_broadcast_rpc_address(gms::inet_address("localhost"));
            locator::i_endpoint_snitch::create_snitch("SimpleSnitch").get();
            auto stop_snitch = defer([] { locator::i_endpoint_snitch::stop_snitch().get(); });

            auto wait_for_background_jobs = defer([] { sstables::await_background_jobs_on_all_shards().get(); });

            auto db = ::make_shared<distributed<database>>();
            auto cfg = make_lw_shared<db::config>(std::move(cfg_in));
            tmpdir data_dir;
            cfg->data_file_directories() = { data_dir.path };
            cfg->commitlog_directory() = data_dir.path + "/commitlog.dir";
            cfg->num_tokens() = 256;
            cfg->ring_delay_ms() = 500;
            cfg->experimental() = true;
            cfg->shutdown_announce_in_ms() = 0;
            boost::filesystem::create_directories((data_dir.path + "/system").c_str());
            boost::filesystem::create_directories(cfg->commitlog_directory().c_str());

            const gms::inet_address listen("127.0.0.1");
            auto& ms = net::get_messaging_service();
            ms.start(listen, std::move(7000)).get();
            auto stop_ms = defer([&ms] { ms.stop().get(); });

            auto& ss = service::get_storage_service();
            ss.start(std::ref(*db));
            auto stop_storage_service = defer([&ss] { ss.stop().get(); });

            db->start(std::move(*cfg)).get();
            auto stop_db = defer([db] { db->stop().get(); });

            // FIXME: split
            tst_init_ms_fd_gossiper(db::config::seed_provider_type()).get();
            auto stop_ms_fd_gossiper = defer([] {
                gms::get_gossiper().stop().get();
                gms::get_failure_detector().stop().get();
            });

            distributed<service::storage_proxy>& proxy = service::get_storage_proxy();
            distributed<service::migration_manager>& mm = service::get_migration_manager();
            distributed<db::batchlog_manager>& bm = db::get_batchlog_manager();

            proxy.start(std::ref(*db)).get();
            auto stop_proxy = defer([&proxy] { proxy.stop().get(); });

            mm.start().get();
            auto stop_mm = defer([&mm] { mm.stop().get(); });

            auto& qp = cql3::get_query_processor();
            qp.start(std::ref(proxy), std::ref(*db)).get();
            auto stop_qp = defer([&qp] { qp.stop().get(); });

            bm.start(std::ref(qp)).get();
            auto stop_bm = defer([&bm] { bm.stop().get(); });

            db->invoke_on_all([] (database& db) {
                return db.init_system_keyspace();
            }).get();

            auto& ks = db->local().find_keyspace(db::system_keyspace::NAME);
            parallel_for_each(ks.metadata()->cf_meta_data(), [&ks] (auto& pair) {
                auto cfm = pair.second;
                return ks.make_directory_for_column_family(cfm->cf_name(), cfm->id());
            }).get();

            // In main.cc we call db::system_keyspace::setup which calls
            // minimal_setup and init_local_cache
            db::system_keyspace::minimal_setup(*db, qp);
            auto stop_system_keyspace = defer([] { db::qctx = {}; });

            db::system_keyspace::init_local_cache().get();
            auto stop_local_cache = defer([] { db::system_keyspace::deinit_local_cache().get(); });

            service::get_local_storage_service().init_server().get();
            auto deinit_storage_service_server = defer([] {
                gms::get_local_gossiper().stop_gossiping().get();
                auth::auth::shutdown().get();
            });

            single_node_cql_env env(db);
            env.start().get();
            auto stop_env = defer([&env] { env.stop().get(); });

            auto query = sprint("create keyspace %s with replication = { 'class' : 'org.apache.cassandra.locator.SimpleStrategy', 'replication_factor' : 1 };", sstring{ks_name});
            env.execute_cql(query).get();

            func(env).get();
        });
    }
};

std::atomic<bool> single_node_cql_env::active = { false };

future<> do_with_cql_env(std::function<future<>(cql_test_env&)> func, const db::config& cfg_in) {
    return single_node_cql_env::do_with(func, cfg_in);
}

future<> do_with_cql_env(std::function<future<>(cql_test_env&)> func) {
    return do_with_cql_env(std::move(func), db::config{});
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2008, 2009 Nokia Corporation.
 *
 * Contact: Marius Vollmer <marius.vollmer@nokia.com>
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public License
 * version 2.1 as published by the Free Software Foundation.
 *
 * This library is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
 * 02110-1301 USA
 *
 */

#include <QtTest/QtTest>
#include <QtCore>
#include "contextpropertyinfo.h"
#include "infobackend.h"

/* Mocked infobackend */

InfoBackend* currentBackend = NULL;

InfoBackend* InfoBackend::instance(const QString &backendName)
{
    if (currentBackend)
        return currentBackend;
    else {
        currentBackend = new InfoBackend();
        return currentBackend;
    }
}

QStringList InfoBackend::listKeys() const
{
    QStringList l;
    l << QString("Battery.Charging");
    l << QString("Media.NowPlaying");
    return l;
}

QStringList InfoBackend::listKeysForPlugin(QString plugin) const
{
    if (plugin == "contextkit-dbus") {
        QStringList l;
        l << QString("Battery.Charging");
        l << QString("Media.NowPlaying");
        return l;
    } else
        return QStringList();
}

QStringList InfoBackend::listPlugins() const
{
    QStringList l;
    l << QString("contextkit-dbus");
    return l;
}

QString InfoBackend::constructionStringForKey(QString key) const
{
    if (key == "Battery.Charging")
        return "system:org.freedesktop.ContextKit.contextd";
    else if (key == "Media.NowPlaying")
        return "system:com.nokia.musicplayer";
    return QString();
}

QString InfoBackend::typeForKey(QString key) const
{
    if (key == "Battery.Charging")
        return "TRUTH";
    else if (key == "Media.NowPlaying")
        return "STRING";
    else
        return "";
}

QString InfoBackend::docForKey(QString key) const
{
    if (key == "Battery.Charging")
        return "Battery.Charging doc";
    else if (key == "Media.NowPlaying")
        return "Media.NowPlaying doc";
    else
        return "";
}

QString InfoBackend::pluginForKey(QString key) const
{
    if (key == "Battery.Charging")
        return "contextkit-dbus";
    else if (key == "Media.NowPlaying")
        return "contextkit-dbus";
    else
        return "";
}

bool InfoBackend::keyExists(QString key) const
{
    if (key == "Battery.Charging")
        return true;
    else if (key == "Media.NowPlaying")
        return true;
    else
        return false;
}

bool InfoBackend::keyProvided(QString key) const
{
    if (key == "Battery.Charging")
        return true;
    else if (key == "Media.NowPlaying")
        return true;
    else
        return false;
}

void InfoBackend::connectNotify(const char *signal)
{
}

void InfoBackend::disconnectNotify(const char *signal)
{
}

/*
void InfoBackend::fireKeysChanged(const QStringList& keys)
{
    emit keysChanged(keys);
}

void InfoBackend::fireKeysAdded(const QStringList& keys)
{
    emit keysAdded(keys);
}

void InfoBackend::fireKeysRemoved(const QStringList& keys)
{
    emit keysRemoved(keys);
}
*/

/* ContextRegistryInfoUnitTest */

class ContextPropertyInfoUnitTest : public QObject
{
    Q_OBJECT

private slots:
    void initTestCase();
    void key();
    void doc();
    void type();
    void exists();
    void provided();
    void providerDBusName();
};

void ContextPropertyInfoUnitTest::initTestCase()
{
}

void ContextPropertyInfoUnitTest::key()
{
    ContextPropertyInfo p("Battery.Charging");
    QCOMPARE(p.key(), QString("Battery.Charging"));
}

void ContextPropertyInfoUnitTest::doc()
{
    ContextPropertyInfo p1("Battery.Charging");
    ContextPropertyInfo p2("Media.NowPlaying");
    ContextPropertyInfo p3("Does.Not.Exist");
    QCOMPARE(p1.doc(), QString("Battery.Charging doc"));
    QCOMPARE(p2.doc(), QString("Media.NowPlaying doc"));
    QCOMPARE(p3.doc(), QString());
}

void ContextPropertyInfoUnitTest::type()
{
    ContextPropertyInfo p1("Battery.Charging");
    ContextPropertyInfo p2("Media.NowPlaying");
    ContextPropertyInfo p3("Does.Not.Exist");
    QCOMPARE(p1.type(), QString("TRUTH"));
    QCOMPARE(p2.type(), QString("STRING"));
    QCOMPARE(p3.type(), QString());
}

void ContextPropertyInfoUnitTest::exists()
{
    ContextPropertyInfo p1("Battery.Charging");
    ContextPropertyInfo p2("Media.NowPlaying");
    ContextPropertyInfo p3("Does.Not.Exist");
    QCOMPARE(p1.exists(), true);
    QCOMPARE(p2.exists(), true);
    QCOMPARE(p3.exists(), false);
}

void ContextPropertyInfoUnitTest::provided()
{
    ContextPropertyInfo p1("Battery.Charging");
    ContextPropertyInfo p2("Media.NowPlaying");
    ContextPropertyInfo p3("Does.Not.Exist");
    QCOMPARE(p1.provided(), true);
    QCOMPARE(p2.provided(), true);
    QCOMPARE(p3.provided(), false);
}

void ContextPropertyInfoUnitTest::providerDBusName()
{
    ContextPropertyInfo p1("Battery.Charging");
    ContextPropertyInfo p2("Media.NowPlaying");
    ContextPropertyInfo p3("Does.Not.Exist");
    QCOMPARE(p1.providerDBusName(), QString("org.freedesktop.ContextKit.contextd"));
    QCOMPARE(p2.providerDBusName(), QString("com.nokia.musicplayer"));
    QCOMPARE(p3.providerDBusName(), QString());
}

#include "contextpropertyinfounittest.moc"
QTEST_MAIN(ContextPropertyInfoUnitTest);
<commit_msg>DBusConnectionType.<commit_after>/*
 * Copyright (C) 2008, 2009 Nokia Corporation.
 *
 * Contact: Marius Vollmer <marius.vollmer@nokia.com>
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public License
 * version 2.1 as published by the Free Software Foundation.
 *
 * This library is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
 * 02110-1301 USA
 *
 */

#include <QtTest/QtTest>
#include <QtCore>
#include "contextpropertyinfo.h"
#include "infobackend.h"

/* Mocked infobackend */

InfoBackend* currentBackend = NULL;

InfoBackend* InfoBackend::instance(const QString &backendName)
{
    if (currentBackend)
        return currentBackend;
    else {
        currentBackend = new InfoBackend();
        return currentBackend;
    }
}

QStringList InfoBackend::listKeys() const
{
    QStringList l;
    l << QString("Battery.Charging");
    l << QString("Media.NowPlaying");
    return l;
}

QStringList InfoBackend::listKeysForPlugin(QString plugin) const
{
    if (plugin == "contextkit-dbus") {
        QStringList l;
        l << QString("Battery.Charging");
        l << QString("Media.NowPlaying");
        return l;
    } else
        return QStringList();
}

QStringList InfoBackend::listPlugins() const
{
    QStringList l;
    l << QString("contextkit-dbus");
    return l;
}

QString InfoBackend::constructionStringForKey(QString key) const
{
    if (key == "Battery.Charging")
        return "system:org.freedesktop.ContextKit.contextd";
    else if (key == "Media.NowPlaying")
        return "session:com.nokia.musicplayer";
    return QString();
}

QString InfoBackend::typeForKey(QString key) const
{
    if (key == "Battery.Charging")
        return "TRUTH";
    else if (key == "Media.NowPlaying")
        return "STRING";
    else
        return "";
}

QString InfoBackend::docForKey(QString key) const
{
    if (key == "Battery.Charging")
        return "Battery.Charging doc";
    else if (key == "Media.NowPlaying")
        return "Media.NowPlaying doc";
    else
        return "";
}

QString InfoBackend::pluginForKey(QString key) const
{
    if (key == "Battery.Charging")
        return "contextkit-dbus";
    else if (key == "Media.NowPlaying")
        return "contextkit-dbus";
    else
        return "";
}

bool InfoBackend::keyExists(QString key) const
{
    if (key == "Battery.Charging")
        return true;
    else if (key == "Media.NowPlaying")
        return true;
    else
        return false;
}

bool InfoBackend::keyProvided(QString key) const
{
    if (key == "Battery.Charging")
        return true;
    else if (key == "Media.NowPlaying")
        return true;
    else
        return false;
}

void InfoBackend::connectNotify(const char *signal)
{
}

void InfoBackend::disconnectNotify(const char *signal)
{
}

/*
void InfoBackend::fireKeysChanged(const QStringList& keys)
{
    emit keysChanged(keys);
}

void InfoBackend::fireKeysAdded(const QStringList& keys)
{
    emit keysAdded(keys);
}

void InfoBackend::fireKeysRemoved(const QStringList& keys)
{
    emit keysRemoved(keys);
}
*/

/* ContextRegistryInfoUnitTest */

class ContextPropertyInfoUnitTest : public QObject
{
    Q_OBJECT

private slots:
    void initTestCase();
    void key();
    void doc();
    void type();
    void exists();
    void provided();
    void providerDBusName();
    void providerDBusType();
};

void ContextPropertyInfoUnitTest::initTestCase()
{
}

void ContextPropertyInfoUnitTest::key()
{
    ContextPropertyInfo p("Battery.Charging");
    QCOMPARE(p.key(), QString("Battery.Charging"));
}

void ContextPropertyInfoUnitTest::doc()
{
    ContextPropertyInfo p1("Battery.Charging");
    ContextPropertyInfo p2("Media.NowPlaying");
    ContextPropertyInfo p3("Does.Not.Exist");
    QCOMPARE(p1.doc(), QString("Battery.Charging doc"));
    QCOMPARE(p2.doc(), QString("Media.NowPlaying doc"));
    QCOMPARE(p3.doc(), QString());
}

void ContextPropertyInfoUnitTest::type()
{
    ContextPropertyInfo p1("Battery.Charging");
    ContextPropertyInfo p2("Media.NowPlaying");
    ContextPropertyInfo p3("Does.Not.Exist");
    QCOMPARE(p1.type(), QString("TRUTH"));
    QCOMPARE(p2.type(), QString("STRING"));
    QCOMPARE(p3.type(), QString());
}

void ContextPropertyInfoUnitTest::exists()
{
    ContextPropertyInfo p1("Battery.Charging");
    ContextPropertyInfo p2("Media.NowPlaying");
    ContextPropertyInfo p3("Does.Not.Exist");
    QCOMPARE(p1.exists(), true);
    QCOMPARE(p2.exists(), true);
    QCOMPARE(p3.exists(), false);
}

void ContextPropertyInfoUnitTest::provided()
{
    ContextPropertyInfo p1("Battery.Charging");
    ContextPropertyInfo p2("Media.NowPlaying");
    ContextPropertyInfo p3("Does.Not.Exist");
    QCOMPARE(p1.provided(), true);
    QCOMPARE(p2.provided(), true);
    QCOMPARE(p3.provided(), false);
}

void ContextPropertyInfoUnitTest::providerDBusName()
{
    ContextPropertyInfo p1("Battery.Charging");
    ContextPropertyInfo p2("Media.NowPlaying");
    ContextPropertyInfo p3("Does.Not.Exist");
    QCOMPARE(p1.providerDBusName(), QString("org.freedesktop.ContextKit.contextd"));
    QCOMPARE(p2.providerDBusName(), QString("com.nokia.musicplayer"));
    QCOMPARE(p3.providerDBusName(), QString());
}

void ContextPropertyInfoUnitTest::providerDBusType()
{
    ContextPropertyInfo p1("Battery.Charging");
    ContextPropertyInfo p2("Media.NowPlaying");
    ContextPropertyInfo p3("Does.Not.Exist");
    QCOMPARE(p1.providerDBusType(), QDBusConnection::SystemBus);
    QCOMPARE(p2.providerDBusType(), QDBusConnection::SessionBus);
    QCOMPARE(p3.providerDBusType(), QDBusConnection::SessionBus);
}

#include "contextpropertyinfounittest.moc"
QTEST_MAIN(ContextPropertyInfoUnitTest);
<|endoftext|>


<commit_before>#ifndef SPLAY_TREE_H_
#define SPLAY_TREE_H_

#include <utility>
#include <iostream>
#include <stack>

namespace dicts {

template<typename K, typename V>
class SplayDict {
  private:
    struct Node; //Forward declaration

  public:
    class ConstIterator {
      public:
        const std::pair<K&, V&> operator*() const;

        ConstIterator& operator++ (int);

        ConstIterator& operator= (const ConstIterator& const_it);

        bool operator== (const ConstIterator& const_it) const;

        bool operator!= (const ConstIterator& const_it) const;

        ConstIterator(Node* root){
          if (root != nullptr){
            current = subtreeMin(root);
            prev = current->parent;
          }
          else {
            current = root;
            prev = nullptr;
          }
        }

        ConstIterator(Node* c, Node* p){
          current = c;
          prev = p;
        }

      private:
        Node* current;
        Node* prev;
    };

    SplayDict();

    SplayDict(const SplayDict& bst);

    ~SplayDict();

    bool exists(const K& key) const;

    V& get(const K& key);

    void set(const K& key, const V& val);

    void erase(const K& key);

    SplayDict& operator= (const SplayDict& bst);

    ConstIterator begin();

    ConstIterator end();

  private:
    struct Node {
      K key;
      V val;

      Node* parent;
      Node* left;
      Node* right;

      Node(Node* p, const K& k, const V& v) {
        key = k;
        val = v;

        parent  = p;
        left    = nullptr;
        right   = nullptr;
      }

    };

    static void printNode(std::ostream& os, Node* node, std::string prefix, bool isTail);

    Node* findNode(Node* subtree, const K& key) const;

    bool hasGrandParent(Node* subtree) const;

    void rightRotate(Node* n);

    void leftRotate(Node* n);

    bool sideOfChild(Node* subtree, bool zag1, bool zag2) const;

    Node* splay(Node* subtree);

    void set(Node*& subtree, Node* parent, const K& key, const V& val);

    static Node* subtreeMin(Node* subtree);

    static Node* subtreeMax(Node* subtree);

    static Node* subtreeMaxKey(Node* subtree, Node* parent);

    static Node* subtreeMinKey(Node* subtree, Node* parent);

    void erase_node(Node* node);

    Node* root;

  friend class ConstIterator;

  friend std::ostream& operator<< (std::ostream& os, const SplayDict<K,V>& spDict){
    printNode(os, spDict.root, "", true);
    return os;
  }

};

template<typename K, typename V>
SplayDict<K, V>::SplayDict() {
    root = nullptr;
}

template<typename K, typename V>
SplayDict<K, V>::~SplayDict() {
  std::stack<Node*> nodes;
  nodes.push(root);

  while(!nodes.empty()){
    Node* current = nodes.top();
    nodes.pop();

    if (current != nullptr){
      nodes.push(current->left);
      nodes.push(current->right);

      delete current;
    }

  }
}

template<typename K, typename V>
SplayDict<K, V>::SplayDict(const SplayDict& bst) {
  for (ConstIterator it = bst.begin(); it != bst.end(); it++) {
    set( (*it).first, (*it).second);
  }
}

template<typename K, typename V>
bool SplayDict<K, V>::exists(const K& key) const {
  return findNode(root, key) != nullptr;
}

template<typename K, typename V>
V& SplayDict<K, V>::get(const K& key){
  Node* node = findNode(root, key);
  root = splay(node);
  return root->val;
}

template<typename K, typename V>
void SplayDict<K, V>::set(const K& key, const V& val) {
  set(root, nullptr, key, val);
}

template<typename K, typename V>
void SplayDict<K, V>::erase(const K& key){
  Node* node = findNode(root, key);

  erase_node(node);
}

template<typename K, typename V>
typename SplayDict<K,V>::ConstIterator SplayDict<K, V>::begin(){
  ConstIterator it = ConstIterator(root);
  return it;
}

template<typename K, typename V>
typename SplayDict<K,V>::ConstIterator SplayDict<K, V>::end(){
  ConstIterator it = ConstIterator(nullptr, root);
  return it;
}

template<typename K, typename V>
typename SplayDict<K, V>::Node* SplayDict<K, V>::findNode(Node* subtree, const K& key) const{
  if (subtree == nullptr){
    return nullptr;
  }
  else if (subtree->key == key) {
    return subtree;
  }
  else if (key < subtree->key) {
    return findNode(subtree->left, key);
  }
  else {
    return findNode(subtree->right, key);
  }
}

template<typename K, typename V>
void SplayDict<K, V>::set(Node*& subtree, Node* parent, const K& key, const V& val){
  if (subtree == nullptr) {
    // by using the reference to pointer, we update parent's pointer
    subtree = new Node(parent, key, val);
    root = splay(subtree);
  }
  else if (subtree->key == key) {
    subtree->val = val;
    root = splay(subtree);

  }
  else if (key < subtree->key) {
    set(subtree->left, subtree, key, val);
  }
  else {
    set(subtree->right, subtree, key, val);
  }
}

template<typename K, typename V>
typename SplayDict<K,V>::Node* SplayDict<K, V>::subtreeMaxKey(Node* subtree, Node* parent){
  if (subtree == nullptr){
    return parent;
  }
  else {
    return subtreeMaxKey(subtree->right, subtree);
  }
}

template<typename K, typename V>
typename SplayDict<K,V>::Node* SplayDict<K, V>::subtreeMin(Node* subtree){
  return subtreeMinKey(subtree->left, subtree);
}

template<typename K, typename V>
typename SplayDict<K,V>::Node* SplayDict<K, V>::subtreeMax(Node* subtree){
  return subtreeMaxKey(subtree->right, subtree);
}

template<typename K, typename V>
typename SplayDict<K,V>::Node* SplayDict<K, V>::subtreeMinKey(Node* subtree, Node* parent){
  if (subtree == nullptr){
    return parent;
  }
  else {
    return subtreeMinKey(subtree->left, subtree);
  }
}

template<typename K, typename V>
void SplayDict<K, V>::erase_node(Node* node){

  splay(node);

  if (node->left != nullptr){
    Node* left_subtree = node->left;
    Node* right_subtree = node->right;

    left_subtree->parent = nullptr;
    right_subtree->parent = nullptr;
    delete node;

    Node* max_left = subtreeMax(left_subtree);

    splay(max_left);
    max_left->right = right_subtree;
    right_subtree->parent = max_left;
    root = max_left;

  }
  else {
    Node* right_subtree = node->right;
    right_subtree->parent = nullptr;
    delete node;

    root = right_subtree;
  }

}

template<typename K, typename V>
const std::pair<K&, V&> SplayDict<K, V>::ConstIterator::operator*() const {
  std::pair<K&, V&> res {current->key, current->val};
  return res;
}

template<typename K, typename V>
typename SplayDict<K,V>::ConstIterator& SplayDict<K, V>::ConstIterator::operator=(const ConstIterator& it) {
  current = it.current;
  prev = it.prev;

  return *this;
}

template<typename K, typename V>
bool SplayDict<K, V>::ConstIterator::operator==(const ConstIterator& it) const {
  return (current == it.current and prev == it.prev);
}

template<typename K, typename V>
bool SplayDict<K, V>::ConstIterator::operator!=(const ConstIterator& it) const {
  return !(*this == it);
}

template<typename K, typename V>
typename SplayDict<K,V>::ConstIterator& SplayDict<K, V>::ConstIterator::operator++(int) {
  if (prev != nullptr and prev->left == current) {
    if (current->right != nullptr){
      current = subtreeMin(current->right);
      prev = current->parent;
    }
    else {
      std::swap(current, prev);
    }
  }
  else if (prev != nullptr and prev->right == current) {
    if (current->right != nullptr) {
      current = subtreeMin(current->right);
      prev = current->parent;
    }
    else {
      do {
        prev = current;
        current = current->parent;
      } while (current != nullptr and current->right == prev);
    }
  }
  else if (current->left == prev){
    if (current->right != nullptr){
      current = subtreeMin(current->right);
      prev = current->parent;
    }
    else {
      do {
        prev = current;
        current = current->parent;
      } while (current != nullptr and current->right == prev);
    }

  }

  return *this;
}

template<typename K, typename V>
bool SplayDict<K, V>::sideOfChild(Node* subtree, bool zag1, bool zag2) const {
  Node* grand_parent = subtree->parent->parent;
  Node* possible_parent;
  if (zag1 == false) {
    possible_parent = grand_parent->left;
  }
  else{
    possible_parent = grand_parent->right;
  }

  if (possible_parent == nullptr){
    return false;
  }

  Node* possible_child;
  if (zag2 == false) {
    possible_child = possible_parent->left;
  }
  else {
    possible_child = possible_parent->right;
  }

  return possible_child == subtree;
}

template<typename K, typename V>
bool SplayDict<K, V>::hasGrandParent(Node* subtree) const {
  if (subtree->parent == nullptr) {
    return false;
  }
  return subtree->parent->parent != nullptr;
}

template<typename K, typename V>
void SplayDict<K, V>::rightRotate(Node* n) {
  Node* new_root = n->left;

  n->left = new_root->right;

  new_root->right = n;

  new_root->parent = n->parent;

  n->parent = new_root;

  if (n->left != nullptr){
    n->left->parent = n;
  }

  if (new_root->parent != nullptr and new_root->parent->left == n){
    new_root->parent->left = new_root;
  }
  else if (new_root->parent != nullptr) {
    new_root->parent->right = new_root;
  }
}

template<typename K, typename V>
void SplayDict<K, V>::leftRotate(Node* n) {
  Node* new_root = n->right;

  n->right = new_root->left;

  new_root->left = n;

  new_root->parent = n->parent;

  n->parent = new_root;

  if (n->right != nullptr) {
    n->right->parent = n;
  }

  if (new_root->parent != nullptr and new_root->parent->left == n) {
    new_root->parent->left = new_root;
  }
  else if (new_root->parent != nullptr) {
    new_root->parent->right = new_root;
  }
}

template<typename K, typename V>
typename SplayDict<K,V>::Node* SplayDict<K, V>::splay(Node* subtree){
  while(hasGrandParent(subtree)) {
    if (sideOfChild(subtree, false, false)) {
      // Zig Zig

      rightRotate(subtree->parent->parent);
      rightRotate(subtree->parent);
    } else if (sideOfChild(subtree, false, true)) {
      // Zig Zag

      leftRotate(subtree->parent);
      rightRotate(subtree->parent);

    } else if (sideOfChild(subtree, true, false)) {
      // Zag Zig

      rightRotate(subtree->parent);
      leftRotate(subtree->parent);
    } else if (sideOfChild(subtree, true, true)) {
      // Zag Zag

      leftRotate(subtree->parent->parent);
      leftRotate(subtree->parent);
    }
  }

  if(subtree->parent != nullptr) {
    if (subtree->parent->left == subtree) {
      rightRotate(subtree->parent);
    }
    else {
      leftRotate(subtree->parent);
    }
  }

  return subtree;

}

template<typename K, typename V>
void SplayDict<K,V>::printNode(std::ostream& os, Node* node, std::string prefix, bool isTail){
    if (node == nullptr){
      return;
    }
    os << (prefix + (isTail ? "└── " : "├── ")) << node->key << std::endl;
    if (node->left != nullptr){
      printNode(os, node->right, (prefix + (isTail ? "    " : "│   ")), false);
      printNode(os, node->left, (prefix + (isTail ? "    " : "│   ")), true);
    }
    else {
      printNode(os, node->right, (prefix + (isTail ? "    " : "│   ")), true);
    }
}

}

#endif
<commit_msg>bug fix deleting key<commit_after>#ifndef SPLAY_TREE_H_
#define SPLAY_TREE_H_

#include <utility>
#include <iostream>
#include <stack>

namespace dicts {

template<typename K, typename V>
class SplayDict {
  private:
    struct Node; //Forward declaration

  public:
    class ConstIterator {
      public:
        const std::pair<K&, V&> operator*() const;

        ConstIterator& operator++ (int);

        ConstIterator& operator= (const ConstIterator& const_it);

        bool operator== (const ConstIterator& const_it) const;

        bool operator!= (const ConstIterator& const_it) const;

        ConstIterator(Node* root){
          if (root != nullptr){
            current = subtreeMin(root);
            prev = current->parent;
          }
          else {
            current = root;
            prev = nullptr;
          }
        }

        ConstIterator(Node* c, Node* p){
          current = c;
          prev = p;
        }

      private:
        Node* current;
        Node* prev;
    };

    SplayDict();

    SplayDict(const SplayDict& bst);

    ~SplayDict();

    bool exists(const K& key) const;

    V& get(const K& key);

    void set(const K& key, const V& val);

    void erase(const K& key);

    SplayDict& operator= (const SplayDict& bst);

    ConstIterator begin();

    ConstIterator end();

  private:
    struct Node {
      K key;
      V val;

      Node* parent;
      Node* left;
      Node* right;

      Node(Node* p, const K& k, const V& v) {
        key = k;
        val = v;

        parent  = p;
        left    = nullptr;
        right   = nullptr;
      }

    };

    static void printNode(std::ostream& os, Node* node, std::string prefix, bool isTail);

    Node* findNode(Node* subtree, const K& key) const;

    bool hasGrandParent(Node* subtree) const;

    void rightRotate(Node* n);

    void leftRotate(Node* n);

    bool sideOfChild(Node* subtree, bool zag1, bool zag2) const;

    Node* splay(Node* subtree);

    void set(Node*& subtree, Node* parent, const K& key, const V& val);

    static Node* subtreeMin(Node* subtree);

    static Node* subtreeMax(Node* subtree);

    static Node* subtreeMaxKey(Node* subtree, Node* parent);

    static Node* subtreeMinKey(Node* subtree, Node* parent);

    void erase_node(Node* node);

    Node* root;

  friend class ConstIterator;

  friend std::ostream& operator<< (std::ostream& os, const SplayDict<K,V>& spDict){
    printNode(os, spDict.root, "", true);
    return os;
  }

};

template<typename K, typename V>
SplayDict<K, V>::SplayDict() {
    root = nullptr;
}

template<typename K, typename V>
SplayDict<K, V>::~SplayDict() {
  std::stack<Node*> nodes;
  nodes.push(root);

  while(!nodes.empty()){
    Node* current = nodes.top();
    nodes.pop();

    if (current != nullptr){
      nodes.push(current->left);
      nodes.push(current->right);

      delete current;
    }

  }
}

template<typename K, typename V>
SplayDict<K, V>::SplayDict(const SplayDict& bst) {
  for (ConstIterator it = bst.begin(); it != bst.end(); it++) {
    set( (*it).first, (*it).second);
  }
}

template<typename K, typename V>
bool SplayDict<K, V>::exists(const K& key) const {
  return findNode(root, key) != nullptr;
}

template<typename K, typename V>
V& SplayDict<K, V>::get(const K& key){
  Node* node = findNode(root, key);
  root = splay(node);
  return root->val;
}

template<typename K, typename V>
void SplayDict<K, V>::set(const K& key, const V& val) {
  set(root, nullptr, key, val);
}

template<typename K, typename V>
void SplayDict<K, V>::erase(const K& key){
  Node* node = findNode(root, key);

  erase_node(node);
}

template<typename K, typename V>
typename SplayDict<K,V>::ConstIterator SplayDict<K, V>::begin(){
  ConstIterator it = ConstIterator(root);
  return it;
}

template<typename K, typename V>
typename SplayDict<K,V>::ConstIterator SplayDict<K, V>::end(){
  ConstIterator it = ConstIterator(nullptr, root);
  return it;
}

template<typename K, typename V>
typename SplayDict<K, V>::Node* SplayDict<K, V>::findNode(Node* subtree, const K& key) const{
  if (subtree == nullptr){
    return nullptr;
  }
  else if (subtree->key == key) {
    return subtree;
  }
  else if (key < subtree->key) {
    return findNode(subtree->left, key);
  }
  else {
    return findNode(subtree->right, key);
  }
}

template<typename K, typename V>
void SplayDict<K, V>::set(Node*& subtree, Node* parent, const K& key, const V& val){
  if (subtree == nullptr) {
    // by using the reference to pointer, we update parent's pointer
    subtree = new Node(parent, key, val);
    root = splay(subtree);
  }
  else if (subtree->key == key) {
    subtree->val = val;
    root = splay(subtree);

  }
  else if (key < subtree->key) {
    set(subtree->left, subtree, key, val);
  }
  else {
    set(subtree->right, subtree, key, val);
  }
}

template<typename K, typename V>
typename SplayDict<K,V>::Node* SplayDict<K, V>::subtreeMaxKey(Node* subtree, Node* parent){
  if (subtree == nullptr){
    return parent;
  }
  else {
    return subtreeMaxKey(subtree->right, subtree);
  }
}

template<typename K, typename V>
typename SplayDict<K,V>::Node* SplayDict<K, V>::subtreeMin(Node* subtree){
  return subtreeMinKey(subtree->left, subtree);
}

template<typename K, typename V>
typename SplayDict<K,V>::Node* SplayDict<K, V>::subtreeMax(Node* subtree){
  return subtreeMaxKey(subtree->right, subtree);
}

template<typename K, typename V>
typename SplayDict<K,V>::Node* SplayDict<K, V>::subtreeMinKey(Node* subtree, Node* parent){
  if (subtree == nullptr){
    return parent;
  }
  else {
    return subtreeMinKey(subtree->left, subtree);
  }
}

template<typename K, typename V>
void SplayDict<K, V>::erase_node(Node* node){

  splay(node);

  if (node->left != nullptr){
    Node* left_subtree = node->left;
    Node* right_subtree = node->right;

    left_subtree->parent = nullptr;
    if(right_subtree != nullptr)
		right_subtree->parent = nullptr;
    delete node;

    Node* max_left = subtreeMax(left_subtree);

    splay(max_left);
    max_left->right = right_subtree;
    if(right_subtree != nullptr)
		right_subtree->parent = max_left;
    root = max_left;

  }
  else {
    Node* right_subtree = node->right;
    right_subtree->parent = nullptr;
    delete node;

    root = right_subtree;
  }

}

template<typename K, typename V>
const std::pair<K&, V&> SplayDict<K, V>::ConstIterator::operator*() const {
  std::pair<K&, V&> res {current->key, current->val};
  return res;
}

template<typename K, typename V>
typename SplayDict<K,V>::ConstIterator& SplayDict<K, V>::ConstIterator::operator=(const ConstIterator& it) {
  current = it.current;
  prev = it.prev;

  return *this;
}

template<typename K, typename V>
bool SplayDict<K, V>::ConstIterator::operator==(const ConstIterator& it) const {
  return (current == it.current and prev == it.prev);
}

template<typename K, typename V>
bool SplayDict<K, V>::ConstIterator::operator!=(const ConstIterator& it) const {
  return !(*this == it);
}

template<typename K, typename V>
typename SplayDict<K,V>::ConstIterator& SplayDict<K, V>::ConstIterator::operator++(int) {
  if (prev != nullptr and prev->left == current) {
    if (current->right != nullptr){
      current = subtreeMin(current->right);
      prev = current->parent;
    }
    else {
      std::swap(current, prev);
    }
  }
  else if (prev != nullptr and prev->right == current) {
    if (current->right != nullptr) {
      current = subtreeMin(current->right);
      prev = current->parent;
    }
    else {
      do {
        prev = current;
        current = current->parent;
      } while (current != nullptr and current->right == prev);
    }
  }
  else if (current->left == prev){
    if (current->right != nullptr){
      current = subtreeMin(current->right);
      prev = current->parent;
    }
    else {
      do {
        prev = current;
        current = current->parent;
      } while (current != nullptr and current->right == prev);
    }

  }

  return *this;
}

template<typename K, typename V>
bool SplayDict<K, V>::sideOfChild(Node* subtree, bool zag1, bool zag2) const {
  Node* grand_parent = subtree->parent->parent;
  Node* possible_parent;
  if (zag1 == false) {
    possible_parent = grand_parent->left;
  }
  else{
    possible_parent = grand_parent->right;
  }

  if (possible_parent == nullptr){
    return false;
  }

  Node* possible_child;
  if (zag2 == false) {
    possible_child = possible_parent->left;
  }
  else {
    possible_child = possible_parent->right;
  }

  return possible_child == subtree;
}

template<typename K, typename V>
bool SplayDict<K, V>::hasGrandParent(Node* subtree) const {
  if (subtree->parent == nullptr) {
    return false;
  }
  return subtree->parent->parent != nullptr;
}

template<typename K, typename V>
void SplayDict<K, V>::rightRotate(Node* n) {
  Node* new_root = n->left;

  n->left = new_root->right;

  new_root->right = n;

  new_root->parent = n->parent;

  n->parent = new_root;

  if (n->left != nullptr){
    n->left->parent = n;
  }

  if (new_root->parent != nullptr and new_root->parent->left == n){
    new_root->parent->left = new_root;
  }
  else if (new_root->parent != nullptr) {
    new_root->parent->right = new_root;
  }
}

template<typename K, typename V>
void SplayDict<K, V>::leftRotate(Node* n) {
  Node* new_root = n->right;

  n->right = new_root->left;

  new_root->left = n;

  new_root->parent = n->parent;

  n->parent = new_root;

  if (n->right != nullptr) {
    n->right->parent = n;
  }

  if (new_root->parent != nullptr and new_root->parent->left == n) {
    new_root->parent->left = new_root;
  }
  else if (new_root->parent != nullptr) {
    new_root->parent->right = new_root;
  }
}

template<typename K, typename V>
typename SplayDict<K,V>::Node* SplayDict<K, V>::splay(Node* subtree){
  while(hasGrandParent(subtree)) {
    if (sideOfChild(subtree, false, false)) {
      // Zig Zig

      rightRotate(subtree->parent->parent);
      rightRotate(subtree->parent);
    } else if (sideOfChild(subtree, false, true)) {
      // Zig Zag

      leftRotate(subtree->parent);
      rightRotate(subtree->parent);

    } else if (sideOfChild(subtree, true, false)) {
      // Zag Zig

      rightRotate(subtree->parent);
      leftRotate(subtree->parent);
    } else if (sideOfChild(subtree, true, true)) {
      // Zag Zag

      leftRotate(subtree->parent->parent);
      leftRotate(subtree->parent);
    }
  }

  if(subtree->parent != nullptr) {
    if (subtree->parent->left == subtree) {
      rightRotate(subtree->parent);
    }
    else {
      leftRotate(subtree->parent);
    }
  }

  return subtree;

}

template<typename K, typename V>
void SplayDict<K,V>::printNode(std::ostream& os, Node* node, std::string prefix, bool isTail){
    if (node == nullptr){
      return;
    }
    os << (prefix + (isTail ? "└── " : "├── ")) << node->key << std::endl;
    if (node->left != nullptr){
      printNode(os, node->right, (prefix + (isTail ? "    " : "│   ")), false);
      printNode(os, node->left, (prefix + (isTail ? "    " : "│   ")), true);
    }
    else {
      printNode(os, node->right, (prefix + (isTail ? "    " : "│   ")), true);
    }
}

}

#endif
<|endoftext|>


<commit_before>/*********************************************************************
 * Software License Agreement (BSD License)
 *
 *  Copyright (c) 2011, Willow Garage, Inc.
 *  All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions
 *  are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer in the documentation and/or other materials provided
 *     with the distribution.
 *   * Neither the name of Willow Garage nor the names of its
 *     contributors may be used to endorse or promote products derived
 *     from this software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 *  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 *  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 *  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 *  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 *  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 *  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 *  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 *  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 *  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 *  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 *********************************************************************/

/* Author: Jon Binney, Ioan Sucan */

#include <moveit/pointcloud_octomap_updater/pointcloud_octomap_updater.h>
#include <moveit/occupancy_map_monitor/occupancy_map_monitor.h>
#include <message_filters/subscriber.h>
#include <pcl_conversions/pcl_conversions.h>
#include <XmlRpcException.h>

namespace occupancy_map_monitor
{

PointCloudOctomapUpdater::PointCloudOctomapUpdater() : OccupancyMapUpdater("PointCloudUpdater"),
                                                       scale_(1.0),
                                                       padding_(0.0),
                                                       max_range_(std::numeric_limits<double>::infinity()),
                                                       point_subsample_(1),
                                                       point_cloud_subscriber_(NULL),
                                                       point_cloud_filter_(NULL),
                               private_nh_("~")
{
}

PointCloudOctomapUpdater::~PointCloudOctomapUpdater()
{
  stopHelper();
}

bool PointCloudOctomapUpdater::setParams(XmlRpc::XmlRpcValue &params)
{
  try
  {
    if (!params.hasMember("point_cloud_topic"))
      return false;
    point_cloud_topic_ = static_cast<const std::string&>(params["point_cloud_topic"]);

    readXmlParam(params, "max_range", &max_range_);
    readXmlParam(params, "shape_padding", &padding_);
    readXmlParam(params, "shape_scale", &scale_);
    readXmlParam(params, "point_subsample", &point_subsample_);
    if (params.hasMember("filtered_cloud_topic"))
      filtered_cloud_topic_ = static_cast<const std::string&>(params["filtered_cloud_topic"]);
  }
  catch (XmlRpc::XmlRpcException &ex)
  {
    ROS_ERROR("XmlRpc Exception: %s", ex.getMessage().c_str());
    return false;
  }

  return true;
}

bool PointCloudOctomapUpdater::initialize()
{
  tf_ = monitor_->getTFClient();
  shape_mask_.reset(new point_containment_filter::ShapeMask());
  shape_mask_->setTransformCallback(boost::bind(&PointCloudOctomapUpdater::getShapeTransform, this, _1, _2));
  if (!filtered_cloud_topic_.empty())
    filtered_cloud_publisher_ = private_nh_.advertise<sensor_msgs::PointCloud2>(filtered_cloud_topic_, 10, false);
  return true;
}

void PointCloudOctomapUpdater::start()
{
  if (point_cloud_subscriber_)
    return;
  /* subscribe to point cloud topic using tf filter*/
  point_cloud_subscriber_ = new message_filters::Subscriber<sensor_msgs::PointCloud2>(root_nh_, point_cloud_topic_, 5);
  if (tf_ && !monitor_->getMapFrame().empty())
  {
    point_cloud_filter_ = new tf::MessageFilter<sensor_msgs::PointCloud2>(*point_cloud_subscriber_, *tf_, monitor_->getMapFrame(), 5);
    point_cloud_filter_->registerCallback(boost::bind(&PointCloudOctomapUpdater::cloudMsgCallback, this, _1));
    ROS_INFO("Listening to '%s' using message filter with target frame '%s'", point_cloud_topic_.c_str(), point_cloud_filter_->getTargetFramesString().c_str());
  }
  else
  {
    point_cloud_subscriber_->registerCallback(boost::bind(&PointCloudOctomapUpdater::cloudMsgCallback, this, _1));
    ROS_INFO("Listening to '%s'", point_cloud_topic_.c_str());
  }
}

void PointCloudOctomapUpdater::stopHelper()
{
  delete point_cloud_filter_;
  delete point_cloud_subscriber_;
}

void PointCloudOctomapUpdater::stop()
{
  stopHelper();
  point_cloud_filter_ = NULL;
  point_cloud_subscriber_ = NULL;
}

ShapeHandle PointCloudOctomapUpdater::excludeShape(const shapes::ShapeConstPtr &shape)
{
  ShapeHandle h = 0;
  if (shape_mask_)
    h = shape_mask_->addShape(shape, scale_, padding_);
  else
    ROS_ERROR("Shape filter not yet initialized!");
  return h;
}

void PointCloudOctomapUpdater::forgetShape(ShapeHandle handle)
{
  if (shape_mask_)
    shape_mask_->removeShape(handle);
}

bool PointCloudOctomapUpdater::getShapeTransform(ShapeHandle h, Eigen::Affine3d &transform) const
{
  ShapeTransformCache::const_iterator it = transform_cache_.find(h);
  if (it == transform_cache_.end())
  {
    ROS_ERROR("Internal error. Shape filter handle %u not found", h);
    return false;
  }
  transform = it->second;
  return true;
}

void PointCloudOctomapUpdater::updateMask(const pcl::PointCloud<pcl::PointXYZ> &cloud, const Eigen::Vector3d &sensor_origin, std::vector<int> &mask)
{
}

void PointCloudOctomapUpdater::cloudMsgCallback(const sensor_msgs::PointCloud2::ConstPtr &cloud_msg)
{
  ROS_DEBUG("Received a new point cloud message");
  ros::WallTime start = ros::WallTime::now();

  if (monitor_->getMapFrame().empty())
    monitor_->setMapFrame(cloud_msg->header.frame_id);

  /* get transform for cloud into map frame */
  tf::StampedTransform map_H_sensor;
  if (monitor_->getMapFrame() == cloud_msg->header.frame_id)
    map_H_sensor.setIdentity();
  else
  {
    if (tf_)
    {
      try
      {
        tf_->lookupTransform(monitor_->getMapFrame(), cloud_msg->header.frame_id, cloud_msg->header.stamp, map_H_sensor);
      }
      catch (tf::TransformException& ex)
      {
        ROS_ERROR_STREAM("Transform error of sensor data: " << ex.what() << "; quitting callback");
        return;
      }
    }
    else
      return;
  }

  /* convert cloud message to pcl cloud object */
  pcl::PointCloud<pcl::PointXYZ> cloud;
  pcl::fromROSMsg(*cloud_msg, cloud);

  /* compute sensor origin in map frame */
  const tf::Vector3 &sensor_origin_tf = map_H_sensor.getOrigin();
  octomap::point3d sensor_origin(sensor_origin_tf.getX(), sensor_origin_tf.getY(), sensor_origin_tf.getZ());
  Eigen::Vector3d sensor_origin_eigen(sensor_origin_tf.getX(), sensor_origin_tf.getY(), sensor_origin_tf.getZ());

  if (!updateTransformCache(cloud_msg->header.frame_id, cloud_msg->header.stamp))
  {
    ROS_ERROR_THROTTLE(1, "Transform cache was not updated. Self-filtering may fail.");
    return;
  }

  /* mask out points on the robot */
  shape_mask_->maskContainment(cloud, sensor_origin_eigen, 0.0, max_range_, mask_);
  updateMask(cloud, sensor_origin_eigen, mask_);

  octomap::KeySet free_cells, occupied_cells, model_cells, clip_cells;
  boost::scoped_ptr<pcl::PointCloud<pcl::PointXYZ> > filtered_cloud;
  if (!filtered_cloud_topic_.empty())
    filtered_cloud.reset(new pcl::PointCloud<pcl::PointXYZ>());

  tree_->lockRead();

  try
  {
    /* do ray tracing to find which cells this point cloud indicates should be free, and which it indicates
     * should be occupied */
    for (unsigned int row = 0; row < cloud.height; row += point_subsample_)
    {
      unsigned int row_c = row * cloud.width;
      for (unsigned int col = 0; col < cloud.width; col += point_subsample_)
      {
        //if (mask_[row_c + col] == point_containment_filter::ShapeMask::CLIP)
        //  continue;
        const pcl::PointXYZ &p = cloud(col, row);

        /* check for NaN */
        if ((p.x == p.x) && (p.y == p.y) && (p.z == p.z))
    {
      /* transform to map frame */
      tf::Vector3 point_tf = map_H_sensor * tf::Vector3(p.x, p.y, p.z);

      /* occupied cell at ray endpoint if ray is shorter than max range and this point
         isn't on a part of the robot*/
      if (mask_[row_c + col] == point_containment_filter::ShapeMask::INSIDE)
        model_cells.insert(tree_->coordToKey(point_tf.getX(), point_tf.getY(), point_tf.getZ()));
      else if (mask_[row_c + col] == point_containment_filter::ShapeMask::CLIP)
        clip_cells.insert(tree_->coordToKey(point_tf.getX(), point_tf.getY(), point_tf.getZ()));
      else
          {
            occupied_cells.insert(tree_->coordToKey(point_tf.getX(), point_tf.getY(), point_tf.getZ()));
            if (filtered_cloud)
              filtered_cloud->push_back(p);
          }
        }
      }
    }

    /* compute the free cells along each ray that ends at an occupied cell */
    for (octomap::KeySet::iterator it = occupied_cells.begin(), end = occupied_cells.end(); it != end; ++it)
      if (tree_->computeRayKeys(sensor_origin, tree_->keyToCoord(*it), key_ray_))
        free_cells.insert(key_ray_.begin(), key_ray_.end());

    /* compute the free cells along each ray that ends at a model cell */
    for (octomap::KeySet::iterator it = model_cells.begin(), end = model_cells.end(); it != end; ++it)
      if (tree_->computeRayKeys(sensor_origin, tree_->keyToCoord(*it), key_ray_))
        free_cells.insert(key_ray_.begin(), key_ray_.end());

    /* compute the free cells along each ray that ends at a clipped cell */
    for (octomap::KeySet::iterator it = clip_cells.begin(), end = clip_cells.end(); it != end; ++it)
      if (tree_->computeRayKeys(sensor_origin, tree_->keyToCoord(*it), key_ray_))
        free_cells.insert(key_ray_.begin(), key_ray_.end());
  }
  catch (...)
  {
    tree_->unlockRead();
    return;
  }

  tree_->unlockRead();

  /* cells that overlap with the model are not occupied */
  for (octomap::KeySet::iterator it = model_cells.begin(), end = model_cells.end(); it != end; ++it)
    occupied_cells.erase(*it);

  /* occupied cells are not free */
  for (octomap::KeySet::iterator it = occupied_cells.begin(), end = occupied_cells.end(); it != end; ++it)
    free_cells.erase(*it);

  tree_->lockWrite();

  try
  {
    /* mark free cells only if not seen occupied in this cloud */
    for (octomap::KeySet::iterator it = free_cells.begin(), end = free_cells.end(); it != end; ++it)
      tree_->updateNode(*it, false);

    /* now mark all occupied cells */
    for (octomap::KeySet::iterator it = occupied_cells.begin(), end = occupied_cells.end(); it != end; ++it)
      tree_->updateNode(*it, true);

    // set the logodds to the minimum for the cells that are part of the model
    const float lg = tree_->getClampingThresMinLog() - tree_->getClampingThresMaxLog();
    for (octomap::KeySet::iterator it = model_cells.begin(), end = model_cells.end(); it != end; ++it)
      tree_->updateNode(*it, lg);
  }
  catch (...)
  {
    ROS_ERROR("Internal error while updating octree");
  }
  tree_->unlockWrite();
  ROS_DEBUG("Processed point cloud in %lf ms", (ros::WallTime::now() - start).toSec() * 1000.0);
  tree_->triggerUpdateCallback();

  if (filtered_cloud)
  {
    sensor_msgs::PointCloud2 filtered_cloud_msg;
    pcl::toROSMsg(*filtered_cloud, filtered_cloud_msg);
    filtered_cloud_msg.header = cloud_msg->header;
    filtered_cloud_publisher_.publish(filtered_cloud_msg);
  }
}

}
<commit_msg>change name of param<commit_after>/*********************************************************************
 * Software License Agreement (BSD License)
 *
 *  Copyright (c) 2011, Willow Garage, Inc.
 *  All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions
 *  are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer in the documentation and/or other materials provided
 *     with the distribution.
 *   * Neither the name of Willow Garage nor the names of its
 *     contributors may be used to endorse or promote products derived
 *     from this software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 *  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 *  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 *  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 *  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 *  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 *  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 *  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 *  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 *  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 *  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 *********************************************************************/

/* Author: Jon Binney, Ioan Sucan */

#include <moveit/pointcloud_octomap_updater/pointcloud_octomap_updater.h>
#include <moveit/occupancy_map_monitor/occupancy_map_monitor.h>
#include <message_filters/subscriber.h>
#include <pcl_conversions/pcl_conversions.h>
#include <XmlRpcException.h>

namespace occupancy_map_monitor
{

PointCloudOctomapUpdater::PointCloudOctomapUpdater() : OccupancyMapUpdater("PointCloudUpdater"),
                                                       scale_(1.0),
                                                       padding_(0.0),
                                                       max_range_(std::numeric_limits<double>::infinity()),
                                                       point_subsample_(1),
                                                       point_cloud_subscriber_(NULL),
                                                       point_cloud_filter_(NULL),
                               private_nh_("~")
{
}

PointCloudOctomapUpdater::~PointCloudOctomapUpdater()
{
  stopHelper();
}

bool PointCloudOctomapUpdater::setParams(XmlRpc::XmlRpcValue &params)
{
  try
  {
    if (!params.hasMember("point_cloud_topic"))
      return false;
    point_cloud_topic_ = static_cast<const std::string&>(params["point_cloud_topic"]);

    readXmlParam(params, "max_range", &max_range_);
    readXmlParam(params, "padding_offset", &padding_);
    readXmlParam(params, "padding_scale", &scale_);
    readXmlParam(params, "point_subsample", &point_subsample_);
    if (params.hasMember("filtered_cloud_topic"))
      filtered_cloud_topic_ = static_cast<const std::string&>(params["filtered_cloud_topic"]);
  }
  catch (XmlRpc::XmlRpcException &ex)
  {
    ROS_ERROR("XmlRpc Exception: %s", ex.getMessage().c_str());
    return false;
  }

  return true;
}

bool PointCloudOctomapUpdater::initialize()
{
  tf_ = monitor_->getTFClient();
  shape_mask_.reset(new point_containment_filter::ShapeMask());
  shape_mask_->setTransformCallback(boost::bind(&PointCloudOctomapUpdater::getShapeTransform, this, _1, _2));
  if (!filtered_cloud_topic_.empty())
    filtered_cloud_publisher_ = private_nh_.advertise<sensor_msgs::PointCloud2>(filtered_cloud_topic_, 10, false);
  return true;
}

void PointCloudOctomapUpdater::start()
{
  if (point_cloud_subscriber_)
    return;
  /* subscribe to point cloud topic using tf filter*/
  point_cloud_subscriber_ = new message_filters::Subscriber<sensor_msgs::PointCloud2>(root_nh_, point_cloud_topic_, 5);
  if (tf_ && !monitor_->getMapFrame().empty())
  {
    point_cloud_filter_ = new tf::MessageFilter<sensor_msgs::PointCloud2>(*point_cloud_subscriber_, *tf_, monitor_->getMapFrame(), 5);
    point_cloud_filter_->registerCallback(boost::bind(&PointCloudOctomapUpdater::cloudMsgCallback, this, _1));
    ROS_INFO("Listening to '%s' using message filter with target frame '%s'", point_cloud_topic_.c_str(), point_cloud_filter_->getTargetFramesString().c_str());
  }
  else
  {
    point_cloud_subscriber_->registerCallback(boost::bind(&PointCloudOctomapUpdater::cloudMsgCallback, this, _1));
    ROS_INFO("Listening to '%s'", point_cloud_topic_.c_str());
  }
}

void PointCloudOctomapUpdater::stopHelper()
{
  delete point_cloud_filter_;
  delete point_cloud_subscriber_;
}

void PointCloudOctomapUpdater::stop()
{
  stopHelper();
  point_cloud_filter_ = NULL;
  point_cloud_subscriber_ = NULL;
}

ShapeHandle PointCloudOctomapUpdater::excludeShape(const shapes::ShapeConstPtr &shape)
{
  ShapeHandle h = 0;
  if (shape_mask_)
    h = shape_mask_->addShape(shape, scale_, padding_);
  else
    ROS_ERROR("Shape filter not yet initialized!");
  return h;
}

void PointCloudOctomapUpdater::forgetShape(ShapeHandle handle)
{
  if (shape_mask_)
    shape_mask_->removeShape(handle);
}

bool PointCloudOctomapUpdater::getShapeTransform(ShapeHandle h, Eigen::Affine3d &transform) const
{
  ShapeTransformCache::const_iterator it = transform_cache_.find(h);
  if (it == transform_cache_.end())
  {
    ROS_ERROR("Internal error. Shape filter handle %u not found", h);
    return false;
  }
  transform = it->second;
  return true;
}

void PointCloudOctomapUpdater::updateMask(const pcl::PointCloud<pcl::PointXYZ> &cloud, const Eigen::Vector3d &sensor_origin, std::vector<int> &mask)
{
}

void PointCloudOctomapUpdater::cloudMsgCallback(const sensor_msgs::PointCloud2::ConstPtr &cloud_msg)
{
  ROS_DEBUG("Received a new point cloud message");
  ros::WallTime start = ros::WallTime::now();

  if (monitor_->getMapFrame().empty())
    monitor_->setMapFrame(cloud_msg->header.frame_id);

  /* get transform for cloud into map frame */
  tf::StampedTransform map_H_sensor;
  if (monitor_->getMapFrame() == cloud_msg->header.frame_id)
    map_H_sensor.setIdentity();
  else
  {
    if (tf_)
    {
      try
      {
        tf_->lookupTransform(monitor_->getMapFrame(), cloud_msg->header.frame_id, cloud_msg->header.stamp, map_H_sensor);
      }
      catch (tf::TransformException& ex)
      {
        ROS_ERROR_STREAM("Transform error of sensor data: " << ex.what() << "; quitting callback");
        return;
      }
    }
    else
      return;
  }

  /* convert cloud message to pcl cloud object */
  pcl::PointCloud<pcl::PointXYZ> cloud;
  pcl::fromROSMsg(*cloud_msg, cloud);

  /* compute sensor origin in map frame */
  const tf::Vector3 &sensor_origin_tf = map_H_sensor.getOrigin();
  octomap::point3d sensor_origin(sensor_origin_tf.getX(), sensor_origin_tf.getY(), sensor_origin_tf.getZ());
  Eigen::Vector3d sensor_origin_eigen(sensor_origin_tf.getX(), sensor_origin_tf.getY(), sensor_origin_tf.getZ());

  if (!updateTransformCache(cloud_msg->header.frame_id, cloud_msg->header.stamp))
  {
    ROS_ERROR_THROTTLE(1, "Transform cache was not updated. Self-filtering may fail.");
    return;
  }

  /* mask out points on the robot */
  shape_mask_->maskContainment(cloud, sensor_origin_eigen, 0.0, max_range_, mask_);
  updateMask(cloud, sensor_origin_eigen, mask_);

  octomap::KeySet free_cells, occupied_cells, model_cells, clip_cells;
  boost::scoped_ptr<pcl::PointCloud<pcl::PointXYZ> > filtered_cloud;
  if (!filtered_cloud_topic_.empty())
    filtered_cloud.reset(new pcl::PointCloud<pcl::PointXYZ>());

  tree_->lockRead();

  try
  {
    /* do ray tracing to find which cells this point cloud indicates should be free, and which it indicates
     * should be occupied */
    for (unsigned int row = 0; row < cloud.height; row += point_subsample_)
    {
      unsigned int row_c = row * cloud.width;
      for (unsigned int col = 0; col < cloud.width; col += point_subsample_)
      {
        //if (mask_[row_c + col] == point_containment_filter::ShapeMask::CLIP)
        //  continue;
        const pcl::PointXYZ &p = cloud(col, row);

        /* check for NaN */
        if ((p.x == p.x) && (p.y == p.y) && (p.z == p.z))
    {
      /* transform to map frame */
      tf::Vector3 point_tf = map_H_sensor * tf::Vector3(p.x, p.y, p.z);

      /* occupied cell at ray endpoint if ray is shorter than max range and this point
         isn't on a part of the robot*/
      if (mask_[row_c + col] == point_containment_filter::ShapeMask::INSIDE)
        model_cells.insert(tree_->coordToKey(point_tf.getX(), point_tf.getY(), point_tf.getZ()));
      else if (mask_[row_c + col] == point_containment_filter::ShapeMask::CLIP)
        clip_cells.insert(tree_->coordToKey(point_tf.getX(), point_tf.getY(), point_tf.getZ()));
      else
          {
            occupied_cells.insert(tree_->coordToKey(point_tf.getX(), point_tf.getY(), point_tf.getZ()));
            if (filtered_cloud)
              filtered_cloud->push_back(p);
          }
        }
      }
    }

    /* compute the free cells along each ray that ends at an occupied cell */
    for (octomap::KeySet::iterator it = occupied_cells.begin(), end = occupied_cells.end(); it != end; ++it)
      if (tree_->computeRayKeys(sensor_origin, tree_->keyToCoord(*it), key_ray_))
        free_cells.insert(key_ray_.begin(), key_ray_.end());

    /* compute the free cells along each ray that ends at a model cell */
    for (octomap::KeySet::iterator it = model_cells.begin(), end = model_cells.end(); it != end; ++it)
      if (tree_->computeRayKeys(sensor_origin, tree_->keyToCoord(*it), key_ray_))
        free_cells.insert(key_ray_.begin(), key_ray_.end());

    /* compute the free cells along each ray that ends at a clipped cell */
    for (octomap::KeySet::iterator it = clip_cells.begin(), end = clip_cells.end(); it != end; ++it)
      if (tree_->computeRayKeys(sensor_origin, tree_->keyToCoord(*it), key_ray_))
        free_cells.insert(key_ray_.begin(), key_ray_.end());
  }
  catch (...)
  {
    tree_->unlockRead();
    return;
  }

  tree_->unlockRead();

  /* cells that overlap with the model are not occupied */
  for (octomap::KeySet::iterator it = model_cells.begin(), end = model_cells.end(); it != end; ++it)
    occupied_cells.erase(*it);

  /* occupied cells are not free */
  for (octomap::KeySet::iterator it = occupied_cells.begin(), end = occupied_cells.end(); it != end; ++it)
    free_cells.erase(*it);

  tree_->lockWrite();

  try
  {
    /* mark free cells only if not seen occupied in this cloud */
    for (octomap::KeySet::iterator it = free_cells.begin(), end = free_cells.end(); it != end; ++it)
      tree_->updateNode(*it, false);

    /* now mark all occupied cells */
    for (octomap::KeySet::iterator it = occupied_cells.begin(), end = occupied_cells.end(); it != end; ++it)
      tree_->updateNode(*it, true);

    // set the logodds to the minimum for the cells that are part of the model
    const float lg = tree_->getClampingThresMinLog() - tree_->getClampingThresMaxLog();
    for (octomap::KeySet::iterator it = model_cells.begin(), end = model_cells.end(); it != end; ++it)
      tree_->updateNode(*it, lg);
  }
  catch (...)
  {
    ROS_ERROR("Internal error while updating octree");
  }
  tree_->unlockWrite();
  ROS_DEBUG("Processed point cloud in %lf ms", (ros::WallTime::now() - start).toSec() * 1000.0);
  tree_->triggerUpdateCallback();

  if (filtered_cloud)
  {
    sensor_msgs::PointCloud2 filtered_cloud_msg;
    pcl::toROSMsg(*filtered_cloud, filtered_cloud_msg);
    filtered_cloud_msg.header = cloud_msg->header;
    filtered_cloud_publisher_.publish(filtered_cloud_msg);
  }
}

}
<|endoftext|>