Cosmetic fixes - minimal refactor