Quick refactor

Cuz I can't help myself