5 files changed, 94 insertions, 15 deletions
diff --git a/app/controllers/api/v1/statuses_controller.rb b/app/controllers/api/v1/statuses_controller.rb
index e98241323..01880565c 100644
--- a/app/controllers/api/v1/statuses_controller.rb
+++ b/app/controllers/api/v1/statuses_controller.rb
@@ -18,7 +18,7 @@ class Api::V1::StatusesController < Api::BaseController
 
   def context
     ancestors_results   = @status.in_reply_to_id.nil? ? [] : @status.ancestors(DEFAULT_STATUSES_LIMIT, current_account)
-    descendants_results = @status.descendants(current_account)
+    descendants_results = @status.descendants(DEFAULT_STATUSES_LIMIT, current_account)
     loaded_ancestors    = cache_collection(ancestors_results, Status)
     loaded_descendants  = cache_collection(descendants_results, Status)
 
diff --git a/app/controllers/statuses_controller.rb b/app/controllers/statuses_controller.rb
index a2943982a..01dac35e4 100644
--- a/app/controllers/statuses_controller.rb
+++ b/app/controllers/statuses_controller.rb
@@ -5,6 +5,8 @@ class StatusesController < ApplicationController
   include Authorization
 
   ANCESTORS_LIMIT = 20
+  DESCENDANTS_LIMIT = 20
+  DESCENDANTS_DEPTH_LIMIT = 4
 
   layout 'public'
 
@@ -19,9 +21,8 @@ class StatusesController < ApplicationController
   def show
     respond_to do |format|
       format.html do
-        @ancestors     = @status.reply? ? cache_collection(@status.ancestors(ANCESTORS_LIMIT, current_account), Status) : []
-        @next_ancestor = @ancestors.size < ANCESTORS_LIMIT ? nil : @ancestors.shift
-        @descendants   = cache_collection(@status.descendants(current_account), Status)
+        set_ancestors
+        set_descendants
 
         render 'stream_entries/show'
       end
@@ -51,10 +52,76 @@ class StatusesController < ApplicationController
 
   private
 
+  def create_descendant_thread(depth, statuses)
+    if depth < DESCENDANTS_DEPTH_LIMIT
+      { statuses: statuses }
+    else
+      next_status = statuses.pop
+      { statuses: statuses, next_status: next_status }
+    end
+  end
+
   def set_account
     @account = Account.find_local!(params[:account_username])
   end
 
