Primitive Low Level Optimization in Ruby

Prologue

I like Ruby. I consider it a beautiful, ergonomic, and eloquent programming language. It obviously can’t compete with faster compiled languages in terms of raw performance, but what I enjoy is that I don’t have to rewrite whole Ruby projects in Go or Rust. Instead, I can identify the hot fragments of the codebase, port those to fast compiled modules, and stay in the pleasant, comfortable Ruby ecosystem for everything else.

In the following text I’ll document a primitive example of writing a low-level C extension for a Ruby module. I will use edit distance (Levenshtein distance) — a measure of how many single-character edits it takes to turn one string into another — as the example.

Step 1. Create a gem

Let’s create a typical Ruby package to hold everything together.

❯ bundle gem --test=rspec --no-ci --no-linter --mit --no-changelog primitive

Creating gem 'primitive'...
rspec is already configured, ignoring --test flag.
MIT License enabled in config
Initializing git repo in /home/vlad/Code/primitive
create  primitive/Gemfile
create  primitive/lib/primitive.rb
create  primitive/lib/primitive/version.rb
create  primitive/sig/primitive.rbs
create  primitive/primitive.gemspec
create  primitive/Rakefile
create  primitive/README.md
create  primitive/bin/console
create  primitive/bin/setup
create  primitive/.gitignore
create  primitive/.rspec
create  primitive/spec/spec_helper.rb
create  primitive/spec/primitive_spec.rb
create  primitive/LICENSE.txt

Gem 'primitive' was successfully created.

Step 2. Add code that does something useful

Specs

It’s often a good idea to start with test examples describing the method’s desired behaviour.

RSpec.describe Primitive do
  describe "#edit_distance" do
    let(:examples) do
      [['', 'abc', 3],
       ['abc', '', 3],
       ['a', 'abcdefgh', 7],
       ['kitten', 'sitting', 3],
       ['saturday', 'sunday', 3],
       ['tuesday', 'thursday', 2],
       ['hematology', 'hepatology', 1],
       # Counts characters rather than bytes: dropping '語' is one edit, not three.
       ['日本語', '日本', 1],
       ['café', 'cafe', 1]]
    end

    specify do
      examples.each do |s1, s2, expected_distance|
        expect(described_class.edit_distance(s1, s2)).to eq(expected_distance)
      end
    end
  end
end

Implementation

The following is a standard dynamic-programming algorithm relying on a table that keeps track of previously computed distances for the strings’ prefixes.

module Primitive
  def edit_distance(s1, s2)
    table = Array.new(s1.size + 1) { |i| Array.new(s2.size + 1) { |j| i + j } }
    1.upto(s1.size) do |i|
      1.upto(s2.size) do |j|
        deletion     = table[i - 1][j] + 1
        insertion    = table[i][j - 1] + 1
        substitution = table[i - 1][j - 1] + (s1[i - 1] == s2[j - 1] ? 0 : 1)

        table[i][j]  = [deletion, insertion, substitution].min
      end
    end
    table[s1.size][s2.size]
  end

  module_function :edit_distance
end

Step 3. Measure

Now that we have test-based evidence of the method working correctly, let’s analyze how fast it is.

Benchmark script

First, add benchmark to Gemfile

gem "benchmark"

And run bundle install

Next, save this as bin/benchmark.

#!/usr/bin/env ruby

require "bundler/setup"
require "primitive"
require "benchmark"

NUCLEOTIDES = %w[A C G T].freeze

# Random nucleotide sequences: DNA sequence alignment is a classic real-world use case for edit distance.
len      = 120
s1       = Array.new(len) { NUCLEOTIDES.sample }.join
s2       = Array.new(len) { NUCLEOTIDES.sample }.join
distance = Primitive.edit_distance(s1, s2)

puts
puts "Sequence 1: #{s1}"
puts "Sequence 2: #{s2}"
puts "Distance  : #{distance}"
puts

rounds = 1000

Benchmark.bmbm(6) do |x|
  x.report("Primitive#edit_distance") { rounds.times { Primitive.edit_distance(s1, s2) } }
end

Results on MacBook Pro

❯ bin/benchmark

