Code improvement and bug fix.