Performance improvements in holding calculation pipeline

app/models/holding/forward_calculator.rb

@@ -18,7 +18,7 @@ def calculate
         trades = portfolio_cache.get_trades(date: date)
         update_cost_basis_tracker(trades)
         next_portfolio = transform_portfolio(current_portfolio, trades, direction: :forward)
-        holdings += build_holdings(next_portfolio, date)
+        holdings.concat(build_holdings(next_portfolio, date))
         current_portfolio = next_portfolio
       end
 

app/models/holding/gapfillable.rb

@@ -9,10 +9,11 @@ def gapfill(holdings)
         next if security_holdings.empty?
 
         sorted = security_holdings.sort_by(&:date)
+        holdings_by_date = security_holdings.index_by(&:date)
         previous_holding = sorted.first
 
         sorted.first.date.upto(Date.current) do |date|
-          holding = security_holdings.find { |h| h.date == date }
+          holding = holdings_by_date[date]
 
           if holding
             filled_holdings << holding

app/models/holding/materializer.rb

@@ -171,7 +171,7 @@ def holding_key(holding)
     end
 
     def purge_stale_holdings
-      portfolio_security_ids = account.entries.trades.map { |entry| entry.entryable.security_id }.uniq
+      portfolio_security_ids = account.trades.distinct.pluck(:security_id)
 
       # Never delete provider-sourced holdings - they're authoritative from the provider
       # If there are no securities in the portfolio, only delete non-provider holdings

app/models/holding/portfolio_cache.rb

@@ -16,22 +16,25 @@ def initialize(account, use_holdings: false, security_ids: nil)
 
   def get_trades(date: nil)
     if date.blank?
-      trades
+      trades.dup
     else
-      trades.select { |t| t.date == date }
+      trades_by_date[date]&.dup || []
     end
   end
 
   def get_price(security_id, date, source: nil)
     security = @security_cache[security_id]
     raise SecurityNotFound.new(security_id, account.id) unless security
 
-    if source.present?
-      price = security[:prices].select { |p| p.price.date == date && p.source == source }.min_by(&:priority)&.price
+    price_with_priority = if source.present?
+      security[:prices_by_date_and_source][[ date, source ]]
     else
-      price = security[:prices].select { |p| p.price.date == date }.min_by(&:priority)&.price
+      security[:prices_by_date][date]
     end
 
+    return nil unless price_with_priority
+
+    price = price_with_priority.price
     return nil unless price
 
     price_money = Money.new(price.price, price.currency)
@@ -61,20 +64,28 @@ def trades
       @trades ||= account.entries.includes(entryable: :security).trades.chronological.to_a
     end
 
+    def trades_by_date
+      @trades_by_date ||= trades.group_by(&:date)
+    end
+
+    def trades_by_security_id
+      @trades_by_security_id ||= trades.group_by { |t| t.entryable.security_id }
+    end
+
     def holdings
       @holdings ||= account.holdings.chronological.to_a
     end
 
-    def collect_unique_securities
-      unique_securities_from_trades = trades.map(&:entryable).map(&:security).uniq
-      unique_securities_from_trades = unique_securities_from_trades.select { |s| @security_ids.include?(s.id) } if @security_ids
-
-      return unique_securities_from_trades unless use_holdings
+    def holdings_by_security_id
+      @holdings_by_security_id ||= holdings.group_by(&:security_id)
+    end
 
-      unique_securities_from_holdings = holdings.map(&:security).uniq
-      unique_securities_from_holdings = unique_securities_from_holdings.select { |s| @security_ids.include?(s.id) } if @security_ids
+    def collect_unique_securities
+      ids = trades_by_security_id.keys
+      ids |= holdings_by_security_id.keys if use_holdings
+      ids &= @security_ids if @security_ids
 
-      (unique_securities_from_trades + unique_securities_from_holdings).uniq
+      Security.where(id: ids).to_a
     end
 
     # Loads all known prices for all securities in the account with priority based on source:
@@ -87,11 +98,19 @@ def load_prices
 
       Rails.logger.info "Preloading #{securities.size} securities for account #{account.id}"
 
+      security_ids = securities.map(&:id)
+
+      # Bulk-load all DB prices for all securities in one query, grouped by security_id
+      db_prices_by_security_id = Security::Price
+        .where(security_id: security_ids, date: account.start_date..Date.current)
+        .select(:security_id, :date, :price, :currency)
+        .group_by(&:security_id)
+
       securities.each do |security|
         Rails.logger.info "Loading security: ID=#{security.id} Ticker=#{security.ticker}"
 
         # High priority prices from DB (synced from provider)
