vibe_zstd 1.0.0

data/ext/vibe_zstd/vibe_zstd.c

@@ -0,0 +1,293 @@
+#include "vibe_zstd_internal.h"
+#include <ruby/thread.h>
+#define ZDICT_STATIC_LINKING_ONLY
+#include <zdict.h>
+
+// Ruby module and class handles
+VALUE rb_mVibeZstd;
+VALUE rb_cVibeZstdCCtx;
+VALUE rb_cVibeZstdDCtx;
+VALUE rb_cVibeZstdCDict;
+VALUE rb_cVibeZstdDDict;
+VALUE rb_cVibeZstdCompressWriter;
+VALUE rb_cVibeZstdDecompressReader;
+
+// Forward declarations for free and mark functions
+static void vibe_zstd_cctx_free(void* ptr);
+static void vibe_zstd_dctx_free(void* ptr);
+static void vibe_zstd_cdict_free(void* ptr);
+static void vibe_zstd_ddict_free(void* ptr);
+static void vibe_zstd_cstream_free(void* ptr);
+static void vibe_zstd_cstream_mark(void* ptr);
+static void vibe_zstd_dstream_free(void* ptr);
+static void vibe_zstd_dstream_mark(void* ptr);
+
+// TypedData type definitions (these are referenced by extern in the split files)
+rb_data_type_t vibe_zstd_cctx_type = {
+    .wrap_struct_name = "vibe_zstd_cctx",
+    .function = {
+        .dmark = NULL,
+        .dfree = (RUBY_DATA_FUNC)vibe_zstd_cctx_free,
+        .dsize = NULL,
+    },
+    .data = NULL,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY,
+};
+
+rb_data_type_t vibe_zstd_dctx_type = {
+    .wrap_struct_name = "vibe_zstd_dctx",
+    .function = {
+        .dmark = NULL,
+        .dfree = (RUBY_DATA_FUNC)vibe_zstd_dctx_free,
+        .dsize = NULL,
+    },
+    .data = NULL,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY,
+};
+
+rb_data_type_t vibe_zstd_cdict_type = {
+    .wrap_struct_name = "vibe_zstd_cdict",
+    .function = {
+        .dmark = NULL,
+        .dfree = (RUBY_DATA_FUNC)vibe_zstd_cdict_free,
+        .dsize = NULL,
+    },
+    .data = NULL,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY,
+};
+
+rb_data_type_t vibe_zstd_ddict_type = {
+    .wrap_struct_name = "vibe_zstd_ddict",
+    .function = {
+        .dmark = NULL,
+        .dfree = (RUBY_DATA_FUNC)vibe_zstd_ddict_free,
+        .dsize = NULL,
+    },
+    .data = NULL,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY,
+};
+
+rb_data_type_t vibe_zstd_cstream_type = {
+    .wrap_struct_name = "vibe_zstd_cstream",
+    .function = {
+        .dmark = (RUBY_DATA_FUNC)vibe_zstd_cstream_mark,
+        .dfree = (RUBY_DATA_FUNC)vibe_zstd_cstream_free,
+        .dsize = NULL,
+    },
+    .data = NULL,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY,
+};
+
+rb_data_type_t vibe_zstd_dstream_type = {
+    .wrap_struct_name = "vibe_zstd_dstream",
+    .function = {
+        .dmark = (RUBY_DATA_FUNC)vibe_zstd_dstream_mark,
+        .dfree = (RUBY_DATA_FUNC)vibe_zstd_dstream_free,
+        .dsize = NULL,
+    },
+    .data = NULL,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY,
+};
+
+// Free functions
+static void
+vibe_zstd_cctx_free(void* ptr) {
+    vibe_zstd_cctx* cctx = ptr;
+    if (cctx->cctx) {
+        ZSTD_freeCCtx(cctx->cctx);
+    }
+    ruby_xfree(cctx);
+}
+
+static void
+vibe_zstd_dctx_free(void* ptr) {
+    vibe_zstd_dctx* dctx = ptr;
+    if (dctx->dctx) {
+        ZSTD_freeDCtx(dctx->dctx);
+    }
+    ruby_xfree(dctx);
+}
+
+static void
+vibe_zstd_cdict_free(void* ptr) {
+    vibe_zstd_cdict* cdict = ptr;
+    if (cdict->cdict) {
+        ZSTD_freeCDict(cdict->cdict);
+    }
+    ruby_xfree(cdict);
+}
+
+static void
+vibe_zstd_ddict_free(void* ptr) {
+    vibe_zstd_ddict* ddict = ptr;
+    if (ddict->ddict) {
+        ZSTD_freeDDict(ddict->ddict);
+    }
+    ruby_xfree(ddict);
+}
+
+static void
+vibe_zstd_cstream_mark(void* ptr) {
+    vibe_zstd_cstream* cstream = ptr;
+    rb_gc_mark(cstream->io);
+}
+
+static void
+vibe_zstd_cstream_free(void* ptr) {
+    vibe_zstd_cstream* cstream = ptr;
+    if (cstream->cstream) {
+        ZSTD_freeCStream(cstream->cstream);
+    }
+    ruby_xfree(cstream);
+}
+
+static void
+vibe_zstd_dstream_mark(void* ptr) {
+    vibe_zstd_dstream* dstream = ptr;
+    rb_gc_mark(dstream->io);
+    rb_gc_mark(dstream->input_data);
+}
+
+static void
+vibe_zstd_dstream_free(void* ptr) {
+    vibe_zstd_dstream* dstream = ptr;
+    if (dstream->dstream) {
+        ZSTD_freeDStream(dstream->dstream);
+    }
+    ruby_xfree(dstream);
+}
+
+// Alloc functions
+static VALUE
+vibe_zstd_cctx_alloc(VALUE klass) {
+    vibe_zstd_cctx* cctx = ALLOC(vibe_zstd_cctx);
+    cctx->cctx = ZSTD_createCCtx();
+    if (!cctx->cctx) {
+        ruby_xfree(cctx);
+        rb_raise(rb_eRuntimeError, "Failed to create ZSTD_CCtx");
+    }
+    return TypedData_Wrap_Struct(klass, &vibe_zstd_cctx_type, cctx);
+}
+
+static VALUE
+vibe_zstd_dctx_alloc(VALUE klass) {
+    vibe_zstd_dctx* dctx = ALLOC(vibe_zstd_dctx);
+    dctx->dctx = ZSTD_createDCtx();
+    if (!dctx->dctx) {
+        ruby_xfree(dctx);
+        rb_raise(rb_eRuntimeError, "Failed to create ZSTD_DCtx");
+    }
+    dctx->initial_capacity = 0;  // 0 = use class default
+    return TypedData_Wrap_Struct(klass, &vibe_zstd_dctx_type, dctx);
+}
+
+static VALUE
+vibe_zstd_cdict_alloc(VALUE klass) {
+    vibe_zstd_cdict* cdict = ALLOC(vibe_zstd_cdict);
+    cdict->cdict = NULL; // Will be set in initialize
+    return TypedData_Wrap_Struct(klass, &vibe_zstd_cdict_type, cdict);
+}
+
+static VALUE
+vibe_zstd_ddict_alloc(VALUE klass) {
+    vibe_zstd_ddict* ddict = ALLOC(vibe_zstd_ddict);
+    ddict->ddict = NULL; // Will be set in initialize
+    return TypedData_Wrap_Struct(klass, &vibe_zstd_ddict_type, ddict);
+}
+
+static VALUE
+vibe_zstd_cstream_alloc(VALUE klass) {
+    vibe_zstd_cstream* cstream = ALLOC(vibe_zstd_cstream);
+    cstream->cstream = NULL;
+    cstream->io = Qnil;
+    return TypedData_Wrap_Struct(klass, &vibe_zstd_cstream_type, cstream);
+}
+
+static VALUE
+vibe_zstd_dstream_alloc(VALUE klass) {
+    vibe_zstd_dstream* dstream = ALLOC(vibe_zstd_dstream);
+    dstream->dstream = NULL;
+    dstream->io = Qnil;
+    dstream->input_data = Qnil;
+    dstream->input.src = NULL;
+    dstream->input.size = 0;
+    dstream->input.pos = 0;
+    return TypedData_Wrap_Struct(klass, &vibe_zstd_dstream_type, dstream);
+}
+
+// Module-level version and compression level functions
+static VALUE
+vibe_zstd_version_number(VALUE self) {
+    (void)self;
+    return UINT2NUM(ZSTD_versionNumber());
+}
+
+static VALUE
+vibe_zstd_version_string(VALUE self) {
+    (void)self;
+    return rb_str_new_cstr(ZSTD_versionString());
+}
+
+static VALUE
+vibe_zstd_min_c_level(VALUE self) {
+    (void)self;
+    return INT2NUM(ZSTD_minCLevel());
+}
+
+static VALUE
+vibe_zstd_max_c_level(VALUE self) {
+    (void)self;
+    return INT2NUM(ZSTD_maxCLevel());
+}
+
+static VALUE
+vibe_zstd_default_c_level(VALUE self) {
+    (void)self;
+    return INT2NUM(ZSTD_defaultCLevel());
+}
+
+// Include the split implementation files
+#include "cctx.c"
+#include "dctx.c"
+#include "dict.c"
+#include "streaming.c"
+#include "frames.c"
+
+// Main initialization function
+RUBY_FUNC_EXPORTED void
+Init_vibe_zstd(void)
+{
+  // Initialize parameter lookup tables
+  init_cctx_param_table();
+  init_dctx_param_table();
+
+  rb_mVibeZstd = rb_define_module("VibeZstd");
+
+  // Define classes
+  rb_cVibeZstdCCtx = rb_define_class_under(rb_mVibeZstd, "CCtx", rb_cObject);
+  rb_cVibeZstdDCtx = rb_define_class_under(rb_mVibeZstd, "DCtx", rb_cObject);
+  rb_cVibeZstdCDict = rb_define_class_under(rb_mVibeZstd, "CDict", rb_cObject);
+  rb_cVibeZstdDDict = rb_define_class_under(rb_mVibeZstd, "DDict", rb_cObject);
+  rb_cVibeZstdCompressWriter = rb_define_class_under(rb_mVibeZstd, "CompressWriter", rb_cObject);
+  rb_cVibeZstdDecompressReader = rb_define_class_under(rb_mVibeZstd, "DecompressReader", rb_cObject);
+
+  // Initialize each subsystem
+  vibe_zstd_cctx_init_class(rb_cVibeZstdCCtx);
+  vibe_zstd_dctx_init_class(rb_cVibeZstdDCtx);
+  vibe_zstd_dict_init_classes(rb_cVibeZstdCDict, rb_cVibeZstdDDict);
+  vibe_zstd_dict_init_module_methods(rb_mVibeZstd);
+  vibe_zstd_streaming_init_classes(rb_cVibeZstdCompressWriter, rb_cVibeZstdDecompressReader);
+  vibe_zstd_frames_init_module_methods(rb_mVibeZstd);
+
+  // Module-level version information
+  rb_define_module_function(rb_mVibeZstd, "version_number", vibe_zstd_version_number, 0);
+  rb_define_module_function(rb_mVibeZstd, "version_string", vibe_zstd_version_string, 0);
+  rb_define_module_function(rb_mVibeZstd, "min_compression_level", vibe_zstd_min_c_level, 0);
+  rb_define_module_function(rb_mVibeZstd, "max_compression_level", vibe_zstd_max_c_level, 0);
+  rb_define_module_function(rb_mVibeZstd, "default_compression_level", vibe_zstd_default_c_level, 0);
+
+  // Aliases
+  rb_define_module_function(rb_mVibeZstd, "min_level", vibe_zstd_min_c_level, 0);
+  rb_define_module_function(rb_mVibeZstd, "max_level", vibe_zstd_max_c_level, 0);
+  rb_define_module_function(rb_mVibeZstd, "default_level", vibe_zstd_default_c_level, 0);
+}

