[ruby/fiddle] Make array access override compatible with base class (#25)

sinisterchipmunk · nobu · commit 4a835621ced0 · 2020-05-23T14:29:16.000+09:00
* Allow access to a struct's underlying memory with `struct[offset, length]`. * Make accessing a struct's underlying memory more convenient. * refactor memory access unit tests for improved clarity ruby/fiddle@c082c81bb5
diff --git a/ext/fiddle/lib/fiddle/struct.rb b/ext/fiddle/lib/fiddle/struct.rb
@@ -54,6 +54,8 @@ def create(klass, types, members)
           @entity = klass.entity_class.new(addr, types)
           @entity.assign_names(members)
         }
+        define_method(:[]) { |*args| @entity.send(:[], *args) }
+        define_method(:[]=) { |*args| @entity.send(:[]=, *args) }
         define_method(:to_ptr){ @entity }
         define_method(:to_i){ @entity.to_i }
         members.each{|name|
@@ -148,8 +150,21 @@ def set_ctypes(types)
       @size = PackInfo.align(offset, max_align)
     end
 
-    # Fetch struct member +name+
-    def [](name)
+    # Fetch struct member +name+ if only one argument is specified. If two
+    # arguments are specified, the first is an offset and the second is a
+    # length and this method returns the string of +length+ bytes beginning at
+    # +offset+.
+    #
+    # Examples:
+    #
+    #     my_struct = struct(['int id']).malloc
+    #     my_struct.id = 1
+    #     my_struct['id'] # => 1
+    #     my_struct[0, 4] # => "\x01\x00\x00\x00".b
+    #
+    def [](*args)
+      return super(*args) if args.size > 1
+      name = args[0]
       idx = @members.index(name)
       if( idx.nil? )
         raise(ArgumentError, "no such member: #{name}")
@@ -182,8 +197,20 @@ def [](name)
       end
     end
 
-    # Set struct member +name+, to value +val+
-    def []=(name, val)
+    # Set struct member +name+, to value +val+. If more arguments are
+    # specified, writes the string of bytes to the memory at the given
+    # +offset+ and +length+.
+    #
+    # Examples:
+    #
+    #     my_struct = struct(['int id']).malloc
+    #     my_struct['id'] = 1
+    #     my_struct[0, 4] = "\x01\x00\x00\x00".b
+    #     my_struct.id # => 1
+    #
+    def []=(*args)
+      return super(*args) if args.size > 2
+      name, val = *args
       idx = @members.index(name)
       if( idx.nil? )
         raise(ArgumentError, "no such member: #{name}")
diff --git a/test/fiddle/test_import.rb b/test/fiddle/test_import.rb
@@ -54,6 +54,28 @@ def test_ensure_call_dlload
       assert_match(/call dlload before/, err.message)
     end
 
+    def test_struct_memory_access()
+      # check memory operations performed directly on struct
+      my_struct = Fiddle::Importer.struct(['int id']).malloc
+      my_struct[0, Fiddle::SIZEOF_INT] = "\x01".b * Fiddle::SIZEOF_INT
+      assert_equal 0x01010101, my_struct.id
+
+      my_struct.id = 0
+      assert_equal "\x00".b * Fiddle::SIZEOF_INT, my_struct[0, Fiddle::SIZEOF_INT]
+    end
+
+    def test_struct_ptr_array_subscript_multiarg()
+      # check memory operations performed on struct#to_ptr
+      struct = Fiddle::Importer.struct([ 'int x' ]).malloc
+      ptr = struct.to_ptr
+
+      struct.x = 0x02020202
+      assert_equal("\x02".b * Fiddle::SIZEOF_INT, ptr[0, Fiddle::SIZEOF_INT])
+
+      ptr[0, Fiddle::SIZEOF_INT] = "\x01".b * Fiddle::SIZEOF_INT
+      assert_equal 0x01010101, struct.x
+    end
+
     def test_malloc()
       s1 = LIBC::Timeval.malloc()
       s2 = LIBC::Timeval.malloc()
