bag_reader.h

Go to the documentation of this file.

// Fast native bag reader
// Author: Max Schwarz <max.schwarz@ais.uni-bonn.de>

#ifndef ROSBAG_FANCY_BAGREADER_H
#define ROSBAG_FANCY_BAGREADER_H

#include <map>
#include <memory>
#include <string>

#include <ros/time.h>
#include <ros/serialization.h>

#include <rosbag/stream.h>

#include <boost/make_shared.hpp>

namespace rosbag_fancy
{

class BagReader
{
public:
        class Exception : public std::runtime_error
        {
        public:
                explicit Exception(const std::string& what)
                 : std::runtime_error{what}
                {}
        };

        struct Connection
        {
                std::uint32_t id;
                std::string topicInBag;

                std::string topicAsPublished;
                std::string type;
                std::string md5sum;
                std::string msgDef;
                std::string callerID;
                bool latching = false;

                std::uint64_t totalCount = 0;
        };

        using ConnectionMap = std::map<std::uint32_t, Connection>;

        struct ConnectionInfo
        {
                std::uint32_t id;
                std::uint32_t msgCount;
        };

        class ChunkIterator;
        class Iterator;

        struct Message
        {
                ros::Time stamp;
                const Connection* connection;

                template<class M>
                boost::shared_ptr<M> instantiate() const
                {
                        if(ros::message_traits::DataType<M>::value() != connection->type)
                                return {};
                        if(ros::message_traits::MD5Sum<M>::value() != connection->md5sum)
                                return {};

                        namespace ser = ros::serialization;
                        ser::IStream stream(const_cast<uint8_t*>(m_data), m_size);

                        auto ret = boost::make_shared<M>();
                        ser::deserialize(stream, *ret);
                        return ret;
                }

                const uint8_t* data() const
                { return m_data; }

                std::size_t size() const
                { return m_size; }

        private:
                friend class BagReader::ChunkIterator;
                friend struct ros::serialization::Serializer<rosbag_fancy::BagReader::Message>;

                const uint8_t* m_data = nullptr;
                std::size_t m_size;
        };

        class ChunkIterator
        {
        public:
                using iterator_category = std::input_iterator_tag;
                using value_type        = Message;
                using reference         = const Message&;
                using pointer           = const Message*;

                ChunkIterator(const ChunkIterator&) = default;
                ChunkIterator& operator=(const ChunkIterator&) = default;

                reference operator*() { return m_msg; }
                pointer operator->() { return &m_msg; }

                ChunkIterator& operator++();
                ChunkIterator operator++(int) { ChunkIterator tmp = *this; ++(*this); return tmp; }

                friend bool operator== (const ChunkIterator& a, const ChunkIterator& b) { return a.m_d == b.m_d && a.m_idx == b.m_idx; };
                friend bool operator!= (const ChunkIterator& a, const ChunkIterator& b) { return a.m_d != b.m_d || a.m_idx != b.m_idx; };

        private:
                friend class BagReader::Iterator;

                ChunkIterator() {}
                explicit ChunkIterator(const BagReader* reader, int chunk);

                class Private;
                std::shared_ptr<Private> m_d;
                std::size_t m_idx = 0;
                Message m_msg;
        };

        class Iterator
        {
        public:
                using iterator_category = std::input_iterator_tag;
                using value_type        = Message;
                using reference         = const Message&;
                using pointer           = const Message*;

                Iterator() {}
                Iterator(const Iterator&) = default;
                Iterator& operator=(const Iterator&) = default;

                reference operator*() { return *m_it; }
                pointer operator->() { return &*m_it; }

                Iterator& operator++();
                Iterator operator++(int) { Iterator tmp = *this; ++(*this); return tmp; }

                void advanceWithPredicates(const std::function<bool(const ConnectionInfo&)>& connPredicate, const std::function<bool(const Message&)>& messagePredicate);
                void findNextWithPredicates(const std::function<bool(const ConnectionInfo&)>& connPredicate, const std::function<bool(const Message&)>& messagePredicate);

                friend bool operator== (const Iterator& a, const Iterator& b) { return a.m_chunk == b.m_chunk && a.m_it == b.m_it; };
                friend bool operator!= (const Iterator& a, const Iterator& b) { return a.m_chunk != b.m_chunk || a.m_it != b.m_it; };

                std::vector<ConnectionInfo>& currentChunkConnections() const;
                rosbag::CompressionType currentChunkCompression() const;

                void skipChunk();
                int chunk() const
                { return m_chunk; }

        private:
                friend class BagReader;

                explicit Iterator(const BagReader* reader, int chunk);

                const BagReader* m_reader;
                int m_chunk = -1;
                ChunkIterator m_it;
        };

        explicit BagReader(const std::string& filename);
        ~BagReader();

        BagReader(const BagReader&) = delete;
        BagReader& operator=(const BagReader&) = delete;

        BagReader(BagReader&& other);

        const ConnectionMap& connections() const;

        ros::Time startTime() const;
        ros::Time endTime() const;

        std::size_t size() const;

        Iterator begin() const;
        Iterator end() const;
        Iterator findTime(const ros::Time& time) const;

        std::size_t numChunks() const;

        int findChunk(const ros::Time& time) const;
        Iterator chunkBegin(int chunk) const;
private:
        class Private;
        std::unique_ptr<Private> m_d;
};

}

namespace ros
{
namespace message_traits
{
template<> struct IsMessage<rosbag_fancy::BagReader::Message> : TrueType { };

template<>
struct MD5Sum<rosbag_fancy::BagReader::Message>
{
        static const char* value(const rosbag_fancy::BagReader::Message& m)
        {
                return m.connection->md5sum.c_str();
        }
};

template<>
struct DataType<rosbag_fancy::BagReader::Message>
{
        static const char* value(const rosbag_fancy::BagReader::Message& m)
        {
                return m.connection->type.c_str();
        }
};

template<>
struct Definition<rosbag_fancy::BagReader::Message>
{
        static const char* value(const rosbag_fancy::BagReader::Message& m)
        {
                return m.connection->msgDef.c_str();
        }
};
}

namespace serialization
{
template<>
struct Serializer<rosbag_fancy::BagReader::Message>
{
        template<typename Stream>
        inline static void write(Stream& stream, const rosbag_fancy::BagReader::Message& t)
        {
                std::memcpy(stream.advance(t.m_size), t.m_data, t.m_size);
        }

        inline static uint32_t serializedLength(const rosbag_fancy::BagReader::Message& t)
        {
                return t.m_size;
        }
};
}
}

#endif