Optimization of REXML::Document#add, #add_element

benchmark/parse_comment.yaml

@@ -24,13 +24,15 @@ contexts:
 prelude: |
   require 'rexml/document'
 
-  SIZE = 100000
+  SIZE = 1000
 
   top_level_xml     = "<!--" + "a" * SIZE + "-->\n<root/>"
   in_doctype_xml    = "<!DOCTYPE foo [<!--" + "a" * SIZE + "-->]><root/>"
   after_doctype_xml = "<root/><!--" + "a" * SIZE + "-->"
+  many_comments_xml = "<!---->" * SIZE + "<a/>"
 
 benchmark:
   'top_level'      : REXML::Document.new(top_level_xml)
   'in_doctype'     : REXML::Document.new(in_doctype_xml)
   'after_doctype'  : REXML::Document.new(after_doctype_xml)
+  'many_comments'  : REXML::Document.new(many_comments_xml)

lib/rexml/document.rb

@@ -197,9 +197,8 @@ def add( child )
         end
         child.parent = self
       else
-        rv = super
-        raise "attempted adding second root element to document" if @elements.size > 1
-        rv
+        raise "attempted adding second root element to document" if child.kind_of?(Element) && root
+        super
       end
     end
     alias :<< :add
@@ -211,9 +210,8 @@ def add( child )
     #
     #   REXML::Element.add_element(name_or_element, attributes)
     def add_element(arg=nil, arg2=nil)
-      rv = super
-      raise "attempted adding second root element to document" if @elements.size > 1
-      rv
+      raise "attempted adding second root element to document" if root
+      super
     end
 
     # :call-seq:

test/test_document.rb

@@ -320,6 +320,50 @@ def test_encoding
       end
     end
 
+    class AddTest < Test::Unit::TestCase
+      def test_add_second_root_element_raises
+        doc = REXML::Document.new("<root/>")
+        assert_raise(RuntimeError, "attempted adding second root element to document") do
+          doc.add(REXML::Element.new("second"))
+        end
+      end
+
+      def test_append_operator_second_root_element_raises
+        doc = REXML::Document.new("<root/>")
+        assert_raise(RuntimeError, "attempted adding second root element to document") do
+          doc << REXML::Element.new("second")
+        end
+      end
+
+      def test_add_element_second_root_raises
+        doc = REXML::Document.new("<root/>")
+        assert_raise(RuntimeError, "attempted adding second root element to document") do
+          doc.add_element("second")
+        end
+      end
+
+      def test_add_element_with_element_second_root_raises
+        doc = REXML::Document.new("<root/>")
+        assert_raise(RuntimeError, "attempted adding second root element to document") do
+          doc.add_element(REXML::Element.new("second"))
+        end
+      end
+
+      def test_add_xml_decl_allowed
+        doc = REXML::Document.new("<root/>")
+        assert_nothing_raised do
+          doc.add(REXML::XMLDecl.new("1.0"))
+        end
+      end
+
+      def test_add_doctype_allowed
+        doc = REXML::Document.new("<root/>")
+        assert_nothing_raised do
+          doc.add(REXML::DocType.new("root"))
+        end
+      end
+    end
+
     class BomTest < Test::Unit::TestCase
       class HaveEncodingTest < self
         def test_utf_8