Improving handling of comments on more edge cases