ruby-fst 0.2.2 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: f41a45f06fbbb333bdd02f3d6fa965152f46293e9c2d9043dbe95525f207e5ca
-  data.tar.gz: 4eb47ab9eaefd2f43e908f0baf6a0f46d8cd7adf66a982a60de5c4fe5df0aa5d
+  metadata.gz: eaf4d02399b3d1c106a6f0de243e84e644982fb3af15cb5b2e95416580a1cf6f
+  data.tar.gz: 247642cba615720e22d7f92ab429ed3a6262bc885fe0e8ef30d09f57a0105251
 SHA512:
-  metadata.gz: 11132c92d5f855fa187a0ff28ecb3b16e25aca77ad6f7165ddf178f70be5837ecdf63bec2f43e76b84e2e941b27e6737d1e9e600d92882ea63d63288502122ec
-  data.tar.gz: '079ec8b103245ffb14c60ea1a36500d02ad61215cc2f99484736cbc59ab93f4901d340d6978380fe33b40dff6310f1d78364817d67b63988487965a26eb4895a'
+  metadata.gz: 981010f25a4534daeee239017575822f1a64365e02cb80d80fc71504b0638c92ead519c6175751c8239c5cb28bf0bd252cde3c6dfabd11c758e5b7bcefcbf05c
+  data.tar.gz: fa2698b0f96ef514d17fc98071d09c9410b648b4e2853ef3aa7533cf49abcdc7490452c5ffb33841648d9cbd7007dfa32df352b77c05339a84f2e30ed0e0ceeb

data/CHANGELOG.md

@@ -6,6 +6,29 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+## [0.3.0] — 2026-10-26
+
+### Changed
+- I/O failures now raise `IOError` instead of `RuntimeError`, and Levenshtein queries too large to
+  compile raise `ArgumentError`. Other errors raise the new `RubyFst::Error`, a `RuntimeError`
+  subclass, so existing `rescue RuntimeError` for builder and format errors keeps working.
+- `new` and `from_path` verify the FST checksum. Corrupt data now raises `RubyFst::Error`; before,
+  it could return wrong results or terminate the process. FSTs written by fst < 0.4 have no
+  checksum and still load unverified. `from_path_mmap` does not verify.
+- Paths accept `Pathname`, and on Unix non-UTF-8 file names.
+- `each` and `search_levenshtein` return an `Enumerator` when called without a block.
+- In-memory FSTs report their size to Ruby's GC.
+
+- `save` writes a temporary file and renames it into place, so readers never see a partial file.
+  Symlinks are followed and permissions kept, but the replacement is a new file: hard links to the
+  old one are broken and ownership becomes the saving user's. Where a rename can't work (a
+  non-regular target such as a FIFO, or an unwritable directory) it writes in place as before. A
+  crash mid-save can leave a `<path>.tmp-<pid>` file behind.
+
+### Fixed
+- Saving over a memory-mapped file no longer crashes the process with `SIGBUS` on Linux and macOS.
+  Windows can't replace a mapped file, so `save` raises `IOError` there.
+
 ## [0.2.2] — 2026-09-24
 
 ### Fixed

data/Cargo.lock

@@ -246,7 +246,7 @@ checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a"
 
 [[package]]
 name = "ruby_fst"
