paper_trail_diff 0.12.0 → 0.13.0

data/lib/paper_trail_diff/nested_comparator.rb

@@ -27,10 +27,16 @@ module PaperTrailDiff
   # one of them, so absence is its own value rather than nil.
   class NestedComparator
     # Stands in for a key that was not there at all.
-    ABSENT = Object.new
-    def ABSENT.inspect = '#<PaperTrailDiff absent>'
-    def ABSENT.to_s = 'absent'
-    ABSENT.freeze
+    class Absent
+      #: () -> String
+      def inspect = '#<PaperTrailDiff absent>'
+
+      #: () -> String
+      def to_s = 'absent'
+    end
+    private_constant :Absent
+
+    ABSENT = Absent.new.freeze
 
     #: (untyped, untyped) -> nested_changes
     def self.call(from_value, to_value)
@@ -49,9 +55,9 @@ module PaperTrailDiff
     #: () -> nested_changes
     def call
       from_structure, to_structure = structures
-      return {} unless from_structure && to_structure
-
       changes = {} #: nested_changes
+      return changes.freeze unless from_structure && to_structure
+
       walk(from_structure, to_structure, [], changes)
       changes.freeze
     end
@@ -76,20 +82,23 @@ module PaperTrailDiff
 
       parsed = begin
         JSON.parse(value)
-      rescue JSON::ParserError, TypeError
+      rescue JSON::ParserError
         nil
       end
-      parsed if parsed.is_a?(Hash)
+      # JSON parses objects into plain Hashes, never a subclass.
+      parsed if parsed.instance_of?(Hash)
     end
 
     #: (Hash[untyped, untyped], Hash[untyped, untyped], Array[String], nested_changes) -> void
     def walk(from_hash, to_hash, path, changes)
+      from_hash = normalize_keys(from_hash)
+      to_hash = normalize_keys(to_hash)
       keys(from_hash, to_hash).each do |key|
-        from_item = from_hash.key?(key) ? from_hash[key] : ABSENT
-        to_item = to_hash.key?(key) ? to_hash[key] : ABSENT
+        from_item = from_hash.fetch(key, ABSENT)
+        to_item = to_hash.fetch(key, ABSENT)
         next if from_item == to_item
 
-        here = [*path, key.to_s].freeze
+        here = (path + [key]).freeze
         if from_item.is_a?(Hash) && to_item.is_a?(Hash)
           walk(from_item, to_item, here, changes)
         else
@@ -98,6 +107,20 @@ module PaperTrailDiff
       end
     end
 
+    #: (Hash[untyped, untyped]) -> Hash[String, untyped]
+    def normalize_keys(hash)
+      normalized = {} #: Hash[String, untyped]
+      hash.each do |key, value|
+        name = key.to_s
+        if normalized.key?(name)
+          raise AmbiguousNestedKeyError, "duplicate nested key after normalization: #{name.inspect}"
+        end
+
+        normalized[name] = value
+      end
+      normalized
+    end
+
     # Two arrays are a membership change; anything else is a change to the
     # value as a whole, including an array facing something that is not one.
     #: (untyped, untyped) -> nested_change
@@ -109,11 +132,12 @@ module PaperTrailDiff
       end
     end
 
-    # Sorted so a report reads the same twice, and stringified because a hash
-    # loaded from JSON and one built in Ruby can key the same field differently.
-    #: (Hash[untyped, untyped], Hash[untyped, untyped]) -> Array[untyped]
+    # Sorted so a report reads the same twice. Both hashes are already keyed by
+    # strings, because a hash loaded from JSON and one built in Ruby can key the
+    # same field differently.
+    #: (Hash[String, untyped], Hash[String, untyped]) -> Array[String]
     def keys(from_hash, to_hash)
-      (from_hash.keys | to_hash.keys).sort_by(&:to_s)
+      (from_hash.keys | to_hash.keys).sort
     end
   end
 end

data/lib/paper_trail_diff/paper_trail_adapter.rb

@@ -10,11 +10,12 @@ module PaperTrailDiff
     # The only thing a wall-clock window can close on besides a later version.
     CLOSE_ON_CURRENT = :current
 
