Optimize preceding-sibling, following-sibling, following, and descendant axes with an unified single mechanism

lib/rexml/xpath_parser.rb

@@ -467,57 +467,18 @@ def following_sibling(nodeset, tester, selector)
 
     def preceding_following_sibling(nodeset, tester, selector, reverse:)
       nodeset = nodeset.select {|node| node.respond_to?(:parent) && node.parent }
-      operator, value = selector
-      case operator
-      when :uniq
-        nodeset.group_by(&:parent).flat_map do |parent, sibling_nodes|
-          sets = Set.new.compare_by_identity
-          sibling_nodes.each {|sibling| sets << sibling }
-          children = parent.children
-          children = children.reverse if reverse
-          children.drop_while {|child| !sets.include?(child) }.drop(1)
-        end.select(&tester)
-      when :index_eq, :index_lt, :index_gt
-        nodeset.group_by(&:parent).flat_map do |parent, sibling_nodes|
-          anchors = Set.new.compare_by_identity
-          sibling_nodes.each {|sibling| anchors << sibling }
-          children = parent.children
-          children = children.reverse if reverse
-          followings = children.drop_while {|child| !anchors.include?(child) }.drop(1)
-          anchor_indexes = Set[0]
-          last_anchor = 0
-          index = 0
-          matched = []
-          followings.each do |node|
-            if tester.call(node)
-              case operator
-              when :index_eq
-                # anchor_indexes only contain values smaller or equal to `index`,
-                # so value < 0 case doesn't accidentally match any node.
-                matched << node if anchor_indexes.include?(index - value)
-              when :index_lt
-                # Position from the last anchor will be the minimum possible position for the node
-                matched << node if index - last_anchor < value
-              when :index_gt
-                # Position from the first anchor(==0) will be the maximum possible position for the node
-                matched << node if index > value
-              end
-              index += 1
-            end
-            if anchors.include?(node)
-              anchor_indexes << index
-              last_anchor = index
-            end
-          end
-          matched
+      nodeset.group_by(&:parent).flat_map do |parent, sibling_nodes|
+        anchors = Set.new.compare_by_identity.replace(sibling_nodes)
+        children = parent.children
+        children = children.reverse if reverse
+        followings = children.drop_while {|child| !anchors.include?(child) }.drop(1)
+        events = [:push]
+        followings.each do |node|
+          events << node if tester.call(node)
+          events << :push if anchors.include?(node)
         end
-      else # Slow path for :nodesets, :reverse_index_eq, :reverse_index_lt, :reverse_index_gt
-        nodesets = nodeset.map do |node|
-          parent = node.parent
-          index = parent.children.index(node)
-          reverse ? parent.children[0...index].reverse : parent.children[index + 1..-1]
-        end
-        non_optimized_nodesets_select(nodesets, tester, selector)
+        anchors.size.times { events << :pop }
+        sequence_positional_scan(events, selector)
       end
     end
 
@@ -829,48 +790,142 @@ def descendant_or_self(nodeset, tester, selector)
     # Scanner for descendant axis
     def descendant(nodeset, tester, selector, include_self: false)
       nodeset = nodeset.select {|node| node.respond_to?(:children) }
-      case selector
-      when :uniq
-        seen = Set.new.compare_by_identity
-        recursive = ->(node) do
-          node_type = node.node_type
-          return if seen.include?(node)
-          seen << node if node_type != :xmldecl
-          return unless node_type == :element || node_type == :document
-          node.children.each do |child|
-            recursive.call(child)
+      targets = Set.new.compare_by_identity.replace(nodeset)
+      descendant_anchor_roots(nodeset).flat_map do |root|
+        descendant_positional_scan(root, targets, tester, selector, include_self)
+      end
+    end
+
+    def descendant_anchor_roots(nodes)
+      seen = Set.new.compare_by_identity
+      nodes.each do |node|
+        descendant_traverse(node, include_self: false) do |n|
+          if !seen.include?(n)
+            seen << n
+            true
           end
         end
-        nodeset.each do |node|
+      end
+      nodes.reject {|node| seen.include?(node) }
+    end
+
+    def descendant_positional_scan(root, targets, tester, selector, include_self)
+      events = []
+      descendant_traverse_event(root) do |type, node|
+        if type == :enter
           if include_self
-            recursive.call(node)
+            events << :push if targets.include?(node)
+            events << node if tester.call(node)
           else
-            node.children.each(&recursive)
+            events << node if !node.equal?(root) && tester.call(node)
+            events << :push if targets.include?(node)
           end
+        elsif type == :leave
+          events << :pop if targets.include?(node)
         end
