utils/byte_instruction.cc

// Copyright 2026 Google LLC
//
// 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
//
//    https://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 "utils/byte_instruction.h"

#include <algorithm>
#include <cstdint>
#include <optional>
#include <string>
#include <vector>

#include "absl/base/nullability.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "riegeli/base/chain.h"
#include "riegeli/bytes/chain_reader.h"
#include "riegeli/bytes/copy_all.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/writer.h"

namespace credentio {

namespace {

constexpr uint64_t kBufferChunkSize = 1024 * 1024;  // 1 MiB

absl::Status ReadDataToBuffer(riegeli::Writer& writer, int64_t length,
                              riegeli::Chain& buffer) {
  // Flush the writer to make sure the data is available for reading.
  if (!writer.Flush()) {
    return writer.status();
  }

  // Create a reader at the current position of the writer.
  riegeli::Reader* reader = writer.ReadMode(writer.pos());
  if (reader == nullptr) {
    return writer.status();
  }

  // Read data from the reader to the buffer.
  if (!reader->ReadAndAppend(length, buffer)) {
    return reader->StatusOrAnnotate(
        absl::DataLossError("Failed to read data from buffer"));
  }

  // Close the reader.
  if (!reader->Close()) {
    return reader->StatusOrAnnotate(
        absl::DataLossError("Failed to close reader"));
  }

  return absl::OkStatus();
}

absl::Status WriteDataFromBuffer(riegeli::Chain& buffer, int64_t length,
                                 riegeli::Writer& writer) {
  // Create a reader for the buffer.
  riegeli::ChainReader cr(&buffer);

  // Read data from the buffer.
  absl::string_view data;
  if (!cr.Read(length, data)) {
    return cr.StatusOrAnnotate(
        absl::DataLossError("Failed to read data from buffer"));
  }

  // Write the data to the writer.
  if (!writer.Write(data)) {
    return writer.status();
  }

  // Close the reader.
  if (!cr.Close()) {
    return cr.StatusOrAnnotate(absl::DataLossError("Failed to close reader"));
  }

  // Remove data from the buffer, done after writing as the absl::string_view
  // data is a ref into the buffer.
  buffer.RemovePrefix(length);

  return absl::OkStatus();
}

absl::Status MoveDataThroughBuffer(riegeli::Chain& buffer, int64_t length,
                                   riegeli::Writer& writer) {
  ABSL_RETURN_IF_ERROR(ReadDataToBuffer(writer, length, buffer));
  return WriteDataFromBuffer(buffer, length, writer);
}

absl::Status MoveDataThroughBuffer(riegeli::Chain& buffer,
                                   absl::string_view data,
                                   riegeli::Writer& writer) {
  buffer.Append(data);
  return WriteDataFromBuffer(buffer, data.size(), writer);
}

absl::Status PushData(riegeli::Writer& writer, riegeli::Chain& buffer,
                      int64_t end_position) {
  if (buffer.empty()) {
    // We didn't insert anything yet, so just seek
    if (!writer.Seek(end_position)) {
      return writer.StatusOrAnnotate(
          absl::DataLossError("Failed to seek writer"));
    }
    return absl::OkStatus();
  }

  uint64_t remaining_byte_count = end_position - writer.pos();

  while (remaining_byte_count > 0) {
    uint64_t count_to_move = std::min(kBufferChunkSize, remaining_byte_count);

    ABSL_RETURN_IF_ERROR(MoveDataThroughBuffer(buffer, count_to_move, writer));

    remaining_byte_count -= count_to_move;
  }
  return absl::OkStatus();
}

}  // namespace

absl::Status ApplyByteInstructions(
    riegeli::Reader* absl_nonnull source,
    absl::Span<const ByteInstruction> instructions,
    riegeli::Writer* absl_nonnull destination) {
  if (!source->Seek(0) || !destination->Seek(0)) {
    return absl::DataLossError("failed to seek to start of file");
  }

  if (instructions.empty()) {
    // No instructions, just copy the source to the destination.
    return riegeli::CopyAll(*source, *destination);
  }

  if (!source->SupportsSize() || !source->Size().has_value()) {
    return absl::InvalidArgumentError("Source does not have a size.");
  }
  uint64_t file_size = source->Size().value_or(0);
  uint64_t instruction_index = 0;

  while (source->pos() < file_size && instruction_index < instructions.size()) {
    const ByteInstruction& ins = instructions[instruction_index];

    if (ins.offset < source->pos()) {
      return absl::InvalidArgumentError(
          "Byte instruction offsets must be in ascending order");
    }
    if (ins.offset >= file_size) {
      // Next instruction is at or past the end of the file.
      break;
    }

    // Copy enough data from the source to reach the instruction offset.
    if (source->pos() < ins.offset) {
      if (!source->Copy(ins.offset - source->pos(), *destination)) {
        ABSL_RETURN_IF_ERROR(source->status());
        ABSL_RETURN_IF_ERROR(destination->status());
        return absl::DataLossError("Unexpected EOF while copying source data");
      }
    }

    // Write the data
    if (!destination->Write(ins.bytes)) {
      return destination->status();
    }

    if (ins.operation == ByteInstruction::Operation::kReplace) {
      uint64_t bytes_to_skip =
          std::min<uint64_t>(ins.bytes.size(), file_size - source->pos());
      if (!source->Skip(bytes_to_skip)) {
        ABSL_RETURN_IF_ERROR(source->status());
        return absl::DataLossError("Unexpected EOF while skipping source data");
      }
    }

    ++instruction_index;
  }

  // Copy any remaining data from the source to the destination.
  if (source->pos() < file_size) {
    ABSL_RETURN_IF_ERROR(riegeli::CopyAll(*source, *destination));
  }

  // Write any remaining instructions.
  while (instruction_index < instructions.size()) {
    const ByteInstruction& ins = instructions[instruction_index];
    if (!destination->Write(ins.bytes)) {
      return destination->status();
    }
    ++instruction_index;
  }

  return absl::OkStatus();
}

absl::Status ApplyByteInstructions(
    absl::Span<const credentio::ByteInstruction> instructions,
    riegeli::Writer* absl_nonnull writer) {
  if (!writer->Seek(0)) {
    return absl::DataLossError("failed to seek to start of file");
  }

  if (instructions.empty()) {
    return absl::OkStatus();
  }

  if (!writer->Size().has_value()) {
    return absl::InvalidArgumentError("Writer does not have a size.");
  }

  uint64_t file_size = writer->Size().value();
  riegeli::Chain buffer;

  uint64_t instruction_index = 0;

  while (writer->pos() < file_size && instruction_index < instructions.size()) {
    const ByteInstruction& ins = instructions[instruction_index];

    if (ins.offset < writer->pos()) {
      return absl::InvalidArgumentError(
          "Byte instruction offsets must be in ascending order");
    }
    if (ins.offset >= file_size) {
      // Next instruction is at or past the end of the file.
      break;
    }

    // Process enough data to reach the instruction offset.
    ABSL_RETURN_IF_ERROR(PushData(*writer, buffer, ins.offset));

    if (ins.operation == credentio::ByteInstruction::Operation::kInsert) {
      // Insert operations do not impact the current offset, just add to the
      // buffer.
      buffer.Append(ins.bytes);
    } else if (ins.operation ==
               credentio::ByteInstruction::Operation::kReplace) {
      if (buffer.empty()) {
        // If the buffer is empty, we can write the data directly to the writer.
        if (!writer->Write(ins.bytes)) {
          return writer->status();
        }
      } else {
        // If the buffer is not empty, we need to move the data through the
        // buffer.
        ABSL_RETURN_IF_ERROR(MoveDataThroughBuffer(buffer, ins.bytes, *writer));
      }
    }

    ++instruction_index;
  }

  // Process all remaining data up to the end of the file.
  if (writer->pos() < file_size) {
    ABSL_RETURN_IF_ERROR(PushData(*writer, buffer, file_size));
  }
  if (!buffer.empty()) {
    // Write any remaining data in the buffer.
    if (!writer->Write(buffer)) {
      return writer->status();
    }
  }

  // Write any remaining instructions.
  while (instruction_index < instructions.size()) {
    if (!writer->Write(instructions[instruction_index].bytes)) {
      return writer->status();
    }
    ++instruction_index;
  }

  return absl::OkStatus();
}

uint64_t CalculateBytesWritten(
    const std::vector<ByteInstruction>& byte_instructions, uint64_t file_size,
    bool is_copy) {
  uint64_t sum_replace_before_insert = 0;
  uint64_t sum_insert = 0;
  std::optional<uint64_t> first_insert_offset =
      is_copy ? std::make_optional<uint64_t>(0) : std::nullopt;

  for (const auto& byte_instruction : byte_instructions) {
    if (byte_instruction.offset >= file_size) {
      // Instructions at or past the end of the original file do not replace
      // any existing bytes; they are simply appended to the destination.
      sum_insert += byte_instruction.bytes.size();
    } else if (byte_instruction.operation ==
               ByteInstruction::Operation::kReplace) {
      if (!first_insert_offset.has_value()) {
        sum_replace_before_insert += byte_instruction.bytes.size();
      }
    } else if (byte_instruction.operation ==
               ByteInstruction::Operation::kInsert) {
      sum_insert += byte_instruction.bytes.size();
      if (!first_insert_offset.has_value()) {
        first_insert_offset = byte_instruction.offset;
      }
    }
  }

  uint64_t bytes_shifted = 0;
  if (!first_insert_offset.has_value() || *first_insert_offset >= file_size) {
    bytes_shifted = 0;
  } else {
    bytes_shifted = file_size - *first_insert_offset;
  }

  return bytes_shifted + sum_insert + sum_replace_before_insert;
}

}  // namespace credentio