ib_memstream.h 2.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  1. #pragma once
  2. #include "ib_mem.h"
  3. #include <library/cpp/binsaver/bin_saver.h>
  4. #include <library/cpp/binsaver/buffered_io.h>
  5. namespace NNetliba {
  6. class TIBMemStream: public IBinaryStream {
  7. struct TBlock {
  8. TIntrusivePtr<TIBMemBlock> Mem;
  9. i64 StartOffset;
  10. int BufSize, DataSize;
  11. TBlock()
  12. : StartOffset(0)
  13. , BufSize(0)
  14. , DataSize(0)
  15. {
  16. }
  17. TBlock(const TBlock& x) {
  18. Copy(x);
  19. }
  20. void operator=(const TBlock& x) {
  21. Copy(x);
  22. }
  23. void Copy(const TBlock& x) {
  24. if (x.BufSize > 0) {
  25. Mem = GetIBMemPool()->Alloc(x.BufSize);
  26. memcpy(Mem->GetData(), x.Mem->GetData(), x.DataSize);
  27. StartOffset = x.StartOffset;
  28. BufSize = x.BufSize;
  29. DataSize = x.DataSize;
  30. } else {
  31. Mem = nullptr;
  32. StartOffset = 0;
  33. BufSize = 0;
  34. DataSize = 0;
  35. }
  36. }
  37. };
  38. TIntrusivePtr<TIBMemPool> MemPool;
  39. TVector<TBlock> Blocks;
  40. int CurBlock;
  41. int CurBlockOffset;
  42. public:
  43. struct TBlockDescr {
  44. ui64 Addr;
  45. int BufSize, DataSize;
  46. ui32 RemoteKey, LocalKey;
  47. };
  48. struct TBlockSizes {
  49. int BufSize, DataSize;
  50. };
  51. public:
  52. TIBMemStream()
  53. : MemPool(GetIBMemPool())
  54. , CurBlock(0)
  55. , CurBlockOffset(0)
  56. {
  57. }
  58. ~TIBMemStream() override {
  59. } // keep gcc happy
  60. bool IsValid() const override {
  61. return true;
  62. }
  63. bool IsFailed() const override {
  64. return false;
  65. }
  66. void Flush() {
  67. }
  68. i64 GetLength();
  69. i64 Seek(i64 pos);
  70. void GetBlocks(TVector<TBlockDescr>* res) const;
  71. void CreateBlocks(const TVector<TBlockSizes>& arr);
  72. void Clear();
  73. private:
  74. int WriteImpl(const void* userBuffer, int size) override;
  75. int ReadImpl(void* userBuffer, int size) override;
  76. };
  77. template <class T>
  78. inline void Serialize(bool bRead, TIBMemStream& ms, T& c) {
  79. IBinSaver bs(ms, bRead);
  80. bs.Add(1, &c);
  81. }
  82. }