Small fixes and optimization