diff --git a/lib/rex/order_book.rb b/lib/rex/order_book.rb index d27f019..ed2c124 100644 --- a/lib/rex/order_book.rb +++ b/lib/rex/order_book.rb @@ -3,29 +3,62 @@ require "rbtree" module Rex class OrderBook def initialize - @sell_side = tree_with_limit_default - @buy_side = tree_with_limit_default + @sell_side = RBTree.new + @buy_side = RBTree.new + @order_ids = {} # order_id => order end def add_order(order) - side = if order.is_buy - @buy_side - else - @sell_side - end + side = side_for_order(order) + 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 - tree = RBTree.new - tree.default_proc = ->(tree, key) { tree[key] = Limit.new(key) } - tree + attr_reader :order_ids, :buy_side, :sell_side + + def side_for_order(order) + if order.is_buy + buy_side + else + sell_side + end end end end diff --git a/spec/order_book_spec.rb b/spec/order_book_spec.rb new file mode 100644 index 0000000..cdafc0f --- /dev/null +++ b/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