fix edge cases :+1: