<commit_before>module Furnace::AVM2
  module Transform
    class NFNormalize
      include AST::Visitor

      def transform(nf)
        @nf = nf.normalize_hierarchy!

        remove_useless_return
        visit @nf

        @nf
      end

      def remove_useless_return
        if @nf.children.last.type == :return_void
          @nf.children.slice! -1
        end
      end

      def on_nop(node)
        node.update(:remove)
      end

      COERCE_MAP = {
        :coerce_a => :any,
        :coerce_b => :bool,
        :coerce_s => :string,
      }

      def on_coerce_imm(node)
        expr, = node.children
        node.update(:coerce, [
          COERCE_MAP[node.type],
          expr
        ])
      end
      alias :on_coerce_a :on_coerce_imm
      alias :on_coerce_b :on_coerce_imm
      alias :on_coerce_s :on_coerce_imm

      CONVERT_MAP = {
        :convert_i => :integer,
        :convert_u => :unsigned,
        :convert_d => :double,
        :convert_s => :string,
        :convert_o => :object,
      }

      def on_convert_imm(node)
        expr, = node.children
        node.update(:convert, [
          CONVERT_MAP[node.type],
          expr
        ])
      end
      alias :on_convert_i :on_convert_imm
      alias :on_convert_u :on_convert_imm
      alias :on_convert_d :on_convert_imm
      alias :on_convert_s :on_convert_imm
      alias :on_convert_o :on_convert_imm

      ForInMatcher = AST::Matcher.new do
        [:while,
          [:has_next2, capture(:object_reg), capture(:index_reg)],
          [:begin,
            [ either_multi[
                [ :set_local, capture(:value_reg) ],
                [ :set_slot, capture(:value_reg), [:get_scope_object, 1] ],
              ],
              [ either[:coerce, :convert], capture(:value_type),
                [ capture(:iterator),
                  [:get_local, backref(:object_reg)],
                  [:get_local, backref(:index_reg)]]]],
            capture_rest(:body)]]
      end

      ForInIndexMatcher = AST::Matcher.new do
        [:set_local, backref(:index_reg), [:integer, 0]]
      end

      ForInObjectMatcher = AST::Matcher.new do
        [:set_local, backref(:object_reg),
          [:coerce, :any,
            capture(:root)]]
      end

      SuperfluousContinueMatcher = AST::Matcher.new do
        [:continue]
      end

      def on_while(node)
        *whatever, code = node.children
        if SuperfluousContinueMatcher.match code.children.last
          code.children.slice! -1
        end

        if captures = ForInMatcher.match(node)
          parent = node.parent

          case captures[:iterator]
          when :next_name
            type = :for_in
          when :next_value
            type = :for_each_in
          else
            return
          end

          index_node = object_node = nil

          loop_index = parent.children.index(node)
          parent.children[0..loop_index].reverse_each do |parent_node|
            if ForInIndexMatcher.match(parent_node, captures)
              index_node  = parent_node
            elsif ForInObjectMatcher.match(parent_node, captures)
              object_node = parent_node
            end

            break if index_node && object_node
          end

          return unless index_node && object_node

          index_node.update(:remove)
          object_node.update(:remove)

          node.update(type, [
            captures[:value_reg],
            captures[:value_type],
            captures[:object_reg],
            AST::Node.new(:begin, captures[:body])
          ])
        end
      end
    end
  end
end<commit_msg>Fix expanded for..in loops.<commit_after>module Furnace::AVM2
  module Transform
    class NFNormalize
      include AST::Visitor

      def transform(nf)
        @nf = nf.normalize_hierarchy!

        remove_useless_return
        visit @nf

        @nf
      end

      def remove_useless_return
        if @nf.children.last.type == :return_void
          @nf.children.slice! -1
        end
      end

      def on_nop(node)
        node.update(:remove)
      end

      COERCE_MAP = {
        :coerce_a => :any,
        :coerce_b => :bool,
        :coerce_s => :string,
      }

      def on_coerce_imm(node)
        expr, = node.children
        node.update(:coerce, [
          COERCE_MAP[node.type],
          expr
        ])
      end
      alias :on_coerce_a :on_coerce_imm
      alias :on_coerce_b :on_coerce_imm
      alias :on_coerce_s :on_coerce_imm

      CONVERT_MAP = {
        :convert_i => :integer,
        :convert_u => :unsigned,
        :convert_d => :double,
        :convert_s => :string,
        :convert_o => :object,
      }

      def on_convert_imm(node)
        expr, = node.children
        node.update(:convert, [
          CONVERT_MAP[node.type],
          expr
        ])
      end
      alias :on_convert_i :on_convert_imm
      alias :on_convert_u :on_convert_imm
      alias :on_convert_d :on_convert_imm
      alias :on_convert_s :on_convert_imm
      alias :on_convert_o :on_convert_imm

      ExpandedForInMatcher = AST::Matcher.new do
        [:if, [:has_next2, skip], skip]
      end

      # Loops can get expanded, but conditionals would never contain
      # has-next2.
      def on_if(node)
        if ExpandedForInMatcher.match node
          condition, body, rest = node.children

          loop = AST::Node.new(:while, [ condition, body ])
          on_while(loop, node.parent, node)

          node.update(:expand, [ loop ] + rest.children)
        end
      end

      ForInMatcher = AST::Matcher.new do
        [:while,
          [:has_next2, capture(:object_reg), capture(:index_reg)],
          [:begin,
            [ either_multi[
                [ :set_local, capture(:value_reg) ],
                [ :set_slot, capture(:value_reg), [:get_scope_object, 1] ],
              ],
              [ either[:coerce, :convert], capture(:value_type),
                [ capture(:iterator),
                  [:get_local, backref(:object_reg)],
                  [:get_local, backref(:index_reg)]]]],
            capture_rest(:body)]]
      end

      ForInIndexMatcher = AST::Matcher.new do
        [:set_local, backref(:index_reg), [:integer, 0]]
      end

      ForInObjectMatcher = AST::Matcher.new do
        [:set_local, backref(:object_reg),
          [:coerce, :any,
            capture(:root)]]
      end

      SuperfluousContinueMatcher = AST::Matcher.new do
        [:continue]
      end

      def on_while(node, parent=node.parent, enclosure=node)
        *whatever, code = node.children
        if SuperfluousContinueMatcher.match code.children.last
          code.children.slice! -1
        end

        if captures = ForInMatcher.match(node)
          case captures[:iterator]
          when :next_name
            type = :for_in
          when :next_value
            type = :for_each_in
          else
            return
          end

          index_node = object_node = nil

          loop_index = parent.children.index(enclosure)
          parent.children[0..loop_index].reverse_each do |parent_node|
            if ForInIndexMatcher.match(parent_node, captures)
              index_node  = parent_node
            elsif ForInObjectMatcher.match(parent_node, captures)
              object_node = parent_node
            end

            break if index_node && object_node
          end

          return unless index_node && object_node

          index_node.update(:remove)
          object_node.update(:remove)

          node.update(type, [
            captures[:value_reg],
            captures[:value_type],
            captures[:object_reg],
            AST::Node.new(:begin, captures[:body])
          ])
        end
      end
    end
  end
end
{'fraction_non_alphanumeric': 0.09808488612836438, 'fraction_numerical': 0.002070393374741201, 'mean_word_length': 1.7792161093131966, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1678757', 'n_tokens_mistral': 2412, 'n_tokens_neox': 2262, 'n_words': 451}