Sequence 1: ACATCCGCGATATTGGTTAGCTTGTCAAAATTTCTCGTCAAGCGAACCCATTATATCCGTACTGAGTATTAAGGAATCTAACCAGGAAGGATCAGAACGGATATAAATTCTCCCTGAATG
Sequence 2: GAATCGAGCCAACATGATCAAGAATTTCGCAGAGAGCATTAGGGAACGCGAGTTGCTTACCGGAAAGAGTGCCTAGAGCATGAGCTTCCAATGGAAATGGTAACTACGGACACTCTCTCC
Distance  : 66

Rehearsal -----------------------------------------------------------
Primitive#edit_distance   3.046411   0.022644   3.069055 (  3.072687)
-------------------------------------------------- total: 3.069055sec

                              user     system      total        real
Primitive#edit_distance   3.056331   0.024423   3.080754 (  3.082315)

That’s roughly 3ms per call in pure Ruby — the number the C extension will need to beat.

Step 4. Adding C

In this part I will manually add the minimum necessary C code wired into the Ruby module. It was possible to automatically generate the C extension files with bundle gem --ext=c but that wouldn’t be educational. Let’s get to it:

❯ touch ext/primitive/main.c     # The main extension code in C
❯ touch ext/primitive/extconf.rb # C extension configuration for the gem to package it all together

And include the configuration into primitive.gemspec

Gem::Specification.new do |spec|
  # ...
  spec.require_paths = ["lib"]
  spec.extensions    = ["ext/primitive/extconf.rb"] # << HERE
  # ...
end

Put the following into the extconf.rb:

require "mkmf"

create_makefile("primitive/main")

In order to compare the low-level implementation against the pure-Ruby baseline I’m going to define a parallel #edit_distance_fast method.

Recall that #edit_distance counts characters rather than bytes, so it stays correct on multi-byte UTF-8 input like '日本語'. C has no notion of a Ruby-style variable-width encoding, so to index codepoints directly by pointer arithmetic on the C side, both strings first need to be re-encoded into something fixed-width — UTF_32LE gives exactly one 4-byte int per codepoint.

# lib/primitive.rb

module Primitive
  def edit_distance(s1, s2) # unchanged, shown in Step 2
    # ...
  end

  def edit_distance_fast(s1, s2) # optimized implementation wrapper
    # Prep the arguments to be 32-bit int sequences
    s1_32le = s1.encode(Encoding::UTF_32LE)
    s2_32le = s2.encode(Encoding::UTF_32LE)

    edit_distance_internal(s1_32le, s2_32le) # Call the method from C
  end

  module_function :edit_distance, :edit_distance_fast
end

Then let’s add the extension’s starter code. It’s not finished but it will allow to wire everything together.

/* ext/primitive/main.c */

#include "ruby.h"

static VALUE rb_edit_distance_internal(VALUE self, VALUE str1, VALUE str2)
{
    /* Full implementation will be added here */

    return UINT2NUM(0);
}

void Init_main(void)
{
    /*
     * The following is basically
     * module Primitive
     *   def edit_distance_internal(s1, s2)
     *   end
     * end
     */
    VALUE primitive_module = rb_define_module("Primitive");
    rb_define_module_function(primitive_module, "edit_distance_internal", rb_edit_distance_internal, 2);
}

Step 5. Compile it

Nothing above is reachable from Ruby yet. rake-compiler builds the shared object and require loads it.

# Gemfile

gem "rake-compiler"

# Rakefile

require "rake/extensiontask"

Rake::ExtensionTask.new("main") do |ext|
  ext.lib_dir = "lib/primitive" # where the compiled object lands
  ext.ext_dir = "ext/primitive" # where the sources live
end

task default: %i[compile spec] # never run specs against a stale build

That "main" has to agree with create_makefile("primitive/main") and with Init_main. Ruby derives the init function’s name from the file it loads, so main.bundle is expected to export Init_main. Get it wrong and the extension compiles happily, then refuses to load.

The compiled object is a build artifact, so it stays out of the repository:

# .gitignore

/lib/primitive/main.bundle

And the module has to load it:

# lib/primitive.rb

require_relative "primitive/version"
require_relative "primitive/main" # << HERE

