How to prevent problems with `return` from block when using Ruby `yield`

Question

As every Ruby programmer eventually discovers, calling blocks or procs that contain return statements can be dangerous as this might exit your current context:

def some_method(&_block)
   puts 1
   yield
   # The following line will never be executed in this example
   # as the yield is actually a `yield-and-return`.
   puts 3
end

def test
  some_method do
    puts 2
    return
  end
end

test

# This prints "1\n2\n" instead of "1\n2\n3\n"    

In cases you want to be absolutely sure some of your code runs after you called a block or proc, you can use a begin ... ensure construct. But since ensure is also called if there is an exception during yield, it requires a little more work.

I've created a tiny module that deals with this problem in two different ways:

1. Using safe_yield, it is detected whether the yielded block or proc actually returns using the return keyword. If so, it raises an exception.

   unknown_block = proc do
     return
   end 
   
   ReturnSafeYield.safe_yield(unknown_block)
   # => Raises a UnexpectedReturnException exception
   

2. Using call_then_yield, you can call a block and then make sure that a second block is executed, even if the first block contains a return statement.

   unknown_block = proc do
     return
   end
   ReturnSafeYield.call_then_yield(unknown_block) do
     # => This line is called even though the above block contains a `return`.
   end
   

I'm considering to create a quick Gem out of this, or is there any built-in solution to prevent quick return from the nested block which I missed?

Answer 1

There is a built-in solution to detect whether a block contains a return statement.

You can use RubyVM::InstructionSequence.disasm to disassemble the block passed in by the user, then search it for throw 1, which represents a return statement.

Here's a sample implementation:

def safe_yield(&block)
  if RubyVM::InstructionSequence.disasm(block) =~ /^\d+ throw +1$/
    raise LocalJumpError
  end

  block.call
end

Here's how you might incorporate it into your library:

def library_method(&block)
  safe_yield(&block)
  puts "library_method succeeded"
rescue LocalJumpError
  puts "library_method encountered illegal return but resumed execution"
end

And here's the user experience for a well-behaved and a misbehaving user:

def nice_user_method
  library_method { 1 + 1 }
end

nice_user_method
# library_method succeeded

def naughty_user_method
  library_method { return false if rand > 0.5 }
end

naughty_user_method
# library_method encountered illegal return but resumed execution

Commentary:

Using raise LocalJumpError/rescue LocalJumpError gets around the issues you encountered when using a blanket ensure.

I chose LocalJumpError because it seems relevant, and because (I think!) there is no possible Ruby code that would result in LocalJumpError being raised "naturally" in this context. If that turns out to be false, you can easily substitute your own new exception class.