Further improve highlight plugin

Handle initial search use-case.