Structured memory deallocation using blocks

ext/fiddle/pointer.c

@@ -200,12 +200,21 @@ rb_fiddle_ptr_initialize(int argc, VALUE argv[], VALUE self)
     return Qnil;
 }
 
+static VALUE
+rb_fiddle_ptr_call_free(VALUE self);
+
 /*
  * call-seq:
  *    Fiddle::Pointer.malloc(size, freefunc = nil)  => fiddle pointer instance
+ *    Fiddle::Pointer.malloc(size, freefunc) { |pointer| ... } => ...
  *
  * == Examples
  *
+ *    # Automatically freeing the pointer when the block is exited - recommended
+ *    Fiddle::Pointer.malloc(size, Fiddle::RUBY_FREE) do |pointer|
+ *      ...
+ *    end
+ *
  *    # Manually freeing but relying on the garbage collector otherwise
  *    pointer = Fiddle::Pointer.malloc(size, Fiddle::RUBY_FREE)
  *    ...
@@ -228,15 +237,16 @@ rb_fiddle_ptr_initialize(int argc, VALUE argv[], VALUE self)
  *    ...
  *
  * Allocate +size+ bytes of memory and associate it with an optional
- * +freefunc+ that will be called when the pointer is garbage collected,
- * if not already manually called via *call_free*.
- *
- * +freefunc+ must be an address pointing to a function or an instance of
- * +Fiddle::Function+. Using +freefunc+ without using +call_free+ may lead to
- * unlimited memory being allocated before any is freed as the native memory
- * the pointer references does not contribute to triggering the Ruby garbage
- * collector. Consider manually freeing the memory as illustrated above. You
- * can combine the techniques as +freefunc+ will not be called twice.
+ * +freefunc+.
+ *
+ * If a block is supplied, the pointer will be yielded to the block instead of
+ * being returned, and the return value of the block will be returned. A
+ * +freefunc+ must be supplied if a block is.
+ *
+ * If a +freefunc+ is supplied it will be called once, when the pointer is
+ * garbage collected or when the block is left if a block is supplied or
+ * when the user calls +call_free+, whichever happens first. +freefunc+ must be
+ * an address pointing to a function or an instance of +Fiddle::Function+.
  */
 static VALUE
 rb_fiddle_ptr_s_malloc(int argc, VALUE argv[], VALUE klass)
@@ -261,7 +271,14 @@ rb_fiddle_ptr_s_malloc(int argc, VALUE argv[], VALUE klass)
     obj = rb_fiddle_ptr_malloc(s,f);
     if (wrap) RPTR_DATA(obj)->wrap[1] = wrap;
 