-version = "0.2.2"
+version = "0.3.0"
 dependencies = [
  "fst",
  "magnus",

data/README.md

@@ -25,7 +25,7 @@ gem 'ruby-fst'
 
 All FST keys are arbitrary byte strings (`Encoding::BINARY`). When a key is returned from the gem (e.g. via `each`, `get_le`, `range`) it always carries the `Encoding::BINARY` encoding — be careful when interpolating into UTF-8 strings or `puts`-ing keys that contain non-ASCII bytes.
 
-Keys must be inserted into a builder in **strictly ascending lexicographic byte order**, with no duplicates. Out-of-order or duplicate inserts raise.
+Keys must be inserted into a builder in **strictly ascending lexicographic byte order**, with no duplicates. Out-of-order or duplicate inserts raise `RubyFst::Error`.
 
 ## Map
 
@@ -156,7 +156,15 @@ map = RubyFst::Map.from_path('/path/to/file.fst')
 map = RubyFst::Map.from_path_mmap('/path/to/file.fst')
 ```
 
-When using `from_path_mmap`, the file must remain unchanged on disk for the lifetime of the resulting `Map` or `Set` — modifying or truncating it causes undefined behavior.
+`new` and `from_path` verify the FST's checksum (FSTs written before fst 0.4 have none and load unverified) and raise `RubyFst::Error` on corrupt or invalid data. `from_path_mmap` skips the check so that opening a large file stays cheap, which means a corrupt file can return wrong results or terminate the process — only memory-map files you trust.
+
+When using `from_path_mmap`, the file must remain unchanged on disk for the lifetime of the resulting `Map` or `Set` — modifying or truncating it in place causes undefined behavior. `save` writes to a temporary file beside the target and renames it into place, so on Linux and macOS saving over a mapped file is safe: the existing mapping keeps reading the old contents. Windows refuses to replace a file while it is mapped, so there `save` raises `IOError` instead.
+
+## Errors
+
+- `IOError` — a file couldn't be read, mapped, or written.
+- `ArgumentError` — a Levenshtein query too large to build an automaton for.
+- `RubyFst::Error` (a `RuntimeError` subclass) — invalid or corrupt FST data, out-of-order builder inserts, or a builder used after `finish`.
 
 ## Development
 

data/ext/ruby_fst/Cargo.toml

@@ -1,7 +1,8 @@
 [package]
 name = "ruby_fst"
-version = "0.2.2"
-edition = "2021"
+version = "0.3.0"
+edition = "2024"
+rust-version = "1.85"
 publish = false
 
 [lib]

data/ext/ruby_fst/src/lib.rs

@@ -1,28 +1,56 @@
 use std::cell::RefCell;
+use std::ffi::OsString;
 use std::fs::{self, File};
+use std::io::{self, Write};
+use std::path::{Path, PathBuf};
 
-use fst::automaton::Levenshtein;
-use fst::raw::{CompiledAddr, Fst, Node, Output};
+use fst::automaton::{Automaton, Levenshtein};
+use fst::raw::{CompiledAddr, Fst, Node, Output, Stream};
 use fst::{IntoStreamer, Streamer};
-use magnus::prelude::*;
-use magnus::{function, method, Error, RArray, RString, Ruby, Value};
+use magnus::value::Lazy;
+use magnus::{
+    DataTypeFunctions, Error, ExceptionClass, RModule, RString, Ruby, TypedData, Value, function,
+    method, prelude::*,
+};
 use memmap2::Mmap;
 
-fn err(msg: impl std::fmt::Display) -> Error {
-    Error::new(ruby().exception_runtime_error(), msg.to_string())
+// ---------------------------------------------------------------------------
+// Errors
+//
+// I/O failures raise IOError, a rejected Levenshtein query raises
+// ArgumentError, and everything else (invalid or corrupt FST data, misuse of
+// a builder) raises RubyFst::Error, a RuntimeError subclass.
+// ---------------------------------------------------------------------------
+
+static FST_ERROR: Lazy<ExceptionClass> = Lazy::new(|ruby| {
+    ruby.class_object()
+        .const_get::<_, RModule>("RubyFst")
+        .and_then(|module| module.const_get("Error"))
+        .expect("RubyFst::Error is defined in init")
+});
+
+fn fst_error(ruby: &Ruby, msg: impl std::fmt::Display) -> Error {
+    Error::new(ruby.get_inner(&FST_ERROR), msg.to_string())
+}
+
+fn io_error(ruby: &Ruby, path: &Path, err: io::Error) -> Error {
+    Error::new(
+        ruby.exception_io_error(),
+        format!("{}: {err}", path.display()),
+    )
 }
 
-// SAFETY: every method exposed to Ruby is invoked by the VM on the GVL-owning
-// thread, so `Ruby::get_unchecked()` is sound — calling code must not invoke
-// this from a non-Ruby thread.
-fn ruby() -> Ruby {
-    unsafe { Ruby::get_unchecked() }
+fn map_fst_error(ruby: &Ruby, err: fst::Error) -> Error {
+    match err {
+        fst::Error::Io(err) => Error::new(ruby.exception_io_error(), err.to_string()),
+        fst::Error::Fst(err) => fst_error(ruby, err),
+    }
 }
 
 // ---------------------------------------------------------------------------
-// Storage: backing buffer for an FST. Either an owned heap allocation or a
-// memory map. Both implement AsRef<[u8]> so fst::Map / fst::Set can accept
-// either without monomorphising the wrapper structs.
+// Storage: the backing buffer for an FST, either an owned heap allocation or a
+// memory map. A heap buffer is reported to Ruby's GC so that large FSTs count
+// towards its memory pressure; a map is backed by the file, so it isn't.
 // ---------------------------------------------------------------------------
 
 enum Storage {
@@ -30,36 +58,122 @@ enum Storage {
     Mmap(Mmap),
 }
 
+impl Storage {
+    fn mem(ruby: &Ruby, bytes: Vec<u8>) -> Self {
+        ruby.gc_adjust_memory_usage(heap_size(&bytes));
+        Storage::Mem(bytes)
+    }
+
+    fn read(ruby: &Ruby, path: &Path) -> Result<Self, Error> {
+        let bytes = fs::read(path).map_err(|e| io_error(ruby, path, e))?;
+        Ok(Self::mem(ruby, bytes))
+    }
+
+    fn mmap(ruby: &Ruby, path: &Path) -> Result<Self, Error> {
+        let file = File::open(path).map_err(|e| io_error(ruby, path, e))?;
+        // SAFETY: modifying or truncating a mapped file is undefined behavior,
+        // and nothing here can prevent another process from doing it. The
+        // README makes keeping the file unchanged the caller's responsibility,
+        // and `save` replaces files by rename, so it never truncates a mapping.
+        let map = unsafe { Mmap::map(&file) }.map_err(|e| io_error(ruby, path, e))?;
+        Ok(Storage::Mmap(map))
+    }
+}
+
 impl AsRef<[u8]> for Storage {
     fn as_ref(&self) -> &[u8] {
         match self {
-            Storage::Mem(v) => v.as_ref(),
-            Storage::Mmap(m) => m.as_ref(),
+            Storage::Mem(bytes) => bytes,
+            Storage::Mmap(map) => map,
         }
     }
 }
 
-fn read_storage(path: &str) -> Result<Storage, Error> {
-    let data = fs::read(path).map_err(err)?;
-    Ok(Storage::Mem(data))
+impl Drop for Storage {
+    fn drop(&mut self) {
+        if let Storage::Mem(bytes) = self {
+            // Dropped during GC finalization, which runs on a Ruby thread.
+            if let Ok(ruby) = Ruby::get() {
+                ruby.gc_adjust_memory_usage(-heap_size(bytes));
+            }
+        }
+    }
 }
 
-fn mmap_storage(path: &str) -> Result<Storage, Error> {
-    let file = File::open(path).map_err(err)?;
-    // SAFETY: callers opt in to mmap and accept the contract that the file
-    // must not be modified or truncated while the FST is alive. We document
-    // this in the Ruby-level docs.
-    let mmap = unsafe { Mmap::map(&file) }.map_err(err)?;
-    Ok(Storage::Mmap(mmap))
+fn heap_size(bytes: &[u8]) -> isize {
+    isize::try_from(bytes.len()).unwrap_or(isize::MAX)
 }
 
-// SAFETY for `RString::as_slice()` callers below: between obtaining the slice
-// and dropping it, no Ruby allocation, GC trigger, or string mutation occurs.
-// We either copy into a Vec immediately or pass the slice into pure-Rust fst
-// operations that never re-enter Ruby.
+// Opens an FST over `storage`. Heap buffers are checksummed first, because the
+// fst crate panics on corrupt data and a panic here terminates the Ruby
+// process. FSTs written before fst 0.4 carry no checksum and are accepted
+// unverified, as they always were. Memory maps skip the check so that opening
+// stays lazy; the README documents that they must only be used on trusted files.
+fn open_fst(ruby: &Ruby, storage: Storage) -> Result<Fst<Storage>, Error> {
+    let verify = matches!(storage, Storage::Mem(_));
+    let fst = Fst::new(storage).map_err(|e| fst_error(ruby, e))?;
+    if verify {
+        match fst.verify() {
+            Ok(()) | Err(fst::Error::Fst(fst::raw::Error::ChecksumMissing)) => {}
+            Err(e) => return Err(fst_error(ruby, e)),
+        }
+    }
+    Ok(fst)
+}
 
-fn rstring_to_vec(s: RString) -> Vec<u8> {
-    unsafe { s.as_slice() }.to_vec()
+// Runs `f` over the bytes of a Ruby string without copying them. `f` must not
+// call into Ruby; pass it pure-Rust work and build Ruby objects from its result.
+fn with_bytes<R>(s: RString, f: impl FnOnce(&[u8]) -> R + Send) -> R {
+    // SAFETY: the slice borrows memory owned by Ruby, which is sound only while
+    // no Ruby code runs, since a yield could mutate the string. The `Send` bound
+    // stops `f` from capturing a `Ruby` handle or any Ruby value (all `!Send`),
+    // and the callers in this file don't fetch one with `Ruby::get` either, so
+    // nothing in `f` can run Ruby code. The higher-ranked bound stops the result
+    // borrowing the slice, so it can't outlive this call; `s` is a method
+    // argument, so Ruby keeps it alive for the duration.
+    f(unsafe { s.as_slice() })
+}
+
+// Writes `bytes` to a temporary file beside the target and renames it into
+// place, so a reader never sees a partial file and a live memory map of the old
+// file is never truncated underneath it. A symlink is resolved first so its
+// target is replaced rather than the link, and an existing file's permissions
+// carry over. Where a rename can't work — the target isn't a regular file, or
+// its directory isn't writable — this writes in place, as `fs::write` would.
+fn write_atomically(path: &Path, bytes: &[u8]) -> io::Result<()> {
+    let target = match fs::canonicalize(path) {
+        Ok(target) => target,
+        Err(e) if e.kind() == io::ErrorKind::NotFound && !path.is_symlink() => path.to_path_buf(),
+        Err(_) => return fs::write(path, bytes),
+    };
+    let existing = fs::metadata(&target).ok();
+    if existing.as_ref().is_some_and(|meta| !meta.is_file()) {
+        return fs::write(&target, bytes);
+    }
+
+    let mut tmp_name = OsString::from(target.as_os_str());
+    tmp_name.push(format!(".tmp-{}", std::process::id()));
+    let tmp_path = PathBuf::from(tmp_name);
+    let mut file = match File::create(&tmp_path) {
+        Ok(file) => file,
+        Err(e) if e.kind() == io::ErrorKind::PermissionDenied => {
+            return fs::write(&target, bytes);
+        }
+        Err(e) => return Err(e),
+    };
+
+    let result = file
+        .write_all(bytes)
+        .and_then(|()| match &existing {
+            Some(meta) => file.set_permissions(meta.permissions()),
+            None => Ok(()),
+        })
+        .and_then(|()| file.sync_all())
+        .and_then(|()| fs::rename(&tmp_path, &target));
+    if result.is_err() {
+        let _ = fs::remove_file(&tmp_path);
+    }
+    result
 }
 
 // Returns the exclusive upper bound for a prefix scan: the smallest byte
@@ -79,291 +193,230 @@ fn prefix_upper_bound(prefix: &[u8]) -> Option<Vec<u8>> {
 }
 
 // ---------------------------------------------------------------------------
-// Map
+// Map and Set share one implementation over the raw FST; they differ only in
+// what they yield for each entry. The methods below are generic over that
+// difference and registered once per class in `init`.
 // ---------------------------------------------------------------------------
 
-#[magnus::wrap(class = "RubyFst::Map", free_immediately, size)]
+#[derive(TypedData)]
+#[magnus(class = "RubyFst::Map", free_immediately, size)]
 struct FstMap {
-    inner: fst::Map<Storage>,
+    fst: Fst<Storage>,
 }
 
-impl FstMap {
-    fn new(bytes: RString) -> Result<Self, Error> {
-        let inner = fst::Map::new(Storage::Mem(rstring_to_vec(bytes))).map_err(err)?;
-        Ok(Self { inner })
-    }
+#[derive(TypedData)]
+#[magnus(class = "RubyFst::Set", free_immediately, size)]
+struct FstSet {
+    fst: Fst<Storage>,
+}
 
-    fn from_path(path: String) -> Result<Self, Error> {
-        let inner = fst::Map::new(read_storage(&path)?).map_err(err)?;
-        Ok(Self { inner })
-    }
+trait FstObject: TypedData + Sized {
+    fn from_fst(fst: Fst<Storage>) -> Self;
+    fn fst(&self) -> &Fst<Storage>;
+    fn yield_entry(ruby: &Ruby, key: &[u8], value: u64) -> Result<(), Error>;
+}
 
-    fn from_path_mmap(path: String) -> Result<Self, Error> {
-        let inner = fst::Map::new(mmap_storage(&path)?).map_err(err)?;
-        Ok(Self { inner })
+impl FstObject for FstMap {
+    fn from_fst(fst: Fst<Storage>) -> Self {
+        Self { fst }
     }
 
-    fn get(&self, key: RString) -> Option<u64> {
-        let key = unsafe { key.as_slice() };
-        self.inner.get(key)
+    fn fst(&self) -> &Fst<Storage> {
+        &self.fst
     }
 
-    fn contains(&self, key: RString) -> bool {
-        let key = unsafe { key.as_slice() };
-        self.inner.contains_key(key)
+    fn yield_entry(ruby: &Ruby, key: &[u8], value: u64) -> Result<(), Error> {
+        ruby.yield_values::<_, Value>((ruby.str_from_slice(key), value))
+            .map(drop)
     }
+}
 
-    fn len(&self) -> usize {
-        self.inner.len()
+impl FstObject for FstSet {
+    fn from_fst(fst: Fst<Storage>) -> Self {
+        Self { fst }
     }
 
-    fn is_empty(&self) -> bool {
-        self.inner.is_empty()
+    fn fst(&self) -> &Fst<Storage> {
+        &self.fst
     }
 
-    fn to_bytes(&self) -> RString {
-        ruby().str_from_slice(self.inner.as_fst().as_bytes())
+    fn yield_entry(ruby: &Ruby, key: &[u8], _value: u64) -> Result<(), Error> {
+        ruby.yield_value::<_, Value>(ruby.str_from_slice(key))
+            .map(drop)
     }
+}
 
-    fn save(&self, path: String) -> Result<(), Error> {
-        fs::write(&path, self.inner.as_fst().as_bytes()).map_err(err)
+impl DataTypeFunctions for FstMap {
+    fn size(&self) -> usize {
+        size_of_val(self) + heap_bytes(&self.fst)
     }
+}
 
-    fn get_le(&self, key: RString) -> Result<Option<RArray>, Error> {
-        let r = ruby();
-        let key = unsafe { key.as_slice() };
-        match floor_lookup(self.inner.as_fst(), key) {
-            Some((found_key, value)) => {
-                let arr = r.ary_new_capa(2);
-                arr.push(r.str_from_slice(&found_key))?;
-                arr.push(value)?;
-                Ok(Some(arr))
-            }
-            None => Ok(None),
-        }
+impl DataTypeFunctions for FstSet {
+    fn size(&self) -> usize {
+        size_of_val(self) + heap_bytes(&self.fst)
     }
+}
 
-    fn get_le_value(&self, key: RString) -> Option<u64> {
-        let key = unsafe { key.as_slice() };
-        floor_value(self.inner.as_fst(), key)
+fn heap_bytes(fst: &Fst<Storage>) -> usize {
+    match fst.as_inner() {
+        Storage::Mem(bytes) => bytes.len(),
+        Storage::Mmap(_) => 0,
     }
+}
 
-    fn get_ge(&self, key: RString) -> Result<Option<RArray>, Error> {
-        let r = ruby();
-        let key = unsafe { key.as_slice() };
-        let mut stream = self.inner.range().ge(key).into_stream();
-        match stream.next() {
-            Some((k, v)) => {
-                let arr = r.ary_new_capa(2);
-                arr.push(r.str_from_slice(k))?;
-                arr.push(v)?;
-                Ok(Some(arr))
-            }
-            None => Ok(None),
-        }
-    }
+fn new<T: FstObject>(ruby: &Ruby, bytes: RString) -> Result<T, Error> {
+    let bytes = with_bytes(bytes, <[u8]>::to_vec);
+    open_fst(ruby, Storage::mem(ruby, bytes)).map(T::from_fst)
+}
 
-    fn get_ge_value(&self, key: RString) -> Option<u64> {
-        let key = unsafe { key.as_slice() };
-        let mut stream = self.inner.range().ge(key).into_stream();
-        stream.next().map(|(_, v)| v)
-    }
+fn from_path<T: FstObject>(ruby: &Ruby, path: PathBuf) -> Result<T, Error> {
+    open_fst(ruby, Storage::read(ruby, &path)?).map(T::from_fst)
+}
 
-    fn each(&self) -> Result<(), Error> {
-        let r = ruby();
-        let mut stream = (&self.inner).into_stream();
-        while let Some((key, value)) = stream.next() {
-            let rb_key = r.str_from_slice(key);
-            let _: Value = r.yield_values((rb_key, value))?;
-        }
-        Ok(())
-    }
+fn from_path_mmap<T: FstObject>(ruby: &Ruby, path: PathBuf) -> Result<T, Error> {
+    open_fst(ruby, Storage::mmap(ruby, &path)?).map(T::from_fst)
+}
 
-    fn range(&self, ge: Option<RString>, le: Option<RString>) -> Result<(), Error> {
-        let r = ruby();
-        let ge_bytes = ge.map(rstring_to_vec);
-        let le_bytes = le.map(rstring_to_vec);
-        let mut builder = self.inner.range();
-        if let Some(ref b) = ge_bytes {
-            builder = builder.ge(b);
-        }
-        if let Some(ref b) = le_bytes {
-            builder = builder.le(b);
-        }
-        let mut stream = builder.into_stream();
-        while let Some((key, value)) = stream.next() {
-            let rb_key = r.str_from_slice(key);
-            let _: Value = r.yield_values((rb_key, value))?;
-        }
-        Ok(())
-    }
+fn contains<T: FstObject>(rb_self: &T, key: RString) -> bool {
+    let fst = rb_self.fst();
+    with_bytes(key, |key| fst.contains_key(key))
+}
 
-    fn starts_with(&self, prefix: RString) -> Result<(), Error> {
-        let r = ruby();
-        let prefix_bytes = rstring_to_vec(prefix);
-        let upper = prefix_upper_bound(&prefix_bytes);
-        let mut builder = self.inner.range().ge(&prefix_bytes);
-        if let Some(ref u) = upper {
-            builder = builder.lt(u);
-        }
-        let mut stream = builder.into_stream();
-        while let Some((key, value)) = stream.next() {
-            let rb_key = r.str_from_slice(key);
-            let _: Value = r.yield_values((rb_key, value))?;
-        }
-        Ok(())
-    }
+fn len<T: FstObject>(rb_self: &T) -> usize {
+    rb_self.fst().len()
+}
 
-    fn search_levenshtein(&self, query: String, distance: u32) -> Result<(), Error> {
-        let r = ruby();
-        let lev = Levenshtein::new(&query, distance).map_err(err)?;
-        let mut stream = self.inner.search(lev).into_stream();
-        while let Some((key, value)) = stream.next() {
-            let rb_key = r.str_from_slice(key);
-            let _: Value = r.yield_values((rb_key, value))?;
-        }
-        Ok(())
-    }
+fn is_empty<T: FstObject>(rb_self: &T) -> bool {
+    rb_self.fst().is_empty()
 }
 
-// ---------------------------------------------------------------------------
-// MapBuilder
-// ---------------------------------------------------------------------------
+fn to_bytes<T: FstObject>(ruby: &Ruby, rb_self: &T) -> RString {
+    ruby.str_from_slice(rb_self.fst().as_bytes())
+}
 
-#[magnus::wrap(class = "RubyFst::MapBuilder", free_immediately, size)]
-struct FstMapBuilder {
-    inner: RefCell<Option<fst::MapBuilder<Vec<u8>>>>,
+fn save<T: FstObject>(ruby: &Ruby, rb_self: &T, path: PathBuf) -> Result<(), Error> {
+    write_atomically(&path, rb_self.fst().as_bytes()).map_err(|e| io_error(ruby, &path, e))
 }
 
-impl FstMapBuilder {
-    fn new() -> Self {
-        Self {
-            inner: RefCell::new(Some(fst::MapBuilder::memory())),
-        }
+fn yield_all<T: FstObject, A: Automaton>(
+    ruby: &Ruby,
+    mut stream: Stream<'_, A>,
+) -> Result<(), Error> {
+    while let Some((key, output)) = stream.next() {
+        T::yield_entry(ruby, key, output.value())?;
     }
+    Ok(())
+}
 
-    fn insert(&self, key: RString, value: u64) -> Result<(), Error> {
-        let key = rstring_to_vec(key);
-        let mut guard = self.inner.borrow_mut();
-        let b = guard.as_mut().ok_or_else(|| err("builder already finished"))?;
-        b.insert(&key, value).map_err(err)
-    }
+fn each<T: FstObject>(ruby: &Ruby, rb_self: &T) -> Result<(), Error> {
+    yield_all::<T, _>(ruby, rb_self.fst().stream())
+}
 
-    fn finish(&self) -> Result<RString, Error> {
-        let mut guard = self.inner.borrow_mut();
-        let b = guard.take().ok_or_else(|| err("builder already finished"))?;
-        let bytes = b.into_inner().map_err(err)?;
-        Ok(ruby().str_from_slice(&bytes))
+fn range<T: FstObject>(
+    ruby: &Ruby,
+    rb_self: &T,
+    ge: Option<RString>,
+    le: Option<RString>,
+) -> Result<(), Error> {
+    let mut builder = rb_self.fst().range();
+    if let Some(ge) = ge {
+        builder = with_bytes(ge, |bound| builder.ge(bound));
     }
+    if let Some(le) = le {
+        builder = with_bytes(le, |bound| builder.le(bound));
+    }
+    yield_all::<T, _>(ruby, builder.into_stream())
 }
 
-// ---------------------------------------------------------------------------
-// Set
-// ---------------------------------------------------------------------------
-
-#[magnus::wrap(class = "RubyFst::Set", free_immediately, size)]
-struct FstSet {
-    inner: fst::Set<Storage>,
+fn starts_with<T: FstObject>(ruby: &Ruby, rb_self: &T, prefix: RString) -> Result<(), Error> {
+    let fst = rb_self.fst();
+    let builder = with_bytes(prefix, |prefix| {
+        let builder = fst.range().ge(prefix);
+        match prefix_upper_bound(prefix) {
+            Some(upper) => builder.lt(upper),
+            None => builder,
+        }
+    });
+    yield_all::<T, _>(ruby, builder.into_stream())
 }
 
-impl FstSet {
-    fn new(bytes: RString) -> Result<Self, Error> {
-        let inner = fst::Set::new(Storage::Mem(rstring_to_vec(bytes))).map_err(err)?;
-        Ok(Self { inner })
-    }
+fn search_levenshtein<T: FstObject>(
+    ruby: &Ruby,
+    rb_self: &T,
+    query: String,
+    distance: u32,
+) -> Result<(), Error> {
+    let automaton = Levenshtein::new(&query, distance)
+        .map_err(|e| Error::new(ruby.exception_arg_error(), e.to_string()))?;
+    yield_all::<T, _>(ruby, rb_self.fst().search(automaton).into_stream())
+}
 
-    fn from_path(path: String) -> Result<Self, Error> {
-        let inner = fst::Set::new(read_storage(&path)?).map_err(err)?;
-        Ok(Self { inner })
-    }
+// Map-only lookups. Keys come back as freshly allocated Ruby strings, so each
+// lookup finishes with the borrowed key before building its result.
 
-    fn from_path_mmap(path: String) -> Result<Self, Error> {
-        let inner = fst::Set::new(mmap_storage(&path)?).map_err(err)?;
-        Ok(Self { inner })
-    }
+fn get(rb_self: &FstMap, key: RString) -> Option<u64> {
+    with_bytes(key, |key| rb_self.fst.get(key).map(Output::value))
+}
 
-    fn contains(&self, key: RString) -> bool {
-        let key = unsafe { key.as_slice() };
-        self.inner.contains(key)
-    }
+fn get_le(ruby: &Ruby, rb_self: &FstMap, key: RString) -> Option<(RString, u64)> {
+    let (found, value) = with_bytes(key, |key| floor_entry(&rb_self.fst, key))?;
+    Some((ruby.str_from_slice(&found), value))
+}
 
-    fn len(&self) -> usize {
-        self.inner.len()
-    }
+fn get_le_value(rb_self: &FstMap, key: RString) -> Option<u64> {
+    with_bytes(key, |key| floor_value(&rb_self.fst, key))
+}
 
-    fn is_empty(&self) -> bool {
-        self.inner.is_empty()
-    }
+fn get_ge(ruby: &Ruby, rb_self: &FstMap, key: RString) -> Option<(RString, u64)> {
+    let mut stream = with_bytes(key, |key| rb_self.fst.range().ge(key)).into_stream();
+    let (found, output) = stream.next()?;
+    Some((ruby.str_from_slice(found), output.value()))
+}
 
-    fn to_bytes(&self) -> RString {
-        ruby().str_from_slice(self.inner.as_fst().as_bytes())
-    }
+fn get_ge_value(rb_self: &FstMap, key: RString) -> Option<u64> {
+    let builder = with_bytes(key, |key| rb_self.fst.range().ge(key));
+    builder
+        .into_stream()
+        .next()
+        .map(|(_, output)| output.value())
+}
 
-    fn save(&self, path: String) -> Result<(), Error> {
-        fs::write(&path, self.inner.as_fst().as_bytes()).map_err(err)
-    }
+// ---------------------------------------------------------------------------
+// Builders
+// ---------------------------------------------------------------------------
 
-    fn each(&self) -> Result<(), Error> {
-        let r = ruby();
-        let mut stream = (&self.inner).into_stream();
-        while let Some(key) = stream.next() {
-            let rb_key = r.str_from_slice(key);
-            let _: Value = r.yield_value(rb_key)?;
-        }
-        Ok(())
-    }
+#[magnus::wrap(class = "RubyFst::MapBuilder", free_immediately, size)]
+struct FstMapBuilder {
+    inner: RefCell<Option<fst::MapBuilder<Vec<u8>>>>,
+}
 
-    fn range(&self, ge: Option<RString>, le: Option<RString>) -> Result<(), Error> {
-        let r = ruby();
-        let ge_bytes = ge.map(rstring_to_vec);
-        let le_bytes = le.map(rstring_to_vec);
-        let mut builder = self.inner.range();
-        if let Some(ref b) = ge_bytes {
-            builder = builder.ge(b);
-        }
-        if let Some(ref b) = le_bytes {
-            builder = builder.le(b);
-        }
-        let mut stream = builder.into_stream();
-        while let Some(key) = stream.next() {
-            let rb_key = r.str_from_slice(key);
-            let _: Value = r.yield_value(rb_key)?;
+impl FstMapBuilder {
+    fn new() -> Self {
+        Self {
+            inner: RefCell::new(Some(fst::MapBuilder::memory())),
         }
-        Ok(())
     }
 
-    fn starts_with(&self, prefix: RString) -> Result<(), Error> {
-        let r = ruby();
-        let prefix_bytes = rstring_to_vec(prefix);
-        let upper = prefix_upper_bound(&prefix_bytes);
-        let mut builder = self.inner.range().ge(&prefix_bytes);
-        if let Some(ref u) = upper {
-            builder = builder.lt(u);
-        }
-        let mut stream = builder.into_stream();
-        while let Some(key) = stream.next() {
-            let rb_key = r.str_from_slice(key);
-            let _: Value = r.yield_value(rb_key)?;
-        }
-        Ok(())
+    fn insert(ruby: &Ruby, rb_self: &Self, key: RString, value: u64) -> Result<(), Error> {
+        let mut guard = rb_self.inner.borrow_mut();
+        let builder = guard
+            .as_mut()
+            .ok_or_else(|| fst_error(ruby, "builder already finished"))?;
+        with_bytes(key, |key| builder.insert(key, value)).map_err(|e| map_fst_error(ruby, e))
     }
 
-    fn search_levenshtein(&self, query: String, distance: u32) -> Result<(), Error> {
-        let r = ruby();
-        let lev = Levenshtein::new(&query, distance).map_err(err)?;
-        let mut stream = self.inner.search(lev).into_stream();
-        while let Some(key) = stream.next() {
-            let rb_key = r.str_from_slice(key);
-            let _: Value = r.yield_value(rb_key)?;
-        }
-        Ok(())
+    fn finish(ruby: &Ruby, rb_self: &Self) -> Result<RString, Error> {
+        let builder = rb_self
+            .inner
+            .borrow_mut()
+            .take()
+            .ok_or_else(|| fst_error(ruby, "builder already finished"))?;
+        let bytes = builder.into_inner().map_err(|e| map_fst_error(ruby, e))?;
+        Ok(ruby.str_from_slice(&bytes))
     }
 }
 
-// ---------------------------------------------------------------------------
-// SetBuilder
-// ---------------------------------------------------------------------------
-
 #[magnus::wrap(class = "RubyFst::SetBuilder", free_immediately, size)]
 struct FstSetBuilder {
     inner: RefCell<Option<fst::SetBuilder<Vec<u8>>>>,
@@ -376,47 +429,91 @@ impl FstSetBuilder {
         }
     }
 
-    fn insert(&self, key: RString) -> Result<(), Error> {
-        let key = rstring_to_vec(key);
-        let mut guard = self.inner.borrow_mut();
-        let b = guard.as_mut().ok_or_else(|| err("builder already finished"))?;
-        b.insert(&key).map_err(err)
+    fn insert(ruby: &Ruby, rb_self: &Self, key: RString) -> Result<(), Error> {
+        let mut guard = rb_self.inner.borrow_mut();
+        let builder = guard
+            .as_mut()
+            .ok_or_else(|| fst_error(ruby, "builder already finished"))?;
+        with_bytes(key, |key| builder.insert(key)).map_err(|e| map_fst_error(ruby, e))
     }
 
-    fn finish(&self) -> Result<RString, Error> {
-        let mut guard = self.inner.borrow_mut();
-        let b = guard.take().ok_or_else(|| err("builder already finished"))?;
-        let bytes = b.into_inner().map_err(err)?;
-        Ok(ruby().str_from_slice(&bytes))
+    fn finish(ruby: &Ruby, rb_self: &Self) -> Result<RString, Error> {
+        let builder = rb_self
+            .inner
+            .borrow_mut()
+            .take()
+            .ok_or_else(|| fst_error(ruby, "builder already finished"))?;
+        let bytes = builder.into_inner().map_err(|e| map_fst_error(ruby, e))?;
+        Ok(ruby.str_from_slice(&bytes))
     }
 }
 
 // ---------------------------------------------------------------------------
 // Floor lookup (get_le): greatest key <= query.
 //
-// The upstream fst crate exposes `range().ge(...)` natively but not a floor
-// operation, so this walks the FST manually: descend matching the query while
-// recording, at each step, the rightmost transition strictly less than the
-// current query byte. If exact match fails we backtrack to the most recent
-// such transition and follow the rightmost path to a leaf.
+// The fst crate exposes `range().ge(...)` natively but not a floor operation,
+// so this walks the FST manually. While descending along the query it keeps
+// the best candidate seen so far: the rightmost transition strictly below the
+// query byte at that depth, or else the prefix itself if it is a key. A deeper
+// candidate always beats a shallower one, because it agrees with the query for
+// longer, so only the latest needs keeping. An exact match beats them all.
 // ---------------------------------------------------------------------------
 
-struct Frame {
-    node_addr: CompiledAddr,
-    output: Output,
-    prefix_len: usize,
-    max_lesser_idx: Option<usize>,
-    is_final: bool,
-    final_value: u64,
+enum Floor {
+    // `key[..len]` is itself in the FST, with this output.
+    Prefix {
+        len: usize,
+        output: Output,
+    },
+    // `key[..len]`, then `input`, then the rightmost path below `addr`.
+    Branch {
+        len: usize,
+        input: u8,
+        addr: CompiledAddr,
+        output: Output,
+    },
 }
 
-fn find_max_lesser(node: &Node, byte: u8) -> Option<usize> {
-    let n = node.len();
-    if n == 0 {
-        return None;
+fn find_floor<D: AsRef<[u8]>>(fst: &Fst<D>, key: &[u8]) -> Option<Floor> {
+    let mut node = fst.root();
+    let mut output = Output::zero();
+    let mut best = None;
+
+    for (depth, &byte) in key.iter().enumerate() {
+        if let Some(i) = max_lesser_transition(&node, byte) {
+            let t = node.transition(i);
+            best = Some(Floor::Branch {
+                len: depth,
+                input: t.inp,
+                addr: t.addr,
+                output: output.cat(t.out),
+            });
+        } else if node.is_final() {
+            best = Some(Floor::Prefix {
+                len: depth,
+                output: output.cat(node.final_output()),
+            });
+        }
+
+        let Some(i) = node.find_input(byte) else {
+            return best;
+        };
+        let t = node.transition(i);
+        output = output.cat(t.out);
+        node = fst.node(t.addr);
+    }
+
+    if node.is_final() {
+        return Some(Floor::Prefix {
+            len: key.len(),
+            output: output.cat(node.final_output()),
+        });
     }
-    let mut lo: usize = 0;
-    let mut hi: usize = n;
+    best
+}
+
+fn max_lesser_transition(node: &Node, byte: u8) -> Option<usize> {
+    let (mut lo, mut hi) = (0, node.len());
     while lo < hi {
         let mid = lo + (hi - lo) / 2;
         if node.transition(mid).inp < byte {
@@ -428,211 +525,97 @@ fn find_max_lesser(node: &Node, byte: u8) -> Option<usize> {
     lo.checked_sub(1)
 }
 
-fn rightmost_to_leaf<D: AsRef<[u8]>>(
+// Follows the last transition from `addr` down to a leaf, passing each input
+// byte to `on_byte`, and returns the accumulated value.
+fn rightmost<D: AsRef<[u8]>>(
     fst: &Fst<D>,
     addr: CompiledAddr,
-    output: Output,
-) -> (Vec<u8>, u64) {
-    let mut node = fst.node(addr);
-    let mut out = output;
-    let mut suffix = Vec::new();
-
-    while !node.is_empty() {
-        let last = node.len() - 1;
-        let t = node.transition(last);
-        suffix.push(t.inp);
-        out = out.cat(t.out);
-        node = fst.node(t.addr);
-    }
-
-    (suffix, out.cat(node.final_output()).value())
-}
-
-fn rightmost_value<D: AsRef<[u8]>>(fst: &Fst<D>, addr: CompiledAddr, output: Output) -> u64 {
+    mut output: Output,
+    mut on_byte: impl FnMut(u8),
+) -> u64 {
     let mut node = fst.node(addr);
-    let mut out = output;
-
     while !node.is_empty() {
-        let last = node.len() - 1;
-        let t = node.transition(last);
-        out = out.cat(t.out);
+        let t = node.transition(node.len() - 1);
+        on_byte(t.inp);
+        output = output.cat(t.out);
         node = fst.node(t.addr);
     }
-
-    out.cat(node.final_output()).value()
+    output.cat(node.final_output()).value()
 }
 
-fn floor_lookup<D: AsRef<[u8]>>(fst: &Fst<D>, key: &[u8]) -> Option<(Vec<u8>, u64)> {
-    let root = fst.root();
-
-    if key.is_empty() {
-        return if root.is_final() {
-            Some((Vec::new(), root.final_output().value()))
-        } else {
-            None
-        };
-    }
-
-    let mut node = root;
-    let mut output = Output::zero();
-    let mut stack: Vec<Frame> = Vec::with_capacity(key.len());
-    let mut matched: usize = 0;
-
-    for &byte in key.iter() {
-        let lesser = find_max_lesser(&node, byte);
-
-        stack.push(Frame {
-            node_addr: node.addr(),
+fn floor_entry<D: AsRef<[u8]>>(fst: &Fst<D>, key: &[u8]) -> Option<(Vec<u8>, u64)> {
+    Some(match find_floor(fst, key)? {
+        Floor::Prefix { len, output } => (key[..len].to_vec(), output.value()),
+        Floor::Branch {
+            len,
+            input,
+            addr,
             output,
-            prefix_len: matched,
-            max_lesser_idx: lesser,
-            is_final: node.is_final(),
-            final_value: output.cat(node.final_output()).value(),
-        });
-
-        match node.find_input(byte) {
-            Some(idx) => {
-                let t = node.transition(idx);
-                output = output.cat(t.out);
-                node = fst.node(t.addr);
-                matched += 1;
-            }
-            None => break,
-        }
-    }
-
-    if matched == key.len() && node.is_final() {
-        return Some((key.to_vec(), output.cat(node.final_output()).value()));
-    }
-
-    while let Some(frame) = stack.pop() {
-        if let Some(j) = frame.max_lesser_idx {
-            let frame_node = fst.node(frame.node_addr);
-            let t = frame_node.transition(j);
-            let mut result = key[..frame.prefix_len].to_vec();
-            result.push(t.inp);
-            let branch_output = frame.output.cat(t.out);
-            let (suffix, val) = rightmost_to_leaf(fst, t.addr, branch_output);
-            result.extend(suffix);
-            return Some((result, val));
-        }
-
-        if frame.is_final {
-            return Some((key[..frame.prefix_len].to_vec(), frame.final_value));
+        } => {
+            let mut found = key[..len].to_vec();
+            found.push(input);
+            let value = rightmost(fst, addr, output, |byte| found.push(byte));
+            (found, value)
         }
-    }
-
-    None
+    })
 }
 
 fn floor_value<D: AsRef<[u8]>>(fst: &Fst<D>, key: &[u8]) -> Option<u64> {
-    let root = fst.root();
-
-    if key.is_empty() {
-        return if root.is_final() {
-            Some(root.final_output().value())
-        } else {
-            None
-        };
-    }
-
-    let mut node = root;
-    let mut output = Output::zero();
-    let mut stack: Vec<Frame> = Vec::with_capacity(key.len());
-    let mut matched: usize = 0;
-
-    for &byte in key.iter() {
-        let lesser = find_max_lesser(&node, byte);
-
-        stack.push(Frame {
-            node_addr: node.addr(),
-            output,
-            prefix_len: matched,
-            max_lesser_idx: lesser,
-            is_final: node.is_final(),
-            final_value: output.cat(node.final_output()).value(),
-        });
-
-        match node.find_input(byte) {
-            Some(idx) => {
-                let t = node.transition(idx);
-                output = output.cat(t.out);
-                node = fst.node(t.addr);
-                matched += 1;
-            }
-            None => break,
-        }
-    }
-
-    if matched == key.len() && node.is_final() {
-        return Some(output.cat(node.final_output()).value());
-    }
-
-    while let Some(frame) = stack.pop() {
-        if let Some(j) = frame.max_lesser_idx {
-            let frame_node = fst.node(frame.node_addr);
-            let t = frame_node.transition(j);
-            let branch_output = frame.output.cat(t.out);
-            return Some(rightmost_value(fst, t.addr, branch_output));
-        }
-
-        if frame.is_final {
-            return Some(frame.final_value);
-        }
-    }
-
-    None
+    Some(match find_floor(fst, key)? {
+        Floor::Prefix { output, .. } => output.value(),
+        Floor::Branch { addr, output, .. } => rightmost(fst, addr, output, |_| {}),
+    })
 }
 
 // ---------------------------------------------------------------------------
 // Init
 // ---------------------------------------------------------------------------
 
+// `method!` can't name a generic parameter from an enclosing function, so the
+// shared methods are registered through a macro instead.
+macro_rules! define_fst_methods {
+    ($class:expr, $type:ty) => {{
+        let class = $class;
+        class.define_singleton_method("new", function!(new::<$type>, 1))?;
+        class.define_singleton_method("from_path", function!(from_path::<$type>, 1))?;
+        class.define_singleton_method("from_path_mmap", function!(from_path_mmap::<$type>, 1))?;
+        class.define_method("contains?", method!(contains::<$type>, 1))?;
+        class.define_method("length", method!(len::<$type>, 0))?;
+        class.define_method("size", method!(len::<$type>, 0))?;
+        class.define_method("empty?", method!(is_empty::<$type>, 0))?;
+        class.define_method("to_bytes", method!(to_bytes::<$type>, 0))?;
+        class.define_method("save", method!(save::<$type>, 1))?;
+        class.define_method("each", method!(each::<$type>, 0))?;
+        class.define_method("range", method!(range::<$type>, 2))?;
+        class.define_method("starts_with", method!(starts_with::<$type>, 1))?;
+        class.define_method(
+            "search_levenshtein",
+            method!(search_levenshtein::<$type>, 2),
+        )?;
+        class
+    }};
+}
+
 #[magnus::init]
 fn init(ruby: &Ruby) -> Result<(), Error> {
     let module = ruby.define_module("RubyFst")?;
+    module.define_error("Error", ruby.exception_runtime_error())?;
 
-    let map_class = module.define_class("Map", ruby.class_object())?;
-    map_class.define_singleton_method("new", function!(FstMap::new, 1))?;
-    map_class.define_singleton_method("from_path", function!(FstMap::from_path, 1))?;
-    map_class.define_singleton_method("from_path_mmap", function!(FstMap::from_path_mmap, 1))?;
-    map_class.define_method("get", method!(FstMap::get, 1))?;
-    map_class.define_method("[]", method!(FstMap::get, 1))?;
-    map_class.define_method("contains?", method!(FstMap::contains, 1))?;
-    map_class.define_method("length", method!(FstMap::len, 0))?;
-    map_class.define_method("size", method!(FstMap::len, 0))?;
-    map_class.define_method("empty?", method!(FstMap::is_empty, 0))?;
-    map_class.define_method("to_bytes", method!(FstMap::to_bytes, 0))?;
-    map_class.define_method("save", method!(FstMap::save, 1))?;
-    map_class.define_method("get_le", method!(FstMap::get_le, 1))?;
-    map_class.define_method("get_le_value", method!(FstMap::get_le_value, 1))?;
-    map_class.define_method("get_ge", method!(FstMap::get_ge, 1))?;
-    map_class.define_method("get_ge_value", method!(FstMap::get_ge_value, 1))?;
-    map_class.define_method("each", method!(FstMap::each, 0))?;
-    map_class.define_method("range", method!(FstMap::range, 2))?;
-    map_class.define_method("starts_with", method!(FstMap::starts_with, 1))?;
-    map_class.define_method("search_levenshtein", method!(FstMap::search_levenshtein, 2))?;
+    let map_class = define_fst_methods!(module.define_class("Map", ruby.class_object())?, FstMap);
+    map_class.define_method("get", method!(get, 1))?;
+    map_class.define_method("[]", method!(get, 1))?;
+    map_class.define_method("get_le", method!(get_le, 1))?;
+    map_class.define_method("get_le_value", method!(get_le_value, 1))?;
+    map_class.define_method("get_ge", method!(get_ge, 1))?;
+    map_class.define_method("get_ge_value", method!(get_ge_value, 1))?;
+
+    define_fst_methods!(module.define_class("Set", ruby.class_object())?, FstSet);
 
     let map_builder = module.define_class("MapBuilder", ruby.class_object())?;
     map_builder.define_singleton_method("new", function!(FstMapBuilder::new, 0))?;
     map_builder.define_method("insert", method!(FstMapBuilder::insert, 2))?;
     map_builder.define_method("finish", method!(FstMapBuilder::finish, 0))?;
 
-    let set_class = module.define_class("Set", ruby.class_object())?;
-    set_class.define_singleton_method("new", function!(FstSet::new, 1))?;
-    set_class.define_singleton_method("from_path", function!(FstSet::from_path, 1))?;
-    set_class.define_singleton_method("from_path_mmap", function!(FstSet::from_path_mmap, 1))?;
-    set_class.define_method("contains?", method!(FstSet::contains, 1))?;
-    set_class.define_method("length", method!(FstSet::len, 0))?;
-    set_class.define_method("size", method!(FstSet::len, 0))?;
-    set_class.define_method("empty?", method!(FstSet::is_empty, 0))?;
-    set_class.define_method("to_bytes", method!(FstSet::to_bytes, 0))?;
-    set_class.define_method("save", method!(FstSet::save, 1))?;
-    set_class.define_method("each", method!(FstSet::each, 0))?;
-    set_class.define_method("range", method!(FstSet::range, 2))?;
-    set_class.define_method("starts_with", method!(FstSet::starts_with, 1))?;
-    set_class.define_method("search_levenshtein", method!(FstSet::search_levenshtein, 2))?;
-
     let set_builder = module.define_class("SetBuilder", ruby.class_object())?;
     set_builder.define_singleton_method("new", function!(FstSetBuilder::new, 0))?;
     set_builder.define_method("insert", method!(FstSetBuilder::insert, 1))?;

data/lib/ruby_fst/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module RubyFst
-  VERSION = '0.2.2'
+  VERSION = '0.3.0'
 end

data/lib/ruby_fst.rb

@@ -9,6 +9,12 @@ end
 
 module RubyFst
   module RangeQuery
+    def each(&block)
+      return enum_for(:each) { size } unless block
+
+      _each(&block)
+    end
+
     def range(ge: nil, le: nil, &block)
       return enum_for(:range, ge:, le:) unless block
 
@@ -20,14 +26,22 @@ module RubyFst
 
       _starts_with(prefix, &block)
     end
+
+    def search_levenshtein(query, distance, &block)
+      return enum_for(:search_levenshtein, query, distance) unless block
+
+      _search_levenshtein(query, distance, &block)
+    end
   end
 
   class Map
     include Enumerable
 
+    alias _each each
     alias _range range
     alias _starts_with starts_with
-    private :_range, :_starts_with
+    alias _search_levenshtein search_levenshtein
+    private :_each, :_range, :_starts_with, :_search_levenshtein
 
     prepend RangeQuery
   end
@@ -35,9 +49,11 @@ module RubyFst
   class Set
     include Enumerable
 
+    alias _each each
     alias _range range
     alias _starts_with starts_with
-    private :_range, :_starts_with
+    alias _search_levenshtein search_levenshtein
+    private :_each, :_range, :_starts_with, :_search_levenshtein
 
     prepend RangeQuery
   end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ruby-fst
 version: !ruby/object:Gem::Version
-  version: 0.2.2
+  version: 0.3.0
 platform: ruby
 authors:
 - Denis Sablic