From 7ec906e657c1a7042b2d71b6fa40747b5d5dfe83 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Dziepak?= <pdziepak@scylladb.com>
Date: Wed, 1 Aug 2018 15:14:31 +0100
Subject: [PATCH] imr: detect lsa migrator mismatch

Each IMR type needs its own LSA migrator. It is possible that user will
provide a migrator for a different type than the one which instance is
being created. This patch adds compile-time detection of that bug.
---
 imr/alloc.hh | 36 ++++++++++++++++++++++++++++++++----
 imr/utils.hh | 11 +++++++++--
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/imr/alloc.hh b/imr/alloc.hh
index 1217b6a2b4..d5a43c402a 100644
--- a/imr/alloc.hh
+++ b/imr/alloc.hh
@@ -84,6 +84,8 @@ template<typename Structure, typename CtxFactory>
 GCC6_CONCEPT(requires ContextFactory<CtxFactory>)
 class lsa_migrate_fn final : public migrate_fn_type, CtxFactory {
 public:
+    using structure = Structure;
+
     explicit lsa_migrate_fn(CtxFactory context_factory)
         : migrate_fn_type(1)
         , CtxFactory(std::move(context_factory))
@@ -201,8 +203,21 @@ public:
         /// arguments are passed to `T::size_when_serialized`.
         ///
         /// \return null pointer of type `uint8_t*`.
+        template<typename T, typename MigrateFn, typename... Args>
+        uint8_t* allocate(MigrateFn* migrate_fn, Args&&... args) noexcept {
+            static_assert(std::is_same_v<typename MigrateFn::structure, T>);
+            return do_allocate<T>(migrate_fn, std::forward<Args>(args)...);
+        }
+
+        template<typename T, typename MigrateFn, typename... Args>
+        auto allocate_nested(MigrateFn* migrate_fn, Args&&... args) noexcept {
+            static_assert(std::is_same_v<typename MigrateFn::structure, T>);
+            return do_allocate_nested<T>(migrate_fn, std::forward<Args>(args)...);
+        }
+
+    private:
         template<typename T, typename... Args>
-        uint8_t* allocate(migrate_fn_type* migrate_fn, Args&& ... args) noexcept {
+        uint8_t* do_allocate(migrate_fn_type* migrate_fn, Args&&... args) noexcept {
             auto size = T::size_when_serialized(std::forward<Args>(args)...);
             _parent.request(size, migrate_fn);
 
@@ -216,7 +231,7 @@ public:
         }
 
         template<typename T, typename... Args>
-        auto allocate_nested(migrate_fn_type* migrate_fn, Args&& ... args) noexcept {
+        auto do_allocate_nested(migrate_fn_type* migrate_fn, Args&& ... args) noexcept {
             auto n = _parent.request(0, migrate_fn);
             return T::get_sizer(continuation(_parent, n),
                                 std::forward<Args>(args)...);
@@ -244,15 +259,28 @@ public:
         /// to the buffer requested in the sizing phase. Arguments are passed
         /// to `T::serialize`.
         /// \return pointer to the IMR object
+        template<typename T, typename MigrateFn, typename... Args>
+        uint8_t* allocate(MigrateFn* migrate_fn, Args&&... args) noexcept {
+            static_assert(std::is_same_v<typename MigrateFn::structure, T>);
+            return do_allocate<T>(migrate_fn, std::forward<Args>(args)...);
+        }
+
+        template<typename T, typename MigrateFn, typename... Args>
+        auto allocate_nested(MigrateFn* migrate_fn, Args&&... args) noexcept {
+            static_assert(std::is_same_v<typename MigrateFn::structure, T>);
+            return do_allocate_nested<T>(migrate_fn, std::forward<Args>(args)...);
+        }
+
+    private:
         template<typename T, typename... Args>
-        uint8_t* allocate(migrate_fn_type* migrate_fn, Args&& ... args) noexcept {
+        uint8_t* do_allocate(migrate_fn_type* migrate_fn, Args&&... args) noexcept {
             auto ptr = _parent.next_object();
             T::serialize(ptr, std::forward<Args>(args)...);
             return ptr;
         }
 
         template<typename T, typename... Args>
-        auto allocate_nested(migrate_fn_type*, Args&& ... args) noexcept {
+        auto do_allocate_nested(migrate_fn_type*, Args&& ... args) noexcept {
             auto ptr = _parent.next_object();
             return T::get_serializer(ptr,
                                      continuation(ptr),
diff --git a/imr/utils.hh b/imr/utils.hh
index 679612b83b..6b6c1d81f3 100644
--- a/imr/utils.hh
+++ b/imr/utils.hh
@@ -162,10 +162,17 @@ public:
     }
 
     /// Create an IMR objects
-    template<typename Writer>
+    template<typename Writer, typename MigrateFn>
     GCC6_CONCEPT(requires WriterAllocator<Writer, Structure>)
     static object make(Writer&& object_writer,
-                       allocation_strategy::migrate_fn migrate = &imr::alloc::default_lsa_migrate_fn<structure>::migrate_fn) {
+                       MigrateFn* migrate = &imr::alloc::default_lsa_migrate_fn<structure>::migrate_fn) {
+        static_assert(std::is_same_v<typename MigrateFn::structure, structure>);
+        return do_make(std::forward<Writer>(object_writer), migrate);
+    }
+private:
+    template<typename Writer>
+    GCC6_CONCEPT(requires WriterAllocator<Writer, Structure>)
+    static object do_make(Writer&& object_writer, allocation_strategy::migrate_fn migrate) {
         struct alloc_deleter {
             size_t _size;