diff --git a/lib/rex/book/limit_order_book.rb b/lib/rex/book/limit_order_book.rb index 8d9effc..005ebff 100644 --- a/lib/rex/book/limit_order_book.rb +++ b/lib/rex/book/limit_order_book.rb @@ -8,6 +8,7 @@ module Rex @sell_side = RBTree.new @buy_side = RBTree.new @order_ids = {} # order_id => order + @user_order_ids = {} # user_id => Set @current_order_id = 0 end @@ -19,6 +20,8 @@ module Rex side[order.price].add_order(order) order_ids[order.id] = order + user_order_ids[order.user_id] ||= Set.new + user_order_ids[order.user_id].add(order.id) end def add_and_match_order(order) @@ -43,6 +46,10 @@ module Rex end order_ids.delete(order.id) + + user_order_ids[order.user_id].delete(order.id) + user_order_ids.delete(order.user_id) if user_order_ids[order.user_id].empty? + order end @@ -72,12 +79,19 @@ module Rex sell_side.first&.[](0) end + def orders_for_user(user_id) + (@user_order_ids[user_id] || []).map do |order_id| + order_ids[order_id] + end + end + private attr_reader( :order_ids, :buy_side, - :sell_side + :sell_side, + :user_order_ids ) def next_order_id diff --git a/spec/book/limit_order_book_spec.rb b/spec/book/limit_order_book_spec.rb index c11dce8..c4e6782 100644 --- a/spec/book/limit_order_book_spec.rb +++ b/spec/book/limit_order_book_spec.rb @@ -101,6 +101,39 @@ RSpec.describe Rex::Book::LimitOrderBook do end end + describe "#orders_for_user" do + context "when user does not have any orders" do + it "returns an empty list" do + expect(instance.orders_for_user("unknown_id")).to eq([]) + end + end + + context "when user has orders" do + let(:order) { build(:order, id: nil, is_buy: true, price: 100) } + + before do + instance.add_order(order) + end + + it "returns a list of the user's orders" do + expect(instance.orders_for_user(order.user_id)).to eq([order]) + end + end + + context "when the user had orders" do + let(:order) { build(:order, id: nil, is_buy: true, price: 100) } + + before do + instance.add_order(order) + instance.cancel_order(order.id) + end + + it "returns an empty list" do + expect(instance.orders_for_user(order.user_id)).to eq([]) + 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) }