fix edge-case logic