Improve while loop logic

No more infinite loops