-        seen.select(&tester)
-      else
-        nodesets = nodeset.map do |node|
-          new_nodeset = []
-          new_nodes = {}
-          descendant_recursive(node, new_nodeset, new_nodes, include_self)
-          new_nodeset
+      end
+      sequence_positional_scan(events, selector)
+    end
+
+    # Select nodes matching the positional predicate from a sequence of events.
+    # Events are: :push, :pop, and node
+    # push/pop is an event that pushes/pops an anchor of position-based predicate
+    def sequence_positional_scan(events, selector)
+      case selector
+      when :uniq
+        return events.grep_v(Symbol)
+      when :nodesets
+        nodes = []
+        indexes = []
+        nodesets = []
+        last_range = nil
+        events.each do |e|
+          case e
+          when :push
+            indexes << nodes.size
+          when :pop
+            start_idx = indexes.pop
+            range = start_idx...nodes.size
+            nodesets << nodes[range] if last_range != range
+            last_range = range
+          else
+            nodes << e
+          end
         end
-        non_optimized_nodesets_select(nodesets, tester, selector)
+        return nodesets
       end
+
+      operator, value = selector
+      reverse = operator == :reverse_index_eq || operator == :reverse_index_lt || operator == :reverse_index_gt
+      events = events.reverse_each if reverse
+      start_event = reverse ? :pop : :push
+      end_event = reverse ? :push : :pop
+      anchor_indexes = []
+      anchor_set = Set.new
+      node_index = 0
+      result = []
+      events.each do |event|
+        if event == start_event
+          anchor_indexes << node_index
+          anchor_set << node_index
+        elsif event == end_event
+          idx = anchor_indexes.pop
+          anchor_set.delete(idx) if anchor_indexes.last != idx
+        else # event is a node that passed the tester
+          case operator
+          when :index_eq, :reverse_index_eq
+            result << event if anchor_set.include?(node_index - value)
+          when :index_lt, :reverse_index_lt
+            result << event if node_index - anchor_indexes.last < value
+          when :index_gt, :reverse_index_gt
+            result << event if node_index > value
+          end
+          node_index += 1
+        end
+      end
+      result
     end
 
-    def descendant_recursive(node, new_nodeset, new_nodes, include_self)
-      if include_self
-        return if new_nodes.key?(node)
-        new_nodeset << node
-        new_nodes[node] = true
+    # Scans the node and its descendants in document order
+    # and dispatches two types of events: :enter and :leave for each node.
+    def descendant_traverse_event(node)
+      stack = [node]
+      until stack.empty?
+        if stack.last
+          node = stack.last
+          # Push nil as a mark that we are entering this node.
+          # If we see this mark again, we know to pop the node and yield :leave
+          stack << nil
+          yield :enter, node
+          node_type = node.node_type
+          if node_type == :element or node_type == :document
+            node.children.reverse_each do |child|
+              stack << child if child.node_type != :xmldecl
+            end
+          end
+        else
+          stack.pop
+          node = stack.pop
+          yield :leave, node
+        end
       end
+    end
 
-      node_type = node.node_type
-      if node_type == :element or node_type == :document
-        node.children.each do |child|
-          descendant_recursive(child, new_nodeset, new_nodes, true)
+    # Scans the descendants of a node in document order and yields each node to the block.
+    # If a block returns falsy value, the node's descendants are not traversed.
+    def descendant_traverse(anchor_node, include_self:)
+      stack = [anchor_node]
+      until stack.empty?
+        node = stack.pop
+        if include_self || !node.equal?(anchor_node)
+          next unless yield node
+        end
+
+        node_type = node.node_type
+        if node_type == :element or node_type == :document
+          node.children.reverse_each do |child|
+            stack << child if child.node_type != :xmldecl
+          end
         end
       end
     end
@@ -923,36 +978,14 @@ def preceding_node_of( node )
 
     # Scanner for following axis
     def following(nodeset, tester, selector)
-      nodesets = nodeset.select {|node| node.respond_to?(:parent) }.map do |node|
-        following_nodes(node)
-      end
-      non_optimized_nodesets_select(nodesets, tester, selector)
-    end
-
-    def following_nodes(node)
-      followings = []
-      following_node = next_sibling_node(node)
-      while following_node
-        followings << following_node
-        following_node = following_node_of(following_node)
+      anchors = Set.new.compare_by_identity.replace(nodeset)
+      events = []
+      descendant_traverse_event(nodeset.first.document || nodeset.first.root) do |type, node|
+        events << :push if type == :leave && anchors.include?(node)
+        events << node if !events.empty? && type == :enter && tester.call(node)
       end
-      followings
-    end
-
-    def following_node_of( node )
-      return node.children[0] if node.kind_of?(Element) and node.children.size > 0
-
-      next_sibling_node(node)
-    end
-
-    def next_sibling_node(node)
-      psn = node.next_sibling_node
-      while psn.nil?
-        return nil if node.parent.nil? or node.parent.class == Document
-        node = node.parent
-        psn = node.next_sibling_node
-      end
-      psn
+      anchors.size.times { events << :pop }
+      sequence_positional_scan(events, selector)
     end
 
     def child(nodeset)

