Improvements per review comments