Improved bake tasks.

Author: samuel-williams-shopify

bake/memory/profiler.rb

@@ -6,34 +6,22 @@
 # Copyright, 2018, by Jonas Peschla.
 # Copyright, 2020-2022, by Samuel Williams.
 
-def check(paths:)
-	require "console"
-	
-	total_size = paths.sum{|path| File.size(path)}
+def initialize(...)
+	super
 
 	require_relative "../../lib/memory"
+end
+
+# Load and print a report from one or more .mprof files.
+#
+# This is a convenience task that combines sampler:load and report:print.
+#
+# @parameter paths [Array<String>] Paths to .mprof files.
+def check(*paths)
+	# Delegate to sampler:load
+	sampler = context.call("memory:sampler:load", *paths)
 
-	report = Memory::Report.general
-	
-	cache = Memory::Cache.new
-	wrapper = Memory::Wrapper.new(cache)
-	
-	progress = Console.logger.progress(report, total_size)
-	
-	paths.each do |path|
-		Console.logger.info(report, "Loading #{path}, #{Memory.formatted_bytes File.size(path)}")
-		
-		File.open(path) do |io|
-			unpacker = wrapper.unpacker(io)
-			count = unpacker.read_array_header
-			
-			report.concat(unpacker)
-			
-			progress.increment(io.size)
-		end
-		
-		Console.logger.info(report, "Loaded allocations, #{report.total_allocated}")
-	end
-	
+	# Generate and print report
+	report = sampler.report
 	report.print($stdout)
 end

bake/memory/report.rb

@@ -0,0 +1,27 @@
+
+def initialize(...)
+	super
+
+	require_relative "../../lib/memory"
+end
+
+# Print a memory report.
+# Accepts either a `Memory::Sampler` or `Memory::Report` as input.
+# If given a Sampler, generates a report first.
+#
+# @parameter input [Memory::Report] The sampler or report to print.
+def print(input:)
+	# Convert Sampler to Report if needed:
+	report = case input
+	when Memory::Sampler
+		input.report
+	when Memory::Report
+		input
+	else
+		raise ArgumentError, "Expected Memory::Sampler or Memory::Report, got #{input.class}"
+	end
+	
+	report.print($stderr)
+
+	return report
+end

bake/memory/sampler.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+# Released under the MIT License.
+# Copyright, 2025, by Samuel Williams.
+
+def initialize(...)
+	super
+	
+	require_relative '../../lib/memory'
+end
+
+# Load a sampler from one or more .mprof files.
+#
+# Multiple files will be combined into a single sampler, useful for
+# analyzing aggregated profiling data.
+#
+# @parameter paths [Array(String)] Paths to .mprof files.
+# @returns [Memory::Sampler] The loaded sampler with all allocations.
+def load(paths:)
+	sampler = Memory::Sampler.new
+	cache = sampler.cache
+	wrapper = sampler.wrapper
+	
+	total_size = paths.sum{|path| File.size(path)}
+	progress = Console.logger.progress(sampler, total_size)
+	
+	paths.each do |path|
+		Console.logger.info(sampler, "Loading #{path}, #{Memory.formatted_bytes File.size(path)}")
+		
+		File.open(path, 'r', encoding: Encoding::BINARY) do |io|
+			unpacker = wrapper.unpacker(io)
+			count = unpacker.read_array_header
+			
+			last_pos = 0
+			
+			# Read allocations directly into the sampler's array:
+			unpacker.each do |allocation|
+				sampler.allocated << allocation
+				
+				# Update progress based on bytes read:
+				current_pos = io.pos
+				progress.increment(current_pos - last_pos)
+				last_pos = current_pos
+			end
+		end
+		
+		Console.logger.info(sampler, "Loaded #{sampler.allocated.size} allocations")
+	end
+	
+	return sampler
+end
+
+# Dump a sampler to a .mprof file.
+#
+# @parameter input [Memory::Sampler] The sampler to dump.
+# @parameter output [String] Path to write the .mprof file.
+# @returns [Memory::Sampler] The input sampler.
+def dump(path, input:)
+	File.open(path, 'w', encoding: Encoding::BINARY) do |io|
+		input.dump(io)
+	end
+	
+	Console.logger.info(self, "Saved sampler to #{path} (#{File.size(path)} bytes)")
+	
+	return input
+end
+
+# Load a sampler from an ObjectSpace heap dump.
+#
+# @parameter path [String] Path to the heap dump JSON file.
+# @returns [Memory::Sampler] A sampler populated with allocations from the heap dump.
+def load_object_space_dump(path)
+	file_size = File.size(path)
+	progress = Console.logger.progress(self, file_size)
+	
+	Console.logger.info(self, "Loading heap dump from #{path} (#{Memory.formatted_bytes(file_size)})")
+	
+	sampler = nil
+	File.open(path, 'r') do |io|
+		sampler = Memory::Sampler.load_object_space_dump(io) do |line_count, object_count|
+			# Update progress based on bytes read:
+			progress.increment(io.pos - progress.current)
+		end
+	end
+	
+	Console.logger.info(self, "Loaded #{sampler.allocated.size} objects from heap dump")
+	
+	return sampler
+end

