1/* COPYRIGHT 2019 Huawei Technologies Co., Ltd.All Rights Reserved. 2 * 3 * Licensed under the Apache License, Version 2.0 (the "License"); 4 * you may not use this file except in compliance with the License. 5 * You may obtain a copy of the License at 6 * 7 * http://www.apache.org/licenses/LICENSE-2.0 8 * 9 * Unless required by applicable law or agreed to in writing, software 10 * distributed under the License is distributed on an "AS IS" BASIS, 11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 * See the License for the specific language governing permissions and 13 * limitations under the License. 14 */ 15#ifndef DATASET_UTIL_BTREE_ITERATOR_H_ 16#define DATASET_UTIL_BTREE_ITERATOR_H_ 17 18#include "securec.h" 19#include "minddata/dataset/util/log_adapter.h" 20#include "btree.h" 21 22namespace mindspore { 23namespace dataset { 24template <typename K, typename V, typename A, typename C, typename T> 25BPlusTree<K, V, A, C, T>::Iterator::~Iterator() { 26 if (locked_) { 27 cur_->rw_lock_.Unlock(); 28 locked_ = false; 29 } 30} 31 32template <typename K, typename V, typename A, typename C, typename T> 33typename BPlusTree<K, V, A, C, T>::Iterator &BPlusTree<K, V, A, C, T>::Iterator::operator++() { 34 if (slot_ + 1u < cur_->slotuse_) { 35 ++slot_; 36 } else if (cur_->link_.next) { 37 if (locked_) { 38 cur_->link_.next->rw_lock_.LockShared(); 39 cur_->rw_lock_.Unlock(); 40 } 41 cur_ = cur_->link_.next; 42 slot_ = 0; 43 } else { 44 slot_ = cur_->slotuse_; 45 } 46 return *this; 47} 48 49template <typename K, typename V, typename A, typename C, typename T> 50typename BPlusTree<K, V, A, C, T>::Iterator BPlusTree<K, V, A, C, T>::Iterator::operator++(int) { 51 Iterator tmp = *this; 52 if (slot_ + 1u < cur_->slotuse_) { 53 ++slot_; 54 } else if (cur_->link_.next) { 55 if (locked_) { 56 cur_->link_.next->rw_lock_.LockShared(); 57 cur_->rw_lock_.Unlock(); 58 } 59 cur_ = cur_->link_.next; 60 slot_ = 0; 61 } else { 62 slot_ = cur_->slotuse_; 63 } 64 return tmp; 65} 66 67template <typename K, typename V, typename A, typename C, typename T> 68typename BPlusTree<K, V, A, C, T>::Iterator &BPlusTree<K, V, A, C, T>::Iterator::operator--() { 69 if (slot_ > 0) { 70 --slot_; 71 } else if (cur_->link_.prev) { 72 if (locked_) { 73 cur_->link_.prev->rw_lock_.LockShared(); 74 cur_->rw_lock_.Unlock(); 75 } 76 cur_ = cur_->link_.prev; 77 slot_ = cur_->slotuse_ - 1; 78 } else { 79 slot_ = 0; 80 } 81 return *this; 82} 83 84template <typename K, typename V, typename A, typename C, typename T> 85typename BPlusTree<K, V, A, C, T>::Iterator BPlusTree<K, V, A, C, T>::Iterator::operator--(int) { 86 Iterator tmp = *this; 87 if (slot_ > 0) { 88 --slot_; 89 } else if (cur_->link_.prev) { 90 if (locked_) { 91 cur_->link_.prev->rw_lock_.LockShared(); 92 cur_->rw_lock_.Unlock(); 93 } 94 cur_ = cur_->link_.prev; 95 slot_ = cur_->slotuse_ - 1; 96 } else { 97 slot_ = 0; 98 } 99 return tmp; 100} 101 102template <typename K, typename V, typename A, typename C, typename T> 103BPlusTree<K, V, A, C, T>::Iterator::Iterator(const BPlusTree<K, V, A, C, T>::Iterator &lhs) { 104 this->cur_ = lhs.cur_; 105 this->slot_ = lhs.slot_; 106 this->locked_ = lhs.locked_; 107 if (this->locked_) { 108 this->cur_->rw_lock_.LockShared(); 109 } 110} 111 112template <typename K, typename V, typename A, typename C, typename T> 113BPlusTree<K, V, A, C, T>::Iterator::Iterator(BPlusTree<K, V, A, C, T>::Iterator &&lhs) noexcept { 114 this->cur_ = lhs.cur_; 115 this->slot_ = lhs.slot_; 116 this->locked_ = lhs.locked_; 117 lhs.locked_ = false; 118 lhs.slot_ = 0; 119 lhs.cur_ = nullptr; 120} 121 122template <typename K, typename V, typename A, typename C, typename T> 123typename BPlusTree<K, V, A, C, T>::Iterator &BPlusTree<K, V, A, C, T>::Iterator::operator=( 124 const BPlusTree<K, V, A, C, T>::Iterator &lhs) { 125 if (*this != lhs) { 126 if (this->locked_) { 127 this->cur_->rw_lock_.Unlock(); 128 } 129 this->cur_ = lhs.cur_; 130 this->slot_ = lhs.slot_; 131 this->locked_ = lhs.locked_; 132 if (this->locked_) { 133 this->cur_->rw_lock_.LockShared(); 134 } 135 } 136 return *this; 137} 138 139template <typename K, typename V, typename A, typename C, typename T> 140typename BPlusTree<K, V, A, C, T>::Iterator &BPlusTree<K, V, A, C, T>::Iterator::operator=( 141 BPlusTree<K, V, A, C, T>::Iterator &&lhs) { 142 if (*this != lhs) { 143 if (this->locked_) { 144 this->cur_->rw_lock_.Unlock(); 145 } 146 this->cur_ = lhs.cur_; 147 this->slot_ = lhs.slot_; 148 this->locked_ = lhs.locked_; 149 lhs.locked_ = false; 150 lhs.slot_ = 0; 151 lhs.cur_ = nullptr; 152 } 153 return *this; 154} 155 156template <typename K, typename V, typename A, typename C, typename T> 157BPlusTree<K, V, A, C, T>::ConstIterator::~ConstIterator() { 158 if (locked_) { 159 cur_->rw_lock_.Unlock(); 160 locked_ = false; 161 } 162} 163 164template <typename K, typename V, typename A, typename C, typename T> 165typename BPlusTree<K, V, A, C, T>::ConstIterator &BPlusTree<K, V, A, C, T>::ConstIterator::operator++() { 166 if (slot_ + 1u < cur_->slotuse_) { 167 ++slot_; 168 } else if (cur_->link_.next) { 169 if (locked_) { 170 cur_->link_.next->rw_lock_.LockShared(); 171 cur_->rw_lock_.Unlock(); 172 } 173 cur_ = cur_->link_.next; 174 slot_ = 0; 175 } else { 176 slot_ = cur_->slotuse_; 177 } 178 return *this; 179} 180 181template <typename K, typename V, typename A, typename C, typename T> 182typename BPlusTree<K, V, A, C, T>::ConstIterator BPlusTree<K, V, A, C, T>::ConstIterator::operator++(int) { 183 Iterator tmp = *this; 184 if (slot_ + 1u < cur_->slotuse_) { 185 ++slot_; 186 } else if (cur_->link_.next) { 187 if (locked_) { 188 cur_->link_.next->rw_lock_.LockShared(); 189 cur_->rw_lock_.Unlock(); 190 } 191 cur_ = cur_->link_.next; 192 slot_ = 0; 193 } else { 194 slot_ = cur_->slotuse_; 195 } 196 return tmp; 197} 198 199template <typename K, typename V, typename A, typename C, typename T> 200typename BPlusTree<K, V, A, C, T>::ConstIterator &BPlusTree<K, V, A, C, T>::ConstIterator::operator--() { 201 if (slot_ > 0) { 202 --slot_; 203 } else if (cur_->link_.prev) { 204 if (locked_) { 205 cur_->link_.prev->rw_lock_.LockShared(); 206 cur_->rw_lock_.Unlock(); 207 } 208 cur_ = cur_->link_.prev; 209 slot_ = cur_->slotuse_ - 1; 210 } else { 211 slot_ = 0; 212 } 213 return *this; 214} 215 216template <typename K, typename V, typename A, typename C, typename T> 217typename BPlusTree<K, V, A, C, T>::ConstIterator BPlusTree<K, V, A, C, T>::ConstIterator::operator--(int) { 218 Iterator tmp = *this; 219 if (slot_ > 0) { 220 --slot_; 221 } else if (cur_->link_.prev) { 222 if (locked_) { 223 cur_->link_.prev->rw_lock_.LockShared(); 224 cur_->rw_lock_.Unlock(); 225 } 226 cur_ = cur_->link_.prev; 227 slot_ = cur_->slotuse_ - 1; 228 } else { 229 slot_ = 0; 230 } 231 return tmp; 232} 233 234template <typename K, typename V, typename A, typename C, typename T> 235BPlusTree<K, V, A, C, T>::ConstIterator::ConstIterator(const BPlusTree<K, V, A, C, T>::ConstIterator &lhs) { 236 this->cur_ = lhs.cur_; 237 this->slot_ = lhs.slot_; 238 this->locked_ = lhs.locked_; 239 if (this->locked_) { 240 this->cur_->rw_lock_.LockShared(); 241 } 242} 243 244template <typename K, typename V, typename A, typename C, typename T> 245BPlusTree<K, V, A, C, T>::ConstIterator::ConstIterator(BPlusTree<K, V, A, C, T>::ConstIterator &&lhs) noexcept { 246 this->cur_ = lhs.cur_; 247 this->slot_ = lhs.slot_; 248 this->locked_ = lhs.locked_; 249 lhs.locked_ = false; 250 lhs.slot_ = 0; 251 lhs.cur_ = nullptr; 252} 253 254template <typename K, typename V, typename A, typename C, typename T> 255typename BPlusTree<K, V, A, C, T>::ConstIterator &BPlusTree<K, V, A, C, T>::ConstIterator::operator=( 256 const BPlusTree<K, V, A, C, T>::ConstIterator &lhs) { 257 if (*this != lhs) { 258 if (this->locked_) { 259 this->cur_->rw_lock_.Unlock(); 260 } 261 this->cur_ = lhs.cur_; 262 this->slot_ = lhs.slot_; 263 this->locked_ = lhs.locked_; 264 if (this->locked_) { 265 this->cur_->rw_lock_.LockShared(); 266 } 267 } 268 return *this; 269} 270 271template <typename K, typename V, typename A, typename C, typename T> 272typename BPlusTree<K, V, A, C, T>::ConstIterator &BPlusTree<K, V, A, C, T>::ConstIterator::operator=( 273 BPlusTree<K, V, A, C, T>::ConstIterator &&lhs) { 274 if (*this != lhs) { 275 if (this->locked_) { 276 this->cur_->rw_lock_.Unlock(); 277 } 278 this->cur_ = lhs.cur_; 279 this->slot_ = lhs.slot_; 280 this->locked_ = lhs.locked_; 281 lhs.locked_ = false; 282 lhs.slot_ = 0; 283 lhs.cur_ = nullptr; 284 } 285 return *this; 286} 287 288template <typename K, typename V, typename A, typename C, typename T> 289std::pair<typename BPlusTree<K, V, A, C, T>::ConstIterator, bool> BPlusTree<K, V, A, C, T>::Search( 290 const key_type &key) const { 291 if (root_ != nullptr) { 292 LeafNode *leaf = nullptr; 293 slot_type slot; 294 RWLock *myLock = nullptr; 295 if (acquire_lock_) { 296 myLock = &this->rw_lock_; 297 // Lock the tree in S, pass the lock to Locate which will unlock it for us underneath. 298 myLock->LockShared(); 299 } 300 IndexRc rc = Locate(myLock, false, root_, key, &leaf, &slot); 301 bool find = (rc == IndexRc::kOk); 302 return std::make_pair(ConstIterator(leaf, slot, find), find); 303 } else { 304 return std::make_pair(cend(), false); 305 } 306} 307 308template <typename K, typename V, typename A, typename C, typename T> 309std::pair<typename BPlusTree<K, V, A, C, T>::Iterator, bool> BPlusTree<K, V, A, C, T>::Search(const key_type &key) { 310 if (root_ != nullptr) { 311 LeafNode *leaf = nullptr; 312 slot_type slot; 313 RWLock *myLock = nullptr; 314 if (acquire_lock_) { 315 myLock = &this->rw_lock_; 316 // Lock the tree in S, pass the lock to Locate which will unlock it for us underneath. 317 myLock->LockShared(); 318 } 319 IndexRc rc = Locate(myLock, false, root_, key, &leaf, &slot); 320 bool find = (rc == IndexRc::kOk); 321 return std::make_pair(Iterator(leaf, slot, find), find); 322 } else { 323 return std::make_pair(end(), false); 324 } 325} 326 327template <typename K, typename V, typename A, typename C, typename T> 328typename BPlusTree<K, V, A, C, T>::value_type BPlusTree<K, V, A, C, T>::operator[](key_type key) { 329 auto r = Search(key); 330 return r.first.value(); 331} 332 333template <typename K, typename V, typename A, typename C, typename T> 334typename BPlusTree<K, V, A, C, T>::Iterator BPlusTree<K, V, A, C, T>::begin() { 335 return Iterator(this); 336} 337 338template <typename K, typename V, typename A, typename C, typename T> 339typename BPlusTree<K, V, A, C, T>::Iterator BPlusTree<K, V, A, C, T>::end() { 340 return Iterator(this->leaf_nodes_.tail, this->leaf_nodes_.tail ? this->leaf_nodes_.tail->slotuse_ : 0); 341} 342 343template <typename K, typename V, typename A, typename C, typename T> 344typename BPlusTree<K, V, A, C, T>::ConstIterator BPlusTree<K, V, A, C, T>::begin() const { 345 return ConstIterator(this); 346} 347 348template <typename K, typename V, typename A, typename C, typename T> 349typename BPlusTree<K, V, A, C, T>::ConstIterator BPlusTree<K, V, A, C, T>::end() const { 350 return ConstIterator(this->leaf_nodes_.tail, this->leaf_nodes_.tail ? this->leaf_nodes_.tail->slotuse_ : 0); 351} 352 353template <typename K, typename V, typename A, typename C, typename T> 354typename BPlusTree<K, V, A, C, T>::ConstIterator BPlusTree<K, V, A, C, T>::cbegin() const { 355 return ConstIterator(this); 356} 357 358template <typename K, typename V, typename A, typename C, typename T> 359typename BPlusTree<K, V, A, C, T>::ConstIterator BPlusTree<K, V, A, C, T>::cend() const { 360 return ConstIterator(this->leaf_nodes_.tail, this->leaf_nodes_.tail ? this->leaf_nodes_.tail->slotuse_ : 0); 361} 362} // namespace dataset 363} // namespace mindspore 364#endif 365