Now let’s compile:

❯ bundle exec rake compile

creating Makefile
compiling ../../../../ext/primitive/main.c
linking shared-object primitive/main.bundle
/usr/bin/install -c -m 0755 main.bundle ../../../../lib/primitive

And check that the stub answers:

irb> Primitive.edit_distance_fast("kitten", "sitting")
=> 0

The answer is wrong, but the call made it into C and back.

Two small things before the real implementation. edit_distance_internal expects UTF-32LE and returns nonsense if given anything else, so I don’t want it public:

private_class_method :edit_distance_internal

And the specs should hold both implementations to the same examples:

RSpec.describe Primitive do
  include described_class

  describe "#edit_distance" do
    # ... examples as before

    specify do
      examples.each do |s1, s2, expected_distance|
        expect(edit_distance(s1, s2)).to eq(expected_distance)
        expect(edit_distance_fast(s1, s2)).to eq(expected_distance)
      end
    end
  end
end

Step 6. The actual implementation

The loop below is a direct transcription of the Ruby one. Everything around it is about crossing the boundary, so that’s what the comments cover.

/* ext/primitive/main.c */

#include "ruby.h"

typedef struct Matrix {
    long nrows;
    long ncols;
    uint32_t *values;
} Matrix;

static inline uint32_t matrix_get(const Matrix *m, long i, long j)
{
    return m->values[i * m->ncols + j];
}

static inline void matrix_set(Matrix *m, long i, long j, uint32_t value)
{
    m->values[i * m->ncols + j] = value;
}

static VALUE rb_edit_distance_internal(VALUE self, VALUE str1, VALUE str2)
{
    /* A VALUE says nothing about what it points at, so unchecked arguments are a segfault, not a TypeError. */
    Check_Type(str1, T_STRING);
    Check_Type(str2, T_STRING);

    /* RSTRING_LEN counts bytes. UTF-32LE spends exactly four per character and carries no BOM. */
    long len1 = RSTRING_LEN(str1) / sizeof(uint32_t);
    long len2 = RSTRING_LEN(str2) / sizeof(uint32_t);

    Matrix table = { .nrows = len1 + 1, .ncols = len2 + 1 };

    /* Like ALLOC_N, but small requests come off the stack and large ones are owned by a hidden Ruby
     * object, so the buffer is reclaimed even if something between here and ALLOCV_END raises. */
    VALUE allocation_handle;
    table.values = ALLOCV_N(uint32_t, allocation_handle, table.nrows * table.ncols);

    /*
     * Keep below ALLOCV_N, the last call here that can trigger a GC.
     * Nothing after it allocates, so these interior pointers stay valid to the end.
     */
    const uint32_t *chars1 = (const uint32_t *)RSTRING_PTR(str1);
    const uint32_t *chars2 = (const uint32_t *)RSTRING_PTR(str2);

    for (long i = 0; i <= len1; i++) matrix_set(&table, i, 0, (uint32_t) i);
    for (long j = 0; j <= len2; j++) matrix_set(&table, 0, j, (uint32_t) j);

    for (long i = 1; i <= len1; i++) {
        for (long j = 1; j <= len2; j++) {
            uint32_t deletion     = matrix_get(&table, i - 1, j) + 1;
            uint32_t insertion    = matrix_get(&table, i, j - 1) + 1;
            uint32_t substitution = matrix_get(&table, i - 1, j - 1) + (chars1[i - 1] == chars2[j - 1] ? 0 : 1);
            uint32_t min          = deletion < insertion ? deletion : insertion;

            matrix_set(&table, i, j, min < substitution ? min : substitution);
        }
    }

    uint32_t distance = matrix_get(&table, len1, len2);

    ALLOCV_END(allocation_handle);

    /* Must stay after the last read of chars1/chars2 to keep the strings live and unmoved until then.
     * Once the only surviving reference is an anonymous uint32_t*, the compiler is free to drop the
     * VALUE itself, and the GC stops seeing it. */
    RB_GC_GUARD(str1);
    RB_GC_GUARD(str2);

    /* A raw number would be read back as a VALUE, pointing at whatever lives at that address. */
    return UINT2NUM(distance);
}

