Small refactoring. Seems like it's working.