Merge 26ac75dac4 into fbe9728f36

app/models/status.rb

@@ -289,18 +289,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?
@@ -386,6 +374,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,135 @@
+class StatusTree < ActiveModelSerializers::Model
+  include PreloadingConcern
+
+  MAX_COUNT = 4_096
+
+  attributes :status, :account, :tree
+
+  class Node < ActiveModelSerializers::Model
+    attributes :status, :tree
+
+    delegate :id, to: :status
+
+    delegate_missing_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 ==(other)
+      other.class.in?([Node, Status]) && id == other.id
+    end
+
+    def inspect
+      "#<StatusTree::Node id: #{id}, in_reply_to_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 status_node
+    find_node(status.id)
+  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)
+    in_reply_to_id = node.in_reply_to_id
+
+    while in_reply_to_id
+      parent_node = find_node(in_reply_to_id)
+      break unless parent_node
+      ancestors << parent_node
+      in_reply_to_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, in_reply_to_id = nil)
+    grouped_nodes = nodes.group_by(&:in_reply_to_id)
+
+    (grouped_nodes[in_reply_to_id] || []).each_with_object({}) do |node, tree|
+      tree[node] = build_tree_from(nodes - [node], node.id)
+    end
+  end
+
+  def descendants_max_depth
+    nil
+  end
+
+  def descendants_max_count
+    MAX_COUNT
+  end
+
+  def 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_serializer.rb

@@ -92,6 +92,10 @@ class REST::StatusSerializer < ActiveModel::Serializer
     object.untrusted_favourites_count || relationships&.attributes_map&.dig(object.id, :favourites_count) || object.favourites_count
   end
 
+  def replies_count
+    StatusTree.new(status: object, account: current_user&.account).status_node.children.size
+  end
+
   def favourited
     if relationships
       relationships.favourites_map[object.id] || false

spec/serializers/rest/status_serializer_spec.rb

@@ -72,4 +72,77 @@ RSpec.describe REST::StatusSerializer do
       end
     end
   end
+
+  describe '#replies_count' do
+    let(:author) { alice }
+    let(:replier) { bob }
+    let!(:status) { Fabricate(:status, account: author, visibility: :public) }
+
+    context 'when being presented to the account that posted the status' do
+      let(:current_user) { Fabricate(:user, account: author) }
+
+      before do
+        Fabricate(:follow, account: replier, target_account: author)
+        Fabricate(:follow, account: author, target_account: replier)
+      end
+
+      context 'when the status has follower-only replies' do
+        let(:reply) { Fabricate(:status, in_reply_to_id: status.id, account: replier, visibility: :private) }
+
+        before do
+          reply
+        end
+
+        it 'counts 1 reply' do
+          expect(subject['replies_count']).to eq(1)
+        end
+
+        context 'when one of the replies has subsequent replies' do
+          before do
+            Fabricate(:status, in_reply_to_id: reply.id, account: author, visibility: :private)
+          end
+
+          it 'does not count that reply' do
+            expect(subject['replies_count']).to eq 1
+          end
+        end
+      end
+    end
+
+    context 'when being presented to a different account' do
+      let(:current_user) { Fabricate(:user) }
+
+      context 'when the status has follower-only replies from an unfollowed account' do
+        before do
+          Fabricate(:status, in_reply_to_id: status.id, account: replier, visibility: :direct)
+        end
+
+        it 'counts 0 replies' do
+          expect(subject['replies_count']).to be 0
+        end
+      end
+
+      context 'when the replies are public' do
+        before do
+          Fabricate(:status, in_reply_to_id: status.id, account: replier, visibility: :public)
+        end
+
+        it 'counts 1 reply' do
+          expect(subject['replies_count']).to eq 1
+        end
+      end
+
+      context 'when there is one public reply and one private' do
+        before do
+          %i[direct public].each do |visibility|
+            Fabricate(:status, in_reply_to_id: status.id, account: replier, visibility: visibility)
+          end
+        end
+
+        it 'counts 1 reply' do
+          expect(subject['replies_count']).to eq 1
+        end
+      end
+    end
+  end
 end