Implement limit volume tracking

lib/rex/matcher.rb

@@ -7,11 +7,8 @@ module Rex
       return trades if highest_buy_order.nil? || lowest_sell_order.nil?
 
       while highest_buy_order.price >= lowest_sell_order.price
-        max_quantity = [highest_buy_order.remaining_quantity, lowest_sell_order.remaining_quantity].min
-        highest_buy_order.remaining_quantity -= max_quantity
-        lowest_sell_order.remaining_quantity -= max_quantity
-
-        trades << Trade.new(
+        max_quantity = min(highest_buy_order.remaining_quantity, lowest_sell_order.remaining_quantity)
+        trade = Trade.new(
           id: order_book.next_trade_id,
           buy_order: highest_buy_order,
           sell_order: lowest_sell_order,
@@ -19,8 +16,8 @@ module Rex
           price: lowest_sell_order.price
         )
 
-        remove_if_filled(highest_buy_order, order_book)
-        remove_if_filled(lowest_sell_order, order_book)
+        order_book.process_trade(trade)
+        trades.push(trade)
 
         # Go for the next run
         highest_buy_order = order_book.highest_buy_order
@@ -34,9 +31,9 @@ module Rex
 
     private
 
-    def remove_if_filled(order, order_book)
-      return unless order.filled?
-      order_book.remove_order(order.id)
+    def min(a, b)
+      return a if a < b
+      b
     end
   end
 end

lib/rex/order_book.rb

@@ -2,6 +2,10 @@ require "rbtree"
 
 module Rex
   class OrderBook
+    attr_reader(
+      :buy_limit_volumes,
+      :sell_limit_volumes
+    )
     def initialize(matcher: Matcher.new)
       @matcher = matcher
       @sell_side = RBTree.new
@@ -9,14 +13,19 @@ module Rex
       @order_ids = {} # order_id => order
       @current_trade_id = 0
       @current_order_id = 0
+      @buy_limit_volumes = {} # price => volume
+      @sell_limit_volumes = {} # price => volume
     end
 
     def add_order(order)
       side = side_for_order(order)
+      limit_volumes = limit_volume_for_order_side(order)
       order.id = next_order_id
 
       side[order.price] ||= Limit.new(order.price)
       side[order.price].add_order(order)
+      limit_volumes[order.price] ||= 0
+      limit_volumes[order.price] += order.remaining_quantity
       order_ids[order.id] = order
     end
 
@@ -30,18 +39,32 @@ module Rex
       return nil if order.nil?
 
       side = side_for_order(order)
-
       limit = side[order.price]
       limit.remove_order(order)
+
       if limit.empty?
         side.delete(limit.price)
       end
 
+      limit_volumes = limit_volume_for_order_side(order)
+      limit_volumes[order.price] -= order.remaining_quantity
+
       order_ids.delete(order.id)
     end
 
     alias_method :cancel_order, :remove_order
 
+    def process_trade(trade)
+      trade.buy_order.remaining_quantity -= trade.quantity
+      trade.sell_order.remaining_quantity -= trade.quantity
+
+      buy_limit_volumes[trade.buy_order.price] -= trade.quantity
+      sell_limit_volumes[trade.sell_order.price] -= trade.quantity
+
+      remove_order(trade.buy_order.id) if trade.buy_order.filled?
+      remove_order(trade.sell_order.id) if trade.sell_order.filled?
+    end
+
     def highest_buy_order
       buy_side.last&.[](1)&.peek_first_order
     end
@@ -64,12 +87,24 @@ module Rex
 
     private
 
-    attr_reader :order_ids, :buy_side, :sell_side
+    attr_reader(
+      :order_ids,
+      :buy_side,
+      :sell_side
+    )
 
     def next_order_id
       @current_order_id += 1
     end
 
+    def limit_volume_for_order_side(order)
+      if order.is_buy
+        @buy_limit_volumes
+      else
+        @sell_limit_volumes
+      end
+    end
+
     def side_for_order(order)
       if order.is_buy
         buy_side

spec/matcher_spec.rb

@@ -41,6 +41,13 @@ RSpec.describe Rex::Matcher do
         expect(order_book.lowest_sell_order).to eq(pricier_sell_order)
         expect(order_book.lowest_sell_order.remaining_quantity).to eq(20)
       end
+
+      it "adjusts the limit volumes" do
+        instance.match(order_book)
+
+        expect(order_book.buy_limit_volumes).to eq({100 => 0})
+        expect(order_book.sell_limit_volumes).to eq({99 => 0, 100 => 20})
+      end
     end
 
     context "when order book is empty" do

spec/order_book_spec.rb

@@ -13,6 +13,13 @@ RSpec.describe Rex::OrderBook do
       expect(instance.best_buy_price).to eq(100)
     end
 
+    it "adjusts the limit volume on the correct side" do
+      order = build(:order, is_buy: true, quantity: 100, remaining_quantity: 99)
+      instance.add_order(order)
+
+      expect(instance.buy_limit_volumes[order.price]).to eq(99)
+    end
+
     it "assigns an order id" do
       expect { instance.add_order(order_a) }.to change(order_a, :id).from(nil).to(1)
       expect { instance.add_order(order_b) }.to change(order_b, :id).from(nil).to(2)
@@ -143,6 +150,13 @@ RSpec.describe Rex::OrderBook do
         instance.add_order(sell_order)
       end
 
+      it "adjusts the limit's volume" do
+        expect do
+          instance.cancel_order(sell_order.id)
+        end.to(change { instance.sell_limit_volumes[sell_order.price] }
+          .by(-sell_order.remaining_quantity))
+      end
+
       it "does not affect the buy side" do
         instance.cancel_order(sell_order.id)