+  def set_ancestors
+    @ancestors     = @status.reply? ? cache_collection(@status.ancestors(ANCESTORS_LIMIT, current_account), Status) : []
+    @next_ancestor = @ancestors.size < ANCESTORS_LIMIT ? nil : @ancestors.shift
+  end
+
+  def set_descendants
+    @max_descendant_thread_id = params[:max_descendant_thread_id]&.to_i
+    @since_descendant_thread_id = params[:since_descendant_thread_id]&.to_i
+
+    descendants = cache_collection(
+      @status.descendants(
+        DESCENDANTS_LIMIT,
+        current_account,
+        @max_descendant_thread_id,
+        @since_descendant_thread_id,
+        DESCENDANTS_DEPTH_LIMIT
+      ),
+      Status
+    )
+    @descendant_threads = []
+
+    if descendants.present?
+      statuses = [descendants.first]
+      depth = 1
+
+      descendants.drop(1).each_with_index do |descendant, index|
+        if descendants[index].id == descendant.in_reply_to_id
+          depth += 1
+          statuses << descendant
+        else
+          @descendant_threads << create_descendant_thread(depth, statuses)
+
+          @descendant_threads.reverse_each do |descendant_thread|
+            statuses = descendant_thread[:statuses]
+
+            index = statuses.find_index do |thread_status|
+              thread_status.id == descendant.in_reply_to_id
+            end
+
+            if index.present?
+              depth += index - statuses.size
+              break
+            end
+
+            depth -= statuses.size
+          end
+
+          statuses = [descendant]
+        end
+      end
+
+      @descendant_threads << create_descendant_thread(depth, statuses)
+    end
+
+    @max_descendant_thread_id = @descendant_threads.pop[:statuses].first.id if descendants.size >= DESCENDANTS_LIMIT
+  end
+
   def set_link_headers
     response.headers['Link'] = LinkHeader.new(
       [
diff --git a/app/models/concerns/status_threading_concern.rb b/app/models/concerns/status_threading_concern.rb
index fffc095ee..a3fd34e94 100644
--- a/app/models/concerns/status_threading_concern.rb
+++ b/app/models/concerns/status_threading_concern.rb
@@ -7,8 +7,8 @@ module StatusThreadingConcern
     find_statuses_from_tree_path(ancestor_ids(limit), account)
   end
 
-  def descendants(account = nil)
-    find_statuses_from_tree_path(descendant_ids, account)
+  def descendants(limit, account = nil, max_child_id = nil, since_child_id = nil, depth = nil)
+    find_statuses_from_tree_path(descendant_ids(limit, max_child_id, since_child_id, depth), account)
   end
 
   private
@@ -46,26 +46,27 @@ module StatusThreadingConcern
     SQL
   end
 
-  def descendant_ids
-    descendant_statuses.pluck(:id)
+  def descendant_ids(limit, max_child_id, since_child_id, depth)
+    descendant_statuses(limit, max_child_id, since_child_id, depth).pluck(:id)
   end
 
-  def descendant_statuses
-    Status.find_by_sql([<<-SQL.squish, id: id])
+  def descendant_statuses(limit, max_child_id, since_child_id, depth)
+    Status.find_by_sql([<<-SQL.squish, id: id, limit: limit, max_child_id: max_child_id, since_child_id: since_child_id, depth: depth])
       WITH RECURSIVE search_tree(id, path)
       AS (
         SELECT id, ARRAY[id]
         FROM statuses
-        WHERE in_reply_to_id = :id
+        WHERE in_reply_to_id = :id AND COALESCE(id < :max_child_id, TRUE) AND COALESCE(id > :since_child_id, TRUE)
         UNION ALL
         SELECT statuses.id, path || statuses.id
         FROM search_tree
         JOIN statuses ON statuses.in_reply_to_id = search_tree.id
-        WHERE NOT statuses.id = ANY(path)
+        WHERE COALESCE(array_length(path, 1) < :depth, TRUE) AND NOT statuses.id = ANY(path)
       )
       SELECT id
       FROM search_tree
       ORDER BY path
+      LIMIT :limit
     SQL
   end
 
diff --git a/app/views/stream_entries/_more.html.haml b/app/views/stream_entries/_more.html.haml
new file mode 100644
index 000000000..9b1dfe4a7
--- /dev/null
+++ b/app/views/stream_entries/_more.html.haml
@@ -0,0 +1,2 @@
+= link_to url, class: 'more light'  do
+  = t('statuses.show_more')
diff --git a/app/views/stream_entries/_status.html.haml b/app/views/stream_entries/_status.html.haml
index 2d0dafcb7..8decdf6b5 100644
--- a/app/views/stream_entries/_status.html.haml
+++ b/app/views/stream_entries/_status.html.haml
@@ -16,8 +16,7 @@
 - if status.reply? && include_threads
   - if @next_ancestor
     .entry{ class: entry_classes }
-      = link_to short_account_status_url(@next_ancestor.account.username, @next_ancestor), class: 'more light'  do
-        = t('statuses.show_more')
+      = render 'stream_entries/more', url: short_account_status_url(@next_ancestor.account.username, @next_ancestor)
   = render partial: 'stream_entries/status', collection: @ancestors, as: :status, locals: { is_predecessor: true, direct_reply_id: status.in_reply_to_id }
 
 .entry{ class: entry_classes }
@@ -40,4 +39,14 @@
   = render (centered ? 'stream_entries/detailed_status' : 'stream_entries/simple_status'), status: status.proper
 
 - if include_threads
-  = render partial: 'stream_entries/status', collection: @descendants, as: :status, locals: { is_successor: true, parent_id: status.id }
+  - if @since_descendant_thread_id
+    .entry{ class: entry_classes }
+      = render 'stream_entries/more', url: short_account_status_url(status.account.username, status, max_descendant_thread_id: @since_descendant_thread_id + 1)
+  - @descendant_threads.each do |thread|
+    = render partial: 'stream_entries/status', collection: thread[:statuses], as: :status, locals: { is_successor: true, parent_id: status.id }
+    - if thread[:next_status]
+      .entry{ class: entry_classes }
+        = render 'stream_entries/more', url: short_account_status_url(thread[:next_status].account.username, thread[:next_status])
+  - if @next_descendant_thread
+    .entry{ class: entry_classes }
+      = render 'stream_entries/more', url: short_account_status_url(status.account.username, status, since_descendant_thread_id: @max_descendant_thread_id - 1)