Bug fixes and more code optimizations