-        db_prices = security.prices.where(date: account.start_date..Date.current).map do |price|
+        db_prices = (db_prices_by_security_id[security.id] || []).map do |price|
           PriceWithPriority.new(
             price: price,
             priority: 1,
@@ -100,8 +119,7 @@ def load_prices
         end
 
         # Medium priority prices from trades
-        trade_prices = trades
-          .select { |t| t.entryable.security_id == security.id }
+        trade_prices = (trades_by_security_id[security.id] || [])
           .map do |trade|
             PriceWithPriority.new(
               price: Security::Price.new(
@@ -117,7 +135,7 @@ def load_prices
 
         # Low priority prices from holdings (if applicable)
         holding_prices = if use_holdings
-          holdings.select { |h| h.security_id == security.id }.map do |holding|
+          (holdings_by_security_id[security.id] || []).map do |holding|
             PriceWithPriority.new(
               price: Security::Price.new(
                 security: security,
@@ -133,9 +151,18 @@ def load_prices
           []
         end
 
+        all_prices = db_prices + trade_prices + holding_prices
+
+        # Index by date for O(1) lookup in get_price instead of O(N) linear scan
+        prices_by_date = all_prices.group_by { |p| p.price.date }
+          .transform_values { |ps| ps.min_by(&:priority) }
+        prices_by_date_and_source = all_prices.group_by { |p| [ p.price.date, p.source ] }
+          .transform_values { |ps| ps.min_by(&:priority) }
+
         @security_cache[security.id] = {
           security: security,
-          prices: db_prices + trade_prices + holding_prices
+          prices_by_date: prices_by_date,
+          prices_by_date_and_source: prices_by_date_and_source
         }
       end
     end

app/models/holding/reverse_calculator.rb

@@ -35,7 +35,7 @@ def calculate_holdings
         previous_portfolio = transform_portfolio(current_portfolio, today_trades, direction: :reverse)
 
         # If current day, always use holding prices (since that's what Plaid gives us).  For historical values, use market data (since Plaid doesn't supply historical prices)
-        holdings += build_holdings(current_portfolio, date, price_source: date == Date.current ? "holding" : nil)
+        holdings.concat(build_holdings(current_portfolio, date, price_source: date == Date.current ? "holding" : nil))
         current_portfolio = previous_portfolio
       end
 
@@ -79,45 +79,46 @@ def build_holdings(portfolio, date, price_source: nil)
       end.compact
     end
 
-    # Pre-compute cost basis for all securities at all dates using forward pass through trades
-    # Stores: { security_id => { date => cost_basis } }
     def precompute_cost_basis
-      @cost_basis_by_date = Hash.new { |h, k| h[k] = {} }
+      @cost_basis_snapshots = Hash.new { |h, k| h[k] = [] }
       tracker = Hash.new { |h, k| h[k] = { total_cost: BigDecimal("0"), total_qty: BigDecimal("0") } }
 
-      trades = portfolio_cache.get_trades.sort_by(&:date)
-      trade_index = 0
-
-      account.start_date.upto(Date.current).each do |date|
-        # Process all trades up to and including this date
-        while trade_index < trades.size && trades[trade_index].date <= date
-          trade_entry = trades[trade_index]
-          trade = trade_entry.entryable
-
-          if trade.qty > 0 # Only track buys
-            security_id = trade.security_id
-            trade_price = Money.new(trade.price, trade.currency)
-            begin
-              converted_price = trade_price.exchange_to(account.currency).amount
-            rescue Money::ConversionError
-              converted_price = trade.price
-            end
-
-            tracker[security_id][:total_cost] += converted_price * trade.qty
-            tracker[security_id][:total_qty] += trade.qty
-          end
-          trade_index += 1
-        end
+      portfolio_cache.get_trades.sort_by(&:date).each do |trade_entry|
+        trade = trade_entry.entryable
+        next unless trade.qty > 0
 
-        # Store current cost basis snapshot for each security at this date
-        tracker.each do |security_id, data|
-          next if data[:total_qty].zero?
-          @cost_basis_by_date[security_id][date] = data[:total_cost] / data[:total_qty]
+        security_id = trade.security_id
+        trade_price = Money.new(trade.price, trade.currency)
+        begin
+          converted_price = trade_price.exchange_to(account.currency).amount
+        rescue Money::ConversionError
+          converted_price = trade.price
         end
+
+        tracker[security_id][:total_cost] += converted_price * trade.qty
+        tracker[security_id][:total_qty] += trade.qty
+
+        @cost_basis_snapshots[security_id] << [
+          trade_entry.date,
+          tracker[security_id][:total_cost] / tracker[security_id][:total_qty]
+        ]
       end
     end
 
     def cost_basis_for(security_id, date)
-      @cost_basis_by_date.dig(security_id, date)
+      snapshots = @cost_basis_snapshots[security_id]
+      return nil if snapshots.empty?
+
+      lo, hi, result = 0, snapshots.size - 1, nil
+      while lo <= hi
+        mid = (lo + hi) / 2
+        if snapshots[mid][0] <= date
+          result = snapshots[mid][1]
+          lo = mid + 1
+        else
+          hi = mid - 1
+        end
+      end
+      result
     end
 end

test/models/holding/reverse_calculator_test.rb

@@ -158,6 +158,84 @@ class Holding::ReverseCalculatorTest < ActiveSupport::TestCase
     end
   end
 
+  # cost_basis_for
+
+  test "cost_basis_for returns nil when there are no buy trades" do
+    security = Security.create!(ticker: "TST", name: "Test")
+    calc = calculator_with_trades(security)
+
+    assert_nil cost_basis_for(calc, security, Date.current)
+  end
+
+  test "cost_basis_for returns nil for dates before the first buy" do
+    security = Security.create!(ticker: "TST", name: "Test")
+    buy_date = 5.days.ago.to_date
+
+    calc = calculator_with_trades(security) do
+      create_trade(security, account: @account, qty: 10, price: 100, date: buy_date)
+    end
+
+    assert_nil cost_basis_for(calc, security, buy_date - 1)
+  end
+
+  test "cost_basis_for returns weighted average cost on buy date" do
+    security = Security.create!(ticker: "TST", name: "Test")
+    buy_date = 5.days.ago.to_date
+
+    calc = calculator_with_trades(security) do
+      create_trade(security, account: @account, qty: 10, price: 100, date: buy_date)
+    end
+
+    assert_in_delta 100.0, cost_basis_for(calc, security, buy_date).to_f, 1e-6
+  end
+
+  test "cost_basis_for carries forward to dates between buys" do
+    security = Security.create!(ticker: "TST", name: "Test")
+    first_buy  = 10.days.ago.to_date
+    second_buy = 3.days.ago.to_date
+
+    calc = calculator_with_trades(security) do
+      create_trade(security, account: @account, qty: 10, price: 100, date: first_buy)
+      create_trade(security, account: @account, qty: 5,  price: 130, date: second_buy)
+    end
+
+    # Between the two buys, cost basis is from the first buy only
+    assert_in_delta 100.0, cost_basis_for(calc, security, first_buy + 1).to_f, 1e-6
+    assert_in_delta 100.0, cost_basis_for(calc, security, second_buy - 1).to_f, 1e-6
+
+    # After second buy: WAC = (10*100 + 5*130) / 15 = 110.0
+    assert_in_delta 110.0, cost_basis_for(calc, security, second_buy).to_f, 1e-6
+    assert_in_delta 110.0, cost_basis_for(calc, security, Date.current).to_f, 1e-6
+  end
+
+  test "cost_basis_for accumulates multiple buys on the same date" do
+    security = Security.create!(ticker: "TST", name: "Test")
+    buy_date = 5.days.ago.to_date
+
+    calc = calculator_with_trades(security) do
+      create_trade(security, account: @account, qty: 10, price: 100, date: buy_date)
+      create_trade(security, account: @account, qty: 5,  price: 130, date: buy_date)
+    end
+
+    # WAC = (10*100 + 5*130) / 15 = 110.0 — not the intermediate value after only the first trade
+    assert_in_delta 110.0, cost_basis_for(calc, security, buy_date).to_f, 1e-6
+  end
+
+  test "cost_basis_for ignores sell trades" do
+    security = Security.create!(ticker: "TST", name: "Test")
+    buy_date  = 10.days.ago.to_date
+    sell_date = 5.days.ago.to_date
+
+    calc = calculator_with_trades(security) do
+      create_trade(security, account: @account, qty: 10,  price: 100, date: buy_date)
+      create_trade(security, account: @account, qty: -5,  price: 120, date: sell_date)
+    end
+
+    # Sell does not change cost basis
+    assert_in_delta 100.0, cost_basis_for(calc, security, sell_date).to_f, 1e-6
+    assert_in_delta 100.0, cost_basis_for(calc, security, Date.current).to_f, 1e-6
+  end
+
   private
     def assert_holdings(expected, calculated)
       expected.each do |expected_entry|
@@ -202,6 +280,18 @@ def load_prices
     # Brokerage Cash: $5,000
     # Holdings Value: $15,000
     # Total Balance: $20,000
+    def calculator_with_trades(security)
+      yield if block_given?
+      snapshot = OpenStruct.new(to_h: { security.id => 10 })
+      calc = Holding::ReverseCalculator.new(@account, portfolio_snapshot: snapshot)
+      calc.send(:precompute_cost_basis)
+      calc
+    end
+
+    def cost_basis_for(calc, security, date)
+      calc.send(:cost_basis_for, security.id, date)
+    end
+
     def load_today_portfolio
       @account.update!(cash_balance: 5000)