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_SERIALIZATION_H
      12              : #define HCCL_SERIALIZATION_H
      13              : 
      14              : #include <cstdint>
      15              : #include <algorithm>
      16              : 
      17              : #include "hccl/base.h"
      18              : #include "exception_util.h"
      19              : #include "log.h"
      20              : 
      21              : namespace Hccl {
      22              : 
      23              : class BinaryStream {
      24              : public:
      25              :     static constexpr std::ios_base::openmode DEFAULT_IOS_MODE = std::ios_base::in | std::ios_base::out;
      26              : 
      27         1591 :     explicit BinaryStream(std::ios_base::openmode mode = DEFAULT_IOS_MODE) : stream(mode | std::ios_base::binary){};
      28              : 
      29          731 :     explicit BinaryStream(std::vector<char> &buf, std::ios_base::openmode mode = DEFAULT_IOS_MODE)
      30          731 :         : stream(mode | std::ios_base::binary)
      31              :     {
      32          731 :         stream.rdbuf()->pubsetbuf(buf.data(), buf.size());
      33          731 :     }
      34              : 
      35              :     template <typename T>
      36        45391 :     BinaryStream &operator<<(const T &t)
      37              :     {
      38        45391 :         stream.write(reinterpret_cast<const char *>(&t), sizeof(T));
      39        45391 :         return *this;
      40              :     }
      41              : 
      42              :     // 多级vector递归序列化
      43              :     template <typename T>
      44          724 :     BinaryStream &operator<<(const std::vector<T> &vec)
      45              :     {
      46          724 :         size_t size = vec.size();
      47          724 :         *this << size;
      48        30038 :         for (const auto &elem : vec) {
      49        29314 :            *this << elem;
      50              :         }
      51          724 :         return *this;
      52              :     }
      53              : 
      54              :     // 对string的输入函数
      55         2343 :     BinaryStream &operator<<(const std::string &s) 
      56              :     {
      57         2343 :         size_t size = s.size();
      58         2343 :         stream.write(reinterpret_cast<const char*>(&size), sizeof(size_t)); // 写入长度
      59         2343 :         stream.write(s.data(), size); // 写入字符数据
      60         2343 :         return *this;
      61              :     }
      62              : 
      63              :     template <typename T>
      64        32788 :     BinaryStream &operator>>(T &t)
      65              :     {
      66        32788 :         stream.read(reinterpret_cast<char *>(&t), sizeof(T));
      67        32788 :         return *this;
      68              :     }
      69              : 
      70              :     // 对string的读取函数
      71         1762 :     BinaryStream &operator>>(std::string &s) 
      72              :     {
      73              :         size_t size;
      74         1762 :         stream.read(reinterpret_cast<char *>(&size), sizeof(size));       // 先从流中读取字符串长度
      75         1762 :         s.resize(size);       // 为string分配足够空间
      76         1762 :         stream.read(&s[0], size); // 直接读取数据到string的缓冲区中，无需再分配内存
      77         1762 :         return *this;
      78              :     }
      79              : 
      80              :     // 多级vector递归反序列化
      81              :     template <typename T>
      82          566 :     BinaryStream &operator>>(std::vector<T> &vec)
      83              :     {
      84              :         size_t size;
      85          566 :         *this >> size;
      86          566 :         vec.resize(size);
      87        20735 :         for (auto &elem : vec) {
      88        20169 :            *this >> elem;
      89              :         }
      90          566 :         return *this;
      91              :     }
      92              : 
      93          709 :     void Dump(std::vector<char> &vec)
      94              :     {
      95          709 :         std::for_each(std::istreambuf_iterator<char>(stream), std::istreambuf_iterator<char>(), [&vec](const char c) {
      96        67079 :             vec.push_back(c);
      97        67079 :         });
      98          709 :     }
      99              : 
     100            1 :     void DumpWithRevert(std::vector<char> &vec)
     101              :     {
     102            1 :         std::streampos originalPos = stream.tellg(); // 保存原始位置
     103            1 :         std::for_each(std::istreambuf_iterator<char>(stream), std::istreambuf_iterator<char>(), [&vec](const char c) {
     104            0 :             vec.push_back(c);
     105            0 :         });
     106            1 :         stream.seekg(originalPos); // 恢复原始位置
     107            1 :     }
     108              : 
     109           12 :     std::uint64_t GetSize()
     110              :     {
     111           12 :         return stream.str().size();
     112              :     }
     113              : 
     114            4 :     std::string GetString()
     115              :     {
     116            4 :         return stream.str();
     117              :     }
     118              : 
     119              :     std::string SplictStream(u64& start, u64& end){
     120              :         std::string temp = stream.str();
     121              :         if (start >= temp.length()) {
     122              :             HCCL_ERROR("[SplictStream]start[%llu] is bigger than stream length[%llu]", start, temp.length());
     123              :             return "";
     124              :         }
     125              :     
     126              :         // 截取子串
     127              :         std::string result = temp.substr(start, end - start);
     128              :     
     129              :         // 返回新的 string
     130              :         return result;
     131              :     }
     132              : 
     133            8 :     void Clear()
     134              :     {
     135            8 :         stream.clear();
     136            8 :     }
     137              : 
     138              : private:
     139              :     std::stringstream stream;
     140              : };
     141              : 
     142              : } // namespace Hccl
     143              : 
     144              : #endif // HCCL_SERIALIZATION_H