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         1612 :     explicit BinaryStream(std::ios_base::openmode mode = DEFAULT_IOS_MODE) : stream(mode | std::ios_base::binary) {};
      28              : 
      29          746 :     explicit BinaryStream(std::vector<char>& buf, std::ios_base::openmode mode = DEFAULT_IOS_MODE)
      30          746 :         : stream(mode | std::ios_base::binary)
      31              :     {
      32          746 :         stream.rdbuf()->pubsetbuf(buf.data(), buf.size());
      33          746 :     }
      34              : 
      35              :     template <typename T>
      36        40700 :     BinaryStream& operator<<(const T& t)
      37              :     {
      38        40700 :         stream.write(reinterpret_cast<const char*>(&t), sizeof(T));
      39        40700 :         return *this;
      40              :     }
      41              : 
      42              :     // 多级vector递归序列化
      43              :     template <typename T>
      44          740 :     BinaryStream& operator<<(const std::vector<T>& vec)
      45              :     {
      46          740 :         size_t size = vec.size();
      47          740 :         *this << size;
      48        31942 :         for (const auto& elem : vec) {
      49        31202 :             *this << elem;
      50              :         }
      51          740 :         return *this;
      52              :     }
      53              : 
      54              :     // 对string的输入函数
      55          944 :     BinaryStream& operator<<(const std::string& s)
      56              :     {
      57          944 :         size_t size = s.size();
      58          944 :         stream.write(reinterpret_cast<const char*>(&size), sizeof(size_t)); // 写入长度
      59          944 :         stream.write(s.data(), size);                                       // 写入字符数据
      60          944 :         return *this;
      61              :     }
      62              : 
      63              :     template <typename T>
      64        27659 :     BinaryStream& operator>>(T& t)
      65              :     {
      66        27659 :         stream.read(reinterpret_cast<char*>(&t), sizeof(T));
      67        27659 :         return *this;
      68              :     }
      69              : 
      70              :     // 对string的读取函数
      71          363 :     BinaryStream& operator>>(std::string& s)
      72              :     {
      73              :         size_t size;
      74          363 :         stream.read(reinterpret_cast<char*>(&size), sizeof(size)); // 先从流中读取字符串长度
      75          363 :         s.resize(size);                                            // 为string分配足够空间
      76          363 :         stream.read(&s[0], size); // 直接读取数据到string的缓冲区中，无需再分配内存
      77          363 :         return *this;
      78              :     }
      79              : 
      80              :     // 多级vector递归反序列化
      81              :     template <typename T>
      82          583 :     BinaryStream& operator>>(std::vector<T>& vec)
      83              :     {
      84              :         size_t size;
      85          583 :         *this >> size;
      86          583 :         vec.resize(size);
      87        22372 :         for (auto& elem : vec) {
      88        21789 :             *this >> elem;
      89              :         }
      90          583 :         return *this;
      91              :     }
      92              : 
      93          723 :     void Dump(std::vector<char>& vec)
      94              :     {
      95          723 :         std::for_each(std::istreambuf_iterator<char>(stream), std::istreambuf_iterator<char>(), [&vec](const char c) {
      96        70667 :             vec.push_back(c);
      97        70667 :         });
      98          723 :     }
      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() { return stream.str().size(); }
     110              : 
     111            6 :     std::string GetString() { return stream.str(); }
     112              : 
     113              :     std::string SplictStream(u64& start, u64& end)
     114              :     {
     115              :         std::string temp = stream.str();
     116              :         if (start >= temp.length()) {
     117              :             HCCL_ERROR("[SplictStream]start[%llu] is bigger than stream length[%llu]", start, temp.length());
     118              :             return "";
     119              :         }
     120              : 
     121              :         // 截取子串
     122              :         std::string result = temp.substr(start, end - start);
     123              : 
     124              :         // 返回新的 string
     125              :         return result;
     126              :     }
     127              : 
     128            8 :     void Clear() { stream.clear(); }
     129              : 
     130              : private:
     131              :     std::stringstream stream;
     132              : };
     133              : 
     134              : } // namespace Hccl
     135              : 
     136              : #endif // HCCL_SERIALIZATION_H