Reading four bytes as a uint32_t gives the codepoint on a little-endian machine. Worth noticing that the algorithm doesn’t depend on that: both strings are encoded the same way and it only ever compares them to each other, so byte-swapped codepoints compare equal exactly when the real ones do.

Step 7. Measure again

Uncomment the second report in bin/benchmark:

ruby, c = Benchmark.bmbm(6) do |x|
  x.report("Primitive#edit_distance")      { rounds.times { Primitive.edit_distance(s1, s2) } }
  x.report("Primitive#edit_distance_fast") { rounds.times { Primitive.edit_distance_fast(s1, s2) } }
end

difference = (ruby.real / c.real).round(1)

puts
puts "#{rounds} rounds, #{len} x #{len} characters: C is #{difference}x faster"

❯ bundle exec rake compile && bin/benchmark

Distance  : 71

Rehearsal ----------------------------------------------------------------
Primitive#edit_distance        2.967991   0.018188   2.986179 (  2.986448)
Primitive#edit_distance_fast   0.013844   0.000652   0.014496 (  0.014508)
------------------------------------------------------- total: 3.000675sec

                                   user     system      total        real
Primitive#edit_distance        2.472894   0.012663   2.485557 (  2.485745)
Primitive#edit_distance_fast   0.013331   0.000009   0.013340 (  0.013340)

1000 rounds, 120 x 120 characters: C is 186.3x faster

13 microseconds per call, including the two #encode calls the wrapper makes every time. The Ruby line drifted down from Step 3 because bmbm rehearses both reports before measuring either, and the C rehearsal leaves the GC heap bigger than it found it.

Step 8. Reduce memory complexity

186x is where it gets tempting to stop, but the algorithm is still carrying a full (n + 1) × (m + 1) table. Here are the inner loop’s three reads with the arithmetic stripped off:

matrix_get(&table, i - 1, j);
matrix_get(&table, i,     j - 1);
matrix_get(&table, i - 1, j - 1);

Every one lands on the current row or the one directly above it. Rows 0 .. i - 2 are never touched again, so two rows are enough if I swap them at the end of each iteration. The distance is symmetric too, which means the shorter string can go on the inner axis to make those two rows as narrow as possible.

typedef struct Vector {
    long size;
    uint32_t *values;
} Vector;

static inline uint32_t vector_get(const Vector *v, long i)
{
    return v->values[i];
}

static inline void vector_set(Vector *v, long i, uint32_t value)
{
    v->values[i] = value;
}

static inline void vector_swap(Vector *a, Vector *b)
{
    Vector swap = *a; *a = *b; *b = swap;
}

static VALUE rb_edit_distance_internal(VALUE self, VALUE str1, VALUE str2)
{
    Check_Type(str1, T_STRING);
    Check_Type(str2, T_STRING);

    long len1 = RSTRING_LEN(str1) / sizeof(uint32_t);
    long len2 = RSTRING_LEN(str2) / sizeof(uint32_t);

    /* Identify the shorter string to narrow the rows. */
    if (len2 > len1) {
        VALUE swap_str = str1; str1 = str2; str2 = swap_str;
        long  swap_len = len1; len1 = len2; len2 = swap_len;
    }

    /* Two rows are enough for the lookback */
    long ncols = len2 + 1;
    VALUE allocation_handle;
    uint32_t *values = ALLOCV_N(uint32_t, allocation_handle, 2 * ncols);

    Vector previous = { .size = ncols, .values = values };
    Vector current  = { .size = ncols, .values = values + ncols };

    const uint32_t *chars1 = (const uint32_t *)RSTRING_PTR(str1);
    const uint32_t *chars2 = (const uint32_t *)RSTRING_PTR(str2);

    for (long j = 0; j < previous.size; j++) vector_set(&previous, j, (uint32_t) j);

    for (long i = 1; i <= len1; i++) {
        vector_set(&current, 0, (uint32_t) i);

        for (long j = 1; j <= len2; j++) {
            uint32_t deletion     = vector_get(&previous, j) + 1;
            uint32_t insertion    = vector_get(&current, j - 1) + 1;
            uint32_t substitution = vector_get(&previous, j - 1) + (chars1[i - 1] == chars2[j - 1] ? 0 : 1);
            uint32_t min          = deletion < insertion ? deletion : insertion;

            vector_set(&current, j, min < substitution ? min : substitution);
        }

        vector_swap(&previous, &current);
    }

    uint32_t distance = vector_get(&previous, len2);

    ALLOCV_END(allocation_handle);

    RB_GC_GUARD(str1);
    RB_GC_GUARD(str2);

    return UINT2NUM(distance);
}

