Show a reply_count that matches what the user sees. Missing specs and polishing.

app/controllers/api/v1/statuses_controller.rb

@@ -14,16 +14,6 @@ class Api::V1::StatusesController < Api::BaseController
   override_rate_limit_headers :create, family: :statuses
   override_rate_limit_headers :update, family: :statuses
 
-  # This API was originally unlimited, pagination cannot be introduced without
-  # breaking backwards-compatibility. Arbitrarily high number to cover most
-  # conversations as quasi-unlimited, it would be too much work to render more
-  # than this anyway
-  CONTEXT_LIMIT = 4_096
-
-  # This remains expensive and we don't want to show everything to logged-out users
-  ANCESTORS_LIMIT         = 40
-  DESCENDANTS_LIMIT       = 60
-  DESCENDANTS_DEPTH_LIMIT = 20
 
   def index
     @statuses = preload_collection(@statuses, Status)
@@ -39,25 +29,10 @@ class Api::V1::StatusesController < Api::BaseController
   def context
     cache_if_unauthenticated!
 
-    ancestors_limit         = CONTEXT_LIMIT
-    descendants_limit       = CONTEXT_LIMIT
-    descendants_depth_limit = nil
+    @status_tree = StatusTree.new(status: @status, account: current_account)
+    @status_node = @status_tree.find_node(@status.id)
 
-    if current_account.nil?
-      ancestors_limit         = ANCESTORS_LIMIT
-      descendants_limit       = DESCENDANTS_LIMIT
-      descendants_depth_limit = DESCENDANTS_DEPTH_LIMIT
-    end
-
-    ancestors_results   = @status.in_reply_to_id.nil? ? [] : @status.ancestors(ancestors_limit, current_account)
-    descendants_results = @status.descendants(descendants_limit, current_account, descendants_depth_limit)
-    loaded_ancestors    = preload_collection(ancestors_results, Status)
-    loaded_descendants  = preload_collection(descendants_results, Status)
-
-    @context = Context.new(ancestors: loaded_ancestors, descendants: loaded_descendants)
-    statuses = [@status] + @context.ancestors + @context.descendants
-
-    render json: @context, serializer: REST::ContextSerializer, relationships: StatusRelationshipsPresenter.new(statuses, current_user&.account_id)
+    render json: @status_node, serializer: REST::StatusTreeNodeSerializer, relationships: StatusRelationshipsPresenter.new(@status_tree.flatten, current_user&.account_id)
   end
 
   def create

app/lib/activitypub/tag_manager.rb

@@ -37,6 +37,8 @@ class ActivityPub::TagManager
     return target.uri if target.respond_to?(:local?) && !target.local?
 
     case target.object_type
+    when :status
+      return activity_account_status_url(target.account, target) if target.reblog?
     when :person
       target.instance_actor? ? instance_actor_url : account_url(target)
     when :note, :comment, :activity

app/models/context.rb

@@ -1,5 +0,0 @@
-# frozen_string_literal: true
-
-class Context < ActiveModelSerializers::Model
-  attributes :ancestors, :descendants
-end

app/models/status.rb

@@ -292,18 +292,6 @@ class Status < ApplicationRecord
     end.take(MEDIA_ATTACHMENTS_LIMIT)
   end
 
-  def replies_count
-    status_stat&.replies_count || 0
-  end
-
-  def reblogs_count
-    status_stat&.reblogs_count || 0
-  end
-
-  def favourites_count
-    status_stat&.favourites_count || 0
-  end
-
   # Reblogs count received from an external instance
   def untrusted_reblogs_count
     status_stat&.untrusted_reblogs_count unless local?
@@ -393,6 +381,7 @@ class Status < ApplicationRecord
   def status_stat
     super || build_status_stat
   end
+  delegate :replies_count, :reblogs_count, :favourites_count, to: :status_stat
 
   def discard_with_reblogs
     discard_time = Time.current

app/models/status_tree.rb