data/ext/vibe_zstd/vibe_zstd.h

@@ -0,0 +1,56 @@
+#ifndef VIBE_ZSTD_H
+#define VIBE_ZSTD_H 1
+
+#include "ruby.h"
+#define ZSTD_STATIC_LINKING_ONLY
+#include <zstd.h>
+
+// TypedData structs
+typedef struct {
+    ZSTD_CCtx* cctx;
+} vibe_zstd_cctx;
+
+typedef struct {
+    ZSTD_DCtx* dctx;
+    size_t initial_capacity;  // Initial capacity for unknown-size decompression (0 = use class default)
+} vibe_zstd_dctx;
+
+typedef struct {
+    ZSTD_CDict* cdict;
+} vibe_zstd_cdict;
+
+typedef struct {
+    ZSTD_DDict* ddict;
+} vibe_zstd_ddict;
+
+typedef struct {
+    ZSTD_CStream* cstream;
+    VALUE io;
+} vibe_zstd_cstream;
+
+typedef struct {
+    ZSTD_DStream* dstream;
+    VALUE io;
+    ZSTD_inBuffer input;  // Zstandard manages the buffer state
+    VALUE input_data;      // Ruby string holding input data
+    int eof;               // Flag to track if we've reached end of stream
+    size_t initial_chunk_size;  // Initial chunk size for unbounded reads (0 = use default)
+} vibe_zstd_dstream;
+
+// TypedData types
+extern rb_data_type_t vibe_zstd_cctx_type;
+extern rb_data_type_t vibe_zstd_dctx_type;
+extern rb_data_type_t vibe_zstd_cdict_type;
+extern rb_data_type_t vibe_zstd_ddict_type;
+extern rb_data_type_t vibe_zstd_cstream_type;
+extern rb_data_type_t vibe_zstd_dstream_type;
+
+// Ruby classes and modules
+extern VALUE rb_cVibeZstdCCtx;
+extern VALUE rb_cVibeZstdDCtx;
+extern VALUE rb_cVibeZstdCDict;
+extern VALUE rb_cVibeZstdDDict;
+extern VALUE rb_cVibeZstdCompressWriter;
+extern VALUE rb_cVibeZstdDecompressReader;
+
+#endif /* VIBE_ZSTD_H */

