pw_kvs: Split out aligned write functionality

- Define the Output interface.
- Create AlignedWriter class, which handles writing aligned and padded
  data to a buffer.
- Define AlignedWriterBuffer, which provides a buffer to an
  AlignedWriter.

Change-Id: I791c5fc2d47617b0866fa779a506cb977905f245
diff --git a/pw_kvs/BUILD b/pw_kvs/BUILD
index b5db383..a059c7b 100644
--- a/pw_kvs/BUILD
+++ b/pw_kvs/BUILD
@@ -25,6 +25,7 @@
 pw_cc_library(
     name = "pw_kvs",
     srcs = [
+        "alignment.cc",
         "checksum.cc",
         "flash_memory.cc",
         "format.cc",
@@ -33,10 +34,12 @@
         "pw_kvs_private/macros.h",
     ],
     hdrs = [
+        "public/pw_kvs/alignment.h",
         "public/pw_kvs/checksum.h",
         "public/pw_kvs/crc16_checksum.h",
         "public/pw_kvs/flash_memory.h",
         "public/pw_kvs/key_value_store.h",
+        "public/pw_kvs/output.h",
     ],
     includes = ["public"],
     deps = [
@@ -71,6 +74,16 @@
 )
 
 pw_cc_test(
+    name = "alignment_test",
+    srcs = [
+        "alignment_test.cc",
+    ],
+    deps = [
+        ":pw_kvs",
+    ],
+)
+
+pw_cc_test(
     name = "checksum_test",
     srcs = ["checksum_test.cc"],
     deps = [
diff --git a/pw_kvs/BUILD.gn b/pw_kvs/BUILD.gn
index 391e3fc..ee09468 100644
--- a/pw_kvs/BUILD.gn
+++ b/pw_kvs/BUILD.gn
@@ -22,11 +22,14 @@
 source_set("pw_kvs") {
   public_configs = [ ":default_config" ]
   public = [
+    "public/pw_kvs/alignment.h",
     "public/pw_kvs/checksum.h",
     "public/pw_kvs/flash_memory.h",
     "public/pw_kvs/key_value_store.h",
+    "public/pw_kvs/output.h",
   ]
   sources = [
+    "alignment.cc",
     "checksum.cc",
     "flash_memory.cc",
     "format.cc",
@@ -85,6 +88,7 @@
 
 pw_test_group("tests") {
   tests = [
+    ":alignment_test",
     ":checksum_test",
     ":entry_test",
     ":key_value_store_test",
@@ -92,6 +96,15 @@
   ]
 }
 
+pw_test("alignment_test") {
+  deps = [
+    ":pw_kvs",
+  ]
+  sources = [
+    "alignment_test.cc",
+  ]
+}
+
 pw_test("checksum_test") {
   deps = [
     ":crc16",
diff --git a/pw_kvs/alignment.cc b/pw_kvs/alignment.cc
new file mode 100644
index 0000000..c470a3c
--- /dev/null
+++ b/pw_kvs/alignment.cc
@@ -0,0 +1,65 @@
+// Copyright 2020 The Pigweed 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
+//
+//     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 "pw_kvs/alignment.h"
+
+namespace pw {
+
+Status AlignedWriter::Write(span<const std::byte> data) {
+  while (!data.empty()) {
+    size_t to_copy = std::min(write_size_ - bytes_in_buffer_, data.size());
+
+    std::memcpy(&buffer_[bytes_in_buffer_], data.data(), to_copy);
+    data = data.subspan(to_copy);
+    bytes_in_buffer_ += to_copy;
+
+    // If the buffer is full, write it out.
+    if (bytes_in_buffer_ == write_size_) {
+      if (Status status = output_.Write(buffer_, write_size_); !status.ok()) {
+        return status;
+      }
+
+      bytes_written_ += write_size_;
+      bytes_in_buffer_ = 0;
+    }
+  }
+
+  return Status::OK;
+}
+
+StatusWithSize AlignedWriter::Flush() {
+  static constexpr std::byte kPadByte = std::byte{0};
+
+  // If data remains in the buffer, pad it to the alignment size and flush the
+  // remaining data.
+  if (bytes_in_buffer_ != 0u) {
+    const size_t remaining_bytes = AlignUp(bytes_in_buffer_, alignment_bytes_);
+    std::memset(&buffer_[bytes_in_buffer_],
+                int(kPadByte),
+                remaining_bytes - bytes_in_buffer_);
+
+    if (Status status = output_.Write(buffer_, remaining_bytes); !status.ok()) {
+      return StatusWithSize(status, bytes_written_);
+    }
+
+    bytes_written_ += remaining_bytes;  // Include padding in the total.
+  }
+
+  const StatusWithSize result(bytes_written_);
+  bytes_written_ = 0;
+  bytes_in_buffer_ = 0;
+  return result;
+}
+
+}  // namespace pw
diff --git a/pw_kvs/alignment_test.cc b/pw_kvs/alignment_test.cc
new file mode 100644
index 0000000..2b0771b
--- /dev/null
+++ b/pw_kvs/alignment_test.cc
@@ -0,0 +1,144 @@
+// Copyright 2020 The Pigweed 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
+//
+//     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 "pw_kvs/alignment.h"
+
+#include <string_view>
+
+#include "gtest/gtest.h"
+
+namespace pw::kvs {
+namespace {
+
+using std::byte;
+
+TEST(AlignUp, Zero) {
+  EXPECT_EQ(0u, AlignUp(0, 1));
+  EXPECT_EQ(0u, AlignUp(0, 2));
+  EXPECT_EQ(0u, AlignUp(0, 15));
+}
+
+TEST(AlignUp, Aligned) {
+  for (size_t i = 1; i < 130; ++i) {
+    EXPECT_EQ(i, AlignUp(i, i));
+    EXPECT_EQ(2 * i, AlignUp(2 * i, i));
+    EXPECT_EQ(3 * i, AlignUp(3 * i, i));
+  }
+}
+
+TEST(AlignUp, NonAligned_PowerOf2) {
+  EXPECT_EQ(32u, AlignUp(1, 32));
+  EXPECT_EQ(32u, AlignUp(31, 32));
+  EXPECT_EQ(64u, AlignUp(33, 32));
+  EXPECT_EQ(64u, AlignUp(45, 32));
+  EXPECT_EQ(64u, AlignUp(63, 32));
+  EXPECT_EQ(128u, AlignUp(127, 32));
+}
+
+TEST(AlignUp, NonAligned_NonPowerOf2) {
+  EXPECT_EQ(2u, AlignUp(1, 2));
+
+  EXPECT_EQ(15u, AlignUp(1, 15));
+  EXPECT_EQ(15u, AlignUp(14, 15));
+  EXPECT_EQ(30u, AlignUp(16, 15));
+}
+
+TEST(AlignDown, Zero) {
+  EXPECT_EQ(0u, AlignDown(0, 1));
+  EXPECT_EQ(0u, AlignDown(0, 2));
+  EXPECT_EQ(0u, AlignDown(0, 15));
+}
+
+TEST(AlignDown, Aligned) {
+  for (size_t i = 1; i < 130; ++i) {
+    EXPECT_EQ(i, AlignDown(i, i));
+    EXPECT_EQ(2 * i, AlignDown(2 * i, i));
+    EXPECT_EQ(3 * i, AlignDown(3 * i, i));
+  }
+}
+
+TEST(AlignDown, NonAligned_PowerOf2) {
+  EXPECT_EQ(0u, AlignDown(1, 32));
+  EXPECT_EQ(0u, AlignDown(31, 32));
+  EXPECT_EQ(32u, AlignDown(33, 32));
+  EXPECT_EQ(32u, AlignDown(45, 32));
+  EXPECT_EQ(32u, AlignDown(63, 32));
+  EXPECT_EQ(96u, AlignDown(127, 32));
+}
+
+TEST(AlignDown, NonAligned_NonPowerOf2) {
+  EXPECT_EQ(0u, AlignDown(1, 2));
+
+  EXPECT_EQ(0u, AlignDown(1, 15));
+  EXPECT_EQ(0u, AlignDown(14, 15));
+  EXPECT_EQ(15u, AlignDown(16, 15));
+}
+
+constexpr std::string_view kData =
+    "123456789_123456789_123456789_123456789_123456789_"   //  50
+    "123456789_123456789_123456789_123456789_123456789_";  // 100
+
+const span<const byte> kBytes = as_bytes(span(kData));
+
+TEST(AlignedWriter, VaryingLengthWriteCalls) {
+  static constexpr size_t kAlignment = 10;
+
+  OutputToFunction output([](span<const byte> data) {
+    EXPECT_EQ(data.size() % kAlignment, 0u);
+    EXPECT_EQ(kData.substr(0, data.size()),
+              std::string_view(reinterpret_cast<const char*>(data.data()),
+                               data.size()));
+    return StatusWithSize(data.size());
+  });
+
+  AlignedWriterBuffer<64> writer(kAlignment, output);
+
+  EXPECT_EQ(Status::OK, writer.Write(kBytes.subspan(0, 1)));
+  EXPECT_EQ(Status::OK, writer.Write(kBytes.subspan(1, 9)));
+  EXPECT_EQ(Status::OK, writer.Write(kBytes.subspan(10, 11)));
+  EXPECT_EQ(Status::OK, writer.Write(kBytes.subspan(21, 20)));
+  EXPECT_EQ(Status::OK, writer.Write(kBytes.subspan(41, 9)));
+  EXPECT_EQ(Status::OK, writer.Write(kBytes.subspan(50, 10)));
+  EXPECT_EQ(Status::OK, writer.Write(kBytes.subspan(60, 30)));
+  EXPECT_EQ(Status::OK, writer.Write(kBytes.subspan(90, 5)));
+  EXPECT_EQ(Status::OK, writer.Write(kBytes.subspan(95, 0)));
+  EXPECT_EQ(Status::OK, writer.Write(kBytes.subspan(95, 4)));
+  EXPECT_EQ(Status::OK, writer.Write(kBytes.subspan(99, 1)));
+
+  auto result = writer.Flush();
+  EXPECT_EQ(Status::OK, result.status());
+  EXPECT_EQ(kData.size(), result.size());
+}
+
+TEST(AlignedWriter, DestructorFlushes) {
+  static size_t called_with_bytes;
+
+  called_with_bytes = 0;
+
+  OutputToFunction output([](span<const byte> data) {
+    called_with_bytes += data.size();
+    return StatusWithSize(data.size());
+  });
+
+  {
+    AlignedWriterBuffer<64> writer(3, output);
+    writer.Write(as_bytes(span("What is this?")));
+    EXPECT_EQ(called_with_bytes, 0u);  // Buffer not full; no output yet.
+  }
+
+  EXPECT_EQ(called_with_bytes, AlignUp(sizeof("What is this?"), 3));
+}
+
+}  // namespace
+}  // namespace pw::kvs
diff --git a/pw_kvs/flash_memory.cc b/pw_kvs/flash_memory.cc
index ca5f5ff..a494c16 100644
--- a/pw_kvs/flash_memory.cc
+++ b/pw_kvs/flash_memory.cc
@@ -20,11 +20,18 @@
 
 #include "pw_kvs_private/macros.h"
 #include "pw_log/log.h"
+#include "pw_status/status_with_size.h"
 
 namespace pw::kvs {
 
 using std::byte;
 
+StatusWithSize FlashPartition::Output::Write(span<const byte> data) {
+  TRY(flash_.Write(address_, data));
+  address_ += data.size();
+  return StatusWithSize(data.size());
+}
+
 Status FlashPartition::Erase(Address address, size_t num_sectors) {
   if (permission_ == PartitionPermission::kReadOnly) {
     return Status::PERMISSION_DENIED;
@@ -47,54 +54,6 @@
   return flash_.Write(PartitionToFlashAddress(address), data);
 }
 
-StatusWithSize FlashPartition::WriteAligned(
-    const Address start_address, std::initializer_list<span<const byte>> data) {
-  byte buffer[64];  // TODO: Configure this?
-
-  Address address = start_address;
-  auto bytes_written = [&]() { return address - start_address; };
-
-  const size_t write_size = AlignDown(sizeof(buffer), alignment_bytes());
-  size_t bytes_in_buffer = 0;
-
-  for (span<const byte> chunk : data) {
-    while (!chunk.empty()) {
-      const size_t to_copy =
-          std::min(write_size - bytes_in_buffer, chunk.size());
-
-      std::memcpy(&buffer[bytes_in_buffer], chunk.data(), to_copy);
-      chunk = chunk.subspan(to_copy);
-      bytes_in_buffer += to_copy;
-
-      // If the buffer is full, write it out.
-      if (bytes_in_buffer == write_size) {
-        Status status = Write(address, span(buffer, write_size));
-        if (!status.ok()) {
-          return StatusWithSize(status, bytes_written());
-        }
-
-        address += write_size;
-        bytes_in_buffer = 0;
-      }
-    }
-  }
-
-  // If data remains in the buffer, pad it to the alignment size and flush
-  // the remaining data.
-  if (bytes_in_buffer != 0u) {
-    size_t remaining_write_size = AlignUp(bytes_in_buffer, alignment_bytes());
-    std::memset(
-        &buffer[bytes_in_buffer], 0, remaining_write_size - bytes_in_buffer);
-    if (Status status = Write(address, span(buffer, remaining_write_size));
-        !status.ok()) {
-      return StatusWithSize(status, bytes_written());
-    }
-    address += remaining_write_size;  // Include padding bytes in the total.
-  }
-
-  return StatusWithSize(bytes_written());
-}
-
 Status FlashPartition::IsRegionErased(Address source_flash_address,
                                       size_t length,
                                       bool* is_erased) {
diff --git a/pw_kvs/key_value_store.cc b/pw_kvs/key_value_store.cc
index b6b8e30..a30bc39 100644
--- a/pw_kvs/key_value_store.cc
+++ b/pw_kvs/key_value_store.cc
@@ -750,12 +750,13 @@
   Address address = NextWritableAddress(sector);
   DBG("Appending to address: %#zx", size_t(address));
 
-  // Handles writing multiple concatenated buffers, while breaking up the writes
-  // into alignment-sized blocks.
-  TRY_ASSIGN(
-      const size_t written,
-      partition_.WriteAligned(
-          address, {as_bytes(span(&header, 1)), as_bytes(span(key)), value}));
+  // Write multiple concatenated buffers and pad the results.
+  FlashPartition::Output flash(partition_, address);
+  TRY_ASSIGN(const size_t written,
+             AlignedWrite<32>(
+                 flash,
+                 header.alignment_bytes(),
+                 {as_bytes(span(&header, 1)), as_bytes(span(key)), value}));
 
   if (options_.verify_on_write) {
     TRY(header.VerifyChecksumInFlash(
diff --git a/pw_kvs/public/pw_kvs/alignment.h b/pw_kvs/public/pw_kvs/alignment.h
new file mode 100644
index 0000000..0829854
--- /dev/null
+++ b/pw_kvs/public/pw_kvs/alignment.h
@@ -0,0 +1,120 @@
+// Copyright 2020 The Pigweed 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
+//
+//     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.
+#pragma once
+
+#include <cstddef>
+#include <cstring>
+#include <initializer_list>
+#include <utility>
+
+#include "pw_kvs/output.h"
+#include "pw_span/span.h"
+#include "pw_status/status_with_size.h"
+
+namespace pw {
+
+// Returns the value rounded down to the nearest multiple of alignment.
+constexpr size_t AlignDown(size_t value, size_t alignment) {
+  return (value / alignment) * alignment;
+}
+
+// Returns the value rounded up to the nearest multiple of alignment.
+constexpr size_t AlignUp(size_t value, size_t alignment) {
+  return (value + alignment - 1) / alignment * alignment;
+}
+
+// Returns the number of padding bytes required to align the provided length.
+constexpr size_t Padding(size_t length, size_t alignment) {
+  return AlignUp(length, alignment) - length;
+}
+
+// Class for managing aligned writes. Stores data in an intermediate buffer and
+// calls an output function with aligned data as the buffer becomes full. Any
+// bytes remaining in the buffer are written to the output when Flush() is
+// called or the AlignedWriter goes out of scope.
+class AlignedWriter {
+ public:
+  AlignedWriter(span<std::byte> buffer, size_t alignment_bytes, Output& writer)
+      : buffer_(buffer.data()),
+        write_size_(AlignDown(buffer.size(), alignment_bytes)),
+        alignment_bytes_(alignment_bytes),
+        output_(writer),
+        bytes_written_(0),
+        bytes_in_buffer_(0) {
+    // TODO(hepler): Add DCHECK to ensure that buffer.size() >= alignment_bytes.
+  }
+
+  ~AlignedWriter() { Flush(); }
+
+  // Writes bytes to the AlignedWriter. The output may be called if the internal
+  // buffer is filled.
+  Status Write(span<const std::byte> data);
+
+  Status Write(const void* data, size_t size) {
+    return Write(span(static_cast<const std::byte*>(data), size));
+  }
+
+  // Flush and reset the AlignedWriter. Any remaining bytes in the buffer are
+  // zero-padded to an alignment boundary and written to the output. Flush is
+  // automatically called when the AlignedWriter goes out of scope.
+  StatusWithSize Flush();
+
+ private:
+  std::byte* const buffer_;
+  const size_t write_size_;
+  const size_t alignment_bytes_;
+
+  Output& output_;
+  size_t bytes_written_;
+  size_t bytes_in_buffer_;
+};
+
+// Declares an AlignedWriter with a built-in buffer.
+template <size_t kBufferSize>
+class AlignedWriterBuffer : public AlignedWriter {
+ public:
+  template <typename... Args>
+  AlignedWriterBuffer(Args&&... aligned_writer_args)
+      : AlignedWriter(buffer_, std::forward<Args>(aligned_writer_args)...) {}
+
+ private:
+  std::byte buffer_[kBufferSize];
+};
+
+// Writes data from multiple buffers using an AlignedWriter.
+template <size_t kBufferSize>
+StatusWithSize AlignedWrite(Output& output,
+                            size_t alignment_bytes,
+                            span<const span<const std::byte>> data) {
+  AlignedWriterBuffer<kBufferSize> buffer(alignment_bytes, output);
+
+  for (const span<const std::byte>& chunk : data) {
+    if (Status status = buffer.Write(chunk); !status.ok()) {
+      return status;
+    }
+  }
+
+  return buffer.Flush();
+}
+
+// Calls AlignedWrite with an initializer list.
+template <size_t kBufferSize>
+StatusWithSize AlignedWrite(Output& output,
+                            size_t alignment_bytes,
+                            std::initializer_list<span<const std::byte>> data) {
+  return AlignedWrite<kBufferSize>(
+      output, alignment_bytes, span(data.begin(), data.size()));
+}
+
+}  // namespace pw
diff --git a/pw_kvs/public/pw_kvs/flash_memory.h b/pw_kvs/public/pw_kvs/flash_memory.h
index 3bc2a2b..a10af48 100644
--- a/pw_kvs/public/pw_kvs/flash_memory.h
+++ b/pw_kvs/public/pw_kvs/flash_memory.h
@@ -17,24 +17,11 @@
 #include <cstdint>
 #include <initializer_list>
 
+#include "pw_kvs/alignment.h"
 #include "pw_span/span.h"
 #include "pw_status/status.h"
 #include "pw_status/status_with_size.h"
 
-namespace pw {
-
-// TODO: These are general-purpose utility functions that should be moved
-//       elsewhere.
-constexpr size_t AlignDown(size_t value, size_t alignment) {
-  return (value / alignment) * alignment;
-}
-
-constexpr size_t AlignUp(size_t value, size_t alignment) {
-  return (value + alignment - 1) / alignment * alignment;
-}
-
-}  // namespace pw
-
 namespace pw::kvs {
 
 enum class PartitionPermission : bool {
@@ -137,6 +124,19 @@
   // The flash address is in the range of: 0 to PartitionSize.
   using Address = uint32_t;
 
+  // Implement Output for the Write method.
+  class Output final : public pw::Output {
+   public:
+    constexpr Output(FlashPartition& flash, FlashPartition::Address address)
+        : flash_(flash), address_(address) {}
+
+    StatusWithSize Write(span<const std::byte> data) override;
+
+   private:
+    FlashPartition& flash_;
+    FlashPartition::Address address_;
+  };
+
   constexpr FlashPartition(
       FlashMemory* flash,
       uint32_t start_sector_index,
@@ -185,10 +185,6 @@
   //          UNKNOWN, on HAL error
   virtual StatusWithSize Write(Address address, span<const std::byte> data);
 
-  // Returns the total number of bytes written, including any padding.
-  StatusWithSize WriteAligned(
-      Address start_address, std::initializer_list<span<const std::byte>> data);
-
   // Check to see if chunk of flash memory is erased. Address and len need to
   // be aligned with FlashMemory.
   // Returns: OK, on success.
diff --git a/pw_kvs/public/pw_kvs/output.h b/pw_kvs/public/pw_kvs/output.h
new file mode 100644
index 0000000..fae001f
--- /dev/null
+++ b/pw_kvs/public/pw_kvs/output.h
@@ -0,0 +1,89 @@
+// Copyright 2020 The Pigweed 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
+//
+//     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.
+#pragma once
+
+#include <cstddef>
+#include <type_traits>
+
+#include "pw_span/span.h"
+#include "pw_status/status_with_size.h"
+
+namespace pw {
+namespace internal {
+
+template <typename T>
+struct FunctionTraits;
+
+template <typename T, typename ReturnType, typename... Args>
+struct FunctionTraits<ReturnType (T::*)(Args...)> {
+  using Class = T;
+  using Return = ReturnType;
+};
+
+}  // namespace internal
+
+// Writes bytes to an unspecified output. Provides a Write function that takes a
+// span of bytes and returns a Status.
+class Output {
+ public:
+  virtual StatusWithSize Write(span<const std::byte> data) = 0;
+
+  // Convenience wrapper for writing data from a pointer and length.
+  StatusWithSize Write(const void* data, size_t size_bytes) {
+    return Write(span(static_cast<const std::byte*>(data), size_bytes));
+  }
+
+ protected:
+  ~Output() = default;
+};
+
+// Output adapter that calls a method on a class with a span of bytes. If the
+// method returns void instead of the expected Status, Write always returns
+// Status::OK.
+template <auto kMethod>
+class OutputToMethod final : public Output {
+  using Class = typename internal::FunctionTraits<decltype(kMethod)>::Class;
+
+ public:
+  constexpr OutputToMethod(Class* object) : object_(*object) {}
+
+  StatusWithSize Write(span<const std::byte> data) override {
+    using Return = typename internal::FunctionTraits<decltype(kMethod)>::Return;
+
+    if constexpr (std::is_void_v<Return>) {
+      (object_.*kMethod)(data);
+      return StatusWithSize(data.size());
+    } else {
+      return (object_.*kMethod)(data);
+    }
+  }
+
+ private:
+  Class& object_;
+};
+
+// Output adapter that calls a free function.
+class OutputToFunction final : public Output {
+ public:
+  OutputToFunction(StatusWithSize (*function)(span<const std::byte>))
+      : function_(function) {}
+
+  StatusWithSize Write(span<const std::byte> data) override {
+    return function_(data);
+  }
+
+  StatusWithSize (*function_)(span<const std::byte>);
+};
+
+}  // namespace pw
diff --git a/pw_kvs/pw_kvs_private/format.h b/pw_kvs/pw_kvs_private/format.h
index 2e9094e..db42b12 100644
--- a/pw_kvs/pw_kvs_private/format.h
+++ b/pw_kvs/pw_kvs_private/format.h
@@ -20,6 +20,7 @@
 #include <cstring>
 #include <string_view>
 
+#include "pw_kvs/alignment.h"
 #include "pw_kvs/checksum.h"
 #include "pw_kvs/flash_memory.h"
 #include "pw_span/span.h"