Spaces:
Running
Running
File size: 3,887 Bytes
5f923cd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | // Copyright 2026 The ODML Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "runtime/util/data_stream.h"
#include <cstdint>
#include <memory>
#include <utility>
#include <vector>
#include "absl/status/status.h" // from @com_google_absl
#include "absl/status/statusor.h" // from @com_google_absl
#include "absl/strings/str_cat.h" // from @com_google_absl
#include "runtime/util/status_macros.h"
namespace litert::lm {
SubStream::~SubStream() {
if (auto parent = parent_.lock()) {
// Note: We ignore errors here as we are in a destructor.
(void)parent->Discard(offset_, size_);
}
}
absl::Status SubStream::ReadAndDiscard(void* buffer, uint64_t offset,
uint64_t size) {
RETURN_IF_ERROR(CheckBounds(offset, size));
if (auto parent = parent_.lock()) {
return parent->ReadAndDiscard(buffer, offset_ + offset, size);
}
return absl::FailedPreconditionError("Parent stream is expired");
}
absl::Status SubStream::ReadAndPreserve(void* buffer, uint64_t offset,
uint64_t size) {
RETURN_IF_ERROR(CheckBounds(offset, size));
if (auto parent = parent_.lock()) {
return parent->ReadAndPreserve(buffer, offset_ + offset, size);
}
return absl::FailedPreconditionError("Parent stream is expired");
}
absl::Status SubStream::Discard(uint64_t offset, uint64_t size) {
RETURN_IF_ERROR(CheckBounds(offset, size));
if (auto parent = parent_.lock()) {
return parent->Discard(offset_ + offset, size);
}
return absl::FailedPreconditionError("Parent stream is expired");
}
absl::Status SubStream::CheckBounds(uint64_t offset, uint64_t size) const {
// Equivalent to `offset + size > size_`
if (size > size_ || offset > size_ - size) {
return absl::OutOfRangeError(
absl::StrCat("Exceeded bounds of substream. offset: ", offset,
", size: ", size, ", Max size: ", size_));
}
return absl::OkStatus();
}
absl::StatusOr<std::shared_ptr<DataStream>> SubStream::OpenSubStream(
uint64_t offset, uint64_t size) {
// Check if the requested substream fits within this SubStream's bounds.
// Note that the parent DataStream::OpenSubStream method doesn't do this for
// us since DataStream doesn't know its own size.
RETURN_IF_ERROR(CheckBounds(offset, size));
// Call the base class implementation to check this SubStream's
// locked_regions_ and create a new SubStream child of this one.
return DataStream::OpenSubStream(offset, size);
}
absl::StatusOr<std::shared_ptr<DataStream>> DataStream::OpenSubStream(
uint64_t offset, uint64_t size) {
for (const auto& region : locked_regions_) {
// Check for overlap: Is [offset, offset + size) overlapping with
// [region.first, region.first + region.second)? Overlap exists if
// offset < region_end AND region_start < offset + size
if (offset < region.first + region.second && region.first < offset + size) {
return absl::AlreadyExistsError(absl::StrCat(
"Failed to open substream: requested region [", offset, ", ",
offset + size, ") overlaps with an existing locked region [",
region.first, ", ", region.first + region.second, ")"));
}
}
locked_regions_.emplace_back(offset, size);
return std::make_shared<SubStream>(shared_from_this(), offset, size);
}
} // namespace litert::lm
|