<commit_before>require 'allocation_tracer'

class Hotch
  def self.memory(name: $0, aggregate: true, &block)
    memory = if aggregate
               $hotch_memory ||= Memory.new(name)
             else
               caller = Kernel.caller_locations(1).first
               name = "#{name}:#{caller.path}:#{caller.lineno}"
               Memory.new(name)
             end

    memory.report_at_exit

    if block
      memory.run(&block)
    else
      memory.start
    end
  end

  class Memory
    def initialize(name, ignore_paths: [], disable_gc: true)
      @name = name
      @ignore_paths = Array(ignore_paths || [])
      @reports = []
      @started = nil
      @disable_gc = disable_gc
    end

    def start
      return if @started
      GC.disable if @disable_gc
      ObjectSpace::AllocationTracer.setup [:path, :line, :type]
      ObjectSpace::AllocationTracer.start
      @started = true
    end

    def stop
      return unless @started
      results = ObjectSpace::AllocationTracer.stop
      @started = nil
      GC.enable if @disable_gc
      @reports << Report.new(results, @ignore_paths)
    end

    def run
      start
      yield
    ensure
      stop
    end

    def report
      # TODO make it persistent (as CSV)
      report = @reports.inject(:sum)
      yield report
      @reports.clear
    end

    def report_at_exit
      return if defined? @at_exit_installed

      at_exit do
        stop

        report do |report|
          report.puts($stdout)
        end
      end

      @at_exit_installed = true
    end

    private

    def name
      @name.gsub(/\W+/, '_')
    end

    class Report
      attr_reader :lines

      def initialize(results, ignore_paths)
        @header = Line.new(*Line.members)
        @lines = results.map do |result|
          Line.from_result(result, ignore_paths)
        end.compact
      end

      def sum(other)
        by_key = Hash[@lines.map { |line| [line.key, line] }]
        other.lines.each do |line|
          if existing = by_key[line.key]
            existing.sum(line)
          else
            by_key[line.key] = line
            @lines << line
          end
        end
        self
      end

      def format
        # TODO refactor
        max_lengths = Array.new(Line.members.size, 0)
        ([@header, @total] + @lines).each do |line|
          line.lengths.each.with_index do |length, i|
            max_lengths[i] = length if length > max_lengths[i]
          end
        end
        max_lengths.map { |len| "%#{len}s" }.join(" ")
      end

      def puts(io)
        total!
        fmt = format
        @header.puts(io, fmt)
        @lines.sort_by(&:count).each { |line| line.puts(io, fmt) }
        @total.puts(io, fmt)
      end

      def to_s
        io = StringIO.new
        puts(io)
        io.string
      end

      private def total!
        return if defined? @total

        @total = Line::Total.new
        @lines.each do |line|
          @total.sum(line)
        end
      end

      class Line < Struct.new(:filename, :type, :count, :old_count, :total_age,
                              :min_age, :max_age, :total_memsize)
        # [
        #   [path, lineno, type],
        #   [count, old_count, total_age, min_age, max_age, total_memsize]
        # ]
        def self.from_result(result, ignore_paths)
          path, line, *args = result.flatten(1)
          return if ignore_paths.any? { |ip| ip == path || ip === path }
          filename = "#{strip_path(path)}:#{line}"
          new(filename, *args)
        end<commit_msg>Ensure that path is always set even if not defined<commit_after>require 'allocation_tracer'

class Hotch
  def self.memory(name: $0, aggregate: true, &block)
    memory = if aggregate
               $hotch_memory ||= Memory.new(name)
             else
               caller = Kernel.caller_locations(1).first
               name = "#{name}:#{caller.path}:#{caller.lineno}"
               Memory.new(name)
             end

    memory.report_at_exit

    if block
      memory.run(&block)
    else
      memory.start
    end
  end

  class Memory
    def initialize(name, ignore_paths: [], disable_gc: true)
      @name = name
      @ignore_paths = Array(ignore_paths || [])
      @reports = []
      @started = nil
      @disable_gc = disable_gc
    end

    def start
      return if @started
      GC.disable if @disable_gc
      ObjectSpace::AllocationTracer.setup [:path, :line, :type]
      ObjectSpace::AllocationTracer.start
      @started = true
    end

    def stop
      return unless @started
      results = ObjectSpace::AllocationTracer.stop
      @started = nil
      GC.enable if @disable_gc
      @reports << Report.new(results, @ignore_paths)
    end

    def run
      start
      yield
    ensure
      stop
    end

    def report
      # TODO make it persistent (as CSV)
      report = @reports.inject(:sum)
      yield report
      @reports.clear
    end

    def report_at_exit
      return if defined? @at_exit_installed

      at_exit do
        stop

        report do |report|
          report.puts($stdout)
        end
      end

      @at_exit_installed = true
    end

    private

    def name
      @name.gsub(/\W+/, '_')
    end

    class Report
      attr_reader :lines

      def initialize(results, ignore_paths)
        @header = Line.new(*Line.members)
        @lines = results.map do |result|
          Line.from_result(result, ignore_paths)
        end.compact
      end

      def sum(other)
        by_key = Hash[@lines.map { |line| [line.key, line] }]
        other.lines.each do |line|
          if existing = by_key[line.key]
            existing.sum(line)
          else
            by_key[line.key] = line
            @lines << line
          end
        end
        self
      end

      def format
        # TODO refactor
        max_lengths = Array.new(Line.members.size, 0)
        ([@header, @total] + @lines).each do |line|
          line.lengths.each.with_index do |length, i|
            max_lengths[i] = length if length > max_lengths[i]
          end
        end
        max_lengths.map { |len| "%#{len}s" }.join(" ")
      end

      def puts(io)
        total!
        fmt = format
        @header.puts(io, fmt)
        @lines.sort_by(&:count).each { |line| line.puts(io, fmt) }
        @total.puts(io, fmt)
      end

      def to_s
        io = StringIO.new
        puts(io)
        io.string
      end

      private def total!
        return if defined? @total

        @total = Line::Total.new
        @lines.each do |line|
          @total.sum(line)
        end
      end

      class Line < Struct.new(:filename, :type, :count, :old_count, :total_age,
                              :min_age, :max_age, :total_memsize)
        # [
        #   [path, lineno, type],
        #   [count, old_count, total_age, min_age, max_age, total_memsize]
        # ]
        def self.from_result(result, ignore_paths)
          path, line, *args = result.flatten(1)
          return if ignore_paths.any? { |ip| ip == path || ip === path }
          filename = "#{strip_path(path || "?")}:#{line}"
          new(filename, *args)
        end
{'fraction_non_alphanumeric': 0.10549173437937798, 'fraction_numerical': 0.0011207621182404036, 'mean_word_length': 1.7227307398932112, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 5, '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': '1505340', 'n_tokens_mistral': 2258, 'n_tokens_neox': 2196, 'n_words': 541}