From fe4034ccf9c3ea019aa820154c7fea29ad4a75b8 Mon Sep 17 00:00:00 2001
From: Stephen Paul Weber (Work) <stephen.weber@shopify.com>
Date: Oct 18 2018 13:42:56 +0000
Subject: Merge pull request #1025 from Shopify/traverse-ast


Liquid::ParseTreeVisitor
---

diff --git a/lib/liquid.rb b/lib/liquid.rb
index 7d9da26..770d2f9 100644
--- a/lib/liquid.rb
+++ b/lib/liquid.rb
@@ -45,6 +45,7 @@ module Liquid
 end
 
 require "liquid/version"
+require 'liquid/parse_tree_visitor'
 require 'liquid/lexer'
 require 'liquid/parser'
 require 'liquid/i18n'
diff --git a/lib/liquid/condition.rb b/lib/liquid/condition.rb
index 3e79849..3b51682 100644
--- a/lib/liquid/condition.rb
+++ b/lib/liquid/condition.rb
@@ -29,7 +29,7 @@ module Liquid
       @@operators
     end
 
-    attr_reader :attachment
+    attr_reader :attachment, :child_condition
     attr_accessor :left, :operator, :right
 
     def initialize(left = nil, operator = nil, right = nil)
@@ -83,7 +83,7 @@ module Liquid
 
     protected
 
-    attr_reader :child_relation, :child_condition
+    attr_reader :child_relation
 
     private
 
@@ -128,6 +128,15 @@ module Liquid
         end
       end
     end
+
+    class ParseTreeVisitor < Liquid::ParseTreeVisitor
+      def children
+        [
+          @node.left, @node.right,
+          @node.child_condition, @node.attachment
+        ].compact
+      end
+    end
   end
 
   class ElseCondition < Condition
