Fix for highlighting issue
Code cleanup