Fixed awful implementation from previous commit