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         3642 :     typename std::unordered_map<keyType, valueType>::iterator begin()
      25              :     {
      26         3642 :         return data_.begin();
      27              :     }
      28              : 
      29         3735 :     typename std::unordered_map<keyType, valueType>::iterator end()
      30              :     {
      31         3735 :         return data_.end();
      32              :     }
      33              : 
      34           83 :     u32 insert(const keyType key, const valueType& value)
      35              :     {
      36           83 :         if (has(key)) {
      37           25 :             ref_[key]++;
      38              :         } else {
      39           58 :             data_.insert(std::make_pair(key, value));
      40           58 :             ref_[key] = 1;
      41              :         }
      42           83 :         return count(key);
      43              :     }
      44              : 
      45           76 :     u32 erase(const keyType key)
      46              :     {
      47           76 :         u32 refCount = count(key);
      48           76 :         if (refCount > 1) {
      49           15 :             ref_[key]--;
      50           61 :         } else if (refCount == 1) {
      51           30 :             data_.erase(key);
      52           30 :             ref_.erase(key);
      53              :         }
      54           76 :         return count(key);
      55              :     }
      56              : 
      57         7293 :     void clear()
      58              :     {
      59         7293 :         data_.clear();
      60         7293 :         ref_.clear();
      61         7293 :     }
      62              : 
      63          275 :     bool has(const keyType key)
      64              :     {
      65          267 :         return (data_.find(key) != data_.end() ? true : false);
      66              :     }
      67              : 
      68          242 :     u32 count(const keyType key)
      69              :     {
      70          258 :         return (has(key) ? ref_[key] : 0);
      71              :     }
      72              : 
      73          272 :     valueType& operator[](const keyType key)
      74              :     {
      75          266 :         return data_[key];
      76              :     }
      77              : 
      78            4 :     HcclResult ref(const keyType key)
      79              :     {
      80            4 :         if (has(key)) {
      81            1 :             ref_[key]++;
      82              :         } else {
      83            3 :             return HCCL_E_PARA;
      84              :         }
      85            1 :         return HCCL_SUCCESS;
      86              :     }
      87              : 
      88            2 :     HcclResult unref(const keyType key)
      89              :     {
      90            2 :         if (has(key)) {
      91            1 :             ref_[key]--;
      92              :         } else {
      93            1 :             return HCCL_E_PARA;
      94              :         }
      95            1 :         return HCCL_SUCCESS;
      96              :     }
      97              : 
      98            1 :     u32 Size()
      99              :     {
     100            1 :         return data_.size();
     101              :     }
     102              : private:
     103              :     std::unordered_map<keyType, valueType> data_;
     104              :     std::unordered_map<keyType, u32> ref_;
     105              : };
     106              : 
     107              : } // namespace hccl
     108              : #endif // HCCL_REFERENCE_MAP_H