Refactor search code