lib/memory/report.rb

@@ -90,10 +90,16 @@ def as_json(options = nil)
 			}
 		end
 
-		# Convert this report to a JSON string.
-		# @returns [String] JSON representation of this report.
-		def to_json(...)
-			as_json.to_json(...)
-		end
+	# Convert this report to a JSON string.
+	# @returns [String] JSON representation of this report.
+	def to_json(...)
+		as_json.to_json(...)
+	end
+
+	# Generate a human-readable representation of this report.
+	# @returns [String] Summary showing allocated and retained totals.
+	def inspect
+		"#<#{self.class}: #{@total_allocated} allocated, #{@total_retained} retained>"
 	end
 end
+end

lib/memory/sampler.rb

@@ -15,6 +15,7 @@
 
 require "objspace"
 require "msgpack"
+require "json"
 require "console"
 
 require_relative "cache"
@@ -65,6 +66,72 @@ def self.unpack(cache, fields)
 	# end
 	# ~~~
 	class Sampler
+	# Load allocations from an ObjectSpace heap dump.
+	#
+	# If a block is given, it will be called periodically with progress information.
+	#
+	# @parameter io [IO] The IO stream containing the heap dump JSON.
+	# @yields [line_count, object_count] Progress callback with current line and object counts.
+	# @returns [Sampler] A new sampler populated with allocations from the heap dump.
+	def self.load_object_space_dump(io, &block)
+		sampler = new
+		cache = sampler.cache
+		
+		line_count = 0
+		object_count = 0
+		report_interval = 10000
+		
+		io.each_line do |line|
+			line_count += 1
+			
+			begin
+				object = JSON.parse(line)
+			rescue JSON::ParserError
+				# Skip invalid JSON lines
+				next
+			end
+			
+			# Skip non-object entries (ROOT, SHAPE, etc.)
+			next unless object['address']
+			
+			# Get allocation information (may be nil if tracing wasn't enabled)
+			file = object['file'] || '(unknown)'
+			line_number = object['line'] || 0
+			
+			# Get object type/class
+			type = object['type'] || 'unknown'
+			
+			# Get memory size
+			memsize = object['memsize'] || 0
+			
+			# Get value for strings
+			value = object['value']
+			
+			allocation = Allocation.new(
+				cache,
+				type,           # class_name
+				file,           # file
+				line_number,    # line
+				memsize,        # memsize
+				value,          # value (for strings)
+				true            # retained (all objects in heap dump are live)
+			)
+			
+			sampler.allocated << allocation
+			object_count += 1
+			
+			# Report progress periodically
+			if block && (object_count % report_interval == 0)
+				block.call(line_count, object_count)
+			end
+		end
+		
+		# Final progress report
+		block.call(line_count, object_count) if block
+					
+		return sampler
+	end
+		
 		# Initialize a new sampler.
 		# @parameter filter [Block | Nil] Optional filter block to select which allocations to track.
 		def initialize(&filter)