Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2025 Huawei Technologies Co., Ltd.
       3              :  * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
       4              :  * CANN Open Software License Agreement Version 2.0 (the "License").
       5              :  * Please refer to the License for details. You may not use this file except in compliance with the License.
       6              :  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
       7              :  * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
       8              :  * See LICENSE in the root of the software repository for the full text of the License.
       9              :  */
      10              : 
      11              : #ifndef HCCL_REFERENCE_MAP_H
      12              : #define HCCL_REFERENCE_MAP_H
      13              : 
      14              : #include <unordered_map>
      15              : 
      16              : #include "hccl/hccl_types.h"
      17              : #include "log.h"
      18              : 
      19              : namespace hccl {
      20              : 
      21              : template <typename keyType, typename valueType>
      22              : class ReferenceMap {
      23              : public:
      24         2078 :     typename std::unordered_map<keyType, valueType>::iterator begin() { return data_.begin(); }
      25              : 
      26         2171 :     typename std::unordered_map<keyType, valueType>::iterator end() { return data_.end(); }
      27              : 
      28           83 :     u32 insert(const keyType key, const valueType& value)
      29              :     {
      30           83 :         if (has(key)) {
      31           25 :             ref_[key]++;
      32              :         } else {
      33           58 :             data_.insert(std::make_pair(key, value));
      34           58 :             ref_[key] = 1;
      35              :         }
      36           83 :         return count(key);
      37              :     }
      38              : 
      39           76 :     u32 erase(const keyType key)
      40              :     {
      41           76 :         u32 refCount = count(key);
      42           76 :         if (refCount > 1) {
      43           15 :             ref_[key]--;
      44           61 :         } else if (refCount == 1) {
      45           30 :             data_.erase(key);
      46           30 :             ref_.erase(key);
      47              :         }
      48           76 :         return count(key);
      49              :     }
      50              : 
      51         2873 :     void clear()
      52              :     {
      53         2873 :         data_.clear();
      54         2873 :         ref_.clear();
      55         2873 :     }
      56              : 
      57          275 :     bool has(const keyType key) { return (data_.find(key) != data_.end() ? true : false); }
      58              : 
      59          258 :     u32 count(const keyType key) { return (has(key) ? ref_[key] : 0); }
      60              : 
      61          272 :     valueType& operator[](const keyType key) { return data_[key]; }
      62              : 
      63            4 :     HcclResult ref(const keyType key)
      64              :     {
      65            4 :         if (has(key)) {
      66            1 :             ref_[key]++;
      67              :         } else {
      68            3 :             return HCCL_E_PARA;
      69              :         }
      70            1 :         return HCCL_SUCCESS;
      71              :     }
      72              : 
      73            2 :     HcclResult unref(const keyType key)
      74              :     {
      75            2 :         if (has(key)) {
      76            1 :             ref_[key]--;
      77              :         } else {
      78            1 :             return HCCL_E_PARA;
      79              :         }
      80            1 :         return HCCL_SUCCESS;
      81              :     }
      82              : 
      83            1 :     u32 Size() { return data_.size(); }
      84              : 
      85              : private:
      86              :     std::unordered_map<keyType, valueType> data_;
      87              :     std::unordered_map<keyType, u32> ref_;
      88              : };
      89              : 
      90              : } // namespace hccl
      91              : #endif // HCCL_REFERENCE_MAP_H