-    #: (associations: Array[String | Symbol], ignore: ignore_option, ?reload_live_endpoints: bool) -> void
-    def initialize(associations:, ignore:, reload_live_endpoints: true)
+    #: (associations: Array[String | Symbol], ignore: ignore_option, ?reload_live_endpoints: bool, ?allow_expensive: bool) -> void
+    def initialize(associations:, ignore:, reload_live_endpoints: true, allow_expensive: false)
       @association_tree = AssociationTree.build(associations)
       @ignore_policy = IgnorePolicy.build(ignore, association_paths: @association_tree.paths)
       @traversal = AssociationTraversal.new(@association_tree)
+      @batch_safety = BatchSafety.new(@traversal, allow_expensive: allow_expensive)
       @traversal_preparer = TraversalPreparer.new(tree: @association_tree, traversal: @traversal)
       @live_endpoints = LiveEndpointProvider.new(
         tree: @association_tree, traversal: @traversal, reload: reload_live_endpoints
@@ -34,7 +35,14 @@ module PaperTrailDiff
       Instrumentation.instrument('compare', payload) do
         Endpoint.validate_pair!(from_endpoint, to_endpoint)
         notify_ambiguous_association_boundary(from_endpoint, to_endpoint)
-        Engine.compare(snapshot_for_endpoint(from_endpoint), snapshot_for_endpoint(to_endpoint))
+        if @association_tree.empty?
+          Engine.compare(snapshot_for_endpoint(from_endpoint), snapshot_for_endpoint(to_endpoint))
+        else
+          # Prepared history orders versions that share a timestamp; reading
+          # associations straight from PT-AT would resolve them by time alone.
+          comparison_batch([{ from: from_endpoint, to: to_endpoint }],
+                           preparer: @traversal_preparer.method(:call)).call.values.fetch(0)
+        end
       end
     end
 
@@ -44,19 +52,14 @@ module PaperTrailDiff
       count = comparisons.is_a?(Array) ? comparisons.length : 0
       payload = @instrumentation_payload.merge(comparison_count: count)
       Instrumentation.instrument('compare_many', payload) do
-        ComparisonBatch.new(
-          comparisons,
-          live_loader: @live_endpoints.method(:call), preparer: @traversal_preparer.method(:call),
-          history_preparer: @historical_store.method(:prepare_batch),
-          historical_snapshotter: method(:historical_snapshot),
-          live_normalizer: method(:normalize_live_snapshot)
-        ).call
+        comparison_batch(comparisons, preparer: method(:prepare_batch_model)).call
       end
     end
 
     #: (untyped, from: untyped, to: untyped, within: untyped, ?version_scope: untyped, ?close_on: Symbol?) -> Array[Step]
     def timeline(record, from:, to:, within:, version_scope: nil, close_on: nil) # rubocop:disable Metrics/ParameterLists
       @traversal_preparer.call(record.class, historical: true)
+      from, to = default_range(record, from, to, within)
       TimelineBuilder.new(
         record,
         from: from,
@@ -73,8 +76,8 @@ module PaperTrailDiff
       payload = @instrumentation_payload.merge(model_type: record.class.base_class.name.to_s)
       Instrumentation.instrument('activity_timeline', payload) do
         @traversal_preparer.call(record.class, historical: true)
+        from, to = default_range(record, from, to, within)
         live = live_endpoint_for(record, close_on, within)
-        reject_live_habtm_activity!(record.class) if Endpoint.record?(to) || live
         steps = activity_builder(
           record, from: from, to: to, within: within, version_scope: version_scope,
                   live_endpoint: live, snapshots: snapshots, group: group
@@ -86,7 +89,9 @@ module PaperTrailDiff
 
     #: (untyped, from: untyped, to: untyped, within: untyped, ?activity: bool, ?version_scope: untyped, ?close_on: Symbol?, ?snapshots: bool, ?group: Symbol?) -> Analysis
     def analyze(record, from:, to:, within:, activity: false, version_scope: nil, close_on: nil, snapshots: false, group: nil) # rubocop:disable Metrics/ParameterLists, Layout/LineLength
+      ActivityGrouping.validate!(group, activity: activity, snapshots: snapshots)
       @traversal_preparer.call(record.class, historical: true)
+      from, to = default_range(record, from, to, within)
       live = live_endpoint_for(record, close_on, within)
       if activity
         return analyze_activity(
@@ -102,19 +107,21 @@ module PaperTrailDiff
     end
 
     # Analyzes many roots over one shared range, preparing their history once.
-    #: (Array[untyped], within: untyped, ?activity: bool, ?version_scope: untyped, ?close_on: Symbol?) -> Hash[identity, Analysis]
-    def analyze_many(records, within:, activity: false, version_scope: nil, close_on: nil)
+    #: (Array[untyped], within: untyped, ?activity: bool, ?version_scope: untyped, ?close_on: Symbol?, ?snapshots: bool, ?group: Symbol?) -> Hash[identity, Analysis]
+    def analyze_many(records, within:, activity: false, version_scope: nil, close_on: nil, snapshots: false, group: nil) # rubocop:disable Metrics/ParameterLists, Layout/LineLength
+      ActivityGrouping.validate!(group, activity: activity, snapshots: snapshots)
       count = records.is_a?(Array) ? records.length : 0
       payload = @instrumentation_payload.merge(comparison_count: count)
       Instrumentation.instrument('analyze_many', payload) do
         AnalysisBatch.new(
           records,
-          time_range: within.nil? ? nil : TimeRange.new(within),
-          version_scope: version_scope,
-          close_on_current: close_on_current?(close_on, within),
-          live_loader: @live_endpoints.method(:call),
+          time_range: within.nil? ? nil : TimeRange.new(within), version_scope: version_scope,
+          close_on_current: close_on_current?(close_on, within) || within.nil?,
+          context_required: activity && !@association_tree.empty?,
+          live_loader: method(:load_batch_endpoints),
           history_preparer: @historical_store.method(:prepare_batch),
-          analyzer: batched_root_analyzer(activity)
+          analyzer: batched_root_analyzer(activity, within: within, snapshots: snapshots,
+                                                    group: group)
         ).call
       end
     end
@@ -123,22 +130,65 @@ module PaperTrailDiff
     # caller reporting on a population does not have to rediscover which of its
     # members changed. The roots the relation could not reach come back named
     # rather than dropped -- see `PaperTrailDiff.analyze_scope`.
-    #: (untyped, limit: Integer?, within: untyped, ?activity: bool, ?version_scope: untyped, ?close_on: Symbol?) -> ScopedAnalysis
-    def analyze_scope(scope, limit:, within:, activity: false, version_scope: nil, close_on: nil) # rubocop:disable Metrics/ParameterLists
-      raise ConfigurationError, 'limit: is required when selecting roots by scope' if limit.nil?
-
-      selection = ScopedRootSelection.new(
-        scope, time_range: within.nil? ? nil : TimeRange.new(within), limit: limit
+    #: (untyped, limit: Integer?, within: untyped, ?activity: bool, ?version_scope: untyped, ?close_on: Symbol?, ?snapshots: bool, ?group: Symbol?, ?historical_filter: untyped, ?max_filter_versions: Integer, ?include_deleted: bool) -> ScopedAnalysis
+    def analyze_scope(scope, limit:, within:, activity: false, version_scope: nil, close_on: nil, snapshots: false, group: nil, historical_filter: nil, max_filter_versions: HistoricalFilter::DEFAULT_VERSION_LIMIT, include_deleted: false) # rubocop:disable Metrics/ParameterLists, Layout/LineLength
+      ActivityGrouping.validate!(group, activity: activity, snapshots: snapshots)
+      @batch_safety.call(scope.all.model) if scope.respond_to?(:all) && scope.respond_to?(:where)
+      # @type var options: report_options
+      options = { within: within, activity: activity, version_scope: version_scope,
+                  close_on: close_on, snapshots: snapshots, group: group }
+      ScopedAnalysisBuilder.new(
+        scope, limit: limit, within: within, historical_filter: historical_filter,
+               max_filter_versions: max_filter_versions, include_deleted: include_deleted,
+               analyzer: ->(records) { analyze_many(records, **options, within: within) },
+               deleted_analyzer: lambda { |identities|
+                 analyze_deleted(scope.all.model, identities, options)
+               }
       ).call
-      ScopedAnalysis.new(
-        analyses: analyze_many(selection.records, within: within, activity: activity,
-                                                  version_scope: version_scope, close_on: close_on),
-        unreachable: selection.unreachable
-      )
     end
 
     private
 
+    #: (Array[comparison_input], preparer: untyped) -> ComparisonBatch
+    def comparison_batch(comparisons, preparer:)
+      ComparisonBatch.new(
+        comparisons,
+        live_loader: @live_endpoints, preparer: preparer,
+        history_preparer: @historical_store.method(:prepare_batch),
+        historical_snapshotter: method(:historical_snapshot),
+        live_normalizer: method(:normalize_live_snapshot)
+      )
+    end
+
+    #: (Array[untyped]) -> Hash[identity, untyped]
+    def load_batch_endpoints(records)
+      records.map(&:class).uniq.each { |model| @batch_safety.call(model) }
+      @live_endpoints.call(records)
+    end
+
+    #: (untyped, historical: bool) -> void
+    def prepare_batch_model(model, historical:)
+      @batch_safety.call(model, historical: historical)
+      @traversal_preparer.call(model, historical: historical)
+    end
+
+    # @rbs @batch_safety: BatchSafety
+
+    #: (untyped, Array[Array[String]], Hash[Symbol, untyped]) -> Hash[Array[String], Analysis?]
+    def analyze_deleted(model, identities, options)
+      within = options.fetch(:within)
+      analyzer = batched_root_analyzer(
+        options.fetch(:activity), within: within, snapshots: options.fetch(:snapshots),
+                                  group: options.fetch(:group)
+      )
+      DeletedAnalysisBatch.new(
+        model, identities: identities, time_range: within.nil? ? nil : TimeRange.new(within),
+               version_scope: options.fetch(:version_scope),
+               history_preparer: @historical_store.method(:prepare_deleted_batch),
+               analyzer: analyzer
+      ).call
+    end
+
     # @rbs @association_tree: AssociationTree
     # @rbs @live_endpoints: LiveEndpointProvider
     # @rbs @instrumentation_payload: Hash[Symbol, untyped]
@@ -153,7 +203,6 @@ module PaperTrailDiff
 
     #: (untyped, from: untyped, to: untyped, within: untyped, version_scope: untyped, live_endpoint: untyped, ?snapshots: bool, ?group: Symbol?) -> Analysis
     def analyze_activity(record, from:, to:, within:, version_scope:, live_endpoint:, snapshots: false, group: nil) # rubocop:disable Metrics/ParameterLists, Layout/LineLength
-      reject_live_habtm_activity!(record.class) if live_endpoint
       activity_builder(
         record, from: from, to: to, within: within, version_scope: version_scope,
                 live_endpoint: live_endpoint, snapshots: snapshots, group: group
@@ -181,25 +230,38 @@ module PaperTrailDiff
       )
     end
 
-    # `close_on:` names what ends a wall-clock window, so it is meaningless for a
-    # range whose endpoints the caller already gave explicitly.
+    # A window with nothing recorded after its last change ends at the live
+    # record, which holds the state that change produced -- whether the window
+    # ended a moment ago or long before. `close_on: :current` asked for this
+    # before it was the default and is still accepted. It names what ends a
+    # wall-clock window, so it is meaningless for explicit endpoints.
     #: (Symbol?, untyped) -> bool
     def close_on_current?(close_on, within)
-      return false if close_on.nil?
-      unless close_on == CLOSE_ON_CURRENT
+      unless close_on.nil? || close_on == CLOSE_ON_CURRENT
         raise ConfigurationError, "close_on: must be #{CLOSE_ON_CURRENT.inspect} or nil"
       end
-      raise ConfigurationError, 'close_on: requires `within`' if within.nil?
+      raise ConfigurationError, 'close_on: requires `within`' if close_on && within.nil?
 
-      true
+      !within.nil?
     end
 
-    # A destroyed root has no current state to close on, and its own destroy
-    # version already terminates the history.
+    # With neither a window nor endpoints, a range is the record's whole
+    # history through its current saved state; a missing endpoint takes that
+    # side of the same default.
+    #: (untyped, untyped, untyped, untyped) -> [untyped, untyped]
+    def default_range(record, from, to, within)
+      return [from, to] unless within.nil?
+
+      [from.nil? ? :first : from, to.nil? ? record : to]
+    end
+
+    # Only a saved row has current state to close on. A destroyed root's own
+    # destroy version already terminates its history, and a reified or unsaved
+    # instance stands for a state that is not in the table.
     #: (untyped, Symbol?, untyped) -> untyped
     def live_endpoint_for(record, close_on, within)
       return unless close_on_current?(close_on, within)
-      return unless Endpoint.record?(record) && !record.destroyed?
+      return unless Endpoint.record?(record) && record.persisted?
 
       record
     end
@@ -213,14 +275,14 @@ module PaperTrailDiff
       @activity_snapshotter = build_activity_snapshotter
     end
 
-    #: (bool) -> BatchedRootAnalyzer
-    def batched_root_analyzer(activity)
+    #: (bool, within: untyped, snapshots: bool, group: Symbol?) -> BatchedRootAnalyzer
+    def batched_root_analyzer(activity, within:, snapshots:, group:)
       BatchedRootAnalyzer.new(
         tree: @association_tree,
         timeline_snapshotter: @timeline_snapshotter,
         activity_snapshotter: @activity_snapshotter,
         preparer: @traversal_preparer.method(:call),
-        activity: activity
+        activity: activity, within: within, snapshots: snapshots, group: group
       )
     end
 
@@ -243,12 +305,13 @@ module PaperTrailDiff
         full_snapshotter: method(:snapshot_at),
         partial_snapshotter: @historical_store.method(:custom),
         association_reader: @historical_store.method(:association_reader),
-        record_transition: @historical_store.method(:record_transition)
+        record_transition: @historical_store.method(:record_transition),
+        prepared_record_after: @historical_store.method(:record_after)
       )
       ActivitySnapshotProvider.new(
-        snapshotter: method(:snapshot_at),
-        refresher: refresher,
-        preparer: @historical_store.method(:prepare)
+        snapshotter: method(:snapshot_at), refresher: refresher,
+        preparer: @historical_store.method(:prepare),
+        activity_index: @historical_store.method(:activity_index)
       )
     end
 
@@ -267,15 +330,6 @@ module PaperTrailDiff
       )
     end
 
-    #: (untyped) -> void
-    def reject_live_habtm_activity!(model_class)
-      paths = @traversal.habtm_paths(model_class)
-      return if paths.empty?
-
-      message = "live-ended activity cannot reconstruct HABTM event boundaries: #{paths.join(', ')}"
-      raise UnsupportedLiveActivityError, message
-    end
-
     #: (untyped) -> RecordSnapshot?
     def snapshot_for_endpoint(endpoint)
       Endpoint.version?(endpoint) ? historical_snapshot(endpoint) : live_snapshot(endpoint)

data/lib/paper_trail_diff/prepared_activity_index.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+# rbs_inline: enabled
+
+module PaperTrailDiff
+  # Shares association event discovery across every root in a prepared graph.
+  class PreparedActivityIndex
+    #: (PreparedRecordIndex, Array[untyped], start_at: untyped, end_at: untyped) -> void
+    def initialize(records, root_versions, start_at:, end_at:)
+      @records = records
+      @range = ActivityRange.new(start_at, end_at)
+      @roots = root_versions.group_by { |version| [version.item_type, version.item_id.to_s] }
+      @nodes = {} #: Hash[String, Array[untyped]]
+      @children = {} #: Hash[Array[String], untyped]
+      @rows = {} #: Hash[Array[String], Hash[String, Array[untyped]]]
+    end
+
+    #: (untyped, Array[untyped]) -> void
+    def register(model, ids)
+      name = model.base_class.name.to_s
+      @nodes[name] = (@nodes.fetch(name, []) | ids)
+    end
+
+    #: (untyped, Array[untyped]) -> Array[untyped]
+    def versions(model, ids)
+      ids.flat_map do |id|
+        key = [model.base_class.name.to_s, id.to_s]
+        @roots[key] || @records.versions_for_id(model, id)
+      end
+    end
+
+    #: (untyped, Array[untyped], untyped, ActivityRange) -> Array[untyped]
+    def child_ids(model, ids, reflection, range)
+      key = [model.base_class.name.to_s, reflection.name.to_s]
+      historical, live = @children[key] ||= child_rows(model, reflection)
+      empty = [] #: Array[untyped]
+      selected = historical_ids(reflection.klass, ids, historical, range)
+      selected |= ids.flat_map { |id| live.fetch(id.to_s, empty) } if range.current_end?
+      selected.uniq
+    end
+
+    #: (untyped, Array[untyped], untyped, ?transactions: bool) -> Array[untyped]
+    def association_rows(model, selected_versions, reflection, transactions: false)
+      key = [model.base_class.name.to_s, reflection.name.to_s]
+      rows = @rows[key] ||= load_association_rows(model, reflection, transactions)
+      version_ids(selected_versions, transactions).flat_map { |id| rows.fetch(id.to_s, []) }
+    end
+
+    private
+
+    # @rbs @records: PreparedRecordIndex
+    # @rbs @range: ActivityRange
+    # @rbs @roots: Hash[Array[String], Array[untyped]]
+    # @rbs @nodes: Hash[String, Array[untyped]]
+    # @rbs @children: Hash[Array[String], untyped]
+    # @rbs @rows: Hash[Array[String], Hash[String, Array[untyped]]]
+
+    #: (untyped, Array[untyped], untyped, ActivityRange) -> Array[untyped]
+    def historical_ids(model, ids, historical, range)
+      empty = [] #: Array[untyped]
+      ids.flat_map do |id|
+        historical.fetch(id.to_s, empty).filter_map do |row|
+          row[1] if eligible_child?(model, row, range)
+        end
+      end
+    end
+
+    #: (untyped) -> Array[untyped]
+    def node_ids(model)
+      name = model.base_class.name.to_s
+      @nodes.fetch(name) do
+        raise ConfigurationError,
+              "prepared activity graph is missing #{name}; " \
+              'register the model before collecting activity'
+      end
+    end
+
+    #: (untyped, untyped) -> untyped
+    def child_rows(model, reflection)
+      loader = ActivityChildCandidateLoader.new(
+        model, node_ids(model), reflection, @range
+      )
+      [loader.historical_rows.group_by { |row| row.first.to_s },
+       loader.current_pairs.group_by { |row| row.first.to_s }
+             .transform_values { |rows| rows.map(&:last) }]
+    end
+
+    # Mirrors VersionAssociationCandidateScope's before-window rule. The loaded
+    # series includes changes to other owners too, so an old membership cannot
+    # survive a later pre-window reassignment just because that owner is absent.
+    #: (untyped, Array[untyped], ActivityRange) -> bool
+    def eligible_child?(model, row, range)
+      _owner, id, timestamp, event = row
+      return timestamp <= range.end_time if timestamp >= range.start_time
+      return false if event == 'destroy'
+
+      later = @records.versions_for_id(model, id)
+                      .bsearch { |version| version.created_at > timestamp }
+      !later || later.created_at >= range.start_time
+    end
+
+    #: (untyped, untyped, bool) -> Hash[String, Array[untyped]]
+    def load_association_rows(model, reflection, transactions)
+      all_versions = versions(model, node_ids(model))
+      foreign_key = transactions ? reflection.name.to_s : reflection.foreign_key.to_s
+      model.paper_trail.version_association_class.where(
+        version_id: version_ids(all_versions, transactions), foreign_key_name: foreign_key
+      ).to_a.group_by { |row| row.version_id.to_s }
+    end
+
+    #: (Array[untyped], bool) -> Array[untyped]
+    def version_ids(versions, transactions)
+      if transactions
+        versions.filter_map do |version|
+          version.transaction_id if version.respond_to?(:transaction_id)
+        end.uniq
+      else
+        versions.map(&:id)
+      end
+    end
+  end
+end

data/lib/paper_trail_diff/prepared_edge_loader.rb

@@ -50,11 +50,15 @@ module PaperTrailDiff
     def belongs_to_groups(owner_class, owner_ids, reflection)
       key = scalar_foreign_key(reflection)
       return {} unless key
+      unless reflection.polymorphic?
+        return group_for(reflection.klass,
+                         @records.attribute_values_for(owner_class, owner_ids, key))
+      end
 
       grouped = {} #: Hash[String, Hash[Symbol, untyped]]
       @records.records_for(owner_class, owner_ids).each do |owner|
         id = owner.public_send(key)
-        model_class = target_class(reflection, owner)
+        model_class = AssociationTraversal.belongs_to_target(reflection, owner)
         add_identity(grouped, model_class, id) if model_class && id
       end
       grouped
@@ -109,7 +113,12 @@ module PaperTrailDiff
         association_class.arel_table[:foreign_key_id],
         version_class.arel_table[:item_id]
       )
-      group_ids_by_owner(pairs)
+      group_ids_by_owner(typed_child_ids(reflection.klass, pairs))
+    end
+
+    #: (untyped, Array[Array[untyped]]) -> Array[Array[untyped]]
+    def typed_child_ids(model_class, pairs)
+      pairs.map { |owner_id, child_id| [owner_id, Support.record_id(model_class, child_id)] }
     end
 
     #: (untyped, Array[untyped], untyped, String) -> untyped
@@ -139,15 +148,6 @@ module PaperTrailDiff
       group_ids_by_owner(relation.pluck(foreign_key, reflection.klass.primary_key))
     end
 
-    #: (untyped, untyped) -> untyped
-    def target_class(reflection, owner)
-      return reflection.klass unless reflection.polymorphic?
-
-      Object.const_get(owner.public_send(reflection.foreign_type).to_s)
-    rescue NameError
-      nil
-    end
-
     #: (Hash[String, Hash[Symbol, untyped]], untyped, untyped) -> void
     def add_identity(groups, model_class, id)
       group = groups[model_class.name.to_s] ||= { model: model_class, ids: [] }

data/lib/paper_trail_diff/prepared_history.rb

@@ -4,17 +4,21 @@
 module PaperTrailDiff
   # Batched record states and relationship candidates for one activity range.
   class PreparedHistory
+    attr_reader :activity_index #: PreparedActivityIndex?
+
     DIRECT_RESOLVERS = {
       belongs_to: :resolve_belongs_to,
       has_one: :resolve_children,
       has_many: :resolve_children,
-      has_and_belongs_to_many: :resolve_habtm
+      has_and_belongs_to_many: :habtm_records
     }.freeze
     private_constant :DIRECT_RESOLVERS
 
-    #: (PreparedRecordIndex) -> void
-    def initialize(records)
+    #: (PreparedRecordIndex, ?activity_index: PreparedActivityIndex?) -> void
+    def initialize(records, activity_index: nil)
       @records = records
+      @state_reader = PreparedStateAssociationReader.new(records, method(:edge))
+      @activity_index = activity_index
       @edges = {} #: Hash[Array[String], Hash[String, Hash[Symbol, untyped]]]
       @unsupported = {} #: Hash[Array[String], bool]
       @habtm = {} #: Hash[Array[String], Hash[String, Array[untyped]]]
@@ -24,7 +28,7 @@ module PaperTrailDiff
     def merge_edge(owner_class, reflection, groups)
       key = edge_key(owner_class, reflection)
       empty = {} #: Hash[String, Hash[Symbol, untyped]]
-      @edges[key] = merge_groups(@edges.fetch(key, empty), groups)
+      @edges[key] = Support.merge_record_groups(@edges.fetch(key, empty), groups)
     end
 
     #: (untyped, untyped) -> Hash[String, Hash[Symbol, untyped]]
@@ -48,6 +52,16 @@ module PaperTrailDiff
       @records.transition(model_class, id, version)
     end
 
+    #: (untyped, untyped, untyped, ?direct: bool) -> [bool, untyped]
+    def record_after(model_class, id, version, direct: false)
+      @records.record_after(model_class, id, version, direct: direct)
+    end
+
+    #: (PreparedRecordState, untyped, untyped) -> [bool, Array[PreparedRecordState]]
+    def snapshot_states(owner, reflection, boundary)
+      @state_reader.call(owner, reflection, boundary)
+    end
+
     #: (untyped, untyped, untyped, habtm_boundary: untyped) -> [bool, Array[untyped]]
     def resolve(record, reflection, boundary, habtm_boundary:)
       return [false, []] unless supported?(record.class, reflection)
@@ -63,6 +77,7 @@ module PaperTrailDiff
     private
 
     # @rbs @records: PreparedRecordIndex
+    # @rbs @state_reader: PreparedStateAssociationReader
     # @rbs @edges: Hash[Array[String], Hash[String, Hash[Symbol, untyped]]]
     # @rbs @unsupported: Hash[Array[String], bool]
     # @rbs @habtm: Hash[Array[String], Hash[String, Array[untyped]]]
@@ -70,12 +85,7 @@ module PaperTrailDiff
     #: (untyped, untyped) -> bool
     def supported?(owner_class, reflection)
       !@unsupported[edge_key(owner_class, reflection)] &&
-        reflection.scope.nil? && !default_scoped?(reflection.klass)
-    end
-
-    #: (untyped) -> bool
-    def default_scoped?(model_class)
-      model_class.default_scopes?
+        reflection.scope.nil? && !reflection.klass.default_scopes?
     end
 
     #: (untyped, untyped, untyped, untyped) -> Array[untyped]?
@@ -94,11 +104,6 @@ module PaperTrailDiff
       child_records(record, reflection, boundary)
     end
 
-    #: (untyped, untyped, untyped, untyped) -> Array[untyped]?
-    def resolve_habtm(record, reflection, boundary, habtm_boundary)
-      habtm_records(record, reflection, boundary, habtm_boundary)
-    end
-
     #: (untyped, untyped, untyped) -> Array[untyped]?
     def belongs_to_records(record, reflection, boundary)
       foreign_key = scalar_foreign_key(reflection)
@@ -107,7 +112,7 @@ module PaperTrailDiff
       id = record.public_send(foreign_key)
       return [] if id.nil?
 
-      model_class = belongs_to_class(record, reflection)
+      model_class = AssociationTraversal.belongs_to_target(reflection, record)
       return unless model_class && @records.loaded?(model_class, id)
 
       Array(@records.record_before(model_class, id, boundary)).compact
@@ -189,15 +194,6 @@ module PaperTrailDiff
       key.to_s unless key.is_a?(Array)
     end
 
-    #: (untyped, untyped) -> untyped
-    def belongs_to_class(record, reflection)
-      return reflection.klass unless reflection.polymorphic?
-
-      Object.const_get(record.public_send(reflection.foreign_type).to_s)
-    rescue NameError
-      nil
-    end
-
     #: (untyped, untyped) -> Array[String]
     def owner_types(reflection, owner)
       types = [owner.class.name, owner.class.base_class.name, reflection.active_record.name]
@@ -218,10 +214,5 @@ module PaperTrailDiff
     def edge_key(owner_class, reflection)
       [owner_class.base_class.name.to_s, reflection.name.to_s]
     end
-
-    #: (Hash[String, Hash[Symbol, untyped]], Hash[String, Hash[Symbol, untyped]]) -> Hash[String, Hash[Symbol, untyped]]
-    def merge_groups(existing, incoming)
-      Support.merge_record_groups(existing, incoming)
-    end
   end
 end