From db39b25546dfb0bef963e51f767500daeb6b8053 Mon Sep 17 00:00:00 2001
From: Yorick Peterse <yorickpeterse@gmail.com>
Date: Fri, 17 Jul 2015 23:24:18 +0200
Subject: [PATCH] XPath compiler support for ancestor-or-self

This also comes with some changes to the specs as the old behaviour of
the Evaluator was incorrect. The Evaluator would bail after matching a
single node but instead it's meant to continue until it runs out of
parent nodes.
---
 lib/oga/ruby/node.rb                          |  2 ++
 lib/oga/xpath/compiler.rb                     | 32 +++++++++++++++++--
 .../compiler/axes/ancestor_or_self_spec.rb    |  5 +++
 3 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/lib/oga/ruby/node.rb b/lib/oga/ruby/node.rb
index da5c1d9..254b78f 100644
--- a/lib/oga/ruby/node.rb
+++ b/lib/oga/ruby/node.rb
@@ -27,6 +27,8 @@ module Oga
     # @private
     #
     class Node < BasicObject
+      undef_method :!, :!=
+
       # @return [Symbol]
       attr_reader :type
 
diff --git a/lib/oga/xpath/compiler.rb b/lib/oga/xpath/compiler.rb
index d9c1329..441e159 100644
--- a/lib/oga/xpath/compiler.rb
+++ b/lib/oga/xpath/compiler.rb
@@ -16,7 +16,11 @@ module Oga
       STAR = '*'
 
       # Node types that require a NodeSet to push nodes into.
-      USE_NODESET = [:path, :absolute_path, :axis]
+      USE_NODESET = [:path, :absolute_path, :axis, :predicate]
+
+      # Node types that require "compile" to define a block to call upon
+      # matching a node.
+      ADD_PUSH_BLOCK = [:axis, :predicate]
 
       ##
       # Compiles and caches an AST.
@@ -37,7 +41,7 @@ module Oga
         document = node_literal
         matched  = matched_literal
 
-        if ast.type == :axis
+        if ADD_PUSH_BLOCK.include?(ast.type)
           ruby_ast = process(ast, document) { |node| matched.push(node) }
         else
           ruby_ast = process(ast, document)
@@ -167,6 +171,30 @@ module Oga
         input.get(string(query))
       end
 
+      ##
+      # Processes the `ancestor-or-self` axis.
+      #
+      # @param [AST::Node] ast
+      # @param [Oga::Ruby::Node] input
+      # @return [Oga::Ruby::Node]
+      #
+      def on_axis_ancestor_or_self(ast, input, &block)
+        parent_var = literal('parent')
+        assign     = parent_var.assign(input)
+        has_parent = parent_var.respond_to?(symbol(:parent))
+          .and(parent_var.parent)
+
+        body = has_parent.while_true do
+          if_statement = process(ast, parent_var, &block).if_true do
+            yield parent_var
+          end
+
+          if_statement.followed_by(parent_var.assign(parent_var.parent))
+        end
+
+        assign.followed_by(body)
+      end
+
       ##
       # Processes a node test predicate.
       #
diff --git a/spec/oga/xpath/compiler/axes/ancestor_or_self_spec.rb b/spec/oga/xpath/compiler/axes/ancestor_or_self_spec.rb
index 58e0b08..f12672f 100644
--- a/spec/oga/xpath/compiler/axes/ancestor_or_self_spec.rb
+++ b/spec/oga/xpath/compiler/axes/ancestor_or_self_spec.rb
@@ -26,6 +26,11 @@ describe Oga::XPath::Compiler do
       evaluate_xpath(@c1, 'ancestor-or-self::*[1]').should == node_set(@c1)
     end
 
+    it 'returns a node set containing all ancestors' do
+      evaluate_xpath(@c1, 'ancestor-or-self::*')
+        .should == node_set(@c1, @b1, @a1)
+    end
+
     it 'returns an empty node set for non existing ancestors' do
       evaluate_xpath(@c1, 'ancestor-or-self::foo').should == node_set
     end