data/ext/vibe_zstd/vibe_zstd_internal.h

@@ -0,0 +1,27 @@
+#ifndef VIBE_ZSTD_INTERNAL_H
+#define VIBE_ZSTD_INTERNAL_H 1
+
+#include "vibe_zstd.h"
+#include <ruby/thread.h>
+#define ZDICT_STATIC_LINKING_ONLY
+#include <zdict.h>
+
+// Function declarations for cross-file usage
+
+// CCtx functions (cctx.c)
+void vibe_zstd_cctx_init_class(VALUE rb_cVibeZstdCCtx);
+
+// DCtx functions (dctx.c)
+void vibe_zstd_dctx_init_class(VALUE rb_cVibeZstdDCtx);
+
+// Dictionary functions (dict.c)
+void vibe_zstd_dict_init_classes(VALUE rb_cVibeZstdCDict, VALUE rb_cVibeZstdDDict);
+void vibe_zstd_dict_init_module_methods(VALUE rb_mVibeZstd);
+
+// Streaming functions (streaming.c)
+void vibe_zstd_streaming_init_classes(VALUE rb_cVibeZstdCompressWriter, VALUE rb_cVibeZstdDecompressReader);
+
+// Frame utility functions (frames.c)
+void vibe_zstd_frames_init_module_methods(VALUE rb_mVibeZstd);
+
+#endif /* VIBE_ZSTD_INTERNAL_H */

