fix custom io reader/writer API

examples/example-safe-tensors.c

@@ -32,7 +32,7 @@ size_t mem_tell(void* desc) {
   mlx_mem_stream* m_desc = desc;
   return m_desc->pos;
 }
-void mem_seek(void* desc, int64_t off, int whence) {
+int mem_seek(void* desc, int64_t off, int whence) {
   printf("SEEK\n");
   mlx_mem_stream* m_desc = desc;
   size_t new_pos;
@@ -48,43 +48,48 @@ void mem_seek(void* desc, int64_t off, int whence) {
       break;
     default:
       m_desc->err = true;
-      return;
+      return -1;
   }
   if (new_pos > m_desc->size) {
     m_desc->err = true;
+    return -1;
   } else {
     m_desc->pos = new_pos;
+    return 0;
   }
 }
-void mem_read(void* desc, char* data, size_t n) {
+size_t mem_read(void* desc, char* data, size_t n) {
   printf("READ %ld\n", n);
   mlx_mem_stream* m_desc = desc;
   if (n + m_desc->pos > m_desc->size) {
     m_desc->err = true;
-    return;
+    return 0;
   }
   memcpy(data, m_desc->data + m_desc->pos, n);
   m_desc->pos += n;
+  return n;
 }
-void mem_read_at_offset(void* desc, char* data, size_t n, size_t off) {
+size_t mem_read_at_offset(void* desc, char* data, size_t n, size_t off) {
   printf("READ@OFFSET %ld @ %ld\n", n, off);
   mlx_mem_stream* m_desc = desc;
   if (off + n > m_desc->size) {
     m_desc->err = true;
-    return;
+    return 0;
   }
   memcpy(data, m_desc->data + off, n);
   m_desc->pos = off;
+  return n;
 }
-void mem_write(void* desc, const char* data, size_t n) {
+size_t mem_write(void* desc, const char* data, size_t n) {
   printf("WRITE %ld\n", n);
   mlx_mem_stream* m_desc = desc;
   if (n + m_desc->pos > m_desc->size) {
     m_desc->err = true;
-    return;
+    return 0;
   }
   memcpy(m_desc->data + m_desc->pos, data, n);
   m_desc->pos += n;
+  return n;
 }
 const char* mem_label(void* desc) {
   printf("LABEL\n");

mlx/c/io_types.h

@@ -37,10 +37,10 @@ typedef struct mlx_io_vtable_ {
   bool (*is_open)(void*);
   bool (*good)(void*);
   size_t (*tell)(void*);
-  void (*seek)(void*, int64_t off, int whence);
-  void (*read)(void*, char* data, size_t n);
-  void (*read_at_offset)(void*, char* data, size_t n, size_t off);
-  void (*write)(void*, const char* data, size_t n);
+  int (*seek)(void*, int64_t off, int whence);
+  size_t (*read)(void*, char* data, size_t n);
+  size_t (*read_at_offset)(void*, char* data, size_t n, size_t off);
+  size_t (*write)(void*, const char* data, size_t n);
   const char* (*label)(void*);
   void (*free)(void*);
 } mlx_io_vtable;

mlx/c/private/io.h

@@ -26,29 +26,45 @@ class CReader : public mlx::core::io::Reader {
   virtual void seek(
       int64_t off,
       std::ios_base::seekdir way = std::ios_base::beg) override {
+    int status = 0;
     switch (way) {
       case std::ios_base::beg:
-        return vtable.seek(desc, off, SEEK_SET);
+        status = vtable.seek(desc, off, SEEK_SET);
         break;
       case std::ios_base::cur:
-        return vtable.seek(desc, off, SEEK_CUR);
+        status = vtable.seek(desc, off, SEEK_CUR);
         break;
       case std::ios_base::end:
-        return vtable.seek(desc, off, SEEK_END);
+        status = vtable.seek(desc, off, SEEK_END);
         break;
       default:
-        throw std::runtime_error("mlx_io_reader: invalid seek way");
+        reader_error("invalid seek way");
+    }
+    if (status < 0) {
+      reader_error("unable to seek");
     }
   }
   virtual void read(char* data, size_t n) override {
-    return vtable.read(desc, data, n);
+    auto read_n = vtable.read(desc, data, n);
+    if (read_n != n) {
+      reader_error(
+          std::format("unable to read {} bytes (read {} instead)", n, read_n));
+    }
   };
   virtual void read(char* data, size_t n, size_t offset) override {
-    return vtable.read_at_offset(desc, data, n, offset);
+    auto read_n = vtable.read_at_offset(desc, data, n, offset);
+    if (read_n != n) {
+      reader_error(
+          std::format("unable to read {} bytes (read {} instead)", n, read_n));
+    }
   };
   virtual std::string label() const override {
     return vtable.label(desc);
   };
+  void reader_error(const std::string& msg) {
+    throw std::runtime_error(
+        std::format("[mlx_io_reader] {} in {}", msg, label()));
+  }
   virtual ~CReader() {
     vtable.free(desc);
   }
@@ -72,26 +88,39 @@ class CWriter : public mlx::core::io::Writer {
   virtual void seek(
       int64_t off,
       std::ios_base::seekdir way = std::ios_base::beg) override {
+    int status = 0;
     switch (way) {
       case std::ios_base::beg:
-        return vtable.seek(desc, off, SEEK_SET);
+        status = vtable.seek(desc, off, SEEK_SET);
         break;
       case std::ios_base::cur:
-        return vtable.seek(desc, off, SEEK_CUR);
+        status = vtable.seek(desc, off, SEEK_CUR);
         break;
       case std::ios_base::end:
-        return vtable.seek(desc, off, SEEK_END);
+        status = vtable.seek(desc, off, SEEK_END);
         break;
       default:
-        throw std::runtime_error("mlx_io_writer: invalid seek way");
+        writer_error("invalid seek way");
+    }
+    if (status < 0) {
+      writer_error("unable to seek");
     }
   }
   virtual void write(const char* data, size_t n) override {
-    return vtable.write(desc, data, n);
+    auto wrote_n = vtable.write(desc, data, n);
+    if (wrote_n != n) {
+      writer_error(
+          std::format(
+              "unable to write {} bytes (wrote {} instead)", n, wrote_n));
+    }
   };
   virtual std::string label() const override {
     return vtable.label(desc);
   };
+  void writer_error(const std::string& msg) {
+    throw std::runtime_error(
+        std::format("[mlx_io_writer] {} in {}", msg, label()));
+  }
   virtual ~CWriter() {
     vtable.free(desc);
   }