Code changes to improve performance and add tests