data/lib/vibe_zstd/constants.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+module VibeZstd
+  # Frame format constants for compression and decompression
+  module Format
+    # Standard zstd frame format with magic number (default)
+    STANDARD = 0
+
+    # Zstd frame format without initial 4-byte magic number
+    # Useful to save 4 bytes per frame, but decoder must be explicitly configured
+    MAGICLESS = 1
+  end
+
+  # Compression strategy constants, listed from fastest to strongest
+  module Strategy
+    FAST = 1
+    DFAST = 2
+    GREEDY = 3
+    LAZY = 4
+    LAZY2 = 5
+    BTLAZY2 = 6
+    BTOPT = 7
+    BTULTRA = 8
+    BTULTRA2 = 9
+  end
+
+  # Literal compression mode constants
+  module LiteralCompressionMode
+    # Automatically determine based on compression level
+    # Negative levels = uncompressed, positive levels = compressed
+    AUTO = 0
+
+    # Always attempt Huffman compression
+    # Uncompressed literals still emitted if not profitable
+    HUFFMAN = 1
+
+    # Always emit uncompressed literals
+    UNCOMPRESSED = 2
+  end
+
+  # Dictionary attachment preference constants
+  module DictAttachPref
+    # Automatically choose attachment method (default)
+    AUTO = 0
+
+    # Force dictionary to be attached by reference
+    FORCE_ATTACH = 1
+
+    # Force dictionary to be copied into working context
+    FORCE_COPY = 2
+
+    # Force dictionary to be reloaded
+    FORCE_LOAD = 3
+  end
+
+  # Context reset modes
+  module ResetDirective
+    # Reset session only, keep parameters
+    SESSION = 1
+
+    # Reset parameters to defaults, keep session state
+    PARAMETERS = 2
+
+    # Reset both session and parameters (full reset)
+    BOTH = 3
+  end
+end