-    return obj;
+    if (rb_block_given_p()) {
+        if (!f) {
+            rb_raise(rb_eArgError, "a free function must be supplied to Fiddle::Pointer.malloc when it is called with a block");
+        }
+        return rb_ensure(rb_yield, obj, rb_fiddle_ptr_call_free, obj);
+    } else {
+        return obj;
+    }
 }
 
 /*

test/fiddle/test_c_struct_entry.rb

@@ -57,21 +57,24 @@ def test_set_ctypes
     def test_aref_pointer_array
       team = CStructEntity.malloc([[TYPE_VOIDP, 2]], Fiddle::RUBY_FREE)
       team.assign_names(["names"])
-      alice = Fiddle::Pointer.malloc(6, Fiddle::RUBY_FREE)
-      alice[0, 6] = "Alice\0"
-      bob = Fiddle::Pointer.malloc(4, Fiddle::RUBY_FREE)
-      bob[0, 4] = "Bob\0"
-      team["names"] = [alice, bob]
-      assert_equal(["Alice", "Bob"], team["names"].map(&:to_s))
+      Fiddle::Pointer.malloc(6, Fiddle::RUBY_FREE) do |alice|
+        alice[0, 6] = "Alice\0"
+        Fiddle::Pointer.malloc(4, Fiddle::RUBY_FREE) do |bob|
+          bob[0, 4] = "Bob\0"
+          team["names"] = [alice, bob]
+          assert_equal(["Alice", "Bob"], team["names"].map(&:to_s))
+        end
+      end
     end
 
     def test_aref_pointer
       user = CStructEntity.malloc([TYPE_VOIDP], Fiddle::RUBY_FREE)
       user.assign_names(["name"])
-      alice = Fiddle::Pointer.malloc(6, Fiddle::RUBY_FREE)
-      alice[0, 6] = "Alice\0"
-      user["name"] = alice
-      assert_equal("Alice", user["name"].to_s)
+      Fiddle::Pointer.malloc(6, Fiddle::RUBY_FREE) do |alice|
+        alice[0, 6] = "Alice\0"
+        user["name"] = alice
+        assert_equal("Alice", user["name"].to_s)
+      end
     end
   end
 end if defined?(Fiddle)

test/fiddle/test_pointer.rb

@@ -32,6 +32,24 @@ def test_malloc_free_func
       assert_equal free.to_i, ptr.free.to_i
     end
 
+    def test_malloc_block
+      escaped_ptr = nil
+      returned = Pointer.malloc(10, Fiddle::RUBY_FREE) do |ptr|
+        assert_equal 10, ptr.size
+        assert_equal Fiddle::RUBY_FREE, ptr.free.to_i
+        escaped_ptr = ptr
+        :returned
+      end
+      assert_equal :returned, returned
+      assert escaped_ptr.freed?
+    end
+
+    def test_malloc_block_no_free
+      assert_raise ArgumentError do 
+        Pointer.malloc(10) { |ptr| }
+      end
+    end
+
     def test_to_str
       str = Marshal.load(Marshal.dump("hello world"))
       ptr = Pointer[str]
@@ -84,17 +102,18 @@ def test_to_ptr_string
     end
 
     def test_to_ptr_io
-      buf = Pointer.malloc(10, Fiddle::RUBY_FREE)
-      File.open(__FILE__, 'r') do |f|
-        ptr = Pointer.to_ptr f
-        fread = Function.new(@libc['fread'],
-                             [TYPE_VOIDP, TYPE_INT, TYPE_INT, TYPE_VOIDP],
-                             TYPE_INT)
-        fread.call(buf.to_i, Fiddle::SIZEOF_CHAR, buf.size - 1, ptr.to_i)
-      end
-
-      File.open(__FILE__, 'r') do |f|
-        assert_equal f.read(9), buf.to_s
+      Pointer.malloc(10, Fiddle::RUBY_FREE) do |buf|
+        File.open(__FILE__, 'r') do |f|
+          ptr = Pointer.to_ptr f
+          fread = Function.new(@libc['fread'],
+                              [TYPE_VOIDP, TYPE_INT, TYPE_INT, TYPE_VOIDP],
+                              TYPE_INT)
+          fread.call(buf.to_i, Fiddle::SIZEOF_CHAR, buf.size - 1, ptr.to_i)
+        end
+
+        File.open(__FILE__, 'r') do |f|
+          assert_equal f.read(9), buf.to_s
+        end
       end
     end
 
@@ -177,7 +196,7 @@ def test_free_with_func
       assert ptr.freed?
       ptr.call_free                 # you can safely run it again
       assert ptr.freed?
-      GC.start                  # you can safely run the GC routine
+      GC.start                      # you can safely run the GC routine
       assert ptr.freed?
     end
 
@@ -203,21 +222,15 @@ def test_null?
     end
 
     def test_size
-      ptr = Pointer.malloc(4)
-      begin
+      Pointer.malloc(4, Fiddle::RUBY_FREE) do |ptr|
         assert_equal 4, ptr.size
-      ensure
-        Fiddle.free ptr
       end
     end
 
     def test_size=
-      ptr = Pointer.malloc(4)
-      begin
+      Pointer.malloc(4, Fiddle::RUBY_FREE) do |ptr|
         ptr.size = 10
         assert_equal 10, ptr.size
-      ensure
-        Fiddle.free ptr
       end
     end