neptools

Modding tools to Neptunia games
git clone https://git.neptards.moe/neptards/neptools.git
Log | Files | Refs | Submodules | README | LICENSE

sink.cpp (12768B)


      1 #include "sink.hpp"
      2 
      3 #include <libshit/except.hpp>
      4 #include <libshit/low_io.hpp>
      5 #include <libshit/lua/boost_endian_traits.hpp>
      6 
      7 #include <iostream>
      8 #include <fstream>
      9 
     10 #include <libshit/doctest.hpp>
     11 
     12 #define LIBSHIT_LOG_NAME "sink"
     13 #include <libshit/logger_helper.hpp>
     14 
     15 namespace Neptools
     16 {
     17   TEST_SUITE_BEGIN("Neptools::Sink");
     18 
     19   namespace
     20   {
     21     struct LIBSHIT_NOLUA MmapSink final : public Sink
     22     {
     23       MmapSink(Libshit::LowIo&& io, FilePosition size);
     24       void Write_(Libshit::StringView data) override;
     25       void Pad_(FileMemSize len) override;
     26 
     27       void MapNext(FileMemSize len);
     28 
     29       Libshit::LowIo io;
     30       Libshit::LowIo::MmapPtr mm;
     31     };
     32 
     33     struct LIBSHIT_NOLUA SimpleSink final : public Sink
     34     {
     35       SimpleSink(Libshit::LowIo io, FilePosition size)
     36         : Sink{size}, io{Libshit::Move(io)}
     37       {
     38         Sink::buf = buf;
     39         buf_size = MEM_CHUNK;
     40       }
     41       ~SimpleSink() override;
     42 
     43       void Write_(Libshit::StringView data) override;
     44       void Pad_(FileMemSize len) override;
     45       void Flush() override;
     46 
     47       Libshit::LowIo io;
     48       Byte buf[MEM_CHUNK];
     49     };
     50   }
     51 
     52   MmapSink::MmapSink(Libshit::LowIo&& io, FilePosition size) : Sink{size}
     53   {
     54     size_t to_map = size < MMAP_LIMIT ? size : MMAP_CHUNK;
     55 
     56     io.Truncate(size);
     57     io.PrepareMmap(true);
     58     mm = io.Mmap(0, to_map, true);
     59     buf_size = to_map;
     60     buf = reinterpret_cast<Neptools::Byte*>(mm.Get());
     61 
     62     this->io = Libshit::Move(io);
     63   }
     64 
     65   void MmapSink::Write_(Libshit::StringView data)
     66   {
     67     LIBSHIT_ASSERT(buf_put == buf_size && offset < size &&
     68                    buf_size == MMAP_CHUNK);
     69 
     70     offset += buf_put;
     71     if (data.length() / MMAP_CHUNK)
     72     {
     73       auto to_write = data.length() / MMAP_CHUNK * MMAP_CHUNK;
     74       io.Pwrite(data.udata(), to_write, offset);
     75       data.remove_prefix(to_write);
     76       offset += to_write;
     77       buf_put = 0;
     78     }
     79 
     80     MapNext(data.length());
     81     if (buf) // https://stackoverflow.com/a/5243068; C99 7.21.1/2
     82       memcpy(buf, data.data(), data.length());
     83     else
     84       LIBSHIT_ASSERT(data.length() == 0);
     85   }
     86 
     87   void MmapSink::Pad_(FileMemSize len)
     88   {
     89     LIBSHIT_ASSERT(buf_put == buf_size && offset < size &&
     90                    buf_size == MMAP_CHUNK);
     91 
     92     offset += buf_put + len / MMAP_CHUNK * MMAP_CHUNK;
     93     LIBSHIT_ASSERT_MSG(offset <= size, "sink overflow");
     94     MapNext(len % MMAP_CHUNK);
     95   }
     96 
     97   void MmapSink::MapNext(FileMemSize len)
     98   {
     99     // wine fails on 0 size
    100     // windows fails if offset+size > file_length...
    101     // (linux doesn't care...)
    102     if (offset < size)
    103     {
    104       auto nbuf_size = std::min<FileMemSize>(MMAP_CHUNK, size-offset);
    105       LIBSHIT_ASSERT(nbuf_size >= len);
    106       mm = io.Mmap(offset, nbuf_size, true);
    107       buf = static_cast<Byte*>(mm.Get());
    108       buf_put = len;
    109       buf_size = nbuf_size;
    110     }
    111     else
    112     {
    113       mm.Reset();
    114       buf = nullptr;
    115     }
    116   }
    117 
    118   SimpleSink::~SimpleSink()
    119   {
    120     try { Flush(); }
    121     catch (std::exception& e)
    122     { ERR << "~SimpleSink " << Libshit::PrintException(true) << std::endl; }
    123   }
    124 
    125   void SimpleSink::Flush()
    126   {
    127     if (buf_put)
    128     {
    129       io.Write(buf, buf_put);
    130       offset += buf_put;
    131       buf_put = 0;
    132     }
    133   }
    134 
    135   void SimpleSink::Write_(Libshit::StringView data)
    136   {
    137     LIBSHIT_ASSERT(buf_size == MEM_CHUNK && buf_put == MEM_CHUNK);
    138     io.Write(buf, MEM_CHUNK);
    139     offset += MEM_CHUNK;
    140 
    141     if (data.length() >= MEM_CHUNK)
    142     {
    143       io.Write(data.data(), data.length());
    144       offset += data.length();
    145       buf_put = 0;
    146     }
    147     else
    148     {
    149       memcpy(buf, data.data(), data.length());
    150       buf_put = data.length();
    151     }
    152   }
    153 
    154   void SimpleSink::Pad_(FileMemSize len)
    155   {
    156     LIBSHIT_ASSERT(buf_size == MEM_CHUNK && buf_put == MEM_CHUNK);
    157     io.Write(buf, MEM_CHUNK);
    158     offset += MEM_CHUNK;
    159 
    160     // assume we're not seekable (I don't care about not mmap-able but seekable
    161     // files)
    162     if (len >= MEM_CHUNK)
    163     {
    164       memset(buf, 0, MEM_CHUNK);
    165       size_t i;
    166       for (i = MEM_CHUNK; i < len; i += MEM_CHUNK)
    167         io.Write(buf, MEM_CHUNK);
    168       offset += i - MEM_CHUNK;
    169     }
    170     else
    171       memset(buf, 0, len);
    172     buf_put = len % MEM_CHUNK;
    173   }
    174 
    175   Libshit::NotNull<Libshit::RefCountedPtr<Sink>> Sink::ToFile(
    176     boost::filesystem::path fname, FilePosition size, bool try_mmap)
    177   {
    178     return Libshit::AddInfo(
    179       [&]() -> Libshit::NotNull<Libshit::RefCountedPtr<Sink>>
    180       {
    181         Libshit::LowIo io{fname.c_str(), Libshit::LowIo::Permission::READ_WRITE,
    182           Libshit::LowIo::Mode::TRUNC_OR_CREATE};
    183         if (LIBSHIT_OS_IS_VITA || !try_mmap)
    184           return Libshit::MakeRefCounted<SimpleSink>(Libshit::Move(io), size);
    185 
    186         try { return Libshit::MakeRefCounted<MmapSink>(Libshit::Move(io), size); }
    187         catch (const Libshit::SystemError& e)
    188         {
    189           WARN << "Mmmap failed, falling back to normal writing: "
    190                << Libshit::PrintException(true) << std::endl;
    191           return Libshit::MakeRefCounted<SimpleSink>(Libshit::Move(io), size);
    192         }
    193       },
    194       [&](auto& e) { Libshit::AddInfos(e, "File name", fname.string()); });
    195   }
    196 
    197   Libshit::NotNull<Libshit::RefCountedPtr<Sink>> Sink::ToStdOut()
    198   {
    199     return Libshit::MakeRefCounted<SimpleSink>(Libshit::LowIo::OpenStdOut(), -1);
    200   }
    201 
    202 #define TRY_MMAP                           \
    203   bool try_mmap;                           \
    204   SUBCASE("mmap") { try_mmap = false; }    \
    205   SUBCASE("no mmap") { try_mmap = true; }  \
    206   CAPTURE(try_mmap)
    207 
    208   TEST_CASE("small simple write")
    209   {
    210     TRY_MMAP;
    211     char buf[16] = {0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15};
    212     {
    213       auto sink = Sink::ToFile("neptools_test.tmp", 16, try_mmap);
    214       REQUIRE(sink->Tell() == 0);
    215       sink->WriteGen(buf);
    216       REQUIRE(sink->Tell() == 16);
    217     }
    218 
    219     char buf2[16];
    220     std::ifstream is{"neptools_test.tmp", std::ios_base::binary};
    221     is.read(buf2, 16);
    222     REQUIRE(is.good());
    223     REQUIRE(memcmp(buf, buf2, 16) == 0);
    224 
    225     is.get();
    226     REQUIRE(is.eof());
    227   }
    228 
    229   TEST_CASE("many small writes")
    230   {
    231     TRY_MMAP;
    232     int buf[6] = {0,77,-123,98,77,-1};
    233     static_assert(sizeof(buf) == 24);
    234 
    235     static constexpr FilePosition SIZE = 2*1024*1024 / 24 * 24;
    236     {
    237       auto sink = Sink::ToFile("neptools_test.tmp", SIZE, try_mmap);
    238       for (FilePosition i = 0; i < SIZE; i += 24)
    239       {
    240         buf[0] = static_cast<int>(i / 24);
    241         sink->WriteGen(buf);
    242       }
    243       REQUIRE(sink->Tell() == SIZE);
    244     }
    245 
    246     std::unique_ptr<char[]> buf_exp{new char[SIZE]};
    247     for (FilePosition i = 0; i < SIZE; i += 24)
    248     {
    249       buf[0] = static_cast<int>(i / 24);
    250       memcpy(buf_exp.get()+i, buf, 24);
    251     }
    252 
    253     std::unique_ptr<char[]> buf_act{new char[SIZE]};
    254     std::ifstream is{"neptools_test.tmp", std::ios_base::binary};
    255     is.read(buf_act.get(), SIZE);
    256     REQUIRE(is.good());
    257     REQUIRE(memcmp(buf_exp.get(), buf_act.get(), SIZE) == 0);
    258 
    259     is.get();
    260     REQUIRE(is.eof());
    261   }
    262 
    263   TEST_CASE("big write")
    264   {
    265     TRY_MMAP;
    266     static constexpr FilePosition SIZE = 2*1024*1024;
    267     std::unique_ptr<Byte[]> buf{new Byte[SIZE]};
    268     for (size_t i = 0; i < SIZE; ++i)
    269       buf[i] = static_cast<Byte>(i);
    270 
    271     {
    272       auto sink = Sink::ToFile("neptools_test.tmp", SIZE, try_mmap);
    273       REQUIRE(sink->Tell() == 0);
    274       sink->Write({buf.get(), SIZE});
    275       REQUIRE(sink->Tell() == SIZE);
    276     }
    277 
    278     std::unique_ptr<char[]> buf2{new char[SIZE]};
    279     std::ifstream is{"neptools_test.tmp", std::ios_base::binary};
    280     is.read(buf2.get(), SIZE);
    281     REQUIRE(is.good());
    282     REQUIRE(memcmp(buf.get(), buf2.get(), SIZE) == 0);
    283 
    284     is.get();
    285     REQUIRE(is.eof());
    286   }
    287 
    288   TEST_CASE("small pad")
    289   {
    290     TRY_MMAP;
    291     char buf[16] = {0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15};
    292     {
    293       auto sink = Sink::ToFile("neptools_test.tmp", 16*3, try_mmap);
    294       REQUIRE(sink->Tell() == 0);
    295       sink->WriteGen(buf);
    296       REQUIRE(sink->Tell() == 16);
    297       sink->Pad(16);
    298       REQUIRE(sink->Tell() == 2*16);
    299       sink->WriteGen(buf);
    300       REQUIRE(sink->Tell() == 3*16);
    301     }
    302 
    303     char buf2[16];
    304     std::ifstream is{"neptools_test.tmp", std::ios_base::binary};
    305     is.read(buf2, 16);
    306     REQUIRE(is.good());
    307     REQUIRE(memcmp(buf, buf2, 16) == 0);
    308 
    309     is.read(buf2, 16);
    310     REQUIRE(is.good());
    311     REQUIRE(std::count(buf2, buf2+16, 0) == 16);
    312 
    313     is.read(buf2, 16);
    314     REQUIRE(is.good());
    315     REQUIRE(memcmp(buf, buf2, 16) == 0);
    316 
    317     is.get();
    318     REQUIRE(is.eof());
    319   }
    320 
    321   TEST_CASE("many small pad")
    322   {
    323     TRY_MMAP;
    324     char buf[10] = {4,5,6,7,8,9,10,11,12,13};
    325     static constexpr FilePosition SIZE = 2*1024*1024 / 24 * 24;
    326     {
    327       auto sink = Sink::ToFile("neptools_test.tmp", SIZE, try_mmap);
    328 
    329       for (FilePosition i = 0; i < SIZE; i += 24)
    330       {
    331         sink->WriteGen(buf);
    332         sink->Pad(14);
    333       }
    334       REQUIRE(sink->Tell() == SIZE);
    335     }
    336 
    337     std::unique_ptr<char[]> buf_exp{new char[SIZE]};
    338     for (FilePosition i = 0; i < SIZE; i += 24)
    339     {
    340       memcpy(buf_exp.get()+i, buf, 10);
    341       memset(buf_exp.get()+i+10, 0, 14);
    342     }
    343 
    344     std::unique_ptr<char[]> buf_act{new char[SIZE]};
    345     std::ifstream is{"neptools_test.tmp", std::ios_base::binary};
    346     is.read(buf_act.get(), SIZE);
    347     REQUIRE(is.good());
    348     REQUIRE(memcmp(buf_exp.get(), buf_act.get(), SIZE) == 0);
    349 
    350     is.get();
    351     REQUIRE(is.eof());
    352   }
    353 
    354   TEST_CASE("large pad")
    355   {
    356     TRY_MMAP;
    357     char buf[16] = {0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15};
    358     static constexpr FilePosition ZERO_SIZE = 2*1024*1024;
    359     {
    360       auto sink = Sink::ToFile("neptools_test.tmp", 16*2 + ZERO_SIZE, try_mmap);
    361       REQUIRE(sink->Tell() == 0);
    362       sink->WriteGen(buf);
    363       REQUIRE(sink->Tell() == 16);
    364       sink->Pad(ZERO_SIZE);
    365       REQUIRE(sink->Tell() == 16+ZERO_SIZE);
    366       sink->WriteGen(buf);
    367       REQUIRE(sink->Tell() == 2*16+ZERO_SIZE);
    368     }
    369 
    370     std::unique_ptr<char[]> buf2{new char[ZERO_SIZE]};
    371     std::ifstream is{"neptools_test.tmp", std::ios_base::binary};
    372     is.read(buf2.get(), 16);
    373     REQUIRE(is.good());
    374     REQUIRE(memcmp(buf, buf2.get(), 16) == 0);
    375 
    376     is.read(buf2.get(), ZERO_SIZE);
    377     REQUIRE(is.good());
    378     REQUIRE(std::count(buf2.get(), buf2.get()+ZERO_SIZE, 0) == ZERO_SIZE);
    379 
    380     is.read(buf2.get(), 16);
    381     REQUIRE(is.good());
    382     REQUIRE(memcmp(buf, buf2.get(), 16) == 0);
    383 
    384     is.get();
    385     REQUIRE(is.eof());
    386   }
    387 
    388   TEST_CASE("sink helpers")
    389   {
    390     TRY_MMAP;
    391     {
    392       auto sink = Sink::ToFile("neptools_test.tmp", 15, try_mmap);
    393       sink->WriteLittleUint8(247);
    394       sink->WriteLittleUint16(1234);
    395       sink->WriteLittleUint32(98765);
    396       sink->WriteCString("asd");
    397       sink->WriteCString(std::string{"def"});
    398     }
    399 
    400     Byte exp[15] = { 247, 0xd2, 0x04, 0xcd, 0x81, 0x01, 0x00,
    401                      'a', 's', 'd', 0, 'd', 'e', 'f', 0 };
    402     char act[15];
    403 
    404     std::ifstream is{"neptools_test.tmp", std::ios_base::binary};
    405     is.read(act, 15);
    406     REQUIRE(is.good());
    407     REQUIRE(memcmp(exp, act, 15) == 0);
    408 
    409     is.get();
    410     REQUIRE(is.eof());
    411   }
    412 
    413 
    414   void MemorySink::Write_(Libshit::StringView)
    415   { LIBSHIT_UNREACHABLE("MemorySink::Write_ called"); }
    416   void MemorySink::Pad_(FileMemSize)
    417   { LIBSHIT_UNREACHABLE("MemorySink::Pad_ called"); }
    418 
    419   TEST_CASE("memory one write")
    420   {
    421     Byte buf[16] = {15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30};
    422     Byte buf2[16];
    423     {
    424       MemorySink sink{buf2, 16};
    425       sink.WriteGen(buf);
    426     }
    427 
    428     REQUIRE(memcmp(buf, buf2, 16) == 0);
    429   }
    430 
    431   TEST_CASE("memory multiple writes")
    432   {
    433     Byte buf[8] = {42,43,44,45,46,47,48,49};
    434     Byte buf_out[32];
    435     Byte buf_exp[32];
    436     {
    437       MemorySink sink{buf_out, 32};
    438       for (size_t i = 0; i < 32; i+=8)
    439       {
    440         memcpy(buf_exp+i, buf, 8);
    441         sink.WriteGen(buf);
    442         buf[0] = i;
    443       }
    444     }
    445 
    446     REQUIRE(memcmp(buf_out, buf_exp, 32) == 0);
    447   }
    448 
    449   TEST_CASE("memory pad")
    450   {
    451     Byte buf[8] = {77,78,79,80,81,82,83,84};
    452     Byte buf_out[32];
    453     Byte buf_exp[32];
    454     {
    455       MemorySink sink{buf_out, 32};
    456       memcpy(buf_exp, buf, 8);
    457       sink.WriteGen(buf);
    458 
    459       memset(buf_exp+8, 0, 16);
    460       sink.Pad(16);
    461 
    462       memcpy(buf_exp+24, buf, 8);
    463       sink.WriteGen(buf);
    464     }
    465 
    466     REQUIRE(memcmp(buf_out, buf_exp, 32) == 0);
    467   }
    468 
    469   TEST_CASE("memory alloc by itself")
    470   {
    471     char buf_exp[4] = { 0x78, 0x56, 0x34, 0x12 };
    472     MemorySink sink{4};
    473 
    474     sink.WriteLittleUint32(0x12345678);
    475     auto buf = sink.Release();
    476 
    477     REQUIRE(memcmp(buf.get(), buf_exp, 4) == 0);
    478   }
    479 
    480 
    481 #if LIBSHIT_WITH_LUA
    482   LIBSHIT_LUAGEN(name: "new", klass: "Neptools::MemorySink")
    483   static Libshit::NotNull<Libshit::SmartPtr<MemorySink>>
    484   MemorySinkFromLua(Libshit::StringView view)
    485   {
    486     std::unique_ptr<Byte[]> buf{new Byte[view.length()]};
    487     memcpy(buf.get(), view.data(), view.length());
    488     return Libshit::MakeSmart<MemorySink>(Libshit::Move(buf), view.length());
    489   }
    490 #endif
    491 
    492   TEST_SUITE_END();
    493 }
    494 
    495 #include "sink.binding.hpp"