Add methods to Sexp to avoid Sexp creation

lib/ruby_parser/bm_sexp.rb

@@ -254,6 +254,35 @@ class Sexp
     end
   end
 
+  def each_arg replace = false
+    expect :call, :attrasgn, :super, :zsuper
+    range = nil
+
+    case self.node_type
+    when :call, :attrasgn
+      if self[3]
+        range = (3...self.length)
+      end
+    when :super, :zsuper
+      if self[1]
+        range = (1...self.length)
+      end
+    end
+
+    if range
+      range.each do |i|
+        res = yield self[i]
+        self[i] = res if replace
+      end
+    end
+
+    self
+  end
+
+  def each_arg! &block
+    self.each_arg true, &block
+  end
+
   #Returns first argument of a method call.
   def first_arg
     expect :call, :attrasgn
@@ -278,6 +307,26 @@ class Sexp
     self[4] = exp
   end
 
+  def third_arg
+    expect :call, :attrasgn
+    self[5]
+  end
+
+  def third_arg= exp
+    expect :call, :attrasgn
+    self[5] = exp
+  end
+
+  def last_arg
+    expect :call, :attrasgn
+
+    if self[3]
+      self[-1]
+    else
+      nil
+    end
+  end
+
   #Returns condition of an if expression:
   #
   #    s(:if,
@@ -462,6 +511,32 @@ class Sexp
     self.body.unshift :rlist
   end
 
+  def each_body replace = false
+    expect :defn, :defs, :methdef, :selfdef, :class, :module
+
+    range = case self.node_type
+            when :defn, :methdef, :class
+              (3...self.length)
+            when :defs, :selfdef
+              (4...self.length)
+            when :module
+              (2...self.length)
+            end
+
+    if range
+      range.each do |i|
+        res = yield self[i]
+        self[i] = res if replace
+      end
+    end
+
+    self
+  end
+
+  def each_body! &block
+    each_body true, &block
+  end
+
   def render_type
     expect :render
     self[1]

test/tests/test_sexp.rb

@@ -18,6 +18,7 @@ class SexpTests < Test::Unit::TestCase
     assert_equal s(:arglist), exp.arglist
     assert_nil exp.first_arg
     assert_nil exp.second_arg
+    assert_nil exp.last_arg
   end
 
   def test_method_call_with_args
@@ -29,6 +30,7 @@ class SexpTests < Test::Unit::TestCase
     assert_equal s(:arglist, s(:lit, 1), s(:lit, 2), s(:lit, 3)), exp.arglist
     assert_equal s(:lit, 1), exp.first_arg
     assert_equal s(:lit, 2), exp.second_arg
+    assert_equal s(:lit, 3), exp.last_arg
   end
 
   def test_method_call_no_target
@@ -40,6 +42,7 @@ class SexpTests < Test::Unit::TestCase
     assert_equal s(:arglist, s(:lit, 1), s(:lit, 2), s(:lit, 3)), exp.arglist
     assert_equal s(:lit, 1), exp.first_arg
     assert_equal s(:lit, 2), exp.second_arg
+    assert_equal s(:lit, 3), exp.last_arg
   end
 
   def test_method_call_set_target
@@ -55,6 +58,7 @@ class SexpTests < Test::Unit::TestCase
 
     assert_equal s(:lit, 1), exp.first_arg
     assert_equal s(:lit, 2), exp.second_arg
+    assert_equal s(:lit, 2), exp.last_arg
     assert_equal s(:arglist, s(:lit, 1), s(:lit, 2)), exp.arglist
     assert_equal s(s(:lit, 1), s(:lit, 2)), exp.args
   end
@@ -69,6 +73,7 @@ class SexpTests < Test::Unit::TestCase
     assert_equal s(s(:lit, 1), s(:lit, 2)), exp.args
     assert_equal s(:lit,1), exp.first_arg
     assert_equal s(:lit, 2), exp.second_arg
+    assert_equal s(:lit, 2), exp.last_arg
   end
 
   def test_method_call_with_block
@@ -277,4 +282,25 @@ class SexpTests < Test::Unit::TestCase
     assert_equal s(:iasgn, :@x, s(:lit, 1)), exp.iasgn(true)
     assert_equal nil, exp.iasgn #Was deleted
   end
+
+  def test_each_arg
+    exp = parse "blah 1, 2, 3"
+
+    args = []
+    exp.each_arg do |a|
+      args << a.value
+    end
+
+    assert_equal [1,2,3], args
+  end
+
+  def test_each_arg!
+    exp = parse "blah 1, 2"
+    exp.each_arg! do |a|
+      s(:lit, a.value + 1)
+    end
+
+    assert_equal s(:lit, 2), exp.first_arg
+    assert_equal s(:lit, 3), exp.second_arg
+  end
 end