Finish basic order book implementation

lib/rex/order_book.rb

@@ -3,29 +3,62 @@ require "rbtree"
 module Rex
   class OrderBook
     def initialize
-      @sell_side = tree_with_limit_default
+      @sell_side = RBTree.new
-      @buy_side = tree_with_limit_default
+      @buy_side = RBTree.new
+      @order_ids = {} # order_id => order
     end
 
     def add_order(order)
-      side = if order.is_buy
+      side = side_for_order(order)
-        @buy_side
-      else
-        @sell_side
-      end
 
+      side[order.price] ||= Limit.new(order.price)
       side[order.price].add_order(order)
+      order_ids[order.id] = order
     end
 
-    def cancel_order(order_id)
+    def remove_order(order_id)
+      order = order_ids[order_id]
+      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
+
+      order_ids.delete(order.id)
+    end
+
+    alias_method :cancel_order, :remove_order
+
+    def highest_buy_order
+      buy_side.last&.[](1)&.peek_first_order
+    end
+
+    def lowest_sell_order
+      sell_side.first&.[](1)&.peek_first_order
+    end
+
+    def best_buy_price
+      buy_side.last&.[](0)
+    end
+
+    def best_sell_price
+      sell_side.first&.[](0)
     end
 
     private
 
-    def tree_with_limit_default
+    attr_reader :order_ids, :buy_side, :sell_side
-      tree = RBTree.new
+
-      tree.default_proc = ->(tree, key) { tree[key] = Limit.new(key) }
+    def side_for_order(order)
-      tree
+      if order.is_buy
+        buy_side
+      else
+        sell_side
+      end
     end
   end
 end

spec/order_book_spec.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+RSpec.describe Rex::OrderBook do
+  let(:instance) { described_class.new }
+
+  describe "#add_order" do
+    let(:order_a) { build(:order, is_buy: true, price: 100) }
+
+    it "adds the order to the order book" do
+      instance.add_order(order_a)
+
+      expect(instance.best_buy_price).to eq(100)
+    end
+
+    context "with multiple orders at the same price" do
+      let(:order_b) { build(:order, is_buy: true, price: 100) }
+
+      it "regards the new order as the best buy price" do
+        instance.add_order(order_a)
+        instance.add_order(order_b)
+
+        expect(instance.best_buy_price).to eq(100)
+        expect(instance.highest_buy_order).to eq(order_a)
+      end
+    end
+
+    context "with multiple orders at the different price levels" do
+      let(:order_b) { build(:order, is_buy: true, price: 120) }
+
+      it "regards the new order as the best buy price" do
+        instance.add_order(order_b)
+        instance.add_order(order_a)
+
+        expect(instance.best_buy_price).to eq(120)
+        expect(instance.highest_buy_order).to eq(order_b)
+      end
+    end
+
+    context "with multiple orders on different sides" do
+      let(:order_b) { build(:order, is_buy: false, price: 120) }
+
+      it "returns the correct prices for both sides" do
+        instance.add_order(order_a)
+        instance.add_order(order_b)
+
+        expect(instance.best_buy_price).to eq(100)
+        expect(instance.best_sell_price).to eq(120)
+      end
+    end
+  end
+
+  describe "#highest_buy_order" do
+    context "when there is nothing in the book" do
+      it "returns nil " do
+        expect(instance.highest_buy_order).to eq(nil)
+      end
+    end
+  end
+
+  describe "#lowest_sell_Order" do
+    context "when there is nothing in the book" do
+      it "returns nil " do
+        expect(instance.highest_buy_order).to eq(nil)
+      end
+    end
+  end
+
+  describe "#cancel_order" do
+    let(:buy_order) { build(:order, is_buy: true, price: 100) }
+    let(:sell_order) { build(:order, is_buy: false, price: 121) }
+
+    before do
+      instance.add_order(buy_order)
+    end
+
+    context "when it is the last order for the given price" do
+      it "removes the order from the side" do
+        instance.cancel_order(buy_order.id)
+
+        expect(instance.best_buy_price).to eq(nil)
+      end
+    end
+
+    context "an order on the sell side is removed" do
+      before do
+        instance.add_order(buy_order)
+        instance.add_order(sell_order)
+      end
+
+      it "does not affect the buy side" do
+        instance.cancel_order(sell_order.id)
+
+        expect(instance.best_buy_price).to eq(100)
+      end
+    end
+  end
+end