test/xpath/test_axis_preceding_sibling.rb

@@ -98,6 +98,9 @@ def test_preceding_following_sibling_multiple_anchors
         </a>
       XML
 
+      assert_equal(%w[1 3 4 5 7 9 11], XPath.match(doc, "/a/anchor/preceding-sibling::b[@id][position() mod 4 = 1]").map {|n| n.attributes["id"] })
+      assert_equal(%w[5 9 10 12], XPath.match(doc, "/a/anchor/following-sibling::b[@id][position() mod 4 = 1]").map {|n| n.attributes["id"] })
+
       assert_equal(%w[2 7 9], XPath.match(doc, "/a/anchor/preceding-sibling::b[position() = 3]").map {|n| n.attributes["id"] })
       assert_equal(%w[2 3 4 7 8 9 10 11], XPath.match(doc, "/a/anchor/preceding-sibling::b[position() <= 3]").map {|n| n.attributes["id"] })
       assert_equal(%w[2 3 4 7 8 9 10 11], XPath.match(doc, "/a/anchor/preceding-sibling::b[4 > position()]").map {|n| n.attributes["id"] })

test/xpath/test_base.rb

@@ -93,6 +93,56 @@ def test_descendant
       assert_equal( 1, XPath.match( doc, "/descendant::z[1]" ).size )
     end
 
+    def test_descendant_positions
+      positions = ['[1]', '[position() <= 2]', '[last()]', '[position() mod 3 = 1]']
+      anchors = {
+        '//a[@id=1]'  => [%w[2], %w[2 3], %w[18], %w[2 5 8 10 13 16]],
+        '//b[@id=2]'  => [%w[3], %w[3 4], %w[4], %w[3]],
+        '//e[@id=9]'  => [%w[10], %w[10 11], %w[14], %w[10 13]],
+        '//f[@id=12]' => [%w[13], %w[13], %w[13], %w[13]],
+      }
+      anchors.each do |anchor_xpath, expecteds|
+        positions.zip(expecteds).each do |position, expected|
+          xpath = "#{anchor_xpath}/descendant::*#{position}"
+          assert_equal(expected, XPath.match(@@doc, xpath).map {|n| n['id'] }, xpath)
+        end
+      end
+      expecteds = anchors.values.transpose.map {|ids| ids.flatten.uniq.sort_by(&:to_i) }
+      positions.zip(expecteds).each do |position, expected|
+        xpath = "(#{anchors.keys.join('|')})/descendant::*#{position}"
+        assert_equal(expected, XPath.match(@@doc, xpath).map {|n| n['id'] }, xpath)
+      end
+    end
+
+    def test_descendant_or_self_positions
+      positions = ['[1]', '[position() <= 2]', '[last()]', '[position() mod 3 = 1]']
+      anchors = {
+        '//a[@id=1]'  => [%w[1], %w[1 2], %w[18], %w[1 4 7 9 12 15 18]],
+        '//b[@id=2]'  => [%w[2], %w[2 3], %w[4], %w[2]],
+        '//e[@id=9]'  => [%w[9], %w[9 10], %w[14], %w[9 12]],
+        '//f[@id=12]' => [%w[12], %w[12 13], %w[13], %w[12]],
+      }
+      anchors.each do |anchor_xpath, expecteds|
+        positions.zip(expecteds).each do |position, expected|
+          xpath = "#{anchor_xpath}/descendant-or-self::*#{position}"
+          assert_equal(expected, XPath.match(@@doc, xpath).map {|n| n['id'] }, xpath)
+        end
+      end
+      expecteds = anchors.values.transpose.map {|ids| ids.flatten.uniq.sort_by(&:to_i) }
+      positions.zip(expecteds).each do |position, expected|
+        xpath = "(#{anchors.keys.join('|')})/descendant-or-self::*#{position}"
+        assert_equal(expected, XPath.match(@@doc, xpath).map {|n| n['id'] }, xpath)
+      end
+    end
+
+    def test_following_positions
+      anchor_xpath = '(//b[@id=2]|//e[@id=9]|//f[@id=12])'
+      assert_equal(%w[5 14 15], XPath.match(@@doc, "#{anchor_xpath}/following::*[1]").map {|n| n['id'] })
+      assert_equal(%w[5 6 14 15 16], XPath.match(@@doc, "#{anchor_xpath}/following::*[position()<=2]").map {|n| n['id'] })
+      assert_equal(%w[18], XPath.match(@@doc, "#{anchor_xpath}/following::*[last()]").map {|n| n['id'] })
+      assert_equal(%w[5 8 10 13 14 15 16 17 18], XPath.match(@@doc, "#{anchor_xpath}/following::*[position() mod 3 = 1]").map {|n| n['id'] })
+    end
+
     def test_root
       source = "<a><b/></a>"
       doc = Document.new( source )