diff --git a/lib/liquid/parse_tree_visitor.rb b/lib/liquid/parse_tree_visitor.rb
new file mode 100644
index 0000000..74f5563
--- /dev/null
+++ b/lib/liquid/parse_tree_visitor.rb
@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module Liquid
+  class ParseTreeVisitor
+    def self.for(node, callbacks = Hash.new(proc {}))
+      if defined?(node.class::ParseTreeVisitor)
+        node.class::ParseTreeVisitor
+      else
+        self
+      end.new(node, callbacks)
+    end
+
+    def initialize(node, callbacks)
+      @node = node
+      @callbacks = callbacks
+    end
+
+    def add_callback_for(*classes, &block)
+      callback = block
+      callback = ->(node, _) { yield node } if block.arity.abs == 1
+      callback = ->(_, _) { yield } if block.arity.zero?
+      classes.each { |klass| @callbacks[klass] = callback }
+      self
+    end
+
+    def visit(context = nil)
+      children.map do |node|
+        item, new_context = @callbacks[node.class].call(node, context)
+        [
+          item,
+          ParseTreeVisitor.for(node, @callbacks).visit(new_context || context)
+        ]
+      end
+    end
+
+    protected
+
+    def children
+      @node.respond_to?(:nodelist) ? Array(@node.nodelist) : []
+    end
+  end
+end
diff --git a/lib/liquid/tags/assign.rb b/lib/liquid/tags/assign.rb
index f6cd5fa..c8d0574 100644
--- a/lib/liquid/tags/assign.rb
+++ b/lib/liquid/tags/assign.rb
@@ -10,6 +10,8 @@ module Liquid
   class Assign < Tag
     Syntax = /(#{VariableSignature}+)\s*=\s*(.*)\s*/om
 
+    attr_reader :to, :from
+
     def initialize(tag_name, markup, options)
       super
       if markup =~ Syntax
@@ -45,6 +47,12 @@ module Liquid
         1
       end
     end
+
+    class ParseTreeVisitor < Liquid::ParseTreeVisitor
+      def children
+        [@node.from]
+      end
+    end
   end
 
   Template.register_tag('assign'.freeze, Assign)
diff --git a/lib/liquid/tags/case.rb b/lib/liquid/tags/case.rb
index 453b4d6..5036b27 100644
--- a/lib/liquid/tags/case.rb
+++ b/lib/liquid/tags/case.rb
@@ -3,6 +3,8 @@ module Liquid
     Syntax     = /(#{QuotedFragment})/o
     WhenSyntax = /(#{QuotedFragment})(?:(?:\s+or\s+|\s*\,\s*)(#{QuotedFragment}.*))?/om
 
+    attr_reader :blocks, :left
+
     def initialize(tag_name, markup, options)
       super
       @blocks = []
@@ -80,6 +82,12 @@ module Liquid
       block.attach(BlockBody.new)
       @blocks << block
     end
+
+    class ParseTreeVisitor < Liquid::ParseTreeVisitor
+      def children
+        [@node.left] + @node.blocks
+      end
+    end
   end
 
   Template.register_tag('case'.freeze, Case)
diff --git a/lib/liquid/tags/cycle.rb b/lib/liquid/tags/cycle.rb
index ad116a6..17aa860 100644
--- a/lib/liquid/tags/cycle.rb
+++ b/lib/liquid/tags/cycle.rb
@@ -15,6 +15,8 @@ module Liquid
     SimpleSyntax = /\A#{QuotedFragment}+/o
     NamedSyntax  = /\A(#{QuotedFragment})\s*\:\s*(.*)/om
 
+    attr_reader :variables
+
     def initialize(tag_name, markup, options)
       super
       case markup
@@ -51,6 +53,12 @@ module Liquid
         $1 ? Expression.parse($1) : nil
       end.compact
     end
+
+    class ParseTreeVisitor < Liquid::ParseTreeVisitor
+      def children
+        Array(@node.variables)
+      end
+    end
   end
 
   Template.register_tag('cycle', Cycle)
diff --git a/lib/liquid/tags/for.rb b/lib/liquid/tags/for.rb
index 6c95624..b69aa78 100644
--- a/lib/liquid/tags/for.rb
+++ b/lib/liquid/tags/for.rb
@@ -46,8 +46,7 @@ module Liquid
   class For < Block
     Syntax = /\A(#{VariableSegment}+)\s+in\s+(#{QuotedFragment}+)\s*(reversed)?/o
 
-    attr_reader :collection_name
-    attr_reader :variable_name
+    attr_reader :collection_name, :variable_name, :limit, :from
 
     def initialize(tag_name, markup, options)
       super
@@ -192,6 +191,12 @@ module Liquid
     def render_else(context)
       @else_block ? @else_block.render(context) : ''.freeze
     end
+
+    class ParseTreeVisitor < Liquid::ParseTreeVisitor
+      def children
+        (super + [@node.limit, @node.from, @node.collection_name]).compact
+      end
+    end
   end
 
   Template.register_tag('for'.freeze, For)
diff --git a/lib/liquid/tags/if.rb b/lib/liquid/tags/if.rb
index 904369d..1451c25 100644
--- a/lib/liquid/tags/if.rb
+++ b/lib/liquid/tags/if.rb
@@ -14,21 +14,23 @@ module Liquid
     ExpressionsAndOperators = /(?:\b(?:\s?and\s?|\s?or\s?)\b|(?:\s*(?!\b(?:\s?and\s?|\s?or\s?)\b)(?:#{QuotedFragment}|\S+)\s*)+)/o
     BOOLEAN_OPERATORS = %w(and or)
 
+    attr_reader :blocks
+
     def initialize(tag_name, markup, options)
       super
       @blocks = []
       push_block('if'.freeze, markup)
     end
 
+    def nodelist
+      @blocks.map(&:attachment)
+    end
+
     def parse(tokens)
       while parse_body(@blocks.last.attachment, tokens)
       end
     end
 
-    def nodelist
-      @blocks.map(&:attachment)
-    end
-
     def unknown_tag(tag, markup, tokens)
       if ['elsif'.freeze, 'else'.freeze].include?(tag)
         push_block(tag, markup)
@@ -108,6 +110,12 @@ module Liquid
         Condition.new(a)
       end
     end
+
+    class ParseTreeVisitor < Liquid::ParseTreeVisitor
+      def children
+        @node.blocks
+      end
+    end
   end
 
   Template.register_tag('if'.freeze, If)
diff --git a/lib/liquid/tags/include.rb b/lib/liquid/tags/include.rb
index a800703..c9f2a28 100644
--- a/lib/liquid/tags/include.rb
+++ b/lib/liquid/tags/include.rb
@@ -16,6 +16,8 @@ module Liquid
   class Include < Tag
     Syntax = /(#{QuotedFragment}+)(\s+(?:with|for)\s+(#{QuotedFragment}+))?/o
 
+    attr_reader :template_name_expr, :variable_name_expr, :attributes
+
     def initialize(tag_name, markup, options)
       super
 
@@ -107,6 +109,15 @@ module Liquid
 
       file_system.read_template_file(context.evaluate(@template_name_expr))
     end
+
+    class ParseTreeVisitor < Liquid::ParseTreeVisitor
+      def children
+        [
+          @node.template_name_expr,
+          @node.variable_name_expr
+        ] + @node.attributes.values
+      end
+    end
   end
 
   Template.register_tag('include'.freeze, Include)
diff --git a/lib/liquid/tags/table_row.rb b/lib/liquid/tags/table_row.rb
index cfdef33..7f391cf 100644
--- a/lib/liquid/tags/table_row.rb
+++ b/lib/liquid/tags/table_row.rb
@@ -2,6 +2,8 @@ module Liquid
   class TableRow < Block
     Syntax = /(\w+)\s+in\s+(#{QuotedFragment}+)/o
 
+    attr_reader :variable_name, :collection_name, :attributes
+
     def initialize(tag_name, markup, options)
       super
       if markup =~ Syntax
@@ -48,6 +50,12 @@ module Liquid
       result << "</tr>\n"
       result
     end
+
+    class ParseTreeVisitor < Liquid::ParseTreeVisitor
+      def children
+        super + @node.attributes.values + [@node.collection_name]
+      end
+    end
   end
 
   Template.register_tag('tablerow'.freeze, TableRow)
diff --git a/lib/liquid/variable.rb b/lib/liquid/variable.rb
index 5f88eb3..c31bffe 100644
--- a/lib/liquid/variable.rb
+++ b/lib/liquid/variable.rb
@@ -138,5 +138,11 @@ module Liquid
         raise error
       end
     end
+
+    class ParseTreeVisitor < Liquid::ParseTreeVisitor
+      def children
+        [@node.name] + @node.filters.flatten
+      end
+    end
   end
 end
diff --git a/lib/liquid/variable_lookup.rb b/lib/liquid/variable_lookup.rb
index 3ed4e4a..8f7ad46 100644
--- a/lib/liquid/variable_lookup.rb
+++ b/lib/liquid/variable_lookup.rb
@@ -78,5 +78,11 @@ module Liquid
     def state
       [@name, @lookups, @command_flags]
     end
+
+    class ParseTreeVisitor < Liquid::ParseTreeVisitor
+      def children
+        @node.lookups
+      end
+    end
   end
 end
diff --git a/test/integration/parse_tree_visitor_test.rb b/test/integration/parse_tree_visitor_test.rb
new file mode 100644
index 0000000..6ad6a2d
--- /dev/null
+++ b/test/integration/parse_tree_visitor_test.rb
@@ -0,0 +1,247 @@
+# frozen_string_literal: true
+
+require 'test_helper'
+
+class ParseTreeVisitorTest < Minitest::Test
+  include Liquid
+
+  def test_variable
+    assert_equal(
+      ["test"],
+      visit(%({{ test }}))
+    )
+  end
+
+  def test_varible_with_filter
+    assert_equal(
+      ["test", "infilter"],
+      visit(%({{ test | split: infilter }}))
+    )
+  end
+
+  def test_dynamic_variable
+    assert_equal(
+      ["test", "inlookup"],
+      visit(%({{ test[inlookup] }}))
+    )
+  end
+
+  def test_if_condition
+    assert_equal(
+      ["test"],
+      visit(%({% if test %}{% endif %}))
+    )
+  end
+
+  def test_complex_if_condition
+    assert_equal(
+      ["test"],
+      visit(%({% if 1 == 1 and 2 == test %}{% endif %}))
+    )
+  end
+
+  def test_if_body
+    assert_equal(
+      ["test"],
+      visit(%({% if 1 == 1 %}{{ test }}{% endif %}))
+    )
+  end
+
+  def test_unless_condition
+    assert_equal(
+      ["test"],
+      visit(%({% unless test %}{% endunless %}))
+    )
+  end
+
+  def test_complex_unless_condition
+    assert_equal(
+      ["test"],
+      visit(%({% unless 1 == 1 and 2 == test %}{% endunless %}))
+    )
+  end
+
+  def test_unless_body
+    assert_equal(
+      ["test"],
+      visit(%({% unless 1 == 1 %}{{ test }}{% endunless %}))
+    )
+  end
+
+  def test_elsif_condition
+    assert_equal(
+      ["test"],
+      visit(%({% if 1 == 1 %}{% elsif test %}{% endif %}))
+    )
+  end
+
+  def test_complex_elsif_condition
+    assert_equal(
+      ["test"],
+      visit(%({% if 1 == 1 %}{% elsif 1 == 1 and 2 == test %}{% endif %}))
+    )
+  end
+
+  def test_elsif_body
+    assert_equal(
+      ["test"],
+      visit(%({% if 1 == 1 %}{% elsif 2 == 2 %}{{ test }}{% endif %}))
+    )
+  end
+
+  def test_else_body
+    assert_equal(
+      ["test"],
+      visit(%({% if 1 == 1 %}{% else %}{{ test }}{% endif %}))
+    )
+  end
+
+  def test_case_left
+    assert_equal(
+      ["test"],
+      visit(%({% case test %}{% endcase %}))
+    )
+  end
+
+  def test_case_condition
+    assert_equal(
+      ["test"],
+      visit(%({% case 1 %}{% when test %}{% endcase %}))
+    )
+  end
+
+  def test_case_when_body
+    assert_equal(
+      ["test"],
+      visit(%({% case 1 %}{% when 2 %}{{ test }}{% endcase %}))
+    )
+  end
+
+  def test_case_else_body
+    assert_equal(
+      ["test"],
+      visit(%({% case 1 %}{% else %}{{ test }}{% endcase %}))
+    )
+  end
+
+  def test_for_in
+    assert_equal(
+      ["test"],
+      visit(%({% for x in test %}{% endfor %}))
+    )
+  end
+
+  def test_for_limit
+    assert_equal(
+      ["test"],
+      visit(%({% for x in (1..5) limit: test %}{% endfor %}))
+    )
+  end
+
+  def test_for_offset
+    assert_equal(
+      ["test"],
+      visit(%({% for x in (1..5) offset: test %}{% endfor %}))
+    )
+  end
+
+  def test_for_body
+    assert_equal(
+      ["test"],
+      visit(%({% for x in (1..5) %}{{ test }}{% endfor %}))
+    )
+  end
+
+  def test_tablerow_in
+    assert_equal(
+      ["test"],
+      visit(%({% tablerow x in test %}{% endtablerow %}))
+    )
+  end
+
+  def test_tablerow_limit
+    assert_equal(
+      ["test"],
+      visit(%({% tablerow x in (1..5) limit: test %}{% endtablerow %}))
+    )
+  end
+
+  def test_tablerow_offset
+    assert_equal(
+      ["test"],
+      visit(%({% tablerow x in (1..5) offset: test %}{% endtablerow %}))
+    )
+  end
+
+  def test_tablerow_body
+    assert_equal(
+      ["test"],
+      visit(%({% tablerow x in (1..5) %}{{ test }}{% endtablerow %}))
+    )
+  end
+
+  def test_cycle
+    assert_equal(
+      ["test"],
+      visit(%({% cycle test %}))
+    )
+  end
+
+  def test_assign
+    assert_equal(
+      ["test"],
+      visit(%({% assign x = test %}))
+    )
+  end
+
+  def test_capture
+    assert_equal(
+      ["test"],
+      visit(%({% capture x %}{{ test }}{% endcapture %}))
+    )
+  end
+
+  def test_include
+    assert_equal(
+      ["test"],
+      visit(%({% include test %}))
+    )
+  end
+
+  def test_include_with
+    assert_equal(
+      ["test"],
+      visit(%({% include "hai" with test %}))
+    )
+  end
+
+  def test_include_for
+    assert_equal(
+      ["test"],
+      visit(%({% include "hai" for test %}))
+    )
+  end
+
+  def test_preserve_tree_structure
+    assert_equal(
+      [[nil, [
+        [nil, [[nil, [["other", []]]]]],
+        ["test", []],
+        ["xs", []]
+      ]]],
+      traversal(%({% for x in xs offset: test %}{{ other }}{% endfor %})).visit
+    )
+  end
+
+  private
+
+  def traversal(template)
+    ParseTreeVisitor
+      .for(Template.parse(template).root)
+      .add_callback_for(VariableLookup, &:name)
+  end
+
+  def visit(template)
+    traversal(template).visit.flatten.compact
+  end
+end