@@ -0,0 +1,144 @@
+class StatusTree < ActiveModelSerializers::Model
+  include PreloadingConcern
+
+  # This API was originally unlimited, pagination cannot be introduced without
+  # breaking backwards-compatibility. Arbitrarily high number to cover most
+  # conversations as quasi-unlimited, it would be too much work to render more
+  # than this anyway
+  MAX_COUNT = 4_096
+
+  # This remains expensive and we don't want to show everything to logged-out users
+  ANCESTORS_MAX_COUNT   = 40
+  DESCENDANTS_MAX_COUNT = 60
+  DESCENDANTS_MAX_DEPTH = 20  
+
+  attributes :status, :account, :tree
+
+  class Node < ActiveModelSerializers::Model
+    attributes :status, :tree
+
+    delegate_missing_to :status
+
+    delegate :id, to: :status
+
+    def object_type = :status
+
+    def ancestors
+      tree.ancestors_for(id)
+    end
+
+    def descendants
+      tree.descendants_for(id)
+    end
+
+    def children
+      tree.children_for(id)
+    end
+
+    def replies_count
+      children.size
+    end
+
+    def ==(other)
+      other.class.in?([Node, Status]) && id == other.id
+    end
+
+    def inspect
+      "#<StatusTree::Node id: #{id}, parent_id: #{in_reply_to_id || 'nil'}>"
+    end
+  end
+
+  def tree
+    @tree ||= begin
+      ancestors = preload_collection(status.in_reply_to_id.nil? ? [] : status.ancestors(ancestors_max_count, account), Status)
+      descendants = preload_collection(status.descendants(descendants_max_count, account, descendants_max_depth), Status)
+      all_nodes = (ancestors + [status] + descendants).map { |status| Node.new(status:, tree: self) }
+      build_tree_from(all_nodes)
+    end
+  end
+
+  def subtree_for(id, subtree = tree)
+    subtree.each do |node, children|
+      return children if node.id == id
+
+      found = subtree_for(id, children)
+      return found if found
+    end
+    nil
+  end
+
+  def flatten
+    collect_descendants(tree)
+  end
+
+  delegate :each, :flat_map, :keys, to: :tree
+
+  def inspect
+    "#<StatusTree #{tree.inspect}>"
+  end
+
+  def find_node(id, subtree = tree)
+    subtree.each do |node, children|
+      return node if node.id == id
+  
+      result = find_node(id, children)
+      return result if result
+    end
+  end
+
+  def ancestors_for(id)
+    ancestors = []
+    node = find_node(id)
+    parent_id = node.in_reply_to_id
+
+    while parent_id
+      parent_node = find_node(parent_id)
+      break unless parent_node
+      ancestors << parent_node
+      parent_id = parent_node.in_reply_to_id
+    end
+
+    ancestors.reverse
+  end
+
+  def descendants_for(id)
+    subtree = subtree_for(id)
+    return [] unless subtree
+
+    collect_descendants(subtree)
+  end
+
+  def children_for(id)
+    subtree = subtree_for(id)
+
+    subtree.keys
+  end
+
+  private
+
+  def build_tree_from(nodes, parent_id = nil)
+    grouped_nodes = nodes.group_by(&:in_reply_to_id)
+
+    (grouped_nodes[parent_id] || []).each_with_object({}) do |node, tree|
+      tree[node] = build_tree_from(nodes - [node], node.id)
+    end
+  end
+
+  def descendants_max_depth
+    account.nil? ? DESCENDANTS_MAX_DEPTH : nil
+  end
+
+  def descendants_max_count
+    account.nil? ? DESCENDANTS_MAX_COUNT : MAX_COUNT
+  end
+
+  def ancestors_max_count
+    account.nil? ? ANCESTORS_MAX_COUNT : MAX_COUNT
+  end
+
+  def collect_descendants(subtree)
+    subtree.flat_map do |node, children|
+      [node] + collect_descendants(children)
+    end
+  end  
+end

app/serializers/rest/status_tree_node_serializer.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-class REST::ContextSerializer < ActiveModel::Serializer
+class REST::StatusTreeNodeSerializer < ActiveModel::Serializer
   has_many :ancestors,   serializer: REST::StatusSerializer
   has_many :descendants, serializer: REST::StatusSerializer
 end