data/lib/vibe_zstd/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module VibeZstd
+  VERSION = "1.0.0"
+end

data/lib/vibe_zstd.rb

@@ -0,0 +1,255 @@
+# frozen_string_literal: true
+
+require_relative "vibe_zstd/version"
+require_relative "vibe_zstd/vibe_zstd"
+require_relative "vibe_zstd/constants"
+
+module VibeZstd
+  class Error < StandardError; end
+
+  # Convenience method for one-off compression
+  # Supports all CCtx#compress options: level, dict, pledged_size
+  def self.compress(data, **options)
+    cctx = CCtx.new
+    cctx.compress(data, **options)
+  end
+
+  # Convenience method for one-off decompression
+  # Supports all DCtx#decompress options: dict, initial_capacity
+  def self.decompress(data, **options)
+    dctx = DCtx.new
+    dctx.decompress(data, **options)
+  end
+
+  # Get the decompressed content size from a compressed frame
+  # Returns nil if size is unknown or data is invalid
+  def self.frame_content_size(data)
+    DCtx.frame_content_size(data)
+  end
+
+  # Iterate over all skippable frames in the data
+  # Yields [content, magic_variant, offset] for each skippable frame
+  def self.each_skippable_frame(data)
+    return enum_for(:each_skippable_frame, data) unless block_given?
+
+    offset = 0
+    while offset < data.bytesize
+      frame_data = data.byteslice(offset..-1)
+      frame_size = find_frame_compressed_size(frame_data)
+
+      # Defense: Prevent infinite loop on malformed data
+      # A valid frame must have non-zero size (at minimum: frame header)
+      if frame_size <= 0
+        raise Error, "Invalid frame: zero or negative size at offset #{offset}"
+      end
+
+      if skippable_frame?(frame_data)
+        content, magic_variant = read_skippable_frame(frame_data)
+        yield content, magic_variant, offset
+      end
+
+      offset += frame_size
+    end
+  end
+
+  # Convenient aliases for version/level methods
+  class << self
+    alias_method :min_level, :min_compression_level
+    alias_method :max_level, :max_compression_level
+    alias_method :default_level, :default_compression_level
+  end
+
+  # Add helper method to CDict for creating matching DDict
+  class CDict
+    # Get or create a matching DDict from this CDict's dictionary data
+    # The DDict is cached so it's only created once
+    #
+    # @return [DDict] Decompression dictionary matching this compression dictionary
+    def to_ddict
+      @ddict ||= DDict.new(@dict_data)
+    end
+    alias_method :ddict, :to_ddict
+  end
+
+  # Thread-local context pooling for high-performance reuse
+  # Ideal for Rails/Puma applications where threads are reused across requests
+  #
+  # Example usage:
+  #   # In Rails model with encrypted attributes
+  #   class User < ApplicationRecord
+  #     encrypts :email
+  #     encrypts :preferences, dict: JSON_DICT
+  #   end
+  #
+  #   # Instead of:
+  #   VibeZstd.decompress(data, dict: dict)  # Creates new DCtx each time
+  #
+  #   # Use:
+  #   VibeZstd::ThreadLocal.decompress(data, dict: dict)  # Reuses DCtx per thread
+  #
+  # Memory footprint: ~128KB per DCtx × unique dictionaries × threads
+  # Example: 3 dicts × 5 Puma threads = 1.9MB total
+  #
+  # Note: Only supports per-operation parameters (level, dict, pledged_size, initial_capacity)
+  # Does NOT support context-level settings (nb_workers, checksum_flag, etc.)
+  module ThreadLocal
+    # Compress data using thread-local context pool
+    # Contexts are keyed by dictionary ID for automatic isolation
+    #
+    # @param data [String] Data to compress
+    # @param level [Integer] Compression level (per-operation, can vary)
+    # @param dict [CDict] Compression dictionary (optional)
+    # @param pledged_size [Integer] Expected input size (optional)
+    # @return [String] Compressed data
+    def self.compress(data, level: nil, dict: nil, pledged_size: nil)
+      # Key by dictionary ID, or :default if no dict
+      key = dict ? dict.dict_id : :default
+
+      # Get or create thread-local context pool
+      Thread.current[:vibe_zstd_cctx_pool] ||= {}
+      cctx = Thread.current[:vibe_zstd_cctx_pool][key] ||= VibeZstd::CCtx.new
+
+      # Build options hash
+      options = {}
+      options[:level] = level if level
+      options[:dict] = dict if dict
+      options[:pledged_size] = pledged_size if pledged_size
+
+      cctx.compress(data, **options)
+    end
+
+    # Decompress data using thread-local context pool
+    # Contexts are keyed by dictionary ID for automatic isolation
+    #
+    # @param data [String] Data to decompress
+    # @param dict [DDict] Decompression dictionary (optional)
+    # @param initial_capacity [Integer] Initial buffer size for unknown-size frames (optional)
+    # @return [String] Decompressed data
+    def self.decompress(data, dict: nil, initial_capacity: nil)
+      key = dict ? dict.dict_id : :default
+
+      # Get or create thread-local context pool
+      Thread.current[:vibe_zstd_dctx_pool] ||= {}
+      dctx = Thread.current[:vibe_zstd_dctx_pool][key] ||= VibeZstd::DCtx.new
+
+      # Build options hash
+      options = {}
+      options[:dict] = dict if dict
+      options[:initial_capacity] = initial_capacity if initial_capacity
+
+      # C code will validate dict matches frame requirements
+      dctx.decompress(data, **options)
+    end
+
+    # Clear all thread-local context pools for the current thread
+    # Useful for testing or explicit memory management
+    def self.clear_thread_cache!
+      Thread.current[:vibe_zstd_cctx_pool] = {}
+      Thread.current[:vibe_zstd_dctx_pool] = {}
+      nil
+    end
+
+    # Get statistics about the current thread's context pools
+    # @return [Hash] Pool statistics
+    def self.thread_cache_stats
+      {
+        compression_contexts: Thread.current[:vibe_zstd_cctx_pool]&.size || 0,
+        decompression_contexts: Thread.current[:vibe_zstd_dctx_pool]&.size || 0,
+        compression_keys: Thread.current[:vibe_zstd_cctx_pool]&.keys || [],
+        decompression_keys: Thread.current[:vibe_zstd_dctx_pool]&.keys || []
+      }
+    end
+  end
+
+  class CompressWriter
+    # Block-based resource management
+    # Automatically calls finish when block completes
+    def self.open(io, **options)
+      writer = new(io, **options)
+      return writer unless block_given?
+
+      begin
+        yield writer
+      ensure
+        writer.finish
+      end
+    end
+  end
+
+  class DecompressReader
+    include Enumerable
+
+    # Block-based resource management
+    # Automatically cleans up when block completes
+    def self.open(io, **options)
+      reader = new(io, **options)
+      return reader unless block_given?
+
+      yield reader
+
+      # Reader doesn't have finish, but this ensures cleanup
+    end
+
+    # Read all remaining data
+    def read_all
+      chunks = []
+      while (chunk = read)
+        chunks << chunk
+      end
+      chunks.join
+    end
+
+    # Alias for eof?
+    def eof
+      eof?
+    end
+
+    # Iterate over chunks (required for Enumerable)
+    def each(chunk_size = nil)
+      return enum_for(:each, chunk_size) unless block_given?
+
+      until eof?
+        chunk = read(chunk_size)
+        yield chunk if chunk
+      end
+    end
+
+    # Read a single line (up to newline or EOF)
+    def gets(sep = $/)
+      return nil if eof?
+
+      line = +""
+      until eof?
+        chunk = read(1)
+        break unless chunk
+
+        line << chunk
+        break if chunk.end_with?(sep)
+      end
+
+      line.empty? ? nil : line
+    end
+
+    # Iterate over lines
+    def each_line(sep = $/)
+      return enum_for(:each_line, sep) unless block_given?
+
+      while (line = gets(sep))
+        yield line
+      end
+    end
+
+    # Alias for gets
+    alias_method :readline, :gets
+
+    # Read exactly n bytes, or raise EOFError
+    def readpartial(maxlen)
+      raise EOFError, "end of file reached" if eof?
+
+      data = read(maxlen)
+      raise EOFError, "end of file reached" if data.nil?
+
+      data
+    end
+  end
+end