For the 120-character sequences that takes the scratch buffer from 121 × 121 × 4 bytes, about 57KB, down to

Which crosses a threshold worth knowing about: ALLOCV_N uses the stack below RUBY_ALLOCV_LIMIT, or 1024 bytes, so for inputs up to 126 characters the allocation disappears entirely. At 4000 × 4000 it’s the difference between 61MB and 32KB.

The same applies to Ruby

The Ruby version is the control group, so it gets the same treatment. Otherwise I stop measuring C against Ruby and start measuring a good algorithm against one I didn’t bother to improve.

def edit_distance(s1, s2)
  chars1 = s1.chars
  chars2 = s2.chars

  # The distance is symmetric, so keep the shorter string on the inner axis to narrow the rows.
  chars1, chars2 = chars2, chars1 if chars2.size > chars1.size

  # Row 0 of the table: turning an empty prefix of chars1 into chars2[0, j] takes j insertions.
  previous = Array.new(chars2.size + 1) { |j| j }
  # Every cell will be set in the outer loop
  current  = Array.new(chars2.size + 1)

  1.upto(chars1.size) do |i|
    current[0] = i

    1.upto(chars2.size) do |j|
      deletion     = previous[j] + 1
      insertion    = current[j - 1] + 1
      substitution = previous[j - 1] + (chars1[i - 1] == chars2[j - 1] ? 0 : 1)
      current[j]   = [deletion, insertion, substitution].min
    end

    # The one just filled in becomes the row above, and the row it displaces is dead
    current, previous = previous, current
  end

  previous[chars2.size]
end

That’s two changes, and the #chars calls are the bigger one. s1[i - 1] allocates a fresh one-character String on every access, and the inner loop reaches into both strings 120 × 120 times — roughly 29,000 throwaway objects per call. Measured separately on the same inputs:

                 user     system      total        real
naive        2.356477   0.013009   2.369486 (  2.369611)
chars        1.503420   0.003926   1.507346 (  1.507438)
tworow       1.150132   0.003041   1.153173 (  1.153278)

Step 9. Measure one last time

❯ bundle exec rake compile && bin/benchmark

Distance  : 66

Rehearsal ----------------------------------------------------------------
Primitive#edit_distance        1.139158   0.002548   1.141706 (  1.141790)
Primitive#edit_distance_fast   0.014832   0.001193   0.016025 (  0.036078)
------------------------------------------------------- total: 1.157731sec

                                   user     system      total        real
Primitive#edit_distance        1.107029   0.002668   1.109697 (  1.109771)
Primitive#edit_distance_fast   0.013272   0.000008   0.013280 (  0.013284)

1000 rounds, 120 x 120 characters: C is 83.5x faster

The headline went down, from 186x to 83x, and that’s the honest number. Nothing got slower — the Ruby baseline got twice as fast, because most of what the first comparison measured was 29,000 string allocations per call rather than the language boundary.

Epilogue

83x is real and it was cheap, but it isn’t free.

There’s no bounds checking in that inner loop either. An off-by-one in Ruby raises NoMethodError on nil; here it corrupts memory, and instead of a backtrace I get a core dump.

The extension also holds the GVL for its whole duration, so it’s faster but not parallel. rb_thread_call_without_gvl can release it, with its own rules about what I may touch while the lock is down.

And crossing the boundary costs a fixed amount per call — two #encode calls plus the dispatch. Against 13 microseconds that’s invisible.

Which is what I wanted out of this: not rewriting the project, just moving a fragment that deserved it, after making sure it was already a decent loop in Ruby.

The full source is at github.com/vlebedeff/primitiverb.