arrow/util/bitmap_generate.h

0001 // Licensed to the Apache Software Foundation (ASF) under one
0002 // or more contributor license agreements.  See the NOTICE file
0003 // distributed with this work for additional information
0004 // regarding copyright ownership.  The ASF licenses this file
0005 // to you under the Apache License, Version 2.0 (the
0006 // "License"); you may not use this file except in compliance
0007 // with the License.  You may obtain a copy of the License at
0008 //
0009 //   http://www.apache.org/licenses/LICENSE-2.0
0010 //
0011 // Unless required by applicable law or agreed to in writing,
0012 // software distributed under the License is distributed on an
0013 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0014 // KIND, either express or implied.  See the License for the
0015 // specific language governing permissions and limitations
0016 // under the License.
0017
0018 #pragma once
0019
0020 #include <cstdint>
0021 #include <memory>
0022
0023 #include "arrow/buffer.h"
0024 #include "arrow/memory_pool.h"
0025 #include "arrow/result.h"
0026 #include "arrow/util/bit_util.h"
0027 #include "arrow/util/visibility.h"
0028
0029 namespace arrow {
0030 namespace internal {
0031
0032 // A std::generate() like function to write sequential bits into a bitmap area.
0033 // Bits preceding the bitmap area are preserved, bits following the bitmap
0034 // area may be clobbered.
0035
0036 template <class Generator>
0037 void GenerateBits(uint8_t* bitmap, int64_t start_offset, int64_t length, Generator&& g) {
0038   if (length == 0) {
0039     return;
0040   }
0041   uint8_t* cur = bitmap + start_offset / 8;
0042   uint8_t bit_mask = bit_util::kBitmask[start_offset % 8];
0043   uint8_t current_byte = *cur & bit_util::kPrecedingBitmask[start_offset % 8];
0044
0045   for (int64_t index = 0; index < length; ++index) {
0046     const bool bit = g();
0047     current_byte = bit ? (current_byte | bit_mask) : current_byte;
0048     bit_mask = static_cast<uint8_t>(bit_mask << 1);
0049     if (bit_mask == 0) {
0050       bit_mask = 1;
0051       *cur++ = current_byte;
0052       current_byte = 0;
0053     }
0054   }
0055   if (bit_mask != 1) {
0056     *cur++ = current_byte;
0057   }
0058 }
0059
0060 // Like GenerateBits(), but unrolls its main loop for higher performance.
0061
0062 template <class Generator>
0063 void GenerateBitsUnrolled(uint8_t* bitmap, int64_t start_offset, int64_t length,
0064                           Generator&& g) {
0065   static_assert(std::is_same<decltype(std::declval<Generator>()()), bool>::value,
0066                 "Functor passed to GenerateBitsUnrolled must return bool");
0067
0068   if (length == 0) {
0069     return;
0070   }
0071   uint8_t current_byte;
0072   uint8_t* cur = bitmap + start_offset / 8;
0073   const uint64_t start_bit_offset = start_offset % 8;
0074   uint8_t bit_mask = bit_util::kBitmask[start_bit_offset];
0075   int64_t remaining = length;
0076
0077   if (bit_mask != 0x01) {
0078     current_byte = *cur & bit_util::kPrecedingBitmask[start_bit_offset];
0079     while (bit_mask != 0 && remaining > 0) {
0080       current_byte |= g() * bit_mask;
0081       bit_mask = static_cast<uint8_t>(bit_mask << 1);
0082       --remaining;
0083     }
0084     *cur++ = current_byte;
0085   }
0086
0087   int64_t remaining_bytes = remaining / 8;
0088   uint8_t out_results[8];
0089   while (remaining_bytes-- > 0) {
0090     for (int i = 0; i < 8; ++i) {
0091       out_results[i] = g();
0092     }
0093     *cur++ = static_cast<uint8_t>(out_results[0] | out_results[1] << 1 |
0094                                   out_results[2] << 2 | out_results[3] << 3 |
0095                                   out_results[4] << 4 | out_results[5] << 5 |
0096                                   out_results[6] << 6 | out_results[7] << 7);
0097   }
0098
0099   int64_t remaining_bits = remaining % 8;
0100   if (remaining_bits) {
0101     current_byte = 0;
0102     bit_mask = 0x01;
0103     while (remaining_bits-- > 0) {
0104       current_byte |= g() * bit_mask;
0105       bit_mask = static_cast<uint8_t>(bit_mask << 1);
0106     }
0107     *cur++ = current_byte;
0108   }
0109 }
0110
0111 }  // namespace internal
0112 }  // namespace arrow