diff --git a/extension/data_loader/shared_ptr_data_loader.h b/extension/data_loader/shared_ptr_data_loader.h index 78fdc2b7157..ff8d0189ae5 100644 --- a/extension/data_loader/shared_ptr_data_loader.h +++ b/extension/data_loader/shared_ptr_data_loader.h @@ -13,6 +13,7 @@ #include #include #include +#include #include namespace executorch { @@ -51,6 +52,24 @@ class SharedPtrDataLoader final : public executorch::runtime::DataLoader { return size_; } + ET_NODISCARD executorch::runtime::Error load_into( + size_t offset, + size_t size, + ET_UNUSED const SegmentInfo& segment_info, + void* buffer) const override { + ET_CHECK_OR_RETURN_ERROR( + buffer != nullptr, + InvalidArgument, + "Destination buffer cannot be null"); + + auto result = load(offset, size, segment_info); + if (!result.ok()) { + return result.error(); + } + std::memcpy(buffer, result->data(), size); + return executorch::runtime::Error::Ok; + } + private: const std::shared_ptr data_; const size_t size_; diff --git a/extension/data_loader/test/shared_ptr_data_loader_test.cpp b/extension/data_loader/test/shared_ptr_data_loader_test.cpp index 62d71ae0560..9fac89914b5 100644 --- a/extension/data_loader/test/shared_ptr_data_loader_test.cpp +++ b/extension/data_loader/test/shared_ptr_data_loader_test.cpp @@ -140,3 +140,60 @@ TEST_F(SharedPtrDataLoaderTest, OutOfBoundsLoadFails) { EXPECT_NE(fb.error(), Error::Ok); } } + +TEST_F(SharedPtrDataLoaderTest, LoadIntoNullDstFails) { + std::shared_ptr data( + new uint8_t[256](), std::default_delete()); + SharedPtrDataLoader edl(data, 256); + + EXPECT_EQ( + edl.load_into( + /*offset=*/0, + /*size=*/1, + /*segment_info=*/ + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program), + nullptr), + Error::InvalidArgument); + + EXPECT_EQ( + edl.load_into( + /*offset=*/0, + /*size=*/0, + /*segment_info=*/ + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program), + nullptr), + Error::InvalidArgument); +} + +TEST_F(SharedPtrDataLoaderTest, LoadIntoCopiesRequestedData) { + constexpr size_t kDataSize = 256; + std::shared_ptr data( + new uint8_t[kDataSize], std::default_delete()); + for (size_t i = 0; i < kDataSize; ++i) { + data[i] = static_cast(i); + } + SharedPtrDataLoader edl(data, kDataSize); + uint8_t buffer[3] = {}; + + EXPECT_EQ( + edl.load_into( + /*offset=*/kDataSize - sizeof(buffer), + /*size=*/sizeof(buffer), + /*segment_info=*/ + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program), + buffer), + Error::Ok); + EXPECT_EQ( + 0, + std::memcmp( + buffer, data.get() + kDataSize - sizeof(buffer), sizeof(buffer))); + + EXPECT_EQ( + edl.load_into( + /*offset=*/0, + /*size=*/kDataSize + 1, + /*segment_info=*/ + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program), + buffer), + Error::InvalidArgument); +}