be/src/vec/common/custom_allocator.h

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#pragma once

#include <cstddef>
#include <type_traits>

#include "vec/common/allocator.h"
#include "vec/common/allocator_fwd.h"

namespace doris {
#include "common/compile_check_begin.h"

template <class T, typename MemoryAllocator = Allocator<false>>
class CustomStdAllocator;

template <typename T>
using DorisVector = std::vector<T, CustomStdAllocator<T>>;

template <class Key, class T, class Compare = std::less<Key>,
          class Allocator = CustomStdAllocator<std::pair<const Key, T>>>
using DorisMap = std::map<Key, T, Compare, Allocator>;

template <typename T>
    requires(std::is_trivial_v<T> && std::is_standard_layout_v<T>)
struct DorisUniqueBufferDeleter : public Allocator<false> {
    size_t count = 0;

    DorisUniqueBufferDeleter() = default;
    DorisUniqueBufferDeleter(size_t n) : count(n) {}

    void operator()(T* ptr) const noexcept {
        if (ptr) {
            Allocator::free(ptr, count * sizeof(T));
        }
    }
};

template <typename T>
    requires(std::is_trivial_v<T> && std::is_standard_layout_v<T>)
class DorisUniqueBufferPtr {
public:
    using Deleter = DorisUniqueBufferDeleter<T>;

    DorisUniqueBufferPtr() = default;
    explicit DorisUniqueBufferPtr(T* ptr) = delete;

    DorisUniqueBufferPtr(size_t size) {
        DorisUniqueBufferDeleter<T> deleter(size);
        void* buf = deleter.alloc(size * sizeof(T));
        if (!buf) {
            return;
        }
        T* arr = static_cast<T*>(buf);
        ptr_ = std::unique_ptr<T[], DorisUniqueBufferDeleter<T>>(arr, std::move(deleter));
    }

    DorisUniqueBufferPtr(std::unique_ptr<T[], Deleter>&& uptr) : ptr_(std::move(uptr)) {}
    DorisUniqueBufferPtr(DorisUniqueBufferPtr&& other) : ptr_(std::move(other.ptr_)) {}

    DorisUniqueBufferPtr& operator=(DorisUniqueBufferPtr&& other) {
        if (this != &other) {
            ptr_ = std::move(other.ptr_);
        }
        return *this;
    }

    DorisUniqueBufferPtr(std::nullptr_t) noexcept {}

    // Delete this function to avoid passing in a pointer allocated by other means(new/malloc).
    void reset(T*) = delete;

    bool operator==(std::nullptr_t) const noexcept { return ptr_ == nullptr; }
    bool operator==(T* other) const noexcept { return ptr_.get() == other; }

    void reset() noexcept { ptr_.reset(); }

    auto release() noexcept { return ptr_.release(); }

    T* get() const noexcept { return ptr_.get(); }
    T& operator*() const noexcept { return *ptr_; }
    T* operator->() const noexcept { return ptr_.get(); }
    T& operator[](size_t i) const { return ptr_[i]; }

private:
    std::unique_ptr<T[], Deleter> ptr_;
};

template <typename T>
    requires(std::is_trivial_v<T> && std::is_standard_layout_v<T>)
DorisUniqueBufferPtr<T> make_unique_buffer(size_t n) {
    return DorisUniqueBufferPtr<T>(n);
}

// NOTE: Even CustomStdAllocator 's allocate/dallocate could modify memory tracker,but it's still stateless,
// because threadcontext owns the memtracker, not CustomStdAllocator.
template <class T, typename MemoryAllocator>
class CustomStdAllocator : private MemoryAllocator {
public:
    using value_type = T;
    using pointer = T*;
    using const_pointer = const T*;
    using size_type = std::size_t;
    using difference_type = std::ptrdiff_t;

    CustomStdAllocator() noexcept = default;

    template <class U>
    struct rebind {
        typedef CustomStdAllocator<U> other;
    };

    template <class Up>
    CustomStdAllocator(const CustomStdAllocator<Up>&) noexcept {}

    T* allocate(size_t n) { return static_cast<T*>(MemoryAllocator::alloc(n * sizeof(T))); }

    void deallocate(T* ptr, size_t n) noexcept { MemoryAllocator::free((void*)ptr, n * sizeof(T)); }

    size_t max_size() const noexcept { return size_t(~0) / sizeof(T); }

    T* allocate(size_t n, const void*) { return allocate(n); }

    // https://en.cppreference.com/w/cpp/memory/allocator/construct.html
    // void construct( pointer p, const_reference val ); (1)	(until C++11)
    //
    // template< class U, class... Args >
    // void construct( U* p, Args&&... args ); (2)	(since C++11)
    //                                              (deprecated in C++17)
    //                                              (removed in C++20)
    // template <class Up, class... Args>
    // void construct(Up* p, Args&&... args) {
    //     ::new ((void*)p) Up(std::forward<Args>(args)...);
    // }

    void destroy(T* p) { p->~T(); }

    T* address(T& t) const noexcept { return std::addressof(t); }

    T* address(const T& t) const noexcept { return std::addressof(t); }
};

template <class T, class Up>
bool operator==(const CustomStdAllocator<T>&, const CustomStdAllocator<Up>&) {
    return true;
}

template <class T, class Up>
bool operator!=(const CustomStdAllocator<T>&, const CustomStdAllocator<Up>&) {
    return false;
}

#include "common/compile_check_end.h"
} // namespace doris